1 files changed, 49 insertions, 6 deletions
diff --git a/src/Agent.Core/Tasks/TaskRegistry.cs b/src/Agent.Core/Tasks/TaskRegistry.cs
index b7a4afb..a89be6b 100644
--- a/src/Agent.Core/Tasks/TaskRegistry.cs
+++ b/src/Agent.Core/Tasks/TaskRegistry.cs
@@ -60,6 +60,12 @@ public sealed class TaskRegistry
     public bool AddUsage(Guid taskId, long inputTokens, long outputTokens) =>
         _tasks.TryGetValue(taskId, out MutableTask? task) && task.AddUsage(inputTokens, outputTokens);
 
+    public bool UpdateProgress(Guid taskId, int completedTasks, int totalTasks, string? currentTask, double criticalPathTokens) =>
+        _tasks.TryGetValue(taskId, out MutableTask? task) && task.UpdateProgress(completedTasks, totalTasks, currentTask, criticalPathTokens);
+
+    public bool UpdateDiagnostics(Guid taskId, int? iterations, int? replans, string? errorCode = null, string? error = null) =>
+        _tasks.TryGetValue(taskId, out MutableTask? task) && task.UpdateDiagnostics(iterations, replans, errorCode, error);
+
     public CancellationToken GetCancellationToken(Guid taskId) =>
         _tasks.TryGetValue(taskId, out MutableTask? task) ? task.CancellationToken : CancellationToken.None;
 
@@ -81,6 +87,13 @@ public sealed class TaskRegistry
         private string? _result;
         private long _inputTokens;
         private long _outputTokens;
+        private int _completedTasks;
+        private int _totalTasks = 1;
+        private double _criticalPathTokens;
+        private int _iterations;
+        private int _replans;
+        private string? _lastErrorCode;
+        private string? _lastError;
         private readonly CancellationTokenSource _cancellation = new();
         private readonly TaskCompletionSource _terminal = new(TaskCreationOptions.RunContinuationsAsynchronously);
 
@@ -93,8 +106,9 @@ public sealed class TaskRegistry
         {
             lock (_sync)
             {
-                int completed = _state == AgentTaskState.Completed ? 1 : 0;
-                return new(TaskId, SessionId, _state, request, temperature, _currentTask, _result, completed, 1, _inputTokens, _outputTokens, 0, images);
+                return new(TaskId, SessionId, _state, request, temperature, _currentTask, _result,
+                    _completedTasks, _totalTasks, _inputTokens, _outputTokens, _criticalPathTokens, images,
+                    _iterations, _replans, _lastErrorCode, _lastError);
             }
         }
 
@@ -122,16 +136,45 @@ public sealed class TaskRegistry
             }
         }
 
+        public bool UpdateProgress(int completedTasks, int totalTasks, string? currentTask, double criticalPathTokens)
+        {
+            if (completedTasks < 0 || totalTasks < 0 || completedTasks > totalTasks || criticalPathTokens < 0) return false;
+            lock (_sync)
+            {
+                _completedTasks = completedTasks;
+                _totalTasks = totalTasks;
+                _currentTask = currentTask;
+                _criticalPathTokens = criticalPathTokens;
+                return true;
+            }
+        }
+
+        public bool UpdateDiagnostics(int? iterations, int? replans, string? errorCode, string? error)
+        {
+            if (iterations < 0 || replans < 0) return false;
+            lock (_sync)
+            {
+                if (iterations.HasValue) _iterations = iterations.Value;
+                if (replans.HasValue) _replans = replans.Value;
+                if (errorCode is not null) _lastErrorCode = errorCode;
+                if (error is not null) _lastError = error;
+                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();
+                if (_state == AgentTaskState.Queued)
+                {
+                    _state = AgentTaskState.Cancelled;
+                    _currentTask = null;
+                    _result = "cancelled";
+                    _terminal.TrySetResult();
+                }
                 return true;
             }
         }