[SSH] Enable agent shims on dedicated clusters

.nextchanges/cli/ssh-agents-dedicated.md

@@ -0,0 +1 @@
+* `databricks ssh connect` now makes `claude` and `codex` available on dedicated clusters. ([#7190](https://github.com/databricks/cli/pull/7190))

acceptance/ssh/agent-launchers/output.txt

@@ -33,3 +33,16 @@ exec [CLI] ssh agent-shim claude "$@"
 /usr/bin:/bin:[TEST_TMP_DIR]/server-home/.agent-shim/bin
 [TEST_TMP_DIR]/server-home/.agent-shim/bin/claude
 [TEST_TMP_DIR]/server-home/.agent-shim/bin/codex
+
+=== Dedicated server
+
+=== SetEnv appends the launcher directory to PATH
+PATH="/usr/bin:/bin:[TEST_TMP_DIR]/server-home/.agent-shim/bin"
+
+>>> ls server-home/.agent-shim/bin
+claude
+codex
+
+=== A non-interactive session finds the launchers on PATH
+[TEST_TMP_DIR]/server-home/.agent-shim/bin/claude
+[TEST_TMP_DIR]/server-home/.agent-shim/bin/codex

acceptance/ssh/agent-launchers/script

@@ -40,3 +40,22 @@ ssh -F /dev/null -i client_key -l "$(id -un)" \
   -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o LogLevel=ERROR \
   -o ProxyCommand="$SSHD -i -f $SERVER_HOME/.ssh-tunnel/sshd_config -o StrictModes=no" \
   localhost 'echo "$PATH"; command -v claude codex'
+
+# Dedicated clusters get the same launchers and PATH entry as serverless.
+rm -rf "$SERVER_HOME"
+mkdir -p "$SERVER_HOME"
+title "Dedicated server\n"
+env -i HOME="$SERVER_HOME" PATH=/usr/bin:/bin DATABRICKS_HOST="$DATABRICKS_HOST" DATABRICKS_TOKEN="$DATABRICKS_TOKEN" \
+  $CLI ssh server --cluster=test-cluster --session-id=test-session --secret-scope-name=ssh-scope \
+  --authorized-key-secret-name=client-public-key --shutdown-delay=1s 2>LOG.dedicated-server
+
+title "SetEnv appends the launcher directory to PATH\n"
+grep '^SetEnv ' server-home/.ssh-tunnel/sshd_config | tr ' ' '\n' | grep '^PATH='
+
+trace ls server-home/.agent-shim/bin
+
+title "A non-interactive session finds the launchers on PATH\n"
+ssh -F /dev/null -i client_key -l "$(id -un)" \
+  -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o LogLevel=ERROR \
+  -o ProxyCommand="$SSHD -i -f $SERVER_HOME/.ssh-tunnel/sshd_config -o StrictModes=no" \
+  localhost 'command -v claude codex'

experimental/ssh/internal/client/agentshim.go

@@ -695,7 +695,8 @@ func prepareAgentLaunch(ctx context.Context, home string, agent agentlauncher.Ag
 
 // put the Databricks context where the agent reads it, returning any extra argv.
 func injectAgentContext(ctx context.Context, home string, agent agentlauncher.AgentSpec) ([]string, error) {
-	systemContext := agentSystemContext(env.Get(ctx, workspaceHomeEnv))
+	serverless := dbr.RuntimeVersion(ctx).Type == dbr.ClusterTypeServerless
+	systemContext := agentSystemContext(env.Get(ctx, workspaceHomeEnv), serverless)
 	switch {
 	case agent.ContextFlag != "":
 		// Keep the scratch context file under the shim's own directory so it never
@@ -823,14 +824,20 @@ func removePath(ctx context.Context, dir string) {
 	_ = os.Setenv("PATH", strings.Join(dirs, string(os.PathListSeparator)))
 }
 
-func agentSystemContext(wsHome string) string {
+func agentSystemContext(wsHome string, serverless bool) string {
 	userHomePart := ""
 	if wsHome != "" {
 		userHomePart = " The user's Databricks workspace home directory is " + wsHome + "."
 	}
+	computePart := "dedicated cluster"
+	persistencePart := "Local files persist across sessions until the cluster restarts or terminates; only paths under /Workspace and /Volumes persist beyond that."
+	if serverless {
+		computePart = "serverless cluster"
+		persistencePart = "This container is ephemeral; only paths under /Workspace and /Volumes persist across sessions."
+	}
 	return fmt.Sprintf(`You are running inside a "databricks ssh connect" session on the driver node of a
-Databricks serverless cluster.
+Databricks %s.
 - The "databricks" CLI is installed and already authenticated: DATABRICKS_HOST and DATABRICKS_TOKEN are set in the environment, so "databricks ..." commands work with no "databricks auth login". The same token governs Unity Catalog and serving-endpoint access.
-- This container is ephemeral; only paths under /Workspace and /Volumes persist across sessions.%s
-- You can run shell commands and use the "databricks" CLI to explore the workspace.`, userHomePart)
+- %s%s
+- You can run shell commands and use the "databricks" CLI to explore the workspace.`, computePart, persistencePart, userHomePart)
 }

experimental/ssh/internal/client/agentshim_test.go

@@ -23,13 +23,19 @@ import (
 	"github.com/databricks/cli/experimental/ssh/internal/agentlauncher"
 
 	"github.com/databricks/cli/libs/cmdio"
+	"github.com/databricks/cli/libs/dbr"
 	"github.com/databricks/cli/libs/telemetry/protos"
 	"github.com/databricks/databricks-sdk-go"
 	"github.com/databricks/databricks-sdk-go/config"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
+var (
+	serverlessRuntime = dbr.Environment{IsDbr: true, Version: "client.4.9"}
+	dedicatedRuntime  = dbr.Environment{IsDbr: true, Version: "16.4"}
+)
+
 func TestPrependPath(t *testing.T) {
 	t.Setenv("PATH", strings.Join([]string{"/a", "/b"}, string(os.PathListSeparator)))
 	prependPath(t.Context(), "/b") // existing entry moves to the front (deduped)
@@ -358,7 +364,7 @@ func TestInjectAgentContext(t *testing.T) {
 
 	t.Run("flag agent writes a scratch file and returns the flag", func(t *testing.T) {
 		home := t.TempDir()
-		args, err := injectAgentContext(t.Context(), home, claude)
+		args, err := injectAgentContext(dbr.MockRuntime(t.Context(), serverlessRuntime), home, claude)
 		require.NoError(t, err)
 		require.Len(t, args, 2)
 		assert.Equal(t, "--append-system-prompt-file", args[0])
@@ -369,7 +375,7 @@ func TestInjectAgentContext(t *testing.T) {
 
 	t.Run("home-file agent writes the global instructions file, no args", func(t *testing.T) {
 		home := t.TempDir()
-		args, err := injectAgentContext(t.Context(), home, codex)
+		args, err := injectAgentContext(dbr.MockRuntime(t.Context(), serverlessRuntime), home, codex)
 		require.NoError(t, err)
 		assert.Nil(t, args)
 		data, err := os.ReadFile(filepath.Join(home, ".codex", "AGENTS.md"))
@@ -382,17 +388,35 @@ func TestInjectAgentContext(t *testing.T) {
 		target := filepath.Join(home, ".codex", "AGENTS.md")
 		require.NoError(t, os.MkdirAll(filepath.Dir(target), 0o755))
 		require.NoError(t, os.WriteFile(target, []byte("user's own instructions"), 0o644))
-		args, err := injectAgentContext(t.Context(), home, codex)
+		args, err := injectAgentContext(dbr.MockRuntime(t.Context(), serverlessRuntime), home, codex)
 		require.NoError(t, err)
 		assert.Nil(t, args)
 		data, err := os.ReadFile(target)
 		require.NoError(t, err)
 		assert.Equal(t, "user's own instructions", string(data))
 	})
 
+	t.Run("context describes the compute the shim runs on", func(t *testing.T) {
+		for _, c := range []struct {
+			runtime dbr.Environment
+			want    string
+		}{
+			{serverlessRuntime, "Databricks serverless cluster."},
+			{dedicatedRuntime, "Databricks dedicated cluster."},
+		} {
+			home := t.TempDir()
+			args, err := injectAgentContext(dbr.MockRuntime(t.Context(), c.runtime), home, claude)
+			require.NoError(t, err)
+			require.Len(t, args, 2)
+			data, err := os.ReadFile(args[1])
+			require.NoError(t, err)
+			assert.Contains(t, string(data), c.want)
+		}
+	})
+
 	t.Run("agent without a mechanism writes nothing", func(t *testing.T) {
 		home := t.TempDir()
-		args, err := injectAgentContext(t.Context(), home, noContext)
+		args, err := injectAgentContext(dbr.MockRuntime(t.Context(), serverlessRuntime), home, noContext)
 		require.NoError(t, err)
 		assert.Nil(t, args)
 		entries, err := os.ReadDir(home)
@@ -403,8 +427,20 @@ func TestInjectAgentContext(t *testing.T) {
 
 func TestAgentSystemContext(t *testing.T) {
 	wsHome := "/Workspace/Users/me@example.com"
-	assert.Contains(t, agentSystemContext(wsHome), "user's Databricks workspace home directory is "+wsHome)
-	assert.NotContains(t, agentSystemContext(""), "user's Databricks workspace home directory")
+	for _, serverless := range []bool{true, false} {
+		assert.Contains(t, agentSystemContext(wsHome, serverless), "user's Databricks workspace home directory is "+wsHome)
+		assert.NotContains(t, agentSystemContext("", serverless), "user's Databricks workspace home directory")
+	}
+
+	serverless := agentSystemContext(wsHome, true)
+	assert.Contains(t, serverless, "Databricks serverless cluster.")
+	assert.Contains(t, serverless, "This container is ephemeral")
+
+	dedicated := agentSystemContext(wsHome, false)
+	assert.Contains(t, dedicated, "Databricks dedicated cluster.")
+	assert.Contains(t, dedicated, "until the cluster restarts or terminates")
+	assert.NotContains(t, dedicated, "serverless")
+	assert.NotContains(t, dedicated, "ephemeral")
 }
 
 func newProbeClient(t *testing.T, host string) *databricks.WorkspaceClient {
@@ -762,7 +798,7 @@ func TestPrepareAgentLaunch(t *testing.T) {
 
 	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"})
+		argv, environ, err := prepareAgentLaunch(dbr.MockRuntime(t.Context(), serverlessRuntime), 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:])
@@ -775,7 +811,7 @@ func TestPrepareAgentLaunch(t *testing.T) {
 		home := filepath.Join(t.TempDir(), "home")
 		require.NoError(t, os.WriteFile(home, nil, 0o644))
 
-		_, _, err := prepareAgentLaunch(t.Context(), home, agent, nil)
+		_, _, err := prepareAgentLaunch(dbr.MockRuntime(t.Context(), serverlessRuntime), home, agent, nil)
 		require.Error(t, err)
 		assert.Equal(t, protos.SshAgentShimErrorCategoryAgentContextSetupFailed, agentShimOutcome{err: err}.category())
 	})

experimental/ssh/internal/server/server.go

@@ -117,13 +117,9 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ServerOpt
 		log.Warnf(ctx, "Failed to seed environment activation, bare python/pip may resolve to the wrong interpreter in interactive sessions: %v", err)
 	}
 
-	// The agent launchers are serverless-only for now: dedicated clusters have no
-	// story for persisting the shims or respecting existing installs. Best-effort
-	// for the same reason as seedEnvActivation: the tunnel works without them.
-	if opts.Serverless {
-		if err := installAgentLaunchers(ctx, wsHome); err != nil {
-			log.Warnf(ctx, "Failed to install agent launchers, the affected agents won't be available in SSH sessions: %v", err)
-		}
+	// Best-effort for the same reason as seedEnvActivation: the tunnel works without them.
+	if err := installAgentLaunchers(ctx, wsHome); err != nil {
+		log.Warnf(ctx, "Failed to install agent launchers, the affected agents won't be available in SSH sessions: %v", err)
 	}
 
 	createServerCommand := func(ctx context.Context) *exec.Cmd {

experimental/ssh/internal/server/sshd.go

@@ -109,7 +109,7 @@ func sshdSetEnv(environ []string, homeDir string, opts ServerOptions) string {
 			// interactive or not, IDE or terminal, bash or sh, with no startup file
 			// involved. Appended (not prepended) so the launchers never shadow a real
 			// binary of the same name earlier on PATH.
-			if opts.Serverless && parts[0] == "PATH" {
+			if parts[0] == "PATH" {
 				value += ":" + path.Join(homeDir, agentlauncher.BinDir)
 			}
 			fmt.Fprintf(&setEnvBuf, ` %s="%s"`, parts[0], escapeEnvValue(value))

experimental/ssh/internal/server/sshd_test.go

@@ -91,10 +91,9 @@ func TestSSHDConfigSetsClientAliveInterval(t *testing.T) {
 func TestSSHDSetEnvAppendsAgentLaunchersToPath(t *testing.T) {
 	environ := []string{"PATH=/usr/bin:/bin", "FOO=bar"}
 
-	serverless := sshdSetEnv(environ, "/home/spark", ServerOptions{Serverless: true})
-	assert.Contains(t, serverless, ` PATH="/usr/bin:/bin:/home/spark/.agent-shim/bin"`)
-	assert.Contains(t, serverless, ` FOO="bar"`)
-
-	dedicated := sshdSetEnv(environ, "/home/spark", ServerOptions{})
-	assert.Contains(t, dedicated, ` PATH="/usr/bin:/bin"`)
+	for _, opts := range []ServerOptions{{Serverless: true}, {}} {
+		setEnv := sshdSetEnv(environ, "/home/spark", opts)
+		assert.Contains(t, setEnv, ` PATH="/usr/bin:/bin:/home/spark/.agent-shim/bin"`, "serverless=%v", opts.Serverless)
+		assert.Contains(t, setEnv, ` FOO="bar"`, "serverless=%v", opts.Serverless)
+	}
 }