1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
|
using System.Collections.Concurrent;
using System.Text.Json;
using ModelContextProtocol.Extensions.Tasks;
using ModelContextProtocol.Protocol;
namespace Agent.Mcp;
public sealed record McpTaskStoreSettings(string Directory, TimeSpan TimeToLive, long PollIntervalMs);
public sealed class PersistentMcpTaskStore : IMcpTaskStore
{
private readonly ConcurrentDictionary<CancellationToken, Guid> _taskIdsByCancellation = new();
private readonly ConcurrentDictionary<string, StoredTask> _tasks = new(StringComparer.Ordinal);
private readonly McpTaskStoreSettings _settings;
private readonly JsonSerializerOptions _json = new(JsonSerializerDefaults.Web) { WriteIndented = true };
public PersistentMcpTaskStore(McpTaskStoreSettings settings)
{
_settings = settings;
System.IO.Directory.CreateDirectory(settings.Directory);
Recover();
}
public event Action<InputResponseReceivedEventArgs>? InputResponseReceived;
public Guid? ResolveTaskId(CancellationToken cancellationToken)
{
if (!_taskIdsByCancellation.TryRemove(cancellationToken, out Guid taskId)) return null;
return taskId;
}
public async Task<McpTaskInfo> CreateTaskAsync(CancellationToken cancellationToken)
{
DateTimeOffset now = DateTimeOffset.UtcNow;
var task = new StoredTask(Guid.NewGuid().ToString("D"), McpTaskStatus.Working, now, now,
_settings.TimeToLive, _settings.PollIntervalMs, "working", null, null, null);
if (!_tasks.TryAdd(task.TaskId, task)) throw new InvalidOperationException("Could not allocate an MCP task ID.");
await SaveAsync(task, cancellationToken).ConfigureAwait(false);
Guid createdId = Guid.Parse(task.TaskId);
_taskIdsByCancellation[cancellationToken] = createdId;
return ToInfo(task);
}
public Task<McpTaskInfo?> GetTaskAsync(string taskId, CancellationToken cancellationToken)
{
cancellationToken.ThrowIfCancellationRequested();
if (!_tasks.TryGetValue(taskId, out StoredTask? task)) return Task.FromResult<McpTaskInfo?>(null);
if (IsExpired(task))
{
_tasks.TryRemove(taskId, out _);
Delete(taskId);
return Task.FromResult<McpTaskInfo?>(null);
}
return Task.FromResult<McpTaskInfo?>(ToInfo(task));
}
public Task SetCompletedAsync(string taskId, JsonElement result, CancellationToken cancellationToken) =>
UpdateAsync(taskId, McpTaskStatus.Completed, "completed", result, null, null, cancellationToken);
public Task SetFailedAsync(string taskId, JsonElement error, CancellationToken cancellationToken) =>
UpdateAsync(taskId, McpTaskStatus.Failed, "failed", null, error, null, cancellationToken);
public async Task<bool> SetCancelledAsync(string taskId, CancellationToken cancellationToken)
{
if (!_tasks.ContainsKey(taskId)) return false;
await UpdateAsync(taskId, McpTaskStatus.Cancelled, "cancelled", null, null, null, cancellationToken).ConfigureAwait(false);
return true;
}
public async Task ResolveInputRequestsAsync(string taskId, IDictionary<string, InputResponse> responses, CancellationToken cancellationToken)
{
if (!_tasks.TryGetValue(taskId, out StoredTask? task)) return;
var pending = task.InputRequests is null
? new Dictionary<string, InputRequest>(StringComparer.Ordinal)
: new Dictionary<string, InputRequest>(task.InputRequests, StringComparer.Ordinal);
var accepted = new List<KeyValuePair<string, InputResponse>>();
foreach ((string requestId, InputResponse response) in responses)
{
if (!pending.Remove(requestId)) continue;
accepted.Add(new KeyValuePair<string, InputResponse>(requestId, response));
}
await UpdateAsync(taskId, pending.Count == 0 ? McpTaskStatus.Working : McpTaskStatus.InputRequired,
pending.Count == 0 ? "working" : "input required", task.Result, task.Error,
pending.Count == 0 ? null : pending, cancellationToken).ConfigureAwait(false);
foreach ((string requestId, InputResponse response) in accepted)
InputResponseReceived?.Invoke(new InputResponseReceivedEventArgs { TaskId = taskId, RequestId = requestId, Response = response });
}
public Task SetInputRequestsAsync(string taskId, IDictionary<string, InputRequest> requests, CancellationToken cancellationToken) =>
UpdateAsync(taskId, McpTaskStatus.InputRequired, "input required", null, null,
new Dictionary<string, InputRequest>(requests, StringComparer.Ordinal), cancellationToken);
private async Task UpdateAsync(string taskId, McpTaskStatus status, string message, JsonElement? result,
JsonElement? error, IReadOnlyDictionary<string, InputRequest>? inputRequests, CancellationToken cancellationToken)
{
if (!_tasks.TryGetValue(taskId, out StoredTask? current)) throw new KeyNotFoundException($"MCP task '{taskId}' was not found.");
var updated = current with
{
Status = status,
LastUpdatedAt = DateTimeOffset.UtcNow,
StatusMessage = message,
Result = result,
Error = error,
InputRequests = inputRequests
};
_tasks[taskId] = updated;
await SaveAsync(updated, cancellationToken).ConfigureAwait(false);
}
private void Recover()
{
foreach (string path in System.IO.Directory.EnumerateFiles(_settings.Directory, "*.json", SearchOption.TopDirectoryOnly))
{
try
{
StoredTask? task = JsonSerializer.Deserialize<StoredTask>(File.ReadAllText(path), _json);
if (task is null || IsExpired(task)) { File.Delete(path); continue; }
if (task.Status is McpTaskStatus.Working or McpTaskStatus.InputRequired)
{
JsonElement error = JsonSerializer.SerializeToElement(new { code = "server_restarted", message = "The server restarted before the task completed." });
task = task with { Status = McpTaskStatus.Failed, StatusMessage = "server_restarted", Error = error, LastUpdatedAt = DateTimeOffset.UtcNow, InputRequests = null };
SaveSynchronously(task);
}
_tasks[task.TaskId] = task;
}
catch (JsonException)
{
File.Move(path, path + ".invalid", true);
}
}
}
private bool IsExpired(StoredTask task) => task.TimeToLive is { } ttl && DateTimeOffset.UtcNow - task.CreatedAt >= ttl;
private async ValueTask SaveAsync(StoredTask task, CancellationToken cancellationToken)
{
string path = PathFor(task.TaskId);
string temp = path + ".tmp-" + Guid.NewGuid().ToString("N");
try
{
await using (var stream = new FileStream(temp, FileMode.CreateNew, FileAccess.Write, FileShare.None, 4096,
FileOptions.Asynchronous | FileOptions.WriteThrough))
{
await JsonSerializer.SerializeAsync(stream, task, _json, cancellationToken).ConfigureAwait(false);
await stream.FlushAsync(cancellationToken).ConfigureAwait(false);
}
File.Move(temp, path, true);
}
finally { if (File.Exists(temp)) File.Delete(temp); }
}
private void SaveSynchronously(StoredTask task)
{
string path = PathFor(task.TaskId);
string temp = path + ".tmp-" + Guid.NewGuid().ToString("N");
try
{
using (var stream = new FileStream(temp, FileMode.CreateNew, FileAccess.Write, FileShare.None, 4096, FileOptions.WriteThrough))
{
JsonSerializer.Serialize(stream, task, _json);
stream.Flush(true);
}
File.Move(temp, path, true);
}
finally { if (File.Exists(temp)) File.Delete(temp); }
}
private void Delete(string taskId)
{
string path = PathFor(taskId);
if (File.Exists(path)) File.Delete(path);
}
private string PathFor(string taskId) => Path.Combine(_settings.Directory, taskId + ".json");
private static McpTaskInfo ToInfo(StoredTask task) => new(task.TaskId, task.Status, task.CreatedAt, task.LastUpdatedAt,
task.TimeToLive, task.PollIntervalMs, task.StatusMessage, task.Result, task.Error, task.InputRequests);
private sealed record StoredTask(
string TaskId,
McpTaskStatus Status,
DateTimeOffset CreatedAt,
DateTimeOffset LastUpdatedAt,
TimeSpan? TimeToLive,
long? PollIntervalMs,
string StatusMessage,
JsonElement? Result,
JsonElement? Error,
IReadOnlyDictionary<string, InputRequest>? InputRequests);
}
|
