From ba952a29a76103f8d60f8993a5f4a23057cf285a Mon Sep 17 00:00:00 2001 From: Robert McLaws <1657085+robertmclaws@users.noreply.github.com> Date: Mon, 26 Oct 2026 06:33:16 -0400 Subject: [PATCH] Fix InMemoryMcpTaskStore TTL clocks, cancel CAS, and missing-task throws (#1963) Split work vs tracking windows, serve-and-drop terminal Gets, insert then sweep, honor CancellationToken, and treat missing mutations as no-ops. --- ...delContextProtocol.Extensions.Tasks.csproj | 8 + .../Server/InMemoryMcpTaskStore.cs | 345 ++++- .../ModelContextProtocol.Tests.csproj | 4 +- .../Server/InMemoryMcpTaskStoreTests.cs | 1222 ++++++++++++++++- .../Server/McpTaskStoreTests.cs | 9 +- .../TaskCancellationIntegrationTests.cs | 16 +- 6 files changed, 1518 insertions(+), 86 deletions(-) diff --git a/src/ModelContextProtocol.Extensions.Tasks/ModelContextProtocol.Extensions.Tasks.csproj b/src/ModelContextProtocol.Extensions.Tasks/ModelContextProtocol.Extensions.Tasks.csproj index cfeef00f4..1c4050348 100644 --- a/src/ModelContextProtocol.Extensions.Tasks/ModelContextProtocol.Extensions.Tasks.csproj +++ b/src/ModelContextProtocol.Extensions.Tasks/ModelContextProtocol.Extensions.Tasks.csproj @@ -43,6 +43,14 @@ + + + + + + + + diff --git a/src/ModelContextProtocol.Extensions.Tasks/Server/InMemoryMcpTaskStore.cs b/src/ModelContextProtocol.Extensions.Tasks/Server/InMemoryMcpTaskStore.cs index d1014c110..f6d31419d 100644 --- a/src/ModelContextProtocol.Extensions.Tasks/Server/InMemoryMcpTaskStore.cs +++ b/src/ModelContextProtocol.Extensions.Tasks/Server/InMemoryMcpTaskStore.cs @@ -15,10 +15,14 @@ namespace ModelContextProtocol.Extensions.Tasks; /// Tasks are not persisted across process restarts. /// /// -/// Tasks created with a are discarded once their time-to-live -/// elapses (as permitted by SEP-2663): an expired task is removed on access, and an opportunistic -/// throttled sweep reclaims expired tasks that are never polled again. Tasks created without a -/// time-to-live are retained until the process exits. +/// Tasks created with a use two clocks (as permitted by SEP-2663): +/// a work window of CreatedAt + TimeToLive during which input and progress are accepted, +/// and a tracking window that keeps a never-terminal snapshot visible until +/// CreatedAt + 2×TimeToLive so a late complete, fail, or cancel can still be recorded. +/// A terminal snapshot stays visible until LastUpdatedAt + TimeToLive, or until +/// serves it and then drops the row. An opportunistic throttled +/// sweep reclaims snapshots nobody collected. Tasks created without a time-to-live are +/// retained until a terminal or until the process exits. /// /// /// For production scenarios requiring durability, session isolation, or more advanced retention @@ -27,9 +31,38 @@ namespace ModelContextProtocol.Extensions.Tasks; /// public class InMemoryMcpTaskStore : IMcpTaskStore { - private readonly ConcurrentDictionary _tasks = new(); - private static readonly long s_sweepIntervalTicks = TimeSpan.FromSeconds(30).Ticks; - private long _lastSweepTicks = DateTimeOffset.UtcNow.UtcTicks; + internal readonly ConcurrentDictionary _tasks = new(); + internal readonly TimeProvider _timeProvider; + internal static readonly long s_sweepIntervalTicks = TimeSpan.FromSeconds(30).Ticks; + internal long _lastSweepTicks; + + /// + /// Initializes a new instance of the class using . + /// + public InMemoryMcpTaskStore() + : this(TimeProvider.System) + { + } + + /// + /// Initializes a new instance of the class. + /// + /// The time provider used for timestamps, time-to-live, and sweep throttling. + /// is . + public InMemoryMcpTaskStore(TimeProvider timeProvider) + { +#if NET + ArgumentNullException.ThrowIfNull(timeProvider); +#else + if (timeProvider is null) + { + throw new ArgumentNullException(nameof(timeProvider)); + } +#endif + + _timeProvider = timeProvider; + _lastSweepTicks = timeProvider.GetUtcNow().UtcTicks; + } /// /// Gets or sets the default poll interval in milliseconds for new tasks. @@ -41,54 +74,79 @@ public class InMemoryMcpTaskStore : IMcpTaskStore /// Gets or sets the default time-to-live for new tasks, or for unlimited. /// /// - /// When set to a positive value, tasks are discarded once this duration elapses from their - /// creation. A or non-positive value keeps tasks until the process exits. + /// When set to a positive value, the work window ends at CreatedAt + TimeToLive and a + /// never-terminal snapshot is discarded at CreatedAt + 2×TimeToLive. A terminal snapshot + /// is discarded at LastUpdatedAt + TimeToLive, or sooner when + /// serves it. A or non-positive value keeps tasks until a terminal get + /// or until the process exits. /// public TimeSpan? DefaultTimeToLive { get; set; } /// public Task CreateTaskAsync(CancellationToken cancellationToken = default) { - var now = DateTimeOffset.UtcNow; - SweepExpired(now); + cancellationToken.ThrowIfCancellationRequested(); + var now = _timeProvider.GetUtcNow(); var taskId = Guid.NewGuid().ToString("N"); - var info = new McpTaskInfo(taskId, McpTaskStatus.Working, now, now, DefaultTimeToLive, DefaultPollIntervalMs); _tasks[taskId] = info; + SweepExpired(now, cancellationToken); + return Task.FromResult(info); } /// public Task GetTaskAsync(string taskId, CancellationToken cancellationToken = default) { - var now = DateTimeOffset.UtcNow; - if (_tasks.TryGetValue(taskId, out var info)) + cancellationToken.ThrowIfCancellationRequested(); + + var now = _timeProvider.GetUtcNow(); + SweepExpired(now, cancellationToken); + + if (!_tasks.TryGetValue(taskId, out var info)) { - if (IsExpired(info, now)) - { - _tasks.TryRemove(taskId, out _); - return Task.FromResult(null); - } + return Task.FromResult(null); + } + + if (HasExpired(info, now)) + { + _tasks.TryRemove(taskId, out _); + return Task.FromResult(null); + } - return Task.FromResult(info); + if (IsTerminal(info.Status)) + { + // The C# client polls tasks/get until it sees a terminal snapshot, then stops. + // Serving that snapshot and dropping the row is the primary prune. + _tasks.TryRemove(taskId, out _); } - return Task.FromResult(null); + return Task.FromResult(info); } /// public Task SetCompletedAsync(string taskId, JsonElement result, CancellationToken cancellationToken = default) { - Update(taskId, entry => IsTerminal(entry.Status) - ? entry - : entry with + cancellationToken.ThrowIfCancellationRequested(); + var now = _timeProvider.GetUtcNow(); + + TryUpdate(taskId, entry => + { + cancellationToken.ThrowIfCancellationRequested(); + if (IsTerminal(entry.Status) || HasExpired(entry, now)) + { + return entry; + } + + return entry with { Status = McpTaskStatus.Completed, Result = result, - LastUpdatedAt = DateTimeOffset.UtcNow, - }); + LastUpdatedAt = now, + }; + }, cancellationToken); return Task.CompletedTask; } @@ -96,14 +154,24 @@ public Task SetCompletedAsync(string taskId, JsonElement result, CancellationTok /// public Task SetFailedAsync(string taskId, JsonElement error, CancellationToken cancellationToken = default) { - Update(taskId, entry => IsTerminal(entry.Status) - ? entry - : entry with + cancellationToken.ThrowIfCancellationRequested(); + var now = _timeProvider.GetUtcNow(); + + TryUpdate(taskId, entry => + { + cancellationToken.ThrowIfCancellationRequested(); + if (IsTerminal(entry.Status) || HasExpired(entry, now)) + { + return entry; + } + + return entry with { Status = McpTaskStatus.Failed, Error = error, - LastUpdatedAt = DateTimeOffset.UtcNow, - }); + LastUpdatedAt = now, + }; + }, cancellationToken); return Task.CompletedTask; } @@ -111,22 +179,25 @@ public Task SetFailedAsync(string taskId, JsonElement error, CancellationToken c /// public Task SetCancelledAsync(string taskId, CancellationToken cancellationToken = default) { - if (!_tasks.TryGetValue(taskId, out var entry) || IsTerminal(entry.Status)) - { - return Task.FromResult(false); - } + cancellationToken.ThrowIfCancellationRequested(); + var now = _timeProvider.GetUtcNow(); bool transitioned = false; - Update(taskId, e => + if (!TryUpdate(taskId, e => { - if (IsTerminal(e.Status)) + cancellationToken.ThrowIfCancellationRequested(); + if (IsTerminal(e.Status) || HasExpired(e, now)) { + transitioned = false; return e; } transitioned = true; - return e with { Status = McpTaskStatus.Cancelled, LastUpdatedAt = DateTimeOffset.UtcNow }; - }); + return e with { Status = McpTaskStatus.Cancelled, LastUpdatedAt = now }; + }, cancellationToken)) + { + return Task.FromResult(false); + } return Task.FromResult(transitioned); } @@ -140,37 +211,39 @@ public Task ResolveInputRequestsAsync( IDictionary inputResponses, CancellationToken cancellationToken = default) { - bool wasTerminal = false; - Update(taskId, entry => + cancellationToken.ThrowIfCancellationRequested(); + var now = _timeProvider.GetUtcNow(); + + bool accepted = false; + TryUpdate(taskId, entry => { - if (IsTerminal(entry.Status)) + cancellationToken.ThrowIfCancellationRequested(); + if (!IsActive(entry, now)) { - wasTerminal = true; + accepted = false; return entry; } - var requests = entry.InputRequests as ImmutableDictionary - ?? entry.InputRequests?.ToImmutableDictionary() - ?? ImmutableDictionary.Empty; - + var requests = GetRequests(entry); foreach (var key in inputResponses.Keys) { requests = requests.Remove(key); } var status = requests.IsEmpty ? McpTaskStatus.Working : entry.Status; + accepted = true; return entry with { InputRequests = requests, Status = status, - LastUpdatedAt = DateTimeOffset.UtcNow, + LastUpdatedAt = now, }; - }); + }, cancellationToken); - if (wasTerminal) + if (!accepted) { - // Drop responses targeting a terminal task — there are no listeners that can act on them. + // Missing, expired, second-window, or already terminal: no listeners can act on these responses. return Task.CompletedTask; } @@ -193,17 +266,18 @@ public Task SetInputRequestsAsync( IDictionary inputRequests, CancellationToken cancellationToken = default) { - Update(taskId, entry => + cancellationToken.ThrowIfCancellationRequested(); + var now = _timeProvider.GetUtcNow(); + + TryUpdate(taskId, entry => { - if (IsTerminal(entry.Status)) + cancellationToken.ThrowIfCancellationRequested(); + if (!IsActive(entry, now)) { return entry; } - var requests = entry.InputRequests as ImmutableDictionary - ?? entry.InputRequests?.ToImmutableDictionary() - ?? ImmutableDictionary.Empty; - + var requests = GetRequests(entry); foreach (var kvp in inputRequests) { requests = requests.SetItem(kvp.Key, kvp.Value); @@ -213,20 +287,133 @@ public Task SetInputRequestsAsync( { InputRequests = requests, Status = McpTaskStatus.InputRequired, - LastUpdatedAt = DateTimeOffset.UtcNow, + LastUpdatedAt = now, }; - }); + }, cancellationToken); return Task.CompletedTask; } - private static bool IsTerminal(McpTaskStatus status) => + /// + /// Returns the task's pending input requests as an . + /// + /// The task snapshot whose requests should be read. + /// + /// 's when it is already an + /// immutable dictionary; a copy when it is a different dictionary type; or an empty dictionary + /// when no requests are stored. + /// + internal static ImmutableDictionary GetRequests(McpTaskInfo entry) => + entry.InputRequests as ImmutableDictionary + ?? entry.InputRequests?.ToImmutableDictionary() + ?? ImmutableDictionary.Empty; + + /// + /// Returns whether is a terminal task status. + /// + /// The status to inspect. + /// + /// when is + /// , , or + /// ; otherwise . + /// + internal static bool IsTerminal(McpTaskStatus status) => status is McpTaskStatus.Completed or McpTaskStatus.Failed or McpTaskStatus.Cancelled; - private static bool IsExpired(McpTaskInfo info, DateTimeOffset now) => - info.TimeToLive is { } ttl && ttl > TimeSpan.Zero && now - info.CreatedAt >= ttl; + /// + /// Returns whether the task is still in its work window and may accept input, progress, or settle. + /// + /// The task snapshot to inspect. + /// The current UTC time from the store's . + /// + /// when the task is terminal. Otherwise when the + /// task has no positive time-to-live, or when is still earlier than + /// CreatedAt + TimeToLive. + /// + internal static bool IsActive(McpTaskInfo info, DateTimeOffset now) + { + if (IsTerminal(info.Status)) + { + return false; + } + + if (!TryGetPositiveTtl(info, out var ttl)) + { + return true; + } + + return now - info.CreatedAt < ttl; + } + + /// + /// Returns whether the snapshot has no remaining reads and should be discarded. + /// + /// The task snapshot to inspect. + /// The current UTC time from the store's . + /// + /// when the task has no positive time-to-live. + /// For a terminal snapshot, when is at or past + /// LastUpdatedAt + TimeToLive. For a never-terminal snapshot, when + /// is at or past CreatedAt + 2×TimeToLive. + /// + internal static bool HasExpired(McpTaskInfo info, DateTimeOffset now) + { + if (!TryGetPositiveTtl(info, out var ttl)) + { + return false; + } + + if (IsTerminal(info.Status)) + { + return now - info.LastUpdatedAt >= ttl; + } + + var elapsed = now - info.CreatedAt; + return elapsed >= ttl && elapsed - ttl >= ttl; + } - private void SweepExpired(DateTimeOffset now) + /// + /// Tries to read a positive time-to-live from . + /// + /// The task snapshot whose time-to-live should be read. + /// + /// When this method returns , the positive time-to-live. + /// When this method returns , . + /// + /// + /// when is a positive duration; + /// otherwise (null, zero, or negative means unlimited retention). + /// + internal static bool TryGetPositiveTtl(McpTaskInfo info, out TimeSpan ttl) + { + if (info.TimeToLive is { } value && value > TimeSpan.Zero) + { + ttl = value; + return true; + } + + ttl = default; + return false; + } + + /// + /// Removes snapshots for which is , at most once + /// per . + /// + /// The current UTC time used for expiry and throttle checks. + /// + /// Observed at each dictionary entry so a cancelled caller does not finish the walk. + /// + /// + /// Create inserts the new snapshot before calling this method, and + /// pulses it so unpolled rows are still reclaimed. Concurrent callers skip the walk when another + /// thread already claimed the interval. Task ids are never reused, so removing by key is safe + /// even if a value was updated after it was judged expired. + /// + /// + /// was cancelled during the walk. + /// + internal void SweepExpired(DateTimeOffset now, CancellationToken cancellationToken) { long last = Interlocked.Read(ref _lastSweepTicks); if (now.UtcTicks - last < s_sweepIntervalTicks) @@ -242,30 +429,48 @@ private void SweepExpired(DateTimeOffset now) foreach (var kvp in _tasks) { - if (IsExpired(kvp.Value, now)) + cancellationToken.ThrowIfCancellationRequested(); + if (HasExpired(kvp.Value, now)) { - // TaskId values are unique GUIDs that are never reused, and CreatedAt/TimeToLive - // are immutable after creation, so an entry judged expired here stays expired. - // Removing by key is therefore safe even if the value was concurrently updated. + // TaskId values are unique GUIDs that are never reused. HasExpired is monotonic + // for a given snapshot (CreatedAt/TimeToLive are immutable; LastUpdatedAt only + // moves forward on settle), so an entry judged expired here stays expired. _tasks.TryRemove(kvp.Key, out _); } } } - private void Update(string taskId, Func transform) + /// + /// Applies to the snapshot for using + /// compare-and-swap until the write succeeds or the key is gone. + /// + /// The task identifier to update. + /// + /// Receives the current snapshot and returns either that same instance (no-op) or a replacement. + /// + /// Observed before each read so a cancelled caller does not commit. + /// + /// when the key existed and the transform was applied or was a no-op; + /// when the key is missing. + /// + /// + /// was cancelled before a commit. + /// + internal bool TryUpdate(string taskId, Func transform, CancellationToken cancellationToken) { SpinWait spin = default; while (true) { + cancellationToken.ThrowIfCancellationRequested(); if (!_tasks.TryGetValue(taskId, out var current)) { - throw new InvalidOperationException($"Task '{taskId}' not found."); + return false; } var updated = transform(current); if (ReferenceEquals(updated, current) || _tasks.TryUpdate(taskId, updated, current)) { - return; + return true; } spin.SpinOnce(); diff --git a/tests/ModelContextProtocol.Tests/ModelContextProtocol.Tests.csproj b/tests/ModelContextProtocol.Tests/ModelContextProtocol.Tests.csproj index 1ac0cbaa7..32f958f37 100644 --- a/tests/ModelContextProtocol.Tests/ModelContextProtocol.Tests.csproj +++ b/tests/ModelContextProtocol.Tests/ModelContextProtocol.Tests.csproj @@ -10,6 +10,8 @@ false true ModelContextProtocol.Tests + true + $(RepoRoot)\Open.snk $(NoWarn);NU1903;NU1902;MCP9006 @@ -35,8 +37,6 @@ - - diff --git a/tests/ModelContextProtocol.Tests/Server/InMemoryMcpTaskStoreTests.cs b/tests/ModelContextProtocol.Tests/Server/InMemoryMcpTaskStoreTests.cs index 236740f61..bd3155096 100644 --- a/tests/ModelContextProtocol.Tests/Server/InMemoryMcpTaskStoreTests.cs +++ b/tests/ModelContextProtocol.Tests/Server/InMemoryMcpTaskStoreTests.cs @@ -1,6 +1,8 @@ +using Microsoft.Extensions.Time.Testing; using ModelContextProtocol.Extensions.Tasks; using ModelContextProtocol.Protocol; -using ModelContextProtocol.Server; +using System.Collections.Concurrent; +using System.Collections.Immutable; using System.Text.Json; #pragma warning disable MCPEXP001 @@ -12,6 +14,10 @@ namespace ModelContextProtocol.Tests.Server; /// public class InMemoryMcpTaskStoreTests { + private static readonly DateTimeOffset s_start = new(2026, 1, 1, 0, 0, 0, TimeSpan.Zero); + private static readonly TimeSpan s_ttl = TimeSpan.FromSeconds(10); + private static readonly TimeSpan s_sweepInterval = TimeSpan.FromSeconds(30); + private CancellationToken CT => TestContext.Current.CancellationToken; private static InputRequest MakeRequest(string payload) => @@ -20,6 +26,25 @@ private static InputRequest MakeRequest(string payload) => private static InputResponse MakeResponse(string payload) => new() { RawValue = JsonSerializer.SerializeToElement(payload, McpJsonUtilities.DefaultOptions) }; + private static JsonElement DonePayload() => + JsonSerializer.SerializeToElement("done", McpJsonUtilities.DefaultOptions); + + private static JsonElement ErrorPayload() => + JsonElement.Parse("""{"message":"boom"}"""); + + private static FakeTimeProvider CreateClock() => new(s_start); + + private static InMemoryMcpTaskStore CreateStore(FakeTimeProvider clock, TimeSpan? ttl = null) => + new(clock) { DefaultTimeToLive = ttl }; + + private static McpTaskInfo Info( + McpTaskStatus status = McpTaskStatus.Working, + DateTimeOffset? createdAt = null, + DateTimeOffset? lastUpdatedAt = null, + TimeSpan? ttl = null, + IReadOnlyDictionary? inputRequests = null) => + new("task", status, createdAt ?? s_start, lastUpdatedAt ?? createdAt ?? s_start, ttl, InputRequests: inputRequests); + [Fact] public async Task CreateTaskAsync_ReturnsWorkingTaskWithUniqueId() { @@ -305,12 +330,13 @@ public async Task ResolveInputRequestsAsync_TransitionsToWorkingWhenAllSatisfied } [Fact] - public async Task SetCompletedAsync_ThrowsForUnknownTask() + public async Task SetCompletedAsync_UnknownTask_IsNoOp() { var store = new InMemoryMcpTaskStore(); - await Assert.ThrowsAsync( - () => store.SetCompletedAsync("nonexistent", JsonSerializer.SerializeToElement("x", McpJsonUtilities.DefaultOptions), CT)); + await store.SetCompletedAsync("nonexistent", JsonSerializer.SerializeToElement("x", McpJsonUtilities.DefaultOptions), CT); + + Assert.Null(await store.GetTaskAsync("nonexistent", CT)); } [Fact] @@ -546,4 +572,1192 @@ public async Task SetInputRequestsAsync_OnTerminalTask_NoOps() Assert.Equal(McpTaskStatus.Cancelled, task.Status); Assert.True(task.InputRequests is null || task.InputRequests.Count == 0); } + + [Fact] + public void Constructor_NullTimeProvider_Throws() + { + Assert.Throws("timeProvider", () => new InMemoryMcpTaskStore(null!)); + } + + [Fact] + public async Task Constructor_TimeProvider_IsUsedForTimestamps() + { + var clock = CreateClock(); + var store = CreateStore(clock); + + var created = await store.CreateTaskAsync(CT); + + Assert.Equal(s_start, created.CreatedAt); + Assert.Equal(s_start, created.LastUpdatedAt); + } + + [Fact] + public async Task GetTaskAsync_AfterOneTtlWhileWorking_StillReturnsTask() + { + var clock = CreateClock(); + var store = CreateStore(clock, s_ttl); + var created = await store.CreateTaskAsync(CT); + + clock.Advance(s_ttl); + + var result = await store.GetTaskAsync(created.TaskId, CT); + + Assert.NotNull(result); + Assert.Equal(McpTaskStatus.Working, result.Status); + Assert.Contains(created.TaskId, store._tasks.Keys); + } + + [Fact] + public async Task GetTaskAsync_AfterTwoTtlWhileWorking_ReturnsNull() + { + var clock = CreateClock(); + var store = CreateStore(clock, s_ttl); + var created = await store.CreateTaskAsync(CT); + + clock.Advance(s_ttl + s_ttl); + + var result = await store.GetTaskAsync(created.TaskId, CT); + + Assert.Null(result); + Assert.False(store._tasks.ContainsKey(created.TaskId)); + } + + [Fact] + public async Task GetTaskAsync_TerminalSnapshot_ServeAndDrop() + { + var store = new InMemoryMcpTaskStore(); + var created = await store.CreateTaskAsync(CT); + await store.SetCompletedAsync(created.TaskId, DonePayload(), CT); + + var first = await store.GetTaskAsync(created.TaskId, CT); + Assert.NotNull(first); + Assert.Equal(McpTaskStatus.Completed, first.Status); + Assert.Equal("done", first.Result!.Value.GetString()); + + var second = await store.GetTaskAsync(created.TaskId, CT); + Assert.Null(second); + Assert.False(store._tasks.ContainsKey(created.TaskId)); + } + + [Fact] + public async Task GetTaskAsync_FailedSnapshot_ServeAndDrop() + { + var store = new InMemoryMcpTaskStore(); + var created = await store.CreateTaskAsync(CT); + await store.SetFailedAsync(created.TaskId, ErrorPayload(), CT); + + var first = await store.GetTaskAsync(created.TaskId, CT); + Assert.NotNull(first); + Assert.Equal(McpTaskStatus.Failed, first.Status); + + Assert.Null(await store.GetTaskAsync(created.TaskId, CT)); + } + + [Fact] + public async Task GetTaskAsync_CancelledSnapshot_ServeAndDrop() + { + var store = new InMemoryMcpTaskStore(); + var created = await store.CreateTaskAsync(CT); + Assert.True(await store.SetCancelledAsync(created.TaskId, CT)); + + var first = await store.GetTaskAsync(created.TaskId, CT); + Assert.NotNull(first); + Assert.Equal(McpTaskStatus.Cancelled, first.Status); + + Assert.Null(await store.GetTaskAsync(created.TaskId, CT)); + } + + [Fact] + public async Task GetTaskAsync_InputRequired_DoesNotDrop() + { + var store = new InMemoryMcpTaskStore(); + var created = await store.CreateTaskAsync(CT); + await store.SetInputRequestsAsync(created.TaskId, new Dictionary + { + ["req1"] = MakeRequest("ask"), + }, CT); + + var first = await store.GetTaskAsync(created.TaskId, CT); + var second = await store.GetTaskAsync(created.TaskId, CT); + + Assert.NotNull(first); + Assert.NotNull(second); + Assert.Equal(McpTaskStatus.InputRequired, second.Status); + } + + [Fact] + public async Task GetTaskAsync_ExpiredTerminal_ReturnsNullWithoutServing() + { + var clock = CreateClock(); + var store = CreateStore(clock, s_ttl); + var created = await store.CreateTaskAsync(CT); + await store.SetCompletedAsync(created.TaskId, DonePayload(), CT); + + clock.Advance(s_ttl); + + var result = await store.GetTaskAsync(created.TaskId, CT); + + Assert.Null(result); + Assert.False(store._tasks.ContainsKey(created.TaskId)); + } + + [Fact] + public async Task GetTaskAsync_TerminalWithinLastUpdatedTtl_IsServed() + { + var clock = CreateClock(); + var store = CreateStore(clock, s_ttl); + var created = await store.CreateTaskAsync(CT); + clock.Advance(TimeSpan.FromSeconds(5)); + await store.SetCompletedAsync(created.TaskId, DonePayload(), CT); + clock.Advance(TimeSpan.FromSeconds(9)); + + var result = await store.GetTaskAsync(created.TaskId, CT); + + Assert.NotNull(result); + Assert.Equal(McpTaskStatus.Completed, result.Status); + Assert.Equal("done", result.Result!.Value.GetString()); + } + + [Fact] + public async Task SetCompletedAsync_InSecondWindow_IsAllowed() + { + var clock = CreateClock(); + var store = CreateStore(clock, s_ttl); + var created = await store.CreateTaskAsync(CT); + + clock.Advance(s_ttl + TimeSpan.FromSeconds(1)); + await store.SetCompletedAsync(created.TaskId, DonePayload(), CT); + + var task = await store.GetTaskAsync(created.TaskId, CT); + Assert.NotNull(task); + Assert.Equal(McpTaskStatus.Completed, task.Status); + Assert.Equal("done", task.Result!.Value.GetString()); + } + + [Fact] + public async Task SetFailedAsync_InSecondWindow_IsAllowed() + { + var clock = CreateClock(); + var store = CreateStore(clock, s_ttl); + var created = await store.CreateTaskAsync(CT); + + clock.Advance(s_ttl + TimeSpan.FromSeconds(1)); + await store.SetFailedAsync(created.TaskId, ErrorPayload(), CT); + + var task = await store.GetTaskAsync(created.TaskId, CT); + Assert.NotNull(task); + Assert.Equal(McpTaskStatus.Failed, task.Status); + } + + [Fact] + public async Task SetCancelledAsync_InSecondWindow_IsAllowed() + { + var clock = CreateClock(); + var store = CreateStore(clock, s_ttl); + var created = await store.CreateTaskAsync(CT); + + clock.Advance(s_ttl + TimeSpan.FromSeconds(1)); + var cancelled = await store.SetCancelledAsync(created.TaskId, CT); + + Assert.True(cancelled); + var task = await store.GetTaskAsync(created.TaskId, CT); + Assert.NotNull(task); + Assert.Equal(McpTaskStatus.Cancelled, task.Status); + } + + [Fact] + public async Task SetInputRequestsAsync_InSecondWindow_IsRefused() + { + var clock = CreateClock(); + var store = CreateStore(clock, s_ttl); + var created = await store.CreateTaskAsync(CT); + + clock.Advance(s_ttl + TimeSpan.FromSeconds(1)); + await store.SetInputRequestsAsync(created.TaskId, new Dictionary + { + ["req1"] = MakeRequest("late"), + }, CT); + + var task = await store.GetTaskAsync(created.TaskId, CT); + Assert.NotNull(task); + Assert.Equal(McpTaskStatus.Working, task.Status); + Assert.True(task.InputRequests is null || task.InputRequests.Count == 0); + } + + [Fact] + public async Task ResolveInputRequestsAsync_InSecondWindow_IsRefusedAndDoesNotFireEvent() + { + var clock = CreateClock(); + var store = CreateStore(clock, s_ttl); + var created = await store.CreateTaskAsync(CT); + await store.SetInputRequestsAsync(created.TaskId, new Dictionary + { + ["req1"] = MakeRequest("ask"), + }, CT); + + clock.Advance(s_ttl + TimeSpan.FromSeconds(1)); + + int eventCount = 0; + store.InputResponseReceived += _ => Interlocked.Increment(ref eventCount); + + await store.ResolveInputRequestsAsync(created.TaskId, new Dictionary + { + ["req1"] = MakeResponse("late"), + }, CT); + + Assert.Equal(0, eventCount); + var task = await store.GetTaskAsync(created.TaskId, CT); + Assert.NotNull(task); + Assert.Equal(McpTaskStatus.InputRequired, task.Status); + Assert.True(task.InputRequests!.ContainsKey("req1")); + } + + [Fact] + public async Task Mutations_AfterHasExpired_AreNoOps() + { + var clock = CreateClock(); + var store = CreateStore(clock, s_ttl); + var created = await store.CreateTaskAsync(CT); + + clock.Advance(s_ttl + s_ttl); + + await store.SetCompletedAsync(created.TaskId, DonePayload(), CT); + await store.SetFailedAsync(created.TaskId, ErrorPayload(), CT); + Assert.False(await store.SetCancelledAsync(created.TaskId, CT)); + await store.SetInputRequestsAsync(created.TaskId, new Dictionary + { + ["req1"] = MakeRequest("late"), + }, CT); + + int eventCount = 0; + store.InputResponseReceived += _ => Interlocked.Increment(ref eventCount); + await store.ResolveInputRequestsAsync(created.TaskId, new Dictionary + { + ["req1"] = MakeResponse("late"), + }, CT); + + Assert.Equal(0, eventCount); + Assert.Null(await store.GetTaskAsync(created.TaskId, CT)); + } + + [Fact] + public async Task Mutations_UnknownTask_AreNoOps() + { + var store = new InMemoryMcpTaskStore(); + + await store.SetFailedAsync("missing", ErrorPayload(), CT); + await store.SetInputRequestsAsync("missing", new Dictionary + { + ["req1"] = MakeRequest("x"), + }, CT); + + int eventCount = 0; + store.InputResponseReceived += _ => Interlocked.Increment(ref eventCount); + await store.ResolveInputRequestsAsync("missing", new Dictionary + { + ["req1"] = MakeResponse("x"), + }, CT); + + Assert.Equal(0, eventCount); + Assert.Null(await store.GetTaskAsync("missing", CT)); + } + + [Fact] + public async Task GetTaskAsync_WithNegativeTimeToLive_DoesNotExpire() + { + var clock = CreateClock(); + var store = CreateStore(clock, TimeSpan.FromSeconds(-5)); + var created = await store.CreateTaskAsync(CT); + + clock.Advance(TimeSpan.FromDays(1)); + + var result = await store.GetTaskAsync(created.TaskId, CT); + Assert.NotNull(result); + Assert.Equal(created.TaskId, result.TaskId); + } + + [Fact] + public async Task DefaultTimeToLive_ChangeAfterCreate_DoesNotAffectExistingTask() + { + var clock = CreateClock(); + var store = CreateStore(clock, s_ttl); + var created = await store.CreateTaskAsync(CT); + + store.DefaultTimeToLive = TimeSpan.FromMilliseconds(1); + clock.Advance(TimeSpan.FromSeconds(5)); + + var result = await store.GetTaskAsync(created.TaskId, CT); + Assert.NotNull(result); + Assert.Equal(s_ttl, result.TimeToLive); + } + + [Fact] + public async Task SetCompletedAsync_UpdatesLastUpdatedAtFromTimeProvider() + { + var clock = CreateClock(); + var store = CreateStore(clock); + var created = await store.CreateTaskAsync(CT); + + clock.Advance(TimeSpan.FromSeconds(3)); + await store.SetCompletedAsync(created.TaskId, DonePayload(), CT); + + var task = await store.GetTaskAsync(created.TaskId, CT); + Assert.NotNull(task); + Assert.Equal(s_start, task.CreatedAt); + Assert.Equal(s_start + TimeSpan.FromSeconds(3), task.LastUpdatedAt); + } + + [Fact] + public async Task CreateTaskAsync_InsertsBeforeSweep_NewTaskSurvives() + { + var clock = CreateClock(); + var store = CreateStore(clock, s_ttl); + var victim = await store.CreateTaskAsync(CT); + + clock.Advance(s_ttl + s_ttl + s_sweepInterval); + var survivor = await store.CreateTaskAsync(CT); + + Assert.False(store._tasks.ContainsKey(victim.TaskId)); + Assert.True(store._tasks.ContainsKey(survivor.TaskId)); + Assert.NotNull(await store.GetTaskAsync(survivor.TaskId, CT)); + } + + [Fact] + public async Task GetTaskAsync_PulsesSweep_RemovesUnpolledExpiredTask() + { + var clock = CreateClock(); + var store = CreateStore(clock, s_ttl); + var victim = await store.CreateTaskAsync(CT); + + store.DefaultTimeToLive = null; + var live = await store.CreateTaskAsync(CT); + + clock.Advance(s_ttl + s_ttl + s_sweepInterval); + Assert.NotNull(await store.GetTaskAsync(live.TaskId, CT)); + + Assert.False(store._tasks.ContainsKey(victim.TaskId)); + Assert.True(store._tasks.ContainsKey(live.TaskId)); + } + + [Fact] + public async Task Sweep_WithinInterval_DoesNotRemoveExpiredTask() + { + var clock = CreateClock(); + var store = CreateStore(clock, s_ttl); + var victim = await store.CreateTaskAsync(CT); + + store.DefaultTimeToLive = null; + var live = await store.CreateTaskAsync(CT); + + clock.Advance(s_ttl + s_ttl); + Assert.NotNull(await store.GetTaskAsync(live.TaskId, CT)); + + Assert.True(store._tasks.ContainsKey(victim.TaskId)); + } + + [Fact] + public async Task Sweep_DoesNotRemoveSecondWindowWorkingTask() + { + var clock = CreateClock(); + var ttl = TimeSpan.FromSeconds(20); + var store = CreateStore(clock, ttl); + var created = await store.CreateTaskAsync(CT); + + // Sweep interval is 30s; with TTL=20s that instant is still in the second window (20s..40s). + clock.Advance(s_sweepInterval); + _ = await store.CreateTaskAsync(CT); + + Assert.True(store._tasks.ContainsKey(created.TaskId)); + var task = await store.GetTaskAsync(created.TaskId, CT); + Assert.NotNull(task); + Assert.Equal(McpTaskStatus.Working, task.Status); + } + + [Fact] + public async Task Sweep_RemovesTerminalPastLastUpdatedTtl() + { + var clock = CreateClock(); + var store = CreateStore(clock, s_ttl); + var created = await store.CreateTaskAsync(CT); + await store.SetCompletedAsync(created.TaskId, DonePayload(), CT); + + clock.Advance(s_ttl + s_sweepInterval); + _ = await store.CreateTaskAsync(CT); + + Assert.False(store._tasks.ContainsKey(created.TaskId)); + } + + [Fact] + public async Task Sweep_DoesNotRemoveTerminalWithinLastUpdatedTtl() + { + var clock = CreateClock(); + var store = CreateStore(clock, TimeSpan.FromSeconds(60)); + var created = await store.CreateTaskAsync(CT); + await store.SetCompletedAsync(created.TaskId, DonePayload(), CT); + + clock.Advance(s_sweepInterval); + _ = await store.CreateTaskAsync(CT); + + Assert.True(store._tasks.ContainsKey(created.TaskId)); + } + + [Fact] + public async Task Sweep_ConcurrentPulses_OnlyOneWalks() + { + var clock = CreateClock(); + var store = CreateStore(clock, s_ttl); + var ids = new string[8]; + for (int i = 0; i < ids.Length; i++) + { + ids[i] = (await store.CreateTaskAsync(CT)).TaskId; + } + + clock.Advance(s_ttl + s_ttl + s_sweepInterval); + + await Task.WhenAll(Enumerable.Range(0, 16).Select(_ => store.CreateTaskAsync(CT))); + + foreach (var id in ids) + { + Assert.False(store._tasks.ContainsKey(id)); + } + } + + [Fact] + public async Task SweepExpired_ConcurrentCallers_CompareExchangeLoserReturns() + { + var clock = CreateClock(); + var store = CreateStore(clock, s_ttl); + _ = await store.CreateTaskAsync(CT); + + int threadCount = Math.Max(16, Environment.ProcessorCount * 2); + for (int round = 0; round < 20; round++) + { + // Reset the throttle so every round's callers all pass the interval check and race the CAS. + store._lastSweepTicks = clock.GetUtcNow().UtcTicks - InMemoryMcpTaskStore.s_sweepIntervalTicks - 1; + var now = clock.GetUtcNow(); + var barrier = new Barrier(threadCount); + var pulses = Enumerable.Range(0, threadCount).Select(_ => Task.Run(() => + { + barrier.SignalAndWait(CT); + store.SweepExpired(now, CancellationToken.None); + }, CT)); + + await Task.WhenAll(pulses); + } + } + + [Fact] + public async Task ConcurrentUpdates_RetryLoop_DoesNotLoseKeys() + { + var store = new InMemoryMcpTaskStore(); + var created = await store.CreateTaskAsync(CT); + var start = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + + var workers = Enumerable.Range(0, 32).Select(async i => + { + await start.Task; + for (int n = 0; n < 40; n++) + { + await store.SetInputRequestsAsync(created.TaskId, new Dictionary + { + [$"req{i}_{n}"] = MakeRequest($"value{i}_{n}") + }, CT); + } + }); + + var running = Task.WhenAll(workers); + start.SetResult(); + await running; + + var task = await store.GetTaskAsync(created.TaskId, CT); + Assert.NotNull(task); + Assert.Equal(32 * 40, task.InputRequests!.Count); + } + + [Fact] + public async Task CreateTaskAsync_CanceledToken_DoesNotInsert() + { + var store = new InMemoryMcpTaskStore(); + + await Assert.ThrowsAnyAsync( + () => store.CreateTaskAsync(new CancellationToken(canceled: true))); + + Assert.Empty(store._tasks); + } + + [Fact] + public async Task PublicMethods_CanceledToken_ThrowWithoutMutating() + { + var store = new InMemoryMcpTaskStore(); + var created = await store.CreateTaskAsync(CT); + var canceled = new CancellationToken(canceled: true); + + await Assert.ThrowsAnyAsync(() => store.GetTaskAsync(created.TaskId, canceled)); + await Assert.ThrowsAnyAsync(() => store.SetCompletedAsync(created.TaskId, DonePayload(), canceled)); + await Assert.ThrowsAnyAsync(() => store.SetFailedAsync(created.TaskId, ErrorPayload(), canceled)); + await Assert.ThrowsAnyAsync(() => store.SetCancelledAsync(created.TaskId, canceled)); + await Assert.ThrowsAnyAsync( + () => store.SetInputRequestsAsync(created.TaskId, new Dictionary { ["r"] = MakeRequest("x") }, canceled)); + await Assert.ThrowsAnyAsync( + () => store.ResolveInputRequestsAsync(created.TaskId, new Dictionary { ["r"] = MakeResponse("x") }, canceled)); + + var task = await store.GetTaskAsync(created.TaskId, CT); + Assert.NotNull(task); + Assert.Equal(McpTaskStatus.Working, task.Status); + } + + [Fact] + public async Task SweepExpired_CanceledDuringWalk_Throws() + { + var inner = CreateClock(); + using var cts = new CancellationTokenSource(); + var clock = new CancelOnNthUtcNowTimeProvider(inner, cts, cancelOnCall: 3); + var store = new InMemoryMcpTaskStore(clock) { DefaultTimeToLive = s_ttl }; + + _ = await store.CreateTaskAsync(CT); + inner.Advance(s_ttl + s_ttl + s_sweepInterval); + + await Assert.ThrowsAnyAsync(() => store.CreateTaskAsync(cts.Token)); + } + + [Fact] + public async Task SetCompletedAsync_CanceledAfterStart_DoesNotCommit() + { + var inner = CreateClock(); + using var cts = new CancellationTokenSource(); + var clock = new CancelOnNthUtcNowTimeProvider(inner, cts, cancelOnCall: 3); + var store = new InMemoryMcpTaskStore(clock); + + var created = await store.CreateTaskAsync(CT); + + await Assert.ThrowsAnyAsync( + () => store.SetCompletedAsync(created.TaskId, DonePayload(), cts.Token)); + + var task = await store.GetTaskAsync(created.TaskId, CT); + Assert.NotNull(task); + Assert.Equal(McpTaskStatus.Working, task.Status); + Assert.Null(task.Result); + } + + [Fact] + public async Task SetCancelledAsync_CanceledAfterStart_DoesNotCommit() + { + var inner = CreateClock(); + using var cts = new CancellationTokenSource(); + var clock = new CancelOnNthUtcNowTimeProvider(inner, cts, cancelOnCall: 3); + var store = new InMemoryMcpTaskStore(clock); + + var created = await store.CreateTaskAsync(CT); + + await Assert.ThrowsAnyAsync( + () => store.SetCancelledAsync(created.TaskId, cts.Token)); + + var task = await store.GetTaskAsync(created.TaskId, CT); + Assert.NotNull(task); + Assert.Equal(McpTaskStatus.Working, task.Status); + } + + [Fact] + public async Task SetFailedAsync_CanceledAfterStart_DoesNotCommit() + { + var inner = CreateClock(); + using var cts = new CancellationTokenSource(); + var clock = new CancelOnNthUtcNowTimeProvider(inner, cts, cancelOnCall: 3); + var store = new InMemoryMcpTaskStore(clock); + + var created = await store.CreateTaskAsync(CT); + + await Assert.ThrowsAnyAsync( + () => store.SetFailedAsync(created.TaskId, ErrorPayload(), cts.Token)); + + var task = await store.GetTaskAsync(created.TaskId, CT); + Assert.NotNull(task); + Assert.Equal(McpTaskStatus.Working, task.Status); + } + + [Fact] + public async Task SetInputRequestsAsync_CanceledAfterStart_DoesNotCommit() + { + var inner = CreateClock(); + using var cts = new CancellationTokenSource(); + var clock = new CancelOnNthUtcNowTimeProvider(inner, cts, cancelOnCall: 3); + var store = new InMemoryMcpTaskStore(clock); + + var created = await store.CreateTaskAsync(CT); + + await Assert.ThrowsAnyAsync( + () => store.SetInputRequestsAsync( + created.TaskId, + new Dictionary { ["r"] = MakeRequest("x") }, + cts.Token)); + + var task = await store.GetTaskAsync(created.TaskId, CT); + Assert.NotNull(task); + Assert.Equal(McpTaskStatus.Working, task.Status); + } + + [Fact] + public async Task ResolveInputRequestsAsync_CanceledAfterStart_DoesNotCommitOrFire() + { + var inner = CreateClock(); + using var cts = new CancellationTokenSource(); + var clock = new CancelOnNthUtcNowTimeProvider(inner, cts, cancelOnCall: 4); + var store = new InMemoryMcpTaskStore(clock); + + var created = await store.CreateTaskAsync(CT); + await store.SetInputRequestsAsync(created.TaskId, new Dictionary + { + ["r"] = MakeRequest("ask"), + }, CT); + + int eventCount = 0; + store.InputResponseReceived += _ => Interlocked.Increment(ref eventCount); + + await Assert.ThrowsAnyAsync( + () => store.ResolveInputRequestsAsync( + created.TaskId, + new Dictionary { ["r"] = MakeResponse("x") }, + cts.Token)); + + Assert.Equal(0, eventCount); + var task = await store.GetTaskAsync(created.TaskId, CT); + Assert.NotNull(task); + Assert.Equal(McpTaskStatus.InputRequired, task.Status); + } + + [Fact] + public async Task SetCancelledAsync_ConcurrentWithComplete_ReturnValueMatchesWinner() + { + var store = new InMemoryMcpTaskStore(); + + for (int i = 0; i < 200; i++) + { + var created = await store.CreateTaskAsync(CT); + var complete = store.SetCompletedAsync(created.TaskId, DonePayload(), CT); + var cancel = store.SetCancelledAsync(created.TaskId, CT); + await Task.WhenAll(complete, cancel); + + bool cancelled = await cancel; + var snapshot = store._tasks[created.TaskId]; + if (cancelled) + { + Assert.Equal(McpTaskStatus.Cancelled, snapshot.Status); + Assert.Null(snapshot.Result); + } + else + { + Assert.Equal(McpTaskStatus.Completed, snapshot.Status); + Assert.Equal("done", snapshot.Result!.Value.GetString()); + } + } + } + + [Fact] + public async Task ConcurrentTerminalGets_AtLeastOneServesThenRowIsGone() + { + var store = new InMemoryMcpTaskStore(); + var created = await store.CreateTaskAsync(CT); + await store.SetCompletedAsync(created.TaskId, DonePayload(), CT); + + var results = await Task.WhenAll( + store.GetTaskAsync(created.TaskId, CT), + store.GetTaskAsync(created.TaskId, CT), + store.GetTaskAsync(created.TaskId, CT)); + + Assert.Contains(results, r => r is { Status: McpTaskStatus.Completed }); + Assert.Null(await store.GetTaskAsync(created.TaskId, CT)); + } + + [Fact] + public async Task ResolveInputRequestsAsync_FiresEventPerResponse() + { + var store = new InMemoryMcpTaskStore(); + var created = await store.CreateTaskAsync(CT); + await store.SetInputRequestsAsync(created.TaskId, new Dictionary + { + ["a"] = MakeRequest("ask-a"), + ["b"] = MakeRequest("ask-b"), + }, CT); + + var seen = new ConcurrentBag(); + store.InputResponseReceived += args => + { + Assert.Equal(created.TaskId, args.TaskId); + seen.Add(args.RequestId); + }; + + await store.ResolveInputRequestsAsync(created.TaskId, new Dictionary + { + ["a"] = MakeResponse("a"), + ["b"] = MakeResponse("b"), + }, CT); + + Assert.Equal(2, seen.Count); + Assert.Contains("a", seen); + Assert.Contains("b", seen); + } + + [Fact] + public async Task ResolveInputRequestsAsync_UnsubscribedHandler_DoesNotFire() + { + var store = new InMemoryMcpTaskStore(); + var created = await store.CreateTaskAsync(CT); + + int eventCount = 0; + Action handler = _ => Interlocked.Increment(ref eventCount); + store.InputResponseReceived += handler; + store.InputResponseReceived -= handler; + + await store.ResolveInputRequestsAsync(created.TaskId, new Dictionary + { + ["r"] = MakeResponse("x"), + }, CT); + + Assert.Equal(0, eventCount); + } + + [Fact] + public async Task ResolveInputRequestsAsync_NonImmutableInputRequests_StillResolves() + { + var store = new InMemoryMcpTaskStore(); + var created = await store.CreateTaskAsync(CT); + + store._tasks[created.TaskId] = created with + { + Status = McpTaskStatus.InputRequired, + InputRequests = new Dictionary + { + ["r"] = MakeRequest("ask"), + }, + }; + + await store.ResolveInputRequestsAsync(created.TaskId, new Dictionary + { + ["r"] = MakeResponse("x"), + }, CT); + + var task = await store.GetTaskAsync(created.TaskId, CT); + Assert.NotNull(task); + Assert.Equal(McpTaskStatus.Working, task.Status); + } + + [Fact] + public async Task SetInputRequestsAsync_NonImmutableInputRequests_Merges() + { + var store = new InMemoryMcpTaskStore(); + var created = await store.CreateTaskAsync(CT); + + store._tasks[created.TaskId] = created with + { + InputRequests = new Dictionary + { + ["a"] = MakeRequest("ask-a"), + }, + }; + + await store.SetInputRequestsAsync(created.TaskId, new Dictionary + { + ["b"] = MakeRequest("ask-b"), + }, CT); + + var task = await store.GetTaskAsync(created.TaskId, CT); + Assert.NotNull(task); + Assert.Equal(McpTaskStatus.InputRequired, task.Status); + Assert.Equal(2, task.InputRequests!.Count); + } + + [Fact] + public async Task ResolveInputRequestsAsync_EmptyResponses_OnWorkingTask_FiresNoEvents() + { + var store = new InMemoryMcpTaskStore(); + var created = await store.CreateTaskAsync(CT); + + int eventCount = 0; + store.InputResponseReceived += _ => Interlocked.Increment(ref eventCount); + + await store.ResolveInputRequestsAsync(created.TaskId, new Dictionary(), CT); + + Assert.Equal(0, eventCount); + var task = await store.GetTaskAsync(created.TaskId, CT); + Assert.NotNull(task); + Assert.Equal(McpTaskStatus.Working, task.Status); + } + + [Fact] + public async Task GetTaskAsync_EmptyTaskId_ReturnsNull() + { + var store = new InMemoryMcpTaskStore(); + Assert.Null(await store.GetTaskAsync(string.Empty, CT)); + } + + [Fact] + public async Task SetFailedAsync_DoesNotOverwriteFailedTask() + { + var store = new InMemoryMcpTaskStore(); + var created = await store.CreateTaskAsync(CT); + await store.SetFailedAsync(created.TaskId, JsonElement.Parse("""{"message":"first"}"""), CT); + await store.SetFailedAsync(created.TaskId, JsonElement.Parse("""{"message":"second"}"""), CT); + + var task = await store.GetTaskAsync(created.TaskId, CT); + Assert.NotNull(task); + Assert.Equal(McpTaskStatus.Failed, task.Status); + Assert.Equal("first", task.Error!.Value.GetProperty("message").GetString()); + } + + [Fact] + public async Task SetCancelledAsync_SecondCall_ReturnsFalse() + { + var store = new InMemoryMcpTaskStore(); + var created = await store.CreateTaskAsync(CT); + + Assert.True(await store.SetCancelledAsync(created.TaskId, CT)); + Assert.False(await store.SetCancelledAsync(created.TaskId, CT)); + } + + [Fact] + public async Task CreateTaskAsync_DefaultPollInterval_Is1000() + { + var store = new InMemoryMcpTaskStore(); + var created = await store.CreateTaskAsync(CT); + Assert.Equal(1000, created.PollIntervalMs); + } + + [Fact] + public async Task ResolveInputRequestsAsync_OnWorkingTask_FiresEvenWhenKeyWasNeverPending() + { + var store = new InMemoryMcpTaskStore(); + var created = await store.CreateTaskAsync(CT); + + InputResponseReceivedEventArgs? args = null; + store.InputResponseReceived += e => args = e; + + var response = MakeResponse("x"); + await store.ResolveInputRequestsAsync(created.TaskId, new Dictionary + { + ["ghost"] = response, + }, CT); + + Assert.NotNull(args); + Assert.Equal(created.TaskId, args.TaskId); + Assert.Equal("ghost", args.RequestId); + } + + [Theory] + [InlineData(McpTaskStatus.Working, false)] + [InlineData(McpTaskStatus.InputRequired, false)] + [InlineData(McpTaskStatus.Completed, true)] + [InlineData(McpTaskStatus.Failed, true)] + [InlineData(McpTaskStatus.Cancelled, true)] + public void IsTerminal_KnownStatuses_ReturnExpected(McpTaskStatus status, bool expected) + { + Assert.Equal(expected, InMemoryMcpTaskStore.IsTerminal(status)); + } + + [Fact] + public void IsTerminal_UnknownStatus_ReturnsFalse() + { + Assert.False(InMemoryMcpTaskStore.IsTerminal((McpTaskStatus)int.MaxValue)); + } + + [Fact] + public void TryGetPositiveTtl_Null_ReturnsFalse() + { + Assert.False(InMemoryMcpTaskStore.TryGetPositiveTtl(Info(ttl: null), out var ttl)); + Assert.Equal(default, ttl); + } + + [Fact] + public void TryGetPositiveTtl_Zero_ReturnsFalse() + { + Assert.False(InMemoryMcpTaskStore.TryGetPositiveTtl(Info(ttl: TimeSpan.Zero), out var ttl)); + Assert.Equal(default, ttl); + } + + [Fact] + public void TryGetPositiveTtl_Negative_ReturnsFalse() + { + Assert.False(InMemoryMcpTaskStore.TryGetPositiveTtl(Info(ttl: TimeSpan.FromSeconds(-1)), out var ttl)); + Assert.Equal(default, ttl); + } + + [Fact] + public void TryGetPositiveTtl_Positive_ReturnsValue() + { + Assert.True(InMemoryMcpTaskStore.TryGetPositiveTtl(Info(ttl: s_ttl), out var ttl)); + Assert.Equal(s_ttl, ttl); + } + + [Fact] + public void IsActive_WorkingWithoutTtl_IsTrue() + { + Assert.True(InMemoryMcpTaskStore.IsActive(Info(), s_start + TimeSpan.FromDays(365))); + } + + [Fact] + public void IsActive_WorkingWithZeroTtl_IsTrue() + { + Assert.True(InMemoryMcpTaskStore.IsActive(Info(ttl: TimeSpan.Zero), s_start + TimeSpan.FromDays(1))); + } + + [Fact] + public void IsActive_WorkingWithNegativeTtl_IsTrue() + { + Assert.True(InMemoryMcpTaskStore.IsActive(Info(ttl: TimeSpan.FromSeconds(-5)), s_start + TimeSpan.FromDays(1))); + } + + [Fact] + public void IsActive_WorkingBeforeTtl_IsTrue() + { + Assert.True(InMemoryMcpTaskStore.IsActive(Info(ttl: s_ttl), s_start + TimeSpan.FromSeconds(9))); + } + + [Fact] + public void IsActive_WorkingAtTtlBoundary_IsFalse() + { + Assert.False(InMemoryMcpTaskStore.IsActive(Info(ttl: s_ttl), s_start + s_ttl)); + } + + [Fact] + public void IsActive_WorkingInSecondWindow_IsFalse() + { + Assert.False(InMemoryMcpTaskStore.IsActive(Info(ttl: s_ttl), s_start + s_ttl + TimeSpan.FromSeconds(1))); + } + + [Fact] + public void IsActive_InputRequiredWithinWindow_IsTrue() + { + Assert.True(InMemoryMcpTaskStore.IsActive(Info(McpTaskStatus.InputRequired, ttl: s_ttl), s_start + TimeSpan.FromSeconds(1))); + } + + [Fact] + public void IsActive_InputRequiredAfterTtl_IsFalse() + { + Assert.False(InMemoryMcpTaskStore.IsActive(Info(McpTaskStatus.InputRequired, ttl: s_ttl), s_start + s_ttl)); + } + + [Theory] + [InlineData(McpTaskStatus.Completed)] + [InlineData(McpTaskStatus.Failed)] + [InlineData(McpTaskStatus.Cancelled)] + public void IsActive_Terminal_IsFalseEvenWithoutTtl(McpTaskStatus status) + { + Assert.False(InMemoryMcpTaskStore.IsActive(Info(status), s_start)); + Assert.False(InMemoryMcpTaskStore.IsActive(Info(status, ttl: s_ttl), s_start + TimeSpan.FromSeconds(1))); + } + + [Fact] + public void HasExpired_NoTtl_IsFalseForWorkingAndTerminal() + { + var far = s_start + TimeSpan.FromDays(365); + Assert.False(InMemoryMcpTaskStore.HasExpired(Info(), far)); + Assert.False(InMemoryMcpTaskStore.HasExpired(Info(ttl: TimeSpan.Zero), far)); + Assert.False(InMemoryMcpTaskStore.HasExpired(Info(ttl: TimeSpan.FromSeconds(-1)), far)); + Assert.False(InMemoryMcpTaskStore.HasExpired(Info(McpTaskStatus.Completed), far)); + } + + [Fact] + public void HasExpired_WorkingAtOneTtl_IsFalse() + { + Assert.False(InMemoryMcpTaskStore.HasExpired(Info(ttl: s_ttl), s_start + s_ttl)); + } + + [Fact] + public void HasExpired_WorkingJustBeforeTwoTtl_IsFalse() + { + Assert.False(InMemoryMcpTaskStore.HasExpired(Info(ttl: s_ttl), s_start + s_ttl + s_ttl - TimeSpan.FromTicks(1))); + } + + [Fact] + public void HasExpired_WorkingAtTwoTtl_IsTrue() + { + Assert.True(InMemoryMcpTaskStore.HasExpired(Info(ttl: s_ttl), s_start + s_ttl + s_ttl)); + } + + [Fact] + public void HasExpired_InputRequiredAtTwoTtl_IsTrue() + { + Assert.True(InMemoryMcpTaskStore.HasExpired(Info(McpTaskStatus.InputRequired, ttl: s_ttl), s_start + s_ttl + s_ttl)); + } + + [Fact] + public void HasExpired_TerminalJustBeforeLastUpdatedPlusTtl_IsFalse() + { + var completed = Info(McpTaskStatus.Completed, lastUpdatedAt: s_start + TimeSpan.FromSeconds(5), ttl: s_ttl); + Assert.False(InMemoryMcpTaskStore.HasExpired(completed, s_start + TimeSpan.FromSeconds(5) + s_ttl - TimeSpan.FromTicks(1))); + } + + [Fact] + public void HasExpired_TerminalAtLastUpdatedPlusTtl_IsTrue() + { + var completed = Info(McpTaskStatus.Completed, lastUpdatedAt: s_start + TimeSpan.FromSeconds(5), ttl: s_ttl); + Assert.True(InMemoryMcpTaskStore.HasExpired(completed, s_start + TimeSpan.FromSeconds(5) + s_ttl)); + } + + [Fact] + public void HasExpired_TerminalClockIsLastUpdatedNotCreated() + { + // Completed at T+15s with TTL=10s: created+2×TTL is T+20s, but LastUpdatedAt+TTL is T+25s. + var completed = Info( + McpTaskStatus.Completed, + createdAt: s_start, + lastUpdatedAt: s_start + TimeSpan.FromSeconds(15), + ttl: s_ttl); + + Assert.False(InMemoryMcpTaskStore.HasExpired(completed, s_start + TimeSpan.FromSeconds(20))); + Assert.True(InMemoryMcpTaskStore.HasExpired(completed, s_start + TimeSpan.FromSeconds(25))); + } + + [Fact] + public void GetRequests_Null_ReturnsEmpty() + { + Assert.Same(ImmutableDictionary.Empty, InMemoryMcpTaskStore.GetRequests(Info())); + } + + [Fact] + public void GetRequests_ImmutableDictionary_ReturnsSameInstance() + { + var requests = ImmutableDictionary.Empty.Add("a", MakeRequest("x")); + var result = InMemoryMcpTaskStore.GetRequests(Info(inputRequests: requests)); + Assert.Same(requests, result); + } + + [Fact] + public void GetRequests_MutableDictionary_CopiesEntries() + { + var requests = new Dictionary { ["a"] = MakeRequest("x") }; + var result = InMemoryMcpTaskStore.GetRequests(Info(inputRequests: requests)); + Assert.False(ReferenceEquals(requests, result)); + Assert.True(result.ContainsKey("a")); + } + + [Fact] + public async Task SweepExpired_WithinInterval_IsNoOp() + { + var clock = CreateClock(); + var store = CreateStore(clock, s_ttl); + var created = await store.CreateTaskAsync(CT); + + store.SweepExpired(clock.GetUtcNow(), CT); + + Assert.True(store._tasks.ContainsKey(created.TaskId)); + } + + [Fact] + public async Task SweepExpired_AfterInterval_RemovesExpiredWorkingTask() + { + var clock = CreateClock(); + var store = CreateStore(clock, s_ttl); + var created = await store.CreateTaskAsync(CT); + + clock.Advance(s_ttl + s_ttl + s_sweepInterval); + store.SweepExpired(clock.GetUtcNow(), CT); + + Assert.False(store._tasks.ContainsKey(created.TaskId)); + } + + [Fact] + public async Task SweepExpired_CanceledToken_ThrowsDuringWalk() + { + var clock = CreateClock(); + var store = CreateStore(clock, s_ttl); + _ = await store.CreateTaskAsync(CT); + + clock.Advance(s_ttl + s_ttl + s_sweepInterval); + + Assert.ThrowsAny(() => + store.SweepExpired(clock.GetUtcNow(), new CancellationToken(canceled: true))); + } + + [Fact] + public async Task SweepExpired_DoesNotRemoveActiveTask() + { + var clock = CreateClock(); + var store = CreateStore(clock, TimeSpan.FromMinutes(10)); + var created = await store.CreateTaskAsync(CT); + + clock.Advance(s_sweepInterval); + store.SweepExpired(clock.GetUtcNow(), CT); + + Assert.True(store._tasks.ContainsKey(created.TaskId)); + } + + [Fact] + public void TryUpdate_MissingKey_ReturnsFalse() + { + var store = new InMemoryMcpTaskStore(); + Assert.False(store.TryUpdate("missing", entry => entry with { StatusMessage = "x" }, CT)); + } + + [Fact] + public async Task TryUpdate_SameInstance_ReturnsTrueWithoutWrite() + { + var store = new InMemoryMcpTaskStore(); + var created = await store.CreateTaskAsync(CT); + var before = store._tasks[created.TaskId]; + + Assert.True(store.TryUpdate(created.TaskId, entry => entry, CT)); + Assert.Same(before, store._tasks[created.TaskId]); + } + + [Fact] + public async Task TryUpdate_NewInstance_ReplacesSnapshot() + { + var store = new InMemoryMcpTaskStore(); + var created = await store.CreateTaskAsync(CT); + + Assert.True(store.TryUpdate(created.TaskId, entry => entry with { StatusMessage = "hi" }, CT)); + Assert.Equal("hi", store._tasks[created.TaskId].StatusMessage); + } + + [Fact] + public async Task TryUpdate_CanceledToken_ThrowsWithoutMutating() + { + var store = new InMemoryMcpTaskStore(); + var created = await store.CreateTaskAsync(CT); + + Assert.ThrowsAny(() => + store.TryUpdate(created.TaskId, entry => entry with { StatusMessage = "nope" }, new CancellationToken(canceled: true))); + + Assert.Null(store._tasks[created.TaskId].StatusMessage); + } + + [Fact] + public void TimeProvider_AndSweepTicks_AreInitializedFromClock() + { + var clock = CreateClock(); + var store = CreateStore(clock); + + Assert.Same(clock, store._timeProvider); + Assert.Equal(s_start.UtcTicks, store._lastSweepTicks); + Assert.Equal(TimeSpan.FromSeconds(30).Ticks, InMemoryMcpTaskStore.s_sweepIntervalTicks); + } + + private sealed class CancelOnNthUtcNowTimeProvider : TimeProvider + { + private readonly FakeTimeProvider _inner; + private readonly CancellationTokenSource _cts; + private readonly int _cancelOnCall; + private int _calls; + + public CancelOnNthUtcNowTimeProvider(FakeTimeProvider inner, CancellationTokenSource cts, int cancelOnCall) + { + _inner = inner; + _cts = cts; + _cancelOnCall = cancelOnCall; + } + + public override DateTimeOffset GetUtcNow() + { + if (Interlocked.Increment(ref _calls) >= _cancelOnCall) + { + _cts.Cancel(); + } + + return _inner.GetUtcNow(); + } + + public override long TimestampFrequency => _inner.TimestampFrequency; + + public override long GetTimestamp() => _inner.GetTimestamp(); + } } diff --git a/tests/ModelContextProtocol.Tests/Server/McpTaskStoreTests.cs b/tests/ModelContextProtocol.Tests/Server/McpTaskStoreTests.cs index 4731e8b39..a9a84cfc3 100644 --- a/tests/ModelContextProtocol.Tests/Server/McpTaskStoreTests.cs +++ b/tests/ModelContextProtocol.Tests/Server/McpTaskStoreTests.cs @@ -565,10 +565,11 @@ public async Task CancelTaskAsync_AlreadyCompleted_AcknowledgesIdempotently_AndD var cancelResult = await client.CancelTaskAsync(taskId, ct); Assert.NotNull(cancelResult); - // The task must remain Completed and the result must not be lost. - var verifyResult = await client.GetTaskAsync(taskId, ct); - var stillCompleted = Assert.IsType(verifyResult); - Assert.NotEqual(default(JsonElement), stillCompleted.Result); + // The completing Get already served-and-dropped the terminal snapshot. Cancel must not + // resurrect a row (as cancelled or otherwise). + var ex = await Assert.ThrowsAsync(async () => + await client.GetTaskAsync(taskId, ct)); + Assert.Contains("Unknown task", ex.Message); } [Fact] diff --git a/tests/ModelContextProtocol.Tests/Server/TaskCancellationIntegrationTests.cs b/tests/ModelContextProtocol.Tests/Server/TaskCancellationIntegrationTests.cs index e751bdcdc..f492ca13c 100644 --- a/tests/ModelContextProtocol.Tests/Server/TaskCancellationIntegrationTests.cs +++ b/tests/ModelContextProtocol.Tests/Server/TaskCancellationIntegrationTests.cs @@ -327,9 +327,11 @@ public async Task CompletedTask_CancelIsAcknowledgedIdempotentlyAndStateUnchange var cancelResult = await client.CancelTaskAsync(taskId, ct); Assert.NotNull(cancelResult); - // Verify status is still completed (not flipped to cancelled). - var verifyResult = await client.GetTaskAsync(taskId, ct); - Assert.IsType(verifyResult); + // The completing Get already served-and-dropped the terminal snapshot. Cancel must not + // resurrect a cancelled row. + var ex = await Assert.ThrowsAsync(async () => + await client.GetTaskAsync(taskId, ct)); + Assert.Contains("Unknown task", ex.Message); } [Fact] @@ -357,8 +359,10 @@ public async Task CompletedWithErrorTask_CancelIsAcknowledgedIdempotently() var cancelResult = await client.CancelTaskAsync(taskId, ct); Assert.NotNull(cancelResult); - // Verify status is still completed (not flipped to cancelled). - var verifyResult = await client.GetTaskAsync(taskId, ct); - Assert.IsType(verifyResult); + // The completing Get already served-and-dropped the terminal snapshot. Cancel must not + // resurrect a cancelled row. + var ex = await Assert.ThrowsAsync(async () => + await client.GetTaskAsync(taskId, ct)); + Assert.Contains("Unknown task", ex.Message); } }