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:
Ishaan Jaffer 2026-03-11 18:36:10 -07:00
parent c5e34313b6
commit 24279c6aaa
2 changed files with 58 additions and 1 deletions

View file

@ -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,

View file

@ -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: