Skip to main content

heartbit_core/agent/flow/
progress.rs

1//! [`ProgressTracker`] — fold [`WorkflowEvent`]s into a live [`RunProgress`]
2//! snapshot, the Rust analog of Claude Code's `/workflows` progress view.
3//!
4//! A tracker installs itself as the ctx's event sink
5//! ([`WorkflowCtxBuilder::on_event`](super::ctx::WorkflowCtxBuilder::on_event))
6//! and accumulates per-phase agent counts, lifecycle tallies, and token totals.
7//! [`snapshot`](ProgressTracker::snapshot) returns a cheap clone for rendering.
8
9use std::sync::{Arc, Mutex};
10
11use crate::llm::types::TokenUsage;
12
13use super::event::{OnWorkflowEvent, WorkflowEvent};
14
15/// Per-phase progress: how many agents were issued under a given phase title.
16#[derive(Debug, Clone, Default, PartialEq, Eq)]
17pub struct PhaseProgress {
18    /// The phase title (from [`WorkflowEvent::PhaseStarted`]).
19    pub title: String,
20    /// Agents issued under this phase so far.
21    pub agents_started: u64,
22}
23
24/// A point-in-time snapshot of a workflow run's progress.
25#[derive(Debug, Clone, Default)]
26pub struct RunProgress {
27    /// The most recently started phase, if any.
28    pub current_phase: Option<String>,
29    /// One entry per phase, in start order.
30    pub phases: Vec<PhaseProgress>,
31    /// Total agents issued (across all phases and the no-phase default).
32    pub agents_started: u64,
33    /// Agents that completed live.
34    pub agents_finished: u64,
35    /// Agents whose output was replayed from the resume journal.
36    pub agents_replayed: u64,
37    /// Agents skipped (e.g. cancelled mid-run).
38    pub agents_skipped: u64,
39    /// Agents that failed with an agent-domain error.
40    pub agents_failed: u64,
41    /// Count of `log()` lines emitted.
42    pub log_lines: u64,
43    /// Accumulated token usage across finished + replayed agents.
44    pub total_tokens: TokenUsage,
45}
46
47impl RunProgress {
48    /// Fold one [`WorkflowEvent`] into this snapshot.
49    ///
50    /// The match is intentionally exhaustive (no `_` arm): because this lives in
51    /// the same crate as [`WorkflowEvent`], adding a variant there makes this
52    /// fail to compile until a handling decision is made here.
53    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/// Accumulates [`WorkflowEvent`]s into a shared [`RunProgress`]. Cheap to clone
86/// (an `Arc` inside); install via [`callback`](Self::callback).
87#[derive(Clone, Default)]
88pub struct ProgressTracker(Arc<Mutex<RunProgress>>);
89
90impl ProgressTracker {
91    /// Create an empty tracker.
92    pub fn new() -> Self {
93        Self::default()
94    }
95
96    /// A point-in-time clone of the accumulated progress.
97    pub fn snapshot(&self) -> RunProgress {
98        self.0.lock().expect("progress lock poisoned").clone()
99    }
100
101    /// An event sink that folds each [`WorkflowEvent`] into this tracker. Pass to
102    /// [`WorkflowCtxBuilder::on_event`](super::ctx::WorkflowCtxBuilder::on_event).
103    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        // An agent with no phase bumps only the global counter.
155        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        // Tokens accumulate across both finished and replayed.
177        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}