Lilly/codex use model catalog

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,
     )
 

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)

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,
     )
 

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),
     }

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(