Fix stale Claude model discovery setting

src/ucode/agents/claude.py

@@ -193,6 +193,9 @@ def _resolve_web_search_model(state: dict) -> str | None:
     "ANTHROPIC_DEFAULT_HAIKU_MODEL",
     "ANTHROPIC_DEFAULT_HAIKU_MODEL_NAME",
 )
+# Launch-scoped feature flags that ucode may write into Claude settings. These
+# must be removed again when the corresponding launch flag is absent.
+CLAUDE_CONDITIONAL_ENV_KEYS = ("CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY",)
 # Env keys ucode used to write but no longer does; stripped from the managed
 # settings file on every launch so stale values never linger.
 CLAUDE_REMOVED_ENV_KEYS = ("CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS",)
@@ -262,6 +265,15 @@ def managed_settings_are_current(state: dict) -> bool:
     return managed_file_is_verified(state, "claude", path, required_scope=required_scope)
 
 
+def gateway_model_discovery_setting_is_absent() -> bool:
+    """Return whether model discovery is absent from persistent Claude settings."""
+    env = read_json_safe(CLAUDE_SETTINGS_PATH).get("env")
+    actual = (
+        env.get("CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY") if isinstance(env, dict) else None
+    )
+    return actual is None
+
+
 def managed_settings_status(state: dict) -> tuple[Path | None, str, str]:
     path = _managed_settings_path()
     status, backup = managed_file_status(state, "claude", path, parser=_parse_managed_settings)
@@ -397,14 +409,6 @@ 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).
-    discovery_enabled = (
-        os.environ.get(GATEWAY_MODEL_DISCOVERY_ENV_VAR) == "1" or smart_routing_v2.enabled()
-    )
-    if discovery_enabled and not provider:
-        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
     # context) ✓") on top of the family-alias row from ANTHROPIC_DEFAULT_OPUS_MODEL.
