refactor(redis): migrate agent MCP permission links to typed store

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__"

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

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")

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",
     [

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"