test(e2e): cover MCP OAuth SSO and cold restart acceptance

2026-09-21 00:21:49 +00:00 · 2026-09-18 22:02:09 -07:00 · 2026-09-18 22:02:09 -07:00 · 7f4dd4eabc
commit 7f4dd4eabc
parent 62051ad9bc
13 changed files with 664 additions and 107 deletions
--- a/.github/e2e-stack/assert_tests_ran.py
+++ b/.github/e2e-stack/assert_tests_ran.py
@ -1,3 +1,4 @@
+import os
 import sys
 import xml.etree.ElementTree as ET
 from pathlib import Path
@ -15,6 +16,12 @@ def main() -> int:
        _ = sys.stdout.write("::error::could not read the test execution report\n")
        return 1
    cases: Final = tuple(report.iter("testcase"))
+    expected_count: Final = os.environ.get("E2E_REQUIRED_TEST_COUNT")
+    if expected_count is not None and (
+        len(cases) != int(expected_count) or any(case.find("skipped") is not None for case in cases)
+    ):
+        _ = sys.stdout.write("::error::required test count was not met or a required case was skipped\n")
+        return 1
    passed: Final = frozenset(
        case.get("file") for case in cases if all(case.find(tag) is None for tag in ("skipped", "failure", "error"))
    )
