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:
Ishaan Jaffer 2026-03-04 21:09:31 -08:00
parent 8f902f925f
commit 2f24a7ebf2
3 changed files with 41 additions and 18 deletions

View file

@ -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">

View file

@ -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={

View file

@ -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>