diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml index bd7e5d55..8a84cd7a 100644 --- a/.github/workflows/integration.yml +++ b/.github/workflows/integration.yml @@ -188,6 +188,7 @@ jobs: AGENT: ${{ matrix.agent }} TEST_MARKER: ${{ inputs.suite == 'tui' && 'live and tui' || 'live' }} and ${{ matrix.agent }} ARTIFACT_NAME: integration-full-${{ matrix.agent }} + UG_INTEGRATION_TRACE_TABLE: main.aigw_tracing.unity_gateway_otel_spans steps: *live-steps managed: @@ -237,7 +238,7 @@ jobs: uv run --no-project --python 3.12 python scripts/run_integration.py \ --python 3.12 --ug-version "$UG_VERSION" --entry-point "$ENTRY_POINT" \ --default-index "$PACKAGE_INDEX" --output "$RUNNER_TEMP/ug-integration" \ - "${args[@]}" -- -m "(managed or managed_fixture) and $AGENT" + "${args[@]}" -- -m "(managed or (managed_fixture and not live)) and $AGENT" - name: Upload managed test evidence if: ${{ !cancelled() }} uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2 diff --git a/scripts/run_integration.py b/scripts/run_integration.py index 9fa73e2b..85a59d98 100644 --- a/scripts/run_integration.py +++ b/scripts/run_integration.py @@ -127,6 +127,8 @@ def arguments(): parser.add_argument("--npm-registry", default="https://registry.npmjs.org") parser.add_argument("--profile", help="Explicit Databricks profile to mint the live bearer.") parser.add_argument("--workspace", default=os.environ.get("UCODE_TEST_WORKSPACE")) + parser.add_argument("--trace-table", default=os.environ.get("UG_INTEGRATION_TRACE_TABLE")) + parser.add_argument("--warehouse-id", default=os.environ.get("UG_INTEGRATION_WAREHOUSE_ID")) parser.add_argument("--output", type=Path, help="New results directory; never reused.") parser.add_argument("--installation-only", action="store_true", help="No workspace calls.") parser.add_argument( @@ -292,6 +294,8 @@ def run(command, *, cwd=output, env=base_env, timeout=600) -> str: "codex_provider_model": args.codex_provider_model, "dependencies": args.dependency, "workspace": args.workspace, + "trace_table": args.trace_table, + "warehouse_id": args.warehouse_id, }, "platform": platform.platform(), "installation_only": args.installation_only, @@ -532,6 +536,8 @@ def run(command, *, cwd=output, env=base_env, timeout=600) -> str: "UG_INTEGRATION_CODEX_PROVIDER_MODEL": args.codex_provider_model, "UCODE_TEST_WORKSPACE": args.workspace or "", "DATABRICKS_BEARER": bearer, + "UG_INTEGRATION_TRACE_TABLE": args.trace_table or "", + "UG_INTEGRATION_WAREHOUSE_ID": args.warehouse_id or "", } ) for agent in agents: diff --git a/tests/README.md b/tests/README.md index 7150f84b..2afc9b8d 100644 --- a/tests/README.md +++ b/tests/README.md @@ -40,6 +40,7 @@ All tests live directly in `integration/`; shared mechanics live in `utils/`. | `test_ug_claude_custom_oauth_cli_boots`, `test_ug_codex_custom_oauth_cli_boots` | Launch with `ENABLE_CUSTOM_OAUTH_FROM_CLI=1`, `--workspace`, and `--client-id databricks-cli` | Real TUI reaches a usable prompt, accepts keyboard input, exits normally, and saves `client_id = databricks-cli` in its generated CLI profile; Claude also reads the OS-managed settings and requires a profile-only `apiKeyHelper` | | `test_ug_claude_headless_prompt_argument`, `test_ug_claude_headless_prompt_stdin`, `test_ug_claude_headless_prompt_after_separator` | Run Claude from a script using each prompt form | Structured final answer contains the file value; exit zero; no routing | | `test_ug_codex_headless_prompt_argument`, `test_ug_codex_headless_prompt_stdin`, `test_ug_codex_headless_prompt_after_separator` | Run Codex from a script using each prompt form | Completed turn and final answer contain the file value; exit zero; no routing | +| `test_ug_codex_exports_trace_to_configured_table` | Configure Codex tracing, complete a headless task carrying a unique trace marker, then wait for ingestion | The configured trace table contains a Codex span with the same trace-safe marker | | `test_ug_claude_headless_explicit_model_bypasses_routing` | Pass `--model VALUE` / `--model=VALUE` with routing enabled | Real file task completes; no routing wrapper | | `test_ug_codex_headless_explicit_model_bypasses_routing` | Pass `--model VALUE` / `--model=VALUE` / `-m VALUE` with routing enabled | Real file task completes; no routing wrapper | | `test_ug_claude_preserves_caller_settings_and_hook` | Pass a settings path containing spaces | Real SessionStart hook executes; caller file unchanged; file task completes | @@ -59,9 +60,9 @@ All tests live directly in `integration/`; shared mechanics live in `utils/`. | `test_ug_and_ucode_auth_helpers_emit_only_the_supplied_bearer` | Run both auth helper commands with the public bearer override, with and without forced refresh | Exact token-only stdout, no warnings or ANSI escapes; no workspace authentication or saved state | | `test_ug_and_ucode_web_search_helpers_preserve_mcp_stdio` | Initialize and list tools through both web-search helper commands | Exactly the MCP JSON-RPC responses; no text/ANSI contamination; existing server/tool identities preserved; no model request | -With both agents selected there are **41 live cases** (6 interactive TUI cases), -**3 managed-workspace cases** (marker `managed`, run against a separate workspace that -publishes a CodingAgentConfig), **5 managed-fixture cases** (marker `managed_fixture`, with only +With both agents selected there are **43 live cases** (6 interactive TUI cases), +**4 managed-workspace cases** (marker `managed`, run against a separate workspace that +publishes a CodingAgentConfig), **11 managed-fixture cases** (marker `managed_fixture`, with only the CodingAgentConfig input injected), and **5 installation checks**. Parametrization varies argument spelling or routing mode, never hides the agent/provider in the test name. Duplicate boot-only cases are incorporated into the Databricks configuration TUI journeys. @@ -85,7 +86,7 @@ dependency graph to reproduce a user's combination. Every relevant same-reposito PR and push to `main` runs both smoke and the full CUJ suite. Smoke covers the Databricks Hosted configure/TUI, custom OAuth CLI TUI, and headless argument journeys for both agents, in two parallel jobs. After smoke finishes, the full -suite runs all 41 cases across two parallel agent jobs: one Claude VM and one +suite runs all 43 cases across two parallel agent jobs: one Claude VM and one Codex VM, each running its configure, headless, and commands/lifecycle cases serially. Each agent is installed once for the full suite, and no two full jobs for the same agent overlap within a run. @@ -116,7 +117,7 @@ pending. The descriptive jobs provide the actual coverage and diagnostics. | Scenario | Status / requirement | | --- | --- | | Live MCP and skills functionality | Deferred; installation tests cover the local web-search MCP handshake and tool listing, not upstream proxying or a real search request | -| Broad configure flags, tracing, multiple workspaces, and PAT flows | Deferred while focusing on basic CUJs | +| Broad configure flags, multiple-workspace lifecycle, and PAT flows | Deferred while focusing on basic CUJs | | Provider switching, relayed/subscription MPS | Not covered by the four provider journeys | | TUI initial prompt supplied on the launch command line | Not yet covered; headless prompt arguments are covered | | Follow-up turns and conversation resume | Not covered; reopen proves startup, not conversation resume | diff --git a/tests/integration/README.md b/tests/integration/README.md index c2bd85f4..4c232fcd 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -93,6 +93,7 @@ test_ug_codex_custom_oauth.py # CLI custom-OAuth launch and profile test_ug_claude_headless.py # script prompts, models, caller settings test_ug_claude_relayed.py # relayed session: subscription + Databricks-hosted models test_ug_codex_headless.py # script prompts and model arguments +test_ug_codex_tracing.py # Codex OTLP export reaches the configured trace table test_ug_claude_commands.py # command help forwarding test_ug_codex_commands.py # command help and parser error forwarding test_ug_codex_app_server.py # actual client/server initialize exchange @@ -156,10 +157,21 @@ service allows a different model. Those choices are recorded in `versions.json`. No service is created or modified. A missing service, permission, or OAuth token fails the selected CUJ, rather than skipping it. -There are **42 live cases** (including 6 TUI journeys) and **5 installation -checks** with both agents. A separate **3 managed-workspace cases** (one per agent -plus an idempotent re-configure, marker `managed`) run against a workspace that publishes a CodingAgentConfig; see -"Managed-workspace journeys" below. A further **9 `managed_fixture` cases** inject the admin config +The Codex tracing journey is part of the Full Codex lane and uses its existing e2e +workspace and bearer. Because that workspace deliberately has no published managed +configuration, the journey injects only a tracing-enabled CodingAgentConfig input through +the suite's `managed_fixture` mechanism; Codex, inference, OTLP export, and table verification +remain real. It adds the prompt's UUID as the trace-safe `ug_integration_marker` span +attribute, waits 30 seconds, and queries +`main.aigw_tracing.unity_gateway_otel_spans` through an existing SQL warehouse in the +workspace. It asserts both that a span with the marker arrived and that the same span +carries the `model` attribute for the model that ran, so the trace is attributable to a +specific model. + +There are **43 live cases** (including 6 TUI journeys) and **5 installation +checks** with both agents. A separate **4 managed-workspace cases** (one per agent, +an idempotent re-configure, and a cache-reuse case; marker `managed`) run against a workspace that publishes a CodingAgentConfig; see +"Managed-workspace journeys" below. A further **11 `managed_fixture` cases** inject the admin config locally (via `UCODE_MANAGED_CONFIG_STUB`) to cover shapes the live workspace does not publish, including a managed MCP server landing in Codex's OS-managed `[mcp_servers]` (interactive configure) while the developer's own config stays untouched, and reaching Claude's `/mcp` view via the @@ -178,6 +190,7 @@ See the named coverage and gaps matrix in -- -m live # default: all live user journeys -- -m smoke # six Hosted, custom OAuth CLI TUI, and headless journeys -- -m tui # six interactive TUI journeys +-- -m 'live and codex' -k trace # installed Codex -> gateway -> configured trace table -- -k test_ug_codex_app_server_client_initializes # one named journey and its variants # Use --installation-only before -- for package checks without credentials. ``` @@ -186,7 +199,7 @@ The old focused checks are now descriptive CUJs with setup and outcomes visible in each test. Duplicate boot-only checks are incorporated into the Databricks configuration TUI journeys. Real failures, including generated config left after revert and banners on app-server stdout, remain assertions. -Live MCP/skills functionality, tracing, the broad configure-option matrix, and other +Live MCP/skills functionality, the broad configure-option matrix, and other agents are outside this focused revision. The configure terminal helper recognizes `[✓]` / `[ ]` agent checkboxes as well @@ -248,13 +261,13 @@ record a discovered `system.ai` model as a test argument. Every same-repository PR and push to `main` runs **Smoke journeys**, followed by **Full journeys** even if smoke fails. Smoke runs the Hosted configure/TUI, headless argument, and custom OAuth CLI TUI journeys for each agent (six cases, -two agent jobs). Full runs all 42 live cases, including those smoke cases, in two +two agent jobs). Full runs all 43 live cases, including those smoke cases, in two disjoint agent lanes: | Agent lane | Marker | Cases | | --- | --- | --- | | Claude | `live and claude` | 17 | -| Codex | `live and codex` | 25 | +| Codex | `live and codex` | 26 | Each lane installs only its agent CLI, once, and runs all its configure, headless, commands, lifecycle, and applicable app-server journeys. Cases remain serial @@ -268,7 +281,7 @@ No test retries or assertion changes compensate for capacity failures. Both matrices use `fail-fast: false` and upload uniquely named evidence even when the other agent fails. The **All integration tests** check requires installation, workspace validation, smoke, and -both full lanes to pass. The **Managed config** lanes run for signal but are temporarily +both full lanes to pass; the tracing journey is included in the Full Codex lane. The **Managed config** lanes run for signal but are temporarily non-blocking (`continue-on-error`): the managed workspace is now runner-reachable, but the lanes stay non-blocking until the managed-config apply path is proven stable. They neither fail the workflow nor gate merges until then. The @@ -494,7 +507,7 @@ uv run --no-project --python 3.12 python scripts/run_integration.py \ unset DATABRICKS_BEARER ``` -This runs all 42 live cases. For the five installation checks, run the same +This runs all 43 live cases. For the five installation checks, run the same runner/version/index arguments with `--installation-only` and omit `-- -m live`; no bearer or workspace is needed. Results remain under `.integration-runs/`. Each invocation needs a new output directory; an existing one is rejected. diff --git a/tests/integration/test_ug_codex_tracing.py b/tests/integration/test_ug_codex_tracing.py new file mode 100644 index 00000000..34ef63eb --- /dev/null +++ b/tests/integration/test_ug_codex_tracing.py @@ -0,0 +1,94 @@ +"""Installed-product customer journey for Codex OpenTelemetry export.""" + +import os +import time +import uuid + +import pytest +from utils.constants import CODEX_TEST_MODEL +from utils.evidence import FileTask +from utils.managed import ( + build_codex_agent_config, + build_coding_agent_config, + set_managed_config_stub, +) +from utils.sql import query_count, resolve_warehouse_id + +pytestmark = [pytest.mark.live, pytest.mark.managed_fixture, pytest.mark.codex] + + +def test_ug_codex_exports_trace_to_configured_table(live_session, workspace, tmp_path): + """Scenario: configure Codex with tracing enabled and run a task with a unique marker. + + Expected: the real agent task completes and, after the ingestion window, the + configured trace table contains a Codex span carrying the same marker. + """ + session = live_session + marker = f"ug-codex-trace-{uuid.uuid4().hex}" + task = FileTask(session) + config = build_coding_agent_config( + "CODING_AGENT_CODEX", + build_codex_agent_config(models=[CODEX_TEST_MODEL], otel_tracing_enabled=True), + ) + set_managed_config_stub(session, tmp_path, config) + session.run( + "configure", + "--workspace", + workspace, + "--skip-validate", + "--skip-upgrade", + "--disable-databricks-ai-tools", + ) + assert session.workspace_state().get("codex_otel_tracing") is True + + result = session.run( + "codex", + "--", + "--config", + f'otel.span_attributes.ug_integration_marker="{marker}"', + "exec", + "--skip-git-repo-check", + "--json", + "--model", + CODEX_TEST_MODEL, + f"{task.prompt} Trace correlation marker: {marker}", + timeout=180, + ) + task.assert_headless_answer("codex", result) + + time.sleep(30) + table = os.environ.get("UG_INTEGRATION_TRACE_TABLE", "").strip() + assert table, "Pass --trace-table for the configured tracing table" + warehouse_id = os.environ.get("UG_INTEGRATION_WAREHOUSE_ID", "").strip() + warehouse_id = warehouse_id or resolve_warehouse_id(workspace, session.env["DATABRICKS_BEARER"]) + bearer = session.env["DATABRICKS_BEARER"] + marker_query = ( + f"SELECT COUNT(*) FROM {table} " + "WHERE time > current_timestamp() - INTERVAL 10 MINUTES " + "AND variant_get(attributes, '$[\"ug_integration_marker\"]', 'STRING') = :marker" + ) + count = query_count( + workspace, + bearer, + warehouse_id, + marker_query, + [{"name": "marker", "value": marker, "type": "STRING"}], + ) + # The span must also carry Codex's `model` attribute matching the model that ran, + # so the trace is attributable to a specific model and not just to this test run. + model_count = query_count( + workspace, + bearer, + warehouse_id, + marker_query + " AND variant_get(attributes, '$[\"model\"]', 'STRING') = :model", + [ + {"name": "marker", "value": marker, "type": "STRING"}, + {"name": "model", "value": CODEX_TEST_MODEL, "type": "STRING"}, + ], + ) + session.record( + "trace-query.json", + {"marker": marker, "table": table, "count": count, "model_count": model_count}, + ) + assert count > 0 + assert model_count > 0, f"Codex span for {marker} lacked model={CODEX_TEST_MODEL}" diff --git a/tests/integration/utils/managed.py b/tests/integration/utils/managed.py index 60b7c521..cb83d8d8 100644 --- a/tests/integration/utils/managed.py +++ b/tests/integration/utils/managed.py @@ -54,13 +54,18 @@ def build_claude_agent_config( } -def build_codex_agent_config(*, models: list[str]) -> dict: +def build_codex_agent_config( + *, models: list[str], otel_tracing_enabled: bool | None = None +) -> dict: + config = { + "models": {"model_services": models}, + "default_models": {"default_model": models[0]}, + } + if otel_tracing_enabled is not None: + config["tracing"] = {"enabled": otel_tracing_enabled} return { "agent": "CODING_AGENT_CODEX", - "config": { - "models": {"model_services": models}, - "default_models": {"default_model": models[0]}, - }, + "config": config, } diff --git a/tests/integration/utils/sql.py b/tests/integration/utils/sql.py new file mode 100644 index 00000000..f37bce14 --- /dev/null +++ b/tests/integration/utils/sql.py @@ -0,0 +1,68 @@ +"""Query a Databricks SQL warehouse through the public Statement Execution API.""" + +from __future__ import annotations + +import json +import time +import urllib.request + + +def resolve_warehouse_id(workspace: str, bearer: str) -> str: + """Choose an existing warehouse, preferring one that is already running.""" + request = urllib.request.Request( + f"{workspace.rstrip('/')}/api/2.0/sql/warehouses", + headers={"Authorization": f"Bearer {bearer}"}, + ) + with urllib.request.urlopen(request, timeout=30) as response: # noqa: S310 + result = json.load(response) + + warehouses = [ + warehouse + for warehouse in result.get("warehouses", []) + if isinstance(warehouse, dict) and warehouse.get("id") + ] + assert warehouses, "The integration workspace has no SQL warehouse" + running = next( + (warehouse for warehouse in warehouses if warehouse.get("state") == "RUNNING"), None + ) + return str((running or warehouses[0])["id"]) + + +def query_count( + workspace: str, + bearer: str, + warehouse_id: str, + statement: str, + parameters: list[dict[str, str]], +) -> int: + request = urllib.request.Request( + f"{workspace.rstrip('/')}/api/2.0/sql/statements", + data=json.dumps( + { + "warehouse_id": warehouse_id, + "statement": statement, + "parameters": parameters, + "wait_timeout": "50s", + "on_wait_timeout": "CONTINUE", + } + ).encode(), + headers={"Authorization": f"Bearer {bearer}", "Content-Type": "application/json"}, + ) + with urllib.request.urlopen(request, timeout=60) as response: # noqa: S310 + result = json.load(response) + + deadline = time.monotonic() + 180 + while result.get("status", {}).get("state") in {"PENDING", "RUNNING"}: + assert time.monotonic() < deadline, "SQL statement did not finish within 180 seconds" + time.sleep(2) + poll = urllib.request.Request( + f"{workspace.rstrip('/')}/api/2.0/sql/statements/{result['statement_id']}", + headers={"Authorization": f"Bearer {bearer}"}, + ) + with urllib.request.urlopen(poll, timeout=30) as response: # noqa: S310 + result = json.load(response) + + assert result.get("status", {}).get("state") == "SUCCEEDED", result.get("status") + rows = result.get("result", {}).get("data_array", []) + assert rows and rows[0], "SQL count query returned no row" + return int(rows[0][0])