Dedup Claude /model list (#165)

dedup model list for claude

Author: asujithan

src/ucode/agents/claude.py

@@ -76,6 +76,18 @@ def _resolve_web_search_model(state: dict) -> str | None:
     "MLFLOW_EXPERIMENT_ID",
     "MLFLOW_TRACING_SQL_WAREHOUSE_ID",
 )
+# Model-selection env keys ucode owns end-to-end. Anything in this tuple that
+# isn't written by render_overlay gets actively pruned from settings.json on
+# every launch, so stale values from older ucode versions never linger.
+CLAUDE_MANAGED_MODEL_ENV_KEYS = (
+    "ANTHROPIC_MODEL",
+    "ANTHROPIC_DEFAULT_OPUS_MODEL",
+    "ANTHROPIC_DEFAULT_OPUS_MODEL_NAME",
+    "ANTHROPIC_DEFAULT_SONNET_MODEL",
+    "ANTHROPIC_DEFAULT_SONNET_MODEL_NAME",
+    "ANTHROPIC_DEFAULT_HAIKU_MODEL",
+    "ANTHROPIC_DEFAULT_HAIKU_MODEL_NAME",
+)
 CLAUDE_TRACING_STOP_HOOK_SUFFIX = " autolog claude stop-hook"
 # Tracing is driven by an `mlflow autolog claude stop-hook` Stop hook, run by
 # the `mlflow` CLI on each session end. Pin to 3.11.x: 3.12 dropped the Unity
@@ -135,13 +147,23 @@ def render_overlay(
         ]
     )
     env: dict[str, str] = {
-        "ANTHROPIC_MODEL": _maybe_add_1m_suffix(model),
         "ANTHROPIC_BASE_URL": base_url,
         "ANTHROPIC_CUSTOM_HEADERS": custom_headers,
         "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS": "1",
         "CLAUDE_CODE_API_KEY_HELPER_TTL_MS": "900000",
     }
+    # Intentionally NOT setting ANTHROPIC_MODEL. 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. Without
+    # it, Default resolves through the pinned family alias and the picker shows
+    # only one row per model. `ucode claude -- --model X` still overrides for a
+    # single session via Claude Code's own --model flag.
+    _ = model  # API stability; no longer pinned via env.
     if claude_models:
+        # Picker rows show the raw routable id (e.g. "system.ai.claude-opus-4-8[1m]")
+        # so users can see which gateway-routable model is behind each shortcut.
+        # We deliberately don't set the `_NAME` companion env vars — the raw id
+        # is more useful than a friendly label for debugging gateway routing.
         if claude_models.get("opus"):
             env["ANTHROPIC_DEFAULT_OPUS_MODEL"] = _maybe_add_1m_suffix(claude_models["opus"])
         if claude_models.get("sonnet"):
@@ -258,6 +280,14 @@ def write_tool_config(state: dict, model: str) -> dict:
                 env_block.pop(key, None)
         # Strip only ucode's tracing Stop hook so user hooks stay intact.
         _remove_tracing_stop_hook(merged)
+    # Prune ucode-managed model env keys we deliberately don't write this run
+    # (e.g. ANTHROPIC_MODEL — see render_overlay).
+    overlay_env = overlay.get("env", {})
+    merged_env = merged.get("env")
+    if isinstance(merged_env, dict):
+        for key in CLAUDE_MANAGED_MODEL_ENV_KEYS:
+            if key not in overlay_env:
+                merged_env.pop(key, None)
     write_json_file(CLAUDE_SETTINGS_PATH, merged)
 
     if web_search_model:

tests/test_agent_claude.py

@@ -21,33 +21,53 @@ def test_display(self):
 
 
 class TestRenderOverlay:
-    def test_sets_anthropic_model(self):
-        overlay, _ = claude.render_overlay(WS, "databricks-claude-sonnet-4")
-        assert overlay["env"]["ANTHROPIC_MODEL"] == "databricks-claude-sonnet-4"
+    def test_does_not_set_anthropic_model_env(self):
+        # We deliberately don't pin ANTHROPIC_MODEL: when set, Claude Code's
+        # /model picker surfaces a duplicate catalog row on top of the family
+        # alias from ANTHROPIC_DEFAULT_OPUS_MODEL. Default falls back to the
+        # active family alias instead.
+        overlay, _ = claude.render_overlay(
+            WS, "databricks-claude-opus-4-7", claude_models={"opus": "databricks-claude-opus-4-7"}
+        )
+        assert "ANTHROPIC_MODEL" not in overlay["env"]
 
     def test_adds_1m_suffix_for_opus_4_6_and_later(self):
-        overlay, _ = claude.render_overlay(WS, "databricks-claude-opus-4-7")
-        assert overlay["env"]["ANTHROPIC_MODEL"] == "databricks-claude-opus-4-7[1m]"
+        overlay, _ = claude.render_overlay(
+            WS, "s4", claude_models={"opus": "databricks-claude-opus-4-7"}
+        )
+        assert overlay["env"]["ANTHROPIC_DEFAULT_OPUS_MODEL"] == "databricks-claude-opus-4-7[1m]"
 
     def test_adds_1m_suffix_for_sonnet_4_6_and_later(self):
-        overlay, _ = claude.render_overlay(WS, "databricks-claude-sonnet-4-7")
-        assert overlay["env"]["ANTHROPIC_MODEL"] == "databricks-claude-sonnet-4-7[1m]"
+        overlay, _ = claude.render_overlay(
+            WS, "s4", claude_models={"sonnet": "databricks-claude-sonnet-4-7"}
+        )
+        assert (
+            overlay["env"]["ANTHROPIC_DEFAULT_SONNET_MODEL"] == "databricks-claude-sonnet-4-7[1m]"
+        )
 
-    def test_does_not_add_1m_suffix_for_other_models(self):
-        overlay, _ = claude.render_overlay(WS, "databricks-claude-haiku-4-6")
-        assert overlay["env"]["ANTHROPIC_MODEL"] == "databricks-claude-haiku-4-6"
+    def test_does_not_add_1m_suffix_for_haiku(self):
+        overlay, _ = claude.render_overlay(
+            WS, "s4", claude_models={"haiku": "databricks-claude-haiku-4-6"}
+        )
+        assert overlay["env"]["ANTHROPIC_DEFAULT_HAIKU_MODEL"] == "databricks-claude-haiku-4-6"
 
     def test_does_not_duplicate_1m_suffix(self):
-        overlay, _ = claude.render_overlay(WS, "databricks-claude-opus-4-7[1m]")
-        assert overlay["env"]["ANTHROPIC_MODEL"] == "databricks-claude-opus-4-7[1m]"
+        overlay, _ = claude.render_overlay(
+            WS, "s4", claude_models={"opus": "databricks-claude-opus-4-7[1m]"}
+        )
+        assert overlay["env"]["ANTHROPIC_DEFAULT_OPUS_MODEL"] == "databricks-claude-opus-4-7[1m]"
 
     def test_adds_1m_suffix_for_model_services_name(self):
-        overlay, _ = claude.render_overlay(WS, "system.ai.claude-opus-4-8")
-        assert overlay["env"]["ANTHROPIC_MODEL"] == "system.ai.claude-opus-4-8[1m]"
+        overlay, _ = claude.render_overlay(
+            WS, "s4", claude_models={"opus": "system.ai.claude-opus-4-8"}
+        )
+        assert overlay["env"]["ANTHROPIC_DEFAULT_OPUS_MODEL"] == "system.ai.claude-opus-4-8[1m]"
 
     def test_no_1m_suffix_for_model_services_haiku(self):
-        overlay, _ = claude.render_overlay(WS, "system.ai.claude-haiku-4-6")
-        assert overlay["env"]["ANTHROPIC_MODEL"] == "system.ai.claude-haiku-4-6"
+        overlay, _ = claude.render_overlay(
+            WS, "s4", claude_models={"haiku": "system.ai.claude-haiku-4-6"}
+        )
+        assert overlay["env"]["ANTHROPIC_DEFAULT_HAIKU_MODEL"] == "system.ai.claude-haiku-4-6"
 
     def test_sets_anthropic_base_url(self):
         overlay, _ = claude.render_overlay(WS, "s4")
@@ -90,6 +110,25 @@ def test_model_overrides_not_set_when_no_models(self):
         env = overlay["env"]
         assert "ANTHROPIC_DEFAULT_SONNET_MODEL" not in env
 
+    def test_picker_labels_show_raw_routable_id(self):
+        # We deliberately don't set the `_NAME` companion env vars. Showing the
+        # raw `system.ai.…` / `databricks-…` id in the picker label tells users
+        # exactly which gateway-routable model is behind each shortcut, which is
+        # more useful than a friendly catalog label for Databricks routing.
+        models = {
+            "opus": "system.ai.claude-opus-4-8",
+            "sonnet": "databricks-claude-sonnet-4-6",
+            "haiku": "system.ai.claude-haiku-4-5",
+        }
+        overlay, _ = claude.render_overlay(WS, "s4", claude_models=models)
+        env = overlay["env"]
+        assert env["ANTHROPIC_DEFAULT_OPUS_MODEL"] == "system.ai.claude-opus-4-8[1m]"
+        assert "ANTHROPIC_DEFAULT_OPUS_MODEL_NAME" not in env
+        assert env["ANTHROPIC_DEFAULT_SONNET_MODEL"] == "databricks-claude-sonnet-4-6[1m]"
+        assert "ANTHROPIC_DEFAULT_SONNET_MODEL_NAME" not in env
+        assert env["ANTHROPIC_DEFAULT_HAIKU_MODEL"] == "system.ai.claude-haiku-4-5"
+        assert "ANTHROPIC_DEFAULT_HAIKU_MODEL_NAME" not in env
+
     def test_managed_keys_include_api_key_helper(self):
         _, keys = claude.render_overlay(WS, "s4")
         assert ["apiKeyHelper"] in keys
@@ -357,3 +396,82 @@ def fake_execvp(binary: str, args: list[str]) -> None:
                 ["claude", "--settings", str(claude.CLAUDE_SETTINGS_PATH), "--debug"],
             )
         ]
+
+
+class TestWriteToolConfigPrunesStaleModelEnv:
+    """Stale ucode-managed model env keys (ANTHROPIC_MODEL, etc.) from earlier
+    ucode versions must be removed on every launch — otherwise they linger in
+    settings.json and re-introduce the duplicate /model picker row that this
+    change is meant to remove.
+    """
+
+    def _patch(self, monkeypatch, existing_settings):
+        monkeypatch.setattr(claude, "backup_existing_file", lambda *a, **kw: True)
+        monkeypatch.setattr(claude, "read_json_safe", lambda path: existing_settings)
+        written: dict = {}
+
+        def fake_write(path, payload):
+            written["payload"] = payload
+
+        monkeypatch.setattr(claude, "write_json_file", fake_write)
+        monkeypatch.setattr(claude, "save_state", lambda state: None)
+        monkeypatch.setattr(claude, "_register_web_search_mcp", lambda *a, **kw: True)
+        return written
+
+    def test_prunes_stale_anthropic_model_from_prior_run(self, monkeypatch):
+        existing = {
+            "env": {
+                "ANTHROPIC_MODEL": "system.ai.claude-opus-4-8[1m]",
+                "ANTHROPIC_DEFAULT_OPUS_MODEL": "system.ai.claude-opus-4-8[1m]",
+                "MY_CUSTOM_VAR": "keep-me",
+            }
+        }
+        written = self._patch(monkeypatch, existing)
+        state = {
+            "workspace": WS,
+            "claude_models": {"opus": "system.ai.claude-opus-4-8"},
+        }
+        claude.write_tool_config(state, "system.ai.claude-opus-4-8")
+        env = written["payload"]["env"]
+        assert "ANTHROPIC_MODEL" not in env
+        # Family default we still write this run is preserved.
+        assert env["ANTHROPIC_DEFAULT_OPUS_MODEL"] == "system.ai.claude-opus-4-8[1m]"
+        # User-owned keys are untouched.
+        assert env["MY_CUSTOM_VAR"] == "keep-me"
+
+    def test_prunes_unused_family_default_when_models_change(self, monkeypatch):
+        # Earlier launch wrote a sonnet default; the new state only has opus.
+        # The stale sonnet keys should be removed.
+        existing = {
+            "env": {
+                "ANTHROPIC_DEFAULT_SONNET_MODEL": "databricks-claude-sonnet-4-6[1m]",
+            }
+        }
+        written = self._patch(monkeypatch, existing)
+        state = {"workspace": WS, "claude_models": {"opus": "system.ai.claude-opus-4-8"}}
+        claude.write_tool_config(state, "system.ai.claude-opus-4-8")
+        env = written["payload"]["env"]
+        assert "ANTHROPIC_DEFAULT_SONNET_MODEL" not in env
+        assert env["ANTHROPIC_DEFAULT_OPUS_MODEL"] == "system.ai.claude-opus-4-8[1m]"
+
+    def test_prunes_stale_name_companion_keys_from_older_ucode(self, monkeypatch):
+        # An older ucode build briefly wrote `_NAME` companion env vars to give
+        # the picker friendly labels. The current build only writes the raw id,
+        # so any leftover `_NAME` keys must be pruned — otherwise users who
+        # tested the in-between version would see stale labels.
+        existing = {
+            "env": {
+                "ANTHROPIC_DEFAULT_OPUS_MODEL": "system.ai.claude-opus-4-8[1m]",
+                "ANTHROPIC_DEFAULT_OPUS_MODEL_NAME": "Opus 4.8 (1M)",
+                "ANTHROPIC_DEFAULT_SONNET_MODEL_NAME": "Sonnet 4.6 (1M)",
+                "ANTHROPIC_DEFAULT_HAIKU_MODEL_NAME": "Haiku 4.5",
+            }
+        }
+        written = self._patch(monkeypatch, existing)
+        state = {"workspace": WS, "claude_models": {"opus": "system.ai.claude-opus-4-8"}}
+        claude.write_tool_config(state, "system.ai.claude-opus-4-8")
+        env = written["payload"]["env"]
+        assert env["ANTHROPIC_DEFAULT_OPUS_MODEL"] == "system.ai.claude-opus-4-8[1m]"
+        assert "ANTHROPIC_DEFAULT_OPUS_MODEL_NAME" not in env
+        assert "ANTHROPIC_DEFAULT_SONNET_MODEL_NAME" not in env
+        assert "ANTHROPIC_DEFAULT_HAIKU_MODEL_NAME" not in env