smart routing: route Codex v2 subagents

src/ucode/agents/codex.py

@@ -29,6 +29,7 @@
 from ucode.managed_files import OS, current_os, write_managed_file
 from ucode.smart_routing.codex_hooks import (
     remove_smart_routing_hooks,
+    routing_models,
     sync_smart_routing_hooks,
 )
 from ucode.state import mark_tool_managed, save_state
@@ -432,9 +433,9 @@ def default_model(state: dict) -> str | None:
     """
     if isinstance(state.get("codex_default_model"), str):
         return state.get("codex_default_model")
-    codex_models = state.get("codex_models") or []
+    models = routing_models(state)
     parsed: list[tuple[str, tuple[int, int | None, int | None, str]]] = [
-        (mid, gpt) for mid in codex_models if (gpt := _parse_gpt(mid)) is not None
+        (mid, gpt) for mid in models if (gpt := _parse_gpt(mid)) is not None
     ]
     if parsed:
 
@@ -449,7 +450,7 @@ def _gpt_version_key(entry: tuple[str, tuple[int, int | None, int | None, str]])
     # after stripping the system.ai. prefix). gpt-oss-* models are confirmed
     # routable through the responses API; non-GPT ids (e.g. moonshotai/kimi-k2.5)
     # would be rejected by the gateway, so they stay excluded.
-    gpt_family = [m for m in codex_models if _is_gpt_family(m)]
+    gpt_family = [m for m in models if _is_gpt_family(m)]
     return gpt_family[0] if gpt_family else None
 
 

src/ucode/smart_routing/claude_pty.py

@@ -25,6 +25,8 @@
 from collections.abc import Callable
 from pathlib import Path
 
+from ucode.smart_routing import routing
+
 MAX_MODEL_NAME_LEN = 200
 CONFIRM_TIMEOUT_S = 3.0
 SWITCH_TIMEOUT_S = 6.0
@@ -58,14 +60,7 @@ def valid_model_name(name: object) -> bool:
 
 def switch_message(model: str, reason: str) -> str:
     """Format the routed-model notice shown in Claude Code."""
-    lines = [
-        "Using Unity Gateway Smart Router.",
-        f"Selected Model : {model}",
-        f"Reason : {reason}",
-    ]
-    width = max(len(line) for line in lines)
-    border = "─" * (width + 2)
-    return "\n".join([f"┌{border}┐", *(f"│ {line:<{width}} │" for line in lines), f"└{border}┘"])
+    return routing.format_switch_message(model, reason)
 
 
 class ConfirmationState:

src/ucode/smart_routing/codex_hooks.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import copy
 import shlex
 import subprocess
 
@@ -11,6 +12,16 @@
 ROUTING_HOOK_COMMAND_MARKER = "codex-router-hook"
 
 
+def routing_models(state: dict) -> list[str]:
+    """Return the configured model services compatible with Codex routing."""
+    models: list[str] = []
+    for key in ("codex_models", "oss_models"):
+        values = state.get(key)
+        if isinstance(values, list):
+            models.extend(value for value in values if isinstance(value, str) and value)
+    return list(dict.fromkeys(models))
+
+
 def sync_smart_routing_hooks(doc: dict, state: dict, *, enabled: bool) -> None:
     """Synchronize ucode-managed routing hooks in a Codex config document."""
     groups = _routing_hook_groups(state) if enabled else {}
@@ -23,16 +34,10 @@ def remove_smart_routing_hooks(doc: dict) -> bool:
 
 
 def _routing_hook_groups(state: dict) -> dict[str, list[dict]]:
-    route_argv = _routing_hook_argv(state, "route-subagent")
     session_argv = _routing_hook_argv(state, "session-start")
     subagent_argv = _routing_hook_argv(state, "record-subagent")
     return {
-        "PreToolUse": [
-            {
-                "matcher": "Agent|.*spawn_agent$",
-                "hooks": [_routing_command_hook(route_argv, status="Routing subagent model")],
-            }
-        ],
+        "PreToolUse": [_pre_tool_use_hook_group(state)],
         "SessionStart": [
             {
                 "matcher": "startup|resume|clear",
@@ -47,7 +52,34 @@ def _routing_hook_groups(state: dict) -> dict[str, list[dict]]:
     }
 
 
-def _routing_hook_argv(state: dict, event: str) -> list[str]:
+def merge_pre_tool_use_hooks(
+    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)]},
+    )
+    return doc["hooks"]["PreToolUse"]
+
+
+def _pre_tool_use_hook_group(state: dict, *, available_models: list[str] | None = None) -> dict:
+    route_argv = _routing_hook_argv(
+        state,
+        "route-subagent",
+        available_models=available_models,
+    )
+    return {
+        "matcher": "Agent|.*spawn_agent$",
+        "hooks": [_routing_command_hook(route_argv, status="Routing subagent model")],
+    }
+
+
+def _routing_hook_argv(
+    state: dict, event: str, *, available_models: list[str] | None = None
+) -> list[str]:
     workspace = str(state.get("workspace") or "")
     argv = [
         build_auth_token_argv(workspace, state.get("profile"), use_pat=bool(state.get("use_pat")))[
@@ -64,7 +96,8 @@ def _routing_hook_argv(state: dict, event: str) -> list[str]:
         argv += ["--profile", profile]
     if state.get("use_pat"):
         argv.append("--use-pat")
-    for model in state.get("codex_models") or []:
+    models = available_models if available_models is not None else routing_models(state)
+    for model in models:
         if isinstance(model, str) and model:
             argv += ["--model", model]
     return argv

src/ucode/smart_routing/codex_routing.py

@@ -18,6 +18,7 @@
 from ucode.config_io import APP_DIR
 from ucode.databricks import get_databricks_token
 from ucode.smart_routing import routing
+from ucode.smart_routing.codex_hooks import routing_models
 from ucode.smart_routing.routing import RoutingDecision
 
 ROUTER_NAME = routing.ROUTER_NAME
@@ -50,8 +51,8 @@ def route_launch_model(state: dict, tool_args: list[str]):
     if task is None:
         return None, None
     workspace = state.get("workspace")
-    models = state.get("codex_models")
-    if not isinstance(workspace, str) or not isinstance(models, list):
+    models = routing_models(state)
+    if not isinstance(workspace, str) or not models:
         return None, "workspace model metadata is unavailable"
     try:
         token = get_databricks_token(workspace, state.get("profile"))

src/ucode/smart_routing/routing.py

@@ -23,6 +23,22 @@
 ROUTER_NAME = "task_v1"
 ROUTING_PATH = "/ai-gateway/routing/v1/routes:select"
 REQUEST_TIMEOUT_S = 30.0
+SUBAGENT_ROUTING_DISCLAIMER = (
+    "Spawned subagents are routed independently based on their own complexity."
+)
+
+
+def format_switch_message(model: str, reason: str) -> str:
+    """Format the routed-model notice shared by Codex and Claude Code."""
+    lines = [
+        "Using Unity Gateway Smart Router.",
+        f"Selected Model : {model}",
+        f"Reason : {reason}",
+        SUBAGENT_ROUTING_DISCLAIMER,
+    ]
+    width = max(len(line) for line in lines)
+    border = "─" * (width + 2)
+    return "\n".join([f"┌{border}┐", *(f"│ {line:<{width}} │" for line in lines), f"└{border}┘"])
 
 
 @dataclass(frozen=True)

src/ucode/smart_routing/v2.py

@@ -14,7 +14,7 @@
 
 import tomlkit
 
-from ucode.config_io import APP_DIR, read_json_safe, write_json_file
+from ucode.config_io import APP_DIR, read_json_safe, read_toml_safe, write_json_file
 from ucode.constants import LOOPBACK_HOST
 from ucode.databricks import (
     build_auth_token_argv,
@@ -23,6 +23,7 @@
 )
 from ucode.smart_routing import codex_interposer, routing
 from ucode.smart_routing.claude_hooks import FIRST_PROMPT_SOCKET_ENV, sync_first_prompt_hook
+from ucode.smart_routing.codex_hooks import merge_pre_tool_use_hooks, routing_models
 from ucode.ui import print_note
 
 ENV_VAR = "ENABLE_SMART_ROUTING_V2"
@@ -71,14 +72,7 @@ def _wait_for_app_server(port: int, timeout: float) -> bool:
 
 
 def _switch_message(model: str, reason: str) -> str:
-    lines = [
-        "Using Unity Gateway Smart Router.",
-        f"Selected Model : {model}",
-        f"Reason : {reason}",
-    ]
-    width = max(len(line) for line in lines)
-    border = "─" * (width + 2)
-    return "\n".join([f"┌{border}┐", *(f"│ {line:<{width}} │" for line in lines), f"└{border}┘"])
+    return routing.format_switch_message(model, reason)
 
 
 def _canonical_claude_model_id(model: str) -> str:
@@ -232,6 +226,11 @@ def _toml_value(value: str | int | float | bool | list[object] | dict[str, objec
         item = tomlkit.inline_table()
         item.update(value)
         return item.as_string()
+    if isinstance(value, list) and any(isinstance(entry, dict) for entry in value):
+        wrapper = tomlkit.inline_table()
+        wrapper["value"] = value
+        rendered = wrapper.as_string()
+        return rendered.removeprefix("{value = ").removesuffix("}")
     return tomlkit.item(value).as_string()
 
 
@@ -240,12 +239,12 @@ def _codex_config_args(overlay: dict) -> list[str]:
     for key, value in overlay.items():
         # This is Codex's AI Gateway transport definition, not Unity Catalog
         # Model Provider Service support; smart routing still cannot use --provider.
-        if key == "model_providers" and isinstance(value, dict):
+        if key in {"hooks", "model_providers"} and isinstance(value, dict):
             for provider_name, provider_config in value.items():
                 args.extend(
                     [
                         "--config",
-                        f"model_providers.{provider_name}={_toml_value(provider_config)}",
+                        f"{key}.{provider_name}={_toml_value(provider_config)}",
                     ]
                 )
         else:
@@ -256,12 +255,25 @@ def _codex_config_args(overlay: dict) -> list[str]:
 # 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."""
