.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
@@ -212,6 +212,7 @@
  </Folder>
  <Folder Name="/Samples/02-agents/ModelContextProtocol/">
    <File Path="samples/02-agents/ModelContextProtocol/README.md" />
+    <Project Path="samples/02-agents/ModelContextProtocol/Agent_MCP_LongRunningTask_Client/Agent_MCP_LongRunningTask_Client.csproj" />
    <Project Path="samples/02-agents/ModelContextProtocol/Agent_MCP_Server/Agent_MCP_Server.csproj" />
    <Project Path="samples/02-agents/ModelContextProtocol/Agent_MCP_Server_Auth/Agent_MCP_Server_Auth.csproj" />
    <Project Path="samples/02-agents/ModelContextProtocol/FoundryAgent_Hosted_MCP/FoundryAgent_Hosted_MCP.csproj" />
@@ -602,6 +603,7 @@
    <Project Path="src/Microsoft.Agents.AI.Hosting.OpenAI/Microsoft.Agents.AI.Hosting.OpenAI.csproj" />
    <Project Path="src/Microsoft.Agents.AI.Hosting/Microsoft.Agents.AI.Hosting.csproj" />
    <Project Path="src/Microsoft.Agents.AI.Hyperlight/Microsoft.Agents.AI.Hyperlight.csproj" />
+    <Project Path="src/Microsoft.Agents.AI.Mcp/Microsoft.Agents.AI.Mcp.csproj" />
    <Project Path="src/Microsoft.Agents.AI.Mem0/Microsoft.Agents.AI.Mem0.csproj" />
    <Project Path="src/Microsoft.Agents.AI.OpenAI/Microsoft.Agents.AI.OpenAI.csproj" />
    <Project Path="src/Microsoft.Agents.AI.Purview/Microsoft.Agents.AI.Purview.csproj" />
@@ -655,6 +657,7 @@
    <Project Path="tests/Microsoft.Agents.AI.Hosting.OpenAI.UnitTests/Microsoft.Agents.AI.Hosting.OpenAI.UnitTests.csproj" />
    <Project Path="tests/Microsoft.Agents.AI.Hosting.UnitTests/Microsoft.Agents.AI.Hosting.UnitTests.csproj" />
    <Project Path="tests/Microsoft.Agents.AI.Hyperlight.UnitTests/Microsoft.Agents.AI.Hyperlight.UnitTests.csproj" />
+    <Project Path="tests/Microsoft.Agents.AI.Mcp.UnitTests/Microsoft.Agents.AI.Mcp.UnitTests.csproj" />
    <Project Path="tests/Microsoft.Agents.AI.Mem0.UnitTests/Microsoft.Agents.AI.Mem0.UnitTests.csproj" />
    <Project Path="tests/Microsoft.Agents.AI.OpenAI.UnitTests/Microsoft.Agents.AI.OpenAI.UnitTests.csproj" />
    <Project Path="tests/Microsoft.Agents.AI.Purview.UnitTests/Microsoft.Agents.AI.Purview.UnitTests.csproj" />
@@ -1151,6 +1151,25 @@ internal static class AgentsSamples
            SkipReason = "Runs as an MCP stdio server that does not exit on its own.",
        },

