From 176995cec27df67145d59178dbe0837f4de5da34 Mon Sep 17 00:00:00 2001
From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com>
Date: Thu, 20 Aug 2026 18:46:14 +0000
Subject: [PATCH 1/9] claude: authenticate gateway model discovery

---
 src/ucode/agents/claude.py | 4 +++-
 tests/test_agent_claude.py | 2 ++
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py
index 5fad5225..58208a21 100644
--- a/src/ucode/agents/claude.py
+++ b/src/ucode/agents/claude.py
@@ -1017,7 +1017,9 @@ def launch(state: dict, tool_args: list[str]) -> None:
         _launch_relayed(state, binary, tool_args)
         return
     if workspace:
-        os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile"))
+        token = get_databricks_token(workspace, state.get("profile"))
+        os.environ["OAUTH_TOKEN"] = token
+        os.environ["ANTHROPIC_AUTH_TOKEN"] = token
     exec_or_spawn(_build_claude_argv(binary, tool_args))
 
 
diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py
index b45b6303..ed494281 100644
--- a/tests/test_agent_claude.py
+++ b/tests/test_agent_claude.py
@@ -642,6 +642,7 @@ def fake_execvp(binary: str, args: list[str]) -> None:
             raise RuntimeError("stop")
 
         monkeypatch.delenv("OAUTH_TOKEN", raising=False)
+        monkeypatch.delenv("ANTHROPIC_AUTH_TOKEN", raising=False)
         monkeypatch.setattr(
             claude, "get_databricks_token", lambda workspace, profile=None: "fresh-token"
         )
@@ -653,6 +654,7 @@ def fake_execvp(binary: str, args: list[str]) -> None:
             assert str(exc) == "stop"
 
         assert os.environ["OAUTH_TOKEN"] == "fresh-token"
