Skip to content

Commit 504ca15

Browse files
committed
Add Claude Code smart routing (stacked on codex smart routing)
Extend smart routing to Claude Code on top of the shared engine. `ucode claude --enable-smart-routing` routes the root session (via ANTHROPIC_MODEL) and subagents (via a PreToolUse hook that rewrites the Agent tool's model) through the workspace task_v1 router. - smart_routing/claude_routing.py + claude_hooks.py: Claude cc-scenario arms (claude-opus-4-8, claude-sonnet-5, harness "claude"), both offered unconditionally per the frozen task_v1 menu; settings.json hook shape. - claude agent + CLI: --enable/--disable-smart-routing, hidden claude-router-hook, reads the shared smart_routing_enabled state key so one opt-in covers codex and claude. Co-authored-by: Isaac
1 parent 8b9957f commit 504ca15

8 files changed

Lines changed: 817 additions & 35 deletions

File tree

src/ucode/agents/__init__.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -312,6 +312,7 @@ def configure_tool(
312312
provider: str | None = None,
313313
provider_models: dict[str, str] | None = None,
314314
relayed: bool = False,
315+
route_root_model: str | None = None,
315316
) -> dict:
316317
result: dict | tuple[dict, str]
317318
if tool == "codex":
@@ -322,7 +323,12 @@ def configure_tool(
322323
if not model and not provider:
323324
raise RuntimeError(f"A {tool} model must be selected before configuration.")
324325
result = claude.write_tool_config(
325-
state, model, provider=provider, provider_models=provider_models, relayed=relayed
326+
state,
327+
model,
328+
provider=provider,
329+
provider_models=provider_models,
330+
relayed=relayed,
331+
route_root_model=route_root_model,
326332
)
327333
else:
328334
# provider routing is claude/codex-only; every other tool needs a model.

src/ucode/agents/claude.py

Lines changed: 64 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,10 @@
2929
get_databricks_token,
3030
)
3131
from ucode.launcher import exec_or_spawn
32+
from ucode.smart_routing.claude_hooks import (
33+
remove_smart_routing_hooks,
34+
sync_smart_routing_hooks,
35+
)
3236
from ucode.state import mark_tool_managed, save_state
3337
from ucode.telemetry import agent_version, ucode_version
3438
from ucode.tracing import tracing_env
@@ -46,6 +50,15 @@
4650
"backup_path": CLAUDE_BACKUP_PATH,
4751
}
4852

53+
# Per-workspace opt-in flag for Claude Code smart routing (state key).
54+
# Shared across agents: one opt-in enables smart routing for every routing-capable
55+
# tool (codex, claude), so a workspace turns it on once. Kept identical to
56+
# codex.SMART_ROUTING_STATE_KEY on purpose.
57+
SMART_ROUTING_STATE_KEY = "smart_routing_enabled"
58+
# Claude Code settings.json hook events ucode manages when routing is enabled;
59+
# marked managed so they're tracked/reverted with the rest of ucode's config.
60+
CLAUDE_ROUTING_HOOK_EVENTS = ("PreToolUse", "SessionStart", "SubagentStart")
61+
4962