-    models: list[str] = []
-    for key in ("codex_models", "oss_models"):
-        values = state.get(key)
-        if isinstance(values, list):
-            models.extend(value for value in values if isinstance(value, str) and value)
-    return list(dict.fromkeys(models))
+    return routing_models(state)
+
+
+def _codex_home_config_path() -> Path:
+    codex_home = os.environ.get("CODEX_HOME")
+    if codex_home:
+        return Path(codex_home).expanduser() / "config.toml"
+    return Path.home() / ".codex" / "config.toml"
+
+
+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
+    return merge_pre_tool_use_hooks(
+        existing if isinstance(existing, list) else [],
+        state,
+        available_models=available_models,
+    )
 
 
 def launch_codex(
@@ -296,6 +308,9 @@ def launch_codex(
         state.get("profile"),
         use_pat=bool(state.get("use_pat")),
     )
+    overlay["hooks"] = {
+        "PreToolUse": _v2_pre_tool_use_hooks(state, available_models),
+    }
     config_args = _codex_config_args(overlay)
     app_port = _free_port()
     app_server_url = _loopback_websocket_url(app_port)

tests/test_agent_codex.py

@@ -242,6 +242,7 @@ def test_smart_routing_writes_profile_scoped_hooks(self, tmp_path, monkeypatch):
                 "workspace": WS,
                 "profile": "prod",
                 "codex_models": ["databricks-gpt-5", "databricks-gpt-5-5"],
+                "oss_models": ["system.ai.glm-5-2"],
                 codex.SMART_ROUTING_STATE_KEY: True,
             }
         )
