mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
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:
parent
29be249965
commit
f5eac90665
1 changed files with 80 additions and 0 deletions
|
|
@ -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]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue