diff --git a/README.md b/README.md index b856249..a4b2e1c 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 @@ -156,6 +171,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..a25439d 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -21,6 +21,10 @@ get_databricks_token, ) from ucode.launcher import exec_or_spawn +from ucode.smart_routing.codex_hooks import ( + remove_smart_routing_hooks, + sync_smart_routing_hooks, +) from ucode.state import mark_tool_managed, save_state from ucode.telemetry import agent_version, ucode_version @@ -33,6 +37,11 @@ CODEX_MODEL_PROVIDER_NAME = "ucode-databricks" MINIMUM_CODEX_VERSION = (0, 134, 0) MINIMUM_CODEX_VERSION_TEXT = "0.134.0" +MINIMUM_ROUTING_CODEX_VERSION = (0, 145, 0) +MINIMUM_ROUTING_CODEX_VERSION_TEXT = "0.145.0" +# Shared across agents: one opt-in enables smart routing for every routing-capable +# tool (codex, claude), so a workspace turns it on once. +SMART_ROUTING_STATE_KEY = "smart_routing_enabled" SPEC: ToolSpec = { @@ -318,6 +327,10 @@ def write_tool_config(state: dict, model: str | None = None, provider: str | Non databricks_profile = state.get("profile") if _use_legacy_layout(): + if smart_routing_enabled(state) and provider is None: + raise RuntimeError( + f"Codex smart routing requires Codex {MINIMUM_ROUTING_CODEX_VERSION_TEXT} or newer." + ) # Codex < 0.134.0 only reads ~/.codex/config.toml. Write the shared # config with [profiles.ucode] + shared [model_providers.ucode-databricks] # and skip the per-profile-file cleanup that would normally strip @@ -358,6 +371,11 @@ def write_tool_config(state: dict, model: str | None = None, provider: str | Non # deep_merge can't drop keys, so clear a `model` pinned by an earlier # non-provider run that the provider overlay omits. doc.pop("model", None) + sync_smart_routing_hooks( + doc, + state, + enabled=smart_routing_enabled(state) and provider is None, + ) write_toml_file(CODEX_CONFIG_PATH, doc) state = mark_tool_managed(state, "codex", MANAGED_KEYS) save_state(state) @@ -399,6 +417,43 @@ def launch(state: dict, tool_args: list[str]) -> None: exec_or_spawn([binary, "--profile", CODEX_PROFILE_NAME, *tool_args]) +def smart_routing_enabled(state: dict) -> bool: + """Return whether the current workspace opted into Codex routing.""" + return state.get(SMART_ROUTING_STATE_KEY) is True + + +def enable_smart_routing(state: dict) -> dict: + """Persist the current workspace's Codex smart-routing opt-in.""" + parsed = _parse_version(agent_version(SPEC["binary"])) + if parsed is not None and parsed < MINIMUM_ROUTING_CODEX_VERSION: + raise RuntimeError( + "Codex smart routing requires Codex " + f"{MINIMUM_ROUTING_CODEX_VERSION_TEXT} or newer; found " + f"{agent_version(SPEC['binary'])}." + ) + state[SMART_ROUTING_STATE_KEY] = True + return state + + +def disable_smart_routing(state: dict) -> bool: + """Disable routing and remove only ucode's Codex routing hooks.""" + state.pop(SMART_ROUTING_STATE_KEY, None) + if state.get("workspace"): + save_state(state) + changed = False + for path in (CODEX_CONFIG_PATH, LEGACY_CODEX_CONFIG_PATH): + if not path.exists(): + continue + doc = read_toml_safe(path) + if remove_smart_routing_hooks(doc): + write_toml_file(path, doc) + changed = True + from ucode.smart_routing.codex_routing import clear_routing_artifacts + + clear_routing_artifacts() + return changed + + def validate_cmd(binary: str) -> list[str]: return [ binary, diff --git a/src/ucode/cli.py b/src/ucode/cli.py index e5bab4e..17a5ede 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,12 @@ from ucode.agents import ( launch as launch_agent, ) -from ucode.agents.codex import revert_legacy_shared_config +from ucode.agents.codex import ( + disable_smart_routing, + enable_smart_routing, + revert_legacy_shared_config, + smart_routing_enabled, +) from ucode.agents.pi import PI_SETTINGS_BACKUP_PATH, PI_SETTINGS_PATH from ucode.config_io import restore_file, set_dry_run from ucode.databricks import ( @@ -61,6 +67,7 @@ revert_mcp_configs, ) from ucode.skills_download import configure_skills_download_command +from ucode.smart_routing.codex_routing import route_launch_model from ucode.state import ( STATE_PATH, clear_state, @@ -952,6 +959,77 @@ def auth_token_cmd( sys.stdout.write(token + "\n") +@app.command("codex-router-hook", hidden=True) +def codex_router_hook_cmd( + event: str, + host: Annotated[str | None, typer.Option("--host")] = None, + profile: Annotated[str | None, typer.Option("--profile")] = None, + use_pat: Annotated[bool, typer.Option("--use-pat")] = False, + model: Annotated[list[str] | None, typer.Option("--model")] = None, +) -> None: + """Run a Codex smart-routing lifecycle hook.""" + import json + import sys + + from ucode.smart_routing.codex_routing import ( + record_session_start, + record_subagent_start, + route_pre_tool_use, + ) + + try: + payload = json.loads(sys.stdin.read() or "{}") + except ValueError: + return + if not isinstance(payload, dict): + return + if event == "session-start": + record_session_start(payload) + return + if event == "record-subagent": + record = record_subagent_start(payload) + matched = record.get("matches_router_decision") + if matched is True: + sys.stdout.write( + json.dumps( + { + "systemMessage": "Smart Routing verified. " + f"Subagent is using {record.get('model')}." + } + ) + ) + elif matched is False: + sys.stdout.write( + json.dumps( + { + "systemMessage": "Smart Routing mismatch: router requested " + f"{record.get('requested_model')}, but Codex started " + f"{record.get('model')}." + } + ) + ) + return + if event != "route-subagent" or not host: + return + token = os.environ.get("OAUTH_TOKEN") or os.environ.get("DATABRICKS_BEARER") + if not token: + if use_pat and not ensure_pat_bearer(profile): + return + try: + token = get_databricks_token(host, profile) + except RuntimeError: + return + output = route_pre_tool_use( + payload, + workspace=host, + token=token, + available_models=model or [], + audit_decision=True, + ) + if output is not None: + sys.stdout.write(json.dumps(output)) + + def _auto_configure_tool(tool: str) -> None: """First-time setup for a single tool — mirrors configure_workspace_command.""" existing = load_state() @@ -995,6 +1073,7 @@ def _launch_tool( provider: str | None = None, skip_preflight: bool = False, workspace: str | None = None, + enable_codex_smart_routing: bool = False, ) -> None: try: tool = normalize_tool(tool_name) @@ -1018,6 +1097,11 @@ def _launch_tool( # An explicit --provider overrides the persisted choice; otherwise fall # back to whatever `ucode configure` saved for this tool. provider = provider or get_provider_service(state, tool) + if tool == "codex" and enable_codex_smart_routing and provider: + raise RuntimeError( + "Codex smart routing cannot be enabled with --provider. " + "Launch without a Model Provider Service and try again." + ) # Validate the provider service before launching — it must exist, be a # provider type this tool can route to (e.g. claude can't use an OpenAI # or Foundry service), and, for Bedrock, expose Claude models to pin. @@ -1041,6 +1125,8 @@ def _launch_tool( skip_model_discovery=bool(provider), skip_preflight=skip_preflight, ) + if tool == "codex" and enable_codex_smart_routing: + state = enable_smart_routing(state) if provider: # Routing through a Model Provider Service pins no Databricks model; # the agent uses its own canonical model names (header selects the @@ -1049,6 +1135,16 @@ def _launch_tool( resolved_model = None else: state, resolved_model = resolve_launch_model(tool, state, None) + if tool == "codex" and smart_routing_enabled(state): + with spinner("Selecting a Codex model with smart routing..."): + decision, routing_error = route_launch_model(state, ctx.args) + if decision is not None: + resolved_model = decision.model + print_note(decision.display_message()) + elif routing_error: + print_warning( + f"Smart routing was unavailable ({routing_error}); using {resolved_model}." + ) state = configure_tool( tool, state, @@ -1062,6 +1158,13 @@ def _launch_tool( print_kv("Provider", provider) elif resolved_model: print_kv("Model", resolved_model) + if tool == "codex" and smart_routing_enabled(state) and not provider: + print_kv("Smart routing", "enabled") + if enable_codex_smart_routing: + print_note( + "Codex requires one-time hook review. Open `/hooks` and trust the " + "ucode routing hooks if prompted." + ) if tool in ("gemini", "opencode", "copilot", "pi"): print_note( f"{TOOL_SPECS[tool]['display']} token refresh is managed automatically " @@ -1116,10 +1219,36 @@ def codex_cmd( ] = None, skip_preflight: SkipPreflightOption = False, workspace: WorkspaceOption = None, + enable_smart_routing_flag: Annotated[ + bool, + typer.Option( + "--enable-smart-routing", + help="Enable AI Gateway model routing for Codex sessions and subagents.", + ), + ] = False, + disable_smart_routing_flag: Annotated[ + bool, + typer.Option( + "--disable-smart-routing", + help="Disable smart routing and remove ucode's Codex routing hooks.", + ), + ] = False, ) -> None: """Launch Codex via Databricks.""" + if enable_smart_routing_flag and disable_smart_routing_flag: + print_err("Use only one of --enable-smart-routing or --disable-smart-routing.") + raise typer.Exit(1) + if disable_smart_routing_flag: + disable_smart_routing(load_state()) + print_success("Codex smart routing disabled; ucode routing hooks removed") + return _launch_tool( - "codex", ctx, provider=provider, skip_preflight=skip_preflight, workspace=workspace + "codex", + ctx, + provider=provider, + skip_preflight=skip_preflight, + workspace=workspace, + enable_codex_smart_routing=enable_smart_routing_flag, ) 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/smart_routing/codex_hooks.py b/src/ucode/smart_routing/codex_hooks.py new file mode 100644 index 0000000..6c6a569 --- /dev/null +++ b/src/ucode/smart_routing/codex_hooks.py @@ -0,0 +1,82 @@ +"""Codex hook configuration for smart subagent routing.""" + +from __future__ import annotations + +import shlex +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_smart_routing_hooks(doc: dict, state: dict, *, enabled: bool) -> None: + """Synchronize ucode-managed routing hooks in a Codex config document.""" + groups = _routing_hook_groups(state) if enabled else {} + hooks.sync_managed_hooks(doc, ROUTING_HOOK_COMMAND_MARKER, groups) + + +def remove_smart_routing_hooks(doc: dict) -> bool: + """Remove only ucode-managed smart-routing hooks.""" + return hooks.remove_managed_hooks(doc, ROUTING_HOOK_COMMAND_MARKER) + + +def _routing_hook_groups(state: dict) -> dict[str, list[dict]]: + 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/smart_routing/codex_routing.py b/src/ucode/smart_routing/codex_routing.py new file mode 100644 index 0000000..68aafa4 --- /dev/null +++ b/src/ucode/smart_routing/codex_routing.py @@ -0,0 +1,233 @@ +"""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 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 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): + 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}" + return request_routing_decision(workspace, token, task, models) + + +# 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( + 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..29ae98a --- /dev/null +++ b/src/ucode/smart_routing/routing.py @@ -0,0 +1,368 @@ +"""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 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. + + ``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 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, + 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) + # 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, + } + 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, + # 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(), + }, + ) + + +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 9a68e03..7208cdb 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.smart_routing import codex_routing WS = "https://example.databricks.com" @@ -217,6 +218,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_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( + "[[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.SMART_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.SMART_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 +319,94 @@ def test_unknown_version_uses_modern_layout(self, monkeypatch): assert codex._use_legacy_layout() is False +class TestCodexSmartRouting: + def test_enable_requires_supported_codex(self, monkeypatch): + monkeypatch.setattr(codex, "agent_version", lambda binary: "0.144.0") + + try: + codex.enable_smart_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(codex_routing, "clear_routing_artifacts", lambda: None) + state = {"workspace": WS, codex.SMART_ROUTING_STATE_KEY: True} + + assert codex.disable_smart_routing(state) is True + + doc = read_toml_safe(config_path) + 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" + + 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): config_dir = tmp_path / ".codex" diff --git a/tests/test_cli.py b/tests/test_cli.py index cb1d128..56677e2 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") @@ -190,6 +191,70 @@ 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_smart_routing_is_consumed_by_ucode(self): + with patch("ucode.cli._launch_tool") as mock_launch: + result = runner.invoke(app, ["codex", "--enable-smart-routing"]) + + assert result.exit_code == 0, result.output + 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_smart_routing") as mock_disable, + patch("ucode.cli._launch_tool") as mock_launch, + ): + result = runner.invoke(app, ["codex", "--disable-smart-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-smart-routing", "--disable-smart-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, + "smart_routing_enabled": True, + "codex_models": ["databricks-gpt-5", "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), + 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" + # 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: 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..ea67baa --- /dev/null +++ b/tests/test_codex_routing.py @@ -0,0 +1,316 @@ +"""Tests for Codex smart-routing hooks.""" + +from __future__ import annotations + +import json +import urllib.error + +from ucode.smart_routing 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"] + # 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", + "message": encrypted_message, + "fork": False, + "model": "gpt-5.5", + } + assert hook["permissionDecisionReason"] == ( + "Using Smart 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 Smart 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": ( + "Smart 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 + + +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."