mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-02 07:13:02 +00:00
refactor(redis): move MCP server manager caches behind a typed store
This commit is contained in:
Devin AI
parent
8907a63d78
commit
ffae582a50
5 changed files with 402 additions and 52 deletions
|
|
@ -172,6 +172,15 @@ from litellm.proxy.mcp_server.result_conversion import (
|
|||
from litellm.proxy.mcp_server.sampling_handler import (
|
||||
MCP_SAMPLING_AVAILABLE,
|
||||
)
|
||||
from litellm.proxy.mcp_server.server_catalog_cache import (
|
||||
get_submitted_server_ids,
|
||||
get_toolset_by_name,
|
||||
get_toolset_permissions,
|
||||
invalidate_submitted_server_ids,
|
||||
put_submitted_server_ids,
|
||||
put_toolset_by_name,
|
||||
put_toolset_permissions,
|
||||
)
|
||||
from litellm.proxy.mcp_server.stdio_gate import (
|
||||
MCP_STDIO_DISABLED_MESSAGE,
|
||||
is_mcp_stdio_blocked,
|
||||
|
|
@ -223,6 +232,7 @@ from litellm.proxy.middleware.per_request_root_path_middleware import (
|
|||
get_request_root_path,
|
||||
)
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
from litellm.redis.stores.mcp_server_catalog import TOOLSET_NOT_FOUND, submitted_servers_key
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.integrations.slack_alerting import AlertType
|
||||
from litellm.types.llms.base import LiteLLMBaseModel
|
||||
|
|
@ -3566,7 +3576,7 @@ class MCPServerManager:
|
|||
|
||||
@staticmethod
|
||||
def get_byom_submitted_servers_cache_key(user_id: str) -> str:
|
||||
return f"byom_submitted_servers:{user_id}"
|
||||
return submitted_servers_key(user_id)
|
||||
|
||||
@with_service_target(MCP_SERVERS_TARGET)
|
||||
async def invalidate_byom_submitted_servers_cache(self, user_id: str | None) -> None:
|
||||
|
|
@ -3575,7 +3585,7 @@ class MCPServerManager:
|
|||
try:
|
||||
from litellm.proxy.proxy_server import user_api_key_cache
|
||||
|
||||
await user_api_key_cache.async_delete_cache(key=self.get_byom_submitted_servers_cache_key(user_id))
|
||||
await invalidate_submitted_server_ids(user_api_key_cache, user_id)
|
||||
except Exception as e: # noqa: BLE001
|
||||
verbose_logger.warning("Failed to invalidate BYOM submitted MCP server cache: %s", e)
|
||||
|
||||
|
|
@ -3596,14 +3606,11 @@ class MCPServerManager:
|
|||
verbose_logger.warning("Failed to load BYOM submitted MCP server cache dependencies: %s", e)
|
||||
return []
|
||||
|
||||
byom_cache_key: Final = self.get_byom_submitted_servers_cache_key(submitter_user_id)
|
||||
submitted_server_ids: list[str] | None = None
|
||||
try:
|
||||
cached_submitted_server_ids: Final = _UNTYPED_VALUE.validate_python(
|
||||
await user_api_key_cache.async_get_cache(key=byom_cache_key)
|
||||
)
|
||||
cached_submitted_server_ids: Final = await get_submitted_server_ids(user_api_key_cache, submitter_user_id)
|
||||
if cached_submitted_server_ids is not None:
|
||||
submitted_server_ids = cast(list[str], cached_submitted_server_ids)
|
||||
submitted_server_ids = list(cached_submitted_server_ids)
|
||||
except Exception as e: # noqa: BLE001
|
||||
verbose_logger.warning("Failed to read BYOM submitted MCP server cache: %s", e)
|
||||
|
||||
|
|
@ -3619,11 +3626,7 @@ class MCPServerManager:
|
|||
verbose_logger.warning("Failed to read BYOM submitted MCP servers from database: %s", e)
|
||||
submitted_server_ids = []
|
||||
try:
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=byom_cache_key,
|
||||
value=submitted_server_ids,
|
||||
ttl=60,
|
||||
)
|
||||
await put_submitted_server_ids(user_api_key_cache, submitter_user_id, submitted_server_ids, ttl=60)
|
||||
except Exception as e: # noqa: BLE001
|
||||
verbose_logger.warning("Failed to write BYOM submitted MCP server cache: %s", e)
|
||||
|
||||
|
|
@ -3851,12 +3854,11 @@ class MCPServerManager:
|
|||
if not toolset_ids or prisma_client is None:
|
||||
return {}
|
||||
|
||||
cache_key: Final = "toolset_perms:" + ",".join(sorted(toolset_ids))
|
||||
cached: Final[dict[str, list[str]] | None] = (
|
||||
None if requires_fresh_policy else await user_api_key_cache.async_get_cache(key=cache_key)
|
||||
cached: Final = (
|
||||
None if requires_fresh_policy else await get_toolset_permissions(user_api_key_cache, toolset_ids)
|
||||
)
|
||||
if cached is not None:
|
||||
return cached
|
||||
return {server_id: list(tools) for server_id, tools in cached.items()}
|
||||
|
||||
try:
|
||||
toolsets: Final = await list_mcp_toolsets(
|
||||
|
|
@ -3868,10 +3870,8 @@ class MCPServerManager:
|
|||
allowed_names = tool_permissions.setdefault(tool["server_id"], [])
|
||||
if tool["tool_name"] not in allowed_names:
|
||||
allowed_names.append(tool["tool_name"])
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=cache_key,
|
||||
value=tool_permissions,
|
||||
ttl=get_management_object_ttl(user_api_key_cache),
|
||||
await put_toolset_permissions(
|
||||
user_api_key_cache, toolset_ids, tool_permissions, ttl=get_management_object_ttl(user_api_key_cache)
|
||||
)
|
||||
return tool_permissions
|
||||
except Exception as e:
|
||||
|
|
@ -3929,29 +3929,18 @@ class MCPServerManager:
|
|||
so callers can always use attribute access (e.g. ``toolset.toolset_id``).
|
||||
"""
|
||||
from litellm.proxy.proxy_server import user_api_key_cache
|
||||
from litellm.types.mcp_server.mcp_toolset import MCPToolset
|
||||
|
||||
cache_key: Final = f"toolset_name:{toolset_name}"
|
||||
cached: Final = await user_api_key_cache.async_get_cache(key=cache_key)
|
||||
cached: Final = await get_toolset_by_name(user_api_key_cache, toolset_name)
|
||||
if cached is not None:
|
||||
# Sentinel value used to cache "not found" so we don't re-query for
|
||||
# names that don't exist.
|
||||
if cached == "__not_found__":
|
||||
return None
|
||||
# Redis deserialises JSON back as a plain dict — reconstruct the model.
|
||||
if isinstance(cached, dict):
|
||||
return MCPToolset(**cached)
|
||||
return cached
|
||||
return None if cached == TOOLSET_NOT_FOUND else cached
|
||||
|
||||
from litellm.proxy.mcp_server.toolset_db import (
|
||||
get_mcp_toolset_by_name,
|
||||
)
|
||||
|
||||
toolset: Final = await get_mcp_toolset_by_name(prisma_client, toolset_name)
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=cache_key,
|
||||
value=(toolset.model_dump(mode="json") if toolset is not None else "__not_found__"),
|
||||
ttl=get_management_object_ttl(user_api_key_cache),
|
||||
await put_toolset_by_name(
|
||||
user_api_key_cache, toolset_name, toolset, ttl=get_management_object_ttl(user_api_key_cache)
|
||||
)
|
||||
return toolset
|
||||
|
||||
|
|
|
|||
158
litellm/proxy/mcp_server/server_catalog_cache.py
Normal file
158
litellm/proxy/mcp_server/server_catalog_cache.py
Normal file
|
|
@ -0,0 +1,158 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
||||
from contextlib import suppress
|
||||
from datetime import timedelta
|
||||
from typing import TYPE_CHECKING, Final, TypeVar
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm.caching.redis_batch import active_request_redis_batch
|
||||
from litellm.redis.stores.mcp_server_catalog import (
|
||||
CACHED_TOOLSET,
|
||||
SUBMITTED_SERVER_IDS,
|
||||
TOOLSET_NOT_FOUND,
|
||||
TOOLSET_PERMISSIONS,
|
||||
CachedToolset,
|
||||
McpServerCatalogStore,
|
||||
TransportMcpServerCatalogStore,
|
||||
submitted_servers_key,
|
||||
toolset_name_key,
|
||||
toolset_permissions_key,
|
||||
)
|
||||
from litellm.redis.transport.contract import RedisFailure
|
||||
from litellm.types.mcp_server.mcp_toolset import MCPToolset
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
_T = TypeVar("_T")
|
||||
_LOCAL_VALUE: Final = TypeAdapter(object)
|
||||
|
||||
|
||||
async def get_submitted_server_ids(cache: UserApiKeyCache, user_id: str) -> tuple[str, ...] | None:
|
||||
return await _read_through(
|
||||
cache,
|
||||
submitted_servers_key(user_id),
|
||||
SUBMITTED_SERVER_IDS,
|
||||
lambda store: store.get_submitted_server_ids(user_id),
|
||||
list,
|
||||
)
|
||||
|
||||
|
||||
async def put_submitted_server_ids(
|
||||
cache: UserApiKeyCache, user_id: str, server_ids: Sequence[str], *, ttl: float
|
||||
) -> None:
|
||||
await _write_through(
|
||||
cache,
|
||||
submitted_servers_key(user_id),
|
||||
list(server_ids),
|
||||
ttl,
|
||||
lambda store, expiry: store.put_submitted_server_ids(user_id, server_ids, ttl=expiry),
|
||||
)
|
||||
|
||||
|
||||
async def invalidate_submitted_server_ids(cache: UserApiKeyCache, user_id: str) -> None:
|
||||
cache.in_memory_cache.delete_cache(submitted_servers_key(user_id))
|
||||
if cache.redis_cache is None:
|
||||
return
|
||||
with suppress(Exception):
|
||||
_ = await _store(cache).invalidate_submitted_server_ids(user_id)
|
||||
|
||||
|
||||
async def get_toolset_permissions(
|
||||
cache: UserApiKeyCache, toolset_ids: Sequence[str]
|
||||
) -> dict[str, tuple[str, ...]] | None:
|
||||
return await _read_through(
|
||||
cache,
|
||||
toolset_permissions_key(toolset_ids),
|
||||
TOOLSET_PERMISSIONS,
|
||||
lambda store: store.get_toolset_permissions(toolset_ids),
|
||||
_permissions_value,
|
||||
)
|
||||
|
||||
|
||||
async def put_toolset_permissions(
|
||||
cache: UserApiKeyCache, toolset_ids: Sequence[str], permissions: Mapping[str, Sequence[str]], *, ttl: float
|
||||
) -> None:
|
||||
await _write_through(
|
||||
cache,
|
||||
toolset_permissions_key(toolset_ids),
|
||||
_permissions_value(permissions),
|
||||
ttl,
|
||||
lambda store, expiry: store.put_toolset_permissions(toolset_ids, permissions, ttl=expiry),
|
||||
)
|
||||
|
||||
|
||||
async def get_toolset_by_name(cache: UserApiKeyCache, toolset_name: str) -> CachedToolset | None:
|
||||
return await _read_through(
|
||||
cache,
|
||||
toolset_name_key(toolset_name),
|
||||
CACHED_TOOLSET,
|
||||
lambda store: store.get_toolset_by_name(toolset_name),
|
||||
_toolset_value,
|
||||
)
|
||||
|
||||
|
||||
async def put_toolset_by_name(
|
||||
cache: UserApiKeyCache, toolset_name: str, toolset: MCPToolset | None, *, ttl: float
|
||||
) -> None:
|
||||
await _write_through(
|
||||
cache,
|
||||
toolset_name_key(toolset_name),
|
||||
_toolset_value(toolset if toolset is not None else TOOLSET_NOT_FOUND),
|
||||
ttl,
|
||||
lambda store, expiry: store.put_toolset_by_name(toolset_name, toolset, ttl=expiry),
|
||||
)
|
||||
|
||||
|
||||
def _permissions_value(permissions: Mapping[str, Sequence[str]]) -> dict[str, list[str]]:
|
||||
return {server_id: list(tools) for server_id, tools in permissions.items()}
|
||||
|
||||
|
||||
def _toolset_value(toolset: CachedToolset) -> object:
|
||||
return toolset if isinstance(toolset, str) else toolset.model_dump(mode="json")
|
||||
|
||||
|
||||
def _store(cache: UserApiKeyCache) -> McpServerCatalogStore:
|
||||
assert cache.redis_cache is not None
|
||||
return TransportMcpServerCatalogStore(cache.redis_cache.transport())
|
||||
|
||||
|
||||
async def _read_through(
|
||||
cache: UserApiKeyCache,
|
||||
key: str,
|
||||
adapter: TypeAdapter[_T],
|
||||
fetch: Callable[[McpServerCatalogStore], Awaitable[_T | None | RedisFailure]],
|
||||
to_local: Callable[[_T], object],
|
||||
) -> _T | None:
|
||||
local: Final = _LOCAL_VALUE.validate_python(await cache.async_get_cache(key=key, local_only=True))
|
||||
if local is not None:
|
||||
with suppress(ValueError):
|
||||
return adapter.validate_python(local)
|
||||
if cache.redis_cache is None:
|
||||
return None
|
||||
batch: Final = active_request_redis_batch(cache.redis_cache)
|
||||
if batch is not None and batch.read_as_missing(key):
|
||||
return None
|
||||
with suppress(Exception):
|
||||
result: Final = await fetch(_store(cache))
|
||||
if result is None or isinstance(result, RedisFailure):
|
||||
return None
|
||||
await cache.async_set_cache(key=key, value=to_local(result), local_only=True)
|
||||
return result
|
||||
return None
|
||||
|
||||
|
||||
async def _write_through(
|
||||
cache: UserApiKeyCache,
|
||||
key: str,
|
||||
local_value: object,
|
||||
ttl: float,
|
||||
put: Callable[[McpServerCatalogStore, timedelta], Awaitable[bool | None | RedisFailure]],
|
||||
) -> None:
|
||||
await cache.async_set_cache(key=key, value=local_value, ttl=ttl, local_only=True)
|
||||
if cache.redis_cache is None:
|
||||
return
|
||||
with suppress(Exception):
|
||||
_ = await put(_store(cache), timedelta(seconds=int(ttl)))
|
||||
121
litellm/redis/stores/mcp_server_catalog.py
Normal file
121
litellm/redis/stores/mcp_server_catalog.py
Normal file
|
|
@ -0,0 +1,121 @@
|
|||
import ast
|
||||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import timedelta
|
||||
from typing import Final, Literal, Protocol, TypeVar
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from litellm.redis.transport.contract import RedisFailure, RedisTransport
|
||||
from litellm.types.mcp_server.mcp_toolset import MCPToolset
|
||||
|
||||
TOOLSET_NOT_FOUND: Final = "__not_found__"
|
||||
|
||||
ToolsetNotFound = Literal["__not_found__"]
|
||||
CachedToolset = MCPToolset | ToolsetNotFound
|
||||
|
||||
SUBMITTED_SERVER_IDS: Final = TypeAdapter(tuple[str, ...])
|
||||
TOOLSET_PERMISSIONS: Final = TypeAdapter(dict[str, tuple[str, ...]])
|
||||
CACHED_TOOLSET: Final[TypeAdapter[CachedToolset]] = TypeAdapter(CachedToolset)
|
||||
_VALUE: Final = TypeAdapter(object)
|
||||
|
||||
_T = TypeVar("_T")
|
||||
|
||||
|
||||
def submitted_servers_key(user_id: str) -> str:
|
||||
return f"byom_submitted_servers:{user_id}"
|
||||
|
||||
|
||||
def toolset_permissions_key(toolset_ids: Sequence[str]) -> str:
|
||||
return "toolset_perms:" + ",".join(sorted(toolset_ids))
|
||||
|
||||
|
||||
def toolset_name_key(toolset_name: str) -> str:
|
||||
return f"toolset_name:{toolset_name}"
|
||||
|
||||
|
||||
class McpServerCatalogStore(Protocol):
|
||||
async def get_submitted_server_ids(self, user_id: str) -> tuple[str, ...] | None | RedisFailure: ...
|
||||
|
||||
async def put_submitted_server_ids(
|
||||
self, user_id: str, server_ids: Sequence[str], *, ttl: timedelta
|
||||
) -> bool | None | RedisFailure: ...
|
||||
|
||||
async def invalidate_submitted_server_ids(self, user_id: str) -> int | RedisFailure: ...
|
||||
|
||||
async def get_toolset_permissions(
|
||||
self, toolset_ids: Sequence[str]
|
||||
) -> dict[str, tuple[str, ...]] | None | RedisFailure: ...
|
||||
|
||||
async def put_toolset_permissions(
|
||||
self, toolset_ids: Sequence[str], permissions: Mapping[str, Sequence[str]], *, ttl: timedelta
|
||||
) -> bool | None | RedisFailure: ...
|
||||
|
||||
async def get_toolset_by_name(self, toolset_name: str) -> CachedToolset | None | RedisFailure: ...
|
||||
|
||||
async def put_toolset_by_name(
|
||||
self, toolset_name: str, toolset: MCPToolset | None, *, ttl: timedelta
|
||||
) -> bool | None | RedisFailure: ...
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TransportMcpServerCatalogStore:
|
||||
transport: RedisTransport
|
||||
|
||||
async def get_submitted_server_ids(self, user_id: str) -> tuple[str, ...] | None | RedisFailure:
|
||||
return await self._get(submitted_servers_key(user_id), SUBMITTED_SERVER_IDS)
|
||||
|
||||
async def put_submitted_server_ids(
|
||||
self, user_id: str, server_ids: Sequence[str], *, ttl: timedelta
|
||||
) -> bool | None | RedisFailure:
|
||||
return await self._put(submitted_servers_key(user_id), json.dumps(list(server_ids)), ttl)
|
||||
|
||||
async def invalidate_submitted_server_ids(self, user_id: str) -> int | RedisFailure:
|
||||
return await self.transport.delete(submitted_servers_key(user_id))
|
||||
|
||||
async def get_toolset_permissions(
|
||||
self, toolset_ids: Sequence[str]
|
||||
) -> dict[str, tuple[str, ...]] | None | RedisFailure:
|
||||
return await self._get(toolset_permissions_key(toolset_ids), TOOLSET_PERMISSIONS)
|
||||
|
||||
async def put_toolset_permissions(
|
||||
self, toolset_ids: Sequence[str], permissions: Mapping[str, Sequence[str]], *, ttl: timedelta
|
||||
) -> bool | None | RedisFailure:
|
||||
payload: Final = json.dumps({server_id: list(tools) for server_id, tools in permissions.items()})
|
||||
return await self._put(toolset_permissions_key(toolset_ids), payload, ttl)
|
||||
|
||||
async def get_toolset_by_name(self, toolset_name: str) -> CachedToolset | None | RedisFailure:
|
||||
return await self._get(toolset_name_key(toolset_name), CACHED_TOOLSET)
|
||||
|
||||
async def put_toolset_by_name(
|
||||
self, toolset_name: str, toolset: MCPToolset | None, *, ttl: timedelta
|
||||
) -> bool | None | RedisFailure:
|
||||
payload: Final = json.dumps(toolset.model_dump(mode="json") if toolset is not None else TOOLSET_NOT_FOUND)
|
||||
return await self._put(toolset_name_key(toolset_name), payload, ttl)
|
||||
|
||||
async def _get(self, key: str, adapter: TypeAdapter[_T]) -> _T | None | RedisFailure:
|
||||
reply: Final = await self.transport.get(key)
|
||||
if reply is None or isinstance(reply, RedisFailure):
|
||||
return reply
|
||||
return _decode(reply, adapter)
|
||||
|
||||
async def _put(self, key: str, payload: str, ttl: timedelta) -> bool | None | RedisFailure:
|
||||
expiry: Final = timedelta(seconds=int(ttl.total_seconds()))
|
||||
batch: Final = self.transport.request_batch()
|
||||
if batch is not None:
|
||||
_ = batch.set_seconds(key, payload, ttl=expiry)
|
||||
return None
|
||||
return await self.transport.set_seconds(key, payload, ttl=expiry)
|
||||
|
||||
|
||||
def _decode(reply: bytes, adapter: TypeAdapter[_T]) -> _T | RedisFailure:
|
||||
try:
|
||||
text: Final = reply.decode("utf-8")
|
||||
try:
|
||||
value: Final = _VALUE.validate_python(json.loads(text))
|
||||
except ValueError:
|
||||
return adapter.validate_python(ast.literal_eval(text))
|
||||
return adapter.validate_python(value)
|
||||
except (UnicodeDecodeError, ValidationError, ValueError, SyntaxError) as error:
|
||||
return RedisFailure(kind="decode", message=f"Invalid MCP server catalog cache entry: {type(error).__name__}")
|
||||
|
|
@ -9032,26 +9032,41 @@ class TestMCPServerManager:
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalidate_byom_submitted_servers_cache_deletes_key(self):
|
||||
from litellm.proxy import proxy_server as proxy_server_module
|
||||
import fakeredis
|
||||
|
||||
cache = MagicMock()
|
||||
cache.async_delete_cache = AsyncMock()
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.proxy import proxy_server as proxy_server_module
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from tests._support.fake_redis import fake_redis_cache
|
||||
|
||||
server = fakeredis.FakeServer()
|
||||
raw = fakeredis.FakeAsyncRedis(server=server)
|
||||
cache = UserApiKeyCache(in_memory_cache=InMemoryCache(), redis_cache=fake_redis_cache(server))
|
||||
await raw.set("byom_submitted_servers:user-123", '["a"]')
|
||||
await raw.set("byom_submitted_servers:other", '["b"]')
|
||||
await cache.async_set_cache(key="byom_submitted_servers:user-123", value=["a"], local_only=True)
|
||||
manager = MCPServerManager()
|
||||
|
||||
with patch.object(proxy_server_module, "user_api_key_cache", cache):
|
||||
await manager.invalidate_byom_submitted_servers_cache("user-123")
|
||||
await manager.invalidate_byom_submitted_servers_cache(None)
|
||||
|
||||
cache.async_delete_cache.assert_awaited_once_with(key="byom_submitted_servers:user-123")
|
||||
assert cache.in_memory_cache.get_cache("byom_submitted_servers:user-123") is None
|
||||
assert await raw.exists("byom_submitted_servers:user-123") == 0
|
||||
assert await raw.get("byom_submitted_servers:other") == b'["b"]'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_active_submitted_ids_cache_miss_queries_db_and_caches(self):
|
||||
import fakeredis
|
||||
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.proxy import proxy_server as proxy_server_module
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from tests._support.fake_redis import fake_redis_cache
|
||||
|
||||
cache = MagicMock()
|
||||
cache.async_get_cache = AsyncMock(return_value=None)
|
||||
cache.async_set_cache = AsyncMock()
|
||||
server = fakeredis.FakeServer()
|
||||
cache = UserApiKeyCache(in_memory_cache=InMemoryCache(), redis_cache=fake_redis_cache(server))
|
||||
manager = MCPServerManager()
|
||||
manager.registry = {
|
||||
"submitted-server": MCPServer(
|
||||
|
|
@ -9073,24 +9088,25 @@ class TestMCPServerManager:
|
|||
result = await manager._get_active_submitted_mcp_server_ids_for_user(user_api_key_auth)
|
||||
|
||||
assert result == ["submitted-server"]
|
||||
cache.async_set_cache.assert_awaited_once_with(
|
||||
key="byom_submitted_servers:user-123",
|
||||
value=["submitted-server", "unknown-server"],
|
||||
ttl=60,
|
||||
)
|
||||
raw = fakeredis.FakeAsyncRedis(server=server)
|
||||
assert await raw.get("byom_submitted_servers:user-123") == b'["submitted-server", "unknown-server"]'
|
||||
assert await raw.ttl("byom_submitted_servers:user-123") == 60
|
||||
peer = UserApiKeyCache(in_memory_cache=InMemoryCache(), redis_cache=fake_redis_cache(server))
|
||||
with patch.object(proxy_server_module, "user_api_key_cache", peer):
|
||||
assert await manager._get_active_submitted_mcp_server_ids_for_user(user_api_key_auth) == ["submitted-server"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_allowed_mcp_servers_fallback_keeps_submitted_byom_servers(self):
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.proxy import proxy_server as proxy_server_module
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.mcp_server.auth.user_api_key_auth_mcp import (
|
||||
MCPRequestHandler,
|
||||
)
|
||||
|
||||
class _Cache:
|
||||
async def async_get_cache(self, key: str):
|
||||
assert key == "byom_submitted_servers:user-123"
|
||||
return ["submitted-server"]
|
||||
cache = UserApiKeyCache(in_memory_cache=InMemoryCache())
|
||||
await cache.async_set_cache(key="byom_submitted_servers:user-123", value=["submitted-server"])
|
||||
|
||||
manager = MCPServerManager()
|
||||
manager.registry = {
|
||||
|
|
@ -9106,7 +9122,7 @@ class TestMCPServerManager:
|
|||
)
|
||||
|
||||
with (
|
||||
patch.object(proxy_server_module, "user_api_key_cache", _Cache()),
|
||||
patch.object(proxy_server_module, "user_api_key_cache", cache),
|
||||
patch.object(proxy_server_module, "prisma_client", None),
|
||||
patch.object(manager, "get_allow_all_keys_server_ids", return_value=["global-server"]),
|
||||
patch.object(
|
||||
|
|
|
|||
66
tests/unit/redis/stores/test_mcp_server_catalog.py
Normal file
66
tests/unit/redis/stores/test_mcp_server_catalog.py
Normal file
|
|
@ -0,0 +1,66 @@
|
|||
import json
|
||||
from datetime import timedelta
|
||||
from typing import Final
|
||||
|
||||
import fakeredis
|
||||
import pytest
|
||||
|
||||
from litellm.redis.stores.mcp_server_catalog import TOOLSET_NOT_FOUND, TransportMcpServerCatalogStore
|
||||
from litellm.redis.transport.contract import RedisFailure
|
||||
from litellm.types.mcp_server.mcp_toolset import MCPToolset
|
||||
from tests._support.fake_redis import fake_redis_cache
|
||||
|
||||
|
||||
def _store(server: fakeredis.FakeServer) -> TransportMcpServerCatalogStore:
|
||||
return TransportMcpServerCatalogStore(fake_redis_cache(server, namespace="ns").transport())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_entries_keep_the_legacy_keys_json_bytes_and_whole_second_ttls(server: fakeredis.FakeServer) -> None:
|
||||
store: Final = _store(server)
|
||||
raw: Final = fakeredis.FakeAsyncRedis(server=server)
|
||||
toolset: Final = MCPToolset(toolset_id="ts-1", toolset_name="ops", tools=[{"server_id": "a", "tool_name": "t"}])
|
||||
|
||||
assert await store.put_submitted_server_ids("u1", ("a", "b"), ttl=timedelta(seconds=60)) is True
|
||||
assert await store.put_toolset_permissions(("t2", "t1"), {"a": ("x",)}, ttl=timedelta(seconds=90.7)) is True
|
||||
assert await store.put_toolset_by_name("ops", toolset, ttl=timedelta(seconds=30)) is True
|
||||
assert await store.put_toolset_by_name("gone", None, ttl=timedelta(seconds=30)) is True
|
||||
|
||||
assert await raw.get("ns:byom_submitted_servers:u1") == b'["a", "b"]'
|
||||
assert await raw.ttl("ns:byom_submitted_servers:u1") == 60
|
||||
assert await raw.get("ns:toolset_perms:t1,t2") == b'{"a": ["x"]}'
|
||||
assert await raw.ttl("ns:toolset_perms:t1,t2") == 90
|
||||
assert await raw.get("ns:toolset_name:gone") == b'"__not_found__"'
|
||||
assert await store.get_submitted_server_ids("u1") == ("a", "b")
|
||||
assert await store.get_toolset_permissions(("t1", "t2")) == {"a": ("x",)}
|
||||
assert await store.get_toolset_by_name("ops") == toolset
|
||||
assert await store.get_toolset_by_name("gone") == TOOLSET_NOT_FOUND
|
||||
assert await store.invalidate_submitted_server_ids("u1") == 1
|
||||
assert await store.get_submitted_server_ids("u1") is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_old_pod_entries_decode_and_wrong_shapes_are_decode_failures(server: fakeredis.FakeServer) -> None:
|
||||
raw: Final = fakeredis.FakeAsyncRedis(server=server)
|
||||
await raw.set("ns:byom_submitted_servers:u1", json.dumps(["a"]))
|
||||
await raw.set("ns:toolset_perms:t1", "{'a': ['x']}")
|
||||
await raw.set("ns:toolset_name:ops", json.dumps({"toolset_id": "ts-1", "toolset_name": "ops"}))
|
||||
await raw.set("ns:byom_submitted_servers:bad", b'{"not": "a list"}')
|
||||
await raw.set("ns:toolset_name:bad", b"{not json")
|
||||
store: Final = _store(server)
|
||||
|
||||
assert await store.get_submitted_server_ids("u1") == ("a",)
|
||||
assert await store.get_toolset_permissions(("t1",)) == {"a": ("x",)}
|
||||
assert await store.get_toolset_by_name("ops") == MCPToolset(toolset_id="ts-1", toolset_name="ops")
|
||||
wrong_shape: Final = await store.get_submitted_server_ids("bad")
|
||||
malformed: Final = await store.get_toolset_by_name("bad")
|
||||
assert isinstance(wrong_shape, RedisFailure) and wrong_shape.kind == "decode"
|
||||
assert isinstance(malformed, RedisFailure) and malformed.kind == "decode"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_outage_is_a_transport_failure(server: fakeredis.FakeServer) -> None:
|
||||
store: Final = _store(server)
|
||||
server.connected = False
|
||||
outage: Final = await store.get_toolset_permissions(("t1",))
|
||||
assert isinstance(outage, RedisFailure) and outage.kind == "transport"
|
||||
Loading…
Add table
Reference in a new issue