+        new SampleDefinition
+        {
+            Name = "Agent_MCP_LongRunningTask_Client",
+            ProjectPath = "samples/02-agents/ModelContextProtocol/Agent_MCP_LongRunningTask_Client",
+            RequiredEnvironmentVariables = ["AZURE_OPENAI_ENDPOINT"],
+            OptionalEnvironmentVariables = ["AZURE_OPENAI_DEPLOYMENT_NAME"],
+            MustContain =
+            [
+                "=== Transparent long-running MCP task (RunAsync) ===",
+                "=== Transparent long-running MCP task (RunStreamingAsync) ===",
+            ],
+            ExpectedOutputDescription =
+            [
+                "The output should show an agent analyzing a dataset named 'sales-2025-q1' and producing a summary mentioning rows, revenue, anomalies, or outliers.",
+                "The output should contain both a non-streaming response (after RunAsync) and a streaming response (after RunStreamingAsync) for the same analysis question.",
+                "The output should not contain error messages or stack traces.",
+            ],
+        },
+
        new SampleDefinition
        {
            Name = "AGUI_Step01_GettingStarted_Client",
@@ -0,0 +1,25 @@
+<Project Sdk="Microsoft.NET.Sdk">
+
+  <PropertyGroup>
+    <OutputType>Exe</OutputType>
+    <TargetFrameworks>net10.0</TargetFrameworks>
+
+    <Nullable>enable</Nullable>
+    <ImplicitUsings>enable</ImplicitUsings>
+    <NoWarn>$(NoWarn);MAAI001;MEAI001;MCPEXP001</NoWarn>
+  </PropertyGroup>
+
+  <ItemGroup>
+    <PackageReference Include="Azure.AI.OpenAI" />
+    <PackageReference Include="Azure.Identity" />
+    <PackageReference Include="Microsoft.Extensions.AI.OpenAI" />
+    <PackageReference Include="Microsoft.Extensions.Hosting" />
+    <PackageReference Include="ModelContextProtocol" />
+  </ItemGroup>
+
+  <ItemGroup>
+    <ProjectReference Include="..\..\..\..\src\Microsoft.Agents.AI.Mcp\Microsoft.Agents.AI.Mcp.csproj" />
+    <ProjectReference Include="..\..\..\..\src\Microsoft.Agents.AI.OpenAI\Microsoft.Agents.AI.OpenAI.csproj" />
+  </ItemGroup>
+
+</Project>
@@ -0,0 +1,145 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+// This sample demonstrates the Microsoft Agent Framework's MCP long-running task support.
+//
+// A small MCP server (hosted in this same executable when launched with "--server") exposes
+// a single task-supporting tool "AnalyzeDataset" that simulates ~15 seconds of work. The
+// client (default mode) connects to it over stdio via Microsoft.Agents.AI.Mcp's
+// McpClientTaskExtensions.ListAgentToolsWithTaskSupportAsync, hands the wrapped tools to a
+// ChatClientAgent, and exercises both invocation styles:
+//   * RunAsync          — blocks until the agent's final response is ready.
+//   * RunStreamingAsync — yields response updates as the model produces them; the model
+//                         still waits for the tool's terminal result before it can begin
+//                         producing the final answer, so the perceived "pause" reflects
+//                         tool execution time, not stream-channel latency.
+//
+// In both cases the wrapper transparently:
+//   1. Calls tools/call with task augmentation (CallToolAsTaskAsync)
+//   2. Polls tasks/get until terminal (PollTaskUntilCompleteAsync)
+//   3. Fetches tasks/result and returns the final result to the function-calling loop
+//
+// No application-level loop or continuation tokens are required in either mode.
+
+using System.ComponentModel;
+using Azure.AI.OpenAI;
+using Azure.Identity;
+using Microsoft.Agents.AI;
+using Microsoft.Agents.AI.Mcp;
+using Microsoft.Extensions.AI;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Hosting;
+using Microsoft.Extensions.Logging;
+using ModelContextProtocol;
+using ModelContextProtocol.Client;
+using ModelContextProtocol.Protocol;
+using ModelContextProtocol.Server;
+using OpenAI.Chat;
+
+if (args.Length > 0 && args[0] == "--server")
+{
+    await RunMcpServerAsync();
+    return;
+}
+
+var endpoint = Environment.GetEnvironmentVariable("AZURE_OPENAI_ENDPOINT") ?? throw new InvalidOperationException("AZURE_OPENAI_ENDPOINT is not set.");
+var deploymentName = Environment.GetEnvironmentVariable("AZURE_OPENAI_DEPLOYMENT_NAME") ?? "gpt-5.4-mini";
+
+// Launch this same assembly as a stdio MCP server in a child process.
+var thisAssemblyPath = typeof(Program).Assembly.Location;
+await using var mcpClient = await McpClient.CreateAsync(new StdioClientTransport(new()
+{
+    Name = "DatasetAnalyzer",
+    Command = "dotnet",
+    Arguments = [thisAssemblyPath, "--server"],
+}));
+
+// Wrap each MCP tool with task-aware behavior. The wrapper inspects the server's
+// execution.taskSupport on each tool and, when it is Required, drives the task lifecycle
+// transparently within the agent's tool loop. Tools that don't require task semantics are
+// returned as-is and invoked inline.
+var taskOptions = new McpTaskOptions
+{
+    DefaultTimeToLive = TimeSpan.FromMinutes(5),
+};
+var mcpTools = await mcpClient.ListAgentToolsWithTaskSupportAsync(taskOptions);
+
+// WARNING: DefaultAzureCredential is convenient for development but requires careful consideration in production.
+// In production, consider using a specific credential (e.g., ManagedIdentityCredential) to avoid
+// latency issues, unintended credential probing, and potential security risks from fallback mechanisms.
+AIAgent agent = new AzureOpenAIClient(
+    new Uri(endpoint),
+    new DefaultAzureCredential())
+     .GetChatClient(deploymentName)
+     .AsAIAgent(
+        instructions: "You answer data-analysis questions by invoking the available tools. Always invoke a tool when one matches the request.",
+        tools: [.. mcpTools.Cast<AITool>()]);
+
+const string Prompt = "Analyze the dataset named 'sales-2025-q1' and summarize the findings.";
+
+Console.WriteLine("=== Transparent long-running MCP task (RunAsync) ===");
+Console.WriteLine("Asking the agent to analyze a dataset; the tool takes ~15s to complete.");
+Console.WriteLine("RunAsync blocks while the wrapper polls the task to completion.");
+Console.WriteLine();
+
+var stopwatch = System.Diagnostics.Stopwatch.StartNew();
+var response = await agent.RunAsync(Prompt);
+stopwatch.Stop();
+
+Console.WriteLine($"Agent response (after {stopwatch.Elapsed.TotalSeconds:F1}s):");
+Console.WriteLine(response.Text);
+
+Console.WriteLine();
+Console.WriteLine("=== Transparent long-running MCP task (RunStreamingAsync) ===");
+Console.WriteLine("Same request via the streaming API. Updates only begin to arrive after the");
+Console.WriteLine("tool's task reaches the Completed state, since the model needs the tool result");
+Console.WriteLine("before it can produce its final answer.");
+Console.WriteLine();
+
+stopwatch.Restart();
+await foreach (var update in agent.RunStreamingAsync(Prompt))
+{
+    Console.Write(update.Text);
+}
+stopwatch.Stop();
+
+Console.WriteLine();
+Console.WriteLine($"(Streaming completed after {stopwatch.Elapsed.TotalSeconds:F1}s.)");
+
+// --- Server mode (launched as a child process via --server) ---------------------------------
+static async Task RunMcpServerAsync()
+{
+    var builder = Host.CreateApplicationBuilder();
+
+    // Critical for stdio transport: any provider that writes to stdout will corrupt the
+    // JSON-RPC channel. Clear all providers; the MCP SDK routes its own diagnostics
+    // appropriately.
+    builder.Logging.ClearProviders();
+    builder.Logging.AddConsole(o => o.LogToStandardErrorThreshold = LogLevel.Trace);
+
+    builder.Services.AddMcpServer(o =>
+    {
+        o.TaskStore = new InMemoryMcpTaskStore();
+        o.ServerInfo = new Implementation { Name = "DatasetAnalyzer", Version = "1.0.0" };
+    })
+    .WithStdioServerTransport()
+    .WithTools<DatasetAnalysisTools>();
+
+    await builder.Build().RunAsync();
+}
+
+#pragma warning disable CA1812 // Discovered by MCP SDK via [McpServerToolType] attribute
+[McpServerToolType]
+internal sealed class DatasetAnalysisTools
+#pragma warning restore CA1812
+{
+    [McpServerTool(Name = "AnalyzeDataset", TaskSupport = ToolTaskSupport.Required)]
+    [Description("Analyze a tabular dataset and return summary statistics. This tool simulates a long-running analytic job (~15 seconds).")]
+    public static async Task<string> AnalyzeDatasetAsync(
+        [Description("The dataset identifier, e.g. 'sales-2025-q1'.")] string datasetName,
+        CancellationToken cancellationToken)
+    {
+        await Task.Delay(TimeSpan.FromSeconds(15), cancellationToken).ConfigureAwait(false);
+
+        return $"Findings for '{datasetName}': 12,403 rows; avg revenue $48,712; 3 anomalies detected in week 7; outliers concentrated in EMEA region.";
+    }
+}
@@ -0,0 +1,60 @@
+# Agent with MCP long-running task (transparent polling)
+
+This sample demonstrates Microsoft Agent Framework's MCP long-running task support: an agent invokes an MCP tool whose execution takes too long for a single request/response cycle, and the framework polls it to completion behind the function-calling loop. From the agent's perspective the tool simply returns its result.
+
+## What this sample shows
+
+- Using `McpClient.ListAgentToolsWithTaskSupportAsync(...)` (in `Microsoft.Agents.AI.Mcp`) to wrap MCP tools with task-aware behavior.
+- Configuring `McpTaskOptions.DefaultTimeToLive` to bound the server-side task.
+- Hosting a small MCP server (in this same executable, launched with `--server`) that advertises `execution.taskSupport=required` on a tool that sleeps for ~15 seconds.
+- No application-level polling, continuation tokens, or `AllowBackgroundResponses` flag are required.
+
+The decorator drives the lifecycle internally:
+
+1. `tools/call` augmented with task metadata (`CallToolAsTaskAsync`)
+2. `tasks/get` polled until terminal (`PollTaskUntilCompleteAsync`)
+3. `tasks/result` retrieved (`GetTaskResultAsync`) and returned to the function-calling loop
+
+The sample exercises both invocation styles against the same wrapper:
+
+- `agent.RunAsync(...)` blocks until the tool completes (~15 seconds in this sample) and returns the final response.
+- `agent.RunStreamingAsync(...)` returns immediately and yields `AgentResponseUpdate` chunks as the model emits them; in this scenario the model only begins streaming its answer once the wrapped tool's task reaches the `Completed` state, so the perceived "pause" before tokens arrive reflects tool execution time, not stream-channel latency.
+
+# Prerequisites
+
+- .NET 10 SDK or later
+- Azure OpenAI service endpoint and a chat-completions deployment
+- Azure CLI installed and authenticated (`az login`)
+
+Set the following environment variables:
+
+```powershell
+$env:AZURE_OPENAI_ENDPOINT="https://your-resource.openai.azure.com/"
+$env:AZURE_OPENAI_DEPLOYMENT_NAME="gpt-5.4-mini"  # optional; defaults to gpt-5.4-mini
+```
+
+# Running
+
+```powershell
+cd Agent_MCP_LongRunningTask_Client
+dotnet run
+```
+
+You should see output similar to:
+
+```
+=== Transparent long-running MCP task (RunAsync) ===
+Asking the agent to analyze a dataset; the tool takes ~15s to complete.
+RunAsync blocks while the wrapper polls the task to completion.
+
+Agent response (after 15.4s):
+The 'sales-2025-q1' dataset contains 12,403 rows ...
+
+=== Transparent long-running MCP task (RunStreamingAsync) ===
+Same request via the streaming API. Updates only begin to arrive after the
+tool's task reaches the Completed state, since the model needs the tool result
+before it can produce its final answer.
+
+The 'sales-2025-q1' dataset contains 12,403 rows ...
+(Streaming completed after 15.7s.)
+```
@@ -22,6 +22,7 @@ Before you begin, ensure you have the following prerequisites:
 |[Agent with MCP server tools](./Agent_MCP_Server/)|This sample demonstrates how to use MCP server tools with a simple agent|
 |[Agent with MCP server tools and authorization](./Agent_MCP_Server_Auth/)|This sample demonstrates how to use MCP Server tools from a protected MCP server with a simple agent|
 |[Responses Agent with Hosted MCP tool](./ResponseAgent_Hosted_MCP/)|This sample demonstrates how to use the Hosted MCP tool with the Responses Service, where the service invokes any MCP tools directly|
+|[Agent with long-running MCP task (transparent polling)](./Agent_MCP_LongRunningTask_Client/)|This sample demonstrates how an agent transparently drives a long-running MCP task (SEP-2663) to completion. The wrapper polls the task internally on both `RunAsync` and `RunStreamingAsync` invocations.|

 ## Running the samples from the console

@@ -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.
+        }
+    }
+}
@@ -0,0 +1,15 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+
+namespace Microsoft.Agents.AI.Mcp.UnitTests;
+
+/// <summary>
+/// Minimal empty <see cref="IServiceProvider"/> for in-memory fixtures that don't use DI.
+/// </summary>
+internal sealed class EmptyServiceProvider : IServiceProvider
+{
+    public static EmptyServiceProvider Instance { get; } = new();
+
+    public object? GetService(Type serviceType) => null;
+}
@@ -0,0 +1,127 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.IO;
+using System.IO.Pipelines;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.Extensions.Logging.Abstractions;
+using ModelContextProtocol;
+using ModelContextProtocol.Client;
+using ModelContextProtocol.Protocol;
+using ModelContextProtocol.Server;
+
+namespace Microsoft.Agents.AI.Mcp.UnitTests;
+
+/// <summary>
+/// In-process MCP server fixture that pairs a <see cref="McpServer"/> and a <see cref="McpClient"/>
+/// over duplex <see cref="Pipe"/>-backed streams so unit tests can exercise the
+/// real task-augmentation protocol without spawning a child process or opening a socket.
+/// </summary>
+internal sealed class InMemoryMcpServerFixture : IAsyncDisposable
+{
+    private readonly McpServer _server;
+    private readonly Task _serverLoop;
+    private readonly CancellationTokenSource _cts;
+
+    public McpClient Client { get; }
+
+    private InMemoryMcpServerFixture(McpServer server, McpClient client, Task serverLoop, CancellationTokenSource cts)
+    {
+        this._server = server;
+        this.Client = client;
+        this._serverLoop = serverLoop;
+        this._cts = cts;
+    }
+
+    public static async Task<InMemoryMcpServerFixture> CreateAsync(
+        McpServerPrimitiveCollection<McpServerTool> tools,
+        CancellationToken cancellationToken = default)
+    {
+        Pipe clientToServer = new();
+        Pipe serverToClient = new();
+
+        // Stream conventions:
+        //   StreamClientTransport(serverInput, serverOutput, ...): serverInput is what the client
+        //   WRITES to (server reads it); serverOutput is what the client READS from (server writes it).
+        //   StreamServerTransport(input, output, ...): input is what the server READS from; output
+        //   is what the server WRITES to.
+        Stream clientWriteStream = clientToServer.Writer.AsStream();
+        Stream clientReadStream = serverToClient.Reader.AsStream();
+        Stream serverReadStream = clientToServer.Reader.AsStream();
+        Stream serverWriteStream = serverToClient.Writer.AsStream();
+
+        StreamServerTransport serverTransport = new(
+            serverReadStream,
+            serverWriteStream,
+            "test-server",
+            NullLoggerFactory.Instance);
+
+        McpServerOptions serverOptions = new()
+        {
+            ServerInfo = new Implementation { Name = "test-server", Version = "1.0.0" },
+            TaskStore = new InMemoryMcpTaskStore(),
+            ToolCollection = tools,
+        };
+
+        McpServer server = McpServer.Create(
+            serverTransport,
+            serverOptions,
+            NullLoggerFactory.Instance,
+            EmptyServiceProvider.Instance);
+
+        CancellationTokenSource cts = new();
+        Task serverLoop = Task.Run(() => server.RunAsync(cts.Token), cts.Token);
+
+        StreamClientTransport clientTransport = new(
+            clientWriteStream,
+            clientReadStream,
+            NullLoggerFactory.Instance);
+
+        McpClient client = await McpClient.CreateAsync(
+            clientTransport,
+            clientOptions: null,
+            NullLoggerFactory.Instance,
+            cancellationToken).ConfigureAwait(false);
+
+        return new InMemoryMcpServerFixture(server, client, serverLoop, cts);
+    }
+
+    public async ValueTask DisposeAsync()
+    {
+        try
+        {
+            await this.Client.DisposeAsync().ConfigureAwait(false);
+        }
+        catch
+        {
+            // Best effort.
+        }
+
+        this._cts.Cancel();
+
+        try
+        {
+            await this._serverLoop.ConfigureAwait(false);
+        }
+        catch (OperationCanceledException)
+        {
+            // Expected.
+        }
+        catch
+        {
+            // Best effort.
+        }
+
+        try
+        {
+            await this._server.DisposeAsync().ConfigureAwait(false);
+        }
+        catch
+        {
+            // Best effort.
+        }
+
+        this._cts.Dispose();
+    }
+}
@@ -0,0 +1,55 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Linq;
+using System.Threading.Tasks;
+using FluentAssertions;
+using Microsoft.Extensions.AI;
+using ModelContextProtocol.Protocol;
+using ModelContextProtocol.Server;
+
+namespace Microsoft.Agents.AI.Mcp.UnitTests;
+
+public class ListAgentToolsWithTaskSupportTests
+{
+    [Fact]
+    public async Task ListAgentToolsWithTaskSupport_WrapsTaskCapableTools_LeavesOthersAsIsAsync()
+    {
+        // Arrange
+        McpServerPrimitiveCollection<McpServerTool> tools = [
+            TestTools.Create("opt", ToolTaskSupport.Optional, () => "opt-result"),
+            TestTools.Create("req", ToolTaskSupport.Required, () => "req-result"),
+            TestTools.Create("forb", ToolTaskSupport.Forbidden, () => "forb-result"),
+            TestTools.Create("none", taskSupport: null, () => "none-result"),
+        ];
+        await using InMemoryMcpServerFixture fixture = await InMemoryMcpServerFixture.CreateAsync(tools);
+
+        // Act
+        var result = await fixture.Client.ListAgentToolsWithTaskSupportAsync();
+
+        // Assert
+        result.Should().HaveCount(4);
+        AIFunction opt = result.Single(f => f.Name == "opt");
+        AIFunction req = result.Single(f => f.Name == "req");
+        AIFunction forb = result.Single(f => f.Name == "forb");
+        AIFunction none = result.Single(f => f.Name == "none");
+
+        req.Should().BeOfType<TaskAwareMcpClientAIFunction>("Required tools must be wrapped");
+        opt.Should().NotBeOfType<TaskAwareMcpClientAIFunction>("Optional tools must not be wrapped; inline invocation is preserved by default");
+        forb.Should().NotBeOfType<TaskAwareMcpClientAIFunction>("Forbidden tools must not be wrapped");
+        none.Should().NotBeOfType<TaskAwareMcpClientAIFunction>("Tools without execution metadata must not be wrapped");
+    }
+
+    [Fact]
+    public async Task ListAgentToolsWithTaskSupport_ThrowsOnNullClientAsync()
+    {
+        // Arrange
+        ModelContextProtocol.Client.McpClient client = null!;
+
+        // Act
+        Func<Task> act = async () => await client.ListAgentToolsWithTaskSupportAsync();
+
+        // Assert
+        await act.Should().ThrowAsync<ArgumentNullException>();
+    }
+}
@@ -0,0 +1,19 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using FluentAssertions;
+
+namespace Microsoft.Agents.AI.Mcp.UnitTests;
+
+public class McpTaskOptionsTests
+{
+    [Fact]
+    public void Defaults_AreSane()
+    {
+        // Act
+        McpTaskOptions options = new();
+
+        // Assert
+        options.DefaultTimeToLive.Should().BeNull();
+        options.CancelRemoteTaskOnLocalCancellation.Should().BeTrue();
+    }
+}
@@ -0,0 +1,18 @@
+<Project Sdk="Microsoft.NET.Sdk">
+
+  <PropertyGroup>
+    <TargetFrameworks>$(TargetFrameworksCore)</TargetFrameworks>
+    <NoWarn>$(NoWarn);MCPEXP001</NoWarn>
+  </PropertyGroup>
+
+  <ItemGroup>
+    <PackageReference Include="FluentAssertions" />
+    <PackageReference Include="Microsoft.Extensions.Logging.Abstractions" />
+    <PackageReference Include="ModelContextProtocol" />
+  </ItemGroup>
+
+  <ItemGroup>
+    <ProjectReference Include="..\..\src\Microsoft.Agents.AI.Mcp\Microsoft.Agents.AI.Mcp.csproj" />
+  </ItemGroup>
+
+</Project>
@@ -0,0 +1,159 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Linq;
+using System.Text.Json;
+using System.Threading;
+using System.Threading.Tasks;
+using FluentAssertions;
+using Microsoft.Extensions.AI;
+using ModelContextProtocol.Protocol;
+using ModelContextProtocol.Server;
+
+namespace Microsoft.Agents.AI.Mcp.UnitTests;
+
+public class TaskAwareMcpClientAIFunctionTests
+{
+    [Fact]
+    public async Task InvokeAsync_RequiredTool_HappyPath_ReturnsResultAsync()
+    {
+        // Arrange
+        McpServerPrimitiveCollection<McpServerTool> tools = [
+            TestTools.Create("req", ToolTaskSupport.Required, () => "required-result"),
+        ];
+        await using InMemoryMcpServerFixture fixture = await InMemoryMcpServerFixture.CreateAsync(tools);
+        var result = await fixture.Client.ListAgentToolsWithTaskSupportAsync();
+        AIFunction req = result.Single(f => f.Name == "req");
+        req.Should().BeOfType<TaskAwareMcpClientAIFunction>();
+
+        // Act
+        object? invokeResult = await req.InvokeAsync(arguments: null, CancellationToken.None);
+
+        // Assert
+        JsonElement payload = invokeResult.Should().BeOfType<JsonElement>().Subject;
+        ExtractTextContent(payload).Should().Be("required-result");
+    }
+
+    [Fact]
+    public async Task InvokeAsync_PropagatesDefaultTimeToLiveAsync()
+    {
+        // Arrange — capture the request meta on the server so we can assert TTL flowed through.
+        TimeSpan? observedTtl = null;
+        McpServerTool tool = McpServerTool.Create(
+            (RequestContext<CallToolRequestParams> ctx) =>
+            {
+                observedTtl = ctx.Params?.Task?.TimeToLive;
+                return "ok";
+            },
+            new McpServerToolCreateOptions
+            {
+                Name = "ttl-tool",
+                Description = "Echoes the requested TTL.",
+                Execution = new ToolExecution { TaskSupport = ToolTaskSupport.Required },
+            });
+        McpServerPrimitiveCollection<McpServerTool> tools = [tool];
+
+        await using InMemoryMcpServerFixture fixture = await InMemoryMcpServerFixture.CreateAsync(tools);
+
+        TimeSpan requestedTtl = TimeSpan.FromMinutes(7);
+        var result = await fixture.Client.ListAgentToolsWithTaskSupportAsync(new McpTaskOptions { DefaultTimeToLive = requestedTtl });
+        AIFunction wrapped = result.Single();
+
+        // Act
+        _ = await wrapped.InvokeAsync(arguments: null, CancellationToken.None);
+
+        // Assert
+        observedTtl.Should().Be(requestedTtl);
+    }
+
+    [Fact]
+    public async Task InvokeAsync_RespectsCancellationAsync()
+    {
+        // Arrange — a tool that never completes until it's cancelled.
+        var serverCancelled = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
+        McpServerTool tool = McpServerTool.Create(
+            async (CancellationToken ct) =>
+            {
+                try
+                {
+                    await Task.Delay(Timeout.Infinite, ct);
+                }
+                catch (OperationCanceledException)
+                {
+                    serverCancelled.TrySetResult(true);
+                    throw;
+                }
+
+                return "should-not-complete";
+            },
+            new McpServerToolCreateOptions
+            {
+                Name = "blocking",
+                Description = "Blocks indefinitely until cancelled.",
+                Execution = new ToolExecution { TaskSupport = ToolTaskSupport.Required },
+            });
+        McpServerPrimitiveCollection<McpServerTool> tools = [tool];
+
+        await using InMemoryMcpServerFixture fixture = await InMemoryMcpServerFixture.CreateAsync(tools);
+        var result = await fixture.Client.ListAgentToolsWithTaskSupportAsync();
+        AIFunction wrapped = result.Single();
+
+        using CancellationTokenSource cts = new();
+
+        // Act — start the invocation, cancel after a brief delay.
+        Task<object?> invocation = wrapped.InvokeAsync(arguments: null, cts.Token).AsTask();
+        await Task.Delay(200);
+        cts.Cancel();
+
+        // Assert — wrapper observes cancellation and signals server-side cancellation.
+        Func<Task> awaitInvocation = async () => await invocation;
+        await awaitInvocation.Should().ThrowAsync<OperationCanceledException>();
+
+        // Server-side handler should have observed cancellation as a result of the wrapper's
+        // tasks/cancel call (best-effort wait — give the server-loop a few seconds).
+        Task observedTask = serverCancelled.Task;
+        Task completed = await Task.WhenAny(observedTask, Task.Delay(TimeSpan.FromSeconds(5)));
+        completed.Should().BeSameAs(observedTask, "the wrapper should have issued tasks/cancel");
+    }
+
+    [Fact]
+    public async Task InvokeAsync_FailedTask_ThrowsInvalidOperationAsync()
+    {
+        // Arrange — a tool whose handler throws, which the server surfaces as a Failed task.
+        McpServerTool tool = McpServerTool.Create(
+            (Func<string>)(() => throw new InvalidOperationException("simulated tool failure")),
+            new McpServerToolCreateOptions
+            {
+                Name = "boom",
+                Description = "Throws unconditionally.",
+                Execution = new ToolExecution { TaskSupport = ToolTaskSupport.Required },
+            });
+        McpServerPrimitiveCollection<McpServerTool> tools = [tool];
+
+        await using InMemoryMcpServerFixture fixture = await InMemoryMcpServerFixture.CreateAsync(tools);
+        var result = await fixture.Client.ListAgentToolsWithTaskSupportAsync();
+        AIFunction wrapped = result.Single();
+
+        // Act
+        Func<Task> act = async () => await wrapped.InvokeAsync(arguments: null, CancellationToken.None);
+
+        // Assert — Phase 1 surfaces non-Completed terminal states as InvalidOperationException
+        // carrying the server's StatusMessage. (See PollAndRetrieveResultAsync.)
+        await act.Should().ThrowAsync<Exception>().Where(ex =>
+            ex is InvalidOperationException
+            || ex.GetType().FullName == "ModelContextProtocol.McpException");
+    }
+
+    /// <summary>
+    /// Extracts the first text-content block from a serialized <c>CallToolResult</c>
+    /// (the JSON shape returned by the wrapper and by <c>McpClientTool.InvokeAsync</c>).
+    /// </summary>
+    private static string ExtractTextContent(JsonElement payload)
+    {
+        payload.ValueKind.Should().Be(JsonValueKind.Object);
+        JsonElement content = payload.GetProperty("content");
+        content.ValueKind.Should().Be(JsonValueKind.Array);
+        JsonElement firstBlock = content.EnumerateArray().First();
+        return firstBlock.GetProperty("text").GetString()!;
+    }
+}
@@ -0,0 +1,30 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using ModelContextProtocol.Protocol;
+using ModelContextProtocol.Server;
+
+namespace Microsoft.Agents.AI.Mcp.UnitTests;
+
+/// <summary>
+/// Helpers to create <see cref="McpServerTool"/> instances with a specific
+/// <see cref="ToolTaskSupport"/> level for in-memory fixtures.
+/// </summary>
+internal static class TestTools
+{
+    public static McpServerTool Create(string name, ToolTaskSupport? taskSupport, Delegate handler)
+    {
+        McpServerToolCreateOptions options = new()
+        {
+            Name = name,
+            Description = $"Test tool {name}.",
+        };
+
+        if (taskSupport is ToolTaskSupport ts)
+        {
+            options.Execution = new ToolExecution { TaskSupport = ts };
+        }
+
+        return McpServerTool.Create(handler, options);
+    }
+}