+        assert os.environ["ANTHROPIC_AUTH_TOKEN"] == "fresh-token"
         assert exec_calls == [
             (
                 "claude",

From e9526a8fa5a50a5a1d5d3ba9f511b9a04d4a295c Mon Sep 17 00:00:00 2001
From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com>
Date: Thu, 20 Aug 2026 19:46:23 +0000
Subject: [PATCH 2/9] claude: initialize gateway before bootstrap

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

diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py
index 58208a21..cd1dabc0 100644
--- a/src/ucode/agents/claude.py
+++ b/src/ucode/agents/claude.py
@@ -1020,6 +1020,8 @@ def launch(state: dict, tool_args: list[str]) -> None:
         token = get_databricks_token(workspace, state.get("profile"))
         os.environ["OAUTH_TOKEN"] = token
         os.environ["ANTHROPIC_AUTH_TOKEN"] = token
+        os.environ["ANTHROPIC_BASE_URL"] = build_tool_base_url("claude", workspace)
+        os.environ["CLAUDE_CODE_USE_GATEWAY"] = "1"
     exec_or_spawn(_build_claude_argv(binary, tool_args))
 
 
diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py
index ed494281..164aea0c 100644
--- a/tests/test_agent_claude.py
+++ b/tests/test_agent_claude.py
@@ -643,6 +643,8 @@ def fake_execvp(binary: str, args: list[str]) -> None:
 
         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 workspace, profile=None: "fresh-token"
         )
@@ -655,6 +657,8 @@ def fake_execvp(binary: str, args: list[str]) -> None:
 
         assert os.environ["OAUTH_TOKEN"] == "fresh-token"
         assert os.environ["ANTHROPIC_AUTH_TOKEN"] == "fresh-token"
+        assert os.environ["ANTHROPIC_BASE_URL"] == f"{WS}/ai-gateway/anthropic"
+        assert os.environ["CLAUDE_CODE_USE_GATEWAY"] == "1"
         assert exec_calls == [
             (
                 "claude",

From 46890ed21b03373be18c6c8e02bfa3ae6d443689 Mon Sep 17 00:00:00 2001
From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com>
Date: Thu, 20 Aug 2026 21:41:03 +0000
Subject: [PATCH 3/9] Refresh Claude gateway credentials through proxy

---
 src/ucode/agents/claude.py  | 52 ++++++++++++++++++++++-----
 src/ucode/gateway_proxy.py  | 25 +++++++------
 tests/test_agent_claude.py  | 71 +++++++++++++++++++++++++++----------
 tests/test_gateway_proxy.py |  7 ++++
 4 files changed, 118 insertions(+), 37 deletions(-)

diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py
index cd1dabc0..206e85d7 100644
--- a/src/ucode/agents/claude.py
+++ b/src/ucode/agents/claude.py
@@ -27,7 +27,6 @@
 from ucode.databricks import (
     build_auth_shell_command,
     build_tool_base_url,
-    get_databricks_token,
 )
 from ucode.launcher import exec_or_spawn
 from ucode.managed_files import OS, current_os, write_managed_file
@@ -868,7 +867,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.
 
@@ -891,7 +895,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]
@@ -901,6 +905,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,
@@ -1010,6 +1016,39 @@ 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:
+    from ucode.gateway_proxy import AUTHORIZATION_HEADER, start_proxy
+
+    workspace = state["workspace"]
+    server, cache, client = start_proxy(
+        workspace, state.get("profile"), 0, token_header=AUTHORIZATION_HEADER
+    )
+    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")
@@ -1017,11 +1056,8 @@ def launch(state: dict, tool_args: list[str]) -> None:
         _launch_relayed(state, binary, tool_args)
         return
     if workspace:
-        token = get_databricks_token(workspace, state.get("profile"))
-        os.environ["OAUTH_TOKEN"] = token
-        os.environ["ANTHROPIC_AUTH_TOKEN"] = token
-        os.environ["ANTHROPIC_BASE_URL"] = build_tool_base_url("claude", workspace)
-        os.environ["CLAUDE_CODE_USE_GATEWAY"] = "1"
+        _launch_gateway(state, binary, tool_args)
+        return
     exec_or_spawn(_build_claude_argv(binary, tool_args))
 
 
diff --git a/src/ucode/gateway_proxy.py b/src/ucode/gateway_proxy.py
index d11b5f1f..995e6865 100644
--- a/src/ucode/gateway_proxy.py
+++ b/src/ucode/gateway_proxy.py
@@ -34,6 +34,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.
 _SWAP_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 +51,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.
@@ -181,11 +179,14 @@ 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 = _SWAP_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
 
 
@@ -193,6 +194,7 @@ class _ProxyHandler(BaseHTTPRequestHandler):
     # Set by the server factory.
     cache: _TokenCache
     client: httpx.Client
+    token_header = _SWAP_HEADER
 
     def log_message(self, format: str, *args: object) -> None:
         return
@@ -219,7 +221,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",
@@ -244,7 +246,7 @@ def _handle(self) -> None:
                 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)
+            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",
@@ -357,7 +359,10 @@ 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 = _SWAP_HEADER,
 ) -> tuple[ThreadingHTTPServer, _TokenCache, httpx.Client]:
     """Start the loopback refresh proxy + its background token refresher.
 
@@ -379,7 +384,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)
diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py
index 164aea0c..5c64e56e 100644
--- a/tests/test_agent_claude.py
+++ b/tests/test_agent_claude.py
@@ -634,36 +634,69 @@ def boom(name, entry, scope=mcp_mod.MCP_USER_SCOPE):
 
 
 class TestClaudeLaunch:
-    def test_sets_oauth_token_before_exec(self, monkeypatch):
-        exec_calls: list[tuple[str, list[str]]] = []
+    def test_runs_through_refresh_proxy(self, monkeypatch):
+        import ucode.gateway_proxy as gateway_proxy
 
-        def fake_execvp(binary: str, args: list[str]) -> None:
-            exec_calls.append((binary, args))
-            raise RuntimeError("stop")
+        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:
+            token = "fresh-token"
+
+            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):
+            calls.append(("proxy", workspace, profile, port, token_header))
+            return Server(), Cache(), Client()
 
         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 workspace, profile=None: "fresh-token"
-        )
-        monkeypatch.setattr(os, "execvp", fake_execvp)
+        monkeypatch.setattr(gateway_proxy, "start_proxy", start_proxy)
+        monkeypatch.setattr(claude.subprocess, "Popen", Process)
 
-        try:
-            claude.launch({"workspace": WS}, ["--debug"])
-        except RuntimeError as exc:
-            assert str(exc) == "stop"
+        with pytest.raises(SystemExit) as exc:
+            claude.launch({"workspace": WS, "profile": "test"}, ["--debug"])
 
+        assert exc.value.code == 0
         assert os.environ["OAUTH_TOKEN"] == "fresh-token"
         assert os.environ["ANTHROPIC_AUTH_TOKEN"] == "fresh-token"
-        assert os.environ["ANTHROPIC_BASE_URL"] == f"{WS}/ai-gateway/anthropic"
+        assert os.environ["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:12345"
         assert os.environ["CLAUDE_CODE_USE_GATEWAY"] == "1"
-        assert exec_calls == [
-            (
-                "claude",
-                ["claude", "--settings", str(claude.CLAUDE_SETTINGS_PATH), "--debug"],
-            )
+        assert calls[:2] == [
+            ("proxy", WS, "test", 0, gateway_proxy.AUTHORIZATION_HEADER),
+            ("serve",),
+        ]
+        assert calls[2][0] == "popen"
+        argv = calls[2][1]
+        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",),
         ]
 
 
diff --git a/tests/test_gateway_proxy.py b/tests/test_gateway_proxy.py
index f103aca2..c438bc7e 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 8459d69e24a52eddfdae54377ce0cb9e91ab3516 Mon Sep 17 00:00:00 2001
From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com>
Date: Thu, 20 Aug 2026 23:53:04 +0000
Subject: [PATCH 4/9] Force refresh expiring gateway tokens

---
 src/ucode/gateway_proxy.py  | 14 ++++++--------
 tests/test_gateway_proxy.py | 13 ++++++-------
 2 files changed, 12 insertions(+), 15 deletions(-)

diff --git a/src/ucode/gateway_proxy.py b/src/ucode/gateway_proxy.py
index 995e6865..1298f449 100644
--- a/src/ucode/gateway_proxy.py
+++ b/src/ucode/gateway_proxy.py
@@ -126,13 +126,11 @@ def __init__(self, workspace: str, profile: str | None) -> None:
         # 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)
+        self._refresh()
 
-    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)
+    def _refresh(self) -> None:
+        """Force-mint a token and record its expiry."""
+        token = get_databricks_token(self._workspace, self._profile, force_refresh=True)
         expiry = _jwt_exp(token) or (time.time() + _DEFAULT_TTL_S)
         with self._state_lock:
             self._token = token
@@ -149,7 +147,7 @@ def _ensure_fresh(self) -> None:
             if self._fresh_enough():  # another thread refreshed while we waited
                 return
             try:
-                self._refresh(force=False)
+                self._refresh()
             except RuntimeError as exc:
                 # Keep serving the current token; a request that then 401s triggers
                 # a forced refresh + retry (see _ProxyHandler._handle).
@@ -164,7 +162,7 @@ def token(self) -> str:
     def refresh(self) -> None:
         """Force a fresh mint now (used by the retry-on-401 path)."""
         with self._refresh_lock:
-            self._refresh(force=True)
+            self._refresh()
 
     def run_refresher(self) -> None:
         while not self._stop.wait(_REFRESHER_POLL_S):
diff --git a/tests/test_gateway_proxy.py b/tests/test_gateway_proxy.py
index c438bc7e..eaff0ad2 100644
--- a/tests/test_gateway_proxy.py
+++ b/tests/test_gateway_proxy.py
@@ -247,15 +247,14 @@ def test_fresh_token_is_not_refreshed(self, monkeypatch):
         _ = 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).
+    def test_near_expiry_triggers_forced_refresh(self, monkeypatch):
+        # First mint expires within the buffer, so reading .token force-refreshes it.
         state = _install_fake_token(monkeypatch, [100, 5000])
         cache = gateway_proxy._TokenCache("ws", None)
         _ = cache.token
-        assert state["forces"] == [True, False]
+        assert state["forces"] == [True, True]
         _ = cache.token  # now fresh again
-        assert state["forces"] == [True, False]
+        assert state["forces"] == [True, True]
 
     def test_refresh_is_single_flighted(self, monkeypatch):
         # A burst of concurrent requests at the expiry boundary must trigger ONE
@@ -267,8 +266,8 @@ def test_refresh_is_single_flighted(self, monkeypatch):
             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]
+        # 1 forced init + exactly 1 forced refresh shared by all 10 readers.
+        assert state["forces"] == [True, True]
 
     def test_ensure_fresh_keeps_token_when_refresh_fails(self, monkeypatch):
         _install_fake_token(monkeypatch, [5000])

From 2639ab80fc3fac93a6012e4bf11a531dd488e544 Mon Sep 17 00:00:00 2001
From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com>
Date: Fri, 21 Aug 2026 16:45:25 +0000
Subject: [PATCH 5/9] Gate Claude gateway refresh proxy

---
 src/ucode/agents/claude.py  | 15 +++++++++----
 src/ucode/gateway_proxy.py  | 45 +++++++++++++++++++++++--------------
 tests/test_agent_claude.py  | 34 +++++++++++++++++++++++-----
 tests/test_gateway_proxy.py | 19 +++++++++++-----
 4 files changed, 81 insertions(+), 32 deletions(-)

diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py
index 206e85d7..a32329e8 100644
--- a/src/ucode/agents/claude.py
+++ b/src/ucode/agents/claude.py
@@ -27,7 +27,9 @@
 from ucode.databricks import (
     build_auth_shell_command,
     build_tool_base_url,
+    get_databricks_token,
 )
+from ucode.gateway_proxy import 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 (
@@ -42,6 +44,7 @@
 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",
@@ -1017,11 +1020,13 @@ def _launch_relayed(state: dict, binary: str, tool_args: list[str]) -> None:
 
 
 def _launch_gateway(state: dict, binary: str, tool_args: list[str]) -> None:
-    from ucode.gateway_proxy import AUTHORIZATION_HEADER, start_proxy
-
     workspace = state["workspace"]
     server, cache, client = start_proxy(
-        workspace, state.get("profile"), 0, token_header=AUTHORIZATION_HEADER
+        workspace,
+        state.get("profile"),
+        0,
+        token_header=AUTHORIZATION_HEADER,
+        force_refresh_near_expiry=True,
     )
     token = cache.token
     os.environ["OAUTH_TOKEN"] = token
@@ -1055,9 +1060,11 @@ def launch(state: dict, tool_args: list[str]) -> None:
     if state.get("claude_relayed"):
         _launch_relayed(state, binary, tool_args)
         return
-    if workspace:
+    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))
 
 
diff --git a/src/ucode/gateway_proxy.py b/src/ucode/gateway_proxy.py
index 1298f449..75073044 100644
--- a/src/ucode/gateway_proxy.py
+++ b/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
@@ -115,9 +114,16 @@ 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()
@@ -126,11 +132,11 @@ def __init__(self, workspace: str, profile: str | None) -> None:
         # 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()
+        self._refresh(force=True)
 
-    def _refresh(self) -> None:
-        """Force-mint a token and record its expiry."""
-        token = get_databricks_token(self._workspace, self._profile, force_refresh=True)
+    def _refresh(self, *, force: bool) -> None:
+        """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:
             self._token = token