5063
def is_update_available() -> tuple[str, str] | None:
5164
return available_npm_package_update(SPEC["package"])
@@ -161,6 +174,7 @@ def render_overlay(
161174
fable_enabled: bool = False,
162175
relayed: bool = False,
163176
relayed_base_url: str | None = None,
177+
route_root_model: str | None = None,
164178
) -> tuple[dict, list[list[str]]]:
165179
"""Return (overlay, managed_key_paths) for Claude settings.json.
166180
@@ -212,13 +226,20 @@ def render_overlay(
212226
"ENABLE_TOOL_SEARCH": "1",
213227
"CLAUDE_CODE_USE_GATEWAY": "1",
214228
}
215-
# Intentionally NOT setting ANTHROPIC_MODEL. Setting it produces a duplicate
216-
# catalog row in Claude Code's /model picker (e.g. "Opus 4.8 (1M context) ✓")
217-
# on top of the family-alias row from ANTHROPIC_DEFAULT_OPUS_MODEL. Without
218-
# it, Default resolves through the pinned family alias and the picker shows
219-
# only one row per model. `ucode claude -- --model X` still overrides for a
220-
# single session via Claude Code's own --model flag.
229+
# Intentionally NOT setting ANTHROPIC_MODEL by default. Setting it produces a
230+
# duplicate catalog row in Claude Code's /model picker (e.g. "Opus 4.8 (1M
231+
# context) ✓") on top of the family-alias row from ANTHROPIC_DEFAULT_OPUS_MODEL.
232+
# Without it, Default resolves through the pinned family alias and the picker
233+
# shows only one row per model. `ucode claude -- --model X` still overrides for
234+
# a single session via Claude Code's own --model flag.
235+
#
236+
# The one exception is smart routing: `route_root_model` pins the
237+
# router's per-launch pick for the root session as ANTHROPIC_MODEL. The
238+
# duplicate-picker-row cost is acceptable because the whole point is to launch
239+
# on the routed model rather than the family default.
221240
_ = model # API stability; no longer pinned via env.
241+
if route_root_model:
242+
env["ANTHROPIC_MODEL"] = route_root_model
222243
# A Bedrock-backed provider needs its provider-side ids pinned verbatim
223244
# (Claude Code's canonical names aren't routable there). These come from the
224245
# service's targets, already de-duped to one id per family upstream.
@@ -336,12 +357,41 @@ def _unregister_web_search_mcp() -> None:
336357
pass
337358

338359

360+
def smart_routing_enabled(state: dict) -> bool:
361+
"""Return whether the current workspace opted into Claude Code routing."""
362+
return state.get(SMART_ROUTING_STATE_KEY) is True
363+
364+
365+
def enable_smart_routing(state: dict) -> dict:
366+
"""Persist the current workspace's Claude Code smart-routing opt-in."""
367+
state[SMART_ROUTING_STATE_KEY] = True
368+
return state
369+
370+
371+
def disable_smart_routing(state: dict) -> bool:
372+
"""Disable routing and remove only ucode's Claude Code routing hooks."""
373+
state.pop(SMART_ROUTING_STATE_KEY, None)
374+
if state.get("workspace"):
375+
save_state(state)
376+
changed = False
377+
if CLAUDE_SETTINGS_PATH.exists():
378+
doc = read_json_safe(CLAUDE_SETTINGS_PATH)
379+
if remove_smart_routing_hooks(doc):
380+
write_json_file(CLAUDE_SETTINGS_PATH, doc)
381+
changed = True
382+
from ucode.smart_routing.claude_routing import clear_routing_artifacts
383+
384+
clear_routing_artifacts()
385+
return changed
386+
387+
339388
def write_tool_config(
340389
state: dict,
341390
model: str | None,
342391
provider: str | None = None,
343392
provider_models: dict[str, str] | None = None,
344393
relayed: bool = False,
394+
route_root_model: str | None = None,
345395
) -> dict:
346396
backup_existing_file(CLAUDE_SETTINGS_PATH, CLAUDE_BACKUP_PATH)
347397
web_search_model = _resolve_web_search_model(state)
@@ -360,6 +410,7 @@ def write_tool_config(
360410
fable_enabled=bool(state.get("fable_enabled")),
361411
relayed=relayed,
362412
relayed_base_url=relayed_base_url,
413+
route_root_model=route_root_model,
363414
)
364415
tracing_env_vars = tracing_env(state, "claude")
365416
stop_hook_command = claude_tracing_stop_hook_command() if tracing_env_vars else None
@@ -404,6 +455,13 @@ def write_tool_config(
404455
if isinstance(merged_env, dict):
405456
for key in CLAUDE_REMOVED_ENV_KEYS:
406457
merged_env.pop(key, None)
458+
# Smart-routing hooks: install ucode's PreToolUse/SessionStart/
459+
# SubagentStart hooks when routing is enabled (and not under a provider,
460+
# which pins no Databricks model), else surgically strip only ucode's own.
461+
routing_enabled = smart_routing_enabled(state) and provider is None
462+
sync_smart_routing_hooks(merged, state, enabled=routing_enabled)
463+
if routing_enabled:
464+
managed_keys = managed_keys + [["hooks", event] for event in CLAUDE_ROUTING_HOOK_EVENTS]
407465
write_json_file(CLAUDE_SETTINGS_PATH, merged)
408466

409467
if web_search_model:

0 commit comments

Comments
 (0)