mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-23 05:34:41 +00:00
test(mcp): replace the live Linear MCP OAuth e2e tests with integration coverage (#47154)
* test(mcp): move the MCP OAuth happy path from a live Linear e2e to an integration test GitHub Actions now runs unit tests only, and the MCP OAuth happy path needed a captured Linear browser login to click through a real consent page. The OAuth flows it exercised are already integration tests against an owned authorization-server double, except one behavior: the stored per-user token must be served by a proxy process that never saw the consent, with no shared Redis to fall back on. That case is now in test_mcp_oauth_flows.py The e2e test, its owned-gateway helper, its workflow, the mcp_oauth_live marker and the JUnit gate script only that workflow used are removed * test(mcp): chat completions run an OAuth server tool with the caller's own stored token * test(e2e): drop the Linear chat-completion MCP OAuth e2e and its helpers The behavior it checked, a chat completion running an OAuth server's tool with the caller's stored per-user token, is now an integration test in test_mcp_llm_endpoints.py that needs no saved Linear browser login * test(mcp): decode the stored per-user credential as urlsafe base64 The plaintext check used b64decode, which drops the - and _ of a urlsafe value and raised on roughly one ciphertext in five * test(mcp): drop the chat OAuth integration test that duplicates test_mcp_slots F2-chat security/test_mcp_slots.py already sends a per-user OAuth token through /v1/chat/completions and asserts the peer receives it. Every mutant the new test killed was also killed by the existing suite, so it added no signal
This commit is contained in:
yuneng-jiang
GitHub
parent
b3c0f3457d
commit
44fcbc30b0
20 changed files with 61 additions and 1674 deletions
69
.github/e2e-stack/assert_tests_ran.py
vendored
69
.github/e2e-stack/assert_tests_ran.py
vendored
|
|
@ -1,69 +0,0 @@
|
|||
import os
|
||||
import re
|
||||
import sys
|
||||
import xml.etree.ElementTree as ET
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "tests/e2e"))
|
||||
from coverage_registry.management_cases import MANAGEMENT_CASES
|
||||
|
||||
|
||||
def main() -> int:
|
||||
selected: Final = tuple(sys.argv[2:])
|
||||
try:
|
||||
report: Final = ET.parse(Path(sys.argv[1])).getroot()
|
||||
except (ET.ParseError, OSError):
|
||||
_ = 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")
|
||||
passed: Final = frozenset(
|
||||
case.get("file") for case in cases if all(case.find(tag) is None for tag in ("skipped", "failure", "error"))
|
||||
)
|
||||
missing: Final = tuple(path for path in selected if path not in passed)
|
||||
required_nodes: Final = frozenset(case.node for case in MANAGEMENT_CASES if case.node.split("::", 1)[0] in selected)
|
||||
passed_nodes: Final = frozenset(
|
||||
prop.get("value")
|
||||
for case in cases
|
||||
if all(case.find(tag) is None for tag in ("skipped", "failure", "error"))
|
||||
for prop in case.findall("./properties/property")
|
||||
if prop.get("name") == "management_node"
|
||||
)
|
||||
missing_nodes: Final = required_nodes - passed_nodes
|
||||
for node in sorted(missing_nodes):
|
||||
_ = sys.stdout.write(f"::error::required management case did not pass: {node}\n")
|
||||
for path in selected:
|
||||
collected: Final = sum(case.get("file") == path for case in cases)
|
||||
skipped: Final = sum(case.get("file") == path and case.find("skipped") is not None for case in cases)
|
||||
_ = sys.stdout.write(f"{path}: {collected} collected, {skipped} skipped\n")
|
||||
for case in cases:
|
||||
if case.get("file") != path or all(case.find(tag) is None for tag in ("failure", "error", "skipped")):
|
||||
continue
|
||||
outcome = "skipped" if case.find("skipped") is not None else "failed"
|
||||
_ = sys.stdout.write(f" {outcome}: {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 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
|
||||
if (
|
||||
selected
|
||||
and not missing
|
||||
and not missing_nodes
|
||||
and not any(case.find(tag) is not None for case in cases for tag in ("failure", "error"))
|
||||
):
|
||||
return 0
|
||||
_ = sys.stdout.write("::error::every selected file must execute a passing test, with no failures or errors\n")
|
||||
return 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
4
.github/workflows/test-linting.yml
vendored
4
.github/workflows/test-linting.yml
vendored
|
|
@ -539,8 +539,8 @@ jobs:
|
|||
- name: test_unit_passed_gate
|
||||
run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_unit_passed_gate.py
|
||||
|
||||
- name: test_e2e_changed_gate
|
||||
run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_e2e_changed_gate.py tests/code_coverage_tests/test_e2e_idp_stack.py
|
||||
- name: test_e2e_idp_stack
|
||||
run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_e2e_idp_stack.py
|
||||
|
||||
- name: test_e2e_metadata
|
||||
env:
|
||||
|
|
|
|||
180
.github/workflows/test-mcp-oauth-e2e.yml
vendored
180
.github/workflows/test-mcp-oauth-e2e.yml
vendored
|
|
@ -1,180 +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/mcp_management_endpoints.py'
|
||||
- 'litellm/proxy/_types.py'
|
||||
- 'litellm/types/mcp_server/mcp_server_manager.py'
|
||||
- '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"
|
||||
|
|
@ -25,7 +25,6 @@ ALLOWED_RAW_CLIENT_FILES = {
|
|||
"e2e/e2e_http.py": ("requests",),
|
||||
"e2e/conftest.py": ("requests",),
|
||||
"e2e/claude_code/pr_gate_version_resolver.py": ("urllib.request",),
|
||||
"e2e/mcp/oauth_chat_client.py": ("httpx",),
|
||||
}
|
||||
|
||||
EXCEPTION_ONLY_NAMES = frozenset({"RequestException", "ConnectionError", "Timeout", "HTTPError"})
|
||||
|
|
|
|||
|
|
@ -1,156 +0,0 @@
|
|||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import xml.etree.ElementTree as ET
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
GATE: Final = Path(__file__).resolve().parents[2] / ".github/e2e-stack/assert_tests_ran.py"
|
||||
SELECTED: Final = ("tests/e2e/access_control/test_a.py", "tests/e2e/access_control/test_b.py")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("second_outcome", "expected_status"),
|
||||
(("passed", 0), ("skipped", 1), ("failure", 1), ("error", 1), ("deselected", 1)),
|
||||
)
|
||||
def test_each_changed_file_must_run(tmp_path: Path, second_outcome: str, expected_status: int) -> None:
|
||||
suite: Final = ET.Element("testsuite")
|
||||
_ = ET.SubElement(suite, "testcase", file=SELECTED[0])
|
||||
if second_outcome != "deselected":
|
||||
second: Final = ET.SubElement(suite, "testcase", file=SELECTED[1])
|
||||
if second_outcome != "passed":
|
||||
_ = ET.SubElement(second, second_outcome)
|
||||
report: Final = tmp_path / "report.xml"
|
||||
ET.ElementTree(suite).write(report)
|
||||
|
||||
result: Final = subprocess.run([sys.executable, str(GATE), str(report), *SELECTED], capture_output=True, text=True)
|
||||
|
||||
assert result.returncode == expected_status, result.stdout
|
||||
|
||||
|
||||
@pytest.mark.parametrize("outcome", ("failure", "error"))
|
||||
def test_passing_case_does_not_hide_a_failure_in_the_same_file(tmp_path: Path, outcome: str) -> None:
|
||||
suite: Final = ET.Element("testsuite")
|
||||
_ = ET.SubElement(suite, "testcase", file=SELECTED[0])
|
||||
failed: Final = ET.SubElement(suite, "testcase", file=SELECTED[0])
|
||||
_ = ET.SubElement(failed, outcome)
|
||||
report: Final = tmp_path / "report.xml"
|
||||
ET.ElementTree(suite).write(report)
|
||||
|
||||
result: Final = subprocess.run(
|
||||
[sys.executable, str(GATE), str(report), SELECTED[0]], capture_output=True, text=True
|
||||
)
|
||||
|
||||
assert result.returncode == 1
|
||||
|
||||
|
||||
def test_failed_cases_are_named_per_selected_file(tmp_path: Path) -> None:
|
||||
suite: Final = ET.Element("testsuite")
|
||||
_ = ET.SubElement(suite, "testcase", file=SELECTED[0], classname="tests.e2e.access_control.test_a", name="test_ok")
|
||||
failed: Final = ET.SubElement(
|
||||
suite, "testcase", file=SELECTED[0], classname="tests.e2e.access_control.test_a", name="test_boom"
|
||||
)
|
||||
_ = ET.SubElement(failed, "failure", message="secret-bearing message")
|
||||
errored: Final = ET.SubElement(
|
||||
suite, "testcase", file=SELECTED[1], classname="tests.e2e.access_control.test_b", name="test_setup"
|
||||
)
|
||||
_ = ET.SubElement(errored, "error")
|
||||
report: Final = tmp_path / "report.xml"
|
||||
ET.ElementTree(suite).write(report)
|
||||
|
||||
result: Final = subprocess.run([sys.executable, str(GATE), str(report), *SELECTED], capture_output=True, text=True)
|
||||
|
||||
assert result.returncode == 1
|
||||
assert " failed: tests.e2e.access_control.test_a::test_boom\n" in result.stdout
|
||||
assert " failed: tests.e2e.access_control.test_b::test_setup\n" in result.stdout
|
||||
assert "test_ok" not in result.stdout
|
||||
assert "secret-bearing message" not in result.stdout
|
||||
|
||||
|
||||
@pytest.mark.parametrize("contents", ("<testsuite/>", "<testsuite", '<testsuite><testcase name="a"/></testsuite>'))
|
||||
def test_missing_execution_evidence_fails(tmp_path: Path, contents: str) -> None:
|
||||
report: Final = tmp_path / "report.xml"
|
||||
_ = report.write_text(contents)
|
||||
|
||||
result: Final = subprocess.run([sys.executable, str(GATE), str(report), *SELECTED], capture_output=True, text=True)
|
||||
|
||||
assert result.returncode == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omitted_role", ("proxy_admin", "team_member", "internal_user_viewer"))
|
||||
def test_one_passing_management_case_cannot_hide_a_missing_actor(tmp_path: Path, omitted_role: str) -> None:
|
||||
suite: Final = ET.Element("testsuite")
|
||||
path: Final = "tests/e2e/management/test_jwt_management_e2e.py"
|
||||
case: Final = ET.SubElement(suite, "testcase", file=path)
|
||||
properties: Final = ET.SubElement(case, "properties")
|
||||
_ = ET.SubElement(
|
||||
properties,
|
||||
"property",
|
||||
name="management_node",
|
||||
value=f"{path}::TestJwtManagement::test_actor_subject_and_database_role[proxy_admin_viewer]",
|
||||
)
|
||||
report: Final = tmp_path / "report.xml"
|
||||
ET.ElementTree(suite).write(report)
|
||||
result: Final = subprocess.run([sys.executable, str(GATE), str(report), path], capture_output=True, text=True)
|
||||
assert result.returncode == 1
|
||||
assert f"test_actor_subject_and_database_role[{omitted_role}]" in result.stdout
|
||||
|
||||
|
||||
@pytest.mark.parametrize("phase", ("setup", "call", "teardown"))
|
||||
@pytest.mark.parametrize("required_count", ("1", "4"))
|
||||
def test_oauth_failure_diagnostics_do_not_publish_private_payloads(
|
||||
tmp_path: Path, phase: str, required_count: 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,
|
||||
env={**os.environ, "E2E_REQUIRED_TEST_COUNT": required_count},
|
||||
)
|
||||
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
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("count", "skip", "expected"), ((0, False, 1), (3, False, 1), (4, False, 0), (5, False, 1), (4, True, 1))
|
||||
)
|
||||
def test_required_count_reports_cases_before_rejecting(tmp_path: Path, count: int, skip: bool, expected: int) -> None:
|
||||
suite = ET.Element("testsuite")
|
||||
for index in range(count):
|
||||
case = ET.SubElement(suite, "testcase", file=SELECTED[0], classname="OAuth", name=f"variant{index}")
|
||||
if skip and index == 0:
|
||||
ET.SubElement(case, "skipped", message="private-skip-reason")
|
||||
report = tmp_path / "report.xml"
|
||||
ET.ElementTree(suite).write(report)
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-I", str(GATE), str(report), SELECTED[0]],
|
||||
env={**os.environ, "E2E_REQUIRED_TEST_COUNT": "4"},
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
assert result.returncode == expected
|
||||
assert f"{count} collected, {int(skip)} skipped" in result.stdout
|
||||
if skip:
|
||||
assert "skipped: OAuth::variant0" in result.stdout
|
||||
assert "private-skip-reason" not in result.stdout + result.stderr
|
||||
|
|
@ -172,12 +172,7 @@ def open_consent() -> None:
|
|||
raise RuntimeError("consent page timed out")
|
||||
|
||||
|
||||
@pytest.mark.mcp_oauth_live
|
||||
def test_oauth_dies_on_consent() -> None:
|
||||
open_consent()
|
||||
|
||||
|
||||
def test_plain_dies_on_consent() -> None:
|
||||
def test_dies_on_consent() -> None:
|
||||
open_consent()
|
||||
"""
|
||||
|
||||
|
|
@ -353,16 +348,13 @@ class TestStepsReachTheReport:
|
|||
|
||||
def test_a_failed_phase_s_own_report_carries_the_steps(tmp_path: Path) -> None:
|
||||
"""Plugins that read the failed setup or call report, not the teardown one
|
||||
junitxml writes from, see where the test died too, oauth-live or not."""
|
||||
junitxml writes from, see where the test died too."""
|
||||
write_suite(tmp_path, {"test_consent.py": FAILED_PHASE_SUITE, "report_spy.py": REPORT_SPY_PLUGIN})
|
||||
child: Final = run_child_pytest(tmp_path, "-p", "report_spy", env={"E2E_MCP_OAUTH_LIVE": "1"})
|
||||
child: Final = run_child_pytest(tmp_path, "-p", "report_spy")
|
||||
seen_path: Final = tmp_path / "failed-reports.jsonl"
|
||||
assert seen_path.exists(), f"no failed report reached the spy:\n{child.stdout}\n{child.stderr}"
|
||||
seen: Final = dict(map(FailedReport.validate_json, seen_path.read_text().splitlines()))
|
||||
assert seen == {
|
||||
"test_oauth_dies_on_consent": ("open the consent page",),
|
||||
"test_plain_dies_on_consent": ("open the consent page",),
|
||||
}, child.stdout
|
||||
assert seen == {"test_dies_on_consent": ("open the consent page",)}, child.stdout
|
||||
|
||||
|
||||
class TestDeclaredPropertiesReachTheReport:
|
||||
|
|
|
|||
|
|
@ -40,7 +40,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)
|
||||
- `logging/` - logging-integration delivery (datadog and friends)
|
||||
- `security/` - secret handling and log-leak protection
|
||||
- `router/` - routing and reliability behavior (fallbacks, cooldowns) plus the memory tests (`test_reliability_memory_e2e.py`: every worker's RSS as read at collection time, before any test traffic, must sit under a fixed idle budget, the release-gate check for a DB-backed boot that idles near the pod limit the way v1.100.x did, and a few hundred failing requests with retries and fallbacks must not grow proxy RSS past a loose ceiling over the measured batches nor store a request snapshot past a fixed size, the release-gate check for the v1.100.0 retry-breadcrumb leak)
|
||||
|
|
@ -53,14 +53,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
|
||||
- Gateway-managed OAuth (an authorize/token dance and per-user stored tokens) is out of scope here, since Datadog only takes static headers. Those flows are integration tests against an owned authorization-server double, in `tests/integration/mcp/test_mcp_oauth_flows.py` and, for a chat completion calling the tool with the caller's stored token, the F2-chat cases of `tests/integration/security/test_mcp_slots.py`. Do not add a third-party OAuth upstream that needs a saved browser login
|
||||
|
||||
## Lay the pattern down in a class
|
||||
|
||||
|
|
@ -243,7 +243,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
|
||||
```
|
||||
|
||||
|
|
|
|||
|
|
@ -261,55 +261,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 `python` job in `.github/workflows/test-linting.yml` runs them alongside the provider-edge and fixture tests at the root of `tests/e2e_harness/` with `--noconftest -o addopts= -o "pythonpath=tests/e2e tests/e2e_harness" -p no:rerunfailures`; 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.
|
||||
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,
|
||||
OTEL_TLS_OPT_IN_ENV,
|
||||
OTEL_V2_OPT_IN_ENV,
|
||||
OWNED_GATEWAY_OPT_IN_ENV,
|
||||
|
|
@ -72,7 +70,6 @@ OPT_IN_MARKERS: Final = MappingProxyType(
|
|||
"redis_chaos": PERF_TESTS_OPT_IN_ENV,
|
||||
"perf": PERF_TESTS_OPT_IN_ENV,
|
||||
"cli_determinism": CLI_DETERMINISM_OPT_IN_ENV,
|
||||
"mcp_oauth_live": MCP_OAUTH_LIVE_OPT_IN_ENV,
|
||||
"provider_edge_host": PROVIDER_EDGE_HOST_OPT_IN_ENV,
|
||||
"owned_gateway": OWNED_GATEWAY_OPT_IN_ENV,
|
||||
"otel_v2": OTEL_V2_OPT_IN_ENV,
|
||||
|
|
@ -170,11 +167,6 @@ def pytest_configure(config: pytest.Config) -> None:
|
|||
"quiet_stack: measures the proxy itself, so it runs while no other test on this host is hitting the stack; "
|
||||
"every other test waits for it to finish",
|
||||
)
|
||||
config.addinivalue_line(
|
||||
"markers",
|
||||
"mcp_oauth_live: real Linear OAuth consent via a captured browser session; deselected unless "
|
||||
"E2E_MCP_OAUTH_LIVE is set",
|
||||
)
|
||||
config.addinivalue_line(
|
||||
"markers",
|
||||
"provider_edge_host: routes provider traffic through the pytest host's edge in every fixture mode, so the "
|
||||
|
|
@ -331,8 +323,6 @@ def pytest_runtest_setup(item: pytest.Item) -> None:
|
|||
item.user_properties.extend(item.config.stash[_IDLE_RSS].junit_properties)
|
||||
if not _reaches_proxy(item):
|
||||
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)
|
||||
|
|
@ -376,12 +366,6 @@ def pytest_runtest_makereport(
|
|||
report = yield
|
||||
if report.when in ("setup", "call"):
|
||||
attach_step_properties(item)
|
||||
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}"))
|
||||
if report.when == "call":
|
||||
item.stash[_CALL_PASSED] = report.passed
|
||||
report.user_properties = list(item.user_properties)
|
||||
|
|
|
|||
|
|
@ -79,14 +79,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
|
||||
|
|
|
|||
|
|
@ -111,10 +111,6 @@ CHEAP_ANTHROPIC_MODEL = os.environ.get("E2E_CHEAP_ANTHROPIC_MODEL", "claude-haik
|
|||
CHEAP_OPENAI_MODEL = os.environ.get("E2E_CHEAP_OPENAI_MODEL", "gpt-5.5")
|
||||
S3_PARTITION_GRANULARITY = os.environ.get("E2E_S3_PARTITION_GRANULARITY", "day")
|
||||
|
||||
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
|
||||
# read exported spans back through it.
|
||||
|
|
@ -207,7 +203,6 @@ MANAGED_FILES_OPT_IN_ENV = "E2E_MANAGED_FILES_STACK"
|
|||
PROMPT_CACHING_OPT_IN_ENV = "E2E_PROMPT_CACHING_STACK"
|
||||
PERF_TESTS_OPT_IN_ENV: Final = "E2E_PERF_TESTS"
|
||||
CLI_DETERMINISM_OPT_IN_ENV = "E2E_CLI_DETERMINISM"
|
||||
MCP_OAUTH_LIVE_OPT_IN_ENV: Final = "E2E_MCP_OAUTH_LIVE"
|
||||
PROVIDER_EDGE_HOST_OPT_IN_ENV: Final = "E2E_PROVIDER_EDGE_HOST_REACHABLE"
|
||||
OWNED_GATEWAY_OPT_IN_ENV: Final = "E2E_OWNED_GATEWAY"
|
||||
OTEL_V2_OPT_IN_ENV: Final = "E2E_OTEL_V2"
|
||||
|
|
|
|||
|
|
@ -1,58 +0,0 @@
|
|||
"""One-time helper to capture a logged-in Linear browser session for the
|
||||
real-Linear MCP e2e test.
|
||||
|
||||
The real-Linear test drives the genuine gateway-managed authorization_code
|
||||
dance against ``mcp.linear.app``. The only step that cannot be scripted is
|
||||
Linear's login (magic link / SSO), so a human authenticates once here and the
|
||||
resulting session (cookies + local storage) is persisted to disk. The e2e test
|
||||
then loads that session in a headless Playwright context and clicks Approve on
|
||||
Linear's consent screen every run, with no human and no login automation.
|
||||
|
||||
Run it with the e2e venv, log into Linear in the window that opens, then return
|
||||
to the terminal and press Enter:
|
||||
|
||||
LITELLM=~/litellm-mcpe2e
|
||||
"$LITELLM"/.venv/bin/python "$LITELLM"/tests/e2e/mcp/linear_session_capture.py
|
||||
|
||||
The session is written to ``E2E_LINEAR_STORAGE_STATE`` (default
|
||||
``~/.litellm-e2e/linear_storage_state.json``), outside the repo. It is a
|
||||
secret: never commit it. Re-run this whenever Linear expires the session.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
from playwright.sync_api import sync_playwright
|
||||
|
||||
DEFAULT_STATE_PATH = Path.home() / ".litellm-e2e" / "linear_storage_state.json"
|
||||
|
||||
|
||||
def capture(state_path: Path) -> None:
|
||||
"""Open a headed browser at Linear, wait for the human to log in, then save
|
||||
the authenticated session to ``state_path``."""
|
||||
state_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with sync_playwright() as playwright:
|
||||
browser = playwright.chromium.launch(headless=False)
|
||||
context = browser.new_context()
|
||||
page = context.new_page()
|
||||
page.goto("https://linear.app/login", wait_until="domcontentloaded")
|
||||
print("\n" + "=" * 72)
|
||||
print("Log into Linear in the browser window that just opened.")
|
||||
print("If Linear emails you a magic link, paste the link into THIS window's")
|
||||
print("address bar (opening it in your default browser won't capture the")
|
||||
print("session). Google SSO works too as long as you complete it here.")
|
||||
print("When your Linear workspace has loaded, come back and press Enter.")
|
||||
print("=" * 72)
|
||||
input("Press Enter once you are logged in... ")
|
||||
page.goto("https://mcp.linear.app/", wait_until="domcontentloaded")
|
||||
context.storage_state(path=str(state_path))
|
||||
browser.close()
|
||||
print(f"\nSaved Linear session to {state_path}")
|
||||
print("Point the e2e test at it with:")
|
||||
print(f' export E2E_LINEAR_STORAGE_STATE="{state_path}"')
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
capture(Path(os.environ.get("E2E_LINEAR_STORAGE_STATE", str(DEFAULT_STATE_PATH))))
|
||||
|
|
@ -1,442 +0,0 @@
|
|||
"""Client for the mcp chat-completion OAuth e2e suite.
|
||||
|
||||
Registers a gateway-managed OAuth (authorization_code) MCP server, seeds the
|
||||
per-user upstream token by driving the interactive authorize dance with the
|
||||
official mcp SDK's OAuthClientProvider (the browser leg is a headless Chromium
|
||||
primed with a human's saved Linear session), then exercises the server through
|
||||
/chat/completions, where the gateway lists and executes its tools with the
|
||||
stored per-user token.
|
||||
|
||||
Management routes (/v1/mcp/server CRUD, /chat/completions) go through the
|
||||
shared ProxyClient transport. The MCP protocol used to seed the token goes through
|
||||
the mcp SDK, the same library production MCP hosts run.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import re
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Final
|
||||
from urllib.parse import parse_qsl
|
||||
|
||||
import httpx
|
||||
import httpx2
|
||||
import pytest
|
||||
from e2e_config import PROXY_BASE_URL, REQUEST_TIMEOUT
|
||||
from e2e_http import AuthHeaders, NoBody, unwrap
|
||||
from e2e_metadata import step
|
||||
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 proxy_client import ProxyClient
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from playwright.async_api import Route
|
||||
|
||||
# Where the "browser" lands at the end of the authorize dance. Nothing listens
|
||||
# here: the route interceptor short-circuits the final redirect and reads the
|
||||
# code/state off its query string, exactly like a desktop MCP host intercepting
|
||||
# its loopback redirect.
|
||||
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"
|
||||
|
||||
|
||||
class InMemoryTokenStorage:
|
||||
"""The mcp SDK's TokenStorage protocol, in memory for one dance: the
|
||||
DCR-registered client and the gateway tokens minted for it."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._tokens: OAuthToken | None = None
|
||||
self._client_info: OAuthClientInformationFull | None = None
|
||||
|
||||
async def get_tokens(self) -> OAuthToken | None:
|
||||
return self._tokens
|
||||
|
||||
async def set_tokens(self, tokens: OAuthToken) -> None:
|
||||
self._tokens = tokens
|
||||
|
||||
async def get_client_info(self) -> OAuthClientInformationFull | None:
|
||||
return self._client_info
|
||||
|
||||
async def set_client_info(self, client_info: OAuthClientInformationFull) -> None:
|
||||
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]:
|
||||
"""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
|
||||
clicks through Linear's consent screens (the mcp.linear.app Approve form,
|
||||
then the linear.app workspace-selection page), riding the rest of the chain
|
||||
(Linear -> gateway callback -> host redirect_uri). The final hop is
|
||||
intercepted and short-circuited, since nothing listens there, and its
|
||||
code/state are read off the query string."""
|
||||
from playwright.async_api import async_playwright
|
||||
|
||||
captured: dict[str, str] = {} # mutable-ok: hand-off from the request listener
|
||||
trail: list[str] = [] # mutable-ok: navigation diagnostics for a failed dance
|
||||
|
||||
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
|
||||
|
||||
async def _swallow_redirect(route: Route) -> None:
|
||||
await route.fulfill(status=200, content_type="text/plain", body="ok")
|
||||
|
||||
async with async_playwright() as playwright:
|
||||
browser = await playwright.chromium.launch(headless=True)
|
||||
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("framenavigated", lambda frame: trail.append(frame.url.split("?", 1)[0]))
|
||||
await page.goto(start_url, wait_until="domcontentloaded")
|
||||
deadline = time.monotonic() + BROWSER_CONSENT_TIMEOUT
|
||||
while "url" not in captured and time.monotonic() < deadline:
|
||||
try:
|
||||
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:
|
||||
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("@")'
|
||||
).first
|
||||
try:
|
||||
await control.click(timeout=5000)
|
||||
except Exception: # noqa: BLE001 - nothing to advance yet; loop and re-check
|
||||
await asyncio.sleep(0.5)
|
||||
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"
|
||||
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:
|
||||
"""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
|
||||
)
|
||||
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"
|
||||
return AuthorizationCodeResult(code=code, state=code_holder.get("state"))
|
||||
|
||||
return OAuthClientProvider(
|
||||
server_url=url,
|
||||
client_metadata=OAuthClientMetadata.model_validate(
|
||||
{
|
||||
"redirect_uris": [OAUTH_CLIENT_REDIRECT_URI],
|
||||
"token_endpoint_auth_method": "none",
|
||||
"grant_types": ["authorization_code", "refresh_token"],
|
||||
"response_types": ["code"],
|
||||
"client_name": "e2e-mcp-host",
|
||||
}
|
||||
),
|
||||
storage=storage,
|
||||
redirect_handler=redirect_handler,
|
||||
callback_handler=callback_handler,
|
||||
)
|
||||
|
||||
|
||||
class _HeaderInjectingTransport(httpx2.AsyncBaseTransport):
|
||||
"""Adds the caller's LiteLLM key header to every outgoing SDK request
|
||||
(discovery, DCR, token exchange), so the gateway resolves which user to
|
||||
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:
|
||||
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]
|
||||
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:
|
||||
return httpx2.AsyncClient(
|
||||
auth=auth,
|
||||
timeout=httpx2.Timeout(REQUEST_TIMEOUT),
|
||||
follow_redirects=True,
|
||||
transport=_HeaderInjectingTransport(httpx2.AsyncHTTPTransport(), headers, gateway_url),
|
||||
)
|
||||
|
||||
|
||||
async def _seed_via_dance(
|
||||
url: str, headers: dict[str, str], storage: InMemoryTokenStorage, storage_state_path: str
|
||||
) -> tuple[str, ...]:
|
||||
async with _oauth_http_client(headers, _oauth_provider(url, storage, storage_state_path)) 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 = await session.list_tools()
|
||||
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
|
||||
|
||||
@step("Register the MCP server with the alias {body.alias}")
|
||||
def create_server(self, body: McpServerCreateBody) -> McpServerInfo:
|
||||
return unwrap(
|
||||
self.proxy.transport.post(
|
||||
"/v1/mcp/server",
|
||||
headers=self.proxy.transport.master,
|
||||
json=body,
|
||||
response_type=McpServerInfo,
|
||||
)
|
||||
)
|
||||
|
||||
@step("Read the MCP server back from /v1/mcp/server")
|
||||
def server_info(self, server_id: str) -> McpServerInfo:
|
||||
return unwrap(
|
||||
self.proxy.transport.get(
|
||||
f"/v1/mcp/server/{server_id}",
|
||||
headers=self.proxy.transport.master,
|
||||
params=NoBody(),
|
||||
response_type=McpServerInfo,
|
||||
)
|
||||
)
|
||||
|
||||
@step("Delete the MCP server")
|
||||
def delete_server(self, server_id: str) -> None:
|
||||
_ = self.proxy.transport.delete(
|
||||
f"/v1/mcp/server/{server_id}",
|
||||
headers=self.proxy.transport.master,
|
||||
json=NoBody(),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
||||
@step("Sign the key's user in to the MCP server {alias} through the OAuth consent flow")
|
||||
def seed_user_token(self, alias: str, key: str, storage_state_path: str) -> tuple[str, ...]:
|
||||
"""Drive the interactive authorize dance for `key`'s user so the gateway
|
||||
stores their upstream token, retried to the shared deadline since the
|
||||
just-created server and key propagate asynchronously. The LiteLLM key
|
||||
rides x-litellm-api-key so the gateway binds the token to that user.
|
||||
Returns the upstream tool names the dance listed, proof the token works."""
|
||||
headers = {"x-litellm-api-key": f"Bearer {key}"}
|
||||
storage = InMemoryTokenStorage()
|
||||
deadline = time.monotonic() + self.proxy.poll_timeout
|
||||
last_error: Exception | None = None
|
||||
while time.monotonic() < deadline:
|
||||
try:
|
||||
return asyncio.run(_seed_via_dance(_mcp_url(alias), headers, storage, storage_state_path))
|
||||
except Exception as exc: # noqa: BLE001 - retried to the deadline; the last error surfaces below
|
||||
last_error = exc
|
||||
time.sleep(self.proxy.poll_interval)
|
||||
pytest.fail(
|
||||
f"authorize dance for {alias!r} never completed within {self.proxy.poll_timeout}s; "
|
||||
f"last error: {last_error!r}"
|
||||
)
|
||||
|
||||
@step("List the tools on the MCP server {alias} and call {tool} over the MCP protocol")
|
||||
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,
|
||||
)
|
||||
)
|
||||
|
||||
@step("List the users with a stored OAuth token for the MCP server")
|
||||
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
|
||||
|
||||
@step("Revoke the user's stored OAuth token for the MCP server")
|
||||
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,
|
||||
)
|
||||
)
|
||||
|
||||
@step("Send a /chat/completions request to {body.model} with an MCP server attached as a tool")
|
||||
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
|
||||
resolves the user from the key and lists/executes the server's tools
|
||||
with that user's stored upstream token."""
|
||||
return unwrap(
|
||||
self.proxy.transport.post(
|
||||
"/chat/completions",
|
||||
headers=headers,
|
||||
json=body,
|
||||
response_type=ChatResponse,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def build_chat_client(proxy: ProxyClient) -> ChatMcpClient:
|
||||
return ChatMcpClient(proxy=proxy)
|
||||
|
|
@ -1,197 +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 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_config import INHERITED_ENV_PREFIXES, available_port
|
||||
from e2e_http import NoBody
|
||||
from e2e_metadata import step
|
||||
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, ValidationError
|
||||
|
||||
|
||||
class StoredOAuth(BaseModel):
|
||||
type: str
|
||||
access_token: SecretStr
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CredentialRow:
|
||||
credential_b64: str = field(repr=False)
|
||||
|
||||
|
||||
@step("Read the user's stored OAuth credential for the MCP server from the database and decrypt it")
|
||||
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"
|
||||
|
||||
|
||||
@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)
|
||||
|
||||
@step("Start the separate LiteLLM proxy and wait for /health/liveliness")
|
||||
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")
|
||||
|
||||
@step("Stop the separate LiteLLM proxy")
|
||||
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"
|
||||
|
||||
@step("Restart the separate LiteLLM proxy process so its in-memory caches start empty")
|
||||
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"
|
||||
|
||||
|
||||
@step("Start a separate LiteLLM proxy from source with JWT auth against Keycloak")
|
||||
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,),
|
||||
control_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,218 +0,0 @@
|
|||
"""On-demand e2e: a chat completion drives a gateway-managed OAuth MCP server.
|
||||
|
||||
The real end-user flow for MCP over an OAuth server: a user registers a Linear
|
||||
authorization_code server, authorizes it once so the gateway stores their
|
||||
upstream token, then sends a normal /chat/completions request with the Linear
|
||||
MCP attached. The gateway resolves the user from the LiteLLM key, lists Linear's
|
||||
tools with the stored per-user token, lets the model call one, executes it
|
||||
upstream with that token, and returns the answer. This is proven against the
|
||||
real Linear MCP server (mcp.linear.app) and a real Anthropic model, once per
|
||||
documented ingress header (x-litellm-api-key and Authorization).
|
||||
|
||||
The authorize dance is seeded through the mcp SDK's OAuthClientProvider; the one
|
||||
step Linear cannot auto-approve is the human consent, so it is captured once out
|
||||
of band (mcp/linear_session_capture.py) into a saved browser session and a
|
||||
headless Chromium clicks Approve every run. The test therefore skips unless
|
||||
E2E_LINEAR_STORAGE_STATE points at that session, so it never runs on the per-PR
|
||||
CI path; it is a nightly/on-demand real-server smoke test.
|
||||
|
||||
Fail-before-fix: without the stored per-user token the gateway lists no Linear
|
||||
tools, so mcp_list_tools comes back empty, nothing is called, and the
|
||||
assertions fail; a served, called, non-empty Linear tool proves the gateway
|
||||
pulled and used the user's token.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from e2e_config import CHEAP_ANTHROPIC_MODEL, LINEAR_MCP_URL, LINEAR_STORAGE_STATE, unique_marker
|
||||
from e2e_http import AuthHeaders
|
||||
from e2e_metadata import Capability, Domain, Mode, Provider, Route, Subject, meta
|
||||
from lifecycle import ResourceManager
|
||||
from models import ChatBody, ChatMessage, KeyGenerateBody, McpChatTool, McpServerCreateBody, ObjectPermission
|
||||
from proxy_client import ProxyClient
|
||||
|
||||
pytest.importorskip("mcp", reason="mcp SDK not installed; run `uv sync --inexact --group e2e-dev`")
|
||||
pytest.importorskip(
|
||||
"playwright.async_api",
|
||||
reason="playwright not installed; run `uv pip install playwright` and `playwright install chromium`",
|
||||
)
|
||||
|
||||
from oauth_chat_client import ChatMcpClient, build_chat_client # noqa: E402 # imports follow the importorskip guards
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.e2e,
|
||||
pytest.mark.skipif(
|
||||
not LINEAR_STORAGE_STATE or not os.path.exists(LINEAR_STORAGE_STATE),
|
||||
reason="set E2E_LINEAR_STORAGE_STATE to a Linear session captured via mcp/linear_session_capture.py",
|
||||
),
|
||||
]
|
||||
|
||||
# Pinned from a live dance during verification (never guessed); the gateway
|
||||
# prefixes every upstream tool name with the server alias. list_teams is a
|
||||
# read-only Linear tool that takes no arguments and returns the caller's teams.
|
||||
LINEAR_READONLY_TOOL = "list_teams"
|
||||
LINEAR_PROMPT = "Use the list_teams tool to list my Linear teams, then reply with the name of one of them."
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def chat_client(proxy: ProxyClient) -> ChatMcpClient:
|
||||
return build_chat_client(proxy)
|
||||
|
||||
|
||||
class TestMcpChatCompletionOauth:
|
||||
"""A scoped internal-user key on a real Linear authorization_code server,
|
||||
used through /chat/completions once per ingress header: the gateway pulls
|
||||
the user's stored upstream token, lists and executes Linear's tools during
|
||||
the completion, and returns the answer."""
|
||||
|
||||
@pytest.mark.covers("mcp.list_tools.oauth.succeeds")
|
||||
@pytest.mark.covers("mcp.call_tool.oauth.succeeds")
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.MCP,
|
||||
route=Route.CHAT_COMPLETIONS,
|
||||
providers=(Provider.ANTHROPIC,),
|
||||
models=(CHEAP_ANTHROPIC_MODEL,),
|
||||
capabilities=(Capability.FUNCTION_CALLING,),
|
||||
mode=Mode.NONSTREAM,
|
||||
)
|
||||
)
|
||||
def test_chat_completion_uses_linear_with_x_litellm_api_key_header(
|
||||
self, chat_client: ChatMcpClient, resources: ResourceManager
|
||||
) -> None:
|
||||
marker = unique_marker()
|
||||
alias = f"e2elinear{marker}"
|
||||
created = chat_client.create_server(
|
||||
McpServerCreateBody(
|
||||
alias=alias,
|
||||
url=LINEAR_MCP_URL,
|
||||
allow_all_keys=False,
|
||||
auth_type="oauth2",
|
||||
oauth2_flow="authorization_code",
|
||||
)
|
||||
)
|
||||
resources.defer(lambda: chat_client.delete_server(created.server_id))
|
||||
|
||||
stored = chat_client.server_info(created.server_id)
|
||||
assert stored.auth_type == "oauth2"
|
||||
assert stored.oauth2_flow == "authorization_code"
|
||||
assert stored.allow_all_keys is False
|
||||
|
||||
key = chat_client.proxy.generate_key(
|
||||
KeyGenerateBody(
|
||||
user_id="e2e-test-user",
|
||||
object_permission=ObjectPermission(mcp_servers=[created.server_id]),
|
||||
)
|
||||
)
|
||||
resources.defer(lambda: chat_client.proxy.delete_key(key))
|
||||
|
||||
seeded = chat_client.seed_user_token(alias, key, LINEAR_STORAGE_STATE)
|
||||
assert f"{alias}-{LINEAR_READONLY_TOOL}" in seeded, (
|
||||
f"the authorize dance listed {seeded}, expected it to include {alias}-{LINEAR_READONLY_TOOL}"
|
||||
)
|
||||
|
||||
response = chat_client.chat_with_mcp(
|
||||
AuthHeaders.model_validate({"x-litellm-api-key": f"Bearer {key}"}),
|
||||
ChatBody(
|
||||
model=CHEAP_ANTHROPIC_MODEL,
|
||||
messages=[ChatMessage(role="user", content=LINEAR_PROMPT)],
|
||||
tools=[
|
||||
McpChatTool(
|
||||
server_url=f"litellm_proxy/mcp/{alias}",
|
||||
server_label=alias,
|
||||
require_approval="never",
|
||||
)
|
||||
],
|
||||
),
|
||||
)
|
||||
|
||||
message = response.choices[0].message
|
||||
assert message is not None and message.content, f"completion returned no answer: {response}"
|
||||
meta = message.provider_specific_fields
|
||||
assert meta is not None, f"no MCP metadata on the completion: {response}"
|
||||
listed = {t.function.name for t in (meta.mcp_list_tools or []) if t.function}
|
||||
assert f"{alias}-{LINEAR_READONLY_TOOL}" in listed, (
|
||||
f"the gateway listed {sorted(listed)}, expected the stored token to surface {alias}-{LINEAR_READONLY_TOOL}"
|
||||
)
|
||||
results = [r for r in (meta.mcp_call_results or []) if r.name == f"{alias}-{LINEAR_READONLY_TOOL}"]
|
||||
assert results and results[0].result, (
|
||||
f"Linear tool {alias}-{LINEAR_READONLY_TOOL} was not executed with a result: {meta.mcp_call_results}"
|
||||
)
|
||||
|
||||
@pytest.mark.covers("mcp.list_tools.oauth.succeeds")
|
||||
@pytest.mark.covers("mcp.call_tool.oauth.succeeds")
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.MCP,
|
||||
route=Route.CHAT_COMPLETIONS,
|
||||
providers=(Provider.ANTHROPIC,),
|
||||
models=(CHEAP_ANTHROPIC_MODEL,),
|
||||
capabilities=(Capability.FUNCTION_CALLING,),
|
||||
mode=Mode.NONSTREAM,
|
||||
)
|
||||
)
|
||||
def test_chat_completion_uses_linear_with_authorization_bearer_header(
|
||||
self, chat_client: ChatMcpClient, resources: ResourceManager
|
||||
) -> None:
|
||||
marker = unique_marker()
|
||||
alias = f"e2elinear{marker}"
|
||||
created = chat_client.create_server(
|
||||
McpServerCreateBody(
|
||||
alias=alias,
|
||||
url=LINEAR_MCP_URL,
|
||||
allow_all_keys=False,
|
||||
auth_type="oauth2",
|
||||
oauth2_flow="authorization_code",
|
||||
)
|
||||
)
|
||||
resources.defer(lambda: chat_client.delete_server(created.server_id))
|
||||
|
||||
stored = chat_client.server_info(created.server_id)
|
||||
assert stored.auth_type == "oauth2"
|
||||
assert stored.oauth2_flow == "authorization_code"
|
||||
assert stored.allow_all_keys is False
|
||||
|
||||
key = chat_client.proxy.generate_key(
|
||||
KeyGenerateBody(
|
||||
user_id="e2e-test-user",
|
||||
object_permission=ObjectPermission(mcp_servers=[created.server_id]),
|
||||
)
|
||||
)
|
||||
resources.defer(lambda: chat_client.proxy.delete_key(key))
|
||||
|
||||
seeded = chat_client.seed_user_token(alias, key, LINEAR_STORAGE_STATE)
|
||||
assert f"{alias}-{LINEAR_READONLY_TOOL}" in seeded, (
|
||||
f"the authorize dance listed {seeded}, expected it to include {alias}-{LINEAR_READONLY_TOOL}"
|
||||
)
|
||||
|
||||
response = chat_client.chat_with_mcp(
|
||||
AuthHeaders.model_validate({"authorization": f"Bearer {key}"}),
|
||||
ChatBody(
|
||||
model=CHEAP_ANTHROPIC_MODEL,
|
||||
messages=[ChatMessage(role="user", content=LINEAR_PROMPT)],
|
||||
tools=[
|
||||
McpChatTool(
|
||||
server_url=f"litellm_proxy/mcp/{alias}",
|
||||
server_label=alias,
|
||||
require_approval="never",
|
||||
)
|
||||
],
|
||||
),
|
||||
)
|
||||
|
||||
message = response.choices[0].message
|
||||
assert message is not None and message.content, f"completion returned no answer: {response}"
|
||||
meta = message.provider_specific_fields
|
||||
assert meta is not None, f"no MCP metadata on the completion: {response}"
|
||||
listed = {t.function.name for t in (meta.mcp_list_tools or []) if t.function}
|
||||
assert f"{alias}-{LINEAR_READONLY_TOOL}" in listed, (
|
||||
f"the gateway listed {sorted(listed)}, expected the stored token to surface {alias}-{LINEAR_READONLY_TOOL}"
|
||||
)
|
||||
results = [r for r in (meta.mcp_call_results or []) if r.name == f"{alias}-{LINEAR_READONLY_TOOL}"]
|
||||
assert results and results[0].result, (
|
||||
f"Linear tool {alias}-{LINEAR_READONLY_TOOL} was not executed with a result: {meta.mcp_call_results}"
|
||||
)
|
||||
|
|
@ -1,216 +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 e2e_metadata import Domain, Route, Subject, meta
|
||||
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):
|
||||
issuer: str
|
||||
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")
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.MCP,
|
||||
route=Route.MCP,
|
||||
)
|
||||
)
|
||||
@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",
|
||||
issuer=metadata.issuer if metadata else None,
|
||||
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)
|
||||
|
|
@ -706,28 +706,14 @@ 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
|
||||
are discovered and registered via DCR when left unset. `allow_all_keys`
|
||||
false scopes the server to keys granted it through object_permission."""
|
||||
"""POST /v1/mcp/server. `allow_all_keys` false scopes the server to keys
|
||||
granted it through object_permission."""
|
||||
|
||||
alias: str
|
||||
url: str
|
||||
transport: str = "http"
|
||||
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
|
||||
issuer: str | None = None
|
||||
authorization_url: str | None = None
|
||||
token_url: str | None = None
|
||||
registration_url: str | None = None
|
||||
credentials: McpOauthCredentials | None = Field(default=None, repr=False)
|
||||
server_name: str | None = None
|
||||
description: str | None = None
|
||||
mcp_info: McpInfo | None = None
|
||||
|
|
@ -767,26 +753,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
|
||||
|
|
|
|||
|
|
@ -14,7 +14,6 @@ markers =
|
|||
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_PERF_TESTS is set
|
||||
quiet_stack: measures the proxy itself, so it runs while no other test on this host is hitting the stack; every other test waits for it to finish
|
||||
mcp_oauth_live: real Linear OAuth consent via a captured browser session; deselected unless E2E_MCP_OAUTH_LIVE is set
|
||||
provider_edge_host: routes provider traffic through the pytest host's edge in every fixture mode, so the gateway must reach the pytest host; deselected unless E2E_PROVIDER_EDGE_HOST_REACHABLE is set
|
||||
owned_gateway: boots its own proxy from source against the stack's Postgres, so it needs DATABASE_URL on the pytest host; deselected unless E2E_OWNED_GATEWAY is set
|
||||
otel_v2: needs a proxy running with LITELLM_OTEL_V2=true; deselected unless E2E_OTEL_V2 is set
|
||||
|
|
|
|||
|
|
@ -697,6 +697,54 @@ def test_per_user_authorization_code_with_pkce_binds_the_token_to_the_authorizin
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.timeout(600)
|
||||
def test_per_user_authorization_code_token_is_read_back_from_the_database_by_a_fresh_proxy_without_redis(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
with mcp_peer() as peer, oauth_server() as auth, gateway.scenario() as scenario:
|
||||
alias: Final = "cold" + uuid.uuid4().hex[:8]
|
||||
identity: Final = _register_oauth(
|
||||
scenario,
|
||||
peer,
|
||||
auth,
|
||||
alias,
|
||||
auth_type="oauth2",
|
||||
oauth2_flow="authorization_code",
|
||||
credentials={"client_id": "cold-client", "client_secret": "cold-secret", "scopes": ["tools.call"]},
|
||||
)
|
||||
owner_id: Final = scenario.user()
|
||||
owner: Final = scenario.key(user_id=owner_id, object_permission={"mcp_servers": [identity]})
|
||||
registered: Final = gateway.client.post(
|
||||
f"/{alias}/register", json={"redirect_uris": [CLIENT_REDIRECT], "client_name": "integration"}
|
||||
)
|
||||
assert registered.status_code in (200, 201), registered.text
|
||||
client_id: Final = str(_response_object(registered)["client_id"])
|
||||
pkce: Final = _Pkce(secrets.token_urlsafe(32))
|
||||
code, _ = _authorize_through_gateway(gateway, auth, alias, owner, client_id, pkce)
|
||||
redeemed: Final = _redeem(gateway, alias, owner, client_id, code, pkce)
|
||||
assert redeemed.status_code == 200, redeemed.text
|
||||
token: Final = _issued_token(redeemed.json())
|
||||
stored: Final = read_rows(
|
||||
'SELECT credential_b64 FROM "LiteLLM_MCPUserCredentials" WHERE user_id = %s AND server_id = %s',
|
||||
(owner_id, identity),
|
||||
)
|
||||
assert len(stored) == 1, stored
|
||||
stored_bytes: Final = base64.urlsafe_b64decode(str(stored[0]["credential_b64"]))
|
||||
assert token.encode() not in stored_bytes, "stored token is plaintext"
|
||||
config: Final = tmp_path / "no_redis.yaml"
|
||||
config.write_text(
|
||||
"general_settings:\n"
|
||||
" master_key: os.environ/LITELLM_MASTER_KEY\n"
|
||||
" database_url: os.environ/DATABASE_URL\n"
|
||||
" store_model_in_db: true\n"
|
||||
)
|
||||
peer.drain()
|
||||
with owned_proxy(gateway, tmp_path, {}, config=config, remove_environment=("REDIS_HOST", "REDIS_PORT")) as cold:
|
||||
called: Final = call_tool(cold, owner, identity, f"{alias}-add", ADD)
|
||||
assert called.status_code == 200, called.text
|
||||
assert _authorizations(peer) == (f"Bearer {token}".encode(),)
|
||||
|
||||
|
||||
def test_authorization_request_without_pkce_is_refused_before_reaching_the_authorization_server(
|
||||
gateway: Gateway,
|
||||
) -> None:
|
||||
|
|
|
|||
|
|
@ -52,7 +52,7 @@ CI = [".github/workflows/test-unit.yml"]
|
|||
("mcp-dependencies", ["pyproject.toml"], "run"),
|
||||
("mcp-dependencies", ["uv.lock"], "run"),
|
||||
("mcp-dependencies", ["litellm/experimental_mcp_client/client.py"], "run"),
|
||||
("mcp-dependencies", ["tests/e2e/mcp/oauth_chat_client.py"], "run"),
|
||||
("mcp-dependencies", ["tests/e2e/mcp/datadog_mcp.py"], "run"),
|
||||
("mcp-dependencies", ["litellm-proxy-extras/pyproject.toml"], "run"),
|
||||
("mcp-dependencies", ["scripts/check_mcp_sdk_install.py"], "run"),
|
||||
("mcp-dependencies", [".github/workflows/test-mcp-dependency-resolution.yml"], "run"),
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue