Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions src/ucode/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -1109,6 +1109,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
Expand Down Expand Up @@ -1180,6 +1182,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
Expand Down
22 changes: 22 additions & 0 deletions src/ucode/databricks.py
Original file line number Diff line number Diff line change
Expand Up @@ -1374,6 +1374,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]:
Expand Down Expand Up @@ -1404,6 +1418,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)
Expand Down Expand Up @@ -2211,6 +2231,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:
Expand Down
11 changes: 8 additions & 3 deletions src/ucode/smart_routing/claude_routing.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
21 changes: 13 additions & 8 deletions src/ucode/smart_routing/routing.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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)
Expand Down
89 changes: 80 additions & 9 deletions tests/test_claude_routing.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."
)


Expand All @@ -192,7 +191,41 @@ 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_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():
Expand Down Expand Up @@ -255,14 +288,52 @@ 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


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"

Expand Down
Loading