[SSH] Support agents in non-interactive sessions

.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))

acceptance/ssh/agent-launchers/out.test.toml

@@ -0,0 +1,4 @@
+Cloud = false
+GOOS.darwin = false
+GOOS.linux = true
+GOOS.windows = false

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

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 -- <cmd>" 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'

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",
+]

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
 }

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 <agent>" 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
+// "<binary> ssh agent-shim <agent>" 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))
+}

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)
+}