From 96ed00e1840669accc54c1fbeb19639ce32023a2 Mon Sep 17 00:00:00 2001 From: Milan Date: Tue, 14 Apr 2026 14:19:31 +0300 Subject: [PATCH 1/4] feat(mcp): gateway InitializeResult.instructions from upstream or YAML - Add optional instructions on MCPServer (config/DB/types) and Prisma migration. - MCPClient: fetch_upstream_initialize_instructions() for one-shot initialize. - Gateway merges per-request instructions: YAML/API overrides; otherwise fetch upstream initialize instructions (skip spec_path/OpenAPI-only servers). - Pass auth headers into instruction merge; ContextVar for gateway Server. - REST: wire instructions on connection-test MCPServer payloads. Made-with: Cursor --- .../migration.sql | 2 + .../litellm_proxy_extras/schema.prisma | 1 + litellm/experimental_mcp_client/client.py | 46 +++++ .../_experimental/mcp_server/mcp_context.py | 5 + .../mcp_server/mcp_server_manager.py | 4 + .../mcp_server/rest_endpoints.py | 1 + .../proxy/_experimental/mcp_server/server.py | 178 ++++++++++++++++-- litellm/proxy/_types.py | 4 + litellm/proxy/schema.prisma | 1 + .../types/mcp_server/mcp_server_manager.py | 2 + schema.prisma | 1 + 11 files changed, 229 insertions(+), 16 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260414140000_add_mcp_server_instructions/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260414140000_add_mcp_server_instructions/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260414140000_add_mcp_server_instructions/migration.sql new file mode 100644 index 00000000000..531024c519f --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260414140000_add_mcp_server_instructions/migration.sql @@ -0,0 +1,2 @@ +-- AlterTable +ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "instructions" TEXT; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index fce95465b55..a728d912715 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -289,6 +289,7 @@ model LiteLLM_MCPServerTable { server_name String? alias String? description String? + instructions String? // MCP InitializeResult.instructions (optional) 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 1423617cac0..fe56a418e7b 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -329,6 +329,52 @@ 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_context.py b/litellm/proxy/_experimental/mcp_server/mcp_context.py index 12830db1d6a..a60138dd340 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_context.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_context.py @@ -14,3 +14,8 @@ from typing import Optional _mcp_active_toolset_id: ContextVar[Optional[str]] = ContextVar( "_mcp_active_toolset_id", default=None ) + +# Per-request merged InitializeResult.instructions; set in MCP HTTP/SSE handlers. +_mcp_gateway_initialize_instructions: ContextVar[Optional[str]] = ContextVar( + "_mcp_gateway_initialize_instructions", default=None +) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 8d3831e75fb..dd7f092cb53 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -351,6 +351,7 @@ class MCPServerManager: aws_service_name=server_config.get("aws_service_name", None), aws_role_name=server_config.get("aws_role_name", None), aws_session_name=server_config.get("aws_session_name", None), + instructions=server_config.get("instructions", None), ) self.config_mcp_servers[server_id] = new_server @@ -693,6 +694,7 @@ class MCPServerManager: aws_service_name=aws_creds.get("aws_service_name"), aws_role_name=aws_creds.get("aws_role_name"), aws_session_name=aws_creds.get("aws_session_name"), + instructions=mcp_server.instructions, ) return new_server @@ -2946,6 +2948,7 @@ class MCPServerManager: token_url=server.token_url, registration_url=server.registration_url, allow_all_keys=server.allow_all_keys, + instructions=server.instructions, ) async def get_all_mcp_servers_with_health_and_teams( @@ -3041,6 +3044,7 @@ class MCPServerManager: is_byok=server.is_byok, byok_description=server.byok_description, byok_api_key_help_url=server.byok_api_key_help_url, + instructions=server.instructions, ) async def get_all_mcp_servers_unfiltered(self) -> List[LiteLLM_MCPServerTable]: diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 32560a2211d..8131c040136 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -933,6 +933,7 @@ if MCP_AVAILABLE: authorization_url=request.authorization_url, registration_url=request.registration_url, oauth2_flow=_oauth2_flow, + instructions=request.instructions, ) stdio_env = global_mcp_server_manager._build_stdio_env( diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 99578d006e1..1402b385a8e 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -6,6 +6,7 @@ LiteLLM MCP Server Routes import asyncio import contextlib +import contextvars import time import traceback import uuid @@ -37,7 +38,10 @@ from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( get_request_base_url, ) -from litellm.proxy._experimental.mcp_server.mcp_context import _mcp_active_toolset_id +from litellm.proxy._experimental.mcp_server.mcp_context import ( + _mcp_active_toolset_id, + _mcp_gateway_initialize_instructions, +) from litellm.proxy._experimental.mcp_server.mcp_debug import MCPDebug from litellm.proxy._experimental.mcp_server.utils import ( LITELLM_MCP_SERVER_DESCRIPTION, @@ -122,6 +126,8 @@ _INITIALIZATION_LOCK = asyncio.Lock() if MCP_AVAILABLE: from mcp.server import Server + from mcp.server.lowlevel.server import NotificationOptions + from mcp.server.models import InitializationOptions # Import auth context variables and middleware from mcp.server.auth.middleware.auth_context import ( @@ -200,10 +206,27 @@ if MCP_AVAILABLE: ) return normalized + class _LitellmMcpGatewayServer(Server): + """Gateway server that injects per-request ``InitializeResult.instructions``.""" + + def create_initialization_options( # type: ignore[override] + 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 + ######################################################## ############ Initialize the MCP Server ################# ######################################################## - server: Server = Server( + server: Server = _LitellmMcpGatewayServer( name=LITELLM_MCP_SERVER_NAME, version=LITELLM_MCP_SERVER_VERSION, ) @@ -814,10 +837,7 @@ if MCP_AVAILABLE: return tools def _get_client_ip_from_context() -> Optional[str]: - """ - Extract client_ip from auth context. - Returns None if context not set (caller should handle this as "no IP filtering"). - """ + """Return ``client_ip`` from MCP auth context (set by HTTP/SSE handlers), or None.""" try: auth_user = auth_context_var.get() if auth_user and isinstance(auth_user, MCPAuthenticatedUser): @@ -836,19 +856,15 @@ 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. 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. + 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. """ - # 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 called without client_ip and no auth context. " - "IP filtering will be skipped. This is expected for internal calls." + "MCP _get_allowed_mcp_servers: client IP unknown; skipping public-internet IP filter." ) allowed_mcp_server_ids = ( @@ -1103,6 +1119,112 @@ if MCP_AVAILABLE: return server_auth_header, extra_headers + async 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.""" + 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]]: + label = ( + server.alias + or server.server_name + or server.name + or server.server_id + or "mcp" + ) + if server.instructions and server.instructions.strip(): + return (label, server.instructions.strip()) + 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, + ) + 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 + + 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: + return texts[0][1] + return "\n\n---\n\n".join(f"[{lbl}]\n{txt}" for lbl, txt in texts) + + async def _set_mcp_gateway_initialize_instructions_token( + 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, + ) + return _mcp_gateway_initialize_instructions.set(merged) + async def _get_tools_from_mcp_servers( # noqa: PLR0915 user_api_key_auth: Optional[UserAPIKeyAuth], mcp_auth_header: Optional[str], @@ -2670,7 +2792,19 @@ if MCP_AVAILABLE: # Request was fully handled (e.g., DELETE on non-existent session) return - await session_manager.handle_request(scope, receive, send) + _instr_tok = await _set_mcp_gateway_initialize_instructions_token( + 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) + finally: + _mcp_gateway_initialize_instructions.reset(_instr_tok) except HTTPException: # Re-raise HTTP exceptions to preserve status codes and details raise @@ -2729,7 +2863,19 @@ if MCP_AVAILABLE: await initialize_session_managers() await asyncio.sleep(0.1) - await sse_session_manager.handle_request(scope, receive, send) + _sse_instr_tok = await _set_mcp_gateway_initialize_instructions_token( + 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) + finally: + _mcp_gateway_initialize_instructions.reset(_sse_instr_tok) except Exception as e: verbose_logger.exception(f"Error handling MCP request: {e}") # Instead of re-raising, try to send a graceful error response diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 0bbee56d5e0..6ba8d0b68a3 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1137,6 +1137,8 @@ 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 args: List[str] = Field(default_factory=list) @@ -1219,6 +1221,7 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase): tool_name_to_description: Optional[Dict[str, str]] = None extra_headers: Optional[List[str]] = None static_headers: Optional[Dict[str, str]] = None + instructions: Optional[str] = None # Stdio-specific fields command: Optional[str] = None args: List[str] = Field(default_factory=list) @@ -1270,6 +1273,7 @@ class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase): transport: MCPTransportType auth_type: Optional[MCPAuthType] = None credentials: Optional[MCPCredentials] = None + instructions: Optional[str] = None created_at: Optional[datetime] = None created_by: Optional[str] = None updated_at: Optional[datetime] = None diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index fce95465b55..a728d912715 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -289,6 +289,7 @@ model LiteLLM_MCPServerTable { server_name String? alias String? description String? + instructions String? // MCP InitializeResult.instructions (optional) 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 a7d0968c0ef..805494b1854 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -27,6 +27,8 @@ 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[ List[str] diff --git a/schema.prisma b/schema.prisma index fce95465b55..a728d912715 100644 --- a/schema.prisma +++ b/schema.prisma @@ -289,6 +289,7 @@ model LiteLLM_MCPServerTable { server_name String? alias String? description String? + instructions String? // MCP InitializeResult.instructions (optional) url String? spec_path String? transport String @default("sse") From 8c2ebee4decd5a9bfe3b8ad5fb46fc8f5ca1e5fb Mon Sep 17 00:00:00 2001 From: Milan Date: Tue, 14 Apr 2026 14:36:06 +0300 Subject: [PATCH 2/4] refactor(mcp): reuse existing sessions for initialize instructions 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 --- .../litellm_proxy_extras/schema.prisma | 2 +- litellm/experimental_mcp_client/client.py | 55 +------ .../mcp_server/mcp_server_manager.py | 18 +++ .../proxy/_experimental/mcp_server/server.py | 143 +++++------------- litellm/proxy/_types.py | 1 - litellm/proxy/schema.prisma | 2 +- .../types/mcp_server/mcp_server_manager.py | 1 - schema.prisma | 2 +- 8 files changed, 71 insertions(+), 153 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index a728d912715..9965c003b0a 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 fe56a418e7b..e703a3956b9 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 dd7f092cb53..750c9204c55 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 1402b385a8e..864b17afe33 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 6ba8d0b68a3..f13dc5efd84 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 a728d912715..9965c003b0a 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 805494b1854..81fb424c153 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 a728d912715..9965c003b0a 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") From 7e656f4329becd0164ccedd56bf102f452fdcd72 Mon Sep 17 00:00:00 2001 From: Milan Date: Tue, 14 Apr 2026 15:54:27 +0300 Subject: [PATCH 3/4] test: add unit tests for MCP initialize instructions feature MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extend existing test modules with coverage for the instructions merge logic, upstream cache, ContextVar-based injection, and client-side capture — following each file's established patterns. Made-with: Cursor --- .../proxy/_experimental/mcp_server/server.py | 26 ++- .../test_mcp_client.py | 75 ++++++++ .../mcp_server/test_mcp_server.py | 175 ++++++++++++++++++ .../mcp_server/test_mcp_server_manager.py | 72 +++++++ 4 files changed, 334 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 864b17afe33..b7dbdeed5fc 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -6,7 +6,6 @@ LiteLLM MCP Server Routes import asyncio import contextlib -import contextvars import time import types import traceback @@ -1161,18 +1160,23 @@ if MCP_AVAILABLE: return texts[0][1] return "\n\n---\n\n".join(f"[{lbl}]\n{txt}" for lbl, txt in texts) - async def _set_mcp_gateway_initialize_instructions_token( + @contextlib.asynccontextmanager + async def _gateway_initialize_instructions_request_scope( user_api_key_auth: Optional[UserAPIKeyAuth], mcp_servers: Optional[List[str]], client_ip: Optional[str], - ) -> contextvars.Token[Optional[str]]: + ) -> AsyncIterator[None]: allowed = await _get_allowed_mcp_servers( user_api_key_auth=user_api_key_auth, mcp_servers=mcp_servers, client_ip=client_ip, ) merged = _merge_gateway_initialize_instructions(allowed_mcp_servers=allowed) - return _mcp_gateway_initialize_instructions.set(merged) + tok = _mcp_gateway_initialize_instructions.set(merged) + try: + yield + finally: + _mcp_gateway_initialize_instructions.reset(tok) async def _get_tools_from_mcp_servers( # noqa: PLR0915 user_api_key_auth: Optional[UserAPIKeyAuth], @@ -2741,15 +2745,12 @@ if MCP_AVAILABLE: # Request was fully handled (e.g., DELETE on non-existent session) return - _instr_tok = await _set_mcp_gateway_initialize_instructions_token( + async with _gateway_initialize_instructions_request_scope( user_api_key_auth, mcp_servers, _client_ip, - ) - try: + ): await session_manager.handle_request(scope, receive, send) - finally: - _mcp_gateway_initialize_instructions.reset(_instr_tok) except HTTPException: # Re-raise HTTP exceptions to preserve status codes and details raise @@ -2808,15 +2809,12 @@ if MCP_AVAILABLE: await initialize_session_managers() await asyncio.sleep(0.1) - _sse_instr_tok = await _set_mcp_gateway_initialize_instructions_token( + async with _gateway_initialize_instructions_request_scope( user_api_key_auth, mcp_servers, _sse_client_ip, - ) - try: + ): await sse_session_manager.handle_request(scope, receive, send) - finally: - _mcp_gateway_initialize_instructions.reset(_sse_instr_tok) except Exception as e: verbose_logger.exception(f"Error handling MCP request: {e}") # Instead of re-raising, try to send a graceful error response diff --git a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py index 13a09f54e68..46d483c248f 100644 --- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py +++ b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py @@ -312,5 +312,80 @@ class TestMCPClient: assert MCPAuth.token.value == "token" +# --------------------------------------------------------------------------- +# _last_initialize_instructions capture +# --------------------------------------------------------------------------- + + +class TestMCPClientInstructionsCapture: + """Tests for _last_initialize_instructions capture during session init.""" + + def test_initial_value_is_none(self): + """Fresh client has no cached instructions.""" + client = MCPClient( + server_url="http://example.com/mcp", + transport_type="http", + ) + assert client._last_initialize_instructions is None + + @pytest.mark.asyncio + @patch("litellm.experimental_mcp_client.client.ClientSession") + async def test_captures_instructions_from_initialize(self, mock_session_cls): + """Instructions from upstream initialize() are captured and stripped.""" + client = MCPClient( + server_url="http://example.com/mcp", + transport_type="http", + ) + + mock_session = AsyncMock() + init_result = MagicMock() + init_result.instructions = " upstream says hello " + mock_session.initialize = AsyncMock(return_value=init_result) + + session_ctx = MagicMock() + session_ctx.__aenter__ = AsyncMock(return_value=mock_session) + session_ctx.__aexit__ = AsyncMock(return_value=False) + mock_session_cls.return_value = session_ctx + + transport_ctx = MagicMock() + transport_ctx.__aenter__ = AsyncMock(return_value=(MagicMock(), MagicMock())) + transport_ctx.__aexit__ = AsyncMock(return_value=False) + + async def _op(session): + return "done" + + await client._execute_session_operation(transport_ctx, _op) + assert client._last_initialize_instructions == "upstream says hello" + + @pytest.mark.asyncio + @patch("litellm.experimental_mcp_client.client.ClientSession") + async def test_none_instructions_stays_none(self, mock_session_cls): + """When upstream returns no instructions the field stays None.""" + client = MCPClient( + server_url="http://example.com/mcp", + transport_type="http", + ) + + mock_session = AsyncMock() + init_result = MagicMock() + init_result.instructions = None + mock_session.initialize = AsyncMock(return_value=init_result) + + session_ctx = MagicMock() + session_ctx.__aenter__ = AsyncMock(return_value=mock_session) + session_ctx.__aexit__ = AsyncMock(return_value=False) + mock_session_cls.return_value = session_ctx + + transport_ctx = MagicMock() + transport_ctx.__aenter__ = AsyncMock(return_value=(MagicMock(), MagicMock())) + transport_ctx.__aexit__ = AsyncMock(return_value=False) + + async def _op(session): + return "done" + + await client._execute_session_operation(transport_ctx, _op) + assert client._last_initialize_instructions is None + + if __name__ == "__main__": pytest.main([__file__]) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 384d428888f..acba06afee4 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -2421,3 +2421,178 @@ async def test_get_tools_from_mcp_servers_injects_stored_oauth2_token(): assert call_kwargs["extra_headers"] == {"Authorization": f"Bearer {STORED_TOKEN}"} assert tools == [tool_1] + + +# --------------------------------------------------------------------------- +# _merge_gateway_initialize_instructions + ContextVar / InitializationOptions +# --------------------------------------------------------------------------- + + +def _make_instruction_server( + server_id="s1", + name="s1", + *, + alias=None, + server_name=None, + instructions=None, + spec_path=None, + url="https://example.com", +): + return MCPServer( + server_id=server_id, + name=name, + alias=alias, + server_name=server_name, + url=url, + transport=MCPTransport.http, + instructions=instructions, + spec_path=spec_path, + ) + + +class TestMergeGatewayInitializeInstructions: + """Tests for _merge_gateway_initialize_instructions.""" + + def _merge(self, servers): + try: + from litellm.proxy._experimental.mcp_server.server import ( + _merge_gateway_initialize_instructions, + ) + except ImportError: + pytest.skip("MCP server not available") + return _merge_gateway_initialize_instructions(servers) + + def test_empty_server_list_returns_none(self): + """No servers yields no instructions.""" + assert self._merge([]) is None + + def test_single_server_yaml_instructions(self): + """A single server with YAML instructions returns them verbatim.""" + s = _make_instruction_server(instructions="Use add() for sums.") + assert self._merge([s]) == "Use add() for sums." + + def test_yaml_instructions_strips_whitespace(self): + """Leading/trailing whitespace is stripped.""" + s = _make_instruction_server(instructions=" padded \n") + assert self._merge([s]) == "padded" + + def test_yaml_override_beats_upstream_cache(self): + """YAML/DB instructions take precedence over upstream cache.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + global_mcp_server_manager._upstream_initialize_instructions_by_server_id["s1"] = "upstream" + try: + s = _make_instruction_server(instructions="yaml wins") + assert self._merge([s]) == "yaml wins" + finally: + global_mcp_server_manager._upstream_initialize_instructions_by_server_id.pop("s1", None) + + def test_upstream_cache_used_when_no_yaml(self): + """Upstream cached instructions are used when no YAML override is set.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + global_mcp_server_manager._upstream_initialize_instructions_by_server_id["s1"] = "from upstream" + try: + s = _make_instruction_server(instructions=None) + assert self._merge([s]) == "from upstream" + finally: + global_mcp_server_manager._upstream_initialize_instructions_by_server_id.pop("s1", None) + + def test_spec_path_servers_skipped(self): + """OpenAPI (spec_path) servers do not contribute instructions.""" + s = _make_instruction_server(spec_path="/openapi.json", url=None) + assert self._merge([s]) is None + + def test_no_instructions_no_cache_returns_none(self): + """Server with no instructions and no cache yields None.""" + s = _make_instruction_server() + assert self._merge([s]) is None + + def test_multiple_servers_merged_with_labels(self): + """Multiple servers get label-prefixed and separator-joined.""" + s1 = _make_instruction_server(server_id="a", name="a", alias="Alpha", instructions="instr A") + s2 = _make_instruction_server(server_id="b", name="b", alias="Beta", instructions="instr B") + result = self._merge([s1, s2]) + assert result is not None + assert "[Alpha]" in result and "[Beta]" in result + assert "instr A" in result and "instr B" in result + assert "---" in result + + def test_single_server_no_label_wrapping(self): + """A single server's instructions are not wrapped with a label.""" + s = _make_instruction_server(alias="MyServer", instructions="single") + result = self._merge([s]) + assert result == "single" + assert "[MyServer]" not in result + + def test_mixed_yaml_cache_specpath(self): + """YAML, upstream-cache, and spec_path servers are handled correctly together.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + global_mcp_server_manager._upstream_initialize_instructions_by_server_id["c"] = "cached C" + try: + s_yaml = _make_instruction_server(server_id="a", name="a", alias="A", instructions="yaml A") + s_spec = _make_instruction_server(server_id="b", name="b", alias="B", spec_path="/spec.json", url=None) + s_cached = _make_instruction_server(server_id="c", name="c", alias="C") + result = self._merge([s_yaml, s_spec, s_cached]) + assert "yaml A" in result + assert "cached C" in result + assert "[B]" not in result + finally: + global_mcp_server_manager._upstream_initialize_instructions_by_server_id.pop("c", None) + + +class TestGatewayCreateInitializationOptions: + """Tests for the patched server.create_initialization_options via ContextVar.""" + + def test_no_contextvar_returns_default_options(self): + """When ContextVar is None, instructions are absent.""" + try: + from litellm.proxy._experimental.mcp_server.mcp_context import ( + _mcp_gateway_initialize_instructions, + ) + from litellm.proxy._experimental.mcp_server.server import server + except ImportError: + pytest.skip("MCP server not available") + + tok = _mcp_gateway_initialize_instructions.set(None) + try: + opts = server.create_initialization_options() + assert getattr(opts, "instructions", None) is None + finally: + _mcp_gateway_initialize_instructions.reset(tok) + + def test_contextvar_set_injects_instructions(self): + """When ContextVar has a value, it appears in InitializationOptions.""" + try: + from litellm.proxy._experimental.mcp_server.mcp_context import ( + _mcp_gateway_initialize_instructions, + ) + from litellm.proxy._experimental.mcp_server.server import server + except ImportError: + pytest.skip("MCP server not available") + + tok = _mcp_gateway_initialize_instructions.set("hello from merge") + try: + opts = server.create_initialization_options() + assert opts.instructions == "hello from merge" + finally: + _mcp_gateway_initialize_instructions.reset(tok) + + def test_contextvar_reset_removes_instructions(self): + """After resetting the ContextVar, instructions disappear.""" + try: + from litellm.proxy._experimental.mcp_server.mcp_context import ( + _mcp_gateway_initialize_instructions, + ) + from litellm.proxy._experimental.mcp_server.server import server + except ImportError: + pytest.skip("MCP server not available") + + tok = _mcp_gateway_initialize_instructions.set("temporary") + _mcp_gateway_initialize_instructions.reset(tok) + opts = server.create_initialization_options() + assert getattr(opts, "instructions", None) is None diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 656a9c616e8..503ef71173f 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -2483,5 +2483,77 @@ class TestHasClientCredentialsOAuth2Flow: assert server.needs_user_oauth_token is False +# --------------------------------------------------------------------------- +# Upstream initialize-instructions cache +# --------------------------------------------------------------------------- + + +class TestMCPServerManagerUpstreamInstructionsCache: + """Tests for the upstream initialize-instructions cache.""" + + def test_get_returns_none_when_empty(self): + """Empty cache returns None for any key.""" + manager = MCPServerManager() + assert manager.get_upstream_initialize_instructions("nonexistent") is None + + def test_remember_stores_stripped_value(self): + """_remember_upstream_initialize_instructions stores a stripped string.""" + manager = MCPServerManager() + fake_server = MagicMock(server_id="srv") + fake_client = MagicMock(_last_initialize_instructions=" hello \n") + manager._remember_upstream_initialize_instructions(fake_server, fake_client) + assert manager.get_upstream_initialize_instructions("srv") == "hello" + + def test_remember_ignores_empty_string(self): + """Whitespace-only instructions are not stored.""" + manager = MCPServerManager() + fake_server = MagicMock(server_id="srv") + fake_client = MagicMock(_last_initialize_instructions=" ") + manager._remember_upstream_initialize_instructions(fake_server, fake_client) + assert manager.get_upstream_initialize_instructions("srv") is None + + def test_remember_ignores_none(self): + """None instructions are not stored.""" + manager = MCPServerManager() + fake_server = MagicMock(server_id="srv") + fake_client = MagicMock(_last_initialize_instructions=None) + manager._remember_upstream_initialize_instructions(fake_server, fake_client) + assert manager.get_upstream_initialize_instructions("srv") is None + + @pytest.mark.asyncio + async def test_load_servers_from_config_clears_cache(self): + """Reloading config clears any previously cached upstream instructions.""" + manager = MCPServerManager() + manager._upstream_initialize_instructions_by_server_id["old"] = "stale" + await manager.load_servers_from_config( + mcp_servers_config={ + "fresh_srv": { + "url": "https://example.com", + "instructions": "from yaml", + } + } + ) + assert manager.get_upstream_initialize_instructions("old") is None + + @pytest.mark.asyncio + async def test_load_servers_reads_instructions_from_config(self): + """instructions field from YAML config is persisted on the MCPServer.""" + manager = MCPServerManager() + await manager.load_servers_from_config( + mcp_servers_config={ + "srv_a": { + "url": "https://a.example.com", + "instructions": "A instructions", + }, + "srv_b": { + "url": "https://b.example.com", + }, + } + ) + by_name = {s.server_name: s for s in manager.config_mcp_servers.values()} + assert "srv_a" in by_name and by_name["srv_a"].instructions == "A instructions" + assert "srv_b" in by_name and by_name["srv_b"].instructions is None + + if __name__ == "__main__": pytest.main([__file__]) From e7c630ed1998174e00e660a104fe96c9829395cc Mon Sep 17 00:00:00 2001 From: Milan Date: Tue, 14 Apr 2026 15:59:00 +0300 Subject: [PATCH 4/4] refactor: inline get_upstream_initialize_instructions Remove the trivial one-line wrapper and access the dict directly. Made-with: Cursor --- .../_experimental/mcp_server/mcp_server_manager.py | 3 --- litellm/proxy/_experimental/mcp_server/server.py | 2 +- .../mcp_server/test_mcp_server_manager.py | 10 +++++----- 3 files changed, 6 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 750c9204c55..68b858868ad 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -186,9 +186,6 @@ class MCPServerManager: """ 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: diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index b7dbdeed5fc..adeacc06f8a 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -1148,7 +1148,7 @@ if MCP_AVAILABLE: continue if server.spec_path: continue - cached = global_mcp_server_manager.get_upstream_initialize_instructions( + cached = global_mcp_server_manager._upstream_initialize_instructions_by_server_id.get( server.server_id ) if cached and cached.strip(): diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 503ef71173f..aa95836a927 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -2494,7 +2494,7 @@ class TestMCPServerManagerUpstreamInstructionsCache: def test_get_returns_none_when_empty(self): """Empty cache returns None for any key.""" manager = MCPServerManager() - assert manager.get_upstream_initialize_instructions("nonexistent") is None + assert manager._upstream_initialize_instructions_by_server_id.get("nonexistent") is None def test_remember_stores_stripped_value(self): """_remember_upstream_initialize_instructions stores a stripped string.""" @@ -2502,7 +2502,7 @@ class TestMCPServerManagerUpstreamInstructionsCache: fake_server = MagicMock(server_id="srv") fake_client = MagicMock(_last_initialize_instructions=" hello \n") manager._remember_upstream_initialize_instructions(fake_server, fake_client) - assert manager.get_upstream_initialize_instructions("srv") == "hello" + assert manager._upstream_initialize_instructions_by_server_id.get("srv") == "hello" def test_remember_ignores_empty_string(self): """Whitespace-only instructions are not stored.""" @@ -2510,7 +2510,7 @@ class TestMCPServerManagerUpstreamInstructionsCache: fake_server = MagicMock(server_id="srv") fake_client = MagicMock(_last_initialize_instructions=" ") manager._remember_upstream_initialize_instructions(fake_server, fake_client) - assert manager.get_upstream_initialize_instructions("srv") is None + assert manager._upstream_initialize_instructions_by_server_id.get("srv") is None def test_remember_ignores_none(self): """None instructions are not stored.""" @@ -2518,7 +2518,7 @@ class TestMCPServerManagerUpstreamInstructionsCache: fake_server = MagicMock(server_id="srv") fake_client = MagicMock(_last_initialize_instructions=None) manager._remember_upstream_initialize_instructions(fake_server, fake_client) - assert manager.get_upstream_initialize_instructions("srv") is None + assert manager._upstream_initialize_instructions_by_server_id.get("srv") is None @pytest.mark.asyncio async def test_load_servers_from_config_clears_cache(self): @@ -2533,7 +2533,7 @@ class TestMCPServerManagerUpstreamInstructionsCache: } } ) - assert manager.get_upstream_initialize_instructions("old") is None + assert manager._upstream_initialize_instructions_by_server_id.get("old") is None @pytest.mark.asyncio async def test_load_servers_reads_instructions_from_config(self):