@@ -147,7 +153,7 @@ def _ensure_fresh(self) -> None:
             if self._fresh_enough():  # another thread refreshed while we waited
                 return
             try:
-                self._refresh()
+                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).
@@ -162,7 +168,7 @@ def token(self) -> str:
     def refresh(self) -> None:
         """Force a fresh mint now (used by the retry-on-401 path)."""
         with self._refresh_lock:
-            self._refresh()
+            self._refresh(force=True)
 
     def run_refresher(self) -> None:
         while not self._stop.wait(_REFRESHER_POLL_S):
@@ -234,9 +240,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.
@@ -361,6 +367,7 @@ def start_proxy(
     profile: str | None,
     port: int,
     token_header: str = _SWAP_HEADER,
+    force_refresh_near_expiry: bool = False,
 ) -> tuple[ThreadingHTTPServer, _TokenCache, httpx.Client]:
     """Start the loopback refresh proxy + its background token refresher.
 
@@ -373,7 +380,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.
diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py
index 5c64e56e..9d33f50b 100644
--- a/tests/test_agent_claude.py
+++ b/tests/test_agent_claude.py
@@ -634,9 +634,21 @@ def boom(name, entry, scope=mcp_mod.MCP_USER_SCOPE):
 
 
 class TestClaudeLaunch:
-    def test_runs_through_refresh_proxy(self, monkeypatch):
-        import ucode.gateway_proxy as gateway_proxy
+    def test_default_launch_keeps_existing_auth_path(self, monkeypatch):
+        calls: list[list[str]] = []
+        monkeypatch.delenv(claude.GATEWAY_MODEL_DISCOVERY_ENV_VAR, raising=False)
+        monkeypatch.delenv("OAUTH_TOKEN", raising=False)
+        monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "token")
+        monkeypatch.setattr(claude, "exec_or_spawn", lambda argv: calls.append(argv))
 
+        claude.launch({"workspace": WS, "profile": "test"}, ["--debug"])
+
+        assert os.environ["OAUTH_TOKEN"] == "token"
+        assert calls == [
+            ["claude", "--settings", str(claude.CLAUDE_SETTINGS_PATH), "--debug"]
+        ]
+
+    def test_runs_through_refresh_proxy(self, monkeypatch):
         calls: list[tuple] = []
 
         class Server:
@@ -665,15 +677,25 @@ def __init__(self, argv):
             def wait(self):
                 return 0
 
-        def start_proxy(workspace, profile, port, token_header):
-            calls.append(("proxy", workspace, profile, port, token_header))
+        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.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(gateway_proxy, "start_proxy", start_proxy)
+        monkeypatch.setattr(claude, "start_proxy", start_proxy)
         monkeypatch.setattr(claude.subprocess, "Popen", Process)
 
         with pytest.raises(SystemExit) as exc:
@@ -685,7 +707,7 @@ def start_proxy(workspace, profile, port, token_header):
         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, gateway_proxy.AUTHORIZATION_HEADER),
+            ("proxy", WS, "test", 0, claude.AUTHORIZATION_HEADER, True),
             ("serve",),
         ]
         assert calls[2][0] == "popen"
diff --git a/tests/test_gateway_proxy.py b/tests/test_gateway_proxy.py
index eaff0ad2..77ecdda8 100644
--- a/tests/test_gateway_proxy.py
+++ b/tests/test_gateway_proxy.py
@@ -247,20 +247,25 @@ def test_fresh_token_is_not_refreshed(self, monkeypatch):
         _ = cache.token
         assert state["forces"] == [True]  # no extra mint while fresh
 
-    def test_near_expiry_triggers_forced_refresh(self, monkeypatch):
-        # First mint expires within the buffer, so reading .token force-refreshes it.
+    def test_near_expiry_preserves_default_nonforce_refresh(self, monkeypatch):
         state = _install_fake_token(monkeypatch, [100, 5000])
         cache = gateway_proxy._TokenCache("ws", None)
         _ = cache.token
-        assert state["forces"] == [True, True]
+        assert state["forces"] == [True, False]
         _ = cache.token  # now fresh again
+        assert state["forces"] == [True, False]
+
+    def test_near_expiry_can_force_refresh(self, monkeypatch):
+        state = _install_fake_token(monkeypatch, [100, 5000])
+        cache = gateway_proxy._TokenCache("ws", None, force_refresh_near_expiry=True)
+        _ = cache.token
         assert state["forces"] == [True, True]
 
     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)
+        cache = gateway_proxy._TokenCache("ws", None, force_refresh_near_expiry=True)
         threads = [threading.Thread(target=lambda: cache.token) for _ in range(10)]
         for t in threads:
             t.start()
@@ -411,7 +416,11 @@ class _StubCache:
             def run_refresher(self):
                 return None
 
-        monkeypatch.setattr(gateway_proxy, "_TokenCache", lambda workspace, profile: _StubCache())
+        monkeypatch.setattr(
+            gateway_proxy,
+            "_TokenCache",
+            lambda workspace, profile, **_kwargs: _StubCache(),
+        )
         # Occupy a port to simulate the leftover proxy holding it.
         occupied = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         occupied.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

From d98d87be51bd26accc2a009d1844b936d4eba584 Mon Sep 17 00:00:00 2001
From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com>
Date: Fri, 21 Aug 2026 17:14:10 +0000
Subject: [PATCH 6/9] Format Claude launch test

---
 tests/test_agent_claude.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py
index 9d33f50b..92ab663d 100644
--- a/tests/test_agent_claude.py
+++ b/tests/test_agent_claude.py
@@ -644,9 +644,7 @@ def test_default_launch_keeps_existing_auth_path(self, monkeypatch):
         claude.launch({"workspace": WS, "profile": "test"}, ["--debug"])
 
         assert os.environ["OAUTH_TOKEN"] == "token"
-        assert calls == [
-            ["claude", "--settings", str(claude.CLAUDE_SETTINGS_PATH), "--debug"]
-        ]
+        assert calls == [["claude", "--settings", str(claude.CLAUDE_SETTINGS_PATH), "--debug"]]
 
     def test_runs_through_refresh_proxy(self, monkeypatch):
         calls: list[tuple] = []

From 987fc78355ecbfe77ace96507e212cdeab29ecd3 Mon Sep 17 00:00:00 2001
From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com>
Date: Fri, 21 Aug 2026 19:45:37 +0000
Subject: [PATCH 7/9] Preserve relayed token refresh behavior

---
 src/ucode/gateway_proxy.py  |  7 +++----
 tests/test_gateway_proxy.py | 10 +++++-----
 2 files changed, 8 insertions(+), 9 deletions(-)

diff --git a/src/ucode/gateway_proxy.py b/src/ucode/gateway_proxy.py
index 75073044..aef1a675 100644
--- a/src/ucode/gateway_proxy.py
+++ b/src/ucode/gateway_proxy.py
@@ -129,10 +129,9 @@ def __init__(
         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."""
