From 8312b8d286858e78517daeb3f9a298d2970a8185 Mon Sep 17 00:00:00 2001
From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com>
Date: Fri, 28 Aug 2026 19:05:25 +0000
Subject: [PATCH 1/4] Require Claude Code 2.1.248 for model discovery

---
 src/ucode/agents/claude.py                    | 72 +++++++++++++++----
 src/ucode/anthropic_model_discovery_proxy.py  | 40 +++++++----
 src/ucode/gateway_proxy.py                    | 21 ++++--
 tests/test_agent_claude.py                    | 62 ++++++++++------
 tests/test_anthropic_model_discovery_proxy.py | 34 +++++++++
 tests/test_gateway_proxy.py                   |  7 --
 6 files changed, 171 insertions(+), 65 deletions(-)

diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py
index 92d18f2b..d7102e27 100644
--- a/src/ucode/agents/claude.py
+++ b/src/ucode/agents/claude.py
@@ -33,10 +33,7 @@
     build_tool_base_url,
     get_databricks_token,
 )
-from ucode.gateway_proxy import (
-    AI_GATEWAY_TOKEN_HEADER,
-    AUTHORIZATION_HEADER,
-)
+from ucode.gateway_proxy import AI_GATEWAY_TOKEN_HEADER
 from ucode.launcher import exec_or_spawn
 from ucode.managed_files import OS, current_os, write_managed_file
 from ucode.smart_routing import v2 as smart_routing_v2
@@ -58,6 +55,8 @@
 CLAUDE_USER_SETTINGS_PATH = CLAUDE_CONFIG_DIR / "settings.json"
 CLAUDE_BACKUP_PATH = APP_DIR / "claude-ucode-settings.backup.json"
 WEB_SEARCH_MCP_STATE_KEY = "claude_web_search_mcp"
+MINIMUM_CLAUDE_VERSION = (2, 1, 248)
+MINIMUM_CLAUDE_VERSION_TEXT = "2.1.248"
 
 SPEC: ToolSpec = {
     "binary": "claude",
@@ -98,6 +97,49 @@ def is_update_available() -> tuple[str, str] | None:
     return available_npm_package_update(SPEC["package"])
 
 
+def _parse_version(value: str) -> tuple[int, int, int] | None:
+    match = re.search(r"(\d+)\.(\d+)\.(\d+)", value)
+    if not match:
+        return None
+    major, minor, patch = match.groups()
+    return int(major), int(minor), int(patch)
+
+
+def _installed_version_status() -> tuple[str, bool] | None:
+    version = agent_version(SPEC["binary"])
+    parsed = _parse_version(version)
+    if parsed is None:
+        return None
+    return version, parsed < MINIMUM_CLAUDE_VERSION
+
+
+def minimum_version_error() -> str | None:
+    status = _installed_version_status()
+    if status is None:
+        return None
+    version, is_too_old = status
+    if not is_too_old:
+        return None
+    return (
+        f"Claude Code {version} is too old for gateway model discovery. "
+        f"Claude Code must be updated to {MINIMUM_CLAUDE_VERSION_TEXT} or newer; "
+        f"run `npm install -g {SPEC['package']}` or `ucode configure`."
+    )
+
+
+def required_update_message() -> str | None:
+    status = _installed_version_status()
+    if status is None:
+        return None
+    version, is_too_old = status
+    if not is_too_old:
+        return None
+    return (
+        f"Claude Code {version} is older than required {MINIMUM_CLAUDE_VERSION_TEXT}; "
+        "updating Claude Code is required for gateway model discovery."
+    )
+
+
 def _resolve_web_search_model(state: dict) -> str | None:
     """Pick the model the web_search MCP server should call. Prefers an
     explicit override in state, otherwise the first endpoint discovered as
@@ -1129,6 +1171,8 @@ def _launch_relayed(state: dict, binary: str, tool_args: list[str]) -> None:
         token_header=AI_GATEWAY_TOKEN_HEADER,
         force_refresh_near_expiry=False,
     )
+    if cache is None:  # defensive: relayed mode always requests swap-token refresh
+        raise RuntimeError("Relayed proxy did not initialize its token cache.")
     # 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.
@@ -1152,21 +1196,20 @@ def _launch_relayed(state: dict, binary: str, tool_args: list[str]) -> None:
     raise SystemExit(returncode)
 
 
-def _launch_claude_with_gateway_proxy(
+def _launch_claude_with_model_discovery_proxy(
     state: dict, binary: str, tool_args: list[str], *, smart_routing: bool
 ) -> None:
-    """Launch Claude through a refreshing gateway proxy."""
+    """Launch Claude through the model-alias proxy using apiKeyHelper auth."""
     workspace = state["workspace"]
     server, cache, client = start_anthropic_model_discovery_proxy(
         workspace,
         state.get("profile"),
         0,
-        token_header=AUTHORIZATION_HEADER,
-        force_refresh_near_expiry=True,
+        token_header=None,
+        force_refresh_near_expiry=False,
     )
-    token = cache.token
-    os.environ["OAUTH_TOKEN"] = token
-    os.environ["ANTHROPIC_AUTH_TOKEN"] = token
+    if not smart_routing:
+        os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile"))
     os.environ["ANTHROPIC_BASE_URL"] = f"http://{LOOPBACK_HOST}:{server.server_address[1]}"
     os.environ["CLAUDE_CODE_USE_GATEWAY"] = "1"
 
@@ -1200,7 +1243,8 @@ def compose_gateway_settings(args: list[str]) -> tuple[dict, list[str]]:
             proc.send_signal(signal.SIGINT)
             returncode = proc.wait()
     finally:
-        cache.stop()
+        if cache is not None:
+            cache.stop()
         server.shutdown()
         client.close()
     raise SystemExit(returncode)
@@ -1227,10 +1271,10 @@ def launch(state: dict, tool_args: list[str]) -> None:
             "Please use Codex or disable smart routing."
         )
     if first_prompt_routing:
-        _launch_claude_with_gateway_proxy(state, binary, tool_args, smart_routing=True)
+        _launch_claude_with_model_discovery_proxy(state, binary, tool_args, smart_routing=True)
         return
     if workspace and os.environ.get(GATEWAY_MODEL_DISCOVERY_ENV_VAR) == "1":
-        _launch_claude_with_gateway_proxy(state, binary, tool_args, smart_routing=False)
+        _launch_claude_with_model_discovery_proxy(state, binary, tool_args, smart_routing=False)
         return
     if workspace:
         os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile"))
diff --git a/src/ucode/anthropic_model_discovery_proxy.py b/src/ucode/anthropic_model_discovery_proxy.py
index 989616f6..67174a34 100644
--- a/src/ucode/anthropic_model_discovery_proxy.py
+++ b/src/ucode/anthropic_model_discovery_proxy.py
@@ -32,14 +32,15 @@
     forwarded_request_headers,
     log_proxy_diagnostic,
     log_token_refresh_failure,
+    passthrough_request_headers,
 )
 
 
 class _ProxyHandler(BaseHTTPRequestHandler):
     # Set by the server factory.
-    cache: TokenCache
+    cache: TokenCache | None
     client: httpx.Client
-    token_header = AI_GATEWAY_TOKEN_HEADER
+    token_header: str | None = AI_GATEWAY_TOKEN_HEADER
 
     def log_message(self, format: str, *args: object) -> None:
         return
@@ -71,8 +72,13 @@ def _handle(self) -> None:
             path=self.path.split("?", 1)[0],
         )
         try:
-            # First attempt with the current token.
-            headers = forwarded_request_headers(self, self.cache.token, self.token_header)
+            # Native gateway discovery uses Claude Code's apiKeyHelper credential
+            # verbatim. Relayed mode still injects its separately refreshed swap
+            # token while preserving the caller's Anthropic subscription OAuth.
+            if self.cache is None or self.token_header is None:
+                headers = passthrough_request_headers(self)
+            else:
+                headers = forwarded_request_headers(self, self.cache.token, self.token_header)
             with self.client.stream(self.command, url, headers=headers, content=body) as resp:
                 log_proxy_diagnostic(
                     "model_discovery_upstream_headers",
@@ -81,7 +87,7 @@ def _handle(self) -> None:
                     status=resp.status_code,
                     elapsed_ms=round((time.monotonic() - started) * 1000),
                 )
-                if resp.status_code not in (401, 403):
+                if resp.status_code not in (401, 403) or self.cache is None:
                     self._relay_response(resp, diagnostic_id=diagnostic_id, started=started)
                     return
                 # Auth rejected. Drain the (small) error body so the pooled
@@ -93,6 +99,7 @@ def _handle(self) -> None:
             except RuntimeError as exc:
                 # Still retry with the existing token after reporting the failure.
                 log_token_refresh_failure(exc)
+            assert self.token_header is not None
             headers = forwarded_request_headers(self, self.cache.token, self.token_header)
             with self.client.stream(self.command, url, headers=headers, content=body) as resp:
                 log_proxy_diagnostic(
@@ -310,15 +317,19 @@ def start_proxy(
     workspace: str,
     profile: str | None,
     port: int,
-    token_header: str,
+    token_header: str | None,
     force_refresh_near_expiry: bool,
-) -> tuple[ThreadingHTTPServer, TokenCache, httpx.Client]:
-    """Start the Anthropic model discovery proxy and token refresher."""
+) -> tuple[ThreadingHTTPServer, TokenCache | None, httpx.Client]:
+    """Start the Anthropic model discovery proxy and optional token refresher."""
     upstream_base = f"{workspace.rstrip('/')}/ai-gateway/anthropic/"
-    cache = TokenCache(
-        workspace,
-        profile,
-        force_refresh_near_expiry=force_refresh_near_expiry,
+    cache = (
+        TokenCache(
+            workspace,
+            profile,
+            force_refresh_near_expiry=force_refresh_near_expiry,
+        )
+        if token_header is not None
+        else None
     )
     client = httpx.Client(base_url=upstream_base, timeout=UPSTREAM_TIMEOUT, follow_redirects=False)
     handler = cast(
@@ -339,6 +350,7 @@ def start_proxy(
     except OSError:
         server = ThreadingHTTPServer((LOOPBACK_HOST, 0), handler)
 
-    refresher = threading.Thread(target=cache.run_refresher, daemon=True)
-    refresher.start()
+    if cache is not None:
+        refresher = threading.Thread(target=cache.run_refresher, daemon=True)
+        refresher.start()
     return server, cache, client
diff --git a/src/ucode/gateway_proxy.py b/src/ucode/gateway_proxy.py
index e5cd4788..17a39557 100644
--- a/src/ucode/gateway_proxy.py
+++ b/src/ucode/gateway_proxy.py
@@ -1,11 +1,9 @@
-"""Loopback refresh proxy for Claude gateway requests.
+"""Loopback refresh proxy for relayed 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. Native gateway discovery instead carries the Databricks credential in
-`Authorization`. The proxy refreshes the applicable header and streams responses
-back verbatim.
+header. The proxy refreshes that header and streams responses back verbatim.
 
 Security invariants (mirroring `databricks.py` token handling):
   - Binds 127.0.0.1 only; never exposed off-host.
@@ -33,7 +31,6 @@
 # Header we overwrite with the freshly-minted Databricks credential. Any
 # client-supplied value is replaced, so a stale settings.json value can't leak.
 AI_GATEWAY_TOKEN_HEADER = "X-Databricks-AI-Gateway-Token"
-AUTHORIZATION_HEADER = "Authorization"
 # Hop-by-hop headers must not be forwarded across a proxy.
 HOP_BY_HOP_HEADERS = frozenset(
     h.lower()
@@ -182,14 +179,24 @@ def stop(self) -> None:
         self._stop.set()
 
 
+def passthrough_request_headers(handler: BaseHTTPRequestHandler) -> dict[str, str]:
+    """Copy end-to-end client headers without forwarding hop-by-hop framing."""
+    return {
+        key: value
+        for key, value in handler.headers.items()
+        if key.lower() not in HOP_BY_HOP_HEADERS
+    }
+
+
 def forwarded_request_headers(
     handler: BaseHTTPRequestHandler,
     token: str,
     token_header: str = AI_GATEWAY_TOKEN_HEADER,
 ) -> dict[str, str]:
-    strip_on_forward = HOP_BY_HOP_HEADERS | {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 passthrough_request_headers(handler).items()
+        if key.lower() != token_header.lower()
     }
     headers[token_header] = f"Bearer {token}"
     return headers
diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py
index 832c1f51..d7b93172 100644
--- a/tests/test_agent_claude.py
+++ b/tests/test_agent_claude.py
@@ -26,6 +26,34 @@ def test_display(self):
         assert claude.SPEC["display"] == "Claude Code"
 
 
+class TestMinimumVersion:
+    @pytest.mark.parametrize("version", ["2.1.248", "2.1.250", "3.0.0"])
+    def test_supported_version(self, monkeypatch, version):
+        monkeypatch.setattr(claude, "agent_version", lambda _binary: version)
+
+        assert claude.minimum_version_error() is None
+        assert claude.required_update_message() is None
+
+    def test_older_version_requires_update(self, monkeypatch):
+        monkeypatch.setattr(claude, "agent_version", lambda _binary: "2.1.247 (Claude Code)")
+
+        assert claude.minimum_version_error() == (
+            "Claude Code 2.1.247 (Claude Code) is too old for gateway model discovery. "
+            "Claude Code must be updated to 2.1.248 or newer; run "
+            "`npm install -g @anthropic-ai/claude-code` or `ucode configure`."
+        )
+        assert claude.required_update_message() == (
+            "Claude Code 2.1.247 (Claude Code) is older than required 2.1.248; "
+            "updating Claude Code is required for gateway model discovery."
+        )
+
+    def test_unknown_version_does_not_block(self, monkeypatch):
+        monkeypatch.setattr(claude, "agent_version", lambda _binary: "unknown")
+
+        assert claude.minimum_version_error() is None
+        assert claude.required_update_message() is None
+
+
 class TestRenderOverlay:
     def test_does_not_set_anthropic_model_env(self):
         # We deliberately don't pin ANTHROPIC_MODEL: when set, Claude Code's
@@ -793,12 +821,6 @@ def serve_forever(self):
             def shutdown(self):
                 calls.append(("shutdown",))
 
-        class Cache:
-            token = "fresh-token"
-
-            def stop(self):
-                calls.append(("stop",))
-
         class Client:
             def close(self):
                 calls.append(("close",))
@@ -821,13 +843,14 @@ def start_proxy(workspace, profile, port, token_header, force_refresh_near_expir
                     force_refresh_near_expiry,
                 )
             )
-            return Server(), Cache(), Client()
+            return Server(), None, Client()
 
         monkeypatch.setenv(claude.GATEWAY_MODEL_DISCOVERY_ENV_VAR, "1")
         monkeypatch.delenv("OAUTH_TOKEN", raising=False)
         monkeypatch.delenv("ANTHROPIC_AUTH_TOKEN", raising=False)
         monkeypatch.delenv("ANTHROPIC_BASE_URL", raising=False)
         monkeypatch.delenv("CLAUDE_CODE_USE_GATEWAY", raising=False)
+        monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "fresh-token")
         monkeypatch.setattr(
             claude,
             "start_anthropic_model_discovery_proxy",
@@ -840,11 +863,11 @@ def start_proxy(workspace, profile, port, token_header, force_refresh_near_expir
 
         assert exc.value.code == 0
         assert os.environ["OAUTH_TOKEN"] == "fresh-token"
-        assert os.environ["ANTHROPIC_AUTH_TOKEN"] == "fresh-token"
+        assert "ANTHROPIC_AUTH_TOKEN" not in os.environ
         assert os.environ["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:12345"
         assert os.environ["CLAUDE_CODE_USE_GATEWAY"] == "1"
         assert calls[:2] == [
-            ("proxy", WS, "test", 0, claude.AUTHORIZATION_HEADER, True),
+            ("proxy", WS, "test", 0, None, False),
             ("serve",),
         ]
         assert calls[2][0] == "popen"
@@ -852,11 +875,7 @@ def start_proxy(workspace, profile, port, token_header, force_refresh_near_expir
         assert argv[:2] == ["claude", "--settings"]
         assert json.loads(argv[2])["env"]["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:12345"
         assert argv[3:] == ["--debug"]
-        assert calls[3:] == [
-            ("stop",),
-            ("shutdown",),
-            ("close",),
-        ]
+        assert calls[3:] == [("shutdown",), ("close",)]
 
     def test_smart_routing_uses_anthropic_proxy(self, monkeypatch):
         calls: list[tuple] = []
@@ -871,12 +890,6 @@ def serve_forever(self):
             def shutdown(self):
                 calls.append(("shutdown",))
 
-        class Cache:
-            token = "fresh-token"
-
-            def stop(self):
-                calls.append(("stop",))
-
         class Client:
             def close(self):
                 calls.append(("close",))
@@ -885,13 +898,15 @@ def start_proxy(workspace, profile, port, token_header, force_refresh_near_expir
             calls.append(
                 ("proxy", workspace, profile, port, token_header, force_refresh_near_expiry)
             )
-            return Server(), Cache(), Client()
+            return Server(), None, Client()
 
         def launch_v2(state, tool_args, **kwargs):
             captured["settings"] = kwargs["compose_settings"](["--debug"])
             raise SystemExit(0)
 
         monkeypatch.setenv(v2.ENV_VAR, "1")
+        monkeypatch.delenv("ANTHROPIC_AUTH_TOKEN", raising=False)
+        monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "fresh-token")
         monkeypatch.setattr(claude, "start_anthropic_model_discovery_proxy", start_proxy)
         monkeypatch.setattr(
             claude,
@@ -904,14 +919,15 @@ def launch_v2(state, tool_args, **kwargs):
             claude.launch({"workspace": WS, "profile": "test"}, ["--debug"])
 
         assert exc.value.code == 0
+        assert "ANTHROPIC_AUTH_TOKEN" not in os.environ
         assert calls[:2] == [
-            ("proxy", WS, "test", 0, claude.AUTHORIZATION_HEADER, True),
+            ("proxy", WS, "test", 0, None, False),
             ("serve",),
         ]
         settings, remaining = captured["settings"]
         assert settings["env"]["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:12345"
         assert remaining == ["--debug"]
-        assert calls[2:] == [("stop",), ("shutdown",), ("close",)]
+        assert calls[2:] == [("shutdown",), ("close",)]
 
 
 class TestWriteToolConfigPrunesStaleModelEnv:
diff --git a/tests/test_anthropic_model_discovery_proxy.py b/tests/test_anthropic_model_discovery_proxy.py
index d44c782a..0b40e1ad 100644
--- a/tests/test_anthropic_model_discovery_proxy.py
+++ b/tests/test_anthropic_model_discovery_proxy.py
@@ -152,6 +152,22 @@ def test_leaves_malformed_discovery_response_unchanged(self):
 
 
 class TestAnthropicModelDiscoveryHandler:
+    def test_forwards_api_key_helper_credential_without_refresh_auth(self):
+        out = _Collect()
+        handler = _handler(out)
+        handler.headers = {"X-Api-Key": "api-key-helper-token"}
+        handler.rfile = io.BytesIO()
+        handler.cache = None
+        handler.token_header = None
+        handler.client = _FakeClient(_FakeResponse(200, {}, b'{"data":[]}'))
+
+        handler._handle()
+
+        _method, _url, headers, _body = handler.client.request
+        assert headers["X-Api-Key"] == "api-key-helper-token"
+        assert "Authorization" not in headers
+        assert "X-Databricks-AI-Gateway-Token" not in headers
+
     def test_inherits_relayed_auth_and_prefixes_models(self):
         out = _Collect()
         handler = _handler(out)
@@ -254,3 +270,21 @@ def run_refresher(self):
     finally:
         server.server_close()
         client.close()
+
+
+def test_start_proxy_skips_token_cache_without_refresh_header(monkeypatch):
+    def unexpected_cache(*_args, **_kwargs):
+        raise AssertionError("passthrough proxy must not create a token cache")
+
+    monkeypatch.setattr(anthropic_model_discovery_proxy, "TokenCache", unexpected_cache)
+
+    server, cache, client = anthropic_model_discovery_proxy.start_proxy(
+        "https://workspace.example.com", "profile", 0, None, False
+    )
+    try:
+        assert cache is None
+        assert server.RequestHandlerClass.cache is None
+        assert server.RequestHandlerClass.token_header is None
+    finally:
+        server.server_close()
+        client.close()
diff --git a/tests/test_gateway_proxy.py b/tests/test_gateway_proxy.py
index 28600045..898343aa 100644
--- a/tests/test_gateway_proxy.py
+++ b/tests/test_gateway_proxy.py
@@ -47,13 +47,6 @@ def test_overwrites_client_supplied_swap_header(self):
         out = gateway_proxy.forwarded_request_headers(handler, "fresh")
         assert out["X-Databricks-AI-Gateway-Token"] == "Bearer fresh"
 
-    def test_overwrites_authorization_header(self):
-        handler = _FakeHandler({"Authorization": "Bearer stale"})
-        out = gateway_proxy.forwarded_request_headers(
-            handler, "fresh", gateway_proxy.AUTHORIZATION_HEADER
-        )
-        assert out["Authorization"] == "Bearer fresh"
-
     def test_strips_hop_by_hop_headers(self):
         handler = _FakeHandler(
             {"Host": "localhost:9", "Content-Length": "5", "Connection": "keep-alive"}

From 85989b5abc5efbefea5a26023012030a268fd26f Mon Sep 17 00:00:00 2001
From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com>
Date: Fri, 28 Aug 2026 19:42:59 +0000
Subject: [PATCH 2/4] Keep API key passthrough local to discovery proxy

---
 src/ucode/anthropic_model_discovery_proxy.py |  7 +++++--
 src/ucode/gateway_proxy.py                   | 21 +++++++-------------
 tests/test_gateway_proxy.py                  |  7 +++++++
 3 files changed, 19 insertions(+), 16 deletions(-)

diff --git a/src/ucode/anthropic_model_discovery_proxy.py b/src/ucode/anthropic_model_discovery_proxy.py
index 67174a34..3b6fce29 100644
--- a/src/ucode/anthropic_model_discovery_proxy.py
+++ b/src/ucode/anthropic_model_discovery_proxy.py
@@ -32,7 +32,6 @@
     forwarded_request_headers,
     log_proxy_diagnostic,
     log_token_refresh_failure,
-    passthrough_request_headers,
 )
 
 
@@ -76,7 +75,11 @@ def _handle(self) -> None:
             # verbatim. Relayed mode still injects its separately refreshed swap
             # token while preserving the caller's Anthropic subscription OAuth.
             if self.cache is None or self.token_header is None:
-                headers = passthrough_request_headers(self)
+                headers = {
+                    key: value
+                    for key, value in self.headers.items()
+                    if key.lower() not in HOP_BY_HOP_HEADERS
+                }
             else:
                 headers = forwarded_request_headers(self, self.cache.token, self.token_header)
             with self.client.stream(self.command, url, headers=headers, content=body) as resp:
diff --git a/src/ucode/gateway_proxy.py b/src/ucode/gateway_proxy.py
index 17a39557..e5cd4788 100644
--- a/src/ucode/gateway_proxy.py
+++ b/src/ucode/gateway_proxy.py
@@ -1,9 +1,11 @@
-"""Loopback refresh proxy for relayed Claude gateway requests.
+"""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. The proxy refreshes that header and streams responses 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.
@@ -31,6 +33,7 @@
 # Header we overwrite with the freshly-minted Databricks credential. Any
 # client-supplied value is replaced, so a stale settings.json value can't leak.
 AI_GATEWAY_TOKEN_HEADER = "X-Databricks-AI-Gateway-Token"
+AUTHORIZATION_HEADER = "Authorization"
 # Hop-by-hop headers must not be forwarded across a proxy.
 HOP_BY_HOP_HEADERS = frozenset(
     h.lower()
@@ -179,24 +182,14 @@ def stop(self) -> None:
         self._stop.set()
 
 
-def passthrough_request_headers(handler: BaseHTTPRequestHandler) -> dict[str, str]:
-    """Copy end-to-end client headers without forwarding hop-by-hop framing."""
-    return {
-        key: value
-        for key, value in handler.headers.items()
-        if key.lower() not in HOP_BY_HOP_HEADERS
-    }
-
-
 def forwarded_request_headers(
     handler: BaseHTTPRequestHandler,
     token: str,
     token_header: str = AI_GATEWAY_TOKEN_HEADER,
 ) -> dict[str, str]:
+    strip_on_forward = HOP_BY_HOP_HEADERS | {token_header.lower()}
     headers = {
-        key: value
-        for key, value in passthrough_request_headers(handler).items()
-        if key.lower() != token_header.lower()
+        key: value for key, value in handler.headers.items() if key.lower() not in strip_on_forward
     }
     headers[token_header] = f"Bearer {token}"
     return headers
diff --git a/tests/test_gateway_proxy.py b/tests/test_gateway_proxy.py
index 898343aa..28600045 100644
--- a/tests/test_gateway_proxy.py
+++ b/tests/test_gateway_proxy.py
@@ -47,6 +47,13 @@ def test_overwrites_client_supplied_swap_header(self):
         out = gateway_proxy.forwarded_request_headers(handler, "fresh")
         assert out["X-Databricks-AI-Gateway-Token"] == "Bearer fresh"
 
+    def test_overwrites_authorization_header(self):
+        handler = _FakeHandler({"Authorization": "Bearer stale"})
+        out = gateway_proxy.forwarded_request_headers(
+            handler, "fresh", gateway_proxy.AUTHORIZATION_HEADER
+        )
+        assert out["Authorization"] == "Bearer fresh"
+
     def test_strips_hop_by_hop_headers(self):
         handler = _FakeHandler(
             {"Host": "localhost:9", "Content-Length": "5", "Connection": "keep-alive"}

From 0e745ab5af1fffb07032567a7f29b02f93892439 Mon Sep 17 00:00:00 2001
From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com>
Date: Fri, 28 Aug 2026 19:48:41 +0000
Subject: [PATCH 3/4] Minimize Claude gateway auth cleanup

---
 src/ucode/agents/claude.py | 15 ++++++---------
 tests/test_agent_claude.py |  1 -
 2 files changed, 6 insertions(+), 10 deletions(-)

diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py
index d7102e27..a540fb72 100644
--- a/src/ucode/agents/claude.py
+++ b/src/ucode/agents/claude.py
@@ -1171,8 +1171,7 @@ def _launch_relayed(state: dict, binary: str, tool_args: list[str]) -> None:
         token_header=AI_GATEWAY_TOKEN_HEADER,
         force_refresh_near_expiry=False,
     )
-    if cache is None:  # defensive: relayed mode always requests swap-token refresh
-        raise RuntimeError("Relayed proxy did not initialize its token cache.")
+    assert cache is not None
     # 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.
@@ -1196,12 +1195,12 @@ def _launch_relayed(state: dict, binary: str, tool_args: list[str]) -> None:
     raise SystemExit(returncode)
 
 
-def _launch_claude_with_model_discovery_proxy(
+def _launch_claude_with_gateway_proxy(
     state: dict, binary: str, tool_args: list[str], *, smart_routing: bool
 ) -> None:
-    """Launch Claude through the model-alias proxy using apiKeyHelper auth."""
+    """Launch Claude through the gateway model-alias proxy."""
     workspace = state["workspace"]
-    server, cache, client = start_anthropic_model_discovery_proxy(
+    server, _cache, client = start_anthropic_model_discovery_proxy(
         workspace,
         state.get("profile"),
         0,
@@ -1243,8 +1242,6 @@ def compose_gateway_settings(args: list[str]) -> tuple[dict, list[str]]:
             proc.send_signal(signal.SIGINT)
             returncode = proc.wait()
     finally:
-        if cache is not None:
-            cache.stop()
         server.shutdown()
         client.close()
     raise SystemExit(returncode)
@@ -1271,10 +1268,10 @@ def launch(state: dict, tool_args: list[str]) -> None:
             "Please use Codex or disable smart routing."
         )
     if first_prompt_routing:
-        _launch_claude_with_model_discovery_proxy(state, binary, tool_args, smart_routing=True)
+        _launch_claude_with_gateway_proxy(state, binary, tool_args, smart_routing=True)
         return
     if workspace and os.environ.get(GATEWAY_MODEL_DISCOVERY_ENV_VAR) == "1":
-        _launch_claude_with_model_discovery_proxy(state, binary, tool_args, smart_routing=False)
+        _launch_claude_with_gateway_proxy(state, binary, tool_args, smart_routing=False)
         return
     if workspace:
         os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile"))
diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py
index d7b93172..f1a316d8 100644
--- a/tests/test_agent_claude.py
+++ b/tests/test_agent_claude.py
@@ -906,7 +906,6 @@ def launch_v2(state, tool_args, **kwargs):
 
         monkeypatch.setenv(v2.ENV_VAR, "1")
         monkeypatch.delenv("ANTHROPIC_AUTH_TOKEN", raising=False)
-        monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "fresh-token")
         monkeypatch.setattr(claude, "start_anthropic_model_discovery_proxy", start_proxy)
         monkeypatch.setattr(
             claude,

From 04217af2522f63d4d48447ec52de0d86e0b42907 Mon Sep 17 00:00:00 2001
From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com>
Date: Fri, 28 Aug 2026 23:51:55 +0000
Subject: [PATCH 4/4] Decouple Claude model discovery from relayed auth

---
 src/ucode/agents/claude.py                    |  17 +--
 src/ucode/anthropic_model_discovery_proxy.py  |  74 ++----------
 tests/test_agent_claude.py                    | 110 ++++++++++++++----
 tests/test_anthropic_model_discovery_proxy.py |  66 ++---------
 4 files changed, 116 insertions(+), 151 deletions(-)

diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py
index a540fb72..9247fe52 100644
--- a/src/ucode/agents/claude.py
+++ b/src/ucode/agents/claude.py
@@ -15,6 +15,7 @@
 from pathlib import Path
 from typing import cast
 
+from ucode import gateway_proxy
 from ucode.agent_updates import available_npm_package_update
 from ucode.anthropic_model_discovery_proxy import (
     start_proxy as start_anthropic_model_discovery_proxy,
@@ -33,7 +34,6 @@
     build_tool_base_url,
     get_databricks_token,
 )
-from ucode.gateway_proxy import AI_GATEWAY_TOKEN_HEADER
 from ucode.launcher import exec_or_spawn
 from ucode.managed_files import OS, current_os, write_managed_file
 from ucode.smart_routing import v2 as smart_routing_v2
@@ -1164,14 +1164,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_anthropic_model_discovery_proxy(
+    server, cache, client = gateway_proxy.start_proxy(
         workspace,
         state.get("profile"),
         port,
-        token_header=AI_GATEWAY_TOKEN_HEADER,
+        token_header=gateway_proxy.AI_GATEWAY_TOKEN_HEADER,
         force_refresh_near_expiry=False,
     )
-    assert cache is not None
     # 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.
@@ -1200,15 +1199,7 @@ def _launch_claude_with_gateway_proxy(
 ) -> None:
     """Launch Claude through the gateway model-alias proxy."""
     workspace = state["workspace"]
-    server, _cache, client = start_anthropic_model_discovery_proxy(
-        workspace,
-        state.get("profile"),
-        0,
-        token_header=None,
-        force_refresh_near_expiry=False,
-    )
-    if not smart_routing:
-        os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile"))
+    server, client = start_anthropic_model_discovery_proxy(workspace, 0)
     os.environ["ANTHROPIC_BASE_URL"] = f"http://{LOOPBACK_HOST}:{server.server_address[1]}"
     os.environ["CLAUDE_CODE_USE_GATEWAY"] = "1"
 
diff --git a/src/ucode/anthropic_model_discovery_proxy.py b/src/ucode/anthropic_model_discovery_proxy.py
index 3b6fce29..5c70ef25 100644
--- a/src/ucode/anthropic_model_discovery_proxy.py
+++ b/src/ucode/anthropic_model_discovery_proxy.py
@@ -1,12 +1,11 @@
 """Loopback proxy for Claude gateway model discovery.
 
-The proxy refreshes the Databricks credential, streams inference responses
-verbatim, and rewrites model discovery responses when needed.
+The proxy forwards Claude Code's apiKeyHelper credential, streams inference
+responses verbatim, and rewrites model discovery responses when needed.
 
 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
-    and is refreshed off the request path.
+  - Never logs header values or bodies.
 """
 
 from __future__ import annotations
@@ -25,21 +24,15 @@
 
 from ucode.constants import LOOPBACK_HOST
 from ucode.gateway_proxy import (
-    AI_GATEWAY_TOKEN_HEADER,
     HOP_BY_HOP_HEADERS,
     UPSTREAM_TIMEOUT,
-    TokenCache,
-    forwarded_request_headers,
     log_proxy_diagnostic,
-    log_token_refresh_failure,
 )
 
 
 class _ProxyHandler(BaseHTTPRequestHandler):
     # Set by the server factory.
-    cache: TokenCache | None
     client: httpx.Client
-    token_header: str | None = AI_GATEWAY_TOKEN_HEADER
 
     def log_message(self, format: str, *args: object) -> None:
         return
@@ -71,17 +64,11 @@ def _handle(self) -> None:
             path=self.path.split("?", 1)[0],
         )
         try:
-            # Native gateway discovery uses Claude Code's apiKeyHelper credential
-            # verbatim. Relayed mode still injects its separately refreshed swap
-            # token while preserving the caller's Anthropic subscription OAuth.
-            if self.cache is None or self.token_header is None:
-                headers = {
-                    key: value
-                    for key, value in self.headers.items()
-                    if key.lower() not in HOP_BY_HOP_HEADERS
-                }
-            else:
-                headers = forwarded_request_headers(self, self.cache.token, self.token_header)
+            headers = {
+                key: value
+                for key, value in self.headers.items()
+                if key.lower() not in HOP_BY_HOP_HEADERS
+            }
             with self.client.stream(self.command, url, headers=headers, content=body) as resp:
                 log_proxy_diagnostic(
                     "model_discovery_upstream_headers",
@@ -90,28 +77,6 @@ def _handle(self) -> None:
                     status=resp.status_code,
                     elapsed_ms=round((time.monotonic() - started) * 1000),
                 )
-                if resp.status_code not in (401, 403) or self.cache is None:
-                    self._relay_response(resp, diagnostic_id=diagnostic_id, started=started)
-                    return
-                # Auth rejected. Drain the (small) error body so the pooled
-                # connection can be reused, then fall through to one retry.
-                resp.read()
-            # Force-refresh the Databricks token and retry once.
-            try:
-                self.cache.refresh()
-            except RuntimeError as exc:
-                # Still retry with the existing token after reporting the failure.
-                log_token_refresh_failure(exc)
-            assert self.token_header is not None
-            headers = forwarded_request_headers(self, self.cache.token, self.token_header)
-            with self.client.stream(self.command, url, headers=headers, content=body) as resp:
-                log_proxy_diagnostic(
-                    "model_discovery_upstream_headers",
-                    request_id=diagnostic_id,
-                    attempt=2,
-                    status=resp.status_code,
-                    elapsed_ms=round((time.monotonic() - started) * 1000),
-                )
                 self._relay_response(resp, diagnostic_id=diagnostic_id, started=started)
         except (BrokenPipeError, ConnectionResetError):
             # Client closed before/while we relayed headers — routine on cancel.
@@ -318,22 +283,10 @@ def _response_chunks(self, resp: httpx.Response) -> tuple[Iterable[bytes], froze
 
 def start_proxy(
     workspace: str,
-    profile: str | None,
     port: int,
-    token_header: str | None,
-    force_refresh_near_expiry: bool,
-) -> tuple[ThreadingHTTPServer, TokenCache | None, httpx.Client]:
-    """Start the Anthropic model discovery proxy and optional token refresher."""
+) -> tuple[ThreadingHTTPServer, httpx.Client]:
+    """Start the Anthropic model discovery proxy."""
     upstream_base = f"{workspace.rstrip('/')}/ai-gateway/anthropic/"
-    cache = (
-        TokenCache(
-            workspace,
-            profile,
-            force_refresh_near_expiry=force_refresh_near_expiry,
-        )
-        if token_header is not None
-        else None
-    )
     client = httpx.Client(base_url=upstream_base, timeout=UPSTREAM_TIMEOUT, follow_redirects=False)
     handler = cast(
         type[BaseHTTPRequestHandler],
@@ -341,9 +294,7 @@ def start_proxy(
             "BoundProxyHandler",
             (_AnthropicModelDiscoveryHandler,),
             {
-                "cache": cache,
                 "client": client,
-                "token_header": token_header,
                 "anthropic_model_aliases": _AnthropicModelAliases(),
             },
         ),
@@ -353,7 +304,4 @@ def start_proxy(
     except OSError:
         server = ThreadingHTTPServer((LOOPBACK_HOST, 0), handler)
 
-    if cache is not None:
-        refresher = threading.Thread(target=cache.run_refresher, daemon=True)
-        refresher.start()
-    return server, cache, client
+    return server, client
diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py
index f1a316d8..1fb4806f 100644
--- a/tests/test_agent_claude.py
+++ b/tests/test_agent_claude.py
@@ -705,6 +705,79 @@ def boom(name, entry, scope=mcp_mod.MCP_USER_SCOPE):
 
 
 class TestClaudeLaunch:
+    def test_relayed_launch_uses_refresh_proxy_not_discovery_proxy(self, monkeypatch):
+        calls: list[tuple] = []
+
+        class Server:
+            server_address = ("127.0.0.1", 12345)
+
+            def serve_forever(self):
+                calls.append(("serve",))
+
+            def shutdown(self):
+                calls.append(("shutdown",))
+
+        class Cache:
+            def stop(self):
+                calls.append(("stop",))
+
+        class Client:
+            def close(self):
+                calls.append(("close",))
+
+        class Process:
+            def __init__(self, argv):
+                calls.append(("popen", argv))
+
+            def wait(self):
+                return 0
+
+        def start_proxy(workspace, profile, port, token_header, force_refresh_near_expiry):
+            calls.append(
+                (
+                    "proxy",
+                    workspace,
+                    profile,
+                    port,
+                    token_header,
+                    force_refresh_near_expiry,
+                )
+            )
+            return Server(), Cache(), Client()
+
+        monkeypatch.setattr(claude, "_managed_relayed_conflicts", lambda: None)
+        monkeypatch.setattr(claude, "_managed_pinned_model", lambda: None)
+        monkeypatch.setattr(claude, "_ensure_subscription_login", lambda: None)
+        monkeypatch.setattr(claude.gateway_proxy, "start_proxy", start_proxy)
+        monkeypatch.setattr(
+            claude,
+            "start_anthropic_model_discovery_proxy",
+            lambda *_args: pytest.fail("relayed auth must not use the discovery proxy"),
+        )
+        monkeypatch.setattr(claude.subprocess, "Popen", Process)
+
+        with pytest.raises(SystemExit) as exc:
+            claude.launch(
+                {
+                    "workspace": WS,
+                    "profile": "test",
+                    "claude_relayed": True,
+                    "relayed_proxy_port": 12345,
+                },
+                ["--debug"],
+            )
+
+        assert exc.value.code == 0
+        assert calls[0] == (
+            "proxy",
+            WS,
+            "test",
+            12345,
+            claude.gateway_proxy.AI_GATEWAY_TOKEN_HEADER,
+            False,
+        )
+        assert calls[-3:] == [("stop",), ("shutdown",), ("close",)]
+
     def test_smart_routing_on_windows_is_not_supported(self, monkeypatch):
         monkeypatch.setenv(v2.ENV_VAR, "1")
         monkeypatch.setattr(claude.os, "name", "nt")
@@ -832,25 +905,20 @@ def __init__(self, argv):
             def wait(self):
                 return 0
 
-        def start_proxy(workspace, profile, port, token_header, force_refresh_near_expiry):
-            calls.append(
-                (
-                    "proxy",
-                    workspace,
-                    profile,
-                    port,
-                    token_header,
-                    force_refresh_near_expiry,
-                )
-            )
-            return Server(), None, Client()
+        def start_proxy(workspace, port):
+            calls.append(("proxy", workspace, port))
+            return Server(), Client()
 
         monkeypatch.setenv(claude.GATEWAY_MODEL_DISCOVERY_ENV_VAR, "1")
         monkeypatch.delenv("OAUTH_TOKEN", raising=False)
         monkeypatch.delenv("ANTHROPIC_AUTH_TOKEN", raising=False)
         monkeypatch.delenv("ANTHROPIC_BASE_URL", raising=False)
         monkeypatch.delenv("CLAUDE_CODE_USE_GATEWAY", raising=False)
-        monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "fresh-token")
+        monkeypatch.setattr(
+            claude,
+            "get_databricks_token",
+            lambda *_args: pytest.fail("model discovery must rely on apiKeyHelper"),
+        )
         monkeypatch.setattr(
             claude,
             "start_anthropic_model_discovery_proxy",
@@ -862,12 +930,12 @@ def start_proxy(workspace, profile, port, token_header, force_refresh_near_expir
             claude.launch({"workspace": WS, "profile": "test"}, ["--debug"])
 
         assert exc.value.code == 0
-        assert os.environ["OAUTH_TOKEN"] == "fresh-token"
+        assert "OAUTH_TOKEN" not in os.environ
         assert "ANTHROPIC_AUTH_TOKEN" not in os.environ
         assert os.environ["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:12345"
         assert os.environ["CLAUDE_CODE_USE_GATEWAY"] == "1"
         assert calls[:2] == [
-            ("proxy", WS, "test", 0, None, False),
+            ("proxy", WS, 0),
             ("serve",),
         ]
         assert calls[2][0] == "popen"
@@ -894,17 +962,16 @@ class Client:
             def close(self):
                 calls.append(("close",))
 
-        def start_proxy(workspace, profile, port, token_header, force_refresh_near_expiry):
-            calls.append(
-                ("proxy", workspace, profile, port, token_header, force_refresh_near_expiry)
-            )
-            return Server(), None, Client()
+        def start_proxy(workspace, port):
+            calls.append(("proxy", workspace, port))
+            return Server(), Client()
 
         def launch_v2(state, tool_args, **kwargs):
             captured["settings"] = kwargs["compose_settings"](["--debug"])
             raise SystemExit(0)
 
         monkeypatch.setenv(v2.ENV_VAR, "1")
+        monkeypatch.delenv("OAUTH_TOKEN", raising=False)
         monkeypatch.delenv("ANTHROPIC_AUTH_TOKEN", raising=False)
         monkeypatch.setattr(claude, "start_anthropic_model_discovery_proxy", start_proxy)
         monkeypatch.setattr(
@@ -918,9 +985,10 @@ def launch_v2(state, tool_args, **kwargs):
             claude.launch({"workspace": WS, "profile": "test"}, ["--debug"])
 
         assert exc.value.code == 0
+        assert "OAUTH_TOKEN" not in os.environ
         assert "ANTHROPIC_AUTH_TOKEN" not in os.environ
         assert calls[:2] == [
-            ("proxy", WS, "test", 0, None, False),
+            ("proxy", WS, 0),
             ("serve",),
         ]
         settings, remaining = captured["settings"]
diff --git a/tests/test_anthropic_model_discovery_proxy.py b/tests/test_anthropic_model_discovery_proxy.py
index 0b40e1ad..b131b6a2 100644
--- a/tests/test_anthropic_model_discovery_proxy.py
+++ b/tests/test_anthropic_model_discovery_proxy.py
@@ -41,13 +41,6 @@ def stream(self, method, url, headers, content):
         return self.response
 
 
-class _FakeCache:
-    token = "databricks-token"
-
-    def refresh(self):
-        return None
-
-
 class _Collect(io.RawIOBase):
     def __init__(self):
         self.data = bytearray()
@@ -157,8 +150,6 @@ def test_forwards_api_key_helper_credential_without_refresh_auth(self):
         handler = _handler(out)
         handler.headers = {"X-Api-Key": "api-key-helper-token"}
         handler.rfile = io.BytesIO()
-        handler.cache = None
-        handler.token_header = None
         handler.client = _FakeClient(_FakeResponse(200, {}, b'{"data":[]}'))
 
         handler._handle()
@@ -168,20 +159,22 @@ def test_forwards_api_key_helper_credential_without_refresh_auth(self):
         assert "Authorization" not in headers
         assert "X-Databricks-AI-Gateway-Token" not in headers
 
-    def test_inherits_relayed_auth_and_prefixes_models(self):
+    def test_prefixes_models_and_strips_hop_by_hop_headers(self):
         out = _Collect()
         handler = _handler(out)
-        handler.headers = {"Authorization": "Bearer subscription-token"}
+        handler.headers = {
+            "X-Api-Key": "api-key-helper-token",
+            "Connection": "keep-alive",
+            "Host": "127.0.0.1",
+        }
         handler.rfile = io.BytesIO()
-        handler.cache = _FakeCache()
         handler.client = _FakeClient(_FakeResponse(200, {}, b'{"data":[{"id":"custom-model"}]}'))
 
         handler._handle()
 
         method, url, headers, body = handler.client.request
         assert (method, url, body) == ("GET", "v1/models", None)
-        assert headers["Authorization"] == "Bearer subscription-token"
-        assert headers["X-Databricks-AI-Gateway-Token"] == "Bearer databricks-token"
+        assert headers == {"X-Api-Key": "api-key-helper-token"}
         assert b"anthropic-aigw-custom-model" in bytes(out.data)
 
     def test_prefixes_successful_model_response_and_drops_content_encoding(self):
@@ -210,20 +203,18 @@ def test_keeps_content_encoding_for_unchanged_error(self):
         assert response.read_calls == 0
         assert response.iter_raw_calls == 1
 
-    def test_streams_relayed_inference_response_without_buffering(self):
+    def test_streams_inference_response_without_buffering(self):
         out = _Collect()
         handler = _handler(out, path="/v1/messages", command="POST")
-        handler.headers = {"Authorization": "Bearer subscription-token", "Content-Length": "2"}
+        handler.headers = {"X-Api-Key": "api-key-helper-token", "Content-Length": "2"}
         handler.rfile = io.BytesIO(b"{}")
-        handler.cache = _FakeCache()
         response = _FakeResponse(200, {"Content-Type": "text/event-stream"}, b"data: event\n\n")
         handler.client = _FakeClient(response)
 
         handler._handle()
 
         _method, _url, headers, _body = handler.client.request
-        assert headers["Authorization"] == "Bearer subscription-token"
-        assert headers["X-Databricks-AI-Gateway-Token"] == "Bearer databricks-token"
+        assert headers == {"X-Api-Key": "api-key-helper-token"}
         assert response.read_calls == 0
         assert response.iter_raw_calls == 1
         assert b"data: event\n\n" in bytes(out.data)
@@ -240,51 +231,18 @@ def test_strips_known_alias_from_message_request(self):
         assert json.loads(body) == {"model": "catalog.schema.custom"}
 
 
-def test_start_proxy_uses_discovery_handler(monkeypatch):
-    class _StubCache:
-        def run_refresher(self):
-            return None
-
-    cache = _StubCache()
-    monkeypatch.setattr(
-        anthropic_model_discovery_proxy,
-        "TokenCache",
-        lambda *_args, **_kwargs: cache,
-    )
-
-    server, actual_cache, client = anthropic_model_discovery_proxy.start_proxy(
-        "https://workspace.example.com", "profile", 0, "header", False
-    )
+def test_start_proxy_uses_discovery_handler():
+    server, client = anthropic_model_discovery_proxy.start_proxy("https://workspace.example.com", 0)
     try:
         handler = server.RequestHandlerClass
         assert issubclass(
             handler,
             anthropic_model_discovery_proxy._AnthropicModelDiscoveryHandler,
         )
-        assert handler.cache is cache
         assert isinstance(
             handler.anthropic_model_aliases,
             anthropic_model_discovery_proxy._AnthropicModelAliases,
         )
-        assert actual_cache is cache
-    finally:
-        server.server_close()
-        client.close()
-
-
-def test_start_proxy_skips_token_cache_without_refresh_header(monkeypatch):
-    def unexpected_cache(*_args, **_kwargs):
-        raise AssertionError("passthrough proxy must not create a token cache")
-
-    monkeypatch.setattr(anthropic_model_discovery_proxy, "TokenCache", unexpected_cache)
-
-    server, cache, client = anthropic_model_discovery_proxy.start_proxy(
-        "https://workspace.example.com", "profile", 0, None, False
-    )
-    try:
-        assert cache is None
-        assert server.RequestHandlerClass.cache is None
-        assert server.RequestHandlerClass.token_header is None
     finally:
         server.server_close()
         client.close()