atman_runtime/
stream.rs

1use crate::notify::{NotifyLevel, NotifyLifecycle, NotifyLocation, NotifyStack};
2use serde::{Deserialize, Serialize};
3
4#[derive(Debug, Clone, Serialize, Deserialize)]
5pub struct NotificationFrame {
6    pub level: NotifyLevel,
7    pub location: NotifyLocation,
8    pub lifecycle: NotifyLifecycle,
9    pub stack: NotifyStack,
10    pub message: String,
11}
12
13impl From<crate::notify::Notification> for NotificationFrame {
14    fn from(n: crate::notify::Notification) -> Self {
15        Self {
16            level: n.level,
17            location: n.location,
18            lifecycle: n.lifecycle,
19            stack: n.stack,
20            message: n.message,
21        }
22    }
23}
24
25#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
26#[serde(rename_all = "snake_case")]
27pub enum CompactionPhase {
28    Running,
29    Finished,
30    Failed,
31}
32
33#[derive(Debug, Clone, Serialize, Deserialize)]
34pub enum StreamFrame {
35    LlmChunk {
36        text: String,
37        model: String,
38        #[serde(default)]
39        run_id: Option<String>,
40    },
41    ThinkingChunk {
42        text: String,
43        #[serde(default)]
44        run_id: Option<String>,
45    },
46    LlmDone {
47        total_tokens: u64,
48        #[serde(default)]
49        run_id: Option<String>,
50    },
51    /// Discard streaming output from the previous attempt before retrying.
52    LlmRetry,
53    LlmCallStats {
54        model: String,
55        input_tokens: u64,
56        output_tokens: u64,
57        cache_read: u64,
58        cache_write: u64,
59        ttft_ms: u64,
60        tokens_per_second: f64,
61        wallclock_ms: u64,
62        run_id: Option<String>,
63        node_id: Option<String>,
64    },
65    ToolUseStart {
66        tool: String,
67        args_preview: String,
68        id: String,
69    },
70    ToolUseDone {
71        tool: String,
72        ok: bool,
73        preview: String,
74        id: String,
75    },
76    Note(String),
77    /// Rich notification with level/location/lifecycle/stack.
78    Notification(NotificationFrame),
79    FlowGraph {
80        run_id: String,
81        graph: crate::nodegraph::FlowGraph,
82    },
83    FlowStart {
84        run_id: String,
85        flow_name: String,
86        #[serde(default)]
87        parent_run_id: Option<String>,
88        #[serde(default)]
89        parent_node_id: Option<String>,
90    },
91    FlowNodeStart {
92        run_id: String,
93        node_id: String,
94        kind: crate::nodegraph::NodeKind,
95        label: String,
96        #[serde(default)]
97        parent_node_id: Option<String>,
98    },
99    FlowNodeEnd {
100        run_id: String,
101        node_id: String,
102        status: crate::event::FlowNodeStatus,
103        output_preview: Option<String>,
104        #[serde(default)]
105        parent_node_id: Option<String>,
106    },
107    FlowDone {
108        run_id: String,
109        flow_name: String,
110        ok: bool,
111        #[serde(default)]
112        cancelled: bool,
113    },
114    ToolNode {
115        run_id: String,
116        parent_node_id: String,
117        tool_use_id: String,
118        tool: String,
119        args_preview: String,
120    },
121    AssistantMsg {
122        flow_run_id: Option<String>,
123        message: crate::message::Message,
124    },
125    ToolResultMsg {
126        flow_run_id: Option<String>,
127        message: crate::message::Message,
128    },
129    ToolPendingApproval {
130        run_id: String,
131        tool_use_id: String,
132        tool_name: String,
133        args_preview: String,
134        level: String,
135        #[serde(default, skip_serializing_if = "Option::is_none")]
136        preview: Option<String>,
137    },
138    ToolApproved {
139        run_id: String,
140        tool_use_id: String,
141        decided_by: String,
142    },
143    ToolDenied {
144        run_id: String,
145        tool_use_id: String,
146        reason: String,
147    },
148    TerminalChunk {
149        handle: String,
150        bytes: Vec<u8>,
151        screen: Option<crate::tools::term::TerminalScreen>,
152        state: crate::tools::term::TermStateSnapshot,
153        #[serde(default)]
154        run_id: Option<String>,
155    },
156    TerminalExited {
157        handle: String,
158        exit_code: Option<i32>,
159        #[serde(default)]
160        run_id: Option<String>,
161    },
162    BashChunk {
163        handle: String,
164        kind: String,
165        line: String,
166        #[serde(default)]
167        run_id: Option<String>,
168    },
169    BashExited {
170        handle: String,
171        exit_code: Option<i32>,
172        #[serde(default)]
173        run_id: Option<String>,
174    },
175    DiffPreview {
176        title: String,
177        old_content: Option<String>,
178        new_content: Option<String>,
179        unified_diff: Option<String>,
180        #[serde(default)]
181        run_id: Option<String>,
182    },
183    CompactionSummary {
184        phase: CompactionPhase,
185        range_start: usize,
186        range_end: usize,
187        summary: String,
188        before_tokens: u64,
189        after_tokens: u64,
190        compacted_count: usize,
191    },
192    MermaidDiagram {
193        source: String,
194    },
195    SubAgentStarted {
196        handle: String,
197        goal: String,
198        child_run_id: String,
199        model: String,
200    },
201    SubAgentDone {
202        handle: String,
203        status: String,
204        final_text: String,
205    },
206    #[serde(other)]
207    Unknown,
208}
209
210/// Extract the run_id (or flow_run_id) from any StreamFrame variant that carries one.
211/// Used to route frames to the correct sub-agent's entry / TUI item.
212pub fn frame_run_id(frame: &StreamFrame) -> Option<&str> {
213    match frame {
214        StreamFrame::FlowStart { run_id, .. }
215        | StreamFrame::FlowNodeStart { run_id, .. }
216        | StreamFrame::FlowNodeEnd { run_id, .. }
217        | StreamFrame::FlowDone { run_id, .. }
218        | StreamFrame::FlowGraph { run_id, .. }
219        | StreamFrame::ToolNode { run_id, .. } => Some(run_id.as_str()),
220        StreamFrame::AssistantMsg {
221            flow_run_id: Some(rid),
222            ..
223        }
224        | StreamFrame::ToolResultMsg {
225            flow_run_id: Some(rid),
226            ..
227        }
228        | StreamFrame::LlmCallStats {
229            run_id: Some(rid), ..
230        } => Some(rid.as_str()),
231        StreamFrame::LlmChunk {
232            run_id: Some(rid), ..
233        }
234        | StreamFrame::ThinkingChunk {
235            run_id: Some(rid), ..
236        }
237        | StreamFrame::LlmDone {
238            run_id: Some(rid), ..
239        }
240        | StreamFrame::TerminalChunk {
241            run_id: Some(rid), ..
242        }
243        | StreamFrame::TerminalExited {
244            run_id: Some(rid), ..
245        }
246        | StreamFrame::BashChunk {
247            run_id: Some(rid), ..
248        }
249        | StreamFrame::BashExited {
250            run_id: Some(rid), ..
251        }
252        | StreamFrame::DiffPreview {
253            run_id: Some(rid), ..
254        } => Some(rid.as_str()),
255        _ => None,
256    }
257}
258
259#[cfg(test)]
260mod tests {
261    use super::*;
262
263    #[test]
264    fn tool_node_round_trips() {
265        let f = StreamFrame::ToolNode {
266            run_id: "r".into(),
267            parent_node_id: "stmt_0".into(),
268            tool_use_id: "tu_1".into(),
269            tool: "fs.read".into(),
270            args_preview: "{}".into(),
271        };
272        let json = serde_json::to_string(&f).unwrap();
273        let back: StreamFrame = serde_json::from_str(&json).unwrap();
274        assert!(matches!(back, StreamFrame::ToolNode { .. }));
275    }
276
277    #[test]
278    fn flow_node_start_serde_carries_parent() {
279        let f = StreamFrame::FlowNodeStart {
280            run_id: "r".into(),
281            node_id: "stmt_1.branch[0]".into(),
282            kind: crate::nodegraph::NodeKind::UserConfirm,
283            label: "b".into(),
284            parent_node_id: Some("stmt_1".into()),
285        };
286        let json = serde_json::to_string(&f).unwrap();
287        assert!(json.contains("\"parent_node_id\":\"stmt_1\""));
288        let back: StreamFrame = serde_json::from_str(&json).unwrap();
289        if let StreamFrame::FlowNodeStart { parent_node_id, .. } = back {
290            assert_eq!(parent_node_id.as_deref(), Some("stmt_1"));
291        } else {
292            panic!("wrong variant");
293        }
294    }
295
296    #[test]
297    fn unknown_bare_variant_falls_back() {
298        let payload = r#""SomeFutureFrame""#;
299        let back: StreamFrame = serde_json::from_str(payload).unwrap();
300        assert!(matches!(back, StreamFrame::Unknown));
301    }
302
303    #[test]
304    fn terminal_chunk_round_trips() {
305        let screen = crate::tools::term::TerminalScreen {
306            rows: 2,
307            cols: 3,
308            cells: vec![
309                crate::tools::term::TerminalCell {
310                    chars: "A".into(),
311                    ..Default::default()
312                },
313                crate::tools::term::TerminalCell::default(),
314                crate::tools::term::TerminalCell::default(),
315                crate::tools::term::TerminalCell::default(),
316                crate::tools::term::TerminalCell::default(),
317                crate::tools::term::TerminalCell::default(),
318            ],
319            cursor: Some((0, 0)),
320            alt_screen: false,
321        };
322        let f = StreamFrame::TerminalChunk {
323            handle: "term_s_0".into(),
324            bytes: b"hi".to_vec(),
325            screen: Some(screen),
326            state: crate::tools::term::TermStateSnapshot::Running,
327            run_id: None,
328        };
329        let json = serde_json::to_string(&f).unwrap();
330        let back: StreamFrame = serde_json::from_str(&json).unwrap();
331        match back {
332            StreamFrame::TerminalChunk {
333                handle,
334                bytes,
335                screen,
336                state,
337                run_id,
338            } => {
339                assert_eq!(handle, "term_s_0");
340                assert_eq!(bytes, b"hi");
341                assert!(run_id.is_none());
342                let screen = screen.expect("screen should be Some");
343                assert_eq!(screen.rows, 2);
344                assert_eq!(screen.cols, 3);
345                assert_eq!(screen.cells.len(), 6);
346                assert_eq!(screen.cells[0].chars, "A");
347                assert!(matches!(
348                    state,
349                    crate::tools::term::TermStateSnapshot::Running
350                ));
351            }
352            _ => panic!("wrong variant"),
353        }
354    }
355
356    #[test]
357    fn terminal_exited_round_trips() {
358        let f = StreamFrame::TerminalExited {
359            handle: "term_s_1".into(),
360            exit_code: Some(0),
361            run_id: None,
362        };
363        let json = serde_json::to_string(&f).unwrap();
364        let back: StreamFrame = serde_json::from_str(&json).unwrap();
365        match back {
366            StreamFrame::TerminalExited {
367                handle, exit_code, ..
368            } => {
369                assert_eq!(handle, "term_s_1");
370                assert_eq!(exit_code, Some(0));
371            }
372            _ => panic!("wrong variant"),
373        }
374    }
375
376    #[test]
377    fn compaction_summary_round_trips() {
378        let f = StreamFrame::CompactionSummary {
379            phase: CompactionPhase::Running,
380            range_start: 3,
381            range_end: 11,
382            summary: String::new(),
383            before_tokens: 42,
384            after_tokens: 0,
385            compacted_count: 8,
386        };
387        let json = serde_json::to_string(&f).unwrap();
388        let back: StreamFrame = serde_json::from_str(&json).unwrap();
389        match back {
390            StreamFrame::CompactionSummary {
391                phase,
392                range_start,
393                range_end,
394                compacted_count,
395                ..
396            } => {
397                assert_eq!(phase, CompactionPhase::Running);
398                assert_eq!(range_start, 3);
399                assert_eq!(range_end, 11);
400                assert_eq!(compacted_count, 8);
401            }
402            _ => panic!("wrong variant"),
403        }
404    }
405}
atman_runtime/stream.rs

atman_runtime/
stream.rs
