from __future__ import annotations import abc import asyncio import inspect import sys from collections.abc import AsyncGenerator, Awaitable, Callable from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager from datetime import timedelta from pathlib import Path from typing import TYPE_CHECKING, Any, Literal, TypeVar, Union, cast import anyio import httpx if sys.version_info < (3, 11): from exceptiongroup import BaseExceptionGroup # pyright: ignore[reportMissingImports] from anyio import ClosedResourceError from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream from mcp import ClientSession, StdioServerParameters, Tool as MCPTool, stdio_client from mcp.client.session import MessageHandlerFnT from mcp.client.sse import sse_client from mcp.client.streamable_http import ( GetSessionIdCallback, StreamableHTTPTransport, streamablehttp_client, ) from mcp.shared.exceptions import McpError from mcp.shared.message import SessionMessage from mcp.types import ( CallToolResult, GetPromptResult, InitializeResult, ListPromptsResult, ListResourcesResult, ListResourceTemplatesResult, ReadResourceResult, ) from typing_extensions import NotRequired, TypedDict from ..exceptions import UserError from ..logger import logger from ..run_context import RunContextWrapper from ..tool import ToolErrorFunction from ..util._types import MaybeAwaitable from .util import ( HttpClientFactory, MCPToolMetaResolver, ToolFilter, ToolFilterContext, ToolFilterStatic, ) class RequireApprovalToolList(TypedDict, total=False): tool_names: list[str] class RequireApprovalObject(TypedDict, total=False): always: RequireApprovalToolList never: RequireApprovalToolList RequireApprovalPolicy = Literal["always", "never"] RequireApprovalMapping = dict[str, RequireApprovalPolicy] if TYPE_CHECKING: LocalMCPApprovalCallable = Callable[ [RunContextWrapper[Any], "AgentBase", MCPTool], MaybeAwaitable[bool], ] else: LocalMCPApprovalCallable = Callable[..., Any] if TYPE_CHECKING: RequireApprovalSetting = ( RequireApprovalPolicy | RequireApprovalObject | RequireApprovalMapping | LocalMCPApprovalCallable | bool | None ) else: RequireApprovalSetting = Union[ # noqa: UP007 RequireApprovalPolicy, RequireApprovalObject, RequireApprovalMapping, LocalMCPApprovalCallable, bool, None, ] T = TypeVar("T") def _create_default_streamable_http_client( headers: dict[str, str] | None = None, timeout: httpx.Timeout | None = None, auth: httpx.Auth | None = None, ) -> httpx.AsyncClient: kwargs: dict[str, Any] = {"follow_redirects": True} if timeout is not None: kwargs["timeout"] = timeout if headers is not None: kwargs["headers"] = headers if auth is not None: kwargs["auth"] = auth return httpx.AsyncClient(**kwargs) class _InitializedNotificationTolerantStreamableHTTPTransport(StreamableHTTPTransport): async def _handle_post_request(self, ctx: Any) -> None: message = ctx.session_message.message if not self._is_initialized_notification(message): await super()._handle_post_request(ctx) return try: await super()._handle_post_request(ctx) except httpx.HTTPError: logger.warning( "Ignoring initialized notification HTTP failure", exc_info=True, ) return @asynccontextmanager async def _streamablehttp_client_with_transport( url: str, *, headers: dict[str, str] | None = None, timeout: float | timedelta = 30, sse_read_timeout: float | timedelta = 60 * 5, terminate_on_close: bool = True, httpx_client_factory: HttpClientFactory = _create_default_streamable_http_client, auth: httpx.Auth | None = None, transport_factory: Callable[[str], StreamableHTTPTransport] = StreamableHTTPTransport, ) -> AsyncGenerator[MCPStreamTransport, None]: timeout_seconds = timeout.total_seconds() if isinstance(timeout, timedelta) else timeout sse_read_timeout_seconds = ( sse_read_timeout.total_seconds() if isinstance(sse_read_timeout, timedelta) else sse_read_timeout ) client = httpx_client_factory( headers=headers, timeout=httpx.Timeout(timeout_seconds, read=sse_read_timeout_seconds), auth=auth, ) transport = transport_factory(url) read_stream_writer, read_stream = anyio.create_memory_object_stream[SessionMessage | Exception]( 0 ) write_stream, write_stream_reader = anyio.create_memory_object_stream[SessionMessage](0) async with client: async with anyio.create_task_group() as tg: try: logger.debug(f"Connecting to StreamableHTTP endpoint: {url}") def start_get_stream() -> None: tg.start_soon(transport.handle_get_stream, client, read_stream_writer) tg.start_soon( transport.post_writer, client, write_stream_reader, read_stream_writer, write_stream, start_get_stream, tg, ) try: yield ( read_stream, write_stream, transport.get_session_id, ) finally: if transport.session_id and terminate_on_close: await transport.terminate_session(client) tg.cancel_scope.cancel() finally: await read_stream_writer.aclose() await write_stream.aclose() class _SharedSessionRequestNeedsIsolation(Exception): """Raised when a shared-session request should be retried on an isolated session.""" class _IsolatedSessionRetryFailed(Exception): """Raised when an isolated-session retry fails after consuming retry budget.""" class _UnsetType: pass _UNSET = _UnsetType() if TYPE_CHECKING: from ..agent import AgentBase MCPStreamTransport = ( tuple[ MemoryObjectReceiveStream[SessionMessage | Exception], MemoryObjectSendStream[SessionMessage], ] | tuple[ MemoryObjectReceiveStream[SessionMessage | Exception], MemoryObjectSendStream[SessionMessage], GetSessionIdCallback | None, ] ) class MCPServer(abc.ABC): """Base class for Model Context Protocol servers.""" def __init__( self, use_structured_content: bool = False, require_approval: RequireApprovalSetting = None, failure_error_function: ToolErrorFunction | None | _UnsetType = _UNSET, tool_meta_resolver: MCPToolMetaResolver | None = None, ): """ Args: use_structured_content: Whether to use `tool_result.structured_content` when calling an MCP tool.Defaults to False for backwards compatibility - most MCP servers still include the structured content in the `tool_result.content`, and using it by default will cause duplicate content. You can set this to True if you know the server will not duplicate the structured content in the `tool_result.content`. require_approval: Approval policy for tools on this server. Accepts "always"/"never", a dict of tool names to those values, a boolean, an object with always/never tool lists (mirroring TS requireApproval), or a sync/async callable that receives `(run_context, agent, tool)` and returns whether the tool call needs approval. Normalized into a needs_approval policy. failure_error_function: Optional function used to convert MCP tool failures into a model-visible error message. If explicitly set to None, tool errors will be raised instead of converted. If left unset, the agent-level configuration (or SDK default) will be used. tool_meta_resolver: Optional callable that produces MCP request metadata (`_meta`) for tool calls. It is invoked by the Agents SDK before calling `call_tool`. """ self.use_structured_content = use_structured_content self._needs_approval_policy = self._normalize_needs_approval( require_approval=require_approval ) self._failure_error_function = failure_error_function self.tool_meta_resolver = tool_meta_resolver @abc.abstractmethod async def connect(self): """Connect to the server. For example, this might mean spawning a subprocess or opening a network connection. The server is expected to remain connected until `cleanup()` is called. """ pass @property @abc.abstractmethod def name(self) -> str: """A readable name for the server.""" pass @abc.abstractmethod async def cleanup(self): """Cleanup the server. For example, this might mean closing a subprocess or closing a network connection. """ pass @abc.abstractmethod async def list_tools( self, run_context: RunContextWrapper[Any] | None = None, agent: AgentBase | None = None, ) -> list[MCPTool]: """List the tools available on the server.""" pass @abc.abstractmethod async def call_tool( self, tool_name: str, arguments: dict[str, Any] | None, meta: dict[str, Any] | None = None, ) -> CallToolResult: """Invoke a tool on the server.""" pass @property def cached_tools(self) -> list[MCPTool] | None: """Return the most recently fetched tools list, if available. Implementations may return `None` when tools have not been fetched yet or caching is disabled. """ return None @abc.abstractmethod async def list_prompts( self, ) -> ListPromptsResult: """List the prompts available on the server.""" pass @abc.abstractmethod async def get_prompt( self, name: str, arguments: dict[str, Any] | None = None ) -> GetPromptResult: """Get a specific prompt from the server.""" pass async def list_resources(self, cursor: str | None = None) -> ListResourcesResult: """List the resources available on the server. Args: cursor: An opaque pagination cursor returned in a previous :class:`~mcp.types.ListResourcesResult` as ``nextCursor``. Pass it here to fetch the next page of results. ``None`` fetches the first page. Returns a :class:`~mcp.types.ListResourcesResult`. When the result contains a ``nextCursor`` field, call this method again with that cursor to retrieve the next page. Subclasses that do not support resources may leave this unimplemented; it will raise :exc:`NotImplementedError` at call time. """ raise NotImplementedError( f"MCP server '{self.name}' does not support list_resources. " "Override this method in your server implementation." ) async def list_resource_templates( self, cursor: str | None = None ) -> ListResourceTemplatesResult: """List the resource templates available on the server. Args: cursor: An opaque pagination cursor returned in a previous :class:`~mcp.types.ListResourceTemplatesResult` as ``nextCursor``. Pass it here to fetch the next page of results. ``None`` fetches the first page. Returns a :class:`~mcp.types.ListResourceTemplatesResult`. When the result contains a ``nextCursor`` field, call this method again with that cursor to retrieve the next page. Subclasses that do not support resource templates may leave this unimplemented; it will raise :exc:`NotImplementedError` at call time. """ raise NotImplementedError( f"MCP server '{self.name}' does not support list_resource_templates. " "Override this method in your server implementation." ) async def read_resource(self, uri: str) -> ReadResourceResult: """Read the contents of a specific resource by URI. Args: uri: The URI of the resource to read. See :class:`~pydantic.networks.AnyUrl` for the supported URI formats. Returns a :class:`~mcp.types.ReadResourceResult`. Subclasses that do not support resources may leave this unimplemented; it will raise :exc:`NotImplementedError` at call time. """ raise NotImplementedError( f"MCP server '{self.name}' does not support read_resource. " "Override this method in your server implementation." ) @staticmethod def _normalize_needs_approval( *, require_approval: RequireApprovalSetting, ) -> ( bool | dict[str, bool] | Callable[[RunContextWrapper[Any], AgentBase, MCPTool], MaybeAwaitable[bool]] ): """Normalize approval inputs to booleans or a name->bool map.""" if require_approval is None: return False def _to_bool(value: str) -> bool: return value == "always" def _is_tool_list_schema(value: object) -> bool: if not isinstance(value, dict): return False for key in ("always", "never"): if key not in value: continue entry = value.get(key) if isinstance(entry, dict) and "tool_names" in entry: return True return False if isinstance(require_approval, dict) and _is_tool_list_schema(require_approval): always_entry: RequireApprovalToolList | Any = require_approval.get("always", {}) never_entry: RequireApprovalToolList | Any = require_approval.get("never", {}) always_names = ( always_entry.get("tool_names", []) if isinstance(always_entry, dict) else [] ) never_names = never_entry.get("tool_names", []) if isinstance(never_entry, dict) else [] tool_list_mapping: dict[str, bool] = {} for name in always_names: tool_list_mapping[str(name)] = True for name in never_names: tool_list_mapping[str(name)] = False return tool_list_mapping if isinstance(require_approval, dict): tool_mapping: dict[str, bool] = {} for name, value in require_approval.items(): if isinstance(value, bool): tool_mapping[str(name)] = value elif isinstance(value, str) and value in ("always", "never"): tool_mapping[str(name)] = _to_bool(value) return tool_mapping if callable(require_approval): return require_approval if isinstance(require_approval, bool): return require_approval return _to_bool(require_approval) def _get_needs_approval_for_tool( self, tool: MCPTool, agent: AgentBase | None, ) -> bool | Callable[[RunContextWrapper[Any], dict[str, Any], str], Awaitable[bool]]: """Return a FunctionTool.needs_approval value for a given MCP tool. Legacy callers may omit ``agent`` when using ``MCPUtil.to_function_tool()`` directly. When approval is configured with a callable policy and no agent is available, this method returns ``True`` to preserve the historical fail-closed behavior. """ policy = self._needs_approval_policy if callable(policy): if agent is None: return True async def _needs_approval( run_context: RunContextWrapper[Any], _args: dict[str, Any], _call_id: str ) -> bool: result = policy(run_context, agent, tool) if inspect.isawaitable(result): result = await result return bool(result) return _needs_approval if isinstance(policy, dict): return bool(policy.get(tool.name, False)) return bool(policy) def _get_failure_error_function( self, agent_failure_error_function: ToolErrorFunction | None ) -> ToolErrorFunction | None: """Return the effective error handler for MCP tool failures.""" if self._failure_error_function is _UNSET: return agent_failure_error_function return cast(ToolErrorFunction | None, self._failure_error_function) class _MCPServerWithClientSession(MCPServer, abc.ABC): """Base class for MCP servers that use a `ClientSession` to communicate with the server.""" @property def cached_tools(self) -> list[MCPTool] | None: return self._tools_list def __init__( self, cache_tools_list: bool, client_session_timeout_seconds: float | None, tool_filter: ToolFilter = None, use_structured_content: bool = False, max_retry_attempts: int = 0, retry_backoff_seconds_base: float = 1.0, message_handler: MessageHandlerFnT | None = None, require_approval: RequireApprovalSetting = None, failure_error_function: ToolErrorFunction | None | _UnsetType = _UNSET, tool_meta_resolver: MCPToolMetaResolver | None = None, ): """ Args: cache_tools_list: Whether to cache the tools list. If `True`, the tools list will be cached and only fetched from the server once. If `False`, the tools list will be fetched from the server on each call to `list_tools()`. The cache can be invalidated by calling `invalidate_tools_cache()`. You should set this to `True` if you know the server will not change its tools list, because it can drastically improve latency (by avoiding a round-trip to the server every time). client_session_timeout_seconds: the read timeout passed to the MCP ClientSession. tool_filter: The tool filter to use for filtering tools. use_structured_content: Whether to use `tool_result.structured_content` when calling an MCP tool. Defaults to False for backwards compatibility - most MCP servers still include the structured content in the `tool_result.content`, and using it by default will cause duplicate content. You can set this to True if you know the server will not duplicate the structured content in the `tool_result.content`. max_retry_attempts: Number of times to retry failed list_tools/call_tool calls. Defaults to no retries. retry_backoff_seconds_base: The base delay, in seconds, used for exponential backoff between retries. message_handler: Optional handler invoked for session messages as delivered by the ClientSession. require_approval: Approval policy for tools on this server. Accepts "always"/"never", a dict of tool names to those values, a boolean, or an object with always/never tool lists. failure_error_function: Optional function used to convert MCP tool failures into a model-visible error message. If explicitly set to None, tool errors will be raised instead of converted. If left unset, the agent-level configuration (or SDK default) will be used. tool_meta_resolver: Optional callable that produces MCP request metadata (`_meta`) for tool calls. It is invoked by the Agents SDK before calling `call_tool`. """ super().__init__( use_structured_content=use_structured_content, require_approval=require_approval, failure_error_function=failure_error_function, tool_meta_resolver=tool_meta_resolver, ) self.session: ClientSession | None = None self.exit_stack: AsyncExitStack = AsyncExitStack() self._cleanup_lock: asyncio.Lock = asyncio.Lock() self._request_lock: asyncio.Lock = asyncio.Lock() self.cache_tools_list = cache_tools_list self.server_initialize_result: InitializeResult | None = None self.client_session_timeout_seconds = client_session_timeout_seconds self.max_retry_attempts = max_retry_attempts self.retry_backoff_seconds_base = retry_backoff_seconds_base self.message_handler = message_handler # The cache is always dirty at startup, so that we fetch tools at least once self._cache_dirty = True self._tools_list: list[MCPTool] | None = None self.tool_filter = tool_filter self._serialize_session_requests = False self._get_session_id: GetSessionIdCallback | None = None async def _maybe_serialize_request(self, func: Callable[[], Awaitable[T]]) -> T: if not self._serialize_session_requests: return await func() async with self._request_lock: return await func() async def _apply_tool_filter( self, tools: list[MCPTool], run_context: RunContextWrapper[Any] | None = None, agent: AgentBase | None = None, ) -> list[MCPTool]: """Apply the tool filter to the list of tools.""" if self.tool_filter is None: return tools # Handle static tool filter if isinstance(self.tool_filter, dict): return self._apply_static_tool_filter(tools, self.tool_filter) # Handle callable tool filter (dynamic filter) else: if run_context is None or agent is None: raise UserError("run_context and agent are required for dynamic tool filtering") return await self._apply_dynamic_tool_filter(tools, run_context, agent) def _apply_static_tool_filter( self, tools: list[MCPTool], static_filter: ToolFilterStatic ) -> list[MCPTool]: """Apply static tool filtering based on allowlist and blocklist.""" filtered_tools = tools # Apply allowed_tool_names filter (whitelist) if "allowed_tool_names" in static_filter: allowed_names = static_filter["allowed_tool_names"] filtered_tools = [t for t in filtered_tools if t.name in allowed_names] # Apply blocked_tool_names filter (blacklist) if "blocked_tool_names" in static_filter: blocked_names = static_filter["blocked_tool_names"] filtered_tools = [t for t in filtered_tools if t.name not in blocked_names] return filtered_tools async def _apply_dynamic_tool_filter( self, tools: list[MCPTool], run_context: RunContextWrapper[Any], agent: AgentBase, ) -> list[MCPTool]: """Apply dynamic tool filtering using a callable filter function.""" # Ensure we have a callable filter if not callable(self.tool_filter): raise ValueError("Tool filter must be callable for dynamic filtering") tool_filter_func = self.tool_filter # Create filter context filter_context = ToolFilterContext( run_context=run_context, agent=agent, server_name=self.name, ) filtered_tools = [] for tool in tools: try: # Call the filter function with context result = tool_filter_func(filter_context, tool) if inspect.isawaitable(result): should_include = await result else: should_include = result if should_include: filtered_tools.append(tool) except Exception as e: logger.error( f"Error applying tool filter to tool '{tool.name}' on server '{self.name}': {e}" ) # On error, exclude the tool for safety continue return filtered_tools @abc.abstractmethod def create_streams( self, ) -> AbstractAsyncContextManager[MCPStreamTransport]: """Create the streams for the server.""" pass async def __aenter__(self): await self.connect() return self async def __aexit__(self, exc_type, exc_value, traceback): await self.cleanup() def invalidate_tools_cache(self): """Invalidate the tools cache.""" self._cache_dirty = True def _extract_http_error_from_exception(self, e: BaseException) -> Exception | None: """Extract HTTP error from exception or ExceptionGroup.""" if isinstance(e, httpx.HTTPStatusError | httpx.ConnectError | httpx.TimeoutException): return e # Check if it's an ExceptionGroup containing HTTP errors if isinstance(e, BaseExceptionGroup): for exc in e.exceptions: if isinstance( exc, httpx.HTTPStatusError | httpx.ConnectError | httpx.TimeoutException ): return exc return None def _raise_user_error_for_http_error(self, http_error: Exception) -> None: """Raise appropriate UserError for HTTP error.""" error_message = f"Failed to connect to MCP server '{self.name}': " if isinstance(http_error, httpx.HTTPStatusError): error_message += f"HTTP error {http_error.response.status_code} ({http_error.response.reason_phrase})" # noqa: E501 elif isinstance(http_error, httpx.ConnectError): error_message += "Could not reach the server." elif isinstance(http_error, httpx.TimeoutException): error_message += "Connection timeout." raise UserError(error_message) from http_error async def _run_with_retries(self, func: Callable[[], Awaitable[T]]) -> T: attempts = 0 while True: try: return await func() except Exception: attempts += 1 if self.max_retry_attempts != -1 and attempts > self.max_retry_attempts: raise backoff = self.retry_backoff_seconds_base * (2 ** (attempts - 1)) await asyncio.sleep(backoff) async def connect(self): """Connect to the server.""" connection_succeeded = False try: transport = await self.exit_stack.enter_async_context(self.create_streams()) # streamablehttp_client returns (read, write, get_session_id) # sse_client returns (read, write) read, write, *rest = transport # Capture the session-id callback when present (streamablehttp_client only). self._get_session_id = rest[0] if rest and callable(rest[0]) else None session = await self.exit_stack.enter_async_context( ClientSession( read, write, timedelta(seconds=self.client_session_timeout_seconds) if self.client_session_timeout_seconds else None, message_handler=self.message_handler, ) ) server_result = await session.initialize() self.server_initialize_result = server_result self.session = session connection_succeeded = True except Exception as e: # Try to extract HTTP error from exception or ExceptionGroup http_error = self._extract_http_error_from_exception(e) if http_error: self._raise_user_error_for_http_error(http_error) # For CancelledError, preserve cancellation semantics - don't wrap it. # If it's masking an HTTP error, cleanup() will extract and raise UserError. if isinstance(e, asyncio.CancelledError): raise # For HTTP-related errors, wrap them if isinstance(e, httpx.HTTPStatusError | httpx.ConnectError | httpx.TimeoutException): self._raise_user_error_for_http_error(e) # For other errors, re-raise as-is (don't wrap non-HTTP errors) raise finally: # Always attempt cleanup on error, but suppress cleanup errors that mask the original if not connection_succeeded: try: await self.cleanup() except UserError: # Re-raise UserError from cleanup (contains the real HTTP error) raise except Exception as cleanup_error: # Suppress RuntimeError about cancel scopes during cleanup - this is a known # issue with the MCP library's async generator cleanup and shouldn't mask the # original error if isinstance(cleanup_error, RuntimeError) and "cancel scope" in str( cleanup_error ): logger.debug( f"Ignoring cancel scope error during cleanup of MCP server " f"'{self.name}': {cleanup_error}" ) else: # Log other cleanup errors but don't raise - original error is more # important logger.warning( f"Error during cleanup of MCP server '{self.name}': {cleanup_error}" ) async def list_tools( self, run_context: RunContextWrapper[Any] | None = None, agent: AgentBase | None = None, ) -> list[MCPTool]: """List the tools available on the server.""" if not self.session: raise UserError("Server not initialized. Make sure you call `connect()` first.") session = self.session assert session is not None try: # Return from cache if caching is enabled, we have tools, and the cache is not dirty if self.cache_tools_list and not self._cache_dirty and self._tools_list: tools = self._tools_list else: # Fetch the tools from the server result = await self._run_with_retries( lambda: self._maybe_serialize_request(lambda: session.list_tools()) ) self._tools_list = result.tools self._cache_dirty = False tools = self._tools_list # Filter tools based on tool_filter filtered_tools = tools if self.tool_filter is not None: filtered_tools = await self._apply_tool_filter(filtered_tools, run_context, agent) return filtered_tools except httpx.HTTPStatusError as e: status_code = e.response.status_code raise UserError( f"Failed to list tools from MCP server '{self.name}': HTTP error {status_code}" ) from e except httpx.ConnectError as e: raise UserError( f"Failed to list tools from MCP server '{self.name}': Connection lost. " f"The server may have disconnected." ) from e async def call_tool( self, tool_name: str, arguments: dict[str, Any] | None, meta: dict[str, Any] | None = None, ) -> CallToolResult: """Invoke a tool on the server.""" if not self.session: raise UserError("Server not initialized. Make sure you call `connect()` first.") session = self.session assert session is not None try: self._validate_required_parameters(tool_name=tool_name, arguments=arguments) if meta is None: return await self._run_with_retries( lambda: self._maybe_serialize_request( lambda: session.call_tool(tool_name, arguments) ) ) return await self._run_with_retries( lambda: self._maybe_serialize_request( lambda: session.call_tool(tool_name, arguments, meta=meta) ) ) except httpx.HTTPStatusError as e: status_code = e.response.status_code raise UserError( f"Failed to call tool '{tool_name}' on MCP server '{self.name}': " f"HTTP error {status_code}" ) from e except httpx.ConnectError as e: raise UserError( f"Failed to call tool '{tool_name}' on MCP server '{self.name}': Connection lost. " f"The server may have disconnected." ) from e def _validate_required_parameters( self, tool_name: str, arguments: dict[str, Any] | None ) -> None: """Validate required tool parameters from cached MCP tool schemas before invocation.""" if self._tools_list is None: return tool = next((item for item in self._tools_list if item.name == tool_name), None) if tool is None or not isinstance(tool.inputSchema, dict): return raw_required = tool.inputSchema.get("required") if not isinstance(raw_required, list) or not raw_required: return if arguments is None: arguments_to_validate: dict[str, Any] = {} elif isinstance(arguments, dict): arguments_to_validate = arguments else: raise UserError( f"Failed to call tool '{tool_name}' on MCP server '{self.name}': " "arguments must be an object." ) required_names = [name for name in raw_required if isinstance(name, str)] missing = [name for name in required_names if name not in arguments_to_validate] if missing: missing_text = ", ".join(sorted(missing)) raise UserError( f"Failed to call tool '{tool_name}' on MCP server '{self.name}': " f"missing required parameters: {missing_text}" ) async def list_prompts( self, ) -> ListPromptsResult: """List the prompts available on the server.""" if not self.session: raise UserError("Server not initialized. Make sure you call `connect()` first.") session = self.session assert session is not None return await self._maybe_serialize_request(lambda: session.list_prompts()) async def get_prompt( self, name: str, arguments: dict[str, Any] | None = None ) -> GetPromptResult: """Get a specific prompt from the server.""" if not self.session: raise UserError("Server not initialized. Make sure you call `connect()` first.") session = self.session assert session is not None return await self._maybe_serialize_request(lambda: session.get_prompt(name, arguments)) async def list_resources(self, cursor: str | None = None) -> ListResourcesResult: """List the resources available on the server.""" if not self.session: raise UserError("Server not initialized. Make sure you call `connect()` first.") session = self.session assert session is not None return await self._maybe_serialize_request(lambda: session.list_resources(cursor)) async def list_resource_templates( self, cursor: str | None = None ) -> ListResourceTemplatesResult: """List the resource templates available on the server.""" if not self.session: raise UserError("Server not initialized. Make sure you call `connect()` first.") session = self.session assert session is not None return await self._maybe_serialize_request(lambda: session.list_resource_templates(cursor)) async def read_resource(self, uri: str) -> ReadResourceResult: """Read the contents of a specific resource by URI. Args: uri: The URI of the resource to read. See :class:`~pydantic.networks.AnyUrl` for the supported URI formats. """ if not self.session: raise UserError("Server not initialized. Make sure you call `connect()` first.") session = self.session assert session is not None from pydantic import AnyUrl return await self._maybe_serialize_request(lambda: session.read_resource(AnyUrl(uri))) async def cleanup(self): """Cleanup the server.""" async with self._cleanup_lock: # Only raise HTTP errors if we're cleaning up after a failed connection. # During normal teardown (via __aexit__), log but don't raise to avoid # masking the original exception. is_failed_connection_cleanup = self.session is None try: await self.exit_stack.aclose() except asyncio.CancelledError as e: logger.debug(f"Cleanup cancelled for MCP server '{self.name}': {e}") raise except BaseExceptionGroup as eg: # Extract HTTP errors from ExceptionGroup raised during cleanup # This happens when background tasks fail (e.g., HTTP errors) http_error = None connect_error = None timeout_error = None error_message = f"Failed to connect to MCP server '{self.name}': " for exc in eg.exceptions: if isinstance(exc, httpx.HTTPStatusError): http_error = exc elif isinstance(exc, httpx.ConnectError): connect_error = exc elif isinstance(exc, httpx.TimeoutException): timeout_error = exc # Only raise HTTP errors if we're cleaning up after a failed connection. # During normal teardown, log them instead. if http_error: if is_failed_connection_cleanup: error_message += f"HTTP error {http_error.response.status_code} ({http_error.response.reason_phrase})" # noqa: E501 raise UserError(error_message) from http_error else: # Normal teardown - log but don't raise logger.warning( f"HTTP error during cleanup of MCP server '{self.name}': {http_error}" ) elif connect_error: if is_failed_connection_cleanup: error_message += "Could not reach the server." raise UserError(error_message) from connect_error else: logger.warning( f"Connection error during cleanup of MCP server '{self.name}': {connect_error}" # noqa: E501 ) elif timeout_error: if is_failed_connection_cleanup: error_message += "Connection timeout." raise UserError(error_message) from timeout_error else: logger.warning( f"Timeout error during cleanup of MCP server '{self.name}': {timeout_error}" # noqa: E501 ) else: # No HTTP error found, suppress RuntimeError about cancel scopes has_cancel_scope_error = any( isinstance(exc, RuntimeError) and "cancel scope" in str(exc) for exc in eg.exceptions ) if has_cancel_scope_error: logger.debug(f"Ignoring cancel scope error during cleanup: {eg}") else: logger.error(f"Error cleaning up server: {eg}") except Exception as e: # Suppress RuntimeError about cancel scopes - this is a known issue with the MCP # library when background tasks fail during async generator cleanup if isinstance(e, RuntimeError) and "cancel scope" in str(e): logger.debug(f"Ignoring cancel scope error during cleanup: {e}") else: logger.error(f"Error cleaning up server: {e}") finally: self.session = None self._get_session_id = None class MCPServerStdioParams(TypedDict): """Mirrors `mcp.client.stdio.StdioServerParameters`, but lets you pass params without another import. """ command: str """The executable to run to start the server. For example, `python` or `node`.""" args: NotRequired[list[str]] """Command line args to pass to the `command` executable. For example, `['foo.py']` or `['server.js', '--port', '8080']`.""" env: NotRequired[dict[str, str]] """The environment variables to set for the server. .""" cwd: NotRequired[str | Path] """The working directory to use when spawning the process.""" encoding: NotRequired[str] """The text encoding used when sending/receiving messages to the server. Defaults to `utf-8`.""" encoding_error_handler: NotRequired[Literal["strict", "ignore", "replace"]] """The text encoding error handler. Defaults to `strict`. See https://docs.python.org/3/library/codecs.html#codec-base-classes for explanations of possible values. """ class MCPServerStdio(_MCPServerWithClientSession): """MCP server implementation that uses the stdio transport. See the [spec] (https://spec.modelcontextprotocol.io/specification/2024-11-05/basic/transports/#stdio) for details. """ def __init__( self, params: MCPServerStdioParams, cache_tools_list: bool = False, name: str | None = None, client_session_timeout_seconds: float | None = 5, tool_filter: ToolFilter = None, use_structured_content: bool = False, max_retry_attempts: int = 0, retry_backoff_seconds_base: float = 1.0, message_handler: MessageHandlerFnT | None = None, require_approval: RequireApprovalSetting = None, failure_error_function: ToolErrorFunction | None | _UnsetType = _UNSET, tool_meta_resolver: MCPToolMetaResolver | None = None, ): """Create a new MCP server based on the stdio transport. Args: params: The params that configure the server. This includes the command to run to start the server, the args to pass to the command, the environment variables to set for the server, the working directory to use when spawning the process, and the text encoding used when sending/receiving messages to the server. cache_tools_list: Whether to cache the tools list. If `True`, the tools list will be cached and only fetched from the server once. If `False`, the tools list will be fetched from the server on each call to `list_tools()`. The cache can be invalidated by calling `invalidate_tools_cache()`. You should set this to `True` if you know the server will not change its tools list, because it can drastically improve latency (by avoiding a round-trip to the server every time). name: A readable name for the server. If not provided, we'll create one from the command. client_session_timeout_seconds: the read timeout passed to the MCP ClientSession. tool_filter: The tool filter to use for filtering tools. use_structured_content: Whether to use `tool_result.structured_content` when calling an MCP tool. Defaults to False for backwards compatibility - most MCP servers still include the structured content in the `tool_result.content`, and using it by default will cause duplicate content. You can set this to True if you know the server will not duplicate the structured content in the `tool_result.content`. max_retry_attempts: Number of times to retry failed list_tools/call_tool calls. Defaults to no retries. retry_backoff_seconds_base: The base delay, in seconds, for exponential backoff between retries. message_handler: Optional handler invoked for session messages as delivered by the ClientSession. require_approval: Approval policy for tools on this server. Accepts "always"/"never", a dict of tool names to those values, or an object with always/never tool lists. failure_error_function: Optional function used to convert MCP tool failures into a model-visible error message. If explicitly set to None, tool errors will be raised instead of converted. If left unset, the agent-level configuration (or SDK default) will be used. tool_meta_resolver: Optional callable that produces MCP request metadata (`_meta`) for tool calls. It is invoked by the Agents SDK before calling `call_tool`. """ super().__init__( cache_tools_list=cache_tools_list, client_session_timeout_seconds=client_session_timeout_seconds, tool_filter=tool_filter, use_structured_content=use_structured_content, max_retry_attempts=max_retry_attempts, retry_backoff_seconds_base=retry_backoff_seconds_base, message_handler=message_handler, require_approval=require_approval, failure_error_function=failure_error_function, tool_meta_resolver=tool_meta_resolver, ) self.params = StdioServerParameters( command=params["command"], args=params.get("args", []), env=params.get("env"), cwd=params.get("cwd"), encoding=params.get("encoding", "utf-8"), encoding_error_handler=params.get("encoding_error_handler", "strict"), ) self._name = name or f"stdio: {self.params.command}" def create_streams( self, ) -> AbstractAsyncContextManager[MCPStreamTransport]: """Create the streams for the server.""" return stdio_client(self.params) @property def name(self) -> str: """A readable name for the server.""" return self._name class MCPServerSseParams(TypedDict): """Mirrors the params in`mcp.client.sse.sse_client`.""" url: str """The URL of the server.""" headers: NotRequired[dict[str, str]] """The headers to send to the server.""" timeout: NotRequired[float] """The timeout for the HTTP request. Defaults to 5 seconds.""" sse_read_timeout: NotRequired[float] """The timeout for the SSE connection, in seconds. Defaults to 5 minutes.""" auth: NotRequired[httpx.Auth | None] """Optional httpx authentication handler (e.g. ``httpx.BasicAuth``, a custom ``httpx.Auth`` subclass for OAuth token refresh, etc.). When provided, it is passed directly to the underlying ``httpx.AsyncClient`` used by the SSE transport. """ httpx_client_factory: NotRequired[HttpClientFactory] """Custom HTTP client factory for configuring httpx.AsyncClient behavior (e.g. to set custom SSL certificates, proxies, or other transport options). """ class MCPServerSse(_MCPServerWithClientSession): """MCP server implementation that uses the HTTP with SSE transport. See the [spec] (https://spec.modelcontextprotocol.io/specification/2024-11-05/basic/transports/#http-with-sse) for details. """ def __init__( self, params: MCPServerSseParams, cache_tools_list: bool = False, name: str | None = None, client_session_timeout_seconds: float | None = 5, tool_filter: ToolFilter = None, use_structured_content: bool = False, max_retry_attempts: int = 0, retry_backoff_seconds_base: float = 1.0, message_handler: MessageHandlerFnT | None = None, require_approval: RequireApprovalSetting = None, failure_error_function: ToolErrorFunction | None | _UnsetType = _UNSET, tool_meta_resolver: MCPToolMetaResolver | None = None, ): """Create a new MCP server based on the HTTP with SSE transport. Args: params: The params that configure the server. This includes the URL of the server, the headers to send to the server, the timeout for the HTTP request, and the timeout for the SSE connection. cache_tools_list: Whether to cache the tools list. If `True`, the tools list will be cached and only fetched from the server once. If `False`, the tools list will be fetched from the server on each call to `list_tools()`. The cache can be invalidated by calling `invalidate_tools_cache()`. You should set this to `True` if you know the server will not change its tools list, because it can drastically improve latency (by avoiding a round-trip to the server every time). name: A readable name for the server. If not provided, we'll create one from the URL. client_session_timeout_seconds: the read timeout passed to the MCP ClientSession. tool_filter: The tool filter to use for filtering tools. use_structured_content: Whether to use `tool_result.structured_content` when calling an MCP tool. Defaults to False for backwards compatibility - most MCP servers still include the structured content in the `tool_result.content`, and using it by default will cause duplicate content. You can set this to True if you know the server will not duplicate the structured content in the `tool_result.content`. max_retry_attempts: Number of times to retry failed list_tools/call_tool calls. Defaults to no retries. retry_backoff_seconds_base: The base delay, in seconds, for exponential backoff between retries. message_handler: Optional handler invoked for session messages as delivered by the ClientSession. require_approval: Approval policy for tools on this server. Accepts "always"/"never", a dict of tool names to those values, or an object with always/never tool lists. failure_error_function: Optional function used to convert MCP tool failures into a model-visible error message. If explicitly set to None, tool errors will be raised instead of converted. If left unset, the agent-level configuration (or SDK default) will be used. tool_meta_resolver: Optional callable that produces MCP request metadata (`_meta`) for tool calls. It is invoked by the Agents SDK before calling `call_tool`. """ super().__init__( cache_tools_list=cache_tools_list, client_session_timeout_seconds=client_session_timeout_seconds, tool_filter=tool_filter, use_structured_content=use_structured_content, max_retry_attempts=max_retry_attempts, retry_backoff_seconds_base=retry_backoff_seconds_base, message_handler=message_handler, require_approval=require_approval, failure_error_function=failure_error_function, tool_meta_resolver=tool_meta_resolver, ) self.params = params self._name = name or f"sse: {self.params['url']}" def create_streams( self, ) -> AbstractAsyncContextManager[MCPStreamTransport]: """Create the streams for the server.""" kwargs: dict[str, Any] = { "url": self.params["url"], "headers": self.params.get("headers", None), "timeout": self.params.get("timeout", 5), "sse_read_timeout": self.params.get("sse_read_timeout", 60 * 5), } if "auth" in self.params: kwargs["auth"] = self.params["auth"] if "httpx_client_factory" in self.params: kwargs["httpx_client_factory"] = self.params["httpx_client_factory"] return sse_client(**kwargs) @property def name(self) -> str: """A readable name for the server.""" return self._name class MCPServerStreamableHttpParams(TypedDict): """Mirrors the params in`mcp.client.streamable_http.streamablehttp_client`.""" url: str """The URL of the server.""" headers: NotRequired[dict[str, str]] """The headers to send to the server.""" timeout: NotRequired[timedelta | float] """The timeout for the HTTP request. Defaults to 5 seconds.""" sse_read_timeout: NotRequired[timedelta | float] """The timeout for the SSE connection, in seconds. Defaults to 5 minutes.""" terminate_on_close: NotRequired[bool] """Terminate on close""" httpx_client_factory: NotRequired[HttpClientFactory] """Custom HTTP client factory for configuring httpx.AsyncClient behavior.""" auth: NotRequired[httpx.Auth | None] """Optional httpx authentication handler (e.g. ``httpx.BasicAuth``, a custom ``httpx.Auth`` subclass for OAuth token refresh, etc.). When provided, it is passed directly to the underlying ``httpx.AsyncClient`` used by the Streamable HTTP transport. """ ignore_initialized_notification_failure: NotRequired[bool] """Whether to ignore failures when sending the best-effort ``notifications/initialized`` POST. Defaults to ``False``. When set to ``True``, initialized-notification failures are logged and ignored so subsequent requests on the same transport can continue. """ class MCPServerStreamableHttp(_MCPServerWithClientSession): """MCP server implementation that uses the Streamable HTTP transport. See the [spec] (https://modelcontextprotocol.io/specification/2025-03-26/basic/transports#streamable-http) for details. """ def __init__( self, params: MCPServerStreamableHttpParams, cache_tools_list: bool = False, name: str | None = None, client_session_timeout_seconds: float | None = 5, tool_filter: ToolFilter = None, use_structured_content: bool = False, max_retry_attempts: int = 0, retry_backoff_seconds_base: float = 1.0, message_handler: MessageHandlerFnT | None = None, require_approval: RequireApprovalSetting = None, failure_error_function: ToolErrorFunction | None | _UnsetType = _UNSET, tool_meta_resolver: MCPToolMetaResolver | None = None, ): """Create a new MCP server based on the Streamable HTTP transport. Args: params: The params that configure the server. This includes the URL of the server, the headers to send to the server, the timeout for the HTTP request, the timeout for the Streamable HTTP connection, whether we need to terminate on close, and an optional custom HTTP client factory. cache_tools_list: Whether to cache the tools list. If `True`, the tools list will be cached and only fetched from the server once. If `False`, the tools list will be fetched from the server on each call to `list_tools()`. The cache can be invalidated by calling `invalidate_tools_cache()`. You should set this to `True` if you know the server will not change its tools list, because it can drastically improve latency (by avoiding a round-trip to the server every time). name: A readable name for the server. If not provided, we'll create one from the URL. client_session_timeout_seconds: the read timeout passed to the MCP ClientSession. tool_filter: The tool filter to use for filtering tools. use_structured_content: Whether to use `tool_result.structured_content` when calling an MCP tool. Defaults to False for backwards compatibility - most MCP servers still include the structured content in the `tool_result.content`, and using it by default will cause duplicate content. You can set this to True if you know the server will not duplicate the structured content in the `tool_result.content`. max_retry_attempts: Number of times to retry failed list_tools/call_tool calls. Defaults to no retries. retry_backoff_seconds_base: The base delay, in seconds, for exponential backoff between retries. message_handler: Optional handler invoked for session messages as delivered by the ClientSession. require_approval: Approval policy for tools on this server. Accepts "always"/"never", a dict of tool names to those values, or an object with always/never tool lists. failure_error_function: Optional function used to convert MCP tool failures into a model-visible error message. If explicitly set to None, tool errors will be raised instead of converted. If left unset, the agent-level configuration (or SDK default) will be used. tool_meta_resolver: Optional callable that produces MCP request metadata (`_meta`) for tool calls. It is invoked by the Agents SDK before calling `call_tool`. """ super().__init__( cache_tools_list=cache_tools_list, client_session_timeout_seconds=client_session_timeout_seconds, tool_filter=tool_filter, use_structured_content=use_structured_content, max_retry_attempts=max_retry_attempts, retry_backoff_seconds_base=retry_backoff_seconds_base, message_handler=message_handler, require_approval=require_approval, failure_error_function=failure_error_function, tool_meta_resolver=tool_meta_resolver, ) self.params = params self._name = name or f"streamable_http: {self.params['url']}" self._serialize_session_requests = True def create_streams( self, ) -> AbstractAsyncContextManager[MCPStreamTransport]: """Create the streams for the server.""" kwargs: dict[str, Any] = { "url": self.params["url"], "headers": self.params.get("headers", None), "timeout": self.params.get("timeout", 5), "sse_read_timeout": self.params.get("sse_read_timeout", 60 * 5), "terminate_on_close": self.params.get("terminate_on_close", True), } httpx_client_factory = self.params.get("httpx_client_factory") if self.params.get("ignore_initialized_notification_failure", False): return _streamablehttp_client_with_transport( **kwargs, httpx_client_factory=httpx_client_factory or _create_default_streamable_http_client, auth=self.params.get("auth"), transport_factory=_InitializedNotificationTolerantStreamableHTTPTransport, ) if httpx_client_factory is not None: kwargs["httpx_client_factory"] = httpx_client_factory if "auth" in self.params: kwargs["auth"] = self.params["auth"] return streamablehttp_client(**kwargs) @asynccontextmanager async def _isolated_client_session(self): async with AsyncExitStack() as exit_stack: transport = await exit_stack.enter_async_context(self.create_streams()) read, write, *_ = transport session = await exit_stack.enter_async_context( ClientSession( read, write, timedelta(seconds=self.client_session_timeout_seconds) if self.client_session_timeout_seconds else None, message_handler=self.message_handler, ) ) await session.initialize() yield session async def _call_tool_with_session( self, session: ClientSession, tool_name: str, arguments: dict[str, Any] | None, meta: dict[str, Any] | None = None, ) -> CallToolResult: if meta is None: return await session.call_tool(tool_name, arguments) return await session.call_tool(tool_name, arguments, meta=meta) def _should_retry_in_isolated_session(self, exc: BaseException) -> bool: if isinstance( exc, asyncio.CancelledError | ClosedResourceError | httpx.ConnectError | httpx.TimeoutException, ): return True if isinstance(exc, httpx.HTTPStatusError): return exc.response.status_code >= 500 if isinstance(exc, McpError): return exc.error.code == httpx.codes.REQUEST_TIMEOUT if isinstance(exc, BaseExceptionGroup): return bool(exc.exceptions) and all( self._should_retry_in_isolated_session(inner) for inner in exc.exceptions ) return False async def _call_tool_with_shared_session( self, tool_name: str, arguments: dict[str, Any] | None, meta: dict[str, Any] | None = None, *, allow_isolated_retry: bool, ) -> CallToolResult: session = self.session assert session is not None try: return await self._maybe_serialize_request( lambda: self._call_tool_with_session(session, tool_name, arguments, meta) ) except BaseException as exc: if allow_isolated_retry and self._should_retry_in_isolated_session(exc): raise _SharedSessionRequestNeedsIsolation from exc raise async def _call_tool_with_isolated_retry( self, tool_name: str, arguments: dict[str, Any] | None, meta: dict[str, Any] | None = None, *, allow_isolated_retry: bool, ) -> tuple[CallToolResult, bool]: request_task = asyncio.create_task( self._call_tool_with_shared_session( tool_name, arguments, meta, allow_isolated_retry=allow_isolated_retry, ) ) try: return await asyncio.shield(request_task), False except _SharedSessionRequestNeedsIsolation: exit_stack = AsyncExitStack() try: session = await exit_stack.enter_async_context(self._isolated_client_session()) except asyncio.CancelledError: await exit_stack.aclose() raise except BaseException as exc: await exit_stack.aclose() raise _IsolatedSessionRetryFailed() from exc try: try: result = await self._call_tool_with_session(session, tool_name, arguments, meta) return result, True except asyncio.CancelledError: raise except BaseException as exc: raise _IsolatedSessionRetryFailed() from exc finally: await exit_stack.aclose() except asyncio.CancelledError: if not request_task.done(): request_task.cancel() try: await request_task except BaseException: pass raise async def call_tool( self, tool_name: str, arguments: dict[str, Any] | None, meta: dict[str, Any] | None = None, ) -> CallToolResult: if not self.session: raise UserError("Server not initialized. Make sure you call `connect()` first.") try: self._validate_required_parameters(tool_name=tool_name, arguments=arguments) retries_used = 0 first_attempt = True while True: if not first_attempt and self.max_retry_attempts != -1: retries_used += 1 allow_isolated_retry = ( self.max_retry_attempts == -1 or retries_used < self.max_retry_attempts ) try: result, used_isolated_retry = await self._call_tool_with_isolated_retry( tool_name, arguments, meta, allow_isolated_retry=allow_isolated_retry, ) if used_isolated_retry and self.max_retry_attempts != -1: retries_used += 1 return result except _IsolatedSessionRetryFailed as exc: retries_used += 1 if self.max_retry_attempts != -1 and retries_used >= self.max_retry_attempts: if exc.__cause__ is not None: raise exc.__cause__ from exc raise exc backoff = self.retry_backoff_seconds_base * (2 ** (retries_used - 1)) await asyncio.sleep(backoff) except Exception: if self.max_retry_attempts != -1 and retries_used >= self.max_retry_attempts: raise backoff = self.retry_backoff_seconds_base * (2**retries_used) await asyncio.sleep(backoff) first_attempt = False except httpx.HTTPStatusError as e: status_code = e.response.status_code raise UserError( f"Failed to call tool '{tool_name}' on MCP server '{self.name}': " f"HTTP error {status_code}" ) from e except httpx.ConnectError as e: raise UserError( f"Failed to call tool '{tool_name}' on MCP server '{self.name}': Connection lost. " f"The server may have disconnected." ) from e except BaseExceptionGroup as e: http_error = self._extract_http_error_from_exception(e) if isinstance(http_error, httpx.HTTPStatusError): status_code = http_error.response.status_code raise UserError( f"Failed to call tool '{tool_name}' on MCP server '{self.name}': " f"HTTP error {status_code}" ) from http_error if isinstance(http_error, httpx.ConnectError): raise UserError( f"Failed to call tool '{tool_name}' on MCP server '{self.name}': " "Connection lost. The server may have disconnected." ) from http_error if isinstance(http_error, httpx.TimeoutException): raise UserError( f"Failed to call tool '{tool_name}' on MCP server '{self.name}': " "Connection timeout." ) from http_error raise @property def name(self) -> str: """A readable name for the server.""" return self._name @property def session_id(self) -> str | None: """The MCP session ID assigned by the server, or None if not yet connected or if the server did not issue a session ID. The session ID is stable for the lifetime of this server instance's connection. You can persist it and pass it back via the Mcp-Session-Id request header (params["headers"]) on a new MCPServerStreamableHttp instance to resume the same server-side session across process restarts or stateless workers. Example:: async with MCPServerStreamableHttp(params={"url": url}) as server: session_id = server.session_id # In a new worker / process: async with MCPServerStreamableHttp( params={"url": url, "headers": {"Mcp-Session-Id": session_id}} ) as server: # Resumes the same server-side session. ... """ if self._get_session_id is None: return None return self._get_session_id()