From 1d4ceedf3b6c76cd69daded7c3f077333d5331f4 Mon Sep 17 00:00:00 2001 From: Rohit Agrawal Date: Wed, 29 Jul 2026 20:13:45 -0400 Subject: [PATCH 1/5] Add opt-in Codex intelligent routing --- README.md | 17 ++ src/ucode/agents/codex.py | 191 ++++++++++++++++++++ src/ucode/cli.py | 134 +++++++++++++- src/ucode/codex_routing.py | 347 ++++++++++++++++++++++++++++++++++++ tests/test_agent_codex.py | 107 +++++++++++ tests/test_cli.py | 58 ++++++ tests/test_codex_routing.py | 278 +++++++++++++++++++++++++++++ 7 files changed, 1130 insertions(+), 2 deletions(-) create mode 100644 src/ucode/codex_routing.py create mode 100644 tests/test_codex_routing.py diff --git a/README.md b/README.md index 909bc1e..0352757 100644 --- a/README.md +++ b/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 @@ -162,6 +177,8 @@ you to run `ucode ` (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 | diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 2a5c657..94df24e 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -4,6 +4,8 @@ import os import re +import shlex +import subprocess from pathlib import Path from ucode.agent_updates import available_npm_package_update @@ -33,6 +35,10 @@ 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" +INTELLIGENT_ROUTING_STATE_KEY = "codex_intelligent_routing_enabled" +ROUTING_HOOK_COMMAND_MARKER = "codex-router-hook" SPEC: ToolSpec = { @@ -318,6 +324,11 @@ def write_tool_config(state: dict, model: str | None = None, provider: str | Non databricks_profile = state.get("profile") if _use_legacy_layout(): + if intelligent_routing_enabled(state) and provider is None: + raise RuntimeError( + "Codex intelligent routing requires Codex " + f"{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 +369,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_intelligent_routing_hooks( + doc, + state, + enabled=intelligent_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 +415,181 @@ def launch(state: dict, tool_args: list[str]) -> None: exec_or_spawn([binary, "--profile", CODEX_PROFILE_NAME, *tool_args]) +def intelligent_routing_enabled(state: dict) -> bool: + """Return whether the current workspace opted into Codex routing.""" + return state.get(INTELLIGENT_ROUTING_STATE_KEY) is True + + +def enable_intelligent_routing(state: dict) -> dict: + """Persist the current workspace's Codex intelligent-routing opt-in.""" + parsed = _parse_version(agent_version(SPEC["binary"])) + if parsed is not None and parsed < MINIMUM_ROUTING_CODEX_VERSION: + raise RuntimeError( + "Codex intelligent routing requires Codex " + f"{MINIMUM_ROUTING_CODEX_VERSION_TEXT} or newer; found " + f"{agent_version(SPEC['binary'])}." + ) + state[INTELLIGENT_ROUTING_STATE_KEY] = True + return state + + +def disable_intelligent_routing(state: dict) -> bool: + """Disable routing and remove only ucode's Codex routing hooks.""" + state.pop(INTELLIGENT_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_intelligent_routing_hooks(doc): + write_toml_file(path, doc) + changed = True + from ucode.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, diff --git a/src/ucode/cli.py b/src/ucode/cli.py index b7b7a6e..409a078 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -3,6 +3,7 @@ from __future__ import annotations +import os import shutil from typing import Annotated @@ -28,7 +29,13 @@ from ucode.agents import ( launch as launch_agent, ) -from ucode.agents.codex import revert_legacy_shared_config +from ucode.agents.codex import ( + disable_intelligent_routing, + 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 from ucode.databricks import ( @@ -953,6 +960,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 intelligent-routing lifecycle hook.""" + import json + import sys + + from ucode.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": "Intelligent Routing verified. " + f"Subagent is using {record.get('model')}." + } + ) + ) + elif matched is False: + sys.stdout.write( + json.dumps( + { + "systemMessage": "Intelligent 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() @@ -996,6 +1074,7 @@ def _launch_tool( provider: str | None = None, skip_preflight: bool = False, workspace: str | None = None, + enable_codex_intelligent_routing: bool = False, ) -> None: try: tool = normalize_tool(tool_name) @@ -1019,6 +1098,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_intelligent_routing and provider: + raise RuntimeError( + "Codex intelligent 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. @@ -1042,6 +1126,8 @@ def _launch_tool( skip_model_discovery=bool(provider), skip_preflight=skip_preflight, ) + if tool == "codex" and enable_codex_intelligent_routing: + state = enable_intelligent_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 @@ -1050,6 +1136,17 @@ def _launch_tool( resolved_model = None else: state, resolved_model = resolve_launch_model(tool, state, None) + if tool == "codex" and intelligent_routing_enabled(state): + with spinner("Selecting a Codex model with intelligent routing..."): + decision, routing_error = route_launch_model(state, ctx.args) + if decision is not None: + resolved_model = decision.model + print_note(f"Using Intelligent Routing. Routing to {resolved_model}.") + elif routing_error: + print_warning( + f"Intelligent routing was unavailable ({routing_error}); " + f"using {resolved_model}." + ) state = configure_tool( tool, state, @@ -1063,6 +1160,13 @@ def _launch_tool( print_kv("Provider", provider) elif resolved_model: print_kv("Model", resolved_model) + if tool == "codex" and intelligent_routing_enabled(state) and not provider: + print_kv("Intelligent routing", "enabled") + if enable_codex_intelligent_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 " @@ -1117,10 +1221,36 @@ def codex_cmd( ] = None, skip_preflight: SkipPreflightOption = False, workspace: WorkspaceOption = None, + enable_intelligent_routing_flag: Annotated[ + bool, + typer.Option( + "--enable-intelligent-routing", + help="Enable AI Gateway model routing for Codex sessions and subagents.", + ), + ] = False, + disable_intelligent_routing_flag: Annotated[ + bool, + typer.Option( + "--disable-intelligent-routing", + help="Disable intelligent routing and remove ucode's Codex routing hooks.", + ), + ] = False, ) -> None: """Launch Codex via Databricks.""" + if enable_intelligent_routing_flag and disable_intelligent_routing_flag: + print_err("Use only one of --enable-intelligent-routing or --disable-intelligent-routing.") + raise typer.Exit(1) + if disable_intelligent_routing_flag: + disable_intelligent_routing(load_state()) + print_success("Codex intelligent 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_intelligent_routing=enable_intelligent_routing_flag, ) diff --git a/src/ucode/codex_routing.py b/src/ucode/codex_routing.py new file mode 100644 index 0000000..6bff5e9 --- /dev/null +++ b/src/ucode/codex_routing.py @@ -0,0 +1,347 @@ +"""Databricks AI Gateway routing helpers for Codex sessions and subagents.""" + +from __future__ import annotations + +import json +import re +import time +import urllib.error +import urllib.request +import uuid +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from ucode.config_io import APP_DIR + +ROUTER_NAME = "task_v1" +ROUTING_PATH = "/ai-gateway/routing/v1/routes:select" +REQUEST_TIMEOUT_S = 30.0 +CODEX_ROUTE_ARMS = ("glm-5-2", "gpt-5-6-sol", "gpt-5-6-luna") +GLM_ROUTE_ARM = "glm-5-2" +GLM_GATEWAY_MODEL = "system.ai.glm-5-2" +SPAWN_AGENT_TOOL_SUFFIX = "spawn_agent" +CANARY_PATH = APP_DIR / "codex-intelligent-routing-canary.json" +AUDIT_PATH = APP_DIR / "codex-intelligent-routing-audit.jsonl" +DECISIONS_PATH = APP_DIR / "codex-intelligent-routing-decisions.jsonl" + +_GPT_RE = re.compile(r"gpt-(\d+)(?:[.-](\d+))?(?:[.-](\d+))?(-.+|[a-z].*)?") + + +@dataclass(frozen=True) +class RoutingDecision: + """One model selection returned by the AI Gateway router.""" + + model: str + raw_model: str + rationale: str = "" + + +def request_routing_decision( + workspace: str, + token: str, + task: str, + available_models: list[str], + *, + timeout: float = REQUEST_TIMEOUT_S, +) -> tuple[RoutingDecision | None, str | None]: + """Ask the workspace ``task_v1`` router for a servable Codex model.""" + candidates = _routing_candidates(available_models) + missing = [ + arm for arm in CODEX_ROUTE_ARMS if arm not in {_normalize_model(m) for m in candidates} + ] + if missing: + return None, f"required Codex routing models are unavailable: {', '.join(missing)}" + + body = { + "route_options": [{"model": model, "harness": "codex"} for model in CODEX_ROUTE_ARMS], + "task": {"prompt": task[:4000]}, + "route_selector": {"router_name": ROUTER_NAME}, + } + request = urllib.request.Request( + workspace.rstrip("/") + ROUTING_PATH, + data=json.dumps(body).encode("utf-8"), + headers={ + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + }, + method="POST", + ) + try: + with urllib.request.urlopen(request, timeout=timeout) as response: + payload = json.loads(response.read().decode("utf-8")) + except urllib.error.HTTPError as exc: + detail = "" + try: + detail = exc.read().decode("utf-8", errors="replace").strip() + except OSError: + pass + reason = f"router returned HTTP {exc.code}" + if detail: + reason = f"{reason}: {detail[:300]}" + return None, reason + except (urllib.error.URLError, OSError, TimeoutError, ValueError) as exc: + return None, f"router request failed: {exc}" + + raw_model = _selected_model(payload) + if raw_model is None: + return None, "router returned no model selection" + model = resolve_routed_model(raw_model, candidates) + if model is None: + return None, f"router selected unsupported model {raw_model!r}" + rationale = payload.get("rationale") if isinstance(payload, dict) else None + return ( + RoutingDecision( + model=model, + raw_model=raw_model, + rationale=rationale if isinstance(rationale, str) else "", + ), + None, + ) + + +def resolve_routed_model(raw_model: str, available_models: list[str]) -> str | None: + """Map a ``task_v1`` arm to a model the configured workspace can serve.""" + candidates = _routing_candidates(available_models) + normalized = {_normalize_model(model): model for model in candidates} + return normalized.get(_normalize_model(raw_model)) + + +def route_pre_tool_use( + payload: dict[str, Any], + *, + workspace: str, + token: str, + available_models: list[str], + timeout: float = REQUEST_TIMEOUT_S, + audit_decision: bool = False, +) -> dict[str, Any] | None: + """Route one Codex ``spawn_agent`` call and rewrite its model.""" + if not is_spawn_agent_tool(payload.get("tool_name")): + return None + tool_input = payload.get("tool_input") + if not isinstance(tool_input, dict): + return None + task_name = tool_input.get("task_name") or tool_input.get("agent_name") + task = task_name if isinstance(task_name, str) and task_name else "Codex subagent task" + decision, _ = request_routing_decision( + workspace, + token, + task, + available_models, + timeout=timeout, + ) + if decision is None: + return None + if decision.raw_model == GLM_ROUTE_ARM: + return { + "systemMessage": ( + "Intelligent Routing selected GLM 5.2, which is not enabled for Codex " + "subagents. Keeping the original subagent model." + ) + } + routed_model = _codex_model_id(decision.model) + if audit_decision: + _record_routing_decision(payload, task, decision, routed_model) + routing_message = f"Using Intelligent Routing. Routing to {routed_model}." + output: dict[str, Any] = { + "hookEventName": "PreToolUse", + "permissionDecision": "allow", + "updatedInput": {**tool_input, "model": routed_model}, + "permissionDecisionReason": routing_message, + } + if decision.rationale: + output["permissionDecisionReason"] += f" {decision.rationale}" + return {"systemMessage": routing_message, "hookSpecificOutput": output} + + +def is_spawn_agent_tool(tool_name: Any) -> bool: + """Return whether a hook payload names Codex's subagent spawn tool.""" + if not isinstance(tool_name, str): + return False + normalized = tool_name.strip().lower() + return normalized == "agent" or normalized.endswith(SPAWN_AGENT_TOOL_SUFFIX) + + +def record_session_start(payload: dict[str, Any]) -> None: + """Write a canary proving Codex trusted and ran the routing hooks.""" + _write_json( + CANARY_PATH, + { + "session_id": payload.get("session_id"), + "model": payload.get("model"), + "at": time.time(), + }, + ) + + +def record_subagent_start(payload: dict[str, Any]) -> dict[str, Any]: + """Append the model Codex actually selected for a routed subagent.""" + actual_model = payload.get("model") + decision = _pending_decision(payload.get("session_id"), actual_model) + record = { + "agent_id": payload.get("agent_id"), + "agent_type": payload.get("agent_type"), + "model": actual_model, + "session_id": payload.get("session_id"), + "at": time.time(), + } + if decision is not None: + record.update( + { + "decision_id": decision.get("decision_id"), + "router_model": decision.get("router_model"), + "requested_model": decision.get("requested_model"), + "matches_router_decision": decision.get("requested_model") == actual_model, + } + ) + _append_jsonl(AUDIT_PATH, record) + return record + + +def clear_routing_artifacts() -> None: + """Remove ucode-owned routing canary and audit files.""" + for path in (CANARY_PATH, AUDIT_PATH, DECISIONS_PATH): + try: + path.unlink(missing_ok=True) + except OSError: + continue + + +def _selected_model(payload: Any) -> str | None: + if not isinstance(payload, dict): + return None + selections = payload.get("route_selection") + if not isinstance(selections, list) or not selections: + return None + selection = selections[0] + if not isinstance(selection, dict): + return None + option = selection.get("route_option") + if not isinstance(option, dict): + return None + model = option.get("model") + return model if isinstance(model, str) and model else None + + +def _routing_candidates(models: list[str]) -> list[str]: + candidates = [model for model in models if isinstance(model, str) and model] + if GLM_ROUTE_ARM not in {_normalize_model(model) for model in candidates}: + candidates.append(GLM_GATEWAY_MODEL) + return candidates + + +def _parse_gpt(model: str) -> tuple[int, int, int, str] | None: + match = _GPT_RE.fullmatch(_normalize_model(model)) + if not match: + return None + major, minor, patch, suffix = match.groups() + return int(major), int(minor or 0), int(patch or 0), suffix or "" + + +def _model_strength(model: str) -> tuple[int, int, int, int]: + parsed = _parse_gpt(model) + if parsed is None: + return (0, 0, 0, 0) + major, minor, patch, suffix = parsed + return major, minor, patch, 1 if not suffix else 0 + + +def _normalize_model(model: str) -> str: + tail = model.rsplit("/", 1)[-1] + for prefix in ("databricks-", "system.ai."): + if tail.startswith(prefix): + tail = tail[len(prefix) :] + break + return tail.lower() + + +def _codex_model_id(model: str) -> str: + tail = model.rsplit("/", 1)[-1] + if tail in {"databricks-gpt-5-2-codex", "databricks-gpt-5-4-nano"}: + return tail + if _normalize_model(model) == GLM_ROUTE_ARM: + return GLM_GATEWAY_MODEL + if model.startswith("system.ai."): + bare = model.removeprefix("system.ai.") + elif tail.startswith("databricks-"): + bare = tail.removeprefix("databricks-") + else: + return model + match = _GPT_RE.fullmatch(bare) + if not match: + return bare + major, minor, patch, suffix = match.groups() + version = major + if minor is not None: + version += f".{minor}" + if patch is not None: + version += f".{patch}" + return f"gpt-{version}{suffix or ''}" + + +def _record_routing_decision( + payload: dict[str, Any], + task_name: str, + decision: RoutingDecision, + requested_model: str, +) -> None: + _append_jsonl( + DECISIONS_PATH, + { + "decision_id": uuid.uuid4().hex, + "session_id": payload.get("session_id"), + "task_name": task_name, + "router_model": decision.raw_model, + "requested_model": requested_model, + "at": time.time(), + }, + ) + + +def _pending_decision(session_id: Any, actual_model: Any) -> dict[str, Any] | None: + used = { + record.get("decision_id") for record in _read_jsonl(AUDIT_PATH) if record.get("decision_id") + } + pending = [ + decision + for decision in _read_jsonl(DECISIONS_PATH) + if decision.get("session_id") == session_id and decision.get("decision_id") not in used + ] + return next( + (decision for decision in pending if decision.get("requested_model") == actual_model), + pending[0] if pending else None, + ) + + +def _read_jsonl(path: Path) -> list[dict[str, Any]]: + try: + lines = path.read_text(encoding="utf-8").splitlines() + except OSError: + return [] + records = [] + for line in lines: + try: + record = json.loads(line) + except ValueError: + continue + if isinstance(record, dict): + records.append(record) + return records + + +def _append_jsonl(path: Path, payload: dict[str, Any]) -> None: + try: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("a", encoding="utf-8") as handle: + handle.write(json.dumps(payload) + "\n") + except OSError: + return + + +def _write_json(path: Path, payload: dict[str, Any]) -> None: + try: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(payload), encoding="utf-8") + except OSError: + return diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index 9a68e03..299f1e7 100644 --- a/tests/test_agent_codex.py +++ b/tests/test_agent_codex.py @@ -217,6 +217,65 @@ def test_writes_legacy_shared_config_when_codex_too_old(self, tmp_path, monkeypa assert provider["base_url"] == f"{WS}/ai-gateway/codex/v1" assert provider["wire_api"] == "responses" + def test_intelligent_routing_writes_profile_scoped_hooks(self, tmp_path, monkeypatch): + config_path = tmp_path / ".codex" / "ucode.config.toml" + config_path.parent.mkdir() + config_path.write_text( + "[[hooks.PreToolUse]]\n" + 'matcher = "Bash"\n' + "[[hooks.PreToolUse.hooks]]\n" + 'type = "command"\n' + 'command = "user-policy"\n', + encoding="utf-8", + ) + monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", config_path) + monkeypatch.setattr(codex, "CODEX_BACKUP_PATH", tmp_path / "backup.toml") + monkeypatch.setattr(codex, "agent_version", lambda binary: "0.145.0") + monkeypatch.setattr(codex, "save_state", lambda state: None) + + codex.write_tool_config( + { + "workspace": WS, + "profile": "prod", + "codex_models": ["databricks-gpt-5", "databricks-gpt-5-5"], + codex.INTELLIGENT_ROUTING_STATE_KEY: True, + } + ) + + doc = read_toml_safe(config_path) + assert set(doc["hooks"]) == {"PreToolUse", "SessionStart", "SubagentStart"} + pre_tool_commands = [ + hook["command"] for group in doc["hooks"]["PreToolUse"] for hook in group["hooks"] + ] + assert "user-policy" in pre_tool_commands + route_command = next( + command for command in pre_tool_commands if "codex-router-hook" in command + ) + assert "route-subagent" in route_command + assert "--host https://example.databricks.com" in route_command + assert "--profile prod" in route_command + assert "--model databricks-gpt-5-5" in route_command + + def test_provider_launch_removes_routing_hooks(self, tmp_path, monkeypatch): + config_path = tmp_path / ".codex" / "ucode.config.toml" + backup_path = tmp_path / "backup.toml" + monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", config_path) + monkeypatch.setattr(codex, "CODEX_BACKUP_PATH", backup_path) + monkeypatch.setattr(codex, "agent_version", lambda binary: "0.145.0") + monkeypatch.setattr(codex, "save_state", lambda state: None) + state = { + "workspace": WS, + "codex_models": ["databricks-gpt-5"], + codex.INTELLIGENT_ROUTING_STATE_KEY: True, + } + + codex.write_tool_config(state) + assert "hooks" in read_toml_safe(config_path) + + codex.write_tool_config(state, provider="main.schema.provider") + + assert "hooks" not in read_toml_safe(config_path) + def test_legacy_write_preserves_other_profiles_in_shared_config(self, tmp_path, monkeypatch): config_dir = tmp_path / ".codex" config_dir.mkdir() @@ -259,6 +318,54 @@ def test_unknown_version_uses_modern_layout(self, monkeypatch): assert codex._use_legacy_layout() is False +class TestCodexIntelligentRouting: + def test_enable_requires_supported_codex(self, monkeypatch): + monkeypatch.setattr(codex, "agent_version", lambda binary: "0.144.0") + + try: + codex.enable_intelligent_routing({}) + assert False + except RuntimeError as exc: + assert "0.145.0 or newer" in str(exc) + + def test_disable_removes_only_ucode_hooks(self, tmp_path, monkeypatch): + config_path = tmp_path / ".codex" / "ucode.config.toml" + legacy_path = tmp_path / ".codex" / "config.toml" + config_path.parent.mkdir() + config_path.write_text( + "[[hooks.PreToolUse]]\n" + 'matcher = "Bash"\n' + "[[hooks.PreToolUse.hooks]]\n" + 'type = "command"\n' + 'command = "user-policy"\n\n' + "[[hooks.PreToolUse]]\n" + 'matcher = "Agent"\n' + "[[hooks.PreToolUse.hooks]]\n" + 'type = "command"\n' + 'command = "ucode codex-router-hook route-subagent"\n\n' + "[[hooks.SessionStart]]\n" + "[[hooks.SessionStart.hooks]]\n" + 'type = "command"\n' + 'command = "ucode codex-router-hook session-start"\n', + encoding="utf-8", + ) + monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", config_path) + monkeypatch.setattr(codex, "LEGACY_CODEX_CONFIG_PATH", legacy_path) + monkeypatch.setattr(codex, "save_state", lambda state: None) + monkeypatch.setattr("ucode.codex_routing.clear_routing_artifacts", lambda: None) + state = {"workspace": WS, codex.INTELLIGENT_ROUTING_STATE_KEY: True} + + assert codex.disable_intelligent_routing(state) is True + + doc = read_toml_safe(config_path) + assert state.get(codex.INTELLIGENT_ROUTING_STATE_KEY) is None + assert list(doc["hooks"]) == ["PreToolUse"] + assert doc["hooks"]["PreToolUse"][0]["hooks"][0]["command"] == "user-policy" + + def test_launch_task_uses_exec_prompt(self): + assert codex._launch_routing_task(["exec", "fix the parser"]) == "fix the parser" + + class TestCodexRemoveLegacyProfile: def test_drops_provider_block_on_modern_path(self, tmp_path, monkeypatch): config_dir = tmp_path / ".codex" diff --git a/tests/test_cli.py b/tests/test_cli.py index f5119c2..ec8e1b5 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -190,6 +190,64 @@ def test_no_workspace_flag_leaves_current_workspace(self): assert result.exit_code == 0, result.output mock_set.assert_not_called() + def test_codex_enable_intelligent_routing_is_consumed_by_ucode(self): + with patch("ucode.cli._launch_tool") as mock_launch: + result = runner.invoke(app, ["codex", "--enable-intelligent-routing"]) + + assert result.exit_code == 0, result.output + assert mock_launch.call_args.kwargs["enable_codex_intelligent_routing"] is True + assert mock_launch.call_args.args[1].args == [] + + def test_codex_disable_removes_hooks_without_launching(self): + with ( + patch("ucode.cli.load_state", return_value=MINIMAL_STATE), + patch("ucode.cli.disable_intelligent_routing") as mock_disable, + patch("ucode.cli._launch_tool") as mock_launch, + ): + result = runner.invoke(app, ["codex", "--disable-intelligent-routing"]) + + assert result.exit_code == 0, result.output + mock_disable.assert_called_once_with(MINIMAL_STATE) + mock_launch.assert_not_called() + assert "routing hooks removed" in result.output + + def test_codex_routing_flags_are_mutually_exclusive(self): + result = runner.invoke( + app, + ["codex", "--enable-intelligent-routing", "--disable-intelligent-routing"], + ) + + assert result.exit_code == 1 + assert "Use only one" in result.output + + def test_enabled_codex_launch_uses_routed_root_model(self): + state = { + **MINIMAL_STATE, + "codex_intelligent_routing_enabled": True, + "codex_models": ["databricks-gpt-5", "databricks-gpt-5-5"], + } + decision = MagicMock(model="databricks-gpt-5-5") + with ( + patch("ucode.cli.ensure_bootstrap_dependencies"), + patch("ucode.cli.load_state", return_value=state), + patch("ucode.cli.ensure_provider_state", return_value=state), + patch("ucode.cli.configure_shared_state", return_value=state), + patch( + "ucode.cli.resolve_launch_model", + return_value=(state, "databricks-gpt-5"), + ), + patch("ucode.cli.route_launch_model", return_value=(decision, None)), + patch("ucode.cli.configure_tool", return_value=state) as mock_configure, + patch("ucode.cli.launch_agent"), + ): + result = runner.invoke(app, ["codex"]) + + assert result.exit_code == 0, result.output + assert mock_configure.call_args.args[2] == "databricks-gpt-5-5" + assert "Using Intelligent Routing. Routing to databricks-gpt-5-5." in _strip_ansi( + result.output + ) + class TestMcpSubcommands: def test_web_search_subcommand_help(self): diff --git a/tests/test_codex_routing.py b/tests/test_codex_routing.py new file mode 100644 index 0000000..3eb2b8c --- /dev/null +++ b/tests/test_codex_routing.py @@ -0,0 +1,278 @@ +"""Tests for Codex intelligent-routing hooks.""" + +from __future__ import annotations + +import json +import urllib.error + +from ucode import codex_routing + +WS = "https://example.databricks.com" + + +class _Response: + def __init__(self, payload: dict): + self.payload = payload + + def __enter__(self): + return self + + def __exit__(self, *args): + return None + + def read(self) -> bytes: + return json.dumps(self.payload).encode("utf-8") + + +def test_routes_with_task_v1_codex_menu(monkeypatch): + captured = {} + + def fake_urlopen(request, timeout): + captured["url"] = request.full_url + captured["headers"] = dict(request.headers) + captured["body"] = json.loads(request.data) + captured["timeout"] = timeout + return _Response( + { + "route_selection": [{"route_option": {"model": "gpt-5-6-sol", "harness": "codex"}}], + "rationale": "Needs the strongest coding model.", + } + ) + + monkeypatch.setattr(codex_routing.urllib.request, "urlopen", fake_urlopen) + + decision, error = codex_routing.request_routing_decision( + WS, + "token", + "Refactor the parser", + ["system.ai.gpt-5-6-luna", "system.ai.gpt-5-6-sol"], + ) + + assert error is None + assert decision == codex_routing.RoutingDecision( + model="system.ai.gpt-5-6-sol", + raw_model="gpt-5-6-sol", + rationale="Needs the strongest coding model.", + ) + assert captured["url"] == f"{WS}/ai-gateway/routing/v1/routes:select" + assert captured["headers"]["Authorization"] == "Bearer token" + assert captured["body"] == { + "route_options": [ + {"model": "glm-5-2", "harness": "codex"}, + {"model": "gpt-5-6-sol", "harness": "codex"}, + {"model": "gpt-5-6-luna", "harness": "codex"}, + ], + "task": {"prompt": "Refactor the parser"}, + "route_selector": {"router_name": "task_v1"}, + } + + +def test_router_model_is_not_substituted_when_exact_model_is_unavailable(): + model = codex_routing.resolve_routed_model( + "gpt-5-6-luna", + ["databricks-gpt-5-4-nano", "databricks-gpt-5", "databricks-gpt-5-5"], + ) + + assert model is None + + +def test_glm_maps_to_databricks_gateway_model(): + model = codex_routing.resolve_routed_model( + "glm-5-2", + ["system.ai.gpt-5-6-luna", "system.ai.gpt-5-6-sol"], + ) + + assert model == "system.ai.glm-5-2" + + +def test_exact_router_model_is_preserved(): + model = codex_routing.resolve_routed_model( + "gpt-5-6-sol", + ["databricks-gpt-5-5", "databricks-gpt-5-6-sol"], + ) + + assert model == "databricks-gpt-5-6-sol" + + +def test_router_failure_fails_open(monkeypatch): + monkeypatch.setattr( + codex_routing.urllib.request, + "urlopen", + lambda *args, **kwargs: (_ for _ in ()).throw(urllib.error.URLError("offline")), + ) + + decision, error = codex_routing.request_routing_decision( + WS, + "token", + "task", + ["system.ai.gpt-5-6-luna", "system.ai.gpt-5-6-sol"], + ) + + assert decision is None + assert "offline" in str(error) + + +def test_spawn_rewrite_preserves_original_input(monkeypatch): + encrypted_message = {"encrypted": "opaque-ciphertext"} + payload = { + "tool_name": "collaborationspawn_agent", + "tool_input": { + "task_name": "reviewer", + "message": encrypted_message, + "fork": False, + }, + } + monkeypatch.setattr( + codex_routing, + "request_routing_decision", + lambda *args, **kwargs: ( + codex_routing.RoutingDecision( + model="databricks-gpt-5-5", + raw_model="gpt-5-6-sol", + rationale="Review needs deeper reasoning.", + ), + None, + ), + ) + + output = codex_routing.route_pre_tool_use( + payload, + workspace=WS, + token="token", + available_models=["databricks-gpt-5-5"], + ) + + hook = output["hookSpecificOutput"] + assert output["systemMessage"] == "Using Intelligent Routing. Routing to gpt-5.5." + assert hook["permissionDecision"] == "allow" + assert hook["updatedInput"] == { + "task_name": "reviewer", + "message": encrypted_message, + "fork": False, + "model": "gpt-5.5", + } + assert hook["permissionDecisionReason"] == ( + "Using Intelligent Routing. Routing to gpt-5.5. Review needs deeper reasoning." + ) + + +def test_spawn_rewrite_uses_codex_model_id_for_uc_endpoint(monkeypatch): + monkeypatch.setattr( + codex_routing, + "request_routing_decision", + lambda *args, **kwargs: ( + codex_routing.RoutingDecision( + model="system.ai.gpt-5-6-luna", + raw_model="gpt-5-6-luna", + ), + 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-6-luna"], + ) + + assert output["systemMessage"] == "Using Intelligent Routing. Routing to gpt-5.6-luna." + assert output["hookSpecificOutput"]["updatedInput"]["model"] == "gpt-5.6-luna" + + +def test_spawn_glm_decision_keeps_original_model(monkeypatch): + monkeypatch.setattr( + codex_routing, + "request_routing_decision", + lambda *args, **kwargs: ( + codex_routing.RoutingDecision( + model="system.ai.glm-5-2", + raw_model="glm-5-2", + ), + None, + ), + ) + + output = codex_routing.route_pre_tool_use( + { + "tool_name": "collaborationspawn_agent", + "tool_input": {"task_name": "glm-task", "message": "encrypted"}, + }, + workspace=WS, + token="token", + available_models=["system.ai.gpt-5-6-luna", "system.ai.gpt-5-6-sol"], + ) + + assert output == { + "systemMessage": ( + "Intelligent Routing selected GLM 5.2, which is not enabled for Codex " + "subagents. Keeping the original subagent model." + ) + } + + +def test_non_spawn_tool_has_no_opinion(): + assert ( + codex_routing.route_pre_tool_use( + {"tool_name": "Bash", "tool_input": {"command": "true"}}, + workspace=WS, + token="token", + available_models=["databricks-gpt-5"], + ) + is None + ) + + +def test_canary_and_audit_are_written(tmp_path, monkeypatch): + canary = tmp_path / "canary.json" + audit = tmp_path / "audit.jsonl" + monkeypatch.setattr(codex_routing, "CANARY_PATH", canary) + monkeypatch.setattr(codex_routing, "AUDIT_PATH", audit) + + codex_routing.record_session_start({"session_id": "s1", "model": "gpt-5.5"}) + codex_routing.record_subagent_start( + {"session_id": "s1", "agent_id": "a1", "agent_type": "reviewer", "model": "gpt-5"} + ) + + assert json.loads(canary.read_text())["session_id"] == "s1" + assert json.loads(audit.read_text().strip())["agent_id"] == "a1" + + +def test_decision_is_reconciled_with_actual_subagent_model(tmp_path, monkeypatch): + decisions = tmp_path / "decisions.jsonl" + audit = tmp_path / "audit.jsonl" + monkeypatch.setattr(codex_routing, "DECISIONS_PATH", decisions) + monkeypatch.setattr(codex_routing, "AUDIT_PATH", audit) + monkeypatch.setattr( + codex_routing, + "request_routing_decision", + lambda *args, **kwargs: ( + codex_routing.RoutingDecision( + model="system.ai.gpt-5-6-luna", + raw_model="gpt-5-6-luna", + ), + None, + ), + ) + + codex_routing.route_pre_tool_use( + { + "session_id": "s1", + "tool_name": "collaborationspawn_agent", + "tool_input": {"task_name": "glm-task", "message": "encrypted"}, + }, + workspace=WS, + token="token", + available_models=["system.ai.gpt-5-6-luna", "system.ai.gpt-5-6-sol"], + audit_decision=True, + ) + record = codex_routing.record_subagent_start( + {"session_id": "s1", "agent_id": "a1", "model": "gpt-5.6-luna"} + ) + + assert record["router_model"] == "gpt-5-6-luna" + assert record["requested_model"] == "gpt-5.6-luna" + assert record["matches_router_decision"] is True From 803941ece8ae9aeaf632166d7e6b3b60ed026231 Mon Sep 17 00:00:00 2001 From: Rohit Agrawal Date: Wed, 29 Jul 2026 20:19:38 -0400 Subject: [PATCH 2/5] Organize Codex intelligent routing modules --- src/ucode/agents/codex.py | 151 +----------------- src/ucode/cli.py | 4 +- src/ucode/intelligent_routing/__init__.py | 1 + src/ucode/intelligent_routing/codex_hooks.py | 124 ++++++++++++++ .../codex_routing.py | 25 +++ tests/test_agent_codex.py | 5 +- tests/test_codex_routing.py | 2 +- 7 files changed, 163 insertions(+), 149 deletions(-) create mode 100644 src/ucode/intelligent_routing/__init__.py create mode 100644 src/ucode/intelligent_routing/codex_hooks.py rename src/ucode/{ => intelligent_routing}/codex_routing.py (91%) diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 94df24e..a571a0b 100644 --- a/src/ucode/agents/codex.py +++ b/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, diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 409a078..cf659ce 100644 --- a/src/ucode/cli.py +++ b/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, @@ -972,7 +972,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, diff --git a/src/ucode/intelligent_routing/__init__.py b/src/ucode/intelligent_routing/__init__.py new file mode 100644 index 0000000..ce6a263 --- /dev/null +++ b/src/ucode/intelligent_routing/__init__.py @@ -0,0 +1 @@ +"""Intelligent-routing integrations for supported coding agents.""" diff --git a/src/ucode/intelligent_routing/codex_hooks.py b/src/ucode/intelligent_routing/codex_hooks.py new file mode 100644 index 0000000..f8a4a54 --- /dev/null +++ b/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 diff --git a/src/ucode/codex_routing.py b/src/ucode/intelligent_routing/codex_routing.py similarity index 91% rename from src/ucode/codex_routing.py rename to src/ucode/intelligent_routing/codex_routing.py index 6bff5e9..eb41e1e 100644 --- a/src/ucode/codex_routing.py +++ b/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, diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index 299f1e7..32eb55d 100644 --- a/tests/test_agent_codex.py +++ b/tests/test_agent_codex.py @@ -6,6 +6,7 @@ from ucode.agents import codex from ucode.config_io import read_toml_safe +from ucode.intelligent_routing import codex_routing WS = "https://example.databricks.com" @@ -352,7 +353,7 @@ def test_disable_removes_only_ucode_hooks(self, tmp_path, monkeypatch): monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", config_path) monkeypatch.setattr(codex, "LEGACY_CODEX_CONFIG_PATH", legacy_path) monkeypatch.setattr(codex, "save_state", lambda state: None) - monkeypatch.setattr("ucode.codex_routing.clear_routing_artifacts", lambda: None) + monkeypatch.setattr(codex_routing, "clear_routing_artifacts", lambda: None) state = {"workspace": WS, codex.INTELLIGENT_ROUTING_STATE_KEY: True} assert codex.disable_intelligent_routing(state) is True @@ -363,7 +364,7 @@ def test_disable_removes_only_ucode_hooks(self, tmp_path, monkeypatch): assert doc["hooks"]["PreToolUse"][0]["hooks"][0]["command"] == "user-policy" def test_launch_task_uses_exec_prompt(self): - assert codex._launch_routing_task(["exec", "fix the parser"]) == "fix the parser" + assert codex_routing._launch_routing_task(["exec", "fix the parser"]) == "fix the parser" class TestCodexRemoveLegacyProfile: diff --git a/tests/test_codex_routing.py b/tests/test_codex_routing.py index 3eb2b8c..ea3cd8c 100644 --- a/tests/test_codex_routing.py +++ b/tests/test_codex_routing.py @@ -5,7 +5,7 @@ import json import urllib.error -from ucode import codex_routing +from ucode.intelligent_routing import codex_routing WS = "https://example.databricks.com" From 4f8797e11eba545a53ad44270b1ddda83818f984 Mon Sep 17 00:00:00 2001 From: Mason Cao Date: Fri, 31 Jul 2026 18:00:30 +0000 Subject: [PATCH 3/5] Generalize smart-routing engine (rename from intelligent routing) Refactor the Codex smart-routing internals into a harness-agnostic core so a second agent can reuse them, and rename the feature "intelligent routing" -> "smart routing" throughout (package dir smart_routing/, symbols, CLI flags --enable/--disable-smart-routing, hidden codex-router-hook, user-facing strings) per naming guidance. - smart_routing/routing.py + hooks.py: shared routes:select call, decision shape, canary/audit bookkeeping, and hook add/remove machinery. - smart_routing/codex_routing.py + codex_hooks.py delegate to the shared core (behavior unchanged). - State key is the agent-neutral "smart_routing_enabled" so a later agent can share one opt-in. Co-authored-by: Isaac --- src/ucode/agents/codex.py | 41 +- src/ucode/cli.py | 63 ++- src/ucode/intelligent_routing/__init__.py | 1 - .../intelligent_routing/codex_routing.py | 372 ------------------ src/ucode/smart_routing/__init__.py | 1 + .../codex_hooks.py | 58 +-- src/ucode/smart_routing/codex_routing.py | 198 ++++++++++ src/ucode/smart_routing/hooks.py | 62 +++ src/ucode/smart_routing/routing.py | 297 ++++++++++++++ tests/test_agent_codex.py | 18 +- tests/test_cli.py | 18 +- tests/test_codex_routing.py | 12 +- 12 files changed, 641 insertions(+), 500 deletions(-) delete mode 100644 src/ucode/intelligent_routing/__init__.py delete mode 100644 src/ucode/intelligent_routing/codex_routing.py create mode 100644 src/ucode/smart_routing/__init__.py rename src/ucode/{intelligent_routing => smart_routing}/codex_hooks.py (53%) create mode 100644 src/ucode/smart_routing/codex_routing.py create mode 100644 src/ucode/smart_routing/hooks.py create mode 100644 src/ucode/smart_routing/routing.py diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index a571a0b..a25439d 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -20,11 +20,11 @@ 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.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 @@ -39,7 +39,9 @@ MINIMUM_CODEX_VERSION_TEXT = "0.134.0" MINIMUM_ROUTING_CODEX_VERSION = (0, 145, 0) MINIMUM_ROUTING_CODEX_VERSION_TEXT = "0.145.0" -INTELLIGENT_ROUTING_STATE_KEY = "codex_intelligent_routing_enabled" +# 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 = { @@ -325,10 +327,9 @@ def write_tool_config(state: dict, model: str | None = None, provider: str | Non databricks_profile = state.get("profile") if _use_legacy_layout(): - if intelligent_routing_enabled(state) and provider is None: + if smart_routing_enabled(state) and provider is None: raise RuntimeError( - "Codex intelligent routing requires Codex " - f"{MINIMUM_ROUTING_CODEX_VERSION_TEXT} or newer." + 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] @@ -370,10 +371,10 @@ 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_smart_routing_hooks( doc, state, - enabled=intelligent_routing_enabled(state) and provider is None, + enabled=smart_routing_enabled(state) and provider is None, ) write_toml_file(CODEX_CONFIG_PATH, doc) state = mark_tool_managed(state, "codex", MANAGED_KEYS) @@ -416,27 +417,27 @@ def launch(state: dict, tool_args: list[str]) -> None: exec_or_spawn([binary, "--profile", CODEX_PROFILE_NAME, *tool_args]) -def intelligent_routing_enabled(state: dict) -> bool: +def smart_routing_enabled(state: dict) -> bool: """Return whether the current workspace opted into Codex routing.""" - return state.get(INTELLIGENT_ROUTING_STATE_KEY) is True + return state.get(SMART_ROUTING_STATE_KEY) is True -def enable_intelligent_routing(state: dict) -> dict: - """Persist the current workspace's Codex intelligent-routing opt-in.""" +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 intelligent routing requires Codex " + "Codex smart routing requires Codex " f"{MINIMUM_ROUTING_CODEX_VERSION_TEXT} or newer; found " f"{agent_version(SPEC['binary'])}." ) - state[INTELLIGENT_ROUTING_STATE_KEY] = True + state[SMART_ROUTING_STATE_KEY] = True return state -def disable_intelligent_routing(state: dict) -> bool: +def disable_smart_routing(state: dict) -> bool: """Disable routing and remove only ucode's Codex routing hooks.""" - state.pop(INTELLIGENT_ROUTING_STATE_KEY, None) + state.pop(SMART_ROUTING_STATE_KEY, None) if state.get("workspace"): save_state(state) changed = False @@ -444,10 +445,10 @@ 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_smart_routing_hooks(doc): write_toml_file(path, doc) changed = True - from ucode.intelligent_routing.codex_routing import clear_routing_artifacts + from ucode.smart_routing.codex_routing import clear_routing_artifacts clear_routing_artifacts() return changed diff --git a/src/ucode/cli.py b/src/ucode/cli.py index cf659ce..3ebf087 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -30,10 +30,10 @@ launch as launch_agent, ) from ucode.agents.codex import ( - disable_intelligent_routing, - enable_intelligent_routing, - intelligent_routing_enabled, + 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 @@ -58,7 +58,6 @@ 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, @@ -68,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, @@ -968,11 +968,11 @@ def codex_router_hook_cmd( use_pat: Annotated[bool, typer.Option("--use-pat")] = False, model: Annotated[list[str] | None, typer.Option("--model")] = None, ) -> None: - """Run a Codex intelligent-routing lifecycle hook.""" + """Run a Codex smart-routing lifecycle hook.""" import json import sys - from ucode.intelligent_routing.codex_routing import ( + from ucode.smart_routing.codex_routing import ( record_session_start, record_subagent_start, route_pre_tool_use, @@ -994,7 +994,7 @@ def codex_router_hook_cmd( sys.stdout.write( json.dumps( { - "systemMessage": "Intelligent Routing verified. " + "systemMessage": "Smart Routing verified. " f"Subagent is using {record.get('model')}." } ) @@ -1003,7 +1003,7 @@ def codex_router_hook_cmd( sys.stdout.write( json.dumps( { - "systemMessage": "Intelligent Routing mismatch: router requested " + "systemMessage": "Smart Routing mismatch: router requested " f"{record.get('requested_model')}, but Codex started " f"{record.get('model')}." } @@ -1074,7 +1074,7 @@ def _launch_tool( provider: str | None = None, skip_preflight: bool = False, workspace: str | None = None, - enable_codex_intelligent_routing: bool = False, + enable_codex_smart_routing: bool = False, ) -> None: try: tool = normalize_tool(tool_name) @@ -1098,9 +1098,9 @@ 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_intelligent_routing and provider: + if tool == "codex" and enable_codex_smart_routing and provider: raise RuntimeError( - "Codex intelligent routing cannot be enabled with --provider. " + "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 @@ -1126,8 +1126,8 @@ def _launch_tool( skip_model_discovery=bool(provider), skip_preflight=skip_preflight, ) - if tool == "codex" and enable_codex_intelligent_routing: - state = enable_intelligent_routing(state) + 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 @@ -1136,16 +1136,15 @@ def _launch_tool( resolved_model = None else: state, resolved_model = resolve_launch_model(tool, state, None) - if tool == "codex" and intelligent_routing_enabled(state): - with spinner("Selecting a Codex model with intelligent routing..."): + 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(f"Using Intelligent Routing. Routing to {resolved_model}.") + print_note(f"Using Smart Routing. Routing to {resolved_model}.") elif routing_error: print_warning( - f"Intelligent routing was unavailable ({routing_error}); " - f"using {resolved_model}." + f"Smart routing was unavailable ({routing_error}); using {resolved_model}." ) state = configure_tool( tool, @@ -1160,9 +1159,9 @@ def _launch_tool( print_kv("Provider", provider) elif resolved_model: print_kv("Model", resolved_model) - if tool == "codex" and intelligent_routing_enabled(state) and not provider: - print_kv("Intelligent routing", "enabled") - if enable_codex_intelligent_routing: + 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." @@ -1221,28 +1220,28 @@ def codex_cmd( ] = None, skip_preflight: SkipPreflightOption = False, workspace: WorkspaceOption = None, - enable_intelligent_routing_flag: Annotated[ + enable_smart_routing_flag: Annotated[ bool, typer.Option( - "--enable-intelligent-routing", + "--enable-smart-routing", help="Enable AI Gateway model routing for Codex sessions and subagents.", ), ] = False, - disable_intelligent_routing_flag: Annotated[ + disable_smart_routing_flag: Annotated[ bool, typer.Option( - "--disable-intelligent-routing", - help="Disable intelligent routing and remove ucode's Codex routing hooks.", + "--disable-smart-routing", + help="Disable smart routing and remove ucode's Codex routing hooks.", ), ] = False, ) -> None: """Launch Codex via Databricks.""" - if enable_intelligent_routing_flag and disable_intelligent_routing_flag: - print_err("Use only one of --enable-intelligent-routing or --disable-intelligent-routing.") + 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_intelligent_routing_flag: - disable_intelligent_routing(load_state()) - print_success("Codex intelligent routing disabled; ucode routing hooks removed") + if disable_smart_routing_flag: + disable_smart_routing(load_state()) + print_success("Codex smart routing disabled; ucode routing hooks removed") return _launch_tool( "codex", @@ -1250,7 +1249,7 @@ def codex_cmd( provider=provider, skip_preflight=skip_preflight, workspace=workspace, - enable_codex_intelligent_routing=enable_intelligent_routing_flag, + enable_codex_smart_routing=enable_smart_routing_flag, ) diff --git a/src/ucode/intelligent_routing/__init__.py b/src/ucode/intelligent_routing/__init__.py deleted file mode 100644 index ce6a263..0000000 --- a/src/ucode/intelligent_routing/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Intelligent-routing integrations for supported coding agents.""" diff --git a/src/ucode/intelligent_routing/codex_routing.py b/src/ucode/intelligent_routing/codex_routing.py deleted file mode 100644 index eb41e1e..0000000 --- a/src/ucode/intelligent_routing/codex_routing.py +++ /dev/null @@ -1,372 +0,0 @@ -"""Databricks AI Gateway routing helpers for Codex sessions and subagents.""" - -from __future__ import annotations - -import json -import re -import time -import urllib.error -import urllib.request -import uuid -from dataclasses import dataclass -from pathlib import Path -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" -REQUEST_TIMEOUT_S = 30.0 -CODEX_ROUTE_ARMS = ("glm-5-2", "gpt-5-6-sol", "gpt-5-6-luna") -GLM_ROUTE_ARM = "glm-5-2" -GLM_GATEWAY_MODEL = "system.ai.glm-5-2" -SPAWN_AGENT_TOOL_SUFFIX = "spawn_agent" -CANARY_PATH = APP_DIR / "codex-intelligent-routing-canary.json" -AUDIT_PATH = APP_DIR / "codex-intelligent-routing-audit.jsonl" -DECISIONS_PATH = APP_DIR / "codex-intelligent-routing-decisions.jsonl" - -_GPT_RE = re.compile(r"gpt-(\d+)(?:[.-](\d+))?(?:[.-](\d+))?(-.+|[a-z].*)?") - - -@dataclass(frozen=True) -class RoutingDecision: - """One model selection returned by the AI Gateway router.""" - - model: str - raw_model: str - 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, - task: str, - available_models: list[str], - *, - timeout: float = REQUEST_TIMEOUT_S, -) -> tuple[RoutingDecision | None, str | None]: - """Ask the workspace ``task_v1`` router for a servable Codex model.""" - candidates = _routing_candidates(available_models) - missing = [ - arm for arm in CODEX_ROUTE_ARMS if arm not in {_normalize_model(m) for m in candidates} - ] - if missing: - return None, f"required Codex routing models are unavailable: {', '.join(missing)}" - - body = { - "route_options": [{"model": model, "harness": "codex"} for model in CODEX_ROUTE_ARMS], - "task": {"prompt": task[:4000]}, - "route_selector": {"router_name": ROUTER_NAME}, - } - request = urllib.request.Request( - workspace.rstrip("/") + ROUTING_PATH, - data=json.dumps(body).encode("utf-8"), - headers={ - "Authorization": f"Bearer {token}", - "Content-Type": "application/json", - }, - method="POST", - ) - try: - with urllib.request.urlopen(request, timeout=timeout) as response: - payload = json.loads(response.read().decode("utf-8")) - except urllib.error.HTTPError as exc: - detail = "" - try: - detail = exc.read().decode("utf-8", errors="replace").strip() - except OSError: - pass - reason = f"router returned HTTP {exc.code}" - if detail: - reason = f"{reason}: {detail[:300]}" - return None, reason - except (urllib.error.URLError, OSError, TimeoutError, ValueError) as exc: - return None, f"router request failed: {exc}" - - raw_model = _selected_model(payload) - if raw_model is None: - return None, "router returned no model selection" - model = resolve_routed_model(raw_model, candidates) - if model is None: - return None, f"router selected unsupported model {raw_model!r}" - rationale = payload.get("rationale") if isinstance(payload, dict) else None - return ( - RoutingDecision( - model=model, - raw_model=raw_model, - rationale=rationale if isinstance(rationale, str) else "", - ), - None, - ) - - -def resolve_routed_model(raw_model: str, available_models: list[str]) -> str | None: - """Map a ``task_v1`` arm to a model the configured workspace can serve.""" - candidates = _routing_candidates(available_models) - normalized = {_normalize_model(model): model for model in candidates} - return normalized.get(_normalize_model(raw_model)) - - -def route_pre_tool_use( - payload: dict[str, Any], - *, - workspace: str, - token: str, - available_models: list[str], - timeout: float = REQUEST_TIMEOUT_S, - audit_decision: bool = False, -) -> dict[str, Any] | None: - """Route one Codex ``spawn_agent`` call and rewrite its model.""" - if not is_spawn_agent_tool(payload.get("tool_name")): - return None - tool_input = payload.get("tool_input") - if not isinstance(tool_input, dict): - return None - task_name = tool_input.get("task_name") or tool_input.get("agent_name") - task = task_name if isinstance(task_name, str) and task_name else "Codex subagent task" - decision, _ = request_routing_decision( - workspace, - token, - task, - available_models, - timeout=timeout, - ) - if decision is None: - return None - if decision.raw_model == GLM_ROUTE_ARM: - return { - "systemMessage": ( - "Intelligent Routing selected GLM 5.2, which is not enabled for Codex " - "subagents. Keeping the original subagent model." - ) - } - routed_model = _codex_model_id(decision.model) - if audit_decision: - _record_routing_decision(payload, task, decision, routed_model) - routing_message = f"Using Intelligent Routing. Routing to {routed_model}." - output: dict[str, Any] = { - "hookEventName": "PreToolUse", - "permissionDecision": "allow", - "updatedInput": {**tool_input, "model": routed_model}, - "permissionDecisionReason": routing_message, - } - if decision.rationale: - output["permissionDecisionReason"] += f" {decision.rationale}" - return {"systemMessage": routing_message, "hookSpecificOutput": output} - - -def is_spawn_agent_tool(tool_name: Any) -> bool: - """Return whether a hook payload names Codex's subagent spawn tool.""" - if not isinstance(tool_name, str): - return False - normalized = tool_name.strip().lower() - return normalized == "agent" or normalized.endswith(SPAWN_AGENT_TOOL_SUFFIX) - - -def record_session_start(payload: dict[str, Any]) -> None: - """Write a canary proving Codex trusted and ran the routing hooks.""" - _write_json( - CANARY_PATH, - { - "session_id": payload.get("session_id"), - "model": payload.get("model"), - "at": time.time(), - }, - ) - - -def record_subagent_start(payload: dict[str, Any]) -> dict[str, Any]: - """Append the model Codex actually selected for a routed subagent.""" - actual_model = payload.get("model") - decision = _pending_decision(payload.get("session_id"), actual_model) - record = { - "agent_id": payload.get("agent_id"), - "agent_type": payload.get("agent_type"), - "model": actual_model, - "session_id": payload.get("session_id"), - "at": time.time(), - } - if decision is not None: - record.update( - { - "decision_id": decision.get("decision_id"), - "router_model": decision.get("router_model"), - "requested_model": decision.get("requested_model"), - "matches_router_decision": decision.get("requested_model") == actual_model, - } - ) - _append_jsonl(AUDIT_PATH, record) - return record - - -def clear_routing_artifacts() -> None: - """Remove ucode-owned routing canary and audit files.""" - for path in (CANARY_PATH, AUDIT_PATH, DECISIONS_PATH): - try: - path.unlink(missing_ok=True) - except OSError: - continue - - -def _selected_model(payload: Any) -> str | None: - if not isinstance(payload, dict): - return None - selections = payload.get("route_selection") - if not isinstance(selections, list) or not selections: - return None - selection = selections[0] - if not isinstance(selection, dict): - return None - option = selection.get("route_option") - if not isinstance(option, dict): - return None - model = option.get("model") - return model if isinstance(model, str) and model else None - - -def _routing_candidates(models: list[str]) -> list[str]: - candidates = [model for model in models if isinstance(model, str) and model] - if GLM_ROUTE_ARM not in {_normalize_model(model) for model in candidates}: - candidates.append(GLM_GATEWAY_MODEL) - return candidates - - -def _parse_gpt(model: str) -> tuple[int, int, int, str] | None: - match = _GPT_RE.fullmatch(_normalize_model(model)) - if not match: - return None - major, minor, patch, suffix = match.groups() - return int(major), int(minor or 0), int(patch or 0), suffix or "" - - -def _model_strength(model: str) -> tuple[int, int, int, int]: - parsed = _parse_gpt(model) - if parsed is None: - return (0, 0, 0, 0) - major, minor, patch, suffix = parsed - return major, minor, patch, 1 if not suffix else 0 - - -def _normalize_model(model: str) -> str: - tail = model.rsplit("/", 1)[-1] - for prefix in ("databricks-", "system.ai."): - if tail.startswith(prefix): - tail = tail[len(prefix) :] - break - return tail.lower() - - -def _codex_model_id(model: str) -> str: - tail = model.rsplit("/", 1)[-1] - if tail in {"databricks-gpt-5-2-codex", "databricks-gpt-5-4-nano"}: - return tail - if _normalize_model(model) == GLM_ROUTE_ARM: - return GLM_GATEWAY_MODEL - if model.startswith("system.ai."): - bare = model.removeprefix("system.ai.") - elif tail.startswith("databricks-"): - bare = tail.removeprefix("databricks-") - else: - return model - match = _GPT_RE.fullmatch(bare) - if not match: - return bare - major, minor, patch, suffix = match.groups() - version = major - if minor is not None: - version += f".{minor}" - if patch is not None: - version += f".{patch}" - return f"gpt-{version}{suffix or ''}" - - -def _record_routing_decision( - payload: dict[str, Any], - task_name: str, - decision: RoutingDecision, - requested_model: str, -) -> None: - _append_jsonl( - DECISIONS_PATH, - { - "decision_id": uuid.uuid4().hex, - "session_id": payload.get("session_id"), - "task_name": task_name, - "router_model": decision.raw_model, - "requested_model": requested_model, - "at": time.time(), - }, - ) - - -def _pending_decision(session_id: Any, actual_model: Any) -> dict[str, Any] | None: - used = { - record.get("decision_id") for record in _read_jsonl(AUDIT_PATH) if record.get("decision_id") - } - pending = [ - decision - for decision in _read_jsonl(DECISIONS_PATH) - if decision.get("session_id") == session_id and decision.get("decision_id") not in used - ] - return next( - (decision for decision in pending if decision.get("requested_model") == actual_model), - pending[0] if pending else None, - ) - - -def _read_jsonl(path: Path) -> list[dict[str, Any]]: - try: - lines = path.read_text(encoding="utf-8").splitlines() - except OSError: - return [] - records = [] - for line in lines: - try: - record = json.loads(line) - except ValueError: - continue - if isinstance(record, dict): - records.append(record) - return records - - -def _append_jsonl(path: Path, payload: dict[str, Any]) -> None: - try: - path.parent.mkdir(parents=True, exist_ok=True) - with path.open("a", encoding="utf-8") as handle: - handle.write(json.dumps(payload) + "\n") - except OSError: - return - - -def _write_json(path: Path, payload: dict[str, Any]) -> None: - try: - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(json.dumps(payload), encoding="utf-8") - except OSError: - return diff --git a/src/ucode/smart_routing/__init__.py b/src/ucode/smart_routing/__init__.py new file mode 100644 index 0000000..c6ba0f9 --- /dev/null +++ b/src/ucode/smart_routing/__init__.py @@ -0,0 +1 @@ +"""Smart-routing integrations for supported coding agents.""" diff --git a/src/ucode/intelligent_routing/codex_hooks.py b/src/ucode/smart_routing/codex_hooks.py similarity index 53% rename from src/ucode/intelligent_routing/codex_hooks.py rename to src/ucode/smart_routing/codex_hooks.py index f8a4a54..6c6a569 100644 --- a/src/ucode/intelligent_routing/codex_hooks.py +++ b/src/ucode/smart_routing/codex_hooks.py @@ -1,4 +1,4 @@ -"""Codex hook configuration for intelligent subagent routing.""" +"""Codex hook configuration for smart subagent routing.""" from __future__ import annotations @@ -6,62 +6,20 @@ import subprocess from ucode.databricks import build_auth_token_argv +from ucode.smart_routing import hooks ROUTING_HOOK_COMMAND_MARKER = "codex-router-hook" -def sync_intelligent_routing_hooks(doc: dict, state: dict, *, enabled: bool) -> None: +def sync_smart_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] + groups = _routing_hook_groups(state) if enabled else {} + hooks.sync_managed_hooks(doc, ROUTING_HOOK_COMMAND_MARKER, 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 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]]: diff --git a/src/ucode/smart_routing/codex_routing.py b/src/ucode/smart_routing/codex_routing.py new file mode 100644 index 0000000..c529bd7 --- /dev/null +++ b/src/ucode/smart_routing/codex_routing.py @@ -0,0 +1,198 @@ +"""Databricks AI Gateway routing helpers for Codex sessions and subagents. + +Codex-specific configuration on top of the shared :mod:`ucode.smart_routing.routing` +core: the frozen ``task_v1`` Codex route arms, the ``spawn_agent`` tool detector, +the Codex model-id translation, and the artifact paths. +""" + +from __future__ import annotations + +import re + +# Re-exported so tests can patch the shared ``urlopen`` seam via +# ``codex_routing.urllib.request`` — the actual call lives in ``routing``, but +# Python modules are singletons so patching this name patches the one call site. +import urllib.request # noqa: F401 +from pathlib import Path +from typing import Any + +from ucode.config_io import APP_DIR +from ucode.databricks import get_databricks_token +from ucode.smart_routing import routing +from ucode.smart_routing.routing import RoutingDecision + +ROUTER_NAME = routing.ROUTER_NAME +ROUTING_PATH = routing.ROUTING_PATH +REQUEST_TIMEOUT_S = routing.REQUEST_TIMEOUT_S +CODEX_ROUTE_ARMS = ("glm-5-2", "gpt-5-6-sol", "gpt-5-6-luna") +GLM_ROUTE_ARM = "glm-5-2" +GLM_GATEWAY_MODEL = "system.ai.glm-5-2" +SPAWN_AGENT_TOOL_SUFFIX = "spawn_agent" +CANARY_PATH = APP_DIR / "codex-smart-routing-canary.json" +AUDIT_PATH = APP_DIR / "codex-smart-routing-audit.jsonl" +DECISIONS_PATH = APP_DIR / "codex-smart-routing-decisions.jsonl" + +GLM_SUBAGENT_SKIP_MESSAGE = ( + "Smart Routing selected GLM 5.2, which is not enabled for Codex " + "subagents. Keeping the original subagent model." +) + +_GPT_RE = re.compile(r"gpt-(\d+)(?:[.-](\d+))?(?:[.-](\d+))?(-.+|[a-z].*)?") + +_normalize_model = routing.normalize_model + + +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, + task: str, + available_models: list[str], + *, + timeout: float = REQUEST_TIMEOUT_S, +) -> tuple[RoutingDecision | None, str | None]: + """Ask the workspace ``task_v1`` router for a servable Codex model.""" + candidates = _routing_candidates(available_models) + missing = [ + arm for arm in CODEX_ROUTE_ARMS if arm not in {_normalize_model(m) for m in candidates} + ] + if missing: + return None, f"required Codex routing models are unavailable: {', '.join(missing)}" + + return routing.select_route( + workspace, + token, + task, + [(arm, "codex") for arm in CODEX_ROUTE_ARMS], + lambda raw_model: resolve_routed_model(raw_model, candidates), + timeout=timeout, + ) + + +def resolve_routed_model(raw_model: str, available_models: list[str]) -> str | None: + """Map a ``task_v1`` arm to a model the configured workspace can serve.""" + candidates = _routing_candidates(available_models) + normalized = {_normalize_model(model): model for model in candidates} + return normalized.get(_normalize_model(raw_model)) + + +def route_pre_tool_use( + payload: dict[str, Any], + *, + workspace: str, + token: str, + available_models: list[str], + timeout: float = REQUEST_TIMEOUT_S, + audit_decision: bool = False, +) -> dict[str, Any] | None: + """Route one Codex ``spawn_agent`` call and rewrite its model.""" + record = None + if audit_decision: + + def record(payload, task, decision, requested): + routing.write_decision_record(DECISIONS_PATH, payload, task, decision, requested) + + return routing.route_spawn_tool( + payload, + is_spawn_agent=is_spawn_agent_tool, + decision_fn=lambda task: request_routing_decision( + workspace, token, task, available_models, timeout=timeout + ), + default_task_label="Codex subagent task", + model_id_mapper=_codex_model_id, + skip_arms={GLM_ROUTE_ARM: GLM_SUBAGENT_SKIP_MESSAGE}, + record_decision=record, + ) + + +def is_spawn_agent_tool(tool_name: Any) -> bool: + """Return whether a hook payload names Codex's subagent spawn tool.""" + if not isinstance(tool_name, str): + return False + normalized = tool_name.strip().lower() + return normalized == "agent" or normalized.endswith(SPAWN_AGENT_TOOL_SUFFIX) + + +def record_session_start(payload: dict[str, Any]) -> None: + """Write a canary proving Codex trusted and ran the routing hooks.""" + routing.record_session_start(CANARY_PATH, payload) + + +def record_subagent_start(payload: dict[str, Any]) -> dict[str, Any]: + """Append the model Codex actually selected for a routed subagent.""" + return routing.record_subagent_start(DECISIONS_PATH, AUDIT_PATH, payload) + + +def clear_routing_artifacts() -> None: + """Remove ucode-owned routing canary and audit files.""" + routing.clear_artifacts((CANARY_PATH, AUDIT_PATH, DECISIONS_PATH)) + + +def _routing_candidates(models: list[str]) -> list[str]: + candidates = [model for model in models if isinstance(model, str) and model] + if GLM_ROUTE_ARM not in {_normalize_model(model) for model in candidates}: + candidates.append(GLM_GATEWAY_MODEL) + return candidates + + +def _parse_gpt(model: str) -> tuple[int, int, int, str] | None: + match = _GPT_RE.fullmatch(_normalize_model(model)) + if not match: + return None + major, minor, patch, suffix = match.groups() + return int(major), int(minor or 0), int(patch or 0), suffix or "" + + +def _model_strength(model: str) -> tuple[int, int, int, int]: + parsed = _parse_gpt(model) + if parsed is None: + return (0, 0, 0, 0) + major, minor, patch, suffix = parsed + return major, minor, patch, 1 if not suffix else 0 + + +def _codex_model_id(model: str) -> str: + tail = model.rsplit("/", 1)[-1] + if tail in {"databricks-gpt-5-2-codex", "databricks-gpt-5-4-nano"}: + return tail + if _normalize_model(model) == GLM_ROUTE_ARM: + return GLM_GATEWAY_MODEL + if model.startswith("system.ai."): + bare = model.removeprefix("system.ai.") + elif tail.startswith("databricks-"): + bare = tail.removeprefix("databricks-") + else: + return model + match = _GPT_RE.fullmatch(bare) + if not match: + return bare + major, minor, patch, suffix = match.groups() + version = major + if minor is not None: + version += f".{minor}" + if patch is not None: + version += f".{patch}" + return f"gpt-{version}{suffix or ''}" diff --git a/src/ucode/smart_routing/hooks.py b/src/ucode/smart_routing/hooks.py new file mode 100644 index 0000000..1c00c69 --- /dev/null +++ b/src/ucode/smart_routing/hooks.py @@ -0,0 +1,62 @@ +"""Shared add/remove machinery for ucode-managed smart-routing hooks. + +Codex config (TOML) and Claude Code settings (JSON) both express hooks as +``{event: [{matcher?, hooks: [{command, ...}]}]}``, and both identify a +ucode-managed hook by a marker substring in its ``command``. This module owns +the surgical merge (add ucode's groups, leave user hooks untouched) and the +surgical removal (drop only hooks whose command carries the marker), so each +harness module supplies just its own per-event hook groups. +""" + +from __future__ import annotations + + +def sync_managed_hooks(doc: dict, marker: str, groups: dict[str, list[dict]]) -> None: + """Replace ucode's marked hooks with ``groups``, preserving user hooks.""" + remove_managed_hooks(doc, marker) + if not groups: + return + hooks = doc.setdefault("hooks", {}) + for event, event_groups in groups.items(): + existing = hooks.get(event) + if not isinstance(existing, list): + existing = [] + hooks[event] = [*existing, *event_groups] + + +def remove_managed_hooks(doc: dict, marker: str) -> bool: + """Remove only hooks whose command contains ``marker``. Returns True if any.""" + 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 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 diff --git a/src/ucode/smart_routing/routing.py b/src/ucode/smart_routing/routing.py new file mode 100644 index 0000000..f157096 --- /dev/null +++ b/src/ucode/smart_routing/routing.py @@ -0,0 +1,297 @@ +"""Shared AI Gateway routing helpers for coding-agent sessions and subagents. + +Both the Codex and Claude Code integrations route through the workspace's +``task_v1`` router at ``/ai-gateway/routing/v1/routes:select``. The +harness-agnostic mechanics live here — the gateway call, the decision shape, +model-name normalization, and the canary/audit/decision bookkeeping. Each +harness module (``codex_routing`` / ``claude_routing``) supplies its own route +arms, spawn-tool detector, model-id translation, and artifact paths. +""" + +from __future__ import annotations + +import json +import time +import urllib.error +import urllib.request +import uuid +from collections.abc import Callable, Iterable +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +ROUTER_NAME = "task_v1" +ROUTING_PATH = "/ai-gateway/routing/v1/routes:select" +REQUEST_TIMEOUT_S = 30.0 + + +@dataclass(frozen=True) +class RoutingDecision: + """One model selection returned by the AI Gateway router.""" + + model: str + raw_model: str + rationale: str = "" + + +def normalize_model(model: str) -> str: + """Strip provider prefixes so router arms and workspace ids compare equal. + + ``system.ai.claude-opus-4-8`` and ``databricks-claude-opus-4-8`` both + normalize to ``claude-opus-4-8`` — the router's canonical arm vocabulary. + """ + tail = model.rsplit("/", 1)[-1] + for prefix in ("databricks-", "system.ai."): + if tail.startswith(prefix): + tail = tail[len(prefix) :] + break + return tail.lower() + + +def select_route( + workspace: str, + token: str, + task: str, + route_options: Iterable[tuple[str, str]], + resolve: Callable[[str], str | None], + *, + router_name: str = ROUTER_NAME, + timeout: float = REQUEST_TIMEOUT_S, +) -> tuple[RoutingDecision | None, str | None]: + """POST one ``routes:select`` request and resolve the router's pick. + + ``route_options`` are ``(model, harness)`` pairs offered to the router (the + frozen menu the caller's harness scenario requires). ``resolve`` maps the + router's chosen arm back to a model id the workspace can serve, returning + None when the arm is unservable. Returns ``(decision, error)``; a failed + call yields ``(None, reason)`` so callers can fail open. + """ + body = { + "route_options": [{"model": model, "harness": harness} for model, harness in route_options], + "task": {"prompt": task[:4000]}, + "route_selector": {"router_name": router_name}, + } + request = urllib.request.Request( + workspace.rstrip("/") + ROUTING_PATH, + data=json.dumps(body).encode("utf-8"), + headers={ + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + }, + method="POST", + ) + try: + with urllib.request.urlopen(request, timeout=timeout) as response: + payload = json.loads(response.read().decode("utf-8")) + except urllib.error.HTTPError as exc: + detail = "" + try: + detail = exc.read().decode("utf-8", errors="replace").strip() + except OSError: + pass + reason = f"router returned HTTP {exc.code}" + if detail: + reason = f"{reason}: {detail[:300]}" + return None, reason + except (urllib.error.URLError, OSError, TimeoutError, ValueError) as exc: + return None, f"router request failed: {exc}" + + raw_model = _selected_model(payload) + if raw_model is None: + return None, "router returned no model selection" + model = resolve(raw_model) + if model is None: + return None, f"router selected unsupported model {raw_model!r}" + rationale = payload.get("rationale") if isinstance(payload, dict) else None + return ( + RoutingDecision( + model=model, + raw_model=raw_model, + rationale=rationale if isinstance(rationale, str) else "", + ), + None, + ) + + +def route_spawn_tool( + payload: dict[str, Any], + *, + is_spawn_agent: Callable[[Any], bool], + decision_fn: Callable[[str], tuple[RoutingDecision | None, str | None]], + default_task_label: str, + model_id_mapper: Callable[[str], str], + skip_arms: dict[str, str] | None = None, + record_decision: Callable[[dict[str, Any], str, RoutingDecision, str], None] | None = None, +) -> dict[str, Any] | None: + """Route one subagent-spawn tool call, rewriting its ``model`` input. + + Returns a PreToolUse hook output that allows the call with the routed model + injected; a bare ``systemMessage`` when the pick is an arm the harness can't + use for subagents (``skip_arms``); or None when the tool is not a spawn or + routing was unavailable — fail open, leaving the original model in place. + """ + if not is_spawn_agent(payload.get("tool_name")): + return None + tool_input = payload.get("tool_input") + if not isinstance(tool_input, dict): + return None + task_name = tool_input.get("task_name") or tool_input.get("agent_name") + task = task_name if isinstance(task_name, str) and task_name else default_task_label + decision, _ = decision_fn(task) + if decision is None: + return None + if skip_arms and decision.raw_model in skip_arms: + return {"systemMessage": skip_arms[decision.raw_model]} + routed_model = model_id_mapper(decision.model) + if record_decision is not None: + record_decision(payload, task, decision, routed_model) + routing_message = f"Using Smart Routing. Routing to {routed_model}." + output: dict[str, Any] = { + "hookEventName": "PreToolUse", + "permissionDecision": "allow", + "updatedInput": {**tool_input, "model": routed_model}, + "permissionDecisionReason": routing_message, + } + if decision.rationale: + output["permissionDecisionReason"] += f" {decision.rationale}" + return {"systemMessage": routing_message, "hookSpecificOutput": output} + + +def record_session_start(canary_path: Path, payload: dict[str, Any]) -> None: + """Write a canary proving the harness trusted and ran the routing hooks.""" + _write_json( + canary_path, + { + "session_id": payload.get("session_id"), + "model": payload.get("model"), + "at": time.time(), + }, + ) + + +def record_subagent_start( + decisions_path: Path, audit_path: Path, payload: dict[str, Any] +) -> dict[str, Any]: + """Append the model the harness actually selected for a routed subagent. + + Reconciles against the pending routing decision for the session (if any) so + the audit row records whether the launched model matched the router's pick. + """ + actual_model = payload.get("model") + decision = _pending_decision( + decisions_path, audit_path, payload.get("session_id"), actual_model + ) + record = { + "agent_id": payload.get("agent_id"), + "agent_type": payload.get("agent_type"), + "model": actual_model, + "session_id": payload.get("session_id"), + "at": time.time(), + } + if decision is not None: + record.update( + { + "decision_id": decision.get("decision_id"), + "router_model": decision.get("router_model"), + "requested_model": decision.get("requested_model"), + "matches_router_decision": decision.get("requested_model") == actual_model, + } + ) + _append_jsonl(audit_path, record) + return record + + +def write_decision_record( + decisions_path: Path, + payload: dict[str, Any], + task_name: str, + decision: RoutingDecision, + requested_model: str, +) -> None: + """Record a routing decision so a later SubagentStart can reconcile it.""" + _append_jsonl( + decisions_path, + { + "decision_id": uuid.uuid4().hex, + "session_id": payload.get("session_id"), + "task_name": task_name, + "router_model": decision.raw_model, + "requested_model": requested_model, + "at": time.time(), + }, + ) + + +def clear_artifacts(paths: Iterable[Path]) -> None: + """Remove ucode-owned routing canary/audit/decision files.""" + for path in paths: + try: + path.unlink(missing_ok=True) + except OSError: + continue + + +def _selected_model(payload: Any) -> str | None: + if not isinstance(payload, dict): + return None + selections = payload.get("route_selection") + if not isinstance(selections, list) or not selections: + return None + selection = selections[0] + if not isinstance(selection, dict): + return None + option = selection.get("route_option") + if not isinstance(option, dict): + return None + model = option.get("model") + return model if isinstance(model, str) and model else None + + +def _pending_decision( + decisions_path: Path, audit_path: Path, session_id: Any, actual_model: Any +) -> dict[str, Any] | None: + used = { + record.get("decision_id") for record in _read_jsonl(audit_path) if record.get("decision_id") + } + pending = [ + decision + for decision in _read_jsonl(decisions_path) + if decision.get("session_id") == session_id and decision.get("decision_id") not in used + ] + return next( + (decision for decision in pending if decision.get("requested_model") == actual_model), + pending[0] if pending else None, + ) + + +def _read_jsonl(path: Path) -> list[dict[str, Any]]: + try: + lines = path.read_text(encoding="utf-8").splitlines() + except OSError: + return [] + records = [] + for line in lines: + try: + record = json.loads(line) + except ValueError: + continue + if isinstance(record, dict): + records.append(record) + return records + + +def _append_jsonl(path: Path, payload: dict[str, Any]) -> None: + try: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("a", encoding="utf-8") as handle: + handle.write(json.dumps(payload) + "\n") + except OSError: + return + + +def _write_json(path: Path, payload: dict[str, Any]) -> None: + try: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(payload), encoding="utf-8") + except OSError: + return diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index 32eb55d..6c03aaa 100644 --- a/tests/test_agent_codex.py +++ b/tests/test_agent_codex.py @@ -6,7 +6,7 @@ from ucode.agents import codex from ucode.config_io import read_toml_safe -from ucode.intelligent_routing import codex_routing +from ucode.smart_routing import codex_routing WS = "https://example.databricks.com" @@ -218,7 +218,7 @@ def test_writes_legacy_shared_config_when_codex_too_old(self, tmp_path, monkeypa assert provider["base_url"] == f"{WS}/ai-gateway/codex/v1" assert provider["wire_api"] == "responses" - def test_intelligent_routing_writes_profile_scoped_hooks(self, tmp_path, monkeypatch): + def test_smart_routing_writes_profile_scoped_hooks(self, tmp_path, monkeypatch): config_path = tmp_path / ".codex" / "ucode.config.toml" config_path.parent.mkdir() config_path.write_text( @@ -239,7 +239,7 @@ def test_intelligent_routing_writes_profile_scoped_hooks(self, tmp_path, monkeyp "workspace": WS, "profile": "prod", "codex_models": ["databricks-gpt-5", "databricks-gpt-5-5"], - codex.INTELLIGENT_ROUTING_STATE_KEY: True, + codex.SMART_ROUTING_STATE_KEY: True, } ) @@ -267,7 +267,7 @@ def test_provider_launch_removes_routing_hooks(self, tmp_path, monkeypatch): state = { "workspace": WS, "codex_models": ["databricks-gpt-5"], - codex.INTELLIGENT_ROUTING_STATE_KEY: True, + codex.SMART_ROUTING_STATE_KEY: True, } codex.write_tool_config(state) @@ -319,12 +319,12 @@ def test_unknown_version_uses_modern_layout(self, monkeypatch): assert codex._use_legacy_layout() is False -class TestCodexIntelligentRouting: +class TestCodexSmartRouting: def test_enable_requires_supported_codex(self, monkeypatch): monkeypatch.setattr(codex, "agent_version", lambda binary: "0.144.0") try: - codex.enable_intelligent_routing({}) + codex.enable_smart_routing({}) assert False except RuntimeError as exc: assert "0.145.0 or newer" in str(exc) @@ -354,12 +354,12 @@ def test_disable_removes_only_ucode_hooks(self, tmp_path, monkeypatch): monkeypatch.setattr(codex, "LEGACY_CODEX_CONFIG_PATH", legacy_path) monkeypatch.setattr(codex, "save_state", lambda state: None) monkeypatch.setattr(codex_routing, "clear_routing_artifacts", lambda: None) - state = {"workspace": WS, codex.INTELLIGENT_ROUTING_STATE_KEY: True} + state = {"workspace": WS, codex.SMART_ROUTING_STATE_KEY: True} - assert codex.disable_intelligent_routing(state) is True + assert codex.disable_smart_routing(state) is True doc = read_toml_safe(config_path) - assert state.get(codex.INTELLIGENT_ROUTING_STATE_KEY) is None + assert state.get(codex.SMART_ROUTING_STATE_KEY) is None assert list(doc["hooks"]) == ["PreToolUse"] assert doc["hooks"]["PreToolUse"][0]["hooks"][0]["command"] == "user-policy" diff --git a/tests/test_cli.py b/tests/test_cli.py index ec8e1b5..cf07a83 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -190,21 +190,21 @@ def test_no_workspace_flag_leaves_current_workspace(self): assert result.exit_code == 0, result.output mock_set.assert_not_called() - def test_codex_enable_intelligent_routing_is_consumed_by_ucode(self): + def test_codex_enable_smart_routing_is_consumed_by_ucode(self): with patch("ucode.cli._launch_tool") as mock_launch: - result = runner.invoke(app, ["codex", "--enable-intelligent-routing"]) + result = runner.invoke(app, ["codex", "--enable-smart-routing"]) assert result.exit_code == 0, result.output - assert mock_launch.call_args.kwargs["enable_codex_intelligent_routing"] is True + assert mock_launch.call_args.kwargs["enable_codex_smart_routing"] is True assert mock_launch.call_args.args[1].args == [] def test_codex_disable_removes_hooks_without_launching(self): with ( patch("ucode.cli.load_state", return_value=MINIMAL_STATE), - patch("ucode.cli.disable_intelligent_routing") as mock_disable, + patch("ucode.cli.disable_smart_routing") as mock_disable, patch("ucode.cli._launch_tool") as mock_launch, ): - result = runner.invoke(app, ["codex", "--disable-intelligent-routing"]) + result = runner.invoke(app, ["codex", "--disable-smart-routing"]) assert result.exit_code == 0, result.output mock_disable.assert_called_once_with(MINIMAL_STATE) @@ -214,7 +214,7 @@ def test_codex_disable_removes_hooks_without_launching(self): def test_codex_routing_flags_are_mutually_exclusive(self): result = runner.invoke( app, - ["codex", "--enable-intelligent-routing", "--disable-intelligent-routing"], + ["codex", "--enable-smart-routing", "--disable-smart-routing"], ) assert result.exit_code == 1 @@ -223,7 +223,7 @@ def test_codex_routing_flags_are_mutually_exclusive(self): def test_enabled_codex_launch_uses_routed_root_model(self): state = { **MINIMAL_STATE, - "codex_intelligent_routing_enabled": True, + "smart_routing_enabled": True, "codex_models": ["databricks-gpt-5", "databricks-gpt-5-5"], } decision = MagicMock(model="databricks-gpt-5-5") @@ -244,9 +244,7 @@ def test_enabled_codex_launch_uses_routed_root_model(self): assert result.exit_code == 0, result.output assert mock_configure.call_args.args[2] == "databricks-gpt-5-5" - assert "Using Intelligent Routing. Routing to databricks-gpt-5-5." in _strip_ansi( - result.output - ) + assert "Using Smart Routing. Routing to databricks-gpt-5-5." in _strip_ansi(result.output) class TestMcpSubcommands: diff --git a/tests/test_codex_routing.py b/tests/test_codex_routing.py index ea3cd8c..bd85041 100644 --- a/tests/test_codex_routing.py +++ b/tests/test_codex_routing.py @@ -1,11 +1,11 @@ -"""Tests for Codex intelligent-routing hooks.""" +"""Tests for Codex smart-routing hooks.""" from __future__ import annotations import json import urllib.error -from ucode.intelligent_routing import codex_routing +from ucode.smart_routing import codex_routing WS = "https://example.databricks.com" @@ -143,7 +143,7 @@ def test_spawn_rewrite_preserves_original_input(monkeypatch): ) hook = output["hookSpecificOutput"] - assert output["systemMessage"] == "Using Intelligent Routing. Routing to gpt-5.5." + assert output["systemMessage"] == "Using Smart Routing. Routing to gpt-5.5." assert hook["permissionDecision"] == "allow" assert hook["updatedInput"] == { "task_name": "reviewer", @@ -152,7 +152,7 @@ def test_spawn_rewrite_preserves_original_input(monkeypatch): "model": "gpt-5.5", } assert hook["permissionDecisionReason"] == ( - "Using Intelligent Routing. Routing to gpt-5.5. Review needs deeper reasoning." + "Using Smart Routing. Routing to gpt-5.5. Review needs deeper reasoning." ) @@ -179,7 +179,7 @@ def test_spawn_rewrite_uses_codex_model_id_for_uc_endpoint(monkeypatch): available_models=["system.ai.gpt-5-6-luna"], ) - assert output["systemMessage"] == "Using Intelligent Routing. Routing to gpt-5.6-luna." + assert output["systemMessage"] == "Using Smart Routing. Routing to gpt-5.6-luna." assert output["hookSpecificOutput"]["updatedInput"]["model"] == "gpt-5.6-luna" @@ -208,7 +208,7 @@ def test_spawn_glm_decision_keeps_original_model(monkeypatch): assert output == { "systemMessage": ( - "Intelligent Routing selected GLM 5.2, which is not enabled for Codex " + "Smart Routing selected GLM 5.2, which is not enabled for Codex " "subagents. Keeping the original subagent model." ) } From f2d6572584685458734b3f8c6b69b0b7db3a6976 Mon Sep 17 00:00:00 2001 From: Mason Cao Date: Fri, 31 Jul 2026 19:32:51 +0000 Subject: [PATCH 4/5] Route root model on the launch-time seed prompt when present MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When the user's first prompt is on the command line (`codex ""`, `codex exec ""`, or after `--`), route the root-session model on that actual prompt. A bare interactive launch carries no prompt, so routing is skipped entirely and the user's configured default model is kept — with no task signal the router could only return its floor arm, so routing there would add a round-trip and silently downgrade the default. Routing a typed-in first prompt is out of scope (no hook/MCP can retarget the root model mid-session on Codex or Claude Code). Adds a shared, conservative `extract_seed_prompt` to routing.py: it skips value-taking options and their values, treats unknown flags as booleans, and honors the `--` positional separator, returning None (→ skip routing) unless it is confident it found the prompt. Co-authored-by: Isaac --- src/ucode/smart_routing/codex_routing.py | 57 +++++++++++++++++++----- src/ucode/smart_routing/routing.py | 54 ++++++++++++++++++++++ tests/test_agent_codex.py | 40 +++++++++++++++++ 3 files changed, 140 insertions(+), 11 deletions(-) diff --git a/src/ucode/smart_routing/codex_routing.py b/src/ucode/smart_routing/codex_routing.py index c529bd7..68aafa4 100644 --- a/src/ucode/smart_routing/codex_routing.py +++ b/src/ucode/smart_routing/codex_routing.py @@ -13,7 +13,6 @@ # ``codex_routing.urllib.request`` — the actual call lives in ``routing``, but # Python modules are singletons so patching this name patches the one call site. import urllib.request # noqa: F401 -from pathlib import Path from typing import Any from ucode.config_io import APP_DIR @@ -43,7 +42,18 @@ def route_launch_model(state: dict, tool_args: list[str]): - """Route a root Codex launch before the Codex process starts.""" + """Route a root Codex launch on the launch-time prompt, if there is one. + + Returns (None, None) when the launch carries no prompt (a bare interactive + session): with no task signal the router can only return its floor arm, so + routing would just add a round-trip and silently override the user's default + model. In that case we don't route and keep the configured default. Routing + on a typed-in first prompt is out of scope — no hook/MCP can retarget the + root model once the session is running. + """ + task = _launch_routing_task(tool_args) + 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): @@ -52,18 +62,43 @@ def route_launch_model(state: dict, tool_args: list[str]): 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}." +# Codex CLI options that consume a following value (from `codex --help`); their +# values must not be mistaken for the seed prompt when parsing launch args. +_CODEX_VALUE_OPTIONS = frozenset( + { + "-c", + "--config", + "-i", + "--image", + "-m", + "--model", + "-p", + "--profile", + "-s", + "--sandbox", + "-a", + "--ask-for-approval", + "-C", + "--cd", + "--add-dir", + "--enable", + "--disable", + "--local-provider", + "--remote", + "--remote-auth-token-env", + } +) + + +def _launch_routing_task(tool_args: list[str]) -> str | None: + # The routing task is the user's real first prompt when it's on the command + # line (`codex ""`, `codex exec ""`, or after `--`). A bare + # interactive launch has no prompt yet → None, and the caller skips routing + # (the root model can't be re-routed once the TUI is running). + return routing.extract_seed_prompt(tool_args, _CODEX_VALUE_OPTIONS) def request_routing_decision( diff --git a/src/ucode/smart_routing/routing.py b/src/ucode/smart_routing/routing.py index f157096..b1c4c08 100644 --- a/src/ucode/smart_routing/routing.py +++ b/src/ucode/smart_routing/routing.py @@ -48,6 +48,60 @@ def normalize_model(model: str) -> str: return tail.lower() +def extract_seed_prompt(tool_args: list[str], value_options: frozenset[str]) -> str | None: + """Recover a launch-time seed prompt from the passthrough CLI args. + + A coding agent launched as `` [OPTIONS] [PROMPT]`` (or with an explicit + ``exec `` subcommand) may carry the user's first prompt on the command + line. When it does, routing the root-session model on that real prompt beats + the generic placeholder. Everything after a ``--`` separator is treated as + positional (the agent's own convention for "stop parsing flags"). + + ``value_options`` is the set of the tool's flags that consume a following + value (e.g. ``-m``/``--model``); their values are skipped so a flag argument + is never mistaken for the prompt. Returns the joined positional tokens, or + None when the args carry no unambiguous prompt (bare launch, or only flags) — + the caller then falls back to its placeholder task. + + Conservative by design: an unrecognized ``--flag`` (not in ``value_options``) + is treated as a boolean and skipped, and ``--flag=value`` forms are skipped + whole. We only return text we are confident is the user's prompt. + """ + if "exec" in tool_args: + after = tool_args[tool_args.index("exec") + 1 :] + positionals = _positional_args(after, value_options) + return " ".join(positionals) if positionals else None + positionals = _positional_args(tool_args, value_options) + return " ".join(positionals) if positionals else None + + +def _positional_args(args: list[str], value_options: frozenset[str]) -> list[str]: + positionals: list[str] = [] + i = 0 + seen_double_dash = False + while i < len(args): + arg = args[i] + if seen_double_dash: + positionals.append(arg) + i += 1 + continue + if arg == "--": + seen_double_dash = True + i += 1 + continue + if arg.startswith("-") and arg != "-": + # `--opt=value` carries its own value; a bare option in value_options + # consumes the next token. Anything else is a boolean flag we skip. + if "=" not in arg and arg in value_options: + i += 2 + else: + i += 1 + continue + positionals.append(arg) + i += 1 + return positionals + + def select_route( workspace: str, token: str, diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index 6c03aaa..7208cdb 100644 --- a/tests/test_agent_codex.py +++ b/tests/test_agent_codex.py @@ -366,6 +366,46 @@ def test_disable_removes_only_ucode_hooks(self, tmp_path, monkeypatch): def test_launch_task_uses_exec_prompt(self): assert codex_routing._launch_routing_task(["exec", "fix the parser"]) == "fix the parser" + def test_launch_task_uses_positional_interactive_prompt(self): + # `codex "fix the parser"` — the seed prompt is routed directly, not + # wrapped in a placeholder. + assert codex_routing._launch_routing_task(["fix the parser"]) == "fix the parser" + + def test_launch_task_skips_value_option_before_prompt(self): + # `-m ` consumes its value; the model id must not be taken as the + # prompt. + assert ( + codex_routing._launch_routing_task(["-m", "gpt-5", "refactor the parser"]) + == "refactor the parser" + ) + + def test_launch_task_honors_double_dash(self): + assert ( + codex_routing._launch_routing_task(["--", "--not-a-flag prompt"]) + == "--not-a-flag prompt" + ) + + def test_launch_task_bare_launch_returns_none(self): + # No prompt on the command line → None, so the caller skips routing and + # keeps the user's default model (root model can't be re-routed once the + # TUI is up). + assert codex_routing._launch_routing_task([]) is None + + def test_launch_task_flags_only_returns_none(self): + assert codex_routing._launch_routing_task(["--search", "-m", "gpt-5"]) is None + + def test_route_launch_model_skips_routing_without_prompt(self, monkeypatch): + # Bare launch: no router call at all, no decision, no error. + def fail(*args, **kwargs): + raise AssertionError("router must not be called on a bare launch") + + monkeypatch.setattr(codex_routing, "request_routing_decision", fail) + decision, error = codex_routing.route_launch_model( + {"workspace": WS, "codex_models": ["system.ai.gpt-5-6-sol"]}, [] + ) + assert decision is None + assert error is None + class TestCodexRemoveLegacyProfile: def test_drops_provider_block_on_modern_path(self, tmp_path, monkeypatch): From dc00331c945cb56d2dd3e1fd4c353c67c8426eee Mon Sep 17 00:00:00 2001 From: Mason Cao Date: Fri, 31 Jul 2026 20:57:12 +0000 Subject: [PATCH 5/5] Surface the router rationale in the shown message, and log it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Both "Using Smart Routing" messages — the launch-time notice and the subagent-routing hook's systemMessage — now include the router's rationale, so the "why" is visible, not just the "what". Previously the launch notice never showed it and the subagent hook put it only in permissionDecisionReason (a field the harness doesn't surface in that line). Centralized in RoutingDecision.display_message() so both paths format identically. Also persist the rationale in the decisions.jsonl record, so an empty value cleanly distinguishes "the gateway returned no rationale" from a display-side placement bug. Co-authored-by: Isaac --- src/ucode/cli.py | 2 +- src/ucode/smart_routing/routing.py | 23 ++++++++++++++--- tests/test_cli.py | 13 ++++++++-- tests/test_codex_routing.py | 40 +++++++++++++++++++++++++++++- 4 files changed, 71 insertions(+), 7 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 3ebf087..60555fe 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -1141,7 +1141,7 @@ def _launch_tool( decision, routing_error = route_launch_model(state, ctx.args) if decision is not None: resolved_model = decision.model - print_note(f"Using Smart Routing. Routing to {resolved_model}.") + print_note(decision.display_message()) elif routing_error: print_warning( f"Smart routing was unavailable ({routing_error}); using {resolved_model}." diff --git a/src/ucode/smart_routing/routing.py b/src/ucode/smart_routing/routing.py index b1c4c08..29ae98a 100644 --- a/src/ucode/smart_routing/routing.py +++ b/src/ucode/smart_routing/routing.py @@ -33,6 +33,19 @@ class RoutingDecision: raw_model: str rationale: str = "" + def display_message(self, model_label: str | None = None) -> str: + """User-facing "Using Smart Routing" line, with the router's rationale. + + Used by both the launch-time notice and the subagent-routing hook so the + "what" (model) and the "why" (rationale) are surfaced consistently. + ``model_label`` overrides the shown model id (e.g. a harness-translated + id); defaults to ``model``. The rationale is appended when present. + """ + message = f"Using Smart Routing. Routing to {model_label or self.model}." + if self.rationale: + message += f" {self.rationale}" + return message + def normalize_model(model: str) -> str: """Strip provider prefixes so router arms and workspace ids compare equal. @@ -199,15 +212,16 @@ def route_spawn_tool( routed_model = model_id_mapper(decision.model) if record_decision is not None: record_decision(payload, task, decision, routed_model) - routing_message = f"Using Smart Routing. Routing to {routed_model}." + # Surface the router's rationale in BOTH the systemMessage (the line the + # harness shows the user) and permissionDecisionReason — the "why", not just + # the "what". The shown model is the harness-translated id (routed_model). + routing_message = decision.display_message(model_label=routed_model) output: dict[str, Any] = { "hookEventName": "PreToolUse", "permissionDecision": "allow", "updatedInput": {**tool_input, "model": routed_model}, "permissionDecisionReason": routing_message, } - if decision.rationale: - output["permissionDecisionReason"] += f" {decision.rationale}" return {"systemMessage": routing_message, "hookSpecificOutput": output} @@ -271,6 +285,9 @@ def write_decision_record( "task_name": task_name, "router_model": decision.raw_model, "requested_model": requested_model, + # Persisted for diagnosis: an empty value means the gateway returned + # no rationale (vs. a placement bug in how we surface it). + "rationale": decision.rationale, "at": time.time(), }, ) diff --git a/tests/test_cli.py b/tests/test_cli.py index cf07a83..0ab437d 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -11,6 +11,7 @@ from typer.testing import CliRunner from ucode.cli import app +from ucode.smart_routing.routing import RoutingDecision _ANSI_RE = re.compile(r"\x1b\[[0-9;]*m") @@ -226,7 +227,11 @@ def test_enabled_codex_launch_uses_routed_root_model(self): "smart_routing_enabled": True, "codex_models": ["databricks-gpt-5", "databricks-gpt-5-5"], } - decision = MagicMock(model="databricks-gpt-5-5") + decision = RoutingDecision( + model="databricks-gpt-5-5", + raw_model="gpt-5-6-sol", + rationale="Cross-cutting refactor.", + ) with ( patch("ucode.cli.ensure_bootstrap_dependencies"), patch("ucode.cli.load_state", return_value=state), @@ -244,7 +249,11 @@ def test_enabled_codex_launch_uses_routed_root_model(self): assert result.exit_code == 0, result.output assert mock_configure.call_args.args[2] == "databricks-gpt-5-5" - assert "Using Smart Routing. Routing to databricks-gpt-5-5." in _strip_ansi(result.output) + # The launch notice surfaces both the routed model and the rationale. + assert ( + "Using Smart Routing. Routing to databricks-gpt-5-5. Cross-cutting refactor." + in _strip_ansi(result.output) + ) class TestMcpSubcommands: diff --git a/tests/test_codex_routing.py b/tests/test_codex_routing.py index bd85041..ea67baa 100644 --- a/tests/test_codex_routing.py +++ b/tests/test_codex_routing.py @@ -143,7 +143,11 @@ def test_spawn_rewrite_preserves_original_input(monkeypatch): ) hook = output["hookSpecificOutput"] - assert output["systemMessage"] == "Using Smart Routing. Routing to gpt-5.5." + # The rationale is surfaced in BOTH the systemMessage (shown to the user) and + # permissionDecisionReason, so the "why" is visible, not just the "what". + assert output["systemMessage"] == ( + "Using Smart Routing. Routing to gpt-5.5. Review needs deeper reasoning." + ) assert hook["permissionDecision"] == "allow" assert hook["updatedInput"] == { "task_name": "reviewer", @@ -276,3 +280,37 @@ def test_decision_is_reconciled_with_actual_subagent_model(tmp_path, monkeypatch assert record["router_model"] == "gpt-5-6-luna" assert record["requested_model"] == "gpt-5.6-luna" assert record["matches_router_decision"] is True + + +def test_decision_record_persists_rationale(tmp_path, monkeypatch): + decisions = tmp_path / "decisions.jsonl" + monkeypatch.setattr(codex_routing, "DECISIONS_PATH", decisions) + monkeypatch.setattr( + codex_routing, + "request_routing_decision", + lambda *args, **kwargs: ( + codex_routing.RoutingDecision( + model="system.ai.gpt-5-6-sol", + raw_model="gpt-5-6-sol", + rationale="Cross-cutting refactor needs the strongest model.", + ), + None, + ), + ) + + codex_routing.route_pre_tool_use( + { + "session_id": "s1", + "tool_name": "collaborationspawn_agent", + "tool_input": {"task_name": "refactor", "message": "encrypted"}, + }, + workspace=WS, + token="token", + available_models=["system.ai.gpt-5-6-sol", "system.ai.gpt-5-6-luna"], + audit_decision=True, + ) + + # The rationale is persisted so an empty value distinguishes "gateway + # returned none" from a display-placement bug. + record = json.loads(decisions.read_text().strip()) + assert record["rationale"] == "Cross-cutting refactor needs the strongest model."