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

src/ucode/agents/pi.py

@@ -21,16 +21,18 @@
 `max_tokens` caps that pi has no global way to honor without per-model
 config we don't currently maintain.
 
-The bearer token is baked into the file and refreshed by a background thread
-while the session runs (same pattern as OpenCode/Copilot).
+Each provider's `apiKey` is pi's `!command` config value rather than a baked
+bearer, so pi mints one per request via `ucode auth-token` and nothing that
+expires is written to `models.json` (the on-demand model OpenCode's auth plugin
+already uses). A token still reaches the process environment: `launch` exports
+`OAUTH_TOKEN` as before.
 """
 
 from __future__ import annotations
 
 import os
 import signal
 import subprocess
-import threading
 
 from ucode.config_io import (
     APP_DIR,
@@ -42,7 +44,7 @@
 )
 from ucode.databricks import (
     ANTHROPIC_FAMILIES,
-    TOKEN_REFRESH_INTERVAL_SECONDS,
+    build_auth_shell_command,
     build_pi_base_urls,
     classify_model_family,
     get_databricks_token,
@@ -101,13 +103,16 @@ def _resolve_model_selector(
 
 def render_overlay(
     model: str,
-    token: str,
+    api_key: str,
     pi_base_urls: dict[str, str],
     claude_models: dict[str, str],
     codex_models: list[str],
     gemini_models: list[str],
 ) -> tuple[dict, list[list[str]]]:
-    """Return (overlay, managed_key_paths) for Pi's private agent config."""
+    """Return (overlay, managed_key_paths) for Pi's private agent config.
+
+    ``api_key`` is a pi config value, not necessarily a literal bearer: see
+    ``build_pi_api_key`` for the `!command` form every provider gets."""
     providers: dict = {}
     keys: list[list[str]] = [["model"]]
     # Pi expands header values that match an env var name. Our UA contains
