mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
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>
This commit is contained in:
parent
6bc9d7820c
commit
eb0ef31fd6
2 changed files with 64 additions and 1 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue