FEAT: Add MCP-wrapped prompt chat target - #2585
Open
Java123456com wants to merge 1 commit into
Open
Conversation
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
|
Java123456com please read the following Contributor License Agreement(CLA). If you agree with the CLA, please reply with the following information.
Contributor License AgreementContribution License AgreementThis Contribution License Agreement (“Agreement”) is agreed to by the party signing below (“You”),
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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, ...)- aPromptTargetwrapper that composes MCP tool use into any inner target that supports system prompts and multi-turn conversation (LiteLLM, Azure ML, Hugging Face, ...). The wrapper:cleanup_target_async);TOOL_CALL: {json}protocol into the system context;TOOL_RESULT/TOOL_ERRORuser turns, and the loop continues until the model answers without a tool call ormax_tool_call_rounds(default 5) is exhausted;_send_prompt_to_target_asyncextension point, so the inner target formats the full conversation each round (consistent with how chat targets consumenormalized_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 officialmcpPython SDK, behind the new optional extrapyrit[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_toolsplus a global allowlist, unknown-tool rejection without execution, per-call timeout, and a round cap with an explicit warning. MalformedTOOL_CALLpayloads bounce back to the model as an error turn instead of being executed.Design decisions (feedback welcome)
PromptTargetwrapper 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 existingMessagePieceType.MCP_CALL/MCP_LIST_TOOLSconstants through OpenAIResponseTarget's native Responses-API MCP support is the natural phase 2 and is deliberately out of scope here.Tests
tests/unit/mcp/(20 tests):test_mcp_client.pydrives a real FastMCP server over stdio (subprocess) end-to-end: initialize handshake,list_toolscatalog,call_toolround-trip, server-side tool errors surfacing asis_error, and unconnected-session errors.test_mcp_wrapped_prompt_chat_target.pycovers 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_targetandtests/unit/models;ruff check/ruff formatclean;uv.lockregenerated for the new extra.