mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
fix(mcp): inject stored OAuth2 token when fetching tools via /responses API
When a user has connected an OAuth2 MCP server (e.g. Atlassian) and then uses the /responses endpoint with that server, tool listing was failing because the stored per-user OAuth token was never injected. Two fixes: 1. server.py: add _get_user_oauth_extra_headers_from_db() helper; call it in _get_tools_from_mcp_servers when oauth2_headers is None for an OAuth2 server, falling back to the user's stored token in LiteLLM_MCPUserCredentials 2. litellm_proxy_mcp_handler.py: also intercept MCP tools whose server_url matches */mcp/<server_name> (e.g. http://localhost:4000/mcp/atlassian_test) by rewriting them to litellm_proxy/mcp/<server_name> so they go through the internal handler (and get the OAuth token injected) instead of being forwarded to OpenAI raw where localhost is unreachable
This commit is contained in:
parent
c5e34313b6
commit
24279c6aaa
2 changed files with 58 additions and 1 deletions
|
|
@ -871,6 +871,35 @@ if MCP_AVAILABLE:
|
|||
|
||||
return allowed_mcp_servers
|
||||
|
||||
async def _get_user_oauth_extra_headers_from_db(
|
||||
server: MCPServer,
|
||||
user_api_key_auth: Optional[UserAPIKeyAuth],
|
||||
) -> Optional[Dict[str, str]]:
|
||||
"""Look up stored OAuth2 token for (user, server) from DB and return as extra_headers dict."""
|
||||
if server.auth_type != MCPAuth.oauth2:
|
||||
return None
|
||||
if user_api_key_auth is None:
|
||||
return None
|
||||
user_id = getattr(user_api_key_auth, "user_id", None)
|
||||
server_id = getattr(server, "server_id", None)
|
||||
if not user_id or not server_id:
|
||||
return None
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415
|
||||
get_user_oauth_credential,
|
||||
)
|
||||
from litellm.proxy.utils import get_prisma_client_or_throw # noqa: PLC0415
|
||||
|
||||
prisma_client = get_prisma_client_or_throw(
|
||||
"Database not connected. Connect a database to use OAuth2 MCP tools."
|
||||
)
|
||||
cred = await get_user_oauth_credential(prisma_client, user_id, server_id)
|
||||
if cred and cred.get("access_token"):
|
||||
return {"Authorization": f"Bearer {cred['access_token']}"}
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
def _prepare_mcp_server_headers(
|
||||
server: MCPServer,
|
||||
mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]],
|
||||
|
|
@ -1030,6 +1059,12 @@ if MCP_AVAILABLE:
|
|||
raw_headers=raw_headers,
|
||||
)
|
||||
|
||||
# If no OAuth2 token came from request headers, fall back to DB lookup
|
||||
if extra_headers is None and server.auth_type == MCPAuth.oauth2:
|
||||
extra_headers = await _get_user_oauth_extra_headers_from_db(
|
||||
server, user_api_key_auth
|
||||
)
|
||||
|
||||
try:
|
||||
tools = await global_mcp_server_manager._get_tools_from_server(
|
||||
server=server,
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import re
|
||||
import traceback
|
||||
from datetime import datetime
|
||||
from typing import (
|
||||
|
|
@ -43,6 +44,10 @@ ToolParam = Any
|
|||
LITELLM_PROXY_MCP_SERVER_URL = "litellm_proxy"
|
||||
LITELLM_PROXY_MCP_SERVER_URL_PREFIX = f"{LITELLM_PROXY_MCP_SERVER_URL}/mcp/"
|
||||
|
||||
# Matches any URL ending in /mcp/<server_name> — e.g. http://localhost:4000/mcp/atlassian_test
|
||||
# Used to auto-route requests targeting the proxy's own MCP endpoint through the internal handler.
|
||||
_PROXY_MCP_PATH_RE = re.compile(r"^.+/mcp/([^/]+)$")
|
||||
|
||||
|
||||
class LiteLLM_Proxy_MCP_Handler:
|
||||
"""
|
||||
|
|
@ -54,7 +59,8 @@ class LiteLLM_Proxy_MCP_Handler:
|
|||
@staticmethod
|
||||
def _should_use_litellm_mcp_gateway(tools: Optional[Iterable[ToolParam]]) -> bool:
|
||||
"""
|
||||
Returns True if the user passed a MCP tool with server_url="litellm_proxy"
|
||||
Returns True if any MCP tool should be handled via the litellm proxy MCP gateway.
|
||||
This includes tools with server_url="litellm_proxy" as well as URLs ending in /mcp/<name>.
|
||||
"""
|
||||
if tools:
|
||||
for tool in tools:
|
||||
|
|
@ -64,6 +70,10 @@ class LiteLLM_Proxy_MCP_Handler:
|
|||
LITELLM_PROXY_MCP_SERVER_URL
|
||||
):
|
||||
return True
|
||||
if isinstance(server_url, str) and _PROXY_MCP_PATH_RE.match(
|
||||
server_url
|
||||
):
|
||||
return True
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -87,6 +97,18 @@ class LiteLLM_Proxy_MCP_Handler:
|
|||
LITELLM_PROXY_MCP_SERVER_URL
|
||||
):
|
||||
mcp_tools_with_litellm_proxy.append(tool)
|
||||
elif isinstance(server_url, str):
|
||||
# Also intercept URLs like http://localhost:4000/mcp/atlassian_test
|
||||
# by rewriting them to the internal litellm_proxy format.
|
||||
m = _PROXY_MCP_PATH_RE.match(server_url)
|
||||
if m:
|
||||
rewritten = dict(tool)
|
||||
rewritten["server_url"] = (
|
||||
f"{LITELLM_PROXY_MCP_SERVER_URL_PREFIX}{m.group(1)}"
|
||||
)
|
||||
mcp_tools_with_litellm_proxy.append(rewritten)
|
||||
else:
|
||||
other_tools.append(tool)
|
||||
else:
|
||||
other_tools.append(tool)
|
||||
else:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue