diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 1f25252f..02f44b5b 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -41,10 +41,8 @@ from ucode.smart_routing import v2 as smart_routing_v2 from ucode.smart_routing.codex_hooks import ( remove_smart_routing_hooks, - routing_models, sync_smart_routing_hooks, ) -from ucode.smart_routing.codex_routing import codex_model_id from ucode.state import mark_tool_managed, save_state from ucode.telemetry import agent_version, ucode_version from ucode.ui import print_warning_err @@ -528,17 +526,16 @@ def _launch_smart_routing(state: dict, tool_args: list[str]) -> None: ) managed_model = default_model(state) - models = routing_models(state) - start_model = ( - managed_model - or (codex_model_id(models[0]) if models else None) - or APP_SERVER_SMART_ROUTING_STARTING_MODEL - ) + models, catalog_path = smart_routing_v2.configured_codex_models(state) + first_model = models[0] if models else None + start_model = managed_model or first_model or APP_SERVER_SMART_ROUTING_STARTING_MODEL smart_routing_v2.launch_codex( state, tool_args, binary=binary, start_model=start_model, + available_models=models, + catalog_path=catalog_path, render_overlay=render_overlay, ) diff --git a/src/ucode/smart_routing/codex_interposer.py b/src/ucode/smart_routing/codex_interposer.py index 5d39b5d1..d463a824 100644 --- a/src/ucode/smart_routing/codex_interposer.py +++ b/src/ucode/smart_routing/codex_interposer.py @@ -7,7 +7,7 @@ import time import uuid from collections.abc import Callable -from dataclasses import dataclass, replace +from dataclasses import dataclass from pathlib import Path from websockets.asyncio.client import connect @@ -99,7 +99,6 @@ def on_tui_frame(self, raw: str) -> TuiFrameResult: if decision is None: self.log(f"[ROUTE] selection failed; keeping current model: {reason}") return TuiFrameResult(raw, needs_settings_update=False) - decision = replace(decision, model=codex_routing.codex_model_id(decision.model)) self.target = decision.model if self.switch_message_fn is not None: self.switch_message = self.switch_message_fn(decision.model, decision.rationale) diff --git a/src/ucode/smart_routing/codex_routing.py b/src/ucode/smart_routing/codex_routing.py index ea048ef9..3c654e8c 100644 --- a/src/ucode/smart_routing/codex_routing.py +++ b/src/ucode/smart_routing/codex_routing.py @@ -99,7 +99,7 @@ def record(payload, task, decision, requested): workspace, token, task, available_models, timeout=timeout ), default_task_label="Codex subagent task", - model_id_mapper=codex_model_id, + model_id_mapper=lambda model: model, record_decision=record, ) diff --git a/src/ucode/smart_routing/v2.py b/src/ucode/smart_routing/v2.py index 9b8a29eb..b404cb43 100644 --- a/src/ucode/smart_routing/v2.py +++ b/src/ucode/smart_routing/v2.py @@ -24,7 +24,7 @@ list_anthropic_model_catalog, list_anthropic_models, ) -from ucode.smart_routing import claude_routing, codex_interposer, routing +from ucode.smart_routing import claude_routing, codex_interposer, codex_routing, routing from ucode.smart_routing.claude_hooks import ( FIRST_PROMPT_SOCKET_ENV, sync_first_prompt_hook, @@ -414,10 +414,47 @@ def route_prompt(prompt: str) -> claude_pty.FirstPromptRoute: sys.exit(returncode) -# TODO: Replace with /codex/v1/models once /codex/v1/models can send GPT models as well. -def _cached_routing_models(state: dict) -> list[str]: - """Return the persisted UC model-service ids usable by Codex routing.""" - return routing_models(state) +def _model_catalog() -> tuple[Path, list[str]] | None: + """Read the first configured Codex model catalog using Codex config precedence.""" + from ucode.agents import codex + + config_paths = ( + codex._managed_config_path(), + codex.CODEX_CONFIG_PATH, + _codex_home_config_path(), + ) + for config_path in dict.fromkeys(config_paths): + if config_path is None: + continue + configured_path = read_toml_safe(config_path).get("model_catalog_json") + if not isinstance(configured_path, str) or not configured_path.strip(): + continue + catalog_path = Path(configured_path).expanduser() + rows = read_json_safe(catalog_path).get("models") + if not isinstance(rows, list): + return catalog_path, [] + models: list[str] = [] + seen_models: set[str] = set() + for row in rows: + slug = row.get("slug") if isinstance(row, dict) else None + if not isinstance(slug, str): + continue + model = slug.strip() + if not model or model in seen_models: + continue + seen_models.add(model) + models.append(model) + return catalog_path, models + return None + + +def configured_codex_models(state: dict) -> tuple[list[str], Path | None]: + """Return catalog models when configured, otherwise the existing routing models.""" + catalog = _model_catalog() + if catalog is not None: + catalog_path, models = catalog + return models, catalog_path + return [codex_routing.codex_model_id(model) for model in routing_models(state)], None def _codex_home_config_path() -> Path: @@ -444,6 +481,8 @@ def launch_codex( *, binary: str, start_model: str | None, + available_models: list[str], + catalog_path: Path | None, render_overlay: Callable[..., dict], ) -> NoReturn: workspace = state.get("workspace") @@ -458,7 +497,6 @@ def launch_codex( profile = state.get("profile") os.environ[OAUTH_TOKEN_ENV_VAR] = get_databricks_token(workspace, profile) - available_models = _cached_routing_models(state) if not available_models: print_note( "Smart routing model metadata is unavailable; starting Codex on gpt-5.6-luna " @@ -470,6 +508,8 @@ def launch_codex( state.get("profile"), use_pat=bool(state.get("use_pat")), ) + if catalog_path is not None: + overlay["model_catalog_json"] = str(catalog_path) overlay["hooks"] = { "PreToolUse": _v2_pre_tool_use_hooks(state, available_models), } diff --git a/tests/test_codex_routing.py b/tests/test_codex_routing.py index 097d8969..55ee7283 100644 --- a/tests/test_codex_routing.py +++ b/tests/test_codex_routing.py @@ -167,7 +167,7 @@ def test_spawn_rewrite_preserves_original_input(monkeypatch): "request_routing_decision", lambda *args, **kwargs: ( codex_routing.RoutingDecision( - model="databricks-gpt-5-5", + model="gpt-5.5", raw_model="gpt-5-6-sol", rationale="Review needs deeper reasoning.", ), @@ -179,7 +179,7 @@ def test_spawn_rewrite_preserves_original_input(monkeypatch): payload, workspace=WS, token="token", - available_models=["databricks-gpt-5-5"], + available_models=["gpt-5.5"], ) hook = output["hookSpecificOutput"] @@ -201,13 +201,13 @@ def test_spawn_rewrite_preserves_original_input(monkeypatch): assert hook["permissionDecisionReason"] == expected_message -def test_spawn_rewrite_uses_codex_model_id_for_uc_endpoint(monkeypatch): +def test_spawn_rewrite_uses_selected_codex_model_id(monkeypatch): monkeypatch.setattr( codex_routing, "request_routing_decision", lambda *args, **kwargs: ( codex_routing.RoutingDecision( - model="system.ai.gpt-5-6-luna", + model="gpt-5.6-luna", raw_model="gpt-5-6-luna", ), None, @@ -221,7 +221,7 @@ def test_spawn_rewrite_uses_codex_model_id_for_uc_endpoint(monkeypatch): }, workspace=WS, token="token", - available_models=["system.ai.gpt-5-6-luna"], + available_models=["gpt-5.6-luna"], ) assert output["systemMessage"] == codex_routing.routing.format_subagent_message( @@ -230,6 +230,32 @@ def test_spawn_rewrite_uses_codex_model_id_for_uc_endpoint(monkeypatch): assert output["hookSpecificOutput"]["updatedInput"]["model"] == "gpt-5.6-luna" +def test_spawn_rewrite_preserves_custom_catalog_model_id(monkeypatch): + monkeypatch.setattr( + codex_routing, + "request_routing_decision", + lambda *args, **kwargs: ( + codex_routing.RoutingDecision( + model="system.ai.gpt-5-5", + raw_model="gpt-5-5", + ), + None, + ), + ) + + output = codex_routing.route_pre_tool_use( + { + "tool_name": "collaborationspawn_agent", + "tool_input": {"task_name": "routing-smoke-test", "message": "encrypted"}, + }, + workspace=WS, + token="token", + available_models=["system.ai.gpt-5-5"], + ) + + assert output["hookSpecificOutput"]["updatedInput"]["model"] == "system.ai.gpt-5-5" + + def test_codex_model_id_maps_uc_gpt_models_to_codex_slugs(): expected = { "system.ai.gpt-5-2": "gpt-5.2", @@ -373,7 +399,7 @@ def test_decision_is_reconciled_with_actual_subagent_model(tmp_path, monkeypatch "request_routing_decision", lambda *args, **kwargs: ( codex_routing.RoutingDecision( - model="system.ai.gpt-5-6-luna", + model="gpt-5.6-luna", raw_model="gpt-5-6-luna", ), None, @@ -388,7 +414,7 @@ def test_decision_is_reconciled_with_actual_subagent_model(tmp_path, monkeypatch }, workspace=WS, token="token", - available_models=["system.ai.gpt-5-6-luna", "system.ai.gpt-5-6-sol"], + available_models=["gpt-5.6-luna", "gpt-5.6-sol"], audit_decision=True, ) record = codex_routing.record_subagent_start( diff --git a/tests/test_codex_smart_routing_v2.py b/tests/test_codex_smart_routing_v2.py index fc1dc582..9180814e 100644 --- a/tests/test_codex_smart_routing_v2.py +++ b/tests/test_codex_smart_routing_v2.py @@ -1,6 +1,7 @@ from __future__ import annotations import json +from pathlib import Path import pytest @@ -23,6 +24,107 @@ def test_smart_routing_switch_message_is_boxed(): ) +class TestCodexModelCatalog: + @staticmethod + def _configure_paths(tmp_path, monkeypatch): + managed = tmp_path / "managed_config.toml" + profile = tmp_path / "ucode.config.toml" + codex_home = tmp_path / "codex-home" + codex_home.mkdir() + monkeypatch.setattr(codex, "_managed_config_path", lambda: managed) + monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", profile) + monkeypatch.setenv("CODEX_HOME", str(codex_home)) + return managed, profile, codex_home / "config.toml" + + @staticmethod + def _write_catalog(path, models): + path.write_text(json.dumps({"models": models}), encoding="utf-8") + + def test_managed_catalog_models_take_priority_and_are_deduplicated(self, tmp_path, monkeypatch): + managed, profile, _user = self._configure_paths(tmp_path, monkeypatch) + managed_catalog = tmp_path / "managed-models.json" + profile_catalog = tmp_path / "profile-models.json" + self._write_catalog( + managed_catalog, + [ + {"slug": " system.ai.gpt-5-5 "}, + {"slug": "system.ai.glm-5-3"}, + {"slug": "system.ai.gpt-5-5"}, + {"display_name": "missing slug"}, + ], + ) + self._write_catalog(profile_catalog, [{"slug": "profile-model"}]) + managed.write_text(f'model_catalog_json = "{managed_catalog}"\n', encoding="utf-8") + profile.write_text(f'model_catalog_json = "{profile_catalog}"\n', encoding="utf-8") + + models, catalog_path = v2.configured_codex_models({"codex_models": ["cached-model"]}) + + assert models == ["system.ai.gpt-5-5", "system.ai.glm-5-3"] + assert catalog_path == managed_catalog + + def test_falls_back_to_profile_catalog(self, tmp_path, monkeypatch): + managed, profile, _user = self._configure_paths(tmp_path, monkeypatch) + catalog = tmp_path / "profile-models.json" + self._write_catalog(catalog, [{"slug": "system.ai.gpt-5-6-sol"}]) + managed.write_text('model_provider = "managed"\n', encoding="utf-8") + profile.write_text(f'model_catalog_json = "{catalog}"\n', encoding="utf-8") + + models, catalog_path = v2.configured_codex_models({}) + + assert models == ["system.ai.gpt-5-6-sol"] + assert catalog_path == catalog + + def test_falls_back_to_user_catalog(self, tmp_path, monkeypatch): + managed, profile, user = self._configure_paths(tmp_path, monkeypatch) + catalog = tmp_path / "user-models.json" + self._write_catalog(catalog, [{"slug": "system.ai.glm-5-3"}]) + managed.write_text('model_provider = "managed"\n', encoding="utf-8") + profile.write_text('model_provider = "profile"\n', encoding="utf-8") + user.write_text(f'model_catalog_json = "{catalog}"\n', encoding="utf-8") + + models, catalog_path = v2.configured_codex_models({}) + + assert models == ["system.ai.glm-5-3"] + assert catalog_path == catalog + + def test_falls_back_to_cached_workspace_models(self, tmp_path, monkeypatch): + self._configure_paths(tmp_path, monkeypatch) + state = { + "codex_models": ["system.ai.gpt-5-6-sol"], + "oss_models": ["system.ai.glm-5-2"], + } + + assert v2.configured_codex_models(state) == ( + ["gpt-5.6-sol", "system.ai.glm-5-2"], + None, + ) + + def test_configured_empty_catalog_does_not_fall_back_to_cached_models( + self, tmp_path, monkeypatch + ): + managed, _profile, _user = self._configure_paths(tmp_path, monkeypatch) + catalog = tmp_path / "empty-models.json" + self._write_catalog(catalog, []) + managed.write_text(f'model_catalog_json = "{catalog}"\n', encoding="utf-8") + + models, catalog_path = v2.configured_codex_models({"codex_models": ["cached-model"]}) + + assert models == [] + assert catalog_path == catalog + + def test_configured_unreadable_catalog_does_not_fall_back_to_cached_models( + self, tmp_path, monkeypatch + ): + managed, _profile, _user = self._configure_paths(tmp_path, monkeypatch) + catalog = tmp_path / "missing-models.json" + managed.write_text(f'model_catalog_json = "{catalog}"\n', encoding="utf-8") + + models, catalog_path = v2.configured_codex_models({"codex_models": ["cached-model"]}) + + assert models == [] + assert catalog_path == catalog + + class TestLaunchCodex: def test_rejects_unsupported_codex_version(self, monkeypatch): monkeypatch.setenv(v2.ENV_VAR, "1") @@ -42,9 +144,12 @@ def test_rejects_unsupported_codex_version(self, monkeypatch): ) def test_codex_smart_routing_launch_dispatches_to_v2(self, monkeypatch, tool_args, options): calls = [] + models = ["system.ai.gpt-5-5"] + catalog_path = Path("/catalog.json") monkeypatch.setenv(v2.ENV_VAR, "1") - monkeypatch.setattr(codex, "default_model", lambda state: "gpt-start") + monkeypatch.setattr(codex, "default_model", lambda state: None) monkeypatch.setattr(codex, "clear_model_preferences", lambda state: False) + monkeypatch.setattr(v2, "configured_codex_models", lambda state: (models, catalog_path)) def launch_v2(state, tool_args, **kwargs): calls.append((state, tool_args, kwargs)) @@ -63,7 +168,9 @@ def launch_v2(state, tool_args, **kwargs): tool_args, { "binary": "codex", - "start_model": "gpt-start", + "start_model": "system.ai.gpt-5-5", + "available_models": models, + "catalog_path": catalog_path, "render_overlay": codex.render_overlay, }, ) @@ -106,6 +213,11 @@ def test_codex_launch_normalizes_cached_bootstrap_model(self, monkeypatch): monkeypatch.setenv(v2.ENV_VAR, "1") monkeypatch.setattr(codex, "clear_model_preferences", lambda state: False) monkeypatch.setattr(codex, "default_model", lambda state: None) + monkeypatch.setattr( + v2, + "configured_codex_models", + lambda state: (["gpt-5.6-luna"], None), + ) def launch_v2(state, tool_args, **kwargs): calls.append(kwargs) @@ -178,6 +290,8 @@ def start_interposer(*args, **kwargs): ["--search"], binary="codex", start_model="gpt-start", + available_models=["system.ai.gpt-5-6-sol", "system.ai.glm-5-2"], + catalog_path=Path("/catalog.json"), render_overlay=codex.render_overlay, ) @@ -192,8 +306,12 @@ def start_interposer(*args, **kwargs): "--config", ] assert processes[0].argv[7].startswith("model_providers.ucode-databricks={") - assert processes[0].argv[8] == "--config" - hook_override = processes[0].argv[9] + assert processes[0].argv[8:10] == [ + "--config", + 'model_catalog_json="/catalog.json"', + ] + assert processes[0].argv[10] == "--config" + hook_override = processes[0].argv[11] assert hook_override.startswith("hooks.PreToolUse=[{") assert 'matcher = "Agent|.*spawn_agent$"' in hook_override assert "codex-router-hook route-subagent" in hook_override @@ -201,7 +319,7 @@ def start_interposer(*args, **kwargs): assert "--profile myprof" in hook_override assert "--model system.ai.gpt-5-6-sol" in hook_override assert "--model system.ai.glm-5-2" in hook_override - assert processes[0].argv[10:] == [ + assert processes[0].argv[12:] == [ "--listen", "ws://127.0.0.1:41001", ] @@ -308,6 +426,8 @@ def test_missing_cached_models_starts_with_bootstrap_model(self, monkeypatch): [], binary="codex", start_model="gpt-5.6-luna", + available_models=[], + catalog_path=None, render_overlay=codex.render_overlay, ) @@ -446,11 +566,11 @@ def select(prompt): assert result.needs_settings_update assert "Task classified as bugfix." in sess.switch_message - def test_maps_selected_uc_gpt_model_and_shows_routing_notice(self): + def test_uses_selected_codex_model_and_shows_routing_notice(self): def select(_prompt): return ( codex_interposer.routing.RoutingDecision( - model="system.ai.gpt-5-6-luna", + model="gpt-5.6-luna", raw_model="gpt-5-6-luna", rationale="Trivial task.", ), @@ -476,6 +596,23 @@ def select(_prompt): ] assert "Selected Model : gpt-5.6-luna" in (injected[1]["params"]["item"]["text"]) + def test_preserves_selected_custom_catalog_slug(self): + sess = codex_interposer._Session( + None, + log=lambda _m: None, + route_decision=lambda _prompt: ( + codex_interposer.routing.RoutingDecision( + model="system.ai.gpt-5-5", + raw_model="gpt-5-5", + ), + None, + ), + ) + + result = sess.on_tui_frame(self._turn_start("system.ai.glm-5-3")) + + assert json.loads(result.frame)["params"]["model"] == "system.ai.gpt-5-5" + def test_routes_first_prompt_to_oss_model(self): def select(_prompt): return (