mcp-go icon indicating copy to clipboard operation
mcp-go copied to clipboard
verified publisher icon mark3labs

Add server-side support for MCP tasks

Open JAORMX opened this issue 1 month ago • 1 comments

This PR implements server-side support for the MCP tasks specification draft.

What's Included

  • Task types and protocol definitions (TaskStatus, request/result types, capabilities)
  • Server-side task storage with session isolation and TTL-based cleanup
  • Handler functions for tasks/get, tasks/list, tasks/result, and tasks/cancel
  • Capability negotiation for task support
  • Code generation integration (hooks and request handlers)
  • Comprehensive test coverage

What's Not Included

  • Client-side task support
  • Task-augmented tool call implementation (two-phase response pattern)
  • Examples demonstrating task usage

This PR focuses on the foundational server infrastructure. The client-side implementation and examples can be added in follow-up PRs.

Implementation Details

All task operations follow existing patterns in the codebase:

  • Functional options for configuration
  • Thread-safe operations with proper mutex usage
  • Session-based isolation for security
  • Integration with the code generation system

Testing

Added comprehensive tests covering:

  • Capability negotiation
  • Task lifecycle management
  • All handler functions
  • Error handling and edge cases
  • TTL cleanup and concurrent operations

All tests pass.

Notes

I can break this into smaller PRs if that would be easier to review. Let me know if you'd prefer that approach.

Summary by CodeRabbit

  • New Features
    • Task management: create, retrieve, list, cancel tasks with multi-state tracking (working, input required, completed, failed, cancelled).
    • Task notifications for status updates and task result retrieval.
    • Server-configurable task capabilities (list, cancel, tool-call) and TTL-based automatic task cleanup.
  • Tests
    • Comprehensive test suite covering lifecycle, listing, cancellation, TTL cleanup, notifications, and status handling.

✏️ Tip: You can customize this high-level summary in your review settings.

JAORMX avatar Nov 18 '25 15:11 JAORMX

Walkthrough

Adds first-class task support: new MCP task types, functional helpers, server-side per-session task management with TTL/cleanup and capability negotiation, request handlers and hooks for task operations, and comprehensive tests exercising lifecycle, cancellation, listing, and notifications.

Changes

Cohort / File(s) Summary
MCP task types & helpers
mcp/types.go, mcp/tasks.go		 Adds TaskStatus, Task, TaskParams, Create/Get/List/Cancel/Result types, TaskStatusNotification, MCP method/notification constants, Marshal/Unmarshal content helpers, functional-option TaskOption builders, and constructors for task results/capabilities.
Server task subsystem
server/server.go		 Adds taskEntry, per-server tasks map and tasksMu, task lifecycle functions (create/get/list/complete/cancel), TTL-based cleanup scheduling, session-aware task isolation, and WithTaskCapabilities ServerOption plus initialize-time capability emission.
Server hooks
server/hooks.go		 Adds per-operation hook types and slices for GetTask/ListTasks/TaskResult/CancelTask, registration methods, and before/after lifecycle invokers wired into existing hook patterns.
Request handling
server/request_handler.go		 Adds cases for tasks/get, tasks/list, tasks/result, and tasks/cancel with capability checks, unmarshalling, header propagation, before/after hook invocations, handler dispatch, and JSON-RPC error conversion.
Tests
server/task_test.go		 Adds tests for capability negotiation, task creation/retrieval/completion/cancellation, TTL expiry and cleanup, terminal-state detection, result waiting, error propagation, and JSON marshal/unmarshal round trips.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

  • Pay extra attention to concurrency and synchronization in server/server.go (use of tasksMu, per-task done channel, cancelFunc, and completed flag).
  • Verify correctness and timing of TTL cleanup and scheduleTaskCleanup.
  • Inspect hook ordering/nil-safety and error propagation in server/hooks.go and server/request_handler.go.
  • Review tests in server/task_test.go for potential timing flakiness and edge-case coverage.

Possibly related issues

  • mark3labs/mcp-go#656 — Matches: implements Tasks capability end-to-end (types, server handlers, hooks, tests) referenced by the issue.

Possibly related PRs

  • mark3labs/mcp-go#164 — Related changes to server hooks and request dispatch; likely overlaps with hook wiring added here.
  • mark3labs/mcp-go#62 — Also extends server hooks/request handling; overlaps with lifecycle hook additions.

Suggested labels

type: enhancement, area: mcp spec, status: needs submitter response

Suggested reviewers

  • robert-jackson-glean
  • pottekkat
  • rwjblue-glean

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)
Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 57.14% which is insufficient. The required threshold is 80.00%. You can run @coderabbitai generate docstrings to improve docstring coverage.
✅ Passed checks (2 passed)
Check name Status Explanation
Title check ✅ Passed The title accurately captures the main change: adding server-side support for MCP tasks.
Description check ✅ Passed The description addresses most required sections including type of change, implementation details, testing, and notes, though Type of Change checkbox is not explicitly marked.
✨ Finishing touches
  • [ ] 📝 Generate docstrings
🧪 Generate unit tests (beta)
  • [ ] Create PR with unit tests
  • [ ] Post copyable unit tests in a comment
📜 Recent review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between b73e4b94c5b28ab1dd225c43bd860cdc959f4a61 and db65292b6bc9246a29928cd448a90469d65c8b41.

⛔ Files ignored due to path filters (1)
  • server/internal/gen/data.go is excluded by !**/gen/**
📒 Files selected for processing (6)
  • mcp/tasks.go (1 hunks)
  • mcp/types.go (5 hunks)
  • server/hooks.go (3 hunks)
  • server/request_handler.go (1 hunks)
  • server/server.go (10 hunks)
  • server/task_test.go (1 hunks)
🧰 Additional context used
📓 Path-based instructions (2)
**/*.go

📄 CodeRabbit inference engine (AGENTS.md)

**/*.go: Order imports: standard library first, then third-party, then local packages (goimports enforces this) Follow Go naming conventions: exported identifiers in PascalCase; unexported in camelCase; acronyms uppercase (HTTP, JSON, MCP) Error handling: return sentinel errors, wrap with fmt.Errorf("context: %w", err), and check with errors.Is/As Prefer explicit types and strongly-typed structs; avoid using any except where protocol flexibility is required (e.g., Arguments any) All exported types and functions must have GoDoc comments starting with the identifier name; avoid inline comments unless necessary Functions that are handlers or long-running must accept context.Context as the first parameter Ensure thread safety for shared state using sync.Mutex and document thread-safety requirements in comments For JSON: use json struct tags with omitempty for optional fields; use json.RawMessage for flexible/deferred parsing

Files:

  • server/request_handler.go
  • server/task_test.go
  • mcp/tasks.go
  • server/server.go
  • mcp/types.go
  • server/hooks.go
**/*_test.go

📄 CodeRabbit inference engine (AGENTS.md)

**/*_test.go: Testing: use testify/assert and testify/require Write table-driven tests using a tests := []struct{ name, ... } pattern Go test files must end with _test.go

Files:

  • server/task_test.go
🧠 Learnings (5)
📚 Learning: 2025-06-30T07:13:17.052Z 
Learnt from: ezynda3
Repo: mark3labs/mcp-go PR: 461
File: server/sampling.go:22-26
Timestamp: 2025-06-30T07:13:17.052Z
Learning: In the mark3labs/mcp-go project, the MCPServer.capabilities field is a struct value (serverCapabilities), not a pointer, so it cannot be nil and doesn't require nil checking. Only pointer fields within the capabilities struct should be checked for nil.

Applied to files:

  • server/request_handler.go
  • server/task_test.go
  • mcp/tasks.go
  • server/server.go
  • mcp/types.go
📚 Learning: 2025-04-21T21:26:32.945Z 
Learnt from: octo
Repo: mark3labs/mcp-go PR: 149
File: mcptest/mcptest.go:0-0
Timestamp: 2025-04-21T21:26:32.945Z
Learning: In the mcptest package, prefer returning errors from helper functions rather than calling t.Fatalf() directly, giving callers flexibility in how to handle errors.

Applied to files:

  • server/task_test.go
📚 Learning: 2025-03-04T07:00:57.111Z 
Learnt from: xinwo
Repo: mark3labs/mcp-go PR: 35
File: mcp/tools.go:0-0
Timestamp: 2025-03-04T07:00:57.111Z
Learning: The Tool struct in the mark3labs/mcp-go project should handle both InputSchema and RawInputSchema consistently between MarshalJSON and UnmarshalJSON methods, even though the tools response from MCP server typically doesn't contain rawInputSchema.

Applied to files:

  • server/task_test.go
📚 Learning: 2025-03-04T06:59:43.882Z 
Learnt from: xinwo
Repo: mark3labs/mcp-go PR: 35
File: mcp/tools.go:107-137
Timestamp: 2025-03-04T06:59:43.882Z
Learning: Tool responses from the MCP server shouldn't contain RawInputSchema, which is why the UnmarshalJSON method for the Tool struct is implemented to handle only the structured InputSchema format.

Applied to files:

  • server/task_test.go
📚 Learning: 2025-10-13T09:35:20.180Z 
Learnt from: CR
Repo: mark3labs/mcp-go PR: 0
File: AGENTS.md:0-0
Timestamp: 2025-10-13T09:35:20.180Z
Learning: Applies to **/*.go : Ensure thread safety for shared state using sync.Mutex and document thread-safety requirements in comments

Applied to files:

  • server/server.go
🧬 Code graph analysis (5)
server/request_handler.go (3)
mcp/types.go (14)
  • MethodTasksGet (68-68)
  • GetTaskRequest (1325-1329)
  • GetTaskResult (1336-1339)
  • METHOD_NOT_FOUND (404-404)
  • INVALID_REQUEST (401-401)
  • MethodTasksList (72-72)
  • ListTasksRequest (1342-1345)
  • ListTasksResult (1348-1351)
  • MethodTasksResult (76-76)
  • TaskResultRequest (1354-1358)
  • TaskResultResult (1366-1370)
  • MethodTasksCancel (80-80)
  • CancelTaskRequest (1373-1377)
  • CancelTaskResult (1384-1387)
server/errors.go (1)
  • ErrUnsupported (10-10)
server/server.go (1)
  • UnparsableMessageError (103-107)
server/task_test.go (3)
server/server.go (3)
  • ServerOption (43-43)
  • WithTaskCapabilities (360-369)
  • NewMCPServer (379-408)
mcp/types.go (15)
  • InitializeResult (477-491)
  • TaskStatusWorking (1278-1278)
  • TaskStatusCompleted (1282-1282)
  • GetTaskResult (1336-1339)
  • JSONRPCError (376-380)
  • INVALID_PARAMS (407-407)
  • ListTasksResult (1348-1351)
  • TaskStatusCancelled (1286-1286)
  • METHOD_NOT_FOUND (404-404)
  • TaskStatus (1274-1274)
  • TaskStatusInputRequired (1280-1280)
  • TaskStatusFailed (1284-1284)
  • JSONRPCMessage (129-129)
  • TaskResultResult (1366-1370)
  • Task (1295-1309)
mcp/tasks.go (8)
  • NewTask (18-30)
  • WithTaskStatus (33-37)
  • WithTaskStatusMessage (40-44)
  • WithTaskTTL (48-52)
  • WithTaskPollInterval (55-59)
  • NewTaskParams (74-78)
  • NewTasksCapability (125-137)
  • NewTasksCapabilityWithToolsOnly (141-151)
mcp/tasks.go (1)
mcp/types.go (15)
  • Task (1295-1309)
  • TaskStatusWorking (1278-1278)
  • TaskStatus (1274-1274)
  • TaskParams (1312-1315)
  • CreateTaskResult (1319-1322)
  • GetTaskResult (1336-1339)
  • ListTasksResult (1348-1351)
  • CancelTaskResult (1384-1387)
  • TaskStatusNotification (1390-1393)
  • Notification (210-213)
  • MethodNotificationTasksStatus (102-102)
  • Params (208-208)
  • TaskStatusNotificationParams (1395-1397)
  • TasksCapability (1245-1252)
  • TaskRequestsCapability (1255-1271)
server/server.go (3)
mcp/types.go (16)
  • Task (1295-1309)
  • TasksCapability (1245-1252)
  • TaskRequestsCapability (1255-1271)
  • GetTaskRequest (1325-1329)
  • GetTaskResult (1336-1339)
  • Params (208-208)
  • ListTasksRequest (1342-1345)
  • ListTasksResult (1348-1351)
  • TaskResultRequest (1354-1358)
  • TaskResultResult (1366-1370)
  • Result (277-281)
  • CancelTaskRequest (1373-1377)
  • CancelTaskResult (1384-1387)
  • TaskStatusFailed (1284-1284)
  • TaskStatusCompleted (1282-1282)
  • TaskStatusCancelled (1286-1286)
mcp/tasks.go (3)
  • TaskOption (9-9)
  • WithTaskTTL (48-52)
  • WithTaskPollInterval (55-59)
server/session.go (1)
  • ClientSessionFromContext (112-117)
server/hooks.go (1)
mcp/types.go (12)
  • GetTaskRequest (1325-1329)
  • GetTaskResult (1336-1339)
  • ListTasksRequest (1342-1345)
  • ListTasksResult (1348-1351)
  • TaskResultRequest (1354-1358)
  • TaskResultResult (1366-1370)
  • CancelTaskRequest (1373-1377)
  • CancelTaskResult (1384-1387)
  • MethodTasksGet (68-68)
  • MethodTasksList (72-72)
  • MethodTasksResult (76-76)
  • MethodTasksCancel (80-80)
🔇 Additional comments (26)
server/hooks.go (3)

94-104: LGTM!

The hook type definitions follow the existing pattern consistently and use the correct MCP task types from mcp/types.go.

133-140: LGTM!

The hook slice fields are properly declared and follow the established pattern used for other operations in the Hooks struct.

553-660: LGTM!

The lifecycle methods follow the existing pattern precisely:

  • Proper nil-safety checks
  • beforeAny/onSuccess invocation in correct order
  • Nil-safe iteration through hook slices

All task-related hook functions are implemented consistently with the rest of the codebase.

server/request_handler.go (4)

332-357: LGTM!

The tasks/get handler follows the established pattern for MCP method handling:

  • Proper capability check against s.capabilities.tasks
  • Correct request unmarshaling and header propagation
  • Hook lifecycle (beforeGetTask → handleGetTask → afterGetTask/onError)
  • Consistent error handling

358-383: LGTM!

The tasks/list handler is correctly implemented following the same pattern as the get handler, with proper type usage for ListTasksRequest and ListTasksResult.

384-409: LGTM!

The tasks/result handler maintains consistency with the other task handlers and correctly uses the TaskResultRequest and TaskResultResult types.

410-435: LGTM!

The tasks/cancel handler completes the task operation support with the same reliable pattern used throughout the request handler.

server/server.go (9)

31-40: LGTM!

The taskEntry struct is well-designed with:

  • Proper session isolation via sessionID
  • Completion signaling via done channel
  • completed flag to guard against double-close (addresses previous review feedback)
  • Clean separation of result and error

220-225: LGTM!

The task capability configuration follows the established pattern in this codebase:

  • Uses functional options pattern
  • Always creates non-nil capability object (consistent with other capability methods)
  • Clear boolean flags for list, cancel, and toolCallTasks features

Also applies to: 359-369

746-769: LGTM!

The task capability exposure in handleInitialize correctly:

  • Checks for nil before constructing capabilities
  • Conditionally sets List, Cancel, and Requests fields based on configured flags
  • Uses proper MCP capability structure with nested empty structs for presence indicators

1403-1526: LGTM!

The task handler functions are well-implemented:

  • handleGetTask and handleListTasks work with task copies (addresses previous data race concerns)
  • handleTaskResult properly waits on the done channel and re-fetches under lock to read resultErr safely
  • handleCancelTask correctly retrieves the updated task after cancellation
  • Error handling is consistent with MCP error codes

1532-1559: LGTM!

The createTask function correctly handles optional parameters (addresses previous review feedback):

  • Builds options slice conditionally to avoid nil dereference
  • Only applies WithTaskTTL and WithTaskPollInterval when pointers are non-nil
  • Properly schedules TTL-based cleanup when applicable

