Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 31 additions & 1 deletion src/ucode/agents/claude.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -135,13 +147,23 @@ def render_overlay(
]
)
env: dict[str, str] = {
"ANTHROPIC_MODEL": _maybe_add_1m_suffix(model),

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

what is the default model now that we remove this?

"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"):
Expand Down Expand Up @@ -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:
Expand Down
150 changes: 134 additions & 16 deletions tests/test_agent_claude.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Loading