From b8cd3a71425a8bb0a220e540d14547222b5e664e Mon Sep 17 00:00:00 2001 From: shrey kharbanda Date: Mon, 26 Oct 2026 18:52:07 -0700 Subject: [PATCH] fix(mcp): keep cached initialize instructions accurate across catalog reloads Every full catalog reload (the 30 second reload_mcp_servers_job, the reload after an admin write, add_deployment) wiped every server's cached upstream initialize instructions. Servers that need per-user auth cannot be prefetched, so gateway initialize returned no instructions for them until a user listed tools again. The reload already drops the entry of each server whose definition changed or was removed, so the global wipe only ever cleared unchanged servers. Remove it and the method that only it used. That wipe was also the only thing that expired two kinds of stale text, so fix both at the source. Each upstream session now records whether initialize finished and what it said, and an answer without instructions drops the cached entry. A per-server generation, bumped whenever a server's definition changes, is read before every upstream session, and a session that outlived an edit of its server no longer writes the old upstream's text back. Resolves LIT-9851 --- litellm/experimental_mcp_client/client.py | 2 +- .../proxy/_experimental/mcp_server/catalog.py | 2 - .../mcp_server/mcp_server_manager.py | 37 +++-- .../test_mcp_client.py | 5 +- .../mcp_server/test_mcp_server_manager.py | 135 ++++++++++++++---- .../test_mcp_server_tool_calls_and_headers.py | 2 +- 6 files changed, 136 insertions(+), 47 deletions(-) diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 79887e4ab99..da250ed4c99 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -792,7 +792,7 @@ class MCPClient: init_result: Final = await self._prepare_session(session) instructions: Final = getattr(init_result, "instructions", None) self._last_initialize_instructions = ( - instructions.strip() or None if isinstance(instructions, str) else None + instructions.strip() if isinstance(instructions, str) else "" ) result: Final = await operation(session) except BaseException as operation_error: diff --git a/litellm/proxy/_experimental/mcp_server/catalog.py b/litellm/proxy/_experimental/mcp_server/catalog.py index 2db548f5863..2e0e5f95798 100644 --- a/litellm/proxy/_experimental/mcp_server/catalog.py +++ b/litellm/proxy/_experimental/mcp_server/catalog.py @@ -525,8 +525,6 @@ class CatalogSnapshots: self.manager.invalidate_oauth_discovery_state(server_id) self._database_identity = database_identity self.manager.registry = registered_registry - if not reuse_unchanged: - self.manager.clear_initialize_instructions() warn_on_shared_identifier_prefixes(registered_registry.values()) # A discovery task may have published into ``previous_registry`` while # this replacement was being staged. Reconcile every published entry diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index db70c7d77de..7189df89385 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2076,6 +2076,7 @@ class MCPServerManager: self._listed_tools_by_server_id: dict[str, _ListedToolsByCaller] = {} # mutable-ok: refreshed per tools/list self._listed_tools_generations: dict[str, int] = {} # mutable-ok: bumped per server save self._upstream_initialize_instructions_by_server_id: dict[str, str] = {} + self._initialize_instructions_generations: Mapping[str, int] = MappingProxyType({}) # 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. @@ -2448,10 +2449,18 @@ class MCPServerManager: return current raise HTTPException(status_code=503, detail="OAuth metadata discovery changed repeatedly; retry shortly") - def _remember_upstream_initialize_instructions(self, server: MCPServer, client: MCPClient) -> None: + def initialize_instructions_generation(self, server_id: str) -> int: + return self._initialize_instructions_generations.get(server_id, 0) + + def _remember_upstream_initialize_instructions(self, server: MCPServer, client: MCPClient, generation: int) -> None: raw: Final[str | None] = 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() + if raw is None or generation != self.initialize_instructions_generation(server.server_id): + return + instructions: Final = raw.strip() + if instructions: + self._upstream_initialize_instructions_by_server_id[server.server_id] = instructions + else: + self._upstream_initialize_instructions_by_server_id.pop(server.server_id, None) async def ensure_upstream_initialize_instructions_cached(self, server: MCPServer) -> None: """ @@ -2495,6 +2504,7 @@ class MCPServerManager: # 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() + generation: Final = self.initialize_instructions_generation(server.server_id) try: resolved_static_headers: Final = await self._resolve_static_headers_with_env_vars( @@ -2514,7 +2524,7 @@ class MCPServerManager: return "ok" await asyncio.wait_for(client.run_with_session(_noop), timeout=MCP_HEALTH_CHECK_TIMEOUT) - self._remember_upstream_initialize_instructions(server, client) + self._remember_upstream_initialize_instructions(server, client, generation) except Exception as e: verbose_logger.debug( "Upstream initialize instructions prefetch failed for %s: %s", @@ -4475,6 +4485,7 @@ class MCPServerManager: oauth2_headers=oauth2_headers, ) listed_generation: Final = self._listed_tools_generations.get(server.server_id, 0) + instructions_generation: Final = self.initialize_instructions_generation(server.server_id) try: # Tool *listing* must not be blocked by missing per-user env vars — @@ -4592,7 +4603,7 @@ class MCPServerManager: else ListToolsResult(tools=await self._fetch_tools_with_timeout(client, server.name)) ) tools = page.tools - self._remember_upstream_initialize_instructions(server, client) + self._remember_upstream_initialize_instructions(server, client, instructions_generation) guarded_tools: Final = await self._guard_tool_catalog( server=server, @@ -4650,10 +4661,6 @@ class MCPServerManager: ) raise_classified_list_failure(e, server.name, suppress_challenge=server.is_dcr_bridge) - def clear_initialize_instructions(self) -> None: - self._upstream_initialize_instructions_by_server_id.clear() - self._upstream_initialize_instructions_probed_at.clear() - def invalidate_discovery_lists(self, server_id: str) -> None: self._upstream_initialize_instructions_by_server_id.pop(server_id, None) self._upstream_initialize_instructions_probed_at.pop(server_id, None) @@ -4667,6 +4674,12 @@ class MCPServerManager: ) self.invalidate_discovery_lists(server_id) + self._initialize_instructions_generations = MappingProxyType( + { + **self._initialize_instructions_generations, + server_id: self.initialize_instructions_generation(server_id) + 1, + } + ) self._drop_listed_tools(server_id) invalidate_oauth_metadata_cache(server_id) @@ -6532,6 +6545,7 @@ class MCPServerManager: GuardrailRaisedException: If guardrails block the call HTTPException: If an HTTP error occurs """ + instructions_generation: Final = self.initialize_instructions_generation(mcp_server.server_id) server_auth_header: Final = _server_auth_header_for(mcp_server, mcp_server_auth_headers, mcp_auth_header) # Extract subject token for OAuth2 Token Exchange (OBO) and ID-JAG flows @@ -6761,7 +6775,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: Final = 1 if proxy_logging_obj else 0 result: Final = mcp_responses[result_index] - self._remember_upstream_initialize_instructions(mcp_server, client) + self._remember_upstream_initialize_instructions(mcp_server, client, instructions_generation) return ( bind_target(result, mcp_server) @@ -7513,6 +7527,7 @@ class MCPServerManager: """ from datetime import datetime + instructions_generation: Final = self.initialize_instructions_generation(server_id) server: Final = self.get_mcp_server_by_id(server_id) if not server: verbose_logger.warning("MCP Server %s not found", server_id) @@ -7574,7 +7589,7 @@ class MCPServerManager: client.run_with_session(_noop, quiet_on_error=True), timeout=MCP_HEALTH_CHECK_TIMEOUT, ) - self._remember_upstream_initialize_instructions(server, client) + self._remember_upstream_initialize_instructions(server, client, instructions_generation) status = "healthy" except asyncio.TimeoutError: health_check_error = f"Health check timed out after {MCP_HEALTH_CHECK_TIMEOUT} seconds" diff --git a/tests/unit/experimental_mcp_client/test_mcp_client.py b/tests/unit/experimental_mcp_client/test_mcp_client.py index 60600afab84..aa8722f1e9c 100644 --- a/tests/unit/experimental_mcp_client/test_mcp_client.py +++ b/tests/unit/experimental_mcp_client/test_mcp_client.py @@ -699,8 +699,7 @@ class TestMCPClientInstructionsCapture: @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.""" + async def test_missing_instructions_are_recorded_as_empty(self, mock_session_cls: MagicMock) -> None: client = MCPClient( server_url="http://example.com/mcp", transport_type="http", @@ -723,7 +722,7 @@ class TestMCPClientInstructionsCapture: return "done" await client._execute_session_operation(transport_ctx, _op) - assert client._last_initialize_instructions is None + assert client._last_initialize_instructions == "" # --------------------------------------------------------------------------- diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py index ecd08e17cda..8e5483d2960 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -10461,29 +10461,82 @@ class TestMCPServerManagerUpstreamInstructionsCache: manager = MCPServerManager() 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.""" - 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._upstream_initialize_instructions_by_server_id.get("srv") == "hello" + @pytest.mark.parametrize( + ("answered", "cached"), + [(" hello \n", "hello"), ("", None), (" ", None), (None, "previous text")], + ids=["replaces_with_stripped_text", "drops_when_upstream_sent_none", "drops_blank", "keeps_without_initialize"], + ) + def test_remember_applies_the_last_completed_initialize(self, answered: str | None, cached: str | None) -> None: + manager: Final = MCPServerManager() + manager._upstream_initialize_instructions_by_server_id["srv"] = "previous text" + manager._remember_upstream_initialize_instructions( + MagicMock(server_id="srv"), + MagicMock(_last_initialize_instructions=answered), + manager.initialize_instructions_generation("srv"), + ) + assert manager._upstream_initialize_instructions_by_server_id.get("srv") == cached - 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._upstream_initialize_instructions_by_server_id.get("srv") is None + @pytest.mark.asyncio + @pytest.mark.parametrize("writer", ["prefetch", "health_check", "tools_list", "tool_call"]) + @pytest.mark.parametrize( + ("mid_flight", "cached"), + [("nothing", "upstream text"), ("discovery_refresh", "upstream text"), ("server_edit", None)], + ids=["unchanged", "discovery_refresh_mid_flight", "edited_mid_flight"], + ) + async def test_writer_skips_instructions_from_a_session_that_outlived_a_server_edit( + self, writer: str, mid_flight: str, cached: str | None + ) -> None: + manager: Final = MCPServerManager() + server: Final = MCPServer( + server_id="srv", + name="srv", + url="https://srv.example/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + ) + manager.registry = {server.server_id: server} - 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._upstream_initialize_instructions_by_server_id.get("srv") is None + def upstream(result: object) -> AsyncMock: + async def respond(*args: object, **kwargs: object) -> object: + match mid_flight: + case "server_edit": + manager.invalidate_server_definition_caches(server.server_id) + case "discovery_refresh": + manager.invalidate_discovery_lists(server.server_id) + return result + + return AsyncMock(side_effect=respond) + + manager.create_mcp_client = AsyncMock( + return_value=SimpleNamespace( + _last_initialize_instructions="upstream text", + run_with_session=upstream("ok"), + list_tools_page=upstream(ListToolsResult(tools=[])), + call_tool=upstream(CallToolResult(content=[], isError=False)), + ) + ) + + match writer: + case "prefetch": + await manager.ensure_upstream_initialize_instructions_cached(server) + case "health_check": + await manager.health_check_server(server.server_id) + case "tools_list": + await manager.get_tools_page(server, params=PaginatedRequestParams()) + case _: + await manager._call_regular_mcp_tool( + mcp_server=server, + original_tool_name="tool", + arguments={}, + tasks=[], + mcp_auth_header=None, + mcp_server_auth_headers=None, + oauth2_headers=None, + raw_headers=None, + proxy_logging_obj=None, + ) + + assert manager._upstream_initialize_instructions_by_server_id.get(server.server_id) == cached @pytest.mark.asyncio async def test_load_servers_from_config_clears_cache(self, config_only_mcp_manager_factory): @@ -18730,16 +18783,40 @@ async def test_catalog_failed_openapi_row_does_not_publish_partial_handlers(monk @pytest.mark.asyncio -async def test_catalog_unchanged_read_preserves_derived_initialize_instructions(monkeypatch): - read_rows = AsyncMock(return_value=[_catalog_row()]) +@pytest.mark.parametrize("full_reload", [False, True], ids=["request_path_refresh", "scheduled_reload"]) +async def test_catalog_refresh_drops_initialize_instructions_only_for_changed_servers( + monkeypatch: pytest.MonkeyPatch, full_reload: bool +) -> None: + def row(server_id: str, updated_day: int) -> LiteLLM_MCPServerTable: + return LiteLLM_MCPServerTable( + server_id=server_id, + server_name=server_id, + transport=MCPTransport.http, + url=f"https://{server_id}.example/mcp", + created_at=datetime(2026, 1, 1), + updated_at=datetime(2026, 1, updated_day), + ) + + read_rows: Final = AsyncMock(return_value=[row("unchanged", 1), row("edited", 1), row("deleted", 1)]) _catalog_database(monkeypatch, read_rows) - manager = MCPServerManager() + manager: Final = MCPServerManager() await manager.catalog.list() - manager._upstream_initialize_instructions_by_server_id = {"catalog-server": "cached instructions"} - manager._upstream_initialize_instructions_probed_at = {"catalog-server": 123.0} - await manager.catalog.list() - assert manager._upstream_initialize_instructions_by_server_id == {"catalog-server": "cached instructions"} - assert manager._upstream_initialize_instructions_probed_at == {"catalog-server": 123.0} + manager._upstream_initialize_instructions_by_server_id = { + "unchanged": "unchanged text", + "edited": "stale edited text", + "deleted": "stale deleted text", + "not-in-db": "config server text", + } + manager._upstream_initialize_instructions_probed_at = {"unchanged": 1.0, "edited": 2.0, "deleted": 3.0} + read_rows.return_value = [row("unchanged", 1), row("edited", 2)] + + await (manager.reload_servers_from_database() if full_reload else manager.catalog.list()) + + assert manager._upstream_initialize_instructions_by_server_id == { + "unchanged": "unchanged text", + "not-in-db": "config server text", + } + assert manager._upstream_initialize_instructions_probed_at == {"unchanged": 1.0} @pytest.mark.asyncio diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py index a873e2ab1ae..3f1740c0a18 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py @@ -6572,7 +6572,7 @@ class TestEnsureUpstreamInitializeInstructionsCached: 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 + fake_client._last_initialize_instructions = "" # upstream sent nothing create = AsyncMock(return_value=fake_client) with patch.object(global_mcp_server_manager, "create_mcp_client", create):