1use std::io::{self, Read, Write};
42
43use serde::{Deserialize, Serialize};
44
45use crate::run_meta::{ContextSnapshot, RegionEntrySnapshot, RegionSnapshot, RunMeta, RunStatus};
46
47pub const RUN_ARCHIVE_MAGIC: &[u8; 4] = b"LVR1";
49
50pub const RUN_ARCHIVE_VERSION: u16 = 1;
52
53#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
60pub struct RunIdentity {
61 pub run_id: String,
63 pub machine_id: String,
65 pub world_id: String,
67 pub created_at: i64,
69}
70
71#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
73pub struct MessageRecord {
74 pub role: String,
76 pub content: String,
78}
79
80#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
82pub struct ToolCallRecord {
83 pub id: String,
85 pub name: String,
87 pub arguments: String,
89 pub result: Option<String>,
91 #[serde(default, skip_serializing_if = "Option::is_none")]
95 pub thought_signature: Option<String>,
96}
97
98#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
101pub struct InferenceRequestRecord {
102 pub model: String,
104 pub system: Vec<String>,
106 pub messages: Vec<MessageRecord>,
108 pub tool_names: Vec<String>,
110 pub temperature: f32,
112 pub max_tokens: usize,
114}
115
116#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
118pub struct InferenceResponseRecord {
119 pub content: String,
121 pub tool_calls: Vec<ToolCallRecord>,
123 pub prompt_tokens: usize,
125 pub completion_tokens: usize,
127 pub cached_tokens: usize,
129 pub cache_write_tokens: usize,
131}
132
133#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
135pub enum RegionDelta {
136 Set(RegionSnapshot),
139 Append {
142 name: String,
144 entries: Vec<RegionEntrySnapshot>,
146 current_tokens: usize,
148 },
149 Clear {
151 name: String,
153 },
154 Remove {
156 name: String,
158 },
159}
160
161#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
163pub struct ContextDelta {
164 pub stage_name: String,
166 pub total_tokens: usize,
168 pub max_tokens: usize,
170 pub regions: Vec<RegionDelta>,
172}
173
174#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
176pub enum RunRecord {
177 Header {
179 identity: RunIdentity,
181 meta: Box<RunMeta>,
183 },
184 OwnershipChanged {
186 machine_id: String,
188 world_id: String,
190 at: i64,
192 },
193 Inference {
195 stage: String,
197 iteration: usize,
199 request: InferenceRequestRecord,
201 response: InferenceResponseRecord,
203 at: i64,
205 },
206 ToolBatch {
214 calls: Vec<ToolCallRecord>,
216 at: i64,
218 #[serde(default)]
220 stage_index: usize,
221 #[serde(default)]
224 iteration: usize,
225 #[serde(default)]
227 response: String,
228 },
229 ToolCallDone {
231 iteration: usize,
233 call_id: String,
235 result: String,
237 at: i64,
239 },
240 ContextCheckpoint {
242 snapshot: ContextSnapshot,
244 at: i64,
246 },
247 ContextDiff {
249 delta: ContextDelta,
251 at: i64,
253 },
254 Message {
256 message: MessageRecord,
258 at: i64,
260 },
261 StatusChanged {
263 status: RunStatus,
265 at: i64,
267 },
268 Checkpoint {
271 meta: Box<RunMeta>,
273 context: ContextSnapshot,
275 at: i64,
277 },
278 Progress {
283 meta: Box<RunMeta>,
285 delta: ContextDelta,
287 at: i64,
289 },
290}
291
292fn is_prefix(prev: &[RegionEntrySnapshot], next: &[RegionEntrySnapshot]) -> bool {
296 prev.len() <= next.len() && next[..prev.len()] == *prev
297}
298
299pub fn diff_context(prev: &ContextSnapshot, next: &ContextSnapshot) -> ContextDelta {
303 let mut regions = Vec::new();
304 for nr in &next.regions {
305 match prev.regions.iter().find(|r| r.name == nr.name) {
306 None => regions.push(RegionDelta::Set(nr.clone())),
307 Some(pr) => {
308 if pr == nr {
309 } else if nr.entries.is_empty() && !pr.entries.is_empty() {
311 regions.push(RegionDelta::Clear {
312 name: nr.name.clone(),
313 });
314 } else if pr.kind == nr.kind
315 && pr.max_tokens == nr.max_tokens
316 && is_prefix(&pr.entries, &nr.entries)
317 {
318 regions.push(RegionDelta::Append {
319 name: nr.name.clone(),
320 entries: nr.entries[pr.entries.len()..].to_vec(),
321 current_tokens: nr.current_tokens,
322 });
323 } else {
324 regions.push(RegionDelta::Set(nr.clone()));
325 }
326 }
327 }
328 }
329 for pr in &prev.regions {
330 if !next.regions.iter().any(|r| r.name == pr.name) {
331 regions.push(RegionDelta::Remove {
332 name: pr.name.clone(),
333 });
334 }
335 }
336 ContextDelta {
337 stage_name: next.stage_name.clone(),
338 total_tokens: next.total_tokens,
339 max_tokens: next.max_tokens,
340 regions,
341 }
342}
343
344pub fn apply_delta(base: &mut ContextSnapshot, delta: &ContextDelta) {
348 base.stage_name = delta.stage_name.clone();
349 base.total_tokens = delta.total_tokens;
350 base.max_tokens = delta.max_tokens;
351 for region_delta in &delta.regions {
352 match region_delta {
353 RegionDelta::Set(snapshot) => {
354 match base.regions.iter_mut().find(|r| r.name == snapshot.name) {
355 Some(existing) => *existing = snapshot.clone(),
356 None => base.regions.push(snapshot.clone()),
357 }
358 }
359 RegionDelta::Append {
360 name,
361 entries,
362 current_tokens,
363 } => {
364 if let Some(region) = base.regions.iter_mut().find(|r| &r.name == name) {
365 region.entries.extend(entries.iter().cloned());
366 region.current_tokens = *current_tokens;
367 }
368 }
369 RegionDelta::Clear { name } => {
370 if let Some(region) = base.regions.iter_mut().find(|r| &r.name == name) {
371 region.entries.clear();
372 region.current_tokens = 0;
373 }
374 }
375 RegionDelta::Remove { name } => {
376 base.regions.retain(|r| &r.name != name);
377 }
378 }
379 }
380}
381
382pub fn write_archive_start(w: &mut dyn Write, version: u16) -> io::Result<()> {
386 w.write_all(RUN_ARCHIVE_MAGIC)?;
387 w.write_all(&version.to_be_bytes())?;
388 Ok(())
389}
390
391pub fn read_archive_start(r: &mut dyn Read) -> io::Result<u16> {
393 let mut magic = [0u8; 4];
394 r.read_exact(&mut magic)?;
395 if &magic != RUN_ARCHIVE_MAGIC {
396 return Err(io::Error::new(
397 io::ErrorKind::InvalidData,
398 "not a leviath run archive (bad magic)",
399 ));
400 }
401 let mut version = [0u8; 2];
402 r.read_exact(&mut version)?;
403 Ok(u16::from_be_bytes(version))
404}
405
406pub fn write_record(w: &mut dyn Write, record: &RunRecord) -> io::Result<()> {
409 let payload = serde_json::to_vec(record).expect("a RunRecord always serializes to JSON");
410 let len = payload.len() as u64;
411 w.write_all(&len.to_be_bytes())?;
412 w.write_all(&payload)?;
413 Ok(())
414}
415
416fn read_exact_or_eof(r: &mut dyn Read, buf: &mut [u8]) -> io::Result<bool> {
419 let mut filled = 0;
420 while filled < buf.len() {
421 match r.read(&mut buf[filled..])? {
422 0 => {
423 if filled == 0 {
424 return Ok(false); }
426 return Err(io::Error::new(
427 io::ErrorKind::UnexpectedEof,
428 "truncated run-archive frame",
429 ));
430 }
431 n => filled += n,
432 }
433 }
434 Ok(true)
435}
436
437pub fn read_record(r: &mut dyn Read) -> io::Result<Option<RunRecord>> {
439 let mut len_bytes = [0u8; 8];
440 if !read_exact_or_eof(r, &mut len_bytes)? {
441 return Ok(None);
442 }
443 let len = u64::from_be_bytes(len_bytes) as usize;
444 let mut payload = vec![0u8; len];
445 if !read_exact_or_eof(r, &mut payload)? {
446 return Err(io::Error::new(
447 io::ErrorKind::UnexpectedEof,
448 "truncated run-archive frame",
449 ));
450 }
451 let record = serde_json::from_slice(&payload)
452 .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
453 Ok(Some(record))
454}
455
456pub fn read_archive(r: &mut dyn Read) -> io::Result<(u16, Vec<RunRecord>)> {
458 let version = read_archive_start(r)?;
459 let mut records = Vec::new();
460 while let Some(record) = read_record(r)? {
461 records.push(record);
462 }
463 Ok((version, records))
464}
465
466pub fn read_archive_lenient(r: &mut dyn Read) -> io::Result<(u16, Vec<RunRecord>)> {
479 let version = read_archive_start(r)?;
480 let mut records = Vec::new();
481 while let Ok(Some(record)) = read_record(r) {
484 records.push(record);
485 }
486 Ok((version, records))
487}
488
489#[derive(Debug, Clone, PartialEq)]
496pub struct PendingToolBatch {
497 pub stage_index: usize,
499 pub iteration: usize,
501 pub response: String,
503 pub calls: Vec<ToolCallRecord>,
505}
506
507#[derive(Debug, Clone, PartialEq)]
510pub struct FoldedRun {
511 pub identity: RunIdentity,
513 pub meta: RunMeta,
515 pub context: ContextSnapshot,
517 pub messages: Vec<MessageRecord>,
519 pub inference_count: usize,
521 pub tool_call_count: usize,
523 pub pending_batch: Option<PendingToolBatch>,
527}
528
529pub fn context_contains_batch(context: &ContextSnapshot, batch: &PendingToolBatch) -> bool {
533 let Some(first_id) = batch.calls.first().map(|c| c.id.as_str()) else {
534 return false;
535 };
536 context.regions.iter().any(|region| {
537 region.entries.iter().any(|entry| {
538 matches!(
539 &entry.kind,
540 crate::region::EntryKind::AssistantTurn { tool_calls }
541 if tool_calls.iter().any(|tc| tc.id == first_id)
542 )
543 })
544 })
545}
546
547pub fn fold(records: &[RunRecord]) -> Option<FoldedRun> {
550 let mut iter = records.iter();
551 let (identity, meta) = match iter.next() {
552 Some(RunRecord::Header { identity, meta }) => (identity.clone(), (**meta).clone()),
553 _ => return None,
554 };
555 let mut folded = FoldedRun {
556 identity,
557 meta,
558 context: ContextSnapshot {
559 stage_name: String::new(),
560 total_tokens: 0,
561 max_tokens: 0,
562 regions: Vec::new(),
563 },
564 messages: Vec::new(),
565 inference_count: 0,
566 tool_call_count: 0,
567 pending_batch: None,
568 };
569 for record in iter {
570 match record {
571 RunRecord::Header { identity, meta } => {
572 folded.identity = identity.clone();
573 folded.meta = (**meta).clone();
574 }
575 RunRecord::OwnershipChanged {
576 machine_id,
577 world_id,
578 ..
579 } => {
580 folded.identity.machine_id = machine_id.clone();
581 folded.identity.world_id = world_id.clone();
582 }
583 RunRecord::Inference { .. } => folded.inference_count += 1,
584 RunRecord::ToolBatch {
585 calls,
586 stage_index,
587 iteration,
588 response,
589 ..
590 } => {
591 folded.tool_call_count += calls.len();
592 folded.pending_batch = Some(PendingToolBatch {
595 stage_index: *stage_index,
596 iteration: *iteration,
597 response: response.clone(),
598 calls: calls.clone(),
599 });
600 }
601 RunRecord::ToolCallDone {
602 iteration,
603 call_id,
604 result,
605 ..
606 } => {
607 if let Some(batch) = folded
610 .pending_batch
611 .as_mut()
612 .filter(|b| b.iteration == *iteration)
613 && let Some(call) = batch.calls.iter_mut().find(|c| c.id == *call_id)
614 {
615 call.result = Some(result.clone());
616 }
617 }
618 RunRecord::ContextCheckpoint { snapshot, .. } => folded.context = snapshot.clone(),
619 RunRecord::ContextDiff { delta, .. } => apply_delta(&mut folded.context, delta),
620 RunRecord::Message { message, .. } => folded.messages.push(message.clone()),
621 RunRecord::StatusChanged { status, .. } => folded.meta.status = status.clone(),
622 RunRecord::Checkpoint { meta, context, .. } => {
623 folded.meta = (**meta).clone();
624 folded.context = context.clone();
625 }
626 RunRecord::Progress { meta, delta, .. } => {
627 folded.meta = (**meta).clone();
628 apply_delta(&mut folded.context, delta);
629 }
630 }
631 }
632 if let Some(batch) = &folded.pending_batch
637 && (folded.meta.iteration != batch.iteration
638 || context_contains_batch(&folded.context, batch))
639 {
640 folded.pending_batch = None;
641 }
642 Some(folded)
643}
644
645#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
648pub struct RunPoint {
649 pub meta: RunMeta,
651 pub context: ContextSnapshot,
653 pub at: i64,
655}
656
657pub fn replay_points(records: &[RunRecord]) -> Vec<RunPoint> {
663 let mut iter = records.iter();
664 let mut meta = match iter.next() {
665 Some(RunRecord::Header { meta, .. }) => (**meta).clone(),
666 _ => return Vec::new(),
667 };
668 let mut context = ContextSnapshot {
669 stage_name: String::new(),
670 total_tokens: 0,
671 max_tokens: 0,
672 regions: Vec::new(),
673 };
674 let mut points = Vec::new();
675 for record in iter {
676 match record {
677 RunRecord::Header { meta: m, .. } => meta = (**m).clone(),
678 RunRecord::StatusChanged { status, .. } => meta.status = status.clone(),
679 RunRecord::ContextCheckpoint { snapshot, at } => {
680 context = snapshot.clone();
681 points.push(RunPoint {
682 meta: meta.clone(),
683 context: context.clone(),
684 at: *at,
685 });
686 }
687 RunRecord::ContextDiff { delta, at } => {
688 apply_delta(&mut context, delta);
689 points.push(RunPoint {
690 meta: meta.clone(),
691 context: context.clone(),
692 at: *at,
693 });
694 }
695 RunRecord::Checkpoint {
696 meta: m,
697 context: c,
698 at,
699 } => {
700 meta = (**m).clone();
701 context = c.clone();
702 points.push(RunPoint {
703 meta: meta.clone(),
704 context: context.clone(),
705 at: *at,
706 });
707 }
708 RunRecord::Progress { meta: m, delta, at } => {
709 meta = (**m).clone();
710 apply_delta(&mut context, delta);
711 points.push(RunPoint {
712 meta: meta.clone(),
713 context: context.clone(),
714 at: *at,
715 });
716 }
717 RunRecord::OwnershipChanged { .. }
719 | RunRecord::Inference { .. }
720 | RunRecord::ToolBatch { .. }
721 | RunRecord::ToolCallDone { .. }
722 | RunRecord::Message { .. } => {}
723 }
724 }
725 points
726}
727
728#[cfg(test)]
729mod tests {
730 use super::*;
731 use crate::run_meta::RunStatus;
732
733 fn identity() -> RunIdentity {
734 RunIdentity {
735 run_id: "run-1".to_string(),
736 machine_id: "machine-a".to_string(),
737 world_id: "world-x".to_string(),
738 created_at: 100,
739 }
740 }
741
742 fn meta() -> RunMeta {
743 RunMeta::new(
744 "run-1".to_string(),
745 "coder".to_string(),
746 "/agents/coder".to_string(),
747 "do it".to_string(),
748 Some("anthropic/claude".to_string()),
749 "/work".to_string(),
750 2,
751 )
752 }
753
754 fn entry(content: &str, tokens: usize) -> RegionEntrySnapshot {
755 RegionEntrySnapshot {
756 content: content.to_string(),
757 tokens,
758 kind: crate::region::EntryKind::Text,
759 metadata: None,
760 key: None,
761 taint: Default::default(),
762 }
763 }
764
765 fn region(name: &str, entries: Vec<RegionEntrySnapshot>) -> RegionSnapshot {
766 let current = entries.iter().map(|e| e.tokens).sum();
767 RegionSnapshot {
768 name: name.to_string(),
769 kind: "clearable".to_string(),
770 current_tokens: current,
771 max_tokens: 1000,
772 entries,
773 }
774 }
775
776 fn snapshot(stage: &str, regions: Vec<RegionSnapshot>) -> ContextSnapshot {
777 let total = regions.iter().map(|r| r.current_tokens).sum();
778 ContextSnapshot {
779 stage_name: stage.to_string(),
780 total_tokens: total,
781 max_tokens: 10_000,
782 regions,
783 }
784 }
785
786 fn header() -> RunRecord {
787 RunRecord::Header {
788 identity: identity(),
789 meta: Box::new(meta()),
790 }
791 }
792
793 fn region_delta_kind(d: &RegionDelta) -> &'static str {
797 match d {
798 RegionDelta::Set(_) => "set",
799 RegionDelta::Append { .. } => "append",
800 RegionDelta::Clear { .. } => "clear",
801 RegionDelta::Remove { .. } => "remove",
802 }
803 }
804
805 fn assert_diff_roundtrip(a: &ContextSnapshot, b: &ContextSnapshot) {
810 let delta = diff_context(a, b);
811 let mut base = a.clone();
812 apply_delta(&mut base, &delta);
813 assert_eq!(&base, b);
814 }
815
816 #[test]
817 fn diff_append_only_growth_is_compact() {
818 let a = snapshot("s1", vec![region("conv", vec![entry("hi", 1)])]);
819 let b = snapshot(
820 "s1",
821 vec![region("conv", vec![entry("hi", 1), entry("there", 2)])],
822 );
823 let delta = diff_context(&a, &b);
824 assert_eq!(region_delta_kind(&delta.regions[0]), "append");
825 assert_diff_roundtrip(&a, &b);
826 }
827
828 #[test]
829 fn diff_new_region_is_set() {
830 let a = snapshot("s1", vec![region("conv", vec![entry("hi", 1)])]);
831 let b = snapshot(
832 "s1",
833 vec![
834 region("conv", vec![entry("hi", 1)]),
835 region("plan", vec![entry("p", 3)]),
836 ],
837 );
838 let delta = diff_context(&a, &b);
839 assert!(delta.regions.iter().any(|d| region_delta_kind(d) == "set"));
840 assert_diff_roundtrip(&a, &b);
841 }
842
843 #[test]
844 fn diff_cleared_region() {
845 let a = snapshot("s1", vec![region("conv", vec![entry("hi", 1)])]);
846 let b = snapshot("s1", vec![region("conv", vec![])]);
847 let delta = diff_context(&a, &b);
848 assert_eq!(region_delta_kind(&delta.regions[0]), "clear");
849 assert_diff_roundtrip(&a, &b);
850 }
851
852 #[test]
853 fn diff_removed_region() {
854 let a = snapshot(
855 "s1",
856 vec![
857 region("conv", vec![entry("hi", 1)]),
858 region("plan", vec![entry("p", 3)]),
859 ],
860 );
861 let b = snapshot("s1", vec![region("conv", vec![entry("hi", 1)])]);
862 let delta = diff_context(&a, &b);
863 assert!(
864 delta
865 .regions
866 .iter()
867 .any(|d| region_delta_kind(d) == "remove")
868 );
869 assert_diff_roundtrip(&a, &b);
870 }
871
872 #[test]
873 fn diff_non_prefix_rewrite_is_set() {
874 let a = snapshot("s1", vec![region("conv", vec![entry("old", 1)])]);
876 let b = snapshot("s1", vec![region("conv", vec![entry("new", 1)])]);
877 let delta = diff_context(&a, &b);
878 assert_eq!(region_delta_kind(&delta.regions[0]), "set");
879 assert_diff_roundtrip(&a, &b);
880 }
881
882 #[test]
883 fn diff_kind_change_is_set_not_append() {
884 let a = snapshot("s1", vec![region("conv", vec![entry("hi", 1)])]);
886 let mut grown = region("conv", vec![entry("hi", 1), entry("more", 1)]);
887 grown.kind = "sliding".to_string();
888 let b = snapshot("s1", vec![grown]);
889 let delta = diff_context(&a, &b);
890 assert_eq!(region_delta_kind(&delta.regions[0]), "set");
891 assert_diff_roundtrip(&a, &b);
892 }
893
894 #[test]
895 fn diff_unchanged_region_emits_nothing() {
896 let a = snapshot("s1", vec![region("conv", vec![entry("hi", 1)])]);
897 let b = a.clone();
898 let delta = diff_context(&a, &b);
899 assert!(delta.regions.is_empty());
900 assert_diff_roundtrip(&a, &b);
901 }
902
903 #[test]
904 fn apply_delta_skips_unknown_regions_leniently() {
905 let mut base = snapshot("s1", vec![]);
907 let delta = ContextDelta {
908 stage_name: "s1".to_string(),
909 total_tokens: 0,
910 max_tokens: 10_000,
911 regions: vec![
912 RegionDelta::Append {
913 name: "ghost".to_string(),
914 entries: vec![entry("x", 1)],
915 current_tokens: 1,
916 },
917 RegionDelta::Clear {
918 name: "ghost".to_string(),
919 },
920 RegionDelta::Remove {
921 name: "ghost".to_string(),
922 },
923 ],
924 };
925 apply_delta(&mut base, &delta);
926 assert!(base.regions.is_empty());
927 }
928
929 fn all_record_kinds() -> Vec<RunRecord> {
932 vec![
933 header(),
934 RunRecord::OwnershipChanged {
935 machine_id: "machine-b".to_string(),
936 world_id: "world-y".to_string(),
937 at: 101,
938 },
939 RunRecord::Inference {
940 stage: "plan".to_string(),
941 iteration: 0,
942 request: InferenceRequestRecord {
943 model: "m".to_string(),
944 system: vec!["sys".to_string()],
945 messages: vec![MessageRecord {
946 role: "user".to_string(),
947 content: "hi".to_string(),
948 }],
949 tool_names: vec!["read_file".to_string()],
950 temperature: 0.7,
951 max_tokens: 1024,
952 },
953 response: InferenceResponseRecord {
954 content: "ok".to_string(),
955 tool_calls: vec![],
956 prompt_tokens: 10,
957 completion_tokens: 5,
958 cached_tokens: 0,
959 cache_write_tokens: 0,
960 },
961 at: 102,
962 },
963 RunRecord::ToolBatch {
964 calls: vec![ToolCallRecord {
965 id: "c1".to_string(),
966 name: "read_file".to_string(),
967 arguments: "{}".to_string(),
968 result: Some("body".to_string()),
969 thought_signature: Some("sig".to_string()),
970 }],
971 at: 103,
972 stage_index: 0,
973 iteration: 0,
974 response: "reading".to_string(),
975 },
976 RunRecord::ToolCallDone {
977 iteration: 0,
978 call_id: "c1".to_string(),
979 result: "body".to_string(),
980 at: 103,
981 },
982 RunRecord::ContextCheckpoint {
983 snapshot: snapshot("plan", vec![region("conv", vec![entry("hi", 1)])]),
984 at: 104,
985 },
986 RunRecord::ContextDiff {
987 delta: ContextDelta {
988 stage_name: "plan".to_string(),
989 total_tokens: 3,
990 max_tokens: 10_000,
991 regions: vec![RegionDelta::Append {
992 name: "conv".to_string(),
993 entries: vec![entry("more", 2)],
994 current_tokens: 3,
995 }],
996 },
997 at: 105,
998 },
999 RunRecord::Message {
1000 message: MessageRecord {
1001 role: "user".to_string(),
1002 content: "another".to_string(),
1003 },
1004 at: 106,
1005 },
1006 RunRecord::StatusChanged {
1007 status: RunStatus::Complete,
1008 at: 107,
1009 },
1010 RunRecord::Checkpoint {
1011 meta: Box::new(meta()),
1012 context: snapshot("plan", vec![region("conv", vec![entry("hi", 1)])]),
1013 at: 108,
1014 },
1015 RunRecord::Progress {
1016 meta: Box::new(meta()),
1017 delta: ContextDelta {
1018 stage_name: "plan".to_string(),
1019 total_tokens: 3,
1020 max_tokens: 10_000,
1021 regions: vec![RegionDelta::Append {
1022 name: "conv".to_string(),
1023 entries: vec![entry("step", 2)],
1024 current_tokens: 3,
1025 }],
1026 },
1027 at: 109,
1028 },
1029 ]
1030 }
1031
1032 #[test]
1033 fn archive_write_then_read_roundtrips_every_record_kind() {
1034 let records = all_record_kinds();
1035 let mut buf = Vec::new();
1036 write_archive_start(&mut buf, RUN_ARCHIVE_VERSION).unwrap();
1037 for r in &records {
1038 write_record(&mut buf, r).unwrap();
1039 }
1040 let (version, read) = read_archive(&mut buf.as_slice()).unwrap();
1041 assert_eq!(version, RUN_ARCHIVE_VERSION);
1042 assert_eq!(read, records);
1043 }
1044
1045 #[test]
1046 fn read_archive_start_rejects_bad_magic() {
1047 let mut bytes: &[u8] = b"XXXX\x00\x01";
1048 let err = read_archive_start(&mut bytes).unwrap_err();
1049 assert_eq!(err.kind(), io::ErrorKind::InvalidData);
1050 }
1051
1052 #[test]
1053 fn read_archive_start_reports_version() {
1054 let mut buf = Vec::new();
1055 write_archive_start(&mut buf, 7).unwrap();
1056 assert_eq!(read_archive_start(&mut buf.as_slice()).unwrap(), 7);
1057 }
1058
1059 #[test]
1060 fn read_record_returns_none_at_clean_eof() {
1061 let empty: &[u8] = &[];
1062 assert!(read_record(&mut { empty }).unwrap().is_none());
1063 }
1064
1065 #[test]
1066 fn read_record_errors_on_truncated_length_prefix() {
1067 let mut bytes: &[u8] = &[0, 0];
1069 let err = read_record(&mut bytes).unwrap_err();
1070 assert_eq!(err.kind(), io::ErrorKind::UnexpectedEof);
1071 }
1072
1073 #[test]
1074 fn read_record_errors_on_truncated_payload() {
1075 let mut bytes: &[u8] = &[0, 0, 0, 0, 0, 0, 0, 10, 1, 2];
1077 let err = read_record(&mut bytes).unwrap_err();
1078 assert_eq!(err.kind(), io::ErrorKind::UnexpectedEof);
1079 }
1080
1081 #[test]
1082 fn read_record_errors_on_empty_payload_at_boundary() {
1083 let mut bytes: &[u8] = &[0, 0, 0, 0, 0, 0, 0, 10];
1086 let err = read_record(&mut bytes).unwrap_err();
1087 assert_eq!(err.kind(), io::ErrorKind::UnexpectedEof);
1088 }
1089
1090 #[test]
1091 fn read_record_errors_on_invalid_json_payload() {
1092 let mut buf = Vec::new();
1094 let bad = b"not json";
1095 buf.extend_from_slice(&(bad.len() as u64).to_be_bytes());
1096 buf.extend_from_slice(bad);
1097 let err = read_record(&mut buf.as_slice()).unwrap_err();
1098 assert_eq!(err.kind(), io::ErrorKind::InvalidData);
1099 }
1100
1101 struct FailingReader;
1104 impl Read for FailingReader {
1105 fn read(&mut self, _buf: &mut [u8]) -> io::Result<usize> {
1106 Err(io::Error::other("device error"))
1107 }
1108 }
1109
1110 #[test]
1111 fn read_record_propagates_reader_errors() {
1112 let err = read_record(&mut FailingReader).unwrap_err();
1113 assert_eq!(err.kind(), io::ErrorKind::Other);
1114 }
1115
1116 #[test]
1117 fn read_archive_propagates_a_bad_preamble() {
1118 let mut bytes: &[u8] = b"LV";
1120 assert!(read_archive(&mut bytes).is_err());
1121 }
1122
1123 #[test]
1124 fn read_archive_propagates_a_bad_frame() {
1125 let mut buf = Vec::new();
1127 write_archive_start(&mut buf, RUN_ARCHIVE_VERSION).unwrap();
1128 buf.extend_from_slice(&[0, 0, 0, 0, 0, 0, 0, 5, 1, 2]); let err = read_archive(&mut buf.as_slice()).unwrap_err();
1130 assert_eq!(err.kind(), io::ErrorKind::UnexpectedEof);
1131 }
1132
1133 struct FailAfter {
1135 remaining: usize,
1136 }
1137 impl Write for FailAfter {
1138 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
1139 if self.remaining == 0 {
1140 return Err(io::Error::other("disk full"));
1141 }
1142 let n = buf.len().min(self.remaining);
1143 self.remaining -= n;
1144 Ok(n)
1145 }
1146 fn flush(&mut self) -> io::Result<()> {
1147 Ok(())
1148 }
1149 }
1150
1151 #[test]
1152 fn fail_after_writer_flush_is_a_noop() {
1153 assert!(FailAfter { remaining: 1 }.flush().is_ok());
1154 }
1155
1156 #[test]
1157 fn write_archive_start_propagates_write_errors() {
1158 assert!(write_archive_start(&mut FailAfter { remaining: 0 }, 1).is_err());
1160 assert!(write_archive_start(&mut FailAfter { remaining: 4 }, 1).is_err());
1161 }
1162
1163 #[test]
1164 fn write_record_propagates_write_errors() {
1165 let rec = header();
1166 assert!(write_record(&mut FailAfter { remaining: 0 }, &rec).is_err());
1168 assert!(write_record(&mut FailAfter { remaining: 8 }, &rec).is_err());
1169 }
1170
1171 #[test]
1172 fn read_archive_lenient_matches_strict_on_a_clean_archive() {
1173 let records = all_record_kinds();
1176 let mut buf = Vec::new();
1177 write_archive_start(&mut buf, RUN_ARCHIVE_VERSION).unwrap();
1178 for r in &records {
1179 write_record(&mut buf, r).unwrap();
1180 }
1181 let (version, read) = read_archive_lenient(&mut buf.as_slice()).unwrap();
1182 assert_eq!(version, RUN_ARCHIVE_VERSION);
1183 assert_eq!(read, records);
1184 }
1185
1186 #[test]
1187 fn read_archive_lenient_keeps_valid_prefix_before_a_torn_tail() {
1188 let mut buf = Vec::new();
1192 write_archive_start(&mut buf, RUN_ARCHIVE_VERSION).unwrap();
1193 write_record(&mut buf, &header()).unwrap();
1194 write_record(
1195 &mut buf,
1196 &RunRecord::ContextCheckpoint {
1197 snapshot: snapshot("plan", vec![region("conv", vec![entry("hi", 1)])]),
1198 at: 1,
1199 },
1200 )
1201 .unwrap();
1202 buf.extend_from_slice(&[0, 0, 0, 0, 0, 0, 0, 10, 1, 2]);
1204
1205 assert!(read_archive(&mut buf.as_slice()).is_err());
1207 let (version, records) = read_archive_lenient(&mut buf.as_slice()).unwrap();
1209 assert_eq!(version, RUN_ARCHIVE_VERSION);
1210 assert_eq!(records.len(), 2);
1211 let folded = fold(&records).expect("prefix starts with a Header");
1212 assert_eq!(folded.context.regions[0].entries.len(), 1);
1213 }
1214
1215 #[test]
1216 fn read_archive_lenient_still_errors_on_a_bad_preamble() {
1217 let mut bad_magic: &[u8] = b"XXXX\x00\x01";
1220 assert!(read_archive_lenient(&mut bad_magic).is_err());
1221 let mut short: &[u8] = b"LVR1";
1223 assert!(read_archive_lenient(&mut short).is_err());
1224 }
1225
1226 #[test]
1227 fn read_archive_start_errors_on_truncated_version() {
1228 let mut bytes: &[u8] = b"LVR1";
1230 let err = read_archive_start(&mut bytes).unwrap_err();
1231 assert_eq!(err.kind(), io::ErrorKind::UnexpectedEof);
1232 }
1233
1234 #[test]
1237 fn fold_requires_a_header_first() {
1238 assert!(fold(&[]).is_none());
1239 assert!(
1240 fold(&[RunRecord::StatusChanged {
1241 status: RunStatus::Complete,
1242 at: 1
1243 }])
1244 .is_none()
1245 );
1246 }
1247
1248 #[test]
1249 fn fold_reconstructs_state_from_the_journal() {
1250 let records = all_record_kinds();
1251 let folded = fold(&records).expect("has header");
1252 assert_eq!(folded.identity.machine_id, "machine-b");
1254 assert_eq!(folded.identity.world_id, "world-y");
1255 assert_eq!(folded.inference_count, 1);
1257 assert_eq!(folded.tool_call_count, 1);
1258 assert_eq!(folded.messages.len(), 1);
1260 assert_eq!(folded.messages[0].content, "another");
1261 assert_eq!(folded.context.regions[0].name, "conv");
1264 assert_eq!(folded.context.regions[0].entries.len(), 2);
1265 assert_eq!(folded.context.total_tokens, 3);
1266 assert_eq!(folded.meta.run_id, "run-1");
1267 let pending = folded.pending_batch.expect("batch never applied");
1270 assert_eq!(pending.calls[0].result.as_deref(), Some("body"));
1271 }
1272
1273 #[test]
1274 fn fold_applies_context_diffs_over_a_checkpoint() {
1275 let records = vec![
1277 header(),
1278 RunRecord::ContextCheckpoint {
1279 snapshot: snapshot("plan", vec![region("conv", vec![entry("hi", 1)])]),
1280 at: 1,
1281 },
1282 RunRecord::ContextDiff {
1283 delta: ContextDelta {
1284 stage_name: "plan".to_string(),
1285 total_tokens: 3,
1286 max_tokens: 10_000,
1287 regions: vec![RegionDelta::Append {
1288 name: "conv".to_string(),
1289 entries: vec![entry("there", 2)],
1290 current_tokens: 3,
1291 }],
1292 },
1293 at: 2,
1294 },
1295 ];
1296 let folded = fold(&records).unwrap();
1297 assert_eq!(folded.context.regions[0].entries.len(), 2);
1298 assert_eq!(folded.context.total_tokens, 3);
1299 }
1300
1301 #[test]
1302 fn fold_later_header_updates_identity_and_meta() {
1303 let mut second_meta = meta();
1305 second_meta.status = RunStatus::Running;
1306 let records = vec![
1307 header(),
1308 RunRecord::Header {
1309 identity: RunIdentity {
1310 run_id: "run-1".to_string(),
1311 machine_id: "machine-c".to_string(),
1312 world_id: "world-z".to_string(),
1313 created_at: 200,
1314 },
1315 meta: Box::new(second_meta),
1316 },
1317 ];
1318 let folded = fold(&records).unwrap();
1319 assert_eq!(folded.identity.machine_id, "machine-c");
1320 assert_eq!(folded.meta.status, RunStatus::Running);
1321 }
1322
1323 #[test]
1324 fn fold_progress_applies_meta_and_context_diff() {
1325 let mut advanced = meta();
1326 advanced.status = RunStatus::Running;
1327 advanced.iteration = 5;
1328 let records = vec![
1329 header(),
1330 RunRecord::ContextCheckpoint {
1331 snapshot: snapshot("plan", vec![region("conv", vec![entry("hi", 1)])]),
1332 at: 1,
1333 },
1334 RunRecord::Progress {
1335 meta: Box::new(advanced),
1336 delta: ContextDelta {
1337 stage_name: "plan".to_string(),
1338 total_tokens: 3,
1339 max_tokens: 10_000,
1340 regions: vec![RegionDelta::Append {
1341 name: "conv".to_string(),
1342 entries: vec![entry("there", 2)],
1343 current_tokens: 3,
1344 }],
1345 },
1346 at: 2,
1347 },
1348 ];
1349 let folded = fold(&records).unwrap();
1350 assert_eq!(folded.meta.iteration, 5);
1351 assert_eq!(folded.meta.status, RunStatus::Running);
1352 assert_eq!(folded.context.regions[0].entries.len(), 2);
1353 }
1354
1355 fn call(id: &str, result: Option<&str>) -> ToolCallRecord {
1358 ToolCallRecord {
1359 id: id.to_string(),
1360 name: "shell".to_string(),
1361 arguments: "{}".to_string(),
1362 result: result.map(str::to_string),
1363 thought_signature: None,
1364 }
1365 }
1366
1367 fn batch(iteration: usize, calls: Vec<ToolCallRecord>) -> RunRecord {
1368 RunRecord::ToolBatch {
1369 calls,
1370 at: 10,
1371 stage_index: 0,
1372 iteration,
1373 response: "running tools".to_string(),
1374 }
1375 }
1376
1377 fn turn_entry(call_ids: &[&str]) -> RegionEntrySnapshot {
1379 let mut e = entry("turn", 1);
1380 e.kind = crate::region::EntryKind::AssistantTurn {
1381 tool_calls: call_ids
1382 .iter()
1383 .map(|id| crate::region::SerializedToolCall {
1384 id: id.to_string(),
1385 name: "shell".to_string(),
1386 arguments: serde_json::Value::Null,
1387 thought_signature: None,
1388 })
1389 .collect(),
1390 };
1391 e
1392 }
1393
1394 #[test]
1395 fn fold_surfaces_a_pending_batch_with_merged_results() {
1396 let records = vec![
1401 header(),
1402 batch(
1403 0,
1404 vec![
1405 call("c1", None),
1406 call("c2", Some("inline")),
1407 call("c3", None),
1408 ],
1409 ),
1410 RunRecord::ToolCallDone {
1411 iteration: 0,
1412 call_id: "c1".to_string(),
1413 result: "ran".to_string(),
1414 at: 11,
1415 },
1416 ];
1417 let folded = fold(&records).unwrap();
1418 let pending = folded.pending_batch.expect("batch is pending");
1419 assert_eq!(pending.iteration, 0);
1420 assert_eq!(pending.response, "running tools");
1421 assert_eq!(pending.calls[0].result.as_deref(), Some("ran"));
1422 assert_eq!(pending.calls[1].result.as_deref(), Some("inline"));
1423 assert_eq!(pending.calls[2].result, None);
1424 assert_eq!(folded.tool_call_count, 3);
1425 }
1426
1427 #[test]
1428 fn fold_keeps_only_the_latest_batch_and_ignores_stale_done_records() {
1429 let mut advanced = meta();
1432 advanced.iteration = 1;
1433 let records = vec![
1434 header(),
1435 batch(0, vec![call("c1", None)]),
1436 RunRecord::Progress {
1437 meta: Box::new(advanced),
1438 delta: ContextDelta {
1439 stage_name: "plan".to_string(),
1440 total_tokens: 0,
1441 max_tokens: 10_000,
1442 regions: vec![],
1443 },
1444 at: 11,
1445 },
1446 batch(1, vec![call("c2", None)]),
1447 RunRecord::ToolCallDone {
1448 iteration: 0,
1449 call_id: "c1".to_string(),
1450 result: "stale".to_string(),
1451 at: 12,
1452 },
1453 RunRecord::ToolCallDone {
1454 iteration: 1,
1455 call_id: "unknown".to_string(),
1456 result: "nowhere to land".to_string(),
1457 at: 13,
1458 },
1459 ];
1460 let folded = fold(&records).unwrap();
1461 let pending = folded.pending_batch.expect("latest batch is pending");
1462 assert_eq!(pending.iteration, 1);
1463 assert_eq!(pending.calls.len(), 1);
1464 assert_eq!(pending.calls[0].id, "c2");
1465 assert_eq!(pending.calls[0].result, None, "stale/unknown dones ignored");
1466 }
1467
1468 #[test]
1469 fn fold_clears_a_batch_once_the_iteration_moves_on() {
1470 let mut advanced = meta();
1473 advanced.iteration = 1;
1474 let records = vec![
1475 header(),
1476 batch(0, vec![call("c1", Some("done"))]),
1477 RunRecord::Progress {
1478 meta: Box::new(advanced),
1479 delta: ContextDelta {
1480 stage_name: "plan".to_string(),
1481 total_tokens: 0,
1482 max_tokens: 10_000,
1483 regions: vec![],
1484 },
1485 at: 11,
1486 },
1487 ];
1488 assert_eq!(fold(&records).unwrap().pending_batch, None);
1489 }
1490
1491 #[test]
1492 fn fold_clears_a_batch_whose_turn_already_landed_in_the_window() {
1493 let records = vec![
1496 header(),
1497 batch(0, vec![call("c1", Some("done"))]),
1498 RunRecord::ContextCheckpoint {
1499 snapshot: snapshot("plan", vec![region("conv", vec![turn_entry(&["c1"])])]),
1500 at: 11,
1501 },
1502 ];
1503 assert_eq!(fold(&records).unwrap().pending_batch, None);
1504 }
1505
1506 #[test]
1507 fn context_contains_batch_matches_only_the_batch_turn() {
1508 let pending = PendingToolBatch {
1509 stage_index: 0,
1510 iteration: 0,
1511 response: String::new(),
1512 calls: vec![call("c1", None)],
1513 };
1514 let other = snapshot("plan", vec![region("conv", vec![turn_entry(&["zz"])])]);
1516 assert!(!context_contains_batch(&other, &pending));
1517 let own = snapshot(
1519 "plan",
1520 vec![region("conv", vec![turn_entry(&["c1", "c2"])])],
1521 );
1522 assert!(context_contains_batch(&own, &pending));
1523 let empty = PendingToolBatch {
1525 calls: vec![],
1526 ..pending
1527 };
1528 assert!(!context_contains_batch(&own, &empty));
1529 }
1530
1531 #[test]
1532 fn old_shape_tool_batch_json_still_parses() {
1533 let json = br#"{"ToolBatch":{"calls":[{"id":"c1","name":"shell","arguments":"{}","result":"ok"}],"at":9}}"#;
1537 let mut buf = Vec::new();
1538 buf.extend_from_slice(&(json.len() as u64).to_be_bytes());
1539 buf.extend_from_slice(json);
1540 let record = read_record(&mut buf.as_slice()).unwrap().unwrap();
1541 assert_eq!(
1542 record,
1543 RunRecord::ToolBatch {
1544 calls: vec![call("c1", Some("ok"))],
1545 at: 9,
1546 stage_index: 0,
1547 iteration: 0,
1548 response: String::new(),
1549 }
1550 );
1551 }
1552
1553 #[test]
1556 fn replay_points_requires_a_header() {
1557 assert!(replay_points(&[]).is_empty());
1558 assert!(
1559 replay_points(&[RunRecord::Message {
1560 message: MessageRecord {
1561 role: "user".to_string(),
1562 content: "x".to_string(),
1563 },
1564 at: 1,
1565 }])
1566 .is_empty()
1567 );
1568 }
1569
1570 #[test]
1571 fn replay_points_emits_a_snapshot_per_context_change() {
1572 let mut running = meta();
1575 running.status = RunStatus::Running;
1576 let records = vec![
1577 header(),
1578 RunRecord::Inference {
1579 stage: "plan".to_string(),
1580 iteration: 0,
1581 request: InferenceRequestRecord {
1582 model: "m".to_string(),
1583 system: vec![],
1584 messages: vec![],
1585 tool_names: vec![],
1586 temperature: 0.7,
1587 max_tokens: 10,
1588 },
1589 response: InferenceResponseRecord {
1590 content: "ok".to_string(),
1591 tool_calls: vec![],
1592 prompt_tokens: 1,
1593 completion_tokens: 1,
1594 cached_tokens: 0,
1595 cache_write_tokens: 0,
1596 },
1597 at: 1,
1598 },
1599 batch(0, vec![call("c1", None)]),
1600 RunRecord::ToolCallDone {
1601 iteration: 0,
1602 call_id: "c1".to_string(),
1603 result: "ran".to_string(),
1604 at: 1,
1605 },
1606 RunRecord::ContextCheckpoint {
1607 snapshot: snapshot("plan", vec![region("conv", vec![entry("hi", 1)])]),
1608 at: 2,
1609 },
1610 RunRecord::StatusChanged {
1611 status: RunStatus::Running,
1612 at: 3,
1613 },
1614 RunRecord::Progress {
1615 meta: Box::new(running),
1616 delta: ContextDelta {
1617 stage_name: "implement".to_string(),
1618 total_tokens: 3,
1619 max_tokens: 10_000,
1620 regions: vec![RegionDelta::Append {
1621 name: "conv".to_string(),
1622 entries: vec![entry("more", 2)],
1623 current_tokens: 3,
1624 }],
1625 },
1626 at: 4,
1627 },
1628 ];
1629 let points = replay_points(&records);
1630 assert_eq!(points.len(), 2, "one point per context change");
1631 assert_eq!(points[0].at, 2);
1633 assert_eq!(points[0].context.regions[0].entries.len(), 1);
1634 assert_eq!(points[1].at, 4);
1637 assert_eq!(points[1].context.regions[0].entries.len(), 2);
1638 assert_eq!(points[1].context.stage_name, "implement");
1639 assert_eq!(points[1].meta.status, RunStatus::Running);
1640 }
1641
1642 #[test]
1643 fn replay_points_handles_context_diff_and_a_later_header() {
1644 let mut relabeled = meta();
1647 relabeled.agent_name = "renamed".to_string();
1648 let records = vec![
1649 header(),
1650 RunRecord::ContextCheckpoint {
1651 snapshot: snapshot("plan", vec![region("conv", vec![entry("hi", 1)])]),
1652 at: 1,
1653 },
1654 RunRecord::Header {
1655 identity: identity(),
1656 meta: Box::new(relabeled),
1657 },
1658 RunRecord::ContextDiff {
1659 delta: ContextDelta {
1660 stage_name: "plan".to_string(),
1661 total_tokens: 3,
1662 max_tokens: 10_000,
1663 regions: vec![RegionDelta::Append {
1664 name: "conv".to_string(),
1665 entries: vec![entry("more", 2)],
1666 current_tokens: 3,
1667 }],
1668 },
1669 at: 2,
1670 },
1671 ];
1672 let points = replay_points(&records);
1673 assert_eq!(points.len(), 2); assert_eq!(points[1].context.regions[0].entries.len(), 2);
1675 assert_eq!(points[1].meta.agent_name, "renamed");
1677 }
1678
1679 #[test]
1680 fn replay_points_over_a_full_checkpoint() {
1681 let records = vec![
1683 header(),
1684 RunRecord::Checkpoint {
1685 meta: Box::new(meta()),
1686 context: snapshot("review", vec![region("conv", vec![entry("x", 4)])]),
1687 at: 9,
1688 },
1689 ];
1690 let points = replay_points(&records);
1691 assert_eq!(points.len(), 1);
1692 assert_eq!(points[0].context.stage_name, "review");
1693 assert_eq!(points[0].context.regions[0].entries[0].tokens, 4);
1694 }
1695}