From 7270f723de7822de9c731b60f5aec69271e7bb2a Mon Sep 17 00:00:00 2001 From: milan-berri Date: Sat, 23 May 2026 02:00:47 +0300 Subject: [PATCH] fix(mcp): forward upstream initialize instructions on cold gateway init (#28231) Prefetch upstream InitializeResult.instructions before merging gateway initialize options when YAML/DB do not set instructions, so clients receive upstream server text on the first MCP initialize without list_tools. Co-authored-by: Cursor --- .../mcp_server/mcp_server_manager.py | 81 +++++++ .../proxy/_experimental/mcp_server/server.py | 16 +- .../mcp_server/test_mcp_server.py | 200 ++++++++++++++++++ 3 files changed, 296 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index a72e8e34a4..d0e9ad7b2a 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -12,6 +12,7 @@ import hashlib import json import os import re +import time from typing import Any, Callable, Dict, List, Literal, Optional, Set, Tuple, Union, cast from urllib.parse import urlparse @@ -250,6 +251,10 @@ class MCPServerManager: } """ self._upstream_initialize_instructions_by_server_id: Dict[str, str] = {} + # Per-server monotonic timestamp of last upstream prefetch attempt (success, + # empty result, or failure). Used to throttle re-probes for servers that do + # not return instructions, and to apply a short cooldown after failures. + self._upstream_initialize_instructions_probed_at: Dict[str, float] = {} def _remember_upstream_initialize_instructions( self, server: MCPServer, client: MCPClient @@ -260,6 +265,80 @@ class MCPServerManager: raw ).strip() + async def _ensure_upstream_initialize_instructions_cached( + self, server: MCPServer + ) -> None: + """ + Open one upstream session and cache InitializeResult.instructions if missing. + + No-op when: + - YAML/DB instructions are set on the server record, + - server is OpenAPI (spec_path), + - non-empty upstream instructions are already cached, + - auth preconditions match health_check_server's skip rules + (per-user auth / missing static auth token), + - a prior probe attempt for this server is within + MCP_HEALTH_CHECK_TIMEOUT seconds (the probe is a health-check-shaped + op and already uses this knob for its inner call timeout; reusing it + as the cooldown avoids reconnecting on every gateway initialize when + upstream returns empty or fails). + """ + if server.spec_path: + return + if server.instructions and server.instructions.strip(): + return + if self._upstream_initialize_instructions_by_server_id.get(server.server_id): + return + if server.requires_per_user_auth: + return + if ( + server.auth_type + and server.auth_type != MCPAuth.none + and server.auth_type != MCPAuth.aws_sigv4 + and not server.authentication_token + ): + return + + last_probed_at = self._upstream_initialize_instructions_probed_at.get( + server.server_id + ) + if ( + last_probed_at is not None + and (time.monotonic() - last_probed_at) < MCP_HEALTH_CHECK_TIMEOUT + ): + return + + # Record the attempt up-front so that a failure / empty response does not + # cause every subsequent initialize request to re-open the upstream session. + self._upstream_initialize_instructions_probed_at[server.server_id] = ( + time.monotonic() + ) + + try: + extra_headers: Optional[Dict[str, str]] = ( + dict(server.static_headers) if server.static_headers else None + ) + client = await self._create_mcp_client( + server=server, + mcp_auth_header=None, + extra_headers=extra_headers, + stdio_env=None, + ) + + async def _noop(_session): + return "ok" + + await asyncio.wait_for( + client.run_with_session(_noop), timeout=MCP_HEALTH_CHECK_TIMEOUT + ) + self._remember_upstream_initialize_instructions(server, client) + except Exception as e: + verbose_logger.debug( + "Upstream initialize instructions prefetch failed for %s: %s", + server.name, + e, + ) + def get_registry(self) -> Dict[str, MCPServer]: """ Get the registered MCP Servers from the registry and union with the config MCP Servers @@ -280,6 +359,7 @@ class MCPServerManager: """ verbose_logger.debug("Loading MCP Servers from config-----") self._upstream_initialize_instructions_by_server_id.clear() + self._upstream_initialize_instructions_probed_at.clear() # Track which aliases have been used to ensure only first occurrence is used used_aliases = set() @@ -3141,6 +3221,7 @@ class MCPServerManager: verbose_logger.debug("Loading MCP servers from database into registry...") self._upstream_initialize_instructions_by_server_id.clear() + self._upstream_initialize_instructions_probed_at.clear() # perform authz check to filter the mcp servers user has access to prisma_client = get_prisma_client_or_throw( diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 5205426edf..f31005be0c 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -1165,7 +1165,7 @@ if MCP_AVAILABLE: def _merge_gateway_initialize_instructions( allowed_mcp_servers: List[MCPServer], ) -> Optional[str]: - """YAML/DB override, else in-memory upstream text from list_tools / health_check / call_tool.""" + """YAML/DB override, else upstream text (prefetch on init, or list_tools / health_check / call_tool cache).""" if not allowed_mcp_servers: return None @@ -1206,6 +1206,20 @@ if MCP_AVAILABLE: mcp_servers=mcp_servers, client_ip=client_ip, ) + if allowed: + # return_exceptions=True: a per-server probe failure (incl. CancelledError + # bubbled from anyio task group teardown on connection refused) must not + # cancel sibling probes or 500 the gateway initialize request. + await asyncio.gather( + *[ + global_mcp_server_manager._ensure_upstream_initialize_instructions_cached( + s + ) + for s in allowed + if s is not None + ], + return_exceptions=True, + ) merged = _merge_gateway_initialize_instructions(allowed_mcp_servers=allowed) tok = _mcp_gateway_initialize_instructions.set(merged) try: 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 f2fd73f3f2..d62720ed36 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 @@ -3036,6 +3036,206 @@ class TestMergeGatewayInitializeInstructions: ) +class TestEnsureUpstreamInitializeInstructionsCached: + @pytest.mark.asyncio + async def test_skips_when_yaml_instructions_set(self): + from unittest.mock import AsyncMock, patch + + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + server = _make_instruction_server( + server_id="yaml-only", instructions="from yaml" + ) + with patch.object( + global_mcp_server_manager, "_create_mcp_client", AsyncMock() + ) as mock_create: + await global_mcp_server_manager._ensure_upstream_initialize_instructions_cached( + server + ) + mock_create.assert_not_awaited() + + @pytest.mark.asyncio + async def test_skips_when_already_cached(self): + from unittest.mock import AsyncMock, patch + + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + server = _make_instruction_server(server_id="cached-only", instructions=None) + global_mcp_server_manager._upstream_initialize_instructions_by_server_id[ + "cached-only" + ] = "warm" + try: + with patch.object( + global_mcp_server_manager, "_create_mcp_client", AsyncMock() + ) as mock_create: + await global_mcp_server_manager._ensure_upstream_initialize_instructions_cached( + server + ) + mock_create.assert_not_awaited() + finally: + global_mcp_server_manager._upstream_initialize_instructions_by_server_id.pop( + "cached-only", None + ) + + @pytest.mark.asyncio + async def test_skips_when_spec_path_set(self): + from unittest.mock import AsyncMock, patch + + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + server = _make_instruction_server( + server_id="openapi-spec", spec_path="/openapi.json", url=None + ) + with patch.object( + global_mcp_server_manager, "_create_mcp_client", AsyncMock() + ) as mock_create: + await global_mcp_server_manager._ensure_upstream_initialize_instructions_cached( + server + ) + mock_create.assert_not_awaited() + + @pytest.mark.asyncio + async def test_runs_upstream_session_and_caches(self): + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + server = _make_instruction_server(server_id="cold-server", instructions=None) + fake_client = MagicMock() + fake_client.run_with_session = AsyncMock(return_value="ok") + fake_client._last_initialize_instructions = " upstream says hi " + + with patch.object( + global_mcp_server_manager, + "_create_mcp_client", + AsyncMock(return_value=fake_client), + ): + try: + await global_mcp_server_manager._ensure_upstream_initialize_instructions_cached( + server + ) + assert ( + global_mcp_server_manager._upstream_initialize_instructions_by_server_id[ + "cold-server" + ] + == "upstream says hi" + ) + finally: + global_mcp_server_manager._upstream_initialize_instructions_by_server_id.pop( + "cold-server", None + ) + global_mcp_server_manager._upstream_initialize_instructions_probed_at.pop( + "cold-server", None + ) + + @pytest.mark.asyncio + async def test_cooldown_after_empty_upstream_response(self): + """Upstream returns no instructions → next call within cooldown must not reconnect.""" + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + server = _make_instruction_server(server_id="empty-server", instructions=None) + fake_client = MagicMock() + fake_client.run_with_session = AsyncMock(return_value="ok") + fake_client._last_initialize_instructions = None # upstream sent nothing + + create = AsyncMock(return_value=fake_client) + with patch.object(global_mcp_server_manager, "_create_mcp_client", create): + try: + await global_mcp_server_manager._ensure_upstream_initialize_instructions_cached( + server + ) + await global_mcp_server_manager._ensure_upstream_initialize_instructions_cached( + server + ) + assert create.await_count == 1, ( + "Second probe within cooldown must not reconnect to upstream" + ) + assert ( + "empty-server" + not in global_mcp_server_manager._upstream_initialize_instructions_by_server_id + ) + assert ( + "empty-server" + in global_mcp_server_manager._upstream_initialize_instructions_probed_at + ) + finally: + global_mcp_server_manager._upstream_initialize_instructions_probed_at.pop( + "empty-server", None + ) + + @pytest.mark.asyncio + async def test_cooldown_after_upstream_failure(self): + """run_with_session raises → cooldown applies, no immediate retry.""" + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + server = _make_instruction_server(server_id="boom-server", instructions=None) + fake_client = MagicMock() + fake_client.run_with_session = AsyncMock(side_effect=RuntimeError("upstream down")) + fake_client._last_initialize_instructions = None + + create = AsyncMock(return_value=fake_client) + with patch.object(global_mcp_server_manager, "_create_mcp_client", create): + try: + await global_mcp_server_manager._ensure_upstream_initialize_instructions_cached( + server + ) + await global_mcp_server_manager._ensure_upstream_initialize_instructions_cached( + server + ) + assert create.await_count == 1, ( + "Second probe within cooldown must not reconnect after failure" + ) + assert ( + "boom-server" + not in global_mcp_server_manager._upstream_initialize_instructions_by_server_id + ) + assert ( + "boom-server" + in global_mcp_server_manager._upstream_initialize_instructions_probed_at + ) + finally: + global_mcp_server_manager._upstream_initialize_instructions_probed_at.pop( + "boom-server", None + ) + + @pytest.mark.asyncio + async def test_reload_resets_probe_cooldown(self): + """load_servers_from_config clears the negative-cache map so reloads re-probe.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + global_mcp_server_manager._upstream_initialize_instructions_probed_at[ + "reload-target" + ] = 1.0 + try: + await global_mcp_server_manager.load_servers_from_config({}) + assert ( + "reload-target" + not in global_mcp_server_manager._upstream_initialize_instructions_probed_at + ) + finally: + global_mcp_server_manager._upstream_initialize_instructions_probed_at.pop( + "reload-target", None + ) + + class TestGatewayCreateInitializationOptions: """Tests for the patched server.create_initialization_options via ContextVar."""