From e24381d90a8e7d18dcddf7e84d913223bf6b7653 Mon Sep 17 00:00:00 2001
From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com>
Date: Mon, 31 Aug 2026 15:55:39 +0000
Subject: [PATCH] feat(claude): discover provider models

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

diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py
index 222e780d..c6d50c51 100644
--- a/src/ucode/agents/claude.py
+++ b/src/ucode/agents/claude.py
@@ -334,13 +334,11 @@ def render_overlay(
         "ENABLE_TOOL_SEARCH": "1",
         "CLAUDE_CODE_USE_GATEWAY": "1",
     }
-    # Native /model discovery: picker lists every gateway Messages-API endpoint,
-    # not just the family aliases. Skipped under a provider (its routing header
-    # would send a discovered gateway id to a provider that can't resolve it).
+    # Native /model discovery uses the MPS header to scope provider launches.
     discovery_enabled = (
         os.environ.get(GATEWAY_MODEL_DISCOVERY_ENV_VAR) == "1" or smart_routing_v2.enabled()
     )
-    if discovery_enabled and not provider:
+    if discovery_enabled:
         env["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] = "1"
     # Intentionally NOT setting ANTHROPIC_MODEL by default. Setting it produces a
     # duplicate catalog row in Claude Code's /model picker (e.g. "Opus 4.8 (1M
diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py
index dc99ab21..9e422ba3 100644
--- a/tests/test_agent_claude.py
+++ b/tests/test_agent_claude.py
@@ -151,13 +151,14 @@ def test_enables_gateway_model_discovery_for_smart_routing_v2(self, monkeypatch)
         overlay, _ = claude.render_overlay(WS, "s4")
         assert overlay["env"]["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1"
 
-    def test_gateway_model_discovery_skipped_under_provider(self, monkeypatch):
-        # A Model Provider Service routes every request to the external provider,
-        # so a discovered gateway endpoint id would reach a provider that can't
-        # resolve it — discovery must be off in that mode.
+    def test_enables_gateway_model_discovery_under_provider(self, monkeypatch):
         monkeypatch.setenv("ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY", "1")
         overlay, _ = claude.render_overlay(WS, "s4", provider="main.x.claude-svc")
-        assert "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY" not in overlay["env"]
+        assert overlay["env"]["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1"
+        assert (
+            "Databricks-Model-Provider-Service: main.x.claude-svc"
+            in overlay["env"]["ANTHROPIC_CUSTOM_HEADERS"]
+        )
 
     def test_sets_api_key_helper(self):
         overlay, _ = claude.render_overlay(WS, "s4")