diff --git a/src/ucode/state.py b/src/ucode/state.py index 6ea8290..7e53d75 100644 --- a/src/ucode/state.py +++ b/src/ucode/state.py @@ -125,6 +125,8 @@ def build_agent_state(state: dict) -> dict[str, dict]: if not isinstance(workspace, str) or not workspace: return {} + from ucode.agents import default_model_for_tool + profile = state.get("profile") if isinstance(state.get("profile"), str) else None base_urls_value = state.get("base_urls") base_urls = base_urls_value if isinstance(base_urls_value, dict) else {} @@ -138,11 +140,18 @@ def build_agent_state(state: dict) -> dict[str, dict]: gemini_models_value = state.get("gemini_models") gemini_models = gemini_models_value if isinstance(gemini_models_value, list) else [] + selection_state = { + **state, + "claude_models": claude_models, + "codex_models": codex_models, + "gemini_models": gemini_models, + } + claude_model = ( claude_models.get("opus") or claude_models.get("sonnet") or claude_models.get("haiku") ) - codex_model = codex_models[0] if codex_models else None - pi_model = claude_model or codex_model or (gemini_models[0] if gemini_models else None) + codex_model = default_model_for_tool("codex", selection_state) + pi_model = default_model_for_tool("pi", selection_state) agents: dict[str, dict] = { "claude": { diff --git a/tests/test_state.py b/tests/test_state.py index 0b1a67f..4555af6 100644 --- a/tests/test_state.py +++ b/tests/test_state.py @@ -109,6 +109,23 @@ def test_round_trip(self): assert loaded["workspace"] == FAKE_WS assert loaded["claude_models"]["sonnet"] == "databricks-claude-sonnet-4" + def test_persists_codex_launcher_default_in_agent_state(self): + save_state( + { + "workspace": FAKE_WS, + "codex_models": [ + "system.ai.gpt-5", + "system.ai.gpt-5-1", + "system.ai.gpt-5-6-luna", + ], + } + ) + + persisted = load_full_state()["workspaces"][FAKE_WS] + assert persisted["codex_models"][0] == "system.ai.gpt-5" + assert persisted["agents"]["codex"]["model"] == "system.ai.gpt-5-6-luna" + assert persisted["agents"]["pi"]["model"] == "system.ai.gpt-5" + def test_save_respects_dry_run(self): import ucode.config_io as config_io_mod