Route only explicit Codex and Claude invocations through smart routing

src/ucode/agents/__init__.py

@@ -41,6 +41,7 @@
 )
 
 from . import claude, codex, copilot, gemini, opencode, pi
+from .args import LaunchOptions as LaunchOptions
 from .args import explicit_model_arg_value as explicit_model_arg_value
 
 _MODULES = {
@@ -415,8 +416,7 @@ def configure_tool(
     elif tool == "claude":
         # A Model Provider Service routes by header and pins no Databricks
         # model, so the usual "model required" guard doesn't apply to claude.
-        # `custom_model` (from `ucode claude --model`) likewise supplies the model.
-        if not model and not provider and not custom_model:
+        if not model and not provider:
             raise RuntimeError(f"A {tool} model must be selected before configuration.")
         result = claude.write_tool_config(
             state,
@@ -447,8 +447,14 @@ def configure_tool(
     return result
 
 
-def launch(tool: str, state: dict, tool_args: list[str]) -> None:
-    _MODULES[tool].launch(state, tool_args)
+def launch(
+    tool: str,
+    state: dict,
+    tool_args: list[str],
+    *,
+    options: LaunchOptions,
+) -> None:
+    _MODULES[tool].launch(state, tool_args, options=options)
 
 
 def check_gateway_endpoint(state: dict, tool: str) -> bool:

src/ucode/agents/args.py

@@ -2,6 +2,18 @@
 
 from __future__ import annotations
 
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class LaunchOptions:
+    """Invocation-scoped options shared by agent launchers."""
+
+    launch_smart_routing: bool = False
+    # Claude's --model is consumed by ucode, so it must be passed separately for this launch.
+    # Codex keeps --model in the forwarded tool arguments instead.
+    claude_launch_model: str | None = None
+
 
 def explicit_model_arg_value(tool_args: list[str]) -> str | None:
     """Return the last model selected before the harness's ``--`` separator."""

src/ucode/agents/claude.py

@@ -48,13 +48,12 @@
     remove_smart_routing_hooks,
     sync_smart_routing_hooks,
 )
-from ucode.smart_routing.claude_routing import CLAUDE_VALUE_OPTIONS
 from ucode.state import MANAGED_OVERLAY_KEY, get_provider_service, mark_tool_managed, save_state
 from ucode.telemetry import agent_version, ucode_version
 from ucode.tracing import tracing_env
 from ucode.ui import print_note, print_success, print_warning
 
-from .args import has_explicit_model_arg
+from .args import LaunchOptions, has_explicit_model_arg
 
 GATEWAY_MODEL_DISCOVERY_ENV_VAR = "ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY"
 CLAUDE_CONFIG_DIR = Path.home() / ".claude"
@@ -77,26 +76,6 @@
 
 # Retained only to identify and remove state written by the legacy persisted opt-in.
 SMART_ROUTING_STATE_KEY = smart_routing_v2.LEGACY_STATE_KEY
-CLAUDE_NONINTERACTIVE_FLAGS = frozenset(
-    {"-p", "--print", "--bg", "--background", "--cloud", "-h", "--help", "-v", "--version"}
-)
-CLAUDE_SUBCOMMANDS = frozenset(
-    {"agents", "auth", "config", "doctor", "install", "mcp", "plugin", "setup-token", "update"}
-)
-CLAUDE_OPTIONAL_VALUE_OPTIONS = frozenset(
-    {
-        "-d",
-        "--debug",
-        "--from-pr",
-        "--prompt-suggestions",
-        "-r",
-        "--resume",
-        "--remote-control",
-        "--teleport",
-        "-w",
-        "--worktree",
-    }
-)
 
 
 def _parse_version(value: str) -> tuple[int, int, int] | None:
@@ -425,19 +404,6 @@ def render_overlay(
     _ = model  # API stability; no longer pinned via env.
     if route_root_model:
         env["ANTHROPIC_MODEL"] = route_root_model
-    # `ucode claude --model <id>` pins an arbitrary Databricks model id for this launch. It CANNOT
-    # go in ANTHROPIC_MODEL: Claude Code validates that value client-side against the models it knows
-    # (via the apiKeyHelper auth path ucode uses) and rejects a raw id with "may not exist ... run
-    # /model". The family-alias vars (ANTHROPIC_DEFAULT_*_MODEL) are passed through unchecked, so pin
-    # the id into all of them — a raw id carries no signal of its family (opus/sonnet/haiku), and
-    # overriding every slot makes the model take effect no matter which one Claude Code resolves
-    # (root session, a tier switch, or a subagent). Wins over the discovered-model aliases below.
-    if custom_model and not provider:
-        env["ANTHROPIC_DEFAULT_OPUS_MODEL"] = custom_model
-        env["ANTHROPIC_DEFAULT_SONNET_MODEL"] = custom_model
-        env["ANTHROPIC_DEFAULT_HAIKU_MODEL"] = custom_model
-        if fable_enabled:
-            env["ANTHROPIC_DEFAULT_FABLE_MODEL"] = custom_model
     # A Bedrock-backed provider needs its provider-side ids pinned verbatim
     # (Claude Code's canonical names aren't routable there). These come from the
     # service's targets, already de-duped to one id per family upstream.
@@ -1195,45 +1161,13 @@ def _original_launch_model(state: dict) -> str | None:
     return default_model(state)
 
 
-def _has_launch_model_override(state: dict) -> bool:
-    override = state.get("_claude_launch_model")
-    return isinstance(override, str) and bool(override.strip())
-
-
 def _has_provider_launch(state: dict) -> bool:
     transient = state.get("_claude_launch_provider")
     return (isinstance(transient, str) and bool(transient.strip())) or bool(
         get_provider_service(state, "claude")
     )
 
 
-def _uses_interactive_tui(tool_args: list[str]) -> bool:
-    if any(arg in CLAUDE_NONINTERACTIVE_FLAGS for arg in tool_args):
-        return False
-
-    index = 0
-    while index < len(tool_args):
-        arg = tool_args[index]
-        if arg == "--":
-            return True
-        if arg in CLAUDE_VALUE_OPTIONS:
-            index += 2
-            continue
-        if arg in CLAUDE_OPTIONAL_VALUE_OPTIONS:
-            if index + 1 < len(tool_args) and not tool_args[index + 1].startswith("-"):
-                index += 2
-            else:
-                index += 1
-            continue
-        if arg.startswith("-"):
-            index += 1
-            continue
-        # Claude accepts an initial prompt positionally and still opens the TUI.
-        # Keep prompts inside the V2 PTY while bypassing utility subcommands.
-        return arg not in CLAUDE_SUBCOMMANDS
-    return True
-
-
 def _launch_model_args(tool_args: list[str], launch_model: str | None) -> list[str]:
     if not launch_model or has_explicit_model_arg(tool_args):
         return []
@@ -1377,27 +1311,24 @@ def _launch_relayed(state: dict, binary: str, tool_args: list[str]) -> None:
     raise SystemExit(returncode)
 
 
-def launch(state: dict, tool_args: list[str]) -> None:
+def launch(
+    state: dict,
+    tool_args: list[str],
+    *,
+    options: LaunchOptions,
+) -> None:
     binary = SPEC["binary"]
     workspace = state.get("workspace")
     if state.get("claude_relayed"):
         _launch_relayed(state, binary, tool_args)
         return
-    first_prompt_routing = (
-        smart_routing_v2.enabled()
-        and bool(workspace)
-        and not _has_launch_model_override(state)
-        and not has_explicit_model_arg(tool_args)
-        and not _has_provider_launch(state)
-        and _uses_interactive_tui(tool_args)
-    )
     # Smart routing v2 needs Unix PTY support, which Windows does not provide.
-    if first_prompt_routing and os.name == "nt":
+    if options.launch_smart_routing and os.name == "nt":
         raise RuntimeError(
             "Smart routing in Claude Code is currently not supported on Windows. "
             "Please use Codex or disable smart routing."
         )
-    if first_prompt_routing:
+    if options.launch_smart_routing:
         smart_routing_v2.launch_claude(
             state,
             tool_args,
@@ -1419,6 +1350,8 @@ def launch(state: dict, tool_args: list[str]) -> None:
         os.environ["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] = "1"
     if workspace:
         os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile"))
+    if options.claude_launch_model:
+        os.environ["ANTHROPIC_MODEL"] = options.claude_launch_model
     exec_or_spawn(_build_claude_argv(binary, tool_args))
 
 

src/ucode/agents/codex.py

@@ -5,9 +5,6 @@
 import copy
 import os
 import re
-import subprocess
-import sys
-import time
 from collections.abc import Callable
 from pathlib import Path
 
@@ -52,6 +49,8 @@
 from ucode.telemetry import agent_version, ucode_version
 from ucode.ui import print_warning_err
 
+from .args import LaunchOptions
+
 CODEX_CONFIG_DIR = Path.home() / ".codex"
 CODEX_PROFILE_NAME = "ucode"
 CODEX_CONFIG_PATH = CODEX_CONFIG_DIR / f"{CODEX_PROFILE_NAME}.config.toml"
@@ -483,96 +482,59 @@ def clear_model_preferences(state: dict) -> bool:
     return changed
 
 
-# codex rejects the global --profile on subcommands that don't accept it
-# (app-server, mcp-server, ...) with a CLI *parse-time* error — before it touches
-# auth, the gateway, or the network — so the rejection exits almost instantly.
-# We use that to decide when to retry without --profile (see launch()). This
-# window is well above codex's ~0.15s cold-start floor and far below the seconds
-# any real session needs to connect and then fail, so it never catches a genuine
-# failure. Its exit code (1) is indistinguishable from an ordinary failure, so
-# elapsed time is the signal we key on rather than stderr text.
-_PROFILE_REJECTED_MAX_SECONDS = 3.0
-
-
-def launch(state: dict, tool_args: list[str]) -> None:
+def launch(
+    state: dict,
+    tool_args: list[str],
+    *,
+    options: LaunchOptions,
+) -> None:
+    if options.launch_smart_routing:
+        _launch_smart_routing(state, tool_args)
+        return
     clear_model_preferences(state)
     binary = SPEC["binary"]
     workspace = state.get("workspace")
-    if smart_routing_v2.enabled():
-        version_text = agent_version(binary)
-        parsed_version = _parse_version(version_text)
-        if parsed_version is not None and parsed_version < MINIMUM_ROUTING_CODEX_VERSION:
-            raise RuntimeError(
-                "Codex smart routing requires Codex "
-                f"{MINIMUM_ROUTING_CODEX_VERSION_TEXT} or newer; found {version_text}."
-            )
-
-        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=_app_server_start_model(),
-            render_overlay=render_overlay,
-        )
     if workspace:
         os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile"))
-    if tool_args[:1] == ["app"]:
-        # `codex app` rejects --profile. Pass the ucode profile as --config
-        # overrides instead, preserving its Databricks provider and auth
-        # settings without changing the user's base config.toml.
-        profile_doc = read_toml_safe(CODEX_CONFIG_PATH)
-        if not profile_doc:
-            raise RuntimeError(
-                f"Cannot launch Codex app with the ucode profile because {CODEX_CONFIG_PATH} "
-                "is missing or empty. Run `ucode configure --agents codex` first."
-            )
-        config_args = codex_config_args(profile_doc)
-        exec_or_spawn([binary, "app", *config_args, *tool_args[1:]])
-        return  # unreachable in production (exec replaces the process)
-    # Run codex with --profile first — the TUI and runtime subcommands
-    # (exec/resume/mcp/...) keep ucode's Databricks routing, including any added
-    # by future codex versions. codex rejects the global --profile on
-    # server-family subcommands (app-server, mcp-server, ...), which are
-    # caller-configured anyway (e.g. omnigent runs `codex app-server` with its
-    # own CODEX_HOME); on that rejection we relaunch without --profile.
-    #
-    # The retry is gated on the attempt failing *fast*: the rejection is a
-    # parse-time error (~0.15s), whereas a session that actually starts can only
-    # fail after a network round-trip (seconds). Without that gate a genuinely
-    # failing `codex exec` would be silently re-run without --profile — i.e. on
-    # the user's own OpenAI login instead of the Databricks gateway (ucode writes
-    # a *named-profile* file, so no --profile means no ucode routing). stdio is
-    # inherited (no capture), so Ctrl-C reaches codex directly and the resulting
-    # KeyboardInterrupt propagates past the retry check — quitting an interactive
-    # session is never mistaken for a --profile rejection.
-    started = time.monotonic()
-    returncode = subprocess.run([binary, "--profile", CODEX_PROFILE_NAME, *tool_args]).returncode
-    if returncode != 0 and time.monotonic() - started < _PROFILE_REJECTED_MAX_SECONDS:
-        # Fast failure: most likely codex rejected --profile on this subcommand.
-        # Relaunch without it, handing over the terminal. (A fast failure for
-        # any other reason — e.g. a bad flag — just re-fails the same way here,
-        # with no ucode routing to lose since the subcommand had none.)
-        #
-        # Warn on *stderr*: this path is reached by `codex app-server`, whose
-        # stdout is a JSON-RPC stream its caller parses. Emit before handing off,
-        # since execvp replaces this process.
-        print_warning_err(
-            "ucode's `--profile` isn't accepted here (error above). Retrying "
-            f"without it: Codex will resolve {LEGACY_CODEX_CONFIG_PATH} and any OS-managed "
-            "settings instead of the ucode profile."
+    # Layer ucode's named profile as ordinary config overrides. Unlike
+    # `--profile`, `--config` is accepted by runtime, utility, and server
+    # commands, so every invocation keeps the same Databricks settings without
+    # classifying Codex subcommands or probing and retrying the real command.
+    profile_doc = read_toml_safe(CODEX_CONFIG_PATH)
+    if not profile_doc:
+        raise RuntimeError(
+            f"Cannot launch Codex with the ucode profile because {CODEX_CONFIG_PATH} "
+            "is missing or empty. Run `ucode configure --agents codex` first."
         )
-        exec_or_spawn([binary, *tool_args])
-        return  # unreachable in production (exec replaces the process)
-    sys.exit(returncode)
+    exec_or_spawn([binary, *codex_config_args(profile_doc), *tool_args])
+
+
+def _launch_smart_routing(state: dict, tool_args: list[str]) -> None:
+    """Launch the Codex TUI through the smart-routing interposer."""
+    clear_model_preferences(state)
+    binary = SPEC["binary"]
+    version_text = agent_version(binary)
+    parsed_version = _parse_version(version_text)
+    if parsed_version is not None and parsed_version < MINIMUM_ROUTING_CODEX_VERSION:
+        raise RuntimeError(
+            "Codex smart routing requires Codex "
+            f"{MINIMUM_ROUTING_CODEX_VERSION_TEXT} or newer; found {version_text}."
+        )
+
+    managed_model = default_model(state)
+    models = routing_models(state)
+    start_model = (
+        managed_model
+        or (codex_model_id(models[0]) if models else None)
+        or APP_SERVER_SMART_ROUTING_STARTING_MODEL
+    )
+    smart_routing_v2.launch_codex(
+        state,
+        tool_args,
+        binary=binary,
+        start_model=start_model,
+        render_overlay=render_overlay,
+    )
 
 
 def disable_smart_routing(state: dict) -> bool:

src/ucode/agents/copilot.py

@@ -36,6 +36,8 @@
 )
 from ucode.state import mark_tool_managed, save_state
 
+from .args import LaunchOptions
+
 COPILOT_CONFIG_DIR = Path.home() / ".copilot"
 COPILOT_ENV_PATH = COPILOT_CONFIG_DIR / "ucode.env"
 COPILOT_MCP_CONFIG_PATH = COPILOT_CONFIG_DIR / "ucode-mcp-config.json"
@@ -179,7 +181,7 @@ def _refresh_forever(state: dict, stop_event: threading.Event) -> None:
             continue
 
 
-def launch(state: dict, tool_args: list[str]) -> None:
+def launch(state: dict, tool_args: list[str], *, options: LaunchOptions) -> None:
     model, token = _refresh_token_once(state)
     env = build_runtime_env(state["workspace"], model, token)