From fafe7af7c7650c55724a329c84842bcce7375a02 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Thu, 20 Aug 2026 23:04:50 +0000 Subject: [PATCH 1/3] 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 2/3] 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 3/3] 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