Support subagent routing in Codex app

src/ucode/agents/codex.py

@@ -71,6 +71,7 @@
 )
 from ucode.smart_routing import v2 as smart_routing_v2
 from ucode.smart_routing.codex_hooks import (
+    reconcile_smart_routing_hooks_file,
     remove_smart_routing_hooks,
     routing_models,
     sync_smart_routing_hooks,
@@ -87,6 +88,7 @@
 CODEX_CONFIG_DIR = Path.home() / ".codex"
 CODEX_PROFILE_NAME = "ucode"
 CODEX_CONFIG_PATH = CODEX_CONFIG_DIR / f"{CODEX_PROFILE_NAME}.config.toml"
+CODEX_HOOKS_PATH = CODEX_CONFIG_DIR / "hooks.json"
 CODEX_BACKUP_PATH = APP_DIR / "codex-ucode-config.backup.toml"
 CODEX_MODEL_CATALOG_PATH = APP_DIR / "codex-model-catalog.json"
 LEGACY_CODEX_CONFIG_PATH = CODEX_CONFIG_DIR / "config.toml"
@@ -107,6 +109,12 @@
 SMART_ROUTING_STATE_KEY = smart_routing_v2.LEGACY_STATE_KEY
 APP_SERVER_SMART_ROUTING_STARTING_MODEL = "gpt-5.6-luna"
 
+
+def _codex_hooks_path() -> Path:
+    codex_home = os.environ.get("CODEX_HOME")
+    return Path(codex_home).expanduser() / "hooks.json" if codex_home else CODEX_HOOKS_PATH
+
+
 SPEC: ToolSpec = {
     "binary": "codex",
     "package": "@openai/codex",
@@ -917,6 +925,11 @@ def launch(
     *,
     options: LaunchOptions,
 ) -> None:
+    reconcile_smart_routing_hooks_file(
+        _codex_hooks_path(),
+        state,
+        enabled=smart_routing_v2.smart_routing_enabled(),
+    )
     if options.launch_smart_routing:
         _launch_smart_routing(state, tool_args)
         return

src/ucode/cli.py

@@ -131,6 +131,7 @@
 from ucode.skills_state import records_for_scope
 from ucode.smart_routing import v2 as smart_routing_v2
 from ucode.smart_routing.claude_hooks import FIRST_PROMPT_SOCKET_ENV, ROUTE_FIRST_PROMPT_EVENT
+from ucode.smart_routing.codex_hooks import SMART_ROUTING_HOOK_OWNER
 from ucode.state import (
     clear_state,
     get_provider_service,
@@ -2033,12 +2034,14 @@ 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,
+    hook_owner: Annotated[str | None, typer.Option("--hook-owner", hidden=True)] = None,
 ) -> None:
     """Run a Codex smart-routing lifecycle hook."""
     import json
     import sys
 
-    if not smart_routing_v2.smart_routing_enabled():
+    owned_routing_hook = hook_owner == SMART_ROUTING_HOOK_OWNER
+    if not owned_routing_hook and not smart_routing_v2.smart_routing_enabled():
         return
 
     from ucode.smart_routing.codex_routing import (
@@ -2255,7 +2258,12 @@ def _disable_smart_routing_for_subcommand(tool: str, ctx: Any) -> Iterator[None]
     must therefore suppress the flag for the whole ucode launch flow. An
     explicit prompt after `--` remains eligible for routing.
     """
-    if _smart_routing_launch_shape(tool, ctx.args, _has_explicit_prompt(ctx)):
+    codex_app_smart_routing = (
+        tool == "codex" and ctx.args[:1] == ["app"] and smart_routing_v2.smart_routing_enabled()
+    )
+    if codex_app_smart_routing or _smart_routing_launch_shape(
+        tool, ctx.args, _has_explicit_prompt(ctx)
+    ):
         yield
         return
     previous = smart_routing_v2.disable_smart_routing()

src/ucode/smart_routing/codex_hooks.py

@@ -3,13 +3,49 @@
 from __future__ import annotations
 
 import copy
+import json
 import shlex
 import subprocess
+from pathlib import Path
 
+from ucode.config_io import is_dry_run, write_json_file
 from ucode.databricks import build_auth_token_argv
 from ucode.smart_routing import hooks
 
 ROUTING_HOOK_COMMAND_MARKER = "codex-router-hook"
+# Keep the original owner value stable so already-trusted hook definitions stay trusted.
+SMART_ROUTING_HOOK_OWNER = "ug-codex-app-subagent-only"
+
+
+def reconcile_smart_routing_hooks_file(path: Path, state: dict, *, enabled: bool) -> None:
+    """Install or remove UG's user-level Codex hooks for the current routing mode."""
+    if not enabled and not path.exists():
+        return
+    try:
+        if path.exists():
+            doc = json.loads(path.read_text(encoding="utf-8"))
+            if not isinstance(doc, dict):
+                raise ValueError("top-level value is not an object")
+        else:
+            doc = {}
+    except (OSError, UnicodeError, ValueError) as exc:
+        raise RuntimeError(
+            f"Cannot update Codex smart-routing hooks because {path} is not valid JSON: {exc}"
+        ) from exc
+
+    original = copy.deepcopy(doc)
+    sync_smart_routing_hooks(
+        doc,
+        state,
+        enabled=enabled,
+        owner=SMART_ROUTING_HOOK_OWNER,
+    )
+    if doc == original:
+        return
+    if doc or is_dry_run():
+        write_json_file(path, doc)
+    else:
+        path.unlink(missing_ok=True)
 
 
 def routing_models(state: dict) -> list[str]:
@@ -22,22 +58,29 @@ def routing_models(state: dict) -> list[str]:
     return list(dict.fromkeys(models))
 
 
-def sync_smart_routing_hooks(doc: dict, state: dict, *, enabled: bool) -> None:
+def sync_smart_routing_hooks(
+    doc: dict,
+    state: dict,
+    *,
+    enabled: bool,
+    owner: str | None = None,
+) -> None:
     """Synchronize ucode-managed routing hooks in a Codex config document."""
-    groups = _routing_hook_groups(state) if enabled else {}
-    hooks.sync_managed_hooks(doc, ROUTING_HOOK_COMMAND_MARKER, groups)
+    marker = owner or ROUTING_HOOK_COMMAND_MARKER
+    groups = _routing_hook_groups(state, owner=owner) if enabled else {}
+    hooks.sync_managed_hooks(doc, marker, groups)
 
 
 def remove_smart_routing_hooks(doc: dict) -> bool:
     """Remove only ucode-managed smart-routing hooks."""
     return hooks.remove_managed_hooks(doc, ROUTING_HOOK_COMMAND_MARKER)
 
 
-def _routing_hook_groups(state: dict) -> dict[str, list[dict]]:
-    session_argv = _routing_hook_argv(state, "session-start")
-    subagent_argv = _routing_hook_argv(state, "record-subagent")
+def _routing_hook_groups(state: dict, *, owner: str | None = None) -> dict[str, list[dict]]:
+    session_argv = _routing_hook_argv(state, "session-start", owner=owner)
+    subagent_argv = _routing_hook_argv(state, "record-subagent", owner=owner)
     return {
-        "PreToolUse": [_pre_tool_use_hook_group(state)],
+        "PreToolUse": [_pre_tool_use_hook_group(state, owner=owner)],
         "SessionStart": [
             {
                 "matcher": "startup|resume|clear",
@@ -52,24 +95,17 @@ 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]
-) -> 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:
+def _pre_tool_use_hook_group(
+    state: dict,
+    *,
+    available_models: list[str] | None = None,
+    owner: str | None = None,
+) -> dict:
     route_argv = _routing_hook_argv(
         state,
         "route-subagent",
         available_models=available_models,
+        owner=owner,
     )
     return {
         "matcher": "Agent|.*spawn_agent$",
@@ -78,7 +114,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,
+    owner: str | None = None,
 ) -> list[str]:
     workspace = str(state.get("workspace") or "")
     argv = [
@@ -88,6 +128,8 @@ def _routing_hook_argv(
         ROUTING_HOOK_COMMAND_MARKER,
         event,
     ]
+    if owner:
+        argv += ["--hook-owner", owner]
     if event != "route-subagent":
         return argv
     argv += ["--host", workspace]

src/ucode/smart_routing/v2.py

@@ -19,7 +19,7 @@
     custom_catalog_models,
     custom_catalog_path,
 )
-from ucode.config_io import APP_DIR, read_json_safe, read_toml_safe, write_json_file
+from ucode.config_io import APP_DIR, read_json_safe, write_json_file
 from ucode.constants import LOOPBACK_HOST
 from ucode.custom_oauth import custom_oauth_cli_enabled, get_custom_client_token
 from ucode.databricks import (
@@ -36,7 +36,7 @@
     sync_first_prompt_hook,
     sync_smart_routing_hooks,
 )
-from ucode.smart_routing.codex_hooks import merge_pre_tool_use_hooks, routing_models
+from ucode.smart_routing.codex_hooks import routing_models
 from ucode.ui import print_warning
 
 ENABLE_SMART_ROUTING_ENV_VAR = "ENABLE_SMART_ROUTING_V2"
@@ -543,24 +543,6 @@ def _cached_routing_models(state: dict) -> list[str]:
     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(
     state: dict,
     tool_args: list[str],
@@ -599,13 +581,10 @@ def launch_codex(
     catalog_path = custom_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),
-    }
     config_args = codex_config_args(overlay)
     if not first_prompt_routing_enabled():
         # Subagent-only routing needs neither the app-server nor the interposer:
-        # the hooks ride in the CLI config, so launch the TUI directly.
+        # the reconciled user hooks handle subagents, so launch the TUI directly.
         exec_or_spawn([binary, *config_args, *tool_args])
     app_port = _free_port()
     app_server_url = _loopback_websocket_url(app_port)

tests/conftest.py

@@ -28,6 +28,7 @@ def _isolate_ucode_state(tmp_path, monkeypatch):
     it can never touch the developer's real ~/.ucode/state.json or invoke the
     privileged writer for an OS-managed agent config.
     """
+    import ucode.codex_config as codex_config_mod
     import ucode.config_io as config_io_mod
     import ucode.databricks as databricks_mod
     import ucode.managed_config as managed_config_mod
@@ -50,6 +51,13 @@ def _isolate_ucode_state(tmp_path, monkeypatch):
         managed_files_mod, "MANAGED_BACKUP_MANIFEST_PATH", backup_dir / "manifest.json"
     )
     monkeypatch.setattr(codex_mod, "codex_managed_config_path", lambda: None)
+    monkeypatch.setattr(codex_mod, "CODEX_HOOKS_PATH", tmp_path / ".codex" / "hooks.json")
+    monkeypatch.setattr(
+        codex_config_mod,
+        "DEFAULT_CODEX_CONFIG_PATH",
+        tmp_path / ".codex" / "ucode.config.toml",
+    )
+    monkeypatch.setattr(codex_config_mod, "codex_managed_config_path", lambda: None)
 
     def reject_privileged_write(path, _desired_text):
         pytest.fail(

tests/integration/test_ug_codex_commands.py

@@ -1,11 +1,17 @@
 """CUJs for inspecting codex commands through ug."""
 
+import json
+
 import pytest
 
 pytestmark = [pytest.mark.live, pytest.mark.codex]
 
 
-@pytest.mark.parametrize("routing", ["0", "1"], ids=["routing-off", "routing-on"])
+@pytest.mark.parametrize(
+    "routing",
+    ["off", "full", "subagent"],
+    ids=["routing-off", "routing-on", "subagent-only"],
+)
 def test_ug_codex_app_help(live_session, workspace, routing):
     """Scenario: configure codex, then ask ug for app subcommand help.
 
@@ -23,10 +29,19 @@ def test_ug_codex_app_help(live_session, workspace, routing):
         "--skip-upgrade",
         "--disable-databricks-ai-tools",
     )
-    session.env["ENABLE_SMART_ROUTING_V2"] = routing
+    session.env["ENABLE_SMART_ROUTING_V2"] = "1" if routing == "full" else "0"
+    session.env["ENABLE_SMART_ROUTING_SUBAGENT_ONLY"] = "1" if routing == "subagent" else "0"
     expected = session.run("app", "--help", binary="codex").stdout.strip()
     actual = session.run("codex", "--", "app", "--help").stdout
     assert expected and expected in actual, actual
+    if routing in {"full", "subagent"}:
+        hooks_path = session.home / ".codex" / "hooks.json"
+        hooks = json.loads(hooks_path.read_text(encoding="utf-8"))["hooks"]
+        assert any(
+            "codex-router-hook route-subagent" in hook.get("command", "")
+            for group in hooks["PreToolUse"]
+            for hook in group["hooks"]
+        )
     session.assert_not_routed()