From 40796fb39cec4a57c4126a9996df90fc65370a56 Mon Sep 17 00:00:00 2001 From: Rohit Agrawal Date: Mon, 17 Aug 2026 16:05:07 -0700 Subject: [PATCH] Fix relayed-auth token refresh: lock contention, resilience, and 401 self-heal Relayed (Claude Max/Team/Enterprise) sessions were intermittently lagging out after ~1h, sometimes needing a manual `databricks auth login` and sometimes a spurious Anthropic re-auth. Diagnostics traced this to the shared Databricks token cache: many ucode helper processes (the gateway proxy, one mcp-proxy per MCP server, web-search, tracing) each mint tokens against one lock-guarded ~/.databricks/token-cache.json, and refreshing the single expiring token at once loses the CLI's cache-write lock ("cache update: exit status 45" -> "run databricks auth login"). The 30-min force-refresh loop and swallowed errors made it worse and invisible. Fixes (all relayed-path only; non-relayed launch unchanged): - databricks.py: get_databricks_token retries transient token-cache lock contention with jittered backoff instead of treating it as a dead session. Benefits every ucode process. Verified: 8-way concurrency goes 6/8 -> 8/8. - gateway_proxy.py: refresh loop is now expiry-aware (JWT exp) and non-force, so it refreshes ~once/hour near expiry rather than writing the cache every 30 min; refreshes are single-flighted so a request burst at expiry triggers one mint. - gateway_proxy.py: refresher thread can no longer be killed by a stray non-RuntimeError, and refresh failures are surfaced instead of swallowed. - gateway_proxy.py: lazy refresh on the request path survives laptop sleep (which freezes the interval timer). - gateway_proxy.py: retry-on-401 force-refreshes the swap token and retries once; a 401 that survives is genuinely the Anthropic layer, so Claude Code's re-auth becomes correct instead of a spurious prompt. Also included (earlier proxy-robustness pass): pooled keep-alive httpx client (no per-request TLS handshake), mid-stream upstream error handling, and explicit per-op upstream timeouts. scripts/diagnose_auth.py: safe-by-default diagnostic that reproduces and confirms the token TTL, duplicate cache keys, rotation behavior, and lock-contention race. Co-authored-by: Isaac --- scripts/diagnose_auth.py | 216 +++++++++++++++++++++++++ src/ucode/agents/claude.py | 3 +- src/ucode/databricks.py | 37 ++++- src/ucode/gateway_proxy.py | 223 +++++++++++++++++++------- tests/test_databricks.py | 21 +++ tests/test_gateway_proxy.py | 306 +++++++++++++++++++++++++++++++++--- 6 files changed, 725 insertions(+), 81 deletions(-) create mode 100644 scripts/diagnose_auth.py diff --git a/scripts/diagnose_auth.py b/scripts/diagnose_auth.py new file mode 100644 index 00000000..6708f36d --- /dev/null +++ b/scripts/diagnose_auth.py @@ -0,0 +1,216 @@ +#!/usr/bin/env python3 +"""Diagnose ucode / Databricks relayed-auth failures. + +Confirms the root causes behind "sessions lag out after ~1h / require a fresh +`databricks auth login` / spuriously trigger Anthropic re-auth": + + 1. Access-token TTL — how long a minted token actually lives. + 2. Duplicate cache entries — the same workspace keyed by BOTH profile-name + and host-URL, each with its own refresh token. + 3. Refresh-token rotation — does using the refresh token invalidate the old + one? (drives the concurrent-refresh race) + 4. Live token validity — does the current token still authenticate? + 5. Concurrency race (--stress) and overnight lapse (--watch). + +Secrets (access/refresh tokens) are NEVER printed — only lengths and claims. + +SAFE BY DEFAULT: a bare run only READS ~/.databricks/token-cache.json and +decodes token claims. `--live`, `--stress`, and `--force` call the CLI and can +ROTATE your real refresh token; `--stress` may even require a re-login. + +Usage: + python scripts/diagnose_auth.py --host [--profile ] + python scripts/diagnose_auth.py --host --profile --live + python scripts/diagnose_auth.py --host --profile --stress 8 + python scripts/diagnose_auth.py --host --profile --watch 300 +""" + +from __future__ import annotations + +import argparse +import base64 +import concurrent.futures +import json +import os +import subprocess +import time +from datetime import UTC, datetime +from pathlib import Path + +CACHE = Path.home() / ".databricks" / "token-cache.json" + + +def _decode_jwt_claims(token: str) -> dict: + """Best-effort decode of a JWT payload (no signature check). {} if not a JWT.""" + try: + payload = token.split(".")[1] + payload += "=" * (-len(payload) % 4) # restore base64 padding + return json.loads(base64.urlsafe_b64decode(payload)) + except Exception: + return {} + + +def _fmt_ttl(exp: float | None) -> str: + if not exp: + return "unknown" + delta = exp - time.time() + sign = "in" if delta >= 0 else "EXPIRED" + mins = abs(delta) / 60 + return f"{sign} {mins:.1f} min" if delta >= 0 else f"{sign} {mins:.1f} min ago" + + +def _load_cache() -> dict: + if not CACHE.exists(): + return {} + return json.loads(CACHE.read_text()).get("tokens", {}) + + +def _matching_keys(host: str, profile: str | None) -> list[str]: + tokens = _load_cache() + host_frag = host.replace("https://", "").rstrip("/") + keys = [k for k in tokens if host_frag in k or (profile and k == profile)] + return keys + + +def inspect_cache(host: str, profile: str | None) -> None: + print("== 1/2. Cache entries + access-token TTL ==") + tokens = _load_cache() + keys = _matching_keys(host, profile) + if not keys: + print(f" no cache entries match host={host!r} profile={profile!r}") + return + for k in keys: + entry = tokens[k] + at = entry.get("access_token", "") + claims = _decode_jwt_claims(at) + exp = claims.get("exp") + iat = claims.get("iat") + lifetime = f"{(exp - iat) / 60:.0f} min" if exp and iat else "?" + print(f" key: {k}") + print( + f" expires_in={entry.get('expires_in')} jwt-lifetime={lifetime} " + f"exp {_fmt_ttl(exp)} refresh_token={len(entry.get('refresh_token', ''))} chars" + ) + if len(keys) > 1: + print( + f" ⚠️ {len(keys)} SEPARATE entries for this workspace — each holds its own\n" + " refresh token and refreshes independently. `--host URL` and\n" + " `--profile NAME` invocations can hit different ones, fragmenting auth." + ) + + +def _cli_token(host: str, profile: str | None, *, force: bool) -> tuple[int, str, str]: + cmd = ["databricks", "auth", "token", "--host", host, "--output", "json"] + if profile: + cmd += ["--profile", profile] + if force: + cmd += ["--force-refresh"] + env = os.environ.copy() + env["DATABRICKS_HOST"] = host + p = subprocess.run(cmd, capture_output=True, text=True, env=env, timeout=30) + return p.returncode, p.stdout, p.stderr + + +def _refresh_token_for(host: str, profile: str | None) -> str: + for k in _matching_keys(host, profile): + rt = _load_cache()[k].get("refresh_token") + if rt: + return rt + return "" + + +def check_rotation(host: str, profile: str | None) -> None: + print("\n== 3. Refresh-token rotation (sequential, safe-ish) ==") + before = _refresh_token_for(host, profile) + rc, out, err = _cli_token(host, profile, force=True) + if rc != 0: + print(f" force-refresh FAILED (rc={rc}): {err.strip()[:200]}") + print(" → this is exactly the failure that forces `databricks auth login`.") + return + after = _refresh_token_for(host, profile) + if before and after and before != after: + print(" ⚠️ refresh token ROTATED on use (old one is now invalid).") + print(" => two concurrent refreshes will race; the loser gets logged out.") + elif before and after and before == after: + print(" refresh token did NOT rotate (reusable) — concurrent refresh is safer.") + else: + print(" could not compare refresh tokens (missing before/after).") + + +def check_live(host: str, profile: str | None) -> None: + print("\n== 4. Live token validity ==") + rc, out, err = _cli_token(host, profile, force=False) + if rc != 0: + print(f" `databricks auth token` FAILED (rc={rc}): {err.strip()[:200]}") + return + token = json.loads(out or "{}").get("access_token", "") + claims = _decode_jwt_claims(token) + print(f" minted token: exp {_fmt_ttl(claims.get('exp'))} ({len(token)} chars)") + try: + import httpx + + r = httpx.get( + f"{host.rstrip('/')}/api/2.0/preview/scim/v2/Me", + headers={"Authorization": f"Bearer {token}"}, + timeout=15, + ) + print(f" workspace API /scim/v2/Me -> HTTP {r.status_code} " + f"({'valid' if r.status_code == 200 else 'REJECTED'})") + except Exception as exc: + print(f" probe skipped: {type(exc).__name__}: {exc}") + + +def stress(host: str, profile: str | None, n: int) -> None: + print(f"\n== 5a. Concurrency race: {n} simultaneous force-refresh ==") + print(" ⚠️ may rotate your token repeatedly and could require re-login.") + with concurrent.futures.ThreadPoolExecutor(max_workers=n) as ex: + results = list(ex.map(lambda _: _cli_token(host, profile, force=True)[0], range(n))) + ok = sum(1 for rc in results if rc == 0) + print(f" succeeded: {ok}/{n} failed: {n - ok}/{n}") + if n - ok: + print( + " ⚠️ concurrent refreshes RACED (expect 'cache update: exit status 45') —\n" + " the CLI serializes token-cache writes with a file lock; the losers\n" + " fail and are told to `databricks auth login`. This is the root cause." + ) + + +def watch(host: str, profile: str | None, interval: int) -> None: + print(f"\n== 5b. Watch: probing every {interval}s (Ctrl-C to stop) ==") + print(" Leave running overnight / during an idle session to catch WHEN auth lapses.") + while True: + ts = datetime.now(UTC).astimezone().strftime("%H:%M:%S") + rc, out, err = _cli_token(host, profile, force=False) + if rc != 0: + print(f" [{ts}] auth token FAILED: {err.strip()[:120]}") + else: + token = json.loads(out or "{}").get("access_token", "") + print(f" [{ts}] ok, exp {_fmt_ttl(_decode_jwt_claims(token).get('exp'))}") + time.sleep(interval) + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--host", required=True, help="workspace URL, e.g. https://x.cloud.databricks.com") + ap.add_argument("--profile", help="Databricks CLI profile name") + ap.add_argument("--live", action="store_true", help="mint a token via the CLI and probe it (mutates cache)") + ap.add_argument("--force", action="store_true", help="run the rotation check (force-refresh; mutates cache)") + ap.add_argument("--stress", type=int, metavar="N", help="N concurrent force-refreshes to reproduce the race") + ap.add_argument("--watch", type=int, metavar="SECONDS", help="poll validity on an interval (overnight repro)") + args = ap.parse_args() + + inspect_cache(args.host, args.profile) + if args.force: + check_rotation(args.host, args.profile) + if args.live: + check_live(args.host, args.profile) + if args.stress: + stress(args.host, args.profile, args.stress) + if args.watch: + watch(args.host, args.profile, args.watch) + if not any([args.force, args.live, args.stress, args.watch]): + print("\n(read-only run — add --live / --force / --stress N / --watch S to probe further)") + + +if __name__ == "__main__": + main() diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index e818cdf9..5fad5225 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -986,7 +986,7 @@ def _launch_relayed(state: dict, binary: str, tool_args: list[str]) -> None: if not isinstance(port, int): raise RuntimeError("Relayed proxy port was not configured; re-run `ucode claude`.") - server, cache = start_proxy(workspace, state.get("profile"), port) + server, cache, client = start_proxy(workspace, state.get("profile"), port) # start_proxy falls back to an OS-assigned port when the cached one is taken # (stale proxy from a killed session). Reconcile settings + state to whatever # it actually bound, so Claude Code connects to the live port. @@ -1006,6 +1006,7 @@ def _launch_relayed(state: dict, binary: str, tool_args: list[str]) -> None: finally: cache.stop() server.shutdown() + client.close() raise SystemExit(returncode) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 88935dd5..377c75c2 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -10,6 +10,7 @@ import logging.handlers import os import platform +import random import re import shlex import shutil @@ -54,6 +55,13 @@ # v1.0.0 is the release that ships `databricks aitools`. MIN_DATABRICKS_CLI_VERSION = (1, 0, 0) TOKEN_REFRESH_INTERVAL_SECONDS = 1800 +# Substrings the Databricks CLI emits when it loses the token-cache write lock +# to a concurrent `databricks auth token` (e.g. another ucode helper process or +# MLflow tracing refreshing the shared ~/.databricks/token-cache.json at the same +# instant). These are transient — the credential is fine, only the local write +# raced — so we retry rather than treat them as an expired session. +_TOKEN_CACHE_LOCK_MARKERS = ("cache update", "exit status 45") +_TOKEN_FETCH_MAX_ATTEMPTS = 4 def _debug_enabled() -> bool: @@ -1066,7 +1074,8 @@ def get_databricks_token( + f" profile={profile or ''}", ) - def _fetch() -> str: + def _fetch() -> tuple[str, str]: + """Return (access_token, stderr). token is '' on any failure.""" try: result = run( cmd, @@ -1078,12 +1087,32 @@ def _fetch() -> str: ) _debug("auth token", _format_subprocess_result(result)) if result.returncode == 0: - return json.loads(result.stdout or "{}").get("access_token", "") + return json.loads(result.stdout or "{}").get("access_token", ""), "" + return "", result.stderr or "" except (subprocess.TimeoutExpired, json.JSONDecodeError) as exc: _debug("auth token", f"exception: {type(exc).__name__}: {exc}") + return "", str(exc) + + def _fetch_with_lock_retry() -> str: + """Mint a token, retrying transient token-cache lock contention. + + Concurrent `databricks auth token` calls racing on the shared cache fail + fast with a lock error (see ``_TOKEN_CACHE_LOCK_MARKERS``). The lock is + held only for the brief cache write, so a short jittered backoff almost + always wins the next attempt. A non-lock failure returns '' immediately + so the caller can fall through to the re-auth path.""" + for attempt in range(_TOKEN_FETCH_MAX_ATTEMPTS): + token, stderr = _fetch() + if token: + return token + if not any(marker in stderr.lower() for marker in _TOKEN_CACHE_LOCK_MARKERS): + return "" + _debug("auth token", f"cache-lock contention (attempt {attempt + 1}); retrying") + if attempt < _TOKEN_FETCH_MAX_ATTEMPTS - 1: + time.sleep(random.uniform(0.05, 0.1 * (2**attempt))) return "" - token = _fetch() + token = _fetch_with_lock_retry() 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") @@ -1106,7 +1135,7 @@ def _fetch() -> str: _debug("auth login", _format_subprocess_result(reauth)) except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as exc: _debug("auth login", f"exception: {type(exc).__name__}: {exc}") - token = _fetch() + token = _fetch_with_lock_retry() if not token: profile_name = profile or find_profile_name_for_host(workspace) diff --git a/src/ucode/gateway_proxy.py b/src/ucode/gateway_proxy.py index 07c2e13c..55a9d8e9 100644 --- a/src/ucode/gateway_proxy.py +++ b/src/ucode/gateway_proxy.py @@ -17,15 +17,17 @@ from __future__ import annotations +import base64 +import binascii +import json +import sys import threading -from email.message import Message +import time 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 +import httpx + +from ucode.databricks import get_databricks_token # Header we overwrite with the freshly-minted Databricks credential. Any # client-supplied value is replaced, so a stale settings.json value can't leak. @@ -46,36 +48,114 @@ "content-length", ) ) +# Request headers the proxy manages itself and must never forward on: hop-by-hop +# plus the swap header (replaced with a freshly-minted value per request). +_STRIP_ON_FORWARD = _HOP_BY_HOP | {_SWAP_HEADER.lower()} _STREAM_CHUNK = 8192 +# Per-operation upstream timeouts. `read` is generous because model turns stream +# over a single response and Anthropic emits SSE pings, so inter-chunk gaps stay +# small; `connect`/`pool` fail fast when the gateway is unreachable. +_UPSTREAM_TIMEOUT = httpx.Timeout(connect=10.0, read=600.0, write=600.0, pool=10.0) +# Refresh once the token has less than this many seconds of life left. Databricks +# access tokens live ~1h; a 10-min buffer leaves ample headroom for a retry. +_REFRESH_BUFFER_S = 600 +# How often the background thread re-checks freshness. Cheap: it only shells out +# to the CLI when actually within the buffer, otherwise it's a bare clock compare. +_REFRESHER_POLL_S = 120 +# Assumed lifetime when a token carries no decodable `exp` (defensive fallback). +_DEFAULT_TTL_S = 3600 + + +def _jwt_exp(token: str) -> float | None: + """Best-effort `exp` (epoch seconds) from a JWT access token, else None.""" + try: + payload = token.split(".")[1] + payload += "=" * (-len(payload) % 4) # restore base64 padding + return float(json.loads(base64.urlsafe_b64decode(payload))["exp"]) + except (IndexError, ValueError, KeyError, binascii.Error, json.JSONDecodeError): + return None + + +def _log_refresh_failure(exc: BaseException) -> None: + """Surface (never silently swallow) a refresh failure, without leaking any + token or header value.""" + sys.stderr.write( + f"[ucode] Databricks token refresh failed: {exc}. If the session stalls, " + "run `databricks auth login` for your workspace profile.\n" + ) class _TokenCache: - """Holds the current Databricks token, refreshed by a background thread so - minting never blocks a request.""" + """Holds the current Databricks token and its expiry, refreshing lazily as it + nears expiry. + + A background thread refreshes proactively so the request path rarely blocks, + but the request path also refreshes on demand — which is what carries the + token across events the timer can't (laptop sleep suspends the monotonic + clock, so a fixed interval silently stops advancing). All refreshes are + single-flighted through ``_refresh_lock`` so a burst of requests at the expiry + boundary triggers exactly one CLI call, not a thundering herd on the shared + token cache.""" 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._state_lock = threading.Lock() # guards _token / _expiry (brief) + self._refresh_lock = threading.Lock() # single-flights the CLI refresh self._stop = threading.Event() + self._token = "" + self._expiry = 0.0 + # Force on start so we begin on a full-TTL token rather than inheriting a + # near-expiry one cached from an earlier CLI call. Raises if auth is dead + # (surfaced by the caller at launch, before Claude Code starts). + self._refresh(force=True) + + def _refresh(self, *, force: bool) -> None: + """Mint a token and record its expiry. Caller holds `_refresh_lock` (or is + __init__). Non-force lets a token another process just refreshed satisfy + this call from the shared cache with no write — shrinking lock contention.""" + token = get_databricks_token(self._workspace, self._profile, force_refresh=force) + expiry = _jwt_exp(token) or (time.time() + _DEFAULT_TTL_S) + with self._state_lock: + self._token = token + self._expiry = expiry + + def _fresh_enough(self) -> bool: + with self._state_lock: + return bool(self._token) and time.time() < self._expiry - _REFRESH_BUFFER_S + + def _ensure_fresh(self) -> None: + if self._fresh_enough(): + return + with self._refresh_lock: + if self._fresh_enough(): # another thread refreshed while we waited + return + try: + self._refresh(force=False) + except RuntimeError as exc: + # Keep serving the current token; a request that then 401s triggers + # a forced refresh + retry (see _ProxyHandler._handle). + _log_refresh_failure(exc) @property def token(self) -> str: - with self._lock: + self._ensure_fresh() + with self._state_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 + """Force a fresh mint now (used by the retry-on-401 path).""" + with self._refresh_lock: + self._refresh(force=True) def run_refresher(self) -> None: - while not self._stop.wait(TOKEN_REFRESH_INTERVAL_SECONDS): + while not self._stop.wait(_REFRESHER_POLL_S): try: - self.refresh() - except RuntimeError: - continue + self._ensure_fresh() + except Exception as exc: # noqa: BLE001 - a stray error must NOT kill the thread + # If this thread dies, nothing refreshes and the session lapses at + # the ~1h mark until restart. Log and keep looping instead. + _log_refresh_failure(exc) def stop(self) -> None: self._stop.set() @@ -83,9 +163,7 @@ def stop(self) -> None: 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() != _SWAP_HEADER.lower() + key: value for key, value in handler.headers.items() if key.lower() not in _STRIP_ON_FORWARD } headers[_SWAP_HEADER] = f"Bearer {token}" return headers @@ -94,55 +172,80 @@ def _forwarded_request_headers(handler: BaseHTTPRequestHandler, token: str) -> d class _ProxyHandler(BaseHTTPRequestHandler): # Set by the server factory. cache: _TokenCache - upstream_base: str + client: httpx.Client def log_message(self, format: str, *args: object) -> None: return + def _safe_send_error(self, code: int, message: str) -> None: + # The client (Claude Code) may already have disconnected, in which case + # reporting the error writes to a dead socket and raises again; swallow it. + try: + self.send_error(code, message) + except OSError: + pass + 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), - ) + url = self.path.lstrip("/") 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/Anthropic) 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): - # The client (Claude Code) may already have disconnected, in which case - # reporting the error writes to a dead socket and raises again; swallow it. + # First attempt with the current token. + headers = _forwarded_request_headers(self, self.cache.token) + with self.client.stream(self.command, url, headers=headers, content=body) as resp: + if resp.status_code not in (401, 403): + self._relay_response(resp) + return + # Auth rejected. Drain the (small) error body so the pooled + # connection can be reused, then fall through to one retry. + resp.read() + # A 401/403 may be a stale Databricks swap token rather than a bad + # Anthropic OAuth — the two are indistinguishable from the status + # alone. Force-refresh the swap token and retry once. If it was the + # Anthropic layer, the retry still 401s and we relay it verbatim, so a + # genuine re-auth is triggered; a stale-Databricks 401 self-heals here + # instead of surfacing to Claude Code as a spurious Anthropic prompt. try: - self.send_error(502, "gateway proxy upstream error") - except OSError: - pass + self.cache.refresh() + except RuntimeError: + pass # refresh failed; retry with the existing token and relay whatever comes + headers = _forwarded_request_headers(self, self.cache.token) + with self.client.stream(self.command, url, headers=headers, content=body) as resp: + self._relay_response(resp) + except (BrokenPipeError, ConnectionResetError): + # Client closed before/while we relayed headers — routine on cancel. + return + except httpx.HTTPError: + # Upstream failed before any bytes reached the client; a 502 is still + # sendable. (An HTTP *status* like 429 is not an error here — httpx + # only raises for transport failures — so real gateway errors are + # relayed verbatim by `_relay_response`.) + self._safe_send_error(502, "gateway proxy upstream error") # Streaming passthrough: forward chunks as they arrive so SSE token streaming # is not buffered (buffering would add full-response latency to first token). - def _relay_response(self, status: int, headers: Message, stream: IO[bytes]) -> None: + # `iter_raw` preserves any Content-Encoding verbatim (we relay that header), + # so the proxy stays byte-transparent. + def _relay_response(self, resp: httpx.Response) -> None: try: - self.send_response(status) - for key, value in headers.items(): + self.send_response(resp.status_code) + for key, value in resp.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() + for chunk in resp.iter_raw(_STREAM_CHUNK): + if chunk: + self.wfile.write(chunk) + self.wfile.flush() except (BrokenPipeError, ConnectionResetError): # Client (Claude Code) closed the connection mid-response — routine on - # cancelled turns / SSE teardown. There is nothing left to relay to, so - # stop quietly rather than crashing the handler thread. + # cancelled turns / SSE teardown. Nothing left to relay to, so stop + # quietly rather than crashing the handler thread. + return + except httpx.HTTPError: + # Upstream dropped mid-stream. Headers (and status) may already be + # sent, so we can't reliably signal a fresh error — stop and let the + # client see a truncated stream rather than corrupt the framing. return # Forward every method: this is a transparent pass-through, so routing any @@ -155,22 +258,28 @@ def __getattr__(self, name: str): def start_proxy( workspace: str, profile: str | None, port: int -) -> tuple[ThreadingHTTPServer, _TokenCache]: +) -> tuple[ThreadingHTTPServer, _TokenCache, httpx.Client]: """Start the loopback refresh proxy + its background token refresher. Binds ``port``, falling back to a fresh OS-assigned port when it is already in use (e.g. a prior session's proxy that was killed before its teardown ran still holds the socket). The caller reads ``server.server_address[1]`` for the - actual port and points Claude Code at it. Returns (server, cache); the caller - runs the server (e.g. in a thread) and calls shutdown()/cache.stop() on exit. + actual port and points Claude Code at it. + + Returns (server, cache, client); the caller runs the server (e.g. in a + thread) and calls shutdown()/cache.stop()/client.close() on exit. """ upstream_base = f"{workspace.rstrip('/')}/ai-gateway/anthropic/" cache = _TokenCache(workspace, profile) + # One pooled, keep-alive client shared across handler threads: reuses TCP+TLS + # to the gateway instead of a fresh handshake per request. Don't follow + # redirects — a proxy relays 3xx verbatim. + client = httpx.Client(base_url=upstream_base, timeout=_UPSTREAM_TIMEOUT, follow_redirects=False) handler = type( "BoundProxyHandler", (_ProxyHandler,), - {"cache": cache, "upstream_base": upstream_base}, + {"cache": cache, "client": client}, ) try: server = ThreadingHTTPServer(("127.0.0.1", port), handler) @@ -181,4 +290,4 @@ def start_proxy( refresher = threading.Thread(target=cache.run_refresher, daemon=True) refresher.start() - return server, cache + return server, cache, client diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 64b4dcc2..7ef8b2da 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -1523,6 +1523,27 @@ def test_reauths_and_retries_when_token_empty(self, tmp_path, monkeypatch): token = get_databricks_token(WS) assert token == "refreshed-token" + def test_retries_on_cache_lock_contention(self, tmp_path, monkeypatch): + # Concurrent `databricks auth token` calls racing on the shared token + # cache fail with "cache update: exit status 45". That's transient (the + # credential is fine), so we must retry — not treat it as a dead session. + call_count = tmp_path / "calls" + call_count.write_text("0") + env = self._fake_databricks( + tmp_path, + f"count=$(cat {call_count})\n" + f"echo $((count + 1)) > {call_count}\n" + 'if [ "$count" -lt 2 ]; then\n' + ' echo "Error: forced token refresh: cache update: exit status 45" >&2\n' + " exit 1\n" + "else\n" + ' echo \'{"access_token": "won-the-lock", "token_type": "Bearer"}\'\n' + "fi", + ) + monkeypatch.setattr("os.environ", env) + token = get_databricks_token(WS) + assert token == "won-the-lock" + def test_raises_when_reauth_also_fails(self, tmp_path, monkeypatch): env = self._fake_databricks( tmp_path, diff --git a/tests/test_gateway_proxy.py b/tests/test_gateway_proxy.py index 68c8eeb8..e367d0f5 100644 --- a/tests/test_gateway_proxy.py +++ b/tests/test_gateway_proxy.py @@ -2,13 +2,25 @@ from __future__ import annotations +import base64 import io +import json import socket -from email.message import Message +import threading +import time + +import httpx from ucode import gateway_proxy +def _make_jwt(exp: float | None) -> str: + """A minimal JWT-shaped token whose payload carries `exp` (or none).""" + claims = {"exp": exp} if exp is not None else {"sub": "x"} + payload = base64.urlsafe_b64encode(json.dumps(claims).encode()).rstrip(b"=").decode() + return f"header.{payload}.sig" + + class _FakeHandler: """Minimal stand-in exposing a `.headers` mapping like BaseHTTPRequestHandler.""" @@ -45,6 +57,18 @@ def test_strips_hop_by_hop_headers(self): assert "Connection" not in out +class _FakeResponse: + """Stand-in for httpx.Response exposing only what `_relay_response` reads.""" + + def __init__(self, status_code: int, headers: dict[str, str], chunks): + self.status_code = status_code + self.headers = headers + self._chunks = chunks + + def iter_raw(self, chunk_size=None): + yield from self._chunks + + class _BrokenPipeWriter(io.RawIOBase): """A wfile stand-in that raises BrokenPipeError on write, mimicking a client (Claude Code) that closed the connection mid-response.""" @@ -53,25 +77,26 @@ def write(self, _data): # type: ignore[override] raise BrokenPipeError(32, "Broken pipe") -class TestRelayResponseClientDisconnect: - def _handler(self, wfile) -> gateway_proxy._ProxyHandler: - # Bypass BaseHTTPRequestHandler.__init__ (which would service a socket); - # we only exercise _relay_response's write path. Set the few attributes the - # send_response/send_header machinery reads (normally populated by __init__). - handler = object.__new__(gateway_proxy._ProxyHandler) - handler.wfile = wfile - handler.request_version = "HTTP/1.1" - handler.requestline = "POST /v1/messages HTTP/1.1" - handler.command = "POST" - handler._headers_buffer = [] - return handler +def _relay_handler(wfile) -> gateway_proxy._ProxyHandler: + # Bypass BaseHTTPRequestHandler.__init__ (which would service a socket); + # we only exercise _relay_response's write path. Set the few attributes the + # send_response/send_header machinery reads (normally populated by __init__). + handler = object.__new__(gateway_proxy._ProxyHandler) + handler.wfile = wfile + handler.request_version = "HTTP/1.1" + handler.requestline = "POST /v1/messages HTTP/1.1" + handler.command = "POST" + handler._headers_buffer = [] + return handler + +class TestRelayResponseClientDisconnect: def test_relay_swallows_broken_pipe_on_headers(self): # Client gone before headers flush: end_headers write raises BrokenPipe. - handler = self._handler(_BrokenPipeWriter()) - stream = io.BytesIO(b'{"ok":true}') + handler = _relay_handler(_BrokenPipeWriter()) + resp = _FakeResponse(200, {}, [b'{"ok":true}']) # Must not raise — a dead client is a routine teardown, not an error. - handler._relay_response(200, Message(), stream) + handler._relay_response(resp) def test_relay_swallows_connection_reset_mid_stream(self): # Headers flush ok, then the client resets while streaming body chunks. @@ -87,8 +112,250 @@ def write(self, data): # type: ignore[override] def flush(self): return None - handler = self._handler(_ResetAfterHeaders()) - handler._relay_response(200, Message(), io.BytesIO(b"chunk-of-sse-data")) + handler = _relay_handler(_ResetAfterHeaders()) + resp = _FakeResponse(200, {}, [b"chunk-of-sse-data"]) + handler._relay_response(resp) + + def test_relay_swallows_upstream_error_mid_stream(self): + # Upstream drops mid-body after headers are already sent — we can't signal + # a fresh error, so stop quietly rather than corrupt the response framing. + class _Ok(io.RawIOBase): + def write(self, data): # type: ignore[override] + return len(data) + + def flush(self): + return None + + def _chunks(): + yield b"partial" + raise httpx.ReadError("upstream dropped") + + handler = _relay_handler(_Ok()) + resp = _FakeResponse(200, {}, _chunks()) + handler._relay_response(resp) # must not raise + + def test_relay_forwards_status_and_skips_hop_by_hop_headers(self): + # A non-200 status (e.g. 429 rate limit) and content headers are relayed; + # hop-by-hop framing headers are dropped. + chunks_written: list[bytes] = [] + + class _Collect(io.RawIOBase): + def write(self, data): # type: ignore[override] + chunks_written.append(bytes(data)) + return len(data) + + def flush(self): + return None + + handler = _relay_handler(_Collect()) + resp = _FakeResponse( + 429, + {"Content-Type": "application/json", "Transfer-Encoding": "chunked"}, + [b'{"type":"error"}'], + ) + handler._relay_response(resp) + blob = b"".join(chunks_written) + assert b"429" in blob + assert b"Content-Type: application/json" in blob + assert b"Transfer-Encoding" not in blob # hop-by-hop, stripped + + +class TestJwtExp: + def test_extracts_exp(self): + assert gateway_proxy._jwt_exp(_make_jwt(1234567890.0)) == 1234567890.0 + + def test_none_on_missing_exp(self): + assert gateway_proxy._jwt_exp(_make_jwt(None)) is None + + def test_none_on_garbage(self): + assert gateway_proxy._jwt_exp("not-a-jwt") is None + + +def _install_fake_token(monkeypatch, exp_offsets, delay=0.0): + """Patch get_databricks_token to hand out JWTs whose exp is now+offset, one + per successive mint (last offset repeats). Records the force flag of each.""" + state = {"i": 0, "forces": []} + + def fake(_ws, _profile, force_refresh=False): + if delay: + time.sleep(delay) + off = exp_offsets[min(state["i"], len(exp_offsets) - 1)] + state["i"] += 1 + state["forces"].append(force_refresh) + return _make_jwt(time.time() + off) + + monkeypatch.setattr(gateway_proxy, "get_databricks_token", fake) + return state + + +class TestTokenCache: + def test_initial_mint_is_forced(self, monkeypatch): + state = _install_fake_token(monkeypatch, [5000]) + gateway_proxy._TokenCache("ws", None) + assert state["forces"] == [True] # full-TTL start + + def test_fresh_token_is_not_refreshed(self, monkeypatch): + state = _install_fake_token(monkeypatch, [5000]) + cache = gateway_proxy._TokenCache("ws", None) + _ = cache.token + _ = cache.token + assert state["forces"] == [True] # no extra mint while fresh + + def test_near_expiry_triggers_nonforce_refresh(self, monkeypatch): + # First mint expires within the buffer -> reading .token refreshes once, + # non-force (so a token another process just wrote can satisfy it). + state = _install_fake_token(monkeypatch, [100, 5000]) + cache = gateway_proxy._TokenCache("ws", None) + _ = cache.token + assert state["forces"] == [True, False] + _ = cache.token # now fresh again + assert state["forces"] == [True, False] + + def test_refresh_is_single_flighted(self, monkeypatch): + # A burst of concurrent requests at the expiry boundary must trigger ONE + # refresh, not a thundering herd on the shared token cache. + state = _install_fake_token(monkeypatch, [100, 5000], delay=0.05) + cache = gateway_proxy._TokenCache("ws", None) + threads = [threading.Thread(target=lambda: cache.token) for _ in range(10)] + for t in threads: + t.start() + for t in threads: + t.join() + # 1 forced init + exactly 1 non-force refresh shared by all 10 readers. + assert state["forces"] == [True, False] + + def test_ensure_fresh_keeps_token_when_refresh_fails(self, monkeypatch): + _install_fake_token(monkeypatch, [5000]) + cache = gateway_proxy._TokenCache("ws", None) + good = cache.token + + def boom(*_a, **_k): + raise RuntimeError("mint failed") + + monkeypatch.setattr(gateway_proxy, "get_databricks_token", boom) + # Force staleness so _ensure_fresh attempts a refresh, which now fails. + cache._expiry = time.time() + assert cache.token == good # last good token retained, no exception + + def test_refresher_loop_survives_unexpected_error(self, monkeypatch): + _install_fake_token(monkeypatch, [5000]) + cache = gateway_proxy._TokenCache("ws", None) + monkeypatch.setattr(gateway_proxy, "_REFRESHER_POLL_S", 0.01) + ticks = [] + + def boom(): + ticks.append(1) + cache.stop() # let the loop exit after this tick + raise ValueError("unexpected, non-RuntimeError") + + monkeypatch.setattr(cache, "_ensure_fresh", boom) + cache.run_refresher() # must RETURN, not propagate — else the thread dies + assert ticks == [1] + + +class _FakeResp: + def __init__(self, status: int, body: bytes = b"", headers: dict | None = None): + self.status_code = status + self._body = body + self.headers = headers or {} + self.read_called = False + + def read(self): + self.read_called = True + return self._body + + def iter_raw(self, _n=None): + yield self._body + + def __enter__(self): + return self + + def __exit__(self, *_a): + return False + + +class _FakeClient: + def __init__(self, responses): + self._responses = list(responses) + self.sent_tokens: list[str | None] = [] + + def stream(self, _method, _url, headers, content): + self.sent_tokens.append(headers.get(gateway_proxy._SWAP_HEADER)) + return self._responses.pop(0) + + +class _FakeCache: + def __init__(self): + self._t = "Bearer-tok1" + self.refreshed = 0 + + @property + def token(self): + return self._t.replace("Bearer-", "") + + def refresh(self): + self.refreshed += 1 + self._t = "Bearer-tok2" + + +def _handle_handler(client, cache, wfile) -> gateway_proxy._ProxyHandler: + h = object.__new__(gateway_proxy._ProxyHandler) + h.client = client + h.cache = cache + h.headers = {} + h.rfile = io.BytesIO(b"") + h.path = "/v1/messages" + h.command = "POST" + h.wfile = wfile + h.request_version = "HTTP/1.1" + h.requestline = "POST /v1/messages HTTP/1.1" + h._headers_buffer = [] + return h + + +class _Collect(io.RawIOBase): + def __init__(self): + self.data = bytearray() + + def write(self, b): # type: ignore[override] + self.data += bytes(b) + return len(b) + + def flush(self): + return None + + +class TestRetryOn401: + def test_401_forces_refresh_and_retries(self): + # A stale swap token yields 401; the proxy force-refreshes and retries, + # this time succeeding, so Claude Code never sees the 401. + client = _FakeClient([_FakeResp(401, b'{"e":1}'), _FakeResp(200, b"ok")]) + cache = _FakeCache() + out = _Collect() + _handle_handler(client, cache, out)._handle() + assert cache.refreshed == 1 + assert client.sent_tokens == ["Bearer tok1", "Bearer tok2"] # retried with fresh token + assert b"200" in bytes(out.data) + assert b"ok" in bytes(out.data) + + def test_persistent_401_is_relayed(self): + # If the retry also 401s, it's genuinely the Anthropic layer — relay it so + # Claude Code re-auths Anthropic (correct), rather than looping forever. + client = _FakeClient([_FakeResp(401, b"a"), _FakeResp(401, b"b")]) + cache = _FakeCache() + out = _Collect() + _handle_handler(client, cache, out)._handle() + assert cache.refreshed == 1 + assert b"401" in bytes(out.data) + + def test_success_first_try_does_not_refresh(self): + client = _FakeClient([_FakeResp(200, b"hi")]) + cache = _FakeCache() + out = _Collect() + _handle_handler(client, cache, out)._handle() + assert cache.refreshed == 0 + assert client.sent_tokens == ["Bearer tok1"] + assert b"hi" in bytes(out.data) class TestStartProxyPortFallback: @@ -107,7 +374,7 @@ def run_refresher(self): occupied.listen(1) busy_port = occupied.getsockname()[1] try: - server, _cache = gateway_proxy.start_proxy( + server, _cache, client = gateway_proxy.start_proxy( "https://x.staging.cloud.databricks.com", None, busy_port ) try: @@ -116,5 +383,6 @@ def run_refresher(self): assert bound != 0 finally: server.server_close() + client.close() finally: occupied.close()