2ca96f548bCo-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Tarek Mahmoud Sayed <tarekms@ntdev.microsoft.com>
549 lines
19 KiB
C#
549 lines
19 KiB
C#
using ModelContextProtocol.Extensions.Tasks;
|
|
using ModelContextProtocol.Protocol;
|
|
using ModelContextProtocol.Server;
|
|
using System.Text.Json;
|
|
|
|
#pragma warning disable MCPEXP001
|
|
|
|
namespace ModelContextProtocol.Tests.Server;
|
|
|
|
/// <summary>
|
|
/// Unit tests for <see cref="InMemoryMcpTaskStore"/>.
|
|
/// </summary>
|
|
public class InMemoryMcpTaskStoreTests
|
|
{
|
|
private CancellationToken CT => TestContext.Current.CancellationToken;
|
|
|
|
private static InputRequest MakeRequest(string payload) =>
|
|
new() { Method = "test/method", Params = JsonSerializer.SerializeToElement(payload, McpJsonUtilities.DefaultOptions) };
|
|
|
|
private static InputResponse MakeResponse(string payload) =>
|
|
new() { RawValue = JsonSerializer.SerializeToElement(payload, McpJsonUtilities.DefaultOptions) };
|
|
|
|
[Fact]
|
|
public async Task CreateTaskAsync_ReturnsWorkingTaskWithUniqueId()
|
|
{
|
|
var store = new InMemoryMcpTaskStore();
|
|
|
|
var result = await store.CreateTaskAsync(CT);
|
|
|
|
Assert.NotNull(result);
|
|
Assert.NotEmpty(result.TaskId);
|
|
Assert.Equal(McpTaskStatus.Working, result.Status);
|
|
Assert.NotEqual(default, result.CreatedAt);
|
|
Assert.NotEqual(default, result.LastUpdatedAt);
|
|
}
|
|
|
|
[Fact]
|
|
public async Task CreateTaskAsync_GeneratesUniqueIds()
|
|
{
|
|
var store = new InMemoryMcpTaskStore();
|
|
|
|
var task1 = await store.CreateTaskAsync(CT);
|
|
var task2 = await store.CreateTaskAsync(CT);
|
|
|
|
Assert.NotEqual(task1.TaskId, task2.TaskId);
|
|
}
|
|
|
|
[Fact]
|
|
public async Task CreateTaskAsync_UsesDefaultPollInterval()
|
|
{
|
|
var store = new InMemoryMcpTaskStore { DefaultPollIntervalMs = 500 };
|
|
|
|
var result = await store.CreateTaskAsync(CT);
|
|
|
|
Assert.Equal(500, result.PollIntervalMs);
|
|
}
|
|
|
|
[Fact]
|
|
public async Task CreateTaskAsync_UsesDefaultTimeToLive()
|
|
{
|
|
var store = new InMemoryMcpTaskStore { DefaultTimeToLive = TimeSpan.FromSeconds(30) };
|
|
|
|
var result = await store.CreateTaskAsync(CT);
|
|
|
|
Assert.Equal(TimeSpan.FromSeconds(30), result.TimeToLive);
|
|
}
|
|
|
|
[Fact]
|
|
public async Task GetTaskAsync_ReturnsWorkingTask()
|
|
{
|
|
var store = new InMemoryMcpTaskStore();
|
|
var created = await store.CreateTaskAsync(CT);
|
|
|
|
var result = await store.GetTaskAsync(created.TaskId, CT);
|
|
|
|
Assert.NotNull(result);
|
|
Assert.Equal(McpTaskStatus.Working, result.Status);
|
|
Assert.Equal(created.TaskId, result.TaskId);
|
|
}
|
|
|
|
[Fact]
|
|
public async Task GetTaskAsync_ReturnsNullForUnknownId()
|
|
{
|
|
var store = new InMemoryMcpTaskStore();
|
|
|
|
var result = await store.GetTaskAsync("nonexistent", CT);
|
|
|
|
Assert.Null(result);
|
|
}
|
|
|
|
[Fact]
|
|
public async Task GetTaskAsync_WithinTimeToLive_ReturnsTask()
|
|
{
|
|
var store = new InMemoryMcpTaskStore { DefaultTimeToLive = TimeSpan.FromMinutes(10) };
|
|
var created = await store.CreateTaskAsync(CT);
|
|
|
|
var result = await store.GetTaskAsync(created.TaskId, CT);
|
|
|
|
Assert.NotNull(result);
|
|
Assert.Equal(created.TaskId, result.TaskId);
|
|
}
|
|
|
|
[Fact]
|
|
public async Task GetTaskAsync_AfterTimeToLiveElapsed_ReturnsNull()
|
|
{
|
|
var store = new InMemoryMcpTaskStore { DefaultTimeToLive = TimeSpan.FromMilliseconds(100) };
|
|
var created = await store.CreateTaskAsync(CT);
|
|
|
|
await Task.Delay(TimeSpan.FromMilliseconds(500), CT);
|
|
|
|
var result = await store.GetTaskAsync(created.TaskId, CT);
|
|
|
|
Assert.Null(result);
|
|
}
|
|
|
|
[Fact]
|
|
public async Task GetTaskAsync_WithoutTimeToLive_DoesNotExpire()
|
|
{
|
|
var store = new InMemoryMcpTaskStore();
|
|
var created = await store.CreateTaskAsync(CT);
|
|
|
|
await Task.Delay(TimeSpan.FromMilliseconds(200), CT);
|
|
|
|
var result = await store.GetTaskAsync(created.TaskId, CT);
|
|
|
|
Assert.NotNull(result);
|
|
Assert.Equal(created.TaskId, result.TaskId);
|
|
}
|
|
|
|
[Fact]
|
|
public async Task GetTaskAsync_WithZeroTimeToLive_DoesNotExpire()
|
|
{
|
|
var store = new InMemoryMcpTaskStore { DefaultTimeToLive = TimeSpan.Zero };
|
|
var created = await store.CreateTaskAsync(CT);
|
|
|
|
await Task.Delay(TimeSpan.FromMilliseconds(200), CT);
|
|
|
|
var result = await store.GetTaskAsync(created.TaskId, CT);
|
|
|
|
Assert.NotNull(result);
|
|
Assert.Equal(created.TaskId, result.TaskId);
|
|
}
|
|
|
|
[Fact]
|
|
public async Task SetCompletedAsync_TransitionsToCompleted()
|
|
{
|
|
var store = new InMemoryMcpTaskStore();
|
|
var created = await store.CreateTaskAsync(CT);
|
|
var resultPayload = JsonDocument.Parse("""{"answer":42}""").RootElement.Clone();
|
|
|
|
await store.SetCompletedAsync(created.TaskId, resultPayload, CT);
|
|
|
|
var task = await store.GetTaskAsync(created.TaskId, CT);
|
|
Assert.NotNull(task);
|
|
Assert.Equal(McpTaskStatus.Completed, task.Status);
|
|
Assert.Equal(42, task.Result!.Value.GetProperty("answer").GetInt32());
|
|
}
|
|
|
|
[Fact]
|
|
public async Task SetFailedAsync_TransitionsToFailed()
|
|
{
|
|
var store = new InMemoryMcpTaskStore();
|
|
var created = await store.CreateTaskAsync(CT);
|
|
var errorPayload = JsonDocument.Parse("""{"message":"boom"}""").RootElement.Clone();
|
|
|
|
await store.SetFailedAsync(created.TaskId, errorPayload, CT);
|
|
|
|
var task = await store.GetTaskAsync(created.TaskId, CT);
|
|
Assert.NotNull(task);
|
|
Assert.Equal(McpTaskStatus.Failed, task.Status);
|
|
Assert.Equal("boom", task.Error!.Value.GetProperty("message").GetString());
|
|
}
|
|
|
|
[Fact]
|
|
public async Task SetCancelledAsync_TransitionsToCancelled()
|
|
{
|
|
var store = new InMemoryMcpTaskStore();
|
|
var created = await store.CreateTaskAsync(CT);
|
|
|
|
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 SetCancelledAsync_ReturnsFalseForTerminalTask()
|
|
{
|
|
var store = new InMemoryMcpTaskStore();
|
|
var created = await store.CreateTaskAsync(CT);
|
|
await store.SetCompletedAsync(created.TaskId, JsonSerializer.SerializeToElement("done", McpJsonUtilities.DefaultOptions), CT);
|
|
|
|
var cancelled = await store.SetCancelledAsync(created.TaskId, CT);
|
|
|
|
Assert.False(cancelled);
|
|
var task = await store.GetTaskAsync(created.TaskId, CT);
|
|
Assert.NotNull(task);
|
|
Assert.Equal(McpTaskStatus.Completed, task.Status);
|
|
}
|
|
|
|
[Fact]
|
|
public async Task SetCancelledAsync_ReturnsFalseForUnknownId()
|
|
{
|
|
var store = new InMemoryMcpTaskStore();
|
|
|
|
var cancelled = await store.SetCancelledAsync("nonexistent", CT);
|
|
|
|
Assert.False(cancelled);
|
|
}
|
|
|
|
[Fact]
|
|
public async Task SetInputRequestsAsync_TransitionsToInputRequired()
|
|
{
|
|
var store = new InMemoryMcpTaskStore();
|
|
var created = await store.CreateTaskAsync(CT);
|
|
|
|
var requests = new Dictionary<string, InputRequest>
|
|
{
|
|
["req1"] = new InputRequest
|
|
{
|
|
Method = "elicitation/create",
|
|
Params = JsonElement.Parse("""{"message":"hello"}"""),
|
|
},
|
|
};
|
|
await store.SetInputRequestsAsync(created.TaskId, requests, CT);
|
|
|
|
var task = await store.GetTaskAsync(created.TaskId, CT);
|
|
Assert.NotNull(task);
|
|
Assert.Equal(McpTaskStatus.InputRequired, task.Status);
|
|
Assert.NotNull(task.InputRequests);
|
|
Assert.Single(task.InputRequests);
|
|
Assert.True(task.InputRequests.ContainsKey("req1"));
|
|
}
|
|
|
|
[Fact]
|
|
public async Task SetInputRequestsAsync_MergesMultipleRequests()
|
|
{
|
|
var store = new InMemoryMcpTaskStore();
|
|
var created = await store.CreateTaskAsync(CT);
|
|
|
|
await store.SetInputRequestsAsync(created.TaskId, new Dictionary<string, InputRequest>
|
|
{
|
|
["req1"] = MakeRequest("first")
|
|
}, CT);
|
|
await store.SetInputRequestsAsync(created.TaskId, new Dictionary<string, InputRequest>
|
|
{
|
|
["req2"] = MakeRequest("second")
|
|
}, CT);
|
|
|
|
var task = await store.GetTaskAsync(created.TaskId, CT);
|
|
Assert.NotNull(task);
|
|
Assert.Equal(McpTaskStatus.InputRequired, task.Status);
|
|
Assert.NotNull(task.InputRequests);
|
|
Assert.Equal(2, task.InputRequests.Count);
|
|
Assert.True(task.InputRequests.ContainsKey("req1"));
|
|
Assert.True(task.InputRequests.ContainsKey("req2"));
|
|
}
|
|
|
|
[Fact]
|
|
public async Task ResolveInputRequestsAsync_RemovesMatchedRequests()
|
|
{
|
|
var store = new InMemoryMcpTaskStore();
|
|
var created = await store.CreateTaskAsync(CT);
|
|
|
|
await store.SetInputRequestsAsync(created.TaskId, new Dictionary<string, InputRequest>
|
|
{
|
|
["req1"] = MakeRequest("request1"),
|
|
["req2"] = MakeRequest("request2"),
|
|
}, CT);
|
|
|
|
await store.ResolveInputRequestsAsync(created.TaskId, new Dictionary<string, InputResponse>
|
|
{
|
|
["req1"] = MakeResponse("response1"),
|
|
}, CT);
|
|
|
|
var task = await store.GetTaskAsync(created.TaskId, CT);
|
|
Assert.NotNull(task);
|
|
Assert.Equal(McpTaskStatus.InputRequired, task.Status);
|
|
Assert.NotNull(task.InputRequests);
|
|
Assert.Single(task.InputRequests);
|
|
Assert.True(task.InputRequests.ContainsKey("req2"));
|
|
}
|
|
|
|
[Fact]
|
|
public async Task ResolveInputRequestsAsync_TransitionsToWorkingWhenAllSatisfied()
|
|
{
|
|
var store = new InMemoryMcpTaskStore();
|
|
var created = await store.CreateTaskAsync(CT);
|
|
|
|
await store.SetInputRequestsAsync(created.TaskId, new Dictionary<string, InputRequest>
|
|
{
|
|
["req1"] = MakeRequest("request1"),
|
|
}, CT);
|
|
|
|
await store.ResolveInputRequestsAsync(created.TaskId, new Dictionary<string, InputResponse>
|
|
{
|
|
["req1"] = MakeResponse("response1"),
|
|
}, CT);
|
|
|
|
var task = await store.GetTaskAsync(created.TaskId, CT);
|
|
Assert.NotNull(task);
|
|
Assert.Equal(McpTaskStatus.Working, task.Status);
|
|
}
|
|
|
|
[Fact]
|
|
public async Task SetCompletedAsync_ThrowsForUnknownTask()
|
|
{
|
|
var store = new InMemoryMcpTaskStore();
|
|
|
|
await Assert.ThrowsAsync<InvalidOperationException>(
|
|
() => store.SetCompletedAsync("nonexistent", JsonSerializer.SerializeToElement("x", McpJsonUtilities.DefaultOptions), CT));
|
|
}
|
|
|
|
[Fact]
|
|
public async Task ConcurrentUpdates_DoNotLoseData()
|
|
{
|
|
var store = new InMemoryMcpTaskStore();
|
|
var created = await store.CreateTaskAsync(CT);
|
|
|
|
var tasks = Enumerable.Range(0, 50).Select(i =>
|
|
store.SetInputRequestsAsync(created.TaskId, new Dictionary<string, InputRequest>
|
|
{
|
|
[$"req{i}"] = MakeRequest($"value{i}")
|
|
}, CT));
|
|
|
|
await Task.WhenAll(tasks);
|
|
|
|
var task = await store.GetTaskAsync(created.TaskId, CT);
|
|
Assert.NotNull(task);
|
|
Assert.Equal(McpTaskStatus.InputRequired, task.Status);
|
|
Assert.NotNull(task.InputRequests);
|
|
Assert.Equal(50, task.InputRequests.Count);
|
|
}
|
|
|
|
[Fact]
|
|
public async Task ResolveInputRequestsAsync_ForExtraKeys_DoesNotThrow()
|
|
{
|
|
var store = new InMemoryMcpTaskStore();
|
|
var created = await store.CreateTaskAsync(CT);
|
|
|
|
await store.ResolveInputRequestsAsync(created.TaskId, new Dictionary<string, InputResponse>
|
|
{
|
|
["unknown-key"] = MakeResponse("response"),
|
|
}, CT);
|
|
|
|
var task = await store.GetTaskAsync(created.TaskId, CT);
|
|
Assert.NotNull(task);
|
|
Assert.Equal(McpTaskStatus.Working, task.Status);
|
|
}
|
|
|
|
[Fact]
|
|
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);
|
|
await store.SetInputRequestsAsync(created.TaskId, new Dictionary<string, InputRequest>
|
|
{
|
|
["a"] = MakeRequest("ask-a"),
|
|
["b"] = MakeRequest("ask-b"),
|
|
}, CT);
|
|
|
|
// First resolve "a" — task should still be InputRequired because "b" remains.
|
|
await store.ResolveInputRequestsAsync(created.TaskId, new Dictionary<string, InputResponse>
|
|
{
|
|
["a"] = MakeResponse("answer-a"),
|
|
}, CT);
|
|
|
|
var afterFirst = await store.GetTaskAsync(created.TaskId, CT);
|
|
Assert.NotNull(afterFirst);
|
|
Assert.Equal(McpTaskStatus.InputRequired, afterFirst.Status);
|
|
Assert.NotNull(afterFirst.InputRequests);
|
|
Assert.Single(afterFirst.InputRequests);
|
|
Assert.Contains("b", afterFirst.InputRequests.Keys);
|
|
|
|
// Re-send "a" — should be a no-op (no exception, no state change).
|
|
await store.ResolveInputRequestsAsync(created.TaskId, new Dictionary<string, InputResponse>
|
|
{
|
|
["a"] = MakeResponse("answer-a-again"),
|
|
}, CT);
|
|
|
|
var afterDup = await store.GetTaskAsync(created.TaskId, CT);
|
|
Assert.NotNull(afterDup);
|
|
Assert.Equal(McpTaskStatus.InputRequired, afterDup.Status);
|
|
Assert.NotNull(afterDup.InputRequests);
|
|
Assert.Single(afterDup.InputRequests);
|
|
Assert.Contains("b", afterDup.InputRequests.Keys);
|
|
|
|
// Resolve the remaining "b" — task should transition back to Working with an empty inputRequests.
|
|
await store.ResolveInputRequestsAsync(created.TaskId, new Dictionary<string, InputResponse>
|
|
{
|
|
["b"] = MakeResponse("answer-b"),
|
|
}, CT);
|
|
|
|
var final = await store.GetTaskAsync(created.TaskId, CT);
|
|
Assert.NotNull(final);
|
|
Assert.Equal(McpTaskStatus.Working, final.Status);
|
|
Assert.True(final.InputRequests is null || final.InputRequests.Count == 0);
|
|
}
|
|
|
|
[Fact]
|
|
public async Task ConcurrentResolveInputRequests_OnDisjointKeys_AllResolveCorrectly()
|
|
{
|
|
// 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 seed = Enumerable.Range(0, 20).ToDictionary(
|
|
i => $"req{i}",
|
|
i => MakeRequest($"ask{i}"));
|
|
await store.SetInputRequestsAsync(created.TaskId, seed, CT);
|
|
|
|
var resolveTasks = Enumerable.Range(0, 20).Select(i =>
|
|
store.ResolveInputRequestsAsync(created.TaskId, new Dictionary<string, InputResponse>
|
|
{
|
|
[$"req{i}"] = MakeResponse($"answer{i}"),
|
|
}, CT));
|
|
|
|
await Task.WhenAll(resolveTasks);
|
|
|
|
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 SetCompletedAsync_DoesNotOverwriteCancelledTask()
|
|
{
|
|
var store = new InMemoryMcpTaskStore();
|
|
var created = await store.CreateTaskAsync(CT);
|
|
|
|
var cancelled = await store.SetCancelledAsync(created.TaskId, CT);
|
|
Assert.True(cancelled);
|
|
|
|
// Background worker finishing after cancellation must not flip the task back to Completed.
|
|
await store.SetCompletedAsync(
|
|
created.TaskId,
|
|
JsonSerializer.SerializeToElement("late-result", McpJsonUtilities.DefaultOptions),
|
|
CT);
|
|
|
|
var task = await store.GetTaskAsync(created.TaskId, CT);
|
|
Assert.NotNull(task);
|
|
Assert.Equal(McpTaskStatus.Cancelled, task.Status);
|
|
Assert.Null(task.Result);
|
|
}
|
|
|
|
[Fact]
|
|
public async Task SetFailedAsync_DoesNotOverwriteCancelledTask()
|
|
{
|
|
var store = new InMemoryMcpTaskStore();
|
|
var created = await store.CreateTaskAsync(CT);
|
|
|
|
await store.SetCancelledAsync(created.TaskId, CT);
|
|
|
|
await store.SetFailedAsync(
|
|
created.TaskId,
|
|
JsonElement.Parse("""{"message":"boom"}"""),
|
|
CT);
|
|
|
|
var task = await store.GetTaskAsync(created.TaskId, CT);
|
|
Assert.NotNull(task);
|
|
Assert.Equal(McpTaskStatus.Cancelled, task.Status);
|
|
Assert.Null(task.Error);
|
|
}
|
|
|
|
[Fact]
|
|
public async Task SetCompletedAsync_DoesNotOverwriteCompletedTask()
|
|
{
|
|
var store = new InMemoryMcpTaskStore();
|
|
var created = await store.CreateTaskAsync(CT);
|
|
|
|
var first = JsonSerializer.SerializeToElement("first", McpJsonUtilities.DefaultOptions);
|
|
await store.SetCompletedAsync(created.TaskId, first, CT);
|
|
|
|
// A second completion attempt must not replace the original result.
|
|
var second = JsonSerializer.SerializeToElement("second", McpJsonUtilities.DefaultOptions);
|
|
await store.SetCompletedAsync(created.TaskId, second, CT);
|
|
|
|
var task = await store.GetTaskAsync(created.TaskId, CT);
|
|
Assert.NotNull(task);
|
|
Assert.Equal(McpTaskStatus.Completed, task.Status);
|
|
Assert.Equal("first", task.Result!.Value.GetString());
|
|
}
|
|
|
|
[Fact]
|
|
public async Task ResolveInputRequestsAsync_OnTerminalTask_DoesNotResurrect()
|
|
{
|
|
var store = new InMemoryMcpTaskStore();
|
|
var created = await store.CreateTaskAsync(CT);
|
|
|
|
await store.SetCompletedAsync(
|
|
created.TaskId,
|
|
JsonSerializer.SerializeToElement("done", McpJsonUtilities.DefaultOptions),
|
|
CT);
|
|
|
|
// A client tasks/update against a Completed task must not flip it back to Working.
|
|
await store.ResolveInputRequestsAsync(created.TaskId, new Dictionary<string, InputResponse>
|
|
{
|
|
["req1"] = MakeResponse("response"),
|
|
}, CT);
|
|
|
|
var task = await store.GetTaskAsync(created.TaskId, CT);
|
|
Assert.NotNull(task);
|
|
Assert.Equal(McpTaskStatus.Completed, task.Status);
|
|
}
|
|
|
|
[Fact]
|
|
public async Task ResolveInputRequestsAsync_OnTerminalTask_DoesNotFireEvent()
|
|
{
|
|
var store = new InMemoryMcpTaskStore();
|
|
var created = await store.CreateTaskAsync(CT);
|
|
|
|
await store.SetCancelledAsync(created.TaskId, CT);
|
|
|
|
int eventCount = 0;
|
|
store.InputResponseReceived += _ => Interlocked.Increment(ref eventCount);
|
|
|
|
await store.ResolveInputRequestsAsync(created.TaskId, new Dictionary<string, InputResponse>
|
|
{
|
|
["req1"] = MakeResponse("response"),
|
|
}, CT);
|
|
|
|
Assert.Equal(0, eventCount);
|
|
}
|
|
|
|
[Fact]
|
|
public async Task SetInputRequestsAsync_OnTerminalTask_NoOps()
|
|
{
|
|
var store = new InMemoryMcpTaskStore();
|
|
var created = await store.CreateTaskAsync(CT);
|
|
|
|
await store.SetCancelledAsync(created.TaskId, CT);
|
|
|
|
await store.SetInputRequestsAsync(created.TaskId, new Dictionary<string, InputRequest>
|
|
{
|
|
["req1"] = MakeRequest("payload"),
|
|
}, CT);
|
|
|
|
var task = await store.GetTaskAsync(created.TaskId, CT);
|
|
Assert.NotNull(task);
|
|
Assert.Equal(McpTaskStatus.Cancelled, task.Status);
|
|
Assert.True(task.InputRequests is null || task.InputRequests.Count == 0);
|
|
}
|
|
}
|