fix: raise 503 in _get_byok_credential when DB unavailable, guard double-prefix in execute_mcp_tool

- _get_byok_credential now raises HTTP 503 when prisma_client is None,
  matching _check_byok_credential behavior. Previously it returned None
  silently, causing callers to surface a misleading 401 instead of the
  correct 503 infrastructure error.
- execute_mcp_tool prefixed-name fallback now checks name doesn't already
  start with the server prefix before calling add_server_prefix_to_name,
  preventing double-prefixed names like "github-github-get_user".
- Add MCP_TOOL_PREFIX_SEPARATOR to server.py utils imports.
- Add regression tests for both fixes.
This commit is contained in:
Ishaan Jaffer 2026-03-07 09:46:56 -08:00
parent ff5585242b
commit 685483e0ec
2 changed files with 102 additions and 5 deletions

View file

@ -42,6 +42,7 @@ from litellm.proxy._experimental.mcp_server.utils import (
LITELLM_MCP_SERVER_DESCRIPTION,
LITELLM_MCP_SERVER_NAME,
LITELLM_MCP_SERVER_VERSION,
MCP_TOOL_PREFIX_SEPARATOR,
add_server_prefix_to_name,
get_server_prefix,
)
@ -1579,7 +1580,17 @@ if MCP_AVAILABLE:
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
return None
# Without a database we cannot fetch the per-user credential.
# Return a 503 (infrastructure problem) so callers can distinguish
# "no credential" (401) from "credential store unavailable" (503).
raise HTTPException(
status_code=503,
detail={
"error": "byok_store_unavailable",
"server_id": mcp_server.server_id,
"message": "Credential store is not available; cannot fetch BYOK credential.",
},
)
raw = await get_user_credential(
prisma_client=prisma_client,
user_id=user_id,
@ -1812,10 +1823,14 @@ if MCP_AVAILABLE:
# caller used the bare (unprefixed) name and we already resolved the
# server, construct the prefixed name and try again.
if local_tool is None and mcp_server is not None and mcp_server.spec_path:
prefixed_name = add_server_prefix_to_name(name, get_server_prefix(mcp_server))
local_tool = global_mcp_tool_registry.get_tool(prefixed_name)
if local_tool:
name = prefixed_name
server_prefix = get_server_prefix(mcp_server)
# Only add the prefix when the tool name doesn't already carry it,
# otherwise we'd produce double-prefixed names like "github-github-get_user".
if not name.startswith(server_prefix + MCP_TOOL_PREFIX_SEPARATOR):
prefixed_name = add_server_prefix_to_name(name, server_prefix)
local_tool = global_mcp_tool_registry.get_tool(prefixed_name)
if local_tool:
name = prefixed_name
if local_tool:
verbose_logger.debug(f"Executing local registry tool: {name}")
# For BYOK servers the credential must be injected via a ContextVar

View file

@ -582,3 +582,85 @@ def test_spec_path_server_uses_tool_registry():
assert server.is_byok is True
# The spec_path short-circuit in _get_tools_from_server is conditional on this field
assert manager is not None
# ---------------------------------------------------------------------------
# Regression: _get_byok_credential raises 503 when prisma_client is None
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_get_byok_credential_raises_503_when_no_db():
"""_get_byok_credential must raise 503 (not return None) when DB unavailable.
Previously it silently returned None, which caused the caller to surface a
401 instead of the correct 503 infrastructure-error response.
"""
from fastapi import HTTPException
from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPTransport
from litellm.proxy._experimental.mcp_server.server import _get_byok_credential
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.mcp_server.mcp_server_manager import MCPServer
mock_server = MCPServer(
server_id="byok-server",
name="testserver",
server_name="TestServer",
transport=MCPTransport.sse,
is_byok=True,
)
mock_user = MagicMock(spec=UserAPIKeyAuth)
mock_user.user_id = "user-db-unavail"
with patch(
"litellm.proxy.proxy_server.prisma_client",
None,
create=True,
):
with pytest.raises(HTTPException) as exc_info:
await _get_byok_credential(mock_server, mock_user)
assert exc_info.value.status_code == 503
assert "byok_store_unavailable" in str(exc_info.value.detail)
# ---------------------------------------------------------------------------
# Regression: no double-prefix in execute_mcp_tool
# ---------------------------------------------------------------------------
def test_no_double_prefix_for_already_prefixed_tool_name():
"""add_server_prefix_to_name must NOT be called when name already has the prefix.
Without the guard, a tool name like "github-get_user" would become
"github-github-get_user" and the registry lookup would always miss.
"""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPTransport
from litellm.proxy._experimental.mcp_server.utils import (
MCP_TOOL_PREFIX_SEPARATOR,
add_server_prefix_to_name,
get_server_prefix,
)
from litellm.types.mcp_server.mcp_server_manager import MCPServer
server = MCPServer(
server_id="s1",
name="github",
server_name="github",
transport=MCPTransport.sse,
spec_path="https://example.com/spec.json",
is_byok=False,
)
server_prefix = get_server_prefix(server)
tool_name = add_server_prefix_to_name("get_user", server_prefix) # "github-get_user"
# Simulate the guard: only prefix when not already prefixed
if not tool_name.startswith(server_prefix + MCP_TOOL_PREFIX_SEPARATOR):
result = add_server_prefix_to_name(tool_name, server_prefix)
else:
result = tool_name
assert result == tool_name, f"Expected no double-prefix, got: {result}"
assert result.count(server_prefix) == 1, f"Prefix appears more than once: {result}"