Add ssh agent-shim subcommand - #6734
Merged
Merged
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>
Collaborator
Integration test reportCommit: 692c221
Top 4 slowest tests (at least 2 minutes):
|
anton-107
reviewed
Sep 18, 2026
anton-107
reviewed
Sep 18, 2026
anton-107
reviewed
Sep 18, 2026
| } | ||
| return []string{agent.contextFlag, f}, nil | ||
| case agent.contextHomeFile != "": | ||
| f := filepath.Join(home, agent.contextHomeFile) |
Contributor
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.
Contributor
Author
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
anton-107
reviewed
Sep 18, 2026
rugpanov
reviewed
Sep 21, 2026
| // 1. uv (installs into ~/.local/bin). | ||
| if _, err := exec.LookPath("uv"); err != nil { | ||
| cmdio.LogString(ctx, "Installing uv...") | ||
| if err := runShell(ctx, "curl -LsSf https://astral.sh/uv/install.sh | sh"); err != nil { |
Contributor
There was a problem hiding this comment.
Supply chain: the bootstrap executes unpinned third-party code with the driver's workspace credentials in scope.
Two commands here run remote code that inherits DATABRICKS_TOKEN and the full driver environment:
- L383 pipes the uv installer straight to
shwith no version pin or checksum:curl -LsSf https://astral.sh/uv/install.sh | sh. - L391 installs ucode from a mutable git tag (
@v0.1.0) rather than an immutable commit SHA.
So an upstream compromise or a tag/installer replacement gets workspace credentials on the driver. Worth noting the inconsistency with the Node path just below, which is SHA256-verified — uv and ucode aren't. Suggestion: pin uv to a versioned installer + known checksum (or vendor it), and pin ucode to a commit SHA.
Minor, same construct: sh -c "curl ... | sh" has no pipefail, so a failed curl still exits 0 (empty stdin to sh) and setup fails later with a misleading "uv/ucode" error rather than at the download.
Contributor
Author
There was a problem hiding this comment.
Fair point, I'll pin the uv installer as well and change to a commit SHA for ucode in a followup PR 👍
anton-107
approved these changes
Sep 22, 2026
rclarey
enabled auto-merge
September 23, 2026 10:56
Collaborator
Integration test reportCommit: ec50a05
359 interesting tests: 190 MISS, 169 FAIL
Top 50 slowest tests (at least 2 minutes):
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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)