mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
.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
This commit is contained in:
Peter Ibekwe
@@ -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();
|
||||
}
|
||||
}
|
||||
+18
@@ -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);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user