--- a/.github/e2e-stack/select_tests.py
+++ b/.github/e2e-stack/select_tests.py
@ -5,6 +5,7 @@ from typing import Final
 SELECTABLE: Final = re.compile(r"^tests/e2e/([A-Za-z0-9_.-]+/)*test_[A-Za-z0-9_.-]+\.py$")
 UNSUPPORTED: Final = re.compile(
    r"^tests/e2e/(ui|claude_code|load)/"
+    r"|^tests/e2e/mcp/test_mcp_oauth_happy_path_e2e\.py$"
    r"|^tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e\.py$"
    r"|^tests/e2e/batches/test_managed_files_enforcement_e2e\.py$"
    r"|^tests/e2e/guardrails/test_presidio_masking_e2e\.py$"
--- a/.github/workflows/test-mcp-oauth-e2e.yml
+++ b/.github/workflows/test-mcp-oauth-e2e.yml
@ -0,0 +1,168 @@
+name: MCP OAuth happy path
+
+on:
+  pull_request:
+    paths:
+      - tests/e2e/idp.py
+      - tests/e2e/provider_edge.py
+      - tests/e2e/models.py
+      - tests/e2e/conftest.py
+      - .github/e2e-stack/assert_tests_ran.py
+      - tests/e2e/mcp/oauth_chat_client.py
+      - tests/e2e/mcp/oauth_gateway.py
+      - tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py
+      - .github/workflows/test-mcp-oauth-e2e.yml
+  workflow_dispatch:
+
+permissions: {}
+
+concurrency:
+  group: mcp-oauth-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  oauth:
+    if: github.event_name == 'workflow_dispatch' || github.event.pull_request.head.repo.full_name == github.repository
+    runs-on: ubuntu-latest
+    environment: e2e-changed
+    timeout-minutes: 45
+    permissions:
+      contents: read
+      id-token: write
+    services:
+      postgres:
+        image: postgres:16.6
+        env:
+          POSTGRES_USER: litellm
+          POSTGRES_PASSWORD: dbpassword9090
+          POSTGRES_DB: litellm
+        ports:
+          - 5432:5432
+        options: >-
+          --health-cmd "pg_isready -U litellm"
+          --health-interval 5s
+          --health-timeout 5s
+          --health-retries 10
+    env:
+      DATABASE_HOST: 127.0.0.1
+      DATABASE_PORT: '5432'
+      DATABASE_USER: litellm
+      DATABASE_PASSWORD: dbpassword9090
+      DATABASE_NAME: litellm
+      DATABASE_URL: postgresql://litellm:dbpassword9090@127.0.0.1:5432/litellm
+      E2E_KEYCLOAK_URL: http://127.0.0.1:8081
+      E2E_KEYCLOAK_ADMIN_USER: admin
+      E2E_KEYCLOAK_ADMIN_PASSWORD: e2e-ephemeral-idp-not-a-secret
+      E2E_FIXTURE_MODE: live
+      E2E_PROVIDER_CACHE: '0'
+      E2E_MCP_OAUTH_LIVE: '1'
+      E2E_REQUIRED_TEST_COUNT: '4'
+    steps:
+      - name: Checkout the tested source
+        uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
+        with:
+          ref: ${{ github.event.pull_request.head.sha || github.sha }}
+          persist-credentials: false
+
+      - name: Require and materialize the upstream login
+        env:
+          STORAGE_STATE: ${{ secrets.E2E_LINEAR_STORAGE_STATE_B64 }}
+        run: |
+          umask 077
+          python3 - <<'PY'
+          import base64
+          import json
+          import os
+          import secrets
+          from pathlib import Path
+          encoded = os.environ.get("STORAGE_STATE", "")
+          if not encoded:
+              raise SystemExit("E2E_LINEAR_STORAGE_STATE_B64 is required; capture and provision a test-account login")
+          state = json.loads(base64.b64decode(encoded, validate=True))
+          if not isinstance(state, dict) or not state.get("cookies"):
+              raise SystemExit("The captured login must contain browser cookies")
+          directory = Path(os.environ["RUNNER_TEMP"]) / "mcp-oauth-private"
+          directory.mkdir(mode=0o700)
+          path = directory / "linear-state.json"
+          path.write_text(json.dumps(state))
+          with open(os.environ["GITHUB_ENV"], "a") as output:
+              output.write(f"E2E_LINEAR_STORAGE_STATE={path}\n")
+              for name in ("LITELLM_MASTER_KEY", "LITELLM_SALT_KEY"):
+                  value = "sk-e2e-" + secrets.token_hex(24)
+                  print(f"::add-mask::{value}")
+                  output.write(f"{name}={value}\n")
+          PY
+
+      - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
+        with:
+          python-version: '3.13'
+      - uses: ./.github/actions/setup-uv-with-retries
+        with:
+          version: '0.10.9'
+      - uses: ./.github/actions/cache-cargo-build
+      - name: Install the frozen E2E environment
+        run: |
+          .github/scripts/uv_sync_with_retries.sh --frozen --extra proxy --extra proxy-runtime --extra extra_proxy --group ci --group proxy-dev --group e2e-dev
+          uv run --no-sync python scripts/prisma_generate_if_needed.py
+          uv run --no-sync playwright install --with-deps chromium
+
+      - name: Configure license access
+        id: aws
+        uses: aws-actions/configure-aws-credentials@e7f100cf4c008499ea8adda475de1042d6975c7b # v6.2.0
+        with:
+          role-to-assume: ${{ vars.E2E_AWS_ROLE_TO_ASSUME }}
+          aws-region: us-east-1
+          role-session-name: mcp-oauth-${{ github.run_id }}
+          role-duration-seconds: 900
+          output-env-credentials: false
+          output-credentials: true
+      - name: Load the E2E license
+        env:
+          AWS_ACCESS_KEY_ID: ${{ steps.aws.outputs.aws-access-key-id }}
+          AWS_SECRET_ACCESS_KEY: ${{ steps.aws.outputs.aws-secret-access-key }}
+          AWS_SESSION_TOKEN: ${{ steps.aws.outputs.aws-session-token }}
+          AWS_DEFAULT_REGION: us-east-1
+        run: |
+          license="$(aws secretsmanager get-secret-value --secret-id litellm-e2e-changed-license --query SecretString --output text)"
+          test -n "${license}"
+          echo "::add-mask::${license}"
+          echo "LITELLM_LICENSE=${license}" >> "${GITHUB_ENV}"
+
+      - uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020 # v4.4.0
+        with:
+          node-version-file: ui/litellm-dashboard/.nvmrc
+      - name: Build the gateway consent UI at the tested commit
+        run: |
+          cd ui/litellm-dashboard
+          ../../scripts/with_dashboard_node.sh npm ci
+          ../../scripts/with_dashboard_node.sh npm run build
+          mkdir -p ../../litellm/proxy/_experimental/out
+          cp -r out/. ../../litellm/proxy/_experimental/out/
+          find ../../litellm/proxy/_experimental/out -name '*.html' ! -name index.html | while read -r page; do
+            mkdir -p "${page%.html}"
+            mv "${page}" "${page%.html}/index.html"
+          done
+
+      - name: Prepare the isolated database and IdP
+        run: |
+          umask 077
+          bash .github/e2e-stack/start-idp.sh
+          uv run --no-sync python migrations/run.py > "${RUNNER_TEMP}/mcp-oauth-private/migrations.log" 2>&1
+
+      - name: Run every required OAuth variant without retries
+        run: |
+          umask 077
+          uv run --no-sync pytest -c tests/e2e/pytest.ini tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py \
+            --rootdir=. --reruns 0 --tb=short -o junit_family=xunit1 \
+            --junitxml="${RUNNER_TEMP}/mcp-oauth-private/results.xml" \
+            > "${RUNNER_TEMP}/mcp-oauth-private/pytest.log" 2>&1
+      - name: Reject skipped or missing cases
+        if: always()
+        run: |
+          uv run --no-sync python .github/e2e-stack/assert_tests_ran.py \
+            "${RUNNER_TEMP}/mcp-oauth-private/results.xml" tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py
+      - name: Remove private login and logs
+        if: always()
+        run: |
+          docker rm -f e2e-keycloak >/dev/null 2>&1 || true
+          rm -rf "${RUNNER_TEMP}/mcp-oauth-private"
--- a/tests/e2e/AGENTS.md
+++ b/tests/e2e/AGENTS.md
@ -33,7 +33,7 @@ Every test under `tests/e2e/mcp/` must exercise the proxy against the real Datad
 - Prefer calling real Datadog tools that prove the product path (e.g. `search_datadog_logs` for list/call and permission denials). Seed a unique marker (`e2e-datadog-mcp-*`) in a chat completion when you need a log the tool can find; dual-read with `dd_logs` from conftest when delivery matters
 - Delete the MCP server (and any keys) through `resources.defer` the same way every other suite tears down
 - If a new MCP behavior cannot be covered with Datadog's tool surface, say so in the PR and get agreement before inventing another upstream; the default is always Datadog
- The two standing exceptions are `test_mcp_chat_completion_oauth_e2e.py` and `test_mcp_oauth_happy_path_e2e.py`. Datadog authenticates with the static `DD-API-KEY` / `DD-APPLICATION-KEY` headers and exposes no authorize/token dance at all, so these tests drive a real Linear MCP server instead; they are still real remote upstreams, so the no-mock, no-fixture rule above holds unchanged
+- The two standing exceptions are `test_mcp_chat_completion_oauth_e2e.py` and `test_mcp_oauth_happy_path_e2e.py`. Datadog authenticates with the static `DD-API-KEY` / `DD-APPLICATION-KEY` headers and exposes no authorize/token dance at all, so these tests drive a real Linear MCP server instead; they are still real remote upstreams, so the no-mock, no-fixture rule above holds unchanged. The direct OAuth test also uses the existing live provider edge to inspect forwarded headers without replay, and owns a separate source-built gateway for cold restarts

 ## Lay the pattern down in a class

--- a/tests/e2e/CONTRIBUTING.md
+++ b/tests/e2e/CONTRIBUTING.md
@ -248,3 +248,49 @@ The semantic header set is `content-type`, `accept`, `anthropic-version`, `anthr
 Excluded transport and telemetry headers are `host`, `content-length`, `connection`, `accept-encoding`, `user-agent`, `traceparent`, `tracestate`, `x-request-id`, `x-client-request-id` and `x-stainless-*`. Inbound transfer-encoding is unsupported; send JSON with content-length framing. The destination represents host identity and the relay carries original body bytes. Replay does not verify credentials, SDK timeout/retry behavior, transport performance, model availability or stateful remote IDs. Live relay uses original request bytes and header values, never the stored identity

 Strict replay harness regression tests live in `tests/code_coverage_tests/test_provider_replay_harness.py`. The CircleCI `provider_replay_harness` job runs them alongside the existing legacy harness files with `--noconftest -o pythonpath=tests/e2e`; they need only synthetic HTTP providers and temporary fixture storage
+
+
+## MCP OAuth happy path
+
+`test_mcp_oauth_happy_path_e2e.py` runs one shared scenario with four variants:
+aggregate gateway SSO and explicitly configured per-server JWT, each directly
+against Linear and through the live provider edge. The edge forwards to real
+Linear without replay and compares the forwarded bearer to the encrypted
+canonical user/server credential. This observes the forwarding boundary, not
+Linear's internal logs. Direct variants independently exercise discovery
+
+Use the existing database preparation, Prisma generation and Keycloak setup.
+Build and stage the dashboard from the tested checkout as in the UI runner.
+Provide `DATABASE_URL`, `LITELLM_MASTER_KEY`, `LITELLM_SALT_KEY`, `LITELLM_LICENSE`,
+and the `E2E_KEYCLOAK_*` settings. Capture a test-account Linear login using
+`mcp/linear_session_capture.py` and set `E2E_LINEAR_STORAGE_STATE` to that private
+file. The test workspace must contain a team. Do not publish browser state or
+raw test/proxy output
+
+```bash
+E2E_MCP_OAUTH_LIVE=1 E2E_FIXTURE_MODE=live E2E_PROVIDER_CACHE=0 \
+  uv run --no-sync pytest tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py \
+  --rootdir=. --reruns 0
+```
+
+The test starts and restarts its own source-built proxy on a free loopback port,
+retaining its database and SSO client but no Redis or process-local cache. It
+does not restart an existing proxy or clear shared databases. Gateway login,
+consent, immediate list/call and post-restart reconnect must all succeed. The
+aggregate client never injects a gateway header; the explicitly labeled JWT
+variant configures `x-litellm-api-key` for the first consent and reconnects with
+only its gateway JWT after restart
+
+`.github/workflows/test-mcp-oauth-e2e.yml` runs the four cases in the protected
+`e2e-changed` environment. Provision `E2E_LINEAR_STORAGE_STATE_B64` as a secret
+there and retain the existing E2E license/AWS role configuration. A missing or
+expired session fails the job; collection, deselection and skips are not passes.
+The generic changed-test job excludes this file because it requires an owned
+proxy and consent UI. No LLM call is needed
+
+Coverage remains limited to authorization-code OAuth over HTTP. M2M, OBO,
+PKCE passthrough, static/BYOK, ID-JAG, forwarding, SigV4 and stdio are outside this
+scenario; consult the registry and LIT-3559 for their existing coverage and gaps.
+LIT-4506 owns broader isolation/failure regressions. LIT-7737 retains ownership
+of dependency/Python compatibility and its matrix; this test reuses its delivered
+environment and does not change dependency constraints or compatibility gates
--- a/tests/e2e/conftest.py
+++ b/tests/e2e/conftest.py
@ -220,6 +220,8 @@ def pytest_runtest_setup(item: pytest.Item) -> None:
    LIVE_PROVIDER_REQUIRED.set(item.get_closest_marker("provider_live") is not None)
    if item.get_closest_marker("e2e") is None:
        return
+    if isinstance(item, pytest.Function) and "oauth_gateway" in item.fixturenames:
+        return
    reason = _proxy_fail_reason()
    if reason is not None:
        pytest.fail(reason)
--- a/tests/e2e/coverage_registry/mcp.yaml
+++ b/tests/e2e/coverage_registry/mcp.yaml
@ -78,7 +78,7 @@
  auth_family: oauth
  assertions: [persists_across_processes]
  source: "outbound_credentials/per_user_oauth_store.py V2PerUserTokenStore"
-  rationale: Stored per-user token is resolved by a gateway process that did not run the consent
+  rationale: Stored per-user token survives a verified restart of an owned gateway with no Redis cache
 - id: mcp.list_tools.none.succeeds
  module: mcp
  tier: P1
--- a/tests/e2e/idp.py
+++ b/tests/e2e/idp.py
@ -435,7 +435,7 @@ def _signal_process_group(process_id: int, signum: int) -> bool:
    return True


-def _stop_process_group(child: subprocess.Popen[bytes]) -> None:
+def stop_process_group(child: subprocess.Popen[bytes]) -> None:
    _signal_process_group(child.pid, signal.SIGTERM)
    deadline: Final = time.monotonic() + 5
    while _process_group_exists(child.pid):
@ -476,7 +476,7 @@ def run_oidc_profile(proxy_url: str, command: list[str]) -> int:
            try:
                return child.wait()
            finally:
-                _stop_process_group(child)
+                stop_process_group(child)


 if __name__ == "__main__":
--- a/tests/e2e/mcp/oauth_chat_client.py
+++ b/tests/e2e/mcp/oauth_chat_client.py
@ -25,6 +25,7 @@ import httpx
 import pytest
 from e2e_config import PROXY_BASE_URL, REQUEST_TIMEOUT
 from e2e_http import AuthHeaders, NoBody, unwrap
+from idp import Identity
 from mcp import ClientSession
 from mcp.client.auth import OAuthClientProvider
 from mcp.client.streamable_http import streamable_http_client
@ -77,7 +78,13 @@ class InMemoryTokenStorage:
        self._client_info = client_info


-async def _browser_follow_authorize(start_url: str, storage_state_path: str) -> tuple[str, str | None]:
+async def _browser_follow_authorize(
+    start_url: str,
+    storage_state_path: str,
+    identity: Identity | None = None,
+    server_alias: str | None = None,
+    allow_upstream_consent: bool = True,
+) -> tuple[str, str | None]:
    """Play the browser's role for a real upstream whose authorize endpoint
    serves an interactive consent page (Linear). A headless Chromium primed
    with a human's saved Linear session opens the gateway authorize URL and
@ -115,6 +122,29 @@ async def _browser_follow_authorize(start_url: str, storage_state_path: str) ->
                pass
            if "url" in captured:
                break
+            if await page.locator("#username").count() and identity is not None:
+                await page.locator("#username").fill(identity.username)
+                await page.locator("#password").fill(identity.password)
+                await page.locator("#kc-login").click()
+                continue
+            if httpx.URL(page.url).host.endswith("linear.app") and not allow_upstream_consent:
+                raise AssertionError("cold reconnect required upstream consent")
+            if "/ui/connect" in page.url and server_alias is not None:
+                card = page.locator("div.cursor-pointer").filter(has=page.get_by_text(server_alias, exact=True))
+                if await card.count() != 1:
+                    await asyncio.sleep(0.5)
+                    continue
+                connect = card.get_by_text("Connect", exact=True)
+                if await connect.count():
+                    await connect.click()
+                    continue
+                if not await card.locator("svg.text-success").count():
+                    await asyncio.sleep(0.5)
+                    continue
+                finish = page.get_by_role("button", name="Finish connecting", exact=True)
+                if await finish.count() and await finish.is_enabled():
+                    await finish.click()
+                    continue
            control = page.locator(
                'button[name="action"][value="approve"], button:has-text("Authorize"), '
                'button:has-text("Allow"), button:has-text("@"), a:has-text("@")'
@ -132,11 +162,18 @@ async def _browser_follow_authorize(start_url: str, storage_state_path: str) ->
        f"final={final_url.split('?', 1)[0]!r}; trail={trail[-6:]}"
    )
    params = dict(parse_qsl(httpx.URL(landing).query.decode()))
-    assert "code" in params, f"client redirect_uri carried no code: {landing}"
+    assert "code" in params, "client redirect_uri carried no authorization code"
    return params["code"], params.get("state")


-def _oauth_provider(url: str, storage: InMemoryTokenStorage, storage_state_path: str | None) -> OAuthClientProvider:
+def _oauth_provider(
+    url: str,
+    storage: InMemoryTokenStorage,
+    storage_state_path: str | None,
+    identity: Identity | None = None,
+    server_alias: str | None = None,
+    allow_upstream_consent: bool = True,
+) -> OAuthClientProvider:
    """The SDK's real OAuth machinery (RFC 9728/8414 discovery, RFC 7591 DCR,
    PKCE, token exchange) with the browser leg driven by Playwright against the
    upstream's consent screen."""
@ -147,7 +184,9 @@ def _oauth_provider(url: str, storage: InMemoryTokenStorage, storage_state_path:

    async def _follow_redirect(authorize_url: str) -> None:
        assert storage_state_path is not None
-        code, state = await _browser_follow_authorize(authorize_url, storage_state_path)
+        code, state = await _browser_follow_authorize(
+            authorize_url, storage_state_path, identity, server_alias, allow_upstream_consent
+        )
        code_holder["code"] = code
        code_holder["state"] = state

@ -202,8 +241,15 @@ class _HeaderInjectingTransport(httpx.AsyncBaseTransport):
            for name, value in self._headers.items():
                if name not in request.headers:
                    request.headers[name] = value
+        else:
+            for name, value in self._headers.items():
+                if request.headers.get(name) == value:
+                    del request.headers[name]
        return await self._inner.handle_async_request(request)

+    async def aclose(self) -> None:
+        await self._inner.aclose()
+

 def _oauth_http_client(
    headers: dict[str, str], auth: OAuthClientProvider, gateway_url: str = PROXY_BASE_URL
@ -242,9 +288,14 @@ async def _list_and_call(
    tool: str,
    arguments: dict[str, str],
    gateway_url: str = PROXY_BASE_URL,
+    identity: Identity | None = None,
+    server_alias: str | None = None,
+    allow_upstream_consent: bool = True,
 ) -> OauthToolRun:
    async with _oauth_http_client(
-        headers, _oauth_provider(url, storage, storage_state_path), gateway_url
+        headers,
+        _oauth_provider(url, storage, storage_state_path, identity, server_alias, allow_upstream_consent),
+        gateway_url,
    ) as http_client:
        async with streamable_http_client(url, http_client=http_client) as (read, write, _):
            async with ClientSession(read, write) as session:
@ -321,29 +372,22 @@ class ChatMcpClient:
        tool: str,
        arguments: dict[str, str],
        base_url: str = PROXY_BASE_URL,
+        identity: Identity | None = None,
+        allow_upstream_consent: bool = True,
    ) -> OauthToolRun:
-        deadline: Final = time.monotonic() + self.proxy.poll_timeout
-        last_error: Exception | None = None
-        while time.monotonic() < deadline:
-            try:
-                return asyncio.run(
-                    _list_and_call(
-                        _mcp_url(alias, base_url),
-                        headers,
-                        storage,
-                        storage_state_path,
-                        tool,
-                        arguments,
-                        base_url,
-                    )
-                )
-            except AssertionError:
-                raise
-            except Exception as exc:  # noqa: BLE001 - retried to the deadline; the last error surfaces below
-                last_error = exc
-                time.sleep(self.proxy.poll_interval)
-        pytest.fail(
-            f"list and call for {alias!r} never completed within {self.proxy.poll_timeout}s; last error: {last_error!r}"
+        return asyncio.run(
+            _list_and_call(
+                f"{base_url.rstrip('/')}/mcp" if identity is not None else _mcp_url(alias, base_url),
+                headers,
+                storage,
+                storage_state_path,
+                tool,
+                arguments,
+                base_url,
+                identity,
+                alias,
+                allow_upstream_consent,
+            )
        )

    def server_user_credentials(self, server_id: str) -> tuple[McpServerUserCredentialRow, ...]:
--- a/tests/e2e/mcp/oauth_gateway.py
+++ b/tests/e2e/mcp/oauth_gateway.py
@ -0,0 +1,197 @@
+"""An owned, source-built OAuth gateway with cold restarts and credential observations.
+
+Only this child process is restarted. Its database and SSO client survive while
+its process-local caches do not; Redis is deliberately absent from its config.
+The optional live edge measures headers without recording credentials or bodies.
+"""
+
+from __future__ import annotations
+
+import os
+import socket
+import subprocess
+import sys
+import threading
+import time
+from collections.abc import Callable, Mapping
+from contextlib import ExitStack
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Final
+
+import psycopg
+from e2e_http import NoBody
+from idp import Keycloak, stop_process_group
+from proxy_client import ProxyClient, build_proxy_client
+from psycopg.rows import class_row
+from pydantic import BaseModel, SecretStr, TypeAdapter, ValidationError
+
+
+class StoredOAuth(BaseModel):
+    type: str
+    access_token: SecretStr
+
+
+@dataclass(frozen=True, slots=True)
+class CredentialRow:
+    credential_b64: str = field(repr=False)
+
+
+def stored_oauth(user_id: str, server_id: str) -> StoredOAuth:
+    from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper
+
+    with psycopg.Connection[CredentialRow].connect(
+        os.environ["DATABASE_URL"], row_factory=class_row(CredentialRow)
+    ) as conn:
+        row: Final = conn.execute(
+            'SELECT credential_b64 FROM "LiteLLM_MCPUserCredentials" WHERE user_id = %s AND server_id = %s',
+            (user_id, server_id),
+        ).fetchone()
+    assert row is not None, "canonical user/server has no persisted credential"
+    plaintext: Final = decrypt_value_helper(
+        row.credential_b64, "e2e_mcp_oauth", exception_type="debug", return_original_value=False
+    )
+    assert plaintext is not None, "persisted credential must decrypt with the gateway salt"
+    assert plaintext != row.credential_b64, "persisted credential must be encrypted"
+    try:
+        credential: Final = StoredOAuth.model_validate_json(plaintext)
+    except ValidationError:
+        raise AssertionError("decrypted credential is not an OAuth payload") from None
+    assert credential.type == "oauth2"
+    assert bool(credential.access_token.get_secret_value()), "stored upstream token is empty"
+    return credential
+
+
+class RpcMethod(BaseModel):
+    method: str = ""
+
+
+@dataclass(slots=True)
+class OAuthObservation:
+    user_id: str
+    server_id: str = ""
+    gateway_token: str = field(default="", repr=False)
+    _seen: tuple[tuple[str, bool, bool], ...] = field(default=(), init=False, repr=False)
+    _lock: threading.Lock = field(default_factory=threading.Lock, init=False, repr=False)
+
+    def observe(self, url: str, headers: Mapping[str, str], body: bytes | None) -> None:
+        if not self.server_id or body is None or not url.endswith("/mcp"):
+            return
+        try:
+            operation: Final = RpcMethod.model_validate_json(body).method
+        except ValidationError:
+            return
+        if operation not in ("tools/list", "tools/call"):
+            return
+        credential: Final = stored_oauth(self.user_id, self.server_id)
+        received: Final = headers.get("authorization", "")
+        matches: Final = received == f"Bearer {credential.access_token.get_secret_value()}"
+        differs: Final = bool(received) and all(
+            value not in (self.gateway_token, f"Bearer {self.gateway_token}") for value in headers.values()
+        )
+        with self._lock:
+            self._seen = (*self._seen, (operation, matches, differs))
+
+    def assert_forwarded(self) -> None:
+        with self._lock:
+            snapshot: Final = self._seen
+            self._seen = ()
+        assert {item[0] for item in snapshot} == {"tools/list", "tools/call"}, "missing upstream observations"
+        assert all(item[1] and item[2] for item in snapshot), "upstream bearer did not match the user's stored token"
+
+
+def available_port() -> int:
+    with socket.socket() as listener:
+        listener.bind(("127.0.0.1", 0))
+        return TypeAdapter(tuple[str, int]).validate_python(listener.getsockname())[1]
+
+
+@dataclass(slots=True)
+class OAuthGateway:
+    base_url: str
+    proxy: ProxyClient
+    _environment: Mapping[str, str] = field(repr=False)
+    _command: tuple[str, ...] = field(repr=False)
+    _log_path: Path
+    _child: subprocess.Popen[bytes] | None = field(default=None, init=False, repr=False)
+
+    def start(self) -> None:
+        with self._log_path.open("ab") as log:
+            self._child = subprocess.Popen(
+                self._command,
+                env=self._environment,
+                stdout=log,
+                stderr=log,
+                start_new_session=True,
+            )
+        deadline: Final = time.monotonic() + 120
+        while time.monotonic() < deadline:
+            assert self._child.poll() is None, "owned OAuth gateway exited; inspect its private log"
+            result = self.proxy.transport.probe("/health/liveliness", params=NoBody())
+            if result.status_code == 200:
+                return
+            time.sleep(0.5)
+        raise AssertionError("owned OAuth gateway did not become ready")
+
+    def stop(self) -> None:
+        if self._child is not None:
+            stop_process_group(self._child)
+            assert self._child.poll() is not None, "old gateway process is still alive"
+
+    def restart(self) -> None:
+        assert self._child is not None
+        previous: Final = self._child.pid
+        self.stop()
+        self.start()
+        assert self._child.pid != previous, "gateway restart did not create a new process"
+
+
+def owned_gateway(idp: Keycloak, directory: Path, cleanup: ExitStack) -> OAuthGateway:
+    for name in ("DATABASE_URL", "LITELLM_LICENSE", "LITELLM_SALT_KEY", "LITELLM_MASTER_KEY"):
+        assert os.environ.get(name), f"{name} is required for the owned OAuth gateway"
+    port: Final = available_port()
+    base_url: Final = f"http://127.0.0.1:{port}"
+
+    def defer(callback: Callable[[], object]) -> None:
+        cleanup.callback(callback)
+
+    browser: Final = idp.browser_client(callback_url=f"{base_url}/sso/callback", defer=defer)
+    config: Final = directory / "oauth-gateway.yaml"
+    config.write_text(
+        "model_list: []\n"
+        "general_settings:\n"
+        "  master_key: os.environ/LITELLM_MASTER_KEY\n"
+        "  database_url: os.environ/DATABASE_URL\n"
+        "  enable_jwt_auth: true\n"
+        "  litellm_jwtauth:\n"
+        "    user_id_jwt_field: sub\n"
+        "    user_email_jwt_field: email\n"
+        "    team_ids_jwt_field: groups\n"
+        "    user_id_upsert: true\n"
+    )
+    environment: Final = {
+        **{key: value for key, value in os.environ.items() if not key.startswith("REDIS_")},
+        **browser.environment(idp.discovery()),
+        "PROXY_BASE_URL": base_url,
+        "JWT_PUBLIC_KEY_URL": idp.jwks_url,
+        "JWT_ISSUER": idp.issuer,
+        "JWT_AUDIENCE": "litellm-e2e",
+        "DISABLE_SCHEMA_UPDATE": "true",
+        "STORE_MODEL_IN_DB": "True",
+        "PYTHONPATH": str(Path(__file__).resolve().parents[3]),
+    }
+    gateway: Final = OAuthGateway(
+        base_url=base_url,
+        proxy=build_proxy_client(
+            base_url=base_url,
+            control_plane_base_url=base_url,
+            replica_urls=(base_url,),
+            master_key=os.environ["LITELLM_MASTER_KEY"],
+        ),
+        _environment=environment,
+        _command=(sys.executable, "-m", "litellm.proxy.proxy_cli", "--config", str(config), "--port", str(port)),
+        _log_path=directory / "oauth-gateway.log",
+    )
+    cleanup.callback(gateway.stop)
+    gateway.start()
+    return gateway
--- a/tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py
+++ b/tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py
@ -1,45 +1,85 @@
-"""Live e2e coverage for the gateway-managed MCP OAuth protocol path.
+"""Real OAuth consent, immediate MCP operations and cold-restart persistence.

-The test creates a JWT-authorized user, completes real Linear authorization
-consent, lists and calls a tool immediately through the per-server MCP route,
-and verifies the canonical per-user credential row. The first run targets the
-first configured gateway replica, and a fresh SDK client then targets a
-different replica to prove that a process which did not run consent resolves
-the stored token.
+Aggregate SSO uses the SDK's normal authentication. The per-server variant is
+explicitly a configured two-header client, not an Authorization-only OAuth host.
+The observed variants forward to the same real Linear upstream and compare its
+bearer at the forwarding boundary; direct variants retain unmodified discovery.
 """

 from __future__ import annotations

 import os
-from typing import Final
+from collections.abc import Iterator
+from contextlib import ExitStack
+from pathlib import Path
+from types import MappingProxyType
+from typing import Final, Literal

 import pytest
-from e2e_config import (
-    LINEAR_MCP_URL,
-    LINEAR_READONLY_TOOL,
-    LINEAR_STORAGE_STATE,
-    PROXY_REPLICA_URLS,
-    unique_marker,
-)
-from e2e_http import AuthHeaders
+from e2e_config import LINEAR_MCP_URL, LINEAR_READONLY_TOOL, LINEAR_STORAGE_STATE, unique_marker
+from e2e_http import AuthHeaders, NoBody, get_external, unwrap
+from idp import Identity, Keycloak
 from lifecycle import ResourceManager
-from models import McpServerCreateBody, ObjectPermission, TeamUpdateBody
-from proxy_client import ProxyClient
-
-pytest.importorskip("mcp", reason="mcp SDK not installed; run `uv sync --inexact --group e2e-dev`")
-pytest.importorskip(
-    "playwright.async_api",
-    reason="playwright not installed; run `uv pip install playwright` and `playwright install chromium`",
+from models import (
+    McpOauthCredentials,
+    McpServerCreateBody,
+    ObjectPermission,
+    TeamMemberAddBody,
+    TeamMemberEntry,
+    TeamUpdateBody,
 )
+from oauth_chat_client import ChatMcpClient, InMemoryTokenStorage, OauthToolRun, build_chat_client
+from oauth_gateway import OAuthGateway, OAuthObservation, owned_gateway, stored_oauth
+from provider_edge import LiveEdge, start_provider_edge
+from proxy_client import ProxyClient
+from pydantic import BaseModel, ValidationError

-from idp import Identity, Keycloak  # noqa: E402
-from oauth_chat_client import ChatMcpClient, InMemoryTokenStorage, build_chat_client  # noqa: E402
-
-pytestmark = [pytest.mark.e2e, pytest.mark.mcp_oauth_live]
+pytestmark = [pytest.mark.e2e, pytest.mark.mcp_oauth_live, pytest.mark.provider_live]


-@pytest.fixture(scope="session")
-def chat_client(proxy: ProxyClient) -> ChatMcpClient:
+class OAuthMetadata(BaseModel):
+    authorization_endpoint: str
+    token_endpoint: str
+    registration_endpoint: str
+
+
+class LinearTeam(BaseModel):
+    id: str
+    name: str
+
+
+class LinearTeams(BaseModel):
+    teams: tuple[LinearTeam, ...]
+
+
+def assert_tool_result(run: OauthToolRun, tool: str) -> None:
+    assert tool in run.tools
+    assert run.is_error is False
+    try:
+        result: Final = LinearTeams.model_validate_json(run.text)
+    except ValidationError:
+        raise AssertionError("list_teams did not return the expected teams payload") from None
+    assert result.teams, "the test workspace must contain at least one team"
+    assert all(team.id and team.name for team in result.teams), "team results must contain identifiers and names"
+
+
+@pytest.fixture(scope="module")
+def oauth_gateway(idp: Keycloak, tmp_path_factory: pytest.TempPathFactory) -> Iterator[OAuthGateway]:
+    assert LINEAR_STORAGE_STATE and Path(LINEAR_STORAGE_STATE).is_file(), (
+        "E2E_LINEAR_STORAGE_STATE must name a captured Linear login; see mcp/linear_session_capture.py"
+    )
+    assert os.environ.get("E2E_FIXTURE_MODE", "live") == "live", "OAuth acceptance cannot use replay"
+    with ExitStack() as cleanup:
+        yield owned_gateway(idp, tmp_path_factory.mktemp("mcp-oauth"), cleanup)
+
+
+@pytest.fixture(scope="module")
+def proxy(oauth_gateway: OAuthGateway) -> ProxyClient:
+    return oauth_gateway.proxy
+
+
+@pytest.fixture(scope="module")
+def client(proxy: ProxyClient) -> ChatMcpClient:
    return build_chat_client(proxy)


@ -47,80 +87,122 @@ class TestMcpOauthHappyPath:
    @pytest.mark.covers("mcp.list_tools.oauth.succeeds")
    @pytest.mark.covers("mcp.call_tool.oauth.succeeds")
    @pytest.mark.covers("mcp.call_tool.oauth.persists_across_processes")
-    def test_jwt_user_lists_and_calls_then_reconnects_from_another_gateway(
+    @pytest.mark.parametrize("route", ("aggregate_sso", "explicit_header_jwt"))
+    @pytest.mark.parametrize("observed", (False, True), ids=("direct", "observed"))
+    def test_consent_list_call_and_cold_restart(
        self,
-        chat_client: ChatMcpClient,
+        client: ChatMcpClient,
        resources: ResourceManager,
        jwt_identity: Identity,
        idp: Keycloak,
+        oauth_gateway: OAuthGateway,
+        route: Literal["aggregate_sso", "explicit_header_jwt"],
+        observed: bool,
    ) -> None:
-        assert LINEAR_STORAGE_STATE and os.path.exists(LINEAR_STORAGE_STATE), (
-            "E2E_MCP_OAUTH_LIVE is set but E2E_LINEAR_STORAGE_STATE does not point at a captured "
-            "Linear session (run mcp/linear_session_capture.py)"
-        )
-
        alias: Final = f"e2elinear{unique_marker()}"
        tool: Final = f"{alias}-{LINEAR_READONLY_TOOL}"
-        assert len(PROXY_REPLICA_URLS) >= 2, (
-            "set LITELLM_PROXY_REPLICA_URLS to at least two gateway URLs; the persistence cell needs a process "
-            "that did not run the consent"
+        token: Final = idp.access_token(jwt_identity)
+        observation: Final = OAuthObservation(user_id=jwt_identity.user_id, gateway_token=token)
+        edge: Final = (
+            start_provider_edge(
+                LiveEdge(observe_request=observation.observe),
+                mounts=MappingProxyType(
+                    {"linear": "https://mcp.linear.app", ".well-known": "https://mcp.linear.app/.well-known"}
+                ),
+            )
+            if observed
+            else None
        )
-        created: Final = chat_client.create_server(
+        if edge is not None:
+            resources.defer(edge.shutdown)
+        metadata: Final = (
+            unwrap(
+                get_external(
+                    "https://mcp.linear.app/.well-known/oauth-authorization-server",
+                    response_type=OAuthMetadata,
+                )
+            )
+            if observed
+            else None
+        )
+        created: Final = client.create_server(
            McpServerCreateBody(
                alias=alias,
-                url=LINEAR_MCP_URL,
+                server_name=alias,
+                url=f"{edge.edge.api_base('linear')}/mcp" if edge is not None else LINEAR_MCP_URL,
+                transport="http",
                allow_all_keys=False,
                auth_type="oauth2",
                oauth2_flow="authorization_code",
-                per_server_oauth_discovery=True,
+                per_server_oauth_discovery=route == "explicit_header_jwt",
+                authorization_url=metadata.authorization_endpoint if metadata else None,
+                token_url=metadata.token_endpoint if metadata else None,
+                registration_url=metadata.registration_endpoint if metadata else None,
+                credentials=McpOauthCredentials(upstream_resource=LINEAR_MCP_URL) if observed else None,
            )
        )
-        resources.defer(lambda: chat_client.delete_server(created.server_id))
-
-        chat_client.proxy.update_team(
+        resources.defer(lambda: client.delete_server(created.server_id))
+        assert client.server_user_credentials(created.server_id) == (), (
+            "scenario must start without upstream credentials"
+        )
+        observation.server_id = created.server_id
+        client.proxy.update_team(
            TeamUpdateBody(
                team_id=jwt_identity.group,
                object_permission=ObjectPermission(mcp_servers=[created.server_id]),
            )
        )
-
-        token: Final = idp.access_token(jwt_identity)
-        headers: Final = {"x-litellm-api-key": f"Bearer {token}"}
-        storage: Final = InMemoryTokenStorage()
-        first_run: Final = chat_client.list_and_call(
+        unwrap(
+            client.proxy.transport.post(
+                "/team/member_add",
+                headers=client.proxy.transport.master,
+                json=TeamMemberAddBody(
+                    team_id=jwt_identity.group, member=TeamMemberEntry(user_id=jwt_identity.user_id, role="user")
+                ),
+                response_type=NoBody,
+            )
+        )
+        headers: Final = {"x-litellm-api-key": f"Bearer {token}"} if route == "explicit_header_jwt" else {}
+        resources.defer(
+            lambda: client.revoke_user_token(
+                created.server_id,
+                AuthHeaders(authorization=f"Bearer {idp.access_token(jwt_identity)}"),
+            )
+        )
+        identity: Final = jwt_identity if route == "aggregate_sso" else None
+        first: Final = client.list_and_call(
            alias,
            headers,
-            storage,
+            InMemoryTokenStorage(),
            LINEAR_STORAGE_STATE,
            tool,
            {},
-            base_url=PROXY_REPLICA_URLS[0],
+            base_url=oauth_gateway.base_url,
+            identity=identity,
        )
-        assert tool in first_run.tools
-        assert first_run.is_error is False
-        assert first_run.text.strip() != ""
-
-        credentials: Final = chat_client.server_user_credentials(created.server_id)
+        assert_tool_result(first, tool)
+        credentials: Final = client.server_user_credentials(created.server_id)
        assert len(credentials) == 1
        assert credentials[0].user_id == jwt_identity.user_id
        assert credentials[0].credential_type == "oauth2"
-        resources.defer(
-            lambda: chat_client.revoke_user_token(
-                created.server_id,
-                AuthHeaders.model_validate(headers),
-            )
-        )
-
-        replica: Final = PROXY_REPLICA_URLS[-1]
-        second_run: Final = chat_client.list_and_call(
+        stored_oauth(jwt_identity.user_id, created.server_id)
+        if observed:
+            observation.assert_forwarded()
+        oauth_gateway.restart()
+        fresh_token: Final = idp.access_token(jwt_identity)
+        observation.gateway_token = fresh_token
+        second: Final = client.list_and_call(
            alias,
-            {"x-litellm-api-key": f"Bearer {idp.access_token(jwt_identity)}"},
+            {"Authorization": f"Bearer {fresh_token}"} if identity is None else {},
            InMemoryTokenStorage(),
-            None,
+            LINEAR_STORAGE_STATE if identity is not None else None,
            tool,
            {},
-            base_url=replica,
+            base_url=oauth_gateway.base_url,
+            identity=identity,
+            allow_upstream_consent=False,
        )
-        assert tool in second_run.tools
-        assert second_run.is_error is False
-        assert second_run.text.strip() != ""
+        assert_tool_result(second, tool)
+        stored_oauth(jwt_identity.user_id, created.server_id)
+        if observed:
+            observation.assert_forwarded()
--- a/tests/e2e/models.py
+++ b/tests/e2e/models.py
@ -572,6 +572,10 @@ class McpInfo(BaseModel):
    logo_url: str | None = None


+class McpOauthCredentials(BaseModel):
+    upstream_resource: str
+
+
 class McpServerCreateBody(BaseModel):
    """POST /v1/mcp/server. For a gateway-managed OAuth server, `auth_type` is
    `oauth2` and `oauth2_flow` is `authorization_code`; the upstream endpoints
@ -587,6 +591,8 @@ class McpServerCreateBody(BaseModel):
    per_server_oauth_discovery: bool | None = None
    authorization_url: str | None = None
    token_url: str | None = None
+    registration_url: str | None = None
+    credentials: McpOauthCredentials | None = None
    server_name: str | None = None
    description: str | None = None
    mcp_info: McpInfo | None = None
--- a/tests/e2e/provider_edge.py
+++ b/tests/e2e/provider_edge.py
@ -46,7 +46,7 @@ import os
 import re
 import threading
 from collections import deque
-from collections.abc import Generator, Mapping, Sequence
+from collections.abc import Callable, Generator, Mapping, Sequence
 from contextlib import closing, contextmanager
 from dataclasses import dataclass, field, replace
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
@ -538,7 +538,7 @@ class ReplayEdge:

@dataclass(frozen=True, slots=True)
 class LiveEdge:
-    pass
+    observe_request: Callable[[str, Mapping[str, str], bytes | None], None] | None = None


 type EdgeBackend = RecordEdge | ReplayEdge | LiveEdge | CacheEdge
@ -787,10 +787,13 @@ def _handle_record(
 def _handle_live(
    method: str, url: str, headers: Mapping[str, str], body: bytes | None, timeout: float,
    cache: CacheEdge | None = None, mount: str = "", test_key: str | None = None,
+    observe_request: Callable[[str, Mapping[str, str], bytes | None], None] | None = None,
 ) -> EdgeOutcome:
    forwarded: Final = {
        name: value for name, value in headers.items() if name.lower() not in _REQUEST_DROPPED_HEADERS
    }
+    if observe_request is not None:
+        observe_request(url, forwarded, body)
    head: Final = (
        forward_stream(method, url, headers=forwarded, body=body, timeout=timeout)
        if cache is None else cache.forward(mount, method, url, forwarded, body, timeout, test_key=test_key)
@ -868,9 +871,10 @@ def handle_edge_request(
                method, _upstream_url(upstream_base, upstream_path, split.query), headers, body, timeout,
                backend, mount, test_key,
            )
-        case LiveEdge():
+        case LiveEdge(observe_request=observe_request):
            return _handle_live(
-                method, _upstream_url(upstream_base, upstream_path, split.query), headers, body, timeout
+                method, _upstream_url(upstream_base, upstream_path, split.query), headers, body, timeout,
+                observe_request=observe_request,
            )
        case RecordEdge():
            return _handle_record(