smart routing: route Codex v2 subagents

src/ucode/agents/codex.py

@@ -29,6 +29,7 @@
 from ucode.managed_files import OS, current_os, write_managed_file
 from ucode.smart_routing.codex_hooks import (
     remove_smart_routing_hooks,
+    routing_models,
     sync_smart_routing_hooks,
 )
 from ucode.state import mark_tool_managed, save_state
@@ -447,9 +448,9 @@ def default_model(state: dict) -> str | None:
     """
     if isinstance(state.get("codex_default_model"), str):
         return state.get("codex_default_model")
-    codex_models = state.get("codex_models") or []
+    models = routing_models(state)
     parsed: list[tuple[str, tuple[int, int | None, int | None, str]]] = [
-        (mid, gpt) for mid in codex_models if (gpt := _parse_gpt(mid)) is not None
+        (mid, gpt) for mid in models if (gpt := _parse_gpt(mid)) is not None
     ]
     if parsed:
 
@@ -464,7 +465,7 @@ def _gpt_version_key(entry: tuple[str, tuple[int, int | None, int | None, str]])
     # after stripping the system.ai. prefix). gpt-oss-* models are confirmed
     # routable through the responses API; non-GPT ids (e.g. moonshotai/kimi-k2.5)
     # would be rejected by the gateway, so they stay excluded.
-    gpt_family = [m for m in codex_models if _is_gpt_family(m)]
+    gpt_family = [m for m in models if _is_gpt_family(m)]
     return gpt_family[0] if gpt_family else None
 
 

src/ucode/cli.py

@@ -642,7 +642,9 @@ def configure_shared_state(
     )
     want_gemini = fetch_all or "gemini" in tools or "opencode" in tools or "pi" in tools
     want_codex = fetch_all or "codex" in tools or "copilot" in tools or "pi" in tools
-    want_oss = fetch_all or "opencode" in tools
+    # Codex smart routing can select OSS models such as GLM, so a Codex-only
+    # configure must persist that discovered family too.
+    want_oss = fetch_all or "opencode" in tools or "codex" in tools
 
     claude_reason: str | None = None
     gemini_reason: str | None = None
@@ -1379,6 +1381,21 @@ def auth_token_cmd(
     sys.stdout.write(token + "\n")
 
 
+def _oauth_token_is_fresh(token: str, buffer_seconds: float = 120) -> bool:
+    import base64
+    import binascii
+    import json
+    import time
+
+    try:
+        payload = token.split(".")[1]
+        payload += "=" * (-len(payload) % 4)
+        expires_at = float(json.loads(base64.urlsafe_b64decode(payload))["exp"])
+    except (IndexError, KeyError, TypeError, ValueError, binascii.Error, json.JSONDecodeError):
+        return False
+    return time.time() < expires_at - buffer_seconds
+
+
 @app.command("codex-router-hook", hidden=True)
 def codex_router_hook_cmd(
     event: str,
@@ -1433,14 +1450,16 @@ def codex_router_hook_cmd(
         return
     if event != "route-subagent" or not host:
         return
-    token = os.environ.get("OAUTH_TOKEN") or os.environ.get("DATABRICKS_BEARER")
+    if use_pat and not ensure_pat_bearer(profile):
+        return
+    token = os.environ.get("DATABRICKS_BEARER", "").strip()
     if not token:
-        if use_pat and not ensure_pat_bearer(profile):
-            return
-        try:
-            token = get_databricks_token(host, profile)
-        except RuntimeError:
-            return
+        token = os.environ.get("OAUTH_TOKEN", "").strip()
+        if not _oauth_token_is_fresh(token):
+            try:
+                token = get_databricks_token(host, profile, force_refresh=True)
+            except RuntimeError:
+                return
     output = route_pre_tool_use(
         payload,
         workspace=host,

src/ucode/config_io.py

@@ -181,7 +181,10 @@ def parse_dotenv(path: Path) -> dict[str, str]:
     """Parse a simple KEY=VALUE / KEY="VALUE" .env file, preserving insertion order.
 
     Comments and blank lines are dropped on round-trip. Lines that don't look