@@ -259,6 +260,7 @@ def test_smart_routing_writes_profile_scoped_hooks(self, tmp_path, monkeypatch):
         assert "--host https://example.databricks.com" in route_command
         assert "--profile prod" in route_command
         assert "--model databricks-gpt-5-5" in route_command
+        assert "--model system.ai.glm-5-2" in route_command
 
     def test_provider_launch_removes_routing_hooks(self, tmp_path, monkeypatch):
         config_path = tmp_path / ".codex" / "ucode.config.toml"
@@ -409,6 +411,27 @@ def fail(*args, **kwargs):
         assert decision is None
         assert error is None
 
+    def test_route_launch_model_includes_codex_and_oss_models(self, monkeypatch):
+        captured = {}
+        monkeypatch.setattr(codex_routing, "get_databricks_token", lambda *args: "token")
+
+        def request(workspace, token, task, models):
+            captured["models"] = models
+            return None, "expected test stop"
+
+        monkeypatch.setattr(codex_routing, "request_routing_decision", request)
+
+        codex_routing.route_launch_model(
+            {
+                "workspace": WS,
+                "codex_models": ["system.ai.gpt-5-6-sol"],
+                "oss_models": ["system.ai.glm-5-2"],
+            },
+            ["Fix the parser"],
+        )
+
+        assert captured["models"] == ["system.ai.gpt-5-6-sol", "system.ai.glm-5-2"]
+
 
 class TestCodexRemoveLegacyProfile:
     def test_drops_provider_block_on_modern_path(self, tmp_path, monkeypatch):
@@ -537,6 +560,14 @@ def test_default_model_falls_back_to_first_when_no_versioned_gpt(self):
         models = ["system.ai.gpt-oss-120b", "system.ai.gpt-oss-20b"]
         assert codex.default_model({"codex_models": models}) == "system.ai.gpt-oss-120b"
 
+    def test_default_model_includes_oss_models(self):
+        state = {
+            "codex_models": [],
+            "oss_models": ["system.ai.kimi-k3", "system.ai.gpt-oss-120b"],
+        }
+
+        assert codex.default_model(state) == "system.ai.gpt-oss-120b"
+
     def test_default_model_prefers_versioned_gpt_over_oss(self):
         # When both versioned and OSS models are present, the versioned one wins.
         models = ["system.ai.gpt-oss-120b", "system.ai.gpt-5"]

tests/test_codex_routing.py

@@ -6,10 +6,24 @@
 import urllib.error
 
 from ucode.smart_routing import codex_routing
+from ucode.smart_routing.codex_hooks import routing_models
 
 WS = "https://example.databricks.com"
 
 
+def test_routing_models_combines_and_deduplicates_codex_and_oss_models():
+    assert routing_models(
+        {
+            "codex_models": ["system.ai.gpt-5-6-sol", "system.ai.gpt-oss-120b"],
+            "oss_models": ["system.ai.gpt-oss-120b", "system.ai.glm-5-2"],
+        }
+    ) == [
+        "system.ai.gpt-5-6-sol",
+        "system.ai.gpt-oss-120b",
+        "system.ai.glm-5-2",
+    ]
+
+
 class _Response:
     def __init__(self, payload: dict):
         self.payload = payload