Fix McpTask double-wrapping: add McpTask case in InvokeAsync switch and bypass ExecuteToolAsTaskAsync for McpTask-returning tools

Agent-Logs-Url: https://github.com/modelcontextprotocol/csharp-sdk/sessions/0b57604e-d75a-4879-a7ed-139c8e03ec6b

Co-authored-by: jeffhandley <1031940+jeffhandley@users.noreply.github.com>

Author: Copilot

src/ModelContextProtocol.Core/Server/AIFunctionMcpServerTool.cs

@@ -244,11 +244,30 @@ private AIFunctionMcpServerTool(AIFunction function, Tool tool, IServiceProvider
 
         _structuredOutputRequiresWrapping = structuredOutputRequiresWrapping;
         _metadata = metadata;
+
+        // Detect if the tool's underlying method returns McpTask (directly or wrapped in Task<>/ValueTask<>).
+        if (function.UnderlyingMethod is { } method)
+        {
+            Type returnType = method.ReturnType;
+            if (returnType.IsGenericType)
+            {
+                Type gt = returnType.GetGenericTypeDefinition();
+                if (gt == typeof(Task<>) || gt == typeof(ValueTask<>))
+                {
+                    returnType = returnType.GetGenericArguments()[0];
+                }
+            }
+
+            ReturnsMcpTask = returnType == typeof(McpTask);
+        }
     }
 
     /// <inheritdoc />
     public override Tool ProtocolTool { get; }
 
+    /// <inheritdoc />
+    internal override bool ReturnsMcpTask { get; }
+
     /// <inheritdoc />
     public override IReadOnlyList<object> Metadata => _metadata;
 
@@ -311,6 +330,11 @@ public override async ValueTask<CallToolResult> InvokeAsync(
 
             CallToolResult callToolResponse => callToolResponse,
 
+            McpTask mcpTask => new()
+            {
+                Task = mcpTask,
+            },
+
             _ => new()
             {
                 Content = [new TextContentBlock { Text = JsonSerializer.Serialize(result, AIFunction.JsonSerializerOptions.GetTypeInfo(typeof(object))) }],

src/ModelContextProtocol.Core/Server/DelegatingMcpServerTool.cs

@@ -23,6 +23,9 @@ protected DelegatingMcpServerTool(McpServerTool innerTool)
     /// <inheritdoc />
     public override Tool ProtocolTool => _innerTool.ProtocolTool;
 
+    /// <inheritdoc />
+    internal override bool ReturnsMcpTask => _innerTool.ReturnsMcpTask;
+
     /// <inheritdoc />
     public override IReadOnlyList<object> Metadata => _innerTool.Metadata;
 

src/ModelContextProtocol.Core/Server/McpServerImpl.cs

@@ -707,6 +707,13 @@ await originalListToolsHandler(request, cancellationToken).ConfigureAwait(false)
                                 McpErrorCode.InvalidParams);
                         }
 
+                        // If the tool manages its own task lifecycle (returns McpTask),
+                        // invoke it directly and return its result without SDK task wrapping.
+                        if (tool.ReturnsMcpTask)
+                        {
+                            return await tool.InvokeAsync(request, cancellationToken).ConfigureAwait(false);
+                        }
+
                         // Task augmentation requested - return CreateTaskResult
                         return await ExecuteToolAsTaskAsync(tool, request, taskMetadata, taskStore, sendNotifications, cancellationToken).ConfigureAwait(false);
                     }

src/ModelContextProtocol.Core/Server/McpServerTool.cs

@@ -157,6 +157,10 @@ protected McpServerTool()
     /// <summary>Gets the protocol <see cref="Tool"/> type for this instance.</summary>
     public abstract Tool ProtocolTool { get; }
 
+    /// <summary>Gets whether the tool's underlying method returns <see cref="McpTask"/>, indicating
+    /// it manages its own task lifecycle and should not be wrapped by the SDK.</summary>
+    internal virtual bool ReturnsMcpTask => false;
+
     /// <summary>
     /// Gets the metadata for this tool instance.
     /// </summary>

tests/ModelContextProtocol.Tests/Server/ToolTaskSupportTests.cs

@@ -657,7 +657,106 @@ public async Task TaskPath_Logs_Error_When_Tool_Throws()
 
     #endregion
 
-    /// <summary>
+    #region Tool Returning McpTask Tests
+
+#pragma warning disable MCPEXP001 // Tasks feature is experimental
+    [Fact]
+    public async Task Tool_ReturningMcpTask_BypassesSdkTaskWrapping()
+    {
+        // Arrange - Server with task store and a tool that creates and returns its own McpTask
+        var taskStore = new InMemoryMcpTaskStore();
+
+        await using var fixture = new ClientServerFixture(
+            LoggerFactory,
+            configureServer: builder =>
+            {
+                builder.WithTools([McpServerTool.Create(
+                    async (IMcpTaskStore store) =>
+                    {
+                        // Tool creates its own task explicitly
+                        var task = await store.CreateTaskAsync(
+                            new McpTaskMetadata(),
+                            new RequestId("tool-req"),
+                            new JsonRpcRequest { Method = "tools/call" });
+                        return task;
+                    },
+                    new McpServerToolCreateOptions
+                    {
+                        Name = "self-managing-tool",
+                        Description = "A tool that creates and returns its own McpTask",
+                        Execution = new ToolExecution { TaskSupport = ToolTaskSupport.Optional }
+                    })]);
+            },
+            configureServices: services =>
+            {
+                services.AddSingleton<IMcpTaskStore>(taskStore);
+                services.Configure<McpServerOptions>(options => options.TaskStore = taskStore);
+            });
+
+        // Act - Call the tool with task metadata (task-augmented request)
+        var callResult = await fixture.Client.CallToolAsync(
+            new CallToolRequestParams
+            {
+                Name = "self-managing-tool",
+                Task = new McpTaskMetadata()
+            },
+            cancellationToken: TestContext.Current.CancellationToken);
+
+        // Assert - Only 1 task should exist (the tool-created one), not 2
+        Assert.NotNull(callResult.Task);
+
+        var allTasks = await fixture.Client.ListTasksAsync(cancellationToken: TestContext.Current.CancellationToken);
+        Assert.Single(allTasks);
+        Assert.Equal(callResult.Task.TaskId, allTasks[0].TaskId);
+    }
+
+    [Fact]
+    public async Task Tool_ReturningMcpTask_WithoutTaskMetadata_ReturnsTaskDirectly()
+    {
+        // Arrange - Server with task store and a tool that returns McpTask
+        var taskStore = new InMemoryMcpTaskStore();
+
+        await using var fixture = new ClientServerFixture(
+            LoggerFactory,
+            configureServer: builder =>
+            {
+                builder.WithTools([McpServerTool.Create(
+                    async (IMcpTaskStore store) =>
+                    {
+                        var task = await store.CreateTaskAsync(
+                            new McpTaskMetadata(),
+                            new RequestId("tool-req"),
+                            new JsonRpcRequest { Method = "tools/call" });
+                        return task;
+                    },
+                    new McpServerToolCreateOptions
+                    {
+                        Name = "self-managing-tool",
+                        Description = "A tool that creates and returns its own McpTask",
+                        Execution = new ToolExecution { TaskSupport = ToolTaskSupport.Optional }
+                    })]);
+            },
+            configureServices: services =>
+            {
+                services.AddSingleton<IMcpTaskStore>(taskStore);
+                services.Configure<McpServerOptions>(options => options.TaskStore = taskStore);
+            });
+
+        // Act - Call without task metadata (normal invocation)
+        var callResult = await fixture.Client.CallToolAsync(
+            "self-managing-tool",
+            cancellationToken: TestContext.Current.CancellationToken);
+
+        // Assert - Tool's McpTask should be on the result
+        Assert.NotNull(callResult.Task);
+
+        var allTasks = await fixture.Client.ListTasksAsync(cancellationToken: TestContext.Current.CancellationToken);
+        Assert.Single(allTasks);
+        Assert.Equal(callResult.Task.TaskId, allTasks[0].TaskId);
+    }
+#pragma warning restore MCPEXP001
+
+    #endregion
     /// A fixture that creates a connected MCP client-server pair for testing.
     /// </summary>
     private sealed class ClientServerFixture : IAsyncDisposable