diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index a728d91271..9965c003b0 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -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") diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index fe56a418e7..e703a3956b 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -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. diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index dd7f092cb5..750c9204c5 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -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 = ( diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 1402b385a8..864b17afe3 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -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) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 6ba8d0b68a..f13dc5efd8 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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 diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index a728d91271..9965c003b0 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -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") diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 805494b185..81fb424c15 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -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[ diff --git a/schema.prisma b/schema.prisma index a728d91271..9965c003b0 100644 --- a/schema.prisma +++ b/schema.prisma @@ -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")