mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
fix(mcp): broadcast BYOK and OAuth credential eviction to peer workers and expire admin session tombstones
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin
parent
aea33b4b50
commit
fd834f6f8b
13 changed files with 366 additions and 74 deletions
|
|
@ -184,7 +184,8 @@ MCP_METADATA_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_METADATA_TIMEOUT", "1
|
|||
MCP_HEALTH_CHECK_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_HEALTH_CHECK_TIMEOUT", "10.0"))
|
||||
MCP_TOOL_LISTING_MAX_PAGES: Final = 1000
|
||||
MCP_GATEWAY_SESSION_ID_PREFIX_LENGTH: Final = 8
|
||||
MCP_ADMIN_TERMINATED_SESSION_IDS_MAX: Final = 1024
|
||||
MCP_BYOK_CREDENTIAL_CACHE_TTL_SECONDS: Final = 60
|
||||
MCP_BYOK_CREDENTIAL_CACHE_MAX_SIZE: Final = 4096
|
||||
|
||||
# Allowlist of commands permitted for MCP stdio transport.
|
||||
# Prevents arbitrary command execution via /mcp-rest/test/* endpoints or server creation.
|
||||
|
|
|
|||
|
|
@ -0,0 +1,38 @@
|
|||
"""Per-worker cache of stored BYOK credentials, keyed so peer workers can evict it over the auth cache pub/sub."""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.constants import MCP_BYOK_CREDENTIAL_CACHE_MAX_SIZE, MCP_BYOK_CREDENTIAL_CACHE_TTL_SECONDS
|
||||
|
||||
_CACHE_KEY_PREFIX: Final = "mcp_byok_credential"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CachedByokCredential:
|
||||
credential: str | None
|
||||
|
||||
|
||||
byok_credential_cache: Final = InMemoryCache(
|
||||
max_size_in_memory=MCP_BYOK_CREDENTIAL_CACHE_MAX_SIZE,
|
||||
default_ttl=MCP_BYOK_CREDENTIAL_CACHE_TTL_SECONDS,
|
||||
)
|
||||
|
||||
|
||||
def byok_credential_cache_key(user_id: str, server_id: str) -> str:
|
||||
return f"{_CACHE_KEY_PREFIX}:{user_id}:{server_id}"
|
||||
|
||||
|
||||
def get_cached_byok_credential(user_id: str, server_id: str) -> CachedByokCredential | None:
|
||||
cached: Final = byok_credential_cache.get_cache( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # InMemoryCache is untyped
|
||||
byok_credential_cache_key(user_id, server_id)
|
||||
)
|
||||
return cached if isinstance(cached, CachedByokCredential) else None
|
||||
|
||||
|
||||
def cache_byok_credential(user_id: str, server_id: str, credential: str | None) -> None:
|
||||
byok_credential_cache.set_cache( # pyright: ignore[reportUnknownMemberType] # InMemoryCache is untyped
|
||||
byok_credential_cache_key(user_id, server_id),
|
||||
CachedByokCredential(credential=credential),
|
||||
)
|
||||
|
|
@ -865,7 +865,7 @@ async def byok_token(
|
|||
_invalidate_byok_cred_cache,
|
||||
)
|
||||
|
||||
_invalidate_byok_cred_cache(user_id, server_id)
|
||||
await _invalidate_byok_cred_cache(user_id, server_id)
|
||||
except Exception as exc:
|
||||
verbose_proxy_logger.error(
|
||||
"byok_token: failed to store user credential for user=%s server=%s: %s",
|
||||
|
|
|
|||
|
|
@ -295,12 +295,15 @@ class MCPPerUserTokenCache:
|
|||
)
|
||||
|
||||
async def delete(self, user_id: str, server_id: str) -> None:
|
||||
"""Invalidate the cached token (removes from both in-memory and Redis layers)."""
|
||||
"""Invalidate the cached token in Redis, here, and in every peer worker's in-memory layer."""
|
||||
try:
|
||||
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( # noqa: PLC0415 # proxy import cycle
|
||||
evict_and_broadcast,
|
||||
)
|
||||
from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415
|
||||
|
||||
key: Final = self._cache_key(user_id, server_id)
|
||||
await user_api_key_cache.async_delete_cache(key)
|
||||
await evict_and_broadcast((key,), user_api_key_cache)
|
||||
except Exception as exc:
|
||||
verbose_logger.debug(
|
||||
"MCPPerUserTokenCache.delete failed for user=%s server=%s: %s",
|
||||
|
|
|
|||
|
|
@ -14,7 +14,7 @@ import time
|
|||
import traceback
|
||||
import types
|
||||
import uuid
|
||||
from collections import Counter, deque
|
||||
from collections import Counter
|
||||
from collections.abc import AsyncIterator, Callable, Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol
|
||||
|
|
@ -30,7 +30,6 @@ from typing_extensions import ReadOnly, TypedDict
|
|||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import (
|
||||
MAXIMUM_TRACEBACK_LINES_TO_LOG,
|
||||
MCP_ADMIN_TERMINATED_SESSION_IDS_MAX,
|
||||
MCP_GATEWAY_SESSION_ID_PREFIX_LENGTH,
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
|
@ -42,6 +41,12 @@ from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
|
|||
MCPRequestHandler,
|
||||
_is_mcp_admitted_user_subject,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.byok_credential_cache import (
|
||||
byok_credential_cache,
|
||||
byok_credential_cache_key,
|
||||
cache_byok_credential,
|
||||
get_cached_byok_credential,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
|
||||
get_request_base_url,
|
||||
)
|
||||
|
|
@ -86,6 +91,9 @@ from litellm.proxy._types import (
|
|||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.auth.ip_address_utils import IPAddressUtils
|
||||
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
|
||||
publish_auth_cache_invalidation,
|
||||
)
|
||||
from litellm.proxy.litellm_pre_call_utils import (
|
||||
LiteLLMProxyRequestSetup,
|
||||
get_chain_id_from_headers,
|
||||
|
|
@ -107,13 +115,6 @@ if TYPE_CHECKING:
|
|||
|
||||
from litellm.proxy._experimental.mcp_server.db import OAuthCredentialPayload
|
||||
|
||||
# Short-lived in-memory cache for BYOK credentials.
|
||||
# Keyed by (user_id, server_id); value is (credential_or_None, monotonic_timestamp).
|
||||
# Storing the credential value (not just a bool) means _get_byok_credential and
|
||||
# _check_byok_credential share a single DB round-trip per TTL window.
|
||||
_byok_cred_cache: Final[dict[tuple[str, str], tuple[str | None, float]]] = {}
|
||||
_BYOK_CRED_CACHE_TTL: Final = 60 # seconds
|
||||
_BYOK_CRED_CACHE_MAX_SIZE: Final = 4096 # cap to prevent unbounded growth
|
||||
_STATEFUL_SESSION_IDLE_TIMEOUT_SECONDS: Final = 30 * 60
|
||||
# Upper bound on concurrent stateful sessions a single caller may hold. Each
|
||||
# `initialize` creates a session that survives until the idle timeout, so
|
||||
|
|
@ -132,20 +133,11 @@ _MCP_TRANSPORT_SPAN_SCOPE_KEY: Final = "litellm_otel_transport_span"
|
|||
_MCP_DESTINATIONS_SCOPE_KEY: Final = "litellm_otel_request_destinations"
|
||||
|
||||
|
||||
def _invalidate_byok_cred_cache(user_id: str, server_id: str) -> None:
|
||||
"""Remove a (user_id, server_id) entry from the BYOK credential cache.
|
||||
|
||||
Call this after storing or deleting a credential so subsequent calls
|
||||
see the fresh value rather than a stale cached result.
|
||||
"""
|
||||
_byok_cred_cache.pop((user_id, server_id), None)
|
||||
|
||||
|
||||
def _write_byok_cred_cache(user_id: str, server_id: str, credential: str | None) -> None:
|
||||
"""Write a credential value to the cache, evicting all entries if at capacity."""
|
||||
if len(_byok_cred_cache) >= _BYOK_CRED_CACHE_MAX_SIZE:
|
||||
_byok_cred_cache.clear()
|
||||
_byok_cred_cache[(user_id, server_id)] = (credential, time.monotonic())
|
||||
async def _invalidate_byok_cred_cache(user_id: str, server_id: str) -> None:
|
||||
"""Drop a stored-or-deleted BYOK credential from this worker's cache and from every peer worker's."""
|
||||
cache_key: Final = byok_credential_cache_key(user_id, server_id)
|
||||
byok_credential_cache.delete_cache(cache_key)
|
||||
await publish_auth_cache_invalidation(cache_key=cache_key)
|
||||
|
||||
|
||||
# Check if MCP is available
|
||||
|
|
@ -623,9 +615,7 @@ if MCP_AVAILABLE:
|
|||
_stateful_session_locks: Final[dict[str, asyncio.Lock]] = {}
|
||||
_stateful_session_active_request_counts: Final[dict[str, int]] = {}
|
||||
_stateful_session_client_info: Final[dict[str, Implementation]] = {} # mutable-ok: cleared on session teardown
|
||||
_admin_terminated_session_ids: Final[deque[str]] = deque( # mutable-ok: bounded ring, appended on admin termination
|
||||
maxlen=MCP_ADMIN_TERMINATED_SESSION_IDS_MAX
|
||||
)
|
||||
_admin_terminated_session_ids: Final[dict[str, float]] = {} # mutable-ok: admin-closed id -> last replay
|
||||
|
||||
class _TerminableTransport(Protocol):
|
||||
async def terminate(self) -> None: ...
|
||||
|
|
@ -697,6 +687,7 @@ if MCP_AVAILABLE:
|
|||
for session_id in list(_stateful_session_auth_context_last_seen):
|
||||
if session_id not in _stateful_session_auth_contexts:
|
||||
_remove_stateful_session_tracking(session_id)
|
||||
_forget_expired_admin_terminated_session_ids(now)
|
||||
|
||||
async def _enforce_stateful_session_cap_for_owner(owner: str) -> bool:
|
||||
"""
|
||||
|
|
@ -2819,35 +2810,28 @@ if MCP_AVAILABLE:
|
|||
mcp_server: MCPServer,
|
||||
user_api_key_auth: UserAPIKeyAuth | None,
|
||||
) -> str | None:
|
||||
"""Retrieve the stored BYOK credential for a user+server pair.
|
||||
|
||||
Uses the shared _byok_cred_cache to avoid a DB round-trip on every
|
||||
tool call within the TTL window.
|
||||
"""
|
||||
"""Retrieve the stored BYOK credential for a user+server pair, served from the worker cache within its TTL."""
|
||||
if not mcp_server.is_byok:
|
||||
return None
|
||||
user_id: Final = (user_api_key_auth.user_id if user_api_key_auth else None) or ""
|
||||
if not user_id:
|
||||
return None
|
||||
|
||||
cache_key: Final = (user_id, mcp_server.server_id)
|
||||
cached: Final = _byok_cred_cache.get(cache_key)
|
||||
cached: Final = get_cached_byok_credential(user_id, mcp_server.server_id)
|
||||
if cached is not None:
|
||||
credential, ts = cached
|
||||
if time.monotonic() - ts < _BYOK_CRED_CACHE_TTL:
|
||||
return credential
|
||||
return cached.credential
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.db import get_user_credential
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
if prisma_client is None:
|
||||
return None
|
||||
credential = await get_user_credential(
|
||||
credential: Final = await get_user_credential(
|
||||
prisma_client=prisma_client,
|
||||
user_id=user_id,
|
||||
server_id=mcp_server.server_id,
|
||||
)
|
||||
_write_byok_cred_cache(user_id, mcp_server.server_id, credential)
|
||||
cache_byok_credential(user_id, mcp_server.server_id, credential)
|
||||
return credential
|
||||
|
||||
async def _check_byok_credential(
|
||||
|
|
@ -2876,27 +2860,23 @@ if MCP_AVAILABLE:
|
|||
headers={"WWW-Authenticate": get_byok_www_authenticate()},
|
||||
)
|
||||
|
||||
# Check shared credential cache before hitting the DB.
|
||||
cache_key: Final = (user_id, mcp_server.server_id)
|
||||
cached: Final = _byok_cred_cache.get(cache_key)
|
||||
cached: Final = get_cached_byok_credential(user_id, mcp_server.server_id)
|
||||
if cached is not None:
|
||||
cached_cred, ts = cached
|
||||
if time.monotonic() - ts < _BYOK_CRED_CACHE_TTL:
|
||||
if cached_cred is None:
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
detail={
|
||||
"error": "byok_auth_required",
|
||||
"server_id": mcp_server.server_id,
|
||||
"server_name": mcp_server.server_name or mcp_server.name,
|
||||
"message": (
|
||||
"No stored credential found for this BYOK server. "
|
||||
"Complete the OAuth authorization flow to provide your API key."
|
||||
),
|
||||
},
|
||||
headers={"WWW-Authenticate": get_byok_www_authenticate()},
|
||||
)
|
||||
return
|
||||
if cached.credential is None:
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
detail={
|
||||
"error": "byok_auth_required",
|
||||
"server_id": mcp_server.server_id,
|
||||
"server_name": mcp_server.server_name or mcp_server.name,
|
||||
"message": (
|
||||
"No stored credential found for this BYOK server. "
|
||||
"Complete the OAuth authorization flow to provide your API key."
|
||||
),
|
||||
},
|
||||
headers={"WWW-Authenticate": get_byok_www_authenticate()},
|
||||
)
|
||||
return
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.db import get_user_credential
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
|
@ -2920,7 +2900,7 @@ if MCP_AVAILABLE:
|
|||
user_id=user_id,
|
||||
server_id=mcp_server.server_id,
|
||||
)
|
||||
_write_byok_cred_cache(user_id, mcp_server.server_id, credential)
|
||||
cache_byok_credential(user_id, mcp_server.server_id, credential)
|
||||
if credential is None:
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
|
|
@ -3906,6 +3886,24 @@ if MCP_AVAILABLE:
|
|||
key_auth: Final = auth_user.user_api_key_auth
|
||||
return key_auth is not None and key_auth.user_id == user_id
|
||||
|
||||
def _forget_expired_admin_terminated_session_ids(now: float) -> None:
|
||||
for session_id in [
|
||||
session_id
|
||||
for session_id, last_replayed in _admin_terminated_session_ids.items()
|
||||
if now - last_replayed >= _STATEFUL_SESSION_IDLE_TIMEOUT_SECONDS
|
||||
]:
|
||||
del _admin_terminated_session_ids[session_id]
|
||||
|
||||
def _is_admin_terminated_session_id(session_id: str, now: float) -> bool:
|
||||
last_replayed: Final = _admin_terminated_session_ids.get(session_id)
|
||||
if last_replayed is None:
|
||||
return False
|
||||
if now - last_replayed >= _STATEFUL_SESSION_IDLE_TIMEOUT_SECONDS:
|
||||
del _admin_terminated_session_ids[session_id]
|
||||
return False
|
||||
_admin_terminated_session_ids[session_id] = now
|
||||
return True
|
||||
|
||||
async def terminate_mcp_gateway_sessions(
|
||||
*,
|
||||
session_id_prefix: str | None = None,
|
||||
|
|
@ -3919,6 +3917,7 @@ if MCP_AVAILABLE:
|
|||
admission. Only sessions held by this worker process are affected.
|
||||
"""
|
||||
now: Final = time.monotonic()
|
||||
_forget_expired_admin_terminated_session_ids(now)
|
||||
server_instances: Final = _stateful_server_instances()
|
||||
targets: Final = tuple(
|
||||
(session_id, auth_user)
|
||||
|
|
@ -3928,7 +3927,7 @@ if MCP_AVAILABLE:
|
|||
)
|
||||
terminated: Final = tuple(_gateway_session_for(session_id, auth_user, now) for session_id, auth_user in targets)
|
||||
for session_id, _ in targets:
|
||||
_admin_terminated_session_ids.append(session_id)
|
||||
_admin_terminated_session_ids[session_id] = now
|
||||
transport = server_instances.pop(session_id, None)
|
||||
_remove_stateful_session_tracking(session_id)
|
||||
if transport is not None:
|
||||
|
|
@ -4064,7 +4063,7 @@ if MCP_AVAILABLE:
|
|||
await success_response(scope, receive, send)
|
||||
return True
|
||||
|
||||
if _session_id in _admin_terminated_session_ids:
|
||||
if _is_admin_terminated_session_id(_session_id, time.monotonic()):
|
||||
terminated_response: Final = JSONResponse(
|
||||
status_code=404,
|
||||
content={ # mutable-ok: JSONResponse content must be a plain dict
|
||||
|
|
|
|||
|
|
@ -2317,7 +2317,7 @@ if MCP_AVAILABLE:
|
|||
_invalidate_byok_cred_cache,
|
||||
)
|
||||
|
||||
_invalidate_byok_cred_cache(user_id, server_id)
|
||||
await _invalidate_byok_cred_cache(user_id, server_id)
|
||||
return MCPUserCredentialResponse(server_id=server_id, has_credential=True)
|
||||
# save=False: credential not persisted
|
||||
return MCPUserCredentialResponse(server_id=server_id, has_credential=False)
|
||||
|
|
@ -2348,7 +2348,7 @@ if MCP_AVAILABLE:
|
|||
_invalidate_byok_cred_cache,
|
||||
)
|
||||
|
||||
_invalidate_byok_cred_cache(target_user_id, server_id)
|
||||
await _invalidate_byok_cred_cache(target_user_id, server_id)
|
||||
return MCPUserCredentialResponse(server_id=server_id, has_credential=False)
|
||||
|
||||
# ── OAuth2 user-credential endpoints ──────────────────────────────────────
|
||||
|
|
|
|||
|
|
@ -307,6 +307,7 @@ from litellm.litellm_core_utils.sensitive_data_masker import (
|
|||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
from litellm.llms.openai_like.model_info import MODEL_INFO_REFRESH_SECONDS
|
||||
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
|
||||
from litellm.proxy._experimental.mcp_server.byok_credential_cache import byok_credential_cache
|
||||
from litellm.proxy._lazy_features import attach_lazy_features, reserve_lazy_slot
|
||||
from litellm.proxy._types import *
|
||||
from litellm.proxy.analytics_endpoints.analytics_endpoints import (
|
||||
|
|
@ -7535,7 +7536,7 @@ class ProxyConfig:
|
|||
subscriber: Final = AuthCacheInvalidationSubscriber(
|
||||
redis_cache=redis_cache,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
additional_in_memory_caches=(spend_counter_cache.in_memory_cache,),
|
||||
additional_in_memory_caches=(spend_counter_cache.in_memory_cache, byok_credential_cache),
|
||||
)
|
||||
self.auth_cache_invalidation_subscriber = subscriber
|
||||
subscriber.start()
|
||||
|
|
|
|||
|
|
@ -0,0 +1,57 @@
|
|||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.byok_credential_cache import (
|
||||
CachedByokCredential,
|
||||
byok_credential_cache,
|
||||
byok_credential_cache_key,
|
||||
cache_byok_credential,
|
||||
get_cached_byok_credential,
|
||||
)
|
||||
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import AuthCacheInvalidationSubscriber
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
|
||||
class _FakeRedisCache:
|
||||
namespace = None
|
||||
|
||||
def init_async_client(self) -> object:
|
||||
return object()
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _empty_cache():
|
||||
byok_credential_cache.flush_cache()
|
||||
yield
|
||||
byok_credential_cache.flush_cache()
|
||||
|
||||
|
||||
def test_a_cached_negative_lookup_is_distinguishable_from_a_miss():
|
||||
assert get_cached_byok_credential("u-1", "srv-1") is None
|
||||
cache_byok_credential("u-1", "srv-1", None)
|
||||
assert get_cached_byok_credential("u-1", "srv-1") == CachedByokCredential(credential=None)
|
||||
cache_byok_credential("u-1", "srv-1", "sk-stored")
|
||||
assert get_cached_byok_credential("u-1", "srv-1") == CachedByokCredential(credential="sk-stored")
|
||||
assert get_cached_byok_credential("u-1", "srv-2") is None
|
||||
|
||||
|
||||
def test_peer_worker_invalidation_message_evicts_the_cached_credential():
|
||||
"""The key a mutating worker broadcasts must be the key every other worker caches under."""
|
||||
cache_byok_credential("mallory", "srv-byok", "sk-revoked")
|
||||
cache_byok_credential("alice", "srv-byok", "sk-kept")
|
||||
subscriber = AuthCacheInvalidationSubscriber(
|
||||
redis_cache=_FakeRedisCache(), # pyright: ignore[reportArgumentType] # subscriber is never started; only its message handler runs
|
||||
user_api_key_cache=UserApiKeyCache(),
|
||||
additional_in_memory_caches=(byok_credential_cache,),
|
||||
)
|
||||
|
||||
subscriber._apply_message( # pyright: ignore[reportPrivateUsage] # exercising the real cross-worker message handler
|
||||
{
|
||||
"type": "message",
|
||||
"data": json.dumps({"cache_key": byok_credential_cache_key("mallory", "srv-byok")}).encode(),
|
||||
}
|
||||
)
|
||||
|
||||
assert get_cached_byok_credential("mallory", "srv-byok") is None
|
||||
assert get_cached_byok_credential("alice", "srv-byok") == CachedByokCredential(credential="sk-kept")
|
||||
|
|
@ -592,7 +592,7 @@ async def test_check_byok_credential_missing_credential(monkeypatch):
|
|||
|
||||
monkeypatch.delenv("PROXY_BASE_URL", raising=False)
|
||||
monkeypatch.delenv("SERVER_ROOT_PATH", raising=False)
|
||||
monkeypatch.setattr(server_module, "_byok_cred_cache", {})
|
||||
server_module.byok_credential_cache.flush_cache()
|
||||
mock_prisma = MagicMock()
|
||||
|
||||
with (
|
||||
|
|
@ -628,7 +628,7 @@ async def test_execute_byok_tool_missing_credential_advertises_api_key_flow(monk
|
|||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
|
||||
monkeypatch.setenv("PROXY_BASE_URL", "https://gateway.example.com/proxy")
|
||||
monkeypatch.setattr(mcp_module, "_byok_cred_cache", {})
|
||||
mcp_module.byok_credential_cache.flush_cache()
|
||||
server = MCPServer(server_id="byok-discovery", name="byok-discovery", transport=MCPTransport.http, is_byok=True)
|
||||
prisma = MagicMock()
|
||||
prisma.db.litellm_mcpusercredentials.find_unique = AsyncMock(return_value=None)
|
||||
|
|
@ -677,6 +677,34 @@ async def test_check_byok_credential_has_credential():
|
|||
await _check_byok_credential(server, user_auth)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalidate_byok_cred_cache_evicts_locally_and_broadcasts_the_same_key():
|
||||
"""A revoked credential must stop being served here and on every peer worker within the TTL."""
|
||||
from litellm.proxy._experimental.mcp_server import server as server_module
|
||||
from litellm.proxy._experimental.mcp_server.byok_credential_cache import byok_credential_cache_key
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
|
||||
server = MCPServer(server_id="byok-revoke", name="byok-server", transport=MCPTransport.http, is_byok=True)
|
||||
user_auth = UserAPIKeyAuth(user_id="mallory", api_key="sk-test")
|
||||
server_module.byok_credential_cache.flush_cache()
|
||||
db_lookup = AsyncMock(side_effect=["sk-before-revoke", None])
|
||||
publish = AsyncMock()
|
||||
|
||||
with (
|
||||
patch("litellm.proxy._experimental.mcp_server.db.get_user_credential", new=db_lookup),
|
||||
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
|
||||
patch.object(server_module, "publish_auth_cache_invalidation", new=publish),
|
||||
):
|
||||
assert await server_module._get_byok_credential(server, user_auth) == "sk-before-revoke"
|
||||
assert await server_module._get_byok_credential(server, user_auth) == "sk-before-revoke"
|
||||
await server_module._invalidate_byok_cred_cache("mallory", "byok-revoke")
|
||||
assert await server_module._get_byok_credential(server, user_auth) is None
|
||||
|
||||
assert db_lookup.await_count == 2
|
||||
publish.assert_awaited_once_with(cache_key=byok_credential_cache_key("mallory", "byok-revoke"))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_byok_credential_db_unavailable_fails_closed():
|
||||
"""BYOK server with no prisma_client → 503, not silent pass.
|
||||
|
|
|
|||
|
|
@ -2989,9 +2989,10 @@ async def test_admin_terminated_session_id_gets_404_instead_of_a_fresh_stateless
|
|||
"""Once an admin closes a session, a client replaying its id must not be silently upgraded to a
|
||||
new stateless session by the stale-header path; it gets 404 and has to initialize again."""
|
||||
try:
|
||||
from starlette.types import Scope
|
||||
|
||||
from litellm.proxy._experimental.mcp_server import server as mcp_server
|
||||
from litellm.proxy._experimental.mcp_server.server import session_manager_stateful
|
||||
from starlette.types import Scope
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
|
|
@ -3048,6 +3049,73 @@ async def test_admin_terminated_session_id_gets_404_instead_of_a_fresh_stateless
|
|||
mcp_server._admin_terminated_session_ids.clear()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_admin_terminated_session_id_stays_refused_while_replayed_and_is_forgotten_like_an_idle_session():
|
||||
"""The refusal window slides on every replay, so a client that keeps retrying is never silently
|
||||
upgraded to a stateless session no matter how many other sessions an admin closes later; an id
|
||||
nobody has replayed for a full idle timeout is dropped from the table by the idle sweep."""
|
||||
try:
|
||||
from starlette.types import Scope
|
||||
|
||||
from litellm.proxy._experimental.mcp_server import server as mcp_server
|
||||
from litellm.proxy._experimental.mcp_server.server import session_manager_stateful
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
idle_timeout = mcp_server._STATEFUL_SESSION_IDLE_TIMEOUT_SECONDS
|
||||
retrying_id, silent_id = "admin-closed-retrying", "admin-closed-silent"
|
||||
contexts = {
|
||||
session_id: mcp_server.MCPAuthenticatedUser(
|
||||
user_api_key_auth=UserAPIKeyAuth(api_key="key-alice", user_id="alice"),
|
||||
)
|
||||
for session_id in (retrying_id, silent_id)
|
||||
}
|
||||
live_transports = {session_id: MagicMock(terminate=AsyncMock()) for session_id in contexts}
|
||||
|
||||
async def replay(session_id: str, now: float) -> tuple[bool, list[bytes]]:
|
||||
scope: Scope = {
|
||||
"type": "http",
|
||||
"method": "POST",
|
||||
"headers": [(b"content-type", b"application/json"), (b"mcp-session-id", session_id.encode())],
|
||||
}
|
||||
with patch.object(mcp_server.time, "monotonic", return_value=now):
|
||||
handled = await mcp_server._handle_stale_mcp_session(
|
||||
scope, AsyncMock(), AsyncMock(), session_manager_stateful
|
||||
)
|
||||
return handled, [k for k, _ in scope["headers"]]
|
||||
|
||||
try:
|
||||
with (
|
||||
patch.object( # test-quality-ok: the transport registry is a module-level singleton; the suite's only seam
|
||||
session_manager_stateful, "_server_instances", live_transports
|
||||
),
|
||||
patch.dict( # test-quality-ok: the session tables are module-level singletons; the suite's only seam
|
||||
mcp_server._stateful_session_auth_contexts, contexts, clear=True
|
||||
),
|
||||
patch.dict( # test-quality-ok: the session tables are module-level singletons; the suite's only seam
|
||||
mcp_server._stateful_session_client_info, {}, clear=True
|
||||
),
|
||||
patch.dict( # test-quality-ok: the session tables are module-level singletons; the suite's only seam
|
||||
mcp_server._stateful_session_auth_context_last_seen, {}, clear=True
|
||||
),
|
||||
):
|
||||
with patch.object(mcp_server.time, "monotonic", return_value=1000.0):
|
||||
closed = await mcp_server.terminate_mcp_gateway_sessions(user_id="alice")
|
||||
assert closed.terminated_sessions == 2
|
||||
|
||||
for elapsed in (idle_timeout - 1, 2 * idle_timeout - 2, 3 * idle_timeout - 3):
|
||||
assert await replay(retrying_id, 1000.0 + elapsed) == (True, [b"content-type", b"mcp-session-id"])
|
||||
|
||||
await mcp_server._purge_expired_stateful_session_auth_contexts(now=1000.0 + idle_timeout)
|
||||
assert set(mcp_server._admin_terminated_session_ids) == {retrying_id}
|
||||
|
||||
assert await replay(silent_id, 1000.0 + idle_timeout) == (False, [b"content-type"])
|
||||
assert await replay(retrying_id, 1000.0 + 4 * idle_timeout) == (False, [b"content-type"])
|
||||
assert mcp_server._admin_terminated_session_ids == {}
|
||||
finally:
|
||||
mcp_server._admin_terminated_session_ids.clear()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_initialize_request_with_existing_session_tracks_new_session():
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -395,6 +395,32 @@ async def test_invalidate_clears_every_identity_for_a_server():
|
|||
assert mock_client.post.call_count == 3
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_per_user_token_delete_evicts_locally_and_broadcasts_to_peer_workers():
|
||||
"""Revoking a user's OAuth token must not leave peer workers serving it from their in-memory layer."""
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy._experimental.mcp_server.oauth2_token_cache import MCPPerUserTokenCache
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
local_cache = UserApiKeyCache()
|
||||
publish = AsyncMock()
|
||||
token_cache = MCPPerUserTokenCache()
|
||||
key = token_cache._cache_key("mallory", "srv-oauth") # pyright: ignore[reportPrivateUsage] # asserting the broadcast names the stored key
|
||||
local_cache.in_memory_cache.set_cache(key, "encrypted-token")
|
||||
|
||||
with (
|
||||
patch.object(proxy_server, "user_api_key_cache", local_cache),
|
||||
patch(
|
||||
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation",
|
||||
new=publish,
|
||||
),
|
||||
):
|
||||
await token_cache.delete("mallory", "srv-oauth")
|
||||
|
||||
assert local_cache.in_memory_cache.get_cache(key) is None
|
||||
publish.assert_awaited_once_with(cache_key=key)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_m2m_mint_uses_admin_entered_token_url_when_issuer_yield_empties_resolved():
|
||||
"""A pinned issuer empties the resolved token_url while configured_token_url keeps the
|
||||
|
|
|
|||
|
|
@ -5152,7 +5152,7 @@ async def test_admin_revokes_another_users_byok_credential():
|
|||
)
|
||||
|
||||
delete_mock = AsyncMock(return_value=None)
|
||||
invalidate_mock = MagicMock()
|
||||
invalidate_mock = AsyncMock()
|
||||
with (
|
||||
patch( # test-quality-ok: endpoint test stubs the Prisma client lookup
|
||||
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
|
||||
|
|
@ -5174,7 +5174,7 @@ async def test_admin_revokes_another_users_byok_credential():
|
|||
|
||||
delete_mock.assert_awaited_once()
|
||||
assert delete_mock.await_args.args[1:] == ("mallory", "srv-byok-admin")
|
||||
invalidate_mock.assert_called_once_with("mallory", "srv-byok-admin")
|
||||
invalidate_mock.assert_awaited_once_with("mallory", "srv-byok-admin")
|
||||
assert result.has_credential is False
|
||||
|
||||
|
||||
|
|
@ -5231,7 +5231,7 @@ async def test_user_naming_themselves_still_deletes_own_byok_credential():
|
|||
new=delete_mock,
|
||||
),
|
||||
patch.object( # test-quality-ok: the cache invalidator is module scoped; the suite's only seam
|
||||
mcp_server, "_invalidate_byok_cred_cache", new=MagicMock()
|
||||
mcp_server, "_invalidate_byok_cred_cache", new=AsyncMock()
|
||||
),
|
||||
):
|
||||
await delete_mcp_user_credential(
|
||||
|
|
|
|||
|
|
@ -13919,3 +13919,74 @@ async def test_token_counter_loads_a_custom_tokenizer_once_per_identifier_revisi
|
|||
]
|
||||
finally:
|
||||
litellm.utils._select_custom_tokenizer_helper.cache_clear()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auth_cache_invalidation_subscriber_evicts_byok_credentials_cached_by_this_worker():
|
||||
"""A peer worker's BYOK revocation broadcast must reach this worker's BYOK credential cache."""
|
||||
from redis.asyncio import Redis
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.byok_credential_cache import (
|
||||
byok_credential_cache,
|
||||
byok_credential_cache_key,
|
||||
cache_byok_credential,
|
||||
get_cached_byok_credential,
|
||||
)
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
class _QueuePubSub:
|
||||
def __init__(self, messages: list[object]) -> None:
|
||||
self.queue: asyncio.Queue[object] = asyncio.Queue()
|
||||
for message in messages:
|
||||
self.queue.put_nowait(message)
|
||||
|
||||
async def subscribe(self, *channels: str) -> None:
|
||||
return None
|
||||
|
||||
async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> object | None:
|
||||
try:
|
||||
return await asyncio.wait_for(self.queue.get(), timeout)
|
||||
except asyncio.TimeoutError:
|
||||
return None
|
||||
|
||||
async def aclose(self) -> None:
|
||||
return None
|
||||
|
||||
class _PubSubRedisClient(Redis):
|
||||
def __init__(self, pubsub: _QueuePubSub) -> None:
|
||||
self._scripted_pubsub = pubsub
|
||||
|
||||
def pubsub(self) -> _QueuePubSub:
|
||||
return self._scripted_pubsub
|
||||
|
||||
class _FakeRedisCache:
|
||||
namespace = None
|
||||
|
||||
def __init__(self, client: object) -> None:
|
||||
self._client = client
|
||||
|
||||
def init_async_client(self) -> object:
|
||||
return self._client
|
||||
|
||||
byok_credential_cache.flush_cache()
|
||||
cache_byok_credential("mallory", "srv-byok", "sk-revoked-elsewhere")
|
||||
message: Final = {
|
||||
"type": "message",
|
||||
"data": json.dumps({"cache_key": byok_credential_cache_key("mallory", "srv-byok")}).encode(),
|
||||
}
|
||||
proxy_config: Final = proxy_server_module.ProxyConfig()
|
||||
proxy_config.start_auth_cache_invalidation_subscriber(
|
||||
redis_cache=_FakeRedisCache(_PubSubRedisClient(_QueuePubSub([message]))), # pyright: ignore[reportArgumentType] # fake pub/sub capable redis; no live redis in this unit test
|
||||
user_api_key_cache=UserApiKeyCache(),
|
||||
)
|
||||
try:
|
||||
for _ in range(200):
|
||||
if get_cached_byok_credential("mallory", "srv-byok") is None:
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
evicted: Final = get_cached_byok_credential("mallory", "srv-byok") is None
|
||||
finally:
|
||||
await proxy_config.stop_auth_cache_invalidation_subscriber()
|
||||
byok_credential_cache.flush_cache()
|
||||
|
||||
assert evicted, "the subscriber does not evict the BYOK credential cache on a peer worker's broadcast"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue