Parallelize smoke, full journeys, and agent CI

.github/workflows/ci.yml

@@ -9,8 +9,13 @@ on:
 permissions:
   contents: read
 
+concurrency:
+  group: ci-${{ github.event.pull_request.number || github.ref }}-${{ github.event_name }}
+  cancel-in-progress: true
+
 jobs:
   test:
+    name: Unit tests
     runs-on: ubuntu-latest
     env:
       # uv.lock pins packages to Databricks' internal pypi-proxy, which hosted
@@ -23,8 +28,41 @@ jobs:
       - run: uv lock
       - run: uv run pytest --ignore=tests/test_e2e.py --ignore=tests/test_e2e_tracing.py
 
-  e2e:
+  e2e-shards:
+    name: ${{ matrix.name }}
     runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - group: gateway
+            name: Gateway API tests
+            keyword: not (TestCodexLaunch or TestClaudeLaunch or TestConfigureSubset or TestModelProviderLaunch or TestGeminiLaunch or TestGeminiFreshInstall or TestGeminiAuthRecovery or TestOpencodeLaunch or TestCopilotLaunch or TestPiLaunch)
+            package: ''
+          - group: claude
+            name: Agent launch tests · Claude
+            keyword: TestClaudeLaunch or TestConfigureSubset or (TestModelProviderLaunch and not codex)
+            package: '@anthropic-ai/claude-code'
+          - group: codex
+            name: Agent launch tests · Codex
+            keyword: TestCodexLaunch or (TestModelProviderLaunch and codex)
+            package: '@openai/codex'
+          - group: gemini
+            name: Agent launch tests · Gemini
+            keyword: TestGeminiLaunch or TestGeminiFreshInstall or TestGeminiAuthRecovery
+            package: '@google/gemini-cli'
+          - group: opencode
+            name: Agent launch tests · OpenCode
+            keyword: TestOpencodeLaunch
+            package: opencode-ai@1
+          - group: copilot
+            name: Agent launch tests · Copilot
+            keyword: TestCopilotLaunch
+            package: '@github/copilot@1.0.80'
+          - group: pi
+            name: Agent launch tests · Pi
+            keyword: TestPiLaunch
+            package: '@earendil-works/pi-coding-agent'
     env:
       # See the test job: hosted runners can't reach the internal pypi-proxy,
       # so resolve against public PyPI (paired with the `uv lock` step below).
@@ -42,34 +80,75 @@ jobs:
       DATABRICKS_BEARER: ${{ secrets.DATABRICKS_BEARER }}
       # Subscription OAuth token for the relayed-launch e2e test; the test skips when unset.
       CLAUDE_CODE_OAUTH_TOKEN: ${{ secrets.CLAUDE_CODE_OAUTH_TOKEN }}
+      AGENT_PACKAGE: ${{ matrix.package }}
+      TEST_KEYWORD: ${{ matrix.keyword }}
     steps:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
       - uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020  # v4
       - uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b  # v8.1.0
       - uses: databricks/setup-cli@bdb89f81c11a5bd647fd55b585b7c396ec68a25a  # v1.0.0
-      # The agent launch tests `_require_binary("codex")` etc. and skip when
-      # the CLI isn't on PATH. Install all six so each TestXxxLaunch test
-      # actually runs instead of skipping.
-      # Pin Copilot because 1.0.81-6 sends unsupported parameters through BYOK providers.
-      - name: Install agent CLIs
-        run: npm install -g
-          @anthropic-ai/claude-code
-          @openai/codex
-          @google/gemini-cli
-          opencode-ai@1
-          @github/copilot@1.0.80
-          @earendil-works/pi-coding-agent
+      - name: Install the shard's agent CLI
+        if: ${{ matrix.package != '' }}
+        run: npm install -g "$AGENT_PACKAGE"
       - run: uv lock
       - run: uv tool install .
       # Redirect stdin so any interactive `databricks auth login --no-browser`
       # fallback EOFs instead of hanging the runner. With DATABRICKS_BEARER
       # set, the auth code path doesn't shell out at all — this is a safety
       # net for any code path we may have missed.
-      - run: uv run pytest tests/test_e2e.py -v < /dev/null
+      - run: uv run pytest tests/test_e2e.py -v -k "$TEST_KEYWORD" < /dev/null
       # MLflow tracing e2e lives in its own file and needs the `tracing`
       # extra so `import mlflow` resolves (otherwise the test importorskips
       # and silently passes as skipped). `!cancelled()` lets this run even
       # when the previous pytest step failed — the two suites are
       # independent and one shouldn't mask the other.
-      - if: ${{ !cancelled() }}
+      - if: ${{ !cancelled() && matrix.group == 'claude' }}
         run: uv run --extra tracing pytest tests/test_e2e_tracing.py -v < /dev/null
+
+  e2e:
+    name: All agent tests
+    if: ${{ !cancelled() }}
+    needs: e2e-shards
+    runs-on: ubuntu-latest
+    timeout-minutes: 5
+    steps:
+      - name: Require every e2e shard to pass
+        env:
+          RESULT: ${{ needs.e2e-shards.result }}
+        run: test "$RESULT" = success
+
+  integration:
+    name: Integration
+    # These suites use the same live workspace. Let agent launch coverage
+    # finish before starting integration requests, even if an agent failed.
+    if: ${{ !cancelled() }}
+    needs: e2e
+    uses: ./.github/workflows/integration.yml
+    secrets: inherit
+
+  # Preserve the exact contexts required by the repository's branch rules.
+  # Keep these aliases until the rules and all active branches migrate together.
+  required-tests:
+    name: test
+    if: ${{ !cancelled() }}
+    needs: test
+    runs-on: ubuntu-latest
+    timeout-minutes: 5
+    steps:
+      - name: Require unit tests to pass
+        env:
+          RESULT: ${{ needs.test.result }}
+        run: test "$RESULT" = success
+
+  required-e2e:
+    name: e2e
+    if: ${{ !cancelled() }}
+    needs: [e2e, integration]
+    runs-on: ubuntu-latest
+    timeout-minutes: 5
+    steps:
+      - name: Require all agent and integration tests to pass
+        env:
+          RESULT: ${{ needs.e2e.result }}
+          INTEGRATION_RESULT: ${{ needs.integration.result }}
+        run: test "$RESULT" = success && test "$INTEGRATION_RESULT" = success

.github/workflows/integration.yml

