diff --git a/.github/scripts/ug-review/review.py b/.github/scripts/ug-review/review.py index c3753dfed..bc86ec729 100644 --- a/.github/scripts/ug-review/review.py +++ b/.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) diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml index 6dee17d30..315243f00 100644 --- a/.github/workflows/integration.yml +++ b/.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: diff --git a/tests/README.md b/tests/README.md index f039e7d4b..60fc2d3e2 100644 --- a/tests/README.md +++ b/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 · ` runs one runner per +`tests/e2e_cuj/test_cuj_*.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` diff --git a/tests/e2e_cuj/AGENTS.md b/tests/e2e_cuj/AGENTS.md index 767198f54..98ddc8e07 100644 --- a/tests/e2e_cuj/AGENTS.md +++ b/tests/e2e_cuj/AGENTS.md @@ -1,5 +1,9 @@ # Dedicated-workspace CUJs +Name CUJ files `test_cuj_.py` and give each a module-level +`CUJ_NAME = "CUJ · "`. CI discovers every file and runs it on its own runner +as `E2E CUJs · `. + 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,6 +51,9 @@ 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` @@ -54,9 +61,10 @@ 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. diff --git a/tests/e2e_cuj/README.md b/tests/e2e_cuj/README.md index 52f5a9fbb..ff6dafdbe 100644 --- a/tests/e2e_cuj/README.md +++ b/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. diff --git a/tests/e2e_cuj/base.py b/tests/e2e_cuj/base.py index cbffca794..0823bef72 100644 --- a/tests/e2e_cuj/base.py +++ b/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 ) diff --git a/tests/e2e_cuj/conftest.py b/tests/e2e_cuj/conftest.py index b11fde12a..04ec5877b 100644 --- a/tests/e2e_cuj/conftest.py +++ b/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) diff --git a/tests/e2e_cuj/helpers/constants.py b/tests/e2e_cuj/helpers/constants.py index 3981e01f3..a46848f61 100644 --- a/tests/e2e_cuj/helpers/constants.py +++ b/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): diff --git a/tests/e2e_cuj/helpers/evidence.py b/tests/e2e_cuj/helpers/evidence.py index 7c955d012..eb75da959 100644 --- a/tests/e2e_cuj/helpers/evidence.py +++ b/tests/e2e_cuj/helpers/evidence.py @@ -6,9 +6,9 @@ from abc import ABC, abstractmethod from dataclasses import asdict, dataclass -from tests.integration.utils.evidence import read_jsonl +from tests.integration.utils.evidence import FileTask, read_jsonl -from .constants import CLAUDE, CODEX +from .constants import CLAUDE, CODEX, NATIVE_MODEL_ALIASES def canonical_model(value): @@ -24,6 +24,33 @@ def canonical_model(value): return "system.ai." + value +def assert_models(models, expected): + """Every canonical model equals `expected`; not for UC models, which canonical_model mangles.""" + actual = {NATIVE_MODEL_ALIASES.get(model, model) for model in map(canonical_model, models)} + assert actual == {canonical_model(expected)}, (models, expected) + + +def assert_served(recorder, request, model): + """The request asked for `model` and its paired response was a non-empty HTTP 200.""" + assert request.payload["model"] == model, request.payload + response = recorder.response_for(request, timeout=240) + assert response.status_code == 200, { + "status": response.status_code, + "model": request.payload["model"], + "output_config": request.payload.get("output_config"), + } + assert response.body, "Inference response was empty" + + +def claude_file_task(session): + """A FileTask naming its absolute path, so Claude reads it without a `find` permission prompt.""" + task = FileTask(session) + task.prompt = ( + f"Use the Read tool to read {session.cwd / task.filename}. Reply with only its contents." + ) + return task + + def message_text(content): if isinstance(content, str): return content @@ -197,6 +224,11 @@ def completed(self, task): assert len(found) <= 1, "Task matched multiple sessions" return found[0] if found else None + def assert_models(self, task, expected): + turn = self.completed(task) + assert turn is not None, f"No completed native turn matched {task.prompt!r}" + assert_models(turn.models, expected) + def observe(self, task): return SessionObservation(helper=self.helper, turn=self.completed(task)) diff --git a/tests/e2e_cuj/helpers/mcp.py b/tests/e2e_cuj/helpers/mcp.py new file mode 100644 index 000000000..aa5caf49b --- /dev/null +++ b/tests/e2e_cuj/helpers/mcp.py @@ -0,0 +1,6 @@ +"""Names under which ug registers managed MCP services with each agent.""" + + +def registered_name(service: str) -> str: + """The server name ug registers for a UC MCP service, for example `ug_e2e-tools-x`.""" + return service.replace(".", "-") diff --git a/tests/e2e_cuj/helpers/poll.py b/tests/e2e_cuj/helpers/poll.py new file mode 100644 index 000000000..26d4f2260 --- /dev/null +++ b/tests/e2e_cuj/helpers/poll.py @@ -0,0 +1,14 @@ +"""Wait for eventually consistent CUJ evidence.""" + +import time + + +def poll(fn, *, timeout, interval, done=bool): + """Call `fn` until `done` accepts its value or `timeout` elapses; return the last value.""" + deadline = time.monotonic() + timeout + while not done(value := fn()): + remaining = deadline - time.monotonic() + if remaining <= 0: + break + time.sleep(min(interval, remaining)) + return value diff --git a/tests/e2e_cuj/helpers/session.py b/tests/e2e_cuj/helpers/session.py index d6537d7d5..2f1a57d2d 100644 --- a/tests/e2e_cuj/helpers/session.py +++ b/tests/e2e_cuj/helpers/session.py @@ -1,21 +1,63 @@ """CUJ authentication and cleanup on the existing integration session.""" import re +from dataclasses import dataclass + +import pytest from tests.integration.utils.harness import UserSession as IntegrationSession +from tests.integration.utils.terminal import TerminalProcess from . import constants +MANAGED_STATE_FILES = ("state.json", "managed-backups/manifest.json") + + +@dataclass(frozen=True) +class MachineWideLeak: + culprit: str + paths: tuple[str, ...] + + +MACHINE_WIDE_LEAK = pytest.StashKey[MachineWideLeak]() + + +def dirty_runner_message(leak: MachineWideLeak | None) -> str: + if leak is None: + return ( + "Existing machine-wide agent settings; use a clean disposable runner. " + "Nothing was changed." + ) + return ( + f"{leak.culprit} left machine-wide agent settings {list(leak.paths)}. Nothing was changed." + ) + class UserSession(IntegrationSession): def __init__(self, root, binary, artifacts, bearer): super().__init__(root, root, binary, artifacts) self.home.chmod(0o700) - self.env.update(DATABRICKS_BEARER=bearer, TERM="xterm-256color") + self.env.update(TERM="xterm-256color") self.env.pop("NO_COLOR", None) + self._installed_bearers = set() + self.install_bearer(bearer) + self._class_bearer = bearer + + def install_bearer(self, bearer): + """Use a bearer for later commands; replaced bearers stay redacted in later artifacts.""" + self._installed_bearers.add(bearer) + self.env["DATABRICKS_BEARER"] = bearer + + def refresh_bearer(self, bearer): + """Replace an expiring class bearer, unless a journey switched to another workspace's.""" + if self.env["DATABRICKS_BEARER"] == self._class_bearer: + self._class_bearer = bearer + self.install_bearer(bearer) def redact(self, text, *, strip_ansi=False): text = super().redact(text, strip_ansi=strip_ansi) + for token in self._installed_bearers: + text = text.replace(token, "") return re.sub(r"(?i)(bearer\s+)[A-Za-z0-9._~+/-]+=*", r"\1", text) def configure(self, args): @@ -24,3 +66,10 @@ def configure(self, args): assert not any(path.exists() for path in constants.MANAGED_PATHS), ( "Non-interactive CUJ configure unexpectedly created machine-wide settings" ) + + def revert_machine_wide(self, name, message): + state_dir = self.home / ".ucode" + if any((state_dir / file).is_file() for file in MANAGED_STATE_FILES): + with TerminalProcess(self, "ug", [str(self.binary), "revert"], name) as terminal: + terminal.finish() + assert not any(path.exists() for path in constants.MANAGED_PATHS), message diff --git a/tests/e2e_cuj/helpers/terminal.py b/tests/e2e_cuj/helpers/terminal.py index 9bd06224d..3ef58715a 100644 --- a/tests/e2e_cuj/helpers/terminal.py +++ b/tests/e2e_cuj/helpers/terminal.py @@ -7,18 +7,38 @@ class Terminal(AgentTerminal): - def __init__(self, session, name, args): - if not args or args[0] not in (CLAUDE, CODEX): - raise ValueError("CUJ terminal requires a Claude or Codex command.") - super().__init__(session, args[0], [str(session.binary), *args], name) + def __init__(self, session, name, args, *, agent=None): + """Run `ug `; bare `ug` (empty args) names the agent it is expected to launch.""" + agent = agent or (args[0] if args else None) + if agent not in (CLAUDE, CODEX) or (args and args[0] != agent): + raise ValueError("CUJ terminal requires a Claude or Codex command or bare ug agent.") + super().__init__(session, agent, [str(session.binary), *args], name) + + def wait_until( + self, + done, + description, + *, + timeout=240, + rejected=("Do you want to proceed?",), + on_screen=None, + ): + """Wait for `done()`, failing on API errors or any `rejected` permission prompt. + + `on_screen(screen)` may answer an expected dialog; it returns True when it sent keys. + """ - def task(self, evidence, task): def completed(screen): assert_no_terminal_api_error(screen) # Do not use wait_for_task's optional tool-permission approval. - assert "Do you want to proceed?" not in screen, ( + assert not any(prompt in screen for prompt in rejected), ( "Unexpected permission request; inspect the actual command:\n" + screen ) - return evidence.completed(task) is not None + if on_screen is not None and on_screen(screen): + return False + return done() - self.wait_for(completed, "completed native file task", timeout=240) + self.wait_for(completed, description, timeout=timeout) + + def task(self, evidence, task): + self.wait_until(lambda: evidence.completed(task) is not None, "completed native file task") diff --git a/tests/e2e_cuj/helpers/workspace.py b/tests/e2e_cuj/helpers/workspace.py index d53ae4a96..8828c4125 100644 --- a/tests/e2e_cuj/helpers/workspace.py +++ b/tests/e2e_cuj/helpers/workspace.py @@ -26,6 +26,16 @@ def __init__(self, client: WorkspaceClient): def policy(cls, config): return {key: value for key, value in config.items() if key not in cls.OUTPUT_FIELDS} + @staticmethod + def agent_configs(config): + """Each enabled agent's settings, keyed by its CodingAgent value.""" + entries = config["enabled_agents"] + configs = {entry["agent"]: entry["config"] for entry in entries} + assert len(configs) == len(entries), ( + f"Duplicate enabled agents: {[entry['agent'] for entry in entries]}" + ) + return configs + def config(self): result = self.client.api_client.do("GET", "/api/ai-gateway/v2/coding-agent-configs") configs = result.get("coding_agent_configs", []) diff --git a/tests/e2e_cuj/test_ug_managed_sentinel.py b/tests/e2e_cuj/test_cuj1_explicit_models.py similarity index 85% rename from tests/e2e_cuj/test_ug_managed_sentinel.py rename to tests/e2e_cuj/test_cuj1_explicit_models.py index cdf367766..0c1d0608e 100644 --- a/tests/e2e_cuj/test_ug_managed_sentinel.py +++ b/tests/e2e_cuj/test_cuj1_explicit_models.py @@ -5,7 +5,6 @@ import json import os import re -import time import tomllib import uuid @@ -16,32 +15,36 @@ from tests.integration.utils.terminal import AgentTerminal from .base import BaseCujTest -from .helpers.constants import CLAUDE, CODEX -from .helpers.evidence import SessionEvidence, canonical_model +from .helpers.constants import ( + AGENT_HEADER, + CLAUDE, + CLAUDE_HAIKU_MODEL, + CLAUDE_OPUS_MODEL, + CLAUDE_SONNET_MODEL, + CODEX, + CODEX_LUNA_MODEL, + CODEX_SOL_MODEL, + INFERENCE_PATHS, + RUN_HEADER, +) +from .helpers.evidence import SessionEvidence +from .helpers.poll import poll from .helpers.tui_request_recorder import TuiRequestRecorder -CLAUDE_MODELS = [ - "system.ai.claude-opus-4-8", - "system.ai.claude-sonnet-4-6", - "system.ai.claude-haiku-4-5", -] +CUJ_NAME = "CUJ 1 · Explicit models, tracing, headers" + +CLAUDE_MODELS = [CLAUDE_OPUS_MODEL, CLAUDE_SONNET_MODEL, CLAUDE_HAIKU_MODEL] CLAUDE_DEFAULTS = { - "default_model": "system.ai.claude-sonnet-4-6", - "default_opus_model": "system.ai.claude-opus-4-8", - "default_sonnet_model": "system.ai.claude-sonnet-4-6", - "default_haiku_model": "system.ai.claude-haiku-4-5", + "default_model": CLAUDE_SONNET_MODEL, + "default_opus_model": CLAUDE_OPUS_MODEL, + "default_sonnet_model": CLAUDE_SONNET_MODEL, + "default_haiku_model": CLAUDE_HAIKU_MODEL, } -CODEX_MODELS = ["system.ai.gpt-5-6-sol", "system.ai.gpt-5-6-luna"] -CODEX_DEFAULT = CODEX_MODELS[0] -CODEX_LUNA = CODEX_MODELS[1] -EXPECTED_HEADER_NAMES = {"x-ug-e2e-run", "x-ug-e2e-agent"} -CLAUDE_REQUEST_PATH = "/ai-gateway/anthropic/v1/messages" -CODEX_REQUEST_PATH = "/ai-gateway/codex/v1/responses" +CODEX_MODELS = [CODEX_SOL_MODEL, CODEX_LUNA_MODEL] +CODEX_DEFAULT = CODEX_SOL_MODEL +EXPECTED_HEADER_NAMES = {RUN_HEADER, AGENT_HEADER} TRACE_WAIT_SECONDS = 360 TRACE_POLL_SECONDS = 10 -NATIVE_MODEL_ALIASES = { - "system.ai.anthropic.claude-haiku-4-5-20251001-v1:0": CLAUDE_MODELS[2], -} def _claude_models_visible(text: str) -> bool: @@ -82,9 +85,9 @@ def _assert_inference_requests( served_models = [] marked_requests = 0 for request in requests: - assert request.headers["x-ug-e2e-run"] == run_id, request.headers + assert request.headers[RUN_HEADER] == run_id, request.headers # AIGTWY-4876: Requests retain the configured agent header. - assert request.headers["x-ug-e2e-agent"] == agent, request.headers + assert request.headers[AGENT_HEADER] == agent, request.headers payload = request.payload assert isinstance(payload, dict), type(payload) assert isinstance(payload.get("model"), str), sorted(payload) @@ -99,16 +102,6 @@ def _assert_inference_requests( ), served_models -def _assert_native_model(evidence: SessionEvidence, task: FileTask, expected: str) -> None: - turn = evidence.completed(task) - assert turn is not None, f"No completed native turn matched {task.prompt!r}" - expected_model = canonical_model(expected) - actual_models = { - NATIVE_MODEL_ALIASES.get(model, model) for model in map(canonical_model, turn.models) - } - assert actual_models == {expected_model}, (turn, expected_model) - - def _managed_config(session) -> dict: cache_path = session.home / ".ucode" / "managed-config.json" cache = json.loads(cache_path.read_text()) @@ -170,15 +163,15 @@ def _assert_published_config(raw: dict, entries: dict[str, dict]) -> dict[str, s assert claude["default_models"] == CLAUDE_DEFAULTS, claude assert claude["tracing"] == {"enabled": True}, claude claude_headers = _normalized_headers(claude["http_headers"]) - assert claude_headers["x-ug-e2e-agent"] == CLAUDE, claude_headers + assert claude_headers[AGENT_HEADER] == CLAUDE, claude_headers assert codex["smart_routing"] == {"enabled": False}, codex assert codex["models"] == {"model_services": CODEX_MODELS}, codex assert codex["default_models"] == {"default_model": CODEX_DEFAULT}, codex assert codex["tracing"] == {"enabled": True}, codex codex_headers = _normalized_headers(codex["http_headers"]) - assert codex_headers["x-ug-e2e-agent"] == CODEX, codex_headers - assert codex_headers["x-ug-e2e-run"] == claude_headers["x-ug-e2e-run"], ( + assert codex_headers[AGENT_HEADER] == CODEX, codex_headers + assert codex_headers[RUN_HEADER] == claude_headers[RUN_HEADER], ( claude_headers, codex_headers, ) @@ -222,8 +215,8 @@ def _assert_generated_configs(session, claude_headers: dict[str, str], workspace codex_provider = codex_profile["model_providers"]["Databricks"] assert codex_provider["base_url"] == f"{workspace}/ai-gateway/codex/v1", codex_profile assert _managed_header_subset(codex_provider["http_headers"]) == { - "x-ug-e2e-run": claude_headers["x-ug-e2e-run"], - "x-ug-e2e-agent": CODEX, + RUN_HEADER: claude_headers[RUN_HEADER], + AGENT_HEADER: CODEX, }, codex_profile codex_catalog_path = session.home / ".ucode" / "codex-model-catalog.json" assert codex_profile["model_catalog_json"] == str(codex_catalog_path), codex_profile @@ -350,39 +343,38 @@ def _wait_for_trace_pairs( pairs: list[tuple[str, str, str]], session, ) -> None: - deadline = time.monotonic() + TRACE_WAIT_SECONDS - observed = 0 - observed_client = 0 - while True: - observed = _trace_pair_count(workspace, bearer, warehouse_id, table, pairs) - observed_client = _trace_client_pair_count(workspace, bearer, warehouse_id, table, pairs) - if observed >= len(pairs) and observed_client >= len(pairs): - session.record( - "trace-query.json", - { - "table": table, - "expected_pairs": [ - {"marker": marker, "model": model, "marker_source": source} - for marker, model, source in pairs - ], - "observed_pairs": observed, - "observed_native_client_pairs": observed_client, - "native_client_model_scope": "expected model per marker; auxiliary models allowed", - }, - ) - return - remaining = deadline - time.monotonic() - if remaining <= 0: - pytest.fail( - "Managed sentinel trace evidence timed out: " - f"observed {observed}/{len(pairs)} internal and " - f"{observed_client}/{len(pairs)} native client marker/model pairs in {table}. " - "The generated headers below are local writer evidence; the public trace table " - "does not expose arbitrary HTTP request headers, so this failure does not claim " - "that X-UG-E2E-* headers reached the gateway. Native client spans may also " - "contain auxiliary model calls for the same task." - ) - time.sleep(min(TRACE_POLL_SECONDS, remaining)) + observed, observed_client = poll( + lambda: ( + _trace_pair_count(workspace, bearer, warehouse_id, table, pairs), + _trace_client_pair_count(workspace, bearer, warehouse_id, table, pairs), + ), + timeout=TRACE_WAIT_SECONDS, + interval=TRACE_POLL_SECONDS, + done=lambda counts: all(count >= len(pairs) for count in counts), + ) + if observed < len(pairs) or observed_client < len(pairs): + pytest.fail( + "Managed sentinel trace evidence timed out: " + f"observed {observed}/{len(pairs)} internal and " + f"{observed_client}/{len(pairs)} native client marker/model pairs in {table}. " + "The generated headers below are local writer evidence; the public trace table " + "does not expose arbitrary HTTP request headers, so this failure does not claim " + "that X-UG-E2E-* headers reached the gateway. Native client spans may also " + "contain auxiliary model calls for the same task." + ) + session.record( + "trace-query.json", + { + "table": table, + "expected_pairs": [ + {"marker": marker, "model": model, "marker_source": source} + for marker, model, source in pairs + ], + "observed_pairs": observed, + "observed_native_client_pairs": observed_client, + "native_client_model_scope": "expected model per marker; auxiliary models allowed", + }, + ) MANAGED_WORKSPACE_URL = "https://dbc-135c115c-c255.cloud.databricks.com" @@ -450,13 +442,13 @@ def test_cuj_managed_claude(self, cuj): claude_headers = _assert_published_config(raw, entries) assert claude["models"] == {"model_services": CLAUDE_MODELS}, claude assert claude["default_models"] == CLAUDE_DEFAULTS, claude - assert claude_headers["x-ug-e2e-agent"] == CLAUDE, claude_headers + assert claude_headers[AGENT_HEADER] == CLAUDE, claude_headers claude_settings_path = session.home / ".claude" / "ucode-settings.json" settings = json.loads(claude_settings_path.read_text()) configured_base_url = f"{workspace.url}/ai-gateway/anthropic" assert settings["env"]["ANTHROPIC_BASE_URL"] == configured_base_url, settings - run_id = claude_headers["x-ug-e2e-run"] + run_id = claude_headers[RUN_HEADER] root_task = FileTask(session) root_marker = f"ug-managed-claude-{uuid.uuid4().hex}-root" @@ -473,11 +465,11 @@ def test_cuj_managed_claude(self, cuj): assert _claude_picker_visible(picker_screen), picker_screen tui.exit_normally() root_task.assert_completed(session, CLAUDE) - _assert_native_model(evidence, root_task, CLAUDE_DEFAULTS["default_model"]) + evidence.assert_models(root_task, CLAUDE_DEFAULTS["default_model"]) _assert_inference_requests( recorder, root_checkpoint, - CLAUDE_REQUEST_PATH, + INFERENCE_PATHS[CLAUDE], root_marker, run_id, CLAUDE, @@ -502,11 +494,11 @@ def test_cuj_managed_claude(self, cuj): tui.wait_for_task(claude_opus_task) tui.exit_normally() claude_opus_task.assert_completed(session, CLAUDE) - _assert_native_model(evidence, claude_opus_task, CLAUDE_DEFAULTS["default_opus_model"]) + evidence.assert_models(claude_opus_task, CLAUDE_DEFAULTS["default_opus_model"]) _assert_inference_requests( recorder, opus_checkpoint, - CLAUDE_REQUEST_PATH, + INFERENCE_PATHS[CLAUDE], claude_opus_marker, run_id, CLAUDE, @@ -531,11 +523,11 @@ def test_cuj_managed_claude(self, cuj): tui.wait_for_task(claude_sonnet_task) tui.exit_normally() claude_sonnet_task.assert_completed(session, CLAUDE) - _assert_native_model(evidence, claude_sonnet_task, CLAUDE_DEFAULTS["default_sonnet_model"]) + evidence.assert_models(claude_sonnet_task, CLAUDE_DEFAULTS["default_sonnet_model"]) _assert_inference_requests( recorder, sonnet_checkpoint, - CLAUDE_REQUEST_PATH, + INFERENCE_PATHS[CLAUDE], claude_sonnet_marker, run_id, CLAUDE, @@ -560,11 +552,11 @@ def test_cuj_managed_claude(self, cuj): tui.wait_for_task(claude_haiku_task) tui.exit_normally() claude_haiku_task.assert_completed(session, CLAUDE) - _assert_native_model(evidence, claude_haiku_task, CLAUDE_DEFAULTS["default_haiku_model"]) + evidence.assert_models(claude_haiku_task, CLAUDE_DEFAULTS["default_haiku_model"]) _assert_inference_requests( recorder, haiku_checkpoint, - CLAUDE_REQUEST_PATH, + INFERENCE_PATHS[CLAUDE], claude_haiku_marker, run_id, CLAUDE, @@ -597,7 +589,7 @@ def test_cuj_managed_codex(self, cuj): claude_headers = _assert_published_config(raw, entries) assert codex["models"] == {"model_services": CODEX_MODELS}, codex assert codex["default_models"] == {"default_model": CODEX_DEFAULT}, codex - assert _normalized_headers(codex["http_headers"])["x-ug-e2e-agent"] == CODEX, codex + assert _normalized_headers(codex["http_headers"])[AGENT_HEADER] == CODEX, codex codex_models = session.codex_model_ids( ["app-server", "--listen", "stdio://"], name="managed-codex-models" ) @@ -607,7 +599,7 @@ def test_cuj_managed_codex(self, cuj): assert profile["model_providers"]["Databricks"]["base_url"] == ( f"{recorder.url}/ai-gateway/codex/v1" ), profile - run_id = claude_headers["x-ug-e2e-run"] + run_id = claude_headers[RUN_HEADER] session.env.pop("OTEL_RESOURCE_ATTRIBUTES", None) codex_default_task = FileTask(session) @@ -633,11 +625,11 @@ def test_cuj_managed_codex(self, cuj): tui.wait_for_task(codex_default_task) tui.exit_normally() codex_default_task.assert_completed(session, CODEX) - _assert_native_model(evidence, codex_default_task, CODEX_DEFAULT) + evidence.assert_models(codex_default_task, CODEX_DEFAULT) _assert_inference_requests( recorder, default_checkpoint, - CODEX_REQUEST_PATH, + INFERENCE_PATHS[CODEX], codex_default_marker, run_id, CODEX, @@ -658,7 +650,7 @@ def test_cuj_managed_codex(self, cuj): CODEX, "--", "--model", - CODEX_LUNA, + CODEX_LUNA_MODEL, "--config", f'otel.span_attributes.ug_integration_marker="{codex_luna_marker}"', ], @@ -669,15 +661,15 @@ def test_cuj_managed_codex(self, cuj): tui.wait_for_task(codex_luna_task) tui.exit_normally() codex_luna_task.assert_completed(session, CODEX) - _assert_native_model(evidence, codex_luna_task, CODEX_LUNA) + evidence.assert_models(codex_luna_task, CODEX_LUNA_MODEL) _assert_inference_requests( recorder, luna_checkpoint, - CODEX_REQUEST_PATH, + INFERENCE_PATHS[CODEX], codex_luna_marker, run_id, CODEX, - CODEX_LUNA, + CODEX_LUNA_MODEL, ) final_profile = tomllib.loads((session.home / ".codex" / "ucode.config.toml").read_text()) assert final_profile["model_providers"]["Databricks"]["base_url"] == ( diff --git a/tests/e2e_cuj/test_cuj2_mps_mcp.py b/tests/e2e_cuj/test_cuj2_mps_mcp.py index 481c9ebd9..1dbc7bc22 100644 --- a/tests/e2e_cuj/test_cuj2_mps_mcp.py +++ b/tests/e2e_cuj/test_cuj2_mps_mcp.py @@ -11,11 +11,9 @@ from tests.integration.utils.evidence import ( agent_sessions, - assert_no_terminal_api_error, assistant_answers, is_child_session, ) -from tests.integration.utils.terminal import AgentTerminal from .base import BaseCujTest from .helpers.constants import ( @@ -25,11 +23,18 @@ FIXTURE_SUMMARY_SKILL_NAME, INFERENCE_PATHS, MODEL_PROVIDER_SERVICE_FIXTURES, + PROVIDER_HEADER, SANDBOX_MCP_SERVICE_NAME, WEB_SEARCH_MCP_SERVICE_NAME, ) +from .helpers.evidence import assert_served +from .helpers.mcp import registered_name +from .helpers.poll import poll +from .helpers.terminal import Terminal from .helpers.workspace import Workspace +CUJ_NAME = "CUJ 2 · MPS and named MCP" + WORKSPACE_URL = "https://dbc-0dcf95cf-e357.cloud.databricks.com" @@ -127,7 +132,7 @@ def _assert_configured_files(session, workspace_url: str) -> None: name, separator, value = line.partition(":") assert separator, line claude_headers[name.strip().casefold()] = value.strip() - assert claude_headers["databricks-model-provider-service"] == claude_mps, claude_headers + assert claude_headers[PROVIDER_HEADER] == claude_mps, claude_headers assert claude_env["ANTHROPIC_DEFAULT_HAIKU_MODEL"] == claude_model, claude_env assert claude_env["ANTHROPIC_BASE_URL"] == workspace_url.rstrip("/") + "/ai-gateway/anthropic" @@ -141,7 +146,7 @@ def _assert_configured_files(session, workspace_url: str) -> None: def _mcp_name_listed(output: str, name: str) -> bool: - dashed = name.replace(".", "-") + dashed = registered_name(name) return any(name in line or dashed in line for line in output.splitlines()) @@ -149,8 +154,18 @@ def _assert_generated_mcp_listings(session) -> None: for agent in (CLAUDE, CODEX): binary = shutil.which(agent, path=session.env["PATH"]) assert binary, f"Required agent is not on the isolated PATH: {agent}" - listing = session.run("mcp", "list", binary=binary, timeout=120) - output = f"{listing.stdout}\n{listing.stderr}" + + def listing(binary=binary): + result = session.run("mcp", "list", binary=binary, timeout=120) + return f"{result.stdout}\n{result.stderr}" + + # `claude mcp list` sometimes exits after its health-check banner without any rows. + output = poll( + listing, + timeout=60, + interval=5, + done=lambda text: _mcp_name_listed(text, SANDBOX_MCP_SERVICE_NAME), + ) assert _mcp_name_listed(output, SANDBOX_MCP_SERVICE_NAME), output assert not _mcp_name_listed(output, WEB_SEARCH_MCP_SERVICE_NAME), output @@ -189,12 +204,9 @@ def _agent_prompt(marker: str) -> str: def _assert_inference_request(recorder, request, *, provider: str, model: str, marker: str) -> None: - assert request.headers["databricks-model-provider-service"] == provider - assert request.payload["model"] == model + assert request.headers[PROVIDER_HEADER] == provider assert marker.encode() in request.body, request.payload - response = recorder.response_for(request, timeout=240) - assert response.status_code == 200, response.status_code - assert response.body, "Inference response was empty" + assert_served(recorder, request, model) class _Cuj2Base(BaseCujTest): @@ -256,16 +268,10 @@ def test_cuj_codex_inference(self, cuj): recorder.prepare_launch() marker = f"CUJ2-CODEX-{uuid.uuid4().hex}" checkpoint = recorder.checkpoint() - with AgentTerminal(session, CODEX, [str(session.binary), CODEX], "cuj2-codex") as tui: + with Terminal(session, "cuj2-codex", [CODEX]) as tui: tui.boot() tui.submit(_agent_prompt(marker)) - - def completed(screen): - assert_no_terminal_api_error(screen) - assert "Do you want to proceed?" not in screen, screen - return _task_complete(session, CODEX, marker) - - tui.wait_for(completed, "completed Codex task", timeout=240) + tui.wait_until(lambda: _task_complete(session, CODEX, marker), "completed Codex task") tui.exit_normally() request = recorder.expect_request( @@ -312,15 +318,10 @@ def test_cuj_claude_inference(self, cuj): recorder.prepare_launch() marker = f"CUJ2-CLAUDE-{uuid.uuid4().hex}" checkpoint = recorder.checkpoint() - with AgentTerminal(session, CLAUDE, [str(session.binary), CLAUDE], "cuj2-claude") as tui: + with Terminal(session, "cuj2-claude", [CLAUDE]) as tui: tui.boot() tui.submit(_agent_prompt(marker)) - - def completed(screen): - assert_no_terminal_api_error(screen) - return _task_complete(session, CLAUDE, marker) - - tui.wait_for(completed, "completed Claude task", timeout=240) + tui.wait_until(lambda: _task_complete(session, CLAUDE, marker), "completed Claude task") tui.exit_normally() request = recorder.expect_request( diff --git a/tests/e2e_cuj/test_cuj2_skills.py b/tests/e2e_cuj/test_cuj2_skills.py index fa38df4c9..2652ed039 100644 --- a/tests/e2e_cuj/test_cuj2_skills.py +++ b/tests/e2e_cuj/test_cuj2_skills.py @@ -24,6 +24,8 @@ ) from .test_cuj2_mps_mcp import _assert_cuj2_config, _Cuj2Base +CUJ_NAME = "CUJ 2 · Named skills" + OTHER_SKILLS = tuple(map(skill_name, (FIXTURE_NOTES_SKILL_NAME, FIXTURE_CATALOG_LOOKUP_SKILL_NAME))) diff --git a/tests/e2e_cuj/test_cuj3_mcp.py b/tests/e2e_cuj/test_cuj3_mcp.py new file mode 100644 index 000000000..ffddb17c8 --- /dev/null +++ b/tests/e2e_cuj/test_cuj3_mcp.py @@ -0,0 +1,281 @@ +"""Managed MCP registration, connection, and tool use in both real agent TUIs.""" + +import hashlib +import json +import re +import uuid + +import pytest + +from tests.integration.utils.evidence import ( + agent_sessions, + assistant_answers, + is_child_session, +) +from tests.integration.utils.mcp import open_claude_mcp_inventory, open_codex_mcp_inventory + +from .base import BaseCujTest +from .helpers.constants import CLAUDE, CODEX +from .helpers.mcp import registered_name +from .helpers.poll import poll +from .helpers.terminal import Terminal + +CUJ_NAME = "CUJ 3 · UC MCP discovery" + +pytestmark = [pytest.mark.managed, pytest.mark.mcp_registration, pytest.mark.workspace_isolated] + +SERVICES = { + "ug_e2e.tools.fixture_metadata": "describe_fixture", + "ug_e2e.tools.fixture_reader": "read_fixture", +} + + +class McpFixtureTask: + def __init__(self): + self.run_id = uuid.uuid4().hex + self.expected = { + tool: hashlib.sha256( + b"unity-gateway-cuj-fixture\0" + + tool.encode("ascii") + + b"\0" + + self.run_id.encode("ascii") + ).hexdigest() + for tool in SERVICES.values() + } + calls = "; ".join( + f"call {tool} on {registered_name(service)} with run_id={self.run_id}" + for service, tool in SERVICES.items() + ) + keys = " and ".join(f'"{tool}"' for tool in SERVICES.values()) + self.prompt = ( + f"Use these read-only MCP tools: {calls}. " + "Do not use shell commands, files, or subagents. Reply with only a JSON object " + f"whose keys are exactly the bare tool names {keys}, with no server name or " + "mcp__ prefix, each mapped to the text that tool returned, without code fences." + ) + # Claude may key a receipt by its native tool name and echo the raw tool result. + self.aliases = { + f"mcp__{registered_name(service)}__{tool}": tool for service, tool in SERVICES.items() + } + + def _receipts(self, observed): + if not isinstance(observed, dict): + return None + receipts = {} + for key, value in observed.items(): + tool = self.aliases.get(key, key) + if isinstance(value, str) and value.startswith("{"): + try: + value = json.loads(value) + except json.JSONDecodeError: + return None + if isinstance(value, dict) and value.keys() == {"result"}: + value = value["result"] + if tool in receipts: + return None + receipts[tool] = value + return receipts + + def answers(self, session, agent: str) -> list[str]: + return [ + answer + for path, records in agent_sessions(session, agent).items() + if not is_child_session(agent, path, records) + for answer in assistant_answers(agent, records) + ] + + def completed(self, session, agent: str) -> bool: + for answer in self.answers(session, agent): + # The TUI renders a fenced reply without its fence. + fenced = re.fullmatch(r"\s*```(?:json)?\n(.*)\n```\s*", answer, re.DOTALL) + try: + observed = json.loads(fenced.group(1) if fenced else answer) + except json.JSONDecodeError: + continue + if self._receipts(observed) == self.expected: + return True + return False + + +def _wait_for_mcp_task(tui, task): + approved_calls = set() + expected_tools = {(registered_name(service), tool) for service, tool in SERVICES.items()} + + def approve_expected_codex_call(screen): + if not ( + "Allow for this session" in screen + and "Always allow" in screen + and "Cancel" in screen + and re.search(r"(?m)^[ \t]*[›❯>][ \t]*[1-4]\.", screen) + ): + return False + questions = list(re.finditer(r"(?m)^[ \t]*(Allow\b[^\n?]+\?)[ \t]*$", screen)) + assert questions, "Unexpected approval menu; inspect the actual command:\n" + screen + question = questions[-1].group(1) + approval = re.fullmatch( + r'Allow the (?P\S+) MCP server to run tool "(?P[^"]+)"\?', + question, + ) + assert tui.agent == CODEX and approval, ( + "Unexpected Codex permission request; inspect the actual command:\n" + screen + ) + tail = screen[questions[-1].end() :] + run_match = re.search(r"(?m)^[ \t]*run_id:[ \t]*(\S+)[ \t]*$", tail) + selected = re.search( + r"(?m)^[ \t]*[›❯>][ \t]*([1-4])\.[ \t]*(Allow(?: for this session)?|Always allow|Cancel)\b", + tail, + ) + assert run_match and selected, ( + "Unexpected Codex approval menu; inspect the actual command:\n" + screen + ) + server = approval.group("server") + tool = approval.group("tool") + run_id = run_match.group(1) + assert (server, tool) in expected_tools and run_id == task.run_id, ( + "Unexpected MCP permission request; inspect the actual command:\n" + screen + ) + assert selected.group(1) == "1" and selected.group(2) == "Allow", ( + "Codex MCP approval must remain a per-call Allow; inspect the actual menu:\n" + screen + ) + call = (server, tool, run_id) + if call in approved_calls: + return False + tui.send("\r", f"allow one MCP call: {server}.{tool}") + approved_calls.add(call) + return True + + try: + tui.wait_until( + lambda: task.completed(tui.session, tui.agent), + "a completed assistant answer with both verified MCP receipts", + rejected=("Do you want to proceed?", "Do you want to allow"), + on_screen=approve_expected_codex_call, + ) + except AssertionError as error: + answers = task.answers(tui.session, tui.agent) + raise AssertionError(f"{error}\nNative assistant answers: {answers!r}") from None + + +def _fixture_server_rows(output): + return { + service: [ + " ".join(line.split()) + for line in output.splitlines() + if line.split()[:1] == [registered_name(service)] + ] + for service in SERVICES + } + + +def _expected_fixture_server_rows(): + return { + service: [f"{registered_name(service)} {service} (managed) {CLAUDE}, {CODEX} connected"] + for service in SERVICES + } + + +def _poll_mcp_list(session): + """Return the first `ug mcp list` output whose fixture rows match, else the last one. + + Claude's single-shot health probe can report a cold-starting server as failed. + """ + + def listing(): + result = session.run("mcp", "list", timeout=120) + return f"{result.stdout}\n{result.stderr}" + + return poll( + listing, + timeout=90, + interval=5, + done=lambda output: _fixture_server_rows(output) == _expected_fixture_server_rows(), + ) + + +@pytest.fixture(autouse=True) +def _revert_mcp_test_state(cuj): + yield + session, _, _ = cuj + session.revert_machine_wide( + "mcp-registration-cleanup-revert", + "MCP registration CUJ teardown left machine-wide agent settings", + ) + + +class TestMcpRegistration(BaseCujTest): + WORKSPACE_URL = "https://dbc-bbdd5508-648e.cloud.databricks.com" + + def test_mcp_list_reports_managed_fixture_servers(self, cuj): + """Scenario: configure both agents and inspect unscoped `ug mcp list` output. + + Expected: both managed fixture rows list Claude/Codex and aggregate `connected` status; + no `ug_e2e.other_tools` server appears. Codex connectivity is verified by its TUI journey. + """ + session, workspace, _recorder = cuj + session.configure( + [ + "configure", + "--agents", + f"{CLAUDE},{CODEX}", + "--workspace", + workspace.url, + "--skip-upgrade", + ] + ) + output = _poll_mcp_list(session) + assert _fixture_server_rows(output) == _expected_fixture_server_rows(), output + assert "ug_e2e.other_tools" not in output, output + assert "ug_e2e-other_tools" not in output, output + + @pytest.mark.claude + @pytest.mark.tui + def test_mcp_registration_claude_servers_connect_and_work(self, cuj): + """Scenario: configure UG, inspect Claude's MCP menu, then call both fixture tools. + + Expected: both scoped servers are registered, connected, and expose their tools; + no ug_e2e.other_tools server or decoy tool appears in the complete inventory; + a parent assistant answer contains the correct receipts for a fresh run ID. + """ + session, workspace, _recorder = cuj + session.configure(["configure", "--workspace", workspace.url, "--skip-upgrade"]) + task = McpFixtureTask() + allowed_tools = ",".join( + f"mcp__{registered_name(service)}__{tool}" for service, tool in SERVICES.items() + ) + with Terminal( + session, + "mcp-registration-claude", + [CLAUDE, "--", "--allowedTools", allowed_tools], + ) as tui: + tui.boot() + inventory = open_claude_mcp_inventory(tui, SERVICES) + assert "ug_e2e-other_tools" not in inventory, inventory + assert "decoy_status" not in inventory, inventory + tui.submit(task.prompt) + _wait_for_mcp_task(tui, task) + tui.exit_normally() + assert task.completed(session, CLAUDE) + session.assert_not_routed() + + @pytest.mark.codex + @pytest.mark.tui + def test_mcp_registration_codex_servers_connect_and_work(self, cuj): + """Scenario: configure UG, inspect Codex's MCP inventory, then call both fixture tools. + + Expected: both scoped servers are registered, connected, and expose their tools; + no ug_e2e.other_tools server or decoy tool appears in the complete inventory; + a completed parent turn returns the correct receipts for a fresh run ID. + """ + session, workspace, _recorder = cuj + session.configure(["configure", "--workspace", workspace.url, "--skip-upgrade"]) + task = McpFixtureTask() + with Terminal(session, "mcp-registration-codex", [CODEX]) as tui: + tui.boot() + inventory = open_codex_mcp_inventory(tui, SERVICES) + assert "ug_e2e-other_tools" not in inventory, inventory + assert "decoy_status" not in inventory, inventory + tui.submit(task.prompt) + _wait_for_mcp_task(tui, task) + tui.exit_normally() + assert task.completed(session, CODEX) + session.assert_not_routed() diff --git a/tests/e2e_cuj/test_catalog_discovery.py b/tests/e2e_cuj/test_cuj3_models.py similarity index 90% rename from tests/e2e_cuj/test_catalog_discovery.py rename to tests/e2e_cuj/test_cuj3_models.py index f77c8c236..c09020179 100644 --- a/tests/e2e_cuj/test_catalog_discovery.py +++ b/tests/e2e_cuj/test_cuj3_models.py @@ -12,7 +12,6 @@ parse_anthropic_provider_page, parse_codex_provider_catalog, ) -from tests.integration.utils.terminal import AgentTerminal, TerminalProcess from .base import BaseCujTest from .catalog_discovery_expectations import ( @@ -26,10 +25,12 @@ MODEL_SCHEMA, OTHER_MODEL_SCHEMA, ) -from .helpers.constants import CLAUDE, CODEX, INFERENCE_PATHS, MANAGED_PATHS -from .helpers.evidence import SessionEvidence, message_text +from .helpers.constants import CLAUDE, CODEX, INFERENCE_PATHS +from .helpers.evidence import SessionEvidence, assert_served, claude_file_task, message_text from .helpers.terminal import Terminal +CUJ_NAME = "CUJ 3 · UC model discovery" + pytestmark = [pytest.mark.managed, pytest.mark.catalog_discovery, pytest.mark.workspace_isolated] @@ -68,14 +69,6 @@ def _catalog_display_names(workspace, agent, schema): raise AssertionError("Anthropic catalog exceeded 20 pages") -def _claude_file_task(session): - task = FileTask(session) - task.prompt = ( - f"Use the Read tool to read {session.cwd / task.filename}. Reply with only its contents." - ) - return task - - def _assert_claude_headless_model(result, expected): final = None for line in result.stdout.splitlines(): @@ -129,33 +122,15 @@ def _assert_inference_evidence(recorder, checkpoint, agent, task, expected): ] assert task_requests, "No inference request contained the submitted task prompt" for request in task_requests: - assert request.payload["model"] == expected_wire_model, request.payload - response = recorder.response_for(request, timeout=240) - assert response.status_code == 200, { - "status": response.status_code, - "model": request.payload["model"], - "output_config": request.payload.get("output_config"), - } - assert response.body, "Inference response was empty" + assert_served(recorder, request, expected_wire_model) @pytest.fixture(autouse=True) def _revert_catalog_test_state(cuj): yield session, _, _ = cuj - 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"], - "catalog-discovery-cleanup-revert", - ) as terminal: - terminal.finish() - assert not any(path.exists() for path in MANAGED_PATHS), ( - "Catalog CUJ teardown left machine-wide agent settings" + session.revert_machine_wide( + "catalog-discovery-cleanup-revert", "Catalog CUJ teardown left machine-wide agent settings" ) @@ -181,7 +156,7 @@ def test_catalog_discovery_claude_picker_preserves_default(self, cuj): assert set(decoy_catalog) == {claude.discovery_model_id(CLAUDE_DECOY)}, decoy_catalog default_display_name = parent_catalog[claude.discovery_model_id(CLAUDE_DEFAULT)] - task = _claude_file_task(session) + task = claude_file_task(session) evidence = SessionEvidence(session.home, CLAUDE) checkpoint = recorder.checkpoint() recorder.prepare_launch() @@ -261,17 +236,15 @@ def test_catalog_discovery_bare_ug_uses_claude_default(self, cuj): default_display_name = _catalog_display_names(workspace, CLAUDE, MODEL_SCHEMA)[ claude.discovery_model_id(CLAUDE_DEFAULT) ] - task = _claude_file_task(session) + task = claude_file_task(session) evidence = SessionEvidence(session.home, CLAUDE) checkpoint = recorder.checkpoint() recorder.prepare_launch() - with AgentTerminal( - session, CLAUDE, [str(session.binary)], "catalog-discovery-bare-ug-default" - ) as tui: + with Terminal(session, "catalog-discovery-bare-ug-default", [], agent=CLAUDE) as tui: tui.boot() assert f"{default_display_name} · API Usage Billing" in tui.visible, tui.visible tui.submit(task.prompt) - Terminal.task(tui, evidence, task) + tui.task(evidence, task) tui.exit_normally() task.assert_completed(session, CLAUDE) _assert_inference_evidence(recorder, checkpoint, CLAUDE, task, CLAUDE_DEFAULT) @@ -289,7 +262,7 @@ def test_catalog_discovery_claude_tui_uses_default(self, cuj): default_display_name = _catalog_display_names(workspace, CLAUDE, MODEL_SCHEMA)[ claude.discovery_model_id(CLAUDE_DEFAULT) ] - task = _claude_file_task(session) + task = claude_file_task(session) evidence = SessionEvidence(session.home, CLAUDE) checkpoint = recorder.checkpoint() recorder.prepare_launch() @@ -333,7 +306,7 @@ def test_catalog_discovery_claude_headless_uses_default(self, cuj): """ session, _, recorder = cuj recorder.configure_session(session, ["configure", "--skip-upgrade"]) - task = _claude_file_task(session) + task = claude_file_task(session) checkpoint = recorder.checkpoint() recorder.prepare_launch() result = session.run( @@ -379,7 +352,7 @@ def test_catalog_discovery_claude_explicit_model_completes_task(self, cuj, model """ session, _, recorder = cuj recorder.configure_session(session, ["configure", "--skip-upgrade"]) - task = _claude_file_task(session) + task = claude_file_task(session) checkpoint = recorder.checkpoint() recorder.prepare_launch() result = session.run( diff --git a/tests/e2e_cuj/test_cuj3_skills.py b/tests/e2e_cuj/test_cuj3_skills.py index b071321be..f1d1c88ad 100644 --- a/tests/e2e_cuj/test_cuj3_skills.py +++ b/tests/e2e_cuj/test_cuj3_skills.py @@ -25,6 +25,8 @@ skills_view, ) +CUJ_NAME = "CUJ 3 · UC skill discovery" + SCHEMA_SKILLS = tuple( map( skill_name, diff --git a/tests/e2e_cuj/test_cuj_smart_routing.py b/tests/e2e_cuj/test_cuj4_smart_routing.py similarity index 95% rename from tests/e2e_cuj/test_cuj_smart_routing.py rename to tests/e2e_cuj/test_cuj4_smart_routing.py index 85936f076..2c6a56423 100644 --- a/tests/e2e_cuj/test_cuj_smart_routing.py +++ b/tests/e2e_cuj/test_cuj4_smart_routing.py @@ -8,9 +8,17 @@ from .base import BaseCujTest from .helpers.constants import CLAUDE, CODEX, INFERENCE_PATHS, CodingAgent -from .helpers.evidence import SessionEvidence, SessionObservation, canonical_model +from .helpers.evidence import ( + SessionEvidence, + SessionObservation, + canonical_model, + claude_file_task, +) from .helpers.terminal import Terminal from .helpers.tui_request_recorder import RecordedRequest, RecordedResponse +from .helpers.workspace import Workspace + +CUJ_NAME = "CUJ 4 · Smart routing" ROUTING_PATH = "/ai-gateway/routing/v1/routes:select" AGENTS = (CLAUDE, CODEX) @@ -45,6 +53,12 @@ class SmartRoutingSessionResults: with_model_override: dict[str, SessionCase] +def _file_task(session, agent): + task = claude_file_task(session) if agent == CLAUDE else FileTask(session) + task.prompt += " Do not delegate." + return task + + def _run_session(session, recorder, agent, task, launch_args): evidence = SessionEvidence(session.home, agent) checkpoint = recorder.checkpoint() @@ -62,7 +76,7 @@ def _assert_published_config_matches_expectations(published): assert published["default_agent"] == CodingAgent.CLAUDE_CODE entries = published["enabled_agents"] assert len(entries) == 2 - configs = {entry["agent"]: entry["config"] for entry in entries} + configs = Workspace.agent_configs(published) assert set(configs) == {CodingAgent.CLAUDE_CODE, CodingAgent.CODEX} agent_configs = { CLAUDE: configs[CodingAgent.CLAUDE_CODE], @@ -103,8 +117,7 @@ def run_smart_routing_journeys(cuj) -> SmartRoutingSessionResults: no_model_override, with_model_override = {}, {} for agent in AGENTS: - task = FileTask(session) - task.prompt += " Do not delegate." + task = _file_task(session, agent) observation, requests = _run_session(session, recorder, agent, task, (agent,)) route_request = next( request @@ -129,8 +142,7 @@ def run_smart_routing_journeys(cuj) -> SmartRoutingSessionResults: inference_response=recorder.response_for(inference_request), ) - task = FileTask(session) - task.prompt += " Do not delegate." + task = _file_task(session, agent) launch_args = (agent, "--model", overrides[agent]) observation, requests = _run_session(session, recorder, agent, task, launch_args) inference_request = next( diff --git a/tests/e2e_cuj/test_ug_budget_defaults.py b/tests/e2e_cuj/test_cuj5_budget_defaults.py similarity index 93% rename from tests/e2e_cuj/test_ug_budget_defaults.py rename to tests/e2e_cuj/test_cuj5_budget_defaults.py index 54b6a63fe..de3357317 100644 --- a/tests/e2e_cuj/test_ug_budget_defaults.py +++ b/tests/e2e_cuj/test_cuj5_budget_defaults.py @@ -8,6 +8,9 @@ from .base import BaseCujTest from .helpers.constants import CLAUDE, CODEX, CodingAgent +from .helpers.workspace import Workspace + +CUJ_NAME = "CUJ 5 · Budget-driven defaults" class _BudgetDefaultsBase(BaseCujTest): @@ -34,13 +37,10 @@ def test_bare_ug_uses_claude_default_below_budget_tier(self, cuj): """ session, workspace, _ = cuj config_path = "/api/ai-gateway/v2/coding-agent-configs" - payload = workspace.client.api_client.do("GET", path=config_path) - config = payload["coding_agent_configs"][0] + config = workspace.config() assert config["default_agent"] == CodingAgent.CLAUDE_CODE, config - claude = next( - agent for agent in config["enabled_agents"] if agent["agent"] == CodingAgent.CLAUDE_CODE - ) - claude_defaults = claude["config"]["default_models"] + claude = Workspace.agent_configs(config)[CodingAgent.CLAUDE_CODE] + claude_defaults = claude["default_models"] assert claude_defaults["default_model"] == self.SONNET_MODEL, config assert config["smart_defaults"]["budget_id"] == self.BUDGET_ID, config assert config["smart_defaults"]["tiers"] == [ @@ -102,15 +102,10 @@ def test_bare_ug_shows_luna_recommendation_and_starts_codex(self, cuj): """ session, workspace, _ = cuj config_path = "/api/ai-gateway/v2/coding-agent-configs" - payload = workspace.client.api_client.do("GET", path=config_path) - config = payload["coding_agent_configs"][0] + config = workspace.config() assert config["default_agent"] == CodingAgent.CLAUDE_CODE, config - codex_agent = next( - agent for agent in config["enabled_agents"] if agent["agent"] == CodingAgent.CODEX - ) - live_default_model = ( - codex_agent.get("config", {}).get("default_models", {}).get("default_model") - ) + codex_agent = Workspace.agent_configs(config)[CodingAgent.CODEX] + live_default_model = codex_agent.get("default_models", {}).get("default_model") assert live_default_model == self.SOL_MODEL, ( f"Live Codex config model mismatch: expected {self.SOL_MODEL!r}, " f"got {live_default_model!r}" diff --git a/tests/e2e_cuj/test_cuj_repeated_config.py b/tests/e2e_cuj/test_cuj6_repeated_config.py similarity index 89% rename from tests/e2e_cuj/test_cuj_repeated_config.py rename to tests/e2e_cuj/test_cuj6_repeated_config.py index 7368a9132..3953c6a7a 100644 --- a/tests/e2e_cuj/test_cuj_repeated_config.py +++ b/tests/e2e_cuj/test_cuj6_repeated_config.py @@ -12,7 +12,6 @@ import hashlib import json -import os import re import shutil import subprocess @@ -24,45 +23,48 @@ from pathlib import Path import pytest -from databricks.sdk import WorkspaceClient -from tests.integration.utils.evidence import FileTask, assert_no_terminal_api_error +from tests.integration.utils.evidence import FileTask from tests.integration.utils.sql import query_count, resolve_trace_table, resolve_warehouse_id -from tests.integration.utils.terminal import AgentTerminal, TerminalProcess +from tests.integration.utils.terminal import TerminalProcess from ucode.agents.claude import CLAUDE_OTEL_TRACE_ENV_KEYS -from .base import BaseCujTest +from .base import BaseCujTest, bearer, make_workspace_client from .helpers.constants import ( + AGENT_HEADER, BEDROCK_PROVIDER_SERVICE_FIXTURE, CLAUDE, + CLAUDE_HAIKU_MODEL, + CLAUDE_OPUS_MODEL, + CLAUDE_SONNET_MODEL, CODEX, + CODEX_LUNA_MODEL, + CODEX_SOL_MODEL, FIXTURE_READER_MCP_SERVICE_NAME, FIXTURE_SUMMARY_SKILL_NAME, INFERENCE_PATHS, MANAGED_PATHS, + PROVIDER_HEADER, + RUN_HEADER, UC_MODEL_LOCATION_FIXTURE, CodingAgent, ) -from .helpers.evidence import SessionEvidence, canonical_model +from .helpers.evidence import SessionEvidence, assert_models, canonical_model +from .helpers.mcp import registered_name +from .helpers.poll import poll from .helpers.skills import SKILL_ROOTS, skill_name, skills_view +from .helpers.terminal import Terminal from .helpers.tui_request_recorder import TuiRequestRecorder from .helpers.workspace import Workspace -CLAUDE_STATIC_MODELS = ( - "system.ai.claude-opus-4-8", - "system.ai.claude-sonnet-4-6", - "system.ai.claude-haiku-4-5", -) -CLAUDE_DEFAULT = "system.ai.claude-sonnet-4-6" -CODEX_SOL = "system.ai.gpt-5-6-sol" -CODEX_LUNA = "system.ai.gpt-5-6-luna" +CUJ_NAME = "CUJ 6 · Repeated config and removal" + +CLAUDE_STATIC_MODELS = (CLAUDE_OPUS_MODEL, CLAUDE_SONNET_MODEL, CLAUDE_HAIKU_MODEL) +CLAUDE_DEFAULT = CLAUDE_SONNET_MODEL CLAUDE_PROVIDER, CLAUDE_PROVIDER_MODEL = BEDROCK_PROVIDER_SERVICE_FIXTURE CODEX_MODEL_LOCATION, CODEX_UC_MODEL = UC_MODEL_LOCATION_FIXTURE -MANAGED_MCP_ENTRY = FIXTURE_READER_MCP_SERVICE_NAME.replace(".", "-") +MANAGED_MCP_ENTRY = registered_name(FIXTURE_READER_MCP_SERVICE_NAME) CODEX_OS_MANAGED_CONFIG = MANAGED_PATHS[1] -RUN_HEADER = "x-ug-e2e-run" -AGENT_HEADER = "x-ug-e2e-agent" -PROVIDER_HEADER = "databricks-model-provider-service" PARENT_SCHEMA_HEADER = "databricks-model-service-parent-schema" CONFIGURE_ARGS = ["configure", "--skip-upgrade", "--disable-databricks-ai-tools"] TRACE_DEADLINE_SECONDS = 300 @@ -106,10 +108,6 @@ def model_leaf(model: str) -> str: return re.sub(r"(-\d{8})?(-v\d+:\d+)?$", "", name) -def agent_config(config: dict, agent: CodingAgent) -> dict: - return next(entry["config"] for entry in config["enabled_agents"] if entry["agent"] == agent) - - def tracing_enabled(agent_settings: dict) -> bool: return bool(agent_settings.get("tracing", {}).get("enabled")) @@ -155,16 +153,6 @@ def skill_bundles(files: "AgentFiles", root: str) -> set[str]: } -def workspace_client(url: str) -> WorkspaceClient: - """The base class authenticates only WORKSPACE_URL; phases B and C need their own clients.""" - return WorkspaceClient( - host=url, - client_id=os.environ["UG_CUJ_SP_CLIENT_ID"], - client_secret=os.environ["UG_CUJ_SP_CLIENT_SECRET"], - auth_type="oauth-m2m", - ) - - @dataclass(frozen=True) class AgentFiles: """ug-written agent settings and MCP registrations in the shared home.""" @@ -252,24 +240,20 @@ def __init__(self, session, workspaces, published, recorders): self.published = published self.recorders = recorders self.run_id = f"cuj6-{uuid.uuid4().hex[:12]}" - bearer = self.bearer("a") + a_bearer = bearer(workspaces["a"].client) # All phase workspaces share one metastore and trace table, so A queries them all. self.trace_url = workspaces["a"].url - self.warehouse_id = resolve_warehouse_id(self.trace_url, bearer) - self.trace_table = resolve_trace_table(self.trace_url, bearer) + self.warehouse_id = resolve_warehouse_id(self.trace_url, a_bearer) + self.trace_table = resolve_trace_table(self.trace_url, a_bearer) self.seeded = self.seed_user_state() @property def home(self) -> Path: return self.session.home - def bearer(self, phase: str) -> str: - # Each workspace issues its own M2M token; refreshed per call because the journey is long. - headers = self.workspaces[phase].client.config.authenticate() - return headers["Authorization"].removeprefix("Bearer ") - def enter(self, phase: str) -> None: - self.session.env["DATABRICKS_BEARER"] = self.bearer(phase) + # Each workspace issues its own M2M token; minted per entry because the journey is long. + self.session.install_bearer(bearer(self.workspaces[phase].client)) def configure(self, phase: str) -> None: self.enter(phase) @@ -404,24 +388,16 @@ def skills_view(self, phase: str, agent: str) -> str: ) def bare_ug_task(self, phase: str) -> AgentTask: - """Bare `ug` in a real terminal; helpers.terminal.Terminal only drives `ug `.""" + """Bare `ug` in a real terminal, which launches the default agent Codex.""" recorder = self.recorders[phase] recorder.prepare_launch() evidence = SessionEvidence(self.home, CODEX) task = FileTask(self.session) checkpoint = recorder.checkpoint() - with AgentTerminal( - self.session, CODEX, [str(self.session.binary)], f"phase-{phase}-bare-ug" - ) as tui: + with Terminal(self.session, f"phase-{phase}-bare-ug", [], agent=CODEX) as tui: tui.boot() tui.submit(task.prompt) - - def completed(screen): - assert_no_terminal_api_error(screen) - assert "Do you want to proceed?" not in screen, screen - return evidence.completed(task) is not None - - tui.wait_for(completed, "completed native file task", timeout=240) + tui.task(evidence, task) tui.exit_normally() requests = self.inference_requests(phase, CODEX, checkpoint) assert requests, "The recorder saw no bare ug Codex inference request" @@ -432,7 +408,7 @@ def trace_count(self, marker: str, *, resource: bool) -> int: attributes = "resource.attributes" if resource else "attributes" return query_count( self.trace_url, - self.bearer("a"), + bearer(self.workspaces["a"].client), self.warehouse_id, f"SELECT COUNT(*) FROM {self.trace_table} " "WHERE time > current_timestamp() - INTERVAL 2 HOURS " @@ -441,12 +417,11 @@ def trace_count(self, marker: str, *, resource: bool) -> int: ) def wait_for_codex_span(self, marker: str) -> int: - deadline = time.monotonic() + TRACE_DEADLINE_SECONDS - while (count := self.trace_count(marker, resource=False)) == 0: - if time.monotonic() > deadline: - break - time.sleep(15) - return count + return poll( + lambda: self.trace_count(marker, resource=False), + timeout=TRACE_DEADLINE_SECONDS, + interval=15, + ) def seed_user_state(self) -> SeededUserState: """Ordinary user-owned prefs, MCP registration, and skill; no ug state.""" @@ -509,8 +484,12 @@ def journey(self, cuj): session, workspace_a, recorder_a = cuj workspaces = { "a": workspace_a, - "b": Workspace(workspace_client(self.PHASE_B_URL)), - "c": Workspace(workspace_client(self.PHASE_C_URL)), + "b": Workspace( + make_workspace_client(self.PHASE_B_URL, self.CLIENT_ID_ENV, self.CLIENT_SECRET_ENV) + ), + "c": Workspace( + make_workspace_client(self.PHASE_C_URL, self.CLIENT_ID_ENV, self.CLIENT_SECRET_ENV) + ), } published = {phase: workspace.config() for phase, workspace in workspaces.items()} with ( @@ -606,11 +585,12 @@ def phase_c(self, journey, phase_b): def test_phase_a_workspace_publishes_plan_config(self, journey): config = journey.published["a"] - claude = agent_config(config, CodingAgent.CLAUDE_CODE) - codex = agent_config(config, CodingAgent.CODEX) + configs = Workspace.agent_configs(config) + claude = configs[CodingAgent.CLAUDE_CODE] + codex = configs[CodingAgent.CODEX] assert config["default_agent"] == CodingAgent.CLAUDE_CODE assert claude["models"]["model_services"] == list(CLAUDE_STATIC_MODELS) - assert codex["models"]["model_services"] == [CODEX_SOL, CODEX_LUNA] + assert codex["models"]["model_services"] == [CODEX_SOL_MODEL, CODEX_LUNA_MODEL] assert tracing_enabled(claude) and tracing_enabled(codex) assert config["mcp_servers"]["names"] == [FIXTURE_READER_MCP_SERVICE_NAME] assert config["skills"]["names"] == [FIXTURE_SUMMARY_SKILL_NAME] @@ -645,18 +625,19 @@ def test_phase_a_codex_sends_codex_headers(self, phase_a): def test_phase_a_claude_runs_its_default(self, phase_a): assert CLAUDE_DEFAULT in {canonical_model(m) for m in phase_a.claude.request_models} - assert {canonical_model(m) for m in phase_a.claude.turn_models} == {CLAUDE_DEFAULT} + assert_models(phase_a.claude.turn_models, CLAUDE_DEFAULT) def test_phase_a_codex_runs_its_default(self, phase_a): - assert {canonical_model(m) for m in phase_a.codex.request_models} == {CODEX_SOL} - assert {canonical_model(m) for m in phase_a.codex.turn_models} == {CODEX_SOL} + assert_models(phase_a.codex.request_models, CODEX_SOL_MODEL) + assert_models(phase_a.codex.turn_models, CODEX_SOL_MODEL) # Phase B: Claude on an MPS, Codex on UC discovery; headers, MCP, and tracing changed. def test_phase_b_workspace_publishes_plan_config(self, journey): config = journey.published["b"] - claude = agent_config(config, CodingAgent.CLAUDE_CODE) - codex = agent_config(config, CodingAgent.CODEX) + configs = Workspace.agent_configs(config) + claude = configs[CodingAgent.CLAUDE_CODE] + codex = configs[CodingAgent.CODEX] assert config["default_agent"] == CodingAgent.CODEX assert claude["models"] == {"model_provider_service": CLAUDE_PROVIDER} assert claude["default_models"]["default_model"] == CLAUDE_PROVIDER_MODEL @@ -750,9 +731,9 @@ def test_phase_b_traces_codex_but_not_claude(self, phase_b): def test_phase_c_workspace_publishes_plan_config(self, journey): config = journey.published["c"] assert [entry["agent"] for entry in config["enabled_agents"]] == [CodingAgent.CODEX] - codex = agent_config(config, CodingAgent.CODEX) - assert codex["models"] == {"model_services": [CODEX_LUNA]} - assert codex["default_models"]["default_model"] == CODEX_LUNA + codex = Workspace.agent_configs(config)[CodingAgent.CODEX] + assert codex["models"] == {"model_services": [CODEX_LUNA_MODEL]} + assert codex["default_models"]["default_model"] == CODEX_LUNA_MODEL assert not tracing_enabled(codex) and not codex.get("http_headers") assert not config.get("skills") @@ -773,13 +754,16 @@ def test_phase_c_codex_drops_phase_b_headers(self, phase_c): assert not {RUN_HEADER, AGENT_HEADER, PARENT_SCHEMA_HEADER} & set(request.headers) def test_phase_c_codex_picker_is_only_luna(self, phase_c): - assert phase_c.picker == (CODEX_LUNA,) + assert phase_c.picker == (CODEX_LUNA_MODEL,) def test_phase_c_codex_and_bare_ug_run_luna(self, phase_c): - assert {canonical_model(m) for m in phase_c.codex.request_models} == {CODEX_LUNA} - assert {canonical_model(m) for m in phase_c.codex.turn_models} == {CODEX_LUNA} - assert {canonical_model(m) for m in phase_c.bare_ug.request_models} == {CODEX_LUNA} - assert {canonical_model(m) for m in phase_c.bare_ug.turn_models} == {CODEX_LUNA} + for models in ( + phase_c.codex.request_models, + phase_c.codex.turn_models, + phase_c.bare_ug.request_models, + phase_c.bare_ug.turn_models, + ): + assert_models(models, CODEX_LUNA_MODEL) def test_phase_c_claude_is_rejected_before_inference(self, phase_c): assert phase_c.disabled_claude_returncode != 0 diff --git a/tests/e2e_cuj/test_mcp_registration.py b/tests/e2e_cuj/test_mcp_registration.py deleted file mode 100644 index 7188004d5..000000000 --- a/tests/e2e_cuj/test_mcp_registration.py +++ /dev/null @@ -1,231 +0,0 @@ -"""Managed MCP registration, connection, and tool use in both real agent TUIs.""" - -import hashlib -import json -import re -import uuid - -import pytest - -from tests.integration.utils.evidence import ( - agent_sessions, - assert_no_terminal_api_error, - assistant_answers, - is_child_session, -) -from tests.integration.utils.mcp import open_claude_mcp_inventory, open_codex_mcp_inventory -from tests.integration.utils.terminal import TerminalProcess - -from .base import BaseCujTest -from .helpers.constants import CLAUDE, CODEX, MANAGED_PATHS -from .helpers.terminal import Terminal - -pytestmark = [pytest.mark.managed, pytest.mark.mcp_registration, pytest.mark.workspace_isolated] - -SERVICES = { - "ug_e2e.tools.fixture_metadata": "describe_fixture", - "ug_e2e.tools.fixture_reader": "read_fixture", -} - - -class McpFixtureTask: - def __init__(self): - self.run_id = uuid.uuid4().hex - self.expected = { - tool: hashlib.sha256( - b"unity-gateway-cuj-fixture\0" - + tool.encode("ascii") - + b"\0" - + self.run_id.encode("ascii") - ).hexdigest() - for tool in SERVICES.values() - } - calls = "; ".join( - f"call {tool} on {service.replace('.', '-')} with run_id={self.run_id}" - for service, tool in SERVICES.items() - ) - self.prompt = ( - f"Use these read-only MCP tools: {calls}. " - "Do not use shell commands, files, or subagents. Reply with only a JSON object " - "mapping each tool name to the text it returned, without code fences." - ) - - def completed(self, session, agent: str) -> bool: - for path, records in agent_sessions(session, agent).items(): - if is_child_session(agent, path, records): - continue - for answer in assistant_answers(agent, records): - try: - observed = json.loads(answer) - except json.JSONDecodeError: - continue - if observed == self.expected: - return True - return False - - -def _wait_for_mcp_task(tui, task): - approved_calls = set() - expected_tools = {(service.replace(".", "-"), tool) for service, tool in SERVICES.items()} - - def completed(screen): - assert_no_terminal_api_error(screen) - assert not any( - prompt in screen for prompt in ("Do you want to proceed?", "Do you want to allow") - ), "Unexpected permission request; inspect the actual command:\n" + screen - - if ( - "Allow for this session" in screen - and "Always allow" in screen - and "Cancel" in screen - and re.search(r"(?m)^[ \t]*[›❯>][ \t]*[1-4]\.", screen) - ): - questions = list(re.finditer(r"(?m)^[ \t]*(Allow\b[^\n?]+\?)[ \t]*$", screen)) - assert questions, "Unexpected approval menu; inspect the actual command:\n" + screen - question = questions[-1].group(1) - approval = re.fullmatch( - r'Allow the (?P\S+) MCP server to run tool "(?P[^"]+)"\?', - question, - ) - assert tui.agent == CODEX and approval, ( - "Unexpected Codex permission request; inspect the actual command:\n" + screen - ) - tail = screen[questions[-1].end() :] - run_match = re.search(r"(?m)^[ \t]*run_id:[ \t]*(\S+)[ \t]*$", tail) - selected = re.search( - r"(?m)^[ \t]*[›❯>][ \t]*([1-4])\.[ \t]*(Allow(?: for this session)?|Always allow|Cancel)\b", - tail, - ) - assert run_match and selected, ( - "Unexpected Codex approval menu; inspect the actual command:\n" + screen - ) - server = approval.group("server") - tool = approval.group("tool") - run_id = run_match.group(1) - assert (server, tool) in expected_tools and run_id == task.run_id, ( - "Unexpected MCP permission request; inspect the actual command:\n" + screen - ) - assert selected.group(1) == "1" and selected.group(2) == "Allow", ( - "Codex MCP approval must remain a per-call Allow; inspect the actual menu:\n" - + screen - ) - call = (server, tool, run_id) - if call not in approved_calls: - tui.send("\r", f"allow one MCP call: {server}.{tool}") - approved_calls.add(call) - return False - return task.completed(tui.session, tui.agent) - - tui.wait_for( - completed, - "a completed assistant answer with both verified MCP receipts", - timeout=240, - ) - - -@pytest.fixture(autouse=True) -def _revert_mcp_test_state(cuj): - yield - session, _, _ = cuj - 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"], - "mcp-registration-cleanup-revert", - ) as terminal: - terminal.finish() - assert not any(path.exists() for path in MANAGED_PATHS), ( - "MCP registration CUJ teardown left machine-wide agent settings" - ) - - -class TestMcpRegistration(BaseCujTest): - WORKSPACE_URL = "https://dbc-bbdd5508-648e.cloud.databricks.com" - - def test_mcp_list_reports_managed_fixture_servers(self, cuj): - """Scenario: configure both agents and inspect unscoped `ug mcp list` output. - - Expected: both managed fixture rows list Claude/Codex and aggregate `connected` status; - no `ug_e2e.other_tools` server appears. Codex connectivity is verified by its TUI journey. - """ - session, workspace, _recorder = cuj - session.configure( - [ - "configure", - "--agents", - f"{CLAUDE},{CODEX}", - "--workspace", - workspace.url, - "--skip-upgrade", - ] - ) - listing = session.run("mcp", "list", timeout=120) - output = f"{listing.stdout}\n{listing.stderr}" - for service in SERVICES: - server_name = service.replace(".", "-") - rows = [ - " ".join(line.split()) - for line in output.splitlines() - if line.split()[:1] == [server_name] - ] - expected_row = f"{server_name} {service} (managed) {CLAUDE}, {CODEX} connected" - assert rows == [expected_row], output - assert "ug_e2e.other_tools" not in output, output - assert "ug_e2e-other_tools" not in output, output - - @pytest.mark.claude - @pytest.mark.tui - def test_mcp_registration_claude_servers_connect_and_work(self, cuj): - """Scenario: configure UG, inspect Claude's MCP menu, then call both fixture tools. - - Expected: both scoped servers are registered, connected, and expose their tools; - no ug_e2e.other_tools server or decoy tool appears in the complete inventory; - a parent assistant answer contains the correct receipts for a fresh run ID. - """ - session, workspace, _recorder = cuj - session.configure(["configure", "--workspace", workspace.url, "--skip-upgrade"]) - task = McpFixtureTask() - allowed_tools = ",".join( - f"mcp__{service.replace('.', '-')}__{tool}" for service, tool in SERVICES.items() - ) - with Terminal( - session, - "mcp-registration-claude", - [CLAUDE, "--", "--allowedTools", allowed_tools], - ) as tui: - tui.boot() - inventory = open_claude_mcp_inventory(tui, SERVICES) - assert "ug_e2e-other_tools" not in inventory, inventory - assert "decoy_status" not in inventory, inventory - tui.submit(task.prompt) - _wait_for_mcp_task(tui, task) - tui.exit_normally() - assert task.completed(session, CLAUDE) - session.assert_not_routed() - - @pytest.mark.codex - @pytest.mark.tui - def test_mcp_registration_codex_servers_connect_and_work(self, cuj): - """Scenario: configure UG, inspect Codex's MCP inventory, then call both fixture tools. - - Expected: both scoped servers are registered, connected, and expose their tools; - no ug_e2e.other_tools server or decoy tool appears in the complete inventory; - a completed parent turn returns the correct receipts for a fresh run ID. - """ - session, workspace, _recorder = cuj - session.configure(["configure", "--workspace", workspace.url, "--skip-upgrade"]) - task = McpFixtureTask() - with Terminal(session, "mcp-registration-codex", [CODEX]) as tui: - tui.boot() - inventory = open_codex_mcp_inventory(tui, SERVICES) - assert "ug_e2e-other_tools" not in inventory, inventory - assert "decoy_status" not in inventory, inventory - tui.submit(task.prompt) - _wait_for_mcp_task(tui, task) - tui.exit_normally() - assert task.completed(session, CODEX) - session.assert_not_routed() diff --git a/tests/integration/README.md b/tests/integration/README.md index b451cac0e..580e0d9c8 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -46,7 +46,7 @@ does not yet assert live `recommendModel` request counts for configs with and wi ## Dedicated-workspace CUJ5 -[`../e2e_cuj/test_ug_budget_defaults.py`](../e2e_cuj/test_ug_budget_defaults.py) runs +[`../e2e_cuj/test_cuj5_budget_defaults.py`](../e2e_cuj/test_cuj5_budget_defaults.py) runs three budget cases in the shared `E2E CUJs` job against the fixed 1% tier. The above-tier cases use `UG_CUJ_SP_CLIENT_ID` / `UG_CUJ_SP_CLIENT_SECRET`; the below-tier case uses `UG_BUDGET_CUJ_SP_CLIENT_ID` / `UG_BUDGET_CUJ_SP_CLIENT_SECRET`. diff --git a/tests/test_e2e_cuj_helpers.py b/tests/test_e2e_cuj_helpers.py new file mode 100644 index 000000000..36784a1b2 --- /dev/null +++ b/tests/test_e2e_cuj_helpers.py @@ -0,0 +1,309 @@ +"""Offline checks for CUJ bearer minting, leak reporting, terminal waits, and task helpers.""" + +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from tests.e2e_cuj import conftest as cuj_conftest +from tests.e2e_cuj import test_cuj2_mps_mcp as cuj2 +from tests.e2e_cuj import test_cuj3_mcp as mcp_registration +from tests.e2e_cuj.base import bearer +from tests.e2e_cuj.helpers import poll as poll_module +from tests.e2e_cuj.helpers.constants import CLAUDE, CLAUDE_HAIKU_MODEL, CODEX, CodingAgent +from tests.e2e_cuj.helpers.evidence import assert_models, claude_file_task +from tests.e2e_cuj.helpers.poll import poll +from tests.e2e_cuj.helpers.session import ( + MACHINE_WIDE_LEAK, + MachineWideLeak, + UserSession, + dirty_runner_message, +) +from tests.e2e_cuj.helpers.terminal import Terminal +from tests.e2e_cuj.helpers.workspace import Workspace +from tests.e2e_cuj.test_cuj3_models import _catalog_display_names + + +def _client(headers): + return SimpleNamespace(config=SimpleNamespace(authenticate=lambda: headers)) + + +def test_bearer_returns_the_token_from_the_authorization_header(): + assert bearer(_client({"Authorization": "Bearer minted-token"})) == "minted-token" + + +@pytest.mark.parametrize( + "headers", + [{}, {"Authorization": "Basic abc"}], +) +def test_bearer_rejects_a_missing_or_non_bearer_authorization(headers): + with pytest.raises(RuntimeError, match="did not return a bearer token"): + bearer(_client(headers)) + + +def test_refresh_bearer_keeps_a_journey_selected_workspace_bearer(tmp_path): + session = UserSession(tmp_path, Path("/installed/ug"), tmp_path / "artifacts", "first") + session.refresh_bearer("second") + assert session.env["DATABRICKS_BEARER"] == "second" + session.install_bearer("phase-b") + session.refresh_bearer("third") + assert session.env["DATABRICKS_BEARER"] == "phase-b" + + +def test_redact_scrubs_bearers_replaced_by_a_refresh_or_phase_switch(tmp_path): + session = UserSession(tmp_path, Path("/installed/ug"), tmp_path / "artifacts", "first-token") + session.refresh_bearer("second-token") + session.install_bearer("phase-b-token") + redacted = session.redact("first-token second-token phase-b-token") + assert redacted == " " + + +def _leak_fixture(tmp_path, monkeypatch, leaked_names): + managed = tmp_path / "etc" + managed.mkdir() + paths = tuple(managed / name for name in ("claude.json", "codex.toml")) + for path in paths: + if path.name in leaked_names: + path.write_text("{}") + monkeypatch.setattr(cuj_conftest, "MANAGED_PATHS", paths) + request = SimpleNamespace( + config=SimpleNamespace(stash=pytest.Stash()), + node=SimpleNamespace(nodeid="tests/e2e_cuj/test_x.py::TestLeaky"), + ) + return request, paths + + +def test_teardown_leak_names_the_culprit(tmp_path, monkeypatch): + request, paths = _leak_fixture(tmp_path, monkeypatch, {"claude.json"}) + + cuj_conftest._record_machine_wide_leak(request) + + leak = request.config.stash[MACHINE_WIDE_LEAK] + assert leak == MachineWideLeak("tests/e2e_cuj/test_x.py::TestLeaky", (str(paths[0]),)) + assert paths[0].exists(), "Leak handling must not touch machine-wide files" + message = dirty_runner_message(leak) + assert "TestLeaky" in message + assert str(paths[0]) in message + + +def test_clean_teardown_records_no_leak(tmp_path, monkeypatch): + request, _ = _leak_fixture(tmp_path, monkeypatch, set()) + cuj_conftest._record_machine_wide_leak(request) + leak = request.config.stash.get(MACHINE_WIDE_LEAK, None) + assert leak is None + assert dirty_runner_message(leak) == ( + "Existing machine-wide agent settings; use a clean disposable runner. Nothing was changed." + ) + + +def _fake_ug(tmp_path, body): + """A real executable standing in for `ug`, so the real PTY driver runs offline.""" + script = tmp_path / "ug" + script.write_text("#!/bin/sh\n" + body) + script.chmod(0o755) + return UserSession(tmp_path, script, tmp_path / "artifacts", "token") + + +@pytest.mark.parametrize( + ("args", "agent"), + [([], None), ([], "gemini"), (["status"], None), ([CLAUDE], CODEX)], +) +def test_terminal_rejects_a_missing_unsupported_or_mismatched_agent(tmp_path, args, agent): + session = UserSession(tmp_path, Path("/installed/ug"), tmp_path / "artifacts", "token") + with pytest.raises(ValueError, match="requires a Claude or Codex command or bare ug agent"): + Terminal(session, "rejected", args, agent=agent) + + +def test_wait_until_rejects_a_permission_prompt(tmp_path): + session = _fake_ug(tmp_path, "echo 'Do you want to proceed?'\nsleep 30\n") + with Terminal(session, "rejects", [], agent=CLAUDE) as tui: + with pytest.raises(AssertionError, match="Unexpected permission request"): + tui.wait_until(lambda: False, "never", timeout=10) + + +def test_wait_until_completes_when_done(tmp_path): + done = tmp_path / "done" + session = _fake_ug(tmp_path, f"echo working\ntouch '{done}'\nsleep 30\n") + with Terminal(session, "completes", [], agent=CODEX) as tui: + tui.wait_until(done.exists, "the done marker", timeout=10) + assert "working" in tui.visible + + +def test_wait_until_lets_on_screen_answer_an_expected_dialog(tmp_path): + done = tmp_path / "done" + session = _fake_ug( + tmp_path, f"echo 'Allow this call?'\nread answer\necho allowed\ntouch '{done}'\nsleep 30\n" + ) + answered = [] + + def approve(screen): + if "Allow this call?" in screen and not answered: + tui.send("\r", "allow the expected call") + answered.append(screen) + return True + return False + + with Terminal(session, "approves", [], agent=CODEX) as tui: + tui.wait_until(done.exists, "the approved call", timeout=10, on_screen=approve) + assert len(answered) == 1 + assert "allowed" in tui.visible + + +def test_mcp_list_poll_retries_a_failed_probe_until_rows_match(monkeypatch): + monkeypatch.setattr(poll_module.time, "sleep", lambda seconds: None) + healthy = "\n".join( + rows[0] for rows in mcp_registration._expected_fixture_server_rows().values() + ) + failed = healthy.replace(f"{CLAUDE}, {CODEX} connected", f"{CLAUDE}:failed {CODEX}:enabled", 1) + outputs = [failed, failed, healthy, "unused"] + session = SimpleNamespace( + run=lambda *args, timeout: SimpleNamespace(stdout=outputs.pop(0), stderr="") + ) + output = mcp_registration._poll_mcp_list(session) + assert output == f"{healthy}\n" + assert outputs == ["unused"] + + +def test_claude_file_task_names_the_absolute_path_without_the_answer(tmp_path): + session = UserSession(tmp_path, Path("/installed/ug"), tmp_path / "artifacts", "token") + task = claude_file_task(session) + path = session.cwd / task.filename + assert path.read_text() == task.value + "\n" + assert str(path) in task.prompt + assert task.value not in task.prompt + + +def test_poll_returns_the_first_accepted_value(): + values = iter([0, 0, 2, 5]) + assert poll(lambda: next(values), timeout=10, interval=0) == 2 + + +def test_poll_returns_the_last_value_at_the_timeout(): + calls = [] + assert poll(lambda: calls.append(1) or (0, 1), timeout=0, interval=0, done=all) == ( + 0, + 1, + ) + assert calls == [1] + + +def test_agent_configs_rejects_a_duplicate_agent(): + entry = {"agent": CodingAgent.CODEX, "config": {}} + with pytest.raises(AssertionError, match="Duplicate enabled agents"): + Workspace.agent_configs({"enabled_agents": [entry, entry]}) + assert Workspace.agent_configs({"enabled_agents": [entry]}) == {CodingAgent.CODEX: {}} + + +def _catalog_workspace(pages): + calls = [] + + def do(method, path, headers, query): + calls.append((path, dict(headers), dict(query))) + return pages.pop(0) + + workspace = SimpleNamespace(client=SimpleNamespace(api_client=SimpleNamespace(do=do))) + return workspace, calls + + +def _anthropic_page(models, last_id=None): + return { + "data": [{"id": model, "display_name": model.upper()} for model in models], + "has_more": last_id is not None, + "last_id": last_id, + } + + +def test_catalog_display_names_follows_the_anthropic_cursor(): + workspace, calls = _catalog_workspace([_anthropic_page(["a"], "a"), _anthropic_page(["b"])]) + assert _catalog_display_names(workspace, CLAUDE, "ug_e2e.models") == {"a": "A", "b": "B"} + assert [query for _, _, query in calls] == [ + {"limit": "1000"}, + {"limit": "1000", "after_id": "a"}, + ] + assert calls[0][1]["Anthropic-Version"] == "2023-06-01" + + +@pytest.mark.parametrize( + ("pages", "message"), + [ + ([_anthropic_page(["a"], "a"), _anthropic_page(["a"], "b")], "Repeated catalog model"), + ([_anthropic_page(["a"], "x"), _anthropic_page(["b"], "x")], "pagination cursor"), + ], +) +def test_catalog_display_names_rejects_repeated_pages(pages, message): + workspace, _ = _catalog_workspace(pages) + with pytest.raises(AssertionError, match=message): + _catalog_display_names(workspace, CLAUDE, "ug_e2e.models") + + +def test_assert_models_maps_native_aliases(): + alias = "anthropic.claude-haiku-4-5-20251001-v1:0" + assert_models([alias, CLAUDE_HAIKU_MODEL], CLAUDE_HAIKU_MODEL) + with pytest.raises(AssertionError): + assert_models([CLAUDE_HAIKU_MODEL, "system.ai.claude-sonnet-4-6"], CLAUDE_HAIKU_MODEL) + + +@pytest.mark.parametrize( + "shape", + ["bare", "native_keys", "raw_result_string", "raw_result_object"], +) +def test_mcp_receipts_accept_native_keys_and_raw_tool_results(shape): + task = mcp_registration.McpFixtureTask() + native = {tool: key for key, tool in task.aliases.items()} + observed = { + "bare": dict(task.expected), + "native_keys": {native[tool]: value for tool, value in task.expected.items()}, + "raw_result_string": { + tool: json.dumps({"result": value}) for tool, value in task.expected.items() + }, + "raw_result_object": {tool: {"result": value} for tool, value in task.expected.items()}, + }[shape] + + assert task._receipts(observed) == task.expected + + +def test_mcp_receipts_reject_wrong_extra_or_duplicate_receipts(): + task = mcp_registration.McpFixtureTask() + describe, read = task.expected + native = {tool: key for key, tool in task.aliases.items()} + + assert task._receipts({**task.expected, read: "0" * 64}) != task.expected + assert task._receipts({**task.expected, "extra": "x"}) != task.expected + assert ( + task._receipts( + {describe: {"result": task.expected[describe], "x": 1}, read: task.expected[read]} + ) + != task.expected + ) + assert task._receipts({**task.expected, native[describe]: task.expected[describe]}) is None + assert task._receipts(["not", "a", "mapping"]) is None + + +def test_mcp_completion_accepts_a_fenced_reply(monkeypatch): + task = mcp_registration.McpFixtureTask() + reply = "```json\n" + json.dumps(task.expected) + "\n```" + monkeypatch.setattr(mcp_registration, "agent_sessions", lambda session, agent: {"s": []}) + monkeypatch.setattr(mcp_registration, "assistant_answers", lambda agent, records: [reply]) + + assert task.completed(None, "claude") + + +def test_cuj2_mcp_listing_retries_a_banner_only_claude_listing(tmp_path, monkeypatch): + for agent in (CLAUDE, CODEX): + (tmp_path / agent).write_text("#!/bin/sh\n") + (tmp_path / agent).chmod(0o755) + banner = "Checking MCP server health…\n\n\n" + rows = f"{cuj2.registered_name(cuj2.SANDBOX_MCP_SERVICE_NAME)}: ug mcp-proxy - ✓ Connected\n" + outputs = {CLAUDE: [banner, rows], CODEX: [rows]} + calls = [] + + def run(*args, binary, timeout): + agent = Path(binary).name + calls.append(agent) + return SimpleNamespace(stdout=outputs[agent].pop(0), stderr="") + + monkeypatch.setattr(poll_module.time, "sleep", lambda seconds: None) + cuj2._assert_generated_mcp_listings(SimpleNamespace(env={"PATH": str(tmp_path)}, run=run)) + + assert calls == [CLAUDE, CLAUDE, CODEX] diff --git a/tests/test_integration_contract.py b/tests/test_integration_contract.py index cce7fdc1a..a1c6491e7 100644 --- a/tests/test_integration_contract.py +++ b/tests/test_integration_contract.py @@ -2,6 +2,7 @@ import ast import json +import os import re import subprocess import sys @@ -54,17 +55,69 @@ def test_managed_integration_ci_is_blocking(): ) in gate -def test_dedicated_cuj_ci_discovers_the_whole_folder(): - workflow = (Path(__file__).parent.parent / ".github/workflows/integration.yml").read_text() - job = workflow.split("\n dedicated-cuj:\n", 1)[1].split("\n cujs:\n", 1)[0] - gate = workflow.split("\n cujs:\n", 1)[1] +WORKFLOW = Path(__file__).parent.parent / ".github/workflows/integration.yml" + + +def _workflow_job(name, next_name): + workflow = WORKFLOW.read_text() + return workflow.split(f"\n {name}:\n", 1)[1].split(f"\n {next_name}:\n", 1)[0] + + +def _discover_cujs(root, output): + plan = _workflow_job("dedicated-cuj-plan", "dedicated-cuj") + script = textwrap.dedent(plan.split(" run: |\n", 1)[1]) + result = subprocess.run( + ["bash", "-c", script], + cwd=root, + env={"PATH": os.environ["PATH"], "GITHUB_OUTPUT": str(output)}, + capture_output=True, + text=True, + timeout=10, + ) + if result.returncode: + return None, result.stderr + return json.loads(output.read_text().removeprefix("cujs=")), result.stderr + + +def test_dedicated_cuj_ci_runs_each_discovered_file_on_its_own_runner(): + job = _workflow_job("dedicated-cuj", "cujs") assert "docs.google.com/document/d/1WKd1fdWD0Y4tAV1H9Si-SGx2UZFL7iZ2S3HtBjidmS0" in job - assert "pytest --confcutdir=tests/e2e_cuj tests/e2e_cuj" in job - assert "find tests/e2e_cuj -name 'test_*.py'" in job + assert "name: E2E CUJs · ${{ matrix.name }}" in job + assert "fail-fast: false" in job + assert "include: ${{ fromJSON(needs.dedicated-cuj-plan.outputs.cujs) }}" in job + assert 'pytest --confcutdir=tests/e2e_cuj "tests/e2e_cuj/${{ matrix.test }}"' in job + assert 'UV_HTTP_TIMEOUT: "120"' in job and 'UV_HTTP_RETRIES: "6"' in job assert "UG_CUJ1_WORKSPACE: ${{ secrets.UG_CUJ1_WORKSPACE }}" in job - assert "test_cuj_" not in job - assert 'result["result"] != "success"' in gate + assert "test_cuj" not in job + + +def test_every_cuj_file_is_discovered_with_a_unique_check_name(tmp_path): + root = Path(__file__).parent.parent + entries, stderr = _discover_cujs(root, tmp_path / "github-output") + + assert entries is not None, stderr + assert [entry["test"] for entry in entries] == sorted( + path.name for path in (root / "tests/e2e_cuj").glob("test_*.py") + ) + names = [entry["name"] for entry in entries] + assert len(set(names)) == len(names), f"Duplicate CUJ_NAME values: {names}" + for entry in entries: + number = re.fullmatch(r"test_cuj(\d+)_[a-z0-9_]+\.py", entry["test"]) + assert number, f"Name CUJ files test_cuj_.py: {entry['test']}" + assert entry["name"].startswith(f"CUJ {number.group(1)} · "), entry + + +def test_cuj_discovery_fails_for_a_file_without_a_check_name(tmp_path): + cujs = tmp_path / "tests/e2e_cuj" + cujs.mkdir(parents=True) + (cujs / "test_cuj1_a.py").write_text('import os\n\nCUJ_NAME = "CUJ 1 · A"\n') + (cujs / "test_cuj2_b.py").write_text("import os\n") + + entries, stderr = _discover_cujs(tmp_path, tmp_path / "github-output") + + assert entries is None + assert "test_cuj2_b.py must define CUJ_NAME" in stderr FORK_GUARD = ( @@ -80,6 +133,7 @@ def test_dedicated_cuj_ci_discovers_the_whole_folder(): "full", "opencode", "managed", + "dedicated-cuj-plan", "dedicated-cuj", "cujs", } @@ -126,7 +180,16 @@ def test_windows_integration_ci_uses_shared_claude_version(): @pytest.mark.parametrize( - "failed_job", ["installation", "workspace", "smoke", "full", "managed", "dedicated-cuj"] + "failed_job", + [ + "installation", + "workspace", + "smoke", + "full", + "managed", + "dedicated-cuj-plan", + "dedicated-cuj", + ], ) @pytest.mark.parametrize("job_result", ["success", "failure", "cancelled", "skipped"]) def test_integration_ci_gate_requires_every_job(failed_job, job_result): @@ -142,6 +205,7 @@ def test_integration_ci_gate_requires_every_job(failed_job, job_result): "smoke", "full", "managed", + "dedicated-cuj-plan", "dedicated-cuj", ) } diff --git a/tests/test_ug_review_bot.py b/tests/test_ug_review_bot.py index 14727fd04..4fa23c4dc 100644 --- a/tests/test_ug_review_bot.py +++ b/tests/test_ug_review_bot.py @@ -122,6 +122,7 @@ def test_responses_request_uses_instructions_and_does_not_store_output(): assert request["input"] == "Review this diff" assert "Review policy" in request["instructions"] assert request["reasoning"] == {"effort": "high"} + assert request["max_output_tokens"] >= 16_000 assert request["store"] is False assert "messages" not in request @@ -143,6 +144,42 @@ def test_extract_response_text_ignores_reasoning_items(): assert review_bot.extract_response_text(payload) == "first\nsecond" +def test_request_review_reports_why_the_response_has_no_text(monkeypatch): + monkeypatch.setenv("DATABRICKS_HOST", "https://example.databricks.com") + monkeypatch.setenv("UG_REVIEW_MODEL", "system.ai.gpt") + monkeypatch.setattr(review_bot, "_oauth_access_token", lambda host: "access-token") + payload = { + "status": "incomplete", + "incomplete_details": {"reason": "max_output_tokens"}, + "output": [{"type": "reasoning", "content": []}], + } + monkeypatch.setattr( + review_bot.urllib.request, + "urlopen", + lambda request, timeout: BytesIO(json.dumps(payload).encode()), + ) + + with pytest.raises( + review_bot.ReviewError, match=r"no output text \(incomplete: max_output_tokens\)" + ): + review_bot._request_review("Title", "", _bundle(), "Review policy", "Repository policy") + + +def test_request_review_ignores_non_object_incomplete_details(monkeypatch): + monkeypatch.setenv("DATABRICKS_HOST", "https://example.databricks.com") + monkeypatch.setenv("UG_REVIEW_MODEL", "system.ai.gpt") + monkeypatch.setattr(review_bot, "_oauth_access_token", lambda host: "access-token") + payload = {"status": "incomplete", "incomplete_details": "max_output_tokens", "output": []} + monkeypatch.setattr( + review_bot.urllib.request, + "urlopen", + lambda request, timeout: BytesIO(json.dumps(payload).encode()), + ) + + with pytest.raises(review_bot.ReviewError, match=r"no output text\.$"): + review_bot._request_review("Title", "", _bundle(), "Review policy", "Repository policy") + + def test_oauth_access_token_uses_service_principal_credentials(monkeypatch): monkeypatch.setenv("DATABRICKS_CLIENT_ID", "client-id") monkeypatch.setenv("DATABRICKS_CLIENT_SECRET", "client-secret")