diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index 27b0458..3afd4fe 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -312,6 +312,7 @@ def configure_tool( provider: str | None = None, provider_models: dict[str, str] | None = None, relayed: bool = False, + route_root_model: str | None = None, ) -> dict: result: dict | tuple[dict, str] if tool == "codex": @@ -322,7 +323,12 @@ def configure_tool( if not model and not provider: raise RuntimeError(f"A {tool} model must be selected before configuration.") result = claude.write_tool_config( - state, model, provider=provider, provider_models=provider_models, relayed=relayed + state, + model, + provider=provider, + provider_models=provider_models, + relayed=relayed, + route_root_model=route_root_model, ) else: # provider routing is claude/codex-only; every other tool needs a model. diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 194210a..1747263 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -29,6 +29,10 @@ get_databricks_token, ) from ucode.launcher import exec_or_spawn +from ucode.smart_routing.claude_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 from ucode.tracing import tracing_env @@ -46,6 +50,15 @@ "backup_path": CLAUDE_BACKUP_PATH, } +# Per-workspace opt-in flag for Claude Code smart routing (state key). +# Shared across agents: one opt-in enables smart routing for every routing-capable +# tool (codex, claude), so a workspace turns it on once. Kept identical to +# codex.SMART_ROUTING_STATE_KEY on purpose. +SMART_ROUTING_STATE_KEY = "smart_routing_enabled" +# Claude Code settings.json hook events ucode manages when routing is enabled; +# marked managed so they're tracked/reverted with the rest of ucode's config. +CLAUDE_ROUTING_HOOK_EVENTS = ("PreToolUse", "SessionStart", "SubagentStart") + def is_update_available() -> tuple[str, str] | None: return available_npm_package_update(SPEC["package"]) @@ -161,6 +174,7 @@ def render_overlay( fable_enabled: bool = False, relayed: bool = False, relayed_base_url: str | None = None, + route_root_model: str | None = None, ) -> tuple[dict, list[list[str]]]: """Return (overlay, managed_key_paths) for Claude settings.json. @@ -212,13 +226,20 @@ def render_overlay( "ENABLE_TOOL_SEARCH": "1", "CLAUDE_CODE_USE_GATEWAY": "1", } - # Intentionally NOT setting ANTHROPIC_MODEL. Setting it produces a duplicate - # catalog row in Claude Code's /model picker (e.g. "Opus 4.8 (1M context) ✓") - # on top of the family-alias row from ANTHROPIC_DEFAULT_OPUS_MODEL. Without - # it, Default resolves through the pinned family alias and the picker shows - # only one row per model. `ucode claude -- --model X` still overrides for a - # single session via Claude Code's own --model flag. + # Intentionally NOT setting ANTHROPIC_MODEL by default. Setting it produces a + # duplicate catalog row in Claude Code's /model picker (e.g. "Opus 4.8 (1M + # context) ✓") on top of the family-alias row from ANTHROPIC_DEFAULT_OPUS_MODEL. + # Without it, Default resolves through the pinned family alias and the picker + # shows only one row per model. `ucode claude -- --model X` still overrides for + # a single session via Claude Code's own --model flag. + # + # The one exception is smart routing: `route_root_model` pins the + # router's per-launch pick for the root session as ANTHROPIC_MODEL. The + # duplicate-picker-row cost is acceptable because the whole point is to launch + # on the routed model rather than the family default. _ = model # API stability; no longer pinned via env. + if route_root_model: + env["ANTHROPIC_MODEL"] = route_root_model # A Bedrock-backed provider needs its provider-side ids pinned verbatim # (Claude Code's canonical names aren't routable there). These come from the # service's targets, already de-duped to one id per family upstream. @@ -336,12 +357,41 @@ def _unregister_web_search_mcp() -> None: pass +def smart_routing_enabled(state: dict) -> bool: + """Return whether the current workspace opted into Claude Code routing.""" + return state.get(SMART_ROUTING_STATE_KEY) is True + + +def enable_smart_routing(state: dict) -> dict: + """Persist the current workspace's Claude Code smart-routing opt-in.""" + state[SMART_ROUTING_STATE_KEY] = True + return state + + +def disable_smart_routing(state: dict) -> bool: + """Disable routing and remove only ucode's Claude Code routing hooks.""" + state.pop(SMART_ROUTING_STATE_KEY, None) + if state.get("workspace"): + save_state(state) + changed = False + if CLAUDE_SETTINGS_PATH.exists(): + doc = read_json_safe(CLAUDE_SETTINGS_PATH) + if remove_smart_routing_hooks(doc): + write_json_file(CLAUDE_SETTINGS_PATH, doc) + changed = True + from ucode.smart_routing.claude_routing import clear_routing_artifacts + + clear_routing_artifacts() + return changed + + def write_tool_config( state: dict, model: str | None, provider: str | None = None, provider_models: dict[str, str] | None = None, relayed: bool = False, + route_root_model: str | None = None, ) -> dict: backup_existing_file(CLAUDE_SETTINGS_PATH, CLAUDE_BACKUP_PATH) web_search_model = _resolve_web_search_model(state) @@ -360,6 +410,7 @@ def write_tool_config( fable_enabled=bool(state.get("fable_enabled")), relayed=relayed, relayed_base_url=relayed_base_url, + route_root_model=route_root_model, ) tracing_env_vars = tracing_env(state, "claude") stop_hook_command = claude_tracing_stop_hook_command() if tracing_env_vars else None @@ -404,6 +455,13 @@ def write_tool_config( if isinstance(merged_env, dict): for key in CLAUDE_REMOVED_ENV_KEYS: merged_env.pop(key, None) + # Smart-routing hooks: install ucode's PreToolUse/SessionStart/ + # SubagentStart hooks when routing is enabled (and not under a provider, + # which pins no Databricks model), else surgically strip only ucode's own. + routing_enabled = smart_routing_enabled(state) and provider is None + sync_smart_routing_hooks(merged, state, enabled=routing_enabled) + if routing_enabled: + managed_keys = managed_keys + [["hooks", event] for event in CLAUDE_ROUTING_HOOK_EVENTS] write_json_file(CLAUDE_SETTINGS_PATH, merged) if web_search_model: diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 17a5ede..acd1610 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -26,15 +26,12 @@ validate_all_tools, validate_tool, ) +from ucode.agents import claude as claude_agent +from ucode.agents import codex as codex_agent from ucode.agents import ( launch as launch_agent, ) -from ucode.agents.codex import ( - disable_smart_routing, - enable_smart_routing, - revert_legacy_shared_config, - smart_routing_enabled, -) +from ucode.agents.codex import revert_legacy_shared_config 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 ( @@ -67,7 +64,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.smart_routing import claude_routing, codex_routing from ucode.state import ( STATE_PATH, clear_state, @@ -1030,6 +1027,77 @@ def codex_router_hook_cmd( sys.stdout.write(json.dumps(output)) +@app.command("claude-router-hook", hidden=True) +def claude_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 Claude Code smart-routing lifecycle hook.""" + import json + import sys + + from ucode.smart_routing.claude_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 Claude Code 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() @@ -1067,13 +1135,20 @@ def _auto_configure_tool(tool: str) -> None: raise RuntimeError(f"{spec['display']} validation failed — config reverted.") +# Agent modules exposing the smart-routing opt-in surface (enable/disable/ +# enabled + launch-model routing), keyed by tool. Both share the identical +# function names via their agent module and their routing module. +_ROUTING_AGENTS = {"codex": codex_agent, "claude": claude_agent} +_ROUTING_MODULES = {"codex": codex_routing, "claude": claude_routing} + + def _launch_tool( tool_name: str, ctx: typer.Context, provider: str | None = None, skip_preflight: bool = False, workspace: str | None = None, - enable_codex_smart_routing: bool = False, + enable_smart_routing_flag: bool = False, ) -> None: try: tool = normalize_tool(tool_name) @@ -1097,10 +1172,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: + routing_agent = _ROUTING_AGENTS.get(tool) + if routing_agent is not None and enable_smart_routing_flag and provider: raise RuntimeError( - "Codex smart routing cannot be enabled with --provider. " - "Launch without a Model Provider Service and try again." + f"{TOOL_SPECS[tool]['display']} 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 @@ -1125,8 +1201,11 @@ 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 routing_agent is not None and enable_smart_routing_flag: + state = routing_agent.enable_smart_routing(state) + # The router's per-launch pick for the root session. Codex pins it as the + # resolved model; claude pins it via ANTHROPIC_MODEL (route_root_model). + route_root_model = None if provider: # Routing through a Model Provider Service pins no Databricks model; # the agent uses its own canonical model names (header selects the @@ -1135,12 +1214,18 @@ 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 routing_agent is not None and routing_agent.smart_routing_enabled(state): + display = TOOL_SPECS[tool]["display"] + with spinner(f"Selecting a {display} model with smart routing..."): + decision, routing_error = _ROUTING_MODULES[tool].route_launch_model( + state, ctx.args + ) if decision is not None: - resolved_model = decision.model print_note(decision.display_message()) + if tool == "codex": + resolved_model = decision.model + else: + route_root_model = decision.model elif routing_error: print_warning( f"Smart routing was unavailable ({routing_error}); using {resolved_model}." @@ -1152,18 +1237,25 @@ def _launch_tool( provider=provider, provider_models=provider_models, relayed=relayed, + route_root_model=route_root_model, ) print_section(f"ucode with {TOOL_SPECS[tool]['display']}") if provider: print_kv("Provider", provider) + elif route_root_model: + print_kv("Model", route_root_model) elif resolved_model: print_kv("Model", resolved_model) - if tool == "codex" and smart_routing_enabled(state) and not provider: + if ( + routing_agent is not None + and routing_agent.smart_routing_enabled(state) + and not provider + ): print_kv("Smart routing", "enabled") - if enable_codex_smart_routing: + if enable_smart_routing_flag: print_note( - "Codex requires one-time hook review. Open `/hooks` and trust the " - "ucode routing hooks if prompted." + f"{TOOL_SPECS[tool]['display']} requires one-time hook review. Open " + "`/hooks` and trust the ucode routing hooks if prompted." ) if tool in ("gemini", "opencode", "copilot", "pi"): print_note( @@ -1239,7 +1331,7 @@ def codex_cmd( 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()) + codex_agent.disable_smart_routing(load_state()) print_success("Codex smart routing disabled; ucode routing hooks removed") return _launch_tool( @@ -1248,7 +1340,7 @@ def codex_cmd( provider=provider, skip_preflight=skip_preflight, workspace=workspace, - enable_codex_smart_routing=enable_smart_routing_flag, + enable_smart_routing_flag=enable_smart_routing_flag, ) @@ -1266,10 +1358,36 @@ def claude_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 Claude Code sessions and subagents.", + ), + ] = False, + disable_smart_routing_flag: Annotated[ + bool, + typer.Option( + "--disable-smart-routing", + help="Disable smart routing and remove ucode's Claude Code routing hooks.", + ), + ] = False, ) -> None: """Launch Claude Code 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: + claude_agent.disable_smart_routing(load_state()) + print_success("Claude Code smart routing disabled; ucode routing hooks removed") + return _launch_tool( - "claude", ctx, provider=provider, skip_preflight=skip_preflight, workspace=workspace + "claude", + ctx, + provider=provider, + skip_preflight=skip_preflight, + workspace=workspace, + enable_smart_routing_flag=enable_smart_routing_flag, ) diff --git a/src/ucode/smart_routing/claude_hooks.py b/src/ucode/smart_routing/claude_hooks.py new file mode 100644 index 0000000..773b770 --- /dev/null +++ b/src/ucode/smart_routing/claude_hooks.py @@ -0,0 +1,91 @@ +"""Claude Code hook configuration for smart subagent routing. + +Written into ``~/.claude/ucode-settings.json`` under Claude Code's hook events. +The ``PreToolUse`` hook matches the subagent-spawn tool (``Agent``, formerly +``Task``) and rewrites its ``model`` input to the router's pick; ``SessionStart`` +and ``SubagentStart`` drive the canary/audit trail. Mirrors ``codex_hooks`` but +emits Claude Code's settings.json hook shape. +""" + +from __future__ import annotations + +import shlex + +from ucode.databricks import build_auth_token_argv +from ucode.smart_routing import hooks + +ROUTING_HOOK_COMMAND_MARKER = "claude-router-hook" + + +def sync_smart_routing_hooks(doc: dict, state: dict, *, enabled: bool) -> None: + """Synchronize ucode-managed routing hooks in a Claude settings 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|Task", + "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") + # The route-subagent hook resolves the router's chosen arm back to a routable + # workspace id, so it needs the discovered claude model ids. + claude_models = state.get("claude_models") + if isinstance(claude_models, dict): + for model in claude_models.values(): + 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), + "timeout": 35, + } + if status: + hook["statusMessage"] = status + return hook diff --git a/src/ucode/smart_routing/claude_routing.py b/src/ucode/smart_routing/claude_routing.py new file mode 100644 index 0000000..005eb28 --- /dev/null +++ b/src/ucode/smart_routing/claude_routing.py @@ -0,0 +1,194 @@ +"""Databricks AI Gateway routing helpers for Claude Code sessions and subagents. + +Claude-specific configuration on top of the shared +:mod:`ucode.smart_routing.routing` core. The ``task_v1`` router infers a +"Claude Code" (``cc``) scenario when it is offered only Claude arms, and that +scenario REQUIRES its full menu — both ``claude-opus-4-8`` and +``claude-sonnet-5`` — or it returns BAD_REQUEST. So both arms are always offered. +""" + +from __future__ import annotations + +# Re-exported so tests can patch the shared ``urlopen`` seam via +# ``claude_routing.urllib.request`` — the 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 +REQUEST_TIMEOUT_S = routing.REQUEST_TIMEOUT_S +# Frozen task_v1 "cc" scenario menu (ai-gateway CanonicalModelNames): the router +# rejects the request unless BOTH are offered as route_options. +CLAUDE_ROUTE_ARMS = ("claude-opus-4-8", "claude-sonnet-5") +# Claude Code's subagent-spawn tool (renamed Task -> Agent); match both. +SPAWN_AGENT_TOOL_NAMES = ("agent", "task") +CANARY_PATH = APP_DIR / "claude-smart-routing-canary.json" +AUDIT_PATH = APP_DIR / "claude-smart-routing-audit.jsonl" +DECISIONS_PATH = APP_DIR / "claude-smart-routing-decisions.jsonl" + +_normalize_model = routing.normalize_model + + +def route_launch_model(state: dict, tool_args: list[str]): + """Route a root Claude Code 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("claude_models") + if not isinstance(workspace, str) or not isinstance(models, dict): + 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}" + available = [m for m in models.values() if isinstance(m, str) and m] + return request_routing_decision(workspace, token, task, available) + + +# Claude Code CLI options that consume a following value (from `claude --help`); +# their values must not be mistaken for the seed prompt. Options whose value is +# optional (`-c`/`-d`/`-r`/`-w`) are intentionally omitted — treating them as +# booleans is safe here (we only skip the flag, never a following positional). +_CLAUDE_VALUE_OPTIONS = frozenset( + { + "-n", + "--name", + "-m", + "--model", + "--add-dir", + "--agent", + "--agents", + "--allowed-tools", + "--disallowed-tools", + "--tools", + "--append-system-prompt", + "--system-prompt", + "--betas", + "--debug-file", + "--effort", + "--fallback-model", + "--file", + "--input-format", + "--output-format", + "--json-schema", + "--max-budget-usd", + "--mcp-config", + "--permission-mode", + "--plugin-dir", + "--plugin-url", + "--remote-control-session-name-prefix", + "--session-id", + "--setting-sources", + "--settings", + } +) + + +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 (`claude ""` or `claude -p ""`, 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 session is running). + return routing.extract_seed_prompt(tool_args, _CLAUDE_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 Claude model. + + Offers the full ``cc`` menu; resolves the router's pick back to the + workspace's routable id (e.g. ``system.ai.claude-opus-4-8``). + """ + available = {_normalize_model(m): m for m in available_models if isinstance(m, str) and m} + missing = [arm for arm in CLAUDE_ROUTE_ARMS if arm not in available] + if missing: + return None, f"required Claude routing models are unavailable: {', '.join(missing)}" + + return routing.select_route( + workspace, + token, + task, + [(arm, "claude") for arm in CLAUDE_ROUTE_ARMS], + lambda raw_model: available.get(_normalize_model(raw_model)), + timeout=timeout, + ) + + +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 Claude Code ``Agent`` (subagent-spawn) call, rewriting 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="Claude Code subagent task", + model_id_mapper=_claude_model_id, + record_decision=record, + ) + + +def is_spawn_agent_tool(tool_name: Any) -> bool: + """Return whether a hook payload names Claude Code's subagent spawn tool.""" + if not isinstance(tool_name, str): + return False + return tool_name.strip().lower() in SPAWN_AGENT_TOOL_NAMES + + +def record_session_start(payload: dict[str, Any]) -> None: + """Write a canary proving Claude Code 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 Claude Code 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 _claude_model_id(model: str) -> str: + """The model id Claude Code should launch the subagent with. + + The router returns a routable workspace id (e.g. + ``system.ai.claude-opus-4-8``); Claude Code's ``Agent`` tool ``model`` field + accepts that id verbatim through the gateway, so pass it through unchanged. + """ + return model diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 60b4a6a..e30c85a 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -8,6 +8,7 @@ import pytest from ucode.agents import claude +from ucode.smart_routing import claude_routing WS = "https://example.databricks.com" @@ -760,3 +761,106 @@ def test_malformed_file_json_raises(self, tmp_path, monkeypatch): monkeypatch.setattr(claude, "read_json_safe", lambda p: {"apiKeyHelper": "u"}) with pytest.raises(RuntimeError, match="not valid JSON"): claude._build_claude_argv("claude", ["--settings", str(bad_file)]) + + +class TestClaudeSmartRouting: + def _capture_write(self, monkeypatch, existing, written): + monkeypatch.setattr(claude, "backup_existing_file", lambda *a, **kw: True) + monkeypatch.setattr(claude, "read_json_safe", lambda path: existing) + monkeypatch.setattr( + claude, "write_json_file", lambda path, payload: written.append(payload) + ) + monkeypatch.setattr(claude, "save_state", lambda state: None) + monkeypatch.setattr(claude, "_register_web_search_mcp", lambda *a, **kw: True) + + def test_enable_sets_state_key(self): + state = claude.enable_smart_routing({}) + assert state[claude.SMART_ROUTING_STATE_KEY] is True + + def test_enabled_reads_state_key(self): + assert claude.smart_routing_enabled({claude.SMART_ROUTING_STATE_KEY: True}) + assert not claude.smart_routing_enabled({}) + + def test_write_config_installs_routing_hooks(self, monkeypatch): + written: list = [] + self._capture_write(monkeypatch, {}, written) + state = { + "workspace": WS, + "claude_models": {"opus": "system.ai.claude-opus-4-8"}, + claude.SMART_ROUTING_STATE_KEY: True, + } + claude.write_tool_config(state, "system.ai.claude-opus-4-8", route_root_model=None) + hooks = written[0]["hooks"] + assert set(hooks) == {"PreToolUse", "SessionStart", "SubagentStart"} + commands = [hook["command"] for group in hooks["PreToolUse"] for hook in group["hooks"]] + assert any("claude-router-hook" in command for command in commands) + + def test_root_model_pins_anthropic_model(self, monkeypatch): + written: list = [] + self._capture_write(monkeypatch, {}, written) + state = { + "workspace": WS, + "claude_models": {"opus": "system.ai.claude-opus-4-8"}, + claude.SMART_ROUTING_STATE_KEY: True, + } + claude.write_tool_config( + state, "system.ai.claude-opus-4-8", route_root_model="system.ai.claude-sonnet-5" + ) + assert written[0]["env"]["ANTHROPIC_MODEL"] == "system.ai.claude-sonnet-5" + + def test_provider_suppresses_routing_hooks(self, monkeypatch): + written: list = [] + self._capture_write(monkeypatch, {}, written) + state = {"workspace": WS, claude.SMART_ROUTING_STATE_KEY: True} + # Under a Model Provider Service no Databricks model is pinned, so routing + # is inapplicable — hooks must not be installed even when the flag is set. + claude.write_tool_config(state, None, provider="cat.sch.svc") + assert "hooks" not in written[0] or "PreToolUse" not in written[0]["hooks"] + + def test_disable_removes_only_ucode_hooks(self, tmp_path, monkeypatch): + settings_path = tmp_path / "ucode-settings.json" + settings_path.write_text( + json.dumps( + { + "hooks": { + "PreToolUse": [ + { + "matcher": "Bash", + "hooks": [{"type": "command", "command": "user-policy"}], + }, + { + "matcher": "Agent|Task", + "hooks": [ + { + "type": "command", + "command": "ucode claude-router-hook route-subagent", + } + ], + }, + ], + "SessionStart": [ + { + "hooks": [ + { + "type": "command", + "command": "ucode claude-router-hook session-start", + } + ] + } + ], + } + } + ), + encoding="utf-8", + ) + monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", settings_path) + monkeypatch.setattr(claude, "save_state", lambda state: None) + monkeypatch.setattr(claude_routing, "clear_routing_artifacts", lambda: None) + state = {"workspace": WS, claude.SMART_ROUTING_STATE_KEY: True} + + assert claude.disable_smart_routing(state) is True + + doc = json.loads(settings_path.read_text()) + assert state.get(claude.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_claude_routing.py b/tests/test_claude_routing.py new file mode 100644 index 0000000..49f77af --- /dev/null +++ b/tests/test_claude_routing.py @@ -0,0 +1,301 @@ +"""Tests for Claude Code smart-routing helpers.""" + +from __future__ import annotations + +import json +import urllib.error + +from ucode.smart_routing import claude_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_claude_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": "claude-opus-4-8", "harness": "claude"}} + ], + "rationale": "Cross-cutting change needs the strongest model.", + } + ) + + monkeypatch.setattr(claude_routing.urllib.request, "urlopen", fake_urlopen) + + decision, error = claude_routing.request_routing_decision( + WS, + "token", + "Refactor the parser", + ["system.ai.claude-opus-4-8", "system.ai.claude-sonnet-5"], + ) + + assert error is None + assert decision == claude_routing.RoutingDecision( + model="system.ai.claude-opus-4-8", + raw_model="claude-opus-4-8", + rationale="Cross-cutting change needs the strongest model.", + ) + assert captured["url"] == f"{WS}/ai-gateway/routing/v1/routes:select" + assert captured["headers"]["Authorization"] == "Bearer token" + # The cc-scenario menu requires BOTH Claude arms be offered. + assert captured["body"] == { + "route_options": [ + {"model": "claude-opus-4-8", "harness": "claude"}, + {"model": "claude-sonnet-5", "harness": "claude"}, + ], + "task": {"prompt": "Refactor the parser"}, + "route_selector": {"router_name": "task_v1"}, + } + + +def test_missing_arm_short_circuits_without_calling_router(monkeypatch): + def fail(*args, **kwargs): + raise AssertionError("router must not be called when an arm is missing") + + monkeypatch.setattr(claude_routing.urllib.request, "urlopen", fail) + + # Only sonnet available; the cc menu also requires opus, so we bail early + # rather than send an under-offered request the router would reject. + decision, error = claude_routing.request_routing_decision( + WS, "token", "task", ["system.ai.claude-sonnet-5"] + ) + + assert decision is None + assert "claude-opus-4-8" in error + + +def test_router_pick_resolves_to_workspace_id(monkeypatch): + monkeypatch.setattr( + claude_routing.urllib.request, + "urlopen", + lambda request, timeout: _Response( + {"route_selection": [{"route_option": {"model": "claude-sonnet-5"}}]} + ), + ) + + decision, error = claude_routing.request_routing_decision( + WS, + "token", + "small fix", + ["databricks-claude-opus-4-8", "databricks-claude-sonnet-5"], + ) + + assert error is None + # The bare router arm maps back to the workspace's routable id. + assert decision.model == "databricks-claude-sonnet-5" + + +def test_router_failure_fails_open(monkeypatch): + monkeypatch.setattr( + claude_routing.urllib.request, + "urlopen", + lambda *args, **kwargs: (_ for _ in ()).throw(urllib.error.URLError("offline")), + ) + + decision, error = claude_routing.request_routing_decision( + WS, + "token", + "task", + ["system.ai.claude-opus-4-8", "system.ai.claude-sonnet-5"], + ) + + assert decision is None + assert "offline" in str(error) + + +def test_spawn_rewrite_injects_routed_model(monkeypatch): + payload = { + "tool_name": "Agent", + "tool_input": { + "subagent_type": "Explore", + "prompt": "map the codebase", + "description": "explore", + }, + } + monkeypatch.setattr( + claude_routing, + "request_routing_decision", + lambda *args, **kwargs: ( + claude_routing.RoutingDecision( + model="system.ai.claude-opus-4-8", + raw_model="claude-opus-4-8", + rationale="Deep exploration needs the strongest model.", + ), + None, + ), + ) + + output = claude_routing.route_pre_tool_use( + payload, + workspace=WS, + token="token", + available_models=["system.ai.claude-opus-4-8", "system.ai.claude-sonnet-5"], + ) + + hook = output["hookSpecificOutput"] + # The rationale is surfaced in the systemMessage (shown to the user), not + # only in permissionDecisionReason. + assert output["systemMessage"] == ( + "Using Smart Routing. Routing to system.ai.claude-opus-4-8. " + "Deep exploration needs the strongest model." + ) + assert hook["permissionDecision"] == "allow" + assert hook["updatedInput"] == { + "subagent_type": "Explore", + "prompt": "map the codebase", + "description": "explore", + "model": "system.ai.claude-opus-4-8", + } + assert hook["permissionDecisionReason"] == ( + "Using Smart Routing. Routing to system.ai.claude-opus-4-8. " + "Deep exploration needs the strongest model." + ) + + +def test_task_tool_alias_is_routed(monkeypatch): + monkeypatch.setattr( + claude_routing, + "request_routing_decision", + lambda *args, **kwargs: ( + claude_routing.RoutingDecision( + model="system.ai.claude-sonnet-5", raw_model="claude-sonnet-5" + ), + None, + ), + ) + + output = claude_routing.route_pre_tool_use( + {"tool_name": "Task", "tool_input": {"subagent_type": "general", "prompt": "x"}}, + workspace=WS, + token="token", + available_models=["system.ai.claude-opus-4-8", "system.ai.claude-sonnet-5"], + ) + + assert output["hookSpecificOutput"]["updatedInput"]["model"] == "system.ai.claude-sonnet-5" + + +def test_non_spawn_tool_has_no_opinion(): + assert ( + claude_routing.route_pre_tool_use( + {"tool_name": "Bash", "tool_input": {"command": "true"}}, + workspace=WS, + token="token", + available_models=["system.ai.claude-opus-4-8", "system.ai.claude-sonnet-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(claude_routing, "CANARY_PATH", canary) + monkeypatch.setattr(claude_routing, "AUDIT_PATH", audit) + + claude_routing.record_session_start({"session_id": "s1", "model": "system.ai.claude-opus-4-8"}) + claude_routing.record_subagent_start( + { + "session_id": "s1", + "agent_id": "a1", + "agent_type": "Explore", + "model": "system.ai.claude-sonnet-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(claude_routing, "DECISIONS_PATH", decisions) + monkeypatch.setattr(claude_routing, "AUDIT_PATH", audit) + monkeypatch.setattr( + claude_routing, + "request_routing_decision", + lambda *args, **kwargs: ( + claude_routing.RoutingDecision( + model="system.ai.claude-opus-4-8", raw_model="claude-opus-4-8" + ), + None, + ), + ) + + claude_routing.route_pre_tool_use( + { + "session_id": "s1", + "tool_name": "Agent", + "tool_input": {"subagent_type": "Explore", "prompt": "x"}, + }, + workspace=WS, + token="token", + available_models=["system.ai.claude-opus-4-8", "system.ai.claude-sonnet-5"], + audit_decision=True, + ) + record = claude_routing.record_subagent_start( + {"session_id": "s1", "agent_id": "a1", "model": "system.ai.claude-opus-4-8"} + ) + + assert record["router_model"] == "claude-opus-4-8" + assert record["requested_model"] == "system.ai.claude-opus-4-8" + assert record["matches_router_decision"] is True + + +def test_launch_task_uses_positional_prompt(): + assert claude_routing._launch_routing_task(["fix the parser"]) == "fix the parser" + + +def test_launch_task_skips_value_option_before_prompt(): + assert ( + claude_routing._launch_routing_task(["--model", "claude-opus-4-8", "refactor it"]) + == "refactor it" + ) + + +def test_launch_task_honors_double_dash(): + assert claude_routing._launch_routing_task(["--", "--literal prompt"]) == "--literal prompt" + + +def test_launch_task_bare_launch_returns_none(): + # No prompt on the command line → None, so the caller skips routing and + # keeps the user's default model. + assert claude_routing._launch_routing_task([]) is None + + +def test_launch_task_flags_only_returns_none(): + assert claude_routing._launch_routing_task(["--model", "claude-sonnet-5", "-p"]) is None + + +def test_route_launch_model_skips_routing_without_prompt(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(claude_routing, "request_routing_decision", fail) + decision, error = claude_routing.route_launch_model( + {"workspace": WS, "claude_models": {"opus": "system.ai.claude-opus-4-8"}}, [] + ) + assert decision is None + assert error is None diff --git a/tests/test_cli.py b/tests/test_cli.py index 56677e2..33720f4 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -196,13 +196,13 @@ def test_codex_enable_smart_routing_is_consumed_by_ucode(self): 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.kwargs["enable_smart_routing_flag"] 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.codex_agent.disable_smart_routing") as mock_disable, patch("ucode.cli._launch_tool") as mock_launch, ): result = runner.invoke(app, ["codex", "--disable-smart-routing"]) @@ -241,7 +241,10 @@ def test_enabled_codex_launch_uses_routed_root_model(self): "ucode.cli.resolve_launch_model", return_value=(state, "databricks-gpt-5"), ), - patch("ucode.cli.route_launch_model", return_value=(decision, None)), + patch( + "ucode.cli.codex_routing.route_launch_model", + return_value=(decision, None), + ), patch("ucode.cli.configure_tool", return_value=state) as mock_configure, patch("ucode.cli.launch_agent"), ):