diff --git a/tests/test_gateway_proxy.py b/tests/test_gateway_proxy.py
index 77ecdda8..a9851747 100644
--- a/tests/test_gateway_proxy.py
+++ b/tests/test_gateway_proxy.py
@@ -235,25 +235,25 @@ def fake(_ws, _profile, force_refresh=False):
 
 
 class TestTokenCache:
-    def test_initial_mint_is_forced(self, monkeypatch):
+    def test_initial_mint_preserves_default_nonforce_refresh(self, monkeypatch):
         state = _install_fake_token(monkeypatch, [5000])
         gateway_proxy._TokenCache("ws", None)
-        assert state["forces"] == [True]  # full-TTL start
+        assert state["forces"] == [False]
 
     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
+        assert state["forces"] == [False]  # no extra mint while fresh
 
     def test_near_expiry_preserves_default_nonforce_refresh(self, monkeypatch):
         state = _install_fake_token(monkeypatch, [100, 5000])
         cache = gateway_proxy._TokenCache("ws", None)
         _ = cache.token
-        assert state["forces"] == [True, False]
+        assert state["forces"] == [False, False]
         _ = cache.token  # now fresh again
-        assert state["forces"] == [True, False]
+        assert state["forces"] == [False, False]
 
     def test_near_expiry_can_force_refresh(self, monkeypatch):
         state = _install_fake_token(monkeypatch, [100, 5000])

