diff --git a/src/ucode/smart_routing/v2.py b/src/ucode/smart_routing/v2.py index 287b37a8..a79c5187 100644 --- a/src/ucode/smart_routing/v2.py +++ b/src/ucode/smart_routing/v2.py @@ -29,6 +29,7 @@ list_anthropic_model_catalog, list_anthropic_models, ) +from ucode.launcher import exec_or_spawn from ucode.smart_routing import claude_routing, codex_interposer, routing from ucode.smart_routing.claude_hooks import ( FIRST_PROMPT_SOCKET_ENV, @@ -39,8 +40,11 @@ from ucode.ui import print_warning ENABLE_SMART_ROUTING_ENV_VAR = "ENABLE_SMART_ROUTING_V2" +ENABLE_SUBAGENT_ROUTING_ENV_VAR = "ENABLE_SMART_ROUTING_SUBAGENT_ONLY" LEGACY_STATE_KEY = "smart_routing_enabled" +_SMART_ROUTING_ENV_VARS = (ENABLE_SMART_ROUTING_ENV_VAR, ENABLE_SUBAGENT_ROUTING_ENV_VAR) + CODEX_INTERPOSER_LOG = APP_DIR / "codex-v2-interposer.log" CLAUDE_TARGET_MODEL = "system.ai.claude-sonnet-4-6[1m]" # TODO(lilly): replace with smart router. @@ -115,32 +119,46 @@ def _model_picker_catalog() -> AnthropicModelCatalog | None: 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" + return any(source.get(var) == "1" for var in _SMART_ROUTING_ENV_VARS) + + +def first_prompt_routing_enabled(env: MutableMapping[str, str] | None = None) -> bool: + """Whether the first prompt is routed. Subagent-only wins over the full V2 flag.""" + source = os.environ if env is None else env + return ( + source.get(ENABLE_SMART_ROUTING_ENV_VAR) == "1" + and source.get(ENABLE_SUBAGENT_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.""" +def enable_smart_routing( + env: MutableMapping[str, str] | None = None, +) -> dict[str, str | None]: + """Set the full smart-routing env var and return the prior value of every routing var.""" target = os.environ if env is None else env - previous = target.get(ENABLE_SMART_ROUTING_ENV_VAR) + previous = {var: target.get(var) for var in _SMART_ROUTING_ENV_VARS} target[ENABLE_SMART_ROUTING_ENV_VAR] = "1" return previous def restore_smart_routing_env( - previous: str | None, env: MutableMapping[str, str] | None = None + previous: dict[str, 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 + for var, value in previous.items(): + if value is None: + target.pop(var, None) + else: + target[var] = value -def disable_smart_routing(env: MutableMapping[str, str] | None = None) -> str | None: - """Temporarily remove the smart-routing env var and return its prior value.""" +def disable_smart_routing( + env: MutableMapping[str, str] | None = None, +) -> dict[str, str | None]: + """Temporarily remove the smart-routing env vars and return their prior values.""" target = os.environ if env is None else env - return target.pop(ENABLE_SMART_ROUTING_ENV_VAR, None) + return {var: target.pop(var, None) for var in _SMART_ROUTING_ENV_VARS} def _loopback_websocket_url(port: int) -> str: @@ -445,6 +463,7 @@ def launch_claude( ) model_ids = catalog.model_ids + route_first_prompt = first_prompt_routing_enabled() run_id = f"{os.getpid()}-{uuid.uuid4().hex[:8]}" socket_path = APP_DIR / f"claude-v2-{run_id}.sock" settings_path = APP_DIR / f"claude-v2-{run_id}.json" @@ -457,8 +476,11 @@ 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) + if route_first_prompt: + env[ENABLE_SMART_ROUTING_ENV_VAR] = "1" + env[FIRST_PROMPT_SOCKET_ENV] = str(socket_path) + else: + env[ENABLE_SUBAGENT_ROUTING_ENV_VAR] = "1" model_overrides = settings.setdefault("modelOverrides", {}) if not isinstance(model_overrides, dict): raise RuntimeError("Claude settings 'modelOverrides' must be an object for smart routing.") @@ -468,12 +490,26 @@ def launch_claude( "claude_models": {str(index): model for index, model in enumerate(model_ids)}, } sync_smart_routing_hooks(settings, routing_state, enabled=True) - sync_first_prompt_hook(settings, hook_executable) + if route_first_prompt: + sync_first_prompt_hook(settings, hook_executable) write_json_file(settings_path, settings) model_args = launch_model_args(remaining, launch_model) routed_agent_args = _with_routed_claude_agents(remaining, model_ids) argv = [binary, "--settings", str(settings_path), *model_args, *routed_agent_args] + if not route_first_prompt: + # Subagent-only routing needs no PTY: the PreToolUse hooks ride in the + # per-launch settings, so spawn Claude directly and clean up after it. + proc = subprocess.Popen(argv) + try: + returncode = proc.wait() + except KeyboardInterrupt: + proc.send_signal(signal.SIGINT) + returncode = proc.wait() + finally: + settings_path.unlink(missing_ok=True) + sys.exit(returncode) + model_setting = _ClaudeModelSettingGuard(user_settings_path) def route_prompt(prompt: str) -> claude_pty.FirstPromptRoute: @@ -566,6 +602,10 @@ def launch_codex( "PreToolUse": _v2_pre_tool_use_hooks(state, available_models), } config_args = codex_config_args(overlay) + if not first_prompt_routing_enabled(): + # Subagent-only routing needs neither the app-server nor the interposer: + # the hooks ride in the CLI config, so launch the TUI directly. + exec_or_spawn([binary, *config_args, *tool_args]) app_port = _free_port() app_server_url = _loopback_websocket_url(app_port) diff --git a/tests/README.md b/tests/README.md index 0490b7e0..df2b94ba 100644 --- a/tests/README.md +++ b/tests/README.md @@ -48,6 +48,8 @@ All tests live directly in `integration/`; shared mechanics live in `utils/`. | `test_ug_codex_app_help`, `test_ug_codex_app_server_help`, `test_ug_codex_exec_help`, `test_ug_codex_mcp_help` | Request subcommand help, routing off/on | Real agent help; no routing wrapper | | `test_ug_codex_app_reports_unknown_argument` | Pass an invalid option directly to `ug codex app`, routing off/on | Real Codex parser error and status preserved | | `test_ug_codex_app_server_client_initializes` | Connect a stdio client, direct/`--` separator, routing off/on | Actual JSON-RPC initialize response; no non-JSON stdout; no routing | +| `test_smart_routing_claude_route_subagent_hook`, `test_smart_routing_codex_route_subagent_hook` | Pipe a real PreToolUse spawn payload to the installed route-subagent hook with subagent-only routing enabled | Allow decision against the live router; requested model replaced by a routed agent definition (Claude) or bundled catalog slug (Codex) from the offered models; one audited decision matching the session and task | +| `test_smart_routing_claude_subagent_only_launch_shows_no_first_prompt_banner`, `test_smart_routing_codex_subagent_only_launch_shows_no_first_prompt_banner` | Configure, then launch the real TUI with both the full and subagent-only routing flags set and submit one file prompt | Subagent-only takes precedence: the prompt completes with no smart-routing banner and no first-prompt routing wrapper (PTY/interposer); Claude's SessionStart canary proves the routing hooks armed; normal exit | | `test_ug_configure_claude_repeat_and_revert`, `test_ug_configure_codex_repeat_and_revert` | Configure twice over user settings; complete a task; revert twice | Settings preserved; no bearer in ug state; generated config removed; status unconfigured | | `test_ug_configure_claude_rejects_invalid_credentials`, `test_ug_configure_codex_rejects_invalid_credentials` | Configure with a rejected bearer against the real workspace | Authentication failure; no successful saved setup | | `test_ug_configure_managed_claude`, `test_ug_configure_managed_codex` | Configure against a workspace that publishes a managed CodingAgentConfig | No agent selector; each agent's generated config exposes exactly the admin's static model_services; real gateway prompt on launch | @@ -65,11 +67,11 @@ All tests live directly in `integration/`; shared mechanics live in `utils/`. | `test_ug_and_ucode_auth_helpers_emit_only_the_supplied_bearer` | Run both auth helper commands with the public bearer override, with and without forced refresh | Exact token-only stdout, no warnings or ANSI escapes; no workspace authentication or saved state | | `test_ug_and_ucode_web_search_helpers_preserve_mcp_stdio` | Initialize and list tools through both web-search helper commands | Exactly the MCP JSON-RPC responses; no text/ANSI contamination; existing server/tool identities preserved; no model request | -With both agents selected there are **41 live cases** (6 interactive TUI cases), +With both agents selected there are **46 live cases** (6 interactive TUI cases), **4 managed-workspace cases** (marker `managed`, run against a separate workspace that -publishes a CodingAgentConfig), **34 managed-fixture cases** (marker `managed_fixture`, with only +publishes a CodingAgentConfig), **36 managed-fixture cases** (marker `managed_fixture`, with only the CodingAgentConfig input injected), and **5 installation checks**. The 24 numbered scenarios -cover configured and fresh state across the Claude and Codex managed-discovery matrix; ten +cover configured and fresh state across the Claude and Codex managed-discovery matrix; twelve existing collected cases cover focused model, MCP, skills, and lifecycle shapes. Parametrization varies argument spelling or routing mode, never hides the agent/provider in the test name. Duplicate boot-only cases @@ -129,7 +131,7 @@ pending. The descriptive jobs provide the actual coverage and diagnostics. | Provider switching, relayed/subscription MPS | Not covered by the four provider journeys | | TUI initial prompt supplied on the launch command line | Not yet covered; headless prompt arguments are covered | | Follow-up turns and conversation resume | Not covered; reopen proves startup, not conversation resume | -| Claude/Codex interactive smart routing | First-prompt routing covered by the `managed_fixture` smart-routing banner journeys; subagent routing, interactive explicit-model bypass, and dedicated routing CI shards remain deferred. Unit/component routing tests do not establish live routing behavior. | +| Claude/Codex interactive smart routing | First-prompt routing covered by the `managed_fixture` smart-routing banner journeys; subagent routing covered at the hook protocol level by the route-subagent hook journeys, which drive the real installed hook commands with a harness-shaped payload against the live router; the subagent-only launch journeys assert the first-prompt banner and routing wrappers stay silent while the routing hooks arm. The agent's interactive spawn decision, interactive explicit-model bypass, and dedicated routing CI shards remain deferred. Unit/component routing tests do not establish live routing behavior. | | Full allow/deny tool-permission matrix | Not covered; onboarding/trust uses actual TUI choices | | Desktop Codex app, Isaac itself, auto-upgrades | Not covered by command forwarding or pinned-version tests | | Native macOS/Windows managed settings, resize/signals | Separate platform coverage needed | diff --git a/tests/integration/README.md b/tests/integration/README.md index 224cfb34..9f5ce3d7 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -96,6 +96,7 @@ test_ug_codex_headless.py # script prompts and model arguments test_ug_claude_commands.py # command help forwarding test_ug_codex_commands.py # command help and parser error forwarding test_ug_codex_app_server.py # actual client/server initialize exchange +test_ug_smart_routing_hooks.py # route-subagent hook contract against the live router test_ug_configure_claude_lifecycle.py # repeat setup, revert, rejected credentials test_ug_configure_codex_lifecycle.py # repeat setup, revert, rejected credentials test_ug_claude_managed_model_discovery.py # fetched/reused Claude MPS policy cases @@ -134,9 +135,13 @@ A fixture file contains an unpredictable value absent from the prompt. Success requires an assistant answer in the real agent transcript containing that value, plus normal TUI exit. Codex evidence requires its task-complete event. Interactive first-prompt routing is covered by the managed_fixture smart-routing -banner journeys below for both agents. Subagent routing, interactive explicit-model -bypass, and dedicated smart-routing CI shards remain deferred; unit/component -routing tests do not establish that live behavior. +banner journeys below for both agents. Subagent routing is covered at the hook protocol +level by the route-subagent hook journeys, which drive the real installed hook commands +with a harness-shaped payload against the live router; the subagent-only launch journeys +assert the first-prompt banner and routing wrappers stay silent while the routing hooks +arm. The agent's interactive spawn decision, interactive explicit-model bypass, and +dedicated smart-routing CI shards remain deferred; unit/component routing tests do not +establish that live behavior. The relayed CUJ launches Claude through a relayed (subscription-relay) MPS and completes a file task on two models: a bare Anthropic id the subscription serves @@ -159,13 +164,13 @@ service allows a different model. Those choices are recorded in `versions.json`. No service is created or modified. A missing service, permission, or OAuth token fails the selected CUJ, rather than skipping it. -There are **42 live cases** (including 6 TUI journeys) and **5 installation +There are **46 live cases** (including 6 TUI journeys) and **5 installation checks** with both agents. A separate **4 managed-workspace cases** (one per agent, an idempotent re-configure, and a cache-TTL journey; marker `managed`) run against a workspace that publishes a -CodingAgentConfig; see "Managed-workspace journeys" below. A further **34 `managed_fixture` +CodingAgentConfig; see "Managed-workspace journeys" below. A further **36 `managed_fixture` cases** use `UCODE_MANAGED_CONFIG_STUB`. Twenty-four explicit configured/fresh Claude and Codex discovery and source-override journeys fetch the published config once per agent, replace that -agent's static source with its dedicated MPS, and reuse the result. Ten existing collected cases +agent's static source with its dedicated MPS, and reuse the result. Twelve existing collected cases cover focused model, MCP, skills, and lifecycle shapes, including per-agent model reconciliation and managed skill cleanup. See the named coverage and gaps matrix in @@ -246,13 +251,13 @@ record a discovered `system.ai` model as a test argument. Every same-repository PR and push to `main` runs **Smoke journeys**, followed by **Full journeys** even if smoke fails. Smoke runs the Hosted configure/TUI, headless argument, and custom OAuth CLI TUI journeys for each agent (six cases, -two agent jobs). Full runs all 42 live cases, including those smoke cases, in two +two agent jobs). Full runs all 46 live cases, including those smoke cases, in two disjoint agent lanes: | Agent lane | Marker | Cases | | --- | --- | --- | -| Claude | `live and claude` | 17 | -| Codex | `live and codex` | 25 | +| Claude | `live and claude` | 19 | +| Codex | `live and codex` | 27 | Each lane installs only its agent CLI, once, and runs all its configure, headless, commands, lifecycle, and applicable app-server journeys. Cases remain serial @@ -504,7 +509,7 @@ uv run --no-project --python 3.12 python scripts/run_integration.py \ unset DATABRICKS_BEARER ``` -This runs all 42 live cases. For the five installation checks, run the same +This runs all 46 live cases. For the five installation checks, run the same runner/version/index arguments with `--installation-only` and omit `-- -m live`; no bearer or workspace is needed. Results remain under `.integration-runs/`. Each invocation needs a new output directory; an existing one is rejected. diff --git a/tests/integration/test_ug_smart_routing_hooks.py b/tests/integration/test_ug_smart_routing_hooks.py new file mode 100644 index 00000000..255aafb3 --- /dev/null +++ b/tests/integration/test_ug_smart_routing_hooks.py @@ -0,0 +1,237 @@ +"""CUJs for the subagent-only smart-routing hook commands against the live workspace router. + +The agent harness invokes ``ug claude-router-hook route-subagent`` / +``ug codex-router-hook route-subagent`` on its PreToolUse event with a JSON payload on +stdin. These journeys drive the real installed hook commands through that stdin contract, +so the routing decision, response shape, and audit trail are asserted without relying on +an agent choosing to spawn a subagent. The interactive spawn decision itself remains +uncovered; see the gaps matrix in tests/README.md. +""" + +import json + +import pytest +from utils.evidence import FileTask +from utils.terminal import AgentTerminal + +# The same model lists as the managed_fixture smart-routing banner journeys, which are +# proven servable route options on the live e2e workspace. +SMART_ROUTING_BANNER = "Using Unity Gateway Smart Router." +SMART_ROUTING_SUBAGENT_NOTICE = "Using Unity Gateway Smart Router - Subagent" +CLAUDE_MODELS = [ + "system.ai.claude-opus-5", + "system.ai.claude-sonnet-5", + "system.ai.claude-haiku-4-5", + "system.ai.glm-5-3", + "system.ai.kimi-k3", +] +CODEX_MODELS = [ + "system.ai.gpt-6-astra", + "system.ai.gpt-5-6-sol", + "system.ai.gpt-5-6-terra", + "system.ai.gpt-5-6-luna", + "system.ai.gpt-5-5", + "system.ai.glm-5-3", + "system.ai.kimi-k3", +] +# Codex launches subagents on its bundled catalog slugs, not the workspace model id, so +# the routed model in the hook response must be one of these slugs. +CODEX_MODEL_SLUGS = { + "gpt-6-astra", + "gpt-5.6-sol", + "gpt-5.6-terra", + "gpt-5.6-luna", + "gpt-5.5", + "glm-5-3", + "kimi-k3", +} + + +@pytest.mark.live +@pytest.mark.claude +def test_smart_routing_claude_route_subagent_hook(live_session, workspace): + """Scenario: with subagent-only routing enabled, Claude Code fires PreToolUse for an + Agent spawn, piping the payload to ``ug claude-router-hook route-subagent``. + + Expected: the hook allows the call against the real workspace router, drops the + requested model in favor of a ``ucode-route-`` agent definition while preserving the + task text, and audits one decision naming an offered model for the session. Only the + hook contract is asserted; no agent decides to spawn. + """ + session = live_session + session.env["ENABLE_SMART_ROUTING_SUBAGENT_ONLY"] = "1" + payload = { + "session_id": "claude-route-subagent-hook", + "tool_name": "Agent", + "tool_input": { + "description": "Refactor the parser", + "prompt": "Refactor the parser module into a package and add unit tests.", + "subagent_type": "general-purpose", + "model": "sonnet", + }, + } + result = session.run( + "claude-router-hook", + "route-subagent", + "--host", + workspace, + *(arg for model in CLAUDE_MODELS for arg in ("--model", model)), + input_text=json.dumps(payload), + timeout=60, + ) + output = json.loads(result.stdout) + hook = output["hookSpecificOutput"] + assert hook["hookEventName"] == "PreToolUse", output + assert hook["permissionDecision"] == "allow", output + assert SMART_ROUTING_SUBAGENT_NOTICE in output["systemMessage"], output + updated = hook["updatedInput"] + assert "model" not in updated, updated + assert updated["subagent_type"].startswith("ucode-route-"), updated + assert updated["prompt"] == payload["tool_input"]["prompt"], updated + assert updated["description"] == payload["tool_input"]["description"], updated + + decisions_path = session.home / ".ucode" / "claude-smart-routing-decisions.jsonl" + assert decisions_path.is_file(), f"hook wrote no routing decision record: {decisions_path}" + rows = [json.loads(line) for line in decisions_path.read_text().splitlines() if line.strip()] + session.record("claude-smart-routing-decisions.jsonl", rows) + assert len(rows) == 1, rows + row = rows[0] + assert row["session_id"] == payload["session_id"], row + assert row["task_name"] == payload["tool_input"]["prompt"], row + assert row["requested_model"] in CLAUDE_MODELS, row + + +@pytest.mark.live +@pytest.mark.codex +def test_smart_routing_codex_route_subagent_hook(live_session, workspace): + """Scenario: with subagent-only routing enabled, Codex fires PreToolUse for a + spawn_agent call, piping the payload to ``ug codex-router-hook route-subagent``. + + Expected: the hook allows the call against the real workspace router, rewrites the + requested model to the bundled catalog slug of an offered model while preserving the + task message, and audits one decision matching the response for the session. Only the + hook contract is asserted; no agent decides to spawn. + """ + session = live_session + session.env["ENABLE_SMART_ROUTING_SUBAGENT_ONLY"] = "1" + payload = { + "session_id": "codex-route-subagent-hook", + "tool_name": "spawn_agent", + "tool_input": { + "task_name": "Refactor the parser", + "message": "Refactor the parser module into a package and add unit tests.", + "model": "gpt-5.5", + }, + } + result = session.run( + "codex-router-hook", + "route-subagent", + "--host", + workspace, + *(arg for model in CODEX_MODELS for arg in ("--model", model)), + input_text=json.dumps(payload), + timeout=60, + ) + output = json.loads(result.stdout) + hook = output["hookSpecificOutput"] + assert hook["hookEventName"] == "PreToolUse", output + assert hook["permissionDecision"] == "allow", output + assert SMART_ROUTING_SUBAGENT_NOTICE in output["systemMessage"], output + updated = hook["updatedInput"] + assert updated["model"] in CODEX_MODEL_SLUGS, updated + assert updated["message"] == payload["tool_input"]["message"], updated + assert updated["task_name"] == payload["tool_input"]["task_name"], updated + + decisions_path = session.home / ".ucode" / "codex-smart-routing-decisions.jsonl" + assert decisions_path.is_file(), f"hook wrote no routing decision record: {decisions_path}" + rows = [json.loads(line) for line in decisions_path.read_text().splitlines() if line.strip()] + session.record("codex-smart-routing-decisions.jsonl", rows) + assert len(rows) == 1, rows + row = rows[0] + assert row["session_id"] == payload["session_id"], row + assert row["task_name"] == payload["tool_input"]["message"], row + assert row["requested_model"] == updated["model"], row + + +@pytest.mark.live +@pytest.mark.claude +def test_smart_routing_claude_subagent_only_launch_shows_no_first_prompt_banner( + live_session, workspace +): + """Scenario: configure Claude, then launch the real TUI with both the full and the + subagent-only routing flags set and submit one file prompt. + + Expected: subagent-only takes precedence over the ambient full flag: the routing + hooks are armed (SessionStart canary), yet the prompt completes with no + smart-routing banner and no first-prompt routing wrapper anywhere in the session, + and the TUI exits normally. Only first-prompt silence is asserted here; subagent + routing engagement is covered by the route-subagent hook journey above. + """ + session = live_session + session.env["ENABLE_SMART_ROUTING_V2"] = "1" + session.env["ENABLE_SMART_ROUTING_SUBAGENT_ONLY"] = "1" + session.run( + "configure", + "--agents", + "claude", + "--workspace", + workspace, + "--skip-validate", + "--skip-upgrade", + "--disable-databricks-ai-tools", + ) + task = FileTask(session) + with AgentTerminal( + session, "claude", [str(session.binary), "claude"], "subagent-only-launch" + ) as tui: + tui.boot() + tui.submit(task.prompt) + tui.wait_for_task(task) + tui.exit_normally() + transcript = "".join(tui.output) + assert SMART_ROUTING_BANNER not in transcript, transcript + session.assert_not_routed() + task.assert_completed(session, "claude") + canary = session.home / ".ucode" / "claude-smart-routing-canary.json" + assert canary.is_file(), f"routing hooks were not armed: {canary}" + + +@pytest.mark.live +@pytest.mark.codex +def test_smart_routing_codex_subagent_only_launch_shows_no_first_prompt_banner( + live_session, workspace +): + """Scenario: configure Codex, then launch the real TUI with both the full and the + subagent-only routing flags set and submit one file prompt. + + Expected: subagent-only takes precedence over the ambient full flag: the prompt + completes with no smart-routing banner and no interposer first-prompt routing wrapper + anywhere in the session, and the TUI exits normally. Only first-prompt silence is + asserted here; subagent routing engagement is covered by the route-subagent hook + journey above. + """ + session = live_session + session.env["ENABLE_SMART_ROUTING_V2"] = "1" + session.env["ENABLE_SMART_ROUTING_SUBAGENT_ONLY"] = "1" + session.run( + "configure", + "--agents", + "codex", + "--workspace", + workspace, + "--skip-validate", + "--skip-upgrade", + "--disable-databricks-ai-tools", + ) + task = FileTask(session) + with AgentTerminal( + session, "codex", [str(session.binary), "codex"], "subagent-only-launch" + ) as tui: + tui.boot() + tui.submit(task.prompt) + tui.wait_for_task(task) + tui.exit_normally() + transcript = "".join(tui.output) + assert SMART_ROUTING_BANNER not in transcript, transcript + session.assert_not_routed() + task.assert_completed(session, "codex") diff --git a/tests/test_claude_smart_routing_v2.py b/tests/test_claude_smart_routing_v2.py index fdd3da5f..e8530ef4 100644 --- a/tests/test_claude_smart_routing_v2.py +++ b/tests/test_claude_smart_routing_v2.py @@ -172,6 +172,28 @@ def test_first_prompt_hook_is_per_launch(self): assert "user-policy" in str(settings["hooks"]["PreToolUse"]) +class TestSmartRoutingEnvVars: + def test_either_flag_enables_smart_routing(self, monkeypatch): + monkeypatch.delenv(v2.ENABLE_SMART_ROUTING_ENV_VAR, raising=False) + monkeypatch.delenv(v2.ENABLE_SUBAGENT_ROUTING_ENV_VAR, raising=False) + assert not v2.smart_routing_enabled() + monkeypatch.setenv(v2.ENABLE_SUBAGENT_ROUTING_ENV_VAR, "1") + assert v2.smart_routing_enabled() + monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_ENV_VAR, "1") + assert v2.smart_routing_enabled() + + def test_subagent_only_flag_suppresses_first_prompt_routing(self, monkeypatch): + monkeypatch.delenv(v2.ENABLE_SMART_ROUTING_ENV_VAR, raising=False) + monkeypatch.setenv(v2.ENABLE_SUBAGENT_ROUTING_ENV_VAR, "1") + assert not v2.first_prompt_routing_enabled() + # Subagent-only wins when both are set: an ambient full flag cannot + # override an explicit subagent-only session. + monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_ENV_VAR, "1") + assert not v2.first_prompt_routing_enabled() + monkeypatch.delenv(v2.ENABLE_SUBAGENT_ROUTING_ENV_VAR) + assert v2.first_prompt_routing_enabled() + + class TestV2Launch: def test_strips_gateway_prefix_for_interposer(self): model = "anthropic-aigw-73ea02b2-system.ai.glm-5-2" @@ -182,6 +204,7 @@ def test_restores_model_captured_immediately_before_switch(self, tmp_path, monke user_settings = tmp_path / "settings.json" ucode_settings.write_text(json.dumps({"env": {"ANTHROPIC_BASE_URL": "https://gw"}})) user_settings.write_text(json.dumps({"model": "opus", "theme": "dark"})) + monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_ENV_VAR, "1") monkeypatch.setattr(claude, "APP_DIR", tmp_path) monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", ucode_settings) monkeypatch.setattr(claude, "CLAUDE_USER_SETTINGS_PATH", user_settings) @@ -290,6 +313,7 @@ def fake_run(argv, **kwargs): def test_does_not_restore_when_wrapper_never_switches(self, tmp_path, monkeypatch): user_settings = tmp_path / "settings.json" user_settings.write_text(json.dumps({"model": "opus"})) + monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_ENV_VAR, "1") monkeypatch.setattr(v2, "APP_DIR", tmp_path) monkeypatch.setattr(v2, "get_databricks_token", lambda *_args, **_kwargs: "token") monkeypatch.setattr(v2, "build_auth_token_argv", lambda *_args, **_kwargs: ["ug"]) @@ -323,6 +347,7 @@ def test_restores_after_routed_model_persists_and_preserves_later_choice( ): user_settings = tmp_path / "settings.json" user_settings.write_text(json.dumps({"model": "haiku", "theme": "dark"})) + monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_ENV_VAR, "1") monkeypatch.setattr(v2, "APP_DIR", tmp_path) monkeypatch.setattr(v2, "get_databricks_token", lambda *_args, **_kwargs: "token") monkeypatch.setattr(v2, "build_auth_token_argv", lambda *_args, **_kwargs: ["ug"]) @@ -361,6 +386,77 @@ def fake_run(_argv, **kwargs): "theme": "dark", } + def test_subagent_only_launch_skips_first_prompt_routing(self, tmp_path, monkeypatch): + user_settings = tmp_path / "settings.json" + user_settings.write_text(json.dumps({"model": "opus"})) + monkeypatch.delenv(v2.ENABLE_SMART_ROUTING_ENV_VAR, raising=False) + monkeypatch.setenv(v2.ENABLE_SUBAGENT_ROUTING_ENV_VAR, "1") + monkeypatch.setenv(claude.GATEWAY_MODEL_DISCOVERY_ENV_VAR, "") + monkeypatch.setenv("CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY", "") + monkeypatch.setattr(v2, "APP_DIR", tmp_path) + monkeypatch.setattr(v2, "_model_picker_catalog", lambda: None) + monkeypatch.setattr(v2, "get_databricks_token", lambda *_args, **_kwargs: "token") + monkeypatch.setattr(v2, "build_auth_token_argv", lambda *_args, **_kwargs: ["ug"]) + monkeypatch.setattr( + v2, + "list_anthropic_model_catalog", + lambda *_args: AnthropicModelCatalog( + model_ids=["system.ai.claude-opus-4-8"], model_id_to_display_name={} + ), + ) + monkeypatch.setattr( + claude_pty, + "run_claude_pty", + lambda *_args, **_kwargs: pytest.fail("subagent-only routing must not use the PTY"), + ) + captured: dict = {} + + class FakeProcess: + def __init__(self, argv, **_kwargs): + captured["argv"] = argv + settings_path = Path(argv[argv.index("--settings") + 1]) + captured["settings_path"] = settings_path + captured["settings"] = json.loads(settings_path.read_text()) + captured["agents"] = json.loads(argv[argv.index("--agents") + 1]) + + def wait(self): + return 4 + + def send_signal(self, _signal): + raise AssertionError("test does not interrupt Claude") + + monkeypatch.setattr(v2.subprocess, "Popen", FakeProcess) + + with pytest.raises(SystemExit) as exc: + v2.launch_claude( + {"workspace": "https://example.com"}, + [], + binary="claude", + user_settings_path=user_settings, + launch_model="opus", + compose_settings=lambda _args: ({}, []), + launch_model_args=claude._launch_model_args, + model_name=claude._maybe_add_1m_suffix, + ) + + assert exc.value.code == 4 + settings = captured["settings"] + env = settings["env"] + assert env[v2.ENABLE_SUBAGENT_ROUTING_ENV_VAR] == "1" + assert v2.ENABLE_SMART_ROUTING_ENV_VAR not in env + assert claude_hooks.FIRST_PROMPT_SOCKET_ENV not in env + # Subagent routing is fully wired; only the first-prompt machinery is absent. + assert "UserPromptSubmit" not in settings["hooks"] + assert "route-subagent" in str(settings["hooks"]["PreToolUse"]) + assert settings["modelOverrides"] == {"claude-opus-4-8": "system.ai.claude-opus-4-8"} + assert {definition["model"] for definition in captured["agents"].values()} == { + "system.ai.claude-opus-4-8" + } + assert captured["argv"][3:5] == ["--model", "opus"] + assert not captured["settings_path"].exists() + # The model-setting guard is a first-prompt concern; user settings stay untouched. + assert json.loads(user_settings.read_text()) == {"model": "opus"} + class TestV2ModelPickerDiscovery: """modelPicker takes priority over gateway model discovery for smart routing.""" @@ -369,6 +465,7 @@ class TestV2ModelPickerDiscovery: def _launch(monkeypatch, tmp_path, *, picker_catalog): user_settings = tmp_path / "settings.json" user_settings.write_text(json.dumps({"model": "opus"})) + monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_ENV_VAR, "1") monkeypatch.setattr(v2, "APP_DIR", tmp_path) monkeypatch.setattr(v2, "CLAUDE_PTY_LOG", tmp_path / "v2.log") monkeypatch.setattr(v2, "get_databricks_token", lambda *_args, **_kwargs: "token") diff --git a/tests/test_cli.py b/tests/test_cli.py index 4a415b58..2a0824a7 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -579,6 +579,25 @@ def test_native_subcommand_suppresses_inherited_smart_routing( assert observed == [None] assert os.environ[cli_mod.smart_routing_v2.ENABLE_SMART_ROUTING_ENV_VAR] == "1" + @pytest.mark.parametrize("tool, subcommand", [("codex", "app"), ("claude", "update")]) + def test_native_subcommand_suppresses_inherited_subagent_routing( + self, monkeypatch, tool, subcommand + ): + monkeypatch.setenv("ENABLE_SMART_ROUTING_SUBAGENT_ONLY", "1") + observed = [] + + with patch( + "ucode.cli._launch_tool", + side_effect=lambda *_args, **_kwargs: observed.append( + os.environ.get(cli_mod.smart_routing_v2.ENABLE_SUBAGENT_ROUTING_ENV_VAR) + ), + ): + result = runner.invoke(app, [tool, subcommand]) + + assert result.exit_code == 0, result.output + assert observed == [None] + assert os.environ[cli_mod.smart_routing_v2.ENABLE_SUBAGENT_ROUTING_ENV_VAR] == "1" + def test_claude_enable_smart_routing_forwards_positional_prompt_to_v2(self): captured = [] diff --git a/tests/test_codex_smart_routing_v2.py b/tests/test_codex_smart_routing_v2.py index 7d8fa12a..fe2f169f 100644 --- a/tests/test_codex_smart_routing_v2.py +++ b/tests/test_codex_smart_routing_v2.py @@ -1,6 +1,7 @@ from __future__ import annotations import json +import os import pytest @@ -187,6 +188,7 @@ def test_owns_app_server_interposer_and_tui_lifecycle(self, monkeypatch): interposer_args = {} stopped = [] token_calls = [] + monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_ENV_VAR, "1") monkeypatch.setenv("CODEX_HOME", "/user/codex-home") monkeypatch.setattr(codex, "ug_version", lambda: "0.1.0") monkeypatch.setattr(codex, "agent_version", lambda binary: "0.148.0") @@ -288,6 +290,51 @@ def start_interposer(*args, **kwargs): assert stopped == [True] assert processes[0].terminated is True + def test_subagent_only_launch_runs_tui_directly(self, tmp_path, monkeypatch): + monkeypatch.setenv(v2.ENABLE_SUBAGENT_ROUTING_ENV_VAR, "1") + monkeypatch.setenv("CODEX_HOME", str(tmp_path)) + monkeypatch.setattr(codex, "ug_version", lambda: "0.1.0") + monkeypatch.setattr(codex, "agent_version", lambda binary: "0.148.0") + monkeypatch.setattr(v2, "get_databricks_token", lambda *_args, **_kwargs: "token") + monkeypatch.setattr( + v2.subprocess, + "Popen", + lambda *_args, **_kwargs: pytest.fail("subagent-only routing spawns no app-server"), + ) + monkeypatch.setattr( + codex_interposer, + "start_interposer_thread", + lambda *_args, **_kwargs: pytest.fail("subagent-only routing must not interpose"), + ) + execd = [] + + def fake_exec(argv): + execd.append(argv) + raise SystemExit(0) + + monkeypatch.setattr(v2, "exec_or_spawn", fake_exec) + + with pytest.raises(SystemExit) as exc: + v2.launch_codex( + {"workspace": WS, "codex_models": ["system.ai.gpt-5-6-sol"]}, + ["--search"], + binary="codex", + start_model="gpt-start", + render_overlay=codex.render_overlay, + ) + + assert exc.value.code == 0 + (argv,) = execd + assert argv[0] == "codex" + assert argv[-1] == "--search" + assert 'model="gpt-start"' in argv + hook_override = next(arg for arg in argv if arg.startswith("hooks.PreToolUse=")) + assert "codex-router-hook route-subagent" in hook_override + assert "--model system.ai.gpt-5-6-sol" in hook_override + # The hook subprocesses inherit the launch environment and pass the routing gate. + assert os.environ[v2.ENABLE_SUBAGENT_ROUTING_ENV_VAR] == "1" + assert os.environ[v2.OAUTH_TOKEN_ENV_VAR] == "token" + def test_v2_pre_tool_hook_preserves_user_hooks(self, tmp_path, monkeypatch): codex_home = tmp_path / ".codex" codex_home.mkdir() @@ -341,6 +388,7 @@ def test_v2_pre_tool_hook_replaces_existing_ucode_hook(self, tmp_path, monkeypat assert "--model old" not in routing_commands[0] def test_missing_cached_models_starts_with_bootstrap_model(self, monkeypatch): + monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_ENV_VAR, "1") monkeypatch.setattr(v2, "get_databricks_token", lambda workspace, profile: "token") monkeypatch.setattr(codex, "agent_version", lambda binary: "unknown") monkeypatch.setattr(v2, "_free_port", lambda: 41001) @@ -482,6 +530,7 @@ def test_launch_prefers_catalog_over_cached_models(self, tmp_path, monkeypatch, monkeypatch, cli=self._catalog(tmp_path / "cli.json", ["gpt-6-astra", "gpt-6-b"]), ) + monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_ENV_VAR, "1") monkeypatch.setattr(v2, "get_databricks_token", lambda *_args: "token") monkeypatch.setattr(v2, "_free_port", lambda: 41001) monkeypatch.setattr(v2, "_wait_for_app_server", lambda port, timeout: True)