mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
fix(mcp): resolve the BYOK credential before keying the listed-tools slot
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
0ab0e49e6a
commit
bef557f3ed
2 changed files with 41 additions and 3 deletions
|
|
@ -4401,9 +4401,16 @@ class MCPServerManager:
|
|||
verbose_logger.info("_get_tools_from_server for %s...", server.name)
|
||||
|
||||
client = None
|
||||
# tools/call resolves the BYOK credential before keying its listed-tools slot; resolve the
|
||||
# same value here or a stored-credential server would list into a slot the call never reads.
|
||||
resolved_mcp_auth_header: Final = (
|
||||
mcp_auth_header
|
||||
if not server.is_byok or isinstance(mcp_auth_header, dict)
|
||||
else await _resolve_byok_mcp_auth_header(server, user_api_key_auth, mcp_auth_header)
|
||||
)
|
||||
listed_caller: Final = ListedToolsCaller(
|
||||
user_api_key_auth=user_api_key_auth,
|
||||
mcp_auth_header=mcp_auth_header,
|
||||
mcp_auth_header=resolved_mcp_auth_header,
|
||||
raw_headers=raw_headers,
|
||||
oauth2_headers=oauth2_headers,
|
||||
)
|
||||
|
|
@ -4449,7 +4456,7 @@ class MCPServerManager:
|
|||
if (
|
||||
get_mcp_jwt_signer() is not None
|
||||
and not has_static_authorization
|
||||
and not mcp_auth_header
|
||||
and not resolved_mcp_auth_header
|
||||
and not has_extra_authorization
|
||||
):
|
||||
extra_headers = await inject_mcp_jwt_headers_for_upstream(
|
||||
|
|
@ -4472,7 +4479,7 @@ class MCPServerManager:
|
|||
|
||||
client = await self._create_mcp_client(
|
||||
server=server,
|
||||
mcp_auth_header=mcp_auth_header,
|
||||
mcp_auth_header=resolved_mcp_auth_header,
|
||||
extra_headers=extra_headers,
|
||||
stdio_env=stdio_env,
|
||||
subject_token=subject_token,
|
||||
|
|
|
|||
|
|
@ -7037,6 +7037,37 @@ class TestMCPServerManager:
|
|||
listed = manager.get_listed_tool(server, "turn", other)
|
||||
assert listed is not None and listed.description == "everyone"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_byok_stored_credential_lists_into_the_slot_tools_call_reads(self):
|
||||
from litellm.proxy._experimental.mcp_server.byok_credential_cache import (
|
||||
byok_credential_cache_key,
|
||||
cache_byok_credential,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.operations import byok_credential_cache
|
||||
|
||||
manager = MCPServerManager()
|
||||
server = MCPServer(
|
||||
server_id="byok-catalog",
|
||||
name="byok_catalog",
|
||||
transport=MCPTransport.http,
|
||||
url="http://byok-catalog",
|
||||
is_byok=True,
|
||||
)
|
||||
user = UserAPIKeyAuth(api_key="sk-litellm", user_id="byok-user")
|
||||
manager._create_mcp_client = AsyncMock(return_value=AsyncMock())
|
||||
manager._fetch_tools_with_timeout = AsyncMock(
|
||||
return_value=[MCPTool(name="turn", description="stored cred catalog", inputSchema={})]
|
||||
)
|
||||
cache_byok_credential("byok-user", "byok-catalog", "stored-secret")
|
||||
try:
|
||||
await manager._get_tools_from_server(server=server, user_api_key_auth=user)
|
||||
finally:
|
||||
byok_credential_cache.delete_cache(byok_credential_cache_key("byok-user", "byok-catalog"))
|
||||
|
||||
call_side = ListedToolsCaller(user_api_key_auth=user, mcp_auth_header="stored-secret")
|
||||
listed = manager.get_listed_tool(server, "turn", call_side)
|
||||
assert listed is not None and listed.description == "stored cred catalog"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("signer", "static_headers", "shared"),
|
||||
[
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue