smart routing: route Claude subagents

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
 

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))
 

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
 

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,

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}