From 2f24a7ebf233c0dac4a40edd696afaf6febe60db Mon Sep 17 00:00:00 2001
From: Ishaan Jaffer <ishaanjaffer0324@gmail.com>
Date: Wed, 4 Mar 2026 21:09:31 -0800
Subject: [PATCH] fix: address greptile review feedback (greploop iteration 4)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- 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
---
 .../mcp_server/byok_oauth_endpoints.py        |  2 +-
 .../proxy/_experimental/mcp_server/server.py  | 55 +++++++++++++------
 .../mcp_tools/ByokCredentialModal.tsx         |  2 +-
 3 files changed, 41 insertions(+), 18 deletions(-)

diff --git a/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py b/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py
index b4d3157b17e..db18885721a 100644
--- a/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py
+++ b/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py
@@ -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">
diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py
index d97193c06d7..5c063839304 100644
--- a/litellm/proxy/_experimental/mcp_server/server.py
+++ b/litellm/proxy/_experimental/mcp_server/server.py
@@ -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={
diff --git a/ui/litellm-dashboard/src/components/mcp_tools/ByokCredentialModal.tsx b/ui/litellm-dashboard/src/components/mcp_tools/ByokCredentialModal.tsx
index 87117afc1cd..473918c1267 100644
--- a/ui/litellm-dashboard/src/components/mcp_tools/ByokCredentialModal.tsx
+++ b/ui/litellm-dashboard/src/components/mcp_tools/ByokCredentialModal.tsx
@@ -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>