Conversation
Adds a hidden `databricks ssh agent-shim <agent>` command that runs on the serverless driver: it probes the workspace AI Gateway, sets up PATH, and execs a ucode-configured coding agent (Claude Code, Codex) with the Databricks session context injected. This first commit assumes the toolchain (uv/ucode/node) is already present on the driver; first-run bootstrap of that toolchain follows in the next commit. Co-authored-by: Isaac <no-reply@databricks.com>
On first launch the shim now installs the toolchain the agent needs before execing ucode: uv (into ~/.local/bin), a pinned ucode release via uv, and Node/npm (downloaded and SHA256-verified) when the serverless image ships none. It also silences npm's update-notifier and records the resolved PATH so later launches skip straight to the agent. Co-authored-by: Isaac <no-reply@databricks.com>
Waiting for approvalBased on git history, these people are best suited to review:
Eligible reviewers: Suggestions based on git history. See OWNERS for ownership rules. |
Integration test reportCommit: 174c5e5
Top 3 slowest tests (at least 2 minutes):
|
| if workspace != "" { | ||
| argv = append(argv, "--workspace", workspace) | ||
| } | ||
| argv = append(argv, contextArgs...) |
There was a problem hiding this comment.
are these meant to be passed to real "claude" and "codex" binaries? if so, don't you need to prepend these with an additional "--" so that they get passed to them instead of to ug?
There was a problem hiding this comment.
ug passes along arguments it doesn't understand, so this works today but you're right I'll add -- to make this more robust 👍
| return nil, fmt.Errorf("failed to acquire setup lock: %w", err) | ||
| } | ||
| // Held by another process: reclaim it if stale, otherwise wait and retry. | ||
| if info, statErr := os.Stat(lockPath); statErr == nil && time.Since(info.ModTime()) > setupLockStaleAfter { |
There was a problem hiding this comment.
The lock becomes “stale” after two minutes, with no heartbeat or owner-liveness check. A legitimate cold installation exceeding that duration lets a second session enter setup concurrently. Worse, the original owner’s unconditional unlock subsequently deletes the second owner’s lock. This defeats the serialization protecting shared installation directories. Suggestion: use an OS-backed lock released on process exit rather than treating elapsed time as proof of death.
There was a problem hiding this comment.
There is no OS-backed lock in go's standard library (or any of the CLI's current dependencies AFAIK), so this is modeled after the existing lock here:
cli/libs/versioncheck/notice.go
Line 124 in ac5988e
| } | ||
| return []string{agent.contextFlag, f}, nil | ||
| case agent.contextHomeFile != "": | ||
| f := filepath.Join(home, agent.contextHomeFile) |
There was a problem hiding this comment.
The shim writes instructions to ~/.codex, while preserving any inherited CODEX_HOME. Pinned UG also writes its Codex configuration to the hardcoded default directory. Codex then reads a different directory, missing both the injected instructions and UG’s profile.
Suggestion: make the shim and upstream configuration honor the same home, or reject the unsupported override early.
There was a problem hiding this comment.
I don't quite follow. You're saying this shim and ug both don't support CODEX_HOME? ug passes a --profile flag to codex to explicitly load the ug managed configuration, I'm not sure how that interacts with CODEX_HOME though.
Either way I'll create a ticket to follow up, but I think this is out of scope for this PR
| if err := client.Config.Authenticate(req); err != nil { | ||
| return nil, fmt.Sprintf("network error: %v", err) | ||
| } | ||
| resp, err := (&http.Client{Transport: client.Config.HTTPTransport, Timeout: 10 * time.Second}).Do(req) |
There was a problem hiding this comment.
Constructing a new http.Client with Config.HTTPTransport is not equivalent to using the configured SDK client: the transport is normally nil, and SDK configuration is applied elsewhere. A configured SDK request succeeds against a local TLS fixture while this preflight rejects the same endpoint because it ignores the configured TLS setting. The hardcoded timeout likewise bypasses SDK timeout configuration.
Question: can we use the configured SDK request path and retain workspace-routing headers?
There was a problem hiding this comment.
The model services endpoint is supported in the SDK but the legacy endpoint is not. I'll see if I can reuse the HTTP client the SDK uses to manually make the legacy request
3 participants
Changes
Add a hidden
ssh agent-shimsubcommand which wraps Unity Gateway CLI to install, configure, and launch coding agents in SSH sessionsWhy
Today it’s possible to get an agent harness working in an SSH session, but the process is entirely manual and undocumented. After connecting, a user must:
x-databricks-use-coding-agent-modeheaderThe limits of this are that nobody knows this recipe exists, and that it is easy to misconfigure. For example, if the small/fast model tier for Claude isn't mapped to a Databricks endpoint, background calls silently fall back to
api.anthropic.comand fail.The
ssh agent-shimsubcommand this PR introduces solves this problem by handling installation of all needed dependencies, then deferring to Unity Gateway CLI to install and configure agents. A later PR will add shims forclaude/codex/etc added toPATHwhich call this subcommand. This will allow users to simply run the command they are used to and get an agent working out-of-the-box.Tests
Added unit tests
Manually testing
./task build snapshot-releasedatabricks ssh connect --releases-dir ./dist./.databricks/ssh-tunnel/<release number>/databricks_cli_<release number>_linux_<CPU arch>/databricks ssh agent-shim [claude|codex]in your workspace user folder (should be default on a new connection)