From f1408f67a2a8d40d89c33906f9297111653b2291 Mon Sep 17 00:00:00 2001 From: xlyoung Date: Sun, 31 May 2026 21:01:46 +0800 Subject: [PATCH] fix: prevent task goroutine context from being canceled by HTTP request lifecycle Fixes #897 When a tool registered with WithTaskSupport(TaskSupportOptional) or AddTaskTool is called asynchronously via streamable HTTP, the task context was derived from the HTTP request context (r.Context()). After the HTTP handler returned the CreateTaskResult response, Go's HTTP server canceled r.Context(), which cascaded into the task context and killed the handler goroutine. Fix: Use context.WithoutCancel(ctx) to detach from the HTTP request lifecycle while preserving context values (session metadata, etc.). The task context is now only canceled by: - The task's own timeout/TTL - Explicit tasks/cancel from the client - Server shutdown Both executeTaskTool and executeRegularToolAsTask are fixed. Updated the corresponding test to verify internal handler cancellation (timeout-based) rather than relying on parent context cancellation, which is no longer propagated to task goroutines. --- server/server.go | 14 +++++-- server/task_tool_test.go | 81 ++++++++++++++++++++-------------------- 2 files changed, 51 insertions(+), 44 deletions(-) diff --git a/server/server.go b/server/server.go index 15a96800b..e5161a8d1 100644 --- a/server/server.go +++ b/server/server.go @@ -2020,8 +2020,11 @@ func (s *MCPServer) executeTaskTool( } }() - // Create cancellable context for this task execution - taskCtx, cancel := context.WithCancel(ctx) + // Create cancellable context for this task execution. + // Use WithoutCancel to detach from the HTTP request lifecycle so the task + // continues running after the HTTP handler returns. Context values (session + // metadata, etc.) are preserved. + taskCtx, cancel := context.WithCancel(context.WithoutCancel(ctx)) defer cancel() // Store cancel func in entry so it can be cancelled via tasks/cancel @@ -2112,8 +2115,11 @@ func (s *MCPServer) executeRegularToolAsTask( regularTool ServerTool, request mcp.CallToolRequest, ) { - // Create cancellable context for this task execution - taskCtx, cancel := context.WithCancel(ctx) + // Create cancellable context for this task execution. + // Use WithoutCancel to detach from the HTTP request lifecycle so the task + // continues running after the HTTP handler returns. Context values (session + // metadata, etc.) are preserved. + taskCtx, cancel := context.WithCancel(context.WithoutCancel(ctx)) defer cancel() // Store cancel func in entry so it can be cancelled via tasks/cancel diff --git a/server/task_tool_test.go b/server/task_tool_test.go index 2f25ebd00..357373b2f 100644 --- a/server/task_tool_test.go +++ b/server/task_tool_test.go @@ -532,9 +532,15 @@ 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 + // (e.g., from its own internal timeout) and returns ctx.Err() before // tasks/cancel is explicitly called, the task is still marked as cancelled // rather than failed. + // + // Note: With the fix for #897 (context.WithoutCancel), the task context + // is now detached from the parent/HTTP request context. This means the + // handler must manage its own cancellation via the task's context, not + // rely on parent context cancellation. We test internal handler cancellation + // by having the handler cancel its own context via a timeout. // Step 1: Create server server := NewMCPServer( @@ -543,11 +549,7 @@ func TestTaskToolTracerBullet(t *testing.T) { 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 + // Step 2: Register a task tool that simulates internal cancellation handlerStarted := make(chan struct{}) selfCancelTool := mcp.NewTool("self_cancel_operation", @@ -556,49 +558,48 @@ 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() + // Simulate internal cancellation by creating a context that cancels quickly + innerCtx, innerCancel := context.WithTimeout(ctx, 10*time.Millisecond) + defer innerCancel() + <-innerCtx.Done() // Return the context error - return nil, ctx.Err() + return nil, innerCtx.Err() }) - // Step 4: Call tool with task augmentation - callRequest := mcp.CallToolRequest{ - Params: mcp.CallToolParams{ - Name: "self_cancel_operation", - Task: &mcp.TaskParams{}, - }, - } + // Step 4: Call tool with task augmentation + callRequest := mcp.CallToolRequest{ + Params: mcp.CallToolParams{ + Name: "self_cancel_operation", + Task: &mcp.TaskParams{}, + }, + } - callResult, callErr := server.handleToolCall(parentCtx, 1, callRequest) - require.Nil(t, callErr) - require.NotNil(t, callResult) + 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 + createTaskResult := callResult.(*mcp.CreateTaskResult) + taskID := createTaskResult.Task.TaskId - // Step 5: Cancel the parent context (simulating external cancellation) - cancelParent() + // Wait for handler to start + <-handlerStarted - // Step 6: Wait for task to complete - var finalTask mcp.Task - for range 20 { - task, _, err := server.getTask(t.Context(), taskID) - require.NoError(t, err) - finalTask = task - if task.Status.IsTerminal() { - break - } - time.Sleep(20 * time.Millisecond) + // Step 5: Wait for task to complete (handler self-cancels via timeout) + var finalTask mcp.Task + for range 20 { + task, _, err := server.getTask(t.Context(), taskID) + require.NoError(t, err) + finalTask = task + if task.Status.IsTerminal() { + break } + 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") - }) + // Step 6: Verify task status is cancelled (not failed) + assert.Equal(t, mcp.TaskStatusCancelled, finalTask.Status) + assert.Contains(t, finalTask.StatusMessage, "context deadline exceeded") +}) t.Run("multiple concurrent task tools", func(t *testing.T) { // Step 1: Create server