mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-06 07:48:06 +00:00
fix(mcp): defer delegated OAuth startup discovery
This commit is contained in:
parent
f6f49d16cc
commit
d4ac495910
2 changed files with 81 additions and 2 deletions
|
|
@ -7249,8 +7249,8 @@ class MCPServerManager:
|
|||
for server in self.get_registry().values():
|
||||
if self.oauth_discovery_slot(server.server_id) is not None:
|
||||
continue
|
||||
if server.needs_user_oauth_token:
|
||||
# Skip OAuth2 servers that rely on user-provided tokens
|
||||
if server.needs_user_oauth_token or server.is_oauth_delegate:
|
||||
# Startup has no caller token for these OAuth servers.
|
||||
continue
|
||||
try:
|
||||
tools = await self.get_tools_from_server(server)
|
||||
|
|
|
|||
|
|
@ -21900,3 +21900,82 @@ def test_id_jag_listed_tool_metadata_is_isolated_between_identity_assertions() -
|
|||
manager.record_listed_tools(server, (first,), callers[0])
|
||||
manager.record_listed_tools(server, (second,), callers[1])
|
||||
assert tuple(manager.get_listed_tool(server, "list_notes", caller) for caller in callers) == (first, second, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("rejected_status", [401, 403])
|
||||
async def test_delegated_oauth_startup_defers_discovery_until_caller_token(
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
rejected_status: int,
|
||||
) -> None:
|
||||
upstream: Final = _DiscoveryUpstream()
|
||||
upstream.tools = (LIST_NOTES.model_dump(by_alias=True),)
|
||||
|
||||
requests: Final[list[httpx2.Request]] = []
|
||||
|
||||
async def respond(request: httpx2.Request) -> httpx2.Response:
|
||||
requests.append(request)
|
||||
if request.headers.get("authorization") != "Bearer caller-token":
|
||||
return httpx2.Response(rejected_status, headers={"WWW-Authenticate": "Bearer"})
|
||||
return await upstream.respond(request)
|
||||
|
||||
manager: Final = MCPServerManager(upstream_transport=httpx2.MockTransport(respond))
|
||||
server: Final = MCPServer(
|
||||
server_id="deferred-oauth",
|
||||
name="deferred_oauth",
|
||||
url="https://upstream.example/mcp",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.oauth_delegate,
|
||||
)
|
||||
manager.registry[server.server_id] = server
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM"):
|
||||
await manager._initialize_tool_name_to_mcp_server_name_mapping()
|
||||
assert not requests
|
||||
assert manager.tool_name_to_mcp_server_name_mapping == {}
|
||||
assert not caplog.records
|
||||
tools: Final = await manager.get_tools_from_server(server, mcp_auth_header="Bearer caller-token")
|
||||
assert tuple(tool.name for tool in tools) == ("deferred-oauth-list_notes",)
|
||||
assert upstream.requests
|
||||
assert all(authorization == "Bearer caller-token" for _, authorization in upstream.requests)
|
||||
with pytest.raises(MCPUpstreamAuthError) as error:
|
||||
await manager.get_tools_from_server(server, mcp_auth_header="Bearer expired-token")
|
||||
assert error.value.status_code == rejected_status
|
||||
assert error.value.www_authenticate == "Bearer"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("status_code", [200, 401, 500])
|
||||
async def test_static_auth_startup_still_discovers_tools_or_reports_failure(
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
status_code: int,
|
||||
) -> None:
|
||||
upstream: Final = _DiscoveryUpstream()
|
||||
upstream.tools = (LIST_NOTES.model_dump(by_alias=True),)
|
||||
requests: Final[list[httpx2.Request]] = []
|
||||
|
||||
async def respond(request: httpx2.Request) -> httpx2.Response:
|
||||
requests.append(request)
|
||||
if status_code != 200:
|
||||
return httpx2.Response(status_code, headers={"WWW-Authenticate": "Bearer"})
|
||||
return await upstream.respond(request)
|
||||
|
||||
manager: Final = MCPServerManager(upstream_transport=httpx2.MockTransport(respond))
|
||||
server: Final = MCPServer(
|
||||
server_id="static-auth",
|
||||
name="static_auth",
|
||||
url="https://upstream.example/mcp",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.bearer_token,
|
||||
authentication_token="static-token",
|
||||
)
|
||||
manager.registry[server.server_id] = server
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM"):
|
||||
await manager._initialize_tool_name_to_mcp_server_name_mapping()
|
||||
assert requests
|
||||
assert all(request.headers.get("authorization") == "Bearer static-token" for request in requests)
|
||||
if status_code == 200:
|
||||
assert manager.get_mcp_server_from_tool_name("list_notes") is server
|
||||
assert not caplog.records
|
||||
else:
|
||||
assert manager.tool_name_to_mcp_server_name_mapping == {}
|
||||
assert any(record.levelno == logging.WARNING for record in caplog.records)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue