test(mcp): add unit test for OAuth2 token injection in _get_tools_from_mcp_servers

Verifies that when _get_tools_from_mcp_servers is called for an OAuth2 MCP
server without oauth2_headers in the request, the implementation:
- calls _prefetch_oauth_creds_for_user once (not per-server) to avoid N+1 queries
- passes the stored token as extra_headers={"Authorization": "Bearer ..."} to
  _get_tools_from_server so the upstream OAuth2 MCP server authenticates correctly
This commit is contained in:
Ishaan Jaffer 2026-03-11 18:47:43 -07:00
parent 29be249965
commit f5eac90665

View file

@ -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]