From fc49c181bca73d63d3b861d7fc2a46a834a3ba8e Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 29 Apr 2026 01:41:24 +0000 Subject: [PATCH 1/7] feat(mcp): opt-in short-ID tool prefix to stay under 60-char tool name limit MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds LITELLM_USE_SHORT_MCP_TOOL_PREFIX. When enabled, tool / prompt / resource / resource-template names emitted from MCP servers are prefixed with a deterministic three-character base62 ID derived from the server's server_id (SHA-256 → base62) instead of the (potentially long) alias / server_name. This keeps namespaced tool names well under the 60-character upper bound enforced by some model APIs while still letting us distinguish MCP-routed tools from local tools. Behavioural notes: - Default off — when the env var is unset, the long-prefix behaviour is unchanged. The plan is to flip the default in a future release and remove the gate after a deprecation window. - Prefix derivation is deterministic, so it is stable across processes, workers and restarts without any persistence layer. - Reverse-lookup is tolerant: _create_prefixed_tools registers every known prefix form (alias / server_name / server_id / short ID) in the routing map and _get_mcp_server_from_tool_name resolves any of them. Old clients holding cached long-prefixed names continue to route correctly even after the flag is enabled. - _get_allowed_mcp_servers_from_mcp_server_names accepts the short prefix in /mcp/{server_name}-style URLs. - The OpenAPI tool-listing path now filters by the active server prefix instead of server.name so spec-backed servers benefit too. Co-authored-by: Mateo Wang --- .../mcp_server/mcp_server_manager.py | 65 +++--- .../proxy/_experimental/mcp_server/server.py | 26 ++- .../proxy/_experimental/mcp_server/utils.py | 104 ++++++++- .../mcp_server/test_short_mcp_tool_prefix.py | 207 ++++++++++++++++++ 4 files changed, 362 insertions(+), 40 deletions(-) create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 251f271903b..a2f0517f81a 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -52,6 +52,7 @@ from litellm.proxy._experimental.mcp_server.utils import ( add_server_prefix_to_name, get_server_prefix, is_tool_name_prefixed, + iter_known_server_prefixes, merge_mcp_headers, normalize_server_name, split_server_prefix_from_name, @@ -1236,7 +1237,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 ) @@ -1838,9 +1843,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 +2610,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) + server = prefix_to_server.get(normalised_prefix) + if 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 server return None diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index c2e998f01e5..3924687a0b1 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -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,14 +712,13 @@ 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 ] + # Always accept server_id even if it isn't part of the + # current prefix form (iter_known_server_prefixes only + # yields it when no other identifier exists). + if server.server_id: + match_list.append(server.server_id.lower()) if server_or_group.lower() in match_list: filtered_server[server.server_id] = server @@ -2031,11 +2031,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: diff --git a/litellm/proxy/_experimental/mcp_server/utils.py b/litellm/proxy/_experimental/mcp_server/utils.py index 146ba10bb76..42910cc790d 100644 --- a/litellm/proxy/_experimental/mcp_server/utils.py +++ b/litellm/proxy/_experimental/mcp_server/utils.py @@ -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,63 @@ 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 base62 ID derived from the server's ``server_id``. +# +# Why three characters and base62 ([0-9A-Za-z])? +# * 62**3 = 238_328 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`` → first three +# base62 chars), which means 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, but for the reverse-lookup path we register every known +# form of the prefix anyway, so a collision only affects the cosmetic +# emitted name, not routing correctness. +# +# 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" + + +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) -> str: + """Derive the deterministic three-character base62 prefix for a server. + + Uses SHA-256 of the server_id and folds the first eight bytes into a + base62 string. 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") + + digest = hashlib.sha256(server_id.encode("utf-8")).digest() + value = int.from_bytes(digest[:8], "big") + chars = [] + for _ in range(SHORT_MCP_TOOL_PREFIX_LENGTH): + value, idx = divmod(value, len(_BASE62_ALPHABET)) + chars.append(_BASE62_ALPHABET[idx]) + return "".join(reversed(chars)) + def is_mcp_available() -> bool: """ @@ -82,7 +140,18 @@ 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 deterministic three-character base62 ID derived from ``server_id`` is + returned. Otherwise we fall back to the historical behaviour: alias if + present, else server_name, else server_id. + """ + if is_short_mcp_tool_prefix_enabled(): + 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 +161,35 @@ 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)) + + 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: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py new file mode 100644 index 00000000000..aef6546c81a --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py @@ -0,0 +1,207 @@ +""" +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_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 From 4e827446d263a228eda7ce8365196574b162322a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 28 Apr 2026 19:58:56 -0700 Subject: [PATCH 2/7] fix: type error --- .../proxy/_experimental/mcp_server/mcp_server_manager.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index a2f0517f81a..8e473c1cb21 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2641,12 +2641,12 @@ class MCPServerManager: server_name_from_prefix, ) = split_server_prefix_from_name(tool_name) normalised_prefix = normalize_server_name(server_name_from_prefix) - server = prefix_to_server.get(normalised_prefix) - if server is not None and ( + 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 server + return matched_server return None From df3dbd18d6fc022267416fbf93993e912029fe9c Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 29 Apr 2026 03:43:13 +0000 Subject: [PATCH 3/7] feat(mcp): rehash short tool prefix on collision and cache per server MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two MCP servers can natural-hash to the same three-character base62 prefix. With 62**3 = 238_328 slots the birthday bound is ~488 servers for 50% collision probability, so a single proxy hosting more than ~100 MCP servers has a non-trivial chance of seeing a collision in practice — and a collision means tool names from two different servers share a routing key, causing silent mis-routing. Mitigation: - compute_short_server_prefix(server_id, attempt=N) folds an attempt counter into the SHA-256 seed, so rehashes are deterministic and produce a fresh three-char prefix space per attempt. - New MCPServer.short_prefix field caches the resolved (post-dedup) prefix on the model so it stays stable across the process lifetime. - MCPServerManager._assign_unique_short_prefix walks attempts 0..N until it finds a prefix not already used by another server in the combined registry. Logs an INFO line when a rehash happens so operators have a breadcrumb if it ever does. - Wired into every registration path: load_servers_from_config, add_server, update_server, reload_servers_from_database. The database reload path also carries the previously-resolved prefix forward so reloads don't churn it. - get_server_prefix prefers the cached short_prefix when set, so the resolved value (not the raw natural hash) is used everywhere. - iter_known_server_prefixes yields the cached short_prefix too, so reverse-lookup tolerance covers the rehashed form. No-op when LITELLM_USE_SHORT_MCP_TOOL_PREFIX is disabled — the field stays None and behaviour is unchanged. Co-authored-by: Mateo Wang --- .../mcp_server/mcp_server_manager.py | 78 ++++++++++++++++++ .../proxy/_experimental/mcp_server/utils.py | 29 +++++-- .../types/mcp_server/mcp_server_manager.py | 6 ++ .../mcp_server/test_short_mcp_tool_prefix.py | 81 +++++++++++++++++++ 4 files changed, 186 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 8e473c1cb21..853208a5382 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -50,7 +50,9 @@ 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, @@ -365,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 @@ -727,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}") @@ -739,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}") @@ -1815,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), + "", + ), + 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]: @@ -2681,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) @@ -2704,10 +2774,18 @@ 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) verbose_logger.debug( "MCP registry refreshed (%s servers in registry)", len(new_registry) diff --git a/litellm/proxy/_experimental/mcp_server/utils.py b/litellm/proxy/_experimental/mcp_server/utils.py index 42910cc790d..a0c278b906a 100644 --- a/litellm/proxy/_experimental/mcp_server/utils.py +++ b/litellm/proxy/_experimental/mcp_server/utils.py @@ -54,17 +54,22 @@ def is_short_mcp_tool_prefix_enabled() -> bool: return raw.strip().lower() in ("1", "true", "yes", "on") -def compute_short_server_prefix(server_id: str) -> str: +def compute_short_server_prefix(server_id: str, attempt: int = 0) -> str: """Derive the deterministic three-character base62 prefix for a server. - Uses SHA-256 of the server_id and folds the first eight bytes into a - base62 string. An empty server_id raises ValueError — short prefixes - require a stable identifier to be deterministic. + Uses SHA-256 of ``f"{server_id}#{attempt}"`` and folds the first eight + bytes into a base62 string. 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") - digest = hashlib.sha256(server_id.encode("utf-8")).digest() + 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") chars = [] for _ in range(SHORT_MCP_TOOL_PREFIX_LENGTH): @@ -143,11 +148,18 @@ def get_server_prefix(server: Any) -> str: """Return the prefix for a server. When the short-prefix mode is enabled (``LITELLM_USE_SHORT_MCP_TOOL_PREFIX``) - a deterministic three-character base62 ID derived from ``server_id`` is - returned. Otherwise we fall back to the historical behaviour: alias if - present, else server_name, else server_id. + 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) @@ -177,6 +189,7 @@ def iter_known_server_prefixes(server: Any) -> Iterator[str]: 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: diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index ace8c8a4188..8f8673b0a7d 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -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 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py index aef6546c81a..cdaecfb2271 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py @@ -205,3 +205,84 @@ class TestManagerShortPrefix: 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" From 6b3f07ba25a375bd3949b6d0fd321a204de4d116 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 29 Apr 2026 03:53:03 +0000 Subject: [PATCH 4/7] fix(mcp): register OpenAPI tools after short prefix collision resolution in reload --- litellm/proxy/_experimental/mcp_server/mcp_server_manager.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 853208a5382..8bb7a5d50d5 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2779,13 +2779,16 @@ class MCPServerManager: 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) From 3fb50563057f88a5291d15e37ada705d942d5320 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 29 Apr 2026 03:59:35 +0000 Subject: [PATCH 5/7] fix(mcp): address greptile review on short tool prefix - server.py: drop the redundant server_id append in _get_filtered_mcp_servers_from_mcp_server_names. iter_known_server_prefixes already yields server_id unconditionally, so the manual append (and its misleading comment) was a no-op duplicate. - utils.py: rewrite the SHORT_MCP_TOOL_PREFIX docstring to accurately describe the collision behaviour. The previous wording said collisions were 'cosmetic only', but a natural-hash collision IS a routing-correctness issue, which is precisely why we already added _assign_unique_short_prefix to rehash deterministically. The new comment cross-references that path. - utils.py: restrict the first character of the short prefix to [A-Za-z] via a 52-char alphabet for position 0 only. The remaining two positions still use the full base62 alphabet. This keeps prefixes valid identifiers on every backend and gives 52*62*62 = 199_888 distinct prefixes (still comfortably more than any realistic deployment). - tests: add coverage proving the first character of the prefix is always alphabetic across many server_ids and rehash attempts. Co-authored-by: Mateo Wang --- .../proxy/_experimental/mcp_server/server.py | 5 -- .../proxy/_experimental/mcp_server/utils.py | 71 ++++++++++++------- .../mcp_server/test_short_mcp_tool_prefix.py | 14 ++++ 3 files changed, 60 insertions(+), 30 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 3924687a0b1..ae6055217b8 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -714,11 +714,6 @@ if MCP_AVAILABLE: match_list = [ s.lower() for s in iter_known_server_prefixes(server) if s ] - # Always accept server_id even if it isn't part of the - # current prefix form (iter_known_server_prefixes only - # yields it when no other identifier exists). - if server.server_id: - match_list.append(server.server_id.lower()) if server_or_group.lower() in match_list: filtered_server[server.server_id] = server diff --git a/litellm/proxy/_experimental/mcp_server/utils.py b/litellm/proxy/_experimental/mcp_server/utils.py index a0c278b906a..df5705c3425 100644 --- a/litellm/proxy/_experimental/mcp_server/utils.py +++ b/litellm/proxy/_experimental/mcp_server/utils.py @@ -21,27 +21,39 @@ MCP_TOOL_PREFIX_FORMAT = "{server_name}{separator}{tool_name}" # 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 base62 ID derived from the server's ``server_id``. +# character ID derived from the server's ``server_id``. # -# Why three characters and base62 ([0-9A-Za-z])? -# * 62**3 = 238_328 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`` → first three -# base62 chars), which means 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, but for the reverse-lookup path we register every known -# form of the prefix anyway, so a collision only affects the cosmetic -# emitted name, not routing correctness. +# 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: @@ -55,15 +67,17 @@ def is_short_mcp_tool_prefix_enabled() -> bool: def compute_short_server_prefix(server_id: str, attempt: int = 0) -> str: - """Derive the deterministic three-character base62 prefix for a server. + """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 base62 string. 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. + 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") @@ -71,10 +85,17 @@ def compute_short_server_prefix(server_id: str, attempt: int = 0) -> str: 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 _ in range(SHORT_MCP_TOOL_PREFIX_LENGTH): - value, idx = divmod(value, len(_BASE62_ALPHABET)) - chars.append(_BASE62_ALPHABET[idx]) + 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)) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py index cdaecfb2271..bf90e9ebef6 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py @@ -57,6 +57,20 @@ class TestShortPrefixHelpers: 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") From 1c9c219a74e060271bb62402c8c05528c3a4055a Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 28 Apr 2026 23:44:34 -0700 Subject: [PATCH 6/7] fix(proxy): self-heal Prisma read paths + harden reconnect state machine MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two related fixes layered on top of the existing reconnect plumbing: 1. Restore reconnect-and-retry on `PrismaClient.get_generic_data` (issue #25143). 1.83.x lost the transport-reconnect-and-retry-once branch that 1.82.6 had on this method, so transient `httpx.ReadError` flaps now surface immediately as `db_exceptions` alerts. `_update_config_from_db` fans out four concurrent `get_generic_data` reads, so a single transport blip used to mark four alerts and a stale config window. Adds `call_with_db_reconnect_retry` to `litellm/proxy/db/exception_handler.py` — a single canonical "try DB read, on transport error reconnect once and retry once" wrapper. Mirrors the inline pattern in `auth_checks._fetch_key_object_from_db_with_reconnect` so we have one implementation rather than three drifting copies, and gives future read paths a clean opt-in. 2. Fix the `_engine_confirmed_dead` flag-reset bug in `_run_reconnect_cycle`. The flag was cleared before `_do_heavy_reconnect()` ran, so any failure inside the heavy reconnect (timeout, missing DATABASE_URL, recreate failure) left the flag False — and the next attempt could silently demote to the lightweight path even though the engine was genuinely dead. Move the reset into the success branch so the flag stays True across heavy-reconnect failures and the next attempt re-enters the heavy branch. Tests: - `tests/test_litellm/proxy/db/test_exception_handler_reconnect_retry.py` (new) — 9 tests covering the helper's contract: happy path, retry on transport error, no retry on data-layer errors, propagation when reconnect fails, propagation after second transport error, `hasattr` guard for partial mocks, fresh-coroutine-per-call invariant, explicit timeout override, default timeouts read off the prisma_client. - `tests/test_litellm/proxy/db/test_prisma_self_heal.py` — adds: - `test_get_generic_data_retries_on_transport_error_for_config_table` - `test_get_generic_data_propagates_when_reconnect_fails` - `test_engine_confirmed_dead_persists_across_failed_heavy_reconnect` (regression test for the flag-reset bug). All 16 self-heal tests + 9 helper tests + 535 auth/exception-handler tests pass locally. --- litellm/proxy/db/exception_handler.py | 123 +++++++++- litellm/proxy/utils.py | 42 +++- .../test_exception_handler_reconnect_retry.py | 225 ++++++++++++++++++ .../proxy/db/test_prisma_self_heal.py | 123 ++++++++++ 4 files changed, 501 insertions(+), 12 deletions(-) create mode 100644 tests/test_litellm/proxy/db/test_exception_handler_reconnect_retry.py diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index cfa90a48eee..f0967732006 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -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,123 @@ 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 `___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, + ) + + did_reconnect = await prisma_client.attempt_db_reconnect( + reason=reason, + timeout_seconds=resolved_timeout, + lock_timeout_seconds=resolved_lock_timeout, + ) + 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() diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 3a1184c434e..93cc8bfd22b 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -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)." diff --git a/tests/test_litellm/proxy/db/test_exception_handler_reconnect_retry.py b/tests/test_litellm/proxy/db/test_exception_handler_reconnect_retry.py new file mode 100644 index 00000000000..22ccb02b7a2 --- /dev/null +++ b/tests/test_litellm/proxy/db/test_exception_handler_reconnect_retry.py @@ -0,0 +1,225 @@ +""" +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 diff --git a/tests/test_litellm/proxy/db/test_prisma_self_heal.py b/tests/test_litellm/proxy/db/test_prisma_self_heal.py index fb215e54777..af65cad09e3 100644 --- a/tests/test_litellm/proxy/db/test_prisma_self_heal.py +++ b/tests/test_litellm/proxy/db/test_prisma_self_heal.py @@ -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 From aa2ef4120098981cb6160d94347743e2e77ffb61 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 28 Apr 2026 23:55:46 -0700 Subject: [PATCH 7/7] fix(proxy): preserve original transport error if reconnect itself raises MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Greptile review on #26756 (P2): if `attempt_db_reconnect` itself raises (e.g. lock cancellation, timer error, unexpected internal failure), the original `httpx.ReadError` / transport error was lost — `failure_handler` and `db_exceptions` alerts then logged the reconnect exception instead of the actual DB transport problem, masking the root cause. Wrap the reconnect call in a try/except. On reconnect failure, re-raise the *original* `first_exc` and chain the reconnect error as `__cause__` so it remains visible for debuggability without becoming the primary exception observers see. Adds `test_call_with_db_reconnect_retry_preserves_original_error_when_reconnect_raises` asserting (a) the propagated exception is the original transport error and (b) the reconnect exception is attached as `__cause__`. --- litellm/proxy/db/exception_handler.py | 25 ++++++++++++---- .../test_exception_handler_reconnect_retry.py | 30 +++++++++++++++++++ 2 files changed, 50 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index f0967732006..ab9d341aa51 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -232,11 +232,26 @@ async def call_with_db_reconnect_retry( first_exc, ) - did_reconnect = await prisma_client.attempt_db_reconnect( - reason=reason, - timeout_seconds=resolved_timeout, - lock_timeout_seconds=resolved_lock_timeout, - ) + # 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 diff --git a/tests/test_litellm/proxy/db/test_exception_handler_reconnect_retry.py b/tests/test_litellm/proxy/db/test_exception_handler_reconnect_retry.py index 22ccb02b7a2..ae0e1f845b0 100644 --- a/tests/test_litellm/proxy/db/test_exception_handler_reconnect_retry.py +++ b/tests/test_litellm/proxy/db/test_exception_handler_reconnect_retry.py @@ -223,3 +223,33 @@ async def test_call_with_db_reconnect_retry_uses_auth_defaults_when_unset(): 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()