From b46733e5b92856732493fd55de36db75238216c5 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Thu, 25 Jun 2026 23:10:41 -0700 Subject: [PATCH] =?UTF-8?q?feat(mcp):=20DualCache-backed=20token=20cache?= =?UTF-8?q?=20backend=20(step=201b=20=C2=A71.5)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The cross-replica TokenCacheBackend implementation that plugs into the foundation's CachedOAuthTokenStore seam: encrypts+serializes the token via the codec and stores it in LiteLLM's shared DualCache under the same per-(user,server) key v1 used, so workers share one refresh and a token cached by v1 or v2 is readable by the other across the cutover. Cache and codec are injected; a non-positive TTL (already-expired token) is not cached, and a missing/corrupt entry reads as a miss. --- .../dual_cache_token_backend.py | 65 +++++++++++++ .../test_dual_cache_token_backend.py | 93 +++++++++++++++++++ 2 files changed, 158 insertions(+) create mode 100644 litellm/proxy/_experimental/mcp_server/outbound_credentials/dual_cache_token_backend.py create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_dual_cache_token_backend.py diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/dual_cache_token_backend.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/dual_cache_token_backend.py new file mode 100644 index 00000000000..eaa4fb71c36 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/dual_cache_token_backend.py @@ -0,0 +1,65 @@ +"""Cross-replica ``TokenCacheBackend``: stores the token in LiteLLM's shared ``DualCache``. + +Plugs into the foundation's ``CachedOAuthTokenStore`` via the ``TokenCacheBackend`` seam. The token is +encrypted + serialized by the injected codec and written under a per-``(user, server)`` key with the +given TTL, so every worker reads one refresh rather than each re-reading and re-refreshing - matching +v1's ``MCPPerUserTokenCache`` (same NaCl encryption and key, so a token cached by either is readable by +the other across the cutover). A missing or undecryptable entry reads as a miss. +""" + +from __future__ import annotations + +from typing import Protocol + +from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( + OAuthToken, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.token_cache_codec import ( + OAuthTokenCacheCodec, +) + + +class AsyncCache(Protocol): + """The slice of LiteLLM's ``DualCache`` this backend needs (Redis-backed, shared across workers).""" + + async def async_get_cache(self, key: str) -> object | None: ... + + async def async_set_cache( + self, key: str, value: str, ttl: float | None = None + ) -> None: ... + + async def async_delete_cache(self, key: str) -> None: ... + + +class DualCacheTokenCacheBackend: + def __init__( + self, + cache: AsyncCache, + codec: OAuthTokenCacheCodec, + *, + key_prefix: str = "mcp:per_user_token:", + ) -> None: + self._cache = cache + self._codec = codec + self._key_prefix = key_prefix + + def _key(self, user_id: str, server_id: str) -> str: + return f"{self._key_prefix}{user_id}:{server_id}" + + async def get(self, user_id: str, server_id: str) -> OAuthToken | None: + blob = await self._cache.async_get_cache(self._key(user_id, server_id)) + return self._codec.decode(blob) if isinstance(blob, str) else None + + async def set( + self, user_id: str, server_id: str, token: OAuthToken, ttl_seconds: float + ) -> None: + if ttl_seconds <= 0: + return + await self._cache.async_set_cache( + self._key(user_id, server_id), + self._codec.encode(token), + ttl=ttl_seconds, + ) + + async def delete(self, user_id: str, server_id: str) -> None: + await self._cache.async_delete_cache(self._key(user_id, server_id)) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_dual_cache_token_backend.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_dual_cache_token_backend.py new file mode 100644 index 00000000000..d79e332f40b --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_dual_cache_token_backend.py @@ -0,0 +1,93 @@ +"""Tests for the DualCache-backed token cache backend: encrypted round-trip, key, TTL, miss.""" + +import pytest + +from litellm.proxy._experimental.mcp_server.outbound_credentials.dual_cache_token_backend import ( + DualCacheTokenCacheBackend, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( + OAuthToken, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.token_cache_codec import ( + OAuthTokenCacheCodec, +) + + +class _FakeCache: + def __init__(self): + self.values = {} + self.ttls = {} + + async def async_get_cache(self, key): + return self.values.get(key) + + async def async_set_cache(self, key, value, ttl=None): + self.values[key] = value + self.ttls[key] = ttl + + async def async_delete_cache(self, key): + self.values.pop(key, None) + + +def _backend(cache): + codec = OAuthTokenCacheCodec( + encrypt=lambda s: f"enc:{s}", + decrypt=lambda b: b[4:] if b.startswith("enc:") else None, + ) + return DualCacheTokenCacheBackend(cache, codec) + + +@pytest.mark.asyncio +async def test_set_then_get_round_trips_encrypted_under_the_per_user_key(): + cache = _FakeCache() + backend = _backend(cache) + await backend.set("alice", "srv", OAuthToken(access_token="at"), 120.0) + + key = "mcp:per_user_token:alice:srv" + assert cache.ttls[key] == 120.0 + # Stored via the codec (encrypted), not the bare token; the codec's own test proves real + # NaCl output hides the secret - here the fake encrypt just wraps, so we check it was applied. + assert cache.values[key] == "enc:at" + + got = await backend.get("alice", "srv") + assert got is not None and got.access_token == "at" + + +@pytest.mark.asyncio +async def test_get_missing_key_is_none(): + assert await _backend(_FakeCache()).get("alice", "srv") is None + + +@pytest.mark.asyncio +async def test_non_str_cache_value_is_a_miss(): + cache = _FakeCache() + cache.values["mcp:per_user_token:alice:srv"] = 12345 # corrupt / wrong type + assert await _backend(cache).get("alice", "srv") is None + + +@pytest.mark.asyncio +async def test_non_positive_ttl_is_not_written(): + cache = _FakeCache() + await _backend(cache).set("alice", "srv", OAuthToken(access_token="at"), 0.0) + assert cache.values == {} # an already-expired token is not cached + + +@pytest.mark.asyncio +async def test_delete_removes_the_entry(): + cache = _FakeCache() + backend = _backend(cache) + await backend.set("alice", "srv", OAuthToken(access_token="at"), 60.0) + await backend.delete("alice", "srv") + assert await backend.get("alice", "srv") is None + + +@pytest.mark.asyncio +async def test_keys_isolate_users_and_servers(): + cache = _FakeCache() + backend = _backend(cache) + await backend.set("alice", "srv", OAuthToken(access_token="a"), 60.0) + await backend.set("bob", "srv", OAuthToken(access_token="b"), 60.0) + alice = await backend.get("alice", "srv") + bob = await backend.get("bob", "srv") + assert alice is not None and alice.access_token == "a" + assert bob is not None and bob.access_token == "b"