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