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..cdf367766 --- /dev/null +++ b/tests/e2e_cuj/test_ug_managed_sentinel.py @@ -0,0 +1,752 @@ +"""End-to-end journeys for managed Claude and Codex configuration.""" + +from __future__ import annotations + +import json +import os +import re +import time +import tomllib +import uuid + +import pytest + +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", + "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_REQUEST_PATH = "/ai-gateway/anthropic/v1/messages" +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: + 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 _assert_inference_requests( + recorder: TuiRequestRecorder, + checkpoint: int, + path: str, + marker: str, + run_id: str, + agent: str, + expected_model: str, +) -> None: + """Check marked requests, headers, models, and tool follow-ups.""" + observed = [] + after = checkpoint + while True: + try: + 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 + observed.append(request) + after = request.sequence + + 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: + assert request.headers["x-ug-e2e-run"] == run_id, request.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"]) + assert marked_requests >= 2, marked_requests + 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 _assert_native_model(evidence: SessionEvidence, task: FileTask, expected: str) -> None: + turn = evidence.completed(task) + assert turn is not None, f"No completed native turn matched {task.prompt!r}" + expected_model = canonical_model(expected) + actual_models = { + NATIVE_MODEL_ALIASES.get(model, model) for model in map(canonical_model, turn.models) + } + assert actual_models == {expected_model}, (turn, expected_model) + + +def _managed_config(session) -> dict: + cache_path = session.home / ".ucode" / "managed-config.json" + cache = json.loads(cache_path.read_text()) + 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], workspace: str) -> None: + claude_settings_path = session.home / ".claude" / "ucode-settings.json" + claude_settings = json.loads(claude_settings_path.read_text()) + # 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"] == [ + {"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_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 + ) + 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()) + 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, + }, 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" + ] == 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 trace marker/model pairs.""" + 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: + """Normalize a gateway client span model.""" + 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 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 = [] + 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)) + + +MANAGED_WORKSPACE_URL = "https://dbc-135c115c-c255.cloud.databricks.com" + + +class TestCujManagedConfiguration(BaseCujTest): + WORKSPACE_URL = MANAGED_WORKSPACE_URL + + 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, workspace, _ = cuj + configured = session.run( + "configure", + "--workspace", + workspace.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_headers = _assert_published_config(raw, entries) + _assert_generated_configs(session, claude_headers, workspace.url) + 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() + }, + }, + ) + + +class TestCujManagedClaude(BaseCujTest): + WORKSPACE_URL = MANAGED_WORKSPACE_URL + + 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, workspace, recorder = cuj + recorder.configure_session(session, ["configure", "--skip-upgrade"]) + + 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_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_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, 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, 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"{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" + 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" + 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, 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, workspace, _ = cuj + configured = session.run( + "configure", + "--workspace", + workspace.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) + _assert_published_config(raw, entries) + bearer = session.env["DATABRICKS_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.url, 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"{claude_task.prompt} Trace marker: {claude_marker}") + tui.wait_for_task(claude_task) + tui.exit_normally() + claude_task.assert_completed(session, CLAUDE) + + 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, + [ + str(session.binary), + CODEX, + "--", + "--config", + f'otel.span_attributes.ug_integration_marker="{codex_marker}"', + ], + "managed-trace-codex", + ) as tui: + tui.boot() + 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) + + _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 9f4c4fda7..38cef65ce 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, @@ -307,7 +314,6 @@ def open_model_picker(self, *, model_visible=None): timeout=60, ) if model_visible is not None: - # Native discovery can finish after the picker shell first renders. 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}) 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" },