opencode: refresh Databricks token via loopback proxy (fixes ~1h mid-session 401s)

src/ucode/agents/opencode.py

@@ -6,6 +6,7 @@
 import signal
 import subprocess
 import threading
+from urllib.parse import urlsplit, urlunsplit
 
 from ucode.agent_updates import available_npm_package_update
 from ucode.config_io import (
@@ -17,7 +18,6 @@
     write_json_file,
 )
 from ucode.databricks import (
-    TOKEN_REFRESH_INTERVAL_SECONDS,
     build_opencode_base_urls,
     get_databricks_token,
     model_token_limits,
@@ -160,14 +160,19 @@ def write_tool_config(
     token: str | None = None,
     *,
     force_refresh: bool = False,
+    base_urls: dict[str, str] | None = None,
 ) -> 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
         )
-    opencode_base_urls = state.get("base_urls", {}).get("opencode") or build_opencode_base_urls(
-        state["workspace"]
+    # ``base_urls`` lets launch() point providers at the loopback refresh proxy
+    # without persisting per-session 127.0.0.1:<port> URLs into workspace state.
+    opencode_base_urls = (
+        base_urls
+        or state.get("base_urls", {}).get("opencode")
+        or build_opencode_base_urls(state["workspace"])
     )
     overlay, managed_keys = render_overlay(
         model,
@@ -241,20 +246,17 @@ def default_model(state: dict) -> str | None:
     return oss[0] if oss else None
 
 
-def _refresh_token_once(state: dict, *, force_refresh: bool = False) -> str:
-    model = default_model(state)
-    if not model:
-        raise RuntimeError("No OpenCode model is configured.")
-    _, token = write_tool_config(state, model, force_refresh=force_refresh)
-    return token
-
-
-def _refresh_forever(state: dict, stop_event: threading.Event) -> None:
-    while not stop_event.wait(TOKEN_REFRESH_INTERVAL_SECONDS):
-        try:
-            _refresh_token_once(state, force_refresh=True)
-        except RuntimeError:
-            continue
+def _loopback_base_urls(real_base_urls: dict[str, str], port: int) -> dict[str, str]:
+    """Rewrite each provider base URL's scheme+host to the loopback proxy, keeping
+    the gateway path so the proxy forwards it verbatim to the workspace host. All
+    OpenCode provider paths share the workspace host, so one proxy serves them all.
+    """
+    loopback = f"127.0.0.1:{port}"
+    out: dict[str, str] = {}
+    for provider, url in real_base_urls.items():
+        parts = urlsplit(url)
+        out[provider] = urlunsplit(("http", loopback, parts.path, parts.query, parts.fragment))
+    return out
 
 
 def build_runtime_env(token: str, state: dict | None = None) -> dict[str, str]:
@@ -265,17 +267,37 @@ 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)
-    env = build_runtime_env(token, state)
+    """Launch opencode behind the loopback refresh proxy.
 
-    stop_event = threading.Event()
-    refresher = threading.Thread(
-        target=_refresh_forever,
-        args=(state, stop_event),
-        daemon=True,
+    OpenCode resolves opencode.json once at process start and never re-reads it, so
+    a token baked into the config goes stale at the ~1h OAuth expiry and the session
+    dies mid-run. Point every provider baseURL at the loopback proxy, which rewrites
+    the Authorization header with a freshly-minted token on every request (the proxy
+    owns its own background refresher). Mirrors claude.py::_launch_relayed.
+    """
+    from ucode.opencode_proxy import start_proxy
+
+    model = default_model(state)
+    if not model:
+        raise RuntimeError("No OpenCode model is configured.")
+
+    workspace = state["workspace"]
+    real_base_urls = state.get("base_urls", {}).get("opencode") or build_opencode_base_urls(
+        workspace
     )
-    refresher.start()
+
+    # One proxy for all providers (they share the workspace host); it overwrites
+    # Authorization with a freshly-minted token on every request.
+    server, cache = start_proxy(workspace, state.get("profile"))
+    port = server.server_address[1]
+
+    # opencode.json points at the proxy; the token in the file is only a bootstrap
+    # the proxy replaces per request, so it never goes stale mid-session.
+    _, token = write_tool_config(state, model, base_urls=_loopback_base_urls(real_base_urls, port))
+    env = build_runtime_env(token, state)
+
+    server_thread = threading.Thread(target=server.serve_forever, daemon=True)
+    server_thread.start()
 
     proc = subprocess.Popen([SPEC["binary"], *tool_args], env=env)
     try:
@@ -284,8 +306,8 @@ def launch(state: dict, tool_args: list[str]) -> None:
         proc.send_signal(signal.SIGINT)
         returncode = proc.wait()
     finally:
-        stop_event.set()
-        refresher.join(timeout=1)
+        cache.stop()
+        server.shutdown()
 
     raise SystemExit(returncode)
 

src/ucode/opencode_proxy.py

@@ -0,0 +1,185 @@
+"""Loopback refresh proxy for OpenCode's Databricks AI Gateway providers.
+
+OpenCode resolves ``opencode.json`` once at process start and never re-reads it,
+so a Databricks OAuth token baked into the config (the provider ``Authorization``
+header) goes stale at the ~1h token lifetime and every request then fails with
+``401 Invalid Token``. ``ucode opencode`` therefore points each provider's
+``baseURL`` at this loopback proxy instead: it forwards every request to the
+workspace gateway unchanged except for overwriting ``Authorization`` with a
+freshly-minted Databricks token, and streams the response back verbatim.
+
+Unlike the relayed-Anthropic proxy (``gateway_proxy``), OpenCode has no
+subscription-OAuth credential to preserve: the Databricks token *is* the
+``Authorization`` credential, so that is the header this proxy swaps. All
+OpenCode provider paths (anthropic ``/v1``, gemini ``/v1beta``, mlflow ``/v1``)
+share the workspace host, so one proxy serves them all — it swaps host +
+``Authorization`` and forwards the client's path unchanged.
+
+Security invariants (mirroring ``databricks.py`` token handling):
+  - Binds 127.0.0.1 only; never exposed off-host.
+  - Never logs header values or bodies. The Databricks token lives in memory,
+    refreshed off the request path.
+"""
+
+from __future__ import annotations
+
+import threading
+from email.message import Message
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from typing import IO
+from urllib import error as urllib_error
+from urllib import request as urllib_request
+from urllib.parse import urljoin
+
+from ucode.databricks import TOKEN_REFRESH_INTERVAL_SECONDS, get_databricks_token
+
+# Header carrying OpenCode's credential; overwritten with the minted token so a
+# stale value baked into opencode.json can never reach the gateway.
+_AUTH_HEADER = "Authorization"
+# Hop-by-hop headers must not be forwarded across the proxy.
+_HOP_BY_HOP = frozenset(
+    h.lower()
+    for h in (
+        "connection",
+        "keep-alive",
+        "proxy-authenticate",
+        "proxy-authorization",
+        "te",
+        "trailers",
+        "transfer-encoding",
+        "upgrade",
+        "host",
+        "content-length",
+    )
+)
+_STREAM_CHUNK = 8192
+
+
+class _TokenCache:
+    """Holds the current Databricks token, refreshed by a background thread so
+    minting never blocks a request."""
+
+    def __init__(self, workspace: str, profile: str | None) -> None:
+        self._workspace = workspace
+        self._profile = profile
+        self._lock = threading.Lock()
+        self._token = get_databricks_token(workspace, profile)
+        self._stop = threading.Event()
+
+    @property
+    def token(self) -> str:
+        with self._lock:
+            return self._token
+
+    def refresh(self) -> None:
+        token = get_databricks_token(self._workspace, self._profile, force_refresh=True)
+        with self._lock:
+            self._token = token
+
+    def run_refresher(self) -> None:
+        while not self._stop.wait(TOKEN_REFRESH_INTERVAL_SECONDS):
+            try:
+                self.refresh()
+            except RuntimeError:
+                continue
+
+    def stop(self) -> None:
+        self._stop.set()
+
+
+def _forwarded_request_headers(handler: BaseHTTPRequestHandler, token: str) -> dict[str, str]:
+    headers = {
+        key: value
+        for key, value in handler.headers.items()
+        if key.lower() not in _HOP_BY_HOP and key.lower() != _AUTH_HEADER.lower()
+    }
+    headers[_AUTH_HEADER] = f"Bearer {token}"
+    return headers
+
+
+class _ProxyHandler(BaseHTTPRequestHandler):
+    # Set by the server factory.
+    cache: _TokenCache
+    upstream_base: str
+
+    def log_message(self, format: str, *args: object) -> None:
+        return
+
+    def _handle(self) -> None:
+        length = int(self.headers.get("Content-Length", 0) or 0)
+        body = self.rfile.read(length) if length else None
+        target = urljoin(self.upstream_base, self.path.lstrip("/"))
+        req = urllib_request.Request(
+            target,
+            data=body,
+            method=self.command,
+            headers=_forwarded_request_headers(self, self.cache.token),
+        )
+        try:
+            with urllib_request.urlopen(req, timeout=600) as resp:
+                self._relay_response(resp.status, resp.headers, resp)
+        except urllib_error.HTTPError as exc:
+            # Upstream (gateway) error — relay status + body verbatim so the agent
+            # sees the real error (e.g. 429 rate_limit_error).
+            self._relay_response(exc.code, exc.headers, exc)
+        except (urllib_error.URLError, OSError):
+            try:
+                self.send_error(502, "opencode proxy upstream error")
+            except OSError:
+                pass
+
+    # Streaming passthrough so SSE token streaming is not buffered.
+    def _relay_response(self, status: int, headers: Message, stream: IO[bytes]) -> None:
+        try:
+            self.send_response(status)
+            for key, value in headers.items():
+                if key.lower() not in _HOP_BY_HOP:
+                    self.send_header(key, value)
+            self.end_headers()
+            while True:
+                chunk = stream.read(_STREAM_CHUNK)
+                if not chunk:
+                    break
+                self.wfile.write(chunk)
+                self.wfile.flush()
+        except (BrokenPipeError, ConnectionResetError):
+            # Client (opencode) closed the connection mid-response — routine on
+            # cancelled turns / SSE teardown; nothing left to relay to.
+            return
+
+    # Forward every method: transparent pass-through, so the gateway rejects
+    # unsupported methods rather than this proxy.
+    def __getattr__(self, name: str):
+        if name.startswith("do_"):
+            return self._handle
+        raise AttributeError(name)
+
+
+def start_proxy(
+    workspace: str, profile: str | None, port: int = 0
+) -> tuple[ThreadingHTTPServer, _TokenCache]:
+    """Start the OpenCode loopback refresh proxy + its background token refresher.
+
+    Forwards to the workspace host (all OpenCode gateway paths live under it) and
+    overwrites ``Authorization`` with a freshly-minted token on every request.
+    Binds ``port`` (default 0 = an OS-assigned free port); the caller reads
+    ``server.server_address[1]`` and points OpenCode's provider baseURLs at it.
+    Returns (server, cache); the caller runs the server (e.g. in a thread) and
+    calls shutdown()/cache.stop() on exit.
+    """
+    upstream_base = f"{workspace.rstrip('/')}/"
+    cache = _TokenCache(workspace, profile)
+
+    handler = type(
+        "BoundOpencodeProxyHandler",
+        (_ProxyHandler,),
+        {"cache": cache, "upstream_base": upstream_base},
+    )
+    try:
+        server = ThreadingHTTPServer(("127.0.0.1", port), handler)
+    except OSError:
+        server = ThreadingHTTPServer(("127.0.0.1", 0), handler)
+
+    refresher = threading.Thread(target=cache.run_refresher, daemon=True)
+    refresher.start()
+    return server, cache

tests/test_agent_opencode.py

@@ -415,3 +415,20 @@ 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"
+
+
+class TestLoopbackBaseUrls:
+    def test_swaps_host_to_loopback_keeps_gateway_path(self):
+        out = opencode._loopback_base_urls(_base_urls(), 5599)
+        assert out["anthropic"] == "http://127.0.0.1:5599/ai-gateway/anthropic/v1"
+        assert out["gemini"] == "http://127.0.0.1:5599/ai-gateway/gemini/v1beta"
+        assert out["oss"] == "http://127.0.0.1:5599/ai-gateway/mlflow/v1"
+
+    def test_all_providers_share_one_loopback_port(self):
+        out = opencode._loopback_base_urls(_base_urls(), 7002)
+        assert {u.split("/")[2] for u in out.values()} == {"127.0.0.1:7002"}
+
+
+class TestLaunchUsesProxyNotFileRefresh:
+    def test_no_background_file_refresher_remains(self):
+        assert not hasattr(opencode, "_refresh_forever")

tests/test_opencode_proxy.py

@@ -0,0 +1,55 @@
+"""Tests for the OpenCode loopback refresh proxy."""
+
+from __future__ import annotations
+
+from ucode import opencode_proxy
+
+
+class _FakeHandler:
+    """Minimal stand-in exposing a `.headers` mapping like BaseHTTPRequestHandler."""
+
+    def __init__(self, headers: dict[str, str]) -> None:
+        self.headers = headers
+
+
+class _StubCache:
+    """Avoids minting a real token when start_proxy runs under test."""
+
+    def run_refresher(self):
+        return None
+
+    def stop(self):
+        return None
+
+
+class TestForwardedRequestHeaders:
+    def test_overwrites_authorization_with_minted_token(self):
+        # OpenCode's baked-in (possibly stale) token must be replaced per request.
+        handler = _FakeHandler({"Authorization": "Bearer stale"})
+        out = opencode_proxy._forwarded_request_headers(handler, "fresh")
+        assert out["Authorization"] == "Bearer fresh"
+
+    def test_injects_authorization_when_client_sent_none(self):
+        handler = _FakeHandler({"Content-Type": "application/json"})
+        out = opencode_proxy._forwarded_request_headers(handler, "tok")
+        assert out["Authorization"] == "Bearer tok"
+
+    def test_strips_hop_by_hop_headers(self):
+        handler = _FakeHandler(
+            {"Host": "localhost:9", "Content-Length": "5", "Connection": "keep-alive"}
+        )
+        out = opencode_proxy._forwarded_request_headers(handler, "t")
+        assert "Host" not in out
+        assert "Content-Length" not in out
+        assert "Connection" not in out
+
+
+class TestStartProxy:
+    def test_forwards_to_workspace_root_so_gateway_path_is_preserved(self, monkeypatch):
+        monkeypatch.setattr(opencode_proxy, "_TokenCache", lambda workspace, profile: _StubCache())
+        server, cache = opencode_proxy.start_proxy("https://ws.databricks.com", None)
+        try:
+            assert server.RequestHandlerClass.upstream_base == "https://ws.databricks.com/"
+        finally:
+            cache.stop()
+            server.server_close()