deepstrike_core/orchestration/
task_graph.rs

1use std::cmp::Ordering;
2use std::collections::{BinaryHeap, HashSet};
3
4use crate::scheduler::policy::SchedulerPolicyConfig;
5use crate::types::error::{DeepStrikeError, Result};
6use crate::types::result::LoopResult;
7use crate::types::task::RuntimeTask;
8
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub enum TaskStatus {
11    Pending,
12    Ready,
13    Running,
14    Completed,
15    CompletedPartial,
16    Failed,
17    SkippedUpstreamFailed,
18}
19
20impl TaskStatus {
21    pub fn is_terminal(self) -> bool {
22        matches!(
23            self,
24            Self::Completed | Self::CompletedPartial | Self::Failed | Self::SkippedUpstreamFailed
25        )
26    }
27}
28
29#[derive(Debug, Clone)]
30pub struct TaskNode {
31    pub id: usize,
32    pub task: RuntimeTask,
33    pub status: TaskStatus,
34    pub result: Option<LoopResult>,
35    pub dependencies: Vec<usize>,
36}
37
38/// DAG of tasks with dependency tracking.
39/// Maintains persistent reverse adjacency and a deterministic ready heap. Completing a node visits
40/// only its outgoing dependents; selecting ready work never scans the graph.
41pub struct TaskGraph {
42    nodes: Vec<TaskNode>,
43    /// Number of dependencies that have not completed successfully per task. Workflow-level
44    /// policies handle partial/failure terminal states explicitly.
45    in_degree: Vec<usize>,
46    /// Persistent dependency → dependents index. Terminal promotion touches only outgoing edges.
47    reverse_adjacency: Vec<Vec<usize>>,
48    ready_heap: BinaryHeap<ReadyEntry>,
49    ready_generation: Vec<u64>,
50    enqueued_round: Vec<u64>,
51    enqueue_sequence: u64,
52    ready_round: u64,
53    scheduling: Vec<SchedulingMetadata>,
54    scheduler_policy: SchedulerPolicyConfig,
55}
56
57#[derive(Debug, Clone, Copy, Default)]
58struct SchedulingMetadata {
59    critical_path_remaining: u64,
60    downstream_fanout: u64,
61    token_cost: u64,
62}
63
64#[derive(Debug, Clone, Copy, PartialEq, Eq)]
65struct ReadyEntry {
66    priority: i128,
67    enqueue_sequence: u64,
68    node_id: usize,
69    generation: u64,
70}
71
72impl Ord for ReadyEntry {
73    fn cmp(&self, other: &Self) -> Ordering {
74        self.priority
75            .cmp(&other.priority)
76            .then_with(|| other.enqueue_sequence.cmp(&self.enqueue_sequence))
77            .then_with(|| other.node_id.cmp(&self.node_id))
78    }
79}
80
81impl PartialOrd for ReadyEntry {
82    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
83        Some(self.cmp(other))
84    }
85}
86
87impl TaskGraph {
88    pub fn new() -> Self {
89        Self {
90            nodes: Vec::new(),
91            in_degree: Vec::new(),
92            reverse_adjacency: Vec::new(),
93            ready_heap: BinaryHeap::new(),
94            ready_generation: Vec::new(),
95            enqueued_round: Vec::new(),
96            enqueue_sequence: 0,
97            ready_round: 0,
98            scheduling: Vec::new(),
99            scheduler_policy: SchedulerPolicyConfig::default(),
100        }
101    }
102
103    /// Add a task, returns its ID. Duplicate dependency entries are collapsed: `in_degree` counts
104    /// entries but [`complete`](Self::complete) decrements once per completed dependency, so a
105    /// duplicated entry would leave the node permanently below its own in-degree (a silent stall).
106    pub fn add(&mut self, task: RuntimeTask, mut dependencies: Vec<usize>) -> usize {
107        let mut seen = std::collections::HashSet::new();
108        dependencies.retain(|d| seen.insert(*d));
109        let id = self.nodes.len();
110        let deg = dependencies.len();
111        let max_index = dependencies.iter().copied().max().unwrap_or(id).max(id);
112        self.reverse_adjacency.resize_with(max_index + 1, Vec::new);
113        for &dependency in &dependencies {
114            self.reverse_adjacency[dependency].push(id);
115        }
116        self.nodes.push(TaskNode {
117            id,
118            task,
119            status: if deg == 0 {
120                TaskStatus::Ready
121            } else {
122                TaskStatus::Pending
123            },
124            result: None,
125            dependencies,
126        });
127        self.in_degree.push(deg);
128        self.ready_generation.push(0);
129        self.enqueued_round.push(self.ready_round);
130        self.scheduling.push(SchedulingMetadata::default());
131        if deg == 0 {
132            self.enqueue_ready(id);
133        }
134        id
135    }
136
137    /// Topological sort — returns ordered IDs or error if cycle detected.
138    pub fn topological_sort(&self) -> Result<Vec<usize>> {
139        let n = self.nodes.len();
140        // `self.in_degree` is the live residual count and is mutated as tasks complete. A
141        // topological validation must always start from the immutable graph shape, otherwise
142        // validating a resumed/partially completed graph double-decrements edges and underflows.
143        let mut in_deg: Vec<usize> = self
144            .nodes
145            .iter()
146            .map(|node| node.dependencies.len())
147            .collect();
148
149        let mut queue: Vec<usize> = (0..n).filter(|&i| in_deg[i] == 0).collect();
150        let mut order = Vec::with_capacity(n);
151
152        while let Some(id) = queue.pop() {
153            order.push(id);
154            for &next in self.reverse_adjacency.get(id).into_iter().flatten() {
155                in_deg[next] -= 1;
156                if in_deg[next] == 0 {
157                    queue.push(next);
158                }
159            }
160        }
161
162        if order.len() != n {
163            return Err(DeepStrikeError::OrchestrationCycle);
164        }
165        Ok(order)
166    }
167
168    /// Return IDs of tasks that are Ready (deps satisfied, not yet started).
169    pub fn ready_tasks(&mut self) -> Vec<usize> {
170        // Drain the live heap so stale generations from loop re-arms are discarded instead of
171        // accumulating for the lifetime of a long workflow. Valid entries are reinserted because
172        // the caller may start only a concurrency-limited prefix of this ordered snapshot.
173        let mut valid_entries = Vec::new();
174        let mut ready = Vec::new();
175        while let Some(entry) = self.ready_heap.pop() {
176            if self.nodes.get(entry.node_id).map(|node| node.status) == Some(TaskStatus::Ready)
177                && self.ready_generation[entry.node_id] == entry.generation
178            {
179                ready.push(entry.node_id);
180                valid_entries.push(entry);
181            }
182        }
183        self.ready_heap.extend(valid_entries);
184        self.ready_round = self.ready_round.saturating_add(1);
185        ready
186    }
187
188    /// Mark a task as running.
189    pub fn start(&mut self, task_id: usize) {
190        if let Some(node) = self.nodes.get_mut(task_id) {
191            node.status = TaskStatus::Running;
192        }
193    }
194
195    /// Re-mark a (running) task as Ready without touching dependents — used to re-arm a loop node
196    /// for its next iteration. Unlike [`complete`](Self::complete), this does NOT decrement any
197    /// in-degree, so the loop node's dependents stay pending until the loop finally `complete`s.
198    pub fn set_ready(&mut self, task_id: usize) {
199        if let Some(node) = self.nodes.get_mut(task_id) {
200            if node.status != TaskStatus::Ready {
201                node.status = TaskStatus::Ready;
202                self.enqueue_ready(task_id);
203            }
204        }
205    }
206
207    /// Mark a task as completed; promote dependents whose in-degree reaches 0.
208    ///
209    /// Idempotent: a task already terminal (Completed/Failed) is left untouched — a duplicate
210    /// completion (at-least-once event delivery, resume replay) must not double-decrement its
211    /// dependents' in-degree, which would underflow (debug panic) or over-promote gated nodes.
212    pub fn complete(&mut self, task_id: usize, result: LoopResult) {
213        {
214            let Some(node) = self.nodes.get_mut(task_id) else {
215                return;
216            };
217            if node.status.is_terminal() {
218                return;
219            }
220            node.status = TaskStatus::Completed;
221            node.result = Some(result);
222        }
223        let dependents = self
224            .reverse_adjacency
225            .get(task_id)
226            .cloned()
227            .unwrap_or_default();
228        for dep_id in dependents {
229            self.in_degree[dep_id] -= 1;
230            if self.in_degree[dep_id] == 0 {
231                let should_enqueue =
232                    self.nodes.get(dep_id).map(|n| n.status) == Some(TaskStatus::Pending);
233                if should_enqueue {
234                    self.nodes[dep_id].status = TaskStatus::Ready;
235                    self.enqueue_ready(dep_id);
236                }
237            }
238        }
239    }
240
241    pub fn complete_partial(&mut self, task_id: usize, result: LoopResult) {
242        if let Some(node) = self.nodes.get_mut(task_id) {
243            if !node.status.is_terminal() {
244                node.status = TaskStatus::CompletedPartial;
245                node.result = Some(result);
246            }
247        }
248    }
249
250    /// Mark a task as failed (dependents remain Pending — caller decides policy). Terminal states
251    /// are sticky: failing an already-completed task must not un-complete it (idempotency twin of
252    /// [`complete`](Self::complete)).
253    pub fn fail(&mut self, task_id: usize) {
254        if let Some(node) = self.nodes.get_mut(task_id) {
255            if !node.status.is_terminal() {
256                node.status = TaskStatus::Failed;
257            }
258        }
259    }
260
261    pub fn fail_with_result(&mut self, task_id: usize, result: LoopResult) {
262        if let Some(node) = self.nodes.get_mut(task_id) {
263            if !node.status.is_terminal() {
264                node.status = TaskStatus::Failed;
265                node.result = Some(result);
266            }
267        }
268    }
269
270    pub fn skip_upstream_failed(&mut self, task_id: usize) {
271        if let Some(node) = self.nodes.get_mut(task_id) {
272            if !node.status.is_terminal() {
273                node.status = TaskStatus::SkippedUpstreamFailed;
274            }
275        }
276    }
277
278    pub fn get(&self, task_id: usize) -> Option<&TaskNode> {
279        self.nodes.get(task_id)
280    }
281
282    pub fn len(&self) -> usize {
283        self.nodes.len()
284    }
285
286    pub fn is_empty(&self) -> bool {
287        self.nodes.is_empty()
288    }
289
290    pub fn all_done(&self) -> bool {
291        self.nodes.iter().all(|n| n.status.is_terminal())
292    }
293
294    pub fn configure_scheduling(&mut self, policy: SchedulerPolicyConfig, token_costs: &[u64]) {
295        self.scheduler_policy = policy;
296        let order = self
297            .topological_sort()
298            .unwrap_or_else(|_| (0..self.nodes.len()).collect());
299        let mut reachable: Vec<HashSet<usize>> = vec![HashSet::new(); self.nodes.len()];
300        for &node in order.iter().rev() {
301            let mut critical = 1u64;
302            let children = self
303                .reverse_adjacency
304                .get(node)
305                .cloned()
306                .unwrap_or_default();
307            for child in children {
308                critical = critical.max(1 + self.scheduling[child].critical_path_remaining);
309                reachable[node].insert(child);
310                let descendants: Vec<usize> = reachable[child].iter().copied().collect();
311                reachable[node].extend(descendants);
312            }
313            self.scheduling[node] = SchedulingMetadata {
314                critical_path_remaining: critical,
315                downstream_fanout: reachable[node].len() as u64,
316                token_cost: token_costs.get(node).copied().unwrap_or(0),
317            };
318        }
319        self.rebuild_ready_heap();
320    }
321
322    fn rebuild_ready_heap(&mut self) {
323        self.ready_heap.clear();
324        for node_id in 0..self.nodes.len() {
325            if self.nodes[node_id].status == TaskStatus::Ready {
326                self.push_ready_entry(node_id);
327            }
328        }
329    }
330
331    fn enqueue_ready(&mut self, task_id: usize) {
332        self.ready_generation[task_id] = self.ready_generation[task_id].saturating_add(1);
333        self.enqueued_round[task_id] = self.ready_round;
334        self.enqueue_sequence = self.enqueue_sequence.saturating_add(1);
335        self.push_ready_entry(task_id);
336    }
337
338    fn push_ready_entry(&mut self, task_id: usize) {
339        let metadata = self.scheduling[task_id];
340        let policy = self.scheduler_policy;
341        let priority = i128::from(policy.critical_path_weight)
342            * i128::from(metadata.critical_path_remaining)
343            + i128::from(policy.fanout_weight) * i128::from(metadata.downstream_fanout)
344            - i128::from(policy.age_weight) * i128::from(self.enqueued_round[task_id])
345            - i128::from(policy.token_cost_weight) * i128::from(metadata.token_cost);
346        self.ready_heap.push(ReadyEntry {
347            priority,
348            enqueue_sequence: self.enqueue_sequence,
349            node_id: task_id,
350            generation: self.ready_generation[task_id],
351        });
352    }
353}
354
355impl Default for TaskGraph {
356    fn default() -> Self {
357        Self::new()
358    }
359}
360
361#[cfg(test)]
362mod tests {
363    use super::*;
364
365    #[test]
366    fn topological_sort_linear() {
367        let mut g = TaskGraph::new();
368        let a = g.add(RuntimeTask::new("A"), vec![]);
369        let b = g.add(RuntimeTask::new("B"), vec![a]);
370        let c = g.add(RuntimeTask::new("C"), vec![b]);
371
372        let order = g.topological_sort().unwrap();
373        assert_eq!(order, vec![0, 1, 2]);
374        let _ = (a, c);
375    }
376
377    #[test]
378    fn detects_cycle() {
379        let mut g = TaskGraph::new();
380        g.nodes.push(TaskNode {
381            id: 0,
382            task: RuntimeTask::new("A"),
383            status: TaskStatus::Pending,
384            result: None,
385            dependencies: vec![1],
386        });
387        g.nodes.push(TaskNode {
388            id: 1,
389            task: RuntimeTask::new("B"),
390            status: TaskStatus::Pending,
391            result: None,
392            dependencies: vec![0],
393        });
394        g.in_degree.push(1);
395        g.in_degree.push(1);
396
397        assert!(g.topological_sort().is_err());
398    }
399
400    #[test]
401    fn ready_tasks_respects_deps() {
402        let mut g = TaskGraph::new();
403        let a = g.add(RuntimeTask::new("A"), vec![]);
404        let _b = g.add(RuntimeTask::new("B"), vec![a]);
405
406        assert_eq!(g.ready_tasks(), vec![0]); // only A is Ready
407    }
408
409    #[test]
410    fn set_ready_rearms_without_promoting_dependents() {
411        let mut g = TaskGraph::new();
412        let a = g.add(RuntimeTask::new("A"), vec![]); // loop node
413        let b = g.add(RuntimeTask::new("B"), vec![a]); // dependent
414        g.start(a);
415        // Re-arm A for its next iteration: A is Ready again, but B stays Pending (no promotion).
416        g.set_ready(a);
417        assert_eq!(g.nodes[a].status, TaskStatus::Ready);
418        assert_eq!(g.nodes[b].status, TaskStatus::Pending);
419        assert_eq!(g.ready_tasks(), vec![a]);
420    }
421
422    #[test]
423    fn complete_promotes_dependent() {
424        use crate::types::result::{LoopResult, TerminationReason};
425        let mut g = TaskGraph::new();
426        let a = g.add(RuntimeTask::new("A"), vec![]);
427        let b = g.add(RuntimeTask::new("B"), vec![a]);
428
429        assert_eq!(g.nodes[b].status, TaskStatus::Pending);
430        g.complete(
431            a,
432            LoopResult {
433                termination: TerminationReason::Completed,
434                final_message: None,
435                turns_used: 1,
436                total_tokens_used: 0,
437                loop_continue: None,
438                classify_branch: None,
439                tournament_winner: None,
440                pace_decision: None,
441            },
442        );
443        assert_eq!(g.nodes[b].status, TaskStatus::Ready);
444    }
445
446    #[test]
447    fn duplicate_complete_is_idempotent() {
448        use crate::types::result::{LoopResult, TerminationReason};
449        let result = || LoopResult {
450            termination: TerminationReason::Completed,
451            final_message: None,
452            turns_used: 1,
453            total_tokens_used: 0,
454            loop_continue: None,
455            classify_branch: None,
456            tournament_winner: None,
457            pace_decision: None,
458        };
459        // b gates on BOTH a and c; a duplicate completion of `a` must not stand in for `c`.
460        let mut g = TaskGraph::new();
461        let a = g.add(RuntimeTask::new("A"), vec![]);
462        let c = g.add(RuntimeTask::new("C"), vec![]);
463        let b = g.add(RuntimeTask::new("B"), vec![a, c]);
464
465        g.complete(a, result());
466        g.complete(a, result()); // duplicate delivery — no double decrement, no panic
467        assert_eq!(g.nodes[b].status, TaskStatus::Pending);
468        g.complete(c, result());
469        assert_eq!(g.nodes[b].status, TaskStatus::Ready);
470        // Terminal states are sticky both ways.
471        g.fail(a);
472        assert_eq!(g.nodes[a].status, TaskStatus::Completed);
473    }
474
475    #[test]
476    fn critical_path_priority_beats_lower_node_id() {
477        let mut g = TaskGraph::new();
478        let wide = g.add(RuntimeTask::new("wide"), vec![]);
479        let chain = g.add(RuntimeTask::new("chain"), vec![]);
480        g.add(RuntimeTask::new("wide-child-a"), vec![wide]);
481        g.add(RuntimeTask::new("wide-child-b"), vec![wide]);
482        let chain_2 = g.add(RuntimeTask::new("chain-2"), vec![chain]);
483        let chain_3 = g.add(RuntimeTask::new("chain-3"), vec![chain_2]);
484        g.add(RuntimeTask::new("chain-4"), vec![chain_3]);
485
486        g.configure_scheduling(SchedulerPolicyConfig::default(), &[]);
487
488        assert_eq!(g.ready_tasks(), vec![chain, wide]);
489    }
490
491    #[test]
492    fn zero_weights_use_fifo_and_loop_rearm_yields() {
493        let mut g = TaskGraph::new();
494        let loop_node = g.add(RuntimeTask::new("loop"), vec![]);
495        let peer = g.add(RuntimeTask::new("peer"), vec![]);
496        let policy = SchedulerPolicyConfig {
497            critical_path_weight: 0,
498            fanout_weight: 0,
499            age_weight: 0,
500            token_cost_weight: 0,
501            ..SchedulerPolicyConfig::default()
502        };
503        g.configure_scheduling(policy, &[]);
504        assert_eq!(g.ready_tasks(), vec![loop_node, peer]);
505
506        g.start(loop_node);
507        g.set_ready(loop_node);
508        assert_eq!(g.ready_tasks(), vec![peer, loop_node]);
509        assert_eq!(
510            g.ready_heap.len(),
511            2,
512            "stale loop generations must be collected"
513        );
514    }
515
516    #[test]
517    fn reverse_adjacency_tracks_only_outgoing_dependents() {
518        let mut g = TaskGraph::new();
519        let root = g.add(RuntimeTask::new("root"), vec![]);
520        let unrelated = g.add(RuntimeTask::new("unrelated"), vec![]);
521        let child = g.add(RuntimeTask::new("child"), vec![root]);
522        g.add(RuntimeTask::new("grandchild"), vec![child]);
523
524        assert_eq!(g.reverse_adjacency[root], vec![child]);
525        assert!(g.reverse_adjacency[unrelated].is_empty());
526    }
527}
deepstrike_core/orchestration/task_graph.rs

deepstrike_core/orchestration/
task_graph.rs
