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                ..
113            } => {
114                let run_id_str = run_id.0.to_string();
115                let node = WorkflowNode {
116                    id: run_id_str.clone(),
117                    kind: if parent_run_id.is_some() {
118                        WorkflowNodeKind::Subflow {
119                            run_id: run_id_str,
120                            flow_name: flow_name.clone(),
121                        }
122                    } else {
123                        WorkflowNodeKind::Flow {
124                            run_id: run_id_str,
125                            flow_name: flow_name.clone(),
126                        }
127                    },
128                    label: flow_name.clone(),
129                    status: NodeStatus::Running,
130                    started_at: Some(chrono::Utc::now()),
131                    ended_at: None,
132                    output_preview: None,
133                    children: Vec::new(),
134                    parallelism: Parallelism::Serial,
135                    approval: None,
136                    llm_stats: None,
137                };
138                match (parent_run_id.as_ref(), parent_node_id.as_deref()) {
139                    (Some(prid), Some(pid)) => {
140                        let scoped = scope_id(&prid.0.to_string(), pid);
141                        if let Some(parent) = find_node_mut(&mut self.root, &scoped) {
142                            parent.children.push(node);
143                        }
144                    }
145                    _ => self.root.push(node),
146                }
147            }
148            Event::FlowEnd { run_id, status, .. } => {
149                let id = run_id.0.to_string();
150                if let Some(n) = find_node_mut(&mut self.root, &id) {
151                    let new_status = match status {
152                        FlowStatus::Ok => NodeStatus::Ok,
153                        FlowStatus::Errored { .. } => NodeStatus::Err,
154                        FlowStatus::Cancelled => NodeStatus::Cancelled,
155                    };
156                    n.status = new_status;
157                    n.ended_at = Some(chrono::Utc::now());
158                    for child in n.children.iter_mut() {
159                        if matches!(child.status, NodeStatus::Running | NodeStatus::Pending) {
160                            child.status = new_status;
161                            child.ended_at = Some(chrono::Utc::now());
162                        }
163                    }
164                }
165            }
166            Event::FlowNodeStart {
167                run_id,
168                node_id,
169                kind: nk,
170                label,
171                parent_node_id,
172                ..
173            } => {
174                let rid = run_id.0.to_string();
175                let scoped_id = scope_id(&rid, node_id);
176                let parent_id = parent_node_id
177                    .as_deref()
178                    .map(|p| scope_id(&rid, p))
179                    .unwrap_or_else(|| rid.clone());
180                let kind = if let Some(idx) = parse_branch_index(node_id) {
181                    WorkflowNodeKind::FanoutBranch { branch_index: idx }
182                } else {
183                    WorkflowNodeKind::Stmt {
184                        node_kind: nk.clone(),
185                    }
186                };
187                let node = WorkflowNode {
188                    id: scoped_id,
189                    kind,
190                    label: label.clone(),
191                    status: NodeStatus::Running,
192                    started_at: Some(chrono::Utc::now()),
193                    ended_at: None,
194                    output_preview: None,
195                    children: Vec::new(),
196                    parallelism: Parallelism::Serial,
197                    approval: None,
198                    llm_stats: None,
199                };
200                if let Some(parent) = find_node_mut(&mut self.root, &parent_id) {
201                    if matches!(node.kind, WorkflowNodeKind::FanoutBranch { .. }) {
202                        parent.parallelism = Parallelism::Parallel;
203                    }
204                    parent.children.push(node);
205                }
206            }
207            Event::FlowNodeEnd {
208                run_id,
209                node_id,
210                status,
211                output_preview,
212                ..
213            } => {
214                let scoped = scope_id(&run_id.0.to_string(), node_id);
215                if let Some(n) = find_node_mut(&mut self.root, &scoped) {
216                    let new_status = match status {
217                        FlowNodeStatus::Ok => NodeStatus::Ok,
218                        FlowNodeStatus::Err => NodeStatus::Err,
219                        FlowNodeStatus::Cancelled => NodeStatus::Cancelled,
220                    };
221                    n.status = new_status;
222                    n.ended_at = Some(chrono::Utc::now());
223                    if let Some(p) = output_preview {
224                        n.output_preview = Some(p.clone());
225                    }
226                    for child in n.children.iter_mut() {
227                        if matches!(child.status, NodeStatus::Running | NodeStatus::Pending) {
228                            child.status = new_status;
229                            child.ended_at = Some(chrono::Utc::now());
230                        }
231                    }
232                }
233            }
234            Event::LlmCall {
235                run_id,
236                node_id,
237                model,
238                usage,
239                wallclock_ms,
240                ttft_ms,
241                tokens_per_second,
242                ..
243            } => {
244                if let (Some(rid), Some(nid)) = (run_id, node_id) {
245                    let scoped = scope_id(&rid.0.to_string(), nid);
246                    if let Some(n) = find_node_mut(&mut self.root, &scoped) {
247                        n.llm_stats = Some(LlmStats {
248                            model: model.clone(),
249                            input_tokens: usage.input,
250                            output_tokens: usage.output,
251                            cache_read: usage.cached_input,
252                            cache_write: usage.cache_write,
253                            ttft_ms: ttft_ms.unwrap_or(0),
254                            tokens_per_second: tokens_per_second.unwrap_or(0.0),
255                            wallclock_ms: *wallclock_ms,
256                        });
257                    }
258                }
259            }
260            Event::ToolNode {
261                run_id,
262                parent_node_id,
263                tool_use_id,
264                tool_name,
265                args_preview,
266                ..
267            } => {
268                let rid = run_id.0.to_string();
269                let scoped_parent = scope_id(&rid, parent_node_id);
270                let id = tool_node_id(&rid, tool_use_id);
271                let node = WorkflowNode {
272                    id,
273                    kind: WorkflowNodeKind::ToolCall {
274                        tool_use_id: tool_use_id.clone(),
275                        tool: tool_name.clone(),
276                        args_preview: args_preview.clone(),
277                        result_preview: None,
278                    },
279                    label: tool_name.clone(),
280                    status: NodeStatus::Running,
281                    started_at: Some(chrono::Utc::now()),
282                    ended_at: None,
283                    output_preview: None,
284                    children: Vec::new(),
285                    parallelism: Parallelism::Serial,
286                    approval: None,
287                    llm_stats: None,
288                };
289                if let Some(parent) = find_node_mut(&mut self.root, &scoped_parent) {
290                    parent.children.push(node);
291                }
292            }
293            Event::AssistantMsg {
294                flow_run_id,
295                message,
296                ..
297            } => {
298                let Some(flow_id) = flow_run_id.as_ref().map(|r| r.0.to_string()) else {
299                    return;
300                };
301                for part in &message.parts {
302                    if let crate::message::MessagePart::ToolUse { id, name, input } = part {
303                        let node_id = tool_node_id(&flow_id, id);
304                        if find_node(&self.root, &node_id).is_some() {
305                            continue;
306                        }
307                        let args_preview = serde_json::to_string(input).unwrap_or_default();
308                        let args_preview: String = args_preview.chars().take(200).collect();
309                        let node = WorkflowNode {
310                            id: node_id,
311                            kind: WorkflowNodeKind::ToolCall {
312                                tool_use_id: id.clone(),
313                                tool: name.clone(),
314                                args_preview: args_preview.clone(),
315                                result_preview: None,
316                            },
317                            label: name.clone(),
318                            status: NodeStatus::Running,
319                            started_at: Some(chrono::Utc::now()),
320                            ended_at: None,
321                            output_preview: None,
322                            children: Vec::new(),
323                            parallelism: Parallelism::Serial,
324                            approval: None,
325                            llm_stats: None,
326                        };
327                        if let Some(parent) = find_node_mut(&mut self.root, &flow_id) {
328                            parent.children.push(node);
329                        }
330                    }
331                }
332            }
333            Event::ToolResultMsg {
334                flow_run_id,
335                message,
336                ..
337            } => {
338                let flow_id = flow_run_id.as_ref().map(|r| r.0.to_string());
339                for part in &message.parts {
340                    if let crate::message::MessagePart::ToolResult {
341                        tool_use_id,
342                        content,
343                        is_error,
344                    } = part
345                    {
346                        let scoped_hit = flow_id.as_deref().and_then(|rid| {
347                            let id = tool_node_id(rid, tool_use_id);
348                            find_node_mut(&mut self.root, &id).map(|_| id)
349                        });
350                        let node = match scoped_hit {
351                            Some(id) => find_node_mut(&mut self.root, &id),
352                            None => find_tool_node_by_tool_use_id(&mut self.root, tool_use_id),
353                        };
354                        if let Some(n) = node {
355                            n.status = if *is_error {
356                                NodeStatus::Err
357                            } else {
358                                NodeStatus::Ok
359                            };
360                            n.ended_at = Some(chrono::Utc::now());
361                            let preview: String = content.chars().take(300).collect();
362                            n.output_preview = Some(preview.clone());
363                            if let WorkflowNodeKind::ToolCall { result_preview, .. } = &mut n.kind {
364                                *result_preview = Some(preview);
365                            }
366                        }
367                    }
368                }
369            }
370            Event::ToolPendingApproval {
371                run_id,
372                tool_use_id,
373                level,
374                preview,
375                ..
376            } => {
377                let rid = run_id.0.to_string();
378                let id = tool_node_id(&rid, tool_use_id);
379                if let Some(n) = find_node_mut(&mut self.root, &id) {
380                    n.approval = Some(ApprovalState::Pending {
381                        level: level.clone(),
382                        preview: preview.clone(),
383                    });
384                }
385            }
386            Event::ToolApproved {
387                run_id,
388                tool_use_id,
389                ..
390            } => {
391                let rid = run_id.0.to_string();
392                let id = tool_node_id(&rid, tool_use_id);
393                if let Some(n) = find_node_mut(&mut self.root, &id) {
394                    n.approval = Some(ApprovalState::Approved);
395                }
396            }
397            Event::ToolDenied {
398                run_id,
399                tool_use_id,
400                reason,
401                ..
402            } => {
403                let rid = run_id.0.to_string();
404                let id = tool_node_id(&rid, tool_use_id);
405                if let Some(n) = find_node_mut(&mut self.root, &id) {
406                    n.approval = Some(ApprovalState::Denied {
407                        reason: reason.clone(),
408                    });
409                }
410            }
411            _ => {}
412        }
413    }
414
415    pub fn find_node(&self, id: &str) -> Option<&WorkflowNode> {
416        find_node(&self.root, id)
417    }
418
419    pub fn find_node_mut(&mut self, id: &str) -> Option<&mut WorkflowNode> {
420        find_node_mut(&mut self.root, id)
421    }
422
423    pub fn apply_stream_frame(&mut self, frame: &crate::stream::StreamFrame) {
424        self.apply_stream_frame_at(frame, None);
425    }
426
427    pub fn apply_stream_frame_at(
428        &mut self,
429        frame: &crate::stream::StreamFrame,
430        override_ts: Option<chrono::DateTime<chrono::Utc>>,
431    ) {
432        use crate::stream::StreamFrame;
433        let now = override_ts.unwrap_or_else(Utc::now);
434        match frame {
435            StreamFrame::FlowGraph { run_id, graph } => {
436                if self.find_node(run_id).is_none() {
437                    self.root.push(WorkflowNode {
438                        id: run_id.clone(),
439                        kind: WorkflowNodeKind::Flow {
440                            run_id: run_id.clone(),
441                            flow_name: graph.flow_name.clone(),
442                        },
443                        label: graph.flow_name.clone(),
444                        status: NodeStatus::Running,
445                        started_at: Some(now),
446                        ended_at: None,
447                        output_preview: None,
448                        children: Vec::new(),
449                        parallelism: Parallelism::Serial,
450                        approval: None,
451                        llm_stats: None,
452                    });
453                }
454            }
455            StreamFrame::FlowStart {
456                run_id,
457                flow_name,
458                parent_run_id,
459                parent_node_id,
460            } => {
461                if self.find_node(run_id).is_some() {
462                    return;
463                }
464                let kind = if parent_run_id.is_some() {
465                    WorkflowNodeKind::Subflow {
466                        run_id: run_id.clone(),
467                        flow_name: flow_name.clone(),
468                    }
469                } else {
470                    WorkflowNodeKind::Flow {
471                        run_id: run_id.clone(),
472                        flow_name: flow_name.clone(),
473                    }
474                };
475                let node = WorkflowNode {
476                    id: run_id.clone(),
477                    kind,
478                    label: flow_name.clone(),
479                    status: NodeStatus::Running,
480                    started_at: Some(now),
481                    ended_at: None,
482                    output_preview: None,
483                    children: Vec::new(),
484                    parallelism: Parallelism::Serial,
485                    approval: None,
486                    llm_stats: None,
487                };
488                match (parent_run_id.as_deref(), parent_node_id.as_deref()) {
489                    (Some(prid), Some(pid)) => {
490                        let scoped = scope_id(prid, pid);
491                        if let Some(parent) = find_node_mut(&mut self.root, &scoped) {
492                            parent.children.push(node);
493                        } else {
494                            self.root.push(node);
495                        }
496                    }
497                    _ => self.root.push(node),
498                }
499            }
500            StreamFrame::FlowNodeStart {
501                run_id,
502                node_id,
503                kind: nk,
504                label,
505                parent_node_id,
506            } => {
507                let scoped_id = scope_id(run_id, node_id);
508                let parent_id = parent_node_id
509                    .as_deref()
510                    .map(|p| scope_id(run_id, p))
511                    .unwrap_or_else(|| run_id.clone());
512                let kind = if let Some(idx) = parse_branch_index(node_id) {
513                    WorkflowNodeKind::FanoutBranch { branch_index: idx }
514                } else {
515                    WorkflowNodeKind::Stmt {
516                        node_kind: nk.clone(),
517                    }
518                };
519                let node = WorkflowNode {
520                    id: scoped_id,
521                    kind,
522                    label: label.clone(),
523                    status: NodeStatus::Running,
524                    started_at: Some(now),
525                    ended_at: None,
526                    output_preview: None,
527                    children: Vec::new(),
528                    parallelism: Parallelism::Serial,
529                    approval: None,
530                    llm_stats: None,
531                };
532                if let Some(parent) = find_node_mut(&mut self.root, &parent_id) {
533                    if matches!(node.kind, WorkflowNodeKind::FanoutBranch { .. }) {
534                        parent.parallelism = Parallelism::Parallel;
535                    }
536                    parent.children.push(node);
537                }
538            }
539            StreamFrame::FlowNodeEnd {
540                run_id,
541                node_id,
542                status,
543                output_preview,
544                ..
545            } => {
546                let scoped = scope_id(run_id, node_id);
547                if let Some(n) = find_node_mut(&mut self.root, &scoped) {
548                    let new_status = match status {
549                        FlowNodeStatus::Ok => NodeStatus::Ok,
550                        FlowNodeStatus::Err => NodeStatus::Err,
551                        FlowNodeStatus::Cancelled => NodeStatus::Cancelled,
552                    };
553                    n.status = new_status;
554                    n.ended_at = Some(now);
555                    if let Some(p) = output_preview {
556                        n.output_preview = Some(p.clone());
557                    }
558                    for child in n.children.iter_mut() {
559                        if matches!(child.status, NodeStatus::Running | NodeStatus::Pending) {
560                            child.status = new_status;
561                            child.ended_at = Some(now);
562                        }
563                    }
564                }
565            }
566            StreamFrame::LlmCallStats {
567                model,
568                input_tokens,
569                output_tokens,
570                cache_read,
571                cache_write,
572                ttft_ms,
573                tokens_per_second,
574                wallclock_ms,
575                run_id,
576                node_id,
577            } => {
578                if let (Some(rid), Some(nid)) = (run_id.as_deref(), node_id.as_deref()) {
579                    let scoped = scope_id(rid, nid);
580                    if let Some(n) = find_node_mut(&mut self.root, &scoped) {
581                        n.llm_stats = Some(LlmStats {
582                            model: model.clone(),
583                            input_tokens: *input_tokens,
584                            output_tokens: *output_tokens,
585                            cache_read: *cache_read,
586                            cache_write: *cache_write,
587                            ttft_ms: *ttft_ms,
588                            tokens_per_second: *tokens_per_second,
589                            wallclock_ms: *wallclock_ms,
590                        });
591                    }
592                }
593            }
594            StreamFrame::ToolNode {
595                run_id,
596                parent_node_id,
597                tool_use_id,
598                tool,
599                args_preview,
600                ..
601            } => {
602                let scoped_parent = scope_id(run_id, parent_node_id);
603                let node = WorkflowNode {
604                    id: tool_node_id(run_id, tool_use_id),
605                    kind: WorkflowNodeKind::ToolCall {
606                        tool_use_id: tool_use_id.clone(),
607                        tool: tool.clone(),
608                        args_preview: args_preview.clone(),
609                        result_preview: None,
610                    },
611                    label: tool.clone(),
612                    status: NodeStatus::Running,
613                    started_at: Some(now),
614                    ended_at: None,
615                    output_preview: None,
616                    children: Vec::new(),
617                    parallelism: Parallelism::Serial,
618                    approval: None,
619                    llm_stats: None,
620                };
621                if let Some(parent) = find_node_mut(&mut self.root, &scoped_parent) {
622                    parent.children.push(node);
623                }
624            }
625            StreamFrame::ToolUseDone {
626                id, ok, preview, ..
627            } => {
628                if let Some(n) = find_tool_node_by_tool_use_id(&mut self.root, id) {
629                    n.status = if *ok { NodeStatus::Ok } else { NodeStatus::Err };
630                    n.ended_at = Some(now);
631                    n.output_preview = Some(preview.clone());
632                }
633            }
634            StreamFrame::FlowDone {
635                run_id,
636                ok,
637                cancelled,
638                ..
639            } => {
640                if let Some(n) = find_node_mut(&mut self.root, run_id) {
641                    let status = if *cancelled {
642                        NodeStatus::Cancelled
643                    } else if *ok {
644                        NodeStatus::Ok
645                    } else {
646                        NodeStatus::Err
647                    };
648                    cascade_terminate(n, status, now);
649                }
650            }
651            StreamFrame::AssistantMsg {
652                flow_run_id,
653                message,
654            } => {
655                let Some(rid_str) = flow_run_id else { return };
656                let Ok(uuid) = uuid::Uuid::parse_str(rid_str) else {
657                    return;
658                };
659                self.apply_event(&Event::AssistantMsg {
660                    turn_id: crate::event::TurnId::now(),
661                    flow_run_id: Some(crate::event::FlowRunId(uuid)),
662                    message: message.clone(),
663                });
664            }
665            StreamFrame::ToolResultMsg { message, .. } => {
666                self.apply_event(&Event::ToolResultMsg {
667                    turn_id: crate::event::TurnId::now(),
668                    flow_run_id: None,
669                    message: message.clone(),
670                });
671            }
672            StreamFrame::ToolPendingApproval {
673                run_id,
674                tool_use_id,
675                level,
676                preview,
677                ..
678            } => {
679                let id = tool_node_id(run_id, tool_use_id);
680                if let Some(n) = find_node_mut(&mut self.root, &id) {
681                    n.approval = Some(ApprovalState::Pending {
682                        level: level.clone(),
683                        preview: preview.clone(),
684                    });
685                }
686            }
687            StreamFrame::ToolApproved {
688                run_id,
689                tool_use_id,
690                ..
691            } => {
692                let id = tool_node_id(run_id, tool_use_id);
693                if let Some(n) = find_node_mut(&mut self.root, &id) {
694                    n.approval = Some(ApprovalState::Approved);
695                }
696            }
697            StreamFrame::ToolDenied {
698                run_id,
699                tool_use_id,
700                reason,
701            } => {
702                let id = tool_node_id(run_id, tool_use_id);
703                if let Some(n) = find_node_mut(&mut self.root, &id) {
704                    n.approval = Some(ApprovalState::Denied {
705                        reason: reason.clone(),
706                    });
707                }
708            }
709            _ => {}
710        }
711    }
712}
713
714fn cascade_terminate(n: &mut WorkflowNode, status: NodeStatus, now: DateTime<Utc>) {
715    if matches!(n.status, NodeStatus::Running | NodeStatus::Pending) {
716        n.status = status;
717        n.ended_at = Some(now);
718    }
719    for child in n.children.iter_mut() {
720        cascade_terminate(child, status, now);
721    }
722}
723
724fn find_node<'a>(nodes: &'a [WorkflowNode], id: &str) -> Option<&'a WorkflowNode> {
725    for n in nodes {
726        if n.id == id {
727            return Some(n);
728        }
729        if let Some(hit) = find_node(&n.children, id) {
730            return Some(hit);
731        }
732    }
733    None
734}
735
736fn find_node_mut<'a>(nodes: &'a mut [WorkflowNode], id: &str) -> Option<&'a mut WorkflowNode> {
737    for n in nodes.iter_mut() {
738        if n.id == id {
739            return Some(n);
740        }
741        if let Some(hit) = find_node_mut(&mut n.children, id) {
742            return Some(hit);
743        }
744    }
745    None
746}
747
748fn scope_id(run_id: &str, node_id: &str) -> String {
749    format!("{run_id}::{node_id}")
750}
751
752fn tool_node_id(run_id: &str, tool_use_id: &str) -> String {
753    format!("tool:{run_id}:{tool_use_id}")
754}
755
756fn find_tool_node_by_tool_use_id<'a>(
757    nodes: &'a mut [WorkflowNode],
758    tool_use_id: &str,
759) -> Option<&'a mut WorkflowNode> {
760    for n in nodes.iter_mut() {
761        if let WorkflowNodeKind::ToolCall {
762            tool_use_id: tid, ..
763        } = &n.kind
764            && tid == tool_use_id
765        {
766            return Some(n);
767        }
768        if let Some(hit) = find_tool_node_by_tool_use_id(&mut n.children, tool_use_id) {
769            return Some(hit);
770        }
771    }
772    None
773}
774
775fn parse_branch_index(node_id: &str) -> Option<usize> {
776    let start = node_id.rfind(".branch[")?;
777    let rest = &node_id[start + ".branch[".len()..];
778    let end = rest.find(']')?;
779    rest[..end].parse().ok()
780}
781
782/// Rebuild a full workflow tree from a session's event log. Replays every
783/// FlowStart / FlowNodeStart / FlowNodeEnd / FlowEnd / FlowGraph event through
784/// a fresh WorkflowGraph so the complete executor tree (root + subflows) is
785/// restored on session reopen.
786pub fn rebuild_workflow_tree(events: &[crate::event::Event]) -> WorkflowGraph {
787    let mut g = WorkflowGraph::new(crate::event::TurnId::now());
788    for ev in events {
789        g.apply_event(ev);
790    }
791    g
792}
793
794/// Rebuild a single FlowRun's message segment from the event log, filtered by
795/// run_id. Returns AssistantMsg + ToolResultMsg messages tagged with the given
796/// flow_run_id, in event order.
797pub fn rebuild_messages_for_run(
798    events: &[crate::event::Event],
799    run_id: &crate::event::FlowRunId,
800) -> Vec<crate::message::Message> {
801    events
802        .iter()
803        .filter_map(|ev| match ev {
804            crate::event::Event::AssistantMsg {
805                flow_run_id: Some(rid),
806                message,
807                ..
808            } if rid == run_id => Some(message.clone()),
809            crate::event::Event::ToolResultMsg {
810                flow_run_id: Some(rid),
811                message,
812                ..
813            } if rid == run_id => Some(message.clone()),
814            _ => None,
815        })
816        .collect()
817}
818
819#[cfg(test)]
820mod tests {
821    use super::*;
822    use crate::event::{FlowRunId, FlowStatus};
823    use crate::nodegraph::NodeKind;
824
825    fn flow_start(run_id: FlowRunId, name: &str) -> Event {
826        Event::FlowStart {
827            run_id,
828            flow_name: name.into(),
829            parent_run_id: None,
830            parent_node_id: None,
831            spawned: false,
832        }
833    }
834
835    fn subflow_start(child: FlowRunId, parent: FlowRunId, parent_node: &str, name: &str) -> Event {
836        Event::FlowStart {
837            run_id: child,
838            flow_name: name.into(),
839            parent_run_id: Some(parent),
840            parent_node_id: Some(parent_node.into()),
841            spawned: false,
842        }
843    }
844
845    fn stmt_start(run_id: FlowRunId, node_id: &str, parent: Option<&str>) -> Event {
846        Event::FlowNodeStart {
847            run_id,
848            node_id: node_id.into(),
849            kind: NodeKind::UserConfirm,
850            label: node_id.into(),
851            parent_node_id: parent.map(String::from),
852        }
853    }
854
855    fn stmt_end(run_id: FlowRunId, node_id: &str, status: FlowNodeStatus) -> Event {
856        Event::FlowNodeEnd {
857            run_id,
858            node_id: node_id.into(),
859            status,
860            output_preview: None,
861        }
862    }
863
864    #[test]
865    fn top_level_flow_becomes_root_child() {
866        let mut g = WorkflowGraph::new(TurnId::now());
867        let rid = FlowRunId::now();
868        g.apply_event(&flow_start(rid.clone(), "main"));
869        assert_eq!(g.root.len(), 1);
870        let flow = &g.root[0];
871        assert!(matches!(flow.kind, WorkflowNodeKind::Flow { .. }));
872        assert_eq!(flow.status, NodeStatus::Running);
873        assert_eq!(flow.id, rid.0.to_string());
874    }
875
876    #[test]
877    fn subflow_attaches_under_parent_node() {
878        let mut g = WorkflowGraph::new(TurnId::now());
879        let parent_flow = FlowRunId::now();
880        let child_flow = FlowRunId::now();
881        g.apply_event(&flow_start(parent_flow.clone(), "outer"));
882        g.apply_event(&stmt_start(parent_flow.clone(), "stmt_0", None));
883        g.apply_event(&subflow_start(
884            child_flow.clone(),
885            parent_flow.clone(),
886            "stmt_0",
887            "inner",
888        ));
889        let scoped = scope_id(&parent_flow.0.to_string(), "stmt_0");
890        let stmt = g.find_node(&scoped).unwrap();
891        assert_eq!(stmt.children.len(), 1);
892        assert!(matches!(
893            stmt.children[0].kind,
894            WorkflowNodeKind::Subflow { .. }
895        ));
896        assert_eq!(stmt.children[0].id, child_flow.0.to_string());
897    }
898
899    #[test]
900    fn tool_node_attaches_and_flow_end_marks_status() {
901        let mut g = WorkflowGraph::new(TurnId::now());
902        let rid = FlowRunId::now();
903        g.apply_event(&flow_start(rid.clone(), "main"));
904        g.apply_event(&stmt_start(rid.clone(), "stmt_0", None));
905        g.apply_event(&Event::ToolNode {
906            run_id: rid.clone(),
907            parent_node_id: "stmt_0".into(),
908            tool_use_id: "tu_1".into(),
909            tool_name: "fs.read".into(),
910            args_preview: "{\"path\":\"a\"}".into(),
911        });
912        g.apply_event(&stmt_end(rid.clone(), "stmt_0", FlowNodeStatus::Ok));
913        g.apply_event(&Event::FlowEnd {
914            run_id: rid.clone(),
915            flow_name: "main".into(),
916            status: FlowStatus::Ok,
917        });
918        let scoped = scope_id(&rid.0.to_string(), "stmt_0");
919        let stmt = g.find_node(&scoped).unwrap();
920        assert_eq!(stmt.status, NodeStatus::Ok);
921        assert_eq!(stmt.children.len(), 1);
922        let tool = &stmt.children[0];
923        assert_eq!(tool.id, tool_node_id(&rid.0.to_string(), "tu_1"));
924        assert!(matches!(tool.kind, WorkflowNodeKind::ToolCall { .. }));
925        assert_eq!(g.root[0].status, NodeStatus::Ok);
926    }
927
928    #[test]
929    fn fanout_branch_marks_parent_parallel() {
930        let mut g = WorkflowGraph::new(TurnId::now());
931        let rid = FlowRunId::now();
932        g.apply_event(&flow_start(rid.clone(), "main"));
933        g.apply_event(&stmt_start(rid.clone(), "stmt_1", None));
934        g.apply_event(&stmt_start(rid.clone(), "stmt_1.branch[0]", Some("stmt_1")));
935        g.apply_event(&stmt_start(rid.clone(), "stmt_1.branch[1]", Some("stmt_1")));
936        let scoped = scope_id(&rid.0.to_string(), "stmt_1");
937        let parent = g.find_node(&scoped).unwrap();
938        assert_eq!(parent.parallelism, Parallelism::Parallel);
939        assert_eq!(parent.children.len(), 2);
940        assert!(matches!(
941            parent.children[0].kind,
942            WorkflowNodeKind::FanoutBranch { branch_index: 0 }
943        ));
944        assert!(matches!(
945            parent.children[1].kind,
946            WorkflowNodeKind::FanoutBranch { branch_index: 1 }
947        ));
948    }
949
950    #[test]
951    fn out_of_order_events_silently_dropped() {
952        let mut g = WorkflowGraph::new(TurnId::now());
953        g.apply_event(&stmt_start(FlowRunId::now(), "stmt_0", Some("missing")));
954        g.apply_event(&Event::ToolNode {
955            run_id: FlowRunId::now(),
956            parent_node_id: "missing".into(),
957            tool_use_id: "tu".into(),
958            tool_name: "t".into(),
959            args_preview: "{}".into(),
960        });
961        assert!(g.root.is_empty());
962    }
963}
964
965#[test]
966fn rebuild_workflow_tree_restores_root_and_subflow() {
967    use crate::event::{Event, FlowRunId};
968    let root = FlowRunId::now();
969    let child = FlowRunId::now();
970    let events = vec![
971        Event::FlowStart {
972            run_id: root.clone(),
973            flow_name: "agent".into(),
974            parent_run_id: None,
975            parent_node_id: None,
976            spawned: false,
977        },
978        Event::FlowNodeStart {
979            run_id: root.clone(),
980            node_id: "stmt_0".into(),
981            kind: crate::nodegraph::NodeKind::Llm { model: None },
982            label: "llm".into(),
983            parent_node_id: None,
984        },
985        Event::FlowStart {
986            run_id: child.clone(),
987            flow_name: "subagent".into(),
988            parent_run_id: Some(root.clone()),
989            parent_node_id: Some("stmt_0".into()),
990            spawned: false,
991        },
992        Event::FlowEnd {
993            run_id: child.clone(),
994            flow_name: "subagent".into(),
995            status: crate::event::FlowStatus::Ok,
996        },
997        Event::FlowEnd {
998            run_id: root.clone(),
999            flow_name: "agent".into(),
1000            status: crate::event::FlowStatus::Ok,
1001        },
1002    ];
1003    let g = rebuild_workflow_tree(&events);
1004    assert_eq!(g.root.len(), 1, "one root flow");
1005    assert_eq!(g.root[0].id, root.0.to_string());
1006    // subflow attaches under stmt_0 (or root if parent node not found)
1007    assert!(!g.root[0].children.is_empty(), "subflow nested under root");
1008}
1009
1010#[test]
1011fn rebuild_messages_for_run_filters_by_run_id() {
1012    use crate::event::{Event, FlowRunId, TurnId};
1013    use crate::message::{Message, MessagePart, MessageRole};
1014    let root = FlowRunId::now();
1015    let child = FlowRunId::now();
1016    let tid = TurnId::now();
1017    let mk = |role, text, rid: Option<FlowRunId>| Event::AssistantMsg {
1018        turn_id: tid.clone(),
1019        flow_run_id: rid,
1020        message: Message {
1021            role,
1022            parts: vec![MessagePart::Text { text }],
1023            turn_id: tid.clone(),
1024        },
1025    };
1026    let events = vec![
1027        mk(MessageRole::Assistant, "root reply".into(), None),
1028        mk(
1029            MessageRole::Assistant,
1030            "child reply".into(),
1031            Some(child.clone()),
1032        ),
1033        mk(MessageRole::Assistant, "root again".into(), None),
1034    ];
1035    // Root agent messages are tagged None (restored via envelope, not per-run rebuild);
1036    // only sub-agent (child) messages carry Some(run_id) and are picked up here.
1037    let root_msgs = rebuild_messages_for_run(&events, &root);
1038    assert_eq!(root_msgs.len(), 0);
1039    let child_msgs = rebuild_messages_for_run(&events, &child);
1040    assert_eq!(child_msgs.len(), 1);
1041    assert_eq!(child_msgs[0].text_concat(), "child reply");
1042}
atman_runtime/workflow.rs

atman_runtime/
workflow.rs
