Organize Codex intelligent routing modules

Author: rohita5l

src/ucode/agents/codex.py

@@ -4,8 +4,6 @@
 
 import os
 import re
-import shlex
-import subprocess
 from pathlib import Path
 
 from ucode.agent_updates import available_npm_package_update
@@ -22,6 +20,10 @@
     build_tool_base_url,
     get_databricks_token,
 )
+from ucode.intelligent_routing.codex_hooks import (
+    remove_intelligent_routing_hooks,
+    sync_intelligent_routing_hooks,
+)
 from ucode.launcher import exec_or_spawn
 from ucode.state import mark_tool_managed, save_state
 from ucode.telemetry import agent_version, ucode_version
@@ -38,7 +40,6 @@
 MINIMUM_ROUTING_CODEX_VERSION = (0, 145, 0)
 MINIMUM_ROUTING_CODEX_VERSION_TEXT = "0.145.0"
 INTELLIGENT_ROUTING_STATE_KEY = "codex_intelligent_routing_enabled"
-ROUTING_HOOK_COMMAND_MARKER = "codex-router-hook"
 
 
 SPEC: ToolSpec = {
@@ -369,7 +370,7 @@ def write_tool_config(state: dict, model: str | None = None, provider: str | Non
         # deep_merge can't drop keys, so clear a `model` pinned by an earlier
         # non-provider run that the provider overlay omits.
         doc.pop("model", None)
-    _sync_intelligent_routing_hooks(
+    sync_intelligent_routing_hooks(
         doc,
         state,
         enabled=intelligent_routing_enabled(state) and provider is None,
@@ -443,153 +444,15 @@ def disable_intelligent_routing(state: dict) -> bool:
         if not path.exists():
             continue
         doc = read_toml_safe(path)
-        if _remove_intelligent_routing_hooks(doc):
+        if remove_intelligent_routing_hooks(doc):
             write_toml_file(path, doc)
             changed = True
-    from ucode.codex_routing import clear_routing_artifacts
+    from ucode.intelligent_routing.codex_routing import clear_routing_artifacts
 
     clear_routing_artifacts()
     return changed
 
 
-def route_launch_model(state: dict, tool_args: list[str]):
-    """Route a root Codex launch before the Codex process starts."""
-    from ucode.codex_routing import request_routing_decision
-
-    workspace = state.get("workspace")
-    models = state.get("codex_models")
-    if not isinstance(workspace, str) or not isinstance(models, list):
-        return None, "workspace model metadata is unavailable"
-    try:
-        token = get_databricks_token(workspace, state.get("profile"))
-    except RuntimeError as exc:
-        return None, f"could not authenticate the routing request: {exc}"
-    task = _launch_routing_task(tool_args)
-    return request_routing_decision(workspace, token, task, models)
-
-
-def _launch_routing_task(tool_args: list[str]) -> str:
-    if "exec" in tool_args:
-        prompt_parts = tool_args[tool_args.index("exec") + 1 :]
-        if prompt_parts:
-            return " ".join(prompt_parts)
-    if tool_args:
-        return "Start a Codex session with options: " + " ".join(tool_args)
-    return f"Start an interactive Codex coding session in {Path.cwd().name}."
-
-
-def _sync_intelligent_routing_hooks(doc: dict, state: dict, *, enabled: bool) -> None:
-    _remove_intelligent_routing_hooks(doc)
-    if not enabled:
-        return
-    hooks = doc.setdefault("hooks", {})
-    for event, groups in _routing_hook_groups(state).items():
-        existing = hooks.get(event)
-        if not isinstance(existing, list):
-            existing = []
-        hooks[event] = [*existing, *groups]
-
-
-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")],
-            }
-        ],
-        "SessionStart": [
-            {
-                "matcher": "startup|resume|clear",
-                "hooks": [_routing_command_hook(session_argv)],
-            }
-        ],
-        "SubagentStart": [
-            {
-                "hooks": [_routing_command_hook(subagent_argv)],
-            }
-        ],
-    }
-
-
-def _routing_hook_argv(state: dict, event: str) -> list[str]:
-    workspace = str(state.get("workspace") or "")
-    argv = [
-        build_auth_token_argv(workspace, state.get("profile"), use_pat=bool(state.get("use_pat")))[
-            0
-        ],
-        ROUTING_HOOK_COMMAND_MARKER,
-        event,
-    ]
-    if event != "route-subagent":
-        return argv
-    argv += ["--host", workspace]
-    profile = state.get("profile")
-    if isinstance(profile, str) and profile:
-        argv += ["--profile", profile]
-    if state.get("use_pat"):
-        argv.append("--use-pat")
-    for model in state.get("codex_models") or []:
-        if isinstance(model, str) and model:
-            argv += ["--model", model]
-    return argv
-
-
-def _routing_command_hook(argv: list[str], *, status: str | None = None) -> dict:
-    hook = {
-        "type": "command",
-        "command": shlex.join(argv),
-        "command_windows": subprocess.list2cmdline(argv),
-        "timeout": 35,
-    }
-    if status:
-        hook["statusMessage"] = status
-    return hook
-
-
-def _remove_intelligent_routing_hooks(doc: dict) -> bool:
-    hooks = doc.get("hooks")
-    if not isinstance(hooks, dict):
-        return False
-    changed = False
-    for event in list(hooks):
-        groups = hooks.get(event)
-        if not isinstance(groups, list):
-            continue
-        kept_groups = []
-        for group in groups:
-            if not isinstance(group, dict):
-                kept_groups.append(group)
-                continue
-            handlers = group.get("hooks")
-            if not isinstance(handlers, list):
-                kept_groups.append(group)
-                continue
-            kept_handlers = [
-                handler
-                for handler in handlers
-                if not (
-                    isinstance(handler, dict)
-                    and ROUTING_HOOK_COMMAND_MARKER in str(handler.get("command") or "")
-                )
-            ]
-            if len(kept_handlers) != len(handlers):
-                changed = True
-            if kept_handlers:
-                group["hooks"] = kept_handlers
-                kept_groups.append(group)
-        if kept_groups:
-            hooks[event] = kept_groups
-        else:
-            hooks.pop(event, None)
-    if not hooks:
-        doc.pop("hooks", None)
-    return changed
-
-
 def validate_cmd(binary: str) -> list[str]:
     return [
         binary,

src/ucode/cli.py

@@ -34,7 +34,6 @@
     enable_intelligent_routing,
     intelligent_routing_enabled,
     revert_legacy_shared_config,
-    route_launch_model,
 )
 from ucode.agents.pi import PI_SETTINGS_BACKUP_PATH, PI_SETTINGS_PATH
 from ucode.config_io import restore_file, set_dry_run
@@ -59,6 +58,7 @@
     resolve_pat_token,
     run_databricks_login,
 )
