[DRAFT] FEAT: Tool Use + MCP

pyproject.toml

@@ -45,6 +45,7 @@ dependencies = [
     "fastapi>=0.115.0",
     "httpx[http2]>=0.27.2",
     "jinja2>=3.1.6",
+    "mcp>=1.0,<2",
     "numpy>=1.26.0; python_version < '3.14'",
     "numpy>=2.3.0; python_version >= '3.14'",
     "openai>=2.2.0",

pyrit/exceptions/__init__.py

@@ -10,6 +10,8 @@
     MissingPromptPlaceholderException,
     PyritException,
     RateLimitException,
+    ToolCallLoopLimitExceeded,
+    ToolCallNotSupported,
     get_retry_max_num_attempts,
     handle_bad_request_exception,
     pyrit_custom_result_retry,
@@ -59,4 +61,6 @@
     "set_execution_context",
     "set_retry_collector",
     "execution_context",
+    "ToolCallLoopLimitExceeded",
+    "ToolCallNotSupported",
 ]

pyrit/exceptions/exception_classes.py

@@ -233,6 +233,70 @@ def __init__(self, *, message: str = "No prompt placeholder") -> None:
         super().__init__(message=message)
 
 
+class ToolCallNotSupported(PyritException):
+    """
+    Raised when a target produces a tool call that the configured
+    :class:`~pyrit.tools.ToolEventPolicy` does not permit to execute
+    (``ToolEventBehavior.RAISE``, or ``EXECUTE`` without a backend).
+
+    The ``partial_conversation`` attribute carries every message produced
+    up to and including the assistant turn that contained the offending
+    tool call(s). Consumers can inspect it to log the surfaced tool-use
+    attempt.
+    """
+
+    def __init__(
+        self,
+        *,
+        message: str = "Tool call not supported by configured policy.",
+        partial_conversation: Optional[list["Message"]] = None,
+    ) -> None:
+        """
+        Initialize the exception.
+
+        Args:
+            message (str): Human-readable error description.
+            partial_conversation (Optional[list[Message]]): Messages produced by
+                the target up to (and including) the assistant turn that
+                contained the disallowed tool call(s).
+        """
+        super().__init__(status_code=400, message=message)
+        self.partial_conversation: list[Message] = (
+            list(partial_conversation) if partial_conversation is not None else []
+        )
+
+
+class ToolCallLoopLimitExceeded(PyritException):
+    """
+    Raised when the tool-use loop runs for more than
+    ``ToolEventPolicy.max_tool_iterations`` iterations without the model
+    producing a stop response.
+
+    The ``partial_conversation`` attribute carries every message produced
+    across all completed iterations. Consumers can inspect it to debug
+    runaway agentic behavior.
+    """
+
+    def __init__(
+        self,
+        *,
+        message: str = "Tool loop exceeded max_tool_iterations without a stop response.",
+        partial_conversation: Optional[list["Message"]] = None,
+    ) -> None:
+        """
+        Initialize the exception.
+
+        Args:
+            message (str): Human-readable error description.
+            partial_conversation (Optional[list[Message]]): Messages produced by
+                the target across every completed iteration of the tool loop.
+        """
+        super().__init__(status_code=400, message=message)
+        self.partial_conversation: list[Message] = (
+            list(partial_conversation) if partial_conversation is not None else []
+        )
+
+
 def pyrit_custom_result_retry(
     retry_function: Callable[..., bool], retry_max_num_attempts: Optional[int] = None
 ) -> Callable[..., Any]:

pyrit/prompt_target/common/discover_target_capabilities.py

@@ -149,6 +149,7 @@ def _permissive_configuration(
         supports_json_output=True,
         supports_editable_history=True,
         supports_system_prompt=True,
+        supports_tool_use=True,
         input_modalities=merged_modalities,
     )
     # Rebuild a fresh configuration from the instance's native capabilities so

pyrit/prompt_target/common/prompt_target.py

@@ -12,6 +12,7 @@
 from pyrit.models.json_response_config import _JsonResponseConfig
 from pyrit.prompt_target.common.target_capabilities import CapabilityName, TargetCapabilities
 from pyrit.prompt_target.common.target_configuration import TargetConfiguration
