From 4d9fabe11db743f328f55750ed9795b4b80ba485 Mon Sep 17 00:00:00 2001
From: Tin Chi Lo <tin@berri.ai>
Date: Wed, 24 Jun 2026 12:09:08 -0700
Subject: [PATCH] fix(mcp): degrade gracefully when BYOK credential lookup
 errors on protocol tools-list

The protocol tools-list aggregates servers in an asyncio.gather without
return_exceptions, so a raise from the per-user BYOK credential lookup would
abort the entire listing, not just the one server. A missing credential
already returns None (lists without a key); guard the lookup so a DB error
does the same instead of propagating, matching the REST helper's behavior.
---
 .../proxy/_experimental/mcp_server/server.py  | 15 +++-
 .../mcp_server/test_mcp_server.py             | 78 +++++++++++++++++++
 2 files changed, 90 insertions(+), 3 deletions(-)

diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py
index 4f25b67944d..3df7b01e57c 100644
--- a/litellm/proxy/_experimental/mcp_server/server.py
+++ b/litellm/proxy/_experimental/mcp_server/server.py
@@ -1771,9 +1771,18 @@ if MCP_AVAILABLE:
                 # does. Without this a BYOK server lists with the shared token,
                 # which is None, so the upstream rejects the listing.
                 if server.is_byok and not server_auth_header:
-                    server_auth_header = await _get_byok_credential(
-                        server, user_api_key_auth
-                    )
+                    # Don't let a credential-lookup error abort the whole
+                    # multi-server listing; degrade to listing without the key,
+                    # matching the REST path's behavior.
+                    try:
+                        server_auth_header = await _get_byok_credential(
+                            server, user_api_key_auth
+                        )
+                    except Exception as e:
+                        verbose_logger.warning(
+                            f"BYOK credential lookup failed for {server.server_id}; "
+                            f"listing without it: {e}"
+                        )
 
                 # Prefer server-stored per-user OAuth when configured, so a stale
                 # Authorization header from the MCP client cannot override Redis/DB
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
index e5bb305f0f2..1c00372fe5f 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
@@ -1023,6 +1023,84 @@ async def test_get_tools_from_mcp_servers_injects_byok_credential():
     assert captured["mcp_auth_header"] == "user-stored-byok-key"
 
 
+@pytest.mark.asyncio
+async def test_get_tools_from_mcp_servers_byok_lookup_error_does_not_abort_listing():
+    """A DB error during the BYOK credential lookup must degrade to listing
+    without the key, not propagate out of the gather and abort the whole list."""
+    try:
+        from litellm.proxy._experimental.mcp_server import server as mcp_server_module
+        from litellm.proxy._experimental.mcp_server.server import (
+            _get_tools_from_mcp_servers,
+            set_auth_context,
+        )
+    except ImportError:
+        pytest.skip("MCP server not available")
+
+    user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="byok_user")
+    set_auth_context(user_api_key_auth)
+
+    byok_server = MagicMock()
+    byok_server.name = "byok_server"
+    byok_server.alias = "byok"
+    byok_server.allowed_tools = None
+    byok_server.disallowed_tools = None
+    byok_server.server_id = "byok_server"
+    byok_server.server_name = "byok_server"
+    byok_server.auth_type = "authorization"
+    byok_server.extra_headers = None
+    byok_server.is_byok = True
+
+    mock_manager = MagicMock()
+    mock_manager.get_allowed_mcp_servers = AsyncMock(return_value=["byok_server"])
+    mock_manager.get_mcp_server_by_id = lambda server_id: byok_server
+    mock_manager.filter_server_ids_by_ip_with_info = lambda server_ids, client_ip: (
+        server_ids,
+        0,
+    )
+
+    captured = {}
+
+    async def mock_get_tools_from_server(
+        server,
+        mcp_auth_header=None,
+        extra_headers=None,
+        add_prefix=True,
+        raw_headers=None,
+        **kwargs,
+    ):
+        captured["mcp_auth_header"] = mcp_auth_header
+        tool = MagicMock()
+        tool.name = "byok_tool"
+        tool.description = "BYOK tool"
+        tool.inputSchema = {}
+        return [tool]
+
+    mock_manager._get_tools_from_server = mock_get_tools_from_server
+
+    with (
+        patch.object(
+            mcp_server_module,
+            "_get_byok_credential",
+            AsyncMock(side_effect=Exception("transient DB error")),
+        ),
+        patch(
+            "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager",
+            mock_manager,
+        ),
+    ):
+        # Must not raise even though the credential lookup blew up.
+        result = await _get_tools_from_mcp_servers(
+            user_api_key_auth=user_api_key_auth,
+            mcp_auth_header=None,
+            mcp_servers=["byok_server"],
+        )
+
+    # The listing still completed; the server was listed without the BYOK key.
+    assert len(result) == 1
+    assert result[0].name == "byok_tool"
+    assert captured["mcp_auth_header"] is None
+
+
 @pytest.mark.asyncio
 async def test_get_tools_from_mcp_servers_handles_all_servers_failing():
     """Test that _get_tools_from_mcp_servers handles all servers failing gracefully"""