FEAT: Add MCP-wrapped prompt chat target

pyproject.toml

@@ -131,6 +131,9 @@ speech = [
     "azure-cognitiveservices-speech>=1.44.0",
 ]
 
+mcp = [
+    "mcp>=1.10.0",
+]
 litellm = [
     # Cap below 1.92.0: 1.92.0+ replaced the universal py3-none-any wheel with a
     # Rust (PyO3) extension and only publishes manylinux_2_28 (x86_64/aarch64) and
@@ -153,6 +156,7 @@ all = [
     "ipykernel>=6.29.5",
     "jupyter>=1.1.1",
     "litellm>=1.84.0,<1.99.0",  # 1.92.0+ drops the universal wheel (no mac/musl/old-glibc/win-arm64); see litellm group
+    "mcp>=1.10.0",
     "ollama>=0.5.1",
     "opencv-python>=4.11.0.86",
     "playwright>=1.49.0",

pyrit/mcp/__init__.py

@@ -0,0 +1,10 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+# ruff: noqa: F401
+
+"""Model Context Protocol (MCP) integration for PyRIT targets."""
+
+from pyrit.mcp.mcp_server_config import MCPServerConfig, MCPTransport
+from pyrit.mcp.mcp_wrapped_prompt_chat_target import MCPWrappedPromptChatTarget
+
+__all__ = ["MCPServerConfig", "MCPTransport", "MCPWrappedPromptChatTarget"]

pyrit/mcp/_client.py

@@ -0,0 +1,187 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+"""
+Thin async client over the official ``mcp`` Python SDK.
+
+Owns the transport plumbing (stdio subprocess / Streamable HTTP sessions) and
+exposes just the two operations the wrapped target needs: listing tools and
+calling one. The ``mcp`` package is an optional dependency (``pyrit[mcp]``) and
+is imported lazily with an actionable error message.
+"""
+
+from __future__ import annotations
+
+import logging
+from contextlib import AsyncExitStack
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+from pyrit.mcp.mcp_server_config import MCPServerConfig, MCPTransport
+
+if TYPE_CHECKING:
+    from types import TracebackType
+
+logger = logging.getLogger(__name__)
+
+
+def _import_mcp() -> Any:
+    """
+    Import the optional ``mcp`` package.
+
+    Returns:
+        The imported module.
+
+    Raises:
+        ModuleNotFoundError: With an actionable install hint if the extra is missing.
+    """
+    try:
+        import mcp  # noqa: PLC0415
+    except ModuleNotFoundError as exc:
+        raise ModuleNotFoundError(
+            "The 'mcp' package is required for MCP support. Install it with: pip install pyrit[mcp]"
+        ) from exc
+    return mcp
+
+
+@dataclass(frozen=True)
+class MCPTool:
+    """A tool declared by an MCP server (catalog entry)."""
+
+    name: str
+    description: str
+    input_schema: dict[str, Any]
+
+
+@dataclass(frozen=True)
+class MCPToolResult:
+    """Result of an MCP tool call, normalized to text."""
+
+    text: str
+    is_error: bool
+
+
+class MCPClientSession:
+    """
+    A live session with a single MCP server.
+
+    Use as an async context manager; entering connects (launching the subprocess
+    for stdio servers) and exiting tears the connection down. Not safe for
+    concurrent use by multiple tasks — the wrapped target serializes access.
+    """
+
+    def __init__(self, *, config: MCPServerConfig) -> None:
+        """
+        Create (but do not start) a session for the given server.
+
+        Args:
+            config (MCPServerConfig): Connection and policy configuration.
+        """
+        self._config = config
+        self._exit_stack: AsyncExitStack | None = None
+        self._session: Any = None  # mcp.ClientSession, untyped until mcp is imported
+
+    @property
+    def server_name(self) -> str:
+        """The friendly server name from the config."""
+        return self._config.name
+
+    async def __aenter__(self) -> MCPClientSession:
+        mcp = _import_mcp()
+        config = self._config
+        self._exit_stack = AsyncExitStack()
+
+        try:
+            if config.transport is MCPTransport.STDIO:
+                stdio_client = mcp.client.stdio.stdio_client
+                server_params = mcp.client.stdio.StdioServerParameters(
+                    command=config.command,
+                    args=list(config.args),
+                    env=dict(config.env) if config.env is not None else None,
+                )
+                read_stream, write_stream = await self._exit_stack.enter_async_context(stdio_client(server_params))
+            else:
+                streamablehttp_client = mcp.client.streamable_http.streamablehttp_client
+                http_transport = await self._exit_stack.enter_async_context(
+                    streamablehttp_client(url=config.url, headers=dict(config.headers) if config.headers else {})
+                )
+                read_stream, write_stream, _get_session_id = http_transport
+
+            client_session = mcp.client.session.ClientSession(read_stream, write_stream)
+            self._session = await self._exit_stack.enter_async_context(client_session)
+            await self._session.initialize()
+        except BaseException:
+            await self._exit_stack.aclose()
+            self._exit_stack = None
+            raise
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        if self._exit_stack is not None:
+            await self._exit_stack.aclose()
+            self._exit_stack = None
+        self._session = None
+
+    def _require_session(self) -> Any:
+        if self._session is None:
+            raise RuntimeError(
+                f"MCP session for server '{self._config.name}' is not connected; use 'async with' to connect."
+            )
+        return self._session
+
+    async def list_tools(self) -> list[MCPTool]:
+        """
+        List the tools the server declares.
+
+        Returns:
+            list[MCPTool]: The server's tool catalog.
+
+        Raises:
+            RuntimeError: If the session is not connected.
+            Exception: Propagates SDK/transport failures after normalizing
+                ``result.isError``-style errors (listing has no error channel,
+                so transport failures propagate).
+        """
+        session = self._require_session()
+        response = await session.list_tools()
+        return [
+            MCPTool(
+                name=tool.name,
+                description=tool.description or "",
+                input_schema=dict(tool.inputSchema) if tool.inputSchema else {},
+            )
+            for tool in response.tools
+        ]
+
+    async def call_tool(self, *, tool_name: str, arguments: dict[str, Any]) -> MCPToolResult:
+        """
+        Call a tool on the server.
+
+        Args:
+            tool_name (str): The tool's registered name.
+            arguments (dict[str, Any]): JSON tool arguments.
+
+        Returns:
+            MCPToolResult: Normalized text result; ``is_error`` is True when the
+                server reported a tool-level error.
+
+        Raises:
+            RuntimeError: If the session is not connected.
+            Exception: Transport failures (timeouts, connection loss) propagate;
+                tool-level errors do not (they are reported via ``is_error``).
+        """
+        session = self._require_session()
+        response = await session.call_tool(name=tool_name, arguments=arguments)
+
+        parts: list[str] = []
+        for content in response.content or []:
+            text = getattr(content, "text", None)
+            if text is not None:
+                parts.append(text)
+        text = "\n".join(parts) if parts else "(no content returned)"
+        return MCPToolResult(text=text, is_error=bool(response.isError))

pyrit/mcp/mcp_server_config.py

@@ -0,0 +1,94 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+"""
+Configuration models for connecting PyRIT to MCP servers.
+
+These models describe *how to reach* an MCP server (transport and connection
+parameters) plus the red-team policy applied to its tools (allowlist, caps).
+They are transport-agnostic inputs to :mod:`pyrit.mcp._client`.
+"""
+
+from __future__ import annotations
+
+from enum import Enum
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class MCPTransport(str, Enum):
+    """
+    Transports supported by :class:`MCPServerConfig`.
+
+    - ``stdio``: launch the server as a subprocess and speak MCP over stdin/stdout.
+      Deterministic and offline-testable; the recommended transport for local tools.
+    - ``streamable_http``: connect to a remote MCP server over the Streamable HTTP
+      transport (the successor of the deprecated HTTP+SSE transport in the MCP spec).
+    """
+
+    STDIO = "stdio"
+    STREAMABLE_HTTP = "streamable_http"
+
+
+class MCPServerConfig(BaseModel):
+    """
+    Connection and policy configuration for a single MCP server.
+
+    Exactly one connection style must be provided per transport: ``command`` for
+    stdio servers, ``url`` for streamable-HTTP servers.
+
+    Args:
+        name (str): Friendly server name used in tool catalog entries, audit logs,
+            and error messages.
+        transport (MCPTransport): Transport used to reach the server.
+        command (str | None): Executable that starts the MCP server (stdio only),
+            e.g. ``"python"``.
+        args (list[str] | None): Arguments forwarded to ``command`` (stdio only).
+        env (dict[str, str] | None): Extra environment variables for the server
+            subprocess (stdio only). Defaults to a minimal environment when unset.
+        url (str | None): Server endpoint URL (streamable-http only).
+        headers (dict[str, str] | None): Extra HTTP headers for the endpoint
+            (streamable-http only), e.g. authorization headers.
+        allowed_tools (list[str] | None): Explicit tool-name allowlist for this
+            server. ``None`` (default) allows every tool the server declares;
+            names listed here are the only ones the wrapped target may execute.
+        tool_call_timeout (float): Per-tool-call timeout in seconds. Defaults to 60.
+    """
+
+    model_config = ConfigDict(extra="forbid")
+
+    name: str = Field(min_length=1, description="Friendly server name used in logs and the tool catalog.")
+    transport: MCPTransport = Field(description="Transport used to reach the MCP server.")
+    command: str | None = Field(None, description="Executable that starts the MCP server (stdio only).")
+    args: list[str] = Field(default_factory=list, description="Arguments for the stdio server command.")
+    env: dict[str, str] | None = Field(None, description="Extra environment variables for the stdio subprocess.")
+    url: str | None = Field(None, description="Server endpoint URL (streamable-http only).")
+    headers: dict[str, str] | None = Field(None, description="Extra HTTP headers (streamable-http only).")
+    allowed_tools: list[str] | None = Field(
+        None,
+        description="Explicit tool-name allowlist for this server; None allows all tools the server declares.",
+    )
+    tool_call_timeout: float = Field(
+        60.0,
+        gt=0,
+        description="Per-tool-call timeout in seconds.",
+    )
+
+    def validate_connection(self) -> None:
+        """
+        Verify the connection parameters match the selected transport.
+
+        Raises:
+            ValueError: If required connection parameters are missing or a
+                parameter is provided that does not belong to the transport.
+        """
+        if self.transport is MCPTransport.STDIO:
+            if not self.command:
+                raise ValueError(f"MCP server '{self.name}': stdio transport requires 'command'.")
+            if self.url:
+                raise ValueError(f"MCP server '{self.name}': 'url' is only valid for the streamable_http transport.")
+        elif self.transport is MCPTransport.STREAMABLE_HTTP:
+            if not self.url:
+                raise ValueError(f"MCP server '{self.name}': streamable_http transport requires 'url'.")
+            if self.command:
+                raise ValueError(f"MCP server '{self.name}': 'command' is only valid for the stdio transport.")