Skip to main content

atman_runtime/
stream.rs

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