Support tool callbacks in MCP sampling#2998
Conversation
|
/review |
docker-agent
left a comment
There was a problem hiding this comment.
Assessment: 🟡 NEEDS ATTENTION
Two medium-severity findings in the newly-added sampling-with-tools code. The core stream-aggregation logic, capability handshake, content building, and limits enforcement are all well-structured and correctly tested.
|
Sorry @EronWright, our project requires all the commit to be signed and verified. Could you sign them? Thanks! |
|
I moved it back to draft @EronWright the time you rewrite the history with signed commits. Sorry for the issue |
b1fad0b to
89a93e6
Compare
|
@aheritier sorry for delay, all commits are now signed/verified. |
89a93e6 to
8e150a2
Compare
|
/review |
docker-agent
left a comment
There was a problem hiding this comment.
Assessment: 🟡 NEEDS ATTENTION
Two medium-severity issues found in the new sampling-with-tools path. Both relate to edge cases in the streaming and reconnect logic.
8e150a2 to
e4ff6d4
Compare
|
The build issue seems for unrelated reasons, re-run? |
|
I think you rebased it on main which was red at that time. Can you rebase again @EronWright ? |
Adds a parallel SamplingWithToolsHandler alongside the existing SamplingHandler so MCP servers can include a tools array in sampling/createMessage requests. The host drives its model with those tools and returns any tool_use blocks as ToolUseContent; the server remains responsible for executing the tool and continuing the loop in a follow-up sampling request. The initialize handshake now advertises sampling.tools capability, and the MCP toolset selects the appropriate go-sdk handler (basic vs. with-tools) based on which handler is registered. Signed-off-by: Eron Wright <eronwright@gmail.com>
Mounts an in-process gomcp.NewServer on an httptest server via StreamableHTTPHandler. Its one tool, ask_with_calculator, runs a sampling loop: sends sampling/createMessage with a calculator tool, gets a tool_use back from the host LLM, "executes" the calculator, sends a follow-up sampling request carrying the tool_result, and returns the final text. The Gemini side is recorded once and replayed on subsequent runs, so the test runs offline in CI. Signed-off-by: Eron Wright <eronwright@gmail.com>
- Reject tool_use blocks under a non-assistant role explicitly rather than silently dropping the message. The MCP spec places tool_use on assistant turns, but a malformed server would previously have its message disappear from the converted chat history with no error. - Document the ordering dependency between SetSampling*Handler and Initialize in both stdio and remote MCP client setups, so future maintainers don't reorder them and silently lose the sampling handler on reconnect. Signed-off-by: Eron Wright <eronwright@gmail.com>
Use errors.New for the static loop-exhausted message and strconv.FormatInt for the calculator result instead of fmt.Errorf / fmt.Sprintf %d. Signed-off-by: Eron Wright <eronwright@gmail.com>
- pkg/runtime/sampling.go: mirror the streaming.go short-circuit in
drainSamplingStreamWithTools so a chunk that carries only tool-call
deltas continues to the next frame instead of falling through into
the FinishReason capture and break. The existing flow happens to be
equivalent, but the explicit guard pre-empts regression if anyone
inserts new processing between the tool-call loop and the break.
- pkg/tools/mcp/{stdio,remote,session_client}.go: collapse the inline
CreateMessage* registration switch into a shared
sessionClient.applySamplingHandlerOpts helper. The helper always
wires the with-tools callback unless the caller has explicitly opted
into the basic-only path, so a SetSamplingWithToolsHandler call
landing after Initialize — including on supervisor-driven reconnect
that races configureToolsetHandlers — still takes effect via the
existing lazy read in handleSamplingWithToolsRequest. Closes the
silent no-handler window the reviewer flagged on stdio.go:53 and
remote.go:104.
Tests:
- TestDrainSamplingStreamWithTools/tool-call-only chunks continue ...
- TestApplySamplingHandlerOpts_RegistrationMatrix
- TestHandleSamplingWithToolsRequest_LateSetterTakesEffect
Signed-off-by: Eron Wright <eronwright@gmail.com>
e4ff6d4 to
182e234
Compare
|
Rebased on |
Summary
Closes the tool callbacks functional gap in MCP sampling support — a follow-up to #2815, addressing one of the remaining items from #2809.
When an MCP server includes a
toolsarray in asampling/createMessagerequest, the host now drives its model with those tools and returns anytool_useblocks back to the server asToolUseContent. The server remains responsible for executing the tool and continuing the loop in a follow-up sampling request.sequenceDiagram participant H as cagent participant S as MCP Server participant L as LLM activate H H->>+S: tools/call {name, arguments} note over S: needs LLM inference S->>+H: sampling/createMessage<br/>{messages, tools: [...]} H->>+L: chat completion L-->>-H: ToolUseContent<br/>stopReason: "toolUse" H-->>-S: CreateMessageResult<br/>{tool_use, stopReason: "toolUse"} note over S: executes tool locally S->>+H: sampling/createMessage<br/>{messages + tool_use + tool_result, tools: [...]} H->>+L: chat completion L-->>-H: TextContent<br/>stopReason: "endTurn" H-->>-S: CreateMessageResult<br/>{text, stopReason: "endTurn"} S-->>-H: tool result deactivate HWhat's new
SamplingWithToolsHandlertype andSampleableWithToolsinterface — additive, parallel to the existingSamplingHandler/Sampleable. No breaking changes to the basic sampling path merged in feat(mcp): add sampling/createMessage support #2815.Initialize, exactly one of the SDK's mutually exclusiveClientOptions.CreateMessage*fields is populated — prefer with-tools when registered, fall back to basic.sampling.toolsso servers know the host can receive tool-enabled requests.pkg/runtime/sampling.go):text,image/audio,tool_use→ assistantToolCalls,tool_result→MessageRoleToolrows (parallel tool_results expand to multiple chat.Message rows).[]*mcp.Tool→[]tools.Toolwith a no-op handler (the server, not the host, executes).model.CreateChatCompletionStream, aggregates streamed tool calls.ContentwithTextContent+ToolUseContentblocks;stopReason: "toolUse"when tool calls are present.maxSamplingTools=64,maxSamplingToolCalls=32.e2e/sampling_test.go): mounts an in-processgomcp.NewServeron anhttptestserver viaStreamableHTTPHandler. The server exposes one tool (ask_with_calculator) whose handler drives a real sampling-with-tools loop against the connecting cagent. The Gemini side is recorded once and replayed on subsequent runs, so the test runs offline in CI.Out of scope (separate gaps from #2809)
Test plan