diff options
Diffstat (limited to 'src/Agent.Mcp')
| -rw-r--r-- | src/Agent.Mcp/Agent.Mcp.csproj | 1 | ||||
| -rw-r--r-- | src/Agent.Mcp/AgentTools.cs | 65 | ||||
| -rw-r--r-- | src/Agent.Mcp/McpHost.cs | 93 | ||||
| -rw-r--r-- | src/Agent.Mcp/McpHostObservability.cs | 37 | ||||
| -rw-r--r-- | src/Agent.Mcp/PersistentMcpTaskStore.cs | 190 | ||||
| -rw-r--r-- | src/Agent.Mcp/QueuedTaskWorker.cs | 10 |
6 files changed, 372 insertions, 24 deletions
diff --git a/src/Agent.Mcp/Agent.Mcp.csproj b/src/Agent.Mcp/Agent.Mcp.csproj index 5b09556..df76e22 100644 --- a/src/Agent.Mcp/Agent.Mcp.csproj +++ b/src/Agent.Mcp/Agent.Mcp.csproj @@ -5,6 +5,7 @@ <ItemGroup> <PackageReference Include="ModelContextProtocol" /> <PackageReference Include="ModelContextProtocol.AspNetCore" /> + <PackageReference Include="ModelContextProtocol.Extensions.Tasks" /> </ItemGroup> <ItemGroup> <FrameworkReference Include="Microsoft.AspNetCore.App" /> diff --git a/src/Agent.Mcp/AgentTools.cs b/src/Agent.Mcp/AgentTools.cs index c44d6ec..6a5759f 100644 --- a/src/Agent.Mcp/AgentTools.cs +++ b/src/Agent.Mcp/AgentTools.cs @@ -3,34 +3,47 @@ using Agent.Core.Localization; using Agent.Core.Tasks; using ModelContextProtocol.Server; using ModelContextProtocol.Protocol; +using ModelContextProtocol.Extensions.Tasks; +using System.Text.Json; +using System.Text.Json.Nodes; namespace Agent.Mcp; [McpServerToolType] -public sealed class AgentTools(TaskRegistry registry, LocaleCatalog locale, VisionLimits visionLimits) +public sealed class AgentTools(TaskRegistry registry, LocaleCatalog locale, VisionLimits visionLimits, PersistentMcpTaskStore taskStore) { [McpServerTool(Name = "agent_run", ReadOnly = false, Destructive = true, OpenWorld = false, UseStructuredContent = false)] [Description("Queue a natural-language request for asynchronous execution.")] - public string AgentRun( + public async Task<string> AgentRun( RequestContext<CallToolRequestParams> context, [Description("Non-empty natural-language work request.")] string request, [Description("Existing session UUID, or omit to create a session.")] Guid? session_id = null, [Description("Task-scoped temperature from 0.0 through 2.0.")] double? temperature = null, - [Description("Optional PNG, JPEG, or WebP images.")] IReadOnlyList<McpImageInput>? images = null) + [Description("Optional PNG, JPEG, or WebP images.")] IReadOnlyList<McpImageInput>? images = null, + CancellationToken cancellationToken = default) { RequireOnly(context, "request", "session_id", "temperature", "images"); ValidateImages(images); - TaskSnapshot task = registry.Enqueue(request, session_id, temperature, images?.Select(static image => image.DataBase64).ToArray()); + Guid? mcpTaskId = GetRelatedTaskId(context.Params?.Meta) ?? GetRelatedTaskId(context.JsonRpcRequest) ?? + taskStore.ResolveTaskId(cancellationToken); + TaskSnapshot task = registry.Enqueue(request, session_id, temperature, images?.Select(static image => image.DataBase64).ToArray(), mcpTaskId); string message = locale.Format("mcp.task_queued", new Dictionary<string, object?> { ["task"] = task.TaskId, ["session"] = task.SessionId }); - return string.Join('\n', + string queued = string.Join('\n', $"task_id: {task.TaskId}", $"session_id: {task.SessionId}", "status: queued", $"message: {message}"); + if (mcpTaskId is null) return queued; + + using CancellationTokenRegistration registration = cancellationToken.Register(() => registry.Cancel(task.TaskId)); + TaskSnapshot terminal = await registry.WaitForTerminalAsync(task.TaskId, cancellationToken).ConfigureAwait(false); + if (terminal.State == AgentTaskState.Failed) throw new InvalidOperationException(terminal.Result ?? "Agent task failed."); + if (terminal.State == AgentTaskState.Cancelled) throw new OperationCanceledException("Agent task was cancelled.", cancellationToken); + return FormatStatus(terminal); } private void ValidateImages(IReadOnlyList<McpImageInput>? images) @@ -66,6 +79,11 @@ public sealed class AgentTools(TaskRegistry registry, LocaleCatalog locale, Visi $"message: {locale.Format("mcp.task_not_found", new Dictionary<string, object?> { ["task"] = task_id })}"); } + return FormatStatus(task); + } + + private string FormatStatus(TaskSnapshot task) + { string state = task.State.ToString().ToLowerInvariant(); bool terminal = task.State is AgentTaskState.Completed or AgentTaskState.Failed or AgentTaskState.Cancelled; return string.Join('\n', @@ -80,6 +98,43 @@ public sealed class AgentTools(TaskRegistry registry, LocaleCatalog locale, Visi terminal ? $"result: {task.Result}" : "result:"); } + private static Guid? GetRelatedTaskId(object? metadata) + { + if (metadata is null) return null; + JsonElement json = JsonSerializer.SerializeToElement(metadata); + return Find(json); + + static Guid? Find(JsonElement json) + { + if (json.ValueKind == JsonValueKind.Object) + { + foreach (JsonProperty property in json.EnumerateObject()) + { + if (property.NameEquals(TasksProtocol.MetaRelatedTask)) return Parse(property.Value); + Guid? nested = Find(property.Value); + if (nested is not null) return nested; + } + } + else if (json.ValueKind == JsonValueKind.Array) + { + foreach (JsonElement item in json.EnumerateArray()) + { + Guid? nested = Find(item); + if (nested is not null) return nested; + } + } + return null; + } + + static Guid? Parse(JsonElement related) + { + if (related.ValueKind == JsonValueKind.String && Guid.TryParse(related.GetString(), out Guid direct)) return direct; + if (related.ValueKind == JsonValueKind.Object && related.TryGetProperty("taskId", out JsonElement taskId) && + Guid.TryParse(taskId.GetString(), out Guid nested)) return nested; + return null; + } + } + private static void RequireOnly(RequestContext<CallToolRequestParams> context, params string[] allowed) { if (context.Params?.Arguments is null) return; diff --git a/src/Agent.Mcp/McpHost.cs b/src/Agent.Mcp/McpHost.cs index ff546e0..9a2ba66 100644 --- a/src/Agent.Mcp/McpHost.cs +++ b/src/Agent.Mcp/McpHost.cs @@ -15,12 +15,16 @@ using Microsoft.AspNetCore.Builder; using Microsoft.AspNetCore.Hosting; using Microsoft.AspNetCore.Http; using ModelContextProtocol.Server; +using ModelContextProtocol.Extensions.Tasks; +using ModelContextProtocol; +using ModelContextProtocol.Protocol; +using System.Diagnostics; namespace Agent.Mcp; public static class McpHost { - public static async Task RunStdioAsync(LocaleCatalog locale, OllamaSettings ollamaSettings, AgentRuntimeSettings runtimeSettings, VisionLimits visionLimits, CancellationToken cancellationToken, InternalToolDispatcher? toolDispatcher = null, IAgentRuntimeObserver? observer = null) + public static async Task RunStdioAsync(LocaleCatalog locale, OllamaSettings ollamaSettings, AgentRuntimeSettings runtimeSettings, VisionLimits visionLimits, CancellationToken cancellationToken, InternalToolDispatcher? toolDispatcher = null, IAgentRuntimeObserver? observer = null, McpTaskStoreSettings? taskStoreSettings = null, IMcpHostObserver? hostObserver = null, IReadOnlyList<McpServerTool>? additionalTools = null, IReadOnlySet<string>? taskCapableTools = null, IReadOnlySet<string>? requiredTaskTools = null, Action<IMcpServerBuilder>? configureBuilder = null) { HostApplicationBuilder builder = Host.CreateApplicationBuilder(new HostApplicationBuilderSettings { @@ -37,18 +41,24 @@ public static class McpHost builder.Services.AddSingleton(runtime); builder.Services.AddHostedService<QueuedTaskWorker>(); - McpServerTool[] tools = CreateTools(registry, locale, visionLimits); + var taskStore = new PersistentMcpTaskStore(taskStoreSettings ?? DefaultTaskStoreSettings()); + McpServerTool[] tools = CreateTools(registry, locale, visionLimits, taskStore).Concat(additionalTools ?? []).ToArray(); - builder.Services + IMcpServerBuilder mcpBuilder = builder.Services .AddMcpServer() .WithStdioServerTransport() .WithTools(tools); + ConfigureAgentTaskFilter(mcpBuilder); + configureBuilder?.Invoke(mcpBuilder); + mcpBuilder.WithTasks(taskStore, options => ConfigureTasks(options, taskCapableTools, requiredTaskTools)); using IHost host = builder.Build(); - await host.RunAsync(cancellationToken).ConfigureAwait(false); + await host.StartAsync(cancellationToken).ConfigureAwait(false); + hostObserver?.Ready(new McpReadyInfo("stdio", null, Environment.ProcessId, "2026-07-28", false, true)); + await host.WaitForShutdownAsync(cancellationToken).ConfigureAwait(false); } - public static async Task RunHttpAsync(LocaleCatalog locale, OllamaSettings ollamaSettings, AgentRuntimeSettings runtimeSettings, VisionLimits visionLimits, string hostName, int port, CancellationToken cancellationToken, InternalToolDispatcher? toolDispatcher = null, IAgentRuntimeObserver? observer = null) + public static async Task RunHttpAsync(LocaleCatalog locale, OllamaSettings ollamaSettings, AgentRuntimeSettings runtimeSettings, VisionLimits visionLimits, string hostName, int port, CancellationToken cancellationToken, InternalToolDispatcher? toolDispatcher = null, IAgentRuntimeObserver? observer = null, McpTaskStoreSettings? taskStoreSettings = null, IMcpHostObserver? hostObserver = null, IReadOnlyList<McpServerTool>? additionalTools = null, IReadOnlySet<string>? taskCapableTools = null, IReadOnlySet<string>? requiredTaskTools = null, Action<IMcpServerBuilder>? configureBuilder = null) { if (!System.Net.IPAddress.TryParse(hostName, out System.Net.IPAddress? address) || !System.Net.IPAddress.IsLoopback(address)) throw new ArgumentException("HTTP transport host must be a loopback IP literal.", nameof(hostName)); @@ -63,29 +73,76 @@ public static class McpHost builder.Services.AddSingleton(registry); builder.Services.AddSingleton(runtime); builder.Services.AddHostedService<QueuedTaskWorker>(); - McpServerTool[] tools = CreateTools(registry, locale, visionLimits); - builder.Services.AddMcpServer().WithHttpTransport().WithTools(tools); + var taskStore = new PersistentMcpTaskStore(taskStoreSettings ?? DefaultTaskStoreSettings()); + McpServerTool[] tools = CreateTools(registry, locale, visionLimits, taskStore).Concat(additionalTools ?? []).ToArray(); + IMcpServerBuilder mcpBuilder = builder.Services.AddMcpServer() + .WithHttpTransport(options => options.Stateless = true) + .WithTools(tools); + ConfigureAgentTaskFilter(mcpBuilder); + configureBuilder?.Invoke(mcpBuilder); + mcpBuilder.WithTasks(taskStore, options => ConfigureTasks(options, taskCapableTools, requiredTaskTools)); await using WebApplication app = builder.Build(); app.Use(async (context, next) => { - if (context.Request.Headers.TryGetValue("Origin", out Microsoft.Extensions.Primitives.StringValues origins) && - origins.Any(origin => !IsLoopbackOrigin(origin))) + var stopwatch = Stopwatch.StartNew(); + DateTimeOffset timestamp = DateTimeOffset.UtcNow; + hostObserver?.HttpRequestStarted(new McpHttpRequestStartInfo(context.Connection.Id, context.TraceIdentifier, + context.Connection.RemoteIpAddress?.ToString(), context.Connection.RemotePort, + context.Request.Method, context.Request.Path.Value ?? string.Empty)); + try { - context.Response.StatusCode = StatusCodes.Status403Forbidden; - return; + if (context.Request.Headers.TryGetValue("Origin", out Microsoft.Extensions.Primitives.StringValues origins) && + origins.Any(origin => !IsLoopbackOrigin(origin))) + { + context.Response.StatusCode = StatusCodes.Status403Forbidden; + } + else await next(context).ConfigureAwait(false); } - if (context.Request.Path == "/mcp" && HttpMethods.IsGet(context.Request.Method)) + finally { - context.Response.StatusCode = StatusCodes.Status405MethodNotAllowed; - return; + stopwatch.Stop(); + hostObserver?.HttpRequest(new McpHttpRequestInfo(timestamp, context.Connection.Id, context.TraceIdentifier, + context.Connection.RemoteIpAddress?.ToString(), context.Connection.RemotePort, + context.Connection.LocalIpAddress?.ToString(), context.Connection.LocalPort, + context.Request.Method, context.Request.Path.Value ?? string.Empty, context.Response.StatusCode, + stopwatch.Elapsed.TotalMilliseconds, context.Request.ContentType, context.Request.ContentLength, + context.Request.Headers.UserAgent.ToString(), context.Request.Headers["MCP-Protocol-Version"].ToString(), + context.Request.Headers["MCP-Method"].ToString(), context.Request.Headers["MCP-Name"].ToString())); } - await next(context).ConfigureAwait(false); }); app.MapMcp("/mcp"); - try { await app.RunAsync(cancellationToken).ConfigureAwait(false); } + try + { + await app.StartAsync(cancellationToken).ConfigureAwait(false); + hostObserver?.Ready(new McpReadyInfo("http", $"http://{listenHost}:{port}/mcp", Environment.ProcessId, "2026-07-28", true, true)); + await app.WaitForShutdownAsync(cancellationToken).ConfigureAwait(false); + } finally { httpClient.Dispose(); } } + private static McpTaskStoreSettings DefaultTaskStoreSettings() => + new(Path.Combine(Path.GetTempPath(), "agent-mcp", "mcp-tasks"), TimeSpan.FromHours(24), 1000); + + private static void ConfigureTasks(McpTasksOptions options, IReadOnlySet<string>? taskCapableTools, IReadOnlySet<string>? requiredTaskTools) + { + IReadOnlySet<string> names = taskCapableTools ?? new HashSet<string>(["agent_run"], StringComparer.Ordinal); + options.ExecutionModeSelector = context => + context.MatchedPrimitive is not McpServerTool tool ? McpTaskExecutionMode.Synchronous : + requiredTaskTools?.Contains(tool.ProtocolTool.Name) == true ? McpTaskExecutionMode.Required : + names.Contains(tool.ProtocolTool.Name) ? McpTaskExecutionMode.Optional : McpTaskExecutionMode.Synchronous; + } + + private static void ConfigureAgentTaskFilter(IMcpServerBuilder builder) + { + builder.WithRequestFilters(filters => filters.AddCallToolFilter(next => async (context, cancellationToken) => + { + CallToolResult result = await next(context, cancellationToken).ConfigureAwait(false); + if (context.Params?.Name == "agent_run" && result.IsError == true) + throw new McpProtocolException("The Agent task failed.", McpErrorCode.InternalError); + return result; + })); + } + private static bool IsLoopbackOrigin(string? origin) { return Uri.TryCreate(origin, UriKind.Absolute, out Uri? uri) && @@ -94,7 +151,7 @@ public static class McpHost (System.Net.IPAddress.TryParse(uri.Host, out System.Net.IPAddress? address) && System.Net.IPAddress.IsLoopback(address))); } - private static McpServerTool[] CreateTools(TaskRegistry registry, LocaleCatalog locale, VisionLimits visionLimits) + private static McpServerTool[] CreateTools(TaskRegistry registry, LocaleCatalog locale, VisionLimits visionLimits, PersistentMcpTaskStore taskStore) { var serializerOptions = new JsonSerializerOptions(JsonSerializerDefaults.Web) { @@ -102,7 +159,7 @@ public static class McpHost TypeInfoResolver = new DefaultJsonTypeInfoResolver(), UnmappedMemberHandling = JsonUnmappedMemberHandling.Disallow }; - var target = new AgentTools(registry, locale, visionLimits); + var target = new AgentTools(registry, locale, visionLimits, taskStore); McpServerTool[] tools = typeof(AgentTools) .GetMethods(BindingFlags.Instance | BindingFlags.Public | BindingFlags.DeclaredOnly) .Where(static method => method.GetCustomAttribute<McpServerToolAttribute>() is not null) diff --git a/src/Agent.Mcp/McpHostObservability.cs b/src/Agent.Mcp/McpHostObservability.cs new file mode 100644 index 0000000..17beb07 --- /dev/null +++ b/src/Agent.Mcp/McpHostObservability.cs @@ -0,0 +1,37 @@ +namespace Agent.Mcp; + +public sealed record McpReadyInfo(string Transport, string? Endpoint, int ProcessId, string ProtocolVersion, bool Stateless, bool LegacyFallback); + +public sealed record McpHttpRequestInfo( + DateTimeOffset Timestamp, + string ConnectionId, + string TraceId, + string? RemoteIp, + int? RemotePort, + string? LocalIp, + int? LocalPort, + string Method, + string Path, + int Status, + double DurationMs, + string? ContentType, + long? ContentLength, + string? UserAgent, + string? ProtocolVersion, + string? McpMethod, + string? McpName); + +public sealed record McpHttpRequestStartInfo( + string ConnectionId, + string TraceId, + string? RemoteIp, + int RemotePort, + string Method, + string Path); + +public interface IMcpHostObserver +{ + void Ready(McpReadyInfo info); + void HttpRequestStarted(McpHttpRequestStartInfo info); + void HttpRequest(McpHttpRequestInfo info); +} 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); +} diff --git a/src/Agent.Mcp/QueuedTaskWorker.cs b/src/Agent.Mcp/QueuedTaskWorker.cs index 23404e3..edea1ad 100644 --- a/src/Agent.Mcp/QueuedTaskWorker.cs +++ b/src/Agent.Mcp/QueuedTaskWorker.cs @@ -12,7 +12,15 @@ public sealed class QueuedTaskWorker(TaskRegistry registry, AgentRuntime runtime { if (!registry.Transition(taskId, AgentTaskState.Planning, "planning")) continue; registry.Transition(taskId, AgentTaskState.Running, "ollama chat"); - await runtime.ExecuteAsync(taskId, stoppingToken).ConfigureAwait(false); + using var linked = CancellationTokenSource.CreateLinkedTokenSource(stoppingToken, registry.GetCancellationToken(taskId)); + try + { + await runtime.ExecuteAsync(taskId, linked.Token).ConfigureAwait(false); + } + catch (OperationCanceledException) when (linked.IsCancellationRequested) + { + registry.Cancel(taskId); + } } } } |