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