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 self.running_leaves
235 .iter()
236 .rev()
237 .chain(self.recent_completed_leaves.iter().rev())
238 .take(limit)
239 .map(|leaf| leaf.path.clone())
240 .collect()
241 }
242
243 pub(super) fn insert_node(&mut self, node: &WorkflowNode, path: &[usize], is_root: bool) {
244 self.counts.nodes = self.counts.nodes.saturating_add(1);
245 if let WorkflowNodeKind::ToolCall { tool, .. } = &node.kind {
246 self.counts.tools = self.counts.tools.saturating_add(1);
247 if tool == "flow.spawn" {
248 self.counts.agents = self.counts.agents.saturating_add(1);
249 }
250 if matches!(
251 tool.as_str(),
252 "fs.edit" | "fs.write" | "hunk.apply" | "hunk.plan_edit"
253 ) {
254 self.counts.edits = self.counts.edits.saturating_add(1);
255 }
256 }
257 self.add_status(node.status);
258 if is_root {
259 insert_time(&mut self.root_started_at, node.started_at);
260 insert_time(&mut self.root_ended_at, node.ended_at);
261 }
262 if let Some(stats) = &node.llm_stats {
263 self.add_llm_stats(stats);
264 }
265 self.node_states
266 .insert(node.id.clone(), NodeSummaryState::from(node));
267 if node.children.is_empty() && is_render_leaf(node) {
268 let order = LeafOrder {
269 started_at: node.started_at.unwrap_or_else(Utc::now),
270 ordinal: self.next_leaf_ordinal,
271 node_id: node.id.clone(),
272 path: path.to_vec(),
273 };
274 self.next_leaf_ordinal = self.next_leaf_ordinal.wrapping_add(1);
275 self.insert_leaf(order, leaf_bucket(node.status));
276 }
277 }
278
279 pub(super) fn remove_leaf(&mut self, node_id: &str) {
280 let Some(record) = self.leaves.remove(node_id) else {
281 return;
282 };
283 self.remove_leaf_order(&record);
284 }
285
286 pub(super) fn sync_node(&mut self, node: &WorkflowNode, is_root: bool) -> bool {
287 let next = NodeSummaryState::from(node);
288 let Some(previous) = self.node_states.insert(node.id.clone(), next.clone()) else {
289 return false;
290 };
291 if previous == next {
292 return false;
293 }
294 if previous.status != next.status {
295 self.remove_status(previous.status);
296 self.add_status(next.status);
297 self.update_leaf_bucket(&node.id, next.status);
298 }
299 if is_root {
300 if previous.started_at != next.started_at {
301 remove_time(&mut self.root_started_at, previous.started_at);
302 insert_time(&mut self.root_started_at, next.started_at);
303 }
304 if previous.ended_at != next.ended_at {
305 remove_time(&mut self.root_ended_at, previous.ended_at);
306 insert_time(&mut self.root_ended_at, next.ended_at);
307 }
308 }
309 if previous.llm_stats != next.llm_stats {
310 if let Some(stats) = &previous.llm_stats {
311 self.remove_llm_stats(stats);
312 }
313 if let Some(stats) = &next.llm_stats {
314 self.add_llm_stats(stats);
315 }
316 }
317 true
318 }
319
320 pub(super) fn sync_subtree(&mut self, node: &WorkflowNode, is_root: bool) -> bool {
321 let mut changed = self.sync_node(node, is_root);
322 for child in &node.children {
323 changed |= self.sync_subtree(child, false);
324 }
325 changed
326 }
327
328 fn add_status(&mut self, status: NodeStatus) {
329 match status {
330 NodeStatus::Running | NodeStatus::Pending => {
331 self.running_nodes = self.running_nodes.saturating_add(1);
332 }
333 NodeStatus::Err => {
334 self.error_nodes = self.error_nodes.saturating_add(1);
335 }
336 NodeStatus::Ok | NodeStatus::Cancelled => {}
337 }
338 }
339
340 fn remove_status(&mut self, status: NodeStatus) {
341 match status {
342 NodeStatus::Running | NodeStatus::Pending => {
343 self.running_nodes = self.running_nodes.saturating_sub(1);
344 }
345 NodeStatus::Err => {
346 self.error_nodes = self.error_nodes.saturating_sub(1);
347 }
348 NodeStatus::Ok | NodeStatus::Cancelled => {}
349 }
350 }
351
352 fn add_llm_stats(&mut self, stats: &LlmStats) {
353 self.llm_routes
354 .entry(llm_route(stats))
355 .or_default()
356 .record(stats);
357 }
358
359 fn remove_llm_stats(&mut self, stats: &LlmStats) {
360 let route = llm_route(stats);
361 let remove_route = self.llm_routes.get_mut(&route).is_some_and(|aggregate| {
362 aggregate.remove(stats);
363 aggregate.calls == 0
364 });
365 if remove_route {
366 self.llm_routes.remove(&route);
367 }
368 }
369
370 fn insert_leaf(&mut self, order: LeafOrder, bucket: LeafBucket) {
371 self.leaves.insert(
372 order.node_id.clone(),
373 LeafRecord {
374 order: order.clone(),
375 bucket,
376 },
377 );
378 match bucket {
379 LeafBucket::Running => {
380 self.running_leaves.insert(order);
381 }
382 LeafBucket::Completed => {
383 self.recent_completed_leaves.insert(order);
384 while self.recent_completed_leaves.len() > RECENT_COMPLETED_LEAF_LIMIT {
385 if let Some(evicted) = self.recent_completed_leaves.pop_first() {
386 self.leaves.remove(&evicted.node_id);
387 }
388 }
389 }
390 }
391 }
392
393 fn remove_leaf_order(&mut self, record: &LeafRecord) {
394 match record.bucket {
395 LeafBucket::Running => {
396 self.running_leaves.remove(&record.order);
397 }
398 LeafBucket::Completed => {
399 self.recent_completed_leaves.remove(&record.order);
400 }
401 }
402 }
403
404 fn update_leaf_bucket(&mut self, node_id: &str, status: NodeStatus) {
405 let Some(mut record) = self.leaves.get(node_id).cloned() else {
406 return;
407 };
408 let next_bucket = leaf_bucket(status);
409 if record.bucket == next_bucket {
410 return;
411 }
412 self.remove_leaf_order(&record);
413 record.bucket = next_bucket;
414 self.leaves.insert(node_id.to_string(), record.clone());
415 match next_bucket {
416 LeafBucket::Running => {
417 self.running_leaves.insert(record.order);
418 }
419 LeafBucket::Completed => {
420 self.recent_completed_leaves.insert(record.order);
421 while self.recent_completed_leaves.len() > RECENT_COMPLETED_LEAF_LIMIT {
422 if let Some(evicted) = self.recent_completed_leaves.pop_first() {
423 self.leaves.remove(&evicted.node_id);
424 }
425 }
426 }
427 }
428 }
429}
430
431fn is_render_leaf(node: &WorkflowNode) -> bool {
432 matches!(
433 node.kind,
434 WorkflowNodeKind::ToolCall { .. }
435 | WorkflowNodeKind::Stmt { .. }
436 | WorkflowNodeKind::FanoutBranch { .. }
437 )
438}
439
440fn leaf_bucket(status: NodeStatus) -> LeafBucket {
441 if matches!(status, NodeStatus::Running | NodeStatus::Pending) {
442 LeafBucket::Running
443 } else {
444 LeafBucket::Completed
445 }
446}
447
448fn llm_route(stats: &LlmStats) -> WorkflowLlmRoute {
449 WorkflowLlmRoute {
450 provider: stats.provider.clone(),
451 model: stats.model.clone(),
452 purpose: stats.context_call_purpose,
453 scope: stats.context_call_scope,
454 }
455}
456
457fn insert_time(times: &mut BTreeMap<DateTime<Utc>, usize>, value: Option<DateTime<Utc>>) {
458 if let Some(value) = value {
459 *times.entry(value).or_default() += 1;
460 }
461}
462
463fn remove_time(times: &mut BTreeMap<DateTime<Utc>, usize>, value: Option<DateTime<Utc>>) {
464 let Some(value) = value else {
465 return;
466 };
467 if let Some(count) = times.get_mut(&value) {
468 *count = count.saturating_sub(1);
469 if *count == 0 {
470 times.remove(&value);
471 }
472 }
473}
474
475#[cfg(test)]
476mod tests {
477 use super::*;
478 use crate::event::TurnId;
479 use crate::workflow::{Parallelism, WorkflowNodeKind};
480
481 fn completed_tool(index: usize, started_at: DateTime<Utc>) -> WorkflowNode {
482 WorkflowNode {
483 id: format!("tool-{index}"),
484 kind: WorkflowNodeKind::ToolCall {
485 tool_use_id: format!("call-{index}"),
486 tool: "fs.read".into(),
487 args_preview: "{}".into(),
488 call_intent: None,
489 result_preview: Some("done".into()),
490 },
491 label: format!("tool-{index}"),
492 status: NodeStatus::Ok,
493 started_at: Some(started_at),
494 ended_at: Some(started_at),
495 output_preview: Some("done".into()),
496 children: Vec::new(),
497 parallelism: Parallelism::Serial,
498 approval: None,
499 llm_stats: None,
500 }
501 }
502
503 #[test]
504 fn recent_completed_leaf_selection_is_bounded_and_newest_first() {
505 let now = Utc::now();
506 let graph = WorkflowGraph {
507 turn_id: TurnId::now(),
508 root: (0..10_000)
509 .map(|index| {
510 completed_tool(index, now + chrono::Duration::milliseconds(index as i64))
511 })
512 .collect(),
513 permission_requests: Default::default(),
514 permission_groups: Default::default(),
515 resolved_permission_groups: Default::default(),
516 };
517
518 let summary = WorkflowSummary::rebuild(&graph);
519 let paths = summary.collapsed_leaf_paths(128);
520
521 assert_eq!(summary.counts().nodes, 10_000);
522 assert_eq!(summary.counts().tools, 10_000);
523 assert_eq!(summary.status(), WorkflowAggregateStatus::Ok);
524 assert_eq!(paths.len(), 128);
525 assert_eq!(paths.first(), Some(&vec![9_999]));
526 assert_eq!(paths.last(), Some(&vec![9_872]));
527 assert_eq!(
528 summary.recent_completed_leaves.len(),
529 RECENT_COMPLETED_LEAF_LIMIT
530 );
531 assert_eq!(summary.leaves.len(), RECENT_COMPLETED_LEAF_LIMIT);
532 }
533
534 #[test]
535 fn running_leaves_are_filled_with_newest_completed_leaves() {
536 let now = Utc::now();
537 let mut root = (0..4)
538 .map(|index| completed_tool(index, now + chrono::Duration::milliseconds(index as i64)))
539 .collect::<Vec<_>>();
540 let mut running = completed_tool(4, now + chrono::Duration::milliseconds(4));
541 running.status = NodeStatus::Running;
542 running.ended_at = None;
543 running.output_preview = None;
544 root.push(running);
545 let graph = WorkflowGraph {
546 turn_id: TurnId::now(),
547 root,
548 permission_requests: Default::default(),
549 permission_groups: Default::default(),
550 resolved_permission_groups: Default::default(),
551 };
552
553 let summary = WorkflowSummary::rebuild(&graph);
554
555 assert_eq!(
556 summary.collapsed_leaf_paths(3),
557 vec![vec![4], vec![3], vec![2]]
558 );
559 }
560}