1 files changed, 190 insertions, 0 deletions
diff --git a/src/Agent.Mcp/PersistentMcpTaskStore.cs b/src/Agent.Mcp/PersistentMcpTaskStore.cs
new file mode 100644
index 0000000..73608d4
--- /dev/null
+++ b/src/Agent.Mcp/PersistentMcpTaskStore.cs
@@ -0,0 +1,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);
+}