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        #[serde(default)]
56        provider: String,
57        #[serde(default)]
58        context_call_purpose: crate::context_plan::ContextCallPurpose,
59        #[serde(default)]
60        context_call_scope: crate::context_plan::ContextCallScope,
61        input_tokens: u64,
62        output_tokens: u64,
63        cache_read: u64,
64        cache_write: u64,
65        ttft_ms: u64,
66        tokens_per_second: f64,
67        wallclock_ms: u64,
68        run_id: Option<String>,
69        node_id: Option<String>,
70    },
71    ToolUseStart {
72        tool: String,
73        args_preview: String,
74        id: String,
75    },
76    ToolUseDone {
77        tool: String,
78        ok: bool,
79        preview: String,
80        id: String,
81    },
82    Note(String),
83    /// Rich notification with level/location/lifecycle/stack.
84    Notification(NotificationFrame),
85    FlowGraph {
86        run_id: String,
87        graph: crate::nodegraph::FlowGraph,
88    },
89    FlowStart {
90        run_id: String,
91        flow_name: String,
92        #[serde(default)]
93        parent_run_id: Option<String>,
94        #[serde(default)]
95        parent_node_id: Option<String>,
96    },
97    FlowNodeStart {
98        run_id: String,
99        node_id: String,
100        kind: crate::nodegraph::NodeKind,
101        label: String,
102        #[serde(default)]
103        parent_node_id: Option<String>,
104    },
105    FlowNodeEnd {
106        run_id: String,
107        node_id: String,
108        status: crate::event::FlowNodeStatus,
109        output_preview: Option<String>,
110        #[serde(default)]
111        parent_node_id: Option<String>,
112    },
113    FlowDone {
114        run_id: String,
115        flow_name: String,
116        ok: bool,
117        #[serde(default)]
118        cancelled: bool,
119        #[serde(default)]
120        suicide: bool,
121    },
122    ToolNode {
123        run_id: String,
124        parent_node_id: String,
125        tool_use_id: String,
126        tool: String,
127        args_preview: String,
128        #[serde(default, skip_serializing_if = "Option::is_none")]
129        call_intent: Option<crate::message::ToolCallIntent>,
130    },
131    AssistantMsg {
132        flow_run_id: Option<String>,
133        message: crate::message::Message,
134    },
135    ToolResultMsg {
136        flow_run_id: Option<String>,
137        message: crate::message::Message,
138    },
139    ToolPendingApproval {
140        run_id: String,
141        tool_use_id: String,
142        tool_name: String,
143        args_preview: String,
144        level: String,
145        #[serde(default, skip_serializing_if = "Option::is_none")]
146        preview: Option<String>,
147    },
148    ToolApproved {
149        run_id: String,
150        tool_use_id: String,
151        decided_by: String,
152    },
153    ToolDenied {
154        run_id: String,
155        tool_use_id: String,
156        reason: String,
157    },
158    PermissionRequestCreated {
159        run_id: String,
160        payload: crate::permission_audit::PermissionRequestAudit,
161    },
162    PermissionRequestTargeted {
163        run_id: String,
164        payload: crate::permission_audit::PermissionRequestAudit,
165    },
166    PermissionRequestDeferred {
167        run_id: String,
168        payload: crate::permission_audit::PermissionRequestAudit,
169    },
170    PermissionRequestApproved {
171        run_id: String,
172        payload: crate::permission_audit::PermissionRequestAudit,
173    },
174    PermissionRequestDenied {
175        run_id: String,
176        payload: crate::permission_audit::PermissionRequestAudit,
177    },
178    PermissionRequestCancelled {
179        run_id: String,
180        payload: crate::permission_audit::PermissionRequestAudit,
181    },
182    PermissionGroupCreated {
183        run_id: String,
184        payload: crate::permission_audit::PermissionGroupAudit,
185    },
186    PermissionGroupUpdated {
187        run_id: String,
188        payload: crate::permission_audit::PermissionGroupAudit,
189    },
190    PermissionGroupResolved {
191        run_id: String,
192        payload: crate::permission_audit::PermissionGroupAudit,
193    },
194    PermissionGrantCreated {
195        run_id: String,
196        payload: crate::permission_audit::PermissionGrantAudit,
197    },
198    PermissionGrantExpired {
199        run_id: String,
200        payload: crate::permission_audit::PermissionGrantAudit,
201    },
202    UnrestrictedExecution {
203        run_id: String,
204        payload: crate::permission_audit::PermissionRequestAudit,
205    },
206    TerminalChunk {
207        handle: String,
208        bytes: Vec<u8>,
209        screen: Option<crate::tools::term::TerminalScreen>,
210        state: crate::tools::term::TermStateSnapshot,
211        #[serde(default, skip_serializing_if = "Option::is_none")]
212        call_intent: Option<crate::message::ToolCallIntent>,
213        #[serde(default)]
214        run_id: Option<String>,
215    },
216    TerminalExited {
217        handle: String,
218        exit_code: Option<i32>,
219        #[serde(default, skip_serializing_if = "Option::is_none")]
220        call_intent: Option<crate::message::ToolCallIntent>,
221        #[serde(default)]
222        run_id: Option<String>,
223    },
224    BashChunk {
225        handle: String,
226        kind: String,
227        line: String,
228        #[serde(default, skip_serializing_if = "Option::is_none")]
229        call_intent: Option<crate::message::ToolCallIntent>,
230        #[serde(default)]
231        run_id: Option<String>,
232    },
233    BashExited {
234        handle: String,
235        exit_code: Option<i32>,
236        #[serde(default)]
237        error: Option<String>,
238        #[serde(default, skip_serializing_if = "Option::is_none")]
239        call_intent: Option<crate::message::ToolCallIntent>,
240        #[serde(default)]
241        run_id: Option<String>,
242    },
243    DiffPreview {
244        title: String,
245        old_content: Option<String>,
246        new_content: Option<String>,
247        unified_diff: Option<String>,
248        #[serde(default)]
249        run_id: Option<String>,
250    },
251    CompactionSummary {
252        phase: CompactionPhase,
253        range_start: usize,
254        range_end: usize,
255        summary: String,
256        before_tokens: u64,
257        after_tokens: u64,
258        compacted_count: usize,
259    },
260    MermaidDiagram {
261        source: String,
262    },
263    SubAgentStarted {
264        handle: String,
265        goal: String,
266        child_run_id: String,
267        model: String,
268    },
269    SubAgentDone {
270        handle: String,
271        status: String,
272        final_text: String,
273    },
274    #[serde(other)]
275    Unknown,
276}
277
278/// Extract the run_id (or flow_run_id) from any StreamFrame variant that carries one.
279/// Used to route frames to the correct sub-agent's entry / TUI item.
280pub fn frame_run_id(frame: &StreamFrame) -> Option<&str> {
281    match frame {
282        StreamFrame::FlowStart { run_id, .. }
283        | StreamFrame::FlowNodeStart { run_id, .. }
284        | StreamFrame::FlowNodeEnd { run_id, .. }
285        | StreamFrame::FlowDone { run_id, .. }
286        | StreamFrame::FlowGraph { run_id, .. }
287        | StreamFrame::ToolNode { run_id, .. }
288        | StreamFrame::ToolPendingApproval { run_id, .. }
289        | StreamFrame::ToolApproved { run_id, .. }
290        | StreamFrame::ToolDenied { run_id, .. }
291        | StreamFrame::PermissionRequestCreated { run_id, .. }
292        | StreamFrame::PermissionRequestTargeted { run_id, .. }
293        | StreamFrame::PermissionRequestDeferred { run_id, .. }
294        | StreamFrame::PermissionRequestApproved { run_id, .. }
295        | StreamFrame::PermissionRequestDenied { run_id, .. }
296        | StreamFrame::PermissionRequestCancelled { run_id, .. }
297        | StreamFrame::PermissionGroupCreated { run_id, .. }
298        | StreamFrame::PermissionGroupUpdated { run_id, .. }
299        | StreamFrame::PermissionGroupResolved { run_id, .. }
300        | StreamFrame::PermissionGrantCreated { run_id, .. }
301        | StreamFrame::PermissionGrantExpired { run_id, .. }
302        | StreamFrame::UnrestrictedExecution { run_id, .. } => Some(run_id.as_str()),
303        StreamFrame::AssistantMsg {
304            flow_run_id: Some(rid),
305            ..
306        }
307        | StreamFrame::ToolResultMsg {
308            flow_run_id: Some(rid),
309            ..
310        }
311        | StreamFrame::LlmCallStats {
312            run_id: Some(rid), ..
313        } => Some(rid.as_str()),
314        StreamFrame::LlmChunk {
315            run_id: Some(rid), ..
316        }
317        | StreamFrame::ThinkingChunk {
318            run_id: Some(rid), ..
319        }
320        | StreamFrame::LlmDone {
321            run_id: Some(rid), ..
322        }
323        | StreamFrame::TerminalChunk {
324            run_id: Some(rid), ..
325        }
326        | StreamFrame::TerminalExited {
327            run_id: Some(rid), ..
328        }
329        | StreamFrame::BashChunk {
330            run_id: Some(rid), ..
331        }
332        | StreamFrame::BashExited {
333            run_id: Some(rid), ..
334        }
335        | StreamFrame::DiffPreview {
336            run_id: Some(rid), ..
337        } => Some(rid.as_str()),
338        _ => None,
339    }
340}
341
342#[cfg(test)]
343mod tests {
344    use super::*;
345
346    #[test]
347    fn tool_node_round_trips() {
348        let f = StreamFrame::ToolNode {
349            run_id: "r".into(),
350            parent_node_id: "stmt_0".into(),
351            tool_use_id: "tu_1".into(),
352            tool: "fs.read".into(),
353            args_preview: "{}".into(),
354            call_intent: None,
355        };
356        let json = serde_json::to_string(&f).unwrap();
357        let back: StreamFrame = serde_json::from_str(&json).unwrap();
358        assert!(matches!(back, StreamFrame::ToolNode { .. }));
359    }
360
361    #[test]
362    fn legacy_llm_stats_default_route_metadata() {
363        let json = r#"{"LlmCallStats":{"model":"m","input_tokens":1,"output_tokens":2,"cache_read":3,"cache_write":4,"ttft_ms":5,"tokens_per_second":6.0,"wallclock_ms":7,"run_id":null,"node_id":null}}"#;
364        let frame: StreamFrame = serde_json::from_str(json).unwrap();
365
366        assert!(matches!(
367            frame,
368            StreamFrame::LlmCallStats {
369                provider,
370                context_call_purpose: crate::context_plan::ContextCallPurpose::General,
371                context_call_scope: crate::context_plan::ContextCallScope::Detached,
372                ..
373            } if provider.is_empty()
374        ));
375    }
376
377    #[test]
378    fn flow_node_start_serde_carries_parent() {
379        let f = StreamFrame::FlowNodeStart {
380            run_id: "r".into(),
381            node_id: "stmt_1.branch[0]".into(),
382            kind: crate::nodegraph::NodeKind::UserConfirm,
383            label: "b".into(),
384            parent_node_id: Some("stmt_1".into()),
385        };
386        let json = serde_json::to_string(&f).unwrap();
387        assert!(json.contains("\"parent_node_id\":\"stmt_1\""));
388        let back: StreamFrame = serde_json::from_str(&json).unwrap();
389        if let StreamFrame::FlowNodeStart { parent_node_id, .. } = back {
390            assert_eq!(parent_node_id.as_deref(), Some("stmt_1"));
391        } else {
392            panic!("wrong variant");
393        }
394    }
395
396    #[test]
397    fn unknown_bare_variant_falls_back() {
398        let payload = r#""SomeFutureFrame""#;
399        let back: StreamFrame = serde_json::from_str(payload).unwrap();
400        assert!(matches!(back, StreamFrame::Unknown));
401    }
402
403    #[test]
404    fn terminal_chunk_round_trips() {
405        let screen = crate::tools::term::TerminalScreen {
406            rows: 2,
407            cols: 3,
408            cells: vec![
409                crate::tools::term::TerminalCell {
410                    chars: "A".into(),
411                    ..Default::default()
412                },
413                crate::tools::term::TerminalCell::default(),
414                crate::tools::term::TerminalCell::default(),
415                crate::tools::term::TerminalCell::default(),
416                crate::tools::term::TerminalCell::default(),
417                crate::tools::term::TerminalCell::default(),
418            ],
419            cursor: Some((0, 0)),
420            alt_screen: false,
421        };
422        let f = StreamFrame::TerminalChunk {
423            handle: "term_s_0".into(),
424            bytes: b"hi".to_vec(),
425            screen: Some(screen),
426            state: crate::tools::term::TermStateSnapshot::Running,
427            call_intent: crate::message::ToolCallIntent::new("Inspect terminal output"),
428            run_id: None,
429        };
430        let json = serde_json::to_string(&f).unwrap();
431        let back: StreamFrame = serde_json::from_str(&json).unwrap();
432        match back {
433            StreamFrame::TerminalChunk {
434                handle,
435                bytes,
436                screen,
437                state,
438                call_intent,
439                run_id,
440            } => {
441                assert_eq!(handle, "term_s_0");
442                assert_eq!(bytes, b"hi");
443                assert!(run_id.is_none());
444                assert_eq!(
445                    call_intent.as_ref().map(|intent| intent.as_str()),
446                    Some("Inspect terminal output")
447                );
448                let screen = screen.expect("screen should be Some");
449                assert_eq!(screen.rows, 2);
450                assert_eq!(screen.cols, 3);
451                assert_eq!(screen.cells.len(), 6);
452                assert_eq!(screen.cells[0].chars, "A");
453                assert!(matches!(
454                    state,
455                    crate::tools::term::TermStateSnapshot::Running
456                ));
457            }
458            _ => panic!("wrong variant"),
459        }
460    }
461
462    #[test]
463    fn bash_exited_error_round_trips_and_legacy_payload_loads() {
464        let frame = StreamFrame::BashExited {
465            handle: "bg_s_1".into(),
466            exit_code: None,
467            error: Some("open log: permission denied".into()),
468            call_intent: crate::message::ToolCallIntent::new("Run verification"),
469            run_id: None,
470        };
471        let json = serde_json::to_string(&frame).unwrap();
472        let back: StreamFrame = serde_json::from_str(&json).unwrap();
473        match back {
474            StreamFrame::BashExited {
475                error, call_intent, ..
476            } => {
477                assert_eq!(error.as_deref(), Some("open log: permission denied"));
478                assert_eq!(
479                    call_intent.as_ref().map(|intent| intent.as_str()),
480                    Some("Run verification")
481                );
482            }
483            _ => panic!("wrong variant"),
484        }
485
486        let legacy = r#"{"BashExited":{"handle":"bg_s_1","exit_code":null,"run_id":null}}"#;
487        let back: StreamFrame = serde_json::from_str(legacy).unwrap();
488        match back {
489            StreamFrame::BashExited {
490                error, call_intent, ..
491            } => {
492                assert!(error.is_none());
493                assert!(call_intent.is_none());
494            }
495            _ => panic!("wrong variant"),
496        }
497    }
498
499    #[test]
500    fn terminal_exited_round_trips() {
501        let f = StreamFrame::TerminalExited {
502            handle: "term_s_1".into(),
503            exit_code: Some(0),
504            call_intent: None,
505            run_id: None,
506        };
507        let json = serde_json::to_string(&f).unwrap();
508        let back: StreamFrame = serde_json::from_str(&json).unwrap();
509        match back {
510            StreamFrame::TerminalExited {
511                handle, exit_code, ..
512            } => {
513                assert_eq!(handle, "term_s_1");
514                assert_eq!(exit_code, Some(0));
515            }
516            _ => panic!("wrong variant"),
517        }
518
519        let legacy = r#"{"TerminalExited":{"handle":"term_s_1","exit_code":0,"run_id":null}}"#;
520        let back: StreamFrame = serde_json::from_str(legacy).unwrap();
521        assert!(matches!(
522            back,
523            StreamFrame::TerminalExited {
524                call_intent: None,
525                ..
526            }
527        ));
528    }
529
530    #[test]
531    fn compaction_summary_round_trips() {
532        let f = StreamFrame::CompactionSummary {
533            phase: CompactionPhase::Running,
534            range_start: 3,
535            range_end: 11,
536            summary: String::new(),
537            before_tokens: 42,
538            after_tokens: 0,
539            compacted_count: 8,
540        };
541        let json = serde_json::to_string(&f).unwrap();
542        let back: StreamFrame = serde_json::from_str(&json).unwrap();
543        match back {
544            StreamFrame::CompactionSummary {
545                phase,
546                range_start,
547                range_end,
548                compacted_count,
549                ..
550            } => {
551                assert_eq!(phase, CompactionPhase::Running);
552                assert_eq!(range_start, 3);
553                assert_eq!(range_end, 11);
554                assert_eq!(compacted_count, 8);
555            }
556            _ => panic!("wrong variant"),
557        }
558    }
559}
atman_runtime/stream.rs

atman_runtime/
stream.rs