@@ -0,0 +1,221 @@
+name: Integration
+
+on:
+  workflow_call:
+  workflow_dispatch:
+    inputs:
+      suite:
+        description: Which integration checks to run
+        type: choice
+        options: [full, smoke, live, tui, installation]
+        default: full
+      ug_version:
+        description: Exact ug release, or checkout
+        default: checkout
+        required: true
+      entry_point:
+        description: Console command (older releases may only have ucode)
+        type: choice
+        options: [ug, ucode]
+        default: ug
+      claude_version:
+        description: Exact Claude Code version
+        default: 2.1.268
+        required: true
+      codex_version:
+        description: Exact Codex version
+        default: 0.154.0
+        required: true
+      dependency:
+        description: Optional exact dependency for reproduction (e.g. tomlkit==0.14.0)
+        default: ''
+        required: false
+      default_index:
+        description: Python package index containing the requested ug version
+        default: https://pypi.org/simple
+        required: true
+
+permissions:
+  contents: read
+
+concurrency:
+  group: integration-${{ github.event.pull_request.number || github.ref }}-${{ github.event_name }}
+  cancel-in-progress: true
+
+env:
+  EXPECTED_WORKSPACE: https://eng-ml-inference-team-us-east-1.cloud.databricks.com
+  UG_VERSION: ${{ inputs.ug_version || 'checkout' }}
+  ENTRY_POINT: ${{ inputs.entry_point || 'ug' }}
+  CLAUDE_VERSION: ${{ inputs.claude_version || '2.1.268' }}
+  CODEX_VERSION: ${{ inputs.codex_version || '0.154.0' }}
+  PACKAGE_INDEX: ${{ inputs.default_index || 'https://pypi.org/simple' }}
+
+jobs:
+  installation:
+    name: Installation tests
+    runs-on: ubuntu-22.04
+    timeout-minutes: 20
+    steps:
+      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
+        with:
+          fetch-depth: 0
+      - uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020  # v4
+        with:
+          node-version: 22.19.0
+      - uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b  # v8.1.0
+        with:
+          version: 0.9.8
+      - name: Test a fresh installed package without credentials
+        run: |
+          uv run --no-project --python 3.12 python scripts/run_integration.py \
+            --python 3.12 --ug-version "$UG_VERSION" --entry-point "$ENTRY_POINT" \
+            --claude-version "$CLAUDE_VERSION" --codex-version "$CODEX_VERSION" \
+            --default-index "$PACKAGE_INDEX" --installation-only --output "$RUNNER_TEMP/ug-integration"
+      - name: Upload installation evidence
+        if: ${{ !cancelled() }}
+        uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02  # v4.6.2
+        with:
+          name: integration-installation
+          include-hidden-files: true
+          path: |
+            ${{ runner.temp }}/ug-integration/*.json
+            ${{ runner.temp }}/ug-integration/*.txt
+            ${{ runner.temp }}/ug-integration/*.xml
+            ${{ runner.temp }}/ug-integration/*.log
+            ${{ runner.temp }}/ug-integration/artifacts/
+            ${{ runner.temp }}/ug-integration/wheels/
+
+  workspace:
+    name: Workspace validation
+    if: ${{ inputs.suite != 'installation' && (github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository) }}
+    runs-on: ubuntu-latest
+    timeout-minutes: 5
+    steps:
+      - name: Verify the existing e2e workspace and credential are configured
+        env:
+          UCODE_TEST_WORKSPACE: ${{ secrets.UCODE_TEST_WORKSPACE }}
+          DATABRICKS_BEARER: ${{ secrets.DATABRICKS_BEARER }}
+        run: |
+          python3 - <<'PY'
+          import os
+          expected = os.environ["EXPECTED_WORKSPACE"].rstrip("/")
+          actual = os.environ["UCODE_TEST_WORKSPACE"].strip().rstrip("/")
+          if actual != expected:
+              raise SystemExit("UCODE_TEST_WORKSPACE does not match the expected CI workspace.")
+          if not os.environ["DATABRICKS_BEARER"].strip():
+              raise SystemExit("The existing e2e DATABRICKS_BEARER secret is missing.")
+          print("CI workspace matches; the existing e2e credential is present.")
+          PY
+
+  smoke:
+    name: Smoke journeys · ${{ matrix.agent == 'claude' && 'Claude' || 'Codex' }}
+    if: ${{ inputs.suite != 'tui' }}
+    needs: workspace
+    runs-on: ubuntu-22.04
+    timeout-minutes: 20
+    strategy:
+      fail-fast: false
+      matrix:
+        agent: [claude, codex]
+    env:
+      UCODE_TEST_WORKSPACE: ${{ secrets.UCODE_TEST_WORKSPACE }}
+      DATABRICKS_BEARER: ${{ secrets.DATABRICKS_BEARER }}
+      DEPENDENCY: ${{ inputs.dependency }}
+      AGENT: ${{ matrix.agent }}
+      TEST_MARKER: live and smoke and ${{ matrix.agent }}
+      TEST_KEYWORD: ''
+      ARTIFACT_NAME: integration-smoke-${{ matrix.agent }}
+    steps: &live-steps
+      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
+        with:
+          fetch-depth: 0
+      - uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020  # v4
+        with:
+          node-version: 22.19.0
+      - uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b  # v8.1.0
+        with:
+          version: 0.9.8
+      - uses: databricks/setup-cli@bdb89f81c11a5bd647fd55b585b7c396ec68a25a  # v1.0.0
+      - name: Run the selected live cases against the e2e workspace
+        shell: bash
+        run: |
+          case "$AGENT" in
+            claude) args=(--claude-version "$CLAUDE_VERSION") ;;
+            codex) args=(--codex-version "$CODEX_VERSION") ;;
+          esac
+          if [[ -n "$DEPENDENCY" ]]; then
+            args+=(--dependency "$DEPENDENCY")
+          fi
+          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 "$TEST_MARKER" -k "$TEST_KEYWORD"
+      - name: Upload live test evidence
+        if: ${{ !cancelled() }}
+        uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02  # v4.6.2
+        with:
+          name: ${{ env.ARTIFACT_NAME }}
+          include-hidden-files: true
+          path: |
+            ${{ runner.temp }}/ug-integration/*.json
+            ${{ runner.temp }}/ug-integration/*.txt
+            ${{ runner.temp }}/ug-integration/*.xml
+            ${{ runner.temp }}/ug-integration/*.log
+            ${{ runner.temp }}/ug-integration/artifacts/
+            ${{ runner.temp }}/ug-integration/wheels/
+
+  full:
+    name: Full journeys · ${{ matrix.agent == 'claude' && 'Claude' || 'Codex' }} · ${{ matrix.group == 'configure' && 'Configure' || matrix.group == 'headless' && 'Headless' || 'Commands' }}
+    if: ${{ !cancelled() && inputs.suite != 'smoke' && needs.workspace.result == 'success' }}
+    needs: [workspace, smoke]
+    runs-on: ubuntu-22.04
+    # The runner enforces its own 3600s pytest deadline; keep the job above it
+    # so a slow run fails through the runner's report instead of a job kill.
+    timeout-minutes: 70
+    strategy:
+      fail-fast: false
+      # All shards share one workspace/model quota. Serial execution retains
+      # every case without a burst of overlapping model requests.
+      max-parallel: 1
+      matrix:
+        agent: [claude, codex]
+        group: ${{ fromJSON(inputs.suite == 'tui' && '["configure"]' || '["configure", "headless", "commands"]') }}
+    env:
+      UCODE_TEST_WORKSPACE: ${{ secrets.UCODE_TEST_WORKSPACE }}
+      DATABRICKS_BEARER: ${{ secrets.DATABRICKS_BEARER }}
+      DEPENDENCY: ${{ inputs.dependency }}
+      AGENT: ${{ matrix.agent }}
+      TEST_MARKER: ${{ matrix.group == 'configure' && 'live and tui' || 'live and not tui' }} and ${{ matrix.agent }}
+      TEST_KEYWORD: ${{ matrix.group == 'configure' && 'configure' || matrix.group == 'headless' && 'headless' || 'not headless' }}
+      ARTIFACT_NAME: integration-full-${{ matrix.agent }}-${{ matrix.group }}
+    steps: *live-steps
+
+  cujs:
+    name: All integration tests
+    if: ${{ !cancelled() && (github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository) }}
+    needs: [installation, workspace, smoke, full]
+    runs-on: ubuntu-latest
+    timeout-minutes: 5
+    steps:
+      - name: Require every selected integration job to pass
+        env:
+          RESULTS: ${{ toJSON(needs) }}
+          SUITE: ${{ inputs.suite || 'full' }}
+        run: |
+          python3 - <<'PY'
+          import json
+          import os
+          results = json.loads(os.environ["RESULTS"])
+          suite = os.environ["SUITE"]
+          required = ["installation"]
+          if suite != "installation":
+              required.append("workspace")
+              if suite != "tui":
+                  required.append("smoke")
+              if suite != "smoke":
+                  required.append("full")
+          failed = [job for job in required if results[job]["result"] != "success"]
+          if failed:
+              raise SystemExit("Integration jobs did not pass: " + ", ".join(failed))
+          print("All selected integration jobs passed: " + ", ".join(required))
+          PY