+from ucode.intelligent_routing.codex_routing import route_launch_model
 from ucode.mcp import (
     MCP_CLIENTS,
     SKILLS_MCP_KIND,
@@ -971,7 +971,7 @@ def codex_router_hook_cmd(
     import json
     import sys
 
-    from ucode.codex_routing import (
+    from ucode.intelligent_routing.codex_routing import (
         record_session_start,
         record_subagent_start,
         route_pre_tool_use,

src/ucode/intelligent_routing/__init__.py

@@ -0,0 +1 @@
+"""Intelligent-routing integrations for supported coding agents."""

src/ucode/intelligent_routing/codex_hooks.py

@@ -0,0 +1,124 @@
+"""Codex hook configuration for intelligent subagent routing."""
+
+from __future__ import annotations
+
+import shlex
+import subprocess
+
+from ucode.databricks import build_auth_token_argv
+
+ROUTING_HOOK_COMMAND_MARKER = "codex-router-hook"
+
+
+def sync_intelligent_routing_hooks(doc: dict, state: dict, *, enabled: bool) -> None:
+    """Synchronize ucode-managed routing hooks in a Codex config document."""
+    remove_intelligent_routing_hooks(doc)
+    if not enabled:
+        return
+    hooks = doc.setdefault("hooks", {})
+    for event, groups in _routing_hook_groups(state).items():
+        existing = hooks.get(event)
+        if not isinstance(existing, list):
+            existing = []
+        hooks[event] = [*existing, *groups]
+
+
+def remove_intelligent_routing_hooks(doc: dict) -> bool:
+    """Remove only ucode-managed intelligent-routing hooks."""
+    hooks = doc.get("hooks")
+    if not isinstance(hooks, dict):
+        return False
+    changed = False
+    for event in list(hooks):
+        groups = hooks.get(event)
+        if not isinstance(groups, list):
+            continue
+        kept_groups = []
+        for group in groups:
+            if not isinstance(group, dict):
+                kept_groups.append(group)
+                continue
+            handlers = group.get("hooks")
+            if not isinstance(handlers, list):
+                kept_groups.append(group)
+                continue
+            kept_handlers = [
+                handler
+                for handler in handlers
+                if not (
+                    isinstance(handler, dict)
+                    and ROUTING_HOOK_COMMAND_MARKER in str(handler.get("command") or "")
+                )
+            ]
+            if len(kept_handlers) != len(handlers):
+                changed = True
+            if kept_handlers:
+                group["hooks"] = kept_handlers
+                kept_groups.append(group)
+        if kept_groups:
+            hooks[event] = kept_groups
+        else:
+            hooks.pop(event, None)
+    if not hooks:
+        doc.pop("hooks", None)
+    return changed
+
+
+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")],
+            }
+        ],
+        "SessionStart": [
+            {
+                "matcher": "startup|resume|clear",
+                "hooks": [_routing_command_hook(session_argv)],
+            }
+        ],
+        "SubagentStart": [
+            {
+                "hooks": [_routing_command_hook(subagent_argv)],
+            }
+        ],
+    }
+
+
+def _routing_hook_argv(state: dict, event: str) -> list[str]:
+    workspace = str(state.get("workspace") or "")
+    argv = [
+        build_auth_token_argv(workspace, state.get("profile"), use_pat=bool(state.get("use_pat")))[
+            0
+        ],
+        ROUTING_HOOK_COMMAND_MARKER,
+        event,
+    ]
+    if event != "route-subagent":
+        return argv
+    argv += ["--host", workspace]
+    profile = state.get("profile")
+    if isinstance(profile, str) and profile:
+        argv += ["--profile", profile]
+    if state.get("use_pat"):
+        argv.append("--use-pat")
+    for model in state.get("codex_models") or []:
+        if isinstance(model, str) and model:
+            argv += ["--model", model]
+    return argv
+
+
+def _routing_command_hook(argv: list[str], *, status: str | None = None) -> dict:
+    hook = {
+        "type": "command",
+        "command": shlex.join(argv),
+        "command_windows": subprocess.list2cmdline(argv),
+        "timeout": 35,
+    }
+    if status:
+        hook["statusMessage"] = status
+    return hook

src/ucode/codex_routing.py …ode/intelligent_routing/codex_routing.pysrc/ucode/codex_routing.py renamed to src/ucode/intelligent_routing/codex_routing.py src/ucode/codex_routing.py renamed to src/ucode/intelligent_routing/codex_routing.py

@@ -13,6 +13,7 @@
 from typing import Any
 
 from ucode.config_io import APP_DIR
+from ucode.databricks import get_databricks_token
 
 ROUTER_NAME = "task_v1"
 ROUTING_PATH = "/ai-gateway/routing/v1/routes:select"
@@ -37,6 +38,30 @@ class RoutingDecision:
     rationale: str = ""
 
 
+def route_launch_model(state: dict, tool_args: list[str]):
+    """Route a root Codex launch before the Codex process starts."""
+    workspace = state.get("workspace")
+    models = state.get("codex_models")
+    if not isinstance(workspace, str) or not isinstance(models, list):
+        return None, "workspace model metadata is unavailable"
+    try:
+        token = get_databricks_token(workspace, state.get("profile"))
+    except RuntimeError as exc:
+        return None, f"could not authenticate the routing request: {exc}"
+    task = _launch_routing_task(tool_args)
+    return request_routing_decision(workspace, token, task, models)
+
+
+def _launch_routing_task(tool_args: list[str]) -> str:
+    if "exec" in tool_args:
+        prompt_parts = tool_args[tool_args.index("exec") + 1 :]
+        if prompt_parts:
+            return " ".join(prompt_parts)
+    if tool_args:
+        return "Start a Codex session with options: " + " ".join(tool_args)
+    return f"Start an interactive Codex coding session in {Path.cwd().name}."
+
+
 def request_routing_decision(
     workspace: str,
     token: str,