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:
Tin 2026-07-16 14:49:37 -07:00 committed by Tin Chi Lo
parent 287a89e2ad
commit f023c819ec

View file

@ -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():
"""