From 202531ea47081908339fd5902b3af1bdbd868a57 Mon Sep 17 00:00:00 2001 From: Mason Cao Date: Thu, 6 Aug 2026 22:28:44 +0000 Subject: [PATCH 1/5] README: document smart routing for Claude Code and fix flag names The flag was renamed from --enable-intelligent-routing to --enable-smart-routing, and Claude Code smart routing was added in #255. Update the README to match the actual CLI. Co-authored-by: Isaac --- README.md | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index b463f5cd..338fe299 100644 --- a/README.md +++ b/README.md @@ -44,19 +44,21 @@ ucode codex --full-auto All agents route through Databricks AI Gateway using your workspace credentials — no API keys required. -Codex intelligent routing is opt-in. Enabling it asks the AI Gateway router to select the -root-session model before launch and installs profile-scoped hooks that route future -`spawn_agent` calls. Codex may require one-time review of the installed hooks through `/hooks`. +Smart routing is opt-in for Codex and Claude Code. Enabling it asks the AI Gateway router +to select the root-session model before launch and installs profile-scoped hooks that route +future subagent calls. Codex may require one-time review of the installed hooks through `/hooks`. ```bash -ucode codex --enable-intelligent-routing +ucode codex --enable-smart-routing +ucode claude --enable-smart-routing ``` -The setting persists for the current workspace. Disable it and remove only ucode's routing +The setting persists per workspace for each agent. Disable and remove only ucode's routing hooks with: ```bash -ucode codex --disable-intelligent-routing +ucode codex --disable-smart-routing +ucode claude --disable-smart-routing ``` To configure all tools at once: @@ -181,8 +183,10 @@ you to run `ucode ` (existing agent sessions need a restart before the MC | `ucode configure --workspaces https://first.databricks.com,https://second.databricks.com` | Configure workspaces without the interactive picker | | `ucode configure --profiles DEFAULT` | Configure using existing Databricks CLI profiles (hosts come from `~/.databrickscfg`) | | `ucode configure --profiles DEFAULT --use-pat` | Authenticate with the profile's personal access token — no browser login | -| `ucode codex --enable-intelligent-routing` | Enable AI Gateway routing for Codex sessions and subagents | -| `ucode codex --disable-intelligent-routing` | Disable routing and remove ucode's Codex routing hooks | +| `ucode codex --enable-smart-routing` | Enable AI Gateway routing for Codex sessions and subagents | +| `ucode codex --disable-smart-routing` | Disable routing and remove ucode's Codex routing hooks | +| `ucode claude --enable-smart-routing` | Enable AI Gateway routing for Claude Code sessions and subagents | +| `ucode claude --disable-smart-routing` | Disable routing and remove ucode's Claude Code routing hooks | | `ucode configure --skip-validate` | Write configs without sending a test message through each agent | | `ucode configure --agents claude --mcp system.ai.slack` | Configure an agent and register its Databricks MCP server(s) in one command | | `ucode configure skills` | Register the skills MCP connection (utility tools only); no skills download | From 9618694e8766adcc56f021431c600f185bc4a8e6 Mon Sep 17 00:00:00 2001 From: Mason Cao Date: Thu, 6 Aug 2026 23:01:27 +0000 Subject: [PATCH 2/5] routing: pin opus-4-8 for Claude smart routing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Discovery's newest-wins family bucketing picks claude-opus-5 over claude-opus-4-8, but CLAUDE_ROUTE_ARMS hardcodes claude-opus-4-8 — so the routing availability check fails with 'required Claude routing models are unavailable: claude-opus-4-8' on workspaces that have both. Pin the opus slot to 4-8 when both exist so routing works with the currently-deployed task_v1 router. Revert once the router accepts opus-5 (PR databricks-eng/universe#2365446). --- src/ucode/databricks.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index b9e48bca..279d7712 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -1344,6 +1344,20 @@ def list_model_services( return [], last_reason or "model-services listing returned no models" +def _prefer_opus_4_8(models: dict[str, str], all_ids: list[str]) -> None: + """Swap the opus slot to claude-opus-4-8 when it's available. + + Discovery picks the newest opus (opus-5) but smart routing's + CLAUDE_ROUTE_ARMS require claude-opus-4-8. Pin to 4-8 when both + exist so the routing availability check passes. + """ + opus = models.get("opus") + if opus and "claude-opus-5" in opus: + opus_48 = next((m for m in all_ids if "claude-opus-4-8" in m), None) + if opus_48: + models["opus"] = opus_48 + + def discover_model_services( workspace: str, token: str ) -> tuple[dict[str, str], list[str], list[str], list[str], str | None]: @@ -1374,6 +1388,12 @@ def discover_model_services( ) if candidates: claude_models[family] = candidates[0] + # Smart routing's CLAUDE_ROUTE_ARMS require claude-opus-4-8, but the + # newest-wins sort above picks opus-5 when both exist — making the + # routing availability check fail. Pin opus-4-8 when it's available so + # routing works with the currently-deployed task_v1 router. Revert to + # newest-wins once the router accepts opus-5 (PR databricks-eng/universe#2365446). + _prefer_opus_4_8(claude_models, ids) codex_models = [m for m in ids if "gpt-" in m] gemini_models = sorted([m for m in ids if "gemini-" in m], key=model_version_sort_key) @@ -2181,6 +2201,8 @@ def discover_claude_models(workspace: str, token: str) -> tuple[dict[str, str], ) if candidates: result[family] = candidates[0] + # Same opus-4-8 pin as discover_model_services — see comment there. + _prefer_opus_4_8(result, raw_ids) if result: return result, None if not raw_ids: From 7983e58cf3899674ad0a4824b7e3a761cb02dde5 Mon Sep 17 00:00:00 2001 From: Mason Cao Date: Thu, 6 Aug 2026 23:13:55 +0000 Subject: [PATCH 3/5] routing: map Claude subagent model to short family name Claude Code's Agent tool model field only accepts short family names (sonnet, opus, haiku, fable), not full workspace ids. The PreToolUse hook was injecting the full id (e.g. system.ai.claude-sonnet-5) into updatedInput.model, which Claude Code rejected with a schema validation error. Map the router's pick back to its family name before injecting. --- src/ucode/smart_routing/claude_routing.py | 11 ++++++++--- tests/test_claude_routing.py | 17 ++++++++--------- 2 files changed, 16 insertions(+), 12 deletions(-) diff --git a/src/ucode/smart_routing/claude_routing.py b/src/ucode/smart_routing/claude_routing.py index 005eb284..79c95064 100644 --- a/src/ucode/smart_routing/claude_routing.py +++ b/src/ucode/smart_routing/claude_routing.py @@ -187,8 +187,13 @@ def clear_routing_artifacts() -> None: 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. + Claude Code's ``Agent`` tool ``model`` field accepts only short family + names (``sonnet``, ``opus``, ``haiku``, ``fable``), not full workspace + ids. Map the router's pick (e.g. ``system.ai.claude-sonnet-5``) back to + its family name. """ + normalized = _normalize_model(model) + for family in ("fable", "opus", "sonnet", "haiku"): + if f"claude-{family}-" in normalized: + return family return model diff --git a/tests/test_claude_routing.py b/tests/test_claude_routing.py index 49f77afc..b85478aa 100644 --- a/tests/test_claude_routing.py +++ b/tests/test_claude_routing.py @@ -155,21 +155,20 @@ def test_spawn_rewrite_injects_routed_model(monkeypatch): hook = output["hookSpecificOutput"] # The rationale is surfaced in the systemMessage (shown to the user), not - # only in permissionDecisionReason. + # only in permissionDecisionReason. The model field is the short family + # name ("opus") that Claude Code's Agent tool schema accepts. assert output["systemMessage"] == ( - "Using Smart Routing. Routing to system.ai.claude-opus-4-8. " - "Deep exploration needs the strongest model." + "Using Smart Routing. Routing to opus. 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", + "model": "opus", } assert hook["permissionDecisionReason"] == ( - "Using Smart Routing. Routing to system.ai.claude-opus-4-8. " - "Deep exploration needs the strongest model." + "Using Smart Routing. Routing to opus. Deep exploration needs the strongest model." ) @@ -192,7 +191,7 @@ def test_task_tool_alias_is_routed(monkeypatch): available_models=["system.ai.claude-opus-4-8", "system.ai.claude-sonnet-5"], ) - assert output["hookSpecificOutput"]["updatedInput"]["model"] == "system.ai.claude-sonnet-5" + assert output["hookSpecificOutput"]["updatedInput"]["model"] == "sonnet" def test_non_spawn_tool_has_no_opinion(): @@ -255,11 +254,11 @@ def test_decision_is_reconciled_with_actual_subagent_model(tmp_path, monkeypatch audit_decision=True, ) record = claude_routing.record_subagent_start( - {"session_id": "s1", "agent_id": "a1", "model": "system.ai.claude-opus-4-8"} + {"session_id": "s1", "agent_id": "a1", "model": "opus"} ) assert record["router_model"] == "claude-opus-4-8" - assert record["requested_model"] == "system.ai.claude-opus-4-8" + assert record["requested_model"] == "opus" assert record["matches_router_decision"] is True From e6450516c3f047c411961f15a81f996ecfb464dd Mon Sep 17 00:00:00 2001 From: Mason Cao Date: Thu, 6 Aug 2026 23:28:49 +0000 Subject: [PATCH 4/5] routing: don't report mismatch when harness omits subagent model MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Claude Code's SubagentStart event does not include the subagent's model field, so the reconciliation audit recorded model=null and reported a false mismatch. When the actual model is unknown, record matches_router_decision as None (unknown) rather than False (mismatch), and emit no SubagentStart message to the user — the PreToolUse hook already injected the routed model. --- src/ucode/cli.py | 4 ++++ src/ucode/smart_routing/routing.py | 7 +++++- tests/test_claude_routing.py | 38 ++++++++++++++++++++++++++++++ 3 files changed, 48 insertions(+), 1 deletion(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 52047c72..7c42e78a 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -1043,6 +1043,8 @@ def codex_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. return if event != "route-subagent" or not host: return @@ -1114,6 +1116,8 @@ 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. return if event != "route-subagent" or not host: return diff --git a/src/ucode/smart_routing/routing.py b/src/ucode/smart_routing/routing.py index 594fc0e9..e8169d16 100644 --- a/src/ucode/smart_routing/routing.py +++ b/src/ucode/smart_routing/routing.py @@ -269,12 +269,17 @@ def record_subagent_start( "at": time.time(), } if decision is not None: + # When the harness doesn't report the subagent's model (actual_model is + # None), we can't verify the match — record None rather than a false + # mismatch. The PreToolUse hook already injected the routed model, so + # routing still worked; the reconciliation is observability, not enforcement. + matches = None if actual_model is None else decision.get("requested_model") == actual_model record.update( { "decision_id": decision.get("decision_id"), "router_model": decision.get("router_model"), "requested_model": decision.get("requested_model"), - "matches_router_decision": decision.get("requested_model") == actual_model, + "matches_router_decision": matches, } ) _append_jsonl(audit_path, record) diff --git a/tests/test_claude_routing.py b/tests/test_claude_routing.py index b85478aa..7b34e307 100644 --- a/tests/test_claude_routing.py +++ b/tests/test_claude_routing.py @@ -262,6 +262,44 @@ def test_decision_is_reconciled_with_actual_subagent_model(tmp_path, monkeypatch assert record["matches_router_decision"] is True +def test_subagent_start_without_model_reports_unknown_not_mismatch(tmp_path, monkeypatch): + # Claude Code's SubagentStart event does not include the subagent's model, + # so reconciliation can't verify the match — record None (unknown) rather + # than a false mismatch. The PreToolUse hook already injected the model. + 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-sonnet-5", raw_model="claude-sonnet-5" + ), + None, + ), + ) + + claude_routing.route_pre_tool_use( + { + "session_id": "s2", + "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": "s2", "agent_id": "a2", "model": None} + ) + + assert record["requested_model"] == "sonnet" + assert record["matches_router_decision"] is None + + def test_launch_task_uses_positional_prompt(): assert claude_routing._launch_routing_task(["fix the parser"]) == "fix the parser" From 2b3c7e464e7c9781ea94a4eef0fdc6b358d4ff08 Mon Sep 17 00:00:00 2001 From: Mason Cao Date: Thu, 6 Aug 2026 23:35:50 +0000 Subject: [PATCH 5/5] routing: send Claude subagent prompt to router instead of generic label MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit route_spawn_tool looked for message/task_name/agent_name to derive the routing task, but Claude Code's Agent tool puts the subagent task in prompt/description — so Claude always fell through to the generic "Claude Code subagent task" label, and the router made decisions on a constant string. Add prompt and description to the front of the field fallback chain. Safe for Codex: its spawn_agent tool uses message, which is absent from Claude payloads, and vice versa. --- src/ucode/smart_routing/routing.py | 14 ++++++------ tests/test_claude_routing.py | 34 ++++++++++++++++++++++++++++++ 2 files changed, 41 insertions(+), 7 deletions(-) diff --git a/src/ucode/smart_routing/routing.py b/src/ucode/smart_routing/routing.py index e8169d16..80df3584 100644 --- a/src/ucode/smart_routing/routing.py +++ b/src/ucode/smart_routing/routing.py @@ -202,16 +202,16 @@ def route_spawn_tool( tool_input = payload.get("tool_input") if not isinstance(tool_input, dict): return None - # Derive the routing task from the first available plaintext field. `message` - # carries the actual subagent task content — prefer it when present and a - # plaintext string (Codex encrypts it at send-time, but the PreToolUse hook - # fires before that, so it may be readable here). When `message` is an - # encrypted dict (or absent), fall back to `task_name` / `agent_name` - # (weaker labels), then the generic default. + # Derive the routing task from the first available plaintext field. The + # harness-specific names are tried in order: `prompt`/`description` (Claude + # Code's Agent tool), `message` (Codex's spawn_agent — encrypted at + # send-time but readable here because the PreToolUse hook fires before + # that), then `task_name` / `agent_name` (weaker labels), then the generic + # default. task = next( ( value - for field in ("message", "task_name", "agent_name") + for field in ("prompt", "description", "message", "task_name", "agent_name") if isinstance(value := tool_input.get(field), str) and value ), default_task_label, diff --git a/tests/test_claude_routing.py b/tests/test_claude_routing.py index 7b34e307..d30ee669 100644 --- a/tests/test_claude_routing.py +++ b/tests/test_claude_routing.py @@ -194,6 +194,40 @@ def test_task_tool_alias_is_routed(monkeypatch): assert output["hookSpecificOutput"]["updatedInput"]["model"] == "sonnet" +def test_spawn_routes_on_prompt_not_generic_label(monkeypatch): + # Claude Code's Agent tool puts the subagent task in `prompt`/`description`, + # not `message`. The router should receive the real prompt, not the generic + # "Claude Code subagent task" fallback. + captured = {} + + def fake_decision(workspace, token, task, models, **kwargs): + captured["task"] = task + return ( + claude_routing.RoutingDecision( + model="system.ai.claude-sonnet-5", raw_model="claude-sonnet-5" + ), + None, + ) + + monkeypatch.setattr(claude_routing, "request_routing_decision", fake_decision) + + claude_routing.route_pre_tool_use( + { + "tool_name": "Agent", + "tool_input": { + "subagent_type": "Explore", + "prompt": "find all Python entry points", + "description": "explore the repo", + }, + }, + workspace=WS, + token="token", + available_models=["system.ai.claude-opus-4-8", "system.ai.claude-sonnet-5"], + ) + + assert captured["task"] == "find all Python entry points" + + def test_non_spawn_tool_has_no_opinion(): assert ( claude_routing.route_pre_tool_use(