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
|
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)
{
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(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);
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;
public Guid TaskId { get; } = taskId;
public Guid SessionId { get; } = sessionId;
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;
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
};
}
}
|
