From 23c89c3d465ebb43706e8f5b86003fb6e51aeff7 Mon Sep 17 00:00:00 2001 From: John Speaks Date: Mon, 5 Oct 2026 16:40:42 +0000 Subject: [PATCH 1/9] Cover managed two-agent CUJ on dedicated workspace --- pyproject.toml | 2 + tests/e2e_cuj/test_ug_managed_sentinel.py | 501 ++++++++++++++++++++++ tests/integration/utils/terminal.py | 7 + uv.lock | 37 ++ 4 files changed, 547 insertions(+) create mode 100644 tests/e2e_cuj/test_ug_managed_sentinel.py diff --git a/pyproject.toml b/pyproject.toml index 39fbd9cf9..2feaf034b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -58,6 +58,8 @@ fallback-version = "0.0.0" [dependency-groups] dev = [ "pytest>=9.0.3", + "pexpect>=4.9.0", + "pyte>=0.8.2", # Pinned so local and CI lint with the identical formatter/type-checker; # ruff/ty output shifts between releases, so bump these deliberately. "ruff==0.16.5", diff --git a/tests/e2e_cuj/test_ug_managed_sentinel.py b/tests/e2e_cuj/test_ug_managed_sentinel.py new file mode 100644 index 000000000..7e5eeca92 --- /dev/null +++ b/tests/e2e_cuj/test_ug_managed_sentinel.py @@ -0,0 +1,501 @@ +"""End-to-end sentinel for the published managed Claude and Codex configuration. + +The selected managed workspace publishes one static configuration for both agents. This journey +fetches that configuration through one public ``ug configure`` command, checks the generated files, +then runs real tasks on the bare managed default and explicit Claude/Codex selections. +""" + +from __future__ import annotations + +import json +import os +import re +import time +import tomllib +import uuid + +import pytest +from base import BaseCujTest +from utils.evidence import FileTask +from utils.model_discovery import claude_model_in_picker +from utils.sql import query_count, resolve_trace_table, resolve_warehouse_id +from utils.terminal import AgentTerminal + +CLAUDE_MODELS = [ + "system.ai.claude-opus-4-8", + "system.ai.claude-sonnet-4-6", + "system.ai.claude-haiku-4-5", +] +CLAUDE_DEFAULTS = { + "default_model": "system.ai.claude-sonnet-4-6", + "default_opus_model": "system.ai.claude-opus-4-8", + "default_sonnet_model": "system.ai.claude-sonnet-4-6", + "default_haiku_model": "system.ai.claude-haiku-4-5", +} +CODEX_MODELS = ["system.ai.gpt-5-6-sol", "system.ai.gpt-5-6-luna"] +CODEX_DEFAULT = CODEX_MODELS[0] +CODEX_LUNA = CODEX_MODELS[1] +EXPECTED_HEADER_NAMES = {"x-ug-e2e-run", "x-ug-e2e-agent"} +CLAUDE_PICKER_LABELS = ("Claude Opus 4.8", "Claude Sonnet 4.6", "Claude Haiku 4.5") +TRACE_WAIT_SECONDS = 360 +TRACE_POLL_SECONDS = 10 + + +def _claude_models_visible(text: str) -> bool: + return all( + claude_model_in_picker(text, model, label) + for model, label in zip(CLAUDE_MODELS, CLAUDE_PICKER_LABELS, strict=True) + ) + + +def _claude_custom_model_ids(screen: str) -> list[str]: + return re.findall( + r"(?m)^\s*(?:[❯›>]\s*)?\d+\.\s+.*Custom model \(([^)]+)\)", + screen, + ) + + +def _managed_config(session) -> dict: + cache_path = session.home / ".ucode" / "managed-config.json" + cache = json.loads(cache_path.read_text()) + raw = cache.get("config") + assert isinstance(raw, dict), cache + return raw + + +def _managed_entries(raw: dict) -> dict[str, dict]: + entries = raw.get("enabled_agents") + assert isinstance(entries, list), raw + assert len(entries) == 2, raw + result = { + entry.get("agent"): entry.get("config") + for entry in entries + if isinstance(entry, dict) and isinstance(entry.get("config"), dict) + } + assert set(result) == {"CODING_AGENT_CLAUDE_CODE", "CODING_AGENT_CODEX"}, raw + return result + + +def _normalized_headers(headers: object) -> dict[str, str]: + assert isinstance(headers, dict), headers + assert all(isinstance(name, str) for name in headers), headers + result = {name.casefold(): value for name, value in headers.items()} + assert set(result) == EXPECTED_HEADER_NAMES, headers + assert all(isinstance(value, str) and value for value in result.values()), headers + return result + + +def _managed_header_subset(headers: object) -> dict[str, str]: + assert isinstance(headers, dict), headers + assert all(isinstance(name, str) for name in headers), headers + normalized = {name.casefold(): value for name, value in headers.items()} + assert EXPECTED_HEADER_NAMES <= set(normalized), headers + assert all( + isinstance(normalized[name], str) and normalized[name] for name in EXPECTED_HEADER_NAMES + ) + return {name: normalized[name] for name in EXPECTED_HEADER_NAMES} + + +def _claude_header_lines(value: object) -> dict[str, str]: + assert isinstance(value, str), value + headers = {} + for line in value.splitlines(): + name, separator, header_value = line.partition(":") + if separator: + headers[name.strip().casefold()] = header_value.strip() + return _managed_header_subset(headers) + + +def _assert_published_config(raw: dict, entries: dict[str, dict]) -> dict[str, str]: + assert raw["default_agent"] == "CODING_AGENT_CLAUDE_CODE", raw + claude = entries["CODING_AGENT_CLAUDE_CODE"] + codex = entries["CODING_AGENT_CODEX"] + + assert claude["smart_routing"] == {"enabled": False}, claude + assert claude["models"] == {"model_services": CLAUDE_MODELS}, claude + assert claude["default_models"] == CLAUDE_DEFAULTS, claude + assert claude["tracing"] == {"enabled": True}, claude + claude_headers = _normalized_headers(claude["http_headers"]) + assert claude_headers["x-ug-e2e-agent"] == "claude", claude_headers + + assert codex["smart_routing"] == {"enabled": False}, codex + assert codex["models"] == {"model_services": CODEX_MODELS}, codex + assert codex["default_models"] == {"default_model": CODEX_DEFAULT}, codex + assert codex["tracing"] == {"enabled": True}, codex + codex_headers = _normalized_headers(codex["http_headers"]) + assert codex_headers["x-ug-e2e-agent"] == "codex", codex_headers + assert codex_headers["x-ug-e2e-run"] == claude_headers["x-ug-e2e-run"], ( + claude_headers, + codex_headers, + ) + return claude_headers + + +def _assert_generated_configs(session, claude_headers: dict[str, str]) -> None: + claude_settings_path = session.home / ".claude" / "ucode-settings.json" + claude_settings = json.loads(claude_settings_path.read_text()) + assert claude_settings["availableModels"] == CLAUDE_MODELS, claude_settings + claude_picker = claude_settings["modelPicker"] + assert claude_picker["replaceBuiltInOptions"] is True, claude_settings + assert claude_picker["options"] == [ + {"model": model, "label": label} + for model, label in zip( + CLAUDE_MODELS, + ("Claude Opus 4.8", "Claude Sonnet 4.6", "Claude Haiku 4.5"), + strict=True, + ) + ], claude_settings + claude_env = claude_settings["env"] + assert claude_env["ANTHROPIC_DEFAULT_OPUS_MODEL"] == CLAUDE_MODELS[0] + "[1m]", claude_settings + assert claude_env["ANTHROPIC_DEFAULT_SONNET_MODEL"] == (CLAUDE_MODELS[1] + "[1m]"), ( + claude_settings + ) + assert claude_env["ANTHROPIC_DEFAULT_HAIKU_MODEL"] == CLAUDE_MODELS[2], claude_settings + assert _claude_header_lines(claude_env["ANTHROPIC_CUSTOM_HEADERS"]) == claude_headers, ( + claude_settings + ) + assert claude_env["CLAUDE_CODE_ENABLE_TELEMETRY"] == "1", claude_settings + assert claude_env["OTEL_TRACES_EXPORTER"] == "otlp", claude_settings + assert claude_env["OTEL_EXPORTER_OTLP_TRACES_ENDPOINT"].endswith( + "/ai-gateway/otel/v1/traces" + ), claude_settings + + codex_profile_path = session.home / ".codex" / "ucode.config.toml" + codex_profile = tomllib.loads(codex_profile_path.read_text()) + assert codex_profile["model"] == CODEX_DEFAULT, codex_profile + codex_provider = codex_profile["model_providers"]["Databricks"] + assert _managed_header_subset(codex_provider["http_headers"]) == { + "x-ug-e2e-run": claude_headers["x-ug-e2e-run"], + "x-ug-e2e-agent": "codex", + }, codex_profile + codex_catalog_path = session.home / ".ucode" / "codex-model-catalog.json" + codex_catalog = json.loads(codex_catalog_path.read_text()) + assert [ + model["slug"] for model in codex_catalog["models"] if model.get("visibility") == "list" + ] == CODEX_MODELS, codex_catalog + codex_app_config = tomllib.loads((session.home / ".codex" / "config.toml").read_text()) + assert codex_app_config["model_catalog_json"] == str(codex_catalog_path), codex_app_config + + +def _trace_pair_count( + workspace: str, + bearer: str, + warehouse_id: str, + table: str, + pairs: list[tuple[str, str, str]], +) -> int: + """Count distinct marker/model pairs in both agents' trace attribute layouts.""" + clauses = [] + parameters = [] + for index, (marker, model, _marker_source) in enumerate(pairs): + clauses.append( + "(marker_source = :marker_source_" + f"{index} AND marker = :marker_{index} AND model = :model_{index})" + ) + parameters.extend( + [ + { + "name": f"marker_source_{index}", + "value": _marker_source, + "type": "STRING", + }, + {"name": f"marker_{index}", "value": marker, "type": "STRING"}, + {"name": f"model_{index}", "value": model, "type": "STRING"}, + ] + ) + + marker_expressions = { + "resource": "variant_get(resource.attributes, '$[\"ug_integration_marker\"]', 'STRING')", + "span": "variant_get(attributes, '$[\"ug_integration_marker\"]', 'STRING')", + } + span_names = { + "resource": "claude_code.llm_request", + "span": "model_client.stream_responses_api", + } + unions = [] + for source in ("resource", "span"): + marker_expression = marker_expressions[source] + model_key = "gen_ai.request.model" if source == "resource" else "model" + unions.append( + "SELECT " + f"'{source}' AS marker_source, " + f"{marker_expression} AS marker, " + f"variant_get(attributes, '$[\"{model_key}\"]', 'STRING') AS model " + f"FROM {table} WHERE kind = 'SPAN_KIND_INTERNAL' " + f"AND name = '{span_names[source]}' " + "AND time > current_timestamp() - INTERVAL 90 MINUTES" + ) + statement = ( + "SELECT COUNT(*) FROM (SELECT marker_source, marker, model FROM (" + + " UNION ALL ".join(unions) + + ") AS candidate_rows WHERE " + + " OR ".join(clauses) + + " GROUP BY marker_source, marker, model) AS observed" + ) + return query_count(workspace, bearer, warehouse_id, statement, parameters) + + +def _native_gateway_model(model: str) -> str: + """Return the model name emitted by the gateway's native client span.""" + return re.sub(r"\[(?:1m|200k)\]$", "", model).removeprefix("system.ai.") + + +def _trace_client_pair_count( + workspace: str, + bearer: str, + warehouse_id: str, + table: str, + pairs: list[tuple[str, str, str]], +) -> int: + """Count expected native gateway client models for the task markers.""" + input_messages = "CAST(variant_get(attributes, '$[\"gen_ai.input.messages\"]') AS STRING)" + model = "variant_get(attributes, '$[\"gen_ai.request.model\"]', 'STRING')" + expected_cases = [] + parameters = [] + for index, (marker, expected_model, _marker_source) in enumerate(pairs): + expected_cases.append( + f"WHEN input_messages LIKE concat('%', :client_marker_{index}, '%') " + f"AND model = :client_model_{index} THEN :client_marker_{index}" + ) + parameters.extend( + [ + { + "name": f"client_marker_{index}", + "value": marker, + "type": "STRING", + }, + { + "name": f"client_model_{index}", + "value": _native_gateway_model(expected_model), + "type": "STRING", + }, + ] + ) + statement = ( + "WITH candidates AS (" + "SELECT " + f"{input_messages} AS input_messages, " + f"{model} AS model " + f"FROM {table} " + "WHERE kind = 'SPAN_KIND_CLIENT' " + "AND time > current_timestamp() - INTERVAL 90 MINUTES) " + "SELECT COUNT(DISTINCT CASE " + " ".join(expected_cases) + " END) FROM candidates" + ) + return query_count(workspace, bearer, warehouse_id, statement, parameters) + + +def _wait_for_trace_pairs( + workspace: str, + bearer: str, + warehouse_id: str, + table: str, + pairs: list[tuple[str, str, str]], + session, +) -> None: + deadline = time.monotonic() + TRACE_WAIT_SECONDS + observed = 0 + observed_client = 0 + while True: + observed = _trace_pair_count(workspace, bearer, warehouse_id, table, pairs) + observed_client = _trace_client_pair_count(workspace, bearer, warehouse_id, table, pairs) + if observed >= len(pairs) and observed_client >= len(pairs): + session.record( + "trace-query.json", + { + "table": table, + "expected_pairs": [ + {"marker": marker, "model": model, "marker_source": source} + for marker, model, source in pairs + ], + "observed_pairs": observed, + "observed_native_client_pairs": observed_client, + "native_client_model_scope": "expected model per marker; auxiliary models allowed", + }, + ) + return + remaining = deadline - time.monotonic() + if remaining <= 0: + pytest.fail( + "Managed sentinel trace evidence timed out: " + f"observed {observed}/{len(pairs)} internal and " + f"{observed_client}/{len(pairs)} native client marker/model pairs in {table}. " + "The generated headers below are local writer evidence; the public trace table " + "does not expose arbitrary HTTP request headers, so this failure does not claim " + "that X-UG-E2E-* headers reached the gateway. Native client spans may also " + "contain auxiliary model calls for the same task." + ) + time.sleep(min(TRACE_POLL_SECONDS, remaining)) + + +class TestCujManagedConfig(BaseCujTest): + WORKSPACE_URL = "https://dbc-135c115c-c255.cloud.databricks.com" + + def test_cuj_managed_config_sentinel(self, live_session): + """Scenario: configure one published two-agent policy and run both agents. + + Expected: no agent selector; configured models, defaults, tracing and headers in generated + files; six completed tasks and matching agent inference plus native gateway client spans. + Auxiliary client calls for a task may use another model. Gateway receipt of headers is not + asserted here. + """ + session = live_session + workspace = self.WORKSPACE_URL + configured_workspace = os.environ.get("UG_CUJ1_WORKSPACE", "").strip().rstrip("/") + if configured_workspace: + assert configured_workspace == workspace, configured_workspace + configured = session.run( + "configure", + "--workspace", + workspace, + "--skip-upgrade", + timeout=300, + ) + assert "Select coding agents to configure:" not in configured.stdout, configured.stdout + + raw = _managed_config(session) + entries = _managed_entries(raw) + claude_headers = _assert_published_config(raw, entries) + _assert_generated_configs(session, claude_headers) + codex_models = session.codex_model_ids( + ["app-server", "--listen", "stdio://"], name="managed-sentinel-codex-models" + ) + assert codex_models == CODEX_MODELS, codex_models + session.record( + "managed-sentinel-wire.json", + { + "default_agent": raw["default_agent"], + "agents": { + agent: { + "models": config["models"], + "default_models": config["default_models"], + "smart_routing": config["smart_routing"], + "tracing": config["tracing"], + "http_headers": sorted(config["http_headers"]), + } + for agent, config in entries.items() + }, + }, + ) + + bearer = session.env["DATABRICKS_BEARER"] + trace_table = resolve_trace_table(workspace, bearer) + warehouse_id = os.environ.get("UG_INTEGRATION_WAREHOUSE_ID", "").strip() + warehouse_id = warehouse_id or resolve_warehouse_id(workspace, bearer) + run_id = uuid.uuid4().hex + trace_pairs: list[tuple[str, str, str]] = [] + + root_task = FileTask(session) + root_marker = f"ug-managed-sentinel-{run_id}-root" + trace_pairs.append((root_marker, CLAUDE_DEFAULTS["default_model"], "resource")) + session.env["OTEL_RESOURCE_ATTRIBUTES"] = f"ug_integration_marker={root_marker}" + with AgentTerminal( + session, "claude", [str(session.binary)], "managed-sentinel-root" + ) as tui: + tui.boot() + tui.submit(f"{root_task.prompt} Trace marker: {root_marker}") + tui.wait_for_task(root_task) + picker_screen = tui.open_model_picker(model_visible=_claude_models_visible) + assert _claude_custom_model_ids(picker_screen) == CLAUDE_MODELS, picker_screen + tui.exit_normally() + root_task.assert_completed(session, "claude") + + claude_opus_task = FileTask(session) + claude_opus_marker = f"ug-managed-sentinel-{run_id}-claude-opus" + trace_pairs.append( + (claude_opus_marker, CLAUDE_DEFAULTS["default_opus_model"] + "[1m]", "resource") + ) + session.env["OTEL_RESOURCE_ATTRIBUTES"] = f"ug_integration_marker={claude_opus_marker}" + with AgentTerminal( + session, + "claude", + [str(session.binary), "claude", "--model", "opus"], + "managed-sentinel-claude-opus", + ) as tui: + tui.boot() + tui.submit(f"{claude_opus_task.prompt} Trace marker: {claude_opus_marker}") + tui.wait_for_task(claude_opus_task) + tui.exit_normally() + claude_opus_task.assert_completed(session, "claude") + + claude_sonnet_task = FileTask(session) + claude_sonnet_marker = f"ug-managed-sentinel-{run_id}-claude-sonnet" + trace_pairs.append( + (claude_sonnet_marker, CLAUDE_DEFAULTS["default_sonnet_model"] + "[1m]", "resource") + ) + session.env["OTEL_RESOURCE_ATTRIBUTES"] = f"ug_integration_marker={claude_sonnet_marker}" + with AgentTerminal( + session, + "claude", + [str(session.binary), "claude", "--model", "sonnet"], + "managed-sentinel-claude-sonnet", + ) as tui: + tui.boot() + tui.submit(f"{claude_sonnet_task.prompt} Trace marker: {claude_sonnet_marker}") + tui.wait_for_task(claude_sonnet_task) + tui.exit_normally() + claude_sonnet_task.assert_completed(session, "claude") + + claude_haiku_task = FileTask(session) + claude_haiku_marker = f"ug-managed-sentinel-{run_id}-claude-haiku" + trace_pairs.append( + (claude_haiku_marker, CLAUDE_DEFAULTS["default_haiku_model"], "resource") + ) + session.env["OTEL_RESOURCE_ATTRIBUTES"] = f"ug_integration_marker={claude_haiku_marker}" + with AgentTerminal( + session, + "claude", + [str(session.binary), "claude", "--model", "haiku"], + "managed-sentinel-claude-haiku", + ) as tui: + tui.boot() + tui.submit(f"{claude_haiku_task.prompt} Trace marker: {claude_haiku_marker}") + tui.wait_for_task(claude_haiku_task) + tui.exit_normally() + claude_haiku_task.assert_completed(session, "claude") + + codex_default_task = FileTask(session) + codex_default_marker = f"ug-managed-sentinel-{run_id}-codex-default" + trace_pairs.append((codex_default_marker, CODEX_DEFAULT, "span")) + session.env.pop("OTEL_RESOURCE_ATTRIBUTES", None) + with AgentTerminal( + session, + "codex", + [ + str(session.binary), + "codex", + "--", + "--config", + f'otel.span_attributes.ug_integration_marker="{codex_default_marker}"', + ], + "managed-sentinel-codex-default", + ) as tui: + tui.boot() + tui.submit(f"{codex_default_task.prompt} Trace marker: {codex_default_marker}") + tui.wait_for_task(codex_default_task) + tui.exit_normally() + codex_default_task.assert_completed(session, "codex") + + codex_luna_task = FileTask(session) + codex_luna_marker = f"ug-managed-sentinel-{run_id}-codex-luna" + trace_pairs.append((codex_luna_marker, CODEX_LUNA, "span")) + with AgentTerminal( + session, + "codex", + [ + str(session.binary), + "codex", + "--", + "--model", + CODEX_LUNA, + "--config", + f'otel.span_attributes.ug_integration_marker="{codex_luna_marker}"', + ], + "managed-sentinel-codex-luna", + ) as tui: + tui.boot() + tui.submit(f"{codex_luna_task.prompt} Trace marker: {codex_luna_marker}") + tui.wait_for_task(codex_luna_task) + tui.exit_normally() + codex_luna_task.assert_completed(session, "codex") + + _wait_for_trace_pairs(workspace, bearer, warehouse_id, trace_table, trace_pairs, session) diff --git a/tests/integration/utils/terminal.py b/tests/integration/utils/terminal.py index 9f4c4fda7..e6a544474 100644 --- a/tests/integration/utils/terminal.py +++ b/tests/integration/utils/terminal.py @@ -232,6 +232,13 @@ def boot(self, timeout=120): # These are interactions with ordinary UI choices, not pre-written # onboarding state. Only the test's disposable project is trusted. dialogs = [ + ( + "external-imports", + "Allow external CLAUDE.md file imports?" in text + and "No, disable external imports" in text + and "Yes, allow external imports" in text, + "\r", + ), ( "theme", "Choose the text style" in text and "Dark mode" in text, diff --git a/uv.lock b/uv.lock index b63869f81..a2e745a30 100644 --- a/uv.lock +++ b/uv.lock @@ -569,6 +569,18 @@ wheels = [ { url = "https://pypi-proxy.cloud.databricks.com/packages/63/34/ba1c580383c9eada3711951fef0795c80b829a078d72188184bcab9dd527/packaging-26.3-py3-none-any.whl", hash = "sha256:d7193f7c8e4e93f444fde0262bf90af30e16fa0ad0ad44cb553c87339b23cd1c", size = 129956, upload-time = "2026-08-04T18:15:27.159Z" }, ] +[[package]] +name = "pexpect" +version = "4.9.0" +source = { registry = "https://pypi-proxy.cloud.databricks.com/simple/" } +dependencies = [ + { name = "ptyprocess" }, +] +sdist = { url = "https://pypi-proxy.cloud.databricks.com/packages/42/92/cc564bf6381ff43ce1f4d06852fc19a2f11d180f23dc32d9588bee2f149d/pexpect-4.9.0.tar.gz", hash = "sha256:ee7d41123f3c9911050ea2c2dac107568dc43b2d3b0c7557a33212c398ead30f", size = 166450, upload-time = "2023-11-25T09:07:26.339Z" } +wheels = [ + { url = "https://pypi-proxy.cloud.databricks.com/packages/9e/c3/059298687310d527a58bb01f3b1965787ee3b40dce76752eda8b44e9a2c5/pexpect-4.9.0-py2.py3-none-any.whl", hash = "sha256:7236d1e080e4936be2dc3e326cec0af72acf9212a7e1d060210e70a47e253523", size = 63772, upload-time = "2023-11-25T06:56:14.81Z" }, +] + [[package]] name = "pluggy" version = "1.6.0" @@ -605,6 +617,15 @@ wheels = [ { url = "https://pypi-proxy.cloud.databricks.com/packages/c4/72/02445137af02769918a93807b2b7890047c32bfb9f90371cbc12688819eb/protobuf-6.33.6-py3-none-any.whl", hash = "sha256:77179e006c476e69bf8e8ce866640091ec42e1beb80b213c3900006ecfba6901", size = 170656, upload-time = "2026-03-18T19:04:59.826Z" }, ] +[[package]] +name = "ptyprocess" +version = "0.7.0" +source = { registry = "https://pypi-proxy.cloud.databricks.com/simple/" } +sdist = { url = "https://pypi-proxy.cloud.databricks.com/packages/20/e5/16ff212c1e452235a90aeb09066144d0c5a6a8c0834397e03f5224495c4e/ptyprocess-0.7.0.tar.gz", hash = "sha256:5c5d0a3b48ceee0b48485e0c26037c0acd7d29765ca3fbb5cb3831d347423220", size = 70762, upload-time = "2020-12-28T15:15:30.155Z" } +wheels = [ + { url = "https://pypi-proxy.cloud.databricks.com/packages/22/a6/858897256d0deac81a172289110f31629fc4cee19b6f01283303e18c8db3/ptyprocess-0.7.0-py2.py3-none-any.whl", hash = "sha256:4b41f3967fce3af57cc7e94b888626c18bf37a083e3651ca8feeb66d492fef35", size = 13993, upload-time = "2020-12-28T15:15:28.35Z" }, +] + [[package]] name = "pyasn1" version = "0.6.4" @@ -748,6 +769,18 @@ crypto = [ { name = "cryptography" }, ] +[[package]] +name = "pyte" +version = "0.8.2" +source = { registry = "https://pypi-proxy.cloud.databricks.com/simple/" } +dependencies = [ + { name = "wcwidth" }, +] +sdist = { url = "https://pypi-proxy.cloud.databricks.com/packages/ab/ab/b599762933eba04de7dc5b31ae083112a6c9a9db15b01d3109ad797559d9/pyte-0.8.2.tar.gz", hash = "sha256:5af970e843fa96a97149d64e170c984721f20e52227a2f57f0a54207f08f083f", size = 92301, upload-time = "2023-11-12T09:33:43.217Z" } +wheels = [ + { url = "https://pypi-proxy.cloud.databricks.com/packages/59/d0/bb522283b90853afbf506cd5b71c650cf708829914efd0003d615cf426cd/pyte-0.8.2-py3-none-any.whl", hash = "sha256:85db42a35798a5aafa96ac4d8da78b090b2c933248819157fc0e6f78876a0135", size = 31627, upload-time = "2023-11-12T09:33:41.096Z" }, +] + [[package]] name = "pytest" version = "9.1.1" @@ -1096,6 +1129,8 @@ dependencies = [ [package.dev-dependencies] dev = [ + { name = "pexpect" }, + { name = "pyte" }, { name = "pytest" }, { name = "ruff" }, { name = "ty" }, @@ -1114,6 +1149,8 @@ requires-dist = [ [package.metadata.requires-dev] dev = [ + { name = "pexpect", specifier = ">=4.9.0" }, + { name = "pyte", specifier = ">=0.8.2" }, { name = "pytest", specifier = ">=9.0.3" }, { name = "ruff", specifier = "==0.16.5" }, { name = "ty", specifier = "==0.0.75" }, From 4a5de2ee56d3608dd806a8f280c9b47ef3423d63 Mon Sep 17 00:00:00 2001 From: John Speaks Date: Mon, 5 Oct 2026 20:46:16 +0000 Subject: [PATCH 2/9] Assert managed Codex model catalog in CUJ1 --- tests/e2e_cuj/test_ug_managed_sentinel.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/e2e_cuj/test_ug_managed_sentinel.py b/tests/e2e_cuj/test_ug_managed_sentinel.py index 7e5eeca92..8fabdb31b 100644 --- a/tests/e2e_cuj/test_ug_managed_sentinel.py +++ b/tests/e2e_cuj/test_ug_managed_sentinel.py @@ -162,13 +162,13 @@ def _assert_generated_configs(session, claude_headers: dict[str, str]) -> None: codex_profile_path = session.home / ".codex" / "ucode.config.toml" codex_profile = tomllib.loads(codex_profile_path.read_text()) - assert codex_profile["model"] == CODEX_DEFAULT, codex_profile codex_provider = codex_profile["model_providers"]["Databricks"] assert _managed_header_subset(codex_provider["http_headers"]) == { "x-ug-e2e-run": claude_headers["x-ug-e2e-run"], "x-ug-e2e-agent": "codex", }, codex_profile codex_catalog_path = session.home / ".ucode" / "codex-model-catalog.json" + assert codex_profile["model_catalog_json"] == str(codex_catalog_path), codex_profile codex_catalog = json.loads(codex_catalog_path.read_text()) assert [ model["slug"] for model in codex_catalog["models"] if model.get("visibility") == "list" From 8b4110afa01010c9d66b15b9e3469e3230489504 Mon Sep 17 00:00:00 2001 From: John Speaks Date: Mon, 5 Oct 2026 21:54:25 +0000 Subject: [PATCH 3/9] Find managed Claude models by visible picker text --- tests/e2e_cuj/test_ug_managed_sentinel.py | 16 ++-------------- tests/integration/utils/terminal.py | 18 ++++++++++-------- 2 files changed, 12 insertions(+), 22 deletions(-) diff --git a/tests/e2e_cuj/test_ug_managed_sentinel.py b/tests/e2e_cuj/test_ug_managed_sentinel.py index 8fabdb31b..4987c8469 100644 --- a/tests/e2e_cuj/test_ug_managed_sentinel.py +++ b/tests/e2e_cuj/test_ug_managed_sentinel.py @@ -17,7 +17,6 @@ import pytest from base import BaseCujTest from utils.evidence import FileTask -from utils.model_discovery import claude_model_in_picker from utils.sql import query_count, resolve_trace_table, resolve_warehouse_id from utils.terminal import AgentTerminal @@ -36,23 +35,12 @@ CODEX_DEFAULT = CODEX_MODELS[0] CODEX_LUNA = CODEX_MODELS[1] EXPECTED_HEADER_NAMES = {"x-ug-e2e-run", "x-ug-e2e-agent"} -CLAUDE_PICKER_LABELS = ("Claude Opus 4.8", "Claude Sonnet 4.6", "Claude Haiku 4.5") TRACE_WAIT_SECONDS = 360 TRACE_POLL_SECONDS = 10 def _claude_models_visible(text: str) -> bool: - return all( - claude_model_in_picker(text, model, label) - for model, label in zip(CLAUDE_MODELS, CLAUDE_PICKER_LABELS, strict=True) - ) - - -def _claude_custom_model_ids(screen: str) -> list[str]: - return re.findall( - r"(?m)^\s*(?:[❯›>]\s*)?\d+\.\s+.*Custom model \(([^)]+)\)", - screen, - ) + return all(model in text for model in CLAUDE_MODELS) def _managed_config(session) -> dict: @@ -395,7 +383,7 @@ def test_cuj_managed_config_sentinel(self, live_session): tui.submit(f"{root_task.prompt} Trace marker: {root_marker}") tui.wait_for_task(root_task) picker_screen = tui.open_model_picker(model_visible=_claude_models_visible) - assert _claude_custom_model_ids(picker_screen) == CLAUDE_MODELS, picker_screen + assert _claude_models_visible(picker_screen), picker_screen tui.exit_normally() root_task.assert_completed(session, "claude") diff --git a/tests/integration/utils/terminal.py b/tests/integration/utils/terminal.py index e6a544474..4601f55bf 100644 --- a/tests/integration/utils/terminal.py +++ b/tests/integration/utils/terminal.py @@ -308,19 +308,21 @@ def check_input_and_exit(self): def open_model_picker(self, *, model_visible=None): """Open Claude's real model picker, record it, then return to the prompt.""" self.submit("/model") - self.wait_for( - lambda text: "Select model" in text and "Switch between Claude models." in text, - "the model picker", - timeout=60, - ) - if model_visible is not None: - # Native discovery can finish after the picker shell first renders. + if model_visible is None: + self.wait_for( + lambda text: "Select model" in text and "Switch between Claude models." in text, + "the model picker", + timeout=60, + ) + else: self.wait_for(model_visible, "a discovered model in the picker", timeout=60) screen = self.visible self.actions.append({"reason": "model-picker-visible", "screen": screen}) self.send("\x1b", "close the model picker") self.wait_for( - lambda text: "Select model" not in text, + lambda text: ( + not model_visible(text) if model_visible is not None else "Select model" not in text + ), "the prompt after closing the model picker", ) return screen From a0b3f992afee8487a42c830167bf49e547d989c6 Mon Sep 17 00:00:00 2001 From: John Speaks Date: Tue, 6 Oct 2026 03:39:05 +0000 Subject: [PATCH 4/9] [AIGTWY-4876] Separate managed CUJ assertions and request evidence --- tests/e2e_cuj/test_ug_managed_sentinel.py | 532 +++++++++++++++++----- 1 file changed, 423 insertions(+), 109 deletions(-) diff --git a/tests/e2e_cuj/test_ug_managed_sentinel.py b/tests/e2e_cuj/test_ug_managed_sentinel.py index 4987c8469..5a251796e 100644 --- a/tests/e2e_cuj/test_ug_managed_sentinel.py +++ b/tests/e2e_cuj/test_ug_managed_sentinel.py @@ -1,8 +1,7 @@ -"""End-to-end sentinel for the published managed Claude and Codex configuration. +"""End-to-end journeys for the published managed Claude and Codex configuration. -The selected managed workspace publishes one static configuration for both agents. This journey -fetches that configuration through one public ``ug configure`` command, checks the generated files, -then runs real tasks on the bare managed default and explicit Claude/Codex selections. +The selected managed workspace publishes one static configuration for both agents. Separate cases +check the generated files, record Claude and Codex inference requests, and query trace evidence. """ from __future__ import annotations @@ -10,12 +9,14 @@ import json import os import re +import shutil import time import tomllib import uuid import pytest from base import BaseCujTest +from helpers.tui_request_recorder import TuiRequestRecorder from utils.evidence import FileTask from utils.sql import query_count, resolve_trace_table, resolve_warehouse_id from utils.terminal import AgentTerminal @@ -35,6 +36,8 @@ CODEX_DEFAULT = CODEX_MODELS[0] CODEX_LUNA = CODEX_MODELS[1] EXPECTED_HEADER_NAMES = {"x-ug-e2e-run", "x-ug-e2e-agent"} +CLAUDE_REQUEST_PATH = "/ai-gateway/anthropic/v1/messages" +CODEX_REQUEST_PATH = "/ai-gateway/codex/v1/responses" TRACE_WAIT_SECONDS = 360 TRACE_POLL_SECONDS = 10 @@ -43,6 +46,103 @@ def _claude_models_visible(text: str) -> bool: return all(model in text for model in CLAUDE_MODELS) +def _claude_picker_visible(text: str) -> bool: + selected_model_row = any( + line.lstrip().startswith(("❯", "›", ">")) and any(model in line for model in CLAUDE_MODELS) + for line in text.splitlines() + ) + return selected_model_row and _claude_models_visible(text) + + +def _install_claude_recorder_launcher(session, recorder_base_url: str) -> None: + """Route the real Claude launch through a disposable settings copy.""" + wrapper_dir = session.home / ".ucode" / "e2e-claude-recorder-bin" + wrapper_dir.mkdir(parents=True, exist_ok=True) + original_path = session.env["PATH"] + claude_binary = shutil.which("claude", path=original_path) + assert claude_binary, original_path + script = f"""#!/usr/bin/env python3 +import json +import os +from pathlib import Path +import sys + +CLAUDE_BINARY = {json.dumps(claude_binary)} +RECORDER_BASE_URL = {json.dumps(recorder_base_url)} +OVERRIDE_DIR = Path({json.dumps(str(wrapper_dir))}) + +args = sys.argv[1:] +try: + settings_index = args.index("--settings") +except ValueError: + settings_index = -1 +if settings_index >= 0: + if settings_index + 1 >= len(args): + raise SystemExit("Claude recorder launcher received --settings without a value") + settings_value = args[settings_index + 1] + if settings_value.lstrip().startswith("{{"): + settings = json.loads(settings_value) + inline = True + else: + settings = json.loads(Path(settings_value).read_text(encoding="utf-8")) + inline = False + settings_env = settings.setdefault("env", {{}}) + if not isinstance(settings_env, dict): + raise SystemExit("Claude settings env must be an object") + settings_env["ANTHROPIC_BASE_URL"] = RECORDER_BASE_URL + if inline: + args[settings_index + 1] = json.dumps(settings, separators=(",", ":")) + else: + override_path = OVERRIDE_DIR / ("settings-" + str(os.getpid()) + ".json") + override_path.write_text(json.dumps(settings), encoding="utf-8") + args[settings_index + 1] = str(override_path) +os.execv(CLAUDE_BINARY, [CLAUDE_BINARY, *args]) +""" + launcher = wrapper_dir / "claude" + launcher.write_text(script, encoding="utf-8") + launcher.chmod(0o755) + session.env["PATH"] = os.pathsep.join((str(wrapper_dir), original_path)) + + +def _assert_inference_requests( + recorder: TuiRequestRecorder, + checkpoint: int, + path: str, + marker: str, + run_id: str, + agent: str, + expected_model: str, +) -> None: + """Assert every recorded inference request for one task and require a tool follow-up.""" + requests = [] + after = checkpoint + while True: + try: + request = recorder.expect_request(method="POST", path=path, after=after, timeout=0.2) + except AssertionError as exc: + assert str(exc).startswith("Timed out waiting for TUI request:"), str(exc) + break + requests.append(request) + after = request.sequence + + assert len(requests) >= 2, [(request.sequence, request.path) for request in requests] + served_models = [] + marked_requests = 0 + for request in requests: + assert request.headers["x-ug-e2e-run"] == run_id, request.headers + # AIGTWY-4876: Each request identifies the active agent, preventing cross-agent headers. + assert request.headers["x-ug-e2e-agent"] == agent, request.headers + payload = request.payload + assert isinstance(payload, dict), type(payload) + assert isinstance(payload.get("model"), str), sorted(payload) + if marker in json.dumps(payload): + marked_requests += 1 + if recorder.response_for(request).status_code == 200: + served_models.append(payload["model"]) + assert marked_requests >= 2, marked_requests + assert expected_model in served_models, served_models + + def _managed_config(session) -> dict: cache_path = session.home / ".ucode" / "managed-config.json" cache = json.loads(cache_path.read_text()) @@ -119,9 +219,10 @@ def _assert_published_config(raw: dict, entries: dict[str, dict]) -> dict[str, s return claude_headers -def _assert_generated_configs(session, claude_headers: dict[str, str]) -> None: +def _assert_generated_configs(session, claude_headers: dict[str, str], workspace: str) -> None: claude_settings_path = session.home / ".claude" / "ucode-settings.json" claude_settings = json.loads(claude_settings_path.read_text()) + # AIGTWY-4876: Each agent picker exposes exactly its configured model list. assert claude_settings["availableModels"] == CLAUDE_MODELS, claude_settings claude_picker = claude_settings["modelPicker"] assert claude_picker["replaceBuiltInOptions"] is True, claude_settings @@ -134,6 +235,7 @@ def _assert_generated_configs(session, claude_headers: dict[str, str]) -> None: ) ], claude_settings claude_env = claude_settings["env"] + assert claude_env["ANTHROPIC_BASE_URL"] == f"{workspace}/ai-gateway/anthropic", claude_settings assert claude_env["ANTHROPIC_DEFAULT_OPUS_MODEL"] == CLAUDE_MODELS[0] + "[1m]", claude_settings assert claude_env["ANTHROPIC_DEFAULT_SONNET_MODEL"] == (CLAUDE_MODELS[1] + "[1m]"), ( claude_settings @@ -151,6 +253,7 @@ def _assert_generated_configs(session, claude_headers: dict[str, str]) -> None: codex_profile_path = session.home / ".codex" / "ucode.config.toml" codex_profile = tomllib.loads(codex_profile_path.read_text()) codex_provider = codex_profile["model_providers"]["Databricks"] + assert codex_provider["base_url"] == f"{workspace}/ai-gateway/codex/v1", codex_profile assert _managed_header_subset(codex_provider["http_headers"]) == { "x-ug-e2e-run": claude_headers["x-ug-e2e-run"], "x-ug-e2e-agent": "codex", @@ -315,16 +418,18 @@ def _wait_for_trace_pairs( time.sleep(min(TRACE_POLL_SECONDS, remaining)) -class TestCujManagedConfig(BaseCujTest): - WORKSPACE_URL = "https://dbc-135c115c-c255.cloud.databricks.com" +MANAGED_WORKSPACE_URL = "https://dbc-135c115c-c255.cloud.databricks.com" + + +class TestCujManagedConfiguration(BaseCujTest): + WORKSPACE_URL = MANAGED_WORKSPACE_URL - def test_cuj_managed_config_sentinel(self, live_session): - """Scenario: configure one published two-agent policy and run both agents. + def test_cuj_managed_configuration(self, live_session): + """Scenario: configure a fresh session against the published two-agent policy. - Expected: no agent selector; configured models, defaults, tracing and headers in generated - files; six completed tasks and matching agent inference plus native gateway client spans. - Auxiliary client calls for a task may use another model. Gateway receipt of headers is not - asserted here. + Expected: ``ug configure`` selects both enabled agents without a picker, writes the + configured models, defaults, aliases, headers and tracing settings, and publishes the + managed Codex catalog. """ session = live_session workspace = self.WORKSPACE_URL @@ -338,12 +443,13 @@ def test_cuj_managed_config_sentinel(self, live_session): "--skip-upgrade", timeout=300, ) + # AIGTWY-4876: Fresh configure selects both managed agents without an agent selector. assert "Select coding agents to configure:" not in configured.stdout, configured.stdout raw = _managed_config(session) entries = _managed_entries(raw) claude_headers = _assert_published_config(raw, entries) - _assert_generated_configs(session, claude_headers) + _assert_generated_configs(session, claude_headers, workspace) codex_models = session.codex_model_ids( ["app-server", "--listen", "stdio://"], name="managed-sentinel-codex-models" ) @@ -365,107 +471,316 @@ def test_cuj_managed_config_sentinel(self, live_session): }, ) - bearer = session.env["DATABRICKS_BEARER"] - trace_table = resolve_trace_table(workspace, bearer) - warehouse_id = os.environ.get("UG_INTEGRATION_WAREHOUSE_ID", "").strip() - warehouse_id = warehouse_id or resolve_warehouse_id(workspace, bearer) - run_id = uuid.uuid4().hex - trace_pairs: list[tuple[str, str, str]] = [] - root_task = FileTask(session) - root_marker = f"ug-managed-sentinel-{run_id}-root" - trace_pairs.append((root_marker, CLAUDE_DEFAULTS["default_model"], "resource")) - session.env["OTEL_RESOURCE_ATTRIBUTES"] = f"ug_integration_marker={root_marker}" - with AgentTerminal( - session, "claude", [str(session.binary)], "managed-sentinel-root" - ) as tui: - tui.boot() - tui.submit(f"{root_task.prompt} Trace marker: {root_marker}") - tui.wait_for_task(root_task) - picker_screen = tui.open_model_picker(model_visible=_claude_models_visible) - assert _claude_models_visible(picker_screen), picker_screen - tui.exit_normally() - root_task.assert_completed(session, "claude") +class TestCujManagedClaude(BaseCujTest): + WORKSPACE_URL = MANAGED_WORKSPACE_URL - claude_opus_task = FileTask(session) - claude_opus_marker = f"ug-managed-sentinel-{run_id}-claude-opus" - trace_pairs.append( - (claude_opus_marker, CLAUDE_DEFAULTS["default_opus_model"] + "[1m]", "resource") + def test_cuj_managed_claude(self, live_session): + """Scenario: configure managed Claude and complete the default and alias TUI tasks. + + Expected: the configured Claude catalog and family defaults are visible, the model picker + shows all three configured models, each task completes, and the recorder sees the managed + headers and requested model on every inference request. + """ + session = live_session + workspace = self.WORKSPACE_URL + configured = session.run( + "configure", + "--workspace", + workspace, + "--skip-upgrade", + timeout=300, ) - session.env["OTEL_RESOURCE_ATTRIBUTES"] = f"ug_integration_marker={claude_opus_marker}" - with AgentTerminal( - session, - "claude", - [str(session.binary), "claude", "--model", "opus"], - "managed-sentinel-claude-opus", - ) as tui: - tui.boot() - tui.submit(f"{claude_opus_task.prompt} Trace marker: {claude_opus_marker}") - tui.wait_for_task(claude_opus_task) - tui.exit_normally() - claude_opus_task.assert_completed(session, "claude") + assert "Select coding agents to configure:" not in configured.stdout, configured.stdout - claude_sonnet_task = FileTask(session) - claude_sonnet_marker = f"ug-managed-sentinel-{run_id}-claude-sonnet" - trace_pairs.append( - (claude_sonnet_marker, CLAUDE_DEFAULTS["default_sonnet_model"] + "[1m]", "resource") + raw = _managed_config(session) + entries = _managed_entries(raw) + claude = entries["CODING_AGENT_CLAUDE_CODE"] + claude_headers = _assert_published_config(raw, entries) + assert claude["models"] == {"model_services": CLAUDE_MODELS}, claude + assert claude["default_models"] == CLAUDE_DEFAULTS, claude + assert claude_headers["x-ug-e2e-agent"] == "claude", claude_headers + + claude_settings_path = session.home / ".claude" / "ucode-settings.json" + settings = json.loads(claude_settings_path.read_text()) + assert settings["env"]["ANTHROPIC_BASE_URL"] == (f"{workspace}/ai-gateway/anthropic"), ( + settings ) - session.env["OTEL_RESOURCE_ATTRIBUTES"] = f"ug_integration_marker={claude_sonnet_marker}" - with AgentTerminal( - session, - "claude", - [str(session.binary), "claude", "--model", "sonnet"], - "managed-sentinel-claude-sonnet", - ) as tui: - tui.boot() - tui.submit(f"{claude_sonnet_task.prompt} Trace marker: {claude_sonnet_marker}") - tui.wait_for_task(claude_sonnet_task) - tui.exit_normally() - claude_sonnet_task.assert_completed(session, "claude") + run_id = claude_headers["x-ug-e2e-run"] + + with TuiRequestRecorder(workspace) as recorder: + _install_claude_recorder_launcher(session, f"{recorder.url}/ai-gateway/anthropic") + + root_task = FileTask(session) + root_marker = f"ug-managed-claude-{uuid.uuid4().hex}-root" + session.env["OTEL_RESOURCE_ATTRIBUTES"] = f"ug_integration_marker={root_marker}" + root_checkpoint = recorder.checkpoint() + with AgentTerminal( + session, "claude", [str(session.binary)], "managed-claude-root" + ) as tui: + tui.boot() + tui.submit(f"{root_task.prompt} Inference marker: {root_marker}") + tui.wait_for_task(root_task) + picker_screen = tui.open_model_picker(model_visible=_claude_picker_visible) + assert _claude_picker_visible(picker_screen), picker_screen + tui.exit_normally() + # AIGTWY-4876: Bare `ug` Claude has no override; the request below requires Sonnet. + root_task.assert_completed(session, "claude") + # AIGTWY-4876: Every Claude request, including tool follow-ups, carries both headers. + _assert_inference_requests( + recorder, + root_checkpoint, + CLAUDE_REQUEST_PATH, + root_marker, + run_id, + "claude", + CLAUDE_DEFAULTS["default_model"], + ) + + claude_opus_task = FileTask(session) + claude_opus_marker = f"ug-managed-claude-{uuid.uuid4().hex}-opus" + session.env["OTEL_RESOURCE_ATTRIBUTES"] = f"ug_integration_marker={claude_opus_marker}" + opus_checkpoint = recorder.checkpoint() + with AgentTerminal( + session, + "claude", + [str(session.binary), "claude", "--model", "opus"], + "managed-claude-opus", + ) as tui: + tui.boot() + tui.submit(f"{claude_opus_task.prompt} Inference marker: {claude_opus_marker}") + tui.wait_for_task(claude_opus_task) + tui.exit_normally() + # AIGTWY-4876: Opus, Sonnet, and Haiku aliases each require their real configured model. + claude_opus_task.assert_completed(session, "claude") + _assert_inference_requests( + recorder, + opus_checkpoint, + CLAUDE_REQUEST_PATH, + claude_opus_marker, + run_id, + "claude", + CLAUDE_DEFAULTS["default_opus_model"] + "[1m]", + ) + + claude_sonnet_task = FileTask(session) + claude_sonnet_marker = f"ug-managed-claude-{uuid.uuid4().hex}-sonnet" + session.env["OTEL_RESOURCE_ATTRIBUTES"] = ( + f"ug_integration_marker={claude_sonnet_marker}" + ) + sonnet_checkpoint = recorder.checkpoint() + with AgentTerminal( + session, + "claude", + [str(session.binary), "claude", "--model", "sonnet"], + "managed-claude-sonnet", + ) as tui: + tui.boot() + tui.submit(f"{claude_sonnet_task.prompt} Inference marker: {claude_sonnet_marker}") + tui.wait_for_task(claude_sonnet_task) + tui.exit_normally() + claude_sonnet_task.assert_completed(session, "claude") + _assert_inference_requests( + recorder, + sonnet_checkpoint, + CLAUDE_REQUEST_PATH, + claude_sonnet_marker, + run_id, + "claude", + CLAUDE_DEFAULTS["default_sonnet_model"] + "[1m]", + ) + + claude_haiku_task = FileTask(session) + claude_haiku_marker = f"ug-managed-claude-{uuid.uuid4().hex}-haiku" + session.env["OTEL_RESOURCE_ATTRIBUTES"] = f"ug_integration_marker={claude_haiku_marker}" + haiku_checkpoint = recorder.checkpoint() + with AgentTerminal( + session, + "claude", + [str(session.binary), "claude", "--model", "haiku"], + "managed-claude-haiku", + ) as tui: + tui.boot() + tui.submit(f"{claude_haiku_task.prompt} Inference marker: {claude_haiku_marker}") + tui.wait_for_task(claude_haiku_task) + tui.exit_normally() + claude_haiku_task.assert_completed(session, "claude") + _assert_inference_requests( + recorder, + haiku_checkpoint, + CLAUDE_REQUEST_PATH, + claude_haiku_marker, + run_id, + "claude", + CLAUDE_DEFAULTS["default_haiku_model"], + ) + final_settings = json.loads(claude_settings_path.read_text()) + assert final_settings["env"]["ANTHROPIC_BASE_URL"] == ( + f"{workspace}/ai-gateway/anthropic" + ), final_settings + assert _claude_header_lines(final_settings["env"]["ANTHROPIC_CUSTOM_HEADERS"]) == ( + claude_headers + ), final_settings + + +class TestCujManagedCodex(BaseCujTest): + WORKSPACE_URL = MANAGED_WORKSPACE_URL + + def test_cuj_managed_codex(self, live_session): + """Scenario: configure managed Codex and complete default and explicit model TUI tasks. + + Expected: the configured Codex catalog and default are visible, each task completes, and + the recorder sees the managed headers and requested Sol or Luna model on each request. + """ + session = live_session + workspace = self.WORKSPACE_URL + configured = session.run( + "configure", + "--workspace", + workspace, + "--skip-upgrade", + timeout=300, + ) + assert "Select coding agents to configure:" not in configured.stdout, configured.stdout - claude_haiku_task = FileTask(session) - claude_haiku_marker = f"ug-managed-sentinel-{run_id}-claude-haiku" - trace_pairs.append( - (claude_haiku_marker, CLAUDE_DEFAULTS["default_haiku_model"], "resource") + raw = _managed_config(session) + entries = _managed_entries(raw) + codex = entries["CODING_AGENT_CODEX"] + claude_headers = _assert_published_config(raw, entries) + assert codex["models"] == {"model_services": CODEX_MODELS}, codex + assert codex["default_models"] == {"default_model": CODEX_DEFAULT}, codex + assert _normalized_headers(codex["http_headers"])["x-ug-e2e-agent"] == "codex", codex + codex_models = session.codex_model_ids( + ["app-server", "--listen", "stdio://"], name="managed-codex-models" ) - session.env["OTEL_RESOURCE_ATTRIBUTES"] = f"ug_integration_marker={claude_haiku_marker}" - with AgentTerminal( - session, - "claude", - [str(session.binary), "claude", "--model", "haiku"], - "managed-sentinel-claude-haiku", - ) as tui: - tui.boot() - tui.submit(f"{claude_haiku_task.prompt} Trace marker: {claude_haiku_marker}") - tui.wait_for_task(claude_haiku_task) - tui.exit_normally() - claude_haiku_task.assert_completed(session, "claude") + assert codex_models == CODEX_MODELS, codex_models - codex_default_task = FileTask(session) - codex_default_marker = f"ug-managed-sentinel-{run_id}-codex-default" - trace_pairs.append((codex_default_marker, CODEX_DEFAULT, "span")) + profile = tomllib.loads((session.home / ".codex" / "ucode.config.toml").read_text()) + assert profile["model_providers"]["Databricks"]["base_url"] == ( + f"{workspace}/ai-gateway/codex/v1" + ), profile + run_id = claude_headers["x-ug-e2e-run"] session.env.pop("OTEL_RESOURCE_ATTRIBUTES", None) - with AgentTerminal( - session, - "codex", - [ - str(session.binary), + + with TuiRequestRecorder(workspace) as recorder: + recorder_base_url = f"{recorder.url}/ai-gateway/codex/v1" + + codex_default_task = FileTask(session) + codex_default_marker = f"ug-managed-codex-{uuid.uuid4().hex}-default" + default_checkpoint = recorder.checkpoint() + with AgentTerminal( + session, "codex", - "--", - "--config", - f'otel.span_attributes.ug_integration_marker="{codex_default_marker}"', - ], - "managed-sentinel-codex-default", - ) as tui: + [ + str(session.binary), + "codex", + "--", + "--config", + f'model_providers.Databricks.base_url="{recorder_base_url}"', + "--config", + f'otel.span_attributes.ug_integration_marker="{codex_default_marker}"', + ], + "managed-codex-default", + ) as tui: + tui.boot() + tui.submit(f"{codex_default_task.prompt} Inference marker: {codex_default_marker}") + tui.wait_for_task(codex_default_task) + tui.exit_normally() + # AIGTWY-4876: Bare `ug codex` has no override; the request below requires Sol. + codex_default_task.assert_completed(session, "codex") + # AIGTWY-4876: Every Codex request, including tool follow-ups, carries both headers. + _assert_inference_requests( + recorder, + default_checkpoint, + CODEX_REQUEST_PATH, + codex_default_marker, + run_id, + "codex", + CODEX_DEFAULT, + ) + + codex_luna_task = FileTask(session) + codex_luna_marker = f"ug-managed-codex-{uuid.uuid4().hex}-luna" + luna_checkpoint = recorder.checkpoint() + with AgentTerminal( + session, + "codex", + [ + str(session.binary), + "codex", + "--", + "--model", + CODEX_LUNA, + "--config", + f'model_providers.Databricks.base_url="{recorder_base_url}"', + "--config", + f'otel.span_attributes.ug_integration_marker="{codex_luna_marker}"', + ], + "managed-codex-luna", + ) as tui: + tui.boot() + tui.submit(f"{codex_luna_task.prompt} Inference marker: {codex_luna_marker}") + tui.wait_for_task(codex_luna_task) + tui.exit_normally() + # AIGTWY-4876: Explicit Luna selection completes and is served by the Luna model. + codex_luna_task.assert_completed(session, "codex") + _assert_inference_requests( + recorder, + luna_checkpoint, + CODEX_REQUEST_PATH, + codex_luna_marker, + run_id, + "codex", + CODEX_LUNA, + ) + + +class TestCujManagedTracing(BaseCujTest): + WORKSPACE_URL = MANAGED_WORKSPACE_URL + + def test_cuj_managed_tracing(self, live_session): + """Scenario: configure both managed agents and run uniquely marked tasks for each agent. + + Expected: one task for each agent completes through a real TUI, then the trace table + contains each expected marker/model pair in the agent and native client span layouts. + Auxiliary native client calls may use another model. + """ + session = live_session + workspace = self.WORKSPACE_URL + configured = session.run( + "configure", + "--workspace", + workspace, + "--skip-upgrade", + timeout=300, + ) + assert "Select coding agents to configure:" not in configured.stdout, configured.stdout + + raw = _managed_config(session) + entries = _managed_entries(raw) + _assert_published_config(raw, entries) + bearer = session.env["DATABRICKS_BEARER"] + trace_table = resolve_trace_table(workspace, bearer) + warehouse_id = os.environ.get("UG_INTEGRATION_WAREHOUSE_ID", "").strip() + warehouse_id = warehouse_id or resolve_warehouse_id(workspace, bearer) + run_id = uuid.uuid4().hex + trace_pairs: list[tuple[str, str, str]] = [] + + claude_task = FileTask(session) + claude_marker = f"ug-managed-trace-{run_id}-claude" + trace_pairs.append((claude_marker, CLAUDE_DEFAULTS["default_model"], "resource")) + session.env["OTEL_RESOURCE_ATTRIBUTES"] = f"ug_integration_marker={claude_marker}" + with AgentTerminal(session, "claude", [str(session.binary)], "managed-trace-claude") as tui: tui.boot() - tui.submit(f"{codex_default_task.prompt} Trace marker: {codex_default_marker}") - tui.wait_for_task(codex_default_task) + tui.submit(f"{claude_task.prompt} Trace marker: {claude_marker}") + tui.wait_for_task(claude_task) tui.exit_normally() - codex_default_task.assert_completed(session, "codex") + claude_task.assert_completed(session, "claude") - codex_luna_task = FileTask(session) - codex_luna_marker = f"ug-managed-sentinel-{run_id}-codex-luna" - trace_pairs.append((codex_luna_marker, CODEX_LUNA, "span")) + codex_task = FileTask(session) + codex_marker = f"ug-managed-trace-{run_id}-codex" + trace_pairs.append((codex_marker, CODEX_DEFAULT, "span")) + session.env.pop("OTEL_RESOURCE_ATTRIBUTES", None) with AgentTerminal( session, "codex", @@ -473,17 +788,16 @@ def test_cuj_managed_config_sentinel(self, live_session): str(session.binary), "codex", "--", - "--model", - CODEX_LUNA, "--config", - f'otel.span_attributes.ug_integration_marker="{codex_luna_marker}"', + f'otel.span_attributes.ug_integration_marker="{codex_marker}"', ], - "managed-sentinel-codex-luna", + "managed-trace-codex", ) as tui: tui.boot() - tui.submit(f"{codex_luna_task.prompt} Trace marker: {codex_luna_marker}") - tui.wait_for_task(codex_luna_task) + tui.submit(f"{codex_task.prompt} Trace marker: {codex_marker}") + tui.wait_for_task(codex_task) tui.exit_normally() - codex_luna_task.assert_completed(session, "codex") + codex_task.assert_completed(session, "codex") + # AIGTWY-4876: Poll for each marker/client model pair until the deadline. _wait_for_trace_pairs(workspace, bearer, warehouse_id, trace_table, trace_pairs, session) From 9a52f133d7a72577d4120206082becdc9b5ad27d Mon Sep 17 00:00:00 2001 From: John Speaks Date: Tue, 6 Oct 2026 04:03:06 +0000 Subject: [PATCH 5/9] [AIGTWY-4876] Report recorder paths on inference assertion failure --- tests/e2e_cuj/test_ug_managed_sentinel.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/tests/e2e_cuj/test_ug_managed_sentinel.py b/tests/e2e_cuj/test_ug_managed_sentinel.py index 5a251796e..076330e24 100644 --- a/tests/e2e_cuj/test_ug_managed_sentinel.py +++ b/tests/e2e_cuj/test_ug_managed_sentinel.py @@ -114,18 +114,19 @@ def _assert_inference_requests( expected_model: str, ) -> None: """Assert every recorded inference request for one task and require a tool follow-up.""" - requests = [] + observed = [] after = checkpoint while True: try: - request = recorder.expect_request(method="POST", path=path, after=after, timeout=0.2) + request = recorder.expect_request(method="POST", after=after, timeout=0.2) except AssertionError as exc: assert str(exc).startswith("Timed out waiting for TUI request:"), str(exc) break - requests.append(request) + observed.append(request) after = request.sequence - assert len(requests) >= 2, [(request.sequence, request.path) for request in requests] + requests = [request for request in observed if request.path == path] + assert len(requests) >= 2, [(request.sequence, request.path) for request in observed] served_models = [] marked_requests = 0 for request in requests: From 485d4f0dd130b3f8206bf88e595e884437386e42 Mon Sep 17 00:00:00 2001 From: John Speaks Date: Tue, 6 Oct 2026 12:20:51 +0000 Subject: [PATCH 6/9] [AIGTWY-4876] Capture managed agent inference with the recorder --- tests/e2e_cuj/test_ug_managed_sentinel.py | 231 ++++++++-------------- 1 file changed, 80 insertions(+), 151 deletions(-) diff --git a/tests/e2e_cuj/test_ug_managed_sentinel.py b/tests/e2e_cuj/test_ug_managed_sentinel.py index 076330e24..72713740c 100644 --- a/tests/e2e_cuj/test_ug_managed_sentinel.py +++ b/tests/e2e_cuj/test_ug_managed_sentinel.py @@ -1,15 +1,10 @@ -"""End-to-end journeys for the published managed Claude and Codex configuration. - -The selected managed workspace publishes one static configuration for both agents. Separate cases -check the generated files, record Claude and Codex inference requests, and query trace evidence. -""" +"""End-to-end journeys for managed Claude and Codex configuration.""" from __future__ import annotations import json import os import re -import shutil import time import tomllib import uuid @@ -54,56 +49,6 @@ def _claude_picker_visible(text: str) -> bool: return selected_model_row and _claude_models_visible(text) -def _install_claude_recorder_launcher(session, recorder_base_url: str) -> None: - """Route the real Claude launch through a disposable settings copy.""" - wrapper_dir = session.home / ".ucode" / "e2e-claude-recorder-bin" - wrapper_dir.mkdir(parents=True, exist_ok=True) - original_path = session.env["PATH"] - claude_binary = shutil.which("claude", path=original_path) - assert claude_binary, original_path - script = f"""#!/usr/bin/env python3 -import json -import os -from pathlib import Path -import sys - -CLAUDE_BINARY = {json.dumps(claude_binary)} -RECORDER_BASE_URL = {json.dumps(recorder_base_url)} -OVERRIDE_DIR = Path({json.dumps(str(wrapper_dir))}) - -args = sys.argv[1:] -try: - settings_index = args.index("--settings") -except ValueError: - settings_index = -1 -if settings_index >= 0: - if settings_index + 1 >= len(args): - raise SystemExit("Claude recorder launcher received --settings without a value") - settings_value = args[settings_index + 1] - if settings_value.lstrip().startswith("{{"): - settings = json.loads(settings_value) - inline = True - else: - settings = json.loads(Path(settings_value).read_text(encoding="utf-8")) - inline = False - settings_env = settings.setdefault("env", {{}}) - if not isinstance(settings_env, dict): - raise SystemExit("Claude settings env must be an object") - settings_env["ANTHROPIC_BASE_URL"] = RECORDER_BASE_URL - if inline: - args[settings_index + 1] = json.dumps(settings, separators=(",", ":")) - else: - override_path = OVERRIDE_DIR / ("settings-" + str(os.getpid()) + ".json") - override_path.write_text(json.dumps(settings), encoding="utf-8") - args[settings_index + 1] = str(override_path) -os.execv(CLAUDE_BINARY, [CLAUDE_BINARY, *args]) -""" - launcher = wrapper_dir / "claude" - launcher.write_text(script, encoding="utf-8") - launcher.chmod(0o755) - session.env["PATH"] = os.pathsep.join((str(wrapper_dir), original_path)) - - def _assert_inference_requests( recorder: TuiRequestRecorder, checkpoint: int, @@ -113,7 +58,7 @@ def _assert_inference_requests( agent: str, expected_model: str, ) -> None: - """Assert every recorded inference request for one task and require a tool follow-up.""" + """Check marked requests, headers, models, and tool follow-ups.""" observed = [] after = checkpoint while True: @@ -131,17 +76,20 @@ def _assert_inference_requests( marked_requests = 0 for request in requests: assert request.headers["x-ug-e2e-run"] == run_id, request.headers - # AIGTWY-4876: Each request identifies the active agent, preventing cross-agent headers. + # AIGTWY-4876: Requests retain the configured agent header. assert request.headers["x-ug-e2e-agent"] == agent, request.headers payload = request.payload assert isinstance(payload, dict), type(payload) assert isinstance(payload.get("model"), str), sorted(payload) if marker in json.dumps(payload): marked_requests += 1 - if recorder.response_for(request).status_code == 200: - served_models.append(payload["model"]) + if recorder.response_for(request).status_code == 200: + served_models.append(payload["model"]) assert marked_requests >= 2, marked_requests - assert expected_model in served_models, served_models + expected = re.sub(r"\[(?:1m|200k)\]$", "", expected_model) + assert ( + sum(re.sub(r"\[(?:1m|200k)\]$", "", model) == expected for model in served_models) >= 2 + ), served_models def _managed_config(session) -> dict: @@ -223,8 +171,9 @@ def _assert_published_config(raw: dict, entries: dict[str, dict]) -> dict[str, s def _assert_generated_configs(session, claude_headers: dict[str, str], workspace: str) -> None: claude_settings_path = session.home / ".claude" / "ucode-settings.json" claude_settings = json.loads(claude_settings_path.read_text()) - # AIGTWY-4876: Each agent picker exposes exactly its configured model list. + # AIGTWY-4876: Pickers expose exactly the configured model list. assert claude_settings["availableModels"] == CLAUDE_MODELS, claude_settings + assert claude_settings["enforceAvailableModels"] is True, claude_settings claude_picker = claude_settings["modelPicker"] assert claude_picker["replaceBuiltInOptions"] is True, claude_settings assert claude_picker["options"] == [ @@ -276,7 +225,7 @@ def _trace_pair_count( table: str, pairs: list[tuple[str, str, str]], ) -> int: - """Count distinct marker/model pairs in both agents' trace attribute layouts.""" + """Count trace marker/model pairs.""" clauses = [] parameters = [] for index, (marker, model, _marker_source) in enumerate(pairs): @@ -328,7 +277,7 @@ def _trace_pair_count( def _native_gateway_model(model: str) -> str: - """Return the model name emitted by the gateway's native client span.""" + """Normalize a gateway client span model.""" return re.sub(r"\[(?:1m|200k)\]$", "", model).removeprefix("system.ai.") @@ -339,7 +288,7 @@ def _trace_client_pair_count( table: str, pairs: list[tuple[str, str, str]], ) -> int: - """Count expected native gateway client models for the task markers.""" + """Count native client marker/model pairs.""" input_messages = "CAST(variant_get(attributes, '$[\"gen_ai.input.messages\"]') AS STRING)" model = "variant_get(attributes, '$[\"gen_ai.request.model\"]', 'STRING')" expected_cases = [] @@ -428,9 +377,7 @@ class TestCujManagedConfiguration(BaseCujTest): def test_cuj_managed_configuration(self, live_session): """Scenario: configure a fresh session against the published two-agent policy. - Expected: ``ug configure`` selects both enabled agents without a picker, writes the - configured models, defaults, aliases, headers and tracing settings, and publishes the - managed Codex catalog. + Expected: configure selects both agents without a picker and writes their managed settings. """ session = live_session workspace = self.WORKSPACE_URL @@ -444,7 +391,6 @@ def test_cuj_managed_configuration(self, live_session): "--skip-upgrade", timeout=300, ) - # AIGTWY-4876: Fresh configure selects both managed agents without an agent selector. assert "Select coding agents to configure:" not in configured.stdout, configured.stdout raw = _managed_config(session) @@ -477,40 +423,37 @@ class TestCujManagedClaude(BaseCujTest): WORKSPACE_URL = MANAGED_WORKSPACE_URL def test_cuj_managed_claude(self, live_session): - """Scenario: configure managed Claude and complete the default and alias TUI tasks. + """Scenario: configure managed Claude and complete default and alias TUI tasks. - Expected: the configured Claude catalog and family defaults are visible, the model picker - shows all three configured models, each task completes, and the recorder sees the managed - headers and requested model on every inference request. + Expected: the picker shows configured models; tasks complete with expected models and + headers on every recorded inference request. """ session = live_session workspace = self.WORKSPACE_URL - configured = session.run( - "configure", - "--workspace", - workspace, - "--skip-upgrade", - timeout=300, - ) - assert "Select coding agents to configure:" not in configured.stdout, configured.stdout - - raw = _managed_config(session) - entries = _managed_entries(raw) - claude = entries["CODING_AGENT_CLAUDE_CODE"] - claude_headers = _assert_published_config(raw, entries) - assert claude["models"] == {"model_services": CLAUDE_MODELS}, claude - assert claude["default_models"] == CLAUDE_DEFAULTS, claude - assert claude_headers["x-ug-e2e-agent"] == "claude", claude_headers - - claude_settings_path = session.home / ".claude" / "ucode-settings.json" - settings = json.loads(claude_settings_path.read_text()) - assert settings["env"]["ANTHROPIC_BASE_URL"] == (f"{workspace}/ai-gateway/anthropic"), ( - settings - ) - run_id = claude_headers["x-ug-e2e-run"] - with TuiRequestRecorder(workspace) as recorder: - _install_claude_recorder_launcher(session, f"{recorder.url}/ai-gateway/anthropic") + configured = session.run( + "configure", + "--workspace", + recorder.url, + "--skip-upgrade", + timeout=300, + ) + assert "Select coding agents to configure:" not in configured.stdout, configured.stdout + + raw = _managed_config(session) + entries = _managed_entries(raw) + claude = entries["CODING_AGENT_CLAUDE_CODE"] + claude_headers = _assert_published_config(raw, entries) + assert claude["models"] == {"model_services": CLAUDE_MODELS}, claude + assert claude["default_models"] == CLAUDE_DEFAULTS, claude + assert claude_headers["x-ug-e2e-agent"] == "claude", claude_headers + + claude_settings_path = session.home / ".claude" / "ucode-settings.json" + settings = json.loads(claude_settings_path.read_text()) + assert settings["env"]["ANTHROPIC_BASE_URL"] == ( + f"{recorder.url}/ai-gateway/anthropic" + ), settings + run_id = claude_headers["x-ug-e2e-run"] root_task = FileTask(session) root_marker = f"ug-managed-claude-{uuid.uuid4().hex}-root" @@ -525,9 +468,7 @@ def test_cuj_managed_claude(self, live_session): picker_screen = tui.open_model_picker(model_visible=_claude_picker_visible) assert _claude_picker_visible(picker_screen), picker_screen tui.exit_normally() - # AIGTWY-4876: Bare `ug` Claude has no override; the request below requires Sonnet. root_task.assert_completed(session, "claude") - # AIGTWY-4876: Every Claude request, including tool follow-ups, carries both headers. _assert_inference_requests( recorder, root_checkpoint, @@ -552,7 +493,6 @@ def test_cuj_managed_claude(self, live_session): tui.submit(f"{claude_opus_task.prompt} Inference marker: {claude_opus_marker}") tui.wait_for_task(claude_opus_task) tui.exit_normally() - # AIGTWY-4876: Opus, Sonnet, and Haiku aliases each require their real configured model. claude_opus_task.assert_completed(session, "claude") _assert_inference_requests( recorder, @@ -615,13 +555,13 @@ def test_cuj_managed_claude(self, live_session): "claude", CLAUDE_DEFAULTS["default_haiku_model"], ) - final_settings = json.loads(claude_settings_path.read_text()) - assert final_settings["env"]["ANTHROPIC_BASE_URL"] == ( - f"{workspace}/ai-gateway/anthropic" - ), final_settings - assert _claude_header_lines(final_settings["env"]["ANTHROPIC_CUSTOM_HEADERS"]) == ( - claude_headers - ), final_settings + final_settings = json.loads(claude_settings_path.read_text()) + assert final_settings["env"]["ANTHROPIC_BASE_URL"] == ( + f"{recorder.url}/ai-gateway/anthropic" + ), final_settings + assert _claude_header_lines(final_settings["env"]["ANTHROPIC_CUSTOM_HEADERS"]) == ( + claude_headers + ), final_settings class TestCujManagedCodex(BaseCujTest): @@ -630,41 +570,39 @@ class TestCujManagedCodex(BaseCujTest): def test_cuj_managed_codex(self, live_session): """Scenario: configure managed Codex and complete default and explicit model TUI tasks. - Expected: the configured Codex catalog and default are visible, each task completes, and - the recorder sees the managed headers and requested Sol or Luna model on each request. + Expected: the catalog is visible; tasks complete with expected Sol or Luna models and + managed headers. """ session = live_session workspace = self.WORKSPACE_URL - configured = session.run( - "configure", - "--workspace", - workspace, - "--skip-upgrade", - timeout=300, - ) - assert "Select coding agents to configure:" not in configured.stdout, configured.stdout - - raw = _managed_config(session) - entries = _managed_entries(raw) - codex = entries["CODING_AGENT_CODEX"] - claude_headers = _assert_published_config(raw, entries) - assert codex["models"] == {"model_services": CODEX_MODELS}, codex - assert codex["default_models"] == {"default_model": CODEX_DEFAULT}, codex - assert _normalized_headers(codex["http_headers"])["x-ug-e2e-agent"] == "codex", codex - codex_models = session.codex_model_ids( - ["app-server", "--listen", "stdio://"], name="managed-codex-models" - ) - assert codex_models == CODEX_MODELS, codex_models - - profile = tomllib.loads((session.home / ".codex" / "ucode.config.toml").read_text()) - assert profile["model_providers"]["Databricks"]["base_url"] == ( - f"{workspace}/ai-gateway/codex/v1" - ), profile - run_id = claude_headers["x-ug-e2e-run"] - session.env.pop("OTEL_RESOURCE_ATTRIBUTES", None) - with TuiRequestRecorder(workspace) as recorder: - recorder_base_url = f"{recorder.url}/ai-gateway/codex/v1" + configured = session.run( + "configure", + "--workspace", + recorder.url, + "--skip-upgrade", + timeout=300, + ) + assert "Select coding agents to configure:" not in configured.stdout, configured.stdout + + raw = _managed_config(session) + entries = _managed_entries(raw) + codex = entries["CODING_AGENT_CODEX"] + claude_headers = _assert_published_config(raw, entries) + assert codex["models"] == {"model_services": CODEX_MODELS}, codex + assert codex["default_models"] == {"default_model": CODEX_DEFAULT}, codex + assert _normalized_headers(codex["http_headers"])["x-ug-e2e-agent"] == "codex", codex + codex_models = session.codex_model_ids( + ["app-server", "--listen", "stdio://"], name="managed-codex-models" + ) + assert codex_models == CODEX_MODELS, codex_models + + profile = tomllib.loads((session.home / ".codex" / "ucode.config.toml").read_text()) + assert profile["model_providers"]["Databricks"]["base_url"] == ( + f"{recorder.url}/ai-gateway/codex/v1" + ), profile + run_id = claude_headers["x-ug-e2e-run"] + session.env.pop("OTEL_RESOURCE_ATTRIBUTES", None) codex_default_task = FileTask(session) codex_default_marker = f"ug-managed-codex-{uuid.uuid4().hex}-default" @@ -677,8 +615,6 @@ def test_cuj_managed_codex(self, live_session): "codex", "--", "--config", - f'model_providers.Databricks.base_url="{recorder_base_url}"', - "--config", f'otel.span_attributes.ug_integration_marker="{codex_default_marker}"', ], "managed-codex-default", @@ -687,9 +623,7 @@ def test_cuj_managed_codex(self, live_session): tui.submit(f"{codex_default_task.prompt} Inference marker: {codex_default_marker}") tui.wait_for_task(codex_default_task) tui.exit_normally() - # AIGTWY-4876: Bare `ug codex` has no override; the request below requires Sol. codex_default_task.assert_completed(session, "codex") - # AIGTWY-4876: Every Codex request, including tool follow-ups, carries both headers. _assert_inference_requests( recorder, default_checkpoint, @@ -713,8 +647,6 @@ def test_cuj_managed_codex(self, live_session): "--model", CODEX_LUNA, "--config", - f'model_providers.Databricks.base_url="{recorder_base_url}"', - "--config", f'otel.span_attributes.ug_integration_marker="{codex_luna_marker}"', ], "managed-codex-luna", @@ -723,7 +655,6 @@ def test_cuj_managed_codex(self, live_session): tui.submit(f"{codex_luna_task.prompt} Inference marker: {codex_luna_marker}") tui.wait_for_task(codex_luna_task) tui.exit_normally() - # AIGTWY-4876: Explicit Luna selection completes and is served by the Luna model. codex_luna_task.assert_completed(session, "codex") _assert_inference_requests( recorder, @@ -740,11 +671,10 @@ class TestCujManagedTracing(BaseCujTest): WORKSPACE_URL = MANAGED_WORKSPACE_URL def test_cuj_managed_tracing(self, live_session): - """Scenario: configure both managed agents and run uniquely marked tasks for each agent. + """Scenario: configure both managed agents and run uniquely marked tasks. - Expected: one task for each agent completes through a real TUI, then the trace table - contains each expected marker/model pair in the agent and native client span layouts. - Auxiliary native client calls may use another model. + Expected: both tasks complete and traces contain every marker/model pair in both span + layouts; auxiliary native client calls may use another model. """ session = live_session workspace = self.WORKSPACE_URL @@ -800,5 +730,4 @@ def test_cuj_managed_tracing(self, live_session): tui.exit_normally() codex_task.assert_completed(session, "codex") - # AIGTWY-4876: Poll for each marker/client model pair until the deadline. _wait_for_trace_pairs(workspace, bearer, warehouse_id, trace_table, trace_pairs, session) From 4f015f0bf2b61a53e3a37fa10a8d6b19939a3a8a Mon Sep 17 00:00:00 2001 From: John Speaks Date: Tue, 6 Oct 2026 14:34:39 +0000 Subject: [PATCH 7/9] [AIGTWY-4876] Restore picker close check for full journeys --- tests/integration/utils/terminal.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/integration/utils/terminal.py b/tests/integration/utils/terminal.py index 4601f55bf..7d7be1932 100644 --- a/tests/integration/utils/terminal.py +++ b/tests/integration/utils/terminal.py @@ -320,9 +320,7 @@ def open_model_picker(self, *, model_visible=None): self.actions.append({"reason": "model-picker-visible", "screen": screen}) self.send("\x1b", "close the model picker") self.wait_for( - lambda text: ( - not model_visible(text) if model_visible is not None else "Select model" not in text - ), + lambda text: "Select model" not in text, "the prompt after closing the model picker", ) return screen From 99097bd108c50c3cde28a9e6ce5509b4c3cfca69 Mon Sep 17 00:00:00 2001 From: John Speaks Date: Wed, 7 Oct 2026 02:00:59 +0000 Subject: [PATCH 8/9] [AIGTWY-4876] Adapt managed CUJ to shared test scaffold --- tests/e2e_cuj/test_ug_managed_sentinel.py | 516 +++++++++++----------- tests/integration/utils/terminal.py | 13 +- 2 files changed, 271 insertions(+), 258 deletions(-) diff --git a/tests/e2e_cuj/test_ug_managed_sentinel.py b/tests/e2e_cuj/test_ug_managed_sentinel.py index 72713740c..506605e71 100644 --- a/tests/e2e_cuj/test_ug_managed_sentinel.py +++ b/tests/e2e_cuj/test_ug_managed_sentinel.py @@ -10,11 +10,15 @@ import uuid import pytest -from base import BaseCujTest -from helpers.tui_request_recorder import TuiRequestRecorder -from utils.evidence import FileTask -from utils.sql import query_count, resolve_trace_table, resolve_warehouse_id -from utils.terminal import AgentTerminal + +from tests.integration.utils.evidence import FileTask +from tests.integration.utils.sql import query_count, resolve_trace_table, resolve_warehouse_id +from tests.integration.utils.terminal import AgentTerminal + +from .base import BaseCujTest +from .helpers.constants import CLAUDE, CODEX +from .helpers.evidence import SessionEvidence, canonical_model +from .helpers.tui_request_recorder import TuiRequestRecorder CLAUDE_MODELS = [ "system.ai.claude-opus-4-8", @@ -92,6 +96,14 @@ def _assert_inference_requests( ), served_models +def _assert_native_model(evidence: SessionEvidence, task: FileTask, expected: str) -> None: + turn = evidence.completed(task) + assert turn is not None, f"No completed native turn matched {task.prompt!r}" + expected_model = canonical_model(expected) + actual_models = {canonical_model(model) for model in turn.models} + assert actual_models == {expected_model}, (turn, expected_model) + + def _managed_config(session) -> dict: cache_path = session.home / ".ucode" / "managed-config.json" cache = json.loads(cache_path.read_text()) @@ -153,14 +165,14 @@ def _assert_published_config(raw: dict, entries: dict[str, dict]) -> dict[str, s assert claude["default_models"] == CLAUDE_DEFAULTS, claude assert claude["tracing"] == {"enabled": True}, claude claude_headers = _normalized_headers(claude["http_headers"]) - assert claude_headers["x-ug-e2e-agent"] == "claude", claude_headers + assert claude_headers["x-ug-e2e-agent"] == CLAUDE, claude_headers assert codex["smart_routing"] == {"enabled": False}, codex assert codex["models"] == {"model_services": CODEX_MODELS}, codex assert codex["default_models"] == {"default_model": CODEX_DEFAULT}, codex assert codex["tracing"] == {"enabled": True}, codex codex_headers = _normalized_headers(codex["http_headers"]) - assert codex_headers["x-ug-e2e-agent"] == "codex", codex_headers + assert codex_headers["x-ug-e2e-agent"] == CODEX, codex_headers assert codex_headers["x-ug-e2e-run"] == claude_headers["x-ug-e2e-run"], ( claude_headers, codex_headers, @@ -206,7 +218,7 @@ def _assert_generated_configs(session, claude_headers: dict[str, str], workspace assert codex_provider["base_url"] == f"{workspace}/ai-gateway/codex/v1", codex_profile assert _managed_header_subset(codex_provider["http_headers"]) == { "x-ug-e2e-run": claude_headers["x-ug-e2e-run"], - "x-ug-e2e-agent": "codex", + "x-ug-e2e-agent": CODEX, }, codex_profile codex_catalog_path = session.home / ".ucode" / "codex-model-catalog.json" assert codex_profile["model_catalog_json"] == str(codex_catalog_path), codex_profile @@ -374,20 +386,16 @@ def _wait_for_trace_pairs( class TestCujManagedConfiguration(BaseCujTest): WORKSPACE_URL = MANAGED_WORKSPACE_URL - def test_cuj_managed_configuration(self, live_session): + def test_cuj_managed_configuration(self, cuj): """Scenario: configure a fresh session against the published two-agent policy. Expected: configure selects both agents without a picker and writes their managed settings. """ - session = live_session - workspace = self.WORKSPACE_URL - configured_workspace = os.environ.get("UG_CUJ1_WORKSPACE", "").strip().rstrip("/") - if configured_workspace: - assert configured_workspace == workspace, configured_workspace + session, workspace, _ = cuj configured = session.run( "configure", "--workspace", - workspace, + workspace.url, "--skip-upgrade", timeout=300, ) @@ -396,7 +404,7 @@ def test_cuj_managed_configuration(self, live_session): raw = _managed_config(session) entries = _managed_entries(raw) claude_headers = _assert_published_config(raw, entries) - _assert_generated_configs(session, claude_headers, workspace) + _assert_generated_configs(session, claude_headers, workspace.url) codex_models = session.codex_model_ids( ["app-server", "--listen", "stdio://"], name="managed-sentinel-codex-models" ) @@ -422,266 +430,270 @@ def test_cuj_managed_configuration(self, live_session): class TestCujManagedClaude(BaseCujTest): WORKSPACE_URL = MANAGED_WORKSPACE_URL - def test_cuj_managed_claude(self, live_session): + def test_cuj_managed_claude(self, cuj): """Scenario: configure managed Claude and complete default and alias TUI tasks. Expected: the picker shows configured models; tasks complete with expected models and headers on every recorded inference request. """ - session = live_session - workspace = self.WORKSPACE_URL - with TuiRequestRecorder(workspace) 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 - - raw = _managed_config(session) - entries = _managed_entries(raw) - claude = entries["CODING_AGENT_CLAUDE_CODE"] - claude_headers = _assert_published_config(raw, entries) - assert claude["models"] == {"model_services": CLAUDE_MODELS}, claude - assert claude["default_models"] == CLAUDE_DEFAULTS, claude - assert claude_headers["x-ug-e2e-agent"] == "claude", claude_headers - - claude_settings_path = session.home / ".claude" / "ucode-settings.json" - settings = json.loads(claude_settings_path.read_text()) - assert settings["env"]["ANTHROPIC_BASE_URL"] == ( - f"{recorder.url}/ai-gateway/anthropic" - ), settings - run_id = claude_headers["x-ug-e2e-run"] - - root_task = FileTask(session) - root_marker = f"ug-managed-claude-{uuid.uuid4().hex}-root" - session.env["OTEL_RESOURCE_ATTRIBUTES"] = f"ug_integration_marker={root_marker}" - root_checkpoint = recorder.checkpoint() - with AgentTerminal( - session, "claude", [str(session.binary)], "managed-claude-root" - ) as tui: - tui.boot() - tui.submit(f"{root_task.prompt} Inference marker: {root_marker}") - tui.wait_for_task(root_task) - picker_screen = tui.open_model_picker(model_visible=_claude_picker_visible) - assert _claude_picker_visible(picker_screen), picker_screen - tui.exit_normally() - root_task.assert_completed(session, "claude") - _assert_inference_requests( - recorder, - root_checkpoint, - CLAUDE_REQUEST_PATH, - root_marker, - run_id, - "claude", - CLAUDE_DEFAULTS["default_model"], - ) + session, workspace, recorder = cuj + recorder.configure_session(session, ["configure", "--skip-upgrade"]) - claude_opus_task = FileTask(session) - claude_opus_marker = f"ug-managed-claude-{uuid.uuid4().hex}-opus" - session.env["OTEL_RESOURCE_ATTRIBUTES"] = f"ug_integration_marker={claude_opus_marker}" - opus_checkpoint = recorder.checkpoint() - with AgentTerminal( - session, - "claude", - [str(session.binary), "claude", "--model", "opus"], - "managed-claude-opus", - ) as tui: - tui.boot() - tui.submit(f"{claude_opus_task.prompt} Inference marker: {claude_opus_marker}") - tui.wait_for_task(claude_opus_task) - tui.exit_normally() - claude_opus_task.assert_completed(session, "claude") - _assert_inference_requests( - recorder, - opus_checkpoint, - CLAUDE_REQUEST_PATH, - claude_opus_marker, - run_id, - "claude", - CLAUDE_DEFAULTS["default_opus_model"] + "[1m]", - ) + raw = _managed_config(session) + entries = _managed_entries(raw) + claude = entries["CODING_AGENT_CLAUDE_CODE"] + claude_headers = _assert_published_config(raw, entries) + assert claude["models"] == {"model_services": CLAUDE_MODELS}, claude + assert claude["default_models"] == CLAUDE_DEFAULTS, claude + assert claude_headers["x-ug-e2e-agent"] == CLAUDE, claude_headers + + claude_settings_path = session.home / ".claude" / "ucode-settings.json" + settings = json.loads(claude_settings_path.read_text()) + configured_base_url = f"{workspace.url}/ai-gateway/anthropic" + assert settings["env"]["ANTHROPIC_BASE_URL"] == configured_base_url, settings + run_id = claude_headers["x-ug-e2e-run"] + + root_task = FileTask(session) + root_marker = f"ug-managed-claude-{uuid.uuid4().hex}-root" + root_task.prompt += f" Inference marker: {root_marker}" + session.env["OTEL_RESOURCE_ATTRIBUTES"] = f"ug_integration_marker={root_marker}" + evidence = SessionEvidence(session.home, CLAUDE) + recorder.prepare_launch() + root_checkpoint = recorder.checkpoint() + with AgentTerminal(session, CLAUDE, [str(session.binary)], "managed-claude-root") as tui: + tui.boot() + tui.submit(root_task.prompt) + tui.wait_for_task(root_task) + picker_screen = tui.open_model_picker(model_visible=_claude_picker_visible) + assert _claude_picker_visible(picker_screen), picker_screen + tui.exit_normally() + root_task.assert_completed(session, CLAUDE) + _assert_native_model(evidence, root_task, CLAUDE_DEFAULTS["default_model"]) + _assert_inference_requests( + recorder, + root_checkpoint, + CLAUDE_REQUEST_PATH, + root_marker, + run_id, + CLAUDE, + CLAUDE_DEFAULTS["default_model"], + ) - claude_sonnet_task = FileTask(session) - claude_sonnet_marker = f"ug-managed-claude-{uuid.uuid4().hex}-sonnet" - session.env["OTEL_RESOURCE_ATTRIBUTES"] = ( - f"ug_integration_marker={claude_sonnet_marker}" - ) - sonnet_checkpoint = recorder.checkpoint() - with AgentTerminal( - session, - "claude", - [str(session.binary), "claude", "--model", "sonnet"], - "managed-claude-sonnet", - ) as tui: - tui.boot() - tui.submit(f"{claude_sonnet_task.prompt} Inference marker: {claude_sonnet_marker}") - tui.wait_for_task(claude_sonnet_task) - tui.exit_normally() - claude_sonnet_task.assert_completed(session, "claude") - _assert_inference_requests( - recorder, - sonnet_checkpoint, - CLAUDE_REQUEST_PATH, - claude_sonnet_marker, - run_id, - "claude", - CLAUDE_DEFAULTS["default_sonnet_model"] + "[1m]", - ) + claude_opus_task = FileTask(session) + claude_opus_marker = f"ug-managed-claude-{uuid.uuid4().hex}-opus" + claude_opus_task.prompt += f" Inference marker: {claude_opus_marker}" + session.env["OTEL_RESOURCE_ATTRIBUTES"] = f"ug_integration_marker={claude_opus_marker}" + evidence = SessionEvidence(session.home, CLAUDE) + recorder.prepare_launch() + opus_checkpoint = recorder.checkpoint() + with AgentTerminal( + session, + CLAUDE, + [str(session.binary), CLAUDE, "--model", "opus"], + "managed-claude-opus", + ) as tui: + tui.boot() + tui.submit(claude_opus_task.prompt) + tui.wait_for_task(claude_opus_task) + tui.exit_normally() + claude_opus_task.assert_completed(session, CLAUDE) + _assert_native_model(evidence, claude_opus_task, CLAUDE_DEFAULTS["default_opus_model"]) + _assert_inference_requests( + recorder, + opus_checkpoint, + CLAUDE_REQUEST_PATH, + claude_opus_marker, + run_id, + CLAUDE, + CLAUDE_DEFAULTS["default_opus_model"] + "[1m]", + ) - claude_haiku_task = FileTask(session) - claude_haiku_marker = f"ug-managed-claude-{uuid.uuid4().hex}-haiku" - session.env["OTEL_RESOURCE_ATTRIBUTES"] = f"ug_integration_marker={claude_haiku_marker}" - haiku_checkpoint = recorder.checkpoint() - with AgentTerminal( - session, - "claude", - [str(session.binary), "claude", "--model", "haiku"], - "managed-claude-haiku", - ) as tui: - tui.boot() - tui.submit(f"{claude_haiku_task.prompt} Inference marker: {claude_haiku_marker}") - tui.wait_for_task(claude_haiku_task) - tui.exit_normally() - claude_haiku_task.assert_completed(session, "claude") - _assert_inference_requests( - recorder, - haiku_checkpoint, - CLAUDE_REQUEST_PATH, - claude_haiku_marker, - run_id, - "claude", - CLAUDE_DEFAULTS["default_haiku_model"], - ) - final_settings = json.loads(claude_settings_path.read_text()) - assert final_settings["env"]["ANTHROPIC_BASE_URL"] == ( - f"{recorder.url}/ai-gateway/anthropic" - ), final_settings - assert _claude_header_lines(final_settings["env"]["ANTHROPIC_CUSTOM_HEADERS"]) == ( - claude_headers - ), final_settings + claude_sonnet_task = FileTask(session) + claude_sonnet_marker = f"ug-managed-claude-{uuid.uuid4().hex}-sonnet" + claude_sonnet_task.prompt += f" Inference marker: {claude_sonnet_marker}" + session.env["OTEL_RESOURCE_ATTRIBUTES"] = f"ug_integration_marker={claude_sonnet_marker}" + evidence = SessionEvidence(session.home, CLAUDE) + recorder.prepare_launch() + sonnet_checkpoint = recorder.checkpoint() + with AgentTerminal( + session, + CLAUDE, + [str(session.binary), CLAUDE, "--model", "sonnet"], + "managed-claude-sonnet", + ) as tui: + tui.boot() + tui.submit(claude_sonnet_task.prompt) + tui.wait_for_task(claude_sonnet_task) + tui.exit_normally() + claude_sonnet_task.assert_completed(session, CLAUDE) + _assert_native_model(evidence, claude_sonnet_task, CLAUDE_DEFAULTS["default_sonnet_model"]) + _assert_inference_requests( + recorder, + sonnet_checkpoint, + CLAUDE_REQUEST_PATH, + claude_sonnet_marker, + run_id, + CLAUDE, + CLAUDE_DEFAULTS["default_sonnet_model"] + "[1m]", + ) + + claude_haiku_task = FileTask(session) + claude_haiku_marker = f"ug-managed-claude-{uuid.uuid4().hex}-haiku" + claude_haiku_task.prompt += f" Inference marker: {claude_haiku_marker}" + session.env["OTEL_RESOURCE_ATTRIBUTES"] = f"ug_integration_marker={claude_haiku_marker}" + evidence = SessionEvidence(session.home, CLAUDE) + recorder.prepare_launch() + haiku_checkpoint = recorder.checkpoint() + with AgentTerminal( + session, + CLAUDE, + [str(session.binary), CLAUDE, "--model", "haiku"], + "managed-claude-haiku", + ) as tui: + tui.boot() + tui.submit(claude_haiku_task.prompt) + tui.wait_for_task(claude_haiku_task) + tui.exit_normally() + claude_haiku_task.assert_completed(session, CLAUDE) + _assert_native_model(evidence, claude_haiku_task, CLAUDE_DEFAULTS["default_haiku_model"]) + _assert_inference_requests( + recorder, + haiku_checkpoint, + CLAUDE_REQUEST_PATH, + claude_haiku_marker, + run_id, + CLAUDE, + CLAUDE_DEFAULTS["default_haiku_model"], + ) + final_settings = json.loads(claude_settings_path.read_text()) + assert final_settings["env"]["ANTHROPIC_BASE_URL"] == ( + f"{recorder.url}/ai-gateway/anthropic" + ), final_settings + assert _claude_header_lines(final_settings["env"]["ANTHROPIC_CUSTOM_HEADERS"]) == ( + claude_headers + ), final_settings class TestCujManagedCodex(BaseCujTest): WORKSPACE_URL = MANAGED_WORKSPACE_URL - def test_cuj_managed_codex(self, live_session): + def test_cuj_managed_codex(self, cuj): """Scenario: configure managed Codex and complete default and explicit model TUI tasks. Expected: the catalog is visible; tasks complete with expected Sol or Luna models and managed headers. """ - session = live_session - workspace = self.WORKSPACE_URL - with TuiRequestRecorder(workspace) 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 - - raw = _managed_config(session) - entries = _managed_entries(raw) - codex = entries["CODING_AGENT_CODEX"] - claude_headers = _assert_published_config(raw, entries) - assert codex["models"] == {"model_services": CODEX_MODELS}, codex - assert codex["default_models"] == {"default_model": CODEX_DEFAULT}, codex - assert _normalized_headers(codex["http_headers"])["x-ug-e2e-agent"] == "codex", codex - codex_models = session.codex_model_ids( - ["app-server", "--listen", "stdio://"], name="managed-codex-models" - ) - assert codex_models == CODEX_MODELS, codex_models - - profile = tomllib.loads((session.home / ".codex" / "ucode.config.toml").read_text()) - assert profile["model_providers"]["Databricks"]["base_url"] == ( - f"{recorder.url}/ai-gateway/codex/v1" - ), profile - run_id = claude_headers["x-ug-e2e-run"] - session.env.pop("OTEL_RESOURCE_ATTRIBUTES", None) - - codex_default_task = FileTask(session) - codex_default_marker = f"ug-managed-codex-{uuid.uuid4().hex}-default" - default_checkpoint = recorder.checkpoint() - with AgentTerminal( - session, - "codex", - [ - str(session.binary), - "codex", - "--", - "--config", - f'otel.span_attributes.ug_integration_marker="{codex_default_marker}"', - ], - "managed-codex-default", - ) as tui: - tui.boot() - tui.submit(f"{codex_default_task.prompt} Inference marker: {codex_default_marker}") - tui.wait_for_task(codex_default_task) - tui.exit_normally() - codex_default_task.assert_completed(session, "codex") - _assert_inference_requests( - recorder, - default_checkpoint, - CODEX_REQUEST_PATH, - codex_default_marker, - run_id, - "codex", - CODEX_DEFAULT, - ) + session, workspace, recorder = cuj + recorder.configure_session(session, ["configure", "--skip-upgrade"]) + + raw = _managed_config(session) + entries = _managed_entries(raw) + codex = entries["CODING_AGENT_CODEX"] + claude_headers = _assert_published_config(raw, entries) + assert codex["models"] == {"model_services": CODEX_MODELS}, codex + assert codex["default_models"] == {"default_model": CODEX_DEFAULT}, codex + assert _normalized_headers(codex["http_headers"])["x-ug-e2e-agent"] == CODEX, codex + codex_models = session.codex_model_ids( + ["app-server", "--listen", "stdio://"], name="managed-codex-models" + ) + assert codex_models == CODEX_MODELS, codex_models + + profile = tomllib.loads((session.home / ".codex" / "ucode.config.toml").read_text()) + assert profile["model_providers"]["Databricks"]["base_url"] == ( + f"{workspace.url}/ai-gateway/codex/v1" + ), profile + run_id = claude_headers["x-ug-e2e-run"] + session.env.pop("OTEL_RESOURCE_ATTRIBUTES", None) + + codex_default_task = FileTask(session) + codex_default_marker = f"ug-managed-codex-{uuid.uuid4().hex}-default" + codex_default_task.prompt += f" Inference marker: {codex_default_marker}" + evidence = SessionEvidence(session.home, CODEX) + recorder.prepare_launch() + default_checkpoint = recorder.checkpoint() + with AgentTerminal( + session, + CODEX, + [ + str(session.binary), + CODEX, + "--", + "--config", + f'otel.span_attributes.ug_integration_marker="{codex_default_marker}"', + ], + "managed-codex-default", + ) as tui: + tui.boot() + tui.submit(codex_default_task.prompt) + tui.wait_for_task(codex_default_task) + tui.exit_normally() + codex_default_task.assert_completed(session, CODEX) + _assert_native_model(evidence, codex_default_task, CODEX_DEFAULT) + _assert_inference_requests( + recorder, + default_checkpoint, + CODEX_REQUEST_PATH, + codex_default_marker, + run_id, + CODEX, + CODEX_DEFAULT, + ) - codex_luna_task = FileTask(session) - codex_luna_marker = f"ug-managed-codex-{uuid.uuid4().hex}-luna" - luna_checkpoint = recorder.checkpoint() - with AgentTerminal( - session, - "codex", - [ - str(session.binary), - "codex", - "--", - "--model", - CODEX_LUNA, - "--config", - f'otel.span_attributes.ug_integration_marker="{codex_luna_marker}"', - ], - "managed-codex-luna", - ) as tui: - tui.boot() - tui.submit(f"{codex_luna_task.prompt} Inference marker: {codex_luna_marker}") - tui.wait_for_task(codex_luna_task) - tui.exit_normally() - codex_luna_task.assert_completed(session, "codex") - _assert_inference_requests( - recorder, - luna_checkpoint, - CODEX_REQUEST_PATH, - codex_luna_marker, - run_id, - "codex", + codex_luna_task = FileTask(session) + codex_luna_marker = f"ug-managed-codex-{uuid.uuid4().hex}-luna" + codex_luna_task.prompt += f" Inference marker: {codex_luna_marker}" + evidence = SessionEvidence(session.home, CODEX) + recorder.prepare_launch() + luna_checkpoint = recorder.checkpoint() + with AgentTerminal( + session, + CODEX, + [ + str(session.binary), + CODEX, + "--", + "--model", CODEX_LUNA, - ) + "--config", + f'otel.span_attributes.ug_integration_marker="{codex_luna_marker}"', + ], + "managed-codex-luna", + ) as tui: + tui.boot() + tui.submit(codex_luna_task.prompt) + tui.wait_for_task(codex_luna_task) + tui.exit_normally() + codex_luna_task.assert_completed(session, CODEX) + _assert_native_model(evidence, codex_luna_task, CODEX_LUNA) + _assert_inference_requests( + recorder, + luna_checkpoint, + CODEX_REQUEST_PATH, + codex_luna_marker, + run_id, + CODEX, + CODEX_LUNA, + ) + final_profile = tomllib.loads((session.home / ".codex" / "ucode.config.toml").read_text()) + assert final_profile["model_providers"]["Databricks"]["base_url"] == ( + f"{recorder.url}/ai-gateway/codex/v1" + ), final_profile class TestCujManagedTracing(BaseCujTest): WORKSPACE_URL = MANAGED_WORKSPACE_URL - def test_cuj_managed_tracing(self, live_session): + def test_cuj_managed_tracing(self, cuj): """Scenario: configure both managed agents and run uniquely marked tasks. Expected: both tasks complete and traces contain every marker/model pair in both span layouts; auxiliary native client calls may use another model. """ - session = live_session - workspace = self.WORKSPACE_URL + session, workspace, _ = cuj configured = session.run( "configure", "--workspace", - workspace, + workspace.url, "--skip-upgrade", timeout=300, ) @@ -691,9 +703,9 @@ def test_cuj_managed_tracing(self, live_session): entries = _managed_entries(raw) _assert_published_config(raw, entries) bearer = session.env["DATABRICKS_BEARER"] - trace_table = resolve_trace_table(workspace, bearer) + trace_table = resolve_trace_table(workspace.url, bearer) warehouse_id = os.environ.get("UG_INTEGRATION_WAREHOUSE_ID", "").strip() - warehouse_id = warehouse_id or resolve_warehouse_id(workspace, bearer) + warehouse_id = warehouse_id or resolve_warehouse_id(workspace.url, bearer) run_id = uuid.uuid4().hex trace_pairs: list[tuple[str, str, str]] = [] @@ -701,12 +713,12 @@ def test_cuj_managed_tracing(self, live_session): claude_marker = f"ug-managed-trace-{run_id}-claude" trace_pairs.append((claude_marker, CLAUDE_DEFAULTS["default_model"], "resource")) session.env["OTEL_RESOURCE_ATTRIBUTES"] = f"ug_integration_marker={claude_marker}" - with AgentTerminal(session, "claude", [str(session.binary)], "managed-trace-claude") as tui: + with AgentTerminal(session, CLAUDE, [str(session.binary)], "managed-trace-claude") as tui: tui.boot() tui.submit(f"{claude_task.prompt} Trace marker: {claude_marker}") tui.wait_for_task(claude_task) tui.exit_normally() - claude_task.assert_completed(session, "claude") + claude_task.assert_completed(session, CLAUDE) codex_task = FileTask(session) codex_marker = f"ug-managed-trace-{run_id}-codex" @@ -714,10 +726,10 @@ def test_cuj_managed_tracing(self, live_session): session.env.pop("OTEL_RESOURCE_ATTRIBUTES", None) with AgentTerminal( session, - "codex", + CODEX, [ str(session.binary), - "codex", + CODEX, "--", "--config", f'otel.span_attributes.ug_integration_marker="{codex_marker}"', @@ -728,6 +740,8 @@ def test_cuj_managed_tracing(self, live_session): tui.submit(f"{codex_task.prompt} Trace marker: {codex_marker}") tui.wait_for_task(codex_task) tui.exit_normally() - codex_task.assert_completed(session, "codex") + codex_task.assert_completed(session, CODEX) - _wait_for_trace_pairs(workspace, bearer, warehouse_id, trace_table, trace_pairs, session) + _wait_for_trace_pairs( + workspace.url, bearer, warehouse_id, trace_table, trace_pairs, session + ) diff --git a/tests/integration/utils/terminal.py b/tests/integration/utils/terminal.py index 7d7be1932..38cef65ce 100644 --- a/tests/integration/utils/terminal.py +++ b/tests/integration/utils/terminal.py @@ -308,13 +308,12 @@ def check_input_and_exit(self): def open_model_picker(self, *, model_visible=None): """Open Claude's real model picker, record it, then return to the prompt.""" self.submit("/model") - if model_visible is None: - self.wait_for( - lambda text: "Select model" in text and "Switch between Claude models." in text, - "the model picker", - timeout=60, - ) - else: + self.wait_for( + lambda text: "Select model" in text and "Switch between Claude models." in text, + "the model picker", + timeout=60, + ) + if model_visible is not None: self.wait_for(model_visible, "a discovered model in the picker", timeout=60) screen = self.visible self.actions.append({"reason": "model-picker-visible", "screen": screen}) From cae93eac27638d149e07351f4334769ebc183d01 Mon Sep 17 00:00:00 2001 From: John Speaks Date: Wed, 7 Oct 2026 02:38:29 +0000 Subject: [PATCH 9/9] [AIGTWY-4876] Align managed CUJ model and recorder checks --- tests/e2e_cuj/test_ug_managed_sentinel.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/tests/e2e_cuj/test_ug_managed_sentinel.py b/tests/e2e_cuj/test_ug_managed_sentinel.py index 506605e71..cdf367766 100644 --- a/tests/e2e_cuj/test_ug_managed_sentinel.py +++ b/tests/e2e_cuj/test_ug_managed_sentinel.py @@ -39,6 +39,9 @@ CODEX_REQUEST_PATH = "/ai-gateway/codex/v1/responses" TRACE_WAIT_SECONDS = 360 TRACE_POLL_SECONDS = 10 +NATIVE_MODEL_ALIASES = { + "system.ai.anthropic.claude-haiku-4-5-20251001-v1:0": CLAUDE_MODELS[2], +} def _claude_models_visible(text: str) -> bool: @@ -100,7 +103,9 @@ def _assert_native_model(evidence: SessionEvidence, task: FileTask, expected: st turn = evidence.completed(task) assert turn is not None, f"No completed native turn matched {task.prompt!r}" expected_model = canonical_model(expected) - actual_models = {canonical_model(model) for model in turn.models} + actual_models = { + NATIVE_MODEL_ALIASES.get(model, model) for model in map(canonical_model, turn.models) + } assert actual_models == {expected_model}, (turn, expected_model) @@ -600,7 +605,7 @@ def test_cuj_managed_codex(self, cuj): profile = tomllib.loads((session.home / ".codex" / "ucode.config.toml").read_text()) assert profile["model_providers"]["Databricks"]["base_url"] == ( - f"{workspace.url}/ai-gateway/codex/v1" + f"{recorder.url}/ai-gateway/codex/v1" ), profile run_id = claude_headers["x-ug-e2e-run"] session.env.pop("OTEL_RESOURCE_ATTRIBUTES", None)