diff --git a/src/ucode/agents/pi.py b/src/ucode/agents/pi.py index a193e8ff..e51e3285 100644 --- a/src/ucode/agents/pi.py +++ b/src/ucode/agents/pi.py @@ -21,8 +21,11 @@ `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 @@ -30,7 +33,6 @@ 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) diff --git a/tests/test_agent_pi.py b/tests/test_agent_pi.py index ff7f172d..416c9769 100644 --- a/tests/test_agent_pi.py +++ b/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}) diff --git a/tests/test_e2e_user_agent.py b/tests/test_e2e_user_agent.py index e6cec214..55f2f1d7 100644 --- a/tests/test_e2e_user_agent.py +++ b/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")