From a0b2f16343f13e1b01ee94507be78256df5346ea Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Tue, 8 Sep 2026 17:33:43 +0000 Subject: [PATCH 1/5] send models in model_catalog.json to smart router if they exist --- src/ucode/agents/codex.py | 14 +-- src/ucode/cli.py | 2 + src/ucode/smart_routing/codex_hooks.py | 32 ++++- src/ucode/smart_routing/codex_interposer.py | 11 +- src/ucode/smart_routing/codex_routing.py | 3 +- src/ucode/smart_routing/v2.py | 73 +++++++++-- tests/test_cli.py | 2 + tests/test_codex_routing.py | 27 ++++ tests/test_codex_smart_routing_v2.py | 133 +++++++++++++++++++- 9 files changed, 272 insertions(+), 25 deletions(-) diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 1f25252f..46c1ea76 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -41,7 +41,6 @@ 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 @@ -528,17 +527,18 @@ 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 - ) + routing_options = smart_routing_v2.codex_routing_options(state) + models = routing_options.models + first_model = None + if models: + first_model = models[0] if routing_options.preserve_model_ids else codex_model_id(models[0]) + 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, + routing_options=routing_options, render_overlay=render_overlay, ) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 02d0be74..d0705136 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -1436,6 +1436,7 @@ def codex_router_hook_cmd( profile: Annotated[str | None, typer.Option("--profile")] = None, use_pat: Annotated[bool, typer.Option("--use-pat")] = False, model: Annotated[list[str] | None, typer.Option("--model")] = None, + preserve_model_ids: Annotated[bool, typer.Option("--preserve-model-ids")] = False, ) -> None: """Run a Codex smart-routing lifecycle hook.""" import json @@ -1502,6 +1503,7 @@ def codex_router_hook_cmd( token=token, available_models=model or [], audit_decision=True, + preserve_model_ids=preserve_model_ids, ) if output is not None: sys.stdout.write(json.dumps(output)) diff --git a/src/ucode/smart_routing/codex_hooks.py b/src/ucode/smart_routing/codex_hooks.py index 421cabda..e79235c1 100644 --- a/src/ucode/smart_routing/codex_hooks.py +++ b/src/ucode/smart_routing/codex_hooks.py @@ -53,23 +53,41 @@ def _routing_hook_groups(state: dict) -> dict[str, list[dict]]: def merge_pre_tool_use_hooks( - existing: list[dict], state: dict, *, available_models: list[str] + existing: list[dict], + state: dict, + *, + available_models: list[str], + preserve_model_ids: bool = False, ) -> list[dict]: """Add the ucode spawn hook to an existing Codex PreToolUse hook list.""" doc = {"hooks": {"PreToolUse": copy.deepcopy(existing)}} hooks.sync_managed_hooks( doc, ROUTING_HOOK_COMMAND_MARKER, - {"PreToolUse": [_pre_tool_use_hook_group(state, available_models=available_models)]}, + { + "PreToolUse": [ + _pre_tool_use_hook_group( + state, + available_models=available_models, + preserve_model_ids=preserve_model_ids, + ) + ] + }, ) return doc["hooks"]["PreToolUse"] -def _pre_tool_use_hook_group(state: dict, *, available_models: list[str] | None = None) -> dict: +def _pre_tool_use_hook_group( + state: dict, + *, + available_models: list[str] | None = None, + preserve_model_ids: bool = False, +) -> dict: route_argv = _routing_hook_argv( state, "route-subagent", available_models=available_models, + preserve_model_ids=preserve_model_ids, ) return { "matcher": "Agent|.*spawn_agent$", @@ -78,7 +96,11 @@ def _pre_tool_use_hook_group(state: dict, *, available_models: list[str] | None def _routing_hook_argv( - state: dict, event: str, *, available_models: list[str] | None = None + state: dict, + event: str, + *, + available_models: list[str] | None = None, + preserve_model_ids: bool = False, ) -> list[str]: workspace = str(state.get("workspace") or "") argv = [ @@ -96,6 +118,8 @@ def _routing_hook_argv( argv += ["--profile", profile] if state.get("use_pat"): argv.append("--use-pat") + if preserve_model_ids: + argv.append("--preserve-model-ids") models = available_models if available_models is not None else routing_models(state) for model in models: if isinstance(model, str) and model: diff --git a/src/ucode/smart_routing/codex_interposer.py b/src/ucode/smart_routing/codex_interposer.py index 5d39b5d1..b8b1dce1 100644 --- a/src/ucode/smart_routing/codex_interposer.py +++ b/src/ucode/smart_routing/codex_interposer.py @@ -60,6 +60,7 @@ def __init__( available_models: list[str] | None = None, route_decision: RouteDecisionFn | None = None, switch_message_fn: SwitchMessageFn | None = None, + preserve_model_ids: bool = False, ) -> None: self.target = target_model self.available_models = list(available_models or []) @@ -67,6 +68,7 @@ def __init__( self.switch_message = switch_message self.route_decision = route_decision self.switch_message_fn = switch_message_fn + self.preserve_model_ids = preserve_model_ids self.thread_id: str | None = None self.settings: dict | None = None self.first_turn_seen = False @@ -99,7 +101,8 @@ 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)) + if not self.preserve_model_ids: + 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) @@ -211,6 +214,7 @@ async def _handle_tui( workspace: str | None = None, token_provider: TokenProvider | None = None, switch_message_fn: SwitchMessageFn | None = None, + preserve_model_ids: bool = False, ) -> None: path = getattr(getattr(tui, "request", None), "path", "/") or "/" uri = upstream_uri.rstrip("/") + path @@ -238,6 +242,7 @@ def route_decision(prompt: str): available_models, route_decision, switch_message_fn, + preserve_model_ids, ) async with connect(uri, max_size=None) as upstream: @@ -300,6 +305,7 @@ async def _serve( workspace: str | None = None, token_provider: TokenProvider | None = None, switch_message_fn: SwitchMessageFn | None = None, + preserve_model_ids: bool = False, ): async def handler(tui): try: @@ -313,6 +319,7 @@ async def handler(tui): workspace, token_provider, switch_message_fn, + preserve_model_ids, ) except Exception as exc: # noqa: BLE001 log(f"[ERR] session: {exc!r}") @@ -333,6 +340,7 @@ def start_interposer_thread( token_provider: TokenProvider | None = None, switch_message_fn: SwitchMessageFn | None = None, switch_message: str | None = None, + preserve_model_ids: bool = False, log_path: Path | None = None, ready_timeout: float = 10.0, ) -> tuple[int, Callable[[], None]]: @@ -365,6 +373,7 @@ def run() -> None: workspace, token_provider, switch_message_fn, + preserve_model_ids, ) ) holder["port"] = holder["server"].sockets[0].getsockname()[1] diff --git a/src/ucode/smart_routing/codex_routing.py b/src/ucode/smart_routing/codex_routing.py index ea048ef9..1213ad91 100644 --- a/src/ucode/smart_routing/codex_routing.py +++ b/src/ucode/smart_routing/codex_routing.py @@ -84,6 +84,7 @@ def route_pre_tool_use( available_models: list[str], timeout: float = REQUEST_TIMEOUT_S, audit_decision: bool = False, + preserve_model_ids: bool = False, ) -> dict[str, Any] | None: """Route one Codex ``spawn_agent`` call and rewrite its model.""" record = None @@ -99,7 +100,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) if preserve_model_ids else codex_model_id, record_decision=record, ) diff --git a/src/ucode/smart_routing/v2.py b/src/ucode/smart_routing/v2.py index 9b8a29eb..7821e598 100644 --- a/src/ucode/smart_routing/v2.py +++ b/src/ucode/smart_routing/v2.py @@ -12,6 +12,7 @@ import urllib.request import uuid from collections.abc import Callable +from dataclasses import dataclass from pathlib import Path from typing import NoReturn, TextIO @@ -58,6 +59,13 @@ _ANTHROPIC_AIGW_MODEL_RE = re.compile(r"^anthropic-aigw-[0-9a-fA-F]{8}-(.+)$") +@dataclass(frozen=True) +class CodexRoutingOptions: + models: list[str] + preserve_model_ids: bool = False + catalog_path: Path | None = None + + def enabled() -> bool: return os.environ.get(ENV_VAR) == "1" @@ -414,10 +422,50 @@ 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.""" + try: + from ucode.agents import codex + + paths = [codex._managed_config_path(), codex.CODEX_CONFIG_PATH, _codex_home_config_path()] + except (ImportError, OSError): + return None + + seen_paths: set[Path] = set() + for config_path in paths: + if config_path is None or config_path in seen_paths: + continue + seen_paths.add(config_path) + 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): + continue + 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) + if models: + return catalog_path, models + return None + + +def codex_routing_options(state: dict) -> CodexRoutingOptions: + """Prefer custom-catalog models, then fall back to persisted workspace models.""" + catalog = _model_catalog() + if catalog is not None: + catalog_path, models = catalog + return CodexRoutingOptions(models, preserve_model_ids=True, catalog_path=catalog_path) + return CodexRoutingOptions(routing_models(state)) def _codex_home_config_path() -> Path: @@ -427,7 +475,9 @@ def _codex_home_config_path() -> Path: return Path.home() / ".codex" / "config.toml" -def _v2_pre_tool_use_hooks(state: dict, available_models: list[str]) -> list[dict]: +def _v2_pre_tool_use_hooks( + state: dict, available_models: list[str], *, preserve_model_ids: bool = False +) -> list[dict]: doc = read_toml_safe(_codex_home_config_path()) configured_hooks = doc.get("hooks") existing = configured_hooks.get("PreToolUse") if isinstance(configured_hooks, dict) else None @@ -435,6 +485,7 @@ def _v2_pre_tool_use_hooks(state: dict, available_models: list[str]) -> list[dic existing if isinstance(existing, list) else [], state, available_models=available_models, + preserve_model_ids=preserve_model_ids, ) @@ -444,6 +495,7 @@ def launch_codex( *, binary: str, start_model: str | None, + routing_options: CodexRoutingOptions, render_overlay: Callable[..., dict], ) -> NoReturn: workspace = state.get("workspace") @@ -458,7 +510,7 @@ def launch_codex( profile = state.get("profile") os.environ[OAUTH_TOKEN_ENV_VAR] = get_databricks_token(workspace, profile) - available_models = _cached_routing_models(state) + available_models = routing_options.models if not available_models: print_note( "Smart routing model metadata is unavailable; starting Codex on gpt-5.6-luna " @@ -470,8 +522,14 @@ def launch_codex( state.get("profile"), use_pat=bool(state.get("use_pat")), ) + if routing_options.catalog_path is not None: + overlay["model_catalog_json"] = str(routing_options.catalog_path) overlay["hooks"] = { - "PreToolUse": _v2_pre_tool_use_hooks(state, available_models), + "PreToolUse": _v2_pre_tool_use_hooks( + state, + available_models, + preserve_model_ids=routing_options.preserve_model_ids, + ), } config_args = codex_config_args(overlay) app_port = _free_port() @@ -499,6 +557,7 @@ def launch_codex( workspace=workspace, token_provider=lambda: get_databricks_token(workspace, profile), switch_message_fn=format_routing_notice, + preserve_model_ids=routing_options.preserve_model_ids, log_path=CODEX_INTERPOSER_LOG, ) tui_url = _loopback_websocket_url(tui_port) diff --git a/tests/test_cli.py b/tests/test_cli.py index 8ef41712..e577fc5f 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -674,6 +674,7 @@ def _invoke_codex_subagent_hook(token_env): "my-profile", "--model", "system.ai.gpt-5-6-sol", + "--preserve-model-ids", ], input='{"tool_name":"collaboration.spawn_agent","tool_input":{"message":"fix it"}}', env={"ENABLE_SMART_ROUTING_V2": "1", **token_env}, @@ -681,6 +682,7 @@ def _invoke_codex_subagent_hook(token_env): assert result.exit_code == 0, result.output assert json.loads(result.output) == routed + assert mock_route.call_args.kwargs["preserve_model_ids"] is True return mock_token, mock_route def test_codex_subagent_hook_reuses_fresh_oauth_token(self, monkeypatch): diff --git a/tests/test_codex_routing.py b/tests/test_codex_routing.py index 097d8969..8c6595eb 100644 --- a/tests/test_codex_routing.py +++ b/tests/test_codex_routing.py @@ -230,6 +230,33 @@ 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"], + preserve_model_ids=True, + ) + + 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", diff --git a/tests/test_codex_smart_routing_v2.py b/tests/test_codex_smart_routing_v2.py index fc1dc582..3e421c40 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,86 @@ 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") + + options = v2.codex_routing_options({"codex_models": ["cached-model"]}) + + assert options == v2.CodexRoutingOptions( + ["system.ai.gpt-5-5", "system.ai.glm-5-3"], + preserve_model_ids=True, + 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") + + options = v2.codex_routing_options({}) + + assert options.models == ["system.ai.gpt-5-6-sol"] + assert options.catalog_path == catalog + assert options.preserve_model_ids is True + + 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") + + options = v2.codex_routing_options({}) + + assert options.models == ["system.ai.glm-5-3"] + assert options.catalog_path == catalog + assert options.preserve_model_ids is True + + 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.codex_routing_options(state) == v2.CodexRoutingOptions( + ["system.ai.gpt-5-6-sol", "system.ai.glm-5-2"] + ) + + class TestLaunchCodex: def test_rejects_unsupported_codex_version(self, monkeypatch): monkeypatch.setenv(v2.ENV_VAR, "1") @@ -42,9 +123,15 @@ def test_rejects_unsupported_codex_version(self, monkeypatch): ) def test_codex_smart_routing_launch_dispatches_to_v2(self, monkeypatch, tool_args, options): calls = [] + routing_options = v2.CodexRoutingOptions( + ["system.ai.gpt-5-5"], + preserve_model_ids=True, + 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, "codex_routing_options", lambda state: routing_options) def launch_v2(state, tool_args, **kwargs): calls.append((state, tool_args, kwargs)) @@ -63,7 +150,8 @@ def launch_v2(state, tool_args, **kwargs): tool_args, { "binary": "codex", - "start_model": "gpt-start", + "start_model": "system.ai.gpt-5-5", + "routing_options": routing_options, "render_overlay": codex.render_overlay, }, ) @@ -106,6 +194,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, + "codex_routing_options", + lambda state: v2.CodexRoutingOptions(["system.ai.gpt-5-6-luna"]), + ) def launch_v2(state, tool_args, **kwargs): calls.append(kwargs) @@ -178,6 +271,11 @@ def start_interposer(*args, **kwargs): ["--search"], binary="codex", start_model="gpt-start", + routing_options=v2.CodexRoutingOptions( + ["system.ai.gpt-5-6-sol", "system.ai.glm-5-2"], + preserve_model_ids=True, + catalog_path=Path("/catalog.json"), + ), render_overlay=codex.render_overlay, ) @@ -192,16 +290,21 @@ 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 assert f"--host {WS}" in hook_override assert "--profile myprof" in hook_override + assert "--preserve-model-ids" 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", ] @@ -225,6 +328,7 @@ def start_interposer(*args, **kwargs): assert interposer_args["kwargs"]["token_provider"]() == "token-2" assert token_calls == [(WS, "myprof"), (WS, "myprof")] assert interposer_args["kwargs"]["switch_message_fn"] is v2.format_routing_notice + assert interposer_args["kwargs"]["preserve_model_ids"] is True assert stopped == [True] assert processes[0].terminated is True @@ -308,6 +412,7 @@ def test_missing_cached_models_starts_with_bootstrap_model(self, monkeypatch): [], binary="codex", start_model="gpt-5.6-luna", + routing_options=v2.CodexRoutingOptions([]), render_overlay=codex.render_overlay, ) @@ -476,6 +581,24 @@ 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, + ), + preserve_model_ids=True, + ) + + 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 ( From 905583080c8099a22fcb0c3628a17eb83bd5a394 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Tue, 8 Sep 2026 17:40:25 +0000 Subject: [PATCH 2/5] update --- src/ucode/agents/codex.py | 5 +---- src/ucode/cli.py | 2 -- src/ucode/smart_routing/codex_hooks.py | 7 ------- src/ucode/smart_routing/codex_interposer.py | 12 +----------- src/ucode/smart_routing/codex_routing.py | 3 +-- src/ucode/smart_routing/v2.py | 21 +++++++-------------- tests/test_cli.py | 2 -- tests/test_codex_routing.py | 15 +++++++-------- tests/test_codex_smart_routing_v2.py | 16 ++++------------ 9 files changed, 21 insertions(+), 62 deletions(-) diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 46c1ea76..395159d0 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -43,7 +43,6 @@ remove_smart_routing_hooks, 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 @@ -529,9 +528,7 @@ def _launch_smart_routing(state: dict, tool_args: list[str]) -> None: managed_model = default_model(state) routing_options = smart_routing_v2.codex_routing_options(state) models = routing_options.models - first_model = None - if models: - first_model = models[0] if routing_options.preserve_model_ids else codex_model_id(models[0]) + 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, diff --git a/src/ucode/cli.py b/src/ucode/cli.py index d0705136..02d0be74 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -1436,7 +1436,6 @@ def codex_router_hook_cmd( profile: Annotated[str | None, typer.Option("--profile")] = None, use_pat: Annotated[bool, typer.Option("--use-pat")] = False, model: Annotated[list[str] | None, typer.Option("--model")] = None, - preserve_model_ids: Annotated[bool, typer.Option("--preserve-model-ids")] = False, ) -> None: """Run a Codex smart-routing lifecycle hook.""" import json @@ -1503,7 +1502,6 @@ def codex_router_hook_cmd( token=token, available_models=model or [], audit_decision=True, - preserve_model_ids=preserve_model_ids, ) if output is not None: sys.stdout.write(json.dumps(output)) diff --git a/src/ucode/smart_routing/codex_hooks.py b/src/ucode/smart_routing/codex_hooks.py index e79235c1..98768324 100644 --- a/src/ucode/smart_routing/codex_hooks.py +++ b/src/ucode/smart_routing/codex_hooks.py @@ -57,7 +57,6 @@ def merge_pre_tool_use_hooks( state: dict, *, available_models: list[str], - preserve_model_ids: bool = False, ) -> list[dict]: """Add the ucode spawn hook to an existing Codex PreToolUse hook list.""" doc = {"hooks": {"PreToolUse": copy.deepcopy(existing)}} @@ -69,7 +68,6 @@ def merge_pre_tool_use_hooks( _pre_tool_use_hook_group( state, available_models=available_models, - preserve_model_ids=preserve_model_ids, ) ] }, @@ -81,13 +79,11 @@ def _pre_tool_use_hook_group( state: dict, *, available_models: list[str] | None = None, - preserve_model_ids: bool = False, ) -> dict: route_argv = _routing_hook_argv( state, "route-subagent", available_models=available_models, - preserve_model_ids=preserve_model_ids, ) return { "matcher": "Agent|.*spawn_agent$", @@ -100,7 +96,6 @@ def _routing_hook_argv( event: str, *, available_models: list[str] | None = None, - preserve_model_ids: bool = False, ) -> list[str]: workspace = str(state.get("workspace") or "") argv = [ @@ -118,8 +113,6 @@ def _routing_hook_argv( argv += ["--profile", profile] if state.get("use_pat"): argv.append("--use-pat") - if preserve_model_ids: - argv.append("--preserve-model-ids") models = available_models if available_models is not None else routing_models(state) for model in models: if isinstance(model, str) and model: diff --git a/src/ucode/smart_routing/codex_interposer.py b/src/ucode/smart_routing/codex_interposer.py index b8b1dce1..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 @@ -60,7 +60,6 @@ def __init__( available_models: list[str] | None = None, route_decision: RouteDecisionFn | None = None, switch_message_fn: SwitchMessageFn | None = None, - preserve_model_ids: bool = False, ) -> None: self.target = target_model self.available_models = list(available_models or []) @@ -68,7 +67,6 @@ def __init__( self.switch_message = switch_message self.route_decision = route_decision self.switch_message_fn = switch_message_fn - self.preserve_model_ids = preserve_model_ids self.thread_id: str | None = None self.settings: dict | None = None self.first_turn_seen = False @@ -101,8 +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) - if not self.preserve_model_ids: - 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) @@ -214,7 +210,6 @@ async def _handle_tui( workspace: str | None = None, token_provider: TokenProvider | None = None, switch_message_fn: SwitchMessageFn | None = None, - preserve_model_ids: bool = False, ) -> None: path = getattr(getattr(tui, "request", None), "path", "/") or "/" uri = upstream_uri.rstrip("/") + path @@ -242,7 +237,6 @@ def route_decision(prompt: str): available_models, route_decision, switch_message_fn, - preserve_model_ids, ) async with connect(uri, max_size=None) as upstream: @@ -305,7 +299,6 @@ async def _serve( workspace: str | None = None, token_provider: TokenProvider | None = None, switch_message_fn: SwitchMessageFn | None = None, - preserve_model_ids: bool = False, ): async def handler(tui): try: @@ -319,7 +312,6 @@ async def handler(tui): workspace, token_provider, switch_message_fn, - preserve_model_ids, ) except Exception as exc: # noqa: BLE001 log(f"[ERR] session: {exc!r}") @@ -340,7 +332,6 @@ def start_interposer_thread( token_provider: TokenProvider | None = None, switch_message_fn: SwitchMessageFn | None = None, switch_message: str | None = None, - preserve_model_ids: bool = False, log_path: Path | None = None, ready_timeout: float = 10.0, ) -> tuple[int, Callable[[], None]]: @@ -373,7 +364,6 @@ def run() -> None: workspace, token_provider, switch_message_fn, - preserve_model_ids, ) ) holder["port"] = holder["server"].sockets[0].getsockname()[1] diff --git a/src/ucode/smart_routing/codex_routing.py b/src/ucode/smart_routing/codex_routing.py index 1213ad91..3c654e8c 100644 --- a/src/ucode/smart_routing/codex_routing.py +++ b/src/ucode/smart_routing/codex_routing.py @@ -84,7 +84,6 @@ def route_pre_tool_use( available_models: list[str], timeout: float = REQUEST_TIMEOUT_S, audit_decision: bool = False, - preserve_model_ids: bool = False, ) -> dict[str, Any] | None: """Route one Codex ``spawn_agent`` call and rewrite its model.""" record = None @@ -100,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=(lambda model: model) if preserve_model_ids else 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 7821e598..a2db07fb 100644 --- a/src/ucode/smart_routing/v2.py +++ b/src/ucode/smart_routing/v2.py @@ -25,7 +25,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, @@ -62,7 +62,6 @@ @dataclass(frozen=True) class CodexRoutingOptions: models: list[str] - preserve_model_ids: bool = False catalog_path: Path | None = None @@ -464,8 +463,10 @@ def codex_routing_options(state: dict) -> CodexRoutingOptions: catalog = _model_catalog() if catalog is not None: catalog_path, models = catalog - return CodexRoutingOptions(models, preserve_model_ids=True, catalog_path=catalog_path) - return CodexRoutingOptions(routing_models(state)) + return CodexRoutingOptions(models, catalog_path=catalog_path) + return CodexRoutingOptions( + [codex_routing.codex_model_id(model) for model in routing_models(state)] + ) def _codex_home_config_path() -> Path: @@ -475,9 +476,7 @@ def _codex_home_config_path() -> Path: return Path.home() / ".codex" / "config.toml" -def _v2_pre_tool_use_hooks( - state: dict, available_models: list[str], *, preserve_model_ids: bool = False -) -> list[dict]: +def _v2_pre_tool_use_hooks(state: dict, available_models: list[str]) -> list[dict]: doc = read_toml_safe(_codex_home_config_path()) configured_hooks = doc.get("hooks") existing = configured_hooks.get("PreToolUse") if isinstance(configured_hooks, dict) else None @@ -485,7 +484,6 @@ def _v2_pre_tool_use_hooks( existing if isinstance(existing, list) else [], state, available_models=available_models, - preserve_model_ids=preserve_model_ids, ) @@ -525,11 +523,7 @@ def launch_codex( if routing_options.catalog_path is not None: overlay["model_catalog_json"] = str(routing_options.catalog_path) overlay["hooks"] = { - "PreToolUse": _v2_pre_tool_use_hooks( - state, - available_models, - preserve_model_ids=routing_options.preserve_model_ids, - ), + "PreToolUse": _v2_pre_tool_use_hooks(state, available_models), } config_args = codex_config_args(overlay) app_port = _free_port() @@ -557,7 +551,6 @@ def launch_codex( workspace=workspace, token_provider=lambda: get_databricks_token(workspace, profile), switch_message_fn=format_routing_notice, - preserve_model_ids=routing_options.preserve_model_ids, log_path=CODEX_INTERPOSER_LOG, ) tui_url = _loopback_websocket_url(tui_port) diff --git a/tests/test_cli.py b/tests/test_cli.py index e577fc5f..8ef41712 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -674,7 +674,6 @@ def _invoke_codex_subagent_hook(token_env): "my-profile", "--model", "system.ai.gpt-5-6-sol", - "--preserve-model-ids", ], input='{"tool_name":"collaboration.spawn_agent","tool_input":{"message":"fix it"}}', env={"ENABLE_SMART_ROUTING_V2": "1", **token_env}, @@ -682,7 +681,6 @@ def _invoke_codex_subagent_hook(token_env): assert result.exit_code == 0, result.output assert json.loads(result.output) == routed - assert mock_route.call_args.kwargs["preserve_model_ids"] is True return mock_token, mock_route def test_codex_subagent_hook_reuses_fresh_oauth_token(self, monkeypatch): diff --git a/tests/test_codex_routing.py b/tests/test_codex_routing.py index 8c6595eb..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( @@ -251,7 +251,6 @@ def test_spawn_rewrite_preserves_custom_catalog_model_id(monkeypatch): workspace=WS, token="token", available_models=["system.ai.gpt-5-5"], - preserve_model_ids=True, ) assert output["hookSpecificOutput"]["updatedInput"]["model"] == "system.ai.gpt-5-5" @@ -400,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, @@ -415,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 3e421c40..c4ce184e 100644 --- a/tests/test_codex_smart_routing_v2.py +++ b/tests/test_codex_smart_routing_v2.py @@ -61,7 +61,6 @@ def test_managed_catalog_models_take_priority_and_are_deduplicated(self, tmp_pat assert options == v2.CodexRoutingOptions( ["system.ai.gpt-5-5", "system.ai.glm-5-3"], - preserve_model_ids=True, catalog_path=managed_catalog, ) @@ -76,7 +75,6 @@ def test_falls_back_to_profile_catalog(self, tmp_path, monkeypatch): assert options.models == ["system.ai.gpt-5-6-sol"] assert options.catalog_path == catalog - assert options.preserve_model_ids is True def test_falls_back_to_user_catalog(self, tmp_path, monkeypatch): managed, profile, user = self._configure_paths(tmp_path, monkeypatch) @@ -90,7 +88,6 @@ def test_falls_back_to_user_catalog(self, tmp_path, monkeypatch): assert options.models == ["system.ai.glm-5-3"] assert options.catalog_path == catalog - assert options.preserve_model_ids is True def test_falls_back_to_cached_workspace_models(self, tmp_path, monkeypatch): self._configure_paths(tmp_path, monkeypatch) @@ -100,7 +97,7 @@ def test_falls_back_to_cached_workspace_models(self, tmp_path, monkeypatch): } assert v2.codex_routing_options(state) == v2.CodexRoutingOptions( - ["system.ai.gpt-5-6-sol", "system.ai.glm-5-2"] + ["gpt-5.6-sol", "system.ai.glm-5-2"] ) @@ -125,7 +122,6 @@ def test_codex_smart_routing_launch_dispatches_to_v2(self, monkeypatch, tool_arg calls = [] routing_options = v2.CodexRoutingOptions( ["system.ai.gpt-5-5"], - preserve_model_ids=True, catalog_path=Path("/catalog.json"), ) monkeypatch.setenv(v2.ENV_VAR, "1") @@ -197,7 +193,7 @@ def test_codex_launch_normalizes_cached_bootstrap_model(self, monkeypatch): monkeypatch.setattr( v2, "codex_routing_options", - lambda state: v2.CodexRoutingOptions(["system.ai.gpt-5-6-luna"]), + lambda state: v2.CodexRoutingOptions(["gpt-5.6-luna"]), ) def launch_v2(state, tool_args, **kwargs): @@ -273,7 +269,6 @@ def start_interposer(*args, **kwargs): start_model="gpt-start", routing_options=v2.CodexRoutingOptions( ["system.ai.gpt-5-6-sol", "system.ai.glm-5-2"], - preserve_model_ids=True, catalog_path=Path("/catalog.json"), ), render_overlay=codex.render_overlay, @@ -301,7 +296,6 @@ def start_interposer(*args, **kwargs): assert "codex-router-hook route-subagent" in hook_override assert f"--host {WS}" in hook_override assert "--profile myprof" in hook_override - assert "--preserve-model-ids" 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[12:] == [ @@ -328,7 +322,6 @@ def start_interposer(*args, **kwargs): assert interposer_args["kwargs"]["token_provider"]() == "token-2" assert token_calls == [(WS, "myprof"), (WS, "myprof")] assert interposer_args["kwargs"]["switch_message_fn"] is v2.format_routing_notice - assert interposer_args["kwargs"]["preserve_model_ids"] is True assert stopped == [True] assert processes[0].terminated is True @@ -551,11 +544,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.", ), @@ -592,7 +585,6 @@ def test_preserves_selected_custom_catalog_slug(self): ), None, ), - preserve_model_ids=True, ) result = sess.on_tui_frame(self._turn_start("system.ai.glm-5-3")) From cba5e5c5235a20a5008d8edffe453acb422db073 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Tue, 8 Sep 2026 17:44:10 +0000 Subject: [PATCH 3/5] update --- src/ucode/smart_routing/codex_hooks.py | 25 ++++--------------------- 1 file changed, 4 insertions(+), 21 deletions(-) diff --git a/src/ucode/smart_routing/codex_hooks.py b/src/ucode/smart_routing/codex_hooks.py index 98768324..421cabda 100644 --- a/src/ucode/smart_routing/codex_hooks.py +++ b/src/ucode/smart_routing/codex_hooks.py @@ -53,33 +53,19 @@ def _routing_hook_groups(state: dict) -> dict[str, list[dict]]: def merge_pre_tool_use_hooks( - existing: list[dict], - state: dict, - *, - available_models: list[str], + existing: list[dict], state: dict, *, available_models: list[str] ) -> list[dict]: """Add the ucode spawn hook to an existing Codex PreToolUse hook list.""" doc = {"hooks": {"PreToolUse": copy.deepcopy(existing)}} hooks.sync_managed_hooks( doc, ROUTING_HOOK_COMMAND_MARKER, - { - "PreToolUse": [ - _pre_tool_use_hook_group( - state, - available_models=available_models, - ) - ] - }, + {"PreToolUse": [_pre_tool_use_hook_group(state, available_models=available_models)]}, ) return doc["hooks"]["PreToolUse"] -def _pre_tool_use_hook_group( - state: dict, - *, - available_models: list[str] | None = None, -) -> dict: +def _pre_tool_use_hook_group(state: dict, *, available_models: list[str] | None = None) -> dict: route_argv = _routing_hook_argv( state, "route-subagent", @@ -92,10 +78,7 @@ def _pre_tool_use_hook_group( def _routing_hook_argv( - state: dict, - event: str, - *, - available_models: list[str] | None = None, + state: dict, event: str, *, available_models: list[str] | None = None ) -> list[str]: workspace = str(state.get("workspace") or "") argv = [ From ba608465c3fdd73ac8f5b66c0213b08071a7cc20 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Tue, 8 Sep 2026 17:57:59 +0000 Subject: [PATCH 4/5] update --- src/ucode/agents/codex.py | 2 +- src/ucode/smart_routing/v2.py | 20 ++++++++------ tests/test_codex_smart_routing_v2.py | 39 ++++++++++++++++++++++++++-- 3 files changed, 50 insertions(+), 11 deletions(-) diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 395159d0..4fe2684e 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -527,7 +527,7 @@ def _launch_smart_routing(state: dict, tool_args: list[str]) -> None: managed_model = default_model(state) routing_options = smart_routing_v2.codex_routing_options(state) - models = routing_options.models + models = smart_routing_v2.codex_models_for_routing(routing_options) 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( diff --git a/src/ucode/smart_routing/v2.py b/src/ucode/smart_routing/v2.py index a2db07fb..1c6ab3c6 100644 --- a/src/ucode/smart_routing/v2.py +++ b/src/ucode/smart_routing/v2.py @@ -441,7 +441,7 @@ def _model_catalog() -> tuple[Path, list[str]] | None: catalog_path = Path(configured_path).expanduser() rows = read_json_safe(catalog_path).get("models") if not isinstance(rows, list): - continue + return catalog_path, [] models: list[str] = [] seen_models: set[str] = set() for row in rows: @@ -453,20 +453,24 @@ def _model_catalog() -> tuple[Path, list[str]] | None: continue seen_models.add(model) models.append(model) - if models: - return catalog_path, models + return catalog_path, models return None def codex_routing_options(state: dict) -> CodexRoutingOptions: - """Prefer custom-catalog models, then fall back to persisted workspace models.""" + """Use a configured catalog exclusively, or the existing persisted model list.""" catalog = _model_catalog() if catalog is not None: catalog_path, models = catalog return CodexRoutingOptions(models, catalog_path=catalog_path) - return CodexRoutingOptions( - [codex_routing.codex_model_id(model) for model in routing_models(state)] - ) + return CodexRoutingOptions(routing_models(state)) + + +def codex_models_for_routing(options: CodexRoutingOptions) -> list[str]: + """Return the model IDs Codex should receive for the selected model source.""" + if options.catalog_path is not None: + return options.models + return [codex_routing.codex_model_id(model) for model in options.models] def _codex_home_config_path() -> Path: @@ -508,7 +512,7 @@ def launch_codex( profile = state.get("profile") os.environ[OAUTH_TOKEN_ENV_VAR] = get_databricks_token(workspace, profile) - available_models = routing_options.models + available_models = codex_models_for_routing(routing_options) if not available_models: print_note( "Smart routing model metadata is unavailable; starting Codex on gpt-5.6-luna " diff --git a/tests/test_codex_smart_routing_v2.py b/tests/test_codex_smart_routing_v2.py index c4ce184e..137df673 100644 --- a/tests/test_codex_smart_routing_v2.py +++ b/tests/test_codex_smart_routing_v2.py @@ -97,9 +97,44 @@ def test_falls_back_to_cached_workspace_models(self, tmp_path, monkeypatch): } assert v2.codex_routing_options(state) == v2.CodexRoutingOptions( - ["gpt-5.6-sol", "system.ai.glm-5-2"] + ["system.ai.gpt-5-6-sol", "system.ai.glm-5-2"] ) + 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") + + options = v2.codex_routing_options({"codex_models": ["cached-model"]}) + + assert options == v2.CodexRoutingOptions([], 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") + + options = v2.codex_routing_options({"codex_models": ["cached-model"]}) + + assert options == v2.CodexRoutingOptions([], catalog_path=catalog) + + def test_models_for_routing_preserves_catalog_slugs_and_maps_cached_models(self): + catalog_options = v2.CodexRoutingOptions( + ["system.ai.gpt-5-5"], catalog_path=Path("/catalog.json") + ) + cached_options = v2.CodexRoutingOptions(["system.ai.gpt-5-6-sol", "system.ai.glm-5-2"]) + + assert v2.codex_models_for_routing(catalog_options) == ["system.ai.gpt-5-5"] + assert v2.codex_models_for_routing(cached_options) == [ + "gpt-5.6-sol", + "system.ai.glm-5-2", + ] + class TestLaunchCodex: def test_rejects_unsupported_codex_version(self, monkeypatch): @@ -193,7 +228,7 @@ def test_codex_launch_normalizes_cached_bootstrap_model(self, monkeypatch): monkeypatch.setattr( v2, "codex_routing_options", - lambda state: v2.CodexRoutingOptions(["gpt-5.6-luna"]), + lambda state: v2.CodexRoutingOptions(["system.ai.gpt-5-6-luna"]), ) def launch_v2(state, tool_args, **kwargs): From dd9b15f6f5269f5dc0c0606df96edff1dae1723f Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Tue, 8 Sep 2026 18:09:41 +0000 Subject: [PATCH 5/5] update --- src/ucode/agents/codex.py | 6 +-- src/ucode/smart_routing/v2.py | 48 +++++++------------ tests/test_codex_smart_routing_v2.py | 71 ++++++++++++---------------- 3 files changed, 48 insertions(+), 77 deletions(-) diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 4fe2684e..02f44b5b 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -526,8 +526,7 @@ def _launch_smart_routing(state: dict, tool_args: list[str]) -> None: ) managed_model = default_model(state) - routing_options = smart_routing_v2.codex_routing_options(state) - models = smart_routing_v2.codex_models_for_routing(routing_options) + 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( @@ -535,7 +534,8 @@ def _launch_smart_routing(state: dict, tool_args: list[str]) -> None: tool_args, binary=binary, start_model=start_model, - routing_options=routing_options, + available_models=models, + catalog_path=catalog_path, render_overlay=render_overlay, ) diff --git a/src/ucode/smart_routing/v2.py b/src/ucode/smart_routing/v2.py index 1c6ab3c6..b404cb43 100644 --- a/src/ucode/smart_routing/v2.py +++ b/src/ucode/smart_routing/v2.py @@ -12,7 +12,6 @@ import urllib.request import uuid from collections.abc import Callable -from dataclasses import dataclass from pathlib import Path from typing import NoReturn, TextIO @@ -59,12 +58,6 @@ _ANTHROPIC_AIGW_MODEL_RE = re.compile(r"^anthropic-aigw-[0-9a-fA-F]{8}-(.+)$") -@dataclass(frozen=True) -class CodexRoutingOptions: - models: list[str] - catalog_path: Path | None = None - - def enabled() -> bool: return os.environ.get(ENV_VAR) == "1" @@ -423,18 +416,16 @@ def route_prompt(prompt: str) -> claude_pty.FirstPromptRoute: def _model_catalog() -> tuple[Path, list[str]] | None: """Read the first configured Codex model catalog using Codex config precedence.""" - try: - from ucode.agents import codex - - paths = [codex._managed_config_path(), codex.CODEX_CONFIG_PATH, _codex_home_config_path()] - except (ImportError, OSError): - return None + from ucode.agents import codex - seen_paths: set[Path] = set() - for config_path in paths: - if config_path is None or config_path in seen_paths: + 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 - seen_paths.add(config_path) configured_path = read_toml_safe(config_path).get("model_catalog_json") if not isinstance(configured_path, str) or not configured_path.strip(): continue @@ -457,20 +448,13 @@ def _model_catalog() -> tuple[Path, list[str]] | None: return None -def codex_routing_options(state: dict) -> CodexRoutingOptions: - """Use a configured catalog exclusively, or the existing persisted model list.""" +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 CodexRoutingOptions(models, catalog_path=catalog_path) - return CodexRoutingOptions(routing_models(state)) - - -def codex_models_for_routing(options: CodexRoutingOptions) -> list[str]: - """Return the model IDs Codex should receive for the selected model source.""" - if options.catalog_path is not None: - return options.models - return [codex_routing.codex_model_id(model) for model in options.models] + return models, catalog_path + return [codex_routing.codex_model_id(model) for model in routing_models(state)], None def _codex_home_config_path() -> Path: @@ -497,7 +481,8 @@ def launch_codex( *, binary: str, start_model: str | None, - routing_options: CodexRoutingOptions, + available_models: list[str], + catalog_path: Path | None, render_overlay: Callable[..., dict], ) -> NoReturn: workspace = state.get("workspace") @@ -512,7 +497,6 @@ def launch_codex( profile = state.get("profile") os.environ[OAUTH_TOKEN_ENV_VAR] = get_databricks_token(workspace, profile) - available_models = codex_models_for_routing(routing_options) if not available_models: print_note( "Smart routing model metadata is unavailable; starting Codex on gpt-5.6-luna " @@ -524,8 +508,8 @@ def launch_codex( state.get("profile"), use_pat=bool(state.get("use_pat")), ) - if routing_options.catalog_path is not None: - overlay["model_catalog_json"] = str(routing_options.catalog_path) + 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_smart_routing_v2.py b/tests/test_codex_smart_routing_v2.py index 137df673..9180814e 100644 --- a/tests/test_codex_smart_routing_v2.py +++ b/tests/test_codex_smart_routing_v2.py @@ -57,12 +57,10 @@ def test_managed_catalog_models_take_priority_and_are_deduplicated(self, tmp_pat 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") - options = v2.codex_routing_options({"codex_models": ["cached-model"]}) + models, catalog_path = v2.configured_codex_models({"codex_models": ["cached-model"]}) - assert options == v2.CodexRoutingOptions( - ["system.ai.gpt-5-5", "system.ai.glm-5-3"], - catalog_path=managed_catalog, - ) + 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) @@ -71,10 +69,10 @@ def test_falls_back_to_profile_catalog(self, tmp_path, monkeypatch): managed.write_text('model_provider = "managed"\n', encoding="utf-8") profile.write_text(f'model_catalog_json = "{catalog}"\n', encoding="utf-8") - options = v2.codex_routing_options({}) + models, catalog_path = v2.configured_codex_models({}) - assert options.models == ["system.ai.gpt-5-6-sol"] - assert options.catalog_path == catalog + 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) @@ -84,10 +82,10 @@ def test_falls_back_to_user_catalog(self, tmp_path, monkeypatch): profile.write_text('model_provider = "profile"\n', encoding="utf-8") user.write_text(f'model_catalog_json = "{catalog}"\n', encoding="utf-8") - options = v2.codex_routing_options({}) + models, catalog_path = v2.configured_codex_models({}) - assert options.models == ["system.ai.glm-5-3"] - assert options.catalog_path == catalog + 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) @@ -96,8 +94,9 @@ def test_falls_back_to_cached_workspace_models(self, tmp_path, monkeypatch): "oss_models": ["system.ai.glm-5-2"], } - assert v2.codex_routing_options(state) == v2.CodexRoutingOptions( - ["system.ai.gpt-5-6-sol", "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( @@ -108,9 +107,10 @@ def test_configured_empty_catalog_does_not_fall_back_to_cached_models( self._write_catalog(catalog, []) managed.write_text(f'model_catalog_json = "{catalog}"\n', encoding="utf-8") - options = v2.codex_routing_options({"codex_models": ["cached-model"]}) + models, catalog_path = v2.configured_codex_models({"codex_models": ["cached-model"]}) - assert options == v2.CodexRoutingOptions([], catalog_path=catalog) + assert models == [] + assert catalog_path == catalog def test_configured_unreadable_catalog_does_not_fall_back_to_cached_models( self, tmp_path, monkeypatch @@ -119,21 +119,10 @@ def test_configured_unreadable_catalog_does_not_fall_back_to_cached_models( catalog = tmp_path / "missing-models.json" managed.write_text(f'model_catalog_json = "{catalog}"\n', encoding="utf-8") - options = v2.codex_routing_options({"codex_models": ["cached-model"]}) + models, catalog_path = v2.configured_codex_models({"codex_models": ["cached-model"]}) - assert options == v2.CodexRoutingOptions([], catalog_path=catalog) - - def test_models_for_routing_preserves_catalog_slugs_and_maps_cached_models(self): - catalog_options = v2.CodexRoutingOptions( - ["system.ai.gpt-5-5"], catalog_path=Path("/catalog.json") - ) - cached_options = v2.CodexRoutingOptions(["system.ai.gpt-5-6-sol", "system.ai.glm-5-2"]) - - assert v2.codex_models_for_routing(catalog_options) == ["system.ai.gpt-5-5"] - assert v2.codex_models_for_routing(cached_options) == [ - "gpt-5.6-sol", - "system.ai.glm-5-2", - ] + assert models == [] + assert catalog_path == catalog class TestLaunchCodex: @@ -155,14 +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 = [] - routing_options = v2.CodexRoutingOptions( - ["system.ai.gpt-5-5"], - catalog_path=Path("/catalog.json"), - ) + models = ["system.ai.gpt-5-5"] + catalog_path = Path("/catalog.json") monkeypatch.setenv(v2.ENV_VAR, "1") monkeypatch.setattr(codex, "default_model", lambda state: None) monkeypatch.setattr(codex, "clear_model_preferences", lambda state: False) - monkeypatch.setattr(v2, "codex_routing_options", lambda state: routing_options) + monkeypatch.setattr(v2, "configured_codex_models", lambda state: (models, catalog_path)) def launch_v2(state, tool_args, **kwargs): calls.append((state, tool_args, kwargs)) @@ -182,7 +169,8 @@ def launch_v2(state, tool_args, **kwargs): { "binary": "codex", "start_model": "system.ai.gpt-5-5", - "routing_options": routing_options, + "available_models": models, + "catalog_path": catalog_path, "render_overlay": codex.render_overlay, }, ) @@ -227,8 +215,8 @@ def test_codex_launch_normalizes_cached_bootstrap_model(self, monkeypatch): monkeypatch.setattr(codex, "default_model", lambda state: None) monkeypatch.setattr( v2, - "codex_routing_options", - lambda state: v2.CodexRoutingOptions(["system.ai.gpt-5-6-luna"]), + "configured_codex_models", + lambda state: (["gpt-5.6-luna"], None), ) def launch_v2(state, tool_args, **kwargs): @@ -302,10 +290,8 @@ def start_interposer(*args, **kwargs): ["--search"], binary="codex", start_model="gpt-start", - routing_options=v2.CodexRoutingOptions( - ["system.ai.gpt-5-6-sol", "system.ai.glm-5-2"], - catalog_path=Path("/catalog.json"), - ), + available_models=["system.ai.gpt-5-6-sol", "system.ai.glm-5-2"], + catalog_path=Path("/catalog.json"), render_overlay=codex.render_overlay, ) @@ -440,7 +426,8 @@ def test_missing_cached_models_starts_with_bootstrap_model(self, monkeypatch): [], binary="codex", start_model="gpt-5.6-luna", - routing_options=v2.CodexRoutingOptions([]), + available_models=[], + catalog_path=None, render_overlay=codex.render_overlay, )