Implement MCP task support (SEP-1686) (#1170)
Co-authored-by: Stephen Halter <halter73@gmail.com>
This commit is contained in:
Eirik Tsarpalis
GitHub
parent
81664c3fce
commit
10e4f5b20f
1 changed files with 11580 additions and 46 deletions
|
|
@ -0,0 +1,478 @@
|
|||
using Microsoft.Extensions.DependencyInjection;
|
||||
using Microsoft.Extensions.Logging;
|
||||
using ModelContextProtocol.Client;
|
||||
using ModelContextProtocol.Protocol;
|
||||
using ModelContextProtocol.Server;
|
||||
using ModelContextProtocol.Tests.Utils;
|
||||
using System.IO.Pipelines;
|
||||
|
||||
namespace ModelContextProtocol.Tests.Server;
|
||||
|
||||
/// <summary>
|
||||
/// Tests for automatic InputRequired status tracking when server-to-client
|
||||
/// requests (SampleAsync, ElicitAsync) are made during task-augmented tool execution.
|
||||
/// </summary>
|
||||
public class AutomaticInputRequiredStatusTests : LoggedTest
|
||||
{
|
||||
public AutomaticInputRequiredStatusTests(ITestOutputHelper testOutputHelper)
|
||||
: base(testOutputHelper)
|
||||
{
|
||||
}
|
||||
|
||||
#pragma warning disable MCPEXP001 // Tasks feature is experimental
|
||||
|
||||
[Fact]
|
||||
public async Task TaskStatus_TransitionsToInputRequired_DuringSampleAsync()
|
||||
{
|
||||
// Arrange
|
||||
var taskStore = new InMemoryMcpTaskStore();
|
||||
var statusesDuringSampling = new List<McpTaskStatus>();
|
||||
var samplingRequestReceived = new TaskCompletionSource<bool>();
|
||||
var continueSampling = new TaskCompletionSource<bool>();
|
||||
|
||||
await using var fixture = new InputRequiredTestFixture(
|
||||
LoggerFactory,
|
||||
configureServer: (services, builder) =>
|
||||
{
|
||||
services.AddSingleton<IMcpTaskStore>(taskStore);
|
||||
services.Configure<McpServerOptions>(options =>
|
||||
{
|
||||
options.TaskStore = taskStore;
|
||||
options.SendTaskStatusNotifications = true; // Enable notifications
|
||||
});
|
||||
|
||||
// Tool that calls SampleAsync during execution
|
||||
builder.WithTools([McpServerTool.Create(
|
||||
async (string prompt, McpServer server, CancellationToken ct) =>
|
||||
{
|
||||
// Call SampleAsync - this should trigger InputRequired status
|
||||
var result = await server.SampleAsync(new CreateMessageRequestParams
|
||||
{
|
||||
Messages = [new SamplingMessage
|
||||
{
|
||||
Role = Role.User,
|
||||
Content = [new TextContentBlock { Text = prompt }]
|
||||
}],
|
||||
MaxTokens = 100
|
||||
}, ct);
|
||||
|
||||
var textContent = result.Content.OfType<TextContentBlock>().FirstOrDefault();
|
||||
return textContent?.Text ?? "No response";
|
||||
},
|
||||
new McpServerToolCreateOptions
|
||||
{
|
||||
Name = "sampling-tool",
|
||||
Description = "A tool that uses sampling"
|
||||
})]);
|
||||
},
|
||||
configureClient: clientOptions =>
|
||||
{
|
||||
clientOptions.Handlers = new McpClientHandlers
|
||||
{
|
||||
SamplingHandler = async (request, progress, ct) =>
|
||||
{
|
||||
// Signal that we received the sampling request
|
||||
samplingRequestReceived.TrySetResult(true);
|
||||
|
||||
// Wait for permission to continue (so we can check status)
|
||||
await continueSampling.Task.WaitAsync(ct);
|
||||
|
||||
return new CreateMessageResult
|
||||
{
|
||||
Content = [new TextContentBlock { Text = "Sampled response" }],
|
||||
Model = "test-model"
|
||||
};
|
||||
}
|
||||
};
|
||||
});
|
||||
|
||||
// Act - Call the tool as a task
|
||||
var mcpTask = await fixture.Client.CallToolAsTaskAsync(
|
||||
"sampling-tool",
|
||||
arguments: new Dictionary<string, object?> { ["prompt"] = "Hello" },
|
||||
taskMetadata: new McpTaskMetadata(),
|
||||
progress: null,
|
||||
cancellationToken: TestContext.Current.CancellationToken);
|
||||
|
||||
// Wait for the sampling request to be received by the client
|
||||
await samplingRequestReceived.Task.WaitAsync(TimeSpan.FromSeconds(5), TestContext.Current.CancellationToken);
|
||||
|
||||
// Check the task status while sampling is in progress
|
||||
var statusDuringSampling = await taskStore.GetTaskAsync(
|
||||
mcpTask.TaskId,
|
||||
cancellationToken: TestContext.Current.CancellationToken);
|
||||
|
||||
if (statusDuringSampling is not null)
|
||||
{
|
||||
statusesDuringSampling.Add(statusDuringSampling.Status);
|
||||
}
|
||||
|
||||
// Allow sampling to complete
|
||||
continueSampling.TrySetResult(true);
|
||||
|
||||
// Wait for task to complete
|
||||
McpTask? finalStatus = null;
|
||||
int maxAttempts = 50;
|
||||
do
|
||||
{
|
||||
await Task.Delay(100, TestContext.Current.CancellationToken);
|
||||
finalStatus = await taskStore.GetTaskAsync(mcpTask.TaskId, cancellationToken: TestContext.Current.CancellationToken);
|
||||
maxAttempts--;
|
||||
}
|
||||
while (finalStatus?.Status == McpTaskStatus.Working && maxAttempts > 0);
|
||||
|
||||
// Assert - Status should have been InputRequired during sampling
|
||||
Assert.Contains(McpTaskStatus.InputRequired, statusesDuringSampling);
|
||||
|
||||
// Final status should be Completed
|
||||
Assert.NotNull(finalStatus);
|
||||
Assert.Equal(McpTaskStatus.Completed, finalStatus.Status);
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task TaskStatus_TransitionsToInputRequired_DuringElicitAsync()
|
||||
{
|
||||
// Arrange
|
||||
var taskStore = new InMemoryMcpTaskStore();
|
||||
var statusesDuringElicitation = new List<McpTaskStatus>();
|
||||
var elicitationRequestReceived = new TaskCompletionSource<bool>();
|
||||
var continueElicitation = new TaskCompletionSource<bool>();
|
||||
|
||||
await using var fixture = new InputRequiredTestFixture(
|
||||
LoggerFactory,
|
||||
configureServer: (services, builder) =>
|
||||
{
|
||||
services.AddSingleton<IMcpTaskStore>(taskStore);
|
||||
services.Configure<McpServerOptions>(options =>
|
||||
{
|
||||
options.TaskStore = taskStore;
|
||||
options.SendTaskStatusNotifications = true;
|
||||
});
|
||||
|
||||
// Tool that calls ElicitAsync during execution
|
||||
builder.WithTools([McpServerTool.Create(
|
||||
async (string message, McpServer server, CancellationToken ct) =>
|
||||
{
|
||||
// Call ElicitAsync - this should trigger InputRequired status
|
||||
var result = await server.ElicitAsync(new ElicitRequestParams
|
||||
{
|
||||
Message = message,
|
||||
RequestedSchema = new()
|
||||
}, ct);
|
||||
|
||||
return result.Action == "confirm" ? "Confirmed" : "Declined";
|
||||
},
|
||||
new McpServerToolCreateOptions
|
||||
{
|
||||
Name = "elicitation-tool",
|
||||
Description = "A tool that uses elicitation"
|
||||
})]);
|
||||
},
|
||||
configureClient: clientOptions =>
|
||||
{
|
||||
clientOptions.Handlers = new McpClientHandlers
|
||||
{
|
||||
ElicitationHandler = async (request, ct) =>
|
||||
{
|
||||
// Signal that we received the elicitation request
|
||||
elicitationRequestReceived.TrySetResult(true);
|
||||
|
||||
// Wait for permission to continue
|
||||
await continueElicitation.Task.WaitAsync(ct);
|
||||
|
||||
return new ElicitResult { Action = "confirm" };
|
||||
}
|
||||
};
|
||||
});
|
||||
|
||||
// Act - Call the tool as a task
|
||||
var mcpTask = await fixture.Client.CallToolAsTaskAsync(
|
||||
"elicitation-tool",
|
||||
arguments: new Dictionary<string, object?> { ["message"] = "Please confirm" },
|
||||
taskMetadata: new McpTaskMetadata(),
|
||||
progress: null,
|
||||
cancellationToken: TestContext.Current.CancellationToken);
|
||||
|
||||
// Wait for the elicitation request to be received
|
||||
await elicitationRequestReceived.Task.WaitAsync(TimeSpan.FromSeconds(5), TestContext.Current.CancellationToken);
|
||||
|
||||
// Check the task status while elicitation is in progress
|
||||
var statusDuringElicitation = await taskStore.GetTaskAsync(
|
||||
mcpTask.TaskId,
|
||||
cancellationToken: TestContext.Current.CancellationToken);
|
||||
|
||||
if (statusDuringElicitation is not null)
|
||||
{
|
||||
statusesDuringElicitation.Add(statusDuringElicitation.Status);
|
||||
}
|
||||
|
||||
// Allow elicitation to complete
|
||||
continueElicitation.TrySetResult(true);
|
||||
|
||||
// Wait for task to complete
|
||||
McpTask? finalStatus = null;
|
||||
int maxAttempts = 50;
|
||||
do
|
||||
{
|
||||
await Task.Delay(100, TestContext.Current.CancellationToken);
|
||||
finalStatus = await taskStore.GetTaskAsync(mcpTask.TaskId, cancellationToken: TestContext.Current.CancellationToken);
|
||||
maxAttempts--;
|
||||
}
|
||||
while (finalStatus?.Status == McpTaskStatus.Working && maxAttempts > 0);
|
||||
|
||||
// Assert - Status should have been InputRequired during elicitation
|
||||
Assert.Contains(McpTaskStatus.InputRequired, statusesDuringElicitation);
|
||||
|
||||
// Final status should be Completed
|
||||
Assert.NotNull(finalStatus);
|
||||
Assert.Equal(McpTaskStatus.Completed, finalStatus.Status);
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task TaskStatus_ReturnsToWorking_AfterSamplingCompletes()
|
||||
{
|
||||
// Arrange
|
||||
var taskStore = new InMemoryMcpTaskStore();
|
||||
var samplingCompleted = new TaskCompletionSource<bool>();
|
||||
var checkStatusAfterSampling = new TaskCompletionSource<bool>();
|
||||
|
||||
await using var fixture = new InputRequiredTestFixture(
|
||||
LoggerFactory,
|
||||
configureServer: (services, builder) =>
|
||||
{
|
||||
services.AddSingleton<IMcpTaskStore>(taskStore);
|
||||
services.Configure<McpServerOptions>(options =>
|
||||
{
|
||||
options.TaskStore = taskStore;
|
||||
});
|
||||
|
||||
// Tool that calls SampleAsync and then waits
|
||||
builder.WithTools([McpServerTool.Create(
|
||||
async (string prompt, McpServer server, CancellationToken ct) =>
|
||||
{
|
||||
// Call SampleAsync
|
||||
var result = await server.SampleAsync(new CreateMessageRequestParams
|
||||
{
|
||||
Messages = [new SamplingMessage
|
||||
{
|
||||
Role = Role.User,
|
||||
Content = [new TextContentBlock { Text = prompt }]
|
||||
}],
|
||||
MaxTokens = 100
|
||||
}, ct);
|
||||
|
||||
// Signal that sampling completed
|
||||
samplingCompleted.TrySetResult(true);
|
||||
|
||||
// Wait so test can check status
|
||||
await checkStatusAfterSampling.Task.WaitAsync(ct);
|
||||
|
||||
var textContent = result.Content.OfType<TextContentBlock>().FirstOrDefault();
|
||||
return textContent?.Text ?? "No response";
|
||||
},
|
||||
new McpServerToolCreateOptions
|
||||
{
|
||||
Name = "sampling-tool",
|
||||
Description = "A tool that uses sampling"
|
||||
})]);
|
||||
},
|
||||
configureClient: clientOptions =>
|
||||
{
|
||||
clientOptions.Handlers = new McpClientHandlers
|
||||
{
|
||||
SamplingHandler = (request, progress, ct) =>
|
||||
{
|
||||
// Return immediately to let sampling complete
|
||||
return new ValueTask<CreateMessageResult>(new CreateMessageResult
|
||||
{
|
||||
Content = [new TextContentBlock { Text = "Response" }],
|
||||
Model = "test-model"
|
||||
});
|
||||
}
|
||||
};
|
||||
});
|
||||
|
||||
// Act - Call the tool as a task
|
||||
var mcpTask = await fixture.Client.CallToolAsTaskAsync(
|
||||
"sampling-tool",
|
||||
arguments: new Dictionary<string, object?> { ["prompt"] = "Hello" },
|
||||
taskMetadata: new McpTaskMetadata(),
|
||||
progress: null,
|
||||
cancellationToken: TestContext.Current.CancellationToken);
|
||||
|
||||
// Wait for sampling to complete inside the tool
|
||||
await samplingCompleted.Task.WaitAsync(TimeSpan.FromSeconds(5), TestContext.Current.CancellationToken);
|
||||
|
||||
// Small delay to ensure status update is processed
|
||||
await Task.Delay(50, TestContext.Current.CancellationToken);
|
||||
|
||||
// Check status after sampling completed (should be back to Working)
|
||||
var taskAfterSampling = await taskStore.GetTaskAsync(mcpTask.TaskId, cancellationToken: TestContext.Current.CancellationToken);
|
||||
|
||||
// Allow tool to complete
|
||||
checkStatusAfterSampling.TrySetResult(true);
|
||||
|
||||
// Assert - Status should be Working after sampling completes (before tool completes)
|
||||
Assert.NotNull(taskAfterSampling);
|
||||
Assert.Equal(McpTaskStatus.Working, taskAfterSampling.Status);
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task TaskStatus_DoesNotChangeToInputRequired_ForNonTaskExecution()
|
||||
{
|
||||
// Arrange - When a tool is NOT executed as a task, SampleAsync should not change any task status
|
||||
var taskStore = new InMemoryMcpTaskStore();
|
||||
var samplingCompleted = new TaskCompletionSource<bool>();
|
||||
|
||||
await using var fixture = new InputRequiredTestFixture(
|
||||
LoggerFactory,
|
||||
configureServer: (services, builder) =>
|
||||
{
|
||||
services.AddSingleton<IMcpTaskStore>(taskStore);
|
||||
services.Configure<McpServerOptions>(options =>
|
||||
{
|
||||
options.TaskStore = taskStore;
|
||||
});
|
||||
|
||||
// Tool that calls SampleAsync - note it doesn't have TaskSupport.Required so can be called directly
|
||||
builder.WithTools([McpServerTool.Create(
|
||||
async (string prompt, McpServer server, CancellationToken ct) =>
|
||||
{
|
||||
var result = await server.SampleAsync(new CreateMessageRequestParams
|
||||
{
|
||||
Messages = [new SamplingMessage
|
||||
{
|
||||
Role = Role.User,
|
||||
Content = [new TextContentBlock { Text = prompt }]
|
||||
}],
|
||||
MaxTokens = 100
|
||||
}, ct);
|
||||
|
||||
samplingCompleted.TrySetResult(true);
|
||||
var textContent = result.Content.OfType<TextContentBlock>().FirstOrDefault();
|
||||
return textContent?.Text ?? "No response";
|
||||
},
|
||||
new McpServerToolCreateOptions
|
||||
{
|
||||
Name = "sampling-tool",
|
||||
Description = "A tool that uses sampling",
|
||||
Execution = new ToolExecution { TaskSupport = ToolTaskSupport.Optional }
|
||||
})]);
|
||||
},
|
||||
configureClient: clientOptions =>
|
||||
{
|
||||
clientOptions.Handlers = new McpClientHandlers
|
||||
{
|
||||
SamplingHandler = (request, progress, ct) =>
|
||||
{
|
||||
return new ValueTask<CreateMessageResult>(new CreateMessageResult
|
||||
{
|
||||
Content = [new TextContentBlock { Text = "Response" }],
|
||||
Model = "test-model"
|
||||
});
|
||||
}
|
||||
};
|
||||
});
|
||||
|
||||
// Act - Call the tool DIRECTLY (not as a task)
|
||||
var result = await fixture.Client.CallToolAsync(
|
||||
"sampling-tool",
|
||||
arguments: new Dictionary<string, object?> { ["prompt"] = "Hello" },
|
||||
cancellationToken: TestContext.Current.CancellationToken);
|
||||
|
||||
await samplingCompleted.Task.WaitAsync(TimeSpan.FromSeconds(5), TestContext.Current.CancellationToken);
|
||||
|
||||
// Assert - No task should exist (tool was not called as a task)
|
||||
var tasks = await taskStore.ListTasksAsync(cancellationToken: TestContext.Current.CancellationToken);
|
||||
Assert.Empty(tasks.Tasks);
|
||||
|
||||
// And the result should still work
|
||||
Assert.NotNull(result);
|
||||
}
|
||||
|
||||
#pragma warning restore MCPEXP001
|
||||
|
||||
/// <summary>
|
||||
/// Test fixture that supports both server and client configuration for InputRequired status tests.
|
||||
/// </summary>
|
||||
private sealed class InputRequiredTestFixture : IAsyncDisposable
|
||||
{
|
||||
private readonly Pipe _clientToServerPipe = new();
|
||||
private readonly Pipe _serverToClientPipe = new();
|
||||
private readonly IServiceProvider _serviceProvider;
|
||||
private readonly McpServer _server;
|
||||
private readonly Task _serverTask;
|
||||
private readonly CancellationTokenSource _cts;
|
||||
|
||||
public McpClient Client { get; }
|
||||
public McpServer Server => _server;
|
||||
|
||||
public InputRequiredTestFixture(
|
||||
ILoggerFactory loggerFactory,
|
||||
Action<ServiceCollection, IMcpServerBuilder>? configureServer = null,
|
||||
Action<McpClientOptions>? configureClient = null)
|
||||
{
|
||||
_cts = new CancellationTokenSource();
|
||||
|
||||
// Configure server
|
||||
var services = new ServiceCollection();
|
||||
services.AddLogging();
|
||||
services.AddSingleton(loggerFactory);
|
||||
|
||||
var builder = services
|
||||
.AddMcpServer()
|
||||
.WithStreamServerTransport(
|
||||
_clientToServerPipe.Reader.AsStream(),
|
||||
_serverToClientPipe.Writer.AsStream());
|
||||
|
||||
configureServer?.Invoke(services, builder);
|
||||
|
||||
_serviceProvider = services.BuildServiceProvider(validateScopes: true);
|
||||
_server = _serviceProvider.GetRequiredService<McpServer>();
|
||||
_serverTask = _server.RunAsync(_cts.Token);
|
||||
|
||||
// Configure client
|
||||
var clientOptions = new McpClientOptions();
|
||||
configureClient?.Invoke(clientOptions);
|
||||
|
||||
// Create client synchronously (test code)
|
||||
Client = McpClient.CreateAsync(
|
||||
new StreamClientTransport(
|
||||
serverInput: _clientToServerPipe.Writer.AsStream(),
|
||||
_serverToClientPipe.Reader.AsStream(),
|
||||
loggerFactory),
|
||||
clientOptions: clientOptions,
|
||||
loggerFactory: loggerFactory,
|
||||
cancellationToken: TestContext.Current.CancellationToken).GetAwaiter().GetResult();
|
||||
}
|
||||
|
||||
public async ValueTask DisposeAsync()
|
||||
{
|
||||
await Client.DisposeAsync();
|
||||
await _cts.CancelAsync();
|
||||
|
||||
_clientToServerPipe.Writer.Complete();
|
||||
_serverToClientPipe.Writer.Complete();
|
||||
|
||||
try
|
||||
{
|
||||
await _serverTask;
|
||||
}
|
||||
catch (OperationCanceledException)
|
||||
{
|
||||
// Expected
|
||||
}
|
||||
|
||||
if (_serviceProvider is IAsyncDisposable asyncDisposable)
|
||||
{
|
||||
await asyncDisposable.DisposeAsync();
|
||||
}
|
||||
else if (_serviceProvider is IDisposable disposable)
|
||||
{
|
||||
disposable.Dispose();
|
||||
}
|
||||
|
||||
_cts.Dispose();
|
||||
}
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue