-
Notifications
You must be signed in to change notification settings - Fork 874
fix: prevent task goroutine context from being canceled by HTTP request lifecycle #900
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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. | ||
|
Comment on lines
533
to
+543
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🛠️ Refactor suggestion | 🟠 Major | ⚡ Quick win This no longer tests the The handler now cancels itself via Also applies to: 577-601 🤖 Prompt for AI Agents |
||
|
|
||
| // 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 | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
🧩 Analysis chain
🌐 Web query:
In Go's standard library, does context.WithoutCancel preserve the parent context's deadline or Done channel?💡 Result:
In Go's standard library, context.WithoutCancel does not preserve the parent context's deadline, and its Done channel is nil [1][2][3]. According to the official documentation, context.WithoutCancel returns a derived context that points to the parent context but is explicitly designed to not be canceled when the parent is canceled [1][3]. Specifically, the returned context has the following properties: - Deadline: It returns no deadline [1][2]. - Done channel: Its Done channel is nil [1][2]. - Err: It returns no error [1][2]. - Cause: Calling Cause on it returns nil [1][2]. While it does not carry over the cancellation or deadline signals from the parent, it does maintain the ability to access values stored in the parent context through its Value method, as it still points to the parent [4]. This makes it useful for scenarios where you need to preserve request-scoped values for background tasks while detaching those tasks from the lifetime of the original request [4].
Citations:
🏁 Script executed:
Repository: mark3labs/mcp-go
Length of output: 12006
🏁 Script executed:
Repository: mark3labs/mcp-go
Length of output: 53361
🏁 Script executed:
Repository: mark3labs/mcp-go
Length of output: 9723
Ensure task TTL/timeout actually cancels background execution (not just bookkeeping)
executeTaskTool/executeRegularToolAsTaskbuildtaskCtxviacontext.WithCancel(context.WithoutCancel(ctx));context.WithoutCancelremoves the parent deadline and makesDonenil, so the task won’t stop when the request context ends.tasks/cancelcallsentry.cancelFunc()correctly (cancelTask), butscheduleTaskCleanup(on TTL expiry) onlydelete(s.tasks, taskID)/ writess.expiredTasksand never cancels the running goroutine.Fix by applying the TTL as a deadline on the task context (from an appropriate non-request parent) and/or invoking
entry.cancelFunc()(orcancelTask) when TTL cleanup triggers, before removing bookkeeping.🤖 Prompt for AI Agents