test

src/ucode/agents/claude.py

@@ -11,6 +11,7 @@
 import socket
 import subprocess
 import threading
+import uuid
 from collections.abc import Callable
 from pathlib import Path
 from typing import cast
@@ -26,13 +27,17 @@
 )
 from ucode.databricks import (
     build_auth_shell_command,
+    build_auth_token_argv,
     build_tool_base_url,
     get_databricks_token,
 )
 from ucode.launcher import exec_or_spawn
 from ucode.managed_files import OS, current_os, write_managed_file
+from ucode.smart_routing import v2 as smart_routing_v2
 from ucode.smart_routing.claude_hooks import (
+    FIRST_PROMPT_SOCKET_ENV,
     remove_smart_routing_hooks,
+    sync_first_prompt_hook,
     sync_smart_routing_hooks,
 )
 from ucode.state import mark_tool_managed, save_state
@@ -62,6 +67,15 @@
 # marked managed so they're tracked/reverted with the rest of ucode's config.
 CLAUDE_ROUTING_HOOK_EVENTS = ("PreToolUse", "SessionStart", "SubagentStart")
 
+# Smart-routing-v2 (model-routing proxy) knobs. The enable flag is shared across agents
+# in smart_routing.v2; only Claude-specific values live here. A loopback proxy sits
+# between Claude Code and the gateway and rewrites each request's `model` to the router's
+# pick. SMART_ROUTING_V2_MODEL is a STUB standing in for a real per-prompt routing
+# decision (see smart_routing.routing.select_route) — swap the router callable in
+# `_v2_router` for genuine dynamic routing.
+SMART_ROUTING_V2_MODEL = "system.ai.claude-sonnet-5"  # stubbed router pick
+SMART_ROUTING_V2_CLAUDE_LOG = APP_DIR / "claude-v2-pty.log"
+
 
 def is_update_available() -> tuple[str, str] | None:
     return available_npm_package_update(SPEC["package"])
@@ -1017,12 +1031,89 @@ def _launch_relayed(state: dict, binary: str, tool_args: list[str]) -> None:
     raise SystemExit(returncode)
 
 
+def _v2_router(state: dict):
+    """Return the routing callable applied to the hook-captured first prompt.
+
+    Today a STUB: ignores the request and always returns SMART_ROUTING_V2_MODEL. For real
+    dynamic routing, classify the prompt here (see ``smart_routing.routing.select_route``)
+    and return the chosen model.
+    ``state`` is accepted so a real router can reach the workspace/token/model metadata."""
+
+    def route(_prompt: str) -> str:
+        return SMART_ROUTING_V2_MODEL
+
+    return route
+
+
+def _launch_smart_routing_v2(state: dict, tool_args: list[str]) -> None:
+    """Launch Claude Code in the first-prompt routing PTY wrapper.
+
+    A per-launch UserPromptSubmit hook captures and blocks the first real prompt. The
+    wrapper switches Claude Code's own model through ``/model`` and replays that prompt,
+    so Claude's label/session state and the model serving the first response agree.
+    """
+    from ucode.smart_routing import claude_pty
+
+    binary = SPEC["binary"]
+    workspace = state["workspace"]
+    os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile"))
+
+    run_id = f"{os.getpid()}-{uuid.uuid4().hex[:8]}"
+    socket_path = APP_DIR / f"claude-v2-{run_id}.sock"
+    settings_path = APP_DIR / f"claude-v2-{run_id}.json"
+
+    # Compose caller settings exactly like the normal launch path, then add the
+    # first-prompt hook only to this process's temporary settings. A unique path avoids
+    # cross-talk between concurrent Claude sessions.
+    caller_values, remaining = _extract_caller_settings(tool_args)
+    settings: dict = {}
+    for value in caller_values:
+        settings = _merge_claude_settings(settings, _load_caller_settings(value))
+    settings = _merge_claude_settings(settings, read_json_safe(CLAUDE_SETTINGS_PATH))
+    hook_executable = build_auth_token_argv(
+        workspace, state.get("profile"), use_pat=bool(state.get("use_pat"))
+    )[0]
+    env = settings.setdefault("env", {})
+    if not isinstance(env, dict):
+        raise RuntimeError("Claude settings 'env' must be an object for smart routing.")
+    env[FIRST_PROMPT_SOCKET_ENV] = str(socket_path)
+    sync_first_prompt_hook(settings, hook_executable)
+    write_json_file(settings_path, settings)
+    argv = [binary, "--settings", str(settings_path), *remaining]
+
+    print_note(
+        f"Smart routing v2: the first submitted prompt will select Claude Code's model "
+        f"(stub -> {SMART_ROUTING_V2_MODEL}); log: {SMART_ROUTING_V2_CLAUDE_LOG}."
+    )
+    try:
+        returncode = claude_pty.run_claude_pty(
+            argv,
+            route_prompt=_v2_router(state),
+            switch_message=(
+                f"✨ Databricks Smart Router selected {SMART_ROUTING_V2_MODEL}. "
+                "Switching Claude Code before running your prompt."
+            ),
+            socket_path=socket_path,
+            log_path=SMART_ROUTING_V2_CLAUDE_LOG,
+        )
+    finally:
+        settings_path.unlink(missing_ok=True)
+        socket_path.unlink(missing_ok=True)
+    raise SystemExit(returncode)
+
+
 def launch(state: dict, tool_args: list[str]) -> None:
     binary = SPEC["binary"]
     workspace = state.get("workspace")
     if state.get("claude_relayed"):
         _launch_relayed(state, binary, tool_args)
         return
