mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
test(mcp): add transport-level M2M regression tests for the preemptive-401 gate
Grafted from PR #33582 (closing as superseded by this PR): drives handle_streamable_http_mcp with real MCPServer objects, parametrized over a stamped client_credentials row and a legacy unstamped M2M-shape row; both must reach the session manager without the per-user token store being consulted
This commit is contained in:
parent
287a89e2ad
commit
f023c819ec
1 changed files with 95 additions and 0 deletions
|
|
@ -664,6 +664,101 @@ async def test_per_user_oauth_missing_stored_token_returns_preemptive_401():
|
|||
assert "Bearer authorization_uri=" in exc_info.value.headers["www-authenticate"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"m2m_fields",
|
||||
[
|
||||
{"oauth2_flow": "client_credentials"},
|
||||
{"client_id": "cid", "client_secret": "csec", "token_url": "https://idp.example.com/token"},
|
||||
],
|
||||
ids=["stamped", "unstamped_m2m_shape"],
|
||||
)
|
||||
async def test_client_credentials_server_is_not_preemptively_challenged(m2m_fields):
|
||||
"""
|
||||
An OAuth2 client_credentials (M2M) server mints its own upstream token;
|
||||
there is no user OAuth flow to bootstrap. The connect-time gate must let
|
||||
the request through to the session manager rather than pushing the client
|
||||
into an interactive OAuth flow it can never complete (the per-user token
|
||||
store is never even consulted for M2M). Covers both a stamped row and a
|
||||
legacy null-flow row with the M2M field shape: the gate must classify the
|
||||
flow through the same request-time chokepoint egress uses, or the two
|
||||
disagree and the unstamped server is challenged for a token egress would
|
||||
never look for.
|
||||
"""
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.server import (
|
||||
handle_streamable_http_mcp,
|
||||
session_manager_stateless,
|
||||
)
|
||||
from litellm.proxy._types import MCPTransport
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
scope = {
|
||||
"type": "http",
|
||||
"method": "POST",
|
||||
"path": "/mcp",
|
||||
"scheme": "http",
|
||||
"query_string": b"",
|
||||
"root_path": "",
|
||||
"server": ("localhost", 8000),
|
||||
"headers": [
|
||||
(b"content-type", b"application/json"),
|
||||
(b"host", b"localhost:8000"),
|
||||
],
|
||||
}
|
||||
receive = AsyncMock(
|
||||
return_value={
|
||||
"type": "http.request",
|
||||
"body": b'{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}',
|
||||
"more_body": False,
|
||||
}
|
||||
)
|
||||
send = AsyncMock()
|
||||
user_auth = MagicMock()
|
||||
user_auth.user_id = "test-user-id"
|
||||
m2m_server = MCPServer(
|
||||
server_id="m2m-server-id",
|
||||
name="m2m_server",
|
||||
server_name="m2m_server",
|
||||
alias="m2m_server",
|
||||
url="https://upstream.example.com/mcp",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.oauth2,
|
||||
**m2m_fields,
|
||||
)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context",
|
||||
new_callable=AsyncMock,
|
||||
return_value=(user_auth, None, ["m2m_server"], None, None, None),
|
||||
),
|
||||
patch("litellm.proxy._experimental.mcp_server.server.set_auth_context"),
|
||||
patch("litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", True),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server._handle_stale_mcp_session",
|
||||
new_callable=AsyncMock,
|
||||
return_value=False,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.has_user_oauth_token",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_has_token,
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name",
|
||||
return_value=m2m_server,
|
||||
),
|
||||
patch.object(session_manager_stateless, "handle_request", new_callable=AsyncMock) as mock_handle_request,
|
||||
patch.object(session_manager_stateless, "_server_instances", {}),
|
||||
):
|
||||
await handle_streamable_http_mcp(scope, receive, send)
|
||||
|
||||
assert mock_handle_request.await_count == 1
|
||||
assert mock_has_token.await_count == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_streamable_http_mcp_delegated_server_surfaces_upstream_challenge():
|
||||
"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue