server: preserve task-augmented execution after request teardown

server/ctx.go

@@ -5,4 +5,5 @@ type contextKey int
 const (
 	// This const is used as key for context value lookup
 	requestHeader contextKey = iota
+	taskExecutionParentContext
 )

server/internal/gen/request_handler.go.tmpl

@@ -80,9 +80,12 @@ func (s *MCPServer) HandleMessage(
 		headers = make(http.Header)
 	}
 
+	taskExecutionCtx := context.WithoutCancel(ctx)
+
 	// Wrap context with cancel for in-flight request cancellation (MCP spec: notifications/cancelled)
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
+	ctx = context.WithValue(ctx, taskExecutionParentContext, taskExecutionCtx)
 
 	// Store cancel func so notifications/cancelled can cancel this request.
 	// Use session-scoped keys to prevent cross-session request ID collisions.

server/server.go

@@ -1985,13 +1985,18 @@ func (s *MCPServer) handleTaskAugmentedToolCall(
 		}
 	}
 
+	taskExecutionCtx := ctx
+	if detachedCtx, ok := ctx.Value(taskExecutionParentContext).(context.Context); ok && detachedCtx != nil {
+		taskExecutionCtx = detachedCtx
+	}
+
 	// Execute tool asynchronously
 	// For regular tools being used as tasks, we need different execution logic
 	if hasTaskHandler {
-		go s.executeTaskTool(ctx, entry, toolToUse, request)
+		go s.executeTaskTool(taskExecutionCtx, entry, toolToUse, request)
 	} else {
 		// Execute regular tool wrapped as a task
-		go s.executeRegularToolAsTask(ctx, entry, regularTool, request)
+		go s.executeRegularToolAsTask(taskExecutionCtx, entry, regularTool, request)
 	}
 
 	// Return CreateTaskResult immediately with task as top-level field

server/streamable_http_test.go

@@ -2322,6 +2322,99 @@
 	}
 }
 
+func TestStreamableHTTP_TaskAugmentedToolSurvivesRequestCompletion(t *testing.T) {
+	mcpServer := NewMCPServer("test-mcp-server", "1.0", WithTaskCapabilities(true, true, true))
+
+	release := make(chan struct{})
+	var releaseOnce sync.Once
+	releaseTool := func() {
+		releaseOnce.Do(func() {
+			close(release)
+		})
+	}
+	t.Cleanup(releaseTool)
+
+	mcpServer.AddTool(mcp.NewTool("async_tool",
+		mcp.WithDescription("A tool that completes asynchronously"),
+		mcp.WithTaskSupport(mcp.TaskSupportOptional),
+	), func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		case <-release:
+		}
+		return mcp.NewToolResultStructured(map[string]any{"status": "ok"}, "ok"), nil
+	})
+
+	server := NewTestStreamableHTTPServer(mcpServer, WithStateful(true))
+	defer server.Close()
+
+	resp, err := postJSON(server.URL, initRequest)
+	require.NoError(t, err)
+	sessionID := resp.Header.Get(HeaderKeySessionID)
+	resp.Body.Close()
+	require.NotEmpty(t, sessionID)
+
+	callResp, err := postSessionJSON(server.URL, sessionID, map[string]any{
+		"jsonrpc": "2.0",
+		"id":      2,
+		"method":  "tools/call",
+		"params": map[string]any{
+			"name": "async_tool",
+			"task": map[string]any{
+				"ttl": 60_000,
+			},
+		},
+	})
+	require.NoError(t, err)
+	defer callResp.Body.Close()
+
+	var created struct {
+		Result struct {
+			Task map[string]any `json:"task"`
+		} `json:"result"`
+	}
+	callBody, err := io.ReadAll(callResp.Body)
+	require.NoError(t, err)
+	if err := json.Unmarshal(callBody, &created); err != nil {
+		var decoded bool
+		for _, line := range strings.Split(string(callBody), "\n") {
+			if data, ok := strings.CutPrefix(line, "data: "); ok {
+				if json.Unmarshal([]byte(data), &created) == nil {
+					decoded = true
+					break
+				}
+			}
+		}
+		require.True(t, decoded, "decode task create response: %s", string(callBody))
+	}
+
+	taskID, _ := created.Result.Task["taskId"].(string)
+	require.NotEmpty(t, taskID)
+	assert.Equal(t, string(mcp.TaskStatusWorking), created.Result.Task["status"])
+
+	getResp, err := postSessionJSON(server.URL, sessionID, map[string]any{
+		"jsonrpc": "2.0",
+		"id":      3,
+		"method":  "tasks/get",
+		"params": map[string]any{
+			"taskId": taskID,
+		},
+	})
+	require.NoError(t, err)
+	defer getResp.Body.Close()
+
+	var taskStatus struct {
+		Result map[string]any `json:"result"`
+	}
+	getBody, err := io.ReadAll(getResp.Body)
+	require.NoError(t, err)
+	require.NoError(t, json.Unmarshal(getBody, &taskStatus), "decode tasks/get response: %s", string(getBody))
+	assert.Equal(t, string(mcp.TaskStatusWorking), taskStatus.Result["status"])
+
+	releaseTool()
+}
+
 // nonFlushingResponseWriter wraps an http.ResponseWriter but does NOT implement http.Flusher.
 // This is used to test the fix for servers/proxies that don't support streaming.
 type nonFlushingResponseWriter struct {