diff --git a/README.md b/README.md index 729e9d2bd..310ca48cd 100644 --- a/README.md +++ b/README.md @@ -419,6 +419,26 @@ s := server.NewMCPServer( When the limit is reached, new task creation requests will fail with an error. Completed, failed, or cancelled tasks don't count toward the limit - only tasks in "working" status. If `WithMaxConcurrentTasks` is not specified or set to 0, there is no limit on concurrent tasks. +##### Transport Considerations for Task Handlers + +Task handlers run as background goroutines. Their context is not cancelled when `HandleMessage` returns. It stays alive until `tasks/cancel` cancels it or the task's TTL expires. Server shutdown does not cancel running tasks. + +Blocking on `ctx.Done()` is therefore safe and is the recommended way to honour cancellation: + +```go +s.AddTaskTool(myTool, func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CreateTaskResult, error) { + for _, item := range items { + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-time.After(perItemDelay * time.Millisecond): + } + processItem(item) + } + return result, nil +}) +``` + For traditional synchronous tools that execute and return results immediately: Simple calculation example: @@ -538,6 +558,26 @@ Each tool should: + +##### Task Authorization + +SEP-2663 requires servers to authenticate and authorize every task-related request (`tasks/get`, `tasks/cancel`, `tasks/result`, `tasks/list`, `tasks/update`). The SDK enforces the following default policy: + +| Task session | Request session | Result | +|---|---|---| +| set | set, **same** | Allowed | +| set | set, **different** | Denied (`task not found`, to avoid enumeration) | +| **empty** | any | Allowed — bearer-token model | +| any | **empty** | Allowed — bearer-token model | + +**Bearer-token default.** When either side has no session (e.g. stdio, in-process, or a task created outside a session), possessing the task ID is the only credential. This is the bearer-token model explicitly permitted by SEP-2663: sufficient entropy that a third party cannot enumerate or guess them. This is the intended default for session-less transports. + +**Enforcing application-level authorization.** Applications that need finer-grained control (per-user or per-tenant access) can use two existing extension points: + +1. **HTTP middleware** on the transport layer. An outer `http.Handler` can validate credentials, attach identity to the request context, and reject unauthorized requests before `HandleMessage` is called. This validates the credential on each request. It does not bind the credential to a session, and a session ID does not identify the caller. To limit a session to one caller, store the authenticated identity for the session and reject requests whose identity differs. + +2. **`OnBefore*` hooks** (`AddBeforeGetTask`, `AddBeforeCancelTask`, `AddBeforeUpdateTask`, etc.) are for **observation only** — logging, auditing, or recording the authenticated identity (which transport middleware can place in `ctx`) alongside the task operation. They are fire-and-forget: they cannot return an error, and the handler runs regardless of what the hook does. **Do not use them to enforce authorization.** Enforce task authorization in HTTP middleware (option 1), which can reject the request before `HandleMessage` is called. (See `TestTaskAccessControl_OnBeforeHookExtensionPoint` in `server/task_auth_contract_test.go`, which demonstrates reading the identity — not enforcing it.) + ### Prompts
diff --git a/client/client.go b/client/client.go index 9f1a4a692..7ff5a4d4d 100644 --- a/client/client.go +++ b/client/client.go @@ -1071,6 +1071,19 @@ func (c *Client) CancelTask( return mcp.ParseCancelTaskResult(response) } +// UpdateTask delivers inputResponses to a task that is waiting for client input +func (c *Client) UpdateTask( + ctx context.Context, + request mcp.UpdateTaskRequest, +) (*mcp.UpdateTaskResult, error) { + response, err := c.sendRequest(ctx, string(mcp.MethodTasksUpdate), request.Params, outboundHeader(request.Header, request.Method)) + if err != nil { + return nil, err + } + + return mcp.ParseUpdateTaskResult(response) +} + // ErrRepeatedTaskListCursor is returned when tasks/list repeats a cursor that // was already requested. Following it would request the same page again. var ErrRepeatedTaskListCursor = errors.New("tasks/list cursor did not advance") diff --git a/examples/task_tool/main.go b/examples/task_tool/main.go index e789f0f41..af1b6ade4 100644 --- a/examples/task_tool/main.go +++ b/examples/task_tool/main.go @@ -1,9 +1,3 @@ -//go:build ignore -// +build ignore - -// NOTE: This example will not compile until task tool implementation is complete. -// Remove the build constraint above once TAS-4 through TAS-10 are implemented. - package main import ( @@ -18,274 +12,171 @@ import ( ) // This example demonstrates task-augmented tools in MCP. -// Task-augmented tools execute asynchronously and return results via polling. -// -// NOTE: This example requires the task tool implementation to be complete. -// Specifically, it requires: -// - TaskToolHandlerFunc type (TAS-4) -// - ServerTaskTool type (TAS-5) -// - AddTaskTool method (TAS-6) -// - handleTaskAugmentedToolCall (TAS-9) -// - executeTaskTool (TAS-10) -// -// The example includes three types of tools: -// 1. process_batch - A TaskSupportRequired tool that processes items in batch -// 2. analyze_data - A TaskSupportOptional tool that can run sync or async -// 3. quick_check - A regular synchronous tool for comparison +// Task-augmented tools execute asynchronously: the server returns a task handle +// immediately, and the client polls tasks/get until the task reaches a terminal +// state (completed, failed, or cancelled). +// The context passed to a task handler is rooted at the server's lifetime, not +// the per-request context. It is only cancelled when tasks/cancel is called or +// the server shuts down. Blocking on ctx.Done() is safe and is the recommended +// way to honour cancellation. // // Usage: -// go run main.go # Start with stdio transport -// go run main.go -t http # Start with HTTP transport on :8080 +// go run main.go # stdio transport +// go run main.go -t http # Streamable HTTP on :8080 func main() { - var transport string - flag.StringVar(&transport, "t", "stdio", "Transport type (stdio or http)") - flag.StringVar(&transport, "transport", "stdio", "Transport type (stdio or http)") + transport := flag.String("t", "stdio", "Transport: stdio or http") flag.Parse() - mcpServer := NewTaskToolServer() + srv := newServer() - if transport == "http" { - httpServer := server.NewStreamableHTTPServer(mcpServer) - log.Printf("HTTP server listening on :8080/mcp") - if err := httpServer.Start(":8080"); err != nil { - log.Fatalf("Server error: %v", err) + switch *transport { + case "http": + log.Println("Streamable HTTP listening on :8080") + if err := server.NewStreamableHTTPServer(srv).Start(":8080"); err != nil { + log.Fatal(err) } - } else { - if err := server.ServeStdio(mcpServer); err != nil { - log.Fatalf("Server error: %v", err) + default: + if err := server.ServeStdio(srv); err != nil { + log.Fatal(err) } } } -// NewTaskToolServer creates a new MCP server with task-augmented tool examples. -func NewTaskToolServer() *server.MCPServer { - // Set up task observability hooks - taskHooks := &server.TaskHooks{} - taskHooks.AddOnTaskCreated(func(ctx context.Context, metrics server.TaskMetrics) { - log.Printf("[METRICS] Task created: %s (tool: %s)", metrics.TaskID, metrics.ToolName) +func newServer() *server.MCPServer { + hooks := &server.TaskHooks{} + hooks.AddOnTaskCreated(func(_ context.Context, m server.TaskMetrics) { + log.Printf("[task] created %s tool=%s", m.TaskID, m.ToolName) }) - taskHooks.AddOnTaskCompleted(func(ctx context.Context, metrics server.TaskMetrics) { - log.Printf("[METRICS] Task completed: %s (tool: %s, duration: %v)", - metrics.TaskID, metrics.ToolName, metrics.Duration) + hooks.AddOnTaskCompleted(func(_ context.Context, m server.TaskMetrics) { + log.Printf("[task] done %s tool=%s duration=%v", m.TaskID, m.ToolName, m.Duration) }) - taskHooks.AddOnTaskFailed(func(ctx context.Context, metrics server.TaskMetrics) { - log.Printf("[METRICS] Task failed: %s (tool: %s, duration: %v, error: %v)", - metrics.TaskID, metrics.ToolName, metrics.Duration, metrics.Error) + hooks.AddOnTaskFailed(func(_ context.Context, m server.TaskMetrics) { + log.Printf("[task] failed %s tool=%s err=%v", m.TaskID, m.ToolName, m.Error) }) - taskHooks.AddOnTaskCancelled(func(ctx context.Context, metrics server.TaskMetrics) { - log.Printf("[METRICS] Task cancelled: %s (tool: %s, duration: %v)", - metrics.TaskID, metrics.ToolName, metrics.Duration) + hooks.AddOnTaskCancelled(func(_ context.Context, m server.TaskMetrics) { + log.Printf("[task] cancelled %s tool=%s", m.TaskID, m.ToolName) }) - // Create server with task capabilities enabled - // listTasks: allows clients to list all tasks - // cancel: allows clients to cancel running tasks - // toolCallTasks: enables task augmentation for tools - mcpServer := server.NewMCPServer( - "example-servers/task-tool", + srv := server.NewMCPServer( + "example-task-tool", "1.0.0", server.WithTaskCapabilities(true, true, true), + // advertise io.modelcontextprotocol/tasks for SEP-2663 clients + server.WithTasksExtension(), server.WithToolCapabilities(true), - server.WithTaskHooks(taskHooks), - server.WithMaxConcurrentTasks(10), // Limit to 10 concurrent tasks - server.WithLogging(), + server.WithTaskHooks(hooks), + server.WithMaxConcurrentTasks(10), ) - // Example 1: Task-Required Tool - // This tool MUST be called with task augmentation - processBatchTool := mcp.NewTool("process_batch", - mcp.WithDescription("Process a batch of items asynchronously. This is a long-running operation that must be executed as a task."), + // Tool 1: must always run as a task. + processBatch := mcp.NewTool("process_batch", + mcp.WithDescription("Process items asynchronously (task required)."), mcp.WithTaskSupport(mcp.TaskSupportRequired), mcp.WithArray("items", - mcp.Description("Array of items to process"), + mcp.Description("Items to process"), mcp.WithStringItems(), mcp.Required(), ), - mcp.WithNumber("delay", - mcp.Description("Delay in seconds per item (simulates processing time)"), - mcp.DefaultNumber(1.0), + mcp.WithNumber("delay_ms", + mcp.Description("Simulated delay per item in ms"), + mcp.DefaultNumber(200), ), ) - mcpServer.AddTaskTool(processBatchTool, handleProcessBatch) + srv.AddTaskTool(processBatch, handleProcessBatch) - // Example 2: Task-Optional Tool - // This tool can be called with or without task augmentation - analyzeDataTool := mcp.NewTool("analyze_data", - mcp.WithDescription("Analyze data. Can run synchronously (fast) or asynchronously (thorough) depending on how it's called."), + // Tool 2: can run synchronously or as a task. + analyzeData := mcp.NewTool("analyze_data", + mcp.WithDescription("Analyse text — quick sync or thorough async."), mcp.WithTaskSupport(mcp.TaskSupportOptional), - mcp.WithString("data", - mcp.Description("Data to analyze"), - mcp.Required(), - ), - mcp.WithBoolean("thorough", - mcp.Description("Whether to perform thorough analysis (slower)"), - mcp.DefaultBool(false), - ), + mcp.WithString("text", mcp.Required()), + mcp.WithBoolean("thorough", mcp.DefaultBool(false)), ) - mcpServer.AddTaskTool(analyzeDataTool, handleAnalyzeData) + // Register as a regular tool with optional task support. + srv.AddTool(analyzeData, handleAnalyzeData) - // Example 3: Regular Synchronous Tool (for comparison) - quickCheckTool := mcp.NewTool("quick_check", - mcp.WithDescription("Perform a quick synchronous check. This tool does not support task augmentation."), - mcp.WithString("input", - mcp.Description("Input to check"), - mcp.Required(), - ), + // Tool 3: plain synchronous tool, no task support. + quickCheck := mcp.NewTool("quick_check", + mcp.WithDescription("Synchronous sanity check."), + mcp.WithString("input", mcp.Required()), ) - mcpServer.AddTool(quickCheckTool, handleQuickCheck) + srv.AddTool(quickCheck, handleQuickCheck) - return mcpServer + return srv } -// handleProcessBatch processes items in a batch asynchronously. -// This demonstrates a task-required tool that performs long-running work. -func handleProcessBatch( - ctx context.Context, - request mcp.CallToolRequest, -) (*mcp.CreateTaskResult, error) { - items := request.GetStringSlice("items", []string{}) - delay := request.GetFloat("delay", 1.0) - +// handleProcessBatch is a TaskToolHandlerFunc, it runs in a background +// goroutine and returns a *CreateTaskResult when the work is done. +func handleProcessBatch(ctx context.Context, req mcp.CallToolRequest) (*mcp.CreateTaskResult, error) { + items := req.GetStringSlice("items", nil) + delayMs := req.GetFloat("delay_ms", 200) if len(items) == 0 { return nil, fmt.Errorf("items array cannot be empty") } + delay := time.Duration(delayMs) * time.Millisecond - log.Printf("Starting batch processing of %d items with %v second delay per item", len(items), delay) + log.Printf("process_batch: starting %d items, delay=%v", len(items), delay) - // Process each item with delay - results := make([]string, 0, len(items)) for i, item := range items { select { case <-ctx.Done(): - // Task was cancelled - log.Printf("Batch processing cancelled after %d/%d items", i, len(items)) + log.Printf("process_batch: cancelled after %d/%d items", i, len(items)) return nil, ctx.Err() - default: - // Simulate processing time - time.Sleep(time.Duration(delay * float64(time.Second))) - result := fmt.Sprintf("Processed '%s' [%d/%d]", item, i+1, len(items)) - results = append(results, result) - log.Printf("Item %d/%d: %s", i+1, len(items), item) + case <-time.After(delay): } - } - log.Printf("Batch processing completed: %d items", len(items)) + log.Printf("process_batch: [%d/%d] %s", i+1, len(items), item) + } - // Return the result - return &mcp.CreateTaskResult{ - Task: mcp.Task{ - // Task fields are managed by the server - }, - }, nil + log.Printf("process_batch: finished %d items", len(items)) + r := mcp.NewCreateTaskResult(mcp.Task{}) + return &r, nil } -// handleAnalyzeData analyzes data either quickly or thoroughly. -// This demonstrates a task-optional tool that adapts based on how it's called. -func handleAnalyzeData( - ctx context.Context, - request mcp.CallToolRequest, -) (*mcp.CreateTaskResult, error) { - data := request.GetString("data", "") - thorough := request.GetBool("thorough", false) - - if data == "" { - return nil, fmt.Errorf("data cannot be empty") - } +// handleAnalyzeData is a ToolHandlerFunc for the optional-task tool. The same +// handler is used whether the tool is called synchronously or as a task. +func handleAnalyzeData(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + text := req.GetString("text", "") + thorough := req.GetBool("thorough", false) - analysisType := "quick" - duration := 500 * time.Millisecond + duration := 100 * time.Millisecond if thorough { - analysisType = "thorough" - duration = 3 * time.Second + duration = 2 * time.Second } - log.Printf("Starting %s analysis of data: %s", analysisType, data) - - // Simulate analysis time + // The same handler runs on both the synchronous and the task path. On the + // task path ctx is cancelled by tasks/cancel or TTL expiry. select { case <-ctx.Done(): - log.Printf("Analysis cancelled") return nil, ctx.Err() case <-time.After(duration): - // Analysis complete } - // Generate analysis results - wordCount := len(data) - charCount := len([]rune(data)) - - result := fmt.Sprintf("Analysis complete (%s mode):\n", analysisType) - result += fmt.Sprintf("- Characters: %d\n", charCount) - result += fmt.Sprintf("- Words: %d\n", wordCount) - + summary := fmt.Sprintf("analysed %d chars", len([]rune(text))) if thorough { - result += fmt.Sprintf("- Uppercase: %d\n", countUppercase(data)) - result += fmt.Sprintf("- Lowercase: %d\n", countLowercase(data)) - result += fmt.Sprintf("- Digits: %d\n", countDigits(data)) + summary += fmt.Sprintf(", upper=%d lower=%d digits=%d", + countRunes(text, 'A', 'Z'), + countRunes(text, 'a', 'z'), + countRunes(text, '0', '9'), + ) } + log.Printf("analyze_data: %s", summary) - log.Printf("Analysis completed: %s mode, %d characters", analysisType, charCount) - - return &mcp.CreateTaskResult{ - Task: mcp.Task{ - // Task fields are managed by the server - }, - }, nil -} - -// handleQuickCheck performs a quick synchronous check. -// This demonstrates a regular tool without task support for comparison. -func handleQuickCheck( - ctx context.Context, - request mcp.CallToolRequest, -) (*mcp.CallToolResult, error) { - input := request.GetString("input", "") - - if input == "" { - return nil, fmt.Errorf("input cannot be empty") - } - - log.Printf("Quick check: %s", input) - - return &mcp.CallToolResult{ - Content: []mcp.Content{ - mcp.TextContent{ - Type: "text", - Text: fmt.Sprintf("Quick check passed for: %s (length: %d)", input, len(input)), - }, - }, - }, nil -} - -// Helper functions for thorough analysis - -func countUppercase(s string) int { - count := 0 - for _, r := range s { - if r >= 'A' && r <= 'Z' { - count++ - } - } - return count + return mcp.NewToolResultText(summary), nil } -func countLowercase(s string) int { - count := 0 - for _, r := range s { - if r >= 'a' && r <= 'z' { - count++ - } - } - return count +// handleQuickCheck is a plain ToolHandlerFunc, synchronous, no task involved. +func handleQuickCheck(_ context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + input := req.GetString("input", "") + return mcp.NewToolResultText(fmt.Sprintf("ok: %q (%d chars)", input, len(input))), nil } -func countDigits(s string) int { - count := 0 +func countRunes(s string, lo, hi rune) int { + n := 0 for _, r := range s { - if r >= '0' && r <= '9' { - count++ + if r >= lo && r <= hi { + n++ } } - return count + return n } diff --git a/mcp/headers.go b/mcp/headers.go index 6bcbe979a..8736dcac7 100644 --- a/mcp/headers.go +++ b/mcp/headers.go @@ -67,7 +67,8 @@ func RequiresStandardHeaders(protocolVersion string) bool { // the given method. func MethodRequiresNameHeader(method MCPMethod) bool { switch method { - case MethodToolsCall, MethodResourcesRead, MethodPromptsGet: + case MethodToolsCall, MethodResourcesRead, MethodPromptsGet, + MethodTasksGet, MethodTasksUpdate, MethodTasksCancel: return true default: return false @@ -76,8 +77,8 @@ func MethodRequiresNameHeader(method MCPMethod) bool { // ExtractHeaderName returns the value that belongs in the Mcp-Name header for // the given method and raw params: params.name for tools/call and prompts/get, -// params.uri for resources/read. It reports false when the method does not -// carry a name. +// params.uri for resources/read, params.taskId for tasks/get, tasks/update, and +// tasks/cancel. It reports false when the method does not carry a name. func ExtractHeaderName(method MCPMethod, params json.RawMessage) (string, bool) { switch method { case MethodToolsCall: @@ -101,6 +102,13 @@ func ExtractHeaderName(method MCPMethod, params json.RawMessage) (string, bool) if json.Unmarshal(params, &p) == nil { return p.URI, true } + case MethodTasksGet, MethodTasksUpdate, MethodTasksCancel: + var p struct { + TaskID string `json:"taskId"` + } + if json.Unmarshal(params, &p) == nil { + return p.TaskID, true + } } return "", false } diff --git a/mcp/headers_test.go b/mcp/headers_test.go index 4b83e6c66..3726db8dc 100644 --- a/mcp/headers_test.go +++ b/mcp/headers_test.go @@ -66,6 +66,9 @@ func TestExtractHeaderName(t *testing.T) { {"tools/call uses the tool name", MethodToolsCall, `{"name":"search"}`, "search", true}, {"prompts/get uses the prompt name", MethodPromptsGet, `{"name":"summarize"}`, "summarize", true}, {"resources/read uses the uri", MethodResourcesRead, `{"uri":"file:///a"}`, "file:///a", true}, + {"tasks/get uses the taskId", MethodTasksGet, `{"taskId":"task-1"}`, "task-1", true}, + {"tasks/update uses the taskId", MethodTasksUpdate, `{"taskId":"task-2"}`, "task-2", true}, + {"tasks/cancel uses the taskId", MethodTasksCancel, `{"taskId":"task-3"}`, "task-3", true}, {"other methods carry no name", MethodToolsList, `{}`, "", false}, } @@ -74,6 +77,9 @@ func TestExtractHeaderName(t *testing.T) { got, ok := ExtractHeaderName(tt.method, json.RawMessage(tt.params)) assert.Equal(t, tt.ok, ok) assert.Equal(t, tt.want, got) + // Methods that carry a name must also require the Mcp-Name header. + assert.Equal(t, tt.ok, MethodRequiresNameHeader(tt.method), + "MethodRequiresNameHeader must agree with ExtractHeaderName for %s", tt.method) }) } } diff --git a/mcp/result.go b/mcp/result.go index 355729643..2e12720e9 100644 --- a/mcp/result.go +++ b/mcp/result.go @@ -16,6 +16,11 @@ const ( // [InputRequiredResult.InputRequests] and retries the original request with // the answers attached. See the multi round-trip request pattern. ResultTypeInputRequired ResultType = "input_required" + + // ResultTypeTask indicates the server has created an asynchronous task to + // handle the request. The client polls tasks/get until the task reaches a + // terminal state. + ResultTypeTask ResultType = "task" ) // IsComplete reports whether the result type denotes a final result. An empty diff --git a/mcp/task_status_parsing_test.go b/mcp/task_status_parsing_test.go index c242c6e2d..52e283761 100644 --- a/mcp/task_status_parsing_test.go +++ b/mcp/task_status_parsing_test.go @@ -8,31 +8,60 @@ import ( ) func TestTaskStatusParsersPreserveEnvelopeAndIntegerFields(t *testing.T) { - for _, kind := range []string{"get", "cancel", "list"} { - t.Run(kind, func(t *testing.T) { - task := `{"taskId":"task","status":"completed","createdAt":"2026-10-09T00:00:00Z","lastUpdatedAt":"2026-10-09T00:00:00Z","ttl":9007199254740993,"pollInterval":9007199254740993}` - raw := json.RawMessage(`{"resultType":"data","_meta":{"key":"value"},` + task[1:]) - if kind == "list" { - raw = json.RawMessage(`{"resultType":"data","_meta":{"key":"value"},"tasks":[` + task + `]}`) - } - var result any - var err error - switch kind { - case "get": - result, err = ParseGetTaskResult(&raw) - case "cancel": - result, err = ParseCancelTaskResult(&raw) - default: - result, err = ParseListTasksResult(&raw) - } - require.NoError(t, err) - encoded, err := json.Marshal(result) - require.NoError(t, err) - var fields map[string]json.RawMessage - require.NoError(t, json.Unmarshal(encoded, &fields)) - require.JSONEq(t, `"data"`, string(fields["resultType"])) - require.Contains(t, string(encoded), `"ttl":9007199254740993`) - require.Contains(t, string(encoded), `"pollInterval":9007199254740993`) - }) + // bigInt is 2^53+1, which is not exactly representable as a float64. Parsing + // must preserve it as an int64, so decoding must not route through float64. + const bigInt = `9007199254740993` + + for _, wire := range []struct { + name string + ttlField string + pollField string + }{ + // Legacy (ttl/pollInterval): decoded by the struct's json tags. + // Modern (ttlMs/pollIntervalMs): decoded by jsonToTask. + {"legacy-input", "ttl", "pollInterval"}, + {"modern-input", "ttlMs", "pollIntervalMs"}, + } { + for _, kind := range []string{"get", "cancel", "list"} { + t.Run(wire.name+"/"+kind, func(t *testing.T) { + task := `{"taskId":"task","status":"completed",` + + `"createdAt":"2026-10-09T00:00:00Z","lastUpdatedAt":"2026-10-09T00:00:00Z",` + + `"` + wire.ttlField + `":` + bigInt + `,"` + wire.pollField + `":` + bigInt + `}` + raw := json.RawMessage(`{"resultType":"data","_meta":{"key":"value"},` + task[1:]) + if kind == "list" { + raw = json.RawMessage(`{"resultType":"data","_meta":{"key":"value"},"tasks":[` + task + `]}`) + } + + var result any + var err error + switch kind { + case "get": + result, err = ParseGetTaskResult(&raw) + case "cancel": + result, err = ParseCancelTaskResult(&raw) + default: + result, err = ParseListTasksResult(&raw) + } + require.NoError(t, err) + + encoded, err := json.Marshal(result) + require.NoError(t, err) + var fields map[string]json.RawMessage + require.NoError(t, json.Unmarshal(encoded, &fields)) + require.JSONEq(t, `"data"`, string(fields["resultType"])) + + // The output wire names depend on the response type, not the + // input. The int64 value must survive the round trip exactly. + if kind == "list" { + // tasks/list serves legacy 2025-11-25 clients (ttl/pollInterval) + require.Contains(t, string(encoded), `"ttl":`+bigInt) + require.Contains(t, string(encoded), `"pollInterval":`+bigInt) + } else { + // tasks/get and tasks/cancel use SEP-2663 names (ttlMs/pollIntervalMs). + require.Contains(t, string(encoded), `"ttlMs":`+bigInt) + require.Contains(t, string(encoded), `"pollIntervalMs":`+bigInt) + } + }) + } } } diff --git a/mcp/tasks.go b/mcp/tasks.go index ca6f1e4d6..fecca5499 100644 --- a/mcp/tasks.go +++ b/mcp/tasks.go @@ -1,9 +1,24 @@ package mcp import ( + "fmt" "time" ) +// ExtensionTasks is the identifier for the MCP Tasks extension. +const ExtensionTasks = "io.modelcontextprotocol/tasks" + +// UnsupportedTaskResultError is returned when a peer responded to a request +// with a task handle (resultType:"task") from the ExtensionTasks extension, +// but the local SDK cannot automatically resolve it. +type UnsupportedTaskResultError struct { + TaskID string +} + +func (e *UnsupportedTaskResultError) Error() string { + return fmt.Sprintf("peer created task %q: the %s extension is not supported by this client", e.TaskID, ExtensionTasks) +} + // TaskOption is a function that configures a Task. // It provides a flexible way to set various properties of a Task using the functional options pattern. type TaskOption func(*Task) @@ -80,9 +95,11 @@ func NewTaskParams(ttlMs *int64) TaskParams { } // NewCreateTaskResult creates a CreateTaskResult with the given task. +// Tasks extension wire format, set Legacy=false. func NewCreateTaskResult(task Task) CreateTaskResult { return CreateTaskResult{ - Task: task, + Task: task, + Legacy: true, } } diff --git a/mcp/types.go b/mcp/types.go index 8fc70a83a..b70cd6bbb 100644 --- a/mcp/types.go +++ b/mcp/types.go @@ -97,8 +97,11 @@ const ( // https://modelcontextprotocol.io/specification/2025-11-25/basic/utilities/tasks MethodTasksGet MCPMethod = "tasks/get" + // MethodTasksUpdate submits responses to outstanding inputRequests for a + // task in the input_required state. + MethodTasksUpdate MCPMethod = "tasks/update" + // MethodTasksList lists all tasks for the current session. - // https://modelcontextprotocol.io/specification/2025-11-25/basic/utilities/tasks MethodTasksList MCPMethod = "tasks/list" // MethodTasksResult retrieves the result of a completed task. @@ -147,6 +150,10 @@ const ( // https://modelcontextprotocol.io/specification/2025-11-25/basic/utilities/tasks MethodNotificationTasksStatus = "notifications/tasks/status" + // MethodNotificationTasks is the SEP-2663 tasks extension notification + // method for task status changes. + MethodNotificationTasks = "notifications/tasks" + // MethodCompletionComplete returns completion suggestions for a given argument // https://modelcontextprotocol.io/specification/2025-11-25/server/utilities/completion MethodCompletionComplete MCPMethod = "completion/complete" @@ -660,6 +667,17 @@ type ClientCapabilities struct { Tasks *TasksCapability `json:"tasks,omitempty"` } +// HasExtension reports whether the client declared support for the named +// extension. The name should be a fully-qualified extension identifier such as +// "io.modelcontextprotocol/tasks". +func (c *ClientCapabilities) HasExtension(name string) bool { + if c == nil { + return false + } + _, ok := c.Extensions[name] + return ok +} + // ServerCapabilities represents capabilities that a server may support. Known // capabilities are defined here, in this schema, but this is not a closed set: any // server can define its own, additional capabilities. @@ -1790,6 +1808,36 @@ func (t Task) GetName() string { return t.TaskId } +// MarshalJSONLegacyTask serializes the Task of legacy 2025-11-25 wire format. +func (r CreateTaskResult) MarshalJSONLegacyTask() ([]byte, error) { + taskFields := map[string]any{ + "taskId": r.Task.TaskId, + "status": r.Task.Status, + "createdAt": r.Task.CreatedAt, + "lastUpdatedAt": r.Task.LastUpdatedAt, + "ttl": r.Task.TTL, + } + if r.Task.StatusMessage != "" { + taskFields["statusMessage"] = r.Task.StatusMessage + } + if r.Task.PollInterval != nil { + taskFields["pollInterval"] = *r.Task.PollInterval + } + raw, err := json.Marshal(taskFields) + if err != nil { + return nil, err + } + m := map[string]json.RawMessage{"task": raw} + if r.Meta != nil { + metaBytes, err := json.Marshal(r.Meta) + if err != nil { + return nil, err + } + m["_meta"] = metaBytes + } + return json.Marshal(m) +} + // TaskParams represents the task metadata included when augmenting a request. type TaskParams struct { // Requested duration in milliseconds to retain task from creation. @@ -1800,10 +1848,84 @@ type TaskParams struct { // It contains task metadata rather than the actual operation result. type CreateTaskResult struct { Result - Task Task `json:"task"` + Task Task `json:"-"` Content []Content `json:"-"` StructuredContent any `json:"-"` IsError bool `json:"-"` + Legacy bool `json:"-"` +} + +// marshalTaskToMap serialises a Task to a map using SEP-2663 wire field names +// (ttlMs, pollIntervalMs). +func marshalTaskToMap(t Task) map[string]any { + m := map[string]any{ + "taskId": t.TaskId, + "status": t.Status, + "createdAt": t.CreatedAt, + "lastUpdatedAt": t.LastUpdatedAt, + } + if t.StatusMessage != "" { + m["statusMessage"] = t.StatusMessage + } + // SEP-2663 requires ttlMs in task responses (number or null), emit null for + // tasks with no expiration rather than omitting the key. + m["ttlMs"] = t.TTL + if t.PollInterval != nil { + m["pollIntervalMs"] = *t.PollInterval + } + return m +} + +// MarshalJSON implements custom JSON marshaling for CreateTaskResult. +func (r CreateTaskResult) MarshalJSON() ([]byte, error) { + if r.Legacy { + return r.MarshalJSONLegacyTask() + } + m := marshalTaskToMap(r.Task) + m["resultType"] = ResultTypeTask + if r.Meta != nil { + m["_meta"] = r.Meta + } + return json.Marshal(m) +} + +// UnmarshalJSON handles both wire formats SEP-2663 and Legacy(2025-11-25). +func (r *CreateTaskResult) UnmarshalJSON(data []byte) error { + // Always decode the common Result fields (_meta, resultType) first. + var base struct { + Meta *Meta `json:"_meta"` + ResultType ResultType `json:"resultType"` + Task *Task `json:"task"` + TaskId string `json:"taskId"` + } + if err := json.Unmarshal(data, &base); err != nil { + return err + } + r.Meta = base.Meta + r.ResultType = base.ResultType + + if base.Task != nil { + // Legacy format: task nested under "task" key. + r.Task = *base.Task + r.Legacy = true + return nil + } + r.Legacy = false + type inlineTask struct { + TaskId string `json:"taskId"` + Status TaskStatus `json:"status"` + StatusMessage string `json:"statusMessage,omitempty"` + CreatedAt string `json:"createdAt"` + LastUpdatedAt string `json:"lastUpdatedAt"` + TTL *int64 `json:"ttlMs"` + PollInterval *int64 `json:"pollIntervalMs,omitempty"` + } + var inline inlineTask + if err := json.Unmarshal(data, &inline); err != nil { + return err + } + r.Task = Task(inline) + return nil } // GetTaskRequest retrieves the current status of a task. @@ -1821,6 +1943,53 @@ type GetTaskParams struct { type GetTaskResult struct { Result Task + TaskResult json.RawMessage `json:"result,omitempty"` + TaskError *JSONRPCErrorDetails `json:"error,omitempty"` + // InputRequests holds outstanding server-to-client requests when Status + // is input_required. + InputRequests InputRequests `json:"inputRequests,omitempty"` + // Legacy instructs MarshalJSON to emit legacy 2025-11-25 wire field names + // (ttl, pollInterval) instead of SEP-2663 names (ttlMs, pollIntervalMs). + Legacy bool `json:"-"` +} + +// MarshalJSON emits SEP-2663 field names (ttlMs, pollIntervalMs) by default, +// or legacy 2025-11-25 field names (ttl, pollInterval) when Legacy is true. +func (r GetTaskResult) MarshalJSON() ([]byte, error) { + var m map[string]any + if r.Legacy { + m = map[string]any{ + "taskId": r.TaskId, + "status": r.Status, + "createdAt": r.CreatedAt, + "lastUpdatedAt": r.LastUpdatedAt, + } + if r.StatusMessage != "" { + m["statusMessage"] = r.StatusMessage + } + m["ttl"] = r.TTL + if r.PollInterval != nil { + m["pollInterval"] = *r.PollInterval + } + } else { + m = marshalTaskToMap(r.Task) + } + if r.Meta != nil { + m["_meta"] = r.Meta + } + if r.ResultType != "" { + m["resultType"] = r.ResultType + } + if r.TaskResult != nil { + m["result"] = r.TaskResult + } + if r.TaskError != nil { + m["error"] = r.TaskError + } + if len(r.InputRequests) > 0 { + m["inputRequests"] = r.InputRequests + } + return json.Marshal(m) } // ListTasksRequest retrieves a paginated list of tasks. @@ -1835,6 +2004,31 @@ type ListTasksResult struct { Tasks []Task `json:"tasks"` } +// MarshalJSON emits legacy 2025-11-25 field names (ttl, pollInterval) for each +// Task in the list. tasks/list is served only to non-modern clients — the +// modern protocol rejects it with METHOD_NOT_FOUND (SEP-2663 removed tasks/list) +// — so those clients read the legacy field names. ttl is emitted unconditionally +// (null for unlimited) per the 2025-11-25 TTL requirement; pollInterval is +// omitted when unset (the legacy schema marks it optional). +func (r ListTasksResult) MarshalJSON() ([]byte, error) { + tasks := make([]map[string]any, len(r.Tasks)) + for i, t := range r.Tasks { + m := marshalTaskToMap(t) + delete(m, "ttlMs") + delete(m, "pollIntervalMs") + m["ttl"] = t.TTL + if t.PollInterval != nil { + m["pollInterval"] = *t.PollInterval + } + tasks[i] = m + } + type alias struct { + PaginatedResult + Tasks []map[string]any `json:"tasks"` + } + return json.Marshal(alias{PaginatedResult: r.PaginatedResult, Tasks: tasks}) +} + // TaskResultRequest retrieves the result of a completed task. type TaskResultRequest struct { Request @@ -1889,6 +2083,45 @@ type CancelTaskParams struct { type CancelTaskResult struct { Result Task + // Legacy instructs MarshalJSON to emit legacy 2025-11-25 field names + // (ttl, pollInterval) for the non-ack cancel response shape. + Legacy bool `json:"-"` +} + +// MarshalJSON serialises CancelTaskResult. +func (r CancelTaskResult) MarshalJSON() ([]byte, error) { + if r.TaskId == "" { + // Ack-only (modern SEP-2663 path). + type resultOnly struct { + Result + } + return json.Marshal(resultOnly{r.Result}) + } + var m map[string]any + if r.Legacy { + m = map[string]any{ + "taskId": r.TaskId, + "status": r.Status, + "createdAt": r.CreatedAt, + "lastUpdatedAt": r.LastUpdatedAt, + } + if r.StatusMessage != "" { + m["statusMessage"] = r.StatusMessage + } + m["ttl"] = r.TTL + if r.PollInterval != nil { + m["pollInterval"] = *r.PollInterval + } + } else { + m = marshalTaskToMap(r.Task) + } + if r.Meta != nil { + m["_meta"] = r.Meta + } + if r.ResultType != "" { + m["resultType"] = r.ResultType + } + return json.Marshal(m) } // TaskStatusNotification is sent when a task's status changes. @@ -1901,6 +2134,25 @@ type TaskStatusNotificationParams struct { Task } +// UpdateTaskRequest submits responses to outstanding inputRequests for a task +// in the input_required state. +type UpdateTaskRequest struct { + Request + Header http.Header `json:"-"` + Params UpdateTaskParams `json:"params"` +} + +// UpdateTaskParams are the parameters for tasks/update. +type UpdateTaskParams struct { + TaskId string `json:"taskId"` + InputResponses InputResponses `json:"inputResponses"` +} + +// UpdateTaskResult is the ack-only result returned by tasks/update. +type UpdateTaskResult struct { + Result +} + // ClientRequest represents any request that can be sent from client to server. type ClientRequest any diff --git a/mcp/utils.go b/mcp/utils.go index d4a4066c3..949321f01 100644 --- a/mcp/utils.go +++ b/mcp/utils.go @@ -868,10 +868,18 @@ func ParseCallToolResult(rawMessage *json.RawMessage) (*CallToolResult, error) { Content json.RawMessage `json:"content"` StructuredContent json.RawMessage `json:"structuredContent"` ResultType ResultType `json:"resultType"` + TaskId string `json:"taskId"` } if err := json.Unmarshal(*rawMessage, &probe); err != nil { return nil, fmt.Errorf("failed to unmarshal response: %w", err) } + // A server-directed task handle (SEP-2663) carries no content; the peer + // created an asynchronous task instead of returning a result. Surface it as + // a typed error so the caller can recover the task ID and poll tasks/get, + // rather than failing with "content is missing". + if probe.ResultType == ResultTypeTask { + return nil, &UnsupportedTaskResultError{TaskID: probe.TaskId} + } // A structured-only result, and a result asking for more input (SEP-2322) // that carries only the input requests and the request state, legitimately // have no content. A null structuredContent counts as missing. @@ -1327,6 +1335,27 @@ func GetTextFromContent(content any) string { } } +// jsonToTask fills TTL and PollInterval from the SEP-2663 wire names (ttlMs, +// pollIntervalMs). The Task struct uses the legacy JSON tags (ttl, pollInterval) +// so plain json.Unmarshal decodes legacy names correctly but skips ttlMs/pollIntervalMs; +// this helper fills the gap. It only overrides fields that json.Unmarshal left nil, +// to avoid clobbering a precisely-decoded integer value with a lossy float64. +func jsonToTask(raw []byte, result *GetTaskResult) { + var modern struct { + TTLMs *int64 `json:"ttlMs"` + PollIntervalMs *int64 `json:"pollIntervalMs"` + } + if err := json.Unmarshal(raw, &modern); err != nil { + return + } + if result.TTL == nil && modern.TTLMs != nil { + result.TTL = modern.TTLMs + } + if result.PollInterval == nil && modern.PollIntervalMs != nil { + result.PollInterval = modern.PollIntervalMs + } +} + // ParseCancelTaskResult parses a JSON message and converts it to a CancelTaskResult. func ParseCancelTaskResult(rawMessage *json.RawMessage) (*CancelTaskResult, error) { if rawMessage == nil { @@ -1336,6 +1365,11 @@ func ParseCancelTaskResult(rawMessage *json.RawMessage) (*CancelTaskResult, erro if err := json.Unmarshal(*rawMessage, &result); err != nil { return nil, fmt.Errorf("failed to unmarshal response: %w", err) } + // Fix up ttlMs/pollIntervalMs: Task struct tags use legacy names so plain + // json.Unmarshal won't decode them from SEP-2663 wire format. + getResult := GetTaskResult{Task: result.Task} + jsonToTask(*rawMessage, &getResult) + result.Task = getResult.Task return &result, nil } @@ -1348,6 +1382,21 @@ func ParseListTasksResult(rawMessage *json.RawMessage) (*ListTasksResult, error) if err := json.Unmarshal(*rawMessage, &result); err != nil { return nil, fmt.Errorf("failed to unmarshal response: %w", err) } + // Fix up ttlMs/pollIntervalMs for each task: Task struct tags use legacy + // names so plain json.Unmarshal skips SEP-2663 wire names. + var envelope struct { + Tasks []json.RawMessage `json:"tasks"` + } + if err := json.Unmarshal(*rawMessage, &envelope); err == nil { + for i, raw := range envelope.Tasks { + if i >= len(result.Tasks) { + break + } + tmp := GetTaskResult{Task: result.Tasks[i]} + jsonToTask(raw, &tmp) + result.Tasks[i] = tmp.Task + } + } return &result, nil } @@ -1402,5 +1451,31 @@ func ParseGetTaskResult(rawMessage *json.RawMessage) (*GetTaskResult, error) { if err := json.Unmarshal(*rawMessage, &result); err != nil { return nil, fmt.Errorf("failed to unmarshal response: %w", err) } + // Fix up ttlMs/pollIntervalMs from SEP-2663 wire format. + jsonToTask(*rawMessage, &result) + // Populate SEP-2663 inline fields (result, error, inputRequests). + var sep struct { + TaskResult json.RawMessage `json:"result"` + TaskError *JSONRPCErrorDetails `json:"error"` + InputRequests InputRequests `json:"inputRequests"` + } + if err := json.Unmarshal(*rawMessage, &sep); err != nil { + return nil, fmt.Errorf("failed to unmarshal SEP-2663 task fields: %w", err) + } + result.TaskResult = sep.TaskResult + result.TaskError = sep.TaskError + result.InputRequests = sep.InputRequests + return &result, nil +} + +// ParseUpdateTaskResult parses a JSON message and converts it to an UpdateTaskResult. +func ParseUpdateTaskResult(rawMessage *json.RawMessage) (*UpdateTaskResult, error) { + if rawMessage == nil { + return nil, fmt.Errorf("response is nil") + } + var result UpdateTaskResult + if err := json.Unmarshal(*rawMessage, &result); err != nil { + return nil, fmt.Errorf("failed to unmarshal response: %w", err) + } return &result, nil } diff --git a/mcp/utils_additional_test.go b/mcp/utils_additional_test.go index 29987b7f2..2d301fefb 100644 --- a/mcp/utils_additional_test.go +++ b/mcp/utils_additional_test.go @@ -1110,3 +1110,20 @@ func TestToolResultContent_IsContent(t *testing.T) { var c Content = ToolResultContent{Type: ContentTypeToolResult, ToolUseID: "tu_1"} assert.NotNil(t, c) } + +// TestParseCallToolResult_TaskHandle verifies that a server-directed task +// handle (resultType:"task", SEP-2663) is surfaced as a typed +// *UnsupportedTaskResultError carrying the task ID, rather than failing with +// "content is missing". The caller can then poll tasks/get with that ID. +func TestParseCallToolResult_TaskHandle(t *testing.T) { + // Wire shape of a modern CreateTaskResult: resultType "task", taskId at the + // top level, and no content. + raw := json.RawMessage(`{"resultType":"task","taskId":"task-abc","status":"working","createdAt":"2026-01-01T00:00:00Z","lastUpdatedAt":"2026-01-01T00:00:00Z"}`) + + _, err := ParseCallToolResult(&raw) + require.Error(t, err) + + var taskErr *UnsupportedTaskResultError + require.True(t, errors.As(err, &taskErr), "expected *UnsupportedTaskResultError, got %T: %v", err, err) + assert.Equal(t, "task-abc", taskErr.TaskID) +} diff --git a/server/discover.go b/server/discover.go index 584f8e42b..31d3e6188 100644 --- a/server/discover.go +++ b/server/discover.go @@ -78,7 +78,7 @@ func (s *MCPServer) handleDiscover( ) (*mcp.DiscoverResult, *requestError) { result := &mcp.DiscoverResult{ SupportedVersions: SupportedProtocolVersionsFromContext(ctx), - Capabilities: s.serverCapabilitiesSnapshot(), + Capabilities: s.serverCapabilitiesSnapshot(IsModernRequest(ctx)), Instructions: s.instructions, } diff --git a/server/handler_errors.go b/server/handler_errors.go index 3cf11c654..27d55573b 100644 --- a/server/handler_errors.go +++ b/server/handler_errors.go @@ -77,6 +77,29 @@ func handlerError(ctx context.Context, id any, err error) *requestError { return &requestError{id: id, code: mcp.INTERNAL_ERROR, err: err} } +// taskErrorDetails converts a failed task's stored handler error into the +// JSON-RPC error details surfaced on the modern tasks/get path. It mirrors +// handlerError's extraction so a task failure preserves the same code, +// client-facing message, and data an exposed handler error would carry on the +// synchronous path, rather than flattening everything to INTERNAL_ERROR. +func taskErrorDetails(err error) *mcp.JSONRPCErrorDetails { + if _, ok := asURLElicitationRequired(err); ok { + return &mcp.JSONRPCErrorDetails{ + Code: mcp.INTERNAL_ERROR, + Message: err.Error(), + } + } + var exposed *exposedHandlerError + if errors.As(err, &exposed) { + details := exposed.provider.JSONRPCError().Error + return &details + } + return &mcp.JSONRPCErrorDetails{ + Code: mcp.INTERNAL_ERROR, + Message: err.Error(), + } +} + // asURLElicitationRequired finds a URLElicitationRequiredError in err's chain, // whether the handler returned the value or a pointer to it. func asURLElicitationRequired(err error) (mcp.URLElicitationRequiredError, bool) { diff --git a/server/hooks.go b/server/hooks.go index 0e5823497..deda2ca02 100644 --- a/server/hooks.go +++ b/server/hooks.go @@ -116,6 +116,9 @@ type OnAfterCallToolFunc func(ctx context.Context, id any, message *mcp.CallTool type OnBeforeGetTaskFunc func(ctx context.Context, id any, message *mcp.GetTaskRequest) type OnAfterGetTaskFunc func(ctx context.Context, id any, message *mcp.GetTaskRequest, result *mcp.GetTaskResult) +type OnBeforeUpdateTaskFunc func(ctx context.Context, id any, message *mcp.UpdateTaskRequest) +type OnAfterUpdateTaskFunc func(ctx context.Context, id any, message *mcp.UpdateTaskRequest, result *mcp.UpdateTaskResult) + type OnBeforeListTasksFunc func(ctx context.Context, id any, message *mcp.ListTasksRequest) type OnAfterListTasksFunc func(ctx context.Context, id any, message *mcp.ListTasksRequest, result *mcp.ListTasksResult) @@ -199,6 +202,8 @@ type Hooks struct { OnAfterCallTool []OnAfterCallToolFunc OnBeforeGetTask []OnBeforeGetTaskFunc OnAfterGetTask []OnAfterGetTaskFunc + OnBeforeUpdateTask []OnBeforeUpdateTaskFunc + OnAfterUpdateTask []OnAfterUpdateTaskFunc OnBeforeListTasks []OnBeforeListTasksFunc OnAfterListTasks []OnAfterListTasksFunc OnBeforeTaskResult []OnBeforeTaskResultFunc @@ -754,6 +759,33 @@ func (c *Hooks) afterGetTask(ctx context.Context, id any, message *mcp.GetTaskRe hook(ctx, id, message, result) } } +func (c *Hooks) AddBeforeUpdateTask(hook OnBeforeUpdateTaskFunc) { + c.OnBeforeUpdateTask = append(c.OnBeforeUpdateTask, hook) +} + +func (c *Hooks) AddAfterUpdateTask(hook OnAfterUpdateTaskFunc) { + c.OnAfterUpdateTask = append(c.OnAfterUpdateTask, hook) +} + +func (c *Hooks) beforeUpdateTask(ctx context.Context, id any, message *mcp.UpdateTaskRequest) { + c.beforeAny(ctx, id, mcp.MethodTasksUpdate, message) + if c == nil { + return + } + for _, hook := range c.OnBeforeUpdateTask { + hook(ctx, id, message) + } +} + +func (c *Hooks) afterUpdateTask(ctx context.Context, id any, message *mcp.UpdateTaskRequest, result *mcp.UpdateTaskResult) { + c.onSuccess(ctx, id, mcp.MethodTasksUpdate, message, result) + if c == nil { + return + } + for _, hook := range c.OnAfterUpdateTask { + hook(ctx, id, message, result) + } +} func (c *Hooks) AddBeforeListTasks(hook OnBeforeListTasksFunc) { c.OnBeforeListTasks = append(c.OnBeforeListTasks, hook) } diff --git a/server/internal/gen/data.go b/server/internal/gen/data.go index 2fb4dc7de..07b5c1449 100644 --- a/server/internal/gen/data.go +++ b/server/internal/gen/data.go @@ -19,6 +19,9 @@ type MCPRequestType struct { // RequiresModern marks methods introduced by protocol version 2026-07-28. // They are rejected with METHOD_NOT_FOUND on legacy requests. RequiresModern bool + // RequiresExtension marks methods that require the client to have declared + // the named MCP extension SEP-2663. + RequiresExtension string } var MCPRequestTypes = []MCPRequestType{ @@ -166,15 +169,26 @@ var MCPRequestTypes = []MCPRequestType{ ResultIsAny: true, // Returns 'any' to support both CallToolResult and CreateTaskResult HasMeta: true, }, { - MethodName: "MethodTasksGet", - ParamType: "GetTaskRequest", - ResultType: "GetTaskResult", - Group: "tasks", - GroupName: "Tasks", - GroupHookName: "Task", - HookName: "GetTask", - UnmarshalError: "invalid get task request", - HandlerFunc: "handleGetTask", + MethodName: "MethodTasksGet", + ParamType: "GetTaskRequest", + ResultType: "GetTaskResult", + Group: "tasks", + GroupName: "Tasks", + GroupHookName: "Task", + HookName: "GetTask", + UnmarshalError: "invalid get task request", + HandlerFunc: "handleGetTask", + RequiresExtension: "mcp.ExtensionTasks", + }, { + MethodName: "MethodTasksUpdate", + ParamType: "UpdateTaskRequest", + ResultType: "UpdateTaskResult", + GroupName: "Tasks", + HookName: "UpdateTask", + UnmarshalError: "invalid update task request", + HandlerFunc: "handleUpdateTask", + RequiresModern: true, + RequiresExtension: "mcp.ExtensionTasks", }, { MethodName: "MethodTasksList", ParamType: "ListTasksRequest", @@ -198,15 +212,16 @@ var MCPRequestTypes = []MCPRequestType{ HandlerFunc: "handleTaskResult", RemovedInModern: true, }, { - MethodName: "MethodTasksCancel", - ParamType: "CancelTaskRequest", - ResultType: "CancelTaskResult", - Group: "tasks", - GroupName: "Tasks", - GroupHookName: "Task", - HookName: "CancelTask", - UnmarshalError: "invalid cancel task request", - HandlerFunc: "handleCancelTask", + MethodName: "MethodTasksCancel", + ParamType: "CancelTaskRequest", + ResultType: "CancelTaskResult", + Group: "tasks", + GroupName: "Tasks", + GroupHookName: "Task", + HookName: "CancelTask", + UnmarshalError: "invalid cancel task request", + HandlerFunc: "handleCancelTask", + RequiresExtension: "mcp.ExtensionTasks", }, { MethodName: "MethodCompletionComplete", ParamType: "CompleteRequest", diff --git a/server/internal/gen/request_handler.go.tmpl b/server/internal/gen/request_handler.go.tmpl index 1fe3245cf..8c52a36f1 100644 --- a/server/internal/gen/request_handler.go.tmpl +++ b/server/internal/gen/request_handler.go.tmpl @@ -188,13 +188,33 @@ func (s *MCPServer) HandleMessage( code: mcp.METHOD_NOT_FOUND, err: fmt.Errorf("%q %w", baseMessage.Method, ErrRemovedInProtocolVersion), } - } else {{ end }}{{ if .RequiresModern }}if !protocolInfo.Modern { + } else {{ end }}{{ if and .RequiresModern .RequiresExtension }}if !protocolInfo.Modern { err = &requestError{ id: baseMessage.ID, code: mcp.METHOD_NOT_FOUND, err: fmt.Errorf("%q %w", baseMessage.Method, ErrRequiresModernProtocol), } - } else {{ end }}{{ if .Group }}if s.capabilities.{{.Group}} == nil { + } else if clientCaps := protocolInfo.ClientCapabilities; clientCaps == nil || !clientCaps.HasExtension({{.RequiresExtension}}) { + err = &requestError{ + id: baseMessage.ID, + code: mcp.MISSING_REQUIRED_CLIENT_CAPABILITY, + err: fmt.Errorf("{{toLower .GroupName}} extension %w", ErrUnsupported), + data: missingExtensionData({{.RequiresExtension}}), + } + } else {{ else }}{{ if .RequiresModern }}if !protocolInfo.Modern { + err = &requestError{ + id: baseMessage.ID, + code: mcp.METHOD_NOT_FOUND, + err: fmt.Errorf("%q %w", baseMessage.Method, ErrRequiresModernProtocol), + } + } else {{ end }}{{ if .RequiresExtension }}if clientCaps := protocolInfo.ClientCapabilities; protocolInfo.Modern && (clientCaps == nil || !clientCaps.HasExtension({{.RequiresExtension}})) { + err = &requestError{ + id: baseMessage.ID, + code: mcp.MISSING_REQUIRED_CLIENT_CAPABILITY, + err: fmt.Errorf("{{toLower .GroupName}} extension %w", ErrUnsupported), + data: missingExtensionData({{.RequiresExtension}}), + } + } else {{ end }}{{ end }}{{ if .Group }}if {{ if .RequiresExtension }}!protocolInfo.Modern && {{ end }}s.capabilities.{{.Group}} == nil { err = &requestError{ id: baseMessage.ID, code: mcp.METHOD_NOT_FOUND, diff --git a/server/request_handler.go b/server/request_handler.go index 97fc8251e..93a102bd0 100644 --- a/server/request_handler.go +++ b/server/request_handler.go @@ -563,7 +563,14 @@ func (s *MCPServer) HandleMessage( case mcp.MethodTasksGet: var request mcp.GetTaskRequest var result *mcp.GetTaskResult - if s.capabilities.tasks == nil { + if clientCaps := protocolInfo.ClientCapabilities; protocolInfo.Modern && (clientCaps == nil || !clientCaps.HasExtension(mcp.ExtensionTasks)) { + err = &requestError{ + id: baseMessage.ID, + code: mcp.MISSING_REQUIRED_CLIENT_CAPABILITY, + err: fmt.Errorf("tasks extension %w", ErrUnsupported), + data: missingExtensionData(mcp.ExtensionTasks), + } + } else if !protocolInfo.Modern && s.capabilities.tasks == nil { err = &requestError{ id: baseMessage.ID, code: mcp.METHOD_NOT_FOUND, @@ -586,6 +593,39 @@ func (s *MCPServer) HandleMessage( } s.hooks.afterGetTask(ctx, baseMessage.ID, &request, result) return createResponse(baseMessage.ID, *result) + case mcp.MethodTasksUpdate: + var request mcp.UpdateTaskRequest + var result *mcp.UpdateTaskResult + if !protocolInfo.Modern { + err = &requestError{ + id: baseMessage.ID, + code: mcp.METHOD_NOT_FOUND, + err: fmt.Errorf("%q %w", baseMessage.Method, ErrRequiresModernProtocol), + } + } else if clientCaps := protocolInfo.ClientCapabilities; clientCaps == nil || !clientCaps.HasExtension(mcp.ExtensionTasks) { + err = &requestError{ + id: baseMessage.ID, + code: mcp.MISSING_REQUIRED_CLIENT_CAPABILITY, + err: fmt.Errorf("tasks extension %w", ErrUnsupported), + data: missingExtensionData(mcp.ExtensionTasks), + } + } else if unmarshalErr := json.Unmarshal(message, &request); unmarshalErr != nil { + err = &requestError{ + id: baseMessage.ID, + code: mcp.INVALID_REQUEST, + err: &UnparsableMessageError{message: message, err: unmarshalErr, method: baseMessage.Method}, + } + } else { + request.Header = headers + s.hooks.beforeUpdateTask(ctx, baseMessage.ID, &request) + result, err = s.handleUpdateTask(ctx, baseMessage.ID, request) + } + if err != nil { + s.hooks.onError(ctx, baseMessage.ID, baseMessage.Method, &request, err) + return err.ToJSONRPCError() + } + s.hooks.afterUpdateTask(ctx, baseMessage.ID, &request, result) + return createResponse(baseMessage.ID, *result) case mcp.MethodTasksList: var request mcp.ListTasksRequest var result *mcp.ListTasksResult @@ -653,7 +693,14 @@ func (s *MCPServer) HandleMessage( case mcp.MethodTasksCancel: var request mcp.CancelTaskRequest var result *mcp.CancelTaskResult - if s.capabilities.tasks == nil { + if clientCaps := protocolInfo.ClientCapabilities; protocolInfo.Modern && (clientCaps == nil || !clientCaps.HasExtension(mcp.ExtensionTasks)) { + err = &requestError{ + id: baseMessage.ID, + code: mcp.MISSING_REQUIRED_CLIENT_CAPABILITY, + err: fmt.Errorf("tasks extension %w", ErrUnsupported), + data: missingExtensionData(mcp.ExtensionTasks), + } + } else if !protocolInfo.Modern && s.capabilities.tasks == nil { err = &requestError{ id: baseMessage.ID, code: mcp.METHOD_NOT_FOUND, diff --git a/server/sep2663_test.go b/server/sep2663_test.go new file mode 100644 index 000000000..ef495ad93 --- /dev/null +++ b/server/sep2663_test.go @@ -0,0 +1,241 @@ +package server + +import ( + "context" + "encoding/json" + "testing" + "time" + + "github.com/mark3labs/mcp-go/mcp" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// modernCallParams builds a _meta block for a modern protocol +// request that declares the tasks extension. +func modernTaskMeta() map[string]any { + return map[string]any{ + "io.modelcontextprotocol/protocolVersion": "2026-07-28", + "io.modelcontextprotocol/clientCapabilities": map[string]any{ + "extensions": map[string]any{ + "io.modelcontextprotocol/tasks": map[string]any{}, + }, + }, + } +} + +// TestSEP2663_ServerDirectedTaskCreation verifies that a modern client can +// trigger task creation without setting the legacy "task" param on the request. +func TestSEP2663_ServerDirectedTaskCreation(t *testing.T) { + srv := NewMCPServer("test", "1.0.0", + WithTaskCapabilities(true, true, true), + WithTasksExtension(), + ) + + tool := mcp.NewTool("async_work", + mcp.WithDescription("Async tool"), + mcp.WithTaskSupport(mcp.TaskSupportOptional), + ) + srv.AddTool(tool, func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return mcp.NewToolResultText("sync result"), nil + }) + + metaBytes, _ := json.Marshal(modernTaskMeta()) + req := `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"async_work","_meta":` + string(metaBytes) + `}}` + + response := srv.HandleMessage(t.Context(), []byte(req)) + resp, ok := response.(mcp.JSONRPCResponse) + require.True(t, ok, "Expected JSONRPCResponse, got %T", response) + + result, ok := resp.Result.(*mcp.CreateTaskResult) + require.True(t, ok, "Expected *CreateTaskResult (server-directed), got %T", resp.Result) + assert.Equal(t, mcp.ResultTypeTask, result.GetResultType()) + assert.NotEmpty(t, result.Task.TaskId) + assert.Equal(t, mcp.TaskStatusWorking, result.Task.Status) + assert.False(t, result.Legacy, "Modern client should get SEP-2663 format") +} + +// TestSEP2663_CreateTaskResultWireFormat verifies the JSON wire shape for +// modern clients matches the SEP-2663 spec (inline fields, resultType:"task"). +func TestSEP2663_CreateTaskResultWireFormat(t *testing.T) { + result := mcp.CreateTaskResult{ + Task: mcp.Task{ + TaskId: "abc123", + Status: mcp.TaskStatusWorking, + CreatedAt: "2026-09-04T12:00:00Z", + LastUpdatedAt: "2026-09-04T12:00:00Z", + }, + } + result.SetResultType(mcp.ResultTypeTask) + + data, err := json.Marshal(result) + require.NoError(t, err) + + var m map[string]any + require.NoError(t, json.Unmarshal(data, &m)) + + assert.Equal(t, "task", m["resultType"]) + assert.Equal(t, "abc123", m["taskId"]) + assert.Equal(t, "working", m["status"]) + assert.Equal(t, "2026-09-04T12:00:00Z", m["createdAt"]) + assert.NotContains(t, m, "task", "task fields must be inline, not wrapped") +} + +// TestSEP2663_LegacyCreateTaskResultWireFormat verifies the JSON wire shape +// for legacy 2025-11-25 clients uses the old "task" wrapper key. +func TestSEP2663_LegacyCreateTaskResultWireFormat(t *testing.T) { + result := mcp.CreateTaskResult{ + Task: mcp.Task{ + TaskId: "abc123", + Status: mcp.TaskStatusWorking, + CreatedAt: "2026-09-04T12:00:00Z", + LastUpdatedAt: "2026-09-04T12:00:00Z", + }, + Legacy: true, + } + + data, err := json.Marshal(result) + require.NoError(t, err) + + var m map[string]any + require.NoError(t, json.Unmarshal(data, &m)) + + assert.Contains(t, m, "task", "legacy format must wrap under 'task' key") + assert.NotContains(t, m, "resultType") + taskMap, ok := m["task"].(map[string]any) + require.True(t, ok) + assert.Equal(t, "abc123", taskMap["taskId"]) +} + +// TestSEP2663_GetTaskInlinesResult verifies that tasks/get for a completed +// task inlines the result field in the response for modern clients. +func TestSEP2663_GetTaskInlinesResult(t *testing.T) { + srv := NewMCPServer("test", "1.0.0", + WithTaskCapabilities(true, true, true), + ) + + modernCtx := buildModernContext(t.Context(), srv) + ttl := int64(60000) + entry, err := srv.createTask(modernCtx, "task-inline", "tool", &ttl, nil) + require.NoError(t, err) + + // Complete the task with a tool result + toolResult := &mcp.CallToolResult{ + Content: []mcp.Content{mcp.NewTextContent("hello world")}, + } + srv.completeTask(entry, toolResult, nil) + + // Wait for completion + for range 100 { + t, _, _ := srv.getTask(modernCtx, "task-inline") + if t.Status.IsTerminal() { + break + } + time.Sleep(5 * time.Millisecond) + } + + getResult, reqErr := srv.handleGetTask(modernCtx, 1, mcp.GetTaskRequest{ + Params: mcp.GetTaskParams{TaskId: "task-inline"}, + }) + require.Nil(t, reqErr) + require.NotNil(t, getResult) + + assert.Equal(t, mcp.TaskStatusCompleted, getResult.Status) + assert.NotNil(t, getResult.TaskResult, "completed task must have inlined result") +} + +// TestSEP2663_CancelTaskAckOnly verifies that tasks/cancel for a modern +// client returns an ack-only result, including for an already-terminal task. +func TestSEP2663_CancelTaskAckOnly(t *testing.T) { + srv := NewMCPServer("test", "1.0.0", + WithTaskCapabilities(true, true, true), + ) + + modernCtx := buildModernContext(t.Context(), srv) + // Create under the same modern context so the task's protocol origin matches + // the canceller (cross-origin isolation rejects mismatches). + entry, err := srv.createTask(modernCtx, "task-cancel-ack", "tool", nil, nil) + require.NoError(t, err) + assert.Equal(t, mcp.TaskStatusWorking, entry.task.Status) + + result, reqErr := srv.handleCancelTask(modernCtx, 1, mcp.CancelTaskRequest{ + Params: mcp.CancelTaskParams{TaskId: "task-cancel-ack"}, + }) + require.Nil(t, reqErr) + require.NotNil(t, result) + + // Verify it serializes as an empty object (no task fields). + data, _ := json.Marshal(result) + var m map[string]any + require.NoError(t, json.Unmarshal(data, &m)) + assert.NotContains(t, m, "taskId", "ack-only result must not have task fields") + + // Cancelling an already-terminal task on the modern path must still ack silently, not error. + result, reqErr = srv.handleCancelTask(modernCtx, 2, mcp.CancelTaskRequest{ + Params: mcp.CancelTaskParams{TaskId: "task-cancel-ack"}, + }) + require.Nil(t, reqErr, "modern cancel of a terminal task must ack, not error") + require.NotNil(t, result) +} + +// TestSEP2663_TasksUpdateAck verifies that tasks/update returns an ack. +func TestSEP2663_TasksUpdateAck(t *testing.T) { + srv := NewMCPServer("test", "1.0.0", + WithTaskCapabilities(true, true, true), + ) + + modernCtx := buildModernContext(t.Context(), srv) + entry, err := srv.createTask(modernCtx, "task-update-ack", "tool", nil, nil) + require.NoError(t, err) + + // tasks/update is only valid when the task is in input_required state. + srv.tasksMu.Lock() + entry.task.Status = mcp.TaskStatusInputRequired + srv.tasksMu.Unlock() + + updateResult, reqErr := srv.handleUpdateTask(modernCtx, 1, mcp.UpdateTaskRequest{ + Params: mcp.UpdateTaskParams{ + TaskId: "task-update-ack", + InputResponses: mcp.InputResponses{}, + }, + }) + require.Nil(t, reqErr) + require.NotNil(t, updateResult) + assert.Equal(t, mcp.ResultTypeComplete, updateResult.GetResultType()) +} + +// TestSEP2663_HasExtension tests the ClientCapabilities.HasExtension helper. +func TestSEP2663_HasExtension(t *testing.T) { + t.Run("nil caps", func(t *testing.T) { + var caps *mcp.ClientCapabilities + assert.False(t, caps.HasExtension(mcp.ExtensionTasks)) + }) + t.Run("no extensions", func(t *testing.T) { + caps := &mcp.ClientCapabilities{} + assert.False(t, caps.HasExtension(mcp.ExtensionTasks)) + }) + t.Run("extension present", func(t *testing.T) { + caps := &mcp.ClientCapabilities{ + Extensions: map[string]any{mcp.ExtensionTasks: map[string]any{}}, + } + assert.True(t, caps.HasExtension(mcp.ExtensionTasks)) + }) + t.Run("different extension", func(t *testing.T) { + caps := &mcp.ClientCapabilities{ + Extensions: map[string]any{"some/other": map[string]any{}}, + } + assert.False(t, caps.HasExtension(mcp.ExtensionTasks)) + }) +} + +// buildModernContext creates a context that simulates a modern protocol request +// with the tasks extension declared. +func buildModernContext(ctx context.Context, _ *MCPServer) context.Context { + return WithRequestProtocolInfo(ctx, &RequestProtocolInfo{ + Modern: true, + ProtocolVersion: "2026-07-28", + ClientCapabilities: &mcp.ClientCapabilities{ + Extensions: map[string]any{mcp.ExtensionTasks: map[string]any{}}, + }, + }) +} diff --git a/server/server.go b/server/server.go index 3bb6ec4ad..6fdd44ecc 100644 --- a/server/server.go +++ b/server/server.go @@ -40,11 +40,19 @@ type taskEntry struct { sessionID string toolName string // Name of the tool that created this task createdAt time.Time // When the task was created (for metrics) + modernTask bool // Whether the task was created by a modern (SEP-2663) request result any // The actual result once completed resultErr error // Error if task failed cancelFunc context.CancelFunc // Function to cancel the task done chan struct{} // Channel to signal task completion completed bool // Whether the task has been completed (guards done channel closure) + // inputResponses holds responses delivered via tasks/update (SEP-2663). + inputResponses mcp.InputResponses + // inputReady is signalled each time tasks/update delivers new inputResponses. + inputReady chan struct{} + // inputRequests holds the pending input requests the task is waiting on + // when it is in the input_required state. + inputRequests mcp.InputRequests // createdHookDone, under tasksMu, says the created hook has returned; // until then the hook of a completion or cancellation waits in // pendingHooks, which createTask runs after the created hook. A task's @@ -195,6 +203,18 @@ func (e *requestError) Unwrap() error { return e.err } +// missingExtensionData builds the error data naming the required extension. +func missingExtensionData(extension string) any { + return map[string]any{ + "capability": extension, + "requiredCapabilities": map[string]any{ + "extensions": map[string]any{ + extension: map[string]any{}, + }, + }, + } +} + // NotificationHandlerFunc handles incoming notifications. type NotificationHandlerFunc func(ctx context.Context, notification mcp.JSONRPCNotification) @@ -255,6 +275,7 @@ type MCPServer struct { tasks map[string]*taskEntry expiredTasks map[string]time.Time // Tracks recently expired task IDs with expiration timestamp maxConcurrentTasks *int // Optional limit on concurrent running tasks + defaultTaskTTLMs *int64 // Fallback TTL applied when client supplies none activeTasks int // Current count of running (non-terminal) tasks inflightCancels sync.Map // Maps request ID -> context.CancelFunc for in-flight requests inputValidator *inputSchemaValidator @@ -677,6 +698,17 @@ func WithTaskCapabilities(list, cancel, toolCallTasks bool) ServerOption { } } +// WithDefaultTaskTTL sets a fallback TTL applied to every task created without +// a client-supplied TTL. Ensures tasks are evicted even when modern SEP-2663 +// clients do not pass a task param. +func WithDefaultTaskTTL(ttlMs int64) ServerOption { + return func(s *MCPServer) { + if ttlMs > 0 { + s.defaultTaskTTLMs = &ttlMs + } + } +} + // WithInstructions sets the server instructions for the client returned in the initialize response func WithInstructions(instructions string) ServerOption { return func(s *MCPServer) { @@ -716,6 +748,19 @@ func WithExtensions(extensions map[string]any) ServerOption { } } +// WithTasksExtension advertises support for the MCP Tasks extension +// in the server's capabilities. +func WithTasksExtension() ServerOption { + return func(s *MCPServer) { + s.capabilitiesMu.Lock() + defer s.capabilitiesMu.Unlock() + if s.capabilities.extensions == nil { + s.capabilities.extensions = make(map[string]any) + } + s.capabilities.extensions[mcp.ExtensionTasks] = map[string]any{} + } +} + // WithIcons sets the server icons for the implementation metadata returned // during initialization. The icons slice and nested Sizes fields are defensively // copied to prevent external mutation. @@ -1266,7 +1311,12 @@ func (s *MCPServer) AddNotificationHandler( // serverCapabilities builds the capability set advertised by this server. // It is shared by the legacy initialize handshake and the modern // server/discover RPC. -func (s *MCPServer) serverCapabilitiesSnapshot() mcp.ServerCapabilities { +// +// When modern is true the legacy tasks.* capability block is suppressed: those +// capabilities are not part of the modern protocol vocabulary (SEP-2663), and +// modern clients discover task support via the io.modelcontextprotocol/tasks +// extension in Extensions instead. +func (s *MCPServer) serverCapabilitiesSnapshot(modern bool) mcp.ServerCapabilities { // Capabilities are registered at runtime - AddTool implicitly enables the // tools capability, for example - so reads must be synchronized against // those writers. @@ -1320,8 +1370,8 @@ func (s *MCPServer) serverCapabilitiesSnapshot() mcp.ServerCapabilities { capabilities.Roots = &struct{}{} } - // Only add task capabilities if they're configured - if s.capabilities.tasks != nil { + // Only add the legacy task capability block on the legacy path. + if s.capabilities.tasks != nil && !modern { tasksCapability := &mcp.TasksCapability{} if s.capabilities.tasks.list { @@ -1377,7 +1427,9 @@ func (s *MCPServer) handleInitialize( _ any, request mcp.InitializeRequest, ) (*mcp.InitializeResult, *requestError) { - capabilities := s.serverCapabilitiesSnapshot() + // The initialize handshake exists only on the legacy protocol, so the + // legacy task capability block is always advertised here. + capabilities := s.serverCapabilitiesSnapshot(false) result := mcp.InitializeResult{ ProtocolVersion: s.protocolVersion(request.Params.ProtocolVersion), @@ -2256,8 +2308,24 @@ func (s *MCPServer) handleToolCall( } // Validate task support requirements + clientDeclaredExtension := false + if info := RequestProtocolInfoFromContext(ctx); info != nil && info.Modern && info.ClientCapabilities != nil { + clientDeclaredExtension = info.ClientCapabilities.HasExtension(mcp.ExtensionTasks) + } + if tool.Tool.Execution != nil && tool.Tool.Execution.TaskSupport == mcp.TaskSupportRequired { - if request.Params.Task == nil { + // Modern protocol: client must declare the extension. + if info := RequestProtocolInfoFromContext(ctx); info != nil && info.Modern { + if !clientDeclaredExtension { + return nil, &requestError{ + id: id, + code: mcp.MISSING_REQUIRED_CLIENT_CAPABILITY, + err: fmt.Errorf("tool '%s' requires the %s extension", request.Params.Name, mcp.ExtensionTasks), + data: missingExtensionData(mcp.ExtensionTasks), + } + } + } else if request.Params.Task == nil { + // Legacy: client must set the task param. return nil, &requestError{ id: id, code: mcp.METHOD_NOT_FOUND, @@ -2266,12 +2334,40 @@ func (s *MCPServer) handleToolCall( } } + // Validate arguments against the tool's input schema (when enabled via + // WithInputSchemaValidation), before task routing so both sync and task + // paths are covered. + if s.inputValidator != nil { + if _, err := s.inputValidator.validate(tool.Tool, request.GetRawArguments()); err != nil { + return validationToolResult(err), nil + } + } + // Check if this should be executed as a task (hybrid mode support) // Tools with TaskSupportOptional or TaskSupportRequired can be executed as tasks - shouldExecuteAsTask := request.Params.Task != nil && - tool.Tool.Execution != nil && - (tool.Tool.Execution.TaskSupport == mcp.TaskSupportOptional || - tool.Tool.Execution.TaskSupport == mcp.TaskSupportRequired) + var shouldExecuteAsTask bool + if info := RequestProtocolInfoFromContext(ctx); info != nil && info.Modern { + shouldExecuteAsTask = clientDeclaredExtension && + tool.Tool.Execution != nil && + (tool.Tool.Execution.TaskSupport == mcp.TaskSupportOptional || + tool.Tool.Execution.TaskSupport == mcp.TaskSupportRequired) + } else { + shouldExecuteAsTask = request.Params.Task != nil && + tool.Tool.Execution != nil && + (tool.Tool.Execution.TaskSupport == mcp.TaskSupportOptional || + tool.Tool.Execution.TaskSupport == mcp.TaskSupportRequired) + } + + // Task-only tools (AddTaskTool) must always go through the task path. + // On modern protocol with extension declared, or legacy with task param set. + // Otherwise fall through to the error below. + if taskToolOnly { + if info := RequestProtocolInfoFromContext(ctx); info != nil && info.Modern { + shouldExecuteAsTask = clientDeclaredExtension + } else if request.Params.Task != nil { + shouldExecuteAsTask = true + } + } if shouldExecuteAsTask { // Route to task-augmented execution handler @@ -2279,6 +2375,18 @@ func (s *MCPServer) handleToolCall( } if taskToolOnly { + // A task-only tool cannot run synchronously. For a modern client that + // has not declared the tasks extension, the request cannot be served + // without returning a CreateTaskResult, so SEP-2663 requires + // MISSING_REQUIRED_CLIENT_CAPABILITY naming the extension. + if info := RequestProtocolInfoFromContext(ctx); info != nil && info.Modern && !clientDeclaredExtension { + return nil, &requestError{ + id: id, + code: mcp.MISSING_REQUIRED_CLIENT_CAPABILITY, + err: fmt.Errorf("tool '%s' requires the %s extension", request.Params.Name, mcp.ExtensionTasks), + data: missingExtensionData(mcp.ExtensionTasks), + } + } return nil, &requestError{ id: id, code: mcp.METHOD_NOT_FOUND, @@ -2286,16 +2394,6 @@ func (s *MCPServer) handleToolCall( } } - // Validate the incoming arguments against the tool's input schema, when - // schema validation has been enabled via WithInputSchemaValidation. A - // validation failure is returned as a SEP-1303 tool execution error so the - // model receives feedback in its context window and can self-correct. - if s.inputValidator != nil { - if _, err := s.inputValidator.validate(tool.Tool, request.GetRawArguments()); err != nil { - return validationToolResult(err), nil - } - } - finalHandler := tool.Handler s.toolMiddlewareMu.RLock() @@ -2430,11 +2528,15 @@ func (s *MCPServer) handleTaskAugmentedToolCall( // Generate task ID (UUID v4) taskID := uuid.New().String() - // Extract TTL from task params + // Extract TTL from task params. var ttl *int64 - if request.Params.Task != nil { + if !IsModernRequest(ctx) && request.Params.Task != nil { ttl = request.Params.Task.TTL } + // Fall back to the server default when no client TTL applies. + if ttl == nil { + ttl = s.defaultTaskTTLMs + } // Create task entry (pollInterval is nil - server doesn't set a default) entry, err := s.createTask(ctx, taskID, request.Params.Name, ttl, nil) @@ -2454,8 +2556,9 @@ func (s *MCPServer) handleTaskAugmentedToolCall( taskCopy := entry.task s.tasksMu.RUnlock() - // Execute tool asynchronously - // For regular tools being used as tasks, we need different execution logic + // Execute tool asynchronously. The goroutine's context is rooted at + // s.serverCtx, not the per-request ctx, so it outlives the originating + // request on every transport. if hasTaskHandler { go s.executeTaskTool(ctx, entry, toolToUse, request) } else { @@ -2463,10 +2566,14 @@ func (s *MCPServer) handleTaskAugmentedToolCall( go s.executeRegularToolAsTask(ctx, entry, regularTool, request) } - // Return CreateTaskResult immediately with task as top-level field - return &mcp.CreateTaskResult{ - Task: taskCopy, - }, nil + // Return CreateTaskResult immediately. + result := &mcp.CreateTaskResult{Task: taskCopy} + if IsModernRequest(ctx) { + result.SetResultType(mcp.ResultTypeTask) + } else { + result.Legacy = true + } + return result, nil } // executeTaskTool executes a task tool handler asynchronously. @@ -2484,8 +2591,10 @@ func (s *MCPServer) executeTaskTool( } }() - // Create cancellable context for this task execution - taskCtx, cancel := context.WithCancel(ctx) + // Preserve ctx values but strip the per-request cancel. + taskCtx, cancel := context.WithCancel( + context.WithoutCancel(ctx), + ) defer cancel() // Register cancellation atomically with TTL cleanup, which may have @@ -2507,7 +2616,7 @@ func (s *MCPServer) executeTaskTool( // The cancelTask method will handle setting the proper status. // However, if cancelTask hasn't been called yet, we should still mark it. if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { - _ = s.cancelTaskEntry(ctx, entry, err.Error()) + _ = s.cancelTaskEntry(ctx, entry, err.Error(), false) return } @@ -2523,22 +2632,26 @@ func (s *MCPServer) executeTaskTool( return } - // Task succeeded - store the CreateTaskResult - // Note: The actual result will be retrieved later via tasks/result - // - // Validate the StructuredContent against the tool's declared output - // schema when WithOutputSchemaValidation is enabled. If validation fails, - // persist a tool execution error in place of the bad result so the - // client cannot retrieve a result that violates the schema via - // tasks/result. handleTaskResult accepts both *CallToolResult and - // *CreateTaskResult, so storing a *CallToolResult here is safe. if s.outputValidator != nil { if _, vErr := s.outputValidator.validateCreateTaskResult(taskTool.Tool, result); vErr != nil { s.completeTask(entry, validationToolResult(vErr), nil) return } } - s.completeTask(entry, s.limitCreateTaskResult(result), nil) + // Apply size limiting via limitCreateTaskResult, then convert to + // *CallToolResult for storage. CreateTaskResult.Content/StructuredContent/IsError + // are tagged json:"-" so marshalling *CreateTaskResult in handleGetTask would + // silently drop all tool output; *CallToolResult marshals those fields correctly. + limited := s.limitCreateTaskResult(result) + callResult := &mcp.CallToolResult{ + Result: mcp.Result{Meta: limited.Meta}, + Content: limited.Content, + IsError: limited.IsError, + } + if limited.StructuredContent != nil { + callResult.StructuredContent = limited.StructuredContent + } + s.completeTask(entry, callResult, nil) } // limitCreateTaskResult measures the payload tasks/result will copy out, @@ -2566,7 +2679,6 @@ func (s *MCPServer) limitCreateTaskResult(result *mcp.CreateTaskResult) *mcp.Cre return &stored } -// executeRegularToolAsTask executes a regular tool handler asynchronously as a task. // This is used for hybrid mode where a tool with TaskSupportOptional is called with task params. func (s *MCPServer) executeRegularToolAsTask( ctx context.Context, @@ -2574,8 +2686,10 @@ func (s *MCPServer) executeRegularToolAsTask( regularTool ServerTool, request mcp.CallToolRequest, ) { - // Create cancellable context for this task execution - taskCtx, cancel := context.WithCancel(ctx) + // Preserve ctx values but strip the per-request cancel. + taskCtx, cancel := context.WithCancel( + context.WithoutCancel(ctx), + ) defer cancel() // Register cancellation atomically with TTL cleanup, which may have @@ -2599,44 +2713,137 @@ func (s *MCPServer) executeRegularToolAsTask( finalHandler = m(finalHandler) } - result, err := finalHandler(taskCtx, request) + // Drive the SEP-2322 multi round-trip on the task path: on an + // input_required result the task parks (surfaced via tasks/get), waits for + // the answer via tasks/update, and re-invokes the handler. Like the + // synchronous resolveMultiRoundTrip, but spread across polls so the + // originating request does not block. + call := request + for attempt := 0; ; attempt++ { + result, err := finalHandler(taskCtx, call) - if err != nil { - // If the error is due to context cancellation, don't mark as failed. - // The cancelTask method will handle setting the proper status. - // However, if cancelTask hasn't been called yet, we should still mark it. - if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { - _ = s.cancelTaskEntry(ctx, entry, err.Error()) + if err != nil { + // If the error is due to context cancellation, don't mark as failed. + // The cancelTask method will handle setting the proper status. + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + _ = s.cancelTaskEntry(ctx, entry, err.Error(), false) + return + } + + // Task failed, complete with error + s.completeTask(entry, nil, err) return } - // Task failed - complete with error - s.completeTask(entry, nil, err) + // tasks/result has nothing to report for a handler that returned neither + // a result nor an error, so the task fails. + if result == nil { + s.completeTask(entry, nil, noResultError("tool", request.Params.Name)) + return + } + + if result.NeedsInput() { + // An input_required result with no requests is the "retry later" / + // load-shedding signal (see resolveMultiRoundTrip). There is nothing + // for the client to answer, so parking would hang the task forever. + // Fail fast, matching the synchronous path's ErrLoadShedding. + if len(result.InputRequests) == 0 { + s.completeTask(entry, nil, fmt.Errorf( + "tool %q returned input_required with no input requests", request.Params.Name)) + return + } + // Parking in input_required is only answerable by modern clients + // via tasks/update. + if info := RequestProtocolInfoFromContext(ctx); info == nil || !info.Modern { + // Fail the task with a clear error instead of parking. + s.completeTask(entry, nil, fmt.Errorf( + "tool %q requires client input, which is not supported for legacy task execution", request.Params.Name)) + return + } + if attempt >= maxLegacyInputRoundTrips { + s.completeTask(entry, nil, fmt.Errorf( + "tool %q exceeded %d input round trips", request.Params.Name, maxLegacyInputRoundTrips)) + return + } + responses, ok := s.awaitTaskInput(taskCtx, entry, result.InputRequests) + if !ok { + // Task was cancelled or expired while awaiting input, the + // cancellation path owns the terminal status. + _ = s.cancelTaskEntry(ctx, entry, "Task cancelled while awaiting input", false) + return + } + call = request + call.Params.MultiRoundTripParams = mcp.MultiRoundTripParams{ + InputResponses: responses, + RequestState: result.RequestState, + } + continue + } + + // Validate the result's StructuredContent against the tool's declared output + // schema when WithOutputSchemaValidation is enabled. + if s.outputValidator != nil { + if _, vErr := s.outputValidator.validate(regularTool.Tool, result); vErr != nil { + // A validation failure replaces the result with a tool execution error + // so the bad payload never reaches the client via tasks/result. + s.completeTask(entry, validationToolResult(vErr), nil) + return + } + } + s.completeTask(entry, s.limitToolResult(result), nil) return } +} - // tasks/result has nothing to report for a handler that returned neither - // a result nor an error, so the task fails. - if result == nil { - s.completeTask(entry, nil, noResultError("tool", request.Params.Name)) - return +// awaitTaskInput parks a task in input_required, publishing requests for the +// client to answer via tasks/update, and blocks until the answer arrives or the +// task is cancelled/expired (taskCtx done). +func (s *MCPServer) awaitTaskInput( + taskCtx context.Context, + entry *taskEntry, + requests mcp.InputRequests, +) (mcp.InputResponses, bool) { + s.tasksMu.Lock() + if entry.completed { + s.tasksMu.Unlock() + return nil, false + } + entry.task.Status = mcp.TaskStatusInputRequired + entry.task.LastUpdatedAt = time.Now().UTC().Format(time.RFC3339) + entry.inputRequests = requests + entry.inputResponses = nil + // Drain any stale ready signal from a previous round so we wait for a fresh + // tasks/update rather than returning immediately. + select { + case <-entry.inputReady: + default: } + task := entry.task + s.tasksMu.Unlock() - // Task succeeded - store the CallToolResult directly - // When retrieved via tasks/result, this will be returned to the client - // - // Mirror the synchronous path: validate the result's StructuredContent - // against the tool's declared output schema when - // WithOutputSchemaValidation is enabled. A validation failure replaces - // the result with a tool execution error so the bad payload never - // reaches the client via tasks/result. - if s.outputValidator != nil { - if _, vErr := s.outputValidator.validate(regularTool.Tool, result); vErr != nil { - s.completeTask(entry, validationToolResult(vErr), nil) - return - } + s.sendTaskStatusNotification(task, entry.modernTask) + + select { + case <-entry.inputReady: + case <-taskCtx.Done(): + return nil, false } - s.completeTask(entry, s.limitToolResult(result), nil) + + s.tasksMu.Lock() + if entry.completed { + s.tasksMu.Unlock() + return nil, false + } + responses := entry.inputResponses + entry.inputResponses = nil + entry.inputRequests = nil + entry.task.Status = mcp.TaskStatusWorking + entry.task.LastUpdatedAt = time.Now().UTC().Format(time.RFC3339) + task = entry.task + s.tasksMu.Unlock() + + s.sendTaskStatusNotification(task, entry.modernTask) + return responses, true } func (s *MCPServer) handleNotification( @@ -2702,6 +2909,55 @@ func (s *MCPServer) handleGetTask( id any, request mcp.GetTaskRequest, ) (*mcp.GetTaskResult, *requestError) { + // For modern protocol (SEP-2663). + if IsModernRequest(ctx) { + entry, err := s.getTaskEntry(ctx, request.Params.TaskId) + if err != nil { + return nil, &requestError{ + id: id, + code: mcp.INVALID_PARAMS, + err: err, + } + } + + s.tasksMu.RLock() + taskCopy := entry.task + storedResult := entry.result + resultErr := entry.resultErr + pendingInputRequests := entry.inputRequests + s.tasksMu.RUnlock() + + result := mcp.NewGetTaskResult(taskCopy) + result.SetResultType(mcp.ResultTypeComplete) + + switch taskCopy.Status { + case mcp.TaskStatusCompleted: + if storedResult != nil { + raw, marshalErr := json.Marshal(storedResult) + if marshalErr != nil { + // The result cannot be encoded. Report it as a task error + // rather than returning a successful but empty result, + // which would silently drop the tool output. + return nil, &requestError{ + id: id, + code: mcp.INTERNAL_ERROR, + err: fmt.Errorf("task %q result is not JSON-encodable: %w", request.Params.TaskId, marshalErr), + } + } + result.TaskResult = json.RawMessage(raw) + } + case mcp.TaskStatusFailed: + if resultErr != nil { + result.TaskError = taskErrorDetails(resultErr) + } + case mcp.TaskStatusInputRequired: + result.InputRequests = pendingInputRequests + } + + return &result, nil + } + + // Legacy path: return task state only. task, _, err := s.getTask(ctx, request.Params.TaskId) if err != nil { return nil, &requestError{ @@ -2711,8 +2967,52 @@ func (s *MCPServer) handleGetTask( } } - result := mcp.NewGetTaskResult(task) - return &result, nil + r := mcp.NewGetTaskResult(task) + r.Legacy = true + return &r, nil +} + +// handleUpdateTask handles tasks/update requests. +func (s *MCPServer) handleUpdateTask( + ctx context.Context, + id any, + request mcp.UpdateTaskRequest, +) (*mcp.UpdateTaskResult, *requestError) { + entry, err := s.getTaskEntry(ctx, request.Params.TaskId) + if err != nil { + return nil, &requestError{ + id: id, + code: mcp.INVALID_PARAMS, + err: err, + } + } + + // Deliver responses to the task's pending input channel. + s.tasksMu.Lock() + if status := entry.task.Status; status != mcp.TaskStatusInputRequired { + s.tasksMu.Unlock() + return nil, &requestError{ + id: id, + code: mcp.INVALID_PARAMS, + err: fmt.Errorf("task %q is not awaiting input (status: %s)", request.Params.TaskId, status), + } + } + if entry.inputResponses == nil { + entry.inputResponses = make(mcp.InputResponses) + } + maps.Copy(entry.inputResponses, request.Params.InputResponses) + // Signal the task handler that responses are available. + if entry.inputReady != nil { + select { + case entry.inputReady <- struct{}{}: + default: + } + } + s.tasksMu.Unlock() + + result := &mcp.UpdateTaskResult{} + result.SetResultType(mcp.ResultTypeComplete) + return result, nil } // handleListTasks handles tasks/list requests to list all tasks. @@ -2859,7 +3159,9 @@ func (s *MCPServer) handleCancelTask( id any, request mcp.CancelTaskRequest, ) (*mcp.CancelTaskResult, *requestError) { - err := s.cancelTask(ctx, request.Params.TaskId) + // Legacy rejects cancelling a terminal task, modern acks. The + // terminal check happens under the task lock to avoid racing completion. + err := s.cancelTask(ctx, request.Params.TaskId, !IsModernRequest(ctx)) if err != nil { return nil, &requestError{ id: id, @@ -2868,7 +3170,14 @@ func (s *MCPServer) handleCancelTask( } } - // Get the updated task + // Modern protocol (SEP-2663): ack-only, no task state in response. + if IsModernRequest(ctx) { + result := mcp.CancelTaskResult{} + result.SetResultType(mcp.ResultTypeComplete) + return &result, nil + } + + // Legacy: return the cancelled task state. task, _, err := s.getTask(ctx, request.Params.TaskId) if err != nil { return nil, &requestError{ @@ -2879,6 +3188,7 @@ func (s *MCPServer) handleCancelTask( } result := mcp.NewCancelTaskResult(task) + result.Legacy = true return &result, nil } @@ -2955,11 +3265,13 @@ func (s *MCPServer) createTask(ctx context.Context, taskID string, toolName stri createdAt := time.Now() entry := &taskEntry{ - task: task, - sessionID: getSessionID(ctx), - toolName: toolName, - createdAt: createdAt, - done: make(chan struct{}), + task: task, + sessionID: getSessionID(ctx), + toolName: toolName, + createdAt: createdAt, + modernTask: IsModernRequest(ctx), + done: make(chan struct{}), + inputReady: make(chan struct{}, 1), } // Single critical section for check + increment + insert @@ -3050,14 +3362,20 @@ func (s *MCPServer) getTask(ctx context.Context, taskID string) (mcp.Task, chan return mcp.Task{}, nil, fmt.Errorf("task not found") } + // Cross-origin isolation, a request must not reach a task created under a + // different protocol surface. This keeps modern (bearer-token) tasks + // invisible to legacy callers and vice versa. + if entry.modernTask != IsModernRequest(ctx) { + s.tasksMu.RUnlock() + return mcp.Task{}, nil, fmt.Errorf("task not found") + } + // Verify session isolation sessionID := getSessionID(ctx) if entry.sessionID != "" && sessionID != "" && entry.sessionID != sessionID { s.tasksMu.RUnlock() return mcp.Task{}, nil, fmt.Errorf("task not found") } - - // Return a copy of the task and the done channel taskCopy := entry.task done := entry.done s.tasksMu.RUnlock() @@ -3080,6 +3398,12 @@ func (s *MCPServer) getTaskEntry(ctx context.Context, taskID string) (*taskEntry } s.tasksMu.RUnlock() + // Cross-origin isolation. A legacy request must not reach a + // modern-origin task or vice versa. + if entry.modernTask != IsModernRequest(ctx) { + return nil, fmt.Errorf("task not found") + } + // Verify session isolation sessionID := getSessionID(ctx) if entry.sessionID != "" && sessionID != "" && entry.sessionID != sessionID { @@ -3091,6 +3415,7 @@ func (s *MCPServer) getTaskEntry(ctx context.Context, taskID string) (*taskEntry // listTasks returns copies of all tasks for the current session. func (s *MCPServer) listTasks(ctx context.Context) []mcp.Task { + modern := IsModernRequest(ctx) sessionID := getSessionID(ctx) s.tasksMu.RLock() @@ -3098,6 +3423,11 @@ func (s *MCPServer) listTasks(ctx context.Context) []mcp.Task { var tasks []mcp.Task for _, entry := range s.tasks { + // Cross-origin isolation, never expose a task created under a different + // protocol surface. + if entry.modernTask != modern { + continue + } // Filter by session if applicable if sessionID == "" || entry.sessionID == "" || entry.sessionID == sessionID { tasks = append(tasks, entry.task) @@ -3153,7 +3483,7 @@ func (s *MCPServer) completeTask(entry *taskEntry, result any, err error) { s.tasksMu.Unlock() // Send task status notification - s.sendTaskStatusNotification(task) + s.sendTaskStatusNotification(task, entry.modernTask) // Fire task hooks, after the created hook if it is still running if s.taskHooks != nil { @@ -3167,22 +3497,29 @@ func (s *MCPServer) completeTask(entry *taskEntry, result any, err error) { } } -// cancelTask cancels a running task. -func (s *MCPServer) cancelTask(ctx context.Context, taskID string) error { +// cancelTask cancels a running task. When rejectTerminal is set (legacy path), +// cancelling an already-terminal task fails with an error instead of acking. +func (s *MCPServer) cancelTask(ctx context.Context, taskID string, rejectTerminal bool) error { entry, err := s.getTaskEntry(ctx, taskID) if err != nil { return err } - return s.cancelTaskEntry(ctx, entry, "Task cancelled by request") + return s.cancelTaskEntry(ctx, entry, "Task cancelled by request", rejectTerminal) } // cancelTaskEntry records a cancellation and invokes observers after unlocking. -func (s *MCPServer) cancelTaskEntry(ctx context.Context, entry *taskEntry, message string) error { +func (s *MCPServer) cancelTaskEntry(ctx context.Context, entry *taskEntry, message string, rejectTerminal bool) error { s.tasksMu.Lock() + + // Already terminal. legacy (rejectTerminal) returns INVALID_PARAMS, modern acks. + // Checked under the lock so a just-completed entry isn't raced. if entry.completed { status := entry.task.Status s.tasksMu.Unlock() - return fmt.Errorf("cannot cancel task in terminal status: %s", status) + if rejectTerminal { + return fmt.Errorf("cannot cancel task in terminal status: %s", status) + } + return nil } cancel := entry.cancelFunc cancelledAt := time.Now() @@ -3208,7 +3545,7 @@ func (s *MCPServer) cancelTaskEntry(ctx context.Context, entry *taskEntry, messa if cancel != nil { cancel() } - s.sendTaskStatusNotification(task) + s.sendTaskStatusNotification(task, entry.modernTask) if s.taskHooks != nil { // After the created hook, if it is still running. s.fireTaskHook(ctx, entry, func(ctx context.Context) { s.taskHooks.taskCancelled(ctx, metrics) }) @@ -3249,7 +3586,14 @@ func (s *MCPServer) scheduleTaskCleanup(taskID string, expected *taskEntry, ttlM } // sendTaskStatusNotification sends a notification when a task's status changes. -func (s *MCPServer) sendTaskStatusNotification(task mcp.Task) { +func (s *MCPServer) sendTaskStatusNotification(task mcp.Task, modernTask bool) { + // Don't broadcast modern tasks, this fans out to every session with no + // origin filtering, so the task id would leak its bearer handle. Modern + // clients poll tasks/get instead. + if modernTask { + return + } + // Convert task to map[string]any for notification params taskMap := map[string]any{ "taskId": task.TaskId, @@ -3258,14 +3602,17 @@ func (s *MCPServer) sendTaskStatusNotification(task mcp.Task) { "lastUpdatedAt": task.LastUpdatedAt, } + if task.TTL != nil { + taskMap["ttl"] = *task.TTL // legacy field name (2025-11-25) + taskMap["ttlMs"] = *task.TTL // SEP-2663 field name + } + if task.StatusMessage != "" { taskMap["statusMessage"] = task.StatusMessage } - if task.TTL != nil { - taskMap["ttl"] = *task.TTL - } if task.PollInterval != nil { - taskMap["pollInterval"] = *task.PollInterval + taskMap["pollInterval"] = *task.PollInterval // legacy field name (2025-11-25) + taskMap["pollIntervalMs"] = *task.PollInterval // SEP-2663 field name } s.SendNotificationToAllClients(mcp.MethodNotificationTasksStatus, taskMap) diff --git a/server/task_auth_contract_test.go b/server/task_auth_contract_test.go new file mode 100644 index 000000000..67f59d91d --- /dev/null +++ b/server/task_auth_contract_test.go @@ -0,0 +1,246 @@ +package server + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/mark3labs/mcp-go/mcp" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestTaskAccessControlContract verifies the bearer-token access model. +func TestTaskAccessControlContract(t *testing.T) { + newServer := func() *MCPServer { + return NewMCPServer("test", "1.0.0", WithTaskCapabilities(true, true, true)) + } + + // sessionCtx returns a context.Context with the given session ID injected, + // registering the session with s so notifications work. + sessionCtx := func(t *testing.T, s *MCPServer, id string) context.Context { + t.Helper() + sess := fakeSess{sessionID: id, notifyChan: make(chan mcp.JSONRPCNotification, 4)} + err := s.RegisterSession(t.Context(), sess) + require.NoError(t, err) + t.Cleanup(func() { s.UnregisterSession(t.Context(), id) }) + return s.WithContext(t.Context(), sess) + } + + t.Run("session match — allowed", func(t *testing.T) { + s := newServer() + ctx := sessionCtx(t, s, "session-A") + + entry, err := s.createTask(ctx, "task-match", "tool", nil, nil) + require.NoError(t, err) + + _, _, err = s.getTask(ctx, entry.task.TaskId) + assert.NoError(t, err, "same session should be allowed to access its own task") + }) + + t.Run("session mismatch — denied", func(t *testing.T) { + s := newServer() + ctxA := sessionCtx(t, s, "session-A") + ctxB := sessionCtx(t, s, "session-B") + + entry, err := s.createTask(ctxA, "task-mismatch", "tool", nil, nil) + require.NoError(t, err) + + _, _, err = s.getTask(ctxB, entry.task.TaskId) + assert.EqualError(t, err, "task not found", + "different session should be denied — returned as not-found to avoid enumeration") + }) + + t.Run("task has no session — bearer-token, any caller allowed", func(t *testing.T) { + s := newServer() + // Create task with no session in context → server-global task. + entry, err := s.createTask(t.Context(), "task-global", "tool", nil, nil) + require.NoError(t, err) + assert.Empty(t, entry.sessionID, "task created without a session should have empty sessionID") + + // A session-bearing request can still access it via task ID alone. + ctxB := sessionCtx(t, s, "session-B") + _, _, err = s.getTask(ctxB, entry.task.TaskId) + assert.NoError(t, err, + "session-less task is bearer-token accessible: any caller who knows the ID may access it") + }) + + t.Run("request has no session — bearer-token, task's session not enforced", func(t *testing.T) { + s := newServer() + ctxA := sessionCtx(t, s, "session-A") + + entry, err := s.createTask(ctxA, "task-sessioned", "tool", nil, nil) + require.NoError(t, err) + assert.Equal(t, "session-A", entry.sessionID) + + // A session-less request (stdio, in-process) can access any task by ID. + _, _, err = s.getTask(t.Context(), entry.task.TaskId) + assert.NoError(t, err, + "session-less request uses bearer-token model: task ID possession grants access") + }) + + t.Run("getTaskEntry applies the same policy as getTask", func(t *testing.T) { + s := newServer() + ctxA := sessionCtx(t, s, "session-A") + ctxB := sessionCtx(t, s, "session-B") + + entry, err := s.createTask(ctxA, "task-entry-policy", "tool", nil, nil) + require.NoError(t, err) + + _, err = s.getTaskEntry(ctxA, entry.task.TaskId) + assert.NoError(t, err, "owner session: allowed") + + _, err = s.getTaskEntry(ctxB, entry.task.TaskId) + assert.EqualError(t, err, "task not found", "different session: denied") + }) +} + +// userKey is a context key used by the auth middleware tests below. +type userKey struct{} + +// TestTaskAccessControl_HTTPMiddleware demonstrates how an HTTP middleware can +// enforce application-level authorization before requests reach the MCP server. +func TestTaskAccessControl_HTTPMiddleware(t *testing.T) { + const validToken = "secret-token" + + mcpServer := NewMCPServer("test", "1.0.0", WithTaskCapabilities(true, true, true)) + + // authMiddleware validates the Authorization header. Unauthorized requests + // are rejected before they reach HandleMessage. + authMiddleware := func(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + token := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") + if token != validToken { + http.Error(w, "unauthorized", http.StatusUnauthorized) + return + } + next.ServeHTTP(w, r) + }) + } + + srv := httptest.NewServer(authMiddleware(NewStreamableHTTPServer(mcpServer))) + t.Cleanup(srv.Close) + + doRequest := func(t *testing.T, token string) int { + t.Helper() + body, _ := json.Marshal(legacyInitRequest) + req, err := http.NewRequestWithContext(t.Context(), http.MethodPost, srv.URL, strings.NewReader(string(body))) + require.NoError(t, err) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + if token != "" { + req.Header.Set("Authorization", "Bearer "+token) + } + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + resp.Body.Close() + return resp.StatusCode + } + + t.Run("valid token — request reaches MCP server", func(t *testing.T) { + status := doRequest(t, validToken) + assert.Equal(t, http.StatusOK, status, + "valid token should pass middleware and receive an MCP response") + }) + + t.Run("missing token — rejected by middleware before reaching MCP server", func(t *testing.T) { + status := doRequest(t, "") + assert.Equal(t, http.StatusUnauthorized, status, + "missing token should be rejected before reaching the MCP server") + }) + + t.Run("wrong token — rejected by middleware before reaching MCP server", func(t *testing.T) { + status := doRequest(t, "wrong-token") + assert.Equal(t, http.StatusUnauthorized, status, + "invalid token should be rejected before reaching the MCP server") + }) +} + +// TestTaskAccessControl_OnBeforeHookExtensionPoint demonstrates how an +// OnBefore* hook reads application identity from ctx to enforce per-user +// or per-tenant task authorization. +func TestTaskAccessControl_OnBeforeHookExtensionPoint(t *testing.T) { + hooks := &Hooks{} + + var observedUser string + hooks.AddBeforeGetTask(func(ctx context.Context, id any, req *mcp.GetTaskRequest) { + // Read the identity attached by any ctx layer. + if user, ok := ctx.Value(userKey{}).(string); ok { + observedUser = user + } + }) + + s := NewMCPServer("test", "1.0.0", + WithTaskCapabilities(true, true, true), + WithHooks(hooks), + ) + + entry, err := s.createTask(t.Context(), "task-hook-ext", "tool", nil, nil) + require.NoError(t, err) + + // Simulate transport middleware having attached "alice" to the context. + ctxWithUser := context.WithValue(t.Context(), userKey{}, "alice") + s.HandleMessage(ctxWithUser, []byte( + `{"jsonrpc":"2.0","id":1,"method":"tasks/get","params":{"taskId":"`+entry.task.TaskId+`"}}`, + )) + + assert.Equal(t, "alice", observedUser, + "hook should receive the identity attached by middleware — application can enforce authz here") +} + +// TestTaskAccessControl_HTTPMiddlewareIdentityFlowsToHook demonstrates the +// complete end-to-end flow: HTTP middleware attaches identity to ctx → +// HandleMessage carries it through → OnBefore* hook reads it for per-request authz. +func TestTaskAccessControl_HTTPMiddlewareIdentityFlowsToHook(t *testing.T) { + hooks := &Hooks{} + + var hookUser string + hooks.AddBeforeGetTask(func(ctx context.Context, id any, req *mcp.GetTaskRequest) { + if user, ok := ctx.Value(userKey{}).(string); ok { + hookUser = user + } + }) + + mcpServer := NewMCPServer("test", "1.0.0", + WithTaskCapabilities(true, true, true), + WithHooks(hooks), + ) + + // Create a task without a session, bearer-token model, any caller can access it. + entry, err := mcpServer.createTask(t.Context(), "task-flow", "tool", nil, nil) + require.NoError(t, err) + + // Middleware attaches the bearer token as the caller's identity. + authMiddleware := func(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + user := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") + ctx := context.WithValue(r.Context(), userKey{}, user) + next.ServeHTTP(w, r.WithContext(ctx)) + }) + } + + srv := httptest.NewServer(authMiddleware(NewStreamableHTTPServer(mcpServer))) + t.Cleanup(srv.Close) + + // Use the legacy initialize handshake to get a valid session, then issue + // tasks/get using that session ID. + sessionID := initStatefulSession(t, srv.URL) + + body := `{"jsonrpc":"2.0","id":2,"method":"tasks/get","params":{"taskId":"` + entry.task.TaskId + `"}}` + req, err := http.NewRequestWithContext(t.Context(), http.MethodPost, srv.URL, strings.NewReader(body)) + require.NoError(t, err) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + req.Header.Set("Authorization", "Bearer bob") + req.Header.Set(HeaderKeySessionID, sessionID) + + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + resp.Body.Close() + + assert.Equal(t, "bob", hookUser, + "identity set by HTTP middleware should flow through HandleMessage to the OnBefore* hook") +} diff --git a/server/task_cancel_result_test.go b/server/task_cancel_result_test.go index 4288fe578..236f4e038 100644 --- a/server/task_cancel_result_test.go +++ b/server/task_cancel_result_test.go @@ -89,7 +89,7 @@ func TestTaskCancellationKeepsResultError(t *testing.T) { t.Fatal("result request did not wait") } } - require.NoError(t, srv.cancelTask(t.Context(), "cancelled")) + require.NoError(t, srv.cancelTask(t.Context(), "cancelled", false)) if waiting { select { case e := <-results: @@ -135,7 +135,7 @@ func TestCancelledTaskResultHTTPEnvelope(t *testing.T) { srv := NewMCPServer("test", "1", WithTaskCapabilities(true, true, true)) _, err := srv.createTask(t.Context(), "cancelled", "work", nil, nil) require.NoError(t, err) - require.NoError(t, srv.cancelTask(t.Context(), "cancelled")) + require.NoError(t, srv.cancelTask(t.Context(), "cancelled", false)) httpServer := httptest.NewServer(NewStreamableHTTPServer(srv, WithStateLess(true))) defer httpServer.Close() ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) diff --git a/server/task_hook_ordering_test.go b/server/task_hook_ordering_test.go index b76e9d3b8..4af573f23 100644 --- a/server/task_hook_ordering_test.go +++ b/server/task_hook_ordering_test.go @@ -61,7 +61,7 @@ func TestTaskHooksCreatedHookPrecedesCancellation(t *testing.T) { taskID := <-entered // The cancellation returns while the created hook still runs: its hook // waits for that one. - require.NoError(t, server.cancelTask(t.Context(), taskID)) + require.NoError(t, server.cancelTask(t.Context(), taskID, false)) close(release) select { case <-cancelled: @@ -92,7 +92,7 @@ func TestTaskHooksCreatedHookMayCancelItsOwnTask(t *testing.T) { hooks := &TaskHooks{} cancelErr := make(chan error, 1) hooks.AddOnTaskCreated(func(ctx context.Context, m TaskMetrics) { - cancelErr <- server.cancelTask(ctx, m.TaskID) + cancelErr <- server.cancelTask(ctx, m.TaskID, false) record("created") }) cancelled := make(chan struct{}) @@ -126,7 +126,7 @@ func TestTaskHooksPanickingCreatedHookDoesNotLoseCancellation(t *testing.T) { var server *MCPServer hooks := &TaskHooks{} hooks.AddOnTaskCreated(func(ctx context.Context, m TaskMetrics) { - require.NoError(t, server.cancelTask(ctx, m.TaskID)) + require.NoError(t, server.cancelTask(ctx, m.TaskID, false)) panic("hook bug") }) cancelled := make(chan struct{}) @@ -164,7 +164,7 @@ func TestTaskHooksWaitingHookOutlivesItsRequest(t *testing.T) { // The cancel request ends, as a served request does, before the created // hook lets its hook run. cancelCtx, endRequest := context.WithCancel(context.WithValue(t.Context(), keyType{}, "request value")) - require.NoError(t, server.cancelTask(cancelCtx, taskID)) + require.NoError(t, server.cancelTask(cancelCtx, taskID, false)) endRequest() close(release) select { diff --git a/server/task_hook_reentrancy_test.go b/server/task_hook_reentrancy_test.go index 7df36cfd8..a6a746b3e 100644 --- a/server/task_hook_reentrancy_test.go +++ b/server/task_hook_reentrancy_test.go @@ -47,7 +47,7 @@ func TestTaskHooksCanReadTaskState(t *testing.T) { s.executeRegularToolAsTask(t.Context(), entry, ServerTool{Handler: func(context.Context, mcp.CallToolRequest) (*mcp.CallToolResult, error) { return nil, context.Canceled }}, mcp.CallToolRequest{}) case "cancelled": hooks.AddOnTaskCancelled(observe) - err = s.cancelTask(t.Context(), "task") + err = s.cancelTask(t.Context(), "task", false) } done <- err }() @@ -79,7 +79,7 @@ func TestTaskHookMetricsRemainEventSnapshotsDuringConcurrentCompletion(t *testin var wg sync.WaitGroup wg.Add(2) go func() { defer wg.Done(); s.completeTask(entry, nil, nil) }() - go func() { defer wg.Done(); _ = s.cancelTask(t.Context(), id) }() + go func() { defer wg.Done(); _ = s.cancelTask(t.Context(), id, false) }() wg.Wait() require.Len(t, events, 2) require.Equal(t, mcp.TaskStatusWorking, events[0].Status) diff --git a/server/task_hooks_test.go b/server/task_hooks_test.go index 53b5121f8..c00d85dd2 100644 --- a/server/task_hooks_test.go +++ b/server/task_hooks_test.go @@ -162,7 +162,7 @@ func TestTaskHooks_TaskCancelled(t *testing.T) { server.tasksMu.Unlock() // Cancel the task - err = server.cancelTask(cancelCtx, "test-task-4") + err = server.cancelTask(cancelCtx, "test-task-4", false) require.NoError(t, err) // Wait for hook to execute with timeout diff --git a/server/task_http_lifecycle_test.go b/server/task_http_lifecycle_test.go new file mode 100644 index 000000000..ce06de52a --- /dev/null +++ b/server/task_http_lifecycle_test.go @@ -0,0 +1,261 @@ +package server + +import ( + "context" + "encoding/json" + "io" + "net/http" + "strings" + "testing" + "time" + + "github.com/mark3labs/mcp-go/mcp" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// parseSSEOrJSON extracts the JSON payload from either a plain JSON response +// or an SSE stream with a single "data: {...}" line. +func parseSSEOrJSON(body []byte) []byte { + for line := range strings.SplitSeq(string(body), "\n") { + if after, ok := strings.CutPrefix(line, "data: "); ok { + after = strings.TrimSpace(after) + if len(after) > 0 { + return []byte(after) + } + } + } + return body +} + +// Streamable HTTP tests can obtain a session ID via the initialize handshake. +var legacyInitRequest = map[string]any{ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": map[string]any{ + "protocolVersion": mcp.LATEST_LEGACY_PROTOCOL_VERSION, + "clientInfo": map[string]any{"name": "test-client", "version": "1.0.0"}, + }, +} + +// initStatefulSession performs the initialize handshake and returns the session ID. +func initStatefulSession(t *testing.T, serverURL string) string { + t.Helper() + resp, err := postJSON(serverURL, legacyInitRequest) + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode) + resp.Body.Close() + sessionID := resp.Header.Get(HeaderKeySessionID) + require.NotEmpty(t, sessionID, "expected session ID from initialize") + return sessionID +} + +// TestTaskHTTPLifecycle_RequestContextDoesNotCancelTask verifies that +// task goroutines are NOT cancelled when the HTTP response is written. +func TestTaskHTTPLifecycle_RequestContextDoesNotCancelTask(t *testing.T) { + srv := NewMCPServer("test", "1.0.0", + WithTaskCapabilities(true, true, true), + ) + + handlerCompleted := make(chan struct{}) + handlerStarted := make(chan struct{}) + + tool := mcp.NewTool("blocking_task", + mcp.WithDescription("A task that completes after a short delay"), + mcp.WithTaskSupport(mcp.TaskSupportRequired), + ) + srv.AddTaskTool(tool, func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CreateTaskResult, error) { + close(handlerStarted) + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-time.After(200 * time.Millisecond): + close(handlerCompleted) + r := mcp.NewCreateTaskResult(mcp.Task{}) + return &r, nil + } + }) + + ts := NewTestStreamableHTTPServer(srv, WithStateful(true)) + defer ts.Close() + + sessionID := initStatefulSession(t, ts.URL) + + callReq := map[string]any{ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/call", + "params": map[string]any{"name": "blocking_task", "task": map[string]any{}}, + } + callResp, err := postSessionJSON(ts.URL, sessionID, callReq) + require.NoError(t, err) + defer callResp.Body.Close() + require.Equal(t, http.StatusOK, callResp.StatusCode) + + var result struct { + Result struct { + Task struct { + TaskId string `json:"taskId"` + Status string `json:"status"` + } `json:"task"` + } `json:"result"` + } + require.NoError(t, json.NewDecoder(callResp.Body).Decode(&result)) + taskID := result.Result.Task.TaskId + require.NotEmpty(t, taskID, "expected a task ID in the response") + assert.Equal(t, "working", result.Result.Task.Status) + + // Wait until the goroutine has started before asserting it completes. + // Without this barrier the 200ms timer may not have begun yet. + <-handlerStarted + + // HTTP response has been sent; handler must still complete. + select { + case <-handlerCompleted: + // Correct: task outlived the HTTP request. + case <-time.After(2 * time.Second): + t.Fatal("task handler was cancelled after HTTP response — " + + "executeTaskTool should use s.serverCtx, not the per-request context") + } +} + +// TestTaskHTTPLifecycle_CancelViaTasksCancel verifies that tasks/cancel works +// correctly over Streamable HTTP regardless of the request-context limitation. +// tasks/cancel operates which is set by the task goroutine and is independent +// of the HTTP request context. +func TestTaskHTTPLifecycle_CancelViaTasksCancel(t *testing.T) { + srv := NewMCPServer("test", "1.0.0", + WithTaskCapabilities(true, true, true), + ) + + handlerRunning := make(chan struct{}) + + tool := mcp.NewTool("cancellable_task", + mcp.WithDescription("A task that can be cancelled via tasks/cancel"), + mcp.WithTaskSupport(mcp.TaskSupportRequired), + ) + srv.AddTaskTool(tool, func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CreateTaskResult, error) { + close(handlerRunning) + <-ctx.Done() + return nil, ctx.Err() + }) + + ts := NewTestStreamableHTTPServer(srv, WithStateful(true)) + defer ts.Close() + + sessionID := initStatefulSession(t, ts.URL) + + callReq := map[string]any{ + "jsonrpc": "2.0", "id": 2, "method": "tools/call", + "params": map[string]any{"name": "cancellable_task", "task": map[string]any{}}, + } + callResp, err := postSessionJSON(ts.URL, sessionID, callReq) + require.NoError(t, err) + var callResult struct { + Result struct { + Task struct { + TaskId string `json:"taskId"` + } `json:"task"` + } `json:"result"` + } + require.NoError(t, json.NewDecoder(callResp.Body).Decode(&callResult)) + callResp.Body.Close() + taskID := callResult.Result.Task.TaskId + require.NotEmpty(t, taskID) + + select { + case <-handlerRunning: + case <-time.After(2 * time.Second): + t.Fatal("handler did not start") + } + + // Cancel via tasks/cancel. + cancelResp, err := postSessionJSON(ts.URL, sessionID, map[string]any{ + "jsonrpc": "2.0", "id": 3, "method": "tasks/cancel", + "params": map[string]any{"taskId": taskID}, + }) + require.NoError(t, err) + require.Equal(t, http.StatusOK, cancelResp.StatusCode) + cancelResp.Body.Close() + + // Poll tasks/get until the task reaches a terminal state. + // Verify the task reaches a terminal state. + var lastStatus string + require.Eventually(t, func() bool { + getResp, err := postSessionJSON(ts.URL, sessionID, map[string]any{ + "jsonrpc": "2.0", "id": 4, "method": "tasks/get", + "params": map[string]any{"taskId": taskID}, + }) + if err != nil { + return false + } + body, _ := io.ReadAll(getResp.Body) + getResp.Body.Close() + var gr struct { + Result struct { + Status string `json:"status"` + } `json:"result"` + } + if err := json.Unmarshal(parseSSEOrJSON(body), &gr); err != nil { + return false + } + lastStatus = gr.Result.Status + s := gr.Result.Status + return s == "cancelled" || s == "failed" || s == "completed" + }, 3*time.Second, 100*time.Millisecond, "task status: %s", lastStatus) +} + +// TestTaskHTTPLifecycle_ContextNotCancelledOnAllTransports verifies that task +// goroutines are not cancelled when the request context is cancelled. +// Cancelling the caller's context after HandleMessage returns must not affect +// the running goroutine. +func TestTaskHTTPLifecycle_ContextNotCancelledOnAllTransports(t *testing.T) { + srv := NewMCPServer("test", "1.0.0", + WithTaskCapabilities(true, true, true), + ) + + handlerCompleted := make(chan struct{}) + + tool := mcp.NewTool("blocking_task_inprocess", + mcp.WithDescription("Completes after a short delay; must not be cancelled by request ctx"), + mcp.WithTaskSupport(mcp.TaskSupportRequired), + ) + srv.AddTaskTool(tool, func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CreateTaskResult, error) { + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-time.After(200 * time.Millisecond): + close(handlerCompleted) + r := mcp.NewCreateTaskResult(mcp.Task{}) + return &r, nil + } + }) + + // Use a cancellable context. We will cancel it right after HandleMessage + // returns to simulate the per-request context being torn down. + reqCtx, cancelReq := context.WithCancel(t.Context()) + notifyChan := make(chan mcp.JSONRPCNotification, 10) + session := fakeSess{sessionID: "inprocess-session", notifyChan: notifyChan} + require.NoError(t, srv.RegisterSession(reqCtx, session)) + sessionCtx := srv.WithContext(reqCtx, session) + + msg, _ := json.Marshal(map[string]any{ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": map[string]any{"name": "blocking_task_inprocess", "task": map[string]any{}}, + }) + srv.HandleMessage(sessionCtx, msg) + + // Cancel the request context immediately. + cancelReq() + + select { + case <-handlerCompleted: + // Correct: goroutine is rooted on s.serverCtx and survives request cancellation. + case <-time.After(2 * time.Second): + t.Fatal("task handler was cancelled after request ctx was cancelled — " + + "executeTaskTool must use s.serverCtx, not the per-request context") + } +} diff --git a/server/task_input_required_test.go b/server/task_input_required_test.go new file mode 100644 index 000000000..76139009d --- /dev/null +++ b/server/task_input_required_test.go @@ -0,0 +1,604 @@ +package server + +import ( + "context" + "encoding/json" + "fmt" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/mark3labs/mcp-go/mcp" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestTaskInputRequired_RegularToolRoundTrip verifies an optional tool run as a +// task (modern, extension declared) that returns input_required. +func TestTaskInputRequired_RegularToolRoundTrip(t *testing.T) { + srv := NewMCPServer("test", "1.0.0", WithTaskCapabilities(true, true, true)) + + const confirmKey = "confirm" + srv.AddTool( + mcp.NewTool("deploy", mcp.WithTaskSupport(mcp.TaskSupportOptional)), + func(_ context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + // Resume, the answer is present on the retry. + if _, ok := req.Params.InputResponses[confirmKey]; ok { + return mcp.NewToolResultText("deployed"), nil + } + // First call, ask for input. + return NewInputRequestBuilder("step=1"). + Elicit(confirmKey, mcp.ElicitationParams{ + Mode: mcp.ElicitationModeForm, + Message: "Proceed?", + RequestedSchema: map[string]any{ + "type": "object", + "properties": map[string]any{"approve": map[string]any{"type": "boolean"}}, + "required": []string{"approve"}, + }, + }).ToolResult(), nil + }) + + ctx := buildModernContext(t.Context(), srv) + + created, reqErr := srv.handleToolCall(ctx, 1, mcp.CallToolRequest{ + Params: mcp.CallToolParams{Name: "deploy"}, + }) + require.Nil(t, reqErr) + task, ok := created.(*mcp.CreateTaskResult) + require.True(t, ok, "expected *CreateTaskResult, got %T", created) + taskID := task.Task.TaskId + + // The async executor should park the task in input_required with the + // pending request surfaced via tasks/get. + require.Eventually(t, func() bool { + got, e := srv.handleGetTask(ctx, 2, mcp.GetTaskRequest{Params: mcp.GetTaskParams{TaskId: taskID}}) + return e == nil && got.Status == mcp.TaskStatusInputRequired + }, 5*time.Second, 10*time.Millisecond) + + got, e := srv.handleGetTask(ctx, 3, mcp.GetTaskRequest{Params: mcp.GetTaskParams{TaskId: taskID}}) + require.Nil(t, e) + require.Contains(t, got.InputRequests, confirmKey, "inputRequests must surface the pending request key") + + // Answer via tasks/update. + _, e = srv.handleUpdateTask(ctx, 4, mcp.UpdateTaskRequest{ + Params: mcp.UpdateTaskParams{ + TaskId: taskID, + InputResponses: mcp.InputResponses{ + confirmKey: mcp.NewElicitationInputResponse(mcp.ElicitationResult{ + ElicitationResponse: mcp.ElicitationResponse{ + Action: mcp.ElicitationResponseActionAccept, + Content: map[string]any{"approve": true}, + }, + }), + }, + }, + }) + require.Nil(t, e) + + // The task should resume and complete with the real result. + require.Eventually(t, func() bool { + got, e := srv.handleGetTask(ctx, 5, mcp.GetTaskRequest{Params: mcp.GetTaskParams{TaskId: taskID}}) + return e == nil && got.Status == mcp.TaskStatusCompleted + }, 5*time.Second, 10*time.Millisecond) + + result, e := srv.handleTaskResult(ctx, 6, mcp.TaskResultRequest{Params: mcp.TaskResultParams{TaskId: taskID}}) + require.Nil(t, e) + require.False(t, result.IsError) + require.Len(t, result.Content, 1) + text, ok := result.Content[0].(mcp.TextContent) + require.True(t, ok) + require.Equal(t, "deployed", text.Text) +} + +// TestTaskInputRequired_CancelWhileParked verifies that cancelling a task that +// is parked in input_required transitions it to cancelled and lets the executor +// goroutine exit cleanly. +func TestTaskInputRequired_CancelWhileParked(t *testing.T) { + srv := NewMCPServer("test", "1.0.0", WithTaskCapabilities(true, true, true)) + + var resumed atomic.Bool + srv.AddTool( + mcp.NewTool("waits", mcp.WithTaskSupport(mcp.TaskSupportOptional)), + func(_ context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + if len(req.Params.InputResponses) > 0 { + resumed.Store(true) + return mcp.NewToolResultText("done"), nil + } + return NewInputRequestBuilder(""). + Elicit("x", mcp.ElicitationParams{Mode: mcp.ElicitationModeForm, Message: "?"}). + ToolResult(), nil + }) + + ctx := buildModernContext(t.Context(), srv) + + created, reqErr := srv.handleToolCall(ctx, 1, mcp.CallToolRequest{ + Params: mcp.CallToolParams{Name: "waits"}, + }) + require.Nil(t, reqErr) + taskID := created.(*mcp.CreateTaskResult).Task.TaskId + + require.Eventually(t, func() bool { + got, e := srv.handleGetTask(ctx, 2, mcp.GetTaskRequest{Params: mcp.GetTaskParams{TaskId: taskID}}) + return e == nil && got.Status == mcp.TaskStatusInputRequired + }, 5*time.Second, 10*time.Millisecond) + + // Cancel while parked. + _, e := srv.handleCancelTask(ctx, 3, mcp.CancelTaskRequest{Params: mcp.CancelTaskParams{TaskId: taskID}}) + require.Nil(t, e) + + require.Eventually(t, func() bool { + got, e := srv.handleGetTask(ctx, 4, mcp.GetTaskRequest{Params: mcp.GetTaskParams{TaskId: taskID}}) + return e == nil && got.Status == mcp.TaskStatusCancelled + }, 5*time.Second, 10*time.Millisecond) + + // Give the executor a moment, it must not have resumed the handler. + time.Sleep(50 * time.Millisecond) + assert.False(t, resumed.Load(), "handler must not resume after cancellation") +} + +// TestTaskInputRequired_SlotFreedOnExpiry verifies that cancelling a task +// parked in input_required frees its concurrency slot. +func TestTaskInputRequired_SlotFreedOnExpiry(t *testing.T) { + srv := NewMCPServer("test", "1.0.0", + WithTaskCapabilities(true, true, true), + WithMaxConcurrentTasks(1), + ) + + srv.AddTool( + mcp.NewTool("ask", mcp.WithTaskSupport(mcp.TaskSupportOptional)), + func(_ context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + if len(req.Params.InputResponses) > 0 { + return mcp.NewToolResultText("done"), nil + } + return NewInputRequestBuilder(""). + Elicit("x", mcp.ElicitationParams{Mode: mcp.ElicitationModeForm, Message: "?"}). + ToolResult(), nil + }) + + ctx := buildModernContext(t.Context(), srv) + + created, reqErr := srv.handleToolCall(ctx, 1, mcp.CallToolRequest{ + Params: mcp.CallToolParams{Name: "ask"}, + }) + require.Nil(t, reqErr) + taskID := created.(*mcp.CreateTaskResult).Task.TaskId + + // Wait until parked in input_required (slot is held). + require.Eventually(t, func() bool { + got, e := srv.handleGetTask(ctx, 2, mcp.GetTaskRequest{Params: mcp.GetTaskParams{TaskId: taskID}}) + return e == nil && got.Status == mcp.TaskStatusInputRequired + }, 5*time.Second, 10*time.Millisecond) + + // Cancel the parked task, the executor must free the slot. + _, e := srv.handleCancelTask(ctx, 3, mcp.CancelTaskRequest{Params: mcp.CancelTaskParams{TaskId: taskID}}) + require.Nil(t, e) + + // A new task must now be creatable despite the 1-task limit. + require.Eventually(t, func() bool { + c2, re := srv.handleToolCall(ctx, 4, mcp.CallToolRequest{Params: mcp.CallToolParams{Name: "ask"}}) + if re != nil { + return false + } + _, ok := c2.(*mcp.CreateTaskResult) + return ok + }, 5*time.Second, 10*time.Millisecond, "concurrency slot must be freed after the parked task is cancelled") +} + +// TestTaskInputRequired_LegacyFailsCleanly verifies a legacy client on the task +// path whose handler returns input_required fails with a clear error rather +// than hanging (tasks/update is modern-only, so it could never answer). +func TestTaskInputRequired_LegacyFailsCleanly(t *testing.T) { + srv := NewMCPServer("test", "1.0.0", WithTaskCapabilities(true, true, true)) + + srv.AddTool( + mcp.NewTool("ask", mcp.WithTaskSupport(mcp.TaskSupportOptional)), + func(context.Context, mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return NewInputRequestBuilder(""). + Elicit("x", mcp.ElicitationParams{Mode: mcp.ElicitationModeForm, Message: "?"}). + ToolResult(), nil + }) + + // Legacy context (no modern protocol info) with an explicit task param. + ctx := t.Context() + created, reqErr := srv.handleToolCall(ctx, 1, mcp.CallToolRequest{ + Params: mcp.CallToolParams{Name: "ask", Task: &mcp.TaskParams{}}, + }) + require.Nil(t, reqErr) + taskID := created.(*mcp.CreateTaskResult).Task.TaskId + + // The task must reach failed (not input_required). + require.Eventually(t, func() bool { + task, _, err := srv.getTask(ctx, taskID) + return err == nil && task.Status == mcp.TaskStatusFailed + }, 5*time.Second, 10*time.Millisecond) + + task, _, err := srv.getTask(ctx, taskID) + require.NoError(t, err) + assert.Equal(t, mcp.TaskStatusFailed, task.Status) + assert.Contains(t, task.StatusMessage, "legacy task execution") +} + +// TestTaskInputRequired_ValidationBlocksTaskCreation verifies that with +// WithInputSchemaValidation(), a modern call missing a required argument is +// rejected synchronously and no task is created. +func TestTaskInputRequired_ValidationBlocksTaskCreation(t *testing.T) { + srv := NewMCPServer("test", "1.0.0", + WithTaskCapabilities(true, true, true), + WithInputSchemaValidation(), + ) + + var handlerRan bool + srv.AddTool( + mcp.NewTool("needs_arg", + mcp.WithTaskSupport(mcp.TaskSupportOptional), + mcp.WithString("text", mcp.Required()), + ), + func(context.Context, mcp.CallToolRequest) (*mcp.CallToolResult, error) { + handlerRan = true + return mcp.NewToolResultText("ok"), nil + }) + + ctx := buildModernContext(t.Context(), srv) + + // Missing the required "text" argument. + result, reqErr := srv.handleToolCall(ctx, 1, mcp.CallToolRequest{ + Params: mcp.CallToolParams{Name: "needs_arg"}, + }) + + // Validation returns a tool execution error result (not a task handle). + require.Nil(t, reqErr) + callResult, ok := result.(*mcp.CallToolResult) + require.True(t, ok, "expected *CallToolResult (validation error), got %T", result) + assert.True(t, callResult.IsError) + + // No task should have been created, and the handler must not have run. + assert.False(t, handlerRan) + srv.tasksMu.RLock() + n := len(srv.tasks) + srv.tasksMu.RUnlock() + assert.Equal(t, 0, n, "no task should be created when input validation fails") +} + +// TestTaskGetModern_FailedPreservesErrorDetails verifies a handler that fails +// with an exposed JSON-RPC error (NewHandlerError) keeps its code, message, and +// data through modern tasks/get instead of flattening to INTERNAL_ERROR. +// testHandlerJSONRPCError (request_error_test.go) supplies INVALID_PARAMS, +// "invalid customer ID", data {field: customer_id}. +func TestTaskGetModern_FailedPreservesErrorDetails(t *testing.T) { + srv := NewMCPServer("test", "1.0.0", WithTaskCapabilities(true, true, true)) + + srv.AddTool( + mcp.NewTool("fails", mcp.WithTaskSupport(mcp.TaskSupportOptional)), + func(context.Context, mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return nil, NewHandlerError(testHandlerJSONRPCError{}) + }) + + ctx := buildModernContext(t.Context(), srv) + + created, reqErr := srv.handleToolCall(ctx, 1, mcp.CallToolRequest{ + Params: mcp.CallToolParams{Name: "fails"}, + }) + require.Nil(t, reqErr) + taskID := created.(*mcp.CreateTaskResult).Task.TaskId + + require.Eventually(t, func() bool { + got, e := srv.handleGetTask(ctx, 2, mcp.GetTaskRequest{Params: mcp.GetTaskParams{TaskId: taskID}}) + return e == nil && got.Status == mcp.TaskStatusFailed + }, 5*time.Second, 10*time.Millisecond) + + got, e := srv.handleGetTask(ctx, 3, mcp.GetTaskRequest{Params: mcp.GetTaskParams{TaskId: taskID}}) + require.Nil(t, e) + require.NotNil(t, got.TaskError, "failed task must carry error details") + assert.Equal(t, mcp.INVALID_PARAMS, got.TaskError.Code, "original JSON-RPC code must be preserved") + assert.Equal(t, "invalid customer ID", got.TaskError.Message, "client-facing message must be preserved") + assert.Equal(t, map[string]any{"field": "customer_id"}, got.TaskError.Data, "error data must be preserved") +} + +// TestTaskModern_IgnoresLegacyTaskParam verifies a modern request's legacy +// params.task.ttl is ignored; the TTL comes from the server default +// (WithDefaultTaskTTL), not the removed-in-modern param. +func TestTaskModern_IgnoresLegacyTaskParam(t *testing.T) { + const serverDefault = int64(60000) + srv := NewMCPServer("test", "1.0.0", + WithTaskCapabilities(true, true, true), + WithDefaultTaskTTL(serverDefault), + ) + srv.AddTool( + mcp.NewTool("work", mcp.WithTaskSupport(mcp.TaskSupportOptional)), + func(context.Context, mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return mcp.NewToolResultText("ok"), nil + }) + + ctx := buildModernContext(t.Context(), srv) + stray := int64(1000) + created, reqErr := srv.handleToolCall(ctx, 1, mcp.CallToolRequest{ + Params: mcp.CallToolParams{Name: "work", Task: &mcp.TaskParams{TTL: &stray}}, + }) + require.Nil(t, reqErr) + taskID := created.(*mcp.CreateTaskResult).Task.TaskId + + task, _, err := srv.getTask(ctx, taskID) + require.NoError(t, err) + require.NotNil(t, task.TTL) + assert.Equal(t, serverDefault, *task.TTL, + "modern request must ignore legacy task.ttl and use the server default") +} + +// TestDiscoverModern_OmitsLegacyTaskCapability verifies that with both +// WithTaskCapabilities and WithTasksExtension, modern discover omits the legacy +// Tasks block and advertises the extension, while legacy initialize keeps it. +func TestDiscoverModern_OmitsLegacyTaskCapability(t *testing.T) { + srv := NewMCPServer("test", "1.0.0", + WithTaskCapabilities(true, true, true), + WithTasksExtension(), + ) + + // Modern discover: legacy Tasks block must be absent, extension present. + modern := srv.serverCapabilitiesSnapshot(true) + assert.Nil(t, modern.Tasks, "modern discovery must not advertise the legacy tasks capability") + require.NotNil(t, modern.Extensions) + _, hasExt := modern.Extensions[mcp.ExtensionTasks] + assert.True(t, hasExt, "modern discovery must advertise the tasks extension") + + // Legacy initialize: legacy Tasks block must still be advertised. + legacy := srv.serverCapabilitiesSnapshot(false) + assert.NotNil(t, legacy.Tasks, "legacy initialize must still advertise the legacy tasks capability") +} + +// TestTaskGetModern_URLElicitationMapsToInternal verifies a handler returning a +// URL-elicitation error surfaces INTERNAL_ERROR via modern tasks/get, not the +// legacy URL_ELICITATION_REQUIRED code (modern clients drive elicitation via +// multi round-trips). +func TestTaskGetModern_URLElicitationMapsToInternal(t *testing.T) { + srv := NewMCPServer("test", "1.0.0", WithTaskCapabilities(true, true, true)) + + srv.AddTool( + mcp.NewTool("needsURL", mcp.WithTaskSupport(mcp.TaskSupportOptional)), + func(context.Context, mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return nil, NewHandlerError(mcp.URLElicitationRequiredError{ + Elicitations: []mcp.ElicitationParams{{Message: "authorize at https://example.com/auth"}}, + }) + }) + + ctx := buildModernContext(t.Context(), srv) + created, reqErr := srv.handleToolCall(ctx, 1, mcp.CallToolRequest{ + Params: mcp.CallToolParams{Name: "needsURL"}, + }) + require.Nil(t, reqErr) + taskID := created.(*mcp.CreateTaskResult).Task.TaskId + + require.Eventually(t, func() bool { + got, e := srv.handleGetTask(ctx, 2, mcp.GetTaskRequest{Params: mcp.GetTaskParams{TaskId: taskID}}) + return e == nil && got.Status == mcp.TaskStatusFailed + }, 5*time.Second, 10*time.Millisecond) + + got, e := srv.handleGetTask(ctx, 3, mcp.GetTaskRequest{Params: mcp.GetTaskParams{TaskId: taskID}}) + require.Nil(t, e) + require.NotNil(t, got.TaskError) + assert.Equal(t, mcp.INTERNAL_ERROR, got.TaskError.Code, + "modern tasks/get must not leak the legacy URL-elicitation code") + assert.NotEqual(t, mcp.URL_ELICITATION_REQUIRED, got.TaskError.Code) +} + +// TestDiscoverModern_OmitsLegacyTaskCapabilityWithoutExtension verifies the +// legacy tasks.* block is suppressed from modern discovery even with +// WithTaskCapabilities ONLY (no extension), it's not modern vocabulary while +// legacy initialize still carries it. +func TestDiscoverModern_OmitsLegacyTaskCapabilityWithoutExtension(t *testing.T) { + srv := NewMCPServer("test", "1.0.0", + WithTaskCapabilities(true, true, true), + // No WithTasksExtension(). + ) + + modern := srv.serverCapabilitiesSnapshot(true) + assert.Nil(t, modern.Tasks, + "modern discovery must not advertise legacy tasks.* even without the extension") + + legacy := srv.serverCapabilitiesSnapshot(false) + assert.NotNil(t, legacy.Tasks, "legacy initialize must still advertise the legacy tasks capability") +} + +// TestTaskInputRequired_EmptyRequestsFailsFast verifies that an input_required +// result with no input requests (the "retry later" signal) fails the task +// rather than parking it forever for a tasks/update that can't answer. +func TestTaskInputRequired_EmptyRequestsFailsFast(t *testing.T) { + srv := NewMCPServer("test", "1.0.0", WithTaskCapabilities(true, true, true)) + + srv.AddTool( + mcp.NewTool("shed", mcp.WithTaskSupport(mcp.TaskSupportOptional)), + func(context.Context, mcp.CallToolRequest) (*mcp.CallToolResult, error) { + // input_required with no InputRequests. + return &mcp.CallToolResult{ + Result: mcp.Result{ResultType: mcp.ResultTypeInputRequired}, + }, nil + }) + + ctx := buildModernContext(t.Context(), srv) + created, reqErr := srv.handleToolCall(ctx, 1, mcp.CallToolRequest{ + Params: mcp.CallToolParams{Name: "shed"}, + }) + require.Nil(t, reqErr) + taskID := created.(*mcp.CreateTaskResult).Task.TaskId + + // Must reach failed, not park in input_required. + require.Eventually(t, func() bool { + task, _, err := srv.getTask(ctx, taskID) + return err == nil && task.Status == mcp.TaskStatusFailed + }, 5*time.Second, 10*time.Millisecond) + + task, _, err := srv.getTask(ctx, taskID) + require.NoError(t, err) + assert.Equal(t, mcp.TaskStatusFailed, task.Status) + assert.Contains(t, task.StatusMessage, "no input requests") +} + +// TestTaskCrossOriginIsolation verifies a task created under one protocol +// surface is invisible to the other: a legacy caller cannot enumerate +// (tasks/list) or access (tasks/get) a modern-origin task, and vice versa. +func TestTaskCrossOriginIsolation(t *testing.T) { + srv := NewMCPServer("test", "1.0.0", WithTaskCapabilities(true, true, true)) + + legacyCtx := t.Context() + modernCtx := buildModernContext(t.Context(), srv) + + // A modern-origin task and a legacy-origin task. + _, err := srv.createTask(modernCtx, "modern-task", "tool", nil, nil) + require.NoError(t, err) + _, err = srv.createTask(legacyCtx, "legacy-task", "tool", nil, nil) + require.NoError(t, err) + + // Legacy caller: getTask must not reach the modern task, but reaches its own. + _, _, err = srv.getTask(legacyCtx, "modern-task") + assert.Error(t, err, "legacy caller must not access a modern-origin task") + _, _, err = srv.getTask(legacyCtx, "legacy-task") + assert.NoError(t, err, "legacy caller must access its own task") + + // Modern caller: the reverse. + _, _, err = srv.getTask(modernCtx, "legacy-task") + assert.Error(t, err, "modern caller must not access a legacy-origin task") + _, _, err = srv.getTask(modernCtx, "modern-task") + assert.NoError(t, err, "modern caller must access its own task") + + // Legacy tasks/list must enumerate only legacy-origin tasks. + listed := srv.listTasks(legacyCtx) + var ids []string + for _, tk := range listed { + ids = append(ids, tk.TaskId) + } + assert.Contains(t, ids, "legacy-task") + assert.NotContains(t, ids, "modern-task", "legacy tasks/list must not enumerate modern tasks") +} + +// TestTaskMissingExtensionErrorData verifies the +// MISSING_REQUIRED_CLIENT_CAPABILITY error for a modern request lacking the +// tasks extension carries the normative requiredCapabilities.extensions data +// identifying io.modelcontextprotocol/tasks. +func TestTaskMissingExtensionErrorData(t *testing.T) { + srv := NewMCPServer("test", "1.0.0", + WithTaskCapabilities(true, true, true), + WithTasksExtension(), + ) + + // Modern request (declares protocol version + capabilities) but WITHOUT + // the tasks extension, so the dispatcher's capability guard fires. + meta, _ := json.Marshal(map[string]any{ + "io.modelcontextprotocol/protocolVersion": "2026-07-28", + "io.modelcontextprotocol/clientCapabilities": map[string]any{}, + }) + msg := `{"jsonrpc":"2.0","id":1,"method":"tasks/get","params":{"taskId":"whatever","_meta":` + string(meta) + `}}` + + response := srv.HandleMessage(t.Context(), []byte(msg)) + errResp, ok := response.(mcp.JSONRPCError) + require.True(t, ok, "expected JSONRPCError, got %T", response) + assert.Equal(t, mcp.MISSING_REQUIRED_CLIENT_CAPABILITY, errResp.Error.Code) + + // The error data must carry requiredCapabilities.extensions[tasks]. + data, ok := errResp.Error.Data.(map[string]any) + require.True(t, ok, "error data must be a structured object, got %T", errResp.Error.Data) + reqCaps, ok := data["requiredCapabilities"].(map[string]any) + require.True(t, ok, "data must contain requiredCapabilities") + exts, ok := reqCaps["extensions"].(map[string]any) + require.True(t, ok, "requiredCapabilities must contain extensions") + _, ok = exts[mcp.ExtensionTasks] + assert.True(t, ok, "extensions must identify %s", mcp.ExtensionTasks) +} + +// TestModernTaskStatusNotSentToOtherSessions verifies a modern task does not +// broadcast its id via notifications/tasks/status (which fans out to every +// session, leaking the bearer handle). Legacy tasks still broadcast. +func TestModernTaskStatusNotSentToOtherSessions(t *testing.T) { + srv := NewMCPServer("test", "1.0.0", WithTaskCapabilities(true, true, true)) + + notifyChan := make(chan mcp.JSONRPCNotification, 10) + require.NoError(t, srv.RegisterSession(t.Context(), + fakeSess{sessionID: "observer", notifyChan: notifyChan})) + + modernCtx := buildModernContext(t.Context(), srv) + ttl := int64(60000) + entry, err := srv.createTask(modernCtx, "modern-task", "tool", &ttl, nil) + require.NoError(t, err) + + srv.completeTask(entry, &mcp.CallToolResult{ + Content: []mcp.Content{mcp.NewTextContent("secret")}, + }, nil) + + select { + case n := <-notifyChan: + t.Fatalf("modern task leaked status broadcast: %s %v", n.Method, n.Params.AdditionalFields) + case <-time.After(100 * time.Millisecond): + // No broadcast — correct. + } + + // A legacy task still broadcasts. + legacyEntry, err := srv.createTask(t.Context(), "legacy-task", "tool", &ttl, nil) + require.NoError(t, err) + srv.completeTask(legacyEntry, &mcp.CallToolResult{ + Content: []mcp.Content{mcp.NewTextContent("ok")}, + }, nil) + + select { + case n := <-notifyChan: + assert.Equal(t, mcp.MethodNotificationTasksStatus, n.Method) + assert.Equal(t, "legacy-task", n.Params.AdditionalFields["taskId"]) + case <-time.After(time.Second): + t.Fatal("expected legacy task status broadcast") + } +} + +// TestLegacyTerminalCancelRejectsAfterCompletion verifies a legacy cancel of an +// already-terminal task errors (INVALID_PARAMS upstream) even when completion +// raced it, while the modern path acks. Decided under the task lock. +func TestLegacyTerminalCancelRejectsAfterCompletion(t *testing.T) { + srv := NewMCPServer("test", "1.0.0", WithTaskCapabilities(true, true, true)) + + ttl := int64(60000) + entry, err := srv.createTask(t.Context(), "legacy-terminal", "tool", &ttl, nil) + require.NoError(t, err) + srv.completeTask(entry, &mcp.CallToolResult{ + Content: []mcp.Content{mcp.NewTextContent("done")}, + }, nil) + + // Legacy cancel (rejectTerminal=true) of a terminal task must error. + err = srv.cancelTask(t.Context(), "legacy-terminal", true) + require.Error(t, err) + assert.Contains(t, err.Error(), "terminal status") + + // Modern cancel (rejectTerminal=false) of the same terminal task acks. + err = srv.cancelTask(t.Context(), "legacy-terminal", false) + assert.NoError(t, err) +} + +// TestConcurrentUpdateAndComplete checks under the race detector that +// tasks/update reading the task status does not race a concurrent completion. +func TestConcurrentUpdateAndComplete(t *testing.T) { + srv := NewMCPServer("test", "1.0.0", WithTaskCapabilities(true, true, true)) + modernCtx := buildModernContext(t.Context(), srv) + + for i := range 50 { + entry, err := srv.createTask(modernCtx, fmt.Sprintf("race-%d", i), "tool", nil, nil) + require.NoError(t, err) + taskID := entry.task.TaskId + + var wg sync.WaitGroup + wg.Add(2) + go func() { + defer wg.Done() + // Rejection branch formats the status; must be race-free. + _, _ = srv.handleUpdateTask(modernCtx, 1, mcp.UpdateTaskRequest{ + Params: mcp.UpdateTaskParams{ + TaskId: taskID, + InputResponses: mcp.InputResponses{}, + }, + }) + }() + go func() { + defer wg.Done() + srv.completeTask(entry, &mcp.CallToolResult{ + Content: []mcp.Content{mcp.NewTextContent("x")}, + }, nil) + }() + wg.Wait() + } +} diff --git a/server/task_limit_test.go b/server/task_limit_test.go index 366fe471d..97e70602d 100644 --- a/server/task_limit_test.go +++ b/server/task_limit_test.go @@ -128,7 +128,7 @@ func TestMaxConcurrentTasks(t *testing.T) { server.tasksMu.Unlock() // Cancel one task - err = server.cancelTask(cancelCtx, "task-1") + err = server.cancelTask(cancelCtx, "task-1", false) require.NoError(t, err) // Now we should be able to create a new task diff --git a/server/task_response_format_test.go b/server/task_response_format_test.go index fd6b4b687..9851f7c5f 100644 --- a/server/task_response_format_test.go +++ b/server/task_response_format_test.go @@ -58,6 +58,7 @@ func TestTaskAugmentedToolCall_ResponseFormat(t *testing.T) { func TestTaskAugmentedToolCall_SpecCompliance(t *testing.T) { server := NewMCPServer("test", "1.0.0", WithTaskCapabilities(true, true, true), + WithExtensions(map[string]any{mcp.ExtensionTasks: map[string]any{}}), ) tool := mcp.NewTool("async_op", @@ -74,9 +75,13 @@ func TestTaskAugmentedToolCall_SpecCompliance(t *testing.T) { "method": "tools/call", "params": { "name": "async_op", - "task": { - "ttl": 60000, - "pollInterval": 5000 + "_meta": { + "io.modelcontextprotocol/protocolVersion": "2026-07-28", + "io.modelcontextprotocol/clientCapabilities": { + "extensions": { + "io.modelcontextprotocol/tasks": {} + } + } } } }`)) @@ -99,22 +104,59 @@ func TestTaskAugmentedToolCall_SpecCompliance(t *testing.T) { result, ok := parsed["result"].(map[string]any) require.True(t, ok, "result should be an object") - // Verify task is a direct field of result (not in _meta) - assert.Contains(t, result, "task", "task should be direct field of result") - assert.NotContains(t, result, "_meta", "task should NOT be in _meta") - - // Verify task structure - task, ok := result["task"].(map[string]any) - require.True(t, ok, "task should be an object") - - assert.Contains(t, task, "taskId") - assert.Contains(t, task, "status") - assert.Equal(t, "working", task["status"]) - assert.Contains(t, task, "createdAt") - assert.Contains(t, task, "lastUpdatedAt") - assert.Contains(t, task, "ttl") - assert.Equal(t, float64(60000), task["ttl"]) - // pollInterval is optional per spec, only check if provided + // Verify SEP-2663 inline structure + assert.Equal(t, "task", result["resultType"], "resultType should be 'task'") + assert.Contains(t, result, "taskId", "taskId should be a direct field of result") + + assert.Contains(t, result, "taskId") + assert.Contains(t, result, "status") + assert.Equal(t, "working", result["status"]) + assert.Contains(t, result, "createdAt") + assert.Contains(t, result, "lastUpdatedAt") + assert.Contains(t, result, "ttlMs") + assert.Nil(t, result["ttlMs"], "ttlMs should be null for a task with no expiration") + // pollIntervalMs is optional per spec, only check if provided // The server sets it to nil when not specified, so it won't appear in JSON - // This is correct behavior per JSON omitempty +} + +// TestLegacyCancelTaskResponseFieldNames verifies that a legacy tasks/cancel response +// carries the cancelled task state using 2025-11-25 field names (ttl, pollInterval), +// not the SEP-2663 names (ttlMs, pollIntervalMs). +func TestLegacyCancelTaskResponseFieldNames(t *testing.T) { + server := NewMCPServer("test", "1.0.0", WithTaskCapabilities(true, true, true)) + + ctx := t.Context() + ttl := int64(60000) + pollInterval := int64(1000) + _, err := server.createTask(ctx, "task-legacy-cancel", "test-tool", &ttl, &pollInterval) + require.NoError(t, err) + + // t.Context() carries no modern protocol info, so this takes the legacy path. + response := server.HandleMessage(ctx, []byte(`{ + "jsonrpc": "2.0", + "id": 1, + "method": "tasks/cancel", + "params": {"taskId": "task-legacy-cancel"} + }`)) + + jsonBytes, err := json.Marshal(response) + require.NoError(t, err) + + var parsed map[string]any + require.NoError(t, json.Unmarshal(jsonBytes, &parsed)) + result, ok := parsed["result"].(map[string]any) + require.True(t, ok, "result should be an object") + + assert.Equal(t, "task-legacy-cancel", result["taskId"]) + assert.Equal(t, "cancelled", result["status"]) + + // Legacy field names must be present. + assert.Contains(t, result, "ttl", "legacy cancel response must use ttl") + assert.Equal(t, float64(60000), result["ttl"]) + assert.Contains(t, result, "pollInterval", "legacy cancel response must use pollInterval") + assert.Equal(t, float64(1000), result["pollInterval"]) + + // SEP-2663 field names must NOT appear on the legacy path. + assert.NotContains(t, result, "ttlMs") + assert.NotContains(t, result, "pollIntervalMs") } diff --git a/server/task_test.go b/server/task_test.go index e2ecc3068..6f95ec59f 100644 --- a/server/task_test.go +++ b/server/task_test.go @@ -265,6 +265,7 @@ func TestMCPServer_HandleCancelTask(t *testing.T) { resp, ok := response.(mcp.JSONRPCResponse) require.True(t, ok, "Expected JSONRPCResponse, got %T", response) + // Legacy (non-modern) request returns CancelTaskResult with the task state. result, ok := resp.Result.(mcp.CancelTaskResult) require.True(t, ok, "Expected CancelTaskResult, got %T", resp.Result) @@ -288,7 +289,8 @@ func TestMCPServer_HandleCancelTerminalTask(t *testing.T) { require.NoError(t, err) server.completeTask(entry, "result", nil) - // Try to cancel completed task + // On the legacy 2025-11-25 protocol, cancelling a terminal task MUST be + // rejected with INVALID_PARAMS. response := server.HandleMessage(ctx, []byte(`{ "jsonrpc": "2.0", "id": 1, @@ -301,6 +303,7 @@ func TestMCPServer_HandleCancelTerminalTask(t *testing.T) { errResp, ok := response.(mcp.JSONRPCError) require.True(t, ok, "Expected JSONRPCError, got %T", response) assert.Equal(t, mcp.INVALID_PARAMS, errResp.Error.Code) + assert.Contains(t, errResp.Error.Message, "terminal status") } func TestMCPServer_TaskWithoutCapabilities(t *testing.T) { @@ -409,7 +412,7 @@ func TestMCPServer_TaskExpiredVsNotFoundErrors(t *testing.T) { assert.Error(t, err) assert.Contains(t, err.Error(), "task has expired") - err = server.cancelTask(ctx, "task-expired") + err = server.cancelTask(ctx, "task-expired", false) assert.Error(t, err) assert.Contains(t, err.Error(), "task has expired") }) @@ -426,7 +429,7 @@ func TestMCPServer_TaskExpiredVsNotFoundErrors(t *testing.T) { assert.Contains(t, err.Error(), "task not found") assert.NotContains(t, err.Error(), "expired") - err = server.cancelTask(ctx, "never-existed") + err = server.cancelTask(ctx, "never-existed", false) assert.Error(t, err) assert.Contains(t, err.Error(), "task not found") assert.NotContains(t, err.Error(), "expired") @@ -856,7 +859,7 @@ func TestMCPServer_TaskLastUpdatedAt(t *testing.T) { time.Sleep(1100 * time.Millisecond) // Cancel the task - cancelErr := server.cancelTask(ctx, "task-cancel") + cancelErr := server.cancelTask(ctx, "task-cancel", false) require.NoError(t, cancelErr) assert.Equal(t, mcp.TaskStatusCancelled, entry.task.Status) @@ -963,8 +966,8 @@ func TestMCPServer_TaskStatusNotifications(t *testing.T) { assert.Equal(t, mcp.TaskStatusCompleted, params["status"]) assert.NotEmpty(t, params["createdAt"]) assert.NotEmpty(t, params["lastUpdatedAt"]) - assert.Equal(t, int64(60000), params["ttl"]) - assert.Equal(t, int64(5000), params["pollInterval"]) + assert.Equal(t, int64(60000), params["ttlMs"]) + assert.Equal(t, int64(5000), params["pollIntervalMs"]) case <-time.After(1 * time.Second): t.Fatal("Expected task status notification but none received") @@ -1013,7 +1016,7 @@ func TestMCPServer_TaskStatusNotifications(t *testing.T) { } // Cancel the task - cancelErr := server.cancelTask(ctx, "task-notify-cancel") + cancelErr := server.cancelTask(ctx, "task-notify-cancel", false) require.NoError(t, cancelErr) // Check for notification @@ -1054,8 +1057,8 @@ func TestMCPServer_TaskStatusNotifications(t *testing.T) { params := notification.Params.AdditionalFields require.NotNil(t, params) - assert.Equal(t, int64(30000), params["ttl"]) - assert.Equal(t, int64(2000), params["pollInterval"]) + assert.Equal(t, int64(30000), params["ttlMs"]) + assert.Equal(t, int64(2000), params["pollIntervalMs"]) case <-time.After(1 * time.Second): t.Fatal("Expected task status notification but none received") @@ -1154,12 +1157,12 @@ func TestMCPServer_TaskStatusNotifications(t *testing.T) { // Complete the task once server.completeTask(entry, "result", nil) - // Wait for and consume first notification + // Wait for and consume the status notification. select { case <-notifyChan: - // First notification received as expected + // notification received as expected case <-time.After(1 * time.Second): - t.Fatal("Expected first notification") + t.Fatal("Expected notification") } // Try to complete again (should be ignored due to guard) diff --git a/server/task_tool_test.go b/server/task_tool_test.go index 2f25ebd00..db15d7edc 100644 --- a/server/task_tool_test.go +++ b/server/task_tool_test.go @@ -447,7 +447,7 @@ func TestTaskToolTracerBullet(t *testing.T) { <-started // Step 4: Cancel the task via tasks/cancel - cancelErr := server.cancelTask(ctx, taskID) + cancelErr := server.cancelTask(ctx, taskID, false) require.NoError(t, cancelErr) // Wait for cancellation to be detected @@ -531,23 +531,16 @@ func TestTaskToolTracerBullet(t *testing.T) { }) t.Run("task tool handler returns context.Canceled before tasks/cancel called", func(t *testing.T) { - // This test verifies that if a handler detects context cancellation - // (e.g., from parent context timeout) and returns ctx.Err() before - // tasks/cancel is explicitly called, the task is still marked as cancelled - // rather than failed. + // This test verifies that when a handler detects context cancellation + // via tasks/cancel and returns ctx.Err(), the task is marked as + // cancelled rather than failed. - // Step 1: Create server server := NewMCPServer( "test-handler-cancellation", "1.0.0", WithTaskCapabilities(true, true, true), ) - // Step 2: Create a parent context that we'll cancel - parentCtx, cancelParent := context.WithCancel(t.Context()) - defer cancelParent() - - // Step 3: Register a task tool that respects context cancellation handlerStarted := make(chan struct{}) selfCancelTool := mcp.NewTool("self_cancel_operation", @@ -556,13 +549,10 @@ func TestTaskToolTracerBullet(t *testing.T) { server.AddTaskTool(selfCancelTool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CreateTaskResult, error) { close(handlerStarted) - // Wait for context cancellation <-ctx.Done() - // Return the context error return nil, ctx.Err() }) - // Step 4: Call tool with task augmentation callRequest := mcp.CallToolRequest{ Params: mcp.CallToolParams{ Name: "self_cancel_operation", @@ -570,20 +560,19 @@ func TestTaskToolTracerBullet(t *testing.T) { }, } - callResult, callErr := server.handleToolCall(parentCtx, 1, callRequest) + callResult, callErr := server.handleToolCall(t.Context(), 1, callRequest) require.Nil(t, callErr) require.NotNil(t, callResult) createTaskResult := callResult.(*mcp.CreateTaskResult) taskID := createTaskResult.Task.TaskId - // Wait for handler to start <-handlerStarted - // Step 5: Cancel the parent context (simulating external cancellation) - cancelParent() + // Cancel via tasks/cancel. + require.NoError(t, server.cancelTask(t.Context(), taskID, false)) - // Step 6: Wait for task to complete + // Wait for task to reach a terminal state. var finalTask mcp.Task for range 20 { task, _, err := server.getTask(t.Context(), taskID) @@ -595,9 +584,7 @@ func TestTaskToolTracerBullet(t *testing.T) { time.Sleep(20 * time.Millisecond) } - // Step 7: Verify task status is cancelled (not failed) assert.Equal(t, mcp.TaskStatusCancelled, finalTask.Status) - assert.Contains(t, finalTask.StatusMessage, "context canceled") }) t.Run("multiple concurrent task tools", func(t *testing.T) { @@ -717,3 +704,120 @@ func TestTaskTool_ModelImmediateResponse(t *testing.T) { assert.Equal(t, message, immediateResponse) }) } + +// TestTaskToolOnly_ModernRequiresExtension verifies that a task-only tool +// (AddTaskTool) whose Execution does not force TaskSupportRequired is NOT +// routed to task creation for a modern request that has not declared the Tasks +// extension — even if the request carries a stray legacy params.task. Creating +// a task there would strand it: tasks/get and tasks/cancel reject modern +// requests without the extension, so the client could never read or cancel it. +// Because the tool cannot run synchronously and the request cannot be served +// without a CreateTaskResult, SEP-2663 requires MISSING_REQUIRED_CLIENT_CAPABILITY. +func TestTaskToolOnly_ModernRequiresExtension(t *testing.T) { + srv := NewMCPServer("test", "1.0.0", WithTaskCapabilities(true, true, true)) + + var handlerRan bool + // Task-only tool with no Execution.TaskSupport set, so the earlier + // TaskSupportRequired guard does not fire. + srv.AddTaskTool( + mcp.NewTool("batch"), + func(context.Context, mcp.CallToolRequest) (*mcp.CreateTaskResult, error) { + handlerRan = true + r := mcp.NewCreateTaskResult(mcp.Task{}) + return &r, nil + }) + + // Modern request WITHOUT the tasks extension, carrying a stray legacy task param. + ctx := WithRequestProtocolInfo(t.Context(), &RequestProtocolInfo{ + Modern: true, + ProtocolVersion: "2026-07-28", + ClientCapabilities: &mcp.ClientCapabilities{}, + }) + + _, reqErr := srv.handleToolCall(ctx, 1, mcp.CallToolRequest{ + Params: mcp.CallToolParams{Name: "batch", Task: &mcp.TaskParams{}}, + }) + + require.NotNil(t, reqErr, "modern no-extension request to a task-only tool must be rejected") + assert.Equal(t, mcp.MISSING_REQUIRED_CLIENT_CAPABILITY, reqErr.code) + assert.False(t, handlerRan, "handler must not run / no task should be created") + + // Sanity: no task was created. + srv.tasksMu.RLock() + n := len(srv.tasks) + srv.tasksMu.RUnlock() + assert.Equal(t, 0, n, "no task should have been created") +} + +// TestTaskToolOnly_ModernWithExtensionRoutes verifies the positive case: the +// same task-only tool IS routed to task creation when the modern client has +// declared the Tasks extension. +func TestTaskToolOnly_ModernWithExtensionRoutes(t *testing.T) { + srv := NewMCPServer("test", "1.0.0", WithTaskCapabilities(true, true, true)) + + srv.AddTaskTool( + mcp.NewTool("batch"), + func(context.Context, mcp.CallToolRequest) (*mcp.CreateTaskResult, error) { + r := mcp.NewCreateTaskResult(mcp.Task{}) + return &r, nil + }) + + ctx := WithRequestProtocolInfo(t.Context(), &RequestProtocolInfo{ + Modern: true, + ProtocolVersion: "2026-07-28", + ClientCapabilities: &mcp.ClientCapabilities{ + Extensions: map[string]any{mcp.ExtensionTasks: map[string]any{}}, + }, + }) + + created, reqErr := srv.handleToolCall(ctx, 1, mcp.CallToolRequest{ + Params: mcp.CallToolParams{Name: "batch"}, + }) + require.Nil(t, reqErr) + _, ok := created.(*mcp.CreateTaskResult) + assert.True(t, ok, "expected *CreateTaskResult, got %T", created) +} + +// TestTaskGetModern_NonEncodableResultReportsError verifies that when a +// completed task's stored result cannot be JSON-encoded (e.g. a tool placed a +// non-serializable value in StructuredContent), the modern tasks/get branch +// returns an error instead of a successful task with empty output. Silently +// dropping the output would make the tool result disappear with no signal. +func TestTaskGetModern_NonEncodableResultReportsError(t *testing.T) { + srv := NewMCPServer("test", "1.0.0", WithTaskCapabilities(true, true, true)) + + srv.AddTool( + mcp.NewTool("bad", mcp.WithTaskSupport(mcp.TaskSupportOptional)), + func(context.Context, mcp.CallToolRequest) (*mcp.CallToolResult, error) { + // A channel cannot be marshaled to JSON; json.Marshal will error. + return &mcp.CallToolResult{StructuredContent: make(chan int)}, nil + }) + + ctx := WithRequestProtocolInfo(t.Context(), &RequestProtocolInfo{ + Modern: true, + ProtocolVersion: "2026-07-28", + ClientCapabilities: &mcp.ClientCapabilities{ + Extensions: map[string]any{mcp.ExtensionTasks: map[string]any{}}, + }, + }) + + created, reqErr := srv.handleToolCall(ctx, 1, mcp.CallToolRequest{ + Params: mcp.CallToolParams{Name: "bad"}, + }) + require.Nil(t, reqErr) + taskID := created.(*mcp.CreateTaskResult).Task.TaskId + + // Wait until the task reaches a terminal (completed) state. + require.Eventually(t, func() bool { + srv.tasksMu.RLock() + defer srv.tasksMu.RUnlock() + e, ok := srv.tasks[taskID] + return ok && e.task.Status.IsTerminal() + }, 5*time.Second, 10*time.Millisecond) + + // Modern tasks/get must surface the encoding failure as an error, not a + // silent empty success. + _, e := srv.handleGetTask(ctx, 2, mcp.GetTaskRequest{Params: mcp.GetTaskParams{TaskId: taskID}}) + require.NotNil(t, e, "non-encodable completed result must produce an error") + assert.Equal(t, mcp.INTERNAL_ERROR, e.code) +}