1 files changed, 34 insertions, 2 deletions
diff --git a/src/Agent.Core/Tasks/TaskRegistry.cs b/src/Agent.Core/Tasks/TaskRegistry.cs
index b61e525..b7a4afb 100644
--- a/src/Agent.Core/Tasks/TaskRegistry.cs
+++ b/src/Agent.Core/Tasks/TaskRegistry.cs
@@ -14,7 +14,7 @@ public sealed class TaskRegistry
 
     public SessionStore Sessions => _sessions;
 
-    public TaskSnapshot Enqueue(string request, Guid? requestedSessionId, double? temperature, IReadOnlyList<string>? images = null)
+    public TaskSnapshot Enqueue(string request, Guid? requestedSessionId, double? temperature, IReadOnlyList<string>? images = null, Guid? requestedTaskId = null)
     {
         if (string.IsNullOrWhiteSpace(request)) throw new ArgumentException("Request cannot be empty.", nameof(request));
         if (temperature is < 0 or > 2) throw new ArgumentOutOfRangeException(nameof(temperature));
@@ -29,7 +29,7 @@ public sealed class TaskRegistry
             sessionId = _sessions.Create().SessionId;
         }
 
-        var task = new MutableTask(Guid.NewGuid(), sessionId, request, temperature, images ?? []);
+        var task = new MutableTask(requestedTaskId ?? Guid.NewGuid(), sessionId, request, temperature, images ?? []);
         if (!_tasks.TryAdd(task.TaskId, task)) throw new InvalidOperationException("Could not register task.");
         if (!_queue.Writer.TryWrite(task.TaskId))
         {
@@ -60,6 +60,19 @@ public sealed class TaskRegistry
     public bool AddUsage(Guid taskId, long inputTokens, long outputTokens) =>
         _tasks.TryGetValue(taskId, out MutableTask? task) && task.AddUsage(inputTokens, outputTokens);
 
+    public CancellationToken GetCancellationToken(Guid taskId) =>
+        _tasks.TryGetValue(taskId, out MutableTask? task) ? task.CancellationToken : CancellationToken.None;
+
+    public bool Cancel(Guid taskId) =>
+        _tasks.TryGetValue(taskId, out MutableTask? task) && task.Cancel();
+
+    public async Task<TaskSnapshot> WaitForTerminalAsync(Guid taskId, CancellationToken cancellationToken)
+    {
+        if (!_tasks.TryGetValue(taskId, out MutableTask? task)) throw new KeyNotFoundException($"Task '{taskId}' was not found.");
+        await task.Terminal.WaitAsync(cancellationToken).ConfigureAwait(false);
+        return task.Snapshot();
+    }
+
     private sealed class MutableTask(Guid taskId, Guid sessionId, string request, double? temperature, IReadOnlyList<string> images)
     {
         private readonly object _sync = new();
@@ -68,9 +81,13 @@ public sealed class TaskRegistry
         private string? _result;
         private long _inputTokens;
         private long _outputTokens;
+        private readonly CancellationTokenSource _cancellation = new();
+        private readonly TaskCompletionSource _terminal = new(TaskCreationOptions.RunContinuationsAsynchronously);
 
         public Guid TaskId { get; } = taskId;
         public Guid SessionId { get; } = sessionId;
+        public CancellationToken CancellationToken => _cancellation.Token;
+        public Task Terminal => _terminal.Task;
 
         public TaskSnapshot Snapshot()
         {
@@ -100,6 +117,21 @@ public sealed class TaskRegistry
                 _state = next;
                 _currentTask = currentTask;
                 _result = result;
+                if (next is AgentTaskState.Completed or AgentTaskState.Failed or AgentTaskState.Cancelled) _terminal.TrySetResult();
+                return true;
+            }
+        }
+
+        public bool Cancel()
+        {
+            lock (_sync)
+            {
+                if (_state is AgentTaskState.Completed or AgentTaskState.Failed or AgentTaskState.Cancelled) return false;
+                _cancellation.Cancel();
+                _state = AgentTaskState.Cancelled;
+                _currentTask = null;
+                _result = "cancelled";
+                _terminal.TrySetResult();
                 return true;
             }
         }