Skip to content

FEAT: Add MCP-wrapped prompt chat target - #2585

Open
Java123456com wants to merge 2 commits into
microsoft:mainfrom
Java123456com:feat/mcp-wrapped-prompt-chat-target
Open

FEAT: Add MCP-wrapped prompt chat target#2585
Java123456com wants to merge 2 commits into
microsoft:mainfrom
Java123456com:feat/mcp-wrapped-prompt-chat-target

Conversation

@Java123456com

Copy link
Copy Markdown

Addresses #1273 (MCP integration), scoped along the refined direction from the issue discussion: composition/wrapping of a generic chat target, not OpenAIResponseTarget-specific. A design summary was posted on the issue before this PR; this PR is deliberately the phase-1 slice.

What this adds

MCPWrappedPromptChatTarget(inner_target, mcp_servers, ...) - a PromptTarget wrapper that composes MCP tool use into any inner target that supports system prompts and multi-turn conversation (LiteLLM, Azure ML, Hugging Face, ...). The wrapper:

  1. connects to the configured MCP servers (sessions established lazily, reused across sends, closed in cleanup_target_async);
  2. lists each server's tools and injects the catalog plus a strict TOOL_CALL: {json} protocol into the system context;
  3. drives a bounded agentic loop: model tool calls are executed through the MCP session, results are fed back as labeled TOOL_RESULT/TOOL_ERROR user turns, and the loop continues until the model answers without a tool call or max_tool_call_rounds (default 5) is exhausted;
  4. delegates every round through the inner target's _send_prompt_to_target_async extension point, so the inner target formats the full conversation each round (consistent with how chat targets consume normalized_conversation).

Intermediate agentic rounds stay inside the wrapper and are audit-logged (server, tool, arguments, result size); only the final response is returned to the caller, so memory and reports reflect conversation-level behavior.

MCPServerConfig / MCPTransport - stdio (subprocess) and streamable-http transports through the official mcp Python SDK, behind the new optional extra pyrit[mcp] (lazy import with an actionable install hint; no hard dependency). Transport connection parameters are validated at construction, and plain SSE is not carried forward (deprecated in favor of streamable-HTTP in the MCP spec).

Policy and audit - per-server allowed_tools plus a global allowlist, unknown-tool rejection without execution, per-call timeout, and a round cap with an explicit warning. Malformed TOOL_CALL payloads bounce back to the model as an error turn instead of being executed.

Design decisions (feedback welcome)

  • Text protocol vs native tool calls: for a generic PromptTarget wrapper there is no native tool-call channel; the protocol instructions are strict and parsing is defensive (last valid call wins, malformed calls are non-fatal). Wiring the existing MessagePieceType.MCP_CALL / MCP_LIST_TOOLS constants through OpenAIResponseTarget's native Responses-API MCP support is the natural phase 2 and is deliberately out of scope here.
  • Memory semantics: each round goes through the inner target's extension point; intermediate turns are not persisted to the wrapper's memory. If maintainers prefer full round-by-round persistence, that is a small change.
  • Example servers: only test fixtures ship in this PR (a tiny FastMCP server over stdio); curated example servers can live in docs/examples once the design is agreed.

Tests

tests/unit/mcp/ (20 tests):

  • test_mcp_client.py drives a real FastMCP server over stdio (subprocess) end-to-end: initialize handshake, list_tools catalog, call_tool round-trip, server-side tool errors surfacing as is_error, and unconnected-session errors.
  • test_mcp_wrapped_prompt_chat_target.py covers the loop protocol with stubbed sessions: catalog/system-context injection, tool-call round-trip, unknown-tool and allowlist denial (rejected without execution), round cap, malformed payload handling, capability validation of the inner target, and session cleanup.

All 20 pass, plus 1,274 passing tests across tests/unit/prompt_target and tests/unit/models; ruff check / ruff format clean; uv.lock regenerated for the new extra.

Adds native MCP (Model Context Protocol) tool use to PyRIT via composition,
per the implementation direction refined on microsoft#1273:

- MCPWrappedPromptChatTarget wraps any PromptTarget that supports system
  prompts and multi-turn conversations; it lists the tools of the configured
  MCP servers, injects the catalog plus a strict TOOL_CALL protocol into the
  system context, executes model-requested calls, and feeds results back as
  labeled turns inside a bounded agentic loop (max_tool_call_rounds).
- MCPServerConfig/MCPTransport support stdio and streamable-http transports
  through the official mcp Python SDK, behind the new optional extra
  pyrit[mcp] (lazy import with an actionable error).
- Policy and audit: per-server and global tool allowlists, unknown-tool
  rejection, per-call timeout, and audit logging of every tool invocation.

Intermediate rounds stay inside the wrapper; only the final response reaches
memory and scoring. Unit tests drive a real FastMCP server over stdio for the
client transport and cover the loop protocol (catalog injection, tool
execution round-trip, unknown/denied tool handling, round cap, cleanup).

Addresses microsoft#1273
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