From 66ce89fc11d3f99cf8206a5e93c2ca1303edd771 Mon Sep 17 00:00:00 2001 From: mrinal Date: Sun, 11 Oct 2026 03:28:31 +0000 Subject: [PATCH] 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