Support tool callbacks in MCP sampling

[EronWright]

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 tools array in a sampling/createMessage request, the host now drives its model with those tools and returns any tool_use blocks back to the server as ToolUseContent. 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 H

Loading

What's new

• New SamplingWithToolsHandler type and SampleableWithTools interface — additive, parallel to the existing SamplingHandler / Sampleable. No breaking changes to the basic sampling path merged in feat(mcp): add sampling/createMessage support #2815.
• MCP toolset wires both handler types. At Initialize, exactly one of the SDK's mutually exclusive ClientOptions.CreateMessage* fields is populated — prefer with-tools when registered, fall back to basic.
• Capability handshake advertises sampling.tools so servers know the host can receive tool-enabled requests.
• Runtime handler (pkg/runtime/sampling.go):
  • Converts V2 multi-block messages: text, image/audio, tool_use → assistant ToolCalls, tool_result → MessageRoleTool rows (parallel tool_results expand to multiple chat.Message rows).
  • Converts []*mcp.Tool → []tools.Tool with a no-op handler (the server, not the host, executes).
  • Drives model.CreateChatCompletionStream, aggregates streamed tool calls.
  • Builds result Content with TextContent + ToolUseContent blocks; stopReason: "toolUse" when tool calls are present.
• New limits: maxSamplingTools=64, maxSamplingToolCalls=32.
• End-to-end test (e2e/sampling_test.go): mounts an in-process gomcp.NewServer on an httptest server via StreamableHTTPHandler. 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)

• Human-in-the-loop approval UI
• Model-preference hints

Test plan

• Unit tests pass: `go test ./pkg/runtime/... ./pkg/tools/...`
• `go build ./...` clean
• `go vet ./...` clean
• `task lint` (0 offenses)
• `gofmt` clean on all changed files
• End-to-end via `TestExec_Gemini_SamplingWithTools` (cassette replays offline; verified once against live Gemini):
  • Handshake includes `sampling.tools` capability (implicit — `ServerSession.CreateMessageWithTools` would refuse otherwise)
  • LLM receives the server-supplied tools
  • Response contains `ToolUseContent` with `stopReason: "toolUse"` when the model emits a tool_use
  • Follow-up sampling request with `tool_result` blocks is converted correctly and the loop terminates with `endTurn`

[aheritier]

/review

[docker-agent]

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.

[dgageot]

Sorry @EronWright, our project requires all the commit to be signed and verified. Could you sign them? Thanks!

[aheritier]

I moved it back to draft @EronWright the time you rewrite the history with signed commits. Sorry for the issue

[EronWright]

@aheritier sorry for delay, all commits are now signed/verified.

[aheritier]

/review

[docker-agent]

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.

[EronWright]

The build issue seems for unrelated reasons, re-run?

[aheritier]

I think you rebased it on main which was red at that time. Can you rebase again @EronWright ?

[EronWright]

Rebased on fbe505064 (current main; green). The previous build-and-test failure was on a red main revision as you suspected. One merge-fix folded into the base commit: main added a second tools.ConfigureHandlers call site in skillSubSessionTools (loop.go) that now also passes r.samplingWithToolsHandler. Local go build ./..., go test ./pkg/runtime/... ./pkg/tools/mcp/..., the sampling cassette e2e, and golangci-lint are all green. Ready for another CI approval.

[Sayt-0]

build-and-test fails: the e2e package panics with SetDefaultTokenStoreFactory called after the default token store was already created, because the new sampling test builds the global token store and the next Execute re-runs keyringstore.Register(). Making Register() idempotent (e.g. a sync.Once) should fix it.