diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index 578aa208..0289b605 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_defaults: dict[str, 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_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 5a98c9c6..01b3a1f9 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -49,7 +49,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 @@ -176,9 +176,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 in the enterprise-managed file +# are preserved unless Coding Agent Config explicitly supplies that family. CLAUDE_MANAGED_MODEL_ENV_KEYS = ( "ANTHROPIC_MODEL", "ANTHROPIC_DEFAULT_FABLE_MODEL", @@ -190,6 +189,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", +} # 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",) @@ -509,6 +514,32 @@ def _maybe_add_1m_suffix(model: str) -> str: return f"{model}[1m]" if should_suffix else model +def _enforce_model_default_hierarchy( + family: str, + *, + 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.""" + 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: + return settings_file_existing_default_model + else: + selected_default_model = ucode_default_model + + if selected_default_model is None: + 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: """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 @@ -589,6 +620,7 @@ def write_tool_config( relayed: bool = False, route_root_model: str | None = None, custom_model: 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) @@ -638,15 +670,44 @@ 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: + 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 ) - # 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) + if enforce_model_default_hierarchy: + settings_file_env = base_env if isinstance(base_env, dict) else {} + target_env = overlay_for_merge["env"] + 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") 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"): + target_env.pop(key, None) + continue + + selected_default_model = _enforce_model_default_hierarchy( + family, + coding_agent_config_defaults=configured_defaults, + settings_file_existing_defaults=settings_file_existing_defaults, + ucode_defaults=ucode_defaults, + ) + 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( existing_custom_headers, overlay_custom_headers ) @@ -665,7 +726,7 @@ 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: @@ -681,9 +742,17 @@ def _compose(base: dict) -> 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, _compose, managed_file_keys, relayed) + _reconcile_managed_settings( + state, + lambda base: _compose(base, enforce_model_default_hierarchy=True), + managed_file_keys, + relayed, + ) if web_search_model: web_search_entry = _web_search_mcp_entry( diff --git a/src/ucode/cli.py b/src/ucode/cli.py index d0789093..6134c9f0 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, @@ -2017,6 +2018,11 @@ def _launch_tool( # or Foundry service), and, for Bedrock, expose Claude models to pin. provider_models = None relayed = False + coding_agent_config_defaults = ( + managed_claude_family_models(managed) or {} + if tool == "claude" and managed is not None + else {} + ) if provider: provider_models, error, relayed = resolve_provider_models(tool, state, provider) if error: @@ -2040,6 +2046,7 @@ def _launch_tool( authored = managed_provider_family_models(managed) if authored: provider_models = 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 @@ -2088,6 +2095,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_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/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 f58fed37..6aa494e8 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" @@ -674,6 +675,51 @@ def fake_write_managed(path, text, **kwargs): monkeypatch.setattr(claude, "reconcile_managed_file", fake_write_managed) + def _write_managed_model_defaults( + self, + monkeypatch, + *, + 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": managed_settings_env}}, + ) + resolved_defaults = coding_agent_config_defaults or ucode_defaults + state = { + "workspace": WS, + "codex_models": [], + "claude_models": resolved_defaults, + "fable_enabled": fable_enabled, + } + if coding_agent_config_defaults: + state[MANAGED_OVERLAY_KEY] = {"claude_models": ucode_defaults} + + claude.write_tool_config( + state, + next(iter(resolved_defaults.values()), "test-model"), + coding_agent_config_defaults=coding_agent_config_defaults, + ) + + _, text = managed_writes[0] + 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 = [] managed_writes: list = [] @@ -741,6 +787,39 @@ 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_applies_model_default_precedence(self, monkeypatch): + managed_defaults = self._write_managed_model_defaults( + monkeypatch, + 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", + }, + ucode_defaults={ + "opus": "system.ai.claude-opus-5", + "sonnet": "system.ai.claude-sonnet-5", + "haiku": "system.ai.claude-haiku-5", + }, + fable_enabled=False, + ) + + assert managed_defaults == { + "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): + managed_defaults = self._write_managed_model_defaults( + monkeypatch, + 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 "fable" not in managed_defaults + def test_managed_file_preserves_enterprise_permission_denies(self, monkeypatch): private_writes: list = [] managed_writes: list = [] @@ -1167,8 +1246,6 @@ def test_prunes_stale_anthropic_model_from_prior_run(self, monkeypatch): 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]", diff --git a/tests/test_cli.py b/tests/test_cli.py index a59f3015..3f201b43 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -3831,6 +3831,27 @@ 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_defaults_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_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 # 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..96d137fb 100644 --- a/tests/test_managed_resolve.py +++ b/tests/test_managed_resolve.py @@ -236,7 +236,15 @@ 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, + 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"] assert env["ANTHROPIC_DEFAULT_OPUS_MODEL"].startswith("system.ai.claude-opus-5")