[SSH] Bump Unity Gateway CLI pin

experimental/ssh/internal/client/agentshim.go

@@ -47,7 +47,9 @@ const (
 	ugRepo = "databricks/unity-gateway"
 
 	// pins the Unity Gateway CLI commit the shim installs.
-	ugCommit = "7f408803f9c0fe4958e6d264fd94595283a49a19" // v0.1.0
+	ugCommit = "034ea0fd11f8e1286b7dd3d945d189be4ff1f42d"
+
+	ugExecutable = "ug"
 
 	// pins the version of uv downloaded
 	uvBaseURL = "https://github.com/astral-sh/uv/releases/download/0.12.18/"
@@ -68,6 +70,8 @@ const (
 	nodeX64LinuxChecksum   = "fd8e59d5a511510f6a298afb548f18c7d2b1be404d8b4a27d94fbe49f56cb2d6"
 	nodeArm64LinuxChecksum = "6ad1325edbdb5649c379b75a237147a666c95d4f9ae8d340fef2d1575d289ad2"
 
+	npmExecutable = "npm"
+
 	// environment variable used to pass the user's workspace home to the shim for the agent context.
 	workspaceHomeEnv = "DATABRICKS_WORKSPACE_HOME"
 
@@ -153,13 +157,23 @@ func RunAgentShim(ctx context.Context, client *databricks.WorkspaceClient, agent
 	if err != nil {
 		return categorize(protos.SshAgentShimErrorCategoryToolchainSetupFailed, fmt.Errorf("failed to resolve home directory: %w", err))
 	}
-	// bootstrapToolchain attributes the per-component install failure; TOOLCHAIN_SETUP_FAILED
-	// is the fallback for a lock or PATH problem.
-	if err := bootstrapToolchain(ctx, home); err != nil {
+
+	ugPath, err := bootstrapToolchain(ctx, home)
+	if err != nil {
 		return ensureCategory(protos.SshAgentShimErrorCategoryToolchainSetupFailed, err)
 	}
 
-	argv0, argv, environ, err := prepareAgentLaunch(ctx, home, agent, workspace, agentArgs)
+	err = setEnvForUnityGateway(ctx)
+	if err != nil {
+		return categorize(protos.SshAgentShimErrorCategoryAgentConfigurationFailed, fmt.Errorf("failed to set environment variables: %w", err))
+	}
+
+	err = runAgentConfiguration(ctx, agent, workspace)
+	if err != nil {
+		return ensureCategory(protos.SshAgentShimErrorCategoryAgentConfigurationFailed, err)
+	}
+
+	argv, environ, err := prepareAgentLaunch(ctx, home, agent, agentArgs)
 	if err != nil {
 		return ensureCategory(protos.SshAgentShimErrorCategoryAgentLaunchFailed, err)
 	}
@@ -169,7 +183,7 @@ func RunAgentShim(ctx context.Context, client *databricks.WorkspaceClient, agent
 	launched = true
 	logAgentShimEvent(ctx, agentName, outcome, time.Since(start))
 	uploadAgentShimTelemetry(ctx, agentName, time.Since(start))
-	return execProcess(argv0, argv, environ)
+	return execProcess(ugPath, argv, environ)
 }
 
 // agentShimOutcome is the observed result of an agent-shim launch, collected by
@@ -457,7 +471,7 @@ func aiGatewayScopeError(host, reason string) error {
 
 // --- end Unity AI Gateway preflight ---
 
-func bootstrapToolchain(ctx context.Context, home string) error {
+func bootstrapToolchain(ctx context.Context, home string) (string, error) {
 	// Reconstruct the PATH tooling runs under on every launch rather than caching
 	// it: every entry is a known location, so deriving it can't restore a stale
 	// path (e.g. a versioned CLI dir from an older session). Drop the shim dir so
@@ -470,19 +484,24 @@ func bootstrapToolchain(ctx context.Context, home string) error {
 
 	err := ensureToolchain(ctx, home, runtime.GOARCH)
 	if err != nil {
-		return err
+		return "", err
 	}
 
 	// npm is on PATH now (shipped or just installed); silence its update-notifier
 	// box so it doesn't clutter the agent session.
 	disableNpmUpdateNotifier(ctx)
 
-	return nil
+	ugPath, err := exec.LookPath(ugExecutable)
+	if err != nil {
+		return "", categorize(protos.SshAgentShimErrorCategoryUnityGatewayCLINotFound, fmt.Errorf("the Unity Gateway CLI was not found on PATH after setup: %w", err))
+	}
+
+	return ugPath, nil
 }
 
 func toolchainReady() bool {
-	_, ugErr := exec.LookPath("ucode")
-	_, npmErr := exec.LookPath("npm")
+	_, ugErr := exec.LookPath(ugExecutable)
+	_, npmErr := exec.LookPath(npmExecutable)
 	return ugErr == nil && npmErr == nil
 }
 
@@ -522,7 +541,7 @@ func ensureToolchain(ctx context.Context, home, arch string) error {
 		}
 	}
 
-	if _, err := exec.LookPath("ucode"); err != nil {
+	if _, err := exec.LookPath(ugExecutable); err != nil {
 		if err := runStep(ctx, "Installing Unity Gateway CLI", func(out io.Writer) error {
 			env := os.Environ()
 			env = append(env, "UV_TOOL_DIR="+uvToolDirAbsolute, "UV_TOOL_BIN_DIR="+uvToolBinDirAbsolute)
@@ -648,31 +667,46 @@ func acquireSetupLock(ctx context.Context, home string) (func(), error) {
 	}
 }
 
-// prepareAgentLaunch resolves the Unity Gateway CLI and builds the argv and environment for
-// launching the agent. It stops short of the exec so the caller can emit telemetry first.
-func prepareAgentLaunch(ctx context.Context, home string, agent agentSpec, workspace string, agentArgs []string) (argv0 string, argv, environ []string, err error) {
-	ugPath, err := exec.LookPath("ucode")
+func setEnvForUnityGateway(ctx context.Context) error {
+	// Pass the session token as DATABRICKS_BEARER so Unity Gateway CLI authenticates headlessly.
+	if token := env.Get(ctx, "DATABRICKS_TOKEN"); token != "" {
+		err := os.Setenv("DATABRICKS_BEARER", token)
+		if err != nil {
+			return err
+		}
+	}
+	err := os.Setenv("UG_SKIP_OS_SETTINGS", "1")
 	if err != nil {
-		return "", nil, nil, categorize(protos.SshAgentShimErrorCategoryUnityGatewayCLINotFound, fmt.Errorf("the Unity Gateway CLI was not found on PATH after setup: %w", err))
+		return err
 	}
-	contextArgs, err := injectAgentContext(ctx, home, agent)
+	err = os.Setenv("AI_GATEWAY_REQUEST_TAGS", "{\"source\":\"databricks-ssh\"}")
 	if err != nil {
-		return "", nil, nil, categorize(protos.SshAgentShimErrorCategoryAgentContextSetupFailed, err)
+		return err
+	}
+	return nil
+}
+
+func runAgentConfiguration(ctx context.Context, agent agentSpec, workspace string) error {
+	if err := runStep(ctx, "Configuring "+agent.name, func(out io.Writer) error {
+		if err := runCommand(ctx, out, nil, ugExecutable, "configure", "--agent", agent.name, "--workspace", workspace, "--enable-databricks-ai-tools"); err != nil {
+			return fmt.Errorf("failed to configure %s: %w", agent.name, err)
+		}
+		return nil
+	}); err != nil {
+		return categorize(protos.SshAgentShimErrorCategoryAgentConfigurationFailed, err)
 	}
-	argv = []string{"ucode", agent.name}
-	if workspace != "" {
-		argv = append(argv, "--workspace", workspace)
+	return nil
+}
+
+func prepareAgentLaunch(ctx context.Context, home string, agent agentSpec, agentArgs []string) (argv, environ []string, err error) {
+	contextArgs, err := injectAgentContext(ctx, home, agent)
+	if err != nil {
+		return nil, nil, categorize(protos.SshAgentShimErrorCategoryAgentContextSetupFailed, err)
 	}
-	argv = append(argv, "--")
+	argv = []string{ugExecutable, agent.name, "--"}
 	argv = append(argv, contextArgs...)
 	argv = append(argv, agentArgs...)
-	// Pass the session token as DATABRICKS_BEARER so Unity Gateway CLI authenticates headlessly.
-	if env.Get(ctx, "DATABRICKS_BEARER") == "" {
-		if token := env.Get(ctx, "DATABRICKS_TOKEN"); token != "" {
-			_ = os.Setenv("DATABRICKS_BEARER", token)
-		}
-	}
-	return ugPath, argv, os.Environ(), nil
+	return argv, os.Environ(), nil
 }
 
 // put the Databricks context where the agent reads it, returning any extra argv.

experimental/ssh/internal/client/agentshim_test.go

@@ -252,7 +252,7 @@ func TestEnsureToolchain(t *testing.T) {
 	for _, staleLock := range []bool{false, true} {
 		t.Run(fmt.Sprintf("no-op when the toolchain is already present (staleLock=%v)", staleLock), func(t *testing.T) {
 			home := t.TempDir()
-			writeFakeExec(t, filepath.Join(home, uvToolBinDir), "ucode", ":")
+			writeFakeExec(t, filepath.Join(home, uvToolBinDir), ugExecutable, ":")
 			writeFakeExec(t, filepath.Join(home, agentDepsDir, "node", "bin"), "npm", ":")
 			writeFakeExec(t, filepath.Join(home, agentDepsDir, "uv"), "uv", fakeUvRecorder)
 			if staleLock {
@@ -270,7 +270,7 @@ func TestEnsureToolchain(t *testing.T) {
 
 	t.Run("installs the Unity Gateway CLI via uv when uv is already present", func(t *testing.T) {
 		home := t.TempDir()
-		// uv and npm are present; ucode is not, so only ucode gets installed and no
+		// uv and npm are present; ug is not, so only ug gets installed and no
 		// download happens
 		writeFakeExec(t, filepath.Join(home, agentDepsDir, "uv"), "uv", fakeUvRecorder)
 		writeFakeExec(t, filepath.Join(home, agentDepsDir, "node", "bin"), "npm", ":")
@@ -283,7 +283,7 @@ func TestEnsureToolchain(t *testing.T) {
 		out, err := os.ReadFile(marker)
 		require.NoError(t, err)
 		assert.Contains(t, string(out), "tool install git+https://github.com/"+ugRepo+"@"+ugCommit)
-		// The tool bin dir is prepended to PATH so the freshly installed ucode is found.
+		// The tool bin dir is prepended to PATH so the freshly installed ug is found.
 		assert.Contains(t, filepath.SplitList(os.Getenv("PATH")), filepath.Join(home, uvToolBinDir))
 		assert.NoFileExists(t, filepath.Join(home, agentRootDir, setupLockName), "the setup lock is released once setup finishes")
 	})
@@ -305,9 +305,9 @@ func TestEnsureToolchain(t *testing.T) {
 		done := make(chan error, 1)
 		go func() { done <- ensureToolchain(ctx, home, "amd64") }()
 
-		// The other client finishes: publish ucode+npm into the PATH dir, then release
+		// The other client finishes: publish ug+npm into the PATH dir, then release
 		// the lock so the waiting ensureToolchain can proceed.
-		writeFakeExec(t, filepath.Join(home, uvToolBinDir), "ucode", ":")
+		writeFakeExec(t, filepath.Join(home, uvToolBinDir), ugExecutable, ":")
 		writeFakeExec(t, filepath.Join(home, agentDepsDir, "node", "bin"), "npm", ":")
 		unlock()
 
@@ -644,14 +644,125 @@ func TestAgentShimOutcomeCategory(t *testing.T) {
 	}
 }
 
-func TestPrepareAgentLaunchUnityGatewayCLINotFound(t *testing.T) {
-	// An empty PATH means exec.LookPath("ucode") can't find the CLI, exercising the
-	// UNITY_GATEWAY_CLI_NOT_FOUND attribution.
-	t.Setenv("PATH", t.TempDir())
+func TestBootstrapToolchain(t *testing.T) {
+	// bootstrapToolchain runs the fake POSIX executables below, and agent-shim is
+	// only intended for Linux, so skip on Windows.
+	if runtime.GOOS == "windows" {
+		return
+	}
+
+	t.Run("returns the path to the Unity Gateway CLI", func(t *testing.T) {
+		home := t.TempDir()
+		writeFakeExec(t, filepath.Join(home, uvToolBinDir), ugExecutable, ":")
+		writeFakeExec(t, filepath.Join(home, agentDepsDir, "node", "bin"), npmExecutable, ":")
+		t.Setenv("PATH", t.TempDir())
+
+		ugPath, err := bootstrapToolchain(cmdio.MockDiscard(t.Context()), home)
+		require.NoError(t, err)
+		assert.Equal(t, filepath.Join(home, uvToolBinDir, ugExecutable), ugPath)
+	})
+
+	t.Run("reports UNITY_GATEWAY_CLI_NOT_FOUND when the install leaves no CLI on PATH", func(t *testing.T) {
+		home := t.TempDir()
+		// uv "succeeds" without installing anything, so the post-setup lookup fails.
+		writeFakeExec(t, filepath.Join(home, agentDepsDir, "uv"), "uv", ":")
+		writeFakeExec(t, filepath.Join(home, agentDepsDir, "node", "bin"), npmExecutable, ":")
+		t.Setenv("PATH", t.TempDir())
+
+		_, err := bootstrapToolchain(cmdio.MockDiscard(t.Context()), home)
+		require.Error(t, err)
+		assert.Equal(t, protos.SshAgentShimErrorCategoryUnityGatewayCLINotFound, agentShimOutcome{err: err}.category())
+	})
+}
+
+func TestSetEnvForUnityGateway(t *testing.T) {
+	t.Run("passes the session token as DATABRICKS_BEARER", func(t *testing.T) {
+		t.Setenv("DATABRICKS_TOKEN", "session-token")
+		t.Setenv("DATABRICKS_BEARER", "stale-bearer")
+		t.Setenv("UG_SKIP_OS_SETTINGS", "")
+		t.Setenv("AI_GATEWAY_REQUEST_TAGS", "")
+
+		require.NoError(t, setEnvForUnityGateway(t.Context()))
+		assert.Equal(t, "session-token", os.Getenv("DATABRICKS_BEARER"))
+		assert.Equal(t, "1", os.Getenv("UG_SKIP_OS_SETTINGS"))
+		assert.JSONEq(t, `{"source":"databricks-ssh"}`, os.Getenv("AI_GATEWAY_REQUEST_TAGS"))
+	})
+
+	t.Run("leaves DATABRICKS_BEARER unset without a session token", func(t *testing.T) {
+		t.Setenv("DATABRICKS_TOKEN", "")
+		t.Setenv("DATABRICKS_BEARER", "")
+		require.NoError(t, os.Unsetenv("DATABRICKS_BEARER"))
+		t.Setenv("UG_SKIP_OS_SETTINGS", "")
+		t.Setenv("AI_GATEWAY_REQUEST_TAGS", "")
+
+		require.NoError(t, setEnvForUnityGateway(t.Context()))
+		_, ok := os.LookupEnv("DATABRICKS_BEARER")
+		assert.False(t, ok)
+		assert.Equal(t, "1", os.Getenv("UG_SKIP_OS_SETTINGS"))
+		assert.JSONEq(t, `{"source":"databricks-ssh"}`, os.Getenv("AI_GATEWAY_REQUEST_TAGS"))
+	})
+}
+
+// fakeUgRecorder is the body of a stand-in `ug` that records its args to
+// $FAKE_UG_LOG. It uses only shell builtins so it needs nothing on PATH.
+const fakeUgRecorder = `echo "args: $*" > "$FAKE_UG_LOG"`
+
+func TestRunAgentConfiguration(t *testing.T) {
+	// runAgentConfiguration runs the fake POSIX `ug` below, and agent-shim is only
+	// intended for Linux, so skip on Windows.
+	if runtime.GOOS == "windows" {
+		return
+	}
 	agent, ok := agentByName("claude")
 	require.True(t, ok)
 
-	_, _, _, err := prepareAgentLaunch(t.Context(), t.TempDir(), agent, "https://example.test", nil)
-	require.Error(t, err)
-	assert.Equal(t, protos.SshAgentShimErrorCategoryUnityGatewayCLINotFound, agentShimOutcome{err: err}.category())
+	t.Run("configures the agent for the workspace", func(t *testing.T) {
+		binDir := t.TempDir()
+		writeFakeExec(t, binDir, ugExecutable, fakeUgRecorder)
+		marker := filepath.Join(t.TempDir(), "ug.log")
+		t.Setenv("FAKE_UG_LOG", marker)
+		t.Setenv("PATH", binDir)
+
+		require.NoError(t, runAgentConfiguration(cmdio.MockDiscard(t.Context()), agent, "https://example.test"))
+
+		out, err := os.ReadFile(marker)
+		require.NoError(t, err)
+		assert.Equal(t, "args: configure --agent claude --workspace https://example.test --enable-databricks-ai-tools\n", string(out))
+	})
+
+	t.Run("a failing configure reports AGENT_CONFIGURATION_FAILED", func(t *testing.T) {
+		binDir := t.TempDir()
+		writeFakeExec(t, binDir, ugExecutable, "exit 1")
+		t.Setenv("PATH", binDir)
+
+		err := runAgentConfiguration(cmdio.MockDiscard(t.Context()), agent, "https://example.test")
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), "failed to configure claude")
+		assert.Equal(t, protos.SshAgentShimErrorCategoryAgentConfigurationFailed, agentShimOutcome{err: err}.category())
+	})
+}
+
+func TestPrepareAgentLaunch(t *testing.T) {
+	agent, ok := agentByName("claude")
+	require.True(t, ok)
+
+	t.Run("passes the context args before the user's agent args", func(t *testing.T) {
+		home := t.TempDir()
+		argv, environ, err := prepareAgentLaunch(t.Context(), home, agent, []string{"--resume"})
+		require.NoError(t, err)
+		require.NotEmpty(t, argv)
+		assert.Equal(t, []string{"claude", "--", "--append-system-prompt-file", filepath.Join(home, agentRootDir, contextFile), "--resume"}, argv[1:])
+		assert.Equal(t, os.Environ(), environ)
+	})
+
+	t.Run("a context setup failure reports AGENT_CONTEXT_SETUP_FAILED", func(t *testing.T) {
+		// A regular file where the home directory should be makes creating the
+		// context directory fail.
+		home := filepath.Join(t.TempDir(), "home")
+		require.NoError(t, os.WriteFile(home, nil, 0o644))
+
+		_, _, err := prepareAgentLaunch(t.Context(), home, agent, nil)
+		require.Error(t, err)
+		assert.Equal(t, protos.SshAgentShimErrorCategoryAgentContextSetupFailed, agentShimOutcome{err: err}.category())
+	})
 }

libs/telemetry/protos/ssh_agent_shim.go

@@ -50,13 +50,13 @@ const (
 	// Installing uv (download, checksum, or extraction) failed.
 	SshAgentShimErrorCategoryUvInstallFailed SshAgentShimErrorCategory = "UV_INSTALL_FAILED"
 
-	// Installing the Unity Gateway CLI (`ucode`, via `uv tool install`) failed.
+	// Installing the Unity Gateway CLI (`ug`, via `uv tool install`) failed.
 	SshAgentShimErrorCategoryGatewayCLIInstallFailed SshAgentShimErrorCategory = "UNITY_GATEWAY_CLI_INSTALL_FAILED"
 
 	// Installing Node.js (download, checksum, or extraction) failed.
 	SshAgentShimErrorCategoryNodeInstallFailed SshAgentShimErrorCategory = "NODE_INSTALL_FAILED"
 
-	// The install steps reported success but the Unity Gateway CLI (`ucode`) was still not on
+	// The install steps reported success but the Unity Gateway CLI (`ug`) was still not on
 	// PATH when the shim went to launch the agent: a setup inconsistency, distinct from the
 	// install itself failing.
 	SshAgentShimErrorCategoryUnityGatewayCLINotFound SshAgentShimErrorCategory = "UNITY_GATEWAY_CLI_NOT_FOUND"
@@ -73,6 +73,9 @@ const (
 	// intentionally does not do).
 	SshAgentShimErrorCategoryAgentLaunchFailed SshAgentShimErrorCategory = "AGENT_LAUNCH_FAILED"
 
+	// The configuration stage failed.
+	SshAgentShimErrorCategoryAgentConfigurationFailed SshAgentShimErrorCategory = "AGENT_CONFIGURATION_FAILED"
+
 	// The user interrupted setup (Ctrl-C or a termination signal) before the agent launched.
 	SshAgentShimErrorCategoryUserAborted SshAgentShimErrorCategory = "USER_ABORTED"