@@ -626,6 +630,7 @@ def write_tool_config(
     managed_file_keys = list(managed_keys)
     for path in (
         [["env", key] for key in CLAUDE_MANAGED_MODEL_ENV_KEYS]
+        + [["env", key] for key in CLAUDE_CONDITIONAL_ENV_KEYS]
         + [["env", key] for key in CLAUDE_REMOVED_ENV_KEYS]
         + [["env", key] for key in CLAUDE_TRACING_ENV_KEYS]
         + [["hooks", "Stop"]]
@@ -669,6 +674,9 @@ def _compose(base: dict) -> dict:
             for key in CLAUDE_MANAGED_MODEL_ENV_KEYS:
                 if key not in overlay_env:
                     merged_env.pop(key, None)
+            for key in CLAUDE_CONDITIONAL_ENV_KEYS:
+                if key not in overlay_env:
+                    merged_env.pop(key, None)
             # deep_merge_dict keeps keys already in the file, so drop the ones ucode no
             # longer writes.
             for key in CLAUDE_REMOVED_ENV_KEYS:
@@ -1206,6 +1214,9 @@ def _build_claude_argv(
     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)
+    merged_env = merged.get("env")
+    if isinstance(merged_env, dict):
+        merged_env.pop("CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY", None)
     return [
         binary,
         *source_args,
@@ -1368,7 +1379,14 @@ def launch(state: dict, tool_args: list[str]) -> None:
     if first_prompt_routing:
         _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":
+    if (
+        workspace
+        and os.environ.get(GATEWAY_MODEL_DISCOVERY_ENV_VAR) == "1"
+        and not _has_provider_launch(state)
+    ):
+        # Discovery is launch-scoped. Pass it in the process environment rather
+        # than persisting it in Claude's private or OS-managed settings.
+        os.environ["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] = "1"
         _launch_claude_with_gateway_proxy(state, binary, tool_args, smart_routing=False)
         return
     if workspace:

src/ucode/cli.py

@@ -1886,6 +1886,7 @@ def _can_launch_from_cached_config(
         return (
             claude_agent.CLAUDE_SETTINGS_PATH.exists()
             and claude_agent.managed_settings_are_current(state)
+            and claude_agent.gateway_model_discovery_setting_is_absent()
         )
     return codex_agent.has_ucode_config() and codex_agent.managed_config_is_current(state)
 

src/ucode/smart_routing/v2.py

@@ -333,6 +333,7 @@ def launch_claude(
     token = get_databricks_token(workspace, state.get("profile"))
     os.environ[OAUTH_TOKEN_ENV_VAR] = token
     os.environ[GATEWAY_MODEL_DISCOVERY_ENV_VAR] = "1"
+    os.environ["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] = "1"
     model_ids, discovery_error = list_anthropic_models(workspace, token)
     if not model_ids:
         raise RuntimeError(discovery_error or "Anthropic models endpoint returned no Claude models")
@@ -348,7 +349,7 @@ def launch_claude(
     env = settings.setdefault("env", {})
     if not isinstance(env, dict):
         raise RuntimeError("Claude settings 'env' must be an object for smart routing.")
-    env["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] = "1"
+    env.pop("CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY", None)
     env[FIRST_PROMPT_SOCKET_ENV] = str(socket_path)
     model_overrides = settings.setdefault("modelOverrides", {})
     if not isinstance(model_overrides, dict):

tests/conftest.py

@@ -52,6 +52,7 @@ def reject_privileged_write(path, _desired_text):
     # `ucode`/`ucode configure` do mid-test. Tests that exercise the managed path set it explicitly.
     monkeypatch.delenv("ENABLE_MANAGED_AGENT_CONFIG", raising=False)
     monkeypatch.delenv("ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY", raising=False)
+    monkeypatch.delenv("CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY", raising=False)
     # The model-services listing is memoized for the life of the process, so without this a cached
     # result would leak into the next test and make a stubbed listing look like it was never called.
     databricks_mod.clear_model_services_cache()

tests/test_agent_claude.py

@@ -201,16 +201,16 @@ def test_gateway_model_discovery_disabled_unless_opted_in(self, monkeypatch, env
         overlay, _ = claude.render_overlay(WS, "s4")
         assert "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY" not in overlay["env"]
 
-    def test_enables_gateway_model_discovery(self, monkeypatch):
+    def test_does_not_persist_gateway_model_discovery(self, monkeypatch):
         monkeypatch.setenv("ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY", "1")
         overlay, _ = claude.render_overlay(WS, "s4")
-        assert overlay["env"]["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1"
+        assert "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY" not in overlay["env"]
 
-    def test_enables_gateway_model_discovery_for_smart_routing_v2(self, monkeypatch):
+    def test_smart_routing_does_not_persist_gateway_model_discovery(self, monkeypatch):
         monkeypatch.setenv(v2.ENV_VAR, "1")
         monkeypatch.delenv(claude.GATEWAY_MODEL_DISCOVERY_ENV_VAR, raising=False)
         overlay, _ = claude.render_overlay(WS, "s4")
-        assert overlay["env"]["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1"
+        assert "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY" not in overlay["env"]
 
     def test_gateway_model_discovery_skipped_under_provider(self, monkeypatch):
         # A Model Provider Service routes every request to the external provider,
@@ -220,6 +220,16 @@ def test_gateway_model_discovery_skipped_under_provider(self, monkeypatch):
         overlay, _ = claude.render_overlay(WS, "s4", provider="main.x.claude-svc")
         assert "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY" not in overlay["env"]
 
+    def test_gateway_model_discovery_setting_detects_stale_opt_in(self, monkeypatch):
+        monkeypatch.delenv(claude.GATEWAY_MODEL_DISCOVERY_ENV_VAR, raising=False)
+        monkeypatch.setattr(
+            claude,
+            "read_json_safe",
+            lambda path: {"env": {"CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY": "1"}},
+        )
+
+        assert claude.gateway_model_discovery_setting_is_absent() is False
+
     def test_sets_api_key_helper(self):
         overlay, _ = claude.render_overlay(WS, "s4")
         assert "apiKeyHelper" in overlay
@@ -613,6 +623,16 @@ def test_strips_stale_disable_experimental_betas(self, monkeypatch):
         assert written[0]["env"]["ENABLE_TOOL_SEARCH"] == "1"
         assert written[0]["env"]["CLAUDE_CODE_USE_GATEWAY"] == "1"
 
+    def test_strips_stale_gateway_model_discovery(self, monkeypatch):
+        existing = {"env": {"CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY": "1"}}
+        written: list = []
+        self._patch(monkeypatch, existing, written)
+        state = {"workspace": WS, "codex_models": []}
+
+        claude.write_tool_config(state, "databricks-claude-sonnet-4")
+
+        assert "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY" not in written[0]["env"]
+
 
 FAKE_MANAGED_PATH = Path("/tmp/ucode-test/managed-settings.json")
 
@@ -678,6 +698,26 @@ def test_managed_file_preserves_other_keys(self, monkeypatch):
         assert written["env"]["ANTHROPIC_BASE_URL"]
         assert written["apiKeyHelper"]
 
+    def test_managed_file_strips_stale_gateway_model_discovery(self, monkeypatch):
+        private_writes: list = []
+        managed_writes: list = []
+        stale = {"env": {"CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY": "1"}}
+        existing = {
+            str(claude.CLAUDE_SETTINGS_PATH): stale,
+            str(FAKE_MANAGED_PATH): stale,
+        }
+        self._patch(monkeypatch, private_writes, managed_writes, existing)
+
+        state = {"workspace": WS, "codex_models": []}
+
+        claude.write_tool_config(state, "databricks-claude-sonnet-4")
+
+        assert "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY" not in private_writes[0][1]["env"]
+        assert (
+            "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"
+            not in json.loads(managed_writes[0][1])["env"]
+        )
+
     def test_managed_file_merges_anthropic_custom_headers(self, monkeypatch):
         private_writes: list = []
         managed_writes: list = []
@@ -1121,6 +1161,7 @@ def start_proxy(workspace, port):
         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 os.environ["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1"
         assert calls[:2] == [
             ("proxy", WS, 0),
             ("serve",),
@@ -1129,6 +1170,7 @@ def start_proxy(workspace, port):
         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 "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY" not in json.loads(argv[2])["env"]
         assert argv[3:] == ["--debug"]
         assert calls[3:] == [("shutdown",), ("close",)]