From 4ab2151014759fa9a770875195d0c6b4957c5db0 Mon Sep 17 00:00:00 2001
From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com>
Date: Mon, 24 Aug 2026 19:34:54 +0000
Subject: [PATCH 8/9] Use Claude model discovery flag name

---
 src/ucode/agents/claude.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py
index a32329e8..94aa2702 100644
--- a/src/ucode/agents/claude.py
+++ b/src/ucode/agents/claude.py
@@ -44,7 +44,7 @@
 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"
+GATEWAY_MODEL_DISCOVERY_ENV_VAR = "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"
 
 SPEC: ToolSpec = {
     "binary": "claude",

From b21bcf305bb9323ae646df7a94669ab7bdc848d9 Mon Sep 17 00:00:00 2001
From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com>
Date: Mon, 24 Aug 2026 20:01:24 +0000
Subject: [PATCH 9/9] Restore ucode model discovery opt-in flag

---
 src/ucode/agents/claude.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py
index 94aa2702..a32329e8 100644
--- a/src/ucode/agents/claude.py
+++ b/src/ucode/agents/claude.py
@@ -44,7 +44,7 @@
 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 = "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"
+GATEWAY_MODEL_DISCOVERY_ENV_VAR = "ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY"
 
 SPEC: ToolSpec = {
     "binary": "claude",