mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
Merge pull request #42040 from BerriAI/revert-41909-litellm_mcp_oauth_happy_path_e2e
Revert "test(e2e): cover MCP OAuth SSO and cold restart persistence"
This commit is contained in:
joshua-berri
GitHub
commit
841193aab2
17 changed files with 30 additions and 942 deletions
15
.github/e2e-stack/assert_tests_ran.py
vendored
15
.github/e2e-stack/assert_tests_ran.py
vendored
|
|
@ -1,5 +1,3 @@
|
|||
import os
|
||||
import re
|
||||
import sys
|
||||
import xml.etree.ElementTree as ET
|
||||
from pathlib import Path
|
||||
|
|
@ -17,12 +15,6 @@ 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"))
|
||||
)
|
||||
|
|
@ -46,13 +38,6 @@ 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
|
||||
|
|
|
|||
1
.github/e2e-stack/select_tests.py
vendored
1
.github/e2e-stack/select_tests.py
vendored
|
|
@ -5,7 +5,6 @@ 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$"
|
||||
|
|
|
|||
177
.github/workflows/test-mcp-oauth-e2e.yml
vendored
177
.github/workflows/test-mcp-oauth-e2e.yml
vendored
|
|
@ -1,177 +0,0 @@
|
|||
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"
|
||||
|
|
@ -226,31 +226,3 @@ 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
|
||||
|
|
|
|||
|
|
@ -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` 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
|
||||
- `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
|
||||
- `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, 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
|
||||
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
|
||||
|
||||
- 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 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
|
||||
- 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
|
||||
|
||||
## Lay the pattern down in a class
|
||||
|
||||
|
|
@ -152,7 +152,7 @@ MCPs - endpoint features with the protocol op as the variant
|
|||
mcp.<operation>.<auth_family>.<assertion>
|
||||
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 | persists_across_processes
|
||||
assertion : succeeds | denied_without_permission
|
||||
e.g. mcp.call_tool.oauth.succeeds
|
||||
```
|
||||
|
||||
|
|
|
|||
|
|
@ -248,56 +248,3 @@ 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
|
||||
|
|
|
|||
|
|
@ -17,7 +17,6 @@ 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
|
||||
|
||||
|
|
@ -29,7 +28,6 @@ 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,
|
||||
|
|
@ -58,7 +56,6 @@ 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,
|
||||
}
|
||||
)
|
||||
|
||||
|
|
@ -135,11 +132,6 @@ 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:
|
||||
|
|
@ -221,8 +213,6 @@ 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)
|
||||
|
|
@ -246,13 +236,6 @@ 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
|
||||
|
|
|
|||
|
|
@ -71,14 +71,6 @@
|
|||
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
|
||||
|
|
|
|||
|
|
@ -52,7 +52,6 @@ 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
|
||||
|
|
@ -145,7 +144,6 @@ 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"))
|
||||
|
|
|
|||
|
|
@ -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__":
|
||||
|
|
|
|||
|
|
@ -18,7 +18,7 @@ import asyncio
|
|||
import re
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Final
|
||||
from typing import TYPE_CHECKING
|
||||
from urllib.parse import parse_qsl
|
||||
|
||||
import httpx
|
||||
|
|
@ -26,21 +26,11 @@ 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 mcp.types import TextContent
|
||||
from models import (
|
||||
ChatBody,
|
||||
ChatResponse,
|
||||
McpOauthUserCredentialStatus,
|
||||
McpServerCreateBody,
|
||||
McpServerInfo,
|
||||
McpServerUserCredentialListResponse,
|
||||
McpServerUserCredentialRow,
|
||||
)
|
||||
from models import ChatBody, ChatResponse, McpServerCreateBody, McpServerInfo
|
||||
from proxy_client import ProxyClient
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -54,8 +44,8 @@ OAUTH_CLIENT_REDIRECT_URI = "http://127.0.0.1:53682/e2e/callback"
|
|||
BROWSER_CONSENT_TIMEOUT = 60.0
|
||||
|
||||
|
||||
def _mcp_url(alias: str, base_url: str = PROXY_BASE_URL) -> str:
|
||||
return f"{base_url.rstrip('/')}/{alias}/mcp"
|
||||
def _mcp_url(alias: str) -> str:
|
||||
return f"{PROXY_BASE_URL}/{alias}/mcp"
|
||||
|
||||
|
||||
class InMemoryTokenStorage:
|
||||
|
|
@ -79,13 +69,7 @@ class InMemoryTokenStorage:
|
|||
self._client_info = client_info
|
||||
|
||||
|
||||
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]:
|
||||
async def _browser_follow_authorize(start_url: str, storage_state_path: str) -> 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
|
||||
|
|
@ -101,9 +85,6 @@ async def _browser_follow_authorize(
|
|||
|
||||
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
|
||||
|
||||
|
|
@ -115,7 +96,7 @@ async def _browser_follow_authorize(
|
|||
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()
|
||||
context.on("request", _note_request)
|
||||
page.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
|
||||
|
|
@ -124,29 +105,8 @@ async def _browser_follow_authorize(
|
|||
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 "upstream_consent" in captured or "url" in captured:
|
||||
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 "/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("@")'
|
||||
|
|
@ -158,44 +118,27 @@ async def _browser_follow_authorize(
|
|||
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, "client redirect_uri carried no authorization code"
|
||||
assert "code" in params, f"client redirect_uri carried no code: {landing}"
|
||||
return params["code"], params.get("state")
|
||||
|
||||
|
||||
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:
|
||||
def _oauth_provider(url: str, storage: InMemoryTokenStorage, storage_state_path: str) -> 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 _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
|
||||
)
|
||||
async def redirect_handler(authorize_url: str) -> None:
|
||||
code, state = await _browser_follow_authorize(authorize_url, storage_state_path)
|
||||
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"
|
||||
|
|
@ -224,45 +167,24 @@ 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], gateway_url: str) -> None:
|
||||
def __init__(self, inner: httpx2.AsyncBaseTransport, headers: dict[str, 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:
|
||||
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]
|
||||
for name, value in self._headers.items():
|
||||
if name not in request.headers:
|
||||
request.headers[name] = value
|
||||
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
|
||||
) -> httpx2.AsyncClient:
|
||||
def _oauth_http_client(headers: dict[str, str], auth: OAuthClientProvider) -> httpx2.AsyncClient:
|
||||
return httpx2.AsyncClient(
|
||||
headers=headers,
|
||||
auth=auth,
|
||||
timeout=httpx2.Timeout(REQUEST_TIMEOUT),
|
||||
follow_redirects=True,
|
||||
transport=_HeaderInjectingTransport(httpx2.AsyncHTTPTransport(), headers, gateway_url),
|
||||
transport=_HeaderInjectingTransport(httpx2.AsyncHTTPTransport(), headers),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -277,43 +199,6 @@ 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
|
||||
|
|
@ -367,53 +252,6 @@ 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
|
||||
|
|
|
|||
|
|
@ -1,198 +0,0 @@
|
|||
"""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
|
||||
|
|
@ -1,207 +0,0 @@
|
|||
"""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)
|
||||
|
|
@ -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,10 +572,6 @@ 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
|
||||
|
|
@ -588,11 +584,8 @@ 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
|
||||
|
|
@ -632,26 +625,6 @@ 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
|
||||
|
|
@ -1199,9 +1172,8 @@ class TeamNewResponse(BaseModel):
|
|||
|
||||
class TeamUpdateBody(BaseModel):
|
||||
team_id: str
|
||||
team_alias: str | None = None
|
||||
team_alias: str
|
||||
models: list[str] | None = None
|
||||
object_permission: ObjectPermission | None = None
|
||||
|
||||
|
||||
class TeamInfoParams(BaseModel):
|
||||
|
|
|
|||
|
|
@ -46,7 +46,7 @@ import os
|
|||
import re
|
||||
import threading
|
||||
from collections import deque
|
||||
from collections.abc import Callable, Generator, Mapping, Sequence
|
||||
from collections.abc import 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:
|
||||
observe_request: Callable[[str, Mapping[str, str], bytes | None], None] | None = None
|
||||
pass
|
||||
|
||||
|
||||
type EdgeBackend = RecordEdge | ReplayEdge | LiveEdge | CacheEdge
|
||||
|
|
@ -787,13 +787,10 @@ 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)
|
||||
|
|
@ -871,10 +868,9 @@ def handle_edge_request(
|
|||
method, _upstream_url(upstream_base, upstream_path, split.query), headers, body, timeout,
|
||||
backend, mount, test_key,
|
||||
)
|
||||
case LiveEdge(observe_request=observe_request):
|
||||
case LiveEdge():
|
||||
return _handle_live(
|
||||
method, _upstream_url(upstream_base, upstream_path, split.query), headers, body, timeout,
|
||||
observe_request=observe_request,
|
||||
method, _upstream_url(upstream_base, upstream_path, split.query), headers, body, timeout
|
||||
)
|
||||
case RecordEdge():
|
||||
return _handle_record(
|
||||
|
|
|
|||
|
|
@ -89,7 +89,6 @@ from models import (
|
|||
TeamDeleteBody,
|
||||
TeamNewBody,
|
||||
TeamNewResponse,
|
||||
TeamUpdateBody,
|
||||
ToolsetCreateBody,
|
||||
ToolsetRow,
|
||||
ToolsetUpdateBody,
|
||||
|
|
@ -872,16 +871,6 @@ 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",
|
||||
|
|
|
|||
|
|
@ -12,4 +12,3 @@ 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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue