diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 24dbaab06..2de406626 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -173,7 +173,17 @@ jobs: # fallback EOFs instead of hanging the runner. With DATABRICKS_BEARER # set, the auth code path doesn't shell out at all — this is a safety # net for any code path we may have missed. - - run: uv run --no-sync pytest tests/test_e2e.py -v -k "$TEST_KEYWORD" < /dev/null + # Rerun a case once only for live-service timeouts and transient upstream errors; + # every other failure fails on the first attempt. `-rfER` lists reruns in the summary. + - run: | + uv run --no-sync pytest tests/test_e2e.py -v -rfER -k "$TEST_KEYWORD" \ + --reruns 1 --reruns-delay 5 \ + --only-rerun TimeoutExpired \ + --only-rerun 'timed out after [0-9]+s' \ + --only-rerun 'TIMEOUT after' \ + --only-rerun TEMPORARILY_UNAVAILABLE \ + --only-rerun REQUEST_LIMIT_EXCEEDED \ + < /dev/null e2e: name: All agent tests diff --git a/pyproject.toml b/pyproject.toml index dfcf91d19..563a75242 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -58,6 +58,7 @@ fallback-version = "0.0.0" [dependency-groups] dev = [ "pytest>=9.0.3", + "pytest-rerunfailures>=16.7", "pexpect>=4.9.0", "pyte>=0.8.2", "pywinpty==3.0.5; sys_platform == 'win32'", diff --git a/scripts/run_integration.py b/scripts/run_integration.py index c40493304..3d1f9cbcd 100644 --- a/scripts/run_integration.py +++ b/scripts/run_integration.py @@ -39,7 +39,9 @@ # Seeds machine-wide managed settings, which ug has no Windows path for. "test_ug_configure_claude_lifecycle.py::test_unmanaged_claude_preserves_preexisting_family_defaults", # Codex's Windows sandbox can't read ug's owner-only session-env dir yet. - "test_ug_smart_routing_hooks.py::test_smart_router_skill_toggles_codex_subagent_routing", + "test_ug_smart_routing_hooks.py::test_smart_router_routes_codex_subagent_at_launch", + "test_ug_smart_routing_hooks.py::test_smart_router_skill_turns_codex_subagent_routing_off", + "test_ug_smart_routing_hooks.py::test_smart_router_skill_turns_codex_subagent_routing_back_on", } HEADLESS_TEST_NODES = { "claude": "test_ug_claude_headless.py::test_ug_claude_headless_prompt_argument", diff --git a/tests/README.md b/tests/README.md index af82fc34e..f453583df 100644 --- a/tests/README.md +++ b/tests/README.md @@ -40,8 +40,9 @@ Claude mod entry-point coverage lives in `test_claude_mod.py` (generated plugin `test_claude_smart_routing_v2.py` (launch wiring), and the explicitly selected `native_claude/test_mod.py` suite. Claude 2.1.290 validates the module and runs its TypeScript event-forwarding and launch-option tests in CI. The existing Claude -skill-toggle integration journey checks the live recipe payload field and Agent/Task -delegation descriptions across on/off/on. Native request-construction checks verify the +skill-toggle integration journeys check the live recipe payload field (launch recipe at +launch and after off/on, DISABLED after off) and Agent/Task delegation descriptions. +Native request-construction checks verify the description switches after both native commands and CLI session toggles, including native toggles after resuming a saved prompt-snapshot fixture. The description includes instructions to end the current turn while awaiting a @@ -119,7 +120,7 @@ under `e2e_cuj/` and `integration/`; no workflow edit is needed. Integration cas retain their agent and workspace markers, including `managed_fixture` when needed. See [the harness opt-in guide](integration/README.md#opting-into-harness-latest). -CUJ 4's preset journeys check Claude's ON/OFF status band, routing scope, and +CUJ 4's Claude-only preset band journey (separate from the delegation journeys) checks Claude's ON/OFF status band, routing scope, and updates after off/on toggles, and that the band is hidden after the launch turn and by the second turn after each toggle. They also check that the prompt hint row under the input ends with `smart router ON` from launch, flips to `OFF` and back to `ON` with each @@ -147,8 +148,8 @@ Offline PTY checks verify that the Claude background-task wait observes "No task running" or a native task menu containing only completed rows before sending `/exit`, without stopping tasks or confirming an exit dialog. Running/scheduled sections, incomplete row counts, and completed-task text outside the native menu cannot satisfy the wait. -The live Claude subagent skill-toggle journey uses this wait after its final calculation. -CUJ4's Claude preset sessions also use it after verifying delegated file tasks: a completed +The live Claude subagent skill-toggle journeys use this wait after their final calculation. +CUJ4's Claude delegation sessions also use it after verifying delegated file tasks: a completed answer does not prove that a resumed child has stopped running. They cover Claude/Codex helper dispatch and rejection of routing decisions without the agent-specific prompt-submission evidence. @@ -187,9 +188,9 @@ CUJ4 accepts both top-level tool definitions (Codex 0.154.0) and Responses Lite task name to the recipient; plaintext prompts still match exactly. Both require tools. The original smart-routing CUJ runs four fresh sessions: routed and explicit model for both Claude and Codex, with explicit-model bypass assertions enabled for both agents. -One additional test has five `SMART_ROUTER_CONFIG_VERSION` cases. -Each case exercises both agents and checks first-prompt routing, orchestrator context in -inference input, and a completed explicitly requested routed subagent. It does not establish +Per-agent tests cover six `SMART_ROUTER_CONFIG_VERSION` cases each; every case checks +first-prompt routing, orchestrator context in inference input, and a completed explicitly +requested routed subagent. A separate Claude-only test covers the status band and hint row. It does not establish automatic orchestrator delegation. Routing-disabled coverage is deferred until a separately preconfigured workspace is assigned. @@ -606,7 +607,7 @@ integration utilities; only CUJ-specific evidence correlation stays in a test fi | `test_ug_codex_app_reports_unknown_argument` | Pass an invalid option directly to `ug codex app` with legacy smart routing off/on | Real Codex parser error and status preserved | | `test_ug_codex_app_server_client_initializes` | Connect a stdio client, direct/`--` separator, with legacy smart 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 native spawn payload to the installed Claude mod command or Codex PreToolUse hook with legacy subagent routing or each supported routing preset | Claude returns model/rationale/updatedInput JSON naming the routed agent the mod labels; Codex returns an allow-hook decision; 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_router_skill_toggles_claude_subagent_routing`, `test_smart_router_skill_toggles_codex_subagent_routing` | Configure, launch a real subagent-only TUI with legacy routing-only/orchestration flags or the three subagent-only presets, then spawn tagged children while invoking the installed Smart Router skill to switch routing on -> off -> on in the same session | Routing-only cases install `smart-router`; orchestration cases also install `smart-router-orchestrator`; all three native children complete; only routing-enabled phases show a native model/rationale row (Claude) or subagent banner (Codex) and produce a live routing decision correlated with the child; Claude's recorded task inference payloads carry `smart_router_recipe_name` as launch recipe -> DISABLED -> launch recipe, omit the recipe header, and receive successful gateway responses; opted-in Codex banners include `[orchestrator on]`; no first-prompt routing wrapper; the customer full-mode preset is intentionally outside this journey; normal exit | +| `test_smart_router_routes_{claude,codex}_subagent_at_launch`, `test_smart_router_skill_turns_{claude,codex}_subagent_routing_off`, `test_smart_router_skill_turns_{claude,codex}_subagent_routing_back_on` | Configure, launch a real subagent-only TUI with legacy routing-only/orchestration flags or the three subagent-only presets, each in its own fresh session: spawn a tagged child at launch; invoke the installed Smart Router skill to turn routing off then spawn a child; or spawn a routed child, turn routing off and spawn an unrouted child, then turn routing back on and spawn a routed child | Routing-only cases install `smart-router`; orchestration cases also install `smart-router-orchestrator`; each tagged native child completes; routed children (launch, before off and after back on) show a native model/rationale row (Claude) or subagent banner (Codex) and produce a live routing decision correlated with the child, while each child after off shows neither and has no decision; Claude's recorded task inference payloads carry `smart_router_recipe_name` as the launch recipe (`at_launch`, `back_on`, where the second child must match the launch child) or DISABLED (after off), omit the recipe header, and receive successful gateway responses; opted-in Codex banners include `[orchestrator on]`; no first-prompt routing wrapper; the customer full-mode preset is intentionally outside this journey; 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_cleans_stale_skills_mcp_on_workspace_switch` | Configure the first workspace, register its skills MCP, switch to a second real workspace, and use Claude | Old registration removed from Claude and the new workspace state; old workspace bucket preserved; repeat configure stays clean; real file task completes on the second workspace | | `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 | @@ -846,7 +847,7 @@ verify that neither a sudo capability probe nor a prompting worker is started. | Fresh provider/parent validation and mixed Bedrock filtering | Not covered after removing the duplicate model-discovery suites | | 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 through `first_prompt_and_subagent_no_orch_v0` is covered by the `managed_fixture` banner journeys and CUJ4's preset journeys. CUJ4's managed journeys and the managed-default skill-toggle cases verify that `smart_routing.enabled` alone leaves the first prompt unrouted and turns on Smart Router Orchestrator; `--enable-smart-routing` has component coverage only. The live subagent-only journeys spawn native children and invoke the installed Smart Router skill to verify routing on -> off -> on within one Claude/Codex session. The scheduled harness-latest workflow repeats CUJ4, routing hooks, skill toggles, and banners with both latest agents. Plugin-refresh survival and native daemon/background dispatch remain deferred. | +| Claude/Codex interactive smart routing | First-prompt routing through `first_prompt_and_subagent_no_orch_v0` is covered by the `managed_fixture` banner journeys and CUJ4's preset journeys. CUJ4's managed journeys and the managed-default skill-toggle cases verify that `smart_routing.enabled` alone leaves the first prompt unrouted and turns on Smart Router Orchestrator; `--enable-smart-routing` has component coverage only. The live subagent-only journeys spawn native children and invoke the installed Smart Router skill to verify routing at launch, off, and off -> on in separate Claude/Codex sessions. The scheduled harness-latest workflow repeats CUJ4, routing hooks, skill toggles, and banners with both latest agents. Plugin-refresh survival and native daemon/background dispatch remain deferred. | | 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. Codex desktop's shared-config catalog pointer (repeat configure and a user-owned pointer) is covered by `integration/test_ug_configure_codex_desktop.py`; desktop tracing, GUI, and inference are not | | Native macOS/Windows live TUI, managed settings, resize/signals | Windows fresh-install, CLI, local helpers, and the Claude headless gateway journey are advisory; Codex Windows CI is blocked on npm proxy access. Live PTY/TUI, managed settings, and signal behavior still need separate platform implementation and coverage | @@ -862,7 +863,7 @@ Native launch-option component tests (`native_claude/test_mod.py`) check two rea processes with different resolved routing baselines. Python cases cover Claude-specific presets, unknown selectors, and absent versus empty flags. These checks do not establish gateway/TUI coverage. -The real session-control process call is covered by the live skill-toggle journey; +The real session-control process call is covered by the live skill-toggle journeys; native component tests stub the subagent routing process protocol. No component check asserts that the launch paths withdraw the skills. The TypeScript launch test also checks event/result forwarding. Separate cases cover @@ -870,10 +871,10 @@ denied-write rollback, invalid request bodies, disabled launches, and body-merge no simulated subprocess success is used to claim toggle coverage. `integration/test_ug_smart_routing_hooks.py` checks `smart_router_recipe_name` in its existing -Claude skill-toggle journey. The recorder forwards to the real workspace, configured +Claude skill-toggle journeys. The recorder forwards to the real workspace, configured through public `ug configure --workspace ` without hand-writing ug state. Checkpoints and each calculation's unique marker correlate parent/child inference -requests to the on/off/on phases. All matched tool-capable inference payloads must carry the +requests to the launch, off, and back-on phases. All matched tool-capable inference payloads must carry the launch recipe, `DISABLED`, then the launch recipe again, omit the recipe header, and reach non-empty HTTP 200 responses, with only the existing verified native compatibility retries accepted. The expected launch recipe comes from the recorded routing request, so a default diff --git a/tests/conftest.py b/tests/conftest.py index 30ded0690..2e95c88c7 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2,6 +2,7 @@ from __future__ import annotations +import functools import json import os from types import SimpleNamespace @@ -156,12 +157,18 @@ def e2e_token(e2e_workspace): return get_databricks_token(e2e_workspace) -@pytest.fixture(scope="session") -def e2e_state(e2e_workspace, e2e_token): - """Full state dict mirroring what configure_shared_state produces.""" - claude_models = fetch_ai_gateway_claude_models(e2e_workspace, e2e_token) - gemini_models = fetch_gemini_models(e2e_workspace, e2e_token) - codex_models = fetch_codex_models(e2e_workspace, e2e_token) +@functools.cache +def e2e_state_for(workspace: str) -> dict: + """Full state dict mirroring what configure_shared_state produces. + + Cached per process so collection-time parametrization (tests/test_e2e.py) and the + ``e2e_state`` fixture query the live catalogs once. Keyed on the workspace only: a bearer + command mints a fresh token per call, so a token key would miss and fetch a second snapshot. + A failed fetch raises and is not cached.""" + token = get_databricks_token(workspace) + claude_models = fetch_ai_gateway_claude_models(workspace, token) + gemini_models = fetch_gemini_models(workspace, token) + codex_models = fetch_codex_models(workspace, token) opencode_models: dict = {} if claude_models: @@ -170,11 +177,18 @@ def e2e_state(e2e_workspace, e2e_token): opencode_models["gemini"] = gemini_models return { - "workspace": e2e_workspace, + "workspace": workspace, "claude_models": claude_models, "gemini_models": gemini_models, "codex_models": codex_models, "opencode_models": opencode_models, - "base_urls": build_shared_base_urls(e2e_workspace), + "base_urls": build_shared_base_urls(workspace), "managed_configs": {}, } + + +@pytest.fixture(scope="session") +def e2e_state(e2e_workspace): + """Full state dict mirroring what configure_shared_state produces.""" + # Copy so a test mutating its state cannot leak into the cached catalog. + return dict(e2e_state_for(e2e_workspace)) diff --git a/tests/e2e_cuj/README.md b/tests/e2e_cuj/README.md index 89349d462..cf3c4e10a 100644 --- a/tests/e2e_cuj/README.md +++ b/tests/e2e_cuj/README.md @@ -73,10 +73,14 @@ They do not launch Desktop or execute the credential helper. `test_cuj4_smart_routing.py` uses its own read-only workspace with managed smart routing enabled. It runs the managed-default and explicit-model journeys and six -`SMART_ROUTER_CONFIG_VERSION` cases using the shared version constants; each launches Claude -and Codex once. With no selector, managed smart routing leaves the first prompt unrouted: the -admin's default model completes the file task, no `routes:select` request is sent, and the -Smart Router Orchestrator context reaches that inference request. +`SMART_ROUTER_CONFIG_VERSION` presets using the shared version constants. With no selector, +managed smart routing leaves the first prompt unrouted: the admin's default model completes the +file task, no `routes:select` request is sent, and the Smart Router Orchestrator context reaches +that inference request. Each preset runs as independent tests so a slow live subagent call +cannot fail unrelated assertions: one first-prompt-then-explicit-subagent journey per agent +(12 cases, with Claude ending in the native background-task wait) and one Claude-only status +band/hint row/toggle journey per preset (6 cases, which needs no delegation or background-task +wait). The CUJ derives routing and orchestration assertions from each agent's resolved preset flags. Component tests independently verify the preset values. The resulting expectations are: @@ -137,5 +141,5 @@ With the same live prerequisites: ```bash uv run --with pexpect==4.9.0 --with pyte==0.8.2 pytest \ --confcutdir=tests/e2e_cuj tests/e2e_cuj/test_cuj4_smart_routing.py \ - -k test_smart_router_config_version -v + -k test_smart_router_preset -v ``` diff --git a/tests/e2e_cuj/test_cuj2_mps_mcp.py b/tests/e2e_cuj/test_cuj2_mps_mcp.py index 446fdaf85..f1bd16eab 100644 --- a/tests/e2e_cuj/test_cuj2_mps_mcp.py +++ b/tests/e2e_cuj/test_cuj2_mps_mcp.py @@ -38,6 +38,8 @@ CUJ_NAME = "CUJ 2 · MPS and named MCP" WORKSPACE_URL = "https://dbc-0dcf95cf-e357.cloud.databricks.com" +# One budget for the inference request to appear and for its response to complete. +INFERENCE_WAIT_SECONDS = 240 def expected_agent_config(agent: str) -> dict: @@ -203,6 +205,29 @@ def _agent_prompt(marker: str) -> str: return f"Reply with the exact marker `{marker}` and no other text. Do not use tools." +def _served_marker_request(recorder, checkpoint: int, agent: str, marker: str): + """The marker-bearing inference request the gateway answered with HTTP 200. + + Claude retries a request whose proxy-to-gateway connection failed (the recorder then + records a 502 for the first attempt), so the first marker request is not always the + one that was served. With no 200, return the first so its failure is asserted. + """ + candidates = [ + request + for request in recorder.requests_after(checkpoint) + if request.method == "POST" + and request.path == INFERENCE_PATHS[agent] + and marker.encode() in request.body + ] + assert candidates, f"No {agent} inference request carried marker {marker}" + # Stop at the first success: later requests resend the history (and so the marker), + # and one the agent abandoned at exit never records a response to wait for. + for request in candidates: + if recorder.response_for(request, timeout=INFERENCE_WAIT_SECONDS).status_code == 200: + return request + return candidates[0] + + def _assert_inference_request(recorder, request, *, provider: str, model: str, marker: str) -> None: assert request.headers[PROVIDER_HEADER] == provider assert marker.encode() in request.body, request.payload @@ -290,7 +315,7 @@ def test_cuj_codex_inference(self, cuj): method="POST", path=INFERENCE_PATHS[CODEX], after=checkpoint, - timeout=240, + timeout=INFERENCE_WAIT_SECONDS, ) _assert_inference_request( recorder, @@ -336,12 +361,13 @@ def test_cuj_claude_inference(self, cuj): tui.wait_until(lambda: _task_complete(session, CLAUDE, marker), "completed Claude task") tui.exit_normally() - request = recorder.expect_request( + recorder.expect_request( method="POST", path=INFERENCE_PATHS[CLAUDE], after=checkpoint, - timeout=240, + timeout=INFERENCE_WAIT_SECONDS, ) + request = _served_marker_request(recorder, checkpoint, CLAUDE, marker) _assert_inference_request( recorder, request, diff --git a/tests/e2e_cuj/test_cuj4_smart_routing.py b/tests/e2e_cuj/test_cuj4_smart_routing.py index 1ce1cda74..9f81bf139 100644 --- a/tests/e2e_cuj/test_cuj4_smart_routing.py +++ b/tests/e2e_cuj/test_cuj4_smart_routing.py @@ -1,6 +1,7 @@ """Smart-routing CUJ: execute once, then report each contract assertion separately.""" import re +from contextlib import contextmanager from dataclasses import dataclass import pytest @@ -42,6 +43,14 @@ ROUTING_PATH = "/ai-gateway/routing/v1/routes:select" AGENTS = (CLAUDE, CODEX) ORCHESTRATOR_CONTEXT = "Smart Router Orchestrator is on for this session." +SMART_ROUTER_CONFIG_VERSIONS = ( + FIRST_PROMPT_AND_SUBAGENT_NO_ORCH_V0, + SUBAGENT_ONLY_V0, + SUBAGENT_ONLY_V1, + SUBAGENT_ORCH_V0, + SUBAGENT_ORCH_V1, + SUBAGENT_ORCH_V0_CLAUDE_ONLY, +) def _smart_router_band_shown(tui, *, enabled, first_prompt): @@ -66,6 +75,47 @@ def _smart_router_band_hidden(tui): return not any(re.match(pattern, line) for line in tui.screen.display) +@dataclass(frozen=True) +class PresetFlags: + subagent_routed: bool + first_prompt_routed: bool + orchestrator_enabled: bool + + +def _preset_flags(version, agent): + """Resolve the preset's routing and orchestration expectations for one agent.""" + env = resolve_environment({"SMART_ROUTER_CONFIG_VERSION": version}, agent=agent) + v2_enabled = env["ENABLE_SMART_ROUTING_V2"] == "1" + subagent_only = env["ENABLE_SMART_ROUTING_SUBAGENT_ONLY"] == "1" + subagent_routed = v2_enabled or subagent_only + return PresetFlags( + subagent_routed=subagent_routed, + first_prompt_routed=v2_enabled and not subagent_only, + orchestrator_enabled=env["ENABLE_SMART_ROUTER_ORCHESTRATOR"] == "1" and subagent_routed, + ) + + +@contextmanager +def smart_router_preset(cuj, version): + """Select a preset for one fresh ug configuration, then restore env and machine-wide settings.""" + session, workspace, recorder = cuj + previous = session.env.get("SMART_ROUTER_CONFIG_VERSION") + session.env["SMART_ROUTER_CONFIG_VERSION"] = version + try: + configs = _assert_published_config_matches_expectations(workspace.config()) + recorder.configure_session(session, ["configure", "--disable-databricks-ai-tools"]) + yield session, workspace, recorder, configs + finally: + if previous is None: + session.env.pop("SMART_ROUTER_CONFIG_VERSION", None) + else: + session.env["SMART_ROUTER_CONFIG_VERSION"] = previous + session.revert_machine_wide( + f"{version}-revert", + "Smart-router preset cleanup left machine-wide agent settings", + ) + + @dataclass(frozen=True) class SessionCase: agent: str @@ -267,254 +317,246 @@ def completed_smart_routing_runs(cuj): class TestCujSmartRouting(BaseCujTest): WORKSPACE_URL = "https://dbc-1a9622fc-2e91.cloud.databricks.com/" - @pytest.mark.parametrize( - "SMART_ROUTER_CONFIG_VERSION", - [ - FIRST_PROMPT_AND_SUBAGENT_NO_ORCH_V0, - SUBAGENT_ONLY_V0, - SUBAGENT_ONLY_V1, - SUBAGENT_ORCH_V0, - SUBAGENT_ORCH_V1, - SUBAGENT_ORCH_V0_CLAUDE_ONLY, - ], - ) - def test_smart_router_config_version(self, cuj, SMART_ROUTER_CONFIG_VERSION): - """Scenario: launch both agents with a preset, then explicitly request a subagent. - - Expected: each agent's resolved flags determine first-prompt routing, - orchestration, and child routing. One native child completes the delegated - task even when routing is disabled. Claude's status band shows ON and the - preset's routing scope, then OFF and ON after toggles; each time it hides once - that turn completes, so it is gone by the second turn. The prompt hint row ends with - smart router ON or OFF from launch and flips with each toggle. + @pytest.mark.parametrize("agent", AGENTS) + @pytest.mark.parametrize("SMART_ROUTER_CONFIG_VERSION", SMART_ROUTER_CONFIG_VERSIONS) + def test_smart_router_preset_routes_first_prompt_then_explicit_native_subagent( + self, cuj, agent, SMART_ROUTER_CONFIG_VERSION + ): + """Scenario: launch one agent with a preset, run a no-delegation file task, then ask for a subagent. + + Expected: the agent's resolved flags determine whether the first prompt is routed + (one HTTP 200 decision matching the prompt, and the routed model serves inference), + whether orchestrator context and the smart-router recipe header/field are present, + and that no child session exists after the first prompt. A second prompt then makes + exactly one native child complete the delegated file task even when routing is + disabled. When subagent routing is on, exactly one decision correlated with the child + task is recorded, the child's inference uses the selected model, and the recipe + correlates with the launch turn; when off, no decision, recipe, or orchestrator + context appears. Claude then reports no running background tasks and exits cleanly. """ - session, workspace, recorder = cuj - previous = session.env.get("SMART_ROUTER_CONFIG_VERSION") - session.env["SMART_ROUTER_CONFIG_VERSION"] = SMART_ROUTER_CONFIG_VERSION - try: - configs = _assert_published_config_matches_expectations(workspace.config()) - recorder.configure_session(session, ["configure", "--disable-databricks-ai-tools"]) - for agent in AGENTS: - expected_env = resolve_environment( - {"SMART_ROUTER_CONFIG_VERSION": SMART_ROUTER_CONFIG_VERSION}, agent=agent + with smart_router_preset(cuj, SMART_ROUTER_CONFIG_VERSION) as ( + session, + workspace, + recorder, + configs, + ): + flags = _preset_flags(SMART_ROUTER_CONFIG_VERSION, agent) + subagent_routed = flags.subagent_routed + first_prompt_routed = flags.first_prompt_routed + orchestrator_enabled = flags.orchestrator_enabled + supported = workspace.model_ids(agent) + evidence = SessionEvidence(session.home, agent) + existing_sessions = set(agent_sessions(session, agent)) + decisions_path = session.home / ".ucode" / f"{agent}-smart-routing-decisions.jsonl" + original_decisions = read_jsonl(decisions_path) + task = FileTask(session) + task.prompt += " Do not delegate." + checkpoint = recorder.checkpoint() + recorder.prepare_launch() + with Terminal(session, f"{SMART_ROUTER_CONFIG_VERSION}-{agent}", [agent]) as tui: + tui.boot(timeout=150) + tui.submit(task.prompt) + tui.task(evidence, task) + requests = recorder.requests_after(checkpoint) + routes = [request for request in requests if request.path == ROUTING_PATH] + assert len(routes) == int(first_prompt_routed), (agent, routes) + expected_model = configs[agent]["default_models"]["default_model"] + if first_prompt_routed: + assert routes[0].payload["task"]["prompt"] == task.prompt + response = recorder.response_for(routes[0]) + assert response.status_code == 200 + selections = response.payload["route_selection"] + assert len(selections) == 1 + expected_model = selections[0]["route_option"]["model"] + options = routes[0].payload["route_options"] + offered = {canonical_model(option["model"]) for option in options} + assert options and all(option["harness"] == agent for option in options) + assert offered <= {canonical_model(model) for model in supported} + selected = canonical_model(expected_model) + assert selected in offered and selected.startswith("system.ai."), selected + inference = _task_inference_request(requests, agent, task.prompt) + inference = served_inference_request(recorder, requests, inference, agent) + assert canonical_model(inference.payload["model"]) == canonical_model( + expected_model ) - v2_enabled = expected_env["ENABLE_SMART_ROUTING_V2"] == "1" - subagent_only = expected_env["ENABLE_SMART_ROUTING_SUBAGENT_ONLY"] == "1" - subagent_routed = v2_enabled or subagent_only - first_prompt_routed = v2_enabled and not subagent_only - orchestrator_enabled = ( - expected_env["ENABLE_SMART_ROUTER_ORCHESTRATOR"] == "1" and subagent_routed + evidence.assert_applied(task, supported, expected=expected_model) + if agent == CLAUDE: + launch_recipe = inference.payload.get("smart_router_recipe_name") + assert ( + isinstance(launch_recipe, str) + and launch_recipe.strip() + and launch_recipe != "DISABLED" + ) + if first_prompt_routed: + assert routes[0].payload["route_selector"]["router_name"] == launch_recipe + assert "databricks-smart-router-recipe" not in inference.headers + else: + assert ( + "databricks-smart-router-recipe" in inference.headers + ) == subagent_routed, agent + assert (ORCHESTRATOR_CONTEXT.encode() in inference.body) == orchestrator_enabled, ( + agent ) - supported = workspace.model_ids(agent) - evidence = SessionEvidence(session.home, agent) - existing_sessions = set(agent_sessions(session, agent)) - decisions_path = session.home / ".ucode" / f"{agent}-smart-routing-decisions.jsonl" - original_decisions = read_jsonl(decisions_path) - task = FileTask(session) - task.prompt += " Do not delegate." + assert not any( + is_child_session(agent, path, records) + for path, records in agent_sessions(session, agent).items() + if path not in existing_sessions + ), agent + + child_task = FileTask(session) + child_task.prompt = child_task.delegate_prompt + decision_count = len(read_jsonl(decisions_path)) checkpoint = recorder.checkpoint() - recorder.prepare_launch() - with Terminal(session, f"{SMART_ROUTER_CONFIG_VERSION}-{agent}", [agent]) as tui: - tui.boot(timeout=150) - if agent == CLAUDE: - tui.wait_for( - lambda _, first_prompt=first_prompt_routed: _smart_router_band_shown( - tui, enabled=True, first_prompt=first_prompt - ), - "the ON Smart Router band with the preset's routing scope", - ) - tui.wait_for( - lambda _: _smart_router_hint_shown(tui, enabled=True), - "smart router ON in the prompt hint row", - ) - tui.submit(task.prompt) - tui.task(evidence, task) + existing_sessions = set(agent_sessions(session, agent)) + tui.submit(child_task.prompt) + # The parent can resume a finished child before replying; wait for its + # reply so /tasks sees that run. + tui.wait_until( + lambda: ( + child_task.completed(session, agent, child=True) + and child_task.completed(session, agent) + ), + "completed native subagent file task and parent reply", + ) + children = { + path: records + for path, records in agent_sessions(session, agent).items() + if path not in existing_sessions and is_child_session(agent, path, records) + } + completed_children = { + path: records + for path, records in children.items() + if any( + child_task.value in answer for answer in assistant_answers(agent, records) + ) + } + assert len(completed_children) == 1, (agent, children.keys()) + decisions = read_jsonl(decisions_path)[decision_count:] + decision_tasks = { + decision.get("task_name") for decision in decisions if decision.get("task_name") + } + requests = recorder.requests_after(checkpoint) + routes = [ + request + for request in requests + if request.path == ROUTING_PATH + and request.payload.get("task", {}).get("prompt") in decision_tasks + ] + assert len(routes) == int(subagent_routed), (agent, routes) + if subagent_routed: + assert_subagent_routed( + session, + agent, + child_task, + decision_ids={decision["decision_id"] for decision in decisions}, + ) + route_prompt = routes[0].payload["task"]["prompt"] + assert isinstance(route_prompt, str) and route_prompt.strip() if agent == CLAUDE: - tui.wait_for( - lambda _: ( - _smart_router_band_hidden(tui) - and _smart_router_hint_shown(tui, enabled=True) - ), - "the Smart Router band hidden after the launch turn, with " - "smart router ON left in the prompt hint row", - ) - requests = recorder.requests_after(checkpoint) - routes = [request for request in requests if request.path == ROUTING_PATH] - assert len(routes) == int(first_prompt_routed), (agent, routes) - expected_model = configs[agent]["default_models"]["default_model"] - if first_prompt_routed: - assert routes[0].payload["task"]["prompt"] == task.prompt - response = recorder.response_for(routes[0]) - assert response.status_code == 200 - selections = response.payload["route_selection"] - assert len(selections) == 1 - expected_model = selections[0]["route_option"]["model"] - options = routes[0].payload["route_options"] - offered = {canonical_model(option["model"]) for option in options} - assert options and all(option["harness"] == agent for option in options) - assert offered <= {canonical_model(model) for model in supported} - selected = canonical_model(expected_model) - assert selected in offered and selected.startswith("system.ai."), selected - inference = _task_inference_request(requests, agent, task.prompt) + assert child_task.filename in route_prompt + response = recorder.response_for(routes[0]) + assert response.status_code == 200 + selections = response.payload["route_selection"] + assert len(selections) == 1 + inference = _task_inference_request( + requests, + agent, + route_prompt, + after=routes[0].sequence, + ) inference = served_inference_request(recorder, requests, inference, agent) assert canonical_model(inference.payload["model"]) == canonical_model( - expected_model + selections[0]["route_option"]["model"] ) - evidence.assert_applied(task, supported, expected=expected_model) if agent == CLAUDE: - launch_recipe = inference.payload.get("smart_router_recipe_name") - assert ( - isinstance(launch_recipe, str) - and launch_recipe.strip() - and launch_recipe != "DISABLED" - ) - if first_prompt_routed: - assert ( - routes[0].payload["route_selector"]["router_name"] == launch_recipe - ) + assert routes[0].payload["route_selector"]["router_name"] == launch_recipe + assert inference.payload.get("smart_router_recipe_name") == launch_recipe assert "databricks-smart-router-recipe" not in inference.headers else: - assert ( - "databricks-smart-router-recipe" in inference.headers - ) == subagent_routed, agent - assert ( - ORCHESTRATOR_CONTEXT.encode() in inference.body - ) == orchestrator_enabled, agent + assert "databricks-smart-router-recipe" in inference.headers + else: + assert read_jsonl(decisions_path) == original_decisions assert not any( - is_child_session(agent, path, records) - for path, records in agent_sessions(session, agent).items() - if path not in existing_sessions + "databricks-smart-router-recipe" in request.headers + or "smart_router_recipe_name" in request.payload + or ORCHESTRATOR_CONTEXT.encode() in request.body + for request in requests ), agent + if agent == CLAUDE: + tui.wait_for_background_tasks() + tui.exit_normally() - child_task = FileTask(session) - child_task.prompt = child_task.delegate_prompt - decision_count = len(read_jsonl(decisions_path)) - checkpoint = recorder.checkpoint() - existing_sessions = set(agent_sessions(session, agent)) - tui.submit(child_task.prompt) - # The parent can resume a finished child before replying; wait for its - # reply so /tasks sees that run. - tui.wait_until( - lambda task=child_task, agent=agent: ( - task.completed(session, agent, child=True) - and task.completed(session, agent) + @pytest.mark.parametrize("SMART_ROUTER_CONFIG_VERSION", SMART_ROUTER_CONFIG_VERSIONS) + def test_smart_router_preset_claude_status_band_and_hint_row_follow_launch_and_toggles( + self, cuj, SMART_ROUTER_CONFIG_VERSION + ): + """Scenario: launch Claude with a preset, run one file task, then toggle /smart-router off and on. + + Expected: at launch the status band shows ON with the preset's routing scope and the + prompt hint row ends with smart router ON. After the launch turn completes the band is + hidden and the hint row still shows ON. Each toggle (off, then on) shows the matching + band, flips the hint row immediately, and after a following completed file task the + band is hidden again while the hint row keeps the toggled state. No delegation or + background-task wait is involved. + """ + with smart_router_preset(cuj, SMART_ROUTER_CONFIG_VERSION) as (session, _, recorder, _): + first_prompt_routed = _preset_flags( + SMART_ROUTER_CONFIG_VERSION, CLAUDE + ).first_prompt_routed + evidence = SessionEvidence(session.home, CLAUDE) + task = FileTask(session) + task.prompt += " Do not delegate." + recorder.prepare_launch() + with Terminal(session, f"{SMART_ROUTER_CONFIG_VERSION}-claude-band", [CLAUDE]) as tui: + tui.boot(timeout=150) + tui.wait_for( + lambda _: _smart_router_band_shown( + tui, enabled=True, first_prompt=first_prompt_routed + ), + "the ON Smart Router band with the preset's routing scope", + ) + tui.wait_for( + lambda _: _smart_router_hint_shown(tui, enabled=True), + "smart router ON in the prompt hint row", + ) + tui.submit(task.prompt) + tui.task(evidence, task) + tui.wait_for( + lambda _: ( + _smart_router_band_hidden(tui) + and _smart_router_hint_shown(tui, enabled=True) + ), + "the Smart Router band hidden after the launch turn, with " + "smart router ON left in the prompt hint row", + ) + for enabled in (False, True): + tui.submit(f"/smart-router {'on' if enabled else 'off'}") + tui.wait_for( + lambda _, enabled=enabled: _smart_router_band_shown( + tui, enabled=enabled, first_prompt=first_prompt_routed ), - "completed native subagent file task and parent reply", + f"the {'ON' if enabled else 'OFF'} Smart Router band", ) - children = { - path: records - for path, records in agent_sessions(session, agent).items() - if path not in existing_sessions and is_child_session(agent, path, records) - } - completed_children = { - path: records - for path, records in children.items() - if any( - child_task.value in answer - for answer in assistant_answers(agent, records) - ) - } - assert len(completed_children) == 1, (agent, children.keys()) - decisions = read_jsonl(decisions_path)[decision_count:] - decision_tasks = { - decision.get("task_name") - for decision in decisions - if decision.get("task_name") - } - requests = recorder.requests_after(checkpoint) - routes = [ - request - for request in requests - if request.path == ROUTING_PATH - and request.payload.get("task", {}).get("prompt") in decision_tasks - ] - assert len(routes) == int(subagent_routed), (agent, routes) - if subagent_routed: - assert_subagent_routed( - session, - agent, - child_task, - decision_ids={decision["decision_id"] for decision in decisions}, - ) - route_prompt = routes[0].payload["task"]["prompt"] - assert isinstance(route_prompt, str) and route_prompt.strip() - if agent == CLAUDE: - assert child_task.filename in route_prompt - response = recorder.response_for(routes[0]) - assert response.status_code == 200 - selections = response.payload["route_selection"] - assert len(selections) == 1 - inference = _task_inference_request( - requests, - agent, - route_prompt, - after=routes[0].sequence, - ) - inference = served_inference_request(recorder, requests, inference, agent) - assert canonical_model(inference.payload["model"]) == canonical_model( - selections[0]["route_option"]["model"] - ) - if agent == CLAUDE: - assert ( - routes[0].payload["route_selector"]["router_name"] == launch_recipe - ) - assert ( - inference.payload.get("smart_router_recipe_name") == launch_recipe - ) - assert "databricks-smart-router-recipe" not in inference.headers - else: - assert "databricks-smart-router-recipe" in inference.headers - else: - assert read_jsonl(decisions_path) == original_decisions - assert not any( - "databricks-smart-router-recipe" in request.headers - or "smart_router_recipe_name" in request.payload - or ORCHESTRATOR_CONTEXT.encode() in request.body - for request in requests - ), agent - if agent == CLAUDE: - tui.wait_for_background_tasks() - for enabled in (False, True): - tui.submit(f"/smart-router {'on' if enabled else 'off'}") - tui.wait_for( - lambda _, enabled=enabled, first_prompt=first_prompt_routed: ( - _smart_router_band_shown( - tui, enabled=enabled, first_prompt=first_prompt - ) - ), - f"the {'ON' if enabled else 'OFF'} Smart Router band", - ) - state = "ON" if enabled else "OFF" - # The hint row flips with the toggle itself, before any turn. - tui.wait_for( - lambda _, enabled=enabled: _smart_router_hint_shown( - tui, enabled=enabled - ), - f"smart router {state} in the prompt hint row after the toggle", - ) - # The band covers the toggle's first turn and is gone by the second; - # the hint row keeps the state. - follow_up = _file_task(session, agent) - tui.submit(follow_up.prompt) - tui.task(evidence, follow_up) - tui.wait_for( - lambda _, enabled=enabled: ( - _smart_router_band_hidden(tui) - and _smart_router_hint_shown(tui, enabled=enabled) - ), - f"the {state} Smart Router band hidden by the second turn, with " - f"smart router {state} left in the prompt hint row", - ) - tui.exit_normally() - finally: - if previous is None: - session.env.pop("SMART_ROUTER_CONFIG_VERSION", None) - else: - session.env["SMART_ROUTER_CONFIG_VERSION"] = previous - session.revert_machine_wide( - f"{SMART_ROUTER_CONFIG_VERSION}-revert", - "Smart-router preset cleanup left machine-wide agent settings", - ) + state = "ON" if enabled else "OFF" + # The hint row flips with the toggle itself, before any turn. + tui.wait_for( + lambda _, enabled=enabled: _smart_router_hint_shown(tui, enabled=enabled), + f"smart router {state} in the prompt hint row after the toggle", + ) + # The band covers the toggle's first turn and is gone by the second; + # the hint row keeps the state. + follow_up = _file_task(session, CLAUDE) + tui.submit(follow_up.prompt) + tui.task(evidence, follow_up) + tui.wait_for( + lambda _, enabled=enabled: ( + _smart_router_band_hidden(tui) + and _smart_router_hint_shown(tui, enabled=enabled) + ), + f"the {state} Smart Router band hidden by the second turn, with " + f"smart router {state} left in the prompt hint row", + ) + tui.exit_normally() @pytest.mark.parametrize("agent", AGENTS) def test_agent_completes_real_first_prompt_file_task_without_model_override( diff --git a/tests/integration/README.md b/tests/integration/README.md index 498eef757..54e671721 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -70,9 +70,10 @@ skill name as data; skill download and invocation are not covered. The dedicated smart-routing CUJ in `../e2e_cuj/test_cuj4_smart_routing.py` runs the original managed-default and explicit-model cases, with explicit-model bypass assertions enabled for -both Claude and Codex. One parameterized test runs all six -`SMART_ROUTER_CONFIG_VERSION` values, checking both agents' first-prompt routing, orchestrator -context in inference input, and completed explicitly requested native subagents. +both Claude and Codex. Two parameterized tests run all six +`SMART_ROUTER_CONFIG_VERSION` values: one per agent checking first-prompt routing, orchestrator +context in inference input, and a completed explicitly requested native subagent, and a +Claude-only one checking the status band and hint row at launch and across off/on toggles. This is not automatic orchestrator-delegation coverage. No workspace configuration is modified. CUJ4 matches both Codex 0.154.0's top-level tools and 0.162.1's Responses Lite `additional_tools` input items. Encrypted child messages match the routed task name @@ -124,7 +125,7 @@ or construct ug state files. The normal test suite checks these boundaries. The Claude mod entry point has separate native component checks in `../native_claude/test_mod.py`. Claude 2.1.290 validates the module and runs its TypeScript event-forwarding and launch-option tests in a dedicated CI job. The -existing Claude skill-toggle journey checks the live recipe payload field across on/off/on. +existing Claude skill-toggle journeys check the live recipe payload field (launch recipe at launch and after off/on, DISABLED after off). The required Isaac feedback form, evidence submission, and background delivery notifications need a separate manual Isaac check. The gateway toggle journey does not cover the form or notification timing. Use the released Isaac skill with the @@ -140,7 +141,7 @@ existing feedback skill without loading the router plugin. The existing unit tests keep their fixtures. Integration has an independent pytest configuration and uses `--confcutdir` so those fixtures cannot leak in. It is not collected by the default `uv run pytest` command. -The Claude subagent skill-toggle journey opens `/tasks` after its final calculation +The Claude subagent skill-toggle journeys open `/tasks` after its final calculation and waits up to 180 seconds for "No tasks currently running" or a native menu containing only completed task rows before submitting `/exit`. Claude 2.1.290 retains finished children in a `Completed (N)` list, so the empty-state message is not required when the list has exactly N @@ -149,7 +150,7 @@ after Escape before typing `/exit`. Offline PTY checks reject active/scheduled t lists, and completed-task text outside the native menu. It neither stops tasks nor confirms an exit dialog. Process exit retains its 30-second timeout. Offline PTY checks cover this ordering. -The shared wait also runs before CUJ4's Claude preset sessions exit, after delegated-task +The shared wait also runs before CUJ4's Claude delegation sessions exit, after delegated-task evidence is checked, because Claude can resume a child after its first completed answer. `../test_startup_imports.py` covers the bare version fast path, deferred imports, @@ -1036,7 +1037,9 @@ existing e2e workspace/bearer. Before marking another agent's cases or tests req different workspace/setup, extend the advisory lanes to run them. The marker alone does not provide that setup or affect unit tests, installation checks, or dedicated CUJs. -Only `test_smart_router_skill_toggles_claude_subagent_routing` and its six variants carry +Only the three Claude Smart Router toggle tests (`test_smart_router_routes_claude_subagent_at_launch`, +`test_smart_router_skill_turns_claude_subagent_routing_off` and +`test_smart_router_skill_turns_claude_subagent_routing_back_on`) and their six variants each carry `soft_fail`. The **Soft-fail journeys · claude · Linux/Windows (non-blocking)** jobs run those cases unchanged, with the pinned Claude version and existing e2e workspace/bearer. They run after Full and Managed config finish, one OS at a time, to avoid adding concurrent inference load. @@ -1395,7 +1398,7 @@ before work and after off/on toggles. Checks ignore terminal styling and trailin UI controls; offline cases reject stale states, wrong scope, and command confirmations. The plugin test host runs without a process capability, so the mod's session-control call -is covered by the live skill-toggle journey rather than `claude plugin test`; native +is covered by the live skill-toggle journeys rather than `claude plugin test`; native component tests stub the subagent routing process protocol. No component check asserts that the launch paths withdraw the skills. The TypeScript launch test also checks event/result forwarding. Separate cases cover @@ -1403,10 +1406,10 @@ denied-write rollback, invalid request bodies, disabled launches, and body-merge no simulated subprocess success is used to claim toggle coverage. `integration/test_ug_smart_routing_hooks.py` checks `smart_router_recipe_name` in its existing -Claude skill-toggle journey. The recorder forwards to the real workspace, configured +Claude skill-toggle journeys. The recorder forwards to the real workspace, configured through public `ug configure --workspace ` without hand-writing ug state. Checkpoints and each calculation's unique marker correlate parent/child inference -requests to the on/off/on phases. All matched tool-capable inference payloads must carry the +requests to the launch, off, and back-on phases. All matched tool-capable inference payloads must carry the launch recipe, `DISABLED`, then the launch recipe again, omit the recipe header, and reach non-empty HTTP 200 responses, with only the existing verified native compatibility retries accepted. The same requests must carry the orchestrator's Agent/Task guidance only while diff --git a/tests/integration/test_ug_claude_tracing.py b/tests/integration/test_ug_claude_tracing.py index c5588224d..b5b3a833f 100644 --- a/tests/integration/test_ug_claude_tracing.py +++ b/tests/integration/test_ug_claude_tracing.py @@ -7,7 +7,7 @@ from utils.constants import CLAUDE_TEST_MODEL from utils.evidence import FileTask, fresh_id from utils.managed import use_managed_config_fixture -from utils.sql import query_count, resolve_trace_table, resolve_warehouse_id +from utils.sql import resolve_trace_table, resolve_warehouse_id, wait_for_count pytestmark = [pytest.mark.live, pytest.mark.claude] @@ -56,14 +56,14 @@ def test_ug_claude_exports_trace_to_configured_table(live_session, workspace): "WHERE time > current_timestamp() - INTERVAL 10 MINUTES " "AND variant_get(resource.attributes, '$[\"ug_integration_marker\"]', 'STRING') = :marker" ) - count = query_count( + count = wait_for_count( workspace, bearer, warehouse_id, marker_query, [{"name": "marker", "value": marker, "type": "STRING"}], ) - model_count = query_count( + model_count = wait_for_count( workspace, bearer, warehouse_id, diff --git a/tests/integration/test_ug_codex_tracing.py b/tests/integration/test_ug_codex_tracing.py index 754029018..d1fcdc8fb 100644 --- a/tests/integration/test_ug_codex_tracing.py +++ b/tests/integration/test_ug_codex_tracing.py @@ -7,7 +7,7 @@ from utils.constants import CODEX_TEST_MODEL from utils.evidence import FileTask, fresh_id from utils.managed import use_managed_config_fixture -from utils.sql import query_count, resolve_trace_table, resolve_warehouse_id +from utils.sql import resolve_trace_table, resolve_warehouse_id, wait_for_count pytestmark = [pytest.mark.live, pytest.mark.codex] @@ -56,7 +56,7 @@ def test_ug_codex_exports_trace_to_configured_table(live_session, workspace): "WHERE time > current_timestamp() - INTERVAL 10 MINUTES " "AND variant_get(attributes, '$[\"ug_integration_marker\"]', 'STRING') = :marker" ) - count = query_count( + count = wait_for_count( workspace, bearer, warehouse_id, @@ -65,7 +65,7 @@ def test_ug_codex_exports_trace_to_configured_table(live_session, workspace): ) # The span must also carry Codex's `model` attribute matching the model that ran, # so the trace is attributable to a specific model and not just to this test run. - model_count = query_count( + model_count = wait_for_count( workspace, bearer, warehouse_id, diff --git a/tests/integration/test_ug_smart_routing_hooks.py b/tests/integration/test_ug_smart_routing_hooks.py index abf5a4be5..067e70476 100644 --- a/tests/integration/test_ug_smart_routing_hooks.py +++ b/tests/integration/test_ug_smart_routing_hooks.py @@ -385,11 +385,7 @@ def test_smart_routing_codex_route_subagent_hook( assert row["requested_model"] == updated["model"], row -@pytest.mark.live -@pytest.mark.claude -@pytest.mark.managed_fixture -@pytest.mark.soft_fail -@pytest.mark.parametrize( +SMART_ROUTER_VERSIONS = pytest.mark.parametrize( "orchestration_enabled, SMART_ROUTER_CONFIG_VERSION", [ (False, None), @@ -408,36 +404,130 @@ def test_smart_routing_codex_route_subagent_hook( "subagent_orch_v0", ], ) -def test_smart_router_skill_toggles_claude_subagent_routing( - live_session, workspace, tmp_path, orchestration_enabled, SMART_ROUTER_CONFIG_VERSION -): - """Scenario: launch Claude with the legacy flags, a subagent selector, or only the managed - config's smart routing default, spawn a child, invoke the installed Smart Router skill to - turn routing off, spawn another child, turn routing back on through the skill, and spawn a - third child in the same real TUI session. - Expected: routing-only cases install Smart Router; enabling orchestration also installs - Smart Router Orchestrator. Each local invocation confirms the new state in the TUI. - All three uniquely tagged calculations complete in native parent and child sessions; - their recorded inference payloads carry smart_router_recipe_name with the launch - recipe, DISABLED, then the launch recipe again, without a recipe header, and receive - non-empty HTTP 200 responses (including verified native - compatibility retries). Only the first and third show native model/rationale rows - and produce live gateway decisions correlated with those children. Claude's native - task view reports no running tasks before /exit. No first-prompt routing wrapper starts. - """ - session = live_session + +def _select_smart_router_version( + session, agent: str, workspace, tmp_path, orchestration_enabled, config_version +) -> None: + """Choose the legacy flags, a subagent selector, or the managed default; install the fixture.""" session.env["TMPDIR"] = str(tmp_path) - if SMART_ROUTER_CONFIG_VERSION is None: + if config_version is None: session.env["ENABLE_SMART_ROUTING_V2"] = "1" session.env["ENABLE_SMART_ROUTING_SUBAGENT_ONLY"] = "1" if orchestration_enabled: session.env["ENABLE_SMART_ROUTER_ORCHESTRATOR"] = "1" - elif SMART_ROUTER_CONFIG_VERSION != MANAGED_DEFAULT: - session.env["SMART_ROUTER_CONFIG_VERSION"] = SMART_ROUTER_CONFIG_VERSION + elif config_version != MANAGED_DEFAULT: + session.env["SMART_ROUTER_CONFIG_VERSION"] = config_version # Managed smart routing would override the legacy env vars set above. - use_smart_routing_fixture( - session, "claude", workspace, smart_routing=SMART_ROUTER_CONFIG_VERSION is not None + use_smart_routing_fixture(session, agent, workspace, smart_routing=config_version is not None) + + +def _assert_claude_calculation_requests( + recorder, + checkpoint, + task, + *, + routed: bool, + launch_recipe: str | None, + orchestration_enabled: bool, +) -> str: + """Assert the recorded gateway traffic for one Claude calculation; return its recipe.""" + requests = recorder.requests_after(checkpoint) + inferences = [ + request + for request in requests + if request.method == "POST" + and request.path == "/ai-gateway/anthropic/v1/messages" + and request.payload.get("tools") + and task.marker in json.dumps(request.payload.get("messages", [])) + ] + assert inferences, f"No recorded inference for {task.marker}" + if routed: + routes = [ + request + for request in requests + if request.method == "POST" + and request.path == "/ai-gateway/routing/v1/routes:select" + and task.marker in request.payload.get("task", {}).get("prompt", "") + ] + assert len(routes) == 1, (task.marker, routes) + recipe = routes[0].payload.get("route_selector", {}).get("router_name") + assert isinstance(recipe, str) and recipe.strip() and recipe != "DISABLED", ( + task.marker, + routes[0].sequence, + routes[0].payload.get("route_selector"), + ) + if launch_recipe is not None: + assert recipe == launch_recipe, (task.marker, recipe, launch_recipe) + else: + recipe = "DISABLED" + for request in inferences: + for tool in request.payload["tools"]: + if tool.get("name") not in {"Agent", "Task"}: + continue + description = tool["description"] + if orchestration_enabled and routed: + assert "before gathering source evidence" in description, description + assert "notification in a later turn" in description + assert "end the current turn and wait for that notification" in description + assert "To continue a previously spawned agent, use SendMessage" in description + assert "starts a fresh agent with no memory of prior runs" in description + assert "When in doubt, don't spawn" not in description, description + assert "A fresh agent costs more than it looks" not in description, description + else: + assert "under the active Smart Router Orchestrator workflow" not in description + assert request.payload.get("smart_router_recipe_name") == recipe, ( + task.marker, + request.sequence, + request.payload.get("smart_router_recipe_name"), + recipe, + ) + assert "databricks-smart-router-recipe" not in request.headers, ( + task.marker, + request.sequence, + ) + served_inference_request(recorder, requests, request, "claude") + return recipe + + +def _assert_claude_session_finished(session, tui) -> str: + """Drain native tasks, exit the TUI, and return its transcript.""" + tui.wait_for_background_tasks() + tui.exit_normally() + return "".join(tui.output) + + +def _assert_claude_not_wrapped(session, transcript: str) -> None: + assert SMART_ROUTING_BANNER not in transcript, transcript + session.assert_not_routed() + 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.claude +@pytest.mark.managed_fixture +@pytest.mark.soft_fail +@SMART_ROUTER_VERSIONS +def test_smart_router_routes_claude_subagent_at_launch( + live_session, workspace, tmp_path, orchestration_enabled, SMART_ROUTER_CONFIG_VERSION +): + """Scenario: launch Claude with the legacy flags, a subagent selector, or only the managed + config's smart routing default and, without touching the Smart Router skill, spawn one child in + the real TUI. + + Expected: routing-only cases install Smart Router; enabling orchestration also installs + Smart Router Orchestrator. The uniquely tagged calculation completes in native parent and + child sessions; its recorded inference payloads carry smart_router_recipe_name with the + launch recipe, without a recipe header, and receive non-empty HTTP 200 responses + (including verified native compatibility retries). The child shows the native + model/rationale row and produces one live gateway decision correlated with it. Claude's + native task view reports no running tasks before /exit. No first-prompt routing wrapper + starts. + """ + session = live_session + _select_smart_router_version( + session, "claude", workspace, tmp_path, orchestration_enabled, SMART_ROUTER_CONFIG_VERSION ) with TuiRequestRecorder(workspace) as recorder: # Configure through the forwarding URL using the public CLI, keeping ug state real. @@ -450,161 +540,275 @@ def test_smart_router_skill_toggles_claude_subagent_routing( "--disable-databricks-ai-tools", ) with AgentTerminal( - session, "claude", [str(session.binary), "claude"], "smart-router-skill-toggle" + session, "claude", [str(session.binary), "claude"], "smart-router-launch" ) as tui: tui.boot() - launch_recipe = None - for expression, result, routed in [ - ("1+1", "2", True), - ("1+2", "3", False), - ("2+2", "4", True), - ]: - if expression != "1+1": - _toggle_with_skill( - tui, - session, - "claude", - enabled=routed, - orchestration_enabled=orchestration_enabled, - ) - checkpoint = recorder.checkpoint() - task = _run_calculation( - tui, - session, - "claude", - expression, - result, - orchestration_enabled if routed else False, - routed=routed, - ) - requests = recorder.requests_after(checkpoint) - inferences = [ - request - for request in requests - if request.method == "POST" - and request.path == "/ai-gateway/anthropic/v1/messages" - and request.payload.get("tools") - and task.marker in json.dumps(request.payload.get("messages", [])) - ] - assert inferences, f"No recorded inference for {task.marker}" - if routed: - routes = [ - request - for request in requests - if request.method == "POST" - and request.path == "/ai-gateway/routing/v1/routes:select" - and task.marker in request.payload.get("task", {}).get("prompt", "") - ] - assert len(routes) == 1, (task.marker, routes) - recipe = routes[0].payload.get("route_selector", {}).get("router_name") - assert isinstance(recipe, str) and recipe.strip() and recipe != "DISABLED", ( - task.marker, - routes[0].sequence, - routes[0].payload.get("route_selector"), - ) - if launch_recipe is None: - launch_recipe = recipe - else: - assert recipe == launch_recipe, (task.marker, recipe, launch_recipe) - else: - recipe = "DISABLED" - for request in inferences: - for tool in request.payload["tools"]: - if tool.get("name") not in {"Agent", "Task"}: - continue - description = tool["description"] - if orchestration_enabled and routed: - assert "before gathering source evidence" in description, description - assert "notification in a later turn" in description - assert ( - "end the current turn and wait for that notification" in description - ) - assert ( - "To continue a previously spawned agent, use SendMessage" - in description - ) - assert ( - "starts a fresh agent with no memory of prior runs" in description - ) - assert "When in doubt, don't spawn" not in description, description - assert "A fresh agent costs more than it looks" not in description, ( - description - ) - else: - assert ( - "under the active Smart Router Orchestrator workflow" - not in description - ) - assert request.payload.get("smart_router_recipe_name") == recipe, ( - task.marker, - request.sequence, - request.payload.get("smart_router_recipe_name"), - recipe, - ) - assert "databricks-smart-router-recipe" not in request.headers, ( - task.marker, - request.sequence, - ) - served_inference_request(recorder, requests, request, "claude") - tui.wait_for_background_tasks() - tui.exit_normally() - transcript = "".join(tui.output) + checkpoint = recorder.checkpoint() + task = _run_calculation( + tui, session, "claude", "1+1", "2", orchestration_enabled, routed=True + ) + _assert_claude_calculation_requests( + recorder, + checkpoint, + task, + routed=True, + launch_recipe=None, + orchestration_enabled=orchestration_enabled, + ) + transcript = _assert_claude_session_finished(session, tui) + _assert_claude_not_wrapped(session, transcript) + + +@pytest.mark.live +@pytest.mark.claude +@pytest.mark.managed_fixture +@pytest.mark.soft_fail +@SMART_ROUTER_VERSIONS +def test_smart_router_skill_turns_claude_subagent_routing_off( + live_session, workspace, tmp_path, orchestration_enabled, SMART_ROUTER_CONFIG_VERSION +): + """Scenario: launch Claude with the legacy flags, a subagent selector, or only the managed + config's smart routing default, invoke the installed Smart Router skill to turn routing off, + then spawn a child in the same real TUI session. + + Expected: routing-only cases install Smart Router; enabling orchestration also installs + Smart Router Orchestrator. The local invocation confirms "Smart Router is off" in the TUI. + The uniquely tagged calculation completes in native parent and child sessions; its recorded + inference payloads carry smart_router_recipe_name DISABLED, without a recipe header, and + receive non-empty HTTP 200 responses (including verified native compatibility retries). + No native model/rationale row is shown, no route is selected, and no live gateway decision + is recorded. Claude's native task view reports no running tasks before /exit. No + first-prompt routing wrapper starts. + """ + session = live_session + _select_smart_router_version( + session, "claude", workspace, tmp_path, orchestration_enabled, SMART_ROUTER_CONFIG_VERSION + ) + with TuiRequestRecorder(workspace) as recorder: + session.run( + "configure", + "--workspace", + recorder.url, + "--skip-validate", + "--skip-upgrade", + "--disable-databricks-ai-tools", + ) + with AgentTerminal( + session, "claude", [str(session.binary), "claude"], "smart-router-skill-off" + ) as tui: + tui.boot() + _toggle_with_skill( + tui, session, "claude", enabled=False, orchestration_enabled=orchestration_enabled + ) + checkpoint = recorder.checkpoint() + task = _run_calculation(tui, session, "claude", "1+2", "3", False, routed=False) + recipe = _assert_claude_calculation_requests( + recorder, + checkpoint, + task, + routed=False, + launch_recipe=None, + orchestration_enabled=orchestration_enabled, + ) + assert recipe == "DISABLED", recipe + transcript = _assert_claude_session_finished(session, tui) + _assert_claude_not_wrapped(session, transcript) + + +@pytest.mark.live +@pytest.mark.claude +@pytest.mark.managed_fixture +@pytest.mark.soft_fail +@SMART_ROUTER_VERSIONS +def test_smart_router_skill_turns_claude_subagent_routing_back_on( + live_session, workspace, tmp_path, orchestration_enabled, SMART_ROUTER_CONFIG_VERSION +): + """Scenario: launch Claude with the legacy flags, a subagent selector, or only the managed + config's smart routing default, spawn a child, invoke the installed Smart Router skill to turn + routing off, spawn a second child, turn routing back on, and spawn a third child in the same + real TUI session. + + Expected: routing-only cases install Smart Router; enabling orchestration also installs + Smart Router Orchestrator. Each local invocation confirms the new state in the TUI. All + uniquely tagged calculations complete in native parent and child sessions. The second + child's recorded inference payloads carry smart_router_recipe_name DISABLED, with no + native model/rationale row and no live gateway decision. The first and third children show + native model/rationale rows with live gateway decisions correlated with them; the third is + routed again with the same recipe as the first, and both carry that recipe. All payloads + omit a recipe header and receive non-empty HTTP 200 responses (including verified native + compatibility retries). Claude's native task view reports no running tasks before /exit. + No first-prompt routing wrapper starts. + """ + session = live_session + _select_smart_router_version( + session, "claude", workspace, tmp_path, orchestration_enabled, SMART_ROUTER_CONFIG_VERSION + ) + with TuiRequestRecorder(workspace) as recorder: + session.run( + "configure", + "--workspace", + recorder.url, + "--skip-validate", + "--skip-upgrade", + "--disable-databricks-ai-tools", + ) + with AgentTerminal( + session, "claude", [str(session.binary), "claude"], "smart-router-skill-back-on" + ) as tui: + tui.boot() + checkpoint = recorder.checkpoint() + launch_task = _run_calculation( + tui, session, "claude", "1+1", "2", orchestration_enabled, routed=True + ) + launch_recipe = _assert_claude_calculation_requests( + recorder, + checkpoint, + launch_task, + routed=True, + launch_recipe=None, + orchestration_enabled=orchestration_enabled, + ) + _toggle_with_skill( + tui, session, "claude", enabled=False, orchestration_enabled=orchestration_enabled + ) + # Prove "off" takes effect after a routed child before turning routing back on. + checkpoint = recorder.checkpoint() + off_task = _run_calculation(tui, session, "claude", "1+2", "3", False, routed=False) + recipe = _assert_claude_calculation_requests( + recorder, + checkpoint, + off_task, + routed=False, + launch_recipe=None, + orchestration_enabled=orchestration_enabled, + ) + assert recipe == "DISABLED", recipe + _toggle_with_skill( + tui, session, "claude", enabled=True, orchestration_enabled=orchestration_enabled + ) + checkpoint = recorder.checkpoint() + task = _run_calculation( + tui, session, "claude", "2+2", "4", orchestration_enabled, routed=True + ) + _assert_claude_calculation_requests( + recorder, + checkpoint, + task, + routed=True, + launch_recipe=launch_recipe, + orchestration_enabled=orchestration_enabled, + ) + transcript = _assert_claude_session_finished(session, tui) + _assert_claude_not_wrapped(session, transcript) + + +@pytest.mark.live +@pytest.mark.codex +@pytest.mark.managed_fixture +@SMART_ROUTER_VERSIONS +def test_smart_router_routes_codex_subagent_at_launch( + live_session, workspace, tmp_path, orchestration_enabled, SMART_ROUTER_CONFIG_VERSION +): + """Scenario: launch Codex with the legacy flags, a subagent selector, or only the managed + config's smart routing default and, without touching the Smart Router skill, spawn one child in + the real TUI. + + Expected: routing-only cases install Smart Router; enabling orchestration also installs + Smart Router Orchestrator. The uniquely tagged calculation completes in a native child + session, shows the subagent-routing banner (including [orchestrator on] when orchestration + is enabled), and produces one live gateway decision correlated with that child. No + first-prompt interposer starts. + """ + session = live_session + _select_smart_router_version( + session, "codex", workspace, tmp_path, orchestration_enabled, SMART_ROUTER_CONFIG_VERSION + ) + session.run( + "configure", + "--workspace", + workspace, + "--skip-validate", + "--skip-upgrade", + "--disable-databricks-ai-tools", + ) + with AgentTerminal( + session, "codex", [str(session.binary), "codex"], "smart-router-launch" + ) as tui: + tui.boot() + _run_calculation(tui, session, "codex", "1+1", "2", orchestration_enabled, routed=True) + tui.exit_normally() + transcript = "".join(tui.output) assert SMART_ROUTING_BANNER not in transcript, transcript session.assert_not_routed() - 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 @pytest.mark.managed_fixture -@pytest.mark.parametrize( - "orchestration_enabled, SMART_ROUTER_CONFIG_VERSION", - [ - (False, None), - (True, None), - (True, MANAGED_DEFAULT), - (False, "subagent_only_v0"), - (False, "subagent_only_v1"), - (True, "subagent_orch_v0"), - ], - ids=[ - "routing-only", - "orchestration", - MANAGED_DEFAULT, - "subagent_only_v0", - "subagent_only_v1", - "subagent_orch_v0", - ], -) -def test_smart_router_skill_toggles_codex_subagent_routing( +@SMART_ROUTER_VERSIONS +def test_smart_router_skill_turns_codex_subagent_routing_off( live_session, workspace, tmp_path, orchestration_enabled, SMART_ROUTER_CONFIG_VERSION ): """Scenario: launch Codex with the legacy flags, a subagent selector, or only the managed - config's smart routing default, spawn a child, invoke the installed Smart Router skill to - turn routing off, spawn another child, turn routing back on through the skill, and spawn a - third child in the same real TUI session. + config's smart routing default, invoke the installed Smart Router skill to turn routing off, + then spawn a child in the same real TUI session. Expected: routing-only cases install Smart Router; enabling orchestration also installs - Smart Router Orchestrator. Each invocation records the CLI - confirmation in the native transcript and changes the saved routing controls, even with - collapsed terminal output; all three uniquely tagged calculations complete in native child - sessions; only the first and third show the subagent-routing banner and produce live gateway - decisions correlated with those children. No first-prompt interposer starts. - With orchestration enabled, both routed banners must include [orchestrator on]. + Smart Router Orchestrator. The invocation records the CLI confirmation in the native + transcript and changes the saved routing controls, even with collapsed terminal output. + The uniquely tagged calculation completes in a native child session, shows no + subagent-routing banner, and produces no live gateway decision. No first-prompt + interposer starts. """ session = live_session - session.env["TMPDIR"] = str(tmp_path) - if SMART_ROUTER_CONFIG_VERSION is None: - session.env["ENABLE_SMART_ROUTING_V2"] = "1" - session.env["ENABLE_SMART_ROUTING_SUBAGENT_ONLY"] = "1" - if orchestration_enabled: - session.env["ENABLE_SMART_ROUTER_ORCHESTRATOR"] = "1" - elif SMART_ROUTER_CONFIG_VERSION != MANAGED_DEFAULT: - session.env["SMART_ROUTER_CONFIG_VERSION"] = SMART_ROUTER_CONFIG_VERSION - # Managed smart routing would override the legacy env vars set above. - use_smart_routing_fixture( - session, "codex", workspace, smart_routing=SMART_ROUTER_CONFIG_VERSION is not None + _select_smart_router_version( + session, "codex", workspace, tmp_path, orchestration_enabled, SMART_ROUTER_CONFIG_VERSION + ) + session.run( + "configure", + "--workspace", + workspace, + "--skip-validate", + "--skip-upgrade", + "--disable-databricks-ai-tools", + ) + with AgentTerminal( + session, "codex", [str(session.binary), "codex"], "smart-router-skill-off" + ) as tui: + tui.boot() + _toggle_with_skill( + tui, session, "codex", enabled=False, orchestration_enabled=orchestration_enabled + ) + _run_calculation(tui, session, "codex", "1+2", "3", False, routed=False) + tui.exit_normally() + transcript = "".join(tui.output) + assert SMART_ROUTING_BANNER not in transcript, transcript + session.assert_not_routed() + + +@pytest.mark.live +@pytest.mark.codex +@pytest.mark.managed_fixture +@SMART_ROUTER_VERSIONS +def test_smart_router_skill_turns_codex_subagent_routing_back_on( + live_session, workspace, tmp_path, orchestration_enabled, SMART_ROUTER_CONFIG_VERSION +): + """Scenario: launch Codex with the legacy flags, a subagent selector, or only the managed + config's smart routing default, spawn a child, invoke the installed Smart Router skill to turn + routing off, spawn a second child, turn routing back on, and spawn a third child in the same + real TUI session. + + Expected: routing-only cases install Smart Router; enabling orchestration also installs + Smart Router Orchestrator. Each invocation records the CLI confirmation in the native + transcript and changes the saved routing controls (off, then restored), even with collapsed + terminal output. Each uniquely tagged calculation completes in a native child session. The + first and third show the subagent-routing banner (including [orchestrator on] when + orchestration is enabled) and produce one live gateway decision correlated with that child; + the second shows no banner and produces no decision. No first-prompt interposer starts. + """ + session = live_session + _select_smart_router_version( + session, "codex", workspace, tmp_path, orchestration_enabled, SMART_ROUTER_CONFIG_VERSION ) session.run( "configure", @@ -615,24 +819,17 @@ def test_smart_router_skill_toggles_codex_subagent_routing( "--disable-databricks-ai-tools", ) with AgentTerminal( - session, "codex", [str(session.binary), "codex"], "smart-router-skill-toggle" + session, "codex", [str(session.binary), "codex"], "smart-router-skill-back-on" ) as tui: tui.boot() _run_calculation(tui, session, "codex", "1+1", "2", orchestration_enabled, routed=True) _toggle_with_skill( - tui, - session, - "codex", - enabled=False, - orchestration_enabled=orchestration_enabled, + tui, session, "codex", enabled=False, orchestration_enabled=orchestration_enabled ) + # Prove "off" takes effect after a routed child before turning routing back on. _run_calculation(tui, session, "codex", "1+2", "3", False, routed=False) _toggle_with_skill( - tui, - session, - "codex", - enabled=True, - orchestration_enabled=orchestration_enabled, + tui, session, "codex", enabled=True, orchestration_enabled=orchestration_enabled ) _run_calculation(tui, session, "codex", "2+2", "4", orchestration_enabled, routed=True) tui.exit_normally() diff --git a/tests/integration/utils/request_recorder.py b/tests/integration/utils/request_recorder.py index e2654f58d..698b098d3 100644 --- a/tests/integration/utils/request_recorder.py +++ b/tests/integration/utils/request_recorder.py @@ -33,6 +33,8 @@ "x-api-key", "x-databricks-ai-gateway-token", } +# How long teardown waits for in-flight forwards before leaving their connections to them. +_DRAIN_SECONDS = 10 @dataclass(frozen=True) @@ -76,6 +78,7 @@ def __init__(self, upstream: str) -> None: self._condition = threading.Condition() self._requests: list[RecordedRequest] = [] self._responses: dict[int, RecordedResponse] = {} + self._in_flight = 0 self._client: httpx.Client | None = None self._server: _Server | None = None self._thread: threading.Thread | None = None @@ -87,7 +90,14 @@ def url(self) -> str: return f"http://127.0.0.1:{self._server.server_address[1]}" def __enter__(self) -> Self: - self._client = httpx.Client(timeout=None, follow_redirects=False) + # One connection per forwarded request: a pooled keep-alive connection that the gateway + # or a load balancer closed while idle fails on reuse with "Server disconnected without + # sending a response". The request is never resent; this only avoids reuse. + self._client = httpx.Client( + timeout=None, + follow_redirects=False, + limits=httpx.Limits(max_keepalive_connections=0), + ) recorder = self class Handler(BaseHTTPRequestHandler): @@ -95,7 +105,14 @@ def log_message(self, format: str, *args: object) -> None: pass def forward(self) -> None: - recorder._forward(self) + with recorder._condition: + recorder._in_flight += 1 + try: + recorder._forward(self) + finally: + with recorder._condition: + recorder._in_flight -= 1 + recorder._condition.notify_all() do_GET = forward do_POST = forward @@ -115,7 +132,15 @@ def __exit__( self._server.shutdown() self._server.server_close() self._thread.join() - self._client.close() + # Handler threads are daemons, so a forward can still be streaming here. Closing its + # socket from this thread frees the descriptor while OpenSSL keeps reading the old + # number, which a later connection in this process may reuse: that thread then reads + # (and corrupts) the new connection's bytes. Close only once forwards have drained; + # a forward that outlives the wait closes its own connection when it finishes. + with self._condition: + drained = self._condition.wait_for(lambda: self._in_flight == 0, timeout=_DRAIN_SECONDS) + if drained: + self._client.close() def checkpoint(self) -> int: with self._condition: @@ -220,7 +245,19 @@ def _forward(self, handler: BaseHTTPRequestHandler) -> None: body=bytes(response_body), ) self._condition.notify_all() - except httpx.HTTPError: + except httpx.HTTPError as error: + # Record the transport failure so a waiter sees the 502 the agent saw + # instead of timing out for a response that can never arrive. + with self._condition: + self._responses.setdefault( + request.sequence, + RecordedResponse( + status_code=502, + headers={}, + body=f"upstream request failed: {type(error).__name__}: {error}".encode(), + ), + ) + self._condition.notify_all() handler.send_error(502, "upstream request failed") diff --git a/tests/integration/utils/sql.py b/tests/integration/utils/sql.py index 03feb222e..8efefe274 100644 --- a/tests/integration/utils/sql.py +++ b/tests/integration/utils/sql.py @@ -6,6 +6,13 @@ import time import urllib.request +# Statement Execution API returns after at most 50s (`wait_timeout`); a cold warehouse can +# take minutes to start, so a statement is then polled until this overall budget. +STATEMENT_DEADLINE_SECONDS = 300 +# Spans reach the table asynchronously after the agent exits; poll for the row this long. +TRACE_ROW_DEADLINE_SECONDS = 300 +TRACE_ROW_POLL_INTERVAL_SECONDS = 15 + DEFAULT_TRACE_TABLE_PREFIX = "unity_gateway" TRACE_TABLE_SUFFIX = "_otel_spans" @@ -71,7 +78,10 @@ def query_count( warehouse_id: str, statement: str, parameters: list[dict[str, str]], + *, + timeout: float = STATEMENT_DEADLINE_SECONDS, ) -> int: + started = time.monotonic() request = urllib.request.Request( f"{workspace.rstrip('/')}/api/2.0/sql/statements", data=json.dumps( @@ -88,9 +98,13 @@ def query_count( with urllib.request.urlopen(request, timeout=60) as response: # noqa: S310 result = json.load(response) - deadline = time.monotonic() + 180 + deadline = started + timeout while result.get("status", {}).get("state") in {"PENDING", "RUNNING"}: - assert time.monotonic() < deadline, "SQL statement did not finish within 180 seconds" + assert time.monotonic() < deadline, ( + f"SQL statement did not finish within {timeout:g} seconds " + f"(elapsed {time.monotonic() - started:.0f}s, " + f"state {result.get('status', {}).get('state')})" + ) time.sleep(2) poll = urllib.request.Request( f"{workspace.rstrip('/')}/api/2.0/sql/statements/{result['statement_id']}", @@ -103,3 +117,29 @@ def query_count( rows = result.get("result", {}).get("data_array", []) assert rows and rows[0], "SQL count query returned no row" return int(rows[0][0]) + + +def wait_for_count( + workspace: str, + bearer: str, + warehouse_id: str, + statement: str, + parameters: list[dict[str, str]], + *, + timeout: float = TRACE_ROW_DEADLINE_SECONDS, + interval: float = TRACE_ROW_POLL_INTERVAL_SECONDS, +) -> int: + """Poll a COUNT query until it is positive or `timeout` elapses; return the last count. + + Callers assert on the result, so a count that never appears still fails honestly. + """ + started = time.monotonic() + deadline = started + timeout + while True: + remaining = deadline - time.monotonic() + count = query_count( + workspace, bearer, warehouse_id, statement, parameters, timeout=max(remaining, 60) + ) + if count > 0 or time.monotonic() + interval >= deadline: + return count + time.sleep(interval) diff --git a/tests/test_e2e.py b/tests/test_e2e.py index 03353873f..87eac3ff3 100644 --- a/tests/test_e2e.py +++ b/tests/test_e2e.py @@ -11,6 +11,7 @@ from __future__ import annotations import atexit +import dataclasses import json import os import secrets @@ -26,6 +27,7 @@ import httpx import pytest +from tests.conftest import e2e_state_for from ucode.agents import resolve_provider_models from ucode.databricks import ( build_shared_base_urls, @@ -465,26 +467,123 @@ def _require_binary(binary: str): pytest.skip(f"`{binary}` is not installed") -class TestCodexLaunch: - """Run codex against every available codex model.""" +# --------------------------------------------------------------------------- +# Per-model parametrization +# --------------------------------------------------------------------------- +# The live catalog is shared by every PR, so one misbehaving model must fail only its own case. +# Each ``*_per_model`` launch test is parametrized over its catalog at collection time via +# ``pytest_generate_tests``; ``e2e_state_for`` caches the fetch so the fixture reuses it. + + +@dataclasses.dataclass(frozen=True) +class _CatalogFetchError: + """Placeholder case for a catalog that could not be fetched; the test body fails with it.""" + + message: str + + +def _codex_launch_models(state: dict) -> list[str]: + return [m for m in (state.get("codex_models") or []) if not _model_is_skipped(m, "codex")] + + +def _claude_launch_models(state: dict) -> list[tuple[str, str]]: + return _launchable_model_items(state.get("claude_models") or {}) + + +def _gemini_launch_models(state: dict) -> list[str]: + models: list = state.get("gemini_models") or [] + # Image-only models can't serve Gemini CLI's tool calls. + if os.name == "nt": + models = [m for m in models if not m.endswith("-image")] + return models + + +def _opencode_launch_models(state: dict) -> list[tuple[str, str]]: + return TestOpencodeLaunch()._all_models(state) + + +def _copilot_launch_models(state: dict) -> list[tuple[str, str]]: + from ucode.agents import copilot + + # Models Copilot has no wire API for are left out of the catalog, as the launch always did. + return [ + (family, model) + for family, model in TestCopilotLaunch()._all_models(state) + if copilot.build_copilot_catalog([model]) + ] + + +def _pi_launch_models(state: dict) -> list[tuple[str, str]]: + return TestPiLaunch()._all_models(state) + + +# test function name -> (catalog -> model cases, "No ..." message when the catalog is empty) +_PER_MODEL_CASES = { + "test_launch_codex_per_model": (_codex_launch_models, "No Codex models available"), + "test_launch_claude_per_model": ( + _claude_launch_models, + "No launchable Claude models available", + ), + "test_launch_gemini_per_model": (_gemini_launch_models, "No Gemini models available"), + "test_launch_opencode_per_model": (_opencode_launch_models, "No OpenCode models available"), + "test_launch_copilot_per_model": ( + _copilot_launch_models, + "No Copilot-compatible models available", + ), + "test_launch_pi_per_model": (_pi_launch_models, "No Pi-compatible models available"), +} + + +def _case_id(case) -> str: + return "-".join(case) if isinstance(case, tuple) else case + - def _codex_models(self, e2e_state: dict) -> list[str]: - models = [ - model - for model in (e2e_state.get("codex_models") or []) - if not _model_is_skipped(model, "codex") +def _per_model_params(lister, empty_message: str) -> list: + workspace = _ws() + if not workspace: + return [ + pytest.param( + None, + id="no-workspace", + marks=pytest.mark.skip( + reason="Set UCODE_TEST_WORKSPACE=https://... to run E2E tests" + ), + ) ] - if not models: - pytest.skip("No Codex models available on this workspace") - return models + try: + cases = lister(e2e_state_for(workspace)) + except Exception as exc: # an auth/catalog outage must stay red, not vanish at collection + error = _CatalogFetchError(f"{type(exc).__name__}: {exc}") + return [pytest.param(error, id="catalog-fetch-error")] + if not cases: + return [pytest.param(None, id="no-models", marks=pytest.mark.skip(reason=empty_message))] + return [pytest.param(case, id=_case_id(case)) for case in cases] + + +def pytest_generate_tests(metafunc): + entry = _PER_MODEL_CASES.get(metafunc.function.__name__) + if entry is not None and "model_case" in metafunc.fixturenames: + metafunc.parametrize("model_case", _per_model_params(*entry)) + - def test_launch_codex_per_model(self, tmp_path, monkeypatch, e2e_state, e2e_workspace): - """Parametrized inline — iterates over all codex models and asserts each works.""" +def _fail_on_catalog_error(model_case) -> None: + if isinstance(model_case, _CatalogFetchError): + pytest.fail(f"Could not fetch the model catalog for this workspace: {model_case.message}") + + +class TestCodexLaunch: + """Run codex against every available codex model.""" + + def test_launch_codex_per_model( + self, model_case, tmp_path, monkeypatch, e2e_state, e2e_workspace + ): + """One case per codex model (see ``_PER_MODEL_CASES``).""" import ucode.config_io as config_io_mod from ucode.agents import codex + _fail_on_catalog_error(model_case) + model = model_case _require_binary("codex") - models = self._codex_models(e2e_state) monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) config_dir = _codex_home_outside_tmp() / ".codex" @@ -494,60 +593,51 @@ def test_launch_codex_per_model(self, tmp_path, monkeypatch, e2e_state, e2e_work monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", config_path) monkeypatch.setattr(codex, "CODEX_BACKUP_PATH", backup_path) - failures = [] timeout_seconds = int(os.environ.get("UCODE_E2E_AGENT_TIMEOUT", "60")) - for model in models: - state = {**e2e_state, "workspace": e2e_workspace} - with pytest.MonkeyPatch().context() as mp: - mp.setattr("ucode.state.save_state", lambda s: None) - codex.write_tool_config(state, model) - - cmd = codex.validate_cmd("codex") - # By default, ug uses the default model of the harness. - # Instead, pin Codex to use the specified model. - cmd[1:1] = ["--model", codex.codex_model_id(model)] - try: - result = _run_agent( - cmd, - env={**os.environ, "CODEX_HOME": str(config_dir)}, - timeout=timeout_seconds, - ) - except subprocess.TimeoutExpired: - failures.append(f"model={model} timed out after {timeout_seconds}s") - continue - - if ( - result.returncode != 0 - or not result.stdout.strip() - or f"model: {codex.codex_model_id(model)}\n" not in result.stderr - ): - # Keep a generous tail of stderr. codex-cli logs a non-fatal model-listing error - # first and the actual cause last, so a short prefix reports the wrong problem — - # at 200 chars the geography failure above read as a `/v1/models` routing error. - failures.append( - f"model={model} rc={result.returncode} " - f"stdout={result.stdout[-500:]!r} stderr={result.stderr[-1500:]!r}" - ) + state = {**e2e_state, "workspace": e2e_workspace} + with pytest.MonkeyPatch().context() as mp: + mp.setattr("ucode.state.save_state", lambda s: None) + codex.write_tool_config(state, model) - assert not failures, "Codex launch failures:\n" + "\n".join(failures) + cmd = codex.validate_cmd("codex") + # By default, ug uses the default model of the harness. + # Instead, pin Codex to use the specified model. + cmd[1:1] = ["--model", codex.codex_model_id(model)] + try: + result = _run_agent( + cmd, + env={**os.environ, "CODEX_HOME": str(config_dir)}, + timeout=timeout_seconds, + ) + except subprocess.TimeoutExpired: + pytest.fail(f"model={model} timed out after {timeout_seconds}s") + + if ( + result.returncode != 0 + or not result.stdout.strip() + or f"model: {codex.codex_model_id(model)}\n" not in result.stderr + ): + # Keep a generous tail of stderr. codex-cli logs a non-fatal model-listing error + # first and the actual cause last, so a short prefix reports the wrong problem — + # at 200 chars the geography failure above read as a `/v1/models` routing error. + pytest.fail( + f"Codex launch failure: model={model} rc={result.returncode} " + f"stdout={result.stdout[-500:]!r} stderr={result.stderr[-1500:]!r}" + ) class TestClaudeLaunch: """Run claude against every available claude model (sonnet, opus, haiku).""" def test_launch_claude_per_model( - self, tmp_path, monkeypatch, e2e_state, e2e_workspace, e2e_token + self, model_case, tmp_path, monkeypatch, e2e_state, e2e_workspace, e2e_token ): import ucode.config_io as config_io_mod from ucode.agents import claude + _fail_on_catalog_error(model_case) + family, model_id = model_case _require_binary("claude") - claude_models: dict = e2e_state.get("claude_models") or {} - if not claude_models: - pytest.skip("No Claude models available on this workspace") - launchable_models = _launchable_model_items(claude_models) - if not launchable_models: - pytest.skip("No launchable Claude models available on this workspace") # Use an isolated config dir so the claude subprocess never reads or # writes ~/.claude/settings.json during this test. @@ -559,29 +649,25 @@ def test_launch_claude_per_model( base_url = build_tool_base_url("claude", e2e_workspace) - failures = [] - for family, model_id in launchable_models: - with pytest.MonkeyPatch().context() as mp: - mp.setattr("ucode.state.save_state", lambda s: None) - claude.write_tool_config({**e2e_state, "workspace": e2e_workspace}, model_id) - - env = { - **os.environ, - "CLAUDE_CONFIG_DIR": str(config_dir), - "ANTHROPIC_MODEL": model_id, - "ANTHROPIC_BASE_URL": base_url, - "ANTHROPIC_API_KEY": e2e_token, - } - cmd = claude.validate_cmd("claude") - result = _run_agent(cmd, env=env, timeout=90) - combined = (result.stdout + result.stderr).strip() - if result.returncode != 0 or not combined: - failures.append( - f"family={family} model={model_id} rc={result.returncode} " - f"stdout={result.stdout[:300]!r} stderr={result.stderr[:300]!r}" - ) + with pytest.MonkeyPatch().context() as mp: + mp.setattr("ucode.state.save_state", lambda s: None) + claude.write_tool_config({**e2e_state, "workspace": e2e_workspace}, model_id) - assert not failures, "Claude launch failures:\n" + "\n".join(failures) + env = { + **os.environ, + "CLAUDE_CONFIG_DIR": str(config_dir), + "ANTHROPIC_MODEL": model_id, + "ANTHROPIC_BASE_URL": base_url, + "ANTHROPIC_API_KEY": e2e_token, + } + cmd = claude.validate_cmd("claude") + result = _run_agent(cmd, env=env, timeout=90) + combined = (result.stdout + result.stderr).strip() + if result.returncode != 0 or not combined: + pytest.fail( + f"Claude launch failure: family={family} model={model_id} rc={result.returncode} " + f"stdout={result.stdout[:300]!r} stderr={result.stderr[:300]!r}" + ) class TestModelProviderLaunch: @@ -765,11 +851,13 @@ class TestGeminiLaunch: """Run the real Gemini CLI against every available gemini model.""" def test_launch_gemini_per_model( - self, tmp_path, monkeypatch, e2e_state, e2e_workspace, e2e_token + self, model_case, tmp_path, monkeypatch, e2e_state, e2e_workspace, e2e_token ): import ucode.config_io as config_io_mod from ucode.agents import gemini + _fail_on_catalog_error(model_case) + model = model_case _require_binary("gemini") # Gemini CLI >= 0.45 rewrites forced flash model ids (e.g. # 'databricks-gemini-3-5-flash') to 'gemini-3.5-flash', which Unity @@ -782,12 +870,6 @@ def test_launch_gemini_per_model( f"Installed Gemini CLI {too_new} is past the supported ceiling " f"({gemini.MAX_GEMINI_VERSION_TEXT}); run `ucode gemini` to downgrade." ) - gemini_models: list = e2e_state.get("gemini_models") or [] - # Image-only models can't serve Gemini CLI's tool calls. - if os.name == "nt": - gemini_models = [m for m in gemini_models if not m.endswith("-image")] - if not gemini_models: - pytest.skip("No Gemini models available on this workspace") monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) monkeypatch.setattr(gemini, "GEMINI_ENV_PATH", tmp_path / "ucode.env") @@ -802,32 +884,27 @@ def test_launch_gemini_per_model( monkeypatch.chdir(tmp_path) timeout_seconds = int(os.environ.get("UCODE_E2E_AGENT_TIMEOUT", "60")) - failures = [] - for model in gemini_models: - with pytest.MonkeyPatch().context() as mp: - mp.setattr("ucode.state.save_state", lambda s: None) - mp.setattr( - "ucode.agents.gemini.get_databricks_token", - lambda ws, profile=None, **kwargs: e2e_token, - ) - state = {**e2e_state, "workspace": e2e_workspace} - gemini.write_tool_config(state, model, token=e2e_token) - # Launch the real Gemini CLI with the exact command ucode uses, - # against the just-written config, to keep actual CLI smoke coverage - # without bringing a validation probe back into `ucode configure`. - env = gemini.build_runtime_env(e2e_workspace, model, e2e_token) - try: - result = _run_agent(gemini.validate_cmd("gemini"), env=env, timeout=timeout_seconds) - except subprocess.TimeoutExpired: - failures.append(f"model={model} timed out after {timeout_seconds}s") - continue - if result.returncode != 0 or not result.stdout.strip(): - failures.append( - f"model={model} rc={result.returncode} " - f"stdout={result.stdout[-500:]!r} stderr={result.stderr[-1500:]!r}" - ) - - assert not failures, "Gemini launch failures:\n" + "\n".join(failures) + with pytest.MonkeyPatch().context() as mp: + mp.setattr("ucode.state.save_state", lambda s: None) + mp.setattr( + "ucode.agents.gemini.get_databricks_token", + lambda ws, profile=None, **kwargs: e2e_token, + ) + state = {**e2e_state, "workspace": e2e_workspace} + gemini.write_tool_config(state, model, token=e2e_token) + # Launch the real Gemini CLI with the exact command ucode uses, + # against the just-written config, to keep actual CLI smoke coverage + # without bringing a validation probe back into `ucode configure`. + env = gemini.build_runtime_env(e2e_workspace, model, e2e_token) + try: + result = _run_agent(gemini.validate_cmd("gemini"), env=env, timeout=timeout_seconds) + except subprocess.TimeoutExpired: + pytest.fail(f"model={model} timed out after {timeout_seconds}s") + if result.returncode != 0 or not result.stdout.strip(): + pytest.fail( + f"Gemini launch failure: model={model} rc={result.returncode} " + f"stdout={result.stdout[-500:]!r} stderr={result.stderr[-1500:]!r}" + ) class TestGeminiFreshInstall: @@ -872,15 +949,16 @@ def _all_models(self, e2e_state: dict) -> list[tuple[str, str]]: return out def test_launch_opencode_per_model( - self, tmp_path, monkeypatch, e2e_state, e2e_workspace, e2e_token + self, model_case, tmp_path, monkeypatch, e2e_state, e2e_workspace, e2e_token ): + import time + import ucode.config_io as config_io_mod from ucode.agents import opencode + _fail_on_catalog_error(model_case) + provider, model = model_case _require_binary("opencode") - models = self._all_models(e2e_state) - if not models: - pytest.skip("No OpenCode models available on this workspace") monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) xdg = tmp_path / "opencode-xdg" @@ -890,60 +968,48 @@ def test_launch_opencode_per_model( monkeypatch.setattr(opencode, "OPENCODE_CONFIG_PATH", config_path) monkeypatch.setattr(opencode, "OPENCODE_BACKUP_PATH", backup_path) - import time + with pytest.MonkeyPatch().context() as mp: + mp.setattr("ucode.state.save_state", lambda s: None) + mp.setattr( + "ucode.agents.opencode.get_databricks_token", + lambda ws, profile=None, **kwargs: e2e_token, + ) + opencode.write_tool_config( + {**e2e_state, "workspace": e2e_workspace}, + model, + token=e2e_token, + ) - print(f"\n[opencode-per-model] {len(models)} models to test", flush=True) - failures = [] - for provider, model in models: - # Reset config file before each model so configs don't bleed together - if config_path.exists(): - config_path.unlink() - - with pytest.MonkeyPatch().context() as mp: - mp.setattr("ucode.state.save_state", lambda s: None) - mp.setattr( - "ucode.agents.opencode.get_databricks_token", - lambda ws, profile=None, **kwargs: e2e_token, - ) - opencode.write_tool_config( - {**e2e_state, "workspace": e2e_workspace}, - model, - token=e2e_token, - ) - - cmd = opencode.validate_cmd("opencode") - print(f"[opencode-per-model] -> {provider}/{model}", flush=True) - t0 = time.monotonic() - try: - result = _run_agent(cmd, env=opencode.build_runtime_env(e2e_token), timeout=180) - except subprocess.TimeoutExpired as exc: - elapsed = time.monotonic() - t0 - partial_stdout = _timeout_output(exc.stdout) - partial_stderr = _timeout_output(exc.stderr) - print( - f"[opencode-per-model] {provider}/{model} TIMEOUT after {elapsed:.1f}s\n" - f" partial stdout: {partial_stdout[:500]!r}\n" - f" partial stderr: {partial_stderr[:500]!r}", - flush=True, - file=sys.stderr, - ) - failures.append( - f"provider={provider} model={model} TIMEOUT after {elapsed:.1f}s " - f"stderr={partial_stderr[:300]!r}" - ) - continue + cmd = opencode.validate_cmd("opencode") + print(f"[opencode-per-model] -> {provider}/{model}", flush=True) + t0 = time.monotonic() + try: + result = _run_agent(cmd, env=opencode.build_runtime_env(e2e_token), timeout=180) + except subprocess.TimeoutExpired as exc: elapsed = time.monotonic() - t0 - combined = (result.stdout + result.stderr).strip() - status = "OK" if result.returncode == 0 and combined else f"FAIL rc={result.returncode}" - print(f"[opencode-per-model] {provider}/{model} {status} ({elapsed:.1f}s)", flush=True) - if result.returncode != 0 or not combined: - failures.append( - f"provider={provider} model={model} rc={result.returncode} " - f"elapsed={elapsed:.1f}s " - f"stdout={result.stdout[:300]!r} stderr={result.stderr[:300]!r}" - ) - - assert not failures, "OpenCode launch failures:\n" + "\n".join(failures) + partial_stdout = _timeout_output(exc.stdout) + partial_stderr = _timeout_output(exc.stderr) + print( + f"[opencode-per-model] {provider}/{model} TIMEOUT after {elapsed:.1f}s\n" + f" partial stdout: {partial_stdout[:500]!r}\n" + f" partial stderr: {partial_stderr[:500]!r}", + flush=True, + file=sys.stderr, + ) + pytest.fail( + f"OpenCode launch failure: provider={provider} model={model} " + f"TIMEOUT after {elapsed:.1f}s stderr={partial_stderr[:300]!r}" + ) + elapsed = time.monotonic() - t0 + combined = (result.stdout + result.stderr).strip() + status = "OK" if result.returncode == 0 and combined else f"FAIL rc={result.returncode}" + print(f"[opencode-per-model] {provider}/{model} {status} ({elapsed:.1f}s)", flush=True) + if result.returncode != 0 or not combined: + pytest.fail( + f"OpenCode launch failure: provider={provider} model={model} " + f"rc={result.returncode} elapsed={elapsed:.1f}s " + f"stdout={result.stdout[:300]!r} stderr={result.stderr[:300]!r}" + ) def test_launch_deepseek_v4_pro( self, tmp_path, monkeypatch, e2e_state, e2e_workspace, e2e_token @@ -1084,19 +1150,19 @@ def test_incompatible_models_are_skipped(self): ] def test_launch_copilot_per_model( - self, tmp_path, monkeypatch, e2e_state, e2e_workspace, e2e_token + self, model_case, tmp_path, monkeypatch, e2e_state, e2e_workspace, e2e_token ): import ucode.config_io as config_io_mod from ucode.agents import copilot + _fail_on_catalog_error(model_case) + _family, model_id = model_case _require_binary("copilot") - models = self._all_models(e2e_state) - if not models: - pytest.skip("No Copilot-compatible models available on this workspace") monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) monkeypatch.setattr(copilot, "COPILOT_CONFIG_DIR", tmp_path / ".copilot") - catalog = copilot.build_copilot_catalog([model for _, model in models]) + catalog = copilot.build_copilot_catalog([model_id]) + model = catalog[0] from ucode import gateway_proxy @@ -1113,28 +1179,24 @@ def test_launch_copilot_per_model( threading.Thread(target=server.serve_forever, daemon=True).start() proxy_base_url = f"http://127.0.0.1:{server.server_address[1]}" - failures = [] config = copilot.render_providers_config(catalog, proxy_base_url, session_secret) try: with copilot.session_registry(config) as registry_path: - for model in catalog: - env = copilot.build_runtime_env(registry_path, model) - cmd = [*copilot.validate_cmd("copilot"), "--model", model.ref] - result = _run_agent(cmd, env=env, timeout=120) - combined = (result.stdout + result.stderr).strip() - if result.returncode != 0 or not combined: - failures.append( - f"model={model.ref} rc={result.returncode} " - f"stdout={result.stdout[:300]!r} stderr={result.stderr[:300]!r}" - ) + env = copilot.build_runtime_env(registry_path, model) + cmd = [*copilot.validate_cmd("copilot"), "--model", model.ref] + result = _run_agent(cmd, env=env, timeout=120) + combined = (result.stdout + result.stderr).strip() + if result.returncode != 0 or not combined: + pytest.fail( + f"Copilot launch failure: model={model.ref} rc={result.returncode} " + f"stdout={result.stdout[:300]!r} stderr={result.stderr[:300]!r}" + ) finally: cache.stop() server.shutdown() server.server_close() client.close() - assert not failures, "Copilot launch failures:\n" + "\n".join(failures) - class TestPiLaunch: """Run pi against every available model across all four providers. @@ -1168,14 +1230,15 @@ def test_incompatible_models_are_skipped(self): } assert self._all_models(state) == [("codex", "databricks-gpt-5-4")] - def test_launch_pi_per_model(self, tmp_path, monkeypatch, e2e_state, e2e_workspace, e2e_token): + def test_launch_pi_per_model( + self, model_case, tmp_path, monkeypatch, e2e_state, e2e_workspace, e2e_token + ): import ucode.config_io as config_io_mod from ucode.agents import pi + _fail_on_catalog_error(model_case) + family, model = model_case _require_binary("pi") - models = self._all_models(e2e_state) - if not models: - pytest.skip("No Pi-compatible models available on this workspace") monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) # Point PI_CODING_AGENT_DIR and ucode's config writer at the same @@ -1191,34 +1254,27 @@ def test_launch_pi_per_model(self, tmp_path, monkeypatch, e2e_state, e2e_workspa monkeypatch.setattr(pi, "PI_SETTINGS_BACKUP_PATH", tmp_path / "pi-settings.backup.json") monkeypatch.setattr(pi, "PI_BACKUP_PATH", backup_path) - failures = [] - for family, model in models: - if config_path.exists(): - config_path.unlink() - - with pytest.MonkeyPatch().context() as mp: - mp.setattr("ucode.state.save_state", lambda s: None) - mp.setattr( - "ucode.agents.pi.get_databricks_token", - lambda ws, profile=None, **kwargs: e2e_token, - ) - pi.write_tool_config( - {**e2e_state, "workspace": e2e_workspace}, - model, - token=e2e_token, - ) - - env = pi.build_runtime_env(e2e_token) - cmd = pi.validate_cmd("pi") - result = _run_agent(cmd, env=env, timeout=120) - combined = (result.stdout + result.stderr).strip() - if result.returncode != 0 or not combined: - failures.append( - f"family={family} model={model} rc={result.returncode} " - f"stdout={result.stdout[:300]!r} stderr={result.stderr[:300]!r}" - ) + with pytest.MonkeyPatch().context() as mp: + mp.setattr("ucode.state.save_state", lambda s: None) + mp.setattr( + "ucode.agents.pi.get_databricks_token", + lambda ws, profile=None, **kwargs: e2e_token, + ) + pi.write_tool_config( + {**e2e_state, "workspace": e2e_workspace}, + model, + token=e2e_token, + ) - assert not failures, "Pi launch failures:\n" + "\n".join(failures) + env = pi.build_runtime_env(e2e_token) + cmd = pi.validate_cmd("pi") + result = _run_agent(cmd, env=env, timeout=120) + combined = (result.stdout + result.stderr).strip() + if result.returncode != 0 or not combined: + pytest.fail( + f"Pi launch failure: family={family} model={model} rc={result.returncode} " + f"stdout={result.stdout[:300]!r} stderr={result.stderr[:300]!r}" + ) # --------------------------------------------------------------------------- diff --git a/tests/test_integration_contract.py b/tests/test_integration_contract.py index 2defe77e6..01f9eb06d 100644 --- a/tests/test_integration_contract.py +++ b/tests/test_integration_contract.py @@ -60,6 +60,12 @@ def test_managed_integration_ci_is_blocking(): WORKFLOW = Path(__file__).parent.parent / ".github/workflows/integration.yml" +# The quarantined Claude Smart Router toggle journey, split into one test per toggle state. +CLAUDE_TOGGLE_TESTS = ( + "test_smart_router_routes_claude_subagent_at_launch", + "test_smart_router_skill_turns_claude_subagent_routing_off", + "test_smart_router_skill_turns_claude_subagent_routing_back_on", +) def _workflow_job(name, next_name): @@ -95,12 +101,7 @@ def test_soft_fail_marks_only_the_affected_claude_journey(): for node in ast.parse(path.read_text()).body: if isinstance(node, ast.FunctionDef) and "soft_fail" in _markers(node.decorator_list): marked.append((path.name, node.name)) - assert marked == [ - ( - "test_ug_smart_routing_hooks.py", - "test_smart_router_skill_toggles_claude_subagent_routing", - ) - ] + assert marked == [("test_ug_smart_routing_hooks.py", name) for name in CLAUDE_TOGGLE_TESTS] @pytest.mark.parametrize( @@ -109,9 +110,9 @@ def test_soft_fail_marks_only_the_affected_claude_journey(): ("smoke", "claude", "test_ug_claude_headless_prompt_argument"), ("smoke", "codex", "test_ug_codex_headless_prompt_argument"), ("full", "claude", "test_smart_routing_claude_route_subagent_hook"), - ("full", "codex", "test_smart_router_skill_toggles_codex_subagent_routing"), + ("full", "codex", "test_smart_router_skill_turns_codex_subagent_routing_back_on"), ("managed", "claude", "test_ug_configure_managed_claude"), - ("managed", "codex", "test_smart_router_skill_toggles_codex_subagent_routing"), + ("managed", "codex", "test_smart_router_skill_turns_codex_subagent_routing_back_on"), ("soft-fail", "claude", None), ], ) @@ -143,13 +144,13 @@ def test_soft_fail_ci_selectors_preserve_coverage(job_name, agent, retained_test assert result.returncode == 0, result.stdout + result.stderr cases = {line for line in result.stdout.splitlines() if line.startswith("test_")} assert cases, result.stdout - toggle = "test_smart_router_skill_toggles_claude_subagent_routing" if retained_test is not None: assert any(f"::{retained_test}" in case for case in cases), cases - assert not any(f"::{toggle}" in case for case in cases), cases + assert not any(f"::{toggle}[" in case for case in cases for toggle in CLAUDE_TOGGLE_TESTS) return assert cases == { f"test_ug_smart_routing_hooks.py::{toggle}[{variant}]" + for toggle in CLAUDE_TOGGLE_TESTS for variant in ( "routing-only", "orchestration", diff --git a/tests/test_tui_request_recorder.py b/tests/test_tui_request_recorder.py index fe2a2d76c..5af20399f 100644 --- a/tests/test_tui_request_recorder.py +++ b/tests/test_tui_request_recorder.py @@ -13,6 +13,7 @@ RecordedResponse, TuiRequestRecorder, ) +from tests.integration.utils import request_recorder @pytest.fixture @@ -144,3 +145,72 @@ def fake_sudo_cp(command, check): entry = json.loads((tmp_path / ".ucode/refresh_ts.json").read_text())[recorder.url] assert "last_checked_at" not in entry assert entry["policy_fingerprint"]["exists"] is True + + +@pytest.fixture +def held_upstream(): + """An upstream whose response body stays open until the test releases it.""" + release = threading.Event() + + class Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def log_message(self, format, *args): + pass + + def do_GET(self): + self.send_response(200) + self.send_header("Content-Length", "4") + self.end_headers() + self.wfile.write(b"he") + self.wfile.flush() + release.wait(10) + self.wfile.write(b"ld") + + with ThreadingHTTPServer(("127.0.0.1", 0), Handler) as server: + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + yield f"http://127.0.0.1:{server.server_address[1]}", release + release.set() + server.shutdown() + thread.join() + + +def _start_held_request(recorder): + bodies = [] + thread = threading.Thread( + target=lambda: bodies.append(httpx.get(f"{recorder.url}/held", timeout=10).content) + ) + thread.start() + recorder.expect_request(method="GET", path="/held", timeout=5) + return thread, bodies + + +def test_teardown_waits_for_in_flight_forward_before_closing_client(held_upstream): + endpoint, release = held_upstream + recorder = TuiRequestRecorder(endpoint).__enter__() + thread, bodies = _start_held_request(recorder) + + threading.Timer(0.3, release.set).start() + recorder.__exit__(None, None, None) + + thread.join(5) + assert bodies == [b"held"] + assert recorder._client.is_closed + + +def test_teardown_leaves_client_open_for_forward_that_outlives_drain(held_upstream, monkeypatch): + # Closing the client under a still-streaming forward frees its socket descriptor while the + # forward's TLS reader keeps using that number; a reused descriptor then corrupts another + # connection in the process. The forward must finish on its own connection instead. + monkeypatch.setattr(request_recorder, "_DRAIN_SECONDS", 0.2) + endpoint, release = held_upstream + recorder = TuiRequestRecorder(endpoint).__enter__() + thread, bodies = _start_held_request(recorder) + + recorder.__exit__(None, None, None) + assert not recorder._client.is_closed + + release.set() + thread.join(5) + assert bodies == [b"held"] diff --git a/uv.lock b/uv.lock index 9c5fd3ccc..2fa86ecec 100644 --- a/uv.lock +++ b/uv.lock @@ -797,6 +797,19 @@ wheels = [ { url = "https://pypi-proxy.cloud.databricks.com/packages/24/25/1de2678b631f5a49215c6c96fff41ba892b0a34df68d6d80292b1b48aa7f/pytest-9.1.1-py3-none-any.whl", hash = "sha256:37a86b45efb9a47a61a36449063e8e18d0cab3161329fc099eb21783169c4f0c", size = 386536, upload-time = "2026-06-19T10:58:31.347Z" }, ] +[[package]] +name = "pytest-rerunfailures" +version = "16.7" +source = { registry = "https://pypi-proxy.cloud.databricks.com/simple/" } +dependencies = [ + { name = "packaging" }, + { name = "pytest" }, +] +sdist = { url = "https://pypi-proxy.cloud.databricks.com/packages/d1/b0/6b5337f9d59b26b0069ea3d5e863c31dc04b69e51bbfb1cf2ea6328fba87/pytest_rerunfailures-16.7.tar.gz", hash = "sha256:6956ddfb65ca1d07e7e3d99e2c5359d82300f4cb062e6049563bd2c106f72d5c", size = 59666, upload-time = "2026-09-17T07:08:48.871Z" } +wheels = [ + { url = "https://pypi-proxy.cloud.databricks.com/packages/c1/d3/07ea35102cf2020ddaaac368d2a6f0bcc63d6523fb8918805bac3eb8b98f/pytest_rerunfailures-16.7-py3-none-any.whl", hash = "sha256:edf1886209c2b7dafe35b5bf1708d6ec40ccf6c6b357f0f02807efcec0204c99", size = 23952, upload-time = "2026-09-17T07:08:47.635Z" }, +] + [[package]] name = "python-multipart" version = "0.0.32" @@ -1150,6 +1163,7 @@ dev = [ { name = "pexpect" }, { name = "pyte" }, { name = "pytest" }, + { name = "pytest-rerunfailures" }, { name = "pywinpty", marker = "sys_platform == 'win32'" }, { name = "ruff" }, { name = "ty" }, @@ -1171,6 +1185,7 @@ dev = [ { name = "pexpect", specifier = ">=4.9.0" }, { name = "pyte", specifier = ">=0.8.2" }, { name = "pytest", specifier = ">=9.0.3" }, + { name = "pytest-rerunfailures", specifier = ">=16.7" }, { name = "pywinpty", marker = "sys_platform == 'win32'", specifier = "==3.0.5" }, { name = "ruff", specifier = "==0.16.5" }, { name = "ty", specifier = "==0.0.75" },