@@ -119,7 +124,7 @@ def render_overlay(
         providers["databricks-claude"] = {
             "baseUrl": pi_base_urls["claude"],
             "api": "anthropic-messages",
-            "apiKey": token,
+            "apiKey": api_key,
             "authHeader": True,
             # Gateway's Anthropic translator rejects per-tool
             # `eager_input_streaming` on the streaming + tools path. Pi sends
@@ -133,7 +138,7 @@ def render_overlay(
         providers["databricks-openai"] = {
             "baseUrl": pi_base_urls["openai"],
             "api": "openai-responses",
-            "apiKey": token,
+            "apiKey": api_key,
             "authHeader": True,
             "headers": ua_headers,
             "models": [{"id": m} for m in codex_models],
@@ -143,7 +148,7 @@ def render_overlay(
         providers["databricks-gemini"] = {
             "baseUrl": pi_base_urls["gemini"],
             "api": "google-generative-ai",
-            "apiKey": token,
+            "apiKey": api_key,
             "authHeader": True,
             "headers": ua_headers,
             "models": [{"id": m} for m in gemini_models],
@@ -157,18 +162,31 @@ def render_overlay(
     return overlay, keys
 
 
+def build_pi_api_key(state: dict) -> str:
+    """Return the `!command` apiKey value pi resolves before every request.
+
+    Pi runs a leading-`!` config value as a command and uses its stdout, and it
+    resolves the provider apiKey per provider request rather than once per
+    process, so the token is minted on demand and never lands in the config.
+
+    No `--force-refresh`: pi has no token cache of its own 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."""
+    return "!" + build_auth_shell_command(
+        state["workspace"],
+        state.get("profile"),
+        use_pat=bool(state.get("use_pat")),
+    )
+
+
 def write_tool_config(
     state: dict,
     model: str,
     token: str | None = None,
-    *,
-    force_refresh: bool = False,
 ) -> tuple[dict, str]:
     backup_existing_file(PI_CONFIG_PATH, PI_BACKUP_PATH)
     if token is None:
-        token = get_databricks_token(
-            state["workspace"], state.get("profile"), force_refresh=force_refresh
-        )
+        token = get_databricks_token(state["workspace"], state.get("profile"))
     pi_base_urls = state.get("base_urls", {}).get("pi") or build_pi_base_urls(state["workspace"])
     managed_families = _managed_model_families(state)
     claude_models, codex_models, gemini_models = managed_families or (
@@ -178,7 +196,7 @@ def write_tool_config(
     )
     overlay, managed_keys = render_overlay(
         model,
-        token,
+        build_pi_api_key(state),
         pi_base_urls,
         claude_models,
         codex_models,
@@ -260,22 +278,14 @@ def default_model(state: dict) -> str | None:
     return gemini_models[0] if gemini_models else None
 
 
-def _refresh_token_once(state: dict, *, force_refresh: bool = False) -> str:
+def _configure_launch(state: dict) -> str:
     model = default_model(state)
     if not model:
         raise RuntimeError("No Pi model is available on this workspace.")
-    _, token = write_tool_config(state, model, force_refresh=force_refresh)
+    _, token = write_tool_config(state, model)
     return token
 
 
-def _refresh_forever(state: dict, stop_event: threading.Event) -> None:
-    while not stop_event.wait(TOKEN_REFRESH_INTERVAL_SECONDS):
-        try:
-            _refresh_token_once(state, force_refresh=True)
-        except RuntimeError:
-            continue
-
-
 def build_runtime_env(token: str) -> dict[str, str]:
     env = os.environ.copy()
     env["OAUTH_TOKEN"] = token
@@ -284,26 +294,16 @@ def build_runtime_env(token: str) -> dict[str, str]:
 
 
 def launch(state: dict, tool_args: list[str], *, options: LaunchOptions) -> None:
-    token = _refresh_token_once(state)
+    """Launch Pi; it re-resolves its apiKey command per request, so no refresher."""
+    token = _configure_launch(state)
     env = build_runtime_env(token)
 
-    stop_event = threading.Event()
-    refresher = threading.Thread(
-        target=_refresh_forever,
-        args=(state, stop_event),
-        daemon=True,
-    )
-    refresher.start()
-
     proc = subprocess.Popen([SPEC["binary"], *tool_args], env=env)
     try:
         returncode = proc.wait()
     except KeyboardInterrupt:
         proc.send_signal(signal.SIGINT)
         returncode = proc.wait()
-    finally:
-        stop_event.set()
-        refresher.join(timeout=1)
 
     raise SystemExit(returncode)
 

tests/test_agent_pi.py

@@ -131,11 +131,12 @@ def test_openai_and_gemini_have_no_compat_flags(self):
 
 
 class TestRenderOverlayAuthAndModels:
-    def test_token_in_api_key(self):
+    def test_api_key_config_value_embedded_verbatim(self):
+        # render_overlay takes a pi config value, so it must not reinterpret it.
         overlay, _ = _overlay(
-            "claude-sonnet", token="mytoken", claude_models={"sonnet": "claude-sonnet"}
+            "claude-sonnet", token="!mint.sh", claude_models={"sonnet": "claude-sonnet"}
         )
-        assert overlay["providers"]["databricks-claude"]["apiKey"] == "mytoken"
+        assert overlay["providers"]["databricks-claude"]["apiKey"] == "!mint.sh"
 
     def test_auth_header_flag_set_on_all_providers(self):
         overlay, _ = _overlay(
@@ -347,7 +348,7 @@ def test_legacy_providers_removed_on_upgrade(self, tmp_path, monkeypatch):
             assert legacy not in written_providers
         assert "databricks-claude" in written_providers
 
-    def test_config_written_with_correct_model_and_token(self, tmp_path, monkeypatch):
+    def test_config_written_with_correct_model_and_auth_command(self, tmp_path, monkeypatch):
         pi_mod, config_file, _, _ = self._setup(tmp_path, monkeypatch)
 
         with (
@@ -358,7 +359,36 @@ def test_config_written_with_correct_model_and_token(self, tmp_path, monkeypatch
 
         written = json.loads(config_file.read_text())
         assert written["model"] == "databricks-claude/claude-sonnet"
-        assert written["providers"]["databricks-claude"]["apiKey"] == "tok"
+        api_key = written["providers"]["databricks-claude"]["apiKey"]
+        assert api_key.startswith("!")
+        assert "auth-token" in api_key
+
+    def test_bearer_never_written_to_the_config(self, tmp_path, monkeypatch):
+        # The whole point: a token on disk is a token that can go stale, and pi
+        # resolves the command per request instead.
+        pi_mod, config_file, _, _ = self._setup(tmp_path, monkeypatch)
+
+        with (
+            patch("ucode.agents.pi.get_databricks_token", return_value="dapi-secret"),
+            patch("ucode.agents.pi.save_state"),
+        ):
+            pi_mod.write_tool_config(self._state(), "claude-sonnet")
+
+        assert "dapi-secret" not in config_file.read_text()
+
+    def test_every_provider_gets_the_auth_command(self, tmp_path, monkeypatch):
+        pi_mod, config_file, _, _ = self._setup(tmp_path, monkeypatch)
+        state = self._state(codex_models=["gpt-5"], gemini_models=["gemini-2"])
+
+        with (
+            patch("ucode.agents.pi.get_databricks_token", return_value="tok"),
+            patch("ucode.agents.pi.save_state"),
+        ):
+            pi_mod.write_tool_config(state, "claude-sonnet", token="tok")
+
+        providers = json.loads(config_file.read_text())["providers"]
+        for name in ("databricks-claude", "databricks-openai", "databricks-gemini"):
+            assert providers[name]["apiKey"].startswith("!"), name
 
     def test_settings_pins_default_provider_and_model(self, tmp_path, monkeypatch):
         # Without this, Pi's `findInitialModel` can fall through to a built-in
@@ -475,3 +505,29 @@ def test_pi_default_model_wins_over_allowlist(self):
     def test_falls_back_to_pi_models_without_default(self):
         state = {"pi_models": ["system.ai.claude-opus-4-8"]}
         assert pi.default_model(state) == "system.ai.claude-opus-4-8"
+
+
+class TestBuildPiApiKey:
+    """Pi resolves a leading-`!` config value as a command before every provider
+    request, so the apiKey is that command rather than a baked bearer."""
+
+    def test_is_a_pi_command_value_running_auth_token(self):
+        api_key = pi.build_pi_api_key({"workspace": WS})
+
+        assert api_key.startswith("!")
+        assert "auth-token" in api_key
+        assert f"--host {WS}" in api_key
+
+    def test_omits_force_refresh(self):
+        # Pi has no token cache on this path, so --force-refresh would round-trip
+        # to the workspace on every turn.
+        assert "--force-refresh" not in pi.build_pi_api_key({"workspace": WS})
+
+    def test_passes_the_profile_through(self):
+        api_key = pi.build_pi_api_key({"workspace": WS, "profile": "stablebox"})
+
+        assert "--profile stablebox" in api_key
+
+    def test_forwards_use_pat(self):
+        assert "--use-pat" in pi.build_pi_api_key({"workspace": WS, "use_pat": True})
+        assert "--use-pat" not in pi.build_pi_api_key({"workspace": WS})

tests/test_e2e_user_agent.py

@@ -353,6 +353,10 @@ def test_user_agent_arrives_at_gateway(self, tmp_path, monkeypatch, capture_serv
             pi.write_tool_config(state, "test-claude-model", token="test-token")
 
         env = pi.build_runtime_env("test-token")
+        # Pi now resolves its apiKey by running `ucode auth-token`. The static
+        # bearer short-circuit gives that command something to print without a
+        # real workspace behind the capture server.
+        env["DATABRICKS_BEARER"] = "test-token"
         result = _run_until_first_request(pi.validate_cmd("pi"), env)
 
         req = capture_server.first_request_with_path_prefix("/ai-gateway/anthropic")