feat(server): implement task-augmented tools capability #707
+5,365
−59
Conversation
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
Add foundational types for task-augmented tool support per MCP spec 2025-11-25: - Add TaskSupport type (string enum) with three constants: - TaskSupportForbidden: tool cannot be invoked as a task (default) - TaskSupportOptional: tool can be invoked either way - TaskSupportRequired: tool must be invoked as a task - Add ToolExecution struct with TaskSupport field - Update Tool struct to include optional Execution field - Update Tool.MarshalJSON to serialize execution field when present - Add comprehensive tests for marshaling behavior and constants Related to #656
- Add Task.GetName() method to implement Named interface for pagination - Update handleListTasks to sort tasks by TaskId and use listByPagination - Add comprehensive pagination tests covering multiple pages and edge cases - Add TestTask_GetName unit test All tests pass with race detector, linter clean.
Add helper functions for associating messages with tasks per MCP spec: - RelatedTaskMetaKey constant for standard metadata key - RelatedTaskMeta() to create task metadata maps - WithRelatedTask() to create Meta objects with task IDs Includes comprehensive test coverage in mcp/tasks_test.go. Part of #656 (TAS-14)
Add WithTaskSupport() function to configure task support modes for tools. The function accepts TaskSupportForbidden, TaskSupportOptional, or TaskSupportRequired values and properly initializes or updates the tool's Execution field. - Implement WithTaskSupport() in mcp/tools.go - Add comprehensive test coverage in mcp/tools_test.go - Test all three task support modes with JSON marshaling - Test Execution initialization when nil - Test preservation of existing Execution instance Related to #656
…d tool calls - Add Task *TaskParams field to CallToolParams struct to support task augmentation - Add ToInt64Ptr() helper function in utils.go for creating int64 pointers - Add comprehensive tests for CallToolRequest with task params: - Test unmarshaling with task params, without task params, and null task params - Test round-trip marshaling/unmarshaling to verify data preservation - All tests pass with race detector, linter clean This enables detection of task-augmented tool calls and unblocks: - TAS-15: Support for optional task tools (hybrid mode) - TAS-16: Validation of task support on required tools Related to #656
Add LastUpdatedAt field to Task struct per MCP spec 2025-11-25 requirement. The field is automatically set on task creation and updated whenever the task status changes (completed, failed, or cancelled). Changes: - Add LastUpdatedAt string field to Task struct in mcp/types.go - Initialize both CreatedAt and LastUpdatedAt in NewTask() (mcp/tasks.go) - Update lastUpdatedAt in completeTask() when task completes or fails - Update lastUpdatedAt in cancelTask() when task is cancelled - Add comprehensive test coverage for timestamp behavior - Enhance JSON marshaling tests to verify field presence All tests pass, linter clean. Complies with MCP spec requirement that tasks MUST include lastUpdatedAt ISO 8601 timestamp.
Implement validation in handleToolCall() to enforce that tools with TaskSupportRequired cannot be called without task augmentation. Returns METHOD_NOT_FOUND error with message 'tool X requires task augmentation'. Add comprehensive test suite TestMCPServer_TaskSupportValidation with 7 test cases covering all task support modes (required, optional, forbidden). Implements MCP spec 2025-11-25 requirement for task support validation. Related to #656
Implement notifications/tasks/status to notify clients when task status changes. Notifications are sent automatically when tasks complete, fail, or are cancelled. Changes: - Add sendTaskStatusNotification() helper method that converts task data to notification params format - Include required fields: taskId, status, createdAt, lastUpdatedAt - Conditionally include optional fields: statusMessage, ttl, pollInterval - Integrate notifications into completeTask() and cancelTask() - Add comprehensive test suite with 7 test cases covering all scenarios Complies with MCP spec 2025-11-25 task notification requirements. Issue: #656
Add routing logic to detect when tools with TaskSupportOptional or TaskSupportRequired are called with task parameters. Tools called with task params now route to task-augmented execution path (placeholder error until TAS-9 implements handleTaskAugmentedToolCall). Tools called without task params continue to execute synchronously. - Add shouldExecuteAsTask detection in handleToolCall() - Add comprehensive test suite TestMCPServer_HybridModeDetection - Update existing tests to expect task execution routing behavior This provides the routing foundation for TAS-9 to implement actual task-augmented tool execution.
- Add detection logic in handleToolCall() to identify task-augmented requests - Route requests with task params to handleTaskAugmentedToolCall() when tool has TaskSupportOptional or TaskSupportRequired - Add handleTaskAugmentedToolCall() stub method with TODO for TAS-9 implementation - Completes TAS-8: routing infrastructure for task-augmented tool execution Related to #656
Add comprehensive example showing task-augmented tools with three types: - process_batch: TaskSupportRequired tool for batch processing - analyze_data: TaskSupportOptional tool with sync/async modes - quick_check: Regular synchronous tool for comparison Includes detailed README with usage examples, HTTP testing instructions, and task workflow explanation. Build constraint prevents compilation until core implementation (TAS-4 through TAS-10) is complete. Related to #656 (TAS-19)
…for async task execution - Add TaskToolHandlerFunc type for asynchronous tool call handlers that return CreateTaskResult - Add ServerTaskTool struct combining Tool with TaskToolHandlerFunc - Add taskTools map to MCPServer for storing task-augmented tools - Initialize taskTools map in NewMCPServer constructor - Add comprehensive test suite validating type definitions and initialization These type definitions provide the foundation for task-augmented tool execution (TAS-4, TAS-5). Related to #656
Implement executeTaskTool method that executes task tool handlers asynchronously. The method creates a cancellable context, stores the cancel function for potential cancellation via tasks/cancel, executes the handler in the background, and completes the task with result or error. Implementation details: - Creates cancellable context using context.WithCancel() - Stores cancel func in task entry under lock before execution - Executes TaskToolHandlerFunc and handles results/errors - Calls completeTask() to update status and send notifications Test coverage includes 6 comprehensive test cases: - Successful task execution stores result - Failed task execution stores error with proper status - Task cancellation via context - Cancel function storage timing (race condition test) - Task status notification on completion - Multiple concurrent task execution All tests pass with race detector, linter clean. This provides the core async execution mechanism for TAS-9 (handleTaskAugmentedToolCall). Related to #656
Implement handleTaskAugmentedToolCall method that enables async execution of tools with TaskSupport. The implementation: - Checks both taskTools and regular tools maps for tool lookup - Validates task support mode (optional/required) - Generates UUID v4 task IDs and creates task entries - Starts async execution via executeTaskTool or executeRegularToolAsTask - Returns CreateTaskResult immediately wrapped in CallToolResult _meta field - Fixes race condition by copying task before returning Add executeRegularToolAsTask method to handle hybrid mode where regular tools with TaskSupport are executed asynchronously. Update test expectations in TestMCPServer_HybridModeDetection and TestMCPServer_TaskSupportValidation to expect successful task creation instead of 'not yet implemented' error. Tools with TaskSupportOptional/Required can now be called with task params and execute asynchronously with proper task tracking.
- Add Task-Augmented Tools section explaining async execution - Document three TaskSupport modes (Forbidden, Optional, Required) - Include code examples for TaskSupportRequired and TaskSupportOptional - Explain task execution flow (call → task ID → async → polling → notifications) - Add context cancellation and WithTaskCapabilities examples - Update Examples section to highlight task_tool example - Reorganize synchronous tools section for clarity
Update handleListTools to return both regular and task-augmented tools in tools/list response. Implementation iterates over both s.tools and s.taskTools maps, sorts all tool names alphabetically, and includes both types in the result. - Extend capacity of tools slice to accommodate both tool types - Collect tool names from both s.tools and s.taskTools maps - Maintain alphabetical sorting across all tools - Protected by existing toolsMu mutex for thread safety - Add comprehensive test suite with 4 test cases: * Mixed regular and task tools with proper sorting * Only task tools scenario * Empty list when no tools exist * Alphabetical ordering verification Task tools identified by Execution.TaskSupport field in metadata.
Implement proper API for registering task-augmented tools following the same pattern as regular tool registration (AddTool/AddTools). - Add AddTaskTool() method for single task tool registration - Add AddTaskTools() method for batch task tool registration - Both methods implicitly register tool capabilities (listChanged=true) - Send tools/list_changed notifications to initialized sessions - Add comprehensive test suite TestMCPServer_AddTaskTool with 4 test cases - Update existing tests to use new methods instead of direct map access All tests pass with race detector, linter clean (0 issues).
Add TestTaskToolTracerBullet in server/task_tool_test.go with 6 end-to-end test scenarios covering: - Complete task tool flow with TaskSupportRequired - TaskSupportOptional sync/async execution modes - Task cancellation via context - Error handling with proper task failure states - Multiple concurrent task tools Fix bug in handleToolCall where taskTools map was not checked during tool lookup. Now checks both s.tools and s.taskTools, allowing tools registered with AddTaskTool to be properly found and validated. All tests pass with race detector, linter reports 0 issues.
- Add explicit checking for context.Canceled and context.DeadlineExceeded errors in executeTaskTool and executeRegularToolAsTask - Mark tasks as 'cancelled' instead of 'failed' when handlers return context cancellation errors - Handle race condition where handler detects cancellation before tasks/cancel is called - Add comprehensive test for handler-detected context cancellation edge case - Update existing test expectations to correctly expect TaskStatusCancelled This ensures proper status differentiation between cancellation and failure, respecting Go's context patterns.
Implement comprehensive observability system for task-augmented tools: - Add TaskMetrics struct with execution metadata (ID, tool name, status, timestamps, duration, session, error) - Add TaskHooks with 5 lifecycle hook types: OnTaskCreated, OnTaskCompleted, OnTaskFailed, OnTaskCancelled, OnTaskStatusChanged - Add WithTaskHooks() ServerOption for hook configuration - Integrate hooks into task lifecycle (createTask, completeTask, cancelTask, executeTaskTool, executeRegularToolAsTask) - Update taskEntry to track toolName and createdAt for metrics - Add comprehensive test suite with 8 test cases covering all hook scenarios - Update task_tool example to demonstrate hook usage with logging This enables developers to monitor task execution, collect metrics, and integrate with monitoring systems for async tool operations. Resolves TAS-24
Add optional resource management for task-augmented tools through WithMaxConcurrentTasks() ServerOption. This prevents resource exhaustion by limiting the number of concurrent running tasks. Changes: - Add maxConcurrentTasks and activeTasks fields to MCPServer - Add WithMaxConcurrentTasks(limit int) ServerOption - Modify createTask() to check limit and return error when exceeded - Update task completion/cancellation to decrement active task counter - Add comprehensive test suite (8 test cases) with race detector - Update examples and documentation The limit only applies to non-terminal (working) tasks. Completed, failed, or cancelled tasks don't count toward the limit. When not configured or set to 0, no limit is enforced. All tests pass with race detector. Linter clean.
Contributor
WalkthroughImplements task-augmented tools and server-side Tasks capability: adds TaskSupport modes, task lifecycle (create/list/result/cancel), concurrency limits, task hooks/observability, TTL/expiration, related metadata helpers, hook signature updates, examples, and extensive tests across server and mcp packages. Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested labels
Suggested reviewers
🚥 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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 8
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
server/server.go (1)
2157-2211:⚠️ Potential issue | 🟠 MajorHooks called while holding lock may cause blocking or deadlock.
completeTaskholdss.tasksMulock (viadefer s.tasksMu.Unlock()at Line 2160) while:
- Calling
s.sendTaskStatusNotification(Line 2190)- Calling
s.taskHooks.taskFailed/taskCompleted(Lines 2207-2209)If hooks perform lengthy operations or try to access task state (which would require the same lock), this could cause:
- Performance degradation (blocking other task operations)
- Potential deadlock if hooks call back into task management
Consider releasing the lock before firing hooks:
🔒 Proposed fix to release lock before hooks
func (s *MCPServer) completeTask(entry *taskEntry, result any, err error) { s.tasksMu.Lock() - defer s.tasksMu.Unlock() // Guard against double completion if entry.completed { + s.tasksMu.Unlock() return } completedAt := time.Now() duration := completedAt.Sub(entry.createdAt) if err != nil { entry.task.Status = mcp.TaskStatusFailed entry.task.StatusMessage = err.Error() entry.resultErr = err } else { entry.task.Status = mcp.TaskStatusCompleted entry.result = result } entry.task.LastUpdatedAt = completedAt.UTC().Format(time.RFC3339) entry.completed = true close(entry.done) s.activeTasks-- + // Copy data needed for hooks before releasing lock + taskCopy := entry.task + toolName := entry.toolName + createdAt := entry.createdAt + sessionID := entry.sessionID + s.tasksMu.Unlock() + // Send task status notification - s.sendTaskStatusNotification(entry.task) + s.sendTaskStatusNotification(taskCopy) // Fire task hooks if s.taskHooks != nil { metrics := TaskMetrics{ - TaskID: entry.task.TaskId, - ToolName: entry.toolName, - Status: entry.task.Status, - StatusMessage: entry.task.StatusMessage, - CreatedAt: entry.createdAt, + TaskID: taskCopy.TaskId, + ToolName: toolName, + Status: taskCopy.Status, + StatusMessage: taskCopy.StatusMessage, + CreatedAt: createdAt, CompletedAt: &completedAt, Duration: duration, - SessionID: entry.sessionID, + SessionID: sessionID, Error: err, } // ... rest of hook calls } }
🤖 Fix all issues with AI agents
In `@examples/task_tool/README.md`:
- Line 112: The README's workflow phrase "pending → completed or failed" is
inconsistent with the code's TaskStatus constant named TaskStatusWorking; update
the README to use "working → completed or failed" (or alternatively rename the
constant to TaskStatusPending in code) so terminology matches the
implementation; locate references to TaskStatusWorking in the code and ensure
the README's status terms exactly match that constant name (or change the
constant and all its usages if you choose to rename it).
- Around line 17-26: Update the README checklist to mark TAS-4 through TAS-10 as
completed and remove the build constraint in main.go: change TAS-4
TaskToolHandlerFunc, TAS-5 ServerTaskTool, TAS-6 AddTaskTool, TAS-9
handleTaskAugmentedToolCall, and TAS-10 executeTaskTool entries to ✅ in
README.md and delete the "//go:build ignore" line from main.go so the program
can be built now that those symbols are implemented.
In `@README.md`:
- Around line 315-316: The sentence in README.md uses the noun form "timeout"
where the verb phrase "time out" is correct; update the sentence that currently
reads "would otherwise block or timeout" to "would otherwise block or time out"
so the verb form is used; ensure the change is made in the paragraph starting
"Task-augmented tools execute asynchronously and return results via polling" to
preserve grammar and spacing.
In `@server/server.go`:
- Around line 2037-2051: The concurrency check and increment in createTask has a
race because s.tasksMu is unlocked too early; hold the mutex for the full
critical section by calling s.tasksMu.Lock() and deferring s.tasksMu.Unlock()
(or explicitly unlocking only after you create/insert the task and increment
s.activeTasks) so the check against s.maxConcurrentTasks and the s.activeTasks++
happen atomically; update any corresponding decrement logic (where s.activeTasks
is decremented) to also use s.tasksMu to keep counts consistent.
- Around line 1545-1600: The code creates a ServerTaskTool wrapper (assigning to
toolToUse with a TaskToolHandlerFunc) but never uses it when hasTaskHandler is
false, causing dead code or wrong behavior; either remove the unused wrapper
block (the ServerTaskTool creation and any TaskToolHandlerFunc code) or change
the async call to use the wrapper by replacing the call to
executeRegularToolAsTask(ctx, entry, regularTool, request) with
executeRegularToolAsTask(ctx, entry, toolToUse, request) (or adjust
executeRegularToolAsTask signature to accept ServerTaskTool), ensuring
references to toolToUse, ServerTaskTool, TaskToolHandlerFunc, hasTaskHandler and
executeRegularToolAsTask/regularTool are updated consistently.
- Around line 1324-1347: Add validation to prevent registering the same tool
name in both s.tools and s.taskTools: in AddTools (and AddTaskTools /
AddTaskTool) check each incoming tool name against the opposite map (e.g., when
adding to s.tools, ensure the name is not already present in s.taskTools, and
vice versa) and return or report an error if a collision is detected instead of
allowing both to be registered; this ensures handleListTools (which currently
prefers s.tools over s.taskTools) won't silently hide task tools and keeps
registration atomic and deterministic.
In `@server/task_hooks_test.go`:
- Around line 19-35: The tests in server/task_hooks_test.go use time.Sleep to
wait for hook execution (e.g., after calling server.createTask in the
TaskHooks/AddOnTaskCreated tests), which causes flakiness; replace each fixed
sleep with a synchronization primitive such as a channel or sync.WaitGroup plus
a timeout context: have the hook callback signal completion (close a done
channel or call wg.Done), then the test should wait on that channel or wg with a
select that also listens on a time.After timeout and fail the test if the
timeout fires. Update all occurrences around TaskHooks/AddOnTaskCreated,
Create/Update/Delete hook tests (the hooks variable,
AddOnTaskCreated/AddOnTaskUpdated/AddOnTaskDeleted callbacks and where
server.createTask/server.updateTask/server.deleteTask is invoked) to use this
pattern rather than time.Sleep.
In `@server/task_tool_test.go`:
- Around line 59-180: The test has a data race: the async task handler writes to
handlerCalled and receivedData while the main goroutine reads them; protect
these shared variables (handlerCalled, receivedData) — e.g., change
handlerCalled to an int32 and use sync/atomic
(atomic.StoreInt32/atomic.LoadInt32) and protect receivedData with either a
sync.Mutex or atomic.Value, or use a channel to signal/transfer the string from
the handler; update the long_running tool handler in the test (the anonymous
func registered with server.AddTaskTool) and the assertion reads to use the
chosen synchronization primitives so reads/writes are race-free (also update
assertions to check atomic/mutex-protected values).
🧹 Nitpick comments (6)
server/task_test.go (2)
617-620: Consider usingrequire.Truefor type assertions to avoid panics.Lines 618-619 perform type assertions without checking the
okvalue. If these assertions fail, the test will panic rather than providing a clear failure message.🔧 Suggested fix
- resp := response.(mcp.JSONRPCResponse) - firstPage := resp.Result.(mcp.ListTasksResult) + resp, ok := response.(mcp.JSONRPCResponse) + require.True(t, ok, "Expected JSONRPCResponse") + firstPage, ok := resp.Result.(mcp.ListTasksResult) + require.True(t, ok, "Expected ListTasksResult")
756-757: Sleep duration for timestamp tests is necessary but slow.The 1100ms sleep ensures RFC3339 timestamps (second precision) will differ. This is correct but makes the test suite slower. Consider adding a comment explaining why this duration is required, or using a test helper that can mock time if test speed becomes a concern.
server/task_hooks_test.go (1)
15-15: Add GoDoc comments for exported TestTaskHooks_ funcs.*These exported test functions don’t have GoDoc comments starting with the identifier name. As per coding guidelines: All exported types and functions must have GoDoc comments starting with the identifier name; avoid inline comments unless necessary.
server/task_tool_test.go (1)
123-123: Reuse mcp.ToInt64Ptr instead of a local helper.There’s already a shared helper (
mcp.ToInt64Ptr) that can replaceptrInt64, avoiding duplication.Possible simplification
- TTL: ptrInt64(60000), + TTL: mcp.ToInt64Ptr(60000), @@ -// ptrInt64 is a helper to get a pointer to an int64 -func ptrInt64(i int64) *int64 { - return &i -}Also applies to: 625-628
server/task_hooks.go (1)
55-78: Consider adding thread-safety documentation or synchronization.The
Add*methods modify slices without synchronization. While hooks are typically registered during server setup before any tasks are processed, concurrent registration during runtime would cause a data race.Consider either:
- Adding a comment documenting that hooks must be registered before the server starts processing requests.
- Adding mutex protection if runtime registration is intended to be supported.
📝 Option 1: Add documentation
+// AddOnTaskCreated registers a hook for task creation events. +// Note: Hooks should be registered during server initialization, before +// the server starts processing requests, to avoid data races. func (h *TaskHooks) AddOnTaskCreated(hook OnTaskCreatedHookFunc) { h.OnTaskCreated = append(h.OnTaskCreated, hook) }server/server.go (1)
1641-1688: Significant code duplication in cancellation handling.The cancellation handling logic in
executeTaskTool(Lines 1641-1688) andexecuteRegularToolAsTask(Lines 1720-1767) is nearly identical (~50 lines). This violates DRY and makes maintenance error-prone.Consider extracting a helper method:
♻️ Proposed refactor to extract common logic
// handleTaskContextCancellation handles the case where a task handler detects // context cancellation before tasks/cancel was called. func (s *MCPServer) handleTaskContextCancellation(ctx context.Context, entry *taskEntry, err error) { s.tasksMu.Lock() alreadyCancelled := entry.task.Status == mcp.TaskStatusCancelled s.tasksMu.Unlock() if alreadyCancelled { return } cancelledAt := time.Now() duration := cancelledAt.Sub(entry.createdAt) s.tasksMu.Lock() if !entry.completed { entry.task.Status = mcp.TaskStatusCancelled entry.task.StatusMessage = err.Error() entry.task.LastUpdatedAt = cancelledAt.UTC().Format(time.RFC3339) entry.completed = true close(entry.done) s.activeTasks-- s.sendTaskStatusNotification(entry.task) if s.taskHooks != nil { metrics := TaskMetrics{ TaskID: entry.task.TaskId, ToolName: entry.toolName, Status: entry.task.Status, StatusMessage: entry.task.StatusMessage, CreatedAt: entry.createdAt, CompletedAt: &cancelledAt, Duration: duration, SessionID: entry.sessionID, } s.taskHooks.taskCancelled(ctx, metrics) } } s.tasksMu.Unlock() }Then both
executeTaskToolandexecuteRegularToolAsTaskcan call:if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { s.handleTaskContextCancellation(ctx, entry, err) return }Also applies to: 1720-1767
…collision validation - Fix race window in createTask() between checking concurrent task limit and inserting task - Restructure to build task entry first without lock - Use single critical section with defer for check + increment + insert - Prevents multiple goroutines from exceeding maxConcurrentTasks under high concurrency - Add tool name collision validation - Prevent same tool name from being registered in both s.tools and s.taskTools - AddTools() now checks for collisions with s.taskTools before registration - AddTaskTools() now checks for collisions with s.tools before registration - Panic with clear error message on collision to fail fast during setup - Prevents silent shadowing where task tools become invisible in handleListTools - Add comprehensive test coverage - Test concurrent task creation respects limit (existing test now passes reliably) - Add TestMCPServer_ToolNameCollisionValidation with 5 collision scenarios - All tests pass with race detector enabled
Shared variables handlerCalled and receivedData were accessed without synchronization from both the async task handler goroutine and the main test thread, causing data races detected by 'go test -race'. Added sync.Mutex protection around all reads and writes to these variables. All tests now pass with race detector enabled.
…iance Task-augmented tool calls now return CreateTaskResult with task as a direct field of result (result.task) instead of nested in _meta (result._meta.task), matching the MCP Tasks Specification (2025-11-25). Changes: - Modified handleToolCall() to return 'any' to support both CallToolResult and CreateTaskResult - Modified handleTaskAugmentedToolCall() to return *CreateTaskResult directly - Updated code generation templates to support ResultIsAny flag for flexible return types - Regenerated hooks.go and request_handler.go with updated templates - Fixed all tests to use pointer type assertions and expect correct format - Added comprehensive spec compliance tests All tests pass with race detector enabled. Resolves: #707 (TAS-1)
Updated doc comment spec version references for all task-related methods from specification/draft to specification/2025-11-25 for spec compliance. Changes: - MethodTasksGet - MethodTasksList - MethodTasksResult - MethodTasksCancel - MethodNotificationTasksStatus All 5 URLs now reference the published 2025-11-25 spec version.
- Add documentation to TaskStatusInputRequired explaining it's not yet implemented - Remove unused wrapper code in handleTaskAugmentedToolCall that was never executed - The wrapper was created but never used when hasTaskHandler=false Related to #707
Replace all time.Sleep calls in task_hooks_test.go with channels and sync.WaitGroup for reliable test synchronization. Update 7 test functions to use select/timeout pattern for single-event tests and WaitGroup for multi-event tests, eliminating race conditions and timing-dependent flakiness. Additionally: - Add model-immediate-response support with helper functions and tests - Improve task error differentiation (expired vs not found) - Remove unused wrapper code in handleTaskAugmentedToolCall - Add comprehensive test coverage for new functionality All tests pass with race detector enabled.
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@server/task_tool_test.go`:
- Around line 185-195: The test currently reads a single notification from
notifyChan and asserts it is mcp.TaskStatusCompleted, which can fail if an
earlier "working" status is received; change the read to loop (or drain) from
notifyChan until a notification with params["taskId"] == taskID and
params["status"] == mcp.TaskStatusCompleted is observed or a timeout elapses.
Use the existing notifyChan, taskID, mcp.TaskStatusCompleted and
statusNotification variables (and time.After for timeout) to repeatedly select
on notifyChan and the timeout, ignore non-matching/working notifications, and
only assert success when the completed status is seen; on timeout call t.Fatal
as before.
🧹 Nitpick comments (5)
server/task_response_format_test.go (1)
13-120: Prefer table-driven tests for these cases. The two similar tests can be merged into a single table-driven test to follow the repo testing guideline and reduce duplicated setup.♻️ Possible refactor sketch
- func TestTaskAugmentedToolCall_ResponseFormat(t *testing.T) { ... } - func TestTaskAugmentedToolCall_SpecCompliance(t *testing.T) { ... } + func TestTaskAugmentedToolCall_ResponseFormat(t *testing.T) { + tests := []struct { + name string + request string + validate func(t *testing.T, resp mcp.JSONRPCMessage) + }{ + { name: "response format", request: `...`, validate: func(t *testing.T, resp mcp.JSONRPCMessage) { /* current assertions */ }}, + { name: "spec compliance", request: `...`, validate: func(t *testing.T, resp mcp.JSONRPCMessage) { /* current assertions */ }}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // shared server/tool setup + response := server.HandleMessage(context.Background(), []byte(tt.request)) + tt.validate(t, response) + }) + } + }As per coding guidelines: Write table-driven tests using a tests := []struct{ name, ... } pattern.
server/task_test.go (1)
811-812: Consider reducing test execution time.The 1100ms sleeps (to ensure RFC3339 timestamps differ) make this test ~3.3 seconds slower. While necessary for second-precision timestamps, consider:
- Using a clock abstraction for faster unit tests
- Marking this as an integration test if speed becomes a concern
Current approach is correct but noted for future optimization.
server/server.go (3)
1608-1765: Significant code duplication between executeTaskTool and executeRegularToolAsTask.These two functions share ~80% identical code, particularly the cancellation handling logic (Lines 1633-1676 vs 1712-1755). Consider extracting the common pattern into a helper function.
♻️ Suggested refactoring approach
// executeToolAsync is a helper that handles common task execution logic func (s *MCPServer) executeToolAsync( ctx context.Context, entry *taskEntry, handler func(ctx context.Context) (any, error), ) { taskCtx, cancel := context.WithCancel(ctx) defer cancel() s.tasksMu.Lock() entry.cancelFunc = cancel s.tasksMu.Unlock() result, err := handler(taskCtx) if err != nil { if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { s.handleTaskCancellation(ctx, entry, err) return } s.completeTask(entry, nil, err) return } s.completeTask(entry, result, nil) }Then
executeTaskToolandexecuteRegularToolAsTaskbecome thin wrappers that just call this helper with the appropriate handler.
2060-2071: Task hooks called while holding tasksMu lock may cause contention.If
taskCreatedhooks are slow or perform I/O, they will block other task operations. Consider calling hooks outside the critical section:♻️ Proposed improvement
// Increment active task counter and insert task atomically s.activeTasks++ s.tasks[taskID] = entry + s.tasksMu.Unlock() // Fire task created hook if s.taskHooks != nil { metrics := TaskMetrics{ TaskID: taskID, ToolName: toolName, Status: task.Status, CreatedAt: createdAt, SessionID: getSessionID(ctx), } s.taskHooks.taskCreated(ctx, metrics) } - s.tasksMu.Unlock() + s.tasksMu.Lock() // Start TTL cleanup if specifiedAlternatively, document that hook implementations must be non-blocking.
2277-2284: Consider batching tombstone cleanup to reduce goroutine count.Each expired task spawns a dedicated goroutine for tombstone cleanup. With high task churn, this could lead to many concurrent goroutines. Consider:
- A single background cleanup goroutine that periodically sweeps
expiredTasks- Store expiration time and clean up during other operations
Current approach works but may not scale well with high task volumes.
- Update task tool test to drain notifyChan until completed status is received, preventing flaky failures when working status arrives first - Fix grammar in README: 'timeout' -> 'time out' (verb form)
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.
Description
This PR implements comprehensive support for task-augmented tools in the MCP Go SDK, allowing tools to execute asynchronously and report progress through the tasks subsystem.
Fixes #656
Type of Change
Checklist
MCP Spec Compliance
Key Features
Core Implementation
Task Management
Observability & Limits
Documentation & Examples
examples/task_tool/Additional Information
This implementation enables long-running tool operations to:
Summary by CodeRabbit
New Features
Documentation
Tests
✏️ Tip: You can customize this high-level summary in your review settings.