1use std::io::{self, Read, Write};
42use std::ops::ControlFlow;
43
44use serde::{Deserialize, Serialize};
45
46use crate::run_meta::{ContextSnapshot, RegionEntrySnapshot, RegionSnapshot, RunMeta, RunStatus};
47
48pub const RUN_ARCHIVE_MAGIC: &[u8; 4] = b"LVR1";
50
51pub const RUN_ARCHIVE_VERSION: u16 = 1;
53
54#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
61pub struct RunIdentity {
62 pub run_id: String,
64 pub machine_id: String,
66 pub world_id: String,
68 pub created_at: i64,
70}
71
72#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
74pub struct MessageRecord {
75 pub role: String,
77 pub content: String,
79}
80
81#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
83pub struct ToolCallRecord {
84 pub id: String,
86 pub name: String,
88 pub arguments: String,
90 pub result: Option<String>,
92 #[serde(default, skip_serializing_if = "Option::is_none")]
96 pub thought_signature: Option<String>,
97}
98
99#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
102pub struct InferenceRequestRecord {
103 pub model: String,
105 pub system: Vec<String>,
107 pub messages: Vec<MessageRecord>,
109 pub tool_names: Vec<String>,
111 pub temperature: f32,
113 pub max_tokens: usize,
115}
116
117#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
119pub struct InferenceResponseRecord {
120 pub content: String,
122 pub tool_calls: Vec<ToolCallRecord>,
124 pub prompt_tokens: usize,
126 pub completion_tokens: usize,
128 pub cached_tokens: usize,
130 pub cache_write_tokens: usize,
132}
133
134#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
136pub enum RegionDelta {
137 Set(RegionSnapshot),
140 Append {
143 name: String,
145 entries: Vec<RegionEntrySnapshot>,
147 current_tokens: usize,
149 },
150 Clear {
152 name: String,
154 },
155 Remove {
157 name: String,
159 },
160}
161
162#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
164pub struct ContextDelta {
165 pub stage_name: String,
167 pub total_tokens: usize,
169 pub max_tokens: usize,
171 pub regions: Vec<RegionDelta>,
173}
174
175#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
177pub enum RunRecord {
178 Header {
180 identity: RunIdentity,
182 meta: Box<RunMeta>,
184 },
185 OwnershipChanged {
187 machine_id: String,
189 world_id: String,
191 at: i64,
193 },
194 Inference {
196 stage: String,
198 iteration: usize,
200 request: InferenceRequestRecord,
202 response: InferenceResponseRecord,
204 at: i64,
206 },
207 ToolBatch {
215 calls: Vec<ToolCallRecord>,
217 at: i64,
219 #[serde(default)]
221 stage_index: usize,
222 #[serde(default)]
225 iteration: usize,
226 #[serde(default)]
228 response: String,
229 },
230 ToolCallDone {
232 iteration: usize,
234 call_id: String,
236 result: String,
238 at: i64,
240 },
241 ContextCheckpoint {
243 snapshot: ContextSnapshot,
245 at: i64,
247 },
248 ContextDiff {
250 delta: ContextDelta,
252 at: i64,
254 },
255 Message {
257 message: MessageRecord,
259 at: i64,
261 },
262 StatusChanged {
264 status: RunStatus,
266 at: i64,
268 },
269 Checkpoint {
272 meta: Box<RunMeta>,
274 context: ContextSnapshot,
276 at: i64,
278 },
279 Progress {
284 meta: Box<RunMeta>,
286 delta: ContextDelta,
288 at: i64,
290 },
291}
292
293fn is_prefix(prev: &[RegionEntrySnapshot], next: &[RegionEntrySnapshot]) -> bool {
297 prev.len() <= next.len() && next[..prev.len()] == *prev
298}
299
300pub fn diff_context(prev: &ContextSnapshot, next: &ContextSnapshot) -> ContextDelta {
304 let mut regions = Vec::new();
305 for nr in &next.regions {
306 match prev.regions.iter().find(|r| r.name == nr.name) {
307 None => regions.push(RegionDelta::Set(nr.clone())),
308 Some(pr) => {
309 if pr == nr {
310 } else if nr.entries.is_empty() && !pr.entries.is_empty() {
312 regions.push(RegionDelta::Clear {
313 name: nr.name.clone(),
314 });
315 } else if pr.kind == nr.kind
316 && pr.max_tokens == nr.max_tokens
317 && is_prefix(&pr.entries, &nr.entries)
318 {
319 regions.push(RegionDelta::Append {
320 name: nr.name.clone(),
321 entries: nr.entries[pr.entries.len()..].to_vec(),
322 current_tokens: nr.current_tokens,
323 });
324 } else {
325 regions.push(RegionDelta::Set(nr.clone()));
326 }
327 }
328 }
329 }
330 for pr in &prev.regions {
331 if !next.regions.iter().any(|r| r.name == pr.name) {
332 regions.push(RegionDelta::Remove {
333 name: pr.name.clone(),
334 });
335 }
336 }
337 ContextDelta {
338 stage_name: next.stage_name.clone(),
339 total_tokens: next.total_tokens,
340 max_tokens: next.max_tokens,
341 regions,
342 }
343}
344
345#[derive(Debug, Clone, PartialEq)]
356pub struct RegionDigest {
357 pub name: String,
359 pub kind: String,
361 pub current_tokens: usize,
363 pub max_tokens: usize,
365 pub entries: Vec<u64>,
367}
368
369#[derive(Debug, Clone, PartialEq, Default)]
371pub struct ContextDigest {
372 pub regions: Vec<RegionDigest>,
374}
375
376fn entry_digest(entry: &RegionEntrySnapshot) -> u64 {
380 use std::hash::{Hash, Hasher};
381 let mut hasher = std::collections::hash_map::DefaultHasher::new();
382 entry.content.hash(&mut hasher);
383 entry.tokens.hash(&mut hasher);
384 entry.key.hash(&mut hasher);
385 serde_json::to_string(&entry.kind)
388 .expect("EntryKind always serializes")
389 .hash(&mut hasher);
390 serde_json::to_string(&entry.metadata)
391 .expect("entry metadata always serializes")
392 .hash(&mut hasher);
393 serde_json::to_string(&entry.taint)
394 .expect("taint always serializes")
395 .hash(&mut hasher);
396 hasher.finish()
397}
398
399pub fn digest_context(snapshot: &ContextSnapshot) -> ContextDigest {
401 ContextDigest {
402 regions: snapshot
403 .regions
404 .iter()
405 .map(|r| RegionDigest {
406 name: r.name.clone(),
407 kind: r.kind.clone(),
408 current_tokens: r.current_tokens,
409 max_tokens: r.max_tokens,
410 entries: r.entries.iter().map(entry_digest).collect(),
411 })
412 .collect(),
413 }
414}
415
416fn is_prefix_digest(prev: &[u64], next: &[RegionEntrySnapshot]) -> bool {
418 prev.len() <= next.len()
419 && prev
420 .iter()
421 .zip(next)
422 .all(|(hash, entry)| *hash == entry_digest(entry))
423}
424
425pub fn diff_context_digest(prev: &ContextDigest, next: &ContextSnapshot) -> ContextDelta {
430 let mut regions = Vec::new();
431 for nr in &next.regions {
432 match prev.regions.iter().find(|r| r.name == nr.name) {
433 None => regions.push(RegionDelta::Set(nr.clone())),
434 Some(pr) => {
435 let unchanged = pr.kind == nr.kind
436 && pr.max_tokens == nr.max_tokens
437 && pr.current_tokens == nr.current_tokens
438 && pr.entries.len() == nr.entries.len()
439 && is_prefix_digest(&pr.entries, &nr.entries);
440 if unchanged {
441 } else if nr.entries.is_empty() && !pr.entries.is_empty() {
443 regions.push(RegionDelta::Clear {
444 name: nr.name.clone(),
445 });
446 } else if pr.kind == nr.kind
447 && pr.max_tokens == nr.max_tokens
448 && is_prefix_digest(&pr.entries, &nr.entries)
449 {
450 regions.push(RegionDelta::Append {
451 name: nr.name.clone(),
452 entries: nr.entries[pr.entries.len()..].to_vec(),
453 current_tokens: nr.current_tokens,
454 });
455 } else {
456 regions.push(RegionDelta::Set(nr.clone()));
457 }
458 }
459 }
460 }
461 for pr in &prev.regions {
462 if !next.regions.iter().any(|r| r.name == pr.name) {
463 regions.push(RegionDelta::Remove {
464 name: pr.name.clone(),
465 });
466 }
467 }
468 ContextDelta {
469 stage_name: next.stage_name.clone(),
470 total_tokens: next.total_tokens,
471 max_tokens: next.max_tokens,
472 regions,
473 }
474}
475
476pub fn apply_delta(base: &mut ContextSnapshot, delta: &ContextDelta) {
480 base.stage_name = delta.stage_name.clone();
481 base.total_tokens = delta.total_tokens;
482 base.max_tokens = delta.max_tokens;
483 for region_delta in &delta.regions {
484 match region_delta {
485 RegionDelta::Set(snapshot) => {
486 match base.regions.iter_mut().find(|r| r.name == snapshot.name) {
487 Some(existing) => *existing = snapshot.clone(),
488 None => base.regions.push(snapshot.clone()),
489 }
490 }
491 RegionDelta::Append {
492 name,
493 entries,
494 current_tokens,
495 } => {
496 if let Some(region) = base.regions.iter_mut().find(|r| &r.name == name) {
497 region.entries.extend(entries.iter().cloned());
498 region.current_tokens = *current_tokens;
499 }
500 }
501 RegionDelta::Clear { name } => {
502 if let Some(region) = base.regions.iter_mut().find(|r| &r.name == name) {
503 region.entries.clear();
504 region.current_tokens = 0;
505 }
506 }
507 RegionDelta::Remove { name } => {
508 base.regions.retain(|r| &r.name != name);
509 }
510 }
511 }
512}
513
514pub fn write_archive_start(w: &mut dyn Write, version: u16) -> io::Result<()> {
518 w.write_all(RUN_ARCHIVE_MAGIC)?;
519 w.write_all(&version.to_be_bytes())?;
520 Ok(())
521}
522
523pub fn read_archive_start(r: &mut dyn Read) -> io::Result<u16> {
525 let mut magic = [0u8; 4];
526 r.read_exact(&mut magic)?;
527 if &magic != RUN_ARCHIVE_MAGIC {
528 return Err(io::Error::new(
529 io::ErrorKind::InvalidData,
530 "not a leviath run archive (bad magic)",
531 ));
532 }
533 let mut version = [0u8; 2];
534 r.read_exact(&mut version)?;
535 Ok(u16::from_be_bytes(version))
536}
537
538pub fn write_record(w: &mut dyn Write, record: &RunRecord) -> io::Result<()> {
541 let payload = serde_json::to_vec(record).expect("a RunRecord always serializes to JSON");
542 let len = payload.len() as u64;
543 w.write_all(&len.to_be_bytes())?;
544 w.write_all(&payload)?;
545 Ok(())
546}
547
548fn read_exact_or_eof(r: &mut dyn Read, buf: &mut [u8]) -> io::Result<bool> {
551 let mut filled = 0;
552 while filled < buf.len() {
553 match r.read(&mut buf[filled..])? {
554 0 => {
555 if filled == 0 {
556 return Ok(false); }
558 return Err(io::Error::new(
559 io::ErrorKind::UnexpectedEof,
560 "truncated run-archive frame",
561 ));
562 }
563 n => filled += n,
564 }
565 }
566 Ok(true)
567}
568
569const MAX_RECORD_BYTES: u64 = 256 * 1024 * 1024;
576
577pub fn read_record(r: &mut dyn Read) -> io::Result<Option<RunRecord>> {
579 let mut len_bytes = [0u8; 8];
580 if !read_exact_or_eof(r, &mut len_bytes)? {
581 return Ok(None);
582 }
583 let len = u64::from_be_bytes(len_bytes);
584 if len > MAX_RECORD_BYTES {
591 return Err(io::Error::new(
592 io::ErrorKind::InvalidData,
593 format!("run-archive frame claims {len} bytes, over the {MAX_RECORD_BYTES} cap"),
594 ));
595 }
596 let mut payload = vec![0u8; len as usize];
597 if !read_exact_or_eof(r, &mut payload)? {
598 return Err(io::Error::new(
599 io::ErrorKind::UnexpectedEof,
600 "truncated run-archive frame",
601 ));
602 }
603 let record = serde_json::from_slice(&payload)
604 .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
605 Ok(Some(record))
606}
607
608pub fn read_archive(r: &mut dyn Read) -> io::Result<(u16, Vec<RunRecord>)> {
610 let version = read_archive_start(r)?;
611 let mut records = Vec::new();
612 while let Some(record) = read_record(r)? {
613 records.push(record);
614 }
615 Ok((version, records))
616}
617
618pub fn read_archive_lenient(r: &mut dyn Read) -> io::Result<(u16, Vec<RunRecord>)> {
631 let version = read_archive_start(r)?;
632 let mut records = Vec::new();
633 while let Ok(Some(record)) = read_record(r) {
636 records.push(record);
637 }
638 Ok((version, records))
639}
640
641#[derive(Debug, Clone, PartialEq)]
648pub struct PendingToolBatch {
649 pub stage_index: usize,
651 pub iteration: usize,
653 pub response: String,
655 pub calls: Vec<ToolCallRecord>,
657}
658
659#[derive(Debug, Clone, PartialEq)]
662pub struct FoldedRun {
663 pub identity: RunIdentity,
665 pub meta: RunMeta,
667 pub context: ContextSnapshot,
669 pub messages: Vec<MessageRecord>,
671 pub inference_count: usize,
673 pub tool_call_count: usize,
675 pub pending_batch: Option<PendingToolBatch>,
679}
680
681pub fn context_contains_batch(context: &ContextSnapshot, batch: &PendingToolBatch) -> bool {
685 let Some(first_id) = batch.calls.first().map(|c| c.id.as_str()) else {
686 return false;
687 };
688 context.regions.iter().any(|region| {
689 region.entries.iter().any(|entry| {
690 matches!(
691 &entry.kind,
692 crate::region::EntryKind::AssistantTurn { tool_calls }
693 if tool_calls.iter().any(|tc| tc.id == first_id)
694 )
695 })
696 })
697}
698
699pub fn fold(records: &[RunRecord]) -> Option<FoldedRun> {
702 let mut iter = records.iter();
703 let (identity, meta) = match iter.next() {
704 Some(RunRecord::Header { identity, meta }) => (identity.clone(), (**meta).clone()),
705 _ => return None,
706 };
707 let mut folded = FoldedRun {
708 identity,
709 meta,
710 context: ContextSnapshot {
711 stage_name: String::new(),
712 total_tokens: 0,
713 max_tokens: 0,
714 regions: Vec::new(),
715 },
716 messages: Vec::new(),
717 inference_count: 0,
718 tool_call_count: 0,
719 pending_batch: None,
720 };
721 for record in iter {
722 match record {
723 RunRecord::Header { identity, meta } => {
724 folded.identity = identity.clone();
725 folded.meta = (**meta).clone();
726 }
727 RunRecord::OwnershipChanged {
728 machine_id,
729 world_id,
730 ..
731 } => {
732 folded.identity.machine_id = machine_id.clone();
733 folded.identity.world_id = world_id.clone();
734 }
735 RunRecord::Inference { .. } => folded.inference_count += 1,
736 RunRecord::ToolBatch {
737 calls,
738 stage_index,
739 iteration,
740 response,
741 ..
742 } => {
743 folded.tool_call_count += calls.len();
744 folded.pending_batch = Some(PendingToolBatch {
747 stage_index: *stage_index,
748 iteration: *iteration,
749 response: response.clone(),
750 calls: calls.clone(),
751 });
752 }
753 RunRecord::ToolCallDone {
754 iteration,
755 call_id,
756 result,
757 ..
758 } => {
759 if let Some(batch) = folded
762 .pending_batch
763 .as_mut()
764 .filter(|b| b.iteration == *iteration)
765 && let Some(call) = batch.calls.iter_mut().find(|c| c.id == *call_id)
766 {
767 call.result = Some(result.clone());
768 }
769 }
770 RunRecord::ContextCheckpoint { snapshot, .. } => folded.context = snapshot.clone(),
771 RunRecord::ContextDiff { delta, .. } => apply_delta(&mut folded.context, delta),
772 RunRecord::Message { message, .. } => folded.messages.push(message.clone()),
773 RunRecord::StatusChanged { status, .. } => folded.meta.status = status.clone(),
774 RunRecord::Checkpoint { meta, context, .. } => {
775 folded.meta = (**meta).clone();
776 folded.context = context.clone();
777 }
778 RunRecord::Progress { meta, delta, .. } => {
779 folded.meta = (**meta).clone();
780 apply_delta(&mut folded.context, delta);
781 }
782 }
783 }
784 if let Some(batch) = &folded.pending_batch
789 && (folded.meta.iteration != batch.iteration
790 || context_contains_batch(&folded.context, batch))
791 {
792 folded.pending_batch = None;
793 }
794 Some(folded)
795}
796
797#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
800pub struct RunPoint {
801 pub meta: RunMeta,
803 pub context: ContextSnapshot,
805 pub at: i64,
807}
808
809#[derive(Debug)]
812pub struct PointRef<'a> {
813 pub index: usize,
817 pub at: i64,
819 pub meta: &'a RunMeta,
821 pub context: &'a ContextSnapshot,
823}
824
825pub fn visit_points(records: &[RunRecord], visit: &mut dyn FnMut(PointRef<'_>) -> ControlFlow<()>) {
844 let mut iter = records.iter();
845 let Some(mut folder) = (match iter.next() {
846 Some(first) => PointFolder::start(first),
847 None => None,
848 }) else {
849 return;
850 };
851 for record in iter {
852 if folder.push(record, visit).is_break() {
853 return;
854 }
855 }
856}
857
858pub fn visit_archive_points(
868 r: &mut dyn Read,
869 visit: &mut dyn FnMut(PointRef<'_>) -> ControlFlow<()>,
870) -> io::Result<()> {
871 read_archive_start(r)?;
872 let mut folder = match read_record(r) {
873 Ok(Some(first)) => match PointFolder::start(&first) {
874 Some(folder) => folder,
875 None => return Ok(()),
876 },
877 _ => return Ok(()),
878 };
879 while let Ok(Some(record)) = read_record(r) {
880 if folder.push(&record, visit).is_break() {
881 return Ok(());
882 }
883 }
884 Ok(())
885}
886
887struct PointFolder {
892 meta: RunMeta,
893 context: ContextSnapshot,
894 index: usize,
895}
896
897impl PointFolder {
898 fn start(first: &RunRecord) -> Option<Self> {
902 match first {
903 RunRecord::Header { meta, .. } => Some(Self {
904 meta: (**meta).clone(),
905 context: ContextSnapshot {
906 stage_name: String::new(),
907 total_tokens: 0,
908 max_tokens: 0,
909 regions: Vec::new(),
910 },
911 index: 0,
912 }),
913 _ => None,
914 }
915 }
916
917 fn push(
919 &mut self,
920 record: &RunRecord,
921 visit: &mut dyn FnMut(PointRef<'_>) -> ControlFlow<()>,
922 ) -> ControlFlow<()> {
923 let at = match record {
924 RunRecord::Header { meta: m, .. } => {
925 self.meta = (**m).clone();
926 return ControlFlow::Continue(());
927 }
928 RunRecord::StatusChanged { status, .. } => {
929 self.meta.status = status.clone();
930 return ControlFlow::Continue(());
931 }
932 RunRecord::ContextCheckpoint { snapshot, at } => {
933 self.context = snapshot.clone();
934 *at
935 }
936 RunRecord::ContextDiff { delta, at } => {
937 apply_delta(&mut self.context, delta);
938 *at
939 }
940 RunRecord::Checkpoint {
941 meta: m,
942 context: c,
943 at,
944 } => {
945 self.meta = (**m).clone();
946 self.context = c.clone();
947 *at
948 }
949 RunRecord::Progress { meta: m, delta, at } => {
950 self.meta = (**m).clone();
951 apply_delta(&mut self.context, delta);
952 *at
953 }
954 RunRecord::OwnershipChanged { .. }
956 | RunRecord::Inference { .. }
957 | RunRecord::ToolBatch { .. }
958 | RunRecord::ToolCallDone { .. }
959 | RunRecord::Message { .. } => return ControlFlow::Continue(()),
960 };
961 let flow = visit(PointRef {
962 index: self.index,
963 at,
964 meta: &self.meta,
965 context: &self.context,
966 });
967 self.index += 1;
968 flow
969 }
970}
971
972pub fn replay_points(records: &[RunRecord]) -> Vec<RunPoint> {
982 let mut points = Vec::new();
983 visit_points(records, &mut |point| {
984 points.push(RunPoint {
985 meta: point.meta.clone(),
986 context: point.context.clone(),
987 at: point.at,
988 });
989 ControlFlow::Continue(())
990 });
991 points
992}
993
994#[cfg(test)]
995mod tests {
996 use super::*;
997 use crate::run_meta::RunStatus;
998
999 fn identity() -> RunIdentity {
1000 RunIdentity {
1001 run_id: "run-1".to_string(),
1002 machine_id: "machine-a".to_string(),
1003 world_id: "world-x".to_string(),
1004 created_at: 100,
1005 }
1006 }
1007
1008 const FIXTURE_NOW: i64 = 1_700_000_000;
1016
1017 fn meta() -> RunMeta {
1018 let mut meta = RunMeta::new(
1019 "run-1".to_string(),
1020 "coder".to_string(),
1021 "/agents/coder".to_string(),
1022 "do it".to_string(),
1023 Some("anthropic/claude".to_string()),
1024 "/work".to_string(),
1025 2,
1026 );
1027 meta.started_at = FIXTURE_NOW;
1028 meta.updated_at = FIXTURE_NOW;
1029 meta
1030 }
1031
1032 #[test]
1036 fn the_fixture_does_not_move_with_the_clock() {
1037 let first = meta();
1038 let mut later = meta();
1039 assert_eq!(first, later, "the fixture is rebuilt identically");
1043 later.started_at += 1;
1044 assert_ne!(
1045 first, later,
1046 "and the comparison is sensitive to the field that used to drift"
1047 );
1048 }
1049
1050 fn entry(content: &str, tokens: usize) -> RegionEntrySnapshot {
1051 RegionEntrySnapshot {
1052 content: content.to_string(),
1053 tokens,
1054 kind: crate::region::EntryKind::Text,
1055 metadata: None,
1056 key: None,
1057 taint: Default::default(),
1058 }
1059 }
1060
1061 fn region(name: &str, entries: Vec<RegionEntrySnapshot>) -> RegionSnapshot {
1062 let current = entries.iter().map(|e| e.tokens).sum();
1063 RegionSnapshot {
1064 name: name.to_string(),
1065 kind: "clearable".to_string(),
1066 current_tokens: current,
1067 max_tokens: 1000,
1068 entries,
1069 }
1070 }
1071
1072 fn snapshot(stage: &str, regions: Vec<RegionSnapshot>) -> ContextSnapshot {
1073 let total = regions.iter().map(|r| r.current_tokens).sum();
1074 ContextSnapshot {
1075 stage_name: stage.to_string(),
1076 total_tokens: total,
1077 max_tokens: 10_000,
1078 regions,
1079 }
1080 }
1081
1082 fn header() -> RunRecord {
1083 RunRecord::Header {
1084 identity: identity(),
1085 meta: Box::new(meta()),
1086 }
1087 }
1088
1089 fn region_delta_kind(d: &RegionDelta) -> &'static str {
1093 match d {
1094 RegionDelta::Set(_) => "set",
1095 RegionDelta::Append { .. } => "append",
1096 RegionDelta::Clear { .. } => "clear",
1097 RegionDelta::Remove { .. } => "remove",
1098 }
1099 }
1100
1101 fn assert_diff_roundtrip(a: &ContextSnapshot, b: &ContextSnapshot) {
1106 let delta = diff_context(a, b);
1107 let mut base = a.clone();
1108 apply_delta(&mut base, &delta);
1109 assert_eq!(&base, b);
1110 }
1111
1112 #[test]
1113 fn diff_append_only_growth_is_compact() {
1114 let a = snapshot("s1", vec![region("conv", vec![entry("hi", 1)])]);
1115 let b = snapshot(
1116 "s1",
1117 vec![region("conv", vec![entry("hi", 1), entry("there", 2)])],
1118 );
1119 let delta = diff_context(&a, &b);
1120 assert_eq!(region_delta_kind(&delta.regions[0]), "append");
1121 assert_diff_roundtrip(&a, &b);
1122 }
1123
1124 #[test]
1125 fn diff_new_region_is_set() {
1126 let a = snapshot("s1", vec![region("conv", vec![entry("hi", 1)])]);
1127 let b = snapshot(
1128 "s1",
1129 vec![
1130 region("conv", vec![entry("hi", 1)]),
1131 region("plan", vec![entry("p", 3)]),
1132 ],
1133 );
1134 let delta = diff_context(&a, &b);
1135 assert!(delta.regions.iter().any(|d| region_delta_kind(d) == "set"));
1136 assert_diff_roundtrip(&a, &b);
1137 }
1138
1139 #[test]
1140 fn diff_cleared_region() {
1141 let a = snapshot("s1", vec![region("conv", vec![entry("hi", 1)])]);
1142 let b = snapshot("s1", vec![region("conv", vec![])]);
1143 let delta = diff_context(&a, &b);
1144 assert_eq!(region_delta_kind(&delta.regions[0]), "clear");
1145 assert_diff_roundtrip(&a, &b);
1146 }
1147
1148 #[test]
1149 fn diff_removed_region() {
1150 let a = snapshot(
1151 "s1",
1152 vec![
1153 region("conv", vec![entry("hi", 1)]),
1154 region("plan", vec![entry("p", 3)]),
1155 ],
1156 );
1157 let b = snapshot("s1", vec![region("conv", vec![entry("hi", 1)])]);
1158 let delta = diff_context(&a, &b);
1159 assert!(
1160 delta
1161 .regions
1162 .iter()
1163 .any(|d| region_delta_kind(d) == "remove")
1164 );
1165 assert_diff_roundtrip(&a, &b);
1166 }
1167
1168 #[test]
1169 fn diff_non_prefix_rewrite_is_set() {
1170 let a = snapshot("s1", vec![region("conv", vec![entry("old", 1)])]);
1172 let b = snapshot("s1", vec![region("conv", vec![entry("new", 1)])]);
1173 let delta = diff_context(&a, &b);
1174 assert_eq!(region_delta_kind(&delta.regions[0]), "set");
1175 assert_diff_roundtrip(&a, &b);
1176 }
1177
1178 #[test]
1179 fn diff_kind_change_is_set_not_append() {
1180 let a = snapshot("s1", vec![region("conv", vec![entry("hi", 1)])]);
1182 let mut grown = region("conv", vec![entry("hi", 1), entry("more", 1)]);
1183 grown.kind = "sliding".to_string();
1184 let b = snapshot("s1", vec![grown]);
1185 let delta = diff_context(&a, &b);
1186 assert_eq!(region_delta_kind(&delta.regions[0]), "set");
1187 assert_diff_roundtrip(&a, &b);
1188 }
1189
1190 fn framed(records: &[RunRecord]) -> Vec<u8> {
1194 let mut buf = Vec::new();
1195 write_archive_start(&mut buf, RUN_ARCHIVE_VERSION).unwrap();
1196 for record in records {
1197 write_record(&mut buf, record).unwrap();
1198 }
1199 buf
1200 }
1201
1202 fn try_collect_streamed(bytes: &[u8]) -> io::Result<Vec<(usize, i64, usize)>> {
1206 let mut seen = Vec::new();
1207 visit_archive_points(&mut &bytes[..], &mut |p| {
1208 seen.push((p.index, p.at, p.context.total_tokens));
1209 ControlFlow::Continue(())
1210 })?;
1211 Ok(seen)
1212 }
1213
1214 fn collect_streamed(bytes: &[u8]) -> Vec<(usize, i64, usize)> {
1216 try_collect_streamed(bytes).unwrap()
1217 }
1218
1219 #[test]
1220 fn visit_archive_points_matches_visit_points() {
1221 let records = vec![
1222 header(),
1223 RunRecord::ContextCheckpoint {
1224 snapshot: snapshot("s1", vec![region("conv", vec![entry("hi", 1)])]),
1225 at: 10,
1226 },
1227 RunRecord::StatusChanged {
1228 status: RunStatus::Running,
1229 at: 11,
1230 },
1231 RunRecord::Progress {
1232 meta: Box::new(meta()),
1233 delta: diff_context(
1234 &snapshot("s1", vec![region("conv", vec![entry("hi", 1)])]),
1235 &snapshot(
1236 "s1",
1237 vec![region("conv", vec![entry("hi", 1), entry("more", 2)])],
1238 ),
1239 ),
1240 at: 12,
1241 },
1242 ];
1243 let mut in_memory = Vec::new();
1244 visit_points(&records, &mut |p| {
1245 in_memory.push((p.index, p.at, p.context.total_tokens));
1246 ControlFlow::Continue(())
1247 });
1248 assert_eq!(collect_streamed(&framed(&records)), in_memory);
1249 assert_eq!(in_memory.len(), 2, "checkpoint + progress = two points");
1250 }
1251
1252 #[test]
1253 fn visit_archive_points_rejects_a_bad_preamble() {
1254 assert!(try_collect_streamed(b"not an archive at all").is_err());
1255 }
1256
1257 #[test]
1258 fn visit_archive_points_is_lenient_about_a_torn_tail() {
1259 let records = vec![
1260 header(),
1261 RunRecord::ContextCheckpoint {
1262 snapshot: snapshot("s1", vec![region("conv", vec![entry("hi", 1)])]),
1263 at: 10,
1264 },
1265 ];
1266 let mut bytes = framed(&records);
1267 bytes.extend_from_slice(&1000u64.to_be_bytes());
1269 bytes.extend_from_slice(b"partial");
1270 assert_eq!(collect_streamed(&bytes).len(), 1, "points before the tear");
1271 }
1272
1273 #[test]
1274 fn visit_archive_points_visits_nothing_without_a_header() {
1275 let records = vec![RunRecord::ContextCheckpoint {
1276 snapshot: snapshot("s1", vec![region("conv", vec![entry("hi", 1)])]),
1277 at: 10,
1278 }];
1279 assert!(collect_streamed(&framed(&records)).is_empty());
1280 assert!(collect_streamed(&framed(&[])).is_empty());
1282 }
1283
1284 #[test]
1285 fn visit_archive_points_stops_on_break() {
1286 let records = vec![
1287 header(),
1288 RunRecord::ContextCheckpoint {
1289 snapshot: snapshot("s1", vec![region("conv", vec![entry("a", 1)])]),
1290 at: 10,
1291 },
1292 RunRecord::ContextCheckpoint {
1293 snapshot: snapshot("s1", vec![region("conv", vec![entry("b", 2)])]),
1294 at: 11,
1295 },
1296 ];
1297 let bytes = framed(&records);
1298 let mut seen = 0;
1299 visit_archive_points(&mut &bytes[..], &mut |_| {
1300 seen += 1;
1301 ControlFlow::Break(())
1302 })
1303 .unwrap();
1304 assert_eq!(seen, 1);
1305 }
1306
1307 fn assert_digest_matches_full_diff(a: &ContextSnapshot, b: &ContextSnapshot) {
1314 let via_digest = diff_context_digest(&digest_context(a), b);
1315 assert_eq!(via_digest, diff_context(a, b));
1316 let mut base = a.clone();
1318 apply_delta(&mut base, &via_digest);
1319 assert_eq!(&base, b);
1320 }
1321
1322 #[test]
1323 fn digest_diff_append_only_growth_is_compact() {
1324 let a = snapshot("s1", vec![region("conv", vec![entry("hi", 1)])]);
1325 let b = snapshot(
1326 "s1",
1327 vec![region("conv", vec![entry("hi", 1), entry("there", 2)])],
1328 );
1329 let delta = diff_context_digest(&digest_context(&a), &b);
1330 assert_eq!(region_delta_kind(&delta.regions[0]), "append");
1331 assert_digest_matches_full_diff(&a, &b);
1332 }
1333
1334 #[test]
1335 fn digest_diff_new_cleared_removed_and_rewritten_regions() {
1336 let a = snapshot(
1337 "s1",
1338 vec![
1339 region("conv", vec![entry("hi", 1)]),
1340 region("gone", vec![entry("bye", 1)]),
1341 region("wiped", vec![entry("w", 1)]),
1342 region("rewritten", vec![entry("old", 1)]),
1343 ],
1344 );
1345 let b = snapshot(
1346 "s1",
1347 vec![
1348 region("conv", vec![entry("hi", 1)]),
1349 region("wiped", vec![]),
1350 region("rewritten", vec![entry("new", 1)]),
1351 region("fresh", vec![entry("f", 2)]),
1352 ],
1353 );
1354 let delta = diff_context_digest(&digest_context(&a), &b);
1355 let kinds: Vec<_> = delta.regions.iter().map(region_delta_kind).collect();
1356 assert_eq!(kinds, vec!["clear", "set", "set", "remove"]);
1357 assert_digest_matches_full_diff(&a, &b);
1358 }
1359
1360 #[test]
1361 fn digest_diff_unchanged_region_emits_nothing() {
1362 let a = snapshot("s1", vec![region("conv", vec![entry("hi", 1)])]);
1363 let delta = diff_context_digest(&digest_context(&a), &a.clone());
1364 assert!(delta.regions.is_empty());
1365 assert_digest_matches_full_diff(&a, &a.clone());
1366 }
1367
1368 #[test]
1369 fn digest_diff_kind_change_is_set_not_append() {
1370 let a = snapshot("s1", vec![region("conv", vec![entry("hi", 1)])]);
1371 let mut grown = region("conv", vec![entry("hi", 1), entry("more", 1)]);
1372 grown.kind = "sliding".to_string();
1373 let b = snapshot("s1", vec![grown]);
1374 let delta = diff_context_digest(&digest_context(&a), &b);
1375 assert_eq!(region_delta_kind(&delta.regions[0]), "set");
1376 assert_digest_matches_full_diff(&a, &b);
1377 }
1378
1379 #[test]
1382 fn digest_diff_token_recount_is_an_empty_append() {
1383 let a = snapshot("s1", vec![region("conv", vec![entry("hi", 1)])]);
1384 let mut recounted = region("conv", vec![entry("hi", 1)]);
1385 recounted.current_tokens = 42;
1386 let b = snapshot("s1", vec![recounted]);
1387 let delta = diff_context_digest(&digest_context(&a), &b);
1388 assert_eq!(region_delta_kind(&delta.regions[0]), "append");
1389 assert_digest_matches_full_diff(&a, &b);
1390 }
1391
1392 #[test]
1395 fn entry_digest_covers_every_field() {
1396 let base = entry("text", 1);
1397 let variants = [
1398 entry("other", 1),
1399 entry("text", 2),
1400 RegionEntrySnapshot {
1401 key: Some("k".to_string()),
1402 ..entry("text", 1)
1403 },
1404 RegionEntrySnapshot {
1405 metadata: Some(serde_json::json!({"a": 1})),
1406 ..entry("text", 1)
1407 },
1408 RegionEntrySnapshot {
1409 kind: crate::region::EntryKind::ToolResult {
1410 tool_call_id: "c1".to_string(),
1411 tool_name: "shell".to_string(),
1412 is_error: false,
1413 },
1414 ..entry("text", 1)
1415 },
1416 ];
1417 let base_hash = entry_digest(&base);
1418 for variant in &variants {
1419 assert_ne!(
1420 entry_digest(variant),
1421 base_hash,
1422 "field change must change the digest: {variant:?}"
1423 );
1424 }
1425 assert_eq!(entry_digest(&base), entry_digest(&entry("text", 1)));
1427 }
1428
1429 #[test]
1430 fn diff_unchanged_region_emits_nothing() {
1431 let a = snapshot("s1", vec![region("conv", vec![entry("hi", 1)])]);
1432 let b = a.clone();
1433 let delta = diff_context(&a, &b);
1434 assert!(delta.regions.is_empty());
1435 assert_diff_roundtrip(&a, &b);
1436 }
1437
1438 #[test]
1439 fn apply_delta_skips_unknown_regions_leniently() {
1440 let mut base = snapshot("s1", vec![]);
1442 let delta = ContextDelta {
1443 stage_name: "s1".to_string(),
1444 total_tokens: 0,
1445 max_tokens: 10_000,
1446 regions: vec![
1447 RegionDelta::Append {
1448 name: "ghost".to_string(),
1449 entries: vec![entry("x", 1)],
1450 current_tokens: 1,
1451 },
1452 RegionDelta::Clear {
1453 name: "ghost".to_string(),
1454 },
1455 RegionDelta::Remove {
1456 name: "ghost".to_string(),
1457 },
1458 ],
1459 };
1460 apply_delta(&mut base, &delta);
1461 assert!(base.regions.is_empty());
1462 }
1463
1464 fn all_record_kinds() -> Vec<RunRecord> {
1467 vec![
1468 header(),
1469 RunRecord::OwnershipChanged {
1470 machine_id: "machine-b".to_string(),
1471 world_id: "world-y".to_string(),
1472 at: 101,
1473 },
1474 RunRecord::Inference {
1475 stage: "plan".to_string(),
1476 iteration: 0,
1477 request: InferenceRequestRecord {
1478 model: "m".to_string(),
1479 system: vec!["sys".to_string()],
1480 messages: vec![MessageRecord {
1481 role: "user".to_string(),
1482 content: "hi".to_string(),
1483 }],
1484 tool_names: vec!["read_file".to_string()],
1485 temperature: 0.7,
1486 max_tokens: 1024,
1487 },
1488 response: InferenceResponseRecord {
1489 content: "ok".to_string(),
1490 tool_calls: vec![],
1491 prompt_tokens: 10,
1492 completion_tokens: 5,
1493 cached_tokens: 0,
1494 cache_write_tokens: 0,
1495 },
1496 at: 102,
1497 },
1498 RunRecord::ToolBatch {
1499 calls: vec![ToolCallRecord {
1500 id: "c1".to_string(),
1501 name: "read_file".to_string(),
1502 arguments: "{}".to_string(),
1503 result: Some("body".to_string()),
1504 thought_signature: Some("sig".to_string()),
1505 }],
1506 at: 103,
1507 stage_index: 0,
1508 iteration: 0,
1509 response: "reading".to_string(),
1510 },
1511 RunRecord::ToolCallDone {
1512 iteration: 0,
1513 call_id: "c1".to_string(),
1514 result: "body".to_string(),
1515 at: 103,
1516 },
1517 RunRecord::ContextCheckpoint {
1518 snapshot: snapshot("plan", vec![region("conv", vec![entry("hi", 1)])]),
1519 at: 104,
1520 },
1521 RunRecord::ContextDiff {
1522 delta: ContextDelta {
1523 stage_name: "plan".to_string(),
1524 total_tokens: 3,
1525 max_tokens: 10_000,
1526 regions: vec![RegionDelta::Append {
1527 name: "conv".to_string(),
1528 entries: vec![entry("more", 2)],
1529 current_tokens: 3,
1530 }],
1531 },
1532 at: 105,
1533 },
1534 RunRecord::Message {
1535 message: MessageRecord {
1536 role: "user".to_string(),
1537 content: "another".to_string(),
1538 },
1539 at: 106,
1540 },
1541 RunRecord::StatusChanged {
1542 status: RunStatus::Complete,
1543 at: 107,
1544 },
1545 RunRecord::Checkpoint {
1546 meta: Box::new(meta()),
1547 context: snapshot("plan", vec![region("conv", vec![entry("hi", 1)])]),
1548 at: 108,
1549 },
1550 RunRecord::Progress {
1551 meta: Box::new(meta()),
1552 delta: ContextDelta {
1553 stage_name: "plan".to_string(),
1554 total_tokens: 3,
1555 max_tokens: 10_000,
1556 regions: vec![RegionDelta::Append {
1557 name: "conv".to_string(),
1558 entries: vec![entry("step", 2)],
1559 current_tokens: 3,
1560 }],
1561 },
1562 at: 109,
1563 },
1564 ]
1565 }
1566
1567 #[test]
1568 fn archive_write_then_read_roundtrips_every_record_kind() {
1569 let records = all_record_kinds();
1570 let mut buf = Vec::new();
1571 write_archive_start(&mut buf, RUN_ARCHIVE_VERSION).unwrap();
1572 for r in &records {
1573 write_record(&mut buf, r).unwrap();
1574 }
1575 let (version, read) = read_archive(&mut buf.as_slice()).unwrap();
1576 assert_eq!(version, RUN_ARCHIVE_VERSION);
1577 assert_eq!(read, records);
1578 }
1579
1580 #[test]
1581 fn read_archive_start_rejects_bad_magic() {
1582 let mut bytes: &[u8] = b"XXXX\x00\x01";
1583 let err = read_archive_start(&mut bytes).unwrap_err();
1584 assert_eq!(err.kind(), io::ErrorKind::InvalidData);
1585 }
1586
1587 #[test]
1588 fn read_archive_start_reports_version() {
1589 let mut buf = Vec::new();
1590 write_archive_start(&mut buf, 7).unwrap();
1591 assert_eq!(read_archive_start(&mut buf.as_slice()).unwrap(), 7);
1592 }
1593
1594 #[test]
1595 fn read_record_returns_none_at_clean_eof() {
1596 let empty: &[u8] = &[];
1597 assert!(read_record(&mut { empty }).unwrap().is_none());
1598 }
1599
1600 #[test]
1601 fn read_record_errors_on_truncated_length_prefix() {
1602 let mut bytes: &[u8] = &[0, 0];
1604 let err = read_record(&mut bytes).unwrap_err();
1605 assert_eq!(err.kind(), io::ErrorKind::UnexpectedEof);
1606 }
1607
1608 #[test]
1609 fn read_record_errors_on_truncated_payload() {
1610 let mut bytes: &[u8] = &[0, 0, 0, 0, 0, 0, 0, 10, 1, 2];
1612 let err = read_record(&mut bytes).unwrap_err();
1613 assert_eq!(err.kind(), io::ErrorKind::UnexpectedEof);
1614 }
1615
1616 #[test]
1617 fn read_record_errors_on_empty_payload_at_boundary() {
1618 let mut bytes: &[u8] = &[0, 0, 0, 0, 0, 0, 0, 10];
1621 let err = read_record(&mut bytes).unwrap_err();
1622 assert_eq!(err.kind(), io::ErrorKind::UnexpectedEof);
1623 }
1624
1625 #[test]
1626 fn read_record_errors_on_invalid_json_payload() {
1627 let mut buf = Vec::new();
1629 let bad = b"not json";
1630 buf.extend_from_slice(&(bad.len() as u64).to_be_bytes());
1631 buf.extend_from_slice(bad);
1632 let err = read_record(&mut buf.as_slice()).unwrap_err();
1633 assert_eq!(err.kind(), io::ErrorKind::InvalidData);
1634 }
1635
1636 struct FailingReader;
1639 impl Read for FailingReader {
1640 fn read(&mut self, _buf: &mut [u8]) -> io::Result<usize> {
1641 Err(io::Error::other("device error"))
1642 }
1643 }
1644
1645 #[test]
1646 fn read_record_propagates_reader_errors() {
1647 let err = read_record(&mut FailingReader).unwrap_err();
1648 assert_eq!(err.kind(), io::ErrorKind::Other);
1649 }
1650
1651 #[test]
1652 fn read_archive_propagates_a_bad_preamble() {
1653 let mut bytes: &[u8] = b"LV";
1655 assert!(read_archive(&mut bytes).is_err());
1656 }
1657
1658 #[test]
1659 fn read_archive_propagates_a_bad_frame() {
1660 let mut buf = Vec::new();
1662 write_archive_start(&mut buf, RUN_ARCHIVE_VERSION).unwrap();
1663 buf.extend_from_slice(&[0, 0, 0, 0, 0, 0, 0, 5, 1, 2]); let err = read_archive(&mut buf.as_slice()).unwrap_err();
1665 assert_eq!(err.kind(), io::ErrorKind::UnexpectedEof);
1666 }
1667
1668 struct FailAfter {
1670 remaining: usize,
1671 }
1672 impl Write for FailAfter {
1673 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
1674 if self.remaining == 0 {
1675 return Err(io::Error::other("disk full"));
1676 }
1677 let n = buf.len().min(self.remaining);
1678 self.remaining -= n;
1679 Ok(n)
1680 }
1681 fn flush(&mut self) -> io::Result<()> {
1682 Ok(())
1683 }
1684 }
1685
1686 #[test]
1687 fn fail_after_writer_flush_is_a_noop() {
1688 assert!(FailAfter { remaining: 1 }.flush().is_ok());
1689 }
1690
1691 #[test]
1692 fn write_archive_start_propagates_write_errors() {
1693 assert!(write_archive_start(&mut FailAfter { remaining: 0 }, 1).is_err());
1695 assert!(write_archive_start(&mut FailAfter { remaining: 4 }, 1).is_err());
1696 }
1697
1698 #[test]
1699 fn write_record_propagates_write_errors() {
1700 let rec = header();
1701 assert!(write_record(&mut FailAfter { remaining: 0 }, &rec).is_err());
1703 assert!(write_record(&mut FailAfter { remaining: 8 }, &rec).is_err());
1704 }
1705
1706 #[test]
1712 fn an_absurd_frame_length_is_an_error_not_an_allocation() {
1713 let mut buf = Vec::new();
1714 write_archive_start(&mut buf, RUN_ARCHIVE_VERSION).unwrap();
1715 write_record(&mut buf, &header()).unwrap();
1716 buf.extend_from_slice(&u64::MAX.to_be_bytes());
1718
1719 let err = read_archive(&mut buf.as_slice())
1720 .expect_err("the strict reader must refuse an impossible frame");
1721 assert_eq!(err.kind(), io::ErrorKind::InvalidData, "{err}");
1722
1723 let (_, records) = read_archive_lenient(&mut buf.as_slice()).unwrap();
1726 assert_eq!(records, vec![header()]);
1727 }
1728
1729 #[test]
1730 fn read_archive_lenient_matches_strict_on_a_clean_archive() {
1731 let records = all_record_kinds();
1734 let mut buf = Vec::new();
1735 write_archive_start(&mut buf, RUN_ARCHIVE_VERSION).unwrap();
1736 for r in &records {
1737 write_record(&mut buf, r).unwrap();
1738 }
1739 let (version, read) = read_archive_lenient(&mut buf.as_slice()).unwrap();
1740 assert_eq!(version, RUN_ARCHIVE_VERSION);
1741 assert_eq!(read, records);
1742 }
1743
1744 #[test]
1745 fn read_archive_lenient_keeps_valid_prefix_before_a_torn_tail() {
1746 let mut buf = Vec::new();
1750 write_archive_start(&mut buf, RUN_ARCHIVE_VERSION).unwrap();
1751 write_record(&mut buf, &header()).unwrap();
1752 write_record(
1753 &mut buf,
1754 &RunRecord::ContextCheckpoint {
1755 snapshot: snapshot("plan", vec![region("conv", vec![entry("hi", 1)])]),
1756 at: 1,
1757 },
1758 )
1759 .unwrap();
1760 buf.extend_from_slice(&[0, 0, 0, 0, 0, 0, 0, 10, 1, 2]);
1762
1763 assert!(read_archive(&mut buf.as_slice()).is_err());
1765 let (version, records) = read_archive_lenient(&mut buf.as_slice()).unwrap();
1767 assert_eq!(version, RUN_ARCHIVE_VERSION);
1768 assert_eq!(records.len(), 2);
1769 let folded = fold(&records).expect("prefix starts with a Header");
1770 assert_eq!(folded.context.regions[0].entries.len(), 1);
1771 }
1772
1773 #[test]
1774 fn read_archive_lenient_still_errors_on_a_bad_preamble() {
1775 let mut bad_magic: &[u8] = b"XXXX\x00\x01";
1778 assert!(read_archive_lenient(&mut bad_magic).is_err());
1779 let mut short: &[u8] = b"LVR1";
1781 assert!(read_archive_lenient(&mut short).is_err());
1782 }
1783
1784 #[test]
1785 fn read_archive_start_errors_on_truncated_version() {
1786 let mut bytes: &[u8] = b"LVR1";
1788 let err = read_archive_start(&mut bytes).unwrap_err();
1789 assert_eq!(err.kind(), io::ErrorKind::UnexpectedEof);
1790 }
1791
1792 #[test]
1795 fn fold_requires_a_header_first() {
1796 assert!(fold(&[]).is_none());
1797 assert!(
1798 fold(&[RunRecord::StatusChanged {
1799 status: RunStatus::Complete,
1800 at: 1
1801 }])
1802 .is_none()
1803 );
1804 }
1805
1806 #[test]
1807 fn fold_reconstructs_state_from_the_journal() {
1808 let records = all_record_kinds();
1809 let folded = fold(&records).expect("has header");
1810 assert_eq!(folded.identity.machine_id, "machine-b");
1812 assert_eq!(folded.identity.world_id, "world-y");
1813 assert_eq!(folded.inference_count, 1);
1815 assert_eq!(folded.tool_call_count, 1);
1816 assert_eq!(folded.messages.len(), 1);
1818 assert_eq!(folded.messages[0].content, "another");
1819 assert_eq!(folded.context.regions[0].name, "conv");
1822 assert_eq!(folded.context.regions[0].entries.len(), 2);
1823 assert_eq!(folded.context.total_tokens, 3);
1824 assert_eq!(folded.meta.run_id, "run-1");
1825 let pending = folded.pending_batch.expect("batch never applied");
1828 assert_eq!(pending.calls[0].result.as_deref(), Some("body"));
1829 }
1830
1831 #[test]
1832 fn fold_applies_context_diffs_over_a_checkpoint() {
1833 let records = vec![
1835 header(),
1836 RunRecord::ContextCheckpoint {
1837 snapshot: snapshot("plan", vec![region("conv", vec![entry("hi", 1)])]),
1838 at: 1,
1839 },
1840 RunRecord::ContextDiff {
1841 delta: ContextDelta {
1842 stage_name: "plan".to_string(),
1843 total_tokens: 3,
1844 max_tokens: 10_000,
1845 regions: vec![RegionDelta::Append {
1846 name: "conv".to_string(),
1847 entries: vec![entry("there", 2)],
1848 current_tokens: 3,
1849 }],
1850 },
1851 at: 2,
1852 },
1853 ];
1854 let folded = fold(&records).unwrap();
1855 assert_eq!(folded.context.regions[0].entries.len(), 2);
1856 assert_eq!(folded.context.total_tokens, 3);
1857 }
1858
1859 #[test]
1860 fn fold_later_header_updates_identity_and_meta() {
1861 let mut second_meta = meta();
1863 second_meta.status = RunStatus::Running;
1864 let records = vec![
1865 header(),
1866 RunRecord::Header {
1867 identity: RunIdentity {
1868 run_id: "run-1".to_string(),
1869 machine_id: "machine-c".to_string(),
1870 world_id: "world-z".to_string(),
1871 created_at: 200,
1872 },
1873 meta: Box::new(second_meta),
1874 },
1875 ];
1876 let folded = fold(&records).unwrap();
1877 assert_eq!(folded.identity.machine_id, "machine-c");
1878 assert_eq!(folded.meta.status, RunStatus::Running);
1879 }
1880
1881 #[test]
1882 fn fold_progress_applies_meta_and_context_diff() {
1883 let mut advanced = meta();
1884 advanced.status = RunStatus::Running;
1885 advanced.iteration = 5;
1886 let records = vec![
1887 header(),
1888 RunRecord::ContextCheckpoint {
1889 snapshot: snapshot("plan", vec![region("conv", vec![entry("hi", 1)])]),
1890 at: 1,
1891 },
1892 RunRecord::Progress {
1893 meta: Box::new(advanced),
1894 delta: ContextDelta {
1895 stage_name: "plan".to_string(),
1896 total_tokens: 3,
1897 max_tokens: 10_000,
1898 regions: vec![RegionDelta::Append {
1899 name: "conv".to_string(),
1900 entries: vec![entry("there", 2)],
1901 current_tokens: 3,
1902 }],
1903 },
1904 at: 2,
1905 },
1906 ];
1907 let folded = fold(&records).unwrap();
1908 assert_eq!(folded.meta.iteration, 5);
1909 assert_eq!(folded.meta.status, RunStatus::Running);
1910 assert_eq!(folded.context.regions[0].entries.len(), 2);
1911 }
1912
1913 #[test]
1917 fn fold_carries_a_submitted_final_output_through_progress() {
1918 let mut answered = meta();
1919 answered.final_output = Some(
1920 crate::output::FinalOutput::new(
1921 "renamed two helpers",
1922 Some("markdown".to_string()),
1923 "summary".to_string(),
1924 9,
1925 )
1926 .descriptor(),
1927 );
1928 answered.output_request = Some(crate::output::OutputSpec {
1929 format: Some("a2ui".to_string()),
1930 ..Default::default()
1931 });
1932 let records = vec![
1933 header(),
1934 RunRecord::Progress {
1935 meta: Box::new(answered),
1936 delta: ContextDelta {
1937 stage_name: "summary".to_string(),
1938 total_tokens: 0,
1939 max_tokens: 10_000,
1940 regions: vec![],
1941 },
1942 at: 2,
1943 },
1944 ];
1945 let folded = fold(&records).unwrap();
1946 let output = folded.meta.final_output.expect("the answer folded through");
1947 assert_eq!(output.bytes, "renamed two helpers".len());
1950 assert_eq!(output.stage, "summary");
1951 assert_eq!(
1952 folded.meta.output_request.and_then(|s| s.format).as_deref(),
1953 Some("a2ui")
1954 );
1955 }
1956
1957 fn call(id: &str, result: Option<&str>) -> ToolCallRecord {
1960 ToolCallRecord {
1961 id: id.to_string(),
1962 name: "shell".to_string(),
1963 arguments: "{}".to_string(),
1964 result: result.map(str::to_string),
1965 thought_signature: None,
1966 }
1967 }
1968
1969 fn batch(iteration: usize, calls: Vec<ToolCallRecord>) -> RunRecord {
1970 RunRecord::ToolBatch {
1971 calls,
1972 at: 10,
1973 stage_index: 0,
1974 iteration,
1975 response: "running tools".to_string(),
1976 }
1977 }
1978
1979 fn turn_entry(call_ids: &[&str]) -> RegionEntrySnapshot {
1981 let mut e = entry("turn", 1);
1982 e.kind = crate::region::EntryKind::AssistantTurn {
1983 tool_calls: call_ids
1984 .iter()
1985 .map(|id| crate::region::SerializedToolCall {
1986 id: id.to_string(),
1987 name: "shell".to_string(),
1988 arguments: serde_json::Value::Null,
1989 thought_signature: None,
1990 })
1991 .collect(),
1992 };
1993 e
1994 }
1995
1996 #[test]
1997 fn fold_surfaces_a_pending_batch_with_merged_results() {
1998 let records = vec![
2003 header(),
2004 batch(
2005 0,
2006 vec![
2007 call("c1", None),
2008 call("c2", Some("inline")),
2009 call("c3", None),
2010 ],
2011 ),
2012 RunRecord::ToolCallDone {
2013 iteration: 0,
2014 call_id: "c1".to_string(),
2015 result: "ran".to_string(),
2016 at: 11,
2017 },
2018 ];
2019 let folded = fold(&records).unwrap();
2020 let pending = folded.pending_batch.expect("batch is pending");
2021 assert_eq!(pending.iteration, 0);
2022 assert_eq!(pending.response, "running tools");
2023 assert_eq!(pending.calls[0].result.as_deref(), Some("ran"));
2024 assert_eq!(pending.calls[1].result.as_deref(), Some("inline"));
2025 assert_eq!(pending.calls[2].result, None);
2026 assert_eq!(folded.tool_call_count, 3);
2027 }
2028
2029 #[test]
2030 fn fold_keeps_only_the_latest_batch_and_ignores_stale_done_records() {
2031 let mut advanced = meta();
2034 advanced.iteration = 1;
2035 let records = vec![
2036 header(),
2037 batch(0, vec![call("c1", None)]),
2038 RunRecord::Progress {
2039 meta: Box::new(advanced),
2040 delta: ContextDelta {
2041 stage_name: "plan".to_string(),
2042 total_tokens: 0,
2043 max_tokens: 10_000,
2044 regions: vec![],
2045 },
2046 at: 11,
2047 },
2048 batch(1, vec![call("c2", None)]),
2049 RunRecord::ToolCallDone {
2050 iteration: 0,
2051 call_id: "c1".to_string(),
2052 result: "stale".to_string(),
2053 at: 12,
2054 },
2055 RunRecord::ToolCallDone {
2056 iteration: 1,
2057 call_id: "unknown".to_string(),
2058 result: "nowhere to land".to_string(),
2059 at: 13,
2060 },
2061 ];
2062 let folded = fold(&records).unwrap();
2063 let pending = folded.pending_batch.expect("latest batch is pending");
2064 assert_eq!(pending.iteration, 1);
2065 assert_eq!(pending.calls.len(), 1);
2066 assert_eq!(pending.calls[0].id, "c2");
2067 assert_eq!(pending.calls[0].result, None, "stale/unknown dones ignored");
2068 }
2069
2070 #[test]
2071 fn fold_clears_a_batch_once_the_iteration_moves_on() {
2072 let mut advanced = meta();
2075 advanced.iteration = 1;
2076 let records = vec![
2077 header(),
2078 batch(0, vec![call("c1", Some("done"))]),
2079 RunRecord::Progress {
2080 meta: Box::new(advanced),
2081 delta: ContextDelta {
2082 stage_name: "plan".to_string(),
2083 total_tokens: 0,
2084 max_tokens: 10_000,
2085 regions: vec![],
2086 },
2087 at: 11,
2088 },
2089 ];
2090 assert_eq!(fold(&records).unwrap().pending_batch, None);
2091 }
2092
2093 #[test]
2094 fn fold_clears_a_batch_whose_turn_already_landed_in_the_window() {
2095 let records = vec![
2098 header(),
2099 batch(0, vec![call("c1", Some("done"))]),
2100 RunRecord::ContextCheckpoint {
2101 snapshot: snapshot("plan", vec![region("conv", vec![turn_entry(&["c1"])])]),
2102 at: 11,
2103 },
2104 ];
2105 assert_eq!(fold(&records).unwrap().pending_batch, None);
2106 }
2107
2108 #[test]
2109 fn context_contains_batch_matches_only_the_batch_turn() {
2110 let pending = PendingToolBatch {
2111 stage_index: 0,
2112 iteration: 0,
2113 response: String::new(),
2114 calls: vec![call("c1", None)],
2115 };
2116 let other = snapshot("plan", vec![region("conv", vec![turn_entry(&["zz"])])]);
2118 assert!(!context_contains_batch(&other, &pending));
2119 let own = snapshot(
2121 "plan",
2122 vec![region("conv", vec![turn_entry(&["c1", "c2"])])],
2123 );
2124 assert!(context_contains_batch(&own, &pending));
2125 let empty = PendingToolBatch {
2127 calls: vec![],
2128 ..pending
2129 };
2130 assert!(!context_contains_batch(&own, &empty));
2131 }
2132
2133 #[test]
2134 fn old_shape_tool_batch_json_still_parses() {
2135 let json = br#"{"ToolBatch":{"calls":[{"id":"c1","name":"shell","arguments":"{}","result":"ok"}],"at":9}}"#;
2139 let mut buf = Vec::new();
2140 buf.extend_from_slice(&(json.len() as u64).to_be_bytes());
2141 buf.extend_from_slice(json);
2142 let record = read_record(&mut buf.as_slice()).unwrap().unwrap();
2143 assert_eq!(
2144 record,
2145 RunRecord::ToolBatch {
2146 calls: vec![call("c1", Some("ok"))],
2147 at: 9,
2148 stage_index: 0,
2149 iteration: 0,
2150 response: String::new(),
2151 }
2152 );
2153 }
2154
2155 fn three_point_records() -> Vec<RunRecord> {
2159 let mut running = meta();
2160 running.status = RunStatus::Running;
2161 vec![
2162 header(),
2163 RunRecord::ContextCheckpoint {
2164 snapshot: snapshot("plan", vec![region("conv", vec![entry("first", 1)])]),
2165 at: 10,
2166 },
2167 RunRecord::ContextDiff {
2168 delta: ContextDelta {
2169 stage_name: "plan".to_string(),
2170 total_tokens: 2,
2171 max_tokens: 10_000,
2172 regions: vec![RegionDelta::Append {
2173 name: "conv".to_string(),
2174 entries: vec![entry("second", 1)],
2175 current_tokens: 2,
2176 }],
2177 },
2178 at: 20,
2179 },
2180 RunRecord::Progress {
2181 meta: Box::new(running),
2182 delta: ContextDelta {
2183 stage_name: "code".to_string(),
2184 total_tokens: 3,
2185 max_tokens: 10_000,
2186 regions: vec![RegionDelta::Append {
2187 name: "conv".to_string(),
2188 entries: vec![entry("third", 1)],
2189 current_tokens: 3,
2190 }],
2191 },
2192 at: 30,
2193 },
2194 ]
2195 }
2196
2197 #[test]
2198 fn visit_points_indexes_points_in_order_and_carries_the_running_window() {
2199 let records = three_point_records();
2200 let mut seen: Vec<(usize, i64, usize)> = Vec::new();
2201 visit_points(&records, &mut |point| {
2202 seen.push((
2203 point.index,
2204 point.at,
2205 point.context.regions[0].entries.len(),
2206 ));
2207 ControlFlow::Continue(())
2208 });
2209 assert_eq!(seen, vec![(0, 10, 1), (1, 20, 2), (2, 30, 3)]);
2211 }
2212
2213 #[test]
2216 fn visit_points_stops_at_the_first_break() {
2217 let records = three_point_records();
2218 let mut visits = 0;
2219 visit_points(&records, &mut |point| {
2220 visits += 1;
2221 if point.index == 1 {
2222 ControlFlow::Break(())
2223 } else {
2224 ControlFlow::Continue(())
2225 }
2226 });
2227 assert_eq!(
2228 visits, 2,
2229 "stopped at the breaking point, did not run the third"
2230 );
2231 }
2232
2233 #[test]
2234 fn visit_points_without_a_header_visits_nothing() {
2235 let mut visits = 0;
2236 {
2237 let mut count = |_: PointRef<'_>| {
2238 visits += 1;
2239 ControlFlow::Continue(())
2240 };
2241
2242 visit_points(&three_point_records(), &mut count);
2246 visit_points(&[], &mut count);
2249 visit_points(
2250 &[RunRecord::ContextCheckpoint {
2251 snapshot: snapshot("plan", vec![]),
2252 at: 1,
2253 }],
2254 &mut count,
2255 );
2256 }
2257 assert_eq!(visits, 3, "only the well-formed journal produced points");
2258 }
2259
2260 #[test]
2264 fn visit_points_and_replay_points_agree() {
2265 for records in [
2266 three_point_records(),
2267 vec![header()],
2268 vec![],
2269 vec![RunRecord::Message {
2270 message: MessageRecord {
2271 role: "user".to_string(),
2272 content: "x".to_string(),
2273 },
2274 at: 1,
2275 }],
2276 ] {
2277 let collected: Vec<RunPoint> = {
2278 let mut out = Vec::new();
2279 visit_points(&records, &mut |point| {
2280 out.push(RunPoint {
2281 meta: point.meta.clone(),
2282 context: point.context.clone(),
2283 at: point.at,
2284 });
2285 ControlFlow::Continue(())
2286 });
2287 out
2288 };
2289 assert_eq!(collected, replay_points(&records));
2290 }
2291 }
2292
2293 #[test]
2294 fn replay_points_requires_a_header() {
2295 assert!(replay_points(&[]).is_empty());
2296 assert!(
2297 replay_points(&[RunRecord::Message {
2298 message: MessageRecord {
2299 role: "user".to_string(),
2300 content: "x".to_string(),
2301 },
2302 at: 1,
2303 }])
2304 .is_empty()
2305 );
2306 }
2307
2308 #[test]
2309 fn replay_points_emits_a_snapshot_per_context_change() {
2310 let mut running = meta();
2313 running.status = RunStatus::Running;
2314 let records = vec![
2315 header(),
2316 RunRecord::Inference {
2317 stage: "plan".to_string(),
2318 iteration: 0,
2319 request: InferenceRequestRecord {
2320 model: "m".to_string(),
2321 system: vec![],
2322 messages: vec![],
2323 tool_names: vec![],
2324 temperature: 0.7,
2325 max_tokens: 10,
2326 },
2327 response: InferenceResponseRecord {
2328 content: "ok".to_string(),
2329 tool_calls: vec![],
2330 prompt_tokens: 1,
2331 completion_tokens: 1,
2332 cached_tokens: 0,
2333 cache_write_tokens: 0,
2334 },
2335 at: 1,
2336 },
2337 batch(0, vec![call("c1", None)]),
2338 RunRecord::ToolCallDone {
2339 iteration: 0,
2340 call_id: "c1".to_string(),
2341 result: "ran".to_string(),
2342 at: 1,
2343 },
2344 RunRecord::ContextCheckpoint {
2345 snapshot: snapshot("plan", vec![region("conv", vec![entry("hi", 1)])]),
2346 at: 2,
2347 },
2348 RunRecord::StatusChanged {
2349 status: RunStatus::Running,
2350 at: 3,
2351 },
2352 RunRecord::Progress {
2353 meta: Box::new(running),
2354 delta: ContextDelta {
2355 stage_name: "implement".to_string(),
2356 total_tokens: 3,
2357 max_tokens: 10_000,
2358 regions: vec![RegionDelta::Append {
2359 name: "conv".to_string(),
2360 entries: vec![entry("more", 2)],
2361 current_tokens: 3,
2362 }],
2363 },
2364 at: 4,
2365 },
2366 ];
2367 let points = replay_points(&records);
2368 assert_eq!(points.len(), 2, "one point per context change");
2369 assert_eq!(points[0].at, 2);
2371 assert_eq!(points[0].context.regions[0].entries.len(), 1);
2372 assert_eq!(points[1].at, 4);
2375 assert_eq!(points[1].context.regions[0].entries.len(), 2);
2376 assert_eq!(points[1].context.stage_name, "implement");
2377 assert_eq!(points[1].meta.status, RunStatus::Running);
2378 }
2379
2380 #[test]
2381 fn replay_points_handles_context_diff_and_a_later_header() {
2382 let mut relabeled = meta();
2385 relabeled.agent_name = "renamed".to_string();
2386 let records = vec![
2387 header(),
2388 RunRecord::ContextCheckpoint {
2389 snapshot: snapshot("plan", vec![region("conv", vec![entry("hi", 1)])]),
2390 at: 1,
2391 },
2392 RunRecord::Header {
2393 identity: identity(),
2394 meta: Box::new(relabeled),
2395 },
2396 RunRecord::ContextDiff {
2397 delta: ContextDelta {
2398 stage_name: "plan".to_string(),
2399 total_tokens: 3,
2400 max_tokens: 10_000,
2401 regions: vec![RegionDelta::Append {
2402 name: "conv".to_string(),
2403 entries: vec![entry("more", 2)],
2404 current_tokens: 3,
2405 }],
2406 },
2407 at: 2,
2408 },
2409 ];
2410 let points = replay_points(&records);
2411 assert_eq!(points.len(), 2); assert_eq!(points[1].context.regions[0].entries.len(), 2);
2413 assert_eq!(points[1].meta.agent_name, "renamed");
2415 }
2416
2417 #[test]
2418 fn replay_points_over_a_full_checkpoint() {
2419 let records = vec![
2421 header(),
2422 RunRecord::Checkpoint {
2423 meta: Box::new(meta()),
2424 context: snapshot("review", vec![region("conv", vec![entry("x", 4)])]),
2425 at: 9,
2426 },
2427 ];
2428 let points = replay_points(&records);
2429 assert_eq!(points.len(), 1);
2430 assert_eq!(points[0].context.stage_name, "review");
2431 assert_eq!(points[0].context.regions[0].entries[0].tokens, 4);
2432 }
2433}