diff --git a/.github/e2e-stack/assert_tests_ran.py b/.github/e2e-stack/assert_tests_ran.py index 2303c42f4fb..1b051f860cc 100644 --- a/.github/e2e-stack/assert_tests_ran.py +++ b/.github/e2e-stack/assert_tests_ran.py @@ -1,3 +1,5 @@ +import os +import re import sys import xml.etree.ElementTree as ET from pathlib import Path @@ -15,6 +17,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")) ) @@ -38,6 +46,13 @@ def main() -> int: if case.get("file") != path or all(case.find(tag) is None for tag in ("failure", "error")): continue _ = sys.stdout.write(f" failed: {case.get('classname', '')}::{case.get('name', '')}\n") + for prop in case.findall("./properties/property"): + name = prop.get("name", "") + value = prop.get("value", "") + if name in ("oauth_failure_phase", "oauth_exception_type", "oauth_frame") and re.fullmatch( + r"[A-Za-z0-9_.:<>-]{1,240}", value + ): + _ = sys.stdout.write(f" {name}: {value}\n") if ( selected and not missing 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..034b9fe49ec --- /dev/null +++ b/.github/workflows/test-mcp-oauth-e2e.yml @@ -0,0 +1,177 @@ +name: MCP OAuth happy path + +on: + pull_request: + paths: + - '.github/workflows/test-mcp-oauth-e2e.yml' + - '.github/e2e-stack/**' + - 'tests/e2e/*.py' + - 'tests/e2e/pytest.ini' + - 'tests/e2e/idp_realm.json' + - 'tests/e2e/mcp/**' + - 'litellm/experimental_mcp_client/**' + - 'litellm/proxy/_experimental/mcp_server/**' + - 'litellm/proxy/auth/**' + - 'litellm/proxy/management_endpoints/*sso*.py' + - 'litellm/proxy/management_endpoints/sso/**' + - 'litellm/proxy/common_utils/encrypt_decrypt_utils.py' + - 'litellm/proxy/proxy_server.py' + - 'litellm/proxy/schema.prisma' + - 'ui/litellm-dashboard/src/app/connect/**' + - 'ui/litellm-dashboard/src/app/mcp/oauth/**' + - 'pyproject.toml' + - 'uv.lock' + workflow_dispatch: + +permissions: {} + +concurrency: + group: mcp-oauth-${{ 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: Report JUnit results and 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/code_coverage_tests/test_e2e_changed_gate.py b/tests/code_coverage_tests/test_e2e_changed_gate.py index 5ae0863baf0..707566c0333 100644 --- a/tests/code_coverage_tests/test_e2e_changed_gate.py +++ b/tests/code_coverage_tests/test_e2e_changed_gate.py @@ -226,3 +226,31 @@ def test_an_unusable_secret_is_named_without_printing_its_value( assert unprintable not in result.stderr assert result.stdout == "" assert not env_path.exists() + + +@pytest.mark.parametrize("phase", ("setup", "call", "teardown")) +def test_oauth_failure_diagnostics_do_not_publish_private_payloads(tmp_path: Path, phase: str) -> None: + suite: Final = ET.Element("testsuite") + case: Final = ET.SubElement(suite, "testcase", file=SELECTED[0]) + private: Final = "private-token-in-exception-message" + failure: Final = ET.SubElement(case, "failure", message=private) + failure.text = private + properties: Final = ET.SubElement(case, "properties") + for name, value in ( + ("oauth_failure_phase", phase), + ("oauth_exception_type", "AssertionError"), + ("oauth_frame", "oauth_gateway.py:120:start"), + ("oauth_frame", f"injected\\n{private}"), + ("unrelated_property", private), + ): + _ = ET.SubElement(properties, "property", name=name, value=value) + report: Final = tmp_path / "report.xml" + ET.ElementTree(suite).write(report) + result: Final = subprocess.run( + [sys.executable, "-I", str(GATE), str(report), SELECTED[0]], capture_output=True, text=True + ) + assert result.returncode == 1 + assert f"oauth_failure_phase: {phase}" in result.stdout + assert "oauth_exception_type: AssertionError" in result.stdout + assert "oauth_frame: oauth_gateway.py:120:start" in result.stdout + assert private not in result.stdout + result.stderr diff --git a/tests/e2e/AGENTS.md b/tests/e2e/AGENTS.md index 8a56e8673c4..9b662e511b8 100644 --- a/tests/e2e/AGENTS.md +++ b/tests/e2e/AGENTS.md @@ -14,7 +14,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `quota_management/` - quota enforcement and accounting, one subfolder per behavior: `ratelimit/` (rpm/tpm blocks, window reset, pacing headers on live traffic), `budgets/` (budget definition, enforcement, and reset windows: key, team, tag, soft, multi-window), and `spend_tracking/` (spend logging and cost attribution on `/spend/*`) - `management/` - key/team/user/organization management routes: create/update/delete persistence via the info routes, team membership, and llm-only-key route denials (API surface; not Playwright) - `a2a/` - the A2A (agent-to-agent) surface: admin registration via `/v1/agents`, proxy-fronted card discovery at `/.well-known/agent-card.json`, and JSON-RPC `message/send` invocation, driving agents backed by the litellm completion bridge (a real provider) and asserting protocol-version normalization (0.3 vs 1.0) -- `mcp/` - the MCP server surface over api_key auth against the real Datadog remote MCP server (see "MCP suite: real Datadog only" below); plus the gateway-managed OAuth (authorization_code) path exercised through `/chat/completions`, the one behavior Datadog's static-header auth cannot reach, seeding the per-user upstream token via the interactive authorize dance driven with the mcp SDK's own OAuth client (headless-browser consent from a saved session) and asserting the completion lists and executes the server's tools with the stored per-user token +- `mcp/` - the MCP server surface over api_key auth against the real Datadog remote MCP server (see "MCP suite: real Datadog only" below); plus the gateway-managed OAuth (authorization_code) path exercised through `/chat/completions` in `test_mcp_chat_completion_oauth_e2e.py` and direct MCP protocol operations in `test_mcp_oauth_happy_path_e2e.py`, the one behavior Datadog's static-header auth cannot reach, seeding the per-user upstream token via the interactive authorize dance driven with the mcp SDK's own OAuth client (headless-browser consent from a saved session) and asserting the completion or protocol call lists and executes the server's tools with the stored per-user token - `logging/` - logging-integration delivery (datadog and friends) - `security/` - secret handling and log-leak protection - `router/` - routing and reliability behavior (fallbacks, cooldowns) plus the memory regression test (`test_reliability_memory_e2e.py`: a few hundred failing requests with retries and fallbacks must not grow proxy RSS past a fixed budget nor store a request snapshot past a fixed size, the release-gate check for the v1.100.0 retry-breadcrumb leak) @@ -26,14 +26,14 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family ## MCP suite: real Datadog only -Every test under `tests/e2e/mcp/` must exercise the proxy against the real Datadog remote MCP server. Do not add a compose service, FastMCP fixture, mock upstream, or any other fake MCP host for this suite +Every test under `tests/e2e/mcp/` must exercise the proxy against the real Datadog remote MCP server, except the two Linear OAuth tests `test_mcp_chat_completion_oauth_e2e.py` and `test_mcp_oauth_happy_path_e2e.py`. Do not add a compose service, FastMCP fixture, mock upstream, or any other fake MCP host for this suite - Register via `register_datadog_mcp` in `tests/e2e/mcp/datadog_mcp.py` (or extend that helper if you need a different `toolsets=` / `allowed_tools` slice of the same Datadog endpoint). That posts `/v1/mcp/server` with `url=datadog_mcp_url(...)` and static headers `DD-API-KEY` / `DD-APPLICATION-KEY` from the process env - Auth is Datadog's documented CI/header path, not a browser OAuth authorize/token dance. Hard-fail when `DD_API_KEY` or `DD_APP_KEY` is missing (`assert_dd_mcp_creds`); never skip for a missing fake upstream - 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 one standing exception is `test_mcp_chat_completion_oauth_e2e.py`. Datadog authenticates with the static `DD-API-KEY` / `DD-APPLICATION-KEY` headers and exposes no authorize/token dance at all, so it cannot exercise gateway-managed OAuth or per-user token seeding in any form. That test drives a real Linear MCP server instead; it is still a real remote upstream, 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 @@ -152,7 +152,7 @@ MCPs - endpoint features with the protocol op as the variant mcp... operation : list_tools | call_tool | list_resources | read_resource | list_prompts | get_prompt auth_family : none | api_key | bearer | oauth - assertion : succeeds | denied_without_permission + assertion : succeeds | denied_without_permission | persists_across_processes e.g. mcp.call_tool.oauth.succeeds ``` diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index 2afcc563824..6c3dc4d0bd1 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -248,3 +248,56 @@ 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` automatically requests a run for +same-repository pull requests changing MCP, gateway authentication/SSO, consent +UI, dependencies or the relevant E2E harness/workflow paths. It retains manual +`workflow_dispatch` for targeted verification. The four cases run in the +protected `e2e-changed` environment after its normal deployment approval; +reviewers should approve and inspect this separate OAuth check when it appears. +Fork pull requests do not run this credentialed job; use a reviewed +same-repository branch for their verification. The workflow's path-filtered +check is not configured here as a globally required branch-protection check. +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 e83827fac74..b0904e39a1f 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -17,6 +17,7 @@ import functools import os from collections.abc import Generator, Iterator from datetime import datetime, timezone +from pathlib import Path from types import MappingProxyType from typing import Final @@ -28,6 +29,7 @@ from e2e_config import ( FIXTURE_DIR, FIXTURE_MODE_RAW, MANAGED_FILES_OPT_IN_ENV, + MCP_OAUTH_LIVE_OPT_IN_ENV, PROMPT_CACHING_OPT_IN_ENV, PROXY_BASE_URL, REDIS_CHAOS_OPT_IN_ENV, @@ -56,6 +58,7 @@ OPT_IN_MARKERS: Final = MappingProxyType( "prompt_caching_stack": PROMPT_CACHING_OPT_IN_ENV, "redis_chaos": REDIS_CHAOS_OPT_IN_ENV, "cli_determinism": CLI_DETERMINISM_OPT_IN_ENV, + "mcp_oauth_live": MCP_OAUTH_LIVE_OPT_IN_ENV, } ) @@ -132,6 +135,11 @@ def pytest_configure(config: pytest.Config) -> None: "redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from " "gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_REDIS_CHAOS is set", ) + config.addinivalue_line( + "markers", + "mcp_oauth_live: real Linear OAuth consent via a captured browser session; deselected unless " + "E2E_MCP_OAUTH_LIVE is set", + ) def pytest_sessionstart(session: pytest.Session) -> None: @@ -213,6 +221,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) @@ -236,6 +246,13 @@ def pytest_runtest_makereport( """Stash the call-phase outcome so teardown can tell a passed test from a failed one without re-deriving it.""" report = yield + if item.get_closest_marker("mcp_oauth_live") is not None and call.excinfo is not None: + # Publish code locations only, never exception messages, source text or locals. + item.user_properties.append(("oauth_failure_phase", report.when)) + item.user_properties.append(("oauth_exception_type", call.excinfo.type.__name__)) + for entry in call.excinfo.traceback: + item.user_properties.append(("oauth_frame", f"{Path(entry.path).name}:{entry.lineno + 1}:{entry.name}")) + report.user_properties = list(item.user_properties) if report.when == "call": item.stash[_CALL_PASSED] = report.passed return report diff --git a/tests/e2e/coverage_registry/mcp.yaml b/tests/e2e/coverage_registry/mcp.yaml index 85ace835144..1cdeac7b77f 100644 --- a/tests/e2e/coverage_registry/mcp.yaml +++ b/tests/e2e/coverage_registry/mcp.yaml @@ -71,6 +71,14 @@ assertions: [succeeds] source: "db.py user_oauth_credential lookup" rationale: OAuth2 token passthrough; per-user credential storage +- id: mcp.call_tool.oauth.persists_across_processes + module: mcp + tier: P1 + operation: call_tool + auth_family: oauth + assertions: [persists_across_processes] + source: "outbound_credentials/per_user_oauth_store.py V2PerUserTokenStore" + 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/e2e_config.py b/tests/e2e/e2e_config.py index 11c52d1398c..a79c158f9c4 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -52,6 +52,7 @@ CHEAP_OPENAI_MODEL = os.environ.get("E2E_CHEAP_OPENAI_MODEL", "gpt-5.5") LINEAR_MCP_URL = os.environ.get("E2E_LINEAR_MCP_URL", "https://mcp.linear.app/mcp") LINEAR_STORAGE_STATE = os.environ.get("E2E_LINEAR_STORAGE_STATE", "") +LINEAR_READONLY_TOOL: Final = "list_teams" # as listed by tools/list on mcp.linear.app when PR #33787 landed # Jaeger query API of the compose stack's OTEL trace destination (the `jaeger` # service in docker-compose.yml maps it to host 16686). Trace-completeness tests @@ -144,6 +145,7 @@ MANAGED_FILES_OPT_IN_ENV = "E2E_MANAGED_FILES_STACK" PROMPT_CACHING_OPT_IN_ENV = "E2E_PROMPT_CACHING_STACK" REDIS_CHAOS_OPT_IN_ENV = "E2E_REDIS_CHAOS" CLI_DETERMINISM_OPT_IN_ENV = "E2E_CLI_DETERMINISM" +MCP_OAUTH_LIVE_OPT_IN_ENV: Final = "E2E_MCP_OAUTH_LIVE" ANOMALY_SESSIONS = int(os.environ.get("E2E_ANOMALY_SESSIONS", "6")) ANOMALY_TURNS_PER_SESSION = int(os.environ.get("E2E_ANOMALY_TURNS_PER_SESSION", "6")) ANOMALY_TURN_ATTEMPTS = int(os.environ.get("E2E_ANOMALY_TURN_ATTEMPTS", "3")) 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 763b348b197..0c5c6106259 100644 --- a/tests/e2e/mcp/oauth_chat_client.py +++ b/tests/e2e/mcp/oauth_chat_client.py @@ -18,7 +18,7 @@ import asyncio import re import time from dataclasses import dataclass -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Final from urllib.parse import parse_qsl import httpx @@ -26,11 +26,21 @@ import httpx2 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 from mcp.shared.auth import AuthorizationCodeResult, OAuthClientInformationFull, OAuthClientMetadata, OAuthToken -from models import ChatBody, ChatResponse, McpServerCreateBody, McpServerInfo +from mcp.types import TextContent +from models import ( + ChatBody, + ChatResponse, + McpOauthUserCredentialStatus, + McpServerCreateBody, + McpServerInfo, + McpServerUserCredentialListResponse, + McpServerUserCredentialRow, +) from proxy_client import ProxyClient if TYPE_CHECKING: @@ -44,8 +54,8 @@ OAUTH_CLIENT_REDIRECT_URI = "http://127.0.0.1:53682/e2e/callback" BROWSER_CONSENT_TIMEOUT = 60.0 -def _mcp_url(alias: str) -> str: - return f"{PROXY_BASE_URL}/{alias}/mcp" +def _mcp_url(alias: str, base_url: str = PROXY_BASE_URL) -> str: + return f"{base_url.rstrip('/')}/{alias}/mcp" class InMemoryTokenStorage: @@ -69,7 +79,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 @@ -85,6 +101,9 @@ async def _browser_follow_authorize(start_url: str, storage_state_path: str) -> def _note_request(request: object) -> None: url = getattr(request, "url", "") + host = httpx.URL(url).host + if not allow_upstream_consent and (host == "linear.app" or host.endswith(".linear.app")): + captured["upstream_consent"] = "seen" if url.startswith(OAUTH_CLIENT_REDIRECT_URI) and "url" not in captured: captured["url"] = url @@ -96,7 +115,7 @@ async def _browser_follow_authorize(start_url: str, storage_state_path: str) -> context = await browser.new_context(storage_state=storage_state_path) await context.route(re.compile(re.escape(OAUTH_CLIENT_REDIRECT_URI) + r".*"), _swallow_redirect) page = await context.new_page() - page.on("request", _note_request) + context.on("request", _note_request) page.on("framenavigated", lambda frame: trail.append(frame.url.split("?", 1)[0])) await page.goto(start_url, wait_until="domcontentloaded") deadline = time.monotonic() + BROWSER_CONSENT_TIMEOUT @@ -105,8 +124,29 @@ async def _browser_follow_authorize(start_url: str, storage_state_path: str) -> await page.wait_for_load_state("networkidle", timeout=8000) except Exception: # noqa: BLE001 - a busy consent page never idles; fall through and try to advance it pass - if "url" in captured: + if "upstream_consent" in captured or "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 "/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("@")' @@ -118,27 +158,44 @@ async def _browser_follow_authorize(start_url: str, storage_state_path: str) -> final_url = page.url await browser.close() + # A redirect chain can finish inside goto/networkidle before the loop checks the page. + assert "upstream_consent" not in captured, "cold reconnect required upstream consent" landing = captured.get("url") assert landing is not None, ( f"consent flow never reached {OAUTH_CLIENT_REDIRECT_URI}; " 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) -> 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.""" code_holder: dict[str, str | None] = {} # mutable-ok: hand-off between the two SDK callbacks - async def redirect_handler(authorize_url: str) -> None: - code, state = await _browser_follow_authorize(authorize_url, storage_state_path) + async def _reject_redirect(_: str) -> None: + raise AssertionError("gateway demanded a fresh upstream consent; stored per-user token was not reused") + + 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, identity, server_alias, allow_upstream_consent + ) code_holder["code"] = code code_holder["state"] = state + redirect_handler: Final = _reject_redirect if storage_state_path is None else _follow_redirect + async def callback_handler() -> AuthorizationCodeResult: code = code_holder.get("code") assert code is not None, "callback_handler ran before the authorize redirect completed" @@ -167,24 +224,45 @@ class _HeaderInjectingTransport(httpx2.AsyncBaseTransport): store the upstream token for from the key on the token exchange, exactly like a production MCP host configured with a LiteLLM key header.""" - def __init__(self, inner: httpx2.AsyncBaseTransport, headers: dict[str, str]) -> None: + def __init__(self, inner: httpx2.AsyncBaseTransport, headers: dict[str, str], gateway_url: str) -> None: self._inner = inner self._headers = headers + self._gateway_url = httpx2.URL(gateway_url) + + @staticmethod + def _port(url: httpx2.URL) -> int | None: + if url.port is not None: + return url.port + return {"http": 80, "https": 443}.get(url.scheme) async def handle_async_request(self, request: httpx2.Request) -> httpx2.Response: - for name, value in self._headers.items(): - if name not in request.headers: - request.headers[name] = value + same_origin: Final = ( + request.url.scheme == self._gateway_url.scheme + and request.url.host == self._gateway_url.host + and self._port(request.url) == self._port(self._gateway_url) + ) + if same_origin: + 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) -> httpx2.AsyncClient: + +def _oauth_http_client( + headers: dict[str, str], auth: OAuthClientProvider, gateway_url: str = PROXY_BASE_URL +) -> httpx2.AsyncClient: return httpx2.AsyncClient( - headers=headers, auth=auth, timeout=httpx2.Timeout(REQUEST_TIMEOUT), follow_redirects=True, - transport=_HeaderInjectingTransport(httpx2.AsyncHTTPTransport(), headers), + transport=_HeaderInjectingTransport(httpx2.AsyncHTTPTransport(), headers, gateway_url), ) @@ -199,6 +277,43 @@ async def _seed_via_dance( return tuple(sorted(tool.name for tool in listed.tools)) +@dataclass(frozen=True, slots=True) +class OauthToolRun: + tools: tuple[str, ...] + is_error: bool + text: str + + +async def _list_and_call( + url: str, + headers: dict[str, str], + storage: InMemoryTokenStorage, + storage_state_path: str | None, + 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, 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: + await session.initialize() + listed: Final = await session.list_tools() + result: Final = await session.call_tool(tool, arguments) + text: Final = "".join(content.text for content in result.content if isinstance(content, TextContent)) + return OauthToolRun( + tools=tuple(sorted(tool_item.name for tool_item in listed.tools)), + is_error=result.is_error, + text=text, + ) + + @dataclass(frozen=True, slots=True) class ChatMcpClient: proxy: ProxyClient @@ -252,6 +367,53 @@ class ChatMcpClient: f"last error: {last_error!r}" ) + def list_and_call( + self, + alias: str, + headers: dict[str, str], + storage: InMemoryTokenStorage, + storage_state_path: str | None, + tool: str, + arguments: dict[str, str], + base_url: str = PROXY_BASE_URL, + identity: Identity | None = None, + allow_upstream_consent: bool = True, + ) -> OauthToolRun: + 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, ...]: + return unwrap( + self.proxy.transport.get( + f"/v1/mcp/server/{server_id}/user-credentials", + headers=self.proxy.transport.master, + params=NoBody(), + response_type=McpServerUserCredentialListResponse, + ) + ).root + + def revoke_user_token(self, server_id: str, headers: AuthHeaders) -> None: + _ = unwrap( + self.proxy.transport.delete( + f"/v1/mcp/server/{server_id}/oauth-user-credential", + headers=headers, + json=NoBody(), + response_type=McpOauthUserCredentialStatus, + ) + ) + def chat_with_mcp(self, headers: AuthHeaders, body: ChatBody) -> ChatResponse: """POST /chat/completions carrying the LiteLLM key in `headers` (either ingress form) with an MCP server attached in `body.tools`. The gateway diff --git a/tests/e2e/mcp/oauth_gateway.py b/tests/e2e/mcp/oauth_gateway.py new file mode 100644 index 00000000000..82bb5f7ba0b --- /dev/null +++ b/tests/e2e/mcp/oauth_gateway.py @@ -0,0 +1,198 @@ +"""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 + +INHERITED_ENV_PREFIXES: Final = ("REDIS_", "MICROSOFT_", "GOOGLE_", "GENERIC_", "PROXY_") + + +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: + """Read the encrypted credential because management APIs omit the plaintext token.""" + 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: + gateway_token: str = field(default="", repr=False) + _seen: tuple[tuple[str, str, 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 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 + received: Final = headers.get("authorization", "") + gateway_leaked: Final = any( + value in (self.gateway_token, f"Bearer {self.gateway_token}") for value in headers.values() + ) + with self._lock: + self._seen = (*self._seen, (operation, received, gateway_leaked)) + + def assert_forwarded(self, expected: StoredOAuth) -> None: + with self._lock: + snapshot: Final = self._seen + self._seen = () + assert {item[0] for item in snapshot} == {"tools/list", "tools/call"}, "missing upstream observations" + expected_header: Final = f"Bearer {expected.access_token.get_secret_value()}" + assert all(item[1] == expected_header for item in snapshot), "upstream bearer did not match the stored token" + assert all(not item[2] for item in snapshot), "gateway bearer leaked to the upstream" + + +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(INHERITED_ENV_PREFIXES)}, + **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 new file mode 100644 index 00000000000..c20b73c0d63 --- /dev/null +++ b/tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py @@ -0,0 +1,207 @@ +"""Real OAuth consent, immediate MCP operations and cold-restart persistence. + +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 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, unique_marker +from e2e_http import AuthHeaders, NoBody, get_external, unwrap +from idp import Identity, Keycloak +from lifecycle import ResourceManager +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 + +pytestmark = [pytest.mark.e2e, pytest.mark.mcp_oauth_live, pytest.mark.provider_live] + + +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) + + +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") + @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, + client: ChatMcpClient, + resources: ResourceManager, + jwt_identity: Identity, + idp: Keycloak, + oauth_gateway: OAuthGateway, + route: Literal["aggregate_sso", "explicit_header_jwt"], + observed: bool, + ) -> None: + alias: Final = f"e2elinear{unique_marker()}" + tool: Final = f"{alias}-{LINEAR_READONLY_TOOL}" + token: Final = idp.access_token(jwt_identity) + observation: Final = OAuthObservation(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 + ) + 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, + 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=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: client.delete_server(created.server_id)) + assert client.server_user_credentials(created.server_id) == (), ( + "scenario must start without upstream credentials" + ) + 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, + ) + ) + client.proxy.update_team( + TeamUpdateBody( + team_id=jwt_identity.group, + object_permission=ObjectPermission(mcp_servers=[created.server_id]), + ) + ) + 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, + InMemoryTokenStorage(), + LINEAR_STORAGE_STATE, + tool, + {}, + base_url=oauth_gateway.base_url, + identity=identity, + ) + 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" + first_stored_oauth: Final = stored_oauth(jwt_identity.user_id, created.server_id) + if observed: + observation.assert_forwarded(first_stored_oauth) + oauth_gateway.restart() + fresh_token: Final = idp.access_token(jwt_identity) + observation.gateway_token = fresh_token + second: Final = client.list_and_call( + alias, + {"Authorization": f"Bearer {fresh_token}"} if identity is None else {}, + InMemoryTokenStorage(), + LINEAR_STORAGE_STATE if identity is not None else None, + tool, + {}, + base_url=oauth_gateway.base_url, + identity=identity, + allow_upstream_consent=False, + ) + assert_tool_result(second, tool) + second_stored_oauth: Final = stored_oauth(jwt_identity.user_id, created.server_id) + if observed: + observation.assert_forwarded(second_stored_oauth) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 9f49c5974d0..4b202e3c663 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -192,7 +192,7 @@ class ImageUrl(BaseModel): class TextContentPart(BaseModel): type: str = "text" text: str - cache_control: "CacheControl | None" = None + cache_control: CacheControl | None = None class ImageContentPart(BaseModel): @@ -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 @@ -584,8 +588,11 @@ class McpServerCreateBody(BaseModel): allow_all_keys: bool = True auth_type: str | None = None oauth2_flow: Literal["client_credentials", "authorization_code"] | None = None + 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 @@ -625,6 +632,26 @@ class McpServerListResponse(RootModel[list[McpServerRow]]): """GET /v1/mcp/server answers with a bare array of servers.""" +class McpServerUserCredentialRow(BaseModel): + user_id: str + credential_type: Literal["oauth2", "byok"] + expires_at: str | None = None + connected_at: str | None = None + updated_at: str + + +class McpServerUserCredentialListResponse(RootModel[tuple[McpServerUserCredentialRow, ...]]): + """GET /v1/mcp/server/{server_id}/user-credentials answers with a bare array.""" + + +class McpOauthUserCredentialStatus(BaseModel): + server_id: str + has_credential: bool + expires_at: str | None = None + is_expired: bool = False + connected_at: str | None = None + + class ToolsetTool(BaseModel): server_id: str tool_name: str @@ -1172,8 +1199,9 @@ class TeamNewResponse(BaseModel): class TeamUpdateBody(BaseModel): team_id: str - team_alias: str + team_alias: str | None = None models: list[str] | None = None + object_permission: ObjectPermission | None = None class TeamInfoParams(BaseModel): 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( diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 44d9df5e5c5..c6ede240c3b 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -89,6 +89,7 @@ from models import ( TeamDeleteBody, TeamNewBody, TeamNewResponse, + TeamUpdateBody, ToolsetCreateBody, ToolsetRow, ToolsetUpdateBody, @@ -871,6 +872,16 @@ class ProxyClient: ) ).team_id + def update_team(self, body: TeamUpdateBody) -> None: + unwrap( + self.transport.post( + "/team/update", + headers=self.transport.master, + json=body, + response_type=NoBody, + ) + ) + def delete_team(self, team_id: str) -> None: result = self.transport.post( "/team/delete", diff --git a/tests/e2e/pytest.ini b/tests/e2e/pytest.ini index f9e5995079b..97acb9ec52b 100644 --- a/tests/e2e/pytest.ini +++ b/tests/e2e/pytest.ini @@ -12,3 +12,4 @@ markers = prompt_caching_stack: needs a proxy running with router_settings.optional_pre_call_checks including prompt_caching; deselected unless E2E_PROMPT_CACHING_STACK is set cli_determinism: drives the real claude CLI for several seconds; deselected unless E2E_CLI_DETERMINISM is set redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_REDIS_CHAOS is set + mcp_oauth_live: real Linear OAuth consent via a captured browser session; deselected unless E2E_MCP_OAUTH_LIVE is set