-    like KEY=... are skipped.
+    like KEY=... are skipped. Leading whitespace (line indentation and spacing
+    after the ``=``) is trimmed, but the exact characters up to the end of the
+    line are preserved — including any trailing spaces — so a value such as
+    ``token = abc123 `` keeps its trailing space.
     """
     if not path.exists():
         return {}
@@ -191,16 +194,20 @@ def parse_dotenv(path: Path) -> dict[str, str]:
     except OSError:
         return {}
     for raw_line in text.splitlines():
-        line = raw_line.strip()
-        if not line or line.startswith("#"):
+        # Detect blank / comment lines on the fully-stripped form so trailing
+        # spaces on value lines don't change which lines are skipped.
+        stripped = raw_line.strip()
+        if not stripped or stripped.startswith("#"):
             continue
-        if "=" not in line:
+        if "=" not in stripped:
             continue
-        key, _, val = line.partition("=")
+        # Only strip *leading* whitespace from the line so the characters
+        # between "=" and end-of-line (including trailing spaces) are preserved.
+        key, _, val = raw_line.lstrip().partition("=")
         key = key.strip()
         if not key:
             continue
-        val = val.strip()
+        val = val.lstrip()
         if len(val) >= 2 and val[0] == val[-1] and val[0] in ('"', "'"):
             val = val[1:-1]
         env[key] = val

src/ucode/smart_routing/codex_hooks.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import copy
 import shlex
 import subprocess
 
@@ -11,6 +12,16 @@
 ROUTING_HOOK_COMMAND_MARKER = "codex-router-hook"
 
 
+def routing_models(state: dict) -> list[str]:
+    """Return the configured model services compatible with Codex routing."""
+    models: list[str] = []
+    for key in ("codex_models", "oss_models"):
+        values = state.get(key)
+        if isinstance(values, list):
+            models.extend(value for value in values if isinstance(value, str) and value)
+    return list(dict.fromkeys(models))
+
+
 def sync_smart_routing_hooks(doc: dict, state: dict, *, enabled: bool) -> None:
     """Synchronize ucode-managed routing hooks in a Codex config document."""
     groups = _routing_hook_groups(state) if enabled else {}
@@ -23,16 +34,10 @@ def remove_smart_routing_hooks(doc: dict) -> bool:
 
 
 def _routing_hook_groups(state: dict) -> dict[str, list[dict]]:
-    route_argv = _routing_hook_argv(state, "route-subagent")
     session_argv = _routing_hook_argv(state, "session-start")
     subagent_argv = _routing_hook_argv(state, "record-subagent")
     return {
-        "PreToolUse": [
-            {
-                "matcher": "Agent|.*spawn_agent$",
-                "hooks": [_routing_command_hook(route_argv, status="Routing subagent model")],
-            }
-        ],
+        "PreToolUse": [_pre_tool_use_hook_group(state)],
         "SessionStart": [
             {
                 "matcher": "startup|resume|clear",
@@ -47,7 +52,34 @@ def _routing_hook_groups(state: dict) -> dict[str, list[dict]]:
     }
 
 
-def _routing_hook_argv(state: dict, event: str) -> list[str]:
+def merge_pre_tool_use_hooks(
+    existing: list[dict], state: dict, *, available_models: list[str]
+) -> list[dict]:
+    """Add the ucode spawn hook to an existing Codex PreToolUse hook list."""
+    doc = {"hooks": {"PreToolUse": copy.deepcopy(existing)}}
+    hooks.sync_managed_hooks(
+        doc,
+        ROUTING_HOOK_COMMAND_MARKER,
+        {"PreToolUse": [_pre_tool_use_hook_group(state, available_models=available_models)]},
+    )
+    return doc["hooks"]["PreToolUse"]
+
+
+def _pre_tool_use_hook_group(state: dict, *, available_models: list[str] | None = None) -> dict:
+    route_argv = _routing_hook_argv(
+        state,
+        "route-subagent",
+        available_models=available_models,
+    )
+    return {
+        "matcher": "Agent|.*spawn_agent$",
+        "hooks": [_routing_command_hook(route_argv, status="Routing subagent model")],
+    }
+
+
+def _routing_hook_argv(
+    state: dict, event: str, *, available_models: list[str] | None = None
+) -> list[str]:
     workspace = str(state.get("workspace") or "")
     argv = [
         build_auth_token_argv(workspace, state.get("profile"), use_pat=bool(state.get("use_pat")))[
@@ -64,7 +96,8 @@ def _routing_hook_argv(state: dict, event: str) -> list[str]:
         argv += ["--profile", profile]
     if state.get("use_pat"):
         argv.append("--use-pat")
-    for model in state.get("codex_models") or []:
+    models = available_models if available_models is not None else routing_models(state)
+    for model in models:
         if isinstance(model, str) and model:
             argv += ["--model", model]
     return argv

src/ucode/smart_routing/codex_interposer.py

@@ -12,7 +12,7 @@
 from websockets.asyncio.client import connect
 from websockets.asyncio.server import serve
 
-from ucode.smart_routing import routing
+from ucode.smart_routing import codex_routing, routing
 
 SETTINGS_UPDATED = "thread/settings/updated"
 ITEM_STARTED = "item/started"
@@ -43,36 +43,6 @@ def _prompt_from_turn(params: dict) -> str | None:
     return prompt or None
 
 
-def _request_routing_decision(
-    workspace: str,
-    token: str,
-    prompt: str,
-    available_models: list[str],
-    log: Callable[[str], None] | None = None,
-) -> tuple[routing.RoutingDecision | None, str | None]:
-    available = {routing.normalize_model(model): model for model in available_models}
-    route_options = [(model, "codex") for model in available]
-    if not route_options:
-        return None, "no cached model services are available"
-    if log is not None:
-        payload = {
-            "route_options": [
-                {"model": model, "harness": harness} for model, harness in route_options
-            ],
-            "task": {"prompt": prompt},
-            "route_selector": {"router_name": routing.ROUTER_NAME},
-        }
-        url = workspace.rstrip("/") + routing.ROUTING_PATH
-        log(f"[ROUTE] request POST {url}: {json.dumps(payload, separators=(',', ':'))}")
-    return routing.select_route(
-        workspace,
-        token,
-        prompt,
-        route_options,
-        lambda selected: available.get(routing.normalize_model(selected)),
-    )
-
-
 class _Session:
     def __init__(
         self,
@@ -93,6 +63,7 @@ def __init__(
         self.settings: dict | None = None
         self.first_turn_seen = False
         self.switch_pending = False
+        self.notice_pending = False
         self.injected = False
 
     def on_tui_frame(self, raw: str) -> str:
@@ -122,6 +93,7 @@ def on_tui_frame(self, raw: str) -> str:
                 self.target = decision.model
                 if self.switch_message_fn is not None:
                     self.switch_message = self.switch_message_fn(decision.model, decision.rationale)
+                self.notice_pending = self.switch_message is not None
                 self.log(f"[ROUTE] selected {decision.model!r}; rationale={decision.rationale!r}")
             old = params.get("model")
             if self.target is not None and old != self.target:
@@ -151,20 +123,24 @@ def on_engine_frame(self, raw: str) -> list[dict]:
         if (
             msg.get("method") == TURN_STARTED
             and not self.injected
-            and self.switch_pending
+            and (self.switch_pending or self.notice_pending)
             and self.thread_id
         ):
             self.injected = True
+            switch_pending = self.switch_pending
             self.switch_pending = False
-            settings = dict(self.settings) if isinstance(self.settings, dict) else {}
-            settings["model"] = self.target
-            self.log(f"[INJECT] {SETTINGS_UPDATED}: model -> {self.target!r} (flip TUI chip)")
-            injected: list[dict] = [
-                {
-                    "method": SETTINGS_UPDATED,
-                    "params": {"threadId": self.thread_id, "threadSettings": settings},
-                }
-            ]
+            self.notice_pending = False
+            injected: list[dict] = []
+            if switch_pending:
+                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)")
+                injected.append(
+                    {
+                        "method": SETTINGS_UPDATED,
+                        "params": {"threadId": self.thread_id, "threadSettings": settings},
+                    }
+                )
             if self.switch_message:
                 turn = params.get("turn")
                 turn_id = turn.get("id") if isinstance(turn, dict) else None
@@ -226,12 +202,12 @@ def route_decision(prompt: str):
                 token = token_provider()
             except RuntimeError as exc:
                 return None, f"could not refresh workspace auth: {exc}"
-            return _request_routing_decision(
+            return codex_routing.request_routing_decision(
                 workspace,
                 token,
                 prompt,
                 list(available_models or []),
-                log,
+                log=log,
             )
 
     sess = _Session(