diff --git a/litellm/proxy/mcp_server/operations.py b/litellm/proxy/mcp_server/operations.py index 66e46c73f98..4a57c0e4947 100644 --- a/litellm/proxy/mcp_server/operations.py +++ b/litellm/proxy/mcp_server/operations.py @@ -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, diff --git a/tests/unit/proxy/mcp_server/test_mcp_tool_search.py b/tests/unit/proxy/mcp_server/test_mcp_tool_search.py index b21202ba227..4b8d2da0113 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_tool_search.py +++ b/tests/unit/proxy/mcp_server/test_mcp_tool_search.py @@ -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) - -