Skip to main content

atman_runtime/
workflow.rs

1use chrono::{DateTime, Utc};
2use serde::{Deserialize, Serialize};
3
4use crate::event::{Event, FlowNodeStatus, FlowStatus, TurnId};
5
6#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
7pub struct WorkflowGraph {
8    pub turn_id: TurnId,
9    pub root: Vec<WorkflowNode>,
10}
11
12#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
13pub struct WorkflowNode {
14    pub id: String,
15    pub kind: WorkflowNodeKind,
16    pub label: String,
17    pub status: NodeStatus,
18    pub started_at: Option<DateTime<Utc>>,
19    pub ended_at: Option<DateTime<Utc>>,
20    pub output_preview: Option<String>,
21    pub children: Vec<WorkflowNode>,
22    pub parallelism: Parallelism,
23    #[serde(default, skip_serializing_if = "Option::is_none")]
24    pub approval: Option<ApprovalState>,
25    #[serde(default, skip_serializing_if = "Option::is_none")]
26    pub llm_stats: Option<LlmStats>,
27}
28
29#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
30pub struct LlmStats {
31    pub model: String,
32    pub input_tokens: u64,
33    pub output_tokens: u64,
34    pub cache_read: u64,
35    pub cache_write: u64,
36    pub ttft_ms: u64,
37    pub tokens_per_second: f64,
38    pub wallclock_ms: u64,
39}
40
41#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
42#[serde(tag = "type", rename_all = "snake_case")]
43pub enum WorkflowNodeKind {
44    Flow {
45        run_id: String,
46        flow_name: String,
47    },
48    Stmt {
49        node_kind: crate::nodegraph::NodeKind,
50    },
51    ToolCall {
52        tool_use_id: String,
53        tool: String,
54        args_preview: String,
55        result_preview: Option<String>,
56    },
57    Subflow {
58        run_id: String,
59        flow_name: String,
60    },
61    FanoutBranch {
62        branch_index: usize,
63    },
64}
65
66#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
67#[serde(rename_all = "snake_case")]
68pub enum NodeStatus {
69    Pending,
70    Running,
71    Ok,
72    Err,
73    Cancelled,
74}
75
76#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
77#[serde(rename_all = "snake_case", tag = "kind")]
78pub enum ApprovalState {
79    Pending {
80        level: String,
81        #[serde(default, skip_serializing_if = "Option::is_none")]
82        preview: Option<String>,
83    },
84    Approved,
85    Denied {
86        reason: String,
87    },
88}
89
90#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
91#[serde(rename_all = "snake_case")]
92pub enum Parallelism {
93    Serial,
94    Parallel,
95}
96
97impl WorkflowGraph {
98    pub fn new(turn_id: TurnId) -> Self {
99        Self {
100            turn_id,
101            root: Vec::new(),
102        }
103    }
104
105    pub fn apply_event(&mut self, event: &Event) {
106        match event {
107            Event::FlowStart {
108                run_id,
109                flow_name,
110                parent_run_id,
111                parent_node_id,
112                ts,
113                ..
114            } => {
115                let run_id_str = run_id.0.to_string();
116                let node = WorkflowNode {
117                    id: run_id_str.clone(),
118                    kind: if parent_run_id.is_some() {
119                        WorkflowNodeKind::Subflow {
120                            run_id: run_id_str,
121                            flow_name: flow_name.clone(),
122                        }
123                    } else {
124                        WorkflowNodeKind::Flow {
125                            run_id: run_id_str,
126                            flow_name: flow_name.clone(),
127                        }
128                    },
129                    label: flow_name.clone(),
130                    status: NodeStatus::Running,
131                    started_at: Some(*ts),
132                    ended_at: None,
133                    output_preview: None,
134                    children: Vec::new(),
135                    parallelism: Parallelism::Serial,
136                    approval: None,
137                    llm_stats: None,
138                };
139                match (parent_run_id.as_ref(), parent_node_id.as_deref()) {
140                    (Some(prid), Some(pid)) => {
141                        let scoped = scope_id(&prid.0.to_string(), pid);
142                        if let Some(parent) = find_node_mut(&mut self.root, &scoped) {
143                            parent.children.push(node);
144                        }
145                    }
146                    _ => self.root.push(node),
147                }
148            }
149            Event::FlowEnd {
150                run_id, status, ts, ..
151            } => {
152                let id = run_id.0.to_string();
153                if let Some(n) = find_node_mut(&mut self.root, &id) {
154                    let new_status = match status {
155                        FlowStatus::Ok => NodeStatus::Ok,
156                        FlowStatus::Errored { .. } => NodeStatus::Err,
157                        FlowStatus::Cancelled => NodeStatus::Cancelled,
158                    };
159                    n.status = new_status;
160                    n.ended_at = Some(*ts);
161                    for child in n.children.iter_mut() {
162                        if matches!(child.status, NodeStatus::Running | NodeStatus::Pending) {
163                            child.status = new_status;
164                            child.ended_at = Some(*ts);
165                        }
166                    }
167                }
168            }
169            Event::FlowNodeStart {
170                run_id,
171                node_id,
172                kind: nk,
173                label,
174                parent_node_id,
175                ts,
176                ..
177            } => {
178                let rid = run_id.0.to_string();
179                let scoped_id = scope_id(&rid, node_id);
180                let parent_id = parent_node_id
181                    .as_deref()
182                    .map(|p| scope_id(&rid, p))
183                    .unwrap_or_else(|| rid.clone());
184                let kind = if let Some(idx) = parse_branch_index(node_id) {
185                    WorkflowNodeKind::FanoutBranch { branch_index: idx }
186                } else {
187                    WorkflowNodeKind::Stmt {
188                        node_kind: nk.clone(),
189                    }
190                };
191                let node = WorkflowNode {
192                    id: scoped_id,
193                    kind,
194                    label: label.clone(),
195                    status: NodeStatus::Running,
196                    started_at: Some(*ts),
197                    ended_at: None,
198                    output_preview: None,
199                    children: Vec::new(),
200                    parallelism: Parallelism::Serial,
201                    approval: None,
202                    llm_stats: None,
203                };
204                if let Some(parent) = find_node_mut(&mut self.root, &parent_id) {
205                    if matches!(node.kind, WorkflowNodeKind::FanoutBranch { .. }) {
206                        parent.parallelism = Parallelism::Parallel;
207                    }
208                    parent.children.push(node);
209                }
210            }
211            Event::FlowNodeEnd {
212                run_id,
213                node_id,
214                status,
215                output_preview,
216                ts,
217                ..
218            } => {
219                let scoped = scope_id(&run_id.0.to_string(), node_id);
220                if let Some(n) = find_node_mut(&mut self.root, &scoped) {
221                    let new_status = match status {
222                        FlowNodeStatus::Ok => NodeStatus::Ok,
223                        FlowNodeStatus::Err => NodeStatus::Err,
224                        FlowNodeStatus::Cancelled => NodeStatus::Cancelled,
225                    };
226                    n.status = new_status;
227                    n.ended_at = Some(*ts);
228                    if let Some(p) = output_preview {
229                        n.output_preview = Some(p.clone());
230                    }
231                    for child in n.children.iter_mut() {
232                        if matches!(child.status, NodeStatus::Running | NodeStatus::Pending) {
233                            child.status = new_status;
234                            child.ended_at = Some(*ts);
235                        }
236                    }
237                }
238            }
239            Event::LlmCall {
240                run_id,
241                node_id,
242                model,
243                usage,
244                wallclock_ms,
245                ttft_ms,
246                tokens_per_second,
247                ..
248            } => {
249                if let (Some(rid), Some(nid)) = (run_id, node_id) {
250                    let scoped = scope_id(&rid.0.to_string(), nid);
251                    if let Some(n) = find_node_mut(&mut self.root, &scoped) {
252                        n.llm_stats = Some(LlmStats {
253                            model: model.clone(),
254                            input_tokens: usage.input,
255                            output_tokens: usage.output,
256                            cache_read: usage.cached_input,
257                            cache_write: usage.cache_write,
258                            ttft_ms: ttft_ms.unwrap_or(0),
259                            tokens_per_second: tokens_per_second.unwrap_or(0.0),
260                            wallclock_ms: *wallclock_ms,
261                        });
262                    }
263                }
264            }
265            Event::ToolNode {
266                run_id,
267                parent_node_id,
268                tool_use_id,
269                tool_name,
270                args_preview,
271                ts,
272                ..
273            } => {
274                let rid = run_id.0.to_string();
275                let scoped_parent = scope_id(&rid, parent_node_id);
276                let id = tool_node_id(&rid, tool_use_id);
277                let node = WorkflowNode {
278                    id,
279                    kind: WorkflowNodeKind::ToolCall {
280                        tool_use_id: tool_use_id.clone(),
281                        tool: tool_name.clone(),
282                        args_preview: args_preview.clone(),
283                        result_preview: None,
284                    },
285                    label: tool_name.clone(),
286                    status: NodeStatus::Running,
287                    started_at: Some(*ts),
288                    ended_at: None,
289                    output_preview: None,
290                    children: Vec::new(),
291                    parallelism: Parallelism::Serial,
292                    approval: None,
293                    llm_stats: None,
294                };
295                if let Some(parent) = find_node_mut(&mut self.root, &scoped_parent) {
296                    parent.children.push(node);
297                }
298            }
299            Event::AssistantMsg {
300                flow_run_id,
301                message,
302                ts,
303                ..
304            } => {
305                let Some(flow_id) = flow_run_id.as_ref().map(|r| r.0.to_string()) else {
306                    return;
307                };
308                for part in &message.parts {
309                    if let crate::message::MessagePart::ToolUse { id, name, input } = part {
310                        let node_id = tool_node_id(&flow_id, id);
311                        if find_node(&self.root, &node_id).is_some() {
312                            continue;
313                        }
314                        let args_preview = serde_json::to_string(input).unwrap_or_default();
315                        let args_preview: String = args_preview.chars().take(200).collect();
316                        let node = WorkflowNode {
317                            id: node_id,
318                            kind: WorkflowNodeKind::ToolCall {
319                                tool_use_id: id.clone(),
320                                tool: name.clone(),
321                                args_preview: args_preview.clone(),
322                                result_preview: None,
323                            },
324                            label: name.clone(),
325                            status: NodeStatus::Running,
326                            started_at: Some(*ts),
327                            ended_at: None,
328                            output_preview: None,
329                            children: Vec::new(),
330                            parallelism: Parallelism::Serial,
331                            approval: None,
332                            llm_stats: None,
333                        };
334                        if let Some(parent) = find_node_mut(&mut self.root, &flow_id) {
335                            parent.children.push(node);
336                        }
337                    }
338                }
339            }
340            Event::ToolResultMsg {
341                flow_run_id,
342                message,
343                ts,
344                ..
345            } => {
346                let flow_id = flow_run_id.as_ref().map(|r| r.0.to_string());
347                for part in &message.parts {
348                    if let crate::message::MessagePart::ToolResult {
349                        tool_use_id,
350                        content,
351                        is_error,
352                    } = part
353                    {
354                        let scoped_hit = flow_id.as_deref().and_then(|rid| {
355                            let id = tool_node_id(rid, tool_use_id);
356                            find_node_mut(&mut self.root, &id).map(|_| id)
357                        });
358                        let node = match scoped_hit {
359                            Some(id) => find_node_mut(&mut self.root, &id),
360                            None => find_tool_node_by_tool_use_id(&mut self.root, tool_use_id),
361                        };
362                        if let Some(n) = node {
363                            n.status = if *is_error {
364                                NodeStatus::Err
365                            } else {
366                                NodeStatus::Ok
367                            };
368                            n.ended_at = Some(*ts);
369                            let preview: String = content.chars().take(300).collect();
370                            n.output_preview = Some(preview.clone());
371                            if let WorkflowNodeKind::ToolCall { result_preview, .. } = &mut n.kind {
372                                *result_preview = Some(preview);
373                            }
374                        }
375                    }
376                }
377            }
378            Event::ToolPendingApproval {
379                run_id,
380                tool_use_id,
381                level,
382                preview,
383                ..
384            } => {
385                let rid = run_id.0.to_string();
386                let id = tool_node_id(&rid, tool_use_id);
387                if let Some(n) = find_node_mut(&mut self.root, &id) {
388                    n.approval = Some(ApprovalState::Pending {
389                        level: level.clone(),
390                        preview: preview.clone(),
391                    });
392                }
393            }
394            Event::ToolApproved {
395                run_id,
396                tool_use_id,
397                ..
398            } => {
399                let rid = run_id.0.to_string();
400                let id = tool_node_id(&rid, tool_use_id);
401                if let Some(n) = find_node_mut(&mut self.root, &id) {
402                    n.approval = Some(ApprovalState::Approved);
403                }
404            }
405            Event::ToolDenied {
406                run_id,
407                tool_use_id,
408                reason,
409                ..
410            } => {
411                let rid = run_id.0.to_string();
412                let id = tool_node_id(&rid, tool_use_id);
413                if let Some(n) = find_node_mut(&mut self.root, &id) {
414                    n.approval = Some(ApprovalState::Denied {
415                        reason: reason.clone(),
416                    });
417                }
418            }
419            _ => {}
420        }
421    }
422
423    pub fn find_node(&self, id: &str) -> Option<&WorkflowNode> {
424        find_node(&self.root, id)
425    }
426
427    pub fn find_node_mut(&mut self, id: &str) -> Option<&mut WorkflowNode> {
428        find_node_mut(&mut self.root, id)
429    }
430
431    pub fn apply_stream_frame(&mut self, frame: &crate::stream::StreamFrame) {
432        self.apply_stream_frame_at(frame, None);
433    }
434
435    pub fn apply_stream_frame_at(
436        &mut self,
437        frame: &crate::stream::StreamFrame,
438        override_ts: Option<chrono::DateTime<chrono::Utc>>,
439    ) {
440        use crate::stream::StreamFrame;
441        let now = override_ts.unwrap_or_else(Utc::now);
442        match frame {
443            StreamFrame::FlowGraph { run_id, graph } => {
444                if self.find_node(run_id).is_none() {
445                    self.root.push(WorkflowNode {
446                        id: run_id.clone(),
447                        kind: WorkflowNodeKind::Flow {
448                            run_id: run_id.clone(),
449                            flow_name: graph.flow_name.clone(),
450                        },
451                        label: graph.flow_name.clone(),
452                        status: NodeStatus::Running,
453                        started_at: Some(now),
454                        ended_at: None,
455                        output_preview: None,
456                        children: Vec::new(),
457                        parallelism: Parallelism::Serial,
458                        approval: None,
459                        llm_stats: None,
460                    });
461                }
462            }
463            StreamFrame::FlowStart {
464                run_id,
465                flow_name,
466                parent_run_id,
467                parent_node_id,
468            } => {
469                if self.find_node(run_id).is_some() {
470                    return;
471                }
472                let kind = if parent_run_id.is_some() {
473                    WorkflowNodeKind::Subflow {
474                        run_id: run_id.clone(),
475                        flow_name: flow_name.clone(),
476                    }
477                } else {
478                    WorkflowNodeKind::Flow {
479                        run_id: run_id.clone(),
480                        flow_name: flow_name.clone(),
481                    }
482                };
483                let node = WorkflowNode {
484                    id: run_id.clone(),
485                    kind,
486                    label: flow_name.clone(),
487                    status: NodeStatus::Running,
488                    started_at: Some(now),
489                    ended_at: None,
490                    output_preview: None,
491                    children: Vec::new(),
492                    parallelism: Parallelism::Serial,
493                    approval: None,
494                    llm_stats: None,
495                };
496                match (parent_run_id.as_deref(), parent_node_id.as_deref()) {
497                    (Some(prid), Some(pid)) => {
498                        let scoped = scope_id(prid, pid);
499                        if let Some(parent) = find_node_mut(&mut self.root, &scoped) {
500                            parent.children.push(node);
501                        } else {
502                            self.root.push(node);
503                        }
504                    }
505                    _ => self.root.push(node),
506                }
507            }
508            StreamFrame::FlowNodeStart {
509                run_id,
510                node_id,
511                kind: nk,
512                label,
513                parent_node_id,
514            } => {
515                let scoped_id = scope_id(run_id, node_id);
516                let parent_id = parent_node_id
517                    .as_deref()
518                    .map(|p| scope_id(run_id, p))
519                    .unwrap_or_else(|| run_id.clone());
520                let kind = if let Some(idx) = parse_branch_index(node_id) {
521                    WorkflowNodeKind::FanoutBranch { branch_index: idx }
522                } else {
523                    WorkflowNodeKind::Stmt {
524                        node_kind: nk.clone(),
525                    }
526                };
527                let node = WorkflowNode {
528                    id: scoped_id,
529                    kind,
530                    label: label.clone(),
531                    status: NodeStatus::Running,
532                    started_at: Some(now),
533                    ended_at: None,
534                    output_preview: None,
535                    children: Vec::new(),
536                    parallelism: Parallelism::Serial,
537                    approval: None,
538                    llm_stats: None,
539                };
540                if let Some(parent) = find_node_mut(&mut self.root, &parent_id) {
541                    if matches!(node.kind, WorkflowNodeKind::FanoutBranch { .. }) {
542                        parent.parallelism = Parallelism::Parallel;
543                    }
544                    parent.children.push(node);
545                }
546            }
547            StreamFrame::FlowNodeEnd {
548                run_id,
549                node_id,
550                status,
551                output_preview,
552                ..
553            } => {
554                let scoped = scope_id(run_id, node_id);
555                if let Some(n) = find_node_mut(&mut self.root, &scoped) {
556                    let new_status = match status {
557                        FlowNodeStatus::Ok => NodeStatus::Ok,
558                        FlowNodeStatus::Err => NodeStatus::Err,
559                        FlowNodeStatus::Cancelled => NodeStatus::Cancelled,
560                    };
561                    n.status = new_status;
562                    n.ended_at = Some(now);
563                    if let Some(p) = output_preview {
564                        n.output_preview = Some(p.clone());
565                    }
566                    for child in n.children.iter_mut() {
567                        if matches!(child.status, NodeStatus::Running | NodeStatus::Pending) {
568                            child.status = new_status;
569                            child.ended_at = Some(now);
570                        }
571                    }
572                }
573            }
574            StreamFrame::LlmCallStats {
575                model,
576                input_tokens,
577                output_tokens,
578                cache_read,
579                cache_write,
580                ttft_ms,
581                tokens_per_second,
582                wallclock_ms,
583                run_id,
584                node_id,
585            } => {
586                if let (Some(rid), Some(nid)) = (run_id.as_deref(), node_id.as_deref()) {
587                    let scoped = scope_id(rid, nid);
588                    if let Some(n) = find_node_mut(&mut self.root, &scoped) {
589                        n.llm_stats = Some(LlmStats {
590                            model: model.clone(),
591                            input_tokens: *input_tokens,
592                            output_tokens: *output_tokens,
593                            cache_read: *cache_read,
594                            cache_write: *cache_write,
595                            ttft_ms: *ttft_ms,
596                            tokens_per_second: *tokens_per_second,
597                            wallclock_ms: *wallclock_ms,
598                        });
599                    }
600                }
601            }
602            StreamFrame::ToolNode {
603                run_id,
604                parent_node_id,
605                tool_use_id,
606                tool,
607                args_preview,
608                ..
609            } => {
610                let scoped_parent = scope_id(run_id, parent_node_id);
611                let node = WorkflowNode {
612                    id: tool_node_id(run_id, tool_use_id),
613                    kind: WorkflowNodeKind::ToolCall {
614                        tool_use_id: tool_use_id.clone(),
615                        tool: tool.clone(),
616                        args_preview: args_preview.clone(),
617                        result_preview: None,
618                    },
619                    label: tool.clone(),
620                    status: NodeStatus::Running,
621                    started_at: Some(now),
622                    ended_at: None,
623                    output_preview: None,
624                    children: Vec::new(),
625                    parallelism: Parallelism::Serial,
626                    approval: None,
627                    llm_stats: None,
628                };
629                if let Some(parent) = find_node_mut(&mut self.root, &scoped_parent) {
630                    parent.children.push(node);
631                }
632            }
633            StreamFrame::ToolUseDone {
634                id, ok, preview, ..
635            } => {
636                if let Some(n) = find_tool_node_by_tool_use_id(&mut self.root, id) {
637                    n.status = if *ok { NodeStatus::Ok } else { NodeStatus::Err };
638                    n.ended_at = Some(now);
639                    n.output_preview = Some(preview.clone());
640                }
641            }
642            StreamFrame::FlowDone {
643                run_id,
644                ok,
645                cancelled,
646                ..
647            } => {
648                if let Some(n) = find_node_mut(&mut self.root, run_id) {
649                    let status = if *cancelled {
650                        NodeStatus::Cancelled
651                    } else if *ok {
652                        NodeStatus::Ok
653                    } else {
654                        NodeStatus::Err
655                    };
656                    cascade_terminate(n, status, now);
657                }
658            }
659            StreamFrame::AssistantMsg {
660                flow_run_id,
661                message,
662            } => {
663                let Some(rid_str) = flow_run_id else { return };
664                let Ok(uuid) = uuid::Uuid::parse_str(rid_str) else {
665                    return;
666                };
667                self.apply_event(&Event::AssistantMsg {
668                    seq: 0,
669                    turn_id: crate::event::TurnId::now(),
670                    flow_run_id: Some(crate::event::FlowRunId(uuid)),
671                    message: message.clone(),
672                    ts: chrono::Utc::now(),
673                });
674            }
675            StreamFrame::ToolResultMsg { message, .. } => {
676                self.apply_event(&Event::ToolResultMsg {
677                    seq: 0,
678                    turn_id: crate::event::TurnId::now(),
679                    flow_run_id: None,
680                    message: message.clone(),
681                    ts: chrono::Utc::now(),
682                });
683            }
684            StreamFrame::ToolPendingApproval {
685                run_id,
686                tool_use_id,
687                level,
688                preview,
689                ..
690            } => {
691                let id = tool_node_id(run_id, tool_use_id);
692                if let Some(n) = find_node_mut(&mut self.root, &id) {
693                    n.approval = Some(ApprovalState::Pending {
694                        level: level.clone(),
695                        preview: preview.clone(),
696                    });
697                }
698            }
699            StreamFrame::ToolApproved {
700                run_id,
701                tool_use_id,
702                ..
703            } => {
704                let id = tool_node_id(run_id, tool_use_id);
705                if let Some(n) = find_node_mut(&mut self.root, &id) {
706                    n.approval = Some(ApprovalState::Approved);
707                }
708            }
709            StreamFrame::ToolDenied {
710                run_id,
711                tool_use_id,
712                reason,
713            } => {
714                let id = tool_node_id(run_id, tool_use_id);
715                if let Some(n) = find_node_mut(&mut self.root, &id) {
716                    n.approval = Some(ApprovalState::Denied {
717                        reason: reason.clone(),
718                    });
719                }
720            }
721            _ => {}
722        }
723    }
724}
725
726fn cascade_terminate(n: &mut WorkflowNode, status: NodeStatus, now: DateTime<Utc>) {
727    if matches!(n.status, NodeStatus::Running | NodeStatus::Pending) {
728        n.status = status;
729        n.ended_at = Some(now);
730    }
731    for child in n.children.iter_mut() {
732        cascade_terminate(child, status, now);
733    }
734}
735
736fn find_node<'a>(nodes: &'a [WorkflowNode], id: &str) -> Option<&'a WorkflowNode> {
737    for n in nodes {
738        if n.id == id {
739            return Some(n);
740        }
741        if let Some(hit) = find_node(&n.children, id) {
742            return Some(hit);
743        }
744    }
745    None
746}
747
748fn find_node_mut<'a>(nodes: &'a mut [WorkflowNode], id: &str) -> Option<&'a mut WorkflowNode> {
749    for n in nodes.iter_mut() {
750        if n.id == id {
751            return Some(n);
752        }
753        if let Some(hit) = find_node_mut(&mut n.children, id) {
754            return Some(hit);
755        }
756    }
757    None
758}
759
760fn scope_id(run_id: &str, node_id: &str) -> String {
761    format!("{run_id}::{node_id}")
762}
763
764fn tool_node_id(run_id: &str, tool_use_id: &str) -> String {
765    format!("tool:{run_id}:{tool_use_id}")
766}
767
768fn find_tool_node_by_tool_use_id<'a>(
769    nodes: &'a mut [WorkflowNode],
770    tool_use_id: &str,
771) -> Option<&'a mut WorkflowNode> {
772    for n in nodes.iter_mut() {
773        if let WorkflowNodeKind::ToolCall {
774            tool_use_id: tid, ..
775        } = &n.kind
776            && tid == tool_use_id
777        {
778            return Some(n);
779        }
780        if let Some(hit) = find_tool_node_by_tool_use_id(&mut n.children, tool_use_id) {
781            return Some(hit);
782        }
783    }
784    None
785}
786
787fn parse_branch_index(node_id: &str) -> Option<usize> {
788    let start = node_id.rfind(".branch[")?;
789    let rest = &node_id[start + ".branch[".len()..];
790    let end = rest.find(']')?;
791    rest[..end].parse().ok()
792}
793
794#[cfg(test)]
795mod tests {
796    use super::*;
797    use crate::event::{FlowRunId, FlowStatus};
798    use crate::nodegraph::NodeKind;
799
800    fn now() -> DateTime<Utc> {
801        Utc::now()
802    }
803
804    fn flow_start(run_id: FlowRunId, name: &str) -> Event {
805        Event::FlowStart {
806            seq: 0,
807            run_id,
808            flow_name: name.into(),
809            parent_run_id: None,
810            parent_node_id: None,
811            ts: now(),
812        }
813    }
814
815    fn subflow_start(child: FlowRunId, parent: FlowRunId, parent_node: &str, name: &str) -> Event {
816        Event::FlowStart {
817            seq: 0,
818            run_id: child,
819            flow_name: name.into(),
820            parent_run_id: Some(parent),
821            parent_node_id: Some(parent_node.into()),
822            ts: now(),
823        }
824    }
825
826    fn stmt_start(run_id: FlowRunId, node_id: &str, parent: Option<&str>) -> Event {
827        Event::FlowNodeStart {
828            seq: 0,
829            run_id,
830            node_id: node_id.into(),
831            kind: NodeKind::UserConfirm,
832            label: node_id.into(),
833            parent_node_id: parent.map(String::from),
834            ts: now(),
835        }
836    }
837
838    fn stmt_end(run_id: FlowRunId, node_id: &str, status: FlowNodeStatus) -> Event {
839        Event::FlowNodeEnd {
840            seq: 0,
841            run_id,
842            node_id: node_id.into(),
843            status,
844            output_preview: None,
845            ts: now(),
846        }
847    }
848
849    #[test]
850    fn top_level_flow_becomes_root_child() {
851        let mut g = WorkflowGraph::new(TurnId::now());
852        let rid = FlowRunId::now();
853        g.apply_event(&flow_start(rid.clone(), "main"));
854        assert_eq!(g.root.len(), 1);
855        let flow = &g.root[0];
856        assert!(matches!(flow.kind, WorkflowNodeKind::Flow { .. }));
857        assert_eq!(flow.status, NodeStatus::Running);
858        assert_eq!(flow.id, rid.0.to_string());
859    }
860
861    #[test]
862    fn subflow_attaches_under_parent_node() {
863        let mut g = WorkflowGraph::new(TurnId::now());
864        let parent_flow = FlowRunId::now();
865        let child_flow = FlowRunId::now();
866        g.apply_event(&flow_start(parent_flow.clone(), "outer"));
867        g.apply_event(&stmt_start(parent_flow.clone(), "stmt_0", None));
868        g.apply_event(&subflow_start(
869            child_flow.clone(),
870            parent_flow.clone(),
871            "stmt_0",
872            "inner",
873        ));
874        let scoped = scope_id(&parent_flow.0.to_string(), "stmt_0");
875        let stmt = g.find_node(&scoped).unwrap();
876        assert_eq!(stmt.children.len(), 1);
877        assert!(matches!(
878            stmt.children[0].kind,
879            WorkflowNodeKind::Subflow { .. }
880        ));
881        assert_eq!(stmt.children[0].id, child_flow.0.to_string());
882    }
883
884    #[test]
885    fn tool_node_attaches_and_flow_end_marks_status() {
886        let mut g = WorkflowGraph::new(TurnId::now());
887        let rid = FlowRunId::now();
888        g.apply_event(&flow_start(rid.clone(), "main"));
889        g.apply_event(&stmt_start(rid.clone(), "stmt_0", None));
890        g.apply_event(&Event::ToolNode {
891            seq: 0,
892            run_id: rid.clone(),
893            parent_node_id: "stmt_0".into(),
894            tool_use_id: "tu_1".into(),
895            tool_name: "fs.read".into(),
896            args_preview: "{\"path\":\"a\"}".into(),
897            ts: now(),
898        });
899        g.apply_event(&stmt_end(rid.clone(), "stmt_0", FlowNodeStatus::Ok));
900        g.apply_event(&Event::FlowEnd {
901            seq: 0,
902            run_id: rid.clone(),
903            flow_name: "main".into(),
904            status: FlowStatus::Ok,
905            ts: now(),
906        });
907        let scoped = scope_id(&rid.0.to_string(), "stmt_0");
908        let stmt = g.find_node(&scoped).unwrap();
909        assert_eq!(stmt.status, NodeStatus::Ok);
910        assert_eq!(stmt.children.len(), 1);
911        let tool = &stmt.children[0];
912        assert_eq!(tool.id, tool_node_id(&rid.0.to_string(), "tu_1"));
913        assert!(matches!(tool.kind, WorkflowNodeKind::ToolCall { .. }));
914        assert_eq!(g.root[0].status, NodeStatus::Ok);
915    }
916
917    #[test]
918    fn fanout_branch_marks_parent_parallel() {
919        let mut g = WorkflowGraph::new(TurnId::now());
920        let rid = FlowRunId::now();
921        g.apply_event(&flow_start(rid.clone(), "main"));
922        g.apply_event(&stmt_start(rid.clone(), "stmt_1", None));
923        g.apply_event(&stmt_start(rid.clone(), "stmt_1.branch[0]", Some("stmt_1")));
924        g.apply_event(&stmt_start(rid.clone(), "stmt_1.branch[1]", Some("stmt_1")));
925        let scoped = scope_id(&rid.0.to_string(), "stmt_1");
926        let parent = g.find_node(&scoped).unwrap();
927        assert_eq!(parent.parallelism, Parallelism::Parallel);
928        assert_eq!(parent.children.len(), 2);
929        assert!(matches!(
930            parent.children[0].kind,
931            WorkflowNodeKind::FanoutBranch { branch_index: 0 }
932        ));
933        assert!(matches!(
934            parent.children[1].kind,
935            WorkflowNodeKind::FanoutBranch { branch_index: 1 }
936        ));
937    }
938
939    #[test]
940    fn out_of_order_events_silently_dropped() {
941        let mut g = WorkflowGraph::new(TurnId::now());
942        g.apply_event(&stmt_start(FlowRunId::now(), "stmt_0", Some("missing")));
943        g.apply_event(&Event::ToolNode {
944            seq: 0,
945            run_id: FlowRunId::now(),
946            parent_node_id: "missing".into(),
947            tool_use_id: "tu".into(),
948            tool_name: "t".into(),
949            args_preview: "{}".into(),
950            ts: now(),
951        });
952        assert!(g.root.is_empty());
953    }
954}