diff --git a/src/ucode/agents/opencode.py b/src/ucode/agents/opencode.py index 19adff71..60fa48ad 100644 --- a/src/ucode/agents/opencode.py +++ b/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: 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) diff --git a/src/ucode/opencode_proxy.py b/src/ucode/opencode_proxy.py new file mode 100644 index 00000000..06917c9a --- /dev/null +++ b/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 diff --git a/tests/test_agent_opencode.py b/tests/test_agent_opencode.py index 9e32e78a..15b9cf9a 100644 --- a/tests/test_agent_opencode.py +++ b/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") diff --git a/tests/test_opencode_proxy.py b/tests/test_opencode_proxy.py new file mode 100644 index 00000000..2dab31b4 --- /dev/null +++ b/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()