mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +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
(cherry picked from commit 44fcbc30b0)
This commit is contained in:
yuneng-jiang
Yuneng Jiang
parent
048b587b0c
commit
48bb30b537
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. The semantic tool filter test turns on a proxy-wide setting, so it is marked `proxy_wide_config` and runs on the redeployed managed-files stack after the main run, never next to it
|
||||
- `mcp/` - the MCP server surface over api_key auth against the real Datadog remote MCP server (see "MCP suite: real Datadog only" below). The semantic tool filter test turns on a proxy-wide setting, so it is marked `proxy_wide_config` and runs on the redeployed managed-files stack after the main run, never next to it
|
||||
- `logging/` - logging-integration delivery (datadog and friends)
|
||||
- `security/` - secret handling and log-leak protection
|
||||
- `router/` - routing and reliability behavior (fallbacks, cooldowns)
|
||||
|
|
@ -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
|
||||
```
|
||||
|
||||
|
|
|
|||
|
|
@ -250,55 +250,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,
|
||||
|
|
@ -166,11 +163,6 @@ def pytest_configure(config: pytest.Config) -> None:
|
|||
"redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from "
|
||||
"gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_PERF_TESTS is set",
|
||||
)
|
||||
config.addinivalue_line(
|
||||
"markers",
|
||||
"mcp_oauth_live: real Linear OAuth consent via a captured browser session; deselected unless "
|
||||
"E2E_MCP_OAUTH_LIVE is set",
|
||||
)
|
||||
config.addinivalue_line(
|
||||
"markers",
|
||||
"provider_edge_host: routes provider traffic through the pytest host's edge in every fixture mode, so the "
|
||||
|
|
@ -326,8 +318,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)
|
||||
|
|
@ -371,12 +361,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)
|
||||
|
|
@ -730,28 +730,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
|
||||
|
|
@ -791,26 +777,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
|
||||
|
|
|
|||
|
|
@ -13,7 +13,6 @@ markers =
|
|||
prompt_caching_stack: needs a proxy running with router_settings.optional_pre_call_checks including prompt_caching; deselected unless E2E_PROMPT_CACHING_STACK is set
|
||||
cli_determinism: drives the real claude CLI for several seconds; deselected unless E2E_CLI_DETERMINISM is set
|
||||
redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_PERF_TESTS is set
|
||||
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