diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 15d17fa..c335948 100644 --- a/src/ucode/agents/claude.py +++ b/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: diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 9888efd..1c801b9 100644 --- a/tests/test_agent_claude.py +++ b/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