1 files changed, 60 insertions, 5 deletions

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;