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>

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)

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

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"

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)

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

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