mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
feat(mcp): cut the call_tool egress over to v2 for authorization_code servers
_resolve_oauth2_headers_for_tool_call steps aside (builds no header) when to_server_spec maps the server, so the v2 resolver drives the token-present case instead of being shadowed by a token v1 places in extra_headers. Non-migrated oauth2 (delegate, client_credentials) and BYOK still build their header on v1. With this, v2 owns the authorization_code egress end to end: inject the refreshed per-user token when present, raise the per-server fail-closed 401 when absent.
This commit is contained in:
parent
40f4a00282
commit
f0cd3b884a
2 changed files with 46 additions and 9 deletions
|
|
@ -3694,6 +3694,12 @@ class MCPServerManager:
|
|||
):
|
||||
return oauth2_headers
|
||||
|
||||
if to_server_spec(mcp_server) is not None:
|
||||
# Migrated to v2: the resolver owns this server's per-user token (inject or fail-closed
|
||||
# 401). Building it into extra_headers here would let the v2 graft defer to it and
|
||||
# shadow the resolver, double-resolving and hiding the per-server challenge.
|
||||
return oauth2_headers
|
||||
|
||||
user_id = getattr(user_api_key_auth, "user_id", None)
|
||||
if not user_id:
|
||||
return oauth2_headers
|
||||
|
|
|
|||
|
|
@ -2221,7 +2221,7 @@ class TestMCPServerManager:
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_oauth2_headers_looks_up_stored_token(self):
|
||||
"""Falls back to stored per-user OAuth headers when no token is supplied."""
|
||||
"""Falls back to stored per-user OAuth headers for a non-migrated (delegate) oauth2 server."""
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
manager = MCPServerManager()
|
||||
|
|
@ -2230,6 +2230,7 @@ class TestMCPServerManager:
|
|||
name="oauth-srv",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.oauth2,
|
||||
delegate_auth_to_upstream=True, # stays on v1, so v1 still builds the header
|
||||
)
|
||||
user_auth = UserAPIKeyAuth(api_key="sk-test", user_id="alice")
|
||||
stored = {"Authorization": "Bearer stored-user-token"}
|
||||
|
|
@ -2245,6 +2246,33 @@ class TestMCPServerManager:
|
|||
assert result == stored
|
||||
mock_lookup.assert_awaited_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_oauth2_headers_steps_aside_for_migrated_server(self):
|
||||
"""A migrated authorization_code server is owned by the v2 resolver, so v1 must not also
|
||||
build the token into extra_headers (which the v2 graft would defer to and shadow).
|
||||
"""
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
manager = MCPServerManager()
|
||||
server = MCPServer(
|
||||
server_id="oauth-srv",
|
||||
name="oauth-srv",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.oauth2, # per-user, not delegate -> migrated to v2
|
||||
)
|
||||
user_auth = UserAPIKeyAuth(api_key="sk-test", user_id="alice")
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db",
|
||||
new=AsyncMock(return_value={"Authorization": "Bearer should-not-be-used"}),
|
||||
) as mock_lookup:
|
||||
result = await manager._resolve_oauth2_headers_for_tool_call(
|
||||
server, oauth2_headers=None, user_api_key_auth=user_auth
|
||||
)
|
||||
|
||||
assert result is None # stepped aside; the v2 resolver handles the token
|
||||
mock_lookup.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_oauth2_headers_swallows_lookup_exception(self):
|
||||
"""Returns supplied headers (None) when the stored-token lookup raises."""
|
||||
|
|
@ -2970,14 +2998,17 @@ class TestMCPServerManager:
|
|||
object_permission_id="perm_no_mcp",
|
||||
)
|
||||
|
||||
with patch.object(
|
||||
manager, "get_allow_all_keys_server_ids", return_value=["global-server"]
|
||||
), patch.object(
|
||||
MCPRequestHandler,
|
||||
"get_allowed_mcp_servers",
|
||||
new_callable=AsyncMock,
|
||||
return_value=["leaked-server"],
|
||||
) as mock_inner:
|
||||
with (
|
||||
patch.object(
|
||||
manager, "get_allow_all_keys_server_ids", return_value=["global-server"]
|
||||
),
|
||||
patch.object(
|
||||
MCPRequestHandler,
|
||||
"get_allowed_mcp_servers",
|
||||
new_callable=AsyncMock,
|
||||
return_value=["leaked-server"],
|
||||
) as mock_inner,
|
||||
):
|
||||
result = await manager.get_allowed_mcp_servers(user_api_key_auth)
|
||||
|
||||
assert result == []
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue