diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index af6faf1..5222e52 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -39,6 +39,7 @@ ) 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 ( remove_smart_routing_hooks, sync_smart_routing_hooks, @@ -48,10 +49,11 @@ from ucode.tracing import tracing_env from ucode.ui import print_err, print_note, print_success, print_warning +GATEWAY_MODEL_DISCOVERY_ENV_VAR = "ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY" CLAUDE_CONFIG_DIR = Path.home() / ".claude" CLAUDE_SETTINGS_PATH = CLAUDE_CONFIG_DIR / "ucode-settings.json" +CLAUDE_USER_SETTINGS_PATH = CLAUDE_CONFIG_DIR / "settings.json" CLAUDE_BACKUP_PATH = APP_DIR / "claude-ucode-settings.backup.json" -GATEWAY_MODEL_DISCOVERY_ENV_VAR = "ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY" SPEC: ToolSpec = { "binary": "claude", @@ -883,10 +885,40 @@ def _merge_claude_settings(base: dict, overlay: dict) -> dict: return merged +def _compose_v2_settings(tool_args: list[str]) -> tuple[dict, list[str]]: + """Compose caller settings with ucode's Claude settings for a v2 launch.""" + caller_values, remaining = _extract_caller_settings(tool_args) + settings: dict = {} + for value in caller_values: + settings = _merge_claude_settings(settings, _load_caller_settings(value)) + return _merge_claude_settings(settings, read_json_safe(CLAUDE_SETTINGS_PATH)), remaining + + +def _original_launch_model(snapshot: dict, state: dict) -> str | None: + override = state.get("_claude_launch_model") + if isinstance(override, str) and override.strip(): + return override.strip() + value = snapshot.get("value") if snapshot.get("present") is True else None + if isinstance(value, str) and value.strip(): + return value.strip() + return default_model(state) + + +def _has_explicit_model_arg(tool_args: list[str]) -> bool: + return any(arg in {"--model", "-m"} or arg.startswith("--model=") for arg in tool_args) + + +def _launch_model_args(tool_args: list[str], launch_model: str | None) -> list[str]: + if not launch_model or _has_explicit_model_arg(tool_args): + return [] + return ["--model", launch_model] + + def _build_claude_argv( binary: str, tool_args: list[str], relayed: bool = False, + launch_model: str | None = None, settings_override: dict | None = None, ) -> list[str]: """Build the ``claude`` argv, composing any caller ``--settings`` with @@ -910,11 +942,19 @@ def _build_claude_argv( filter through and shadow the subscription OAuth. """ source_args = ["--setting-sources", _RELAYED_SETTING_SOURCES] if relayed else [] + model_args = _launch_model_args(tool_args, launch_model) caller_values, remaining = _extract_caller_settings(tool_args) if not caller_values and settings_override is None: # No caller --settings: hand Claude ucode's settings file directly (the # common path; behavior unchanged). - return [binary, *source_args, "--settings", str(CLAUDE_SETTINGS_PATH), *tool_args] + return [ + binary, + *source_args, + "--settings", + str(CLAUDE_SETTINGS_PATH), + *model_args, + *tool_args, + ] caller_settings: dict = {} for value in caller_values: caller_settings = _merge_claude_settings(caller_settings, _load_caller_settings(value)) @@ -928,6 +968,7 @@ def _build_claude_argv( *source_args, "--settings", json.dumps(merged, separators=(",", ":")), + *model_args, *remaining, ] @@ -1036,7 +1077,10 @@ def _launch_relayed(state: dict, binary: str, tool_args: list[str]) -> None: raise SystemExit(returncode) -def _launch_gateway(state: dict, binary: str, tool_args: list[str]) -> None: +def _launch_gateway( + state: dict, binary: str, tool_args: list[str], launch_model: str | None +) -> None: + """Launch discovery-enabled Claude through a refreshing gateway proxy.""" workspace = state["workspace"] server, cache, client = start_anthropic_model_discovery_proxy( workspace, @@ -1053,11 +1097,14 @@ def _launch_gateway(state: dict, binary: str, tool_args: list[str]) -> None: server_thread = threading.Thread(target=server.serve_forever, daemon=True) server_thread.start() - settings_override = { - "env": {"ANTHROPIC_BASE_URL": os.environ["ANTHROPIC_BASE_URL"]}, - } + settings_override = {"env": {"ANTHROPIC_BASE_URL": os.environ["ANTHROPIC_BASE_URL"]}} proc = subprocess.Popen( - _build_claude_argv(binary, tool_args, settings_override=settings_override) + _build_claude_argv( + binary, + tool_args, + launch_model=launch_model, + settings_override=settings_override, + ) ) try: returncode = proc.wait() @@ -1074,15 +1121,37 @@ def _launch_gateway(state: dict, binary: str, tool_args: list[str]) -> None: def launch(state: dict, tool_args: list[str]) -> None: binary = SPEC["binary"] workspace = state.get("workspace") + # Recover a prior switch interrupted before its surgical settings restore. + smart_routing_v2.recover_claude_model_snapshots(CLAUDE_USER_SETTINGS_PATH) + model_snapshot = smart_routing_v2.snapshot_claude_model_setting(CLAUDE_USER_SETTINGS_PATH) + launch_model = _original_launch_model(model_snapshot, state) if state.get("claude_relayed"): _launch_relayed(state, binary, tool_args) return + # Smart routing v2 needs Unix PTY support, which Windows does not provide. + if smart_routing_v2.enabled() and os.name == "nt": + raise RuntimeError( + "Smart routing in Claude Code is currently not supported on Windows. " + "Please use Codex or disable smart routing." + ) + if smart_routing_v2.enabled() and workspace: + smart_routing_v2.launch_claude( + state, + tool_args, + binary=binary, + user_settings_path=CLAUDE_USER_SETTINGS_PATH, + model_snapshot=model_snapshot, + launch_model=launch_model, + compose_settings=_compose_v2_settings, + launch_model_args=_launch_model_args, + ) + return if workspace and os.environ.get(GATEWAY_MODEL_DISCOVERY_ENV_VAR) == "1": - _launch_gateway(state, binary, tool_args) + _launch_gateway(state, binary, tool_args, launch_model) return if workspace: os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile")) - exec_or_spawn(_build_claude_argv(binary, tool_args)) + exec_or_spawn(_build_claude_argv(binary, tool_args, launch_model=launch_model)) def validate_cmd(binary: str) -> list[str]: diff --git a/src/ucode/cli.py b/src/ucode/cli.py index c8d7118..d1ece25 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -107,6 +107,7 @@ download_managed_skills_on_launch, ) from ucode.smart_routing import claude_routing, codex_routing +from ucode.smart_routing import v2 as smart_routing_v2 from ucode.state import ( STATE_PATH, clear_state, @@ -1403,6 +1404,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 +1422,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 @@ -1862,7 +1882,12 @@ def _launch_tool( managed_launch_model(managed, recommendation, tool) if managed is not None else None ) state, resolved_model = resolve_launch_model(tool, state, managed_model) - if routing_agent is not None and routing_agent.smart_routing_enabled(state): + first_prompt_routes_claude = tool == "claude" and smart_routing_v2.enabled() + if ( + routing_agent is not None + and routing_agent.smart_routing_enabled(state) + and not first_prompt_routes_claude + ): display = TOOL_SPECS[tool]["display"] with spinner(f"Selecting a {display} model with smart routing..."): decision, routing_error = _ROUTING_MODULES[tool].route_launch_model( @@ -1957,6 +1982,13 @@ def _launch_tool( if managed is not None and not is_dry_run(): _register_managed_mcp_servers(managed, tool, state) _apply_managed_skills(managed, tool, state) + if tool == "claude": + # Transient launch precedence for claude.py's universal --model flag. + # An explicit choice wins, followed by a routed/managed root pick; + # neither value is persisted into workspace state. + launch_model = model or route_root_model + if launch_model: + state["_claude_launch_model"] = launch_model print_success(f"Starting {TOOL_SPECS[tool]['display']}") launch_agent(tool, state, ctx.args) except RuntimeError as exc: diff --git a/src/ucode/smart_routing/claude_hooks.py b/src/ucode/smart_routing/claude_hooks.py index 773b770..6bc0534 100644 --- a/src/ucode/smart_routing/claude_hooks.py +++ b/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,23 @@ 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 first-prompt hook to a per-launch settings document.""" + groups = { + "UserPromptSubmit": [ + { + "hooks": [ + _routing_command_hook( + [executable, ROUTING_HOOK_COMMAND_MARKER, "route-first-prompt"], + 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") diff --git a/src/ucode/smart_routing/claude_pty.py b/src/ucode/smart_routing/claude_pty.py new file mode 100644 index 0000000..efca9d6 --- /dev/null +++ b/src/ucode/smart_routing/claude_pty.py @@ -0,0 +1,428 @@ +"""PTY wrapper that routes Claude Code's first prompt with ``/model ``. + +Claude Code has no runtime control protocol, so a per-launch UserPromptSubmit +hook captures and blocks the first prompt. While the TUI is idle, this wrapper +types a direct model command, restores the user's persisted default model, and +replays the prompt. Terminal traffic is otherwise forwarded unchanged. +""" + +from __future__ import annotations + +import contextlib +import fcntl +import json +import os +import pty +import re +import select +import signal +import socket +import struct +import termios +import threading +import time +import tty +from collections.abc import Callable +from pathlib import Path + +MAX_MODEL_NAME_LEN = 200 +CONFIRM_TIMEOUT_S = 3.0 +SWITCH_TIMEOUT_S = 6.0 +READY_QUIET_S = 0.75 +SELECT_TIMEOUT_S = 0.2 +_MODEL_NAME_RE = re.compile(r"^[A-Za-z0-9._:/\-\[\]]+$") +_ANSI_RE = re.compile(rb"\x1b\[[0-9;?]*[ -/]*[@-~]|\x1b\][^\x07]*(?:\x07|\x1b\\)|\x1b[@-Z\\-_]") + + +def strip_ansi(data: bytes) -> str: + return _ANSI_RE.sub(b"", data).decode("utf-8", "replace") + + +def _squash(text: str) -> str: + return "".join(text.split()) + + +def _match_text(data: bytes) -> str: + return _squash(strip_ansi(data)) + + +def valid_model_name(name: object) -> bool: + """Return whether *name* is safe to type as one slash-command argument.""" + return ( + isinstance(name, str) + and 0 < len(name) <= MAX_MODEL_NAME_LEN + and bool(_MODEL_NAME_RE.fullmatch(name)) + ) + + +def switch_message(model: str, reason: str) -> str: + """Format the routed-model notice shown in Claude Code.""" + lines = [ + "Using Unity Gateway Smart Router.", + f"Selected Model : {model}", + f"Reason : {reason}", + ] + width = max(len(line) for line in lines) + border = "─" * (width + 2) + return "\n".join( + [f"┌{border}┐", *(f"│ {line:<{width}} │" for line in lines), f"└{border}┘"] + ) + + +class ConfirmationState: + """Detect and accept Claude's optional cache-cost confirmation dialog.""" + + PROMPT_MARKERS = ("Switch model?", "Yes, switch to", "No, go back") + + def __init__(self, window: int = 4096) -> None: + self._buf = "" + self._armed_until = 0.0 + self._window = window + + def arm(self, deadline: float) -> None: + self._armed_until = deadline + self._buf = "" + + def clear(self) -> None: + self._armed_until = 0.0 + self._buf = "" + + def observe(self, chunk: bytes, now: float) -> bytes | None: + if self._armed_until == 0.0: + return None + if now > self._armed_until: + self.clear() + return None + self._buf = (self._buf + _match_text(chunk))[-self._window :] + if all(_squash(marker) in self._buf for marker in self.PROMPT_MARKERS): + self.clear() + return b"\r" + return None + + +class OutputMarkerDetector: + """Latching ANSI-insensitive substring detector.""" + + def __init__(self, markers: tuple[str, ...], window: int = 4096) -> None: + self._markers = markers + self._buf = "" + self._window = window + self.triggered = False + + def observe(self, chunk: bytes) -> bool: + if self.triggered: + return True + self._buf = (self._buf + _match_text(chunk))[-self._window :] + if any(_squash(marker) in self._buf for marker in self._markers): + self.triggered = True + return self.triggered + + +def inject_model_switch(master_fd: int, model: str) -> None: + """Type Claude Code's direct, persistent model command.""" + if not valid_model_name(model): + raise ValueError(f"Unsafe Claude model name: {model!r}") + os.write(master_fd, f"/model {model}\r".encode()) + + +def inject_prompt(master_fd: int, prompt: str, *, submit: bool = True) -> None: + """Replay a captured prompt as one bracketed paste.""" + clean = prompt.replace("\r\n", "\n").replace("\r", "\n") + clean = clean.replace("\x00", "").replace("\x1b", "") + suffix = b"\r" if submit else b"" + os.write(master_fd, b"\x1b[200~" + clean.encode() + b"\x1b[201~" + suffix) + + +def inject_note(out_fd: int, message: str) -> None: + os.write(out_fd, ("\r\n\x1b[36m" + message + "\x1b[0m\r\n").encode()) + + +def request_first_prompt_route(path: Path, payload: dict, *, timeout: float = 5.0) -> dict | None: + """Send a UserPromptSubmit payload to the owning PTY wrapper.""" + prompt = payload.get("prompt") + if not isinstance(prompt, str) or not prompt.strip(): + return None + request = { + "method": "route_first_prompt", + "prompt": prompt, + "session_id": payload.get("session_id"), + } + try: + client = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + client.settimeout(timeout) + with client: + client.connect(str(path)) + client.sendall((json.dumps(request) + "\n").encode()) + with client.makefile("rb") as stream: + raw = stream.readline() + response = json.loads(raw) if raw else None + except (OSError, ValueError): + return None + return response if isinstance(response, dict) else None + + +def first_prompt_hook_output(response: dict | None) -> dict | None: + """Translate the wrapper response into Claude hook output.""" + if not isinstance(response, dict) or response.get("action") != "block": + return None + model = response.get("model") + if not valid_model_name(model): + return None + return { + "decision": "block", + "reason": switch_message( + model, "Low complexity, unclear intent, and no code reference." + ), + } + + +def serve_first_prompt_socket( + path: Path, + route_prompt: Callable[[str], str], + on_blocked_prompt: Callable[[str, str], None], + stop: threading.Event, + *, + log: Callable[[str], None] = lambda _message: None, +) -> threading.Thread: + """Serve the hook protocol, blocking exactly one non-command prompt.""" + + def serve() -> None: + claimed = False + try: + path.unlink(missing_ok=True) + server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + server.bind(str(path)) + os.chmod(path, 0o600) + server.listen(4) + server.settimeout(0.5) + except OSError as exc: + log(f"[ERR] first-prompt socket bind failed: {exc!r}") + return + log(f"[READY] first-prompt socket {path}") + try: + while not stop.is_set(): + try: + conn, _ = server.accept() + except TimeoutError: + continue + except OSError: + break + with conn, conn.makefile("rwb") as stream: + response: dict = {"action": "allow"} + blocked: tuple[str, str] | None = None + try: + request = json.loads(stream.readline()) + prompt = request.get("prompt") if isinstance(request, dict) else None + is_route = ( + isinstance(request, dict) + and request.get("method") == "route_first_prompt" + ) + is_command = isinstance(prompt, str) and prompt.lstrip().startswith("/") + if ( + is_route + and not claimed + and isinstance(prompt, str) + and prompt.strip() + and not is_command + ): + model = route_prompt(prompt) + if valid_model_name(model): + claimed = True + response = {"action": "block", "model": model} + blocked = (prompt, model) + except Exception as exc: # noqa: BLE001 - hooks must fail open + log(f"[ERR] first-prompt request: {exc!r}") + stream.write((json.dumps(response) + "\n").encode()) + stream.flush() + if blocked is not None: + on_blocked_prompt(*blocked) + finally: + server.close() + + thread = threading.Thread(target=serve, name="claude-first-prompt", daemon=True) + thread.start() + return thread + + +class TerminalModeGuard: + """Put stdin in raw mode for the PTY session and restore it on exit.""" + + def __init__(self, fd: int = 0) -> None: + self.fd = fd + self._saved: list | None = None + + def __enter__(self) -> TerminalModeGuard: + if os.isatty(self.fd): + self._saved = termios.tcgetattr(self.fd) + tty.setraw(self.fd) + return self + + def __exit__(self, *_exc: object) -> None: + if self._saved is not None: + termios.tcsetattr(self.fd, termios.TCSADRAIN, self._saved) + self._saved = None + + +def sync_winsize(master_fd: int, stdin_fd: int = 0) -> None: + if not os.isatty(stdin_fd): + return + try: + packed = fcntl.ioctl(stdin_fd, termios.TIOCGWINSZ, struct.pack("HHHH", 0, 0, 0, 0)) + fcntl.ioctl(master_fd, termios.TIOCSWINSZ, packed) + except OSError: + pass + + +def run_claude_pty( + argv: list[str], + *, + route_prompt: Callable[[str], str], + switch_message: str, + socket_path: Path, + log_path: Path | None = None, +) -> int: + """Run Claude in a PTY, switch its model, and replay the first prompt.""" + + def log(message: str) -> None: + if log_path is None: + return + try: + with open(log_path, "a", encoding="utf-8") as handle: + handle.write(f"{time.strftime('%H:%M:%S')} {message}\n") + except OSError: + pass + + debug = os.environ.get("UCODE_CLAUDE_PTY_DEBUG") == "1" + gate_read, gate_write = os.pipe() + pid, master_fd = pty.fork() + if pid == 0: + os.close(gate_write) + try: + os.read(gate_read, 1) + finally: + os.close(gate_read) + os.execvp(argv[0], argv) + os._exit(127) + os.close(gate_read) + + confirm = ConfirmationState() + stop = threading.Event() + pending_lock = threading.Lock() + pending: dict[str, tuple[str, str] | None] = {"value": None} + + def on_blocked_prompt(prompt: str, model: str) -> None: + with pending_lock: + pending["value"] = (prompt, model) + log(f"[ROUTE] first prompt -> {model!r}") + + server_thread = serve_first_prompt_socket( + socket_path, route_prompt, on_blocked_prompt, stop, log=log + ) + socket_deadline = time.monotonic() + 2.0 + while ( + not socket_path.exists() and server_thread.is_alive() and time.monotonic() < socket_deadline + ): + time.sleep(0.01) + if not socket_path.exists(): + log("[ERR] first-prompt socket was not ready before Claude launch") + os.write(gate_write, b"1") + os.close(gate_write) + + previous_winch = signal.getsignal(signal.SIGWINCH) + + def on_winch(_signum: int, _frame: object) -> None: + sync_winsize(master_fd) + + try: + with TerminalModeGuard(0): + signal.signal(signal.SIGWINCH, on_winch) + sync_winsize(master_fd) + stdin_open = True + last_output = 0.0 + phase = "waiting_prompt" + routed_prompt = "" + routed_model = "" + switch_started = 0.0 + switch_complete: OutputMarkerDetector | None = None + while True: + readable = [master_fd, 0] if stdin_open else [master_fd] + try: + ready_fds, _, _ = select.select(readable, [], [], SELECT_TIMEOUT_S) + except InterruptedError: + continue + + if 0 in ready_fds: + try: + data = os.read(0, 4096) + except OSError: + data = b"" + if not data: + stdin_open = False + else: + os.write(master_fd, data) + + if master_fd in ready_fds: + try: + chunk = os.read(master_fd, 8192) + except OSError: + chunk = b"" + if not chunk: + break + os.write(1, chunk) + last_output = time.monotonic() + if debug: + log(f"[OUT] {strip_ansi(chunk)[:400]!r}") + keystroke = confirm.observe(chunk, last_output) + if keystroke is not None: + os.write(master_fd, keystroke) + if phase == "switching" and switch_complete is not None: + switch_complete.observe(chunk) + + if phase == "waiting_prompt": + with pending_lock: + captured = pending["value"] + if captured is not None: + routed_prompt, routed_model = captured + phase = "waiting_to_switch" + + now = time.monotonic() + idle = last_output > 0.0 and now - last_output >= READY_QUIET_S + if phase == "waiting_to_switch" and idle: + inject_note(1, switch_message) + inject_model_switch(master_fd, routed_model) + confirm.arm(now + CONFIRM_TIMEOUT_S) + switch_complete = OutputMarkerDetector(("Set model to", "Model set to")) + switch_started = now + phase = "switching" + log(f"[SWITCH] /model {routed_model}") + elif ( + phase == "switching" + and switch_complete is not None + and switch_complete.triggered + ): + inject_prompt(master_fd, routed_prompt) + phase = "done" + log("[REPLAY] first prompt submitted") + elif phase == "switching" and now - switch_started >= SWITCH_TIMEOUT_S: + os.write(master_fd, b"\x1b") + inject_note( + 1, + "Smart Routing could not confirm the model switch. " + "Your prompt was restored but not submitted.", + ) + inject_prompt(master_fd, routed_prompt, submit=False) + phase = "failed" + log(f"[ERR] direct model switch timed out for {routed_model!r}") + finally: + signal.signal(signal.SIGWINCH, previous_winch) + stop.set() + with contextlib.suppress(OSError): + os.close(master_fd) + socket_path.unlink(missing_ok=True) + + _waited_pid, status = os.waitpid(pid, 0) + if os.WIFEXITED(status): + return os.WEXITSTATUS(status) + if os.WIFSIGNALED(status): + return 128 + os.WTERMSIG(status) + return 1 diff --git a/src/ucode/smart_routing/v2.py b/src/ucode/smart_routing/v2.py index af3e4c3..a6541b8 100644 --- a/src/ucode/smart_routing/v2.py +++ b/src/ucode/smart_routing/v2.py @@ -7,25 +7,33 @@ import sys import time import urllib.request +import uuid from collections.abc import Callable +from pathlib import Path from typing import NoReturn import tomlkit -from ucode.config_io import APP_DIR -from ucode.databricks import get_databricks_token -from ucode.smart_routing import codex_interposer +from ucode.config_io import APP_DIR, read_json_safe, write_json_file +from ucode.constants import LOOPBACK_HOST +from ucode.databricks import build_auth_token_argv, get_databricks_token +from ucode.smart_routing import claude_pty, codex_interposer +from ucode.smart_routing.claude_hooks import FIRST_PROMPT_SOCKET_ENV, sync_first_prompt_hook +from ucode.ui import print_note ENV_VAR = "ENABLE_SMART_ROUTING_V2" CODEX_TARGET_MODEL = "system.ai.glm-5-2" # TODO(lilly): replace with smart router. CODEX_INTERPOSER_LOG = APP_DIR / "codex-v2-interposer.log" -CODEX_SWITCH_REASON = "Low complexity, unclear intent, and no code reference." # TODO(lilly): replace with smart router rationale. +STUBBED_SWITCH_REASON = "Low complexity, unclear intent, and no code reference." # TODO(lilly): replace with smart router rationale. + +CLAUDE_TARGET_MODEL = "system.ai.claude-sonnet-4-6[1m]" # TODO(lilly): replace with smart router. +CLAUDE_PTY_LOG = APP_DIR / "claude-v2-pty.log" +CLAUDE_MODEL_SNAPSHOT_PATH = APP_DIR / "claude-default-model.snapshot.json" APP_SERVER_READY_TIMEOUT_SECONDS = 30 PROCESS_SHUTDOWN_TIMEOUT_SECONDS = 5 OAUTH_TOKEN_ENV_VAR = "OAUTH_TOKEN" -LOOPBACK_HOST = "127.0.0.1" HEALTH_REQUEST_TIMEOUT_SECONDS = 1 HEALTH_POLL_INTERVAL_SECONDS = 0.25 @@ -34,6 +42,48 @@ def enabled() -> bool: return os.environ.get(ENV_VAR) == "1" +def snapshot_claude_model_setting(user_settings_path: Path) -> dict: + """Capture only Claude's user-level ``model`` setting.""" + settings = read_json_safe(user_settings_path) + return {"present": "model" in settings, "value": settings.get("model")} + + +def _save_claude_model_snapshot(snapshot: dict, snapshot_path: Path) -> None: + """Journal the pre-switch model so the next launch can recover after a crash.""" + write_json_file(snapshot_path, snapshot) + + +def restore_claude_model_snapshot( + user_settings_path: Path, + snapshot_path: Path | None = None, + snapshot: dict | None = None, +) -> bool: + """Restore only the journaled ``model`` field, preserving every sibling setting.""" + snapshot_path = snapshot_path or CLAUDE_MODEL_SNAPSHOT_PATH + if snapshot is None and not snapshot_path.exists(): + return False + snapshot = snapshot if snapshot is not None else read_json_safe(snapshot_path) + present = snapshot.get("present") + if not isinstance(present, bool): + raise RuntimeError(f"Claude model recovery snapshot is invalid: {snapshot_path}") + settings = read_json_safe(user_settings_path) + if present: + settings["model"] = snapshot.get("value") + else: + settings.pop("model", None) + write_json_file(user_settings_path, settings) + snapshot_path.unlink(missing_ok=True) + return True + + +def recover_claude_model_snapshots(user_settings_path: Path) -> None: + """Repair defaults left by interrupted Claude PTY launches.""" + candidates = {CLAUDE_MODEL_SNAPSHOT_PATH} + candidates.update(APP_DIR.glob("claude-default-model.*.snapshot.json")) + for path in sorted(candidates): + restore_claude_model_snapshot(user_settings_path, path) + + def _loopback_websocket_url(port: int) -> str: return f"ws://{LOOPBACK_HOST}:{port}" @@ -70,6 +120,81 @@ def _switch_message(model: str, reason: str) -> str: return "\n".join([f"┌{border}┐", *(f"│ {line:<{width}} │" for line in lines), f"└{border}┘"]) +def _route_claude_prompt(_prompt: str) -> str: + return CLAUDE_TARGET_MODEL + + +def launch_claude( + state: dict, + tool_args: list[str], + *, + binary: str, + user_settings_path: Path, + model_snapshot: dict, + launch_model: str | None, + compose_settings: Callable[[list[str]], tuple[dict, list[str]]], + launch_model_args: Callable[[list[str], str | None], list[str]], +) -> NoReturn: + """Launch Claude in the first-prompt routing PTY wrapper.""" + workspace = state.get("workspace") + if not workspace: + raise RuntimeError( + "Smart routing v2 needs a configured workspace; run `ucode configure claude` first." + ) + os.environ[OAUTH_TOKEN_ENV_VAR] = 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" + model_snapshot_path = APP_DIR / f"claude-default-model.{run_id}.snapshot.json" + + settings, remaining = compose_settings(tool_args) + 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) + model_args = launch_model_args(remaining, launch_model) + argv = [binary, "--settings", str(settings_path), *model_args, *remaining] + + _save_claude_model_snapshot(model_snapshot, model_snapshot_path) + restored = False + + def restore_default_model() -> bool: + nonlocal restored + if restored: + return False + result = restore_claude_model_snapshot( + user_settings_path, model_snapshot_path, model_snapshot + ) + restored = True + return result + + print_note( + "Smart routing v2: the first submitted prompt will select Claude Code's " + f"model ({CLAUDE_TARGET_MODEL}); log: {CLAUDE_PTY_LOG}." + ) + try: + returncode = claude_pty.run_claude_pty( + argv, + route_prompt=_route_claude_prompt, + switch_message=claude_pty.switch_message( + CLAUDE_TARGET_MODEL, STUBBED_SWITCH_REASON + ), + socket_path=socket_path, + log_path=CLAUDE_PTY_LOG, + ) + finally: + restore_default_model() + settings_path.unlink(missing_ok=True) + socket_path.unlink(missing_ok=True) + sys.exit(returncode) + + def _toml_value(value: str | int | float | bool | list[object] | dict[str, object]) -> str: if isinstance(value, dict): item = tomlkit.inline_table() @@ -144,7 +269,7 @@ def launch_codex( LOOPBACK_HOST, app_server_url, CODEX_TARGET_MODEL, - switch_message=_switch_message(CODEX_TARGET_MODEL, CODEX_SWITCH_REASON), + switch_message=_switch_message(CODEX_TARGET_MODEL, STUBBED_SWITCH_REASON), log_path=CODEX_INTERPOSER_LOG, ) tui_url = _loopback_websocket_url(tui_port) diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 112ea05..c3c9bf2 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -9,7 +9,7 @@ import pytest from ucode.agents import claude -from ucode.smart_routing import claude_routing +from ucode.smart_routing import claude_routing, v2 WS = "https://example.databricks.com" @@ -654,21 +654,44 @@ def boom(name, entry, scope=mcp_mod.MCP_USER_SCOPE): class TestClaudeLaunch: + def test_smart_routing_on_windows_is_not_supported(self, monkeypatch): + monkeypatch.setenv(v2.ENV_VAR, "1") + monkeypatch.setattr(claude.os, "name", "nt") + monkeypatch.setattr(v2, "recover_claude_model_snapshots", lambda _path: None) + monkeypatch.setattr(v2, "snapshot_claude_model_setting", lambda _path: {"present": False}) + + with pytest.raises( + RuntimeError, + match="Smart routing in Claude Code is currently not supported on Windows", + ): + claude.launch({"workspace": WS, "profile": "test"}, ["--debug"]) + def test_default_launch_keeps_existing_auth_path(self, monkeypatch): calls: list[list[str]] = [] monkeypatch.delenv(claude.GATEWAY_MODEL_DISCOVERY_ENV_VAR, raising=False) monkeypatch.delenv("OAUTH_TOKEN", raising=False) + monkeypatch.setattr( + v2, "recover_claude_model_snapshots", lambda _path: None + ) + monkeypatch.setattr( + v2, "snapshot_claude_model_setting", lambda _path: {"present": True, "value": "opus"} + ) monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "token") monkeypatch.setattr(claude, "exec_or_spawn", lambda argv: calls.append(argv)) claude.launch({"workspace": WS, "profile": "test"}, ["--debug"]) assert os.environ["OAUTH_TOKEN"] == "token" - assert calls == [["claude", "--settings", str(claude.CLAUDE_SETTINGS_PATH), "--debug"]] + assert calls == [ + ["claude", "--settings", str(claude.CLAUDE_SETTINGS_PATH), "--model", "opus", "--debug"] + ] def test_gateway_discovery_uses_anthropic_proxy(self, monkeypatch): calls: list[tuple] = [] + monkeypatch.setattr(v2, "recover_claude_model_snapshots", lambda _path: None) + monkeypatch.setattr(v2, "snapshot_claude_model_setting", lambda _path: {"present": False}) + class Server: server_address = ("127.0.0.1", 12345) diff --git a/tests/test_claude_smart_routing_v2.py b/tests/test_claude_smart_routing_v2.py new file mode 100644 index 0000000..0a1dcbf --- /dev/null +++ b/tests/test_claude_smart_routing_v2.py @@ -0,0 +1,219 @@ +"""Tests for Claude's experimental first-prompt PTY routing path.""" + +from __future__ import annotations + +import json +import os +import sys +import threading +import time +from pathlib import Path + +import pytest + +from ucode.agents import claude +from ucode.smart_routing import claude_hooks, claude_pty, v2 + + +class TestDirectModelCommand: + @pytest.mark.parametrize( + "name", + ["system.ai.claude-opus-4-8[1m]", "databricks-claude-sonnet-5", "opus"], + ) + def test_accepts_model_names(self, name): + assert claude_pty.valid_model_name(name) + + @pytest.mark.parametrize("name", ["", "a b", "a\nb", "x" * 201, None]) + def test_rejects_unsafe_model_names(self, name): + assert not claude_pty.valid_model_name(name) + + def test_types_direct_model_command(self): + read_fd, write_fd = os.pipe() + try: + claude_pty.inject_model_switch(write_fd, "system.ai.claude-sonnet-5") + assert os.read(read_fd, 200) == b"/model system.ai.claude-sonnet-5\r" + finally: + os.close(read_fd) + os.close(write_fd) + + +class TestFirstPromptHook: + def test_renders_boxed_router_notice(self): + model = "system.ai.claude-sonnet-4-6[1m]" + reason = "Low complexity, unclear intent, and no code reference." + result = claude_pty.first_prompt_hook_output( + {"action": "block", "model": model} + ) + + assert result == {"decision": "block", "reason": v2._switch_message(model, reason)} + assert claude_pty.switch_message(model, reason) == v2._switch_message(model, reason) + + def test_blocks_once_then_allows_replay(self, tmp_path): + socket_path = tmp_path / "first.sock" + blocked: list[tuple[str, str]] = [] + stop = threading.Event() + claude_pty.serve_first_prompt_socket( + socket_path, + lambda _prompt: "sonnet", + lambda prompt, model: blocked.append((prompt, model)), + stop, + ) + try: + deadline = time.monotonic() + 5 + while not socket_path.exists() and time.monotonic() < deadline: + time.sleep(0.01) + first = claude_pty.request_first_prompt_route( + socket_path, {"session_id": "s1", "prompt": "fix the parser"} + ) + replay = claude_pty.request_first_prompt_route( + socket_path, {"session_id": "s1", "prompt": "fix the parser"} + ) + assert first == {"action": "block", "model": "sonnet"} + assert replay == {"action": "allow"} + assert blocked == [("fix the parser", "sonnet")] + finally: + stop.set() + + def test_first_prompt_hook_is_per_launch(self): + settings = {"hooks": {"PreToolUse": [{"hooks": [{"command": "user-policy"}]}]}} + claude_hooks.sync_first_prompt_hook(settings, "/bin/ucode") + claude_hooks.sync_first_prompt_hook(settings, "/bin/ucode") + command = settings["hooks"]["UserPromptSubmit"][0]["hooks"][0]["command"] + assert command == "/bin/ucode claude-router-hook route-first-prompt" + assert len(settings["hooks"]["UserPromptSubmit"]) == 1 + assert "user-policy" in str(settings["hooks"]["PreToolUse"]) + + +class TestV2Launch: + def test_saves_default_passes_model_and_restores_only_model(self, tmp_path, monkeypatch): + ucode_settings = tmp_path / "ucode-settings.json" + user_settings = tmp_path / "settings.json" + ucode_settings.write_text(json.dumps({"env": {"ANTHROPIC_BASE_URL": "https://gw"}})) + user_settings.write_text(json.dumps({"model": "opus", "theme": "dark"})) + monkeypatch.setattr(claude, "APP_DIR", tmp_path) + monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", ucode_settings) + monkeypatch.setattr(claude, "CLAUDE_USER_SETTINGS_PATH", user_settings) + monkeypatch.setattr(v2, "APP_DIR", tmp_path) + monkeypatch.setattr(v2, "CLAUDE_PTY_LOG", tmp_path / "v2.log") + monkeypatch.setattr(v2, "get_databricks_token", lambda *_args, **_kwargs: "token") + monkeypatch.setattr(v2, "build_auth_token_argv", lambda *_args, **_kwargs: ["ucode"]) + captured: dict = {} + + def fake_run(argv, **kwargs): + captured["argv"] = argv + generated = Path(argv[argv.index("--settings") + 1]) + captured["settings"] = json.loads(generated.read_text()) + # Simulate `/model` changing the user file, plus an unrelated concurrent edit. + user_settings.write_text(json.dumps({"model": "routed", "theme": "light", "new": True})) + return 0 + + monkeypatch.setattr(claude_pty, "run_claude_pty", fake_run) + snapshot = v2.snapshot_claude_model_setting(user_settings) + with pytest.raises(SystemExit) as exc: + v2.launch_claude( + {"workspace": "https://example.com"}, + ["--debug"], + binary="claude", + user_settings_path=user_settings, + model_snapshot=snapshot, + launch_model="opus", + compose_settings=claude._compose_v2_settings, + launch_model_args=claude._launch_model_args, + ) + + assert exc.value.code == 0 + assert captured["argv"][-3:] == ["--model", "opus", "--debug"] + assert claude_hooks.FIRST_PROMPT_SOCKET_ENV in captured["settings"]["env"] + assert "modelPicker" not in captured["settings"] + assert json.loads(user_settings.read_text()) == { + "model": "opus", + "theme": "light", + "new": True, + } + assert not list(tmp_path.glob("claude-default-model.*.snapshot.json")) + + +class TestModelRecovery: + def test_restore_changes_only_model_field(self, tmp_path): + user = tmp_path / "settings.json" + snapshot_path = tmp_path / "model-snapshot.json" + user.write_text(json.dumps({"model": "opus", "theme": "dark"})) + original = v2.snapshot_claude_model_setting(user) + v2._save_claude_model_snapshot(original, snapshot_path) + + user.write_text(json.dumps({"model": "routed", "theme": "light", "new": True})) + assert v2.restore_claude_model_snapshot(user, snapshot_path) is True + assert json.loads(user.read_text()) == { + "model": "opus", + "theme": "light", + "new": True, + } + assert v2.restore_claude_model_snapshot(user, snapshot_path) is False + + def test_restore_removes_model_when_original_was_absent(self, tmp_path): + user = tmp_path / "settings.json" + snapshot_path = tmp_path / "model-snapshot.json" + user.write_text(json.dumps({"theme": "dark"})) + v2._save_claude_model_snapshot(v2.snapshot_claude_model_setting(user), snapshot_path) + user.write_text(json.dumps({"model": "routed", "theme": "light"})) + + v2.restore_claude_model_snapshot(user, snapshot_path) + assert json.loads(user.read_text()) == {"theme": "light"} + + +class TestPtyFlow: + def test_direct_switch_restore_and_replay(self, tmp_path): + fake_claude = tmp_path / "fake_claude.py" + capture = tmp_path / "capture.json" + socket_path = tmp_path / "first.sock" + fake_claude.write_text( + """ +import json +import os +import socket +import sys +import tty +from pathlib import Path + +socket_path = sys.argv[1] +capture_path = Path(sys.argv[2]) +client = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) +client.connect(socket_path) +client.sendall((json.dumps({ + "method": "route_first_prompt", + "prompt": "fix\\nthe parser", + "session_id": "s1", +}) + "\\n").encode()) +response = client.makefile("rb").readline() +client.close() +assert json.loads(response)["action"] == "block" +print("Smart Routing blocked the prompt", flush=True) +tty.setraw(0) + +def read_until(suffix): + data = b"" + while not data.endswith(suffix): + data += os.read(0, 1) + return data + +model_command = read_until(b"\\r") +print("Set model to system.ai.claude-sonnet-5", flush=True) +replayed = read_until(b"\\x1b[201~\\r") +capture_path.write_text(json.dumps({ + "command": model_command.decode(), + "replayed": replayed.decode(), +})) +""".lstrip() + ) + result = claude_pty.run_claude_pty( + [sys.executable, str(fake_claude), str(socket_path), str(capture)], + route_prompt=lambda _prompt: "system.ai.claude-sonnet-5", + switch_message="router selected sonnet", + socket_path=socket_path, + ) + + assert result == 0 + assert json.loads(capture.read_text()) == { + "command": "/model system.ai.claude-sonnet-5\r", + "replayed": "\x1b[200~fix\nthe parser\x1b[201~\r", + } diff --git a/tests/test_cli.py b/tests/test_cli.py index 99b99dc..1dabb3b 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -280,6 +280,34 @@ def test_enabled_codex_launch_uses_routed_root_model(self): in _strip_ansi(result.output) ) + def test_claude_v2_skips_legacy_prelaunch_routing(self, monkeypatch): + monkeypatch.setenv("ENABLE_SMART_ROUTING_V2", "1") + state = { + **MINIMAL_STATE, + "smart_routing_enabled": True, + "claude_models": {"opus": "system.ai.claude-opus-4-8"}, + } + with ( + patch("ucode.cli.ensure_bootstrap_dependencies"), + patch("ucode.cli.load_state", return_value=state), + patch("ucode.cli.ensure_provider_state", return_value=state), + patch("ucode.cli.configure_shared_state", return_value=state), + patch( + "ucode.cli.resolve_launch_model", + return_value=(state, "system.ai.claude-opus-4-8"), + ), + patch("ucode.cli.claude_routing.route_launch_model") as mock_route, + patch("ucode.cli.configure_tool", return_value=state) as mock_configure, + patch("ucode.cli._fetch_managed_config", return_value=(None, False)), + patch("ucode.cli.launch_agent") as mock_launch, + ): + result = runner.invoke(app, ["claude"]) + + assert result.exit_code == 0, result.output + mock_route.assert_not_called() + assert mock_configure.call_args.kwargs["route_root_model"] is None + assert "_claude_launch_model" not in mock_launch.call_args.args[1] + class TestClaudeModelFlag: """`ucode claude --model ` pins the id into the family aliases so the gateway resolves any