1561-1603: LGTM!

The task retrieval functions are correctly implemented:

  • getTask returns a copy of the task (line 1579), preventing external mutation and data races
  • Both functions enforce session isolation
  • Appropriate use of RLock for read operations
  • getTaskEntry clearly documented for internal use

1605-1621: LGTM!

The listTasks function correctly returns task copies (addresses previous review feedback about leaking internal pointers):

  • Returns []mcp.Task instead of []*taskEntry
  • Proper session-based filtering
  • Thread-safe with RLock

1623-1645: LGTM!

The completeTask function is well-protected against double completion:

  • Guards with completed flag check (addresses previous review feedback)
  • Properly sets task status (completed/failed) based on error presence
  • Closes done channel exactly once under mutex protection

1647-1692: LGTM!

The remaining task management functions are correctly implemented:

  • cancelTask properly guards against double completion and closes the channel exactly once (addresses previous review feedback)
  • scheduleTaskCleanup provides simple TTL-based cleanup
  • getSessionID helper cleanly extracts session information from context
server/task_test.go (2)

14-343: LGTM!

The test suite provides comprehensive coverage:

  • Capability negotiation with various flag combinations
  • Task lifecycle (create, retrieve, complete)
  • All handler functions via JSON-RPC messages
  • Error cases (not found, terminal cancellation, missing capabilities)
  • Follows Go testing best practices with table-driven tests

345-543: LGTM!

The additional test cases cover important functionality:

  • TTL-based cleanup with appropriate timing
  • Terminal status detection for all status values
  • Blocking behavior for incomplete tasks (with timeout protection)
  • Error propagation through task completion
  • Helper function correctness
  • JSON marshaling round-trip integrity

Excellent test coverage overall.

mcp/tasks.go (3)

7-30: LGTM!

The functional options pattern for Task creation is well-implemented:

  • Clear TaskOption type definition
  • Sensible defaults (Status: Working, CreatedAt: current time)
  • Options applied in order for flexibility

32-118: LGTM!

The option builders and helper constructors are clean and consistent:

  • All WithTask* functions follow the same pattern
  • Proper handling of pointer types for optional fields
  • Result constructors provide convenient APIs

139-151: LGTM!

The NewTasksCapabilityWithToolsOnly function now correctly implements "tools only" semantics (addresses previous review feedback):

  • Only sets Requests.Tools.Call
  • Does NOT set List or Cancel fields
  • Implementation matches the comment
mcp/types.go (5)

66-80: LGTM!

The MCP method constants for tasks are properly defined:

  • Consistent naming pattern (tasks/get, tasks/list, etc.)
  • Include spec URLs for reference
  • Notification method follows naming convention

Also applies to: 100-102

514-515: LGTM!

The Tasks capability fields are symmetrically added to both ClientCapabilities and ServerCapabilities:

  • Uses pointer type for optional presence (follows existing pattern)
  • Proper JSON omitempty tag
  • Good documentation

Also applies to: 550-551

1243-1271: LGTM!

The task capability structures follow MCP conventions:

  • Empty struct pointers for presence indicators
  • Proper nesting for request-specific capabilities
  • Clear documentation of each field's purpose

1273-1309: LGTM!

The core task types are well-defined:

  • TaskStatus enum with all required states
  • IsTerminal() method correctly identifies terminal states
  • Task struct with proper field types and JSON tags
  • Optional fields (TTL, PollInterval) correctly use pointers

1311-1397: LGTM!

The request and result types for task operations are complete and well-structured:

  • All request types properly embed base types and include Header fields
  • Parameter structs contain required fields
  • Result types follow MCP conventions
  • TaskResultResult appropriately notes that actual result varies by request type
  • TaskStatusNotification properly structured for status change notifications

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

coderabbitai[bot] avatar Nov 18 '25 15:11 coderabbitai[bot]

@JAORMX can you fix the linting errors?

ezynda3 avatar Dec 10 '25 11:12 ezynda3

@ezynda3 done

JAORMX avatar Dec 11 '25 17:12 JAORMX