claude: authenticate gateway model discovery

src/ucode/agents/claude.py

@@ -29,6 +29,7 @@
     build_tool_base_url,
     get_databricks_token,
 )
+from ucode.gateway_proxy import AI_GATEWAY_TOKEN_HEADER, AUTHORIZATION_HEADER, start_proxy
 from ucode.launcher import exec_or_spawn
 from ucode.managed_files import OS, current_os, write_managed_file
 from ucode.smart_routing.claude_hooks import (
@@ -40,10 +41,10 @@
 from ucode.tracing import tracing_env
 from ucode.ui import print_err, print_note, print_success, print_warning
 
-GATEWAY_MODEL_DISCOVERY_ENV_VAR = "ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY"
 CLAUDE_CONFIG_DIR = Path.home() / ".claude"
 CLAUDE_SETTINGS_PATH = CLAUDE_CONFIG_DIR / "ucode-settings.json"
 CLAUDE_BACKUP_PATH = APP_DIR / "claude-ucode-settings.backup.json"
+GATEWAY_MODEL_DISCOVERY_ENV_VAR = "ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY"
 
 SPEC: ToolSpec = {
     "binary": "claude",
@@ -875,7 +876,12 @@ def _merge_claude_settings(base: dict, overlay: dict) -> dict:
     return merged
 
 
-def _build_claude_argv(binary: str, tool_args: list[str], relayed: bool = False) -> list[str]:
+def _build_claude_argv(
+    binary: str,
+    tool_args: list[str],
+    relayed: bool = False,
+    settings_override: dict | None = None,
+) -> list[str]:
     """Build the ``claude`` argv, composing any caller ``--settings`` with
     ucode's managed settings.
 
@@ -898,7 +904,7 @@ def _build_claude_argv(binary: str, tool_args: list[str], relayed: bool = False)
     """
     source_args = ["--setting-sources", _RELAYED_SETTING_SOURCES] if relayed else []
     caller_values, remaining = _extract_caller_settings(tool_args)
-    if not caller_values:
+    if not caller_values and settings_override is None:
         # No caller --settings: hand Claude ucode's settings file directly (the
         # common path; behavior unchanged).
         return [binary, *source_args, "--settings", str(CLAUDE_SETTINGS_PATH), *tool_args]
@@ -908,6 +914,8 @@ def _build_claude_argv(binary: str, tool_args: list[str], relayed: bool = False)
     # ucode wins over the caller for conflicting keys (protects gateway auth);
     # hooks from both sides survive.
     merged = _merge_claude_settings(caller_settings, read_json_safe(CLAUDE_SETTINGS_PATH))
+    if settings_override is not None:
+        merged = _merge_claude_settings(merged, settings_override)
     return [
         binary,
         *source_args,
@@ -966,8 +974,6 @@ def _launch_relayed(state: dict, binary: str, tool_args: list[str]) -> None:
     """Relayed launch: sign into the Claude subscription, start the loopback
     refresh proxy, then run Claude Code alongside it (the proxy must outlive the
     exec, so we spawn-and-wait rather than replacing the process)."""
-    from ucode.gateway_proxy import start_proxy
-
     conflict = _managed_relayed_conflicts()
     if conflict is not None:
         managed_path, keys = conflict
@@ -993,7 +999,13 @@ 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, client = start_proxy(workspace, state.get("profile"), port)
+    server, cache, client = start_proxy(
+        workspace,
+        state.get("profile"),
+        port,
+        token_header=AI_GATEWAY_TOKEN_HEADER,
+        force_refresh_near_expiry=False,
+    )
     # 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.
@@ -1017,12 +1029,50 @@ def _launch_relayed(state: dict, binary: str, tool_args: list[str]) -> None:
     raise SystemExit(returncode)
 
 
+def _launch_gateway(state: dict, binary: str, tool_args: list[str]) -> None:
+    workspace = state["workspace"]
+    server, cache, client = start_proxy(
+        workspace,
+        state.get("profile"),
+        0,
+        token_header=AUTHORIZATION_HEADER,
+        force_refresh_near_expiry=True,
+    )
+    token = cache.token
+    os.environ["OAUTH_TOKEN"] = token
+    os.environ["ANTHROPIC_AUTH_TOKEN"] = token
+    os.environ["ANTHROPIC_BASE_URL"] = f"http://127.0.0.1:{server.server_address[1]}"
+    os.environ["CLAUDE_CODE_USE_GATEWAY"] = "1"
+
+    server_thread = threading.Thread(target=server.serve_forever, daemon=True)
+    server_thread.start()
+    settings_override = {
+        "env": {"ANTHROPIC_BASE_URL": os.environ["ANTHROPIC_BASE_URL"]},
+    }
+    proc = subprocess.Popen(
+        _build_claude_argv(binary, tool_args, settings_override=settings_override)
+    )
+    try:
+        returncode = proc.wait()
+    except KeyboardInterrupt:
+        proc.send_signal(signal.SIGINT)
+        returncode = proc.wait()
+    finally:
+        cache.stop()
+        server.shutdown()
+        client.close()
+    raise SystemExit(returncode)
+
+
 def launch(state: dict, tool_args: list[str]) -> None:
     binary = SPEC["binary"]
     workspace = state.get("workspace")
     if state.get("claude_relayed"):
         _launch_relayed(state, binary, tool_args)
         return
+    if workspace and os.environ.get(GATEWAY_MODEL_DISCOVERY_ENV_VAR) == "1":
+        _launch_gateway(state, binary, tool_args)
+        return
     if workspace:
         os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile"))
     exec_or_spawn(_build_claude_argv(binary, tool_args))

src/ucode/gateway_proxy.py

@@ -1,18 +1,17 @@
-"""Loopback refresh proxy for relayed Anthropic (Claude Max/Team/Enterprise).
+"""Loopback refresh proxy for Claude gateway requests.
 
 A relayed Model Provider Service authenticates the caller's own Anthropic
 subscription OAuth (which Claude Code owns in the `Authorization` header) and
 carries a Databricks credential in the `X-Databricks-AI-Gateway-Token` swap
-header. That Databricks token is short-lived and a static settings.json header
-can't be refreshed, so `ucode claude` points `ANTHROPIC_BASE_URL` at this proxy
-instead: it forwards every request to the workspace gateway unchanged except for
-adding a freshly-minted swap header, and streams the response back verbatim.
+header. Native gateway discovery instead carries the Databricks credential in
+`Authorization`. The proxy refreshes the applicable header and streams responses
+back verbatim.
 
 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; the Anthropic OAuth in `Authorization` is
-    passed through untouched and never read, stored, or logged.
+    passed through untouched in relayed mode and never logged.
 """
 
 from __future__ import annotations
@@ -33,7 +32,8 @@
 
 # Header we overwrite with the freshly-minted Databricks credential. Any
 # client-supplied value is replaced, so a stale settings.json value can't leak.
-_SWAP_HEADER = "X-Databricks-AI-Gateway-Token"
+AI_GATEWAY_TOKEN_HEADER = "X-Databricks-AI-Gateway-Token"
+AUTHORIZATION_HEADER = "Authorization"
 # Hop-by-hop headers must not be forwarded across the proxy.
 _HOP_BY_HOP = frozenset(
     h.lower()
@@ -50,9 +50,6 @@
         "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()}
 # 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.
@@ -117,23 +114,27 @@ class _TokenCache:
     boundary triggers exactly one CLI call, not a thundering herd on the shared
     token cache."""
 
-    def __init__(self, workspace: str, profile: str | None) -> None:
+    def __init__(
+        self,
+        workspace: str,
+        profile: str | None,
+        *,
+        force_refresh_near_expiry: bool = False,
+    ) -> None:
         self._workspace = workspace
         self._profile = profile
+        self._force_refresh_near_expiry = force_refresh_near_expiry
         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)
+        # Preserve the existing non-forced relayed-auth fetch. Gateway discovery
+        # opts into a forced fetch so its static client token starts with a full TTL.
+        self._refresh(force=force_refresh_near_expiry)
 
     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."""
+        """Mint a token and record its expiry."""
         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:
@@ -151,7 +152,7 @@ def _ensure_fresh(self) -> None:
             if self._fresh_enough():  # another thread refreshed while we waited
                 return
             try:
-                self._refresh(force=False)
+                self._refresh(force=self._force_refresh_near_expiry)
             except RuntimeError as exc:
                 # Keep serving the current token; a request that then 401s triggers
                 # a forced refresh + retry (see _ProxyHandler._handle).
@@ -181,18 +182,24 @@ def stop(self) -> None:
         self._stop.set()
 
 
-def _forwarded_request_headers(handler: BaseHTTPRequestHandler, token: str) -> dict[str, str]:
+def _forwarded_request_headers(
+    handler: BaseHTTPRequestHandler,
+    token: str,
+    token_header: str = AI_GATEWAY_TOKEN_HEADER,
+) -> dict[str, str]:
+    strip_on_forward = _HOP_BY_HOP | {token_header.lower()}
     headers = {
-        key: value for key, value in handler.headers.items() if key.lower() not in _STRIP_ON_FORWARD
+        key: value for key, value in handler.headers.items() if key.lower() not in strip_on_forward
     }
-    headers[_SWAP_HEADER] = f"Bearer {token}"
+    headers[token_header] = f"Bearer {token}"
     return headers
 
 
 class _ProxyHandler(BaseHTTPRequestHandler):
     # Set by the server factory.
     cache: _TokenCache
     client: httpx.Client
+    token_header = AI_GATEWAY_TOKEN_HEADER
 
     def log_message(self, format: str, *args: object) -> None:
         return
@@ -219,7 +226,7 @@ def _handle(self) -> None:
         )
         try:
             # First attempt with the current token.
-            headers = _forwarded_request_headers(self, self.cache.token)
+            headers = _forwarded_request_headers(self, self.cache.token, self.token_header)
             with self.client.stream(self.command, url, headers=headers, content=body) as resp:
                 _diagnostic_log(
                     "upstream_headers",
@@ -234,9 +241,9 @@ def _handle(self) -> None:
                 # 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
+            # A relayed 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 Databricks 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.
@@ -249,7 +256,7 @@ def _handle(self) -> None:
                 # which otherwise reads as an Anthropic `/login` prompt and sends the
                 # user to the wrong re-auth. Still retry + relay with the existing token.
                 _log_refresh_failure(exc)
-            headers = _forwarded_request_headers(self, self.cache.token)
+            headers = _forwarded_request_headers(self, self.cache.token, self.token_header)
             with self.client.stream(self.command, url, headers=headers, content=body) as resp:
                 _diagnostic_log(
                     "upstream_headers",
@@ -362,7 +369,11 @@ def __getattr__(self, name: str):
 
 
 def start_proxy(
-    workspace: str, profile: str | None, port: int
+    workspace: str,
+    profile: str | None,
+    port: int,
+    token_header: str,
+    force_refresh_near_expiry: bool,
 ) -> tuple[ThreadingHTTPServer, _TokenCache, httpx.Client]:
     """Start the loopback refresh proxy + its background token refresher.
 
@@ -375,7 +386,11 @@ def start_proxy(
     thread) and calls shutdown()/cache.stop()/client.close() on exit.
     """
     upstream_base = f"{workspace.rstrip('/')}/ai-gateway/anthropic/"
-    cache = _TokenCache(workspace, profile)
+    cache = _TokenCache(
+        workspace,
+        profile,
+        force_refresh_near_expiry=force_refresh_near_expiry,
+    )
     # 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.
@@ -384,7 +399,7 @@ def start_proxy(
     handler = type(
         "BoundProxyHandler",
         (_ProxyHandler,),
-        {"cache": cache, "client": client},
+        {"cache": cache, "client": client, "token_header": token_header},
     )
     try:
         server = ThreadingHTTPServer(("127.0.0.1", port), handler)