mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-10 03:28:53 +00:00
feat(byok): inject stored BYOK credential as mcp_auth_header on tool execution
This commit is contained in:
parent
77c54a9dfc
commit
bc9466588f
1 changed files with 28 additions and 0 deletions
|
|
@ -1504,6 +1504,27 @@ if MCP_AVAILABLE:
|
||||||
)
|
)
|
||||||
return name
|
return name
|
||||||
|
|
||||||
|
async def _get_byok_credential(
|
||||||
|
mcp_server: MCPServer,
|
||||||
|
user_api_key_auth: Optional[UserAPIKeyAuth],
|
||||||
|
) -> Optional[str]:
|
||||||
|
"""Retrieve the stored BYOK credential for a user+server pair."""
|
||||||
|
if not mcp_server.is_byok:
|
||||||
|
return None
|
||||||
|
user_id = (user_api_key_auth.user_id if user_api_key_auth else None) or ""
|
||||||
|
if not user_id:
|
||||||
|
return None
|
||||||
|
from litellm.proxy._experimental.mcp_server.db import get_user_credential
|
||||||
|
from litellm.proxy.proxy_server import prisma_client
|
||||||
|
|
||||||
|
if prisma_client is None:
|
||||||
|
return None
|
||||||
|
return await get_user_credential(
|
||||||
|
prisma_client=prisma_client,
|
||||||
|
user_id=user_id,
|
||||||
|
server_id=mcp_server.server_id,
|
||||||
|
)
|
||||||
|
|
||||||
async def _check_byok_credential(
|
async def _check_byok_credential(
|
||||||
mcp_server: MCPServer,
|
mcp_server: MCPServer,
|
||||||
user_api_key_auth: Optional[UserAPIKeyAuth],
|
user_api_key_auth: Optional[UserAPIKeyAuth],
|
||||||
|
|
@ -1693,6 +1714,13 @@ if MCP_AVAILABLE:
|
||||||
# that an MCP client can trigger the authorization flow.
|
# that an MCP client can trigger the authorization flow.
|
||||||
await _check_byok_credential(mcp_server, user_api_key_auth)
|
await _check_byok_credential(mcp_server, user_api_key_auth)
|
||||||
|
|
||||||
|
# For BYOK servers, inject the user's stored credential as the
|
||||||
|
# auth header if no explicit override was provided by the caller.
|
||||||
|
if mcp_server.is_byok and not mcp_auth_header:
|
||||||
|
mcp_auth_header = await _get_byok_credential(
|
||||||
|
mcp_server, user_api_key_auth
|
||||||
|
)
|
||||||
|
|
||||||
response = await _handle_managed_mcp_tool(
|
response = await _handle_managed_mcp_tool(
|
||||||
server_name=server_name,
|
server_name=server_name,
|
||||||
name=original_tool_name, # Pass the full name (potentially prefixed)
|
name=original_tool_name, # Pass the full name (potentially prefixed)
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue