mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-06 02:48:13 +00:00
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:
parent
ff5585242b
commit
685483e0ec
2 changed files with 102 additions and 5 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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}"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue