.NET: Add MCP long-running task support for MCP client tools (#5994)

* Add MCP long-running task support for MCP client tools * Fixed project file formatting issue. * Removed experimentation tag from MCP alpha project. * Addressed PR comments
2026-06-16 21:04:09 +08:00 · 2026-05-22 12:09:54 -07:00
parent 9fdd7429a8
commit 793403f3db
17 changed files with 960 additions and 0 deletions
@@ -0,0 +1,61 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System.Collections.Generic;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.Extensions.AI;
+using Microsoft.Shared.Diagnostics;
+using ModelContextProtocol.Client;
+using ModelContextProtocol.Protocol;
+
+namespace Microsoft.Agents.AI.Mcp;
+
+/// <summary>
+/// Extension methods on <see cref="McpClient"/> that expose MCP server tools to a Microsoft
+/// Agent Framework agent with optional long-running task (SEP-2663) handling.
+/// </summary>
+public static class McpClientTaskExtensions
+{
+    /// <summary>
+    /// Lists tools advertised by the connected MCP server and returns each as an
+    /// <see cref="AIFunction"/>. Tools that declare <see cref="ToolTaskSupport.Required"/>
+    /// are wrapped with task-aware behavior so an agent can transparently drive long-running
+    /// invocations. All other tools — including those that declare
+    /// <see cref="ToolTaskSupport.Optional"/> — are returned as-is, preserving inline
+    /// (synchronous) invocation semantics by default.
+    /// </summary>
+    /// <param name="client">The connected MCP client.</param>
+    /// <param name="options">
+    /// Options that control the task lifecycle for task-capable tools.
+    /// When <see langword="null"/>, defaults described on <see cref="McpTaskOptions"/> apply.
+    /// </param>
+    /// <param name="cancellationToken">Token used to cancel listing the server's tools.</param>
+    /// <returns>The tools, ready to pass to <c>AsAIAgent(tools: …)</c>.</returns>
+    public static async Task<IReadOnlyList<AIFunction>> ListAgentToolsWithTaskSupportAsync(
+        this McpClient client,
+        McpTaskOptions? options = null,
+        CancellationToken cancellationToken = default)
+    {
+        _ = Throw.IfNull(client);
+
+        McpTaskOptions effectiveOptions = options ?? new McpTaskOptions();
+
+        IList<McpClientTool> tools = await client.ListToolsAsync(cancellationToken: cancellationToken).ConfigureAwait(false);
+
+        AIFunction[] result = new AIFunction[tools.Count];
+        for (int i = 0; i < tools.Count; i++)
+        {
+            ToolTaskSupport? taskSupport = tools[i].ProtocolTool.Execution?.TaskSupport;
+            if (taskSupport is ToolTaskSupport.Required)
+            {
+                result[i] = new TaskAwareMcpClientAIFunction(client, tools[i], effectiveOptions);
+            }
+            else
+            {
+                result[i] = tools[i];
+            }
+        }
+
+        return result;
+    }
+}
@@ -0,0 +1,39 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+
+namespace Microsoft.Agents.AI.Mcp;
+
+/// <summary>
+/// Configures how an MCP client wrapper drives the
+/// <see href="https://modelcontextprotocol.io/specification/2025-11-25/basic/utilities/tasks">MCP tasks</see>
+/// lifecycle when an underlying server tool returns a <c>CreateTaskResult</c>.
+/// </summary>
+/// <remarks>
+/// <para>
+/// All members of this type are subject to change. The MCP task surface is experimental
+/// and tracks the in-flight specification.
+/// </para>
+/// </remarks>
+public sealed class McpTaskOptions
+{
+    /// <summary>
+    /// Gets or sets the time-to-live the wrapper attaches to a newly created server-side task.
+    /// </summary>
+    /// <remarks>
+    /// When <see langword="null"/> the wrapper omits the <c>ttl</c> hint and lets the server
+    /// pick its own value. The server's chosen TTL is always authoritative.
+    /// </remarks>
+    public TimeSpan? DefaultTimeToLive { get; set; }
+
+    /// <summary>
+    /// Gets or sets a value indicating whether the wrapper should send
+    /// <c>tasks/cancel</c> when the local <see cref="System.Threading.CancellationToken"/>
+    /// fires during a tool invocation.
+    /// </summary>
+    /// <remarks>
+    /// Defaults to <see langword="true"/>: a local cancellation means "the caller is giving up
+    /// on this tool invocation" and the server-side task has no further consumer.
+    /// </remarks>
+    public bool CancelRemoteTaskOnLocalCancellation { get; set; } = true;
+}
@@ -0,0 +1,37 @@
+<Project Sdk="Microsoft.NET.Sdk">
+
+  <PropertyGroup>
+    <TargetFrameworks>$(TargetFrameworksCore)</TargetFrameworks>
+    <RootNamespace>Microsoft.Agents.AI.Mcp</RootNamespace>
+    <VersionSuffix>alpha</VersionSuffix>
+    <NoWarn>$(NoWarn);MEAI001;MCPEXP001</NoWarn>
+  </PropertyGroup>
+
+  <Import Project="$(RepoRoot)/dotnet/nuget/nuget-package.props" />
+
+  <PropertyGroup>
+    <InjectSharedThrow>true</InjectSharedThrow>
+    <InjectTrimAttributesOnLegacy>true</InjectTrimAttributesOnLegacy>
+  </PropertyGroup>
+
+  <PropertyGroup>
+    <Title>Microsoft Agent Framework MCP</Title>
+    <Description>Provides Microsoft Agent Framework support for Model Context Protocol (MCP), including long-running task (SEP-2663) integration for MCP clients.</Description>
+  </PropertyGroup>
+
+  <!-- Disable package validation baseline until the first release -->
+  <PropertyGroup>
+    <PackageValidationBaselineVersion />
+    <EnablePackageValidation>false</EnablePackageValidation>
+  </PropertyGroup>
+
+  <ItemGroup>
+    <PackageReference Include="Microsoft.Extensions.AI" />
+    <PackageReference Include="ModelContextProtocol" />
+  </ItemGroup>
+
+  <ItemGroup>
+    <InternalsVisibleTo Include="Microsoft.Agents.AI.Mcp.UnitTests" />
+  </ItemGroup>
+
+</Project>
@@ -0,0 +1,147 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Text.Json;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.Extensions.AI;
+using Microsoft.Shared.Diagnostics;
+using ModelContextProtocol;
+using ModelContextProtocol.Client;
+using ModelContextProtocol.Protocol;
+
+namespace Microsoft.Agents.AI.Mcp;
+
+/// <summary>
+/// An <see cref="AIFunction"/> wrapper around an <see cref="McpClientTool"/> that drives the
+/// <see href="https://modelcontextprotocol.io/specification/2025-11-25/basic/utilities/tasks">MCP long-running task</see>
+/// lifecycle (SEP-2663) on behalf of the agent's tool loop.
+/// </summary>
+/// <remarks>
+/// <para>
+/// The wrapper invokes the tool with task augmentation via
+/// <see cref="McpClient.CallToolAsTaskAsync"/>, polls to completion via
+/// <see cref="McpClient.PollTaskUntilCompleteAsync"/>, and fetches the result via
+/// <see cref="McpClient.GetTaskResultAsync"/>. The result is returned to the caller as a
+/// <see cref="JsonElement"/> containing the serialized <see cref="CallToolResult"/> — the
+/// same wire shape produced by <see cref="McpClientTool"/>.<see cref="AIFunction.InvokeAsync(AIFunctionArguments, CancellationToken)"/>
+/// so that downstream <see cref="FunctionResultContent"/> serialization is byte-identical to
+/// a non-task-augmented MCP tool call. The agent's function-calling loop is unaware that a
+/// task was used.
+/// </para>
+/// <para>
+/// This wrapper is intended to be applied only to tools whose
+/// <see cref="ToolExecution.TaskSupport"/> is <see cref="ToolTaskSupport.Required"/>
+/// (selected by <see cref="McpClientTaskExtensions.ListAgentToolsWithTaskSupportAsync"/>).
+/// As a defensive fallback, if the server still rejects the task-augmented call with
+/// <see cref="McpErrorCode.MethodNotFound"/> (e.g. because tool-level capabilities changed
+/// between <c>tools/list</c> and invocation), the wrapper transparently falls back to a
+/// non-augmented call through the inner <see cref="McpClientTool"/>.
+/// </para>
+/// </remarks>
+internal sealed class TaskAwareMcpClientAIFunction : AIFunction
+{
+    private readonly McpClient _client;
+    private readonly McpClientTool _inner;
+    private readonly McpTaskOptions _options;
+
+    internal TaskAwareMcpClientAIFunction(McpClient client, McpClientTool inner, McpTaskOptions options)
+    {
+        _ = Throw.IfNull(client);
+        _ = Throw.IfNull(inner);
+        _ = Throw.IfNull(options);
+
+        this._client = client;
+        this._inner = inner;
+        this._options = options;
+    }
+
+    /// <inheritdoc />
+    public override string Name => this._inner.Name;
+
+    /// <inheritdoc />
+    public override string Description => this._inner.Description;
+
+    /// <inheritdoc />
+    public override JsonElement JsonSchema => this._inner.JsonSchema;
+
+    /// <inheritdoc />
+    public override JsonElement? ReturnJsonSchema => this._inner.ReturnJsonSchema;
+
+    /// <inheritdoc />
+    public override JsonSerializerOptions JsonSerializerOptions => this._inner.JsonSerializerOptions;
+
+    /// <inheritdoc />
+    protected override async ValueTask<object?> InvokeCoreAsync(
+        AIFunctionArguments arguments,
+        CancellationToken cancellationToken)
+    {
+        _ = Throw.IfNull(arguments);
+
+        McpTaskMetadata? metadata = null;
+        if (this._options.DefaultTimeToLive is TimeSpan ttl)
+        {
+            metadata = new McpTaskMetadata { TimeToLive = ttl };
+        }
+
+        McpTask task;
+        try
+        {
+            task = await this._client.CallToolAsTaskAsync(
+                this._inner.Name,
+                arguments,
+                taskMetadata: metadata,
+                progress: null,
+                options: null,
+                cancellationToken: cancellationToken).ConfigureAwait(false);
+        }
+        catch (McpProtocolException ex) when (ex.ErrorCode == McpErrorCode.MethodNotFound)
+        {
+            // Defensive fallback: the server's advertised TaskSupport indicated this tool
+            // could be invoked as a task, but the server now rejects task augmentation for it
+            // (e.g. capability changed between tools/list and invocation). Fall back to a
+            // non-augmented call through the inner McpClientTool.
+            return await this._inner.InvokeAsync(arguments, cancellationToken).ConfigureAwait(false);
+        }
+
+        return await this.PollAndRetrieveResultAsync(task.TaskId, cancellationToken).ConfigureAwait(false);
+    }
+
+    private async Task<JsonElement> PollAndRetrieveResultAsync(string taskId, CancellationToken cancellationToken)
+    {
+        try
+        {
+            McpTask terminal = await this._client.PollTaskUntilCompleteAsync(taskId, options: null, cancellationToken).ConfigureAwait(false);
+
+            return terminal.Status switch
+            {
+                McpTaskStatus.Completed => await this._client.GetTaskResultAsync(taskId, options: null, cancellationToken).ConfigureAwait(false),
+                McpTaskStatus.Cancelled => throw new OperationCanceledException(FormatTerminalStatusMessage(taskId, terminal)),
+                _ => throw new InvalidOperationException(FormatTerminalStatusMessage(taskId, terminal)),// Failed (or any future non-terminal-but-unhandled status that the poll loop returns).
+            };
+        }
+        catch (OperationCanceledException) when (this._options.CancelRemoteTaskOnLocalCancellation && cancellationToken.IsCancellationRequested)
+        {
+            await this.TryCancelTaskAsync(taskId).ConfigureAwait(false);
+            throw;
+        }
+    }
+
+    private static string FormatTerminalStatusMessage(string taskId, McpTask terminal)
+        => string.IsNullOrEmpty(terminal.StatusMessage)
+            ? $"MCP task '{taskId}' ended in terminal status '{terminal.Status}'."
+            : $"MCP task '{taskId}' ended in terminal status '{terminal.Status}': {terminal.StatusMessage}";
+
+    private async Task TryCancelTaskAsync(string taskId)
+    {
+        try
+        {
+            using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(5));
+            _ = await this._client.CancelTaskAsync(taskId, options: null, cts.Token).ConfigureAwait(false);
+        }
+        catch
+        {
+            // Best-effort cancellation; do not mask the original cancellation reason.
+        }
+    }
+}