1 files changed, 115 insertions, 0 deletions
diff --git a/src/Agent.Core/Tasks/TaskRegistry.cs b/src/Agent.Core/Tasks/TaskRegistry.cs
new file mode 100644
index 0000000..b61e525
--- /dev/null
+++ b/src/Agent.Core/Tasks/TaskRegistry.cs
@@ -0,0 +1,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
+        };
+    }
+}