Skip to content
14 changes: 13 additions & 1 deletion src/ucode/agents/codex.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,8 +29,10 @@
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.smart_routing.codex_routing import codex_model_id
from ucode.state import mark_tool_managed, save_state
from ucode.telemetry import agent_version, ucode_version
from ucode.ui import print_warning_err
Expand All @@ -49,6 +51,7 @@
# Shared across agents: one opt-in enables smart routing for every routing-capable
# tool (codex, claude), so a workspace turns it on once.
SMART_ROUTING_STATE_KEY = "smart_routing_enabled"
APP_SERVER_SMART_ROUTING_STARTING_MODEL = "gpt-5.6-luna"

SPEC: ToolSpec = {
"binary": "codex",
Expand Down Expand Up @@ -456,11 +459,20 @@ def launch(state: dict, tool_args: list[str]) -> None:
# path so flag-off launches retain their existing dependencies and behavior.
from ucode.smart_routing import v2 as smart_routing_v2

def _app_server_start_model() -> str:
managed_model = default_model(state)
if managed_model:
return managed_model
models = routing_models(state)
if models:
return codex_model_id(models[0])
return APP_SERVER_SMART_ROUTING_STARTING_MODEL

smart_routing_v2.launch_codex(
state,
tool_args,
binary=binary,
start_model=default_model(state),
start_model=_app_server_start_model(),
render_overlay=render_overlay,
)
if workspace:
Expand Down
2 changes: 2 additions & 0 deletions src/ucode/smart_routing/codex_interposer.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
import time
import uuid
from collections.abc import Callable
from dataclasses import replace
from pathlib import Path

from websockets.asyncio.client import connect
Expand Down Expand Up @@ -90,6 +91,7 @@ def on_tui_frame(self, raw: str) -> str:
if decision is None:
self.log(f"[ROUTE] selection failed; keeping current model: {reason}")
return raw
decision = replace(decision, model=codex_routing.codex_model_id(decision.model))
self.target = decision.model
if self.switch_message_fn is not None:
self.switch_message = self.switch_message_fn(decision.model, decision.rationale)
Expand Down
10 changes: 8 additions & 2 deletions src/ucode/smart_routing/codex_routing.py
Original file line number Diff line number Diff line change
Expand Up @@ -159,7 +159,7 @@ def record(payload, task, decision, requested):
workspace, token, task, available_models, timeout=timeout
),
default_task_label="Codex subagent task",
model_id_mapper=_codex_model_id,
model_id_mapper=codex_model_id,
record_decision=record,
)

Expand Down Expand Up @@ -203,7 +203,13 @@ def _model_strength(model: str) -> tuple[int, int, int, int]:
return major, minor, patch, 1 if not suffix else 0


def _codex_model_id(model: str) -> str:
def codex_model_id(model: str) -> str:
"""Map a UC GPT service ID to Codex's bundled catalog slug.

Codex's bundled GPT catalog owns the model metadata for these aliases,
while the AI Gateway resolves them back to the matching ``system.ai`` service.
Leave non-GPT models unchanged because their metadata comes from the gateway catalog.
"""
tail = model.rsplit("/", 1)[-1]
if tail in {"databricks-gpt-5-2-codex", "databricks-gpt-5-4-nano"}:
return tail
Expand Down
6 changes: 3 additions & 3 deletions src/ucode/smart_routing/v2.py
Original file line number Diff line number Diff line change
Expand Up @@ -468,9 +468,9 @@ def launch_codex(
os.environ[OAUTH_TOKEN_ENV_VAR] = get_databricks_token(workspace, profile)
available_models = _cached_routing_models(state)
if not available_models:
raise RuntimeError(
"Smart routing v2 has no cached Unity Catalog model services; "
"run `ucode configure codex` to refresh them."
print_note(
"Smart routing model metadata is unavailable; starting Codex on gpt-5.6-luna "
"without automatic model switching. Run `ucode configure codex` to enable routing."
)
overlay = render_overlay(
workspace,
Expand Down
15 changes: 15 additions & 0 deletions tests/test_codex_routing.py
Original file line number Diff line number Diff line change
Expand Up @@ -204,6 +204,21 @@ def test_spawn_rewrite_uses_codex_model_id_for_uc_endpoint(monkeypatch):
assert output["hookSpecificOutput"]["updatedInput"]["model"] == "gpt-5.6-luna"


def test_codex_model_id_maps_uc_gpt_models_to_codex_slugs():
expected = {
"system.ai.gpt-5-2": "gpt-5.2",
"system.ai.gpt-5-4": "gpt-5.4",
"system.ai.gpt-5-4-mini": "gpt-5.4-mini",
"system.ai.gpt-5-5": "gpt-5.5",
"system.ai.gpt-5-6-luna": "gpt-5.6-luna",
"system.ai.gpt-5-6-sol": "gpt-5.6-sol",
"system.ai.gpt-5-6-terra": "gpt-5.6-terra",
}
assert {model: codex_routing.codex_model_id(model) for model in expected} == expected
assert codex_routing.codex_model_id("system.ai.gpt-5-6-experimental") == "gpt-5.6-experimental"
assert codex_routing.codex_model_id("system.ai.glm-5-2") == "system.ai.glm-5-2"


def test_spawn_glm_decision_applies_glm_model(monkeypatch):
# GLM is no longer skipped for Codex subagents: a GLM routing decision is
# applied like any other arm.
Expand Down
58 changes: 52 additions & 6 deletions tests/test_codex_smart_routing_v2.py
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@ def test_codex_launch_dispatches_when_flag_enabled(self, monkeypatch):
calls = []
monkeypatch.setenv(v2.ENV_VAR, "1")
monkeypatch.setattr(codex, "default_model", lambda state: "gpt-start")
monkeypatch.setattr(codex, "clear_model_preferences", lambda state: False)

def launch_v2(state, tool_args, **kwargs):
calls.append((state, tool_args, kwargs))
Expand All @@ -77,6 +78,26 @@ def launch_v2(state, tool_args, **kwargs):
)
]

def test_codex_launch_normalizes_cached_bootstrap_model(self, monkeypatch):
calls = []
monkeypatch.setenv(v2.ENV_VAR, "1")
monkeypatch.setattr(codex, "clear_model_preferences", lambda state: False)
monkeypatch.setattr(codex, "default_model", lambda state: None)

def launch_v2(state, tool_args, **kwargs):
calls.append(kwargs)
raise SystemExit(0)

monkeypatch.setattr(v2, "launch_codex", launch_v2)

with pytest.raises(SystemExit):
codex.launch(
{"workspace": WS, "codex_models": ["system.ai.gpt-5-6-luna"]},
[],
)

assert calls[0]["start_model"] == "gpt-5.6-luna"

def test_owns_app_server_interposer_and_tui_lifecycle(self, monkeypatch):
processes = []
interposer_args = {}
Expand Down Expand Up @@ -233,18 +254,41 @@ def test_v2_pre_tool_hook_replaces_existing_ucode_hook(self, tmp_path, monkeypat
assert "--model system.ai.gpt-5-6-sol" in routing_commands[0]
assert "--model old" not in routing_commands[0]

def test_missing_cached_models_blocks_launch(self, monkeypatch):
def test_missing_cached_models_starts_with_bootstrap_model(self, monkeypatch):
monkeypatch.setattr(v2, "get_databricks_token", lambda workspace, profile: "token")
monkeypatch.setattr(codex, "agent_version", lambda binary: "unknown")
monkeypatch.setattr(v2, "_free_port", lambda: 41001)
monkeypatch.setattr(v2, "_wait_for_app_server", lambda port, timeout: True)
monkeypatch.setattr(
v2.subprocess,
"Popen",
lambda *args, **kwargs: type(
"Process",
(),
{
"wait": lambda self, timeout=None: 0,
"terminate": lambda self: None,
"kill": lambda self: None,
},
)(),
)
monkeypatch.setattr(
codex_interposer,
"start_interposer_thread",
lambda *args, **kwargs: (41002, lambda: None),
)

with pytest.raises(RuntimeError, match="ucode configure codex"):
with pytest.raises(SystemExit) as exc:
v2.launch_codex(
{"workspace": WS},
[],
binary="codex",
start_model="gpt-start",
start_model="gpt-5.6-luna",
render_overlay=codex.render_overlay,
)

assert exc.value.code == 0


def test_interposer_startup_failure_is_propagated(monkeypatch):
async def fail_to_serve(*args, **kwargs):
Expand Down Expand Up @@ -368,7 +412,7 @@ def select(prompt):
assert json.loads(output)["params"]["model"] == "claude-opus-4-8"
assert "Task classified as bugfix." in sess.switch_message

def test_shows_routing_notice_when_selected_model_is_already_active(self):
def test_maps_selected_uc_gpt_model_and_shows_routing_notice(self):
def select(_prompt):
return (
codex_interposer.routing.RoutingDecision(
Expand All @@ -387,14 +431,16 @@ def select(_prompt):
)
frame = self._turn_start("system.ai.gpt-5-6-luna")

assert sess.on_tui_frame(frame) == frame
output = sess.on_tui_frame(frame)
assert json.loads(output)["params"]["model"] == "gpt-5.6-luna"
injected = sess.on_engine_frame(self._turn_started("turn-1"))

assert [message["method"] for message in injected] == [
codex_interposer.SETTINGS_UPDATED,
codex_interposer.ITEM_STARTED,
codex_interposer.ITEM_COMPLETED,
]
assert "Selected Model : system.ai.gpt-5-6-luna" in (injected[0]["params"]["item"]["text"])
assert "Selected Model : gpt-5.6-luna" in (injected[1]["params"]["item"]["text"])

def test_routes_first_prompt_to_oss_model(self):
def select(_prompt):
Expand Down
Loading