Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 10 additions & 4 deletions server/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Comment on lines +2023 to +2027

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | 🏗️ Heavy lift

🧩 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:

#!/bin/bash
set -euo pipefail

# Show the relevant segments around the referenced line ranges
sed -n '2000,2060p' server/server.go
echo "----"
sed -n '2080,2160p' server/server.go
echo "----"
sed -n '1920,2005p' server/server.go
echo "----"
sed -n '2480,2605p' server/server.go

# Locate createTask and task execution/cancellation usage
echo "==== createTask definition/uses ===="
rg -n "func (createTask|\\w*createTask)\\b|\\bcreateTask\\b" server/server.go
echo "----"
# Print the function(s) containing createTask, if found
rg -n "createTask\\(" server/server.go | head -n 20

Repository: mark3labs/mcp-go

Length of output: 12006


🏁 Script executed:

#!/bin/bash
set -euo pipefail

rg -n "scheduleTaskCleanup\\b|cancelTask\\b|tasks/cancel|TaskStatusCancelled|cancelFunc\\b" server/server.go
echo "----"
# Show likely relevant implementations around those matches
# (Using approximate line windows from rg results via grep for line numbers)
python3 - <<'PY'
import re,subprocess,os,sys,json,math
from pathlib import Path
text=subprocess.check_output(["bash","-lc","rg -n \"scheduleTaskCleanup\\\\b|cancelTask\\\\b|tasks/cancel|TaskStatusCancelled|cancelFunc\\\\b\" server/server.go"],text=True)
lines=[]
for line in text.strip().splitlines():
    m=re.match(r"(\d+):", line)
    if m:
        lines.append(int(m.group(1)))
# print windows around each match
for ln in sorted(set(lines)):
    start=max(1,ln-30); end=ln+80
    print(f"\n===== server/server.go:{start}-{end} (around {ln}) =====")
    subprocess.run(["bash","-lc", f"sed -n '{start},{end}p' server/server.go"], check=False)
PY

Repository: mark3labs/mcp-go

Length of output: 53361


🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Narrow, targeted inspection to avoid truncation
echo "===== executeTaskTool context-withoutcancel sites ====="
sed -n '2010,2055p' server/server.go
echo "----"
sed -n '2105,2175p' server/server.go

echo "===== createTask TTL scheduling call ====="
sed -n '2470,2565p' server/server.go

echo "===== cancelTask implementation ====="
sed -n '2660,2765p' server/server.go

Repository: mark3labs/mcp-go

Length of output: 9723


Ensure task TTL/timeout actually cancels background execution (not just bookkeeping)

  • executeTaskTool / executeRegularToolAsTask build taskCtx via context.WithCancel(context.WithoutCancel(ctx)); context.WithoutCancel removes the parent deadline and makes Done nil, so the task won’t stop when the request context ends.
  • Cancellation via tasks/cancel calls entry.cancelFunc() correctly (cancelTask), but scheduleTaskCleanup (on TTL expiry) only delete(s.tasks, taskID) / writes s.expiredTasks and never cancels the running goroutine.
  • This means work can continue running after TTL expiry and even after the task record is removed, leaving no reliable way to stop it via TTL.

Fix by applying the TTL as a deadline on the task context (from an appropriate non-request parent) and/or invoking entry.cancelFunc() (or cancelTask) when TTL cleanup triggers, before removing bookkeeping.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@server/server.go` around lines 2023 - 2027, The task context is currently
created with context.WithCancel(context.WithoutCancel(ctx)) which removes parent
deadlines so TTL cleanup only deletes bookkeeping; fix by making the TTL
actually cancel execution: in executeTaskTool / executeRegularToolAsTask create
taskCtx from a non-request parent (e.g., context.Background()) and apply the TTL
via context.WithTimeout or context.WithDeadline to produce a cancellable context
assigned to taskCtx (or ensure scheduleTaskCleanup calls entry.cancelFunc() /
cancelTask before deleting s.tasks); update scheduleTaskCleanup to invoke
entry.cancelFunc() when TTL expires (then remove bookkeeping) so background
goroutines are actually stopped.

defer cancel()

// Store cancel func in entry so it can be cancelled via tasks/cancel
Expand Down Expand Up @@ -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
Expand Down
81 changes: 41 additions & 40 deletions server/task_tool_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion | 🟠 Major | ⚡ Quick win

This no longer tests the #897 regression.

The handler now cancels itself via innerCtx, but the parent/request context is never canceled. If the production code regressed to context.WithCancel(ctx), this subtest would still pass because the inner timeout fires first. Please keep a case that cancels the parent right after handleToolCall and asserts the task stays alive until its own cancellation path fires.

Also applies to: 577-601

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@server/task_tool_test.go` around lines 533 - 543, The test "task tool handler
returns context.Canceled before tasks/cancel called" no longer verifies the `#897`
regression because the handler cancels an inner context (innerCtx) while the
parent/request context is never cancelled; add an additional subcase that
explicitly cancels the parent/request context immediately after calling
handleToolCall (simulate parent cancellation), then assert the task is not
marked cancelled/failed until the handler's own cancellation path runs (i.e.,
the task remains alive until innerCtx times out); apply the same change to the
similar test block around lines 577-601 so both scenarios (parent-cancel and
handler-internal-cancel) are covered using the existing test helpers and the
same assertions on task state.


// Step 1: Create server
server := NewMCPServer(
Expand All @@ -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",
Expand All @@ -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
Expand Down
Loading