diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index cefe2ef763..5f32065593 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -2671,13 +2671,19 @@ if MCP_AVAILABLE: server_name, client_ip=_client_ip ) if server and server.auth_type == MCPAuth.oauth2 and not oauth2_headers: - # For servers that store per-user tokens server-side, skip the - # pre-emptive 401 — the call_tool / list_tools dispatch will look - # up the stored token from Redis / DB and only fail at the MCP - # protocol level if none is found, giving the client a proper - # tool-execution error rather than an HTTP 401. + # For per-user OAuth servers, only skip the pre-emptive 401 when + # a stored token actually exists for this user+server pair. + # 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: - continue + stored_oauth_headers = ( + await _get_user_oauth_extra_headers_from_db( + server=server, + user_api_key_auth=user_api_key_auth, + ) + ) + if stored_oauth_headers: + continue request = StarletteRequest(scope) base_url = get_request_base_url(request) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py index 5eb8c1e51a..2fbcd5d0d7 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py @@ -10,6 +10,9 @@ they may send a stale `mcp-session-id` header. This test verifies that: import pytest from unittest.mock import AsyncMock, MagicMock, patch +from fastapi import HTTPException +from litellm.types.mcp import MCPAuth + class TestHandleStaleMcpSession: """Unit tests for the _handle_stale_mcp_session helper.""" @@ -438,3 +441,129 @@ async def test_no_mcp_session_id_header_works_normally(): header_names = [k for k, v in captured_scope.get("headers", [])] assert b"mcp-session-id" not in header_names assert b"content-type" in header_names + + +@pytest.mark.asyncio +async def test_per_user_oauth_missing_stored_token_returns_preemptive_401(): + """ + Per-user OAuth server with no stored token should fail fast with 401 + + WWW-Authenticate so PKCE can start. + """ + try: + from litellm.proxy._experimental.mcp_server.server import ( + handle_streamable_http_mcp, + session_manager, + ) + except ImportError: + pytest.skip("MCP server not available") + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp", + "headers": [ + (b"content-type", b"application/json"), + ], + } + receive = AsyncMock() + 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 + + with patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=(user_auth, None, ["repro_oauth_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._get_user_oauth_extra_headers_from_db", + new_callable=AsyncMock, + return_value=None, + ) as mock_get_stored_token, patch( + "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name", + return_value=oauth_server, + ), patch.object( + session_manager, + "handle_request", + new_callable=AsyncMock, + ) as mock_handle_request: + with pytest.raises(HTTPException) as exc_info: + await handle_streamable_http_mcp(scope, receive, send) + + exc = exc_info.value + assert exc.status_code == 401 + assert "www-authenticate" in exc.headers + assert mock_get_stored_token.await_count == 1 + assert mock_handle_request.await_count == 0 + + +@pytest.mark.asyncio +async def test_per_user_oauth_with_stored_token_skips_preemptive_401(): + """ + Per-user OAuth server with an existing stored token should skip pre-emptive + 401 and continue to session manager request handling. + """ + try: + from litellm.proxy._experimental.mcp_server.server import ( + handle_streamable_http_mcp, + session_manager, + ) + except ImportError: + pytest.skip("MCP server not available") + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp", + "headers": [ + (b"content-type", b"application/json"), + ], + } + receive = AsyncMock() + 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 + + with patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=(user_auth, None, ["repro_oauth_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._get_user_oauth_extra_headers_from_db", + new_callable=AsyncMock, + return_value={"Authorization": "Bearer cached-token"}, + ) as mock_get_stored_token, patch( + "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name", + return_value=oauth_server, + ), patch.object( + session_manager, + "handle_request", + new_callable=AsyncMock, + ) as mock_handle_request: + await handle_streamable_http_mcp(scope, receive, send) + + assert mock_get_stored_token.await_count == 1 + assert mock_handle_request.await_count == 1