From fafe7af7c7650c55724a329c84842bcce7375a02 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Thu, 20 Aug 2026 23:04:50 +0000 Subject: [PATCH 01/18] codex switcher --- pyproject.toml | 4 + scripts/codex_arch_b_probe.py | 222 ++++++++++++ scripts/codex_model_router_poc.py | 355 ++++++++++++++++++++ scripts/codex_tui_interposer.py | 340 +++++++++++++++++++ src/ucode/agents/codex.py | 122 ++++++- src/ucode/smart_routing/codex_interposer.py | 234 +++++++++++++ tests/test_codex_smart_routing_v2.py | 127 +++++++ uv.lock | 91 +++++ 8 files changed, 1494 insertions(+), 1 deletion(-) create mode 100755 scripts/codex_arch_b_probe.py create mode 100644 scripts/codex_model_router_poc.py create mode 100644 scripts/codex_tui_interposer.py create mode 100644 src/ucode/smart_routing/codex_interposer.py create mode 100644 tests/test_codex_smart_routing_v2.py diff --git a/pyproject.toml b/pyproject.toml index 27eb0619..17f3f683 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,6 +26,10 @@ dependencies = [ "questionary>=2.0.0", "tomlkit>=0.13.0", "typer>=0.12.0", + # WebSocket client+server for the experimental `ENABLE_SMART_ROUTING_V2` Codex launch path: + # ucode interposes on Codex's `--remote` WebSocket transport to switch the model at runtime + # (see ucode.smart_routing.codex_interposer). + "websockets>=13", ] [project.optional-dependencies] diff --git a/scripts/codex_arch_b_probe.py b/scripts/codex_arch_b_probe.py new file mode 100755 index 00000000..aaa23201 --- /dev/null +++ b/scripts/codex_arch_b_probe.py @@ -0,0 +1,222 @@ +#!/usr/bin/env python3 +"""Arch B probe: MITM proxy between TUI and app-server, rewriting model in turn/start. + +This demonstrates feasibility of interposing on the real TUI without modifying it. +The proxy: +1. Listens on a unix socket that the TUI connects to (via --remote) +2. Forwards all messages to a real app-server +3. Rewrites turn/start.model to a fixed value (proving router capability) +4. Passes everything else through unchanged +""" +import json +import os +import socket +import subprocess +import sys +import threading +import time +import argparse + +class CodexInterposer: + """MITM proxy for Codex messages.""" + + def __init__(self, listen_sock_path, app_server_sock_path, target_model): + self.listen_sock_path = listen_sock_path + self.app_server_sock_path = app_server_sock_path + self.target_model = target_model + self.listener = None + self.running = True + + def cleanup(self): + """Clean up listener socket.""" + if os.path.exists(self.listen_sock_path): + try: + os.remove(self.listen_sock_path) + except: + pass + if self.listener: + try: + self.listener.close() + except: + pass + + def rewrite_message(self, msg): + """Rewrite turn/start to force target model.""" + if not isinstance(msg, dict): + return msg + + method = msg.get("method") + if method == "turn/start": + params = msg.get("params", {}) + if isinstance(params, dict): + old_model = params.get("model") + if old_model != self.target_model: + print(f"[REWRITE] turn/start: {old_model!r} -> {self.target_model!r}") + params["model"] = self.target_model + msg["params"] = params + + return msg + + def relay_messages(self, client_sock, as_sock): + """Relay messages bidirectionally, rewriting turn/start.""" + + def tui_to_as(): + """TUI -> app-server (with rewriting)""" + buffer = "" + while self.running: + try: + data = client_sock.recv(1024) + if not data: + print("[TUI->AS] Connection closed by TUI") + break + + buffer += data.decode('utf-8', errors='replace') + + # Process complete lines + while '\n' in buffer: + line, buffer = buffer.split('\n', 1) + line = line.strip() + if not line: + continue + + try: + msg = json.loads(line) + msg = self.rewrite_message(msg) + rewritten = json.dumps(msg) + as_sock.sendall((rewritten + '\n').encode('utf-8')) + print(f"[TUI->AS] {msg.get('method', msg.get('type', '?'))}") + except Exception as e: + print(f"[TUI->AS] Error: {e}") + as_sock.sendall((line + '\n').encode('utf-8')) + except Exception as e: + print(f"[TUI->AS] Exception: {e}") + break + + def as_to_tui(): + """app-server -> TUI (pass-through)""" + buffer = "" + while self.running: + try: + data = as_sock.recv(1024) + if not data: + print("[AS->TUI] Connection closed by app-server") + break + + buffer += data.decode('utf-8', errors='replace') + + # Process complete lines + while '\n' in buffer: + line, buffer = buffer.split('\n', 1) + line = line.strip() + if not line: + continue + + try: + msg = json.loads(line) + method = msg.get('method') + if method in ('turn/start', 'turn/completed', 'item/completed'): + print(f"[AS->TUI] {method}") + except: + pass + + client_sock.sendall((line + '\n').encode('utf-8')) + except Exception as e: + print(f"[AS->TUI] Exception: {e}") + break + + t1 = threading.Thread(target=tui_to_as, daemon=True) + t2 = threading.Thread(target=as_to_tui, daemon=True) + t1.start() + t2.start() + + # Wait for either thread to finish + t1.join(timeout=300) + t2.join(timeout=300) + + self.running = False + + def handle_client(self, client_sock, addr): + """Handle a single TUI connection.""" + print(f"[CLIENT] Connected from {addr}") + + try: + # Connect to real app-server + as_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + as_sock.connect(self.app_server_sock_path) + print(f"[RELAY] Connected to app-server at {self.app_server_sock_path}") + + # Relay bidirectionally + self.relay_messages(client_sock, as_sock) + + as_sock.close() + except Exception as e: + print(f"[ERROR] Failed to relay: {e}") + finally: + client_sock.close() + print(f"[CLIENT] Disconnected") + + def run(self): + """Start the interposer listening for TUI connections.""" + self.cleanup() + + print(f"[STARTUP] Creating listener at {self.listen_sock_path}") + + self.listener = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + self.listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + self.listener.bind(self.listen_sock_path) + self.listener.listen(1) + + print(f"[STARTUP] Listening for TUI connections") + print(f"[INFO] Target model: {self.target_model!r}") + print(f"[INFO] To connect TUI: codex --remote unix://{self.listen_sock_path}") + + try: + while self.running: + try: + self.listener.settimeout(1.0) + client_sock, addr = self.listener.accept() + + # Handle in a thread + t = threading.Thread( + target=self.handle_client, + args=(client_sock, addr), + daemon=True + ) + t.start() + except socket.timeout: + continue + except KeyboardInterrupt: + print("\n[SHUTDOWN] Interrupted") + break + finally: + self.cleanup() + + +def main(): + parser = argparse.ArgumentParser( + description="Codex model interposer: MITM proxy to rewrite turn/start.model" + ) + parser.add_argument( + "--listen", + default="/home/lilly.luo/.cache/codex-b/tui-remote.sock", + help="Socket for TUI to connect to" + ) + parser.add_argument( + "--app-server", + default="/home/lilly.luo/.cache/codex-b/as.sock", + help="Real app-server socket" + ) + parser.add_argument( + "--model", + default="gpt-5.5", + help="Model to force for all turns" + ) + + args = parser.parse_args() + + interposer = CodexInterposer(args.listen, args.app_server, args.model) + interposer.run() + + +if __name__ == "__main__": + sys.exit(main() or 0) diff --git a/scripts/codex_model_router_poc.py b/scripts/codex_model_router_poc.py new file mode 100644 index 00000000..b0134d1f --- /dev/null +++ b/scripts/codex_model_router_poc.py @@ -0,0 +1,355 @@ +#!/usr/bin/env python3 +"""POC: a minimal interactive Codex client that can switch models mid-session. + +Launches `codex app-server` under the hood, gives you a prompt, and lets you +change the model live with `/model ` — the switch happens by setting the +per-turn `model` field on `turn/start`, so history is preserved across it. + +This is arch A from the plan (a thin app-server client). It is NOT Codex's +polished TUI; it's the smallest thing that proves "launch, type, switch". + +Run it with the repo's Python 3.12 venv (system python3 here is 3.6): + + /home/lilly.luo/ucode/.venv/bin/python scripts/codex_model_router_poc.py + +For interactive mode, omit all flags. For a self-test that proves mid-session +model switching with context preservation: + + /home/lilly.luo/ucode/.venv/bin/python scripts/codex_model_router_poc.py --selftest + +Auth/gateway config is generated from your existing ~/.codex/ucode.config.toml +provider block into an isolated CODEX_HOME, so it uses the same Databricks +gateway + `ucode auth-token` refresh that `ucode codex` uses. + +In-session commands: + /model switch the model for subsequent turns (e.g. /model gpt-5.5) + /model show the current model + /quit exit +""" +from __future__ import annotations + +import json +import os +import queue +import subprocess +import sys +import threading +import time +import tomllib +from pathlib import Path + +import tomlkit + +UCODE_CODEX_CONFIG = Path.home() / ".codex" / "ucode.config.toml" +POC_HOME = Path.home() / ".cache" / "ucode-codex-router-poc" +DEFAULT_MODEL = "system.ai.gpt-5-6-luna" +EXAMPLE_MODELS = ["system.ai.gpt-5-6-luna", "gpt-5.5"] + + +def build_codex_home() -> Path: + """Generate an isolated CODEX_HOME whose config.toml carries ONLY the ucode + gateway provider block (model_provider + model + [model_providers.*]), copied + from ~/.codex/ucode.config.toml. Keeps the app-server pointed at the same + Databricks gateway + auth-token refresh, without the hooks/tui cruft.""" + if not UCODE_CODEX_CONFIG.exists(): + sys.exit( + f"Missing {UCODE_CODEX_CONFIG}. Run `ucode configure codex` (or `ucode codex`) first " + "so the Databricks provider block exists." + ) + src = tomllib.loads(UCODE_CODEX_CONFIG.read_text()) + minimal = tomlkit.document() + if "model_provider" in src: + minimal["model_provider"] = src["model_provider"] + minimal["model"] = src.get("model", DEFAULT_MODEL) + if "model_reasoning_effort" in src: + minimal["model_reasoning_effort"] = src["model_reasoning_effort"] + if "model_providers" in src: + minimal["model_providers"] = src["model_providers"] + POC_HOME.mkdir(parents=True, exist_ok=True) + (POC_HOME / "config.toml").write_text(tomlkit.dumps(minimal)) + return POC_HOME + + +class AppServer: + """Thin newline-delimited-JSON stdio client for `codex app-server`.""" + + def __init__(self, codex_home: Path) -> None: + env = dict(os.environ) + env["CODEX_HOME"] = str(codex_home) + self.proc = subprocess.Popen( + ["codex", "app-server"], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + bufsize=1, + env=env, + ) + self._q: queue.Queue = queue.Queue() + self._id = 0 + threading.Thread(target=self._read_stdout, daemon=True).start() + threading.Thread(target=self._drain_stderr, daemon=True).start() + + def _read_stdout(self) -> None: + for line in self.proc.stdout: # type: ignore[union-attr] + line = line.strip() + if line: + try: + self._q.put(json.loads(line)) + except ValueError: + pass + + def _drain_stderr(self) -> None: + # app-server logs benign catalog-refresh 404s here; keep them out of the UI + # but available if the user wants them (uncomment to debug). + for _line in self.proc.stderr: # type: ignore[union-attr] + pass + + def _send(self, method: str, params: dict | None = None, *, notify: bool = False): + msg: dict = {"method": method} + if not notify: + self._id += 1 + msg["id"] = self._id + if params is not None: + msg["params"] = params + self.proc.stdin.write(json.dumps(msg) + "\n") # type: ignore[union-attr] + self.proc.stdin.flush() # type: ignore[union-attr] + return msg.get("id") + + def _wait(self, pred, timeout: float): + end = time.time() + timeout + while time.time() < end: + try: + msg = self._q.get(timeout=min(1.0, max(0.05, end - time.time()))) + except queue.Empty: + continue + if pred(msg): + return msg + return None + + def request(self, method: str, params: dict | None = None, *, timeout: float = 60.0): + rid = self._send(method, params) + return self._wait(lambda m: m.get("id") == rid and ("result" in m or "error" in m), timeout) + + def initialize(self) -> None: + self.request( + "initialize", + {"clientInfo": {"name": "ucode-codex-router-poc", "version": "0.1"}, "capabilities": {}}, + timeout=30, + ) + self._send("initialized", {}, notify=True) + + def start_thread(self, model: str) -> str: + resp = self.request( + "thread/start", {"model": model, "cwd": os.getcwd(), "approvalPolicy": "never"}, timeout=60 + ) + result = (resp or {}).get("result", {}) + tid = result.get("thread", {}).get("id") or result.get("threadId") + if not tid: + sys.exit(f"thread/start failed: {json.dumps(resp)[:400]}") + return tid + + def run_turn(self, thread_id: str, text: str, model: str, *, timeout: float = 300.0) -> None: + """Send one user turn on `model`, streaming assistant text to stdout live.""" + rid = self._send( + "turn/start", + {"threadId": thread_id, "input": [{"type": "text", "text": text}], "model": model}, + ) + # Ack (status inProgress) — then stream until turn/completed. + self._wait(lambda m: m.get("id") == rid and ("result" in m or "error" in m), 30) + end = time.time() + timeout + printed_any = False + while time.time() < end: + try: + msg = self._q.get(timeout=min(1.0, max(0.05, end - time.time()))) + except queue.Empty: + continue + method = msg.get("method") + params = msg.get("params") or {} + if method == "item/agentMessage/delta": + delta = _find_str(params, ("delta", "text")) + if delta: + sys.stdout.write(delta) + sys.stdout.flush() + printed_any = True + elif method == "turn/completed": + turn = params.get("turn", {}) + if turn.get("status") == "failed": + err = turn.get("error", {}) + print(f"\n [turn failed: {err.get('message', err)}]") + elif not printed_any: + # No deltas seen (some models don't stream) — print final items. + print(_final_text(turn) or " [no text returned]") + print() + return + print("\n [timed out waiting for the turn to complete]") + + def close(self) -> None: + try: + self.proc.stdin.close() # type: ignore[union-attr] + except Exception: + pass + self.proc.terminate() + try: + self.proc.wait(timeout=5) + except Exception: + self.proc.kill() + + +def _find_str(obj, keys) -> str | None: + if isinstance(obj, dict): + for k, v in obj.items(): + if k in keys and isinstance(v, str): + return v + r = _find_str(v, keys) + if r: + return r + elif isinstance(obj, list): + for v in obj: + r = _find_str(v, keys) + if r: + return r + return None + + +def _final_text(turn: dict) -> str: + out = [] + for item in turn.get("items", []) or []: + if isinstance(item, dict) and item.get("type") == "agentMessage": + t = item.get("text") + if t: + out.append(t) + return "\n".join(out) + + +def _capture_turn_text(server: AppServer, thread_id: str, text: str, model: str) -> str: + """Run a turn and capture the full assistant response text.""" + rid = server._send( + "turn/start", + {"threadId": thread_id, "input": [{"type": "text", "text": text}], "model": model}, + ) + # Ack (status inProgress) — then stream until turn/completed. + server._wait(lambda m: m.get("id") == rid and ("result" in m or "error" in m), 30) + captured_text = [] + end = time.time() + 300.0 + while time.time() < end: + try: + msg = server._q.get(timeout=min(1.0, max(0.05, end - time.time()))) + except queue.Empty: + continue + method = msg.get("method") + params = msg.get("params") or {} + if method == "item/agentMessage/delta": + delta = _find_str(params, ("delta", "text")) + if delta: + captured_text.append(delta) + elif method == "turn/completed": + turn = params.get("turn", {}) + if turn.get("status") == "failed": + err = turn.get("error", {}) + return f"[FAILED: {err.get('message', err)}]" + # Collect any remaining text from final items + final = _final_text(turn) + if final and not captured_text: + captured_text.append(final) + return "".join(captured_text) + return "[TIMEOUT]" + + +def selftest() -> int: + """Non-interactive self-test: prove mid-session model switch with context.""" + home = build_codex_home() + server = AppServer(home) + try: + print("Starting codex app-server for self-test…") + server.initialize() + thread_id = server.start_thread(DEFAULT_MODEL) + print(f"Thread created with model {DEFAULT_MODEL}") + + # Turn 1: simple model A request + print("\n=== Turn 1 (model A) ===") + t1_prompt = "Reply with exactly: TURN1_OK" + print(f"Prompt: {t1_prompt}") + t1_response = _capture_turn_text(server, thread_id, t1_prompt, "system.ai.gpt-5-6-luna") + print(f"Response: {t1_response!r}") + if "TURN1_OK" not in t1_response: + print(f"ERROR: Turn 1 did not contain TURN1_OK") + return 1 + + # Turn 2: switch model and test context preservation + print("\n=== Turn 2 (model B, testing context) ===") + t2_prompt = "What token did you reply on the previous turn? Then say TURN2_OK." + print(f"Switching to gpt-5.5…") + print(f"Prompt: {t2_prompt}") + t2_response = _capture_turn_text(server, thread_id, t2_prompt, "gpt-5.5") + print(f"Response: {t2_response!r}") + + # Verify context was preserved: t2 should mention TURN1_OK + if "TURN1_OK" not in t2_response: + print(f"ERROR: Turn 2 did not contain TURN1_OK (context not preserved)") + return 1 + + if "TURN2_OK" not in t2_response: + print(f"WARNING: Turn 2 did not contain TURN2_OK (but context was preserved)") + + print("\n=== SUCCESS ===") + print("Mid-session model switch with context preservation verified!") + return 0 + except Exception as e: + print(f"ERROR: {e}", file=sys.stderr) + import traceback + traceback.print_exc(file=sys.stderr) + return 1 + finally: + server.close() + + +def main() -> int: + # Parse command-line arguments + if len(sys.argv) > 1: + if sys.argv[1] == "--selftest": + return selftest() + elif sys.argv[1] in ("--help", "-h"): + print(__doc__) + return 0 + else: + print(f"Unknown argument: {sys.argv[1]}", file=sys.stderr) + print(f"Use: {sys.argv[0]} [--selftest] [--help]", file=sys.stderr) + return 1 + + # Interactive mode + home = build_codex_home() + server = AppServer(home) + current_model = DEFAULT_MODEL + try: + print("Starting codex app-server…") + server.initialize() + thread_id = server.start_thread(current_model) + print(f"\nCodex ready. model = {current_model}") + print(f"Commands: /model /quit (try: {', '.join(EXAMPLE_MODELS)})\n") + while True: + try: + line = input(f"[{current_model}] › ").strip() + except (EOFError, KeyboardInterrupt): + print() + break + if not line: + continue + if line == "/quit": + break + if line.startswith("/model"): + arg = line[len("/model"):].strip() + if not arg: + print(f" current model: {current_model}") + else: + current_model = arg + print(f" → switched to {current_model} (applies to the next turn; history kept)") + continue + server.run_turn(thread_id, line, current_model) + return 0 + finally: + server.close() + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/codex_tui_interposer.py b/scripts/codex_tui_interposer.py new file mode 100644 index 00000000..081f8795 --- /dev/null +++ b/scripts/codex_tui_interposer.py @@ -0,0 +1,340 @@ +#!/usr/bin/env python3 +"""Arch B: a WebSocket MITM that lets you keep the REAL Codex TUI while the model +is switched under program control. + +Codex's remote transport (`codex --remote ws://…`) is WebSocket (a plain-JSONL +client is rejected with HTTP 400 "Connection header did not include 'upgrade'"; +a proper upgrade returns 101). Each JSON-RPC message is one WebSocket text frame. +This proxy sits between the TUI and a real `codex app-server`, forwarding every +frame untouched except: + + - `turn/start` (TUI->engine): after an initial hold of `--after` turns, its + `model` is rewritten to `--model`. `turn/start.model` is documented as + "override the model for this turn and subsequent turns", so the live session + retargets with history preserved. + - When the hold expires (right after your Nth prompt completes) it INJECTS a + `thread/settings/updated` notification (engine->TUI) carrying the new model, + so the TUI's on-screen model indicator follows the switch. + +So the demo is: start the TUI on model X, submit your first prompt (answered by +X), and from then on the session runs on `--model` (and the chip flips to it). + +Topology: + codex app-server --listen ws://127.0.0.1:8801 (real engine) + this interposer ws://127.0.0.1:8802 -> ws://127.0.0.1:8801 (switches model) + codex --remote ws://127.0.0.1:8802 --model system.ai.gpt-5-6-luna (real TUI) + +Run via uv so nothing is installed globally: + + uv run --with websockets python scripts/codex_tui_interposer.py \ + --listen 127.0.0.1:8802 --upstream ws://127.0.0.1:8801 \ + --model gpt-5.5 --after 1 + +Self-test (spawns its own app-server + a simulated TUI; proves hold + switch end +to end against the gateway): + + uv run --with websockets --with tomlkit python \ + scripts/codex_tui_interposer.py --selftest +""" +from __future__ import annotations + +import argparse +import asyncio +import contextlib +import json +import os +import socket +import subprocess +import sys +import time +from pathlib import Path + +from websockets.asyncio.client import connect +from websockets.asyncio.server import serve + +SETTINGS_UPDATED = "thread/settings/updated" + + +class Session: + """Per-TUI-connection state: hold the first `after` turns, then switch model.""" + + def __init__(self, target_model: str, after: int, log) -> None: + self.target = target_model + self.after = after + self.log = log + self.turns = 0 + self.thread_id: str | None = None + self.settings: dict | None = None + self.injected = False + + def on_tui_frame(self, raw: str) -> str: + """TUI->engine: rewrite turn/start.model once past the hold.""" + try: + msg = json.loads(raw) + except ValueError: + return raw + if not isinstance(msg, dict): + return raw + params = msg.get("params") + if msg.get("method") == "turn/start" and isinstance(params, dict): + self.turns += 1 + if isinstance(params.get("threadId"), str): + self.thread_id = params["threadId"] + if self.turns > self.after: + old = params.get("model") + if old != self.target: + params["model"] = self.target + self.log(f"[REWRITE] turn #{self.turns}: model {old!r} -> {self.target!r}") + return json.dumps(msg) + return raw + + def on_engine_frame(self, raw: str): + """engine->TUI: capture thread id/settings; after the hold's last turn + completes, return an injected settings-updated notification (or None).""" + try: + msg = json.loads(raw) + except ValueError: + return None + if not isinstance(msg, dict): + return None + params = msg.get("params") if isinstance(msg.get("params"), dict) else {} + result = msg.get("result") if isinstance(msg.get("result"), dict) else {} + # Capture threadId + a real threadSettings object wherever it appears. + for src in (params, result): + tid = src.get("threadId") or (src.get("thread") or {}).get("id") + if isinstance(tid, str): + self.thread_id = tid + ts = src.get("threadSettings") + if isinstance(ts, dict): + self.settings = ts + # When the hold's final turn completes, flip the on-screen model. + if ( + msg.get("method") == "turn/completed" + and not self.injected + and self.turns >= self.after + and self.thread_id + ): + self.injected = True + settings = dict(self.settings) if isinstance(self.settings, dict) else {} + settings["model"] = self.target + self.log(f"[INJECT] {SETTINGS_UPDATED}: model -> {self.target!r} (flip TUI chip)") + return { + "method": SETTINGS_UPDATED, + "params": {"threadId": self.thread_id, "threadSettings": settings}, + } + return None + + +async def _handle_tui(tui, upstream_uri: str, target_model: str, after: int, log) -> None: + path = getattr(getattr(tui, "request", None), "path", "/") or "/" + uri = upstream_uri.rstrip("/") + path + log(f"[CONN] TUI connected (path={path}); dialing app-server {uri}") + sess = Session(target_model, after, log) + async with connect(uri, max_size=None) as upstream: + + async def tui_to_app(): + async for frame in tui: + if isinstance(frame, str): + frame = sess.on_tui_frame(frame) + await upstream.send(frame) + + async def app_to_tui(): + async for frame in upstream: + await tui.send(frame) + if isinstance(frame, str): + inj = sess.on_engine_frame(frame) + if inj is not None: + await tui.send(json.dumps(inj)) + + a = asyncio.create_task(tui_to_app()) + b = asyncio.create_task(app_to_tui()) + _done, pending = await asyncio.wait({a, b}, return_when=asyncio.FIRST_COMPLETED) + for t in pending: + t.cancel() + with contextlib.suppress(asyncio.CancelledError): + await t + log("[CONN] TUI session closed") + + +async def serve_interposer(host: str, port: int, upstream_uri: str, model: str, after: int, *, quiet=False): + def log(m: str) -> None: + if not quiet: + print(m, file=sys.stderr, flush=True) + + async def handler(tui): + try: + await _handle_tui(tui, upstream_uri, model, after, log) + except Exception as exc: # noqa: BLE001 - one session must not kill the server + log(f"[ERR] session: {exc!r}") + + server = await serve(handler, host, port, max_size=None) + log(f"[READY] ws://{host}:{port} -> {upstream_uri} (hold {after} turn(s), then switch to {model!r})") + return server + + +# --------------------------------------------------------------------------- # +# Self-test +# --------------------------------------------------------------------------- # + +UCODE_CODEX_CONFIG = Path.home() / ".codex" / "ucode.config.toml" +SELFTEST_HOME = Path.home() / ".cache" / "ucode-codex-interposer" +START_MODEL = "system.ai.gpt-5-6-luna" +TARGET_MODEL = "gpt-5.5" +BOGUS_MODEL = "totally-bogus-model-zzz" + + +def _free_port() -> int: + s = socket.socket(); s.bind(("127.0.0.1", 0)); p = s.getsockname()[1]; s.close(); return p + + +def _build_codex_home() -> Path: + import tomllib + import tomlkit + + if not UCODE_CODEX_CONFIG.exists(): + sys.exit(f"Missing {UCODE_CODEX_CONFIG}; run `ucode codex` once so the provider block exists.") + src = tomllib.loads(UCODE_CODEX_CONFIG.read_text()) + doc = tomlkit.document() + for k in ("model_provider", "model", "model_reasoning_effort", "model_providers"): + if k in src: + doc[k] = src[k] + SELFTEST_HOME.mkdir(parents=True, exist_ok=True) + (SELFTEST_HOME / "config.toml").write_text(tomlkit.dumps(doc)) + return SELFTEST_HOME + + +async def _wait_healthz(port: int, timeout: float = 30.0) -> bool: + import urllib.request + + end = time.time() + timeout + while time.time() < end: + try: + with urllib.request.urlopen(f"http://127.0.0.1:{port}/healthz", timeout=1) as r: + if r.status == 200: + return True + except Exception: + await asyncio.sleep(0.25) + return False + + +async def _simulated_tui(port: int) -> dict: + """Turn 1 uses START_MODEL (should pass through). Turn 2 sends a BOGUS model + (should be rewritten to TARGET_MODEL and therefore succeed). Also watches for + the injected settings-updated frame after turn 1.""" + out = {"t1": None, "t2": None, "injected_model": None, "error": None} + nid = 0 + async with connect(f"ws://127.0.0.1:{port}", max_size=None) as ws: + async def send(method, params=None, notify=False): + nonlocal nid + m = {"method": method} + if not notify: + nid += 1 + m["id"] = nid + if params is not None: + m["params"] = params + await ws.send(json.dumps(m)) + return m.get("id") + + async def until(pred, timeout=180): + end = time.time() + timeout + while time.time() < end: + try: + frame = await asyncio.wait_for(ws.recv(), timeout=min(5, end - time.time())) + except asyncio.TimeoutError: + continue + if not isinstance(frame, str): + continue + try: + msg = json.loads(frame) + except ValueError: + continue + if msg.get("method") == SETTINGS_UPDATED: + out["injected_model"] = (msg.get("params", {}).get("threadSettings", {}) or {}).get("model") + if pred(msg): + return msg + return None + + await send("initialize", {"clientInfo": {"name": "sim", "version": "0"}, "capabilities": {}}) + await until(lambda m: m.get("id") == 1 and ("result" in m or "error" in m), 30) + await send("initialized", {}, notify=True) + rid = await send("thread/start", {"model": START_MODEL, "cwd": os.getcwd(), "approvalPolicy": "never"}) + ts = await until(lambda m: m.get("id") == rid and "result" in m, 60) + tid = ((ts or {}).get("result", {}).get("thread", {}) or {}).get("id") + if not tid: + out["error"] = f"thread/start failed: {json.dumps(ts)[:200]}" + return out + await send("turn/start", {"threadId": tid, "input": [{"type": "text", "text": "Say A"}], "model": START_MODEL}) + tc1 = await until(lambda m: m.get("method") == "turn/completed", 180) + out["t1"] = (tc1 or {}).get("params", {}).get("turn", {}).get("status") + await send("turn/start", {"threadId": tid, "input": [{"type": "text", "text": "Say B"}], "model": BOGUS_MODEL}) + tc2 = await until(lambda m: m.get("method") == "turn/completed", 180) + out["t2"] = (tc2 or {}).get("params", {}).get("turn", {}).get("status") + return out + + +async def _selftest() -> int: + home = _build_codex_home() + port_a, port_b = _free_port(), _free_port() + env = dict(os.environ); env["CODEX_HOME"] = str(home) + print(f"Starting codex app-server on ws://127.0.0.1:{port_a} …") + proc = subprocess.Popen( + ["codex", "app-server", "--listen", f"ws://127.0.0.1:{port_a}"], + stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, env=env, + ) + server = None + try: + if not await _wait_healthz(port_a): + print("app-server did not become healthy", file=sys.stderr) + return 1 + server = await serve_interposer("127.0.0.1", port_b, f"ws://127.0.0.1:{port_a}", TARGET_MODEL, after=1) + print(f"Interposer up: hold 1 turn on the TUI's model, then switch -> {TARGET_MODEL!r}\n") + r = await _simulated_tui(port_b) + print() + ok = ( + r["t1"] == "completed" # turn 1 ran on the pass-through START_MODEL + and r["t2"] == "completed" # turn 2 sent BOGUS but was rewritten -> succeeded + and r["injected_model"] == TARGET_MODEL # settings-updated injected to flip the chip + ) + print("=== RESULT ===") + print(f" turn1 (start model, passthrough): {r['t1']}") + print(f" turn2 (client sent BOGUS -> rewritten): {r['t2']}") + print(f" injected settings-updated model: {r['injected_model']!r}") + print(f" error: {r['error']!r}") + print("=== SUCCESS ===" if ok else "=== FAILED ===") + return 0 if ok else 1 + finally: + if server is not None: + server.close() + with contextlib.suppress(Exception): + await server.wait_closed() + proc.terminate() + with contextlib.suppress(Exception): + proc.wait(timeout=5) + + +def main() -> int: + ap = argparse.ArgumentParser(description="WebSocket MITM interposer for the Codex TUI (arch B).") + ap.add_argument("--listen", default="127.0.0.1:8802", help="host:port for the TUI to connect to") + ap.add_argument("--upstream", default="ws://127.0.0.1:8801", help="real app-server ws:// URI") + ap.add_argument("--model", default=TARGET_MODEL, help="model to switch to after the hold") + ap.add_argument("--after", type=int, default=1, help="pass through this many turns before switching (default 1)") + ap.add_argument("--selftest", action="store_true", help="spawn app-server + simulated TUI and prove hold+switch") + args = ap.parse_args() + + if args.selftest: + return asyncio.run(_selftest()) + + host, _, port = args.listen.partition(":") + + async def _run(): + await serve_interposer(host, int(port), args.upstream, args.model, args.after) + await asyncio.Future() + + try: + return asyncio.run(_run()) or 0 + except KeyboardInterrupt: + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 2b7415e2..071f897f 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -4,6 +4,7 @@ import os import re +import signal import subprocess import sys import time @@ -33,7 +34,7 @@ ) from ucode.state import mark_tool_managed, save_state from ucode.telemetry import agent_version, ucode_version -from ucode.ui import print_warning_err +from ucode.ui import print_note, print_warning_err CODEX_CONFIG_DIR = Path.home() / ".codex" CODEX_PROFILE_NAME = "ucode" @@ -50,6 +51,19 @@ # tool (codex, claude), so a workspace turns it on once. SMART_ROUTING_STATE_KEY = "smart_routing_enabled" +# Smart routing v2 (experimental, env-gated). When ENABLE_SMART_ROUTING_V2=1, a single +# `ucode codex` launches the REAL Codex TUI against a ucode-run `codex app-server`, with a +# WebSocket interposer (see smart_routing.codex_interposer) that holds the first turn on the +# normal model then switches to a fixed target. ucode owns all three processes and tears the +# app-server + interposer down when the TUI exits. +SMART_ROUTING_V2_ENV_VAR = "ENABLE_SMART_ROUTING_V2" +SMART_ROUTING_V2_TARGET_MODEL = "gpt-5.5" # hardcoded switch-to model for now +SMART_ROUTING_V2_AFTER = 1 # pass through this many turns before switching +SMART_ROUTING_V2_HOME = APP_DIR / "codex-v2-home" # CODEX_HOME for the ucode-run app-server +SMART_ROUTING_V2_LOG = ( + APP_DIR / "codex-v2-interposer.log" +) # interposer log (not stdout: TUI owns it) + SPEC: ToolSpec = { "binary": "codex", @@ -465,9 +479,115 @@ def _gpt_version_key(entry: tuple[str, tuple[int, int | None, int | None, str]]) _PROFILE_REJECTED_MAX_SECONDS = 3.0 +def smart_routing_v2_enabled() -> bool: + """Return whether the experimental smart-routing-v2 launch path is enabled.""" + return os.environ.get(SMART_ROUTING_V2_ENV_VAR) == "1" + + +def _generate_v2_app_server_home(state: dict, model: str) -> Path: + """Write an isolated CODEX_HOME whose config.toml carries the ucode gateway + provider block, for the ucode-run `codex app-server`. + + The app-server rejects the global `--profile`, so a default-config CODEX_HOME + is how it inherits ucode's gateway (base_url + `ucode auth-token` refresh). + Reuses `render_overlay` — the same provider block `ucode configure codex` writes.""" + home = SMART_ROUTING_V2_HOME + home.mkdir(parents=True, exist_ok=True) + config_path = home / "config.toml" + overlay = render_overlay( + state["workspace"], + model, + state.get("profile"), + use_pat=bool(state.get("use_pat")), + ) + doc = read_toml_safe(config_path) + deep_merge_dict(doc, overlay) + write_toml_file(config_path, doc) + return home + + +def _launch_smart_routing_v2(state: dict, tool_args: list[str]) -> None: + """Experimental single-command launch of the real Codex TUI with runtime model switching. + + ucode owns three processes: a `codex app-server` subprocess, the WebSocket interposer + (daemon thread), and the `codex --remote` TUI (foreground). The interposer holds the first + turn on the normal model, then rewrites subsequent turns to SMART_ROUTING_V2_TARGET_MODEL and + injects a settings update so the TUI reflects the switch. The app-server + interposer are torn + down when the TUI exits. Mirrors the lifecycle of `claude.py::_launch_relayed`. + """ + from ucode.smart_routing import codex_interposer + + binary = SPEC["binary"] + workspace = state.get("workspace") + if not workspace: + raise RuntimeError( + "Smart routing v2 needs a configured workspace; run `ucode configure codex` first." + ) + start_model = default_model(state) + if not start_model: + raise RuntimeError( + "Smart routing v2 could not determine a starting Codex model for this workspace." + ) + + os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile")) + home = _generate_v2_app_server_home(state, start_model) + app_port = codex_interposer.free_port() + tui_port = codex_interposer.free_port() + + print_note( + f"Smart routing v2: starting on {start_model}, switching to " + f"{SMART_ROUTING_V2_TARGET_MODEL} after the first prompt " + f"(interposer log: {SMART_ROUTING_V2_LOG})." + ) + + app_server = subprocess.Popen( + [binary, "app-server", "--listen", f"ws://127.0.0.1:{app_port}"], + env={**os.environ, "CODEX_HOME": str(home)}, + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + stop_interposer = None + try: + if not codex_interposer.wait_healthz(app_port, timeout=30): + raise RuntimeError( + "Codex app-server did not become ready for smart routing v2; check workspace auth." + ) + _thread, stop_interposer = codex_interposer.start_interposer_thread( + "127.0.0.1", + tui_port, + f"ws://127.0.0.1:{app_port}", + SMART_ROUTING_V2_TARGET_MODEL, + SMART_ROUTING_V2_AFTER, + log_path=SMART_ROUTING_V2_LOG, + ) + # Foreground TUI. Popen (not exec) so this process stays alive to tear down the + # app-server + interposer when the TUI exits (see claude.py::_launch_relayed). + tui = subprocess.Popen( + [binary, "--remote", f"ws://127.0.0.1:{tui_port}", "--model", start_model, *tool_args] + ) + try: + returncode = tui.wait() + except KeyboardInterrupt: + tui.send_signal(signal.SIGINT) + returncode = tui.wait() + finally: + if stop_interposer is not None: + stop_interposer() + app_server.terminate() + try: + app_server.wait(timeout=5) + except Exception: # noqa: BLE001 - the app-server must never linger + app_server.kill() + sys.exit(returncode) + + def launch(state: dict, tool_args: list[str]) -> None: binary = SPEC["binary"] workspace = state.get("workspace") + if smart_routing_v2_enabled(): + _launch_smart_routing_v2(state, tool_args) + return if workspace: os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile")) # Run codex with --profile first — the TUI and runtime subcommands diff --git a/src/ucode/smart_routing/codex_interposer.py b/src/ucode/smart_routing/codex_interposer.py new file mode 100644 index 00000000..a729635c --- /dev/null +++ b/src/ucode/smart_routing/codex_interposer.py @@ -0,0 +1,234 @@ +"""WebSocket interposer for the Codex TUI's ``--remote`` transport (smart routing v2). + +Codex's remote transport (``codex --remote ws://…``) is WebSocket: a plain-JSONL +client is rejected with HTTP 400 ("Connection header did not include 'upgrade'"), +a proper upgrade returns 101, and each JSON-RPC message is one WebSocket text +frame. This module sits between the real TUI and a real ``codex app-server``, +forwarding every frame untouched except: + + - ``turn/start`` (TUI->engine): after an initial hold of ``after`` turns, its + ``model`` is rewritten. ``turn/start.model`` is documented as "override the + model for this turn and subsequent turns", so the live session retargets with + history preserved. + - When the hold expires (right after the Nth prompt completes) an injected + ``thread/settings/updated`` notification (engine->TUI) carries the new model, + so the TUI's on-screen model indicator follows the switch. + +``ucode.agents.codex`` runs :func:`start_interposer_thread` in a daemon thread +while it owns the app-server subprocess and the ``codex --remote`` TUI, so the +whole thing launches from the single ``ucode codex`` command. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import json +import socket +import threading +import time +import urllib.request +from collections.abc import Callable +from pathlib import Path + +from websockets.asyncio.client import connect +from websockets.asyncio.server import serve + +SETTINGS_UPDATED = "thread/settings/updated" + + +class _Session: + """Per-TUI-connection state: hold the first ``after`` turns, then switch model.""" + + def __init__(self, target_model: str, after: int, log: Callable[[str], None]) -> None: + self.target = target_model + self.after = after + self.log = log + self.turns = 0 + self.thread_id: str | None = None + self.settings: dict | None = None + self.injected = False + + def on_tui_frame(self, raw: str) -> str: + """TUI->engine: rewrite ``turn/start.model`` once past the hold.""" + try: + msg = json.loads(raw) + except ValueError: + return raw + if not isinstance(msg, dict): + return raw + params = msg.get("params") + if msg.get("method") == "turn/start" and isinstance(params, dict): + self.turns += 1 + if isinstance(params.get("threadId"), str): + self.thread_id = params["threadId"] + if self.turns > self.after: + old = params.get("model") + if old != self.target: + params["model"] = self.target + self.log(f"[REWRITE] turn #{self.turns}: model {old!r} -> {self.target!r}") + return json.dumps(msg) + return raw + + def on_engine_frame(self, raw: str) -> dict | None: + """engine->TUI: capture thread id/settings; after the hold's last turn + completes, return an injected settings-updated notification (or None).""" + try: + msg = json.loads(raw) + except ValueError: + return None + if not isinstance(msg, dict): + return None + params = msg.get("params") if isinstance(msg.get("params"), dict) else {} + result = msg.get("result") if isinstance(msg.get("result"), dict) else {} + for src in (params, result): + tid = src.get("threadId") or (src.get("thread") or {}).get("id") + if isinstance(tid, str): + self.thread_id = tid + ts = src.get("threadSettings") + if isinstance(ts, dict): + self.settings = ts + if ( + msg.get("method") == "turn/completed" + and not self.injected + and self.turns >= self.after + and self.thread_id + ): + self.injected = True + settings = dict(self.settings) if isinstance(self.settings, dict) else {} + settings["model"] = self.target + self.log(f"[INJECT] {SETTINGS_UPDATED}: model -> {self.target!r} (flip TUI chip)") + return { + "method": SETTINGS_UPDATED, + "params": {"threadId": self.thread_id, "threadSettings": settings}, + } + return None + + +async def _handle_tui(tui, upstream_uri: str, target_model: str, after: int, log) -> None: + path = getattr(getattr(tui, "request", None), "path", "/") or "/" + uri = upstream_uri.rstrip("/") + path + log(f"[CONN] TUI connected (path={path}); dialing app-server {uri}") + sess = _Session(target_model, after, log) + async with connect(uri, max_size=None) as upstream: + + async def tui_to_app(): + async for frame in tui: + if isinstance(frame, str): + frame = sess.on_tui_frame(frame) + await upstream.send(frame) + + async def app_to_tui(): + async for frame in upstream: + await tui.send(frame) + if isinstance(frame, str): + inj = sess.on_engine_frame(frame) + if inj is not None: + await tui.send(json.dumps(inj)) + + a = asyncio.create_task(tui_to_app()) + b = asyncio.create_task(app_to_tui()) + _done, pending = await asyncio.wait({a, b}, return_when=asyncio.FIRST_COMPLETED) + for t in pending: + t.cancel() + with contextlib.suppress(asyncio.CancelledError): + await t + log("[CONN] TUI session closed") + + +async def _serve(host: str, port: int, upstream_uri: str, model: str, after: int, log): + async def handler(tui): + try: + await _handle_tui(tui, upstream_uri, model, after, log) + except Exception as exc: # noqa: BLE001 - one session must never kill the server + log(f"[ERR] session: {exc!r}") + + server = await serve(handler, host, port, max_size=None) + log(f"[READY] ws://{host}:{port} -> {upstream_uri} (hold {after} turn(s), then -> {model!r})") + return server + + +def start_interposer_thread( + host: str, + port: int, + upstream_uri: str, + model: str, + after: int, + *, + log_path: Path | None = None, + ready_timeout: float = 10.0, +) -> tuple[threading.Thread, Callable[[], None]]: + """Run the interposer's asyncio server in a daemon thread. + + Returns ``(thread, stop)``; ``stop()`` shuts the server down and stops the + loop. Logs go to ``log_path`` (appended) when given — never to stdout/stderr, + which the foreground TUI owns. Blocks until the server is listening (or + ``ready_timeout`` elapses).""" + + def log(message: str) -> None: + if log_path is None: + return + line = f"{time.strftime('%H:%M:%S')} {message}\n" + try: + with open(log_path, "a", encoding="utf-8") as handle: + handle.write(line) + except OSError: + pass + + loop = asyncio.new_event_loop() + holder: dict = {} + ready = threading.Event() + + def run() -> None: + asyncio.set_event_loop(loop) + try: + holder["server"] = loop.run_until_complete( + _serve(host, port, upstream_uri, model, after, log) + ) + except Exception as exc: # noqa: BLE001 - surface bind/connect failures to the log + log(f"[ERR] failed to start interposer: {exc!r}") + ready.set() + loop.close() + return + ready.set() + loop.run_forever() + # Stopped: close the server and drain. + server = holder.get("server") + if server is not None: + server.close() + with contextlib.suppress(Exception): + loop.run_until_complete(server.wait_closed()) + loop.close() + + thread = threading.Thread(target=run, name="codex-interposer", daemon=True) + thread.start() + ready.wait(timeout=ready_timeout) + + def stop() -> None: + with contextlib.suppress(RuntimeError): + loop.call_soon_threadsafe(loop.stop) + + return thread, stop + + +def free_port() -> int: + """Grab an unused loopback TCP port (races are irrelevant for local ephemeral use).""" + sock = socket.socket() + sock.bind(("127.0.0.1", 0)) + port = sock.getsockname()[1] + sock.close() + return port + + +def wait_healthz(port: int, timeout: float = 30.0) -> bool: + """Poll the app-server's ``/healthz`` until it returns 200, or timeout.""" + url = f"http://127.0.0.1:{port}/healthz" + end = time.time() + timeout + while time.time() < end: + try: + with urllib.request.urlopen(url, timeout=1) as resp: # noqa: S310 - fixed localhost URL + if resp.status == 200: + return True + except Exception: # noqa: BLE001 - not ready yet; keep polling + time.sleep(0.25) + return False diff --git a/tests/test_codex_smart_routing_v2.py b/tests/test_codex_smart_routing_v2.py new file mode 100644 index 00000000..1ebffd88 --- /dev/null +++ b/tests/test_codex_smart_routing_v2.py @@ -0,0 +1,127 @@ +"""Tests for the experimental ENABLE_SMART_ROUTING_V2 Codex launch path.""" + +from __future__ import annotations + +import json + +from ucode.agents import codex +from ucode.config_io import read_toml_safe +from ucode.smart_routing import codex_interposer + +WS = "https://example.databricks.com" + + +class TestV2FlagGating: + def test_disabled_by_default(self, monkeypatch): + monkeypatch.delenv("ENABLE_SMART_ROUTING_V2", raising=False) + assert codex.smart_routing_v2_enabled() is False + + def test_enabled_when_1(self, monkeypatch): + monkeypatch.setenv("ENABLE_SMART_ROUTING_V2", "1") + assert codex.smart_routing_v2_enabled() is True + + def test_other_values_do_not_enable(self, monkeypatch): + monkeypatch.setenv("ENABLE_SMART_ROUTING_V2", "true") + assert codex.smart_routing_v2_enabled() is False + + def test_launch_dispatches_to_v2(self, monkeypatch): + monkeypatch.setenv("ENABLE_SMART_ROUTING_V2", "1") + called = {} + monkeypatch.setattr( + codex, + "_launch_smart_routing_v2", + lambda state, args: called.setdefault("hit", (state, args)), + ) + + # Should return via the v2 branch before touching normal launch/auth. + def _fail_if_normal_path(*_a, **_k): # pragma: no cover - only if v2 branch is skipped + raise AssertionError("normal launch path ran despite ENABLE_SMART_ROUTING_V2=1") + + monkeypatch.setattr(codex, "get_databricks_token", _fail_if_normal_path) + codex.launch({"workspace": WS}, ["--foo"]) + assert called["hit"] == ({"workspace": WS}, ["--foo"]) + + +class TestGenerateV2Home: + def test_writes_provider_config(self, tmp_path, monkeypatch): + monkeypatch.setattr(codex, "SMART_ROUTING_V2_HOME", tmp_path / "v2home") + monkeypatch.setattr(codex, "ucode_version", lambda: "0.1.0") + monkeypatch.setattr(codex, "agent_version", lambda binary: "0.148.0") + + home = codex._generate_v2_app_server_home( + {"workspace": WS, "profile": "myprof"}, "gpt-5.6-luna" + ) + + assert home == tmp_path / "v2home" + doc = read_toml_safe(home / "config.toml") + assert doc["model_provider"] == codex.CODEX_MODEL_PROVIDER_NAME + assert doc["model"] == "gpt-5.6-luna" + provider = doc["model_providers"][codex.CODEX_MODEL_PROVIDER_NAME] + assert provider["base_url"].endswith("/ai-gateway/codex/v1") + # Self-refreshing auth command is preserved (app-server rejects --profile). + assert provider["auth"]["command"].endswith("ucode") + assert "myprof" in provider["auth"]["args"] + + +class TestInterposerSession: + def _turn_start(self, model: str, thread_id: str = "t1") -> str: + return json.dumps( + { + "method": "turn/start", + "id": 1, + "params": {"threadId": thread_id, "input": [], "model": model}, + } + ) + + def test_holds_first_turn_then_switches(self): + sess = codex_interposer._Session("gpt-5.5", after=1, log=lambda _m: None) + # Turn 1 passes through unchanged (still on the TUI's model). + out1 = sess.on_tui_frame(self._turn_start("system.ai.gpt-5-6-luna")) + assert json.loads(out1)["params"]["model"] == "system.ai.gpt-5-6-luna" + # Turn 2 is rewritten to the target. + out2 = sess.on_tui_frame(self._turn_start("system.ai.gpt-5-6-luna")) + assert json.loads(out2)["params"]["model"] == "gpt-5.5" + + def test_after_zero_switches_immediately(self): + sess = codex_interposer._Session("gpt-5.5", after=0, log=lambda _m: None) + out1 = sess.on_tui_frame(self._turn_start("luna")) + assert json.loads(out1)["params"]["model"] == "gpt-5.5" + + def test_non_turn_frames_pass_through(self): + sess = codex_interposer._Session("gpt-5.5", after=1, log=lambda _m: None) + frame = json.dumps({"method": "initialize", "id": 1, "params": {}}) + assert sess.on_tui_frame(frame) == frame + + def test_injects_settings_update_after_hold(self): + sess = codex_interposer._Session("gpt-5.5", after=1, log=lambda _m: None) + sess.on_tui_frame(self._turn_start("luna")) # turn 1 (the hold) + inj = sess.on_engine_frame( + json.dumps( + { + "method": "turn/completed", + "params": {"threadId": "t1", "turn": {"status": "completed"}}, + } + ) + ) + assert inj is not None + assert inj["method"] == codex_interposer.SETTINGS_UPDATED + assert inj["params"]["threadId"] == "t1" + assert inj["params"]["threadSettings"]["model"] == "gpt-5.5" + + def test_injects_only_once(self): + sess = codex_interposer._Session("gpt-5.5", after=1, log=lambda _m: None) + sess.on_tui_frame(self._turn_start("luna")) + done = json.dumps({"method": "turn/completed", "params": {"threadId": "t1", "turn": {}}}) + assert sess.on_engine_frame(done) is not None + assert sess.on_engine_frame(done) is None # second completion: no re-inject + + +class TestInterposerHelpers: + def test_free_port_returns_usable_port(self): + port = codex_interposer.free_port() + assert isinstance(port, int) + assert 1024 <= port <= 65535 + + def test_wait_healthz_false_on_dead_port(self): + dead = codex_interposer.free_port() + assert codex_interposer.wait_healthz(dead, timeout=1.0) is False diff --git a/uv.lock b/uv.lock index 0b7d205b..bb61ef81 100644 --- a/uv.lock +++ b/uv.lock @@ -3163,6 +3163,7 @@ dependencies = [ { name = "questionary" }, { name = "tomlkit" }, { name = "typer" }, + { name = "websockets" }, ] [package.optional-dependencies] @@ -3186,6 +3187,7 @@ requires-dist = [ { name = "questionary", specifier = ">=2.0.0" }, { name = "tomlkit", specifier = ">=0.13.0" }, { name = "typer", specifier = ">=0.12.0" }, + { name = "websockets", specifier = ">=13" }, ] provides-extras = ["tracing"] @@ -3236,6 +3238,95 @@ wheels = [ { url = "https://pypi-proxy.cloud.databricks.com/packages/68/5a/199c59e0a824a3db2b89c5d2dade7ab5f9624dbf6448dc291b46d5ec94d3/wcwidth-0.6.0-py3-none-any.whl", hash = "sha256:1a3a1e510b553315f8e146c54764f4fb6264ffad731b3d78088cdb1478ffbdad", size = 94189, upload-time = "2026-02-06T19:19:39.646Z" }, ] +[[package]] +name = "websockets" +version = "17.0.1" +source = { registry = "https://pypi-proxy.cloud.databricks.com/simple" } +sdist = { url = "https://pypi-proxy.cloud.databricks.com/packages/f7/96/e01084f83a64bcb3a27994bd0cb0db68ff29d9c6707fae37ec19b18ba990/websockets-17.0.1.tar.gz", hash = "sha256:5baa9bc0dfbae8c507e51c8cf1b6d4628086f7a87bbd3a9952bd5f035451f1cc", size = 183298, upload-time = "2026-07-31T11:31:27.665Z" } +wheels = [ + { url = "https://pypi-proxy.cloud.databricks.com/packages/50/ff/6199a52d864215750af8668d84b0274775011a90052081f5a9495807a92b/websockets-17.0.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:10f461191125c63902ea7394ae9e752b1b5785641850c1d365bb30b0f88bc53f", size = 212603, upload-time = "2026-07-31T11:29:30.771Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/54/7e/439a962bcada88dcf586da77a1b2385f91e2d2910e9359540934c827156b/websockets-17.0.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:cffc84ddec6da7f447677266fee2a3c40ecc78172f00752aa1150b8a8d65df1d", size = 210286, upload-time = "2026-07-31T11:29:32.157Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/d9/82/123660edc759c225626b3b91952c7625f85c77a8362acbc35a4623120f7d/websockets-17.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:c23e532c8a2325a1e7486de8763a60dc43e83f01bcaeca07e3ba79652c156db1", size = 210549, upload-time = "2026-07-31T11:29:33.388Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/cd/2f/2940e57080cf56f28190287516400126d5a76b52b9a61dc10ba6f6400dbe/websockets-17.0.1-cp312-cp312-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:c09e097d0e46e3c289bedab9a475ae344b70c30ff5646e46af22b4e6fdc97b21", size = 219874, upload-time = "2026-07-31T11:29:34.608Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/42/28/9ec976c16d63cc51c28dfec74b66854048c0b8b6579946e902f91b69e8bf/websockets-17.0.1-cp312-cp312-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:f47b0815af3948ec6a440b3afa02f05b18cc0939549e91b5c677b5d9c2c8472a", size = 220150, upload-time = "2026-07-31T11:29:35.831Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/f4/a4/850c699a16bbc451723856360c59bd997bec075e637154f3fa96e80d5760/websockets-17.0.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8848c207049ad49d318e5f64a3d4d7bb189f8328d0d98e65647788f2a085785c", size = 221389, upload-time = "2026-07-31T11:29:37.189Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/47/e1/f60a891c1a4b3420d5052333a84eb7241e1fb4a71866dec1562f5fa30027/websockets-17.0.1-cp312-cp312-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:2604de7228506b13a44a256a9d223943340c0e725af5d367dc068e192b027761", size = 224169, upload-time = "2026-07-31T11:29:38.61Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/26/fb/e2a893be6fae4fddfe50ddc3035a331d3f381103d5467b7900026bdb3a64/websockets-17.0.1-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:07abc3bd196a48af476a82fd47f3f79a6a3f70937a9f930cef703cfa0c9d83b6", size = 222025, upload-time = "2026-07-31T11:29:39.897Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/d9/72/e3144b2d79276fab9798ed7d4aea2f0847434f186800b6f56a1eddcb3114/websockets-17.0.1-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:769ce7e2acfd9a89f2bed3a9c0da229459516bbc00bd4c9e2ca492c613ae4861", size = 220779, upload-time = "2026-07-31T11:29:41.084Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/c0/5e/69c02174fbcf1c40c6adc45d3c316a401558392fe7bab8969ef8c46f1689/websockets-17.0.1-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:07d78a509c3333f5908c83d7f78144ea68a6c9ec28110f5c54d81d8fcdc262c4", size = 218053, upload-time = "2026-07-31T11:29:42.322Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/b3/09/7574778b095b99cfa0856583462f56568df784f9b41485145169b2ec9c64/websockets-17.0.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:ffad64ce7ad3703d652a3fd9af26238377d24ce52c6ad8ff35d26d82f61f493f", size = 220825, upload-time = "2026-07-31T11:29:43.553Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/c5/e6/f46571f38765dbc4cbc0d0b47de8db65768006dbbd4340e6f5f51bc1d895/websockets-17.0.1-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:e95e321d0d763f2b6633512605f6112ebd70d5746f3ce05c941909d4a25233f2", size = 219427, upload-time = "2026-07-31T11:29:44.731Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/9d/31/6ff1fee057bd7e9dd5237fc064a749615378d003aa045b5bfc2d12b2f4f7/websockets-17.0.1-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:cd526c8228e759c1006c4b7c9ac71dc4e925ced1a6a6a5a8e94643709738f63e", size = 220198, upload-time = "2026-07-31T11:29:45.997Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/7a/4f/d41847227a44b9ad87c3d5a9fddbfad8b7c4d6032878d8460d9d37c2d44f/websockets-17.0.1-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:8cd3369e42c0246afaf9d669cfc19797e3a49e8c0a639544459c57597108b966", size = 221304, upload-time = "2026-07-31T11:29:47.314Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/63/30/21a7e326c6ad2eb526cd5b816383d59cdeb28b8805b65a543c3cfbd8e8ce/websockets-17.0.1-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:b580794e926cab7ff42ee4371ef14e0b22cb2bb722a607f77769136468f49a3f", size = 218858, upload-time = "2026-07-31T11:29:48.587Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/2d/48/55b0331cd5bec9ce29748f79edc00075805450a47011d0c8e3b1c61dbf04/websockets-17.0.1-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:5033ffe6804dd53afafa7d08e8c3eef2d2431f34d58ca30507a8442dd04a033a", size = 219840, upload-time = "2026-07-31T11:29:49.791Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/78/6e/2e8bc06e546f49b32a58a2bc2957902d1809ecc37552d3d7ccd6639a126e/websockets-17.0.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:6db9e5bf3649ab506c6ae8a3ac85a00fb1ae3816d75962771b2df8adbc5d40d2", size = 220116, upload-time = "2026-07-31T11:29:51.026Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/b0/ad/4bac01fa41aca54307157b9c9f68b066a6bb51fb18716ba618078a67b283/websockets-17.0.1-cp312-cp312-win32.whl", hash = "sha256:bc0bca48ba24c6c866847fd20478a51dd547fa0ad258dab9615c414ec534bbc0", size = 213050, upload-time = "2026-07-31T11:29:52.328Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/82/d8/c3a78cccc74a554780e9e76e323d5cde891048627025f0f82623e22dc3df/websockets-17.0.1-cp312-cp312-win_amd64.whl", hash = "sha256:2b3f3020171202b135ca078e20434977c6b2b02af647130d6980c9e39b9462e3", size = 213348, upload-time = "2026-07-31T11:29:53.891Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/7b/25/e1b8824bd632c8a5a62d504b61e9e35e470b67e4be0206f5c28f90c7f86d/websockets-17.0.1-cp312-cp312-win_arm64.whl", hash = "sha256:41d6aa06b5ab832aee72fedf47a149535b121ac900b6bb4d3fe14712afac9a79", size = 213276, upload-time = "2026-07-31T11:29:55.299Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/ba/a8/79c577bc2f874ee22f6f5ccdab97ba9ce6b96806be3fcc3a6d8490f88a21/websockets-17.0.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:55b12e47dcee83673a40d07686cfb6f9d6dfc285976ade9463f61d2bef3fad22", size = 212593, upload-time = "2026-07-31T11:29:56.518Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/db/99/e1cfaf419bb3b2fcfd6792a846f1d936293132b0b9a56530ced016c83c7b/websockets-17.0.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:c1c118a6b0e25bfc9a6802075d748fa6321714ffbdf3c88d29d9a0e3c7386c75", size = 210280, upload-time = "2026-07-31T11:29:57.768Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/a2/ef/cc994494bf7d97e41833f6ff55c24f535e4d527a10370b9631737e9c2f00/websockets-17.0.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:734d20364dc2cfe03674883cafcf580b6e431c5ce42b476312b9285310230cf9", size = 210538, upload-time = "2026-07-31T11:29:59.021Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/87/32/fbf2d132f63ba3e67f675bccf333469786a24e0418969ce1d8e6ff9e6f02/websockets-17.0.1-cp313-cp313-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:9493314a99e599163c854fb5900ad7f7ea38c5cb9d9103aa30b3c6b8181c01fa", size = 219925, upload-time = "2026-07-31T11:30:00.298Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/16/50/64eee3d25a47fe744a9490e0627cc373dca096755db740f91c28bd61cd35/websockets-17.0.1-cp313-cp313-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:18ded646ce98cdd3c0235825b3252f1df55765ba49b616bb10282f758667b4d0", size = 220206, upload-time = "2026-07-31T11:30:01.52Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/15/56/10ed4bc4dd75f204e3c62bd4898e44a8742a27773c80b188cfa7888aad2d/websockets-17.0.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c1bec5d6a19f5fbe87e4940739cfc65e7bb53d8b353e1029b8037a1653b321bc", size = 221445, upload-time = "2026-07-31T11:30:02.788Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/bf/be/bb14328614c068ab09569962fbf218fc00413ce3febc6d2684c764b6f37e/websockets-17.0.1-cp313-cp313-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:872273e629ca7e3d35f16a2dc6ede84e1d5c831e616b8277de6e4f83114e7c58", size = 222887, upload-time = "2026-07-31T11:30:03.943Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/32/1b/4cb0eec2fee310007104687493175af190019f705940c864f9c523fe9f6f/websockets-17.0.1-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:1df81d174c1561292de9e40b141cafc04f69077272f6c352afe1d743e20810df", size = 222072, upload-time = "2026-07-31T11:30:05.258Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/6c/9c/14e6391de777ddb39c439c450deb551406d445e25a5877d6fa25c49d4544/websockets-17.0.1-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:759adeb5b0c5775b563254ec63b5b79089fc0045b479143a0b1b8c0ebaae1253", size = 220826, upload-time = "2026-07-31T11:30:06.53Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/cb/57/96e94e384442247bbed5d3ab67381c7257355c2d66b62c3ad33a17f5d385/websockets-17.0.1-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:4d1d99db29b5444e3982f1ce2ba8a833508ad44b2f1fbd0bd99e81d825c0b461", size = 218107, upload-time = "2026-07-31T11:30:07.766Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/c0/8c/9c9dedd14c3919435df9b35cdee7111268c751252b87652f3a6a4f56e760/websockets-17.0.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:02ed63bf26dda9fa27df730a41f6664586c4ee05972c8fb667ce1725b3fd13d3", size = 220889, upload-time = "2026-07-31T11:30:09.035Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/94/4d/ca73c2ac82c00f50c529784bacb323e42da4816333211bc1543d90c9cf11/websockets-17.0.1-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:eab6de8a98b9a7772cf686d00b4de439fc7efb8ab05ae106ef227291d06f87c5", size = 219486, upload-time = "2026-07-31T11:30:10.289Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/5b/da/fb37ac09dcd7c69dd73bac979ed393df35f78a3c232e293d1ff3bd586d24/websockets-17.0.1-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:2a855b6dfe21c4d3420be265ae031829ba8ba0be0ea350d9f7c3ef30ae63ebe2", size = 220258, upload-time = "2026-07-31T11:30:11.605Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/fc/04/9693f191d968a93f37326a17301a101d49580889c688f466699f89ecdee1/websockets-17.0.1-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:7002d5f9e1c3ddd991cdfdbfee18cc8c8b196b2445022892badacd6cb338bbbc", size = 221358, upload-time = "2026-07-31T11:30:12.858Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/cf/29/ad0d85c01db5dcf22898d51648bd2c25af0dd0a4a41c550b11acddeeeba7/websockets-17.0.1-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:c395bda8e7d8f51a02e80261fb57127979e5c472675d9a96b2860619ad47da48", size = 218921, upload-time = "2026-07-31T11:30:14.064Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/18/3b/bf8e855e495dcca63f2b8aa019cf2ada3160e1fa66d833c7417f3b1f7f38/websockets-17.0.1-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:aadc298969ad229d8e3029fc5cc751fdad286696230f9cf014e90ff9cd8e6ea0", size = 219871, upload-time = "2026-07-31T11:30:15.358Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/3b/db/c7abd6639a93a40279cd1ddc57e09e1c4f8381c4cfccdb775aa5aac9770a/websockets-17.0.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:f11a398d8170b7ac5000baf7f258dcda579ef3ea744e0cc6a165e0dfbc0d3198", size = 220154, upload-time = "2026-07-31T11:30:16.96Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/f6/2a/25a9f8f2e5a6ef34e911d2f55d9f756bdeb92b4c28cfb77b8430bbc73cb1/websockets-17.0.1-cp313-cp313-win32.whl", hash = "sha256:846a4a8b0833e3cad57523d9e3bd50ec8ea05ab9d06c582f82a1340ba096af5f", size = 213038, upload-time = "2026-07-31T11:30:18.434Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/81/2f/ea1380f72bb11b64fc5bc7ae0d42de5bbf3e6dc13b965706b2a1d4e17cdf/websockets-17.0.1-cp313-cp313-win_amd64.whl", hash = "sha256:409d93efcaa14f7a99592c5baaef5ec6ca94fba0f5aec1a86f693977c69c9c1c", size = 213348, upload-time = "2026-07-31T11:30:19.693Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/e6/c7/b956ed9151c3c74530ebc62d716fbfdbde7507a6acc6423a64f9ecfb6b8a/websockets-17.0.1-cp313-cp313-win_arm64.whl", hash = "sha256:90246fa9e6cb192a778ce6ce024057ec54317a894db7899c922dcdc1f4cbf6a5", size = 213282, upload-time = "2026-07-31T11:30:21.045Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/98/dc/cadab608924ac605647031472fb1f8792d7d4ea07565ba1899ec42028e0d/websockets-17.0.1-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:53b90c00bc6201ab6695c7ff51a04d0e425514c37515e9eeecd2c1b978ac6c0e", size = 212640, upload-time = "2026-07-31T11:30:22.436Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/16/7a/b034d13ca181211bbd58bb50835cb196a7784cd505b5a2079d4d03374f9f/websockets-17.0.1-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:5f33a649bfcb8312524173cc4bbafa7dbb236e18eee9aa31a1d324ca0ddda28c", size = 210332, upload-time = "2026-07-31T11:30:23.608Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/2f/4d/943ede39b53744768edf1ed84a3f9401527388228a3d6c1249c02c3d6bd7/websockets-17.0.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:cddc675ec31bca65473321f9a9794e488b43b3b8de5d02c8ef4810c5d5792163", size = 210546, upload-time = "2026-07-31T11:30:24.932Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/cf/e3/88dc159d2ae66743c669443246243f28d873b0c5e58271b8cc1ca0440334/websockets-17.0.1-cp314-cp314-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:b3ff0ad440ad52dda64138f16895f66403f40192365e39b1010e889f289746b0", size = 219928, upload-time = "2026-07-31T11:30:26.221Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/fe/f2/ff27eaefa15851a5cf7f004ab827a022bf2d6632cb520f89cb100db7e84b/websockets-17.0.1-cp314-cp314-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:72d7f2a5aeb4e82daa4ee18f125b4277f427033359be5c745ad709608446cc2c", size = 220279, upload-time = "2026-07-31T11:30:27.49Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/19/2e/a5166149f363d2449c1cb2dde6486a245521979509d53b87a09f3e79662b/websockets-17.0.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6fd88365da261c53d3e943fb37e0d0721b9cde119f6b2e3fc84369b6ab234d63", size = 221525, upload-time = "2026-07-31T11:30:28.872Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/9b/b5/f46931269b3ff3bde65d27c65ddb22f9bb8ce92ac2c6c4df0910128f6219/websockets-17.0.1-cp314-cp314-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:ab9f962a5b64a5c3c845d556b7dc4e6fb683f7b67179f8205e814bb2e0213ffe", size = 222897, upload-time = "2026-07-31T11:30:30.164Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/42/f4/deccf3439f35df953ec35e13fe07986821c5f1ab5785d69614283bdb9034/websockets-17.0.1-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:8c07f145d0b9e90cbd96035f31fb79199aef4da1872854e36ebeb258e3d57594", size = 222129, upload-time = "2026-07-31T11:30:31.489Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/35/a5/e1b57a59da92ade37fd021567a17b518ea8267b28e5530075844cdb525fe/websockets-17.0.1-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:9f7747d3daa41a11f25f7cca5dc988fc51da97b311bed4c9d843860f79779283", size = 220875, upload-time = "2026-07-31T11:30:32.805Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/f0/30/e7d0889c790a854156de424575fd67af79ddbaed9ff3157ae863dfd1c1dc/websockets-17.0.1-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:2abb1ba0a5133b7d2ef3c1c9f4b0c1e8a101012dce0b594ab2b2888d9a64820e", size = 218160, upload-time = "2026-07-31T11:30:34.512Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/09/2e/43db785d6ed9ae7594fae7b62bbc9cb4dfee2b015e06a1005f1e5ce283b6/websockets-17.0.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:f3fd9a1f87f8f0f3f8e9f9bd0195f7516562d13f5b178db8c5784d1f60b60bed", size = 220951, upload-time = "2026-07-31T11:30:35.806Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/eb/f5/4ac3cab3d5e8a830657a822f64a8910e3803229c6783e59c3fd9a3487427/websockets-17.0.1-cp314-cp314-musllinux_1_2_armv7l.whl", hash = "sha256:2bc14b481e05e331811108daa1aeb41a5e237a5564ef2f02ec5a356a0f102f78", size = 219460, upload-time = "2026-07-31T11:30:37.273Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/03/0e/c3a4020673ffc17c82cf1a467835038a196a555d3b4f2a50f0f063cf8ccc/websockets-17.0.1-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:57d2ee9b24b404ce75f3814f92073c0ed88106c950148d2427fe8d25ca254d1f", size = 220248, upload-time = "2026-07-31T11:30:38.527Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/7d/87/e47a6a278cc1dfade38444c893ce18322943c25d4b780a74450d9d164be1/websockets-17.0.1-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:1b363bfd72a52c0658a3154a4cff219f15a474b35a235057d38853bf151acce7", size = 221421, upload-time = "2026-07-31T11:30:39.879Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/da/8f/473d5fc4e3836e375b0233c6ef26777e6e5e3f7bfc84ccd524eae4090ed5/websockets-17.0.1-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:10b1587c599fa0f2c89154587c80e0fda98ade6c9fa8c0260a2823fb1800b685", size = 218975, upload-time = "2026-07-31T11:30:41.192Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/d4/b9/819ec2dcdf69031d7e9cab11247f3a6ff9bbc8c7c53ada1dbcb9055b227b/websockets-17.0.1-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:d7d72843691f50b91127c50688df10cb72ec6f4c4b1d7e2c11ab33b16acf8e51", size = 219925, upload-time = "2026-07-31T11:30:42.524Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/85/b9/6c0da301f6118502e079cf92f4e864adf28e56b3f8c0f6085076ced7b876/websockets-17.0.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:90973a3a00f23afdfd1c9b06fb84289bf0220f247ef8a62501a1967c7af54f7b", size = 220218, upload-time = "2026-07-31T11:30:44.04Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/55/f9/cba32dc9dd856565263d6272f594255bd0e2781deb8cd982c026a54760ad/websockets-17.0.1-cp314-cp314-win32.whl", hash = "sha256:599b03beb77633bffc095334338fad79cafc2b01fbd58953838130a9ae967d7b", size = 212626, upload-time = "2026-07-31T11:30:45.579Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/fa/95/91cdd8c192287d7ea741f37cf7d64fdc1a14410f06f73805e428a1a590af/websockets-17.0.1-cp314-cp314-win_amd64.whl", hash = "sha256:81ce19c6046ace11da7001781be7317bb1dc389f399af4b2ed962190f76f9add", size = 212969, upload-time = "2026-07-31T11:30:46.983Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/8e/fd/8c98a1e431960661c5769ab1a4dd66494e87ab02d791cc79e51e0d9a289f/websockets-17.0.1-cp314-cp314-win_arm64.whl", hash = "sha256:efe0ae052a8d023b87198921e8a7ce1dc7768816bcd2fbc20df171ac73a04891", size = 212850, upload-time = "2026-07-31T11:30:48.304Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/13/c1/142f5186ee7dc3beee0426b998a79e223e067b7689afcaa95890d64aa800/websockets-17.0.1-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:ab56439c9f74c52770690c7b2f616b3bf775cb3920453ee355ac765c032d8bbf", size = 212967, upload-time = "2026-07-31T11:30:49.688Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/02/de/4b03ed316c9dee180365286c298219809ff247be39beeaaf9958b21167ab/websockets-17.0.1-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:20a92f78ac8250984ed459faa9ca48c285adbfc0038ddc3fdac6046990a9c9ed", size = 210504, upload-time = "2026-07-31T11:30:50.987Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/cd/d8/ad2b3e8f867e1e8cac3077e2f33ffb60b71bd763d6cfc71bd916f113c3bf/websockets-17.0.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:6a434e59962a4fb9016bea327e1d14d6cd67670ecfb8942b4f4a0c24036634ce", size = 210702, upload-time = "2026-07-31T11:30:52.261Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/48/18/7a77a82ce9d6f831c07b176da3942f7e71acd0f115f3ecdb1d00a040eb01/websockets-17.0.1-cp314-cp314t-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:2503c7e2a5049a12d5dac917a46d5d52591283a766165b8176bb167560421b38", size = 220290, upload-time = "2026-07-31T11:30:53.582Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/f5/af/43c3e3c3ea7ba4693c2181743f3221957df28bededadcbd9fc8a0661bde0/websockets-17.0.1-cp314-cp314t-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:28012a54510fe8301bb893ef143cec30a2780a2d3bc20b7bbdf4379d7a63945d", size = 220573, upload-time = "2026-07-31T11:30:54.966Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/1f/bd/ed48eca15725743ee7e2dc172e15c61de29e85ec98dace7b14257f366836/websockets-17.0.1-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:22bd00f8bae2bccdb5dbe41e20f58ba44ca9fff0b4b561aaf39099c35da762ed", size = 221747, upload-time = "2026-07-31T11:30:56.762Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/99/50/838deb7937a8225c4925dd4a977eafea473fabf444178a99de0bc7e92bb0/websockets-17.0.1-cp314-cp314t-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:e98ec9ec61cce5bc4b8b218322ad090b0994eb060bb04da704c62ef0a3d864e6", size = 223891, upload-time = "2026-07-31T11:30:58.127Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/f8/e9/657fb70c6eb6bcd01adfa5d2b06496e9911e1c1a8813d353b8c00f7591cd/websockets-17.0.1-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:8e387adb0c692c6b5571bdeafc8ac9d1901ea30f10309134780b16ecd35e6605", size = 222317, upload-time = "2026-07-31T11:30:59.416Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/07/4c/82cb722afa5428fed981331210c4c07600570db01bb1620579f655b5adaf/websockets-17.0.1-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:bd1470d2c53fe53269bf5619da7725d30dd9b9693f1689f7a85eab8dea734442", size = 221047, upload-time = "2026-07-31T11:31:00.757Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/90/84/bd6d67d6bc65f0de0cb50de55dab42f256a9876a351c4736522eb168fda0/websockets-17.0.1-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:884af729b8ab50486acd94d9768c2b60914bf39b579ebba0a5cb73bfdfd61fd2", size = 218626, upload-time = "2026-07-31T11:31:02.48Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/96/cb/6a372c8553976f0d8f97f5115826ed47e34b3be6b8bf0d0249af249a7416/websockets-17.0.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:a60fa1a25cca1bcc2bf87b8d6be37a741f0a3239fb5e9cfb7a37173b68ffcf87", size = 221299, upload-time = "2026-07-31T11:31:03.795Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/bc/31/f966e8472337974f74d788b3ef6c6f3b8b9a5f201efd16a843c91d269fa5/websockets-17.0.1-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:e8208f2729cba030ff872a92064c97584eeb9502f53d32a05a0f05d5a17ca6c6", size = 219789, upload-time = "2026-07-31T11:31:05.08Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/57/34/404e83a6cc7b0efcac810b7041bffd72ff76900e6fd0aa45a26c92fb2ffe/websockets-17.0.1-cp314-cp314t-musllinux_1_2_i686.whl", hash = "sha256:54cdcaa56f5d3eafd57058f0fa4a3de93a310b43a3c4699f06efc4c0bd054a5a", size = 220678, upload-time = "2026-07-31T11:31:06.635Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/e5/70/8946188c2a68d67251859b589a3634918cf7867bf0b891347a5ecaa43d30/websockets-17.0.1-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:4d41c0a1d47a478bc432b3b9068097bee1ce0c5b19327ea6f75c2ab34ab1f2fb", size = 221697, upload-time = "2026-07-31T11:31:08.023Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/04/16/ee73fc2083a2938ac6209f4ec804960496835b20a0068dbcfe8424957c04/websockets-17.0.1-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:f991247276797d0c61ab7770bc9791eadc16f683b4d83517f624932adc1a8bab", size = 219390, upload-time = "2026-07-31T11:31:09.378Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/94/6a/d5f88033c69932af6cdaa72da62516ade47c257e3bf69f4c0ba5f40e12a2/websockets-17.0.1-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:733e3cc7171fa1b899edbe725ef9382d0e960657dc1fd933f3281ae910c01dab", size = 220161, upload-time = "2026-07-31T11:31:10.915Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/e1/2e/6183dd2c0370287ecf4afe0bb33aca364208e5e7b0e1a286adcaecc0c78b/websockets-17.0.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:810cb3fb5fa6e447216f4e82d9a85cb8aed0929ae3538153ddfe8a6e3121a58d", size = 220591, upload-time = "2026-07-31T11:31:12.289Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/26/93/70f6516d85b9744f7eac224c4b1b9ef4e84133f80b53be02080cb1c3e663/websockets-17.0.1-cp314-cp314t-win32.whl", hash = "sha256:17ac37716c0244e82c9e384c41653c090b1864c6610224ca3857e7f7b58fce10", size = 212755, upload-time = "2026-07-31T11:31:13.883Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/f1/2c/9d9c1da5a7ea9af307b386d25f64d1dead4729644198d2b92e36db5dfd41/websockets-17.0.1-cp314-cp314t-win_amd64.whl", hash = "sha256:bb31f42ea095ea826463c770829aa188a86c9a5c976b1467cbbf583c811de833", size = 213094, upload-time = "2026-07-31T11:31:15.367Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/5b/24/a585e7573e128070605d003b5544729bcd58d9756c7e99d550818ca4b916/websockets-17.0.1-cp314-cp314t-win_arm64.whl", hash = "sha256:dbfae8e75b342e31fc6fd1a8bbb393b7cbb91d6cfd581650300a94381e7b7e2b", size = 213009, upload-time = "2026-07-31T11:31:16.776Z" }, + { url = "https://pypi-proxy.cloud.databricks.com/packages/09/ce/3929538b2b9918f5eee623fbf3346893973191f6df93f19bbda097bd7bb7/websockets-17.0.1-py3-none-any.whl", hash = "sha256:c6be9cba65c65cc76dfa3d4619e359ff02a4476c74e179b215236c11a0b32345", size = 206718, upload-time = "2026-07-31T11:31:26.037Z" }, +] + [[package]] name = "werkzeug" version = "3.1.8" From e1492a4f0360de06b89e4dca4920f88a89cc905f Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Thu, 20 Aug 2026 23:27:25 +0000 Subject: [PATCH 02/18] rm --- scripts/codex_arch_b_probe.py | 222 ------------------- scripts/codex_model_router_poc.py | 355 ------------------------------ scripts/codex_tui_interposer.py | 340 ---------------------------- 3 files changed, 917 deletions(-) delete mode 100755 scripts/codex_arch_b_probe.py delete mode 100644 scripts/codex_model_router_poc.py delete mode 100644 scripts/codex_tui_interposer.py diff --git a/scripts/codex_arch_b_probe.py b/scripts/codex_arch_b_probe.py deleted file mode 100755 index aaa23201..00000000 --- a/scripts/codex_arch_b_probe.py +++ /dev/null @@ -1,222 +0,0 @@ -#!/usr/bin/env python3 -"""Arch B probe: MITM proxy between TUI and app-server, rewriting model in turn/start. - -This demonstrates feasibility of interposing on the real TUI without modifying it. -The proxy: -1. Listens on a unix socket that the TUI connects to (via --remote) -2. Forwards all messages to a real app-server -3. Rewrites turn/start.model to a fixed value (proving router capability) -4. Passes everything else through unchanged -""" -import json -import os -import socket -import subprocess -import sys -import threading -import time -import argparse - -class CodexInterposer: - """MITM proxy for Codex messages.""" - - def __init__(self, listen_sock_path, app_server_sock_path, target_model): - self.listen_sock_path = listen_sock_path - self.app_server_sock_path = app_server_sock_path - self.target_model = target_model - self.listener = None - self.running = True - - def cleanup(self): - """Clean up listener socket.""" - if os.path.exists(self.listen_sock_path): - try: - os.remove(self.listen_sock_path) - except: - pass - if self.listener: - try: - self.listener.close() - except: - pass - - def rewrite_message(self, msg): - """Rewrite turn/start to force target model.""" - if not isinstance(msg, dict): - return msg - - method = msg.get("method") - if method == "turn/start": - params = msg.get("params", {}) - if isinstance(params, dict): - old_model = params.get("model") - if old_model != self.target_model: - print(f"[REWRITE] turn/start: {old_model!r} -> {self.target_model!r}") - params["model"] = self.target_model - msg["params"] = params - - return msg - - def relay_messages(self, client_sock, as_sock): - """Relay messages bidirectionally, rewriting turn/start.""" - - def tui_to_as(): - """TUI -> app-server (with rewriting)""" - buffer = "" - while self.running: - try: - data = client_sock.recv(1024) - if not data: - print("[TUI->AS] Connection closed by TUI") - break - - buffer += data.decode('utf-8', errors='replace') - - # Process complete lines - while '\n' in buffer: - line, buffer = buffer.split('\n', 1) - line = line.strip() - if not line: - continue - - try: - msg = json.loads(line) - msg = self.rewrite_message(msg) - rewritten = json.dumps(msg) - as_sock.sendall((rewritten + '\n').encode('utf-8')) - print(f"[TUI->AS] {msg.get('method', msg.get('type', '?'))}") - except Exception as e: - print(f"[TUI->AS] Error: {e}") - as_sock.sendall((line + '\n').encode('utf-8')) - except Exception as e: - print(f"[TUI->AS] Exception: {e}") - break - - def as_to_tui(): - """app-server -> TUI (pass-through)""" - buffer = "" - while self.running: - try: - data = as_sock.recv(1024) - if not data: - print("[AS->TUI] Connection closed by app-server") - break - - buffer += data.decode('utf-8', errors='replace') - - # Process complete lines - while '\n' in buffer: - line, buffer = buffer.split('\n', 1) - line = line.strip() - if not line: - continue - - try: - msg = json.loads(line) - method = msg.get('method') - if method in ('turn/start', 'turn/completed', 'item/completed'): - print(f"[AS->TUI] {method}") - except: - pass - - client_sock.sendall((line + '\n').encode('utf-8')) - except Exception as e: - print(f"[AS->TUI] Exception: {e}") - break - - t1 = threading.Thread(target=tui_to_as, daemon=True) - t2 = threading.Thread(target=as_to_tui, daemon=True) - t1.start() - t2.start() - - # Wait for either thread to finish - t1.join(timeout=300) - t2.join(timeout=300) - - self.running = False - - def handle_client(self, client_sock, addr): - """Handle a single TUI connection.""" - print(f"[CLIENT] Connected from {addr}") - - try: - # Connect to real app-server - as_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) - as_sock.connect(self.app_server_sock_path) - print(f"[RELAY] Connected to app-server at {self.app_server_sock_path}") - - # Relay bidirectionally - self.relay_messages(client_sock, as_sock) - - as_sock.close() - except Exception as e: - print(f"[ERROR] Failed to relay: {e}") - finally: - client_sock.close() - print(f"[CLIENT] Disconnected") - - def run(self): - """Start the interposer listening for TUI connections.""" - self.cleanup() - - print(f"[STARTUP] Creating listener at {self.listen_sock_path}") - - self.listener = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) - self.listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - self.listener.bind(self.listen_sock_path) - self.listener.listen(1) - - print(f"[STARTUP] Listening for TUI connections") - print(f"[INFO] Target model: {self.target_model!r}") - print(f"[INFO] To connect TUI: codex --remote unix://{self.listen_sock_path}") - - try: - while self.running: - try: - self.listener.settimeout(1.0) - client_sock, addr = self.listener.accept() - - # Handle in a thread - t = threading.Thread( - target=self.handle_client, - args=(client_sock, addr), - daemon=True - ) - t.start() - except socket.timeout: - continue - except KeyboardInterrupt: - print("\n[SHUTDOWN] Interrupted") - break - finally: - self.cleanup() - - -def main(): - parser = argparse.ArgumentParser( - description="Codex model interposer: MITM proxy to rewrite turn/start.model" - ) - parser.add_argument( - "--listen", - default="/home/lilly.luo/.cache/codex-b/tui-remote.sock", - help="Socket for TUI to connect to" - ) - parser.add_argument( - "--app-server", - default="/home/lilly.luo/.cache/codex-b/as.sock", - help="Real app-server socket" - ) - parser.add_argument( - "--model", - default="gpt-5.5", - help="Model to force for all turns" - ) - - args = parser.parse_args() - - interposer = CodexInterposer(args.listen, args.app_server, args.model) - interposer.run() - - -if __name__ == "__main__": - sys.exit(main() or 0) diff --git a/scripts/codex_model_router_poc.py b/scripts/codex_model_router_poc.py deleted file mode 100644 index b0134d1f..00000000 --- a/scripts/codex_model_router_poc.py +++ /dev/null @@ -1,355 +0,0 @@ -#!/usr/bin/env python3 -"""POC: a minimal interactive Codex client that can switch models mid-session. - -Launches `codex app-server` under the hood, gives you a prompt, and lets you -change the model live with `/model ` — the switch happens by setting the -per-turn `model` field on `turn/start`, so history is preserved across it. - -This is arch A from the plan (a thin app-server client). It is NOT Codex's -polished TUI; it's the smallest thing that proves "launch, type, switch". - -Run it with the repo's Python 3.12 venv (system python3 here is 3.6): - - /home/lilly.luo/ucode/.venv/bin/python scripts/codex_model_router_poc.py - -For interactive mode, omit all flags. For a self-test that proves mid-session -model switching with context preservation: - - /home/lilly.luo/ucode/.venv/bin/python scripts/codex_model_router_poc.py --selftest - -Auth/gateway config is generated from your existing ~/.codex/ucode.config.toml -provider block into an isolated CODEX_HOME, so it uses the same Databricks -gateway + `ucode auth-token` refresh that `ucode codex` uses. - -In-session commands: - /model switch the model for subsequent turns (e.g. /model gpt-5.5) - /model show the current model - /quit exit -""" -from __future__ import annotations - -import json -import os -import queue -import subprocess -import sys -import threading -import time -import tomllib -from pathlib import Path - -import tomlkit - -UCODE_CODEX_CONFIG = Path.home() / ".codex" / "ucode.config.toml" -POC_HOME = Path.home() / ".cache" / "ucode-codex-router-poc" -DEFAULT_MODEL = "system.ai.gpt-5-6-luna" -EXAMPLE_MODELS = ["system.ai.gpt-5-6-luna", "gpt-5.5"] - - -def build_codex_home() -> Path: - """Generate an isolated CODEX_HOME whose config.toml carries ONLY the ucode - gateway provider block (model_provider + model + [model_providers.*]), copied - from ~/.codex/ucode.config.toml. Keeps the app-server pointed at the same - Databricks gateway + auth-token refresh, without the hooks/tui cruft.""" - if not UCODE_CODEX_CONFIG.exists(): - sys.exit( - f"Missing {UCODE_CODEX_CONFIG}. Run `ucode configure codex` (or `ucode codex`) first " - "so the Databricks provider block exists." - ) - src = tomllib.loads(UCODE_CODEX_CONFIG.read_text()) - minimal = tomlkit.document() - if "model_provider" in src: - minimal["model_provider"] = src["model_provider"] - minimal["model"] = src.get("model", DEFAULT_MODEL) - if "model_reasoning_effort" in src: - minimal["model_reasoning_effort"] = src["model_reasoning_effort"] - if "model_providers" in src: - minimal["model_providers"] = src["model_providers"] - POC_HOME.mkdir(parents=True, exist_ok=True) - (POC_HOME / "config.toml").write_text(tomlkit.dumps(minimal)) - return POC_HOME - - -class AppServer: - """Thin newline-delimited-JSON stdio client for `codex app-server`.""" - - def __init__(self, codex_home: Path) -> None: - env = dict(os.environ) - env["CODEX_HOME"] = str(codex_home) - self.proc = subprocess.Popen( - ["codex", "app-server"], - stdin=subprocess.PIPE, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - text=True, - bufsize=1, - env=env, - ) - self._q: queue.Queue = queue.Queue() - self._id = 0 - threading.Thread(target=self._read_stdout, daemon=True).start() - threading.Thread(target=self._drain_stderr, daemon=True).start() - - def _read_stdout(self) -> None: - for line in self.proc.stdout: # type: ignore[union-attr] - line = line.strip() - if line: - try: - self._q.put(json.loads(line)) - except ValueError: - pass - - def _drain_stderr(self) -> None: - # app-server logs benign catalog-refresh 404s here; keep them out of the UI - # but available if the user wants them (uncomment to debug). - for _line in self.proc.stderr: # type: ignore[union-attr] - pass - - def _send(self, method: str, params: dict | None = None, *, notify: bool = False): - msg: dict = {"method": method} - if not notify: - self._id += 1 - msg["id"] = self._id - if params is not None: - msg["params"] = params - self.proc.stdin.write(json.dumps(msg) + "\n") # type: ignore[union-attr] - self.proc.stdin.flush() # type: ignore[union-attr] - return msg.get("id") - - def _wait(self, pred, timeout: float): - end = time.time() + timeout - while time.time() < end: - try: - msg = self._q.get(timeout=min(1.0, max(0.05, end - time.time()))) - except queue.Empty: - continue - if pred(msg): - return msg - return None - - def request(self, method: str, params: dict | None = None, *, timeout: float = 60.0): - rid = self._send(method, params) - return self._wait(lambda m: m.get("id") == rid and ("result" in m or "error" in m), timeout) - - def initialize(self) -> None: - self.request( - "initialize", - {"clientInfo": {"name": "ucode-codex-router-poc", "version": "0.1"}, "capabilities": {}}, - timeout=30, - ) - self._send("initialized", {}, notify=True) - - def start_thread(self, model: str) -> str: - resp = self.request( - "thread/start", {"model": model, "cwd": os.getcwd(), "approvalPolicy": "never"}, timeout=60 - ) - result = (resp or {}).get("result", {}) - tid = result.get("thread", {}).get("id") or result.get("threadId") - if not tid: - sys.exit(f"thread/start failed: {json.dumps(resp)[:400]}") - return tid - - def run_turn(self, thread_id: str, text: str, model: str, *, timeout: float = 300.0) -> None: - """Send one user turn on `model`, streaming assistant text to stdout live.""" - rid = self._send( - "turn/start", - {"threadId": thread_id, "input": [{"type": "text", "text": text}], "model": model}, - ) - # Ack (status inProgress) — then stream until turn/completed. - self._wait(lambda m: m.get("id") == rid and ("result" in m or "error" in m), 30) - end = time.time() + timeout - printed_any = False - while time.time() < end: - try: - msg = self._q.get(timeout=min(1.0, max(0.05, end - time.time()))) - except queue.Empty: - continue - method = msg.get("method") - params = msg.get("params") or {} - if method == "item/agentMessage/delta": - delta = _find_str(params, ("delta", "text")) - if delta: - sys.stdout.write(delta) - sys.stdout.flush() - printed_any = True - elif method == "turn/completed": - turn = params.get("turn", {}) - if turn.get("status") == "failed": - err = turn.get("error", {}) - print(f"\n [turn failed: {err.get('message', err)}]") - elif not printed_any: - # No deltas seen (some models don't stream) — print final items. - print(_final_text(turn) or " [no text returned]") - print() - return - print("\n [timed out waiting for the turn to complete]") - - def close(self) -> None: - try: - self.proc.stdin.close() # type: ignore[union-attr] - except Exception: - pass - self.proc.terminate() - try: - self.proc.wait(timeout=5) - except Exception: - self.proc.kill() - - -def _find_str(obj, keys) -> str | None: - if isinstance(obj, dict): - for k, v in obj.items(): - if k in keys and isinstance(v, str): - return v - r = _find_str(v, keys) - if r: - return r - elif isinstance(obj, list): - for v in obj: - r = _find_str(v, keys) - if r: - return r - return None - - -def _final_text(turn: dict) -> str: - out = [] - for item in turn.get("items", []) or []: - if isinstance(item, dict) and item.get("type") == "agentMessage": - t = item.get("text") - if t: - out.append(t) - return "\n".join(out) - - -def _capture_turn_text(server: AppServer, thread_id: str, text: str, model: str) -> str: - """Run a turn and capture the full assistant response text.""" - rid = server._send( - "turn/start", - {"threadId": thread_id, "input": [{"type": "text", "text": text}], "model": model}, - ) - # Ack (status inProgress) — then stream until turn/completed. - server._wait(lambda m: m.get("id") == rid and ("result" in m or "error" in m), 30) - captured_text = [] - end = time.time() + 300.0 - while time.time() < end: - try: - msg = server._q.get(timeout=min(1.0, max(0.05, end - time.time()))) - except queue.Empty: - continue - method = msg.get("method") - params = msg.get("params") or {} - if method == "item/agentMessage/delta": - delta = _find_str(params, ("delta", "text")) - if delta: - captured_text.append(delta) - elif method == "turn/completed": - turn = params.get("turn", {}) - if turn.get("status") == "failed": - err = turn.get("error", {}) - return f"[FAILED: {err.get('message', err)}]" - # Collect any remaining text from final items - final = _final_text(turn) - if final and not captured_text: - captured_text.append(final) - return "".join(captured_text) - return "[TIMEOUT]" - - -def selftest() -> int: - """Non-interactive self-test: prove mid-session model switch with context.""" - home = build_codex_home() - server = AppServer(home) - try: - print("Starting codex app-server for self-test…") - server.initialize() - thread_id = server.start_thread(DEFAULT_MODEL) - print(f"Thread created with model {DEFAULT_MODEL}") - - # Turn 1: simple model A request - print("\n=== Turn 1 (model A) ===") - t1_prompt = "Reply with exactly: TURN1_OK" - print(f"Prompt: {t1_prompt}") - t1_response = _capture_turn_text(server, thread_id, t1_prompt, "system.ai.gpt-5-6-luna") - print(f"Response: {t1_response!r}") - if "TURN1_OK" not in t1_response: - print(f"ERROR: Turn 1 did not contain TURN1_OK") - return 1 - - # Turn 2: switch model and test context preservation - print("\n=== Turn 2 (model B, testing context) ===") - t2_prompt = "What token did you reply on the previous turn? Then say TURN2_OK." - print(f"Switching to gpt-5.5…") - print(f"Prompt: {t2_prompt}") - t2_response = _capture_turn_text(server, thread_id, t2_prompt, "gpt-5.5") - print(f"Response: {t2_response!r}") - - # Verify context was preserved: t2 should mention TURN1_OK - if "TURN1_OK" not in t2_response: - print(f"ERROR: Turn 2 did not contain TURN1_OK (context not preserved)") - return 1 - - if "TURN2_OK" not in t2_response: - print(f"WARNING: Turn 2 did not contain TURN2_OK (but context was preserved)") - - print("\n=== SUCCESS ===") - print("Mid-session model switch with context preservation verified!") - return 0 - except Exception as e: - print(f"ERROR: {e}", file=sys.stderr) - import traceback - traceback.print_exc(file=sys.stderr) - return 1 - finally: - server.close() - - -def main() -> int: - # Parse command-line arguments - if len(sys.argv) > 1: - if sys.argv[1] == "--selftest": - return selftest() - elif sys.argv[1] in ("--help", "-h"): - print(__doc__) - return 0 - else: - print(f"Unknown argument: {sys.argv[1]}", file=sys.stderr) - print(f"Use: {sys.argv[0]} [--selftest] [--help]", file=sys.stderr) - return 1 - - # Interactive mode - home = build_codex_home() - server = AppServer(home) - current_model = DEFAULT_MODEL - try: - print("Starting codex app-server…") - server.initialize() - thread_id = server.start_thread(current_model) - print(f"\nCodex ready. model = {current_model}") - print(f"Commands: /model /quit (try: {', '.join(EXAMPLE_MODELS)})\n") - while True: - try: - line = input(f"[{current_model}] › ").strip() - except (EOFError, KeyboardInterrupt): - print() - break - if not line: - continue - if line == "/quit": - break - if line.startswith("/model"): - arg = line[len("/model"):].strip() - if not arg: - print(f" current model: {current_model}") - else: - current_model = arg - print(f" → switched to {current_model} (applies to the next turn; history kept)") - continue - server.run_turn(thread_id, line, current_model) - return 0 - finally: - server.close() - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/scripts/codex_tui_interposer.py b/scripts/codex_tui_interposer.py deleted file mode 100644 index 081f8795..00000000 --- a/scripts/codex_tui_interposer.py +++ /dev/null @@ -1,340 +0,0 @@ -#!/usr/bin/env python3 -"""Arch B: a WebSocket MITM that lets you keep the REAL Codex TUI while the model -is switched under program control. - -Codex's remote transport (`codex --remote ws://…`) is WebSocket (a plain-JSONL -client is rejected with HTTP 400 "Connection header did not include 'upgrade'"; -a proper upgrade returns 101). Each JSON-RPC message is one WebSocket text frame. -This proxy sits between the TUI and a real `codex app-server`, forwarding every -frame untouched except: - - - `turn/start` (TUI->engine): after an initial hold of `--after` turns, its - `model` is rewritten to `--model`. `turn/start.model` is documented as - "override the model for this turn and subsequent turns", so the live session - retargets with history preserved. - - When the hold expires (right after your Nth prompt completes) it INJECTS a - `thread/settings/updated` notification (engine->TUI) carrying the new model, - so the TUI's on-screen model indicator follows the switch. - -So the demo is: start the TUI on model X, submit your first prompt (answered by -X), and from then on the session runs on `--model` (and the chip flips to it). - -Topology: - codex app-server --listen ws://127.0.0.1:8801 (real engine) - this interposer ws://127.0.0.1:8802 -> ws://127.0.0.1:8801 (switches model) - codex --remote ws://127.0.0.1:8802 --model system.ai.gpt-5-6-luna (real TUI) - -Run via uv so nothing is installed globally: - - uv run --with websockets python scripts/codex_tui_interposer.py \ - --listen 127.0.0.1:8802 --upstream ws://127.0.0.1:8801 \ - --model gpt-5.5 --after 1 - -Self-test (spawns its own app-server + a simulated TUI; proves hold + switch end -to end against the gateway): - - uv run --with websockets --with tomlkit python \ - scripts/codex_tui_interposer.py --selftest -""" -from __future__ import annotations - -import argparse -import asyncio -import contextlib -import json -import os -import socket -import subprocess -import sys -import time -from pathlib import Path - -from websockets.asyncio.client import connect -from websockets.asyncio.server import serve - -SETTINGS_UPDATED = "thread/settings/updated" - - -class Session: - """Per-TUI-connection state: hold the first `after` turns, then switch model.""" - - def __init__(self, target_model: str, after: int, log) -> None: - self.target = target_model - self.after = after - self.log = log - self.turns = 0 - self.thread_id: str | None = None - self.settings: dict | None = None - self.injected = False - - def on_tui_frame(self, raw: str) -> str: - """TUI->engine: rewrite turn/start.model once past the hold.""" - try: - msg = json.loads(raw) - except ValueError: - return raw - if not isinstance(msg, dict): - return raw - params = msg.get("params") - if msg.get("method") == "turn/start" and isinstance(params, dict): - self.turns += 1 - if isinstance(params.get("threadId"), str): - self.thread_id = params["threadId"] - if self.turns > self.after: - old = params.get("model") - if old != self.target: - params["model"] = self.target - self.log(f"[REWRITE] turn #{self.turns}: model {old!r} -> {self.target!r}") - return json.dumps(msg) - return raw - - def on_engine_frame(self, raw: str): - """engine->TUI: capture thread id/settings; after the hold's last turn - completes, return an injected settings-updated notification (or None).""" - try: - msg = json.loads(raw) - except ValueError: - return None - if not isinstance(msg, dict): - return None - params = msg.get("params") if isinstance(msg.get("params"), dict) else {} - result = msg.get("result") if isinstance(msg.get("result"), dict) else {} - # Capture threadId + a real threadSettings object wherever it appears. - for src in (params, result): - tid = src.get("threadId") or (src.get("thread") or {}).get("id") - if isinstance(tid, str): - self.thread_id = tid - ts = src.get("threadSettings") - if isinstance(ts, dict): - self.settings = ts - # When the hold's final turn completes, flip the on-screen model. - if ( - msg.get("method") == "turn/completed" - and not self.injected - and self.turns >= self.after - and self.thread_id - ): - self.injected = True - settings = dict(self.settings) if isinstance(self.settings, dict) else {} - settings["model"] = self.target - self.log(f"[INJECT] {SETTINGS_UPDATED}: model -> {self.target!r} (flip TUI chip)") - return { - "method": SETTINGS_UPDATED, - "params": {"threadId": self.thread_id, "threadSettings": settings}, - } - return None - - -async def _handle_tui(tui, upstream_uri: str, target_model: str, after: int, log) -> None: - path = getattr(getattr(tui, "request", None), "path", "/") or "/" - uri = upstream_uri.rstrip("/") + path - log(f"[CONN] TUI connected (path={path}); dialing app-server {uri}") - sess = Session(target_model, after, log) - async with connect(uri, max_size=None) as upstream: - - async def tui_to_app(): - async for frame in tui: - if isinstance(frame, str): - frame = sess.on_tui_frame(frame) - await upstream.send(frame) - - async def app_to_tui(): - async for frame in upstream: - await tui.send(frame) - if isinstance(frame, str): - inj = sess.on_engine_frame(frame) - if inj is not None: - await tui.send(json.dumps(inj)) - - a = asyncio.create_task(tui_to_app()) - b = asyncio.create_task(app_to_tui()) - _done, pending = await asyncio.wait({a, b}, return_when=asyncio.FIRST_COMPLETED) - for t in pending: - t.cancel() - with contextlib.suppress(asyncio.CancelledError): - await t - log("[CONN] TUI session closed") - - -async def serve_interposer(host: str, port: int, upstream_uri: str, model: str, after: int, *, quiet=False): - def log(m: str) -> None: - if not quiet: - print(m, file=sys.stderr, flush=True) - - async def handler(tui): - try: - await _handle_tui(tui, upstream_uri, model, after, log) - except Exception as exc: # noqa: BLE001 - one session must not kill the server - log(f"[ERR] session: {exc!r}") - - server = await serve(handler, host, port, max_size=None) - log(f"[READY] ws://{host}:{port} -> {upstream_uri} (hold {after} turn(s), then switch to {model!r})") - return server - - -# --------------------------------------------------------------------------- # -# Self-test -# --------------------------------------------------------------------------- # - -UCODE_CODEX_CONFIG = Path.home() / ".codex" / "ucode.config.toml" -SELFTEST_HOME = Path.home() / ".cache" / "ucode-codex-interposer" -START_MODEL = "system.ai.gpt-5-6-luna" -TARGET_MODEL = "gpt-5.5" -BOGUS_MODEL = "totally-bogus-model-zzz" - - -def _free_port() -> int: - s = socket.socket(); s.bind(("127.0.0.1", 0)); p = s.getsockname()[1]; s.close(); return p - - -def _build_codex_home() -> Path: - import tomllib - import tomlkit - - if not UCODE_CODEX_CONFIG.exists(): - sys.exit(f"Missing {UCODE_CODEX_CONFIG}; run `ucode codex` once so the provider block exists.") - src = tomllib.loads(UCODE_CODEX_CONFIG.read_text()) - doc = tomlkit.document() - for k in ("model_provider", "model", "model_reasoning_effort", "model_providers"): - if k in src: - doc[k] = src[k] - SELFTEST_HOME.mkdir(parents=True, exist_ok=True) - (SELFTEST_HOME / "config.toml").write_text(tomlkit.dumps(doc)) - return SELFTEST_HOME - - -async def _wait_healthz(port: int, timeout: float = 30.0) -> bool: - import urllib.request - - end = time.time() + timeout - while time.time() < end: - try: - with urllib.request.urlopen(f"http://127.0.0.1:{port}/healthz", timeout=1) as r: - if r.status == 200: - return True - except Exception: - await asyncio.sleep(0.25) - return False - - -async def _simulated_tui(port: int) -> dict: - """Turn 1 uses START_MODEL (should pass through). Turn 2 sends a BOGUS model - (should be rewritten to TARGET_MODEL and therefore succeed). Also watches for - the injected settings-updated frame after turn 1.""" - out = {"t1": None, "t2": None, "injected_model": None, "error": None} - nid = 0 - async with connect(f"ws://127.0.0.1:{port}", max_size=None) as ws: - async def send(method, params=None, notify=False): - nonlocal nid - m = {"method": method} - if not notify: - nid += 1 - m["id"] = nid - if params is not None: - m["params"] = params - await ws.send(json.dumps(m)) - return m.get("id") - - async def until(pred, timeout=180): - end = time.time() + timeout - while time.time() < end: - try: - frame = await asyncio.wait_for(ws.recv(), timeout=min(5, end - time.time())) - except asyncio.TimeoutError: - continue - if not isinstance(frame, str): - continue - try: - msg = json.loads(frame) - except ValueError: - continue - if msg.get("method") == SETTINGS_UPDATED: - out["injected_model"] = (msg.get("params", {}).get("threadSettings", {}) or {}).get("model") - if pred(msg): - return msg - return None - - await send("initialize", {"clientInfo": {"name": "sim", "version": "0"}, "capabilities": {}}) - await until(lambda m: m.get("id") == 1 and ("result" in m or "error" in m), 30) - await send("initialized", {}, notify=True) - rid = await send("thread/start", {"model": START_MODEL, "cwd": os.getcwd(), "approvalPolicy": "never"}) - ts = await until(lambda m: m.get("id") == rid and "result" in m, 60) - tid = ((ts or {}).get("result", {}).get("thread", {}) or {}).get("id") - if not tid: - out["error"] = f"thread/start failed: {json.dumps(ts)[:200]}" - return out - await send("turn/start", {"threadId": tid, "input": [{"type": "text", "text": "Say A"}], "model": START_MODEL}) - tc1 = await until(lambda m: m.get("method") == "turn/completed", 180) - out["t1"] = (tc1 or {}).get("params", {}).get("turn", {}).get("status") - await send("turn/start", {"threadId": tid, "input": [{"type": "text", "text": "Say B"}], "model": BOGUS_MODEL}) - tc2 = await until(lambda m: m.get("method") == "turn/completed", 180) - out["t2"] = (tc2 or {}).get("params", {}).get("turn", {}).get("status") - return out - - -async def _selftest() -> int: - home = _build_codex_home() - port_a, port_b = _free_port(), _free_port() - env = dict(os.environ); env["CODEX_HOME"] = str(home) - print(f"Starting codex app-server on ws://127.0.0.1:{port_a} …") - proc = subprocess.Popen( - ["codex", "app-server", "--listen", f"ws://127.0.0.1:{port_a}"], - stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, env=env, - ) - server = None - try: - if not await _wait_healthz(port_a): - print("app-server did not become healthy", file=sys.stderr) - return 1 - server = await serve_interposer("127.0.0.1", port_b, f"ws://127.0.0.1:{port_a}", TARGET_MODEL, after=1) - print(f"Interposer up: hold 1 turn on the TUI's model, then switch -> {TARGET_MODEL!r}\n") - r = await _simulated_tui(port_b) - print() - ok = ( - r["t1"] == "completed" # turn 1 ran on the pass-through START_MODEL - and r["t2"] == "completed" # turn 2 sent BOGUS but was rewritten -> succeeded - and r["injected_model"] == TARGET_MODEL # settings-updated injected to flip the chip - ) - print("=== RESULT ===") - print(f" turn1 (start model, passthrough): {r['t1']}") - print(f" turn2 (client sent BOGUS -> rewritten): {r['t2']}") - print(f" injected settings-updated model: {r['injected_model']!r}") - print(f" error: {r['error']!r}") - print("=== SUCCESS ===" if ok else "=== FAILED ===") - return 0 if ok else 1 - finally: - if server is not None: - server.close() - with contextlib.suppress(Exception): - await server.wait_closed() - proc.terminate() - with contextlib.suppress(Exception): - proc.wait(timeout=5) - - -def main() -> int: - ap = argparse.ArgumentParser(description="WebSocket MITM interposer for the Codex TUI (arch B).") - ap.add_argument("--listen", default="127.0.0.1:8802", help="host:port for the TUI to connect to") - ap.add_argument("--upstream", default="ws://127.0.0.1:8801", help="real app-server ws:// URI") - ap.add_argument("--model", default=TARGET_MODEL, help="model to switch to after the hold") - ap.add_argument("--after", type=int, default=1, help="pass through this many turns before switching (default 1)") - ap.add_argument("--selftest", action="store_true", help="spawn app-server + simulated TUI and prove hold+switch") - args = ap.parse_args() - - if args.selftest: - return asyncio.run(_selftest()) - - host, _, port = args.listen.partition(":") - - async def _run(): - await serve_interposer(host, int(port), args.upstream, args.model, args.after) - await asyncio.Future() - - try: - return asyncio.run(_run()) or 0 - except KeyboardInterrupt: - return 0 - - -if __name__ == "__main__": - sys.exit(main()) From e19698a28455a6be53cea73e739458c184f8b00e Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Thu, 20 Aug 2026 23:37:19 +0000 Subject: [PATCH 03/18] update --- pyproject.toml | 4 +-- src/ucode/agents/codex.py | 21 +++++-------- src/ucode/smart_routing/v2.py | 23 +++++++++++++++ tests/test_codex_smart_routing_v2.py | 44 ++-------------------------- 4 files changed, 34 insertions(+), 58 deletions(-) create mode 100644 src/ucode/smart_routing/v2.py diff --git a/pyproject.toml b/pyproject.toml index 17f3f683..52564f5c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,9 +26,7 @@ dependencies = [ "questionary>=2.0.0", "tomlkit>=0.13.0", "typer>=0.12.0", - # WebSocket client+server for the experimental `ENABLE_SMART_ROUTING_V2` Codex launch path: - # ucode interposes on Codex's `--remote` WebSocket transport to switch the model at runtime - # (see ucode.smart_routing.codex_interposer). + # enable changing codex model after first prompt "websockets>=13", ] diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 071f897f..a966aecd 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -28,6 +28,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.codex_hooks import ( remove_smart_routing_hooks, sync_smart_routing_hooks, @@ -51,14 +52,13 @@ # tool (codex, claude), so a workspace turns it on once. SMART_ROUTING_STATE_KEY = "smart_routing_enabled" -# Smart routing v2 (experimental, env-gated). When ENABLE_SMART_ROUTING_V2=1, a single -# `ucode codex` launches the REAL Codex TUI against a ucode-run `codex app-server`, with a -# WebSocket interposer (see smart_routing.codex_interposer) that holds the first turn on the -# normal model then switches to a fixed target. ucode owns all three processes and tears the +# Codex-specific smart-routing-v2 settings. The shared enable flag + hold-turns live in +# `smart_routing.v2`; here we keep only what is Codex-specific: the switch-to model and the +# app-server's CODEX_HOME / interposer log paths. When enabled, a single `ucode codex` +# launches the REAL Codex TUI against a ucode-run `codex app-server` with a WebSocket +# interposer (smart_routing.codex_interposer); ucode owns all three processes and tears the # app-server + interposer down when the TUI exits. -SMART_ROUTING_V2_ENV_VAR = "ENABLE_SMART_ROUTING_V2" SMART_ROUTING_V2_TARGET_MODEL = "gpt-5.5" # hardcoded switch-to model for now -SMART_ROUTING_V2_AFTER = 1 # pass through this many turns before switching SMART_ROUTING_V2_HOME = APP_DIR / "codex-v2-home" # CODEX_HOME for the ucode-run app-server SMART_ROUTING_V2_LOG = ( APP_DIR / "codex-v2-interposer.log" @@ -479,11 +479,6 @@ def _gpt_version_key(entry: tuple[str, tuple[int, int | None, int | None, str]]) _PROFILE_REJECTED_MAX_SECONDS = 3.0 -def smart_routing_v2_enabled() -> bool: - """Return whether the experimental smart-routing-v2 launch path is enabled.""" - return os.environ.get(SMART_ROUTING_V2_ENV_VAR) == "1" - - def _generate_v2_app_server_home(state: dict, model: str) -> Path: """Write an isolated CODEX_HOME whose config.toml carries the ucode gateway provider block, for the ucode-run `codex app-server`. @@ -558,7 +553,7 @@ def _launch_smart_routing_v2(state: dict, tool_args: list[str]) -> None: tui_port, f"ws://127.0.0.1:{app_port}", SMART_ROUTING_V2_TARGET_MODEL, - SMART_ROUTING_V2_AFTER, + smart_routing_v2.SWITCH_AFTER_TURNS, log_path=SMART_ROUTING_V2_LOG, ) # Foreground TUI. Popen (not exec) so this process stays alive to tear down the @@ -585,7 +580,7 @@ def _launch_smart_routing_v2(state: dict, tool_args: list[str]) -> None: def launch(state: dict, tool_args: list[str]) -> None: binary = SPEC["binary"] workspace = state.get("workspace") - if smart_routing_v2_enabled(): + if smart_routing_v2.enabled(): _launch_smart_routing_v2(state, tool_args) return if workspace: diff --git a/src/ucode/smart_routing/v2.py b/src/ucode/smart_routing/v2.py new file mode 100644 index 00000000..6b954f0d --- /dev/null +++ b/src/ucode/smart_routing/v2.py @@ -0,0 +1,23 @@ +"""Shared configuration for smart routing v2 — the runtime model-switching launch path. + +Smart routing v2 launches an agent's real TUI against a ucode-run app-server with a +WebSocket interposer that switches the model mid-session (see e.g. +``smart_routing.codex_interposer``). The enable flag and cross-agent knobs live here so +every routing-capable agent (Codex today, Claude Code next) reads them from one place; +each agent keeps its own target model, paths, and launch wiring. +""" + +from __future__ import annotations + +import os + +# Single env var that enables the v2 launch path for every routing-capable agent. +ENV_VAR = "ENABLE_SMART_ROUTING_V2" + +# Turns to keep the session on its starting model before switching (0 = switch immediately). +SWITCH_AFTER_TURNS = 1 + + +def enabled() -> bool: + """Return whether the smart-routing-v2 launch path is enabled via the env var.""" + return os.environ.get(ENV_VAR) == "1" diff --git a/tests/test_codex_smart_routing_v2.py b/tests/test_codex_smart_routing_v2.py index 1ebffd88..e62e336e 100644 --- a/tests/test_codex_smart_routing_v2.py +++ b/tests/test_codex_smart_routing_v2.py @@ -11,37 +11,6 @@ WS = "https://example.databricks.com" -class TestV2FlagGating: - def test_disabled_by_default(self, monkeypatch): - monkeypatch.delenv("ENABLE_SMART_ROUTING_V2", raising=False) - assert codex.smart_routing_v2_enabled() is False - - def test_enabled_when_1(self, monkeypatch): - monkeypatch.setenv("ENABLE_SMART_ROUTING_V2", "1") - assert codex.smart_routing_v2_enabled() is True - - def test_other_values_do_not_enable(self, monkeypatch): - monkeypatch.setenv("ENABLE_SMART_ROUTING_V2", "true") - assert codex.smart_routing_v2_enabled() is False - - def test_launch_dispatches_to_v2(self, monkeypatch): - monkeypatch.setenv("ENABLE_SMART_ROUTING_V2", "1") - called = {} - monkeypatch.setattr( - codex, - "_launch_smart_routing_v2", - lambda state, args: called.setdefault("hit", (state, args)), - ) - - # Should return via the v2 branch before touching normal launch/auth. - def _fail_if_normal_path(*_a, **_k): # pragma: no cover - only if v2 branch is skipped - raise AssertionError("normal launch path ran despite ENABLE_SMART_ROUTING_V2=1") - - monkeypatch.setattr(codex, "get_databricks_token", _fail_if_normal_path) - codex.launch({"workspace": WS}, ["--foo"]) - assert called["hit"] == ({"workspace": WS}, ["--foo"]) - - class TestGenerateV2Home: def test_writes_provider_config(self, tmp_path, monkeypatch): monkeypatch.setattr(codex, "SMART_ROUTING_V2_HOME", tmp_path / "v2home") @@ -64,6 +33,8 @@ def test_writes_provider_config(self, tmp_path, monkeypatch): class TestInterposerSession: + """The interposer's hold-then-switch + settings-injection logic (the novel behavior).""" + def _turn_start(self, model: str, thread_id: str = "t1") -> str: return json.dumps( { @@ -114,14 +85,3 @@ def test_injects_only_once(self): done = json.dumps({"method": "turn/completed", "params": {"threadId": "t1", "turn": {}}}) assert sess.on_engine_frame(done) is not None assert sess.on_engine_frame(done) is None # second completion: no re-inject - - -class TestInterposerHelpers: - def test_free_port_returns_usable_port(self): - port = codex_interposer.free_port() - assert isinstance(port, int) - assert 1024 <= port <= 65535 - - def test_wait_healthz_false_on_dead_port(self): - dead = codex_interposer.free_port() - assert codex_interposer.wait_healthz(dead, timeout=1.0) is False From 3fe3df921be77dd7dd46e6c23ea0e99c6e99f84f Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Mon, 24 Aug 2026 13:35:05 +0000 Subject: [PATCH 04/18] update --- src/ucode/agents/codex.py | 6 +- src/ucode/smart_routing/codex_interposer.py | 87 +++++++++++++++------ tests/test_codex_smart_routing_v2.py | 33 ++++++-- 3 files changed, 93 insertions(+), 33 deletions(-) diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index a966aecd..be63143c 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -58,7 +58,7 @@ # launches the REAL Codex TUI against a ucode-run `codex app-server` with a WebSocket # interposer (smart_routing.codex_interposer); ucode owns all three processes and tears the # app-server + interposer down when the TUI exits. -SMART_ROUTING_V2_TARGET_MODEL = "gpt-5.5" # hardcoded switch-to model for now +SMART_ROUTING_V2_TARGET_MODEL = "system.ai.glm-5-2" # hardcoded switch-to model for now SMART_ROUTING_V2_HOME = APP_DIR / "codex-v2-home" # CODEX_HOME for the ucode-run app-server SMART_ROUTING_V2_LOG = ( APP_DIR / "codex-v2-interposer.log" @@ -554,6 +554,10 @@ def _launch_smart_routing_v2(state: dict, tool_args: list[str]) -> None: f"ws://127.0.0.1:{app_port}", SMART_ROUTING_V2_TARGET_MODEL, smart_routing_v2.SWITCH_AFTER_TURNS, + switch_message=( + f"Databricks Smart Router selected model {SMART_ROUTING_V2_TARGET_MODEL} " + "due to low complexity, unclear intent, and no code reference." + ), log_path=SMART_ROUTING_V2_LOG, ) # Foreground TUI. Popen (not exec) so this process stays alive to tear down the diff --git a/src/ucode/smart_routing/codex_interposer.py b/src/ucode/smart_routing/codex_interposer.py index a729635c..10e54e3f 100644 --- a/src/ucode/smart_routing/codex_interposer.py +++ b/src/ucode/smart_routing/codex_interposer.py @@ -10,9 +10,11 @@ ``model`` is rewritten. ``turn/start.model`` is documented as "override the model for this turn and subsequent turns", so the live session retargets with history preserved. - - When the hold expires (right after the Nth prompt completes) an injected - ``thread/settings/updated`` notification (engine->TUI) carries the new model, - so the TUI's on-screen model indicator follows the switch. + - When the hold expires (right after the Nth prompt completes) two notifications + are injected (engine->TUI): a ``thread/settings/updated`` carrying the new + model, so the TUI's on-screen model indicator follows the switch, and — when a + ``switch_message`` is configured — a ``warning`` notification that surfaces a + one-line explanation of why the model was switched in the TUI's chat log. ``ucode.agents.codex`` runs :func:`start_interposer_thread` in a daemon thread while it owns the app-server subprocess and the ``codex --remote`` TUI, so the @@ -35,15 +37,24 @@ from websockets.asyncio.server import serve SETTINGS_UPDATED = "thread/settings/updated" +WARNING = "warning" class _Session: """Per-TUI-connection state: hold the first ``after`` turns, then switch model.""" - def __init__(self, target_model: str, after: int, log: Callable[[str], None]) -> None: + def __init__( + self, + target_model: str, + after: int, + log: Callable[[str], None], + switch_message: str | None = None, + ) -> None: self.target = target_model self.after = after self.log = log + # One-line explanation surfaced in the TUI when the switch fires; None skips it. + self.switch_message = switch_message self.turns = 0 self.thread_id: str | None = None self.settings: dict | None = None @@ -70,15 +81,19 @@ def on_tui_frame(self, raw: str) -> str: return json.dumps(msg) return raw - def on_engine_frame(self, raw: str) -> dict | None: + def on_engine_frame(self, raw: str) -> list[dict]: """engine->TUI: capture thread id/settings; after the hold's last turn - completes, return an injected settings-updated notification (or None).""" + completes, return the notifications to inject (empty list = none). + + On the switch this yields a ``thread/settings/updated`` (flips the TUI's + model chip) and, when ``switch_message`` is set, a ``warning`` that + explains why the model changed.""" try: msg = json.loads(raw) except ValueError: - return None + return [] if not isinstance(msg, dict): - return None + return [] params = msg.get("params") if isinstance(msg.get("params"), dict) else {} result = msg.get("result") if isinstance(msg.get("result"), dict) else {} for src in (params, result): @@ -98,18 +113,31 @@ def on_engine_frame(self, raw: str) -> dict | None: settings = dict(self.settings) if isinstance(self.settings, dict) else {} settings["model"] = self.target self.log(f"[INJECT] {SETTINGS_UPDATED}: model -> {self.target!r} (flip TUI chip)") - return { - "method": SETTINGS_UPDATED, - "params": {"threadId": self.thread_id, "threadSettings": settings}, - } - return None - - -async def _handle_tui(tui, upstream_uri: str, target_model: str, after: int, log) -> None: + injected: list[dict] = [ + { + "method": SETTINGS_UPDATED, + "params": {"threadId": self.thread_id, "threadSettings": settings}, + } + ] + if self.switch_message: + self.log(f"[INJECT] {WARNING}: {self.switch_message!r} (why model switched)") + injected.append( + { + "method": WARNING, + "params": {"threadId": self.thread_id, "message": self.switch_message}, + } + ) + return injected + return [] + + +async def _handle_tui( + tui, upstream_uri: str, target_model: str, after: int, log, switch_message: str | None = None +) -> None: path = getattr(getattr(tui, "request", None), "path", "/") or "/" uri = upstream_uri.rstrip("/") + path log(f"[CONN] TUI connected (path={path}); dialing app-server {uri}") - sess = _Session(target_model, after, log) + sess = _Session(target_model, after, log, switch_message) async with connect(uri, max_size=None) as upstream: async def tui_to_app(): @@ -122,8 +150,7 @@ async def app_to_tui(): async for frame in upstream: await tui.send(frame) if isinstance(frame, str): - inj = sess.on_engine_frame(frame) - if inj is not None: + for inj in sess.on_engine_frame(frame): await tui.send(json.dumps(inj)) a = asyncio.create_task(tui_to_app()) @@ -136,10 +163,18 @@ async def app_to_tui(): log("[CONN] TUI session closed") -async def _serve(host: str, port: int, upstream_uri: str, model: str, after: int, log): +async def _serve( + host: str, + port: int, + upstream_uri: str, + model: str, + after: int, + log, + switch_message: str | None = None, +): async def handler(tui): try: - await _handle_tui(tui, upstream_uri, model, after, log) + await _handle_tui(tui, upstream_uri, model, after, log, switch_message) except Exception as exc: # noqa: BLE001 - one session must never kill the server log(f"[ERR] session: {exc!r}") @@ -155,15 +190,17 @@ def start_interposer_thread( model: str, after: int, *, + switch_message: str | None = None, log_path: Path | None = None, ready_timeout: float = 10.0, ) -> tuple[threading.Thread, Callable[[], None]]: """Run the interposer's asyncio server in a daemon thread. Returns ``(thread, stop)``; ``stop()`` shuts the server down and stops the - loop. Logs go to ``log_path`` (appended) when given — never to stdout/stderr, - which the foreground TUI owns. Blocks until the server is listening (or - ``ready_timeout`` elapses).""" + loop. ``switch_message``, when set, is surfaced in the TUI as a ``warning`` + explaining why the model switched. Logs go to ``log_path`` (appended) when + given — never to stdout/stderr, which the foreground TUI owns. Blocks until + the server is listening (or ``ready_timeout`` elapses).""" def log(message: str) -> None: if log_path is None: @@ -183,7 +220,7 @@ def run() -> None: asyncio.set_event_loop(loop) try: holder["server"] = loop.run_until_complete( - _serve(host, port, upstream_uri, model, after, log) + _serve(host, port, upstream_uri, model, after, log, switch_message) ) except Exception as exc: # noqa: BLE001 - surface bind/connect failures to the log log(f"[ERR] failed to start interposer: {exc!r}") diff --git a/tests/test_codex_smart_routing_v2.py b/tests/test_codex_smart_routing_v2.py index e62e336e..e7bd7325 100644 --- a/tests/test_codex_smart_routing_v2.py +++ b/tests/test_codex_smart_routing_v2.py @@ -66,7 +66,7 @@ def test_non_turn_frames_pass_through(self): def test_injects_settings_update_after_hold(self): sess = codex_interposer._Session("gpt-5.5", after=1, log=lambda _m: None) sess.on_tui_frame(self._turn_start("luna")) # turn 1 (the hold) - inj = sess.on_engine_frame( + injected = sess.on_engine_frame( json.dumps( { "method": "turn/completed", @@ -74,14 +74,33 @@ def test_injects_settings_update_after_hold(self): } ) ) - assert inj is not None - assert inj["method"] == codex_interposer.SETTINGS_UPDATED - assert inj["params"]["threadId"] == "t1" - assert inj["params"]["threadSettings"]["model"] == "gpt-5.5" + settings = next(m for m in injected if m["method"] == codex_interposer.SETTINGS_UPDATED) + assert settings["params"]["threadId"] == "t1" + assert settings["params"]["threadSettings"]["model"] == "gpt-5.5" + + def test_injects_switch_warning_when_message_set(self): + sess = codex_interposer._Session( + "gpt-5.5", after=1, log=lambda _m: None, switch_message="switched to gpt-5.5 because X" + ) + sess.on_tui_frame(self._turn_start("luna")) # turn 1 (the hold) + injected = sess.on_engine_frame( + json.dumps({"method": "turn/completed", "params": {"threadId": "t1", "turn": {}}}) + ) + warning = next(m for m in injected if m["method"] == codex_interposer.WARNING) + assert warning["params"]["threadId"] == "t1" + assert warning["params"]["message"] == "switched to gpt-5.5 because X" + + def test_no_warning_without_message(self): + sess = codex_interposer._Session("gpt-5.5", after=1, log=lambda _m: None) + sess.on_tui_frame(self._turn_start("luna")) + injected = sess.on_engine_frame( + json.dumps({"method": "turn/completed", "params": {"threadId": "t1", "turn": {}}}) + ) + assert [m["method"] for m in injected] == [codex_interposer.SETTINGS_UPDATED] def test_injects_only_once(self): sess = codex_interposer._Session("gpt-5.5", after=1, log=lambda _m: None) sess.on_tui_frame(self._turn_start("luna")) done = json.dumps({"method": "turn/completed", "params": {"threadId": "t1", "turn": {}}}) - assert sess.on_engine_frame(done) is not None - assert sess.on_engine_frame(done) is None # second completion: no re-inject + assert sess.on_engine_frame(done) # first completion: injects + assert sess.on_engine_frame(done) == [] # second completion: no re-inject From d4f53a241f9ae897baece71ac8de4a6794ec629d Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Mon, 24 Aug 2026 14:05:35 +0000 Subject: [PATCH 05/18] update --- src/ucode/agents/codex.py | 2 +- src/ucode/smart_routing/codex_interposer.py | 89 +++++++++++++++++---- src/ucode/smart_routing/v2.py | 5 +- tests/test_codex_smart_routing_v2.py | 70 ++++++++++------ 4 files changed, 121 insertions(+), 45 deletions(-) diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index be63143c..ea2a5cee 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -555,7 +555,7 @@ def _launch_smart_routing_v2(state: dict, tool_args: list[str]) -> None: SMART_ROUTING_V2_TARGET_MODEL, smart_routing_v2.SWITCH_AFTER_TURNS, switch_message=( - f"Databricks Smart Router selected model {SMART_ROUTING_V2_TARGET_MODEL} " + f"✨ Databricks Smart Router selected model {SMART_ROUTING_V2_TARGET_MODEL} " "due to low complexity, unclear intent, and no code reference." ), log_path=SMART_ROUTING_V2_LOG, diff --git a/src/ucode/smart_routing/codex_interposer.py b/src/ucode/smart_routing/codex_interposer.py index 10e54e3f..4bf0a3ff 100644 --- a/src/ucode/smart_routing/codex_interposer.py +++ b/src/ucode/smart_routing/codex_interposer.py @@ -10,11 +10,18 @@ ``model`` is rewritten. ``turn/start.model`` is documented as "override the model for this turn and subsequent turns", so the live session retargets with history preserved. - - When the hold expires (right after the Nth prompt completes) two notifications - are injected (engine->TUI): a ``thread/settings/updated`` carrying the new - model, so the TUI's on-screen model indicator follows the switch, and — when a - ``switch_message`` is configured — a ``warning`` notification that surfaces a - one-line explanation of why the model was switched in the TUI's chat log. + - When the first switched turn starts — on that turn's ``turn/started``, before + any response items stream — two things are injected (engine->TUI): a + ``thread/settings/updated`` carrying the new model, so the TUI's on-screen + model indicator follows the switch, and — when a ``switch_message`` is + configured — an ``agentMessage`` item (as an ``item/started`` + ``item/completed`` + pair) that surfaces a one-line explanation of why the model was switched, ahead + of the model's reply. An ``agentMessage`` renders as ordinary chat text (no + warning styling); Codex's protocol has no neutral free-text notification + (``warning``, ``configWarning``, ``deprecationNotice`` all render as warnings), + so an item is the way to show an informational note. The ``item/started`` is + required: the TUI creates the message widget on ``item/started``, so a lone + ``item/completed`` has no widget to finalize and renders nothing. ``ucode.agents.codex`` runs :func:`start_interposer_thread` in a daemon thread while it owns the app-server subprocess and the ``codex --remote`` TUI, so the @@ -26,10 +33,12 @@ import asyncio import contextlib import json +import os import socket import threading import time import urllib.request +import uuid from collections.abc import Callable from pathlib import Path @@ -37,7 +46,12 @@ from websockets.asyncio.server import serve SETTINGS_UPDATED = "thread/settings/updated" -WARNING = "warning" +ITEM_STARTED = "item/started" +ITEM_COMPLETED = "item/completed" + +# Set UCODE_INTERPOSER_DEBUG=1 to dump raw engine->TUI item/turn frames (and the frames we +# inject) to the interposer log, for comparing our synthetic note against real assistant frames. +_DEBUG = os.environ.get("UCODE_INTERPOSER_DEBUG") == "1" class _Session: @@ -69,6 +83,8 @@ def on_tui_frame(self, raw: str) -> str: if not isinstance(msg, dict): return raw params = msg.get("params") + if _DEBUG and isinstance(msg.get("method"), str) and msg["method"].startswith("turn/"): + self.log(f"[DEBUG->engine] {msg['method']}: {raw[:1200]}") if msg.get("method") == "turn/start" and isinstance(params, dict): self.turns += 1 if isinstance(params.get("threadId"), str): @@ -82,18 +98,24 @@ def on_tui_frame(self, raw: str) -> str: return raw def on_engine_frame(self, raw: str) -> list[dict]: - """engine->TUI: capture thread id/settings; after the hold's last turn - completes, return the notifications to inject (empty list = none). - - On the switch this yields a ``thread/settings/updated`` (flips the TUI's - model chip) and, when ``switch_message`` is set, a ``warning`` that - explains why the model changed.""" + """engine->TUI: capture thread id/settings; when the first switched turn + starts, return the frames to inject (empty list = none). + + On the switched turn's ``turn/started`` — before its response streams — + this yields a ``thread/settings/updated`` (flips the TUI's model chip) + and, when ``switch_message`` is set, an ``item/completed`` carrying an + ``agentMessage`` — plain chat text (no warning styling) that explains why + the model changed, shown ahead of the model's reply.""" try: msg = json.loads(raw) except ValueError: return [] if not isinstance(msg, dict): return [] + if _DEBUG: + method = msg.get("method") + if isinstance(method, str) and (method.startswith("item/") or method.startswith("turn/")): + self.log(f"[DEBUG<-engine] {method}: {raw[:1800]}") params = msg.get("params") if isinstance(msg.get("params"), dict) else {} result = msg.get("result") if isinstance(msg.get("result"), dict) else {} for src in (params, result): @@ -104,9 +126,9 @@ def on_engine_frame(self, raw: str) -> list[dict]: if isinstance(ts, dict): self.settings = ts if ( - msg.get("method") == "turn/completed" + msg.get("method") == "turn/started" and not self.injected - and self.turns >= self.after + and self.turns > self.after and self.thread_id ): self.injected = True @@ -120,13 +142,46 @@ def on_engine_frame(self, raw: str) -> list[dict]: } ] if self.switch_message: - self.log(f"[INJECT] {WARNING}: {self.switch_message!r} (why model switched)") + params_obj = msg.get("params") if isinstance(msg.get("params"), dict) else {} + turn = params_obj.get("turn") if isinstance(params_obj, dict) else {} + turn_id = turn.get("id") if isinstance(turn, dict) else None + now_ms = int(time.time() * 1000) + item = { + "type": "agentMessage", + "id": f"ucode-smart-router-{uuid.uuid4().hex}", + "text": self.switch_message, + "phase": None, + "memoryCitation": None, + } + self.log(f"[INJECT] agentMessage note (started+completed): {self.switch_message!r}") + # The TUI creates the message widget on item/started; a lone item/completed + # has no widget to finalize and renders nothing. Send the full lifecycle with + # the text already populated (no deltas needed for a static note). + injected.append( + { + "method": ITEM_STARTED, + "params": { + "item": item, + "threadId": self.thread_id, + "turnId": turn_id, + "startedAtMs": now_ms, + }, + } + ) injected.append( { - "method": WARNING, - "params": {"threadId": self.thread_id, "message": self.switch_message}, + "method": ITEM_COMPLETED, + "params": { + "item": item, + "threadId": self.thread_id, + "turnId": turn_id, + "completedAtMs": now_ms, + }, } ) + if _DEBUG: + for frame in injected: + self.log(f"[DEBUG-inject] {json.dumps(frame)[:1800]}") return injected return [] diff --git a/src/ucode/smart_routing/v2.py b/src/ucode/smart_routing/v2.py index 6b954f0d..04523348 100644 --- a/src/ucode/smart_routing/v2.py +++ b/src/ucode/smart_routing/v2.py @@ -14,8 +14,9 @@ # Single env var that enables the v2 launch path for every routing-capable agent. ENV_VAR = "ENABLE_SMART_ROUTING_V2" -# Turns to keep the session on its starting model before switching (0 = switch immediately). -SWITCH_AFTER_TURNS = 1 +# Turns to keep the session on its starting model before switching (0 = switch immediately, +# i.e. the very first prompt already routes to the target model and shows the switch note). +SWITCH_AFTER_TURNS = 0 def enabled() -> bool: diff --git a/tests/test_codex_smart_routing_v2.py b/tests/test_codex_smart_routing_v2.py index e7bd7325..66dd8620 100644 --- a/tests/test_codex_smart_routing_v2.py +++ b/tests/test_codex_smart_routing_v2.py @@ -63,44 +63,64 @@ def test_non_turn_frames_pass_through(self): frame = json.dumps({"method": "initialize", "id": 1, "params": {}}) assert sess.on_tui_frame(frame) == frame - def test_injects_settings_update_after_hold(self): - sess = codex_interposer._Session("gpt-5.5", after=1, log=lambda _m: None) - sess.on_tui_frame(self._turn_start("luna")) # turn 1 (the hold) - injected = sess.on_engine_frame( - json.dumps( - { - "method": "turn/completed", - "params": {"threadId": "t1", "turn": {"status": "completed"}}, - } - ) + def _turn_started(self, turn_id: str, thread_id: str = "t1") -> str: + return json.dumps( + {"method": "turn/started", "params": {"threadId": thread_id, "turn": {"id": turn_id}}} ) + + def test_holds_note_until_switched_turn_starts(self): + # The note/chip-flip fire on the SWITCHED turn's turn/started (before its + # response), never on the held turn. + sess = codex_interposer._Session("gpt-5.5", after=1, log=lambda _m: None) + sess.on_tui_frame(self._turn_start("luna")) # turn 1 (held) + assert sess.on_engine_frame(self._turn_started("turn-1")) == [] # held: no inject + sess.on_tui_frame(self._turn_start("luna")) # turn 2 (switched) + injected = sess.on_engine_frame(self._turn_started("turn-2")) settings = next(m for m in injected if m["method"] == codex_interposer.SETTINGS_UPDATED) assert settings["params"]["threadId"] == "t1" assert settings["params"]["threadSettings"]["model"] == "gpt-5.5" - def test_injects_switch_warning_when_message_set(self): + def test_injects_switch_note_as_agent_message_when_message_set(self): sess = codex_interposer._Session( - "gpt-5.5", after=1, log=lambda _m: None, switch_message="switched to gpt-5.5 because X" + "gpt-5.5", after=1, log=lambda _m: None, switch_message="selected glm-5-2 because X" ) - sess.on_tui_frame(self._turn_start("luna")) # turn 1 (the hold) - injected = sess.on_engine_frame( - json.dumps({"method": "turn/completed", "params": {"threadId": "t1", "turn": {}}}) + sess.on_tui_frame(self._turn_start("luna")) # turn 1 (held) + sess.on_engine_frame(self._turn_started("turn-1")) + sess.on_tui_frame(self._turn_start("luna")) # turn 2 (switched) + injected = sess.on_engine_frame(self._turn_started("turn-2")) + # The note is a full agentMessage lifecycle: item/started THEN item/completed, + # both carrying the same item (a lone item/completed renders nothing in the TUI). + started = next(m for m in injected if m["method"] == codex_interposer.ITEM_STARTED) + completed = next(m for m in injected if m["method"] == codex_interposer.ITEM_COMPLETED) + assert started["params"]["turnId"] == "turn-2" + assert completed["params"]["turnId"] == "turn-2" + for frame in (started, completed): + item = frame["params"]["item"] + # An agentMessage renders as plain chat text, not a yellow warning banner. + assert item["type"] == "agentMessage" + assert item["text"] == "selected glm-5-2 because X" + assert started["params"]["item"]["id"] == completed["params"]["item"]["id"] + + def test_after_zero_injects_on_first_turn_start(self): + sess = codex_interposer._Session( + "gpt-5.5", after=0, log=lambda _m: None, switch_message="switched" ) - warning = next(m for m in injected if m["method"] == codex_interposer.WARNING) - assert warning["params"]["threadId"] == "t1" - assert warning["params"]["message"] == "switched to gpt-5.5 because X" + sess.on_tui_frame(self._turn_start("luna")) # turn 1 (switched immediately) + injected = sess.on_engine_frame(self._turn_started("turn-1")) + methods = [m["method"] for m in injected] + assert codex_interposer.ITEM_STARTED in methods + assert codex_interposer.ITEM_COMPLETED in methods - def test_no_warning_without_message(self): + def test_no_note_without_message(self): sess = codex_interposer._Session("gpt-5.5", after=1, log=lambda _m: None) sess.on_tui_frame(self._turn_start("luna")) - injected = sess.on_engine_frame( - json.dumps({"method": "turn/completed", "params": {"threadId": "t1", "turn": {}}}) - ) + sess.on_tui_frame(self._turn_start("luna")) + injected = sess.on_engine_frame(self._turn_started("turn-2")) assert [m["method"] for m in injected] == [codex_interposer.SETTINGS_UPDATED] def test_injects_only_once(self): sess = codex_interposer._Session("gpt-5.5", after=1, log=lambda _m: None) sess.on_tui_frame(self._turn_start("luna")) - done = json.dumps({"method": "turn/completed", "params": {"threadId": "t1", "turn": {}}}) - assert sess.on_engine_frame(done) # first completion: injects - assert sess.on_engine_frame(done) == [] # second completion: no re-inject + sess.on_tui_frame(self._turn_start("luna")) + assert sess.on_engine_frame(self._turn_started("turn-2")) # switched turn: injects + assert sess.on_engine_frame(self._turn_started("turn-3")) == [] # later turn: no re-inject From 87bddcb2cbd82da973d12e224fb3684c06df385c Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Mon, 24 Aug 2026 14:13:12 +0000 Subject: [PATCH 06/18] update --- src/ucode/agents/codex.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index ea2a5cee..95812cfd 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -58,6 +58,7 @@ # launches the REAL Codex TUI against a ucode-run `codex app-server` with a WebSocket # interposer (smart_routing.codex_interposer); ucode owns all three processes and tears the # app-server + interposer down when the TUI exits. +SMART_ROUTING_V2_START_MODEL = "gpt-5.5" # hardcoded start model for now SMART_ROUTING_V2_TARGET_MODEL = "system.ai.glm-5-2" # hardcoded switch-to model for now SMART_ROUTING_V2_HOME = APP_DIR / "codex-v2-home" # CODEX_HOME for the ucode-run app-server SMART_ROUTING_V2_LOG = ( @@ -518,11 +519,7 @@ def _launch_smart_routing_v2(state: dict, tool_args: list[str]) -> None: raise RuntimeError( "Smart routing v2 needs a configured workspace; run `ucode configure codex` first." ) - start_model = default_model(state) - if not start_model: - raise RuntimeError( - "Smart routing v2 could not determine a starting Codex model for this workspace." - ) + start_model = SMART_ROUTING_V2_START_MODEL os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile")) home = _generate_v2_app_server_home(state, start_model) From cd074719553bee88fdb4d56e7abc2f406953a5ce Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Tue, 25 Aug 2026 20:10:31 +0000 Subject: [PATCH 07/18] Refine Codex smart routing launch message --- src/ucode/agents/codex.py | 41 ++++++++++++++++------------ tests/test_codex_smart_routing_v2.py | 12 ++++++++ 2 files changed, 35 insertions(+), 18 deletions(-) diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 95812cfd..d99ef8ae 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -35,7 +35,7 @@ ) from ucode.state import mark_tool_managed, save_state from ucode.telemetry import agent_version, ucode_version -from ucode.ui import print_note, print_warning_err +from ucode.ui import print_warning_err CODEX_CONFIG_DIR = Path.home() / ".codex" CODEX_PROFILE_NAME = "ucode" @@ -52,18 +52,25 @@ # tool (codex, claude), so a workspace turns it on once. SMART_ROUTING_STATE_KEY = "smart_routing_enabled" -# Codex-specific smart-routing-v2 settings. The shared enable flag + hold-turns live in -# `smart_routing.v2`; here we keep only what is Codex-specific: the switch-to model and the -# app-server's CODEX_HOME / interposer log paths. When enabled, a single `ucode codex` -# launches the REAL Codex TUI against a ucode-run `codex app-server` with a WebSocket -# interposer (smart_routing.codex_interposer); ucode owns all three processes and tears the -# app-server + interposer down when the TUI exits. -SMART_ROUTING_V2_START_MODEL = "gpt-5.5" # hardcoded start model for now SMART_ROUTING_V2_TARGET_MODEL = "system.ai.glm-5-2" # hardcoded switch-to model for now SMART_ROUTING_V2_HOME = APP_DIR / "codex-v2-home" # CODEX_HOME for the ucode-run app-server SMART_ROUTING_V2_LOG = ( APP_DIR / "codex-v2-interposer.log" ) # interposer log (not stdout: TUI owns it) +SMART_ROUTING_V2_REASON = "Low complexity, unclear intent, and no code reference." + + +def _smart_routing_switch_message(model: str, reason: str) -> str: + lines = [ + "Using Unity Gateway Smart Router.", + f"Selected Model : {model}", + f"Reason : {reason}", + ] + width = max(len(line) for line in lines) + border = "─" * (width + 2) + return "\n".join( + [f"┌{border}┐", *(f"│ {line:<{width}} │" for line in lines), f"└{border}┘"] + ) SPEC: ToolSpec = { @@ -519,19 +526,17 @@ def _launch_smart_routing_v2(state: dict, tool_args: list[str]) -> None: raise RuntimeError( "Smart routing v2 needs a configured workspace; run `ucode configure codex` first." ) - start_model = SMART_ROUTING_V2_START_MODEL + start_model = default_model(state) + if not start_model: + raise RuntimeError( + "Smart routing v2 could not determine a starting Codex model for this workspace." + ) os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile")) home = _generate_v2_app_server_home(state, start_model) app_port = codex_interposer.free_port() tui_port = codex_interposer.free_port() - print_note( - f"Smart routing v2: starting on {start_model}, switching to " - f"{SMART_ROUTING_V2_TARGET_MODEL} after the first prompt " - f"(interposer log: {SMART_ROUTING_V2_LOG})." - ) - app_server = subprocess.Popen( [binary, "app-server", "--listen", f"ws://127.0.0.1:{app_port}"], env={**os.environ, "CODEX_HOME": str(home)}, @@ -551,9 +556,9 @@ def _launch_smart_routing_v2(state: dict, tool_args: list[str]) -> None: f"ws://127.0.0.1:{app_port}", SMART_ROUTING_V2_TARGET_MODEL, smart_routing_v2.SWITCH_AFTER_TURNS, - switch_message=( - f"✨ Databricks Smart Router selected model {SMART_ROUTING_V2_TARGET_MODEL} " - "due to low complexity, unclear intent, and no code reference." + switch_message=_smart_routing_switch_message( + SMART_ROUTING_V2_TARGET_MODEL, + SMART_ROUTING_V2_REASON, ), log_path=SMART_ROUTING_V2_LOG, ) diff --git a/tests/test_codex_smart_routing_v2.py b/tests/test_codex_smart_routing_v2.py index 66dd8620..ac5a33dc 100644 --- a/tests/test_codex_smart_routing_v2.py +++ b/tests/test_codex_smart_routing_v2.py @@ -32,6 +32,18 @@ def test_writes_provider_config(self, tmp_path, monkeypatch): assert "myprof" in provider["auth"]["args"] +def test_smart_routing_switch_message_is_boxed(): + message = codex._smart_routing_switch_message("model-x", "Because X.") + + assert message == ( + "┌───────────────────────────────────┐\n" + "│ Using Unity Gateway Smart Router. │\n" + "│ Selected Model : model-x │\n" + "│ Reason : Because X. │\n" + "└───────────────────────────────────┘" + ) + + class TestInterposerSession: """The interposer's hold-then-switch + settings-injection logic (the novel behavior).""" From ce910e124a6402732f61de7f49727dedd7e22f66 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Tue, 25 Aug 2026 20:20:27 +0000 Subject: [PATCH 08/18] Move Codex smart routing lifecycle to v2 --- src/ucode/agents/codex.py | 127 ++----------------- src/ucode/smart_routing/codex_interposer.py | 28 +++-- src/ucode/smart_routing/v2.py | 128 +++++++++++++++++++- tests/test_codex_smart_routing_v2.py | 95 ++++++++++++++- 4 files changed, 237 insertions(+), 141 deletions(-) diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index d99ef8ae..28afe107 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -4,7 +4,6 @@ import os import re -import signal import subprocess import sys import time @@ -52,27 +51,6 @@ # tool (codex, claude), so a workspace turns it on once. SMART_ROUTING_STATE_KEY = "smart_routing_enabled" -SMART_ROUTING_V2_TARGET_MODEL = "system.ai.glm-5-2" # hardcoded switch-to model for now -SMART_ROUTING_V2_HOME = APP_DIR / "codex-v2-home" # CODEX_HOME for the ucode-run app-server -SMART_ROUTING_V2_LOG = ( - APP_DIR / "codex-v2-interposer.log" -) # interposer log (not stdout: TUI owns it) -SMART_ROUTING_V2_REASON = "Low complexity, unclear intent, and no code reference." - - -def _smart_routing_switch_message(model: str, reason: str) -> str: - lines = [ - "Using Unity Gateway Smart Router.", - f"Selected Model : {model}", - f"Reason : {reason}", - ] - width = max(len(line) for line in lines) - border = "─" * (width + 2) - return "\n".join( - [f"┌{border}┐", *(f"│ {line:<{width}} │" for line in lines), f"└{border}┘"] - ) - - SPEC: ToolSpec = { "binary": "codex", "package": "@openai/codex", @@ -487,108 +465,17 @@ def _gpt_version_key(entry: tuple[str, tuple[int, int | None, int | None, str]]) _PROFILE_REJECTED_MAX_SECONDS = 3.0 -def _generate_v2_app_server_home(state: dict, model: str) -> Path: - """Write an isolated CODEX_HOME whose config.toml carries the ucode gateway - provider block, for the ucode-run `codex app-server`. - - The app-server rejects the global `--profile`, so a default-config CODEX_HOME - is how it inherits ucode's gateway (base_url + `ucode auth-token` refresh). - Reuses `render_overlay` — the same provider block `ucode configure codex` writes.""" - home = SMART_ROUTING_V2_HOME - home.mkdir(parents=True, exist_ok=True) - config_path = home / "config.toml" - overlay = render_overlay( - state["workspace"], - model, - state.get("profile"), - use_pat=bool(state.get("use_pat")), - ) - doc = read_toml_safe(config_path) - deep_merge_dict(doc, overlay) - write_toml_file(config_path, doc) - return home - - -def _launch_smart_routing_v2(state: dict, tool_args: list[str]) -> None: - """Experimental single-command launch of the real Codex TUI with runtime model switching. - - ucode owns three processes: a `codex app-server` subprocess, the WebSocket interposer - (daemon thread), and the `codex --remote` TUI (foreground). The interposer holds the first - turn on the normal model, then rewrites subsequent turns to SMART_ROUTING_V2_TARGET_MODEL and - injects a settings update so the TUI reflects the switch. The app-server + interposer are torn - down when the TUI exits. Mirrors the lifecycle of `claude.py::_launch_relayed`. - """ - from ucode.smart_routing import codex_interposer - - binary = SPEC["binary"] - workspace = state.get("workspace") - if not workspace: - raise RuntimeError( - "Smart routing v2 needs a configured workspace; run `ucode configure codex` first." - ) - start_model = default_model(state) - if not start_model: - raise RuntimeError( - "Smart routing v2 could not determine a starting Codex model for this workspace." - ) - - os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile")) - home = _generate_v2_app_server_home(state, start_model) - app_port = codex_interposer.free_port() - tui_port = codex_interposer.free_port() - - app_server = subprocess.Popen( - [binary, "app-server", "--listen", f"ws://127.0.0.1:{app_port}"], - env={**os.environ, "CODEX_HOME": str(home)}, - stdin=subprocess.DEVNULL, - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - ) - stop_interposer = None - try: - if not codex_interposer.wait_healthz(app_port, timeout=30): - raise RuntimeError( - "Codex app-server did not become ready for smart routing v2; check workspace auth." - ) - _thread, stop_interposer = codex_interposer.start_interposer_thread( - "127.0.0.1", - tui_port, - f"ws://127.0.0.1:{app_port}", - SMART_ROUTING_V2_TARGET_MODEL, - smart_routing_v2.SWITCH_AFTER_TURNS, - switch_message=_smart_routing_switch_message( - SMART_ROUTING_V2_TARGET_MODEL, - SMART_ROUTING_V2_REASON, - ), - log_path=SMART_ROUTING_V2_LOG, - ) - # Foreground TUI. Popen (not exec) so this process stays alive to tear down the - # app-server + interposer when the TUI exits (see claude.py::_launch_relayed). - tui = subprocess.Popen( - [binary, "--remote", f"ws://127.0.0.1:{tui_port}", "--model", start_model, *tool_args] - ) - try: - returncode = tui.wait() - except KeyboardInterrupt: - tui.send_signal(signal.SIGINT) - returncode = tui.wait() - finally: - if stop_interposer is not None: - stop_interposer() - app_server.terminate() - try: - app_server.wait(timeout=5) - except Exception: # noqa: BLE001 - the app-server must never linger - app_server.kill() - sys.exit(returncode) - - def launch(state: dict, tool_args: list[str]) -> None: binary = SPEC["binary"] workspace = state.get("workspace") if smart_routing_v2.enabled(): - _launch_smart_routing_v2(state, tool_args) - return + smart_routing_v2.launch_codex( + state, + tool_args, + binary=binary, + start_model=default_model(state), + render_overlay=render_overlay, + ) if workspace: os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile")) # Run codex with --profile first — the TUI and runtime subcommands diff --git a/src/ucode/smart_routing/codex_interposer.py b/src/ucode/smart_routing/codex_interposer.py index 4bf0a3ff..e6488560 100644 --- a/src/ucode/smart_routing/codex_interposer.py +++ b/src/ucode/smart_routing/codex_interposer.py @@ -15,7 +15,7 @@ ``thread/settings/updated`` carrying the new model, so the TUI's on-screen model indicator follows the switch, and — when a ``switch_message`` is configured — an ``agentMessage`` item (as an ``item/started`` + ``item/completed`` - pair) that surfaces a one-line explanation of why the model was switched, ahead + pair) that surfaces an explanation of why the model was switched, ahead of the model's reply. An ``agentMessage`` renders as ordinary chat text (no warning styling); Codex's protocol has no neutral free-text notification (``warning``, ``configWarning``, ``deprecationNotice`` all render as warnings), @@ -23,9 +23,9 @@ required: the TUI creates the message widget on ``item/started``, so a lone ``item/completed`` has no widget to finalize and renders nothing. -``ucode.agents.codex`` runs :func:`start_interposer_thread` in a daemon thread -while it owns the app-server subprocess and the ``codex --remote`` TUI, so the -whole thing launches from the single ``ucode codex`` command. +``ucode.smart_routing.v2`` runs :func:`start_interposer_thread` in a daemon +thread while it owns the app-server subprocess and the ``codex --remote`` TUI, +so the whole thing launches from the single ``ucode codex`` command. """ from __future__ import annotations @@ -48,6 +48,7 @@ SETTINGS_UPDATED = "thread/settings/updated" ITEM_STARTED = "item/started" ITEM_COMPLETED = "item/completed" +LOOPBACK_HOST = "127.0.0.1" # Set UCODE_INTERPOSER_DEBUG=1 to dump raw engine->TUI item/turn frames (and the frames we # inject) to the interposer log, for comparing our synthetic note against real assistant frames. @@ -103,9 +104,10 @@ def on_engine_frame(self, raw: str) -> list[dict]: On the switched turn's ``turn/started`` — before its response streams — this yields a ``thread/settings/updated`` (flips the TUI's model chip) - and, when ``switch_message`` is set, an ``item/completed`` carrying an - ``agentMessage`` — plain chat text (no warning styling) that explains why - the model changed, shown ahead of the model's reply.""" + and, when ``switch_message`` is set, an ``item/started`` + + ``item/completed`` pair carrying an ``agentMessage`` — plain chat text + (no warning styling) that explains why the model changed, shown ahead of + the model's reply.""" try: msg = json.loads(raw) except ValueError: @@ -114,7 +116,9 @@ def on_engine_frame(self, raw: str) -> list[dict]: return [] if _DEBUG: method = msg.get("method") - if isinstance(method, str) and (method.startswith("item/") or method.startswith("turn/")): + if isinstance(method, str) and ( + method.startswith("item/") or method.startswith("turn/") + ): self.log(f"[DEBUG<-engine] {method}: {raw[:1800]}") params = msg.get("params") if isinstance(msg.get("params"), dict) else {} result = msg.get("result") if isinstance(msg.get("result"), dict) else {} @@ -252,8 +256,8 @@ def start_interposer_thread( """Run the interposer's asyncio server in a daemon thread. Returns ``(thread, stop)``; ``stop()`` shuts the server down and stops the - loop. ``switch_message``, when set, is surfaced in the TUI as a ``warning`` - explaining why the model switched. Logs go to ``log_path`` (appended) when + loop. ``switch_message``, when set, is surfaced in the TUI as an + ``agentMessage`` explaining why the model switched. Logs go to ``log_path`` (appended) when given — never to stdout/stderr, which the foreground TUI owns. Blocks until the server is listening (or ``ready_timeout`` elapses).""" @@ -306,7 +310,7 @@ def stop() -> None: def free_port() -> int: """Grab an unused loopback TCP port (races are irrelevant for local ephemeral use).""" sock = socket.socket() - sock.bind(("127.0.0.1", 0)) + sock.bind((LOOPBACK_HOST, 0)) port = sock.getsockname()[1] sock.close() return port @@ -314,7 +318,7 @@ def free_port() -> int: def wait_healthz(port: int, timeout: float = 30.0) -> bool: """Poll the app-server's ``/healthz`` until it returns 200, or timeout.""" - url = f"http://127.0.0.1:{port}/healthz" + url = f"http://{LOOPBACK_HOST}:{port}/healthz" end = time.time() + timeout while time.time() < end: try: diff --git a/src/ucode/smart_routing/v2.py b/src/ucode/smart_routing/v2.py index 04523348..c2f55d85 100644 --- a/src/ucode/smart_routing/v2.py +++ b/src/ucode/smart_routing/v2.py @@ -1,15 +1,24 @@ -"""Shared configuration for smart routing v2 — the runtime model-switching launch path. +"""Runtime model-switching launch path for smart routing v2. Smart routing v2 launches an agent's real TUI against a ucode-run app-server with a WebSocket interposer that switches the model mid-session (see e.g. -``smart_routing.codex_interposer``). The enable flag and cross-agent knobs live here so -every routing-capable agent (Codex today, Claude Code next) reads them from one place; -each agent keeps its own target model, paths, and launch wiring. +``smart_routing.codex_interposer``). The enable flag, launch configuration, and process +lifecycle live here so agent modules only need to supply their provider overlay. """ from __future__ import annotations import os +import signal +import subprocess +import sys +from collections.abc import Callable +from pathlib import Path +from typing import NoReturn + +from ucode.config_io import APP_DIR, deep_merge_dict, read_toml_safe, write_toml_file +from ucode.databricks import get_databricks_token +from ucode.smart_routing import codex_interposer # Single env var that enables the v2 launch path for every routing-capable agent. ENV_VAR = "ENABLE_SMART_ROUTING_V2" @@ -18,7 +27,118 @@ # i.e. the very first prompt already routes to the target model and shows the switch note). SWITCH_AFTER_TURNS = 0 +CODEX_TARGET_MODEL = "system.ai.glm-5-2" # TODO(lilly): replace with smart router. +CODEX_APP_SERVER_HOME = APP_DIR / "codex-v2-home" +CODEX_INTERPOSER_LOG = APP_DIR / "codex-v2-interposer.log" +CODEX_SWITCH_REASON = "Low complexity, unclear intent, and no code reference." # TODO(lilly): replace with smart router rationale. + +APP_SERVER_READY_TIMEOUT_SECONDS = 30 +PROCESS_SHUTDOWN_TIMEOUT_SECONDS = 5 +OAUTH_TOKEN_ENV_VAR = "OAUTH_TOKEN" +CODEX_HOME_ENV_VAR = "CODEX_HOME" + def enabled() -> bool: """Return whether the smart-routing-v2 launch path is enabled via the env var.""" return os.environ.get(ENV_VAR) == "1" + + +def _loopback_websocket_url(port: int) -> str: + return f"ws://{codex_interposer.LOOPBACK_HOST}:{port}" + + +def _switch_message(model: str, reason: str) -> str: + lines = [ + "Using Unity Gateway Smart Router.", + f"Selected Model : {model}", + f"Reason : {reason}", + ] + width = max(len(line) for line in lines) + border = "─" * (width + 2) + return "\n".join([f"┌{border}┐", *(f"│ {line:<{width}} │" for line in lines), f"└{border}┘"]) + + +def _generate_codex_app_server_home( + state: dict, + model: str, + render_overlay: Callable[..., dict], +) -> Path: + """Write the isolated CODEX_HOME used by the ucode-run app-server.""" + CODEX_APP_SERVER_HOME.mkdir(parents=True, exist_ok=True) + config_path = CODEX_APP_SERVER_HOME / "config.toml" + overlay = render_overlay( + state["workspace"], + model, + state.get("profile"), + use_pat=bool(state.get("use_pat")), + ) + doc = read_toml_safe(config_path) + deep_merge_dict(doc, overlay) + write_toml_file(config_path, doc) + return CODEX_APP_SERVER_HOME + + +def launch_codex( + state: dict, + tool_args: list[str], + *, + binary: str, + start_model: str | None, + render_overlay: Callable[..., dict], +) -> NoReturn: + """Launch the Codex app-server, interposer, and remote TUI as one lifecycle.""" + workspace = state.get("workspace") + if not workspace: + raise RuntimeError( + "Smart routing v2 needs a configured workspace; run `ucode configure codex` first." + ) + if not start_model: + raise RuntimeError( + "Smart routing v2 could not determine a starting Codex model for this workspace." + ) + + os.environ[OAUTH_TOKEN_ENV_VAR] = get_databricks_token(workspace, state.get("profile")) + home = _generate_codex_app_server_home(state, start_model, render_overlay) + app_port = codex_interposer.free_port() + tui_port = codex_interposer.free_port() + app_server_url = _loopback_websocket_url(app_port) + tui_url = _loopback_websocket_url(tui_port) + + app_server = subprocess.Popen( + [binary, "app-server", "--listen", app_server_url], + env={**os.environ, CODEX_HOME_ENV_VAR: str(home)}, + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + stop_interposer = None + try: + if not codex_interposer.wait_healthz(app_port, timeout=APP_SERVER_READY_TIMEOUT_SECONDS): + raise RuntimeError( + "Codex app-server did not become ready for smart routing v2; check workspace auth." + ) + _thread, stop_interposer = codex_interposer.start_interposer_thread( + codex_interposer.LOOPBACK_HOST, + tui_port, + app_server_url, + CODEX_TARGET_MODEL, + SWITCH_AFTER_TURNS, + switch_message=_switch_message(CODEX_TARGET_MODEL, CODEX_SWITCH_REASON), + log_path=CODEX_INTERPOSER_LOG, + ) + # Keep ucode alive while the TUI runs so it can tear down the app-server and interposer. + tui = subprocess.Popen([binary, "--remote", tui_url, "--model", start_model, *tool_args]) + try: + returncode = tui.wait() + except KeyboardInterrupt: + tui.send_signal(signal.SIGINT) + returncode = tui.wait() + finally: + if stop_interposer is not None: + stop_interposer() + app_server.terminate() + try: + app_server.wait(timeout=PROCESS_SHUTDOWN_TIMEOUT_SECONDS) + except Exception: # noqa: BLE001 - the app-server must never linger + app_server.kill() + sys.exit(returncode) diff --git a/tests/test_codex_smart_routing_v2.py b/tests/test_codex_smart_routing_v2.py index ac5a33dc..10389f58 100644 --- a/tests/test_codex_smart_routing_v2.py +++ b/tests/test_codex_smart_routing_v2.py @@ -4,21 +4,25 @@ import json +import pytest + from ucode.agents import codex from ucode.config_io import read_toml_safe -from ucode.smart_routing import codex_interposer +from ucode.smart_routing import codex_interposer, v2 WS = "https://example.databricks.com" class TestGenerateV2Home: def test_writes_provider_config(self, tmp_path, monkeypatch): - monkeypatch.setattr(codex, "SMART_ROUTING_V2_HOME", tmp_path / "v2home") + monkeypatch.setattr(v2, "CODEX_APP_SERVER_HOME", tmp_path / "v2home") monkeypatch.setattr(codex, "ucode_version", lambda: "0.1.0") monkeypatch.setattr(codex, "agent_version", lambda binary: "0.148.0") - home = codex._generate_v2_app_server_home( - {"workspace": WS, "profile": "myprof"}, "gpt-5.6-luna" + home = v2._generate_codex_app_server_home( + {"workspace": WS, "profile": "myprof"}, + "gpt-5.6-luna", + codex.render_overlay, ) assert home == tmp_path / "v2home" @@ -33,7 +37,7 @@ def test_writes_provider_config(self, tmp_path, monkeypatch): def test_smart_routing_switch_message_is_boxed(): - message = codex._smart_routing_switch_message("model-x", "Because X.") + message = v2._switch_message("model-x", "Because X.") assert message == ( "┌───────────────────────────────────┐\n" @@ -44,6 +48,87 @@ def test_smart_routing_switch_message_is_boxed(): ) +class TestLaunchCodex: + def test_owns_app_server_interposer_and_tui_lifecycle(self, tmp_path, monkeypatch): + processes = [] + interposer_args = {} + stopped = [] + + class FakeProcess: + def __init__(self, argv, **kwargs): + self.argv = argv + self.kwargs = kwargs + self.terminated = False + processes.append(self) + + def wait(self, timeout=None): + return 0 if timeout is not None else 7 + + def terminate(self): + self.terminated = True + + def kill(self): + raise AssertionError("clean shutdown should not need kill") + + def send_signal(self, _signal): + raise AssertionError("test does not interrupt the TUI") + + ports = iter([41001, 41002]) + monkeypatch.setattr(v2.subprocess, "Popen", FakeProcess) + monkeypatch.setattr(v2, "get_databricks_token", lambda workspace, profile: "token") + monkeypatch.setattr( + v2, + "_generate_codex_app_server_home", + lambda state, model, render_overlay: tmp_path, + ) + monkeypatch.setattr(codex_interposer, "free_port", lambda: next(ports)) + monkeypatch.setattr(codex_interposer, "wait_healthz", lambda port, timeout: True) + + def start_interposer(*args, **kwargs): + interposer_args["args"] = args + interposer_args["kwargs"] = kwargs + return object(), lambda: stopped.append(True) + + monkeypatch.setattr(codex_interposer, "start_interposer_thread", start_interposer) + + with pytest.raises(SystemExit) as exc: + v2.launch_codex( + {"workspace": WS, "profile": "myprof"}, + ["--search"], + binary="codex", + start_model="gpt-start", + render_overlay=codex.render_overlay, + ) + + assert exc.value.code == 7 + assert processes[0].argv == [ + "codex", + "app-server", + "--listen", + "ws://127.0.0.1:41001", + ] + assert processes[0].kwargs["env"][v2.OAUTH_TOKEN_ENV_VAR] == "token" + assert processes[0].kwargs["env"][v2.CODEX_HOME_ENV_VAR] == str(tmp_path) + assert processes[1].argv == [ + "codex", + "--remote", + "ws://127.0.0.1:41002", + "--model", + "gpt-start", + "--search", + ] + assert interposer_args["args"] == ( + codex_interposer.LOOPBACK_HOST, + 41002, + "ws://127.0.0.1:41001", + v2.CODEX_TARGET_MODEL, + v2.SWITCH_AFTER_TURNS, + ) + assert "Using Unity Gateway Smart Router." in interposer_args["kwargs"]["switch_message"] + assert stopped == [True] + assert processes[0].terminated is True + + class TestInterposerSession: """The interposer's hold-then-switch + settings-injection logic (the novel behavior).""" From e71f2a7266af58cefa329b63acf01515c6956e8f Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Tue, 25 Aug 2026 20:30:10 +0000 Subject: [PATCH 09/18] Simplify Codex smart routing interposer --- src/ucode/smart_routing/codex_interposer.py | 126 +++++++------------- src/ucode/smart_routing/v2.py | 48 ++++++-- tests/test_codex_smart_routing_v2.py | 102 ++++++++-------- 3 files changed, 129 insertions(+), 147 deletions(-) diff --git a/src/ucode/smart_routing/codex_interposer.py b/src/ucode/smart_routing/codex_interposer.py index e6488560..d83c62a5 100644 --- a/src/ucode/smart_routing/codex_interposer.py +++ b/src/ucode/smart_routing/codex_interposer.py @@ -6,10 +6,9 @@ frame. This module sits between the real TUI and a real ``codex app-server``, forwarding every frame untouched except: - - ``turn/start`` (TUI->engine): after an initial hold of ``after`` turns, its - ``model`` is rewritten. ``turn/start.model`` is documented as "override the - model for this turn and subsequent turns", so the live session retargets with - history preserved. + - ``turn/start`` (TUI->engine): its ``model`` is rewritten. + ``turn/start.model`` is documented as "override the model for this turn and + subsequent turns", so the live session retargets with history preserved. - When the first switched turn starts — on that turn's ``turn/started``, before any response items stream — two things are injected (engine->TUI): a ``thread/settings/updated`` carrying the new model, so the TUI's on-screen @@ -33,11 +32,8 @@ import asyncio import contextlib import json -import os -import socket import threading import time -import urllib.request import uuid from collections.abc import Callable from pathlib import Path @@ -48,35 +44,29 @@ SETTINGS_UPDATED = "thread/settings/updated" ITEM_STARTED = "item/started" ITEM_COMPLETED = "item/completed" -LOOPBACK_HOST = "127.0.0.1" - -# Set UCODE_INTERPOSER_DEBUG=1 to dump raw engine->TUI item/turn frames (and the frames we -# inject) to the interposer log, for comparing our synthetic note against real assistant frames. -_DEBUG = os.environ.get("UCODE_INTERPOSER_DEBUG") == "1" +TURN_START = "turn/start" +TURN_STARTED = "turn/started" class _Session: - """Per-TUI-connection state: hold the first ``after`` turns, then switch model.""" + """Per-TUI-connection state for switching the model once.""" def __init__( self, target_model: str, - after: int, log: Callable[[str], None], switch_message: str | None = None, ) -> None: self.target = target_model - self.after = after self.log = log - # One-line explanation surfaced in the TUI when the switch fires; None skips it. self.switch_message = switch_message - self.turns = 0 self.thread_id: str | None = None self.settings: dict | None = None + self.switch_pending = False self.injected = False def on_tui_frame(self, raw: str) -> str: - """TUI->engine: rewrite ``turn/start.model`` once past the hold.""" + """TUI->engine: rewrite ``turn/start.model`` to the selected model.""" try: msg = json.loads(raw) except ValueError: @@ -84,18 +74,15 @@ def on_tui_frame(self, raw: str) -> str: if not isinstance(msg, dict): return raw params = msg.get("params") - if _DEBUG and isinstance(msg.get("method"), str) and msg["method"].startswith("turn/"): - self.log(f"[DEBUG->engine] {msg['method']}: {raw[:1200]}") - if msg.get("method") == "turn/start" and isinstance(params, dict): - self.turns += 1 + if msg.get("method") == TURN_START and isinstance(params, dict): if isinstance(params.get("threadId"), str): self.thread_id = params["threadId"] - if self.turns > self.after: - old = params.get("model") - if old != self.target: - params["model"] = self.target - self.log(f"[REWRITE] turn #{self.turns}: model {old!r} -> {self.target!r}") - return json.dumps(msg) + old = params.get("model") + if old != self.target: + params["model"] = self.target + self.switch_pending = not self.injected + self.log(f"[REWRITE] model {old!r} -> {self.target!r}") + return json.dumps(msg) return raw def on_engine_frame(self, raw: str) -> list[dict]: @@ -114,28 +101,24 @@ def on_engine_frame(self, raw: str) -> list[dict]: return [] if not isinstance(msg, dict): return [] - if _DEBUG: - method = msg.get("method") - if isinstance(method, str) and ( - method.startswith("item/") or method.startswith("turn/") - ): - self.log(f"[DEBUG<-engine] {method}: {raw[:1800]}") params = msg.get("params") if isinstance(msg.get("params"), dict) else {} result = msg.get("result") if isinstance(msg.get("result"), dict) else {} for src in (params, result): - tid = src.get("threadId") or (src.get("thread") or {}).get("id") + thread = src.get("thread") + tid = src.get("threadId") or (thread.get("id") if isinstance(thread, dict) else None) if isinstance(tid, str): self.thread_id = tid ts = src.get("threadSettings") if isinstance(ts, dict): self.settings = ts if ( - msg.get("method") == "turn/started" + msg.get("method") == TURN_STARTED and not self.injected - and self.turns > self.after + and self.switch_pending and self.thread_id ): self.injected = True + self.switch_pending = False settings = dict(self.settings) if isinstance(self.settings, dict) else {} settings["model"] = self.target self.log(f"[INJECT] {SETTINGS_UPDATED}: model -> {self.target!r} (flip TUI chip)") @@ -146,8 +129,7 @@ def on_engine_frame(self, raw: str) -> list[dict]: } ] if self.switch_message: - params_obj = msg.get("params") if isinstance(msg.get("params"), dict) else {} - turn = params_obj.get("turn") if isinstance(params_obj, dict) else {} + turn = params.get("turn") turn_id = turn.get("id") if isinstance(turn, dict) else None now_ms = int(time.time() * 1000) item = { @@ -183,20 +165,17 @@ def on_engine_frame(self, raw: str) -> list[dict]: }, } ) - if _DEBUG: - for frame in injected: - self.log(f"[DEBUG-inject] {json.dumps(frame)[:1800]}") return injected return [] async def _handle_tui( - tui, upstream_uri: str, target_model: str, after: int, log, switch_message: str | None = None + tui, upstream_uri: str, target_model: str, log, switch_message: str | None = None ) -> None: path = getattr(getattr(tui, "request", None), "path", "/") or "/" uri = upstream_uri.rstrip("/") + path log(f"[CONN] TUI connected (path={path}); dialing app-server {uri}") - sess = _Session(target_model, after, log, switch_message) + sess = _Session(target_model, log, switch_message) async with connect(uri, max_size=None) as upstream: async def tui_to_app(): @@ -214,11 +193,12 @@ async def app_to_tui(): a = asyncio.create_task(tui_to_app()) b = asyncio.create_task(app_to_tui()) - _done, pending = await asyncio.wait({a, b}, return_when=asyncio.FIRST_COMPLETED) + done, pending = await asyncio.wait({a, b}, return_when=asyncio.FIRST_COMPLETED) for t in pending: t.cancel() - with contextlib.suppress(asyncio.CancelledError): - await t + await asyncio.gather(*pending, return_exceptions=True) + for task in done: + task.result() log("[CONN] TUI session closed") @@ -227,36 +207,34 @@ async def _serve( port: int, upstream_uri: str, model: str, - after: int, log, switch_message: str | None = None, ): async def handler(tui): try: - await _handle_tui(tui, upstream_uri, model, after, log, switch_message) + await _handle_tui(tui, upstream_uri, model, log, switch_message) except Exception as exc: # noqa: BLE001 - one session must never kill the server log(f"[ERR] session: {exc!r}") server = await serve(handler, host, port, max_size=None) - log(f"[READY] ws://{host}:{port} -> {upstream_uri} (hold {after} turn(s), then -> {model!r})") + bound_port = server.sockets[0].getsockname()[1] + log(f"[READY] ws://{host}:{bound_port} -> {upstream_uri} (switch -> {model!r})") return server def start_interposer_thread( host: str, - port: int, upstream_uri: str, model: str, - after: int, *, switch_message: str | None = None, log_path: Path | None = None, ready_timeout: float = 10.0, -) -> tuple[threading.Thread, Callable[[], None]]: +) -> tuple[int, Callable[[], None]]: """Run the interposer's asyncio server in a daemon thread. - Returns ``(thread, stop)``; ``stop()`` shuts the server down and stops the - loop. ``switch_message``, when set, is surfaced in the TUI as an + Binds an OS-assigned loopback port and returns ``(port, stop)``. ``stop()`` + shuts the server down and joins its thread. ``switch_message``, when set, is surfaced as an ``agentMessage`` explaining why the model switched. Logs go to ``log_path`` (appended) when given — never to stdout/stderr, which the foreground TUI owns. Blocks until the server is listening (or ``ready_timeout`` elapses).""" @@ -279,9 +257,11 @@ def run() -> None: asyncio.set_event_loop(loop) try: holder["server"] = loop.run_until_complete( - _serve(host, port, upstream_uri, model, after, log, switch_message) + _serve(host, 0, upstream_uri, model, log, switch_message) ) + holder["port"] = holder["server"].sockets[0].getsockname()[1] except Exception as exc: # noqa: BLE001 - surface bind/connect failures to the log + holder["error"] = exc log(f"[ERR] failed to start interposer: {exc!r}") ready.set() loop.close() @@ -298,33 +278,17 @@ def run() -> None: thread = threading.Thread(target=run, name="codex-interposer", daemon=True) thread.start() - ready.wait(timeout=ready_timeout) + + if not ready.wait(timeout=ready_timeout): + loop.call_soon_threadsafe(loop.stop) + thread.join(timeout=ready_timeout) + raise RuntimeError("Codex interposer did not become ready in time.") + if error := holder.get("error"): + raise RuntimeError("Codex interposer failed to start.") from error def stop() -> None: with contextlib.suppress(RuntimeError): loop.call_soon_threadsafe(loop.stop) + thread.join(timeout=ready_timeout) - return thread, stop - - -def free_port() -> int: - """Grab an unused loopback TCP port (races are irrelevant for local ephemeral use).""" - sock = socket.socket() - sock.bind((LOOPBACK_HOST, 0)) - port = sock.getsockname()[1] - sock.close() - return port - - -def wait_healthz(port: int, timeout: float = 30.0) -> bool: - """Poll the app-server's ``/healthz`` until it returns 200, or timeout.""" - url = f"http://{LOOPBACK_HOST}:{port}/healthz" - end = time.time() + timeout - while time.time() < end: - try: - with urllib.request.urlopen(url, timeout=1) as resp: # noqa: S310 - fixed localhost URL - if resp.status == 200: - return True - except Exception: # noqa: BLE001 - not ready yet; keep polling - time.sleep(0.25) - return False + return holder["port"], stop diff --git a/src/ucode/smart_routing/v2.py b/src/ucode/smart_routing/v2.py index c2f55d85..3b5bbb62 100644 --- a/src/ucode/smart_routing/v2.py +++ b/src/ucode/smart_routing/v2.py @@ -10,8 +10,11 @@ import os import signal +import socket import subprocess import sys +import time +import urllib.request from collections.abc import Callable from pathlib import Path from typing import NoReturn @@ -23,10 +26,6 @@ # Single env var that enables the v2 launch path for every routing-capable agent. ENV_VAR = "ENABLE_SMART_ROUTING_V2" -# Turns to keep the session on its starting model before switching (0 = switch immediately, -# i.e. the very first prompt already routes to the target model and shows the switch note). -SWITCH_AFTER_TURNS = 0 - CODEX_TARGET_MODEL = "system.ai.glm-5-2" # TODO(lilly): replace with smart router. CODEX_APP_SERVER_HOME = APP_DIR / "codex-v2-home" CODEX_INTERPOSER_LOG = APP_DIR / "codex-v2-interposer.log" @@ -36,6 +35,9 @@ PROCESS_SHUTDOWN_TIMEOUT_SECONDS = 5 OAUTH_TOKEN_ENV_VAR = "OAUTH_TOKEN" CODEX_HOME_ENV_VAR = "CODEX_HOME" +LOOPBACK_HOST = "127.0.0.1" +HEALTH_REQUEST_TIMEOUT_SECONDS = 1 +HEALTH_POLL_INTERVAL_SECONDS = 0.25 def enabled() -> bool: @@ -44,7 +46,30 @@ def enabled() -> bool: def _loopback_websocket_url(port: int) -> str: - return f"ws://{codex_interposer.LOOPBACK_HOST}:{port}" + return f"ws://{LOOPBACK_HOST}:{port}" + + +def _free_port() -> int: + """Return an available loopback port for the app-server to bind.""" + with socket.socket() as sock: + sock.bind((LOOPBACK_HOST, 0)) + return sock.getsockname()[1] + + +def _wait_for_app_server(port: int, timeout: float) -> bool: + """Poll the app-server's health endpoint until it is ready or times out.""" + url = f"http://{LOOPBACK_HOST}:{port}/healthz" + end = time.monotonic() + timeout + while time.monotonic() < end: + try: + with urllib.request.urlopen( # noqa: S310 - fixed loopback URL + url, timeout=HEALTH_REQUEST_TIMEOUT_SECONDS + ) as response: + if response.status == 200: + return True + except Exception: # noqa: BLE001 - the app-server is not ready yet + time.sleep(HEALTH_POLL_INTERVAL_SECONDS) + return False def _switch_message(model: str, reason: str) -> str: @@ -99,10 +124,8 @@ def launch_codex( os.environ[OAUTH_TOKEN_ENV_VAR] = get_databricks_token(workspace, state.get("profile")) home = _generate_codex_app_server_home(state, start_model, render_overlay) - app_port = codex_interposer.free_port() - tui_port = codex_interposer.free_port() + app_port = _free_port() app_server_url = _loopback_websocket_url(app_port) - tui_url = _loopback_websocket_url(tui_port) app_server = subprocess.Popen( [binary, "app-server", "--listen", app_server_url], @@ -113,19 +136,18 @@ def launch_codex( ) stop_interposer = None try: - if not codex_interposer.wait_healthz(app_port, timeout=APP_SERVER_READY_TIMEOUT_SECONDS): + if not _wait_for_app_server(app_port, timeout=APP_SERVER_READY_TIMEOUT_SECONDS): raise RuntimeError( "Codex app-server did not become ready for smart routing v2; check workspace auth." ) - _thread, stop_interposer = codex_interposer.start_interposer_thread( - codex_interposer.LOOPBACK_HOST, - tui_port, + tui_port, stop_interposer = codex_interposer.start_interposer_thread( + LOOPBACK_HOST, app_server_url, CODEX_TARGET_MODEL, - SWITCH_AFTER_TURNS, switch_message=_switch_message(CODEX_TARGET_MODEL, CODEX_SWITCH_REASON), log_path=CODEX_INTERPOSER_LOG, ) + tui_url = _loopback_websocket_url(tui_port) # Keep ucode alive while the TUI runs so it can tear down the app-server and interposer. tui = subprocess.Popen([binary, "--remote", tui_url, "--model", start_model, *tool_args]) try: diff --git a/tests/test_codex_smart_routing_v2.py b/tests/test_codex_smart_routing_v2.py index 10389f58..dbcdc1e2 100644 --- a/tests/test_codex_smart_routing_v2.py +++ b/tests/test_codex_smart_routing_v2.py @@ -73,7 +73,6 @@ def kill(self): def send_signal(self, _signal): raise AssertionError("test does not interrupt the TUI") - ports = iter([41001, 41002]) monkeypatch.setattr(v2.subprocess, "Popen", FakeProcess) monkeypatch.setattr(v2, "get_databricks_token", lambda workspace, profile: "token") monkeypatch.setattr( @@ -81,13 +80,13 @@ def send_signal(self, _signal): "_generate_codex_app_server_home", lambda state, model, render_overlay: tmp_path, ) - monkeypatch.setattr(codex_interposer, "free_port", lambda: next(ports)) - monkeypatch.setattr(codex_interposer, "wait_healthz", lambda port, timeout: True) + monkeypatch.setattr(v2, "_free_port", lambda: 41001) + monkeypatch.setattr(v2, "_wait_for_app_server", lambda port, timeout: True) def start_interposer(*args, **kwargs): interposer_args["args"] = args interposer_args["kwargs"] = kwargs - return object(), lambda: stopped.append(True) + return 41002, lambda: stopped.append(True) monkeypatch.setattr(codex_interposer, "start_interposer_thread", start_interposer) @@ -118,79 +117,87 @@ def start_interposer(*args, **kwargs): "--search", ] assert interposer_args["args"] == ( - codex_interposer.LOOPBACK_HOST, - 41002, + v2.LOOPBACK_HOST, "ws://127.0.0.1:41001", v2.CODEX_TARGET_MODEL, - v2.SWITCH_AFTER_TURNS, ) assert "Using Unity Gateway Smart Router." in interposer_args["kwargs"]["switch_message"] assert stopped == [True] assert processes[0].terminated is True +def test_interposer_startup_failure_is_propagated(monkeypatch): + async def fail_to_serve(*args, **kwargs): + raise OSError("bind failed") + + monkeypatch.setattr(codex_interposer, "_serve", fail_to_serve) + + with pytest.raises(RuntimeError, match="failed to start") as exc: + codex_interposer.start_interposer_thread( + v2.LOOPBACK_HOST, + "ws://127.0.0.1:41001", + "model-x", + ) + + assert isinstance(exc.value.__cause__, OSError) + + class TestInterposerSession: """The interposer's hold-then-switch + settings-injection logic (the novel behavior).""" def _turn_start(self, model: str, thread_id: str = "t1") -> str: return json.dumps( { - "method": "turn/start", + "method": codex_interposer.TURN_START, "id": 1, "params": {"threadId": thread_id, "input": [], "model": model}, } ) - def test_holds_first_turn_then_switches(self): - sess = codex_interposer._Session("gpt-5.5", after=1, log=lambda _m: None) - # Turn 1 passes through unchanged (still on the TUI's model). - out1 = sess.on_tui_frame(self._turn_start("system.ai.gpt-5-6-luna")) - assert json.loads(out1)["params"]["model"] == "system.ai.gpt-5-6-luna" - # Turn 2 is rewritten to the target. - out2 = sess.on_tui_frame(self._turn_start("system.ai.gpt-5-6-luna")) - assert json.loads(out2)["params"]["model"] == "gpt-5.5" + def test_switches_first_turn(self): + sess = codex_interposer._Session("gpt-5.5", log=lambda _m: None) + output = sess.on_tui_frame(self._turn_start("system.ai.gpt-5-6-luna")) + assert json.loads(output)["params"]["model"] == "gpt-5.5" - def test_after_zero_switches_immediately(self): - sess = codex_interposer._Session("gpt-5.5", after=0, log=lambda _m: None) - out1 = sess.on_tui_frame(self._turn_start("luna")) - assert json.loads(out1)["params"]["model"] == "gpt-5.5" + def test_does_not_schedule_notification_when_model_is_already_selected(self): + sess = codex_interposer._Session("gpt-5.5", log=lambda _m: None) + frame = self._turn_start("gpt-5.5") + assert sess.on_tui_frame(frame) == frame + assert sess.on_engine_frame(self._turn_started("turn-1")) == [] def test_non_turn_frames_pass_through(self): - sess = codex_interposer._Session("gpt-5.5", after=1, log=lambda _m: None) + sess = codex_interposer._Session("gpt-5.5", log=lambda _m: None) frame = json.dumps({"method": "initialize", "id": 1, "params": {}}) assert sess.on_tui_frame(frame) == frame def _turn_started(self, turn_id: str, thread_id: str = "t1") -> str: return json.dumps( - {"method": "turn/started", "params": {"threadId": thread_id, "turn": {"id": turn_id}}} + { + "method": codex_interposer.TURN_STARTED, + "params": {"threadId": thread_id, "turn": {"id": turn_id}}, + } ) - def test_holds_note_until_switched_turn_starts(self): - # The note/chip-flip fire on the SWITCHED turn's turn/started (before its - # response), never on the held turn. - sess = codex_interposer._Session("gpt-5.5", after=1, log=lambda _m: None) - sess.on_tui_frame(self._turn_start("luna")) # turn 1 (held) - assert sess.on_engine_frame(self._turn_started("turn-1")) == [] # held: no inject - sess.on_tui_frame(self._turn_start("luna")) # turn 2 (switched) - injected = sess.on_engine_frame(self._turn_started("turn-2")) + def test_injects_note_when_switched_turn_starts(self): + sess = codex_interposer._Session("gpt-5.5", log=lambda _m: None) + sess.on_tui_frame(self._turn_start("luna")) + injected = sess.on_engine_frame(self._turn_started("turn-1")) settings = next(m for m in injected if m["method"] == codex_interposer.SETTINGS_UPDATED) assert settings["params"]["threadId"] == "t1" assert settings["params"]["threadSettings"]["model"] == "gpt-5.5" def test_injects_switch_note_as_agent_message_when_message_set(self): sess = codex_interposer._Session( - "gpt-5.5", after=1, log=lambda _m: None, switch_message="selected glm-5-2 because X" + "gpt-5.5", log=lambda _m: None, switch_message="selected glm-5-2 because X" ) - sess.on_tui_frame(self._turn_start("luna")) # turn 1 (held) - sess.on_engine_frame(self._turn_started("turn-1")) - sess.on_tui_frame(self._turn_start("luna")) # turn 2 (switched) - injected = sess.on_engine_frame(self._turn_started("turn-2")) + sess.on_tui_frame(self._turn_start("luna")) + injected = sess.on_engine_frame(self._turn_started("turn-1")) # The note is a full agentMessage lifecycle: item/started THEN item/completed, # both carrying the same item (a lone item/completed renders nothing in the TUI). started = next(m for m in injected if m["method"] == codex_interposer.ITEM_STARTED) completed = next(m for m in injected if m["method"] == codex_interposer.ITEM_COMPLETED) - assert started["params"]["turnId"] == "turn-2" - assert completed["params"]["turnId"] == "turn-2" + assert started["params"]["turnId"] == "turn-1" + assert completed["params"]["turnId"] == "turn-1" for frame in (started, completed): item = frame["params"]["item"] # An agentMessage renders as plain chat text, not a yellow warning banner. @@ -198,26 +205,15 @@ def test_injects_switch_note_as_agent_message_when_message_set(self): assert item["text"] == "selected glm-5-2 because X" assert started["params"]["item"]["id"] == completed["params"]["item"]["id"] - def test_after_zero_injects_on_first_turn_start(self): - sess = codex_interposer._Session( - "gpt-5.5", after=0, log=lambda _m: None, switch_message="switched" - ) - sess.on_tui_frame(self._turn_start("luna")) # turn 1 (switched immediately) - injected = sess.on_engine_frame(self._turn_started("turn-1")) - methods = [m["method"] for m in injected] - assert codex_interposer.ITEM_STARTED in methods - assert codex_interposer.ITEM_COMPLETED in methods - def test_no_note_without_message(self): - sess = codex_interposer._Session("gpt-5.5", after=1, log=lambda _m: None) + sess = codex_interposer._Session("gpt-5.5", log=lambda _m: None) sess.on_tui_frame(self._turn_start("luna")) - sess.on_tui_frame(self._turn_start("luna")) - injected = sess.on_engine_frame(self._turn_started("turn-2")) + injected = sess.on_engine_frame(self._turn_started("turn-1")) assert [m["method"] for m in injected] == [codex_interposer.SETTINGS_UPDATED] def test_injects_only_once(self): - sess = codex_interposer._Session("gpt-5.5", after=1, log=lambda _m: None) + sess = codex_interposer._Session("gpt-5.5", log=lambda _m: None) sess.on_tui_frame(self._turn_start("luna")) + assert sess.on_engine_frame(self._turn_started("turn-1")) sess.on_tui_frame(self._turn_start("luna")) - assert sess.on_engine_frame(self._turn_started("turn-2")) # switched turn: injects - assert sess.on_engine_frame(self._turn_started("turn-3")) == [] # later turn: no re-inject + assert sess.on_engine_frame(self._turn_started("turn-2")) == [] From 21be1d36d3ddc149e36a35c05254e3ea1ed83a59 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Tue, 25 Aug 2026 20:34:15 +0000 Subject: [PATCH 10/18] Trim smart routing documentation --- src/ucode/smart_routing/codex_interposer.py | 58 ++------------------- src/ucode/smart_routing/v2.py | 21 ++------ tests/test_codex_smart_routing_v2.py | 8 --- 3 files changed, 6 insertions(+), 81 deletions(-) diff --git a/src/ucode/smart_routing/codex_interposer.py b/src/ucode/smart_routing/codex_interposer.py index d83c62a5..f3ae3929 100644 --- a/src/ucode/smart_routing/codex_interposer.py +++ b/src/ucode/smart_routing/codex_interposer.py @@ -1,32 +1,3 @@ -"""WebSocket interposer for the Codex TUI's ``--remote`` transport (smart routing v2). - -Codex's remote transport (``codex --remote ws://…``) is WebSocket: a plain-JSONL -client is rejected with HTTP 400 ("Connection header did not include 'upgrade'"), -a proper upgrade returns 101, and each JSON-RPC message is one WebSocket text -frame. This module sits between the real TUI and a real ``codex app-server``, -forwarding every frame untouched except: - - - ``turn/start`` (TUI->engine): its ``model`` is rewritten. - ``turn/start.model`` is documented as "override the model for this turn and - subsequent turns", so the live session retargets with history preserved. - - When the first switched turn starts — on that turn's ``turn/started``, before - any response items stream — two things are injected (engine->TUI): a - ``thread/settings/updated`` carrying the new model, so the TUI's on-screen - model indicator follows the switch, and — when a ``switch_message`` is - configured — an ``agentMessage`` item (as an ``item/started`` + ``item/completed`` - pair) that surfaces an explanation of why the model was switched, ahead - of the model's reply. An ``agentMessage`` renders as ordinary chat text (no - warning styling); Codex's protocol has no neutral free-text notification - (``warning``, ``configWarning``, ``deprecationNotice`` all render as warnings), - so an item is the way to show an informational note. The ``item/started`` is - required: the TUI creates the message widget on ``item/started``, so a lone - ``item/completed`` has no widget to finalize and renders nothing. - -``ucode.smart_routing.v2`` runs :func:`start_interposer_thread` in a daemon -thread while it owns the app-server subprocess and the ``codex --remote`` TUI, -so the whole thing launches from the single ``ucode codex`` command. -""" - from __future__ import annotations import asyncio @@ -49,8 +20,6 @@ class _Session: - """Per-TUI-connection state for switching the model once.""" - def __init__( self, target_model: str, @@ -66,7 +35,6 @@ def __init__( self.injected = False def on_tui_frame(self, raw: str) -> str: - """TUI->engine: rewrite ``turn/start.model`` to the selected model.""" try: msg = json.loads(raw) except ValueError: @@ -86,15 +54,6 @@ def on_tui_frame(self, raw: str) -> str: return raw def on_engine_frame(self, raw: str) -> list[dict]: - """engine->TUI: capture thread id/settings; when the first switched turn - starts, return the frames to inject (empty list = none). - - On the switched turn's ``turn/started`` — before its response streams — - this yields a ``thread/settings/updated`` (flips the TUI's model chip) - and, when ``switch_message`` is set, an ``item/started`` + - ``item/completed`` pair carrying an ``agentMessage`` — plain chat text - (no warning styling) that explains why the model changed, shown ahead of - the model's reply.""" try: msg = json.loads(raw) except ValueError: @@ -140,9 +99,7 @@ def on_engine_frame(self, raw: str) -> list[dict]: "memoryCitation": None, } self.log(f"[INJECT] agentMessage note (started+completed): {self.switch_message!r}") - # The TUI creates the message widget on item/started; a lone item/completed - # has no widget to finalize and renders nothing. Send the full lifecycle with - # the text already populated (no deltas needed for a static note). + # Codex renders the message only when it receives both lifecycle events. injected.append( { "method": ITEM_STARTED, @@ -213,7 +170,7 @@ async def _serve( async def handler(tui): try: await _handle_tui(tui, upstream_uri, model, log, switch_message) - except Exception as exc: # noqa: BLE001 - one session must never kill the server + except Exception as exc: # noqa: BLE001 log(f"[ERR] session: {exc!r}") server = await serve(handler, host, port, max_size=None) @@ -231,14 +188,6 @@ def start_interposer_thread( log_path: Path | None = None, ready_timeout: float = 10.0, ) -> tuple[int, Callable[[], None]]: - """Run the interposer's asyncio server in a daemon thread. - - Binds an OS-assigned loopback port and returns ``(port, stop)``. ``stop()`` - shuts the server down and joins its thread. ``switch_message``, when set, is surfaced as an - ``agentMessage`` explaining why the model switched. Logs go to ``log_path`` (appended) when - given — never to stdout/stderr, which the foreground TUI owns. Blocks until - the server is listening (or ``ready_timeout`` elapses).""" - def log(message: str) -> None: if log_path is None: return @@ -260,7 +209,7 @@ def run() -> None: _serve(host, 0, upstream_uri, model, log, switch_message) ) holder["port"] = holder["server"].sockets[0].getsockname()[1] - except Exception as exc: # noqa: BLE001 - surface bind/connect failures to the log + except Exception as exc: # noqa: BLE001 holder["error"] = exc log(f"[ERR] failed to start interposer: {exc!r}") ready.set() @@ -268,7 +217,6 @@ def run() -> None: return ready.set() loop.run_forever() - # Stopped: close the server and drain. server = holder.get("server") if server is not None: server.close() diff --git a/src/ucode/smart_routing/v2.py b/src/ucode/smart_routing/v2.py index 3b5bbb62..b129ff3d 100644 --- a/src/ucode/smart_routing/v2.py +++ b/src/ucode/smart_routing/v2.py @@ -1,11 +1,3 @@ -"""Runtime model-switching launch path for smart routing v2. - -Smart routing v2 launches an agent's real TUI against a ucode-run app-server with a -WebSocket interposer that switches the model mid-session (see e.g. -``smart_routing.codex_interposer``). The enable flag, launch configuration, and process -lifecycle live here so agent modules only need to supply their provider overlay. -""" - from __future__ import annotations import os @@ -23,7 +15,6 @@ from ucode.databricks import get_databricks_token from ucode.smart_routing import codex_interposer -# Single env var that enables the v2 launch path for every routing-capable agent. ENV_VAR = "ENABLE_SMART_ROUTING_V2" CODEX_TARGET_MODEL = "system.ai.glm-5-2" # TODO(lilly): replace with smart router. @@ -41,7 +32,6 @@ def enabled() -> bool: - """Return whether the smart-routing-v2 launch path is enabled via the env var.""" return os.environ.get(ENV_VAR) == "1" @@ -50,24 +40,22 @@ def _loopback_websocket_url(port: int) -> str: def _free_port() -> int: - """Return an available loopback port for the app-server to bind.""" with socket.socket() as sock: sock.bind((LOOPBACK_HOST, 0)) return sock.getsockname()[1] def _wait_for_app_server(port: int, timeout: float) -> bool: - """Poll the app-server's health endpoint until it is ready or times out.""" url = f"http://{LOOPBACK_HOST}:{port}/healthz" end = time.monotonic() + timeout while time.monotonic() < end: try: - with urllib.request.urlopen( # noqa: S310 - fixed loopback URL + with urllib.request.urlopen( # noqa: S310 url, timeout=HEALTH_REQUEST_TIMEOUT_SECONDS ) as response: if response.status == 200: return True - except Exception: # noqa: BLE001 - the app-server is not ready yet + except Exception: # noqa: BLE001 time.sleep(HEALTH_POLL_INTERVAL_SECONDS) return False @@ -88,7 +76,6 @@ def _generate_codex_app_server_home( model: str, render_overlay: Callable[..., dict], ) -> Path: - """Write the isolated CODEX_HOME used by the ucode-run app-server.""" CODEX_APP_SERVER_HOME.mkdir(parents=True, exist_ok=True) config_path = CODEX_APP_SERVER_HOME / "config.toml" overlay = render_overlay( @@ -111,7 +98,6 @@ def launch_codex( start_model: str | None, render_overlay: Callable[..., dict], ) -> NoReturn: - """Launch the Codex app-server, interposer, and remote TUI as one lifecycle.""" workspace = state.get("workspace") if not workspace: raise RuntimeError( @@ -148,7 +134,6 @@ def launch_codex( log_path=CODEX_INTERPOSER_LOG, ) tui_url = _loopback_websocket_url(tui_port) - # Keep ucode alive while the TUI runs so it can tear down the app-server and interposer. tui = subprocess.Popen([binary, "--remote", tui_url, "--model", start_model, *tool_args]) try: returncode = tui.wait() @@ -161,6 +146,6 @@ def launch_codex( app_server.terminate() try: app_server.wait(timeout=PROCESS_SHUTDOWN_TIMEOUT_SECONDS) - except Exception: # noqa: BLE001 - the app-server must never linger + except Exception: # noqa: BLE001 app_server.kill() sys.exit(returncode) diff --git a/tests/test_codex_smart_routing_v2.py b/tests/test_codex_smart_routing_v2.py index dbcdc1e2..624d73da 100644 --- a/tests/test_codex_smart_routing_v2.py +++ b/tests/test_codex_smart_routing_v2.py @@ -1,5 +1,3 @@ -"""Tests for the experimental ENABLE_SMART_ROUTING_V2 Codex launch path.""" - from __future__ import annotations import json @@ -31,7 +29,6 @@ def test_writes_provider_config(self, tmp_path, monkeypatch): assert doc["model"] == "gpt-5.6-luna" provider = doc["model_providers"][codex.CODEX_MODEL_PROVIDER_NAME] assert provider["base_url"].endswith("/ai-gateway/codex/v1") - # Self-refreshing auth command is preserved (app-server rejects --profile). assert provider["auth"]["command"].endswith("ucode") assert "myprof" in provider["auth"]["args"] @@ -143,8 +140,6 @@ async def fail_to_serve(*args, **kwargs): class TestInterposerSession: - """The interposer's hold-then-switch + settings-injection logic (the novel behavior).""" - def _turn_start(self, model: str, thread_id: str = "t1") -> str: return json.dumps( { @@ -192,15 +187,12 @@ def test_injects_switch_note_as_agent_message_when_message_set(self): ) sess.on_tui_frame(self._turn_start("luna")) injected = sess.on_engine_frame(self._turn_started("turn-1")) - # The note is a full agentMessage lifecycle: item/started THEN item/completed, - # both carrying the same item (a lone item/completed renders nothing in the TUI). started = next(m for m in injected if m["method"] == codex_interposer.ITEM_STARTED) completed = next(m for m in injected if m["method"] == codex_interposer.ITEM_COMPLETED) assert started["params"]["turnId"] == "turn-1" assert completed["params"]["turnId"] == "turn-1" for frame in (started, completed): item = frame["params"]["item"] - # An agentMessage renders as plain chat text, not a yellow warning banner. assert item["type"] == "agentMessage" assert item["text"] == "selected glm-5-2 because X" assert started["params"]["item"]["id"] == completed["params"]["item"]["id"] From cb3ff5193155fa1a2a3417d00cc3aad1db57ebbb Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Tue, 25 Aug 2026 21:42:23 +0000 Subject: [PATCH 11/18] Route Claude prompts with direct model commands --- src/ucode/agents/claude.py | 196 ++++++++++- src/ucode/cli.py | 26 ++ src/ucode/smart_routing/claude_hooks.py | 19 ++ src/ucode/smart_routing/claude_pty.py | 418 ++++++++++++++++++++++++ tests/test_agent_claude.py | 75 ++++- tests/test_claude_smart_routing_v2.py | 184 +++++++++++ 6 files changed, 911 insertions(+), 7 deletions(-) create mode 100644 src/ucode/smart_routing/claude_pty.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..292a05c8 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 @@ -43,7 +48,9 @@ GATEWAY_MODEL_DISCOVERY_ENV_VAR = "ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY" CLAUDE_CONFIG_DIR = Path.home() / ".claude" CLAUDE_SETTINGS_PATH = CLAUDE_CONFIG_DIR / "ucode-settings.json" +CLAUDE_USER_SETTINGS_PATH = CLAUDE_CONFIG_DIR / "settings.json" CLAUDE_BACKUP_PATH = APP_DIR / "claude-ucode-settings.backup.json" +CLAUDE_MODEL_SNAPSHOT_PATH = APP_DIR / "claude-default-model.snapshot.json" SPEC: ToolSpec = { "binary": "claude", @@ -62,6 +69,10 @@ # marked managed so they're tracked/reverted with the rest of ucode's config. CLAUDE_ROUTING_HOOK_EVENTS = ("PreToolUse", "SessionStart", "SubagentStart") +# Smart-routing-v2 currently uses a fixed route while its dynamic router is wired up. +SMART_ROUTING_V2_MODEL = "system.ai.claude-sonnet-4-6[1m]" +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"]) @@ -875,7 +886,74 @@ def _merge_claude_settings(base: dict, overlay: dict) -> dict: return merged -def _build_claude_argv(binary: str, tool_args: list[str], relayed: bool = False) -> list[str]: +def _snapshot_user_model_setting() -> dict: + """Capture only Claude's user-level ``model`` setting.""" + settings = read_json_safe(CLAUDE_USER_SETTINGS_PATH) + return {"present": "model" in settings, "value": settings.get("model")} + + +def _save_user_model_snapshot(snapshot: dict, snapshot_path: Path | None = None) -> None: + """Journal the pre-switch model so the next launch can recover after a crash.""" + snapshot_path = snapshot_path or CLAUDE_MODEL_SNAPSHOT_PATH + write_json_file(snapshot_path, snapshot) + + +def _restore_user_model_snapshot( + snapshot_path: Path | None = None, + snapshot: dict | None = None, +) -> bool: + """Restore only the journaled ``model`` field, preserving every sibling setting.""" + snapshot_path = snapshot_path or CLAUDE_MODEL_SNAPSHOT_PATH + if snapshot is None and not snapshot_path.exists(): + return False + snapshot = snapshot if snapshot is not None else read_json_safe(snapshot_path) + present = snapshot.get("present") + if not isinstance(present, bool): + raise RuntimeError(f"Claude model recovery snapshot is invalid: {snapshot_path}") + settings = read_json_safe(CLAUDE_USER_SETTINGS_PATH) + if present: + settings["model"] = snapshot.get("value") + else: + settings.pop("model", None) + write_json_file(CLAUDE_USER_SETTINGS_PATH, settings) + snapshot_path.unlink(missing_ok=True) + return True + + +def _recover_user_model_snapshots() -> None: + """Repair defaults left by interrupted PTY launches.""" + candidates = {CLAUDE_MODEL_SNAPSHOT_PATH} + candidates.update(APP_DIR.glob("claude-default-model.*.snapshot.json")) + for path in sorted(candidates): + _restore_user_model_snapshot(path) + + +def _original_launch_model(snapshot: dict, state: dict) -> str | None: + override = state.get("_claude_launch_model") + if isinstance(override, str) and override.strip(): + return override.strip() + value = snapshot.get("value") if snapshot.get("present") is True else None + if isinstance(value, str) and value.strip(): + return value.strip() + return default_model(state) + + +def _has_explicit_model_arg(tool_args: list[str]) -> bool: + return any(arg in {"--model", "-m"} or arg.startswith("--model=") for arg in tool_args) + + +def _launch_model_args(tool_args: list[str], launch_model: str | None) -> list[str]: + if not launch_model or _has_explicit_model_arg(tool_args): + return [] + return ["--model", launch_model] + + +def _build_claude_argv( + binary: str, + tool_args: list[str], + relayed: bool = False, + launch_model: str | None = None, +) -> list[str]: """Build the ``claude`` argv, composing any caller ``--settings`` with ucode's managed settings. @@ -897,11 +975,19 @@ def _build_claude_argv(binary: str, tool_args: list[str], relayed: bool = False) filter through and shadow the subscription OAuth. """ source_args = ["--setting-sources", _RELAYED_SETTING_SOURCES] if relayed else [] + model_args = _launch_model_args(tool_args, launch_model) caller_values, remaining = _extract_caller_settings(tool_args) if not caller_values: # No caller --settings: hand Claude ucode's settings file directly (the # common path; behavior unchanged). - return [binary, *source_args, "--settings", str(CLAUDE_SETTINGS_PATH), *tool_args] + return [ + binary, + *source_args, + "--settings", + str(CLAUDE_SETTINGS_PATH), + *model_args, + *tool_args, + ] caller_settings: dict = {} for value in caller_values: caller_settings = _merge_claude_settings(caller_settings, _load_caller_settings(value)) @@ -913,6 +999,7 @@ def _build_claude_argv(binary: str, tool_args: list[str], relayed: bool = False) *source_args, "--settings", json.dumps(merged, separators=(",", ":")), + *model_args, *remaining, ] @@ -962,7 +1049,9 @@ def _rewrite_relayed_port(state: dict, port: int) -> None: write_json_file(CLAUDE_SETTINGS_PATH, settings) -def _launch_relayed(state: dict, binary: str, tool_args: list[str]) -> None: +def _launch_relayed( + state: dict, binary: str, tool_args: list[str], launch_model: str | None = None +) -> None: """Relayed launch: sign into the Claude subscription, start the loopback refresh proxy, then run Claude Code alongside it (the proxy must outlive the exec, so we spawn-and-wait rather than replacing the process).""" @@ -1004,7 +1093,9 @@ def _launch_relayed(state: dict, binary: str, tool_args: list[str]) -> None: server_thread = threading.Thread(target=server.serve_forever, daemon=True) server_thread.start() - proc = subprocess.Popen(_build_claude_argv(binary, tool_args, relayed=True)) + proc = subprocess.Popen( + _build_claude_argv(binary, tool_args, relayed=True, launch_model=launch_model) + ) try: returncode = proc.wait() except KeyboardInterrupt: @@ -1017,15 +1108,108 @@ def _launch_relayed(state: dict, binary: str, tool_args: list[str]) -> None: raise SystemExit(returncode) +def _v2_router(state: dict): + """Return the first-prompt router; currently fixed while routing is prototyped.""" + + def route(_prompt: str) -> str: + return SMART_ROUTING_V2_MODEL + + return route + + +def _launch_smart_routing_v2( + state: dict, + tool_args: list[str], + *, + model_snapshot: dict, + launch_model: str | None, +) -> None: + """Launch Claude in the first-prompt routing PTY wrapper.""" + 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" + model_snapshot_path = APP_DIR / f"claude-default-model.{run_id}.snapshot.json" + + 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) + model_args = _launch_model_args(remaining, launch_model) + argv = [binary, "--settings", str(settings_path), *model_args, *remaining] + + _save_user_model_snapshot(model_snapshot, model_snapshot_path) + restored = False + + def restore_default_model() -> bool: + nonlocal restored + if restored: + return False + result = _restore_user_model_snapshot(model_snapshot_path, model_snapshot) + restored = True + return result + + print_note( + "Smart routing v2: the first submitted prompt will select Claude Code's " + f"model (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, + restore_default_model=restore_default_model, + log_path=SMART_ROUTING_V2_CLAUDE_LOG, + ) + finally: + # The PTY restores immediately after a confirmed switch. This fallback + # covers startup failures, timeouts, signals, and normal child exit. + restore_default_model() + 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") + # Recover a prior switch interrupted before its surgical settings restore. + _recover_user_model_snapshots() + model_snapshot = _snapshot_user_model_setting() + launch_model = _original_launch_model(model_snapshot, state) if state.get("claude_relayed"): - _launch_relayed(state, binary, tool_args) + _launch_relayed(state, binary, tool_args, launch_model) + return + if smart_routing_v2.enabled() and workspace and os.name != "nt": + _launch_smart_routing_v2( + state, + tool_args, + model_snapshot=model_snapshot, + launch_model=launch_model, + ) return if workspace: os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile")) - exec_or_spawn(_build_claude_argv(binary, tool_args)) + exec_or_spawn(_build_claude_argv(binary, tool_args, launch_model=launch_model)) def validate_cmd(binary: str) -> list[str]: diff --git a/src/ucode/cli.py b/src/ucode/cli.py index ca71dd6a..1fdb53b9 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 @@ -1957,6 +1976,13 @@ def _launch_tool( if managed is not None and not is_dry_run(): _register_managed_mcp_servers(managed, tool, state) _apply_managed_skills(managed, tool, state) + if tool == "claude": + # Transient launch precedence for claude.py's universal --model flag. + # An explicit choice wins, followed by a routed/managed root pick; + # neither value is persisted into workspace state. + launch_model = model or route_root_model + if launch_model: + state["_claude_launch_model"] = launch_model print_success(f"Starting {TOOL_SPECS[tool]['display']}") launch_agent(tool, state, ctx.args) except RuntimeError as exc: diff --git a/src/ucode/smart_routing/claude_hooks.py b/src/ucode/smart_routing/claude_hooks.py index 773b7707..6bc05348 100644 --- a/src/ucode/smart_routing/claude_hooks.py +++ b/src/ucode/smart_routing/claude_hooks.py @@ -15,6 +15,8 @@ from ucode.smart_routing import hooks ROUTING_HOOK_COMMAND_MARKER = "claude-router-hook" +FIRST_PROMPT_HOOK_MARKER = "claude-router-hook route-first-prompt" +FIRST_PROMPT_SOCKET_ENV = "UCODE_CLAUDE_V2_SOCKET" def sync_smart_routing_hooks(doc: dict, state: dict, *, enabled: bool) -> None: @@ -28,6 +30,23 @@ def remove_smart_routing_hooks(doc: dict) -> bool: return hooks.remove_managed_hooks(doc, ROUTING_HOOK_COMMAND_MARKER) +def sync_first_prompt_hook(doc: dict, executable: str) -> None: + """Add the first-prompt hook to a per-launch settings document.""" + groups = { + "UserPromptSubmit": [ + { + "hooks": [ + _routing_command_hook( + [executable, ROUTING_HOOK_COMMAND_MARKER, "route-first-prompt"], + status="Selecting a model with Smart Routing", + ) + ] + } + ] + } + hooks.sync_managed_hooks(doc, FIRST_PROMPT_HOOK_MARKER, groups) + + def _routing_hook_groups(state: dict) -> dict[str, list[dict]]: route_argv = _routing_hook_argv(state, "route-subagent") session_argv = _routing_hook_argv(state, "session-start") diff --git a/src/ucode/smart_routing/claude_pty.py b/src/ucode/smart_routing/claude_pty.py new file mode 100644 index 00000000..0f3585db --- /dev/null +++ b/src/ucode/smart_routing/claude_pty.py @@ -0,0 +1,418 @@ +"""PTY wrapper that routes Claude Code's first prompt with ``/model ``. + +Claude Code has no runtime control protocol, so a per-launch UserPromptSubmit +hook captures and blocks the first prompt. While the TUI is idle, this wrapper +types a direct model command, restores the user's persisted default model, and +replays the prompt. Terminal traffic is otherwise forwarded unchanged. +""" + +from __future__ import annotations + +import contextlib +import fcntl +import json +import os +import pty +import re +import select +import signal +import socket +import struct +import termios +import threading +import time +import tty +from collections.abc import Callable +from pathlib import Path + +MAX_MODEL_NAME_LEN = 200 +CONFIRM_TIMEOUT_S = 3.0 +SWITCH_TIMEOUT_S = 6.0 +READY_QUIET_S = 0.75 +SELECT_TIMEOUT_S = 0.2 +_MODEL_NAME_RE = re.compile(r"^[A-Za-z0-9._:/\-\[\]]+$") +_ANSI_RE = re.compile(rb"\x1b\[[0-9;?]*[ -/]*[@-~]|\x1b\][^\x07]*(?:\x07|\x1b\\)|\x1b[@-Z\\-_]") + + +def strip_ansi(data: bytes) -> str: + return _ANSI_RE.sub(b"", data).decode("utf-8", "replace") + + +def _squash(text: str) -> str: + return "".join(text.split()) + + +def _match_text(data: bytes) -> str: + return _squash(strip_ansi(data)) + + +def valid_model_name(name: object) -> bool: + """Return whether *name* is safe to type as one slash-command argument.""" + return ( + isinstance(name, str) + and 0 < len(name) <= MAX_MODEL_NAME_LEN + and bool(_MODEL_NAME_RE.fullmatch(name)) + ) + + +class ConfirmationState: + """Detect and accept Claude's optional cache-cost confirmation dialog.""" + + PROMPT_MARKERS = ("Switch model?", "Yes, switch to", "No, go back") + + def __init__(self, window: int = 4096) -> None: + self._buf = "" + self._armed_until = 0.0 + self._window = window + + def arm(self, deadline: float) -> None: + self._armed_until = deadline + self._buf = "" + + def clear(self) -> None: + self._armed_until = 0.0 + self._buf = "" + + def observe(self, chunk: bytes, now: float) -> bytes | None: + if self._armed_until == 0.0: + return None + if now > self._armed_until: + self.clear() + return None + self._buf = (self._buf + _match_text(chunk))[-self._window :] + if all(_squash(marker) in self._buf for marker in self.PROMPT_MARKERS): + self.clear() + return b"\r" + return None + + +class OutputMarkerDetector: + """Latching ANSI-insensitive substring detector.""" + + def __init__(self, markers: tuple[str, ...], window: int = 4096) -> None: + self._markers = markers + self._buf = "" + self._window = window + self.triggered = False + + def observe(self, chunk: bytes) -> bool: + if self.triggered: + return True + self._buf = (self._buf + _match_text(chunk))[-self._window :] + if any(_squash(marker) in self._buf for marker in self._markers): + self.triggered = True + return self.triggered + + +def inject_model_switch(master_fd: int, model: str) -> None: + """Type Claude Code's direct, persistent model command.""" + if not valid_model_name(model): + raise ValueError(f"Unsafe Claude model name: {model!r}") + os.write(master_fd, f"/model {model}\r".encode()) + + +def inject_prompt(master_fd: int, prompt: str, *, submit: bool = True) -> None: + """Replay a captured prompt as one bracketed paste.""" + clean = prompt.replace("\r\n", "\n").replace("\r", "\n") + clean = clean.replace("\x00", "").replace("\x1b", "") + suffix = b"\r" if submit else b"" + os.write(master_fd, b"\x1b[200~" + clean.encode() + b"\x1b[201~" + suffix) + + +def inject_note(out_fd: int, message: str) -> None: + os.write(out_fd, ("\r\n\x1b[36m" + message + "\x1b[0m\r\n").encode()) + + +def request_first_prompt_route(path: Path, payload: dict, *, timeout: float = 5.0) -> dict | None: + """Send a UserPromptSubmit payload to the owning PTY wrapper.""" + prompt = payload.get("prompt") + if not isinstance(prompt, str) or not prompt.strip(): + return None + request = { + "method": "route_first_prompt", + "prompt": prompt, + "session_id": payload.get("session_id"), + } + try: + client = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + client.settimeout(timeout) + with client: + client.connect(str(path)) + client.sendall((json.dumps(request) + "\n").encode()) + with client.makefile("rb") as stream: + raw = stream.readline() + response = json.loads(raw) if raw else None + except (OSError, ValueError): + return None + return response if isinstance(response, dict) else None + + +def first_prompt_hook_output(response: dict | None) -> dict | None: + """Translate the wrapper response into Claude hook output.""" + if not isinstance(response, dict) or response.get("action") != "block": + return None + model = response.get("model") + if not valid_model_name(model): + return None + return { + "decision": "block", + "reason": ( + 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 _message: None, +) -> threading.Thread: + """Serve the hook protocol, blocking exactly one non-command prompt.""" + + def serve() -> None: + claimed = False + try: + path.unlink(missing_ok=True) + server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + server.bind(str(path)) + os.chmod(path, 0o600) + server.listen(4) + server.settimeout(0.5) + except OSError as exc: + log(f"[ERR] first-prompt socket bind failed: {exc!r}") + return + log(f"[READY] first-prompt socket {path}") + try: + while not stop.is_set(): + try: + conn, _ = server.accept() + except TimeoutError: + continue + except OSError: + break + with conn, conn.makefile("rwb") as stream: + response: dict = {"action": "allow"} + blocked: tuple[str, str] | None = None + try: + request = json.loads(stream.readline()) + prompt = request.get("prompt") if isinstance(request, dict) else None + is_route = ( + isinstance(request, dict) + and request.get("method") == "route_first_prompt" + ) + is_command = isinstance(prompt, str) and prompt.lstrip().startswith("/") + if ( + is_route + and not claimed + and isinstance(prompt, str) + and prompt.strip() + and not is_command + ): + model = route_prompt(prompt) + if valid_model_name(model): + claimed = True + response = {"action": "block", "model": model} + blocked = (prompt, model) + except Exception as exc: # noqa: BLE001 - hooks must fail open + log(f"[ERR] first-prompt request: {exc!r}") + stream.write((json.dumps(response) + "\n").encode()) + stream.flush() + if blocked is not None: + on_blocked_prompt(*blocked) + finally: + server.close() + + thread = threading.Thread(target=serve, name="claude-first-prompt", daemon=True) + thread.start() + return thread + + +class TerminalModeGuard: + """Put stdin in raw mode for the PTY session and restore it on exit.""" + + def __init__(self, fd: int = 0) -> None: + self.fd = fd + self._saved: list | None = None + + def __enter__(self) -> TerminalModeGuard: + if os.isatty(self.fd): + self._saved = termios.tcgetattr(self.fd) + tty.setraw(self.fd) + return self + + def __exit__(self, *_exc: object) -> None: + if self._saved is not None: + termios.tcsetattr(self.fd, termios.TCSADRAIN, self._saved) + self._saved = None + + +def sync_winsize(master_fd: int, stdin_fd: int = 0) -> None: + if not os.isatty(stdin_fd): + return + try: + packed = fcntl.ioctl(stdin_fd, termios.TIOCGWINSZ, struct.pack("HHHH", 0, 0, 0, 0)) + fcntl.ioctl(master_fd, termios.TIOCSWINSZ, packed) + except OSError: + pass + + +def run_claude_pty( + argv: list[str], + *, + route_prompt: Callable[[str], str], + switch_message: str, + socket_path: Path, + restore_default_model: Callable[[], object], + log_path: Path | None = None, +) -> int: + """Run Claude in a PTY, switch its model, restore the default, and replay.""" + + def log(message: str) -> None: + if log_path is None: + return + try: + with open(log_path, "a", encoding="utf-8") as handle: + handle.write(f"{time.strftime('%H:%M:%S')} {message}\n") + except OSError: + pass + + debug = os.environ.get("UCODE_CLAUDE_PTY_DEBUG") == "1" + gate_read, gate_write = os.pipe() + pid, master_fd = pty.fork() + if pid == 0: + os.close(gate_write) + try: + os.read(gate_read, 1) + finally: + os.close(gate_read) + os.execvp(argv[0], argv) + os._exit(127) + os.close(gate_read) + + confirm = ConfirmationState() + stop = threading.Event() + pending_lock = threading.Lock() + pending: dict[str, tuple[str, str] | None] = {"value": None} + + def on_blocked_prompt(prompt: str, model: str) -> None: + with pending_lock: + pending["value"] = (prompt, model) + log(f"[ROUTE] first prompt -> {model!r}") + + server_thread = serve_first_prompt_socket( + socket_path, route_prompt, on_blocked_prompt, stop, log=log + ) + socket_deadline = time.monotonic() + 2.0 + while ( + not socket_path.exists() and server_thread.is_alive() and time.monotonic() < socket_deadline + ): + time.sleep(0.01) + if not socket_path.exists(): + log("[ERR] first-prompt socket was not ready before Claude launch") + os.write(gate_write, b"1") + os.close(gate_write) + + previous_winch = signal.getsignal(signal.SIGWINCH) + + def on_winch(_signum: int, _frame: object) -> None: + sync_winsize(master_fd) + + try: + with TerminalModeGuard(0): + signal.signal(signal.SIGWINCH, on_winch) + sync_winsize(master_fd) + stdin_open = True + last_output = 0.0 + phase = "waiting_prompt" + routed_prompt = "" + routed_model = "" + switch_started = 0.0 + switch_complete: OutputMarkerDetector | None = None + while True: + readable = [master_fd, 0] if stdin_open else [master_fd] + try: + ready_fds, _, _ = select.select(readable, [], [], SELECT_TIMEOUT_S) + except InterruptedError: + continue + + if 0 in ready_fds: + try: + data = os.read(0, 4096) + except OSError: + data = b"" + if not data: + stdin_open = False + else: + os.write(master_fd, data) + + if master_fd in ready_fds: + try: + chunk = os.read(master_fd, 8192) + except OSError: + chunk = b"" + if not chunk: + break + os.write(1, chunk) + last_output = time.monotonic() + if debug: + log(f"[OUT] {strip_ansi(chunk)[:400]!r}") + keystroke = confirm.observe(chunk, last_output) + if keystroke is not None: + os.write(master_fd, keystroke) + if phase == "switching" and switch_complete is not None: + switch_complete.observe(chunk) + + if phase == "waiting_prompt": + with pending_lock: + captured = pending["value"] + if captured is not None: + routed_prompt, routed_model = captured + phase = "waiting_to_switch" + + now = time.monotonic() + idle = last_output > 0.0 and now - last_output >= READY_QUIET_S + if phase == "waiting_to_switch" and idle: + inject_note(1, switch_message) + inject_model_switch(master_fd, routed_model) + confirm.arm(now + CONFIRM_TIMEOUT_S) + switch_complete = OutputMarkerDetector(("Set model to", "Model set to")) + switch_started = now + phase = "switching" + log(f"[SWITCH] /model {routed_model}") + elif ( + phase == "switching" + and switch_complete is not None + and switch_complete.triggered + ): + restore_default_model() + inject_prompt(master_fd, routed_prompt) + phase = "done" + log("[RESTORE] default model restored; first prompt replayed") + elif phase == "switching" and now - switch_started >= SWITCH_TIMEOUT_S: + os.write(master_fd, b"\x1b") + restore_default_model() + inject_note( + 1, + "Smart Routing could not confirm the model switch. " + "Your prompt was restored but not submitted.", + ) + inject_prompt(master_fd, routed_prompt, submit=False) + phase = "failed" + log(f"[ERR] direct model switch timed out for {routed_model!r}") + finally: + signal.signal(signal.SIGWINCH, previous_winch) + stop.set() + with contextlib.suppress(OSError): + os.close(master_fd) + socket_path.unlink(missing_ok=True) + + _waited_pid, status = os.waitpid(pid, 0) + if os.WIFEXITED(status): + return os.WEXITSTATUS(status) + if os.WIFSIGNALED(status): + return 128 + os.WTERMSIG(status) + return 1 diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 37b74012..c2664d69 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -665,6 +665,11 @@ def fake_execvp(binary: str, args: list[str]) -> None: monkeypatch.setattr( claude, "get_databricks_token", lambda workspace, profile=None: "fresh-token" ) + monkeypatch.setattr( + claude, + "_snapshot_user_model_setting", + lambda: {"present": True, "value": "opus"}, + ) monkeypatch.setattr(os, "execvp", fake_execvp) try: @@ -676,7 +681,14 @@ def fake_execvp(binary: str, args: list[str]) -> None: assert exec_calls == [ ( "claude", - ["claude", "--settings", str(claude.CLAUDE_SETTINGS_PATH), "--debug"], + [ + "claude", + "--settings", + str(claude.CLAUDE_SETTINGS_PATH), + "--model", + "opus", + "--debug", + ], ) ] @@ -773,6 +785,26 @@ def test_non_relayed_does_not_set_setting_sources(self, monkeypatch): argv = claude._build_claude_argv("claude", ["-p", "hi"], relayed=False) assert "--setting-sources" not in argv + def test_saved_default_is_passed_as_launch_model(self): + argv = claude._build_claude_argv("claude", ["-p", "hi"], launch_model="opus") + assert argv == [ + "claude", + "--settings", + str(claude.CLAUDE_SETTINGS_PATH), + "--model", + "opus", + "-p", + "hi", + ] + + @pytest.mark.parametrize( + "explicit", [["--model", "sonnet"], ["--model=sonnet"], ["-m", "sonnet"]] + ) + def test_explicit_model_overrides_saved_default(self, explicit): + argv = claude._build_claude_argv("claude", explicit, launch_model="opus") + assert "opus" not in argv + assert argv[-len(explicit) :] == explicit + def test_relayed_excludes_user_scope_via_setting_sources(self, monkeypatch): # Relayed must drop the user scope so a stale ~/.claude/settings.json # apiKeyHelper can't merge through and shadow the subscription OAuth. @@ -888,6 +920,47 @@ def test_malformed_file_json_raises(self, tmp_path, monkeypatch): claude._build_claude_argv("claude", ["--settings", str(bad_file)]) +class TestClaudeDefaultModelRecovery: + def _paths(self, monkeypatch, tmp_path): + user = tmp_path / "settings.json" + snapshot = tmp_path / "model-snapshot.json" + monkeypatch.setattr(claude, "CLAUDE_USER_SETTINGS_PATH", user) + monkeypatch.setattr(claude, "CLAUDE_MODEL_SNAPSHOT_PATH", snapshot) + return user, snapshot + + def test_restore_changes_only_model_field(self, monkeypatch, tmp_path): + user, _snapshot = self._paths(monkeypatch, tmp_path) + user.write_text(json.dumps({"model": "opus", "theme": "dark"})) + original = claude._snapshot_user_model_setting() + claude._save_user_model_snapshot(original) + + user.write_text(json.dumps({"model": "routed", "theme": "light", "new": True})) + assert claude._restore_user_model_snapshot() is True + assert json.loads(user.read_text()) == {"model": "opus", "theme": "light", "new": True} + assert claude._restore_user_model_snapshot() is False + + def test_restore_removes_model_when_original_was_absent(self, monkeypatch, tmp_path): + user, _snapshot = self._paths(monkeypatch, tmp_path) + user.write_text(json.dumps({"theme": "dark"})) + claude._save_user_model_snapshot(claude._snapshot_user_model_setting()) + user.write_text(json.dumps({"model": "routed", "theme": "light"})) + + claude._restore_user_model_snapshot() + assert json.loads(user.read_text()) == {"theme": "light"} + + def test_missing_user_default_falls_back_to_ucode_model(self): + state = {"claude_models": {"opus": "system.ai.claude-opus-4-8"}} + assert ( + claude._original_launch_model({"present": False, "value": None}, state) + == "system.ai.claude-opus-4-8" + ) + + def test_transient_launch_override_wins_over_saved_default(self): + state = {"_claude_launch_model": "system.ai.claude-sonnet-5"} + snapshot = {"present": True, "value": "opus"} + assert claude._original_launch_model(snapshot, state) == "system.ai.claude-sonnet-5" + + class TestClaudeSmartRouting: def _capture_write(self, monkeypatch, existing, written): monkeypatch.setattr(claude, "backup_existing_file", lambda *a, **kw: True) diff --git a/tests/test_claude_smart_routing_v2.py b/tests/test_claude_smart_routing_v2.py new file mode 100644 index 00000000..474344e2 --- /dev/null +++ b/tests/test_claude_smart_routing_v2.py @@ -0,0 +1,184 @@ +"""Tests for Claude's experimental first-prompt PTY routing path.""" + +from __future__ import annotations + +import json +import os +import sys +import threading +import time +from pathlib import Path + +import pytest + +from ucode.agents import claude +from ucode.smart_routing import claude_hooks, claude_pty + + +class TestDirectModelCommand: + @pytest.mark.parametrize( + "name", + ["system.ai.claude-opus-4-8[1m]", "databricks-claude-sonnet-5", "opus"], + ) + def test_accepts_model_names(self, name): + assert claude_pty.valid_model_name(name) + + @pytest.mark.parametrize("name", ["", "a b", "a\nb", "x" * 201, None]) + def test_rejects_unsafe_model_names(self, name): + assert not claude_pty.valid_model_name(name) + + def test_types_direct_model_command(self): + read_fd, write_fd = os.pipe() + try: + claude_pty.inject_model_switch(write_fd, "system.ai.claude-sonnet-5") + assert os.read(read_fd, 200) == b"/model system.ai.claude-sonnet-5\r" + finally: + os.close(read_fd) + os.close(write_fd) + + +class TestFirstPromptHook: + def test_blocks_once_then_allows_replay(self, tmp_path): + socket_path = tmp_path / "first.sock" + blocked: list[tuple[str, str]] = [] + stop = threading.Event() + claude_pty.serve_first_prompt_socket( + socket_path, + lambda _prompt: "sonnet", + lambda prompt, model: blocked.append((prompt, model)), + stop, + ) + try: + deadline = time.monotonic() + 5 + while not socket_path.exists() and time.monotonic() < deadline: + time.sleep(0.01) + first = claude_pty.request_first_prompt_route( + socket_path, {"session_id": "s1", "prompt": "fix the parser"} + ) + replay = claude_pty.request_first_prompt_route( + socket_path, {"session_id": "s1", "prompt": "fix the parser"} + ) + assert first == {"action": "block", "model": "sonnet"} + assert replay == {"action": "allow"} + assert blocked == [("fix the parser", "sonnet")] + finally: + stop.set() + + def test_first_prompt_hook_is_per_launch(self): + settings = {"hooks": {"PreToolUse": [{"hooks": [{"command": "user-policy"}]}]}} + claude_hooks.sync_first_prompt_hook(settings, "/bin/ucode") + claude_hooks.sync_first_prompt_hook(settings, "/bin/ucode") + command = settings["hooks"]["UserPromptSubmit"][0]["hooks"][0]["command"] + assert command == "/bin/ucode claude-router-hook route-first-prompt" + assert len(settings["hooks"]["UserPromptSubmit"]) == 1 + assert "user-policy" in str(settings["hooks"]["PreToolUse"]) + + +class TestV2Launch: + def test_saves_default_passes_model_and_restores_only_model(self, tmp_path, monkeypatch): + ucode_settings = tmp_path / "ucode-settings.json" + user_settings = tmp_path / "settings.json" + snapshot_path = tmp_path / "model-snapshot.json" + ucode_settings.write_text(json.dumps({"env": {"ANTHROPIC_BASE_URL": "https://gw"}})) + user_settings.write_text(json.dumps({"model": "opus", "theme": "dark"})) + monkeypatch.setattr(claude, "APP_DIR", tmp_path) + monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", ucode_settings) + monkeypatch.setattr(claude, "CLAUDE_USER_SETTINGS_PATH", user_settings) + monkeypatch.setattr(claude, "CLAUDE_MODEL_SNAPSHOT_PATH", snapshot_path) + monkeypatch.setattr(claude, "SMART_ROUTING_V2_CLAUDE_LOG", tmp_path / "v2.log") + monkeypatch.setattr(claude, "get_databricks_token", lambda *_args, **_kwargs: "token") + monkeypatch.setattr(claude, "build_auth_token_argv", lambda *_args, **_kwargs: ["ucode"]) + captured: dict = {} + + def fake_run(argv, **kwargs): + captured["argv"] = argv + generated = Path(argv[argv.index("--settings") + 1]) + captured["settings"] = json.loads(generated.read_text()) + # Simulate `/model` changing the user file, plus an unrelated concurrent edit. + user_settings.write_text(json.dumps({"model": "routed", "theme": "light", "new": True})) + kwargs["restore_default_model"]() + return 0 + + monkeypatch.setattr(claude_pty, "run_claude_pty", fake_run) + snapshot = claude._snapshot_user_model_setting() + with pytest.raises(SystemExit) as exc: + claude._launch_smart_routing_v2( + {"workspace": "https://example.com"}, + ["--debug"], + model_snapshot=snapshot, + launch_model="opus", + ) + + assert exc.value.code == 0 + assert captured["argv"][-3:] == ["--model", "opus", "--debug"] + assert claude_hooks.FIRST_PROMPT_SOCKET_ENV in captured["settings"]["env"] + assert "modelPicker" not in captured["settings"] + assert json.loads(user_settings.read_text()) == { + "model": "opus", + "theme": "light", + "new": True, + } + assert not snapshot_path.exists() + assert not list(tmp_path.glob("claude-default-model.*.snapshot.json")) + + +class TestPtyFlow: + def test_direct_switch_restore_and_replay(self, tmp_path): + fake_claude = tmp_path / "fake_claude.py" + capture = tmp_path / "capture.json" + socket_path = tmp_path / "first.sock" + fake_claude.write_text( + """ +import json +import os +import socket +import sys +import tty +from pathlib import Path + +socket_path = sys.argv[1] +capture_path = Path(sys.argv[2]) +client = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) +client.connect(socket_path) +client.sendall((json.dumps({ + "method": "route_first_prompt", + "prompt": "fix\\nthe parser", + "session_id": "s1", +}) + "\\n").encode()) +response = client.makefile("rb").readline() +client.close() +assert json.loads(response)["action"] == "block" +print("Smart Routing blocked the prompt", flush=True) +tty.setraw(0) + +def read_until(suffix): + data = b"" + while not data.endswith(suffix): + data += os.read(0, 1) + return data + +model_command = read_until(b"\\r") +print("Set model to system.ai.claude-sonnet-5", flush=True) +replayed = read_until(b"\\x1b[201~\\r") +capture_path.write_text(json.dumps({ + "command": model_command.decode(), + "replayed": replayed.decode(), +})) +""".lstrip() + ) + restored: list[bool] = [] + + 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, + restore_default_model=lambda: restored.append(True), + ) + + assert result == 0 + assert restored == [True] + assert json.loads(capture.read_text()) == { + "command": "/model system.ai.claude-sonnet-5\r", + "replayed": "\x1b[200~fix\nthe parser\x1b[201~\r", + } From fcf6f639d293f626c88529c4a76a2563150f68f2 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Tue, 25 Aug 2026 22:01:04 +0000 Subject: [PATCH 12/18] Restore Claude default after process exit --- src/ucode/agents/claude.py | 6 +++--- src/ucode/smart_routing/claude_pty.py | 7 ++----- tests/test_claude_smart_routing_v2.py | 5 ----- 3 files changed, 5 insertions(+), 13 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 292a05c8..fb20dd89 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -1177,12 +1177,12 @@ def restore_default_model() -> bool: "low complexity, unclear intent, and no code reference." ), socket_path=socket_path, - restore_default_model=restore_default_model, log_path=SMART_ROUTING_V2_CLAUDE_LOG, ) finally: - # The PTY restores immediately after a confirmed switch. This fallback - # covers startup failures, timeouts, signals, and normal child exit. + # Restore only after Claude exits. Claude persists `/model` asynchronously; + # restoring as soon as its success message renders races that delayed write. + # The journal repairs hard-killed and concurrent launches before they start. restore_default_model() settings_path.unlink(missing_ok=True) socket_path.unlink(missing_ok=True) diff --git a/src/ucode/smart_routing/claude_pty.py b/src/ucode/smart_routing/claude_pty.py index 0f3585db..43e90fdc 100644 --- a/src/ucode/smart_routing/claude_pty.py +++ b/src/ucode/smart_routing/claude_pty.py @@ -266,10 +266,9 @@ def run_claude_pty( route_prompt: Callable[[str], str], switch_message: str, socket_path: Path, - restore_default_model: Callable[[], object], log_path: Path | None = None, ) -> int: - """Run Claude in a PTY, switch its model, restore the default, and replay.""" + """Run Claude in a PTY, switch its model, and replay the first prompt.""" def log(message: str) -> None: if log_path is None: @@ -388,13 +387,11 @@ def on_winch(_signum: int, _frame: object) -> None: and switch_complete is not None and switch_complete.triggered ): - restore_default_model() inject_prompt(master_fd, routed_prompt) phase = "done" - log("[RESTORE] default model restored; first prompt replayed") + log("[REPLAY] first prompt submitted") elif phase == "switching" and now - switch_started >= SWITCH_TIMEOUT_S: os.write(master_fd, b"\x1b") - restore_default_model() inject_note( 1, "Smart Routing could not confirm the model switch. " diff --git a/tests/test_claude_smart_routing_v2.py b/tests/test_claude_smart_routing_v2.py index 474344e2..1260272c 100644 --- a/tests/test_claude_smart_routing_v2.py +++ b/tests/test_claude_smart_routing_v2.py @@ -96,7 +96,6 @@ def fake_run(argv, **kwargs): captured["settings"] = json.loads(generated.read_text()) # Simulate `/model` changing the user file, plus an unrelated concurrent edit. user_settings.write_text(json.dumps({"model": "routed", "theme": "light", "new": True})) - kwargs["restore_default_model"]() return 0 monkeypatch.setattr(claude_pty, "run_claude_pty", fake_run) @@ -166,18 +165,14 @@ def read_until(suffix): })) """.lstrip() ) - restored: list[bool] = [] - 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, - restore_default_model=lambda: restored.append(True), ) assert result == 0 - assert restored == [True] assert json.loads(capture.read_text()) == { "command": "/model system.ai.claude-sonnet-5\r", "replayed": "\x1b[200~fix\nthe parser\x1b[201~\r", From 88b278986fe28f51c65bfae877b66645e09c7e54 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Tue, 25 Aug 2026 22:08:17 +0000 Subject: [PATCH 13/18] Skip legacy routing for Claude v2 launches --- src/ucode/cli.py | 8 +++++++- tests/test_cli.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 1fdb53b9..36b52d61 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -107,6 +107,7 @@ download_managed_skills_on_launch, ) from ucode.smart_routing import claude_routing, codex_routing +from ucode.smart_routing import v2 as smart_routing_v2 from ucode.state import ( STATE_PATH, clear_state, @@ -1881,7 +1882,12 @@ def _launch_tool( managed_launch_model(managed, recommendation, tool) if managed is not None else None ) state, resolved_model = resolve_launch_model(tool, state, managed_model) - if routing_agent is not None and routing_agent.smart_routing_enabled(state): + first_prompt_routes_claude = tool == "claude" and smart_routing_v2.enabled() + if ( + routing_agent is not None + and routing_agent.smart_routing_enabled(state) + and not first_prompt_routes_claude + ): display = TOOL_SPECS[tool]["display"] with spinner(f"Selecting a {display} model with smart routing..."): decision, routing_error = _ROUTING_MODULES[tool].route_launch_model( diff --git a/tests/test_cli.py b/tests/test_cli.py index 99b99dcc..1dabb3b8 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -280,6 +280,34 @@ def test_enabled_codex_launch_uses_routed_root_model(self): in _strip_ansi(result.output) ) + def test_claude_v2_skips_legacy_prelaunch_routing(self, monkeypatch): + monkeypatch.setenv("ENABLE_SMART_ROUTING_V2", "1") + state = { + **MINIMAL_STATE, + "smart_routing_enabled": True, + "claude_models": {"opus": "system.ai.claude-opus-4-8"}, + } + with ( + patch("ucode.cli.ensure_bootstrap_dependencies"), + patch("ucode.cli.load_state", return_value=state), + patch("ucode.cli.ensure_provider_state", return_value=state), + patch("ucode.cli.configure_shared_state", return_value=state), + patch( + "ucode.cli.resolve_launch_model", + return_value=(state, "system.ai.claude-opus-4-8"), + ), + patch("ucode.cli.claude_routing.route_launch_model") as mock_route, + patch("ucode.cli.configure_tool", return_value=state) as mock_configure, + patch("ucode.cli._fetch_managed_config", return_value=(None, False)), + patch("ucode.cli.launch_agent") as mock_launch, + ): + result = runner.invoke(app, ["claude"]) + + assert result.exit_code == 0, result.output + mock_route.assert_not_called() + assert mock_configure.call_args.kwargs["route_root_model"] is None + assert "_claude_launch_model" not in mock_launch.call_args.args[1] + class TestClaudeModelFlag: """`ucode claude --model ` pins the id into the family aliases so the gateway resolves any From 4208e0d924cda38ebd501f37edff100a8e99d8bf Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Tue, 25 Aug 2026 22:29:35 +0000 Subject: [PATCH 14/18] Move Claude v2 launch orchestration --- src/ucode/agents/claude.py | 148 +++----------------------- src/ucode/smart_routing/v2.py | 141 +++++++++++++++++++++++- tests/test_agent_claude.py | 37 +------ tests/test_claude_smart_routing_v2.py | 48 +++++++-- 4 files changed, 196 insertions(+), 178 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index fb20dd89..ad7dcdf1 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -11,7 +11,6 @@ import socket import subprocess import threading -import uuid from collections.abc import Callable from pathlib import Path from typing import cast @@ -27,7 +26,6 @@ ) from ucode.databricks import ( build_auth_shell_command, - build_auth_token_argv, build_tool_base_url, get_databricks_token, ) @@ -35,9 +33,7 @@ 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 @@ -50,7 +46,6 @@ CLAUDE_SETTINGS_PATH = CLAUDE_CONFIG_DIR / "ucode-settings.json" CLAUDE_USER_SETTINGS_PATH = CLAUDE_CONFIG_DIR / "settings.json" CLAUDE_BACKUP_PATH = APP_DIR / "claude-ucode-settings.backup.json" -CLAUDE_MODEL_SNAPSHOT_PATH = APP_DIR / "claude-default-model.snapshot.json" SPEC: ToolSpec = { "binary": "claude", @@ -69,10 +64,6 @@ # marked managed so they're tracked/reverted with the rest of ucode's config. CLAUDE_ROUTING_HOOK_EVENTS = ("PreToolUse", "SessionStart", "SubagentStart") -# Smart-routing-v2 currently uses a fixed route while its dynamic router is wired up. -SMART_ROUTING_V2_MODEL = "system.ai.claude-sonnet-4-6[1m]" -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"]) @@ -886,46 +877,13 @@ def _merge_claude_settings(base: dict, overlay: dict) -> dict: return merged -def _snapshot_user_model_setting() -> dict: - """Capture only Claude's user-level ``model`` setting.""" - settings = read_json_safe(CLAUDE_USER_SETTINGS_PATH) - return {"present": "model" in settings, "value": settings.get("model")} - - -def _save_user_model_snapshot(snapshot: dict, snapshot_path: Path | None = None) -> None: - """Journal the pre-switch model so the next launch can recover after a crash.""" - snapshot_path = snapshot_path or CLAUDE_MODEL_SNAPSHOT_PATH - write_json_file(snapshot_path, snapshot) - - -def _restore_user_model_snapshot( - snapshot_path: Path | None = None, - snapshot: dict | None = None, -) -> bool: - """Restore only the journaled ``model`` field, preserving every sibling setting.""" - snapshot_path = snapshot_path or CLAUDE_MODEL_SNAPSHOT_PATH - if snapshot is None and not snapshot_path.exists(): - return False - snapshot = snapshot if snapshot is not None else read_json_safe(snapshot_path) - present = snapshot.get("present") - if not isinstance(present, bool): - raise RuntimeError(f"Claude model recovery snapshot is invalid: {snapshot_path}") - settings = read_json_safe(CLAUDE_USER_SETTINGS_PATH) - if present: - settings["model"] = snapshot.get("value") - else: - settings.pop("model", None) - write_json_file(CLAUDE_USER_SETTINGS_PATH, settings) - snapshot_path.unlink(missing_ok=True) - return True - - -def _recover_user_model_snapshots() -> None: - """Repair defaults left by interrupted PTY launches.""" - candidates = {CLAUDE_MODEL_SNAPSHOT_PATH} - candidates.update(APP_DIR.glob("claude-default-model.*.snapshot.json")) - for path in sorted(candidates): - _restore_user_model_snapshot(path) +def _compose_v2_settings(tool_args: list[str]) -> tuple[dict, list[str]]: + """Compose caller settings with ucode's Claude settings for a v2 launch.""" + caller_values, remaining = _extract_caller_settings(tool_args) + settings: dict = {} + for value in caller_values: + settings = _merge_claude_settings(settings, _load_caller_settings(value)) + return _merge_claude_settings(settings, read_json_safe(CLAUDE_SETTINGS_PATH)), remaining def _original_launch_model(snapshot: dict, state: dict) -> str | None: @@ -1108,105 +1066,27 @@ def _launch_relayed( raise SystemExit(returncode) -def _v2_router(state: dict): - """Return the first-prompt router; currently fixed while routing is prototyped.""" - - def route(_prompt: str) -> str: - return SMART_ROUTING_V2_MODEL - - return route - - -def _launch_smart_routing_v2( - state: dict, - tool_args: list[str], - *, - model_snapshot: dict, - launch_model: str | None, -) -> None: - """Launch Claude in the first-prompt routing PTY wrapper.""" - 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" - model_snapshot_path = APP_DIR / f"claude-default-model.{run_id}.snapshot.json" - - 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) - model_args = _launch_model_args(remaining, launch_model) - argv = [binary, "--settings", str(settings_path), *model_args, *remaining] - - _save_user_model_snapshot(model_snapshot, model_snapshot_path) - restored = False - - def restore_default_model() -> bool: - nonlocal restored - if restored: - return False - result = _restore_user_model_snapshot(model_snapshot_path, model_snapshot) - restored = True - return result - - print_note( - "Smart routing v2: the first submitted prompt will select Claude Code's " - f"model (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: - # Restore only after Claude exits. Claude persists `/model` asynchronously; - # restoring as soon as its success message renders races that delayed write. - # The journal repairs hard-killed and concurrent launches before they start. - restore_default_model() - 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") # Recover a prior switch interrupted before its surgical settings restore. - _recover_user_model_snapshots() - model_snapshot = _snapshot_user_model_setting() + smart_routing_v2.recover_claude_model_snapshots(CLAUDE_USER_SETTINGS_PATH) + model_snapshot = smart_routing_v2.snapshot_claude_model_setting(CLAUDE_USER_SETTINGS_PATH) launch_model = _original_launch_model(model_snapshot, state) if state.get("claude_relayed"): _launch_relayed(state, binary, tool_args, launch_model) return if smart_routing_v2.enabled() and workspace and os.name != "nt": - _launch_smart_routing_v2( + smart_routing_v2.launch_claude( state, tool_args, + binary=binary, + user_settings_path=CLAUDE_USER_SETTINGS_PATH, model_snapshot=model_snapshot, launch_model=launch_model, + compose_settings=_compose_v2_settings, + launch_model_args=_launch_model_args, ) - return if workspace: os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile")) exec_or_spawn(_build_claude_argv(binary, tool_args, launch_model=launch_model)) diff --git a/src/ucode/smart_routing/v2.py b/src/ucode/smart_routing/v2.py index b129ff3d..71d84f29 100644 --- a/src/ucode/smart_routing/v2.py +++ b/src/ucode/smart_routing/v2.py @@ -7,13 +7,23 @@ import sys import time import urllib.request +import uuid from collections.abc import Callable from pathlib import Path from typing import NoReturn -from ucode.config_io import APP_DIR, deep_merge_dict, read_toml_safe, write_toml_file -from ucode.databricks import get_databricks_token -from ucode.smart_routing import codex_interposer +from ucode.config_io import ( + APP_DIR, + deep_merge_dict, + read_json_safe, + read_toml_safe, + write_json_file, + write_toml_file, +) +from ucode.databricks import build_auth_token_argv, get_databricks_token +from ucode.smart_routing import claude_pty, codex_interposer +from ucode.smart_routing.claude_hooks import FIRST_PROMPT_SOCKET_ENV, sync_first_prompt_hook +from ucode.ui import print_note ENV_VAR = "ENABLE_SMART_ROUTING_V2" @@ -22,6 +32,10 @@ CODEX_INTERPOSER_LOG = APP_DIR / "codex-v2-interposer.log" CODEX_SWITCH_REASON = "Low complexity, unclear intent, and no code reference." # TODO(lilly): replace with smart router rationale. +CLAUDE_TARGET_MODEL = "system.ai.claude-sonnet-4-6[1m]" # TODO(lilly): replace with smart router. +CLAUDE_PTY_LOG = APP_DIR / "claude-v2-pty.log" +CLAUDE_MODEL_SNAPSHOT_PATH = APP_DIR / "claude-default-model.snapshot.json" + APP_SERVER_READY_TIMEOUT_SECONDS = 30 PROCESS_SHUTDOWN_TIMEOUT_SECONDS = 5 OAUTH_TOKEN_ENV_VAR = "OAUTH_TOKEN" @@ -35,6 +49,48 @@ def enabled() -> bool: return os.environ.get(ENV_VAR) == "1" +def snapshot_claude_model_setting(user_settings_path: Path) -> dict: + """Capture only Claude's user-level ``model`` setting.""" + settings = read_json_safe(user_settings_path) + return {"present": "model" in settings, "value": settings.get("model")} + + +def _save_claude_model_snapshot(snapshot: dict, snapshot_path: Path) -> None: + """Journal the pre-switch model so the next launch can recover after a crash.""" + write_json_file(snapshot_path, snapshot) + + +def restore_claude_model_snapshot( + user_settings_path: Path, + snapshot_path: Path | None = None, + snapshot: dict | None = None, +) -> bool: + """Restore only the journaled ``model`` field, preserving every sibling setting.""" + snapshot_path = snapshot_path or CLAUDE_MODEL_SNAPSHOT_PATH + if snapshot is None and not snapshot_path.exists(): + return False + snapshot = snapshot if snapshot is not None else read_json_safe(snapshot_path) + present = snapshot.get("present") + if not isinstance(present, bool): + raise RuntimeError(f"Claude model recovery snapshot is invalid: {snapshot_path}") + settings = read_json_safe(user_settings_path) + if present: + settings["model"] = snapshot.get("value") + else: + settings.pop("model", None) + write_json_file(user_settings_path, settings) + snapshot_path.unlink(missing_ok=True) + return True + + +def recover_claude_model_snapshots(user_settings_path: Path) -> None: + """Repair defaults left by interrupted Claude PTY launches.""" + candidates = {CLAUDE_MODEL_SNAPSHOT_PATH} + candidates.update(APP_DIR.glob("claude-default-model.*.snapshot.json")) + for path in sorted(candidates): + restore_claude_model_snapshot(user_settings_path, path) + + def _loopback_websocket_url(port: int) -> str: return f"ws://{LOOPBACK_HOST}:{port}" @@ -90,6 +146,85 @@ def _generate_codex_app_server_home( return CODEX_APP_SERVER_HOME +def _route_claude_prompt(_prompt: str) -> str: + return CLAUDE_TARGET_MODEL + + +def launch_claude( + state: dict, + tool_args: list[str], + *, + binary: str, + user_settings_path: Path, + model_snapshot: dict, + launch_model: str | None, + compose_settings: Callable[[list[str]], tuple[dict, list[str]]], + launch_model_args: Callable[[list[str], str | None], list[str]], +) -> NoReturn: + """Launch Claude in the first-prompt routing PTY wrapper.""" + workspace = state.get("workspace") + if not workspace: + raise RuntimeError( + "Smart routing v2 needs a configured workspace; run `ucode configure claude` first." + ) + os.environ[OAUTH_TOKEN_ENV_VAR] = get_databricks_token(workspace, state.get("profile")) + + run_id = f"{os.getpid()}-{uuid.uuid4().hex[:8]}" + socket_path = APP_DIR / f"claude-v2-{run_id}.sock" + settings_path = APP_DIR / f"claude-v2-{run_id}.json" + model_snapshot_path = APP_DIR / f"claude-default-model.{run_id}.snapshot.json" + + settings, remaining = compose_settings(tool_args) + hook_executable = build_auth_token_argv( + workspace, state.get("profile"), use_pat=bool(state.get("use_pat")) + )[0] + env = settings.setdefault("env", {}) + if not isinstance(env, dict): + raise RuntimeError("Claude settings 'env' must be an object for smart routing.") + env[FIRST_PROMPT_SOCKET_ENV] = str(socket_path) + sync_first_prompt_hook(settings, hook_executable) + write_json_file(settings_path, settings) + model_args = launch_model_args(remaining, launch_model) + argv = [binary, "--settings", str(settings_path), *model_args, *remaining] + + _save_claude_model_snapshot(model_snapshot, model_snapshot_path) + restored = False + + def restore_default_model() -> bool: + nonlocal restored + if restored: + return False + result = restore_claude_model_snapshot( + user_settings_path, model_snapshot_path, model_snapshot + ) + restored = True + return result + + print_note( + "Smart routing v2: the first submitted prompt will select Claude Code's " + f"model ({CLAUDE_TARGET_MODEL}); log: {CLAUDE_PTY_LOG}." + ) + try: + returncode = claude_pty.run_claude_pty( + argv, + route_prompt=_route_claude_prompt, + switch_message=( + f"✨ Databricks Smart Router selected {CLAUDE_TARGET_MODEL} due to " + "low complexity, unclear intent, and no code reference." + ), + socket_path=socket_path, + log_path=CLAUDE_PTY_LOG, + ) + finally: + # Restore only after Claude exits. Claude persists `/model` asynchronously; + # restoring as soon as its success message renders races that delayed write. + # The journal repairs hard-killed and concurrent launches before they start. + restore_default_model() + settings_path.unlink(missing_ok=True) + socket_path.unlink(missing_ok=True) + sys.exit(returncode) + + def launch_codex( state: dict, tool_args: list[str], diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index c2664d69..fb132354 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -9,7 +9,7 @@ import pytest from ucode.agents import claude -from ucode.smart_routing import claude_routing +from ucode.smart_routing import claude_routing, v2 WS = "https://example.databricks.com" @@ -666,9 +666,9 @@ def fake_execvp(binary: str, args: list[str]) -> None: claude, "get_databricks_token", lambda workspace, profile=None: "fresh-token" ) monkeypatch.setattr( - claude, - "_snapshot_user_model_setting", - lambda: {"present": True, "value": "opus"}, + v2, + "snapshot_claude_model_setting", + lambda _path: {"present": True, "value": "opus"}, ) monkeypatch.setattr(os, "execvp", fake_execvp) @@ -920,34 +920,7 @@ def test_malformed_file_json_raises(self, tmp_path, monkeypatch): claude._build_claude_argv("claude", ["--settings", str(bad_file)]) -class TestClaudeDefaultModelRecovery: - def _paths(self, monkeypatch, tmp_path): - user = tmp_path / "settings.json" - snapshot = tmp_path / "model-snapshot.json" - monkeypatch.setattr(claude, "CLAUDE_USER_SETTINGS_PATH", user) - monkeypatch.setattr(claude, "CLAUDE_MODEL_SNAPSHOT_PATH", snapshot) - return user, snapshot - - def test_restore_changes_only_model_field(self, monkeypatch, tmp_path): - user, _snapshot = self._paths(monkeypatch, tmp_path) - user.write_text(json.dumps({"model": "opus", "theme": "dark"})) - original = claude._snapshot_user_model_setting() - claude._save_user_model_snapshot(original) - - user.write_text(json.dumps({"model": "routed", "theme": "light", "new": True})) - assert claude._restore_user_model_snapshot() is True - assert json.loads(user.read_text()) == {"model": "opus", "theme": "light", "new": True} - assert claude._restore_user_model_snapshot() is False - - def test_restore_removes_model_when_original_was_absent(self, monkeypatch, tmp_path): - user, _snapshot = self._paths(monkeypatch, tmp_path) - user.write_text(json.dumps({"theme": "dark"})) - claude._save_user_model_snapshot(claude._snapshot_user_model_setting()) - user.write_text(json.dumps({"model": "routed", "theme": "light"})) - - claude._restore_user_model_snapshot() - assert json.loads(user.read_text()) == {"theme": "light"} - +class TestClaudeLaunchModel: def test_missing_user_default_falls_back_to_ucode_model(self): state = {"claude_models": {"opus": "system.ai.claude-opus-4-8"}} assert ( diff --git a/tests/test_claude_smart_routing_v2.py b/tests/test_claude_smart_routing_v2.py index 1260272c..a67f78d5 100644 --- a/tests/test_claude_smart_routing_v2.py +++ b/tests/test_claude_smart_routing_v2.py @@ -12,7 +12,7 @@ import pytest from ucode.agents import claude -from ucode.smart_routing import claude_hooks, claude_pty +from ucode.smart_routing import claude_hooks, claude_pty, v2 class TestDirectModelCommand: @@ -78,16 +78,15 @@ class TestV2Launch: def test_saves_default_passes_model_and_restores_only_model(self, tmp_path, monkeypatch): ucode_settings = tmp_path / "ucode-settings.json" user_settings = tmp_path / "settings.json" - snapshot_path = tmp_path / "model-snapshot.json" ucode_settings.write_text(json.dumps({"env": {"ANTHROPIC_BASE_URL": "https://gw"}})) user_settings.write_text(json.dumps({"model": "opus", "theme": "dark"})) monkeypatch.setattr(claude, "APP_DIR", tmp_path) monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", ucode_settings) monkeypatch.setattr(claude, "CLAUDE_USER_SETTINGS_PATH", user_settings) - monkeypatch.setattr(claude, "CLAUDE_MODEL_SNAPSHOT_PATH", snapshot_path) - monkeypatch.setattr(claude, "SMART_ROUTING_V2_CLAUDE_LOG", tmp_path / "v2.log") - monkeypatch.setattr(claude, "get_databricks_token", lambda *_args, **_kwargs: "token") - monkeypatch.setattr(claude, "build_auth_token_argv", lambda *_args, **_kwargs: ["ucode"]) + monkeypatch.setattr(v2, "APP_DIR", tmp_path) + monkeypatch.setattr(v2, "CLAUDE_PTY_LOG", tmp_path / "v2.log") + monkeypatch.setattr(v2, "get_databricks_token", lambda *_args, **_kwargs: "token") + monkeypatch.setattr(v2, "build_auth_token_argv", lambda *_args, **_kwargs: ["ucode"]) captured: dict = {} def fake_run(argv, **kwargs): @@ -99,13 +98,17 @@ def fake_run(argv, **kwargs): return 0 monkeypatch.setattr(claude_pty, "run_claude_pty", fake_run) - snapshot = claude._snapshot_user_model_setting() + snapshot = v2.snapshot_claude_model_setting(user_settings) with pytest.raises(SystemExit) as exc: - claude._launch_smart_routing_v2( + v2.launch_claude( {"workspace": "https://example.com"}, ["--debug"], + binary="claude", + user_settings_path=user_settings, model_snapshot=snapshot, launch_model="opus", + compose_settings=claude._compose_v2_settings, + launch_model_args=claude._launch_model_args, ) assert exc.value.code == 0 @@ -117,10 +120,37 @@ def fake_run(argv, **kwargs): "theme": "light", "new": True, } - assert not snapshot_path.exists() assert not list(tmp_path.glob("claude-default-model.*.snapshot.json")) +class TestModelRecovery: + def test_restore_changes_only_model_field(self, tmp_path): + user = tmp_path / "settings.json" + snapshot_path = tmp_path / "model-snapshot.json" + user.write_text(json.dumps({"model": "opus", "theme": "dark"})) + original = v2.snapshot_claude_model_setting(user) + v2._save_claude_model_snapshot(original, snapshot_path) + + user.write_text(json.dumps({"model": "routed", "theme": "light", "new": True})) + assert v2.restore_claude_model_snapshot(user, snapshot_path) is True + assert json.loads(user.read_text()) == { + "model": "opus", + "theme": "light", + "new": True, + } + assert v2.restore_claude_model_snapshot(user, snapshot_path) is False + + def test_restore_removes_model_when_original_was_absent(self, tmp_path): + user = tmp_path / "settings.json" + snapshot_path = tmp_path / "model-snapshot.json" + user.write_text(json.dumps({"theme": "dark"})) + v2._save_claude_model_snapshot(v2.snapshot_claude_model_setting(user), snapshot_path) + user.write_text(json.dumps({"model": "routed", "theme": "light"})) + + v2.restore_claude_model_snapshot(user, snapshot_path) + assert json.loads(user.read_text()) == {"theme": "light"} + + class TestPtyFlow: def test_direct_switch_restore_and_replay(self, tmp_path): fake_claude = tmp_path / "fake_claude.py" From 3c7e803b22056e7dd499ce09e8e31633afbdd28b Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Wed, 26 Aug 2026 02:17:28 +0000 Subject: [PATCH 15/18] update --- src/ucode/agents/claude.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index cb811345..85f83ea7 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -1011,9 +1011,7 @@ def _rewrite_relayed_port(state: dict, port: int) -> None: write_json_file(CLAUDE_SETTINGS_PATH, settings) -def _launch_relayed( - state: dict, binary: str, tool_args: list[str], launch_model: str | None = None -) -> None: +def _launch_relayed(state: dict, binary: str, tool_args: list[str]) -> None: """Relayed launch: sign into the Claude subscription, start the loopback refresh proxy, then run Claude Code alongside it (the proxy must outlive the exec, so we spawn-and-wait rather than replacing the process).""" @@ -1059,9 +1057,7 @@ def _launch_relayed( server_thread = threading.Thread(target=server.serve_forever, daemon=True) server_thread.start() - proc = subprocess.Popen( - _build_claude_argv(binary, tool_args, relayed=True, launch_model=launch_model) - ) + proc = subprocess.Popen(_build_claude_argv(binary, tool_args, relayed=True)) try: returncode = proc.wait() except KeyboardInterrupt: @@ -1123,7 +1119,7 @@ def launch(state: dict, tool_args: list[str]) -> None: model_snapshot = smart_routing_v2.snapshot_claude_model_setting(CLAUDE_USER_SETTINGS_PATH) launch_model = _original_launch_model(model_snapshot, state) if state.get("claude_relayed"): - _launch_relayed(state, binary, tool_args, launch_model) + _launch_relayed(state, binary, tool_args) return if smart_routing_v2.enabled() and workspace and os.name != "nt": smart_routing_v2.launch_claude( From 23c4f3e7582ce8649b1b8ff6a8056e8042ea2e7f Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Wed, 26 Aug 2026 16:05:29 +0000 Subject: [PATCH 16/18] Format Claude smart routing notices --- src/ucode/smart_routing/claude_pty.py | 19 ++++++++++++++++--- src/ucode/smart_routing/v2.py | 9 ++++----- tests/test_claude_smart_routing_v2.py | 10 ++++++++++ 3 files changed, 30 insertions(+), 8 deletions(-) diff --git a/src/ucode/smart_routing/claude_pty.py b/src/ucode/smart_routing/claude_pty.py index 43e90fdc..efca9d65 100644 --- a/src/ucode/smart_routing/claude_pty.py +++ b/src/ucode/smart_routing/claude_pty.py @@ -55,6 +55,20 @@ def valid_model_name(name: object) -> bool: ) +def switch_message(model: str, reason: str) -> str: + """Format the routed-model notice shown in Claude Code.""" + lines = [ + "Using Unity Gateway Smart Router.", + f"Selected Model : {model}", + f"Reason : {reason}", + ] + width = max(len(line) for line in lines) + border = "─" * (width + 2) + return "\n".join( + [f"┌{border}┐", *(f"│ {line:<{width}} │" for line in lines), f"└{border}┘"] + ) + + class ConfirmationState: """Detect and accept Claude's optional cache-cost confirmation dialog.""" @@ -156,9 +170,8 @@ def first_prompt_hook_output(response: dict | None) -> dict | None: return None return { "decision": "block", - "reason": ( - f"✨ Smart Router selected {model} due to low complexity, unclear intent, " - "and no code reference." + "reason": switch_message( + model, "Low complexity, unclear intent, and no code reference." ), } diff --git a/src/ucode/smart_routing/v2.py b/src/ucode/smart_routing/v2.py index 789d4a5f..32b2ca58 100644 --- a/src/ucode/smart_routing/v2.py +++ b/src/ucode/smart_routing/v2.py @@ -24,7 +24,7 @@ CODEX_TARGET_MODEL = "system.ai.glm-5-2" # TODO(lilly): replace with smart router. CODEX_INTERPOSER_LOG = APP_DIR / "codex-v2-interposer.log" -CODEX_SWITCH_REASON = "Low complexity, unclear intent, and no code reference." # TODO(lilly): replace with smart router rationale. +STUBBED_SWITCH_REASON = "Low complexity, unclear intent, and no code reference." # TODO(lilly): replace with smart router rationale. CLAUDE_TARGET_MODEL = "system.ai.claude-sonnet-4-6[1m]" # TODO(lilly): replace with smart router. CLAUDE_PTY_LOG = APP_DIR / "claude-v2-pty.log" @@ -182,9 +182,8 @@ def restore_default_model() -> bool: returncode = claude_pty.run_claude_pty( argv, route_prompt=_route_claude_prompt, - switch_message=( - f"✨ Databricks Smart Router selected {CLAUDE_TARGET_MODEL} due to " - "low complexity, unclear intent, and no code reference." + switch_message=claude_pty.switch_message( + CLAUDE_TARGET_MODEL, STUBBED_SWITCH_REASON ), socket_path=socket_path, log_path=CLAUDE_PTY_LOG, @@ -270,7 +269,7 @@ def launch_codex( LOOPBACK_HOST, app_server_url, CODEX_TARGET_MODEL, - switch_message=_switch_message(CODEX_TARGET_MODEL, CODEX_SWITCH_REASON), + switch_message=_switch_message(CODEX_TARGET_MODEL, STUBBED_SWITCH_REASON), log_path=CODEX_INTERPOSER_LOG, ) tui_url = _loopback_websocket_url(tui_port) diff --git a/tests/test_claude_smart_routing_v2.py b/tests/test_claude_smart_routing_v2.py index a67f78d5..0a1dcbfa 100644 --- a/tests/test_claude_smart_routing_v2.py +++ b/tests/test_claude_smart_routing_v2.py @@ -38,6 +38,16 @@ def test_types_direct_model_command(self): class TestFirstPromptHook: + def test_renders_boxed_router_notice(self): + model = "system.ai.claude-sonnet-4-6[1m]" + reason = "Low complexity, unclear intent, and no code reference." + result = claude_pty.first_prompt_hook_output( + {"action": "block", "model": model} + ) + + assert result == {"decision": "block", "reason": v2._switch_message(model, reason)} + assert claude_pty.switch_message(model, reason) == v2._switch_message(model, reason) + def test_blocks_once_then_allows_replay(self, tmp_path): socket_path = tmp_path / "first.sock" blocked: list[tuple[str, str]] = [] From 3cefefa45db082b3cc3092f03ad35660e6ede7cf Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Wed, 26 Aug 2026 16:36:29 +0000 Subject: [PATCH 17/18] update --- src/ucode/agents/claude.py | 8 +++++++- tests/test_agent_claude.py | 12 ++++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 85f83ea7..dea3ce80 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -1121,7 +1121,13 @@ def launch(state: dict, tool_args: list[str]) -> None: if state.get("claude_relayed"): _launch_relayed(state, binary, tool_args) return - if smart_routing_v2.enabled() and workspace and os.name != "nt": + # Smart routing v2 needs Unix PTY support, which Windows does not provide. + if smart_routing_v2.enabled() and os.name == "nt": + raise RuntimeError( + "Smart routing in Claude Code is currently not supported on Windows. " + "Please use Codex or disable smart routing." + ) + if smart_routing_v2.enabled() and workspace: smart_routing_v2.launch_claude( state, tool_args, diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 28b7e4f5..0bd0a889 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -654,6 +654,18 @@ def boom(name, entry, scope=mcp_mod.MCP_USER_SCOPE): class TestClaudeLaunch: + def test_smart_routing_on_windows_is_not_supported(self, monkeypatch): + monkeypatch.setenv(v2.ENV_VAR, "1") + monkeypatch.setattr(claude.os, "name", "nt") + monkeypatch.setattr(v2, "recover_claude_model_snapshots", lambda _path: None) + monkeypatch.setattr(v2, "snapshot_claude_model_setting", lambda _path: {"present": False}) + + with pytest.raises( + RuntimeError, + match="Smart routing in Claude Code is currently not supported on Windows", + ): + claude.launch({"workspace": WS, "profile": "test"}, ["--debug"]) + def test_default_launch_keeps_existing_auth_path(self, monkeypatch): calls: list[list[str]] = [] monkeypatch.delenv(claude.GATEWAY_MODEL_DISCOVERY_ENV_VAR, raising=False) From 90516c30325e84fb66652e7aa87dc0513cebefc3 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Wed, 26 Aug 2026 20:04:28 +0000 Subject: [PATCH 18/18] Reuse shared loopback host constant --- src/ucode/smart_routing/v2.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/ucode/smart_routing/v2.py b/src/ucode/smart_routing/v2.py index 32b2ca58..a6541b8f 100644 --- a/src/ucode/smart_routing/v2.py +++ b/src/ucode/smart_routing/v2.py @@ -15,6 +15,7 @@ import tomlkit from ucode.config_io import APP_DIR, read_json_safe, write_json_file +from ucode.constants import LOOPBACK_HOST from ucode.databricks import build_auth_token_argv, get_databricks_token from ucode.smart_routing import claude_pty, codex_interposer from ucode.smart_routing.claude_hooks import FIRST_PROMPT_SOCKET_ENV, sync_first_prompt_hook @@ -33,7 +34,6 @@ APP_SERVER_READY_TIMEOUT_SECONDS = 30 PROCESS_SHUTDOWN_TIMEOUT_SECONDS = 5 OAUTH_TOKEN_ENV_VAR = "OAUTH_TOKEN" -LOOPBACK_HOST = "127.0.0.1" HEALTH_REQUEST_TIMEOUT_SECONDS = 1 HEALTH_POLL_INTERVAL_SECONDS = 0.25