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
|
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 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 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)
{
int completed = _state == AgentTaskState.Completed ? 1 : 0;
return new(TaskId, SessionId, _state, request, temperature, _currentTask, _result, completed, 1, _inputTokens, _outputTokens, 0, images);
}
}
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 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;
}
}
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
};
}
}
|
