diff --git a/.nextchanges/cli/ssh-agents-dedicated.md b/.nextchanges/cli/ssh-agents-dedicated.md new file mode 100644 index 00000000000..f77a919251e --- /dev/null +++ b/.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)) 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) + } }