Render subagent routing using claude code mod

docs/claude-mod.md

@@ -11,6 +11,18 @@ input with the permission mode) always ends with `smart router ON` or `smart rou
 routing launches use this path, and the existing temporary-directory lifecycle
 cleans up the files when Claude exits.
 
+The same entry point routes Agent/Task calls in a native `tool.call` hook. It pipes the
+call as JSON to `ug claude-router-hook route-subagent-mod`, whose argument list arrives in
+the `routeCommand` plugin option, and continues the call with the returned input. The
+matching Agent row shows a spinner and elapsed seconds while the router answers, then the
+selected model, routing time, and rationale. Results live in memory by `tool_use_id`, so a
+plugin reload drops them. Disabled routing, command failures, and invalid responses keep
+the original call; failures also log a diagnostic.
+
+The old routing PreToolUse hook is no longer written, so a call is not routed twice. A
+`route-subagent` hook that an earlier UG version persisted still gets its PreToolUse
+response and plain-text notice.
+
 `constants.ts` holds the values the mod shares with UG's CLI and the gateway: the recipe
 body field, the disabled recipe, the CLI module, agent, and routing flags, and the
 routing keys used to read session overrides. It ships
@@ -41,8 +53,8 @@ routing hooks apply over the environment, so an "off" saved earlier by the CLI o
 cannot outlast the command. The mod writes no file itself: it runs the existing
 `--enable-smart-routing`/`--disable-smart-routing` entry point, which validates and replaces
 that file atomically. Session start restores only the launch baseline, so an override saved
-before `/clear` still survives it. The plugin test host provides no process capability, so
-this path is covered by the live skill-toggle journey rather than `claude plugin test`.
+before `/clear` still survives it. The live skill-toggle journey covers the real CLI write;
+native component tests stub the process calls.
 
 The existing Claude skill-toggle integration journey checks recorded inference
 payloads during its three calculation phases: launch recipe, `DISABLED` after off,

src/ucode/agents/claude_mods/register.ts

@@ -1,5 +1,6 @@
 import type { EngineInterface, On, PluginOptions } from 'claude-code';
 import { registerStatusBand, showStatusBand } from './smart-routing-status.ts';
+import { registerSubagentRouting } from './subagent-routing.ts';
 import {
   CLI_AGENT,
   CLI_MODULE,
@@ -109,4 +110,5 @@ export function register(on: On, options: PluginOptions): void {
     }
   });
   registerStatusBand(on);
+  registerSubagentRouting(on, options);
 }

src/ucode/agents/claude_mods/subagent-routing.ts

@@ -0,0 +1,118 @@
+import type { On, PluginOptions } from 'claude-code'
+
+type RoutingRecord = { model: string; rationale: string; seconds?: number }
+
+const shortModel = (model: string): string => model.replace(/^system\.ai\./, '').replace(/^databricks-/, '')
+// Claude's own spinner glyphs, played forwards then back as Claude does.
+const SPINNER_GLYPHS = ['·', '✢', '✳', '✶', '✻', '✽']
+const SPINNER = [...SPINNER_GLYPHS, ...SPINNER_GLYPHS.slice(1, -1).reverse()]
+const SPINNER_INTERVAL_MS = 100
+
+// The timer is cosmetic: without a clock it is omitted rather than blocking routing.
+const now = async ($: any): Promise<number | undefined> => {
+  try {
+    return await $.clock.now()
+  } catch {
+    return undefined
+  }
+}
+const secondsSince = async ($: any, start: number | undefined): Promise<number | undefined> => {
+  const current = start === undefined ? undefined : await now($)
+  return current === undefined || start === undefined ? undefined : Math.max(0, current - start) / 1000
+}
+
+export const registerSubagentRouting = (on: On, options: PluginOptions): void => {
+  const records = new Map<string, RoutingRecord>()
+  // Calls waiting on the router, by when routing started, shown with a spinner and timer.
+  const pending = new Map<string, number | undefined>()
+
+  on('tool.call', { tool: ['Agent', 'Task'] }, async ($, e, next) => {
+    const command = options.routeCommand as readonly string[]
+    if (!command.length) return next(e)
+    const { tool, tool_use_id, agentId, ...input } = e
+    let route
+    let spinner: { cancel: () => void } | undefined
+    if (tool_use_id) {
+      pending.set(tool_use_id, await now($))
+      // Mods have no global timers; $.clock redraws the spinner while the router answers.
+      // The frame comes from the clock, so parallel calls do not speed it up.
+      try {
+        spinner = $.clock.every(SPINNER_INTERVAL_MS, () => $.ui.invalidate('ui.render'))
+      } catch {}
+      $.ui.invalidate('ui.render')
+    }
+    try {
+      const session_id = await $.session.id()
+      const result = await $.process.run(command, {
+        stdin: JSON.stringify({ session_id, tool_use_id, tool_name: tool, tool_input: input }),
+        timeoutMs: 35000,
+      })
+      if (result.exitCode !== 0) throw new Error('Router command failed')
+      // An empty response means routing is disabled or no model could be selected.
+      route = result.stdout.trim() ? JSON.parse(result.stdout) : undefined
+      if (route !== undefined && (
+          typeof route?.model !== 'string' || !route.model ||
+          typeof route.rationale !== 'string' ||
+          typeof route.updatedInput?.subagent_type !== 'string' || !route.updatedInput.subagent_type)) {
+        throw new Error('Invalid routing response')
+      }
+    } catch {
+      route = undefined
+      $.ui.log('Smart Router could not route this subagent; using its original settings.')
+    } finally {
+      spinner?.cancel()
+    }
+    if (tool_use_id) {
+      const seconds = await secondsSince($, pending.get(tool_use_id))
+      pending.delete(tool_use_id)
+      if (route) records.set(tool_use_id, { model: route.model, rationale: route.rationale, seconds })
+      $.ui.invalidate('ui.render')
+    }
+    if (!route) return next(e)
+    // Preserve Claude's permission checks and every other mod's tool handlers.
+    return next({ ...route.updatedInput, tool, tool_use_id, agentId })
+  })
+
+  on('ui.render', { component: 'ToolUse' }, async ($, e, next) => {
+    const original = await next(e)
+    const record = records.get(e.requestId)
+    if (!record && !pending.has(e.requestId)) return original
+
+    const { Box, Text } = $.ui.resolve(e)
+    if (!record) {
+      const glyph = SPINNER[Math.floor(((await now($)) ?? 0) / SPINNER_INTERVAL_MS) % SPINNER.length]
+      const seconds = await secondsSince($, pending.get(e.requestId))
+      const elapsed = seconds === undefined ? '' : ` ${Math.floor(seconds)}s`
+      return Box({ flexDirection: 'column', children: [original, Box({
+        flexDirection: 'row',
+        children: [
+          Text({ dimColor: true, children: ['  ↳ '] }),
+          Text({ color: 'green', children: [`${glyph} `] }),
+          Text({ dimColor: true, children: [`Selecting a model…${elapsed}`] }),
+        ],
+      })] })
+    }
+    // Claude's row can name only the agent type, so the routed model is always shown.
+    const children = [original, Box({
+      flexDirection: 'row',
+      children: [
+        // Labels recede so the routed model stands out.
+        Text({ dimColor: true, children: ['  ↳ Selected Model: '] }),
+        // The same green as the status band's ON pill.
+        Text({ color: 'green', bold: true, children: [shortModel(record.model)] }),
+        // How long routing took; omitted when no clock was available.
+        ...(record.seconds === undefined
+          ? []
+          : [Text({ dimColor: true, children: [` (routed in ${record.seconds.toFixed(1)}s)`] })]),
+      ],
+    })]
+    // Align the reason, including its wrapped lines, under "Selected Model" after the arrow.
+    if (record.rationale) {
+      children.push(Box({
+        paddingLeft: 4,
+        children: [Text({ children: [Text({ dimColor: true, children: ['Reason: '] }), record.rationale] })],
+      }))
+    }
+    return Box({ flexDirection: 'column', children })
+  })
+}

src/ucode/cli.py

@@ -139,7 +139,11 @@
 from ucode.skills_list import configured_skill_counts_by_agent, list_configured_skills_command
 from ucode.skills_state import records_for_scope
 from ucode.smart_routing import v2 as smart_routing_v2
-from ucode.smart_routing.claude_hooks import FIRST_PROMPT_SOCKET_ENV, ROUTE_FIRST_PROMPT_EVENT
+from ucode.smart_routing.claude_hooks import (
+    FIRST_PROMPT_SOCKET_ENV,
+    ROUTE_FIRST_PROMPT_EVENT,
+    ROUTE_SUBAGENT_MOD_EVENT,
+)
 from ucode.smart_routing.session_env import (
     effective_environment,
     session_env_path,
@@ -2621,9 +2625,9 @@ def claude_router_hook_cmd(
                 )
             )
         # When matched is None the harness didn't report the subagent model —
-        # the PreToolUse hook already injected the routed model, so emit nothing.
+        # routing already injected the selected model, so emit nothing.
         return
-    if event != "route-subagent" or not host:
+    if event not in {"route-subagent", ROUTE_SUBAGENT_MOD_EVENT} or not host:
         return
     token = os.environ.get("OAUTH_TOKEN") or os.environ.get("DATABRICKS_BEARER")
     if not token:
@@ -2633,12 +2637,15 @@ def claude_router_hook_cmd(
             token = get_databricks_token(host, profile)
         except RuntimeError:
             return
-    output = smart_routing_v2.route_claude_pre_tool_use(
+    output = smart_routing_v2.route_claude_subagent(
         payload,
         workspace=host,
         token=token,
         available_models=model or [],
         audit_decision=True,
+        # Hooks persisted by earlier UG versions still call route-subagent and expect a
+        # PreToolUse response; the mod's route-subagent-mod event takes the plain details.
+        hook_output=event == "route-subagent",
     )
     if output is not None:
         sys.stdout.write(json.dumps(output))

src/ucode/smart_routing/claude_hooks.py

@@ -1,10 +1,8 @@
 """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.
+The native Claude mod routes subagent calls. ``SessionStart`` and
+``SubagentStart`` hooks drive the canary/audit trail.
 """
 
 from __future__ import annotations
@@ -16,12 +14,13 @@
 
 ROUTING_HOOK_COMMAND_MARKER = "claude-router-hook"
 ROUTE_FIRST_PROMPT_EVENT = "route-first-prompt"
+ROUTE_SUBAGENT_MOD_EVENT = "route-subagent-mod"
 FIRST_PROMPT_HOOK_MARKER = f"{ROUTING_HOOK_COMMAND_MARKER} {ROUTE_FIRST_PROMPT_EVENT}"
 FIRST_PROMPT_SOCKET_ENV = "UCODE_CLAUDE_V2_SOCKET"
 
 
 def sync_smart_routing_hooks(doc: dict, state: dict, *, enabled: bool) -> None:
-    """Synchronize ucode-managed routing hooks in a Claude settings document."""
+    """Synchronize ucode-managed lifecycle hooks, removing legacy routing hooks."""
     groups = _routing_hook_groups(state) if enabled else {}
     hooks.sync_managed_hooks(doc, ROUTING_HOOK_COMMAND_MARKER, groups)
 
@@ -49,16 +48,9 @@ def sync_first_prompt_hook(doc: dict, executable: str) -> None:
 
 
 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")
+    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",
@@ -73,7 +65,8 @@ def _routing_hook_groups(state: dict) -> dict[str, list[dict]]:
     }
 
 
-def _routing_hook_argv(state: dict, event: str) -> list[str]:
+def routing_hook_argv(state: dict, event: str) -> list[str]:
+    """Build the shared CLI command for a lifecycle hook or native mod request."""
     workspace = str(state.get("workspace") or "")
     argv = [
         build_auth_token_argv(workspace, state.get("profile"), use_pat=bool(state.get("use_pat")))[
@@ -82,7 +75,7 @@ def _routing_hook_argv(state: dict, event: str) -> list[str]:
         ROUTING_HOOK_COMMAND_MARKER,
         event,
     ]
-    if event != "route-subagent":
+    if event != ROUTE_SUBAGENT_MOD_EVENT:
         return argv
     argv += ["--host", workspace]
     profile = state.get("profile")

src/ucode/smart_routing/v2.py

@@ -59,6 +59,8 @@
 from ucode.smart_routing import claude_routing, orchestrator, routing
 from ucode.smart_routing.claude_hooks import (
     FIRST_PROMPT_SOCKET_ENV,
+    ROUTE_SUBAGENT_MOD_EVENT,
+    routing_hook_argv,
     sync_first_prompt_hook,
     sync_smart_routing_hooks,
 )
@@ -342,6 +344,7 @@ def _claude_mod_launch_options(env: Mapping[str, str] | None = None) -> dict:
         # Native plugin options cannot represent null; preserve absent versus empty values.
         "unset": [key for key in SMART_ROUTING_ENV_KEYS if key not in source],
         "recipe": source.get(routing.ROUTER_NAME_ENV_VAR, "").strip() or routing.ROUTER_NAME,
+        "routeCommand": [],
     }
 
 
@@ -355,6 +358,13 @@ def _write_routed_claude_plugin(plugin_dir: Path, model_ids: list[str]) -> None:
             "description": "Launch-scoped agents for Unity Gateway smart routing.",
             "author": {"name": "Databricks"},
             "userConfig": {
+                "routeCommand": {
+                    "type": "string",
+                    "multiple": True,
+                    "default": [],
+                    "title": "Subagent router",
+                    "description": "CLI argument list for a subagent routing decision",
+                },
                 **{
                     key: {"type": "string", "title": key, "description": "Original routing flag"}
                     for key in SMART_ROUTING_ENV_KEYS
@@ -395,7 +405,13 @@ def _write_routed_claude_plugin(plugin_dir: Path, model_ids: list[str]) -> None:
     orchestrator.add_claude_agents(plugin_dir)
     mods = Path(__file__).parents[1] / "agents" / "claude_mods"
     # These are imported by the entry point, so they ship beside it rather than as modules.
-    for name in ("register.ts", "constants.ts", "extra-body.ts", "smart-routing-status.ts"):
+    for name in (
+        "register.ts",
+        "constants.ts",
+        "extra-body.ts",
+        "smart-routing-status.ts",
+        "subagent-routing.ts",
+    ):
         write_text_file(plugin_dir / "hooks" / name, (mods / name).read_text(encoding="utf-8"))
     write_json_file(plugin_dir / "hooks" / "hooks.json", {"modules": ["./register.ts"]})
 
@@ -463,15 +479,20 @@ def _route_print_mode_prompt(
     return model_name(_unwrapped_claude_model_id(decision.model))
 
 
-def route_claude_pre_tool_use(
+def route_claude_subagent(
     payload: dict,
     *,
     workspace: str,
     token: str,
     available_models: list[str],
     audit_decision: bool = False,
+    hook_output: bool = False,
 ) -> dict | None:
-    """Route a Claude Agent call through a transient exact-model agent definition."""
+    """Return a routed Agent input and UI details using the shared Python client.
+
+    ``hook_output`` returns the legacy PreToolUse hook response instead, with its plain-text
+    notice, for hooks that earlier UG versions persisted in settings files.
+    """
     route = routing.resolve_spawn_route(
         payload,
         is_spawn_agent=claude_routing.is_spawn_agent_tool,
@@ -495,21 +516,29 @@ def route_claude_pre_tool_use(
         **{key: value for key, value in route.tool_input.items() if key != "model"},
         "subagent_type": _routed_claude_agent_name(route.routed_model),
     }
-    agent_name = updated_input.get("subagent_type", "")
-    if orchestrator.enabled(agent=AGENT_CLAUDE):
-        agent_name += " [orchestrator on]"
-    routing_message = claude_routing.SUBAGENT_NOTICE_CONFIG.message(
-        route.decision,
-        route.routed_model,
-        {**updated_input, "subagent_type": agent_name},
-    )
-    hook_output = {
-        "hookEventName": "PreToolUse",
-        "permissionDecision": "allow",
+    if hook_output:
+        agent_name = updated_input["subagent_type"]
+        if orchestrator.enabled(agent=AGENT_CLAUDE):
+            agent_name += " [orchestrator on]"
+        routing_message = claude_routing.SUBAGENT_NOTICE_CONFIG.message(
+            route.decision,
+            route.routed_model,
+            {**updated_input, "subagent_type": agent_name},
+        )
+        return {
+            "systemMessage": routing_message,
+            "hookSpecificOutput": {
+                "hookEventName": "PreToolUse",
+                "permissionDecision": "allow",
+                "updatedInput": updated_input,
+                "permissionDecisionReason": routing_message,
+            },
+        }
+    return {
         "updatedInput": updated_input,
-        "permissionDecisionReason": routing_message,
+        "model": route.routed_model,
+        "rationale": route.decision.rationale,
     }
-    return {"systemMessage": routing_message, "hookSpecificOutput": hook_output}
 
 
 def _is_claude_target_model(value: object) -> bool:
@@ -689,7 +718,10 @@ def route_prompt(prompt: str) -> claude_pty.FirstPromptRoute:
             env[SESSION_PYTHON_ENV_VAR] = os.environ[SESSION_PYTHON_ENV_VAR]
             settings.setdefault("pluginConfigs", {}).setdefault(CLAUDE_ROUTING_PLUGIN_NAME, {})[
                 "options"
-            ] = _claude_mod_launch_options({**resolve_environment(agent=AGENT_CLAUDE), **env})
+            ] = {
+                **_claude_mod_launch_options({**resolve_environment(agent=AGENT_CLAUDE), **env}),
+                "routeCommand": routing_hook_argv(routing_state, ROUTE_SUBAGENT_MOD_EVENT),
+            }
             try:
                 write_json_file(settings_path, settings)
                 _write_routed_claude_plugin(plugin_dir, model_ids)