From 586eeef1bfa27ce472463968e44c1e2e8e1cc293 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Fri, 4 Sep 2026 23:09:52 +0000 Subject: [PATCH 1/5] only do smart routing for 2 cases: no subcommand or the -- --- src/ucode/agents/__init__.py | 11 +++- src/ucode/agents/args.py | 10 ++++ src/ucode/agents/claude.py | 13 +++-- src/ucode/agents/codex.py | 81 +++++++++++++++++++--------- src/ucode/agents/copilot.py | 4 +- src/ucode/agents/gemini.py | 4 +- src/ucode/agents/opencode.py | 4 +- src/ucode/agents/pi.py | 4 +- src/ucode/cli.py | 60 ++++++++++++++++++--- tests/test_agent_claude.py | 39 +++++++++++--- tests/test_agent_codex.py | 47 ++++++++++++---- tests/test_agents_init.py | 15 ++++++ tests/test_cli.py | 23 ++++++++ tests/test_codex_smart_routing_v2.py | 56 +++++++++++++++++-- 14 files changed, 306 insertions(+), 65 deletions(-) diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index 232c65cf..0c82e4fa 100644 --- a/src/ucode/agents/__init__.py +++ b/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 = { @@ -447,8 +448,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: diff --git a/src/ucode/agents/args.py b/src/ucode/agents/args.py index b7092e97..05d70909 100644 --- a/src/ucode/agents/args.py +++ b/src/ucode/agents/args.py @@ -2,6 +2,16 @@ from __future__ import annotations +from dataclasses import dataclass + + +@dataclass(frozen=True) +class LaunchOptions: + """Invocation-scoped options shared by agent launchers.""" + + smart_routing: bool = False + explicit_prompt: bool = False + def explicit_model_arg_value(tool_args: list[str]) -> str | None: """Return the last model selected before the harness's ``--`` separator.""" diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 01b3a1f9..01212922 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -54,7 +54,7 @@ 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" @@ -1377,19 +1377,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() + options.smart_routing 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) + and (options.explicit_prompt or _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": diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 625329f3..5649e570 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -52,6 +52,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" @@ -494,35 +496,34 @@ def clear_model_preferences(state: dict) -> bool: _PROFILE_REJECTED_MAX_SECONDS = 3.0 -def launch(state: dict, tool_args: list[str]) -> None: +def should_use_smart_routing(tool_args: list[str], *, explicit_prompt: bool = False) -> bool: + """Return whether this invocation explicitly selects the routed TUI path. + + Smart routing is intentionally opt-in by invocation shape: a bare Codex + launch, or a single prompt passed after ucode's explicit ``--`` boundary. + Commands and options (including ``--model``) use the ordinary launcher. + """ + return not tool_args or ( + explicit_prompt + and len(tool_args) <= 1 + and not any(arg in {"-m", "--model"} or arg.startswith("--model=") for arg in tool_args) + ) + + +def launch( + state: dict, + tool_args: list[str], + *, + options: LaunchOptions, +) -> None: + if options.smart_routing and should_use_smart_routing( + tool_args, explicit_prompt=options.explicit_prompt + ): + _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"]: @@ -575,6 +576,34 @@ def _app_server_start_model() -> str: sys.exit(returncode) +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: """Disable routing and remove only ucode's Codex routing hooks.""" state.pop(SMART_ROUTING_STATE_KEY, None) diff --git a/src/ucode/agents/copilot.py b/src/ucode/agents/copilot.py index 3e876fbc..44b1df83 100644 --- a/src/ucode/agents/copilot.py +++ b/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) diff --git a/src/ucode/agents/gemini.py b/src/ucode/agents/gemini.py index 5f14f1a3..5dc81b5c 100644 --- a/src/ucode/agents/gemini.py +++ b/src/ucode/agents/gemini.py @@ -33,6 +33,8 @@ ) from ucode.telemetry import agent_version, ucode_version +from .args import LaunchOptions + GEMINI_CONFIG_DIR = Path.home() / ".gemini" GEMINI_ENV_PATH = GEMINI_CONFIG_DIR / "ucode.env" GEMINI_BACKUP_PATH = APP_DIR / "gemini-ucode-env.backup" @@ -228,7 +230,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: provider = get_provider_service(state, "gemini") token = _refresh_token_once(state) model = _launch_model(state, provider) diff --git a/src/ucode/agents/opencode.py b/src/ucode/agents/opencode.py index b07aef56..114e28f4 100644 --- a/src/ucode/agents/opencode.py +++ b/src/ucode/agents/opencode.py @@ -27,6 +27,8 @@ from ucode.state import mark_tool_managed, save_state from ucode.telemetry import agent_version, ucode_version +from .args import LaunchOptions + OPENCODE_XDG_CONFIG_HOME = APP_DIR / "opencode-xdg" OPENCODE_CONFIG_DIR = OPENCODE_XDG_CONFIG_HOME / "opencode" OPENCODE_CONFIG_PATH = OPENCODE_CONFIG_DIR / "opencode.json" @@ -391,7 +393,7 @@ def build_runtime_env(token: str, state: dict | None = None) -> dict[str, str]: return env -def launch(state: dict, tool_args: list[str]) -> None: +def launch(state: dict, tool_args: list[str], *, options: LaunchOptions) -> None: """Launch OpenCode with on-demand token refresh from its local plugin.""" token = _configure_launch(state) env = build_runtime_env(token, state) diff --git a/src/ucode/agents/pi.py b/src/ucode/agents/pi.py index a673a548..a193e8ff 100644 --- a/src/ucode/agents/pi.py +++ b/src/ucode/agents/pi.py @@ -50,6 +50,8 @@ from ucode.state import mark_tool_managed, save_state from ucode.telemetry import agent_version, ucode_version +from .args import LaunchOptions + PI_UCODE_HOME = APP_DIR / "pi-home" PI_CONFIG_DIR = PI_UCODE_HOME / ".pi" / "agent" PI_CONFIG_PATH = PI_CONFIG_DIR / "models.json" @@ -281,7 +283,7 @@ def build_runtime_env(token: str) -> dict[str, str]: return env -def launch(state: dict, tool_args: list[str]) -> None: +def launch(state: dict, tool_args: list[str], *, options: LaunchOptions) -> None: token = _refresh_token_once(state) env = build_runtime_env(token) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 1f217b0c..7ec50b61 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -12,10 +12,13 @@ from typing import Annotated import typer +from click import Context as ClickContext from rich.panel import Panel +from typer.core import TyperCommand from ucode.agents import ( TOOL_SPECS, + LaunchOptions, check_gateway_endpoint, configure_selected_tools, configure_single_tool, @@ -1983,6 +1986,7 @@ def _can_launch_from_cached_config( model: str | None, explicit_provider: str | None, workspace_url: str | None, + smart_routing_enabled: bool | None = None, ) -> bool: """Return whether a normal Claude/Codex launch can use its cached config.""" if tool not in CAN_USE_CACHED_CONFIG_AGENTS: @@ -1991,7 +1995,10 @@ def _can_launch_from_cached_config( if refresh or model or explicit_provider is not None: return False - if tool == "codex" and smart_routing_v2.enabled(): + if smart_routing_enabled is None: + smart_routing_enabled = smart_routing_v2.enabled() + + if tool == "codex" and smart_routing_enabled: if not state.get("codex_models") or not state.get("oss_models"): return False @@ -2029,6 +2036,12 @@ def _launch_tool( ) -> None: try: tool = normalize_tool(tool_name) + explicit_prompt = _has_explicit_prompt(ctx) + smart_routing_enabled = smart_routing_v2.enabled() + launch_options = LaunchOptions( + smart_routing=smart_routing_enabled, + explicit_prompt=explicit_prompt, + ) # 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. @@ -2069,12 +2082,13 @@ def _launch_tool( model=model, explicit_provider=explicit_provider, workspace_url=workspace_url, + smart_routing_enabled=smart_routing_enabled, ): print_section(_launch_title(tool)) if forwarded_model: print_kv("Model", forwarded_model) print_success(f"Starting {TOOL_SPECS[tool]['display']}") - launch_agent(tool, state, ctx.args) + launch_agent(tool, state, ctx.args, options=launch_options) return # Fetched before `configure_shared_state` because it decides whether this agent may launch # at all and whether the model discovery below can be skipped. @@ -2131,7 +2145,7 @@ def _launch_tool( provider = managed_provider # Checked after the managed config settles `provider`: an admin-set provider must trip this # guard too, or routing would be persisted as on while a provider is active. - if tool in CAN_USE_CACHED_CONFIG_AGENTS and smart_routing_v2.enabled() and provider: + if tool in CAN_USE_CACHED_CONFIG_AGENTS and smart_routing_enabled and provider: raise RuntimeError( f"{TOOL_SPECS[tool]['display']} smart routing cannot be enabled with " "--provider. Launch without a Model Provider Service and try again." @@ -2254,7 +2268,7 @@ def _launch_tool( print_kv("Model", route_root_model) elif resolved_model: print_kv("Model", resolved_model) - if tool in CAN_USE_CACHED_CONFIG_AGENTS and smart_routing_v2.enabled() and not provider: + if tool in CAN_USE_CACHED_CONFIG_AGENTS and smart_routing_enabled and not provider: print_kv("Smart routing", "enabled") print_note( f"{TOOL_SPECS[tool]['display']} may require one-time hook review. Open " @@ -2285,7 +2299,7 @@ def _launch_tool( if provider: state["_claude_launch_provider"] = provider print_success(f"Starting {TOOL_SPECS[tool]['display']}") - launch_agent(tool, state, ctx.args) + launch_agent(tool, state, ctx.args, options=launch_options) except RuntimeError as exc: print_err(str(exc)) raise typer.Exit(1) from None @@ -2351,6 +2365,30 @@ def _disable_managed_config_if_requested(skip_managed_config: bool) -> None: ] +_PROMPT_SUFFIX_KEY = "ucode_explicit_prompt_suffix" + + +class _PromptAwareCommand(TyperCommand): + """Record an agent's ``--`` prompt separator before Click removes it.""" + + def parse_args(self, ctx: ClickContext, args: list[str]) -> list[str]: + try: + separator = args.index("--") + except ValueError: + pass + else: + ctx.meta[_PROMPT_SUFFIX_KEY] = tuple(args[separator + 1 :]) + return super().parse_args(ctx, args) + + +def _has_explicit_prompt(ctx: typer.Context) -> bool: + suffix = ctx.meta.get(_PROMPT_SUFFIX_KEY) + if not isinstance(suffix, tuple): + return False + suffix_args = list(suffix) + return ctx.args == suffix_args and len(suffix_args) <= 1 + + @app.callback(invoke_without_command=True) def default( ctx: typer.Context, @@ -2469,7 +2507,11 @@ def _print_no_managed_config_guidance(workspace: str, profile: str | None) -> No print_note("Run `ug setup` to configure one for your workspace, then `ug publish`.") -@app.command("codex", context_settings={"allow_extra_args": True, "ignore_unknown_options": True}) +@app.command( + "codex", + cls=_PromptAwareCommand, + context_settings={"allow_extra_args": True, "ignore_unknown_options": True}, +) def codex_cmd( ctx: typer.Context, provider: Annotated[ @@ -2527,7 +2569,11 @@ def codex_cmd( ) -@app.command("claude", context_settings={"allow_extra_args": True, "ignore_unknown_options": True}) +@app.command( + "claude", + cls=_PromptAwareCommand, + context_settings={"allow_extra_args": True, "ignore_unknown_options": True}, +) def claude_cmd( ctx: typer.Context, provider: Annotated[ diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 61f07438..25e36351 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -9,7 +9,7 @@ import pytest -from ucode.agents import claude +from ucode.agents import LaunchOptions, claude from ucode.smart_routing import claude_routing, v2 from ucode.state import MANAGED_OVERLAY_KEY @@ -1073,6 +1073,7 @@ def start_proxy(workspace, profile, port, token_header, force_refresh_near_expir "relayed_proxy_port": 12345, }, ["--debug"], + options=LaunchOptions(), ) assert exc.value.code == 0 @@ -1094,7 +1095,11 @@ def test_smart_routing_on_windows_is_not_supported(self, monkeypatch): RuntimeError, match="Smart routing in Claude Code is currently not supported on Windows", ): - claude.launch({"workspace": WS, "profile": "test"}, ["--debug"]) + claude.launch( + {"workspace": WS, "profile": "test"}, + ["--debug"], + options=LaunchOptions(smart_routing=True), + ) def test_default_launch_keeps_existing_auth_path(self, monkeypatch): calls: list[list[str]] = [] @@ -1104,7 +1109,7 @@ def test_default_launch_keeps_existing_auth_path(self, monkeypatch): monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "token") monkeypatch.setattr(claude, "exec_or_spawn", lambda argv: calls.append(argv)) - claude.launch({"workspace": WS, "profile": "test"}, ["--debug"]) + claude.launch({"workspace": WS, "profile": "test"}, ["--debug"], options=LaunchOptions()) assert os.environ["OAUTH_TOKEN"] == "token" assert calls == [["claude", "--settings", str(claude.CLAUDE_SETTINGS_PATH), "--debug"]] @@ -1119,6 +1124,7 @@ def test_v2_launch_override_bypasses_first_prompt_routing(self, monkeypatch): claude.launch( {"workspace": WS, "_claude_launch_model": "system.ai.glm-5-2"}, ["--debug"], + options=LaunchOptions(smart_routing=True), ) assert calls == [["claude", "--settings", str(claude.CLAUDE_SETTINGS_PATH), "--debug"]] @@ -1138,7 +1144,7 @@ def test_v2_explicit_claude_model_bypasses_first_prompt_routing(self, monkeypatc monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "token") monkeypatch.setattr(claude, "exec_or_spawn", lambda argv: calls.append(argv)) - claude.launch({"workspace": WS}, tool_args) + claude.launch({"workspace": WS}, tool_args, options=LaunchOptions(smart_routing=True)) assert calls == [["claude", "--settings", str(claude.CLAUDE_SETTINGS_PATH), *tool_args]] v2.launch_claude.assert_not_called() @@ -1157,7 +1163,11 @@ def test_v2_provider_launch_bypasses_first_prompt_routing(self, monkeypatch, pro monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "token") monkeypatch.setattr(claude, "exec_or_spawn", lambda argv: calls.append(argv)) - claude.launch({"workspace": WS, **provider_state}, ["--debug"]) + claude.launch( + {"workspace": WS, **provider_state}, + ["--debug"], + options=LaunchOptions(smart_routing=True), + ) assert calls == [["claude", "--settings", str(claude.CLAUDE_SETTINGS_PATH), "--debug"]] v2.launch_claude.assert_not_called() @@ -1176,7 +1186,7 @@ def test_v2_noninteractive_launch_bypasses_first_prompt_routing(self, monkeypatc monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "token") monkeypatch.setattr(claude, "exec_or_spawn", lambda argv: calls.append(argv)) - claude.launch({"workspace": WS}, tool_args) + claude.launch({"workspace": WS}, tool_args, options=LaunchOptions(smart_routing=True)) assert calls == [["claude", "--settings", str(claude.CLAUDE_SETTINGS_PATH), *tool_args]] v2.launch_claude.assert_not_called() @@ -1188,7 +1198,7 @@ def test_v2_positional_prompt_uses_first_prompt_routing(self, monkeypatch, tool_ monkeypatch.setattr(claude, "_original_launch_model", lambda _state: None) monkeypatch.setattr(v2, "launch_claude", launch_v2) - claude.launch({"workspace": WS}, tool_args) + claude.launch({"workspace": WS}, tool_args, options=LaunchOptions(smart_routing=True)) launch_v2.assert_called_once_with( {"workspace": WS}, @@ -1201,6 +1211,19 @@ def test_v2_positional_prompt_uses_first_prompt_routing(self, monkeypatch, tool_ model_name=claude._maybe_add_1m_suffix, ) + def test_v2_explicit_prompt_overrides_subcommand_shaped_text(self, monkeypatch): + launch_v2 = Mock() + monkeypatch.setattr(claude, "_original_launch_model", lambda _state: None) + monkeypatch.setattr(v2, "launch_claude", launch_v2) + + claude.launch( + {"workspace": WS}, + ["doctor"], + options=LaunchOptions(smart_routing=True, explicit_prompt=True), + ) + + launch_v2.assert_called_once() + def test_v2_does_not_treat_option_value_as_positional_argument(self): assert claude._uses_interactive_tui(["--name", "doctor"]) is True @@ -1215,7 +1238,7 @@ def test_gateway_discovery_uses_direct_gateway(self, monkeypatch): monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "token") monkeypatch.setattr(claude, "exec_or_spawn", lambda argv: calls.append(argv)) - claude.launch({"workspace": WS, "profile": "test"}, ["--debug"]) + claude.launch({"workspace": WS, "profile": "test"}, ["--debug"], options=LaunchOptions()) assert os.environ["OAUTH_TOKEN"] == "token" assert os.environ["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1" diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index 2cd11dc3..91f35f42 100644 --- a/tests/test_agent_codex.py +++ b/tests/test_agent_codex.py @@ -7,7 +7,7 @@ import pytest -from ucode.agents import codex +from ucode.agents import LaunchOptions, codex from ucode.config_io import read_toml_safe from ucode.smart_routing import codex_routing @@ -25,6 +25,35 @@ def test_display(self): assert codex.SPEC["display"] == "Codex" +class TestShouldUseSmartRouting: + @pytest.mark.parametrize( + ("tool_args", "explicit_prompt"), + [([], False), ([], True), (["fix this"], True)], + ) + def test_accepts_only_explicit_tui_shapes(self, tool_args, explicit_prompt): + assert codex.should_use_smart_routing(tool_args, explicit_prompt=explicit_prompt) is True + + @pytest.mark.parametrize( + "tool_args", + [ + ["exec", "fix this"], + ["review"], + ["app-server"], + ["update"], + ["--model", "gpt-5.6-sol"], + ["--model", "gpt-5.6-sol", "--", "fix this"], + ["fix this"], + ["one", "two"], + ], + ) + def test_rejects_all_other_invocation_shapes(self, tool_args): + assert codex.should_use_smart_routing(tool_args) is False + + @pytest.mark.parametrize("tool_args", [["--model=x"], ["--model"], ["-m"]]) + def test_explicit_model_bypasses_separator_opt_in(self, tool_args): + assert codex.should_use_smart_routing(tool_args, explicit_prompt=True) is False + + class TestHasUcodeConfig: def test_detects_profile_config(self, tmp_path, monkeypatch): config_path = tmp_path / "ucode.config.toml" @@ -575,7 +604,7 @@ def test_sets_oauth_token_and_runs_with_profile(self, monkeypatch): codex, "get_databricks_token", lambda workspace, profile=None: "fresh-token" ) with pytest.raises(SystemExit) as exc: - codex.launch({"workspace": WS}, ["--search"]) + codex.launch({"workspace": WS}, ["--search"], options=LaunchOptions()) assert exc.value.code == 0 assert os.environ["OAUTH_TOKEN"] == "fresh-token" assert runs == [["codex", "--profile", "ucode", "--search"]] @@ -584,7 +613,7 @@ def test_sets_oauth_token_and_runs_with_profile(self, monkeypatch): def test_success_propagates_exit_without_retry(self, monkeypatch): runs, fallbacks = self._patch(monkeypatch, returncode=0, elapsed=0.2) with pytest.raises(SystemExit) as exc: - codex.launch({"workspace": WS}, ["exec", "hi"]) + codex.launch({"workspace": WS}, ["exec", "hi"], options=LaunchOptions()) assert exc.value.code == 0 assert runs == [["codex", "--profile", "ucode", "exec", "hi"]] assert fallbacks == [] @@ -602,7 +631,7 @@ def test_app_layers_ucode_profile_as_config_overrides(self, tmp_path, monkeypatc monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", profile_path) runs, launches = self._patch(monkeypatch, returncode=0, elapsed=0.2) - codex.launch({"workspace": WS}, ["app", "--new-window"]) + codex.launch({"workspace": WS}, ["app", "--new-window"], options=LaunchOptions()) assert runs == [] assert launches[0][:2] == ["codex", "app"] @@ -619,7 +648,7 @@ def test_app_requires_populated_ucode_profile(self, tmp_path, monkeypatch): runs, launches = self._patch(monkeypatch, returncode=0, elapsed=0.2) with pytest.raises(RuntimeError, match="ucode configure --agents codex"): - codex.launch({"workspace": WS}, ["app"]) + codex.launch({"workspace": WS}, ["app"], options=LaunchOptions()) assert runs == [] assert launches == [] @@ -629,7 +658,7 @@ def test_fast_failure_relaunches_without_profile(self, monkeypatch): # a fast nonzero exit → relaunch without --profile. for args in (["app-server", "--listen", "u"], ["mcp-server"]): runs, fallbacks = self._patch(monkeypatch, returncode=1, elapsed=0.15) - codex.launch({"workspace": WS}, args) + codex.launch({"workspace": WS}, args, options=LaunchOptions()) assert runs == [["codex", "--profile", "ucode", *args]] assert fallbacks == [["codex", *args]] @@ -644,7 +673,7 @@ def test_fallback_warns_on_stderr_before_handoff(self, monkeypatch, capsys): "exec_or_spawn", lambda argv: warned_before_handoff.append(capsys.readouterr()), ) - codex.launch({"workspace": WS}, ["app-server"]) + codex.launch({"workspace": WS}, ["app-server"], options=LaunchOptions()) # execvp replaces the process, so the warning must already be out by then. assert len(warned_before_handoff) == 1 @@ -666,7 +695,7 @@ def test_slow_failure_does_not_retry(self, monkeypatch): # (relaunching the user's prompt on their own OpenAI login). runs, fallbacks = self._patch(monkeypatch, returncode=1, elapsed=8.0) with pytest.raises(SystemExit) as exc: - codex.launch({"workspace": WS}, ["exec", "hi"]) + codex.launch({"workspace": WS}, ["exec", "hi"], options=LaunchOptions()) assert exc.value.code == 1 assert fallbacks == [] @@ -674,7 +703,7 @@ def test_fast_success_does_not_retry(self, monkeypatch): # The retry is gated on a *nonzero* exit; a fast clean exit just returns. runs, fallbacks = self._patch(monkeypatch, returncode=0, elapsed=0.15) with pytest.raises(SystemExit) as exc: - codex.launch({"workspace": WS}, []) + codex.launch({"workspace": WS}, [], options=LaunchOptions()) assert exc.value.code == 0 assert fallbacks == [] diff --git a/tests/test_agents_init.py b/tests/test_agents_init.py index 1b98ce57..0935ac94 100644 --- a/tests/test_agents_init.py +++ b/tests/test_agents_init.py @@ -11,6 +11,7 @@ from ucode.agents import ( DEFAULT_TOOL, TOOL_SPECS, + LaunchOptions, check_gateway_endpoint, configure_selected_tools, default_model_for_tool, @@ -96,6 +97,20 @@ def test_tool_update_available_uses_agent_override(self, monkeypatch): assert agents_mod.tool_update_available("opencode") == ("1.18.15", "1.18.16") +def test_launch_dispatches_invocation_options(monkeypatch): + calls = [] + options = LaunchOptions(smart_routing=True) + monkeypatch.setattr( + agents_mod.codex, + "launch", + lambda state, tool_args, *, options: calls.append((state, tool_args, options)), + ) + + agents_mod.launch("codex", {"workspace": "ws"}, ["prompt"], options=options) + + assert calls == [({"workspace": "ws"}, ["prompt"], options)] + + class TestInstallAiToolsForAgents: def _capture(self, monkeypatch): captured = {} diff --git a/tests/test_cli.py b/tests/test_cli.py index 9e3e05df..e9d68e53 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -17,6 +17,7 @@ import pytest from typer.testing import CliRunner +import ucode.cli as cli_mod import ucode.databricks as db_mod from ucode.cli import app from ucode.databricks import GatewayProbe @@ -464,6 +465,28 @@ def capture(_tool, ctx, **_kwargs): assert result.exit_code == 0, result.output assert captured == [("1", ["fix the parser"])] + @pytest.mark.parametrize( + ("args", "forwarded", "has_separator"), + [ + (["codex", "--", "fix the parser"], ["fix the parser"], True), + (["codex", "--"], [], True), + (["claude", "--", "doctor"], ["doctor"], True), + ( + ["codex", "--model", "gpt-5.6-sol", "--", "fix the parser"], + ["--model", "gpt-5.6-sol", "fix the parser"], + False, + ), + ], + ) + def test_agent_records_prompt_separator(self, args, forwarded, has_separator): + with patch("ucode.cli._launch_tool") as mock_launch: + result = runner.invoke(app, args) + + assert result.exit_code == 0, result.output + ctx = mock_launch.call_args.args[1] + assert ctx.args == forwarded + assert cli_mod._has_explicit_prompt(ctx) is has_separator + def test_codex_refresh_is_consumed_by_ucode(self): with patch("ucode.cli._launch_tool") as mock_launch: result = runner.invoke(app, ["codex", "--refresh"]) diff --git a/tests/test_codex_smart_routing_v2.py b/tests/test_codex_smart_routing_v2.py index 58e72126..15c8e033 100644 --- a/tests/test_codex_smart_routing_v2.py +++ b/tests/test_codex_smart_routing_v2.py @@ -4,7 +4,7 @@ import pytest -from ucode.agents import codex +from ucode.agents import LaunchOptions, codex from ucode.smart_routing import codex_interposer, codex_routing, v2 WS = "https://example.databricks.com" @@ -31,9 +31,19 @@ def test_rejects_unsupported_codex_version(self, monkeypatch): monkeypatch.setattr(v2, "launch_codex", lambda *args, **kwargs: pytest.fail("launched")) with pytest.raises(RuntimeError, match="requires Codex 0.145.0 or newer; found 0.144.0"): - codex.launch({"workspace": WS}, []) + codex.launch({"workspace": WS}, [], options=LaunchOptions(smart_routing=True)) - def test_codex_launch_dispatches_when_flag_enabled(self, monkeypatch): + @pytest.mark.parametrize( + ("tool_args", "options"), + [ + ([], LaunchOptions(smart_routing=True)), + ( + ["fix the parser"], + LaunchOptions(smart_routing=True, explicit_prompt=True), + ), + ], + ) + def test_codex_smart_routing_launch_dispatches_to_v2(self, monkeypatch, tool_args, options): calls = [] monkeypatch.setenv(v2.ENV_VAR, "1") monkeypatch.setattr(codex, "default_model", lambda state: "gpt-start") @@ -47,13 +57,13 @@ def launch_v2(state, tool_args, **kwargs): state = {"workspace": WS} with pytest.raises(SystemExit) as exc: - codex.launch(state, ["--search"]) + codex.launch(state, tool_args, options=options) assert exc.value.code == 0 assert calls == [ ( state, - ["--search"], + tool_args, { "binary": "codex", "start_model": "gpt-start", @@ -62,6 +72,41 @@ def launch_v2(state, tool_args, **kwargs): ) ] + @pytest.mark.parametrize( + "tool_args", + [ + ["exec", "fix this"], + ["review"], + ["app-server"], + ["update"], + ["--model", "gpt-5.6-sol"], + ["--model", "gpt-5.6-sol", "--", "fix this"], + ["fix this"], + ], + ) + def test_codex_launch_bypasses_routing_for_other_shapes(self, monkeypatch, tool_args): + runs = [] + monkeypatch.setenv(v2.ENV_VAR, "1") + monkeypatch.setattr(codex, "clear_model_preferences", lambda state: False) + monkeypatch.setattr(codex, "agent_version", lambda binary: "0.144.0") + monkeypatch.setattr(codex, "get_databricks_token", lambda *_args: "token") + monkeypatch.setattr(v2, "launch_codex", lambda *args, **kwargs: pytest.fail("launched")) + monkeypatch.setattr( + codex.subprocess, + "run", + lambda argv: runs.append(argv) or codex.subprocess.CompletedProcess(argv, 0), + ) + + with pytest.raises(SystemExit) as exc: + codex.launch( + {"workspace": WS}, + tool_args, + options=LaunchOptions(smart_routing=True), + ) + + assert exc.value.code == 0 + assert runs == [["codex", "--profile", "ucode", *tool_args]] + def test_codex_launch_normalizes_cached_bootstrap_model(self, monkeypatch): calls = [] monkeypatch.setenv(v2.ENV_VAR, "1") @@ -78,6 +123,7 @@ def launch_v2(state, tool_args, **kwargs): codex.launch( {"workspace": WS, "codex_models": ["system.ai.gpt-5-6-luna"]}, [], + options=LaunchOptions(smart_routing=True), ) assert calls[0]["start_model"] == "gpt-5.6-luna" From 9cd4c7ed8b14fd25e901a91ab3801123d44a3063 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Sat, 5 Sep 2026 00:19:42 +0000 Subject: [PATCH 2/5] lint --- src/ucode/agents/args.py | 3 +- src/ucode/agents/claude.py | 65 +--------- src/ucode/agents/codex.py | 89 ++------------ src/ucode/cli.py | 50 +++++++- src/ucode/codex_config.py | 32 +++-- tests/test_agent_claude.py | 87 ++----------- tests/test_agent_codex.py | 176 ++++++--------------------- tests/test_agents_init.py | 2 +- tests/test_cli.py | 27 ++++ tests/test_codex_config.py | 31 +++++ tests/test_codex_smart_routing_v2.py | 38 +++--- 11 files changed, 200 insertions(+), 400 deletions(-) diff --git a/src/ucode/agents/args.py b/src/ucode/agents/args.py index 05d70909..064eceda 100644 --- a/src/ucode/agents/args.py +++ b/src/ucode/agents/args.py @@ -9,8 +9,7 @@ class LaunchOptions: """Invocation-scoped options shared by agent launchers.""" - smart_routing: bool = False - explicit_prompt: bool = False + launch_smart_routing: bool = False 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 01212922..08263f65 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -48,7 +48,6 @@ 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 @@ -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: @@ -1195,11 +1174,6 @@ 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( @@ -1207,33 +1181,6 @@ def _has_provider_launch(state: dict) -> bool: ) -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 [] @@ -1388,21 +1335,13 @@ def launch( if state.get("claude_relayed"): _launch_relayed(state, binary, tool_args) return - first_prompt_routing = ( - options.smart_routing - 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 (options.explicit_prompt or _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, diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 5649e570..e3854901 100644 --- a/src/ucode/agents/codex.py +++ b/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 @@ -485,40 +482,13 @@ 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 should_use_smart_routing(tool_args: list[str], *, explicit_prompt: bool = False) -> bool: - """Return whether this invocation explicitly selects the routed TUI path. - - Smart routing is intentionally opt-in by invocation shape: a bare Codex - launch, or a single prompt passed after ucode's explicit ``--`` boundary. - Commands and options (including ``--model``) use the ordinary launcher. - """ - return not tool_args or ( - explicit_prompt - and len(tool_args) <= 1 - and not any(arg in {"-m", "--model"} or arg.startswith("--model=") for arg in tool_args) - ) - - def launch( state: dict, tool_args: list[str], *, options: LaunchOptions, ) -> None: - if options.smart_routing and should_use_smart_routing( - tool_args, explicit_prompt=options.explicit_prompt - ): + if options.launch_smart_routing: _launch_smart_routing(state, tool_args) return clear_model_preferences(state) @@ -526,54 +496,17 @@ def launch( workspace = state.get("workspace") 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: diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 0275b447..b4c4062c 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -12,8 +12,8 @@ from typing import Annotated import typer -from click import Context as ClickContext from rich.panel import Panel +from typer._click import Context as ClickContext from typer.core import TyperCommand from ucode.agents import ( @@ -2023,6 +2023,32 @@ def _can_launch_from_cached_config( return codex_agent.has_ucode_config() and codex_agent.managed_config_is_current(state) +def _launch_options( + tool: str, + tool_args: list[str], + *, + smart_routing_enabled: bool, + explicit_prompt: bool, + model: str | None, + provider: str | None, +) -> LaunchOptions: + has_model_override = model is not None or explicit_model_arg_value(tool_args) is not None + return LaunchOptions( + launch_smart_routing=( + # Smart routing is enabled globally. + smart_routing_enabled + # Only Claude Code and Codex currently support smart routing. + and tool in CAN_USE_CACHED_CONFIG_AGENTS + # An explicit model selection must bypass smart routing. + and not has_model_override + # Smart routing does not currently support Model Provider Services. + and provider is None + # Route a bare agent launch or a prompt explicitly passed after `--`. + and (not tool_args or explicit_prompt) + ) + ) + + def _launch_tool( tool_name: str, ctx: typer.Context, @@ -2038,10 +2064,6 @@ def _launch_tool( tool = normalize_tool(tool_name) explicit_prompt = _has_explicit_prompt(ctx) smart_routing_enabled = smart_routing_v2.enabled() - launch_options = LaunchOptions( - smart_routing=smart_routing_enabled, - explicit_prompt=explicit_prompt, - ) # 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. @@ -2075,6 +2097,14 @@ def _launch_tool( # back to whatever `ug configure` saved for this tool. provider = provider or get_provider_service(state, tool) state = _migrate_legacy_smart_routing(state) + launch_options = _launch_options( + tool, + ctx.args, + smart_routing_enabled=smart_routing_enabled, + explicit_prompt=explicit_prompt, + model=model, + provider=provider, + ) if _can_launch_from_cached_config( tool, state, @@ -2082,7 +2112,7 @@ def _launch_tool( model=model, explicit_provider=explicit_provider, workspace_url=workspace_url, - smart_routing_enabled=smart_routing_enabled, + smart_routing_enabled=launch_options.launch_smart_routing, ): print_section(_launch_title(tool)) if forwarded_model: @@ -2298,6 +2328,14 @@ def _launch_tool( state["_claude_launch_model"] = launch_model if provider: state["_claude_launch_provider"] = provider + launch_options = _launch_options( + tool, + ctx.args, + smart_routing_enabled=smart_routing_enabled, + explicit_prompt=explicit_prompt, + model=model or (route_root_model if tool == "claude" else None), + provider=provider, + ) print_success(f"Starting {TOOL_SPECS[tool]['display']}") launch_agent(tool, state, ctx.args, options=launch_options) except RuntimeError as exc: diff --git a/src/ucode/codex_config.py b/src/ucode/codex_config.py index 25070cce..bbf3ca90 100644 --- a/src/ucode/codex_config.py +++ b/src/ucode/codex_config.py @@ -2,20 +2,30 @@ from __future__ import annotations +from collections.abc import Mapping + import tomlkit +from tomlkit.items import Item + + +def _toml_item(value: object) -> Item: + if isinstance(value, Mapping): + inline = tomlkit.inline_table() + for key, child in value.items(): + inline[str(key)] = _toml_item(child) + return inline + if isinstance(value, list): + array = tomlkit.array() + for child in value: + array.append(_toml_item(child)) + return array + if isinstance(value, Item): + return value + return tomlkit.item(value) -def _toml_value(value: str | int | float | bool | list[object] | dict[str, object]) -> str: - if isinstance(value, dict): - 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() +def _toml_value(value: object) -> str: + return _toml_item(value).as_string() def codex_config_args(config: dict) -> list[str]: diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 25e36351..11a9c5b4 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -1098,7 +1098,7 @@ def test_smart_routing_on_windows_is_not_supported(self, monkeypatch): claude.launch( {"workspace": WS, "profile": "test"}, ["--debug"], - options=LaunchOptions(smart_routing=True), + options=LaunchOptions(launch_smart_routing=True), ) def test_default_launch_keeps_existing_auth_path(self, monkeypatch): @@ -1114,64 +1114,6 @@ def test_default_launch_keeps_existing_auth_path(self, monkeypatch): assert os.environ["OAUTH_TOKEN"] == "token" assert calls == [["claude", "--settings", str(claude.CLAUDE_SETTINGS_PATH), "--debug"]] - def test_v2_launch_override_bypasses_first_prompt_routing(self, monkeypatch): - calls: list[list[str]] = [] - monkeypatch.setenv(v2.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)) - - claude.launch( - {"workspace": WS, "_claude_launch_model": "system.ai.glm-5-2"}, - ["--debug"], - options=LaunchOptions(smart_routing=True), - ) - - assert calls == [["claude", "--settings", str(claude.CLAUDE_SETTINGS_PATH), "--debug"]] - v2.launch_claude.assert_not_called() - - @pytest.mark.parametrize( - "tool_args", - [ - ["-m", "opus"], - ["--model=opus"], - ], - ) - def test_v2_explicit_claude_model_bypasses_first_prompt_routing(self, monkeypatch, tool_args): - calls: list[list[str]] = [] - monkeypatch.setenv(v2.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)) - - claude.launch({"workspace": WS}, tool_args, options=LaunchOptions(smart_routing=True)) - - assert calls == [["claude", "--settings", str(claude.CLAUDE_SETTINGS_PATH), *tool_args]] - v2.launch_claude.assert_not_called() - - @pytest.mark.parametrize( - "provider_state", - [ - {"provider_services": {"claude": "main.default.anthropic"}}, - {"_claude_launch_provider": "main.default.anthropic"}, - ], - ) - def test_v2_provider_launch_bypasses_first_prompt_routing(self, monkeypatch, provider_state): - calls: list[list[str]] = [] - monkeypatch.setenv(v2.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)) - - claude.launch( - {"workspace": WS, **provider_state}, - ["--debug"], - options=LaunchOptions(smart_routing=True), - ) - - assert calls == [["claude", "--settings", str(claude.CLAUDE_SETTINGS_PATH), "--debug"]] - v2.launch_claude.assert_not_called() - @pytest.mark.parametrize( "tool_args", [ @@ -1186,7 +1128,7 @@ def test_v2_noninteractive_launch_bypasses_first_prompt_routing(self, monkeypatc monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "token") monkeypatch.setattr(claude, "exec_or_spawn", lambda argv: calls.append(argv)) - claude.launch({"workspace": WS}, tool_args, options=LaunchOptions(smart_routing=True)) + claude.launch({"workspace": WS}, tool_args, options=LaunchOptions()) assert calls == [["claude", "--settings", str(claude.CLAUDE_SETTINGS_PATH), *tool_args]] v2.launch_claude.assert_not_called() @@ -1198,7 +1140,11 @@ def test_v2_positional_prompt_uses_first_prompt_routing(self, monkeypatch, tool_ monkeypatch.setattr(claude, "_original_launch_model", lambda _state: None) monkeypatch.setattr(v2, "launch_claude", launch_v2) - claude.launch({"workspace": WS}, tool_args, options=LaunchOptions(smart_routing=True)) + claude.launch( + {"workspace": WS}, + tool_args, + options=LaunchOptions(launch_smart_routing=True), + ) launch_v2.assert_called_once_with( {"workspace": WS}, @@ -1211,25 +1157,6 @@ def test_v2_positional_prompt_uses_first_prompt_routing(self, monkeypatch, tool_ model_name=claude._maybe_add_1m_suffix, ) - def test_v2_explicit_prompt_overrides_subcommand_shaped_text(self, monkeypatch): - launch_v2 = Mock() - monkeypatch.setattr(claude, "_original_launch_model", lambda _state: None) - monkeypatch.setattr(v2, "launch_claude", launch_v2) - - claude.launch( - {"workspace": WS}, - ["doctor"], - options=LaunchOptions(smart_routing=True, explicit_prompt=True), - ) - - launch_v2.assert_called_once() - - def test_v2_does_not_treat_option_value_as_positional_argument(self): - assert claude._uses_interactive_tui(["--name", "doctor"]) is True - - def test_v2_treats_optional_option_value_as_interactive(self): - assert claude._uses_interactive_tui(["--resume", "session-id"]) is True - def test_gateway_discovery_uses_direct_gateway(self, monkeypatch): calls: list[list[str]] = [] monkeypatch.delenv(v2.ENV_VAR, raising=False) diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index 91f35f42..2e429202 100644 --- a/tests/test_agent_codex.py +++ b/tests/test_agent_codex.py @@ -25,35 +25,6 @@ def test_display(self): assert codex.SPEC["display"] == "Codex" -class TestShouldUseSmartRouting: - @pytest.mark.parametrize( - ("tool_args", "explicit_prompt"), - [([], False), ([], True), (["fix this"], True)], - ) - def test_accepts_only_explicit_tui_shapes(self, tool_args, explicit_prompt): - assert codex.should_use_smart_routing(tool_args, explicit_prompt=explicit_prompt) is True - - @pytest.mark.parametrize( - "tool_args", - [ - ["exec", "fix this"], - ["review"], - ["app-server"], - ["update"], - ["--model", "gpt-5.6-sol"], - ["--model", "gpt-5.6-sol", "--", "fix this"], - ["fix this"], - ["one", "two"], - ], - ) - def test_rejects_all_other_invocation_shapes(self, tool_args): - assert codex.should_use_smart_routing(tool_args) is False - - @pytest.mark.parametrize("tool_args", [["--model=x"], ["--model"], ["-m"]]) - def test_explicit_model_bypasses_separator_opt_in(self, tool_args): - assert codex.should_use_smart_routing(tool_args, explicit_prompt=True) is False - - class TestHasUcodeConfig: def test_detects_profile_config(self, tmp_path, monkeypatch): config_path = tmp_path / "ucode.config.toml" @@ -572,53 +543,10 @@ def test_skips_git_repo_check(self): class TestCodexLaunch: - """launch() runs codex with --profile first and relaunches without it only - when that attempt fails *fast* — codex's --profile rejection is a parse-time - error, so a fast nonzero exit means the subcommand didn't accept --profile. - A slow failure is a real session error and is propagated unchanged.""" + """Normal launches layer the ucode profile as universal config overrides.""" @staticmethod - def _patch(monkeypatch, *, returncode: int, elapsed: float): - """Stub subprocess.run to return `returncode` and make launch() perceive - `elapsed` seconds between its two time.monotonic() reads.""" - runs: list[list[str]] = [] - fallbacks: list[list[str]] = [] - - def fake_run(argv, **kwargs): - runs.append(argv) - return codex.subprocess.CompletedProcess(argv, returncode) - - # launch() reads time.monotonic() once before run and once after. - clock = iter([100.0, 100.0 + elapsed]) - monkeypatch.setattr(codex.subprocess, "run", fake_run) - monkeypatch.setattr(codex.time, "monotonic", lambda: next(clock)) - monkeypatch.setattr(codex, "exec_or_spawn", lambda argv: fallbacks.append(argv)) - monkeypatch.setattr(codex, "get_databricks_token", lambda workspace, profile=None: "tok") - monkeypatch.setattr(codex, "clear_model_preferences", lambda state: False) - return runs, fallbacks - - def test_sets_oauth_token_and_runs_with_profile(self, monkeypatch): - monkeypatch.delenv("OAUTH_TOKEN", raising=False) - runs, fallbacks = self._patch(monkeypatch, returncode=0, elapsed=0.5) - monkeypatch.setattr( - codex, "get_databricks_token", lambda workspace, profile=None: "fresh-token" - ) - with pytest.raises(SystemExit) as exc: - codex.launch({"workspace": WS}, ["--search"], options=LaunchOptions()) - assert exc.value.code == 0 - assert os.environ["OAUTH_TOKEN"] == "fresh-token" - assert runs == [["codex", "--profile", "ucode", "--search"]] - assert fallbacks == [] - - def test_success_propagates_exit_without_retry(self, monkeypatch): - runs, fallbacks = self._patch(monkeypatch, returncode=0, elapsed=0.2) - with pytest.raises(SystemExit) as exc: - codex.launch({"workspace": WS}, ["exec", "hi"], options=LaunchOptions()) - assert exc.value.code == 0 - assert runs == [["codex", "--profile", "ucode", "exec", "hi"]] - assert fallbacks == [] - - def test_app_layers_ucode_profile_as_config_overrides(self, tmp_path, monkeypatch): + def _patch(tmp_path, monkeypatch): profile_path = tmp_path / "ucode.config.toml" profile_path.write_text( 'model_provider = "ucode-databricks"\n\n' @@ -628,85 +556,59 @@ def test_app_layers_ucode_profile_as_config_overrides(self, tmp_path, monkeypatc 'wire_api = "responses"\n', encoding="utf-8", ) + launches: list[list[str]] = [] monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", profile_path) - runs, launches = self._patch(monkeypatch, returncode=0, elapsed=0.2) + monkeypatch.setattr(codex, "exec_or_spawn", lambda argv: launches.append(argv)) + monkeypatch.setattr(codex, "get_databricks_token", lambda workspace, profile=None: "tok") + monkeypatch.setattr(codex, "clear_model_preferences", lambda state: False) + return launches + + def test_sets_oauth_token(self, tmp_path, monkeypatch): + monkeypatch.delenv("OAUTH_TOKEN", raising=False) + launches = self._patch(tmp_path, monkeypatch) + monkeypatch.setattr( + codex, "get_databricks_token", lambda workspace, profile=None: "fresh-token" + ) + codex.launch({"workspace": WS}, ["--search"], options=LaunchOptions()) + + assert os.environ["OAUTH_TOKEN"] == "fresh-token" + assert launches[0][-1] == "--search" + + @pytest.mark.parametrize( + "tool_args", + [ + ["exec", "hi"], + ["update"], + ["app-server", "--listen", "stdio://"], + ["app", "--new-window"], + ], + ) + def test_layers_profile_as_config_overrides(self, tmp_path, monkeypatch, tool_args): + launches = self._patch(tmp_path, monkeypatch) - codex.launch({"workspace": WS}, ["app", "--new-window"], options=LaunchOptions()) + codex.launch({"workspace": WS}, tool_args, options=LaunchOptions()) - assert runs == [] - assert launches[0][:2] == ["codex", "app"] + assert launches[0][0] == "codex" assert "--profile" not in launches[0] - assert launches[0][-1] == "--new-window" + assert launches[0][-len(tool_args) :] == tool_args assert 'model_provider="ucode-databricks"' in launches[0] provider_arg = next( arg for arg in launches[0] if arg.startswith("model_providers.ucode-databricks=") ) assert 'base_url = "https://example.databricks.com/ai-gateway/codex/v1"' in provider_arg - def test_app_requires_populated_ucode_profile(self, tmp_path, monkeypatch): + def test_requires_populated_ucode_profile(self, tmp_path, monkeypatch): monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", tmp_path / "missing.config.toml") - runs, launches = self._patch(monkeypatch, returncode=0, elapsed=0.2) + launches = [] + monkeypatch.setattr(codex, "exec_or_spawn", lambda argv: launches.append(argv)) + monkeypatch.setattr(codex, "get_databricks_token", lambda *_args: "tok") + monkeypatch.setattr(codex, "clear_model_preferences", lambda state: False) with pytest.raises(RuntimeError, match="ucode configure --agents codex"): - codex.launch({"workspace": WS}, ["app"], options=LaunchOptions()) + codex.launch({"workspace": WS}, ["update"], options=LaunchOptions()) - assert runs == [] assert launches == [] - def test_fast_failure_relaunches_without_profile(self, monkeypatch): - # codex rejects --profile on server-family subcommands at parse time — - # a fast nonzero exit → relaunch without --profile. - for args in (["app-server", "--listen", "u"], ["mcp-server"]): - runs, fallbacks = self._patch(monkeypatch, returncode=1, elapsed=0.15) - codex.launch({"workspace": WS}, args, options=LaunchOptions()) - assert runs == [["codex", "--profile", "ucode", *args]] - assert fallbacks == [["codex", *args]] - - def test_fallback_warns_on_stderr_before_handoff(self, monkeypatch, capsys): - # The fallback drops ucode's Databricks routing, so it must say so. The - # warning goes to *stderr*: `codex app-server` speaks JSON-RPC on stdout, - # and a warning there would corrupt the stream its caller parses. - warned_before_handoff = [] - runs, fallbacks = self._patch(monkeypatch, returncode=1, elapsed=0.15) - monkeypatch.setattr( - codex, - "exec_or_spawn", - lambda argv: warned_before_handoff.append(capsys.readouterr()), - ) - codex.launch({"workspace": WS}, ["app-server"], options=LaunchOptions()) - - # execvp replaces the process, so the warning must already be out by then. - assert len(warned_before_handoff) == 1 - captured = warned_before_handoff[0] - err = " ".join(captured.err.split()) # unwrap Rich's width-based wrapping - # Attributes the flag to ucode (users never type --profile themselves) - # and points at codex's own error so it doesn't read as their mistake. - assert "ucode's `--profile`" in err - assert "error above" in err - # Names both config scopes Codex will resolve without the ucode profile. - assert str(codex.LEGACY_CODEX_CONFIG_PATH) in err - assert "OS-managed settings" in err - assert "instead of the ucode profile" in err - assert captured.out == "" - - def test_slow_failure_does_not_retry(self, monkeypatch): - # A session that started and then failed (seconds) must NOT be re-run - # without --profile — that would silently drop ucode's Databricks routing - # (relaunching the user's prompt on their own OpenAI login). - runs, fallbacks = self._patch(monkeypatch, returncode=1, elapsed=8.0) - with pytest.raises(SystemExit) as exc: - codex.launch({"workspace": WS}, ["exec", "hi"], options=LaunchOptions()) - assert exc.value.code == 1 - assert fallbacks == [] - - def test_fast_success_does_not_retry(self, monkeypatch): - # The retry is gated on a *nonzero* exit; a fast clean exit just returns. - runs, fallbacks = self._patch(monkeypatch, returncode=0, elapsed=0.15) - with pytest.raises(SystemExit) as exc: - codex.launch({"workspace": WS}, [], options=LaunchOptions()) - assert exc.value.code == 0 - assert fallbacks == [] - class TestCodexManagedConfig: """Every normal configuration also reconciles Codex's OS-managed config.""" diff --git a/tests/test_agents_init.py b/tests/test_agents_init.py index 0935ac94..8bef1b7c 100644 --- a/tests/test_agents_init.py +++ b/tests/test_agents_init.py @@ -99,7 +99,7 @@ def test_tool_update_available_uses_agent_override(self, monkeypatch): def test_launch_dispatches_invocation_options(monkeypatch): calls = [] - options = LaunchOptions(smart_routing=True) + options = LaunchOptions(launch_smart_routing=True) monkeypatch.setattr( agents_mod.codex, "launch", diff --git a/tests/test_cli.py b/tests/test_cli.py index e9d68e53..96829955 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -487,6 +487,33 @@ def test_agent_records_prompt_separator(self, args, forwarded, has_separator): assert ctx.args == forwarded assert cli_mod._has_explicit_prompt(ctx) is has_separator + @pytest.mark.parametrize("tool", ["codex", "claude"]) + @pytest.mark.parametrize( + ("tool_args", "explicit_prompt", "model", "provider", "expected"), + [ + ([], False, None, None, True), + (["fix this"], True, None, None, True), + (["fix this"], False, None, None, False), + (["update"], False, None, None, False), + (["--model", "fixed"], False, None, None, False), + ([], False, "fixed", None, False), + ([], False, None, "catalog.schema.service", False), + ], + ) + def test_codex_and_claude_share_smart_routing_policy( + self, tool, tool_args, explicit_prompt, model, provider, expected + ): + options = cli_mod._launch_options( + tool, + tool_args, + smart_routing_enabled=True, + explicit_prompt=explicit_prompt, + model=model, + provider=provider, + ) + + assert options.launch_smart_routing is expected + def test_codex_refresh_is_consumed_by_ucode(self): with patch("ucode.cli._launch_tool") as mock_launch: result = runner.invoke(app, ["codex", "--refresh"]) diff --git a/tests/test_codex_config.py b/tests/test_codex_config.py index 445b95fb..061169e7 100644 --- a/tests/test_codex_config.py +++ b/tests/test_codex_config.py @@ -1,5 +1,7 @@ from __future__ import annotations +import tomlkit + from ucode.agents import codex from ucode.codex_config import codex_config_args @@ -29,3 +31,32 @@ def test_layers_provider_overrides_without_replacing_user_config(self, monkeypat assert "/ai-gateway/codex/v1" in provider_override assert 'command = "' in provider_override assert '"myprof"' in provider_override + + def test_renders_nested_tables_from_parsed_profile(self): + profile = tomlkit.parse( + """ +model_provider = "ucode-databricks" + +[model_providers.ucode-databricks] +name = "Databricks AI Gateway" + +[model_providers.ucode-databricks.http_headers] +User-Agent = "ucode" + +[model_providers.ucode-databricks.auth] +command = "ucode" +args = ["codex-token"] + +[tui.model_availability_nux] +"gpt-5.6-sol" = 1 +""" + ) + + args = codex_config_args(profile) + + provider_override = next( + arg for arg in args if arg.startswith("model_providers.ucode-databricks=") + ) + assert 'http_headers = {User-Agent = "ucode"}' in provider_override + assert 'auth = {command = "ucode", args = ["codex-token"]}' in provider_override + assert 'tui={model_availability_nux = {"gpt-5.6-sol" = 1}}' in args diff --git a/tests/test_codex_smart_routing_v2.py b/tests/test_codex_smart_routing_v2.py index 15c8e033..818fa553 100644 --- a/tests/test_codex_smart_routing_v2.py +++ b/tests/test_codex_smart_routing_v2.py @@ -31,16 +31,13 @@ def test_rejects_unsupported_codex_version(self, monkeypatch): monkeypatch.setattr(v2, "launch_codex", lambda *args, **kwargs: pytest.fail("launched")) with pytest.raises(RuntimeError, match="requires Codex 0.145.0 or newer; found 0.144.0"): - codex.launch({"workspace": WS}, [], options=LaunchOptions(smart_routing=True)) + codex.launch({"workspace": WS}, [], options=LaunchOptions(launch_smart_routing=True)) @pytest.mark.parametrize( ("tool_args", "options"), [ - ([], LaunchOptions(smart_routing=True)), - ( - ["fix the parser"], - LaunchOptions(smart_routing=True, explicit_prompt=True), - ), + ([], LaunchOptions(launch_smart_routing=True)), + (["fix the parser"], LaunchOptions(launch_smart_routing=True)), ], ) def test_codex_smart_routing_launch_dispatches_to_v2(self, monkeypatch, tool_args, options): @@ -84,28 +81,25 @@ def launch_v2(state, tool_args, **kwargs): ["fix this"], ], ) - def test_codex_launch_bypasses_routing_for_other_shapes(self, monkeypatch, tool_args): - runs = [] + def test_codex_launch_bypasses_routing_for_other_shapes(self, tmp_path, monkeypatch, tool_args): + 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.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") monkeypatch.setattr(codex, "get_databricks_token", lambda *_args: "token") monkeypatch.setattr(v2, "launch_codex", lambda *args, **kwargs: pytest.fail("launched")) - monkeypatch.setattr( - codex.subprocess, - "run", - lambda argv: runs.append(argv) or codex.subprocess.CompletedProcess(argv, 0), - ) + monkeypatch.setattr(codex, "exec_or_spawn", lambda argv: launches.append(argv)) - with pytest.raises(SystemExit) as exc: - codex.launch( - {"workspace": WS}, - tool_args, - options=LaunchOptions(smart_routing=True), - ) + codex.launch( + {"workspace": WS}, + tool_args, + options=LaunchOptions(), + ) - assert exc.value.code == 0 - assert runs == [["codex", "--profile", "ucode", *tool_args]] + assert launches == [["codex", "--config", 'model_provider="ucode-databricks"', *tool_args]] def test_codex_launch_normalizes_cached_bootstrap_model(self, monkeypatch): calls = [] @@ -123,7 +117,7 @@ def launch_v2(state, tool_args, **kwargs): codex.launch( {"workspace": WS, "codex_models": ["system.ai.gpt-5-6-luna"]}, [], - options=LaunchOptions(smart_routing=True), + options=LaunchOptions(launch_smart_routing=True), ) assert calls[0]["start_model"] == "gpt-5.6-luna" From 678c43a83153d4e84d0b5f74ab959e104dec5333 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Sat, 5 Sep 2026 00:28:08 +0000 Subject: [PATCH 3/5] small tweak --- tests/test_cli.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_cli.py b/tests/test_cli.py index 96829955..a05754bc 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -634,6 +634,7 @@ 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): monkeypatch.delenv("ENABLE_SMART_ROUTING_V2", raising=False) From 248df952f35692dbb521d2c55f700368fde26f14 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Sat, 5 Sep 2026 01:02:34 +0000 Subject: [PATCH 4/5] update --- src/ucode/agents/__init__.py | 3 +-- src/ucode/agents/args.py | 3 +++ src/ucode/agents/claude.py | 15 ++------------- src/ucode/cli.py | 12 +++++------- tests/test_agent_claude.py | 36 +++++++++++++++++++++++++----------- tests/test_cli.py | 12 +++++++----- 6 files changed, 43 insertions(+), 38 deletions(-) diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index 0c82e4fa..cafdd347 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -416,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, diff --git a/src/ucode/agents/args.py b/src/ucode/agents/args.py index 064eceda..7cfb9e1c 100644 --- a/src/ucode/agents/args.py +++ b/src/ucode/agents/args.py @@ -10,6 +10,9 @@ 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: diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 08263f65..7e37ccbf 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -404,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 ` 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. @@ -1363,6 +1350,8 @@ def launch( 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)) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index b4c4062c..ad6a8869 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -1998,6 +1998,7 @@ def _can_launch_from_cached_config( if smart_routing_enabled is None: smart_routing_enabled = smart_routing_v2.enabled() + # TODO(lilly): replace with codex/v1/models or custom catalog if tool == "codex" and smart_routing_enabled: if not state.get("codex_models") or not state.get("oss_models"): return False @@ -2034,6 +2035,7 @@ def _launch_options( ) -> LaunchOptions: has_model_override = model is not None or explicit_model_arg_value(tool_args) is not None return LaunchOptions( + claude_launch_model=model if tool == "claude" and provider is None else None, launch_smart_routing=( # Smart routing is enabled globally. smart_routing_enabled @@ -2253,9 +2255,7 @@ def _launch_tool( resolved_model = managed_model # An explicit `--model` is the user's own choice and outranks everything above (managed # default, smart-routing pick). Non-claude agents take it as the resolved model, which - # their CLIs pass to the gateway verbatim. Claude is special (see custom_model below): - # Claude Code validates ANTHROPIC_MODEL client-side and rejects a raw Databricks id, so - # the id can't ride `resolved_model` — it is threaded separately as `custom_model`. + # Codex keeps an explicit --model in ctx.args and passes it to its CLI verbatim. if model and tool != "claude": resolved_model = model state = configure_tool( @@ -2266,10 +2266,8 @@ def _launch_tool( provider_models=provider_models, relayed=relayed, route_root_model=route_root_model, - # Under a provider, --model is honored via route_root_model (above), not custom_model — - # the latter pins a raw id into every family alias, which would clobber the service's - # per-family target pins. - custom_model=model if (tool == "claude" and not provider) else None, + # Claude's explicit model is launch-scoped and is passed through LaunchOptions below. + custom_model=None, coding_agent_config_defaults=coding_agent_config_defaults, ) # Relayed = a Claude subscription: forward --model to Claude Code's own flag, like `-- --model X`. diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 11a9c5b4..8a73d950 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -142,10 +142,8 @@ def test_no_1m_suffix_for_model_services_haiku(self): ) assert overlay["env"]["ANTHROPIC_DEFAULT_HAIKU_MODEL"] == "system.ai.claude-haiku-4-6" - def test_custom_model_pins_all_family_aliases(self): - # `ucode claude --model` pins the id into every family alias so it takes effect whichever - # slot Claude Code resolves — and NOT into ANTHROPIC_MODEL, which Claude Code validates and - # rejects for a raw Databricks id. It overrides the discovered-model aliases. + def test_custom_model_does_not_persist_model_selection(self): + # Explicit model selection is launch-scoped and must not be written to settings. overlay, _ = claude.render_overlay( WS, "s4", @@ -153,22 +151,23 @@ def test_custom_model_pins_all_family_aliases(self): custom_model="main.aarushi.claude-opus-5", ) env = overlay["env"] - assert env["ANTHROPIC_DEFAULT_OPUS_MODEL"] == "main.aarushi.claude-opus-5" - assert env["ANTHROPIC_DEFAULT_SONNET_MODEL"] == "main.aarushi.claude-opus-5" - assert env["ANTHROPIC_DEFAULT_HAIKU_MODEL"] == "main.aarushi.claude-opus-5" assert "ANTHROPIC_MODEL" not in env + assert env["ANTHROPIC_DEFAULT_OPUS_MODEL"] == "system.ai.claude-opus-4-8[1m]" + assert env["ANTHROPIC_DEFAULT_SONNET_MODEL"] == "system.ai.sonnet" + assert "ANTHROPIC_DEFAULT_HAIKU_MODEL" not in env # No [1m] suffix is appended to the custom id — it's passed through verbatim. - assert "[1m]" not in env["ANTHROPIC_DEFAULT_OPUS_MODEL"] + assert "main.aarushi.claude-opus-5" not in env.values() - def test_custom_model_pins_fable_alias_only_when_fable_enabled(self): + def test_custom_model_does_not_persist_fable_selection(self): without = claude.render_overlay(WS, "s4", claude_models={}, custom_model="main.x.m")[0][ "env" ] - assert "ANTHROPIC_DEFAULT_FABLE_MODEL" not in without + assert "ANTHROPIC_MODEL" not in without with_fable = claude.render_overlay( WS, "s4", claude_models={}, custom_model="main.x.m", fable_enabled=True )[0]["env"] - assert with_fable["ANTHROPIC_DEFAULT_FABLE_MODEL"] == "main.x.m" + assert "ANTHROPIC_MODEL" not in with_fable + assert "ANTHROPIC_DEFAULT_FABLE_MODEL" not in with_fable def test_sets_anthropic_base_url(self): overlay, _ = claude.render_overlay(WS, "s4") @@ -1114,6 +1113,21 @@ def test_default_launch_keeps_existing_auth_path(self, monkeypatch): assert os.environ["OAUTH_TOKEN"] == "token" assert calls == [["claude", "--settings", str(claude.CLAUDE_SETTINGS_PATH), "--debug"]] + def test_launch_model_is_only_set_for_current_process(self, monkeypatch): + calls: list[list[str]] = [] + monkeypatch.delenv("ANTHROPIC_MODEL", raising=False) + monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "token") + monkeypatch.setattr(claude, "exec_or_spawn", lambda argv: calls.append(argv)) + + claude.launch( + {"workspace": WS, "profile": "test"}, + [], + options=LaunchOptions(claude_launch_model="cat.schema.model"), + ) + + assert os.environ["ANTHROPIC_MODEL"] == "cat.schema.model" + assert calls == [["claude", "--settings", str(claude.CLAUDE_SETTINGS_PATH)]] + @pytest.mark.parametrize( "tool_args", [ diff --git a/tests/test_cli.py b/tests/test_cli.py index a05754bc..9189d329 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -808,7 +808,7 @@ def test_forwarded_model_is_reported_in_launch_summary(self, forwarded_args): assert "Model: system.ai.claude-opus-4-8" not in output assert mock_launch.call_args.args[2] == forwarded_args - def test_model_threads_to_claude_as_custom_model(self, monkeypatch): + def test_model_is_launch_scoped_for_claude(self, monkeypatch): monkeypatch.delenv("ENABLE_SMART_ROUTING_V2", raising=False) with ( patch("ucode.cli.ensure_bootstrap_dependencies"), @@ -822,11 +822,13 @@ def test_model_threads_to_claude_as_custom_model(self, monkeypatch): ): result = runner.invoke(app, ["claude", "--model", "cat.schema.claude-opus-5"]) assert result.exit_code == 0, result.output - # Claude routes --model as custom_model (pinned into the family aliases by render_overlay), - # NOT as ANTHROPIC_MODEL — Claude Code validates that value and rejects a raw id. - assert mock_configure.call_args.kwargs["custom_model"] == "cat.schema.claude-opus-5" + # The model is passed through invocation-scoped LaunchOptions, not persisted in settings. + assert mock_configure.call_args.kwargs["custom_model"] is None 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"].claude_launch_model + == "cat.schema.claude-opus-5" + ) def test_v2_model_sets_transient_launch_override(self, monkeypatch): monkeypatch.setenv("ENABLE_SMART_ROUTING_V2", "1") From 2dd0c6747634e05201422ccc93569aa3f2066a63 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Sat, 5 Sep 2026 01:05:33 +0000 Subject: [PATCH 5/5] ruff --- src/ucode/cli.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index ad6a8869..00152e4a 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -2047,7 +2047,7 @@ def _launch_options( and provider is None # Route a bare agent launch or a prompt explicitly passed after `--`. and (not tool_args or explicit_prompt) - ) + ), )