feat(mcp): add MCP Tasks support at the SDK level (#21)#22
Merged
Conversation
Implement Phase 1 of the MCP Tasks spec so long-running tools/call
requests can run asynchronously, survive proxy timeouts, and be
cancelled mid-flight.
- connection pool now advertises mcp.NewTasksCapability() during
initialize and captures the InitializeResult so callers can detect
per-server task support
- new MCPServerConfig.TasksMode (auto|never|always, default auto)
parsed from both new and legacy mcp.json shapes
- ExecuteTool augments tools/call with TaskParams when policy and
capability allow, polls tasks/get / tasks/result until terminal,
and best-effort tasks/cancel on context cancellation
- new MCPToolManager methods: SetTaskConfig, ListServerTasks,
GetServerTask, CancelServerTask
- public SDK surface in pkg/kit: MCPTask, MCPTaskStatus, MCPTaskMode,
MCPTaskProgress, MCPTaskProgressHandler, plus Options fields
(MCPTaskMode, MCPTaskTimeout, MCPTaskTTL, MCPTaskPollInterval,
MCPTaskMaxPollInterval, MCPTaskProgress) and Kit.{List,Get,Cancel}
MCPTask methods
- works around two upstream mcp-go v0.51.0 parser bugs
(ParseCallToolResult rejects task responses; ParseTaskResultResult
looks for content under a non-existent nested key) by decoding the
wire shape directly via the transport
- defaults to MCPTaskModeAuto so servers that don't advertise task
support behave exactly as before
Fixes #21
|
Connected to Huly®: KIT-23 |
|
Warning Rate limit exceeded
To keep reviews running without waiting, you can enable usage-based add-on for your organization. This allows additional reviews beyond the hourly cap. Account admins can enable it under billing. ⌛ How to resolve this issue?After the wait time has elapsed, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout. Please see our FAQ for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (10)
📝 WalkthroughWalkthroughThis PR adds MCP Tasks support to Kit's MCP client and SDK. It enables asynchronous, pollable tool execution via task-augmented ChangesMCP Tasks Feature
Sequence Diagram(s)sequenceDiagram
participant Client as Client<br/>(Kit Agent)
participant Pool as Connection<br/>Pool
participant Server as MCP Server
participant Manager as Tool<br/>Manager
rect rgba(0, 100, 200, 0.5)
Note over Client,Server: Initialization: Advertise Task Capability
Client->>Pool: Initialize with<br/>TasksCapability
Pool->>Server: Initialize Request
Server-->>Pool: InitializeResult with<br/>Tasks capability
Pool->>Pool: Store InitializeResult<br/>Set SupportsToolTasks=true
end
rect rgba(100, 150, 100, 0.5)
Note over Client,Manager: Tool Execution: Detect & Route
Client->>Manager: ExecuteTool(toolName, args)
Manager->>Pool: Check server<br/>task support
Manager->>Manager: Resolve task mode<br/>(auto/never/always)
end
rect rgba(200, 100, 100, 0.5)
Note over Manager,Server: Task Path: Async Dispatch & Poll
Manager->>Server: tools/call with<br/>TaskParams{TTL}
Server-->>Manager: CreateTaskResult<br/>{taskId, status}
loop Poll Until Terminal
Manager->>Server: tasks/get{taskId}
Server-->>Manager: Task{status, TTL,<br/>pollInterval}
Manager->>Manager: Emit Progress
alt status is terminal
Manager->>Server: tasks/result{taskId}
Server-->>Manager: TaskResult
end
end
end
rect rgba(150, 100, 150, 0.5)
Note over Manager,Client: Return & Progress
Manager-->>Client: MCPToolResult with<br/>final content
alt Progress Callback
Manager->>Client: Progress(status,<br/>statusMessage)
end
end
sequenceDiagram
participant Agent as Agent/Kit
participant SDK as SDK Config
participant Manager as Tool Manager
participant Pool as Connection Pool
rect rgba(100, 150, 200, 0.5)
Note over Agent,Pool: Configuration Flow at Agent Setup
Agent->>SDK: New(WithMCPTaskMode{...},<br/>WithMCPTaskTTL, ...)
SDK->>SDK: toToolsConfig() converts<br/>SDK → internal types
Agent->>Manager: SetTaskConfig(MCPTaskConfig)
Manager->>Manager: Store taskCfg
end
rect rgba(150, 100, 100, 0.5)
Note over Agent,Manager: Lazy Server Addition (AddMCPServer)
Agent->>Manager: AddMCPServer(serverName)
Manager->>Pool: Create new connection
Pool->>Manager: SetTaskConfig(same cfg)
end
rect rgba(100, 100, 150, 0.5)
Note over SDK,Manager: Task Management APIs (ListMCPTasks, etc.)
SDK->>Manager: ListMCPTasks(ctx, serverName)
Manager->>Pool: Get task client for server
Manager->>Pool: Call tasks/list RPC
Manager->>SDK: Convert MCPTaskInfo → MCPTask
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~75 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Review rate limit: 0/1 reviews remaining, refill in 46 minutes and 47 seconds.Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (3)
internal/tools/mcp_tasks.go (1)
304-363: 💤 Low valueWorkaround for upstream mcp-go parsing issue is well-documented.
The custom
fetchTaskResultdecoder correctly works around the mcp-go v0.51.0 issue whereParseTaskResultResultlooks for content under a nested "result" key that doesn't exist in the wire format.Consider adding a TODO or tracking issue reference so this can be removed when mcp-go is fixed.
💡 Suggested improvement
// fetchTaskResult issues tasks/result on the transport and parses the raw // response. The upstream client.TaskResult helper delegates to // mcp.ParseTaskResultResult which (as of mcp-go v0.51.0) looks for the // content array under a nested "result" key that never exists in the // wire format — leading to systematically empty Content. Doing the // parse here keeps the polling path working until that is fixed upstream. +// +// TODO: Remove this workaround when mcp-go fixes ParseTaskResultResult. +// See: https://github.com/mark3labs/mcp-go/issues/XXX (file upstream issue) func fetchTaskResult(ctx context.Context, c *client.Client, taskID string) (*mcp.TaskResultResult, error) {🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@internal/tools/mcp_tasks.go` around lines 304 - 363, Add a TODO comment and tracking issue reference to the custom fetchTaskResult implementation so future maintainers know this is a temporary workaround for mcp-go v0.51.0's ParseTaskResultResult bug; update the comment block above the fetchTaskResult function to include a clear "TODO: remove when fixed" line plus either the upstream issue URL or a placeholder (e.g. GH-xxxx) and the mcp-go version (v0.51.0) so it can be removed once the upstream fix lands.internal/tools/mcp.go (2)
726-729: 💤 Low valueRedundant progress handler parameter.
pollTaskUntilTerminalreceivesm.taskCfgwhich already contains theProgresshandler, but it's also passed separately as the last argument. The function signature inmcp_tasks.go:221-228shows bothcfg MCPTaskConfig(which hascfg.Progress) and a separateprogress MCPTaskProgressHandlerparameter.This redundancy could lead to confusion about which handler is used. Consider either:
- Using only
cfg.ProgressinsidepollTaskUntilTerminaland removing the separate parameter, or- Documenting that the explicit
progressparameter overridescfg.Progress🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@internal/tools/mcp.go` around lines 726 - 729, The call to pollTaskUntilTerminal passes m.taskCfg and also m.taskCfg.Progress redundantly; change pollTaskUntilTerminal (in mcp_tasks.go: function pollTaskUntilTerminal) to remove the separate progress MCPTaskProgressHandler parameter and have the function use cfg.Progress from the MCPTaskConfig type (MCPTaskConfig and MCPTaskProgressHandler), then update all call sites (including this call that currently passes mapping.serverName, taskResult.Task, m.taskCfg, m.taskCfg.Progress) to omit the extra progress argument so only cfg is passed; ensure any documentation/comments reflect that pollTaskUntilTerminal uses cfg.Progress exclusively.
673-691: 💤 Low valueConsider logging when task capability is unavailable due to client type.
When
conn.clientis not a*client.Client, the code silently falls back to synchronous execution. This could surprise users who configuredMCPTaskModeAlwaysbut don't see task behavior.A debug-level log would help diagnose such situations without breaking existing behavior.
💡 Suggested improvement
rawClient, ok := conn.client.(*client.Client) if !ok { // Older client implementations — fall back to the synchronous shape. + if m.debugLogger != nil && m.debugLogger.IsDebugEnabled() { + m.debugLogger.LogDebug(fmt.Sprintf("[DEBUG] Task mode enabled but client type %T does not support tasks; falling back to sync", conn.client)) + } callParams.Task = nil🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@internal/tools/mcp.go` around lines 673 - 691, When conn.client is not a *client.Client the code silently falls back to synchronous execution; add a debug-level log before the fallback that states the task capability is unavailable and we are falling back to synchronous "tools/call" execution, including identifying info such as mapping.serverName and the attempted task setting (e.g., callParams.Task or the configured MCPTaskModeAlways), so operators can see why tasks weren’t used; keep behavior identical otherwise and use the existing logger used in this scope for the message.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@internal/config/config.go`:
- Around line 41-56: In Config.Validate(), add a check for the TasksMode field
to fail-fast on invalid values: verify TasksMode is either "" (empty), "auto",
"never", or "always" and return a validation error if it's anything else; update
the Config.Validate() method to perform this check early and produce a clear
error message mentioning the invalid TasksMode value so callers know which
setting is wrong.
In `@pkg/kit/kit.go`:
- Around line 1425-1432: The MCP task options set into MCPTaskConfig (via
mcpTaskOptions -> toToolsConfig) are only applied to the top-level Kit and not
carried into child agents; update Subagent() so when it constructs the child
Options (and its MCPConfig) it also copies/propagates the MCP task fields
(perServer/MCPTaskMode, defaultTTL/MCPTaskTTL, pollInterval/MCPTaskPollInterval,
maxPollInterval/MCPTaskMaxPollInterval, timeout/MCPTaskTimeout,
progress/MCPTaskProgress) — e.g., reuse mcpTaskOptions.toToolsConfig() or
explicitly copy those fields into the child Options.MCPConfig/MCPTaskConfig
construction so child subagents get the same task overrides as the parent Kit.
---
Nitpick comments:
In `@internal/tools/mcp_tasks.go`:
- Around line 304-363: Add a TODO comment and tracking issue reference to the
custom fetchTaskResult implementation so future maintainers know this is a
temporary workaround for mcp-go v0.51.0's ParseTaskResultResult bug; update the
comment block above the fetchTaskResult function to include a clear "TODO:
remove when fixed" line plus either the upstream issue URL or a placeholder
(e.g. GH-xxxx) and the mcp-go version (v0.51.0) so it can be removed once the
upstream fix lands.
In `@internal/tools/mcp.go`:
- Around line 726-729: The call to pollTaskUntilTerminal passes m.taskCfg and
also m.taskCfg.Progress redundantly; change pollTaskUntilTerminal (in
mcp_tasks.go: function pollTaskUntilTerminal) to remove the separate progress
MCPTaskProgressHandler parameter and have the function use cfg.Progress from the
MCPTaskConfig type (MCPTaskConfig and MCPTaskProgressHandler), then update all
call sites (including this call that currently passes mapping.serverName,
taskResult.Task, m.taskCfg, m.taskCfg.Progress) to omit the extra progress
argument so only cfg is passed; ensure any documentation/comments reflect that
pollTaskUntilTerminal uses cfg.Progress exclusively.
- Around line 673-691: When conn.client is not a *client.Client the code
silently falls back to synchronous execution; add a debug-level log before the
fallback that states the task capability is unavailable and we are falling back
to synchronous "tools/call" execution, including identifying info such as
mapping.serverName and the attempted task setting (e.g., callParams.Task or the
configured MCPTaskModeAlways), so operators can see why tasks weren’t used; keep
behavior identical otherwise and use the existing logger used in this scope for
the message.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 740cd4b5-2a35-4ea9-a0bc-1e1273e45f2d
📒 Files selected for processing (12)
internal/agent/agent.gointernal/agent/factory.gointernal/config/config.gointernal/config/config_test.gointernal/kitsetup/setup.gointernal/tools/connection_pool.gointernal/tools/mcp.gointernal/tools/mcp_tasks.gointernal/tools/mcp_tasks_test.gopkg/kit/kit.gopkg/kit/mcp_tasks.gopkg/kit/mcp_tasks_test.go
- README: add tasksMode YAML example and MCP Tasks subsection with SDK opt-in snippet - pkg/kit/README: add MCP Tasks subsection covering MCPTaskMode, progress callbacks, and List/Get/Cancel methods - www/configuration: document the tasksMode server field plus a per-mode behaviour table - www/sdk/options: extend the Compaction & MCP table with the six new Options fields and add a top-level MCP Tasks section - www/sdk/overview: add a brief MCP Tasks section between MCP prompts/resources and Context & compaction All examples verified against the public symbols in pkg/kit/mcp_tasks.go; docs site builds cleanly via npx tome build.
Address two review findings on the MCP Tasks PR. - Config.Validate() now rejects unknown tasksMode values with a clear error naming the server and bad value. Without this a typo (e.g. "alwasy") was silently downgraded to "auto" by the runtime parser. - Kit.Subagent() now propagates the parent's six MCP task options (mode map, timeout, TTL, poll interval, max poll interval, progress callback) onto the child via a new inheritMCPTaskOptions helper. Without this, child subagents always saw default polling and no progress feedback regardless of parent configuration. The propagation logic lives in a helper so the test exercises the real code path instead of duplicating it; future task fields only need to be added in one place.
11 tasks
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
1 participant
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.
Description
Implements Phase 1 of MCP Tasks (spec) so long-running MCP tools can run asynchronously, survive HTTP/SSE proxy timeouts, and be cancelled mid-flight. The full design and motivation are in #21 — this PR delivers the MVP slice (capability handshake, polling-based execution, context-aware cancel) and the public SDK surface; Phase 2 push notifications and Phase 3 UX/extension hooks are deferred to follow-ups.
The connection pool now advertises
mcp.NewTasksCapability()duringinitializeand captures the server'sInitializeResultso we can detect per-server task support.MCPToolManager.ExecuteToolaugmentstools/callwithTaskParamswhen policy + capability allow, pollstasks/get/tasks/resultuntil the task reaches a terminal state, and best-efforttasks/cancels on context cancellation. Defaults toMCPTaskModeAutoso any server that doesn't advertise task support keeps its existing synchronous behaviour bit-for-bit.The public SDK gets a clean dependency-agnostic surface (
MCPTask,MCPTaskStatus,MCPTaskMode,MCPTaskProgress,MCPTaskProgressHandler) onpkg/kit/, six newOptionsfields (MCPTaskMode,MCPTaskTimeout,MCPTaskTTL,MCPTaskPollInterval,MCPTaskMaxPollInterval,MCPTaskProgress), and threeKitmethods (ListMCPTasks,GetMCPTask,CancelMCPTask).mcp.jsonlearns one new field,tasksMode(auto/never/always), recognised in both new and legacy formats.Fixes #21
Type of Change
Checklist
go fmt,go vet,golangci-lint runall clean)go test ./... -race)Additional Information
Backward compatibility. Defaulting
TasksModetoautomeans servers that don't advertise task support duringinitializeget exactly the previous synchronous code path — no change in JSON wire shape, no change in result content typing. The new SDK fields are all zero-valued opt-ins.Files added
internal/tools/mcp_tasks.go— task modes, config, low-levelcallToolWithTask,pollTaskUntilTerminal,fetchTaskResultinternal/tools/mcp_tasks_test.go— capability advertisement/detection, mode parsing, end-to-end polling against an in-process server, progress emission, "never" mode, list/get/cancel error pathspkg/kit/mcp_tasks.go— public SDK types andKit.{List,Get,Cancel}MCPTaskmethodspkg/kit/mcp_tasks_test.go— terminality, options conversion, nil-Kit safetyFiles modified
internal/config/config.go(+ test) —MCPServerConfig.TasksModefield, parsed in both new and legacy unmarshal pathsinternal/tools/connection_pool.go— advertise tasks capability, captureInitializeResult, addSupportsToolTasks/ServerSupportsToolTasksinternal/tools/mcp.go— task-awareExecuteTool,SetTaskConfig,ListServerTasks/GetServerTask/CancelServerTaskinternal/agent/agent.go,internal/agent/factory.go,internal/kitsetup/setup.go— propagateMCPTaskConfigfrom SDK options through to the tool manager (including the lazyAddMCPServerpath)pkg/kit/kit.go— six newOptionsfields wired intosetupOptsUpstream workarounds (commented inline at the call sites)
mcp.ParseCallToolResultrequires acontentfield, so aCreateTaskResultresponse would error out — bypassed via raw transportSendRequest+ manual decode.mcp.ParseTaskResultResult(mcp-go v0.51.0) looks forcontentunder a nested"result"key that doesn't exist in the wire format, returning an emptyContentslice — same workaround.AddTaskTool. Real-world transports (stdio, SSE, streamable HTTP) keep the connection's context alive across the async work, so this only affects the test fixture; documented in the test's helper comment.Out of scope (per #21, deferred to follow-ups)
notifications/tasks/statusevent-driven updates (this PR uses polling)./tasksslash command,_meta.modelImmediateResponsehandling, and extension hooks (OnMCPTaskStarted/Completed/Cancelled).Summary by CodeRabbit
Release Notes