[AIGTWY-4972] Consolidate CUJ helpers, fix flakes, parallelize CUJs

.github/scripts/ug-review/review.py

@@ -272,7 +272,7 @@ def build_responses_request(
         "instructions": build_system_prompt(review_policy, repository_policy),
         "input": user_prompt,
         "reasoning": {"effort": "high"},
-        "max_output_tokens": 4_000,
+        "max_output_tokens": 16_000,
         "store": False,
     }
 
@@ -405,7 +405,10 @@ def _request_review(
         raise ReviewError("Databricks returned an unexpected review response.")
     content = extract_response_text(payload)
     if not content:
-        raise ReviewError("Databricks Responses API returned no output text.")
+        incomplete = payload.get("incomplete_details")
+        cause = incomplete.get("reason") if isinstance(incomplete, dict) else None
+        reason = f" (incomplete: {cause})" if cause else ""
+        raise ReviewError(f"Databricks Responses API returned no output text{reason}.")
     return parse_review(content, bundle)
 
 

.github/workflows/integration.yml

@@ -420,14 +420,43 @@ jobs:
             ${{ runner.temp }}/ug-integration/artifacts/
             ${{ runner.temp }}/ug-integration/wheels/
 
+  dedicated-cuj-plan:
+    name: Plan E2E CUJs
+    if: ${{ github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository }}
+    runs-on: ubuntu-latest
+    timeout-minutes: 5
+    outputs:
+      cujs: ${{ steps.discover.outputs.cujs }}
+    steps:
+      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
+      - name: Discover every E2E CUJ file
+        id: discover
+        shell: bash
+        run: |
+          set -euo pipefail
+          cujs=$(for path in tests/e2e_cuj/test_*.py; do
+            name=$(sed -n 's/^CUJ_NAME = "\(.*\)"$/\1/p' "$path")
+            [ -n "$name" ] || { echo "$path must define CUJ_NAME" >&2; exit 1; }
+            jq -nc --arg test "${path##*/}" --arg name "$name" '{test: $test, name: $name}'
+          done | jq -sc .)
+          echo "cujs=$cujs" >> "$GITHUB_OUTPUT"
+
   dedicated-cuj:
-    name: E2E CUJs
+    name: E2E CUJs · ${{ matrix.name }}
     # CUJ definitions: https://docs.google.com/document/d/1WKd1fdWD0Y4tAV1H9Si-SGx2UZFL7iZ2S3HtBjidmS0/edit?usp=drive_web&ouid=117446678071138243519
-    if: ${{ github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository }}
+    # One runner per file: CUJs own separate workspaces but share machine-wide /etc agent settings.
+    needs: dedicated-cuj-plan
     runs-on: ubuntu-22.04
     timeout-minutes: 70
+    strategy:
+      fail-fast: false
+      matrix:
+        include: ${{ fromJSON(needs.dedicated-cuj-plan.outputs.cujs) }}
     env:
       UV_INDEX_URL: https://pypi.org/simple
+      # Every CUJ runner installs from PyPI at once; ride out slow responses instead of failing a shard.
+      UV_HTTP_TIMEOUT: "120"
+      UV_HTTP_RETRIES: "6"
       UG_CUJ1_WORKSPACE: ${{ secrets.UG_CUJ1_WORKSPACE }}
       UG_CUJ_SP_CLIENT_ID: ${{ secrets.UG_CUJ_SP_CLIENT_ID }}
       UG_CUJ_SP_CLIENT_SECRET: ${{ secrets.UG_CUJ_SP_CLIENT_SECRET }}
@@ -451,20 +480,16 @@ jobs:
           npm install -g \
             "@anthropic-ai/claude-code@$CLAUDE_VERSION" \
             "@openai/codex@$CODEX_VERSION"
-      - name: Run every E2E CUJ
+      - name: Run ${{ matrix.test }}
         run: |
-          if ! find tests/e2e_cuj -name 'test_*.py' -print -quit | grep -q .; then
-            echo "No E2E CUJs found."
-            exit 0
-          fi
           uv run --with pexpect==4.9.0 --with pyte==0.8.2 \
-            pytest --confcutdir=tests/e2e_cuj tests/e2e_cuj -v
+            pytest --confcutdir=tests/e2e_cuj "tests/e2e_cuj/${{ matrix.test }}" -v
 
   cujs:
     name: All integration tests
     if: ${{ always() && (github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository) }}
     # installation-windows remains advisory until its initial platform issues are fixed.
-    needs: [installation, workspace, smoke, full, managed, dedicated-cuj]
+    needs: [installation, workspace, smoke, full, managed, dedicated-cuj-plan, dedicated-cuj]
     runs-on: ubuntu-latest
     timeout-minutes: 5
     steps:

tests/README.md

@@ -203,7 +203,7 @@ integration utilities; only CUJ-specific evidence correlation stays in a test fi
 | `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 |
 | `test_ug_configure_managed_claude`, `test_ug_configure_managed_codex` | Configure on the managed workspace with the stubbed `managed_workspace_default.json` CodingAgentConfig | No agent selector; each agent's generated config exposes exactly the admin's static model_services; real gateway prompt on launch. The Codex case also checks the shared catalog pointer, restart guidance, and a fresh bare app-server's visible model list |
-| `e2e_cuj/test_ug_budget_defaults.py` | Launch bare `ug` with separate low-spend and above-tier principals against the fixed 1% tier (`spending_percentage=0.01`); explicitly launch `ug claude` above the tier | Bare launches check Claude/Sonnet below the tier and the Codex/Luna recommendation above it; `ug usage` agrees with backend spend, threshold, and percentage. Explicit `ug claude` displays the backend's Codex/Luna recommendation while its generated model setting and native header select Sonnet. No budget writes or inference tasks. Run in the shared `E2E CUJs` job with both credential pairs documented in `integration/README.md`. |
+| `e2e_cuj/test_cuj5_budget_defaults.py` | Launch bare `ug` with separate low-spend and above-tier principals against the fixed 1% tier (`spending_percentage=0.01`); explicitly launch `ug claude` above the tier | Bare launches check Claude/Sonnet below the tier and the Codex/Luna recommendation above it; `ug usage` agrees with backend spend, threshold, and percentage. Explicit `ug claude` displays the backend's Codex/Luna recommendation while its generated model setting and native header select Sonnet. No budget writes or inference tasks. Runs on its own `E2E CUJs · CUJ 5 · Budget-driven defaults` runner with both credential pairs documented in `integration/README.md`. |
 | `test_case_01_*` | Launch managed Claude without defaults after configure and from fresh state | Claude receives the admin MPS header; its gateway cache and replacement picker match the independently fetched provider model IDs; catalog labels are preserved and a model appears in a numbered picker row |
 | `test_case_03_*`, `test_case_05_*` | Pass a provider or model-location override to managed Claude after configure and from fresh state | ug rejects the override before Claude starts and preserves agent-owned state |
 | `test_case_02_*` | Launch managed Codex after configure and from fresh state | The scoped and stable catalogs, ug-launched app server, and fresh bare app server match the independently fetched admin MPS model IDs. The configured case uses real `ug revert` to remove ug's shared pointer and stable file while preserving a user setting |
@@ -316,8 +316,9 @@ locally; it does not establish live inference or in-session model switching.
 Check names describe the coverage: `Unit tests`, `Gateway API tests`,
 `Agent launch tests · Claude`, `Smoke journeys · Claude`, and
 `Full journeys · Claude` (with the other agents named likewise).
-Unit tests still run as one job. Both matrices use `fail-fast: false` so one
-failure does not cancel other coverage.
+Unit tests still run as one job. `E2E CUJs · CUJ <N> · <topic>` runs one runner per
+`tests/e2e_cuj/test_cuj<N>_*.py` file, because CUJs share machine-wide `/etc` agent
+settings. All matrices use `fail-fast: false` so one failure does not cancel other coverage.
 
 The small `test` and `e2e` compatibility gates retain the exact status contexts
 required by the repository's branch rules. `test` requires `Unit tests`; `e2e`

tests/e2e_cuj/AGENTS.md

@@ -1,5 +1,9 @@
 # Dedicated-workspace CUJs
 
+Name CUJ files `test_cuj<N>_<topic>.py` and give each a module-level
+`CUJ_NAME = "CUJ <N> · <topic>"`. CI discovers every file and runs it on its own runner
+as `E2E CUJs · <CUJ_NAME>`.
+
 Each CUJ owns a separate workspace. Subclass `BaseCujTest` from `base.py` and set
 `WORKSPACE_URL`. Setup provides `self.workspace`, a Databricks SDK client using
 `UG_CUJ_SP_CLIENT_ID` and `UG_CUJ_SP_CLIENT_SECRET` with OAuth M2M authentication.
@@ -47,16 +51,20 @@ In `helpers/evidence.py`, `ClaudeCujHelper` and `CodexCujHelper` implement
 
 Model evidence combines observed gateway requests/responses with native completed-turn
 records; neither routing banners nor native records alone prove an applied decision.
+`assert_models` checks a canonical native model set and `assert_served` a paired HTTP 200.
+`Workspace.agent_configs` splits a published config by agent.
+`helpers/mcp.py` names the dashed server a managed MCP service registers as.
 
 `test_cuj2_mps_mcp.py` runs three independently collected cases against a permanently
 preconfigured workspace: configuration, Codex inference, and Claude inference. `test_cuj2_skills.py`
 adds the named-skill cases against the same workspace. Each case
 reads and validates the exact two-agent MPS/MCP shape. No case creates, updates, deletes,
 or restores CodingAgentConfig. The fixture's final read-only equality check is the only
 workspace mutation guard. The test uses the shared `UserSession`, `Workspace`, and
-`TuiRequestRecorder`, `AgentTerminal`, and integration transcript readers directly. The
-configuration case verifies the selected sandbox service and generated MCP listings; live
-`run_code` execution remains deferred because MAS cannot downscope the CI service principal.
+`TuiRequestRecorder`, `AgentTerminal` (via `helpers/terminal.py`), and integration
+transcript readers directly. The configuration case verifies the selected sandbox service
+and generated MCP listings; live `run_code` execution remains deferred because MAS cannot
+downscope the CI service principal.
 
 Use a clean disposable POSIX runner without existing machine-wide agent settings.
 Install the intended versions of `ug`, `claude`, `codex`, and `databricks` on PATH.

tests/e2e_cuj/README.md

@@ -38,13 +38,12 @@ workspace configuration remains read-only. Each case cleans up with public `ug r
 the next case configures, including after failures. Set `UG_CUJ_SP_CLIENT_ID` and
 `UG_CUJ_SP_CLIENT_SECRET` before a live run.
 
-CI discovers these tests through the shared `dedicated-cuj` job, which installs both
-pinned agents. No catalog-specific workflow or workspace secret is needed. Locally,
+CI discovers each test file and runs it on its own runner with both pinned agents. No catalog-specific workflow or workspace secret is needed. Locally,
 install ug, both agents, and the Databricks CLI on a clean POSIX host, then run:
 
 ```bash
 uv run --with pexpect==4.9.0 --with pyte==0.8.2 pytest \
-  --confcutdir=tests/e2e_cuj tests/e2e_cuj/test_catalog_discovery.py -v
+  --confcutdir=tests/e2e_cuj tests/e2e_cuj/test_cuj3_models.py -v
 ```
 
 MCP/skills discovery remains separate.

tests/e2e_cuj/base.py

@@ -7,6 +7,28 @@
 from databricks.sdk import WorkspaceClient
 
 
+def make_workspace_client(url: str, client_id_env: str, client_secret_env: str) -> WorkspaceClient:
+    client_id = os.environ.get(client_id_env, "")
+    client_secret = os.environ.get(client_secret_env, "")
+    if not client_id or not client_secret:
+        pytest.fail(f"Set {client_id_env} and {client_secret_env}.", pytrace=False)
+    return WorkspaceClient(
+        host=url,
+        client_id=client_id,
+        client_secret=client_secret,
+        auth_type="oauth-m2m",
+    )
+
+
+def bearer(client: WorkspaceClient) -> str:
+    """Mint a short-lived bearer; only it, never the SP secret, reaches agent processes."""
+    authorization = client.config.authenticate().get("Authorization", "")
+    token = authorization.removeprefix("Bearer ")
+    if not authorization.startswith("Bearer ") or not token:
+        raise RuntimeError("Service-principal authentication did not return a bearer token.")
+    return token
+
+
 class BaseCujTest:
     WORKSPACE_URL: ClassVar[str] = ""
     CLIENT_ID_ENV: ClassVar[str] = "UG_CUJ_SP_CLIENT_ID"
@@ -19,16 +41,6 @@ def setup_workspace(self, request):
         if not workspace_url:
             pytest.fail("Set WORKSPACE_URL on your CUJ test class.", pytrace=False)
 
-        client_id_env = request.cls.CLIENT_ID_ENV
-        client_secret_env = request.cls.CLIENT_SECRET_ENV
-        client_id = os.environ.get(client_id_env, "")
-        client_secret = os.environ.get(client_secret_env, "")
-        if not client_id or not client_secret:
-            pytest.fail(f"Set {client_id_env} and {client_secret_env}.", pytrace=False)
-
-        request.cls.workspace = WorkspaceClient(
-            host=workspace_url,
-            client_id=client_id,
-            client_secret=client_secret,
-            auth_type="oauth-m2m",
+        request.cls.workspace = make_workspace_client(
+            workspace_url, request.cls.CLIENT_ID_ENV, request.cls.CLIENT_SECRET_ENV
         )

tests/e2e_cuj/conftest.py

@@ -8,10 +8,14 @@
 import pytest
 from databricks.sdk.errors import DatabricksError
 
-from tests.integration.utils.terminal import TerminalProcess
-
+from .base import bearer
 from .helpers.constants import CLAUDE, CODEX, MANAGED_PATHS
-from .helpers.session import UserSession
+from .helpers.session import (
+    MACHINE_WIDE_LEAK,
+    MachineWideLeak,
+    UserSession,
+    dirty_runner_message,
+)
 from .helpers.tui_request_recorder import TuiRequestRecorder
 from .helpers.workspace import Workspace
 
@@ -22,26 +26,37 @@ def pytest_configure(config):
         raise pytest.UsageError("Run with --confcutdir=tests/e2e_cuj to exclude unit-test mocks.")
 
 
+def _scrubbed(error: DatabricksError) -> RuntimeError:
+    # Server messages may echo credentials; retain only the SDK error type.
+    return RuntimeError(f"Workspace API failed: {type(error).__name__}")
+
+
+def _record_machine_wide_leak(request):
+    """Name the class that leaked machine-wide settings to the later classes it blocks."""
+    leaked = tuple(str(path) for path in MANAGED_PATHS if path.exists())
+    if leaked:
+        request.config.stash[MACHINE_WIDE_LEAK] = MachineWideLeak(request.node.nodeid, leaked)
+
+
 @pytest.fixture(scope="class")
 def cuj(request, setup_workspace, tmp_path_factory):
     assert os.name == "posix", "Full TUI CUJs require a disposable POSIX runner"
-    assert not any(path.exists() for path in MANAGED_PATHS), (
-        "Existing machine-wide agent settings; use a clean disposable runner. Nothing was changed."
+    assert not any(path.exists() for path in MANAGED_PATHS), dirty_runner_message(
+        request.config.stash.get(MACHINE_WIDE_LEAK, None)
     )
     for tool in ("ug", CLAUDE, CODEX, "databricks"):
         assert shutil.which(tool), f"Install the required CLI before running this CUJ: {tool}"
 
-    authorization = request.cls.workspace.config.authenticate().get("Authorization", "")
-    if not authorization.startswith("Bearer ") or not authorization.removeprefix("Bearer "):
-        raise RuntimeError("Service-principal authentication did not return a bearer token.")
-    bearer = authorization.removeprefix("Bearer ")
+    class_bearer = bearer(request.cls.workspace)
     run_directory = tmp_path_factory.mktemp(request.cls.__name__)
     artifacts = run_directory / "artifacts"
     print(f"CUJ artifacts: {artifacts}")
     with tempfile.TemporaryDirectory(prefix="ug-cuj-", dir="/tmp") as temporary:
         workspace = Workspace(request.cls.workspace)
         with TuiRequestRecorder(workspace.url) as recorder:
-            session = UserSession(Path(temporary), Path(shutil.which("ug")), artifacts, bearer)
+            session = UserSession(
+                Path(temporary), Path(shutil.which("ug")), artifacts, class_bearer
+            )
             try:
                 published = workspace.config()
                 try:
@@ -50,21 +65,23 @@ def cuj(request, setup_workspace, tmp_path_factory):
                     try:
                         workspace.assert_unchanged(published)
                     finally:
-                        state_dir = session.home / ".ucode"
-                        if any(
-                            (state_dir / name).is_file()
-                            for name in ("state.json", "managed-backups/manifest.json")
-                        ):
-                            with TerminalProcess(
-                                session,
-                                "ug",
-                                [str(session.binary), "revert"],
-                                "cleanup-revert",
-                            ) as terminal:
-                                terminal.finish()
-                        assert not any(path.exists() for path in MANAGED_PATHS), (
-                            "CUJ teardown left machine-wide agent settings"
-                        )
+                        try:
+                            session.revert_machine_wide(
+                                "cleanup-revert", "CUJ teardown left machine-wide agent settings"
+                            )
+                        finally:
+                            _record_machine_wide_leak(request)
             except DatabricksError as error:
-                # Server messages may echo credentials; retain only the SDK error type.
-                raise RuntimeError(f"Workspace API failed: {type(error).__name__}") from None
+                raise _scrubbed(error) from None
+
+
+@pytest.fixture(autouse=True)
+def _refresh_cuj_bearer(request):
+    # The class bearer is minted once, but a long class can outlive the ~1h M2M token.
+    if "cuj" in request.fixturenames:
+        session, _, _ = request.getfixturevalue("cuj")
+        try:
+            refreshed = bearer(request.cls.workspace)
+        except DatabricksError as error:
+            raise _scrubbed(error) from None
+        session.refresh_bearer(refreshed)

tests/e2e_cuj/helpers/constants.py

@@ -1,4 +1,4 @@
-"""Agent, metastore fixture, and transcript values shared by CUJs."""
+"""Agent, model, header, metastore fixture, and transcript values shared by CUJs."""
 
 from enum import StrEnum
 from pathlib import Path
@@ -15,6 +15,19 @@
     CLAUDE: "/ai-gateway/anthropic/v1/messages",
     CODEX: "/ai-gateway/codex/v1/responses",
 }
+CLAUDE_OPUS_MODEL = "system.ai.claude-opus-4-8"
+CLAUDE_SONNET_MODEL = "system.ai.claude-sonnet-4-6"
+CLAUDE_HAIKU_MODEL = "system.ai.claude-haiku-4-5"
+# Claude reports Bedrock-routed Haiku under this native id.
+NATIVE_MODEL_ALIASES = {
+    "system.ai.anthropic.claude-haiku-4-5-20251001-v1:0": CLAUDE_HAIKU_MODEL,
+}
+CODEX_SOL_MODEL = "system.ai.gpt-5-6-sol"
+CODEX_LUNA_MODEL = "system.ai.gpt-5-6-luna"
+# TuiRequestRecorder lowercases recorded header names.
+RUN_HEADER = "x-ug-e2e-run"
+AGENT_HEADER = "x-ug-e2e-agent"
+PROVIDER_HEADER = "databricks-model-provider-service"
 
 
 class CodingAgent(StrEnum):