diff --git a/litellm/llms/anthropic/pass_through/messages/handler.py b/litellm/llms/anthropic/pass_through/messages/handler.py index d0f0fa4d901..e0d10cd37f4 100644 --- a/litellm/llms/anthropic/pass_through/messages/handler.py +++ b/litellm/llms/anthropic/pass_through/messages/handler.py @@ -28,6 +28,7 @@ from litellm.llms.base_llm.anthropic_messages.transformation import ( from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.responses.mcp.tool_precheck import may_name_litellm_mcp_gateway from litellm.types.llms.anthropic import AllAnthropicToolsValues from litellm.types.llms.anthropic_messages.anthropic_request import AnthropicMetadata from litellm.types.llms.anthropic_messages.anthropic_response import ( @@ -603,38 +604,34 @@ def anthropic_messages_handler( # downstream path (native passthrough and both bridges) gets real tools rather than a # reference the provider cannot resolve. Popped from kwargs so it never reaches the provider. skip_mcp_handler: Final = kwargs.pop("_skip_mcp_handler", False) - if not skip_mcp_handler and tools: + if not skip_mcp_handler and tools and may_name_litellm_mcp_gateway(tools): # pyright: ignore[reportUnknownArgumentType] # untyped tools list from litellm.llms.anthropic.pass_through.messages.mcp_handler import ( anthropic_messages_with_mcp, ) - from litellm.responses.mcp.litellm_proxy_mcp_handler import ( - LiteLLM_Proxy_MCP_Handler, - ) - if LiteLLM_Proxy_MCP_Handler.should_use_litellm_mcp_gateway(tools=tools): - return anthropic_messages_with_mcp( - max_tokens=max_tokens, - messages=messages, - model=model, - metadata=metadata, - stop_sequences=stop_sequences, - stream=stream, - system=cast( # cast-ok: anthropic accepts list system blocks; handler signature is narrower - "str | None", system - ), - temperature=temperature, - thinking=thinking, - tool_choice=tool_choice, - tools=tools, - top_k=top_k, - top_p=top_p, - container=container, - api_key=api_key, - api_base=api_base, - client=client, - custom_llm_provider=custom_llm_provider, - **kwargs, - ) + return anthropic_messages_with_mcp( + max_tokens=max_tokens, + messages=messages, + model=model, + metadata=metadata, + stop_sequences=stop_sequences, + stream=stream, + system=cast( # cast-ok: anthropic accepts list system blocks; handler signature is narrower + "str | None", system + ), + temperature=temperature, + thinking=thinking, + tool_choice=tool_choice, + tools=tools, + top_k=top_k, + top_p=top_p, + container=container, + api_key=api_key, + api_base=api_base, + client=client, + custom_llm_provider=custom_llm_provider, + **kwargs, + ) anthropic_messages_provider_config: Final = _native_messages_config( model, custom_llm_provider, kwargs.get("model_info"), api_base diff --git a/litellm/main.py b/litellm/main.py index 9430af34401..3ecb325c5e6 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -135,6 +135,7 @@ from litellm.llms.vertex_ai.common_utils import ( get_vertex_ai_model_route, ) from litellm.realtime_api.main import realtime_health_check +from litellm.responses.mcp.tool_precheck import may_name_litellm_mcp_gateway from litellm.secret_managers.main import get_secret_bool, get_secret_str from litellm.types.completion import ( CompletionDispatchContext, @@ -5434,63 +5435,55 @@ def completion( control_options: Final = _resolve_control_options(kwargs, model) skip_mcp_handler: Final = kwargs.pop("_skip_mcp_handler", False) - if not skip_mcp_handler and tools: + if not skip_mcp_handler and tools and may_name_litellm_mcp_gateway(tools): # pyright: ignore[reportUnknownArgumentType] # untyped tools list from litellm.responses.mcp.chat_completions_handler import acompletion_with_mcp - from litellm.responses.mcp.litellm_proxy_mcp_handler import ( - LiteLLM_Proxy_MCP_Handler, - ) - from litellm.types.llms.openai import ToolParam - # Check if MCP tools are present (following responses pattern) - # Cast tools to Optional[Iterable[ToolParam]] for type checking - tools_for_mcp: Final = cast(Iterable[ToolParam] | None, tools) - if LiteLLM_Proxy_MCP_Handler.should_use_litellm_mcp_gateway(tools=tools_for_mcp): - return acompletion_with_mcp( # pyright: ignore[reportReturnType] # MCP path returns a coroutine that acompletion() awaits; completion()'s sync return type omits it - model=model, - messages=messages, - functions=functions, - function_call=function_call, - timeout=timeout, - temperature=temperature, - top_p=top_p, - n=n, - stream=stream, - stream_options=stream_options, - stop=stop, - max_tokens=max_tokens, - max_completion_tokens=max_completion_tokens, - modalities=modalities, - prediction=prediction, - audio=audio, - presence_penalty=presence_penalty, - frequency_penalty=frequency_penalty, - logit_bias=logit_bias, - user=user, - response_format=response_format, - seed=seed, - tools=tools, - tool_choice=tool_choice, - parallel_tool_calls=parallel_tool_calls, - logprobs=logprobs, - top_logprobs=top_logprobs, - deployment_id=deployment_id, - reasoning_effort=reasoning_effort, - verbosity=verbosity, - safety_identifier=safety_identifier, - service_tier=service_tier, - store=store, - prompt_cache_key=prompt_cache_key, - base_url=base_url, - api_version=api_version, - api_key=api_key, - model_list=model_list, - extra_headers=extra_headers, - thinking=thinking, - web_search_options=web_search_options, - shared_session=shared_session, - enable_json_schema_validation=enable_json_schema_validation, - **kwargs, - ) + return acompletion_with_mcp( # pyright: ignore[reportReturnType] # MCP path returns a coroutine that acompletion() awaits; completion()'s sync return type omits it + model=model, + messages=messages, + functions=functions, + function_call=function_call, + timeout=timeout, + temperature=temperature, + top_p=top_p, + n=n, + stream=stream, + stream_options=stream_options, + stop=stop, + max_tokens=max_tokens, + max_completion_tokens=max_completion_tokens, + modalities=modalities, + prediction=prediction, + audio=audio, + presence_penalty=presence_penalty, + frequency_penalty=frequency_penalty, + logit_bias=logit_bias, + user=user, + response_format=response_format, + seed=seed, + tools=tools, + tool_choice=tool_choice, + parallel_tool_calls=parallel_tool_calls, + logprobs=logprobs, + top_logprobs=top_logprobs, + deployment_id=deployment_id, + reasoning_effort=reasoning_effort, + verbosity=verbosity, + safety_identifier=safety_identifier, + service_tier=service_tier, + store=store, + prompt_cache_key=prompt_cache_key, + base_url=base_url, + api_version=api_version, + api_key=api_key, + model_list=model_list, + extra_headers=extra_headers, + thinking=thinking, + web_search_options=web_search_options, + shared_session=shared_session, + enable_json_schema_validation=enable_json_schema_validation, + **kwargs, + ) api_base = kwargs.get("api_base", None) mock_response: Final[MOCK_RESPONSE_TYPE | None] = kwargs.get("mock_response", None) mock_tool_calls: Final = kwargs.get("mock_tool_calls", None) diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 727a5ccfedf..98475e42b99 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -49,6 +49,7 @@ from litellm.responses.mcp.round_billing import ( first_round_billed_on_failure, summed_cost_breakdown, ) +from litellm.responses.mcp.tool_precheck import may_name_litellm_mcp_gateway from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ( PromptObject, @@ -1220,11 +1221,7 @@ def _responses_try_dispatch_mcp_gateway( | None ): """Return a response when MCP gateway handles the call; otherwise None.""" - from litellm.responses.mcp.litellm_proxy_mcp_handler import ( - LiteLLM_Proxy_MCP_Handler, - ) - - if skip_mcp_handler or not LiteLLM_Proxy_MCP_Handler.should_use_litellm_mcp_gateway(tools=tools): + if skip_mcp_handler or not may_name_litellm_mcp_gateway(tools): return None mcp_call_kwargs: Final = { "model": model, diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index b4dbf4c3d4e..e66a7d1102b 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -1,4 +1,3 @@ -import re import traceback from collections.abc import Awaitable, Callable, Collection, Iterable, Mapping, Sequence from datetime import datetime @@ -20,6 +19,12 @@ from litellm.proxy._experimental.mcp_server.utils import ( strip_known_server_prefix, ) from litellm.responses.main import aresponses # noqa: TID251 # inner call must skip the MCP gateway that invoked it +from litellm.responses.mcp.tool_precheck import ( + LITELLM_PROXY_MCP_SERVER_URL_PREFIX, + may_name_litellm_mcp_gateway, + names_gateway_explicitly, + proxy_path_mcp_name, +) from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator from litellm.types.llms.openai import ( ResponseInputParam, @@ -65,29 +70,6 @@ class _AresponsesOptions(ExtraItemsTypedDict, total=False, extra_items=object): pass -LITELLM_PROXY_MCP_SERVER_URL: Final = "litellm_proxy" -LITELLM_PROXY_MCP_SERVER_URL_PREFIX: Final = f"{LITELLM_PROXY_MCP_SERVER_URL}/mcp/" - -_PROXY_MCP_PATH_RE: Final = re.compile(r"^https?://.+/mcp/([^/]+)$") - - -def _mcp_server_url(tool: ToolParam) -> str | None: - if not isinstance(tool, dict) or tool.get("type") != "mcp": - return None - server_url: Final = tool.get("server_url") - return server_url if isinstance(server_url, str) else None - - -def _names_gateway_explicitly(tool: ToolParam) -> bool: - return (_mcp_server_url(tool) or "").startswith(LITELLM_PROXY_MCP_SERVER_URL) - - -def _proxy_path_mcp_name(tool: ToolParam) -> str | None: - server_url: Final = _mcp_server_url(tool) - match: Final = None if server_url is None else _PROXY_MCP_PATH_RE.match(server_url) - return None if match is None else match.group(1) - - def _registered_mcp_servers() -> Collection[MCPServer]: from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, @@ -133,7 +115,7 @@ async def _gateway_served_names( async def _served_mcp_path_names( tools: Collection[ToolParam], served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] ) -> frozenset[str]: - names: Final = frozenset(name for name in map(_proxy_path_mcp_name, tools) if name is not None) + names: Final = frozenset(name for name in map(proxy_path_mcp_name, tools) if name is not None) return await served_names(names) if names else frozenset[str]() @@ -162,17 +144,15 @@ class LiteLLM_Proxy_MCP_Handler: @staticmethod def should_use_litellm_mcp_gateway(tools: Iterable[ToolParam] | None) -> bool: - """True when a tool may name this gateway: server_url "litellm_proxy..." or an http(s) URL ending in - /mcp/. `_split_mcp_tools` then settles which of the latter the gateway actually serves.""" - return any(_names_gateway_explicitly(tool) or _proxy_path_mcp_name(tool) is not None for tool in tools or ()) + return may_name_litellm_mcp_gateway(tools) _should_use_litellm_mcp_gateway = should_use_litellm_mcp_gateway @staticmethod def _parse_mcp_tools(tools: Iterable[Mapping[str, object]] | None) -> SplitTools: items: Final = tuple(tools or ()) - gateway_tools: Final[list[ToolParam]] = [tool for tool in items if _names_gateway_explicitly(tool)] - other_tools: Final[list[ToolParam]] = [tool for tool in items if not _names_gateway_explicitly(tool)] + gateway_tools: Final[list[ToolParam]] = [tool for tool in items if names_gateway_explicitly(tool)] + other_tools: Final[list[ToolParam]] = [tool for tool in items if not names_gateway_explicitly(tool)] return gateway_tools, other_tools @staticmethod @@ -185,7 +165,7 @@ class LiteLLM_Proxy_MCP_Handler: return LiteLLM_Proxy_MCP_Handler._parse_mcp_tools( [ {**tool, "server_url": f"{LITELLM_PROXY_MCP_SERVER_URL_PREFIX}{name}"} - if (name := _proxy_path_mcp_name(tool)) in served + if (name := proxy_path_mcp_name(tool)) in served else tool for tool in items ] @@ -201,7 +181,7 @@ class LiteLLM_Proxy_MCP_Handler: ) -> tuple[bool, ...]: items: Final = tuple(tools or ()) served: Final = await _served_mcp_path_names(items, served_names) - return tuple(_names_gateway_explicitly(tool) or _proxy_path_mcp_name(tool) in served for tool in items) + return tuple(names_gateway_explicitly(tool) or proxy_path_mcp_name(tool) in served for tool in items) @staticmethod async def _apply_toolset_permissions( diff --git a/litellm/responses/mcp/tool_precheck.py b/litellm/responses/mcp/tool_precheck.py new file mode 100644 index 00000000000..b76f4066e6d --- /dev/null +++ b/litellm/responses/mcp/tool_precheck.py @@ -0,0 +1,43 @@ +import importlib +import re +from collections.abc import Iterable +from typing import Final + +LITELLM_PROXY_MCP_SERVER_URL: Final = "litellm_proxy" +LITELLM_PROXY_MCP_SERVER_URL_PREFIX: Final = f"{LITELLM_PROXY_MCP_SERVER_URL}/mcp/" + +_PROXY_MCP_PATH_RE: Final = re.compile(r"^https?://.+/mcp/([^/]+)$") +_GATEWAY_REGISTRY_MODULE: Final = "litellm.proxy._experimental.mcp_server.mcp_server_manager" + + +def _mcp_server_url(tool: object) -> str | None: + match tool: + case {"type": "mcp", "server_url": str() as server_url} if isinstance(tool, dict): + return server_url + case _: + return None + + +def names_gateway_explicitly(tool: object) -> bool: + return (_mcp_server_url(tool) or "").startswith(LITELLM_PROXY_MCP_SERVER_URL) + + +def proxy_path_mcp_name(tool: object) -> str | None: + server_url: Final = _mcp_server_url(tool) + match: Final = None if server_url is None else _PROXY_MCP_PATH_RE.match(server_url) + return None if match is None else match.group(1) + + +def _gateway_registry_loads() -> bool: + try: + importlib.import_module(_GATEWAY_REGISTRY_MODULE) + except ImportError: + return False + return True + + +def may_name_litellm_mcp_gateway(tools: Iterable[object] | None) -> bool: + items: Final = tuple(tools or ()) + if any(map(names_gateway_explicitly, items)): + return True + return any(proxy_path_mcp_name(tool) is not None for tool in items) and _gateway_registry_loads() diff --git a/tests/base_sdk_tests/check_base_sdk_install.py b/tests/base_sdk_tests/check_base_sdk_install.py index 3abea2a9707..4abd71e39a6 100644 --- a/tests/base_sdk_tests/check_base_sdk_install.py +++ b/tests/base_sdk_tests/check_base_sdk_install.py @@ -102,6 +102,29 @@ def check_completion() -> str: return "mock completion round-trips" +def check_completion_with_function_tool() -> str: + import litellm + + response: Final = litellm.completion( + model="gpt-4o", + messages=[{"role": "user", "content": "ping"}], + tools=[ + { + "type": "function", + "function": { + "name": "noop", + "description": "No operation", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + mock_response="pong", + ) + content: Final = response.choices[0].message.content + _require(content == "pong", f"mock completion with a function tool returned {content!r}") + return "mock completion with a function tool round-trips" + + def check_mcp_install_guidance() -> str: try: import litellm.experimental_mcp_client @@ -209,6 +232,7 @@ CHECKS: tuple[tuple[str, Callable[[], str]], ...] = ( ("import litellm", check_import), ("optional MCP installation guidance", check_mcp_install_guidance), ("chat completion", check_completion), + ("chat completion with a function tool", check_completion_with_function_tool), ("embedding", check_embedding), ("bundled model metadata", check_bundled_model_metadata), ("token counter", check_token_counter), diff --git a/tests/base_sdk_tests/check_sdk_http.py b/tests/base_sdk_tests/check_sdk_http.py index c93306cdced..f7a1e30ac99 100644 --- a/tests/base_sdk_tests/check_sdk_http.py +++ b/tests/base_sdk_tests/check_sdk_http.py @@ -89,7 +89,6 @@ def payload(response: object) -> Mapping[str, JsonValue]: return JSON_OBJECT.validate_json(response.model_dump_json(exclude_none=True)) -@skipIf(importlib.util.find_spec("mcp") is None, "BUG: LIT-9333 completion(tools=...) imports optional mcp") def check_tools(base: str, asynchronous: bool = False) -> None: tools: Final = [ { @@ -122,7 +121,6 @@ def check_tools(base: str, asynchronous: bool = False) -> None: assert json.loads(body) == {"model": "gpt-6", "messages": MESSAGES, "tools": tools} -@skipIf(importlib.util.find_spec("mcp") is None, "BUG: LIT-9333 responses() imports optional mcp") def check_responses(base: str, asynchronous: bool = False) -> None: output: Final[Sequence[JsonValue]] = [ { @@ -144,14 +142,62 @@ def check_responses(base: str, asynchronous: bool = False) -> None: "usage": {"input_tokens": 3, "output_tokens": 2, "total_tokens": 5}, } ) - arguments: Final = dict(model="openai/gpt-6", input="ping", api_key="test-key", api_base=base + "/v1") + tools: Final = [ + {"type": "mcp", "server_label": "docs", "server_url": "https://mcp.example.com/sse"}, + {"type": "mcp", "server_label": "repos", "server_url": "https://mcp.example.com/mcp/repos"}, + ] + arguments: Final = dict(model="openai/gpt-6", input="ping", tools=tools, api_key="test-key", api_base=base + "/v1") response: Final = asyncio.run(litellm.aresponses(**arguments)) if asynchronous else litellm.responses(**arguments) assert payload(response)["output"] == output assert payload(response)["usage"] == {"input_tokens": 3, "output_tokens": 2, "total_tokens": 5} path, headers, body = REQUESTS.get(timeout=10) assert path == "/v1/responses" assert headers["authorization"] == "Bearer test-key" - assert json.loads(body) == {"model": "gpt-6", "input": "ping"} + assert json.loads(body) == {"model": "gpt-6", "input": "ping", "tools": tools} + + +def check_messages_tools(base: str) -> None: + tools: Final = [ + { + "name": "lookup", + "description": "Look up the forecast for a city", + "input_schema": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + } + ] + content: Final = [{"type": "tool_use", "id": "toolu-http-check", "name": "lookup", "input": {"city": "Paris"}}] + enqueue( + { + "id": "msg-tools-http-check", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": content, + "stop_reason": "tool_use", + "stop_sequence": None, + "usage": {"input_tokens": 3, "output_tokens": 2}, + } + ) + + async def send() -> object: + create: Final = cast( # cast-ok: SDK facade exposes untyped kwargs + Callable[..., Awaitable[object]], litellm.anthropic.messages.acreate + ) + return await create( + model="anthropic/claude-sonnet-5-5", + messages=MESSAGES, + tools=tools, + max_tokens=16, + api_key="test-key", + api_base=base, + ) + + response: Final = JSON_OBJECT.validate_python(asyncio.run(send())) + assert response["content"] == content, response + path, headers, body = REQUESTS.get(timeout=10) + assert path == "/v1/messages" + assert headers["x-api-key"] == "test-key" + sent: Final = JSON_OBJECT.validate_json(body) + assert (sent["model"], sent["messages"], sent["tools"]) == ("claude-sonnet-5-5", MESSAGES, tools), sent @skipIf(importlib.util.find_spec("fastapi") is None, "BUG: LIT-9333 cache hits import proxy-only fastapi") @@ -392,6 +438,7 @@ def main() -> None: partial(check_tools, asynchronous=True), check_responses, partial(check_responses, asynchronous=True), + check_messages_tools, check_cache_hit, check_messages_stream_and_callback, ): diff --git a/tests/integration/mcp/test_mcp_llm_endpoints.py b/tests/integration/mcp/test_mcp_llm_endpoints.py index 9f6101423bf..4edf7b620f2 100644 --- a/tests/integration/mcp/test_mcp_llm_endpoints.py +++ b/tests/integration/mcp/test_mcp_llm_endpoints.py @@ -566,6 +566,61 @@ def test_server_scoped_gateway_url_exposes_only_that_servers_tools(gateway: Gate assert [call["body"]["params"]["name"] for call in _peer_add_calls(rig.peer)] == ["add"] +@pytest.mark.parametrize("surface", ("chat", "responses", "messages")) +def test_proxy_path_url_of_a_served_server_is_listed_and_executed_by_the_gateway( + gateway: Gateway, surface: Surface +) -> None: + with _rig(gateway, surface) as rig: + server_url: Final = f"{str(gateway.client.base_url).rstrip('/')}/mcp/{rig.alias}" + response: Final = rig.send(_granted_key(rig), [{**AUTO, "server_url": server_url}]) + assert response.status_code == 200, response.text + requests: Final = rig.upstream_tools() + assert len(requests) == 2 and all(rig.tool in names for names in requests), requests + peer_calls: Final = tuple(JSON_OBJECT.validate_python(call["body"]) for call in _peer_add_calls(rig.peer)) + assert [object_value(call["params"])["name"] for call in peer_calls] == ["add"], peer_calls + assert rig.final_text(JSON_OBJECT.validate_json(response.content)) == ANSWER, response.text + + +LOOKUP_PARAMETERS: Final = {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]} +CALLER_FUNCTION_TOOLS: Final[Mapping[Surface, Mapping[str, object]]] = { + "chat": {"type": "function", "function": {"name": "lookup", "parameters": LOOKUP_PARAMETERS}}, + "responses": {"type": "function", "name": "lookup", "parameters": LOOKUP_PARAMETERS}, + "messages": {"name": "lookup", "input_schema": LOOKUP_PARAMETERS}, +} + + +@pytest.mark.parametrize("surface", ("chat", "responses", "messages")) +def test_tools_that_do_not_name_the_gateway_reach_the_provider_unchanged_and_never_the_peer( + gateway: Gateway, surface: Surface +) -> None: + with _rig(gateway, surface) as rig: + remote_mcp_tools: Final = [ + { + "type": "mcp", + "server_label": "docs", + "server_url": "https://mcp.example.com/sse", + "require_approval": "never", + }, + { + "type": "mcp", + "server_label": "repos", + "server_url": f"https://mcp.example.com/mcp/unserved{uuid.uuid4().hex[:8]}", + "require_approval": "never", + }, + ] + response: Final = rig.send(_granted_key(rig), [CALLER_FUNCTION_TOOLS[surface], *remote_mcp_tools]) + assert response.status_code == 200, response.text + bodies: Final = tuple( + JSON_OBJECT.validate_json(request.body) for request in rig.wire.drain() if request.method == "POST" + ) + assert len(bodies) == 1, bodies + forwarded: Final = bodies[0].get("tools") + assert isinstance(forwarded, list) and len(forwarded) == 3, forwarded + assert [tool for tool in forwarded if isinstance(tool, dict) and tool.get("type") == "mcp"] == remote_mcp_tools + assert "lookup" in json.dumps([tool for tool in forwarded if tool not in remote_mcp_tools]), forwarded + assert tool_calls(rig.peer.drain()) == (), "a tool that does not name the gateway reached the peer" + + def test_streaming_chat_executes_the_tool_once_and_streams_the_follow_up(gateway: Gateway) -> None: with _rig(gateway, "chat") as rig: key: Final = _granted_key(rig) diff --git a/tests/unit/responses/mcp/test_tool_precheck.py b/tests/unit/responses/mcp/test_tool_precheck.py new file mode 100644 index 00000000000..832151ce2a6 --- /dev/null +++ b/tests/unit/responses/mcp/test_tool_precheck.py @@ -0,0 +1,318 @@ +import json +import sys +from collections import deque +from collections.abc import Callable, Iterable, Mapping, Sequence +from importlib.abc import MetaPathFinder +from importlib.machinery import ModuleSpec +from types import ModuleType +from typing import Final + +import httpx +import pytest +import respx + +import litellm +from litellm.responses.mcp.tool_precheck import may_name_litellm_mcp_gateway + +GATEWAY_REGISTRY_MODULE: Final = "litellm.proxy._experimental.mcp_server.mcp_server_manager" +GATEWAY_HANDLER_MODULES: Final = ( + "litellm.responses.mcp.litellm_proxy_mcp_handler", + "litellm.responses.mcp.chat_completions_handler", + "litellm.llms.anthropic.pass_through.messages.mcp_handler", +) + +FUNCTION_TOOL: Final = { + "type": "function", + "function": {"name": "noop", "description": "No operation", "parameters": {"type": "object", "properties": {}}}, +} +RESPONSES_FUNCTION_TOOL: Final = { + "type": "function", + "name": "noop", + "description": "No operation", + "parameters": {"type": "object", "properties": {}}, +} +ANTHROPIC_FUNCTION_TOOL: Final = { + "name": "noop", + "description": "No operation", + "input_schema": {"type": "object", "properties": {}}, +} +ANTHROPIC_MESSAGE_PAYLOAD: Final = { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, +} +GATEWAY_MCP_TOOL: Final = {"type": "mcp", "server_label": "litellm", "server_url": "litellm_proxy"} +GATEWAY_SERVER_MCP_TOOL: Final = {"type": "mcp", "server_label": "x", "server_url": "litellm_proxy/mcp/github"} +GATEWAY_PATH_MCP_TOOL: Final = {"type": "mcp", "server_label": "x", "server_url": "https://gw.example.com/mcp/github"} +REMOTE_MCP_TOOL: Final = { + "type": "mcp", + "server_label": "deepwiki", + "server_url": "https://mcp.deepwiki.com/sse", + "require_approval": "never", +} + +RESPONSES_PAYLOAD: Final = { + "id": "resp_1", + "object": "response", + "created_at": 1734366691, + "status": "completed", + "model": "gpt-4o-mini", + "output": [ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Done.", "annotations": []}], + } + ], + "parallel_tool_calls": True, + "usage": { + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 15, + "output_tokens_details": {"reasoning_tokens": 0}, + }, + "tools": [], +} + + +@pytest.fixture +def gateway_handler_unavailable(monkeypatch: pytest.MonkeyPatch) -> None: + for module_name in GATEWAY_HANDLER_MODULES: + monkeypatch.setitem(sys.modules, module_name, None) + + +class _FailingImport(MetaPathFinder): + def __init__(self, module_name: str, error: ImportError) -> None: + self.module_name: Final = module_name + self.error: Final = error + + def find_spec( + self, fullname: str, path: Sequence[str] | None, target: ModuleType | None = None + ) -> ModuleSpec | None: + if fullname == self.module_name: + raise self.error + return None + + +@pytest.fixture( + params=[ + ModuleNotFoundError("No module named 'mcp'"), + ImportError("cannot import name 'CacheableResult' from 'mcp.types'"), + ModuleNotFoundError("No module named 'fastapi'"), + ], + ids=["mcp-missing", "mcp-1x", "proxy-extras-missing"], +) +def gateway_registry_unavailable( + request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch, gateway_handler_unavailable: None +) -> None: + monkeypatch.delitem(sys.modules, GATEWAY_REGISTRY_MODULE, raising=False) + monkeypatch.setattr(sys, "meta_path", [_FailingImport(GATEWAY_REGISTRY_MODULE, request.param), *sys.meta_path]) + + +class _ToolList(list[Mapping[str, object]]): + pass + + +@pytest.mark.parametrize( + "tools", + [[FUNCTION_TOOL], [REMOTE_MCP_TOOL], _ToolList([FUNCTION_TOOL]), (FUNCTION_TOOL, REMOTE_MCP_TOOL)], + ids=["function-tool", "remote-mcp-tool", "list-subclass", "tuple"], +) +@pytest.mark.usefixtures("gateway_handler_unavailable") +def test_completion_without_a_gateway_tool_does_not_import_the_gateway_handler(tools: Iterable[Mapping[str, object]]): + response: Final = litellm.completion( + model="openai/gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + tools=tools, + mock_response="ok", + ) + + assert response.choices[0].message.content == "ok" + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("gateway_handler_unavailable") +async def test_acompletion_with_function_tool_does_not_import_the_gateway_handler(): + response: Final = await litellm.acompletion( + model="openai/gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + tools=[FUNCTION_TOOL], + mock_response="ok", + ) + + assert response.choices[0].message.content == "ok" + + +@pytest.mark.parametrize( + "tool", + [GATEWAY_MCP_TOOL, GATEWAY_SERVER_MCP_TOOL, GATEWAY_PATH_MCP_TOOL], + ids=["litellm-proxy", "litellm-proxy-server", "gateway-path-url"], +) +@pytest.mark.usefixtures("gateway_handler_unavailable") +def test_completion_with_gateway_mcp_tool_still_reaches_the_handler(tool: Mapping[str, object]): + with pytest.raises(ModuleNotFoundError, match=r"litellm_proxy_mcp_handler|chat_completions_handler"): + litellm.completion( + model="openai/gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + tools=[tool], + mock_response="ok", + ) + + +@pytest.mark.parametrize( + "tools", + [None, [RESPONSES_FUNCTION_TOOL], [REMOTE_MCP_TOOL]], + ids=["no-tools", "function-tool", "remote-mcp-tool"], +) +@pytest.mark.usefixtures("gateway_handler_unavailable") +def test_responses_without_a_gateway_tool_does_not_import_the_gateway_handler( + tools: list[Mapping[str, object]] | None, respx_mock: respx.MockRouter +): + upstream: Final = respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=httpx.Response(200, json=RESPONSES_PAYLOAD) + ) + + response: Final = litellm.responses(model="openai/gpt-4o-mini", input="hi", tools=tools, api_key="sk-test") + + assert upstream.called + assert response.output[0].content[0].text == "Done." + + +@pytest.mark.usefixtures("gateway_handler_unavailable") +def test_responses_with_gateway_mcp_tool_still_reaches_the_handler(): + with pytest.raises(litellm.APIConnectionError, match=r"litellm_proxy_mcp_handler"): + litellm.responses(model="openai/gpt-4o-mini", input="hi", tools=[GATEWAY_MCP_TOOL], api_key="sk-test") + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("gateway_handler_unavailable") +async def test_anthropic_messages_with_function_tool_does_not_import_the_gateway_handler( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + upstream: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock( + return_value=httpx.Response(200, json=ANTHROPIC_MESSAGE_PAYLOAD) + ) + + response: Final = await litellm.anthropic_interface.messages.acreate( + model="anthropic/claude-sonnet-5-5", + messages=[{"role": "user", "content": "hi"}], + max_tokens=16, + tools=[ANTHROPIC_FUNCTION_TOOL], + api_key="sk-test", + ) + + assert upstream.called + assert response["content"][0]["text"] == "ok" + + +@pytest.mark.parametrize( + "tools, expected", + [ + (None, False), + ([], False), + ([FUNCTION_TOOL], False), + ([RESPONSES_FUNCTION_TOOL], False), + ([{"type": "web_search"}], False), + (["not-a-mapping"], False), + ([REMOTE_MCP_TOOL], False), + ([{"type": "mcp", "server_label": "x"}], False), + ([{"type": "mcp", "server_label": "x", "server_url": None}], False), + ([{"server_url": "litellm_proxy"}], False), + ([FUNCTION_TOOL, GATEWAY_MCP_TOOL], True), + ([GATEWAY_SERVER_MCP_TOOL], True), + ([GATEWAY_PATH_MCP_TOOL], True), + ([REMOTE_MCP_TOOL, GATEWAY_PATH_MCP_TOOL], True), + ], + ids=[ + "none", + "empty", + "chat-function", + "responses-function", + "web-search", + "non-mapping", + "remote-mcp", + "mcp-without-url", + "mcp-null-url", + "gateway-url-without-mcp-type", + "gateway-among-functions", + "gateway-server", + "gateway-path-url", + "gateway-path-after-remote-mcp", + ], +) +def test_may_name_litellm_mcp_gateway_matches_only_gateway_mcp_tools(tools: list[object] | None, expected: bool): + assert may_name_litellm_mcp_gateway(tools) is expected + + +@pytest.mark.parametrize( + "tools, expected", + [ + ([GATEWAY_PATH_MCP_TOOL], False), + ([FUNCTION_TOOL, GATEWAY_PATH_MCP_TOOL], False), + ([GATEWAY_MCP_TOOL], True), + ([GATEWAY_SERVER_MCP_TOOL, GATEWAY_PATH_MCP_TOOL], True), + ], + ids=["gateway-path-url", "gateway-path-among-functions", "litellm-proxy", "litellm-proxy-server-and-path-url"], +) +@pytest.mark.usefixtures("gateway_registry_unavailable") +def test_proxy_path_url_names_the_gateway_only_when_its_registry_loads(tools: list[object], expected: bool): + assert may_name_litellm_mcp_gateway(tools) is expected + + +@pytest.mark.usefixtures("gateway_registry_unavailable") +def test_responses_sends_a_proxy_path_mcp_url_to_the_provider_without_the_gateway_registry( + respx_mock: respx.MockRouter, +): + upstream: Final = respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=httpx.Response(200, json=RESPONSES_PAYLOAD) + ) + + response: Final = litellm.responses( + model="openai/gpt-4o-mini", input="hi", tools=[GATEWAY_PATH_MCP_TOOL], api_key="sk-test" + ) + + assert response.output[0].content[0].text == "Done." + assert json.loads(upstream.calls.last.request.content)["tools"] == [GATEWAY_PATH_MCP_TOOL] + + +@pytest.mark.usefixtures("gateway_registry_unavailable") +def test_completion_sends_a_proxy_path_mcp_url_to_the_provider_without_the_gateway_registry(): + response: Final = litellm.completion( + model="openai/gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + tools=[GATEWAY_PATH_MCP_TOOL], + mock_response="ok", + ) + + assert response.choices[0].message.content == "ok" + + +class _ToolTuple(tuple[Mapping[str, object], ...]): + pass + + +@pytest.mark.parametrize( + "make_tools", + [ + lambda items: (item for item in items), + lambda items: iter(items), + lambda items: deque(items), + lambda items: _ToolList(items), + lambda items: _ToolTuple(items), + ], + ids=["generator", "iterator", "deque", "list-subclass", "tuple-subclass"], +) +@pytest.mark.usefixtures("gateway_handler_unavailable") +def test_precheck_reads_every_iterable_shape_without_the_gateway_handler( + make_tools: Callable[[list[Mapping[str, object]]], Iterable[Mapping[str, object]]], +): + assert may_name_litellm_mcp_gateway(make_tools([FUNCTION_TOOL, REMOTE_MCP_TOOL])) is False + assert may_name_litellm_mcp_gateway(make_tools([FUNCTION_TOOL, GATEWAY_MCP_TOOL])) is True diff --git a/tests/unit/responses/test_responses_prompt_management.py b/tests/unit/responses/test_responses_prompt_management.py index 1e8880424ea..9af92e7f9ca 100644 --- a/tests/unit/responses/test_responses_prompt_management.py +++ b/tests/unit/responses/test_responses_prompt_management.py @@ -70,10 +70,6 @@ def _patch_responses_dispatch(): import_module("litellm.responses.main").litellm, "get_llm_provider", side_effect=_provider_by_model, ), - patch.object( - import_module("litellm.responses.mcp.litellm_proxy_mcp_handler").LiteLLM_Proxy_MCP_Handler, "should_use_litellm_mcp_gateway", - return_value=False, - ), patch.object( import_module("litellm.responses.main").ProviderConfigManager, "get_provider_responses_api_config", return_value=None, @@ -158,7 +154,7 @@ class TestResponsesAPIPromptManagement: ) patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3]: + with patches[0], patches[1], patches[2]: import litellm litellm.responses( @@ -193,7 +189,7 @@ class TestResponsesAPIPromptManagement: ) patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3]: + with patches[0], patches[1], patches[2]: import litellm litellm.responses( @@ -216,7 +212,7 @@ class TestResponsesAPIPromptManagement: ) patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3]: + with patches[0], patches[1], patches[2]: import litellm litellm.responses( @@ -247,7 +243,7 @@ class TestResponsesAPIPromptManagement: logging_obj.model_call_details = {} patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3] as mock_handler: + with patches[0], patches[1], patches[2] as mock_handler: import litellm litellm.responses( @@ -273,7 +269,7 @@ class TestResponsesAPIPromptManagement: ) patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3] as mock_handler: + with patches[0], patches[1], patches[2] as mock_handler: import litellm litellm.responses( @@ -305,7 +301,7 @@ class TestResponsesAPIPromptManagement: ) patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3]: + with patches[0], patches[1], patches[2]: import litellm litellm.responses( @@ -328,7 +324,7 @@ class TestResponsesAPIPromptManagement: ) patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3] as mock_handler: + with patches[0], patches[1], patches[2] as mock_handler: import litellm litellm.responses( @@ -368,7 +364,7 @@ class TestResponsesAPIPromptManagement: ) patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3] as mock_handler: + with patches[0], patches[1], patches[2] as mock_handler: import litellm litellm.responses( @@ -399,8 +395,7 @@ class TestResponsesAPIPromptManagement: side_effect=_provider_by_model, ), patches[1], - patches[2], - patches[3] as mock_handler, + patches[2] as mock_handler, ): import litellm @@ -437,7 +432,7 @@ class TestAsyncResponsesAPIPromptManagement: ) patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3]: + with patches[0], patches[1], patches[2]: import litellm await litellm.aresponses( @@ -467,7 +462,7 @@ class TestAsyncResponsesAPIPromptManagement: ) patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3] as mock_handler: + with patches[0], patches[1], patches[2] as mock_handler: import litellm await litellm.aresponses( @@ -498,7 +493,7 @@ class TestAsyncResponsesAPIPromptManagement: ) patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3]: + with patches[0], patches[1], patches[2]: import litellm await litellm.aresponses( @@ -524,7 +519,7 @@ class TestAsyncResponsesAPIPromptManagement: ) patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3] as mock_handler: + with patches[0], patches[1], patches[2] as mock_handler: import litellm await litellm.aresponses( @@ -657,7 +652,7 @@ async def test_async_guardrail_block_from_prompt_hook_reaches_caller_unwrapped() logging_obj.async_get_chat_completion_prompt = AsyncMock(side_effect=block) patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3], pytest.raises(HTTPException) as exc_info: + with patches[0], patches[1], patches[2], pytest.raises(HTTPException) as exc_info: await litellm.aresponses(input="Hi", model="gpt-4o", prompt_id="blocked", litellm_logging_obj=logging_obj) assert exc_info.value is block @@ -669,7 +664,7 @@ def test_sync_guardrail_block_from_prompt_hook_reaches_caller_unwrapped(): logging_obj.get_chat_completion_prompt.side_effect = block patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3], pytest.raises(HTTPException) as exc_info: + with patches[0], patches[1], patches[2], pytest.raises(HTTPException) as exc_info: litellm.responses(input="Hi", model="gpt-4o", prompt_id="blocked", litellm_logging_obj=logging_obj) assert exc_info.value is block