Merge remote-tracking branch 'origin/main' into pr-257-resolve

Signed-off-by: Anthony Ivan <anthony.ivan@databricks.com>

# Conflicts:
#	src/ucode/cli.py

Author: Anthony Ivan

README.md

@@ -40,6 +40,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`.
+
+```bash
+ucode codex --enable-intelligent-routing
+```
+
+The setting persists for the current workspace. Disable it and remove only ucode's routing
+hooks with:
+
+```bash
+ucode codex --disable-intelligent-routing
+```
+
 To configure all tools at once:
 
 ```bash
@@ -162,6 +177,8 @@ you to run `ucode <agent>` (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 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 |

src/ucode/agents/__init__.py

@@ -312,6 +312,7 @@ def configure_tool(
     provider: str | None = None,
     provider_models: dict[str, str] | None = None,
     relayed: bool = False,
+    route_root_model: str | None = None,
 ) -> dict:
     result: dict | tuple[dict, str]
     if tool == "codex":
@@ -322,7 +323,12 @@ def configure_tool(
         if not model and not provider:
             raise RuntimeError(f"A {tool} model must be selected before configuration.")
         result = claude.write_tool_config(
-            state, model, provider=provider, provider_models=provider_models, relayed=relayed
+            state,
+            model,
+            provider=provider,
+            provider_models=provider_models,
+            relayed=relayed,
+            route_root_model=route_root_model,
         )
     else:
         # provider routing is claude/codex-only; every other tool needs a model.

src/ucode/agents/claude.py

@@ -30,6 +30,10 @@
     get_databricks_token,
 )
 from ucode.launcher import exec_or_spawn
+from ucode.smart_routing.claude_hooks import (
+    remove_smart_routing_hooks,
+    sync_smart_routing_hooks,
+)
 from ucode.state import mark_tool_managed, save_state
 from ucode.telemetry import agent_version, ucode_version
 from ucode.tracing import tracing_env
@@ -47,6 +51,15 @@
     "backup_path": CLAUDE_BACKUP_PATH,
 }
 
+# Per-workspace opt-in flag for Claude Code smart routing (state key).
+# Shared across agents: one opt-in enables smart routing for every routing-capable
+# tool (codex, claude), so a workspace turns it on once. Kept identical to
+# codex.SMART_ROUTING_STATE_KEY on purpose.
+SMART_ROUTING_STATE_KEY = "smart_routing_enabled"
+# Claude Code settings.json hook events ucode manages when routing is enabled;
+# marked managed so they're tracked/reverted with the rest of ucode's config.
+CLAUDE_ROUTING_HOOK_EVENTS = ("PreToolUse", "SessionStart", "SubagentStart")
+
 
 def is_update_available() -> tuple[str, str] | None:
     return available_npm_package_update(SPEC["package"])
