Skip to content

[DRAFT] FEAT: Tool Use + MCP#1811

Draft
ValbuenaVC wants to merge 7 commits into
microsoft:mainfrom
ValbuenaVC:MCP
Draft

[DRAFT] FEAT: Tool Use + MCP#1811
ValbuenaVC wants to merge 7 commits into
microsoft:mainfrom
ValbuenaVC:MCP

Conversation

@ValbuenaVC
Copy link
Copy Markdown
Contributor

@ValbuenaVC ValbuenaVC commented May 26, 2026

Description

Today PyRIT's tool-calling story is fragmented and incomplete:

  • OpenAIChatTarget parses tool_calls into function_call pieces and stops — no execution, no loop.
  • OpenAIResponseTarget hand-rolls a complete agentic loop inside _send_prompt_to_target_async (pyrit/prompt_target/openai/openai_response_target.py:590-626), accepting a custom_functions registry of Python callables. It handles one tool call per turn.
  • MCP servers are not a recognized concept.

This PR introduces a single, target-agnostic tool-use primitive:

  • New pyrit/tools/ package with a tool_loop decorator, a ToolCallParser Protocol (per-target detection), and a ToolBackend ABC with two concrete backends in v1 — LocalToolBackend (Python callables, preserves today's custom_functions behavior) and MCPToolBackend (stdio MCP servers via the official mcp SDK, composing one MCPClient per MCPServerSpec through a shared AsyncExitStack).
  • New TargetCapabilities.tool_use flag and ToolEventPolicy (EXECUTE / RAISE / RETURN_RAW) on TargetConfiguration.
  • OpenAIResponseTarget migrated onto the decorator (behavior-preserving, except multiple tool calls in one turn are now dispatched all-at-once sequentially — the protocol-intended behavior).
  • OpenAIChatTarget gains end-to-end tool calling.
  • custom_functions kwarg on OpenAIResponseTarget deprecated (warns, removed_in="0.16.0"); internally rewrapped as a LocalToolBackend.

Future MCP transports (HTTP/SSE, Docker sandbox), additional sandbox providers, and streaming all plug in behind the existing ToolBackend / MCPServerSpec interfaces with no abstraction changes. The MCPServerSpec union ships with three variants in v1: LocalMCPServerSpec (the only one with a working transport) plus stub declarations of RemoteMCPServerSpec and DockerMCPServerSpec that raise NotImplementedError in connect_async. The follow-up sandbox PR's diff becomes "implement an already-declared variant" rather than "expand the union + add an implementation."

Closes nothing existing; tracks future work in TODOs marked # TODO(streaming-v2), # TODO(mcp-http-transport), # TODO(mcp-resources), # TODO(sandbox-provider).

Compatibility

This PR is not breaking for the standard tool-calling path. A short list of source- and behavior-compat caveats reviewers should know about:

  1. Source-compat — PromptTarget.send_prompt_async is @final (C5). External subclasses that override the public entrypoint (not just _send_prompt_to_target_async) will fail to import. No in-tree target overrides it today.
  2. Source-compat — OpenAIChatTarget function_call envelope rename (C6). The Chat target's function_call piece switches to the canonical envelope (call_id / name / arguments at the top level) shared with OpenAIResponseTarget. Callers that introspected the previous nested {"function": {...}} JSON shape will need to update. Today's Chat target only parses tool calls — it does not dispatch — so callers that forwarded pieces verbatim downstream are unaffected.
  3. Deprecation — OpenAIResponseTarget(custom_functions=...) (C7). The kwarg now emits DeprecationWarning(removed_in="0.16.0") and is internally rewrapped as a LocalToolBackend. No runtime behavior change in 0.15.x.
  4. Intentional behavior change — multi-call-per-turn dispatch (C7). When a model response contains N>1 tool calls, the new loop dispatches all N sequentially in declaration order. The old hand-rolled loop in OpenAIResponseTarget._perform_async_with_tools only dispatched the last call per turn. This is strictly more dispatching, not less, so it cannot regress any working code; it matches the OpenAI protocol's actual intent.
  5. Private API removal (C7). _find_last_pending_tool_call, _execute_call_section, and _make_tool_piece on OpenAIResponseTarget are removed; their logic moves into pyrit/tools/openai_response_helpers.py. Listed for changelog completeness — these were always private.

Tests and Documentation

  • New tests/unit/tools/ directory covering the decorator, parsing, LocalToolBackend, MCPClient (real stdio subprocess against a deterministic FastMCP fixture), and MCPToolBackend (multi-server routing, name-collision detection, name_prefix disambiguation, allowed_tools filtering, and concurrent-dispatch serialization).
  • New tests/unit/prompt_target/common/test_prompt_target_tool_loop.py asserting decorator-order-of-execution against _FakeToolTarget and using patch_central_database to verify per-message insert ordering, per-role labeling (assistant, tool), and per-data-type labeling (function_call, function_call_output) in the actual DB schema.
  • Updated tests/unit/prompt_target/target/test_openai_response_target_function_chaining.py and new tests/unit/prompt_target/target/test_openai_chat_target_tool_loop.py covering both targets' parsers, _tool_schemas() outbound translation, deprecation warning on custom_functions, and multi-call-per-turn sequential dispatch.
  • New tests/integration/tools/test_red_teaming_with_tools.py running the real RedTeamingAttack against both targets with only the HTTP layer mocked. Tools served by the real echo_mcp_server subprocess.
  • New tests/integration/scenarios/test_scenario_with_mcp_tools.py (C10) running a lightweight scenario end-to-end against a real echo_mcp_server subprocess and verifying the Memory DB transcript for function_call / function_call_output pieces in declaration order.
  • No notebook/doc additions in this PR — follow-up scenarios PR will exercise the public API.

JupyText: not applicable (no notebook changes).

Victor Valbuena and others added 4 commits May 26, 2026 13:59
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@ValbuenaVC ValbuenaVC changed the title [DRAFT] FEAT: Tool Use via MCP [DRAFT] FEAT: Tool Use + MCP May 26, 2026
Victor Valbuena and others added 3 commits May 26, 2026 15:56
… into PromptTarget.send_prompt_async

C4 lands the in-tree wiring for the generic tool-use loop introduced by C2/C3:

- TargetCapabilities gains supports_tool_use: bool (default False) and
  CapabilityName.TOOL_USE for the corresponding enum value, matching the
  existing supports_X / "supports_X" naming convention used by every
  other capability.
- TargetConfiguration grows tool_event_policy + tool_backend kwargs,
  both gettable/settable properties. The setter (and constructor)
  validate that a non-None tool_backend requires supports_tool_use=True;
  otherwise they raise ValueError immediately. ToolBackend /
  ToolEventPolicy imports are quoted + behind TYPE_CHECKING to keep
  pyrit.prompt_target.common from importing pyrit.tools eagerly.
- PromptTarget.send_prompt_async picks up @tool_loop (below the existing
  @Final). The wrapper is a no-op when tool_event_policy is None, so
  every existing target keeps its current behavior. _tool_parser
  (property, default None) and _tool_schemas() (default []) are added
  on the base class as the two collaborators @tool_loop reads.
- _permissive_configuration is updated to flip supports_tool_use=True
  alongside the other supports_X flags so the all-flags-on probe loop
  in test_discover_target_capabilities still sees every CapabilityName
  value as supported.

tests/unit/tools/conftest.py drops the hand-decorated @tool_loop on
_FakeToolTarget.send_prompt_async (which would now violate the base
class's @Final) and instead wires policy + backend through
TargetConfiguration. _tool_parser becomes a subclass property since
the base class now defines one.

Tests:
- test_tool_event_policy.py adds U7 (capability flag wiring through the
  wrapper) plus dataclass field defaults and the TargetConfiguration
  validator.
- test_prompt_target_tool_loop.py adds U1 / U2 (DB-end) / U8 / U9 / U11
  exercised against a _ProductionShapedTarget that uses the real
  base-class _get_normalized_conversation_async (memory round-trip via
  patch_central_database). Plus default-_tool_parser / -_tool_schemas
  assertions.

Validation: 8104 unit tests pass; pre-commit clean.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant