From 61244d356a6e72bcbb80179d16ac2c60e6d62fb1 Mon Sep 17 00:00:00 2001 From: mrinal Date: Sat, 10 Oct 2026 20:28:13 +0000 Subject: [PATCH] fix(mcp): constant store reads per build poll and cancel orphaned index builds Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/semantic_tool_filter.py | 11 ++-- .../mcp_server/semantic_tool_index_store.py | 5 +- .../proxy/hooks/mcp_semantic_filter/hook.py | 2 +- litellm/proxy/proxy_server.py | 59 ++++++++++++------- .../test_semantic_tool_index_store.py | 52 +++++++++++++++- .../proxy/proxy_server/test_proxy_config.py | 37 ++++++++++++ 6 files changed, 134 insertions(+), 32 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py b/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py index 16398f4d035..dbd39a5d077 100644 --- a/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py +++ b/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py @@ -138,13 +138,12 @@ class SemanticMCPToolFilter: ) @catalog_operation(global_manager) - async def build_router_from_mcp_registry(self, *, async_index: bool = False) -> None: + async def build_router_from_mcp_registry(self) -> None: """Build semantic router from all MCP tools in the registry (no auth checks). With a vector store the shared Redis path runs: acquire the build lock or wait for the holding pod's vectors, then embed only the misses. - Without one, the sync build blocks startup exactly as before; the - keyword arg lets a background task run the async (aadd) path instead. + Without one, the sync build blocks startup exactly as before. """ from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, @@ -174,7 +173,7 @@ class SemanticMCPToolFilter: return verbose_logger.info("Fetched %s tools from %s MCP servers", len(all_tools), len(registry)) - if self._vector_store is not None or async_index: + if self._vector_store is not None: await self._abuild_router(all_tools) else: self._build_router(all_tools) @@ -308,9 +307,9 @@ class SemanticMCPToolFilter: return False descriptions: Final = tuple(self.extract_tool_info(tool)[1] for tool in tools) - keys: Final = tuple(tool_vector_key(self.embedding_identity, text) for text in set(descriptions)) + keys_set: Final = frozenset(tool_vector_key(self.embedding_identity, text) for text in set(descriptions)) - while {key for key in keys if key not in await store.get_many(keys)}: + while keys_set - (await store.get_many(tuple(keys_set))).keys(): if await store.try_acquire_build_lock(): return True await asyncio.sleep(self._lock_poll_interval_s) diff --git a/litellm/proxy/_experimental/mcp_server/semantic_tool_index_store.py b/litellm/proxy/_experimental/mcp_server/semantic_tool_index_store.py index 9429ecb2989..78a9528b14e 100644 --- a/litellm/proxy/_experimental/mcp_server/semantic_tool_index_store.py +++ b/litellm/proxy/_experimental/mcp_server/semantic_tool_index_store.py @@ -79,10 +79,7 @@ class RedisToolVectorStore: verbose_logger.warning("MCP semantic tool index store put_many failed; vectors stay local: %s", e) async def try_acquire_build_lock(self) -> bool: - # PodLockManager returns False both when another pod holds the lock and - # when Redis errors internally, so a degraded store cannot tell them - # apart; err toward embedding locally rather than waiting on writes - # that a dead Redis can never deliver. + # PodLockManager returns False on Redis errors too; once degraded, embed locally instead of waiting. try: acquired: Final = await self._lock_manager.acquire_lock( _BUILD_LOCK_ID, ttl=_BUILD_LOCK_TTL_S, allow_reentrant=True diff --git a/litellm/proxy/hooks/mcp_semantic_filter/hook.py b/litellm/proxy/hooks/mcp_semantic_filter/hook.py index 8cedea69c6f..54f8d9f3ba6 100644 --- a/litellm/proxy/hooks/mcp_semantic_filter/hook.py +++ b/litellm/proxy/hooks/mcp_semantic_filter/hook.py @@ -52,7 +52,7 @@ def _make_vector_store(redis_cache: Optional["RedisCache"]) -> Optional["ToolVec async def _run_deferred_index_build(semantic_filter: "SemanticMCPToolFilter") -> None: """Background task for a shared (Redis-backed) index build; swallows failures so requests keep passing through.""" try: - await semantic_filter.build_router_from_mcp_registry(async_index=True) + await semantic_filter.build_router_from_mcp_registry() except Exception: # noqa: BLE001 # background task must not die; requests keep passing through verbose_proxy_logger.exception( "MCP semantic tool index deferred build failed; request-time indexing remains active" diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7d04b4f016f..5289c9ccea5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -8527,6 +8527,41 @@ class ProxyConfig: db_values: Final = self._prepared_db_settings_values("litellm_settings", raw_settings) self._apply_litellm_settings_db_values(db_values) + def _semantic_filter_settings_unchanged(self, mcp_semantic_filter_config: object, hook_type: type) -> bool: + """True when the stored settings match the in-memory ones and a live hook already covers them.""" + import litellm + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + + if not hasattr(self, "_last_semantic_filter_config"): + return False + if self._last_semantic_filter_config != mcp_semantic_filter_config: + return False + for active_hook in litellm.logging_callback_manager.get_custom_loggers_for_type(hook_type): + if not isinstance(active_hook, SemanticToolFilterHook): + continue + if active_hook.index_build_task is not None: + verbose_proxy_logger.debug( + "Semantic filter settings unchanged, index build already started; skipping reinitialization" + ) + return True + if active_hook.filter is not None and active_hook.filter.tool_router is not None: + verbose_proxy_logger.debug("Semantic filter settings unchanged, skipping reinitialization") + return True + verbose_proxy_logger.info( + "Semantic filter settings unchanged, but hook is missing or uninitialized. Reinitializing." + ) + return False + + @staticmethod + def _cancel_orphaned_semantic_filter_index_builds() -> None: + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + + for active_hook in litellm.logging_callback_manager.get_custom_loggers_for_type(SemanticToolFilterHook): + if not isinstance(active_hook, SemanticToolFilterHook): + continue + if (build_task := active_hook.index_build_task) is not None and not build_task.done(): + build_task.cancel() + async def init_semantic_filter_settings_in_db(self, prisma_client: PrismaClient) -> None: """ Initialize MCP semantic filter settings from database. @@ -8553,26 +8588,10 @@ class ProxyConfig: return # Check if settings have changed (compare with in-memory state) - if hasattr(self, "_last_semantic_filter_config"): - if self._last_semantic_filter_config == mcp_semantic_filter_config: - # If hook is missing or router isn't built yet, reinitialize anyway - active_hooks = litellm.logging_callback_manager.get_custom_loggers_for_type(SemanticToolFilterHook) - if active_hooks: - for active_hook in active_hooks: - if isinstance(active_hook, SemanticToolFilterHook): - if active_hook.index_build_task is not None: - verbose_proxy_logger.debug( - "Semantic filter settings unchanged, index build already started; skipping reinitialization" - ) - return - if active_hook.filter is not None and active_hook.filter.tool_router is not None: - verbose_proxy_logger.debug( - "Semantic filter settings unchanged, skipping reinitialization" - ) - return - verbose_proxy_logger.info( - "Semantic filter settings unchanged, but hook is missing or uninitialized. Reinitializing." - ) + if self._semantic_filter_settings_unchanged(mcp_semantic_filter_config, SemanticToolFilterHook): + return + + self._cancel_orphaned_semantic_filter_index_builds() # Remove old hooks using logging callback manager litellm.logging_callback_manager.remove_callbacks_by_type(litellm.callbacks, SemanticToolFilterHook) diff --git a/tests/unit/proxy/_experimental/mcp_server/test_semantic_tool_index_store.py b/tests/unit/proxy/_experimental/mcp_server/test_semantic_tool_index_store.py index 0530de440ad..caf10d469e1 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_semantic_tool_index_store.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_semantic_tool_index_store.py @@ -6,8 +6,9 @@ cross-pod build lock, and the CachingToolEncoder serving document embeddings from the store instead of re-embedding. """ +import asyncio import sys -from unittest.mock import Mock +from unittest.mock import AsyncMock, Mock, patch import pytest @@ -57,6 +58,9 @@ class _InMemoryToolVectorStore: self.release_calls = 0 self.lock_result = True self._fill_on_call = None + self.block_put_many = False + self.put_started = asyncio.Event() + self.put_block = asyncio.Event() def fill_on_get_many_call(self, call_number): self._fill_on_call = call_number @@ -71,6 +75,9 @@ class _InMemoryToolVectorStore: self._pending_fill = dict(keys_vectors) async def put_many(self, vectors): + if self.block_put_many: + self.put_started.set() + await self.put_block.wait() self.vectors.update({k: tuple(v) for k, v in vectors.items()}) async def try_acquire_build_lock(self): @@ -287,3 +294,46 @@ async def test_lock_held_by_other_pod_waits_for_shared_vectors(): assert store.acquire_calls >= 1 assert store.get_many_calls >= 3 assert recorded == [["test"]] + + +@requires_semantic_router +@pytest.mark.asyncio +async def test_full_build_reads_the_store_a_constant_number_of_times(): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + tools = [_tool(f"srv-a_get_tool_{i}", f"do thing {i}") for i in range(20)] + store = _InMemoryToolVectorStore() + filter_instance = _make_filter([], store=store) + + with ( + patch.object(global_mcp_server_manager, "get_registry", return_value={"srv-a": object()}), + patch.object( + global_mcp_server_manager, "get_tools_for_server", new=AsyncMock(return_value=tools) + ), + ): + await filter_instance.build_router_from_mcp_registry() + + assert filter_instance.tool_router is not None + assert store.get_many_calls <= 3, ( + f"{store.get_many_calls} get_many calls for {len(tools)} tools; reads must not scale with the tool count" + ) + + +@requires_semantic_router +@pytest.mark.asyncio +async def test_cancelling_a_build_while_holding_the_lock_releases_it(): + tools = [_tool("srv-a_get_ticket", "get a ticket"), _tool("srv-b_get_weather", "weather forecast")] + store = _InMemoryToolVectorStore() + store.block_put_many = True + filter_instance = _make_filter([], store=store) + + build_task = asyncio.create_task(filter_instance._abuild_router(tools)) + await asyncio.wait_for(store.put_started.wait(), timeout=5) + assert store.acquire_calls == 1 + + build_task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(build_task, timeout=5) + assert store.release_calls == 1 diff --git a/tests/unit/proxy/proxy_server/test_proxy_config.py b/tests/unit/proxy/proxy_server/test_proxy_config.py index 66b2f2787ac..069f2e1bf2b 100644 --- a/tests/unit/proxy/proxy_server/test_proxy_config.py +++ b/tests/unit/proxy/proxy_server/test_proxy_config.py @@ -5640,3 +5640,40 @@ async def test_proxy_config_semantic_filter_db_sync_skips_reinit_while_index_bui finally: gate.set() await build_task + + +@pytest.mark.asyncio +async def test_proxy_config_semantic_filter_db_sync_cancels_orphaned_index_build(monkeypatch): + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + + settings: Final = { + "enabled": True, + "embedding_model": "embed", + "top_k": 3, + "similarity_threshold": 0.5, + } + changed_settings: Final = {**settings, "top_k": 5} + gate: Final = asyncio.Event() + build_task: Final = asyncio.create_task(gate.wait()) + hook: Final = MagicMock(spec=SemanticToolFilterHook) + hook.index_build_task = build_task + hook.filter = SimpleNamespace(tool_router=None) + removed: Final = MagicMock() + monkeypatch.setattr( + "litellm.proxy.proxy_server.get_config_param", + AsyncMock(return_value=SimpleNamespace(param_value={"mcp_semantic_tool_filter": changed_settings})), + ) + monkeypatch.setattr(litellm.logging_callback_manager, "get_custom_loggers_for_type", lambda _: [hook]) + monkeypatch.setattr(litellm.logging_callback_manager, "remove_callbacks_by_type", removed) + pc: Final = ProxyConfig() + pc._last_semantic_filter_config = settings + try: + await pc._init_semantic_filter_settings_in_db(prisma_client=MagicMock()) + + await asyncio.sleep(0) + assert build_task.cancelled() + removed.assert_called_once() + finally: + gate.set() + if not build_task.done(): + build_task.cancel()