Skip to main content

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    },
39    ThinkingChunk {
40        text: String,
41    },
42    LlmDone {
43        total_tokens: u64,
44    },
45    /// Discard streaming output from the previous attempt before retrying.
46    LlmRetry,
47    LlmCallStats {
48        model: String,
49        input_tokens: u64,
50        output_tokens: u64,
51        cache_read: u64,
52        cache_write: u64,
53        ttft_ms: u64,
54        tokens_per_second: f64,
55        wallclock_ms: u64,
56        run_id: Option<String>,
57        node_id: Option<String>,
58    },
59    ToolUseStart {
60        tool: String,
61        args_preview: String,
62        id: String,
63    },
64    ToolUseDone {
65        tool: String,
66        ok: bool,
67        preview: String,
68        id: String,
69    },
70    Note(String),
71    /// Rich notification with level/location/lifecycle/stack.
72    Notification(NotificationFrame),
73    FlowGraph {
74        run_id: String,
75        graph: crate::nodegraph::FlowGraph,
76    },
77    FlowStart {
78        run_id: String,
79        flow_name: String,
80        #[serde(default)]
81        parent_run_id: Option<String>,
82        #[serde(default)]
83        parent_node_id: Option<String>,
84    },
85    FlowNodeStart {
86        run_id: String,
87        node_id: String,
88        kind: crate::nodegraph::NodeKind,
89        label: String,
90        #[serde(default)]
91        parent_node_id: Option<String>,
92    },
93    FlowNodeEnd {
94        run_id: String,
95        node_id: String,
96        status: crate::event::FlowNodeStatus,
97        output_preview: Option<String>,
98        #[serde(default)]
99        parent_node_id: Option<String>,
100    },
101    FlowDone {
102        run_id: String,
103        flow_name: String,
104        ok: bool,
105        #[serde(default)]
106        cancelled: bool,
107    },
108    ToolNode {
109        run_id: String,
110        parent_node_id: String,
111        tool_use_id: String,
112        tool: String,
113        args_preview: String,
114    },
115    AssistantMsg {
116        flow_run_id: Option<String>,
117        message: crate::message::Message,
118    },
119    ToolResultMsg {
120        flow_run_id: Option<String>,
121        message: crate::message::Message,
122    },
123    ToolPendingApproval {
124        run_id: String,
125        tool_use_id: String,
126        tool_name: String,
127        args_preview: String,
128        level: String,
129        #[serde(default, skip_serializing_if = "Option::is_none")]
130        preview: Option<String>,
131    },
132    ToolApproved {
133        run_id: String,
134        tool_use_id: String,
135        decided_by: String,
136    },
137    ToolDenied {
138        run_id: String,
139        tool_use_id: String,
140        reason: String,
141    },
142    TerminalChunk {
143        handle: String,
144        bytes: Vec<u8>,
145        screen: Option<crate::tools::term::TerminalScreen>,
146        state: crate::tools::term::TermStateSnapshot,
147    },
148    TerminalExited {
149        handle: String,
150        exit_code: Option<i32>,
151    },
152    BashChunk {
153        handle: String,
154        kind: String,
155        line: String,
156    },
157    BashExited {
158        handle: String,
159        exit_code: Option<i32>,
160    },
161    DiffPreview {
162        title: String,
163        old_content: Option<String>,
164        new_content: Option<String>,
165        unified_diff: Option<String>,
166    },
167    CompactionSummary {
168        phase: CompactionPhase,
169        range_start: usize,
170        range_end: usize,
171        summary: String,
172        before_tokens: u64,
173        after_tokens: u64,
174        compacted_count: usize,
175    },
176    MermaidDiagram {
177        source: String,
178    },
179    #[serde(other)]
180    Unknown,
181}
182
183#[cfg(test)]
184mod tests {
185    use super::*;
186
187    #[test]
188    fn tool_node_round_trips() {
189        let f = StreamFrame::ToolNode {
190            run_id: "r".into(),
191            parent_node_id: "stmt_0".into(),
192            tool_use_id: "tu_1".into(),
193            tool: "fs.read".into(),
194            args_preview: "{}".into(),
195        };
196        let json = serde_json::to_string(&f).unwrap();
197        let back: StreamFrame = serde_json::from_str(&json).unwrap();
198        assert!(matches!(back, StreamFrame::ToolNode { .. }));
199    }
200
201    #[test]
202    fn flow_node_start_serde_carries_parent() {
203        let f = StreamFrame::FlowNodeStart {
204            run_id: "r".into(),
205            node_id: "stmt_1.branch[0]".into(),
206            kind: crate::nodegraph::NodeKind::UserConfirm,
207            label: "b".into(),
208            parent_node_id: Some("stmt_1".into()),
209        };
210        let json = serde_json::to_string(&f).unwrap();
211        assert!(json.contains("\"parent_node_id\":\"stmt_1\""));
212        let back: StreamFrame = serde_json::from_str(&json).unwrap();
213        if let StreamFrame::FlowNodeStart { parent_node_id, .. } = back {
214            assert_eq!(parent_node_id.as_deref(), Some("stmt_1"));
215        } else {
216            panic!("wrong variant");
217        }
218    }
219
220    #[test]
221    fn unknown_bare_variant_falls_back() {
222        let payload = r#""SomeFutureFrame""#;
223        let back: StreamFrame = serde_json::from_str(payload).unwrap();
224        assert!(matches!(back, StreamFrame::Unknown));
225    }
226
227    #[test]
228    fn terminal_chunk_round_trips() {
229        let screen = crate::tools::term::TerminalScreen {
230            rows: 2,
231            cols: 3,
232            cells: vec![
233                crate::tools::term::TerminalCell {
234                    chars: "A".into(),
235                    ..Default::default()
236                },
237                crate::tools::term::TerminalCell::default(),
238                crate::tools::term::TerminalCell::default(),
239                crate::tools::term::TerminalCell::default(),
240                crate::tools::term::TerminalCell::default(),
241                crate::tools::term::TerminalCell::default(),
242            ],
243            cursor: Some((0, 0)),
244            alt_screen: false,
245        };
246        let f = StreamFrame::TerminalChunk {
247            handle: "term_s_0".into(),
248            bytes: b"hi".to_vec(),
249            screen: Some(screen),
250            state: crate::tools::term::TermStateSnapshot::Running,
251        };
252        let json = serde_json::to_string(&f).unwrap();
253        let back: StreamFrame = serde_json::from_str(&json).unwrap();
254        match back {
255            StreamFrame::TerminalChunk {
256                handle,
257                bytes,
258                screen,
259                state,
260            } => {
261                assert_eq!(handle, "term_s_0");
262                assert_eq!(bytes, b"hi");
263                let screen = screen.expect("screen should be Some");
264                assert_eq!(screen.rows, 2);
265                assert_eq!(screen.cols, 3);
266                assert_eq!(screen.cells.len(), 6);
267                assert_eq!(screen.cells[0].chars, "A");
268                assert!(matches!(
269                    state,
270                    crate::tools::term::TermStateSnapshot::Running
271                ));
272            }
273            _ => panic!("wrong variant"),
274        }
275    }
276
277    #[test]
278    fn terminal_exited_round_trips() {
279        let f = StreamFrame::TerminalExited {
280            handle: "term_s_1".into(),
281            exit_code: Some(0),
282        };
283        let json = serde_json::to_string(&f).unwrap();
284        let back: StreamFrame = serde_json::from_str(&json).unwrap();
285        match back {
286            StreamFrame::TerminalExited { handle, exit_code } => {
287                assert_eq!(handle, "term_s_1");
288                assert_eq!(exit_code, Some(0));
289            }
290            _ => panic!("wrong variant"),
291        }
292    }
293
294    #[test]
295    fn compaction_summary_round_trips() {
296        let f = StreamFrame::CompactionSummary {
297            phase: CompactionPhase::Running,
298            range_start: 3,
299            range_end: 11,
300            summary: String::new(),
301            before_tokens: 42,
302            after_tokens: 0,
303            compacted_count: 8,
304        };
305        let json = serde_json::to_string(&f).unwrap();
306        let back: StreamFrame = serde_json::from_str(&json).unwrap();
307        match back {
308            StreamFrame::CompactionSummary {
309                phase,
310                range_start,
311                range_end,
312                compacted_count,
313                ..
314            } => {
315                assert_eq!(phase, CompactionPhase::Running);
316                assert_eq!(range_start, 3);
317                assert_eq!(range_end, 11);
318                assert_eq!(compacted_count, 8);
319            }
320            _ => panic!("wrong variant"),
321        }
322    }
323}