From c0971b5deb2b55a716cc715dce95c9704933a383 Mon Sep 17 00:00:00 2001 From: Russell Clarey Date: Tue, 27 Oct 2026 10:30:24 +0000 Subject: [PATCH 1/2] ssh: enable agent launchers on dedicated clusters Install the claude and codex launchers and append their directory to the session PATH on dedicated clusters too, not just serverless. The agent system context now names the compute type and describes what persists on it. Co-authored-by: Isaac --- .nextchanges/cli/ssh-agents-dedicated.md | 1 + acceptance/ssh/agent-launchers/output.txt | 13 +++++ acceptance/ssh/agent-launchers/script | 19 +++++++ experimental/ssh/internal/client/agentshim.go | 17 ++++-- .../ssh/internal/client/agentshim_test.go | 52 ++++++++++++++++--- experimental/ssh/internal/server/server.go | 10 ++-- experimental/ssh/internal/server/sshd.go | 2 +- experimental/ssh/internal/server/sshd_test.go | 11 ++-- 8 files changed, 98 insertions(+), 27 deletions(-) create mode 100644 .nextchanges/cli/ssh-agents-dedicated.md diff --git a/.nextchanges/cli/ssh-agents-dedicated.md b/.nextchanges/cli/ssh-agents-dedicated.md new file mode 100644 index 00000000000..de84ff41a42 --- /dev/null +++ b/.nextchanges/cli/ssh-agents-dedicated.md @@ -0,0 +1 @@ +* `databricks ssh connect` now makes `claude` and `codex` available on dedicated clusters, as it already did on serverless. diff --git a/acceptance/ssh/agent-launchers/output.txt b/acceptance/ssh/agent-launchers/output.txt index 261c6f2acb5..3a4bd48d14f 100644 --- a/acceptance/ssh/agent-launchers/output.txt +++ b/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 diff --git a/acceptance/ssh/agent-launchers/script b/acceptance/ssh/agent-launchers/script index 1a40666e908..84c72a3d375 100644 --- a/acceptance/ssh/agent-launchers/script +++ b/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' diff --git a/experimental/ssh/internal/client/agentshim.go b/experimental/ssh/internal/client/agentshim.go index 6361a4ca6ec..29974efd32d 100644 --- a/experimental/ssh/internal/client/agentshim.go +++ b/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) } diff --git a/experimental/ssh/internal/client/agentshim_test.go b/experimental/ssh/internal/client/agentshim_test.go index a2f0c5db50d..a175f8ddb5d 100644 --- a/experimental/ssh/internal/client/agentshim_test.go +++ b/experimental/ssh/internal/client/agentshim_test.go @@ -23,6 +23,7 @@ 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" @@ -30,6 +31,11 @@ import ( "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,7 +388,7 @@ 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) @@ -390,9 +396,27 @@ func TestInjectAgentContext(t *testing.T) { 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()) }) diff --git a/experimental/ssh/internal/server/server.go b/experimental/ssh/internal/server/server.go index 87e3e6d87e5..426c1194947 100644 --- a/experimental/ssh/internal/server/server.go +++ b/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 { diff --git a/experimental/ssh/internal/server/sshd.go b/experimental/ssh/internal/server/sshd.go index 3ee26ac3ecd..82fa40c43b9 100644 --- a/experimental/ssh/internal/server/sshd.go +++ b/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)) diff --git a/experimental/ssh/internal/server/sshd_test.go b/experimental/ssh/internal/server/sshd_test.go index d8f3419a78e..2ee2ba17995 100644 --- a/experimental/ssh/internal/server/sshd_test.go +++ b/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) + } } From 099fb328ae8d5d87b80106583f59bafb44db7283 Mon Sep 17 00:00:00 2001 From: Russell Clarey Date: Tue, 27 Oct 2026 11:06:31 +0000 Subject: [PATCH 2/2] changelog --- .nextchanges/cli/ssh-agents-dedicated.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.nextchanges/cli/ssh-agents-dedicated.md b/.nextchanges/cli/ssh-agents-dedicated.md index de84ff41a42..f77a919251e 100644 --- a/.nextchanges/cli/ssh-agents-dedicated.md +++ b/.nextchanges/cli/ssh-agents-dedicated.md @@ -1 +1 @@ -* `databricks ssh connect` now makes `claude` and `codex` available on dedicated clusters, as it already did on serverless. +* `databricks ssh connect` now makes `claude` and `codex` available on dedicated clusters. ([#7190](https://github.com/databricks/cli/pull/7190))