diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 4b1f555a..40568a6e 100644 --- a/src/ucode/agents/codex.py +++ b/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 diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 6a6ae3b0..661d5db2 100644 --- a/src/ucode/cli.py +++ b/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 diff --git a/src/ucode/config_io.py b/src/ucode/config_io.py index 6a92066d..f67f3f32 100644 --- a/src/ucode/config_io.py +++ b/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 diff --git a/src/ucode/smart_routing/codex_hooks.py b/src/ucode/smart_routing/codex_hooks.py index 6c6a5691..421cabda 100644 --- a/src/ucode/smart_routing/codex_hooks.py +++ b/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 diff --git a/src/ucode/smart_routing/codex_interposer.py b/src/ucode/smart_routing/codex_interposer.py index 9b1a51ea..2841b8f6 100644 --- a/src/ucode/smart_routing/codex_interposer.py +++ b/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( diff --git a/src/ucode/smart_routing/codex_routing.py b/src/ucode/smart_routing/codex_routing.py index f37ecd48..87906f2a 100644 --- a/src/ucode/smart_routing/codex_routing.py +++ b/src/ucode/smart_routing/codex_routing.py @@ -1,31 +1,31 @@ """Databricks AI Gateway routing helpers for Codex sessions and subagents. Codex-specific configuration on top of the shared :mod:`ucode.smart_routing.routing` -core: the frozen ``task_v1`` Codex route arms, the ``spawn_agent`` tool detector, -the Codex model-id translation, and the artifact paths. +core: the workspace-backed ``task_v1`` route options, the ``spawn_agent`` tool +detector, the Codex model-id translation, and the artifact paths. """ from __future__ import annotations +import json import re # Re-exported so tests can patch the shared ``urlopen`` seam via # ``codex_routing.urllib.request`` — the actual call lives in ``routing``, but # Python modules are singletons so patching this name patches the one call site. import urllib.request # noqa: F401 +from collections.abc import Callable from typing import Any from ucode.config_io import APP_DIR from ucode.databricks import get_databricks_token from ucode.smart_routing import routing +from ucode.smart_routing.codex_hooks import routing_models from ucode.smart_routing.routing import RoutingDecision ROUTER_NAME = routing.ROUTER_NAME ROUTING_PATH = routing.ROUTING_PATH REQUEST_TIMEOUT_S = routing.REQUEST_TIMEOUT_S -CODEX_ROUTE_ARMS = ("glm-5-2", "gpt-5-6-sol", "gpt-5-6-luna") -GLM_ROUTE_ARM = "glm-5-2" -GLM_GATEWAY_MODEL = "system.ai.glm-5-2" SPAWN_AGENT_TOOL_SUFFIX = "spawn_agent" CANARY_PATH = APP_DIR / "codex-smart-routing-canary.json" AUDIT_PATH = APP_DIR / "codex-smart-routing-audit.jsonl" @@ -50,8 +50,8 @@ def route_launch_model(state: dict, tool_args: list[str]): if task is None: return None, None workspace = state.get("workspace") - models = state.get("codex_models") - if not isinstance(workspace, str) or not isinstance(models, list): + models = routing_models(state) + if not isinstance(workspace, str) or not models: return None, "workspace model metadata is unavailable" try: token = get_databricks_token(workspace, state.get("profile")) @@ -103,29 +103,36 @@ def request_routing_decision( available_models: list[str], *, timeout: float = REQUEST_TIMEOUT_S, + log: Callable[[str], None] | None = None, ) -> tuple[RoutingDecision | None, str | None]: """Ask the workspace ``task_v1`` router for a servable Codex model.""" - candidates = _routing_candidates(available_models) - missing = [ - arm for arm in CODEX_ROUTE_ARMS if arm not in {_normalize_model(m) for m in candidates} - ] - if missing: - return None, f"required Codex routing models are unavailable: {', '.join(missing)}" - + available = {_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": task}, + "route_selector": {"router_name": ROUTER_NAME}, + } + url = workspace.rstrip("/") + ROUTING_PATH + log(f"[ROUTE] request POST {url}: {json.dumps(payload, separators=(',', ':'))}") return routing.select_route( workspace, token, task, - [(arm, "codex") for arm in CODEX_ROUTE_ARMS], - lambda raw_model: resolve_routed_model(raw_model, candidates), + route_options, + lambda raw_model: available.get(_normalize_model(raw_model)), timeout=timeout, ) def resolve_routed_model(raw_model: str, available_models: list[str]) -> str | None: """Map a ``task_v1`` arm to a model the configured workspace can serve.""" - candidates = _routing_candidates(available_models) - normalized = {_normalize_model(model): model for model in candidates} + normalized = {_normalize_model(model): model for model in available_models} return normalized.get(_normalize_model(raw_model)) @@ -180,13 +187,6 @@ def clear_routing_artifacts() -> None: routing.clear_artifacts((CANARY_PATH, AUDIT_PATH, DECISIONS_PATH)) -def _routing_candidates(models: list[str]) -> list[str]: - candidates = [model for model in models if isinstance(model, str) and model] - if GLM_ROUTE_ARM not in {_normalize_model(model) for model in candidates}: - candidates.append(GLM_GATEWAY_MODEL) - return candidates - - def _parse_gpt(model: str) -> tuple[int, int, int, str] | None: match = _GPT_RE.fullmatch(_normalize_model(model)) if not match: @@ -207,8 +207,6 @@ def _codex_model_id(model: str) -> str: tail = model.rsplit("/", 1)[-1] if tail in {"databricks-gpt-5-2-codex", "databricks-gpt-5-4-nano"}: return tail - if _normalize_model(model) == GLM_ROUTE_ARM: - return GLM_GATEWAY_MODEL if model.startswith("system.ai."): bare = model.removeprefix("system.ai.") elif tail.startswith("databricks-"): @@ -217,7 +215,7 @@ def _codex_model_id(model: str) -> str: return model match = _GPT_RE.fullmatch(bare) if not match: - return bare + return model major, minor, patch, suffix = match.groups() version = major if minor is not None: diff --git a/src/ucode/smart_routing/routing.py b/src/ucode/smart_routing/routing.py index 2fecacd5..d2f578de 100644 --- a/src/ucode/smart_routing/routing.py +++ b/src/ucode/smart_routing/routing.py @@ -23,6 +23,36 @@ ROUTER_NAME = "task_v1" ROUTING_PATH = "/ai-gateway/routing/v1/routes:select" REQUEST_TIMEOUT_S = 30.0 +SUBAGENT_ROUTING_DISCLAIMER = ( + "Spawned subagents are routed independently based on their own complexity." +) + + +def format_switch_message(model: str, reason: str | None) -> str: + """Format the first-prompt routed-model notice.""" + lines = [ + "Using Unity Gateway Smart Router.", + f"Selected Model : {model}", + *([f"Reason : {reason}"] if reason else []), + SUBAGENT_ROUTING_DISCLAIMER, + ] + return _format_box(lines) + + +def format_subagent_message(model: str, reason: str | None) -> str: + """Format a routed-subagent notice without the first-prompt disclaimer.""" + lines = [ + "Using Unity Gateway Smart Router - Subagent", + f"Selected Model : {model}", + *([f"Reason : {reason}"] if reason else []), + ] + return _format_box(lines) + + +def _format_box(lines: list[str]) -> str: + width = max(len(line) for line in lines) + border = "─" * (width + 2) + return "\n".join([f"┌{border}┐", *(f"│ {line:<{width}} │" for line in lines), f"└{border}┘"]) @dataclass(frozen=True) @@ -33,18 +63,16 @@ class RoutingDecision: raw_model: str rationale: str = "" - def display_message(self, model_label: str | None = None) -> str: - """User-facing "Using Smart Routing" line, with the router's rationale. + def display_message(self, model_label: str | None = None, *, subagent: bool = False) -> str: + """Return the boxed smart-routing notice with the router's rationale. Used by both the launch-time notice and the subagent-routing hook so the "what" (model) and the "why" (rationale) are surfaced consistently. ``model_label`` overrides the shown model id (e.g. a harness-translated - id); defaults to ``model``. The rationale is appended when present. + id); defaults to ``model``. """ - message = f"Using Smart Routing. Routing to {model_label or self.model}." - if self.rationale: - message += f" {self.rationale}" - return message + formatter = format_subagent_message if subagent else format_switch_message + return formatter(model_label or self.model, self.rationale) @dataclass(frozen=True) @@ -250,7 +278,10 @@ def route_spawn_tool( # Surface the router's rationale in BOTH the systemMessage (the line the # harness shows the user) and permissionDecisionReason — the "why", not just # the "what". The shown model is the harness-translated id (routed_model). - routing_message = route.decision.display_message(model_label=route.routed_model) + routing_message = route.decision.display_message( + model_label=route.routed_model, + subagent=True, + ) output: dict[str, Any] = { "hookEventName": "PreToolUse", "permissionDecision": "allow", diff --git a/src/ucode/smart_routing/v2.py b/src/ucode/smart_routing/v2.py index 09d5dcb4..c4f99ebc 100644 --- a/src/ucode/smart_routing/v2.py +++ b/src/ucode/smart_routing/v2.py @@ -16,7 +16,7 @@ import tomlkit -from ucode.config_io import APP_DIR, read_json_safe, write_json_file +from ucode.config_io import APP_DIR, read_json_safe, read_toml_safe, write_json_file from ucode.constants import LOOPBACK_HOST from ucode.databricks import ( build_auth_token_argv, @@ -29,6 +29,7 @@ sync_first_prompt_hook, sync_smart_routing_hooks, ) +from ucode.smart_routing.codex_hooks import merge_pre_tool_use_hooks, routing_models from ucode.ui import print_note ENV_VAR = "ENABLE_SMART_ROUTING_V2" @@ -80,17 +81,12 @@ def _wait_for_app_server(port: int, timeout: float) -> bool: return False -def format_routing_notice(model: str, reason: str | None, *, title: str | None = None) -> str: - lines = [ - *([title] if title else []), - "Using Unity Gateway Smart Router.", - f"Selected Model : {model}", - ] - if reason: - lines.append(f"Reason : {reason}") - width = max(map(len, lines)) - border = "─" * (width + 2) - return "\n".join([f"┌{border}┐", *(f"│ {line:<{width}} │" for line in lines), f"└{border}┘"]) +def _switch_message(model: str, reason: str) -> str: + return routing.format_switch_message(model, reason) + + +def format_routing_notice(model: str, reason: str | None) -> str: + return routing.format_switch_message(model, reason or "") def _canonical_claude_model_id(model: str) -> str: @@ -249,10 +245,9 @@ def route_claude_pre_tool_use( route.decision, route.routed_model, ) - routing_message = format_routing_notice( + routing_message = routing.format_subagent_message( route.routed_model, route.decision.rationale, - title="Subagent Smart Routing", ) updated_input = { **{key: value for key, value in route.tool_input.items() if key != "model"}, @@ -401,6 +396,11 @@ def _toml_value(value: str | int | float | bool | list[object] | dict[str, objec item = tomlkit.inline_table() item.update(value) return item.as_string() + if isinstance(value, list) and any(isinstance(entry, dict) for entry in value): + wrapper = tomlkit.inline_table() + wrapper["value"] = value + rendered = wrapper.as_string() + return rendered.removeprefix("{value = ").removesuffix("}") return tomlkit.item(value).as_string() @@ -409,12 +409,12 @@ def _codex_config_args(overlay: dict) -> list[str]: for key, value in overlay.items(): # This is Codex's AI Gateway transport definition, not Unity Catalog # Model Provider Service support; smart routing still cannot use --provider. - if key == "model_providers" and isinstance(value, dict): + if key in {"hooks", "model_providers"} and isinstance(value, dict): for provider_name, provider_config in value.items(): args.extend( [ "--config", - f"model_providers.{provider_name}={_toml_value(provider_config)}", + f"{key}.{provider_name}={_toml_value(provider_config)}", ] ) else: @@ -425,12 +425,25 @@ def _codex_config_args(overlay: dict) -> list[str]: # TODO: Replace with /codex/v1/models once /codex/v1/models can send GPT models as well. def _cached_routing_models(state: dict) -> list[str]: """Return the persisted UC model-service ids usable by 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)) + return routing_models(state) + + +def _codex_home_config_path() -> Path: + codex_home = os.environ.get("CODEX_HOME") + if codex_home: + return Path(codex_home).expanduser() / "config.toml" + return Path.home() / ".codex" / "config.toml" + + +def _v2_pre_tool_use_hooks(state: dict, available_models: list[str]) -> list[dict]: + doc = read_toml_safe(_codex_home_config_path()) + configured_hooks = doc.get("hooks") + existing = configured_hooks.get("PreToolUse") if isinstance(configured_hooks, dict) else None + return merge_pre_tool_use_hooks( + existing if isinstance(existing, list) else [], + state, + available_models=available_models, + ) def launch_codex( @@ -465,6 +478,9 @@ def launch_codex( state.get("profile"), use_pat=bool(state.get("use_pat")), ) + overlay["hooks"] = { + "PreToolUse": _v2_pre_tool_use_hooks(state, available_models), + } config_args = _codex_config_args(overlay) app_port = _free_port() app_server_url = _loopback_websocket_url(app_port) diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index 9560edb2..b3a37aae 100644 --- a/tests/test_agent_codex.py +++ b/tests/test_agent_codex.py @@ -265,6 +265,7 @@ def test_smart_routing_writes_profile_scoped_hooks(self, tmp_path, monkeypatch): "workspace": WS, "profile": "prod", "codex_models": ["databricks-gpt-5", "databricks-gpt-5-5"], + "oss_models": ["system.ai.glm-5-2"], codex.SMART_ROUTING_STATE_KEY: True, } ) @@ -282,6 +283,7 @@ def test_smart_routing_writes_profile_scoped_hooks(self, tmp_path, monkeypatch): assert "--host https://example.databricks.com" in route_command assert "--profile prod" in route_command assert "--model databricks-gpt-5-5" in route_command + assert "--model system.ai.glm-5-2" in route_command def test_provider_launch_removes_routing_hooks(self, tmp_path, monkeypatch): config_path = tmp_path / ".codex" / "ucode.config.toml" @@ -432,6 +434,27 @@ def fail(*args, **kwargs): assert decision is None assert error is None + def test_route_launch_model_includes_codex_and_oss_models(self, monkeypatch): + captured = {} + monkeypatch.setattr(codex_routing, "get_databricks_token", lambda *args: "token") + + def request(workspace, token, task, models): + captured["models"] = models + return None, "expected test stop" + + monkeypatch.setattr(codex_routing, "request_routing_decision", request) + + codex_routing.route_launch_model( + { + "workspace": WS, + "codex_models": ["system.ai.gpt-5-6-sol"], + "oss_models": ["system.ai.glm-5-2"], + }, + ["Fix the parser"], + ) + + assert captured["models"] == ["system.ai.gpt-5-6-sol", "system.ai.glm-5-2"] + class TestCodexRemoveLegacyProfile: def test_drops_provider_block_on_modern_path(self, tmp_path, monkeypatch): @@ -560,6 +583,14 @@ def test_default_model_falls_back_to_first_when_no_versioned_gpt(self): models = ["system.ai.gpt-oss-120b", "system.ai.gpt-oss-20b"] assert codex.default_model({"codex_models": models}) == "system.ai.gpt-oss-120b" + def test_default_model_includes_oss_models(self): + state = { + "codex_models": [], + "oss_models": ["system.ai.kimi-k3", "system.ai.gpt-oss-120b"], + } + + assert codex.default_model(state) == "system.ai.gpt-oss-120b" + def test_default_model_prefers_versioned_gpt_over_oss(self): # When both versioned and OSS models are present, the versioned one wins. models = ["system.ai.gpt-oss-120b", "system.ai.gpt-5"] diff --git a/tests/test_claude_routing.py b/tests/test_claude_routing.py index d30ee669..e8abcc8c 100644 --- a/tests/test_claude_routing.py +++ b/tests/test_claude_routing.py @@ -157,9 +157,10 @@ def test_spawn_rewrite_injects_routed_model(monkeypatch): # The rationale is surfaced in the systemMessage (shown to the user), not # only in permissionDecisionReason. The model field is the short family # name ("opus") that Claude Code's Agent tool schema accepts. - assert output["systemMessage"] == ( - "Using Smart Routing. Routing to opus. Deep exploration needs the strongest model." + expected_message = claude_routing.routing.format_subagent_message( + "opus", "Deep exploration needs the strongest model." ) + assert output["systemMessage"] == expected_message assert hook["permissionDecision"] == "allow" assert hook["updatedInput"] == { "subagent_type": "Explore", @@ -167,9 +168,7 @@ def test_spawn_rewrite_injects_routed_model(monkeypatch): "description": "explore", "model": "opus", } - assert hook["permissionDecisionReason"] == ( - "Using Smart Routing. Routing to opus. Deep exploration needs the strongest model." - ) + assert hook["permissionDecisionReason"] == expected_message def test_task_tool_alias_is_routed(monkeypatch): diff --git a/tests/test_claude_smart_routing_v2.py b/tests/test_claude_smart_routing_v2.py index b5b4e803..cbc4af43 100644 --- a/tests/test_claude_smart_routing_v2.py +++ b/tests/test_claude_smart_routing_v2.py @@ -310,10 +310,9 @@ def fake_select(workspace, token, task, route_options, resolve, **kwargs): assert updated_input["subagent_type"] == v2._routed_claude_agent_name( "system.ai.claude-opus-4-8" ) - expected_message = v2.format_routing_notice( + expected_message = routing.format_subagent_message( "system.ai.claude-opus-4-8", "", - title="Subagent Smart Routing", ) assert output["systemMessage"] == expected_message assert output["hookSpecificOutput"]["permissionDecisionReason"] == expected_message diff --git a/tests/test_cli.py b/tests/test_cli.py index a4f90a80..19efb486 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -297,10 +297,10 @@ def test_enabled_codex_launch_uses_routed_root_model(self): assert result.exit_code == 0, result.output assert mock_configure.call_args.args[2] == "databricks-gpt-5-5" # The launch notice surfaces both the routed model and the rationale. - assert ( - "Using Smart Routing. Routing to databricks-gpt-5-5. Cross-cutting refactor." - in _strip_ansi(result.output) - ) + output = _strip_ansi(result.output) + assert "Using Unity Gateway Smart Router." in output + assert "Selected Model : databricks-gpt-5-5" in output + assert "Reason : Cross-cutting refactor." in output def test_claude_v2_skips_legacy_prelaunch_routing(self, monkeypatch): monkeypatch.setenv("ENABLE_SMART_ROUTING_V2", "1") @@ -2305,6 +2305,29 @@ def test_uc_models_used_without_legacy_fallback(self, monkeypatch): assert legacy_called == [] assert "uc_enabled" not in state + def test_codex_only_configure_persists_discovered_oss_models(self, monkeypatch): + cli_mod, *_ = self._stub_deps(monkeypatch, pat_token="dapi-pat") + monkeypatch.setattr( + cli_mod, + "discover_model_services", + lambda w, t: ( + {}, + ["system.ai.gpt-5-6-sol"], + [], + ["system.ai.glm-5-2"], + None, + ), + ) + + state = cli_mod.configure_shared_state( + self.WS, + profile="DEFAULT", + tools=["codex"], + ) + + assert state["codex_models"] == ["system.ai.gpt-5-6-sol"] + assert state["oss_models"] == ["system.ai.glm-5-2"] + def _stub_with_fable(self, monkeypatch): cli_mod, *_ = self._stub_deps(monkeypatch, pat_token="dapi-pat") monkeypatch.setattr( diff --git a/tests/test_codex_routing.py b/tests/test_codex_routing.py index f5d84c0c..8973eaf1 100644 --- a/tests/test_codex_routing.py +++ b/tests/test_codex_routing.py @@ -6,10 +6,24 @@ import urllib.error from ucode.smart_routing import codex_routing +from ucode.smart_routing.codex_hooks import routing_models WS = "https://example.databricks.com" +def test_routing_models_combines_and_deduplicates_codex_and_oss_models(): + assert routing_models( + { + "codex_models": ["system.ai.gpt-5-6-sol", "system.ai.gpt-oss-120b"], + "oss_models": ["system.ai.gpt-oss-120b", "system.ai.glm-5-2"], + } + ) == [ + "system.ai.gpt-5-6-sol", + "system.ai.gpt-oss-120b", + "system.ai.glm-5-2", + ] + + class _Response: def __init__(self, payload: dict): self.payload = payload @@ -24,7 +38,7 @@ def read(self) -> bytes: return json.dumps(self.payload).encode("utf-8") -def test_routes_with_task_v1_codex_menu(monkeypatch): +def test_routes_with_models_from_stored_state(monkeypatch): captured = {} task = "Refactor the parser" + "x" * 5000 @@ -59,9 +73,8 @@ def fake_urlopen(request, timeout): assert captured["headers"]["Authorization"] == "Bearer token" assert captured["body"] == { "route_options": [ - {"model": "glm-5-2", "harness": "codex"}, - {"model": "gpt-5-6-sol", "harness": "codex"}, {"model": "gpt-5-6-luna", "harness": "codex"}, + {"model": "gpt-5-6-sol", "harness": "codex"}, ], "task": {"prompt": task}, "route_selector": {"router_name": "task_v1"}, @@ -80,7 +93,7 @@ def test_router_model_is_not_substituted_when_exact_model_is_unavailable(): def test_glm_maps_to_databricks_gateway_model(): model = codex_routing.resolve_routed_model( "glm-5-2", - ["system.ai.gpt-5-6-luna", "system.ai.gpt-5-6-sol"], + ["system.ai.gpt-5-6-luna", "system.ai.gpt-5-6-sol", "system.ai.glm-5-2"], ) assert model == "system.ai.glm-5-2" @@ -146,9 +159,12 @@ def test_spawn_rewrite_preserves_original_input(monkeypatch): hook = output["hookSpecificOutput"] # The rationale is surfaced in BOTH the systemMessage (shown to the user) and # permissionDecisionReason, so the "why" is visible, not just the "what". - assert output["systemMessage"] == ( - "Using Smart Routing. Routing to gpt-5.5. Review needs deeper reasoning." + expected_message = codex_routing.routing.format_subagent_message( + "gpt-5.5", "Review needs deeper reasoning." ) + assert output["systemMessage"] == expected_message + assert "Using Unity Gateway Smart Router - Subagent" in expected_message + assert codex_routing.routing.SUBAGENT_ROUTING_DISCLAIMER not in expected_message assert hook["permissionDecision"] == "allow" assert hook["updatedInput"] == { "task_name": "reviewer", @@ -156,9 +172,7 @@ def test_spawn_rewrite_preserves_original_input(monkeypatch): "fork": False, "model": "gpt-5.5", } - assert hook["permissionDecisionReason"] == ( - "Using Smart Routing. Routing to gpt-5.5. Review needs deeper reasoning." - ) + assert hook["permissionDecisionReason"] == expected_message def test_spawn_rewrite_uses_codex_model_id_for_uc_endpoint(monkeypatch): @@ -184,7 +198,9 @@ def test_spawn_rewrite_uses_codex_model_id_for_uc_endpoint(monkeypatch): available_models=["system.ai.gpt-5-6-luna"], ) - assert output["systemMessage"] == "Using Smart Routing. Routing to gpt-5.6-luna." + assert output["systemMessage"] == codex_routing.routing.format_subagent_message( + "gpt-5.6-luna", "" + ) assert output["hookSpecificOutput"]["updatedInput"]["model"] == "gpt-5.6-luna" @@ -210,11 +226,15 @@ def test_spawn_glm_decision_applies_glm_model(monkeypatch): }, workspace=WS, token="token", - available_models=["system.ai.gpt-5-6-luna", "system.ai.gpt-5-6-sol"], + available_models=[ + "system.ai.gpt-5-6-luna", + "system.ai.gpt-5-6-sol", + "system.ai.glm-5-2", + ], ) assert output["hookSpecificOutput"]["updatedInput"]["model"] == "system.ai.glm-5-2" - assert "Using Smart Routing. Routing to system.ai.glm-5-2." in output["systemMessage"] + assert "Selected Model : system.ai.glm-5-2" in output["systemMessage"] def test_non_spawn_tool_has_no_opinion(): diff --git a/tests/test_codex_smart_routing_v2.py b/tests/test_codex_smart_routing_v2.py index 3e4e7ef7..8c90a621 100644 --- a/tests/test_codex_smart_routing_v2.py +++ b/tests/test_codex_smart_routing_v2.py @@ -5,7 +5,7 @@ import pytest from ucode.agents import codex -from ucode.smart_routing import codex_interposer, v2 +from ucode.smart_routing import codex_interposer, codex_routing, v2 WS = "https://example.databricks.com" @@ -39,11 +39,12 @@ def test_smart_routing_switch_message_is_boxed(): message = v2.format_routing_notice("model-x", "Because X.") assert message == ( - "┌───────────────────────────────────┐\n" - "│ Using Unity Gateway Smart Router. │\n" - "│ Selected Model : model-x │\n" - "│ Reason : Because X. │\n" - "└───────────────────────────────────┘" + "┌───────────────────────────────────────────────────────────────────────────┐\n" + "│ Using Unity Gateway Smart Router. │\n" + "│ Selected Model : model-x │\n" + "│ Reason : Because X. │\n" + "│ Spawned subagents are routed independently based on their own complexity. │\n" + "└───────────────────────────────────────────────────────────────────────────┘" ) @@ -145,11 +146,20 @@ def start_interposer(*args, **kwargs): 'model="gpt-start"', "--config", ] - assert processes[0].argv[8:] == [ + assert processes[0].argv[7].startswith("model_providers.ucode-databricks={") + assert processes[0].argv[8] == "--config" + hook_override = processes[0].argv[9] + assert hook_override.startswith("hooks.PreToolUse=[{") + assert 'matcher = "Agent|.*spawn_agent$"' in hook_override + assert "codex-router-hook route-subagent" in hook_override + assert f"--host {WS}" in hook_override + assert "--profile myprof" in hook_override + assert "--model system.ai.gpt-5-6-sol" in hook_override + assert "--model system.ai.glm-5-2" in hook_override + assert processes[0].argv[10:] == [ "--listen", "ws://127.0.0.1:41001", ] - assert processes[0].argv[7].startswith("model_providers.ucode-databricks={") assert processes[0].kwargs["env"][v2.OAUTH_TOKEN_ENV_VAR] == "token-1" assert processes[0].kwargs["env"]["CODEX_HOME"] == "/user/codex-home" assert processes[1].argv == [ @@ -173,6 +183,28 @@ def start_interposer(*args, **kwargs): assert stopped == [True] assert processes[0].terminated is True + def test_v2_pre_tool_hook_preserves_user_hooks(self, tmp_path, monkeypatch): + codex_home = tmp_path / ".codex" + codex_home.mkdir() + (codex_home / "config.toml").write_text( + "[[hooks.PreToolUse]]\n" + 'matcher = "Bash"\n' + "[[hooks.PreToolUse.hooks]]\n" + 'type = "command"\n' + 'command = "user-policy"\n', + encoding="utf-8", + ) + monkeypatch.setenv("CODEX_HOME", str(codex_home)) + + configured = v2._v2_pre_tool_use_hooks( + {"workspace": WS, "profile": "myprof"}, + ["system.ai.gpt-5-6-sol"], + ) + + assert configured[0]["hooks"][0]["command"] == "user-policy" + assert configured[1]["matcher"] == "Agent|.*spawn_agent$" + assert "--model system.ai.gpt-5-6-sol" in configured[1]["hooks"][0]["command"] + def test_missing_cached_models_blocks_launch(self, monkeypatch): monkeypatch.setattr(v2, "get_databricks_token", lambda workspace, profile: "token") @@ -308,6 +340,58 @@ 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 select(_prompt): + return ( + codex_interposer.routing.RoutingDecision( + model="system.ai.gpt-5-6-luna", + raw_model="gpt-5-6-luna", + rationale="Trivial task.", + ), + None, + ) + + sess = codex_interposer._Session( + None, + log=lambda _m: None, + route_decision=select, + switch_message_fn=v2._switch_message, + ) + frame = self._turn_start("system.ai.gpt-5-6-luna") + + assert sess.on_tui_frame(frame) == frame + injected = sess.on_engine_frame(self._turn_started("turn-1")) + + assert [message["method"] for message in injected] == [ + codex_interposer.ITEM_STARTED, + codex_interposer.ITEM_COMPLETED, + ] + assert "Selected Model : system.ai.gpt-5-6-luna" in (injected[0]["params"]["item"]["text"]) + + def test_routes_first_prompt_to_oss_model(self): + def select(_prompt): + return ( + codex_interposer.routing.RoutingDecision( + model="system.ai.glm-5-2", + raw_model="glm-5-2", + rationale="Short isolated task.", + ), + None, + ) + + sess = codex_interposer._Session( + None, + log=lambda _m: None, + available_models=["system.ai.gpt-5-6-sol", "system.ai.glm-5-2"], + route_decision=select, + switch_message_fn=v2._switch_message, + ) + + output = sess.on_tui_frame(self._turn_start("system.ai.gpt-5-6-sol")) + + assert json.loads(output)["params"]["model"] == "system.ai.glm-5-2" + assert "Selected Model : system.ai.glm-5-2" in sess.switch_message + def test_router_failure_keeps_original_model(self): sess = codex_interposer._Session( None, @@ -323,12 +407,13 @@ def test_routing_request_uses_models_prompt_and_same_token(monkeypatch): captured = {} logged = [] - def select_route(workspace, token, task, route_options, resolve): + def select_route(workspace, token, task, route_options, resolve, *, timeout): captured.update( workspace=workspace, token=token, task=task, route_options=list(route_options), + timeout=timeout, ) return ( codex_interposer.routing.RoutingDecision( @@ -339,9 +424,9 @@ def select_route(workspace, token, task, route_options, resolve): None, ) - monkeypatch.setattr(codex_interposer.routing, "select_route", select_route) + monkeypatch.setattr(codex_routing.routing, "select_route", select_route) - decision, reason = codex_interposer._request_routing_decision( + decision, reason = codex_routing.request_routing_decision( WS, "same-oauth-token", "Fix the parser", @@ -351,7 +436,7 @@ def select_route(workspace, token, task, route_options, resolve): "system.ai.gpt-5-6-luna", "system.ai.glm-5-2", ], - logged.append, + log=logged.append, ) assert reason is None @@ -360,6 +445,7 @@ def select_route(workspace, token, task, route_options, resolve): "workspace": WS, "token": "same-oauth-token", "task": "Fix the parser", + "timeout": codex_routing.REQUEST_TIMEOUT_S, "route_options": [ ("kimi-k3-neo", "codex"), ("gpt-5-6-sol", "codex"), diff --git a/tests/test_config_io.py b/tests/test_config_io.py index d8852592..f2af8559 100644 --- a/tests/test_config_io.py +++ b/tests/test_config_io.py @@ -271,6 +271,19 @@ def test_value_with_equals(self, tmp_path): p.write_text("URL=http://example.com?a=1\n", encoding="utf-8") assert parse_dotenv(p) == {"URL": "http://example.com?a=1"} + def test_preserves_trailing_spaces_in_value(self, tmp_path): + p = tmp_path / ".env" + p.write_text("TOKEN=abc123 \n", encoding="utf-8") + result = parse_dotenv(p) + assert result == {"TOKEN": "abc123 "} + assert result["TOKEN"].endswith(" ") + + def test_preserves_trailing_spaces_around_delimiter(self, tmp_path): + p = tmp_path / ".env" + # Leading whitespace after "=" is still trimmed, trailing spaces survive. + p.write_text("TOKEN = abc123 \n", encoding="utf-8") + assert parse_dotenv(p) == {"TOKEN": "abc123 "} + # --------------------------------------------------------------------------- # deep_merge_dict