scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Merge pull request #26733 from BerriAI/litellm_mcp-short-prefix-id-0e42

Browse source 
feat(mcp): opt-in short-ID tool prefix to keep MCP tool names under the 60-char limit
This commit is contained in:
Mateo Wang 2026-04-29 13:48:01 -07:00 • committed by GitHub
commit 97a3bd5ff4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 574 additions and 41 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

148
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

@ -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)

21
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -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:

138
litellm/proxy/_experimental/mcp_server/utils.py
View file

@ -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:

6
litellm/types/mcp_server/mcp_server_manager.py
View file

@ -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

302
tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py Normal file
View file

@ -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"