mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
Merge branch 'BerriAI:litellm_internal_staging' into fix/cached-token-custom-cost
This commit is contained in:
Gabriel Zirondi
GitHub
commit
dbb7f93b2f
9 changed files with 1120 additions and 53 deletions
|
|
@ -50,8 +50,11 @@ from litellm.proxy._experimental.mcp_server.oauth2_token_cache import resolve_mc
|
|||
from litellm.proxy._experimental.mcp_server.utils import (
|
||||
MCP_TOOL_PREFIX_SEPARATOR,
|
||||
add_server_prefix_to_name,
|
||||
compute_short_server_prefix,
|
||||
get_server_prefix,
|
||||
is_short_mcp_tool_prefix_enabled,
|
||||
is_tool_name_prefixed,
|
||||
iter_known_server_prefixes,
|
||||
merge_mcp_headers,
|
||||
normalize_server_name,
|
||||
split_server_prefix_from_name,
|
||||
|
|
@ -364,6 +367,7 @@ class MCPServerManager:
|
|||
aws_session_name=server_config.get("aws_session_name", None),
|
||||
instructions=server_config.get("instructions", None),
|
||||
)
|
||||
self._assign_unique_short_prefix(new_server)
|
||||
self.config_mcp_servers[server_id] = new_server
|
||||
|
||||
# Check if this is an OpenAPI-based server
|
||||
|
|
@ -726,6 +730,7 @@ class MCPServerManager:
|
|||
try:
|
||||
if mcp_server.server_id not in self.registry:
|
||||
new_server = await self.build_mcp_server_from_table(mcp_server)
|
||||
self._assign_unique_short_prefix(new_server)
|
||||
self.registry[mcp_server.server_id] = new_server
|
||||
await self._maybe_register_openapi_tools(new_server)
|
||||
verbose_logger.debug(f"Added MCP Server: {new_server.name}")
|
||||
|
|
@ -738,6 +743,12 @@ class MCPServerManager:
|
|||
try:
|
||||
if mcp_server.server_id in self.registry:
|
||||
new_server = await self.build_mcp_server_from_table(mcp_server)
|
||||
# Carry the previously-resolved short prefix across so the
|
||||
# tool names stay stable for clients holding cached lists.
|
||||
existing_prefix = self.registry[mcp_server.server_id].short_prefix
|
||||
if existing_prefix and not new_server.short_prefix:
|
||||
new_server.short_prefix = existing_prefix
|
||||
self._assign_unique_short_prefix(new_server)
|
||||
self.registry[mcp_server.server_id] = new_server
|
||||
await self._maybe_register_openapi_tools(new_server)
|
||||
verbose_logger.debug(f"Updated MCP Server: {new_server.name}")
|
||||
|
|
@ -1236,7 +1247,11 @@ class MCPServerManager:
|
|||
|
||||
## HANDLE OPENAPI TOOLS
|
||||
if server.spec_path:
|
||||
_tools = global_mcp_tool_registry.list_tools(tool_prefix=server.name)
|
||||
# OpenAPI tools were stored in the registry under the prefix
|
||||
# active at registration time — fetch by that same prefix.
|
||||
_tools = global_mcp_tool_registry.list_tools(
|
||||
tool_prefix=get_server_prefix(server)
|
||||
)
|
||||
tools = global_mcp_tool_registry.convert_tools_to_mcp_sdk_tool_type(
|
||||
_tools
|
||||
)
|
||||
|
|
@ -1810,6 +1825,63 @@ class MCPServerManager:
|
|||
verbose_logger.warning(f"Error listing tools from {server_name}: {str(e)}")
|
||||
return []
|
||||
|
||||
_SHORT_PREFIX_MAX_REHASH_ATTEMPTS = 1024
|
||||
|
||||
def _assign_unique_short_prefix(self, server: MCPServer) -> None:
|
||||
"""Resolve and cache a collision-free short tool prefix on ``server``.
|
||||
|
||||
Called at registration time for every MCP server entering the
|
||||
registry. Mutates ``server.short_prefix`` in place. No-ops when
|
||||
``LITELLM_USE_SHORT_MCP_TOOL_PREFIX`` is disabled, when the server
|
||||
has no ``server_id`` (synthetic temp-server objects), or when a
|
||||
prefix is already cached.
|
||||
|
||||
Collision strategy: take the natural hash; if it's already used by
|
||||
a *different* server in the combined registry, rehash with an
|
||||
incrementing attempt counter until we find an unused slot. The
|
||||
attempt counter is folded into the hash so the resulting prefix is
|
||||
still deterministic for a given (server_id, set-of-other-server-ids)
|
||||
pair within one process.
|
||||
"""
|
||||
if not is_short_mcp_tool_prefix_enabled():
|
||||
return
|
||||
if server.short_prefix:
|
||||
return
|
||||
if not server.server_id:
|
||||
return
|
||||
|
||||
used: Dict[str, str] = {}
|
||||
for other in self.get_registry().values():
|
||||
if other.server_id == server.server_id:
|
||||
continue
|
||||
if other.short_prefix:
|
||||
used[other.short_prefix] = other.server_id
|
||||
|
||||
for attempt in range(self._SHORT_PREFIX_MAX_REHASH_ATTEMPTS):
|
||||
candidate = compute_short_server_prefix(server.server_id, attempt=attempt)
|
||||
if candidate not in used:
|
||||
server.short_prefix = candidate
|
||||
if attempt > 0:
|
||||
verbose_logger.info(
|
||||
"MCP short-prefix collision resolved for server %s: "
|
||||
"natural hash collided with %s, using rehashed prefix "
|
||||
"%s (attempt=%d).",
|
||||
server.server_id,
|
||||
used.get(
|
||||
compute_short_server_prefix(server.server_id, attempt=0),
|
||||
"<unknown>",
|
||||
),
|
||||
candidate,
|
||||
attempt,
|
||||
)
|
||||
return
|
||||
|
||||
raise RuntimeError(
|
||||
f"Unable to assign a unique short MCP tool prefix for server "
|
||||
f"{server.server_id} after {self._SHORT_PREFIX_MAX_REHASH_ATTEMPTS} "
|
||||
"attempts; the 3-character prefix space is too crowded."
|
||||
)
|
||||
|
||||
def _create_prefixed_tools(
|
||||
self, tools: List[MCPTool], server: MCPServer, add_prefix: bool = True
|
||||
) -> List[MCPTool]:
|
||||
|
|
@ -1838,9 +1910,13 @@ class MCPServerManager:
|
|||
tool_copy.name = name_to_use
|
||||
prefixed_tools.append(tool_copy)
|
||||
|
||||
# Update tool to server mapping for resolution (support both forms)
|
||||
# Register every known prefix form (alias, server_name, server_id,
|
||||
# short ID) so call_tool can resolve regardless of which form a
|
||||
# caller / cached client is using.
|
||||
self.tool_name_to_mcp_server_name_mapping[original_name] = prefix
|
||||
self.tool_name_to_mcp_server_name_mapping[prefixed_name] = prefix
|
||||
for known_prefix in iter_known_server_prefixes(server):
|
||||
qualified = add_server_prefix_to_name(original_name, known_prefix)
|
||||
self.tool_name_to_mcp_server_name_mapping[qualified] = prefix
|
||||
|
||||
verbose_logger.info(
|
||||
f"Successfully fetched {len(prefixed_tools)} tools from server {server.name}"
|
||||
|
|
@ -2601,37 +2677,43 @@ class MCPServerManager:
|
|||
Returns:
|
||||
MCPServer if found, None otherwise
|
||||
"""
|
||||
registry_servers = list(self.get_registry().values())
|
||||
|
||||
# Build prefix → server lookup covering every known form a tool name
|
||||
# may take (alias / server_name / server_id / short ID). This is what
|
||||
# makes the short-prefix mode work without breaking historical names.
|
||||
prefix_to_server: Dict[str, MCPServer] = {}
|
||||
for server in registry_servers:
|
||||
for known_prefix in iter_known_server_prefixes(server):
|
||||
normalised = normalize_server_name(known_prefix)
|
||||
prefix_to_server.setdefault(normalised, server)
|
||||
|
||||
# First try with the original tool name
|
||||
if tool_name in self.tool_name_to_mcp_server_name_mapping:
|
||||
server_name = self.tool_name_to_mcp_server_name_mapping[tool_name]
|
||||
for server in self.get_registry().values():
|
||||
if normalize_server_name(server.name) == normalize_server_name(
|
||||
server_name
|
||||
):
|
||||
normalised_lookup = normalize_server_name(server_name)
|
||||
if normalised_lookup in prefix_to_server:
|
||||
return prefix_to_server[normalised_lookup]
|
||||
for server in registry_servers:
|
||||
if normalize_server_name(server.name) == normalised_lookup:
|
||||
return server
|
||||
|
||||
# If not found and tool name is prefixed, try extracting server name from prefix
|
||||
known_prefixes = {
|
||||
normalize_server_name(get_server_prefix(s))
|
||||
for s in self.get_registry().values()
|
||||
if get_server_prefix(s)
|
||||
}
|
||||
if is_tool_name_prefixed(tool_name, known_server_prefixes=known_prefixes):
|
||||
# If not found and tool name is prefixed, extract the prefix and
|
||||
# match against any known form.
|
||||
if is_tool_name_prefixed(
|
||||
tool_name, known_server_prefixes=set(prefix_to_server.keys())
|
||||
):
|
||||
(
|
||||
original_tool_name,
|
||||
server_name_from_prefix,
|
||||
) = split_server_prefix_from_name(tool_name)
|
||||
if original_tool_name in self.tool_name_to_mcp_server_name_mapping:
|
||||
for server in self.get_registry().values():
|
||||
if server.server_name is None:
|
||||
if normalize_server_name(server.name) == normalize_server_name(
|
||||
server_name_from_prefix
|
||||
):
|
||||
return server
|
||||
elif normalize_server_name(
|
||||
server.server_name
|
||||
) == normalize_server_name(server_name_from_prefix):
|
||||
return server
|
||||
normalised_prefix = normalize_server_name(server_name_from_prefix)
|
||||
matched_server = prefix_to_server.get(normalised_prefix)
|
||||
if matched_server is not None and (
|
||||
original_tool_name in self.tool_name_to_mcp_server_name_mapping
|
||||
or tool_name in self.tool_name_to_mcp_server_name_mapping
|
||||
):
|
||||
return matched_server
|
||||
|
||||
return None
|
||||
|
||||
|
|
@ -2666,6 +2748,9 @@ class MCPServerManager:
|
|||
previous_registry = self.registry
|
||||
new_registry: Dict[str, MCPServer] = {}
|
||||
|
||||
# Stage one: build every server. Stage two assigns short prefixes
|
||||
# against the *full* set so dedup is deterministic regardless of
|
||||
# iteration order.
|
||||
for server in db_mcp_servers:
|
||||
existing_server = previous_registry.get(server.server_id)
|
||||
|
||||
|
|
@ -2689,10 +2774,21 @@ class MCPServerManager:
|
|||
f"Building server from DB: {server.server_id} ({server.server_name})"
|
||||
)
|
||||
new_server = await self.build_mcp_server_from_table(server)
|
||||
# Carry the cached short_prefix from the previous registry entry
|
||||
# (if any) so the prefix is stable across reloads.
|
||||
if existing_server is not None and existing_server.short_prefix:
|
||||
new_server.short_prefix = existing_server.short_prefix
|
||||
new_registry[server.server_id] = new_server
|
||||
await self._maybe_register_openapi_tools(new_server)
|
||||
|
||||
# Swap in the new registry first so _assign_unique_short_prefix
|
||||
# sees the complete set when checking for collisions.
|
||||
self.registry = new_registry
|
||||
for new_server in new_registry.values():
|
||||
self._assign_unique_short_prefix(new_server)
|
||||
# Register OpenAPI tools *after* the final short prefix is assigned
|
||||
# so the tools are stored in the global registry under the same
|
||||
# prefix that lookups will use.
|
||||
await self._maybe_register_openapi_tools(new_server)
|
||||
|
||||
verbose_logger.debug(
|
||||
"MCP registry refreshed (%s servers in registry)", len(new_registry)
|
||||
|
|
|
|||
|
|
@ -49,6 +49,7 @@ from litellm.proxy._experimental.mcp_server.utils import (
|
|||
LITELLM_MCP_SERVER_VERSION,
|
||||
add_server_prefix_to_name,
|
||||
get_server_prefix,
|
||||
iter_known_server_prefixes,
|
||||
)
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.ip_address_utils import IPAddressUtils
|
||||
|
|
@ -711,13 +712,7 @@ if MCP_AVAILABLE:
|
|||
for server in allowed_mcp_servers:
|
||||
if server:
|
||||
match_list = [
|
||||
s.lower()
|
||||
for s in [
|
||||
server.alias,
|
||||
server.server_name,
|
||||
server.server_id,
|
||||
]
|
||||
if s is not None
|
||||
s.lower() for s in iter_known_server_prefixes(server) if s
|
||||
]
|
||||
|
||||
if server_or_group.lower() in match_list:
|
||||
|
|
@ -2031,11 +2026,13 @@ if MCP_AVAILABLE:
|
|||
# Remove prefix from tool name for logging and processing
|
||||
original_tool_name, server_name = split_server_prefix_from_name(name)
|
||||
|
||||
# If tool name is unprefixed, resolve its server so we can enforce permissions
|
||||
if not server_name:
|
||||
mcp_server = global_mcp_server_manager._get_mcp_server_from_tool_name(name)
|
||||
if mcp_server:
|
||||
server_name = mcp_server.name
|
||||
# Resolve the actual MCP server up-front so the permission check uses
|
||||
# the canonical server.name even when the tool name is prefixed with a
|
||||
# short ID (LITELLM_USE_SHORT_MCP_TOOL_PREFIX) that doesn't match the
|
||||
# server's display name directly.
|
||||
mcp_server = global_mcp_server_manager._get_mcp_server_from_tool_name(name)
|
||||
if mcp_server is not None:
|
||||
server_name = mcp_server.name
|
||||
|
||||
# Only enforce server-level permissions when we can resolve a server
|
||||
if server_name:
|
||||
|
|
|
|||
|
|
@ -2,10 +2,11 @@
|
|||
MCP Server Utilities
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, Mapping, Optional, Tuple
|
||||
from typing import Any, Dict, Iterator, Mapping, Optional, Tuple
|
||||
|
||||
import os
|
||||
import hashlib
|
||||
import importlib
|
||||
import os
|
||||
|
||||
# Constants
|
||||
LITELLM_MCP_SERVER_NAME = "litellm-mcp-server"
|
||||
|
|
@ -14,6 +15,89 @@ LITELLM_MCP_SERVER_DESCRIPTION = "MCP Server for LiteLLM"
|
|||
MCP_TOOL_PREFIX_SEPARATOR = os.environ.get("MCP_TOOL_PREFIX_SEPARATOR", "-")
|
||||
MCP_TOOL_PREFIX_FORMAT = "{server_name}{separator}{tool_name}"
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Short-ID tool prefix (opt-in)
|
||||
# ---------------------------------------------------------------------------
|
||||
# When LITELLM_USE_SHORT_MCP_TOOL_PREFIX is truthy the prefix attached to MCP
|
||||
# tool / prompt / resource / resource-template names switches from the
|
||||
# (potentially long) human-readable server name to a deterministic three
|
||||
# character ID derived from the server's ``server_id``.
|
||||
#
|
||||
# Why three characters?
|
||||
# * The first character is restricted to 52 alphabetic characters
|
||||
# ([A-Za-z]) and the remaining two characters use the full base62
|
||||
# alphabet ([0-9A-Za-z]). That guarantees the prefix never starts
|
||||
# with a digit so it remains a valid identifier for every model API
|
||||
# (some providers historically required a leading alphabetic char).
|
||||
# * 52 * 62 * 62 = 199_888 distinct IDs. The chance of a real local
|
||||
# tool name happening to begin with the exact prefix LiteLLM assigned
|
||||
# to a given MCP server is negligible in practice.
|
||||
# * The IDs are short enough that prefixed tool names stay well under
|
||||
# the 60-character upper bound enforced by some model APIs (Anthropic
|
||||
# etc.) even for long upstream tool names.
|
||||
# * The mapping is deterministic (SHA-256 of ``server_id`` → three
|
||||
# characters drawn from the alphabets above), so the prefix is stable
|
||||
# across processes, workers and restarts without any persistence
|
||||
# layer. Two servers with different ``server_id`` values can in
|
||||
# principle hash to the same three chars; that natural-hash collision
|
||||
# IS a routing-correctness issue (the second registrant would otherwise
|
||||
# have its tools misrouted to the first), so registration goes through
|
||||
# ``MCPServerManager._assign_unique_short_prefix`` which rehashes with
|
||||
# a deterministic attempt counter until it finds an unused prefix and
|
||||
# caches the result on ``MCPServer.short_prefix``. A collision is
|
||||
# logged at INFO when it happens.
|
||||
#
|
||||
# This flag is intentionally opt-in for the first release so customers can
|
||||
# migrate. It will become the default in a future release.
|
||||
SHORT_MCP_TOOL_PREFIX_LENGTH = 3
|
||||
_BASE62_ALPHABET = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
|
||||
# Subset of _BASE62_ALPHABET used for the *first* character only, to
|
||||
# guarantee the prefix never starts with a digit.
|
||||
_BASE52_ALPHA_ALPHABET = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
|
||||
|
||||
|
||||
def is_short_mcp_tool_prefix_enabled() -> bool:
|
||||
"""Return True when the short-ID tool prefix mode is enabled.
|
||||
|
||||
Read at call time (not import time) so tests and runtime config changes
|
||||
take effect without reimporting the module.
|
||||
"""
|
||||
raw = os.environ.get("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "")
|
||||
return raw.strip().lower() in ("1", "true", "yes", "on")
|
||||
|
||||
|
||||
def compute_short_server_prefix(server_id: str, attempt: int = 0) -> str:
|
||||
"""Derive the deterministic three-character prefix for a server.
|
||||
|
||||
Uses SHA-256 of ``f"{server_id}#{attempt}"`` and folds the first eight
|
||||
bytes into a fixed-length string whose first character is drawn from
|
||||
``_BASE52_ALPHA_ALPHABET`` (so the prefix never starts with a digit)
|
||||
and whose remaining characters are drawn from the full base62
|
||||
alphabet. Pass ``attempt > 0`` to rehash to a different prefix when
|
||||
the natural hash collides with a prefix already assigned to another
|
||||
server (see ``MCPServerManager._assign_unique_short_prefix``). An
|
||||
empty ``server_id`` raises ``ValueError`` — short prefixes require a
|
||||
stable identifier to be deterministic.
|
||||
"""
|
||||
if not server_id:
|
||||
raise ValueError("compute_short_server_prefix requires a non-empty server_id")
|
||||
|
||||
seed = server_id if attempt == 0 else f"{server_id}#{attempt}"
|
||||
digest = hashlib.sha256(seed.encode("utf-8")).digest()
|
||||
value = int.from_bytes(digest[:8], "big")
|
||||
|
||||
# Build chars from least-significant to most-significant; we reverse
|
||||
# at the end so the first emitted char comes from the high-order
|
||||
# bits of the digest (which is the position we constrain to be
|
||||
# alphabetic).
|
||||
chars = []
|
||||
for position in range(SHORT_MCP_TOOL_PREFIX_LENGTH):
|
||||
is_first_char = position == SHORT_MCP_TOOL_PREFIX_LENGTH - 1
|
||||
alphabet = _BASE52_ALPHA_ALPHABET if is_first_char else _BASE62_ALPHABET
|
||||
value, idx = divmod(value, len(alphabet))
|
||||
chars.append(alphabet[idx])
|
||||
return "".join(reversed(chars))
|
||||
|
||||
|
||||
def is_mcp_available() -> bool:
|
||||
"""
|
||||
|
|
@ -82,7 +166,25 @@ def add_server_prefix_to_name(name: str, server_name: str) -> str:
|
|||
|
||||
|
||||
def get_server_prefix(server: Any) -> str:
|
||||
"""Return the prefix for a server: alias if present, else server_name, else server_id"""
|
||||
"""Return the prefix for a server.
|
||||
|
||||
When the short-prefix mode is enabled (``LITELLM_USE_SHORT_MCP_TOOL_PREFIX``)
|
||||
a three-character base62 ID is returned. We prefer the cached
|
||||
``server.short_prefix`` value when set — that field is populated at
|
||||
registration time by ``MCPServerManager._assign_unique_short_prefix``
|
||||
and resolves natural-hash collisions deterministically — and only fall
|
||||
back to the natural hash for ad-hoc / temp-server objects without a
|
||||
cached value. In default mode the historical behaviour is preserved:
|
||||
alias if present, else server_name, else server_id.
|
||||
"""
|
||||
if is_short_mcp_tool_prefix_enabled():
|
||||
cached = getattr(server, "short_prefix", None)
|
||||
if cached:
|
||||
return cached
|
||||
server_id = getattr(server, "server_id", None)
|
||||
if server_id:
|
||||
return compute_short_server_prefix(server_id)
|
||||
|
||||
if hasattr(server, "alias") and server.alias:
|
||||
return server.alias
|
||||
if hasattr(server, "server_name") and server.server_name:
|
||||
|
|
@ -92,6 +194,36 @@ def get_server_prefix(server: Any) -> str:
|
|||
return ""
|
||||
|
||||
|
||||
def iter_known_server_prefixes(server: Any) -> Iterator[str]:
|
||||
"""Yield every prefix form that may appear in tool names for ``server``.
|
||||
|
||||
Always includes the *current* prefix returned by ``get_server_prefix``.
|
||||
Additionally yields the historical (alias / server_name / server_id) and
|
||||
short-ID forms so the routing layer can resolve tool names regardless of
|
||||
which prefix mode was active when the client first observed them.
|
||||
"""
|
||||
seen = set()
|
||||
|
||||
def _emit(value: Optional[str]) -> Iterator[str]:
|
||||
if value and value not in seen:
|
||||
seen.add(value)
|
||||
yield value
|
||||
|
||||
yield from _emit(get_server_prefix(server))
|
||||
yield from _emit(getattr(server, "short_prefix", None))
|
||||
|
||||
server_id = getattr(server, "server_id", None)
|
||||
if server_id:
|
||||
try:
|
||||
yield from _emit(compute_short_server_prefix(server_id))
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
yield from _emit(getattr(server, "alias", None))
|
||||
yield from _emit(getattr(server, "server_name", None))
|
||||
yield from _emit(server_id)
|
||||
|
||||
|
||||
def split_server_prefix_from_name(prefixed_name: str) -> Tuple[str, str]:
|
||||
"""Return the unprefixed name plus the server name used as prefix."""
|
||||
if MCP_TOOL_PREFIX_SEPARATOR in prefixed_name:
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
from typing import Union
|
||||
from typing import Any, Awaitable, Callable, Optional, Union
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy._types import (
|
||||
DB_CONNECTION_ERROR_TYPES,
|
||||
ProxyErrorTypes,
|
||||
|
|
@ -123,3 +124,138 @@ class PrismaDBExceptionHandler:
|
|||
):
|
||||
return None
|
||||
raise e
|
||||
|
||||
|
||||
# Default fallback timeouts when neither the caller nor the prisma_client
|
||||
# expose `_db_auth_reconnect_timeout_seconds` / `_db_auth_reconnect_lock_timeout_seconds`.
|
||||
# Match the auth path's existing defaults so behavior is uniform across read paths.
|
||||
_DEFAULT_RECONNECT_TIMEOUT_SECONDS = 2.0
|
||||
_DEFAULT_RECONNECT_LOCK_TIMEOUT_SECONDS = 0.1
|
||||
|
||||
|
||||
def _coerce_timeout(value: Any, fallback: float) -> float:
|
||||
"""Return `value` if it is a real int/float, else `fallback`. Guards
|
||||
against tests that mock `prisma_client` and leave the timeout slots as
|
||||
MagicMock instances."""
|
||||
if isinstance(value, (int, float)) and not isinstance(value, bool):
|
||||
return float(value)
|
||||
return fallback
|
||||
|
||||
|
||||
async def call_with_db_reconnect_retry(
|
||||
prisma_client: Any,
|
||||
coro_factory: Callable[[], Awaitable[Any]],
|
||||
*,
|
||||
reason: str,
|
||||
timeout_seconds: Optional[float] = None,
|
||||
lock_timeout_seconds: Optional[float] = None,
|
||||
) -> Any:
|
||||
"""Run a Prisma read coroutine with one transport-reconnect-and-retry.
|
||||
|
||||
The canonical "self-heal a transient DB transport blip" wrapper used by
|
||||
`PrismaClient.get_generic_data` and other read paths. Mirrors the inline
|
||||
pattern in `auth_checks._fetch_key_object_from_db_with_reconnect` so we
|
||||
have a single implementation rather than three drifting copies.
|
||||
|
||||
Behavior:
|
||||
1. Await `coro_factory()`. On success, return its value.
|
||||
2. On exception, if it is NOT a transport error (per
|
||||
`is_database_transport_error`), re-raise — data-layer errors like
|
||||
`UniqueViolationError` mean the DB is reachable, reconnect would be
|
||||
pointless.
|
||||
3. If `prisma_client` does not expose `attempt_db_reconnect`, re-raise.
|
||||
This guards against partial stand-ins / older clients in tests.
|
||||
4. Call `prisma_client.attempt_db_reconnect(reason=...)`. If it returns
|
||||
False (cooldown / lock contention / reconnect failure), re-raise.
|
||||
5. Otherwise await `coro_factory()` a second time and return / propagate
|
||||
its result. At-most-one retry by construction — no infinite loop.
|
||||
|
||||
`coro_factory` MUST be a zero-arg callable that returns a fresh awaitable
|
||||
on each call. Passing an already-awaited coroutine would fail on retry
|
||||
with `RuntimeError: cannot reuse already awaited coroutine`.
|
||||
|
||||
`reason` should follow `<subsystem>_<operation>_<table>_failure` so
|
||||
telemetry distinguishes between fan-out callers (e.g.
|
||||
`_update_config_from_db` issues four concurrent reads).
|
||||
|
||||
Args:
|
||||
prisma_client: The `PrismaClient` (or stand-in) that owns
|
||||
`attempt_db_reconnect` and the `_db_auth_reconnect_*` defaults.
|
||||
coro_factory: Zero-arg callable returning the read awaitable.
|
||||
reason: Telemetry tag forwarded to `attempt_db_reconnect`.
|
||||
timeout_seconds: Optional override for the reconnect cycle timeout.
|
||||
Defaults to `prisma_client._db_auth_reconnect_timeout_seconds`,
|
||||
then to 2.0s.
|
||||
lock_timeout_seconds: Optional override for how long the helper will
|
||||
wait to acquire the reconnect lock. Defaults to
|
||||
`prisma_client._db_auth_reconnect_lock_timeout_seconds`, then to
|
||||
0.1s.
|
||||
|
||||
Returns:
|
||||
Whatever `coro_factory()` returns (on first or second attempt).
|
||||
|
||||
Raises:
|
||||
Whatever `coro_factory()` raises if the failure is not a transport
|
||||
error, or if the reconnect attempt does not succeed, or if the retry
|
||||
also fails.
|
||||
"""
|
||||
try:
|
||||
return await coro_factory()
|
||||
except Exception as first_exc:
|
||||
if not PrismaDBExceptionHandler.is_database_transport_error(first_exc):
|
||||
raise
|
||||
if not hasattr(prisma_client, "attempt_db_reconnect"):
|
||||
raise
|
||||
|
||||
resolved_timeout = _coerce_timeout(
|
||||
(
|
||||
timeout_seconds
|
||||
if timeout_seconds is not None
|
||||
else getattr(prisma_client, "_db_auth_reconnect_timeout_seconds", None)
|
||||
),
|
||||
_DEFAULT_RECONNECT_TIMEOUT_SECONDS,
|
||||
)
|
||||
resolved_lock_timeout = _coerce_timeout(
|
||||
(
|
||||
lock_timeout_seconds
|
||||
if lock_timeout_seconds is not None
|
||||
else getattr(
|
||||
prisma_client, "_db_auth_reconnect_lock_timeout_seconds", None
|
||||
)
|
||||
),
|
||||
_DEFAULT_RECONNECT_LOCK_TIMEOUT_SECONDS,
|
||||
)
|
||||
|
||||
verbose_proxy_logger.warning(
|
||||
"DB transport error on read; attempting reconnect-and-retry. reason=%s error=%s",
|
||||
reason,
|
||||
first_exc,
|
||||
)
|
||||
|
||||
# Preserve the original transport error in telemetry. If
|
||||
# `attempt_db_reconnect` itself raises (e.g. lock cancellation, timer
|
||||
# error, unexpected internal failure), surfacing that exception
|
||||
# instead of `first_exc` would mask the actual DB transport problem
|
||||
# in `failure_handler` / `db_exceptions` alerts. Chain the reconnect
|
||||
# error as the cause for debuggability without losing the original.
|
||||
try:
|
||||
did_reconnect = await prisma_client.attempt_db_reconnect(
|
||||
reason=reason,
|
||||
timeout_seconds=resolved_timeout,
|
||||
lock_timeout_seconds=resolved_lock_timeout,
|
||||
)
|
||||
except Exception as reconnect_exc:
|
||||
verbose_proxy_logger.warning(
|
||||
"DB reconnect attempt raised; preserving original transport error. "
|
||||
"reason=%s reconnect_error=%s",
|
||||
reason,
|
||||
reconnect_exc,
|
||||
)
|
||||
raise first_exc from reconnect_exc
|
||||
if not did_reconnect:
|
||||
raise
|
||||
|
||||
# At most one retry. If the retry also raises a transport error, we
|
||||
# propagate — repeated reconnect-loops are the watchdog's job, not
|
||||
# this helper's.
|
||||
return await coro_factory()
|
||||
|
|
|
|||
|
|
@ -106,7 +106,10 @@ from litellm.proxy.db.create_views import (
|
|||
should_create_missing_views,
|
||||
)
|
||||
from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
|
||||
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
|
||||
from litellm.proxy.db.exception_handler import (
|
||||
PrismaDBExceptionHandler,
|
||||
call_with_db_reconnect_retry,
|
||||
)
|
||||
from litellm.proxy.db.log_db_metrics import log_db_metrics
|
||||
from litellm.proxy.db.prisma_client import PrismaWrapper
|
||||
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
|
||||
|
|
@ -2779,30 +2782,42 @@ class PrismaClient:
|
|||
table_name: Literal["users", "keys", "config", "spend"],
|
||||
):
|
||||
"""
|
||||
Generic implementation of get data
|
||||
Generic implementation of get data.
|
||||
|
||||
Self-heals across a single transient transport blip via
|
||||
`call_with_db_reconnect_retry`: on `httpx.ReadError` /
|
||||
`ClientNotConnectedError` / similar, attempt one DB reconnect and
|
||||
retry once before surfacing the failure. Restores the 1.82.6 behavior
|
||||
that was lost in 1.83.x — see issue #25143.
|
||||
"""
|
||||
start_time = time.time()
|
||||
try:
|
||||
|
||||
async def _do_query():
|
||||
if table_name == "users":
|
||||
response = await self.db.litellm_usertable.find_first(
|
||||
return await self.db.litellm_usertable.find_first(
|
||||
where={key: value} # type: ignore
|
||||
)
|
||||
elif table_name == "keys":
|
||||
response = await self.db.litellm_verificationtoken.find_first( # type: ignore
|
||||
return await self.db.litellm_verificationtoken.find_first( # type: ignore
|
||||
where={key: value} # type: ignore
|
||||
)
|
||||
elif table_name == "config":
|
||||
response = await self.db.litellm_config.find_first( # type: ignore
|
||||
return await self.db.litellm_config.find_first( # type: ignore
|
||||
where={key: value} # type: ignore
|
||||
)
|
||||
elif table_name == "spend":
|
||||
response = await self.db.l.find_first( # type: ignore
|
||||
return await self.db.l.find_first( # type: ignore
|
||||
where={key: value} # type: ignore
|
||||
)
|
||||
return response
|
||||
except Exception as e:
|
||||
import traceback
|
||||
return None
|
||||
|
||||
try:
|
||||
return await call_with_db_reconnect_retry(
|
||||
self,
|
||||
_do_query,
|
||||
reason=f"prisma_get_generic_data_{table_name}_lookup_failure",
|
||||
)
|
||||
except Exception as e:
|
||||
error_msg = f"LiteLLM Prisma Client Exception get_generic_data: {str(e)}"
|
||||
verbose_proxy_logger.error(error_msg)
|
||||
error_msg = error_msg + "\nException Type: {}".format(type(e))
|
||||
|
|
@ -4204,7 +4219,6 @@ class PrismaClient:
|
|||
)
|
||||
self._reap_all_zombies()
|
||||
self._cleanup_engine_watcher()
|
||||
self._engine_confirmed_dead = False
|
||||
|
||||
async def _do_heavy_reconnect() -> None:
|
||||
db_url = os.getenv("DATABASE_URL", "")
|
||||
|
|
@ -4217,6 +4231,12 @@ class PrismaClient:
|
|||
await self._start_engine_watcher()
|
||||
|
||||
await asyncio.wait_for(_do_heavy_reconnect(), timeout=effective_timeout)
|
||||
# Only clear the "dead engine" flag after the heavy reconnect
|
||||
# actually completed. If `_do_heavy_reconnect()` raises (timeout,
|
||||
# missing DATABASE_URL, recreate failure), the flag stays True so
|
||||
# the next attempt re-enters the heavy branch instead of silently
|
||||
# demoting to the lightweight path.
|
||||
self._engine_confirmed_dead = False
|
||||
else:
|
||||
verbose_proxy_logger.debug(
|
||||
"Performing Prisma DB reconnect (engine alive or unknown)."
|
||||
|
|
|
|||
|
|
@ -81,6 +81,12 @@ class MCPServer(BaseModel):
|
|||
# Defaults to the token's expires_in minus the expiry buffer, or
|
||||
# MCP_PER_USER_TOKEN_DEFAULT_TTL when expires_in is absent.
|
||||
token_storage_ttl_seconds: Optional[int] = None
|
||||
# Resolved short-ID tool prefix when LITELLM_USE_SHORT_MCP_TOOL_PREFIX is
|
||||
# enabled. Set by ``MCPServerManager._assign_unique_short_prefix`` at
|
||||
# registration time so that natural-hash collisions between two
|
||||
# different ``server_id`` values are bumped deterministically. Left
|
||||
# ``None`` in default-prefix mode.
|
||||
short_prefix: Optional[str] = None
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
@property
|
||||
|
|
|
|||
|
|
@ -0,0 +1,302 @@
|
|||
"""
|
||||
Tests for the short-ID MCP tool prefix (LITELLM_USE_SHORT_MCP_TOOL_PREFIX).
|
||||
|
||||
The short-prefix mode swaps the historical alias/server_name prefix on
|
||||
tool names for a deterministic three-character base62 ID derived from the
|
||||
server's ``server_id``. This keeps tool names well below the 60-char
|
||||
upper bound enforced by some model APIs while remaining stable across
|
||||
processes/restarts and tolerant of mixed-version clients.
|
||||
"""
|
||||
|
||||
from typing import List
|
||||
|
||||
import pytest
|
||||
from mcp.types import Tool as MCPTool
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager
|
||||
from litellm.proxy._experimental.mcp_server.utils import (
|
||||
SHORT_MCP_TOOL_PREFIX_LENGTH,
|
||||
add_server_prefix_to_name,
|
||||
compute_short_server_prefix,
|
||||
get_server_prefix,
|
||||
is_short_mcp_tool_prefix_enabled,
|
||||
iter_known_server_prefixes,
|
||||
)
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
|
||||
|
||||
def _make_server(
|
||||
*,
|
||||
server_id: str = "abcdef-1234",
|
||||
server_name: str = "github_onprem",
|
||||
alias: str = "github_onprem",
|
||||
) -> MCPServer:
|
||||
return MCPServer(
|
||||
server_id=server_id,
|
||||
name=alias or server_name,
|
||||
alias=alias,
|
||||
server_name=server_name,
|
||||
transport="http",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _reset_env(monkeypatch):
|
||||
monkeypatch.delenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", raising=False)
|
||||
yield
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Pure helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestShortPrefixHelpers:
|
||||
def test_short_prefix_is_three_base62_chars(self):
|
||||
prefix = compute_short_server_prefix("any-server-id")
|
||||
assert len(prefix) == SHORT_MCP_TOOL_PREFIX_LENGTH
|
||||
assert prefix.isalnum() and prefix.isascii()
|
||||
|
||||
def test_short_prefix_first_char_is_alphabetic(self):
|
||||
"""The first char must be [A-Za-z] so the prefix is a valid identifier
|
||||
on every model API (some providers historically required the first
|
||||
character of a function name to be alphabetic)."""
|
||||
# Sweep many server_ids and rehash attempts to give us coverage of
|
||||
# every position the high-order bits can land on.
|
||||
for i in range(200):
|
||||
for attempt in range(4):
|
||||
prefix = compute_short_server_prefix(f"server-{i}", attempt=attempt)
|
||||
assert prefix[0].isalpha(), (
|
||||
f"prefix {prefix!r} for server-{i} (attempt={attempt}) "
|
||||
f"starts with a non-alphabetic character"
|
||||
)
|
||||
|
||||
def test_short_prefix_is_deterministic(self):
|
||||
assert compute_short_server_prefix("abc") == compute_short_server_prefix("abc")
|
||||
assert compute_short_server_prefix("abc") != compute_short_server_prefix("abd")
|
||||
|
||||
def test_short_prefix_requires_server_id(self):
|
||||
with pytest.raises(ValueError):
|
||||
compute_short_server_prefix("")
|
||||
|
||||
def test_flag_defaults_to_false(self):
|
||||
assert is_short_mcp_tool_prefix_enabled() is False
|
||||
|
||||
@pytest.mark.parametrize("value", ["1", "true", "TRUE", "yes", "On"])
|
||||
def test_flag_truthy_values(self, monkeypatch, value):
|
||||
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", value)
|
||||
assert is_short_mcp_tool_prefix_enabled() is True
|
||||
|
||||
@pytest.mark.parametrize("value", ["0", "false", "no", "off", ""])
|
||||
def test_flag_falsey_values(self, monkeypatch, value):
|
||||
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", value)
|
||||
assert is_short_mcp_tool_prefix_enabled() is False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# get_server_prefix behaviour
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestGetServerPrefix:
|
||||
def test_default_mode_uses_alias(self):
|
||||
server = _make_server(alias="github_onprem", server_name="github_onprem")
|
||||
assert get_server_prefix(server) == "github_onprem"
|
||||
|
||||
def test_short_mode_uses_short_id(self, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
|
||||
server = _make_server(server_id="abcdef-1234")
|
||||
prefix = get_server_prefix(server)
|
||||
assert prefix == compute_short_server_prefix("abcdef-1234")
|
||||
assert len(prefix) == SHORT_MCP_TOOL_PREFIX_LENGTH
|
||||
|
||||
def test_short_mode_falls_back_when_no_server_id(self, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
|
||||
|
||||
class _Bare:
|
||||
alias = "fallback_alias"
|
||||
server_name = None
|
||||
server_id = None
|
||||
|
||||
assert get_server_prefix(_Bare()) == "fallback_alias"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# iter_known_server_prefixes — covers reverse-lookup tolerance
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestIterKnownServerPrefixes:
|
||||
def test_default_mode_includes_short_id_too(self):
|
||||
server = _make_server()
|
||||
prefixes = list(iter_known_server_prefixes(server))
|
||||
# Contains the live prefix and every known form so that mixed-mode
|
||||
# clients can be resolved.
|
||||
assert "github_onprem" in prefixes
|
||||
assert compute_short_server_prefix(server.server_id) in prefixes
|
||||
|
||||
def test_short_mode_still_yields_long_forms(self, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
|
||||
server = _make_server()
|
||||
prefixes = list(iter_known_server_prefixes(server))
|
||||
assert "github_onprem" in prefixes
|
||||
assert compute_short_server_prefix(server.server_id) in prefixes
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Manager-level behaviour: list + reverse-lookup
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _stub_tools() -> List[MCPTool]:
|
||||
return [
|
||||
MCPTool(name="get_repo", description="", inputSchema={"type": "object"}),
|
||||
MCPTool(name="list_issues", description="", inputSchema={"type": "object"}),
|
||||
]
|
||||
|
||||
|
||||
class TestManagerShortPrefix:
|
||||
def test_list_tools_uses_short_prefix_when_flag_on(self, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
|
||||
manager = MCPServerManager()
|
||||
server = _make_server()
|
||||
|
||||
out = manager._create_prefixed_tools(_stub_tools(), server)
|
||||
|
||||
short = compute_short_server_prefix(server.server_id)
|
||||
assert {t.name for t in out} == {f"{short}-get_repo", f"{short}-list_issues"}
|
||||
|
||||
def test_call_tool_lookup_resolves_short_prefix(self, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
|
||||
manager = MCPServerManager()
|
||||
server = _make_server()
|
||||
manager.registry[server.server_id] = server
|
||||
manager._create_prefixed_tools(_stub_tools(), server)
|
||||
|
||||
short = compute_short_server_prefix(server.server_id)
|
||||
resolved = manager._get_mcp_server_from_tool_name(f"{short}-get_repo")
|
||||
assert resolved is server
|
||||
|
||||
def test_call_tool_lookup_resolves_long_prefix_in_short_mode(self, monkeypatch):
|
||||
"""Old clients that cached the long-prefix name must still route."""
|
||||
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
|
||||
manager = MCPServerManager()
|
||||
server = _make_server()
|
||||
manager.registry[server.server_id] = server
|
||||
manager._create_prefixed_tools(_stub_tools(), server)
|
||||
|
||||
resolved = manager._get_mcp_server_from_tool_name("github_onprem-get_repo")
|
||||
assert resolved is server
|
||||
|
||||
def test_default_mode_unchanged(self):
|
||||
manager = MCPServerManager()
|
||||
server = _make_server()
|
||||
|
||||
out = manager._create_prefixed_tools(_stub_tools(), server)
|
||||
|
||||
assert {t.name for t in out} == {
|
||||
"github_onprem-get_repo",
|
||||
"github_onprem-list_issues",
|
||||
}
|
||||
assert (
|
||||
manager._get_mcp_server_from_tool_name("github_onprem-get_repo") is None
|
||||
) # registry empty
|
||||
manager.registry[server.server_id] = server
|
||||
assert (
|
||||
manager._get_mcp_server_from_tool_name("github_onprem-get_repo") is server
|
||||
)
|
||||
|
||||
def test_total_tool_name_length_short_enough(self, monkeypatch):
|
||||
"""The short prefix keeps tool names under the 60-char limit even
|
||||
when the upstream tool name is itself reasonably long."""
|
||||
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
|
||||
long_server_name = "a" * 50
|
||||
server = _make_server(
|
||||
server_id="server-id-1",
|
||||
server_name=long_server_name,
|
||||
alias=long_server_name,
|
||||
)
|
||||
prefix = get_server_prefix(server)
|
||||
full = add_server_prefix_to_name("get_repo", prefix)
|
||||
assert len(full) < 60
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Collision-resolution at registration time
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestShortPrefixCollisionResolution:
|
||||
"""``_assign_unique_short_prefix`` must rehash on collision.
|
||||
|
||||
The dedup path is exercised by forcing two distinct ``server_id``
|
||||
values to both hash to the same natural prefix via a monkeypatched
|
||||
``compute_short_server_prefix``.
|
||||
"""
|
||||
|
||||
def test_no_op_when_flag_off(self):
|
||||
manager = MCPServerManager()
|
||||
server = _make_server(server_id="abc")
|
||||
manager._assign_unique_short_prefix(server)
|
||||
assert server.short_prefix is None
|
||||
|
||||
def test_assigns_natural_hash_when_no_collision(self, monkeypatch):
|
||||
from litellm.proxy._experimental.mcp_server import utils as mcp_utils
|
||||
|
||||
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
|
||||
manager = MCPServerManager()
|
||||
server = _make_server(server_id="abc")
|
||||
manager._assign_unique_short_prefix(server)
|
||||
|
||||
assert server.short_prefix == mcp_utils.compute_short_server_prefix("abc")
|
||||
|
||||
def test_rehashes_when_natural_hash_collides(self, monkeypatch):
|
||||
"""Two server_ids that natural-hash to the same prefix get
|
||||
deterministic, distinct short prefixes."""
|
||||
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
|
||||
|
||||
# Force every attempt=0 hash to "AAA" and attempt=1 to "AAB".
|
||||
# That way the second server registered must rehash to "AAB".
|
||||
from litellm.proxy._experimental.mcp_server import utils as mcp_utils
|
||||
|
||||
def _fake_hash(server_id: str, attempt: int = 0) -> str:
|
||||
return "AAA" if attempt == 0 else f"AA{chr(ord('A') + attempt)}"
|
||||
|
||||
monkeypatch.setattr(mcp_utils, "compute_short_server_prefix", _fake_hash)
|
||||
# Also patch the symbol that the manager imported at module load.
|
||||
from litellm.proxy._experimental.mcp_server import (
|
||||
mcp_server_manager as mgr_module,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(mgr_module, "compute_short_server_prefix", _fake_hash)
|
||||
|
||||
manager = MCPServerManager()
|
||||
first = _make_server(server_id="server-1", alias="srv1")
|
||||
second = _make_server(server_id="server-2", alias="srv2")
|
||||
|
||||
# Pretend both are already in the registry so dedup sees both.
|
||||
manager.registry[first.server_id] = first
|
||||
manager._assign_unique_short_prefix(first)
|
||||
manager.registry[second.server_id] = second
|
||||
manager._assign_unique_short_prefix(second)
|
||||
|
||||
assert first.short_prefix == "AAA"
|
||||
assert second.short_prefix == "AAB"
|
||||
assert first.short_prefix != second.short_prefix
|
||||
|
||||
def test_cached_prefix_is_reused(self, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
|
||||
manager = MCPServerManager()
|
||||
server = _make_server(server_id="abc")
|
||||
server.short_prefix = "ZZZ" # pretend a previous registration set this
|
||||
|
||||
manager._assign_unique_short_prefix(server)
|
||||
|
||||
assert server.short_prefix == "ZZZ"
|
||||
|
||||
def test_get_server_prefix_prefers_cached(self, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
|
||||
server = _make_server(server_id="abc")
|
||||
server.short_prefix = "Q9q"
|
||||
|
||||
assert get_server_prefix(server) == "Q9q"
|
||||
|
|
@ -0,0 +1,255 @@
|
|||
"""
|
||||
Unit tests for `call_with_db_reconnect_retry` — the canonical "try DB read,
|
||||
on transport error reconnect once and retry once" helper.
|
||||
|
||||
Covers the regression in issue #25143 where read paths (e.g.
|
||||
`PrismaClient.get_generic_data`) lost their reconnect-and-retry-once branch in
|
||||
LiteLLM 1.83.x and started emitting `db_exceptions` alerts on transient
|
||||
`httpx.ReadError` flaps that used to self-heal in 1.82.6.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from prisma.errors import UniqueViolationError
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../.."))
|
||||
|
||||
from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry
|
||||
|
||||
|
||||
def _make_client(
|
||||
*,
|
||||
attempt_db_reconnect_return: bool = True,
|
||||
has_attempt_db_reconnect: bool = True,
|
||||
):
|
||||
"""Build a minimal stand-in for PrismaClient that exposes only the surface
|
||||
`call_with_db_reconnect_retry` actually pokes at."""
|
||||
client = MagicMock()
|
||||
if has_attempt_db_reconnect:
|
||||
client.attempt_db_reconnect = AsyncMock(
|
||||
return_value=attempt_db_reconnect_return
|
||||
)
|
||||
else:
|
||||
# `hasattr(client, "attempt_db_reconnect")` must return False — MagicMock
|
||||
# auto-creates attributes, so we wipe it out via `spec`.
|
||||
client = MagicMock(spec=[])
|
||||
client._db_auth_reconnect_timeout_seconds = 2.0
|
||||
client._db_auth_reconnect_lock_timeout_seconds = 0.1
|
||||
return client
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_db_reconnect_retry_returns_value_on_first_success():
|
||||
"""Happy path: factory succeeds first call, no reconnect attempted."""
|
||||
client = _make_client()
|
||||
|
||||
async def _factory():
|
||||
return {"id": 1}
|
||||
|
||||
result = await call_with_db_reconnect_retry(client, _factory, reason="happy_path")
|
||||
|
||||
assert result == {"id": 1}
|
||||
client.attempt_db_reconnect.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_db_reconnect_retry_retries_after_transport_error():
|
||||
"""Transport error on first call → reconnect → second call succeeds."""
|
||||
client = _make_client(attempt_db_reconnect_return=True)
|
||||
|
||||
invocations = []
|
||||
|
||||
async def _factory():
|
||||
invocations.append(None)
|
||||
if len(invocations) == 1:
|
||||
raise httpx.ReadError("transport blip")
|
||||
return {"id": 1}
|
||||
|
||||
result = await call_with_db_reconnect_retry(
|
||||
client, _factory, reason="prisma_get_generic_data_config_lookup_failure"
|
||||
)
|
||||
|
||||
assert result == {"id": 1}
|
||||
assert len(invocations) == 2
|
||||
client.attempt_db_reconnect.assert_awaited_once()
|
||||
call_kwargs = client.attempt_db_reconnect.await_args.kwargs
|
||||
assert call_kwargs["reason"] == "prisma_get_generic_data_config_lookup_failure"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_db_reconnect_retry_does_not_retry_on_data_layer_error():
|
||||
"""Data-layer errors (e.g. UniqueViolationError) are NOT transport errors —
|
||||
propagate immediately, do not reconnect."""
|
||||
client = _make_client()
|
||||
|
||||
async def _factory():
|
||||
raise UniqueViolationError(
|
||||
data={"user_facing_error": {"meta": {}}},
|
||||
message="Unique constraint failed",
|
||||
)
|
||||
|
||||
with pytest.raises(UniqueViolationError):
|
||||
await call_with_db_reconnect_retry(client, _factory, reason="data_layer_test")
|
||||
|
||||
client.attempt_db_reconnect.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_db_reconnect_retry_propagates_when_reconnect_fails():
|
||||
"""Transport error, but reconnect returns False → propagate the original
|
||||
exception. Do not call factory a second time."""
|
||||
client = _make_client(attempt_db_reconnect_return=False)
|
||||
|
||||
invocations = []
|
||||
|
||||
async def _factory():
|
||||
invocations.append(None)
|
||||
raise httpx.ReadError("transport blip")
|
||||
|
||||
with pytest.raises(httpx.ReadError):
|
||||
await call_with_db_reconnect_retry(client, _factory, reason="reconnect_fails")
|
||||
|
||||
assert len(invocations) == 1
|
||||
client.attempt_db_reconnect.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_db_reconnect_retry_propagates_after_second_transport_error():
|
||||
"""Transport error, reconnect succeeds, retry also raises transport error →
|
||||
propagate. At most one retry by construction (no infinite loop)."""
|
||||
client = _make_client(attempt_db_reconnect_return=True)
|
||||
|
||||
invocations = []
|
||||
|
||||
async def _factory():
|
||||
invocations.append(None)
|
||||
raise httpx.ReadError("still failing")
|
||||
|
||||
with pytest.raises(httpx.ReadError):
|
||||
await call_with_db_reconnect_retry(
|
||||
client, _factory, reason="second_transport_error"
|
||||
)
|
||||
|
||||
assert len(invocations) == 2
|
||||
client.attempt_db_reconnect.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_db_reconnect_retry_skips_when_no_attempt_db_reconnect_attr():
|
||||
"""Older PrismaClient stand-ins / partial mocks may not expose
|
||||
`attempt_db_reconnect`. The helper must not crash — just propagate the
|
||||
original exception. Mirrors the `hasattr` guard from
|
||||
`auth_checks._fetch_key_object_from_db_with_reconnect`."""
|
||||
client = _make_client(has_attempt_db_reconnect=False)
|
||||
|
||||
async def _factory():
|
||||
raise httpx.ReadError("transport blip")
|
||||
|
||||
with pytest.raises(httpx.ReadError):
|
||||
await call_with_db_reconnect_retry(client, _factory, reason="no_reconnect_attr")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_db_reconnect_retry_invokes_factory_twice_not_same_coro():
|
||||
"""Guard against the obvious bug of awaiting the same coroutine twice
|
||||
(`RuntimeError: cannot reuse already awaited coroutine`). The helper must
|
||||
call the factory a fresh time on retry, not cache an awaitable."""
|
||||
client = _make_client(attempt_db_reconnect_return=True)
|
||||
|
||||
factory_call_count = 0
|
||||
|
||||
async def _factory():
|
||||
nonlocal factory_call_count
|
||||
factory_call_count += 1
|
||||
if factory_call_count == 1:
|
||||
raise httpx.ReadError("transport blip")
|
||||
return "ok"
|
||||
|
||||
result = await call_with_db_reconnect_retry(
|
||||
client, _factory, reason="fresh_coro_on_retry"
|
||||
)
|
||||
|
||||
assert result == "ok"
|
||||
assert factory_call_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_db_reconnect_retry_passes_explicit_timeouts():
|
||||
"""Explicit timeout_seconds / lock_timeout_seconds override the auth
|
||||
defaults read off the prisma_client object."""
|
||||
client = _make_client(attempt_db_reconnect_return=True)
|
||||
|
||||
async def _factory():
|
||||
if not hasattr(_factory, "_called"):
|
||||
_factory._called = True # type: ignore[attr-defined]
|
||||
raise httpx.ReadError("transport blip")
|
||||
return "ok"
|
||||
|
||||
result = await call_with_db_reconnect_retry(
|
||||
client,
|
||||
_factory,
|
||||
reason="explicit_timeouts",
|
||||
timeout_seconds=5.5,
|
||||
lock_timeout_seconds=0.25,
|
||||
)
|
||||
|
||||
assert result == "ok"
|
||||
call_kwargs = client.attempt_db_reconnect.await_args.kwargs
|
||||
assert call_kwargs["timeout_seconds"] == 5.5
|
||||
assert call_kwargs["lock_timeout_seconds"] == 0.25
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_db_reconnect_retry_uses_auth_defaults_when_unset():
|
||||
"""When timeouts are not provided, helper reads
|
||||
`_db_auth_reconnect_timeout_seconds` / `_db_auth_reconnect_lock_timeout_seconds`
|
||||
off the prisma_client (matching the auth path's existing convention)."""
|
||||
client = _make_client(attempt_db_reconnect_return=True)
|
||||
client._db_auth_reconnect_timeout_seconds = 3.0
|
||||
client._db_auth_reconnect_lock_timeout_seconds = 0.5
|
||||
|
||||
async def _factory():
|
||||
if not hasattr(_factory, "_called"):
|
||||
_factory._called = True # type: ignore[attr-defined]
|
||||
raise httpx.ReadError("transport blip")
|
||||
return "ok"
|
||||
|
||||
await call_with_db_reconnect_retry(client, _factory, reason="defaults")
|
||||
|
||||
call_kwargs = client.attempt_db_reconnect.await_args.kwargs
|
||||
assert call_kwargs["timeout_seconds"] == 3.0
|
||||
assert call_kwargs["lock_timeout_seconds"] == 0.5
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_db_reconnect_retry_preserves_original_error_when_reconnect_raises():
|
||||
"""If `attempt_db_reconnect` itself raises (lock cancellation, timer
|
||||
error, unexpected internal failure), the helper must surface the
|
||||
*original* transport error to telemetry — not the reconnect exception.
|
||||
Otherwise `failure_handler` / `db_exceptions` alerts log the wrong
|
||||
error string and the actual DB transport problem becomes invisible.
|
||||
|
||||
The reconnect error is chained as the `__cause__` for debuggability."""
|
||||
client = MagicMock()
|
||||
reconnect_exc = RuntimeError("simulated reconnect lock cancellation")
|
||||
client.attempt_db_reconnect = AsyncMock(side_effect=reconnect_exc)
|
||||
client._db_auth_reconnect_timeout_seconds = 2.0
|
||||
client._db_auth_reconnect_lock_timeout_seconds = 0.1
|
||||
|
||||
original_exc = httpx.ReadError("transport blip")
|
||||
|
||||
async def _factory():
|
||||
raise original_exc
|
||||
|
||||
with pytest.raises(httpx.ReadError) as exc_info:
|
||||
await call_with_db_reconnect_retry(
|
||||
client, _factory, reason="reconnect_itself_raises"
|
||||
)
|
||||
|
||||
assert exc_info.value is original_exc
|
||||
assert exc_info.value.__cause__ is reconnect_exc
|
||||
client.attempt_db_reconnect.assert_awaited_once()
|
||||
|
|
@ -5,6 +5,7 @@ import sys
|
|||
import time
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
sys.path.insert(
|
||||
|
|
@ -358,3 +359,125 @@ async def test_lightweight_reconnect_skips_kill_on_successful_disconnect(
|
|||
await client._run_reconnect_cycle(timeout_seconds=5.0)
|
||||
|
||||
mock_kill.assert_not_called()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# get_generic_data: transport-reconnect-and-retry coverage (issue #25143)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_generic_data_retries_on_transport_error_for_config_table(
|
||||
mock_proxy_logging,
|
||||
):
|
||||
"""`get_generic_data(table_name="config")` self-heals on a transient
|
||||
`httpx.ReadError`: reconnect once, retry once, return the row.
|
||||
|
||||
Regression for issue #25143 — the 1.83.x line lost the reconnect-and-retry
|
||||
branch that 1.82.6 had on this method. `_update_config_from_db` fans out
|
||||
four concurrent `get_generic_data` calls, so a single transport flap used
|
||||
to surface as four `db_exceptions` alerts and a stale config window.
|
||||
"""
|
||||
client = PrismaClient(
|
||||
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
|
||||
)
|
||||
|
||||
expected_row = {"param_name": "general_settings", "param_value": {"foo": "bar"}}
|
||||
invocations: list[None] = []
|
||||
|
||||
async def _flaky_find_first(**kwargs):
|
||||
invocations.append(None)
|
||||
if len(invocations) == 1:
|
||||
raise httpx.ReadError("simulated transport blip")
|
||||
return expected_row
|
||||
|
||||
client.db.litellm_config.find_first = AsyncMock(side_effect=_flaky_find_first)
|
||||
client.attempt_db_reconnect = AsyncMock(return_value=True)
|
||||
|
||||
result = await client.get_generic_data(
|
||||
key="param_name",
|
||||
value="general_settings",
|
||||
table_name="config",
|
||||
)
|
||||
|
||||
assert result == expected_row
|
||||
assert len(invocations) == 2
|
||||
client.attempt_db_reconnect.assert_awaited_once()
|
||||
reconnect_kwargs = client.attempt_db_reconnect.await_args.kwargs
|
||||
assert reconnect_kwargs["reason"] == "prisma_get_generic_data_config_lookup_failure"
|
||||
|
||||
# The failure_handler telemetry side-effect must NOT fire on the first
|
||||
# transport blip — only if the post-retry call also fails. Drain the
|
||||
# event loop so any spuriously-spawned task would have run by now.
|
||||
await asyncio.sleep(0)
|
||||
mock_proxy_logging.failure_handler.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_generic_data_propagates_when_reconnect_fails(mock_proxy_logging):
|
||||
"""If reconnect itself does not succeed, propagate the original transport
|
||||
error and let the existing failure_handler / db_exceptions telemetry fire."""
|
||||
client = PrismaClient(
|
||||
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
|
||||
)
|
||||
|
||||
client.db.litellm_config.find_first = AsyncMock(
|
||||
side_effect=httpx.ReadError("simulated transport blip")
|
||||
)
|
||||
client.attempt_db_reconnect = AsyncMock(return_value=False)
|
||||
|
||||
with pytest.raises(httpx.ReadError):
|
||||
await client.get_generic_data(
|
||||
key="param_name",
|
||||
value="general_settings",
|
||||
table_name="config",
|
||||
)
|
||||
|
||||
client.attempt_db_reconnect.assert_awaited_once()
|
||||
# Failure telemetry IS expected here — the read genuinely failed.
|
||||
await asyncio.sleep(0)
|
||||
mock_proxy_logging.failure_handler.assert_called_once()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _engine_confirmed_dead flag-reset bug (B2)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_engine_confirmed_dead_persists_across_failed_heavy_reconnect(
|
||||
mock_proxy_logging,
|
||||
):
|
||||
"""Regression test for the flag-reset bug.
|
||||
|
||||
Before the fix, `_run_reconnect_cycle` cleared
|
||||
`self._engine_confirmed_dead = False` *before* awaiting
|
||||
`_do_heavy_reconnect()`. If the heavy reconnect raised (e.g. timeout,
|
||||
missing DATABASE_URL, recreate failure), the flag was left cleared and the
|
||||
next attempt could demote to the lightweight path even though the engine
|
||||
was genuinely dead.
|
||||
|
||||
The fix moves the reset into the success branch — the flag must stay True
|
||||
when heavy reconnect raises.
|
||||
"""
|
||||
client = PrismaClient(
|
||||
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
|
||||
)
|
||||
client._engine_confirmed_dead = True
|
||||
client._engine_pid = 0 # so `_is_engine_alive` is not consulted
|
||||
|
||||
# Make the heavy reconnect path raise.
|
||||
client.db.recreate_prisma_client = AsyncMock(
|
||||
side_effect=RuntimeError("simulated heavy reconnect failure")
|
||||
)
|
||||
client._start_engine_watcher = AsyncMock()
|
||||
client._cleanup_engine_watcher = MagicMock()
|
||||
client._reap_all_zombies = MagicMock()
|
||||
|
||||
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
|
||||
with pytest.raises(Exception):
|
||||
await client._run_reconnect_cycle(timeout_seconds=5.0)
|
||||
|
||||
# The flag must STILL be True so the next attempt re-enters the heavy
|
||||
# branch instead of silently demoting to the lightweight path.
|
||||
assert client._engine_confirmed_dead is True
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue