Refresh OpenCode Databricks tokens on demand

.github/workflows/ci.yml

@@ -54,7 +54,7 @@ jobs:
           @anthropic-ai/claude-code
           @openai/codex
           @google/gemini-cli
-          opencode-ai
+          opencode-ai@1
           @github/copilot@1.0.80
           @earendil-works/pi-coding-agent
       - run: uv lock

src/ucode/agents/__init__.py

@@ -251,6 +251,9 @@ def tool_binary_installed(tool: str) -> bool:
 def tool_update_available(tool: str) -> tuple[str, str] | None:
     """Return ``(current, latest)`` when a newer agent CLI is published, else None.
     Read-only wrapper over the npm update check — for ``ucode doctor``."""
+    checker = getattr(_MODULES[tool], "is_update_available", None)
+    if callable(checker):
+        return checker()
     return available_npm_package_update(TOOL_SPECS[tool]["package"])
 
 

src/ucode/agents/opencode.py

@@ -2,21 +2,24 @@
 
 from __future__ import annotations
 
+import json
 import os
+import re
 import signal
 import subprocess
-import threading
 
+from ucode.agent_updates import available_npm_package_update
 from ucode.config_io import (
     APP_DIR,
     ToolSpec,
     backup_existing_file,
     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,
@@ -28,10 +31,16 @@
 OPENCODE_CONFIG_DIR = OPENCODE_XDG_CONFIG_HOME / "opencode"
 OPENCODE_CONFIG_PATH = OPENCODE_CONFIG_DIR / "opencode.json"
 OPENCODE_BACKUP_PATH = APP_DIR / "opencode-config.backup.json"
+OPENCODE_AUTH_PLUGIN_PATH = OPENCODE_CONFIG_DIR / "plugin" / "ucode-auth.js"
+OPENCODE_NPM_PACKAGE = "opencode-ai"
+MINIMUM_OPENCODE_VERSION = (1, 0, 220)
+MINIMUM_OPENCODE_VERSION_TEXT = "1.0.220"
 
 SPEC: ToolSpec = {
     "binary": "opencode",
-    "package": "opencode-ai",
+    # npm's `latest` tag currently points at an old 0.0.0 beta. Stay on the
+    # stable v1 line for both fresh installs and updates.
+    "package": f"{OPENCODE_NPM_PACKAGE}@1",
     "display": "OpenCode",
     "config_path": OPENCODE_CONFIG_PATH,
     "backup_path": OPENCODE_BACKUP_PATH,
@@ -43,6 +52,140 @@
     ["provider", "databricks-oss"],
 ]
 
+_AUTH_PLUGIN_TEMPLATE = """// Generated by ucode. Keep Databricks auth fresh for model requests.
+import { execFile } from "node:child_process"
+import { promisify } from "node:util"
+
+const DATABRICKS_PROVIDERS = new Set([
+  "databricks-anthropic",
+  "databricks-google",
+  "databricks-oss",
+])
+const AUTH_COMMAND = __AUTH_COMMAND__
+const REFRESH_SKEW_MS = 120_000
+const run = promisify(execFile)
+
+let accessToken
+let expiresAt = 0
+let refreshPromise
+
+function cacheToken(value) {
+  accessToken = value
+  expiresAt = Infinity
+  try {
+    const claims = JSON.parse(Buffer.from(value.split(".")[1], "base64url").toString())
+    if (typeof claims.exp === "number") expiresAt = claims.exp * 1000
+  } catch {}
+}
+
+async function mintToken() {
+  try {
+    const { stdout } = await run(AUTH_COMMAND[0], AUTH_COMMAND.slice(1), { encoding: "utf8" })
+    const token = stdout.trim()
+    if (!token) throw new Error("returned an empty token")
+    cacheToken(token)
+  } catch (error) {
+    const detail = String(error.stderr || error.message || "").trim()
+    throw new Error("ucode auth-token failed" + (detail ? ": " + detail : ""))
+  }
+}
+
+function refreshToken() {
+  if (!refreshPromise) refreshPromise = mintToken().finally(() => { refreshPromise = undefined })
+  return refreshPromise
+}
+
+function requestWithToken(input, init, token) {
+  const headers = new Headers(input instanceof Request ? input.headers : undefined)
+  new Headers(init?.headers).forEach((value, key) => headers.set(key, value))
+  headers.set("Authorization", "Bearer " + token)
+  return { ...init, headers }
+}
+
+async function databricksFetch(input, init) {
+  if (!accessToken || expiresAt <= Date.now() + REFRESH_SKEW_MS) await refreshToken()
+  const attemptedToken = accessToken
+  const response = await fetch(input, requestWithToken(input, init, attemptedToken))
+  if (response.status !== 401) return response
+
+  // Another concurrent request may already have refreshed this token.
+  if (accessToken === attemptedToken) await refreshToken()
+  return fetch(input, requestWithToken(input, init, accessToken))
+}
+
+export const UcodeDatabricksAuth = async () => ({
+  config: async (config) => {
+    for (const providerID of DATABRICKS_PROVIDERS) {
+      const options = config.provider?.[providerID]?.options
+      if (!options) continue
+      if (!accessToken && typeof options.apiKey === "string") cacheToken(options.apiKey)
+      options.fetch = databricksFetch
+    }
+  },
+})
+"""
+
+
+def _parse_version(value: str) -> tuple[int, int, int] | None:
+    match = re.search(r"(\d+)\.(\d+)\.(\d+)", value)
+    if match is None:
+        return None
+    major, minor, patch = match.groups()
+    return int(major), int(minor), int(patch)
+
+
+def _minimum_version_message() -> str | None:
+    installed = agent_version(SPEC["binary"])
+    parsed = _parse_version(installed)
+    if parsed is None or parsed >= MINIMUM_OPENCODE_VERSION:
+        return None
+    return (
+        f"OpenCode {installed} is too old. ucode requires OpenCode "
+        f"{MINIMUM_OPENCODE_VERSION_TEXT} or newer for renewable Databricks authentication."
+    )
+
+
+def required_update_message() -> str | None:
+    return _minimum_version_message()
+
+
+def minimum_version_error() -> str | None:
+    message = _minimum_version_message()
+    if message is None:
+        return None
+    return f"{message} Update it with `npm install -g {SPEC['package']}`."
+
+
+def is_update_available() -> tuple[str, str] | None:
+    """Offer only stable OpenCode v1 updates, never npm's beta `latest` tag."""
+    update = available_npm_package_update(OPENCODE_NPM_PACKAGE)
+    if update is None:
+        return None
+    _, target = update
+    parsed = _parse_version(target)
+    if "-" in target or parsed is None or parsed[0] != 1:
+        return None
+    return update
+
+
+def render_auth_plugin(state: dict) -> str:
+    """Render the local OpenCode plugin that refreshes Databricks auth on demand."""
+    argv = build_auth_token_argv(
+        state["workspace"],
+        state.get("profile"),
+        use_pat=bool(state.get("use_pat")),
+    )
+    # A 401 must not return the same still-unexpired cached credential.
+    argv.append("--force-refresh")
+    return _AUTH_PLUGIN_TEMPLATE.replace("__AUTH_COMMAND__", json.dumps(argv))
+
+
+def write_auth_plugin(state: dict) -> None:
+    """Install the auto-discovered hook in ucode's isolated OpenCode config."""
+    # Derive the path so tests that redirect OPENCODE_CONFIG_PATH stay isolated.
+    path = OPENCODE_CONFIG_PATH.parent / "plugin" / OPENCODE_AUTH_PLUGIN_PATH.name
+    write_text_file(path, render_auth_plugin(state))
+
 
 def _resolve_model_selector(model: str, opencode_models: dict[str, list[str]]) -> str:
     """Return an OpenCode model selector in provider/model form when possible."""
@@ -153,14 +296,10 @@ def write_tool_config(
     state: dict,
     model: str,
     token: str | None = None,
-    *,
-    force_refresh: bool = False,
 ) -> tuple[dict, str]:
     backup_existing_file(OPENCODE_CONFIG_PATH, OPENCODE_BACKUP_PATH)
     if token is None:
-        token = get_databricks_token(
-            state["workspace"], state.get("profile"), force_refresh=force_refresh
-        )
+        token = get_databricks_token(state["workspace"], state.get("profile"))
     opencode_base_urls = state.get("base_urls", {}).get("opencode") or build_opencode_base_urls(
         state["workspace"]
     )
@@ -171,6 +310,7 @@ def write_tool_config(
         state.get("opencode_models") or {},
     )
     existing = read_json_safe(OPENCODE_CONFIG_PATH)
+    write_auth_plugin(state)
     providers = existing.get("provider")
     if isinstance(providers, dict):
         for stale in (
@@ -236,22 +376,14 @@ def default_model(state: dict) -> str | None:
     return oss[0] if oss else None
 
 
-def _refresh_token_once(state: dict, *, force_refresh: bool = False) -> str:
+def _configure_launch(state: dict) -> str:
     model = default_model(state)
     if not model:
         raise RuntimeError("No OpenCode model is configured.")
-    _, token = write_tool_config(state, model, force_refresh=force_refresh)
+    _, token = write_tool_config(state, model)
     return token
 
 
-def _refresh_forever(state: dict, stop_event: threading.Event) -> None:
-    while not stop_event.wait(TOKEN_REFRESH_INTERVAL_SECONDS):
-        try:
-            _refresh_token_once(state, force_refresh=True)
-        except RuntimeError:
-            continue
-
-
 def build_runtime_env(token: str, state: dict | None = None) -> dict[str, str]:
     env = os.environ.copy()
     env["OAUTH_TOKEN"] = token
@@ -260,27 +392,16 @@ def build_runtime_env(token: str, state: dict | None = None) -> dict[str, str]:
 
 
 def launch(state: dict, tool_args: list[str]) -> None:
-    """Launch opencode with background token refresh (same pattern as Gemini)."""
-    token = _refresh_token_once(state)
+    """Launch OpenCode with on-demand token refresh from its local plugin."""
+    token = _configure_launch(state)
     env = build_runtime_env(token, state)
 
-    stop_event = threading.Event()
-    refresher = threading.Thread(
-        target=_refresh_forever,
-        args=(state, stop_event),
-        daemon=True,
-    )
-    refresher.start()
-
     proc = subprocess.Popen([SPEC["binary"], *tool_args], env=env)
     try:
         returncode = proc.wait()
     except KeyboardInterrupt:
         proc.send_signal(signal.SIGINT)
         returncode = proc.wait()
-    finally:
-        stop_event.set()
-        refresher.join(timeout=1)
 
     raise SystemExit(returncode)
 

src/ucode/cli.py

@@ -1368,6 +1368,10 @@ def auth_token_cmd(
     use_pat: Annotated[
         bool, typer.Option("--use-pat", help="Read the profile's static PAT instead of OAuth.")
     ] = False,
+    force_refresh: Annotated[
+        bool,
+        typer.Option("--force-refresh", help="Force the Databricks CLI to mint a new token."),
+    ] = False,
 ) -> None:
     """Print a Databricks bearer token to stdout, then exit.
 
@@ -1398,7 +1402,7 @@ def auth_token_cmd(
             )
             raise typer.Exit(1)
     try:
-        token = get_databricks_token(workspace, profile)
+        token = get_databricks_token(workspace, profile, force_refresh=force_refresh)
     except RuntimeError as exc:
         print_err(str(exc))
         raise typer.Exit(1) from None