diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index f86920ebd92..5e3a3946604 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -16,6 +16,12 @@ from litellm.caching.redis_batch import active_request_redis_batch from litellm.caching.redis_cache import RedisCache from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec +from litellm.redis.stores.mcp_agent_permission_links import ( + AGENT_NO_MCP_PERMISSION_SENTINEL, + McpAgentPermissionLinkStore, + TransportMcpAgentPermissionLinkStore, + agent_permission_link_key, +) from litellm.redis.stores.mcp_user_permission_links import ( USER_NO_MCP_PERMISSION_SENTINEL as _USER_NO_PERMISSION, ) @@ -368,6 +374,44 @@ def object_permission_cache_key(object_permission_id: str) -> str: return f"object_permission_id:{object_permission_id}" +async def get_agent_permission_link(cache: UserApiKeyCache, agent_id: str) -> str | None: + key: Final = agent_permission_link_key(agent_id) + local: Final = _PERMISSION_LINK_VALUE.validate_python(await cache.async_get_cache(key=key, local_only=True)) + if local is not None: + return local if isinstance(local, str) else None + 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): + store: Final[McpAgentPermissionLinkStore] = TransportMcpAgentPermissionLinkStore(cache.redis_cache.transport()) + result: Final = await store.get(agent_id) + if isinstance(result, RedisFailure): + return None + if result is not None: + await cache.async_set_cache(key=key, value=result, local_only=True) + return result + return None + + +async def fill_agent_permission_link( + cache: UserApiKeyCache, agent_id: str, permission_id: str | None, *, ttl: float +) -> None: + await cache.async_set_cache( + key=agent_permission_link_key(agent_id), + value=permission_id or AGENT_NO_MCP_PERMISSION_SENTINEL, + ttl=ttl, + local_only=True, + ) + if cache.redis_cache is not None: + with suppress(Exception): + store: Final[McpAgentPermissionLinkStore] = TransportMcpAgentPermissionLinkStore( + cache.redis_cache.transport() + ) + _ = await store.fill(agent_id, permission_id, ttl=timedelta(seconds=int(ttl))) + + #: Cached under ``tag_registry_cache_key`` when the table exceeds ``TAG_REGISTRY_MAX_SIZE``: #: registry unusable, fall back to the per-tag lookup. TAG_REGISTRY_OVERFLOW_SENTINEL: Final = "__tag_registry_overflow__" diff --git a/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py index 67a3f2fb590..c671ad7d8c4 100644 --- a/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py @@ -61,7 +61,9 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg from litellm.proxy.common_utils.user_api_key_cache import ( AUTH_OBJECTS_TARGET, USER_NO_MCP_PERMISSION_SENTINEL, + fill_agent_permission_link, fill_user_permission_link, + get_agent_permission_link, get_management_object_ttl, get_user_permission_link, ) @@ -91,6 +93,7 @@ from litellm.proxy.mcp_server.utils import ( asgi_scope_fields, strip_known_server_prefix, ) +from litellm.redis.stores.mcp_agent_permission_links import AGENT_NO_MCP_PERMISSION_SENTINEL from litellm.repositories.table_repositories import ( AgentsRepository, MCPServerRepository, @@ -4447,9 +4450,7 @@ class MCPRequestHandler: return list(end_user_tools) return list(set(allowed_tools) & set(end_user_tools)) - # Sentinel stored in cache when an agent has no object_permission, so we - # don't re-query the DB on every MCP request for that agent. - _AGENT_NO_PERMISSION_SENTINEL = "__agent_no_mcp_permission__" + _AGENT_NO_PERMISSION_SENTINEL: Final = AGENT_NO_MCP_PERMISSION_SENTINEL @staticmethod @with_service_target(AUTH_OBJECTS_TARGET) @@ -4462,9 +4463,8 @@ class MCPRequestHandler: a link we DID resolve can make the caller deny.""" from litellm.proxy.proxy_server import user_api_key_cache - cache_key: Final = f"agent_object_permission_id:{agent_id}" try: - cached: Final = _CACHED_VALUE.validate_python(await user_api_key_cache.async_get_cache(key=cache_key)) + cached: Final = await get_agent_permission_link(user_api_key_cache, agent_id) if cached == MCPRequestHandler._AGENT_NO_PERMISSION_SENTINEL: return None if isinstance(cached, str) and cached: @@ -4472,9 +4472,10 @@ class MCPRequestHandler: agent_row: Final = await AgentsRepository(prisma_client).table.find_unique(where={"agent_id": agent_id}) linked: Final[object] = getattr(agent_row, "object_permission_id", None) if agent_row is not None else None object_permission_id: Final = linked if isinstance(linked, str) and linked else None - await user_api_key_cache.async_set_cache( - key=cache_key, - value=object_permission_id or MCPRequestHandler._AGENT_NO_PERMISSION_SENTINEL, + await fill_agent_permission_link( + user_api_key_cache, + agent_id, + object_permission_id, ttl=get_management_object_ttl(user_api_key_cache), ) return object_permission_id diff --git a/litellm/redis/stores/mcp_agent_permission_links.py b/litellm/redis/stores/mcp_agent_permission_links.py new file mode 100644 index 00000000000..fe28319614d --- /dev/null +++ b/litellm/redis/stores/mcp_agent_permission_links.py @@ -0,0 +1,59 @@ +import ast +import json +from dataclasses import dataclass +from datetime import timedelta +from typing import Final, Protocol + +from pydantic import TypeAdapter, ValidationError + +from litellm.redis.transport.contract import RedisFailure, RedisTransport + +AGENT_NO_MCP_PERMISSION_SENTINEL: Final = "__agent_no_mcp_permission__" +_VALUE: Final = TypeAdapter(object) + + +def agent_permission_link_key(agent_id: str) -> str: + return f"agent_object_permission_id:{agent_id}" + + +class McpAgentPermissionLinkStore(Protocol): + async def get(self, agent_id: str) -> str | None | RedisFailure: ... + + async def fill(self, agent_id: str, permission_id: str | None, *, ttl: timedelta) -> bool | None | RedisFailure: ... + + +@dataclass(frozen=True, slots=True) +class TransportMcpAgentPermissionLinkStore: + transport: RedisTransport + + async def get(self, agent_id: str) -> str | None | RedisFailure: + reply: Final = await self.transport.get(agent_permission_link_key(agent_id)) + if reply is None or isinstance(reply, RedisFailure): + return reply + try: + decoded: Final = reply.decode("utf-8") + try: + value: Final = _VALUE.validate_json(decoded) + except ValidationError: + literal: Final = _VALUE.validate_python(ast.literal_eval(decoded)) + return _permission_id(literal) + return _permission_id(value) + except (UnicodeDecodeError, ValueError, SyntaxError): + return RedisFailure(kind="decode", message="Invalid agent permission link") + + async def fill(self, agent_id: str, permission_id: str | None, *, ttl: timedelta) -> bool | None | RedisFailure: + seconds: Final = int(ttl.total_seconds()) + key: Final = agent_permission_link_key(agent_id) + value: Final = json.dumps(permission_id or AGENT_NO_MCP_PERMISSION_SENTINEL) + expiry: Final = timedelta(seconds=seconds) + batch: Final = self.transport.request_batch() + if batch is not None: + _ = batch.set_seconds(key, value, ttl=expiry) + return None + return await self.transport.set_seconds(key, value, ttl=expiry) + + +def _permission_id(value: object) -> str | RedisFailure: + if isinstance(value, str): + return value + return RedisFailure(kind="decode", message="Agent permission link is not a string") diff --git a/tests/unit/proxy/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/unit/proxy/mcp_server/auth/test_user_api_key_auth_mcp.py index 230a3c84790..b912b399058 100644 --- a/tests/unit/proxy/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/unit/proxy/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -79,6 +79,45 @@ async def test_fresh_user_policy_does_not_use_redis_sentinel(monkeypatch: pytest assert error.value.status_code == 503 +@pytest.mark.asyncio +@pytest.mark.parametrize("permission_id", ["agent-policy", None]) +async def test_agent_permission_link_fills_and_hits_redis_without_reloading_agent( + monkeypatch: pytest.MonkeyPatch, permission_id: str | None +) -> None: + server: Final = fakeredis.FakeServer() + redis: Final = fake_redis_cache(server, namespace="auth") + cache: Final = UserApiKeyCache(redis_cache=redis, default_in_memory_ttl=41) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache) + prisma: Final = MagicMock() + prisma.db.litellm_agentstable.find_unique = AsyncMock( + return_value=SimpleNamespace(object_permission_id=permission_id) + ) + assert await MCPRequestHandler._agent_object_permission_id("link-agent", prisma) == permission_id # pyright: ignore[reportPrivateUsage] # resolver internals under test + assert await redis.async_get_cache("agent_object_permission_id:link-agent") == ( + permission_id or "__agent_no_mcp_permission__" + ) + cache.in_memory_cache.delete_cache("agent_object_permission_id:link-agent") + assert await MCPRequestHandler._agent_object_permission_id("link-agent", prisma) == permission_id # pyright: ignore[reportPrivateUsage] # resolver internals under test + server.connected = False + assert await MCPRequestHandler._agent_object_permission_id("link-agent", prisma) == permission_id # pyright: ignore[reportPrivateUsage] # resolver internals under test + prisma.db.litellm_agentstable.find_unique.assert_awaited_once_with(where={"agent_id": "link-agent"}) + + +@pytest.mark.asyncio +async def test_unavailable_agent_policy_keeps_no_ceiling_without_negative_cache( + monkeypatch: pytest.MonkeyPatch, +) -> None: + server: Final = fakeredis.FakeServer() + cache: Final = UserApiKeyCache(redis_cache=fake_redis_cache(server, namespace="auth")) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache) + prisma: Final = MagicMock() + prisma.db.litellm_agentstable.find_unique = AsyncMock(side_effect=RuntimeError("database unavailable")) + server.connected = False + assert await MCPRequestHandler._agent_object_permission_id("link-agent", prisma) is None # pyright: ignore[reportPrivateUsage] # resolver internals under test + assert await cache.async_get_cache("agent_object_permission_id:link-agent", local_only=True) is None + prisma.db.litellm_agentstable.find_unique.assert_awaited_once_with(where={"agent_id": "link-agent"}) + + @pytest.mark.parametrize( "headers,expected", [ diff --git a/tests/unit/redis/stores/test_mcp_agent_permission_links.py b/tests/unit/redis/stores/test_mcp_agent_permission_links.py new file mode 100644 index 00000000000..0b02b50af2c --- /dev/null +++ b/tests/unit/redis/stores/test_mcp_agent_permission_links.py @@ -0,0 +1,209 @@ +from datetime import timedelta +from typing import Final + +import fakeredis +import pytest +import pytest_asyncio +from pytest import MonkeyPatch +from redis.exceptions import ConnectionError as RedisConnectionError + +from litellm.caching.redis_batch import active_request_redis_batch, request_redis_batch_scope +from litellm.caching.redis_cache import RedisCache +from litellm.proxy.common_utils.user_api_key_cache import ( + UserApiKeyCache, + fill_agent_permission_link, + get_agent_permission_link, +) +from litellm.redis.stores.mcp_agent_permission_links import TransportMcpAgentPermissionLinkStore +from litellm.redis.transport.contract import RedisFailure +from tests._support.fake_redis import fake_redis_cache + + +@pytest_asyncio.fixture(loop_scope="function") +async def link_cache(monkeypatch: MonkeyPatch) -> tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache]: + monkeypatch.setattr("time.time", lambda: 1000.0) + server: Final = fakeredis.FakeServer() + redis: Final = fake_redis_cache(server, namespace="links") + cache: Final = UserApiKeyCache(redis_cache=redis, default_in_memory_ttl=37.9) + return server, redis, cache + + +@pytest.mark.asyncio +@pytest.mark.parametrize("value", ['policy é "quoted"', "__agent_no_mcp_permission__", ""]) +async def test_old_writer_new_reader( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], value: str +) -> None: + _, redis, cache = link_cache + await redis.async_set_cache("agent_object_permission_id:u", value, ttl=37.9) + assert await get_agent_permission_link(cache, "u") == value + assert await cache.async_get_cache("agent_object_permission_id:u", local_only=True) == value + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "value,wire", + [(None, b'"__agent_no_mcp_permission__"'), ('policy é "quoted"', b'"policy \\u00e9 \\"quoted\\""')], +) +@pytest.mark.parametrize("ttl", [37.9, 37.999999999]) +async def test_new_writer_old_reader_and_whole_second_expiry( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], value: str | None, wire: bytes, ttl: float +) -> None: + server, redis, cache = link_cache + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await fill_agent_permission_link(cache, "u", value, ttl=ttl) + assert await raw.get("links:agent_object_permission_id:u") == wire + assert await raw.pttl("links:agent_object_permission_id:u") == 37000 + assert await redis.async_get_cache("agent_object_permission_id:u") == (value or "__agent_no_mcp_permission__") + assert await raw.keys("*") == [b"links:agent_object_permission_id:u"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("wire", [b"broken", b"42", b"null", b"{}", b"[]", b"\xff"]) +async def test_bad_payload_is_decode_failure_and_caller_cache_miss( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], wire: bytes +) -> None: + server, redis, cache = link_cache + await fakeredis.FakeAsyncRedis(server=server).set("links:agent_object_permission_id:u", wire) + result: Final = await TransportMcpAgentPermissionLinkStore(redis.transport()).get("u") + assert isinstance(result, RedisFailure) and result.kind == "decode" + assert await get_agent_permission_link(cache, "u") is None + + +@pytest.mark.asyncio +async def test_literal_legacy_string_and_local_non_string( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], +) -> None: + server, _, cache = link_cache + await fakeredis.FakeAsyncRedis(server=server).set("links:agent_object_permission_id:u", b"'legacy'") + assert await get_agent_permission_link(cache, "u") == "legacy" + await cache.async_set_cache("agent_object_permission_id:u", {"invalid": True}, local_only=True) + assert await get_agent_permission_link(cache, "u") is None + + +@pytest.mark.asyncio +async def test_agent_link_does_not_overwrite_user_link_or_permission_row( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], +) -> None: + server, redis, cache = link_cache + await redis.async_set_cache("user_object_permission_id:u", "user-policy", ttl=60) + await redis.async_set_cache("object_permission_id:u", {"mcp_servers": ["server"]}, ttl=60) + await fill_agent_permission_link(cache, "u", "agent-policy", ttl=37.9) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + assert await raw.mget( + "links:agent_object_permission_id:u", "links:user_object_permission_id:u", "links:object_permission_id:u" + ) == [b'"agent-policy"', b'"user-policy"', b'{"mcp_servers": ["server"]}'] + + +@pytest.mark.asyncio +async def test_outage_preserves_local_fill_and_open_breaker( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], +) -> None: + server, redis, cache = link_cache + server.connected = False + assert await get_agent_permission_link(cache, "u") is None + await fill_agent_permission_link(cache, "u", "policy", ttl=30) + assert await get_agent_permission_link(cache, "u") == "policy" + cache.in_memory_cache.delete_cache("agent_object_permission_id:u") + assert await cache.async_get_cache("agent_object_permission_id:u", local_only=True) is None + store: Final = TransportMcpAgentPermissionLinkStore(redis.transport()) + for _ in range(10): + assert isinstance(await store.get("u"), RedisFailure) + result: Final = await store.get("u") + assert isinstance(result, RedisFailure) and result.kind == "breaker_open" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ttl", [37.9, 37.999999999]) +async def test_request_fill_stays_queued_and_clears_prior_miss( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], ttl: float +) -> None: + server, redis, cache = link_cache + raw: Final = fakeredis.FakeAsyncRedis(server=server) + with request_redis_batch_scope(): + batch: Final = active_request_redis_batch(redis) + assert batch is not None + assert await batch.mget(("agent_object_permission_id:u",)) == {"agent_object_permission_id:u": None} + await redis.async_set_cache("agent_object_permission_id:u", "external", ttl=60) + assert await get_agent_permission_link(cache, "u") is None + await fill_agent_permission_link(cache, "u", "policy", ttl=ttl) + assert await get_agent_permission_link(cache, "u") == "policy" + assert await raw.get("links:agent_object_permission_id:u") == b'"external"' + await batch.flush() + assert await raw.get("links:agent_object_permission_id:u") == b'"policy"' + assert await raw.pttl("links:agent_object_permission_id:u") == 37000 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ttl", [0.9, 0.999999999, 0, -1]) +async def test_rejected_expiry_keeps_existing_redis_value_and_new_local_value( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], ttl: float +) -> None: + _, redis, cache = link_cache + await redis.async_set_cache("agent_object_permission_id:u", "old", ttl=60) + await fill_agent_permission_link(cache, "u", "new", ttl=ttl) + assert await cache.async_get_cache("agent_object_permission_id:u", local_only=True) == ("new" if ttl >= 0 else None) + assert await redis.async_get_cache("agent_object_permission_id:u") == "old" + result: Final = await TransportMcpAgentPermissionLinkStore(redis.transport()).fill( + "u", "new", ttl=timedelta(seconds=int(ttl)) + ) + assert isinstance(result, RedisFailure) and result.kind == "server" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ttl", [0.999999999, 0, -1]) +async def test_rejected_batched_expiry_keeps_existing_value( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], ttl: float +) -> None: + _, redis, cache = link_cache + await redis.async_set_cache("agent_object_permission_id:u", "old", ttl=60) + with request_redis_batch_scope() as scope: + await fill_agent_permission_link(cache, "u", "new", ttl=ttl) + await scope.flush_all() + assert await redis.async_get_cache("agent_object_permission_id:u") == "old" + + +@pytest.mark.asyncio +async def test_unclassified_client_failure_keeps_cache_optional( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], monkeypatch: MonkeyPatch +) -> None: + _, redis, cache = link_cache + client: Final = redis.init_async_client() + assert isinstance(client, fakeredis.FakeAsyncRedis) + await client.ping() + connection: Final = await client.connection_pool.get_connection() + read_response: Final = connection.read_response + + async def unavailable_reply(**kwargs: object) -> object: + await read_response(**kwargs) + raise RuntimeError("client callback failed") + + monkeypatch.setattr(connection, "read_response", unavailable_reply) + await client.connection_pool.release(connection) + assert await cache.async_get_cache("agent_object_permission_id:a") is None + assert await get_agent_permission_link(cache, "a") is None + await fill_agent_permission_link(cache, "a", "policy", ttl=30) + assert await get_agent_permission_link(cache, "a") == "policy" + + +@pytest.mark.asyncio +async def test_lost_set_reply_does_not_roll_back_applied_write( + link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], monkeypatch: MonkeyPatch +) -> None: + server, redis, cache = link_cache + client: Final = redis.init_async_client() + assert isinstance(client, fakeredis.FakeAsyncRedis) + await client.ping() + connection: Final = await client.connection_pool.get_connection() + read_response: Final = connection.read_response + + async def lost_reply(**kwargs: object) -> object: + await read_response(**kwargs) + raise RedisConnectionError("reply lost after write") + + monkeypatch.setattr(connection, "read_response", lost_reply) + await client.connection_pool.release(connection) + await fill_agent_permission_link(cache, "u", "applied", ttl=37.9) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + assert await raw.get("links:agent_object_permission_id:u") == b'"applied"' + assert await raw.pttl("links:agent_object_permission_id:u") == 37000 + assert await get_agent_permission_link(cache, "u") == "applied"