scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

fix(mcp): don't block MCP initialize on a per-user OAuth token refresh

Browse source 
The pre-flight 401 check on the streamable-HTTP initialize path called
_get_user_oauth_extra_headers_from_db purely to decide whether to emit a
PKCE-triggering 401. That function is not a pure read: when it finds an
expired-but-refreshable stored token it performs a synchronous token
refresh (an httpx POST to the server's token_url) that inherits litellm's
600s default read timeout. Because this only runs for servers named in the
path, a single-server endpoint like /slack_obo/mcp could hang the entire
initialize call until the MCP client aborted, while the aggregator /mcp
endpoint (which never enters this loop) stayed responsive.

Replace the refresh-triggering lookup with _user_has_resolvable_oauth_token,
which answers the only question the pre-flight check needs ("does the caller
have a usable token, or should we 401 for PKCE?") without any network call.
A valid token, or an expired one that still holds a refresh_token, counts as
present; the lazy list_tools / call_tool path performs the actual refresh,
matching how the aggregator already behaves. As defense-in-depth, bound
refresh_user_oauth_token's POST with MCP_METADATA_TIMEOUT so a stalled token
endpoint can never block for ten minutes on any path.
This commit is contained in:
Tin Chi Lo 2026-06-11 12:21:11 -07:00
parent a992ed18df
commit 58857988a2
3 changed files with 244 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
litellm/proxy/_experimental/mcp_server/db.py
View file

@ -7,7 +7,10 @@ from typing import Any, Dict, Iterable, List, Optional, Set, Union, cast
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
from litellm.constants import MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS
from litellm.constants import (
MCP_METADATA_TIMEOUT,
MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS,
)
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.proxy._types import (
LiteLLM_MCPServerTable,
@ -1121,6 +1124,7 @@ async def refresh_user_oauth_token(
token_url,
headers={"Accept": "application/json"},
data=token_data,
timeout=MCP_METADATA_TIMEOUT,
)
response.raise_for_status()
body: Dict[str, Any] = response.json()

63
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -1414,6 +1414,64 @@ if MCP_AVAILABLE:
)
return None
async def _user_has_resolvable_oauth_token(
server: MCPServer,
user_api_key_auth: Optional[UserAPIKeyAuth],
) -> bool:
"""Whether the caller already has a usable per-user OAuth2 token for this
server, decided WITHOUT performing a network refresh.
Used by the pre-flight 401 check on the initialize path. A stored token
that is still valid — or expired but holding a refresh_token — means we
must not fire the PKCE-triggering 401; the lazy list_tools / call_tool
path refreshes it. Triggering the refresh here would block initialize on
a token-endpoint round-trip and time out the MCP client. Fails open so a
transient lookup error neither 401s a legitimate caller nor stalls the
connection.
"""
if server.auth_type != MCPAuth.oauth2:
return False
if user_api_key_auth is None:
return False
user_id = getattr(user_api_key_auth, "user_id", None)
server_id = getattr(server, "server_id", None)
if not user_id or not server_id:
return False
try:
from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415
get_user_oauth_credential,
is_oauth_credential_expired,
)
from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( # noqa: PLC0415
mcp_per_user_token_cache,
)
if await mcp_per_user_token_cache.get(user_id, server_id) is not None:
return True
from litellm.proxy.utils import ( # noqa: PLC0415
get_prisma_client_or_throw,
)
prisma_client = get_prisma_client_or_throw(
"Database not connected. Connect a database to use OAuth2 MCP tools."
)
cred = await get_user_oauth_credential(prisma_client, user_id, server_id)
if not cred or not cred.get("access_token"):
return False
if not is_oauth_credential_expired(cred):
return True
return bool(cred.get("refresh_token"))
except Exception as e:
verbose_logger.warning(
"_user_has_resolvable_oauth_token: lookup failed for "
"user=%s server=%s: %s",
user_id,
server_id,
e,
)
return True
async def _prefetch_oauth_creds_for_user(
user_api_key_auth: Optional[UserAPIKeyAuth],
) -> Dict[str, Dict[str, Any]]:
@ -3420,11 +3478,10 @@ if MCP_AVAILABLE:
# If no stored token exists, fail fast with 401 so clients can
# kick off PKCE/interactive OAuth flow immediately.
if server.needs_user_oauth_token:
stored_oauth_headers = await _get_user_oauth_extra_headers_from_db(
if await _user_has_resolvable_oauth_token(
server=server,
user_api_key_auth=user_api_key_auth,
)
if stored_oauth_headers:
):
continue
request = StarletteRequest(scope)

187
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py
View file

@ -648,10 +648,10 @@ async def test_per_user_oauth_missing_stored_token_returns_preemptive_401():
return_value=False,
),
patch(
"litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db",
"litellm.proxy._experimental.mcp_server.server._user_has_resolvable_oauth_token",
new_callable=AsyncMock,
return_value=None,
) as mock_get_stored_token,
return_value=False,
) as mock_has_token,
patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name",
return_value=oauth_server,
@ -666,7 +666,7 @@ async def test_per_user_oauth_missing_stored_token_returns_preemptive_401():
await handle_streamable_http_mcp(scope, receive, send)
# Verify a 401 was raised
assert mock_get_stored_token.await_count == 1
assert mock_has_token.await_count == 1
assert mock_handle_request.await_count == 0
assert exc_info.value.status_code == 401
assert "www-authenticate" in exc_info.value.headers
@ -733,10 +733,10 @@ async def test_per_user_oauth_with_stored_token_skips_preemptive_401():
return_value=False,
),
patch(
"litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db",
"litellm.proxy._experimental.mcp_server.server._user_has_resolvable_oauth_token",
new_callable=AsyncMock,
return_value={"Authorization": "Bearer cached-token"},
) as mock_get_stored_token,
return_value=True,
) as mock_has_token,
patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name",
return_value=oauth_server,
@ -754,10 +754,181 @@ async def test_per_user_oauth_with_stored_token_skips_preemptive_401():
):
await handle_streamable_http_mcp(scope, receive, send)
assert mock_get_stored_token.await_count == 1
assert mock_has_token.await_count == 1
assert mock_handle_request.await_count == 1
@pytest.mark.asyncio
async def test_preemptive_check_does_not_refresh_expired_token_on_initialize():
"""
Regression: the pre-flight 401 check on a per-user OAuth server must NOT
trigger a network token refresh. An expired-but-refreshable stored token
counts as present (the lazy list_tools / call_tool path refreshes it), so
the request proceeds to the session manager instead of blocking initialize
on a token-endpoint round-trip (which timed out individual /server/mcp
endpoints while the aggregator stayed unaffected).
"""
from datetime import datetime, timedelta, timezone
try:
from litellm.proxy._experimental.mcp_server.server import (
handle_streamable_http_mcp,
session_manager_stateless,
)
except ImportError:
pytest.skip("MCP server not available")
scope = {
"type": "http",
"method": "POST",
"path": "/slack_obo/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"
oauth_server = MagicMock()
oauth_server.auth_type = MCPAuth.oauth2
oauth_server.needs_user_oauth_token = True
oauth_server.server_id = "slack-obo-server-id"
expired_cred = {
"access_token": "expired-access-token",
"refresh_token": "stored-refresh-token",
"expires_at": (datetime.now(timezone.utc) - timedelta(hours=1)).isoformat(),
}
with (
patch(
"litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context",
new_callable=AsyncMock,
return_value=(user_auth, None, ["slack_obo"], 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.get_mcp_server_by_name",
return_value=oauth_server,
),
patch(
"litellm.proxy._experimental.mcp_server.oauth2_token_cache.mcp_per_user_token_cache.get",
new_callable=AsyncMock,
return_value=None,
),
patch(
"litellm.proxy.utils.get_prisma_client_or_throw",
return_value=MagicMock(),
),
patch(
"litellm.proxy._experimental.mcp_server.db.get_user_oauth_credential",
new_callable=AsyncMock,
return_value=expired_cred,
),
patch(
"litellm.proxy._experimental.mcp_server.db.refresh_user_oauth_token",
new_callable=AsyncMock,
) as mock_refresh,
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_refresh.await_count == 0
assert mock_handle_request.await_count == 1
@pytest.mark.asyncio
async def test_user_has_resolvable_oauth_token_classification():
"""
_user_has_resolvable_oauth_token resolves token presence WITHOUT a network
refresh: valid or expired-with-refresh-token counts as present; missing or
expired-without-refresh-token does not.
"""
from datetime import datetime, timedelta, timezone
try:
from litellm.proxy._experimental.mcp_server.server import (
_user_has_resolvable_oauth_token,
)
except ImportError:
pytest.skip("MCP server not available")
user_auth = MagicMock()
user_auth.user_id = "test-user-id"
server = MagicMock()
server.auth_type = MCPAuth.oauth2
server.server_id = "srv-1"
past = (datetime.now(timezone.utc) - timedelta(hours=1)).isoformat()
future = (datetime.now(timezone.utc) + timedelta(hours=1)).isoformat()
cases = {
"valid": ({"access_token": "a", "expires_at": future}, True),
"expired_with_refresh": (
{"access_token": "a", "expires_at": past, "refresh_token": "r"},
True,
),
"expired_no_refresh": ({"access_token": "a", "expires_at": past}, False),
"missing": (None, False),
}
for label, (cred, expected) in cases.items():
with (
patch(
"litellm.proxy._experimental.mcp_server.oauth2_token_cache.mcp_per_user_token_cache.get",
new_callable=AsyncMock,
return_value=None,
),
patch(
"litellm.proxy.utils.get_prisma_client_or_throw",
return_value=MagicMock(),
),
patch(
"litellm.proxy._experimental.mcp_server.db.get_user_oauth_credential",
new_callable=AsyncMock,
return_value=cred,
),
patch(
"litellm.proxy._experimental.mcp_server.db.refresh_user_oauth_token",
new_callable=AsyncMock,
) as mock_refresh,
):
result = await _user_has_resolvable_oauth_token(server, user_auth)
assert result is expected, f"case {label}: expected {expected}, got {result}"
assert mock_refresh.await_count == 0, f"case {label} triggered a refresh"
@pytest.mark.asyncio
async def test_handle_streamable_http_mcp_emits_401_for_delegated_server_without_token():
"""