1use std::ops::ControlFlow;
11use std::path::Path;
12
13use zeph_llm::provider::{Message, MessagePart, Role};
14
15use crate::error::SessionError;
16use crate::event::{SessionEvent, SessionEventEnvelope};
17use crate::log::SessionEventLog;
18
19#[derive(Debug, Clone, Default)]
21pub struct ReconstructedState {
22 pub messages: Vec<Message>,
24 pub last_seq: Option<u64>,
26 pub provider_name: String,
27 pub model: String,
28 pub cwd: String,
29}
30
31pub struct ReplayEngine;
33
34impl ReplayEngine {
35 #[tracing::instrument(name = "session.replay.run", skip_all, level = "debug")]
52 pub async fn replay(
53 session_dir: &Path,
54 up_to: Option<u64>,
55 ) -> Result<ReconstructedState, SessionError> {
56 let log = SessionEventLog::open(session_dir).await?;
57
58 let mut messages: Vec<Message> = Vec::new();
59 let mut origin_seqs: Vec<u64> = Vec::new();
60 let mut state = ReconstructedState::default();
61
62 log.read_chunked(|chunk| {
63 for envelope in chunk {
64 if fold_step(&mut state, &mut messages, &mut origin_seqs, envelope, up_to)
65 .is_break()
66 {
67 return ControlFlow::Break(());
68 }
69 }
70 ControlFlow::Continue(())
71 })
72 .await?;
73
74 state.messages = messages;
75 Ok(state)
76 }
77
78 #[must_use]
83 #[tracing::instrument(
84 name = "session.replay.fold",
85 skip_all,
86 level = "debug",
87 fields(event_count = events.len())
88 )]
89 pub fn fold(events: Vec<SessionEventEnvelope>, up_to: Option<u64>) -> ReconstructedState {
90 let mut messages: Vec<Message> = Vec::new();
91 let mut origin_seqs: Vec<u64> = Vec::new();
94 let mut state = ReconstructedState::default();
95
96 for envelope in events {
97 if fold_step(&mut state, &mut messages, &mut origin_seqs, envelope, up_to).is_break() {
98 break;
99 }
100 }
101
102 state.messages = messages;
103 state
104 }
105}
106
107fn fold_step(
115 state: &mut ReconstructedState,
116 messages: &mut Vec<Message>,
117 origin_seqs: &mut Vec<u64>,
118 envelope: SessionEventEnvelope,
119 up_to: Option<u64>,
120) -> ControlFlow<()> {
121 if let Some(bound) = up_to
122 && envelope.seq >= bound
123 {
124 return ControlFlow::Break(());
125 }
126 let seq = envelope.seq;
127 state.last_seq = Some(seq);
128
129 match envelope.kind {
130 SessionEvent::SessionStarted {
131 cwd,
132 provider_name,
133 model,
134 ..
135 } => {
136 state.cwd = cwd;
137 state.provider_name = provider_name;
138 state.model = model;
139 }
140 SessionEvent::UserMessage { text, .. } => {
141 messages.push(Message::from_legacy(Role::User, text));
142 origin_seqs.push(seq);
143 }
144 SessionEvent::AssistantMessage { parts } => {
145 messages.push(Message::from_parts(Role::Assistant, parts));
146 origin_seqs.push(seq);
147 }
148 SessionEvent::ToolCall { id, name, input } => {
149 push_part_to_last_assistant(
150 messages,
151 origin_seqs,
152 seq,
153 MessagePart::ToolUse { id, name, input },
154 );
155 }
156 SessionEvent::ToolResult {
157 id,
158 output,
159 is_error,
160 ..
161 } => {
162 push_part_to_tool_result_batch(
163 messages,
164 origin_seqs,
165 seq,
166 MessagePart::ToolResult {
167 tool_use_id: id,
168 content: output,
169 is_error,
170 },
171 );
172 }
173 SessionEvent::Condensation {
174 replaced_seq_range: (lo, hi),
175 summary,
176 ..
177 } => {
178 replace_range(messages, origin_seqs, lo, hi, summary.to_markdown());
179 }
180 SessionEvent::Compaction { summary, .. } => {
181 if let Some(summary) = summary {
189 let hi = origin_seqs.last().copied().unwrap_or(seq);
190 replace_range(messages, origin_seqs, 0, hi, summary.to_markdown());
191 }
192 }
193 SessionEvent::ModelChanged {
194 provider_name,
195 model,
196 } => {
197 state.provider_name = provider_name;
198 state.model = model;
199 }
200 SessionEvent::ForkPoint { .. } | SessionEvent::SessionEnded { .. } => {}
201 }
202
203 ControlFlow::Continue(())
204}
205
206fn push_part_to_last_assistant(
211 messages: &mut Vec<Message>,
212 origin_seqs: &mut Vec<u64>,
213 seq: u64,
214 part: MessagePart,
215) {
216 if let Some(last) = messages.last_mut()
217 && last.role == Role::Assistant
218 {
219 last.parts.push(part);
220 return;
221 }
222 messages.push(Message::from_parts(Role::Assistant, vec![part]));
223 origin_seqs.push(seq);
224}
225
226fn push_part_to_tool_result_batch(
238 messages: &mut Vec<Message>,
239 origin_seqs: &mut Vec<u64>,
240 seq: u64,
241 part: MessagePart,
242) {
243 let is_open_batch = messages.last().is_some_and(|m| {
244 m.role == Role::User
245 && !m.parts.is_empty()
246 && m.parts
247 .iter()
248 .all(|p| matches!(p, MessagePart::ToolResult { .. }))
249 });
250 if is_open_batch {
251 let last = messages.last_mut().expect("checked by is_open_batch above");
252 last.parts.push(part);
253 last.rebuild_content();
254 return;
255 }
256 messages.push(Message::from_parts(Role::User, vec![part]));
257 origin_seqs.push(seq);
258}
259
260fn replace_range(
263 messages: &mut Vec<Message>,
264 origin_seqs: &mut Vec<u64>,
265 lo: u64,
266 hi: u64,
267 summary_text: String,
268) {
269 let mut new_messages = Vec::with_capacity(messages.len());
270 let mut new_seqs = Vec::with_capacity(origin_seqs.len());
271 let mut inserted = false;
272
273 for (message, seq) in messages.drain(..).zip(origin_seqs.drain(..)) {
274 if seq >= lo && seq <= hi {
275 if !inserted {
276 new_messages.push(Message::from_parts(
277 Role::System,
278 vec![MessagePart::Summary {
279 text: summary_text.clone(),
280 }],
281 ));
282 new_seqs.push(lo);
283 inserted = true;
284 }
285 continue;
286 }
287 new_messages.push(message);
288 new_seqs.push(seq);
289 }
290
291 if !inserted {
292 new_messages.push(Message::from_parts(
293 Role::System,
294 vec![MessagePart::Summary { text: summary_text }],
295 ));
296 new_seqs.push(lo);
297 }
298
299 *messages = new_messages;
300 *origin_seqs = new_seqs;
301}
302
303#[cfg(test)]
304mod tests {
305 use super::*;
306 use zeph_common::memory::AnchoredSummary;
307
308 fn envelope(seq: u64, kind: SessionEvent) -> SessionEventEnvelope {
309 SessionEventEnvelope::new(seq, None, None, kind)
310 }
311
312 #[tokio::test]
313 #[serial_test::serial(session_history_integrity)]
314 async fn test_replay_empty_session() {
315 let dir = tempfile::tempdir().unwrap();
316 let state = ReplayEngine::replay(dir.path(), None).await.unwrap();
317 assert!(state.messages.is_empty());
318 assert!(state.last_seq.is_none());
319 }
320
321 #[tokio::test]
322 #[serial_test::serial(session_history_integrity)]
323 async fn test_replay_basic_turn() {
324 let dir = tempfile::tempdir().unwrap();
325 let log = SessionEventLog::open(dir.path()).await.unwrap();
326 log.append(
327 None,
328 None,
329 SessionEvent::SessionStarted {
330 session_id: "s1".to_owned(),
331 cwd: "/repo".to_owned(),
332 provider_name: "claude".to_owned(),
333 model: "opus".to_owned(),
334 forked_from: None,
335 },
336 )
337 .await
338 .unwrap();
339 log.append(
340 Some(1),
341 None,
342 SessionEvent::UserMessage {
343 text: "hi".to_owned(),
344 image_refs: vec![],
345 },
346 )
347 .await
348 .unwrap();
349 log.append(
350 Some(1),
351 None,
352 SessionEvent::AssistantMessage {
353 parts: vec![MessagePart::Text {
354 text: "hello".to_owned(),
355 }],
356 },
357 )
358 .await
359 .unwrap();
360
361 let state = ReplayEngine::replay(dir.path(), None).await.unwrap();
362 assert_eq!(state.messages.len(), 2);
363 assert_eq!(state.messages[0].role, Role::User);
364 assert_eq!(state.messages[1].role, Role::Assistant);
365 assert_eq!(state.provider_name, "claude");
366 assert_eq!(state.cwd, "/repo");
367 assert_eq!(state.last_seq, Some(2));
368 }
369
370 #[test]
371 #[serial_test::serial(session_history_integrity)]
372 fn test_replay_tool_roundtrip() {
373 let events = vec![
374 envelope(
375 0,
376 SessionEvent::UserMessage {
377 text: "run ls".to_owned(),
378 image_refs: vec![],
379 },
380 ),
381 envelope(
382 1,
383 SessionEvent::ToolCall {
384 id: "tc1".to_owned(),
385 name: "shell".to_owned(),
386 input: serde_json::json!({"cmd": "ls"}),
387 },
388 ),
389 envelope(
390 2,
391 SessionEvent::ToolResult {
392 id: "tc1".to_owned(),
393 name: "shell".to_owned(),
394 output: "file.txt".to_owned(),
395 is_error: false,
396 duration_ms: 5,
397 },
398 ),
399 ];
400 let state = ReplayEngine::fold(events, None);
401 assert_eq!(
402 state.messages.len(),
403 3,
404 "user message + assistant ToolUse message + user ToolResult message (#5464: a \
405 ToolResult must never merge into the preceding Assistant message — OpenAI/Claude \
406 both require it in a separate Role::User message)"
407 );
408 let assistant = &state.messages[1];
409 assert_eq!(assistant.role, Role::Assistant);
410 assert_eq!(assistant.parts.len(), 1);
411 assert!(matches!(assistant.parts[0], MessagePart::ToolUse { .. }));
412
413 let tool_result_msg = &state.messages[2];
414 assert_eq!(tool_result_msg.role, Role::User);
415 assert_eq!(tool_result_msg.parts.len(), 1);
416 assert!(matches!(
417 tool_result_msg.parts[0],
418 MessagePart::ToolResult { .. }
419 ));
420 }
421
422 #[test]
423 #[serial_test::serial(session_history_integrity)]
424 fn test_replay_tool_result_batch_merges_into_one_user_message() {
425 let events = vec![
430 envelope(
431 0,
432 SessionEvent::AssistantMessage {
433 parts: vec![
434 MessagePart::ToolUse {
435 id: "tc1".to_owned(),
436 name: "shell".to_owned(),
437 input: serde_json::json!({}),
438 },
439 MessagePart::ToolUse {
440 id: "tc2".to_owned(),
441 name: "shell".to_owned(),
442 input: serde_json::json!({}),
443 },
444 ],
445 },
446 ),
447 envelope(
448 1,
449 SessionEvent::ToolResult {
450 id: "tc1".to_owned(),
451 name: "shell".to_owned(),
452 output: "a".to_owned(),
453 is_error: false,
454 duration_ms: 1,
455 },
456 ),
457 envelope(
458 2,
459 SessionEvent::ToolResult {
460 id: "tc2".to_owned(),
461 name: "shell".to_owned(),
462 output: "b".to_owned(),
463 is_error: false,
464 duration_ms: 1,
465 },
466 ),
467 ];
468 let state = ReplayEngine::fold(events, None);
469 assert_eq!(state.messages.len(), 2);
470 assert_eq!(state.messages[1].role, Role::User);
471 assert_eq!(state.messages[1].parts.len(), 2);
472 }
473
474 #[test]
475 #[serial_test::serial(session_history_integrity)]
476 fn test_replay_tool_result_never_merges_into_plain_user_message() {
477 let events = vec![
480 envelope(
481 0,
482 SessionEvent::UserMessage {
483 text: "hello".to_owned(),
484 image_refs: vec![],
485 },
486 ),
487 envelope(
488 1,
489 SessionEvent::ToolResult {
490 id: "tc1".to_owned(),
491 name: "shell".to_owned(),
492 output: "a".to_owned(),
493 is_error: false,
494 duration_ms: 1,
495 },
496 ),
497 ];
498 let state = ReplayEngine::fold(events, None);
499 assert_eq!(state.messages.len(), 2);
500 assert!(state.messages[0].parts.is_empty());
501 assert_eq!(state.messages[1].parts.len(), 1);
502 }
503
504 #[test]
505 #[serial_test::serial(session_history_integrity)]
506 fn test_replay_condensation_folds() {
507 let summary = AnchoredSummary {
508 session_intent: "test".to_owned(),
509 files_modified: vec![],
510 decisions_made: vec![],
511 open_questions: vec![],
512 next_steps: vec!["continue".to_owned()],
513 };
514 let events = vec![
515 envelope(
516 0,
517 SessionEvent::UserMessage {
518 text: "a".to_owned(),
519 image_refs: vec![],
520 },
521 ),
522 envelope(
523 1,
524 SessionEvent::AssistantMessage {
525 parts: vec![MessagePart::Text {
526 text: "b".to_owned(),
527 }],
528 },
529 ),
530 envelope(
531 2,
532 SessionEvent::Condensation {
533 replaced_seq_range: (0, 1),
534 summary,
535 tokens_before: 100,
536 tokens_after: 10,
537 },
538 ),
539 envelope(
540 3,
541 SessionEvent::UserMessage {
542 text: "c".to_owned(),
543 image_refs: vec![],
544 },
545 ),
546 ];
547 let state = ReplayEngine::fold(events, None);
548 assert_eq!(state.messages.len(), 2);
550 assert_eq!(state.messages[0].role, Role::System);
551 assert!(matches!(
552 state.messages[0].parts[0],
553 MessagePart::Summary { .. }
554 ));
555 assert_eq!(state.messages[1].role, Role::User);
556 }
557
558 #[test]
559 #[serial_test::serial(session_history_integrity)]
560 fn test_replay_stop_at_seq() {
561 let events = vec![
562 envelope(
563 0,
564 SessionEvent::UserMessage {
565 text: "a".to_owned(),
566 image_refs: vec![],
567 },
568 ),
569 envelope(
570 1,
571 SessionEvent::UserMessage {
572 text: "b".to_owned(),
573 image_refs: vec![],
574 },
575 ),
576 envelope(
577 2,
578 SessionEvent::UserMessage {
579 text: "c".to_owned(),
580 image_refs: vec![],
581 },
582 ),
583 ];
584 let state = ReplayEngine::fold(events, Some(2));
585 assert_eq!(state.messages.len(), 2);
586 assert_eq!(state.last_seq, Some(1));
587 }
588
589 #[allow(clippy::too_many_lines)] async fn seed_large_synthetic_log(dir: &Path, n_turns: u64) -> SessionEventLog {
598 use crate::event::CompactionTier;
599 use zeph_common::memory::AnchoredSummary;
600
601 let log = SessionEventLog::open(dir).await.unwrap();
602
603 log.append(
604 None,
605 None,
606 SessionEvent::SessionStarted {
607 session_id: "s1".to_owned(),
608 cwd: "/repo".to_owned(),
609 provider_name: "claude".to_owned(),
610 model: "opus".to_owned(),
611 forked_from: None,
612 },
613 )
614 .await
615 .unwrap();
616
617 for turn in 0..n_turns {
618 log.append(
619 Some(turn),
620 None,
621 SessionEvent::UserMessage {
622 text: format!("user turn {turn}"),
623 image_refs: vec![],
624 },
625 )
626 .await
627 .unwrap();
628
629 if turn % 7 == 0 {
630 log.append(
632 Some(turn),
633 None,
634 SessionEvent::AssistantMessage { parts: vec![] },
635 )
636 .await
637 .unwrap();
638 log.append(
639 Some(turn),
640 None,
641 SessionEvent::ToolCall {
642 id: format!("tc-{turn}"),
643 name: "shell".to_owned(),
644 input: serde_json::json!({"cmd": "ls"}),
645 },
646 )
647 .await
648 .unwrap();
649 log.append(
650 Some(turn),
651 None,
652 SessionEvent::ToolResult {
653 id: format!("tc-{turn}"),
654 name: "shell".to_owned(),
655 output: format!("output-{turn}"),
656 is_error: false,
657 duration_ms: 3,
658 },
659 )
660 .await
661 .unwrap();
662 } else {
663 log.append(
664 Some(turn),
665 None,
666 SessionEvent::AssistantMessage {
667 parts: vec![MessagePart::Text {
668 text: format!("assistant reply {turn}"),
669 }],
670 },
671 )
672 .await
673 .unwrap();
674 }
675
676 if turn == 100 {
677 log.append(
679 Some(turn),
680 None,
681 SessionEvent::Condensation {
682 replaced_seq_range: (0, 10),
683 summary: AnchoredSummary {
684 session_intent: "test".to_owned(),
685 files_modified: vec![],
686 decisions_made: vec![],
687 open_questions: vec![],
688 next_steps: vec!["continue".to_owned()],
689 },
690 tokens_before: 500,
691 tokens_after: 50,
692 },
693 )
694 .await
695 .unwrap();
696 }
697
698 if turn == 150 {
699 log.append(
703 Some(turn),
704 None,
705 SessionEvent::Compaction {
706 tier: CompactionTier::Hard,
707 cleared_count: 42,
708 summary: Some(AnchoredSummary {
709 session_intent: "test".to_owned(),
710 files_modified: vec![],
711 decisions_made: vec![],
712 open_questions: vec![],
713 next_steps: vec!["keep going".to_owned()],
714 }),
715 },
716 )
717 .await
718 .unwrap();
719 }
720 }
721
722 log.append(
725 None,
726 None,
727 SessionEvent::ForkPoint {
728 new_session_id: "child-of-s1".to_owned(),
729 },
730 )
731 .await
732 .unwrap();
733 log.append(
734 None,
735 None,
736 SessionEvent::SessionEnded {
737 reason: "user_quit".to_owned(),
738 },
739 )
740 .await
741 .unwrap();
742 log.append(
743 None,
744 None,
745 SessionEvent::ModelChanged {
746 provider_name: "openai".to_owned(),
747 model: "gpt-5.4".to_owned(),
748 },
749 )
750 .await
751 .unwrap();
752
753 log
754 }
755
756 #[tokio::test]
760 #[serial_test::serial(session_history_integrity)]
761 async fn test_replay_streaming_matches_vec_based_fold_on_large_log() {
762 const N_TURNS: u64 = 250; let dir = tempfile::tempdir().unwrap();
765 let log = seed_large_synthetic_log(dir.path(), N_TURNS).await;
766
767 let all_events = log.read_all().await.unwrap();
769 assert!(
770 all_events.len() > 100,
771 "synthetic log must exceed one REPLAY_CHUNK_SIZE to exercise multi-chunk streaming"
772 );
773 let vec_based = ReplayEngine::fold(all_events, None);
774
775 let streamed = ReplayEngine::replay(dir.path(), None).await.unwrap();
777
778 assert_eq!(streamed.last_seq, vec_based.last_seq);
779 assert_eq!(streamed.provider_name, vec_based.provider_name);
780 assert_eq!(streamed.model, vec_based.model);
781 assert_eq!(streamed.cwd, vec_based.cwd);
782 assert_eq!(streamed.messages.len(), vec_based.messages.len());
783 assert_eq!(
784 serde_json::to_string(&streamed.messages).unwrap(),
785 serde_json::to_string(&vec_based.messages).unwrap(),
786 "streaming replay must be byte-identical to the old Vec-based fold"
787 );
788 }
789
790 #[tokio::test]
797 #[serial_test::serial(session_history_integrity)]
798 async fn test_replay_torn_tail_across_chunk_boundary() {
799 const N_TURNS: u64 = 250; let dir = tempfile::tempdir().unwrap();
802 let log = seed_large_synthetic_log(dir.path(), N_TURNS).await;
803 let path = log.path().to_path_buf();
804 drop(log);
805
806 let full = tokio::fs::read(&path).await.unwrap();
808 let cut = full.len() - 5;
809 tokio::fs::write(&path, &full[..cut]).await.unwrap();
810
811 let whole_file_log = SessionEventLog::open(dir.path()).await.unwrap();
813 let whole_file_events = whole_file_log.read_all().await.unwrap();
814 assert!(
815 whole_file_events.len() > 100,
816 "test must still exceed one REPLAY_CHUNK_SIZE after dropping the torn line"
817 );
818 let expected = ReplayEngine::fold(whole_file_events, None);
819 assert_eq!(
820 expected.last_seq,
821 Some(576),
822 "the torn line must be the trailing seq-577 ModelChanged event (verified against the \
823 real fixture) — a differential-only assertion below would pass vacuously if both \
824 paths regressed identically and stopped dropping the torn tail at all"
825 );
826
827 let actual = ReplayEngine::replay(dir.path(), None).await.unwrap();
829
830 assert_eq!(
831 actual.last_seq, expected.last_seq,
832 "chunked replay must drop the torn tail at the same seq as the whole-file path"
833 );
834 assert_eq!(
835 serde_json::to_string(&actual.messages).unwrap(),
836 serde_json::to_string(&expected.messages).unwrap(),
837 "chunked replay must drop a torn tail beyond a chunk boundary identically to the \
838 whole-file read+fold path"
839 );
840 }
841
842 #[tokio::test]
848 #[serial_test::serial(session_history_integrity)]
849 async fn test_replay_up_to_matches_fold_at_chunk_boundaries() {
850 const N_TURNS: u64 = 250; let dir = tempfile::tempdir().unwrap();
853 let log = seed_large_synthetic_log(dir.path(), N_TURNS).await;
854 let all_events = log.read_all().await.unwrap();
855 assert!(
856 all_events.len() > 200,
857 "synthetic log must exceed 200 events to exercise several REPLAY_CHUNK_SIZE boundaries"
858 );
859
860 for up_to in [99u64, 100, 101, 199, 200, 201] {
861 let expected = ReplayEngine::fold(all_events.clone(), Some(up_to));
862 let actual = ReplayEngine::replay(dir.path(), Some(up_to)).await.unwrap();
863
864 assert_eq!(
865 actual.last_seq, expected.last_seq,
866 "up_to={up_to}: last_seq mismatch between streamed replay and Vec-based fold"
867 );
868 assert_eq!(
869 serde_json::to_string(&actual.messages).unwrap(),
870 serde_json::to_string(&expected.messages).unwrap(),
871 "up_to={up_to}: streamed replay must match Vec-based fold exactly at/near chunk \
872 boundaries"
873 );
874 }
875 }
876}