+from pyrit.tools import ToolCallParser, tool_loop
 
 logger = logging.getLogger(__name__)
 
@@ -85,6 +86,7 @@ def __init__(
             logging.basicConfig(level=logging.INFO)
 
     @final
+    @tool_loop
     async def send_prompt_async(self, *, message: Message) -> list[Message]:
         """
         Validate, normalize, and send a prompt to the target.
@@ -97,6 +99,13 @@ async def send_prompt_async(self, *, message: Message) -> list[Message]:
         3. Delegates to ``_send_prompt_to_target_async`` with the normalized
            conversation.
 
+        When the target's :attr:`configuration.tool_event_policy` is set, the
+        :func:`pyrit.tools.tool_loop` decorator replaces this body with the
+        agentic loop and re-enters :meth:`_send_prompt_to_target_async`
+        repeatedly until the model issues a stop response (or the configured
+        ``max_tool_iterations`` is hit). When no policy is set, the decorator
+        is a no-op and the body below runs unchanged.
+
         Subclasses MUST NOT override this method. Override
         ``_send_prompt_to_target_async`` instead.
 
@@ -132,6 +141,39 @@ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Me
             list[Message]: Response messages from the target.
         """
 
+    @property
+    def _tool_parser(self) -> ToolCallParser | None:
+        """
+        Per-target :class:`ToolCallParser` consulted by :func:`pyrit.tools.tool_loop`.
+
+        Targets that participate in the tool-use loop override this property
+        to return a parser that walks their response messages and extracts
+        :class:`~pyrit.tools.ToolCall` instances. The base default of
+        ``None`` signals "this target does not participate" -- the wrapper
+        short-circuits after the first response.
+
+        Returns:
+            ToolCallParser | None: The parser, or ``None`` for the default
+            no-tool-use behavior.
+        """
+        return None
+
+    def _tool_schemas(self) -> list[dict[str, Any]]:
+        """
+        Outbound tool-schema list sent on the next request to the model.
+
+        Targets that participate in the tool-use loop override this method
+        to translate the active :class:`~pyrit.tools.ToolBackend.schemas`
+        into the wire format their model expects (Responses API vs. Chat
+        Completions API vs. anything else). The base default returns an
+        empty list, which means no schemas are advertised.
+
+        Returns:
+            list[dict[str, Any]]: One schema per advertised tool, in the
+            target-specific wire format. Empty by default.
+        """
+        return []
+
     def _validate_request(self, *, normalized_conversation: list[Message]) -> None:
         """
         Validate the normalized conversation before sending to the target.

pyrit/prompt_target/common/target_capabilities.py

@@ -24,6 +24,7 @@ class CapabilityName(str, Enum):
     JSON_OUTPUT = "supports_json_output"
     EDITABLE_HISTORY = "supports_editable_history"
     SYSTEM_PROMPT = "supports_system_prompt"
+    TOOL_USE = "supports_tool_use"
 
 
 class UnsupportedCapabilityBehavior(str, Enum):
@@ -138,6 +139,14 @@ class attribute. Users can override individual capabilities per instance
     # Whether the target natively supports system prompts.
     supports_system_prompt: bool = False
 
+    # Whether the target natively supports model-issued tool calls (the
+    # canonical OpenAI ``function_call`` / ``function_call_output`` envelopes
+    # plus an outbound tool-schema list). Targets without this capability
+    # cannot host a tool-use loop -- attempting to configure a
+    # :class:`TargetConfiguration` with a ``tool_backend`` on a target whose
+    # capabilities have ``supports_tool_use=False`` raises at construction.
+    supports_tool_use: bool = False
+
     # The input modalities supported by the target (e.g., "text", "image").
     input_modalities: frozenset[frozenset[PromptDataType]] = frozenset({frozenset(["text"])})
 

pyrit/prompt_target/common/target_configuration.py

@@ -4,7 +4,7 @@
 import logging
 from collections.abc import Mapping
 from dataclasses import fields
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from pyrit.message_normalizer import MessageListNormalizer
 from pyrit.models import Message
@@ -16,6 +16,10 @@
     UnsupportedCapabilityBehavior,
 )
 
+if TYPE_CHECKING:
+    from pyrit.tools.backend import ToolBackend
+    from pyrit.tools.models import ToolEventPolicy
+
 logger = logging.getLogger(__name__)
 
 
@@ -39,6 +43,15 @@ class TargetConfiguration:
 
     Each target defines defaults; callers can override policy or individual
     normalizers at creation time.
+
+    Tool use is configured by setting :attr:`tool_event_policy` (mandatory
+    when a target's response contains tool calls; controls EXECUTE / RAISE /
+    RETURN\\_RAW behavior) and optionally :attr:`tool_backend` (required only
+    when ``tool_event_policy.behavior`` is ``EXECUTE``). Both default to
+    ``None`` and are read by :func:`pyrit.tools.tool_loop` at runtime;
+    constructing a configuration with a ``tool_backend`` on a target that
+    does not declare ``capabilities.supports_tool_use=True`` raises
+    immediately.
     """
 
     def __init__(
@@ -47,6 +60,8 @@ def __init__(
         capabilities: TargetCapabilities,
         policy: CapabilityHandlingPolicy | None = None,
         normalizer_overrides: Mapping[CapabilityName, MessageListNormalizer[Any]] | None = None,
+        tool_event_policy: "ToolEventPolicy | None" = None,
+        tool_backend: "ToolBackend | None" = None,
     ) -> None:
         """
         Build a target configuration and resolve the normalization pipeline.
@@ -57,14 +72,34 @@ def __init__(
                 capability. Defaults to RAISE for all adaptable capabilities.
             normalizer_overrides (Mapping[CapabilityName, MessageListNormalizer[Any]] | None):
                 Optional overrides for specific capability normalizers.
+            tool_event_policy (ToolEventPolicy | None): How
+                :func:`pyrit.tools.tool_loop` should react to a pending tool
+                call from the target. ``None`` means the loop is disabled and
+                the wrapper short-circuits.
+            tool_backend (ToolBackend | None): Dispatch table used when
+                ``tool_event_policy.behavior`` is ``EXECUTE``. ``None`` is
+                valid only for the RAISE / RETURN\\_RAW policies and the
+                no-policy passthrough.
+
+        Raises:
+            ValueError: If ``tool_backend`` is set on a target whose
+                capabilities do not include ``supports_tool_use``.
         """
+        if tool_backend is not None and not capabilities.includes(capability=CapabilityName.TOOL_USE):
+            raise ValueError(
+                "tool_backend is set but capabilities.supports_tool_use is False. "
+                "Either declare supports_tool_use=True on the target's capabilities, "
+                "or remove the tool_backend."
+            )
         self._capabilities = capabilities
         self._policy = policy or _DEFAULT_POLICY
         self._pipeline = ConversationNormalizationPipeline.from_capabilities(
             capabilities=self._capabilities,
             policy=self._policy,
             normalizer_overrides=normalizer_overrides,
         )
+        self._tool_event_policy = tool_event_policy
+        self._tool_backend = tool_backend
 
     @property
     def capabilities(self) -> TargetCapabilities:
@@ -81,6 +116,41 @@ def pipeline(self) -> ConversationNormalizationPipeline:
         """The resolved normalization pipeline."""
         return self._pipeline
 
+    @property
+    def tool_event_policy(self) -> "ToolEventPolicy | None":
+        """The tool-use policy consulted by :func:`pyrit.tools.tool_loop`."""
+        return self._tool_event_policy
+
+    @tool_event_policy.setter
+    def tool_event_policy(self, value: "ToolEventPolicy | None") -> None:
+        """Allow runtime updates so callers can opt a configured target into tool use."""
+        self._tool_event_policy = value
+
+    @property
+    def tool_backend(self) -> "ToolBackend | None":
+        """The tool dispatch backend used when the loop's behavior is ``EXECUTE``."""
+        return self._tool_backend
+
+    @tool_backend.setter
+    def tool_backend(self, value: "ToolBackend | None") -> None:
+        """
+        Allow runtime updates to the backend.
+
+        Re-runs the ``supports_tool_use`` validator so a backend can never be
+        installed onto a configuration that does not declare the capability.
+
+        Raises:
+            ValueError: If ``value`` is not ``None`` and the configuration's
+                capabilities do not include ``supports_tool_use``.
+        """
+        if value is not None and not self._capabilities.includes(capability=CapabilityName.TOOL_USE):
+            raise ValueError(
+                "tool_backend is set but capabilities.supports_tool_use is False. "
+                "Either declare supports_tool_use=True on the target's capabilities, "
+                "or remove the tool_backend."
+            )
+        self._tool_backend = value
+
     def includes(self, *, capability: CapabilityName) -> bool:
         """
         Check whether the target includes support for the given capability.

pyrit/tools/__init__.py

@@ -0,0 +1,72 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+"""
+Generic tool-use scaffolding for :class:`~pyrit.prompt_target.PromptTarget`.
+
+This package provides a transport-agnostic tool-calling loop. The
+:func:`tool_loop` decorator, when applied to ``send_prompt_async``, runs
+the standard PyRIT validate+normalize work once and then repeatedly
+re-enters the target's protected ``_send_prompt_to_target_async`` until
+the model issues a stop response (or a configured limit is hit).
+
+A target opts in by declaring two collaborators:
+
+* ``self._tool_parser`` — a :class:`ToolCallParser` that walks a
+  response message and extracts pending :class:`ToolCall` instances.
+* ``self.configuration.tool_event_policy`` — a :class:`ToolEventPolicy`
+  whose :class:`ToolEventBehavior` decides whether to ``EXECUTE``,
+  ``RAISE``, or ``RETURN_RAW`` on each detected call.
+
+When the policy is ``EXECUTE``, calls are dispatched through
+``self.configuration.tool_backend``, an implementation of
+:class:`ToolBackend`. :class:`LocalToolBackend` is the in-process
+backend shipped here; :class:`MCPToolBackend` ships in C3 and proxies
+through one or more MCP servers.
+
+The :class:`ToolBackend` Protocol is intentionally distinct from
+:mod:`pyrit.registry` — that namespace is reserved for framework-level
+identity registries (``TargetRegistry``, ``ScorerRegistry``) that
+register named singletons for CLI lookup, which a per-target tool
+dispatch table is not.
+
+Wiring of ``@tool_loop`` onto :class:`PromptTarget.send_prompt_async`
+and of the ``tool_event_policy`` / ``tool_backend`` fields onto
+:class:`TargetConfiguration` lands in C4/C5.
+
+The two exception types the loop raises
+(:class:`~pyrit.exceptions.ToolCallNotSupported` and
+:class:`~pyrit.exceptions.ToolCallLoopLimitExceeded`) live in
+:mod:`pyrit.exceptions` alongside the rest of PyRIT's exception
+catalog, so non-tools callers (attacks, normalizers) can import them
+without taking a subsystem-level dependency on ``pyrit.tools``.
+"""
+
+from pyrit.tools.backend import ToolBackend
+from pyrit.tools.local_backend import LocalToolBackend
+from pyrit.tools.mcp_backend import MCPToolBackend
+from pyrit.tools.mcp_client import (
+    DockerMCPServerSpec,
+    LocalMCPServerSpec,
+    MCPClient,
+    MCPServerSpec,
+    RemoteMCPServerSpec,
+)
+from pyrit.tools.models import ToolCall, ToolEventBehavior, ToolEventPolicy, tool_loop
+from pyrit.tools.parsers import ToolCallParser
+
+__all__ = [
+    "DockerMCPServerSpec",
+    "LocalMCPServerSpec",
+    "LocalToolBackend",
+    "MCPClient",
+    "MCPServerSpec",
+    "MCPToolBackend",
+    "RemoteMCPServerSpec",
+    "ToolBackend",
+    "ToolCall",
+    "ToolCallParser",
+    "ToolEventBehavior",
+    "ToolEventPolicy",
+    "tool_loop",
+]