diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 91209b36c90..e64626013ed 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -3426,6 +3426,77 @@ class MCPServerManager: GuardrailRaisedException: If guardrails block the call HTTPException: If an HTTP error occurs """ + tasks.append( + asyncio.create_task( + self._open_and_call_tool( + mcp_server, + original_tool_name, + arguments, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + hook_extra_headers=hook_extra_headers, + host_progress_callback=host_progress_callback, + user_api_key_auth=user_api_key_auth, + ) + ) + ) + + _timeout = ( + mcp_server.timeout if mcp_server.timeout is not None else MCP_CLIENT_TIMEOUT + ) + try: + mcp_responses = await asyncio.wait_for( + asyncio.gather(*tasks), timeout=_timeout + ) + except asyncio.TimeoutError: + raise HTTPException( + status_code=504, + detail={ + "error": "timeout", + "message": f"MCP tool call timed out after {_timeout}s", + }, + ) + except ( + BlockedPiiEntityError, + GuardrailRaisedException, + HTTPException, + ) as e: + verbose_logger.error( + f"Guardrail blocked MCP tool call during result check: {str(e)}" + ) + raise e + + # If proxy_logging_obj is None, the tool call result is at index 0 + # If proxy_logging_obj is not None, the tool call result is at index 1 (after the during hook task) + result_index = 1 if proxy_logging_obj else 0 + result = mcp_responses[result_index] + + return cast(CallToolResult, result) + + async def _open_and_call_tool( + self, + mcp_server: MCPServer, + original_tool_name: str, + arguments: Dict[str, Any], + *, + mcp_auth_header: Optional[str], + mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]], + oauth2_headers: Optional[Dict[str, str]], + raw_headers: Optional[Dict[str, str]], + hook_extra_headers: Optional[Dict[str, str]], + host_progress_callback: Optional["ProgressFnT"], + user_api_key_auth: Optional[UserAPIKeyAuth], + ) -> CallToolResult: + """Open the upstream connection and invoke one tool, returning its result. + + The transport seam factored out of _call_regular_mcp_tool so v2 can override only the + connect-and-call step (resolve() + UpstreamConnection) while inheriting the + guardrail/hook, gather, timeout, and result-extraction tail. This v1 implementation + builds the outbound headers, creates the MCPClient, calls the tool, and caches the + upstream's initialize instructions. + """ # Get server-specific auth header if available (case-insensitive) # FIX: Added case-insensitive matching to handle auth header keys that may not match # the exact case of server alias/name (e.g., '1litellmagcgateway' vs '1LiteLLMAGCGateway') @@ -3538,44 +3609,9 @@ class MCPServerManager: arguments=arguments, ) - async def _call_tool_via_client(client, params): - return await client.call_tool( - params, host_progress_callback=host_progress_callback - ) - - tasks.append( - asyncio.create_task(_call_tool_via_client(client, call_tool_params)) + result = await client.call_tool( + call_tool_params, host_progress_callback=host_progress_callback ) - - _timeout = ( - mcp_server.timeout if mcp_server.timeout is not None else MCP_CLIENT_TIMEOUT - ) - try: - mcp_responses = await asyncio.wait_for( - asyncio.gather(*tasks), timeout=_timeout - ) - except asyncio.TimeoutError: - raise HTTPException( - status_code=504, - detail={ - "error": "timeout", - "message": f"MCP tool call timed out after {_timeout}s", - }, - ) - except ( - BlockedPiiEntityError, - GuardrailRaisedException, - HTTPException, - ) as e: - verbose_logger.error( - f"Guardrail blocked MCP tool call during result check: {str(e)}" - ) - raise e - - # If proxy_logging_obj is None, the tool call result is at index 0 - # If proxy_logging_obj is not None, the tool call result is at index 1 (after the during hook task) - result_index = 1 if proxy_logging_obj else 0 - result = mcp_responses[result_index] self._remember_upstream_initialize_instructions(mcp_server, client) return cast(CallToolResult, result) @@ -4538,6 +4574,8 @@ def _make_global_mcp_server_manager() -> MCPServerManager: _global_mcp_server_manager: Optional[MCPServerManager] = None if TYPE_CHECKING: + from mcp.shared.session import ProgressFnT + global_mcp_server_manager: MCPServerManager diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager_v2.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager_v2.py index 8c348fb1e72..39d6001b89d 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager_v2.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager_v2.py @@ -27,7 +27,14 @@ from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerM from litellm.proxy._types import MCPTransport if TYPE_CHECKING: - from mcp.types import GetPromptResult, Prompt, ReadResourceResult, Resource + from mcp.shared.session import ProgressFnT + from mcp.types import ( + CallToolResult, + GetPromptResult, + Prompt, + ReadResourceResult, + Resource, + ) from mcp.types import Tool as MCPTool from pydantic import AnyUrl @@ -82,6 +89,8 @@ class MCPServerManagerV2(MCPServerManager): raw_headers: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, subject_token: Optional[str] = None, + forward_caller_headers: bool = False, + raise_on_missing_env: bool = False, ) -> Result[UpstreamConnection, CredError]: """The shared egress seam: resolve auth via resolve() and build the UpstreamConnection. @@ -106,10 +115,14 @@ class MCPServerManagerV2(MCPServerManager): ) if isinstance(auth, Error): return Error(auth.error) - resolved_static = await self._resolve_static_headers_with_env_vars( - server, user_api_key_auth, raise_on_missing=False + egress = await self._build_egress_headers( + server, + user_api_key_auth, + raw_headers, + forward_caller_headers=forward_caller_headers, + raise_on_missing_env=raise_on_missing_env, ) - headers = {**(extra_headers or {}), **(resolved_static or {})} or None + headers = {**(extra_headers or {}), **egress} or None is_stdio = server.transport == MCPTransport.stdio return Ok( UpstreamConnection( @@ -123,6 +136,53 @@ class MCPServerManagerV2(MCPServerManager): ) ) + async def _build_egress_headers( + self, + server: MCPServer, + user_api_key_auth: Optional[UserAPIKeyAuth], + raw_headers: Optional[Dict[str, str]], + *, + forward_caller_headers: bool, + raise_on_missing_env: bool, + ) -> Dict[str, str]: + """The non-credential egress headers: configured static headers (with per-user env-var + interpolation) plus, on the call path, the caller headers the server forwards upstream. + + The credential (Authorization) is resolve()'s, never this dict. raise_on_missing_env + propagates MCPMissingUserEnvVarsError (412 + setup URL) on the call path; list ops degrade. + forward_caller_headers gates caller-header forwarding (call only; list/prompts/resources do + not forward, matching v1). + """ + static = await self._resolve_static_headers_with_env_vars( + server, user_api_key_auth, raise_on_missing=raise_on_missing_env + ) + forwarded = ( + self._forwarded_request_headers(server, raw_headers) + if forward_caller_headers + else {} + ) + return {**forwarded, **(static or {})} + + @staticmethod + def _forwarded_request_headers( + server: MCPServer, + raw_headers: Optional[Dict[str, str]], + ) -> Dict[str, str]: + """Caller request headers the server is configured to forward (server.extra_headers pulled + from raw_headers). Authorization is always stripped: the upstream credential is resolve()'s, + so the forwarder never ships inbound auth upstream (the credential-isolation invariant; the + passthrough/override paths that would forward an inbound token defer to v1). + """ + if not server.extra_headers or not raw_headers: + return {} + normalized = {k.lower(): v for k, v in raw_headers.items()} + return { + header: normalized[header.lower()] + for header in server.extra_headers + if header.lower() != "authorization" + and normalized.get(header.lower()) is not None + } + async def _get_tools_from_server( self, server: MCPServer, @@ -266,6 +326,69 @@ class MCPServerManagerV2(MCPServerManager): self._egress_item_failure(server, result.error) return result.ok + async def _open_and_call_tool( + self, + mcp_server: MCPServer, + original_tool_name: str, + arguments: Dict[str, object], + *, + mcp_auth_header: Optional[str], + mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]], + oauth2_headers: Optional[Dict[str, str]], + raw_headers: Optional[Dict[str, str]], + hook_extra_headers: Optional[Dict[str, str]], + host_progress_callback: Optional[ProgressFnT], + user_api_key_auth: Optional[UserAPIKeyAuth], + ) -> CallToolResult: + from litellm.proxy.gateway.mcp.result import Error + + # The effective inbound credential is the per-server header if present, else the deprecated + # mcp_auth_header (matches v1); an override here means the request defers to v1. + server_auth_header: Optional[Union[str, Dict[str, str]]] = mcp_auth_header + if mcp_server_auth_headers: + from litellm.proxy._experimental.mcp_server.utils import ( + lookup_mcp_server_auth_in_headers, + ) + + found = lookup_mcp_server_auth_in_headers( + mcp_server_auth_headers, + alias=mcp_server.alias, + server_name=mcp_server.server_name, + ) + if found is not None: + server_auth_header = found + + # Defer to v1 for guardrail-injected headers (JWT signer etc.), which the v2 path does not + # apply, in addition to the usual _should_defer cases (unmapped mode, OpenAPI, override). + if hook_extra_headers or self._should_defer(mcp_server, server_auth_header): + return await super()._open_and_call_tool( + mcp_server, + original_tool_name, + arguments, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + hook_extra_headers=hook_extra_headers, + host_progress_callback=host_progress_callback, + user_api_key_auth=user_api_key_auth, + ) + conn = await self._v2_connection( + mcp_server, + user_api_key_auth, + raw_headers=raw_headers, + forward_caller_headers=True, + raise_on_missing_env=True, + ) + if isinstance(conn, Error): + self._egress_item_failure(mcp_server, conn.error) + result = await conn.ok.call_tool( + original_tool_name, arguments, host_progress_callback + ) + if isinstance(result, Error): + self._egress_item_failure(mcp_server, result.error) + return result.ok + def _egress_list_failure( self, server: MCPServer, error: "CredError | ConnError" ) -> None: diff --git a/litellm/proxy/_experimental/mcp_server/v2_egress.py b/litellm/proxy/_experimental/mcp_server/v2_egress.py index 338d3adb007..f2d4b1c536b 100644 --- a/litellm/proxy/_experimental/mcp_server/v2_egress.py +++ b/litellm/proxy/_experimental/mcp_server/v2_egress.py @@ -41,6 +41,7 @@ if TYPE_CHECKING: ) from mcp import ReadResourceResult, Resource from mcp.shared.message import SessionMessage + from mcp.shared.session import ProgressFnT from mcp.types import CallToolResult, GetPromptResult, Prompt from mcp.types import Tool as MCPTool from pydantic import AnyUrl @@ -193,10 +194,15 @@ class UpstreamConnection: return await self._run(op) async def call_tool( - self, name: str, arguments: Dict[str, object] + self, + name: str, + arguments: Dict[str, object], + progress_callback: Optional[ProgressFnT] = None, ) -> Result[CallToolResult, ConnError]: async def op(session: ClientSession) -> CallToolResult: - return await session.call_tool(name, arguments) + return await session.call_tool( + name, arguments, progress_callback=progress_callback + ) return await self._run(op) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager_v2.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager_v2.py index 3da6fd47daa..9b5ba53ce05 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager_v2.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager_v2.py @@ -159,3 +159,30 @@ async def test_v2_override_gets_prompt_via_upstream_connection(echo_server_url): ) result = await manager.get_prompt_from_server(server, "greeting", {"name": "Tin"}) assert result.messages + + +@pytest.mark.asyncio +async def test_v2_override_calls_tool_via_upstream_connection(echo_server_url): + # Tool-call path: the factored _open_and_call_tool seam routes through resolve() + + # UpstreamConnection.call_tool (v2) for the none mode, returning the upstream CallToolResult. + manager = MCPServerManagerV2() + server = MCPServer( + server_id="echo1", + name="echo1", + transport=MCPTransport.http, + url=echo_server_url, + auth_type=MCPAuth.none, + ) + result = await manager._open_and_call_tool( + server, + "echo", + {"text": "hi"}, + mcp_auth_header=None, + mcp_server_auth_headers=None, + oauth2_headers=None, + raw_headers=None, + hook_extra_headers=None, + host_progress_callback=None, + user_api_key_auth=None, + ) + assert any("echo: hi" in getattr(c, "text", "") for c in result.content)