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