@@ -214,6 +227,7 @@ def render_overlay(
     fable_enabled: bool = False,
     relayed: bool = False,
     relayed_base_url: str | None = None,
+    route_root_model: str | None = None,
 ) -> tuple[dict, list[list[str]]]:
     """Return (overlay, managed_key_paths) for Claude settings.json.
 
@@ -265,13 +279,20 @@ def render_overlay(
         "ENABLE_TOOL_SEARCH": "1",
         "CLAUDE_CODE_USE_GATEWAY": "1",
     }
-    # Intentionally NOT setting ANTHROPIC_MODEL. Setting it produces a duplicate
-    # catalog row in Claude Code's /model picker (e.g. "Opus 4.8 (1M context) ✓")
-    # on top of the family-alias row from ANTHROPIC_DEFAULT_OPUS_MODEL. Without
-    # it, Default resolves through the pinned family alias and the picker shows
-    # only one row per model. `ucode claude -- --model X` still overrides for a
-    # single session via Claude Code's own --model flag.
+    # Intentionally NOT setting ANTHROPIC_MODEL by default. Setting it produces a
+    # duplicate catalog row in Claude Code's /model picker (e.g. "Opus 4.8 (1M
+    # context) ✓") on top of the family-alias row from ANTHROPIC_DEFAULT_OPUS_MODEL.
+    # Without it, Default resolves through the pinned family alias and the picker
+    # shows only one row per model. `ucode claude -- --model X` still overrides for
+    # a single session via Claude Code's own --model flag.
+    #
+    # The one exception is smart routing: `route_root_model` pins the
+    # router's per-launch pick for the root session as ANTHROPIC_MODEL. The
+    # duplicate-picker-row cost is acceptable because the whole point is to launch
+    # on the routed model rather than the family default.
     _ = model  # API stability; no longer pinned via env.
+    if route_root_model:
+        env["ANTHROPIC_MODEL"] = route_root_model
     # A Bedrock-backed provider needs its provider-side ids pinned verbatim
     # (Claude Code's canonical names aren't routable there). These come from the
     # service's targets, already de-duped to one id per family upstream.
@@ -389,12 +410,41 @@ def _unregister_web_search_mcp() -> None:
             pass
 
 
+def smart_routing_enabled(state: dict) -> bool:
+    """Return whether the current workspace opted into Claude Code routing."""
+    return state.get(SMART_ROUTING_STATE_KEY) is True
+
+
+def enable_smart_routing(state: dict) -> dict:
+    """Persist the current workspace's Claude Code smart-routing opt-in."""
+    state[SMART_ROUTING_STATE_KEY] = True
+    return state
+
+
+def disable_smart_routing(state: dict) -> bool:
+    """Disable routing and remove only ucode's Claude Code routing hooks."""
+    state.pop(SMART_ROUTING_STATE_KEY, None)
+    if state.get("workspace"):
+        save_state(state)
+    changed = False
+    if CLAUDE_SETTINGS_PATH.exists():
+        doc = read_json_safe(CLAUDE_SETTINGS_PATH)
+        if remove_smart_routing_hooks(doc):
+            write_json_file(CLAUDE_SETTINGS_PATH, doc)
+            changed = True
+    from ucode.smart_routing.claude_routing import clear_routing_artifacts
+
+    clear_routing_artifacts()
+    return changed
+
+
 def write_tool_config(
     state: dict,
     model: str | None,
     provider: str | None = None,
     provider_models: dict[str, str] | None = None,
     relayed: bool = False,
+    route_root_model: str | None = None,
 ) -> dict:
     backup_existing_file(CLAUDE_SETTINGS_PATH, CLAUDE_BACKUP_PATH)
     web_search_model = _resolve_web_search_model(state)
@@ -413,6 +463,7 @@ def write_tool_config(
         fable_enabled=bool(state.get("fable_enabled")),
         relayed=relayed,
         relayed_base_url=relayed_base_url,
+        route_root_model=route_root_model,
     )
     tracing_env_vars = tracing_env(state, "claude")
     stop_hook_command = claude_tracing_stop_hook_command() if tracing_env_vars else None
@@ -457,6 +508,13 @@ def write_tool_config(
     if isinstance(merged_env, dict):
         for key in CLAUDE_REMOVED_ENV_KEYS:
             merged_env.pop(key, None)
+    # Smart-routing hooks: install ucode's PreToolUse/SessionStart/
+    # SubagentStart hooks when routing is enabled (and not under a provider,
+    # which pins no Databricks model), else surgically strip only ucode's own.
+    routing_enabled = smart_routing_enabled(state) and provider is None
+    sync_smart_routing_hooks(merged, state, enabled=routing_enabled)
+    if routing_enabled:
+        managed_keys = managed_keys + [["hooks", event] for event in CLAUDE_ROUTING_HOOK_EVENTS]
     write_json_file(CLAUDE_SETTINGS_PATH, merged)
 
     if web_search_model:

src/ucode/agents/codex.py

@@ -21,6 +21,10 @@
     get_databricks_token,
 )
 from ucode.launcher import exec_or_spawn
