smart routing: switch model in codex after first prompt

pyproject.toml

@@ -26,6 +26,8 @@ dependencies = [
   "questionary>=2.0.0",
   "tomlkit>=0.13.0",
   "typer>=0.12.0",
+  # enable changing codex model after first prompt
+  "websockets>=13",
 ]
 
 [project.optional-dependencies]

src/ucode/agents/codex.py

@@ -4,6 +4,7 @@
 
 import os
 import re
+import signal
 import subprocess
 import sys
 import time
@@ -27,13 +28,14 @@
 )
 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,
 )
 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 +52,18 @@
 # tool (codex, claude), so a workspace turns it on once.
 SMART_ROUTING_STATE_KEY = "smart_routing_enabled"
 
+# Codex-specific smart-routing-v2 settings. The shared enable flag + hold-turns live in
+# `smart_routing.v2`; here we keep only what is Codex-specific: the switch-to model and the
+# app-server's CODEX_HOME / interposer log paths. When enabled, a single `ucode codex`
+# launches the REAL Codex TUI against a ucode-run `codex app-server` with a WebSocket
+# interposer (smart_routing.codex_interposer); ucode owns all three processes and tears the
+# app-server + interposer down when the TUI exits.
+SMART_ROUTING_V2_TARGET_MODEL = "gpt-5.5"  # hardcoded switch-to model for now
+SMART_ROUTING_V2_HOME = APP_DIR / "codex-v2-home"  # CODEX_HOME for the ucode-run app-server
+SMART_ROUTING_V2_LOG = (
+    APP_DIR / "codex-v2-interposer.log"
+)  # interposer log (not stdout: TUI owns it)
+
 
 SPEC: ToolSpec = {
     "binary": "codex",
@@ -465,9 +479,110 @@ def _gpt_version_key(entry: tuple[str, tuple[int, int | None, int | None, str]])
 _PROFILE_REJECTED_MAX_SECONDS = 3.0
 
 
+def _generate_v2_app_server_home(state: dict, model: str) -> Path:
+    """Write an isolated CODEX_HOME whose config.toml carries the ucode gateway
+    provider block, for the ucode-run `codex app-server`.
+
+    The app-server rejects the global `--profile`, so a default-config CODEX_HOME
+    is how it inherits ucode's gateway (base_url + `ucode auth-token` refresh).
+    Reuses `render_overlay` — the same provider block `ucode configure codex` writes."""
+    home = SMART_ROUTING_V2_HOME
+    home.mkdir(parents=True, exist_ok=True)
+    config_path = home / "config.toml"
+    overlay = render_overlay(
+        state["workspace"],
+        model,
+        state.get("profile"),
+        use_pat=bool(state.get("use_pat")),
+    )
+    doc = read_toml_safe(config_path)
+    deep_merge_dict(doc, overlay)
+    write_toml_file(config_path, doc)
+    return home
+
+
+def _launch_smart_routing_v2(state: dict, tool_args: list[str]) -> None:
+    """Experimental single-command launch of the real Codex TUI with runtime model switching.
+
+    ucode owns three processes: a `codex app-server` subprocess, the WebSocket interposer
+    (daemon thread), and the `codex --remote` TUI (foreground). The interposer holds the first
+    turn on the normal model, then rewrites subsequent turns to SMART_ROUTING_V2_TARGET_MODEL and
+    injects a settings update so the TUI reflects the switch. The app-server + interposer are torn
+    down when the TUI exits. Mirrors the lifecycle of `claude.py::_launch_relayed`.
+    """
+    from ucode.smart_routing import codex_interposer
+
+    binary = SPEC["binary"]
+    workspace = state.get("workspace")
+    if not workspace:
+        raise RuntimeError(
+            "Smart routing v2 needs a configured workspace; run `ucode configure codex` first."
+        )
+    start_model = default_model(state)
+    if not start_model:
+        raise RuntimeError(
+            "Smart routing v2 could not determine a starting Codex model for this workspace."
+        )
+
+    os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile"))
+    home = _generate_v2_app_server_home(state, start_model)
+    app_port = codex_interposer.free_port()
+    tui_port = codex_interposer.free_port()
+
+    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.SWITCH_AFTER_TURNS,
+            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

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