fix(mcp): inject per-user BYOK credential on protocol tools-list path
Some checks failed
LiteLLM Rust / rustfmt, clippy, test (push) Has been cancelled

The REST tools-list injection only covered rest_endpoints.py; the MCP
protocol /tools/list path (_fetch_and_filter_server_tools) still listed a
BYOK server with the shared token, which is None, so the upstream rejected
the listing. Inject the caller's stored per-user key there too, reusing
_get_byok_credential so all five static auth types (including authorization)
format uniformly downstream.
This commit is contained in:
Tin Chi Lo 2026-06-24 10:48:43 -07:00
parent b1affbd72e
commit 21c69751c2
2 changed files with 89 additions and 0 deletions

View file

@ -1766,6 +1766,15 @@ if MCP_AVAILABLE:
user_api_key_auth=user_api_key_auth,
)
# BYOK: inject the caller's stored per-user key so the protocol
# tools-list path authenticates the same way execute_mcp_tool
# 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
)
# Prefer server-stored per-user OAuth when configured, so a stale
# Authorization header from the MCP client cannot override Redis/DB
# (same issue as call_tool in mcp_server_manager: VS Code caches tokens).

View file

@ -943,6 +943,86 @@ async def test_get_tools_from_mcp_servers_continues_when_one_server_fails():
)
@pytest.mark.asyncio
async def test_get_tools_from_mcp_servers_injects_byok_credential():
"""The protocol tools-list path must inject the caller's stored per-user
BYOK credential as the upstream auth header, mirroring execute_mcp_tool.
Without the injection the BYOK server lists with the shared token (None)."""
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(return_value="user-stored-byok-key"),
) as mock_byok,
patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager",
mock_manager,
),
):
result = await _get_tools_from_mcp_servers(
user_api_key_auth=user_api_key_auth,
mcp_auth_header=None,
mcp_servers=["byok_server"],
)
assert len(result) == 1
assert result[0].name == "byok_tool"
# The per-user BYOK key was looked up for this user+server...
mock_byok.assert_awaited_once_with(byok_server, user_api_key_auth)
# ...and injected as the upstream auth header for the listing call.
assert captured["mcp_auth_header"] == "user-stored-byok-key"
@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"""