diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index de2ec13b4a3..314eed95985 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -2093,3 +2093,83 @@ async def test_get_tools_from_mcp_servers_logs_list_tools_to_spendlogs_when_enab assert spend_meta["tool_count_total"] == 1 assert spend_meta["allowed_server_count"] == 1 assert spend_meta["per_server_tool_counts"]["server_a"] == 1 + + +@pytest.mark.asyncio +async def test_get_tools_from_mcp_servers_injects_stored_oauth2_token(): + """ + When _get_tools_from_mcp_servers is called for an OAuth2 MCP server and no + oauth2_headers are provided in the request (e.g. a /responses API call from a + chat UI), the per-user stored token must be fetched from the DB and passed as + extra_headers to _get_tools_from_server. + + The implementation pre-fetches all user credentials in a single bulk query + (_prefetch_oauth_creds_for_user) to avoid N+1 queries in the gather loop. + + This covers the bug where OAuth2 MCP tools were always empty in the /responses + API because the stored credential was never injected. + """ + try: + from litellm.proxy._experimental.mcp_server.server import ( + _get_tools_from_mcp_servers, + ) + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.mcp import MCPAuth + except ImportError: + pytest.skip("MCP server not available") + + STORED_TOKEN = "atlassian-oauth-access-token-xyz" + SERVER_ID = "srv-oauth2-id" + USER_ID = "user-123" + + user_auth = UserAPIKeyAuth(api_key="test-key", user_id=USER_ID) + + oauth2_server = MagicMock(name="atlassian_server") + oauth2_server.name = "atlassian_test" + oauth2_server.alias = "atlassian_test" + oauth2_server.server_name = "atlassian_test" + oauth2_server.server_id = SERVER_ID + oauth2_server.auth_type = MCPAuth.oauth2 + oauth2_server.extra_headers = None + + # Simulate the DB returning a valid credential for this user+server + prefetched_creds = {SERVER_ID: {"access_token": STORED_TOKEN, "server_id": SERVER_ID}} + + tool_1 = MagicMock() + tool_1.name = "atlassian_test-search" + + with patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(return_value=[oauth2_server]), + ), patch( + # Patch the bulk prefetch so no real DB connection is needed + "litellm.proxy._experimental.mcp_server.server._prefetch_oauth_creds_for_user", + new=AsyncMock(return_value=prefetched_creds), + ) as mock_prefetch, patch( + "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + ) as mock_manager, patch( + "litellm.proxy._experimental.mcp_server.server.filter_tools_by_allowed_tools", + side_effect=lambda tools, _server: tools, + ), patch( + "litellm.proxy._experimental.mcp_server.server.filter_tools_by_key_team_permissions", + new=AsyncMock(side_effect=lambda tools, **_: tools), + ): + mock_manager._get_tools_from_server = AsyncMock(return_value=[tool_1]) + + tools = await _get_tools_from_mcp_servers( + user_api_key_auth=user_auth, + mcp_auth_header=None, + mcp_servers=["atlassian_test"], + mcp_server_auth_headers=None, + oauth2_headers=None, # No token from request — must fall back to DB + ) + + # Bulk credential prefetch was called once (not once per server) + mock_prefetch.assert_awaited_once_with(user_auth) + + # The stored token was forwarded to the MCP transport layer as extra_headers + mock_manager._get_tools_from_server.assert_awaited_once() + call_kwargs = mock_manager._get_tools_from_server.await_args.kwargs + assert call_kwargs["extra_headers"] == {"Authorization": f"Bearer {STORED_TOKEN}"} + + assert tools == [tool_1]