heartbit_core/agent/flow/
progress.rs1use std::sync::{Arc, Mutex};
10
11use crate::llm::types::TokenUsage;
12
13use super::event::{OnWorkflowEvent, WorkflowEvent};
14
15#[derive(Debug, Clone, Default, PartialEq, Eq)]
17pub struct PhaseProgress {
18 pub title: String,
20 pub agents_started: u64,
22}
23
24#[derive(Debug, Clone, Default)]
26pub struct RunProgress {
27 pub current_phase: Option<String>,
29 pub phases: Vec<PhaseProgress>,
31 pub agents_started: u64,
33 pub agents_finished: u64,
35 pub agents_replayed: u64,
37 pub agents_skipped: u64,
39 pub agents_failed: u64,
41 pub log_lines: u64,
43 pub total_tokens: TokenUsage,
45}
46
47impl RunProgress {
48 pub fn apply(&mut self, event: &WorkflowEvent) {
54 match event {
55 WorkflowEvent::PhaseStarted { title } => {
56 self.current_phase = Some(title.clone());
57 self.phases.push(PhaseProgress {
58 title: title.clone(),
59 agents_started: 0,
60 });
61 }
62 WorkflowEvent::AgentStarted { phase, .. } => {
63 self.agents_started += 1;
64 if let Some(phase) = phase
65 && let Some(p) = self.phases.iter_mut().rev().find(|p| &p.title == phase)
66 {
67 p.agents_started += 1;
68 }
69 }
70 WorkflowEvent::AgentFinished { usage, .. } => {
71 self.agents_finished += 1;
72 self.total_tokens += *usage;
73 }
74 WorkflowEvent::AgentReplayed { usage, .. } => {
75 self.agents_replayed += 1;
76 self.total_tokens += *usage;
77 }
78 WorkflowEvent::AgentSkipped { .. } => self.agents_skipped += 1,
79 WorkflowEvent::AgentFailed { .. } => self.agents_failed += 1,
80 WorkflowEvent::LogLine { .. } => self.log_lines += 1,
81 }
82 }
83}
84
85#[derive(Clone, Default)]
88pub struct ProgressTracker(Arc<Mutex<RunProgress>>);
89
90impl ProgressTracker {
91 pub fn new() -> Self {
93 Self::default()
94 }
95
96 pub fn snapshot(&self) -> RunProgress {
98 self.0.lock().expect("progress lock poisoned").clone()
99 }
100
101 pub fn callback(&self) -> Arc<OnWorkflowEvent> {
104 let inner = Arc::clone(&self.0);
105 Arc::new(move |event: WorkflowEvent| {
106 if let Ok(mut progress) = inner.lock() {
107 progress.apply(&event);
108 }
109 })
110 }
111}
112
113#[cfg(test)]
114mod tests {
115 use super::*;
116
117 fn usage(input: u32, output: u32) -> TokenUsage {
118 TokenUsage {
119 input_tokens: input,
120 output_tokens: output,
121 ..Default::default()
122 }
123 }
124
125 #[test]
126 fn phase_started_tracks_current_and_appends() {
127 let mut p = RunProgress::default();
128 p.apply(&WorkflowEvent::PhaseStarted {
129 title: "research".into(),
130 });
131 p.apply(&WorkflowEvent::PhaseStarted {
132 title: "verify".into(),
133 });
134 assert_eq!(p.current_phase.as_deref(), Some("verify"));
135 assert_eq!(p.phases.len(), 2);
136 assert_eq!(p.phases[0].title, "research");
137 assert_eq!(p.phases[1].title, "verify");
138 }
139
140 #[test]
141 fn agent_started_bumps_global_and_matching_phase() {
142 let mut p = RunProgress::default();
143 p.apply(&WorkflowEvent::PhaseStarted {
144 title: "research".into(),
145 });
146 p.apply(&WorkflowEvent::AgentStarted {
147 label: "a".into(),
148 phase: Some("research".into()),
149 });
150 p.apply(&WorkflowEvent::AgentStarted {
151 label: "b".into(),
152 phase: Some("research".into()),
153 });
154 p.apply(&WorkflowEvent::AgentStarted {
156 label: "c".into(),
157 phase: None,
158 });
159 assert_eq!(p.agents_started, 3);
160 assert_eq!(p.phases[0].agents_started, 2);
161 }
162
163 #[test]
164 fn finished_and_replayed_both_accumulate_tokens_but_count_separately() {
165 let mut p = RunProgress::default();
166 p.apply(&WorkflowEvent::AgentFinished {
167 label: "a".into(),
168 usage: usage(100, 50),
169 });
170 p.apply(&WorkflowEvent::AgentReplayed {
171 label: "b".into(),
172 usage: usage(10, 5),
173 });
174 assert_eq!(p.agents_finished, 1);
175 assert_eq!(p.agents_replayed, 1);
176 assert_eq!(p.total_tokens.input_tokens, 110);
178 assert_eq!(p.total_tokens.output_tokens, 55);
179 }
180
181 #[test]
182 fn skipped_failed_and_log_counters() {
183 let mut p = RunProgress::default();
184 p.apply(&WorkflowEvent::AgentSkipped { label: "a".into() });
185 p.apply(&WorkflowEvent::AgentFailed {
186 label: "b".into(),
187 error: "boom".into(),
188 });
189 p.apply(&WorkflowEvent::AgentFailed {
190 label: "c".into(),
191 error: "bang".into(),
192 });
193 p.apply(&WorkflowEvent::LogLine { msg: "tick".into() });
194 assert_eq!(p.agents_skipped, 1);
195 assert_eq!(p.agents_failed, 2);
196 assert_eq!(p.log_lines, 1);
197 }
198
199 #[test]
200 fn callback_folds_events_into_snapshot() {
201 let tracker = ProgressTracker::new();
202 let cb = tracker.callback();
203 cb(WorkflowEvent::PhaseStarted { title: "p1".into() });
204 cb(WorkflowEvent::AgentStarted {
205 label: "a".into(),
206 phase: Some("p1".into()),
207 });
208 cb(WorkflowEvent::AgentFinished {
209 label: "a".into(),
210 usage: usage(7, 3),
211 });
212 let snap = tracker.snapshot();
213 assert_eq!(snap.current_phase.as_deref(), Some("p1"));
214 assert_eq!(snap.agents_started, 1);
215 assert_eq!(snap.agents_finished, 1);
216 assert_eq!(snap.total_tokens.input_tokens, 7);
217 }
218}