Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
namespace ModelContextProtocol.Extensions.Tasks;

/// <summary>
/// Specifies how a tool call participates in the MCP Tasks extension.
/// </summary>
public enum McpTaskExecutionMode
{
/// <summary>The tool always executes synchronously.</summary>
Synchronous,

/// <summary>The tool executes as a task when the client declares the Tasks extension.</summary>
Optional,

/// <summary>The tool requires the client to declare the Tasks extension.</summary>
Required,
}
Original file line number Diff line number Diff line change
Expand Up @@ -29,13 +29,36 @@ public static class McpTasksBuilderExtensions
/// <param name="store">The task store.</param>
/// <returns>The builder provided in <paramref name="builder"/>.</returns>
public static IMcpServerBuilder WithTasks(this IMcpServerBuilder builder, IMcpTaskStore store)
=> WithTasks(builder, store, static _ => { });

/// <summary>
/// Enables MCP Tasks support backed by the specified task store.
/// </summary>
/// <param name="builder">The server builder.</param>
/// <param name="store">The task store.</param>
/// <param name="configure">A callback that configures per-call task execution behavior.</param>
/// <returns>The builder provided in <paramref name="builder"/>.</returns>
public static IMcpServerBuilder WithTasks(
this IMcpServerBuilder builder,
IMcpTaskStore store,
Action<McpTasksOptions> configure)
{
#if NET
ArgumentNullException.ThrowIfNull(builder);
ArgumentNullException.ThrowIfNull(store);
ArgumentNullException.ThrowIfNull(configure);
#else
if (builder is null) throw new ArgumentNullException(nameof(builder));
if (store is null) throw new ArgumentNullException(nameof(store));
if (configure is null) throw new ArgumentNullException(nameof(configure));
#endif

McpTasksOptions taskOptions = new();
configure(taskOptions);
#if NET
ArgumentNullException.ThrowIfNull(taskOptions.ExecutionModeSelector);
#else
if (taskOptions.ExecutionModeSelector is null) throw new ArgumentNullException(nameof(taskOptions.ExecutionModeSelector));
#endif

// Resolve ILoggerFactory from the provider (rather than requiring the caller to pass one) so the
Expand All @@ -45,18 +68,21 @@ public static IMcpServerBuilder WithTasks(this IMcpServerBuilder builder, IMcpTa
sp => new McpTasksConfigureOptions(
store,
sp.GetRequiredService<IServiceScopeFactory>(),
sp.GetService<ILoggerFactory>()));
sp.GetService<ILoggerFactory>(),
taskOptions));
return builder;
}

private sealed class McpTasksConfigureOptions(
IMcpTaskStore store,
IServiceScopeFactory serviceScopeFactory,
ILoggerFactory? loggerFactory) : IConfigureOptions<McpServerOptions>
ILoggerFactory? loggerFactory,
McpTasksOptions taskOptions) : IConfigureOptions<McpServerOptions>
{
private readonly IMcpTaskStore _store = store;
private readonly IServiceScopeFactory _serviceScopeFactory = serviceScopeFactory;
private readonly ILogger _logger = (loggerFactory ?? NullLoggerFactory.Instance).CreateLogger<McpTasksConfigureOptions>();
private readonly McpTasksOptions _taskOptions = taskOptions;
private readonly ConcurrentDictionary<string, CancellationTokenSource> _cancellationSources = new(StringComparer.Ordinal);

public void Configure(McpServerOptions options)
Expand Down Expand Up @@ -108,9 +134,28 @@ public void Configure(McpServerOptions options)
options.Filters.Request.CallToolWithAlternateFilters.Count,
async (request, next, cancellationToken) =>
{
var executionMode = _taskOptions.ExecutionModeSelector(request);
if (executionMode is not McpTaskExecutionMode.Synchronous and
not McpTaskExecutionMode.Optional and
not McpTaskExecutionMode.Required)
{
throw new InvalidOperationException(
$"{nameof(McpTasksOptions.ExecutionModeSelector)} returned an invalid {nameof(McpTaskExecutionMode)} value.");
}

if (executionMode == McpTaskExecutionMode.Synchronous)
{
return await next(request, cancellationToken).ConfigureAwait(false);
}

if (!IsJuly2026OrLaterProtocolRequest(request.JsonRpcRequest) ||
!HasTaskExtensionOptIn(request.JsonRpcRequest))
{
if (executionMode == McpTaskExecutionMode.Required)
{
throw CreateMissingTasksCapabilityException();
}

return await next(request, cancellationToken).ConfigureAwait(false);
}

Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,22 @@
using ModelContextProtocol.Protocol;
using ModelContextProtocol.Server;

namespace ModelContextProtocol.Extensions.Tasks;

/// <summary>
/// Configures server-side MCP Tasks behavior.
/// </summary>
public sealed class McpTasksOptions
{
/// <summary>
/// Gets or sets the callback that selects task execution behavior for each tool call.
/// </summary>
/// <remarks>
/// The default treats every tool as task-capable, preserving the behavior of
/// <c>WithTasks</c>.
/// The callback may inspect <see cref="RequestContext{TParams}.MatchedPrimitive"/> to select
/// behavior from tool metadata without requiring Core to reference the Tasks extension.
/// </remarks>
public Func<RequestContext<CallToolRequestParams>, McpTaskExecutionMode> ExecutionModeSelector { get; set; } =
static _ => McpTaskExecutionMode.Optional;
}
Original file line number Diff line number Diff line change
Expand Up @@ -129,18 +129,15 @@ public async Task RunMrtrConformanceTest(string scenario)

[Theory]
[InlineData("tasks-wire-fields")]
[InlineData("tasks-lifecycle")]
[InlineData("tasks-capability-negotiation")]
[InlineData("tasks-request-headers")]
[InlineData("tasks-dispatch-and-envelope")]
[InlineData("tasks-status-notifications")]
[InlineData("tasks-required-task-error")]
[InlineData("tasks-request-state-removal")]
[InlineData("tasks-mrtr-input")]
// Most remaining scenarios require per-tool task execution configuration that the SDK
// does not currently expose; status notifications await an upstream harness rewrite.
// Keep them listed here for incremental enablement.
// [InlineData("tasks-lifecycle")]
// [InlineData("tasks-capability-negotiation")]
// [InlineData("tasks-request-headers")]
// [InlineData("tasks-dispatch-and-envelope")]
// [InlineData("tasks-status-notifications")]
// [InlineData("tasks-required-task-error")]
// [InlineData("tasks-mrtr-composition")]
[InlineData("tasks-mrtr-composition")]
public async Task RunTasksExtensionConformanceTest(string scenario)
{
Assert.SkipWhen(!NodeHelpers.IsNodeInstalled(), "Node.js is not installed. Skipping conformance tests.");
Expand Down
13 changes: 12 additions & 1 deletion tests/ModelContextProtocol.ConformanceServer/Program.cs
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,18 @@ private static void ConfigureConformanceMcpServer(
{
DefaultPollIntervalMs = 50,
DefaultTimeToLive = TimeSpan.FromMinutes(5),
})
},
options => options.ExecutionModeSelector = request =>
request.Params?.Name switch
{
"slow_compute" or "protocol_error_job" or "confirm_delete" or "multi_input"
=> McpTaskExecutionMode.Optional,
"failing_job"
=> McpTaskExecutionMode.Required,
"test_tool_with_task" when request.Params.InputResponses is { Count: > 0 }
=> McpTaskExecutionMode.Required,
_ => McpTaskExecutionMode.Synchronous,
})
.WithTools<ConformanceTools>()
.WithTools<ConformanceTaskTools>()
.WithTools<IncompleteResultTools>()
Expand Down
87 changes: 86 additions & 1 deletion tests/ModelContextProtocol.Tests/Server/McpTaskStoreTests.cs
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,13 @@ protected override void ConfigureServices(ServiceCollection services, IMcpServer
.WithTasks(new InMemoryMcpTaskStore
{
DefaultPollIntervalMs = 50,
});
}, options => options.ExecutionModeSelector = request =>
request.Params?.Name switch
{
"sync-tool" => McpTaskExecutionMode.Synchronous,
"required-tool" => McpTaskExecutionMode.Required,
_ => McpTaskExecutionMode.Optional,
});
}

