Resolve Pi's Databricks auth per request instead of baking a bearer

[dhruv0811]

Summary

Pi's models.json carried the bearer as a literal apiKey, so a background thread rewrote the file every 30 minutes to keep it from going stale (_refresh_forever, TOKEN_REFRESH_INTERVAL_SECONDS). Three costs: ucode had to stay alive as the supervising parent, a live token sat on disk, and the config could still hold a 30-minute-old token at request time.

Pi has a native hook for this. A config value starting with ! is run as a command and its stdout used as the value, and the provider auth path resolves it through resolveConfigValueOrThrow (the uncached path) rather than the process-lifetime cache in resolve-config-value.js.

So each provider's apiKey becomes !ucode auth-token --host ..., and the refresher goes away:

 providers["databricks-claude"] = {
     "baseUrl": pi_base_urls["claude"],
     "api": "anthropic-messages",
-    "apiKey": token,
+    "apiKey": api_key,      # "!" + build_auth_shell_command(...)
     "authHeader": True,

Net effect: Pi joins Claude Code and Codex in resolving auth through ucode auth-token, and lands on the same shape OpenCode already moved to with its auth plugin. No token on disk, no supervisor requirement, no staleness window.

Verified, not assumed

The load-bearing claim is "resolved per request, not once per process." Checked against pi 0.84.3 with a loopback fake gateway and an apiKey of !mint.sh, where mint.sh increments a counter and prints tok-N. Turn 1 returns a tool_use, so pi runs the tool and issues a second provider request in the same process:

req1 auth=Bearer tok-1 stream=True tools=4
req2 auth=Bearer tok-2 stream=True tools=4
mints=2

Two requests, two distinct mints. pi auth print-api-key --provider <name> independently confirms the !command resolution.

Notes

• No --force-refresh, unlike OpenCode's plugin. That plugin caches the token itself and needs the flag for its 401 retry; Pi has no cache on this path, so forcing a mint would round-trip to the workspace every turn. Plain auth-token serves the CLI's cached token until it nears expiry.
• --profile and --use-pat are forwarded, so PAT-configured workspaces keep working through the same command.
• Pi does not re-mint and retry on a 401. My first probe returned 401 to everything: one request, one mint, error surfaced verbatim, no retry. The exposure shrinks from ~30 minutes to the gap between mint and request, so it isn't reachable in practice, but it's worth knowing.
• OAUTH_TOKEN is unchanged, still exported at launch. Pi doesn't read it for model auth (the probe ran with only the !command apiKey), so the launch-time mint may be droppable, but that's a separate question from this fix.

Test plan

uv run pytest tests/test_agent_pi.py -q                     # 55 passed
uv run pytest tests/ -q --ignore=tests/test_e2e.py \
    --ignore=tests/test_e2e_tracing.py                      # 2243 passed, 1 pre-existing failure
uv run ruff check . && uv run ruff format --check src tests  # clean

Updated:

• test_token_in_api_key → test_api_key_config_value_embedded_verbatim: render_overlay must embed the config value it's handed without reinterpreting it.
• test_config_written_with_correct_model_and_token → ..._and_auth_command: the written apiKey is a !-prefixed auth-token command.
• tests/test_e2e_user_agent.py::TestPiUserAgent: sets DATABRICKS_BEARER in pi's env so the real ucode auth-token the config now invokes has something to print without a workspace behind the capture server. This makes the test stronger: it drives the real command end to end through the real pi binary, and it went from ~30s to ~1.5s because the request now arrives promptly.

Added:

• test_bearer_never_written_to_the_config: a real-looking token passed to write_tool_config must not appear anywhere in models.json.
• test_every_provider_gets_the_auth_command: all three providers, not just claude.
• TestBuildPiApiKey: leading !, runs auth-token, carries --host, omits --force-refresh, forwards --profile and --use-pat.

Pre-existing failure, not from this change

tests/test_e2e_user_agent.py::TestClaudeUserAgent::test_user_agent_arrives_at_gateway fails on clean origin/main too (verified by stashing this change and re-running: 1 failed, 4 passed). Untouched here.

Update: CI is green on both jobs, so that Claude failure is local to my machine (missing or mismatched claude binary), not a break on main.

---

Copilot review

"Docstring says pi never holds a token" — fair, the wording overclaimed. The bearer no longer lands in models.json, but launch still exports OAUTH_TOKEN, so a token does reach the process environment. Docstring now states the real guarantee.

"The models.json backup can persist an old baked bearer" — real, but pre-existing and not pi-specific, so leaving it out of this PR. backup_existing_file returns early when the backup already exists (config_io.py), so it captures the user's config once, on ucode's first write, and never refreshes. A ucode-written bearer only lands there on the narrow path where someone deletes the backup and re-runs an older ucode. The same mechanism backs up token-bearing configs for OpenCode, Gemini, Codex, and Claude today on main, so the fix belongs in backup_existing_file, covering all six, rather than as a pi-only patch here. Happy to send that separately.

---

Verified in a live Kubernetes sandbox

Built this branch (merged with the DATABRICKS_BEARER_COMMAND branch) into a wheel, installed it in an agent-sandbox Pod launched by an Omnigent server, and configured all four harnesses with no credential in the Pod other than a bearer command.

Pi's models.json, written by this branch:

databricks-claude      !/home/omnigent/bin/ucode auth-token --host https://<workspace>
databricks-openai      !/home/omnigent/bin/ucode auth-token --host https://<workspace>
databricks-gemini      !/home/omnigent/bin/ucode auth-token --host https://<workspace>

A command on every provider, no --force-refresh, and grepping the actual bearer against every file ucode wrote under its app dir found it in none of Pi's. ucode configure --agents claude,codex,pi,opencode then reported Pi is working from its live validation message, so a real turn went through Databricks AI Gateway resolving auth per request.

Worth noting for contrast, since it is the same problem this PR fixes and is unchanged on main: OpenCode's opencode.json still carries the bearer as a literal apiKey (a 791-char token on disk, for all three providers). Its plugin mints on demand when accessToken is unset, so that seed looks like an optimization paid for with a token on disk. Separate change; noting it here only because Pi and OpenCode were the two baked-bearer harnesses and only one of them is fixed by this PR.

[Copilot]

🟡 Changes recommended

The Pi config backup path can still persist an old baked bearer token on disk, which undermines the security goal of avoiding tokens on disk.

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

Pull request overview

This PR updates the Pi agent integration to avoid embedding Databricks bearer tokens in models.json by switching each provider’s apiKey to a Pi !command that runs ucode auth-token per request, eliminating the background token refresher thread and reducing token staleness and on-disk exposure.

Changes:

• Replace Pi providers’ literal bearer apiKey with a !ucode auth-token ... command via build_pi_api_key.
• Remove the periodic token refresh thread from Pi launch, relying on per-request resolution instead.
• Update and extend tests to validate the new !command apiKey behavior and ensure tokens are not written to models.json.

File summaries

File	Description
src/ucode/agents/pi.py	Switch Pi provider auth to a per-request !command apiKey and remove the token refresher thread.
tests/test_agent_pi.py	Update existing tests and add coverage for the new Pi apiKey command generation and “no bearer in models.json” behavior.
tests/test_e2e_user_agent.py	Adjust Pi e2e harness env so ucode auth-token can resolve via DATABRICKS_BEARER without a real workspace behind the capture server.

Review details

• Files reviewed: 3/3 changed files
• Comments generated: 2
• 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]

🟢 Approval recommended

The changes consistently remove on-disk bearer persistence for Pi, eliminate the refresh thread, and add/adjust tests to cover the new on-demand auth-command behavior across providers.

Review details

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