summaryrefslogtreecommitdiff
path: root/src/Agent.Planner/TaskGraph.cs
blob: a4cdc5fff07bea12d97570cda67b691e846ec36e (plain)
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
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,
        IEnumerable<string>? requiredTools = 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() ?? [];
        HashSet<string> tools = requiredTools?.ToHashSet(StringComparer.Ordinal) ?? [];
        if (requiresExternalEffect && tools.Count == 0) throw new PlannerValidationException("An external-effect Task must name at least one required tool.");
        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, tools));
        _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, HashSet<string> requiredTools)
    {
        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 HashSet<string> RequiredTools { get; } = requiredTools;
        public PlannerTask Snapshot() => new(Id, Description, Dependencies.ToHashSet(), State, BudgetClass, Estimate, ActualInputTokens, ActualOutputTokens, RequiresExternalEffect, RequiredTools.ToHashSet(StringComparer.Ordinal));
    }
}