mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-06 02:48:13 +00:00
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:
parent
c3387c476f
commit
19ab870f32
2 changed files with 90 additions and 2 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue