From e4f26ad04ba846a95d1647acc2107b1ee82fc7d8 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Tue, 1 Sep 2026 22:36:30 +0000 Subject: [PATCH 1/9] Merge user Anthropic custom headers --- src/ucode/agents/claude.py | 41 +++++++++++++++++++++++ tests/test_agent_claude.py | 67 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 108 insertions(+) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 4fea68e8..5ef42e4b 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -193,6 +193,14 @@ def _resolve_web_search_model(state: dict) -> str | None: # 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",) +ANTHROPIC_CUSTOM_HEADERS_ENV_KEY = "ANTHROPIC_CUSTOM_HEADERS" +CLAUDE_MANAGED_CUSTOM_HEADER_NAMES = frozenset( + { + "x-databricks-use-coding-agent-mode", + "user-agent", + "databricks-model-provider-service", + } +) 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 @@ -626,9 +634,19 @@ def write_tool_config( # V2 installs routing hooks in a transient per-launch settings file. Persistent settings must # contain no ucode routing hooks; surgically strip legacy ones while preserving user hooks. def _compose(base: dict) -> dict: + base_env = base.get("env") + existing_custom_headers = ( + base_env.get(ANTHROPIC_CUSTOM_HEADERS_ENV_KEY) + if isinstance(base_env, dict) + else None + ) # deepcopy the overlay per file so merging into one base can't alias nested dicts into # the other (deep_merge_dict grafts overlay's own dict objects onto a base missing the key). merged = deep_merge_dict(base, copy.deepcopy(overlay)) + overlay_custom_headers = overlay["env"][ANTHROPIC_CUSTOM_HEADERS_ENV_KEY] + merged["env"][ANTHROPIC_CUSTOM_HEADERS_ENV_KEY] = _merge_anthropic_custom_headers( + existing_custom_headers, overlay_custom_headers + ) # Drop any apiKeyHelper a prior non-relayed launch left in the file; relayed # must not carry one (it would outrank the subscription OAuth). if relayed: @@ -687,6 +705,29 @@ def _compose(base: dict) -> dict: return state +def _merge_anthropic_custom_headers(existing: object, ucode_headers: str) -> str: + """Preserve user headers while replacing the header names managed by ucode. + + Claude's ``ANTHROPIC_CUSTOM_HEADERS`` value is a newline-delimited string rather than a map, + so a normal deep merge replaces it wholesale. Header names are case-insensitive; any existing + line with a name ucode writes is dropped before ucode's current value is appended. Non-header + lines are retained to avoid silently discarding user configuration we do not understand. + """ + + if not isinstance(existing, str) or not existing: + return ucode_headers + + preserved: list[str] = [] + for line in existing.splitlines(): + name, separator, _value = line.partition(":") + if separator and name.strip().casefold() in CLAUDE_MANAGED_CUSTOM_HEADER_NAMES: + continue + if line: + preserved.append(line) + preserved.extend(ucode_headers.splitlines()) + return "\n".join(preserved) + + def _reconcile_managed_settings( state: dict, compose: Callable[[dict], dict], diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 29088326..3d3ba971 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -382,6 +382,51 @@ def test_headers_newline_delimited(self, monkeypatch): assert "\n" in self._ua(monkeypatch) +class TestMergeAnthropicCustomHeaders: + def test_preserves_existing_user_headers(self): + existing = "X-User-Header: keep-me\nopaque-user-value" + managed = "User-Agent: ucode/1.0 claude/2.0" + + merged = claude._merge_anthropic_custom_headers(existing, managed) + + assert "X-User-Header: keep-me" in merged + assert "opaque-user-value" in merged + + def test_overrides_existing_ucode_owned_headers(self): + existing = "\n".join( + [ + "user-agent: custom-agent", + "X-Databricks-Use-Coding-Agent-Mode: false", + "Databricks-Model-Provider-Service: stale.provider.service", + ] + ) + managed = "\n".join( + [ + "x-databricks-use-coding-agent-mode: true", + "User-Agent: ucode/1.0 claude/2.0", + ] + ) + + merged = claude._merge_anthropic_custom_headers(existing, managed) + + assert "custom-agent" not in merged + assert "false" not in merged + assert "stale.provider.service" not in merged + assert merged.endswith(managed) + + def test_adds_managed_headers_when_none_exist(self): + managed = "\n".join( + [ + "x-databricks-use-coding-agent-mode: true", + "User-Agent: ucode/1.0 claude/2.0", + ] + ) + + merged = claude._merge_anthropic_custom_headers(None, managed) + + assert merged == managed + + class TestRenderOverlayWebSearchDisable: def test_settings_overlay_never_includes_mcp_servers(self): # MCP servers belong in ~/.claude.json, not settings.json. @@ -599,6 +644,28 @@ def test_managed_file_preserves_other_keys(self, monkeypatch): assert written["env"]["ANTHROPIC_BASE_URL"] assert written["apiKeyHelper"] + def test_managed_file_merges_anthropic_custom_headers(self, monkeypatch): + private_writes: list = [] + managed_writes: list = [] + existing = { + str(FAKE_MANAGED_PATH): { + "env": { + "ANTHROPIC_CUSTOM_HEADERS": "X-Enterprise-Header: retain\nUser-Agent: old" + } + } + } + self._patch(monkeypatch, private_writes, managed_writes, existing) + state = {"workspace": WS, "codex_models": []} + + claude.write_tool_config(state, "databricks-claude-sonnet-4") + + _, text = managed_writes[0] + headers = json.loads(text)["env"]["ANTHROPIC_CUSTOM_HEADERS"] + assert "X-Enterprise-Header: retain" in headers + assert "User-Agent: old" not in headers + assert "User-Agent: ucode/" in headers + assert "x-databricks-use-coding-agent-mode: true" in headers + def test_managed_file_preserves_enterprise_permission_denies(self, monkeypatch): private_writes: list = [] managed_writes: list = [] From 3b5bb53c73e9e8137a61dff470f2c3e636f6131e Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Tue, 1 Sep 2026 22:57:02 +0000 Subject: [PATCH 2/9] fix --- src/ucode/agents/claude.py | 18 ++++++----- tests/test_agent_claude.py | 62 ++++++++++++++------------------------ 2 files changed, 33 insertions(+), 47 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 5ef42e4b..3ce53822 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -636,9 +636,7 @@ def write_tool_config( def _compose(base: dict) -> dict: base_env = base.get("env") existing_custom_headers = ( - base_env.get(ANTHROPIC_CUSTOM_HEADERS_ENV_KEY) - if isinstance(base_env, dict) - else None + base_env.get(ANTHROPIC_CUSTOM_HEADERS_ENV_KEY) if isinstance(base_env, dict) else None ) # deepcopy the overlay per file so merging into one base can't alias nested dicts into # the other (deep_merge_dict grafts overlay's own dict objects onto a base missing the key). @@ -708,10 +706,16 @@ def _compose(base: dict) -> dict: def _merge_anthropic_custom_headers(existing: object, ucode_headers: str) -> str: """Preserve user headers while replacing the header names managed by ucode. - Claude's ``ANTHROPIC_CUSTOM_HEADERS`` value is a newline-delimited string rather than a map, - so a normal deep merge replaces it wholesale. Header names are case-insensitive; any existing - line with a name ucode writes is dropped before ucode's current value is appended. Non-header - lines are retained to avoid silently discarding user configuration we do not understand. + Claude's ``ANTHROPIC_CUSTOM_HEADERS`` value is a newline-delimited string. To merge it, we: + + 1. Split the existing custom headers by newline into individual header items. + 2. Split each item on ``:`` to identify its header name. + 3. Replace headers in ``CLAUDE_MANAGED_CUSTOM_HEADER_NAMES`` with ucode's values, while + preserving all other existing headers. + 4. Combine the preserved existing headers with ucode's managed header values. + + Header names are compared case-insensitively. Non-header lines are also preserved to avoid + silently discarding user configuration we do not understand. """ if not isinstance(existing, str) or not existing: diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 3d3ba971..eea05850 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -383,48 +383,29 @@ def test_headers_newline_delimited(self, monkeypatch): class TestMergeAnthropicCustomHeaders: - def test_preserves_existing_user_headers(self): - existing = "X-User-Header: keep-me\nopaque-user-value" - managed = "User-Agent: ucode/1.0 claude/2.0" - - merged = claude._merge_anthropic_custom_headers(existing, managed) - - assert "X-User-Header: keep-me" in merged - assert "opaque-user-value" in merged - - def test_overrides_existing_ucode_owned_headers(self): - existing = "\n".join( + def test_merges_existing_settings_with_ucode_managed_headers(self): + headers_from_existing_settings = "\n".join( [ + "X-User-Header: keep-me", "user-agent: custom-agent", - "X-Databricks-Use-Coding-Agent-Mode: false", - "Databricks-Model-Provider-Service: stale.provider.service", ] ) - managed = "\n".join( + headers_managed_by_ucode = "\n".join( [ "x-databricks-use-coding-agent-mode: true", "User-Agent: ucode/1.0 claude/2.0", ] ) - merged = claude._merge_anthropic_custom_headers(existing, managed) - - assert "custom-agent" not in merged - assert "false" not in merged - assert "stale.provider.service" not in merged - assert merged.endswith(managed) - - def test_adds_managed_headers_when_none_exist(self): - managed = "\n".join( - [ - "x-databricks-use-coding-agent-mode: true", - "User-Agent: ucode/1.0 claude/2.0", - ] + merged_headers = claude._merge_anthropic_custom_headers( + headers_from_existing_settings, headers_managed_by_ucode ) - merged = claude._merge_anthropic_custom_headers(None, managed) - - assert merged == managed + assert merged_headers.splitlines() == [ + "X-User-Header: keep-me", # Preserved from existing settings. + "x-databricks-use-coding-agent-mode: true", # Newly added by ucode. + "User-Agent: ucode/1.0 claude/2.0", # From ucode; overwrites existing. + ] class TestRenderOverlayWebSearchDisable: @@ -647,24 +628,25 @@ def test_managed_file_preserves_other_keys(self, monkeypatch): def test_managed_file_merges_anthropic_custom_headers(self, monkeypatch): private_writes: list = [] managed_writes: list = [] - existing = { + existing_managed_settings = { str(FAKE_MANAGED_PATH): { - "env": { - "ANTHROPIC_CUSTOM_HEADERS": "X-Enterprise-Header: retain\nUser-Agent: old" - } + "env": {"ANTHROPIC_CUSTOM_HEADERS": "X-Enterprise-Header: retain\nUser-Agent: old"} } } - self._patch(monkeypatch, private_writes, managed_writes, existing) + self._patch(monkeypatch, private_writes, managed_writes, existing_managed_settings) + monkeypatch.setattr(claude, "ucode_version", lambda: "1.0") + monkeypatch.setattr(claude, "agent_version", lambda _binary: "2.0") state = {"workspace": WS, "codex_models": []} claude.write_tool_config(state, "databricks-claude-sonnet-4") _, text = managed_writes[0] - headers = json.loads(text)["env"]["ANTHROPIC_CUSTOM_HEADERS"] - assert "X-Enterprise-Header: retain" in headers - assert "User-Agent: old" not in headers - assert "User-Agent: ucode/" in headers - assert "x-databricks-use-coding-agent-mode: true" in headers + merged_headers = json.loads(text)["env"]["ANTHROPIC_CUSTOM_HEADERS"] + assert merged_headers.splitlines() == [ + "X-Enterprise-Header: retain", # Preserved from existing managed settings. + "x-databricks-use-coding-agent-mode: true", # Newly added by ucode. + "User-Agent: ucode/1.0 claude/2.0", # From ucode; overwrites existing. + ] def test_managed_file_preserves_enterprise_permission_denies(self, monkeypatch): private_writes: list = [] From cd039a4ebd5733ed2307d438aa44abe1bcd421fb Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Wed, 2 Sep 2026 01:17:50 +0000 Subject: [PATCH 3/9] Preserve managed Claude model defaults --- src/ucode/agents/__init__.py | 2 + src/ucode/agents/claude.py | 29 +++++++++----- src/ucode/cli.py | 8 ++++ src/ucode/managed_resolve.py | 7 ++++ tests/test_agent_claude.py | 75 +++++++++++++++++++++++++++++++---- tests/test_cli.py | 19 +++++++++ tests/test_managed_resolve.py | 7 +++- 7 files changed, 130 insertions(+), 17 deletions(-) diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index 578aa208..7765f4c4 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -350,6 +350,7 @@ def configure_tool( relayed: bool = False, route_root_model: str | None = None, custom_model: str | None = None, + coding_agent_config_families: set[str] | None = None, ) -> dict: result: dict | tuple[dict, str] if tool == "codex": @@ -368,6 +369,7 @@ def configure_tool( relayed=relayed, route_root_model=route_root_model, custom_model=custom_model, + coding_agent_config_families=coding_agent_config_families, ) else: # provider routing is claude/codex-only; every other tool needs a model. diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index f7333140..d6129c22 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -179,9 +179,8 @@ 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. +# Model-selection env keys ucode manages. Existing family defaults are preserved unless Coding +# Agent Config explicitly supplies that family. CLAUDE_MANAGED_MODEL_ENV_KEYS = ( "ANTHROPIC_MODEL", "ANTHROPIC_DEFAULT_FABLE_MODEL", @@ -193,6 +192,12 @@ def _resolve_web_search_model(state: dict) -> str | None: "ANTHROPIC_DEFAULT_HAIKU_MODEL", "ANTHROPIC_DEFAULT_HAIKU_MODEL_NAME", ) +CLAUDE_DEFAULT_MODEL_ENV_KEYS = { + "fable": "ANTHROPIC_DEFAULT_FABLE_MODEL", + "opus": "ANTHROPIC_DEFAULT_OPUS_MODEL", + "sonnet": "ANTHROPIC_DEFAULT_SONNET_MODEL", + "haiku": "ANTHROPIC_DEFAULT_HAIKU_MODEL", +} # 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",) @@ -588,6 +593,7 @@ def write_tool_config( relayed: bool = False, route_root_model: str | None = None, custom_model: str | None = None, + coding_agent_config_families: set[str] | None = None, ) -> dict: backup_existing_file(CLAUDE_SETTINGS_PATH, CLAUDE_BACKUP_PATH) web_search_model = _resolve_web_search_model(state) @@ -641,10 +647,15 @@ def _compose(base: dict) -> dict: existing_custom_headers = ( base_env.get(ANTHROPIC_CUSTOM_HEADERS_ENV_KEY) if isinstance(base_env, dict) else None ) - # deepcopy the overlay per file so merging into one base can't alias nested dicts into - # the other (deep_merge_dict grafts overlay's own dict objects onto a base missing the key). - merged = deep_merge_dict(base, copy.deepcopy(overlay)) - overlay_custom_headers = overlay["env"][ANTHROPIC_CUSTOM_HEADERS_ENV_KEY] + # Copy the overlay per file so merging into one base cannot affect the other. + overlay_for_merge = copy.deepcopy(overlay) + # Only Coding Agent Config may add or replace family defaults. + configured_families = coding_agent_config_families or set() + for family, key in CLAUDE_DEFAULT_MODEL_ENV_KEYS.items(): + if family not in configured_families: + overlay_for_merge["env"].pop(key, None) + merged = deep_merge_dict(base, overlay_for_merge) + overlay_custom_headers = overlay_for_merge["env"][ANTHROPIC_CUSTOM_HEADERS_ENV_KEY] merged["env"][ANTHROPIC_CUSTOM_HEADERS_ENV_KEY] = _merge_anthropic_custom_headers( existing_custom_headers, overlay_custom_headers ) @@ -663,11 +674,11 @@ def _compose(base: dict) -> dict: _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", {}) + overlay_env = overlay_for_merge.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: + if key not in overlay_env and key not in CLAUDE_DEFAULT_MODEL_ENV_KEYS.values(): merged_env.pop(key, None) # deep_merge_dict keeps keys already in the file, so drop the ones ucode no # longer writes. diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 2168e8a6..68545e4b 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -74,6 +74,7 @@ refresh_managed_config, ) from ucode.managed_resolve import ( + managed_claude_family_models, managed_default_model, managed_enabled_tools, managed_launch_model, @@ -2015,6 +2016,11 @@ def _launch_tool( # or Foundry service), and, for Bedrock, expose Claude models to pin. provider_models = None relayed = False + coding_agent_config_families = ( + set(managed_claude_family_models(managed) or {}) + if tool == "claude" and managed is not None + else set() + ) if provider: provider_models, error, relayed = resolve_provider_models(tool, state, provider) if error: @@ -2038,6 +2044,7 @@ def _launch_tool( authored = managed_provider_family_models(managed) if authored: provider_models = authored + coding_agent_config_families = set(authored) # The router's per-launch pick for the root session. Codex pins it as the # resolved model; claude pins it via ANTHROPIC_MODEL (route_root_model). route_root_model = None @@ -2091,6 +2098,7 @@ def _launch_tool( # the latter pins a raw id into every family alias, which would clobber the service's # per-family target pins. custom_model=model if (tool == "claude" and not provider) else None, + coding_agent_config_families=coding_agent_config_families, ) print_section(_launch_title(tool)) if managed is not None: diff --git a/src/ucode/managed_resolve.py b/src/ucode/managed_resolve.py index b6658d4e..947b174b 100644 --- a/src/ucode/managed_resolve.py +++ b/src/ucode/managed_resolve.py @@ -186,6 +186,13 @@ def managed_default_model(managed: dict, tool: str) -> str | None: return _str(_agent_model_config(managed, tool).get("default_model")) +def managed_claude_family_models(managed: dict) -> dict[str, str] | None: + """Claude family models explicitly authored by Coding Agent Config.""" + + models = _manifest_models(managed, "claude") + return cast("dict[str, str]", models) if isinstance(models, dict) else None + + def managed_provider_family_models(managed: dict) -> dict[str, str] | None: """Claude's authored per-family models for launch, when a managed config routes it through a Model Provider Service. diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index bd504df4..8b8a76ed 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -654,6 +654,33 @@ def fake_write_managed(path, text, **kwargs): monkeypatch.setattr(claude, "reconcile_managed_file", fake_write_managed) + def _write_managed_model_defaults( + self, + monkeypatch, + *, + existing: dict[str, str], + discovered: dict[str, str], + configured_families: set[str] | None = None, + ) -> dict: + private_writes: list = [] + managed_writes: list = [] + self._patch( + monkeypatch, + private_writes, + managed_writes, + {str(FAKE_MANAGED_PATH): {"env": existing}}, + ) + state = {"workspace": WS, "codex_models": [], "claude_models": discovered} + + claude.write_tool_config( + state, + next(iter(discovered.values())), + coding_agent_config_families=configured_families, + ) + + _, text = managed_writes[0] + return json.loads(text)["env"] + def test_writes_managed_file_by_default(self, monkeypatch): private_writes: list = [] managed_writes: list = [] @@ -701,6 +728,40 @@ def test_managed_file_merges_anthropic_custom_headers(self, monkeypatch): "x-databricks-use-coding-agent-mode: true", # Newly added by ucode. ] + def test_managed_file_preserves_existing_model_defaults(self, monkeypatch): + managed_env = self._write_managed_model_defaults( + monkeypatch, + existing={ + "ANTHROPIC_DEFAULT_FABLE_MODEL": "", + "ANTHROPIC_DEFAULT_OPUS_MODEL": "enterprise-opus", + }, + discovered={"opus": "discovered-opus"}, + ) + + assert managed_env["ANTHROPIC_DEFAULT_FABLE_MODEL"] == "" # Preserved as empty. + assert ( + managed_env["ANTHROPIC_DEFAULT_OPUS_MODEL"] == "enterprise-opus" + ) # Preserved from managed settings. + + def test_managed_file_does_not_add_discovered_model_defaults(self, monkeypatch): + managed_env = self._write_managed_model_defaults( + monkeypatch, + existing={}, + discovered={"haiku": "discovered-haiku"}, + ) + + assert "ANTHROPIC_DEFAULT_HAIKU_MODEL" not in managed_env + + def test_coding_agent_config_overwrites_managed_model_default(self, monkeypatch): + managed_env = self._write_managed_model_defaults( + monkeypatch, + existing={"ANTHROPIC_DEFAULT_SONNET_MODEL": "enterprise-sonnet"}, + discovered={"sonnet": "system.ai.claude-sonnet-4-6"}, + configured_families={"sonnet"}, + ) + + assert managed_env["ANTHROPIC_DEFAULT_SONNET_MODEL"] == "system.ai.claude-sonnet-4-6[1m]" + def test_managed_file_preserves_enterprise_permission_denies(self, monkeypatch): private_writes: list = [] managed_writes: list = [] @@ -1184,7 +1245,7 @@ def launch_v2(state, tool_args, **kwargs): assert calls[2:] == [("shutdown",), ("close",)] -class TestWriteToolConfigPrunesStaleModelEnv: +class TestWriteToolConfigModelEnv: """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 @@ -1220,14 +1281,12 @@ def test_prunes_stale_anthropic_model_from_prior_run(self, monkeypatch): 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. + # Family defaults are preserved unless Coding Agent Config owns the family. 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. + def test_preserves_family_defaults_when_discovered_models_change(self, monkeypatch): existing = { "env": { "ANTHROPIC_DEFAULT_SONNET_MODEL": "databricks-claude-sonnet-4-6[1m]", @@ -1237,8 +1296,10 @@ def test_prunes_unused_family_default_when_models_change(self, monkeypatch): 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]" + assert ( + env["ANTHROPIC_DEFAULT_SONNET_MODEL"] == "databricks-claude-sonnet-4-6[1m]" + ) # Preserved from existing settings. + assert "ANTHROPIC_DEFAULT_OPUS_MODEL" not in env # Discovery does not add defaults. def test_prunes_stale_name_companion_keys_from_older_ucode(self, monkeypatch): # An older ucode build briefly wrote `_NAME` companion env vars to give diff --git a/tests/test_cli.py b/tests/test_cli.py index 5d9e53ce..908fe567 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -3697,6 +3697,25 @@ def test_the_recommended_agent_gets_the_recommended_model(self, monkeypatch): ) assert cfg.call_args.args[2] == "system.ai.claude-haiku-4-5" + def test_passes_configured_claude_families_to_writer(self, monkeypatch): + monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") + managed = { + "enabled_agents": { + "claude": { + "model_config": { + "models": { + "default_sonnet_model": "system.ai.claude-sonnet-4-6", + } + } + } + } + } + + result, _calls, cfg = self._launch(monkeypatch, managed=managed) + + assert result.exit_code == 0, result.output + assert cfg.call_args.kwargs["coding_agent_config_families"] == {"sonnet"} + def test_another_agent_keeps_its_own_model_and_is_told_why(self, monkeypatch): # A tier's model belongs to the tier's agent; pinning it on claude would land a Kimi id in # ANTHROPIC_MODEL, which the Anthropic-dialect endpoint cannot serve. diff --git a/tests/test_managed_resolve.py b/tests/test_managed_resolve.py index cf5d5cdb..97f4c0cd 100644 --- a/tests/test_managed_resolve.py +++ b/tests/test_managed_resolve.py @@ -236,7 +236,12 @@ def test_developers_state_file_keeps_their_own_model(self, real_state_file): def test_settings_file_gets_the_managed_model(self, real_state_file): # The other half of the contract: precedence must actually reach the generated file. resolved_state = resolve_state(MANAGED, state_mod.load_state(), "claude") - claude.write_tool_config(resolved_state, None) + claude.write_tool_config( + resolved_state, + None, + # MANAGED explicitly configures these three Claude family slots. + coding_agent_config_families={"opus", "sonnet", "haiku"}, + ) env = json.loads((real_state_file / "ucode-settings.json").read_text())["env"] assert env["ANTHROPIC_DEFAULT_OPUS_MODEL"].startswith("system.ai.claude-opus-5") From 63a0b2c24371f915ac43974417906c1022bd245e Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Wed, 2 Sep 2026 13:39:35 +0000 Subject: [PATCH 4/9] ruff --- tests/test_agent_claude.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 09deeab2..e1be3fbb 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -743,9 +743,7 @@ def test_managed_file_applies_model_default_precedence(self, monkeypatch): configured_families={"opus"}, ) - assert managed_env["ANTHROPIC_DEFAULT_OPUS_MODEL"] == ( - "system.ai.claude-opus-4-8[1m]" - ) + assert managed_env["ANTHROPIC_DEFAULT_OPUS_MODEL"] == ("system.ai.claude-opus-4-8[1m]") assert managed_env["ANTHROPIC_DEFAULT_SONNET_MODEL"] == "system.ai.claude-sonnet-4-6" assert managed_env["ANTHROPIC_DEFAULT_HAIKU_MODEL"] == "system.ai.claude-haiku-5" From fe612cb6a3977802942f359b2b0e7dffa9533c55 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Wed, 2 Sep 2026 14:07:32 +0000 Subject: [PATCH 5/9] update for fable --- src/ucode/agents/claude.py | 41 +++++++++++++++++++++----- tests/test_agent_claude.py | 59 +++++++++++++++++++++++++++++++++++++- 2 files changed, 92 insertions(+), 8 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 034e69db..6ca8eb8a 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -52,7 +52,7 @@ sync_smart_routing_hooks, ) from ucode.smart_routing.claude_routing import CLAUDE_VALUE_OPTIONS -from ucode.state import get_provider_service, mark_tool_managed, save_state +from ucode.state import MANAGED_OVERLAY_KEY, get_provider_service, mark_tool_managed, save_state from ucode.telemetry import agent_version, ucode_version from ucode.tracing import tracing_env from ucode.ui import print_note, print_success, print_warning @@ -642,19 +642,36 @@ def write_tool_config( # V2 installs routing hooks in a transient per-launch settings file. Persistent settings must # contain no ucode routing hooks; surgically strip legacy ones while preserving user hooks. - def _compose(base: dict, *, preserve_model_defaults: bool = False) -> dict: + def _compose(base: dict, *, enforce_model_default_hierarchy: bool) -> dict: base_env = base.get("env") existing_custom_headers = ( base_env.get(ANTHROPIC_CUSTOM_HEADERS_ENV_KEY) if isinstance(base_env, dict) else None ) # Copy the overlay per file so merging into one base cannot affect the other. overlay_for_merge = copy.deepcopy(overlay) - if preserve_model_defaults: + if enforce_model_default_hierarchy: + fable_key = CLAUDE_DEFAULT_MODEL_ENV_KEYS["fable"] + configured_families = coding_agent_config_families or set() + managed_overlay = state.get(MANAGED_OVERLAY_KEY) + local_models = ( + managed_overlay.get("claude_models") if isinstance(managed_overlay, dict) else None + ) + local_fable = local_models.get("fable") if isinstance(local_models, dict) else None + if ( + state.get("fable_enabled") + and "fable" not in configured_families + and fable_key not in overlay_for_merge["env"] + and isinstance(local_fable, str) + and local_fable + ): + # Managed resolution drops families omitted by Coding Agent Config. Restore the + # locally discovered Fable model only for the managed-file fallback. + overlay_for_merge["env"][fable_key] = local_fable # Default-model precedence in the enterprise-managed file, highest to lowest: # 1. Coding Agent Config default # 2. Existing default in the managed file # 3. Ucode-discovered default - configured_families = coding_agent_config_families or set() + # Fable participates only when explicitly enabled; otherwise it is removed. for family, key in CLAUDE_DEFAULT_MODEL_ENV_KEYS.items(): if ( family not in configured_families @@ -687,7 +704,12 @@ def _compose(base: dict, *, preserve_model_defaults: bool = False) -> dict: if isinstance(merged_env, dict): for key in CLAUDE_MANAGED_MODEL_ENV_KEYS: if key not in overlay_env and not ( - preserve_model_defaults and key in CLAUDE_DEFAULT_MODEL_ENV_KEYS.values() + enforce_model_default_hierarchy + and key in CLAUDE_DEFAULT_MODEL_ENV_KEYS.values() + and not ( + key == CLAUDE_DEFAULT_MODEL_ENV_KEYS["fable"] + and not state.get("fable_enabled") + ) ): merged_env.pop(key, None) # deep_merge_dict keeps keys already in the file, so drop the ones ucode no @@ -697,11 +719,16 @@ def _compose(base: dict, *, preserve_model_defaults: bool = False) -> dict: sync_smart_routing_hooks(merged, state, enabled=False) return merged - write_json_file(CLAUDE_SETTINGS_PATH, _compose(read_json_safe(CLAUDE_SETTINGS_PATH))) + write_json_file( + CLAUDE_SETTINGS_PATH, + _compose( + read_json_safe(CLAUDE_SETTINGS_PATH), enforce_model_default_hierarchy=False + ), + ) _reconcile_managed_settings( state, - lambda base: _compose(base, preserve_model_defaults=True), + lambda base: _compose(base, enforce_model_default_hierarchy=True), managed_file_keys, relayed, ) diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index e1be3fbb..90a57e0f 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -11,6 +11,7 @@ from ucode.agents import claude from ucode.smart_routing import claude_routing, v2 +from ucode.state import MANAGED_OVERLAY_KEY WS = "https://example.databricks.com" @@ -661,6 +662,8 @@ def _write_managed_model_defaults( existing: dict[str, str], discovered: dict[str, str], configured_families: set[str] | None = None, + fable_enabled: bool = False, + local_models: dict[str, str] | None = None, ) -> dict: private_writes: list = [] managed_writes: list = [] @@ -670,7 +673,14 @@ def _write_managed_model_defaults( managed_writes, {str(FAKE_MANAGED_PATH): {"env": existing}}, ) - state = {"workspace": WS, "codex_models": [], "claude_models": discovered} + state = { + "workspace": WS, + "codex_models": [], + "claude_models": discovered, + "fable_enabled": fable_enabled, + } + if local_models is not None: + state[MANAGED_OVERLAY_KEY] = {"claude_models": local_models} claude.write_tool_config( state, @@ -747,6 +757,53 @@ def test_managed_file_applies_model_default_precedence(self, monkeypatch): assert managed_env["ANTHROPIC_DEFAULT_SONNET_MODEL"] == "system.ai.claude-sonnet-4-6" assert managed_env["ANTHROPIC_DEFAULT_HAIKU_MODEL"] == "system.ai.claude-haiku-5" + def test_managed_file_removes_fable_default_when_fable_is_disabled(self, monkeypatch): + managed_env = self._write_managed_model_defaults( + monkeypatch, + existing={"ANTHROPIC_DEFAULT_FABLE_MODEL": "enterprise-fable"}, + discovered={"fable": "coding-agent-config-fable"}, + configured_families={"fable"}, + ) + + assert "ANTHROPIC_DEFAULT_FABLE_MODEL" not in managed_env + + @pytest.mark.parametrize( + ("existing", "discovered", "configured_families", "expected"), + [ + ("enterprise-fable", "coding-agent-config-fable", {"fable"}, "coding-agent-config-fable"), + ("enterprise-fable", "ucode-fable", set(), "enterprise-fable"), + (None, "ucode-fable", set(), "ucode-fable"), + ], + ) + def test_managed_file_applies_fable_precedence_when_enabled( + self, monkeypatch, existing, discovered, configured_families, expected + ): + managed_env = self._write_managed_model_defaults( + monkeypatch, + existing=( + {"ANTHROPIC_DEFAULT_FABLE_MODEL": existing} if existing is not None else {} + ), + discovered={"fable": discovered}, + configured_families=configured_families, + fable_enabled=True, + ) + + assert managed_env["ANTHROPIC_DEFAULT_FABLE_MODEL"] == expected + + def test_managed_file_uses_discovered_fable_when_coding_agent_config_omits_it( + self, monkeypatch + ): + managed_env = self._write_managed_model_defaults( + monkeypatch, + existing={}, + discovered={"opus": "coding-agent-config-opus"}, + configured_families={"opus"}, + fable_enabled=True, + local_models={"fable": "discovered-fable"}, + ) + + assert managed_env["ANTHROPIC_DEFAULT_FABLE_MODEL"] == "discovered-fable" + def test_managed_file_preserves_enterprise_permission_denies(self, monkeypatch): private_writes: list = [] managed_writes: list = [] From 58d37c4b13e66be95b6c14772012eb669b0f4f39 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Wed, 2 Sep 2026 14:19:12 +0000 Subject: [PATCH 6/9] update --- src/ucode/agents/claude.py | 98 +++++++++++++++++++++++++------------- tests/test_agent_claude.py | 5 +- 2 files changed, 69 insertions(+), 34 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 6ca8eb8a..ecb4968a 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -513,6 +513,34 @@ def _maybe_add_1m_suffix(model: str) -> str: return f"{model}[1m]" if should_suffix else model +def _enforce_model_default_hierarchy( + family: str, + *, + target_env: dict, + coding_agent_config_families: set[str], + settings_file_env: dict, + ucode_default_models: dict[str, str], +) -> None: + """Apply managed-file model precedence for one Claude family.""" + key = CLAUDE_DEFAULT_MODEL_ENV_KEYS[family] + coding_agent_config_default_model = ( + target_env.get(key) if family in coding_agent_config_families else None + ) + settings_file_existing_default_model = settings_file_env.get(key) + ucode_default_model = ucode_default_models.get(family) + + if coding_agent_config_default_model is not None: + selected_default_model = coding_agent_config_default_model + elif settings_file_existing_default_model is not None: + selected_default_model = settings_file_existing_default_model + else: + selected_default_model = ucode_default_model + if selected_default_model is None: + target_env.pop(key, None) + else: + target_env[key] = selected_default_model + + def _register_web_search_mcp(workspace: str, search_model: str, profile: str | None = None) -> bool: """Register (or replace) the web_search MCP server in Claude Code's user scope via `claude mcp add-json`. Removes any prior entry first so re-runs @@ -650,35 +678,48 @@ def _compose(base: dict, *, enforce_model_default_hierarchy: bool) -> dict: # Copy the overlay per file so merging into one base cannot affect the other. overlay_for_merge = copy.deepcopy(overlay) if enforce_model_default_hierarchy: - fable_key = CLAUDE_DEFAULT_MODEL_ENV_KEYS["fable"] configured_families = coding_agent_config_families or set() + settings_file_env = base_env if isinstance(base_env, dict) else {} + target_env = overlay_for_merge["env"] managed_overlay = state.get(MANAGED_OVERLAY_KEY) - local_models = ( + ucode_models = ( managed_overlay.get("claude_models") if isinstance(managed_overlay, dict) else None ) - local_fable = local_models.get("fable") if isinstance(local_models, dict) else None - if ( - state.get("fable_enabled") - and "fable" not in configured_families - and fable_key not in overlay_for_merge["env"] - and isinstance(local_fable, str) - and local_fable - ): - # Managed resolution drops families omitted by Coding Agent Config. Restore the - # locally discovered Fable model only for the managed-file fallback. - overlay_for_merge["env"][fable_key] = local_fable - # Default-model precedence in the enterprise-managed file, highest to lowest: - # 1. Coding Agent Config default - # 2. Existing default in the managed file - # 3. Ucode-discovered default - # Fable participates only when explicitly enabled; otherwise it is removed. + ucode_default_models: dict[str, str] = {} for family, key in CLAUDE_DEFAULT_MODEL_ENV_KEYS.items(): - if ( - family not in configured_families - and isinstance(base_env, dict) - and key in base_env - ): - overlay_for_merge["env"].pop(key, None) + if isinstance(ucode_models, dict): + model = ucode_models.get(family) + if isinstance(model, str) and model: + ucode_default_models[family] = ( + _maybe_add_1m_suffix(model) + if family in ("opus", "sonnet") + else model + ) + else: + model = target_env.get(key) + if isinstance(model, str) and model: + ucode_default_models[family] = model + + for family in ("opus", "sonnet", "haiku"): + _enforce_model_default_hierarchy( + family, + target_env=target_env, + coding_agent_config_families=configured_families, + settings_file_env=settings_file_env, + ucode_default_models=ucode_default_models, + ) + + # Fable uses the same hierarchy only when explicitly enabled. Otherwise, remove it. + if state.get("fable_enabled"): + _enforce_model_default_hierarchy( + "fable", + target_env=target_env, + coding_agent_config_families=configured_families, + settings_file_env=settings_file_env, + ucode_default_models=ucode_default_models, + ) + else: + target_env.pop(CLAUDE_DEFAULT_MODEL_ENV_KEYS["fable"], None) merged = deep_merge_dict(base, overlay_for_merge) overlay_custom_headers = overlay_for_merge["env"][ANTHROPIC_CUSTOM_HEADERS_ENV_KEY] merged["env"][ANTHROPIC_CUSTOM_HEADERS_ENV_KEY] = _merge_anthropic_custom_headers( @@ -703,14 +744,7 @@ def _compose(base: dict, *, enforce_model_default_hierarchy: bool) -> dict: merged_env = merged.get("env") if isinstance(merged_env, dict): for key in CLAUDE_MANAGED_MODEL_ENV_KEYS: - if key not in overlay_env and not ( - enforce_model_default_hierarchy - and key in CLAUDE_DEFAULT_MODEL_ENV_KEYS.values() - and not ( - key == CLAUDE_DEFAULT_MODEL_ENV_KEYS["fable"] - and not state.get("fable_enabled") - ) - ): + 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. diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 90a57e0f..d6582035 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -745,8 +745,9 @@ def test_managed_file_applies_model_default_precedence(self, monkeypatch): "ANTHROPIC_DEFAULT_OPUS_MODEL": "system.ai.claude-opus-5", "ANTHROPIC_DEFAULT_SONNET_MODEL": "system.ai.claude-sonnet-4-6", }, - discovered={ - "opus": "system.ai.claude-opus-4-8", + discovered={"opus": "system.ai.claude-opus-4-8"}, + local_models={ + "opus": "system.ai.claude-opus-5", "sonnet": "system.ai.claude-sonnet-5", "haiku": "system.ai.claude-haiku-5", }, From 60a5a1a2290fb4f2d5c2e29a496ca9e369a8a9bb Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Wed, 2 Sep 2026 15:43:31 +0000 Subject: [PATCH 7/9] test --- src/ucode/agents/__init__.py | 4 +- src/ucode/agents/claude.py | 90 ++++++++++++----------------- src/ucode/cli.py | 10 ++-- tests/test_agent_claude.py | 105 +++++++++++++--------------------- tests/test_cli.py | 6 +- tests/test_managed_resolve.py | 7 ++- 6 files changed, 92 insertions(+), 130 deletions(-) diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index 7765f4c4..0289b605 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -350,7 +350,7 @@ def configure_tool( relayed: bool = False, route_root_model: str | None = None, custom_model: str | None = None, - coding_agent_config_families: set[str] | None = None, + coding_agent_config_defaults: dict[str, str] | None = None, ) -> dict: result: dict | tuple[dict, str] if tool == "codex": @@ -369,7 +369,7 @@ def configure_tool( relayed=relayed, route_root_model=route_root_model, custom_model=custom_model, - coding_agent_config_families=coding_agent_config_families, + coding_agent_config_defaults=coding_agent_config_defaults, ) else: # provider routing is claude/codex-only; every other tool needs a model. diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index ecb4968a..a6cc6031 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -516,29 +516,27 @@ def _maybe_add_1m_suffix(model: str) -> str: def _enforce_model_default_hierarchy( family: str, *, - target_env: dict, - coding_agent_config_families: set[str], - settings_file_env: dict, - ucode_default_models: dict[str, str], -) -> None: + coding_agent_config_defaults: dict[str, str], + settings_file_existing_defaults: dict[str, str], + ucode_defaults: dict[str, str], +) -> str | None: """Apply managed-file model precedence for one Claude family.""" - key = CLAUDE_DEFAULT_MODEL_ENV_KEYS[family] - coding_agent_config_default_model = ( - target_env.get(key) if family in coding_agent_config_families else None - ) - settings_file_existing_default_model = settings_file_env.get(key) - ucode_default_model = ucode_default_models.get(family) + coding_agent_config_default_model = coding_agent_config_defaults.get(family) + settings_file_existing_default_model = settings_file_existing_defaults.get(family) + ucode_default_model = ucode_defaults.get(family) if coding_agent_config_default_model is not None: selected_default_model = coding_agent_config_default_model elif settings_file_existing_default_model is not None: - selected_default_model = settings_file_existing_default_model + return settings_file_existing_default_model else: selected_default_model = ucode_default_model + if selected_default_model is None: - target_env.pop(key, None) - else: - target_env[key] = selected_default_model + return None + if family in ("opus", "sonnet"): + return _maybe_add_1m_suffix(selected_default_model) + return selected_default_model def _register_web_search_mcp(workspace: str, search_model: str, profile: str | None = None) -> bool: @@ -621,7 +619,7 @@ def write_tool_config( relayed: bool = False, route_root_model: str | None = None, custom_model: str | None = None, - coding_agent_config_families: set[str] | None = None, + coding_agent_config_defaults: dict[str, str] | None = None, ) -> dict: backup_existing_file(CLAUDE_SETTINGS_PATH, CLAUDE_BACKUP_PATH) web_search_model = _resolve_web_search_model(state) @@ -678,48 +676,34 @@ def _compose(base: dict, *, enforce_model_default_hierarchy: bool) -> dict: # Copy the overlay per file so merging into one base cannot affect the other. overlay_for_merge = copy.deepcopy(overlay) if enforce_model_default_hierarchy: - configured_families = coding_agent_config_families or set() settings_file_env = base_env if isinstance(base_env, dict) else {} target_env = overlay_for_merge["env"] - managed_overlay = state.get(MANAGED_OVERLAY_KEY) - ucode_models = ( - managed_overlay.get("claude_models") if isinstance(managed_overlay, dict) else None + configured_defaults = coding_agent_config_defaults or {} + settings_file_existing_defaults = { + family: model + for family, key in CLAUDE_DEFAULT_MODEL_ENV_KEYS.items() + if isinstance((model := settings_file_env.get(key)), str) + } + managed_overlay = state.get(MANAGED_OVERLAY_KEY, {}) + ucode_defaults = managed_overlay.get( + "claude_models", state.get("claude_models", {}) ) - ucode_default_models: dict[str, str] = {} + for family, key in CLAUDE_DEFAULT_MODEL_ENV_KEYS.items(): - if isinstance(ucode_models, dict): - model = ucode_models.get(family) - if isinstance(model, str) and model: - ucode_default_models[family] = ( - _maybe_add_1m_suffix(model) - if family in ("opus", "sonnet") - else model - ) - else: - model = target_env.get(key) - if isinstance(model, str) and model: - ucode_default_models[family] = model + if family == "fable" and not state.get("fable_enabled"): + target_env.pop(key, None) + continue - for family in ("opus", "sonnet", "haiku"): - _enforce_model_default_hierarchy( + selected_default_model = _enforce_model_default_hierarchy( family, - target_env=target_env, - coding_agent_config_families=configured_families, - settings_file_env=settings_file_env, - ucode_default_models=ucode_default_models, - ) - - # Fable uses the same hierarchy only when explicitly enabled. Otherwise, remove it. - if state.get("fable_enabled"): - _enforce_model_default_hierarchy( - "fable", - target_env=target_env, - coding_agent_config_families=configured_families, - settings_file_env=settings_file_env, - ucode_default_models=ucode_default_models, + coding_agent_config_defaults=configured_defaults, + settings_file_existing_defaults=settings_file_existing_defaults, + ucode_defaults=ucode_defaults, ) - else: - target_env.pop(CLAUDE_DEFAULT_MODEL_ENV_KEYS["fable"], None) + if selected_default_model is None: + target_env.pop(key, None) + else: + target_env[key] = selected_default_model merged = deep_merge_dict(base, overlay_for_merge) overlay_custom_headers = overlay_for_merge["env"][ANTHROPIC_CUSTOM_HEADERS_ENV_KEY] merged["env"][ANTHROPIC_CUSTOM_HEADERS_ENV_KEY] = _merge_anthropic_custom_headers( @@ -755,9 +739,7 @@ def _compose(base: dict, *, enforce_model_default_hierarchy: bool) -> dict: write_json_file( CLAUDE_SETTINGS_PATH, - _compose( - read_json_safe(CLAUDE_SETTINGS_PATH), enforce_model_default_hierarchy=False - ), + _compose(read_json_safe(CLAUDE_SETTINGS_PATH), enforce_model_default_hierarchy=False), ) _reconcile_managed_settings( diff --git a/src/ucode/cli.py b/src/ucode/cli.py index e7594446..8e6f9b7a 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -2016,10 +2016,10 @@ def _launch_tool( # or Foundry service), and, for Bedrock, expose Claude models to pin. provider_models = None relayed = False - coding_agent_config_families = ( - set(managed_claude_family_models(managed) or {}) + coding_agent_config_defaults = ( + managed_claude_family_models(managed) or {} if tool == "claude" and managed is not None - else set() + else {} ) if provider: provider_models, error, relayed = resolve_provider_models(tool, state, provider) @@ -2044,7 +2044,7 @@ def _launch_tool( authored = managed_provider_family_models(managed) if authored: provider_models = authored - coding_agent_config_families = set(authored) + coding_agent_config_defaults = authored # The router's per-launch pick for the root session. Codex pins it as the # resolved model; claude pins it via ANTHROPIC_MODEL (route_root_model). route_root_model = None @@ -2093,7 +2093,7 @@ def _launch_tool( # the latter pins a raw id into every family alias, which would clobber the service's # per-family target pins. custom_model=model if (tool == "claude" and not provider) else None, - coding_agent_config_families=coding_agent_config_families, + coding_agent_config_defaults=coding_agent_config_defaults, ) # Relayed = a Claude subscription: forward --model to Claude Code's own flag, like `-- --model X`. if tool == "claude" and provider and relayed and model and not forwarded_model: diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index d6582035..9b925e2e 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -659,37 +659,46 @@ def _write_managed_model_defaults( self, monkeypatch, *, - existing: dict[str, str], - discovered: dict[str, str], - configured_families: set[str] | None = None, - fable_enabled: bool = False, - local_models: dict[str, str] | None = None, - ) -> dict: + coding_agent_config_defaults: dict[str, str], + managed_settings_defaults: dict[str, str], + ucode_defaults: dict[str, str], + fable_enabled: bool, + ) -> dict[str, str]: private_writes: list = [] managed_writes: list = [] + managed_settings_env = { + claude.CLAUDE_DEFAULT_MODEL_ENV_KEYS[family]: model + for family, model in managed_settings_defaults.items() + } self._patch( monkeypatch, private_writes, managed_writes, - {str(FAKE_MANAGED_PATH): {"env": existing}}, + {str(FAKE_MANAGED_PATH): {"env": managed_settings_env}}, ) + resolved_defaults = coding_agent_config_defaults or ucode_defaults state = { "workspace": WS, "codex_models": [], - "claude_models": discovered, + "claude_models": resolved_defaults, "fable_enabled": fable_enabled, } - if local_models is not None: - state[MANAGED_OVERLAY_KEY] = {"claude_models": local_models} + if coding_agent_config_defaults: + state[MANAGED_OVERLAY_KEY] = {"claude_models": ucode_defaults} claude.write_tool_config( state, - next(iter(discovered.values())), - coding_agent_config_families=configured_families, + next(iter(resolved_defaults.values()), "test-model"), + coding_agent_config_defaults=coding_agent_config_defaults, ) _, text = managed_writes[0] - return json.loads(text)["env"] + written_env = json.loads(text)["env"] + return { + family: written_env[key] + for family, key in claude.CLAUDE_DEFAULT_MODEL_ENV_KEYS.items() + if key in written_env + } def test_writes_managed_file_by_default(self, monkeypatch): private_writes: list = [] @@ -739,71 +748,37 @@ def test_managed_file_merges_anthropic_custom_headers(self, monkeypatch): ] def test_managed_file_applies_model_default_precedence(self, monkeypatch): - managed_env = self._write_managed_model_defaults( + managed_defaults = self._write_managed_model_defaults( monkeypatch, - existing={ - "ANTHROPIC_DEFAULT_OPUS_MODEL": "system.ai.claude-opus-5", - "ANTHROPIC_DEFAULT_SONNET_MODEL": "system.ai.claude-sonnet-4-6", + coding_agent_config_defaults={"opus": "system.ai.claude-opus-4-8"}, + managed_settings_defaults={ + "opus": "system.ai.claude-opus-5", + "sonnet": "system.ai.claude-sonnet-4-6", }, - discovered={"opus": "system.ai.claude-opus-4-8"}, - local_models={ + ucode_defaults={ "opus": "system.ai.claude-opus-5", "sonnet": "system.ai.claude-sonnet-5", "haiku": "system.ai.claude-haiku-5", }, - configured_families={"opus"}, + fable_enabled=False, ) - assert managed_env["ANTHROPIC_DEFAULT_OPUS_MODEL"] == ("system.ai.claude-opus-4-8[1m]") - assert managed_env["ANTHROPIC_DEFAULT_SONNET_MODEL"] == "system.ai.claude-sonnet-4-6" - assert managed_env["ANTHROPIC_DEFAULT_HAIKU_MODEL"] == "system.ai.claude-haiku-5" + assert managed_defaults == { + "opus": "system.ai.claude-opus-4-8[1m]", + "sonnet": "system.ai.claude-sonnet-4-6", + "haiku": "system.ai.claude-haiku-5", + } def test_managed_file_removes_fable_default_when_fable_is_disabled(self, monkeypatch): - managed_env = self._write_managed_model_defaults( - monkeypatch, - existing={"ANTHROPIC_DEFAULT_FABLE_MODEL": "enterprise-fable"}, - discovered={"fable": "coding-agent-config-fable"}, - configured_families={"fable"}, - ) - - assert "ANTHROPIC_DEFAULT_FABLE_MODEL" not in managed_env - - @pytest.mark.parametrize( - ("existing", "discovered", "configured_families", "expected"), - [ - ("enterprise-fable", "coding-agent-config-fable", {"fable"}, "coding-agent-config-fable"), - ("enterprise-fable", "ucode-fable", set(), "enterprise-fable"), - (None, "ucode-fable", set(), "ucode-fable"), - ], - ) - def test_managed_file_applies_fable_precedence_when_enabled( - self, monkeypatch, existing, discovered, configured_families, expected - ): - managed_env = self._write_managed_model_defaults( - monkeypatch, - existing=( - {"ANTHROPIC_DEFAULT_FABLE_MODEL": existing} if existing is not None else {} - ), - discovered={"fable": discovered}, - configured_families=configured_families, - fable_enabled=True, - ) - - assert managed_env["ANTHROPIC_DEFAULT_FABLE_MODEL"] == expected - - def test_managed_file_uses_discovered_fable_when_coding_agent_config_omits_it( - self, monkeypatch - ): - managed_env = self._write_managed_model_defaults( + managed_defaults = self._write_managed_model_defaults( monkeypatch, - existing={}, - discovered={"opus": "coding-agent-config-opus"}, - configured_families={"opus"}, - fable_enabled=True, - local_models={"fable": "discovered-fable"}, + coding_agent_config_defaults={"fable": "coding-agent-config-fable"}, + managed_settings_defaults={"fable": "managed-settings-fable"}, + ucode_defaults={"fable": "ucode-fable"}, + fable_enabled=False, ) - assert managed_env["ANTHROPIC_DEFAULT_FABLE_MODEL"] == "discovered-fable" + assert "fable" not in managed_defaults def test_managed_file_preserves_enterprise_permission_denies(self, monkeypatch): private_writes: list = [] diff --git a/tests/test_cli.py b/tests/test_cli.py index 5809f492..a867e8bd 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -3711,7 +3711,7 @@ def test_the_recommended_agent_gets_the_recommended_model(self, monkeypatch): ) assert cfg.call_args.args[2] == "system.ai.claude-haiku-4-5" - def test_passes_configured_claude_families_to_writer(self, monkeypatch): + def test_passes_configured_claude_defaults_to_writer(self, monkeypatch): monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") managed = { "enabled_agents": { @@ -3728,7 +3728,9 @@ def test_passes_configured_claude_families_to_writer(self, monkeypatch): result, _calls, cfg = self._launch(monkeypatch, managed=managed) assert result.exit_code == 0, result.output - assert cfg.call_args.kwargs["coding_agent_config_families"] == {"sonnet"} + assert cfg.call_args.kwargs["coding_agent_config_defaults"] == { + "sonnet": "system.ai.claude-sonnet-4-6" + } def test_another_agent_keeps_its_own_model_and_is_told_why(self, monkeypatch): # A tier's model belongs to the tier's agent; pinning it on claude would land a Kimi id in diff --git a/tests/test_managed_resolve.py b/tests/test_managed_resolve.py index 97f4c0cd..96d137fb 100644 --- a/tests/test_managed_resolve.py +++ b/tests/test_managed_resolve.py @@ -239,8 +239,11 @@ def test_settings_file_gets_the_managed_model(self, real_state_file): claude.write_tool_config( resolved_state, None, - # MANAGED explicitly configures these three Claude family slots. - coding_agent_config_families={"opus", "sonnet", "haiku"}, + coding_agent_config_defaults={ + "opus": "system.ai.claude-opus-5", + "sonnet": "system.ai.claude-sonnet-4-6", + "haiku": "system.ai.claude-haiku-4-5", + }, ) env = json.loads((real_state_file / "ucode-settings.json").read_text())["env"] From ae53bdfa9b50357318bcc04943c7b8c9e3c4985b Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Wed, 2 Sep 2026 15:48:35 +0000 Subject: [PATCH 8/9] update --- src/ucode/agents/claude.py | 4 +--- tests/test_agent_claude.py | 6 +++--- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index a6cc6031..5f179cbd 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -685,9 +685,7 @@ def _compose(base: dict, *, enforce_model_default_hierarchy: bool) -> dict: if isinstance((model := settings_file_env.get(key)), str) } managed_overlay = state.get(MANAGED_OVERLAY_KEY, {}) - ucode_defaults = managed_overlay.get( - "claude_models", state.get("claude_models", {}) - ) + ucode_defaults = managed_overlay.get("claude_models", state.get("claude_models", {})) for family, key in CLAUDE_DEFAULT_MODEL_ENV_KEYS.items(): if family == "fable" and not state.get("fable_enabled"): diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 9b925e2e..513ae65a 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -764,9 +764,9 @@ def test_managed_file_applies_model_default_precedence(self, monkeypatch): ) assert managed_defaults == { - "opus": "system.ai.claude-opus-4-8[1m]", - "sonnet": "system.ai.claude-sonnet-4-6", - "haiku": "system.ai.claude-haiku-5", + "opus": "system.ai.claude-opus-4-8[1m]", # Coding Agent Config took priority. + "sonnet": "system.ai.claude-sonnet-4-6", # Existing managed setting took priority. + "haiku": "system.ai.claude-haiku-5", # Ucode default took priority. } def test_managed_file_removes_fable_default_when_fable_is_disabled(self, monkeypatch): From 1407d29c21cb07aed24caa2de890c7787ca4569b Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Wed, 2 Sep 2026 18:00:26 +0000 Subject: [PATCH 9/9] clean up --- src/ucode/agents/claude.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 81b8e46e..3e773d70 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -690,7 +690,9 @@ def _compose(base: dict, *, enforce_model_default_hierarchy: bool) -> dict: if isinstance((model := settings_file_env.get(key)), str) } managed_overlay = state.get(MANAGED_OVERLAY_KEY, {}) - ucode_defaults = managed_overlay.get("claude_models", state.get("claude_models", {})) + ucode_defaults = ( + managed_overlay.get("claude_models") or state.get("claude_models") or {} + ) for family, key in CLAUDE_DEFAULT_MODEL_ENV_KEYS.items(): if family == "fable" and not state.get("fable_enabled"):