diff --git a/docs/concepts/tasks/tasks.md b/docs/concepts/tasks/tasks.md index c1ab0c23a..4e6dabb9c 100644 --- a/docs/concepts/tasks/tasks.md +++ b/docs/concepts/tasks/tasks.md @@ -282,6 +282,108 @@ public sealed class MyTaskStore : IMcpTaskStore } ``` +### Delegating execution to an external runtime + +By default, `WithTasks` executes the tool in-process on the .NET thread pool. To delegate +execution to a durable system such as Temporal, Orleans, Hangfire, or an external queue, +register an : + +```csharp +builder.WithTasks( + myDurableTaskStore, + options => + { + options.TaskExecutor = new TemporalTaskExecutor(workflowClient); + }); +``` + +An executor can also be resolved from the service provider — register `IMcpTaskExecutor` in DI +and omit `TaskExecutor`. The executor is resolved from each task's execution scope, so scoped +registrations get one instance per task; singleton registrations behave as usual. When neither +is configured, tasks run in-process exactly as before. + +The executor is invoked after the task record is durably created in the store. + must return only +after execution has been durably started — for example, after the external runtime has +accepted the job — mirroring the durability requirement SEP-2663 §306 places on +. It must not wait +for the task to complete. If `StartAsync` throws, the exception is not returned as an error +from the original `tools/call`: that call still succeeds with +, the task is marked failed via +`SetFailedAsync`, and the client discovers the failure on its first `tasks/get` poll. By +contrast, failures before the task record exists — resolving the executor or + — do fail the +original `tools/call`. After a successful `StartAsync`, the SDK stops tracking the task and +the store is the single source of truth for its state. + +One boundary to be aware of is the window between `CreateTaskAsync` completing and +`StartAsync` returning: if the process exits during it, the store is left with a `Working` +task whose work was never submitted to the external runtime. The SDK performs no +reconciliation of such tasks, so integrations that must recover from a crash in this window +need their own strategy — for example TTL cleanup, startup reconciliation, an outbox, or a +durable execution intent. Using the task ID as an idempotency key makes resubmission safe, +but it does not by itself retry a submission that was never attempted. + +The passed to the +executor exposes the task identity, the matched tool request bound to a fresh execution +scope, and a token that fires on `tasks/cancel`. Executors that want the tool to run +locally call +, +which runs the remaining request filters and the tool, records the outcome in the store, +and releases the execution scope. Executors that hand execution off to an external system +should read what they need from + and then call + to release +the scope-bound services. + +Primitive matching and the filters registered before Tasks — including ASP.NET Core +authorization — have already run by the time `StartAsync` is called. The remaining +alternate-result filters and the ordinary call-tool filters run only inside +`RunToolPipelineAsync`, so an executor that performs a pure handoff to an external runtime +bypasses them. When the tool pipeline will not run locally, validation, auditing, +transformations, and other cross-cutting policies must be applied by the external runtime — +or by a filter registered before Tasks — instead. + +```csharp +public sealed class TemporalTaskExecutor(ITemporalClient workflowClient) : IMcpTaskExecutor +{ + public async ValueTask StartAsync( + McpTaskExecutionContext context, CancellationToken cancellationToken) + { + // Submit the tool request to the durable runtime. The workflow communicates with + // IMcpTaskStore directly to record progress and results. + await workflowClient.StartWorkflowAsync( + "run-mcp-task", + new McpTaskPayload(context.TaskId, context.Request.Params), + id: context.TaskId, + cancellationToken); + + // The scope-bound services are no longer needed in this process. + await context.DisposeAsync(); + } +} +``` + +`tasks/get`, `tasks/update`, and `tasks/cancel` continue to be served entirely from the +`IMcpTaskStore`, so a different server instance can serve polling clients after the process +that started the task exits — the acceptance scenario for durable execution. + +Note that +signals cancellation only within the process that created the task. A `tasks/cancel` handled +by a different server instance can update the shared store, but it cannot signal that +process's token. External-runtime integrations whose cancellation must survive server +replacement therefore need to propagate it through their store or another durable mechanism; +the token remains the cancellation signal for the default in-process executor and other +same-process execution paths. + +Note that elicitation and sampling issued from *outside* the process that owns the client +session cannot be routed through the task's input-request channel; an external worker that +needs multi-round-trip input should rely on the store's + +event, or run the pipeline locally via + +from the process that owns the session. + ### Status semantics is the terminal status whenever the diff --git a/src/ModelContextProtocol.Extensions.Tasks/Server/IMcpTaskExecutor.cs b/src/ModelContextProtocol.Extensions.Tasks/Server/IMcpTaskExecutor.cs new file mode 100644 index 000000000..1ca98723c --- /dev/null +++ b/src/ModelContextProtocol.Extensions.Tasks/Server/IMcpTaskExecutor.cs @@ -0,0 +1,51 @@ +using ModelContextProtocol.Protocol; +using ModelContextProtocol.Server; + +namespace ModelContextProtocol.Extensions.Tasks; + +/// +/// Executes a task created by the Tasks extension after the task record has been +/// durably created in the . +/// +/// +/// +/// By default, tasks execute in-process via the .NET thread pool. Registering a custom +/// executor delegates execution to an external system such as Temporal, Orleans, Hangfire, +/// or a distributed queue. The executor is invoked once per task, after the task record is +/// created and the execution context is fully wired. +/// +/// +/// must return only after execution has been durably started +/// (e.g., the external runtime has accepted the job), mirroring the durability requirement +/// SEP-2663 §306 places on . It must not wait +/// for the task to complete; completion is recorded in the store by whichever system +/// performs the execution. +/// +/// +/// If throws, the exception is not returned as an error from the +/// original tools/call: that call still succeeds with , +/// the task is marked failed via , and the client +/// discovers the failure on its first poll. After a successful , +/// the SDK no longer tracks the task; the store is the single source of truth for its state. +/// +/// +/// See the SEP-2663 +/// specification for details on the tasks extension. +/// +/// +public interface IMcpTaskExecutor +{ + /// + /// Starts execution of a task. + /// + /// + /// The execution context for the task, providing the task identity, the matched tool + /// request bound to a fresh execution scope, and a helper for running the normal tool + /// invocation pipeline locally. + /// + /// + /// A token that fires when the task is cancelled via tasks/cancel. + /// + /// A that completes when execution has been durably started. + ValueTask StartAsync(McpTaskExecutionContext context, CancellationToken cancellationToken); +} diff --git a/src/ModelContextProtocol.Extensions.Tasks/Server/McpTaskExecutionContext.cs b/src/ModelContextProtocol.Extensions.Tasks/Server/McpTaskExecutionContext.cs new file mode 100644 index 000000000..44296e29d --- /dev/null +++ b/src/ModelContextProtocol.Extensions.Tasks/Server/McpTaskExecutionContext.cs @@ -0,0 +1,118 @@ +using ModelContextProtocol.Protocol; +using ModelContextProtocol.Server; + +namespace ModelContextProtocol.Extensions.Tasks; + +/// +/// The execution context handed to an when a task starts. +/// +/// +/// +/// The context owns the request-scoped services and the cancellation registration for task +/// execution. Executors that run the tool locally via +/// never dispose anything explicitly; the context releases its resources when the pipeline +/// completes. Executors that hand execution off to an external system should extract what +/// they need from and then call once the +/// scope-bound services are no longer needed. +/// +/// +/// carries a server whose outgoing requests (elicitation, sampling) +/// are intercepted and routed through the task's pending input requests, so responses +/// submitted via tasks/update are delivered even when a different server instance +/// serves the polling client. +/// +/// +/// and coordinate through an +/// atomic state transition: concurrent callers cannot both win, so the pipeline runs at most +/// once and disposal cannot race with it. Every caller but the winner observes the context +/// as disposed. +/// +/// +public sealed class McpTaskExecutionContext : IAsyncDisposable +{ + private readonly Func, CancellationToken, Task> _pipelineRunner; + private readonly Func _disposer; + private int _disposed; + + internal McpTaskExecutionContext( + McpTaskInfo taskInfo, + RequestContext request, + CancellationToken cancellation, + Func, CancellationToken, Task> pipelineRunner, + Func disposer) + { + TaskInfo = taskInfo; + Request = request; + CancellationToken = cancellation; + _pipelineRunner = pipelineRunner; + _disposer = disposer; + } + + /// + /// Gets the unique identifier of the created task. + /// + public string TaskId => TaskInfo.TaskId; + + /// + /// Gets the store record for the created task, with an initial status of + /// . + /// + public McpTaskInfo TaskInfo { get; } + + /// + /// Gets the matched tool request, bound to the task's execution scope, with the task + /// outgoing-request interceptor already attached. + /// + public RequestContext Request { get; } + + /// + /// Gets a token that fires when the task is cancelled via tasks/cancel. + /// + public CancellationToken CancellationToken { get; } + + /// + /// Runs the normal tool invocation pipeline (the remaining request filters and the tool + /// itself), records the outcome in the task store, and releases the context's resources. + /// + /// + /// + /// Use this when the executor wants the tool to run locally, preserving the behavior of + /// WithTasks without a custom executor. Outcomes — including cancellation, protocol + /// errors, and unhandled exceptions — are recorded in the store by this call; it does not + /// rethrow them. + /// + /// + /// After calling this method, the executor must not use again, and + /// becomes a no-op. + /// + /// + /// A token to cancel pipeline execution. + public async ValueTask RunToolPipelineAsync(CancellationToken cancellationToken) + { + if (Interlocked.CompareExchange(ref _disposed, 1, 0) != 0) + { + throw new ObjectDisposedException(nameof(McpTaskExecutionContext)); + } + + await _pipelineRunner(Request, cancellationToken).ConfigureAwait(false); + } + + /// + /// Releases the context's resources: the request-scoped services of the execution scope + /// and the cancellation registration for the task. + /// + /// + /// Executors that hand execution off to an external system call this once they no longer + /// need . It is called automatically when + /// completes; calling it afterwards is a no-op. + /// + public async ValueTask DisposeAsync() + { + if (Interlocked.CompareExchange(ref _disposed, 1, 0) != 0) + { + return; + } + + await _disposer().ConfigureAwait(false); + } +} diff --git a/src/ModelContextProtocol.Extensions.Tasks/Server/McpTasksBuilderExtensions.cs b/src/ModelContextProtocol.Extensions.Tasks/Server/McpTasksBuilderExtensions.cs index e61466a69..2e4cea9ad 100644 --- a/src/ModelContextProtocol.Extensions.Tasks/Server/McpTasksBuilderExtensions.cs +++ b/src/ModelContextProtocol.Extensions.Tasks/Server/McpTasksBuilderExtensions.cs @@ -171,8 +171,17 @@ private async ValueTask> RunAsTaskAsync( }; McpTaskInfo taskInfo; + IMcpTaskExecutor executor; try { + // Resolve the executor from the execution scope (falling back to the process-local + // default) rather than the root provider so scoped and transient registrations get + // correct lifetimes; singleton registrations still yield the same instance. Resolving + // before the task record is created keeps a DI misconfiguration from leaving a + // durably-created task stuck at Working: the resolution error fails tools/call instead. + executor = _taskOptions.TaskExecutor + ?? executionScope.ServiceProvider.GetService() + ?? ProcessLocalMcpTaskExecutor.Instance; taskInfo = await _store.CreateTaskAsync(cancellationToken).ConfigureAwait(false); } catch @@ -189,9 +198,24 @@ private async ValueTask> RunAsTaskAsync( // Capture the token before dispatching. Cancellation can remove and dispose the source // before the background delegate starts. var taskCancellationToken = cts.Token; - _ = Task.Run( - () => ExecuteTaskAsync(next, executionRequest, taskId, taskCancellationToken, executionScope), - CancellationToken.None); + var context = new McpTaskExecutionContext( + taskInfo, + executionRequest, + taskCancellationToken, + (req, ct) => ExecuteTaskAsync(next, req, taskId, ct, executionScope), + () => ReleaseExecutionResourcesAsync(executionScope, taskId)); + + try + { + await executor.StartAsync(context, taskCancellationToken).ConfigureAwait(false); + } + catch (Exception ex) + { + // The task record exists, so the client will poll it. Record the start failure as + // the task's failure before returning the task alternate, so the client's first + // poll observes the terminal state rather than racing it. + await RecordStartFailureAsync(context, ex).ConfigureAwait(false); + } return ResultOrAlternate.FromAlternate( ToCreateTaskResult(taskInfo), @@ -235,10 +259,52 @@ private async Task ExecuteTaskAsync( } finally { - if (_cancellationSources.TryRemove(taskId, out var registeredCts)) - { - registeredCts.Dispose(); - } + RemoveCancellationSource(taskId); + } + } + + private async Task RecordStartFailureAsync(McpTaskExecutionContext context, Exception exception) + { + _logger.LogError(exception, "Starting execution of task '{TaskId}' failed.", context.TaskId); + + try + { + await context.DisposeAsync().ConfigureAwait(false); + } + catch (Exception disposeEx) + { + _logger.LogError(disposeEx, "Failed to release resources of task '{TaskId}' after a failed start.", context.TaskId); + } + + try + { + var error = new JsonRpcErrorDetail { Code = (int)McpErrorCode.InternalError, Message = exception.Message }; + var errorJson = JsonSerializer.SerializeToElement(error, McpJsonUtilities.DefaultOptions.GetTypeInfo()); + await _store.SetFailedAsync(context.TaskId, errorJson).ConfigureAwait(false); + } + catch (Exception storeEx) + { + _logger.LogError(storeEx, "Failed to record the failure of task '{TaskId}'.", context.TaskId); + } + } + + private async Task ReleaseExecutionResourcesAsync(AsyncServiceScope executionScope, string taskId) + { + try + { + await executionScope.DisposeAsync().ConfigureAwait(false); + } + finally + { + RemoveCancellationSource(taskId); + } + } + + private void RemoveCancellationSource(string taskId) + { + if (_cancellationSources.TryRemove(taskId, out var registeredCts)) + { + registeredCts.Dispose(); } } diff --git a/src/ModelContextProtocol.Extensions.Tasks/Server/McpTasksOptions.cs b/src/ModelContextProtocol.Extensions.Tasks/Server/McpTasksOptions.cs index 818a53fcf..e040857ef 100644 --- a/src/ModelContextProtocol.Extensions.Tasks/Server/McpTasksOptions.cs +++ b/src/ModelContextProtocol.Extensions.Tasks/Server/McpTasksOptions.cs @@ -19,4 +19,15 @@ public sealed class McpTasksOptions /// public Func, McpTaskExecutionMode> ExecutionModeSelector { get; set; } = static _ => McpTaskExecutionMode.Optional; + + /// + /// Gets or sets the executor that starts task execution. + /// + /// + /// When (the default), the extension resolves a single registered + /// from the service provider, if one exists. If neither is + /// present, tasks execute in-process on the .NET thread pool, preserving the behavior of + /// WithTasks without a custom executor. + /// + public IMcpTaskExecutor? TaskExecutor { get; set; } } diff --git a/src/ModelContextProtocol.Extensions.Tasks/Server/ProcessLocalMcpTaskExecutor.cs b/src/ModelContextProtocol.Extensions.Tasks/Server/ProcessLocalMcpTaskExecutor.cs new file mode 100644 index 000000000..5f54a7a25 --- /dev/null +++ b/src/ModelContextProtocol.Extensions.Tasks/Server/ProcessLocalMcpTaskExecutor.cs @@ -0,0 +1,40 @@ +namespace ModelContextProtocol.Extensions.Tasks; + +/// +/// The default that runs the tool invocation pipeline +/// in-process on the .NET thread pool. +/// +/// +/// This is the executor used by WithTasks when no custom executor is configured, +/// and the base behavior a custom executor can fall back to via +/// . It is exposed as a type so +/// decorators can identify or wrap the default, but it cannot be constructed externally; +/// use . +/// +public sealed class ProcessLocalMcpTaskExecutor : IMcpTaskExecutor +{ + private ProcessLocalMcpTaskExecutor() + { + } + + /// + /// Gets the singleton instance of the process-local executor. + /// + public static ProcessLocalMcpTaskExecutor Instance { get; } = new(); + + /// + public ValueTask StartAsync(McpTaskExecutionContext context, CancellationToken cancellationToken) + { +#if NET + ArgumentNullException.ThrowIfNull(context); +#else + if (context is null) throw new ArgumentNullException(nameof(context)); +#endif + + _ = Task.Run( + () => context.RunToolPipelineAsync(context.CancellationToken).AsTask(), + CancellationToken.None); + + return default; + } +} diff --git a/tests/ModelContextProtocol.Tests/Server/McpServerTaskExecutorTests.cs b/tests/ModelContextProtocol.Tests/Server/McpServerTaskExecutorTests.cs new file mode 100644 index 000000000..e43a52c46 --- /dev/null +++ b/tests/ModelContextProtocol.Tests/Server/McpServerTaskExecutorTests.cs @@ -0,0 +1,408 @@ +using Microsoft.Extensions.DependencyInjection; +using ModelContextProtocol.Client; +using ModelContextProtocol.Extensions.Tasks; +using ModelContextProtocol.Protocol; +using ModelContextProtocol.Server; +using ModelContextProtocol.Tests.Utils; +using System.Runtime.InteropServices; +using System.Text.Json; + +namespace ModelContextProtocol.Tests.Server; + +/// +/// Tests for , the extension point that lets a server +/// delegate task execution to an external system instead of running the tool +/// in-process. +/// +public class McpServerTaskExecutorTests : ClientServerTestBase +{ + private readonly InMemoryMcpTaskStore _taskStore = new() { DefaultPollIntervalMs = 10 }; + private readonly TaskCompletionSource _executorInvoked = new(TaskCreationOptions.RunContinuationsAsynchronously); + private readonly TaskCompletionSource _scopeDisposed = new(TaskCreationOptions.RunContinuationsAsynchronously); + private readonly TaskCompletionSource _toolStarted = new(TaskCreationOptions.RunContinuationsAsynchronously); + private readonly TaskCompletionSource _toolCancellationFired = new(TaskCreationOptions.RunContinuationsAsynchronously); + private readonly TaskCompletionSource _executorCancelled = new(TaskCreationOptions.RunContinuationsAsynchronously); + private Exception? _startException; + private bool _runPipelineLocally; + private int _toolStartCount; + + public McpServerTaskExecutorTests(ITestOutputHelper testOutputHelper) : base(testOutputHelper) + { +#if !NET + Assert.SkipWhen(RuntimeInformation.IsOSPlatform(OSPlatform.Windows), "https://github.com/modelcontextprotocol/csharp-sdk/issues/587"); +#endif + } + + protected override void ConfigureServices(ServiceCollection services, IMcpServerBuilder mcpServerBuilder) + { + services.AddScoped(_ => new ScopedDependency(_scopeDisposed)); + + mcpServerBuilder + .WithTasks( + _taskStore, + options => + { + options.TaskExecutor = new CallbackTaskExecutor(this); + }) + .WithTools([McpServerTool.Create( + async (CancellationToken ct) => + { + Interlocked.Increment(ref _toolStartCount); + _toolStarted.TrySetResult(true); + try + { + await Task.Delay(Timeout.Infinite, ct); + return "completed"; + } + catch (OperationCanceledException) + { + _toolCancellationFired.TrySetResult(true); + throw; + } + }, + new McpServerToolCreateOptions { Name = "long-running-tool" }), + McpServerTool.Create( + () => "local result", + new McpServerToolCreateOptions { Name = "local-tool" })]); + } + + [Fact] + public async Task CustomExecutor_ReceivesTaskAndRequestBoundToExecutionScope() + { + await using var client = await CreateMcpClientForServer(); + var cancellationToken = TestContext.Current.CancellationToken; + + var augmented = await client.CallToolAsTaskAsync( + new CallToolRequestParams { Name = "long-running-tool" }, + cancellationToken); + + Assert.True(augmented.IsTask); + var context = await _executorInvoked.Task.WaitAsync(TestConstants.DefaultTimeout, cancellationToken); + Assert.Equal(augmented.TaskCreated!.TaskId, context.TaskId); + Assert.Equal(McpTaskStatus.Working, context.TaskInfo.Status); + Assert.Equal("long-running-tool", context.Request.MatchedPrimitive?.Id); + Assert.NotNull(context.Request.Services); + Assert.Same( + context.Request.Services!.GetRequiredService(), + context.Request.Services.GetRequiredService()); + + // The tool body must not run until the executor starts the pipeline. + Assert.False(_toolStarted.Task.IsCompleted); + } + + [Fact] + public async Task CustomExecutor_NotRunningPipeline_ToolBodyNeverRunsInProcess() + { + await using var client = await CreateMcpClientForServer(); + var cancellationToken = TestContext.Current.CancellationToken; + + var augmented = await client.CallToolAsTaskAsync( + new CallToolRequestParams { Name = "long-running-tool" }, + cancellationToken); + + Assert.True(augmented.IsTask); + var context = await _executorInvoked.Task.WaitAsync(TestConstants.DefaultTimeout, cancellationToken); + + // Simulate the external runtime completing the task directly through the store. + var result = JsonSerializer.SerializeToElement( + new CallToolResult { Content = [new TextContentBlock { Text = "external result" }] }, + McpJsonUtilities.DefaultOptions.GetTypeInfo()); + await _taskStore.SetCompletedAsync(context.TaskId, result, cancellationToken); + + var task = await PollUntilTerminalAsync(client, augmented.TaskCreated!.TaskId, cancellationToken); + Assert.IsType(task); + Assert.False(_toolStarted.Task.IsCompleted); + } + + [Fact] + public async Task CustomExecutor_RunsPipelineLocally_ResultRecordedInStore() + { + _runPipelineLocally = true; + + await using var client = await CreateMcpClientForServer(); + var cancellationToken = TestContext.Current.CancellationToken; + + var result = await client.CallToolWithPollingAsync( + new CallToolRequestParams { Name = "local-tool" }, + cancellationToken: cancellationToken); + + Assert.Equal("local result", Assert.IsType(Assert.Single(result.Content)).Text); + } + + [Fact] + public async Task CustomExecutor_RunsPipelineLocally_ScopeDisposedAfterCompletion() + { + _runPipelineLocally = true; + + await using var client = await CreateMcpClientForServer(); + var cancellationToken = TestContext.Current.CancellationToken; + + var augmented = await client.CallToolAsTaskAsync( + new CallToolRequestParams { Name = "local-tool" }, + cancellationToken); + Assert.True(augmented.IsTask); + + var task = await PollUntilTerminalAsync(client, augmented.TaskCreated!.TaskId, cancellationToken); + Assert.IsType(task); + Assert.True(await _scopeDisposed.Task.WaitAsync(TestConstants.DefaultTimeout, cancellationToken)); + } + + [Fact] + public async Task CustomExecutor_ThrowingFromStartAsync_MarksTaskFailed() + { + _startException = new InvalidOperationException("external runtime unavailable"); + + await using var client = await CreateMcpClientForServer(); + var cancellationToken = TestContext.Current.CancellationToken; + + var augmented = await client.CallToolAsTaskAsync( + new CallToolRequestParams { Name = "long-running-tool" }, + cancellationToken); + + Assert.True(augmented.IsTask); + + var task = await PollUntilTerminalAsync(client, augmented.TaskCreated!.TaskId, cancellationToken); + var failed = Assert.IsType(task); + Assert.Contains("external runtime unavailable", failed.Error.GetRawText()); + Assert.False(_toolStarted.Task.IsCompleted); + } + + [Fact] + public async Task CustomExecutor_TasksCancel_FiresExecutorCancellationToken() + { + await using var client = await CreateMcpClientForServer(); + var cancellationToken = TestContext.Current.CancellationToken; + + var augmented = await client.CallToolAsTaskAsync( + new CallToolRequestParams { Name = "long-running-tool" }, + cancellationToken); + Assert.True(augmented.IsTask); + await _executorInvoked.Task.WaitAsync(TestConstants.DefaultTimeout, cancellationToken); + + await client.CancelTaskAsync(augmented.TaskCreated!.TaskId, cancellationToken); + + Assert.True(await _executorCancelled.Task.WaitAsync(TestConstants.DefaultTimeout, cancellationToken)); + } + + [Fact] + public async Task CustomExecutor_DisposesContext_ReleasesScopeWithoutRunningTool() + { + await using var client = await CreateMcpClientForServer(); + var cancellationToken = TestContext.Current.CancellationToken; + + var augmented = await client.CallToolAsTaskAsync( + new CallToolRequestParams { Name = "long-running-tool" }, + cancellationToken); + Assert.True(augmented.IsTask); + var context = await _executorInvoked.Task.WaitAsync(TestConstants.DefaultTimeout, cancellationToken); + + await context.DisposeAsync(); + + await _scopeDisposed.Task.WaitAsync(TestConstants.DefaultTimeout, cancellationToken); + Assert.False(_toolStarted.Task.IsCompleted); + + // DisposeAsync is idempotent. + await context.DisposeAsync(); + } + + [Fact] + public async Task CustomExecutor_RunsPipelineAfterDispose_ThrowsObjectDisposed() + { + await using var client = await CreateMcpClientForServer(); + var cancellationToken = TestContext.Current.CancellationToken; + + var augmented = await client.CallToolAsTaskAsync( + new CallToolRequestParams { Name = "long-running-tool" }, + cancellationToken); + Assert.True(augmented.IsTask); + var context = await _executorInvoked.Task.WaitAsync(TestConstants.DefaultTimeout, cancellationToken); + + await context.DisposeAsync(); + + await Assert.ThrowsAsync( + () => context.RunToolPipelineAsync(cancellationToken).AsTask()); + } + + [Fact] + public async Task CustomExecutor_StartsPipelineLater_TokenCancelsPipeline() + { + await using var client = await CreateMcpClientForServer(); + var cancellationToken = TestContext.Current.CancellationToken; + + var augmented = await client.CallToolAsTaskAsync( + new CallToolRequestParams { Name = "long-running-tool" }, + cancellationToken); + Assert.True(augmented.IsTask); + var context = await _executorInvoked.Task.WaitAsync(TestConstants.DefaultTimeout, cancellationToken); + + // The executor hands the task to an external runtime, which later runs the pipeline + // locally. Cancelling via tasks/cancel must fire the context token and cancel the + // pipeline. + _ = Task.Run(() => context.RunToolPipelineAsync(context.CancellationToken).AsTask(), CancellationToken.None); + await _toolStarted.Task.WaitAsync(TestConstants.DefaultTimeout, cancellationToken); + + await client.CancelTaskAsync(augmented.TaskCreated!.TaskId, cancellationToken); + + Assert.True(await _toolCancellationFired.Task.WaitAsync(TestConstants.DefaultTimeout, cancellationToken)); + + var task = await PollUntilTerminalAsync(client, augmented.TaskCreated!.TaskId, cancellationToken); + Assert.IsType(task); + } + + [Fact] + public async Task CustomExecutor_ConcurrentPipelineStarts_RunToolOnlyOnce() + { + await using var client = await CreateMcpClientForServer(); + var cancellationToken = TestContext.Current.CancellationToken; + + var augmented = await client.CallToolAsTaskAsync( + new CallToolRequestParams { Name = "long-running-tool" }, + cancellationToken); + Assert.True(augmented.IsTask); + var context = await _executorInvoked.Task.WaitAsync(TestConstants.DefaultTimeout, cancellationToken); + + // Two concurrent starts of the pipeline: exactly one may win; the loser must observe + // the context as disposed rather than starting a second execution of the tool. + var first = Task.Run(() => context.RunToolPipelineAsync(context.CancellationToken).AsTask(), CancellationToken.None); + var second = Task.Run(() => context.RunToolPipelineAsync(context.CancellationToken).AsTask(), CancellationToken.None); + + await _toolStarted.Task.WaitAsync(TestConstants.DefaultTimeout, cancellationToken); + Assert.Equal(1, _toolStartCount); + + var loser = await Task.WhenAny(first, second); + await Assert.ThrowsAsync(() => loser); + + // The winning execution is still the one recorded in the store. + await client.CancelTaskAsync(augmented.TaskCreated!.TaskId, cancellationToken); + var task = await PollUntilTerminalAsync(client, augmented.TaskCreated!.TaskId, cancellationToken); + Assert.IsType(task); + Assert.Equal(1, _toolStartCount); + } + + private static async Task PollUntilTerminalAsync( + McpClient client, string taskId, CancellationToken cancellationToken) + { + while (true) + { + cancellationToken.ThrowIfCancellationRequested(); + var task = await client.GetTaskAsync(taskId, cancellationToken); + if (task is not WorkingTaskResult) + { + return task; + } + + await Task.Delay(10, cancellationToken); + } + } + + private sealed class CallbackTaskExecutor(McpServerTaskExecutorTests test) : IMcpTaskExecutor + { + public async ValueTask StartAsync(McpTaskExecutionContext context, CancellationToken cancellationToken) + { + test._executorInvoked.TrySetResult(context); + context.CancellationToken.Register(() => test._executorCancelled.TrySetResult(true)); + + // Materialize a scoped dependency so the tests can observe scope disposal, the way + // an external executor reads scope-bound services before handing work off. + _ = context.Request.Services!.GetRequiredService(); + + if (test._startException is { } exception) + { + throw exception; + } + + if (test._runPipelineLocally) + { + await context.RunToolPipelineAsync(context.CancellationToken).ConfigureAwait(false); + } + } + } + + private sealed class ScopedDependency(TaskCompletionSource disposed) : IAsyncDisposable + { + public ValueTask DisposeAsync() + { + disposed.TrySetResult(true); + return default; + } + } +} + +public class McpServerTaskExecutorDiResolutionTests : ClientServerTestBase +{ + private readonly TaskCompletionSource _executorInvoked = new(TaskCreationOptions.RunContinuationsAsynchronously); + private int _executorInstances; + + public McpServerTaskExecutorDiResolutionTests(ITestOutputHelper testOutputHelper) : base(testOutputHelper) + { +#if !NET + Assert.SkipWhen(RuntimeInformation.IsOSPlatform(OSPlatform.Windows), "https://github.com/modelcontextprotocol/csharp-sdk/issues/587"); +#endif + } + + protected override void ConfigureServices(ServiceCollection services, IMcpServerBuilder mcpServerBuilder) + { + services.AddScoped(_ => + { + Interlocked.Increment(ref _executorInstances); + return new DiTaskExecutor(this); + }); + + mcpServerBuilder + .WithTasks(new InMemoryMcpTaskStore { DefaultPollIntervalMs = 10 }) + .WithTools([McpServerTool.Create( + () => "local result", + new McpServerToolCreateOptions { Name = "local-tool" })]); + } + + [Fact] + public async Task ScopedExecutor_RegisteredInDi_IsResolvedPerTaskAndRunsPipeline() + { + await using var client = await CreateMcpClientForServer(); + var cancellationToken = TestContext.Current.CancellationToken; + + var augmented = await client.CallToolAsTaskAsync( + new CallToolRequestParams { Name = "local-tool" }, + cancellationToken); + Assert.True(augmented.IsTask); + + await _executorInvoked.Task.WaitAsync(TestConstants.DefaultTimeout, cancellationToken); + var task = await PollUntilTerminalAsync(client, augmented.TaskCreated!.TaskId, cancellationToken); + Assert.IsType(task); + Assert.Equal(1, _executorInstances); + + // A second task resolves a fresh scoped executor instance. + var second = await client.CallToolAsTaskAsync( + new CallToolRequestParams { Name = "local-tool" }, + cancellationToken); + Assert.True(second.IsTask); + var secondTask = await PollUntilTerminalAsync(client, second.TaskCreated!.TaskId, cancellationToken); + Assert.IsType(secondTask); + Assert.Equal(2, _executorInstances); + } + + private static async Task PollUntilTerminalAsync( + McpClient client, string taskId, CancellationToken cancellationToken) + { + while (true) + { + cancellationToken.ThrowIfCancellationRequested(); + var task = await client.GetTaskAsync(taskId, cancellationToken); + if (task is not WorkingTaskResult) + { + return task; + } + + await Task.Delay(10, cancellationToken); + } + } + + private sealed class DiTaskExecutor(McpServerTaskExecutorDiResolutionTests test) : IMcpTaskExecutor + { + public ValueTask StartAsync(McpTaskExecutionContext context, CancellationToken cancellationToken) + { + test._executorInvoked.TrySetResult(context); + return context.RunToolPipelineAsync(context.CancellationToken); + } + } +}