atman_runtime/
message_stream.rs

1//! Incrementally maintained message accumulator.  `window()` returns a
2//! `MessageWindow` anchored at the last compaction summary (zero-copy);
3//! `full_messages()` returns a shared `Arc<Vec<Message>>` of every message.
4
5use std::ops::Deref;
6use std::sync::{Arc, Mutex};
7
8use crate::compaction::is_compaction_summary;
9use crate::event::EventEnvelope;
10use crate::message::Message;
11
12#[derive(Clone)]
13pub struct MessageWindow {
14    messages: Arc<Vec<Message>>,
15    start: usize,
16}
17
18impl MessageWindow {
19    pub fn to_vec(&self) -> Vec<Message> {
20        self.as_slice().to_vec()
21    }
22
23    fn as_slice(&self) -> &[Message] {
24        &self.messages[self.start..]
25    }
26}
27
28impl Deref for MessageWindow {
29    type Target = [Message];
30
31    fn deref(&self) -> &Self::Target {
32        self.as_slice()
33    }
34}
35
36struct Acc {
37    compacted: Vec<(u64, Message)>,
38    full_raw: Vec<(u64, Message)>,
39    replayed: usize,
40    full_cache: Arc<Vec<Message>>,
41    window_cache: MessageWindow,
42}
43
44pub struct MessageStream {
45    events: Arc<Mutex<Vec<EventEnvelope>>>,
46    initial_compacted: Vec<(u64, Message)>,
47    initial_raw: Vec<(u64, Message)>,
48    acc: Mutex<Acc>,
49}
50
51impl MessageStream {
52    pub fn new(events: Arc<Mutex<Vec<EventEnvelope>>>) -> Self {
53        let empty = Arc::new(Vec::new());
54        Self {
55            events,
56            initial_compacted: Vec::new(),
57            initial_raw: Vec::new(),
58            acc: Mutex::new(Acc {
59                compacted: Vec::new(),
60                full_raw: Vec::new(),
61                replayed: 0,
62                full_cache: Arc::clone(&empty),
63                window_cache: MessageWindow {
64                    messages: empty,
65                    start: 0,
66                },
67            }),
68        }
69    }
70
71    pub fn with_initial(
72        events: Arc<Mutex<Vec<EventEnvelope>>>,
73        compacted: Vec<(u64, Message)>,
74        raw: Vec<(u64, Message)>,
75    ) -> Self {
76        let full: Arc<Vec<Message>> = Arc::new(raw.iter().map(|(_, msg)| msg.clone()).collect());
77        let window_messages: Arc<Vec<Message>> =
78            Arc::new(compacted.iter().map(|(_, msg)| msg.clone()).collect());
79        let start = window_messages
80            .iter()
81            .rposition(is_compaction_summary)
82            .unwrap_or(0);
83        let window = MessageWindow {
84            messages: window_messages,
85            start,
86        };
87        Self {
88            events,
89            initial_compacted: compacted.clone(),
90            initial_raw: raw.clone(),
91            acc: Mutex::new(Acc {
92                compacted,
93                full_raw: raw,
94                replayed: 0,
95                full_cache: full,
96                window_cache: window,
97            }),
98        }
99    }
100
101    pub fn full_messages(&self) -> Arc<Vec<Message>> {
102        let events = self.events.lock().expect("events poisoned");
103        let mut acc = self.acc.lock().expect("acc poisoned");
104        self.ensure_fresh_locked(&events, &mut acc);
105        Arc::clone(&acc.full_cache)
106    }
107
108    pub fn window(&self) -> MessageWindow {
109        let events = self.events.lock().expect("events poisoned");
110        let mut acc = self.acc.lock().expect("acc poisoned");
111        self.ensure_fresh_locked(&events, &mut acc);
112        acc.window_cache.clone()
113    }
114
115    fn ensure_fresh_locked(&self, events: &[EventEnvelope], acc: &mut Acc) {
116        if acc.compacted.is_empty() {
117            acc.compacted = self.initial_compacted.clone();
118            acc.full_raw = self.initial_raw.clone();
119        }
120        if acc.replayed >= events.len() {
121            return;
122        }
123        for ev in &events[acc.replayed..] {
124            crate::projection::message_window::apply_envelope_to_messages(ev, &mut acc.compacted);
125            match &ev.event {
126                crate::event::Event::UserMsg {
127                    message,
128                    flow_run_id: None,
129                    ..
130                }
131                | crate::event::Event::AssistantMsg {
132                    message,
133                    flow_run_id: None,
134                    ..
135                }
136                | crate::event::Event::ToolResultMsg {
137                    message,
138                    flow_run_id: None,
139                    ..
140                }
141                | crate::event::Event::SystemMsg { message, .. } => {
142                    acc.full_raw.push((ev.seq, message.clone()));
143                }
144                _ => {}
145            }
146        }
147        acc.replayed = events.len();
148
149        let compacted: Vec<Message> = acc.compacted.iter().map(|(_, msg)| msg).cloned().collect();
150        let start = compacted
151            .iter()
152            .rposition(is_compaction_summary)
153            .unwrap_or(0);
154        let compacted_arc = Arc::new(compacted);
155        acc.window_cache = MessageWindow {
156            messages: Arc::clone(&compacted_arc),
157            start,
158        };
159
160        let raw: Vec<Message> = acc.full_raw.iter().map(|(_, msg)| msg).cloned().collect();
161        acc.full_cache = Arc::new(raw);
162    }
163}
164
165#[cfg(test)]
166mod tests {
167    use super::*;
168    use crate::event::TurnId;
169    use crate::event::{Event, EventEnvelope};
170    use crate::message::{MessageOrigin, MessagePart, MessageRole};
171
172    fn user(text: &str) -> Message {
173        Message {
174            role: MessageRole::User,
175            parts: vec![MessagePart::Text {
176                text: text.to_string(),
177            }],
178            turn_id: TurnId::now(),
179            origin: MessageOrigin::User,
180        }
181    }
182
183    fn assistant(text: &str) -> Message {
184        Message {
185            role: MessageRole::Assistant,
186            parts: vec![MessagePart::Text {
187                text: text.to_string(),
188            }],
189            turn_id: TurnId::now(),
190            origin: MessageOrigin::User,
191        }
192    }
193
194    fn compact_summary(text: &str) -> Message {
195        Message::system_compact_summary(TurnId::now(), text, 0, 1, 2)
196    }
197
198    fn make_msg_event(ty: &str, msg: &Message, _seq: u64) -> Event {
199        match ty {
200            "user_msg" => Event::UserMsg {
201                turn_id: msg.turn_id.clone(),
202                flow_run_id: None,
203                message: msg.clone(),
204            },
205            "assistant_msg" => Event::AssistantMsg {
206                turn_id: msg.turn_id.clone(),
207                flow_run_id: None,
208                message: msg.clone(),
209            },
210            "system_msg" => Event::SystemMsg {
211                turn_id: msg.turn_id.clone(),
212                message: msg.clone(),
213            },
214            _ => unreachable!(),
215        }
216    }
217
218    fn make_context_compact(
219        range_start: u64,
220        range_end: u64,
221        before_tokens: u64,
222        after_tokens: u64,
223        summary_text: &str,
224        replacement_msg_seq: u64,
225    ) -> Event {
226        Event::ContextCompact {
227            session_id: "test".into(),
228            before_tokens,
229            after_tokens,
230            compacted_range_start: range_start,
231            compacted_range_end: range_end,
232            summary_text: Some(summary_text.into()),
233            replacement_msg_seq: Some(replacement_msg_seq),
234        }
235    }
236
237    fn event_envelopes(events: Vec<Event>) -> Arc<Mutex<Vec<EventEnvelope>>> {
238        Arc::new(Mutex::new(
239            events
240                .into_iter()
241                .enumerate()
242                .map(|(i, event)| EventEnvelope::new((i + 1) as u64, event))
243                .collect(),
244        ))
245    }
246
247    #[test]
248    fn full_messages_filters_only_message_events() {
249        let u1 = user("hello");
250        let a1 = assistant("hi there");
251        let events = event_envelopes(vec![
252            make_msg_event("user_msg", &u1, 1),
253            Event::TurnStart {
254                turn_id: TurnId::now(),
255            },
256            make_msg_event("assistant_msg", &a1, 2),
257            Event::LlmCall {
258                model: "m".into(),
259                provider: "p".into(),
260                usage: crate::provider::TokenUsage::default(),
261                wallclock_ms: 0,
262                ttft_ms: None,
263                tokens_per_second: None,
264                status: crate::event::LlmCallStatus::Ok,
265                run_id: None,
266                node_id: None,
267            },
268        ]);
269        let ms = MessageStream::new(events);
270        let msgs = ms.full_messages();
271        assert_eq!(msgs.len(), 2);
272        assert_eq!(msgs[0].text_concat(), "hello");
273        assert_eq!(msgs[1].text_concat(), "hi there");
274    }
275
276    #[test]
277    fn window_no_summary_returns_all() {
278        let events = vec![
279            make_msg_event("user_msg", &user("a"), 1),
280            make_msg_event("assistant_msg", &assistant("b"), 2),
281            make_msg_event("user_msg", &user("c"), 3),
282        ];
283        let ms = MessageStream::new(event_envelopes(events));
284        assert_eq!(ms.window().len(), 3);
285    }
286
287    #[test]
288    fn window_single_summary_starts_from_it() {
289        let s1 = compact_summary("summary 1");
290        let events = vec![
291            make_msg_event("user_msg", &user("old"), 1),
292            make_msg_event("assistant_msg", &assistant("old"), 2),
293            make_msg_event("system_msg", &s1, 3),
294            make_msg_event("user_msg", &user("new"), 4),
295            make_msg_event("assistant_msg", &assistant("new"), 5),
296        ];
297        let ms = MessageStream::new(event_envelopes(events));
298        let w = ms.window();
299        assert_eq!(w.len(), 3);
300        assert!(matches!(w[0].parts[0], MessagePart::CompactSummary { .. }));
301    }
302
303    #[test]
304    fn window_multiple_summaries_uses_last() {
305        let s1 = compact_summary("summary 1");
306        let s2 = compact_summary("summary 2");
307        let events = vec![
308            make_msg_event("system_msg", &s1, 1),
309            make_msg_event("user_msg", &user("m1"), 2),
310            make_msg_event("system_msg", &s2, 3),
311            make_msg_event("user_msg", &user("m2"), 4),
312        ];
313        let ms = MessageStream::new(event_envelopes(events));
314        let w = ms.window();
315        assert_eq!(w.len(), 2);
316        assert!(matches!(w[0].parts[0], MessagePart::CompactSummary { .. }));
317        if let MessagePart::CompactSummary { summary, .. } = &w[0].parts[0] {
318            assert_eq!(summary, "summary 2");
319        }
320    }
321
322    #[test]
323    fn window_no_prefix_before_summary() {
324        let s1 = compact_summary("summary");
325        let events = vec![
326            make_msg_event("user_msg", &user("very old"), 1),
327            make_msg_event("assistant_msg", &assistant("very old"), 2),
328            make_msg_event("system_msg", &s1, 3),
329            make_msg_event("user_msg", &user("new"), 4),
330        ];
331        let ms = MessageStream::new(event_envelopes(events));
332        let w = ms.window();
333        assert_eq!(w.len(), 2);
334        assert!(matches!(w[0].parts[0], MessagePart::CompactSummary { .. }));
335        assert_eq!(w[1].text_concat(), "new");
336    }
337
338    #[test]
339    fn window_empty_stream_returns_empty() {
340        let ms = MessageStream::new(event_envelopes(Vec::new()));
341        assert!(ms.window().is_empty());
342    }
343
344    #[test]
345    fn context_compact_replaces_range_with_summary() {
346        let events = vec![
347            make_msg_event("user_msg", &user("old u1"), 1),
348            make_msg_event("assistant_msg", &assistant("old a1"), 2),
349            make_msg_event("user_msg", &user("old u2"), 3),
350            make_msg_event("system_msg", &compact_summary("summary"), 4),
351            make_context_compact(0, 2, 100, 50, "compaction summary text", 4),
352            make_msg_event("user_msg", &user("after compact"), 5),
353        ];
354        let ms = MessageStream::new(event_envelopes(events));
355        let w = ms.window();
356        assert_eq!(w.len(), 2);
357        assert!(matches!(w[0].parts[0], MessagePart::CompactSummary { .. }));
358    }
359
360    #[test]
361    fn multiple_compactions_applied_in_order() {
362        let events = vec![
363            make_msg_event("user_msg", &user("a"), 1),
364            make_msg_event("assistant_msg", &assistant("b"), 2),
365            make_msg_event("system_msg", &compact_summary("s1"), 3),
366            make_context_compact(0, 1, 200, 100, "first summary", 3),
367            make_msg_event("user_msg", &user("c"), 4),
368            make_msg_event("assistant_msg", &assistant("d"), 5),
369            make_msg_event("system_msg", &compact_summary("s2"), 6),
370            make_context_compact(1, 2, 150, 80, "second summary", 6),
371            make_msg_event("user_msg", &user("e"), 7),
372        ];
373        let ms = MessageStream::new(event_envelopes(events));
374        let w = ms.window();
375        assert_eq!(w.len(), 2);
376        assert!(matches!(w[0].parts[0], MessagePart::CompactSummary { .. }));
377        if let MessagePart::CompactSummary { summary, .. } = &w[0].parts[0] {
378            assert_eq!(summary, "second summary");
379        }
380    }
381
382    #[test]
383    fn compact_then_user_message_produces_summary_plus_user() {
384        let events = vec![
385            make_msg_event("user_msg", &user("old u1"), 1),
386            make_msg_event("assistant_msg", &assistant("old a1"), 2),
387            make_msg_event("user_msg", &user("old u2"), 3),
388            make_msg_event("system_msg", &compact_summary("compact summary"), 4),
389            make_context_compact(0, 2, 200, 100, "compact summary", 4),
390            make_msg_event("user_msg", &user("new message after compact"), 5),
391        ];
392        let ms = MessageStream::new(event_envelopes(events));
393        let w = ms.window();
394        assert_eq!(w.len(), 2);
395        assert!(matches!(w[0].parts[0], MessagePart::CompactSummary { .. }));
396        assert_eq!(w[1].text_concat(), "new message after compact");
397    }
398
399    #[test]
400    fn no_compaction_window_equals_full_messages() {
401        let events = vec![
402            make_msg_event("user_msg", &user("first"), 1),
403            make_msg_event("assistant_msg", &assistant("second"), 2),
404            make_msg_event("user_msg", &user("third"), 3),
405        ];
406        let ms = MessageStream::new(event_envelopes(events));
407        assert_eq!(ms.full_messages().len(), 3);
408        assert_eq!(ms.window().len(), 3);
409    }
410
411    #[test]
412    fn full_messages_retains_compacted_history() {
413        let events = vec![
414            make_msg_event("user_msg", &user("old u1"), 1),
415            make_msg_event("assistant_msg", &assistant("old a1"), 2),
416            make_msg_event("user_msg", &user("old u2"), 3),
417            make_msg_event("system_msg", &compact_summary("summary"), 4),
418            make_context_compact(0, 2, 200, 100, "summary", 4),
419            make_msg_event("user_msg", &user("after compact"), 5),
420        ];
421        let ms = MessageStream::new(event_envelopes(events));
422        // Window: only compact summary + messages after it
423        let w = ms.window();
424        assert_eq!(w.len(), 2);
425        // Full: all messages including compacted ones
426        let f = ms.full_messages();
427        assert_eq!(f.len(), 5, "full must retain compacted messages");
428        assert_eq!(f[0].text_concat(), "old u1");
429        assert_eq!(f[1].text_concat(), "old a1");
430        assert_eq!(f[2].text_concat(), "old u2");
431    }
432
433    #[test]
434    fn third_compaction_replaces_second_summary() {
435        let events = vec![
436            make_msg_event("user_msg", &user("a"), 1),
437            make_msg_event("assistant_msg", &assistant("b"), 2),
438            make_msg_event("system_msg", &compact_summary("s1"), 3),
439            make_context_compact(0, 1, 100, 50, "s1 text", 3),
440            make_msg_event("user_msg", &user("c"), 4),
441            make_msg_event("system_msg", &compact_summary("s2"), 5),
442            make_context_compact(0, 1, 80, 40, "s2 text", 6),
443            make_msg_event("user_msg", &user("d"), 6),
444            make_msg_event("system_msg", &compact_summary("s3"), 7),
445            make_context_compact(0, 1, 70, 30, "s3 text", 9),
446            make_msg_event("user_msg", &user("final"), 8),
447        ];
448        let ms = MessageStream::new(event_envelopes(events));
449        let w = ms.window();
450        assert_eq!(w.len(), 2);
451        if let MessagePart::CompactSummary { summary, .. } = &w[0].parts[0] {
452            assert_eq!(summary, "s3 text");
453        }
454        assert_eq!(w[1].text_concat(), "final");
455    }
456
457    #[test]
458    fn checkpoint_replaces_live_window_and_accepts_following_messages() {
459        let old = user("old");
460        let summary = compact_summary("checkpoint summary");
461        let retained = user("retained current user");
462        let events = event_envelopes(vec![
463            make_msg_event("user_msg", &old, 1),
464            Event::Checkpoint {
465                session_id: "test".into(),
466                messages: vec![summary.clone(), retained.clone()],
467                window_tokens: 10,
468            },
469            make_msg_event("assistant_msg", &assistant("next provider output"), 3),
470        ]);
471        let ms = MessageStream::new(events);
472
473        let window = ms.window();
474        assert_eq!(window.len(), 3);
475        assert!(matches!(
476            window[0].parts[0],
477            MessagePart::CompactSummary { .. }
478        ));
479        assert_eq!(window[1].text_concat(), "retained current user");
480        assert_eq!(window[2].text_concat(), "next provider output");
481        assert!(!window.iter().any(|message| message.text_concat() == "old"));
482    }
483
484    #[test]
485    fn reopened_session_uses_compacted_window_before_new_events() {
486        let initial_compacted = vec![
487            (10, compact_summary("checkpoint summary")),
488            (11, assistant("retained tail")),
489        ];
490        let initial_raw = vec![(1, user("dead user")), (2, assistant("dead assistant"))];
491        let ms = MessageStream::with_initial(
492            Arc::new(Mutex::new(Vec::new())),
493            initial_compacted,
494            initial_raw,
495        );
496
497        let window = ms.window();
498        assert_eq!(window.len(), 2);
499        assert!(matches!(
500            window[0].parts[0],
501            MessagePart::CompactSummary { .. }
502        ));
503        assert_eq!(window[1].text_concat(), "retained tail");
504
505        let full = ms.full_messages();
506        assert_eq!(full.len(), 2);
507        assert_eq!(full[0].text_concat(), "dead user");
508        assert_eq!(full[1].text_concat(), "dead assistant");
509    }
510
511    #[test]
512    fn reopened_session_keeps_initial_messages_after_new_events() {
513        let initial_compacted = vec![
514            (1, compact_summary("compaction summary")),
515            (2, assistant("tail assistant")),
516        ];
517        let initial_raw = vec![
518            (1, compact_summary("compaction summary")),
519            (2, assistant("tail assistant")),
520        ];
521        let events = Arc::new(Mutex::new(Vec::new()));
522        let ms = MessageStream::with_initial(events.clone(), initial_compacted, initial_raw);
523
524        events.lock().unwrap().push(EventEnvelope::new(
525            1,
526            Event::TurnStart {
527                turn_id: TurnId::now(),
528            },
529        ));
530        events.lock().unwrap().push(EventEnvelope::new(
531            2,
532            Event::UserMsg {
533                turn_id: TurnId::now(),
534                flow_run_id: None,
535                message: user("latest user"),
536            },
537        ));
538
539        let w = ms.window();
540        assert_eq!(w.len(), 3);
541        assert!(matches!(w[0].parts[0], MessagePart::CompactSummary { .. }));
542        assert_eq!(w[1].text_concat(), "tail assistant");
543        assert_eq!(w[2].text_concat(), "latest user");
544    }
545
546    /// Regression: after a runtime ContextCompact, full_messages() must still
547    /// contain the pre-compact messages.
548    #[test]
549    fn full_messages_retains_pre_compact_history_after_runtime_compact() {
550        let initial_compacted = vec![
551            (1, compact_summary("prior summary")),
552            (2, user("old user")),
553            (3, assistant("old assistant")),
554        ];
555        let initial_raw = vec![
556            (1, compact_summary("prior summary")),
557            (2, user("old user")),
558            (3, assistant("old assistant")),
559        ];
560        let events = Arc::new(Mutex::new(Vec::new()));
561        let ms = MessageStream::with_initial(events.clone(), initial_compacted, initial_raw);
562
563        events.lock().unwrap().push(EventEnvelope::new(
564            10,
565            Event::UserMsg {
566                turn_id: TurnId::now(),
567                flow_run_id: None,
568                message: user("new user before compact"),
569            },
570        ));
571        events.lock().unwrap().push(EventEnvelope::new(
572            11,
573            Event::AssistantMsg {
574                turn_id: TurnId::now(),
575                flow_run_id: None,
576                message: assistant("new assistant before compact"),
577            },
578        ));
579
580        let before = ms.full_messages();
581        assert_eq!(before.len(), 5, "pre-compact full should have all 5 msgs");
582
583        events.lock().unwrap().push(EventEnvelope::new(
584            12,
585            Event::SystemMsg {
586                turn_id: TurnId::now(),
587                message: compact_summary("runtime summary"),
588            },
589        ));
590        events.lock().unwrap().push(EventEnvelope::new(
591            13,
592            Event::ContextCompact {
593                session_id: "test".into(),
594                before_tokens: 1000,
595                after_tokens: 100,
596                compacted_range_start: 1,
597                compacted_range_end: 2,
598                summary_text: Some("runtime summary".into()),
599                replacement_msg_seq: Some(12),
600            },
601        ));
602
603        events.lock().unwrap().push(EventEnvelope::new(
604            14,
605            Event::UserMsg {
606                turn_id: TurnId::now(),
607                flow_run_id: None,
608                message: user("after compact user"),
609            },
610        ));
611
612        let w = ms.window();
613        assert_eq!(w.len(), 4, "window after compact");
614        assert!(matches!(w[0].parts[0], MessagePart::CompactSummary { .. }));
615
616        let full = ms.full_messages();
617        assert!(
618            full.len() >= 6,
619            "full must retain pre-compact history, got {} msgs: {:?}",
620            full.len(),
621            full.iter().map(|m| m.text_concat()).collect::<Vec<_>>()
622        );
623        let texts: Vec<String> = full.iter().map(|m| m.text_concat()).collect();
624        assert!(
625            texts.iter().any(|t| t.contains("old user")),
626            "full must contain pre-compact 'old user', got: {:?}",
627            texts
628        );
629        assert!(
630            texts.iter().any(|t| t.contains("old assistant")),
631            "full must contain pre-compact 'old assistant', got: {:?}",
632            texts
633        );
634    }
635}
atman_runtime/message_stream.rs

atman_runtime/
message_stream.rs
