scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

Merge f4f172633e into 7f1b0a4633

Browse source
This commit is contained in:
Felix Pan 2026-11-05 23:12:44 +00:00 • committed by GitHub
commit 00f7a75666
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 320 additions and 23 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

160
litellm/proxy/mcp_server/operations.py
View file

@ -48,7 +48,7 @@ from mcp.types import (
ToolsCapability,
)
from mcp.types import Tool as MCPTool
from pydantic import AnyUrl, ConfigDict, Field, JsonValue, SecretStr, TypeAdapter
from pydantic import AnyUrl, ConfigDict, Field, JsonValue, SecretStr, TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict, assert_never
from litellm._logging import verbose_logger
@ -167,6 +167,7 @@ from litellm.types.mcp import (
DEFAULT_CREDENTIAL_HEADER,
MCP_LEGACY_VERSIONS,
MCPAuth,
MCPToolSearchSettings,
MCPTransport,
without_header,
)
@ -185,6 +186,7 @@ __all__ = (
"_check_byok_credential",
"_client_has_passthrough_authorization",
"_client_has_per_server_auth_header",
"_dispatch_mcp_tool_search",
"_dispatch_virtual_mcp_tool",
"_fire_mcp_tool_call_logging",
"_get_allowed_mcp_servers",
@ -315,6 +317,48 @@ async def _mcp_server_rate_limit_rejection(
return admission.rejected_servers[0][1] if admission.rejected_servers else None
async def _build_admitted_mcp_logging_obj(
name: str,
arguments: dict[str, object],
user_api_key_auth: UserAPIKeyAuth,
logging_auth: UserAPIKeyAuth | None = None,
raw_headers: Mapping[str, str] | None = None,
client_ip: str | None = None,
) -> LiteLLMLoggingObj | None:
"""Run shared key admission and logging setup for virtual MCP operations."""
from litellm.proxy.common_request_processing import (
ProxyBaseLLMRequestProcessing,
)
from litellm.proxy.proxy_server import (
general_settings,
proxy_config,
proxy_logging_obj,
)
request: Final = build_synthetic_mcp_request(
path="/mcp/tools/call",
raw_headers=raw_headers,
client_ip=client_ip,
)
pre_call_result: Final = await ProxyBaseLLMRequestProcessing(
data={"name": name, "arguments": arguments}
).common_processing_pre_call_logic(
request=request,
user_api_key_dict=user_api_key_auth,
proxy_config=proxy_config,
route_type=CallTypes.call_mcp_tool.value,
proxy_logging_obj=proxy_logging_obj,
general_settings=general_settings,
logging_auth=logging_auth,
)
_, logging_candidate = TypeAdapter(tuple[dict[str, object], object]).validate_python(pre_call_result)
if not isinstance(logging_candidate, LiteLLMLoggingObj):
verbose_logger.warning("MCP request processing did not return a logging object")
return None
virtual_logging_obj: Final = logging_candidate
return virtual_logging_obj
async def _build_virtual_call_logging_obj(
name: str,
arguments: dict[str, object],
@ -358,6 +402,27 @@ async def _build_virtual_call_logging_obj(
return virtual_logging_obj
async def _build_keyword_search_logging_obj(
name: str,
arguments: Mapping[str, object],
user_api_key_auth: UserAPIKeyAuth,
raw_headers: Mapping[str, str] | None = None,
client_ip: str | None = None,
) -> LiteLLMLoggingObj | None:
"""Admit and build the logger for a keyword tool search.
Keyword search has no downstream model or real-tool call to perform key
admission, so it owns this step independently from ``mcp_tool_call``.
"""
return await _build_admitted_mcp_logging_obj(
name=name,
arguments=dict(arguments),
user_api_key_auth=user_api_key_auth,
raw_headers=raw_headers,
client_ip=client_ip,
)
async def _log_virtual_mcp_tool_failure(
*,
logging_obj: LiteLLMLoggingObj | None,
@ -409,7 +474,93 @@ async def _log_virtual_mcp_tool_failure(
verbose_logger.exception("Error logging failed MCP proxy tool call")
async def _dispatch_virtual_mcp_tool(
async def _dispatch_mcp_tool_search( # pyright: ignore[reportUnknownParameterType] # MCP SDK result type is unresolved
query: str,
top_k: int,
user_api_key_auth: UserAPIKeyAuth,
client_ip: str | None = None,
mcp_servers: list[str] | None = None, # mutable-ok: mirrors the MCP call request container
mcp_auth_header: str | None = None,
mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, # mutable-ok: forwarded request headers
oauth2_headers: dict[str, str] | None = None, # mutable-ok: forwarded request headers
raw_headers: dict[str, str] | None = None, # mutable-ok: forwarded request headers
) -> CallToolResult:
from litellm.proxy.mcp_server.tool_search import (
MCP_TOOL_SEARCH_TOOL_NAME,
handle_mcp_tool_search, # pyright: ignore[reportUnknownVariableType] # MCP SDK result type is unresolved
mcp_tool_search_settings,
)
settings: Final[MCPToolSearchSettings | ValidationError] = mcp_tool_search_settings()
if not isinstance(settings, ValidationError) and settings.embedding_model is not None:
return await handle_mcp_tool_search( # pyright: ignore[reportUnknownVariableType] # MCP SDK result type is unresolved
query=query,
top_k=top_k,
user_api_key_dict=user_api_key_auth,
client_ip=client_ip,
mcp_servers=mcp_servers,
mcp_auth_header=mcp_auth_header,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
)
start_time: Final = datetime.now() # noqa: DTZ005 # logging timestamps are naive
request_data: Final[dict[str, object]] = { # mutable-ok: matches the virtual tool request contract
"query": query,
"top_k": top_k,
}
try:
logging_obj: Final = await _build_keyword_search_logging_obj(
name=MCP_TOOL_SEARCH_TOOL_NAME,
arguments=request_data,
user_api_key_auth=user_api_key_auth,
raw_headers=raw_headers,
client_ip=client_ip,
)
except Exception as exc:
await fire_mcp_tool_call_failure_logging(
None,
exc,
start_time,
user_api_key_auth,
request_data,
)
raise
try:
result: Final[CallToolResult] = await handle_mcp_tool_search( # pyright: ignore[reportUnknownVariableType] # MCP SDK result type is unresolved
query=query,
top_k=top_k,
user_api_key_dict=user_api_key_auth,
client_ip=client_ip,
mcp_servers=mcp_servers,
mcp_auth_header=mcp_auth_header,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
)
except Exception as exc:
await fire_mcp_tool_call_failure_logging(
logging_obj,
exc,
start_time,
user_api_key_auth,
request_data,
)
raise
if logging_obj is None:
return result # pyright: ignore[reportUnknownVariableType] # narrowed typed MCP result
return await _fire_mcp_tool_call_logging( # pyright: ignore[reportUnknownVariableType] # MCP SDK result type is unresolved
logging_obj=logging_obj,
result=result, # pyright: ignore[reportUnknownArgumentType] # narrowed typed MCP result
start_time=start_time,
end_time=datetime.now(), # noqa: DTZ005 # logging timestamps are naive
user_api_key_auth=user_api_key_auth,
request_data=request_data,
)
async def _dispatch_virtual_mcp_tool( # pyright: ignore[reportUnknownParameterType] # MCP SDK result type is unresolved
name: str,
arguments: dict[str, object] | None,
user_api_key_auth: UserAPIKeyAuth | None,
@ -440,7 +591,6 @@ async def _dispatch_virtual_mcp_tool(
handle_agent_search,
handle_mcp_proxy_tool,
handle_mcp_tool_call,
handle_mcp_tool_search,
handle_skill_search,
)
@ -564,10 +714,10 @@ async def _dispatch_virtual_mcp_tool(
args: Final = arguments or {}
if name == MCP_TOOL_SEARCH_TOOL_NAME:
return await handle_mcp_tool_search(
return await _dispatch_mcp_tool_search(
query=TypeAdapter(str).validate_python(args.get("query", "")),
top_k=coerce_top_k(args.get("top_k", 5)),
user_api_key_dict=user_api_key_auth,
user_api_key_auth=user_api_key_auth,
client_ip=client_ip,
mcp_servers=mcp_servers,
mcp_auth_header=mcp_auth_header,

183
tests/unit/proxy/mcp_server/test_mcp_tool_search.py
View file

@ -179,13 +179,12 @@ FAKE_VECTORS: dict[str, Vector] = {
}
def _paged_params():
from mcp.types import PaginatedRequestParams
return PaginatedRequestParams()
class RecordingEmbedder:
def __init__(self) -> None:
self.calls: list[tuple[str, ...]] = []
@ -1058,11 +1057,9 @@ class TestDispatchVirtualMcpTool:
@pytest.mark.asyncio
async def test_routes_search_with_client_ip(self) -> None:
from litellm.proxy.mcp_server import server as srv
uak = UserAPIKeyAuth(api_key="k", object_permission=_make_perm(mcp_tool_search_enabled=True))
with patch(
"litellm.proxy.mcp_server.tool_search.handle_mcp_tool_search",
"litellm.proxy.mcp_server.operations._dispatch_mcp_tool_search",
new_callable=AsyncMock,
return_value="SEARCH_RESULT",
) as mock_search:
@ -1080,8 +1077,6 @@ class TestDispatchVirtualMcpTool:
@pytest.mark.asyncio
async def test_routes_agent_search_to_its_handler(self) -> None:
from litellm.proxy.mcp_server import server as srv
uak = UserAPIKeyAuth(api_key="k", object_permission=_make_perm(mcp_tool_search_enabled=True))
with patch( # test-quality-ok: dispatch routing is the subject; the handler is faked like its siblings here
"litellm.proxy.mcp_server.tool_search.handle_agent_search",
@ -1115,8 +1110,6 @@ class TestDispatchVirtualMcpTool:
@pytest.mark.asyncio
async def test_routes_call_with_client_ip(self) -> None:
from litellm.proxy.mcp_server import server as srv
uak = UserAPIKeyAuth(api_key="k", object_permission=_make_perm(mcp_tool_search_enabled=True))
with patch(
"litellm.proxy.mcp_server.tool_search.handle_mcp_tool_call",
@ -1148,13 +1141,12 @@ class TestDispatchVirtualMcpTool:
"""Regression: the SSE dispatch must run the pre-call pipeline and forward
the resulting logging object to handle_mcp_tool_call, otherwise mcp_tool_call
over /mcp/ skips spend logging and guardrails (unlike the REST path)."""
from litellm.proxy.mcp_server import server as srv
uak = UserAPIKeyAuth(api_key="k", object_permission=_make_perm(mcp_tool_search_enabled=True))
sentinel_logging_obj = object()
with (
patch.object(
mcp_operations, "_build_virtual_call_logging_obj",
mcp_operations,
"_build_virtual_call_logging_obj",
new_callable=AsyncMock,
return_value=sentinel_logging_obj,
) as mock_build,
@ -1178,11 +1170,9 @@ class TestDispatchVirtualMcpTool:
async def test_search_coerces_non_int_top_k(self) -> None:
"""Regression: a non-integer top_k from an MCP client must not raise; it
falls back to the default instead of ValueError propagating out."""
from litellm.proxy.mcp_server import server as srv
uak = UserAPIKeyAuth(api_key="k", object_permission=_make_perm(mcp_tool_search_enabled=True))
with patch(
"litellm.proxy.mcp_server.tool_search.handle_mcp_tool_search",
"litellm.proxy.mcp_server.operations._dispatch_mcp_tool_search",
new_callable=AsyncMock,
return_value="SEARCH_RESULT",
) as mock_search:
@ -1195,6 +1185,163 @@ class TestDispatchVirtualMcpTool:
assert mock_search.await_args.kwargs["top_k"] == 5
@pytest.mark.asyncio
async def test_keyword_search_runs_admission_and_completed_logging(self, monkeypatch) -> None:
import litellm
from litellm.litellm_core_utils.litellm_logging import (
Logging as LiteLLMLoggingObj,
)
from litellm.proxy.mcp_server.tool_search import (
MCP_TOOL_SEARCH_TOOL_NAME,
)
monkeypatch.setattr(litellm, "mcp_tool_search", {})
uak = UserAPIKeyAuth(api_key="k")
sentinel_logging_obj = MagicMock(spec=LiteLLMLoggingObj)
processors = []
class FakeProcessor:
def __init__(self, data):
self.data = data
processors.append(self)
async def common_processing_pre_call_logic(self, **kwargs):
return self.data, sentinel_logging_obj
with (
patch(
"litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing",
FakeProcessor,
),
patch(
"litellm.proxy.mcp_server.tool_search.handle_mcp_tool_search",
new_callable=AsyncMock,
return_value="SEARCH_RESULT",
),
patch.object(
mcp_operations,
"_fire_mcp_tool_call_logging",
new_callable=AsyncMock,
return_value="LOGGED_SEARCH_RESULT",
) as mock_fire,
):
result = await mcp_operations._dispatch_mcp_tool_search(
query="ping",
top_k=3,
user_api_key_auth=uak,
)
assert processors[0].data["name"] == MCP_TOOL_SEARCH_TOOL_NAME
assert processors[0].data["arguments"] == {
"query": "ping",
"top_k": 3,
}
assert mock_fire.await_count == 1
assert result == "LOGGED_SEARCH_RESULT"
@pytest.mark.asyncio
async def test_keyword_search_logs_admission_failure(self, monkeypatch) -> None:
import litellm
monkeypatch.setattr(litellm, "mcp_tool_search", {})
uak = UserAPIKeyAuth(api_key="k")
failure = RuntimeError("rate limited")
with (
patch.object(
mcp_operations,
"_build_keyword_search_logging_obj",
new_callable=AsyncMock,
side_effect=failure,
),
patch.object(
mcp_operations,
"fire_mcp_tool_call_failure_logging",
new_callable=AsyncMock,
) as mock_failure_log,
):
with pytest.raises(RuntimeError, match="rate limited"):
await mcp_operations._dispatch_mcp_tool_search(
query="ping",
top_k=3,
user_api_key_auth=uak,
)
mock_failure_log.assert_awaited_once()
@pytest.mark.asyncio
async def test_keyword_search_admission_is_independent_of_mcp_tool_call_logger(self, monkeypatch) -> None:
import litellm
from litellm.litellm_core_utils.litellm_logging import (
Logging as LiteLLMLoggingObj,
)
monkeypatch.setattr(litellm, "mcp_tool_search", {})
sentinel_logging_obj = MagicMock(spec=LiteLLMLoggingObj)
class FakeProcessor:
def __init__(self, data):
self.data = data
async def common_processing_pre_call_logic(self, **kwargs):
return kwargs, sentinel_logging_obj
with (
patch(
"litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing",
FakeProcessor,
),
patch.object(
mcp_operations,
"_build_virtual_call_logging_obj",
side_effect=AssertionError("mcp_tool_call logger must not own keyword search admission"),
),
patch(
"litellm.proxy.mcp_server.tool_search.handle_mcp_tool_search",
new_callable=AsyncMock,
return_value="SEARCH_RESULT",
),
patch.object(
mcp_operations,
"_fire_mcp_tool_call_logging",
new_callable=AsyncMock,
return_value="LOGGED_SEARCH_RESULT",
),
):
result = await mcp_operations._dispatch_mcp_tool_search(
query="ping",
top_k=3,
user_api_key_auth=UserAPIKeyAuth(api_key="k"),
)
assert result == "LOGGED_SEARCH_RESULT"
@pytest.mark.asyncio
async def test_semantic_search_keeps_embedding_request_accounting(self, monkeypatch) -> None:
import litellm
monkeypatch.setattr(litellm, "mcp_tool_search", {"embedding_model": "embedding-model"})
uak = UserAPIKeyAuth(api_key="k")
with (
patch(
"litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing",
side_effect=AssertionError("semantic path must not add another request"),
),
patch(
"litellm.proxy.mcp_server.tool_search.handle_mcp_tool_search",
new_callable=AsyncMock,
return_value="SEARCH_RESULT",
) as mock_search,
):
result = await mcp_operations._dispatch_mcp_tool_search(
query="ping",
top_k=3,
user_api_key_auth=uak,
)
assert mock_search.await_count == 1
assert result == "SEARCH_RESULT"
@pytest.mark.asyncio
async def test_call_handler_forwards_auth_headers_to_execute(self) -> None:
"""Regression: per-request auth headers must reach execute_mcp_tool so
@ -1287,7 +1434,9 @@ class TestCaptureHostProgressCallback:
@pytest.mark.parametrize("token", ["tok12345", 12345, 0])
@pytest.mark.parametrize("message", [None, "halfway"])
@pytest.mark.parametrize("delivery_error", [None, RuntimeError("request stream closed")])
async def test_forwards_progress_on_originating_request(self, _mcp_request_ctx, token, message, delivery_error) -> None:
async def test_forwards_progress_on_originating_request(
self, _mcp_request_ctx, token, message, delivery_error
) -> None:
from mcp.server.connection import Connection
from mcp.server.session import ServerSession
from mcp.shared.dispatcher import DispatchContext
@ -1467,5 +1616,3 @@ async def test_cold_concurrent_schema_validation_accepts_valid_arguments() -> No
finally:
if original_affinity is not None:
os.sched_setaffinity(0, original_affinity)