diff --git a/src/ucode/agents/opencode.py b/src/ucode/agents/opencode.py index 4b614f3..dadf404 100644 --- a/src/ucode/agents/opencode.py +++ b/src/ucode/agents/opencode.py @@ -2,6 +2,7 @@ from __future__ import annotations +import json import os import signal import subprocess @@ -15,9 +16,11 @@ 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, @@ -31,6 +34,18 @@ OPENCODE_BACKUP_PATH = APP_DIR / "opencode-config.backup.json" OPENCODE_MCP_AUTH_HEADER_VALUE = "Bearer {env:OAUTH_TOKEN}" +# Axis B of #190: opencode bakes provider.options.headers.Authorization into +# opencode.json once at `ucode configure`/refresh time and never re-reads it +# for the life of the process, so a long-lived session keeps sending a +# bearer token that may since have been rotated or revoked. This +# ucode-managed plugin's `chat.headers` hook overrides that static header on +# every chat request with a freshly fetched (and short-TTL-cached) token. +# The static opencode.json header stays in place as the bootstrap/fallback +# the plugin fails open to if the refresh call errors. +OPENCODE_PLUGIN_DIR = OPENCODE_CONFIG_DIR / "plugin" +OPENCODE_PLUGIN_PATH = OPENCODE_PLUGIN_DIR / "ucode-databricks-auth.js" +AUTH_PLUGIN_TOKEN_TTL_SECONDS = 60 + SPEC: ToolSpec = { "binary": "opencode", "package": "opencode-ai", @@ -155,6 +170,127 @@ def render_overlay( return overlay, keys +def _managed_provider_ids(managed_keys: list[list[str]]) -> list[str]: + """Provider ids actually configured, derived from render_overlay's managed_keys. + + `managed_keys` entries are key paths like `["provider", "databricks-anthropic"]` + for each provider render_overlay populated (only providers with at least one + configured model are included); this just picks those out.""" + return [key[1] for key in managed_keys if len(key) == 2 and key[0] == "provider"] + + +def render_auth_plugin(workspace: str, profile: str | None, provider_ids: list[str]) -> str: + """Return the JS source for the ucode-managed opencode `chat.headers` plugin. + + On every chat request, the plugin sets a freshly fetched + `Authorization: Bearer ` header for our Databricks providers + (identified by `provider_ids`), overriding the static header baked into + opencode.json. The token comes from `ucode auth-token` -- the same + cross-platform helper used elsewhere (see `build_auth_token_argv`) -- but + is cached in-process for `AUTH_PLUGIN_TOKEN_TTL_SECONDS` so it doesn't + shell out on every request; a single in-flight promise coalesces + concurrent requests within the process. On any error, the hook leaves + `output.headers` untouched (fail open to the static bootstrap token).""" + argv = build_auth_token_argv(workspace, profile) + argv_literal = json.dumps(argv) + provider_ids_literal = json.dumps(sorted(provider_ids)) + ttl_ms = AUTH_PLUGIN_TOKEN_TTL_SECONDS * 1000 + + return f"""\ +// Generated and managed by ucode -- do not edit by hand. Overwritten on every +// `ucode configure` / token refresh. +// +// Fixes #190: opencode reads provider.options.headers.Authorization from +// opencode.json once at startup and never again, so a long-lived session +// keeps sending a token that may since have been rotated or revoked. This +// `chat.headers` hook sets a fresh Authorization header on every chat +// request for our Databricks providers, overriding the static header. + +import {{ execFile }} from "node:child_process"; +import {{ promisify }} from "node:util"; + +const execFileAsync = promisify(execFile); + +const UCODE_AUTH_TOKEN_ARGV = {argv_literal}; +const MANAGED_PROVIDER_IDS = new Set({provider_ids_literal}); +const TTL_MS = {ttl_ms}; + +let cachedToken = null; +let cachedAt = 0; +let inflight = null; + +async function fetchToken() {{ + const [cmd, ...args] = UCODE_AUTH_TOKEN_ARGV; + const {{ stdout }} = await execFileAsync(cmd, args); + const token = stdout.trim(); + if (!token) {{ + // `ucode auth-token` exited 0 but printed nothing usable -- treat as a + // failure so the caller's catch block fails open to the static + // bootstrap token instead of caching/sending an empty bearer. + throw new Error("ucode auth-token returned empty output"); + }} + return token; +}} + +async function getToken() {{ + const now = Date.now(); + if (cachedToken && now - cachedAt < TTL_MS) {{ + return cachedToken; + }} + if (inflight) {{ + return inflight; + }} + inflight = fetchToken() + .then((token) => {{ + cachedToken = token; + cachedAt = Date.now(); + return token; + }}) + .finally(() => {{ + inflight = null; + }}); + return inflight; +}} + +export const UcodeDatabricksAuth = async (_ctx) => ({{ + "chat.headers": async (input, output) => {{ + // NOTE: the hook's `input.provider` is the runtime Provider record + // itself (id/name/env/options/source/models), not the ProviderContext + // shape ({{source, info, options}}) some type surfaces suggest -- `id` + // lives directly on `input.provider`, not `input.provider.info.id`. + // Confirmed empirically against opencode 1.17.10; verify against a real + // hook invocation before changing this path again (#190 follow-up). + const providerId = input.provider?.id; + if (!providerId || !MANAGED_PROVIDER_IDS.has(providerId)) {{ + return; + }} + try {{ + const token = await getToken(); + // Defensive: ensure `output.headers` exists even if this opencode + // version ever invokes the hook without a pre-populated headers + // object, so the header still gets set rather than throwing (which + // would be caught below and silently no-op the refresh). + output.headers = output.headers || {{}}; + output.headers["Authorization"] = "Bearer " + token; + }} catch (err) {{ + console.error("[ucode] failed to refresh Databricks token:", err); + // Fail open: leave the static bootstrap token from opencode.json. + }} + }}, +}}); +""" + + +def write_auth_plugin(state: dict, provider_ids: list[str]) -> None: + """Write the ucode-managed opencode auth plugin (see `render_auth_plugin`). + + Always writes -- including with an empty `provider_ids` -- so a + workspace/model change that drops the last provider still overwrites a + stale plugin rather than leaving one referencing removed provider ids.""" + plugin_js = render_auth_plugin(state["workspace"], state.get("profile"), provider_ids) + write_text_file(OPENCODE_PLUGIN_PATH, plugin_js) + + def write_tool_config( state: dict, model: str, @@ -188,6 +324,9 @@ def write_tool_config( providers.pop(stale, None) merged = deep_merge_dict(existing, overlay) write_json_file(OPENCODE_CONFIG_PATH, merged) + # Plugin is the live/per-request override; opencode.json (above) remains + # the static bootstrap/fallback the plugin fails open to. See #190. + write_auth_plugin(state, _managed_provider_ids(managed_keys)) state = mark_tool_managed(state, "opencode", managed_keys) save_state(state) return state, token diff --git a/src/ucode/cli.py b/src/ucode/cli.py index b2f23be..64e7627 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -28,6 +28,7 @@ launch as launch_agent, ) from ucode.agents.codex import revert_legacy_shared_config +from ucode.agents.opencode import OPENCODE_PLUGIN_PATH from ucode.agents.pi import PI_SETTINGS_BACKUP_PATH, PI_SETTINGS_PATH from ucode.config_io import restore_file, set_dry_run from ucode.databricks import ( @@ -676,6 +677,29 @@ def status() -> int: return 0 +def _remove_opencode_auth_plugin() -> bool: + """Delete the ucode-managed opencode auth plugin file (see agents/opencode.py). + + Unlike opencode.json, this file has no pre-ucode content to restore -- + it's entirely ucode-generated -- so revert just deletes it rather than + going through the backup/restore_file machinery (#190). + + Raises ``RuntimeError`` if the file exists but can't be removed (e.g. + permissions), matching ``restore_file``'s convention elsewhere in + revert -- silently swallowing that would leave `ucode revert` reporting + "unchanged" and continuing as if nothing were wrong, with no actionable + error for a state that's actually only partially reverted.""" + try: + if OPENCODE_PLUGIN_PATH.exists(): + OPENCODE_PLUGIN_PATH.unlink() + return True + return False + except OSError as exc: + raise RuntimeError( + f"Failed to remove OpenCode auth plugin at {OPENCODE_PLUGIN_PATH}" + ) from exc + + def revert() -> int: state = load_state() managed_configs = state.get("managed_configs") or {} @@ -690,6 +714,7 @@ def revert() -> int: pi_settings_restored = restore_file( PI_SETTINGS_PATH, PI_SETTINGS_BACKUP_PATH, bool(managed_configs.get("pi")) ) + opencode_plugin_removed = _remove_opencode_auth_plugin() # Older Codex (< 0.134.0) had ucode edit the shared ~/.codex/config.toml in # place; restoring the per-profile file above does not undo that. legacy_codex_stripped = revert_legacy_shared_config() @@ -702,6 +727,7 @@ def revert() -> int: if legacy_codex_stripped: print_kv("Codex shared config", "ucode entries removed") print_kv("Pi settings", "restored" if pi_settings_restored else "unchanged") + print_kv("OpenCode auth plugin", "removed" if opencode_plugin_removed else "unchanged") for client, spec in MCP_CLIENTS.items(): print_kv( f"{spec['display']} MCP config", diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 9f89c6b..e867f54 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -4,7 +4,10 @@ from __future__ import annotations import configparser +import contextlib import functools +import hashlib +import importlib import json import logging import logging.handlers @@ -15,6 +18,7 @@ import shutil import subprocess import time +from collections.abc import Callable from concurrent.futures import ( ThreadPoolExecutor, as_completed, @@ -23,6 +27,7 @@ TimeoutError as FutureTimeoutError, ) from pathlib import Path +from types import ModuleType from typing import Literal, cast, overload from urllib import error as urllib_error from urllib import request as urllib_request @@ -42,6 +47,15 @@ spinner, ) +# `fcntl` (POSIX file locking) doesn't exist on Windows. Load it dynamically +# so we degrade to no cross-process locking there instead of failing to +# import this module at all; see `_refresh_lock`. +fcntl: ModuleType | None +try: + fcntl = importlib.import_module("fcntl") +except ImportError: + fcntl = None + UNIX_DATABRICKS_INSTALL_URL = ( "https://raw.githubusercontent.com/databricks/setup-cli/main/install.sh" ) @@ -51,6 +65,19 @@ AI_GATEWAY_V2_DOCS_URL = "https://docs.databricks.com/aws/en/ai-gateway/overview-beta" MIN_DATABRICKS_CLI_VERSION = (0, 298, 0) TOKEN_REFRESH_INTERVAL_SECONDS = 1800 +# A `--force-refresh` within this window of a peer process's successful +# force-refresh (same workspace+profile) is considered redundant: the OAuth +# refresh token is rotating, so redeeming it twice in quick succession from +# concurrent `ucode opencode` sessions can revoke the token family (#190). +TOKEN_REFRESH_COALESCE_SECONDS = 60 +# Must exceed the `databricks auth token` subprocess timeout (`timeout=15` at +# the `run(...)` call in the `_fetch` closure below), so a live peer that is +# still inside its own (up to 15s) --force-refresh subprocess is guaranteed +# to finish and touch the sentinel before any waiter gives up on the lock and +# proceeds unlocked (#190: an unlocked waiter racing a still-refreshing peer +# can redeem the rotating refresh token twice). +_LOCK_ACQUIRE_TIMEOUT_SECONDS = 30 +_LOCK_POLL_INTERVAL_SECONDS = 0.2 def _debug_enabled() -> bool: @@ -776,6 +803,155 @@ def ensure_databricks_auth(workspace: str, profile: str | None = None) -> None: run_databricks_login(workspace, profile) +def _refresh_lock_paths(workspace: str, profile: str | None) -> tuple[Path, Path]: + """Per-workspace+profile lock file and sentinel path, both under APP_DIR. + + Hashing the workspace+profile pair keeps concurrent refreshes for + *different* workspaces from serializing against each other, while still + giving every process refreshing the *same* workspace+profile a shared, + stable path to flock.""" + digest = hashlib.sha256(f"{workspace}|{profile or ''}".encode()).hexdigest()[:16] + lock_path = APP_DIR / f".token-refresh-{digest}.lock" + sentinel_path = Path(f"{lock_path}.ts") + return lock_path, sentinel_path + + +def _sentinel_is_fresh(sentinel_path: Path) -> bool: + """Whether `sentinel_path`'s mtime is within TOKEN_REFRESH_COALESCE_SECONDS.""" + try: + mtime = sentinel_path.stat().st_mtime + except OSError: + return False + return (time.time() - mtime) < TOKEN_REFRESH_COALESCE_SECONDS + + +def _touch_sentinel(sentinel_path: Path) -> None: + try: + # The sentinel's parent (APP_DIR) is normally created by + # `_refresh_lock` before this runs, but that only happens on the + # `fcntl`-available branch -- on Windows (no `fcntl`) `_refresh_lock` + # yields immediately without ever creating it, which would make + # `.touch()` raise `FileNotFoundError` and silently disable + # coalescing on every such platform. Ensuring the directory here + # keeps best-effort coalescing working even without file locking. + sentinel_path.parent.mkdir(parents=True, exist_ok=True) + sentinel_path.touch() + except OSError as exc: + _debug("refresh_lock", f"failed to touch sentinel {sentinel_path}: {exc}") + + +@contextlib.contextmanager +def _refresh_lock(lock_path: Path): + """Best-effort cross-process exclusive lock guarding a force-refresh. + + Degrades to a no-op (yields ``False`` without locking anything) when + ``fcntl`` is unavailable (Windows) or the lock can't be opened or + acquired for any ``OSError`` — a missing or stuck lock must never + prevent a token refresh from proceeding. Acquisition is a bounded, + non-blocking retry loop (flock has no native timeout) capped at + ``_LOCK_ACQUIRE_TIMEOUT_SECONDS`` before falling back to proceeding + without the lock.""" + if fcntl is None: + yield False + return + + try: + APP_DIR.mkdir(parents=True, exist_ok=True) + handle = lock_path.open("a+") + except OSError as exc: + _debug("refresh_lock", f"could not open lock file {lock_path}: {exc}") + yield False + return + + acquired = False + unexpected_error = False + try: + deadline = time.monotonic() + _LOCK_ACQUIRE_TIMEOUT_SECONDS + while time.monotonic() < deadline: + try: + fcntl.flock(handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + acquired = True + break + except BlockingIOError: + # Expected contention (a peer holds the lock) -- keep + # polling until acquired or the deadline passes. + time.sleep(_LOCK_POLL_INTERVAL_SECONDS) + except OSError as exc: + # Anything else (e.g. EBADF, or a filesystem that doesn't + # support flock at all) is not going to resolve by waiting + # -- fall back to unlocked immediately instead of spinning + # for the full timeout. + unexpected_error = True + _debug( + "refresh_lock", + f"unexpected flock error on {lock_path}, proceeding without lock: {exc}", + ) + break + if not acquired and not unexpected_error: + _debug("refresh_lock", "acquire timeout, proceeding without lock") + _debug("refresh_lock", f"acquired={acquired} path={lock_path}") + yield acquired + finally: + if acquired: + try: + fcntl.flock(handle.fileno(), fcntl.LOCK_UN) + except OSError: + pass + handle.close() + + +def _perform_force_refresh( + workspace: str, + profile: str | None, + cmd: list[str], + fetch: Callable[[], str], +) -> str: + """Serialize `--force-refresh` across processes and coalesce redundant ones. + + Holds a cross-process flock (see `_refresh_lock`) for the duration of an + actual force-refresh so two processes can never redeem the rotating + refresh token concurrently. If a peer already force-refreshed within + `TOKEN_REFRESH_COALESCE_SECONDS`, drops `--force-refresh` from `cmd` (in + place, so retries via the same `fetch` closure stay coalesced too) and + just fetches the now-current cached token instead. + + The sentinel is touched on ATTEMPT, not on success: Databricks rotates + the OAuth refresh token when the `--force-refresh` redemption reaches the + server, not when we successfully parse a token out of our subprocess's + stdout. If the redemption subprocess actually rotated the token + server-side but returned a non-zero exit code or unparseable stdout + (both of which `_fetch` swallows into `""`), failing to touch the + sentinel would let a concurrent peer redeem the same (now-stale) + refresh token again and revoke the whole token family (#190). Marking on + attempt instead means we only risk occasionally serving a slightly-stale + cached token to a coalesced peer — strictly safer than a double + redemption — and any transient failure that gets coalesced away here + will simply be retried after the `TOKEN_REFRESH_COALESCE_SECONDS` window + elapses. Only touch it here, in the branch that actually sent + `--force-refresh`; the coalesced branch above reads the cached token and + performs no redemption, so it must never touch the sentinel.""" + lock_path, sentinel_path = _refresh_lock_paths(workspace, profile) + with _refresh_lock(lock_path) as locked: + _debug( + "get_databricks_token", + f"force-refresh lock acquired={locked} workspace={workspace} profile={profile or ''}", + ) + if _sentinel_is_fresh(sentinel_path): + _debug( + "get_databricks_token", + f"coalescing --force-refresh: peer refreshed within " + f"{TOKEN_REFRESH_COALESCE_SECONDS}s", + ) + if "--force-refresh" in cmd: + cmd.remove("--force-refresh") + return fetch() + + _debug("get_databricks_token", "performing --force-refresh") + token = fetch() + _touch_sentinel(sentinel_path) + return token + + def get_databricks_token( workspace: str, profile: str | None = None, @@ -835,7 +1011,10 @@ def _fetch() -> str: _debug("auth token", f"exception: {type(exc).__name__}: {exc}") return "" - token = _fetch() + if force_refresh: + token = _perform_force_refresh(workspace, profile, cmd, _fetch) + else: + token = _fetch() if not token: # Session may have expired — attempt non-interactive re-auth and retry once. _debug("auth token", "empty on first fetch; attempting auth login --no-browser") diff --git a/tests/test_agent_opencode.py b/tests/test_agent_opencode.py index 0960c64..dc6d405 100644 --- a/tests/test_agent_opencode.py +++ b/tests/test_agent_opencode.py @@ -5,6 +5,8 @@ import json from unittest.mock import patch +import pytest + from ucode.agents import opencode WS = "https://example.databricks.com" @@ -412,3 +414,292 @@ def test_config_written_with_correct_model(self, tmp_path, monkeypatch): written = json.loads(config_file.read_text()) assert written["model"] == "databricks-anthropic/claude-sonnet" + + +ALL_PROVIDER_IDS = ["databricks-anthropic", "databricks-google", "databricks-oss"] + + +class TestManagedProviderIds: + def test_picks_out_provider_keys_only(self): + managed_keys = [ + ["model"], + ["provider", "databricks-anthropic"], + ["provider", "databricks-oss"], + ] + assert opencode._managed_provider_ids(managed_keys) == [ + "databricks-anthropic", + "databricks-oss", + ] + + def test_empty_when_no_provider_keys(self): + assert opencode._managed_provider_ids([["model"]]) == [] + + +class TestRenderAuthPlugin: + def test_contains_chat_headers_hook(self): + js = opencode.render_auth_plugin(WS, None, ALL_PROVIDER_IDS) + assert '"chat.headers"' in js + + def test_embeds_workspace_in_auth_token_argv(self): + js = opencode.render_auth_plugin(WS, None, ALL_PROVIDER_IDS) + assert "auth-token" in js + assert "--host" in js + assert WS in js + + def test_embeds_profile_flag_when_provided(self): + js = opencode.render_auth_plugin(WS, "my-profile", ALL_PROVIDER_IDS) + assert "--profile" in js + assert "my-profile" in js + + def test_omits_profile_flag_when_not_provided(self): + js = opencode.render_auth_plugin(WS, None, ALL_PROVIDER_IDS) + assert "--profile" not in js + + def test_embeds_provider_id_guard_set_with_all_providers(self): + js = opencode.render_auth_plugin(WS, None, ALL_PROVIDER_IDS) + for provider_id in ALL_PROVIDER_IDS: + assert provider_id in js + assert "MANAGED_PROVIDER_IDS" in js + + def test_embeds_only_configured_subset_of_provider_ids(self): + js = opencode.render_auth_plugin(WS, None, ["databricks-anthropic"]) + assert "databricks-anthropic" in js + assert "databricks-google" not in js + assert "databricks-oss" not in js + + def test_embeds_ttl_constant(self): + js = opencode.render_auth_plugin(WS, None, ALL_PROVIDER_IDS) + assert "TTL_MS" in js + assert str(opencode.AUTH_PLUGIN_TOKEN_TTL_SECONDS * 1000) in js + + def test_exports_plugin_function(self): + js = opencode.render_auth_plugin(WS, None, ALL_PROVIDER_IDS) + assert "export const UcodeDatabricksAuth" in js + + def test_provider_scoping_guard_reads_provider_id_directly(self): + # `chat.headers`'s `input.provider` is opencode's runtime Provider + # record (id/name/env/options/source/models) -- NOT a nested + # `{source, info, options}` wrapper. `id` lives directly on + # `input.provider`, never `input.provider.info.id`. This was + # confirmed empirically against a real opencode 1.17.10 hook + # invocation (see #190 follow-up) after a version of this plugin + # that read `input.provider?.info?.id` silently no-opped on every + # request -- the guard's `!providerId` was always true, so the + # hook returned immediately and never refreshed the Authorization + # header. See TestRenderAuthPluginRuntimeBehavior below for a live + # Node execution of the generated JS against the real hook shape, + # which is what actually catches a regression here -- a + # string-content assertion alone would not. + js = opencode.render_auth_plugin(WS, None, ALL_PROVIDER_IDS) + assert "input.provider?.id" in js + assert "input.provider?.info?.id" not in js + + def test_fails_open_on_error(self): + js = opencode.render_auth_plugin(WS, None, ALL_PROVIDER_IDS) + assert "catch" in js + assert "console.error" in js + + def test_uses_execfile_not_shell(self): + js = opencode.render_auth_plugin(WS, None, ALL_PROVIDER_IDS) + assert "execFile" in js + assert "node:child_process" in js + + +class TestRenderAuthPluginRuntimeBehavior: + """Executes the generated plugin under real Node against the actual + opencode `chat.headers` hook input shape, rather than asserting on the + generated JS's string content. A version of this plugin that read + `input.provider?.info?.id` (instead of `input.provider?.id`) passed + every prior test in this file -- because those tests only checked what + string literals appear in the generated source -- while silently + no-opping on every real request (the guard's `!providerId` was always + true). This class exists specifically to catch that class of bug: + a mismatch between our assumed hook-input shape and opencode's real + runtime shape, confirmed via `opencode`'s own `@opencode-ai/plugin` + type surface and an isolated end-to-end run against opencode 1.17.10 + (#190 follow-up). Skips if `node` isn't on PATH.""" + + def _run_plugin( + self, + tmp_path, + provider_id: str | None, + fake_token: str | None = "fresh-token-xyz", + headers_preset: bool = True, + ): + import shutil + import subprocess + import sys + + node = shutil.which("node") + if not node: + pytest.skip("`node` is not installed") + + fake_ucode = tmp_path / "fake_ucode.py" + # `fake_token=None` simulates `ucode auth-token` exiting 0 but + # printing nothing (or only whitespace) -- a successful subprocess + # with unusable output. + print_stmt = f"print({fake_token!r})" if fake_token is not None else 'print(" ")' + # No shebang needed -- invoked directly via `sys.executable` below, + # not executed as a standalone script. + fake_ucode.write_text(f"{print_stmt}\n") + + with patch( + "ucode.agents.opencode.build_auth_token_argv", + # Use the same interpreter running this test rather than a bare + # `python3`, which isn't guaranteed to be on PATH in every + # environment (even one running Python tests). + return_value=[sys.executable, str(fake_ucode)], + ): + js = opencode.render_auth_plugin(WS, None, ["databricks-anthropic"]) + + plugin_path = tmp_path / "plugin.mjs" + plugin_path.write_text(js) + + provider_literal = f"{{ id: {provider_id!r} }}" if provider_id else "undefined" + output_literal = ( + '{ headers: { Authorization: "Bearer STALE-BOOTSTRAP-TOKEN" } }' + if headers_preset + else "{}" + ) + harness = tmp_path / "harness.mjs" + # A `file://` URL is a portable ESM import specifier on every OS; + # `.as_posix()` produces a bare path, which isn't a valid ESM + # specifier on Windows (`C:\\...`). + harness.write_text(f""" +import {{ UcodeDatabricksAuth }} from {plugin_path.as_uri()!r}; + +const hooks = await UcodeDatabricksAuth({{}}); +const output = {output_literal}; +const input = {{ + sessionID: "ses_test", + agent: "build", + model: {{}}, + provider: {provider_literal}, + message: {{}}, +}}; +await hooks["chat.headers"](input, output); +console.log(JSON.stringify(output.headers ?? null)); +""") + + result = subprocess.run( + [node, str(harness)], + capture_output=True, + text=True, + timeout=15, + ) + return result + + def test_managed_provider_id_directly_on_input_provider_refreshes_header(self, tmp_path): + # This is the real opencode runtime shape: `input.provider` IS the + # Provider record, `id` is a direct property. + result = self._run_plugin(tmp_path, provider_id="databricks-anthropic") + assert result.returncode == 0, result.stderr + headers = json.loads(result.stdout.strip().splitlines()[-1]) + assert headers["Authorization"] == "Bearer fresh-token-xyz" + + def test_unmanaged_provider_id_leaves_static_header_untouched(self, tmp_path): + result = self._run_plugin(tmp_path, provider_id="some-other-provider") + assert result.returncode == 0, result.stderr + headers = json.loads(result.stdout.strip().splitlines()[-1]) + assert headers["Authorization"] == "Bearer STALE-BOOTSTRAP-TOKEN" + + def test_missing_provider_leaves_static_header_untouched(self, tmp_path): + result = self._run_plugin(tmp_path, provider_id=None) + assert result.returncode == 0, result.stderr + headers = json.loads(result.stdout.strip().splitlines()[-1]) + assert headers["Authorization"] == "Bearer STALE-BOOTSTRAP-TOKEN" + + def test_empty_token_output_fails_open_instead_of_setting_empty_bearer(self, tmp_path): + # `ucode auth-token` exiting 0 with empty/whitespace-only stdout must + # not be treated as a successful fetch -- otherwise the hook would + # cache and send `Authorization: Bearer ` (empty), clobbering a + # possibly-still-valid bootstrap token for the full TTL window. + result = self._run_plugin(tmp_path, provider_id="databricks-anthropic", fake_token=None) + assert result.returncode == 0, result.stderr + headers = json.loads(result.stdout.strip().splitlines()[-1]) + assert headers["Authorization"] == "Bearer STALE-BOOTSTRAP-TOKEN" + + def test_initializes_missing_output_headers_instead_of_throwing(self, tmp_path): + # If opencode ever invokes the hook with no pre-populated + # `output.headers`, the hook must still set Authorization rather + # than throwing (which the try/catch would otherwise silently + # swallow, no-opping the refresh). + result = self._run_plugin( + tmp_path, provider_id="databricks-anthropic", headers_preset=False + ) + assert result.returncode == 0, result.stderr + headers = json.loads(result.stdout.strip().splitlines()[-1]) + assert headers["Authorization"] == "Bearer fresh-token-xyz" + + +class TestWriteAuthPlugin: + def test_writes_plugin_file_containing_hook(self, tmp_path, monkeypatch): + import ucode.agents.opencode as oc_mod + + plugin_path = tmp_path / "ucode-databricks-auth.js" + monkeypatch.setattr(oc_mod, "OPENCODE_PLUGIN_PATH", plugin_path) + + state = {"workspace": WS, "profile": None} + oc_mod.write_auth_plugin(state, ["databricks-anthropic"]) + + assert plugin_path.exists() + content = plugin_path.read_text() + assert '"chat.headers"' in content + assert "databricks-anthropic" in content + + +class TestWriteToolConfigAuthPlugin: + def _write_and_return(self, tmp_path, monkeypatch, opencode_models): + import ucode.agents.opencode as oc_mod + import ucode.config_io as config_io_mod + + 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 / "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_PLUGIN_PATH", plugin_file) + + state = { + "workspace": WS, + "base_urls": {"opencode": _base_urls()}, + "opencode_models": opencode_models, + "managed_configs": {}, + } + + with ( + patch("ucode.agents.opencode.get_databricks_token", return_value="tok"), + patch("ucode.agents.opencode.save_state"), + ): + oc_mod.write_tool_config(state, "claude-sonnet", token="tok") + + return config_file, plugin_file + + def test_writes_plugin_file_to_configured_path(self, tmp_path, monkeypatch): + _config_file, plugin_file = self._write_and_return( + tmp_path, monkeypatch, {"anthropic": ["claude-sonnet"]} + ) + + assert plugin_file.exists() + assert '"chat.headers"' in plugin_file.read_text() + + def test_plugin_scoped_to_configured_providers_only(self, tmp_path, monkeypatch): + _config_file, plugin_file = self._write_and_return( + tmp_path, monkeypatch, {"anthropic": ["claude-sonnet"]} + ) + + content = plugin_file.read_text() + assert "databricks-anthropic" in content + assert "databricks-google" not in content + assert "databricks-oss" not in content + + def test_static_bootstrap_config_still_present_in_opencode_json(self, tmp_path, monkeypatch): + config_file, _plugin_file = self._write_and_return( + tmp_path, monkeypatch, {"anthropic": ["claude-sonnet"]} + ) + + written = json.loads(config_file.read_text()) + anthropic_options = written["provider"]["databricks-anthropic"]["options"] + assert anthropic_options["apiKey"] == "tok" + assert anthropic_options["headers"]["Authorization"] == "Bearer tok" diff --git a/tests/test_cli.py b/tests/test_cli.py index 3c5d404..ca122f9 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -341,6 +341,48 @@ def test_status_treats_available_tools_as_configured_agents(self): assert "https://example.databricks.com/ai-gateway/gemini" not in result.output +class TestRemoveOpencodeAuthPlugin: + """Review comment on #197: _remove_opencode_auth_plugin previously + swallowed OSError silently, which would leave `ucode revert` reporting + "unchanged" and continuing as if nothing were wrong on a real removal + failure (permissions, filesystem issues) rather than surfacing an + actionable error -- matching restore_file's existing raise convention.""" + + def test_returns_false_when_file_absent(self, tmp_path, monkeypatch): + import ucode.cli as cli_mod + + monkeypatch.setattr(cli_mod, "OPENCODE_PLUGIN_PATH", tmp_path / "missing.js") + + assert cli_mod._remove_opencode_auth_plugin() is False + + def test_returns_true_and_deletes_file_when_present(self, tmp_path, monkeypatch): + import ucode.cli as cli_mod + + plugin_path = tmp_path / "ucode-databricks-auth.js" + plugin_path.write_text("// generated") + monkeypatch.setattr(cli_mod, "OPENCODE_PLUGIN_PATH", plugin_path) + + assert cli_mod._remove_opencode_auth_plugin() is True + assert not plugin_path.exists() + + def test_raises_runtime_error_instead_of_silently_swallowing_removal_failure( + self, tmp_path, monkeypatch + ): + import ucode.cli as cli_mod + + plugin_path = tmp_path / "ucode-databricks-auth.js" + plugin_path.write_text("// generated") + monkeypatch.setattr(cli_mod, "OPENCODE_PLUGIN_PATH", plugin_path) + + def _raise_permission_error(self): + raise PermissionError("denied") + + monkeypatch.setattr(type(plugin_path), "unlink", _raise_permission_error) + + with pytest.raises(RuntimeError, match="Failed to remove OpenCode auth plugin"): + cli_mod._remove_opencode_auth_plugin() + + class TestRevert: def test_reverts_mcp_configs_before_clearing_state(self): state = { diff --git a/tests/test_databricks.py b/tests/test_databricks.py index f59aa1d..1565ed3 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -2,9 +2,11 @@ from __future__ import annotations +import errno import json import os import subprocess +import time import pytest @@ -1163,6 +1165,206 @@ def test_error_suggests_logout_when_matching_profile_exists(self, tmp_path, monk assert f"databricks auth login --host {WS} --profile example-profile" in message +class TestForceRefreshLockAndCoalescing: + """Axis A of #190: serialize `--force-refresh` across processes and + coalesce redundant refreshes within TOKEN_REFRESH_COALESCE_SECONDS.""" + + def _fake_databricks(self, tmp_path, script: str) -> dict: + fake = tmp_path / "databricks" + fake.write_text(f"#!/bin/sh\n{script}\n") + fake.chmod(0o755) + return {**os.environ, "PATH": f"{tmp_path}:{os.environ['PATH']}"} + + def _argv_recording_env(self, tmp_path, argv_log) -> dict: + return self._fake_databricks( + tmp_path, + f'printf "%s\\n" "$@" >> {argv_log}\n' + 'echo \'{"access_token": "good-token", "token_type": "Bearer"}\'', + ) + + def test_coalesces_force_refresh_when_sentinel_is_fresh(self, tmp_path, monkeypatch): + monkeypatch.setattr(db_mod, "APP_DIR", tmp_path) + argv_log = tmp_path / "argv" + env = self._argv_recording_env(tmp_path, argv_log) + monkeypatch.setattr("os.environ", env) + + _, sentinel_path = db_mod._refresh_lock_paths(WS, None) + sentinel_path.touch() # fresh mtime == a peer just force-refreshed + mtime_before = sentinel_path.stat().st_mtime_ns + + token = get_databricks_token(WS, force_refresh=True) + + assert token == "good-token" + argv = argv_log.read_text().splitlines() + assert "--force-refresh" not in argv + # M2: the coalesced branch performs no redemption, so it must never + # re-touch the sentinel. + assert sentinel_path.stat().st_mtime_ns == mtime_before + + def test_performs_force_refresh_and_touches_sentinel_when_stale(self, tmp_path, monkeypatch): + monkeypatch.setattr(db_mod, "APP_DIR", tmp_path) + argv_log = tmp_path / "argv" + env = self._argv_recording_env(tmp_path, argv_log) + monkeypatch.setattr("os.environ", env) + + _, sentinel_path = db_mod._refresh_lock_paths(WS, None) + assert not sentinel_path.exists() + + token = get_databricks_token(WS, force_refresh=True) + + assert token == "good-token" + argv = argv_log.read_text().splitlines() + assert "--force-refresh" in argv + assert sentinel_path.exists() + + def test_force_refresh_attempt_touches_sentinel_even_when_fetch_returns_empty( + self, tmp_path, monkeypatch + ): + """M2: mark-on-attempt, not mark-on-success. The OAuth server rotates + the refresh token when a `--force-refresh` redemption reaches it, not + when we successfully parse a token out of the subprocess's stdout. If + the subprocess exits non-zero or returns unparseable stdout, `_fetch` + swallows that into "" — but the sentinel must still be touched so a + concurrent peer doesn't redeem the same now-rotated refresh token + again. `_perform_force_refresh` must also still return "" so the + caller's outer empty-token re-auth + retry path keeps firing.""" + monkeypatch.setattr(db_mod, "APP_DIR", tmp_path) + _, sentinel_path = db_mod._refresh_lock_paths(WS, None) + assert not sentinel_path.exists() + + cmd = ["databricks", "auth", "token", "--host", WS, "--output", "json", "--force-refresh"] + token = db_mod._perform_force_refresh(WS, None, cmd, fetch=lambda: "") + + assert token == "" + assert sentinel_path.exists() + + def test_performs_force_refresh_when_sentinel_is_stale(self, tmp_path, monkeypatch): + monkeypatch.setattr(db_mod, "APP_DIR", tmp_path) + argv_log = tmp_path / "argv" + env = self._argv_recording_env(tmp_path, argv_log) + monkeypatch.setattr("os.environ", env) + + _, sentinel_path = db_mod._refresh_lock_paths(WS, None) + sentinel_path.touch() + old = time.time() - db_mod.TOKEN_REFRESH_COALESCE_SECONDS - 5 + os.utime(sentinel_path, (old, old)) + + token = get_databricks_token(WS, force_refresh=True) + + assert token == "good-token" + argv = argv_log.read_text().splitlines() + assert "--force-refresh" in argv + + def test_degrades_to_no_lock_when_fcntl_missing(self, tmp_path, monkeypatch): + monkeypatch.setattr(db_mod, "APP_DIR", tmp_path) + monkeypatch.setattr(db_mod, "fcntl", None) + env = self._fake_databricks( + tmp_path, + 'echo \'{"access_token": "good-token", "token_type": "Bearer"}\'', + ) + monkeypatch.setattr("os.environ", env) + + token = get_databricks_token(WS, force_refresh=True) + + assert token == "good-token" + + def test_lock_file_path_differs_for_different_workspace_or_profile(self, tmp_path, monkeypatch): + monkeypatch.setattr(db_mod, "APP_DIR", tmp_path) + + lock_a, sentinel_a = db_mod._refresh_lock_paths(WS, None) + lock_b, sentinel_b = db_mod._refresh_lock_paths("https://other.databricks.com", None) + lock_c, sentinel_c = db_mod._refresh_lock_paths(WS, "some-profile") + + assert lock_a != lock_b + assert lock_a != lock_c + assert lock_b != lock_c + assert sentinel_a != sentinel_b + assert sentinel_a != sentinel_c + assert lock_a.parent == tmp_path + + def test_non_force_refresh_path_ignores_lock_and_sentinel(self, tmp_path, monkeypatch): + # force_refresh=False must stay cheap/unlocked: no lock file should + # even be created. + monkeypatch.setattr(db_mod, "APP_DIR", tmp_path) + env = self._fake_databricks( + tmp_path, + 'echo \'{"access_token": "good-token", "token_type": "Bearer"}\'', + ) + monkeypatch.setattr("os.environ", env) + + token = get_databricks_token(WS) + + assert token == "good-token" + lock_path, sentinel_path = db_mod._refresh_lock_paths(WS, None) + assert not lock_path.exists() + assert not sentinel_path.exists() + + def test_touch_sentinel_creates_missing_parent_directory(self, tmp_path, monkeypatch): + # Review comment on #197: on the fcntl-unavailable branch, + # _refresh_lock returns without ever creating APP_DIR, so + # _touch_sentinel used to raise (swallowed) FileNotFoundError and + # silently disable coalescing on every such platform. It must create + # its own parent directory rather than depending on the caller. + missing_dir = tmp_path / "does-not-exist-yet" + assert not missing_dir.exists() + _, sentinel_path = db_mod._refresh_lock_paths(WS, None) + sentinel_path = missing_dir / sentinel_path.name + + db_mod._touch_sentinel(sentinel_path) + + assert sentinel_path.exists() + + def test_falls_back_immediately_on_unexpected_flock_error(self, tmp_path, monkeypatch): + # Review comment on #197: only lock *contention* (BlockingIOError) + # should be retried in the acquisition loop. Any other OSError + # (e.g. EBADF, or a filesystem that doesn't support flock at all) + # is not going to resolve by waiting, and used to spin for the full + # _LOCK_ACQUIRE_TIMEOUT_SECONDS before falling back. + monkeypatch.setattr(db_mod, "APP_DIR", tmp_path) + monkeypatch.setattr(db_mod, "_LOCK_ACQUIRE_TIMEOUT_SECONDS", 30) + + def raise_unexpected(fd, op): + raise OSError(errno.EBADF, "Bad file descriptor") + + monkeypatch.setattr(db_mod.fcntl, "flock", raise_unexpected) + lock_path, _ = db_mod._refresh_lock_paths(WS, None) + + start = time.monotonic() + with db_mod._refresh_lock(lock_path) as locked: + elapsed = time.monotonic() - start + assert locked is False + + # Bailed out almost immediately, not after spinning for anywhere + # near the 30s acquisition timeout. + assert elapsed < 2.0 + + def test_retries_through_contention_then_acquires(self, tmp_path, monkeypatch): + # BlockingIOError (lock contention) must still be retried, not + # treated as the "give up immediately" unexpected-error case. + monkeypatch.setattr(db_mod, "APP_DIR", tmp_path) + attempts: list[int] = [] + real_flock = db_mod.fcntl.flock + acquire_op = db_mod.fcntl.LOCK_EX | db_mod.fcntl.LOCK_NB + + def flaky_flock(fd, op): + if op != acquire_op: + # Pass unlock calls straight through -- only the acquire + # attempts are what's under test here. + return real_flock(fd, op) + attempts.append(1) + if len(attempts) < 3: + raise BlockingIOError(errno.EAGAIN, "Resource temporarily unavailable") + return real_flock(fd, op) + + monkeypatch.setattr(db_mod, "_LOCK_POLL_INTERVAL_SECONDS", 0.01) + monkeypatch.setattr(db_mod.fcntl, "flock", flaky_flock) + lock_path, _ = db_mod._refresh_lock_paths(WS, None) + + with db_mod._refresh_lock(lock_path) as locked: + assert locked is True + assert len(attempts) == 3 + + class TestListDatabricksConnections: def test_lists_paginated_connections_with_workspace_env(self, monkeypatch): calls: list[dict] = []