diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index 3afd4fe..a63fcca 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -342,7 +342,7 @@ def configure_tool( result = pi.write_tool_config(state, model) else: result = opencode.write_tool_config(state, model) - # gemini/opencode/copilot/pi return (state, token); codex/claude return state + # gemini/copilot return (state, token); the other tools return state. if isinstance(result, tuple): return result[0] return result diff --git a/src/ucode/agents/opencode.py b/src/ucode/agents/opencode.py index b89536e..1cffc02 100644 --- a/src/ucode/agents/opencode.py +++ b/src/ucode/agents/opencode.py @@ -1,11 +1,9 @@ -"""OpenCode agent: writes opencode.json with Databricks-backed providers.""" +"""OpenCode agent: writes Databricks providers and a request-time auth plugin.""" from __future__ import annotations +import json import os -import signal -import subprocess -import threading from ucode.agent_updates import available_npm_package_update from ucode.config_io import ( @@ -15,19 +13,21 @@ deep_merge_dict, read_json_safe, write_json_file, + write_text_file, ) from ucode.databricks import ( - TOKEN_REFRESH_INTERVAL_SECONDS, + build_auth_token_argv, build_opencode_base_urls, - get_databricks_token, model_token_limits, ) +from ucode.launcher import exec_or_spawn from ucode.state import mark_tool_managed, save_state from ucode.telemetry import agent_version, ucode_version OPENCODE_XDG_CONFIG_HOME = APP_DIR / "opencode-xdg" OPENCODE_CONFIG_DIR = OPENCODE_XDG_CONFIG_HOME / "opencode" OPENCODE_CONFIG_PATH = OPENCODE_CONFIG_DIR / "opencode.json" +OPENCODE_AUTH_PLUGIN_PATH = OPENCODE_CONFIG_DIR / "plugins" / "ucode-databricks-auth.js" OPENCODE_BACKUP_PATH = APP_DIR / "opencode-config.backup.json" SPEC: ToolSpec = { @@ -38,11 +38,8 @@ "backup_path": OPENCODE_BACKUP_PATH, } -PROVIDER_KEYS: list[list[str]] = [ - ["provider", "databricks-anthropic"], - ["provider", "databricks-google"], - ["provider", "databricks-oss"], -] +PROVIDER_NAMES = ("databricks-anthropic", "databricks-google", "databricks-oss") +PROVIDER_KEYS: list[list[str]] = [["provider", name] for name in PROVIDER_NAMES] def is_update_available() -> tuple[str, str] | None: @@ -85,12 +82,10 @@ def _oss_model_overlay(model: str, ua_header: dict[str, str]) -> dict: def render_overlay( model: str, - token: str, opencode_base_urls: dict[str, str], opencode_models: dict[str, list[str]], ) -> tuple[dict, list[list[str]]]: """Return (overlay, managed_key_paths) for opencode.json.""" - auth_headers = {"Authorization": f"Bearer {token}"} # OpenCode hardcodes `User-Agent: opencode/` in session/llm.ts for # every provider, after the AI SDK's combineHeaders. The provider-level # `headers` are clobbered by that injection, but per-model `headers` are @@ -119,8 +114,6 @@ def render_overlay( "npm": "@ai-sdk/anthropic", "options": { "baseURL": opencode_base_urls["anthropic"], - "apiKey": token, - "headers": auth_headers, }, "models": dict.fromkeys(anthropic_models, anthropic_model_overlay), } @@ -130,8 +123,6 @@ def render_overlay( "npm": "@ai-sdk/google", "options": { "baseURL": opencode_base_urls["gemini"], - "apiKey": token, - "headers": auth_headers, }, "models": {m: {"headers": ua_header} for m in gemini_models}, } @@ -141,8 +132,6 @@ def render_overlay( "npm": "@ai-sdk/openai", "options": { "baseURL": opencode_base_urls["oss"], - "apiKey": token, - "headers": auth_headers, }, "models": {m: _oss_model_overlay(m, ua_header) for m in oss_models}, } @@ -154,42 +143,63 @@ def render_overlay( return overlay, keys -def write_tool_config( - state: dict, - model: str, - token: str | None = None, - *, - force_refresh: bool = False, -) -> tuple[dict, str]: +def render_auth_plugin(auth_argv: list[str]) -> str: + """Render an OpenCode plugin that obtains a fresh bearer per request.""" + return f"""import {{ execFileSync }} from "node:child_process" + +const [command, ...args] = {json.dumps(auth_argv)} +const providers = {json.dumps(PROVIDER_NAMES)} + +function token() {{ + const value = execFileSync(command, args, {{ encoding: "utf8" }}).trim() + if (!value) throw new Error("ucode auth-token returned no access token") + return value +}} + +async function authenticatedFetch(input, init) {{ + const headers = new Headers(init?.headers) + headers.set("Authorization", `Bearer ${{token()}}`) + return fetch(input, {{ ...init, headers }}) +}} + +export const UcodeDatabricksAuth = async () => ({{ + async config(config) {{ + for (const providerName of providers) {{ + const provider = config.provider?.[providerName] + if (!provider) continue + provider.options ??= {{}} + provider.options.apiKey = "databricks-cli" + provider.options.fetch = authenticatedFetch + }} + }}, +}}) +""" + + +def write_tool_config(state: dict, model: str) -> dict: backup_existing_file(OPENCODE_CONFIG_PATH, OPENCODE_BACKUP_PATH) - if token is None: - token = get_databricks_token( - state["workspace"], state.get("profile"), force_refresh=force_refresh - ) + auth_argv = build_auth_token_argv( + state["workspace"], state.get("profile"), use_pat=bool(state.get("use_pat")) + ) + write_text_file(OPENCODE_AUTH_PLUGIN_PATH, render_auth_plugin(auth_argv)) opencode_base_urls = state.get("base_urls", {}).get("opencode") or build_opencode_base_urls( state["workspace"] ) overlay, managed_keys = render_overlay( model, - token, opencode_base_urls, state.get("opencode_models") or {}, ) existing = read_json_safe(OPENCODE_CONFIG_PATH) providers = existing.get("provider") if isinstance(providers, dict): - for stale in ( - "databricks-anthropic", - "databricks-google", - "databricks-openai", - "databricks-oss", - ): + for stale in (*PROVIDER_NAMES, "databricks-openai"): providers.pop(stale, None) merged = deep_merge_dict(existing, overlay) write_json_file(OPENCODE_CONFIG_PATH, merged) state = mark_tool_managed(state, "opencode", managed_keys) save_state(state) - return state, token + return state def build_mcp_server_entry(argv: list[str]) -> dict: @@ -239,53 +249,16 @@ def default_model(state: dict) -> str | None: return oss[0] if oss else None -def _refresh_token_once(state: dict, *, force_refresh: bool = False) -> str: - model = default_model(state) - if not model: - raise RuntimeError("No OpenCode model is configured.") - _, token = write_tool_config(state, model, force_refresh=force_refresh) - 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, state: dict | None = None) -> dict[str, str]: +def build_runtime_env() -> dict[str, str]: env = os.environ.copy() - env["OAUTH_TOKEN"] = token env["XDG_CONFIG_HOME"] = str(OPENCODE_XDG_CONFIG_HOME) return env def launch(state: dict, tool_args: list[str]) -> None: - """Launch opencode with background token refresh (same pattern as Gemini).""" - token = _refresh_token_once(state) - env = build_runtime_env(token, state) - - 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) + _ = state + os.environ["XDG_CONFIG_HOME"] = str(OPENCODE_XDG_CONFIG_HOME) + exec_or_spawn([SPEC["binary"], *tool_args]) def validate_cmd(binary: str) -> list[str]: @@ -296,4 +269,4 @@ def validate_env(state: dict) -> dict[str, str]: workspace = state.get("workspace") if not workspace: raise RuntimeError("No workspace configured.") - return build_runtime_env(get_databricks_token(workspace, state.get("profile")), state) + return build_runtime_env() diff --git a/src/ucode/agents/pi.py b/src/ucode/agents/pi.py index e7c1760..ce4f803 100644 --- a/src/ucode/agents/pi.py +++ b/src/ucode/agents/pi.py @@ -21,16 +21,14 @@ `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 delegates authentication to ``ucode auth-token`` through Pi's +request-time ``!command`` API-key resolution. Tokens stay out of configuration +and the Databricks CLI owns OAuth refresh. """ from __future__ import annotations import os -import signal -import subprocess -import threading from ucode.agent_updates import available_npm_package_update from ucode.config_io import ( @@ -42,15 +40,14 @@ write_json_file, ) from ucode.databricks import ( - TOKEN_REFRESH_INTERVAL_SECONDS, + build_auth_shell_command, build_pi_base_urls, - get_databricks_token, ) +from ucode.launcher import exec_or_spawn from ucode.state import mark_tool_managed, save_state from ucode.telemetry import agent_version, ucode_version -PI_UCODE_HOME = APP_DIR / "pi-home" -PI_CONFIG_DIR = PI_UCODE_HOME / ".pi" / "agent" +PI_CONFIG_DIR = APP_DIR / "pi-home" / ".pi" / "agent" PI_CONFIG_PATH = PI_CONFIG_DIR / "models.json" PI_SETTINGS_PATH = PI_CONFIG_DIR / "settings.json" PI_BACKUP_PATH = APP_DIR / "pi-models.backup.json" @@ -102,7 +99,7 @@ def _resolve_model_selector( def render_overlay( model: str, - token: str, + auth_command: str, pi_base_urls: dict[str, str], claude_models: dict[str, str], codex_models: list[str], @@ -120,7 +117,7 @@ def render_overlay( providers["databricks-claude"] = { "baseUrl": pi_base_urls["claude"], "api": "anthropic-messages", - "apiKey": token, + "apiKey": f"!{auth_command}", "authHeader": True, # Gateway's Anthropic translator rejects per-tool # `eager_input_streaming` on the streaming + tools path. Pi sends @@ -134,7 +131,7 @@ def render_overlay( providers["databricks-openai"] = { "baseUrl": pi_base_urls["openai"], "api": "openai-responses", - "apiKey": token, + "apiKey": f"!{auth_command}", "authHeader": True, "headers": ua_headers, "models": [{"id": m} for m in codex_models], @@ -144,7 +141,7 @@ def render_overlay( providers["databricks-gemini"] = { "baseUrl": pi_base_urls["gemini"], "api": "google-generative-ai", - "apiKey": token, + "apiKey": f"!{auth_command}", "authHeader": True, "headers": ua_headers, "models": [{"id": m} for m in gemini_models], @@ -158,22 +155,15 @@ def render_overlay( return overlay, keys -def write_tool_config( - state: dict, - model: str, - token: str | None = None, - *, - force_refresh: bool = False, -) -> tuple[dict, str]: +def write_tool_config(state: dict, model: str) -> dict: 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 - ) + auth_command = build_auth_shell_command( + state["workspace"], state.get("profile"), use_pat=bool(state.get("use_pat")) + ) pi_base_urls = state.get("base_urls", {}).get("pi") or build_pi_base_urls(state["workspace"]) overlay, managed_keys = render_overlay( model, - token, + auth_command, pi_base_urls, state.get("claude_models") or {}, state.get("codex_models") or [], @@ -189,7 +179,7 @@ def write_tool_config( _write_settings(overlay["model"]) state = mark_tool_managed(state, "pi", managed_keys) save_state(state) - return state, token + return state def _write_settings(model_selector: str) -> None: @@ -218,52 +208,16 @@ 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: - 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) - 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]: +def build_runtime_env() -> dict[str, str]: env = os.environ.copy() - env["OAUTH_TOKEN"] = token - env["HOME"] = str(PI_UCODE_HOME) + env["PI_CODING_AGENT_DIR"] = str(PI_CONFIG_DIR) return env def launch(state: dict, tool_args: list[str]) -> None: - token = _refresh_token_once(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) + _ = state + os.environ["PI_CODING_AGENT_DIR"] = str(PI_CONFIG_DIR) + exec_or_spawn([SPEC["binary"], *tool_args]) def validate_cmd(binary: str) -> list[str]: @@ -274,4 +228,4 @@ def validate_env(state: dict) -> dict[str, str]: workspace = state.get("workspace") if not workspace: raise RuntimeError("No workspace configured.") - return build_runtime_env(get_databricks_token(workspace, state.get("profile"))) + return build_runtime_env() diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 97f8ae7..0f51151 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -921,11 +921,12 @@ def auth_token_cmd( ) -> None: """Print a Databricks bearer token to stdout, then exit. - This is the cross-platform helper invoked by Claude Code's `apiKeyHelper` - and Codex's auth command on every token refresh. It is not meant for - interactive use. All token logic (DATABRICKS_BEARER short-circuit, PAT - profiles, OAuth refresh) lives in `get_databricks_token`, so the same - binary works on macOS, Linux, and Windows without any POSIX shell.""" + This is the cross-platform helper invoked by Claude, Codex, Pi, and the + ucode-managed OpenCode auth plugin whenever they need a bearer token. It is + not meant for interactive use. All token logic (DATABRICKS_BEARER + short-circuit, PAT profiles, OAuth refresh) lives in + `get_databricks_token`, so the same binary works on macOS, Linux, and + Windows without any POSIX shell.""" import sys state = load_state() @@ -1258,7 +1259,12 @@ def _launch_tool( f"{TOOL_SPECS[tool]['display']} requires one-time hook review. Open " "`/hooks` and trust the ucode routing hooks if prompted." ) - if tool in ("gemini", "opencode", "copilot", "pi"): + if tool in ("opencode", "pi"): + print_note( + f"{TOOL_SPECS[tool]['display']} obtains a fresh Databricks token " + "through `ucode auth-token` for every model request." + ) + elif tool in ("gemini", "copilot"): print_note( f"{TOOL_SPECS[tool]['display']} token refresh is managed automatically " f"every 30 minutes while the session is running." diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 1d32f31..f0a6e9f 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -1177,8 +1177,9 @@ def build_auth_shell_command( """Single-line, shell-quoted form of :func:`build_auth_token_argv`. Used where a tool wants the helper as one command *string* (Claude Code's - `apiKeyHelper`). On every platform this resolves to the `ucode auth-token` - executable rather than a POSIX shell pipeline, so no `sh`/`jq` is required.""" + `apiKeyHelper` and Pi's `!command` API key). On every platform this resolves + to the `ucode auth-token` executable rather than a POSIX shell pipeline, so + no `sh`/`jq` is required.""" argv = build_auth_token_argv(workspace, profile, use_pat=use_pat) if platform.system() == "Windows": return subprocess.list2cmdline(argv) diff --git a/tests/test_agent_opencode.py b/tests/test_agent_opencode.py index 71ca2bc..5a04bf1 100644 --- a/tests/test_agent_opencode.py +++ b/tests/test_agent_opencode.py @@ -36,88 +36,76 @@ def test_config_path_is_under_ucode_xdg_home(self): class TestRenderOverlay: def test_sets_model(self): - overlay, _ = opencode.render_overlay("claude-sonnet", "tok", _base_urls(), {}) + overlay, _ = opencode.render_overlay("claude-sonnet", _base_urls(), {}) assert overlay["model"] == "claude-sonnet" def test_anthropic_provider_added_when_models_present(self): models = {"anthropic": ["claude-sonnet"], "gemini": []} - overlay, _ = opencode.render_overlay("claude-sonnet", "tok", _base_urls(), models) + overlay, _ = opencode.render_overlay("claude-sonnet", _base_urls(), models) assert "databricks-anthropic" in overlay["provider"] def test_gemini_provider_added_when_models_present(self): models = {"anthropic": [], "gemini": ["gemini-2"]} - overlay, _ = opencode.render_overlay("gemini-2", "tok", _base_urls(), models) + overlay, _ = opencode.render_overlay("gemini-2", _base_urls(), models) assert "databricks-google" in overlay["provider"] def test_oss_provider_added_when_models_present(self): models = {"oss": ["system.ai.kimi-k2-7-code"]} - overlay, _ = opencode.render_overlay( - "system.ai.kimi-k2-7-code", "tok", _base_urls(), models - ) + overlay, _ = opencode.render_overlay("system.ai.kimi-k2-7-code", _base_urls(), models) assert "databricks-oss" in overlay["provider"] def test_oss_provider_uses_ai_sdk_openai_package(self): models = {"oss": ["system.ai.kimi-k2-7-code"]} - overlay, _ = opencode.render_overlay( - "system.ai.kimi-k2-7-code", "tok", _base_urls(), models - ) + overlay, _ = opencode.render_overlay("system.ai.kimi-k2-7-code", _base_urls(), models) assert overlay["provider"]["databricks-oss"]["npm"] == "@ai-sdk/openai" def test_both_providers_when_both_present(self): models = {"anthropic": ["claude-sonnet"], "gemini": ["gemini-2"]} - overlay, _ = opencode.render_overlay("claude-sonnet", "tok", _base_urls(), models) + overlay, _ = opencode.render_overlay("claude-sonnet", _base_urls(), models) assert "databricks-anthropic" in overlay["provider"] assert "databricks-google" in overlay["provider"] def test_no_provider_key_when_no_models(self): - overlay, _ = opencode.render_overlay("model", "tok", _base_urls(), {}) + overlay, _ = opencode.render_overlay("model", _base_urls(), {}) assert "provider" not in overlay def test_anthropic_base_url(self): models = {"anthropic": ["claude-sonnet"]} - overlay, _ = opencode.render_overlay("claude-sonnet", "tok", _base_urls(), models) + overlay, _ = opencode.render_overlay("claude-sonnet", _base_urls(), models) options = overlay["provider"]["databricks-anthropic"]["options"] assert options["baseURL"] == f"{WS}/ai-gateway/anthropic/v1" def test_gemini_base_url(self): models = {"gemini": ["gemini-2"]} - overlay, _ = opencode.render_overlay("gemini-2", "tok", _base_urls(), models) + overlay, _ = opencode.render_overlay("gemini-2", _base_urls(), models) options = overlay["provider"]["databricks-google"]["options"] assert options["baseURL"] == f"{WS}/ai-gateway/gemini/v1beta" def test_oss_base_url(self): models = {"oss": ["system.ai.kimi-k2-7-code"]} - overlay, _ = opencode.render_overlay( - "system.ai.kimi-k2-7-code", "tok", _base_urls(), models - ) + overlay, _ = opencode.render_overlay("system.ai.kimi-k2-7-code", _base_urls(), models) options = overlay["provider"]["databricks-oss"]["options"] assert options["baseURL"] == f"{WS}/ai-gateway/mlflow/v1" def test_glm_gets_token_limits(self): models = {"oss": ["system.ai.glm-5-2"]} - overlay, _ = opencode.render_overlay("system.ai.glm-5-2", "tok", _base_urls(), models) + overlay, _ = opencode.render_overlay("system.ai.glm-5-2", _base_urls(), models) glm = overlay["provider"]["databricks-oss"]["models"]["system.ai.glm-5-2"] # OpenCode's schema requires both context and output on `limit`. assert glm["limit"] == {"context": 200000, "output": 25000} def test_non_glm_oss_model_has_no_output_cap(self): models = {"oss": ["system.ai.kimi-k2-7-code"]} - overlay, _ = opencode.render_overlay( - "system.ai.kimi-k2-7-code", "tok", _base_urls(), models - ) + overlay, _ = opencode.render_overlay("system.ai.kimi-k2-7-code", _base_urls(), models) kimi = overlay["provider"]["databricks-oss"]["models"]["system.ai.kimi-k2-7-code"] assert "limit" not in kimi - def test_token_in_api_key(self): - models = {"anthropic": ["claude-sonnet"]} - overlay, _ = opencode.render_overlay("claude-sonnet", "mytoken", _base_urls(), models) - assert overlay["provider"]["databricks-anthropic"]["options"]["apiKey"] == "mytoken" - - def test_authorization_header(self): + def test_provider_options_do_not_persist_credentials(self): models = {"anthropic": ["claude-sonnet"]} - overlay, _ = opencode.render_overlay("claude-sonnet", "tok", _base_urls(), models) - headers = overlay["provider"]["databricks-anthropic"]["options"]["headers"] - assert headers["Authorization"] == "Bearer tok" + overlay, _ = opencode.render_overlay("claude-sonnet", _base_urls(), models) + options = overlay["provider"]["databricks-anthropic"]["options"] + assert "apiKey" not in options + assert "headers" not in options def test_anthropic_tool_streaming_disabled(self): # @ai-sdk/anthropic injects `eager_input_streaming: true` on tool defs, @@ -125,7 +113,7 @@ def test_anthropic_tool_streaming_disabled(self): # Claude models, so we opt out per-model. The setting must live in # `models..options` — per-call providerOptions — not provider options. models = {"anthropic": ["claude-sonnet"]} - overlay, _ = opencode.render_overlay("claude-sonnet", "tok", _base_urls(), models) + overlay, _ = opencode.render_overlay("claude-sonnet", _base_urls(), models) model_entry = overlay["provider"]["databricks-anthropic"]["models"]["claude-sonnet"] assert model_entry["options"]["toolStreaming"] is False @@ -135,7 +123,7 @@ def test_user_agent_header_anthropic(self, monkeypatch): monkeypatch.setattr(opencode, "ucode_version", lambda: "0.1.0") monkeypatch.setattr(opencode, "agent_version", lambda binary: "0.74.0") models = {"anthropic": ["claude-sonnet"]} - overlay, _ = opencode.render_overlay("claude-sonnet", "tok", _base_urls(), models) + overlay, _ = opencode.render_overlay("claude-sonnet", _base_urls(), models) model_headers = overlay["provider"]["databricks-anthropic"]["models"]["claude-sonnet"][ "headers" ] @@ -145,60 +133,49 @@ def test_user_agent_header_gemini(self, monkeypatch): monkeypatch.setattr(opencode, "ucode_version", lambda: "0.1.0") monkeypatch.setattr(opencode, "agent_version", lambda binary: "0.74.0") models = {"gemini": ["gemini-2"]} - overlay, _ = opencode.render_overlay("gemini-2", "tok", _base_urls(), models) + overlay, _ = opencode.render_overlay("gemini-2", _base_urls(), models) model_headers = overlay["provider"]["databricks-google"]["models"]["gemini-2"]["headers"] assert model_headers["User-Agent"] == "ucode/0.1.0 opencode/0.74.0" - def test_provider_level_headers_only_authorization(self, monkeypatch): - # Sanity: provider-level headers should NOT include User-Agent (since - # it's clobbered there) — only Authorization. - models = {"anthropic": ["claude-sonnet"]} - overlay, _ = opencode.render_overlay("claude-sonnet", "tok", _base_urls(), models) - provider_headers = overlay["provider"]["databricks-anthropic"]["options"]["headers"] - assert "User-Agent" not in provider_headers - assert provider_headers["Authorization"] == "Bearer tok" - def test_managed_keys_include_model(self): - _, keys = opencode.render_overlay("model", "tok", _base_urls(), {}) + _, keys = opencode.render_overlay("model", _base_urls(), {}) assert ["model"] in keys def test_managed_keys_include_anthropic_provider(self): models = {"anthropic": ["claude-sonnet"]} - _, keys = opencode.render_overlay("claude-sonnet", "tok", _base_urls(), models) + _, keys = opencode.render_overlay("claude-sonnet", _base_urls(), models) assert ["provider", "databricks-anthropic"] in keys def test_managed_keys_include_gemini_provider(self): models = {"gemini": ["gemini-2"]} - _, keys = opencode.render_overlay("gemini-2", "tok", _base_urls(), models) + _, keys = opencode.render_overlay("gemini-2", _base_urls(), models) assert ["provider", "databricks-google"] in keys def test_managed_keys_include_oss_provider(self): models = {"oss": ["system.ai.kimi-k2-7-code"]} - _, keys = opencode.render_overlay("system.ai.kimi-k2-7-code", "tok", _base_urls(), models) + _, keys = opencode.render_overlay("system.ai.kimi-k2-7-code", _base_urls(), models) assert ["provider", "databricks-oss"] in keys def test_anthropic_models_listed(self): models = {"anthropic": ["claude-sonnet", "claude-haiku"]} - overlay, _ = opencode.render_overlay("claude-sonnet", "tok", _base_urls(), models) + overlay, _ = opencode.render_overlay("claude-sonnet", _base_urls(), models) provider_models = overlay["provider"]["databricks-anthropic"]["models"] assert "claude-sonnet" in provider_models assert "claude-haiku" in provider_models def test_prefixes_anthropic_model_with_provider_id(self): models = {"anthropic": ["claude-sonnet"], "gemini": []} - overlay, _ = opencode.render_overlay("claude-sonnet", "tok", _base_urls(), models) + overlay, _ = opencode.render_overlay("claude-sonnet", _base_urls(), models) assert overlay["model"] == "databricks-anthropic/claude-sonnet" def test_prefixes_gemini_model_with_provider_id(self): models = {"anthropic": [], "gemini": ["gemini-2"]} - overlay, _ = opencode.render_overlay("gemini-2", "tok", _base_urls(), models) + overlay, _ = opencode.render_overlay("gemini-2", _base_urls(), models) assert overlay["model"] == "databricks-google/gemini-2" def test_prefixes_oss_model_with_provider_id(self): models = {"oss": ["system.ai.kimi-k2-7-code"]} - overlay, _ = opencode.render_overlay( - "system.ai.kimi-k2-7-code", "tok", _base_urls(), models - ) + overlay, _ = opencode.render_overlay("system.ai.kimi-k2-7-code", _base_urls(), models) assert overlay["model"] == "databricks-oss/system.ai.kimi-k2-7-code" @@ -294,17 +271,29 @@ def test_removes_mcp_server_without_clobbering_others(self, tmp_path, monkeypatc class TestBuildRuntimeEnv: - def test_sets_oauth_token_for_mcp(self): - env = opencode.build_runtime_env("tok") - - assert env["OAUTH_TOKEN"] == "tok" - def test_sets_ucode_xdg_config_home(self): - env = opencode.build_runtime_env("tok") + env = opencode.build_runtime_env() assert env["XDG_CONFIG_HOME"] == str(opencode.OPENCODE_XDG_CONFIG_HOME) +class TestRenderAuthPlugin: + def test_calls_ucode_auth_token_and_injects_bearer(self): + plugin = opencode.render_auth_plugin( + ["/opt/ucode", "auth-token", "--host", WS, "--profile", "DEFAULT"] + ) + + assert '[command, ...args] = ["/opt/ucode", "auth-token"' in plugin + assert '"auth-token", "--host"' in plugin + assert 'headers.set("Authorization", `Bearer ${token()}`)' in plugin + + def test_targets_all_ucode_provider_names(self): + plugin = opencode.render_auth_plugin(["ucode", "auth-token", "--host", WS]) + + for provider_name in opencode.PROVIDER_NAMES: + assert provider_name in plugin + + class TestOpencodeDefaultModel: def test_prefers_anthropic(self): state = {"opencode_models": {"anthropic": ["claude-sonnet"], "gemini": ["gemini-2"]}} @@ -351,8 +340,10 @@ def test_stale_providers_removed_before_merge(self, tmp_path, monkeypatch): monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) config_file = tmp_path / "opencode.json" backup_file = tmp_path / "opencode-backup.json" + plugin_file = tmp_path / "plugins" / "ucode-databricks-auth.js" monkeypatch.setattr(oc_mod, "OPENCODE_CONFIG_PATH", config_file) monkeypatch.setattr(oc_mod, "OPENCODE_BACKUP_PATH", backup_file) + monkeypatch.setattr(oc_mod, "OPENCODE_AUTH_PLUGIN_PATH", plugin_file) stale = { "provider": { @@ -371,10 +362,13 @@ def test_stale_providers_removed_before_merge(self, tmp_path, monkeypatch): } with ( - patch("ucode.agents.opencode.get_databricks_token", return_value="tok"), + patch( + "ucode.agents.opencode.build_auth_token_argv", + return_value=["/opt/ucode", "auth-token", "--host", WS], + ), patch("ucode.agents.opencode.save_state"), ): - oc_mod.write_tool_config(state, "claude-sonnet", token="tok") + oc_mod.write_tool_config(state, "claude-sonnet") written = json.loads(config_file.read_text()) providers = written.get("provider", {}) @@ -382,6 +376,7 @@ def test_stale_providers_removed_before_merge(self, tmp_path, monkeypatch): assert providers.get("databricks-anthropic") != {"old": True} # unmanaged provider entry survives assert providers.get("other-provider") == {"keep": True} + assert "/opt/ucode" in plugin_file.read_text(encoding="utf-8") def test_config_written_with_correct_model(self, tmp_path, monkeypatch): import ucode.agents.opencode as oc_mod @@ -390,8 +385,10 @@ def test_config_written_with_correct_model(self, tmp_path, monkeypatch): monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) config_file = tmp_path / "opencode.json" backup_file = tmp_path / "opencode-backup.json" + plugin_file = tmp_path / "plugins" / "ucode-databricks-auth.js" monkeypatch.setattr(oc_mod, "OPENCODE_CONFIG_PATH", config_file) monkeypatch.setattr(oc_mod, "OPENCODE_BACKUP_PATH", backup_file) + monkeypatch.setattr(oc_mod, "OPENCODE_AUTH_PLUGIN_PATH", plugin_file) state = { "workspace": WS, @@ -401,10 +398,13 @@ def test_config_written_with_correct_model(self, tmp_path, monkeypatch): } with ( - patch("ucode.agents.opencode.get_databricks_token", return_value="tok"), + patch( + "ucode.agents.opencode.build_auth_token_argv", + return_value=["/opt/ucode", "auth-token", "--host", WS], + ), patch("ucode.agents.opencode.save_state"), ): - oc_mod.write_tool_config(state, "claude-sonnet", token="tok") + oc_mod.write_tool_config(state, "claude-sonnet") written = json.loads(config_file.read_text()) assert written["model"] == "databricks-anthropic/claude-sonnet" diff --git a/tests/test_agent_pi.py b/tests/test_agent_pi.py index 0afc5fb..7ebce55 100644 --- a/tests/test_agent_pi.py +++ b/tests/test_agent_pi.py @@ -29,12 +29,12 @@ def _empty() -> dict: } -def _overlay(model: str, token: str = "tok", **kwargs): +def _overlay(model: str, auth_command: str = "ucode auth-token", **kwargs): """Wrapper to call render_overlay with sensible defaults so tests stay terse.""" bundle = {**_empty(), **kwargs} return pi.render_overlay( model, - token, + auth_command, _base_urls(), bundle["claude_models"], bundle["codex_models"], @@ -55,7 +55,7 @@ def test_display(self): def test_config_path_under_pi_agent_dir(self): assert pi.SPEC["config_path"].name == "models.json" assert pi.SPEC["config_path"].parent.name == "agent" - assert pi.PI_UCODE_HOME in pi.SPEC["config_path"].parents + assert pi.SPEC["config_path"].parent == pi.PI_CONFIG_DIR class TestRenderOverlayProviders: @@ -131,11 +131,15 @@ def test_openai_and_gemini_have_no_compat_flags(self): class TestRenderOverlayAuthAndModels: - def test_token_in_api_key(self): + def test_auth_command_in_api_key(self): overlay, _ = _overlay( - "claude-sonnet", token="mytoken", claude_models={"sonnet": "claude-sonnet"} + "claude-sonnet", + auth_command="/opt/ucode auth-token --host https://example.databricks.com", + claude_models={"sonnet": "claude-sonnet"}, + ) + assert overlay["providers"]["databricks-claude"]["apiKey"] == ( + "!/opt/ucode auth-token --host https://example.databricks.com" ) - assert overlay["providers"]["databricks-claude"]["apiKey"] == "mytoken" def test_auth_header_flag_set_on_all_providers(self): overlay, _ = _overlay( @@ -236,13 +240,19 @@ def test_returns_none_when_empty(self): class TestBuildRuntimeEnv: - def test_sets_oauth_token(self): - env = pi.build_runtime_env("tok") - assert env["OAUTH_TOKEN"] == "tok" + def test_does_not_set_oauth_token(self, monkeypatch): + monkeypatch.delenv("OAUTH_TOKEN", raising=False) + env = pi.build_runtime_env() + assert "OAUTH_TOKEN" not in env - def test_sets_ucode_home(self): - env = pi.build_runtime_env("tok") - assert env["HOME"] == str(pi.PI_UCODE_HOME) + def test_sets_pi_config_dir(self): + env = pi.build_runtime_env() + assert env["PI_CODING_AGENT_DIR"] == str(pi.PI_CONFIG_DIR) + + def test_preserves_home_for_databricks_auth(self, monkeypatch): + monkeypatch.setenv("HOME", "/Users/example") + env = pi.build_runtime_env() + assert env["HOME"] == "/Users/example" class TestPiValidateCmd: @@ -302,10 +312,13 @@ def test_stale_managed_providers_removed_before_merge(self, tmp_path, monkeypatc config_file.write_text(json.dumps(stale), encoding="utf-8") with ( - patch("ucode.agents.pi.get_databricks_token", return_value="tok"), + patch( + "ucode.agents.pi.build_auth_shell_command", + return_value="/opt/ucode auth-token --host https://example.databricks.com", + ), patch("ucode.agents.pi.save_state"), ): - pi_mod.write_tool_config(self._state(), "claude-sonnet", token="tok") + pi_mod.write_tool_config(self._state(), "claude-sonnet") written = json.loads(config_file.read_text()) providers = written.get("providers", {}) @@ -333,28 +346,36 @@ def test_legacy_providers_removed_on_upgrade(self, tmp_path, monkeypatch): ) with ( - patch("ucode.agents.pi.get_databricks_token", return_value="tok"), + patch( + "ucode.agents.pi.build_auth_shell_command", + return_value="/opt/ucode auth-token --host https://example.databricks.com", + ), patch("ucode.agents.pi.save_state"), ): - pi_mod.write_tool_config(self._state(), "claude-sonnet", token="tok") + pi_mod.write_tool_config(self._state(), "claude-sonnet") written_providers = json.loads(config_file.read_text()).get("providers", {}) for legacy in ("databricks-anthropic", "databricks-codex", "databricks-oss"): 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 ( - patch("ucode.agents.pi.get_databricks_token", return_value="tok"), + patch( + "ucode.agents.pi.build_auth_shell_command", + return_value="/opt/ucode auth-token --host https://example.databricks.com", + ), patch("ucode.agents.pi.save_state"), ): - pi_mod.write_tool_config(self._state(), "claude-sonnet", token="tok") + pi_mod.write_tool_config(self._state(), "claude-sonnet") written = json.loads(config_file.read_text()) assert written["model"] == "databricks-claude/claude-sonnet" - assert written["providers"]["databricks-claude"]["apiKey"] == "tok" + assert written["providers"]["databricks-claude"]["apiKey"] == ( + "!/opt/ucode auth-token --host https://example.databricks.com" + ) def test_settings_pins_default_provider_and_model(self, tmp_path, monkeypatch): # Without this, Pi's `findInitialModel` can fall through to a built-in @@ -363,10 +384,13 @@ def test_settings_pins_default_provider_and_model(self, tmp_path, monkeypatch): pi_mod, _, settings_file, _ = self._setup(tmp_path, monkeypatch) with ( - patch("ucode.agents.pi.get_databricks_token", return_value="tok"), + patch( + "ucode.agents.pi.build_auth_shell_command", + return_value="/opt/ucode auth-token --host https://example.databricks.com", + ), patch("ucode.agents.pi.save_state"), ): - pi_mod.write_tool_config(self._state(), "claude-sonnet", token="tok") + pi_mod.write_tool_config(self._state(), "claude-sonnet") settings = json.loads(settings_file.read_text()) assert settings["defaultProvider"] == "databricks-claude" @@ -380,10 +404,13 @@ def test_pre_existing_settings_are_backed_up_before_first_write(self, tmp_path, settings_file.write_text(original, encoding="utf-8") with ( - patch("ucode.agents.pi.get_databricks_token", return_value="tok"), + patch( + "ucode.agents.pi.build_auth_shell_command", + return_value="/opt/ucode auth-token --host https://example.databricks.com", + ), patch("ucode.agents.pi.save_state"), ): - pi_mod.write_tool_config(self._state(), "claude-sonnet", token="tok") + pi_mod.write_tool_config(self._state(), "claude-sonnet") assert settings_backup_file.read_text(encoding="utf-8") == original # The on-disk settings still get the ucode pin applied via deep_merge. diff --git a/tests/test_e2e.py b/tests/test_e2e.py index 025ced4..3ff6277 100644 --- a/tests/test_e2e.py +++ b/tests/test_e2e.py @@ -746,9 +746,11 @@ def test_launch_opencode_per_model( monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) xdg = tmp_path / "opencode-xdg" config_path = xdg / "opencode" / "opencode.json" + plugin_path = xdg / "opencode" / "plugins" / "ucode-databricks-auth.js" backup_path = tmp_path / "opencode-config.backup.json" monkeypatch.setattr(opencode, "OPENCODE_XDG_CONFIG_HOME", xdg) monkeypatch.setattr(opencode, "OPENCODE_CONFIG_PATH", config_path) + monkeypatch.setattr(opencode, "OPENCODE_AUTH_PLUGIN_PATH", plugin_path) monkeypatch.setattr(opencode, "OPENCODE_BACKUP_PATH", backup_path) import sys @@ -763,21 +765,18 @@ def test_launch_opencode_per_model( with pytest.MonkeyPatch().context() as mp: mp.setattr("ucode.state.save_state", lambda s: None) - mp.setattr( - "ucode.agents.opencode.get_databricks_token", - lambda ws, profile=None, **kwargs: e2e_token, - ) opencode.write_tool_config( {**e2e_state, "workspace": e2e_workspace}, model, - token=e2e_token, ) cmd = opencode.validate_cmd("opencode") print(f"[opencode-per-model] -> {provider}/{model}", flush=True) t0 = time.monotonic() try: - result = _run_agent(cmd, env=opencode.build_runtime_env(e2e_token), timeout=180) + env = opencode.build_runtime_env() + env["DATABRICKS_BEARER"] = e2e_token + result = _run_agent(cmd, env=env, timeout=180) except subprocess.TimeoutExpired as exc: elapsed = time.monotonic() - t0 partial_stdout = (exc.stdout or b"").decode("utf-8", errors="replace") @@ -911,15 +910,16 @@ def test_launch_pi_per_model(self, tmp_path, monkeypatch, e2e_state, e2e_workspa pytest.skip("No Pi-compatible models available on this workspace") monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) - # Pi reads models.json below HOME/.pi/agent. Point both pi's runtime - # HOME and our writer at the same isolated tmp home. - pi_home = tmp_path / "pi-home" - pi_dir = pi_home / ".pi" / "agent" + # Point Pi's native config-directory override and our writer at the + # same isolated directory. + pi_dir = tmp_path / "pi-home" / ".pi" / "agent" config_path = pi_dir / "models.json" backup_path = tmp_path / "pi-models.backup.json" - monkeypatch.setattr(pi, "PI_UCODE_HOME", pi_home) + monkeypatch.setattr(pi, "PI_CONFIG_DIR", pi_dir) monkeypatch.setattr(pi, "PI_CONFIG_PATH", config_path) + monkeypatch.setattr(pi, "PI_SETTINGS_PATH", pi_dir / "settings.json") monkeypatch.setattr(pi, "PI_BACKUP_PATH", backup_path) + monkeypatch.setattr(pi, "PI_SETTINGS_BACKUP_PATH", tmp_path / "pi-settings.backup.json") failures = [] for family, model in models: @@ -928,17 +928,13 @@ def test_launch_pi_per_model(self, tmp_path, monkeypatch, e2e_state, e2e_workspa with pytest.MonkeyPatch().context() as mp: mp.setattr("ucode.state.save_state", lambda s: None) - mp.setattr( - "ucode.agents.pi.get_databricks_token", - lambda ws, profile=None, **kwargs: e2e_token, - ) pi.write_tool_config( {**e2e_state, "workspace": e2e_workspace}, model, - token=e2e_token, ) - env = pi.build_runtime_env(e2e_token) + env = pi.build_runtime_env() + env["DATABRICKS_BEARER"] = e2e_token cmd = pi.validate_cmd("pi") result = _run_agent(cmd, env=env, timeout=120) combined = (result.stdout + result.stderr).strip() @@ -1055,8 +1051,8 @@ def test_subprocess_initialize_list_and_call(self, e2e_state, e2e_workspace): # These tests verify that when Databricks auth fails (empty token), the agents # recover by re-authenticating rather than hanging or crashing. # -# Claude uses apiKeyHelper (shell command called by Claude Code on each refresh). -# Gemini/OpenCode/Copilot use get_databricks_token() at launch and on refresh. +# Claude, OpenCode, and Pi invoke `ucode auth-token` from their request path. +# Gemini and Copilot use get_databricks_token() at launch and on refresh. # --------------------------------------------------------------------------- diff --git a/tests/test_e2e_user_agent.py b/tests/test_e2e_user_agent.py index 884e663..f781507 100644 --- a/tests/test_e2e_user_agent.py +++ b/tests/test_e2e_user_agent.py @@ -248,8 +248,11 @@ def test_user_agent_arrives_at_gateway(self, tmp_path, monkeypatch, capture_serv opencode_dir = xdg / "opencode" opencode_dir.mkdir(parents=True) config_path = opencode_dir / "opencode.json" + plugin_path = opencode_dir / "plugins" / "ucode-databricks-auth.js" monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) + monkeypatch.setattr(opencode, "OPENCODE_XDG_CONFIG_HOME", xdg) monkeypatch.setattr(opencode, "OPENCODE_CONFIG_PATH", config_path) + monkeypatch.setattr(opencode, "OPENCODE_AUTH_PLUGIN_PATH", plugin_path) monkeypatch.setattr(opencode, "OPENCODE_BACKUP_PATH", tmp_path / "opencode.backup.json") # Construct a state with localhost base URLs so render_overlay points @@ -266,14 +269,12 @@ def test_user_agent_arrives_at_gateway(self, tmp_path, monkeypatch, capture_serv } with pytest.MonkeyPatch().context() as mp: mp.setattr("ucode.state.save_state", lambda s: None) - mp.setattr( - "ucode.agents.opencode.get_databricks_token", - lambda ws, profile=None, **kwargs: "test-token", - ) - opencode.write_tool_config(state, "test-claude-model", token="test-token") + opencode.write_tool_config(state, "test-claude-model") - env = {**os.environ, "OAUTH_TOKEN": "test-token", "XDG_CONFIG_HOME": str(xdg)} - result = _run_until_first_request(opencode.validate_cmd("opencode"), env) + env = opencode.build_runtime_env() + env["DATABRICKS_BEARER"] = "test-token" + # A fresh OpenCode XDG home installs its plugin runtime on first use. + result = _run_until_first_request(opencode.validate_cmd("opencode"), env, timeout=45) req = capture_server.first_request_with_path_prefix("/ai-gateway/anthropic") assert req is not None, _no_request_msg(capture_server, result) @@ -286,6 +287,7 @@ def test_user_agent_arrives_at_gateway(self, tmp_path, monkeypatch, capture_serv assert ua.startswith(expected_prefix), ( f"OpenCode UA missing ucode prefix.\n got: {ua!r}\n prefix: {expected_prefix!r}" ) + assert req.headers.get("Authorization") == "Bearer test-token" class TestGeminiUserAgent: @@ -319,14 +321,15 @@ def test_user_agent_arrives_at_gateway(self, tmp_path, monkeypatch, capture_serv from ucode.agents import pi _require_binary("pi") - pi_home = tmp_path / "pi-home" - pi_dir = pi_home / ".pi" / "agent" + pi_dir = tmp_path / "pi-home" / ".pi" / "agent" config_path = pi_dir / "models.json" monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) - monkeypatch.setattr(pi, "PI_UCODE_HOME", pi_home) + monkeypatch.setattr(pi, "PI_CONFIG_DIR", pi_dir) monkeypatch.setattr(pi, "PI_CONFIG_PATH", config_path) + monkeypatch.setattr(pi, "PI_SETTINGS_PATH", pi_dir / "settings.json") monkeypatch.setattr(pi, "PI_BACKUP_PATH", tmp_path / "pi.backup.json") + monkeypatch.setattr(pi, "PI_SETTINGS_BACKUP_PATH", tmp_path / "pi-settings.backup.json") state = { "workspace": capture_server.base_url, @@ -343,15 +346,13 @@ def test_user_agent_arrives_at_gateway(self, tmp_path, monkeypatch, capture_serv } with pytest.MonkeyPatch().context() as mp: mp.setattr("ucode.state.save_state", lambda s: None) - mp.setattr( - "ucode.agents.pi.get_databricks_token", - lambda ws, profile=None, **kwargs: "test-token", - ) - pi.write_tool_config(state, "test-claude-model", token="test-token") + pi.write_tool_config(state, "test-claude-model") - env = pi.build_runtime_env("test-token") + env = pi.build_runtime_env() + 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") assert req is not None, _no_request_msg(capture_server, result) _assert_ua(req, _expected_ua("pi", "pi")) + assert req.headers.get("Authorization") == "Bearer test-token"