mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-10 03:28:53 +00:00
fix: address greptile review feedback (greploop iteration 4)
- Unified credential cache: store actual credential value (Optional[str]) instead of just bool so _get_byok_credential also benefits from caching — eliminates the DB hit on every BYOK tool call within the 60s TTL window - Extracted _write_byok_cred_cache() helper for consistent cache writes - Replaced has_user_credential with get_user_credential in _check_byok_credential so one DB call satisfies both existence check and value retrieval - Remove false 'encrypted at rest' claim from OAuth HTML and ByokCredentialModal
This commit is contained in:
parent
8f902f925f
commit
2f24a7ebf2
3 changed files with 41 additions and 18 deletions
|
|
@ -480,7 +480,7 @@ def _build_authorize_html(
|
|||
<span class="sec-icon">
|
||||
<svg xmlns="http://www.w3.org/2000/svg" width="14" height="14" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><rect x="3" y="11" width="18" height="11" rx="2" ry="2"/><path d="M7 11V7a5 5 0 0 1 10 0v4"/></svg>
|
||||
</span>
|
||||
<p>Your key is encrypted at rest and transmitted securely. It is never shared with third parties.</p>
|
||||
<p>Your key is stored securely and transmitted over HTTPS. It is never shared with third parties.</p>
|
||||
</div>
|
||||
|
||||
<button type="submit" class="btn-connect" id="connectBtn">
|
||||
|
|
|
|||
|
|
@ -55,9 +55,11 @@ from litellm.types.mcp_server.mcp_server_manager import MCPInfo, MCPServer
|
|||
from litellm.types.utils import CallTypes, StandardLoggingMCPToolCall
|
||||
from litellm.utils import Rules, client, function_setup
|
||||
|
||||
# Short-lived in-memory cache for BYOK credential existence checks.
|
||||
# Keyed by (user_id, server_id); value is (credential_exists, monotonic_timestamp).
|
||||
_byok_cred_cache: Dict[Tuple[str, str], Tuple[bool, float]] = {}
|
||||
# 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: Dict[Tuple[str, str], Tuple[Optional[str], float]] = {}
|
||||
_BYOK_CRED_CACHE_TTL = 60 # seconds
|
||||
_BYOK_CRED_CACHE_MAX_SIZE = 4096 # cap to prevent unbounded growth
|
||||
|
||||
|
|
@ -65,11 +67,20 @@ _BYOK_CRED_CACHE_MAX_SIZE = 4096 # cap to prevent unbounded growth
|
|||
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 a new credential so the user isn't blocked by a
|
||||
previously cached negative result.
|
||||
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: Optional[str]
|
||||
) -> 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())
|
||||
|
||||
# Check if MCP is available
|
||||
# "mcp" requires python 3.10 or higher, but several litellm users use python 3.8
|
||||
# We're making this conditional import to avoid breaking users who use python 3.8.
|
||||
|
|
@ -1521,22 +1532,36 @@ if MCP_AVAILABLE:
|
|||
mcp_server: MCPServer,
|
||||
user_api_key_auth: Optional[UserAPIKeyAuth],
|
||||
) -> Optional[str]:
|
||||
"""Retrieve the stored BYOK credential for a user+server pair."""
|
||||
"""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.
|
||||
"""
|
||||
if not mcp_server.is_byok:
|
||||
return None
|
||||
user_id = (user_api_key_auth.user_id if user_api_key_auth else None) or ""
|
||||
if not user_id:
|
||||
return None
|
||||
|
||||
cache_key = (user_id, mcp_server.server_id)
|
||||
cached = _byok_cred_cache.get(cache_key)
|
||||
if cached is not None:
|
||||
credential, ts = cached
|
||||
if time.monotonic() - ts < _BYOK_CRED_CACHE_TTL:
|
||||
return 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
|
||||
return await get_user_credential(
|
||||
credential = 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)
|
||||
return credential
|
||||
|
||||
async def _check_byok_credential(
|
||||
mcp_server: MCPServer,
|
||||
|
|
@ -1566,13 +1591,13 @@ if MCP_AVAILABLE:
|
|||
},
|
||||
)
|
||||
|
||||
# Check short-lived in-memory cache before hitting the DB on every tool call
|
||||
# Check shared credential cache before hitting the DB.
|
||||
cache_key = (user_id, mcp_server.server_id)
|
||||
cached = _byok_cred_cache.get(cache_key)
|
||||
if cached is not None:
|
||||
credential_exists, ts = cached
|
||||
cached_cred, ts = cached
|
||||
if time.monotonic() - ts < _BYOK_CRED_CACHE_TTL:
|
||||
if not credential_exists:
|
||||
if cached_cred is None:
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
detail={
|
||||
|
|
@ -1590,21 +1615,19 @@ if MCP_AVAILABLE:
|
|||
)
|
||||
return
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.db import has_user_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
|
||||
|
||||
credential_exists = await has_user_credential(
|
||||
credential = await get_user_credential(
|
||||
prisma_client=prisma_client,
|
||||
user_id=user_id,
|
||||
server_id=mcp_server.server_id,
|
||||
)
|
||||
if len(_byok_cred_cache) >= _BYOK_CRED_CACHE_MAX_SIZE:
|
||||
_byok_cred_cache.clear()
|
||||
_byok_cred_cache[cache_key] = (credential_exists, time.monotonic())
|
||||
if not credential_exists:
|
||||
_write_byok_cred_cache(user_id, mcp_server.server_id, credential)
|
||||
if credential is None:
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
detail={
|
||||
|
|
|
|||
|
|
@ -233,7 +233,7 @@ export const ByokCredentialModal: React.FC<ByokCredentialModalProps> = ({
|
|||
<div className="bg-blue-50 rounded-xl p-4 flex items-start gap-3 mb-6">
|
||||
<LockOutlined className="text-blue-400 mt-0.5 flex-shrink-0" />
|
||||
<p className="text-sm text-blue-700">
|
||||
Your key is encrypted at rest and transmitted securely. It is never shared with third parties.
|
||||
Your key is stored securely and transmitted over HTTPS. It is never shared with third parties.
|
||||
</p>
|
||||
</div>
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue