1use std::cmp::Ordering;
2use std::collections::{BTreeMap, BTreeSet, HashMap};
3
4use chrono::{DateTime, Utc};
5
6use crate::context_plan::{ContextCallPurpose, ContextCallScope};
7use crate::workflow::{LlmStats, NodeStatus, WorkflowGraph, WorkflowNode, WorkflowNodeKind};
8
9const RECENT_COMPLETED_LEAF_LIMIT: usize = 256;
10
11#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
12pub struct WorkflowCounts {
13 pub nodes: usize,
14 pub agents: usize,
15 pub tools: usize,
16 pub edits: usize,
17}
18
19#[derive(Clone, Copy, Debug, PartialEq, Eq)]
20pub enum WorkflowAggregateStatus {
21 Empty,
22 Running,
23 Error,
24 Ok,
25}
26
27#[derive(Clone, Debug, PartialEq, Eq, Hash)]
28pub struct WorkflowLlmRoute {
29 pub provider: String,
30 pub model: String,
31 pub purpose: ContextCallPurpose,
32 pub scope: ContextCallScope,
33}
34
35impl WorkflowLlmRoute {
36 pub fn is_primary(&self) -> bool {
37 self.purpose == ContextCallPurpose::General && self.scope == ContextCallScope::Root
38 }
39}
40
41#[derive(Clone, Copy, Debug, Default, PartialEq)]
42pub struct WorkflowLlmAggregate {
43 pub calls: usize,
44 pub total_in: u64,
45 pub total_out: u64,
46 pub cache_read: u64,
47 pub cache_write: u64,
48 pub total_ttft_ms: u64,
49 pub speed_sum: f64,
50 pub speed_count: usize,
51}
52
53impl WorkflowLlmAggregate {
54 pub fn record(&mut self, stats: &LlmStats) {
55 self.calls = self.calls.saturating_add(1);
56 self.total_in = self
57 .total_in
58 .saturating_add(stats.input_tokens)
59 .saturating_add(stats.cache_read)
60 .saturating_add(stats.cache_write);
61 self.total_out = self.total_out.saturating_add(stats.output_tokens);
62 self.cache_read = self.cache_read.saturating_add(stats.cache_read);
63 self.cache_write = self.cache_write.saturating_add(stats.cache_write);
64 self.total_ttft_ms = self.total_ttft_ms.saturating_add(stats.ttft_ms);
65 if stats.tokens_per_second > 0.0 {
66 self.speed_sum += stats.tokens_per_second;
67 self.speed_count = self.speed_count.saturating_add(1);
68 }
69 }
70
71 pub fn merge(&mut self, other: Self) {
72 self.calls = self.calls.saturating_add(other.calls);
73 self.total_in = self.total_in.saturating_add(other.total_in);
74 self.total_out = self.total_out.saturating_add(other.total_out);
75 self.cache_read = self.cache_read.saturating_add(other.cache_read);
76 self.cache_write = self.cache_write.saturating_add(other.cache_write);
77 self.total_ttft_ms = self.total_ttft_ms.saturating_add(other.total_ttft_ms);
78 self.speed_sum += other.speed_sum;
79 self.speed_count = self.speed_count.saturating_add(other.speed_count);
80 }
81
82 pub fn average_speed(self) -> f64 {
83 if self.speed_count == 0 {
84 0.0
85 } else {
86 self.speed_sum / self.speed_count as f64
87 }
88 }
89
90 fn remove(&mut self, stats: &LlmStats) {
91 self.calls = self.calls.saturating_sub(1);
92 self.total_in = self
93 .total_in
94 .saturating_sub(stats.input_tokens)
95 .saturating_sub(stats.cache_read)
96 .saturating_sub(stats.cache_write);
97 self.total_out = self.total_out.saturating_sub(stats.output_tokens);
98 self.cache_read = self.cache_read.saturating_sub(stats.cache_read);
99 self.cache_write = self.cache_write.saturating_sub(stats.cache_write);
100 self.total_ttft_ms = self.total_ttft_ms.saturating_sub(stats.ttft_ms);
101 if stats.tokens_per_second > 0.0 {
102 self.speed_sum -= stats.tokens_per_second;
103 self.speed_count = self.speed_count.saturating_sub(1);
104 }
105 }
106}
107
108#[derive(Clone, Debug, PartialEq)]
109struct NodeSummaryState {
110 status: NodeStatus,
111 started_at: Option<DateTime<Utc>>,
112 ended_at: Option<DateTime<Utc>>,
113 llm_stats: Option<LlmStats>,
114}
115
116impl From<&WorkflowNode> for NodeSummaryState {
117 fn from(node: &WorkflowNode) -> Self {
118 Self {
119 status: node.status,
120 started_at: node.started_at,
121 ended_at: node.ended_at,
122 llm_stats: node.llm_stats.clone(),
123 }
124 }
125}
126
127#[derive(Clone, Debug, PartialEq, Eq)]
128struct LeafOrder {
129 started_at: DateTime<Utc>,
130 ordinal: u64,
131 node_id: String,
132 path: Vec<usize>,
133}
134
135impl Ord for LeafOrder {
136 fn cmp(&self, other: &Self) -> Ordering {
137 self.started_at
138 .cmp(&other.started_at)
139 .then_with(|| other.ordinal.cmp(&self.ordinal))
140 .then_with(|| self.node_id.cmp(&other.node_id))
141 }
142}
143
144impl PartialOrd for LeafOrder {
145 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
146 Some(self.cmp(other))
147 }
148}
149
150#[derive(Clone, Copy, Debug, PartialEq, Eq)]
151enum LeafBucket {
152 Running,
153 Completed,
154}
155
156#[derive(Clone, Debug)]
157struct LeafRecord {
158 order: LeafOrder,
159 bucket: LeafBucket,
160}
161
162#[derive(Clone, Debug, Default)]
163pub struct WorkflowSummary {
164 counts: WorkflowCounts,
165 running_nodes: usize,
166 error_nodes: usize,
167 root_started_at: BTreeMap<DateTime<Utc>, usize>,
168 root_ended_at: BTreeMap<DateTime<Utc>, usize>,
169 llm_routes: HashMap<WorkflowLlmRoute, WorkflowLlmAggregate>,
170 node_states: HashMap<String, NodeSummaryState>,
171 leaves: HashMap<String, LeafRecord>,
172 running_leaves: BTreeSet<LeafOrder>,
173 recent_completed_leaves: BTreeSet<LeafOrder>,
174 next_leaf_ordinal: u64,
175}
176
177impl WorkflowSummary {
178 pub(super) fn rebuild(graph: &WorkflowGraph) -> Self {
179 fn visit(summary: &mut WorkflowSummary, nodes: &[WorkflowNode], path: &mut Vec<usize>) {
180 for (index, node) in nodes.iter().enumerate() {
181 path.push(index);
182 summary.insert_node(node, path, path.len() == 1);
183 visit(summary, &node.children, path);
184 path.pop();
185 }
186 }
187
188 let mut summary = Self::default();
189 visit(&mut summary, &graph.root, &mut Vec::new());
190 summary
191 }
192
193 pub fn counts(&self) -> WorkflowCounts {
194 self.counts
195 }
196
197 pub fn status(&self) -> WorkflowAggregateStatus {
198 if self.running_nodes > 0 {
199 WorkflowAggregateStatus::Running
200 } else if self.error_nodes > 0 {
201 WorkflowAggregateStatus::Error
202 } else if self.counts.nodes == 0 {
203 WorkflowAggregateStatus::Empty
204 } else {
205 WorkflowAggregateStatus::Ok
206 }
207 }
208
209 pub fn started_at(&self) -> Option<DateTime<Utc>> {
210 self.root_started_at.first_key_value().map(|(at, _)| *at)
211 }
212
213 pub fn ended_at(&self) -> Option<DateTime<Utc>> {
214 self.root_ended_at.last_key_value().map(|(at, _)| *at)
215 }
216
217 pub fn elapsed_secs(&self, now: DateTime<Utc>) -> i64 {
218 let Some(started_at) = self.started_at() else {
219 return 0;
220 };
221 let ended_at = if self.status() == WorkflowAggregateStatus::Running {
222 now
223 } else {
224 self.ended_at().unwrap_or(started_at)
225 };
226 (ended_at - started_at).num_seconds().max(0)
227 }
228
229 pub fn llm_routes(&self) -> &HashMap<WorkflowLlmRoute, WorkflowLlmAggregate> {
230 &self.llm_routes
231 }
232
233 pub fn collapsed_leaf_paths(&self, limit: usize) -> Vec<Vec<usize>> {
234 let leaves = if self.running_leaves.is_empty() {
235 &self.recent_completed_leaves
236 } else {
237 &self.running_leaves
238 };
239 leaves
240 .iter()
241 .rev()
242 .take(limit)
243 .map(|leaf| leaf.path.clone())
244 .collect()
245 }
246
247 pub(super) fn insert_node(&mut self, node: &WorkflowNode, path: &[usize], is_root: bool) {
248 self.counts.nodes = self.counts.nodes.saturating_add(1);
249 if let WorkflowNodeKind::ToolCall { tool, .. } = &node.kind {
250 self.counts.tools = self.counts.tools.saturating_add(1);
251 if tool == "flow.spawn" {
252 self.counts.agents = self.counts.agents.saturating_add(1);
253 }
254 if matches!(
255 tool.as_str(),
256 "fs.edit" | "fs.write" | "hunk.apply" | "hunk.plan_edit"
257 ) {
258 self.counts.edits = self.counts.edits.saturating_add(1);
259 }
260 }
261 self.add_status(node.status);
262 if is_root {
263 insert_time(&mut self.root_started_at, node.started_at);
264 insert_time(&mut self.root_ended_at, node.ended_at);
265 }
266 if let Some(stats) = &node.llm_stats {
267 self.add_llm_stats(stats);
268 }
269 self.node_states
270 .insert(node.id.clone(), NodeSummaryState::from(node));
271 if node.children.is_empty() && is_render_leaf(node) {
272 let order = LeafOrder {
273 started_at: node.started_at.unwrap_or_else(Utc::now),
274 ordinal: self.next_leaf_ordinal,
275 node_id: node.id.clone(),
276 path: path.to_vec(),
277 };
278 self.next_leaf_ordinal = self.next_leaf_ordinal.wrapping_add(1);
279 self.insert_leaf(order, leaf_bucket(node.status));
280 }
281 }
282
283 pub(super) fn remove_leaf(&mut self, node_id: &str) {
284 let Some(record) = self.leaves.remove(node_id) else {
285 return;
286 };
287 self.remove_leaf_order(&record);
288 }
289
290 pub(super) fn sync_node(&mut self, node: &WorkflowNode, is_root: bool) -> bool {
291 let next = NodeSummaryState::from(node);
292 let Some(previous) = self.node_states.insert(node.id.clone(), next.clone()) else {
293 return false;
294 };
295 if previous == next {
296 return false;
297 }
298 if previous.status != next.status {
299 self.remove_status(previous.status);
300 self.add_status(next.status);
301 self.update_leaf_bucket(&node.id, next.status);
302 }
303 if is_root {
304 if previous.started_at != next.started_at {
305 remove_time(&mut self.root_started_at, previous.started_at);
306 insert_time(&mut self.root_started_at, next.started_at);
307 }
308 if previous.ended_at != next.ended_at {
309 remove_time(&mut self.root_ended_at, previous.ended_at);
310 insert_time(&mut self.root_ended_at, next.ended_at);
311 }
312 }
313 if previous.llm_stats != next.llm_stats {
314 if let Some(stats) = &previous.llm_stats {
315 self.remove_llm_stats(stats);
316 }
317 if let Some(stats) = &next.llm_stats {
318 self.add_llm_stats(stats);
319 }
320 }
321 true
322 }
323
324 pub(super) fn sync_subtree(&mut self, node: &WorkflowNode, is_root: bool) -> bool {
325 let mut changed = self.sync_node(node, is_root);
326 for child in &node.children {
327 changed |= self.sync_subtree(child, false);
328 }
329 changed
330 }
331
332 fn add_status(&mut self, status: NodeStatus) {
333 match status {
334 NodeStatus::Running | NodeStatus::Pending => {
335 self.running_nodes = self.running_nodes.saturating_add(1);
336 }
337 NodeStatus::Err => {
338 self.error_nodes = self.error_nodes.saturating_add(1);
339 }
340 NodeStatus::Ok | NodeStatus::Cancelled => {}
341 }
342 }
343
344 fn remove_status(&mut self, status: NodeStatus) {
345 match status {
346 NodeStatus::Running | NodeStatus::Pending => {
347 self.running_nodes = self.running_nodes.saturating_sub(1);
348 }
349 NodeStatus::Err => {
350 self.error_nodes = self.error_nodes.saturating_sub(1);
351 }
352 NodeStatus::Ok | NodeStatus::Cancelled => {}
353 }
354 }
355
356 fn add_llm_stats(&mut self, stats: &LlmStats) {
357 self.llm_routes
358 .entry(llm_route(stats))
359 .or_default()
360 .record(stats);
361 }
362
363 fn remove_llm_stats(&mut self, stats: &LlmStats) {
364 let route = llm_route(stats);
365 let remove_route = self.llm_routes.get_mut(&route).is_some_and(|aggregate| {
366 aggregate.remove(stats);
367 aggregate.calls == 0
368 });
369 if remove_route {
370 self.llm_routes.remove(&route);
371 }
372 }
373
374 fn insert_leaf(&mut self, order: LeafOrder, bucket: LeafBucket) {
375 self.leaves.insert(
376 order.node_id.clone(),
377 LeafRecord {
378 order: order.clone(),
379 bucket,
380 },
381 );
382 match bucket {
383 LeafBucket::Running => {
384 self.running_leaves.insert(order);
385 }
386 LeafBucket::Completed => {
387 self.recent_completed_leaves.insert(order);
388 while self.recent_completed_leaves.len() > RECENT_COMPLETED_LEAF_LIMIT {
389 if let Some(evicted) = self.recent_completed_leaves.pop_first() {
390 self.leaves.remove(&evicted.node_id);
391 }
392 }
393 }
394 }
395 }
396
397 fn remove_leaf_order(&mut self, record: &LeafRecord) {
398 match record.bucket {
399 LeafBucket::Running => {
400 self.running_leaves.remove(&record.order);
401 }
402 LeafBucket::Completed => {
403 self.recent_completed_leaves.remove(&record.order);
404 }
405 }
406 }
407
408 fn update_leaf_bucket(&mut self, node_id: &str, status: NodeStatus) {
409 let Some(mut record) = self.leaves.get(node_id).cloned() else {
410 return;
411 };
412 let next_bucket = leaf_bucket(status);
413 if record.bucket == next_bucket {
414 return;
415 }
416 self.remove_leaf_order(&record);
417 record.bucket = next_bucket;
418 self.leaves.insert(node_id.to_string(), record.clone());
419 match next_bucket {
420 LeafBucket::Running => {
421 self.running_leaves.insert(record.order);
422 }
423 LeafBucket::Completed => {
424 self.recent_completed_leaves.insert(record.order);
425 while self.recent_completed_leaves.len() > RECENT_COMPLETED_LEAF_LIMIT {
426 if let Some(evicted) = self.recent_completed_leaves.pop_first() {
427 self.leaves.remove(&evicted.node_id);
428 }
429 }
430 }
431 }
432 }
433}
434
435fn is_render_leaf(node: &WorkflowNode) -> bool {
436 matches!(
437 node.kind,
438 WorkflowNodeKind::ToolCall { .. }
439 | WorkflowNodeKind::Stmt { .. }
440 | WorkflowNodeKind::FanoutBranch { .. }
441 )
442}
443
444fn leaf_bucket(status: NodeStatus) -> LeafBucket {
445 if matches!(status, NodeStatus::Running | NodeStatus::Pending) {
446 LeafBucket::Running
447 } else {
448 LeafBucket::Completed
449 }
450}
451
452fn llm_route(stats: &LlmStats) -> WorkflowLlmRoute {
453 WorkflowLlmRoute {
454 provider: stats.provider.clone(),
455 model: stats.model.clone(),
456 purpose: stats.context_call_purpose,
457 scope: stats.context_call_scope,
458 }
459}
460
461fn insert_time(times: &mut BTreeMap<DateTime<Utc>, usize>, value: Option<DateTime<Utc>>) {
462 if let Some(value) = value {
463 *times.entry(value).or_default() += 1;
464 }
465}
466
467fn remove_time(times: &mut BTreeMap<DateTime<Utc>, usize>, value: Option<DateTime<Utc>>) {
468 let Some(value) = value else {
469 return;
470 };
471 if let Some(count) = times.get_mut(&value) {
472 *count = count.saturating_sub(1);
473 if *count == 0 {
474 times.remove(&value);
475 }
476 }
477}
478
479#[cfg(test)]
480mod tests {
481 use super::*;
482 use crate::event::TurnId;
483 use crate::workflow::{Parallelism, WorkflowNodeKind};
484
485 fn completed_tool(index: usize, started_at: DateTime<Utc>) -> WorkflowNode {
486 WorkflowNode {
487 id: format!("tool-{index}"),
488 kind: WorkflowNodeKind::ToolCall {
489 tool_use_id: format!("call-{index}"),
490 tool: "fs.read".into(),
491 args_preview: "{}".into(),
492 call_intent: None,
493 result_preview: Some("done".into()),
494 },
495 label: format!("tool-{index}"),
496 status: NodeStatus::Ok,
497 started_at: Some(started_at),
498 ended_at: Some(started_at),
499 output_preview: Some("done".into()),
500 children: Vec::new(),
501 parallelism: Parallelism::Serial,
502 approval: None,
503 llm_stats: None,
504 }
505 }
506
507 #[test]
508 fn recent_completed_leaf_selection_is_bounded_and_newest_first() {
509 let now = Utc::now();
510 let graph = WorkflowGraph {
511 turn_id: TurnId::now(),
512 root: (0..10_000)
513 .map(|index| {
514 completed_tool(index, now + chrono::Duration::milliseconds(index as i64))
515 })
516 .collect(),
517 permission_requests: Default::default(),
518 permission_groups: Default::default(),
519 resolved_permission_groups: Default::default(),
520 };
521
522 let summary = WorkflowSummary::rebuild(&graph);
523 let paths = summary.collapsed_leaf_paths(128);
524
525 assert_eq!(summary.counts().nodes, 10_000);
526 assert_eq!(summary.counts().tools, 10_000);
527 assert_eq!(summary.status(), WorkflowAggregateStatus::Ok);
528 assert_eq!(paths.len(), 128);
529 assert_eq!(paths.first(), Some(&vec![9_999]));
530 assert_eq!(paths.last(), Some(&vec![9_872]));
531 assert_eq!(
532 summary.recent_completed_leaves.len(),
533 RECENT_COMPLETED_LEAF_LIMIT
534 );
535 assert_eq!(summary.leaves.len(), RECENT_COMPLETED_LEAF_LIMIT);
536 }
537}