1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
|
using System.Collections.Concurrent;
using System.Threading.Channels;
using Agent.Core.Sessions;
namespace Agent.Core.Tasks;
public sealed class TaskRegistry
{
private readonly ConcurrentDictionary<Guid, MutableTask> _tasks = new();
private readonly SessionStore _sessions;
private readonly Channel<Guid> _queue = Channel.CreateUnbounded<Guid>(new UnboundedChannelOptions { SingleReader = true, SingleWriter = false });
public TaskRegistry(SessionStore? sessions = null) => _sessions = sessions ?? new SessionStore();
public SessionStore Sessions => _sessions;
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));
Guid sessionId;
if (requestedSessionId is { } existing)
{
if (!_sessions.TryGet(existing, out _)) throw new KeyNotFoundException($"Session '{existing}' was not found.");
sessionId = existing;
}
else
{
sessionId = _sessions.Create().SessionId;
}
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))
{
_tasks.TryRemove(task.TaskId, out _);
throw new InvalidOperationException("Task queue is not accepting work.");
}
return task.Snapshot();
}
public bool TryGet(Guid taskId, Guid? sessionId, out TaskSnapshot? snapshot)
{
if (_tasks.TryGetValue(taskId, out MutableTask? task) && (sessionId is null || task.SessionId == sessionId))
{
snapshot = task.Snapshot();
return true;
}
snapshot = null;
return false;
}
public IAsyncEnumerable<Guid> ReadQueuedAsync(CancellationToken cancellationToken) => _queue.Reader.ReadAllAsync(cancellationToken);
public bool Transition(Guid taskId, AgentTaskState state, string? currentTask = null, string? result = null)
{
return _tasks.TryGetValue(taskId, out MutableTask? task) && task.Transition(state, currentTask, result);
}
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;
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();
private AgentTaskState _state = AgentTaskState.Queued;
private string? _currentTask;
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);
public Guid TaskId { get; } = taskId;
public Guid SessionId { get; } = sessionId;
public CancellationToken CancellationToken => _cancellation.Token;
public Task Terminal => _terminal.Task;
public TaskSnapshot Snapshot()
{
lock (_sync)
{
return new(TaskId, SessionId, _state, request, temperature, _currentTask, _result,
_completedTasks, _totalTasks, _inputTokens, _outputTokens, _criticalPathTokens, images,
_iterations, _replans, _lastErrorCode, _lastError);
}
}
public bool AddUsage(long inputTokens, long outputTokens)
{
if (inputTokens < 0 || outputTokens < 0) return false;
lock (_sync)
{
_inputTokens += inputTokens;
_outputTokens += outputTokens;
return true;
}
}
public bool Transition(AgentTaskState next, string? currentTask, string? result)
{
lock (_sync)
{
if (!Allowed(_state, next)) return false;
_state = next;
_currentTask = currentTask;
_result = result;
if (next is AgentTaskState.Completed or AgentTaskState.Failed or AgentTaskState.Cancelled) _terminal.TrySetResult();
return true;
}
}
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();
if (_state == AgentTaskState.Queued)
{
_state = AgentTaskState.Cancelled;
_currentTask = null;
_result = "cancelled";
_terminal.TrySetResult();
}
return true;
}
}
private static bool Allowed(AgentTaskState current, AgentTaskState next) => current switch
{
AgentTaskState.Queued => next is AgentTaskState.Planning or AgentTaskState.Failed or AgentTaskState.Cancelled,
AgentTaskState.Planning => next is AgentTaskState.Running or AgentTaskState.Failed or AgentTaskState.Cancelled,
AgentTaskState.Running => next is AgentTaskState.Planning or AgentTaskState.Completed or AgentTaskState.Failed or AgentTaskState.Cancelled,
_ => false
};
}
}
|
