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):