diff --git a/src/ucode/agents/args.py b/src/ucode/agents/args.py index 7cfb9e1c9..faca49dfe 100644 --- a/src/ucode/agents/args.py +++ b/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: diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 23dabb0fc..db46a8ace 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -101,7 +101,7 @@ 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}." @@ -109,7 +109,10 @@ def _minimum_version_requirement_message(version: str) -> str: 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,7 +1499,8 @@ 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, @@ -1514,9 +1508,16 @@ def launch( 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]: diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 671fbb749..572d98989 100644 --- a/src/ucode/agents/codex.py +++ b/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 diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 5efe1831a..2788c3afe 100644 --- a/src/ucode/cli.py +++ b/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 `.`.") 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 diff --git a/src/ucode/managed_config.py b/src/ucode/managed_config.py index 652147f29..368a7dcd2 100644 --- a/src/ucode/managed_config.py +++ b/src/ucode/managed_config.py @@ -139,6 +139,7 @@ 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 @@ -146,12 +147,14 @@ 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 diff --git a/src/ucode/managed_setup.py b/src/ucode/managed_setup.py index 097d016cf..7d1fd3f6f 100644 --- a/src/ucode/managed_setup.py +++ b/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) diff --git a/src/ucode/smart_routing/codex_routing.py b/src/ucode/smart_routing/codex_routing.py index 154275c62..9569bc14d 100644 --- a/src/ucode/smart_routing/codex_routing.py +++ b/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,7 +81,7 @@ 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, ) @@ -77,8 +89,8 @@ def request_routing_decision( 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( diff --git a/src/ucode/smart_routing/v2.py b/src/ucode/smart_routing/v2.py index e2afd2722..00d897279 100644 --- a/src/ucode/smart_routing/v2.py +++ b/src/ucode/smart_routing/v2.py @@ -10,7 +10,7 @@ import time import urllib.request import uuid -from collections.abc import Callable +from collections.abc import Callable, MutableMapping from pathlib import Path from typing import NoReturn, TextIO @@ -36,7 +36,7 @@ from ucode.smart_routing.codex_hooks import merge_pre_tool_use_hooks, routing_models from ucode.ui import print_warning -ENV_VAR = "ENABLE_SMART_ROUTING_V2" +ENABLE_SMART_ROUTING_ENV_VAR = "ENABLE_SMART_ROUTING_V2" LEGACY_STATE_KEY = "smart_routing_enabled" CODEX_INTERPOSER_LOG = APP_DIR / "codex-v2-interposer.log" @@ -98,8 +98,34 @@ def _model_picker_catalog() -> AnthropicModelCatalog | None: return None -def enabled() -> bool: - return os.environ.get(ENV_VAR) == "1" +def smart_routing_enabled(env: MutableMapping[str, str] | None = None) -> bool: + source = os.environ if env is None else env + return source.get(ENABLE_SMART_ROUTING_ENV_VAR) == "1" + + +def enable_smart_routing(env: MutableMapping[str, str] | None = None) -> str | None: + """Set the only supported smart-routing env var and return its prior value.""" + target = os.environ if env is None else env + previous = target.get(ENABLE_SMART_ROUTING_ENV_VAR) + target[ENABLE_SMART_ROUTING_ENV_VAR] = "1" + return previous + + +def restore_smart_routing_env( + previous: str | None, env: MutableMapping[str, str] | None = None +) -> None: + """Restore the env state captured when smart routing was enabled or disabled.""" + target = os.environ if env is None else env + if previous is None: + target.pop(ENABLE_SMART_ROUTING_ENV_VAR, None) + else: + target[ENABLE_SMART_ROUTING_ENV_VAR] = previous + + +def disable_smart_routing(env: MutableMapping[str, str] | None = None) -> str | None: + """Temporarily remove the smart-routing env var and return its prior value.""" + target = os.environ if env is None else env + return target.pop(ENABLE_SMART_ROUTING_ENV_VAR, None) def _loopback_websocket_url(port: int) -> str: @@ -416,6 +442,7 @@ def launch_claude( if not isinstance(env, dict): raise RuntimeError("Claude settings 'env' must be an object for smart routing.") env.pop("CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY", None) + env[ENABLE_SMART_ROUTING_ENV_VAR] = "1" env[FIRST_PROMPT_SOCKET_ENV] = str(socket_path) model_overrides = settings.setdefault("modelOverrides", {}) if not isinstance(model_overrides, dict): diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 34e0388d4..8f484cc1b 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -47,13 +47,13 @@ def test_display(self): class TestMinimumVersion: @pytest.mark.parametrize("version", ["2.1.248", "2.1.250", "3.0.0"]) def test_supported_version(self, monkeypatch, version): - monkeypatch.setenv(v2.ENV_VAR, "1") + monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_ENV_VAR, "1") monkeypatch.setattr(claude, "agent_version", lambda _binary: version) assert claude.minimum_version_error() is None def test_older_version_requires_update(self, monkeypatch): - monkeypatch.setenv(v2.ENV_VAR, "1") + monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_ENV_VAR, "1") monkeypatch.setattr(claude, "agent_version", lambda _binary: "2.1.247") assert claude.minimum_version_error() == ( @@ -72,20 +72,20 @@ def test_older_version_requires_update_for_model_discovery(self, monkeypatch): assert claude.minimum_version_error() == expected def test_smart_routing_message_wins_when_both_features_are_enabled(self, monkeypatch): - monkeypatch.setenv(v2.ENV_VAR, "1") + monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_ENV_VAR, "1") monkeypatch.setenv(claude.GATEWAY_MODEL_DISCOVERY_ENV_VAR, "1") monkeypatch.setattr(claude, "agent_version", lambda _binary: "2.1.247") assert claude.minimum_version_error().startswith("Smart routing requires") def test_unknown_version_does_not_block(self, monkeypatch): - monkeypatch.setenv(v2.ENV_VAR, "1") + monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_ENV_VAR, "1") monkeypatch.setattr(claude, "agent_version", lambda _binary: "unknown") assert claude.minimum_version_error() is None def test_older_version_is_not_validated_without_discovery_features(self, monkeypatch): - monkeypatch.delenv(v2.ENV_VAR, raising=False) + monkeypatch.delenv(v2.ENABLE_SMART_ROUTING_ENV_VAR, raising=False) monkeypatch.delenv(claude.GATEWAY_MODEL_DISCOVERY_ENV_VAR, raising=False) monkeypatch.setattr(claude, "agent_version", lambda _binary: "2.1.247") @@ -207,7 +207,7 @@ def test_does_not_persist_gateway_model_discovery(self, monkeypatch): assert "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY" not in overlay["env"] def test_smart_routing_does_not_persist_gateway_model_discovery(self, monkeypatch): - monkeypatch.setenv(v2.ENV_VAR, "1") + monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_ENV_VAR, "1") monkeypatch.delenv(claude.GATEWAY_MODEL_DISCOVERY_ENV_VAR, raising=False) overlay, _ = claude.render_overlay(WS, "s4") assert "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY" not in overlay["env"] @@ -1522,7 +1522,7 @@ def start_proxy(workspace, profile, port, token_header, force_refresh_near_expir assert calls[-3:] == [("stop",), ("shutdown",), ("close",)] def test_smart_routing_on_windows_is_not_supported(self, monkeypatch): - monkeypatch.setenv(v2.ENV_VAR, "1") + monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_ENV_VAR, "1") monkeypatch.setattr(claude.os, "name", "nt") with pytest.raises( @@ -1537,7 +1537,7 @@ def test_smart_routing_on_windows_is_not_supported(self, monkeypatch): def test_default_launch_keeps_existing_auth_path(self, monkeypatch): calls: list[list[str]] = [] - monkeypatch.delenv(v2.ENV_VAR, raising=False) + monkeypatch.delenv(v2.ENABLE_SMART_ROUTING_ENV_VAR, raising=False) monkeypatch.delenv(claude.GATEWAY_MODEL_DISCOVERY_ENV_VAR, raising=False) monkeypatch.delenv("OAUTH_TOKEN", raising=False) monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "token") @@ -1557,11 +1557,13 @@ def test_launch_model_is_only_set_for_current_process(self, monkeypatch): claude.launch( {"workspace": WS, "profile": "test"}, [], - options=LaunchOptions(claude_launch_model="cat.schema.model"), + options=LaunchOptions(user_pinned_model="cat.schema.model"), ) assert os.environ["ANTHROPIC_MODEL"] == "cat.schema.model" - assert calls == [["claude", "--settings", str(claude.CLAUDE_SETTINGS_PATH)]] + assert calls[0][:2] == ["claude", "--settings"] + settings = json.loads(calls[0][2]) + assert settings["env"]["ANTHROPIC_MODEL"] == "cat.schema.model" @pytest.mark.parametrize( "tool_args", @@ -1572,7 +1574,7 @@ def test_launch_model_is_only_set_for_current_process(self, monkeypatch): ) def test_v2_noninteractive_launch_bypasses_first_prompt_routing(self, monkeypatch, tool_args): calls: list[list[str]] = [] - monkeypatch.setenv(v2.ENV_VAR, "1") + monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_ENV_VAR, "1") monkeypatch.setattr(v2, "launch_claude", Mock()) monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "token") monkeypatch.setattr(claude, "exec_or_spawn", lambda argv: calls.append(argv)) @@ -1584,9 +1586,8 @@ def test_v2_noninteractive_launch_bypasses_first_prompt_routing(self, monkeypatc @pytest.mark.parametrize("tool_args", [["fix this bug"], ["--", "fix this bug"]]) def test_v2_positional_prompt_uses_first_prompt_routing(self, monkeypatch, tool_args): - monkeypatch.setenv(v2.ENV_VAR, "1") + monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_ENV_VAR, "1") launch_v2 = Mock() - monkeypatch.setattr(claude, "_original_launch_model", lambda _state: None) monkeypatch.setattr(v2, "launch_claude", launch_v2) claude.launch( @@ -1608,7 +1609,7 @@ def test_v2_positional_prompt_uses_first_prompt_routing(self, monkeypatch, tool_ def test_gateway_discovery_uses_direct_gateway(self, monkeypatch): calls: list[list[str]] = [] - monkeypatch.delenv(v2.ENV_VAR, raising=False) + monkeypatch.delenv(v2.ENABLE_SMART_ROUTING_ENV_VAR, raising=False) monkeypatch.setenv(claude.GATEWAY_MODEL_DISCOVERY_ENV_VAR, "1") monkeypatch.delenv("OAUTH_TOKEN", raising=False) monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "token") @@ -1622,7 +1623,7 @@ def test_gateway_discovery_uses_direct_gateway(self, monkeypatch): def test_gateway_discovery_enabled_under_provider(self, monkeypatch): calls: list[list[str]] = [] - monkeypatch.delenv(v2.ENV_VAR, raising=False) + monkeypatch.delenv(v2.ENABLE_SMART_ROUTING_ENV_VAR, raising=False) monkeypatch.setenv(claude.GATEWAY_MODEL_DISCOVERY_ENV_VAR, "1") monkeypatch.delenv("OAUTH_TOKEN", raising=False) monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "token") diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index 93aaa2972..a95a1fe61 100644 --- a/tests/test_agent_codex.py +++ b/tests/test_agent_codex.py @@ -29,14 +29,14 @@ def test_display(self): class TestMinimumVersion: def test_smart_routing_old_version_requires_update(self, monkeypatch): - monkeypatch.setenv(codex.smart_routing_v2.ENV_VAR, "1") + monkeypatch.setenv(codex.smart_routing_v2.ENABLE_SMART_ROUTING_ENV_VAR, "1") monkeypatch.setattr(codex, "agent_version", lambda _binary: "0.144.0") expected = "Codex smart routing requires Codex 0.145.0 or newer; found 0.144.0." assert codex.minimum_version_error() == expected def test_old_version_is_not_blocked_without_smart_routing(self, monkeypatch): - monkeypatch.delenv(codex.smart_routing_v2.ENV_VAR, raising=False) + monkeypatch.delenv(codex.smart_routing_v2.ENABLE_SMART_ROUTING_ENV_VAR, raising=False) monkeypatch.setattr(codex, "agent_version", lambda _binary: "0.144.0") assert codex.minimum_version_error() is None @@ -204,7 +204,7 @@ def test_smart_routing_preserves_configured_startup_model(self, tmp_path, monkey monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", config_path) monkeypatch.setattr(codex, "CODEX_BACKUP_PATH", tmp_path / "backup.toml") monkeypatch.setattr(codex, "agent_version", lambda _: "0.145.0") - monkeypatch.setenv(codex.smart_routing_v2.ENV_VAR, "1") + monkeypatch.setenv(codex.smart_routing_v2.ENABLE_SMART_ROUTING_ENV_VAR, "1") monkeypatch.delenv("CODEX_HOME", raising=False) state = {"workspace": WS} diff --git a/tests/test_claude_smart_routing_v2.py b/tests/test_claude_smart_routing_v2.py index 2afba0281..fdd3da5fd 100644 --- a/tests/test_claude_smart_routing_v2.py +++ b/tests/test_claude_smart_routing_v2.py @@ -259,6 +259,7 @@ def fake_run(argv, **kwargs): "claude-opus-4-8": "system.ai.claude-opus-4-8", "claude-sonnet-5": "system.ai.claude-sonnet-5", } + assert captured["settings"]["env"][v2.ENABLE_SMART_ROUTING_ENV_VAR] == "1" assert claude_hooks.FIRST_PROMPT_SOCKET_ENV in captured["settings"]["env"] first_prompt_command = captured["settings"]["hooks"]["UserPromptSubmit"][0]["hooks"][0][ "command" diff --git a/tests/test_cli.py b/tests/test_cli.py index 931fe2b2a..9c9912d85 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -470,14 +470,14 @@ def test_codex_enable_smart_routing_is_consumed_by_ucode(self): with patch( "ucode.cli._launch_tool", side_effect=lambda *_args, **_kwargs: enabled_during_launch.append( - os.environ.get("ENABLE_SMART_ROUTING_V2") + os.environ.get(cli_mod.smart_routing_v2.ENABLE_SMART_ROUTING_ENV_VAR) ), ) as mock_launch: result = runner.invoke(app, ["codex", "--enable-smart-routing"]) assert result.exit_code == 0, result.output assert enabled_during_launch == ["1"] - assert "ENABLE_SMART_ROUTING_V2" not in os.environ + assert cli_mod.smart_routing_v2.ENABLE_SMART_ROUTING_ENV_VAR not in os.environ assert mock_launch.call_args.args[1].args == [] @pytest.mark.parametrize("tool, subcommand", [("codex", "app"), ("claude", "update")]) @@ -490,20 +490,22 @@ def test_native_subcommand_suppresses_inherited_smart_routing( with patch( "ucode.cli._launch_tool", side_effect=lambda *_args, **_kwargs: observed.append( - os.environ.get("ENABLE_SMART_ROUTING_V2") + os.environ.get(cli_mod.smart_routing_v2.ENABLE_SMART_ROUTING_ENV_VAR) ), ): result = runner.invoke(app, [tool, subcommand]) assert result.exit_code == 0, result.output assert observed == [None] - assert os.environ["ENABLE_SMART_ROUTING_V2"] == "1" + assert os.environ[cli_mod.smart_routing_v2.ENABLE_SMART_ROUTING_ENV_VAR] == "1" def test_claude_enable_smart_routing_forwards_positional_prompt_to_v2(self): captured = [] def capture(_tool, ctx, **_kwargs): - captured.append((os.environ.get("ENABLE_SMART_ROUTING_V2"), ctx.args)) + captured.append( + (os.environ.get(cli_mod.smart_routing_v2.ENABLE_SMART_ROUTING_ENV_VAR), ctx.args) + ) with patch("ucode.cli._launch_tool", side_effect=capture): result = runner.invoke( @@ -556,7 +558,7 @@ def test_codex_and_claude_share_smart_routing_policy( tool_args, smart_routing_enabled=True, explicit_prompt=explicit_prompt, - model=model, + user_pinned_model=model, provider=provider, ) @@ -580,12 +582,148 @@ def test_claude_options_allow_smart_routing_except_model(self, tool_args, expect tool_args, smart_routing_enabled=True, explicit_prompt=False, - model=None, + user_pinned_model=None, provider=None, ) assert options.launch_smart_routing is expected + def test_claude_without_user_pin_leaves_model_to_harness(self): + options = cli_mod._launch_options( + "claude", + [], + smart_routing_enabled=True, + explicit_prompt=False, + user_pinned_model=None, + provider=None, + ) + + assert options.launch_smart_routing is True + assert options.user_pinned_model is None + + def test_launch_options_carry_user_pin(self): + options = cli_mod.LaunchOptions(user_pinned_model="user-model") + assert options.user_pinned_model == "user-model" + + @pytest.mark.parametrize("tool", ["claude", "codex"]) + @pytest.mark.parametrize("smart_routing_enabled", [False, True]) + @pytest.mark.parametrize("managed_default", [None, "managed-model"]) + def test_managed_model_and_routing_launch_matrix( + self, monkeypatch, tool, smart_routing_enabled, managed_default + ): + """A managed default selects the starting model without suppressing managed routing.""" + state = dict(MINIMAL_STATE) + agent_config = {"smart_routing_enabled": smart_routing_enabled} + if managed_default is not None: + agent_config["model_config"] = {"default_model": managed_default} + managed = {"enabled_agents": {tool: agent_config}} + configured: list[dict] = [] + + def capture_configure(_tool, configured_state, model, **kwargs): + configured.append({"state": configured_state, "model": model, **kwargs}) + return configured_state + + monkeypatch.delenv("ENABLE_SMART_ROUTING_V2", raising=False) + with ( + patch("ucode.cli.load_state", return_value=state), + patch("ucode.cli.apply_pat_environment"), + patch("ucode.cli.ensure_bootstrap_dependencies"), + patch("ucode.cli.ensure_provider_state", return_value=state), + patch("ucode.cli.configure_shared_state", return_value=state), + patch("ucode.cli.configure_tool", side_effect=capture_configure), + patch("ucode.cli._fetch_managed_config", return_value=(managed, False)), + patch("ucode.cli._fetch_budget_recommendation", return_value=None), + patch("ucode.cli._register_managed_mcp_servers"), + patch("ucode.cli._download_managed_skills"), + patch("ucode.cli.launch_agent") as launch, + ): + result = runner.invoke(app, [tool]) + + assert result.exit_code == 0, result.output + options = launch.call_args.kwargs["options"] + assert options.user_pinned_model is None + assert options.launch_smart_routing is smart_routing_enabled + if managed_default is not None: + assert configured[0]["model"] == managed_default + if tool == "claude": + assert configured[0]["route_root_model"] == managed_default + else: + assert configured[0]["route_root_model"] is None + + @pytest.mark.parametrize("tool", ["claude", "codex"]) + @pytest.mark.parametrize( + "pin_args", + [ + pytest.param(["--model", "pinned-model"], id="direct-model"), + pytest.param(["--", "--model", "pinned-model"], id="forwarded-model"), + pytest.param(["--", "--model=pinned-model"], id="forwarded-model-equals"), + ], + ) + def test_user_pinned_model_wins_over_managed_default_and_routing( + self, monkeypatch, tool, pin_args + ): + """Every supported spelling of a user model pin wins and bypasses smart routing.""" + state = dict(MINIMAL_STATE) + managed = { + "enabled_agents": { + tool: { + "smart_routing_enabled": True, + "model_config": {"default_model": "managed-model"}, + } + } + } + monkeypatch.delenv("ENABLE_SMART_ROUTING_V2", raising=False) + with ( + patch("ucode.cli.load_state", return_value=state), + patch("ucode.cli.apply_pat_environment"), + patch("ucode.cli.ensure_bootstrap_dependencies"), + patch("ucode.cli.ensure_provider_state", return_value=state), + patch("ucode.cli.configure_shared_state", return_value=state), + patch("ucode.cli.configure_tool", return_value=state), + patch("ucode.cli._fetch_managed_config", return_value=(managed, False)), + patch("ucode.cli._fetch_budget_recommendation", return_value=None), + patch("ucode.cli._register_managed_mcp_servers"), + patch("ucode.cli._download_managed_skills"), + patch("ucode.cli.launch_agent") as launch, + ): + result = runner.invoke(app, [tool, *pin_args]) + + assert result.exit_code == 0, result.output + options = launch.call_args.kwargs["options"] + assert options.user_pinned_model == "pinned-model" + assert options.launch_smart_routing is False + + def test_codex_short_user_pin_wins_over_managed_default_and_routing(self, monkeypatch): + state = dict(MINIMAL_STATE) + managed = { + "enabled_agents": { + "codex": { + "smart_routing_enabled": True, + "model_config": {"default_model": "managed-model"}, + } + } + } + monkeypatch.delenv("ENABLE_SMART_ROUTING_V2", raising=False) + with ( + patch("ucode.cli.load_state", return_value=state), + patch("ucode.cli.apply_pat_environment"), + patch("ucode.cli.ensure_bootstrap_dependencies"), + patch("ucode.cli.ensure_provider_state", return_value=state), + patch("ucode.cli.configure_shared_state", return_value=state), + patch("ucode.cli.configure_tool", return_value=state), + patch("ucode.cli._fetch_managed_config", return_value=(managed, False)), + patch("ucode.cli._fetch_budget_recommendation", return_value=None), + patch("ucode.cli._register_managed_mcp_servers"), + patch("ucode.cli._download_managed_skills"), + patch("ucode.cli.launch_agent") as launch, + ): + result = runner.invoke(app, ["codex", "--", "-m", "pinned-model"]) + + assert result.exit_code == 0, result.output + options = launch.call_args.kwargs["options"] + assert options.user_pinned_model == "pinned-model" + assert options.launch_smart_routing is False + def test_codex_refresh_is_consumed_by_ucode(self): with patch("ucode.cli._launch_tool") as mock_launch: result = runner.invoke(app, ["codex", "--refresh"]) @@ -741,7 +879,6 @@ def test_claude_v2_skips_legacy_prelaunch_routing(self, monkeypatch): assert result.exit_code == 0, result.output assert mock_configure.call_args.kwargs["route_root_model"] is None - assert "_claude_launch_model" not in mock_launch.call_args.args[1] assert mock_launch.call_args.kwargs["options"].launch_smart_routing is True def test_claude_v2_first_prompt_hook_is_disabled_without_flag(self, monkeypatch): @@ -934,8 +1071,7 @@ def test_model_is_launch_scoped_for_claude(self, monkeypatch): assert mock_configure.call_args.kwargs["custom_model"] is None assert mock_configure.call_args.kwargs["route_root_model"] is None assert ( - mock_launch.call_args.kwargs["options"].claude_launch_model - == "cat.schema.claude-opus-5" + mock_launch.call_args.kwargs["options"].user_pinned_model == "cat.schema.claude-opus-5" ) def test_v2_model_sets_transient_launch_override(self, monkeypatch): @@ -954,7 +1090,7 @@ def test_v2_model_sets_transient_launch_override(self, monkeypatch): result = runner.invoke(app, ["claude", "--model", "system.ai.glm-5-2"]) assert result.exit_code == 0, result.output - assert mock_launch.call_args.args[1]["_claude_launch_model"] == "system.ai.glm-5-2" + assert mock_launch.call_args.kwargs["options"].user_pinned_model == "system.ai.glm-5-2" assert mock_launch.call_args.kwargs["options"].launch_smart_routing is False @staticmethod diff --git a/tests/test_codex_smart_routing_v2.py b/tests/test_codex_smart_routing_v2.py index 11b672df1..f58760135 100644 --- a/tests/test_codex_smart_routing_v2.py +++ b/tests/test_codex_smart_routing_v2.py @@ -40,7 +40,7 @@ def test_smart_routing_switch_message_wraps_to_fixed_width(): class TestLaunchCodex: def test_rejects_unsupported_codex_version(self, monkeypatch): - monkeypatch.setenv(v2.ENV_VAR, "1") + monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_ENV_VAR, "1") monkeypatch.setattr(codex, "clear_model_preferences", lambda state: False) monkeypatch.setattr(codex, "agent_version", lambda binary: "0.144.0") monkeypatch.setattr(v2, "launch_codex", lambda *args, **kwargs: pytest.fail("launched")) @@ -57,7 +57,7 @@ def test_rejects_unsupported_codex_version(self, monkeypatch): ) def test_codex_smart_routing_launch_dispatches_to_v2(self, monkeypatch, tool_args, options): calls = [] - monkeypatch.setenv(v2.ENV_VAR, "1") + monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_ENV_VAR, "1") monkeypatch.setattr(codex, "_smart_routing_config_model", lambda state: "gpt-start") monkeypatch.setattr(codex, "clear_model_preferences", lambda state: False) @@ -100,7 +100,7 @@ def test_codex_launch_bypasses_routing_for_other_shapes(self, tmp_path, monkeypa launches = [] profile_path = tmp_path / "ucode.config.toml" profile_path.write_text('model_provider = "ucode-databricks"\n', encoding="utf-8") - monkeypatch.setenv(v2.ENV_VAR, "1") + monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_ENV_VAR, "1") monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", profile_path) monkeypatch.setattr(codex, "clear_model_preferences", lambda state: False) monkeypatch.setattr(codex, "agent_version", lambda binary: "0.144.0") @@ -118,7 +118,7 @@ def test_codex_launch_bypasses_routing_for_other_shapes(self, tmp_path, monkeypa def test_codex_launch_normalizes_cached_bootstrap_model(self, monkeypatch): calls = [] - monkeypatch.setenv(v2.ENV_VAR, "1") + monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_ENV_VAR, "1") monkeypatch.setattr(codex, "clear_model_preferences", lambda state: False) monkeypatch.setattr(codex, "_smart_routing_config_model", lambda state: None) @@ -157,7 +157,7 @@ def test_startup_config_precedence( managed_path = tmp_path / "managed_config.toml" profile_path = config_home / "ucode.config.toml" user_path = config_home / "config.toml" - monkeypatch.setenv(v2.ENV_VAR, "1") + monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_ENV_VAR, "1") monkeypatch.delenv("CODEX_HOME", raising=False) monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", profile_path) monkeypatch.setattr(codex, "codex_managed_config_path", lambda: managed_path) @@ -496,7 +496,7 @@ def start_interposer(*args, **kwargs): def test_start_model_comes_from_custom_catalog(self, monkeypatch): calls = [] - monkeypatch.setenv(v2.ENV_VAR, "1") + monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_ENV_VAR, "1") monkeypatch.setattr(codex, "clear_model_preferences", lambda state: False) monkeypatch.setattr(codex, "_smart_routing_config_model", lambda state: None) monkeypatch.setattr(codex, "custom_catalog_models", lambda: ["gpt-6-astra", "gpt-6-b"]) @@ -826,3 +826,30 @@ def select_route(workspace, token, task, route_options, resolve, *, router_name, "route_selector": {"router_name": codex_routing.routing.ROUTER_NAME}, } assert "same-oauth-token" not in logged[0] + + +def test_routing_request_deduplicates_equivalent_gpt_spellings(monkeypatch): + captured = {} + + def select_route(workspace, token, task, route_options, resolve, *, router_name, timeout): + captured["route_options"] = list(route_options) + return None, "not selected" + + monkeypatch.setattr(codex_routing.routing, "select_route", select_route) + + codex_routing.request_routing_decision( + WS, + "token", + "Fix the parser", + [ + "system.ai.gpt-5-6-sol", + "gpt-5.6-sol", + "system.ai.gpt-5-6-luna", + "gpt-5.6-luna", + ], + ) + + assert captured["route_options"] == [ + ("gpt-5-6-sol", "codex"), + ("gpt-5-6-luna", "codex"), + ] diff --git a/tests/test_managed_config.py b/tests/test_managed_config.py index 9d0a86e47..02593889f 100644 --- a/tests/test_managed_config.py +++ b/tests/test_managed_config.py @@ -96,6 +96,10 @@ def test_http_headers_map_to_http_headers(self): claude = normalize_managed_config(RAW_MANIFEST)["enabled_agents"]["claude"] assert claude["http_headers"] == {"x-databricks-workspace": "eng-ml-inference"} + def test_smart_routing_maps_to_agent_switch(self): + claude = normalize_managed_config(RAW_MANIFEST)["enabled_agents"]["claude"] + assert claude["smart_routing_enabled"] is True + def test_per_agent_tracing_enabled_is_carried(self): claude = normalize_managed_config(RAW_MANIFEST)["enabled_agents"]["claude"] assert claude["otel_tracing_enabled"] is True diff --git a/tests/test_managed_setup.py b/tests/test_managed_setup.py index 4af3ed8fe..76b38f5f9 100644 --- a/tests/test_managed_setup.py +++ b/tests/test_managed_setup.py @@ -422,6 +422,13 @@ def test_slots_serialize_into_the_default_models_map(self): "default_opus_model": "system.ai.claude-opus-4-8", } + def test_smart_routing_serializes_into_agent_config(self): + manifest = { + "enabled_agents": {"codex": {"smart_routing_enabled": True}}, + } + payload = serialize_managed_config(manifest) + assert payload["enabled_agents"][0]["config"]["smart_routing"] == {"enabled": True} + class TestClaudeFamilyCandidates: """Discovery keeps one id per family for the launch path; authoring needs the alternatives."""