From 7f4dd4eabcce3c53a413e4697e96a8ca03834928 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 22:02:09 -0700 Subject: [PATCH] test(e2e): cover MCP OAuth SSO and cold restart acceptance --- .github/e2e-stack/assert_tests_ran.py | 7 + .github/e2e-stack/select_tests.py | 1 + .github/workflows/test-mcp-oauth-e2e.yml | 168 +++++++++++++ tests/e2e/AGENTS.md | 2 +- tests/e2e/CONTRIBUTING.md | 46 ++++ tests/e2e/conftest.py | 2 + tests/e2e/coverage_registry/mcp.yaml | 2 +- tests/e2e/idp.py | 4 +- tests/e2e/mcp/oauth_chat_client.py | 98 +++++--- tests/e2e/mcp/oauth_gateway.py | 197 +++++++++++++++ .../e2e/mcp/test_mcp_oauth_happy_path_e2e.py | 226 ++++++++++++------ tests/e2e/models.py | 6 + tests/e2e/provider_edge.py | 12 +- 13 files changed, 664 insertions(+), 107 deletions(-) create mode 100644 .github/workflows/test-mcp-oauth-e2e.yml create mode 100644 tests/e2e/mcp/oauth_gateway.py diff --git a/.github/e2e-stack/assert_tests_ran.py b/.github/e2e-stack/assert_tests_ran.py index 2303c42f4fb..bc299b14af8 100644 --- 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")) ) diff --git a/.github/e2e-stack/select_tests.py b/.github/e2e-stack/select_tests.py index 982e93cf642..a9ca1f88660 100644 --- 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$" diff --git a/.github/workflows/test-mcp-oauth-e2e.yml b/.github/workflows/test-mcp-oauth-e2e.yml new file mode 100644 index 00000000000..5fc9b711fd5 --- /dev/null +++ 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" diff --git a/tests/e2e/AGENTS.md b/tests/e2e/AGENTS.md index 967a85f1255..9b662e511b8 100644 --- 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 diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index 2afcc563824..2adac08329f 100644 --- 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 diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index d7d173c93d4..268d517a7fe 100644 --- 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) diff --git a/tests/e2e/coverage_registry/mcp.yaml b/tests/e2e/coverage_registry/mcp.yaml index bf511ad6b06..1cdeac7b77f 100644 --- 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 diff --git a/tests/e2e/idp.py b/tests/e2e/idp.py index 2dc7c2ad71b..a89a036baeb 100644 --- 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__": diff --git a/tests/e2e/mcp/oauth_chat_client.py b/tests/e2e/mcp/oauth_chat_client.py index ebae8029a47..f6e72fb37dc 100644 --- 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, ...]: diff --git a/tests/e2e/mcp/oauth_gateway.py b/tests/e2e/mcp/oauth_gateway.py new file mode 100644 index 00000000000..cd71502aad5 --- /dev/null +++ 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 diff --git a/tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py b/tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py index 2755629421a..305989850f1 100644 --- 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() diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 4308984c3be..4b202e3c663 100644 --- 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 diff --git a/tests/e2e/provider_edge.py b/tests/e2e/provider_edge.py index 7bbb1375623..136b00208f7 100644 --- 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(