diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 9f4df255..c76cda43 100644 --- a/src/ucode/agents/claude.py +++ b/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-4-6[1m]" # 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} due to " + "low complexity, unclear intent, and no code reference." + ), + 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)) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index ca71dd6a..0b1b24aa 100644 --- a/src/ucode/cli.py +++ b/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 diff --git a/src/ucode/smart_routing/claude_hooks.py b/src/ucode/smart_routing/claude_hooks.py index 773b7707..dd96baaf 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,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") diff --git a/src/ucode/smart_routing/claude_pty.py b/src/ucode/smart_routing/claude_pty.py new file mode 100644 index 00000000..e7aa129b --- /dev/null +++ b/src/ucode/smart_routing/claude_pty.py @@ -0,0 +1,720 @@ +"""PTY wrapper for Claude Code's TUI (smart routing v2). + +Claude Code has no ``app-server``/JSON-RPC seam like Codex, so there is nothing to +interpose on the wire. Instead this module runs the real ``claude`` TUI inside a PTY: +it forwards stdin<->master and master<->stdout untouched, and drives a *model switch* +through Claude Code's ``/model`` picker, using its ``s`` (session-only) action, and +auto-confirming the optional "Switch model?" cache dialog by watching the PTY output. + +``ucode.agents.claude`` owns the lifecycle: it enters this from the single ``ucode claude`` +command when ``ENABLE_SMART_ROUTING_V2=1``. Logs go to ``log_path`` (appended) only — never +stdout/stderr, which the foreground TUI owns (same discipline as ``codex_interposer``). + +The first prompt comes from a ``UserPromptSubmit`` hook over an owner-only Unix socket. +The hook blocks that one submission, allowing the wrapper to type ``/model`` while the +TUI is idle and then replay the exact prompt. A second hook invocation (the replay) is +allowed through. This makes the selected model the real Claude Code session model before +the first inference request, rather than merely rewriting the request below the client. +""" + +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 + +# Output markers are whitespace-squashed because Claude Code renders spaces as cursor +# moves that vanish under strip_ansi. READY_MARKERS remains useful to runtime probes/tests; +# the first-prompt path itself waits for the hook-blocked TUI to go idle. +READY_MARKERS = ("? for shortcuts", "Welcome to Claude Code") + +MAX_MODEL_NAME_LEN = 200 +CONFIRM_TIMEOUT_S = 3.0 +SWITCH_TIMEOUT_S = 6.0 +# Once the TUI stays quiet this long after rendering the hook block/model result, it is +# idle and safe for the wrapper to type the next command. +READY_QUIET_S = 0.75 +# select() wake interval while waiting to switch, so the quiet period is observable. +SELECT_TIMEOUT_S = 0.2 +_MODEL_NAME_RE = re.compile(r"^[A-Za-z0-9._:/\-\[\]]+$") +_CLAUDE_MODEL_RE = re.compile( + r"^(?:system\.ai\.)?claude-(opus|sonnet|haiku)-(\d+)(?:-(\d+))?(\[1m\])?$", + re.IGNORECASE, +) +# CSI/OSC/simple escape sequences — enough to make substring matching robust across +# the styled bytes Claude Code's Ink renderer emits. +_ANSI_RE = re.compile(rb"\x1b\[[0-9;?]*[ -/]*[@-~]|\x1b\][^\x07]*(?:\x07|\x1b\\)|\x1b[@-Z\\-_]") + + +def strip_ansi(data: bytes) -> str: + """Drop ANSI escape sequences and decode to text for substring matching.""" + return _ANSI_RE.sub(b"", data).decode("utf-8", "replace") + + +def _squash(text: str) -> str: + """Remove ALL whitespace. Claude Code's TUI positions words with cursor-move + escapes rather than literal spaces, so after stripping ANSI the words run + together (e.g. 'esc to interrupt' -> 'esctointerrupt'). Squashing both the + observed text and the markers makes multi-word markers match regardless.""" + return "".join(text.split()) + + +def _match_text(data: bytes) -> str: + """ANSI-stripped, whitespace-squashed text for robust marker matching.""" + return _squash(strip_ansi(data)) + + +def valid_model_name(name: object) -> bool: + """Guard the value before it is typed into a ``/model`` command. + + Rejects anything that could break out of the slash command (spaces, ``;``, + carriage returns/newlines, control chars) or is absurdly long. Accepts the + gateway/UC ids Claude Code's picker shows, e.g. ``system.ai.claude-opus-4-8[1m]``. + """ + return ( + isinstance(name, str) + and 0 < len(name) <= MAX_MODEL_NAME_LEN + and bool(_MODEL_NAME_RE.match(name)) + ) + + +def model_picker_labels(model: str) -> tuple[str, ...]: + """Return raw and friendly labels Claude Code may render for a model. + + Claude Code 2.1.243 changed gateway-backed picker rows from raw endpoint IDs such + as ``system.ai.claude-sonnet-5`` to friendly labels such as ``Sonnet 5``. Keep the + raw ID first for older versions and add the derived label for newer versions. + """ + labels = [model] + match = _CLAUDE_MODEL_RE.fullmatch(model) + if match is not None: + family, major, minor, long_context = match.groups() + version = major if minor is None else f"{major}.{minor}" + friendly = f"{family.title()} {version}" + if long_context: + friendly += " (1M)" + labels.append(friendly) + return tuple(labels) + + +class ConfirmationState: + """Watch PTY output for Claude's "Switch model?" dialog and auto-confirm it. + + ``arm`` after choosing the model with ``s`` (with a deadline), feed each output chunk to + ``observe``; when the complete confirmation prompt is seen it returns Enter and + disarms. The session-only choice has already been made with ``s`` in the model + picker; this dialog only acknowledges the cache cost of changing models. A rolling + buffer keeps the last ``window`` chars so + a marker split across two PTY reads (or interrupted by an ANSI escape) still matches. + """ + + # Waiting for the final option is important: reacting to the title alone can send + # Enter while Ink is still mounting the dialog, which leaks into the preceding UI. + 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 substring detector over an ANSI-stripped rolling buffer. + + Returns ``True`` from ``observe`` once any configured marker has been seen, and + stays ``True`` thereafter (``triggered``). Used for TUI-readiness and turn-start + detection where there is no JSON-RPC frame to key off. + """ + + 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 + + +class ModelPickerRows: + """Discover the routed and currently focused row numbers from picker output.""" + + def __init__(self, model: str, window: int = 16384) -> None: + self._targets = tuple(_squash(label) for label in model_picker_labels(model)) + self._buf = "" + self._window = window + self.target_row: int | None = None + self.focused_row: int | None = None + + def observe(self, chunk: bytes) -> None: + self._buf = (self._buf + _match_text(chunk))[-self._window :] + targets = [ + match + for target in self._targets + for match in re.finditer(rf"(\d+)\.{re.escape(target)}", self._buf) + ] + focused = list(re.finditer(r"❯(\d+)\.", self._buf)) + if targets: + self.target_row = int(targets[-1].group(1)) + if focused: + self.focused_row = int(focused[-1].group(1)) + + @property + def navigation(self) -> bytes | None: + """Arrow keys required to move from the focused row to the routed row.""" + if self.target_row is None or self.focused_row is None: + return None + delta = self.target_row - self.focused_row + key = b"\x1b[B" if delta > 0 else b"\x1b[A" + return key * abs(delta) + + +def inject_model_switch(master_fd: int) -> None: + """Open Claude Code's model picker. + + Passing the model directly (``/model ``) always persists it in interactive + Claude Code. Only the full picker exposes the ``s`` session-only action. + """ + os.write(master_fd, b"/model\r") + + +def inject_prompt(master_fd: int, prompt: str, *, submit: bool = True) -> None: + """Replay a hook-captured prompt using terminal bracketed-paste mode. + + Bracketed paste preserves multiline input and prevents embedded newlines from + submitting partial prompts. Escape/NUL bytes cannot be meaningful prompt text here + and are removed so captured content cannot terminate the paste envelope. + """ + 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: + """Splice a wrapper-authored, cyan note into the output stream (see module docstring §5). + + The wrapper owns stdout, so it can write its own bytes; timing this at the readiness + boundary (before the first prompt) lands it in static scroll-back above the input box. + """ + os.write(out_fd, ("\r\n\x1b[36m" + message + "\x1b[0m\r\n").encode()) + + +# --- first-prompt hook channel -------------------------------------------------------- + + +def request_first_prompt_route(path: Path, payload: dict, *, timeout: float = 5.0) -> dict | None: + """Send a ``UserPromptSubmit`` payload to the PTY wrapper. + + Hook failures deliberately fail open: returning ``None`` makes the hook emit no + blocking decision, so Claude processes the user's prompt on its existing model. + """ + 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's UserPromptSubmit 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": ( + f"✨ Smart Router selected {model} due to 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 _m: None, +) -> threading.Thread: + """Serve the hook protocol, blocking exactly one non-command prompt. + + Slash commands are allowed without claiming the first-prompt slot. After the first + prompt is blocked, every later request—including the wrapper's replay—is allowed. + """ + + def serve() -> None: + claimed = False + try: + if path.exists(): + path.unlink() + srv = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + srv.bind(str(path)) + os.chmod(path, 0o600) + srv.listen(4) + srv.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, _ = srv.accept() + except TimeoutError: + continue + except OSError: + break + with conn, conn.makefile("rwb") as stream: + raw = stream.readline() + response: dict = {"action": "allow"} + blocked: tuple[str, str] | None = None + try: + request = json.loads(raw) + 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 isinstance(prompt, str) + and prompt.strip() + and not is_command + ): + if not claimed: + 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: + srv.close() + + thread = threading.Thread(target=serve, name="claude-first-prompt", daemon=True) + thread.start() + return thread + + +# --- JSON-RPC control channel --------------------------------------------------------- + +_PARSE_ERROR = -32700 +_INVALID_REQUEST = -32600 +_METHOD_NOT_FOUND = -32601 +_INVALID_PARAMS = -32602 + + +def _rpc_error(rid: object, code: int, message: str) -> str: + return json.dumps({"jsonrpc": "2.0", "id": rid, "error": {"code": code, "message": message}}) + + +def handle_jsonrpc_line(line: str, on_model_set: Callable[[str], None]) -> str | None: + """Handle one JSON-RPC request line. + + Returns the response JSON string, or ``None`` for a notification (no ``id``). + Dispatches ``model.set`` -> ``on_model_set(name)`` after validating the name. + """ + if not line.strip(): + return None + try: + request = json.loads(line) + except ValueError: + return _rpc_error(None, _PARSE_ERROR, "Parse error") + if not isinstance(request, dict): + return _rpc_error(None, _INVALID_REQUEST, "Invalid Request") + + rid = request.get("id") + is_notification = "id" not in request + + def reply(response: str) -> str | None: + return None if is_notification else response + + if request.get("jsonrpc") != "2.0" or not isinstance(request.get("method"), str): + return reply(_rpc_error(rid, _INVALID_REQUEST, "Invalid Request")) + if request.get("method") != "model.set": + return reply(_rpc_error(rid, _METHOD_NOT_FOUND, "Method not found")) + + params = request.get("params") + name = params.get("name") if isinstance(params, dict) else None + if not isinstance(name, str) or not valid_model_name(name): + return reply(_rpc_error(rid, _INVALID_PARAMS, "Invalid params: model name")) + + on_model_set(name) # name narrowed to str by the isinstance guard above + if is_notification: + return None + return json.dumps({"jsonrpc": "2.0", "id": rid, "result": {"model": name, "injected": True}}) + + +def serve_control_socket( + path: Path, + on_model_set: Callable[[str], None], + stop: threading.Event, + *, + log: Callable[[str], None] = lambda _m: None, +) -> threading.Thread: + """Start a daemon AF_UNIX server thread that dispatches line-delimited JSON-RPC. + + Unlinks a stale socket file, binds with owner-only perms, and accepts connections + until ``stop`` is set. Returns the started thread. + """ + + def serve() -> None: + try: + if path.exists(): + path.unlink() + srv = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + srv.bind(str(path)) + os.chmod(path, 0o600) + srv.listen(4) + srv.settimeout(0.5) + except OSError as exc: + log(f"[ERR] control socket bind failed: {exc!r}") + return + log(f"[READY] control socket {path}") + try: + while not stop.is_set(): + try: + conn, _ = srv.accept() + except TimeoutError: + continue + except OSError: + break + with conn, conn.makefile("rwb") as stream: + for raw in stream: + try: + response = handle_jsonrpc_line( + raw.decode("utf-8", "replace"), on_model_set + ) + except Exception as exc: # noqa: BLE001 - one bad line must not kill the server + log(f"[ERR] control line: {exc!r}") + continue + if response is not None: + stream.write((response + "\n").encode()) + stream.flush() + finally: + srv.close() + + thread = threading.Thread(target=serve, name="claude-pty-control", daemon=True) + thread.start() + return thread + + +# --- terminal / window-size plumbing -------------------------------------------------- + + +class TerminalModeGuard: + """Put stdin in raw mode for the PTY session; always restore on exit. + + A no-op when stdin is not a TTY (piped input under tests/CI). Ports the Rust POC's + ``TerminalModeGuard``. + """ + + 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: + """Propagate the controlling terminal's window size onto the PTY master.""" + 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 + + +# --- orchestrator --------------------------------------------------------------------- + + +def run_claude_pty( + argv: list[str], + *, + route_prompt: Callable[[str], str], + switch_message: str, + socket_path: Path, + log_path: Path | None = None, +) -> int: + """Spawn *argv* in a PTY and run the smart-router CUJ, returning the child exit code. + + A UserPromptSubmit hook sends the first prompt to ``socket_path`` and blocks it. + Once Claude returns to an idle prompt box, this wrapper types ``/model``, confirms + the switch, and replays the captured prompt. The replay is allowed by the socket's + one-shot gate, so the first inference runs on Claude Code's newly selected model. + """ + + 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 + + # UCODE_CLAUDE_PTY_DEBUG=1 dumps each ANSI-stripped output chunk to the log so the + # readiness / confirm-dialog / turn markers can be grepped from a real launch. + debug = os.environ.get("UCODE_CLAUDE_PTY_DEBUG") == "1" + + # Hold the child immediately before exec until the hook socket is listening. Without + # this small gate, a positional prompt on a very fast launch could invoke the hook + # before the server thread has bound and fail open on the starting model. + gate_read, gate_write = os.pipe() + pid, master_fd = pty.fork() + if pid == 0: # child: become claude + os.close(gate_write) + try: + os.read(gate_read, 1) + finally: + os.close(gate_read) + os.execvp(argv[0], argv) + os._exit(127) # unreachable if execvp succeeds + 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) + + 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 # monotonic of the most recent TUI output (0 = none yet) + phase = "waiting_prompt" + routed_prompt = "" + routed_model = "" + switch_started = 0.0 + picker_ready: OutputMarkerDetector | None = None + picker_rows: ModelPickerRows | None = None + switch_complete: OutputMarkerDetector | None = None + switch_step = "" + navigation_output_seen = False + while True: + readable = [master_fd, 0] if stdin_open else [master_fd] + try: + ready_fds, _, _ = select.select(readable, [], [], SELECT_TIMEOUT_S) + except InterruptedError: # SIGWINCH etc. + continue + + if 0 in ready_fds: + try: + data = os.read(0, 4096) + except OSError: + data = b"" + if not data: + stdin_open = False # EOF on stdin: stop selecting it, keep pumping + else: + os.write(master_fd, data) + + if master_fd in ready_fds: + try: + chunk = os.read(master_fd, 8192) + except OSError: # child exited -> EIO on Linux + 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": + if picker_ready is not None: + picker_ready.observe(chunk) + if picker_rows is not None: + picker_rows.observe(chunk) + if switch_step == "navigating": + navigation_output_seen = True + if 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) + picker_ready = OutputMarkerDetector(("use this session only",)) + picker_rows = ModelPickerRows(routed_model) + switch_started = now + switch_step = "opening_picker" + phase = "switching" + log(f"[SWITCH] -> {routed_model!r}") + 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 switch_step == "opening_picker" + and picker_ready is not None + and picker_ready.triggered + and picker_rows is not None + and picker_rows.navigation is not None + ): + navigation = picker_rows.navigation + log( + f"[PICKER] row {picker_rows.focused_row} -> " + f"{picker_rows.target_row}" + ) + if navigation: + os.write(master_fd, navigation) + navigation_output_seen = False + switch_step = "navigating" + else: + os.write(master_fd, b"s") + confirm.arm(now + CONFIRM_TIMEOUT_S) + switch_complete = OutputMarkerDetector(("for this session only",)) + switch_step = "selected" + elif ( + phase == "switching" + and switch_step == "navigating" + and navigation_output_seen + ): + # `s` is Claude Code's model-picker action for this session only. + os.write(master_fd, b"s") + confirm.arm(now + CONFIRM_TIMEOUT_S) + switch_complete = OutputMarkerDetector(("for this session only",)) + switch_step = "selected" + elif phase == "switching" and now - switch_started >= SWITCH_TIMEOUT_S: + # Do not silently run on the wrong model. Dismiss any open picker and put + # the original text back in the editor without submitting it, so the user + # can recover manually without losing their prompt. + 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" + reason = ( + "model picker did not render the selected model" + if picker_rows is not None and picker_rows.target_row is None + else "model switch confirmation timed out" + ) + log(f"[ERR] {reason} for {routed_model!r}") + finally: + stop.set() + with contextlib.suppress(OSError): + os.close(master_fd) + if socket_path.exists(): + with contextlib.suppress(OSError): + socket_path.unlink() + + _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/tests/test_claude_smart_routing_v2.py b/tests/test_claude_smart_routing_v2.py new file mode 100644 index 00000000..165bf090 --- /dev/null +++ b/tests/test_claude_smart_routing_v2.py @@ -0,0 +1,524 @@ +"""Tests for the experimental ENABLE_SMART_ROUTING_V2 Claude Code (PTY wrapper) launch path.""" + +from __future__ import annotations + +import fcntl +import json +import os +import pty +import socket +import struct +import sys +import termios +import threading +import time +from pathlib import Path + +import pytest + +from ucode.agents import claude +from ucode.smart_routing import claude_hooks, claude_pty + + +class TestStripAnsi: + def test_strips_csi_and_leaves_text(self): + assert claude_pty.strip_ansi(b"\x1b[1mSwitch\x1b[0m model?") == "Switch model?" + + def test_partial_sequence_is_lenient(self): + # A dangling ESC without a full sequence should not raise or eat later text. + assert "hello" in claude_pty.strip_ansi(b"hello\x1b") + + +class TestValidModelName: + @pytest.mark.parametrize( + "name", + [ + "system.ai.claude-opus-4-8[1m]", + "databricks-claude-sonnet-4-6", + "opus", + "claude-3-5-haiku", + ], + ) + def test_accepts_gateway_ids(self, name): + assert claude_pty.valid_model_name(name) + + @pytest.mark.parametrize( + "name", + ["", "a b", "a;b", "a\nb", "a\rb", "x" * 201, 123, None, "cmd`whoami`"], + ) + def test_rejects_injection_and_junk(self, name): + assert not claude_pty.valid_model_name(name) + + +class TestModelPickerLabels: + @pytest.mark.parametrize( + ("model", "friendly"), + [ + ("system.ai.claude-sonnet-5", "Sonnet 5"), + ("system.ai.claude-haiku-4-5", "Haiku 4.5"), + ("system.ai.claude-opus-4-8[1m]", "Opus 4.8 (1M)"), + ], + ) + def test_derives_current_claude_code_picker_label(self, model, friendly): + assert claude_pty.model_picker_labels(model) == (model, friendly) + + def test_keeps_unknown_model_as_raw_label(self): + assert claude_pty.model_picker_labels("custom-model") == ("custom-model",) + + +class TestConfirmationState: + def test_full_prompt_confirms_and_disarms(self): + state = claude_pty.ConfirmationState() + now = time.monotonic() + state.arm(now + 5) + assert state.observe(b"Switch model? Yes, switch to Opus / No, go back", now) == b"\r" + # Disarmed after firing: a second identical chunk is ignored. + assert state.observe(b"Switch model?", now) is None + + def test_marker_split_across_chunks(self): + state = claude_pty.ConfirmationState() + now = time.monotonic() + state.arm(now + 5) + assert state.observe(b"\x1b[1mSwitch mo", now) is None + assert state.observe(b"del?\x1b[0m Yes, switch to Opus", now) is None + assert state.observe(b" / No, go back", now) == b"\r" + + def test_waits_for_complete_dialog(self): + state = claude_pty.ConfirmationState() + now = time.monotonic() + state.arm(now + 5) + assert state.observe(b"Switch model? Yes, switch to Opus", now) is None + assert state.observe(b"No, go back", now) == b"\r" + + def test_past_deadline_clears(self): + state = claude_pty.ConfirmationState() + state.arm(time.monotonic() - 1) # already expired + assert state.observe(b"Switch model? Yes, switch to X", time.monotonic()) is None + + def test_not_armed_returns_none(self): + state = claude_pty.ConfirmationState() + assert state.observe(b"Switch model? Yes, switch to X", time.monotonic()) is None + + +class TestOutputMarkerDetector: + def test_latches_after_marker(self): + det = claude_pty.OutputMarkerDetector(("? for shortcuts",)) + assert det.observe(b"booting up") is False + assert det.observe(b"ready: ? for shortcuts") is True + # Latched: stays True even on unrelated later output. + assert det.observe(b"nothing here") is True + + def test_marker_split_across_chunks(self): + det = claude_pty.OutputMarkerDetector(("esc to interrupt",)) + assert det.observe(b"\x1b[2mesc to ") is False + assert det.observe(b"interr") is False + assert det.observe(b"upt\x1b[0m") is True + + +class TestHandleJsonRpc: + def test_valid_model_set(self): + seen: list[str] = [] + resp = claude_pty.handle_jsonrpc_line( + json.dumps( + {"jsonrpc": "2.0", "id": 1, "method": "model.set", "params": {"name": "opus"}} + ), + seen.append, + ) + assert seen == ["opus"] + assert json.loads(resp)["result"]["model"] == "opus" + + def test_notification_has_no_response_but_dispatches(self): + seen: list[str] = [] + resp = claude_pty.handle_jsonrpc_line( + json.dumps({"jsonrpc": "2.0", "method": "model.set", "params": {"name": "opus"}}), + seen.append, + ) + assert resp is None + assert seen == ["opus"] + + def test_unknown_method(self): + resp = claude_pty.handle_jsonrpc_line( + json.dumps({"jsonrpc": "2.0", "id": 2, "method": "nope"}), lambda _m: None + ) + assert json.loads(resp)["error"]["code"] == -32601 + + def test_missing_model_is_invalid_params(self): + resp = claude_pty.handle_jsonrpc_line( + json.dumps({"jsonrpc": "2.0", "id": 3, "method": "model.set", "params": {}}), + lambda _m: pytest.fail("must not dispatch"), + ) + assert json.loads(resp)["error"]["code"] == -32602 + + def test_injection_model_name_rejected(self): + resp = claude_pty.handle_jsonrpc_line( + json.dumps( + {"jsonrpc": "2.0", "id": 4, "method": "model.set", "params": {"name": "x\r/help"}} + ), + lambda _m: pytest.fail("must not dispatch"), + ) + assert json.loads(resp)["error"]["code"] == -32602 + + def test_bad_json_is_parse_error(self): + resp = claude_pty.handle_jsonrpc_line("{not json", lambda _m: None) + assert json.loads(resp)["error"]["code"] == -32700 + + def test_non_object_is_invalid_request(self): + resp = claude_pty.handle_jsonrpc_line("123", lambda _m: None) + assert json.loads(resp)["error"]["code"] == -32600 + + +class TestInjectors: + def test_model_picker_rows_navigates_to_target(self): + rows = claude_pty.ModelPickerRows("system.ai.claude-sonnet-5") + rows.observe(b" 3. system.ai.claude-sonnet-5 Custom Sonnet model\r") + rows.observe("\x1b[1m ❯ 4. system.ai.claude-haiku-4-5\x1b[0m".encode()) + assert rows.target_row == 3 + assert rows.focused_row == 4 + assert rows.navigation == b"\x1b[A" + + def test_model_picker_rows_can_move_down_or_stay(self): + rows = claude_pty.ModelPickerRows("target") + rows.observe("❯ 2. current\r 4. target".encode()) + assert rows.navigation == b"\x1b[B\x1b[B" + rows.observe("❯ 4. target".encode()) + assert rows.navigation == b"" + + def test_model_picker_rows_accepts_friendly_label(self): + rows = claude_pty.ModelPickerRows("system.ai.claude-sonnet-5") + rows.observe(" 2. Haiku 4.5\r❯ 3. Sonnet 5".encode()) + assert rows.target_row == 3 + assert rows.focused_row == 3 + assert rows.navigation == b"" + + def test_inject_model_switch_opens_picker(self): + read_fd, write_fd = os.pipe() + try: + claude_pty.inject_model_switch(write_fd) + assert os.read(read_fd, 100) == b"/model\r" + finally: + os.close(read_fd) + os.close(write_fd) + + def test_inject_note_writes_message(self): + read_fd, write_fd = os.pipe() + try: + claude_pty.inject_note(write_fd, "router picked opus") + out = os.read(read_fd, 200).decode() + assert "router picked opus" in out + finally: + os.close(read_fd) + os.close(write_fd) + + def test_replays_multiline_prompt_as_one_bracketed_paste(self): + read_fd, write_fd = os.pipe() + try: + claude_pty.inject_prompt(write_fd, "first\nsecond") + assert os.read(read_fd, 200) == b"\x1b[200~first\nsecond\x1b[201~\r" + finally: + os.close(read_fd) + os.close(write_fd) + + def test_can_restore_prompt_without_submitting(self): + read_fd, write_fd = os.pipe() + try: + claude_pty.inject_prompt(write_fd, "try again", submit=False) + assert os.read(read_fd, 200) == b"\x1b[200~try again\x1b[201~" + finally: + os.close(read_fd) + os.close(write_fd) + + +class TestTerminalModeGuard: + def test_noop_when_not_a_tty(self): + read_fd, write_fd = os.pipe() + try: + # A pipe fd is not a TTY: entering/exiting must not touch termios or raise. + with claude_pty.TerminalModeGuard(read_fd) as guard: + assert guard._saved is None + finally: + os.close(read_fd) + os.close(write_fd) + + def test_restores_tty_attrs(self): + master_fd, slave_fd = pty.openpty() + try: + before = termios.tcgetattr(slave_fd) + with claude_pty.TerminalModeGuard(slave_fd): + pass + assert termios.tcgetattr(slave_fd) == before + finally: + os.close(master_fd) + os.close(slave_fd) + + +class TestSyncWinsize: + def test_propagates_window_size(self): + stdin_master, stdin_slave = pty.openpty() + out_master, out_slave = pty.openpty() + try: + want = struct.pack("HHHH", 40, 120, 0, 0) # rows, cols + fcntl.ioctl(stdin_slave, termios.TIOCSWINSZ, want) + claude_pty.sync_winsize(out_master, stdin_slave) + got = fcntl.ioctl(out_slave, termios.TIOCGWINSZ, struct.pack("HHHH", 0, 0, 0, 0)) + assert struct.unpack("HHHH", got)[:2] == (40, 120) + finally: + for fd in (stdin_master, stdin_slave, out_master, out_slave): + os.close(fd) + + +class TestServeControlSocket: + def test_dispatches_model_set_over_socket(self, tmp_path): + sock_path = tmp_path / "ctl.sock" + seen: list[str] = [] + stop = threading.Event() + claude_pty.serve_control_socket(sock_path, seen.append, stop) + try: + deadline = time.monotonic() + 5 + while not sock_path.exists() and time.monotonic() < deadline: + time.sleep(0.01) + client = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + client.connect(str(sock_path)) + request = ( + json.dumps( + {"jsonrpc": "2.0", "id": 1, "method": "model.set", "params": {"name": "opus"}} + ) + + "\n" + ) + client.sendall(request.encode()) + response = client.makefile("rb").readline() + client.close() + assert seen == ["opus"] + assert json.loads(response)["result"]["model"] == "opus" + finally: + stop.set() + + +class TestFirstPromptHook: + def test_blocks_once_then_allows_replay(self, tmp_path): + sock_path = tmp_path / "first.sock" + blocked: list[tuple[str, str]] = [] + stop = threading.Event() + claude_pty.serve_first_prompt_socket( + sock_path, + lambda prompt: "sonnet" if prompt else "opus", + lambda prompt, model: blocked.append((prompt, model)), + stop, + ) + try: + deadline = time.monotonic() + 5 + while not sock_path.exists() and time.monotonic() < deadline: + time.sleep(0.01) + first = claude_pty.request_first_prompt_route( + sock_path, {"session_id": "s1", "prompt": "fix the parser"} + ) + replay = claude_pty.request_first_prompt_route( + sock_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_slash_command_does_not_claim_first_prompt(self, tmp_path): + sock_path = tmp_path / "first.sock" + stop = threading.Event() + claude_pty.serve_first_prompt_socket( + sock_path, lambda _prompt: "opus", lambda *_args: None, stop + ) + try: + deadline = time.monotonic() + 5 + while not sock_path.exists() and time.monotonic() < deadline: + time.sleep(0.01) + command = claude_pty.request_first_prompt_route(sock_path, {"prompt": "/hooks"}) + prompt = claude_pty.request_first_prompt_route(sock_path, {"prompt": "do work"}) + assert command == {"action": "allow"} + assert prompt == {"action": "block", "model": "opus"} + finally: + stop.set() + + def test_hook_output_blocks_with_actionable_message(self): + output = claude_pty.first_prompt_hook_output( + {"action": "block", "model": "system.ai.claude-sonnet-5"} + ) + assert output["decision"] == "block" + assert output["reason"] == ( + "✨ Smart Router selected system.ai.claude-sonnet-5 due to low complexity, " + "unclear intent, and no code reference." + ) + + def test_hook_failure_allows_prompt(self, tmp_path): + assert ( + claude_pty.request_first_prompt_route( + tmp_path / "missing.sock", {"prompt": "do work"}, timeout=0.01 + ) + is None + ) + assert claude_pty.first_prompt_hook_output(None) is None + + +class TestFirstPromptHookSettings: + def test_adds_stable_hook_without_removing_other_routing_hooks(self): + doc = { + "hooks": { + "PreToolUse": [ + { + "hooks": [ + { + "type": "command", + "command": "ucode claude-router-hook route-subagent", + } + ] + } + ] + } + } + claude_hooks.sync_first_prompt_hook(doc, "/bin/ucode") + claude_hooks.sync_first_prompt_hook(doc, "/bin/ucode") + command = doc["hooks"]["UserPromptSubmit"][0]["hooks"][0]["command"] + assert command == "/bin/ucode claude-router-hook route-first-prompt" + assert len(doc["hooks"]["UserPromptSubmit"]) == 1 + assert "route-subagent" in str(doc["hooks"]["PreToolUse"]) + + +class TestV2Launch: + def test_uses_unique_hook_settings_and_cleans_them_up(self, tmp_path, monkeypatch): + settings_path = tmp_path / "ucode-settings.json" + settings_path.write_text(json.dumps({"env": {"ANTHROPIC_BASE_URL": "https://gateway"}})) + monkeypatch.setattr(claude, "APP_DIR", tmp_path) + monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", settings_path) + monkeypatch.setattr(claude, "SMART_ROUTING_V2_CLAUDE_LOG", tmp_path / "v2.log") + monkeypatch.setattr(claude, "get_databricks_token", lambda *_a, **_k: "token") + monkeypatch.setattr(claude, "build_auth_token_argv", lambda *_a, **_k: ["/bin/ucode"]) + captured: dict = {} + + def fake_run(argv, **kwargs): + generated = argv[argv.index("--settings") + 1] + captured["path"] = generated + captured["settings"] = json.loads(Path(generated).read_text()) + captured["kwargs"] = kwargs + return 0 + + monkeypatch.setattr(claude_pty, "run_claude_pty", fake_run) + with pytest.raises(SystemExit) as exc: + claude._launch_smart_routing_v2({"workspace": "https://example.com"}, ["--debug"]) + + assert exc.value.code == 0 + assert not Path(captured["path"]).exists() + env = captured["settings"]["env"] + assert env["ANTHROPIC_BASE_URL"] == "https://gateway" + assert env[claude_hooks.FIRST_PROMPT_SOCKET_ENV].endswith(".sock") + hook = captured["settings"]["hooks"]["UserPromptSubmit"][0]["hooks"][0] + assert hook["command"] == "/bin/ucode claude-router-hook route-first-prompt" + assert captured["kwargs"]["route_prompt"]("anything") == claude.SMART_ROUTING_V2_MODEL + + +class TestPtyFlow: + def test_hook_block_switch_confirm_and_replay(self, tmp_path): + fake_claude = tmp_path / "fake_claude.py" + capture = tmp_path / "capture.bin" + 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", + "model": "system.ai.claude-sonnet-5", +} +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 + +def read_exact(size): + data = b"" + while len(data) < size: + data += os.read(0, size - len(data)) + return data + +model_command = read_until(b"\\r") +print( + "Select model\\n" + " 3. system.ai.claude-sonnet-5 Custom Sonnet model\\n" + "❯ 4. system.ai.claude-haiku-4-5 Custom Haiku model\\n" + "Enter to set as default s use this session only", + flush=True, +) +navigation = read_exact(3) +print("model picker focus moved", flush=True) +choice = read_exact(1) +print("Switch model? Yes, switch to Sonnet / No, go back", flush=True) +confirmation = read_exact(1) +print("Set model to Sonnet for this session only", flush=True) +replayed = read_until(b"\\x1b[201~\\r") +capture_path.write_bytes( + model_command + b"|" + navigation + b"|" + choice + b"|" + confirmation + + b"|" + replayed +) +""".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 capture.read_bytes() == ( + b"/model\r|\x1b[A|s|\r|\x1b[200~fix\nthe parser\x1b[201~\r" + ) + + +class TestV2Router: + def test_stub_routes_everything_to_fixed_model(self): + route = claude._v2_router({}) + assert route("fix the parser") == claude.SMART_ROUTING_V2_MODEL + assert claude.SMART_ROUTING_V2_MODEL == "system.ai.claude-sonnet-4-6[1m]" + + +class TestLaunchGate: + def test_v2_gate_routes_to_pty_when_enabled(self, monkeypatch): + monkeypatch.setenv("ENABLE_SMART_ROUTING_V2", "1") + calls: dict[str, bool] = {} + monkeypatch.setattr( + claude, "_launch_smart_routing_v2", lambda *_a: calls.__setitem__("v2", True) + ) + monkeypatch.setattr(claude, "exec_or_spawn", lambda *_a: calls.__setitem__("exec", True)) + claude.launch({"workspace": "https://example.databricks.com"}, []) + assert calls == {"v2": True} + + def test_normal_launch_when_disabled(self, monkeypatch): + monkeypatch.delenv("ENABLE_SMART_ROUTING_V2", raising=False) + calls: dict[str, bool] = {} + monkeypatch.setattr( + claude, "_launch_smart_routing_v2", lambda *_a: calls.__setitem__("v2", True) + ) + monkeypatch.setattr(claude, "get_databricks_token", lambda *_a, **_k: "tok") + monkeypatch.setattr(claude, "exec_or_spawn", lambda *_a: calls.__setitem__("exec", True)) + claude.launch({"workspace": "https://example.databricks.com"}, []) + assert calls == {"exec": True}