+    # Experimental single-command launch with runtime model switching. Relayed is
+    # excluded for now (it already spawns-and-waits behind a loopback proxy). No real
+    # exec seam on Windows, so POSIX only.
+    if smart_routing_v2.enabled() and workspace and os.name != "nt":
+        _launch_smart_routing_v2(state, tool_args)
+        return
     if workspace:
         os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile"))
     exec_or_spawn(_build_claude_argv(binary, tool_args))

src/ucode/cli.py

@@ -1403,6 +1403,7 @@ def claude_router_hook_cmd(
     profile: Annotated[str | None, typer.Option("--profile")] = None,
     use_pat: Annotated[bool, typer.Option("--use-pat")] = False,
     model: Annotated[list[str] | None, typer.Option("--model")] = None,
+    socket_path: Annotated[str | None, typer.Option("--socket")] = None,
 ) -> None:
     """Run a Claude Code smart-routing lifecycle hook."""
     import json
@@ -1420,6 +1421,24 @@ def claude_router_hook_cmd(
         return
     if not isinstance(payload, dict):
         return
+    if event == "route-first-prompt":
+        if not socket_path:
+            from ucode.smart_routing.claude_hooks import FIRST_PROMPT_SOCKET_ENV
+
+            socket_path = os.environ.get(FIRST_PROMPT_SOCKET_ENV)
+        if not socket_path:
+            return
+        from pathlib import Path
+
+        from ucode.smart_routing.claude_pty import (
+            first_prompt_hook_output,
+            request_first_prompt_route,
+        )
+
+        output = first_prompt_hook_output(request_first_prompt_route(Path(socket_path), payload))
+        if output is not None:
+            sys.stdout.write(json.dumps(output))
+        return
     if event == "session-start":
         record_session_start(payload)
         return

src/ucode/smart_routing/claude_hooks.py

@@ -15,6 +15,8 @@
 from ucode.smart_routing import hooks
 
 ROUTING_HOOK_COMMAND_MARKER = "claude-router-hook"
+FIRST_PROMPT_HOOK_MARKER = "claude-router-hook route-first-prompt"
+FIRST_PROMPT_SOCKET_ENV = "UCODE_CLAUDE_V2_SOCKET"
 
 
 def sync_smart_routing_hooks(doc: dict, state: dict, *, enabled: bool) -> None:
@@ -28,6 +30,33 @@ def remove_smart_routing_hooks(doc: dict) -> bool:
     return hooks.remove_managed_hooks(doc, ROUTING_HOOK_COMMAND_MARKER)
 
 
+def sync_first_prompt_hook(doc: dict, executable: str) -> None:
+    """Add the per-launch first-prompt routing hook to a settings document.
+
+    The hook is written only to the temporary settings passed to the PTY-wrapped
+    Claude process.  Using a more specific marker than the persistent smart-routing
+    hooks preserves the existing SessionStart/PreToolUse/SubagentStart handlers.
+    """
+    argv = [
+        executable,
+        ROUTING_HOOK_COMMAND_MARKER,
+        "route-first-prompt",
+    ]
+    groups = {
+        "UserPromptSubmit": [
+            {
+                "hooks": [
+                    _routing_command_hook(
+                        argv,
+                        status="Selecting a model with Smart Routing",
+                    )
+                ]
+            }
+        ]
+    }
+    hooks.sync_managed_hooks(doc, FIRST_PROMPT_HOOK_MARKER, groups)
+
+
 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")