From ffae582a50ad0ea1d26be4aefd6bf81565e35bcd Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 22:30:22 +0000 Subject: [PATCH] refactor(redis): move MCP server manager caches behind a typed store --- .../proxy/mcp_server/mcp_server_manager.py | 59 +++---- .../proxy/mcp_server/server_catalog_cache.py | 158 ++++++++++++++++++ litellm/redis/stores/mcp_server_catalog.py | 121 ++++++++++++++ .../mcp_server/test_mcp_server_manager.py | 50 ++++-- .../redis/stores/test_mcp_server_catalog.py | 66 ++++++++ 5 files changed, 402 insertions(+), 52 deletions(-) create mode 100644 litellm/proxy/mcp_server/server_catalog_cache.py create mode 100644 litellm/redis/stores/mcp_server_catalog.py create mode 100644 tests/unit/redis/stores/test_mcp_server_catalog.py diff --git a/litellm/proxy/mcp_server/mcp_server_manager.py b/litellm/proxy/mcp_server/mcp_server_manager.py index c772e2e9d09..98f466c2012 100644 --- a/litellm/proxy/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/mcp_server/mcp_server_manager.py @@ -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 diff --git a/litellm/proxy/mcp_server/server_catalog_cache.py b/litellm/proxy/mcp_server/server_catalog_cache.py new file mode 100644 index 00000000000..7f0db4de2ab --- /dev/null +++ b/litellm/proxy/mcp_server/server_catalog_cache.py @@ -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))) diff --git a/litellm/redis/stores/mcp_server_catalog.py b/litellm/redis/stores/mcp_server_catalog.py new file mode 100644 index 00000000000..1809c6fde9f --- /dev/null +++ b/litellm/redis/stores/mcp_server_catalog.py @@ -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__}") diff --git a/tests/unit/proxy/mcp_server/test_mcp_server_manager.py b/tests/unit/proxy/mcp_server/test_mcp_server_manager.py index 13e4704bd90..bfa888dd77a 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_server_manager.py +++ b/tests/unit/proxy/mcp_server/test_mcp_server_manager.py @@ -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( diff --git a/tests/unit/redis/stores/test_mcp_server_catalog.py b/tests/unit/redis/stores/test_mcp_server_catalog.py new file mode 100644 index 00000000000..afaa714519b --- /dev/null +++ b/tests/unit/redis/stores/test_mcp_server_catalog.py @@ -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"