diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 92d18f2b..222e780d 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -118,7 +118,7 @@ def _resolve_web_search_model(state: dict) -> str | None: # Matches both the AI Gateway form (`databricks-claude-opus-4-8`) and the UC # model-services form (`system.ai.claude-opus-4-8`). _CLAUDE_MODEL_RE = re.compile( - r"^(?:system\.ai\.)?(?:databricks-)?claude-(opus|sonnet)-(\d+)-(\d+)(.*)$" + r"^(?:system\.ai\.)?(?:databricks-)?claude-(opus|sonnet)-(\d+)(?:-(\d+))?(.*)$" ) # Env keys the MLflow Stop hook reads to route traces. Written into the @@ -438,7 +438,7 @@ def _maybe_add_1m_suffix(model: str) -> str: family, major_raw, minor_raw, _ = match.groups() major = int(major_raw) - minor = int(minor_raw) + minor = int(minor_raw or 0) should_suffix = (family == "opus" and (major, minor) >= (4, 6)) or ( family == "sonnet" and (major, minor) >= (4, 6) ) @@ -1188,6 +1188,7 @@ def compose_gateway_settings(args: list[str]) -> tuple[dict, list[str]]: launch_model=_original_launch_model(state), compose_settings=compose_gateway_settings, launch_model_args=_launch_model_args, + model_name=_maybe_add_1m_suffix, ) return diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 0e31e163..6a6ae3b0 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -1534,13 +1534,22 @@ def claude_router_hook_cmd( 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 smart_routing_v2.enabled(): + output = smart_routing_v2.route_claude_pre_tool_use( + payload, + workspace=host, + token=token, + available_models=model or [], + audit_decision=True, + ) + else: + 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)) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 3e4d4507..fc7ab3bd 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -52,6 +52,7 @@ "https://raw.githubusercontent.com/databricks/setup-cli/main/install.ps1" ) AI_GATEWAY_V2_DOCS_URL = "https://docs.databricks.com/aws/en/ai-gateway/overview-beta" +ANTHROPIC_MODELS_PATH = "/ai-gateway/anthropic/v1/models" # v1.0.0 is the release that ships `databricks aitools`. MIN_DATABRICKS_CLI_VERSION = (1, 0, 0) TOKEN_REFRESH_INTERVAL_SECONDS = 1800 @@ -2738,7 +2739,7 @@ def list_anthropic_models(workspace: str, token: str) -> tuple[list[str], str | validation. """ hostname = workspace_hostname(workspace) - payload, reason = _http_get_json(f"https://{hostname}/ai-gateway/anthropic/v1/models", token) + payload, reason = _http_get_json(f"https://{hostname}{ANTHROPIC_MODELS_PATH}", token) if payload is None: return [], reason @@ -2765,7 +2766,7 @@ def discover_claude_models(workspace: str, token: str) -> tuple[dict[str, str], matching the expected naming convention). """ hostname = workspace_hostname(workspace) - payload, reason = _http_get_json(f"https://{hostname}/ai-gateway/anthropic/v1/models", token) + payload, reason = _http_get_json(f"https://{hostname}{ANTHROPIC_MODELS_PATH}", token) if payload is None: return {}, reason diff --git a/src/ucode/smart_routing/claude_pty.py b/src/ucode/smart_routing/claude_pty.py index 9b455015..933b7cef 100644 --- a/src/ucode/smart_routing/claude_pty.py +++ b/src/ucode/smart_routing/claude_pty.py @@ -56,18 +56,6 @@ def valid_model_name(name: object) -> bool: ) -def switch_message(model: str, reason: str) -> str: - """Format the routed-model notice shown in Claude Code.""" - lines = [ - "Using Unity Gateway Smart Router.", - f"Selected Model : {model}", - f"Reason : {reason}", - ] - width = max(len(line) for line in lines) - border = "─" * (width + 2) - return "\n".join([f"┌{border}┐", *(f"│ {line:<{width}} │" for line in lines), f"└{border}┘"]) - - class ConfirmationState: """Detect and accept Claude's optional cache-cost confirmation dialog.""" @@ -162,21 +150,24 @@ def request_first_prompt_route(path: Path, payload: dict, *, timeout: float = 5. def first_prompt_hook_output(response: dict | None) -> dict | None: """Translate the wrapper response into Claude hook output.""" + from ucode.smart_routing.v2 import format_routing_notice + if not isinstance(response, dict) or response.get("action") != "block": return None model = response.get("model") if not valid_model_name(model): return None assert isinstance(model, str) + rationale = response.get("rationale") return { "decision": "block", - "reason": switch_message(model, "Low complexity, unclear intent, and no code reference."), + "reason": format_routing_notice(model, rationale if isinstance(rationale, str) else None), } def serve_first_prompt_socket( path: Path, - route_prompt: Callable[[str], str], + route_prompt: Callable[[str], tuple[str, str]], on_blocked_prompt: Callable[[str, str], None], stop: threading.Event, *, @@ -223,10 +214,14 @@ def serve() -> None: and prompt.strip() and not is_command ): - model = route_prompt(prompt) + model, rationale = route_prompt(prompt) if valid_model_name(model): claimed = True - response = {"action": "block", "model": model} + response = { + "action": "block", + "model": model, + "rationale": rationale, + } blocked = (prompt, model) except Exception as exc: # noqa: BLE001 - hooks must fail open log(f"[ERR] first-prompt request: {exc!r}") @@ -274,7 +269,7 @@ def sync_winsize(master_fd: int, stdin_fd: int = 0) -> None: def run_claude_pty( argv: list[str], *, - route_prompt: Callable[[str], str], + route_prompt: Callable[[str], tuple[str, str]], socket_path: Path, prepare_model_switch: Callable[[str], None] = lambda _model: None, model_switch_persisted: Callable[[], bool] = lambda: True, diff --git a/src/ucode/smart_routing/routing.py b/src/ucode/smart_routing/routing.py index 3f4ca2d0..2fecacd5 100644 --- a/src/ucode/smart_routing/routing.py +++ b/src/ucode/smart_routing/routing.py @@ -47,6 +47,14 @@ def display_message(self, model_label: str | None = None) -> str: return message +@dataclass(frozen=True) +class SpawnRoute: + tool_input: dict[str, Any] + task: str + decision: RoutingDecision + routed_model: str + + def normalize_model(model: str) -> str: """Strip provider prefixes so router arms and workspace ids compare equal. @@ -180,23 +188,15 @@ def select_route( ) -def route_spawn_tool( +def resolve_spawn_route( 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. - """ +) -> SpawnRoute | None: + """Resolve one subagent-spawn payload to a routed model.""" if not is_spawn_agent(payload.get("tool_name")): return None tool_input = payload.get("tool_input") @@ -219,19 +219,42 @@ def route_spawn_tool( 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) + return SpawnRoute(tool_input, task, decision, routed_model) + + +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.""" + route = resolve_spawn_route( + payload, + is_spawn_agent=is_spawn_agent, + decision_fn=decision_fn, + default_task_label=default_task_label, + model_id_mapper=model_id_mapper, + ) + if route is None: + return None + if skip_arms and route.decision.raw_model in skip_arms: + return {"systemMessage": skip_arms[route.decision.raw_model]} if record_decision is not None: - record_decision(payload, task, decision, routed_model) + record_decision(payload, route.task, route.decision, route.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) + routing_message = route.decision.display_message(model_label=route.routed_model) output: dict[str, Any] = { "hookEventName": "PreToolUse", "permissionDecision": "allow", - "updatedInput": {**tool_input, "model": routed_model}, + "updatedInput": {**route.tool_input, "model": route.routed_model}, "permissionDecisionReason": routing_message, } return {"systemMessage": routing_message, "hookSpecificOutput": output} diff --git a/src/ucode/smart_routing/v2.py b/src/ucode/smart_routing/v2.py index 66e8cf8e..09d5dcb4 100644 --- a/src/ucode/smart_routing/v2.py +++ b/src/ucode/smart_routing/v2.py @@ -1,5 +1,7 @@ from __future__ import annotations +import hashlib +import json import os import signal import socket @@ -21,14 +23,17 @@ get_databricks_token, list_anthropic_models, ) -from ucode.smart_routing import codex_interposer, routing -from ucode.smart_routing.claude_hooks import FIRST_PROMPT_SOCKET_ENV, sync_first_prompt_hook +from ucode.smart_routing import claude_routing, codex_interposer, routing +from ucode.smart_routing.claude_hooks import ( + FIRST_PROMPT_SOCKET_ENV, + sync_first_prompt_hook, + sync_smart_routing_hooks, +) from ucode.ui import print_note ENV_VAR = "ENABLE_SMART_ROUTING_V2" CODEX_INTERPOSER_LOG = APP_DIR / "codex-v2-interposer.log" -STUBBED_SWITCH_REASON = "Low complexity, unclear intent, and no code reference." # TODO(lilly): replace with smart router rationale. CLAUDE_TARGET_MODEL = "system.ai.claude-sonnet-4-6[1m]" # TODO(lilly): replace with smart router. CLAUDE_PTY_LOG = APP_DIR / "claude-v2-pty.log" @@ -39,6 +44,11 @@ HEALTH_REQUEST_TIMEOUT_SECONDS = 1 HEALTH_POLL_INTERVAL_SECONDS = 0.25 CLAUDE_ROUTE_SELECTION_TIMEOUT_S = 20.0 +CLAUDE_ROUTED_AGENT_PREFIX = "ucode-route-" +CLAUDE_ROUTED_AGENT_PROMPT = ( + "Complete the delegated task exactly as requested. Follow the parent agent's instructions and " + "return a concise report of your findings or changes." +) def enabled() -> bool: @@ -70,13 +80,15 @@ def _wait_for_app_server(port: int, timeout: float) -> bool: return False -def _switch_message(model: str, reason: str) -> str: +def format_routing_notice(model: str, reason: str | None, *, title: str | None = None) -> str: lines = [ + *([title] if title else []), "Using Unity Gateway Smart Router.", f"Selected Model : {model}", - f"Reason : {reason}", ] - width = max(len(line) for line in lines) + if reason: + lines.append(f"Reason : {reason}") + width = max(map(len, lines)) border = "─" * (width + 2) return "\n".join([f"┌{border}┐", *(f"│ {line:<{width}} │" for line in lines), f"└{border}┘"]) @@ -89,21 +101,95 @@ def _canonical_claude_model_id(model: str) -> str: return model -def _route_claude_prompt(state: dict, token: str, prompt: str) -> routing.RoutingDecision: - workspace = state.get("workspace") - if not isinstance(workspace, str): - raise RuntimeError("workspace metadata is unavailable") +def _canonical_claude_models(model_ids: list[str]) -> list[str]: + return list( + dict.fromkeys( + _canonical_claude_model_id(model) + for model in model_ids + if isinstance(model, str) and model + ) + ) - model_ids, discovery_error = list_anthropic_models(workspace, token) - if not model_ids: - raise RuntimeError(discovery_error or "Anthropic models endpoint returned no Claude models") +def _claude_model_overrides(model_ids: list[str]) -> dict[str, str]: + overrides: dict[str, str] = {} + prefix = "system.ai." + for model in _canonical_claude_models(model_ids): + if model.startswith(f"{prefix}claude-"): + overrides[model[len(prefix) :]] = model + return overrides + + +def _routed_claude_agent_name(model: str) -> str: + canonical = _canonical_claude_model_id(model) + normalized = routing.normalize_model(canonical) + safe = "".join(character if character.isalnum() else "-" for character in normalized) + slug = "-".join(part for part in safe.split("-") if part) + digest = hashlib.sha256(canonical.encode()).hexdigest()[:8] + return f"{CLAUDE_ROUTED_AGENT_PREFIX}{slug[:36]}-{digest}" + + +def _routed_claude_agent_definitions(model_ids: list[str]) -> dict[str, dict[str, str]]: + return { + _routed_claude_agent_name(model): { + "description": f"Smart-routed coding agent using {model}", + "prompt": CLAUDE_ROUTED_AGENT_PROMPT, + "model": model, + } + for model in _canonical_claude_models(model_ids) + } + + +def _with_routed_claude_agents(tool_args: list[str], model_ids: list[str]) -> list[str]: + definitions = _routed_claude_agent_definitions(model_ids) + caller_definitions: dict = {} + remaining: list[str] = [] + index = 0 + while index < len(tool_args): + arg = tool_args[index] + if arg == "--": + remaining.extend(tool_args[index:]) + break + if arg == "--agents": + if index + 1 >= len(tool_args): + raise RuntimeError("Claude's --agents option requires a JSON object.") + raw = tool_args[index + 1] + index += 2 + elif arg.startswith("--agents="): + raw = arg.partition("=")[2] + index += 1 + else: + remaining.append(arg) + index += 1 + continue + try: + parsed = json.loads(raw) + except ValueError as exc: + raise RuntimeError("Claude's --agents option must contain valid JSON.") from exc + if not isinstance(parsed, dict): + raise RuntimeError("Claude's --agents option must contain a JSON object.") + caller_definitions.update(parsed) + + collisions = definitions.keys() & caller_definitions.keys() + if collisions: + names = ", ".join(sorted(collisions)) + raise RuntimeError(f"Claude --agents names conflict with smart routing: {names}.") + combined = {**caller_definitions, **definitions} + return ["--agents", json.dumps(combined, separators=(",", ":")), *remaining] + + +def _request_claude_routing_decision( + workspace: str, + token: str, + prompt: str, + model_ids: list[str], +) -> tuple[routing.RoutingDecision | None, str | None]: available: dict[str, str] = {} - for model in model_ids: - if isinstance(model, str) and model: - canonical = _canonical_claude_model_id(model) - available.setdefault(routing.normalize_model(canonical), canonical) - decision, error = routing.select_route( + for model in _canonical_claude_models(model_ids): + available.setdefault(routing.normalize_model(model), model) + if not available: + return None, "Anthropic models endpoint returned no Claude models" + return routing.select_route( workspace, token, prompt, @@ -111,11 +197,76 @@ def _route_claude_prompt(state: dict, token: str, prompt: str) -> routing.Routin lambda selected: available.get(routing.normalize_model(selected)), timeout=CLAUDE_ROUTE_SELECTION_TIMEOUT_S, ) + + +def _route_claude_prompt( + state: dict, + token: str, + prompt: str, + model_ids: list[str] | None = None, +) -> routing.RoutingDecision: + workspace = state.get("workspace") + if not isinstance(workspace, str): + raise RuntimeError("workspace metadata is unavailable") + + if model_ids is None: + model_ids, discovery_error = list_anthropic_models(workspace, token) + if not model_ids: + raise RuntimeError( + discovery_error or "Anthropic models endpoint returned no Claude models" + ) + decision, error = _request_claude_routing_decision(workspace, token, prompt, model_ids) if decision is None: raise RuntimeError(error or "router returned no Claude model selection") return decision +def route_claude_pre_tool_use( + payload: dict, + *, + workspace: str, + token: str, + available_models: list[str], + audit_decision: bool = False, +) -> dict | None: + """Route a Claude Agent call through a transient exact-model agent definition.""" + route = routing.resolve_spawn_route( + payload, + is_spawn_agent=claude_routing.is_spawn_agent_tool, + decision_fn=lambda task: _request_claude_routing_decision( + workspace, token, task, available_models + ), + default_task_label="Claude Code subagent task", + model_id_mapper=lambda model: model, + ) + if route is None: + return None + if audit_decision: + routing.write_decision_record( + claude_routing.DECISIONS_PATH, + payload, + route.task, + route.decision, + route.routed_model, + ) + routing_message = format_routing_notice( + route.routed_model, + route.decision.rationale, + title="Subagent Smart Routing", + ) + updated_input = { + **{key: value for key, value in route.tool_input.items() if key != "model"}, + "subagent_type": _routed_claude_agent_name(route.routed_model), + } + hook_output = { + "hookEventName": "PreToolUse", + "permissionDecision": "allow", + "updatedInput": updated_input, + "permissionDecisionReason": routing_message, + } + return {"systemMessage": routing_message, "hookSpecificOutput": hook_output} + + def _is_claude_target_model(value: object) -> bool: if not isinstance(value, str): return False @@ -172,6 +323,7 @@ def launch_claude( launch_model: str | None, compose_settings: Callable[[list[str]], tuple[dict, list[str]]], launch_model_args: Callable[[list[str], str | None], list[str]], + model_name: Callable[[str], str], ) -> NoReturn: """Launch Claude in the first-prompt routing PTY wrapper.""" from ucode.agents.claude import GATEWAY_MODEL_DISCOVERY_ENV_VAR @@ -185,6 +337,9 @@ def launch_claude( token = get_databricks_token(workspace, state.get("profile")) os.environ[OAUTH_TOKEN_ENV_VAR] = token os.environ[GATEWAY_MODEL_DISCOVERY_ENV_VAR] = "1" + model_ids, discovery_error = list_anthropic_models(workspace, token) + if not model_ids: + raise RuntimeError(discovery_error or "Anthropic models endpoint returned no Claude models") run_id = f"{os.getpid()}-{uuid.uuid4().hex[:8]}" socket_path = APP_DIR / f"claude-v2-{run_id}.sock" @@ -199,13 +354,27 @@ def launch_claude( raise RuntimeError("Claude settings 'env' must be an object for smart routing.") env["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] = "1" env[FIRST_PROMPT_SOCKET_ENV] = str(socket_path) + model_overrides = settings.setdefault("modelOverrides", {}) + if not isinstance(model_overrides, dict): + raise RuntimeError("Claude settings 'modelOverrides' must be an object for smart routing.") + model_overrides.update(_claude_model_overrides(model_ids)) + routing_state = { + **state, + "claude_models": {str(index): model for index, model in enumerate(model_ids)}, + } + sync_smart_routing_hooks(settings, routing_state, enabled=True) sync_first_prompt_hook(settings, hook_executable) write_json_file(settings_path, settings) model_args = launch_model_args(remaining, launch_model) - argv = [binary, "--settings", str(settings_path), *model_args, *remaining] + routed_agent_args = _with_routed_claude_agents(remaining, model_ids) + argv = [binary, "--settings", str(settings_path), *model_args, *routed_agent_args] model_setting = _ClaudeModelSettingGuard(user_settings_path) + def route_prompt(prompt: str) -> tuple[str, str]: + decision = _route_claude_prompt(state, token, prompt, model_ids) + return model_name(decision.model), decision.rationale + print_note( "Smart routing v2: the first submitted prompt will select Claude Code's " f"model; log: {CLAUDE_PTY_LOG}." @@ -213,7 +382,7 @@ def launch_claude( try: returncode = claude_pty.run_claude_pty( argv, - route_prompt=lambda prompt: _route_claude_prompt(state, token, prompt).model, + route_prompt=route_prompt, socket_path=socket_path, prepare_model_switch=model_setting.begin, model_switch_persisted=model_setting.is_routed, @@ -321,7 +490,7 @@ def launch_codex( available_models=available_models, workspace=workspace, token_provider=lambda: get_databricks_token(workspace, profile), - switch_message_fn=_switch_message, + switch_message_fn=format_routing_notice, log_path=CODEX_INTERPOSER_LOG, ) tui_url = _loopback_websocket_url(tui_port) diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 832c1f51..dc99ab21 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -27,6 +27,11 @@ def test_display(self): class TestRenderOverlay: + def test_long_context_suffix_supports_major_only_claude_versions(self): + assert claude._maybe_add_1m_suffix("system.ai.claude-sonnet-5") == ( + "system.ai.claude-sonnet-5[1m]" + ) + def test_does_not_set_anthropic_model_env(self): # We deliberately don't pin ANTHROPIC_MODEL: when set, Claude Code's # /model picker surfaces a duplicate catalog row on top of the family diff --git a/tests/test_claude_smart_routing_v2.py b/tests/test_claude_smart_routing_v2.py index 9fa8398b..b5b4e803 100644 --- a/tests/test_claude_smart_routing_v2.py +++ b/tests/test_claude_smart_routing_v2.py @@ -11,7 +11,7 @@ import pytest from ucode.agents import claude -from ucode.smart_routing import claude_hooks, claude_pty, v2 +from ucode.smart_routing import claude_hooks, claude_pty, routing, v2 class TestDirectModelCommand: @@ -30,11 +30,22 @@ def test_rejects_unsafe_model_names(self, name): class TestFirstPromptHook: def test_renders_boxed_router_notice(self): model = "system.ai.claude-sonnet-4-6[1m]" - reason = "Low complexity, unclear intent, and no code reference." - result = claude_pty.first_prompt_hook_output({"action": "block", "model": model}) + reason = "Routed to Sonnet because the task is narrowly scoped." + result = claude_pty.first_prompt_hook_output( + {"action": "block", "model": model, "rationale": reason} + ) + + assert result == { + "decision": "block", + "reason": v2.format_routing_notice(model, reason), + } + + def test_omits_reason_when_router_returns_none(self): + result = claude_pty.first_prompt_hook_output( + {"action": "block", "model": "system.ai.claude-sonnet-5"} + ) - assert result == {"decision": "block", "reason": v2._switch_message(model, reason)} - assert claude_pty.switch_message(model, reason) == v2._switch_message(model, reason) + assert "Reason" not in result["reason"] def test_blocks_once_then_allows_replay(self, tmp_path): socket_path = tmp_path / "first.sock" @@ -42,7 +53,7 @@ def test_blocks_once_then_allows_replay(self, tmp_path): stop = threading.Event() claude_pty.serve_first_prompt_socket( socket_path, - lambda _prompt: "sonnet", + lambda _prompt: ("sonnet", "Selected for a narrow task."), lambda prompt, model: blocked.append((prompt, model)), stop, ) @@ -56,7 +67,11 @@ def test_blocks_once_then_allows_replay(self, tmp_path): replay = claude_pty.request_first_prompt_route( socket_path, {"session_id": "s1", "prompt": "fix the parser"} ) - assert first == {"action": "block", "model": "sonnet"} + assert first == { + "action": "block", + "model": "sonnet", + "rationale": "Selected for a narrow task.", + } assert replay == {"action": "allow"} assert blocked == [("fix the parser", "sonnet")] finally: @@ -85,18 +100,29 @@ def test_restores_model_captured_immediately_before_switch(self, tmp_path, monke monkeypatch.setattr(v2, "CLAUDE_PTY_LOG", tmp_path / "v2.log") monkeypatch.setattr(v2, "get_databricks_token", lambda *_args, **_kwargs: "token") monkeypatch.setattr(v2, "build_auth_token_argv", lambda *_args, **_kwargs: ["ucode"]) + monkeypatch.setattr( + v2, + "list_anthropic_models", + lambda *_args: ( + ["system.ai.claude-opus-4-8", "system.ai.claude-sonnet-5"], + None, + ), + ) monkeypatch.setattr( v2, "_route_claude_prompt", lambda *_args: v2.routing.RoutingDecision( model="system.ai.claude-sonnet-5", raw_model="claude-sonnet-5", + rationale="Selected for the parser task.", ), ) captured: dict = {} def fake_run(argv, **kwargs): captured["argv"] = argv + agents_index = argv.index("--agents") + captured["agents"] = json.loads(argv[agents_index + 1]) captured["routed_model"] = kwargs["route_prompt"]("fix the parser") generated = Path(argv[argv.index("--settings") + 1]) captured["settings"] = json.loads(generated.read_text()) @@ -123,12 +149,38 @@ def fake_run(argv, **kwargs): launch_model="opus", compose_settings=claude._compose_v2_settings, launch_model_args=claude._launch_model_args, + model_name=claude._maybe_add_1m_suffix, ) assert exc.value.code == 0 - assert captured["argv"][-3:] == ["--model", "opus", "--debug"] - assert captured["routed_model"] == "system.ai.claude-sonnet-5" + assert captured["argv"][3:5] == ["--model", "opus"] + assert captured["argv"][-1] == "--debug" + assert captured["routed_model"] == ( + "system.ai.claude-sonnet-5[1m]", + "Selected for the parser task.", + ) + assert {definition["model"] for definition in captured["agents"].values()} == { + "system.ai.claude-opus-4-8", + "system.ai.claude-sonnet-5", + } + assert captured["settings"]["modelOverrides"] == { + "claude-opus-4-8": "system.ai.claude-opus-4-8", + "claude-sonnet-5": "system.ai.claude-sonnet-5", + } assert claude_hooks.FIRST_PROMPT_SOCKET_ENV in captured["settings"]["env"] + first_prompt_command = captured["settings"]["hooks"]["UserPromptSubmit"][0]["hooks"][0][ + "command" + ] + assert first_prompt_command == "ucode claude-router-hook route-first-prompt" + route_commands = [ + hook["command"] + for group in captured["settings"]["hooks"]["PreToolUse"] + for hook in group["hooks"] + if "route-subagent" in hook["command"] + ] + assert len(route_commands) == 1 + assert "--model system.ai.claude-opus-4-8" in route_commands[0] + assert "--model system.ai.claude-sonnet-5" in route_commands[0] assert "modelPicker" not in captured["settings"] assert captured["restored_during_run"] == { "model": "haiku", @@ -147,6 +199,7 @@ def test_does_not_restore_when_wrapper_never_switches(self, tmp_path, monkeypatc monkeypatch.setattr(v2, "APP_DIR", tmp_path) monkeypatch.setattr(v2, "get_databricks_token", lambda *_args, **_kwargs: "token") monkeypatch.setattr(v2, "build_auth_token_argv", lambda *_args, **_kwargs: ["ucode"]) + monkeypatch.setattr(v2, "list_anthropic_models", lambda *_args: (["opus"], None)) def fake_run(_argv, **_kwargs): user_settings.write_text(json.dumps({"model": "user-selected"})) @@ -162,6 +215,7 @@ def fake_run(_argv, **_kwargs): launch_model="opus", compose_settings=lambda _args: ({}, []), launch_model_args=claude._launch_model_args, + model_name=claude._maybe_add_1m_suffix, ) assert json.loads(user_settings.read_text()) == {"model": "user-selected"} @@ -174,6 +228,7 @@ def test_restores_after_routed_model_persists_and_preserves_later_choice( monkeypatch.setattr(v2, "APP_DIR", tmp_path) monkeypatch.setattr(v2, "get_databricks_token", lambda *_args, **_kwargs: "token") monkeypatch.setattr(v2, "build_auth_token_argv", lambda *_args, **_kwargs: ["ucode"]) + monkeypatch.setattr(v2, "list_anthropic_models", lambda *_args: (["opus"], None)) def fake_run(_argv, **kwargs): routed_model = "system.ai.claude-opus-4-8" @@ -196,6 +251,7 @@ def fake_run(_argv, **kwargs): launch_model=None, compose_settings=lambda _args: ({}, []), launch_model_args=claude._launch_model_args, + model_name=claude._maybe_add_1m_suffix, ) assert json.loads(user_settings.read_text()) == { @@ -203,6 +259,108 @@ def fake_run(_argv, **kwargs): "theme": "dark", } + +class TestSubagentRouting: + def test_routes_agent_prompt_with_initialized_model_menu(self, tmp_path, monkeypatch): + captured = {} + decisions_path = tmp_path / "decisions.jsonl" + monkeypatch.setattr(v2.claude_routing, "DECISIONS_PATH", decisions_path) + + def fake_select(workspace, token, task, route_options, resolve, **kwargs): + captured.update( + workspace=workspace, + token=token, + task=task, + route_options=list(route_options), + ) + return ( + routing.RoutingDecision( + model=resolve("claude-opus-4-8"), + raw_model="claude-opus-4-8", + ), + None, + ) + + monkeypatch.setattr(routing, "select_route", fake_select) + output = v2.route_claude_pre_tool_use( + { + "tool_name": "Agent", + "tool_input": {"prompt": "inspect the parser", "model": "sonnet"}, + }, + workspace="https://example.com", + token="token", + available_models=[ + "system.ai.claude-opus-4-8", + "databricks-claude-sonnet-5", + ], + audit_decision=True, + ) + + assert captured == { + "workspace": "https://example.com", + "token": "token", + "task": "inspect the parser", + "route_options": [ + ("claude-opus-4-8", "claude"), + ("claude-sonnet-5", "claude"), + ], + } + updated_input = output["hookSpecificOutput"]["updatedInput"] + assert "model" not in updated_input + assert updated_input["subagent_type"] == v2._routed_claude_agent_name( + "system.ai.claude-opus-4-8" + ) + expected_message = v2.format_routing_notice( + "system.ai.claude-opus-4-8", + "", + title="Subagent Smart Routing", + ) + assert output["systemMessage"] == expected_message + assert output["hookSpecificOutput"]["permissionDecisionReason"] == expected_message + decision_record = json.loads(decisions_path.read_text()) + assert decision_record["requested_model"] == "system.ai.claude-opus-4-8" + + def test_merges_caller_agents_with_transient_routed_agents(self): + args = v2._with_routed_claude_agents( + [ + "--agents", + json.dumps( + { + "reviewer": { + "description": "Reviews code", + "prompt": "Review the requested code.", + } + } + ), + "--debug", + ], + ["databricks-claude-opus-4-8"], + ) + + assert args[0] == "--agents" + definitions = json.loads(args[1]) + assert definitions["reviewer"]["prompt"] == "Review the requested code." + routed = definitions[v2._routed_claude_agent_name("system.ai.claude-opus-4-8")] + assert routed["model"] == "system.ai.claude-opus-4-8" + assert args[2:] == ["--debug"] + + def test_leaves_non_claude_custom_agent_model_unchanged(self): + definitions = v2._routed_claude_agent_definitions(["catalog.schema.gpt-5"]) + + assert next(iter(definitions.values()))["model"] == "catalog.schema.gpt-5" + + def test_maps_gateway_claude_ids_to_known_model_metadata(self): + assert v2._claude_model_overrides( + [ + "system.ai.claude-opus-4-8", + "databricks-claude-sonnet-5", + "catalog.schema.gpt-5", + ] + ) == { + "claude-opus-4-8": "system.ai.claude-opus-4-8", + "claude-sonnet-5": "system.ai.claude-sonnet-5", + } + def test_model_switch_lock_serializes_routed_sessions(self, tmp_path, monkeypatch): user_settings = tmp_path / "settings.json" user_settings.write_text(json.dumps({"model": "haiku"})) @@ -250,7 +408,7 @@ def is_alive(): with pytest.raises(RuntimeError, match="Claude was not launched"): claude_pty.run_claude_pty( ["claude"], - route_prompt=lambda _prompt: "sonnet", + route_prompt=lambda _prompt: ("sonnet", ""), socket_path=tmp_path / "missing.sock", ) @@ -308,7 +466,7 @@ def read_until(suffix): str(capture), str(restored), ], - route_prompt=lambda _prompt: "system.ai.claude-sonnet-5", + route_prompt=lambda _prompt: ("system.ai.claude-sonnet-5", ""), socket_path=socket_path, restore_model_setting=lambda: restored.write_text("restored"), ) diff --git a/tests/test_cli.py b/tests/test_cli.py index 57ca6642..a4f90a80 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -3,6 +3,7 @@ from __future__ import annotations import contextlib +import json import os import re from unittest.mock import MagicMock, patch @@ -342,6 +343,41 @@ def test_claude_v2_first_prompt_hook_is_disabled_without_flag(self, monkeypatch) assert result.output == "" mock_request.assert_not_called() + def test_claude_v2_subagent_hook_uses_v2_router(self, monkeypatch): + monkeypatch.setenv("ENABLE_SMART_ROUTING_V2", "1") + routed = { + "hookSpecificOutput": { + "hookEventName": "PreToolUse", + "permissionDecision": "allow", + "updatedInput": {"prompt": "fix it", "model": "opus"}, + } + } + with ( + patch( + "ucode.cli.smart_routing_v2.route_claude_pre_tool_use", + return_value=routed, + ) as mock_v2_route, + patch("ucode.cli.claude_routing.route_pre_tool_use") as mock_legacy_route, + ): + result = runner.invoke( + app, + [ + "claude-router-hook", + "route-subagent", + "--host", + "https://example.com", + "--model", + "system.ai.claude-opus-4-8", + ], + input='{"tool_name":"Agent","tool_input":{"prompt":"fix it"}}', + env={"OAUTH_TOKEN": "token"}, + ) + + assert result.exit_code == 0, result.output + assert json.loads(result.output) == routed + mock_v2_route.assert_called_once() + mock_legacy_route.assert_not_called() + class TestClaudeModelFlag: """`ucode claude --model ` pins the id into the family aliases so the gateway resolves any diff --git a/tests/test_codex_smart_routing_v2.py b/tests/test_codex_smart_routing_v2.py index 3eb8c863..3e4e7ef7 100644 --- a/tests/test_codex_smart_routing_v2.py +++ b/tests/test_codex_smart_routing_v2.py @@ -36,7 +36,7 @@ def test_layers_provider_overrides_without_replacing_user_config(self, monkeypat def test_smart_routing_switch_message_is_boxed(): - message = v2._switch_message("model-x", "Because X.") + message = v2.format_routing_notice("model-x", "Because X.") assert message == ( "┌───────────────────────────────────┐\n" @@ -169,7 +169,7 @@ def start_interposer(*args, **kwargs): assert token_calls == [(WS, "myprof")] assert interposer_args["kwargs"]["token_provider"]() == "token-2" assert token_calls == [(WS, "myprof"), (WS, "myprof")] - assert interposer_args["kwargs"]["switch_message_fn"] is v2._switch_message + assert interposer_args["kwargs"]["switch_message_fn"] is v2.format_routing_notice assert stopped == [True] assert processes[0].terminated is True @@ -299,7 +299,7 @@ def select(prompt): log=lambda _m: None, available_models=["claude-opus-4-8", "gpt-5.5"], route_decision=select, - switch_message_fn=v2._switch_message, + switch_message_fn=v2.format_routing_notice, ) output = sess.on_tui_frame(self._turn_start("gpt-5.5", prompt="Fix issue #42"))