From c6ceb197e09b8b2e54df088dd1f2c4f9a5b1b53d Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Mon, 24 Aug 2026 22:07:47 +0000 Subject: [PATCH 1/2] test --- src/ucode/agents/claude.py | 77 ++++ src/ucode/smart_routing/claude_proxy.py | 202 +++++++++++ src/ucode/smart_routing/claude_pty.py | 450 ++++++++++++++++++++++++ tests/test_claude_proxy.py | 58 +++ tests/test_claude_smart_routing_v2.py | 260 ++++++++++++++ 5 files changed, 1047 insertions(+) create mode 100644 src/ucode/smart_routing/claude_proxy.py create mode 100644 src/ucode/smart_routing/claude_pty.py create mode 100644 tests/test_claude_proxy.py create mode 100644 tests/test_claude_smart_routing_v2.py diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 9f4df255..b2755efa 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -31,6 +31,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, @@ -62,6 +63,16 @@ # 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-router.log" +SMART_ROUTING_V2_SETTINGS = APP_DIR / "claude-v2-settings.json" + def is_update_available() -> tuple[str, str] | None: return available_npm_package_update(SPEC["package"]) @@ -1017,12 +1028,78 @@ def _launch_relayed(state: dict, binary: str, tool_args: list[str]) -> None: raise SystemExit(returncode) +def _v2_router(state: dict): + """Return the routing callable the proxy applies to each Messages API request. + + Today a STUB: ignores the request and always returns SMART_ROUTING_V2_MODEL. For real + dynamic routing, classify the prompt here (the request body's ``messages`` carry it; + see ``smart_routing.routing.select_route``) and return the chosen model per request. + ``state`` is accepted so a real router can reach the workspace/token/model metadata.""" + + def route(_body: dict) -> str: + return SMART_ROUTING_V2_MODEL + + return route + + +def _launch_smart_routing_v2(state: dict, tool_args: list[str]) -> None: + """Launch Claude Code behind a loopback model-routing proxy. + + A local proxy sits between Claude Code and the gateway (via ANTHROPIC_BASE_URL) and + rewrites each request's ``model`` to the router's pick — so the FIRST prompt already + routes to the chosen model, every turn can route independently, and Claude Code never + needs to know the target exists (no ``/model``, no discovery, no mutated default). + Spawn-and-wait (not exec) so this process stays alive to run + tear down the proxy — + mirroring ``_launch_relayed``.""" + from ucode.smart_routing import claude_proxy + + binary = SPEC["binary"] + workspace = state["workspace"] + os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile")) + + server, client, port = claude_proxy.start_proxy( + workspace, _v2_router(state), log_path=SMART_ROUTING_V2_CLAUDE_LOG + ) + # Point Claude Code at the proxy without touching the shared settings file: write a + # per-launch settings file that overrides only ANTHROPIC_BASE_URL, and pass it via + # --settings. Cleaned up on exit; the user's ~/.claude/ucode-settings.json is untouched. + settings = read_json_safe(CLAUDE_SETTINGS_PATH) + env = settings.setdefault("env", {}) + if isinstance(env, dict): + env["ANTHROPIC_BASE_URL"] = f"http://127.0.0.1:{port}" + write_json_file(SMART_ROUTING_V2_SETTINGS, settings) + argv = [binary, "--settings", str(SMART_ROUTING_V2_SETTINGS), *tool_args] + + print_note( + f"Smart routing v2: routing every prompt through the local model router " + f"(stub -> {SMART_ROUTING_V2_MODEL}); router log: {SMART_ROUTING_V2_CLAUDE_LOG}." + ) + server_thread = threading.Thread(target=server.serve_forever, daemon=True) + server_thread.start() + proc = subprocess.Popen(argv) + try: + returncode = proc.wait() + except KeyboardInterrupt: + proc.send_signal(signal.SIGINT) + returncode = proc.wait() + finally: + server.shutdown() + client.close() + 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/smart_routing/claude_proxy.py b/src/ucode/smart_routing/claude_proxy.py new file mode 100644 index 00000000..8d34f037 --- /dev/null +++ b/src/ucode/smart_routing/claude_proxy.py @@ -0,0 +1,202 @@ +"""Model-routing proxy for Claude Code (smart routing v2). + +Claude Code has no ``app-server``/JSON-RPC seam like Codex, and — as +``claude_routing.route_launch_model`` notes — "no hook/MCP can retarget the root +model once the session is running." But every turn Claude Code makes an Anthropic +Messages API request to ``ANTHROPIC_BASE_URL``, and that request body carries BOTH +the prompt and the ``model``. So the routing seam is the HTTP request: point +``ANTHROPIC_BASE_URL`` at this loopback proxy, and it reads the prompt out of each +request, picks a model, rewrites the body's ``model`` field, and forwards to the +workspace gateway. This is the direct analog of the Codex interposer (which rewrites +``turn/start.model`` on the WebSocket) — just one layer down, at the HTTP request. + +Because it routes below Claude Code: + - the FIRST prompt is routed correctly (the request carries it), + - every turn can route independently, + - Claude Code never has to know the target model exists (no ``/model``, no gateway + model discovery, no mutated default), and + - nothing scrapes or drives the TUI. + +Auth is passthrough: in the normal (non-relayed) launch Claude Code's ``apiKeyHelper`` +already mints the gateway credential and sends it, so — unlike ``gateway_proxy`` — this +proxy needs no token management. It forwards headers verbatim (minus hop-by-hop) and +streams the response back byte-for-byte (SSE token streaming is not buffered). + +Security: binds 127.0.0.1 only; never logs header values or bodies (the routing log +records only the model ids swapped, never prompt text). +""" + +from __future__ import annotations + +import json +import time +from collections.abc import Callable +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path + +import httpx + +# Hop-by-hop headers must not be forwarded across the proxy; content-length is dropped +# too because we may rewrite the body (httpx recomputes it from the content we pass). +_HOP_BY_HOP = frozenset( + h.lower() + for h in ( + "connection", + "keep-alive", + "proxy-authenticate", + "proxy-authorization", + "te", + "trailers", + "transfer-encoding", + "upgrade", + "host", + "content-length", + ) +) +# Generous read timeout: a turn streams over one response with SSE pings between chunks. +_UPSTREAM_TIMEOUT = httpx.Timeout(connect=10.0, read=600.0, write=600.0, pool=10.0) + +# A router maps a parsed request body to the model id to use. The stub ignores the body +# and always returns the fixed pick; a real router inspects body["messages"]. +RouteFn = Callable[[dict], str] + + +def rewrite_model( + body: bytes, *, is_json: bool, route: RouteFn +) -> tuple[bytes, str | None, str | None]: + """Rewrite the ``model`` field of a Messages API request body. + + Returns ``(new_body, original_model, chosen_model)``. When the body is not a JSON + object with a ``model`` key (or the router picks the same model), ``new_body`` is the + input unchanged. Pure and side-effect-free, so it is unit-testable without a server. + """ + if not body or not is_json: + return body, None, None + try: + parsed = json.loads(body) + except ValueError: + return body, None, None + if not isinstance(parsed, dict) or "model" not in parsed: + return body, None, None + original = parsed.get("model") + chosen = route(parsed) + if not chosen or chosen == original: + return body, original if isinstance(original, str) else None, chosen + parsed["model"] = chosen + return json.dumps(parsed).encode(), original if isinstance(original, str) else None, chosen + + +def _is_json_request(handler: BaseHTTPRequestHandler) -> bool: + """True when the request body is plain (uncompressed) JSON we can safely rewrite.""" + if "content-encoding" in {k.lower() for k in handler.headers}: + return False # compressed body: don't touch it + return "json" in handler.headers.get("Content-Type", "").lower() + + +def _forwarded_headers(handler: BaseHTTPRequestHandler) -> dict[str, str]: + return {k: v for k, v in handler.headers.items() if k.lower() not in _HOP_BY_HOP} + + +class _RouterProxyHandler(BaseHTTPRequestHandler): + # Set by the server factory. + client: httpx.Client + route: RouteFn + log_fn: Callable[[str], None] + + def log_message(self, format: str, *args: object) -> None: + return + + def _safe_send_error(self, code: int, message: str) -> None: + try: + self.send_error(code, message) + except OSError: + pass + + def _handle(self) -> None: + length = int(self.headers.get("Content-Length", 0) or 0) + body = self.rfile.read(length) if length else b"" + url = self.path.lstrip("/") + + new_body = body + original = chosen = None + if body and self.command == "POST": + new_body, original, chosen = rewrite_model( + body, is_json=_is_json_request(self), route=self.route + ) + # One line per request (model ids + path only, never prompt text) so the log shows + # ALL traffic — confirming requests reach the proxy, not just rewrites. + action = "ROUTE" if (chosen and chosen != original) else "PASS" + self.log_fn(f"[{action}] {self.command} /{url} model={original!r}->{chosen!r}") + + try: + with self.client.stream( + self.command, url, headers=_forwarded_headers(self), content=new_body or None + ) as resp: + self.log_fn(f"[UPSTREAM] /{url} {resp.status_code}") + self._relay_response(resp) + except (BrokenPipeError, ConnectionResetError): + return # client (Claude Code) closed before/while relaying — routine on cancel + except httpx.HTTPError as exc: + self.log_fn(f"[ERR] /{url} {type(exc).__name__}: {exc}") + self._safe_send_error(502, "model router proxy upstream error") + + def _relay_response(self, resp: httpx.Response) -> None: + # Stream chunks as they arrive so SSE token streaming is not buffered; iter_raw + # preserves any Content-Encoding verbatim (we relay that header). + try: + self.send_response(resp.status_code) + for key, value in resp.headers.items(): + if key.lower() not in _HOP_BY_HOP: + self.send_header(key, value) + self.end_headers() + for chunk in resp.iter_raw(): + if chunk: + self.wfile.write(chunk) + self.wfile.flush() + except (BrokenPipeError, ConnectionResetError): + return # client closed mid-response — routine on cancelled turns / SSE teardown + except httpx.HTTPError: + return # upstream dropped mid-stream; headers already sent, stop cleanly + + # Transparent pass-through for every method (GET model discovery, POST messages, …). + def __getattr__(self, name: str): + if name.startswith("do_"): + return self._handle + raise AttributeError(name) + + +def start_proxy( + workspace: str, + route: RouteFn, + *, + log_path: Path | None = None, +) -> tuple[ThreadingHTTPServer, httpx.Client, int]: + """Start the loopback model-routing proxy on an OS-assigned port. + + Forwards to ``{workspace}/ai-gateway/anthropic/`` with headers passed through + verbatim (Claude Code's apiKeyHelper credential included). Returns + ``(server, client, port)``; the caller runs ``server`` (e.g. in a thread) and calls + ``server.shutdown()`` / ``client.close()`` on exit. Point ``ANTHROPIC_BASE_URL`` at + ``http://127.0.0.1:{port}``. + """ + + 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 + + upstream_base = f"{workspace.rstrip('/')}/ai-gateway/anthropic/" + client = httpx.Client(base_url=upstream_base, timeout=_UPSTREAM_TIMEOUT, follow_redirects=False) + handler = type( + "BoundRouterProxyHandler", + (_RouterProxyHandler,), + {"client": client, "route": staticmethod(route), "log_fn": staticmethod(log)}, + ) + server = ThreadingHTTPServer(("127.0.0.1", 0), handler) + port = server.server_address[1] + log(f"[READY] 127.0.0.1:{port} -> {upstream_base}") + return server, client, port diff --git a/src/ucode/smart_routing/claude_pty.py b/src/ucode/smart_routing/claude_pty.py new file mode 100644 index 00000000..6d95fb9a --- /dev/null +++ b/src/ucode/smart_routing/claude_pty.py @@ -0,0 +1,450 @@ +"""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* +by typing ``/model `` into the TUI and auto-confirming the "Switch model?" dialog +by watching the PTY output. A Unix-domain control socket accepts line-delimited JSON-RPC +``model.set`` requests (parity with the reference POC and a clean test surface); the +automatic first-prompt CUJ drives the same ``on_model_set`` path internally. + +``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``). + +Detecting *when the user submits their first prompt* is done from **stdin** (the Enter +keystroke after typed content), not by scraping output — Claude Code's TUI renders spaces +as cursor moves and uses randomized spinner text, so output markers are unreliable. Marker +matching that remains (readiness, the confirm dialog) is whitespace-insensitive (see +:func:`_squash`). The only runtime-spike output constants left are :data:`READY_MARKERS` +(pre-emptive mode only) and :data:`ConfirmationState.PROMPT_MARKERS`. +""" + +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 + +# --- runtime-spike markers (whitespace-squashed substrings; tune from a real claude launch) --- +# TUI is ready for input. A *bonus* signal for the pre-emptive switch mode; the primary +# readiness signal is version-independent (idle after the initial paint — see READY_QUIET_S). +# Matching is whitespace-insensitive (see _squash), because Claude Code renders spaces as +# cursor moves that vanish under strip_ansi. +READY_MARKERS = ("? for shortcuts", "Welcome to Claude Code") + +MAX_MODEL_NAME_LEN = 200 +CONFIRM_TIMEOUT_S = 5.0 +# Pre-emptive readiness heuristic: once the TUI has produced output and then stays quiet +# this long, its initial paint is done and it is idle waiting for input — safe to type +# `/model`. Marker-independent, so it survives Claude Code TUI changes. +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._:/\-\[\]]+$") +# 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)) + ) + + +class ConfirmationState: + """Watch PTY output for Claude's "Switch model?" dialog and auto-confirm it. + + ``arm`` after typing ``/model`` (with a deadline), feed each output chunk to + ``observe``; when the confirmation prompt is seen it returns the keystrokes to + send (``b"\\r"``) and disarms. A rolling buffer keeps the last ``window`` chars so + a marker split across two PTY reads (or interrupted by an ANSI escape) still matches. + """ + + PROMPT_MARKERS = ("Switch model?", "Yes, switch to") + + 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 any(_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 + + +def inject_model_switch(master_fd: int, model: str, confirm: ConfirmationState, now: float) -> None: + """Type ``/model \\r`` into the TUI and arm the confirm watcher.""" + os.write(master_fd, f"/model {model}\r".encode()) + confirm.arm(now + CONFIRM_TIMEOUT_S) + + +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()) + + +# --- 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], + *, + target_model: str, + switch_message: str, + socket_path: Path, + log_path: Path | None = None, + switch_mode: str = "preemptive", +) -> int: + """Spawn *argv* in a PTY and run the smart-router CUJ, returning the child exit code. + + Pumps stdin<->master and master<->stdout; runs the JSON-RPC control socket; and + auto-switches to ``target_model`` — in ``"reactive"`` mode on the first prompt the + user submits (detected from the Enter keystroke on stdin), or in ``"preemptive"`` + mode once the TUI paints and goes idle — surfacing ``switch_message`` in the transcript. + """ + + 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" + + pid, master_fd = pty.fork() + if pid == 0: # child: become claude + os.execvp(argv[0], argv) + os._exit(127) # unreachable if execvp succeeds + + confirm = ConfirmationState() + ready = OutputMarkerDetector(READY_MARKERS) + lock = threading.Lock() + switched = {"done": False} + stop = threading.Event() + + def switch_to(model: str) -> None: + with lock: + if switched["done"]: + return + switched["done"] = True + inject_note(1, switch_message) + inject_model_switch(master_fd, model, confirm, time.monotonic()) + log(f"[SWITCH] -> {model!r}") + + serve_control_socket(socket_path, switch_to, stop, log=log) + + 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) + typed_content = False # printable input seen since the last Enter (reactive mode) + pending_switch = False # first prompt submitted; switch when Claude next goes idle + while True: + readable = [master_fd, 0] if stdin_open else [master_fd] + # Both modes fire on an idle gap, so they need periodic wakes while waiting: + # preemptive from the start, reactive only after the first prompt is submitted. + need_idle_wake = not switched["done"] and ( + switch_mode == "preemptive" + or (switch_mode == "reactive" and pending_switch) + ) + timeout = SELECT_TIMEOUT_S if need_idle_wake else None + try: + ready_fds, _, _ = select.select(readable, [], [], timeout) + 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) + # Reactive: the user submits their first prompt when they press Enter + # after typing. This is a stdin keystroke (fully under our control) — + # far more reliable than scraping Claude's rendered output. Arm here; + # the actual /model keystroke fires once Claude goes idle (below), so it + # lands in an idle input box rather than queued mid-response. + if switch_mode == "reactive" and not switched["done"] and not pending_switch: + if any(byte >= 0x20 and byte != 0x7F for byte in data): + typed_content = True + if typed_content and (b"\r" in data or b"\n" in data): + pending_switch = True + + 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}") + with lock: + keystroke = confirm.observe(chunk, last_output) + if keystroke is not None: + os.write(master_fd, keystroke) + ready.observe(chunk) + + # Fire on an idle gap: preemptive before any prompt (TUI painted + idle); + # reactive only after the first prompt was submitted (turn done + idle), so the + # /model command types into an idle input box. + if not switched["done"]: + now = time.monotonic() + idle = last_output > 0.0 and (now - last_output) >= READY_QUIET_S + if switch_mode == "preemptive" and (ready.triggered or idle): + switch_to(target_model) + elif switch_mode == "reactive" and pending_switch and idle: + switch_to(target_model) + 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_proxy.py b/tests/test_claude_proxy.py new file mode 100644 index 00000000..f6439837 --- /dev/null +++ b/tests/test_claude_proxy.py @@ -0,0 +1,58 @@ +"""Tests for the smart-routing-v2 model-routing proxy's request rewriting.""" + +from __future__ import annotations + +import json + +from ucode.smart_routing import claude_proxy + + +def _stub(model: str): + return lambda _body: model + + +class TestRewriteModel: + def test_rewrites_model_in_json_body(self): + body = json.dumps({"model": "opus", "messages": [{"role": "user", "content": "hi"}]}).encode() + new_body, original, chosen = claude_proxy.rewrite_model( + body, is_json=True, route=_stub("system.ai.claude-sonnet-5") + ) + assert original == "opus" + assert chosen == "system.ai.claude-sonnet-5" + assert json.loads(new_body)["model"] == "system.ai.claude-sonnet-5" + # Everything else in the body is preserved. + assert json.loads(new_body)["messages"] == [{"role": "user", "content": "hi"}] + + def test_unchanged_when_router_picks_same_model(self): + body = json.dumps({"model": "system.ai.claude-sonnet-5"}).encode() + new_body, original, chosen = claude_proxy.rewrite_model( + body, is_json=True, route=_stub("system.ai.claude-sonnet-5") + ) + assert new_body == body + assert original == "system.ai.claude-sonnet-5" + assert chosen == "system.ai.claude-sonnet-5" + + def test_unchanged_when_not_json(self): + body = b"\x00\x01 not json" + new_body, original, chosen = claude_proxy.rewrite_model( + body, is_json=False, route=_stub("x") + ) + assert new_body == body + assert original is None and chosen is None + + def test_unchanged_when_body_has_no_model(self): + body = json.dumps({"messages": []}).encode() + new_body, original, chosen = claude_proxy.rewrite_model(body, is_json=True, route=_stub("x")) + assert new_body == body + assert original is None and chosen is None + + def test_unchanged_on_empty_body(self): + new_body, original, chosen = claude_proxy.rewrite_model(b"", is_json=True, route=_stub("x")) + assert new_body == b"" + assert original is None and chosen is None + + def test_malformed_json_passes_through(self): + body = b'{"model": "opus"' # truncated + new_body, original, chosen = claude_proxy.rewrite_model(body, is_json=True, route=_stub("y")) + assert new_body == body + assert original is None and chosen is None diff --git a/tests/test_claude_smart_routing_v2.py b/tests/test_claude_smart_routing_v2.py new file mode 100644 index 00000000..7f3f6741 --- /dev/null +++ b/tests/test_claude_smart_routing_v2.py @@ -0,0 +1,260 @@ +"""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 termios +import threading +import time + +import pytest + +from ucode.agents import claude +from ucode.smart_routing import 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 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) == 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_inject_model_switch_types_command_and_arms(self): + read_fd, write_fd = os.pipe() + try: + confirm = claude_pty.ConfirmationState() + claude_pty.inject_model_switch(write_fd, "opus", confirm, time.monotonic()) + assert os.read(read_fd, 100) == b"/model opus\r" + # Armed: it now auto-confirms the dialog. + assert confirm.observe(b"Switch model? Yes, switch to X", time.monotonic()) == b"\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) + + +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 TestV2Router: + def test_stub_routes_everything_to_fixed_model(self): + route = claude._v2_router({}) + assert route({"model": "opus", "messages": []}) == claude.SMART_ROUTING_V2_MODEL + assert claude.SMART_ROUTING_V2_MODEL == "system.ai.claude-sonnet-5" + + +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} From 77e684e3cf3bb2833314e5478611e6421d084c15 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Mon, 24 Aug 2026 23:27:32 +0000 Subject: [PATCH 2/2] hi --- src/ucode/agents/claude.py | 86 ++--- src/ucode/cli.py | 19 ++ src/ucode/smart_routing/claude_hooks.py | 29 ++ src/ucode/smart_routing/claude_proxy.py | 202 ------------ src/ucode/smart_routing/claude_pty.py | 403 +++++++++++++++++++----- tests/test_claude_proxy.py | 58 ---- tests/test_claude_smart_routing_v2.py | 259 ++++++++++++++- 7 files changed, 671 insertions(+), 385 deletions(-) delete mode 100644 src/ucode/smart_routing/claude_proxy.py delete mode 100644 tests/test_claude_proxy.py diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index b2755efa..d6fabfe3 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,6 +27,7 @@ ) from ucode.databricks import ( build_auth_shell_command, + build_auth_token_argv, build_tool_base_url, get_databricks_token, ) @@ -33,7 +35,9 @@ 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 @@ -70,8 +74,7 @@ # 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-router.log" -SMART_ROUTING_V2_SETTINGS = APP_DIR / "claude-v2-settings.json" +SMART_ROUTING_V2_CLAUDE_LOG = APP_DIR / "claude-v2-pty.log" def is_update_available() -> tuple[str, str] | None: @@ -1029,62 +1032,73 @@ def _launch_relayed(state: dict, binary: str, tool_args: list[str]) -> None: def _v2_router(state: dict): - """Return the routing callable the proxy applies to each Messages API request. + """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 (the request body's ``messages`` carry it; - see ``smart_routing.routing.select_route``) and return the chosen model per request. + 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(_body: dict) -> str: + 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 behind a loopback model-routing proxy. + """Launch Claude Code in the first-prompt routing PTY wrapper. - A local proxy sits between Claude Code and the gateway (via ANTHROPIC_BASE_URL) and - rewrites each request's ``model`` to the router's pick — so the FIRST prompt already - routes to the chosen model, every turn can route independently, and Claude Code never - needs to know the target exists (no ``/model``, no discovery, no mutated default). - Spawn-and-wait (not exec) so this process stays alive to run + tear down the proxy — - mirroring ``_launch_relayed``.""" - from ucode.smart_routing import claude_proxy + 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")) - server, client, port = claude_proxy.start_proxy( - workspace, _v2_router(state), log_path=SMART_ROUTING_V2_CLAUDE_LOG - ) - # Point Claude Code at the proxy without touching the shared settings file: write a - # per-launch settings file that overrides only ANTHROPIC_BASE_URL, and pass it via - # --settings. Cleaned up on exit; the user's ~/.claude/ucode-settings.json is untouched. - settings = read_json_safe(CLAUDE_SETTINGS_PATH) + 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 isinstance(env, dict): - env["ANTHROPIC_BASE_URL"] = f"http://127.0.0.1:{port}" - write_json_file(SMART_ROUTING_V2_SETTINGS, settings) - argv = [binary, "--settings", str(SMART_ROUTING_V2_SETTINGS), *tool_args] + 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: routing every prompt through the local model router " - f"(stub -> {SMART_ROUTING_V2_MODEL}); router log: {SMART_ROUTING_V2_CLAUDE_LOG}." + 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}." ) - server_thread = threading.Thread(target=server.serve_forever, daemon=True) - server_thread.start() - proc = subprocess.Popen(argv) try: - returncode = proc.wait() - except KeyboardInterrupt: - proc.send_signal(signal.SIGINT) - returncode = proc.wait() + 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: - server.shutdown() - client.close() + settings_path.unlink(missing_ok=True) + socket_path.unlink(missing_ok=True) raise SystemExit(returncode) 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_proxy.py b/src/ucode/smart_routing/claude_proxy.py deleted file mode 100644 index 8d34f037..00000000 --- a/src/ucode/smart_routing/claude_proxy.py +++ /dev/null @@ -1,202 +0,0 @@ -"""Model-routing proxy for Claude Code (smart routing v2). - -Claude Code has no ``app-server``/JSON-RPC seam like Codex, and — as -``claude_routing.route_launch_model`` notes — "no hook/MCP can retarget the root -model once the session is running." But every turn Claude Code makes an Anthropic -Messages API request to ``ANTHROPIC_BASE_URL``, and that request body carries BOTH -the prompt and the ``model``. So the routing seam is the HTTP request: point -``ANTHROPIC_BASE_URL`` at this loopback proxy, and it reads the prompt out of each -request, picks a model, rewrites the body's ``model`` field, and forwards to the -workspace gateway. This is the direct analog of the Codex interposer (which rewrites -``turn/start.model`` on the WebSocket) — just one layer down, at the HTTP request. - -Because it routes below Claude Code: - - the FIRST prompt is routed correctly (the request carries it), - - every turn can route independently, - - Claude Code never has to know the target model exists (no ``/model``, no gateway - model discovery, no mutated default), and - - nothing scrapes or drives the TUI. - -Auth is passthrough: in the normal (non-relayed) launch Claude Code's ``apiKeyHelper`` -already mints the gateway credential and sends it, so — unlike ``gateway_proxy`` — this -proxy needs no token management. It forwards headers verbatim (minus hop-by-hop) and -streams the response back byte-for-byte (SSE token streaming is not buffered). - -Security: binds 127.0.0.1 only; never logs header values or bodies (the routing log -records only the model ids swapped, never prompt text). -""" - -from __future__ import annotations - -import json -import time -from collections.abc import Callable -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -from pathlib import Path - -import httpx - -# Hop-by-hop headers must not be forwarded across the proxy; content-length is dropped -# too because we may rewrite the body (httpx recomputes it from the content we pass). -_HOP_BY_HOP = frozenset( - h.lower() - for h in ( - "connection", - "keep-alive", - "proxy-authenticate", - "proxy-authorization", - "te", - "trailers", - "transfer-encoding", - "upgrade", - "host", - "content-length", - ) -) -# Generous read timeout: a turn streams over one response with SSE pings between chunks. -_UPSTREAM_TIMEOUT = httpx.Timeout(connect=10.0, read=600.0, write=600.0, pool=10.0) - -# A router maps a parsed request body to the model id to use. The stub ignores the body -# and always returns the fixed pick; a real router inspects body["messages"]. -RouteFn = Callable[[dict], str] - - -def rewrite_model( - body: bytes, *, is_json: bool, route: RouteFn -) -> tuple[bytes, str | None, str | None]: - """Rewrite the ``model`` field of a Messages API request body. - - Returns ``(new_body, original_model, chosen_model)``. When the body is not a JSON - object with a ``model`` key (or the router picks the same model), ``new_body`` is the - input unchanged. Pure and side-effect-free, so it is unit-testable without a server. - """ - if not body or not is_json: - return body, None, None - try: - parsed = json.loads(body) - except ValueError: - return body, None, None - if not isinstance(parsed, dict) or "model" not in parsed: - return body, None, None - original = parsed.get("model") - chosen = route(parsed) - if not chosen or chosen == original: - return body, original if isinstance(original, str) else None, chosen - parsed["model"] = chosen - return json.dumps(parsed).encode(), original if isinstance(original, str) else None, chosen - - -def _is_json_request(handler: BaseHTTPRequestHandler) -> bool: - """True when the request body is plain (uncompressed) JSON we can safely rewrite.""" - if "content-encoding" in {k.lower() for k in handler.headers}: - return False # compressed body: don't touch it - return "json" in handler.headers.get("Content-Type", "").lower() - - -def _forwarded_headers(handler: BaseHTTPRequestHandler) -> dict[str, str]: - return {k: v for k, v in handler.headers.items() if k.lower() not in _HOP_BY_HOP} - - -class _RouterProxyHandler(BaseHTTPRequestHandler): - # Set by the server factory. - client: httpx.Client - route: RouteFn - log_fn: Callable[[str], None] - - def log_message(self, format: str, *args: object) -> None: - return - - def _safe_send_error(self, code: int, message: str) -> None: - try: - self.send_error(code, message) - except OSError: - pass - - def _handle(self) -> None: - length = int(self.headers.get("Content-Length", 0) or 0) - body = self.rfile.read(length) if length else b"" - url = self.path.lstrip("/") - - new_body = body - original = chosen = None - if body and self.command == "POST": - new_body, original, chosen = rewrite_model( - body, is_json=_is_json_request(self), route=self.route - ) - # One line per request (model ids + path only, never prompt text) so the log shows - # ALL traffic — confirming requests reach the proxy, not just rewrites. - action = "ROUTE" if (chosen and chosen != original) else "PASS" - self.log_fn(f"[{action}] {self.command} /{url} model={original!r}->{chosen!r}") - - try: - with self.client.stream( - self.command, url, headers=_forwarded_headers(self), content=new_body or None - ) as resp: - self.log_fn(f"[UPSTREAM] /{url} {resp.status_code}") - self._relay_response(resp) - except (BrokenPipeError, ConnectionResetError): - return # client (Claude Code) closed before/while relaying — routine on cancel - except httpx.HTTPError as exc: - self.log_fn(f"[ERR] /{url} {type(exc).__name__}: {exc}") - self._safe_send_error(502, "model router proxy upstream error") - - def _relay_response(self, resp: httpx.Response) -> None: - # Stream chunks as they arrive so SSE token streaming is not buffered; iter_raw - # preserves any Content-Encoding verbatim (we relay that header). - try: - self.send_response(resp.status_code) - for key, value in resp.headers.items(): - if key.lower() not in _HOP_BY_HOP: - self.send_header(key, value) - self.end_headers() - for chunk in resp.iter_raw(): - if chunk: - self.wfile.write(chunk) - self.wfile.flush() - except (BrokenPipeError, ConnectionResetError): - return # client closed mid-response — routine on cancelled turns / SSE teardown - except httpx.HTTPError: - return # upstream dropped mid-stream; headers already sent, stop cleanly - - # Transparent pass-through for every method (GET model discovery, POST messages, …). - def __getattr__(self, name: str): - if name.startswith("do_"): - return self._handle - raise AttributeError(name) - - -def start_proxy( - workspace: str, - route: RouteFn, - *, - log_path: Path | None = None, -) -> tuple[ThreadingHTTPServer, httpx.Client, int]: - """Start the loopback model-routing proxy on an OS-assigned port. - - Forwards to ``{workspace}/ai-gateway/anthropic/`` with headers passed through - verbatim (Claude Code's apiKeyHelper credential included). Returns - ``(server, client, port)``; the caller runs ``server`` (e.g. in a thread) and calls - ``server.shutdown()`` / ``client.close()`` on exit. Point ``ANTHROPIC_BASE_URL`` at - ``http://127.0.0.1:{port}``. - """ - - 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 - - upstream_base = f"{workspace.rstrip('/')}/ai-gateway/anthropic/" - client = httpx.Client(base_url=upstream_base, timeout=_UPSTREAM_TIMEOUT, follow_redirects=False) - handler = type( - "BoundRouterProxyHandler", - (_RouterProxyHandler,), - {"client": client, "route": staticmethod(route), "log_fn": staticmethod(log)}, - ) - server = ThreadingHTTPServer(("127.0.0.1", 0), handler) - port = server.server_address[1] - log(f"[READY] 127.0.0.1:{port} -> {upstream_base}") - return server, client, port diff --git a/src/ucode/smart_routing/claude_pty.py b/src/ucode/smart_routing/claude_pty.py index 6d95fb9a..0042d436 100644 --- a/src/ucode/smart_routing/claude_pty.py +++ b/src/ucode/smart_routing/claude_pty.py @@ -3,21 +3,18 @@ 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* -by typing ``/model `` into the TUI and auto-confirming the "Switch model?" dialog -by watching the PTY output. A Unix-domain control socket accepts line-delimited JSON-RPC -``model.set`` requests (parity with the reference POC and a clean test surface); the -automatic first-prompt CUJ drives the same ``on_model_set`` path internally. +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``). -Detecting *when the user submits their first prompt* is done from **stdin** (the Enter -keystroke after typed content), not by scraping output — Claude Code's TUI renders spaces -as cursor moves and uses randomized spinner text, so output markers are unreliable. Marker -matching that remains (readiness, the confirm dialog) is whitespace-insensitive (see -:func:`_squash`). The only runtime-spike output constants left are :data:`READY_MARKERS` -(pre-emptive mode only) and :data:`ConfirmationState.PROMPT_MARKERS`. +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 @@ -39,18 +36,16 @@ from collections.abc import Callable from pathlib import Path -# --- runtime-spike markers (whitespace-squashed substrings; tune from a real claude launch) --- -# TUI is ready for input. A *bonus* signal for the pre-emptive switch mode; the primary -# readiness signal is version-independent (idle after the initial paint — see READY_QUIET_S). -# Matching is whitespace-insensitive (see _squash), because Claude Code renders spaces as -# cursor moves that vanish under strip_ansi. +# 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 = 5.0 -# Pre-emptive readiness heuristic: once the TUI has produced output and then stays quiet -# this long, its initial paint is done and it is idle waiting for input — safe to type -# `/model`. Marker-independent, so it survives Claude Code TUI changes. +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 @@ -95,13 +90,17 @@ def valid_model_name(name: object) -> bool: class ConfirmationState: """Watch PTY output for Claude's "Switch model?" dialog and auto-confirm it. - ``arm`` after typing ``/model`` (with a deadline), feed each output chunk to - ``observe``; when the confirmation prompt is seen it returns the keystrokes to - send (``b"\\r"``) and disarms. A rolling buffer keeps the last ``window`` chars so + ``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. """ - PROMPT_MARKERS = ("Switch model?", "Yes, switch to") + # 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 = "" @@ -123,7 +122,7 @@ def observe(self, chunk: bytes, now: float) -> bytes | None: self.clear() return None self._buf = (self._buf + _match_text(chunk))[-self._window :] - if any(_squash(marker) in self._buf for marker in self.PROMPT_MARKERS): + if all(_squash(marker) in self._buf for marker in self.PROMPT_MARKERS): self.clear() return b"\r" return None @@ -152,10 +151,55 @@ def observe(self, chunk: bytes) -> bool: return self.triggered -def inject_model_switch(master_fd: int, model: str, confirm: ConfirmationState, now: float) -> None: - """Type ``/model \\r`` into the TUI and arm the confirm watcher.""" - os.write(master_fd, f"/model {model}\r".encode()) - confirm.arm(now + CONFIRM_TIMEOUT_S) +class ModelPickerRows: + """Discover the routed and currently focused row numbers from picker output.""" + + def __init__(self, model: str, window: int = 16384) -> None: + self._target = _squash(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 = list(re.finditer(rf"(\d+)\.{re.escape(self._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: @@ -167,6 +211,127 @@ def inject_note(out_fd: int, message: str) -> None: 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 @@ -313,18 +478,17 @@ def sync_winsize(master_fd: int, stdin_fd: int = 0) -> None: def run_claude_pty( argv: list[str], *, - target_model: str, + route_prompt: Callable[[str], str], switch_message: str, socket_path: Path, log_path: Path | None = None, - switch_mode: str = "preemptive", ) -> int: """Spawn *argv* in a PTY and run the smart-router CUJ, returning the child exit code. - Pumps stdin<->master and master<->stdout; runs the JSON-RPC control socket; and - auto-switches to ``target_model`` — in ``"reactive"`` mode on the first prompt the - user submits (detected from the Enter keystroke on stdin), or in ``"preemptive"`` - mode once the TUI paints and goes idle — surfacing ``switch_message`` in the transcript. + 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: @@ -340,27 +504,43 @@ def log(message: str) -> None: # 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() - ready = OutputMarkerDetector(READY_MARKERS) - lock = threading.Lock() - switched = {"done": False} stop = threading.Event() + pending_lock = threading.Lock() + pending: dict[str, tuple[str, str] | None] = {"value": None} - def switch_to(model: str) -> None: - with lock: - if switched["done"]: - return - switched["done"] = True - inject_note(1, switch_message) - inject_model_switch(master_fd, model, confirm, time.monotonic()) - log(f"[SWITCH] -> {model!r}") + def on_blocked_prompt(prompt: str, model: str) -> None: + with pending_lock: + pending["value"] = (prompt, model) + log(f"[ROUTE] first prompt -> {model!r}") - serve_control_socket(socket_path, switch_to, stop, log=log) + 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) @@ -371,19 +551,19 @@ def on_winch(_signum: int, _frame: object) -> None: sync_winsize(master_fd) stdin_open = True last_output = 0.0 # monotonic of the most recent TUI output (0 = none yet) - typed_content = False # printable input seen since the last Enter (reactive mode) - pending_switch = False # first prompt submitted; switch when Claude next goes idle + 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] - # Both modes fire on an idle gap, so they need periodic wakes while waiting: - # preemptive from the start, reactive only after the first prompt is submitted. - need_idle_wake = not switched["done"] and ( - switch_mode == "preemptive" - or (switch_mode == "reactive" and pending_switch) - ) - timeout = SELECT_TIMEOUT_S if need_idle_wake else None try: - ready_fds, _, _ = select.select(readable, [], [], timeout) + ready_fds, _, _ = select.select(readable, [], [], SELECT_TIMEOUT_S) except InterruptedError: # SIGWINCH etc. continue @@ -396,16 +576,6 @@ def on_winch(_signum: int, _frame: object) -> None: stdin_open = False # EOF on stdin: stop selecting it, keep pumping else: os.write(master_fd, data) - # Reactive: the user submits their first prompt when they press Enter - # after typing. This is a stdin keystroke (fully under our control) — - # far more reliable than scraping Claude's rendered output. Arm here; - # the actual /model keystroke fires once Claude goes idle (below), so it - # lands in an idle input box rather than queued mid-response. - if switch_mode == "reactive" and not switched["done"] and not pending_switch: - if any(byte >= 0x20 and byte != 0x7F for byte in data): - typed_content = True - if typed_content and (b"\r" in data or b"\n" in data): - pending_switch = True if master_fd in ready_fds: try: @@ -418,22 +588,95 @@ def on_winch(_signum: int, _frame: object) -> None: last_output = time.monotonic() if debug: log(f"[OUT] {strip_ansi(chunk)[:400]!r}") - with lock: - keystroke = confirm.observe(chunk, last_output) + keystroke = confirm.observe(chunk, last_output) if keystroke is not None: os.write(master_fd, keystroke) - ready.observe(chunk) - - # Fire on an idle gap: preemptive before any prompt (TUI painted + idle); - # reactive only after the first prompt was submitted (turn done + idle), so the - # /model command types into an idle input box. - if not switched["done"]: - now = time.monotonic() - idle = last_output > 0.0 and (now - last_output) >= READY_QUIET_S - if switch_mode == "preemptive" and (ready.triggered or idle): - switch_to(target_model) - elif switch_mode == "reactive" and pending_switch and idle: - switch_to(target_model) + 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): diff --git a/tests/test_claude_proxy.py b/tests/test_claude_proxy.py deleted file mode 100644 index f6439837..00000000 --- a/tests/test_claude_proxy.py +++ /dev/null @@ -1,58 +0,0 @@ -"""Tests for the smart-routing-v2 model-routing proxy's request rewriting.""" - -from __future__ import annotations - -import json - -from ucode.smart_routing import claude_proxy - - -def _stub(model: str): - return lambda _body: model - - -class TestRewriteModel: - def test_rewrites_model_in_json_body(self): - body = json.dumps({"model": "opus", "messages": [{"role": "user", "content": "hi"}]}).encode() - new_body, original, chosen = claude_proxy.rewrite_model( - body, is_json=True, route=_stub("system.ai.claude-sonnet-5") - ) - assert original == "opus" - assert chosen == "system.ai.claude-sonnet-5" - assert json.loads(new_body)["model"] == "system.ai.claude-sonnet-5" - # Everything else in the body is preserved. - assert json.loads(new_body)["messages"] == [{"role": "user", "content": "hi"}] - - def test_unchanged_when_router_picks_same_model(self): - body = json.dumps({"model": "system.ai.claude-sonnet-5"}).encode() - new_body, original, chosen = claude_proxy.rewrite_model( - body, is_json=True, route=_stub("system.ai.claude-sonnet-5") - ) - assert new_body == body - assert original == "system.ai.claude-sonnet-5" - assert chosen == "system.ai.claude-sonnet-5" - - def test_unchanged_when_not_json(self): - body = b"\x00\x01 not json" - new_body, original, chosen = claude_proxy.rewrite_model( - body, is_json=False, route=_stub("x") - ) - assert new_body == body - assert original is None and chosen is None - - def test_unchanged_when_body_has_no_model(self): - body = json.dumps({"messages": []}).encode() - new_body, original, chosen = claude_proxy.rewrite_model(body, is_json=True, route=_stub("x")) - assert new_body == body - assert original is None and chosen is None - - def test_unchanged_on_empty_body(self): - new_body, original, chosen = claude_proxy.rewrite_model(b"", is_json=True, route=_stub("x")) - assert new_body == b"" - assert original is None and chosen is None - - def test_malformed_json_passes_through(self): - body = b'{"model": "opus"' # truncated - new_body, original, chosen = claude_proxy.rewrite_model(body, is_json=True, route=_stub("y")) - assert new_body == body - assert original is None and chosen is None diff --git a/tests/test_claude_smart_routing_v2.py b/tests/test_claude_smart_routing_v2.py index 7f3f6741..937a50c6 100644 --- a/tests/test_claude_smart_routing_v2.py +++ b/tests/test_claude_smart_routing_v2.py @@ -8,14 +8,16 @@ 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_pty +from ucode.smart_routing import claude_hooks, claude_pty class TestStripAnsi: @@ -62,7 +64,15 @@ def test_marker_split_across_chunks(self): 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) == b"\r" + 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() @@ -142,14 +152,26 @@ def test_non_object_is_invalid_request(self): class TestInjectors: - def test_inject_model_switch_types_command_and_arms(self): + 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_inject_model_switch_opens_picker(self): read_fd, write_fd = os.pipe() try: - confirm = claude_pty.ConfirmationState() - claude_pty.inject_model_switch(write_fd, "opus", confirm, time.monotonic()) - assert os.read(read_fd, 100) == b"/model opus\r" - # Armed: it now auto-confirms the dialog. - assert confirm.observe(b"Switch model? Yes, switch to X", time.monotonic()) == b"\r" + 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) @@ -164,6 +186,24 @@ def test_inject_note_writes_message(self): 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): @@ -230,10 +270,211 @@ def test_dispatches_model_set_over_socket(self, tmp_path): 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({"model": "opus", "messages": []}) == claude.SMART_ROUTING_V2_MODEL + assert route("fix the parser") == claude.SMART_ROUTING_V2_MODEL assert claude.SMART_ROUTING_V2_MODEL == "system.ai.claude-sonnet-5"