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        error: Option<String>,
174        #[serde(default)]
175        run_id: Option<String>,
176    },
177    DiffPreview {
178        title: String,
179        old_content: Option<String>,
180        new_content: Option<String>,
181        unified_diff: Option<String>,
182        #[serde(default)]
183        run_id: Option<String>,
184    },
185    CompactionSummary {
186        phase: CompactionPhase,
187        range_start: usize,
188        range_end: usize,
189        summary: String,
190        before_tokens: u64,
191        after_tokens: u64,
192        compacted_count: usize,
193    },
194    MermaidDiagram {
195        source: String,
196    },
197    SubAgentStarted {
198        handle: String,
199        goal: String,
200        child_run_id: String,
201        model: String,
202    },
203    SubAgentDone {
204        handle: String,
205        status: String,
206        final_text: String,
207    },
208    #[serde(other)]
209    Unknown,
210}
211
212/// Extract the run_id (or flow_run_id) from any StreamFrame variant that carries one.
213/// Used to route frames to the correct sub-agent's entry / TUI item.
214pub fn frame_run_id(frame: &StreamFrame) -> Option<&str> {
215    match frame {
216        StreamFrame::FlowStart { run_id, .. }
217        | StreamFrame::FlowNodeStart { run_id, .. }
218        | StreamFrame::FlowNodeEnd { run_id, .. }
219        | StreamFrame::FlowDone { run_id, .. }
220        | StreamFrame::FlowGraph { run_id, .. }
221        | StreamFrame::ToolNode { run_id, .. } => Some(run_id.as_str()),
222        StreamFrame::AssistantMsg {
223            flow_run_id: Some(rid),
224            ..
225        }
226        | StreamFrame::ToolResultMsg {
227            flow_run_id: Some(rid),
228            ..
229        }
230        | StreamFrame::LlmCallStats {
231            run_id: Some(rid), ..
232        } => Some(rid.as_str()),
233        StreamFrame::LlmChunk {
234            run_id: Some(rid), ..
235        }
236        | StreamFrame::ThinkingChunk {
237            run_id: Some(rid), ..
238        }
239        | StreamFrame::LlmDone {
240            run_id: Some(rid), ..
241        }
242        | StreamFrame::TerminalChunk {
243            run_id: Some(rid), ..
244        }
245        | StreamFrame::TerminalExited {
246            run_id: Some(rid), ..
247        }
248        | StreamFrame::BashChunk {
249            run_id: Some(rid), ..
250        }
251        | StreamFrame::BashExited {
252            run_id: Some(rid), ..
253        }
254        | StreamFrame::DiffPreview {
255            run_id: Some(rid), ..
256        } => Some(rid.as_str()),
257        _ => None,
258    }
259}
260
261#[cfg(test)]
262mod tests {
263    use super::*;
264
265    #[test]
266    fn tool_node_round_trips() {
267        let f = StreamFrame::ToolNode {
268            run_id: "r".into(),
269            parent_node_id: "stmt_0".into(),
270            tool_use_id: "tu_1".into(),
271            tool: "fs.read".into(),
272            args_preview: "{}".into(),
273        };
274        let json = serde_json::to_string(&f).unwrap();
275        let back: StreamFrame = serde_json::from_str(&json).unwrap();
276        assert!(matches!(back, StreamFrame::ToolNode { .. }));
277    }
278
279    #[test]
280    fn flow_node_start_serde_carries_parent() {
281        let f = StreamFrame::FlowNodeStart {
282            run_id: "r".into(),
283            node_id: "stmt_1.branch[0]".into(),
284            kind: crate::nodegraph::NodeKind::UserConfirm,
285            label: "b".into(),
286            parent_node_id: Some("stmt_1".into()),
287        };
288        let json = serde_json::to_string(&f).unwrap();
289        assert!(json.contains("\"parent_node_id\":\"stmt_1\""));
290        let back: StreamFrame = serde_json::from_str(&json).unwrap();
291        if let StreamFrame::FlowNodeStart { parent_node_id, .. } = back {
292            assert_eq!(parent_node_id.as_deref(), Some("stmt_1"));
293        } else {
294            panic!("wrong variant");
295        }
296    }
297
298    #[test]
299    fn unknown_bare_variant_falls_back() {
300        let payload = r#""SomeFutureFrame""#;
301        let back: StreamFrame = serde_json::from_str(payload).unwrap();
302        assert!(matches!(back, StreamFrame::Unknown));
303    }
304
305    #[test]
306    fn terminal_chunk_round_trips() {
307        let screen = crate::tools::term::TerminalScreen {
308            rows: 2,
309            cols: 3,
310            cells: vec![
311                crate::tools::term::TerminalCell {
312                    chars: "A".into(),
313                    ..Default::default()
314                },
315                crate::tools::term::TerminalCell::default(),
316                crate::tools::term::TerminalCell::default(),
317                crate::tools::term::TerminalCell::default(),
318                crate::tools::term::TerminalCell::default(),
319                crate::tools::term::TerminalCell::default(),
320            ],
321            cursor: Some((0, 0)),
322            alt_screen: false,
323        };
324        let f = StreamFrame::TerminalChunk {
325            handle: "term_s_0".into(),
326            bytes: b"hi".to_vec(),
327            screen: Some(screen),
328            state: crate::tools::term::TermStateSnapshot::Running,
329            run_id: None,
330        };
331        let json = serde_json::to_string(&f).unwrap();
332        let back: StreamFrame = serde_json::from_str(&json).unwrap();
333        match back {
334            StreamFrame::TerminalChunk {
335                handle,
336                bytes,
337                screen,
338                state,
339                run_id,
340            } => {
341                assert_eq!(handle, "term_s_0");
342                assert_eq!(bytes, b"hi");
343                assert!(run_id.is_none());
344                let screen = screen.expect("screen should be Some");
345                assert_eq!(screen.rows, 2);
346                assert_eq!(screen.cols, 3);
347                assert_eq!(screen.cells.len(), 6);
348                assert_eq!(screen.cells[0].chars, "A");
349                assert!(matches!(
350                    state,
351                    crate::tools::term::TermStateSnapshot::Running
352                ));
353            }
354            _ => panic!("wrong variant"),
355        }
356    }
357
358    #[test]
359    fn bash_exited_error_round_trips_and_legacy_payload_loads() {
360        let frame = StreamFrame::BashExited {
361            handle: "bg_s_1".into(),
362            exit_code: None,
363            error: Some("open log: permission denied".into()),
364            run_id: None,
365        };
366        let json = serde_json::to_string(&frame).unwrap();
367        let back: StreamFrame = serde_json::from_str(&json).unwrap();
368        match back {
369            StreamFrame::BashExited { error, .. } => {
370                assert_eq!(error.as_deref(), Some("open log: permission denied"));
371            }
372            _ => panic!("wrong variant"),
373        }
374
375        let legacy = r#"{"BashExited":{"handle":"bg_s_1","exit_code":null,"run_id":null}}"#;
376        let back: StreamFrame = serde_json::from_str(legacy).unwrap();
377        match back {
378            StreamFrame::BashExited { error, .. } => assert!(error.is_none()),
379            _ => panic!("wrong variant"),
380        }
381    }
382
383    #[test]
384    fn terminal_exited_round_trips() {
385        let f = StreamFrame::TerminalExited {
386            handle: "term_s_1".into(),
387            exit_code: Some(0),
388            run_id: None,
389        };
390        let json = serde_json::to_string(&f).unwrap();
391        let back: StreamFrame = serde_json::from_str(&json).unwrap();
392        match back {
393            StreamFrame::TerminalExited {
394                handle, exit_code, ..
395            } => {
396                assert_eq!(handle, "term_s_1");
397                assert_eq!(exit_code, Some(0));
398            }
399            _ => panic!("wrong variant"),
400        }
401    }
402
403    #[test]
404    fn compaction_summary_round_trips() {
405        let f = StreamFrame::CompactionSummary {
406            phase: CompactionPhase::Running,
407            range_start: 3,
408            range_end: 11,
409            summary: String::new(),
410            before_tokens: 42,
411            after_tokens: 0,
412            compacted_count: 8,
413        };
414        let json = serde_json::to_string(&f).unwrap();
415        let back: StreamFrame = serde_json::from_str(&json).unwrap();
416        match back {
417            StreamFrame::CompactionSummary {
418                phase,
419                range_start,
420                range_end,
421                compacted_count,
422                ..
423            } => {
424                assert_eq!(phase, CompactionPhase::Running);
425                assert_eq!(range_start, 3);
426                assert_eq!(range_end, 11);
427                assert_eq!(compacted_count, 8);
428            }
429            _ => panic!("wrong variant"),
430        }
431    }
432}
atman_runtime/stream.rs

atman_runtime/
stream.rs
