Add opt-in Codex intelligent routing

README.md

@@ -40,6 +40,21 @@ ucode codex --full-auto
 
 All agents route through Databricks AI Gateway using your workspace credentials — no API keys required.
 
+Codex intelligent routing is opt-in. Enabling it asks the AI Gateway router to select the
+root-session model before launch and installs profile-scoped hooks that route future
+`spawn_agent` calls. Codex may require one-time review of the installed hooks through `/hooks`.
+
+```bash
+ucode codex --enable-intelligent-routing
+```
+
+The setting persists for the current workspace. Disable it and remove only ucode's routing
+hooks with:
+
+```bash
+ucode codex --disable-intelligent-routing
+```
+
 To configure all tools at once:
 
 ```bash
@@ -156,6 +171,8 @@ you to run `ucode <agent>` (existing agent sessions need a restart before the MC
 | `ucode configure --workspaces https://first.databricks.com,https://second.databricks.com` | Configure workspaces without the interactive picker |
 | `ucode configure --profiles DEFAULT` | Configure using existing Databricks CLI profiles (hosts come from `~/.databrickscfg`) |
 | `ucode configure --profiles DEFAULT --use-pat` | Authenticate with the profile's personal access token — no browser login |
+| `ucode codex --enable-intelligent-routing` | Enable AI Gateway routing for Codex sessions and subagents |
+| `ucode codex --disable-intelligent-routing` | Disable routing and remove ucode's Codex routing hooks |
 | `ucode configure --skip-validate` | Write configs without sending a test message through each agent |
 | `ucode configure --agents claude --mcp system.ai.slack` | Configure an agent and register its Databricks MCP server(s) in one command |
 | `ucode configure skills` | Register the skills MCP connection (utility tools only); no skills download |

src/ucode/agents/codex.py

@@ -21,6 +21,10 @@
     get_databricks_token,
 )
 from ucode.launcher import exec_or_spawn
+from ucode.smart_routing.codex_hooks import (
+    remove_smart_routing_hooks,
+    sync_smart_routing_hooks,
+)
 from ucode.state import mark_tool_managed, save_state
 from ucode.telemetry import agent_version, ucode_version
 
@@ -33,6 +37,11 @@
 CODEX_MODEL_PROVIDER_NAME = "ucode-databricks"
 MINIMUM_CODEX_VERSION = (0, 134, 0)
 MINIMUM_CODEX_VERSION_TEXT = "0.134.0"