[Fact]
Expand All @@ -52,6 +58,79 @@ public async Task CallToolAsTaskAsync_WithTaskCapability_ReturnsCreateTaskResult
Assert.Equal(McpTaskStatus.Working, augmented.TaskCreated.Status);
}

[Fact]
public async Task CallToolAsync_SynchronousExecutionMode_ReturnsResultDirectly()
{
await using McpClient client = await CreateMcpClientForServer();

CallToolResult result = await client.CallToolAsync(
"sync-tool",
cancellationToken: TestContext.Current.CancellationToken);

Assert.Equal("sync", Assert.IsType<TextContentBlock>(result.Content[0]).Text);
}

[Fact]
public async Task CallToolAsync_RequiredExecutionModeWithoutCapability_Throws()
{
await using McpClient client = await CreateMcpClientForServer();

MissingRequiredClientCapabilityException exception = await Assert.ThrowsAsync<MissingRequiredClientCapabilityException>(async () =>
await client.CallToolAsync(
"required-tool",
cancellationToken: TestContext.Current.CancellationToken));

Assert.Equal(McpErrorCode.MissingRequiredClientCapability, exception.ErrorCode);
}

[Fact]
public async Task CallToolAsync_RequiredExecutionModeWithLegacyProtocol_Throws()
{
await using McpClient client = await CreateMcpClientForServer(new McpClientOptions
{
ProtocolVersion = McpProtocolVersions.November2025ProtocolVersion,
});

MissingRequiredClientCapabilityException exception = await Assert.ThrowsAsync<MissingRequiredClientCapabilityException>(async () =>
await client.CallToolAsync(
"required-tool",
cancellationToken: TestContext.Current.CancellationToken));

Assert.Equal(McpErrorCode.MissingRequiredClientCapability, exception.ErrorCode);
}

[Fact]
public async Task CallToolAsTaskAsync_RequiredExecutionMode_ReturnsTask()
{
await using McpClient client = await CreateMcpClientForServer();

ResultOrCreatedTask<CallToolResult> result = await client.CallToolAsTaskAsync(
new CallToolRequestParams { Name = "required-tool" },
TestContext.Current.CancellationToken);

Assert.True(result.IsTask);

CallToolResult completedResult = await client.CallToolWithPollingAsync(
new CallToolRequestParams { Name = "required-tool" },
cancellationToken: TestContext.Current.CancellationToken);

Assert.Equal("required", Assert.IsType<TextContentBlock>(completedResult.Content[0]).Text);
}

[Fact]
public void WithTasks_NullExecutionModeSelector_Throws()
{
ServiceCollection services = new();

Assert.Throws<ArgumentNullException>(() =>
services
.AddMcpServer()
.WithStreamServerTransport(Stream.Null, Stream.Null)
.WithTasks(
new InMemoryMcpTaskStore(),
options => options.ExecutionModeSelector = null!));
}

[Fact]
public async Task CallToolAsync_WithTaskStore_PollsToCompletion()
{
Expand Down Expand Up @@ -609,6 +688,12 @@ public async Task MultiElicit_ViaTask_HandlerCalledExactlyOncePerUniqueKey_Acros
[McpServerToolType]
private sealed class TaskStoreTestTools
{
[McpServerTool(Name = "sync-tool")]
public static string SyncTool() => "sync";

[McpServerTool(Name = "required-tool")]
public static string RequiredTool() => "required";
Comment thread
PranavSenthilnathan marked this conversation as resolved.

[McpServerTool(Name = "slow-tool"), System.ComponentModel.Description("A tool that takes time")]
public static async Task<string> SlowTool(CancellationToken cancellationToken)
{
Expand Down