From 8d70a267eec8f93eb9837c84b8314e2d3569a6cc Mon Sep 17 00:00:00 2001 From: mrinal Date: Sat, 10 Oct 2026 20:04:32 +0000 Subject: [PATCH 1/9] perf(mcp): share semantic tool embeddings across pods via redis Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/_experimental/mcp_server/AGENTS.md | 1 + .../mcp_server/semantic_tool_filter.py | 251 +++++++++++---- .../mcp_server/semantic_tool_index_store.py | 172 +++++++++++ .../proxy/hooks/mcp_semantic_filter/hook.py | 64 +++- litellm/proxy/proxy_server.py | 7 + .../mcp_server/test_semantic_tool_filter.py | 106 +++++++ .../test_semantic_tool_index_store.py | 289 ++++++++++++++++++ .../proxy/proxy_server/test_proxy_config.py | 33 ++ 8 files changed, 848 insertions(+), 75 deletions(-) create mode 100644 litellm/proxy/_experimental/mcp_server/semantic_tool_index_store.py create mode 100644 tests/unit/proxy/_experimental/mcp_server/test_semantic_tool_index_store.py diff --git a/litellm/proxy/_experimental/mcp_server/AGENTS.md b/litellm/proxy/_experimental/mcp_server/AGENTS.md index 8c53216f77a..809a2086a06 100644 --- a/litellm/proxy/_experimental/mcp_server/AGENTS.md +++ b/litellm/proxy/_experimental/mcp_server/AGENTS.md @@ -41,6 +41,7 @@ litellm/proxy/_experimental/mcp_server/ sampling_handler.py # MCP sampling to LiteLLM completion flow elicitation_handler.py # MCP elicitation relay flow semantic_tool_filter.py # semantic filtering of available MCP tools + semantic_tool_index_store.py # Redis-shared embeddings + build lock for the semantic tool index tool_search.py # opt-in virtual tools (mcp_tool_search + mcp_tool_call) for large catalogs guardrail_translation/ handler.py # MCP guardrail result translation diff --git a/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py b/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py index 2f88d8c361f..16398f4d035 100644 --- a/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py +++ b/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py @@ -8,17 +8,26 @@ import asyncio from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final +from mcp.types import Tool as MCPTool +from pydantic import TypeAdapter + from litellm._logging import verbose_logger from litellm.exceptions import ContextWindowExceededError from litellm.litellm_core_utils.exception_mapping_utils import ExceptionCheckers from litellm.proxy._experimental.mcp_server.catalog import catalog_operation, global_manager from litellm.proxy._experimental.mcp_server.faults import iter_exception_tree +from litellm.proxy._experimental.mcp_server.semantic_tool_index_store import ( + CachingToolEncoder, + ToolVectorStore, + tool_vector_key, +) from litellm.proxy._experimental.mcp_server.utils import MCP_TOOL_PREFIX_SEPARATOR if TYPE_CHECKING: from semantic_router.routers import SemanticRouter from litellm.router import Router + from litellm.router_strategy.auto_router.litellm_encoder import LiteLLMRouterEncoder class SemanticToolFilterContextWindowError(Exception): @@ -58,6 +67,8 @@ class SemanticMCPToolFilter: top_k: int = 10, similarity_threshold: float = 0.3, enabled: bool = True, + vector_store: ToolVectorStore | None = None, + lock_poll_interval_s: float = 5.0, ): """ Initialize the semantic tool filter. @@ -68,6 +79,11 @@ class SemanticMCPToolFilter: top_k: Maximum number of tools to return similarity_threshold: Minimum similarity score for filtering enabled: Whether filtering is enabled + vector_store: Shared embedding store; when set the index builds + cross-pod in the background and requests pass through unfiltered + until startup_index_ready flips + lock_poll_interval_s: How often to re-check the store while another + pod holds the build lock """ self.enabled = enabled self.top_k = top_k @@ -78,10 +94,58 @@ class SemanticMCPToolFilter: self.context_window_error: str | None = None self._tool_map: dict[str, object] = {} # MCPTool objects or OpenAI function dicts self._index_sync_lock = asyncio.Lock() + self._vector_store = vector_store + self._lock_poll_interval_s = lock_poll_interval_s + self.embedding_identity = self._compute_embedding_identity() + self.startup_index_ready = vector_store is None + + def _compute_embedding_identity(self) -> str: + """embedding_model plus every deployment model it routes to, so repointing an alias invalidates the shared vectors.""" + try: + deployments: Final = self.router_instance.get_model_list(model_name=self.embedding_model) or () + deployment_maps: Final = TypeAdapter(tuple[Mapping[str, object], ...]).validate_python(deployments) + litellm_params: Final = tuple( + TypeAdapter(Mapping[str, object]).validate_python(deployment.get("litellm_params") or {}) + for deployment in deployment_maps + ) + deployment_models: Final = sorted( + {str(model) for model in (p.get("model") for p in litellm_params) if model} + ) + except (AttributeError, TypeError, KeyError, ValueError): + return self.embedding_model + if not deployment_models: + return self.embedding_model + return self.embedding_model + "|" + "|".join(deployment_models) + + def _new_encoder(self) -> "LiteLLMRouterEncoder": + """Encoder used for every index build; cached in Redis when a store is set.""" + from litellm.router_strategy.auto_router.litellm_encoder import ( + LiteLLMRouterEncoder, + ) + + if self._vector_store is not None: + return CachingToolEncoder( + litellm_router_instance=self.router_instance, + model_name=self.embedding_model, + score_threshold=self.similarity_threshold, + store=self._vector_store, + embedding_identity=self.embedding_identity, + ) + return LiteLLMRouterEncoder( + litellm_router_instance=self.router_instance, + model_name=self.embedding_model, + score_threshold=self.similarity_threshold, + ) @catalog_operation(global_manager) - async def build_router_from_mcp_registry(self) -> None: - """Build semantic router from all MCP tools in the registry (no auth checks).""" + async def build_router_from_mcp_registry(self, *, async_index: bool = False) -> 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. + """ from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) @@ -95,7 +159,7 @@ class SemanticMCPToolFilter: return # Fetch tools from all servers in parallel - all_tools: Final = [] + all_tools: Final[Sequence[object]] = [] for server_id, server in registry.items(): try: tools = await global_mcp_server_manager.get_tools_for_server(server_id) @@ -110,39 +174,66 @@ class SemanticMCPToolFilter: return verbose_logger.info("Fetched %s tools from %s MCP servers", len(all_tools), len(registry)) - self._build_router(all_tools) + if self._vector_store is not None or async_index: + await self._abuild_router(all_tools) + else: + self._build_router(all_tools) except Exception as e: verbose_logger.error("Failed to build router from MCP registry: %s", e) self.tool_router = None raise + finally: + self.startup_index_ready = True - def extract_tool_info(self, tool) -> tuple[str, str]: + def extract_tool_info(self, tool: object) -> tuple[str, str]: """Extract name and description from MCP tool or OpenAI function dict.""" name: str description: str if isinstance(tool, dict): # OpenAI function format - name = tool.get("name", "") - description = tool.get("description", name) + tool_dict: Final = TypeAdapter(Mapping[str, object]).validate_python(tool) + name = str(tool_dict.get("name", "")) + description = str(tool_dict.get("description", name)) else: # MCPTool object - name = str(tool.name) - description = str(tool.description) if tool.description else str(tool.name) + mcp_tool: Final = MCPTool.model_validate(tool, from_attributes=True) + name = str(mcp_tool.name) + description = str(mcp_tool.description) if mcp_tool.description else str(mcp_tool.name) return name, description _extract_tool_info = extract_tool_info - def _build_router(self, tools: list) -> None: - """Build semantic router with tools (MCPTool objects or OpenAI function dicts).""" - from semantic_router.routers import SemanticRouter + def _tools_to_routes_and_map(self, tools: Sequence[object]) -> tuple[Sequence[object], Mapping[str, object]]: from semantic_router.routers.base import Route - from litellm.router_strategy.auto_router.litellm_encoder import ( - LiteLLMRouterEncoder, - ) + extracted: Final = tuple((self.extract_tool_info(tool), tool) for tool in tools) + tool_map: Final = {name: tool for (name, _), tool in extracted} + routes: Final[Sequence[object]] = [ + Route( + name=name, + description=description, + utterances=[description], + score_threshold=self.similarity_threshold, + ) + for (name, description), _ in extracted + ] + return routes, tool_map + + def _handle_build_error(self, error: Exception) -> bool: + """True when the failure was recorded as a context-window overflow; False when the caller should re-raise.""" + verbose_logger.error("Failed to build semantic router: %s", error) + self.tool_router = None + if _is_context_window_error(error): + self.context_window_error = str(error) + return True + return False + + def _build_router(self, tools: Sequence[object]) -> None: + """Build semantic router with tools (MCPTool objects or OpenAI function dicts).""" + from semantic_router.routers import SemanticRouter if not tools: self.tool_router = None @@ -150,53 +241,95 @@ class SemanticMCPToolFilter: try: self.context_window_error = None - # Convert tools to routes - routes: Final = [] - self._tool_map = {} - - for tool in tools: - name, description = self.extract_tool_info(tool) - self._tool_map[name] = tool - - routes.append( - Route( - name=name, - description=description, - utterances=[description], - score_threshold=self.similarity_threshold, - ) - ) - + routes, tool_map = self._tools_to_routes_and_map(tools) self.tool_router = SemanticRouter( - routes=routes, - encoder=LiteLLMRouterEncoder( - litellm_router_instance=self.router_instance, - model_name=self.embedding_model, - score_threshold=self.similarity_threshold, - ), + routes=list(routes), + encoder=self._new_encoder(), auto_sync="local", ) + self._tool_map = dict(tool_map) verbose_logger.info("Built semantic router with %s tools", len(routes)) except Exception as e: - verbose_logger.error("Failed to build semantic router: %s", e) - self.tool_router = None - if _is_context_window_error(e): - self.context_window_error = str(e) + if self._handle_build_error(e): return raise + async def _abuild_router(self, tools: Sequence[object]) -> None: + """ + Async index build for the background path. + + When a store is set, this first waits on the cross-pod build lock so + only one pod embeds; the caching encoder then fills any vectors the + holder did not write. + """ + from semantic_router.routers import SemanticRouter + + if not tools: + self.tool_router = None + return + + lock_held: Final = await self._await_store_or_lock(tools) if self._vector_store is not None else False + try: + self.context_window_error = None + routes, tool_map = self._tools_to_routes_and_map(tools) + encoder: Final = self._new_encoder() + router: Final = SemanticRouter( # pyright: ignore[reportUnknownVariableType] # semantic_router ships no type stubs + routes=[], + encoder=encoder, + init_async_index=True, + top_k=self.top_k, + ) + + # aadd's lazy init probes index dimensions through the sync encoder, + # so dimensions are seeded with an async probe to keep the build async. + dims_probe: Final = await encoder.aencode_queries(["test"]) # pyright: ignore[reportUnknownMemberType] # semantic_router ships no type stubs + router.index.dimensions = len(dims_probe[0]) # pyright: ignore[reportUnknownMemberType] # semantic_router ships no type stubs + await router.aadd(list(routes)) # pyright: ignore[reportUnknownMemberType] # semantic_router ships no type stubs + + self._tool_map = dict(tool_map) + self.tool_router = router + + verbose_logger.info("Built semantic router with %s tools via async index build", len(routes)) + + except Exception as e: + if self._handle_build_error(e): + return + raise + finally: + if lock_held and self._vector_store is not None: + await self._vector_store.release_build_lock() + + async def _await_store_or_lock(self, tools: Sequence[object]) -> bool: + """True when this pod holds the build lock after the wait; False when every vector arrived without it.""" + store: Final = self._vector_store + if store is None: + 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)) + + while {key for key in keys if key not in await store.get_many(keys)}: + if await store.try_acquire_build_lock(): + return True + await asyncio.sleep(self._lock_poll_interval_s) + return False + def _has_tools_missing_from_index(self, tools: Sequence[object]) -> bool: """Allocation-free check for any named tool not yet in the semantic index.""" - return any(name and name not in self._tool_map for name in (self.extract_tool_info(t)[0] for t in tools)) + return any( + name and (name not in self._tool_map or self.extract_tool_info(self._tool_map[name])[1] != description) + for name, description in (self.extract_tool_info(t) for t in tools) + ) def _tools_missing_from_index(self, tools: Sequence[object]) -> Mapping[str, object]: """Map name -> tool for every named tool not yet in the semantic index.""" + infos: Final = ((self.extract_tool_info(t), t) for t in tools) return { name: tool - for name, tool in ((self.extract_tool_info(t)[0], t) for t in tools) - if name and name not in self._tool_map + for (name, description), tool in infos + if name and (name not in self._tool_map or self.extract_tool_info(self._tool_map[name])[1] != description) } async def _ensure_tools_indexed(self, available_tools: Sequence[object]) -> None: @@ -216,11 +349,6 @@ class SemanticMCPToolFilter: description cannot poison the filter for other users on the worker. """ from semantic_router.routers import SemanticRouter - from semantic_router.routers.base import Route - - from litellm.router_strategy.auto_router.litellm_encoder import ( - LiteLLMRouterEncoder, - ) if not self._has_tools_missing_from_index(available_tools): return @@ -230,32 +358,21 @@ class SemanticMCPToolFilter: if not missing: return - descriptions: Final = {name: self.extract_tool_info(tool)[1] for name, tool in missing.items()} - routes: Final = [ - Route( - name=name, - description=description, - utterances=[description], - score_threshold=self.similarity_threshold, - ) - for name, description in descriptions.items() - ] + routes, _ = self._tools_to_routes_and_map(tuple(missing.values())) if self.tool_router is None: router: Final = SemanticRouter( routes=[], - encoder=LiteLLMRouterEncoder( - litellm_router_instance=self.router_instance, - model_name=self.embedding_model, - score_threshold=self.similarity_threshold, - ), + encoder=self._new_encoder(), auto_sync="local", top_k=self.top_k, ) - await router.aadd(routes) + await router.aadd(list(routes)) self.tool_router = router else: - await self.tool_router.aadd(routes) + for name in (name for name in missing if name in self._tool_map): + await self.tool_router.adelete(name) # pyright: ignore[reportUnknownMemberType] # semantic_router ships no type stubs + await self.tool_router.aadd(list(routes)) self._tool_map.update(missing) verbose_logger.info( diff --git a/litellm/proxy/_experimental/mcp_server/semantic_tool_index_store.py b/litellm/proxy/_experimental/mcp_server/semantic_tool_index_store.py new file mode 100644 index 00000000000..9429ecb2989 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/semantic_tool_index_store.py @@ -0,0 +1,172 @@ +""" +Redis-shared vector store for the MCP semantic tool index. + +Lets every pod reuse embeddings another pod already computed instead of +re-embedding every tool description at startup. +""" + +import base64 +import sys +from array import array +from collections.abc import Mapping, Sequence +from hashlib import sha256 +from typing import TYPE_CHECKING, Final, Protocol + +from pydantic import TypeAdapter + +from litellm._logging import verbose_logger +from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager +from litellm.router_strategy.auto_router.litellm_encoder import LiteLLMRouterEncoder + +if TYPE_CHECKING: + from litellm.caching.redis_cache import RedisCache + from litellm.router import Router + +_KEY_PREFIX: Final = "litellm:mcp_semantic_tool_index:v1:" +_BUILD_LOCK_ID: Final = "mcp_semantic_tool_index_build" +_BUILD_LOCK_TTL_S: Final = 300 +_VECTOR_TTL_S: Final = 30 * 24 * 60 * 60 +_PUT_CHUNK_SIZE: Final = 50 + + +def tool_vector_key(embedding_identity: str, text: str) -> str: + """Cache key for one embedded description; identity changes invalidate every entry.""" + return _KEY_PREFIX + sha256((embedding_identity + "\0" + text).encode("utf-8")).hexdigest() + + +class ToolVectorStore(Protocol): + """Shared embedding cache plus a cross-pod build lock.""" + + async def get_many(self, keys: Sequence[str]) -> Mapping[str, tuple[float, ...]]: + """Return only the keys that hit.""" + ... + + async def put_many(self, vectors: Mapping[str, Sequence[float]]) -> None: ... + + async def try_acquire_build_lock(self) -> bool: ... + + async def release_build_lock(self) -> None: ... + + +class RedisToolVectorStore: + """ToolVectorStore over RedisCache; every Redis failure degrades to a local embed.""" + + def __init__(self, redis_cache: "RedisCache", lock_manager: PodLockManager) -> None: + self._redis_cache = redis_cache + self._lock_manager = lock_manager + self._degraded = False + + async def get_many(self, keys: Sequence[str]) -> Mapping[str, tuple[float, ...]]: + try: + raw: Final = TypeAdapter(Mapping[str, object]).validate_python( + await self._redis_cache.async_batch_get_cache(list(keys)) # pyright: ignore[reportUnknownMemberType] # RedisCache declares a partially untyped signature + ) + except Exception as e: # noqa: BLE001 # any redis failure must degrade to a local embed, never raise + self._degraded = True + verbose_logger.warning("MCP semantic tool index store get_many failed; embedding locally: %s", e) + return {} + return {key: vector for key in keys if (vector := _decode_vector(raw.get(key))) is not None} + + async def put_many(self, vectors: Mapping[str, Sequence[float]]) -> None: + try: + entries: Final = tuple((key, _encode_vector(vector)) for key, vector in vectors.items()) + for start in range(0, len(entries), _PUT_CHUNK_SIZE): + await self._redis_cache.async_set_cache_pipeline( # pyright: ignore[reportUnknownMemberType] # RedisCache declares a partially untyped signature + entries[start : start + _PUT_CHUNK_SIZE], ttl=_VECTOR_TTL_S + ) + except Exception as e: # noqa: BLE001 # any redis failure must degrade to a local embed, never raise + self._degraded = True + 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. + try: + acquired: Final = await self._lock_manager.acquire_lock( + _BUILD_LOCK_ID, ttl=_BUILD_LOCK_TTL_S, allow_reentrant=True + ) + return bool(acquired) or self._degraded + except Exception as e: # noqa: BLE001 # any redis failure must degrade to a local embed, never raise + self._degraded = True + verbose_logger.warning("MCP semantic tool index store lock acquire failed; embedding locally: %s", e) + return True + + async def release_build_lock(self) -> None: + try: + await self._lock_manager.release_lock(_BUILD_LOCK_ID) + except Exception as e: # noqa: BLE001 # any redis failure must degrade to a local embed, never raise + verbose_logger.warning("MCP semantic tool index store lock release failed: %s", e) + + +def _encode_vector(vector: Sequence[float]) -> str: + packed: Final = array("f", vector) + if sys.byteorder == "big": + packed.byteswap() + return base64.b64encode(packed.tobytes()).decode("ascii") + + +def _decode_vector(value: object) -> tuple[float, ...] | None: + try: + if not isinstance(value, (str, bytes, bytearray)): + return None + raw: Final = base64.b64decode(value) + vector: Final = array("f") + vector.frombytes(raw) + if sys.byteorder == "big": + vector.byteswap() + return tuple(vector) + except Exception: # noqa: BLE001 # corrupt payloads are treated as a plain cache miss + return None + + +class CachingToolEncoder(LiteLLMRouterEncoder): + """LiteLLMRouterEncoder that serves document embeddings from a ToolVectorStore. + + Queries are never cached; only add/aadd's document utterances go through + the store. Duplicate texts share one key and one embed. + """ + + _tool_vector_store: ToolVectorStore + _embedding_identity: str + + def __init__( + self, + *, + store: ToolVectorStore, + embedding_identity: str, + litellm_router_instance: "Router", + model_name: str, + score_threshold: float | None = None, + ) -> None: + super().__init__( + litellm_router_instance=litellm_router_instance, + model_name=model_name, + score_threshold=score_threshold, + ) + self._tool_vector_store = store + self._embedding_identity = embedding_identity + + def encode_documents( # pyright: ignore[reportIncompatibleMethodOverride] # semantic_router calls encode_documents without kwargs + self, docs: Sequence[str] + ) -> list[list[float]]: # mutable-ok: mirrors the parent DenseEncoder contract + return super().encode_documents(list(docs)) + + async def aencode_documents( # pyright: ignore[reportIncompatibleMethodOverride] # semantic_router calls aencode_documents without kwargs + self, docs: Sequence[str] + ) -> list[list[float]]: # mutable-ok: mirrors the parent DenseEncoder contract + unique_docs: Final = tuple(dict.fromkeys(docs)) + keys: Final = tuple(tool_vector_key(self._embedding_identity, text) for text in unique_docs) + hits: Final = await self._tool_vector_store.get_many(keys) + + miss_pairs: Final = tuple((text, key) for text, key in zip(unique_docs, keys) if key not in hits) + embedded: Final = dict(hits) + if miss_pairs: + vectors: Final = await super().aencode_documents([text for text, _ in miss_pairs]) + new_vectors: Final = {key: tuple(vector) for (_, key), vector in zip(miss_pairs, vectors)} + await self._tool_vector_store.put_many(new_vectors) + embedded.update(new_vectors) + + vectors_by_text: Final = dict(zip(unique_docs, (embedded[key] for key in keys))) + return [list(vectors_by_text[text]) for text in docs] diff --git a/litellm/proxy/hooks/mcp_semantic_filter/hook.py b/litellm/proxy/hooks/mcp_semantic_filter/hook.py index d33811d3dea..8cedea69c6f 100644 --- a/litellm/proxy/hooks/mcp_semantic_filter/hook.py +++ b/litellm/proxy/hooks/mcp_semantic_filter/hook.py @@ -5,6 +5,8 @@ Pre-call hook that filters MCP tools semantically before LLM inference. Reduces context window size and improves tool selection accuracy. """ +import asyncio +import importlib.util from collections.abc import Awaitable, Callable, Collection, Iterable, Mapping, Sequence from typing import TYPE_CHECKING, Final, Optional @@ -24,13 +26,39 @@ from litellm.proxy._experimental.mcp_server.semantic_tool_filter import ( if TYPE_CHECKING: from litellm.caching.caching import DualCache + from litellm.caching.redis_cache import RedisCache from litellm.proxy._experimental.mcp_server.semantic_tool_filter import ( SemanticMCPToolFilter, ) + from litellm.proxy._experimental.mcp_server.semantic_tool_index_store import ToolVectorStore from litellm.proxy._types import UserAPIKeyAuth from litellm.router import Router +def _make_vector_store(redis_cache: Optional["RedisCache"]) -> Optional["ToolVectorStore"]: + """Redis-backed shared vector store, or None when no Redis is configured.""" + if redis_cache is None: + return None + from litellm.proxy._experimental.mcp_server.semantic_tool_index_store import ( + RedisToolVectorStore, + ) + from litellm.proxy.db.db_transaction_queue.pod_lock_manager import ( + PodLockManager, + ) + + return RedisToolVectorStore(redis_cache, PodLockManager(redis_cache)) + + +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) + 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" + ) + + class SemanticToolFilterConfig(TypedDict, total=False): enabled: ReadOnly[bool] embedding_model: ReadOnly[str] @@ -69,6 +97,7 @@ class SemanticToolFilterHook(CustomLogger): """ super().__init__() self.filter = semantic_filter + self.index_build_task: asyncio.Task[None] | None = None verbose_proxy_logger.debug( "Initialized SemanticToolFilterHook with filter: enabled=%s, top_k=%s", @@ -314,8 +343,11 @@ class SemanticToolFilterHook(CustomLogger): if self._should_expand_mcp_tools(tools): verbose_proxy_logger.debug("Detected litellm_proxy MCP references, expanding before semantic filtering") - if not self.filter.enabled: - verbose_proxy_logger.debug("Semantic filter disabled, leaving MCP references untouched") + if not self.filter.enabled or not self.filter.startup_index_ready: + verbose_proxy_logger.debug( + "Semantic filter %s, leaving MCP references untouched", + "disabled" if not self.filter.enabled else "index still building", + ) return None try: @@ -360,8 +392,11 @@ class SemanticToolFilterHook(CustomLogger): verbose_proxy_logger.debug("No messages in request, skipping semantic filter") return None - if not self.filter.enabled: - verbose_proxy_logger.debug("Semantic filter disabled, skipping") + if not self.filter.enabled or not self.filter.startup_index_ready: + verbose_proxy_logger.debug( + "Semantic filter %s, skipping", + "disabled" if not self.filter.enabled else "index still building", + ) return None try: @@ -476,6 +511,7 @@ class SemanticToolFilterHook(CustomLogger): async def initialize_from_config( config: SemanticToolFilterConfig | None, llm_router: Optional["Router"], + redis_cache: Optional["RedisCache"] = None, ) -> Optional["SemanticToolFilterHook"]: """ Initialize semantic tool filter from proxy config. @@ -483,6 +519,9 @@ class SemanticToolFilterHook(CustomLogger): Args: config: Proxy configuration dict (litellm_settings.mcp_semantic_tool_filter) llm_router: LiteLLM router instance for embeddings + redis_cache: Usage-tracking Redis cache; when configured the tool + index is built in the background with embeddings shared + cross-pod through it Returns: SemanticToolFilterHook instance if enabled, None otherwise @@ -500,6 +539,9 @@ class SemanticToolFilterHook(CustomLogger): return None try: + if redis_cache is not None and importlib.util.find_spec("semantic_router") is None: + raise ImportError("No module named 'semantic_router'") + embedding_model: Final = config.get("embedding_model", DEFAULT_MCP_SEMANTIC_FILTER_EMBEDDING_MODEL) top_k: Final = config.get("top_k", DEFAULT_MCP_SEMANTIC_FILTER_TOP_K) similarity_threshold = config.get("similarity_threshold", DEFAULT_MCP_SEMANTIC_FILTER_SIMILARITY_THRESHOLD) @@ -510,18 +552,24 @@ class SemanticToolFilterHook(CustomLogger): top_k=top_k, similarity_threshold=similarity_threshold, enabled=True, + vector_store=_make_vector_store(redis_cache), ) - # Build router from MCP registry on startup - await semantic_filter.build_router_from_mcp_registry() - hook: Final = SemanticToolFilterHook(semantic_filter) + if redis_cache is not None: + # Index builds in the background; requests pass through unfiltered until it finishes. + hook.index_build_task = asyncio.create_task(_run_deferred_index_build(semantic_filter)) + else: + # Build router from MCP registry on startup + await semantic_filter.build_router_from_mcp_registry() + verbose_proxy_logger.info( - "✅ MCP Semantic Tool Filter enabled: embedding_model=%s, top_k=%s, similarity_threshold=%s", + "✅ MCP Semantic Tool Filter enabled: embedding_model=%s, top_k=%s, similarity_threshold=%s, shared_index=%s", embedding_model, top_k, similarity_threshold, + redis_cache is not None, ) return hook diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index dac46855fd6..7d04b4f016f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -8560,6 +8560,11 @@ class ProxyConfig: 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" @@ -8578,6 +8583,7 @@ class ProxyConfig: hook: Final = await SemanticToolFilterHook.initialize_from_config( config=mcp_semantic_filter_config, llm_router=llm_router, + redis_cache=redis_usage_cache, ) if hook: litellm.logging_callback_manager.add_litellm_callback(hook) @@ -10596,6 +10602,7 @@ class ProxyStartupEvent: hook: Final = await SemanticToolFilterHook.initialize_from_config( config=mcp_semantic_filter_config, llm_router=llm_router, + redis_cache=redis_usage_cache, ) if hook: diff --git a/tests/unit/proxy/_experimental/mcp_server/test_semantic_tool_filter.py b/tests/unit/proxy/_experimental/mcp_server/test_semantic_tool_filter.py index f0b4e94f72f..1ba1b33b094 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_semantic_tool_filter.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_semantic_tool_filter.py @@ -2239,3 +2239,109 @@ async def test_semantic_filter_hook_narrows_only_references_the_gateway_serves() ) assert narrowed == [{**gateway_reference, "allowed_tools": ["srv-tool_1"]}, external_tool] + + +@requires_semantic_router +@pytest.mark.asyncio +async def test_semantic_filter_reindexes_tool_when_description_changes(): + """ + Regression test: a tool that keeps its name but gets a new description + must be re-embedded, otherwise the index keeps serving the stale vector. + """ + recorded_inputs = [] + filter_instance = _make_keyword_filter(recorded_inputs) + original = MCPTool( + name="srv-a-get_thing", + description="get the weather for a city", + inputSchema={"type": "object"}, + ) + edited = MCPTool( + name="srv-a-get_thing", + description="get a Linear ticket by id", + inputSchema={"type": "object"}, + ) + filter_instance._build_router([original]) + recorded_inputs.clear() + + filtered = await filter_instance.filter_tools( + query="linear ticket", + available_tools=[edited], + ) + + assert [edited] == filtered + assert any("get a Linear ticket by id" in text for batch in recorded_inputs for text in batch) + indexed_tool = filter_instance._tool_map["srv-a-get_thing"] + assert filter_instance.extract_tool_info(indexed_tool)[1] == "get a Linear ticket by id" + + +@requires_semantic_router +@pytest.mark.asyncio +async def test_hook_returns_before_shared_index_build_finishes(): + """ + With Redis configured the hook registers immediately and the shared + index build runs as a background task that finishes once the store + responds; without Redis the build stays inline. + """ + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + + gate = asyncio.Event() + + class _GatedRedis: + def __init__(self): + self.backing = {} + + async def async_batch_get_cache(self, key_list, parent_otel_span=None): + await gate.wait() + return {key: self.backing[key] for key in key_list if key in self.backing} + + async def async_set_cache(self, key, value, **kwargs): + self.backing[key] = value + return True + + async def async_get_cache(self, key): + return self.backing.get(key) + + async def async_set_cache_pipeline(self, cache_list, ttl=None, **kwargs): + self.backing.update(dict(cache_list)) + + recorded_inputs = [] + tool = MCPTool( + name="srv-a-get_thing", + description="get a Linear ticket by id", + inputSchema={"type": "object"}, + ) + 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=[tool]) + ), + ): + hook = await SemanticToolFilterHook.initialize_from_config( + config={"enabled": True}, + llm_router=_make_keyword_embedding_router(recorded_inputs), + redis_cache=_GatedRedis(), + ) + + assert hook is not None + assert hook.index_build_task is not None + await asyncio.sleep(0) + assert not hook.index_build_task.done() + assert not hook.filter.startup_index_ready + + gate.set() + await asyncio.wait_for(hook.index_build_task, timeout=30) + assert hook.filter.startup_index_ready + assert hook.filter.tool_router is not None + + inline_hook = await SemanticToolFilterHook.initialize_from_config( + config={"enabled": True}, + llm_router=_make_keyword_embedding_router([]), + redis_cache=None, + ) + assert inline_hook is not None + assert inline_hook.index_build_task is None + assert inline_hook.filter.startup_index_ready + assert inline_hook.filter.tool_router is not None 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 new file mode 100644 index 00000000000..0530de440ad --- /dev/null +++ b/tests/unit/proxy/_experimental/mcp_server/test_semantic_tool_index_store.py @@ -0,0 +1,289 @@ +""" +Unit tests for the Redis-shared MCP semantic tool index store. + +Covers the float32 base64 round trip, degradation on Redis failure, the +cross-pod build lock, and the CachingToolEncoder serving document +embeddings from the store instead of re-embedding. +""" + +import sys +from unittest.mock import Mock + +import pytest + +if sys.version_info < (3, 11): # BaseExceptionGroup is a builtin only from 3.11 + from exceptiongroup import BaseExceptionGroup + + +from mcp.types import Tool as MCPTool + + +requires_semantic_router = pytest.mark.skipif( + sys.version_info >= (3, 14), reason="The semantic-router extra excludes Python 3.14" +) + + +def _fake_embedding_router(recorded_inputs): + """Router stand-in returning deterministic 2-dim keyword one-hots, like the filter tests.""" + from litellm.types.utils import Embedding, EmbeddingResponse + + def _vector(text): + return [1.0, 0.0] if "ticket" in text.lower() else [0.0, 1.0] + + def mock_embedding_sync(*args, **kwargs): + texts = kwargs["input"] + recorded_inputs.append(list(texts)) + return EmbeddingResponse( + data=[Embedding(embedding=_vector(t), index=i, object="embedding") for i, t in enumerate(texts)], + model="text-embedding-3-small", + object="list", + usage={"prompt_tokens": 10, "total_tokens": 10}, + ) + + async def mock_embedding_async(*args, **kwargs): + return mock_embedding_sync(*args, **kwargs) + + router = Mock() + router.embedding = mock_embedding_sync + router.aembedding = mock_embedding_async + return router + + +class _InMemoryToolVectorStore: + def __init__(self, vectors=None): + self.vectors = dict(vectors or {}) + self.get_many_calls = 0 + self.acquire_calls = 0 + self.release_calls = 0 + self.lock_result = True + self._fill_on_call = None + + def fill_on_get_many_call(self, call_number): + self._fill_on_call = call_number + + async def get_many(self, keys): + self.get_many_calls += 1 + if self._fill_on_call is not None and self.get_many_calls >= self._fill_on_call: + self.vectors.update(self._pending_fill) + return {key: self.vectors[key] for key in keys if key in self.vectors} + + def seed_from(self, keys_vectors): + self._pending_fill = dict(keys_vectors) + + async def put_many(self, vectors): + self.vectors.update({k: tuple(v) for k, v in vectors.items()}) + + async def try_acquire_build_lock(self): + self.acquire_calls += 1 + return self.lock_result + + async def release_build_lock(self): + self.release_calls += 1 + + +class _RaisingRedisCache: + async def async_batch_get_cache(self, key_list, parent_otel_span=None): + raise ConnectionError("redis is down") + + async def async_set_cache_pipeline(self, cache_list, ttl=None, **kwargs): + raise ConnectionError("redis is down") + + async def async_get_cache(self, key): + raise ConnectionError("redis is down") + + async def async_set_cache(self, key, value, **kwargs): + raise ConnectionError("redis is down") + + +class _DictRedisCache: + def __init__(self): + self.backing = {} + self.pipeline_calls = 0 + + async def async_batch_get_cache(self, key_list, parent_otel_span=None): + return {key: self.backing[key] for key in key_list if key in self.backing} + + async def async_set_cache_pipeline(self, cache_list, ttl=None, **kwargs): + self.pipeline_calls += 1 + self.backing.update(dict(cache_list)) + + +def _tool(name, description): + return MCPTool(name=name, description=description, inputSchema={"type": "object"}) + + +def _make_filter(recorded_inputs, store=None, embedding_model="text-embedding-3-small"): + from litellm.proxy._experimental.mcp_server.semantic_tool_filter import ( + SemanticMCPToolFilter, + ) + + return SemanticMCPToolFilter( + embedding_model=embedding_model, + litellm_router_instance=_fake_embedding_router(recorded_inputs), + top_k=3, + similarity_threshold=0.3, + enabled=True, + vector_store=store, + lock_poll_interval_s=0.0, + ) + + +@requires_semantic_router +@pytest.mark.asyncio +async def test_store_put_get_round_trip_within_float32_tolerance(): + from litellm.proxy._experimental.mcp_server.semantic_tool_index_store import ( + RedisToolVectorStore, + tool_vector_key, + ) + + redis = _DictRedisCache() + store = RedisToolVectorStore(redis, Mock()) + + key = tool_vector_key("embed-a", "get a ticket") + await store.put_many({key: [0.6000000238, -0.8123456789, 3.14]}) + + hits = await store.get_many([key]) + assert key in hits + assert list(hits[key]) == pytest.approx([0.6000000238, -0.8123456789, 3.14], abs=1e-6) + + +@requires_semantic_router +@pytest.mark.asyncio +async def test_store_treats_corrupt_value_as_miss(): + from litellm.proxy._experimental.mcp_server.semantic_tool_index_store import ( + RedisToolVectorStore, + tool_vector_key, + ) + + redis = _DictRedisCache() + key = tool_vector_key("embed-a", "get a ticket") + redis.backing[key] = "!!! not base64 / not float32 bytes !!!" + store = RedisToolVectorStore(redis, Mock()) + + assert await store.get_many([key]) == {} + + +@requires_semantic_router +@pytest.mark.asyncio +async def test_raising_redis_still_builds_the_index(): + from litellm.proxy._experimental.mcp_server.semantic_tool_index_store import ( + RedisToolVectorStore, + ) + from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager + + redis = _RaisingRedisCache() + recorded = [] + tools = [_tool("srv-a_get_ticket", "get a ticket"), _tool("srv-b_get_weather", "weather forecast")] + store = RedisToolVectorStore(redis, PodLockManager(redis)) + filter_instance = _make_filter(recorded, store=store) + + await filter_instance._abuild_router(tools) + + assert filter_instance.tool_router is not None + assert [batch for batch in recorded if batch != ["test"]] == [ + ["get a ticket", "weather forecast"] + ] + + +@requires_semantic_router +@pytest.mark.asyncio +async def test_build_with_warm_store_never_embeds_descriptions(): + from litellm.proxy._experimental.mcp_server.semantic_tool_index_store import ( + tool_vector_key, + ) + + tools = [_tool("srv-a_get_ticket", "get a ticket"), _tool("srv-b_get_weather", "weather forecast")] + identity = "text-embedding-3-small" + warm = _InMemoryToolVectorStore( + {tool_vector_key(identity, description): [0.0, 1.0] for _, description in ((t.name, t.description) for t in tools)} + ) + recorded = [] + filter_instance = _make_filter(recorded, store=warm) + + await filter_instance._abuild_router(tools) + + assert filter_instance.tool_router is not None + # Only the dims probe embeds; the 300-tool warm path makes zero embedding calls. + assert recorded == [["test"]] + assert warm.acquire_calls == 0 + + +@requires_semantic_router +@pytest.mark.asyncio +async def test_cold_store_embeds_once_then_second_filter_embeds_nothing(): + tools = [_tool("srv-a_get_ticket", "get a ticket"), _tool("srv-b_get_weather", "weather forecast")] + store = _InMemoryToolVectorStore() + + first_recorded = [] + first = _make_filter(first_recorded, store=store) + await first._abuild_router(tools) + assert first.tool_router is not None + assert [batch for batch in first_recorded if batch != ["test"]] == [ + ["get a ticket", "weather forecast"] + ] + + second_recorded = [] + second = _make_filter(second_recorded, store=store) + await second._abuild_router(tools) + assert second.tool_router is not None + assert second_recorded == [["test"]] + + +@requires_semantic_router +@pytest.mark.asyncio +async def test_edited_description_reembeds_only_that_tool(): + tools = [_tool("srv-a_get_ticket", "get a ticket"), _tool("srv-b_get_weather", "weather forecast")] + store = _InMemoryToolVectorStore() + recorded = [] + filter_instance = _make_filter(recorded, store=store) + await filter_instance._abuild_router(tools) + + recorded.clear() + edited = [_tool("srv-a_get_ticket", "close a support ticket fast"), _tool("srv-b_get_weather", "weather forecast")] + await filter_instance._abuild_router(edited) + + assert [batch for batch in recorded if batch != ["test"]] == [["close a support ticket fast"]] + + +@requires_semantic_router +@pytest.mark.asyncio +async def test_different_embedding_identity_reembeds_everything(): + tools = [_tool("srv-a_get_ticket", "get a ticket"), _tool("srv-b_get_weather", "weather forecast")] + store = _InMemoryToolVectorStore() + + first_recorded = [] + await _make_filter(first_recorded, store=store, embedding_model="text-embedding-3-small")._abuild_router(tools) + + second_recorded = [] + second = _make_filter(second_recorded, store=store, embedding_model="text-embedding-3-large") + await second._abuild_router(tools) + + assert second.tool_router is not None + assert [batch for batch in second_recorded if batch != ["test"]] == [ + ["get a ticket", "weather forecast"] + ] + + +@requires_semantic_router +@pytest.mark.asyncio +async def test_lock_held_by_other_pod_waits_for_shared_vectors(): + from litellm.proxy._experimental.mcp_server.semantic_tool_index_store import ( + tool_vector_key, + ) + + tools = [_tool("srv-a_get_ticket", "get a ticket"), _tool("srv-b_get_weather", "weather forecast")] + store = _InMemoryToolVectorStore() + store.lock_result = False + identity = "text-embedding-3-small" + # The "holder" pod's vectors appear on the third get_many poll. + store.seed_from({tool_vector_key(identity, t.description): [0.0, 1.0] for t in tools}) + store.fill_on_get_many_call(3) + + recorded = [] + filter_instance = _make_filter(recorded, store=store) + await filter_instance._abuild_router(tools) + + assert filter_instance.tool_router is not None + assert store.acquire_calls >= 1 + assert store.get_many_calls >= 3 + assert recorded == [["test"]] diff --git a/tests/unit/proxy/proxy_server/test_proxy_config.py b/tests/unit/proxy/proxy_server/test_proxy_config.py index 720695b0f42..66b2f2787ac 100644 --- a/tests/unit/proxy/proxy_server/test_proxy_config.py +++ b/tests/unit/proxy/proxy_server/test_proxy_config.py @@ -5607,3 +5607,36 @@ async def test_proxy_config_validates_advertised_mcp_versions_at_load(tmp_path, return with pytest.raises(ValidationError): await ProxyConfig().load_config(router=None, config_file_path=str(config)) + + +@pytest.mark.asyncio +async def test_proxy_config_semantic_filter_db_sync_skips_reinit_while_index_build_in_flight(monkeypatch): + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + + settings: Final = { + "enabled": True, + "embedding_model": "embed", + "top_k": 3, + "similarity_threshold": 0.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": 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()) + + removed.assert_not_called() + finally: + gate.set() + await build_task From 61244d356a6e72bcbb80179d16ac2c60e6d62fb1 Mon Sep 17 00:00:00 2001 From: mrinal Date: Sat, 10 Oct 2026 20:28:13 +0000 Subject: [PATCH 2/9] 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() From bd30e9dfb88310a071e820bda8382cdd6cab1468 Mon Sep 17 00:00:00 2001 From: mrinal Date: Sat, 10 Oct 2026 21:11:32 +0000 Subject: [PATCH 3/9] test(mcp): integration coverage for redis-shared semantic tool index Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp/test_mcp_semantic_filter_startup.py | 526 ++++++++++++++++++ 1 file changed, 526 insertions(+) create mode 100644 tests/integration/mcp/test_mcp_semantic_filter_startup.py diff --git a/tests/integration/mcp/test_mcp_semantic_filter_startup.py b/tests/integration/mcp/test_mcp_semantic_filter_startup.py new file mode 100644 index 00000000000..b23e85c88cc --- /dev/null +++ b/tests/integration/mcp/test_mcp_semantic_filter_startup.py @@ -0,0 +1,526 @@ +import json +import os +import sys +import threading +import time +import uuid +from collections import Counter +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from pathlib import Path +from typing import Final, Literal + +import httpx +import pytest +import yaml +from integration._support.client import GATEWAY_LIMITS, Gateway, eventually +from integration._support.mcp import ScriptedTool, scripted_peer, text_result +from integration._support.process import ( + DB_PUSH, + OwnedProxy, + _launch, + _proxy_environment, + _proxy_root, + _stop_launch, + owned_proxy_process, +) +from integration._support.redis_process import OwnedRedis, owned_redis +from integration._support.wire import Reply, Request, wire_server + +PROMPT: Final = "please refund the invoice for ACME" +MCP_REF: Final = { + "type": "mcp", + "server_url": "litellm_proxy", + "server_label": "litellm", + "require_approval": "never", +} +SEMANTIC_SETTINGS: Final = { + "enabled": True, + "embedding_model": "semantic-embed", + "top_k": 3, + "similarity_threshold": 0.5, +} +NO_REDIS_ENV: Final = ( + "REDIS_HOST", + "REDIS_PORT", + "REDIS_URL", + "REDIS_USERNAME", + "REDIS_PASSWORD", + "REDIS_CLUSTER_NODES", + "REDIS_SENTINEL_NODES", + "REDIS_SSL", + "DATABASE_URL_READ_REPLICA", +) +Surface = Literal["chat", "responses"] +TOOL_COUNT: Final = 300 + + +class EmbeddingDouble: + def __init__(self, mode: Literal["healthy", "held"], delay: float = 0.0) -> None: + self.release = threading.Event() + self.mode = mode + self.delay = delay + self.arrived: list[tuple[tuple[str, ...], float]] = [] + self.answered: list[tuple[tuple[str, ...], float]] = [] + + def document_inputs(self, descriptions: frozenset[str]) -> list[str]: + """Embedding inputs that are tool descriptions; query embeddings ('test', prompts) excluded.""" + return [text for inputs, _ in self.arrived for text in inputs if text in descriptions] + + def respond(self, request: Request) -> Reply: + if request.method == "GET" and request.target.endswith("/models"): + return Reply(body=b'{"object": "list", "data": []}') + assert request.method == "POST" and request.target.endswith("/embeddings"), request.target + body: Final = json.loads(request.body) + raw: Final = body.get("input", []) + inputs: Final = tuple(str(text) for text in raw) if isinstance(raw, list) else (str(raw),) + self.arrived.append((inputs, time.monotonic())) + if self.mode == "held": + self.release.wait(timeout=120) + if self.delay: + time.sleep(self.delay) + data: Final = [ + { + "object": "embedding", + "index": index, + "embedding": [1.0, 0.0] if "invoice" in text.lower() else [0.0, 1.0], + } + for index, text in enumerate(inputs) + ] + self.answered.append((inputs, time.monotonic())) + return Reply( + body=json.dumps( + { + "object": "list", + "data": data, + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 1, "total_tokens": 1}, + } + ).encode() + ) + + +def _chat_double() -> tuple[Callable[[Request], Reply], list[dict[str, object]]]: + bodies: list[dict[str, object]] = [] + + def respond(request: Request) -> Reply: + if request.method == "GET" and request.target.endswith("/models"): + return Reply(body=b'{"object": "list", "data": []}') + body: Final = json.loads(request.body) + bodies.append(body) + identity: Final = uuid.uuid4().hex[:12] + usage: Final = {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15} + if request.target.endswith("/chat/completions"): + return Reply( + body=json.dumps( + { + "id": f"chatcmpl-{identity}", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "message": {"role": "assistant", "content": "done"}, + } + ], + "usage": usage, + } + ).encode() + ) + assert request.target.endswith("/responses"), request.target + return Reply( + body=json.dumps( + { + "id": f"resp_{identity}", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [ + { + "type": "message", + "id": f"msg_{identity}", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "done", "annotations": []}], + } + ], + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + } + ).encode() + ) + + return respond, bodies + + +def _tool_names(body: dict[str, object]) -> tuple[str, ...]: + tools: Final = body.get("tools") + if not isinstance(tools, list): + return () + return tuple( + str(tool["name"] if "name" in tool else tool["function"]["name"]) for tool in tools if isinstance(tool, dict) + ) + + +def _peer_names(bodies: list[dict[str, object]], alias: str) -> frozenset[str]: + prefix: Final = f"{alias}-" + return frozenset(name for body in bodies for name in _tool_names(body) if name.startswith(prefix)) + + +def _tools(invoice_description: str = "Refund an invoice for a customer") -> tuple[ScriptedTool, ...]: + refund: Final = ScriptedTool( + "refund_invoice", + lambda params: text_result("refunded"), + description=invoice_description, + ) + weather: Final = tuple( + ScriptedTool( + f"weather_{n}", + lambda params: text_result("sunny"), + description=f"Look up the weather forecast for city {n}", + ) + for n in range(TOOL_COUNT - 1) + ) + return (refund, *weather) + + +def _all_names(alias: str) -> frozenset[str]: + return frozenset((f"{alias}-refund_invoice", *(f"{alias}-weather_{n}" for n in range(TOOL_COUNT - 1)))) + + +def _descriptions(invoice_description: str = "Refund an invoice for a customer") -> frozenset[str]: + return frozenset(tool.description for tool in _tools(invoice_description)) + + +def _config( + directory: Path, + *, + alias: str, + peer_url: str, + embeddings_url: str, + chat_url: str, + redis: OwnedRedis | None, +) -> Path: + document: Final = { + "model_list": [ + { + "model_name": "semantic-embed", + "litellm_params": { + "model": "openai/text-embedding-3-small", + "api_base": f"{embeddings_url}/v1", + "api_key": "synthetic-embed-key", + }, + }, + { + "model_name": "chat-model", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_base": f"{chat_url}/v1", + "api_key": "synthetic-chat-key", + }, + }, + { + "model_name": "responses-model", + "litellm_params": { + "model": "openai/responses/gpt-4o-mini", + "api_base": f"{chat_url}/v1", + "api_key": "synthetic-chat-key", + }, + }, + ], + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + "store_model_in_db": True, + **({"coordination_redis": {"host": redis.host, "port": redis.port}} if redis is not None else {}), + }, + "mcp_servers": {alias: {"transport": "http", "url": peer_url}}, + "litellm_settings": {"mcp_semantic_tool_filter": SEMANTIC_SETTINGS}, + } + path: Final = directory / f"semantic-filter-{uuid.uuid4().hex}.yaml" + path.write_text(yaml.dump(document)) + return path + + +def _send(proxy: Gateway, surface: Surface) -> httpx.Response: + if surface == "chat": + return proxy.request( + "POST", + "/v1/chat/completions", + {"model": "chat-model", "messages": [{"role": "user", "content": PROMPT}], "tools": [MCP_REF]}, + ) + return proxy.request( + "POST", + "/v1/responses", + {"model": "responses-model", "input": PROMPT, "tools": [MCP_REF]}, + ) + + +def _request_is_filtered(proxy: Gateway, surface: Surface, bodies: list[dict[str, object]], alias: str) -> bool: + try: + response: Final = _send(proxy, surface) + except httpx.TransportError: + return False + return response.status_code == 200 and _peer_names(bodies[-1:], alias) == frozenset((f"{alias}-refund_invoice",)) + + +@contextmanager +def _proxy_without_readiness_gate( + gateway: Gateway, directory: Path, overrides: dict[str, str], *, config: Path +) -> Iterator[OwnedProxy]: + """owned_proxy_process without the readiness wait, for builds that must stay unready.""" + root: Final = _proxy_root() + environment: Final = _proxy_environment(gateway, overrides, NO_REDIS_ENV) + output: Final = Path(os.environ.get("INTEGRATION_RESULTS_DIR", str(directory))) + output.mkdir(parents=True, exist_ok=True) + command: Final = ( + sys.executable, + "-m", + "integration._support.proxy", + "--config", + str(config), + "--host", + "127.0.0.1", + "--num_workers", + "1", + "--timeout_worker_healthcheck", + "5", + *DB_PUSH, + ) + launch: Final = _launch(command, root, environment, output) + try: + with httpx.Client( + base_url=f"http://127.0.0.1:{launch.port}", timeout=15, trust_env=False, limits=GATEWAY_LIMITS + ) as client: + yield OwnedProxy(Gateway(client, gateway.key, gateway.upstream_url), launch.process, launch.log) + finally: + _stop_launch(launch) + + +def test_redis_backed_proxy_is_live_while_300_tool_index_builds(gateway: Gateway, tmp_path: Path) -> None: + alias: Final = "sf" + uuid.uuid4().hex[:8] + embeddings: Final = EmbeddingDouble("held") + respond, bodies = _chat_double() + with ( + owned_redis(tmp_path) as redis, + scripted_peer(*_tools()) as peer, + wire_server(respond) as chat_wire, + wire_server(embeddings.respond) as embedding_wire, + ): + try: + config: Final = _config( + tmp_path, + alias=alias, + peer_url=peer.url, + embeddings_url=embedding_wire.url, + chat_url=chat_wire.url, + redis=redis, + ) + with owned_proxy_process(gateway, tmp_path, {}, config=config) as proxy: + liveliness: Final = proxy.gateway.client.get("/health/liveliness") + assert liveliness.status_code == 200, liveliness.text + response: Final = _send(proxy.gateway, "chat") + assert response.status_code == 200, response.text + assert _peer_names(bodies[-1:], alias) == _all_names(alias), _peer_names(bodies[-1:], alias) + embeddings.release.set() + eventually( + lambda: _request_is_filtered(proxy.gateway, "chat", bodies, alias), + lambda filtered: filtered, + seconds=60, + ) + finally: + embeddings.release.set() + + +def test_second_proxy_on_warm_redis_embeds_no_descriptions(gateway: Gateway, tmp_path: Path) -> None: + alias: Final = "sf" + uuid.uuid4().hex[:8] + descriptions: Final = _descriptions() + with owned_redis(tmp_path) as redis, scripted_peer(*_tools()) as peer: + first_embeddings: Final = EmbeddingDouble("healthy") + respond, bodies = _chat_double() + with ( + wire_server(respond) as chat_wire, + wire_server(first_embeddings.respond) as first_embedding_wire, + ): + config: Final = _config( + tmp_path, + alias=alias, + peer_url=peer.url, + embeddings_url=first_embedding_wire.url, + chat_url=chat_wire.url, + redis=redis, + ) + with owned_proxy_process(gateway, tmp_path, {}, config=config) as first_proxy: + eventually( + lambda: _request_is_filtered(first_proxy.gateway, "chat", bodies, alias), + lambda filtered: filtered, + seconds=60, + ) + second_embeddings: Final = EmbeddingDouble("healthy") + second_respond, second_bodies = _chat_double() + with ( + wire_server(second_respond) as second_chat_wire, + wire_server(second_embeddings.respond) as second_embedding_wire, + ): + config = _config( + tmp_path, + alias=alias, + peer_url=peer.url, + embeddings_url=second_embedding_wire.url, + chat_url=second_chat_wire.url, + redis=redis, + ) + with owned_proxy_process(gateway, tmp_path, {}, config=config) as second_proxy: + assert second_embeddings.document_inputs(descriptions) == [] + assert _request_is_filtered(second_proxy.gateway, "chat", second_bodies, alias) + + +def test_two_proxies_cold_starting_together_embed_each_description_once(gateway: Gateway, tmp_path: Path) -> None: + alias: Final = "sf" + uuid.uuid4().hex[:8] + descriptions: Final = _descriptions() + first_embeddings: Final = EmbeddingDouble("healthy", delay=10.0) + second_embeddings: Final = EmbeddingDouble("healthy") + respond_a, bodies_a = _chat_double() + respond_b, bodies_b = _chat_double() + with ( + owned_redis(tmp_path) as redis, + scripted_peer(*_tools()) as peer, + wire_server(respond_a) as chat_wire_a, + wire_server(respond_b) as chat_wire_b, + wire_server(first_embeddings.respond) as embedding_wire_a, + wire_server(second_embeddings.respond) as embedding_wire_b, + ): + config_a: Final = _config( + tmp_path, + alias=alias, + peer_url=peer.url, + embeddings_url=embedding_wire_a.url, + chat_url=chat_wire_a.url, + redis=redis, + ) + config_b: Final = _config( + tmp_path, + alias=alias, + peer_url=peer.url, + embeddings_url=embedding_wire_b.url, + chat_url=chat_wire_b.url, + redis=redis, + ) + with ( + owned_proxy_process(gateway, tmp_path, {}, config=config_a) as proxy_a, + owned_proxy_process(gateway, tmp_path, {}, config=config_b) as proxy_b, + ): + eventually( + lambda: _request_is_filtered(proxy_a.gateway, "chat", bodies_a, alias), + lambda filtered: filtered, + seconds=120, + ) + eventually( + lambda: _request_is_filtered(proxy_b.gateway, "chat", bodies_b, alias), + lambda filtered: filtered, + seconds=60, + ) + document_inputs: Final = Counter( + first_embeddings.document_inputs(descriptions) + second_embeddings.document_inputs(descriptions) + ) + assert document_inputs == Counter({description: 1 for description in descriptions}) + + +def test_edited_description_reembeds_only_that_tool(gateway: Gateway, tmp_path: Path) -> None: + alias: Final = "sf" + uuid.uuid4().hex[:8] + descriptions_v1: Final = _descriptions() + edited: Final = "Cancel and refund an invoice instantly" + with owned_redis(tmp_path) as redis: + first_embeddings: Final = EmbeddingDouble("healthy") + respond, bodies = _chat_double() + with ( + scripted_peer(*_tools()) as peer_v1, + wire_server(respond) as chat_wire, + wire_server(first_embeddings.respond) as first_embedding_wire, + ): + config: Final = _config( + tmp_path, + alias=alias, + peer_url=peer_v1.url, + embeddings_url=first_embedding_wire.url, + chat_url=chat_wire.url, + redis=redis, + ) + with owned_proxy_process(gateway, tmp_path, {}, config=config) as first_proxy: + eventually( + lambda: _request_is_filtered(first_proxy.gateway, "chat", bodies, alias), + lambda filtered: filtered, + seconds=60, + ) + second_embeddings: Final = EmbeddingDouble("healthy") + second_respond, second_bodies = _chat_double() + with ( + scripted_peer(*_tools(edited)) as peer_v2, + wire_server(second_respond) as second_chat_wire, + wire_server(second_embeddings.respond) as second_embedding_wire, + ): + config = _config( + tmp_path, + alias=alias, + peer_url=peer_v2.url, + embeddings_url=second_embedding_wire.url, + chat_url=second_chat_wire.url, + redis=redis, + ) + with owned_proxy_process(gateway, tmp_path, {}, config=config) as second_proxy: + assert second_embeddings.document_inputs(descriptions_v1) == [] + document_inputs: Final = [ + text + for inputs, _ in second_embeddings.arrived + for text in inputs + if text not in descriptions_v1 and text not in ("test", PROMPT) + ] + assert document_inputs == [edited] + eventually( + lambda: _request_is_filtered(second_proxy.gateway, "chat", second_bodies, alias), + lambda filtered: filtered, + seconds=60, + ) + + +def test_without_redis_startup_waits_for_the_index(gateway: Gateway, tmp_path: Path) -> None: + alias: Final = "sf" + uuid.uuid4().hex[:8] + embeddings: Final = EmbeddingDouble("held") + respond, bodies = _chat_double() + with ( + scripted_peer(*_tools()) as peer, + wire_server(respond) as chat_wire, + wire_server(embeddings.respond) as embedding_wire, + ): + try: + config: Final = _config( + tmp_path, + alias=alias, + peer_url=peer.url, + embeddings_url=embedding_wire.url, + chat_url=chat_wire.url, + redis=None, + ) + with _proxy_without_readiness_gate(gateway, tmp_path, {}, config=config) as proxy: + + def liveliness() -> int | None: + try: + return proxy.gateway.client.get("/health/liveliness", timeout=2).status_code + except httpx.TransportError: + return None + + deadline: Final = time.monotonic() + 5 + while time.monotonic() < deadline: + assert liveliness() != 200, "blocking build should keep liveliness from 200 while held" + time.sleep(0.5) + embeddings.release.set() + eventually(liveliness, lambda status: status == 200, seconds=60) + eventually( + lambda: _request_is_filtered(proxy.gateway, "chat", bodies, alias), + lambda filtered: filtered, + seconds=60, + ) + finally: + embeddings.release.set() From fafdf7a82cbe6b2e180d289cabdf96dc08b1dea5 Mon Sep 17 00:00:00 2001 From: mrinal Date: Sat, 10 Oct 2026 23:38:08 +0000 Subject: [PATCH 4/9] fix(mcp): only reindex changed tool descriptions from registry-expanded tools Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/semantic_tool_filter.py | 51 ++-- .../proxy/hooks/mcp_semantic_filter/hook.py | 4 +- .../mcp_server/test_semantic_tool_filter.py | 258 +++++++++++------- 3 files changed, 198 insertions(+), 115 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py b/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py index dbd39a5d077..4409c60cb0a 100644 --- a/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py +++ b/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py @@ -315,45 +315,53 @@ class SemanticMCPToolFilter: await asyncio.sleep(self._lock_poll_interval_s) return False - def _has_tools_missing_from_index(self, tools: Sequence[object]) -> bool: + def _has_tools_missing_from_index(self, tools: Sequence[object], *, reindex_changed_descriptions: bool) -> bool: """Allocation-free check for any named tool not yet in the semantic index.""" return any( - name and (name not in self._tool_map or self.extract_tool_info(self._tool_map[name])[1] != description) + name + and ( + name not in self._tool_map + or (reindex_changed_descriptions and self.extract_tool_info(self._tool_map[name])[1] != description) + ) for name, description in (self.extract_tool_info(t) for t in tools) ) - def _tools_missing_from_index(self, tools: Sequence[object]) -> Mapping[str, object]: + def _tools_missing_from_index( + self, tools: Sequence[object], *, reindex_changed_descriptions: bool + ) -> Mapping[str, object]: """Map name -> tool for every named tool not yet in the semantic index.""" infos: Final = ((self.extract_tool_info(t), t) for t in tools) return { name: tool for (name, description), tool in infos - if name and (name not in self._tool_map or self.extract_tool_info(self._tool_map[name])[1] != description) + if name + and ( + name not in self._tool_map + or (reindex_changed_descriptions and self.extract_tool_info(self._tool_map[name])[1] != description) + ) } - async def _ensure_tools_indexed(self, available_tools: Sequence[object]) -> None: + async def _ensure_tools_indexed( + self, available_tools: Sequence[object], *, reindex_changed_descriptions: bool + ) -> None: """ - Index request-time tools the startup build never saw. + Index request-time tools missing from the startup index. - The startup index lists every registered MCP server WITHOUT per-user - credentials, so servers requiring per-user auth (interactive OAuth - tokens, user-scoped env vars) contribute zero routes. Tools reaching - the filter came through an authenticated expansion; without indexing - them here they can never be selected, so requests either bypass - filtering entirely (N->N) or lose every tool to unrelated matches. - - Runs async-only (no synchronous embedding on the request path) and - never writes shared error state: an embedding failure here raises and - is scoped to the requesting call, so one request's oversized tool - description cannot poison the filter for other users on the worker. + With ``reindex_changed_descriptions`` a tool whose description changed + also counts as missing: its old route is dropped before re-embedding so + a failed embed leaves it "missing" and the next request retries. """ from semantic_router.routers import SemanticRouter - if not self._has_tools_missing_from_index(available_tools): + if not self._has_tools_missing_from_index( + available_tools, reindex_changed_descriptions=reindex_changed_descriptions + ): return async with self._index_sync_lock: - missing: Final = self._tools_missing_from_index(available_tools) + missing: Final = self._tools_missing_from_index( + available_tools, reindex_changed_descriptions=reindex_changed_descriptions + ) if not missing: return @@ -371,6 +379,7 @@ class SemanticMCPToolFilter: else: for name in (name for name in missing if name in self._tool_map): await self.tool_router.adelete(name) # pyright: ignore[reportUnknownMemberType] # semantic_router ships no type stubs + self._tool_map.pop(name) await self.tool_router.aadd(list(routes)) self._tool_map.update(missing) @@ -383,6 +392,8 @@ class SemanticMCPToolFilter: query: str, available_tools: list[Any], top_k: int | None = None, + *, + reindex_changed_descriptions: bool = False, ) -> list[Any]: """ Filter tools semantically based on query. @@ -414,7 +425,7 @@ class SemanticMCPToolFilter: # Run semantic filtering try: - await self._ensure_tools_indexed(available_tools) + await self._ensure_tools_indexed(available_tools, reindex_changed_descriptions=reindex_changed_descriptions) if self.tool_router is None: verbose_logger.warning("Semantic router could not be built from the request's tools") diff --git a/litellm/proxy/hooks/mcp_semantic_filter/hook.py b/litellm/proxy/hooks/mcp_semantic_filter/hook.py index 54f8d9f3ba6..3682213309b 100644 --- a/litellm/proxy/hooks/mcp_semantic_filter/hook.py +++ b/litellm/proxy/hooks/mcp_semantic_filter/hook.py @@ -194,7 +194,9 @@ class SemanticToolFilterHook(CustomLogger): verbose_proxy_logger.debug("No user query found, skipping semantic filter on expanded MCP tools") return expanded_tools - return await self.filter.filter_tools(query=user_query, available_tools=expanded_tools) + return await self.filter.filter_tools( + query=user_query, available_tools=expanded_tools, reindex_changed_descriptions=True + ) def _selected_tool_names(self, filtered_tools: Sequence[object]) -> list[str]: """Names of the semantically selected tools, as produced by the MCP expansion.""" diff --git a/tests/unit/proxy/_experimental/mcp_server/test_semantic_tool_filter.py b/tests/unit/proxy/_experimental/mcp_server/test_semantic_tool_filter.py index 1ba1b33b094..1f578dc99ad 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_semantic_tool_filter.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_semantic_tool_filter.py @@ -129,25 +129,17 @@ async def test_semantic_filter_basic_filtering(): ) # Assertions - validate filtering mechanics work - assert ( - len(filtered) <= 3 - ), f"Should return at most 3 tools (top_k), got {len(filtered)}" + assert len(filtered) <= 3, f"Should return at most 3 tools (top_k), got {len(filtered)}" assert len(filtered) > 0, "Should return at least some tools" - assert len(filtered) < len( - tools - ), f"Should filter down from {len(tools)} tools, got {len(filtered)}" + assert len(filtered) < len(tools), f"Should filter down from {len(tools)} tools, got {len(filtered)}" # Validate tools are actual MCPTool objects for tool in filtered: assert hasattr(tool, "name"), "Filtered result should be MCPTool with name" - assert hasattr( - tool, "description" - ), "Filtered result should be MCPTool with description" + assert hasattr(tool, "description"), "Filtered result should be MCPTool with description" filtered_names = [t.name for t in filtered] - print( - f"✅ Successfully filtered {len(tools)} tools down to top {len(filtered)}: {filtered_names}" - ) + print(f"✅ Successfully filtered {len(tools)} tools down to top {len(filtered)}: {filtered_names}") print(f" Filter respects top_k parameter correctly") @@ -226,12 +218,7 @@ async def test_semantic_filter_disabled(): SemanticMCPToolFilter, ) - tools = [ - MCPTool( - name=f"tool_{i}", description=f"Tool {i}", inputSchema={"type": "object"} - ) - for i in range(10) - ] + tools = [MCPTool(name=f"tool_{i}", description=f"Tool {i}", inputSchema={"type": "object"}) for i in range(10)] mock_router = Mock() @@ -251,9 +238,7 @@ async def test_semantic_filter_disabled(): ) # Should return all tools when disabled - assert len(filtered) == len( - tools - ), f"Expected all {len(tools)} tools, got {len(filtered)}" + assert len(filtered) == len(tools), f"Expected all {len(tools)} tools, got {len(filtered)}" @pytest.mark.asyncio @@ -373,12 +358,7 @@ async def test_semantic_filter_hook_triggers_on_completion(): ) # Prepare data - completion request with tools - tools = [ - MCPTool( - name=f"tool_{i}", description=f"Tool {i}", inputSchema={"type": "object"} - ) - for i in range(10) - ] + tools = [MCPTool(name=f"tool_{i}", description=f"Tool {i}", inputSchema={"type": "object"}) for i in range(10)] # Build router with the tools before filtering filter_instance._build_router(tools) @@ -408,9 +388,7 @@ async def test_semantic_filter_hook_triggers_on_completion(): # Assertions assert result is not None, "Hook should return modified data" assert "tools" in result, "Result should contain tools" - assert len(result["tools"]) < len( - tools - ), f"Hook should filter tools, got {len(result['tools'])}/{len(tools)}" + assert len(result["tools"]) < len(tools), f"Hook should filter tools, got {len(result['tools'])}/{len(tools)}" print(f"✅ Hook triggered correctly: {len(tools)} -> {len(result['tools'])} tools") @@ -557,18 +535,12 @@ async def test_semantic_filter_hook_preserves_native_tools(): filtered = result["tools"] # Native tools must survive - native_in_result = [ - t for t in filtered if isinstance(t, dict) and t.get("type") == "function" - ] - assert ( - len(native_in_result) == 2 - ), f"Both native tools must survive, got {len(native_in_result)}" + native_in_result = [t for t in filtered if isinstance(t, dict) and t.get("type") == "function"] + assert len(native_in_result) == 2, f"Both native tools must survive, got {len(native_in_result)}" # MCP tools should be filtered (top_k=2) mcp_in_result = [t for t in filtered if not isinstance(t, dict)] - assert ( - len(mcp_in_result) <= 2 - ), f"MCP tools should be filtered to top_k=2, got {len(mcp_in_result)}" + assert len(mcp_in_result) <= 2, f"MCP tools should be filtered to top_k=2, got {len(mcp_in_result)}" # Total should be native + filtered MCP assert len(filtered) <= 4, f"Expected at most 4 tools, got {len(filtered)}" @@ -579,12 +551,8 @@ async def test_semantic_filter_hook_preserves_native_tools(): # Stats should report MCP-only counts, not inflated with native tools stats = result["metadata"]["litellm_semantic_filter_stats"] mcp_before, mcp_after = stats.split("->") - assert ( - int(mcp_before) == 5 - ), f"Stats 'from' should be MCP count (5), got {mcp_before}" - assert int(mcp_after) == len( - mcp_in_result - ), f"Stats 'to' should match filtered MCP count, got {mcp_after}" + assert int(mcp_before) == 5, f"Stats 'from' should be MCP count (5), got {mcp_before}" + assert int(mcp_after) == len(mcp_in_result), f"Stats 'to' should match filtered MCP count, got {mcp_after}" print( f"✅ Hook preserves native tools: {len(all_tools)} -> {len(filtered)} " @@ -679,20 +647,15 @@ async def test_semantic_filter_hook_all_native_tools(): filtered = result["tools"] # All native tools must pass through - assert ( - len(filtered) == 3 - ), f"All 3 native tools must pass through, got {len(filtered)}" + assert len(filtered) == 3, f"All 3 native tools must pass through, got {len(filtered)}" # No spurious semantic filter stats (P2 fix) - assert ( - "litellm_semantic_filter_stats" not in result["metadata"] - ), "Should NOT emit semantic filter stats for all-native-tool requests" - - print( - f"✅ Hook passes through all {len(filtered)} native tools, " - f"no spurious filter headers emitted" + assert "litellm_semantic_filter_stats" not in result["metadata"], ( + "Should NOT emit semantic filter stats for all-native-tool requests" ) + print(f"✅ Hook passes through all {len(filtered)} native tools, no spurious filter headers emitted") + @requires_semantic_router @pytest.mark.asyncio @@ -758,8 +721,7 @@ async def test_semantic_filter_hook_responses_api_name_collision(): # Verify classification: should be native, not MCP assert not hook._is_mcp_tool(responses_api_tool), ( - "Responses API tool with type=function + top-level name " - "should be classified as native, not MCP" + "Responses API tool with type=function + top-level name should be classified as native, not MCP" ) # Full hook test: all-native request should preserve tools @@ -890,9 +852,9 @@ async def test_semantic_filter_hook_filters_expanded_litellm_proxy_tools(): for name in allowed_tools: assert name in expanded_names, "Selected tool names must come from the expanded tools" - assert ( - "litellm_semantic_filter_stats" in result["metadata"] - ), "Filter stats must be emitted for the litellm_proxy expansion path" + assert "litellm_semantic_filter_stats" in result["metadata"], ( + "Filter stats must be emitted for the litellm_proxy expansion path" + ) stats = result["metadata"]["litellm_semantic_filter_stats"] total, selected = stats.split("->") assert int(total) == 5, f"Stats 'from' should be pre-filter expanded count (5), got {total}" @@ -1274,9 +1236,9 @@ async def test_semantic_filter_hook_expansion_skips_filter_when_disabled(): "require_approval": "never", } ], "The MCP reference must be left intact for the MCP gateway to expand" - assert ( - "litellm_semantic_filter_stats" not in data["metadata"] - ), "No filter stats may be emitted when the filter is disabled" + assert "litellm_semantic_filter_stats" not in data["metadata"], ( + "No filter stats may be emitted when the filter is disabled" + ) print("✅ Disabled filter: MCP reference untouched, no spurious stats") @@ -1380,9 +1342,7 @@ async def test_semantic_filter_hook_preserves_tool_order(): assert filtered[1] is native_tool, "Second tool should be native_tool" assert filtered[2] is mcp_tool_b, "Third tool should be mcp_tool_b" - print( - "✅ Tool ordering preserved: [mcp_A, native, mcp_B] maintained after filtering" - ) + print("✅ Tool ordering preserved: [mcp_A, native, mcp_B] maintained after filtering") class TestGetToolsByNames: @@ -1492,9 +1452,7 @@ class TestGetToolsByNames: filter_instance = self._make_filter() available_tools = [{"name": "litellm_api-fs-read_file", "description": "read"}] - matched = filter_instance._get_tools_by_names( - ["fs-read_file", "api-fs-read_file"], available_tools - ) + matched = filter_instance._get_tools_by_names(["fs-read_file", "api-fs-read_file"], available_tools) assert len(matched) == 1 @@ -1525,9 +1483,7 @@ class TestGetToolsByNames: {"name": "litellm_fs-delete", "description": "delete"}, ] - matched = filter_instance._get_tools_by_names( - ["fs-write", "fs-delete", "fs-read"], available_tools - ) + matched = filter_instance._get_tools_by_names(["fs-write", "fs-delete", "fs-read"], available_tools) names = [t["name"] for t in matched] assert names == [ @@ -1682,10 +1638,7 @@ async def test_semantic_filter_fails_closed_on_query_time_context_window_error() state = {"raise_context_error": False} filter_instance = _make_context_window_filter(state) - tools = [ - MCPTool(name=f"tool_{i}", description=f"Tool {i}", inputSchema={"type": "object"}) - for i in range(5) - ] + tools = [MCPTool(name=f"tool_{i}", description=f"Tool {i}", inputSchema={"type": "object"}) for i in range(5)] filter_instance._build_router(tools) assert filter_instance.tool_router is not None @@ -1715,10 +1668,7 @@ async def test_semantic_filter_records_build_time_context_window_error(): state = {"raise_context_error": True} filter_instance = _make_context_window_filter(state) - tools = [ - MCPTool(name=f"tool_{i}", description=f"Tool {i}", inputSchema={"type": "object"}) - for i in range(5) - ] + tools = [MCPTool(name=f"tool_{i}", description=f"Tool {i}", inputSchema={"type": "object"}) for i in range(5)] filter_instance._build_router(tools) assert filter_instance.tool_router is None @@ -1749,10 +1699,7 @@ async def test_semantic_filter_hook_fails_closed_on_context_window_error(): state = {"raise_context_error": False} filter_instance = _make_context_window_filter(state) - tools = [ - MCPTool(name=f"tool_{i}", description=f"Tool {i}", inputSchema={"type": "object"}) - for i in range(5) - ] + tools = [MCPTool(name=f"tool_{i}", description=f"Tool {i}", inputSchema={"type": "object"}) for i in range(5)] filter_instance._build_router(tools) hook = SemanticToolFilterHook(filter_instance) @@ -1861,10 +1808,7 @@ async def test_semantic_filter_hook_ignores_build_error_for_native_only_tools(): state = {"raise_context_error": True} filter_instance = _make_context_window_filter(state) - mcp_tools = [ - MCPTool(name=f"tool_{i}", description=f"Tool {i}", inputSchema={"type": "object"}) - for i in range(3) - ] + mcp_tools = [MCPTool(name=f"tool_{i}", description=f"Tool {i}", inputSchema={"type": "object"}) for i in range(3)] filter_instance._build_router(mcp_tools) assert filter_instance.context_window_error is not None @@ -1921,9 +1865,7 @@ def test_is_context_window_error_detection_variants(): def wrap_without_explicit_chaining(): try: - raise litellm.ContextWindowExceededError( - message="overflow", model="m", llm_provider="openai" - ) + raise litellm.ContextWindowExceededError(message="overflow", model="m", llm_provider="openai") except litellm.ContextWindowExceededError: raise ValueError("wrapper without explicit chaining") @@ -2266,6 +2208,7 @@ async def test_semantic_filter_reindexes_tool_when_description_changes(): filtered = await filter_instance.filter_tools( query="linear ticket", available_tools=[edited], + reindex_changed_descriptions=True, ) assert [edited] == filtered @@ -2274,6 +2217,135 @@ async def test_semantic_filter_reindexes_tool_when_description_changes(): assert filter_instance.extract_tool_info(indexed_tool)[1] == "get a Linear ticket by id" +@requires_semantic_router +@pytest.mark.asyncio +async def test_hook_request_body_tool_cannot_trigger_description_reindex(): + """ + A caller-supplied tool with a registered name but a made-up description + must not re-embed or replace the worker-wide index entry. Only the + registry-expanded path is trusted to reindex changed descriptions. + """ + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + + recorded_inputs = [] + filter_instance = _make_keyword_filter(recorded_inputs) + original = MCPTool( + name="srv-a-get_thing", + description="get the weather for a city", + inputSchema={"type": "object"}, + ) + filter_instance._build_router([original]) + recorded_inputs.clear() + + spoofed = { + "name": "srv-a-get_thing", + "description": "totally made up description", + "inputSchema": {"type": "object"}, + } + hook = SemanticToolFilterHook(filter_instance) + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "weather"}], + "tools": [spoofed], + "metadata": {}, + } + + await hook.async_pre_call_hook(user_api_key_dict=Mock(), cache=Mock(), data=data, call_type="completion") + + assert recorded_inputs == [["weather"]], "request-body tools must not trigger document embeddings" + indexed_tool = filter_instance._tool_map["srv-a-get_thing"] + assert filter_instance.extract_tool_info(indexed_tool)[1] == "get the weather for a city" + + +@requires_semantic_router +@pytest.mark.asyncio +async def test_hook_expanded_path_reindexes_changed_description(): + """ + Tools the proxy expanded from the MCP registry are trusted: a changed + description re-embeds exactly that one tool through _filter_expanded_tools. + """ + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + + recorded_inputs = [] + filter_instance = _make_keyword_filter(recorded_inputs) + original = MCPTool( + name="srv-a-get_thing", + description="get the weather for a city", + inputSchema={"type": "object"}, + ) + filter_instance._build_router([original]) + recorded_inputs.clear() + + edited = { + "name": "srv-a-get_thing", + "description": "get a Linear ticket by id", + "inputSchema": {"type": "object"}, + } + hook = SemanticToolFilterHook(filter_instance) + filtered = await hook._filter_expanded_tools( + data={"messages": [{"role": "user", "content": "linear ticket"}]}, + expanded_tools=[edited], + ) + + assert [edited] == filtered + assert recorded_inputs == [["get a Linear ticket by id"], ["linear ticket"]] + indexed_tool = filter_instance._tool_map["srv-a-get_thing"] + assert filter_instance.extract_tool_info(indexed_tool)[1] == "get a Linear ticket by id" + + +@requires_semantic_router +@pytest.mark.asyncio +async def test_failed_reindex_retries_on_next_request(): + """ + When reindexing a changed description raises mid-embed, the tool must stay + 'missing' so the next request retries it instead of silently dropping it + from matching (its route was already deleted from the router). + """ + recorded_inputs = [] + filter_instance = _make_keyword_filter(recorded_inputs) + original = MCPTool( + name="srv-a-get_thing", + description="get the weather for a city", + inputSchema={"type": "object"}, + ) + filter_instance._build_router([original]) + recorded_inputs.clear() + + original_aembedding = filter_instance.router_instance.aembedding + calls = {"n": 0} + + async def fail_once(*args, **kwargs): + calls["n"] += 1 + if calls["n"] == 1: + raise RuntimeError("embedding backend down") + return await original_aembedding(*args, **kwargs) + + filter_instance.router_instance.aembedding = fail_once + + edited = MCPTool( + name="srv-a-get_thing", + description="get a Linear ticket by id", + inputSchema={"type": "object"}, + ) + await filter_instance.filter_tools( + query="linear ticket", + available_tools=[edited], + reindex_changed_descriptions=True, + ) + assert "srv-a-get_thing" not in filter_instance._tool_map, ( + "failed re-embed must leave the tool missing, not half-indexed" + ) + + filtered = await filter_instance.filter_tools( + query="linear ticket", + available_tools=[edited], + reindex_changed_descriptions=True, + ) + assert [edited] == filtered + assert any("get a Linear ticket by id" in text for batch in recorded_inputs for text in batch) + assert "srv-a-get_thing" in filter_instance._tool_map + + @requires_semantic_router @pytest.mark.asyncio async def test_hook_returns_before_shared_index_build_finishes(): @@ -2315,9 +2387,7 @@ async def test_hook_returns_before_shared_index_build_finishes(): ) 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=[tool]) - ), + patch.object(global_mcp_server_manager, "get_tools_for_server", new=AsyncMock(return_value=[tool])), ): hook = await SemanticToolFilterHook.initialize_from_config( config={"enabled": True}, From 845c4a9a4501da53d4510b112ceee8e353b1be93 Mon Sep 17 00:00:00 2001 From: mrinal Date: Sat, 10 Oct 2026 23:47:04 +0000 Subject: [PATCH 5/9] fix(mcp): declare the semantic tool index redis key family Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../_experimental/mcp_server/semantic_tool_index_store.py | 4 ++++ 1 file changed, 4 insertions(+) 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 78a9528b14e..dc21e0c965b 100644 --- a/litellm/proxy/_experimental/mcp_server/semantic_tool_index_store.py +++ b/litellm/proxy/_experimental/mcp_server/semantic_tool_index_store.py @@ -14,6 +14,7 @@ from typing import TYPE_CHECKING, Final, Protocol from pydantic import TypeAdapter +from litellm._internal_context import with_service_target from litellm._logging import verbose_logger from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager from litellm.router_strategy.auto_router.litellm_encoder import LiteLLMRouterEncoder @@ -27,6 +28,7 @@ _BUILD_LOCK_ID: Final = "mcp_semantic_tool_index_build" _BUILD_LOCK_TTL_S: Final = 300 _VECTOR_TTL_S: Final = 30 * 24 * 60 * 60 _PUT_CHUNK_SIZE: Final = 50 +_CACHE_TARGET: Final = "mcp_semantic_tool_index" def tool_vector_key(embedding_identity: str, text: str) -> str: @@ -56,6 +58,7 @@ class RedisToolVectorStore: self._lock_manager = lock_manager self._degraded = False + @with_service_target(_CACHE_TARGET) async def get_many(self, keys: Sequence[str]) -> Mapping[str, tuple[float, ...]]: try: raw: Final = TypeAdapter(Mapping[str, object]).validate_python( @@ -67,6 +70,7 @@ class RedisToolVectorStore: return {} return {key: vector for key in keys if (vector := _decode_vector(raw.get(key))) is not None} + @with_service_target(_CACHE_TARGET) async def put_many(self, vectors: Mapping[str, Sequence[float]]) -> None: try: entries: Final = tuple((key, _encode_vector(vector)) for key, vector in vectors.items()) From 66ce89fc11d3f99cf8206a5e93c2ca1303edd771 Mon Sep 17 00:00:00 2001 From: mrinal Date: Sun, 11 Oct 2026 03:28:31 +0000 Subject: [PATCH 6/9] fix(mcp): harden the shared semantic tool index cache Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/_experimental/mcp_server/AGENTS.md | 1 + .../mcp_server/semantic_tool_encoder.py | 83 +++++++ .../mcp_server/semantic_tool_filter.py | 30 ++- .../mcp_server/semantic_tool_index_store.py | 53 ----- .../mcp_server/test_semantic_tool_encoder.py | 208 ++++++++++++++++++ .../mcp_server/test_semantic_tool_filter.py | 32 +++ .../test_semantic_tool_index_store.py | 149 ++++++------- 7 files changed, 413 insertions(+), 143 deletions(-) create mode 100644 litellm/proxy/_experimental/mcp_server/semantic_tool_encoder.py create mode 100644 tests/unit/proxy/_experimental/mcp_server/test_semantic_tool_encoder.py diff --git a/litellm/proxy/_experimental/mcp_server/AGENTS.md b/litellm/proxy/_experimental/mcp_server/AGENTS.md index 809a2086a06..2abab418446 100644 --- a/litellm/proxy/_experimental/mcp_server/AGENTS.md +++ b/litellm/proxy/_experimental/mcp_server/AGENTS.md @@ -42,6 +42,7 @@ litellm/proxy/_experimental/mcp_server/ elicitation_handler.py # MCP elicitation relay flow semantic_tool_filter.py # semantic filtering of available MCP tools semantic_tool_index_store.py # Redis-shared embeddings + build lock for the semantic tool index + semantic_tool_encoder.py # CachingToolEncoder bridging the store and LiteLLMRouterEncoder tool_search.py # opt-in virtual tools (mcp_tool_search + mcp_tool_call) for large catalogs guardrail_translation/ handler.py # MCP guardrail result translation diff --git a/litellm/proxy/_experimental/mcp_server/semantic_tool_encoder.py b/litellm/proxy/_experimental/mcp_server/semantic_tool_encoder.py new file mode 100644 index 00000000000..338b5f06459 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/semantic_tool_encoder.py @@ -0,0 +1,83 @@ +""" +Caching encoder for the shared MCP semantic tool index. + +Lives apart from semantic_tool_index_store because LiteLLMRouterEncoder pulls +in semantic_router, an optional extra that must stay out of the store's import +chain. +""" + +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Final + +from litellm.proxy._experimental.mcp_server.semantic_tool_index_store import ( + ToolVectorStore, + tool_vector_key, +) +from litellm.router_strategy.auto_router.litellm_encoder import LiteLLMRouterEncoder + +if TYPE_CHECKING: + from litellm.router import Router + + +class CachingToolEncoder(LiteLLMRouterEncoder): + """LiteLLMRouterEncoder that serves document embeddings from a ToolVectorStore. + + Queries are never cached; only add/aadd's document utterances go through + the store. Duplicate texts share one key and one embed. Cached vectors + whose length disagrees with a live probe are treated as misses, so a + stale entry can never feed the index a wrong-shape vector. + """ + + _tool_vector_store: ToolVectorStore + _embedding_identity: str + _dimensions: int | None + + def __init__( + self, + *, + store: ToolVectorStore, + embedding_identity: str, + litellm_router_instance: "Router", + model_name: str, + score_threshold: float | None = None, + ) -> None: + super().__init__( + litellm_router_instance=litellm_router_instance, + model_name=model_name, + score_threshold=score_threshold, + ) + self._tool_vector_store = store + self._embedding_identity = embedding_identity + self._dimensions = None + + async def adimensions(self) -> int: + if self._dimensions is None: + probe: Final = await self.aencode_queries(["test"]) # pyright: ignore[reportUnknownMemberType] # semantic_router ships no type stubs + self._dimensions = len(probe[0]) + return self._dimensions + + def encode_documents( # pyright: ignore[reportIncompatibleMethodOverride] # semantic_router calls encode_documents without kwargs + self, docs: Sequence[str] + ) -> list[list[float]]: # mutable-ok: mirrors the parent DenseEncoder contract + return super().encode_documents(list(docs)) + + async def aencode_documents( # pyright: ignore[reportIncompatibleMethodOverride] # semantic_router calls aencode_documents without kwargs + self, docs: Sequence[str] + ) -> list[list[float]]: # mutable-ok: mirrors the parent DenseEncoder contract + unique_docs: Final = tuple(dict.fromkeys(docs)) + keys: Final = tuple(tool_vector_key(self._embedding_identity, text) for text in unique_docs) + hits: Final[Mapping[str, tuple[float, ...]]] = await self._tool_vector_store.get_many(keys) + valid_hits: Final = ( + {key: vector for key, vector in hits.items() if len(vector) == await self.adimensions()} if hits else {} + ) + + miss_pairs: Final = tuple((text, key) for text, key in zip(unique_docs, keys) if key not in valid_hits) + embedded: Final = dict(valid_hits) + if miss_pairs: + vectors: Final = await super().aencode_documents([text for text, _ in miss_pairs]) + new_vectors: Final = {key: tuple(vector) for (_, key), vector in zip(miss_pairs, vectors)} + await self._tool_vector_store.put_many(new_vectors) + embedded.update(new_vectors) + + vectors_by_text: Final = dict(zip(unique_docs, (embedded[key] for key in keys))) + return [list(vectors_by_text[text]) for text in docs] diff --git a/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py b/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py index 4409c60cb0a..2141d95ed21 100644 --- a/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py +++ b/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py @@ -17,7 +17,6 @@ from litellm.litellm_core_utils.exception_mapping_utils import ExceptionCheckers from litellm.proxy._experimental.mcp_server.catalog import catalog_operation, global_manager from litellm.proxy._experimental.mcp_server.faults import iter_exception_tree from litellm.proxy._experimental.mcp_server.semantic_tool_index_store import ( - CachingToolEncoder, ToolVectorStore, tool_vector_key, ) @@ -69,6 +68,7 @@ class SemanticMCPToolFilter: enabled: bool = True, vector_store: ToolVectorStore | None = None, lock_poll_interval_s: float = 5.0, + lock_wait_timeout_s: float = 300.0, ): """ Initialize the semantic tool filter. @@ -84,6 +84,8 @@ class SemanticMCPToolFilter: until startup_index_ready flips lock_poll_interval_s: How often to re-check the store while another pod holds the build lock + lock_wait_timeout_s: Give up waiting for another pod's vectors after + this long and embed locally (matches the build-lock TTL) """ self.enabled = enabled self.top_k = top_k @@ -96,6 +98,7 @@ class SemanticMCPToolFilter: self._index_sync_lock = asyncio.Lock() self._vector_store = vector_store self._lock_poll_interval_s = lock_poll_interval_s + self._lock_wait_timeout_s = lock_wait_timeout_s self.embedding_identity = self._compute_embedding_identity() self.startup_index_ready = vector_store is None @@ -109,7 +112,11 @@ class SemanticMCPToolFilter: for deployment in deployment_maps ) deployment_models: Final = sorted( - {str(model) for model in (p.get("model") for p in litellm_params) if model} + { + str(model) + ("@dimensions=" + str(dimensions) if dimensions is not None else "") + for model, dimensions in ((p.get("model"), p.get("dimensions")) for p in litellm_params) + if model + } ) except (AttributeError, TypeError, KeyError, ValueError): return self.embedding_model @@ -119,6 +126,7 @@ class SemanticMCPToolFilter: def _new_encoder(self) -> "LiteLLMRouterEncoder": """Encoder used for every index build; cached in Redis when a store is set.""" + from litellm.proxy._experimental.mcp_server.semantic_tool_encoder import CachingToolEncoder from litellm.router_strategy.auto_router.litellm_encoder import ( LiteLLMRouterEncoder, ) @@ -265,6 +273,8 @@ class SemanticMCPToolFilter: """ from semantic_router.routers import SemanticRouter + from litellm.proxy._experimental.mcp_server.semantic_tool_encoder import CachingToolEncoder + if not tools: self.tool_router = None return @@ -283,8 +293,12 @@ class SemanticMCPToolFilter: # aadd's lazy init probes index dimensions through the sync encoder, # so dimensions are seeded with an async probe to keep the build async. - dims_probe: Final = await encoder.aencode_queries(["test"]) # pyright: ignore[reportUnknownMemberType] # semantic_router ships no type stubs - router.index.dimensions = len(dims_probe[0]) # pyright: ignore[reportUnknownMemberType] # semantic_router ships no type stubs + dimensions: Final = ( + await encoder.adimensions() + if isinstance(encoder, CachingToolEncoder) + else len((await encoder.aencode_queries(["test"]))[0]) # pyright: ignore[reportUnknownMemberType] # semantic_router ships no type stubs + ) + router.index.dimensions = dimensions # pyright: ignore[reportUnknownMemberType] # semantic_router ships no type stubs await router.aadd(list(routes)) # pyright: ignore[reportUnknownMemberType] # semantic_router ships no type stubs self._tool_map = dict(tool_map) @@ -309,9 +323,17 @@ class SemanticMCPToolFilter: descriptions: Final = tuple(self.extract_tool_info(tool)[1] for tool in tools) keys_set: Final = frozenset(tool_vector_key(self.embedding_identity, text) for text in set(descriptions)) + loop: Final = asyncio.get_running_loop() + deadline: Final = loop.time() + self._lock_wait_timeout_s while keys_set - (await store.get_many(tuple(keys_set))).keys(): if await store.try_acquire_build_lock(): return True + if loop.time() >= deadline: + verbose_logger.warning( + "MCP semantic tool index wait exceeded %.1fs; embedding locally", + self._lock_wait_timeout_s, + ) + return False await asyncio.sleep(self._lock_poll_interval_s) return False 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 dc21e0c965b..f290fa87c20 100644 --- a/litellm/proxy/_experimental/mcp_server/semantic_tool_index_store.py +++ b/litellm/proxy/_experimental/mcp_server/semantic_tool_index_store.py @@ -17,11 +17,9 @@ from pydantic import TypeAdapter from litellm._internal_context import with_service_target from litellm._logging import verbose_logger from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager -from litellm.router_strategy.auto_router.litellm_encoder import LiteLLMRouterEncoder if TYPE_CHECKING: from litellm.caching.redis_cache import RedisCache - from litellm.router import Router _KEY_PREFIX: Final = "litellm:mcp_semantic_tool_index:v1:" _BUILD_LOCK_ID: Final = "mcp_semantic_tool_index_build" @@ -120,54 +118,3 @@ def _decode_vector(value: object) -> tuple[float, ...] | None: return tuple(vector) except Exception: # noqa: BLE001 # corrupt payloads are treated as a plain cache miss return None - - -class CachingToolEncoder(LiteLLMRouterEncoder): - """LiteLLMRouterEncoder that serves document embeddings from a ToolVectorStore. - - Queries are never cached; only add/aadd's document utterances go through - the store. Duplicate texts share one key and one embed. - """ - - _tool_vector_store: ToolVectorStore - _embedding_identity: str - - def __init__( - self, - *, - store: ToolVectorStore, - embedding_identity: str, - litellm_router_instance: "Router", - model_name: str, - score_threshold: float | None = None, - ) -> None: - super().__init__( - litellm_router_instance=litellm_router_instance, - model_name=model_name, - score_threshold=score_threshold, - ) - self._tool_vector_store = store - self._embedding_identity = embedding_identity - - def encode_documents( # pyright: ignore[reportIncompatibleMethodOverride] # semantic_router calls encode_documents without kwargs - self, docs: Sequence[str] - ) -> list[list[float]]: # mutable-ok: mirrors the parent DenseEncoder contract - return super().encode_documents(list(docs)) - - async def aencode_documents( # pyright: ignore[reportIncompatibleMethodOverride] # semantic_router calls aencode_documents without kwargs - self, docs: Sequence[str] - ) -> list[list[float]]: # mutable-ok: mirrors the parent DenseEncoder contract - unique_docs: Final = tuple(dict.fromkeys(docs)) - keys: Final = tuple(tool_vector_key(self._embedding_identity, text) for text in unique_docs) - hits: Final = await self._tool_vector_store.get_many(keys) - - miss_pairs: Final = tuple((text, key) for text, key in zip(unique_docs, keys) if key not in hits) - embedded: Final = dict(hits) - if miss_pairs: - vectors: Final = await super().aencode_documents([text for text, _ in miss_pairs]) - new_vectors: Final = {key: tuple(vector) for (_, key), vector in zip(miss_pairs, vectors)} - await self._tool_vector_store.put_many(new_vectors) - embedded.update(new_vectors) - - vectors_by_text: Final = dict(zip(unique_docs, (embedded[key] for key in keys))) - return [list(vectors_by_text[text]) for text in docs] diff --git a/tests/unit/proxy/_experimental/mcp_server/test_semantic_tool_encoder.py b/tests/unit/proxy/_experimental/mcp_server/test_semantic_tool_encoder.py new file mode 100644 index 00000000000..a5df27c2520 --- /dev/null +++ b/tests/unit/proxy/_experimental/mcp_server/test_semantic_tool_encoder.py @@ -0,0 +1,208 @@ +""" +Unit tests for CachingToolEncoder, the LiteLLMRouterEncoder that serves +document embeddings from the shared ToolVectorStore. +""" + +import asyncio +import sys +from unittest.mock import Mock + +import pytest + +from mcp.types import Tool as MCPTool + +requires_semantic_router = pytest.mark.skipif( + sys.version_info >= (3, 14), reason="The semantic-router extra excludes Python 3.14" +) + + +def _fake_embedding_router(recorded_inputs): + """Router stand-in returning deterministic 2-dim keyword one-hots, like the filter tests.""" + from litellm.types.utils import Embedding, EmbeddingResponse + + def _vector(text): + return [1.0, 0.0] if "ticket" in text.lower() else [0.0, 1.0] + + def mock_embedding_sync(*args, **kwargs): + texts = kwargs["input"] + recorded_inputs.append(list(texts)) + return EmbeddingResponse( + data=[Embedding(embedding=_vector(t), index=i, object="embedding") for i, t in enumerate(texts)], + model="text-embedding-3-small", + object="list", + usage={"prompt_tokens": 10, "total_tokens": 10}, + ) + + async def mock_embedding_async(*args, **kwargs): + return mock_embedding_sync(*args, **kwargs) + + router = Mock() + router.embedding = mock_embedding_sync + router.aembedding = mock_embedding_async + return router + + +class _InMemoryToolVectorStore: + def __init__(self, vectors=None): + self.vectors = dict(vectors or {}) + self.get_many_calls = 0 + self.acquire_calls = 0 + self.release_calls = 0 + self.lock_result = True + self._fill_on_call = None + + def fill_on_get_many_call(self, call_number): + self._fill_on_call = call_number + + async def get_many(self, keys): + self.get_many_calls += 1 + if self._fill_on_call is not None and self.get_many_calls >= self._fill_on_call: + self.vectors.update(self._pending_fill) + return {key: self.vectors[key] for key in keys if key in self.vectors} + + def seed_from(self, keys_vectors): + self._pending_fill = dict(keys_vectors) + + async def put_many(self, vectors): + self.vectors.update({k: tuple(v) for k, v in vectors.items()}) + + async def try_acquire_build_lock(self): + self.acquire_calls += 1 + return self.lock_result + + async def release_build_lock(self): + self.release_calls += 1 + + +def _tool(name, description): + return MCPTool(name=name, description=description, inputSchema={"type": "object"}) + + +def _make_filter(recorded_inputs, store=None, embedding_model="text-embedding-3-small"): + from litellm.proxy._experimental.mcp_server.semantic_tool_filter import ( + SemanticMCPToolFilter, + ) + + return SemanticMCPToolFilter( + embedding_model=embedding_model, + litellm_router_instance=_fake_embedding_router(recorded_inputs), + top_k=3, + similarity_threshold=0.3, + enabled=True, + vector_store=store, + lock_poll_interval_s=0.0, + ) + + +def _make_encoder(recorded_inputs, store, embedding_identity="text-embedding-3-small"): + from litellm.proxy._experimental.mcp_server.semantic_tool_encoder import CachingToolEncoder + + return CachingToolEncoder( + litellm_router_instance=_fake_embedding_router(recorded_inputs), + model_name="text-embedding-3-small", + score_threshold=0.3, + store=store, + embedding_identity=embedding_identity, + ) + + +@requires_semantic_router +@pytest.mark.asyncio +async def test_build_with_warm_store_never_embeds_descriptions(): + from litellm.proxy._experimental.mcp_server.semantic_tool_index_store import ( + tool_vector_key, + ) + + tools = [_tool("srv-a_get_ticket", "get a ticket"), _tool("srv-b_get_weather", "weather forecast")] + identity = "text-embedding-3-small" + warm = _InMemoryToolVectorStore( + { + tool_vector_key(identity, description): [0.0, 1.0] + for _, description in ((t.name, t.description) for t in tools) + } + ) + recorded = [] + filter_instance = _make_filter(recorded, store=warm) + + await filter_instance._abuild_router(tools) + + assert filter_instance.tool_router is not None + # Only the dims probe embeds; the 300-tool warm path makes zero embedding calls. + assert recorded == [["test"]] + assert warm.acquire_calls == 0 + + +@requires_semantic_router +@pytest.mark.asyncio +async def test_cold_store_embeds_once_then_second_filter_embeds_nothing(): + tools = [_tool("srv-a_get_ticket", "get a ticket"), _tool("srv-b_get_weather", "weather forecast")] + store = _InMemoryToolVectorStore() + + first_recorded = [] + first = _make_filter(first_recorded, store=store) + await first._abuild_router(tools) + assert first.tool_router is not None + assert [batch for batch in first_recorded if batch != ["test"]] == [["get a ticket", "weather forecast"]] + + second_recorded = [] + second = _make_filter(second_recorded, store=store) + await second._abuild_router(tools) + assert second.tool_router is not None + assert second_recorded == [["test"]] + + +@requires_semantic_router +@pytest.mark.asyncio +async def test_edited_description_reembeds_only_that_tool(): + tools = [_tool("srv-a_get_ticket", "get a ticket"), _tool("srv-b_get_weather", "weather forecast")] + store = _InMemoryToolVectorStore() + recorded = [] + filter_instance = _make_filter(recorded, store=store) + await filter_instance._abuild_router(tools) + + recorded.clear() + edited = [_tool("srv-a_get_ticket", "close a support ticket fast"), _tool("srv-b_get_weather", "weather forecast")] + await filter_instance._abuild_router(edited) + + assert [batch for batch in recorded if batch != ["test"]] == [["close a support ticket fast"]] + + +@requires_semantic_router +@pytest.mark.asyncio +async def test_different_embedding_identity_reembeds_everything(): + tools = [_tool("srv-a_get_ticket", "get a ticket"), _tool("srv-b_get_weather", "weather forecast")] + store = _InMemoryToolVectorStore() + + first_recorded = [] + await _make_filter(first_recorded, store=store, embedding_model="text-embedding-3-small")._abuild_router(tools) + + second_recorded = [] + second = _make_filter(second_recorded, store=store, embedding_model="text-embedding-3-large") + await second._abuild_router(tools) + + assert second.tool_router is not None + assert [batch for batch in second_recorded if batch != ["test"]] == [["get a ticket", "weather forecast"]] + + +@requires_semantic_router +@pytest.mark.asyncio +async def test_wrong_length_cached_vector_is_reembedded_and_overwritten(): + """ + A cached vector whose length disagrees with a live dimension probe is + treated as a miss: it is re-embedded, served at the right length, and the + store entry is overwritten so the stale shape cannot leak into the index. + """ + from litellm.proxy._experimental.mcp_server.semantic_tool_index_store import ( + tool_vector_key, + ) + + key = tool_vector_key("text-embedding-3-small", "get a ticket") + store = _InMemoryToolVectorStore({key: (9.9, 9.9, 9.9, 9.9)}) + recorded = [] + encoder = _make_encoder(recorded, store) + + vectors = await encoder.aencode_documents(["get a ticket"]) + + assert vectors == [[1.0, 0.0]] + assert [batch for batch in recorded if batch != ["test"]] == [["get a ticket"]] + assert store.vectors[key] == (1.0, 0.0) diff --git a/tests/unit/proxy/_experimental/mcp_server/test_semantic_tool_filter.py b/tests/unit/proxy/_experimental/mcp_server/test_semantic_tool_filter.py index 1f578dc99ad..7f8f491825b 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_semantic_tool_filter.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_semantic_tool_filter.py @@ -2415,3 +2415,35 @@ async def test_hook_returns_before_shared_index_build_finishes(): assert inline_hook.index_build_task is None assert inline_hook.filter.startup_index_ready assert inline_hook.filter.tool_router is not None + + +def test_embedding_identity_marks_deployment_dimensions(): + """ + The shared-vector key identity must change when a deployment's embedding + dimensions change, and stay byte-identical to the pre-dimensions format + when no deployment sets them (warm caches must keep hitting). + """ + from litellm.proxy._experimental.mcp_server.semantic_tool_filter import ( + SemanticMCPToolFilter, + ) + + router_without = Mock() + router_without.get_model_list = lambda model_name: [{"litellm_params": {"model": "amazon.titan-embed-text-v2:0"}}] + router_with = Mock() + router_with.get_model_list = lambda model_name: [ + {"litellm_params": {"model": "amazon.titan-embed-text-v2:0", "dimensions": 512}} + ] + + base = SemanticMCPToolFilter( + embedding_model="titan-alias", + litellm_router_instance=router_without, + enabled=True, + ) + sized = SemanticMCPToolFilter( + embedding_model="titan-alias", + litellm_router_instance=router_with, + enabled=True, + ) + + assert base.embedding_identity == "titan-alias|amazon.titan-embed-text-v2:0" + assert sized.embedding_identity == "titan-alias|amazon.titan-embed-text-v2:0@dimensions=512" 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 caf10d469e1..3c19fb4a1e4 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 @@ -7,6 +7,8 @@ embeddings from the store instead of re-embedding. """ import asyncio +import os +import subprocess import sys from unittest.mock import AsyncMock, Mock, patch @@ -119,7 +121,7 @@ def _tool(name, description): return MCPTool(name=name, description=description, inputSchema={"type": "object"}) -def _make_filter(recorded_inputs, store=None, embedding_model="text-embedding-3-small"): +def _make_filter(recorded_inputs, store=None, embedding_model="text-embedding-3-small", lock_wait_timeout_s=300.0): from litellm.proxy._experimental.mcp_server.semantic_tool_filter import ( SemanticMCPToolFilter, ) @@ -132,6 +134,7 @@ def _make_filter(recorded_inputs, store=None, embedding_model="text-embedding-3- enabled=True, vector_store=store, lock_poll_interval_s=0.0, + lock_wait_timeout_s=lock_wait_timeout_s, ) @@ -187,88 +190,7 @@ async def test_raising_redis_still_builds_the_index(): await filter_instance._abuild_router(tools) assert filter_instance.tool_router is not None - assert [batch for batch in recorded if batch != ["test"]] == [ - ["get a ticket", "weather forecast"] - ] - - -@requires_semantic_router -@pytest.mark.asyncio -async def test_build_with_warm_store_never_embeds_descriptions(): - from litellm.proxy._experimental.mcp_server.semantic_tool_index_store import ( - tool_vector_key, - ) - - tools = [_tool("srv-a_get_ticket", "get a ticket"), _tool("srv-b_get_weather", "weather forecast")] - identity = "text-embedding-3-small" - warm = _InMemoryToolVectorStore( - {tool_vector_key(identity, description): [0.0, 1.0] for _, description in ((t.name, t.description) for t in tools)} - ) - recorded = [] - filter_instance = _make_filter(recorded, store=warm) - - await filter_instance._abuild_router(tools) - - assert filter_instance.tool_router is not None - # Only the dims probe embeds; the 300-tool warm path makes zero embedding calls. - assert recorded == [["test"]] - assert warm.acquire_calls == 0 - - -@requires_semantic_router -@pytest.mark.asyncio -async def test_cold_store_embeds_once_then_second_filter_embeds_nothing(): - tools = [_tool("srv-a_get_ticket", "get a ticket"), _tool("srv-b_get_weather", "weather forecast")] - store = _InMemoryToolVectorStore() - - first_recorded = [] - first = _make_filter(first_recorded, store=store) - await first._abuild_router(tools) - assert first.tool_router is not None - assert [batch for batch in first_recorded if batch != ["test"]] == [ - ["get a ticket", "weather forecast"] - ] - - second_recorded = [] - second = _make_filter(second_recorded, store=store) - await second._abuild_router(tools) - assert second.tool_router is not None - assert second_recorded == [["test"]] - - -@requires_semantic_router -@pytest.mark.asyncio -async def test_edited_description_reembeds_only_that_tool(): - tools = [_tool("srv-a_get_ticket", "get a ticket"), _tool("srv-b_get_weather", "weather forecast")] - store = _InMemoryToolVectorStore() - recorded = [] - filter_instance = _make_filter(recorded, store=store) - await filter_instance._abuild_router(tools) - - recorded.clear() - edited = [_tool("srv-a_get_ticket", "close a support ticket fast"), _tool("srv-b_get_weather", "weather forecast")] - await filter_instance._abuild_router(edited) - - assert [batch for batch in recorded if batch != ["test"]] == [["close a support ticket fast"]] - - -@requires_semantic_router -@pytest.mark.asyncio -async def test_different_embedding_identity_reembeds_everything(): - tools = [_tool("srv-a_get_ticket", "get a ticket"), _tool("srv-b_get_weather", "weather forecast")] - store = _InMemoryToolVectorStore() - - first_recorded = [] - await _make_filter(first_recorded, store=store, embedding_model="text-embedding-3-small")._abuild_router(tools) - - second_recorded = [] - second = _make_filter(second_recorded, store=store, embedding_model="text-embedding-3-large") - await second._abuild_router(tools) - - assert second.tool_router is not None - assert [batch for batch in second_recorded if batch != ["test"]] == [ - ["get a ticket", "weather forecast"] - ] + assert [batch for batch in recorded if batch != ["test"]] == [["get a ticket", "weather forecast"]] @requires_semantic_router @@ -309,9 +231,7 @@ async def test_full_build_reads_the_store_a_constant_number_of_times(): 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) - ), + patch.object(global_mcp_server_manager, "get_tools_for_server", new=AsyncMock(return_value=tools)), ): await filter_instance.build_router_from_mcp_registry() @@ -337,3 +257,60 @@ async def test_cancelling_a_build_while_holding_the_lock_releases_it(): with pytest.raises(asyncio.CancelledError): await asyncio.wait_for(build_task, timeout=5) assert store.release_calls == 1 + + +@requires_semantic_router +@pytest.mark.asyncio +async def test_lock_wait_times_out_and_the_build_falls_back_to_local_embed(): + """ + A store that never fills and never grants the lock must not wait forever: + after lock_wait_timeout_s the build embeds locally, the startup index + flips ready, and the filter still selects tools. + """ + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + tools = [_tool("srv-a_get_ticket", "get a ticket"), _tool("srv-b_get_weather", "weather forecast")] + store = _InMemoryToolVectorStore() + store.lock_result = False + recorded = [] + filter_instance = _make_filter(recorded, store=store, lock_wait_timeout_s=0.05) + filter_instance._lock_poll_interval_s = 0.01 + + 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.startup_index_ready is True + assert filter_instance.tool_router is not None + assert store.release_calls == 0 + assert [batch for batch in recorded if batch != ["test"]] == [["get a ticket", "weather forecast"]] + filtered = await filter_instance.filter_tools( + query="refund a ticket", available_tools=tools, reindex_changed_descriptions=True + ) + assert filtered == [tools[0]] + + +def test_store_and_hook_import_without_semantic_router(): + """semantic_router is an optional extra; the store and the hook must import when it is absent.""" + root = os.path.dirname( + os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))) + ) + result = subprocess.run( + [ + sys.executable, + "-I", + "-c", + "import sys; sys.path.insert(0, " + repr(root) + "); " + "sys.modules['semantic_router'] = None; " + "import litellm.proxy._experimental.mcp_server.semantic_tool_index_store; " + "import litellm.proxy.hooks.mcp_semantic_filter", + ], + capture_output=True, + text=True, + timeout=120, + ) + assert result.returncode == 0, result.stderr From 6bc9d7820c310ed918b3a05e39941458d21be587 Mon Sep 17 00:00:00 2001 From: mrinal Date: Sun, 11 Oct 2026 04:05:37 +0000 Subject: [PATCH 7/9] fix(mcp): retry failed semantic index builds and key embeddings by endpoint Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/semantic_tool_filter.py | 11 +- litellm/proxy/proxy_server.py | 2 +- .../mcp/test_mcp_semantic_filter_startup.py | 121 ++++++++---------- .../mcp_server/test_semantic_tool_filter.py | 19 +++ .../proxy/proxy_server/test_proxy_config.py | 46 ++++++- 5 files changed, 124 insertions(+), 75 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py b/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py index 2141d95ed21..6b2b93f2626 100644 --- a/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py +++ b/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py @@ -56,6 +56,13 @@ def _is_context_window_error(error: BaseException | None) -> bool: ) +def _embedding_params_suffix(litellm_params: Mapping[str, object]) -> str: + """Deployment traits that make cached vectors non-interchangeable across endpoints.""" + return (f"@dimensions={litellm_params['dimensions']}" if litellm_params.get("dimensions") is not None else "") + ( + f"@api_base={litellm_params['api_base']}" if litellm_params.get("api_base") is not None else "" + ) + + class SemanticMCPToolFilter: """Filters MCP tools using semantic similarity to reduce context window size.""" @@ -113,8 +120,8 @@ class SemanticMCPToolFilter: ) deployment_models: Final = sorted( { - str(model) + ("@dimensions=" + str(dimensions) if dimensions is not None else "") - for model, dimensions in ((p.get("model"), p.get("dimensions")) for p in litellm_params) + str(model) + _embedding_params_suffix(params) + for model, params in ((p.get("model"), p) for p in litellm_params) if model } ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 223aced9be5..5494aef3faf 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -8545,7 +8545,7 @@ class ProxyConfig: 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: + if active_hook.index_build_task is not None and not active_hook.index_build_task.done(): verbose_proxy_logger.debug( "Semantic filter settings unchanged, index build already started; skipping reinitialization" ) diff --git a/tests/integration/mcp/test_mcp_semantic_filter_startup.py b/tests/integration/mcp/test_mcp_semantic_filter_startup.py index b23e85c88cc..64ed44c6a90 100644 --- a/tests/integration/mcp/test_mcp_semantic_filter_startup.py +++ b/tests/integration/mcp/test_mcp_semantic_filter_startup.py @@ -338,17 +338,19 @@ def test_second_proxy_on_warm_redis_embeds_no_descriptions(gateway: Gateway, tmp alias: Final = "sf" + uuid.uuid4().hex[:8] descriptions: Final = _descriptions() with owned_redis(tmp_path) as redis, scripted_peer(*_tools()) as peer: - first_embeddings: Final = EmbeddingDouble("healthy") + embeddings: Final = EmbeddingDouble("healthy") respond, bodies = _chat_double() + second_respond, second_bodies = _chat_double() with ( wire_server(respond) as chat_wire, - wire_server(first_embeddings.respond) as first_embedding_wire, + wire_server(second_respond) as second_chat_wire, + wire_server(embeddings.respond) as embedding_wire, ): config: Final = _config( tmp_path, alias=alias, peer_url=peer.url, - embeddings_url=first_embedding_wire.url, + embeddings_url=embedding_wire.url, chat_url=chat_wire.url, redis=redis, ) @@ -358,30 +360,24 @@ def test_second_proxy_on_warm_redis_embeds_no_descriptions(gateway: Gateway, tmp lambda filtered: filtered, seconds=60, ) - second_embeddings: Final = EmbeddingDouble("healthy") - second_respond, second_bodies = _chat_double() - with ( - wire_server(second_respond) as second_chat_wire, - wire_server(second_embeddings.respond) as second_embedding_wire, - ): + embedded_so_far: Final = embeddings.document_inputs(descriptions) config = _config( tmp_path, alias=alias, peer_url=peer.url, - embeddings_url=second_embedding_wire.url, + embeddings_url=embedding_wire.url, chat_url=second_chat_wire.url, redis=redis, ) with owned_proxy_process(gateway, tmp_path, {}, config=config) as second_proxy: - assert second_embeddings.document_inputs(descriptions) == [] assert _request_is_filtered(second_proxy.gateway, "chat", second_bodies, alias) + assert embeddings.document_inputs(descriptions) == embedded_so_far def test_two_proxies_cold_starting_together_embed_each_description_once(gateway: Gateway, tmp_path: Path) -> None: alias: Final = "sf" + uuid.uuid4().hex[:8] descriptions: Final = _descriptions() - first_embeddings: Final = EmbeddingDouble("healthy", delay=10.0) - second_embeddings: Final = EmbeddingDouble("healthy") + embeddings: Final = EmbeddingDouble("healthy", delay=10.0) respond_a, bodies_a = _chat_double() respond_b, bodies_b = _chat_double() with ( @@ -389,14 +385,13 @@ def test_two_proxies_cold_starting_together_embed_each_description_once(gateway: scripted_peer(*_tools()) as peer, wire_server(respond_a) as chat_wire_a, wire_server(respond_b) as chat_wire_b, - wire_server(first_embeddings.respond) as embedding_wire_a, - wire_server(second_embeddings.respond) as embedding_wire_b, + wire_server(embeddings.respond) as embedding_wire, ): config_a: Final = _config( tmp_path, alias=alias, peer_url=peer.url, - embeddings_url=embedding_wire_a.url, + embeddings_url=embedding_wire.url, chat_url=chat_wire_a.url, redis=redis, ) @@ -404,7 +399,7 @@ def test_two_proxies_cold_starting_together_embed_each_description_once(gateway: tmp_path, alias=alias, peer_url=peer.url, - embeddings_url=embedding_wire_b.url, + embeddings_url=embedding_wire.url, chat_url=chat_wire_b.url, redis=redis, ) @@ -422,9 +417,7 @@ def test_two_proxies_cold_starting_together_embed_each_description_once(gateway: lambda filtered: filtered, seconds=60, ) - document_inputs: Final = Counter( - first_embeddings.document_inputs(descriptions) + second_embeddings.document_inputs(descriptions) - ) + document_inputs: Final = Counter(embeddings.document_inputs(descriptions)) assert document_inputs == Counter({description: 1 for description in descriptions}) @@ -433,56 +426,50 @@ def test_edited_description_reembeds_only_that_tool(gateway: Gateway, tmp_path: descriptions_v1: Final = _descriptions() edited: Final = "Cancel and refund an invoice instantly" with owned_redis(tmp_path) as redis: - first_embeddings: Final = EmbeddingDouble("healthy") + embeddings: Final = EmbeddingDouble("healthy") respond, bodies = _chat_double() - with ( - scripted_peer(*_tools()) as peer_v1, - wire_server(respond) as chat_wire, - wire_server(first_embeddings.respond) as first_embedding_wire, - ): - config: Final = _config( - tmp_path, - alias=alias, - peer_url=peer_v1.url, - embeddings_url=first_embedding_wire.url, - chat_url=chat_wire.url, - redis=redis, - ) - with owned_proxy_process(gateway, tmp_path, {}, config=config) as first_proxy: - eventually( - lambda: _request_is_filtered(first_proxy.gateway, "chat", bodies, alias), - lambda filtered: filtered, - seconds=60, - ) - second_embeddings: Final = EmbeddingDouble("healthy") second_respond, second_bodies = _chat_double() - with ( - scripted_peer(*_tools(edited)) as peer_v2, - wire_server(second_respond) as second_chat_wire, - wire_server(second_embeddings.respond) as second_embedding_wire, - ): - config = _config( - tmp_path, - alias=alias, - peer_url=peer_v2.url, - embeddings_url=second_embedding_wire.url, - chat_url=second_chat_wire.url, - redis=redis, - ) - with owned_proxy_process(gateway, tmp_path, {}, config=config) as second_proxy: - assert second_embeddings.document_inputs(descriptions_v1) == [] - document_inputs: Final = [ - text - for inputs, _ in second_embeddings.arrived - for text in inputs - if text not in descriptions_v1 and text not in ("test", PROMPT) - ] - assert document_inputs == [edited] - eventually( - lambda: _request_is_filtered(second_proxy.gateway, "chat", second_bodies, alias), - lambda filtered: filtered, - seconds=60, + with wire_server(embeddings.respond) as embedding_wire: + with ( + scripted_peer(*_tools()) as peer_v1, + wire_server(respond) as chat_wire, + ): + config: Final = _config( + tmp_path, + alias=alias, + peer_url=peer_v1.url, + embeddings_url=embedding_wire.url, + chat_url=chat_wire.url, + redis=redis, ) + with owned_proxy_process(gateway, tmp_path, {}, config=config) as first_proxy: + eventually( + lambda: _request_is_filtered(first_proxy.gateway, "chat", bodies, alias), + lambda filtered: filtered, + seconds=60, + ) + with ( + scripted_peer(*_tools(edited)) as peer_v2, + wire_server(second_respond) as second_chat_wire, + ): + config = _config( + tmp_path, + alias=alias, + peer_url=peer_v2.url, + embeddings_url=embedding_wire.url, + chat_url=second_chat_wire.url, + redis=redis, + ) + with owned_proxy_process(gateway, tmp_path, {}, config=config) as second_proxy: + eventually( + lambda: _request_is_filtered(second_proxy.gateway, "chat", second_bodies, alias), + lambda filtered: filtered, + seconds=60, + ) + document_inputs: Final = Counter( + text for inputs, _ in embeddings.arrived for text in inputs if text not in ("test", PROMPT) + ) + assert document_inputs == Counter({**{d: 1 for d in descriptions_v1}, edited: 1}) def test_without_redis_startup_waits_for_the_index(gateway: Gateway, tmp_path: Path) -> None: diff --git a/tests/unit/proxy/_experimental/mcp_server/test_semantic_tool_filter.py b/tests/unit/proxy/_experimental/mcp_server/test_semantic_tool_filter.py index 7f8f491825b..ef593481c27 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_semantic_tool_filter.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_semantic_tool_filter.py @@ -2433,6 +2433,16 @@ def test_embedding_identity_marks_deployment_dimensions(): router_with.get_model_list = lambda model_name: [ {"litellm_params": {"model": "amazon.titan-embed-text-v2:0", "dimensions": 512}} ] + router_with_api_base = Mock() + router_with_api_base.get_model_list = lambda model_name: [ + { + "litellm_params": { + "model": "amazon.titan-embed-text-v2:0", + "dimensions": 512, + "api_base": "https://bedrock-endpoint.example.com", + } + } + ] base = SemanticMCPToolFilter( embedding_model="titan-alias", @@ -2444,6 +2454,15 @@ def test_embedding_identity_marks_deployment_dimensions(): litellm_router_instance=router_with, enabled=True, ) + based = SemanticMCPToolFilter( + embedding_model="titan-alias", + litellm_router_instance=router_with_api_base, + enabled=True, + ) assert base.embedding_identity == "titan-alias|amazon.titan-embed-text-v2:0" assert sized.embedding_identity == "titan-alias|amazon.titan-embed-text-v2:0@dimensions=512" + assert ( + based.embedding_identity + == "titan-alias|amazon.titan-embed-text-v2:0@dimensions=512@api_base=https://bedrock-endpoint.example.com" + ) diff --git a/tests/unit/proxy/proxy_server/test_proxy_config.py b/tests/unit/proxy/proxy_server/test_proxy_config.py index 069f2e1bf2b..d4f9e6dd6e0 100644 --- a/tests/unit/proxy/proxy_server/test_proxy_config.py +++ b/tests/unit/proxy/proxy_server/test_proxy_config.py @@ -96,7 +96,9 @@ def test_proxy_config_upserts_existing_deployment_without_duplicating_it(monkeyp @pytest.mark.asyncio -async def test_proxy_config_preserves_models_when_config_read_fails(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: +async def test_proxy_config_preserves_models_when_config_read_fails( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: from litellm.proxy import proxy_server router = litellm.Router(model_list=[_deployment("config-deployment").to_json(exclude_none=True)]) @@ -129,6 +131,8 @@ async def test_proxy_config_keeps_database_models_and_deletes_stale_models( assert desired == frozenset({"kept-deployment"}) assert router.get_model_ids() == ["kept-deployment"] + + from tests._master_key import MASTER_KEY @@ -144,7 +148,6 @@ async def test_proxy_config_loads_lens_store_from_yaml(tmp_path, monkeypatch) -> assert is_lens_tracing_enabled(settings["tracing"], {}) is True - @pytest.mark.asyncio @pytest.mark.parametrize("shutdown_error", [False, True]) async def test_tracing_config_automatically_exports_spend_without_a_storage_dependency( @@ -170,7 +173,9 @@ async def test_tracing_config_automatically_exports_spend_without_a_storage_depe outcome: Final = pytest.raises(RuntimeError, match="shutdown failure") if shutdown_error else nullcontext() with outcome: async with manage_tracing(enabled=True, client_factory=client): - logger: Final = next(callback for callback in litellm._async_success_callback if isinstance(callback, LensExporter)) + logger: Final = next( + callback for callback in litellm._async_success_callback if isinstance(callback, LensExporter) + ) await logger.async_log_success_event( {"standard_logging_object": {"id": "response-1", "response_cost": 0.25}}, None, None, None ) @@ -186,7 +191,6 @@ async def test_tracing_config_automatically_exports_spend_without_a_storage_depe assert logger.task is not None and logger.task.done() and not logger.task.cancelled() - # --------------------------------------------------------------------------- # _is_remote_module_url # --------------------------------------------------------------------------- @@ -4835,7 +4839,9 @@ async def test_ProxyConfig__update_config_from_db_keeps_keys_the_config_file_omi @pytest.mark.asyncio @pytest.mark.parametrize("ui_settings_already_synced", [False, True]) -async def test_ProxyConfig_add_deployment_continues_after_null_pass_through_endpoints(monkeypatch, ui_settings_already_synced): +async def test_ProxyConfig_add_deployment_continues_after_null_pass_through_endpoints( + monkeypatch, ui_settings_already_synced +): from litellm.proxy import proxy_server pc = ProxyConfig() @@ -5642,6 +5648,36 @@ async def test_proxy_config_semantic_filter_db_sync_skips_reinit_while_index_bui await build_task +@pytest.mark.asyncio +async def test_proxy_config_semantic_filter_db_sync_reinits_after_finished_failed_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, + } + build_task: Final = asyncio.create_task(asyncio.sleep(0)) + await build_task + 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": 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 + + await pc._init_semantic_filter_settings_in_db(prisma_client=MagicMock()) + + removed.assert_called_once() + + @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 From eb0ef31fd689273629a0e49dc30c9fa089b733c6 Mon Sep 17 00:00:00 2001 From: mrinal Date: Sun, 11 Oct 2026 18:04:06 +0000 Subject: [PATCH 8/9] fix(mcp): embed locally when the semantic index lock write fails Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/semantic_tool_index_store.py | 15 +++++- .../test_semantic_tool_index_store.py | 50 +++++++++++++++++++ 2 files changed, 64 insertions(+), 1 deletion(-) 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 f290fa87c20..4a85fe3a023 100644 --- a/litellm/proxy/_experimental/mcp_server/semantic_tool_index_store.py +++ b/litellm/proxy/_experimental/mcp_server/semantic_tool_index_store.py @@ -86,11 +86,24 @@ class RedisToolVectorStore: acquired: Final = await self._lock_manager.acquire_lock( _BUILD_LOCK_ID, ttl=_BUILD_LOCK_TTL_S, allow_reentrant=True ) - return bool(acquired) or self._degraded + held_by: Final = ( + None + if acquired or self._degraded + else await self._redis_cache.async_get_cache( + PodLockManager.get_redis_lock_key(_BUILD_LOCK_ID) # pyright: ignore[reportUnknownArgumentType] # RedisCache declares a partially untyped signature + ) + ) except Exception as e: # noqa: BLE001 # any redis failure must degrade to a local embed, never raise self._degraded = True verbose_logger.warning("MCP semantic tool index store lock acquire failed; embedding locally: %s", e) return True + if acquired or self._degraded: + return True + if held_by is not None: + return False + self._degraded = True + verbose_logger.warning("MCP semantic tool index store lock write failed with no holder; embedding locally") + return True async def release_build_lock(self) -> None: try: 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 3c19fb4a1e4..b53bab95e9d 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 @@ -104,6 +104,25 @@ class _RaisingRedisCache: raise ConnectionError("redis is down") +class _ReadOnlyRedisCache: + """Reads work; every write reports failure like a read-only replica.""" + + def __init__(self): + self.backing = {} + + async def async_batch_get_cache(self, key_list, parent_otel_span=None): + return {key: self.backing[key] for key in key_list if key in self.backing} + + async def async_set_cache_pipeline(self, cache_list, ttl=None, **kwargs): + raise ConnectionError("NOPERM read only replica") + + async def async_get_cache(self, key, parent_otel_span=None, **kwargs): + return self.backing.get(key) + + async def async_set_cache(self, key, value, **kwargs): + return False + + class _DictRedisCache: def __init__(self): self.backing = {} @@ -294,6 +313,37 @@ async def test_lock_wait_times_out_and_the_build_falls_back_to_local_embed(): assert filtered == [tools[0]] +@requires_semantic_router +@pytest.mark.asyncio +async def test_lock_acquire_embeds_locally_when_writes_fail_and_no_pod_holds_the_lock(): + from litellm.proxy._experimental.mcp_server.semantic_tool_index_store import ( + RedisToolVectorStore, + ) + from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager + + redis = _ReadOnlyRedisCache() + store = RedisToolVectorStore(redis, PodLockManager(redis)) + + assert await store.try_acquire_build_lock() is True + + +@requires_semantic_router +@pytest.mark.asyncio +async def test_lock_acquire_returns_false_when_another_pod_holds_the_lock(): + from litellm.proxy._experimental.mcp_server.semantic_tool_index_store import ( + RedisToolVectorStore, + ) + from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager + + redis = _ReadOnlyRedisCache() + lock_manager = PodLockManager(redis) + redis.backing[PodLockManager.get_redis_lock_key("mcp_semantic_tool_index_build")] = "other-pod" + store = RedisToolVectorStore(redis, lock_manager) + + assert await store.try_acquire_build_lock() is False + assert store._degraded is False + + def test_store_and_hook_import_without_semantic_router(): """semantic_router is an optional extra; the store and the hook must import when it is absent.""" root = os.path.dirname( From 4da2587912387fc63a60583f2cbb5b4d3f55a519 Mon Sep 17 00:00:00 2001 From: mrinal Date: Sun, 11 Oct 2026 18:13:36 +0000 Subject: [PATCH 9/9] fix(mcp): declare the redis key family on the semantic index lock check Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/_experimental/mcp_server/semantic_tool_index_store.py | 1 + 1 file changed, 1 insertion(+) 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 4a85fe3a023..212bee108f8 100644 --- a/litellm/proxy/_experimental/mcp_server/semantic_tool_index_store.py +++ b/litellm/proxy/_experimental/mcp_server/semantic_tool_index_store.py @@ -80,6 +80,7 @@ class RedisToolVectorStore: self._degraded = True verbose_logger.warning("MCP semantic tool index store put_many failed; vectors stay local: %s", e) + @with_service_target(_CACHE_TARGET) async def try_acquire_build_lock(self) -> bool: # PodLockManager returns False on Redis errors too; once degraded, embed locally instead of waiting. try: