From f7c597fe58965b570c80139da07f62d8652292e5 Mon Sep 17 00:00:00 2001 From: Jesse Escobedo Date: Wed, 8 Jul 2026 11:19:01 -0400 Subject: [PATCH 1/7] fix(databricks): serialize token --force-refresh with cross-process flock (#190) Concurrent long-lived ucode sessions each ran `databricks auth token --force-refresh` independently, racing on the shared OAuth cache and redeeming (and revoking) the rotating refresh token concurrently. Serialize --force-refresh calls per workspace+profile behind a cross-process flock under APP_DIR, and coalesce redundant refreshes: if a peer force-refreshed within TOKEN_REFRESH_COALESCE_SECONDS (60s), drop --force-refresh and reuse the now-current cached token instead of redeeming the refresh token again. Degrades gracefully to no locking when fcntl is unavailable (Windows) or lock acquisition fails for any reason -- a stuck/missing lock must never block a token refresh. The non-force-refresh (cheap read) path is unchanged and unlocked. --- src/ucode/databricks.py | 136 ++++++++++++++++++++++++++++++++++++++- tests/test_databricks.py | 111 ++++++++++++++++++++++++++++++++ 2 files changed, 246 insertions(+), 1 deletion(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 9f89c6b..613d483 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,13 @@ 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 +_LOCK_ACQUIRE_TIMEOUT_SECONDS = 15 +_LOCK_POLL_INTERVAL_SECONDS = 0.2 def _debug_enabled() -> bool: @@ -776,6 +797,116 @@ 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: + 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 + 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 OSError: + time.sleep(_LOCK_POLL_INTERVAL_SECONDS) + _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.""" + 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() + if token: + _touch_sentinel(sentinel_path) + return token + + def get_databricks_token( workspace: str, profile: str | None = None, @@ -835,7 +966,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_databricks.py b/tests/test_databricks.py index f59aa1d..c7a5ac5 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -5,6 +5,7 @@ import json import os import subprocess +import time import pytest @@ -1163,6 +1164,116 @@ 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 + + token = get_databricks_token(WS, force_refresh=True) + + assert token == "good-token" + argv = argv_log.read_text().splitlines() + assert "--force-refresh" not in argv + + 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_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() + + class TestListDatabricksConnections: def test_lists_paginated_connections_with_workspace_env(self, monkeypatch): calls: list[dict] = [] From 1c183929766242be4256cc2a057699b2f9e775e7 Mon Sep 17 00:00:00 2001 From: Jesse Escobedo Date: Wed, 8 Jul 2026 11:30:47 -0400 Subject: [PATCH 2/7] fix(opencode): inject fresh Authorization per-request via chat.headers plugin (#190) opencode reads provider.options.headers.Authorization from opencode.json once at startup and never again, so a long-lived session keeps sending a bearer token that may since have been rotated or revoked. Ship a ucode-managed opencode plugin (chat.headers hook, auto-loaded from the opencode plugin dir) that sets a fresh Authorization header on every chat request for our Databricks providers, sourcing the token from `ucode auth-token` but cached in-process with a short TTL (60s) plus an in-flight promise guard so it doesn't shell out on every request or spawn concurrent subprocesses. The static apiKey/headers in opencode.json are left untouched as the bootstrap/fallback the plugin fails open to on any refresh error. write_tool_config now also writes the plugin, scoped to only the provider ids actually configured. `ucode revert` deletes the (entirely ucode-generated, no-backup-needed) plugin file. --- src/ucode/agents/opencode.py | 121 ++++++++++++++++++++++++++++ src/ucode/cli.py | 18 +++++ tests/test_agent_opencode.py | 148 +++++++++++++++++++++++++++++++++++ 3 files changed, 287 insertions(+) diff --git a/src/ucode/agents/opencode.py b/src/ucode/agents/opencode.py index 4b614f3..839f443 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,109 @@ 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); + return stdout.trim(); +}} + +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) => {{ + const providerId = input.provider?.info?.id; + if (!providerId || !MANAGED_PROVIDER_IDS.has(providerId)) {{ + return; + }} + try {{ + const token = await getToken(); + 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 +306,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..6b36c6d 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,21 @@ 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).""" + try: + if OPENCODE_PLUGIN_PATH.exists(): + OPENCODE_PLUGIN_PATH.unlink() + return True + except OSError: + pass + return False + + def revert() -> int: state = load_state() managed_configs = state.get("managed_configs") or {} @@ -690,6 +706,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 +719,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/tests/test_agent_opencode.py b/tests/test_agent_opencode.py index 0960c64..2dcfcc9 100644 --- a/tests/test_agent_opencode.py +++ b/tests/test_agent_opencode.py @@ -412,3 +412,151 @@ 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_checks_provider_info_id(self): + js = opencode.render_auth_plugin(WS, None, ALL_PROVIDER_IDS) + assert "input.provider?.info?.id" 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 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" From 98478ab122c934d9fee0bd9c7053eb220c3cbc4c Mon Sep 17 00:00:00 2001 From: Jesse Escobedo Date: Wed, 8 Jul 2026 12:54:04 -0400 Subject: [PATCH 3/7] fix(databricks): mark refresh sentinel on attempt + widen lock timeout (review #190) --- src/ucode/databricks.py | 31 +++++++++++++++++++++++++++---- tests/test_databricks.py | 25 +++++++++++++++++++++++++ 2 files changed, 52 insertions(+), 4 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 613d483..e7307a4 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -70,7 +70,13 @@ # 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 -_LOCK_ACQUIRE_TIMEOUT_SECONDS = 15 +# 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 @@ -859,6 +865,8 @@ def _refresh_lock(lock_path: Path): break except OSError: time.sleep(_LOCK_POLL_INTERVAL_SECONDS) + if not acquired: + _debug("refresh_lock", "acquire timeout, proceeding without lock") _debug("refresh_lock", f"acquired={acquired} path={lock_path}") yield acquired finally: @@ -883,7 +891,23 @@ def _perform_force_refresh( 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.""" + 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( @@ -902,8 +926,7 @@ def _perform_force_refresh( _debug("get_databricks_token", "performing --force-refresh") token = fetch() - if token: - _touch_sentinel(sentinel_path) + _touch_sentinel(sentinel_path) return token diff --git a/tests/test_databricks.py b/tests/test_databricks.py index c7a5ac5..dc02581 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -1189,12 +1189,16 @@ def test_coalesces_force_refresh_when_sentinel_is_fresh(self, tmp_path, monkeypa _, 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) @@ -1212,6 +1216,27 @@ def test_performs_force_refresh_and_touches_sentinel_when_stale(self, tmp_path, 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" From 7fce61867ee2daf3c93548443723a4b58311eec5 Mon Sep 17 00:00:00 2001 From: Jesse Escobedo Date: Wed, 8 Jul 2026 15:35:58 -0400 Subject: [PATCH 4/7] fix(opencode): read provider id directly off input.provider (#190) The generated chat.headers plugin guarded provider scoping with input.provider?.info?.id, but opencode's real hook input.provider is the runtime Provider record itself (id/name/env/options/source/models) -- id lives directly on input.provider, never under a nested .info. The guard's !providerId check was therefore always true, so the hook returned immediately on every request and never refreshed Authorization for any provider. Axis B was a complete no-op in production despite passing all prior tests, because those tests only asserted on the generated JS's string content, never its behavior against a real hook invocation. Found via an isolated end-to-end test against the real opencode 1.17.10 binary (separate $HOME/XDG dirs, real Databricks auth, no contact with the live ucode install): a well-formed-but-expired JWT baked into the static provider.options.headers.Authorization/apiKey still produced 'Bad Request: Invalid Token' with the plugin loaded and firing -- diagnostic instrumentation showed the guard's providerId was always undefined. Fixed to input.provider?.id and reverified against the same isolated harness: the plugin now fetches a fresh token and the request succeeds despite the stale static config. Adds TestRenderAuthPluginRuntimeBehavior, which shells out to a real node process and invokes the generated plugin's chat.headers hook against the actual runtime input shape (not a mocked one) -- verified to fail against the old input.provider?.info?.id line and pass against the fix, so this class of bug is now covered by an assertion that actually exercises behavior rather than string content. --- src/ucode/agents/opencode.py | 8 ++- tests/test_agent_opencode.py | 105 ++++++++++++++++++++++++++++++++++- 2 files changed, 110 insertions(+), 3 deletions(-) diff --git a/src/ucode/agents/opencode.py b/src/ucode/agents/opencode.py index 839f443..da454dd 100644 --- a/src/ucode/agents/opencode.py +++ b/src/ucode/agents/opencode.py @@ -247,7 +247,13 @@ def render_auth_plugin(workspace: str, profile: str | None, provider_ids: list[s export const UcodeDatabricksAuth = async (_ctx) => ({{ "chat.headers": async (input, output) => {{ - const providerId = input.provider?.info?.id; + // 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; }} diff --git a/tests/test_agent_opencode.py b/tests/test_agent_opencode.py index 2dcfcc9..f376ed4 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" @@ -474,9 +476,23 @@ 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_checks_provider_info_id(self): + 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?.info?.id" in js + 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) @@ -489,6 +505,91 @@ def test_uses_execfile_not_shell(self): 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 = "fresh-token-xyz"): + import shutil + import subprocess + + node = shutil.which("node") + if not node: + pytest.skip("`node` is not installed") + + fake_ucode = tmp_path / "fake_ucode.py" + fake_ucode.write_text( + "#!/usr/bin/env python3\nimport sys\nprint(sys.argv[-1] if False else " + f"{fake_token!r})\n" + ) + fake_ucode.chmod(0o755) + + with patch( + "ucode.agents.opencode.build_auth_token_argv", + return_value=["python3", 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" + harness = tmp_path / "harness.mjs" + harness.write_text(f""" +import {{ UcodeDatabricksAuth }} from {str(plugin_path.as_posix())!r}; + +const hooks = await UcodeDatabricksAuth({{}}); +const output = {{ headers: {{ Authorization: "Bearer STALE-BOOTSTRAP-TOKEN" }} }}; +const input = {{ + sessionID: "ses_test", + agent: "build", + model: {{}}, + provider: {provider_literal}, + message: {{}}, +}}; +await hooks["chat.headers"](input, output); +console.log(JSON.stringify(output.headers)); +""") + + 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" + + class TestWriteAuthPlugin: def test_writes_plugin_file_containing_hook(self, tmp_path, monkeypatch): import ucode.agents.opencode as oc_mod From d25993485ff5b8e4fe2b41567d9cd09b77ff29e3 Mon Sep 17 00:00:00 2001 From: Jesse Escobedo Date: Thu, 9 Jul 2026 09:58:03 -0400 Subject: [PATCH 5/7] fix(opencode): fail open on empty token output, guard missing output.headers (review) Addresses two Copilot review comments on #197: - fetchToken() returned stdout.trim() unconditionally. If `ucode auth-token` ever exits 0 with empty/whitespace-only stdout, the plugin would cache and send Authorization: Bearer , clobbering a possibly-still-valid bootstrap token for the full TTL window instead of failing open. Empty output now throws, so the hook's catch block fails open as intended. - The chat.headers hook assumed output.headers always exists. If opencode ever invokes the hook without a pre-populated headers object, the assignment would throw -- caught silently by the existing try/catch, but the refresh would never apply. Defensively initialize output.headers before writing to it. Both are proven by new Node-executed runtime tests (matching the class of test that caught the input.provider?.id bug): each new test was verified to fail against the pre-fix code and pass against the fix, not just pass trivially. --- src/ucode/agents/opencode.py | 14 ++++++++++- tests/test_agent_opencode.py | 48 ++++++++++++++++++++++++++++++------ 2 files changed, 54 insertions(+), 8 deletions(-) diff --git a/src/ucode/agents/opencode.py b/src/ucode/agents/opencode.py index da454dd..dadf404 100644 --- a/src/ucode/agents/opencode.py +++ b/src/ucode/agents/opencode.py @@ -222,7 +222,14 @@ def render_auth_plugin(workspace: str, profile: str | None, provider_ids: list[s async function fetchToken() {{ const [cmd, ...args] = UCODE_AUTH_TOKEN_ARGV; const {{ stdout }} = await execFileAsync(cmd, args); - return stdout.trim(); + 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() {{ @@ -259,6 +266,11 @@ def render_auth_plugin(workspace: str, profile: str | None, provider_ids: list[s }} 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); diff --git a/tests/test_agent_opencode.py b/tests/test_agent_opencode.py index f376ed4..7a2528a 100644 --- a/tests/test_agent_opencode.py +++ b/tests/test_agent_opencode.py @@ -519,7 +519,13 @@ class TestRenderAuthPluginRuntimeBehavior: 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 = "fresh-token-xyz"): + 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 @@ -528,10 +534,11 @@ def _run_plugin(self, tmp_path, provider_id: str | None, fake_token: str = "fres pytest.skip("`node` is not installed") fake_ucode = tmp_path / "fake_ucode.py" - fake_ucode.write_text( - "#!/usr/bin/env python3\nimport sys\nprint(sys.argv[-1] if False else " - f"{fake_token!r})\n" - ) + # `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(" ")' + fake_ucode.write_text(f"#!/usr/bin/env python3\n{print_stmt}\n") fake_ucode.chmod(0o755) with patch( @@ -544,12 +551,17 @@ def _run_plugin(self, tmp_path, provider_id: str | None, fake_token: str = "fres 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" harness.write_text(f""" import {{ UcodeDatabricksAuth }} from {str(plugin_path.as_posix())!r}; const hooks = await UcodeDatabricksAuth({{}}); -const output = {{ headers: {{ Authorization: "Bearer STALE-BOOTSTRAP-TOKEN" }} }}; +const output = {output_literal}; const input = {{ sessionID: "ses_test", agent: "build", @@ -558,7 +570,7 @@ def _run_plugin(self, tmp_path, provider_id: str | None, fake_token: str = "fres message: {{}}, }}; await hooks["chat.headers"](input, output); -console.log(JSON.stringify(output.headers)); +console.log(JSON.stringify(output.headers ?? null)); """) result = subprocess.run( @@ -589,6 +601,28 @@ def test_missing_provider_leaves_static_header_untouched(self, tmp_path): 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): From 7e0084b4a53ecae063111f40e38f9f11439b28c8 Mon Sep 17 00:00:00 2001 From: Jesse Escobedo Date: Thu, 9 Jul 2026 10:08:50 -0400 Subject: [PATCH 6/7] test(opencode): use sys.executable and Path.as_uri() in runtime harness (review) Addresses two more Copilot review comments on #197 (both in the TestRenderAuthPluginRuntimeBehavior test harness, not production code): - _run_plugin() hard-coded `python3` when patching build_auth_token_argv for the fake ucode binary. python3 isn't guaranteed to be on PATH in every environment. Use sys.executable so the fake script runs under the same interpreter as the test. - The Node harness imported the generated plugin via plugin_path.as_posix(), a bare filesystem path. That's not a valid ESM import specifier on Windows (C:\...). Use Path.as_uri() for a portable file:// URL. The other two comments on this PR (opencode.py lines 233/275) are stale duplicates of the two already fixed in d259934 -- same comment IDs, diff_hunk frozen at review-creation time before that fix landed. No action needed there beyond the earlier reply. --- tests/test_agent_opencode.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/tests/test_agent_opencode.py b/tests/test_agent_opencode.py index 7a2528a..dc6d405 100644 --- a/tests/test_agent_opencode.py +++ b/tests/test_agent_opencode.py @@ -528,6 +528,7 @@ def _run_plugin( ): import shutil import subprocess + import sys node = shutil.which("node") if not node: @@ -538,12 +539,16 @@ def _run_plugin( # 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(" ")' - fake_ucode.write_text(f"#!/usr/bin/env python3\n{print_stmt}\n") - fake_ucode.chmod(0o755) + # 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", - return_value=["python3", str(fake_ucode)], + # 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"]) @@ -557,8 +562,11 @@ def _run_plugin( 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 {str(plugin_path.as_posix())!r}; +import {{ UcodeDatabricksAuth }} from {plugin_path.as_uri()!r}; const hooks = await UcodeDatabricksAuth({{}}); const output = {output_literal}; From d8143905bf51e2bcb3088e6d7075637f3b5bbf4a Mon Sep 17 00:00:00 2001 From: Jesse Escobedo Date: Thu, 9 Jul 2026 10:25:55 -0400 Subject: [PATCH 7/7] fix(databricks,cli): harden lock/sentinel edge cases, don't swallow revert errors (review) Addresses three more Copilot review comments on #197: - _touch_sentinel() could raise (swallowed) FileNotFoundError when its parent directory doesn't exist yet -- true whenever fcntl is unavailable (Windows), since _refresh_lock returns before ever creating APP_DIR on that branch. Silently disabled coalescing on every such platform. Now creates its own parent directory. - _refresh_lock()'s acquisition loop retried on *any* OSError from flock(), not just contention. An unexpected error (EBADF, or a filesystem that doesn't support flock at all) would spin for the full _LOCK_ACQUIRE_TIMEOUT_SECONDS (30s) before falling back to unlocked, instead of failing back immediately. Now only retries on BlockingIOError (the actual contention case) and bails out right away on anything else. - _remove_opencode_auth_plugin() silently swallowed OSError, so a real removal failure (permissions, filesystem issues) would leave `ucode revert` reporting "unchanged" and proceeding as if nothing were wrong. Now raises RuntimeError, matching restore_file's existing convention elsewhere in revert. All three are covered by new tests, each verified to fail against the pre-fix code (including the 30s-spin one, confirmed to actually take ~30s before the fix and <2s after) and pass against the fix. --- src/ucode/cli.py | 16 +++++++--- src/ucode/databricks.py | 26 ++++++++++++++-- tests/test_cli.py | 42 +++++++++++++++++++++++++ tests/test_databricks.py | 66 ++++++++++++++++++++++++++++++++++++++++ 4 files changed, 144 insertions(+), 6 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 6b36c6d..64e7627 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -682,14 +682,22 @@ def _remove_opencode_auth_plugin() -> bool: 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).""" + 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 - except OSError: - pass - return False + return False + except OSError as exc: + raise RuntimeError( + f"Failed to remove OpenCode auth plugin at {OPENCODE_PLUGIN_PATH}" + ) from exc def revert() -> int: diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index e7307a4..e867f54 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -827,6 +827,14 @@ def _sentinel_is_fresh(sentinel_path: Path) -> bool: 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}") @@ -856,6 +864,7 @@ def _refresh_lock(lock_path: Path): return acquired = False + unexpected_error = False try: deadline = time.monotonic() + _LOCK_ACQUIRE_TIMEOUT_SECONDS while time.monotonic() < deadline: @@ -863,9 +872,22 @@ def _refresh_lock(lock_path: Path): fcntl.flock(handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) acquired = True break - except OSError: + except BlockingIOError: + # Expected contention (a peer holds the lock) -- keep + # polling until acquired or the deadline passes. time.sleep(_LOCK_POLL_INTERVAL_SECONDS) - if not acquired: + 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 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 dc02581..1565ed3 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -2,6 +2,7 @@ from __future__ import annotations +import errno import json import os import subprocess @@ -1298,6 +1299,71 @@ def test_non_force_refresh_path_ignores_lock_and_sentinel(self, tmp_path, monkey 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):