+MINIMUM_ROUTING_CODEX_VERSION = (0, 145, 0)
+MINIMUM_ROUTING_CODEX_VERSION_TEXT = "0.145.0"
+# Shared across agents: one opt-in enables smart routing for every routing-capable
+# tool (codex, claude), so a workspace turns it on once.
+SMART_ROUTING_STATE_KEY = "smart_routing_enabled"
 
 
 SPEC: ToolSpec = {
@@ -318,6 +327,10 @@ def write_tool_config(state: dict, model: str | None = None, provider: str | Non
     databricks_profile = state.get("profile")
 
     if _use_legacy_layout():
+        if smart_routing_enabled(state) and provider is None:
+            raise RuntimeError(
+                f"Codex smart routing requires Codex {MINIMUM_ROUTING_CODEX_VERSION_TEXT} or newer."
+            )
         # Codex < 0.134.0 only reads ~/.codex/config.toml. Write the shared
         # config with [profiles.ucode] + shared [model_providers.ucode-databricks]
         # and skip the per-profile-file cleanup that would normally strip
@@ -358,6 +371,11 @@ 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_smart_routing_hooks(
+        doc,
+        state,
+        enabled=smart_routing_enabled(state) and provider is None,
+    )
     write_toml_file(CODEX_CONFIG_PATH, doc)
     state = mark_tool_managed(state, "codex", MANAGED_KEYS)
     save_state(state)
@@ -399,6 +417,43 @@ def launch(state: dict, tool_args: list[str]) -> None:
     exec_or_spawn([binary, "--profile", CODEX_PROFILE_NAME, *tool_args])
 
 
+def smart_routing_enabled(state: dict) -> bool:
+    """Return whether the current workspace opted into Codex routing."""
+    return state.get(SMART_ROUTING_STATE_KEY) is True
+
+
+def enable_smart_routing(state: dict) -> dict:
+    """Persist the current workspace's Codex smart-routing opt-in."""
+    parsed = _parse_version(agent_version(SPEC["binary"]))
+    if parsed is not None and parsed < MINIMUM_ROUTING_CODEX_VERSION:
+        raise RuntimeError(
+            "Codex smart routing requires Codex "
+            f"{MINIMUM_ROUTING_CODEX_VERSION_TEXT} or newer; found "
+            f"{agent_version(SPEC['binary'])}."
+        )
+    state[SMART_ROUTING_STATE_KEY] = True
+    return state
+
+
+def disable_smart_routing(state: dict) -> bool:
+    """Disable routing and remove only ucode's Codex routing hooks."""
+    state.pop(SMART_ROUTING_STATE_KEY, None)
+    if state.get("workspace"):
+        save_state(state)
+    changed = False
+    for path in (CODEX_CONFIG_PATH, LEGACY_CODEX_CONFIG_PATH):
+        if not path.exists():
+            continue
+        doc = read_toml_safe(path)
+        if remove_smart_routing_hooks(doc):
+            write_toml_file(path, doc)
+            changed = True
+    from ucode.smart_routing.codex_routing import clear_routing_artifacts
+
+    clear_routing_artifacts()
+    return changed
+
+
 def validate_cmd(binary: str) -> list[str]:
     return [
         binary,

src/ucode/cli.py

@@ -3,6 +3,7 @@
 
 from __future__ import annotations
 
+import os
 import shutil
 from typing import Annotated
 
@@ -28,7 +29,12 @@
 from ucode.agents import (
     launch as launch_agent,
 )
-from ucode.agents.codex import revert_legacy_shared_config
+from ucode.agents.codex import (
+    disable_smart_routing,
+    enable_smart_routing,
+    revert_legacy_shared_config,
+    smart_routing_enabled,
+)
 from ucode.agents.pi import PI_SETTINGS_BACKUP_PATH, PI_SETTINGS_PATH
 from ucode.config_io import restore_file, set_dry_run
 from ucode.databricks import (
@@ -61,6 +67,7 @@
     revert_mcp_configs,
 )
 from ucode.skills_download import configure_skills_download_command
+from ucode.smart_routing.codex_routing import route_launch_model
 from ucode.state import (
     STATE_PATH,
     clear_state,
@@ -952,6 +959,77 @@ def auth_token_cmd(
     sys.stdout.write(token + "\n")
 
 
+@app.command("codex-router-hook", hidden=True)
+def codex_router_hook_cmd(
+    event: str,
+    host: Annotated[str | None, typer.Option("--host")] = None,
+    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,
+) -> None:
+    """Run a Codex smart-routing lifecycle hook."""
+    import json
+    import sys
+
+    from ucode.smart_routing.codex_routing import (
+        record_session_start,
+        record_subagent_start,
+        route_pre_tool_use,
+    )
+
+    try:
+        payload = json.loads(sys.stdin.read() or "{}")
+    except ValueError:
+        return
+    if not isinstance(payload, dict):
+        return
+    if event == "session-start":
+        record_session_start(payload)
+        return
+    if event == "record-subagent":
+        record = record_subagent_start(payload)
+        matched = record.get("matches_router_decision")
+        if matched is True:
+            sys.stdout.write(
+                json.dumps(
+                    {
+                        "systemMessage": "Smart Routing verified. "
+                        f"Subagent is using {record.get('model')}."
+                    }
+                )
+            )
+        elif matched is False:
+            sys.stdout.write(
+                json.dumps(
+                    {
+                        "systemMessage": "Smart Routing mismatch: router requested "
+                        f"{record.get('requested_model')}, but Codex started "
+                        f"{record.get('model')}."
+                    }
+                )
+            )
+        return
+    if event != "route-subagent" or not host:
+        return
+    token = os.environ.get("OAUTH_TOKEN") or os.environ.get("DATABRICKS_BEARER")
+    if not token:
+        if use_pat and not ensure_pat_bearer(profile):
+            return
+        try:
+            token = get_databricks_token(host, profile)
+        except RuntimeError:
+            return
+    output = route_pre_tool_use(
+        payload,
+        workspace=host,
+        token=token,
+        available_models=model or [],
+        audit_decision=True,
+    )
+    if output is not None:
+        sys.stdout.write(json.dumps(output))
+
+
 def _auto_configure_tool(tool: str) -> None:
     """First-time setup for a single tool — mirrors configure_workspace_command."""
     existing = load_state()
@@ -995,6 +1073,7 @@ def _launch_tool(
     provider: str | None = None,
     skip_preflight: bool = False,
     workspace: str | None = None,
+    enable_codex_smart_routing: bool = False,
 ) -> None:
     try:
         tool = normalize_tool(tool_name)
@@ -1018,6 +1097,11 @@ def _launch_tool(
         # An explicit --provider overrides the persisted choice; otherwise fall
         # back to whatever `ucode configure` saved for this tool.
         provider = provider or get_provider_service(state, tool)
+        if tool == "codex" and enable_codex_smart_routing and provider:
+            raise RuntimeError(
+                "Codex smart routing cannot be enabled with --provider. "
+                "Launch without a Model Provider Service and try again."
+            )
         # Validate the provider service before launching — it must exist, be a
         # provider type this tool can route to (e.g. claude can't use an OpenAI
         # or Foundry service), and, for Bedrock, expose Claude models to pin.
@@ -1041,6 +1125,8 @@ def _launch_tool(
             skip_model_discovery=bool(provider),
             skip_preflight=skip_preflight,
         )
+        if tool == "codex" and enable_codex_smart_routing:
+            state = enable_smart_routing(state)
         if provider:
             # Routing through a Model Provider Service pins no Databricks model;
             # the agent uses its own canonical model names (header selects the
@@ -1049,6 +1135,16 @@ def _launch_tool(
             resolved_model = None
         else:
             state, resolved_model = resolve_launch_model(tool, state, None)
+            if tool == "codex" and smart_routing_enabled(state):
+                with spinner("Selecting a Codex model with smart routing..."):
+                    decision, routing_error = route_launch_model(state, ctx.args)
+                if decision is not None:
+                    resolved_model = decision.model
+                    print_note(decision.display_message())
+                elif routing_error:
+                    print_warning(
+                        f"Smart routing was unavailable ({routing_error}); using {resolved_model}."
+                    )
         state = configure_tool(
             tool,
             state,
@@ -1062,6 +1158,13 @@ def _launch_tool(
             print_kv("Provider", provider)
         elif resolved_model:
             print_kv("Model", resolved_model)
+        if tool == "codex" and smart_routing_enabled(state) and not provider:
+            print_kv("Smart routing", "enabled")
+            if enable_codex_smart_routing:
+                print_note(
+                    "Codex requires one-time hook review. Open `/hooks` and trust the "
+                    "ucode routing hooks if prompted."
+                )
         if tool in ("gemini", "opencode", "copilot", "pi"):
             print_note(
                 f"{TOOL_SPECS[tool]['display']} token refresh is managed automatically "
@@ -1116,10 +1219,36 @@ def codex_cmd(
     ] = None,
     skip_preflight: SkipPreflightOption = False,
     workspace: WorkspaceOption = None,
+    enable_smart_routing_flag: Annotated[
+        bool,
+        typer.Option(
+            "--enable-smart-routing",
+            help="Enable AI Gateway model routing for Codex sessions and subagents.",
+        ),
+    ] = False,
+    disable_smart_routing_flag: Annotated[
+        bool,
+        typer.Option(
+            "--disable-smart-routing",
+            help="Disable smart routing and remove ucode's Codex routing hooks.",
+        ),
+    ] = False,
 ) -> None:
     """Launch Codex via Databricks."""
+    if enable_smart_routing_flag and disable_smart_routing_flag:
+        print_err("Use only one of --enable-smart-routing or --disable-smart-routing.")
+        raise typer.Exit(1)
+    if disable_smart_routing_flag:
+        disable_smart_routing(load_state())
+        print_success("Codex smart routing disabled; ucode routing hooks removed")
+        return
     _launch_tool(
-        "codex", ctx, provider=provider, skip_preflight=skip_preflight, workspace=workspace
+        "codex",
+        ctx,
+        provider=provider,
+        skip_preflight=skip_preflight,
+        workspace=workspace,
+        enable_codex_smart_routing=enable_smart_routing_flag,
     )
 
 

src/ucode/smart_routing/__init__.py

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