From ce0099035fa9f6def18206d8c0d93fef3d38c025 Mon Sep 17 00:00:00 2001 From: John Speaks Date: Mon, 5 Oct 2026 16:40:26 +0000 Subject: [PATCH] Honor managed Claude aliases and Codex defaults --- src/ucode/agents/args.py | 1 + src/ucode/agents/claude.py | 68 ++++++++++++++++-- src/ucode/agents/codex.py | 1 - src/ucode/cli.py | 9 +++ tests/test_agent_claude.py | 140 +++++++++++++++++++++++++++++++++++++ tests/test_agent_codex.py | 29 +++++++- tests/test_cli.py | 41 +++++++++++ 7 files changed, 280 insertions(+), 9 deletions(-) diff --git a/src/ucode/agents/args.py b/src/ucode/agents/args.py index faca49dfe..29f9d92c3 100644 --- a/src/ucode/agents/args.py +++ b/src/ucode/agents/args.py @@ -11,6 +11,7 @@ class LaunchOptions: launch_smart_routing: bool = False user_pinned_model: str | None = None + managed_claude_model: str | None = None def explicit_model_arg_value(tool_args: list[str]) -> str | None: diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 83910d65d..70efc7ec9 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -1266,6 +1266,12 @@ def write_tool_config( # this launch's warm managed-config cache (no extra round trip); a failed fetch degrades to None # (treated as unmanaged), never blocking the write. managed_config_present = refresh_managed_config(state).manifest is not None + if coding_agent_config_defaults: + managed_overlay = state.get(MANAGED_OVERLAY_KEY) + if not isinstance(managed_overlay, dict): + managed_overlay = {} + state[MANAGED_OVERLAY_KEY] = managed_overlay + managed_overlay.setdefault("claude_models", state.get("claude_models")) previous_keys = ((state.get("managed_configs") or {}).get("claude") or {}).get("keys", []) web_search_model = _resolve_web_search_model(state) should_write_tracing_settings = ( @@ -1827,6 +1833,47 @@ def _launch_model_args(tool_args: list[str], launch_model: str | None) -> list[s return ["--model", launch_model] +def _managed_family_alias_model(state: dict, model: str | None) -> str | None: + """Resolve a managed family alias without losing its context-window suffix.""" + family_env_key = CLAUDE_DEFAULT_MODEL_ENV_KEYS.get(model or "") + if family_env_key is None: + return None + provider_services = state.get("provider_services") + if state.get("_claude_launch_provider") or ( + not state.get("_claude_launch_parent_schema") + and isinstance(provider_services, dict) + and provider_services.get("claude") + ): + return None + overlay = state.get(MANAGED_OVERLAY_KEY) + if not isinstance(overlay, dict): + return None + if "claude_models" not in overlay and "claude_static_models" not in overlay: + return None + env = read_json_safe(CLAUDE_SETTINGS_PATH).get("env") + if not isinstance(env, dict): + return None + configured_model = env.get(family_env_key) + return configured_model if isinstance(configured_model, str) and configured_model else None + + +def _without_model_args(tool_args: list[str]) -> list[str]: + """Remove forwarded Claude model flags after a managed family alias is resolved.""" + remaining: list[str] = [] + index = 0 + while index < len(tool_args): + arg = tool_args[index] + if arg in {"--model", "-m"} and index + 1 < len(tool_args): + index += 2 + continue + if arg.startswith("--model="): + index += 1 + continue + remaining.append(arg) + index += 1 + return remaining + + def _resolve_picker_model_id(model: str, settings_env: dict) -> str: """Resolve configured aliases and context suffixes for comparisons only.""" model = re.sub(r"\[(?:1m|200k)\]$", "", model) @@ -2064,12 +2111,21 @@ def launch( settings_override = None launch_args = list(tool_args) if options.user_pinned_model: - os.environ["ANTHROPIC_MODEL"] = options.user_pinned_model - settings_override = {"env": {"ANTHROPIC_MODEL": options.user_pinned_model}} - launch_args = [ - *_launch_model_args(tool_args, options.user_pinned_model), - *tool_args, - ] + managed_alias_model = options.managed_claude_model or _managed_family_alias_model( + state, options.user_pinned_model + ) + launch_model = managed_alias_model or options.user_pinned_model + if options.managed_claude_model and not launch_model.endswith(("[1m]", "[200k]")): + launch_model = _maybe_add_1m_suffix(launch_model) + os.environ["ANTHROPIC_MODEL"] = launch_model + settings_override = {"env": {"ANTHROPIC_MODEL": launch_model}} + if managed_alias_model: + launch_args = _without_model_args(tool_args) + else: + launch_args = [ + *_launch_model_args(tool_args, options.user_pinned_model), + *tool_args, + ] else: picker_models = state.get("_claude_launch_picker_models") if isinstance(picker_models, list) and picker_models: diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 75606811d..8b0bcd3b0 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -788,7 +788,6 @@ def default_model(state: dict) -> str | None: return state["codex_default_model"] if smart_routing_v2.smart_routing_enabled(): return _smart_routing_config_model(state) - clear_model_preferences(state) return None diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 396d93791..e51b0a390 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -2551,10 +2551,12 @@ def _launch_options( explicit_prompt: bool, user_pinned_model: str | None, provider: str | None, + managed_claude_model: str | None = None, ) -> LaunchOptions: return LaunchOptions( # Pinned models for providers are resolved above through the provider-specific launch path. user_pinned_model=user_pinned_model if provider is None else None, + managed_claude_model=managed_claude_model if provider is None else None, launch_smart_routing=( # Smart routing is enabled globally. smart_routing_enabled @@ -2928,6 +2930,8 @@ def _launch_tool( if tool == "claude": if provider: state["_claude_launch_provider"] = provider + elif parent_schema: + state["_claude_launch_parent_schema"] = parent_schema elif tool == "codex": if provider: state["_codex_launch_provider"] = provider @@ -2942,6 +2946,11 @@ def _launch_tool( # initial/fallback model and still participates in a routed session. user_pinned_model=model or forwarded_model, provider=provider, + managed_claude_model=( + coding_agent_config_defaults.get(model) + if tool == "claude" and model and coding_agent_config_defaults + else None + ), ) print_success(f"Starting {TOOL_SPECS[tool]['display']}") with _smart_routing_v2_flag( diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 148543d1e..53723ee70 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -2482,6 +2482,146 @@ def test_launch_model_is_only_set_for_current_process(self, monkeypatch): settings = json.loads(calls[0][2]) assert settings["env"]["ANTHROPIC_MODEL"] == "cat.schema.model" + def test_managed_family_alias_uses_configured_model_without_native_alias( + self, monkeypatch, tmp_path + ): + calls: list[list[str]] = [] + settings_path = tmp_path / "ucode-settings.json" + settings_path.write_text( + json.dumps({"env": {"ANTHROPIC_DEFAULT_OPUS_MODEL": "system.ai.claude-opus-4-8[1m]"}}) + ) + monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", settings_path) + monkeypatch.delenv("ANTHROPIC_MODEL", raising=False) + monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "token") + monkeypatch.setattr(claude, "exec_or_spawn", lambda argv: calls.append(argv)) + + claude.launch( + { + "workspace": WS, + MANAGED_OVERLAY_KEY: {"claude_models": {"opus": "developer-opus"}}, + }, + [], + options=LaunchOptions(user_pinned_model="opus"), + ) + + assert os.environ["ANTHROPIC_MODEL"] == "system.ai.claude-opus-4-8[1m]" + launch_settings = json.loads(calls[0][2]) + assert launch_settings["env"]["ANTHROPIC_MODEL"] == "system.ai.claude-opus-4-8[1m]" + assert "--model" not in calls[0] + + def test_launch_boundary_managed_family_alias_uses_actual_model(self, monkeypatch): + calls: list[list[str]] = [] + monkeypatch.delenv("ANTHROPIC_MODEL", raising=False) + monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "token") + monkeypatch.setattr(claude, "exec_or_spawn", lambda argv: calls.append(argv)) + + claude.launch( + {"workspace": WS}, + [], + options=LaunchOptions( + user_pinned_model="opus", + managed_claude_model="system.ai.claude-opus-4-8", + ), + ) + + assert os.environ["ANTHROPIC_MODEL"] == "system.ai.claude-opus-4-8[1m]" + launch_settings = json.loads(calls[0][2]) + assert launch_settings["env"]["ANTHROPIC_MODEL"] == "system.ai.claude-opus-4-8[1m]" + assert "--model" not in calls[0] + + def test_managed_parent_overrides_stale_provider_for_forwarded_alias( + self, monkeypatch, tmp_path + ): + calls: list[list[str]] = [] + settings_path = tmp_path / "ucode-settings.json" + settings_path.write_text( + json.dumps({"env": {"ANTHROPIC_DEFAULT_OPUS_MODEL": "system.ai.claude-opus-4-8[1m]"}}) + ) + monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", settings_path) + monkeypatch.delenv("ANTHROPIC_MODEL", raising=False) + monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "token") + monkeypatch.setattr(claude, "exec_or_spawn", lambda argv: calls.append(argv)) + + claude.launch( + { + "workspace": WS, + "provider_services": {"claude": "main.default.stale"}, + "_claude_launch_parent_schema": "main.models", + MANAGED_OVERLAY_KEY: {"claude_models": {"opus": "developer-opus"}}, + }, + ["--model", "opus"], + options=LaunchOptions(user_pinned_model="opus"), + ) + + assert os.environ["ANTHROPIC_MODEL"] == "system.ai.claude-opus-4-8[1m]" + launch_settings = json.loads(calls[0][2]) + assert launch_settings["env"]["ANTHROPIC_MODEL"] == "system.ai.claude-opus-4-8[1m]" + assert "--model" not in calls[0] + + def test_equal_managed_family_value_still_marks_managed_alias(self, monkeypatch, tmp_path): + calls: list[list[str]] = [] + settings_path = tmp_path / "ucode-settings.json" + monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", settings_path) + monkeypatch.setattr(claude, "CLAUDE_BACKUP_PATH", tmp_path / "backup.json") + monkeypatch.setattr(claude, "save_state", lambda state: None) + monkeypatch.setattr( + claude, + "write_json_file", + lambda path, payload: path.write_text(json.dumps(payload)), + ) + monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "token") + monkeypatch.setattr(claude, "exec_or_spawn", lambda argv: calls.append(argv)) + state = { + "workspace": WS, + "claude_models": {"opus": "system.ai.claude-opus-4-8"}, + } + + claude.write_tool_config( + state, + "system.ai.claude-opus-4-8", + coding_agent_config_defaults={"opus": "system.ai.claude-opus-4-8"}, + ) + claude.launch(state, [], options=LaunchOptions(user_pinned_model="opus")) + + assert state[MANAGED_OVERLAY_KEY]["claude_models"] == state["claude_models"] + launch_settings = json.loads(calls[0][2]) + assert launch_settings["env"]["ANTHROPIC_MODEL"] == "system.ai.claude-opus-4-8[1m]" + assert "--model" not in calls[0] + + def test_unmanaged_family_alias_keeps_native_alias(self, monkeypatch): + calls: list[list[str]] = [] + monkeypatch.delenv("ANTHROPIC_MODEL", raising=False) + monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "token") + monkeypatch.setattr(claude, "exec_or_spawn", lambda argv: calls.append(argv)) + + claude.launch( + {"workspace": WS}, + [], + options=LaunchOptions(user_pinned_model="opus"), + ) + + assert os.environ["ANTHROPIC_MODEL"] == "opus" + assert calls[0][-2:] == ["--model", "opus"] + + def test_provider_family_alias_keeps_native_alias(self, monkeypatch): + calls: list[list[str]] = [] + monkeypatch.delenv("ANTHROPIC_MODEL", raising=False) + monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "token") + monkeypatch.setattr(claude, "exec_or_spawn", lambda argv: calls.append(argv)) + + claude.launch( + { + "workspace": WS, + "_claude_launch_provider": "main.default.anthropic", + MANAGED_OVERLAY_KEY: {"claude_models": {"opus": "managed-opus"}}, + }, + [], + options=LaunchOptions(user_pinned_model="opus"), + ) + + assert os.environ["ANTHROPIC_MODEL"] == "opus" + assert calls[0][-2:] == ["--model", "opus"] + def test_launch_default_model_is_inherited_by_smart_routing(self, monkeypatch): monkeypatch.delenv("ANTHROPIC_DEFAULT_MODEL", raising=False) monkeypatch.setattr(v2, "launch_claude", Mock()) diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index c810a1b6f..51f85abb9 100644 --- a/tests/test_agent_codex.py +++ b/tests/test_agent_codex.py @@ -9,9 +9,11 @@ import pytest +import ucode.state as state_mod from ucode import managed_files from ucode.agents import LaunchOptions, codex from ucode.config_io import read_toml_safe +from ucode.managed_resolve import resolve_state from ucode.smart_routing import codex_routing WS = "https://example.databricks.com" @@ -909,16 +911,39 @@ def _isolate_profile_config(self, tmp_path, monkeypatch): monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", tmp_path / "ucode.config.toml") monkeypatch.setattr(codex, "CODEX_BACKUP_PATH", tmp_path / "backup.toml") - def test_clears_profile_model_preferences(self, tmp_path): + def test_configure_clears_profile_model_preferences(self, tmp_path, monkeypatch): codex.CODEX_CONFIG_PATH.write_text( 'model = "gpt-5.6-sol"\nmodel_reasoning_effort = "medium"\n', encoding="utf-8" ) + monkeypatch.setattr(codex, "agent_version", lambda _: "0.145.0") + monkeypatch.setattr(codex, "save_state", lambda _: None) + + codex.write_tool_config({"workspace": WS, "codex_models": ["system.ai.gpt-5-6-luna"]}) - assert codex.default_model({"codex_models": ["system.ai.gpt-5-6-luna"]}) is None doc = read_toml_safe(codex.CODEX_CONFIG_PATH) assert "model" not in doc assert "model_reasoning_effort" not in doc + def test_managed_default_survives_saved_state_hydration(self, tmp_path, monkeypatch): + monkeypatch.setattr(codex, "agent_version", lambda _: "0.145.0") + state_dir = tmp_path / ".ucode" + monkeypatch.setattr(state_mod, "APP_DIR", state_dir) + monkeypatch.setattr(state_mod, "STATE_PATH", state_dir / "state.json") + developer_state = { + "workspace": WS, + "codex_models": ["system.ai.gpt-5-6-luna"], + } + managed = { + "enabled_agents": {"codex": {"model_config": {"default_model": "admin-chosen-default"}}} + } + resolved = resolve_state(managed, developer_state, "codex") + + codex.write_tool_config(resolved) + + assert read_toml_safe(codex.CODEX_CONFIG_PATH)["model"] == "admin-chosen-default" + persisted = state_mod.load_full_state()["workspaces"][WS] + assert "codex_default_model" not in persisted + def test_none_when_no_configured_model(self): assert codex.default_model({}) is None diff --git a/tests/test_cli.py b/tests/test_cli.py index baee1974f..c809b753b 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1484,6 +1484,47 @@ def test_model_is_launch_scoped_for_claude(self, monkeypatch): mock_launch.call_args.kwargs["options"].user_pinned_model == "cat.schema.claude-opus-5" ) + def test_managed_family_alias_resolves_at_launch_boundary(self): + managed = { + "enabled_agents": { + "claude": { + "model_config": { + "default_models_by_model_family": { + "default_opus_model": "system.ai.claude-opus-4-8" + } + } + } + } + } + + with _launch_policy_patches(managed) as calls: + result = runner.invoke(app, ["claude", "--model", "opus"]) + + assert result.exit_code == 0, result.output + options = calls["launch"].call_args.kwargs["options"] + assert options.user_pinned_model == "opus" + assert options.managed_claude_model == "system.ai.claude-opus-4-8" + assert calls["launch"].call_args.args[2] == [] + + def test_managed_parent_marks_active_source_for_forwarded_alias(self): + managed = { + "enabled_agents": { + "claude": {"model_config": {"unity_catalog_location": "main.models"}} + } + } + + with _launch_policy_patches(managed, persisted_provider="main.default.stale") as calls: + calls["state"]["provider_services"] = {"claude": "main.default.stale"} + result = runner.invoke(app, ["claude", "--", "--model", "opus"]) + + assert result.exit_code == 0, result.output + launch_state = calls["launch"].call_args.args[1] + assert launch_state["_claude_launch_parent_schema"] == "main.models" + options = calls["launch"].call_args.kwargs["options"] + assert options.user_pinned_model == "opus" + assert options.managed_claude_model is None + assert calls["launch"].call_args.args[2] == ["--model", "opus"] + def test_v2_model_sets_transient_launch_override(self, monkeypatch): monkeypatch.setenv("ENABLE_SMART_ROUTING_V2", "1") state = {**MINIMAL_STATE, "claude_models": {"opus": "system.ai.claude-opus-4-8"}}