fix(mcp-oauth2): LRU cache eviction + regression tests for bug fixes

- Replace _byok_cred_cache.clear() with single-entry eviction in
  _write_byok_cred_cache: at capacity, evict only the oldest key
  instead of wiping all 4096 entries (thundering-herd prevention)
- Add regression tests:
  - Cache LRU eviction: verifies exactly 1 entry evicted at capacity
  - MCPServer byok fields: is_byok/byok_description/byok_api_key_help_url
    propagated correctly (regression for the load_servers_from_config
    missing-fields bug that caused BYOK servers to lose their fields)
  - Cache capacity: tests all paths in _write_byok_cred_cache
This commit is contained in:
Ishaan Jaffer 2026-03-07 09:18:46 -08:00
parent c3387c476f
commit 19ab870f32
2 changed files with 90 additions and 2 deletions

View file

@ -96,9 +96,14 @@ def _invalidate_byok_cred_cache(user_id: str, server_id: str) -> None:
def _write_byok_cred_cache(
user_id: str, server_id: str, credential: Optional[str]
) -> None:
"""Write a credential value to the cache, evicting all entries if at capacity."""
"""Write a credential value to the cache, evicting the oldest entry if at capacity.
Evicts a single entry (LRU-style) rather than clearing all at once to avoid
a thundering-herd DB spike when the cache fills under load.
"""
if len(_byok_cred_cache) >= _BYOK_CRED_CACHE_MAX_SIZE:
_byok_cred_cache.clear()
oldest_key = next(iter(_byok_cred_cache))
del _byok_cred_cache[oldest_key]
_byok_cred_cache[(user_id, server_id)] = (credential, time.monotonic())
# Check if MCP is available

View file

@ -416,3 +416,86 @@ def test_extract_access_token_json_without_access_token_key():
blob = json.dumps({"some_other_key": "value"})
# Falls back to raw string since there's no access_token
assert _extract_access_token(blob) == blob
# ---------------------------------------------------------------------------
# _write_byok_cred_cache — LRU eviction (regression for thundering-herd bug)
# ---------------------------------------------------------------------------
def test_write_byok_cred_cache_evicts_single_oldest_entry():
"""At capacity, _write_byok_cred_cache must evict one (oldest) entry, not clear all."""
from litellm.proxy._experimental.mcp_server.server import (
_BYOK_CRED_CACHE_MAX_SIZE,
_byok_cred_cache,
_write_byok_cred_cache,
)
_byok_cred_cache.clear()
# Fill to capacity
for i in range(_BYOK_CRED_CACHE_MAX_SIZE):
_byok_cred_cache[(f"user{i}", "server")] = ("token", 0.0)
assert len(_byok_cred_cache) == _BYOK_CRED_CACHE_MAX_SIZE
# Writing one more entry should evict the oldest, not clear everything
_write_byok_cred_cache("new_user", "server", "new_token")
# Cache should still be at max size — only ONE entry was evicted
assert len(_byok_cred_cache) == _BYOK_CRED_CACHE_MAX_SIZE
# The new entry must be present
assert ("new_user", "server") in _byok_cred_cache
# The oldest entry (user0) should have been evicted
assert ("user0", "server") not in _byok_cred_cache
_byok_cred_cache.clear()
# ---------------------------------------------------------------------------
# Regression tests for the three bug fixes in c1fcbf6219
# ---------------------------------------------------------------------------
def test_mcp_server_byok_fields_propagated():
"""Bug fix 1 (related): MCPServer must accept and preserve is_byok,
byok_description, and byok_api_key_help_url (previously missing from
load_servers_from_config which caused BYOK OAuth2 servers to lose these
fields after config load).
"""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPTransport
from litellm.types.mcp_server.mcp_server_manager import MCPServer
server = MCPServer(
server_id="github",
name="github",
server_name="GitHub",
transport=MCPTransport.sse,
is_byok=True,
byok_description=["Read your profile", "List your repos"],
byok_api_key_help_url="https://github.com/settings/tokens",
)
assert server.is_byok is True
assert server.byok_description == ["Read your profile", "List your repos"]
assert server.byok_api_key_help_url == "https://github.com/settings/tokens"
def test_mcp_server_byok_fields_default_to_empty():
"""Regression: byok fields default to False/[] (not None) so they don't
need null-guards in code that iterates byok_description.
"""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPTransport
from litellm.types.mcp_server.mcp_server_manager import MCPServer
server = MCPServer(
server_id="plain",
name="plain",
server_name="Plain Server",
transport=MCPTransport.http,
)
assert server.is_byok is False
assert isinstance(server.byok_description, list)