1 files changed, 179 insertions, 0 deletions
diff --git a/src/Agent.Planner/TaskGraph.cs b/src/Agent.Planner/TaskGraph.cs
new file mode 100644
index 0000000..7a4d45b
--- /dev/null
+++ b/src/Agent.Planner/TaskGraph.cs
@@ -0,0 +1,179 @@
+namespace Agent.Planner;
+
+public sealed class TaskGraph(int maxTasks = 32)
+{
+    private readonly Dictionary<Guid, MutableTask> _tasks = [];
+    private readonly List<Guid> _insertionOrder = [];
+    private CriticalPathResult _criticalPath = new([], 0);
+
+    public int MaxTasks { get; } = maxTasks > 0 ? maxTasks : throw new ArgumentOutOfRangeException(nameof(maxTasks));
+    public int Count => _tasks.Count;
+    public CriticalPathResult CriticalPath => _criticalPath;
+    public long TotalTokenBudget => _tasks.Values.Aggregate(0L, static (total, task) => checked(total + task.BudgetClass.MaxTokens()));
+    public IReadOnlyList<PlannerTask> Tasks => _insertionOrder.Select(id => _tasks[id].Snapshot()).ToArray();
+
+    public PlannerTask AddTask(
+        string description,
+        TokenBudgetClass budgetClass,
+        TokenEstimate estimate,
+        IEnumerable<Guid>? dependencies = null,
+        bool requiresExternalEffect = false,
+        Guid? id = null)
+    {
+        if (string.IsNullOrWhiteSpace(description)) throw new PlannerValidationException("Task description is required.");
+        if (_tasks.Count >= MaxTasks) throw new PlannerValidationException($"Task limit of {MaxTasks} exceeded.");
+        estimate.Validate(budgetClass);
+        Guid taskId = id ?? Guid.NewGuid();
+        if (_tasks.ContainsKey(taskId)) throw new PlannerValidationException($"Task '{taskId}' already exists.");
+        HashSet<Guid> dependencySet = dependencies?.ToHashSet() ?? [];
+        EnsureDependenciesExist(dependencySet);
+        if (dependencySet.Contains(taskId)) throw new PlannerValidationException("A task cannot depend on itself.");
+
+        _tasks.Add(taskId, new MutableTask(taskId, description, dependencySet, budgetClass, estimate, requiresExternalEffect));
+        _insertionOrder.Add(taskId);
+        try
+        {
+            Recalculate();
+        }
+        catch
+        {
+            _tasks.Remove(taskId);
+            _insertionOrder.Remove(taskId);
+            throw;
+        }
+        return _tasks[taskId].Snapshot();
+    }
+
+    public void AddDependency(Guid taskId, Guid dependencyId)
+    {
+        MutableTask task = GetMutable(taskId);
+        EnsureMutable(task);
+        _ = GetMutable(dependencyId);
+        if (!task.Dependencies.Add(dependencyId)) return;
+        try
+        {
+            Recalculate();
+        }
+        catch
+        {
+            task.Dependencies.Remove(dependencyId);
+            throw;
+        }
+    }
+
+    public void SetEstimate(Guid taskId, TokenBudgetClass budgetClass, TokenEstimate estimate)
+    {
+        MutableTask task = GetMutable(taskId);
+        EnsureMutable(task);
+        estimate.Validate(budgetClass);
+        (TokenBudgetClass previousBudget, TokenEstimate previousEstimate) = (task.BudgetClass, task.Estimate);
+        task.BudgetClass = budgetClass;
+        task.Estimate = estimate;
+        try { Recalculate(); }
+        catch
+        {
+            task.BudgetClass = previousBudget;
+            task.Estimate = previousEstimate;
+            throw;
+        }
+    }
+
+    public void SetState(Guid taskId, PlannerTaskState state)
+    {
+        MutableTask task = GetMutable(taskId);
+        EnsureMutable(task);
+        task.State = state;
+    }
+
+    public void AddUsage(Guid taskId, long inputTokens, long outputTokens)
+    {
+        if (inputTokens < 0 || outputTokens < 0) throw new ArgumentOutOfRangeException(nameof(inputTokens), "Token usage cannot be negative.");
+        MutableTask task = GetMutable(taskId);
+        EnsureMutable(task);
+        task.ActualInputTokens = checked(task.ActualInputTokens + inputTokens);
+        task.ActualOutputTokens = checked(task.ActualOutputTokens + outputTokens);
+    }
+
+    public PlannerTask Get(Guid taskId) => GetMutable(taskId).Snapshot();
+
+    private void Recalculate()
+    {
+        IReadOnlyList<Guid> order = TopologicalOrder();
+        var costs = new Dictionary<Guid, double>();
+        var predecessors = new Dictionary<Guid, Guid?>();
+        foreach (Guid id in order)
+        {
+            MutableTask task = _tasks[id];
+            Guid? best = null;
+            double bestCost = 0;
+            foreach (Guid dependency in task.Dependencies.OrderBy(InsertionIndex))
+            {
+                double candidate = costs[dependency];
+                if (best is null || candidate > bestCost)
+                {
+                    best = dependency;
+                    bestCost = candidate;
+                }
+            }
+            costs[id] = bestCost + task.Estimate.ExpectedTokens;
+            predecessors[id] = best;
+        }
+        if (order.Count == 0)
+        {
+            _criticalPath = new([], 0);
+            return;
+        }
+        Guid end = order[0];
+        foreach (Guid id in order.Skip(1)) if (costs[id] > costs[end]) end = id;
+        var path = new List<Guid>();
+        for (Guid? current = end; current is not null; current = predecessors[current.Value]) path.Add(current.Value);
+        path.Reverse();
+        _criticalPath = new(path, costs[end]);
+    }
+
+    private IReadOnlyList<Guid> TopologicalOrder()
+    {
+        var indegree = _tasks.ToDictionary(static pair => pair.Key, static pair => pair.Value.Dependencies.Count);
+        var ready = new SortedSet<int>(_insertionOrder.Select((id, index) => (id, index)).Where(item => indegree[item.id] == 0).Select(static item => item.index));
+        var result = new List<Guid>(_tasks.Count);
+        while (ready.Count > 0)
+        {
+            int index = ready.Min;
+            ready.Remove(index);
+            Guid id = _insertionOrder[index];
+            result.Add(id);
+            foreach (Guid dependent in _insertionOrder.Where(candidate => _tasks[candidate].Dependencies.Contains(id)))
+            {
+                indegree[dependent]--;
+                if (indegree[dependent] == 0) ready.Add(InsertionIndex(dependent));
+            }
+        }
+        if (result.Count != _tasks.Count) throw new PlannerValidationException("TaskGraph must be a DAG; the update creates a cycle.");
+        return result;
+    }
+
+    private int InsertionIndex(Guid id) => _insertionOrder.IndexOf(id);
+    private MutableTask GetMutable(Guid id) => _tasks.TryGetValue(id, out MutableTask? task) ? task : throw new KeyNotFoundException($"Task '{id}' was not found.");
+    private void EnsureDependenciesExist(IEnumerable<Guid> dependencies)
+    {
+        foreach (Guid dependency in dependencies) _ = GetMutable(dependency);
+    }
+    private static void EnsureMutable(MutableTask task)
+    {
+        if (task.State == PlannerTaskState.Completed) throw new PlannerValidationException("Completed tasks are immutable.");
+    }
+
+    private sealed class MutableTask(Guid id, string description, HashSet<Guid> dependencies, TokenBudgetClass budgetClass, TokenEstimate estimate, bool requiresExternalEffect)
+    {
+        public Guid Id { get; } = id;
+        public string Description { get; } = description;
+        public HashSet<Guid> Dependencies { get; } = dependencies;
+        public PlannerTaskState State { get; set; } = PlannerTaskState.Queued;
+        public TokenBudgetClass BudgetClass { get; set; } = budgetClass;
+        public TokenEstimate Estimate { get; set; } = estimate;
+        public long ActualInputTokens { get; set; }
+        public long ActualOutputTokens { get; set; }
+        public bool RequiresExternalEffect { get; } = requiresExternalEffect;
+        public PlannerTask Snapshot() => new(Id, Description, Dependencies.ToHashSet(), State, BudgetClass, Estimate, ActualInputTokens, ActualOutputTokens, RequiresExternalEffect);
+    }
+}