Test managed model routing precedence

src/ucode/agents/args.py

@@ -10,9 +10,7 @@ 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
+    user_pinned_model: str | None = None
 
 
 def explicit_model_arg_value(tool_args: list[str]) -> str | None:

src/ucode/agents/claude.py

@@ -101,15 +101,18 @@ def _parse_version(value: str) -> tuple[int, int, int] | None:
 
 
 def _minimum_version_requirement_message(version: str) -> str:
-    feature = "Smart routing" if smart_routing_v2.enabled() else "Model discovery"
+    feature = "Smart routing" if smart_routing_v2.smart_routing_enabled() else "Model discovery"
     return (
         f"{feature} requires Claude Code {MINIMUM_CLAUDE_VERSION_TEXT} or newer. "
         f"Your current version is Claude Code {version}."
     )
 
 
 def minimum_version_error() -> str | None:
-    if os.environ.get(GATEWAY_MODEL_DISCOVERY_ENV_VAR) != "1" and not smart_routing_v2.enabled():
+    if (
+        os.environ.get(GATEWAY_MODEL_DISCOVERY_ENV_VAR) != "1"
+        and not smart_routing_v2.smart_routing_enabled()
+    ):
         return None
     version = agent_version(SPEC["binary"])
     parsed = _parse_version(version)
@@ -1321,16 +1324,6 @@ def _compose_v2_settings(tool_args: list[str]) -> tuple[dict, list[str]]:
     return _merge_claude_settings(settings, read_json_safe(CLAUDE_SETTINGS_PATH)), remaining
 
 
-def _original_launch_model(state: dict) -> str | None:
-    override = state.get("_claude_launch_model")
-    if isinstance(override, str) and override.strip():
-        return override.strip()
-    value = read_json_safe(CLAUDE_USER_SETTINGS_PATH).get("model")
-    if isinstance(value, str) and value.strip():
-        return value.strip()
-    return default_model(state)
-
-
 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 []
@@ -1506,17 +1499,25 @@ def launch(
             tool_args,
             binary=binary,
             user_settings_path=CLAUDE_USER_SETTINGS_PATH,
-            launch_model=_original_launch_model(state),
+            # With no user pin, let Claude resolve its starting model from its own settings.
+            launch_model=options.user_pinned_model,
             compose_settings=_compose_v2_settings,
             launch_model_args=_launch_model_args,
             model_name=_maybe_add_1m_suffix,
         )
         return
     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))
+    settings_override = None
+    launch_args = list(tool_args)
+    if options.user_pinned_model:
+        os.environ["ANTHROPIC_MODEL"] = options.user_pinned_model
+        settings_override = {"env": {"ANTHROPIC_MODEL": options.user_pinned_model}}
+        launch_args = [
+            *_launch_model_args(tool_args, options.user_pinned_model),
+            *tool_args,
+        ]
+    exec_or_spawn(_build_claude_argv(binary, launch_args, settings_override=settings_override))
 
 
 def validate_cmd(binary: str) -> list[str]:

src/ucode/agents/codex.py

@@ -135,7 +135,7 @@ def _parse_version(value: str) -> tuple[int, int, int] | None:
 
 def minimum_version_error() -> str | None:
     """Return the active smart-routing version blocker, if any."""
-    if not smart_routing_v2.enabled():
+    if not smart_routing_v2.smart_routing_enabled():
         return None
     version = agent_version(SPEC["binary"])
     parsed = _parse_version(version)
@@ -458,7 +458,7 @@ def compose(base: dict, *, include_catalog: bool = True) -> dict:
         prune_key_paths(base, _MODEL_SERVICE_ROUTING_KEY_PATHS)
         deep_merge_dict(base, copy.deepcopy(overlay))
         # deep_merge can't drop keys, so clear model preferences from an earlier run.
-        if chosen_model is None and not smart_routing_v2.enabled():
+        if chosen_model is None and not smart_routing_v2.smart_routing_enabled():
             for key in ("model", "model_reasoning_effort"):
                 base.pop(key, None)
         if include_catalog:
@@ -586,7 +586,7 @@ def default_model(state: dict) -> str | None:
     """Return a managed Codex model, or leave selection to Codex."""
     if isinstance(state.get("codex_default_model"), str):
         return state["codex_default_model"]
-    if smart_routing_v2.enabled():
+    if smart_routing_v2.smart_routing_enabled():
         return _smart_routing_config_model(state)
     clear_model_preferences(state)
     return None
@@ -615,7 +615,7 @@ def config_precedence_paths() -> tuple[Path, ...]:
 
 def clear_model_preferences(state: dict) -> bool:
     """Remove ucode profile model preferences so Codex selects its default."""
-    if smart_routing_v2.enabled():
+    if smart_routing_v2.smart_routing_enabled():
         return False
     if isinstance(state.get("codex_default_model"), str):
         return False

src/ucode/cli.py

@@ -1586,7 +1586,7 @@ def codex_router_hook_cmd(
     import json
     import sys
 
-    if not smart_routing_v2.enabled():
+    if not smart_routing_v2.smart_routing_enabled():
         return
 
     from ucode.smart_routing.codex_routing import (
@@ -1665,7 +1665,7 @@ def claude_router_hook_cmd(
     import json
     import sys
 
-    if not smart_routing_v2.enabled():
+    if not smart_routing_v2.smart_routing_enabled():
         return
 
     from ucode.smart_routing.claude_routing import (
@@ -1691,7 +1691,12 @@ def claude_router_hook_cmd(
             request_first_prompt_route,
         )
 
-        output = first_prompt_hook_output(request_first_prompt_route(Path(socket_path), payload))
+        response = request_first_prompt_route(
+            Path(socket_path),
+            payload,
+            timeout=smart_routing_v2.CLAUDE_ROUTE_SELECTION_TIMEOUT_S + 5.0,
+        )
+        output = first_prompt_hook_output(response)
         if output is not None:
             sys.stdout.write(json.dumps(output))
         return
@@ -1782,15 +1787,11 @@ def _smart_routing_v2_flag(enabled: bool) -> Iterator[None]:
     if not enabled:
         yield
         return
-    previous = os.environ.get(smart_routing_v2.ENV_VAR)
-    os.environ[smart_routing_v2.ENV_VAR] = "1"
+    previous = smart_routing_v2.enable_smart_routing()
     try:
         yield
     finally:
-        if previous is None:
-            os.environ.pop(smart_routing_v2.ENV_VAR, None)
-        else:
-            os.environ[smart_routing_v2.ENV_VAR] = previous
+        smart_routing_v2.restore_smart_routing_env(previous)
 
 
 @contextmanager
@@ -1805,12 +1806,11 @@ def _disable_smart_routing_for_subcommand(tool: str, ctx: Any) -> Iterator[None]
     if _smart_routing_launch_shape(tool, ctx.args, _has_explicit_prompt(ctx)):
         yield
         return
-    previous = os.environ.pop(smart_routing_v2.ENV_VAR, None)
+    previous = smart_routing_v2.disable_smart_routing()
     try:
         yield
     finally:
-        if previous is not None:
-            os.environ[smart_routing_v2.ENV_VAR] = previous
+        smart_routing_v2.restore_smart_routing_env(previous)
 
 
 def _migrate_legacy_smart_routing(state: dict) -> dict:
@@ -2017,11 +2017,12 @@ def _launch_options(
     *,
     smart_routing_enabled: bool,
     explicit_prompt: bool,
-    model: str | None,
+    user_pinned_model: str | None,
     provider: str | None,
 ) -> LaunchOptions:
     return LaunchOptions(
-        claude_launch_model=model if tool == "claude" and provider is None else None,
+        # Pinned models for providers are resolved above through the provider-specific launch path.
+        user_pinned_model=user_pinned_model if provider is None else None,
         launch_smart_routing=(
             # Smart routing is enabled globally.
             smart_routing_enabled
@@ -2034,12 +2035,32 @@ def _launch_options(
                 tool,
                 tool_args,
                 explicit_prompt=explicit_prompt,
-                model=model,
+                model=user_pinned_model,
             )
         ),
     )
 
 
+@contextmanager
+def _managed_smart_routing_environment(managed: dict | None, tool: str) -> Iterator[None]:
+    """Expose an agent's managed smart-routing switch only to its launched session."""
+    if not _managed_smart_routing_enabled(managed, tool):
+        yield
+        return
+
+    previous = smart_routing_v2.enable_smart_routing()
+    try:
+        yield
+    finally:
+        smart_routing_v2.restore_smart_routing_env(previous)
+
+
+def _managed_smart_routing_enabled(managed: dict | None, tool: str) -> bool:
+    """Whether the workspace enabled smart routing for this specific agent."""
+    agent_config = ((managed or {}).get("enabled_agents") or {}).get(tool) or {}
+    return agent_config.get("smart_routing_enabled") is True
+
+
 def _launch_tool(
     tool_name: str,
     ctx: typer.Context,
@@ -2064,7 +2085,7 @@ def _launch_tool(
         if parent_schema is not None and not is_valid_catalog_schema(parent_schema):
             raise RuntimeError("--parent must be `<catalog>.<schema>`.")
         explicit_prompt = _has_explicit_prompt(ctx)
-        smart_routing_enabled = smart_routing_v2.enabled()
+        smart_routing_enabled = smart_routing_v2.smart_routing_enabled()
         # Launchers such as isaac put their harness arguments after `--`, so the harness's own
         # `--model` lands in ctx.args instead of a ucode option. It still determines the effective
         # launch model and should therefore win in the launch summary.
@@ -2110,6 +2131,10 @@ def _launch_tool(
             managed, coding_agent_config_feature_disabled = _fetch_managed_config(state)
         # Checked before discovery, which can take tens of seconds, so a blocked launch fails fast.
         _reject_disabled_agent(managed, tool)
+        # The environment switch remains a developer override; managed config is the workspace
+        # policy equivalent and must take effect before launch options are computed.
+        managed_smart_routing_enabled = _managed_smart_routing_enabled(managed, tool)
+        smart_routing_enabled = smart_routing_enabled or managed_smart_routing_enabled
         # Discovery exists to find models and isn't needed for managed config that already names them.
         managed_models_known = managed_supplies_models(managed, tool)
         # Re-fetch model lists on every launch so newly-added Databricks
@@ -2205,6 +2230,7 @@ def _launch_tool(
         # The router's per-launch pick for the root session. Codex pins it as the
         # resolved model; claude pins it via ANTHROPIC_MODEL (route_root_model).
         route_root_model = None
+        managed_model = None
         relayed_forward_model = None  # forwarded to Claude Code's --model for a relayed provider
         if provider:
             # Routing through a Model Provider Service pins no Databricks model;
@@ -2299,13 +2325,6 @@ def _launch_tool(
             _register_managed_mcp_servers(managed, tool, state)
             _download_managed_skills(managed, state)
         if tool == "claude":
-            if smart_routing_v2.enabled():
-                # Transient launch precedence for the v2 PTY's initial --model flag.
-                # An explicit choice wins, followed by a routed/managed root pick;
-                # neither value is persisted into workspace state.
-                launch_model = model or route_root_model
-                if launch_model:
-                    state["_claude_launch_model"] = launch_model
             if provider:
                 state["_claude_launch_provider"] = provider
         elif tool == "codex":
@@ -2318,11 +2337,14 @@ def _launch_tool(
             ctx.args,
             smart_routing_enabled=smart_routing_enabled,
             explicit_prompt=explicit_prompt,
-            model=model or (route_root_model if tool == "claude" else None),
+            # Only a developer's explicit model disables routing. A managed default is the
+            # initial/fallback model and still participates in a routed session.
+            user_pinned_model=model or forwarded_model,
             provider=provider,
         )
         print_success(f"Starting {TOOL_SPECS[tool]['display']}")
-        launch_agent(tool, state, ctx.args, options=launch_options)
+        with _managed_smart_routing_environment(managed, tool):
+            launch_agent(tool, state, ctx.args, options=launch_options)
     except RuntimeError as exc:
         print_err(str(exc))
         raise typer.Exit(1) from None

src/ucode/managed_config.py

@@ -139,19 +139,22 @@ class AgentConfig:
 
     http_headers: dict[str, str] | None = None
     models: AgentModels | None = None
+    smart_routing_enabled: bool = False
     otel_tracing_enabled: bool | None = None
 
     @classmethod
     def from_wire(cls, config: object) -> AgentConfig:
         """Parse wire format AgentConfig into normalized AgentConfig."""
         config_dict = _as_dict(config)
         headers = _clean_str_dict(config_dict.get("http_headers"))
+        smart_routing = _as_dict(config_dict.get("smart_routing"))
         agent_models = AgentModels.from_wire(
             config_dict.get("default_models"), config_dict.get("models")
         )
         return cls(
             http_headers=headers or None,
             models=agent_models,
+            smart_routing_enabled=smart_routing.get("enabled") is True,
             otel_tracing_enabled=_tracing_enabled(config_dict.get("tracing")),
         )
 
@@ -160,6 +163,8 @@ def to_internal(self) -> dict:
         result: dict = {}
         if self.http_headers:
             result["http_headers"] = self.http_headers
+        if self.smart_routing_enabled:
+            result["smart_routing_enabled"] = True
         model_config = self.models.to_internal() if self.models else None
         if model_config is not None:
             result["model_config"] = model_config

src/ucode/managed_setup.py

@@ -219,6 +219,8 @@ def _enabled_agent_payload(tool: str, agent_config: dict) -> dict:
         clean = {k: v for k, v in headers.items() if isinstance(k, str) and isinstance(v, str)}
         if clean:
             config["http_headers"] = clean
+    if agent_config.get("smart_routing_enabled") is True:
+        config["smart_routing"] = {"enabled": True}
     model_config = agent_config.get("model_config")
     if isinstance(model_config, dict):
         payload = _model_config_payload(tool, model_config)

src/ucode/smart_routing/codex_routing.py

@@ -39,6 +39,18 @@
 _normalize_model = routing.normalize_model
 
 
+# TODO: Remove once server-side smart routing is more robust to duplicate model names.
+def _normalize_route_model(model: str) -> str:
+    """Canonicalize equivalent Codex GPT spellings to one router option."""
+    normalized = _normalize_model(model)
+    match = _GPT_RE.fullmatch(normalized)
+    if match is None:
+        return normalized
+    major, minor, patch, suffix = match.groups()
+    version = "-".join(part for part in (major, minor, patch) if part is not None)
+    return f"gpt-{version}{suffix or ''}"
+
+
 def request_routing_decision(
     workspace: str,
     token: str,
@@ -49,7 +61,7 @@ def request_routing_decision(
     log: Callable[[str], None] | None = None,
 ) -> tuple[RoutingDecision | None, str | None]:
     """Ask the router for a servable Codex model."""
-    available = {_normalize_model(model): model for model in available_models}
+    available = {_normalize_route_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"
@@ -69,16 +81,16 @@ def request_routing_decision(
         token,
         task,
         route_options,
-        lambda raw_model: available.get(_normalize_model(raw_model)),
+        lambda raw_model: available.get(_normalize_route_model(raw_model)),
         router_name=router_name,
         timeout=timeout,
     )
 
 
 def resolve_routed_model(raw_model: str, available_models: list[str]) -> str | None:
     """Map a router arm to a model the configured workspace can serve."""
-    normalized = {_normalize_model(model): model for model in available_models}
-    return normalized.get(_normalize_model(raw_model))
+    normalized = {_normalize_route_model(model): model for model in available_models}
+    return normalized.get(_normalize_route_model(raw_model))
 
 
 def route_pre_tool_use(