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
38pub struct TaskGraph {
42 nodes: Vec<TaskNode>,
43 in_degree: Vec<usize>,
46 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 factors: SchedulingFactors,
63}
64
65#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
68#[serde(deny_unknown_fields)]
69pub struct SchedulingFactors {
70 pub deadline_urgency: u64,
71 pub process_priority: u64,
72 pub resource_pressure: u64,
73 pub budget_pressure: u64,
74}
75
76impl SchedulingFactors {
77 pub fn is_zero(&self) -> bool {
78 *self == Self::default()
79 }
80}
81
82#[derive(Debug, Clone, Copy, PartialEq, Eq)]
83struct ReadyEntry {
84 priority: i128,
85 enqueue_sequence: u64,
86 node_id: usize,
87 generation: u64,
88}
89
90impl Ord for ReadyEntry {
91 fn cmp(&self, other: &Self) -> Ordering {
92 self.priority
93 .cmp(&other.priority)
94 .then_with(|| other.enqueue_sequence.cmp(&self.enqueue_sequence))
95 .then_with(|| other.node_id.cmp(&self.node_id))
96 }
97}
98
99impl PartialOrd for ReadyEntry {
100 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
101 Some(self.cmp(other))
102 }
103}
104
105impl TaskGraph {
106 pub fn new() -> Self {
107 Self {
108 nodes: Vec::new(),
109 in_degree: Vec::new(),
110 reverse_adjacency: Vec::new(),
111 ready_heap: BinaryHeap::new(),
112 ready_generation: Vec::new(),
113 enqueued_round: Vec::new(),
114 enqueue_sequence: 0,
115 ready_round: 0,
116 scheduling: Vec::new(),
117 scheduler_policy: SchedulerPolicyConfig::default(),
118 }
119 }
120
121 pub fn add(&mut self, task: RuntimeTask, mut dependencies: Vec<usize>) -> usize {
125 let mut seen = std::collections::HashSet::new();
126 dependencies.retain(|d| seen.insert(*d));
127 let id = self.nodes.len();
128 let deg = dependencies.len();
129 let max_index = dependencies.iter().copied().max().unwrap_or(id).max(id);
130 self.reverse_adjacency.resize_with(max_index + 1, Vec::new);
131 for &dependency in &dependencies {
132 self.reverse_adjacency[dependency].push(id);
133 }
134 self.nodes.push(TaskNode {
135 id,
136 task,
137 status: if deg == 0 {
138 TaskStatus::Ready
139 } else {
140 TaskStatus::Pending
141 },
142 result: None,
143 dependencies,
144 });
145 self.in_degree.push(deg);
146 self.ready_generation.push(0);
147 self.enqueued_round.push(self.ready_round);
148 self.scheduling.push(SchedulingMetadata::default());
149 if deg == 0 {
150 self.enqueue_ready(id);
151 }
152 id
153 }
154
155 pub fn topological_sort(&self) -> Result<Vec<usize>> {
157 let n = self.nodes.len();
158 let mut in_deg: Vec<usize> = self
162 .nodes
163 .iter()
164 .map(|node| node.dependencies.len())
165 .collect();
166
167 let mut queue: Vec<usize> = (0..n).filter(|&i| in_deg[i] == 0).collect();
168 let mut order = Vec::with_capacity(n);
169
170 while let Some(id) = queue.pop() {
171 order.push(id);
172 for &next in self.reverse_adjacency.get(id).into_iter().flatten() {
173 in_deg[next] -= 1;
174 if in_deg[next] == 0 {
175 queue.push(next);
176 }
177 }
178 }
179
180 if order.len() != n {
181 return Err(DeepStrikeError::OrchestrationCycle);
182 }
183 Ok(order)
184 }
185
186 pub fn ready_tasks(&mut self) -> Vec<usize> {
188 let mut valid_entries = Vec::new();
192 let mut ready = Vec::new();
193 while let Some(entry) = self.ready_heap.pop() {
194 if self.nodes.get(entry.node_id).map(|node| node.status) == Some(TaskStatus::Ready)
195 && self.ready_generation[entry.node_id] == entry.generation
196 {
197 ready.push(entry.node_id);
198 valid_entries.push(entry);
199 }
200 }
201 self.ready_heap.extend(valid_entries);
202 self.ready_round = self.ready_round.saturating_add(1);
203 ready
204 }
205
206 pub fn start(&mut self, task_id: usize) {
208 if let Some(node) = self.nodes.get_mut(task_id) {
209 node.status = TaskStatus::Running;
210 }
211 }
212
213 pub fn set_ready(&mut self, task_id: usize) {
217 if let Some(node) = self.nodes.get_mut(task_id) {
218 if node.status != TaskStatus::Ready {
219 node.status = TaskStatus::Ready;
220 self.enqueue_ready(task_id);
221 }
222 }
223 }
224
225 pub fn complete(&mut self, task_id: usize, result: LoopResult) {
231 {
232 let Some(node) = self.nodes.get_mut(task_id) else {
233 return;
234 };
235 if node.status.is_terminal() {
236 return;
237 }
238 node.status = TaskStatus::Completed;
239 node.result = Some(result);
240 }
241 let dependents = self
242 .reverse_adjacency
243 .get(task_id)
244 .cloned()
245 .unwrap_or_default();
246 for dep_id in dependents {
247 self.in_degree[dep_id] -= 1;
248 if self.in_degree[dep_id] == 0 {
249 let should_enqueue =
250 self.nodes.get(dep_id).map(|n| n.status) == Some(TaskStatus::Pending);
251 if should_enqueue {
252 self.nodes[dep_id].status = TaskStatus::Ready;
253 self.enqueue_ready(dep_id);
254 }
255 }
256 }
257 }
258
259 pub fn complete_partial(&mut self, task_id: usize, result: LoopResult) {
260 if let Some(node) = self.nodes.get_mut(task_id) {
261 if !node.status.is_terminal() {
262 node.status = TaskStatus::CompletedPartial;
263 node.result = Some(result);
264 }
265 }
266 }
267
268 pub fn fail(&mut self, task_id: usize) {
272 if let Some(node) = self.nodes.get_mut(task_id) {
273 if !node.status.is_terminal() {
274 node.status = TaskStatus::Failed;
275 }
276 }
277 }
278
279 pub fn fail_with_result(&mut self, task_id: usize, result: LoopResult) {
280 if let Some(node) = self.nodes.get_mut(task_id) {
281 if !node.status.is_terminal() {
282 node.status = TaskStatus::Failed;
283 node.result = Some(result);
284 }
285 }
286 }
287
288 pub fn skip_upstream_failed(&mut self, task_id: usize) {
289 if let Some(node) = self.nodes.get_mut(task_id) {
290 if !node.status.is_terminal() {
291 node.status = TaskStatus::SkippedUpstreamFailed;
292 }
293 }
294 }
295
296 pub(crate) fn restore_runtime_state(
302 &mut self,
303 states: &[(TaskStatus, Option<LoopResult>)],
304 ) -> std::result::Result<(), String> {
305 if states.len() != self.nodes.len() {
306 return Err(format!(
307 "workflow checkpoint carries {} node states for a {} node graph",
308 states.len(),
309 self.nodes.len()
310 ));
311 }
312
313 for (node, (status, result)) in self.nodes.iter_mut().zip(states) {
314 node.status = *status;
315 node.result = result.clone();
316 }
317 self.in_degree = self
318 .nodes
319 .iter()
320 .map(|node| {
321 node.dependencies
322 .iter()
323 .filter(|&&dependency| {
324 self.nodes.get(dependency).map(|node| node.status)
325 != Some(TaskStatus::Completed)
326 })
327 .count()
328 })
329 .collect();
330 self.ready_heap.clear();
331 self.ready_generation.fill(0);
332 self.enqueued_round.fill(0);
333 self.enqueue_sequence = 0;
334 self.ready_round = 0;
335 for node in 0..self.nodes.len() {
336 if self.nodes[node].status == TaskStatus::Ready {
337 self.enqueue_ready(node);
338 }
339 }
340 Ok(())
341 }
342
343 pub fn get(&self, task_id: usize) -> Option<&TaskNode> {
344 self.nodes.get(task_id)
345 }
346
347 pub fn len(&self) -> usize {
348 self.nodes.len()
349 }
350
351 pub fn is_empty(&self) -> bool {
352 self.nodes.is_empty()
353 }
354
355 pub fn all_done(&self) -> bool {
356 self.nodes.iter().all(|n| n.status.is_terminal())
357 }
358
359 pub fn configure_scheduling(&mut self, policy: SchedulerPolicyConfig, token_costs: &[u64]) {
360 self.configure_scheduling_with_factors(policy, token_costs, &[]);
361 }
362
363 pub fn configure_scheduling_with_factors(
364 &mut self,
365 policy: SchedulerPolicyConfig,
366 token_costs: &[u64],
367 factors: &[SchedulingFactors],
368 ) {
369 self.scheduler_policy = policy;
370 let order = self
371 .topological_sort()
372 .unwrap_or_else(|_| (0..self.nodes.len()).collect());
373 let mut reachable: Vec<HashSet<usize>> = vec![HashSet::new(); self.nodes.len()];
374 for &node in order.iter().rev() {
375 let mut critical = 1u64;
376 let children = self
377 .reverse_adjacency
378 .get(node)
379 .cloned()
380 .unwrap_or_default();
381 for child in children {
382 critical = critical.max(1 + self.scheduling[child].critical_path_remaining);
383 reachable[node].insert(child);
384 let descendants: Vec<usize> = reachable[child].iter().copied().collect();
385 reachable[node].extend(descendants);
386 }
387 self.scheduling[node] = SchedulingMetadata {
388 critical_path_remaining: critical,
389 downstream_fanout: reachable[node].len() as u64,
390 token_cost: token_costs.get(node).copied().unwrap_or(0),
391 factors: factors.get(node).copied().unwrap_or_default(),
392 };
393 }
394 self.rebuild_ready_heap();
395 }
396
397 fn rebuild_ready_heap(&mut self) {
398 self.ready_heap.clear();
399 for node_id in 0..self.nodes.len() {
400 if self.nodes[node_id].status == TaskStatus::Ready {
401 self.push_ready_entry(node_id);
402 }
403 }
404 }
405
406 fn enqueue_ready(&mut self, task_id: usize) {
407 self.ready_generation[task_id] = self.ready_generation[task_id].saturating_add(1);
408 self.enqueued_round[task_id] = self.ready_round;
409 self.enqueue_sequence = self.enqueue_sequence.saturating_add(1);
410 self.push_ready_entry(task_id);
411 }
412
413 fn push_ready_entry(&mut self, task_id: usize) {
414 let metadata = self.scheduling[task_id];
415 let policy = self.scheduler_policy;
416 let priority = i128::from(policy.critical_path_weight)
417 * i128::from(metadata.critical_path_remaining)
418 + i128::from(policy.fanout_weight) * i128::from(metadata.downstream_fanout)
419 - i128::from(policy.age_weight) * i128::from(self.enqueued_round[task_id])
420 - i128::from(policy.token_cost_weight) * i128::from(metadata.token_cost);
421 let priority = priority
422 + i128::from(policy.deadline_weight) * i128::from(metadata.factors.deadline_urgency)
423 + i128::from(policy.process_priority_weight)
424 * i128::from(metadata.factors.process_priority)
425 + i128::from(policy.resource_pressure_weight)
426 * i128::from(metadata.factors.resource_pressure)
427 + i128::from(policy.budget_pressure_weight)
428 * i128::from(metadata.factors.budget_pressure);
429 self.ready_heap.push(ReadyEntry {
430 priority,
431 enqueue_sequence: self.enqueue_sequence,
432 node_id: task_id,
433 generation: self.ready_generation[task_id],
434 });
435 }
436}
437
438impl Default for TaskGraph {
439 fn default() -> Self {
440 Self::new()
441 }
442}
443
444#[cfg(test)]
445mod tests {
446 use super::*;
447
448 #[test]
449 fn topological_sort_linear() {
450 let mut g = TaskGraph::new();
451 let a = g.add(RuntimeTask::new("A"), vec![]);
452 let b = g.add(RuntimeTask::new("B"), vec![a]);
453 let c = g.add(RuntimeTask::new("C"), vec![b]);
454
455 let order = g.topological_sort().unwrap();
456 assert_eq!(order, vec![0, 1, 2]);
457 let _ = (a, c);
458 }
459
460 #[test]
461 fn detects_cycle() {
462 let mut g = TaskGraph::new();
463 g.nodes.push(TaskNode {
464 id: 0,
465 task: RuntimeTask::new("A"),
466 status: TaskStatus::Pending,
467 result: None,
468 dependencies: vec![1],
469 });
470 g.nodes.push(TaskNode {
471 id: 1,
472 task: RuntimeTask::new("B"),
473 status: TaskStatus::Pending,
474 result: None,
475 dependencies: vec![0],
476 });
477 g.in_degree.push(1);
478 g.in_degree.push(1);
479
480 assert!(g.topological_sort().is_err());
481 }
482
483 #[test]
484 fn ready_tasks_respects_deps() {
485 let mut g = TaskGraph::new();
486 let a = g.add(RuntimeTask::new("A"), vec![]);
487 let _b = g.add(RuntimeTask::new("B"), vec![a]);
488
489 assert_eq!(g.ready_tasks(), vec![0]); }
491
492 #[test]
493 fn set_ready_rearms_without_promoting_dependents() {
494 let mut g = TaskGraph::new();
495 let a = g.add(RuntimeTask::new("A"), vec![]); let b = g.add(RuntimeTask::new("B"), vec![a]); g.start(a);
498 g.set_ready(a);
500 assert_eq!(g.nodes[a].status, TaskStatus::Ready);
501 assert_eq!(g.nodes[b].status, TaskStatus::Pending);
502 assert_eq!(g.ready_tasks(), vec![a]);
503 }
504
505 #[test]
506 fn complete_promotes_dependent() {
507 use crate::types::result::{LoopResult, TerminationReason};
508 let mut g = TaskGraph::new();
509 let a = g.add(RuntimeTask::new("A"), vec![]);
510 let b = g.add(RuntimeTask::new("B"), vec![a]);
511
512 assert_eq!(g.nodes[b].status, TaskStatus::Pending);
513 g.complete(
514 a,
515 LoopResult {
516 termination: TerminationReason::Completed,
517 final_message: None,
518 turns_used: 1,
519 total_tokens_used: 0,
520 loop_continue: None,
521 classify_branch: None,
522 tournament_winner: None,
523 pace_decision: None,
524 },
525 );
526 assert_eq!(g.nodes[b].status, TaskStatus::Ready);
527 }
528
529 #[test]
530 fn duplicate_complete_is_idempotent() {
531 use crate::types::result::{LoopResult, TerminationReason};
532 let result = || LoopResult {
533 termination: TerminationReason::Completed,
534 final_message: None,
535 turns_used: 1,
536 total_tokens_used: 0,
537 loop_continue: None,
538 classify_branch: None,
539 tournament_winner: None,
540 pace_decision: None,
541 };
542 let mut g = TaskGraph::new();
544 let a = g.add(RuntimeTask::new("A"), vec![]);
545 let c = g.add(RuntimeTask::new("C"), vec![]);
546 let b = g.add(RuntimeTask::new("B"), vec![a, c]);
547
548 g.complete(a, result());
549 g.complete(a, result()); assert_eq!(g.nodes[b].status, TaskStatus::Pending);
551 g.complete(c, result());
552 assert_eq!(g.nodes[b].status, TaskStatus::Ready);
553 g.fail(a);
555 assert_eq!(g.nodes[a].status, TaskStatus::Completed);
556 }
557
558 #[test]
559 fn critical_path_priority_beats_lower_node_id() {
560 let mut g = TaskGraph::new();
561 let wide = g.add(RuntimeTask::new("wide"), vec![]);
562 let chain = g.add(RuntimeTask::new("chain"), vec![]);
563 g.add(RuntimeTask::new("wide-child-a"), vec![wide]);
564 g.add(RuntimeTask::new("wide-child-b"), vec![wide]);
565 let chain_2 = g.add(RuntimeTask::new("chain-2"), vec![chain]);
566 let chain_3 = g.add(RuntimeTask::new("chain-3"), vec![chain_2]);
567 g.add(RuntimeTask::new("chain-4"), vec![chain_3]);
568
569 g.configure_scheduling(SchedulerPolicyConfig::default(), &[]);
570
571 assert_eq!(g.ready_tasks(), vec![chain, wide]);
572 }
573
574 #[test]
575 fn zero_weights_use_fifo_and_loop_rearm_yields() {
576 let mut g = TaskGraph::new();
577 let loop_node = g.add(RuntimeTask::new("loop"), vec![]);
578 let peer = g.add(RuntimeTask::new("peer"), vec![]);
579 let policy = SchedulerPolicyConfig {
580 critical_path_weight: 0,
581 fanout_weight: 0,
582 age_weight: 0,
583 token_cost_weight: 0,
584 ..SchedulerPolicyConfig::default()
585 };
586 g.configure_scheduling(policy, &[]);
587 assert_eq!(g.ready_tasks(), vec![loop_node, peer]);
588
589 g.start(loop_node);
590 g.set_ready(loop_node);
591 assert_eq!(g.ready_tasks(), vec![peer, loop_node]);
592 assert_eq!(
593 g.ready_heap.len(),
594 2,
595 "stale loop generations must be collected"
596 );
597 }
598
599 #[test]
600 fn reverse_adjacency_tracks_only_outgoing_dependents() {
601 let mut g = TaskGraph::new();
602 let root = g.add(RuntimeTask::new("root"), vec![]);
603 let unrelated = g.add(RuntimeTask::new("unrelated"), vec![]);
604 let child = g.add(RuntimeTask::new("child"), vec![root]);
605 g.add(RuntimeTask::new("grandchild"), vec![child]);
606
607 assert_eq!(g.reverse_adjacency[root], vec![child]);
608 assert!(g.reverse_adjacency[unrelated].is_empty());
609 }
610
611 #[test]
612 fn scheduler_trace_uses_integer_deadline_priority_and_pressure_without_starving_a_peer() {
613 let mut g = TaskGraph::new();
614 let deadline = g.add(RuntimeTask::new("deadline"), vec![]);
615 let peer = g.add(RuntimeTask::new("peer"), vec![]);
616 let mut policy = SchedulerPolicyConfig::default();
617 policy.critical_path_weight = 0;
618 policy.fanout_weight = 0;
619 policy.age_weight = 1_000;
620 policy.token_cost_weight = 0;
621 policy.deadline_weight = 100;
622 policy.process_priority_weight = 10;
623 policy.resource_pressure_weight = 10;
624 policy.budget_pressure_weight = 10;
625 g.configure_scheduling_with_factors(
626 policy,
627 &[],
628 &[
629 SchedulingFactors {
630 deadline_urgency: 5,
631 process_priority: 1,
632 resource_pressure: 0,
633 budget_pressure: 0,
634 },
635 SchedulingFactors {
636 deadline_urgency: 0,
637 process_priority: 0,
638 resource_pressure: 0,
639 budget_pressure: 0,
640 },
641 ],
642 );
643 assert_eq!(g.ready_tasks(), vec![deadline, peer]);
644
645 g.start(deadline);
648 g.set_ready(deadline);
649 assert_eq!(g.ready_tasks(), vec![peer, deadline]);
650
651 let mut replay = TaskGraph::new();
652 replay.add(RuntimeTask::new("deadline"), vec![]);
653 replay.add(RuntimeTask::new("peer"), vec![]);
654 replay.configure_scheduling_with_factors(
655 policy,
656 &[],
657 &[
658 SchedulingFactors {
659 deadline_urgency: 5,
660 process_priority: 1,
661 resource_pressure: 0,
662 budget_pressure: 0,
663 },
664 SchedulingFactors {
665 deadline_urgency: 0,
666 process_priority: 0,
667 resource_pressure: 0,
668 budget_pressure: 0,
669 },
670 ],
671 );
672 assert_eq!(replay.ready_tasks(), vec![0, 1]);
673 }
674
675 #[test]
676 fn scheduler_factor_mutation_matrix_changes_only_its_configured_integer_axis() {
677 let mut g = TaskGraph::new();
678 let first = g.add(RuntimeTask::new("first"), vec![]);
679 let second = g.add(RuntimeTask::new("second"), vec![]);
680 let weights = |deadline, priority, resource, budget| SchedulerPolicyConfig {
681 critical_path_weight: 0,
682 fanout_weight: 0,
683 age_weight: 0,
684 token_cost_weight: 0,
685 deadline_weight: deadline,
686 process_priority_weight: priority,
687 resource_pressure_weight: resource,
688 budget_pressure_weight: budget,
689 ..SchedulerPolicyConfig::default()
690 };
691 let factors = [
692 SchedulingFactors {
693 deadline_urgency: 1,
694 process_priority: 1,
695 resource_pressure: 1,
696 budget_pressure: 1,
697 },
698 SchedulingFactors::default(),
699 ];
700 for policy in [
701 weights(1, 0, 0, 0),
702 weights(0, 1, 0, 0),
703 weights(0, 0, 1, 0),
704 weights(0, 0, 0, 1),
705 ] {
706 g.configure_scheduling_with_factors(policy, &[], &factors);
707 assert_eq!(g.ready_tasks(), vec![first, second]);
708 }
709 }
710
711 #[test]
712 fn restored_ready_state_replays_the_same_integer_factor_order() {
713 let policy = SchedulerPolicyConfig {
714 critical_path_weight: 0,
715 fanout_weight: 0,
716 age_weight: 0,
717 token_cost_weight: 0,
718 deadline_weight: 10,
719 process_priority_weight: 1,
720 resource_pressure_weight: 1,
721 budget_pressure_weight: 1,
722 ..SchedulerPolicyConfig::default()
723 };
724 let factors = [
725 SchedulingFactors {
726 deadline_urgency: 0,
727 process_priority: 1,
728 resource_pressure: 0,
729 budget_pressure: 0,
730 },
731 SchedulingFactors {
732 deadline_urgency: 1,
733 process_priority: 0,
734 resource_pressure: 0,
735 budget_pressure: 0,
736 },
737 ];
738 let mut original = TaskGraph::new();
739 original.add(RuntimeTask::new("priority"), vec![]);
740 original.add(RuntimeTask::new("deadline"), vec![]);
741 original.configure_scheduling_with_factors(policy, &[], &factors);
742 let expected = original.ready_tasks();
743
744 let states = [(TaskStatus::Ready, None), (TaskStatus::Ready, None)];
745 let mut restored = TaskGraph::new();
746 restored.add(RuntimeTask::new("priority"), vec![]);
747 restored.add(RuntimeTask::new("deadline"), vec![]);
748 restored.configure_scheduling_with_factors(policy, &[], &factors);
749 restored.restore_runtime_state(&states).unwrap();
750 restored.configure_scheduling_with_factors(policy, &[], &factors);
753 assert_eq!(restored.ready_tasks(), expected);
754 }
755}