diff --git a/docs/concepts/tasks/tasks.md b/docs/concepts/tasks/tasks.md index c1ab0c23a..756653c0d 100644 --- a/docs/concepts/tasks/tasks.md +++ b/docs/concepts/tasks/tasks.md @@ -246,6 +246,14 @@ requirements drawn from the SEP and the SDK contract: resolves immediately — even from a different process or node. Stores backed by eventually consistent storage must wait for the write to become visible (quorum acknowledgement, write-through, etc.) before returning. Required by SEP-2663 §306. + When the call includes a non-null `executionIntent`, persist it atomically with the task + record so it can be read back via + ; copy the + `JsonElement` (for example with `Clone()`) instead of retaining the executor's original + backing document — the executor may dispose that document once execution is handed off, + and a retained reference surfaces later as an `ObjectDisposedException`. The intent is + server-only — never surfaced in protocol responses, notifications, or errors — and a + store that cannot persist it must throw rather than silently dropping it. 5. **Singleton under stateless HTTP** — when the server runs in stateless mode (each request spins up a fresh server instance), the same `IMcpTaskStore` instance must be shared across requests — either by registering it as a singleton in DI, or by backing it with external @@ -282,6 +290,232 @@ 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's method +is invoked after the task record is durably created in the store, and 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, creating the +execution intent, or + — fail the +original `tools/call` with an error result, and nothing is persisted: no task record is left +orphaned at `Working`, and no intent exists without its task. +After a successful `StartAsync`, the SDK stops tracking the task and the store is the single +source of truth for its state. + +#### Persisted execution intent + +There is still a crash 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 +durably submitted to the external runtime. To close that window, executors that delegate to an +external runtime also implement +. +The full ordering is: + +1. Authorization and validation (the filters registered before Tasks) run. +2. The SDK calls `CreateExecutionIntentAsync`, which returns a portable, side-effect-free + description of how the task will be started — or `null` for stateless executors like + , which run tools + in-process and have nothing to reconstruct. +3. The SDK passes the intent to + , which persists + the task record and the intent atomically, as a single write. +4. The SDK calls `StartAsync` to start the external work. + +The intent contract: + +- **Opaque and executor-owned.** Its schema and versioning belong to the executor, not to the + SDK or the protocol. An intent can outlive the process that wrote it — an orphan recovered + at startup may have been persisted by an earlier deployment of the executor — so version + the payload (for example with a `version` field) so reconciliation code can tell + generations apart and evolve safely. +- **Portable.** Only data the external runtime needs to reconstruct the submission — no + runtime objects, services, credentials, clients, transports, or delegates. The intent is + persisted in the task store alongside the task record, so treat it like any other durable + data: anything with store access can read it, and secrets must stay out of it. +- **Server-only.** It never surfaces in MCP responses, notifications, or errors; it is + readable only through the store's + property (and + for + the task's executor). +- **Copy and reject in the store.** Stores must copy the `JsonElement` (for example with + `Clone()`) rather than retaining a reference to the executor's backing document — the + executor may dispose that document once execution is handed off, and a retained reference + surfaces later as an `ObjectDisposedException`. A store that cannot persist a non-null + intent must throw, rejecting the task creation, rather than silently dropping it. + +The intent captures the *submission*, not the execution. The DI execution scope and its +request-scoped services, the matched tool primitive, the context token wired to +`tasks/cancel`, and the task's input-request channel for elicitation and sampling all die +with the crashed process. A recovered execution therefore behaves like any other external +worker: it runs in the external runtime, records progress and results in the store, and has +no +to fall back on. Cancellation for a recovered task cannot reach the dead process's token, so +it must be propagated through the external runtime or the store, as with any cross-process +execution; multi-round-trip input relies on the store's + +event. + +```csharp +public sealed class TemporalTaskExecutor(ITemporalClient workflowClient) : IMcpTaskExecutor +{ + public ValueTask CreateExecutionIntentAsync( + RequestContext request, CancellationToken cancellationToken) + { + // Portable, versioned data only: what the workflow needs to reconstruct the + // submission. No services, clients, transports, credentials, or delegates — + // this element must survive a process restart inside the task store. + JsonObject intent = new() + { + ["version"] = 1, + ["request"] = JsonSerializer.SerializeToNode( + request.Params, + McpJsonUtilities.DefaultOptions.GetTypeInfo()), + }; + + return ValueTask.FromResult( + JsonSerializer.SerializeToElement( + intent, + McpJsonUtilities.DefaultOptions.GetTypeInfo())); + } + + 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(); + } +} +``` + +#### Recovering orphaned tasks + +Because the task record and its intent are persisted together, a restarting integration can +discover orphaned `Working` tasks through its own durable store, read the persisted intent, +and reconstruct the submission. Reconciliation must be safe at either crash point: if the +external runtime never accepted the job, resubmission starts it; if the runtime *had* +accepted the job before the crash, resubmitting with the same task ID deduplicates — most +external runtimes treat the ID as a unique workflow or job key. Either way, re-read the +task's state immediately before resubmitting and leave tasks that already reached a terminal +state — completed, failed, or cancelled by a client while the process was down — untouched. + +Until reconciliation runs, clients polling an orphan observe `Working` for as long as the +task's TTL permits; reconciliation is what moves the orphan to a terminal state. An orphan +can also linger in + when the process +died later in execution — the same intent-based reconstruction applies, but the pending +input exchange additionally needs a live `InputResponseReceived` subscriber before it can +resume. + +The SDK performs no reconciliation itself; this is integration code, typically run at +startup. retains the +intent alongside the task record, but like all of its state it does not survive process +restarts — intent-based recovery requires a store backed by durable storage. For tasks +created without an intent (a stateless executor was configured), reconciliation cannot +reconstruct a submission, so integrations fall back to their own strategy — TTL cleanup, +for example. + +```csharp +foreach (var task in await durableStore.FindWorkingTasksAsync()) +{ + if (task.ExecutionIntent is not { } intent) + { + continue; // Created by a stateless executor; nothing to reconstruct. + } + + var payload = JsonNode.Parse(intent.GetRawText())!; + if (payload["version"]?.GetValue() is not 1) + { + continue; // Unknown intent generation: skip or migrate explicitly. + } + + var requestParams = payload["request"]!.Deserialize( + McpJsonUtilities.DefaultOptions.GetTypeInfo())!; + + // Re-check the task's state if discovery ran earlier — terminal tasks stay + // untouched — then resubmit. The workflow ID is the task ID: a deduplicated + // no-op if the job was already accepted before the crash, a reconstruction + // if it never was. + await workflowClient.StartWorkflowAsync( + "run-mcp-task", + new McpTaskPayload(task.TaskId, requestParams), + id: task.TaskId, + cancellationToken); +} +``` + +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. + +`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 @@ -350,6 +584,10 @@ compatibility bridge for the previous experimental API. - **Server-push task status notifications (SEP-2575)**: not yet implemented. Clients rely on polling exclusively. +- **Orphaned-task reconciliation**: the SDK persists the executor's execution intent with the + task record, but discovering orphaned tasks and resubmitting them is integration code; the + SDK performs no reconciliation itself (see + [Persisted execution intent](#persisted-execution-intent)). - **Lazy task creation**: when a tool runs through the task store, the store's is invoked eagerly before the inner handler runs, so tools that complete inline still incur a store write. There is diff --git a/src/ModelContextProtocol.Extensions.Tasks/Server/IMcpTaskExecutor.cs b/src/ModelContextProtocol.Extensions.Tasks/Server/IMcpTaskExecutor.cs new file mode 100644 index 000000000..8f9841588 --- /dev/null +++ b/src/ModelContextProtocol.Extensions.Tasks/Server/IMcpTaskExecutor.cs @@ -0,0 +1,100 @@ +using ModelContextProtocol.Protocol; +using ModelContextProtocol.Server; +using System.Text.Json; + +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. +/// +/// +/// Executors that delegate to an external runtime should also implement +/// so a persisted intent is available for +/// recovering tasks whose process exited between task creation and a completed start. +/// See that method for the intent contract. +/// +/// +/// See the SEP-2663 +/// specification for details on the tasks extension. +/// +/// +public interface IMcpTaskExecutor +{ + /// + /// Creates a portable, side-effect-free description of how the task will be started, + /// which the SDK persists atomically with the task record. + /// + /// The tool request that is being converted into a task. + /// Cancellation token for the operation. + /// + /// An opaque carrying the execution intent, or + /// for stateless executors (such as ) that need no + /// recovery metadata. + /// + /// + /// + /// The intent closes the crash window between task creation and start: the store persists it + /// atomically alongside the task record, so if the process exits before + /// completes, an integration can discover the orphaned task through + /// its own durable store and reconstruct the submission from the persisted intent — using the + /// task ID as an idempotency key so a resubmission is safe. + /// + /// + /// The intent is executor-owned and opaque to the SDK: its schema and versioning belong to + /// the executor, not to the SDK or the protocol. It must be portable — only data the external + /// runtime needs to reconstruct the submission, such as the tool name, arguments, and + /// executor-specific routing hints. It must not capture runtime objects, services, + /// credentials, clients, transports, or delegates. Because the intent is persisted and may + /// be recovered by a later deployment of the executor, embed a version marker in the payload + /// so reconciliation code can distinguish generations and evolve safely. + /// + /// + /// This method must be side-effect-free: no external submission or state mutation happens + /// here. It runs after authorization and validation, but before the task record is created. + /// If it throws, no task is created and the exception fails the original tools/call. + /// + /// + /// The persisted intent is server-only: it never surfaces in MCP responses, notifications, + /// or errors. + /// + /// + ValueTask CreateExecutionIntentAsync( + RequestContext request, + CancellationToken cancellationToken); + + /// + /// 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/IMcpTaskStore.cs b/src/ModelContextProtocol.Extensions.Tasks/Server/IMcpTaskStore.cs index 6851f21ac..f664a5e36 100644 --- a/src/ModelContextProtocol.Extensions.Tasks/Server/IMcpTaskStore.cs +++ b/src/ModelContextProtocol.Extensions.Tasks/Server/IMcpTaskStore.cs @@ -35,6 +35,11 @@ public interface IMcpTaskStore /// /// Creates a new task for tracking an asynchronous execution. /// + /// + /// The executor-owned execution intent to persist atomically with the task record, or + /// when the task's executor is stateless and needs no recovery + /// metadata. + /// /// Cancellation token for the operation. /// /// A with a unique task ID, initial status of , @@ -54,8 +59,22 @@ public interface IMcpTaskStore /// write to be visible (e.g., quorum acknowledgement, write-through, or an equivalent /// barrier) before returning. /// + /// + /// When is non-null, it MUST be persisted atomically with + /// the task record — a later must be able to return it via + /// — and implementations MUST copy the + /// (e.g., via ) rather than retaining + /// a reference to the executor's original backing document, which the executor may dispose + /// once execution has been handed off; a retained reference surfaces later as an + /// . The intent is server-only and MUST NOT surface in + /// protocol responses, notifications, or errors. Implementations that cannot persist a + /// non-null intent MUST reject it by throwing from this method, rather than silently + /// dropping it, so no task is created without its recovery metadata. + /// /// - Task CreateTaskAsync(CancellationToken cancellationToken = default); + Task CreateTaskAsync( + JsonElement? executionIntent = null, + CancellationToken cancellationToken = default); /// /// Retrieves the current state of a task. diff --git a/src/ModelContextProtocol.Extensions.Tasks/Server/InMemoryMcpTaskStore.cs b/src/ModelContextProtocol.Extensions.Tasks/Server/InMemoryMcpTaskStore.cs index d1014c110..10555f299 100644 --- a/src/ModelContextProtocol.Extensions.Tasks/Server/InMemoryMcpTaskStore.cs +++ b/src/ModelContextProtocol.Extensions.Tasks/Server/InMemoryMcpTaskStore.cs @@ -24,6 +24,11 @@ namespace ModelContextProtocol.Extensions.Tasks; /// For production scenarios requiring durability, session isolation, or more advanced retention /// policies, implement a custom . /// +/// +/// The execution intent supplied to is copied and retained with the +/// task record, available via . Like all state in this +/// store, it does not survive process restarts. +/// /// public class InMemoryMcpTaskStore : IMcpTaskStore { @@ -47,14 +52,21 @@ public class InMemoryMcpTaskStore : IMcpTaskStore public TimeSpan? DefaultTimeToLive { get; set; } /// - public Task CreateTaskAsync(CancellationToken cancellationToken = default) + public Task CreateTaskAsync( + JsonElement? executionIntent = null, + CancellationToken cancellationToken = default) { var now = DateTimeOffset.UtcNow; SweepExpired(now); var taskId = Guid.NewGuid().ToString("N"); - var info = new McpTaskInfo(taskId, McpTaskStatus.Working, now, now, DefaultTimeToLive, DefaultPollIntervalMs); + // Clone the intent so the stored record doesn't hold a reference to the executor's + // original backing document, which the executor may dispose. + var info = new McpTaskInfo(taskId, McpTaskStatus.Working, now, now, DefaultTimeToLive, DefaultPollIntervalMs) + { + ExecutionIntent = executionIntent?.Clone(), + }; _tasks[taskId] = info; return Task.FromResult(info); diff --git a/src/ModelContextProtocol.Extensions.Tasks/Server/McpTaskExecutionContext.cs b/src/ModelContextProtocol.Extensions.Tasks/Server/McpTaskExecutionContext.cs new file mode 100644 index 000000000..720820496 --- /dev/null +++ b/src/ModelContextProtocol.Extensions.Tasks/Server/McpTaskExecutionContext.cs @@ -0,0 +1,132 @@ +using ModelContextProtocol.Protocol; +using ModelContextProtocol.Server; +using System.Text.Json; + +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 executor-owned execution intent recovered from the store for this task, or + /// when the task's executor created no intent. + /// + /// + /// This is the intent returned by + /// as persisted (and copied) by the store. In the normal flow the executor that created the + /// intent already has it; the property is useful to executors that wrap or decorate another + /// executor, and to reconciliation code that reads the recovered intent from the store's + /// after a crash. + /// + public JsonElement? ExecutionIntent => TaskInfo.ExecutionIntent; + + /// + /// 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/McpTaskInfo.cs b/src/ModelContextProtocol.Extensions.Tasks/Server/McpTaskInfo.cs index a9a86ed54..4b567caf6 100644 --- a/src/ModelContextProtocol.Extensions.Tasks/Server/McpTaskInfo.cs +++ b/src/ModelContextProtocol.Extensions.Tasks/Server/McpTaskInfo.cs @@ -23,4 +23,23 @@ public sealed record McpTaskInfo( string? StatusMessage = null, JsonElement? Result = null, JsonElement? Error = null, - IReadOnlyDictionary? InputRequests = null); + IReadOnlyDictionary? InputRequests = null) +{ + /// + /// Gets the executor-owned execution intent persisted with the task, or + /// when the task's executor is stateless. + /// + /// + /// + /// The intent is opaque to the SDK — its schema and versioning belong to the + /// that created it — and exists so that an integration + /// can reconstruct an external submission for a task whose process exited between + /// task creation and a completed start (see + /// ). + /// + /// + /// It is server-only: it never surfaces in MCP responses, notifications, or errors. + /// + /// + public JsonElement? ExecutionIntent { get; init; } +} diff --git a/src/ModelContextProtocol.Extensions.Tasks/Server/McpTasksBuilderExtensions.cs b/src/ModelContextProtocol.Extensions.Tasks/Server/McpTasksBuilderExtensions.cs index e61466a69..ccfeab31e 100644 --- a/src/ModelContextProtocol.Extensions.Tasks/Server/McpTasksBuilderExtensions.cs +++ b/src/ModelContextProtocol.Extensions.Tasks/Server/McpTasksBuilderExtensions.cs @@ -171,9 +171,25 @@ private async ValueTask> RunAsTaskAsync( }; McpTaskInfo taskInfo; + IMcpTaskExecutor executor; try { - taskInfo = await _store.CreateTaskAsync(cancellationToken).ConfigureAwait(false); + // 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; + + // Create the intent first, then persist it atomically with the task record so a + // crash between creation and a completed start can be recovered from the store. + // Both steps precede the task record: a failure here fails tools/call instead of + // leaving a durably-created task orphaned at Working. + JsonElement? executionIntent = await executor + .CreateExecutionIntentAsync(request, cancellationToken).ConfigureAwait(false); + taskInfo = await _store.CreateTaskAsync(executionIntent, cancellationToken).ConfigureAwait(false); } catch { @@ -189,9 +205,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 +266,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..05d430d89 --- /dev/null +++ b/src/ModelContextProtocol.Extensions.Tasks/Server/ProcessLocalMcpTaskExecutor.cs @@ -0,0 +1,53 @@ +using ModelContextProtocol.Protocol; +using ModelContextProtocol.Server; +using System.Text.Json; + +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(); + + /// + /// + /// The process-local executor is stateless: execution runs in this process, so there is + /// no submission to reconstruct and no intent to persist. + /// + public ValueTask CreateExecutionIntentAsync( + RequestContext request, + CancellationToken cancellationToken) => default; + + /// + 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.AspNetCore.Tests/HttpTaskIntegrationTests.cs b/tests/ModelContextProtocol.AspNetCore.Tests/HttpTaskIntegrationTests.cs index d64a55987..feb980d0d 100644 --- a/tests/ModelContextProtocol.AspNetCore.Tests/HttpTaskIntegrationTests.cs +++ b/tests/ModelContextProtocol.AspNetCore.Tests/HttpTaskIntegrationTests.cs @@ -8,6 +8,7 @@ using ModelContextProtocol.Server; using Moq; using System.Security.Claims; +using System.Text.Json; namespace ModelContextProtocol.AspNetCore.Tests; @@ -178,7 +179,7 @@ public async Task WithTasks_UnauthorizedTool_DoesNotCreateTask(bool registerTask Assert.Equal(McpErrorCode.InvalidRequest, exception.ErrorCode); taskStore.Verify( - store => store.CreateTaskAsync(It.IsAny()), + store => store.CreateTaskAsync(It.IsAny(), It.IsAny()), Times.Never); } diff --git a/tests/ModelContextProtocol.Tests/Server/InMemoryMcpTaskStoreTests.cs b/tests/ModelContextProtocol.Tests/Server/InMemoryMcpTaskStoreTests.cs index 236740f61..b1efa4286 100644 --- a/tests/ModelContextProtocol.Tests/Server/InMemoryMcpTaskStoreTests.cs +++ b/tests/ModelContextProtocol.Tests/Server/InMemoryMcpTaskStoreTests.cs @@ -25,7 +25,7 @@ public async Task CreateTaskAsync_ReturnsWorkingTaskWithUniqueId() { var store = new InMemoryMcpTaskStore(); - var result = await store.CreateTaskAsync(CT); + var result = await store.CreateTaskAsync(cancellationToken: CT); Assert.NotNull(result); Assert.NotEmpty(result.TaskId); @@ -39,8 +39,8 @@ public async Task CreateTaskAsync_GeneratesUniqueIds() { var store = new InMemoryMcpTaskStore(); - var task1 = await store.CreateTaskAsync(CT); - var task2 = await store.CreateTaskAsync(CT); + var task1 = await store.CreateTaskAsync(cancellationToken: CT); + var task2 = await store.CreateTaskAsync(cancellationToken: CT); Assert.NotEqual(task1.TaskId, task2.TaskId); } @@ -50,7 +50,7 @@ public async Task CreateTaskAsync_UsesDefaultPollInterval() { var store = new InMemoryMcpTaskStore { DefaultPollIntervalMs = 500 }; - var result = await store.CreateTaskAsync(CT); + var result = await store.CreateTaskAsync(cancellationToken: CT); Assert.Equal(500, result.PollIntervalMs); } @@ -60,16 +60,59 @@ public async Task CreateTaskAsync_UsesDefaultTimeToLive() { var store = new InMemoryMcpTaskStore { DefaultTimeToLive = TimeSpan.FromSeconds(30) }; - var result = await store.CreateTaskAsync(CT); + var result = await store.CreateTaskAsync(cancellationToken: CT); Assert.Equal(TimeSpan.FromSeconds(30), result.TimeToLive); } + [Fact] + public async Task CreateTaskAsync_PersistsExecutionIntent() + { + var store = new InMemoryMcpTaskStore(); + var intent = JsonSerializer.SerializeToElement( + new { queue = "mcp-tasks", tool = "long-running-tool" }, + McpJsonUtilities.DefaultOptions); + + var created = await store.CreateTaskAsync(intent, CT); + var retrieved = await store.GetTaskAsync(created.TaskId, CT); + + Assert.NotNull(retrieved); + Assert.NotNull(retrieved.ExecutionIntent); + Assert.Equal(intent.GetRawText(), retrieved.ExecutionIntent.Value.GetRawText()); + Assert.Equal(intent.GetRawText(), created.ExecutionIntent!.Value.GetRawText()); + } + + [Fact] + public async Task CreateTaskAsync_WithoutExecutionIntent_LeavesIntentNull() + { + var store = new InMemoryMcpTaskStore(); + + var created = await store.CreateTaskAsync(cancellationToken: CT); + + Assert.Null(created.ExecutionIntent); + } + + [Fact] + public async Task CreateTaskAsync_CopiesExecutionIntentElement() + { + const string IntentJson = """{ "queue": "mcp-tasks" }"""; + var store = new InMemoryMcpTaskStore(); + using var document = JsonDocument.Parse(IntentJson); + + var created = await store.CreateTaskAsync(document.RootElement, CT); + document.Dispose(); + + // The store must not retain a reference to the executor's original backing document. + var retrieved = await store.GetTaskAsync(created.TaskId, CT); + Assert.NotNull(retrieved); + Assert.Equal(IntentJson, retrieved.ExecutionIntent!.Value.GetRawText()); + } + [Fact] public async Task GetTaskAsync_ReturnsWorkingTask() { var store = new InMemoryMcpTaskStore(); - var created = await store.CreateTaskAsync(CT); + var created = await store.CreateTaskAsync(cancellationToken: CT); var result = await store.GetTaskAsync(created.TaskId, CT); @@ -92,7 +135,7 @@ public async Task GetTaskAsync_ReturnsNullForUnknownId() public async Task GetTaskAsync_WithinTimeToLive_ReturnsTask() { var store = new InMemoryMcpTaskStore { DefaultTimeToLive = TimeSpan.FromMinutes(10) }; - var created = await store.CreateTaskAsync(CT); + var created = await store.CreateTaskAsync(cancellationToken: CT); var result = await store.GetTaskAsync(created.TaskId, CT); @@ -104,7 +147,7 @@ public async Task GetTaskAsync_WithinTimeToLive_ReturnsTask() public async Task GetTaskAsync_AfterTimeToLiveElapsed_ReturnsNull() { var store = new InMemoryMcpTaskStore { DefaultTimeToLive = TimeSpan.FromMilliseconds(100) }; - var created = await store.CreateTaskAsync(CT); + var created = await store.CreateTaskAsync(cancellationToken: CT); await Task.Delay(TimeSpan.FromMilliseconds(500), CT); @@ -117,7 +160,7 @@ public async Task GetTaskAsync_AfterTimeToLiveElapsed_ReturnsNull() public async Task GetTaskAsync_WithoutTimeToLive_DoesNotExpire() { var store = new InMemoryMcpTaskStore(); - var created = await store.CreateTaskAsync(CT); + var created = await store.CreateTaskAsync(cancellationToken: CT); await Task.Delay(TimeSpan.FromMilliseconds(200), CT); @@ -131,7 +174,7 @@ public async Task GetTaskAsync_WithoutTimeToLive_DoesNotExpire() public async Task GetTaskAsync_WithZeroTimeToLive_DoesNotExpire() { var store = new InMemoryMcpTaskStore { DefaultTimeToLive = TimeSpan.Zero }; - var created = await store.CreateTaskAsync(CT); + var created = await store.CreateTaskAsync(cancellationToken: CT); await Task.Delay(TimeSpan.FromMilliseconds(200), CT); @@ -145,7 +188,7 @@ public async Task GetTaskAsync_WithZeroTimeToLive_DoesNotExpire() public async Task SetCompletedAsync_TransitionsToCompleted() { var store = new InMemoryMcpTaskStore(); - var created = await store.CreateTaskAsync(CT); + var created = await store.CreateTaskAsync(cancellationToken: CT); var resultPayload = JsonDocument.Parse("""{"answer":42}""").RootElement.Clone(); await store.SetCompletedAsync(created.TaskId, resultPayload, CT); @@ -160,7 +203,7 @@ public async Task SetCompletedAsync_TransitionsToCompleted() public async Task SetFailedAsync_TransitionsToFailed() { var store = new InMemoryMcpTaskStore(); - var created = await store.CreateTaskAsync(CT); + var created = await store.CreateTaskAsync(cancellationToken: CT); var errorPayload = JsonDocument.Parse("""{"message":"boom"}""").RootElement.Clone(); await store.SetFailedAsync(created.TaskId, errorPayload, CT); @@ -175,7 +218,7 @@ public async Task SetFailedAsync_TransitionsToFailed() public async Task SetCancelledAsync_TransitionsToCancelled() { var store = new InMemoryMcpTaskStore(); - var created = await store.CreateTaskAsync(CT); + var created = await store.CreateTaskAsync(cancellationToken: CT); var cancelled = await store.SetCancelledAsync(created.TaskId, CT); @@ -189,7 +232,7 @@ public async Task SetCancelledAsync_TransitionsToCancelled() public async Task SetCancelledAsync_ReturnsFalseForTerminalTask() { var store = new InMemoryMcpTaskStore(); - var created = await store.CreateTaskAsync(CT); + var created = await store.CreateTaskAsync(cancellationToken: CT); await store.SetCompletedAsync(created.TaskId, JsonSerializer.SerializeToElement("done", McpJsonUtilities.DefaultOptions), CT); var cancelled = await store.SetCancelledAsync(created.TaskId, CT); @@ -214,7 +257,7 @@ public async Task SetCancelledAsync_ReturnsFalseForUnknownId() public async Task SetInputRequestsAsync_TransitionsToInputRequired() { var store = new InMemoryMcpTaskStore(); - var created = await store.CreateTaskAsync(CT); + var created = await store.CreateTaskAsync(cancellationToken: CT); var requests = new Dictionary { @@ -238,7 +281,7 @@ public async Task SetInputRequestsAsync_TransitionsToInputRequired() public async Task SetInputRequestsAsync_MergesMultipleRequests() { var store = new InMemoryMcpTaskStore(); - var created = await store.CreateTaskAsync(CT); + var created = await store.CreateTaskAsync(cancellationToken: CT); await store.SetInputRequestsAsync(created.TaskId, new Dictionary { @@ -262,7 +305,7 @@ public async Task SetInputRequestsAsync_MergesMultipleRequests() public async Task ResolveInputRequestsAsync_RemovesMatchedRequests() { var store = new InMemoryMcpTaskStore(); - var created = await store.CreateTaskAsync(CT); + var created = await store.CreateTaskAsync(cancellationToken: CT); await store.SetInputRequestsAsync(created.TaskId, new Dictionary { @@ -287,7 +330,7 @@ public async Task ResolveInputRequestsAsync_RemovesMatchedRequests() public async Task ResolveInputRequestsAsync_TransitionsToWorkingWhenAllSatisfied() { var store = new InMemoryMcpTaskStore(); - var created = await store.CreateTaskAsync(CT); + var created = await store.CreateTaskAsync(cancellationToken: CT); await store.SetInputRequestsAsync(created.TaskId, new Dictionary { @@ -317,7 +360,7 @@ await Assert.ThrowsAsync( public async Task ConcurrentUpdates_DoNotLoseData() { var store = new InMemoryMcpTaskStore(); - var created = await store.CreateTaskAsync(CT); + var created = await store.CreateTaskAsync(cancellationToken: CT); var tasks = Enumerable.Range(0, 50).Select(i => store.SetInputRequestsAsync(created.TaskId, new Dictionary @@ -338,7 +381,7 @@ public async Task ConcurrentUpdates_DoNotLoseData() public async Task ResolveInputRequestsAsync_ForExtraKeys_DoesNotThrow() { var store = new InMemoryMcpTaskStore(); - var created = await store.CreateTaskAsync(CT); + var created = await store.CreateTaskAsync(cancellationToken: CT); await store.ResolveInputRequestsAsync(created.TaskId, new Dictionary { @@ -356,7 +399,7 @@ public async Task ResolveInputRequestsAsync_AlreadyResolvedKey_IsNoOp() // SEP-2663: "Each entry key SHOULD be unique across the lifetime of a given task" and // servers should tolerate clients re-sending an inputResponse for an already-resolved key. var store = new InMemoryMcpTaskStore(); - var created = await store.CreateTaskAsync(CT); + var created = await store.CreateTaskAsync(cancellationToken: CT); await store.SetInputRequestsAsync(created.TaskId, new Dictionary { ["a"] = MakeRequest("ask-a"), @@ -407,7 +450,7 @@ public async Task ConcurrentResolveInputRequests_OnDisjointKeys_AllResolveCorrec // Verifies the optimistic-concurrency loop in InMemoryMcpTaskStore handles parallel // tasks/update calls that each resolve a distinct subset of pending input requests. var store = new InMemoryMcpTaskStore(); - var created = await store.CreateTaskAsync(CT); + var created = await store.CreateTaskAsync(cancellationToken: CT); var seed = Enumerable.Range(0, 20).ToDictionary( i => $"req{i}", @@ -432,7 +475,7 @@ public async Task ConcurrentResolveInputRequests_OnDisjointKeys_AllResolveCorrec public async Task SetCompletedAsync_DoesNotOverwriteCancelledTask() { var store = new InMemoryMcpTaskStore(); - var created = await store.CreateTaskAsync(CT); + var created = await store.CreateTaskAsync(cancellationToken: CT); var cancelled = await store.SetCancelledAsync(created.TaskId, CT); Assert.True(cancelled); @@ -453,7 +496,7 @@ await store.SetCompletedAsync( public async Task SetFailedAsync_DoesNotOverwriteCancelledTask() { var store = new InMemoryMcpTaskStore(); - var created = await store.CreateTaskAsync(CT); + var created = await store.CreateTaskAsync(cancellationToken: CT); await store.SetCancelledAsync(created.TaskId, CT); @@ -472,7 +515,7 @@ await store.SetFailedAsync( public async Task SetCompletedAsync_DoesNotOverwriteCompletedTask() { var store = new InMemoryMcpTaskStore(); - var created = await store.CreateTaskAsync(CT); + var created = await store.CreateTaskAsync(cancellationToken: CT); var first = JsonSerializer.SerializeToElement("first", McpJsonUtilities.DefaultOptions); await store.SetCompletedAsync(created.TaskId, first, CT); @@ -491,7 +534,7 @@ public async Task SetCompletedAsync_DoesNotOverwriteCompletedTask() public async Task ResolveInputRequestsAsync_OnTerminalTask_DoesNotResurrect() { var store = new InMemoryMcpTaskStore(); - var created = await store.CreateTaskAsync(CT); + var created = await store.CreateTaskAsync(cancellationToken: CT); await store.SetCompletedAsync( created.TaskId, @@ -513,7 +556,7 @@ await store.SetCompletedAsync( public async Task ResolveInputRequestsAsync_OnTerminalTask_DoesNotFireEvent() { var store = new InMemoryMcpTaskStore(); - var created = await store.CreateTaskAsync(CT); + var created = await store.CreateTaskAsync(cancellationToken: CT); await store.SetCancelledAsync(created.TaskId, CT); @@ -532,7 +575,7 @@ public async Task ResolveInputRequestsAsync_OnTerminalTask_DoesNotFireEvent() public async Task SetInputRequestsAsync_OnTerminalTask_NoOps() { var store = new InMemoryMcpTaskStore(); - var created = await store.CreateTaskAsync(CT); + var created = await store.CreateTaskAsync(cancellationToken: CT); await store.SetCancelledAsync(created.TaskId, CT); diff --git a/tests/ModelContextProtocol.Tests/Server/McpServerTaskExecutorTests.cs b/tests/ModelContextProtocol.Tests/Server/McpServerTaskExecutorTests.cs new file mode 100644 index 000000000..4c3bc919d --- /dev/null +++ b/tests/ModelContextProtocol.Tests/Server/McpServerTaskExecutorTests.cs @@ -0,0 +1,684 @@ +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; +using System.Text.Json.Nodes; + +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, including the persisted execution intent used for crash recovery. +/// +public class McpServerTaskExecutorTests : ClientServerTestBase +{ + private const string IntentMarker = "intent-leak-marker"; + + private readonly InMemoryMcpTaskStore _taskStore = new() { DefaultPollIntervalMs = 10 }; + private readonly List _events = []; + private RecordingTaskStore? _recordingStore; + 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 Exception? _intentException; + private JsonElement? _executionIntent; + private bool _runPipelineLocally; + private int _toolStartCount; + + // The test base class invokes ConfigureServices from its constructor, so the wrapper + // cannot be assigned in this class's constructor body; resolve it lazily instead. + private RecordingTaskStore RecordingStore => _recordingStore ??= new RecordingTaskStore(_taskStore, _events); + + 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( + RecordingStore, + 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); + } + + [Fact] + public async Task CustomExecutor_ExecutionIntent_CreatedBeforeTaskRecordAndPersisted() + { + _executionIntent = MakeIntent(); + + 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); + + // Authorization and validation already ran; then create intent, persist task + intent + // atomically, and start execution — in that order. + Assert.Equal(["intent", "create", "start"], _events); + + // The context exposes the intent recovered from the store, and the store persisted it + // with the task record. + var context = await _executorInvoked.Task.WaitAsync(TestConstants.DefaultTimeout, cancellationToken); + Assert.True(context.ExecutionIntent.HasValue); + Assert.Equal(_executionIntent!.Value.GetRawText(), context.ExecutionIntent.Value.GetRawText()); + + var stored = await _taskStore.GetTaskAsync(context.TaskId, cancellationToken); + Assert.NotNull(stored); + Assert.Equal(_executionIntent.Value.GetRawText(), stored!.ExecutionIntent!.Value.GetRawText()); + } + + [Fact] + public async Task CustomExecutor_ExecutionIntent_NeverSurfacesInProtocolResponses() + { + _executionIntent = MakeIntent(); + + await using var client = await CreateMcpClientForServer(); + var cancellationToken = TestContext.Current.CancellationToken; + + // Drive tools/call over the wire and inspect the raw response so the check covers what + // the server actually sent, not just the deserialized client types. + JsonRpcRequest callRequest = new() + { + Method = RequestMethods.ToolsCall, + Params = JsonSerializer.SerializeToNode( + new CallToolRequestParams { Name = "long-running-tool", Meta = CreateTaskCapabilityMeta() }, + McpJsonUtilities.DefaultOptions.GetTypeInfo()), + }; + JsonRpcResponse callResponse = await client.SendRequestAsync(callRequest, cancellationToken); + var callResult = Assert.IsType(callResponse.Result); + Assert.Equal("task", callResult["resultType"]?.GetValue()); + Assert.DoesNotContain(IntentMarker, callResult.ToJsonString()); + + var taskId = callResult["taskId"]!.GetValue(); + + // Working state: tasks/get must not leak the intent. + var working = await GetTaskOverWireAsync(client, taskId, cancellationToken); + Assert.Equal("working", working["status"]?.GetValue()); + Assert.DoesNotContain(IntentMarker, working.ToJsonString()); + + // Terminal states: complete the task through the store and poll to completion. + var result = JsonSerializer.SerializeToElement( + new CallToolResult { Content = [new TextContentBlock { Text = "external result" }] }, + McpJsonUtilities.DefaultOptions.GetTypeInfo()); + await _taskStore.SetCompletedAsync(taskId, result, cancellationToken); + + while (true) + { + var terminal = await GetTaskOverWireAsync(client, taskId, cancellationToken); + if (terminal["status"]?.GetValue() is not "working") + { + Assert.DoesNotContain(IntentMarker, terminal.ToJsonString()); + break; + } + + await Task.Delay(10, cancellationToken); + } + } + + [Fact] + public async Task CustomExecutor_ExecutionIntent_NeverSurfacesInFailedTaskResult() + { + _executionIntent = MakeIntent(); + _startException = new InvalidOperationException("external runtime unavailable"); + + await using var client = await CreateMcpClientForServer(); + var cancellationToken = TestContext.Current.CancellationToken; + + JsonRpcRequest callRequest = new() + { + Method = RequestMethods.ToolsCall, + Params = JsonSerializer.SerializeToNode( + new CallToolRequestParams { Name = "long-running-tool", Meta = CreateTaskCapabilityMeta() }, + McpJsonUtilities.DefaultOptions.GetTypeInfo()), + }; + JsonRpcResponse callResponse = await client.SendRequestAsync(callRequest, cancellationToken); + var callResult = Assert.IsType(callResponse.Result); + Assert.Equal("task", callResult["resultType"]?.GetValue()); + Assert.DoesNotContain(IntentMarker, callResult.ToJsonString()); + + var taskId = callResult["taskId"]!.GetValue(); + + while (true) + { + var terminal = await GetTaskOverWireAsync(client, taskId, cancellationToken); + if (terminal["status"]?.GetValue() is not "working") + { + Assert.Equal("failed", terminal["status"]?.GetValue()); + Assert.DoesNotContain(IntentMarker, terminal.ToJsonString()); + break; + } + + await Task.Delay(10, cancellationToken); + } + } + + [Fact] + public async Task CustomExecutor_CreateIntentThrows_FailsToolsCallWithoutCreatingTask() + { + _intentException = new InvalidOperationException("intent construction failed"); + + await using var client = await CreateMcpClientForServer(); + var cancellationToken = TestContext.Current.CancellationToken; + + // A failure before the task record exists fails the original tools/call — the caller + // gets an error result, never a task alternate. + var augmented = await client.CallToolAsTaskAsync( + new CallToolRequestParams { Name = "long-running-tool" }, + cancellationToken); + Assert.False(augmented.IsTask); + Assert.True(augmented.Result!.IsError); + + Assert.Equal(["intent"], _events); + Assert.False(_executorInvoked.Task.IsCompleted); + Assert.False(_toolStarted.Task.IsCompleted); + } + + [Fact] + public async Task CustomExecutor_StoreRejectingIntent_FailsToolsCallWithoutStartingExecution() + { + _executionIntent = MakeIntent(); + RecordingStore.RejectExecutionIntent = true; + + await using var client = await CreateMcpClientForServer(); + var cancellationToken = TestContext.Current.CancellationToken; + + // A store that cannot persist the intent must reject it rather than silently dropping + // it; the caller gets an error result, no task is created, and execution never starts. + var augmented = await client.CallToolAsTaskAsync( + new CallToolRequestParams { Name = "long-running-tool" }, + cancellationToken); + Assert.False(augmented.IsTask); + Assert.True(augmented.Result!.IsError); + + Assert.Equal(["intent"], _events); + Assert.False(_executorInvoked.Task.IsCompleted); + Assert.False(_toolStarted.Task.IsCompleted); + } + + [Fact] + public async Task ProcessLocalExecutor_CreateExecutionIntent_ReturnsNull() + { + Assert.Null(await ProcessLocalMcpTaskExecutor.Instance.CreateExecutionIntentAsync( + null!, TestContext.Current.CancellationToken)); + } + + private static JsonElement MakeIntent() => JsonSerializer.SerializeToElement( + new { queue = IntentMarker, tool = "long-running-tool" }, + McpJsonUtilities.DefaultOptions); + + private static JsonObject CreateTaskCapabilityMeta() => new() + { + [MetaKeys.ClientCapabilities] = new JsonObject + { + ["extensions"] = new JsonObject + { + [TasksProtocol.ExtensionId] = new JsonObject(), + }, + }, + }; + + private static async Task GetTaskOverWireAsync( + McpClient client, string taskId, CancellationToken cancellationToken) + { + JsonRpcRequest request = new() + { + Method = TasksProtocol.MethodTasksGet, + Params = JsonSerializer.SerializeToNode( + new GetTaskRequestParams { TaskId = taskId, Meta = CreateTaskCapabilityMeta() }, + McpTasksJsonContext.Default.GetTaskRequestParams), + }; + + JsonRpcResponse response = await client.SendRequestAsync(request, cancellationToken); + return Assert.IsType(response.Result); + } + + 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 ValueTask CreateExecutionIntentAsync( + RequestContext request, CancellationToken cancellationToken) + { + test._events.Add("intent"); + + if (test._intentException is { } intentException) + { + throw intentException; + } + + return ValueTask.FromResult(test._executionIntent); + } + + public async ValueTask StartAsync(McpTaskExecutionContext context, CancellationToken cancellationToken) + { + test._events.Add("start"); + 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; + } + } + + /// + /// Wraps a task store to record when runs, so + /// tests can pin the intent → create → start ordering, and to simulate a store that cannot + /// persist an execution intent. + /// + private sealed class RecordingTaskStore(IMcpTaskStore inner, List events) : IMcpTaskStore + { + public bool RejectExecutionIntent { get; set; } + + public event Action? InputResponseReceived + { + add => inner.InputResponseReceived += value; + remove => inner.InputResponseReceived -= value; + } + + public async Task CreateTaskAsync( + JsonElement? executionIntent = null, + CancellationToken cancellationToken = default) + { + if (RejectExecutionIntent && executionIntent is not null) + { + throw new NotSupportedException("This store does not support persisting an execution intent."); + } + + var info = await inner.CreateTaskAsync(executionIntent, cancellationToken); + events.Add("create"); + return info; + } + + public Task GetTaskAsync(string taskId, CancellationToken cancellationToken = default) + => inner.GetTaskAsync(taskId, cancellationToken); + + public Task SetCompletedAsync(string taskId, JsonElement result, CancellationToken cancellationToken = default) + => inner.SetCompletedAsync(taskId, result, cancellationToken); + + public Task SetFailedAsync(string taskId, JsonElement error, CancellationToken cancellationToken = default) + => inner.SetFailedAsync(taskId, error, cancellationToken); + + public Task SetCancelledAsync(string taskId, CancellationToken cancellationToken = default) + => inner.SetCancelledAsync(taskId, cancellationToken); + + public Task ResolveInputRequestsAsync( + string taskId, + IDictionary inputResponses, + CancellationToken cancellationToken = default) + => inner.ResolveInputRequestsAsync(taskId, inputResponses, cancellationToken); + + public Task SetInputRequestsAsync( + string taskId, + IDictionary inputRequests, + CancellationToken cancellationToken = default) + => inner.SetInputRequestsAsync(taskId, inputRequests, cancellationToken); + } +} + +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 CreateExecutionIntentAsync( + RequestContext request, CancellationToken cancellationToken) + => ValueTask.FromResult(null); + + public ValueTask StartAsync(McpTaskExecutionContext context, CancellationToken cancellationToken) + { + test._executorInvoked.TrySetResult(context); + return context.RunToolPipelineAsync(context.CancellationToken); + } + } +}