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

atman_runtime/
workflow.rs