+from ucode.smart_routing.codex_hooks import (
+    remove_smart_routing_hooks,
+    sync_smart_routing_hooks,
+)
 from ucode.state import mark_tool_managed, save_state
 from ucode.telemetry import agent_version, ucode_version
 
@@ -33,6 +37,11 @@
 CODEX_MODEL_PROVIDER_NAME = "ucode-databricks"
 MINIMUM_CODEX_VERSION = (0, 134, 0)
 MINIMUM_CODEX_VERSION_TEXT = "0.134.0"
+MINIMUM_ROUTING_CODEX_VERSION = (0, 145, 0)
+MINIMUM_ROUTING_CODEX_VERSION_TEXT = "0.145.0"
+# Shared across agents: one opt-in enables smart routing for every routing-capable
+# tool (codex, claude), so a workspace turns it on once.
+SMART_ROUTING_STATE_KEY = "smart_routing_enabled"
 
 
 SPEC: ToolSpec = {
@@ -318,6 +327,10 @@ def write_tool_config(state: dict, model: str | None = None, provider: str | Non
     databricks_profile = state.get("profile")
 
     if _use_legacy_layout():
+        if smart_routing_enabled(state) and provider is None:
+            raise RuntimeError(
+                f"Codex smart routing requires Codex {MINIMUM_ROUTING_CODEX_VERSION_TEXT} or newer."
+            )
         # Codex < 0.134.0 only reads ~/.codex/config.toml. Write the shared
         # config with [profiles.ucode] + shared [model_providers.ucode-databricks]
         # and skip the per-profile-file cleanup that would normally strip
@@ -358,6 +371,11 @@ def write_tool_config(state: dict, model: str | None = None, provider: str | Non
         # deep_merge can't drop keys, so clear a `model` pinned by an earlier
         # non-provider run that the provider overlay omits.
         doc.pop("model", None)
+    sync_smart_routing_hooks(
+        doc,
+        state,
+        enabled=smart_routing_enabled(state) and provider is None,
+    )
     write_toml_file(CODEX_CONFIG_PATH, doc)
     state = mark_tool_managed(state, "codex", MANAGED_KEYS)
     save_state(state)
@@ -399,6 +417,43 @@ def launch(state: dict, tool_args: list[str]) -> None:
     exec_or_spawn([binary, "--profile", CODEX_PROFILE_NAME, *tool_args])
 
 
+def smart_routing_enabled(state: dict) -> bool:
+    """Return whether the current workspace opted into Codex routing."""
+    return state.get(SMART_ROUTING_STATE_KEY) is True
+
+
+def enable_smart_routing(state: dict) -> dict:
+    """Persist the current workspace's Codex smart-routing opt-in."""
+    parsed = _parse_version(agent_version(SPEC["binary"]))
+    if parsed is not None and parsed < MINIMUM_ROUTING_CODEX_VERSION:
+        raise RuntimeError(
+            "Codex smart routing requires Codex "
+            f"{MINIMUM_ROUTING_CODEX_VERSION_TEXT} or newer; found "
+            f"{agent_version(SPEC['binary'])}."
+        )
+    state[SMART_ROUTING_STATE_KEY] = True
+    return state
+
+
+def disable_smart_routing(state: dict) -> bool:
+    """Disable routing and remove only ucode's Codex routing hooks."""
+    state.pop(SMART_ROUTING_STATE_KEY, None)
+    if state.get("workspace"):
+        save_state(state)
+    changed = False
+    for path in (CODEX_CONFIG_PATH, LEGACY_CODEX_CONFIG_PATH):
+        if not path.exists():
+            continue
+        doc = read_toml_safe(path)
+        if remove_smart_routing_hooks(doc):
+            write_toml_file(path, doc)
+            changed = True
+    from ucode.smart_routing.codex_routing import clear_routing_artifacts
+
+    clear_routing_artifacts()
+    return changed
+
+
 def validate_cmd(binary: str) -> list[str]:
     return [
         binary,