From 6a71a3b1a0fb7b9697b8c290490452c1f7fa8efd Mon Sep 17 00:00:00 2001 From: yucheng Date: Sun, 27 Sep 2026 01:12:38 +0000 Subject: [PATCH] test(mcp): add integration coverage for the caller sign-in gates Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp/test_mcp_caller_sign_in.py | 218 ++++++++++++++++++ 1 file changed, 218 insertions(+) create mode 100644 tests/integration/mcp/test_mcp_caller_sign_in.py diff --git a/tests/integration/mcp/test_mcp_caller_sign_in.py b/tests/integration/mcp/test_mcp_caller_sign_in.py new file mode 100644 index 00000000000..f80a0c086af --- /dev/null +++ b/tests/integration/mcp/test_mcp_caller_sign_in.py @@ -0,0 +1,218 @@ +import json +import uuid +from collections.abc import Mapping +from pathlib import Path +from typing import Final + +import httpx +import yaml +from cryptography.hazmat.primitives.asymmetric import rsa +from integration._support.client import Gateway +from integration._support.mcp import ( + INITIALIZE, + McpCaller, + forget_mcp, + mcp_peer, + register_mcp, + tool_calls, +) +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, wire_server +from jwt import algorithms as jwt_algorithms + +ADD: Final = {"a": 2, "b": 3} +ACCEPT: Final = {"Accept": "application/json, text/event-stream"} + + +def _rpc(gateway: Gateway, path: str, key: str, headers: dict[str, str]) -> httpx.Response: + return gateway.client.post( + path, + json={"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": INITIALIZE}, + headers={"x-litellm-api-key": key, **ACCEPT, **headers}, + ) + + +def _sign_in_config(guardrail_params: dict[str, object], path: Path) -> Path: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["guardrails"] = [{"guardrail_name": "signin" + uuid.uuid4().hex, "litellm_params": guardrail_params}] + path.write_text(yaml.safe_dump(config)) + return path + + +def test_multi_server_connect_with_a_litellm_key_in_authorization_admits_an_obo_server( + gateway: Gateway, +) -> None: + with mcp_peer() as obo_peer, mcp_peer() as math_peer, gateway.scenario() as scenario: + obo_alias: Final = "obo" + uuid.uuid4().hex[:8] + math_alias: Final = "math" + uuid.uuid4().hex[:8] + obo_id: Final = register_mcp( + scenario, + obo_peer, + obo_alias, + auth_type="oauth2_token_exchange", + token_exchange_endpoint="http://127.0.0.1:9/token", + credentials={"client_id": "obo-client", "client_secret": "obo-secret"}, + ) + math_id: Final = register_mcp(scenario, math_peer, math_alias) + key: Final = scenario.key(object_permission={"mcp_servers": [obo_id, math_id]}) + caller: Final = McpCaller( + gateway, + key, + "root", + headers={ + "x-mcp-servers": f"{obo_alias},{math_alias}", + "Authorization": f"Bearer {key}", + }, + ) + init: Final = caller.initialize() + assert init.ok, init.raw + listed: Final = caller.list_tools() + assert listed.ok, listed.raw + assert any(name.endswith("add") for name in listed.tools), listed.tools + + +def test_alias_first_lookup_wins_over_a_server_whose_name_matches_the_alias(gateway: Gateway, tmp_path: Path) -> None: + with mcp_peer() as math_peer, mcp_peer() as obo_peer: + name: Final = "gh" + uuid.uuid4().hex[:6] + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["mcp_servers"] = { + name: { + "transport": "http", + "url": obo_peer.url, + "auth_type": "oauth2_token_exchange", + "token_exchange_endpoint": "http://127.0.0.1:9/token", + "credentials": {"client_id": "obo-client", "client_secret": "obo-secret"}, + } + } + path: Final = tmp_path / "collision.yaml" + path.write_text(yaml.safe_dump(config)) + with ( + owned_proxy(gateway, tmp_path, {}, config=path) as candidate, + candidate.scenario() as scenario, + ): + second: Final = candidate.request( + "POST", + "/v1/mcp/server", + {"server_name": name + "_public", "alias": name, **math_peer.registration()}, + ) + assert second.status_code == 201, second.text + second_id: Final = second.json()["server_id"] + scenario.cleanups.callback(forget_mcp, candidate, second_id) + key: Final = scenario.key(object_permission={"mcp_servers": [second_id]}) + + init: Final = _rpc(candidate, f"/mcp/{name}", key, {}) + assert init.status_code == 200, init.text + listing: Final = candidate.client.post( + f"/mcp/{name}", + json={"jsonrpc": "2.0", "id": 2, "method": "tools/list", "params": {}}, + headers={"x-litellm-api-key": key, **ACCEPT}, + ) + assert listing.status_code == 200, listing.text + body: Final = json.loads( + next(line[5:].strip() for line in listing.text.splitlines() if line.startswith("data:")) + ) + names: Final = {tool["name"] for tool in body["result"]["tools"]} + assert any(name.endswith("add") for name in names), names + + +def test_jwt_signer_verifies_the_bearer_that_admitted_the_call(gateway: Gateway, tmp_path: Path) -> None: + signer_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + jwk: Final = json.loads(jwt_algorithms.RSAAlgorithm.to_jwk(signer_key.public_key())) + jwk["kid"] = "idp" + holder: Final = [] + + def idp(request: Request) -> Reply: + issuer: Final = holder[0].url + if request.target.endswith("/.well-known/openid-configuration"): + return Reply(body=json.dumps({"issuer": issuer, "jwks_uri": issuer + "/jwks"}).encode()) + return Reply(body=json.dumps({"keys": [jwk]}).encode()) + + def introspect(request: Request) -> Reply: + return Reply(body=b'{"active": true, "sub": "subject-1"}') + + with ( + mcp_peer() as peer, + wire_server(idp) as verify_idp, + wire_server(introspect) as introspect_stub, + gateway.scenario() as scenario, + ): + holder.append(verify_idp) + alias: Final = "sig" + uuid.uuid4().hex[:8] + identity: Final = register_mcp(scenario, peer, alias) + config: Final = _sign_in_config( + { + "guardrail": "mcp_jwt_signer", + "mode": "pre_mcp_call", + "default_on": True, + "access_token_discovery_uri": verify_idp.url + "/.well-known/openid-configuration", + "token_introspection_endpoint": introspect_stub.url, + "required_claims": ["sub"], + }, + tmp_path / "signer.yaml", + ) + with owned_proxy(gateway, tmp_path, {}, config=config) as candidate, candidate.scenario() as owned: + key: Final = owned.key(object_permission={"mcp_servers": [identity]}) + + def call(headers: Mapping[str, str]) -> httpx.Response: + return candidate.client.post( + "/mcp-rest/tools/call", + headers=dict(headers), + json={"name": "add", "arguments": ADD, "server_id": identity}, + ) + + admitted: Final = call({"Authorization": f"Bearer {key}"}) + assert admitted.status_code == 200, admitted.text + assert len(tool_calls(peer.drain())) == 1 + probes: Final = tuple(item for item in introspect_stub.drain() if item.body) + assert any(key.encode() in probe.body for probe in probes), probes + + split: Final = call({"x-litellm-api-key": key}) + assert split.status_code == 403, split.text + assert introspect_stub.drain() == () + + +def test_agent_365_gated_server_challenges_at_connect_and_advertises_entra(gateway: Gateway, tmp_path: Path) -> None: + def nothing(request: Request) -> Reply: + return Reply(status=500) + + with wire_server(nothing) as api: + config: Final = _sign_in_config( + { + "guardrail": "agent_365", + "mode": "pre_mcp_call", + "default_on": True, + "tenant_id": "00000000-0000-0000-0000-000000000000", + "client_id": "22222222-2222-2222-2222-222222222222", + "client_secret": "secret", + "api_base": api.url, + }, + tmp_path / "agent365.yaml", + ) + with ( + owned_proxy(gateway, tmp_path, {}, config=config) as candidate, + mcp_peer() as peer, + candidate.scenario() as scenario, + ): + alias: Final = "a365" + uuid.uuid4().hex[:8] + identity: Final = register_mcp(scenario, peer, alias) + granted: Final = scenario.key(object_permission={"mcp_servers": [identity]}) + denied: Final = scenario.key(object_permission={"mcp_servers": ["no-mcp-servers"]}) + + challenged: Final = _rpc(candidate, f"/mcp/{alias}", granted, {}) + assert challenged.status_code == 401, challenged.text + authenticate: Final = challenged.headers.get("www-authenticate", "") + assert f'resource_metadata="/.well-known/oauth-protected-resource/mcp/{alias}"' in authenticate + assert 'error="invalid_token"' in authenticate + + discovery: Final = candidate.client.get(f"/.well-known/oauth-protected-resource/mcp/{alias}") + assert discovery.status_code == 200, discovery.text + document: Final = discovery.json() + assert document["authorization_servers"] == [ + "https://login.microsoftonline.com/00000000-0000-0000-0000-000000000000/v2.0" + ] + assert document["scopes_supported"] == ["api://22222222-2222-2222-2222-222222222222/access_as_user"] + + refused: Final = _rpc(candidate, f"/mcp/{alias}", denied, {}) + assert refused.status_code == 403, refused.text + assert "www-authenticate" not in refused.headers + assert api.drain() == ()