deepstrike_core/orchestration/
task_graph.rs

1use crate::types::error::{DeepStrikeError, Result};
2use crate::types::result::LoopResult;
3use crate::types::task::RuntimeTask;
4
5#[derive(Debug, Clone, Copy, PartialEq, Eq)]
6pub enum TaskStatus {
7    Pending,
8    Ready,
9    Running,
10    Completed,
11    Failed,
12}
13
14#[derive(Debug, Clone)]
15pub struct TaskNode {
16    pub id: usize,
17    pub task: RuntimeTask,
18    pub status: TaskStatus,
19    pub result: Option<LoopResult>,
20    pub dependencies: Vec<usize>,
21}
22
23/// DAG of tasks with dependency tracking.
24/// Maintains an in-degree counter so `ready_tasks()` is O(1) amortized.
25pub struct TaskGraph {
26    nodes: Vec<TaskNode>,
27    /// Number of incomplete dependencies per task.
28    in_degree: Vec<usize>,
29}
30
31impl TaskGraph {
32    pub fn new() -> Self {
33        Self {
34            nodes: Vec::new(),
35            in_degree: Vec::new(),
36        }
37    }
38
39    /// Add a task, returns its ID. Duplicate dependency entries are collapsed: `in_degree` counts
40    /// entries but [`complete`](Self::complete) decrements once per completed dependency, so a
41    /// duplicated entry would leave the node permanently below its own in-degree (a silent stall).
42    pub fn add(&mut self, task: RuntimeTask, mut dependencies: Vec<usize>) -> usize {
43        let mut seen = std::collections::HashSet::new();
44        dependencies.retain(|d| seen.insert(*d));
45        let id = self.nodes.len();
46        let deg = dependencies.len();
47        self.nodes.push(TaskNode {
48            id,
49            task,
50            status: if deg == 0 {
51                TaskStatus::Ready
52            } else {
53                TaskStatus::Pending
54            },
55            result: None,
56            dependencies,
57        });
58        self.in_degree.push(deg);
59        id
60    }
61
62    /// Topological sort — returns ordered IDs or error if cycle detected.
63    pub fn topological_sort(&self) -> Result<Vec<usize>> {
64        let n = self.nodes.len();
65        let mut in_deg = self.in_degree.clone();
66        let mut adj: Vec<Vec<usize>> = vec![Vec::new(); n];
67
68        for node in &self.nodes {
69            for &dep in &node.dependencies {
70                adj[dep].push(node.id);
71            }
72        }
73
74        let mut queue: Vec<usize> = (0..n).filter(|&i| in_deg[i] == 0).collect();
75        let mut order = Vec::with_capacity(n);
76
77        while let Some(id) = queue.pop() {
78            order.push(id);
79            for &next in &adj[id] {
80                in_deg[next] -= 1;
81                if in_deg[next] == 0 {
82                    queue.push(next);
83                }
84            }
85        }
86
87        if order.len() != n {
88            return Err(DeepStrikeError::OrchestrationCycle);
89        }
90        Ok(order)
91    }
92
93    /// Return IDs of tasks that are Ready (deps satisfied, not yet started).
94    pub fn ready_tasks(&self) -> Vec<usize> {
95        self.nodes
96            .iter()
97            .filter(|n| n.status == TaskStatus::Ready)
98            .map(|n| n.id)
99            .collect()
100    }
101
102    /// Mark a task as running.
103    pub fn start(&mut self, task_id: usize) {
104        if let Some(node) = self.nodes.get_mut(task_id) {
105            node.status = TaskStatus::Running;
106        }
107    }
108
109    /// Re-mark a (running) task as Ready without touching dependents — used to re-arm a loop node
110    /// for its next iteration. Unlike [`complete`](Self::complete), this does NOT decrement any
111    /// in-degree, so the loop node's dependents stay pending until the loop finally `complete`s.
112    pub fn set_ready(&mut self, task_id: usize) {
113        if let Some(node) = self.nodes.get_mut(task_id) {
114            node.status = TaskStatus::Ready;
115        }
116    }
117
118    /// Mark a task as completed; promote dependents whose in-degree reaches 0.
119    ///
120    /// Idempotent: a task already terminal (Completed/Failed) is left untouched — a duplicate
121    /// completion (at-least-once event delivery, resume replay) must not double-decrement its
122    /// dependents' in-degree, which would underflow (debug panic) or over-promote gated nodes.
123    pub fn complete(&mut self, task_id: usize, result: LoopResult) {
124        let Some(node) = self.nodes.get_mut(task_id) else {
125            return;
126        };
127        if matches!(node.status, TaskStatus::Completed | TaskStatus::Failed) {
128            return;
129        }
130        node.status = TaskStatus::Completed;
131        node.result = Some(result);
132        // Collect dependents first to avoid borrow conflict
133        let dependents: Vec<usize> = self
134            .nodes
135            .iter()
136            .filter(|n| n.dependencies.contains(&task_id))
137            .map(|n| n.id)
138            .collect();
139        for dep_id in dependents {
140            self.in_degree[dep_id] -= 1;
141            if self.in_degree[dep_id] == 0 {
142                if let Some(n) = self.nodes.get_mut(dep_id) {
143                    if n.status == TaskStatus::Pending {
144                        n.status = TaskStatus::Ready;
145                    }
146                }
147            }
148        }
149    }
150
151    /// Mark a task as failed (dependents remain Pending — caller decides policy). Terminal states
152    /// are sticky: failing an already-completed task must not un-complete it (idempotency twin of
153    /// [`complete`](Self::complete)).
154    pub fn fail(&mut self, task_id: usize) {
155        if let Some(node) = self.nodes.get_mut(task_id) {
156            if !matches!(node.status, TaskStatus::Completed | TaskStatus::Failed) {
157                node.status = TaskStatus::Failed;
158            }
159        }
160    }
161
162    pub fn get(&self, task_id: usize) -> Option<&TaskNode> {
163        self.nodes.get(task_id)
164    }
165
166    pub fn len(&self) -> usize {
167        self.nodes.len()
168    }
169
170    pub fn is_empty(&self) -> bool {
171        self.nodes.is_empty()
172    }
173
174    pub fn all_done(&self) -> bool {
175        self.nodes
176            .iter()
177            .all(|n| matches!(n.status, TaskStatus::Completed | TaskStatus::Failed))
178    }
179}
180
181impl Default for TaskGraph {
182    fn default() -> Self {
183        Self::new()
184    }
185}
186
187#[cfg(test)]
188mod tests {
189    use super::*;
190
191    #[test]
192    fn topological_sort_linear() {
193        let mut g = TaskGraph::new();
194        let a = g.add(RuntimeTask::new("A"), vec![]);
195        let b = g.add(RuntimeTask::new("B"), vec![a]);
196        let c = g.add(RuntimeTask::new("C"), vec![b]);
197
198        let order = g.topological_sort().unwrap();
199        assert_eq!(order, vec![0, 1, 2]);
200        let _ = (a, c);
201    }
202
203    #[test]
204    fn detects_cycle() {
205        let mut g = TaskGraph::new();
206        g.nodes.push(TaskNode {
207            id: 0,
208            task: RuntimeTask::new("A"),
209            status: TaskStatus::Pending,
210            result: None,
211            dependencies: vec![1],
212        });
213        g.nodes.push(TaskNode {
214            id: 1,
215            task: RuntimeTask::new("B"),
216            status: TaskStatus::Pending,
217            result: None,
218            dependencies: vec![0],
219        });
220        g.in_degree.push(1);
221        g.in_degree.push(1);
222
223        assert!(g.topological_sort().is_err());
224    }
225
226    #[test]
227    fn ready_tasks_respects_deps() {
228        let mut g = TaskGraph::new();
229        let a = g.add(RuntimeTask::new("A"), vec![]);
230        let _b = g.add(RuntimeTask::new("B"), vec![a]);
231
232        assert_eq!(g.ready_tasks(), vec![0]); // only A is Ready
233    }
234
235    #[test]
236    fn set_ready_rearms_without_promoting_dependents() {
237        let mut g = TaskGraph::new();
238        let a = g.add(RuntimeTask::new("A"), vec![]); // loop node
239        let b = g.add(RuntimeTask::new("B"), vec![a]); // dependent
240        g.start(a);
241        // Re-arm A for its next iteration: A is Ready again, but B stays Pending (no promotion).
242        g.set_ready(a);
243        assert_eq!(g.nodes[a].status, TaskStatus::Ready);
244        assert_eq!(g.nodes[b].status, TaskStatus::Pending);
245        assert_eq!(g.ready_tasks(), vec![a]);
246    }
247
248    #[test]
249    fn complete_promotes_dependent() {
250        use crate::types::result::{LoopResult, TerminationReason};
251        let mut g = TaskGraph::new();
252        let a = g.add(RuntimeTask::new("A"), vec![]);
253        let b = g.add(RuntimeTask::new("B"), vec![a]);
254
255        assert_eq!(g.nodes[b].status, TaskStatus::Pending);
256        g.complete(
257            a,
258            LoopResult {
259                termination: TerminationReason::Completed,
260                final_message: None,
261                turns_used: 1,
262                total_tokens_used: 0,
263                loop_continue: None,
264                classify_branch: None,
265                tournament_winner: None,
266                pace_decision: None,
267            },
268        );
269        assert_eq!(g.nodes[b].status, TaskStatus::Ready);
270    }
271
272    #[test]
273    fn duplicate_complete_is_idempotent() {
274        use crate::types::result::{LoopResult, TerminationReason};
275        let result = || LoopResult {
276            termination: TerminationReason::Completed,
277            final_message: None,
278            turns_used: 1,
279            total_tokens_used: 0,
280            loop_continue: None,
281            classify_branch: None,
282            tournament_winner: None,
283            pace_decision: None,
284        };
285        // b gates on BOTH a and c; a duplicate completion of `a` must not stand in for `c`.
286        let mut g = TaskGraph::new();
287        let a = g.add(RuntimeTask::new("A"), vec![]);
288        let c = g.add(RuntimeTask::new("C"), vec![]);
289        let b = g.add(RuntimeTask::new("B"), vec![a, c]);
290
291        g.complete(a, result());
292        g.complete(a, result()); // duplicate delivery — no double decrement, no panic
293        assert_eq!(g.nodes[b].status, TaskStatus::Pending);
294        g.complete(c, result());
295        assert_eq!(g.nodes[b].status, TaskStatus::Ready);
296        // Terminal states are sticky both ways.
297        g.fail(a);
298        assert_eq!(g.nodes[a].status, TaskStatus::Completed);
299    }
300}
deepstrike_core/orchestration/task_graph.rs

deepstrike_core/orchestration/
task_graph.rs
