tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-10 05:15:19 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(mcp): reuse existing sessions for initialize instructions

Browse Source 
Remove the gateway-specific initialize fetch path and reuse instructions captured during existing MCP calls (list_tools/health_check/call_tool), while keeping YAML/DB instructions as immediate overrides.

Made-with: Cursor
This commit is contained in:
Milan
2026-04-14 14:36:06 +03:00
parent 96ed00e184
commit 8c2ebee4de
8 changed files with 71 additions and 153 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
+1 -1
View File
@@ -289,7 +289,7 @@ model LiteLLM_MCPServerTable {
server_name String?
alias String?
description String?
instructions String? // MCP InitializeResult.instructions (optional)
instructions String?
url String?
spec_path String?
transport String @default("sse")
litellm/experimental_mcp_client/client.py
+8 -47
View File
@@ -221,6 +221,7 @@ class MCPClient:
self.extra_headers: Optional[Dict[str, str]] = extra_headers
self.ssl_verify: Optional[VerifyTypes] = ssl_verify
self._aws_auth: Optional[httpx.Auth] = aws_auth
self._last_initialize_instructions: Optional[str] = None
# handle the basic auth value if provided
if auth_value:
self.update_auth_value(auth_value)
@@ -296,7 +297,12 @@ class MCPClient:
session_ctx = ClientSession(read_stream, write_stream)
session = await session_ctx.__aenter__()
try:
await session.initialize()
init_result = await session.initialize()
self._last_initialize_instructions = None
if init_result is not None:
ins = getattr(init_result, "instructions", None)
if isinstance(ins, str) and ins.strip():
self._last_initialize_instructions = ins.strip()
return await operation(session)
finally:
try:
@@ -315,6 +321,7 @@ class MCPClient:
"""Open a session, run the provided coroutine, and clean up."""
http_client: Optional[httpx.AsyncClient] = None
try:
self._last_initialize_instructions = None
transport_ctx, http_client = self._create_transport_context()
return await self._execute_session_operation(transport_ctx, operation)
except Exception:
@@ -329,52 +336,6 @@ class MCPClient:
except BaseException as e:
verbose_logger.debug(f"Error during http_client cleanup: {e}")
async def fetch_upstream_initialize_instructions(self) -> Optional[str]:
"""Open a transport, run ``initialize`` once, return upstream ``instructions``."""
http_client: Optional[httpx.AsyncClient] = None
try:
transport_ctx, http_client = self._create_transport_context()
transport = await transport_ctx.__aenter__()
try:
read_stream, write_stream = transport[0], transport[1]
session_ctx = ClientSession(read_stream, write_stream)
session = await session_ctx.__aenter__()
try:
init = await session.initialize()
return init.instructions
finally:
try:
await session_ctx.__aexit__(None, None, None)
except BaseException as e:
verbose_logger.debug(
"Error during session context exit (instructions fetch): %s",
e,
)
finally:
try:
await transport_ctx.__aexit__(None, None, None)
except BaseException as e:
verbose_logger.debug(
"Error during transport context exit (instructions fetch): %s",
e,
)
except Exception as e:
verbose_logger.debug(
"fetch_upstream_initialize_instructions failed for %s: %s",
self.server_url or "stdio",
e,
)
return None
finally:
if http_client is not None:
try:
await http_client.aclose()
except BaseException as e:
verbose_logger.debug(
"Error during http_client cleanup (instructions fetch): %s",
e,
)
def update_auth_value(self, mcp_auth_value: Union[str, Dict[str, str]]):
"""
Set the authentication header for the MCP client.
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
+18
View File
@@ -184,6 +184,19 @@ class MCPServerManager:
"gmail_send_email": "zapier_mcp_server",
}
"""
self._upstream_initialize_instructions_by_server_id: Dict[str, str] = {}
def get_upstream_initialize_instructions(self, server_id: str) -> Optional[str]:
return self._upstream_initialize_instructions_by_server_id.get(server_id)
def _remember_upstream_initialize_instructions(
self, server: MCPServer, client: MCPClient
) -> None:
raw = getattr(client, "_last_initialize_instructions", None)
if raw and str(raw).strip():
self._upstream_initialize_instructions_by_server_id[server.server_id] = (
str(raw).strip()
)
def get_registry(self) -> Dict[str, MCPServer]:
"""
@@ -204,6 +217,7 @@ class MCPServerManager:
mcp_aliases: Optional dictionary mapping aliases to server names from litellm_settings
"""
verbose_logger.debug("Loading MCP Servers from config-----")
self._upstream_initialize_instructions_by_server_id.clear()
# Track which aliases have been used to ensure only first occurrence is used
used_aliases = set()
@@ -1249,6 +1263,7 @@ class MCPServerManager:
return tools
else:
tools = await self._fetch_tools_with_timeout(client, server.name)
self._remember_upstream_initialize_instructions(server, client)
prefixed_or_original_tools = self._create_prefixed_tools(
tools, server, add_prefix=add_prefix
@@ -2385,6 +2400,7 @@ class MCPServerManager:
# If proxy_logging_obj is not None, the tool call result is at index 1 (after the during hook task)
result_index = 1 if proxy_logging_obj else 0
result = mcp_responses[result_index]
self._remember_upstream_initialize_instructions(mcp_server, client)
return cast(CallToolResult, result)
@@ -2624,6 +2640,7 @@ class MCPServerManager:
)
verbose_logger.debug("Loading MCP servers from database into registry...")
self._upstream_initialize_instructions_by_server_id.clear()
# perform authz check to filter the mcp servers user has access to
prisma_client = get_prisma_client_or_throw(
@@ -2907,6 +2924,7 @@ class MCPServerManager:
await asyncio.wait_for(
client.run_with_session(_noop), timeout=MCP_HEALTH_CHECK_TIMEOUT
)
self._remember_upstream_initialize_instructions(server, client)
status = "healthy"
except asyncio.TimeoutError:
health_check_error = (
litellm/proxy/_experimental/mcp_server/server.py
+42 -101
View File
@@ -8,6 +8,7 @@ import asyncio
import contextlib
import contextvars
import time
import types
import traceback
import uuid
from datetime import datetime
@@ -206,30 +207,31 @@ if MCP_AVAILABLE:
)
return normalized
class _LitellmMcpGatewayServer(Server):
"""Gateway server that injects per-request ``InitializeResult.instructions``."""
def create_initialization_options( # type: ignore[override]
def _gateway_create_initialization_options(
self,
notification_options: Optional[NotificationOptions] = None,
experimental_capabilities: Optional[Dict[str, Dict[str, Any]]] = None,
) -> InitializationOptions:
opts = Server.create_initialization_options(
self,
notification_options: Optional[NotificationOptions] = None,
experimental_capabilities: Optional[Dict[str, Dict[str, Any]]] = None,
) -> InitializationOptions:
opts = super().create_initialization_options(
notification_options=notification_options,
experimental_capabilities=experimental_capabilities or {},
)
merged = _mcp_gateway_initialize_instructions.get()
if merged is not None:
return opts.model_copy(update={"instructions": merged})
return opts
notification_options=notification_options,
experimental_capabilities=experimental_capabilities or {},
)
merged = _mcp_gateway_initialize_instructions.get()
if merged is not None:
return opts.model_copy(update={"instructions": merged})
return opts
########################################################
############ Initialize the MCP Server #################
########################################################
server: Server = _LitellmMcpGatewayServer(
server: Server = Server(
name=LITELLM_MCP_SERVER_NAME,
version=LITELLM_MCP_SERVER_VERSION,
)
server.create_initialization_options = types.MethodType( # type: ignore[method-assign]
_gateway_create_initialization_options, server
)
sse: SseServerTransport = SseServerTransport("/mcp/sse/messages")
# Create session managers
@@ -837,7 +839,10 @@ if MCP_AVAILABLE:
return tools
def _get_client_ip_from_context() -> Optional[str]:
"""Return ``client_ip`` from MCP auth context (set by HTTP/SSE handlers), or None."""
"""
Extract client_ip from auth context.
Returns None if context not set (caller should handle this as "no IP filtering").
"""
try:
auth_user = auth_context_var.get()
if auth_user and isinstance(auth_user, MCPAuthenticatedUser):
@@ -856,15 +861,19 @@ if MCP_AVAILABLE:
Args:
user_api_key_auth: The authenticated user's API key info.
mcp_servers: Optional list of server names to filter to.
client_ip: Client IP for IP-based access control. MCP HTTP/SSE handlers set auth context (including ``client_ip``) before MCP work; when this is
``None``, ``client_ip`` is taken from that context. Callers may still
pass ``client_ip`` explicitly when already computed.
client_ip: Client IP for IP-based access control. If None, falls back to
auth context. Pass explicitly from request handlers for safety.
Note: If client_ip is None and auth context is not set, IP filtering is skipped.
This is intentional for internal callers but may indicate a bug if called
from a request handler without proper context setup.
"""
# Use explicit client_ip if provided, otherwise try auth context
if client_ip is None:
client_ip = _get_client_ip_from_context()
if client_ip is None:
verbose_logger.debug(
"MCP _get_allowed_mcp_servers: client IP unknown; skipping public-internet IP filter."
"MCP _get_allowed_mcp_servers called without client_ip and no auth context. "
"IP filtering will be skipped. This is expected for internal calls."
)
allowed_mcp_server_ids = (
@@ -1119,29 +1128,15 @@ if MCP_AVAILABLE:
return server_auth_header, extra_headers
async def _merge_gateway_initialize_instructions(
def _merge_gateway_initialize_instructions(
allowed_mcp_servers: List[MCPServer],
user_api_key_auth: Optional[UserAPIKeyAuth],
mcp_auth_header: Optional[str],
mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]],
oauth2_headers: Optional[Dict[str, str]],
raw_headers: Optional[Dict[str, str]],
) -> Optional[str]:
"""Merge ``instructions`` for gateway ``initialize``: YAML/API overrides upstream."""
"""YAML/DB override, else in-memory upstream text from list_tools / health_check / call_tool."""
if not allowed_mcp_servers:
return None
_has_oauth2_server = any(
getattr(s, "auth_type", None) == MCPAuth.oauth2
for s in allowed_mcp_servers
)
_prefetched_oauth_creds = (
await _prefetch_oauth_creds_for_user(user_api_key_auth)
if _has_oauth2_server
else {}
)
async def _one(server: MCPServer) -> Optional[Tuple[str, str]]:
texts: List[Tuple[str, str]] = []
for server in allowed_mcp_servers:
label = (
server.alias
or server.server_name
@@ -1150,50 +1145,16 @@ if MCP_AVAILABLE:
or "mcp"
)
if server.instructions and server.instructions.strip():
return (label, server.instructions.strip())
texts.append((label, server.instructions.strip()))
continue
if server.spec_path:
return None
server_auth_header, extra_headers = _prepare_mcp_server_headers(
server=server,
mcp_server_auth_headers=mcp_server_auth_headers,
mcp_auth_header=mcp_auth_header,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
continue
cached = global_mcp_server_manager.get_upstream_initialize_instructions(
server.server_id
)
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,
prefetched_creds=_prefetched_oauth_creds,
)
try:
if server.static_headers:
if extra_headers is None:
extra_headers = {}
extra_headers.update(server.static_headers)
stdio_env = global_mcp_server_manager._build_stdio_env(
server, raw_headers
)
client = await global_mcp_server_manager._create_mcp_client(
server=server,
mcp_auth_header=server_auth_header,
extra_headers=extra_headers,
stdio_env=stdio_env,
)
text = await client.fetch_upstream_initialize_instructions()
if text and text.strip():
return (label, text.strip())
except Exception as e:
verbose_logger.debug(
"MCP gateway: upstream instructions fetch failed for %s: %s",
server.name,
e,
)
return None
if cached and cached.strip():
texts.append((label, cached.strip()))
pairs = await asyncio.gather(*(_one(s) for s in allowed_mcp_servers))
texts = [p for p in pairs if p is not None]
if not texts:
return None
if len(texts) == 1:
@@ -1204,25 +1165,13 @@ if MCP_AVAILABLE:
user_api_key_auth: Optional[UserAPIKeyAuth],
mcp_servers: Optional[List[str]],
client_ip: Optional[str],
mcp_auth_header: Optional[str],
mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]],
oauth2_headers: Optional[Dict[str, str]],
raw_headers: Optional[Dict[str, str]],
) -> contextvars.Token[Optional[str]]:
"""Resolve merged gateway ``instructions``; return ContextVar token to reset."""
allowed = await _get_allowed_mcp_servers(
user_api_key_auth=user_api_key_auth,
mcp_servers=mcp_servers,
client_ip=client_ip,
)
merged = await _merge_gateway_initialize_instructions(
allowed_mcp_servers=allowed,
user_api_key_auth=user_api_key_auth,
mcp_auth_header=mcp_auth_header,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
)
merged = _merge_gateway_initialize_instructions(allowed_mcp_servers=allowed)
return _mcp_gateway_initialize_instructions.set(merged)
async def _get_tools_from_mcp_servers( # noqa: PLR0915
@@ -2796,10 +2745,6 @@ if MCP_AVAILABLE:
user_api_key_auth,
mcp_servers,
_client_ip,
mcp_auth_header,
mcp_server_auth_headers,
oauth2_headers,
raw_headers,
)
try:
await session_manager.handle_request(scope, receive, send)
@@ -2867,10 +2812,6 @@ if MCP_AVAILABLE:
user_api_key_auth,
mcp_servers,
_sse_client_ip,
mcp_auth_header,
mcp_server_auth_headers,
oauth2_headers,
raw_headers,
)
try:
await sse_session_manager.handle_request(scope, receive, send)
litellm/proxy/_types.py
-1
View File
@@ -1137,7 +1137,6 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase):
tool_name_to_description: Optional[Dict[str, str]] = None
extra_headers: Optional[List[str]] = None
static_headers: Optional[Dict[str, str]] = None
# Shown to MCP clients in InitializeResult.instructions (optional)
instructions: Optional[str] = None
# Stdio-specific fields
command: Optional[str] = None
litellm/proxy/schema.prisma
+1 -1
View File
@@ -289,7 +289,7 @@ model LiteLLM_MCPServerTable {
server_name String?
alias String?
description String?
instructions String? // MCP InitializeResult.instructions (optional)
instructions String?
url String?
spec_path String?
transport String @default("sse")
litellm/types/mcp_server/mcp_server_manager.py
-1
View File
@@ -27,7 +27,6 @@ class MCPServer(BaseModel):
spec_path: Optional[str] = None
auth_type: Optional[MCPAuthType] = None
authentication_token: Optional[str] = None
# Optional text returned on MCP `initialize` (InitializeResult.instructions)
instructions: Optional[str] = None
mcp_info: Optional[MCPInfo] = None
extra_headers: Optional[
schema.prisma
+1 -1
View File
@@ -289,7 +289,7 @@ model LiteLLM_MCPServerTable {
server_name String?
alias String?
description String?
instructions String? // MCP InitializeResult.instructions (optional)
instructions String?
url String?
spec_path String?
transport String @default("sse")