Add DATABRICKS_BEARER_COMMAND for externally brokered tokens

[dhruv0811]

Summary

DATABRICKS_BEARER lets a caller supply a pre-fetched bearer and skip the OAuth path. It is a value, though, and a value can't be rewritten in a running process. So it doesn't work for any caller whose bearer expires and has to be re-minted mid-session: an external credential broker, or a sidecar holding the refresh token on the caller's behalf.

This adds the command form of the same hatch. When DATABRICKS_BEARER_COMMAND is set, get_databricks_token runs it and returns what it prints, on every fetch rather than once.

The reason it's this small: every token consumer already funnels through that one function, so they all inherit it.

claude apiKeyHelper ─┐
codex auth.command  ─┼→ ucode auth-token ─┐
mcp servers ────────┘                     │
gateway_proxy ────────────────────────────┼→ get_databricks_token()
opencode/gemini/pi/copilot launchers ─────┘     ├─ DATABRICKS_BEARER          (existing)
                                                ├─ DATABRICKS_BEARER_COMMAND (new)
                                                ├─ PAT via ensure_pat_bearer (existing)
                                                └─ databricks auth token      (existing)

Design points

• shlex.split, not sh -c. Same reason build_auth_token_argv moved off the POSIX databricks ... | jq pipeline in Windows: ucode claude fails with apiKeyHelper POSIX shell error #116: plain argv runs identically on macOS, Linux, and Windows.
• Fails closed, no fall-through to OAuth. A broker-backed profile carries no OAuth cache to refresh, so falling through would surface a misleading stale-login error instead of the real cause. This mirrors how auth-token --use-pat already fails closed (cli.py), for the same reason.
• has_valid_databricks_auth short-circuits on it too. Otherwise ensure_databricks_auth probes the CLI and can open a browser for a workspace whose auth the caller already owns.
• Precedence unchanged. A non-empty DATABRICKS_BEARER still wins, so --use-pat (which exports one via ensure_pat_bearer) is unaffected.
• No token in the debug log. This command's stdout is the bearer, so the debug line records only the exit code and stderr rather than going through _format_subprocess_result, which includes stdout on a non-zero exit.

Motivation

We're wiring Databricks credentials into sandboxed coding-agent sessions in omnigent, where the sandbox deliberately holds no long-lived credential. It gets a short-lived handle plus the coordinates of a broker that vends a workspace token on demand, and the token is re-minted per request rather than written to disk.

ucode already does the right thing structurally: apiKeyHelper and codex's auth.command are re-invoked commands, not baked values. The only missing piece was a way to point that resolution at something other than the local OAuth cache. With this, the whole integration on our side is exporting one env var, and no per-harness config code.

Useful outside that case too: any CI or M2M setup whose bearer outlives a single fetch but not the session.

Test plan

tests/test_databricks.py::TestBearerCommand, 7 cases. Each puts a recording fake databricks on PATH that would happily serve a token, so asserting the marker file is absent asserts the OAuth path was never reached.

• serves the command's output, CLI untouched
• re-runs the command on every fetch (token-1, then token-2), which is the point of the change
• passes arguments without a shell (--coords 'a path' arrives as two argv entries)
• fails closed when the command prints no token, CLI untouched
• reports an unrunnable command
• static DATABRICKS_BEARER still wins
• has_valid_databricks_auth short-circuits, CLI untouched

uv run pytest tests/test_databricks.py -q   # 273 passed
uv run pytest tests/test_cli.py -q          # 270 passed
uv run ruff check . && uv run ruff format --check src tests

Not run locally: tests/test_e2e.py (needs a workspace). Inert unless the variable is set, so every existing path is unchanged.

---

Copilot review

"shlex.split breaks Windows paths" — confirmed and fixed. shlex.split(r"C:\bin\broker.exe --arg") returns ['C:binbroker.exe', '--arg'], which defeats the cross-platform point of keeping this path shell-free.

Took a different fix than suggested, though. posix=False swaps one Windows failure for another: it keeps the backslashes but leaves the quotes inside the token, so "C:\Program Files\ucode\ucode.exe" auth-token tokenizes as ['"C:\\Program Files\\ucode\\ucode.exe"', ...]. Instead, Windows takes the whole command line as one string and lets CreateProcess split it, which is the exact inverse of the string build_auth_shell_command already emits there via list2cmdline. So the string passes through on Windows and shlex.split stays on POSIX, and both a bare C:\... path and a quoted path with spaces round-trip correctly.

run's annotation widens to list[str] | str to match what subprocess.run already accepts. Covered by test_windows_hands_the_command_line_over_verbatim.

Second Copilot pass

"Returns stdout as the token even on a non-zero exit" — correct, and a real hole: a broker that printed its error to stdout and exited non-zero had that error forwarded as a bearer, resurfacing as a 401 far from the cause. Now requires a zero exit as well as non-empty stdout, matching _fetch in get_databricks_token, which already ignores output on a non-zero return. Covered by test_fails_closed_when_the_command_exits_non_zero.

"_env hardcodes : when prepending to PATH" — switched to os.pathsep. Worth being clear that this does not make the tests Windows-capable: the fake databricks and broker are #!/bin/sh scripts, so they are POSIX-only regardless of separator, as is the existing TestGetDatabricksToken._fake_databricks helper they mirror. CI runs ubuntu-latest for both jobs. Making this file cross-platform means replacing the shell-script fakes repo-wide, which is a separate change.

---

Found by running this branch in a Kubernetes sandbox

configure --workspaces called run_databricks_login unconditionally, so ucode configure inside a sandbox Pod sat on an interactive databricks auth login forever. That login cannot help a caller who brings their own bearer: it needs a browser the Pod does not have, and get_databricks_token returns before it would ever reach the OAuth path.

Both hatches were affected, so this also fixes the pre-existing DATABRICKS_BEARER case, and the two short-circuit checks collapse into one predicate:

def external_bearer_configured() -> bool:
    """Whether something outside ucode owns auth for this process."""

now used by has_valid_databricks_auth and by configure's forced-login branch. --use-pat already had a non-interactive path; this gives the same property to a caller supplying a bearer or a bearer command. Covered by TestForcedLoginWithExternalBearer (skips for both hatches, still logs in when neither is set).

One more gap this surfaced, left alone here: ucode's bootstrap hard-requires the databricks CLI binary even when DATABRICKS_BEARER_COMMAND makes it unnecessary for minting, and its auto-install shells sudo, which fails in a rootless Pod (sh: 1: sudo: not found). Worth a separate look.

[Copilot]

🟡 Changes recommended

The command parsing uses POSIX shlex.split() semantics which can break Windows-style paths containing backslashes, undermining the PR’s cross-platform intent.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Adds support for externally brokered Databricks bearer tokens by introducing DATABRICKS_BEARER_COMMAND, allowing get_databricks_token() to re-run a caller-provided command on each fetch (instead of relying on a static env var that can’t be updated mid-process).

Changes:

• Add DATABRICKS_BEARER_COMMAND short-circuit in get_databricks_token() and has_valid_databricks_auth(), with fail-closed behavior and stderr-only debug logging.
• Add a focused test suite covering precedence, re-execution per fetch, argv parsing (no shell), and failure modes.

File summaries

File	Description
src/ucode/databricks.py	Implements the new command-based bearer resolution and auth short-circuiting logic.
tests/test_databricks.py	Adds tests validating behavior, precedence, and “no OAuth CLI fallback” guarantees.

Review details

• Files reviewed: 2/2 changed files
• Comments generated: 1
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

🔵 Needs a closer look

The bearer-command path currently accepts stdout as a token even on non-zero exit codes, and the new tests hardcode PATH separators in a way that breaks Windows compatibility.

Review details

Suppressed comments (2)

Previously missed (2) — in code that hasn't changed since the last review.

src/ucode/databricks.py:1172

• _bearer_from_command returns stdout as the token even when the command exits non-zero. That can accidentally treat an error message printed to stdout as a bearer (and masks genuine failures). It should fail closed on non-zero exit codes, regardless of stdout contents.
  tests/test_databricks.py:3471
• _env hardcodes : when prepending to PATH. This will break these tests on Windows (PATH uses os.pathsep, typically ;) and is easy to make platform-correct.

• Files reviewed: 2/2 changed files
• Comments generated: 0 new
• Review effort level: Lite

[Copilot]

🟡 Changes recommended

The new tests include a brittle PATH assumption and one test’s behavior doesn’t actually exercise the intended “printed no token” branch.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

• Files reviewed: 2/2 changed files
• Comments generated: 2
• Review effort level: Lite

[Copilot]

🔵 Needs a closer look

The new error paths include the full DATABRICKS_BEARER_COMMAND value in exceptions, which can leak sensitive command-line arguments into logs.

Review details

Suppressed comments (2)

Previously missed (1) — in code that hasn't changed since the last review.

src/ucode/databricks.py:1160

• The raised RuntimeError includes the full value of DATABRICKS_BEARER_COMMAND (a user-provided command line). If the command embeds sensitive arguments (client secrets, refresh tokens, etc.), this will leak into CLI output/CI logs. Consider omitting the command string (or only including the executable name) in the exception message.

This issue also appears on line 1170 of the same file.

src/ucode/databricks.py:1172

• This error message also echoes the full DATABRICKS_BEARER_COMMAND string, which can inadvertently expose secrets if they are passed as command arguments. Prefer leaving the command out (or redacting it) while still including the exit code and stderr.

    reason = f"exited {result.returncode}" if result.returncode else "printed no token"
    detail = f" Stderr: {stderr}" if stderr else ""
    raise RuntimeError(f"DATABRICKS_BEARER_COMMAND {reason}. Command: {command}.{detail}")

• Files reviewed: 2/2 changed files
• Comments generated: 0 new
• Review effort level: Lite

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.