diff --git a/.nextchanges/cli/ssh-agents-non-interactive.md b/.nextchanges/cli/ssh-agents-non-interactive.md new file mode 100644 index 00000000000..fefc6bc05dd --- /dev/null +++ b/.nextchanges/cli/ssh-agents-non-interactive.md @@ -0,0 +1 @@ +* `databricks ssh connect` on serverless now makes `claude` and `codex` available in IDE sessions (`--ide`) and to non-interactive commands (e.g. `databricks ssh connect -- 'claude -p "..."'`). ([#7123](https://github.com/databricks/cli/pull/7123)) diff --git a/acceptance/ssh/agent-launchers/out.test.toml b/acceptance/ssh/agent-launchers/out.test.toml new file mode 100644 index 00000000000..44cce516e73 --- /dev/null +++ b/acceptance/ssh/agent-launchers/out.test.toml @@ -0,0 +1,4 @@ +Cloud = false +GOOS.darwin = false +GOOS.linux = true +GOOS.windows = false diff --git a/acceptance/ssh/agent-launchers/output.txt b/acceptance/ssh/agent-launchers/output.txt new file mode 100644 index 00000000000..261c6f2acb5 --- /dev/null +++ b/acceptance/ssh/agent-launchers/output.txt @@ -0,0 +1,35 @@ + +=== Serverless server + +=== SetEnv appends the launcher directory to PATH +SetEnv +HOME="[TEST_TMP_DIR]/server-home" +PATH="/usr/bin:/bin:[TEST_TMP_DIR]/server-home/.agent-shim/bin" +DATABRICKS_HOST="[DATABRICKS_URL]" +DATABRICKS_TOKEN="[DATABRICKS_TOKEN]" +DATABRICKS_CLI_PATH="[CLI]" +DATABRICKS_CLI_UPSTREAM=databricks_ssh_tunnel +DATABRICKS_CLI_UPSTREAM_VERSION= +DATABRICKS_SDK_UPSTREAM=databricks_ssh_tunnel +DATABRICKS_SDK_UPSTREAM_VERSION= +GIT_CONFIG_GLOBAL=/Workspace/.proc/self/git/config +ENABLE_DATABRICKS_CLI=true +PYTHONPYCACHEPREFIX=/tmp/pycache +DATABRICKS_JUPYTER_SERVERLESS=true + +=== Each supported agent has a launcher that delegates to the agent shim + +>>> ls server-home/.agent-shim/bin +claude +codex + +>>> cat server-home/.agent-shim/bin/claude +#!/usr/bin/env bash +# "claude" launcher from "databricks ssh connect"; delegates to "ssh agent-shim claude". +export DATABRICKS_WORKSPACE_HOME='/Workspace/Users/[USERNAME]' +exec [CLI] ssh agent-shim claude "$@" + +=== A non-interactive session finds the launchers on PATH +/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 diff --git a/acceptance/ssh/agent-launchers/script b/acceptance/ssh/agent-launchers/script new file mode 100644 index 00000000000..1a40666e908 --- /dev/null +++ b/acceptance/ssh/agent-launchers/script @@ -0,0 +1,42 @@ +# Runs the real "ssh server" (normally started on the cluster by the bootstrap job) +# against the test server, then opens a non-interactive SSH session through the +# sshd config it wrote. That session never reads a shell startup file, so it only +# finds the agent launchers if sshd's SetEnv puts them on PATH. +SSHD="$(command -v sshd || ls /usr/sbin/sshd /usr/local/sbin/sshd /sbin/sshd 2>/dev/null | head -n 1)" +if [ -z "$SSHD" ]; then + echo "SKIP_TEST sshd (openssh-server) not installed" + exit 0 +fi + +# The client's half of what "ssh connect" sets up: a key pair whose public half the +# server authorizes from the session's secret scope. +ssh-keygen -q -t ed25519 -N '' -f client_key +$CLI secrets create-scope ssh-scope +$CLI secrets put-secret ssh-scope client-public-key --string-value "$(cat client_key.pub)" + +# The server copies its whole environment into SetEnv, so start it from a fixed one: +# that keeps the generated config independent of the machine running the test. HOME +# is where the server writes its sshd config and the launchers. A short shutdown +# delay makes the server exit on its own once it is set up. +SERVER_HOME="$PWD/server-home" +mkdir -p "$SERVER_HOME" +title "Serverless server\n" +env -i HOME="$SERVER_HOME" PATH=/usr/bin:/bin DATABRICKS_HOST="$DATABRICKS_HOST" DATABRICKS_TOKEN="$DATABRICKS_TOKEN" \ + $CLI ssh server --serverless --cluster=test-cluster --session-id=test-session --secret-scope-name=ssh-scope \ + --authorized-key-secret-name=client-public-key --shutdown-delay=1s 2>LOG.server + +# One variable per line; no value here contains a space. +title "SetEnv appends the launcher directory to PATH\n" +grep '^SetEnv ' server-home/.ssh-tunnel/sshd_config | tr ' ' '\n' + +title "Each supported agent has a launcher that delegates to the agent shim\n" +trace ls server-home/.agent-shim/bin +trace cat server-home/.agent-shim/bin/claude + +# A non-interactive session, as "ssh connect -- " or an IDE's remote exec opens: +# ssh runs the command through the login shell without a PTY. +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 'echo "$PATH"; command -v claude codex' diff --git a/acceptance/ssh/agent-launchers/test.toml b/acceptance/ssh/agent-launchers/test.toml new file mode 100644 index 00000000000..05181600419 --- /dev/null +++ b/acceptance/ssh/agent-launchers/test.toml @@ -0,0 +1,16 @@ +Cloud = false + +# Linux-only: the test drives a real sshd, reliable only there (and where +# task test-exp-ssh provisions it). Absent GOOS keys default to enabled, so the +# other OSes must be disabled explicitly. +GOOS.linux = true +GOOS.darwin = false +GOOS.windows = false + +# The client key pair is generated per run; the server's home holds its sshd +# config, host key and the files it seeds, which the script prints selectively. +Ignore = [ + "client_key", + "client_key.pub", + "server-home", +] diff --git a/experimental/ssh/cmd/agent_shim.go b/experimental/ssh/cmd/agent_shim.go index 4eb97189201..2573050f4ae 100644 --- a/experimental/ssh/cmd/agent_shim.go +++ b/experimental/ssh/cmd/agent_shim.go @@ -2,6 +2,7 @@ package ssh import ( "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/experimental/ssh/internal/agentlauncher" "github.com/databricks/cli/experimental/ssh/internal/client" "github.com/databricks/cli/libs/cmdctx" "github.com/spf13/cobra" @@ -13,8 +14,8 @@ func newAgentShimCommand() *cobra.Command { Short: "Launch a coding agent configured against Unity AI Gateway (invoked on the remote)", Hidden: true, } - for _, agent := range client.SupportedAgentNames() { - cmd.AddCommand(newAgentShimAgentCommand(agent)) + for _, agent := range agentlauncher.SupportedAgents { + cmd.AddCommand(newAgentShimAgentCommand(agent.Name)) } return cmd } diff --git a/experimental/ssh/internal/agentlauncher/agentlauncher.go b/experimental/ssh/internal/agentlauncher/agentlauncher.go new file mode 100644 index 00000000000..2c7daac8ade --- /dev/null +++ b/experimental/ssh/internal/agentlauncher/agentlauncher.go @@ -0,0 +1,83 @@ +// Package agentlauncher installs the on-PATH launchers that start a coding agent +// through "databricks ssh agent-shim " on the remote compute. It is shared +// by the server, which installs the launchers, and the client, whose shim they run. +package agentlauncher + +import ( + "context" + "errors" + "fmt" + "path/filepath" + + "github.com/databricks/cli/libs/atomicfile" + "github.com/databricks/cli/libs/env" + "github.com/databricks/cli/libs/shellquote" + "github.com/databricks/cli/libs/telemetry/protos" +) + +const ( + // RootDir is the home-relative directory holding files/dependencies for the shim. + RootDir = ".agent-shim" + + // BinDir holds the per-agent launchers. + BinDir = RootDir + "/bin" + + // WorkspaceHomeEnv passes the user's workspace home to the shim for the agent context. + WorkspaceHomeEnv = "DATABRICKS_WORKSPACE_HOME" +) + +type AgentSpec struct { + Name string + TelemetryName protos.SshAgentShimAgent + // ContextFlag passes the context file via this CLI flag (Claude's --append-system-prompt-file). + ContextFlag string + // ContextHomeFile writes the context to this $HOME instructions file (codex's AGENTS.md). + ContextHomeFile string +} + +// Currently limited to the agents whose Unity Gateway CLI command supports --workspace +var SupportedAgents = []AgentSpec{ + {Name: "claude", TelemetryName: protos.SshAgentShimAgentClaudeCode, ContextFlag: "--append-system-prompt-file"}, + {Name: "codex", TelemetryName: protos.SshAgentShimAgentCodex, ContextHomeFile: ".codex/AGENTS.md"}, +} + +func AgentByName(name string) (AgentSpec, bool) { + for _, a := range SupportedAgents { + if a.Name == name { + return a, true + } + } + return AgentSpec{}, false +} + +// Install writes a launcher per supported agent to $HOME/BinDir. Each launcher runs +// " ssh agent-shim " and forwards wsHome for the agent context. +// Each launcher is written atomically: bash reads a script incrementally, so +// truncating the live file in place could cut off a launcher that a running +// session is still executing. A failed launcher doesn't stop the others from +// being written; the returned error names every agent that failed. +func Install(ctx context.Context, binary, wsHome string) error { + home, err := env.UserHomeDir(ctx) + if err != nil { + return fmt.Errorf("failed to get home directory: %w", err) + } + var errs []error + for _, agent := range SupportedAgents { + path := filepath.Join(home, BinDir, agent.Name) + err := atomicfile.Write(path, []byte(script(binary, wsHome, agent.Name)), 0o755, atomicfile.MkDir(0o755)) + if err != nil { + errs = append(errs, fmt.Errorf("failed to write the %s launcher: %w", agent.Name, err)) + } + } + return errors.Join(errs...) +} + +// script returns one agent's launcher. binary and wsHome are shell-quoted so a +// path or username with shell metacharacters can't break (or inject into) it. +func script(binary, wsHome, agent string) string { + return fmt.Sprintf(`#!/usr/bin/env bash +# "%[1]s" launcher from "databricks ssh connect"; delegates to "ssh agent-shim %[1]s". +export %[2]s=%[3]s +exec %[4]s ssh agent-shim %[1]s "$@" +`, agent, WorkspaceHomeEnv, shellquote.BashArg(wsHome), shellquote.BashArg(binary)) +} diff --git a/experimental/ssh/internal/agentlauncher/agentlauncher_test.go b/experimental/ssh/internal/agentlauncher/agentlauncher_test.go new file mode 100644 index 00000000000..b502e604f92 --- /dev/null +++ b/experimental/ssh/internal/agentlauncher/agentlauncher_test.go @@ -0,0 +1,81 @@ +package agentlauncher + +import ( + "os" + "path/filepath" + "runtime" + "testing" + + "github.com/databricks/cli/libs/env" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestScript(t *testing.T) { + got := script("/Workspace/Users/o'brien@example.com/databricks", "/Workspace/Users/o'brien@example.com", "codex") + assert.Equal(t, `#!/usr/bin/env bash +# "codex" launcher from "databricks ssh connect"; delegates to "ssh agent-shim codex". +export DATABRICKS_WORKSPACE_HOME='/Workspace/Users/o'\''brien@example.com' +exec '/Workspace/Users/o'\''brien@example.com/databricks' ssh agent-shim codex "$@" +`, got) +} + +func TestInstall(t *testing.T) { + const wsHome = "/Workspace/Users/me@example.com" + home := t.TempDir() + t.Setenv(env.HomeEnvVar(), home) + require.NoError(t, Install(t.Context(), "/old/databricks", wsHome)) + + for _, agent := range SupportedAgents { + path := filepath.Join(home, BinDir, agent.Name) + content, err := os.ReadFile(path) + require.NoError(t, err) + assert.Equal(t, script("/old/databricks", wsHome, agent.Name), string(content)) + if runtime.GOOS != "windows" { + info, err := os.Stat(path) + require.NoError(t, err) + assert.Equal(t, os.FileMode(0o755), info.Mode().Perm()) + } + } + + // A server restart from a new binary replaces the launchers in place. + require.NoError(t, Install(t.Context(), "/new/databricks", wsHome)) + content, err := os.ReadFile(filepath.Join(home, BinDir, "claude")) + require.NoError(t, err) + assert.Equal(t, script("/new/databricks", wsHome, "claude"), string(content)) +} + +func TestInstallContinuesPastFailure(t *testing.T) { + const wsHome = "/Workspace/Users/me@example.com" + home := t.TempDir() + t.Setenv(env.HomeEnvVar(), home) + // A non-empty directory where the claude launcher goes can't be replaced by a file. + require.NoError(t, os.MkdirAll(filepath.Join(home, BinDir, "claude", "blocker"), 0o755)) + + err := Install(t.Context(), "/databricks", wsHome) + assert.ErrorContains(t, err, "failed to write the claude launcher") + assert.NotContains(t, err.Error(), "codex") + + content, err := os.ReadFile(filepath.Join(home, BinDir, "codex")) + require.NoError(t, err) + assert.Equal(t, script("/databricks", wsHome, "codex"), string(content)) +} + +func TestSupportedAgents(t *testing.T) { + // Each agent injects context by exactly one mechanism (flag OR home file); + // the two are never both set. + expect := map[string]struct{ flag, home string }{ + "claude": {flag: "--append-system-prompt-file"}, + "codex": {home: ".codex/AGENTS.md"}, + } + require.Len(t, SupportedAgents, len(expect)) + for name, want := range expect { + a, ok := AgentByName(name) + require.True(t, ok) + assert.Equal(t, want.flag, a.ContextFlag, "%s ContextFlag", name) + assert.Equal(t, want.home, a.ContextHomeFile, "%s ContextHomeFile", name) + } + + _, ok := AgentByName("not-an-agent") + assert.False(t, ok) +} diff --git a/experimental/ssh/internal/client/agentshim.go b/experimental/ssh/internal/client/agentshim.go index b021aceb21b..17352791963 100644 --- a/experimental/ssh/internal/client/agentshim.go +++ b/experimental/ssh/internal/client/agentshim.go @@ -16,6 +16,7 @@ import ( "strings" "time" + "github.com/databricks/cli/experimental/ssh/internal/agentlauncher" "github.com/databricks/cli/internal/build" "github.com/databricks/cli/libs/auth" "github.com/databricks/cli/libs/cmdctx" @@ -35,10 +36,7 @@ import ( const ( // home-relative directory to hold files/dependencies for the shim. - agentRootDir = ".agent-shim" - - // directory which holds the per-agent wrappers; must match remoteShimDir in client.go. - agentBinDir = agentRootDir + "/bin" + agentRootDir = agentlauncher.RootDir // directory which holds toolchain deps (npm, uv) fetched when the image ships none. agentDepsDir = agentRootDir + "/deps" @@ -73,7 +71,7 @@ const ( npmExecutable = "npm" // environment variable used to pass the user's workspace home to the shim for the agent context. - workspaceHomeEnv = "DATABRICKS_WORKSPACE_HOME" + workspaceHomeEnv = agentlauncher.WorkspaceHomeEnv // the file holding the Databricks session context. contextFile = "agent-system-context.md" @@ -88,39 +86,6 @@ const ( setupLockStaleAfter = 2 * time.Minute ) -type agentSpec struct { - name string - telemetryName protos.SshAgentShimAgent - // contextFlag passes the context file via this CLI flag (Claude's --append-system-prompt-file). - contextFlag string - // contextHomeFile writes the context to this $HOME instructions file (codex's AGENTS.md). - contextHomeFile string -} - -// Currently limited to the agents whose Unity Gateway CLI command supports --workspace -var supportedAgents = []agentSpec{ - {name: "claude", telemetryName: protos.SshAgentShimAgentClaudeCode, contextFlag: "--append-system-prompt-file"}, - {name: "codex", telemetryName: protos.SshAgentShimAgentCodex, contextHomeFile: ".codex/AGENTS.md"}, -} - -// SupportedAgentNames lists the agents the ssh command registers a subcommand for. -func SupportedAgentNames() []string { - names := make([]string, len(supportedAgents)) - for i, a := range supportedAgents { - names[i] = a.name - } - return names -} - -func agentByName(name string) (agentSpec, bool) { - for _, a := range supportedAgents { - if a.name == name { - return a, true - } - } - return agentSpec{}, false -} - func RunAgentShim(ctx context.Context, client *databricks.WorkspaceClient, agentName string, agentArgs []string) (retErr error) { // Report the outcome of the launch. Setup runs from here to the exec below, so time it // from here too. Each failing step sets outcome.errorCategory; the returned error is @@ -142,7 +107,7 @@ func RunAgentShim(ctx context.Context, client *databricks.WorkspaceClient, agent logAgentShimEvent(ctx, agentName, outcome, time.Since(start)) }() - agent, ok := agentByName(agentName) + agent, ok := agentlauncher.AgentByName(agentName) if !ok { return categorize(protos.SshAgentShimErrorCategoryUnsupportedAgent, fmt.Errorf("unsupported agent %q", agentName)) } @@ -255,9 +220,9 @@ func buildAgentShimEvent(agentName string, outcome agentShimOutcome, setupDurati // lookup agent by name so this function is the single decision point for // telemetry and is easily testable telemetryName := protos.SshAgentShimAgentUnspecified - agent, ok := agentByName(agentName) + agent, ok := agentlauncher.AgentByName(agentName) if ok { - telemetryName = agent.telemetryName + telemetryName = agent.TelemetryName } return &protos.SshAgentShimEvent{ Agent: telemetryName, @@ -476,7 +441,7 @@ func bootstrapToolchain(ctx context.Context, home string) (string, error) { // 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 // Unity Gateway CLI execs the real agent rather than this wrapper - removePath(ctx, filepath.Join(home, agentBinDir)) + removePath(ctx, filepath.Join(home, agentlauncher.BinDir)) // add this CLI to PATH so agents can use it if self, err := os.Executable(); err == nil { prependPath(ctx, filepath.Dir(self)) @@ -686,10 +651,10 @@ func setEnvForUnityGateway(ctx context.Context) error { 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) +func runAgentConfiguration(ctx context.Context, agent agentlauncher.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 { @@ -698,22 +663,22 @@ func runAgentConfiguration(ctx context.Context, agent agentSpec, workspace strin return nil } -func prepareAgentLaunch(ctx context.Context, home string, agent agentSpec, agentArgs []string) (argv, environ []string, err error) { +func prepareAgentLaunch(ctx context.Context, home string, agent agentlauncher.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 = []string{ugExecutable, agent.name, "--"} + argv = []string{ugExecutable, agent.Name, "--"} argv = append(argv, contextArgs...) argv = append(argv, agentArgs...) return argv, os.Environ(), nil } // put the Databricks context where the agent reads it, returning any extra argv. -func injectAgentContext(ctx context.Context, home string, agent agentSpec) ([]string, error) { +func injectAgentContext(ctx context.Context, home string, agent agentlauncher.AgentSpec) ([]string, error) { systemContext := agentSystemContext(env.Get(ctx, workspaceHomeEnv)) switch { - case agent.contextFlag != "": + case agent.ContextFlag != "": // Keep the scratch context file under the shim's own directory so it never // clobbers an unrelated file the user happens to have in their home. dir := filepath.Join(home, agentRootDir) @@ -724,9 +689,9 @@ func injectAgentContext(ctx context.Context, home string, agent agentSpec) ([]st if err := os.WriteFile(f, []byte(systemContext), 0o644); err != nil { return nil, fmt.Errorf("failed to write agent context file: %w", err) } - return []string{agent.contextFlag, f}, nil - case agent.contextHomeFile != "": - f := filepath.Join(home, agent.contextHomeFile) + return []string{agent.ContextFlag, f}, nil + case agent.ContextHomeFile != "": + f := filepath.Join(home, agent.ContextHomeFile) if _, err := os.Stat(f); err == nil { return nil, nil // already present — don't clobber } else if !errors.Is(err, os.ErrNotExist) { @@ -840,13 +805,13 @@ func removePath(ctx context.Context, dir string) { } func agentSystemContext(wsHome string) string { - cwd := "the user's Databricks workspace home directory" + userHomePart := "" if wsHome != "" { - cwd = wsHome + userHomePart = " The user's Databricks workspace home directory is " + wsHome + "." } return fmt.Sprintf(`You are running inside a "databricks ssh connect" session on the driver node of a Databricks serverless cluster. - 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. Your working directory is %s. -- You can run shell commands and use the "databricks" CLI to explore the workspace.`, cwd) +- 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) } diff --git a/experimental/ssh/internal/client/agentshim_test.go b/experimental/ssh/internal/client/agentshim_test.go index 4cd486f25e3..9981a2c2f90 100644 --- a/experimental/ssh/internal/client/agentshim_test.go +++ b/experimental/ssh/internal/client/agentshim_test.go @@ -20,6 +20,8 @@ import ( "testing" "time" + "github.com/databricks/cli/experimental/ssh/internal/agentlauncher" + "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/telemetry/protos" "github.com/databricks/databricks-sdk-go" @@ -316,31 +318,10 @@ func TestEnsureToolchain(t *testing.T) { }) } -func TestSupportedAgents(t *testing.T) { - names := SupportedAgentNames() - assert.Equal(t, []string{"claude", "codex"}, names) - - // Each agent injects context by exactly one mechanism (flag OR home file); - // the two are never both set. - expect := map[string]struct{ flag, home string }{ - "claude": {flag: "--append-system-prompt-file"}, - "codex": {home: ".codex/AGENTS.md"}, - } - for name, want := range expect { - a, ok := agentByName(name) - require.True(t, ok) - assert.Equal(t, want.flag, a.contextFlag, "%s contextFlag", name) - assert.Equal(t, want.home, a.contextHomeFile, "%s contextHomeFile", name) - } - - _, ok := agentByName("not-an-agent") - assert.False(t, ok) -} - func TestInjectAgentContext(t *testing.T) { - claude, _ := agentByName("claude") - codex, _ := agentByName("codex") - noContext := agentSpec{name: "none"} // an agent with neither mechanism + claude, _ := agentlauncher.AgentByName("claude") + codex, _ := agentlauncher.AgentByName("codex") + noContext := agentlauncher.AgentSpec{Name: "none"} // an agent with neither mechanism t.Run("flag agent writes a scratch file and returns the flag", func(t *testing.T) { home := t.TempDir() @@ -388,9 +369,9 @@ func TestInjectAgentContext(t *testing.T) { } func TestAgentSystemContext(t *testing.T) { - assert.Contains(t, agentSystemContext("/Workspace/Users/me@example.com"), "working directory is /Workspace/Users/me@example.com") - // Falls back to a generic phrase when the workspace home is unknown. - assert.Contains(t, agentSystemContext(""), "working directory is the user's Databricks workspace home directory") + 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") } func newProbeClient(t *testing.T, host string) *databricks.WorkspaceClient { @@ -565,9 +546,9 @@ func TestProbeModelServicesPageError(t *testing.T) { } func TestBuildAgentShimEvent(t *testing.T) { - for _, agent := range supportedAgents { - t.Run(agent.name+" has an appropriate enum for telemetry", func(t *testing.T) { - got := buildAgentShimEvent(agent.name, agentShimOutcome{}, 0) + for _, agent := range agentlauncher.SupportedAgents { + t.Run(agent.Name+" has an appropriate enum for telemetry", func(t *testing.T) { + got := buildAgentShimEvent(agent.Name, agentShimOutcome{}, 0) assert.NotEqual(t, protos.SshAgentShimAgentUnspecified, got.Agent) }) } @@ -713,7 +694,7 @@ func TestRunAgentConfiguration(t *testing.T) { if runtime.GOOS == "windows" { return } - agent, ok := agentByName("claude") + agent, ok := agentlauncher.AgentByName("claude") require.True(t, ok) t.Run("configures the agent for the workspace", func(t *testing.T) { @@ -743,7 +724,7 @@ func TestRunAgentConfiguration(t *testing.T) { } func TestPrepareAgentLaunch(t *testing.T) { - agent, ok := agentByName("claude") + agent, ok := agentlauncher.AgentByName("claude") require.True(t, ok) t.Run("passes the context args before the user's agent args", func(t *testing.T) { diff --git a/experimental/ssh/internal/client/client.go b/experimental/ssh/internal/client/client.go index 3ac21fab8f9..80142d54d4c 100644 --- a/experimental/ssh/internal/client/client.go +++ b/experimental/ssh/internal/client/client.go @@ -538,7 +538,7 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt return runIDE(ctx, client, userName, keyPath, knownHostsPath, serverPort, clusterID, opts) } else { log.Infof(ctx, "Additional SSH arguments: %v", opts.AdditionalArgs) - return spawnSSHClient(ctx, client, userName, keyPath, knownHostsPath, serverPort, clusterID, version, opts) + return spawnSSHClient(ctx, client, userName, keyPath, knownHostsPath, serverPort, clusterID, opts) } } @@ -948,34 +948,6 @@ func shellSingleQuote(s string) string { return "'" + strings.ReplaceAll(s, "'", `'\''`) + "'" } -// agentWrapperScript returns one agent's on-PATH launcher: it resolves the uploaded -// databricks CLI for the driver's arch, then delegates to `ssh agent-shim `. -// wsHome anchors that binary path and is forwarded (via env) for the agent context. -func agentWrapperScript(wsHome, version, agent string) string { - versionedDir := wsHome + "/.databricks/ssh-tunnel/" + version - // Mirror getReleaseName (minus .zip); ${_arch} is expanded on the remote. - subdir := strings.TrimSuffix(getReleaseName("${_arch}", version), ".zip") - // Single-quote the workspace/version-derived path so a username or version with - // shell metacharacters can't break (or inject into) the exec line, while leaving - // ${_arch} unquoted so the remote shell still expands it (getReleaseName always - // embeds it, so Cut always splits here). - before, after, _ := strings.Cut(subdir, "${_arch}") - binary := shellSingleQuote(versionedDir+"/"+before) + "${_arch}" + shellSingleQuote(after+"/databricks") - - // buildRemoteShellArgs only installs launchers when wsHome is set, so it is - // always non-empty here. - exports := "export " + workspaceHomeEnv + "=" + shellSingleQuote(wsHome) + "\n" - - return fmt.Sprintf(`#!/usr/bin/env bash -# "%[3]s" launcher from "databricks ssh connect"; delegates to "ssh agent-shim %[3]s". -case "$(uname -m)" in - x86_64|amd64) _arch=amd64 ;; - aarch64|arm64) _arch=arm64 ;; - *) echo "%[3]s: unsupported architecture $(uname -m)" >&2; exit 1 ;; -esac -%[1]sexec %[2]s ssh agent-shim %[3]s "$@"`, exports, binary, agent) -} - // buildRemoteShellArgs returns the ssh arguments that follow the hostname. // // For the interactive case (no remote command given), it forces PTY allocation @@ -987,8 +959,9 @@ esac // would resolve to the system interpreter instead of $DATABRICKS_VIRTUAL_ENV. Using // -i avoids that reset; the server's ~/.bashrc snippet (see seedEnvActivation) then // re-prepends the environment bin after /etc/bash.bashrc runs, so bare `python`/`pip` -// resolve to the environment interpreter. It first installs a per-agent launcher -// (see agentWrapperScript) on PATH, then cds into wsHome when set (else $HOME). +// resolve to the environment interpreter. When wsHome is set, the shell first changes +// into the user's workspace home folder; if that directory is missing the cd is +// ignored and the shell still launches from $HOME. // // For the non-interactive case (e.g. `databricks ssh connect ... -- ls -la`), // the user's command is returned verbatim so behavior is unchanged. @@ -996,39 +969,15 @@ esac // Note: this returns the remote command only. PTY allocation (-t) is added to // the ssh options *before* the destination by the caller; -t placed after the // host would be parsed as part of the remote command, not as ssh's flag. -func buildRemoteShellArgs(opts ClientOptions, wsHome, version string) []string { +func buildRemoteShellArgs(opts ClientOptions, wsHome string) []string { if len(opts.AdditionalArgs) > 0 { return opts.AdditionalArgs } - shell := `command -v bash >/dev/null 2>&1 && exec bash -i || exec "${SHELL:-/bin/sh}" -i` - // For any interactive session, land in the user's workspace home when known; - // if that directory is missing the cd is ignored and the shell still launches - // from $HOME. This is independent of the launcher install below, so dedicated - // clusters keep the workspace cwd they had before agent launchers existed. + cmd := `command -v bash >/dev/null 2>&1 && exec bash -i || exec "${SHELL:-/bin/sh}" -i` if wsHome != "" { - shell = "cd " + shellSingleQuote(wsHome) + " 2>/dev/null; " + shell - } - // The agent launchers are serverless-only for now: dedicated clusters have no - // story for persisting the shims or respecting existing installs. They also - // need the workspace path to locate the uploaded CLI. Otherwise just open the shell. - if !opts.IsServerlessMode() || wsHome == "" { - return []string{shell} - } - // Write each agent's wrapper (quoted heredoc keeps its $VARS literal), then open the shell. - remoteShimDir := "$HOME/" + agentBinDir - var cmd strings.Builder - fmt.Fprintf(&cmd, "mkdir -p \"%s\"\n", remoteShimDir) - for _, agent := range SupportedAgentNames() { - fmt.Fprintf(&cmd, `cat > "%[1]s/%[2]s" <<'AGENT_SHIM_WRAPPER_EOF' -%[3]s -AGENT_SHIM_WRAPPER_EOF -chmod +x "%[1]s/%[2]s" -`, remoteShimDir, agent, agentWrapperScript(wsHome, version, agent)) - } - // Append the shim dir (don't prepend) so the wrappers never shadow a real - // binary of the same name earlier on PATH. - fmt.Fprintf(&cmd, "export PATH=\"$PATH:%s\"\n", remoteShimDir) - return []string{cmd.String() + shell} + cmd = "cd " + shellSingleQuote(wsHome) + " 2>/dev/null; " + cmd + } + return []string{cmd} } // buildSSHArgs assembles the argument list for the ssh client. Options come @@ -1040,7 +989,7 @@ chmod +x "%[1]s/%[2]s" // Host key checking is strict rather than accept-new because the caller has already // pinned the server's key in knownHostsPath (see pinServerHostKey), so there is nothing // left to accept on trust. -func buildSSHArgs(userName, privateKeyPath, knownHostsPath, proxyCommand, hostName, wsHome, version string, opts ClientOptions) []string { +func buildSSHArgs(userName, privateKeyPath, knownHostsPath, proxyCommand, hostName, wsHome string, opts ClientOptions) []string { sshArgs := []string{ "-l", userName, "-i", privateKeyPath, @@ -1055,11 +1004,11 @@ func buildSSHArgs(userName, privateKeyPath, knownHostsPath, proxyCommand, hostNa sshArgs = append(sshArgs, "-t") } sshArgs = append(sshArgs, hostName) - sshArgs = append(sshArgs, buildRemoteShellArgs(opts, wsHome, version)...) + sshArgs = append(sshArgs, buildRemoteShellArgs(opts, wsHome)...) return sshArgs } -func spawnSSHClient(ctx context.Context, client *databricks.WorkspaceClient, userName, privateKeyPath, knownHostsPath string, serverPort int, clusterID, version string, opts ClientOptions) error { +func spawnSSHClient(ctx context.Context, client *databricks.WorkspaceClient, userName, privateKeyPath, knownHostsPath string, serverPort int, clusterID string, opts ClientOptions) error { // Create a copy with metadata for the ProxyCommand optsWithMetadata := opts optsWithMetadata.ServerMetadata = FormatMetadata(userName, serverPort, clusterID) @@ -1076,14 +1025,14 @@ func spawnSSHClient(ctx context.Context, client *databricks.WorkspaceClient, use // OS home. Only needed for an interactive session; skip the lookup otherwise. var wsHome string if len(opts.AdditionalArgs) == 0 { - if currentUser, err := client.CurrentUser.Me(ctx, iam.MeRequest{}); err != nil { - log.Warnf(ctx, "Failed to resolve current user for workspace home directory: %v", err) + if homeDir, err := sshWorkspace.GetWorkspaceHomeDir(ctx, client); err != nil { + log.Warnf(ctx, "Failed to resolve workspace home directory: %v", err) } else { - wsHome = "/Workspace/Users/" + currentUser.UserName + wsHome = homeDir } } - sshArgs := buildSSHArgs(userName, privateKeyPath, knownHostsPath, proxyCommand, hostName, wsHome, version, opts) + sshArgs := buildSSHArgs(userName, privateKeyPath, knownHostsPath, proxyCommand, hostName, wsHome, opts) log.Debugf(ctx, "Launching SSH client: ssh %s", strings.Join(sshArgs, " ")) sshCmd := exec.CommandContext(ctx, "ssh", sshArgs...) diff --git a/experimental/ssh/internal/client/client_internal_test.go b/experimental/ssh/internal/client/client_internal_test.go index 746f5d58f98..1db7adbd88a 100644 --- a/experimental/ssh/internal/client/client_internal_test.go +++ b/experimental/ssh/internal/client/client_internal_test.go @@ -368,51 +368,30 @@ func TestHostKeyChangedHint(t *testing.T) { func TestBuildRemoteShellArgs(t *testing.T) { const bashCmd = `command -v bash >/dev/null 2>&1 && exec bash -i || exec "${SHELL:-/bin/sh}" -i` - const version = "1.2.3" + const wsHome = "/Workspace/Users/me@example.com" // A serverless session has a connection name and no cluster ID. serverless := ClientOptions{ConnectionName: "my-conn"} - t.Run("serverless installs a launcher per agent, cds home, then opens bash", func(t *testing.T) { - const wsHome = "/Workspace/Users/me@example.com" - args := buildRemoteShellArgs(serverless, wsHome, version) - require.Len(t, args, 1) - // A wrapper is written for every supported agent before the shell opens. - for _, agent := range SupportedAgentNames() { - assert.Contains(t, args[0], `cat > "$HOME/.agent-shim/bin/`+agent+`"`) - assert.Contains(t, args[0], agentWrapperScript(wsHome, version, agent)) - } - assert.Contains(t, args[0], `export PATH="$PATH:$HOME/.agent-shim/bin"`) - // The cd into the workspace home precedes the shell launch. - assert.True(t, strings.HasSuffix(args[0], `cd '`+wsHome+`' 2>/dev/null; `+bashCmd)) - }) - - t.Run("dedicated cluster skips the launchers but keeps the workspace cwd", func(t *testing.T) { - const wsHome = "/Workspace/Users/me@example.com" - args := buildRemoteShellArgs(ClientOptions{ClusterID: "abc-123"}, wsHome, version) - require.Len(t, args, 1) - // No launchers on a dedicated cluster... - assert.NotContains(t, args[0], "agent-shim/bin") - // ...but the cd into the workspace home is preserved (regression guard). - assert.Equal(t, `cd '`+wsHome+`' 2>/dev/null; `+bashCmd, args[0]) + t.Run("interactive cds into the workspace home, then opens bash", func(t *testing.T) { + args := buildRemoteShellArgs(serverless, wsHome) + assert.Equal(t, []string{`cd '` + wsHome + `' 2>/dev/null; ` + bashCmd}, args) }) - t.Run("serverless without a workspace home skips the launchers", func(t *testing.T) { - // Without wsHome the wrappers can't locate the CLI, so only the shell opens. - args := buildRemoteShellArgs(serverless, "", version) - require.Len(t, args, 1) - assert.Equal(t, bashCmd, args[0]) + t.Run("interactive without a workspace home just opens bash", func(t *testing.T) { + args := buildRemoteShellArgs(serverless, "") + assert.Equal(t, []string{bashCmd}, args) }) t.Run("non-interactive passes additional args verbatim", func(t *testing.T) { additional := []string{"ls", "-la"} - args := buildRemoteShellArgs(ClientOptions{ConnectionName: "my-conn", AdditionalArgs: additional}, "/Workspace/Users/me@example.com", version) + args := buildRemoteShellArgs(ClientOptions{ConnectionName: "my-conn", AdditionalArgs: additional}, wsHome) assert.Equal(t, additional, args) }) } func TestBuildSSHArgsSetsServerAliveInterval(t *testing.T) { - args := buildSSHArgs("user", "/key", "/pins/myhost", "proxy command", "myhost", "", "1.2.3", ClientOptions{}) + args := buildSSHArgs("user", "/key", "/pins/myhost", "proxy command", "myhost", "", ClientOptions{}) // ssh stops parsing options at the destination, so an option placed after the host would be // treated as part of the remote command rather than as an ssh option. @@ -423,7 +402,7 @@ func TestBuildSSHArgsSetsServerAliveInterval(t *testing.T) { } func TestBuildSSHArgsPinsHostKey(t *testing.T) { - args := buildSSHArgs("user", "/key", "/pins/myhost", "proxy command", "myhost", "", "1.2.3", ClientOptions{}) + args := buildSSHArgs("user", "/key", "/pins/myhost", "proxy command", "myhost", "", ClientOptions{}) // The pinned file is the whole point of strict checking here: without it ssh would // fall back to ~/.ssh/known_hosts, where an entry for this name may be left over from @@ -437,27 +416,6 @@ func TestBuildSSHArgsPinsHostKey(t *testing.T) { } } -func TestAgentWrapperScript(t *testing.T) { - const wsHome = "/Workspace/Users/me@example.com" - - t.Run("resolves the uploaded binary and delegates to agent-shim ", func(t *testing.T) { - wrapper := agentWrapperScript(wsHome, "1.2.3", "codex") - assert.True(t, strings.HasPrefix(wrapper, "#!/usr/bin/env bash")) - // Architecture is resolved on the remote via uname -m. - assert.Contains(t, wrapper, "case \"$(uname -m)\" in") - // Binary path uses the version dir + arch subdir; the workspace/version-derived - // portion is single-quoted while ${_arch} stays expandable on the remote. - assert.Contains(t, wrapper, `exec '`+wsHome+`/.databricks/ssh-tunnel/1.2.3/databricks_cli_1.2.3_linux_'${_arch}'/databricks' ssh agent-shim codex "$@"`) - // The workspace home is forwarded so the shim can name the working directory. - assert.Contains(t, wrapper, "export "+workspaceHomeEnv+"='"+wsHome+"'") - }) - - t.Run("dev builds use the version-less release subdir", func(t *testing.T) { - wrapper := agentWrapperScript(wsHome, "1.2.3-dev+abc", "claude") - assert.Contains(t, wrapper, `/ssh-tunnel/1.2.3-dev+abc/databricks_cli_linux_'${_arch}'/databricks`) - }) -} - func TestBuildSSHArgsPTYPlacement(t *testing.T) { indexOf := func(args []string, want string) int { for i, a := range args { @@ -469,7 +427,7 @@ func TestBuildSSHArgsPTYPlacement(t *testing.T) { } t.Run("interactive forces a PTY before the destination", func(t *testing.T) { - args := buildSSHArgs("user", "/key", "/pins/myhost", "proxy command", "myhost", "/Workspace/Users/me@example.com", "1.2.3", ClientOptions{}) + args := buildSSHArgs("user", "/key", "/pins/myhost", "proxy command", "myhost", "/Workspace/Users/me@example.com", ClientOptions{}) ptyIdx := indexOf(args, "-t") hostIdx := indexOf(args, "myhost") require.NotEqual(t, -1, ptyIdx, "-t must be present for interactive sessions") @@ -481,7 +439,7 @@ func TestBuildSSHArgsPTYPlacement(t *testing.T) { }) t.Run("non-interactive does not force a PTY", func(t *testing.T) { - args := buildSSHArgs("user", "/key", "/pins/myhost", "proxy command", "myhost", "", "1.2.3", ClientOptions{AdditionalArgs: []string{"ls", "-la"}}) + args := buildSSHArgs("user", "/key", "/pins/myhost", "proxy command", "myhost", "", ClientOptions{AdditionalArgs: []string{"ls", "-la"}}) assert.Equal(t, -1, indexOf(args, "-t"), "no PTY for non-interactive passthrough") hostIdx := indexOf(args, "myhost") require.NotEqual(t, -1, hostIdx) diff --git a/experimental/ssh/internal/server/server.go b/experimental/ssh/internal/server/server.go index 39f0938b773..87e3e6d87e5 100644 --- a/experimental/ssh/internal/server/server.go +++ b/experimental/ssh/internal/server/server.go @@ -18,6 +18,7 @@ import ( "path/filepath" "time" + "github.com/databricks/cli/experimental/ssh/internal/agentlauncher" "github.com/databricks/cli/experimental/ssh/internal/proxy" "github.com/databricks/cli/experimental/ssh/internal/workspace" "github.com/databricks/cli/libs/env" @@ -80,6 +81,12 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ServerOpt listenAddr := fmt.Sprintf("0.0.0.0:%d", port) log.Info(ctx, "Starting server on "+listenAddr) + // Resolved once and shared by the metadata write and the agent launchers. + wsHome, err := workspace.GetWorkspaceHomeDir(ctx, client) + if err != nil { + return fmt.Errorf("failed to resolve the workspace home directory: %w", err) + } + // Save metadata including ClusterID (required for Driver Proxy connections in serverless mode) metadata := &workspace.WorkspaceMetadata{ Port: port, @@ -87,7 +94,7 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ServerOpt UsagePolicyID: opts.UsagePolicyID, KeepDetachedProcesses: opts.KeepDetachedProcesses, } - err = workspace.SaveWorkspaceMetadata(ctx, client, opts.Version, opts.SessionID, metadata) + err = workspace.SaveWorkspaceMetadata(ctx, client, wsHome, opts.Version, opts.SessionID, metadata) if err != nil { return fmt.Errorf("failed to save metadata to the workspace: %w", err) } @@ -110,6 +117,15 @@ 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) + } + } + createServerCommand := func(ctx context.Context) *exec.Cmd { return createSSHDProcess(ctx, sshdConfigPath) } @@ -293,3 +309,14 @@ func seedEnvActivation(ctx context.Context) error { log.Info(ctx, "Seeded environment activation in "+bashrcPath) return nil } + +// installAgentLaunchers writes the agent launchers under $HOME; prepareSSHDConfig +// puts them on PATH for every session. The launchers run the agent shim from the +// same uploaded CLI binary this server runs from. +func installAgentLaunchers(ctx context.Context, wsHome string) error { + binary, err := os.Executable() + if err != nil { + return fmt.Errorf("failed to resolve the CLI binary: %w", err) + } + return agentlauncher.Install(ctx, binary, wsHome) +} diff --git a/experimental/ssh/internal/server/server_test.go b/experimental/ssh/internal/server/server_test.go index ae3ca5ba0a0..9ec4aae6bf2 100644 --- a/experimental/ssh/internal/server/server_test.go +++ b/experimental/ssh/internal/server/server_test.go @@ -6,7 +6,9 @@ import ( "strings" "testing" + "github.com/databricks/cli/experimental/ssh/internal/agentlauncher" "github.com/databricks/cli/libs/env" + "github.com/databricks/cli/libs/shellquote" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -83,3 +85,17 @@ func TestSeedEnvActivation(t *testing.T) { assert.Contains(t, string(again), "export FOO=bar") }) } + +func TestInstallAgentLaunchers(t *testing.T) { + home := t.TempDir() + t.Setenv(env.HomeEnvVar(), home) + require.NoError(t, installAgentLaunchers(t.Context(), "/Workspace/Users/me@example.com")) + + // The launcher runs the shim from this binary and forwards the user's workspace home. + self, err := os.Executable() + require.NoError(t, err) + launcher, err := os.ReadFile(filepath.Join(home, agentlauncher.BinDir, "claude")) + require.NoError(t, err) + assert.Contains(t, string(launcher), "export DATABRICKS_WORKSPACE_HOME='/Workspace/Users/me@example.com'\n") + assert.Contains(t, string(launcher), "exec "+shellquote.BashArg(self)+" ssh agent-shim claude") +} diff --git a/experimental/ssh/internal/server/sshd.go b/experimental/ssh/internal/server/sshd.go index 01fd0e3cdb0..3ee26ac3ecd 100644 --- a/experimental/ssh/internal/server/sshd.go +++ b/experimental/ssh/internal/server/sshd.go @@ -12,6 +12,7 @@ import ( "strconv" "strings" + "github.com/databricks/cli/experimental/ssh/internal/agentlauncher" "github.com/databricks/cli/experimental/ssh/internal/keys" "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/log" @@ -68,27 +69,7 @@ func prepareSSHDConfig(ctx context.Context, client *databricks.WorkspaceClient, return "", err } - // Set all available env vars, wrapping values in quotes, escaping quotes, and stripping newlines - var setEnvBuf strings.Builder - setEnvBuf.WriteString("SetEnv") - for _, env := range os.Environ() { - parts := strings.SplitN(env, "=", 2) - if len(parts) == 2 { - fmt.Fprintf(&setEnvBuf, ` %s="%s"`, parts[0], escapeEnvValue(parts[1])) - } - } - setEnvBuf.WriteString(" DATABRICKS_CLI_UPSTREAM=databricks_ssh_tunnel") - setEnvBuf.WriteString(" DATABRICKS_CLI_UPSTREAM_VERSION=" + opts.Version) - setEnvBuf.WriteString(" DATABRICKS_SDK_UPSTREAM=databricks_ssh_tunnel") - setEnvBuf.WriteString(" DATABRICKS_SDK_UPSTREAM_VERSION=" + opts.Version) - setEnvBuf.WriteString(" GIT_CONFIG_GLOBAL=/Workspace/.proc/self/git/config") - setEnvBuf.WriteString(" ENABLE_DATABRICKS_CLI=true") - setEnvBuf.WriteString(" PYTHONPYCACHEPREFIX=/tmp/pycache") - if opts.Serverless { - setEnvBuf.WriteString(" DATABRICKS_JUPYTER_SERVERLESS=true") - } - setEnv := setEnvBuf.String() - + setEnv := sshdSetEnv(os.Environ(), homeDir, opts) if err := os.WriteFile(sshdConfig, []byte(sshdConfigContent(keyPath, authKeysPath, setEnv)), 0o600); err != nil { return "", err } @@ -114,6 +95,39 @@ func sshdConfigContent(hostKeyPath, authorizedKeysPath, setEnv string) string { setEnv + "\n" } +// sshdSetEnv returns the SetEnv directive that seeds every session's environment. +func sshdSetEnv(environ []string, homeDir string, opts ServerOptions) string { + // Set all available env vars, wrapping values in quotes, escaping quotes, and stripping newlines + var setEnvBuf strings.Builder + setEnvBuf.WriteString("SetEnv") + for _, env := range environ { + parts := strings.SplitN(env, "=", 2) + if len(parts) == 2 { + value := parts[1] + // sshd applies SetEnv before any shell starts, so this puts the agent + // launchers (see installAgentLaunchers) on PATH for every session: + // 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" { + value += ":" + path.Join(homeDir, agentlauncher.BinDir) + } + fmt.Fprintf(&setEnvBuf, ` %s="%s"`, parts[0], escapeEnvValue(value)) + } + } + setEnvBuf.WriteString(" DATABRICKS_CLI_UPSTREAM=databricks_ssh_tunnel") + setEnvBuf.WriteString(" DATABRICKS_CLI_UPSTREAM_VERSION=" + opts.Version) + setEnvBuf.WriteString(" DATABRICKS_SDK_UPSTREAM=databricks_ssh_tunnel") + setEnvBuf.WriteString(" DATABRICKS_SDK_UPSTREAM_VERSION=" + opts.Version) + setEnvBuf.WriteString(" GIT_CONFIG_GLOBAL=/Workspace/.proc/self/git/config") + setEnvBuf.WriteString(" ENABLE_DATABRICKS_CLI=true") + setEnvBuf.WriteString(" PYTHONPYCACHEPREFIX=/tmp/pycache") + if opts.Serverless { + setEnvBuf.WriteString(" DATABRICKS_JUPYTER_SERVERLESS=true") + } + return setEnvBuf.String() +} + func createSSHDProcess(ctx context.Context, configPath string) *exec.Cmd { return exec.CommandContext(ctx, "/usr/sbin/sshd", "-f", configPath, "-i") } diff --git a/experimental/ssh/internal/server/sshd_test.go b/experimental/ssh/internal/server/sshd_test.go index 4887f9895d8..d8f3419a78e 100644 --- a/experimental/ssh/internal/server/sshd_test.go +++ b/experimental/ssh/internal/server/sshd_test.go @@ -87,3 +87,14 @@ func TestSSHDConfigSetsClientAliveInterval(t *testing.T) { assert.Less(t, clientAliveIntervalSeconds, 8*60) assert.Greater(t, 3*clientAliveIntervalSeconds, 30) } + +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"`) +} diff --git a/experimental/ssh/internal/workspace/workspace.go b/experimental/ssh/internal/workspace/workspace.go index c67adb6bac4..a6a29ca7bb2 100644 --- a/experimental/ssh/internal/workspace/workspace.go +++ b/experimental/ssh/internal/workspace/workspace.go @@ -29,30 +29,39 @@ type WorkspaceMetadata struct { KeepDetachedProcesses bool `json:"keep_detached_processes,omitempty"` } -func getWorkspaceRootDir(ctx context.Context, client *databricks.WorkspaceClient) (string, error) { +// GetWorkspaceHomeDir returns the current user's workspace home folder. +func GetWorkspaceHomeDir(ctx context.Context, client *databricks.WorkspaceClient) (string, error) { me, err := client.CurrentUser.Me(ctx, iam.MeRequest{}) if err != nil { return "", fmt.Errorf("failed to get current user: %w", err) } - return fmt.Sprintf("/Workspace/Users/%s/.databricks/ssh-tunnel", me.UserName), nil + return "/Workspace/Users/" + me.UserName, nil +} + +func versionedDir(homeDir, version string) string { + return filepath.ToSlash(filepath.Join(homeDir, ".databricks/ssh-tunnel", version)) +} + +func contentDir(homeDir, version, sessionID string) string { + return filepath.ToSlash(filepath.Join(versionedDir(homeDir, version), sessionID)) } func GetWorkspaceVersionedDir(ctx context.Context, client *databricks.WorkspaceClient, version string) (string, error) { - contentDir, err := getWorkspaceRootDir(ctx, client) + homeDir, err := GetWorkspaceHomeDir(ctx, client) if err != nil { return "", fmt.Errorf("failed to get workspace root directory: %w", err) } - return filepath.ToSlash(filepath.Join(contentDir, version)), nil + return versionedDir(homeDir, version), nil } // GetWorkspaceContentDir returns the directory for storing session content. // sessionID is the unique identifier for the session (cluster ID for dedicated clusters, connection name for serverless). func GetWorkspaceContentDir(ctx context.Context, client *databricks.WorkspaceClient, version, sessionID string) (string, error) { - contentDir, err := GetWorkspaceVersionedDir(ctx, client, version) + homeDir, err := GetWorkspaceHomeDir(ctx, client) if err != nil { return "", fmt.Errorf("failed to get versioned workspace directory: %w", err) } - return filepath.ToSlash(filepath.Join(contentDir, sessionID)), nil + return contentDir(homeDir, version, sessionID), nil } // GetWorkspaceMetadata loads session metadata from the workspace. @@ -85,20 +94,15 @@ func GetWorkspaceMetadata(ctx context.Context, client *databricks.WorkspaceClien return &metadata, nil } -// SaveWorkspaceMetadata saves session metadata to the workspace. +// SaveWorkspaceMetadata saves session metadata under the given workspace home. // sessionID is the unique identifier for the session (cluster ID for dedicated clusters, connection name for serverless). -func SaveWorkspaceMetadata(ctx context.Context, client *databricks.WorkspaceClient, version, sessionID string, metadata *WorkspaceMetadata) error { +func SaveWorkspaceMetadata(ctx context.Context, client *databricks.WorkspaceClient, homeDir, version, sessionID string, metadata *WorkspaceMetadata) error { metadataBytes, err := json.Marshal(metadata) if err != nil { return fmt.Errorf("failed to marshal metadata: %w", err) } - contentDir, err := GetWorkspaceContentDir(ctx, client, version, sessionID) - if err != nil { - return fmt.Errorf("failed to get workspace content directory: %w", err) - } - - workspaceFiler, err := filer.NewWorkspaceFilesClient(client, contentDir) + workspaceFiler, err := filer.NewWorkspaceFilesClient(client, contentDir(homeDir, version, sessionID)) if err != nil { return fmt.Errorf("failed to create workspace files client: %w", err) }