From 0d8afa65761ba08aa253296aa7885576d2c066b7 Mon Sep 17 00:00:00 2001 From: Mason Cao Date: Mon, 5 Oct 2026 20:00:29 +0000 Subject: [PATCH 01/19] [AIGTWY-4877] Add MPS and MCP two-agent CUJ --- tests/README.md | 21 +- tests/e2e_cuj/helpers/__init__.py | 792 +++++++++++++++++++++++++++++ tests/e2e_cuj/pytest.ini | 5 + tests/e2e_cuj/test_cuj2_mps_mcp.py | 428 ++++++++++++++++ 4 files changed, 1244 insertions(+), 2 deletions(-) create mode 100644 tests/e2e_cuj/helpers/__init__.py create mode 100644 tests/e2e_cuj/pytest.ini create mode 100644 tests/e2e_cuj/test_cuj2_mps_mcp.py diff --git a/tests/README.md b/tests/README.md index f9e9a8130..070516916 100644 --- a/tests/README.md +++ b/tests/README.md @@ -10,6 +10,7 @@ mocks, monkeypatching, fake binaries/services, or fabricated ug state. | Dedicated-workspace CUJ | `e2e_cuj/test_cuj_*.py` | Real user journeys against read-only, preconfigured workspaces | | Existing e2e | `test_e2e*.py` | Real workspace behavior with some patched setup/internal calls | | Integration CUJs | `integration/test_*.py` | Public configure, TUI, script, command, protocol, and lifecycle journeys | +| Dedicated-workspace CUJs | `e2e_cuj/test_*.py` | Real published CodingAgentConfig, MPS, MCP, and two-agent inference | | Installation | `integration/test_installation.py` | Fresh installed package, CLI, and local helpers without credentials on Linux and advisory native Windows | Dedicated CUJs reuse `integration/utils` session/terminal mechanics and file-task @@ -131,10 +132,12 @@ These are **implemented assertions**, not a claim that every version passes. Consult the run's JUnit report and artifacts for results. Each function states its **Scenario** and **Expected** outcome and shows its configure and launch commands. Fixtures supply fresh environments and credentials, never configured ug. -All tests live directly in `integration/`; shared mechanics live in `utils/`. +Integration tests live directly in `integration/`; shared mechanics live in `utils/`. +Dedicated-workspace CUJs live in `e2e_cuj/` and use their own black-box helpers. | Test | User action | Expected evidence | | --- | --- | --- | +| `test_configuration`, `test_codex`, `test_claude` in `test_cuj2_mps_mcp.py` | Temporarily publish the two-agent MPS/MCP config; run separate configure, Codex TUI, and Claude TUI cases | Configuration verifies both native MPS APIs, generated settings, and sandbox inclusion with web_search excluded; each TUI case records its exact MPS target header and model, paired HTTP 200 response, correlated `run_code` call/result, and structured final marker | | `test_ug_configure_claude_databricks` | Configure Databricks Hosted; execute the generated auth helper; launch plain `ug claude`, read a file, and open `/model` | Generated helper invokes `ug` with clean token stdout; assistant returns an unpredictable file value; native discovery caches `system.ai` models and the picker shows a discovered model without an opt-in flag; normal exit; reopen with working keyboard input | | `test_ug_configure_claude_anthropic_mps` | Select Anthropic MPS in the real configure picker; launch Claude | Saved provider in status; completed TUI file task; normal exit | | `test_ug_configure_codex_databricks` | Configure Databricks Hosted; execute the generated auth helper; open Codex TUI and read a file | Generated helper invokes `ug` with clean token stdout; completed assistant answer contains the file value; normal exit and reopen | @@ -209,6 +212,19 @@ deduplication, and raw catalog ID/display-name rows for scoped pickers. Managed discovery expectations come from separate read-only, provider-scoped model-list requests; they do not rely solely on ug's generated catalog. +The dedicated-workspace CUJ suite adds three independently runnable CUJ2 MPS+MCP cases: +one configuration case and separate Codex and Claude TUI task cases. Each temporarily +publishes the two-agent config and restores the original workspace config. The TUI cases +forward agent requests through a per-test loopback recorder to the real workspace and +assert the provider-service header, model, paired HTTP 200 response, correlated MCP tool +result, and final answer. Failed TUI runs retain terminal output/actions/screen and +best-effort agent transcript snapshots under the runner artifact directory before the +temporary homes are deleted. The suite requires +`UG_CUJ_SP_CLIENT_ID`, `UG_CUJ_SP_CLIENT_SECRET`, `UG_INTEGRATION_BIN`, +`UG_INTEGRATION_RUN_DIR`, and runner-provided `DATABRICKS_BEARER`; run it with +`uv run pytest --confcutdir=tests/e2e_cuj tests/e2e_cuj`. It owns the published +config for the test duration and does not use `UCODE_MANAGED_CONFIG_STUB`. + ug no longer runs a post-configure agent probe; the deprecated `--skip-validate` flag is accepted as a no-op where older journeys still pass it. Tests retain `--skip-upgrade` as a deprecated no-op too; UG only upgrades agents below its @@ -288,7 +304,8 @@ These are unit/component checks; they do not establish live sudo password-prompt | Scenario | Status / requirement | | --- | --- | -| Live MCP and skills functionality | Deferred; installation tests cover the local web-search MCP handshake and tool listing, not upstream proxying or a real search request. Custom OAuth search dispatch/refresh has component coverage; live parent/child search and permission decisions remain unverified | +| Broad live MCP functionality | CUJ2 covers the sandbox MCP and two-agent tool use. Installation tests cover the local web-search MCP handshake and tool listing, not upstream proxying or a real search request. Other MCP services, live parent/child search, and permission decisions remain deferred | +| Dedicated CUJ Skill coverage | Deferred because of a backend storage-path issue; CUJ2 has no Skill configuration, download, invocation, or assertion path | | Broad configure flags, multiple workspaces, and PAT flows | Deferred while focusing on basic CUJs | | Workspace-switch MCP cleanup | The `workspace_switch` CUJ covers real registration, cleanup, repeat configure, and a completed Claude task. Unit/component tests cover duplicate attempts and injected removal failures; the CUJ does not force an agent timeout. It runs in the required managed CI lane for full/live runs. | | Relayed/subscription MPS discovery | Not covered by the scoped discovery journeys | diff --git a/tests/e2e_cuj/helpers/__init__.py b/tests/e2e_cuj/helpers/__init__.py new file mode 100644 index 000000000..8713dbb03 --- /dev/null +++ b/tests/e2e_cuj/helpers/__init__.py @@ -0,0 +1,792 @@ +"""Black-box helpers for dedicated-workspace CUJs. + +This module talks to the installed commands and the Databricks SDK only. It deliberately does +not import ``ucode``: a CUJ must exercise the installed wheel just as a user does. +""" + +from __future__ import annotations + +import contextlib +import importlib +import json +import os +import re +import signal +import subprocess +import tempfile +import time +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +_ANSI = re.compile(r"\x1b\[[0-9;?]*[ -/]*[@-~]") +_CALL_TYPES = { + "custom_tool_call", + "function_call", + "mcp_tool_call", + "tool_call", + "tool_use", +} +_RESULT_TYPES = { + "custom_tool_call_output", + "function_call_output", + "mcp_tool_result", + "mcp_tool_call_output", + "tool_result", + "tool_call_output", +} + + +def _kill_process_group(process: subprocess.Popen[str]) -> None: + if process.poll() is not None: + return + if os.name == "posix": + with contextlib.suppress(ProcessLookupError): + os.killpg(process.pid, signal.SIGTERM) + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + with contextlib.suppress(ProcessLookupError): + os.killpg(process.pid, signal.SIGKILL) + elif os.name == "nt": + process.kill() + process.wait(timeout=5) + + +def _process_group_options() -> dict[str, Any]: + if os.name == "posix": + return {"start_new_session": True} + if os.name == "nt": + return {"creationflags": subprocess.CREATE_NEW_PROCESS_GROUP} + return {} + + +def _json_strings(value: object) -> list[str]: + if isinstance(value, str): + return [value] + if isinstance(value, dict): + return [text for item in value.values() for text in _json_strings(item)] + if isinstance(value, list): + return [text for item in value for text in _json_strings(item)] + return [] + + +def _dict_nodes(value: object): + if isinstance(value, dict): + yield value + for child in value.values(): + yield from _dict_nodes(child) + elif isinstance(value, list): + for child in value: + yield from _dict_nodes(child) + + +def _first_string(mapping: dict, *keys: str) -> str | None: + for key in keys: + value = mapping.get(key) + if isinstance(value, str) and value: + return value + return None + + +def _call_id(mapping: dict) -> str | None: + return _first_string(mapping, "id", "call_id", "callId", "tool_use_id", "toolUseId") + + +@dataclass(frozen=True) +class CommandResult: + argv: tuple[str, ...] + returncode: int + stdout: str + stderr: str + timed_out: bool = False + + +class CujSession: + """An isolated home/project and subprocess runner for one CUJ case.""" + + def __init__(self, run_root: Path, binary: Path, case_name: str): + self._temporary = tempfile.TemporaryDirectory(prefix="cuj-", dir=run_root) + self.root = Path(self._temporary.name) + self.home = self.root / "home" + self.project = self.root / "project with spaces" + self.home.mkdir() + self.project.mkdir() + self.binary = binary + self.artifacts = run_root / "artifacts" / case_name + self.artifacts.mkdir(parents=True, exist_ok=True) + self._commands = 0 + self.env = self._environment() + + def _environment(self) -> dict[str, str]: + keep = ( + "PATH", + "SYSTEMROOT", + "COMSPEC", + "LANG", + "LC_ALL", + "SSL_CERT_FILE", + "SSL_CERT_DIR", + "REQUESTS_CA_BUNDLE", + "NODE_EXTRA_CA_CERTS", + "HTTPS_PROXY", + "HTTP_PROXY", + "NO_PROXY", + ) + env = {key: os.environ[key] for key in keep if key in os.environ} + bearer = os.environ.get("DATABRICKS_BEARER", "").strip() + if not bearer: + raise AssertionError("DATABRICKS_BEARER must be provided by the CUJ runner") + env.update( + { + "HOME": str(self.home), + "USERPROFILE": str(self.home), + "XDG_CONFIG_HOME": str(self.home / ".config"), + "XDG_CACHE_HOME": str(self.home / ".cache"), + "XDG_DATA_HOME": str(self.home / ".local/share"), + "CLAUDE_CONFIG_DIR": str(self.home / ".claude"), + "CODEX_HOME": str(self.home / ".codex"), + "DATABRICKS_CONFIG_FILE": str(self.home / ".databrickscfg"), + "DATABRICKS_BEARER": bearer, + "UG_INTEGRATION_BIN": str(self.binary), + "DISABLE_AUTOUPDATER": "1", + "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1", + "NO_COLOR": "1", + "TERM": "dumb", + "COLUMNS": "160", + "ENABLE_SMART_ROUTING_V2": "0", + "PYTHONNOUSERSITE": "1", + } + ) + return env + + def _redact(self, text: str) -> str: + for name in ("DATABRICKS_BEARER", "UG_CUJ_SP_CLIENT_SECRET"): + token = os.environ.get(name) or self.env.get(name) + if token: + text = text.replace(token, "") + return _ANSI.sub("", text) + + def record(self, name: str, value: object) -> None: + payload = json.dumps(value, indent=2, sort_keys=True, default=str) + self.artifacts.joinpath(name).write_text(self._redact(payload), encoding="utf-8") + + def run( + self, + *args: str, + timeout: int = 180, + binary: str | Path | None = None, + input_text: str | None = None, + check: bool = True, + ) -> CommandResult: + command = tuple(str(binary or self.binary) for _ in [0]) + tuple(args) + process = subprocess.Popen( + list(command), + cwd=self.project, + env=self.env, + stdin=subprocess.PIPE if input_text is not None else subprocess.DEVNULL, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + **_process_group_options(), + ) + timed_out = False + try: + stdout, stderr = process.communicate(input=input_text, timeout=timeout) + except subprocess.TimeoutExpired: + timed_out = True + _kill_process_group(process) + stdout, stderr = process.communicate(timeout=5) + finally: + _kill_process_group(process) + result = CommandResult( + command, + process.returncode, + self._redact(stdout), + self._redact(stderr), + timed_out, + ) + self._commands += 1 + self.record( + f"command-{self._commands}.json", + { + "argv": command, + "returncode": result.returncode, + "stderr": result.stderr, + "stdout": result.stdout, + "timed_out": result.timed_out, + }, + ) + if result.timed_out: + raise AssertionError(f"Command exceeded {timeout}s: {command}") + if check and result.returncode != 0: + raise AssertionError( + f"Command failed ({result.returncode}): {command}\n{result.stderr}" + ) + return result + + def external(self, program: str, *args: str, timeout: int = 120) -> CommandResult: + return self.run(*args, binary=program, timeout=timeout) + + def _record_tui_artifact( + self, + command: tuple[str, ...], + *, + stdout: str, + returncode: int | None, + timed_out: bool, + actions: list[dict[str, str]], + screen: str, + failure: BaseException | None, + ) -> None: + self._commands += 1 + try: + self.record( + f"command-{self._commands}.json", + { + "argv": command, + "returncode": returncode, + "stderr": "", + "stdout": stdout, + "timed_out": timed_out, + "interactive": True, + "actions": actions, + "screen": screen, + "error": f"{type(failure).__name__}: {failure}" if failure else None, + }, + ) + except BaseException: + if failure is None: + raise + + def run_tui( + self, + agent: str, + prompt: str, + *, + completion: Callable[[CujSession], bool], + timeout: int = 240, + ) -> CommandResult: + """Run a real interactive agent through a PTY and return its terminal transcript.""" + if os.name != "posix": + raise AssertionError("The dedicated CUJ requires POSIX PTY support.") + try: + pexpect = importlib.import_module("pexpect") + pyte = importlib.import_module("pyte") + except ImportError as exc: + raise AssertionError("The dedicated CUJ runner must install pexpect and pyte.") from exc + + if agent not in {"claude", "codex"}: + raise AssertionError(f"Unsupported TUI agent: {agent}") + # Bare ug resolves the workspace default (Codex in CUJ2); Claude needs its explicit + # subcommand so this path exercises the configured Claude provider too. + command = (str(self.binary),) if agent == "codex" else (str(self.binary), "claude") + env = {**self.env, "TERM": "xterm-256color", "COLUMNS": "140", "LINES": "60"} + try: + child = pexpect.spawn( + command[0], + command[1:], + cwd=str(self.project), + env=env, + encoding="utf-8", + codec_errors="replace", + dimensions=(60, 140), + timeout=1, + ) + except BaseException as exc: + self._record_tui_artifact( + command, + stdout="", + returncode=None, + timed_out=False, + actions=[], + screen="", + failure=exc, + ) + raise + + def send_terminal_reply(data: str) -> None: + child.send(data) + + # pyte's screen emits terminal-device replies through this stream hook. The callback is + # deliberately limited to replies generated from the rendered screen state. + class TerminalScreen(pyte.Screen): + def write_process_input(self, data: str) -> None: + send_terminal_reply(data) + + screen = TerminalScreen(140, 60) + stream = pyte.Stream(screen) + output: list[str] = [] + actions: list[dict[str, str]] = [] + ended = False + timed_out = False + + def visible() -> str: + return "\n".join(line.rstrip() for line in screen.display) + + def read() -> None: + nonlocal ended + try: + chunk = child.read_nonblocking(size=65536, timeout=0.2) + except pexpect.TIMEOUT: + return + except pexpect.EOF: + before = child.before + if before: + output.append(before) + stream.feed(before) + ended = True + return + output.append(chunk) + stream.feed(chunk) + + def send(keys: str, reason: str) -> None: + actions.append({"reason": reason, "keys": keys, "screen_before": visible()}) + child.send(keys) + + def selected_line() -> str: + return next( + (line.strip() for line in visible().splitlines() if re.match(r"^\s*[›❯>]", line)), + "", + ) + + def choose(prompt_text: str, label: str) -> None: + nonlocal timed_out + deadline = time.monotonic() + 120 + visited: set[str] = set() + while time.monotonic() < deadline: + read() + current_screen = visible() + assert not ended, f"TUI exited while waiting for {prompt_text}:\n{current_screen}" + current = selected_line() + if prompt_text in current_screen and current: + if label in current: + send("\r", f"choose {label}") + return + assert current not in visited, f"Menu does not offer {label}:\n{current_screen}" + visited.add(current) + send("\x1b[B", f"move towards {label}") + time.sleep(0.05) + timed_out = True + raise AssertionError(f"TUI did not show {prompt_text} within 120s:\n{visible()}") + + def wait_for_prompt() -> None: + nonlocal timed_out + deadline = time.monotonic() + 120 + handled: set[str] = set() + ready_since: float | None = None + while time.monotonic() < deadline: + read() + assert not ended, f"TUI exited before its prompt:\n{visible()}" + current_screen = visible() + if "Accessing workspace:" in current_screen and self.project.name in current_screen: + choose("Accessing workspace:", "Yes, I trust this folder") + continue + if "Hooks need review" in current_screen: + choose("Hooks need review", "Trust all and continue") + continue + dialogs = [ + ( + "theme", + "Choose the text style" in current_screen and "Dark mode" in current_screen, + ), + ( + "security-notes", + "Security notes" in current_screen + and "Enter to continue" in current_screen, + ), + ( + "trust-folder", + self.project.name in current_screen + and bool( + re.search(r"1[.)]\s+Yes, I trust (?:this|the) folder", current_screen) + ), + ), + ( + "trust-directory", + self.project.name in current_screen + and "trust" in current_screen.lower() + and bool(re.search(r"1[.)]\s+Yes, (?:continue|proceed)", current_screen)), + ), + ] + handled_dialog = False + for label, shown in dialogs: + if shown: + handled_dialog = True + if label not in handled: + send("\r", label) + handled.add(label) + break + if handled_dialog: + continue + assert "Select login method:" not in current_screen, ( + f"Configured ug launched {agent}'s account-login flow instead of its " + "gateway session:\n" + current_screen + ) + if agent == "codex": + assert not ( + "Sign in with ChatGPT" in current_screen + and "Provide your own API key" in current_screen + ), "Configured ug launched Codex's account-login flow:\n" + current_screen + title = "Claude Code" if agent == "claude" else "Codex" + title_ready = title in current_screen and "loading" not in current_screen.lower() + if title_ready and re.search(r"(?m)^\s*[❯›>]\s*(?!\d+[.)])", current_screen): + ready_since = ready_since or time.monotonic() + if time.monotonic() - ready_since >= 1: + actions.append({"reason": "prompt-ready", "screen": current_screen}) + return + else: + ready_since = None + time.sleep(0.05) + timed_out = True + raise AssertionError(f"TUI did not reach a usable prompt within 120s:\n{visible()}") + + def wait_for_completion() -> None: + nonlocal timed_out + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + read() + if ended: + raise AssertionError(f"TUI exited before the task completed:\n{visible()}") + if completion(self): + return + time.sleep(0.05) + timed_out = True + raise AssertionError(f"TUI did not complete within {timeout}s:\n{visible()}") + + failure: BaseException | None = None + returncode: int | None = None + try: + wait_for_prompt() + send(prompt, "type task prompt") + deadline = time.monotonic() + 30 + while time.monotonic() < deadline: + read() + if prompt.split()[-1] in visible(): + break + assert not ended, f"TUI exited while rendering the task prompt:\n{visible()}" + time.sleep(0.05) + else: + timed_out = True + raise AssertionError(f"TUI did not render the task prompt:\n{visible()}") + send("\r", "submit task prompt") + wait_for_completion() + send("/exit", "type exit command") + send("\r", "submit exit command") + deadline = time.monotonic() + 30 + while not ended and time.monotonic() < deadline: + read() + time.sleep(0.05) + if not ended: + timed_out = True + raise AssertionError(f"TUI did not exit within 30s:\n{visible()}") + child.close(force=False) + returncode = child.exitstatus if child.exitstatus is not None else 1 + if returncode != 0: + raise AssertionError( + f"Interactive command failed ({returncode}): {command}\n{visible()}" + ) + return CommandResult( + command, + returncode, + self._redact("".join(output)), + "", + ) + except BaseException as exc: + failure = exc + raise + finally: + if not ended: + with contextlib.suppress(OSError): + os.killpg(child.pid, signal.SIGTERM) + with contextlib.suppress(OSError): + os.killpg(child.pid, signal.SIGKILL) + with contextlib.suppress(Exception): + child.close(force=True) + self._record_tui_artifact( + command, + stdout="".join(output), + returncode=returncode if returncode is not None else child.exitstatus, + timed_out=timed_out, + actions=actions, + screen=visible(), + failure=failure, + ) + + def read_json(self, relative_path: str) -> object: + return json.loads(self.home.joinpath(relative_path).read_text(encoding="utf-8")) + + def transcripts(self, agent: str) -> dict[str, list[dict]]: + root = self.home / (".claude/projects" if agent == "claude" else ".codex/sessions") + result: dict[str, list[dict]] = {} + paths = sorted(root.rglob("*.jsonl")) if root.is_dir() else [] + for path in paths: + records: list[dict] = [] + for line_number, line in enumerate(path.read_text(encoding="utf-8").splitlines(), 1): + if not line.strip(): + continue + try: + row = json.loads(line) + except json.JSONDecodeError as exc: + raise AssertionError( + f"Invalid {agent} transcript {path}:{line_number}" + ) from exc + if isinstance(row, dict): + records.append(row) + result[str(path.relative_to(root))] = records + return result + + def transcript_contains(self, agent: str, *needles: str) -> bool: + """Return true when one complete transcript record contains every needle.""" + root = self.home / (".claude/projects" if agent == "claude" else ".codex/sessions") + if not root.is_dir(): + return False + for path in root.rglob("*.jsonl"): + for line in path.read_text(encoding="utf-8").splitlines(): + if not line.strip(): + continue + with contextlib.suppress(json.JSONDecodeError): + record = json.loads(line) + if isinstance(record, dict): + text = "\n".join(_json_strings(record)) + if all(needle in text for needle in needles): + return True + return False + + def close(self) -> None: + self._temporary.cleanup() + + +def _jsonl_records(transcripts: dict[str, list[dict]]) -> list[dict]: + return [record for rows in transcripts.values() for record in rows] + + +def _tool_events(transcripts: dict[str, list[dict]]) -> tuple[dict[str, dict], dict[str, dict]]: + calls: dict[str, dict] = {} + results: dict[str, dict] = {} + for record in _jsonl_records(transcripts): + for node in _dict_nodes(record): + event_type = str(node.get("type", "")).lower() + if event_type in _CALL_TYPES: + identifier = _call_id(node) + if identifier: + calls[identifier] = node + elif event_type in _RESULT_TYPES: + identifier = _call_id(node) + if identifier: + results[identifier] = node + return calls, results + + +def correlated_tool_evidence( + transcripts: dict[str, list[dict]], tool_fragment: str, marker: str +) -> dict[str, object]: + calls, results = _tool_events(transcripts) + for identifier, call in calls.items(): + call_text = "\n".join(_json_strings(call)) + result = results.get(identifier) + result_text = "\n".join(_json_strings(result)) if result else "" + if tool_fragment in call_text and marker in call_text and marker in result_text: + return {"call_id": identifier, "call": call, "result": result} + raise AssertionError( + f"No correlated {tool_fragment!r} call/result carried marker {marker!r}; " + f"calls={list(calls)} results={list(results)}" + ) + + +def structured_final_answers( + agent: str, command: CommandResult, transcripts: dict[str, list[dict]] +) -> list[str]: + answers: list[str] = [] + for line in command.stdout.splitlines(): + with contextlib.suppress(json.JSONDecodeError): + row = json.loads(line) + if not isinstance(row, dict): + continue + if ( + agent == "claude" + and row.get("type") == "result" + and isinstance(row.get("result"), str) + ): + answers.append(row["result"]) + if agent == "codex" and row.get("type") == "item.completed": + item = row.get("item") + if isinstance(item, dict) and item.get("type") == "agent_message": + text = item.get("text") + if isinstance(text, str): + answers.append(text) + for record in _jsonl_records(transcripts): + if agent == "claude" and record.get("type") == "assistant": + message = record.get("message") + if isinstance(message, dict): + answers.extend( + text + for part in message.get("content", []) + if isinstance(part, dict) + for text in [part.get("text")] + if isinstance(text, str) + ) + if agent == "codex" and record.get("type") == "event_msg": + payload = record.get("payload") + if isinstance(payload, dict) and payload.get("type") == "task_complete": + text = payload.get("last_agent_message") + if isinstance(text, str): + answers.append(text) + return answers + + +class WorkspaceApi: + """Small JSON/raw adapter over a ``WorkspaceClient``'s public API client.""" + + def __init__(self, workspace_client): + self.client = workspace_client + + def request( + self, + method: str, + path: str, + *, + query: dict[str, object] | None = None, + body: dict | None = None, + headers: dict[str, str] | None = None, + ) -> object: + return self.client.api_client.do( + method, + path=path, + query=query, + body=body, + headers={"Accept": "application/json", **(headers or {})}, + ) + + def json( + self, + method: str, + path: str, + *, + query: dict[str, object] | None = None, + body: dict | None = None, + headers: dict[str, str] | None = None, + ) -> dict | list: + response = self.request( + method, + query=query, + body=body, + path=path, + headers=headers, + ) + if not isinstance(response, (dict, list)): + raise AssertionError(f"Expected JSON from {method} {path}, got {type(response)!r}") + return response + + def raw_bytes(self, path: str) -> bytes: + response = self.client.api_client.do("GET", path=path, raw=True) + stream = response.get("contents") if isinstance(response, dict) else response + if not hasattr(stream, "read"): + raise AssertionError(f"Expected streamed bytes from {path}, got {type(response)!r}") + try: + value = stream.read() + finally: + stream.close() + if not isinstance(value, bytes): + raise AssertionError(f"Expected bytes from {path}, got {type(value)!r}") + return value + + +def _sse_json(data: bytes) -> list[dict]: + text = data.decode("utf-8", errors="strict") + stripped = text.strip() + if stripped: + with contextlib.suppress(json.JSONDecodeError): + value = json.loads(stripped) + if isinstance(value, dict): + return [value] + events: list[dict] = [] + for line in text.splitlines(): + if not line.startswith("data:"): + continue + payload = line.removeprefix("data:").strip() + if not payload or payload == "[DONE]": + continue + value = json.loads(payload) + if isinstance(value, dict): + events.append(value) + return events + + +class McpRpc: + def __init__(self, api: WorkspaceApi, path: str): + self.api = api + self.path = path + self.session_id: str | None = None + self._next_id = 1 + + def request(self, method: str, params: dict | None = None) -> dict: + request_id = self._next_id + self._next_id += 1 + body = {"jsonrpc": "2.0", "id": request_id, "method": method, "params": params or {}} + headers = { + "Accept": "application/json, text/event-stream", + "Content-Type": "application/json", + } + if self.session_id: + headers["Mcp-Session-Id"] = self.session_id + response = self.api.client.api_client.do( + "POST", + path=self.path, + headers=headers, + body=body, + raw=True, + response_headers=["mcp-session-id"], + ) + if isinstance(response, dict): + self.session_id = response.get("mcp-session-id") or response.get("Mcp-Session-Id") + stream = response.get("contents") + else: + stream = response + if not hasattr(stream, "read"): + raise AssertionError(f"MCP {method} did not return a response stream") + try: + events = _sse_json(stream.read()) + finally: + stream.close() + matches = [event for event in events if event.get("id") == request_id] + if not matches: + raise AssertionError(f"MCP {method} returned no JSON-RPC response: {events!r}") + result = matches[-1] + if "error" in result: + raise AssertionError(f"MCP {method} failed: {result['error']!r}") + return result + + def notify(self, method: str, params: dict | None = None) -> None: + headers = { + "Accept": "application/json, text/event-stream", + "Content-Type": "application/json", + } + if self.session_id: + headers["Mcp-Session-Id"] = self.session_id + response = self.api.client.api_client.do( + "POST", + path=self.path, + headers=headers, + body={"jsonrpc": "2.0", "method": method, "params": params or {}}, + ) + # Streamable-HTTP servers commonly acknowledge notifications with HTTP 202 and no body; + # api_client.do returns None for that successful response and raises on HTTP failures. + if response is not None and not isinstance(response, (dict, list)): + raise AssertionError(f"MCP notification returned {type(response)!r}") + + +def mcp_tools(result: dict) -> list[dict]: + payload = result.get("result") + if not isinstance(payload, dict): + raise AssertionError(f"MCP response had no result object: {result!r}") + tools = payload.get("tools") + if not isinstance(tools, list) or not all(isinstance(tool, dict) for tool in tools): + raise AssertionError(f"MCP tools/list returned an invalid tool list: {result!r}") + return tools + + +def listed_mcp_names(output: str, dotted_name: str) -> bool: + dashed = dotted_name.replace(".", "-") + return any(dotted_name in line or dashed in line for line in output.splitlines()) diff --git a/tests/e2e_cuj/pytest.ini b/tests/e2e_cuj/pytest.ini new file mode 100644 index 000000000..39ddadd4f --- /dev/null +++ b/tests/e2e_cuj/pytest.ini @@ -0,0 +1,5 @@ +[pytest] +testpaths = . +addopts = -ra +markers = + cuj2: MPS plus explicitly selected MCP resources diff --git a/tests/e2e_cuj/test_cuj2_mps_mcp.py b/tests/e2e_cuj/test_cuj2_mps_mcp.py new file mode 100644 index 000000000..bb8cc0736 --- /dev/null +++ b/tests/e2e_cuj/test_cuj2_mps_mcp.py @@ -0,0 +1,428 @@ +"""CUJ2: configure two real agents from per-agent MPS resources with one MCP.""" + +from __future__ import annotations + +import json +import tomllib +import uuid +from collections.abc import Iterator +from contextlib import contextmanager + +import pytest +from databricks.sdk.service.catalog import ListModelProviderServicesRequestView + +from .base import BaseCujTest +from .helpers import ( + McpRpc, + WorkspaceApi, + correlated_tool_evidence, + listed_mcp_names, + mcp_tools, + structured_final_answers, +) +from .helpers.tui_request_recorder import TuiRequestRecorder + +WORKSPACE_URL = "https://dbc-0dcf95cf-e357.cloud.databricks.com" +AGENT_MPS: dict[str, tuple[str, str]] = { + "claude": ("ug_e2e.providers.anthropic", "claude-haiku-4-5-20251001"), + "codex": ("ug_e2e.providers.openai", "gpt-5-nano"), +} +MCP_NAME = "system.ai.sandbox" +NEGATIVE_MCP_NAME = "system.ai.web_search" +CONFIG_PATH = "/api/ai-gateway/v2/coding-agent-configs" +CONFIG_UPDATE_MASK = "default_agent,enabled_agents,mcp_servers" +CONFIG_MUTABLE_FIELDS = ("default_agent", "enabled_agents", "mcp_servers") +AGENT_NAMES = { + "claude": "CODING_AGENT_CLAUDE_CODE", + "codex": "CODING_AGENT_CODEX", +} + + +def _config_rows(payload: dict | list) -> list[dict]: + rows = payload.get("coding_agent_configs") if isinstance(payload, dict) else payload + if not isinstance(rows, list) or not all(isinstance(row, dict) for row in rows): + raise AssertionError(f"Invalid CodingAgentConfig listing: {payload!r}") + return rows + + +def _agent_config(agent: str) -> dict: + mps_name, model = AGENT_MPS[agent] + return { + "agent": AGENT_NAMES[agent], + "config": { + "models": {"model_provider_service": mps_name}, + "default_models": {"default_model": model}, + "smart_routing": {"enabled": False}, + "tracing": {"enabled": False}, + }, + } + + +def _desired_config() -> dict: + return { + "spec_version": 1, + "default_agent": AGENT_NAMES["codex"], + "enabled_agents": [_agent_config("claude"), _agent_config("codex")], + "mcp_servers": {"names": [MCP_NAME]}, + } + + +def _single_config(api: WorkspaceApi) -> dict | None: + rows = _config_rows(api.json("GET", CONFIG_PATH)) + if len(rows) > 1: + raise AssertionError(f"Workspace has more than one CodingAgentConfig: {rows!r}") + return rows[0] if rows else None + + +def _config_resource_path(name: str) -> str: + return f"{CONFIG_PATH}/{name.removeprefix('coding-agent-configs/').strip('/')}" + + +def _ensure_single_config(api: WorkspaceApi, desired: dict) -> dict: + original = _single_config(api) + if original is None: + api.json("POST", CONFIG_PATH, body=desired) + else: + name = original.get("name") + if not isinstance(name, str) or not name: + raise AssertionError(f"Existing CodingAgentConfig has no resource name: {original!r}") + api.json( + "PATCH", + _config_resource_path(name), + query={"update_mask": CONFIG_UPDATE_MASK}, + body={**desired, "name": name}, + ) + updated = _single_config(api) + if updated is None: + raise AssertionError("Create/update did not leave a CodingAgentConfig") + return updated + + +def _restore_config(api: WorkspaceApi, original: dict | None) -> None: + if original is None: + created = _single_config(api) + if created is None: + return + name = created.get("name") + if not isinstance(name, str) or not name: + raise AssertionError(f"Temporary CodingAgentConfig has no resource name: {created!r}") + api.request("DELETE", _config_resource_path(name)) + remaining = _single_config(api) + if remaining is not None: + raise AssertionError(f"Temporary CodingAgentConfig was not deleted: {remaining!r}") + return + + name = original.get("name") + if not isinstance(name, str) or not name: + raise AssertionError(f"Original CodingAgentConfig has no resource name: {original!r}") + body = { + key: original[key] for key in ("spec_version", *CONFIG_MUTABLE_FIELDS) if key in original + } + body["name"] = name + api.json( + "PATCH", + _config_resource_path(name), + query={"update_mask": CONFIG_UPDATE_MASK}, + body=body, + ) + restored = _single_config(api) + if restored is None: + raise AssertionError("Original CodingAgentConfig was not restored") + for field in CONFIG_MUTABLE_FIELDS: + assert restored.get(field) == original.get(field), ( + f"CodingAgentConfig field {field!r} was not restored: " + f"expected {original.get(field)!r}, got {restored.get(field)!r}" + ) + + +@contextmanager +def _published_config(api: WorkspaceApi) -> Iterator[dict]: + original = _single_config(api) + try: + yield _ensure_single_config(api, _desired_config()) + finally: + _restore_config(api, original) + + +def _assert_published_config(config: dict) -> None: + assert config.get("default_agent") == AGENT_NAMES["codex"], config + assert config.get("mcp_servers") == {"names": [MCP_NAME]}, config + + entries = config.get("enabled_agents") + assert isinstance(entries, list) and len(entries) == 2, config + by_agent = {entry.get("agent"): entry for entry in entries if isinstance(entry, dict)} + assert set(by_agent) == set(AGENT_NAMES.values()), config + for agent, agent_name in AGENT_NAMES.items(): + mps_name, model = AGENT_MPS[agent] + entry_config = by_agent[agent_name].get("config") + assert isinstance(entry_config, dict), by_agent[agent_name] + assert entry_config.get("models") == {"model_provider_service": mps_name}, entry_config + assert entry_config.get("default_models") == {"default_model": model}, entry_config + assert entry_config.get("smart_routing") == {"enabled": False}, entry_config + assert entry_config.get("tracing") == {"enabled": False}, entry_config + + +def _assert_mps_catalog(workspace_client) -> None: + listed = list( + workspace_client.ai_gateway.list_model_provider_services( + parent="schemas/ug_e2e.providers", + view=ListModelProviderServicesRequestView.FULL, + ) + ) + listed_by_name = { + item.name.removeprefix("model-provider-services/"): item for item in listed if item.name + } + for agent, (mps_name, model) in AGENT_MPS.items(): + service = workspace_client.ai_gateway.get_model_provider_service( + f"model-provider-services/{mps_name}" + ) + assert service.name and service.name.endswith(mps_name), service + config = service.config + assert config is not None, service + targets = config.targets or [] + models = [target.model for target in targets] + assert models == [model], service + assert config.allow_all_targets is not True, service + expected_type = ( + "EXTERNAL_MODEL_PROVIDER_TYPE_ANTHROPIC" + if agent == "claude" + else "EXTERNAL_MODEL_PROVIDER_TYPE_OPENAI" + ) + provider_type = getattr(config.provider_type, "value", config.provider_type) + assert provider_type == expected_type, config + expected_api = "anthropic/v1/messages" if agent == "claude" else "openai/v1/responses" + assert [api for target in targets for api in (target.native_api_types or [])] == [ + expected_api + ], targets + + matching = listed_by_name.get(mps_name) + assert matching is not None, [item.name for item in listed] + listed_targets = matching.config.targets if matching.config else None + assert [target.model for target in listed_targets or []] == [model], matching + + +def _assert_mcp_prerequisite(workspace_client, api: WorkspaceApi) -> None: + sandbox = workspace_client.ai_gateway.get_mcp_service(f"mcp-services/{MCP_NAME}") + web_search = workspace_client.ai_gateway.get_mcp_service(f"mcp-services/{NEGATIVE_MCP_NAME}") + assert sandbox.name and sandbox.name.endswith(MCP_NAME), sandbox + assert web_search.name and web_search.name.endswith(NEGATIVE_MCP_NAME), web_search + + endpoint = f"/ai-gateway/mcp-services/{MCP_NAME}" + rpc = McpRpc(api, endpoint) + initialize = rpc.request( + "initialize", + { + "protocolVersion": "2025-03-26", + "capabilities": {}, + "clientInfo": {"name": "ug-cuj2", "version": "1"}, + }, + ) + assert isinstance(initialize.get("result"), dict), initialize + rpc.notify("notifications/initialized") + tools = mcp_tools(rpc.request("tools/list")) + run_code = [tool for tool in tools if tool.get("name") == "run_code"] + assert len(run_code) == 1, tools + + +def _assert_configured_files(session, workspace: str) -> None: + managed_cache = json.loads( + (session.home / ".ucode" / "managed-config.json").read_text(encoding="utf-8") + ) + managed_config = managed_cache.get("config") + assert isinstance(managed_config, dict), managed_cache + _assert_published_config(managed_config) + + claude_settings = session.home / ".claude" / "ucode-settings.json" + assert claude_settings.is_file(), claude_settings + settings = json.loads(claude_settings.read_text(encoding="utf-8")) + claude_env = settings.get("env") + assert isinstance(claude_env, dict), settings + claude_mps, claude_model = AGENT_MPS["claude"] + custom_headers = claude_env.get("ANTHROPIC_CUSTOM_HEADERS") + assert isinstance(custom_headers, str), claude_env + claude_headers = {} + for line in custom_headers.splitlines(): + 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_env["ANTHROPIC_DEFAULT_HAIKU_MODEL"] == claude_model, claude_env + assert claude_env["ANTHROPIC_BASE_URL"] == workspace.rstrip("/") + "/ai-gateway/anthropic" + + _, codex_model = AGENT_MPS["codex"] + codex_path = session.home / ".codex" / "ucode.config.toml" + assert codex_path.is_file(), codex_path + codex_config = tomllib.loads(codex_path.read_text(encoding="utf-8")) + assert codex_config["model_provider"] == "Databricks", codex_config + assert codex_config["model"] == codex_model, codex_config + codex_provider = codex_config["model_providers"]["Databricks"] + assert codex_provider["wire_api"] == "responses", codex_provider + assert codex_provider["base_url"] == workspace.rstrip("/") + "/ai-gateway/codex/v1" + + +def _agent_prompt(marker: str) -> str: + return ( + "Complete this task using the configured model and MCP; do not answer from the prompt. " + f"Call the configured `run_code` MCP tool with a Python program that prints the exact " + f"marker `{marker}`. Wait for the real tool result. " + f"Final answer must contain the marker `{marker}`, with no invented tool output." + ) + + +def _run_agent_task(session, agent: str, marker: str): + prompt = _agent_prompt(marker) + + def completed(current) -> bool: + if not current.transcript_contains(agent, marker): + return False + try: + correlated_tool_evidence(current.transcripts(agent), "run_code", marker) + except AssertionError: + return False + return True + + return session.run_tui( + agent, + prompt, + completion=completed, + timeout=240, + ) + + +@pytest.mark.cuj2 +class TestCuj2MpsExplicitMcp(BaseCujTest): + WORKSPACE_URL = WORKSPACE_URL + + def test_configuration(self, live_session): + """Scenario: publish both agent-specific MPS targets and the sandbox MCP, then configure ug. + + Expected: the published config selects gpt-5-nano for Codex and + claude-haiku-4-5-20251001 for Claude, with routing and tracing off; + both native MPS APIs advertise their selected model, ug writes the two + agent configs, and both agents list sandbox but exclude web_search. + """ + session = live_session + api = WorkspaceApi(self.workspace) + _assert_mps_catalog(self.workspace) + _assert_mcp_prerequisite(self.workspace, api) + + with _published_config(api) as published: + _assert_published_config(published) + configured = session.run( + "configure", + "--workspace", + WORKSPACE_URL, + "--skip-upgrade", + timeout=300, + ) + assert "Select coding agents to configure:" not in configured.stdout, configured.stdout + _assert_configured_files(session, WORKSPACE_URL) + + claude_listing = session.external("claude", "mcp", "list", timeout=120) + claude_output = f"{claude_listing.stdout}\n{claude_listing.stderr}" + assert listed_mcp_names(claude_output, MCP_NAME), claude_output + assert not listed_mcp_names(claude_output, NEGATIVE_MCP_NAME), claude_output + + codex_listing = session.external("codex", "mcp", "list", timeout=120) + codex_output = f"{codex_listing.stdout}\n{codex_listing.stderr}" + assert listed_mcp_names(codex_output, MCP_NAME), codex_output + assert not listed_mcp_names(codex_output, NEGATIVE_MCP_NAME), codex_output + + def test_codex(self, live_session): + """Scenario: configure ug from the dedicated workspace and complete a real Codex task. + + Expected: Codex sends a Responses request for gpt-5-nano with the + ug_e2e.providers.openai target header and receives its paired HTTP 200 + response; its transcript contains a correlated sandbox run_code result + and its final answer contains the task marker. + """ + session = live_session + api = WorkspaceApi(self.workspace) + with _published_config(api) as published: + _assert_published_config(published) + with TuiRequestRecorder(WORKSPACE_URL) as recorder: + configured = session.run( + "configure", + "--workspace", + recorder.url, + "--skip-upgrade", + timeout=300, + ) + assert "Select coding agents to configure:" not in configured.stdout, ( + configured.stdout + ) + _assert_configured_files(session, recorder.url) + + marker = f"CUJ2-CODEX-{uuid.uuid4().hex}" + checkpoint = recorder.checkpoint() + result = _run_agent_task(session, "codex", marker) + request = recorder.expect_request( + method="POST", + path="/ai-gateway/codex/v1/responses", + after=checkpoint, + timeout=240, + ) + assert request.headers["databricks-model-provider-service"] == ( + "ug_e2e.providers.openai" + ) + assert request.payload["model"] == "gpt-5-nano" + 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, "Codex inference response was empty" + + transcripts = session.transcripts("codex") + assert transcripts, "No Codex transcript was written" + assert correlated_tool_evidence(transcripts, "run_code", marker) + answers = structured_final_answers("codex", result, transcripts) + assert any(marker in answer for answer in answers), answers + + def test_claude(self, live_session): + """Scenario: configure ug from the dedicated workspace and complete a real Claude task. + + Expected: Claude sends a Messages request for + claude-haiku-4-5-20251001 with the ug_e2e.providers.anthropic target + header and receives its paired HTTP 200 response; its transcript + contains a correlated sandbox run_code result and its final answer + contains the task marker. + """ + session = live_session + api = WorkspaceApi(self.workspace) + with _published_config(api) as published: + _assert_published_config(published) + with TuiRequestRecorder(WORKSPACE_URL) as recorder: + configured = session.run( + "configure", + "--workspace", + recorder.url, + "--skip-upgrade", + timeout=300, + ) + assert "Select coding agents to configure:" not in configured.stdout, ( + configured.stdout + ) + _assert_configured_files(session, recorder.url) + + marker = f"CUJ2-CLAUDE-{uuid.uuid4().hex}" + checkpoint = recorder.checkpoint() + result = _run_agent_task(session, "claude", marker) + request = recorder.expect_request( + method="POST", + path="/ai-gateway/anthropic/v1/messages", + after=checkpoint, + timeout=240, + ) + assert request.headers["databricks-model-provider-service"] == ( + "ug_e2e.providers.anthropic" + ) + assert request.payload["model"] == "claude-haiku-4-5-20251001" + 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, "Claude inference response was empty" + + transcripts = session.transcripts("claude") + assert transcripts, "No Claude transcript was written" + assert correlated_tool_evidence(transcripts, "run_code", marker) + answers = structured_final_answers("claude", result, transcripts) + assert any(marker in answer for answer in answers), answers From 7bb19cd2dc317d28e048a21c5f3a5a6f9f3c086f Mon Sep 17 00:00:00 2001 From: Mason Cao Date: Tue, 6 Oct 2026 00:07:42 +0000 Subject: [PATCH 02/19] [AIGTWY-4877] Use automatic CUJ discovery --- tests/README.md | 2 +- tests/e2e_cuj/test_cuj2_mps_mcp.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/README.md b/tests/README.md index 070516916..11a2495d7 100644 --- a/tests/README.md +++ b/tests/README.md @@ -137,7 +137,7 @@ Dedicated-workspace CUJs live in `e2e_cuj/` and use their own black-box helpers. | Test | User action | Expected evidence | | --- | --- | --- | -| `test_configuration`, `test_codex`, `test_claude` in `test_cuj2_mps_mcp.py` | Temporarily publish the two-agent MPS/MCP config; run separate configure, Codex TUI, and Claude TUI cases | Configuration verifies both native MPS APIs, generated settings, and sandbox inclusion with web_search excluded; each TUI case records its exact MPS target header and model, paired HTTP 200 response, correlated `run_code` call/result, and structured final marker | +| `test_cuj_configuration`, `test_cuj_codex_inference`, `test_cuj_claude_inference` in `test_cuj2_mps_mcp.py` | Temporarily publish the two-agent MPS/MCP config; run separate configure, Codex TUI, and Claude TUI cases | Configuration verifies both native MPS APIs, generated settings, and sandbox inclusion with web_search excluded; each TUI case records its exact MPS target header and model, paired HTTP 200 response, correlated `run_code` call/result, and structured final marker | | `test_ug_configure_claude_databricks` | Configure Databricks Hosted; execute the generated auth helper; launch plain `ug claude`, read a file, and open `/model` | Generated helper invokes `ug` with clean token stdout; assistant returns an unpredictable file value; native discovery caches `system.ai` models and the picker shows a discovered model without an opt-in flag; normal exit; reopen with working keyboard input | | `test_ug_configure_claude_anthropic_mps` | Select Anthropic MPS in the real configure picker; launch Claude | Saved provider in status; completed TUI file task; normal exit | | `test_ug_configure_codex_databricks` | Configure Databricks Hosted; execute the generated auth helper; open Codex TUI and read a file | Generated helper invokes `ug` with clean token stdout; completed assistant answer contains the file value; normal exit and reopen | diff --git a/tests/e2e_cuj/test_cuj2_mps_mcp.py b/tests/e2e_cuj/test_cuj2_mps_mcp.py index bb8cc0736..90a56446b 100644 --- a/tests/e2e_cuj/test_cuj2_mps_mcp.py +++ b/tests/e2e_cuj/test_cuj2_mps_mcp.py @@ -293,7 +293,7 @@ def completed(current) -> bool: class TestCuj2MpsExplicitMcp(BaseCujTest): WORKSPACE_URL = WORKSPACE_URL - def test_configuration(self, live_session): + def test_cuj_configuration(self, live_session): """Scenario: publish both agent-specific MPS targets and the sandbox MCP, then configure ug. Expected: the published config selects gpt-5-nano for Codex and @@ -328,7 +328,7 @@ def test_configuration(self, live_session): assert listed_mcp_names(codex_output, MCP_NAME), codex_output assert not listed_mcp_names(codex_output, NEGATIVE_MCP_NAME), codex_output - def test_codex(self, live_session): + def test_cuj_codex_inference(self, live_session): """Scenario: configure ug from the dedicated workspace and complete a real Codex task. Expected: Codex sends a Responses request for gpt-5-nano with the @@ -377,7 +377,7 @@ def test_codex(self, live_session): answers = structured_final_answers("codex", result, transcripts) assert any(marker in answer for answer in answers), answers - def test_claude(self, live_session): + def test_cuj_claude_inference(self, live_session): """Scenario: configure ug from the dedicated workspace and complete a real Claude task. Expected: Claude sends a Messages request for From 4e047963c4a829c2c9b1e8f21125a92f5200662f Mon Sep 17 00:00:00 2001 From: Mason Cao Date: Tue, 6 Oct 2026 00:31:22 +0000 Subject: [PATCH 03/19] [AIGTWY-4877] Reuse dedicated CUJ framework --- tests/README.md | 47 +- tests/e2e_cuj/AGENTS.md | 10 + tests/e2e_cuj/helpers/__init__.py | 792 ----------------------------- tests/e2e_cuj/helpers/constants.py | 37 +- tests/e2e_cuj/pytest.ini | 5 - tests/e2e_cuj/test_cuj2_mps_mcp.py | 610 +++++++++++----------- tests/integration/README.md | 10 +- 7 files changed, 370 insertions(+), 1141 deletions(-) delete mode 100644 tests/e2e_cuj/helpers/__init__.py delete mode 100644 tests/e2e_cuj/pytest.ini diff --git a/tests/README.md b/tests/README.md index 11a2495d7..ec3692f59 100644 --- a/tests/README.md +++ b/tests/README.md @@ -7,19 +7,19 @@ mocks, monkeypatching, fake binaries/services, or fabricated ug state. | Category | Location | What it proves | | --- | --- | --- | | Unit/component | Existing `test_*.py` files | Individual behavior; dependencies may be mocked | -| Dedicated-workspace CUJ | `e2e_cuj/test_cuj_*.py` | Real user journeys against read-only, preconfigured workspaces | +| Dedicated-workspace CUJs | `e2e_cuj/test_*.py` | Real routing, MPS, MCP, and two-agent inference with read-only workspace configuration | | Existing e2e | `test_e2e*.py` | Real workspace behavior with some patched setup/internal calls | | Integration CUJs | `integration/test_*.py` | Public configure, TUI, script, command, protocol, and lifecycle journeys | -| Dedicated-workspace CUJs | `e2e_cuj/test_*.py` | Real published CodingAgentConfig, MPS, MCP, and two-agent inference | | Installation | `integration/test_installation.py` | Fresh installed package, CLI, and local helpers without credentials on Linux and advisory native Windows | -Dedicated CUJs reuse `integration/utils` session/terminal mechanics and file-task -and transcript readers, not its config stubs or pytest fixtures. Prompt/model -correlation remains CUJ-specific. Concurrent runs may read the same CUJ workspace. -Workspace operations reuse the base class's authenticated Databricks SDK client; -offline tests require GET-only API calls and verify config changes fail without repair. -The live fixture compares configuration before and after the journey, even on failure. -CUJs never republish configuration or create a remote reservation. +Dedicated CUJs reuse `integration/utils` session/terminal mechanics, file-task, and transcript +readers, not its config stubs or pytest fixtures. Prompt/model correlation remains +CUJ-specific. Concurrent runs may read the same CUJ workspace. Workspace operations reuse +the base class's authenticated Databricks SDK client; offline tests require GET-only API +calls and verify config changes fail without repair. The live fixture compares configuration +before and after every journey, even on failure. +Routing-disabled coverage is deferred until a separately preconfigured workspace is assigned; +these CUJs never republish configuration or create remote reservations. CUJ helper tests also verify that unsupported agent names fail rather than defaulting to Codex. They cover Claude/Codex helper dispatch and rejection of routing decisions without the agent-specific prompt-submission evidence. @@ -133,11 +133,12 @@ Consult the run's JUnit report and artifacts for results. Each function states its **Scenario** and **Expected** outcome and shows its configure and launch commands. Fixtures supply fresh environments and credentials, never configured ug. Integration tests live directly in `integration/`; shared mechanics live in `utils/`. -Dedicated-workspace CUJs live in `e2e_cuj/` and use their own black-box helpers. +Dedicated-workspace CUJs live in `e2e_cuj/` and use the shared `cuj` fixture plus +integration utilities; only CUJ-specific evidence correlation stays in a test file. | Test | User action | Expected evidence | | --- | --- | --- | -| `test_cuj_configuration`, `test_cuj_codex_inference`, `test_cuj_claude_inference` in `test_cuj2_mps_mcp.py` | Temporarily publish the two-agent MPS/MCP config; run separate configure, Codex TUI, and Claude TUI cases | Configuration verifies both native MPS APIs, generated settings, and sandbox inclusion with web_search excluded; each TUI case records its exact MPS target header and model, paired HTTP 200 response, correlated `run_code` call/result, and structured final marker | +| `test_cuj_configuration`, `test_cuj_codex_inference`, `test_cuj_claude_inference` in `test_cuj2_mps_mcp.py` | Read the permanently preconfigured two-agent MPS/MCP config; run separate configure, Codex TUI, and Claude TUI cases | Configuration verifies both native MPS APIs, generated settings, and sandbox inclusion with web_search excluded; each TUI case records its exact MPS target header and model, paired HTTP 200 response, correlated `run_code` call/result, and final marker. No workspace config CRUD is performed | | `test_ug_configure_claude_databricks` | Configure Databricks Hosted; execute the generated auth helper; launch plain `ug claude`, read a file, and open `/model` | Generated helper invokes `ug` with clean token stdout; assistant returns an unpredictable file value; native discovery caches `system.ai` models and the picker shows a discovered model without an opt-in flag; normal exit; reopen with working keyboard input | | `test_ug_configure_claude_anthropic_mps` | Select Anthropic MPS in the real configure picker; launch Claude | Saved provider in status; completed TUI file task; normal exit | | `test_ug_configure_codex_databricks` | Configure Databricks Hosted; execute the generated auth helper; open Codex TUI and read a file | Generated helper invokes `ug` with clean token stdout; completed assistant answer contains the file value; normal exit and reopen | @@ -213,17 +214,19 @@ Managed discovery expectations come from separate read-only, provider-scoped model-list requests; they do not rely solely on ug's generated catalog. The dedicated-workspace CUJ suite adds three independently runnable CUJ2 MPS+MCP cases: -one configuration case and separate Codex and Claude TUI task cases. Each temporarily -publishes the two-agent config and restores the original workspace config. The TUI cases -forward agent requests through a per-test loopback recorder to the real workspace and -assert the provider-service header, model, paired HTTP 200 response, correlated MCP tool -result, and final answer. Failed TUI runs retain terminal output/actions/screen and -best-effort agent transcript snapshots under the runner artifact directory before the -temporary homes are deleted. The suite requires -`UG_CUJ_SP_CLIENT_ID`, `UG_CUJ_SP_CLIENT_SECRET`, `UG_INTEGRATION_BIN`, -`UG_INTEGRATION_RUN_DIR`, and runner-provided `DATABRICKS_BEARER`; run it with -`uv run pytest --confcutdir=tests/e2e_cuj tests/e2e_cuj`. It owns the published -config for the test duration and does not use `UCODE_MANAGED_CONFIG_STUB`. +one configuration case and separate Codex and Claude TUI task cases. The workspace is +preconfigured for CUJ2 and is read-only; every case validates it, and the fixture compares +it again during teardown. The TUI cases forward agent requests through the shared per-test +loopback recorder to the real workspace and assert the provider-service header, model, +paired HTTP 200 response, correlated MCP tool result, and final answer. Failed TUI runs +retain terminal output/actions/screen and best-effort agent transcript snapshots under the +runner artifact directory before the temporary homes are deleted. The suite requires +`UG_CUJ_SP_CLIENT_ID` and `UG_CUJ_SP_CLIENT_SECRET`, with `ug`, `claude`, `codex`, and +`databricks` on PATH; run it with +`uv run --with pexpect==4.9.0 --with pyte==0.8.2 pytest --confcutdir=tests/e2e_cuj tests/e2e_cuj`. +Pytest prints the per-test artifact directory; the fixture forwards only its short-lived +bearer to the isolated session. It reads the published config as an input only and does not +use `UCODE_MANAGED_CONFIG_STUB`. ug no longer runs a post-configure agent probe; the deprecated `--skip-validate` flag is accepted as a no-op where older journeys still pass it. Tests retain diff --git a/tests/e2e_cuj/AGENTS.md b/tests/e2e_cuj/AGENTS.md index 448158862..9471024dc 100644 --- a/tests/e2e_cuj/AGENTS.md +++ b/tests/e2e_cuj/AGENTS.md @@ -3,6 +3,8 @@ 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. +The workspaces share one metastore, so metastore-scoped fixture names belong in +`helpers/constants.py`; keep each workspace URL in its test class. All workspace configurations are READ ONLY. Never create, publish, update, delete, or restore workspace configuration, including during setup or teardown. Validate the preconfigured workspace and fail on a mismatch; never repair it. Concurrent @@ -49,6 +51,14 @@ model, for both Claude and Codex. It uses the published models/defaults and asse configuration is unchanged, including during teardown after failures. Routing-disabled coverage requires a separately preconfigured workspace and is not part of this CUJ. +`test_cuj2_mps_mcp.py` runs three independently collected cases against a permanently +preconfigured workspace: configuration, Codex inference, and Claude inference. 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; only +CUJ2's `run_code` call/result correlation remains local to the test. + Use a clean disposable POSIX runner without existing machine-wide agent settings. Install the intended versions of `ug`, `claude`, `codex`, and `databricks` on PATH. The terminal helpers also require `pexpect==4.9.0` and `pyte==0.8.2`: diff --git a/tests/e2e_cuj/helpers/__init__.py b/tests/e2e_cuj/helpers/__init__.py deleted file mode 100644 index 8713dbb03..000000000 --- a/tests/e2e_cuj/helpers/__init__.py +++ /dev/null @@ -1,792 +0,0 @@ -"""Black-box helpers for dedicated-workspace CUJs. - -This module talks to the installed commands and the Databricks SDK only. It deliberately does -not import ``ucode``: a CUJ must exercise the installed wheel just as a user does. -""" - -from __future__ import annotations - -import contextlib -import importlib -import json -import os -import re -import signal -import subprocess -import tempfile -import time -from collections.abc import Callable -from dataclasses import dataclass -from pathlib import Path -from typing import Any - -_ANSI = re.compile(r"\x1b\[[0-9;?]*[ -/]*[@-~]") -_CALL_TYPES = { - "custom_tool_call", - "function_call", - "mcp_tool_call", - "tool_call", - "tool_use", -} -_RESULT_TYPES = { - "custom_tool_call_output", - "function_call_output", - "mcp_tool_result", - "mcp_tool_call_output", - "tool_result", - "tool_call_output", -} - - -def _kill_process_group(process: subprocess.Popen[str]) -> None: - if process.poll() is not None: - return - if os.name == "posix": - with contextlib.suppress(ProcessLookupError): - os.killpg(process.pid, signal.SIGTERM) - try: - process.wait(timeout=5) - except subprocess.TimeoutExpired: - with contextlib.suppress(ProcessLookupError): - os.killpg(process.pid, signal.SIGKILL) - elif os.name == "nt": - process.kill() - process.wait(timeout=5) - - -def _process_group_options() -> dict[str, Any]: - if os.name == "posix": - return {"start_new_session": True} - if os.name == "nt": - return {"creationflags": subprocess.CREATE_NEW_PROCESS_GROUP} - return {} - - -def _json_strings(value: object) -> list[str]: - if isinstance(value, str): - return [value] - if isinstance(value, dict): - return [text for item in value.values() for text in _json_strings(item)] - if isinstance(value, list): - return [text for item in value for text in _json_strings(item)] - return [] - - -def _dict_nodes(value: object): - if isinstance(value, dict): - yield value - for child in value.values(): - yield from _dict_nodes(child) - elif isinstance(value, list): - for child in value: - yield from _dict_nodes(child) - - -def _first_string(mapping: dict, *keys: str) -> str | None: - for key in keys: - value = mapping.get(key) - if isinstance(value, str) and value: - return value - return None - - -def _call_id(mapping: dict) -> str | None: - return _first_string(mapping, "id", "call_id", "callId", "tool_use_id", "toolUseId") - - -@dataclass(frozen=True) -class CommandResult: - argv: tuple[str, ...] - returncode: int - stdout: str - stderr: str - timed_out: bool = False - - -class CujSession: - """An isolated home/project and subprocess runner for one CUJ case.""" - - def __init__(self, run_root: Path, binary: Path, case_name: str): - self._temporary = tempfile.TemporaryDirectory(prefix="cuj-", dir=run_root) - self.root = Path(self._temporary.name) - self.home = self.root / "home" - self.project = self.root / "project with spaces" - self.home.mkdir() - self.project.mkdir() - self.binary = binary - self.artifacts = run_root / "artifacts" / case_name - self.artifacts.mkdir(parents=True, exist_ok=True) - self._commands = 0 - self.env = self._environment() - - def _environment(self) -> dict[str, str]: - keep = ( - "PATH", - "SYSTEMROOT", - "COMSPEC", - "LANG", - "LC_ALL", - "SSL_CERT_FILE", - "SSL_CERT_DIR", - "REQUESTS_CA_BUNDLE", - "NODE_EXTRA_CA_CERTS", - "HTTPS_PROXY", - "HTTP_PROXY", - "NO_PROXY", - ) - env = {key: os.environ[key] for key in keep if key in os.environ} - bearer = os.environ.get("DATABRICKS_BEARER", "").strip() - if not bearer: - raise AssertionError("DATABRICKS_BEARER must be provided by the CUJ runner") - env.update( - { - "HOME": str(self.home), - "USERPROFILE": str(self.home), - "XDG_CONFIG_HOME": str(self.home / ".config"), - "XDG_CACHE_HOME": str(self.home / ".cache"), - "XDG_DATA_HOME": str(self.home / ".local/share"), - "CLAUDE_CONFIG_DIR": str(self.home / ".claude"), - "CODEX_HOME": str(self.home / ".codex"), - "DATABRICKS_CONFIG_FILE": str(self.home / ".databrickscfg"), - "DATABRICKS_BEARER": bearer, - "UG_INTEGRATION_BIN": str(self.binary), - "DISABLE_AUTOUPDATER": "1", - "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1", - "NO_COLOR": "1", - "TERM": "dumb", - "COLUMNS": "160", - "ENABLE_SMART_ROUTING_V2": "0", - "PYTHONNOUSERSITE": "1", - } - ) - return env - - def _redact(self, text: str) -> str: - for name in ("DATABRICKS_BEARER", "UG_CUJ_SP_CLIENT_SECRET"): - token = os.environ.get(name) or self.env.get(name) - if token: - text = text.replace(token, "") - return _ANSI.sub("", text) - - def record(self, name: str, value: object) -> None: - payload = json.dumps(value, indent=2, sort_keys=True, default=str) - self.artifacts.joinpath(name).write_text(self._redact(payload), encoding="utf-8") - - def run( - self, - *args: str, - timeout: int = 180, - binary: str | Path | None = None, - input_text: str | None = None, - check: bool = True, - ) -> CommandResult: - command = tuple(str(binary or self.binary) for _ in [0]) + tuple(args) - process = subprocess.Popen( - list(command), - cwd=self.project, - env=self.env, - stdin=subprocess.PIPE if input_text is not None else subprocess.DEVNULL, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - text=True, - **_process_group_options(), - ) - timed_out = False - try: - stdout, stderr = process.communicate(input=input_text, timeout=timeout) - except subprocess.TimeoutExpired: - timed_out = True - _kill_process_group(process) - stdout, stderr = process.communicate(timeout=5) - finally: - _kill_process_group(process) - result = CommandResult( - command, - process.returncode, - self._redact(stdout), - self._redact(stderr), - timed_out, - ) - self._commands += 1 - self.record( - f"command-{self._commands}.json", - { - "argv": command, - "returncode": result.returncode, - "stderr": result.stderr, - "stdout": result.stdout, - "timed_out": result.timed_out, - }, - ) - if result.timed_out: - raise AssertionError(f"Command exceeded {timeout}s: {command}") - if check and result.returncode != 0: - raise AssertionError( - f"Command failed ({result.returncode}): {command}\n{result.stderr}" - ) - return result - - def external(self, program: str, *args: str, timeout: int = 120) -> CommandResult: - return self.run(*args, binary=program, timeout=timeout) - - def _record_tui_artifact( - self, - command: tuple[str, ...], - *, - stdout: str, - returncode: int | None, - timed_out: bool, - actions: list[dict[str, str]], - screen: str, - failure: BaseException | None, - ) -> None: - self._commands += 1 - try: - self.record( - f"command-{self._commands}.json", - { - "argv": command, - "returncode": returncode, - "stderr": "", - "stdout": stdout, - "timed_out": timed_out, - "interactive": True, - "actions": actions, - "screen": screen, - "error": f"{type(failure).__name__}: {failure}" if failure else None, - }, - ) - except BaseException: - if failure is None: - raise - - def run_tui( - self, - agent: str, - prompt: str, - *, - completion: Callable[[CujSession], bool], - timeout: int = 240, - ) -> CommandResult: - """Run a real interactive agent through a PTY and return its terminal transcript.""" - if os.name != "posix": - raise AssertionError("The dedicated CUJ requires POSIX PTY support.") - try: - pexpect = importlib.import_module("pexpect") - pyte = importlib.import_module("pyte") - except ImportError as exc: - raise AssertionError("The dedicated CUJ runner must install pexpect and pyte.") from exc - - if agent not in {"claude", "codex"}: - raise AssertionError(f"Unsupported TUI agent: {agent}") - # Bare ug resolves the workspace default (Codex in CUJ2); Claude needs its explicit - # subcommand so this path exercises the configured Claude provider too. - command = (str(self.binary),) if agent == "codex" else (str(self.binary), "claude") - env = {**self.env, "TERM": "xterm-256color", "COLUMNS": "140", "LINES": "60"} - try: - child = pexpect.spawn( - command[0], - command[1:], - cwd=str(self.project), - env=env, - encoding="utf-8", - codec_errors="replace", - dimensions=(60, 140), - timeout=1, - ) - except BaseException as exc: - self._record_tui_artifact( - command, - stdout="", - returncode=None, - timed_out=False, - actions=[], - screen="", - failure=exc, - ) - raise - - def send_terminal_reply(data: str) -> None: - child.send(data) - - # pyte's screen emits terminal-device replies through this stream hook. The callback is - # deliberately limited to replies generated from the rendered screen state. - class TerminalScreen(pyte.Screen): - def write_process_input(self, data: str) -> None: - send_terminal_reply(data) - - screen = TerminalScreen(140, 60) - stream = pyte.Stream(screen) - output: list[str] = [] - actions: list[dict[str, str]] = [] - ended = False - timed_out = False - - def visible() -> str: - return "\n".join(line.rstrip() for line in screen.display) - - def read() -> None: - nonlocal ended - try: - chunk = child.read_nonblocking(size=65536, timeout=0.2) - except pexpect.TIMEOUT: - return - except pexpect.EOF: - before = child.before - if before: - output.append(before) - stream.feed(before) - ended = True - return - output.append(chunk) - stream.feed(chunk) - - def send(keys: str, reason: str) -> None: - actions.append({"reason": reason, "keys": keys, "screen_before": visible()}) - child.send(keys) - - def selected_line() -> str: - return next( - (line.strip() for line in visible().splitlines() if re.match(r"^\s*[›❯>]", line)), - "", - ) - - def choose(prompt_text: str, label: str) -> None: - nonlocal timed_out - deadline = time.monotonic() + 120 - visited: set[str] = set() - while time.monotonic() < deadline: - read() - current_screen = visible() - assert not ended, f"TUI exited while waiting for {prompt_text}:\n{current_screen}" - current = selected_line() - if prompt_text in current_screen and current: - if label in current: - send("\r", f"choose {label}") - return - assert current not in visited, f"Menu does not offer {label}:\n{current_screen}" - visited.add(current) - send("\x1b[B", f"move towards {label}") - time.sleep(0.05) - timed_out = True - raise AssertionError(f"TUI did not show {prompt_text} within 120s:\n{visible()}") - - def wait_for_prompt() -> None: - nonlocal timed_out - deadline = time.monotonic() + 120 - handled: set[str] = set() - ready_since: float | None = None - while time.monotonic() < deadline: - read() - assert not ended, f"TUI exited before its prompt:\n{visible()}" - current_screen = visible() - if "Accessing workspace:" in current_screen and self.project.name in current_screen: - choose("Accessing workspace:", "Yes, I trust this folder") - continue - if "Hooks need review" in current_screen: - choose("Hooks need review", "Trust all and continue") - continue - dialogs = [ - ( - "theme", - "Choose the text style" in current_screen and "Dark mode" in current_screen, - ), - ( - "security-notes", - "Security notes" in current_screen - and "Enter to continue" in current_screen, - ), - ( - "trust-folder", - self.project.name in current_screen - and bool( - re.search(r"1[.)]\s+Yes, I trust (?:this|the) folder", current_screen) - ), - ), - ( - "trust-directory", - self.project.name in current_screen - and "trust" in current_screen.lower() - and bool(re.search(r"1[.)]\s+Yes, (?:continue|proceed)", current_screen)), - ), - ] - handled_dialog = False - for label, shown in dialogs: - if shown: - handled_dialog = True - if label not in handled: - send("\r", label) - handled.add(label) - break - if handled_dialog: - continue - assert "Select login method:" not in current_screen, ( - f"Configured ug launched {agent}'s account-login flow instead of its " - "gateway session:\n" + current_screen - ) - if agent == "codex": - assert not ( - "Sign in with ChatGPT" in current_screen - and "Provide your own API key" in current_screen - ), "Configured ug launched Codex's account-login flow:\n" + current_screen - title = "Claude Code" if agent == "claude" else "Codex" - title_ready = title in current_screen and "loading" not in current_screen.lower() - if title_ready and re.search(r"(?m)^\s*[❯›>]\s*(?!\d+[.)])", current_screen): - ready_since = ready_since or time.monotonic() - if time.monotonic() - ready_since >= 1: - actions.append({"reason": "prompt-ready", "screen": current_screen}) - return - else: - ready_since = None - time.sleep(0.05) - timed_out = True - raise AssertionError(f"TUI did not reach a usable prompt within 120s:\n{visible()}") - - def wait_for_completion() -> None: - nonlocal timed_out - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - read() - if ended: - raise AssertionError(f"TUI exited before the task completed:\n{visible()}") - if completion(self): - return - time.sleep(0.05) - timed_out = True - raise AssertionError(f"TUI did not complete within {timeout}s:\n{visible()}") - - failure: BaseException | None = None - returncode: int | None = None - try: - wait_for_prompt() - send(prompt, "type task prompt") - deadline = time.monotonic() + 30 - while time.monotonic() < deadline: - read() - if prompt.split()[-1] in visible(): - break - assert not ended, f"TUI exited while rendering the task prompt:\n{visible()}" - time.sleep(0.05) - else: - timed_out = True - raise AssertionError(f"TUI did not render the task prompt:\n{visible()}") - send("\r", "submit task prompt") - wait_for_completion() - send("/exit", "type exit command") - send("\r", "submit exit command") - deadline = time.monotonic() + 30 - while not ended and time.monotonic() < deadline: - read() - time.sleep(0.05) - if not ended: - timed_out = True - raise AssertionError(f"TUI did not exit within 30s:\n{visible()}") - child.close(force=False) - returncode = child.exitstatus if child.exitstatus is not None else 1 - if returncode != 0: - raise AssertionError( - f"Interactive command failed ({returncode}): {command}\n{visible()}" - ) - return CommandResult( - command, - returncode, - self._redact("".join(output)), - "", - ) - except BaseException as exc: - failure = exc - raise - finally: - if not ended: - with contextlib.suppress(OSError): - os.killpg(child.pid, signal.SIGTERM) - with contextlib.suppress(OSError): - os.killpg(child.pid, signal.SIGKILL) - with contextlib.suppress(Exception): - child.close(force=True) - self._record_tui_artifact( - command, - stdout="".join(output), - returncode=returncode if returncode is not None else child.exitstatus, - timed_out=timed_out, - actions=actions, - screen=visible(), - failure=failure, - ) - - def read_json(self, relative_path: str) -> object: - return json.loads(self.home.joinpath(relative_path).read_text(encoding="utf-8")) - - def transcripts(self, agent: str) -> dict[str, list[dict]]: - root = self.home / (".claude/projects" if agent == "claude" else ".codex/sessions") - result: dict[str, list[dict]] = {} - paths = sorted(root.rglob("*.jsonl")) if root.is_dir() else [] - for path in paths: - records: list[dict] = [] - for line_number, line in enumerate(path.read_text(encoding="utf-8").splitlines(), 1): - if not line.strip(): - continue - try: - row = json.loads(line) - except json.JSONDecodeError as exc: - raise AssertionError( - f"Invalid {agent} transcript {path}:{line_number}" - ) from exc - if isinstance(row, dict): - records.append(row) - result[str(path.relative_to(root))] = records - return result - - def transcript_contains(self, agent: str, *needles: str) -> bool: - """Return true when one complete transcript record contains every needle.""" - root = self.home / (".claude/projects" if agent == "claude" else ".codex/sessions") - if not root.is_dir(): - return False - for path in root.rglob("*.jsonl"): - for line in path.read_text(encoding="utf-8").splitlines(): - if not line.strip(): - continue - with contextlib.suppress(json.JSONDecodeError): - record = json.loads(line) - if isinstance(record, dict): - text = "\n".join(_json_strings(record)) - if all(needle in text for needle in needles): - return True - return False - - def close(self) -> None: - self._temporary.cleanup() - - -def _jsonl_records(transcripts: dict[str, list[dict]]) -> list[dict]: - return [record for rows in transcripts.values() for record in rows] - - -def _tool_events(transcripts: dict[str, list[dict]]) -> tuple[dict[str, dict], dict[str, dict]]: - calls: dict[str, dict] = {} - results: dict[str, dict] = {} - for record in _jsonl_records(transcripts): - for node in _dict_nodes(record): - event_type = str(node.get("type", "")).lower() - if event_type in _CALL_TYPES: - identifier = _call_id(node) - if identifier: - calls[identifier] = node - elif event_type in _RESULT_TYPES: - identifier = _call_id(node) - if identifier: - results[identifier] = node - return calls, results - - -def correlated_tool_evidence( - transcripts: dict[str, list[dict]], tool_fragment: str, marker: str -) -> dict[str, object]: - calls, results = _tool_events(transcripts) - for identifier, call in calls.items(): - call_text = "\n".join(_json_strings(call)) - result = results.get(identifier) - result_text = "\n".join(_json_strings(result)) if result else "" - if tool_fragment in call_text and marker in call_text and marker in result_text: - return {"call_id": identifier, "call": call, "result": result} - raise AssertionError( - f"No correlated {tool_fragment!r} call/result carried marker {marker!r}; " - f"calls={list(calls)} results={list(results)}" - ) - - -def structured_final_answers( - agent: str, command: CommandResult, transcripts: dict[str, list[dict]] -) -> list[str]: - answers: list[str] = [] - for line in command.stdout.splitlines(): - with contextlib.suppress(json.JSONDecodeError): - row = json.loads(line) - if not isinstance(row, dict): - continue - if ( - agent == "claude" - and row.get("type") == "result" - and isinstance(row.get("result"), str) - ): - answers.append(row["result"]) - if agent == "codex" and row.get("type") == "item.completed": - item = row.get("item") - if isinstance(item, dict) and item.get("type") == "agent_message": - text = item.get("text") - if isinstance(text, str): - answers.append(text) - for record in _jsonl_records(transcripts): - if agent == "claude" and record.get("type") == "assistant": - message = record.get("message") - if isinstance(message, dict): - answers.extend( - text - for part in message.get("content", []) - if isinstance(part, dict) - for text in [part.get("text")] - if isinstance(text, str) - ) - if agent == "codex" and record.get("type") == "event_msg": - payload = record.get("payload") - if isinstance(payload, dict) and payload.get("type") == "task_complete": - text = payload.get("last_agent_message") - if isinstance(text, str): - answers.append(text) - return answers - - -class WorkspaceApi: - """Small JSON/raw adapter over a ``WorkspaceClient``'s public API client.""" - - def __init__(self, workspace_client): - self.client = workspace_client - - def request( - self, - method: str, - path: str, - *, - query: dict[str, object] | None = None, - body: dict | None = None, - headers: dict[str, str] | None = None, - ) -> object: - return self.client.api_client.do( - method, - path=path, - query=query, - body=body, - headers={"Accept": "application/json", **(headers or {})}, - ) - - def json( - self, - method: str, - path: str, - *, - query: dict[str, object] | None = None, - body: dict | None = None, - headers: dict[str, str] | None = None, - ) -> dict | list: - response = self.request( - method, - query=query, - body=body, - path=path, - headers=headers, - ) - if not isinstance(response, (dict, list)): - raise AssertionError(f"Expected JSON from {method} {path}, got {type(response)!r}") - return response - - def raw_bytes(self, path: str) -> bytes: - response = self.client.api_client.do("GET", path=path, raw=True) - stream = response.get("contents") if isinstance(response, dict) else response - if not hasattr(stream, "read"): - raise AssertionError(f"Expected streamed bytes from {path}, got {type(response)!r}") - try: - value = stream.read() - finally: - stream.close() - if not isinstance(value, bytes): - raise AssertionError(f"Expected bytes from {path}, got {type(value)!r}") - return value - - -def _sse_json(data: bytes) -> list[dict]: - text = data.decode("utf-8", errors="strict") - stripped = text.strip() - if stripped: - with contextlib.suppress(json.JSONDecodeError): - value = json.loads(stripped) - if isinstance(value, dict): - return [value] - events: list[dict] = [] - for line in text.splitlines(): - if not line.startswith("data:"): - continue - payload = line.removeprefix("data:").strip() - if not payload or payload == "[DONE]": - continue - value = json.loads(payload) - if isinstance(value, dict): - events.append(value) - return events - - -class McpRpc: - def __init__(self, api: WorkspaceApi, path: str): - self.api = api - self.path = path - self.session_id: str | None = None - self._next_id = 1 - - def request(self, method: str, params: dict | None = None) -> dict: - request_id = self._next_id - self._next_id += 1 - body = {"jsonrpc": "2.0", "id": request_id, "method": method, "params": params or {}} - headers = { - "Accept": "application/json, text/event-stream", - "Content-Type": "application/json", - } - if self.session_id: - headers["Mcp-Session-Id"] = self.session_id - response = self.api.client.api_client.do( - "POST", - path=self.path, - headers=headers, - body=body, - raw=True, - response_headers=["mcp-session-id"], - ) - if isinstance(response, dict): - self.session_id = response.get("mcp-session-id") or response.get("Mcp-Session-Id") - stream = response.get("contents") - else: - stream = response - if not hasattr(stream, "read"): - raise AssertionError(f"MCP {method} did not return a response stream") - try: - events = _sse_json(stream.read()) - finally: - stream.close() - matches = [event for event in events if event.get("id") == request_id] - if not matches: - raise AssertionError(f"MCP {method} returned no JSON-RPC response: {events!r}") - result = matches[-1] - if "error" in result: - raise AssertionError(f"MCP {method} failed: {result['error']!r}") - return result - - def notify(self, method: str, params: dict | None = None) -> None: - headers = { - "Accept": "application/json, text/event-stream", - "Content-Type": "application/json", - } - if self.session_id: - headers["Mcp-Session-Id"] = self.session_id - response = self.api.client.api_client.do( - "POST", - path=self.path, - headers=headers, - body={"jsonrpc": "2.0", "method": method, "params": params or {}}, - ) - # Streamable-HTTP servers commonly acknowledge notifications with HTTP 202 and no body; - # api_client.do returns None for that successful response and raises on HTTP failures. - if response is not None and not isinstance(response, (dict, list)): - raise AssertionError(f"MCP notification returned {type(response)!r}") - - -def mcp_tools(result: dict) -> list[dict]: - payload = result.get("result") - if not isinstance(payload, dict): - raise AssertionError(f"MCP response had no result object: {result!r}") - tools = payload.get("tools") - if not isinstance(tools, list) or not all(isinstance(tool, dict) for tool in tools): - raise AssertionError(f"MCP tools/list returned an invalid tool list: {result!r}") - return tools - - -def listed_mcp_names(output: str, dotted_name: str) -> bool: - dashed = dotted_name.replace(".", "-") - return any(dotted_name in line or dashed in line for line in output.splitlines()) diff --git a/tests/e2e_cuj/helpers/constants.py b/tests/e2e_cuj/helpers/constants.py index fbd144e36..4e193da88 100644 --- a/tests/e2e_cuj/helpers/constants.py +++ b/tests/e2e_cuj/helpers/constants.py @@ -1,4 +1,4 @@ -"""CLI agent names and public API wire values shared by CUJs.""" +"""Agent, metastore fixture, and transcript values shared by CUJs.""" from enum import StrEnum from pathlib import Path @@ -20,3 +20,38 @@ class CodingAgent(StrEnum): CLAUDE_CODE = "CODING_AGENT_CLAUDE_CODE" CODEX = "CODING_AGENT_CODEX" + + +CODING_AGENT_BY_CLI_NAME = { + CLAUDE: CodingAgent.CLAUDE_CODE, + CODEX: CodingAgent.CODEX, +} + +MODEL_PROVIDER_SERVICE_FIXTURES = { + CLAUDE: ("ug_e2e.providers.anthropic", "claude-haiku-4-5-20251001"), + CODEX: ("ug_e2e.providers.openai", "gpt-5-nano"), +} + +SANDBOX_MCP_SERVICE_NAME = "system.ai.sandbox" +WEB_SEARCH_MCP_SERVICE_NAME = "system.ai.web_search" +SANDBOX_MCP_TOOL_IDENTIFIERS = frozenset( + { + SANDBOX_MCP_SERVICE_NAME, + SANDBOX_MCP_SERVICE_NAME.replace(".", "-"), + SANDBOX_MCP_SERVICE_NAME.replace(".", "_"), + } +) + +TOOL_CALL_EVENT_TYPES = frozenset( + {"custom_tool_call", "function_call", "mcp_tool_call", "tool_call", "tool_use"} +) +TOOL_RESULT_EVENT_TYPES = frozenset( + { + "custom_tool_call_output", + "function_call_output", + "mcp_tool_result", + "mcp_tool_call_output", + "tool_result", + "tool_call_output", + } +) diff --git a/tests/e2e_cuj/pytest.ini b/tests/e2e_cuj/pytest.ini deleted file mode 100644 index 39ddadd4f..000000000 --- a/tests/e2e_cuj/pytest.ini +++ /dev/null @@ -1,5 +0,0 @@ -[pytest] -testpaths = . -addopts = -ra -markers = - cuj2: MPS plus explicitly selected MCP resources diff --git a/tests/e2e_cuj/test_cuj2_mps_mcp.py b/tests/e2e_cuj/test_cuj2_mps_mcp.py index 90a56446b..2ce1ad55f 100644 --- a/tests/e2e_cuj/test_cuj2_mps_mcp.py +++ b/tests/e2e_cuj/test_cuj2_mps_mcp.py @@ -3,52 +3,41 @@ from __future__ import annotations import json +import shutil import tomllib import uuid -from collections.abc import Iterator -from contextlib import contextmanager -import pytest from databricks.sdk.service.catalog import ListModelProviderServicesRequestView +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 import ( - McpRpc, - WorkspaceApi, - correlated_tool_evidence, - listed_mcp_names, - mcp_tools, - structured_final_answers, +from .helpers.constants import ( + CLAUDE, + CODEX, + CODING_AGENT_BY_CLI_NAME, + MODEL_PROVIDER_SERVICE_FIXTURES, + SANDBOX_MCP_SERVICE_NAME, + SANDBOX_MCP_TOOL_IDENTIFIERS, + TOOL_CALL_EVENT_TYPES, + TOOL_RESULT_EVENT_TYPES, + WEB_SEARCH_MCP_SERVICE_NAME, ) -from .helpers.tui_request_recorder import TuiRequestRecorder +from .helpers.workspace import Workspace WORKSPACE_URL = "https://dbc-0dcf95cf-e357.cloud.databricks.com" -AGENT_MPS: dict[str, tuple[str, str]] = { - "claude": ("ug_e2e.providers.anthropic", "claude-haiku-4-5-20251001"), - "codex": ("ug_e2e.providers.openai", "gpt-5-nano"), -} -MCP_NAME = "system.ai.sandbox" -NEGATIVE_MCP_NAME = "system.ai.web_search" -CONFIG_PATH = "/api/ai-gateway/v2/coding-agent-configs" -CONFIG_UPDATE_MASK = "default_agent,enabled_agents,mcp_servers" -CONFIG_MUTABLE_FIELDS = ("default_agent", "enabled_agents", "mcp_servers") -AGENT_NAMES = { - "claude": "CODING_AGENT_CLAUDE_CODE", - "codex": "CODING_AGENT_CODEX", -} - - -def _config_rows(payload: dict | list) -> list[dict]: - rows = payload.get("coding_agent_configs") if isinstance(payload, dict) else payload - if not isinstance(rows, list) or not all(isinstance(row, dict) for row in rows): - raise AssertionError(f"Invalid CodingAgentConfig listing: {payload!r}") - return rows def _agent_config(agent: str) -> dict: - mps_name, model = AGENT_MPS[agent] + mps_name, model = MODEL_PROVIDER_SERVICE_FIXTURES[agent] return { - "agent": AGENT_NAMES[agent], + "agent": CODING_AGENT_BY_CLI_NAME[agent], "config": { "models": {"model_provider_service": mps_name}, "default_models": {"default_model": model}, @@ -58,113 +47,21 @@ def _agent_config(agent: str) -> dict: } -def _desired_config() -> dict: - return { +def _assert_cuj2_config(config: dict) -> None: + actual = Workspace.policy(config) + actual.pop("name", None) + expected = { "spec_version": 1, - "default_agent": AGENT_NAMES["codex"], - "enabled_agents": [_agent_config("claude"), _agent_config("codex")], - "mcp_servers": {"names": [MCP_NAME]}, + "default_agent": CODING_AGENT_BY_CLI_NAME[CODEX], + "enabled_agents": [_agent_config(CLAUDE), _agent_config(CODEX)], + "mcp_servers": {"names": [SANDBOX_MCP_SERVICE_NAME]}, } + assert actual == expected, config -def _single_config(api: WorkspaceApi) -> dict | None: - rows = _config_rows(api.json("GET", CONFIG_PATH)) - if len(rows) > 1: - raise AssertionError(f"Workspace has more than one CodingAgentConfig: {rows!r}") - return rows[0] if rows else None - - -def _config_resource_path(name: str) -> str: - return f"{CONFIG_PATH}/{name.removeprefix('coding-agent-configs/').strip('/')}" - - -def _ensure_single_config(api: WorkspaceApi, desired: dict) -> dict: - original = _single_config(api) - if original is None: - api.json("POST", CONFIG_PATH, body=desired) - else: - name = original.get("name") - if not isinstance(name, str) or not name: - raise AssertionError(f"Existing CodingAgentConfig has no resource name: {original!r}") - api.json( - "PATCH", - _config_resource_path(name), - query={"update_mask": CONFIG_UPDATE_MASK}, - body={**desired, "name": name}, - ) - updated = _single_config(api) - if updated is None: - raise AssertionError("Create/update did not leave a CodingAgentConfig") - return updated - - -def _restore_config(api: WorkspaceApi, original: dict | None) -> None: - if original is None: - created = _single_config(api) - if created is None: - return - name = created.get("name") - if not isinstance(name, str) or not name: - raise AssertionError(f"Temporary CodingAgentConfig has no resource name: {created!r}") - api.request("DELETE", _config_resource_path(name)) - remaining = _single_config(api) - if remaining is not None: - raise AssertionError(f"Temporary CodingAgentConfig was not deleted: {remaining!r}") - return - - name = original.get("name") - if not isinstance(name, str) or not name: - raise AssertionError(f"Original CodingAgentConfig has no resource name: {original!r}") - body = { - key: original[key] for key in ("spec_version", *CONFIG_MUTABLE_FIELDS) if key in original - } - body["name"] = name - api.json( - "PATCH", - _config_resource_path(name), - query={"update_mask": CONFIG_UPDATE_MASK}, - body=body, - ) - restored = _single_config(api) - if restored is None: - raise AssertionError("Original CodingAgentConfig was not restored") - for field in CONFIG_MUTABLE_FIELDS: - assert restored.get(field) == original.get(field), ( - f"CodingAgentConfig field {field!r} was not restored: " - f"expected {original.get(field)!r}, got {restored.get(field)!r}" - ) - - -@contextmanager -def _published_config(api: WorkspaceApi) -> Iterator[dict]: - original = _single_config(api) - try: - yield _ensure_single_config(api, _desired_config()) - finally: - _restore_config(api, original) - - -def _assert_published_config(config: dict) -> None: - assert config.get("default_agent") == AGENT_NAMES["codex"], config - assert config.get("mcp_servers") == {"names": [MCP_NAME]}, config - - entries = config.get("enabled_agents") - assert isinstance(entries, list) and len(entries) == 2, config - by_agent = {entry.get("agent"): entry for entry in entries if isinstance(entry, dict)} - assert set(by_agent) == set(AGENT_NAMES.values()), config - for agent, agent_name in AGENT_NAMES.items(): - mps_name, model = AGENT_MPS[agent] - entry_config = by_agent[agent_name].get("config") - assert isinstance(entry_config, dict), by_agent[agent_name] - assert entry_config.get("models") == {"model_provider_service": mps_name}, entry_config - assert entry_config.get("default_models") == {"default_model": model}, entry_config - assert entry_config.get("smart_routing") == {"enabled": False}, entry_config - assert entry_config.get("tracing") == {"enabled": False}, entry_config - - -def _assert_mps_catalog(workspace_client) -> None: +def _assert_mps_catalog(workspace) -> None: listed = list( - workspace_client.ai_gateway.list_model_provider_services( + workspace.client.ai_gateway.list_model_provider_services( parent="schemas/ug_e2e.providers", view=ListModelProviderServicesRequestView.FULL, ) @@ -172,28 +69,30 @@ def _assert_mps_catalog(workspace_client) -> None: listed_by_name = { item.name.removeprefix("model-provider-services/"): item for item in listed if item.name } - for agent, (mps_name, model) in AGENT_MPS.items(): - service = workspace_client.ai_gateway.get_model_provider_service( + for agent, (mps_name, model) in MODEL_PROVIDER_SERVICE_FIXTURES.items(): + service = workspace.client.ai_gateway.get_model_provider_service( f"model-provider-services/{mps_name}" ) assert service.name and service.name.endswith(mps_name), service config = service.config assert config is not None, service targets = config.targets or [] - models = [target.model for target in targets] - assert models == [model], service + assert [target.model for target in targets] == [model], service assert config.allow_all_targets is not True, service expected_type = ( "EXTERNAL_MODEL_PROVIDER_TYPE_ANTHROPIC" - if agent == "claude" + if agent == CLAUDE else "EXTERNAL_MODEL_PROVIDER_TYPE_OPENAI" ) provider_type = getattr(config.provider_type, "value", config.provider_type) assert provider_type == expected_type, config - expected_api = "anthropic/v1/messages" if agent == "claude" else "openai/v1/responses" - assert [api for target in targets for api in (target.native_api_types or [])] == [ - expected_api - ], targets + expected_api = "anthropic/v1/messages" if agent == CLAUDE else "openai/v1/responses" + native_apis = [ + getattr(api, "value", api) + for target in targets + for api in (target.native_api_types or []) + ] + assert native_apis == [expected_api], targets matching = listed_by_name.get(mps_name) assert matching is not None, [item.name for item in listed] @@ -201,43 +100,26 @@ def _assert_mps_catalog(workspace_client) -> None: assert [target.model for target in listed_targets or []] == [model], matching -def _assert_mcp_prerequisite(workspace_client, api: WorkspaceApi) -> None: - sandbox = workspace_client.ai_gateway.get_mcp_service(f"mcp-services/{MCP_NAME}") - web_search = workspace_client.ai_gateway.get_mcp_service(f"mcp-services/{NEGATIVE_MCP_NAME}") - assert sandbox.name and sandbox.name.endswith(MCP_NAME), sandbox - assert web_search.name and web_search.name.endswith(NEGATIVE_MCP_NAME), web_search - - endpoint = f"/ai-gateway/mcp-services/{MCP_NAME}" - rpc = McpRpc(api, endpoint) - initialize = rpc.request( - "initialize", - { - "protocolVersion": "2025-03-26", - "capabilities": {}, - "clientInfo": {"name": "ug-cuj2", "version": "1"}, - }, - ) - assert isinstance(initialize.get("result"), dict), initialize - rpc.notify("notifications/initialized") - tools = mcp_tools(rpc.request("tools/list")) - run_code = [tool for tool in tools if tool.get("name") == "run_code"] - assert len(run_code) == 1, tools +def _assert_mcp_services(workspace) -> None: + for name in (SANDBOX_MCP_SERVICE_NAME, WEB_SEARCH_MCP_SERVICE_NAME): + service = workspace.client.ai_gateway.get_mcp_service(f"mcp-services/{name}") + assert service.name and service.name.endswith(name), service -def _assert_configured_files(session, workspace: str) -> None: +def _assert_configured_files(session, workspace_url: str) -> None: managed_cache = json.loads( (session.home / ".ucode" / "managed-config.json").read_text(encoding="utf-8") ) managed_config = managed_cache.get("config") assert isinstance(managed_config, dict), managed_cache - _assert_published_config(managed_config) + _assert_cuj2_config(managed_config) claude_settings = session.home / ".claude" / "ucode-settings.json" assert claude_settings.is_file(), claude_settings settings = json.loads(claude_settings.read_text(encoding="utf-8")) claude_env = settings.get("env") assert isinstance(claude_env, dict), settings - claude_mps, claude_model = AGENT_MPS["claude"] + claude_mps, claude_model = MODEL_PROVIDER_SERVICE_FIXTURES[CLAUDE] custom_headers = claude_env.get("ANTHROPIC_CUSTOM_HEADERS") assert isinstance(custom_headers, str), claude_env claude_headers = {} @@ -247,17 +129,126 @@ def _assert_configured_files(session, workspace: str) -> None: claude_headers[name.strip().casefold()] = value.strip() assert claude_headers["databricks-model-provider-service"] == claude_mps, claude_headers assert claude_env["ANTHROPIC_DEFAULT_HAIKU_MODEL"] == claude_model, claude_env - assert claude_env["ANTHROPIC_BASE_URL"] == workspace.rstrip("/") + "/ai-gateway/anthropic" + assert claude_env["ANTHROPIC_BASE_URL"] == workspace_url.rstrip("/") + "/ai-gateway/anthropic" - _, codex_model = AGENT_MPS["codex"] + _, codex_model = MODEL_PROVIDER_SERVICE_FIXTURES[CODEX] codex_path = session.home / ".codex" / "ucode.config.toml" assert codex_path.is_file(), codex_path codex_config = tomllib.loads(codex_path.read_text(encoding="utf-8")) assert codex_config["model_provider"] == "Databricks", codex_config - assert codex_config["model"] == codex_model, codex_config codex_provider = codex_config["model_providers"]["Databricks"] assert codex_provider["wire_api"] == "responses", codex_provider - assert codex_provider["base_url"] == workspace.rstrip("/") + "/ai-gateway/codex/v1" + assert codex_provider["base_url"] == workspace_url.rstrip("/") + "/ai-gateway/codex/v1" + assert session.workspace_state()["codex_default_model"] == codex_model, ( + session.workspace_state() + ) + + +def _mcp_name_listed(output: str, name: str) -> bool: + dashed = name.replace(".", "-") + return any(name in line or dashed in line for line in output.splitlines()) + + +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}" + assert _mcp_name_listed(output, SANDBOX_MCP_SERVICE_NAME), output + assert not _mcp_name_listed(output, WEB_SEARCH_MCP_SERVICE_NAME), output + + +def _dict_nodes(value: object): + if isinstance(value, dict): + yield value + for child in value.values(): + yield from _dict_nodes(child) + elif isinstance(value, list): + for child in value: + yield from _dict_nodes(child) + + +def _call_id(mapping: dict) -> str | None: + for key in ("id", "call_id", "callId", "tool_use_id", "toolUseId"): + value = mapping.get(key) + if isinstance(value, str) and value: + return value + return None + + +def _parent_transcripts(session, agent: str) -> dict[str, list[dict]]: + return { + path: records + for path, records in agent_sessions(session, agent).items() + if not is_child_session(agent, path, records) + } + + +def _correlated_run_code(transcripts: dict[str, list[dict]], marker: str) -> dict: + calls: dict[str, dict] = {} + results: dict[str, dict] = {} + for records in transcripts.values(): + for record in records: + for node in _dict_nodes(record): + event_type = str(node.get("type", "")).lower() + identifier = _call_id(node) + if not identifier: + continue + if event_type in TOOL_CALL_EVENT_TYPES: + calls[identifier] = node + elif event_type in TOOL_RESULT_EVENT_TYPES: + results[identifier] = node + + for identifier, call in calls.items(): + call_text = json.dumps(call, sort_keys=True, default=str).lower() + result = results.get(identifier) + result_text = json.dumps(result, sort_keys=True, default=str) if result else "" + failed = any( + node.get("is_error") is True + or str(node.get("status", "")).lower() in {"error", "failed", "failure"} + for node in _dict_nodes(result) + ) + if ( + "run_code" in call_text + and any(name in call_text for name in SANDBOX_MCP_TOOL_IDENTIFIERS) + and marker.lower() in call_text + and marker in result_text + and not failed + ): + return {"call_id": identifier, "call": call, "result": result} + raise AssertionError( + f"No successful sandbox run_code call/result carried marker {marker!r}; " + f"calls={list(calls)} results={list(results)}" + ) + + +def _task_complete(session, agent: str, marker: str) -> bool: + transcripts = _parent_transcripts(session, agent) + if not any( + marker in answer + for records in transcripts.values() + for answer in assistant_answers(agent, records) + ): + return False + try: + _correlated_run_code(transcripts, marker) + except AssertionError: + return False + return True + + +def _assert_task_evidence(session, agent: str, marker: str) -> None: + transcripts = _parent_transcripts(session, agent) + tool_evidence = _correlated_run_code(transcripts, marker) + answers = [ + answer for records in transcripts.values() for answer in assistant_answers(agent, records) + ] + assert any(marker in answer for answer in answers), answers + session.record( + f"cuj2-{agent}-evidence", + {"marker": marker, "tool": tool_evidence, "answers": answers}, + ) def _agent_prompt(marker: str) -> str: @@ -269,160 +260,147 @@ def _agent_prompt(marker: str) -> str: ) -def _run_agent_task(session, agent: str, marker: str): - prompt = _agent_prompt(marker) +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 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" - def completed(current) -> bool: - if not current.transcript_contains(agent, marker): - return False - try: - correlated_tool_evidence(current.transcripts(agent), "run_code", marker) - except AssertionError: - return False - return True - return session.run_tui( - agent, - prompt, - completion=completed, - timeout=240, - ) - - -@pytest.mark.cuj2 class TestCuj2MpsExplicitMcp(BaseCujTest): WORKSPACE_URL = WORKSPACE_URL - def test_cuj_configuration(self, live_session): - """Scenario: publish both agent-specific MPS targets and the sandbox MCP, then configure ug. + def test_cuj_configuration(self, cuj): + """Scenario: configure ug from the preconfigured two-agent MPS/MCP workspace. - Expected: the published config selects gpt-5-nano for Codex and - claude-haiku-4-5-20251001 for Claude, with routing and tracing off; - both native MPS APIs advertise their selected model, ug writes the two - agent configs, and both agents list sandbox but exclude web_search. + Expected: the read-only CodingAgentConfig selects the two exact MPS resources and sandbox + MCP; both native MPS APIs advertise their selected model, generated agent settings use + the exact provider/model values, and agent MCP listings exclude web_search. """ - session = live_session - api = WorkspaceApi(self.workspace) - _assert_mps_catalog(self.workspace) - _assert_mcp_prerequisite(self.workspace, api) - - with _published_config(api) as published: - _assert_published_config(published) - configured = session.run( + session, workspace, _ = cuj + published = workspace.config() + _assert_cuj2_config(published) + _assert_mps_catalog(workspace) + _assert_mcp_services(workspace) + + configured = session.run( + "configure", + "--workspace", + workspace.url, + "--skip-upgrade", + "--disable-databricks-ai-tools", + timeout=300, + ) + assert "Select coding agents to configure:" not in configured.stdout, configured.stdout + _assert_configured_files(session, workspace.url) + _assert_generated_mcp_listings(session) + + def test_cuj_codex_inference(self, cuj): + """Scenario: configure through the recorder and complete a real Codex MCP task. + + Expected: Codex sends a Responses request for gpt-5-nano with the + ug_e2e.providers.openai target header and receives its paired HTTP 200 response; its + parent transcript contains a correlated sandbox run_code call/result and final marker. + """ + session, workspace, recorder = cuj + published = workspace.config() + _assert_cuj2_config(published) + _assert_mps_catalog(workspace) + _assert_mcp_services(workspace) + + session.configure( + [ "configure", "--workspace", - WORKSPACE_URL, + recorder.url, "--skip-upgrade", - timeout=300, - ) - assert "Select coding agents to configure:" not in configured.stdout, configured.stdout - _assert_configured_files(session, WORKSPACE_URL) - - claude_listing = session.external("claude", "mcp", "list", timeout=120) - claude_output = f"{claude_listing.stdout}\n{claude_listing.stderr}" - assert listed_mcp_names(claude_output, MCP_NAME), claude_output - assert not listed_mcp_names(claude_output, NEGATIVE_MCP_NAME), claude_output - - codex_listing = session.external("codex", "mcp", "list", timeout=120) - codex_output = f"{codex_listing.stdout}\n{codex_listing.stderr}" - assert listed_mcp_names(codex_output, MCP_NAME), codex_output - assert not listed_mcp_names(codex_output, NEGATIVE_MCP_NAME), codex_output + "--disable-databricks-ai-tools", + ] + ) + _assert_configured_files(session, recorder.url) + _assert_generated_mcp_listings(session) + + marker = f"CUJ2-CODEX-{uuid.uuid4().hex}" + checkpoint = recorder.checkpoint() + with AgentTerminal(session, CODEX, [str(session.binary), CODEX], "cuj2-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 MCP task", timeout=240) + tui.exit_normally() + + request = recorder.expect_request( + method="POST", + path="/ai-gateway/codex/v1/responses", + after=checkpoint, + timeout=240, + ) + _assert_inference_request( + recorder, + request, + provider=MODEL_PROVIDER_SERVICE_FIXTURES[CODEX][0], + model=MODEL_PROVIDER_SERVICE_FIXTURES[CODEX][1], + marker=marker, + ) + _assert_task_evidence(session, CODEX, marker) - def test_cuj_codex_inference(self, live_session): - """Scenario: configure ug from the dedicated workspace and complete a real Codex task. + def test_cuj_claude_inference(self, cuj): + """Scenario: configure through the recorder and complete a real Claude MCP task. - Expected: Codex sends a Responses request for gpt-5-nano with the - ug_e2e.providers.openai target header and receives its paired HTTP 200 - response; its transcript contains a correlated sandbox run_code result - and its final answer contains the task marker. + Expected: Claude sends a Messages request for claude-haiku-4-5-20251001 with the + ug_e2e.providers.anthropic target header and receives its paired HTTP 200 response; its + parent transcript contains a correlated sandbox run_code call/result and final marker. """ - session = live_session - api = WorkspaceApi(self.workspace) - with _published_config(api) as published: - _assert_published_config(published) - with TuiRequestRecorder(WORKSPACE_URL) as recorder: - configured = session.run( - "configure", - "--workspace", - recorder.url, - "--skip-upgrade", - timeout=300, - ) - assert "Select coding agents to configure:" not in configured.stdout, ( - configured.stdout - ) - _assert_configured_files(session, recorder.url) - - marker = f"CUJ2-CODEX-{uuid.uuid4().hex}" - checkpoint = recorder.checkpoint() - result = _run_agent_task(session, "codex", marker) - request = recorder.expect_request( - method="POST", - path="/ai-gateway/codex/v1/responses", - after=checkpoint, - timeout=240, - ) - assert request.headers["databricks-model-provider-service"] == ( - "ug_e2e.providers.openai" - ) - assert request.payload["model"] == "gpt-5-nano" - 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, "Codex inference response was empty" - - transcripts = session.transcripts("codex") - assert transcripts, "No Codex transcript was written" - assert correlated_tool_evidence(transcripts, "run_code", marker) - answers = structured_final_answers("codex", result, transcripts) - assert any(marker in answer for answer in answers), answers - - def test_cuj_claude_inference(self, live_session): - """Scenario: configure ug from the dedicated workspace and complete a real Claude task. - - Expected: Claude sends a Messages request for - claude-haiku-4-5-20251001 with the ug_e2e.providers.anthropic target - header and receives its paired HTTP 200 response; its transcript - contains a correlated sandbox run_code result and its final answer - contains the task marker. - """ - session = live_session - api = WorkspaceApi(self.workspace) - with _published_config(api) as published: - _assert_published_config(published) - with TuiRequestRecorder(WORKSPACE_URL) as recorder: - configured = session.run( - "configure", - "--workspace", - recorder.url, - "--skip-upgrade", - timeout=300, - ) - assert "Select coding agents to configure:" not in configured.stdout, ( - configured.stdout - ) - _assert_configured_files(session, recorder.url) - - marker = f"CUJ2-CLAUDE-{uuid.uuid4().hex}" - checkpoint = recorder.checkpoint() - result = _run_agent_task(session, "claude", marker) - request = recorder.expect_request( - method="POST", - path="/ai-gateway/anthropic/v1/messages", - after=checkpoint, - timeout=240, - ) - assert request.headers["databricks-model-provider-service"] == ( - "ug_e2e.providers.anthropic" - ) - assert request.payload["model"] == "claude-haiku-4-5-20251001" - 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, "Claude inference response was empty" - - transcripts = session.transcripts("claude") - assert transcripts, "No Claude transcript was written" - assert correlated_tool_evidence(transcripts, "run_code", marker) - answers = structured_final_answers("claude", result, transcripts) - assert any(marker in answer for answer in answers), answers + session, workspace, recorder = cuj + published = workspace.config() + _assert_cuj2_config(published) + _assert_mps_catalog(workspace) + _assert_mcp_services(workspace) + + session.configure( + [ + "configure", + "--workspace", + recorder.url, + "--skip-upgrade", + "--disable-databricks-ai-tools", + ] + ) + _assert_configured_files(session, recorder.url) + _assert_generated_mcp_listings(session) + + marker = f"CUJ2-CLAUDE-{uuid.uuid4().hex}" + checkpoint = recorder.checkpoint() + with AgentTerminal(session, CLAUDE, [str(session.binary), CLAUDE], "cuj2-claude") 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, CLAUDE, marker) + + tui.wait_for(completed, "completed Claude MCP task", timeout=240) + tui.exit_normally() + + request = recorder.expect_request( + method="POST", + path="/ai-gateway/anthropic/v1/messages", + after=checkpoint, + timeout=240, + ) + _assert_inference_request( + recorder, + request, + provider=MODEL_PROVIDER_SERVICE_FIXTURES[CLAUDE][0], + model=MODEL_PROVIDER_SERVICE_FIXTURES[CLAUDE][1], + marker=marker, + ) + _assert_task_evidence(session, CLAUDE, marker) diff --git a/tests/integration/README.md b/tests/integration/README.md index 5f631a5eb..67448660e 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -3,11 +3,11 @@ The separate [dedicated-workspace CUJ](../e2e_cuj/AGENTS.md) lives in `tests/e2e_cuj/` and runs directly with pytest. It does not use this suite's runner or config fixtures. It reuses the session, terminal, file-task, and transcript helpers in `utils/`. -Its Claude/Codex evidence helpers keep scenario-specific assertions separate from shared mechanics. -Workspace config/catalog reads use its base class's Databricks SDK client. Configuration -is read-only and checked for changes at teardown; concurrent readers need no reservation. -The smart-routing CUJ's four sessions cover routing and explicit-model bypass, not -routing-disabled behavior. +CUJ-specific evidence stays separate from shared mechanics. Workspace config/catalog reads use +its base class's Databricks SDK client. Configuration is read-only and checked for changes at +teardown; concurrent readers need no reservation. One journey covers four routing and +explicit-model sessions. Another has three separately collected cases for exact MPS/MCP +configuration, Codex inference, and Claude inference. Routing-disabled behavior remains uncovered. This suite runs the **installed product** through subprocesses, against the same `UCODE_TEST_WORKSPACE` used by the existing e2e tests. It does not import `ucode`, From 1ef1b1669b7a762876aeb5f7c6156ceb150ec4fc Mon Sep 17 00:00:00 2001 From: Mason Cao Date: Tue, 6 Oct 2026 01:51:47 +0000 Subject: [PATCH 04/19] [AIGTWY-4877] Preserve CUJ recorder origin for MPS lookup --- src/ucode/databricks.py | 10 ++++----- tests/e2e_cuj/test_cuj2_mps_mcp.py | 4 ---- tests/test_databricks.py | 34 ++++++++++++++++++++++++++++++ 3 files changed, 38 insertions(+), 10 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 970f51ecc..9c720bca9 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -2384,7 +2384,7 @@ def list_model_provider_services( # reach the next. return [dict(service) for service in cached], None - hostname = workspace_hostname(workspace) + origin = workspace_origin(workspace) services: list[dict] = [] page_token: str | None = None seen_tokens: set[str] = set() @@ -2395,9 +2395,7 @@ def list_model_provider_services( params["parent"] = f"schemas/{parent}" if page_token: params["page_token"] = page_token - url = ( - f"https://{hostname}/api/2.1/unity-catalog/model-provider-services?{urlencode(params)}" - ) + url = f"{origin}/api/2.1/unity-catalog/model-provider-services?{urlencode(params)}" payload, reason = _http_get_json(url, token, timeout=30) if payload is None: # Surface the failure only if we have nothing yet; a mid-pagination blip still @@ -2470,8 +2468,8 @@ def get_model_provider_service( server-side filter) makes a service that plainly exists look absent. Addressing it directly removes that whole class of false negative. """ - hostname = workspace_hostname(workspace) - url = f"https://{hostname}/api/2.1/unity-catalog/model-provider-services/{service_name}" + origin = workspace_origin(workspace) + url = f"{origin}/api/2.1/unity-catalog/model-provider-services/{service_name}" payload, reason = _http_get_json(url, token, timeout=30) if payload is None: return None, reason diff --git a/tests/e2e_cuj/test_cuj2_mps_mcp.py b/tests/e2e_cuj/test_cuj2_mps_mcp.py index 2ce1ad55f..be3a8c2ae 100644 --- a/tests/e2e_cuj/test_cuj2_mps_mcp.py +++ b/tests/e2e_cuj/test_cuj2_mps_mcp.py @@ -131,7 +131,6 @@ def _assert_configured_files(session, workspace_url: str) -> None: assert claude_env["ANTHROPIC_DEFAULT_HAIKU_MODEL"] == claude_model, claude_env assert claude_env["ANTHROPIC_BASE_URL"] == workspace_url.rstrip("/") + "/ai-gateway/anthropic" - _, codex_model = MODEL_PROVIDER_SERVICE_FIXTURES[CODEX] codex_path = session.home / ".codex" / "ucode.config.toml" assert codex_path.is_file(), codex_path codex_config = tomllib.loads(codex_path.read_text(encoding="utf-8")) @@ -139,9 +138,6 @@ def _assert_configured_files(session, workspace_url: str) -> None: codex_provider = codex_config["model_providers"]["Databricks"] assert codex_provider["wire_api"] == "responses", codex_provider assert codex_provider["base_url"] == workspace_url.rstrip("/") + "/ai-gateway/codex/v1" - assert session.workspace_state()["codex_default_model"] == codex_model, ( - session.workspace_state() - ) def _mcp_name_listed(output: str, name: str) -> bool: diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 40a47fcc0..eda29e87b 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -1035,6 +1035,20 @@ def fake_get(url, token, **kwargs): assert "page_size=" in seen["url"] + def test_preserves_loopback_workspace_origin(self, monkeypatch): + seen: dict = {} + + def fake_get(url, token, **kwargs): + seen["url"] = url + return self._page(["main.s.one"]), None + + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) + db_mod.list_model_provider_services("http://127.0.0.1:54321", "tok", use_cache=False) + + assert seen["url"].startswith( + "http://127.0.0.1:54321/api/2.1/unity-catalog/model-provider-services?" + ) + class TestGetModelProviderService: def test_addresses_the_service_directly(self, monkeypatch): @@ -1061,6 +1075,26 @@ def fake_get(url, token, **kwargs): assert service["allow_all_targets"] is True assert seen["url"].endswith("/model-provider-services/main.tien_le.openai_all") + def test_preserves_loopback_workspace_origin(self, monkeypatch): + seen: dict = {} + + def fake_get(url, token, **kwargs): + seen["url"] = url + return { + "name": "model-provider-services/main.tien_le.openai_all", + "config": {"provider_type": "EXTERNAL_MODEL_PROVIDER_TYPE_OPENAI"}, + }, None + + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) + db_mod.get_model_provider_service( + "main.tien_le.openai_all", "http://127.0.0.1:54321", "tok" + ) + + assert seen["url"] == ( + "http://127.0.0.1:54321/api/2.1/unity-catalog/" + "model-provider-services/main.tien_le.openai_all" + ) + def test_missing_service_returns_the_reason(self, monkeypatch): monkeypatch.setattr( db_mod, "_http_get_json", lambda url, token, **kw: (None, "HTTP 404 Not Found") From 518fabbebbb7c0a4b31c1c84fc6fa0602dac0870 Mon Sep 17 00:00:00 2001 From: Mason Cao Date: Tue, 6 Oct 2026 02:15:18 +0000 Subject: [PATCH 05/19] [AIGTWY-4877] Isolate CUJ2 cases --- tests/e2e_cuj/test_cuj2_mps_mcp.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/e2e_cuj/test_cuj2_mps_mcp.py b/tests/e2e_cuj/test_cuj2_mps_mcp.py index be3a8c2ae..d39949365 100644 --- a/tests/e2e_cuj/test_cuj2_mps_mcp.py +++ b/tests/e2e_cuj/test_cuj2_mps_mcp.py @@ -265,9 +265,11 @@ def _assert_inference_request(recorder, request, *, provider: str, model: str, m assert response.body, "Inference response was empty" -class TestCuj2MpsExplicitMcp(BaseCujTest): +class _Cuj2Base(BaseCujTest): WORKSPACE_URL = WORKSPACE_URL + +class TestCuj2Configuration(_Cuj2Base): def test_cuj_configuration(self, cuj): """Scenario: configure ug from the preconfigured two-agent MPS/MCP workspace. @@ -293,6 +295,8 @@ def test_cuj_configuration(self, cuj): _assert_configured_files(session, workspace.url) _assert_generated_mcp_listings(session) + +class TestCuj2CodexInference(_Cuj2Base): def test_cuj_codex_inference(self, cuj): """Scenario: configure through the recorder and complete a real Codex MCP task. @@ -347,6 +351,8 @@ def completed(screen): ) _assert_task_evidence(session, CODEX, marker) + +class TestCuj2ClaudeInference(_Cuj2Base): def test_cuj_claude_inference(self, cuj): """Scenario: configure through the recorder and complete a real Claude MCP task. From b4e5a5d65675c78fc37680c5db8d70661dc4b35f Mon Sep 17 00:00:00 2001 From: Mason Cao Date: Tue, 6 Oct 2026 02:30:21 +0000 Subject: [PATCH 06/19] [AIGTWY-4877] Fix MCP auth and Claude permission handling --- src/ucode/agents/codex.py | 11 ++++++++++- tests/e2e_cuj/test_cuj2_mps_mcp.py | 19 ++++++++++++++++++- tests/test_agent_codex.py | 10 ++++++++++ 3 files changed, 38 insertions(+), 2 deletions(-) diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 8b0bcd3b0..e63463113 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -613,6 +613,11 @@ def _reconcile_managed_config(state: dict, compose: Callable[[dict], dict]) -> N MANAGED_MCP_CONFIG_KEY = "mcp_servers" +_CODEX_MCP_AUTH_ENV_VARS = [ + "DATABRICKS_BEARER", + "DATABRICKS_BEARER_COMMAND", + "DATABRICKS_CONFIG_FILE", +] def managed_mcp_uses_managed_file() -> bool: @@ -627,7 +632,11 @@ def managed_mcp_uses_managed_file() -> bool: def managed_mcp_entry(argv: list[str]) -> dict: """A ``[mcp_servers.]`` stdio entry from the ``ug mcp-proxy`` argv (same as user scope).""" - return {"command": argv[0], "args": list(argv[1:])} + return { + "command": argv[0], + "args": list(argv[1:]), + "env_vars": list(_CODEX_MCP_AUTH_ENV_VARS), + } def managed_mcp_http_entry(url: str, client_id: str) -> dict: diff --git a/tests/e2e_cuj/test_cuj2_mps_mcp.py b/tests/e2e_cuj/test_cuj2_mps_mcp.py index d39949365..43c0c11a3 100644 --- a/tests/e2e_cuj/test_cuj2_mps_mcp.py +++ b/tests/e2e_cuj/test_cuj2_mps_mcp.py @@ -383,10 +383,27 @@ def test_cuj_claude_inference(self, cuj): with AgentTerminal(session, CLAUDE, [str(session.binary), CLAUDE], "cuj2-claude") as tui: tui.boot() tui.submit(_agent_prompt(marker)) + permission_approved = False def completed(screen): + nonlocal permission_approved assert_no_terminal_api_error(screen) - assert "Do you want to proceed?" not in screen, screen + if "Do you want to proceed?" in screen: + expected_prompt = ( + "system-ai-sandbox — Run Code Tool: (MCP)", + f'code: "print(\\"{marker}\\")"', + 'language: "python"', + "Do you want to proceed?", + "❯ 1. Yes", + ) + assert all(part in screen for part in expected_prompt), ( + "Claude requested an unrecognized permission:\n" + screen + ) + if permission_approved: + return False + tui.send("\r", f"approve sandbox run_code for marker {marker}") + permission_approved = True + return False return _task_complete(session, CLAUDE, marker) tui.wait_for(completed, "completed Claude MCP task", timeout=240) diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index bde491007..c7f1e866a 100644 --- a/tests/test_agent_codex.py +++ b/tests/test_agent_codex.py @@ -1926,6 +1926,11 @@ def test_writes_mcp_servers_preserving_model_keys(self, monkeypatch): assert doc["model_provider"] == "Databricks" assert doc["mcp_servers"]["system-ai-github"]["command"] == "/opt/ug" assert doc["mcp_servers"]["system-ai-github"]["args"][:2] == ["mcp-proxy", "--url"] + assert doc["mcp_servers"]["system-ai-github"]["env_vars"] == [ + "DATABRICKS_BEARER", + "DATABRICKS_BEARER_COMMAND", + "DATABRICKS_CONFIG_FILE", + ] assert captured["owned_paths"] == [["mcp_servers"]] assert captured["tool"] == "codex" @@ -2055,6 +2060,11 @@ def test_adds_and_preserves_other_tables(self, tmp_path, monkeypatch): assert dict(doc["mcp_servers"]["system-ai-github"]) == { "command": "ug", "args": ["mcp-proxy", "https://ws/svc"], + "env_vars": [ + "DATABRICKS_BEARER", + "DATABRICKS_BEARER_COMMAND", + "DATABRICKS_CONFIG_FILE", + ], } def test_removes_named_entries_only(self, tmp_path, monkeypatch): From 01ee8224a9eda068e3319341079b9a465eff6c1f Mon Sep 17 00:00:00 2001 From: Mason Cao Date: Tue, 6 Oct 2026 02:55:32 +0000 Subject: [PATCH 07/19] [AIGTWY-4877] Capture Codex MCP startup failures --- tests/e2e_cuj/test_cuj2_mps_mcp.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/tests/e2e_cuj/test_cuj2_mps_mcp.py b/tests/e2e_cuj/test_cuj2_mps_mcp.py index 43c0c11a3..b54f71bd0 100644 --- a/tests/e2e_cuj/test_cuj2_mps_mcp.py +++ b/tests/e2e_cuj/test_cuj2_mps_mcp.py @@ -256,6 +256,20 @@ def _agent_prompt(marker: str) -> str: ) +def _assert_codex_mcp_started(tui: AgentTerminal) -> None: + if "startup issue" not in tui.visible: + return + previous = tui.visible + tui.send("\x14", "open Codex startup issue details") + tui.wait_for( + lambda screen: screen != previous, + "Codex startup issue details", + timeout=10, + stable_for=1, + ) + raise AssertionError("Codex reported a startup issue:\n" + tui.visible) + + 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 @@ -326,6 +340,7 @@ def test_cuj_codex_inference(self, cuj): checkpoint = recorder.checkpoint() with AgentTerminal(session, CODEX, [str(session.binary), CODEX], "cuj2-codex") as tui: tui.boot() + _assert_codex_mcp_started(tui) tui.submit(_agent_prompt(marker)) def completed(screen): From 916e27a9058337e5be4ba83244d2f4d6b37ce1b8 Mon Sep 17 00:00:00 2001 From: Mason Cao Date: Tue, 6 Oct 2026 03:00:31 +0000 Subject: [PATCH 08/19] [AIGTWY-4877] Handle CUJ agent startup prompts --- tests/e2e_cuj/test_cuj2_mps_mcp.py | 16 ---------------- 1 file changed, 16 deletions(-) diff --git a/tests/e2e_cuj/test_cuj2_mps_mcp.py b/tests/e2e_cuj/test_cuj2_mps_mcp.py index b54f71bd0..ed2075e86 100644 --- a/tests/e2e_cuj/test_cuj2_mps_mcp.py +++ b/tests/e2e_cuj/test_cuj2_mps_mcp.py @@ -256,20 +256,6 @@ def _agent_prompt(marker: str) -> str: ) -def _assert_codex_mcp_started(tui: AgentTerminal) -> None: - if "startup issue" not in tui.visible: - return - previous = tui.visible - tui.send("\x14", "open Codex startup issue details") - tui.wait_for( - lambda screen: screen != previous, - "Codex startup issue details", - timeout=10, - stable_for=1, - ) - raise AssertionError("Codex reported a startup issue:\n" + tui.visible) - - 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 @@ -340,7 +326,6 @@ def test_cuj_codex_inference(self, cuj): checkpoint = recorder.checkpoint() with AgentTerminal(session, CODEX, [str(session.binary), CODEX], "cuj2-codex") as tui: tui.boot() - _assert_codex_mcp_started(tui) tui.submit(_agent_prompt(marker)) def completed(screen): @@ -407,7 +392,6 @@ def completed(screen): expected_prompt = ( "system-ai-sandbox — Run Code Tool: (MCP)", f'code: "print(\\"{marker}\\")"', - 'language: "python"', "Do you want to proceed?", "❯ 1. Yes", ) From 3dff52ed891bac2c3f2f57721d364aef9e47a6d3 Mon Sep 17 00:00:00 2001 From: Mason Cao Date: Tue, 6 Oct 2026 03:29:52 +0000 Subject: [PATCH 09/19] [AIGTWY-4877] Defer unsupported sandbox execution --- tests/README.md | 8 ++- tests/e2e_cuj/AGENTS.md | 5 +- tests/e2e_cuj/helpers/constants.py | 21 ------ tests/e2e_cuj/test_cuj2_mps_mcp.py | 108 +++-------------------------- 4 files changed, 18 insertions(+), 124 deletions(-) diff --git a/tests/README.md b/tests/README.md index ec3692f59..87a62f453 100644 --- a/tests/README.md +++ b/tests/README.md @@ -138,7 +138,7 @@ integration utilities; only CUJ-specific evidence correlation stays in a test fi | Test | User action | Expected evidence | | --- | --- | --- | -| `test_cuj_configuration`, `test_cuj_codex_inference`, `test_cuj_claude_inference` in `test_cuj2_mps_mcp.py` | Read the permanently preconfigured two-agent MPS/MCP config; run separate configure, Codex TUI, and Claude TUI cases | Configuration verifies both native MPS APIs, generated settings, and sandbox inclusion with web_search excluded; each TUI case records its exact MPS target header and model, paired HTTP 200 response, correlated `run_code` call/result, and final marker. No workspace config CRUD is performed | +| `test_cuj_configuration`, `test_cuj_codex_inference`, `test_cuj_claude_inference` in `test_cuj2_mps_mcp.py` | Read the permanently preconfigured two-agent MPS/MCP config; run separate configure, Codex TUI, and Claude TUI cases | Configuration verifies both native MPS APIs, generated settings, and sandbox inclusion with web_search excluded; each TUI case records its exact MPS target header and model, paired HTTP 200 response, and final marker. Live sandbox execution is deferred because MAS cannot downscope the CI service principal. No workspace config CRUD is performed | | `test_ug_configure_claude_databricks` | Configure Databricks Hosted; execute the generated auth helper; launch plain `ug claude`, read a file, and open `/model` | Generated helper invokes `ug` with clean token stdout; assistant returns an unpredictable file value; native discovery caches `system.ai` models and the picker shows a discovered model without an opt-in flag; normal exit; reopen with working keyboard input | | `test_ug_configure_claude_anthropic_mps` | Select Anthropic MPS in the real configure picker; launch Claude | Saved provider in status; completed TUI file task; normal exit | | `test_ug_configure_codex_databricks` | Configure Databricks Hosted; execute the generated auth helper; open Codex TUI and read a file | Generated helper invokes `ug` with clean token stdout; completed assistant answer contains the file value; normal exit and reopen | @@ -218,7 +218,9 @@ one configuration case and separate Codex and Claude TUI task cases. The workspa preconfigured for CUJ2 and is read-only; every case validates it, and the fixture compares it again during teardown. The TUI cases forward agent requests through the shared per-test loopback recorder to the real workspace and assert the provider-service header, model, -paired HTTP 200 response, correlated MCP tool result, and final answer. Failed TUI runs +paired HTTP 200 response, and final answer. The configuration case verifies sandbox MCP +selection and generated client listings; live sandbox execution remains deferred because +MAS cannot downscope the CI service principal. Failed TUI runs retain terminal output/actions/screen and best-effort agent transcript snapshots under the runner artifact directory before the temporary homes are deleted. The suite requires `UG_CUJ_SP_CLIENT_ID` and `UG_CUJ_SP_CLIENT_SECRET`, with `ug`, `claude`, `codex`, and @@ -307,7 +309,7 @@ These are unit/component checks; they do not establish live sudo password-prompt | Scenario | Status / requirement | | --- | --- | -| Broad live MCP functionality | CUJ2 covers the sandbox MCP and two-agent tool use. Installation tests cover the local web-search MCP handshake and tool listing, not upstream proxying or a real search request. Other MCP services, live parent/child search, and permission decisions remain deferred | +| Broad live MCP functionality | CUJ2 covers sandbox MCP configuration and generated client listings. Live sandbox execution is deferred because MAS cannot downscope the CI service principal. Installation tests cover the local web-search MCP handshake and tool listing, not upstream proxying or a real search request. Other MCP services, live parent/child search, and permission decisions remain deferred | | Dedicated CUJ Skill coverage | Deferred because of a backend storage-path issue; CUJ2 has no Skill configuration, download, invocation, or assertion path | | Broad configure flags, multiple workspaces, and PAT flows | Deferred while focusing on basic CUJs | | Workspace-switch MCP cleanup | The `workspace_switch` CUJ covers real registration, cleanup, repeat configure, and a completed Claude task. Unit/component tests cover duplicate attempts and injected removal failures; the CUJ does not force an agent timeout. It runs in the required managed CI lane for full/live runs. | diff --git a/tests/e2e_cuj/AGENTS.md b/tests/e2e_cuj/AGENTS.md index 9471024dc..bd136dba6 100644 --- a/tests/e2e_cuj/AGENTS.md +++ b/tests/e2e_cuj/AGENTS.md @@ -56,8 +56,9 @@ preconfigured workspace: configuration, Codex inference, and Claude inference. E 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; only -CUJ2's `run_code` call/result correlation remains local to the test. +`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. 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/helpers/constants.py b/tests/e2e_cuj/helpers/constants.py index 4e193da88..3a934d550 100644 --- a/tests/e2e_cuj/helpers/constants.py +++ b/tests/e2e_cuj/helpers/constants.py @@ -34,24 +34,3 @@ class CodingAgent(StrEnum): SANDBOX_MCP_SERVICE_NAME = "system.ai.sandbox" WEB_SEARCH_MCP_SERVICE_NAME = "system.ai.web_search" -SANDBOX_MCP_TOOL_IDENTIFIERS = frozenset( - { - SANDBOX_MCP_SERVICE_NAME, - SANDBOX_MCP_SERVICE_NAME.replace(".", "-"), - SANDBOX_MCP_SERVICE_NAME.replace(".", "_"), - } -) - -TOOL_CALL_EVENT_TYPES = frozenset( - {"custom_tool_call", "function_call", "mcp_tool_call", "tool_call", "tool_use"} -) -TOOL_RESULT_EVENT_TYPES = frozenset( - { - "custom_tool_call_output", - "function_call_output", - "mcp_tool_result", - "mcp_tool_call_output", - "tool_result", - "tool_call_output", - } -) diff --git a/tests/e2e_cuj/test_cuj2_mps_mcp.py b/tests/e2e_cuj/test_cuj2_mps_mcp.py index ed2075e86..d367f564f 100644 --- a/tests/e2e_cuj/test_cuj2_mps_mcp.py +++ b/tests/e2e_cuj/test_cuj2_mps_mcp.py @@ -24,9 +24,6 @@ CODING_AGENT_BY_CLI_NAME, MODEL_PROVIDER_SERVICE_FIXTURES, SANDBOX_MCP_SERVICE_NAME, - SANDBOX_MCP_TOOL_IDENTIFIERS, - TOOL_CALL_EVENT_TYPES, - TOOL_RESULT_EVENT_TYPES, WEB_SEARCH_MCP_SERVICE_NAME, ) from .helpers.workspace import Workspace @@ -155,24 +152,6 @@ def _assert_generated_mcp_listings(session) -> None: assert not _mcp_name_listed(output, WEB_SEARCH_MCP_SERVICE_NAME), output -def _dict_nodes(value: object): - if isinstance(value, dict): - yield value - for child in value.values(): - yield from _dict_nodes(child) - elif isinstance(value, list): - for child in value: - yield from _dict_nodes(child) - - -def _call_id(mapping: dict) -> str | None: - for key in ("id", "call_id", "callId", "tool_use_id", "toolUseId"): - value = mapping.get(key) - if isinstance(value, str) and value: - return value - return None - - def _parent_transcripts(session, agent: str) -> dict[str, list[dict]]: return { path: records @@ -181,79 +160,29 @@ def _parent_transcripts(session, agent: str) -> dict[str, list[dict]]: } -def _correlated_run_code(transcripts: dict[str, list[dict]], marker: str) -> dict: - calls: dict[str, dict] = {} - results: dict[str, dict] = {} - for records in transcripts.values(): - for record in records: - for node in _dict_nodes(record): - event_type = str(node.get("type", "")).lower() - identifier = _call_id(node) - if not identifier: - continue - if event_type in TOOL_CALL_EVENT_TYPES: - calls[identifier] = node - elif event_type in TOOL_RESULT_EVENT_TYPES: - results[identifier] = node - - for identifier, call in calls.items(): - call_text = json.dumps(call, sort_keys=True, default=str).lower() - result = results.get(identifier) - result_text = json.dumps(result, sort_keys=True, default=str) if result else "" - failed = any( - node.get("is_error") is True - or str(node.get("status", "")).lower() in {"error", "failed", "failure"} - for node in _dict_nodes(result) - ) - if ( - "run_code" in call_text - and any(name in call_text for name in SANDBOX_MCP_TOOL_IDENTIFIERS) - and marker.lower() in call_text - and marker in result_text - and not failed - ): - return {"call_id": identifier, "call": call, "result": result} - raise AssertionError( - f"No successful sandbox run_code call/result carried marker {marker!r}; " - f"calls={list(calls)} results={list(results)}" - ) - - def _task_complete(session, agent: str, marker: str) -> bool: transcripts = _parent_transcripts(session, agent) - if not any( + return any( marker in answer for records in transcripts.values() for answer in assistant_answers(agent, records) - ): - return False - try: - _correlated_run_code(transcripts, marker) - except AssertionError: - return False - return True + ) def _assert_task_evidence(session, agent: str, marker: str) -> None: transcripts = _parent_transcripts(session, agent) - tool_evidence = _correlated_run_code(transcripts, marker) answers = [ answer for records in transcripts.values() for answer in assistant_answers(agent, records) ] assert any(marker in answer for answer in answers), answers session.record( f"cuj2-{agent}-evidence", - {"marker": marker, "tool": tool_evidence, "answers": answers}, + {"marker": marker, "answers": answers}, ) def _agent_prompt(marker: str) -> str: - return ( - "Complete this task using the configured model and MCP; do not answer from the prompt. " - f"Call the configured `run_code` MCP tool with a Python program that prints the exact " - f"marker `{marker}`. Wait for the real tool result. " - f"Final answer must contain the marker `{marker}`, with no invented tool output." - ) + return f"Reply with the exact marker `{marker}` and no other text. Do not use tools." def _assert_inference_request(recorder, request, *, provider: str, model: str, marker: str) -> None: @@ -298,11 +227,11 @@ def test_cuj_configuration(self, cuj): class TestCuj2CodexInference(_Cuj2Base): def test_cuj_codex_inference(self, cuj): - """Scenario: configure through the recorder and complete a real Codex MCP task. + """Scenario: configure through the recorder and complete a real Codex inference. Expected: Codex sends a Responses request for gpt-5-nano with the ug_e2e.providers.openai target header and receives its paired HTTP 200 response; its - parent transcript contains a correlated sandbox run_code call/result and final marker. + parent transcript contains the final marker. """ session, workspace, recorder = cuj published = workspace.config() @@ -333,7 +262,7 @@ def completed(screen): assert "Do you want to proceed?" not in screen, screen return _task_complete(session, CODEX, marker) - tui.wait_for(completed, "completed Codex MCP task", timeout=240) + tui.wait_for(completed, "completed Codex task", timeout=240) tui.exit_normally() request = recorder.expect_request( @@ -354,11 +283,11 @@ def completed(screen): class TestCuj2ClaudeInference(_Cuj2Base): def test_cuj_claude_inference(self, cuj): - """Scenario: configure through the recorder and complete a real Claude MCP task. + """Scenario: configure through the recorder and complete a real Claude inference. Expected: Claude sends a Messages request for claude-haiku-4-5-20251001 with the ug_e2e.providers.anthropic target header and receives its paired HTTP 200 response; its - parent transcript contains a correlated sandbox run_code call/result and final marker. + parent transcript contains the final marker. """ session, workspace, recorder = cuj published = workspace.config() @@ -383,29 +312,12 @@ def test_cuj_claude_inference(self, cuj): with AgentTerminal(session, CLAUDE, [str(session.binary), CLAUDE], "cuj2-claude") as tui: tui.boot() tui.submit(_agent_prompt(marker)) - permission_approved = False def completed(screen): - nonlocal permission_approved assert_no_terminal_api_error(screen) - if "Do you want to proceed?" in screen: - expected_prompt = ( - "system-ai-sandbox — Run Code Tool: (MCP)", - f'code: "print(\\"{marker}\\")"', - "Do you want to proceed?", - "❯ 1. Yes", - ) - assert all(part in screen for part in expected_prompt), ( - "Claude requested an unrecognized permission:\n" + screen - ) - if permission_approved: - return False - tui.send("\r", f"approve sandbox run_code for marker {marker}") - permission_approved = True - return False return _task_complete(session, CLAUDE, marker) - tui.wait_for(completed, "completed Claude MCP task", timeout=240) + tui.wait_for(completed, "completed Claude task", timeout=240) tui.exit_normally() request = recorder.expect_request( From aecb84a8a00557df228b90aaf884e2d6c061c52d Mon Sep 17 00:00:00 2001 From: Mason Cao Date: Tue, 6 Oct 2026 03:29:53 +0000 Subject: [PATCH 10/19] [AIGTWY-4877] Fix Claude model selection in CUJs --- src/ucode/agents/claude.py | 28 +++++++++++++++++++- src/ucode/cli.py | 21 ++++++++++++--- tests/e2e_cuj/helpers/evidence.py | 5 ++++ tests/test_agent_claude.py | 43 +++++++++++++++++++++++++++++++ tests/test_cli.py | 28 +++++++++++++++++++- tests/test_cuj_evidence.py | 8 ++++++ 6 files changed, 127 insertions(+), 6 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 83910d65d..e07f3da40 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -209,6 +209,10 @@ def _otel_trace_env(workspace: str) -> dict[str, str]: "sonnet": "ANTHROPIC_DEFAULT_SONNET_MODEL", "haiku": "ANTHROPIC_DEFAULT_HAIKU_MODEL", } +# Raw launch ids have historically been pinned to Claude's three standard family aliases. Fable is +# discovered/configured separately and should not silently turn an explicit model launch into a +# Fable selection. +CLAUDE_CUSTOM_MODEL_FAMILIES = ("opus", "sonnet", "haiku") # Launch-scoped feature flags that ucode may write into Claude settings. These # must be removed again when the corresponding launch flag is absent. CLAUDE_CONDITIONAL_ENV_KEYS = ("CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY",) @@ -1827,6 +1831,21 @@ def _launch_model_args(tool_args: list[str], launch_model: str | None) -> list[s return ["--model", launch_model] +def _launch_custom_model_settings(model: str) -> dict: + """Pin a raw Databricks model through Claude's launch-scoped family aliases. + + Claude Code validates ``ANTHROPIC_MODEL`` and its native ``--model`` value against the + client-side catalog, so neither can carry a Databricks model-service id. The family aliases + are passed through to the gateway unchanged. Keep this override in the inline launch settings; + ``write_tool_config`` must continue to leave the user's persistent model selection alone. + """ + return { + "env": { + CLAUDE_DEFAULT_MODEL_ENV_KEYS[family]: model for family in CLAUDE_CUSTOM_MODEL_FAMILIES + }, + } + + def _resolve_picker_model_id(model: str, settings_env: dict) -> str: """Resolve configured aliases and context suffixes for comparisons only.""" model = re.sub(r"\[(?:1m|200k)\]$", "", model) @@ -2063,7 +2082,14 @@ def launch( os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile")) settings_override = None launch_args = list(tool_args) - if options.user_pinned_model: + launch_custom_model = state.get("_claude_launch_custom_model") + if isinstance(launch_custom_model, str) and launch_custom_model: + # A ucode-owned --model is a raw Databricks id. Do not pass it through Claude Code's + # client-side model validation or persist it in ucode-settings.json. Native model options + # already present in tool_args remain untouched below. + os.environ.pop("ANTHROPIC_MODEL", None) + settings_override = _launch_custom_model_settings(launch_custom_model) + elif options.user_pinned_model: os.environ["ANTHROPIC_MODEL"] = options.user_pinned_model settings_override = {"env": {"ANTHROPIC_MODEL": options.user_pinned_model}} launch_args = [ diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 6c377cd35..1c3c2be16 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -2855,6 +2855,11 @@ def _launch_tool( # Codex keeps an explicit --model in ctx.args and passes it to its CLI verbatim. if model and tool != "claude": resolved_model = model + if model and tool == "claude" and not provider: + # A direct Claude model is applied through the launch-scoped family aliases below. + # Clear a managed/default route-root pin so it cannot win over the user's explicit + # choice in the persistent settings file. + route_root_model = None if coding_agent_config_defaults and not state.get("claude_static_models") and not relayed: picker_catalog = claude_agent.default_model_picker_catalog( coding_agent_config_defaults, @@ -2871,7 +2876,8 @@ def _launch_tool( picker_catalog=picker_catalog, relayed=relayed, route_root_model=route_root_model, - # Claude's explicit model is launch-scoped and is passed through LaunchOptions below. + # Claude's explicit model is added to a transient state marker after configuration; + # passing it here would write the selection into persistent settings. custom_model=None, coding_agent_config_defaults=coding_agent_config_defaults, parent_schema=parent_schema, @@ -2880,7 +2886,7 @@ def _launch_tool( # Claude re-adds an out-of-catalog saved model to /model even when built-ins are # replaced. Keep the catalog launch-scoped and leave the user's settings alone. state["_claude_launch_picker_models"] = picker_catalog.model_ids - if managed is None and (explicit_provider or parent_schema): + if managed is None and (explicit_provider or parent_schema) and not model: # The permanent Default row should also resolve within the selected catalog. state["_claude_launch_default_model"] = ( claude_agent.default_model( @@ -2921,6 +2927,13 @@ def _launch_tool( if tool == "claude": if provider: state["_claude_launch_provider"] = provider + elif model: + # Claude Code validates ANTHROPIC_MODEL and its native --model value against its + # own catalog. Keep ucode's explicit Databricks model out of both paths: launch() + # applies it to the family aliases in a per-process settings override instead. + # This marker is deliberately added after configure_tool saved state, so it can + # never become a persisted model selection. + state["_claude_launch_custom_model"] = model elif tool == "codex": if provider: state["_codex_launch_provider"] = provider @@ -3215,8 +3228,8 @@ def claude_cmd( typer.Option( "--model", help="Launch on a specific Databricks model id (e.g. a UC " - "`..`). Pinned via ANTHROPIC_MODEL so the gateway " - "resolves it — unlike Claude Code's own --model, which rejects non-catalog ids. " + "`..`). Pinned via launch-scoped Claude family aliases so the " + "gateway resolves it — unlike Claude Code's own --model, which rejects non-catalog ids. " "With --provider, pass a family (opus/sonnet/haiku) or a target the service allows to " "start on that tier instead of Claude Code's opus default. Pass before any `--` separator.", ), diff --git a/tests/e2e_cuj/helpers/evidence.py b/tests/e2e_cuj/helpers/evidence.py index 7c955d012..09deb2fda 100644 --- a/tests/e2e_cuj/helpers/evidence.py +++ b/tests/e2e_cuj/helpers/evidence.py @@ -10,6 +10,10 @@ from .constants import CLAUDE, CODEX +_EXPLICIT_MODEL_ALIASES = { + "anthropic.claude-haiku-4-5-20251001-v1:0": "system.ai.claude-haiku-4-5", +} + def canonical_model(value): """Only documented gateway aliases, not fuzzy family/date matching.""" @@ -17,6 +21,7 @@ def canonical_model(value): # `anthropic-aigw-1234abcd-system.ai.glm-5-2`. value = re.sub(r"^anthropic-aigw-[0-9a-f]{8}-", "", value) value = value.removesuffix("[1m]") + value = _EXPLICIT_MODEL_ALIASES.get(value, value) value = value.removeprefix("system.ai.") # Codex writes dotted major/minor versions (`gpt-5.6-sol`) while the # gateway model identifier uses a hyphen (`system.ai.gpt-5-6-sol`). diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 148543d1e..005190d2a 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -2482,6 +2482,49 @@ def test_launch_model_is_only_set_for_current_process(self, monkeypatch): settings = json.loads(calls[0][2]) assert settings["env"]["ANTHROPIC_MODEL"] == "cat.schema.model" + def test_launch_custom_model_uses_family_aliases_without_native_model( + self, monkeypatch, tmp_path + ): + calls: list[list[str]] = [] + monkeypatch.setenv("ANTHROPIC_MODEL", "stale-model") + monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "token") + monkeypatch.setattr(claude, "exec_or_spawn", lambda argv: calls.append(argv)) + settings_path = tmp_path / "ucode-settings.json" + settings_path.write_text(json.dumps({"env": {"USER_SETTING": "keep"}})) + monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", settings_path) + original_settings = settings_path.read_bytes() + + custom_model = "system.ai.claude-opus-4-8" + claude.launch( + {"workspace": WS, "_claude_launch_custom_model": custom_model}, + ["--debug"], + options=LaunchOptions(user_pinned_model=custom_model), + ) + + assert os.environ.get("ANTHROPIC_MODEL") is None + assert calls[0][0:2] == ["claude", "--settings"] + settings = json.loads(calls[0][2]) + for family in claude.CLAUDE_CUSTOM_MODEL_FAMILIES: + key = claude.CLAUDE_DEFAULT_MODEL_ENV_KEYS[family] + assert settings["env"][key] == custom_model + assert "ANTHROPIC_DEFAULT_FABLE_MODEL" not in settings["env"] + assert "--model" not in calls[0] + assert calls[0][-1] == "--debug" + assert settings_path.read_bytes() == original_settings + + def test_launch_custom_model_preserves_forwarded_native_model(self, monkeypatch): + calls: list[list[str]] = [] + monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "token") + monkeypatch.setattr(claude, "exec_or_spawn", lambda argv: calls.append(argv)) + + claude.launch( + {"workspace": WS, "_claude_launch_custom_model": "system.ai.claude-opus-4-8"}, + ["--model", "claude-sonnet-5"], + options=LaunchOptions(user_pinned_model="system.ai.claude-opus-4-8"), + ) + + assert calls[0][-2:] == ["--model", "claude-sonnet-5"] + def test_launch_default_model_is_inherited_by_smart_routing(self, monkeypatch): monkeypatch.delenv("ANTHROPIC_DEFAULT_MODEL", raising=False) monkeypatch.setattr(v2, "launch_claude", Mock()) diff --git a/tests/test_cli.py b/tests/test_cli.py index 1c44bed56..c654cdcab 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -813,6 +813,11 @@ def test_claude_model_location_preserves_explicit_model(self): assert calls["launch"].call_args.kwargs["options"].user_pinned_model == ( "main.default.claude-opus-5" ) + assert "_claude_launch_default_model" not in calls["launch"].call_args.args[1] + assert ( + calls["launch"].call_args.args[1]["_claude_launch_custom_model"] + == "main.default.claude-opus-5" + ) assert calls["launch"].call_args.args[2] == [] @pytest.mark.parametrize( @@ -1504,13 +1509,34 @@ def test_model_is_launch_scoped_for_claude(self, monkeypatch): ): result = runner.invoke(app, ["claude", "--model", "cat.schema.claude-opus-5"]) assert result.exit_code == 0, result.output - # The model is passed through invocation-scoped LaunchOptions, not persisted in settings. + # The model is passed through invocation-scoped launch state/options, not persisted in + # settings. assert mock_configure.call_args.kwargs["custom_model"] is None assert mock_configure.call_args.kwargs["route_root_model"] is None + assert ( + mock_launch.call_args.args[1]["_claude_launch_custom_model"] + == "cat.schema.claude-opus-5" + ) assert ( mock_launch.call_args.kwargs["options"].user_pinned_model == "cat.schema.claude-opus-5" ) + def test_explicit_model_clears_managed_default_route_root(self): + managed = { + "enabled_agents": { + "claude": {"model_config": {"default_model": "system.ai.claude-sonnet-5"}} + } + } + with _launch_policy_patches(managed) as calls: + result = runner.invoke(app, ["claude", "--model", "system.ai.claude-opus-4-8"]) + + assert result.exit_code == 0, result.output + assert calls["configure"].call_args.kwargs["route_root_model"] is None + assert ( + calls["launch"].call_args.args[1]["_claude_launch_custom_model"] + == "system.ai.claude-opus-4-8" + ) + def test_v2_model_sets_transient_launch_override(self, monkeypatch): monkeypatch.setenv("ENABLE_SMART_ROUTING_V2", "1") state = {**MINIMAL_STATE, "claude_models": {"opus": "system.ai.claude-opus-4-8"}} diff --git a/tests/test_cuj_evidence.py b/tests/test_cuj_evidence.py index fa50ef2ef..76d857c14 100644 --- a/tests/test_cuj_evidence.py +++ b/tests/test_cuj_evidence.py @@ -163,6 +163,14 @@ def test_cuj_evidence_ignores_existing_session(tmp_path): ("anthropic-aigw-1234abcd-system.ai.glm-5-2[1m]", "system.ai.glm-5-2"), ("gpt-5.6-sol", "system.ai.gpt-5-6-sol"), ("system.ai.gpt-6-sol", "system.ai.gpt-6-sol"), + ( + "anthropic.claude-haiku-4-5-20251001-v1:0", + "system.ai.claude-haiku-4-5", + ), + ( + "anthropic.claude-haiku-4-5-20251002-v1:0", + "system.ai.anthropic.claude-haiku-4-5-20251002-v1:0", + ), ( "anthropic-aigw-invalid-claude-opus-4-8", "system.ai.anthropic-aigw-invalid-claude-opus-4-8", From 92dad14d2820e0fc8608f3d71aa9270ddb040c58 Mon Sep 17 00:00:00 2001 From: Mason Cao Date: Tue, 6 Oct 2026 03:56:02 +0000 Subject: [PATCH 11/19] [AIGTWY-4877] Override saved Claude model for explicit launches --- src/ucode/agents/claude.py | 7 ++++--- tests/test_agent_claude.py | 8 ++++++-- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index e07f3da40..f7d691b7f 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -213,6 +213,7 @@ def _otel_trace_env(workspace: str) -> dict[str, str]: # discovered/configured separately and should not silently turn an explicit model launch into a # Fable selection. CLAUDE_CUSTOM_MODEL_FAMILIES = ("opus", "sonnet", "haiku") +CLAUDE_CUSTOM_MODEL_SELECTOR = "opus" # Launch-scoped feature flags that ucode may write into Claude settings. These # must be removed again when the corresponding launch flag is absent. CLAUDE_CONDITIONAL_ENV_KEYS = ("CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY",) @@ -2085,9 +2086,9 @@ def launch( launch_custom_model = state.get("_claude_launch_custom_model") if isinstance(launch_custom_model, str) and launch_custom_model: # A ucode-owned --model is a raw Databricks id. Do not pass it through Claude Code's - # client-side model validation or persist it in ucode-settings.json. Native model options - # already present in tool_args remain untouched below. - os.environ.pop("ANTHROPIC_MODEL", None) + # client-side model validation or persist it in ucode-settings.json. A native family alias + # must still outrank any saved picker model; every alias resolves to the custom id below. + os.environ["ANTHROPIC_MODEL"] = CLAUDE_CUSTOM_MODEL_SELECTOR settings_override = _launch_custom_model_settings(launch_custom_model) elif options.user_pinned_model: os.environ["ANTHROPIC_MODEL"] = options.user_pinned_model diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 005190d2a..d299410d0 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -2490,7 +2490,10 @@ def test_launch_custom_model_uses_family_aliases_without_native_model( monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "token") monkeypatch.setattr(claude, "exec_or_spawn", lambda argv: calls.append(argv)) settings_path = tmp_path / "ucode-settings.json" - settings_path.write_text(json.dumps({"env": {"USER_SETTING": "keep"}})) + saved_model = "system.ai.claude-haiku-4-5" + settings_path.write_text( + json.dumps({"model": saved_model, "env": {"USER_SETTING": "keep"}}) + ) monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", settings_path) original_settings = settings_path.read_bytes() @@ -2501,9 +2504,10 @@ def test_launch_custom_model_uses_family_aliases_without_native_model( options=LaunchOptions(user_pinned_model=custom_model), ) - assert os.environ.get("ANTHROPIC_MODEL") is None + assert os.environ["ANTHROPIC_MODEL"] == claude.CLAUDE_CUSTOM_MODEL_SELECTOR assert calls[0][0:2] == ["claude", "--settings"] settings = json.loads(calls[0][2]) + assert settings["model"] == saved_model for family in claude.CLAUDE_CUSTOM_MODEL_FAMILIES: key = claude.CLAUDE_DEFAULT_MODEL_ENV_KEYS[family] assert settings["env"][key] == custom_model From d0ca0ff5f18600cdbc3e24f42e19f802cce43579 Mon Sep 17 00:00:00 2001 From: Mason Cao Date: Tue, 6 Oct 2026 04:20:11 +0000 Subject: [PATCH 12/19] [AIGTWY-4877] Make explicit Claude catalog model override saved model --- src/ucode/agents/claude.py | 35 ++++++++++++++++++++++++++++++----- tests/test_agent_claude.py | 24 ++++++++++++++++++++++++ 2 files changed, 54 insertions(+), 5 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index f7d691b7f..5b4053f3d 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -1858,6 +1858,24 @@ def _resolve_picker_model_id(model: str, settings_env: dict) -> str: return re.sub(r"\[(?:1m|200k)\]$", "", model) +def _is_managed_launch_model(state: dict, model: str) -> bool: + """Return whether *model* is present in Claude's managed model catalog.""" + picker_models = state.get("_claude_launch_picker_models") + if not isinstance(picker_models, list) or not picker_models: + picker_models = state.get("claude_static_models") + if not isinstance(picker_models, list) or not picker_models: + return False + + settings_env = read_json_safe(CLAUDE_SETTINGS_PATH).get("env") + settings_env = settings_env if isinstance(settings_env, dict) else {} + resolved_model = _resolve_picker_model_id(model, settings_env) + return any( + isinstance(picker_model, str) + and _resolve_picker_model_id(picker_model, settings_env) == resolved_model + for picker_model in picker_models + ) + + def _resolve_launch_binary(binary: str) -> str: """Resolve Claude's native executable without sending arguments through a batch shim.""" if os.name != "nt": @@ -2085,11 +2103,18 @@ def launch( launch_args = list(tool_args) launch_custom_model = state.get("_claude_launch_custom_model") if isinstance(launch_custom_model, str) and launch_custom_model: - # A ucode-owned --model is a raw Databricks id. Do not pass it through Claude Code's - # client-side model validation or persist it in ucode-settings.json. A native family alias - # must still outrank any saved picker model; every alias resolves to the custom id below. - os.environ["ANTHROPIC_MODEL"] = CLAUDE_CUSTOM_MODEL_SELECTOR - settings_override = _launch_custom_model_settings(launch_custom_model) + if _is_managed_launch_model(state, launch_custom_model): + # Managed model ids are in Claude's catalog, so native --model gives them precedence + # over an enterprise or user settings model saved by Claude Code. + launch_args = [ + *_launch_model_args(tool_args, launch_custom_model), + *tool_args, + ] + else: + # Arbitrary Databricks ids are not in Claude Code's client-side catalog. Route those + # through the family aliases, which the gateway receives unchanged. + os.environ["ANTHROPIC_MODEL"] = CLAUDE_CUSTOM_MODEL_SELECTOR + settings_override = _launch_custom_model_settings(launch_custom_model) elif options.user_pinned_model: os.environ["ANTHROPIC_MODEL"] = options.user_pinned_model settings_override = {"env": {"ANTHROPIC_MODEL": options.user_pinned_model}} diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index d299410d0..038ba02e3 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -2529,6 +2529,30 @@ def test_launch_custom_model_preserves_forwarded_native_model(self, monkeypatch) assert calls[0][-2:] == ["--model", "claude-sonnet-5"] + def test_launch_managed_custom_model_uses_native_model_over_stale_settings( + self, monkeypatch, tmp_path + ): + calls: list[list[str]] = [] + settings_path = tmp_path / "settings.json" + settings_path.write_text(json.dumps({"model": "system.ai.claude-haiku-4-5"})) + monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", settings_path) + monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "token") + monkeypatch.setattr(claude, "exec_or_spawn", lambda argv: calls.append(argv)) + + custom_model = "system.ai.claude-opus-4-8" + claude.launch( + { + "workspace": WS, + "_claude_launch_custom_model": custom_model, + "claude_static_models": ["system.ai.claude-opus-4-8", "system.ai.claude-haiku-4-5"], + }, + [], + options=LaunchOptions(user_pinned_model=custom_model), + ) + + assert calls[0][1:3] == ["--settings", str(settings_path)] + assert calls[0][-2:] == ["--model", custom_model] + def test_launch_default_model_is_inherited_by_smart_routing(self, monkeypatch): monkeypatch.delenv("ANTHROPIC_DEFAULT_MODEL", raising=False) monkeypatch.setattr(v2, "launch_claude", Mock()) From d04b9958cf8690727621c8732f42dbf23f187801 Mon Sep 17 00:00:00 2001 From: Mason Cao Date: Tue, 6 Oct 2026 12:37:44 +0000 Subject: [PATCH 13/19] [AIGTWY-4877] Fix Claude model pinning and skill confirmation --- src/ucode/agents/claude.py | 36 ++++--------------- .../test_ug_smart_routing_hooks.py | 5 ++- tests/test_agent_claude.py | 31 +++++++++++----- 3 files changed, 32 insertions(+), 40 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 5b4053f3d..c9d3155b2 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -1858,24 +1858,6 @@ def _resolve_picker_model_id(model: str, settings_env: dict) -> str: return re.sub(r"\[(?:1m|200k)\]$", "", model) -def _is_managed_launch_model(state: dict, model: str) -> bool: - """Return whether *model* is present in Claude's managed model catalog.""" - picker_models = state.get("_claude_launch_picker_models") - if not isinstance(picker_models, list) or not picker_models: - picker_models = state.get("claude_static_models") - if not isinstance(picker_models, list) or not picker_models: - return False - - settings_env = read_json_safe(CLAUDE_SETTINGS_PATH).get("env") - settings_env = settings_env if isinstance(settings_env, dict) else {} - resolved_model = _resolve_picker_model_id(model, settings_env) - return any( - isinstance(picker_model, str) - and _resolve_picker_model_id(picker_model, settings_env) == resolved_model - for picker_model in picker_models - ) - - def _resolve_launch_binary(binary: str) -> str: """Resolve Claude's native executable without sending arguments through a batch shim.""" if os.name != "nt": @@ -2103,18 +2085,12 @@ def launch( launch_args = list(tool_args) launch_custom_model = state.get("_claude_launch_custom_model") if isinstance(launch_custom_model, str) and launch_custom_model: - if _is_managed_launch_model(state, launch_custom_model): - # Managed model ids are in Claude's catalog, so native --model gives them precedence - # over an enterprise or user settings model saved by Claude Code. - launch_args = [ - *_launch_model_args(tool_args, launch_custom_model), - *tool_args, - ] - else: - # Arbitrary Databricks ids are not in Claude Code's client-side catalog. Route those - # through the family aliases, which the gateway receives unchanged. - os.environ["ANTHROPIC_MODEL"] = CLAUDE_CUSTOM_MODEL_SELECTOR - settings_override = _launch_custom_model_settings(launch_custom_model) + os.environ["ANTHROPIC_MODEL"] = CLAUDE_CUSTOM_MODEL_SELECTOR + settings_override = _launch_custom_model_settings(launch_custom_model) + launch_args = [ + *_launch_model_args(tool_args, CLAUDE_CUSTOM_MODEL_SELECTOR), + *tool_args, + ] elif options.user_pinned_model: os.environ["ANTHROPIC_MODEL"] = options.user_pinned_model settings_override = {"env": {"ANTHROPIC_MODEL": options.user_pinned_model}} diff --git a/tests/integration/test_ug_smart_routing_hooks.py b/tests/integration/test_ug_smart_routing_hooks.py index 67f55076c..58ddcffe3 100644 --- a/tests/integration/test_ug_smart_routing_hooks.py +++ b/tests/integration/test_ug_smart_routing_hooks.py @@ -131,7 +131,10 @@ def _toggle_with_skill(tui, session, agent: str, enabled: bool) -> None: lambda screen: ( "Smart Router" in screen and confirmation in screen - and assistant_answer_contains(session, agent, confirmation) + and ( + assistant_answer_contains(session, agent, confirmation) + or assistant_answer_contains(session, agent, f"**{state}** for this session") + ) ), f"the installed Smart Router skill to turn routing {state}", timeout=120, diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 038ba02e3..b7dc01ecc 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -2482,7 +2482,7 @@ def test_launch_model_is_only_set_for_current_process(self, monkeypatch): settings = json.loads(calls[0][2]) assert settings["env"]["ANTHROPIC_MODEL"] == "cat.schema.model" - def test_launch_custom_model_uses_family_aliases_without_native_model( + def test_launch_custom_model_uses_family_aliases_and_native_selector( self, monkeypatch, tmp_path ): calls: list[list[str]] = [] @@ -2511,9 +2511,10 @@ def test_launch_custom_model_uses_family_aliases_without_native_model( for family in claude.CLAUDE_CUSTOM_MODEL_FAMILIES: key = claude.CLAUDE_DEFAULT_MODEL_ENV_KEYS[family] assert settings["env"][key] == custom_model + assert settings["env"]["USER_SETTING"] == "keep" assert "ANTHROPIC_DEFAULT_FABLE_MODEL" not in settings["env"] - assert "--model" not in calls[0] - assert calls[0][-1] == "--debug" + assert calls[0][3:] == ["--model", claude.CLAUDE_CUSTOM_MODEL_SELECTOR, "--debug"] + assert custom_model not in calls[0][3:] assert settings_path.read_bytes() == original_settings def test_launch_custom_model_preserves_forwarded_native_model(self, monkeypatch): @@ -2527,14 +2528,16 @@ def test_launch_custom_model_preserves_forwarded_native_model(self, monkeypatch) options=LaunchOptions(user_pinned_model="system.ai.claude-opus-4-8"), ) - assert calls[0][-2:] == ["--model", "claude-sonnet-5"] + assert calls[0][3:] == ["--model", "claude-sonnet-5"] - def test_launch_managed_custom_model_uses_native_model_over_stale_settings( + def test_launch_custom_model_uses_native_selector_over_stale_settings( self, monkeypatch, tmp_path ): calls: list[list[str]] = [] settings_path = tmp_path / "settings.json" - settings_path.write_text(json.dumps({"model": "system.ai.claude-haiku-4-5"})) + saved_model = "system.ai.claude-haiku-4-5" + settings_path.write_text(json.dumps({"model": saved_model})) + original_settings = settings_path.read_bytes() monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", settings_path) monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "token") monkeypatch.setattr(claude, "exec_or_spawn", lambda argv: calls.append(argv)) @@ -2544,14 +2547,24 @@ def test_launch_managed_custom_model_uses_native_model_over_stale_settings( { "workspace": WS, "_claude_launch_custom_model": custom_model, - "claude_static_models": ["system.ai.claude-opus-4-8", "system.ai.claude-haiku-4-5"], + "claude_static_models": [ + "system.ai.claude-opus-4-8", + "system.ai.claude-haiku-4-5", + ], }, [], options=LaunchOptions(user_pinned_model=custom_model), ) - assert calls[0][1:3] == ["--settings", str(settings_path)] - assert calls[0][-2:] == ["--model", custom_model] + assert calls[0][0:2] == ["claude", "--settings"] + settings = json.loads(calls[0][2]) + assert settings["model"] == saved_model + for family in claude.CLAUDE_CUSTOM_MODEL_FAMILIES: + key = claude.CLAUDE_DEFAULT_MODEL_ENV_KEYS[family] + assert settings["env"][key] == custom_model + assert calls[0][3:] == ["--model", claude.CLAUDE_CUSTOM_MODEL_SELECTOR] + assert custom_model not in calls[0][3:] + assert settings_path.read_bytes() == original_settings def test_launch_default_model_is_inherited_by_smart_routing(self, monkeypatch): monkeypatch.delenv("ANTHROPIC_DEFAULT_MODEL", raising=False) From e920a2a5fa605333ef36e358746c7f84a5a73109 Mon Sep 17 00:00:00 2001 From: Mason Cao Date: Tue, 6 Oct 2026 12:44:42 +0000 Subject: [PATCH 14/19] [AIGTWY-4877] Reuse CUJ request recorder helpers --- src/ucode/databricks.py | 4 ++-- tests/e2e_cuj/test_cuj2_mps_mcp.py | 25 +++++++++++++------------ 2 files changed, 15 insertions(+), 14 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 9c720bca9..8e7b64aba 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -2384,7 +2384,7 @@ def list_model_provider_services( # reach the next. return [dict(service) for service in cached], None - origin = workspace_origin(workspace) + origin = normalize_workspace_url(workspace) services: list[dict] = [] page_token: str | None = None seen_tokens: set[str] = set() @@ -2468,7 +2468,7 @@ def get_model_provider_service( server-side filter) makes a service that plainly exists look absent. Addressing it directly removes that whole class of false negative. """ - origin = workspace_origin(workspace) + origin = normalize_workspace_url(workspace) url = f"{origin}/api/2.1/unity-catalog/model-provider-services/{service_name}" payload, reason = _http_get_json(url, token, timeout=30) if payload is None: diff --git a/tests/e2e_cuj/test_cuj2_mps_mcp.py b/tests/e2e_cuj/test_cuj2_mps_mcp.py index d367f564f..2ee4500aa 100644 --- a/tests/e2e_cuj/test_cuj2_mps_mcp.py +++ b/tests/e2e_cuj/test_cuj2_mps_mcp.py @@ -22,6 +22,7 @@ CLAUDE, CODEX, CODING_AGENT_BY_CLI_NAME, + INFERENCE_PATHS, MODEL_PROVIDER_SERVICE_FIXTURES, SANDBOX_MCP_SERVICE_NAME, WEB_SEARCH_MCP_SERVICE_NAME, @@ -239,18 +240,18 @@ def test_cuj_codex_inference(self, cuj): _assert_mps_catalog(workspace) _assert_mcp_services(workspace) - session.configure( + recorder.configure_session( + session, [ "configure", - "--workspace", - recorder.url, "--skip-upgrade", "--disable-databricks-ai-tools", - ] + ], ) - _assert_configured_files(session, recorder.url) + _assert_configured_files(session, workspace.url) _assert_generated_mcp_listings(session) + 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: @@ -267,7 +268,7 @@ def completed(screen): request = recorder.expect_request( method="POST", - path="/ai-gateway/codex/v1/responses", + path=INFERENCE_PATHS[CODEX], after=checkpoint, timeout=240, ) @@ -295,18 +296,18 @@ def test_cuj_claude_inference(self, cuj): _assert_mps_catalog(workspace) _assert_mcp_services(workspace) - session.configure( + recorder.configure_session( + session, [ "configure", - "--workspace", - recorder.url, "--skip-upgrade", "--disable-databricks-ai-tools", - ] + ], ) - _assert_configured_files(session, recorder.url) + _assert_configured_files(session, workspace.url) _assert_generated_mcp_listings(session) + 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: @@ -322,7 +323,7 @@ def completed(screen): request = recorder.expect_request( method="POST", - path="/ai-gateway/anthropic/v1/messages", + path=INFERENCE_PATHS[CLAUDE], after=checkpoint, timeout=240, ) From 71f681fd5cc6ce2b91e6c26c840ac226f51982a8 Mon Sep 17 00:00:00 2001 From: Mason Cao Date: Tue, 6 Oct 2026 12:57:28 +0000 Subject: [PATCH 15/19] [AIGTWY-4877] Revert CUJ machine settings between cases --- tests/e2e_cuj/conftest.py | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/tests/e2e_cuj/conftest.py b/tests/e2e_cuj/conftest.py index 2306b8c12..b11fde12a 100644 --- a/tests/e2e_cuj/conftest.py +++ b/tests/e2e_cuj/conftest.py @@ -8,6 +8,8 @@ import pytest from databricks.sdk.errors import DatabricksError +from tests.integration.utils.terminal import TerminalProcess + from .helpers.constants import CLAUDE, CODEX, MANAGED_PATHS from .helpers.session import UserSession from .helpers.tui_request_recorder import TuiRequestRecorder @@ -45,7 +47,24 @@ def cuj(request, setup_workspace, tmp_path_factory): try: yield session, workspace, recorder finally: - workspace.assert_unchanged(published) + 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" + ) 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 From 19793614693202447a1c3127576ada696ecc677f Mon Sep 17 00:00:00 2001 From: Mason Cao Date: Tue, 6 Oct 2026 13:13:02 +0000 Subject: [PATCH 16/19] [AIGTWY-4877] Preserve recorder origin for Claude discovery --- src/ucode/databricks.py | 4 ++-- tests/test_databricks.py | 15 +++++++++++++++ 2 files changed, 17 insertions(+), 2 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 8e7b64aba..b67ba8f7e 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -2954,14 +2954,14 @@ def _get_anthropic_models_json( parent_schema: str | None = None, provider: str | None = None, ) -> tuple[dict | list | None, str | None]: - hostname = workspace_hostname(workspace) + origin = normalize_workspace_url(workspace) headers = None if provider is not None: headers = {MODEL_PROVIDER_SERVICE_HEADER: provider} elif parent_schema is not None: headers = {MODEL_SERVICE_PARENT_SCHEMA_HEADER: parent_schema} return _http_get_json( - f"https://{hostname}{ANTHROPIC_MODELS_PATH}?limit=1000", + f"{origin}{ANTHROPIC_MODELS_PATH}?limit=1000", token, max_retries=_ANTHROPIC_MODEL_DISCOVERY_SETUP_MAX_RETRIES, **({"headers": headers} if headers is not None else {}), diff --git a/tests/test_databricks.py b/tests/test_databricks.py index eda29e87b..ca619be8b 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -318,6 +318,21 @@ def test_multiple_locations_preserve_order(self): class TestDiscoverClaudeModels: + def test_preserves_explicit_origin(self, monkeypatch): + captured = {} + + def fake_get(url, token, **kwargs): + captured["request"] = (url, token, kwargs) + return {"data": []}, None + + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) + + db_mod.list_anthropic_models("http://127.0.0.1:43123", "token") + + assert captured["request"][0] == ( + "http://127.0.0.1:43123/ai-gateway/anthropic/v1/models?limit=1000" + ) + def test_lists_all_anthropic_model_ids_without_legacy_validation(self, monkeypatch): captured = {} payload = { From 934904c8836fc2c86160615f32eb7b448e0f90a9 Mon Sep 17 00:00:00 2001 From: Mason Cao Date: Tue, 6 Oct 2026 13:48:44 +0000 Subject: [PATCH 17/19] [AIGTWY-4877] Trim stale CUJ changes --- src/ucode/agents/claude.py | 30 +------ src/ucode/cli.py | 21 +---- tests/README.md | 17 ++-- tests/e2e_cuj/AGENTS.md | 5 -- tests/e2e_cuj/helpers/evidence.py | 5 -- tests/integration/README.md | 10 +-- .../test_ug_smart_routing_hooks.py | 5 +- tests/test_agent_claude.py | 84 ------------------- tests/test_cli.py | 28 +------ tests/test_cuj_evidence.py | 8 -- 10 files changed, 20 insertions(+), 193 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index c9d3155b2..83910d65d 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -209,11 +209,6 @@ def _otel_trace_env(workspace: str) -> dict[str, str]: "sonnet": "ANTHROPIC_DEFAULT_SONNET_MODEL", "haiku": "ANTHROPIC_DEFAULT_HAIKU_MODEL", } -# Raw launch ids have historically been pinned to Claude's three standard family aliases. Fable is -# discovered/configured separately and should not silently turn an explicit model launch into a -# Fable selection. -CLAUDE_CUSTOM_MODEL_FAMILIES = ("opus", "sonnet", "haiku") -CLAUDE_CUSTOM_MODEL_SELECTOR = "opus" # Launch-scoped feature flags that ucode may write into Claude settings. These # must be removed again when the corresponding launch flag is absent. CLAUDE_CONDITIONAL_ENV_KEYS = ("CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY",) @@ -1832,21 +1827,6 @@ def _launch_model_args(tool_args: list[str], launch_model: str | None) -> list[s return ["--model", launch_model] -def _launch_custom_model_settings(model: str) -> dict: - """Pin a raw Databricks model through Claude's launch-scoped family aliases. - - Claude Code validates ``ANTHROPIC_MODEL`` and its native ``--model`` value against the - client-side catalog, so neither can carry a Databricks model-service id. The family aliases - are passed through to the gateway unchanged. Keep this override in the inline launch settings; - ``write_tool_config`` must continue to leave the user's persistent model selection alone. - """ - return { - "env": { - CLAUDE_DEFAULT_MODEL_ENV_KEYS[family]: model for family in CLAUDE_CUSTOM_MODEL_FAMILIES - }, - } - - def _resolve_picker_model_id(model: str, settings_env: dict) -> str: """Resolve configured aliases and context suffixes for comparisons only.""" model = re.sub(r"\[(?:1m|200k)\]$", "", model) @@ -2083,15 +2063,7 @@ def launch( os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile")) settings_override = None launch_args = list(tool_args) - launch_custom_model = state.get("_claude_launch_custom_model") - if isinstance(launch_custom_model, str) and launch_custom_model: - os.environ["ANTHROPIC_MODEL"] = CLAUDE_CUSTOM_MODEL_SELECTOR - settings_override = _launch_custom_model_settings(launch_custom_model) - launch_args = [ - *_launch_model_args(tool_args, CLAUDE_CUSTOM_MODEL_SELECTOR), - *tool_args, - ] - elif options.user_pinned_model: + if options.user_pinned_model: os.environ["ANTHROPIC_MODEL"] = options.user_pinned_model settings_override = {"env": {"ANTHROPIC_MODEL": options.user_pinned_model}} launch_args = [ diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 1c3c2be16..6c377cd35 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -2855,11 +2855,6 @@ def _launch_tool( # Codex keeps an explicit --model in ctx.args and passes it to its CLI verbatim. if model and tool != "claude": resolved_model = model - if model and tool == "claude" and not provider: - # A direct Claude model is applied through the launch-scoped family aliases below. - # Clear a managed/default route-root pin so it cannot win over the user's explicit - # choice in the persistent settings file. - route_root_model = None if coding_agent_config_defaults and not state.get("claude_static_models") and not relayed: picker_catalog = claude_agent.default_model_picker_catalog( coding_agent_config_defaults, @@ -2876,8 +2871,7 @@ def _launch_tool( picker_catalog=picker_catalog, relayed=relayed, route_root_model=route_root_model, - # Claude's explicit model is added to a transient state marker after configuration; - # passing it here would write the selection into persistent settings. + # Claude's explicit model is launch-scoped and is passed through LaunchOptions below. custom_model=None, coding_agent_config_defaults=coding_agent_config_defaults, parent_schema=parent_schema, @@ -2886,7 +2880,7 @@ def _launch_tool( # Claude re-adds an out-of-catalog saved model to /model even when built-ins are # replaced. Keep the catalog launch-scoped and leave the user's settings alone. state["_claude_launch_picker_models"] = picker_catalog.model_ids - if managed is None and (explicit_provider or parent_schema) and not model: + if managed is None and (explicit_provider or parent_schema): # The permanent Default row should also resolve within the selected catalog. state["_claude_launch_default_model"] = ( claude_agent.default_model( @@ -2927,13 +2921,6 @@ def _launch_tool( if tool == "claude": if provider: state["_claude_launch_provider"] = provider - elif model: - # Claude Code validates ANTHROPIC_MODEL and its native --model value against its - # own catalog. Keep ucode's explicit Databricks model out of both paths: launch() - # applies it to the family aliases in a per-process settings override instead. - # This marker is deliberately added after configure_tool saved state, so it can - # never become a persisted model selection. - state["_claude_launch_custom_model"] = model elif tool == "codex": if provider: state["_codex_launch_provider"] = provider @@ -3228,8 +3215,8 @@ def claude_cmd( typer.Option( "--model", help="Launch on a specific Databricks model id (e.g. a UC " - "`..`). Pinned via launch-scoped Claude family aliases so the " - "gateway resolves it — unlike Claude Code's own --model, which rejects non-catalog ids. " + "`..`). Pinned via ANTHROPIC_MODEL so the gateway " + "resolves it — unlike Claude Code's own --model, which rejects non-catalog ids. " "With --provider, pass a family (opus/sonnet/haiku) or a target the service allows to " "start on that tier instead of Claude Code's opus default. Pass before any `--` separator.", ), diff --git a/tests/README.md b/tests/README.md index 87a62f453..c125b45bc 100644 --- a/tests/README.md +++ b/tests/README.md @@ -7,19 +7,18 @@ mocks, monkeypatching, fake binaries/services, or fabricated ug state. | Category | Location | What it proves | | --- | --- | --- | | Unit/component | Existing `test_*.py` files | Individual behavior; dependencies may be mocked | -| Dedicated-workspace CUJs | `e2e_cuj/test_*.py` | Real routing, MPS, MCP, and two-agent inference with read-only workspace configuration | +| Dedicated-workspace CUJ | `e2e_cuj/test_cuj_*.py` | Real user journeys against read-only, preconfigured workspaces | | Existing e2e | `test_e2e*.py` | Real workspace behavior with some patched setup/internal calls | | Integration CUJs | `integration/test_*.py` | Public configure, TUI, script, command, protocol, and lifecycle journeys | | Installation | `integration/test_installation.py` | Fresh installed package, CLI, and local helpers without credentials on Linux and advisory native Windows | -Dedicated CUJs reuse `integration/utils` session/terminal mechanics, file-task, and transcript -readers, not its config stubs or pytest fixtures. Prompt/model correlation remains -CUJ-specific. Concurrent runs may read the same CUJ workspace. Workspace operations reuse -the base class's authenticated Databricks SDK client; offline tests require GET-only API -calls and verify config changes fail without repair. The live fixture compares configuration -before and after every journey, even on failure. -Routing-disabled coverage is deferred until a separately preconfigured workspace is assigned; -these CUJs never republish configuration or create remote reservations. +Dedicated CUJs reuse `integration/utils` session/terminal mechanics and file-task +and transcript readers, not its config stubs or pytest fixtures. Prompt/model +correlation remains CUJ-specific. Concurrent runs may read the same CUJ workspace. +Workspace operations reuse the base class's authenticated Databricks SDK client; +offline tests require GET-only API calls and verify config changes fail without repair. +The live fixture compares configuration before and after the journey, even on failure. +CUJs never republish configuration or create a remote reservation. CUJ helper tests also verify that unsupported agent names fail rather than defaulting to Codex. They cover Claude/Codex helper dispatch and rejection of routing decisions without the agent-specific prompt-submission evidence. diff --git a/tests/e2e_cuj/AGENTS.md b/tests/e2e_cuj/AGENTS.md index bd136dba6..6961d9093 100644 --- a/tests/e2e_cuj/AGENTS.md +++ b/tests/e2e_cuj/AGENTS.md @@ -46,11 +46,6 @@ 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. -`test_cuj_smart_routing.py` runs four fresh interactive sessions: routed and explicit -model, for both Claude and Codex. It uses the published models/defaults and asserts -configuration is unchanged, including during teardown after failures. Routing-disabled -coverage requires a separately preconfigured workspace and is not part of this CUJ. - `test_cuj2_mps_mcp.py` runs three independently collected cases against a permanently preconfigured workspace: configuration, Codex inference, and Claude inference. Each case reads and validates the exact two-agent MPS/MCP shape. No case creates, updates, deletes, diff --git a/tests/e2e_cuj/helpers/evidence.py b/tests/e2e_cuj/helpers/evidence.py index 09deb2fda..7c955d012 100644 --- a/tests/e2e_cuj/helpers/evidence.py +++ b/tests/e2e_cuj/helpers/evidence.py @@ -10,10 +10,6 @@ from .constants import CLAUDE, CODEX -_EXPLICIT_MODEL_ALIASES = { - "anthropic.claude-haiku-4-5-20251001-v1:0": "system.ai.claude-haiku-4-5", -} - def canonical_model(value): """Only documented gateway aliases, not fuzzy family/date matching.""" @@ -21,7 +17,6 @@ def canonical_model(value): # `anthropic-aigw-1234abcd-system.ai.glm-5-2`. value = re.sub(r"^anthropic-aigw-[0-9a-f]{8}-", "", value) value = value.removesuffix("[1m]") - value = _EXPLICIT_MODEL_ALIASES.get(value, value) value = value.removeprefix("system.ai.") # Codex writes dotted major/minor versions (`gpt-5.6-sol`) while the # gateway model identifier uses a hyphen (`system.ai.gpt-5-6-sol`). diff --git a/tests/integration/README.md b/tests/integration/README.md index 67448660e..3fb8c5457 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -3,11 +3,11 @@ The separate [dedicated-workspace CUJ](../e2e_cuj/AGENTS.md) lives in `tests/e2e_cuj/` and runs directly with pytest. It does not use this suite's runner or config fixtures. It reuses the session, terminal, file-task, and transcript helpers in `utils/`. -CUJ-specific evidence stays separate from shared mechanics. Workspace config/catalog reads use -its base class's Databricks SDK client. Configuration is read-only and checked for changes at -teardown; concurrent readers need no reservation. One journey covers four routing and -explicit-model sessions. Another has three separately collected cases for exact MPS/MCP -configuration, Codex inference, and Claude inference. Routing-disabled behavior remains uncovered. +Its Claude/Codex evidence helpers keep scenario-specific assertions separate from shared mechanics. +Workspace config/catalog reads use its base class's Databricks SDK client. Configuration +is read-only and checked for changes at teardown; concurrent readers need no reservation. +CUJ2 adds three separately collected cases for exact MPS/MCP configuration, Codex inference, +and Claude inference. This suite runs the **installed product** through subprocesses, against the same `UCODE_TEST_WORKSPACE` used by the existing e2e tests. It does not import `ucode`, diff --git a/tests/integration/test_ug_smart_routing_hooks.py b/tests/integration/test_ug_smart_routing_hooks.py index 58ddcffe3..67f55076c 100644 --- a/tests/integration/test_ug_smart_routing_hooks.py +++ b/tests/integration/test_ug_smart_routing_hooks.py @@ -131,10 +131,7 @@ def _toggle_with_skill(tui, session, agent: str, enabled: bool) -> None: lambda screen: ( "Smart Router" in screen and confirmation in screen - and ( - assistant_answer_contains(session, agent, confirmation) - or assistant_answer_contains(session, agent, f"**{state}** for this session") - ) + and assistant_answer_contains(session, agent, confirmation) ), f"the installed Smart Router skill to turn routing {state}", timeout=120, diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index b7dc01ecc..148543d1e 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -2482,90 +2482,6 @@ def test_launch_model_is_only_set_for_current_process(self, monkeypatch): settings = json.loads(calls[0][2]) assert settings["env"]["ANTHROPIC_MODEL"] == "cat.schema.model" - def test_launch_custom_model_uses_family_aliases_and_native_selector( - self, monkeypatch, tmp_path - ): - calls: list[list[str]] = [] - monkeypatch.setenv("ANTHROPIC_MODEL", "stale-model") - monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "token") - monkeypatch.setattr(claude, "exec_or_spawn", lambda argv: calls.append(argv)) - settings_path = tmp_path / "ucode-settings.json" - saved_model = "system.ai.claude-haiku-4-5" - settings_path.write_text( - json.dumps({"model": saved_model, "env": {"USER_SETTING": "keep"}}) - ) - monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", settings_path) - original_settings = settings_path.read_bytes() - - custom_model = "system.ai.claude-opus-4-8" - claude.launch( - {"workspace": WS, "_claude_launch_custom_model": custom_model}, - ["--debug"], - options=LaunchOptions(user_pinned_model=custom_model), - ) - - assert os.environ["ANTHROPIC_MODEL"] == claude.CLAUDE_CUSTOM_MODEL_SELECTOR - assert calls[0][0:2] == ["claude", "--settings"] - settings = json.loads(calls[0][2]) - assert settings["model"] == saved_model - for family in claude.CLAUDE_CUSTOM_MODEL_FAMILIES: - key = claude.CLAUDE_DEFAULT_MODEL_ENV_KEYS[family] - assert settings["env"][key] == custom_model - assert settings["env"]["USER_SETTING"] == "keep" - assert "ANTHROPIC_DEFAULT_FABLE_MODEL" not in settings["env"] - assert calls[0][3:] == ["--model", claude.CLAUDE_CUSTOM_MODEL_SELECTOR, "--debug"] - assert custom_model not in calls[0][3:] - assert settings_path.read_bytes() == original_settings - - def test_launch_custom_model_preserves_forwarded_native_model(self, monkeypatch): - calls: list[list[str]] = [] - monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "token") - monkeypatch.setattr(claude, "exec_or_spawn", lambda argv: calls.append(argv)) - - claude.launch( - {"workspace": WS, "_claude_launch_custom_model": "system.ai.claude-opus-4-8"}, - ["--model", "claude-sonnet-5"], - options=LaunchOptions(user_pinned_model="system.ai.claude-opus-4-8"), - ) - - assert calls[0][3:] == ["--model", "claude-sonnet-5"] - - def test_launch_custom_model_uses_native_selector_over_stale_settings( - self, monkeypatch, tmp_path - ): - calls: list[list[str]] = [] - settings_path = tmp_path / "settings.json" - saved_model = "system.ai.claude-haiku-4-5" - settings_path.write_text(json.dumps({"model": saved_model})) - original_settings = settings_path.read_bytes() - monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", settings_path) - monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "token") - monkeypatch.setattr(claude, "exec_or_spawn", lambda argv: calls.append(argv)) - - custom_model = "system.ai.claude-opus-4-8" - claude.launch( - { - "workspace": WS, - "_claude_launch_custom_model": custom_model, - "claude_static_models": [ - "system.ai.claude-opus-4-8", - "system.ai.claude-haiku-4-5", - ], - }, - [], - options=LaunchOptions(user_pinned_model=custom_model), - ) - - assert calls[0][0:2] == ["claude", "--settings"] - settings = json.loads(calls[0][2]) - assert settings["model"] == saved_model - for family in claude.CLAUDE_CUSTOM_MODEL_FAMILIES: - key = claude.CLAUDE_DEFAULT_MODEL_ENV_KEYS[family] - assert settings["env"][key] == custom_model - assert calls[0][3:] == ["--model", claude.CLAUDE_CUSTOM_MODEL_SELECTOR] - assert custom_model not in calls[0][3:] - assert settings_path.read_bytes() == original_settings - def test_launch_default_model_is_inherited_by_smart_routing(self, monkeypatch): monkeypatch.delenv("ANTHROPIC_DEFAULT_MODEL", raising=False) monkeypatch.setattr(v2, "launch_claude", Mock()) diff --git a/tests/test_cli.py b/tests/test_cli.py index c654cdcab..1c44bed56 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -813,11 +813,6 @@ def test_claude_model_location_preserves_explicit_model(self): assert calls["launch"].call_args.kwargs["options"].user_pinned_model == ( "main.default.claude-opus-5" ) - assert "_claude_launch_default_model" not in calls["launch"].call_args.args[1] - assert ( - calls["launch"].call_args.args[1]["_claude_launch_custom_model"] - == "main.default.claude-opus-5" - ) assert calls["launch"].call_args.args[2] == [] @pytest.mark.parametrize( @@ -1509,34 +1504,13 @@ def test_model_is_launch_scoped_for_claude(self, monkeypatch): ): result = runner.invoke(app, ["claude", "--model", "cat.schema.claude-opus-5"]) assert result.exit_code == 0, result.output - # The model is passed through invocation-scoped launch state/options, not persisted in - # settings. + # The model is passed through invocation-scoped LaunchOptions, not persisted in settings. assert mock_configure.call_args.kwargs["custom_model"] is None assert mock_configure.call_args.kwargs["route_root_model"] is None - assert ( - mock_launch.call_args.args[1]["_claude_launch_custom_model"] - == "cat.schema.claude-opus-5" - ) assert ( mock_launch.call_args.kwargs["options"].user_pinned_model == "cat.schema.claude-opus-5" ) - def test_explicit_model_clears_managed_default_route_root(self): - managed = { - "enabled_agents": { - "claude": {"model_config": {"default_model": "system.ai.claude-sonnet-5"}} - } - } - with _launch_policy_patches(managed) as calls: - result = runner.invoke(app, ["claude", "--model", "system.ai.claude-opus-4-8"]) - - assert result.exit_code == 0, result.output - assert calls["configure"].call_args.kwargs["route_root_model"] is None - assert ( - calls["launch"].call_args.args[1]["_claude_launch_custom_model"] - == "system.ai.claude-opus-4-8" - ) - def test_v2_model_sets_transient_launch_override(self, monkeypatch): monkeypatch.setenv("ENABLE_SMART_ROUTING_V2", "1") state = {**MINIMAL_STATE, "claude_models": {"opus": "system.ai.claude-opus-4-8"}} diff --git a/tests/test_cuj_evidence.py b/tests/test_cuj_evidence.py index 76d857c14..fa50ef2ef 100644 --- a/tests/test_cuj_evidence.py +++ b/tests/test_cuj_evidence.py @@ -163,14 +163,6 @@ def test_cuj_evidence_ignores_existing_session(tmp_path): ("anthropic-aigw-1234abcd-system.ai.glm-5-2[1m]", "system.ai.glm-5-2"), ("gpt-5.6-sol", "system.ai.gpt-5-6-sol"), ("system.ai.gpt-6-sol", "system.ai.gpt-6-sol"), - ( - "anthropic.claude-haiku-4-5-20251001-v1:0", - "system.ai.claude-haiku-4-5", - ), - ( - "anthropic.claude-haiku-4-5-20251002-v1:0", - "system.ai.anthropic.claude-haiku-4-5-20251002-v1:0", - ), ( "anthropic-aigw-invalid-claude-opus-4-8", "system.ai.anthropic-aigw-invalid-claude-opus-4-8", From 7ba473a63094549c6f5ad928fba3ed8d95612c54 Mon Sep 17 00:00:00 2001 From: Mason Cao Date: Tue, 6 Oct 2026 14:01:24 +0000 Subject: [PATCH 18/19] [AIGTWY-4877] Accept formatted Smart Router confirmation --- tests/integration/test_ug_smart_routing_hooks.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/integration/test_ug_smart_routing_hooks.py b/tests/integration/test_ug_smart_routing_hooks.py index 67f55076c..58ddcffe3 100644 --- a/tests/integration/test_ug_smart_routing_hooks.py +++ b/tests/integration/test_ug_smart_routing_hooks.py @@ -131,7 +131,10 @@ def _toggle_with_skill(tui, session, agent: str, enabled: bool) -> None: lambda screen: ( "Smart Router" in screen and confirmation in screen - and assistant_answer_contains(session, agent, confirmation) + and ( + assistant_answer_contains(session, agent, confirmation) + or assistant_answer_contains(session, agent, f"**{state}** for this session") + ) ), f"the installed Smart Router skill to turn routing {state}", timeout=120, From f187532969232941fca0a845bae8dcba54eafda6 Mon Sep 17 00:00:00 2001 From: Mason Cao Date: Tue, 6 Oct 2026 17:46:25 +0000 Subject: [PATCH 19/19] [AIGTWY-4877] Clarify CUJ expected config helper --- tests/e2e_cuj/test_cuj2_mps_mcp.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/e2e_cuj/test_cuj2_mps_mcp.py b/tests/e2e_cuj/test_cuj2_mps_mcp.py index 2ee4500aa..0af6311df 100644 --- a/tests/e2e_cuj/test_cuj2_mps_mcp.py +++ b/tests/e2e_cuj/test_cuj2_mps_mcp.py @@ -32,7 +32,7 @@ WORKSPACE_URL = "https://dbc-0dcf95cf-e357.cloud.databricks.com" -def _agent_config(agent: str) -> dict: +def expected_agent_config(agent: str) -> dict: mps_name, model = MODEL_PROVIDER_SERVICE_FIXTURES[agent] return { "agent": CODING_AGENT_BY_CLI_NAME[agent], @@ -51,7 +51,7 @@ def _assert_cuj2_config(config: dict) -> None: expected = { "spec_version": 1, "default_agent": CODING_AGENT_BY_CLI_NAME[CODEX], - "enabled_agents": [_agent_config(CLAUDE), _agent_config(CODEX)], + "enabled_agents": [expected_agent_config(CLAUDE), expected_agent_config(CODEX)], "mcp_servers": {"names": [SANDBOX_MCP_SERVICE_NAME]}, } assert actual == expected, config