Addressed PR comments

This commit is contained in:
Peter Ibekwe
2026-05-22 09:46:29 -07:00
Unverified
parent 42b2954e07
commit fed973b6b3
6 changed files with 19 additions and 38 deletions
@@ -54,8 +54,9 @@ await using var mcpClient = await McpClient.CreateAsync(new StdioClientTransport
}));
// Wrap each MCP tool with task-aware behavior. The wrapper inspects the server's
// execution.taskSupport on each tool and, if Optional/Required, drives the task lifecycle
// transparently within the agent's tool loop.
// 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),
@@ -131,7 +132,7 @@ static async Task RunMcpServerAsync()
internal sealed class DatasetAnalysisTools
#pragma warning restore CA1812
{
[McpServerTool(Name = "AnalyzeDataset", TaskSupport = ToolTaskSupport.Optional)]
[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,
@@ -6,7 +6,7 @@ This sample demonstrates Microsoft Agent Framework's MCP long-running task suppo
- 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=optional` on a tool that sleeps for ~15 seconds.
- 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:
@@ -18,9 +18,11 @@ 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.Optional"/> or
/// <see cref="ToolTaskSupport.Required"/> are wrapped with task-aware behavior so an agent
/// can transparently drive long-running invocations. All other tools are returned as-is.
/// <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">
@@ -44,7 +46,7 @@ public static class McpClientTaskExtensions
for (int i = 0; i < tools.Count; i++)
{
ToolTaskSupport? taskSupport = tools[i].ProtocolTool.Execution?.TaskSupport;
if (taskSupport is ToolTaskSupport.Optional or ToolTaskSupport.Required)
if (taskSupport is ToolTaskSupport.Required)
{
result[i] = new TaskAwareMcpClientAIFunction(client, tools[i], effectiveOptions);
}
@@ -31,13 +31,12 @@ namespace Microsoft.Agents.AI.Mcp;
/// </para>
/// <para>
/// This wrapper is intended to be applied only to tools whose
/// <see cref="ToolExecution.TaskSupport"/> is <see cref="ToolTaskSupport.Optional"/> or
/// <see cref="ToolTaskSupport.Required"/> (selected by
/// <see cref="McpClientTaskExtensions.ListAgentToolsWithTaskSupportAsync"/>). As a defensive
/// fallback, if the server still rejects the task-augmented call with JSON-RPC error
/// <c>-32601</c> (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"/>.
/// <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
@@ -96,7 +95,7 @@ internal sealed class TaskAwareMcpClientAIFunction : AIFunction
options: null,
cancellationToken: cancellationToken).ConfigureAwait(false);
}
catch (McpProtocolException ex) when ((int)ex.ErrorCode == -32601)
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
@@ -34,8 +34,8 @@ public class ListAgentToolsWithTaskSupportTests
AIFunction forb = result.Single(f => f.Name == "forb");
AIFunction none = result.Single(f => f.Name == "none");
opt.Should().BeOfType<TaskAwareMcpClientAIFunction>("Optional tools must be wrapped");
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");
}
@@ -14,27 +14,6 @@ namespace Microsoft.Agents.AI.Mcp.UnitTests;
public class TaskAwareMcpClientAIFunctionTests
{
[Fact]
public async Task InvokeAsync_OptionalTool_HappyPath_ReturnsResultAsync()
{
// Arrange
McpServerPrimitiveCollection<McpServerTool> tools = [
TestTools.Create("opt", ToolTaskSupport.Optional, () => "optional-result"),
];
await using InMemoryMcpServerFixture fixture = await InMemoryMcpServerFixture.CreateAsync(tools);
var result = await fixture.Client.ListAgentToolsWithTaskSupportAsync();
AIFunction opt = result.Single(f => f.Name == "opt");
opt.Should().BeOfType<TaskAwareMcpClientAIFunction>();
// Act
object? invokeResult = await opt.InvokeAsync(arguments: null, CancellationToken.None);
// Assert — wrapper returns a JsonElement containing the serialized CallToolResult
// (same wire shape as McpClientTool.InvokeAsync).
JsonElement payload = invokeResult.Should().BeOfType<JsonElement>().Subject;
ExtractTextContent(payload).Should().Be("optional-result");
}
[Fact]
public async Task InvokeAsync_RequiredTool_HappyPath_ReturnsResultAsync()
{