1use crate::contract::backend::AgentStatus;
19use crate::contract::event::{AgentEvent, EventSender};
20
21use crate::contract::finding::Finding;
22use crate::contract::ids::{AgentId, PhaseId, RunId, TokenUsage};
23use crate::scheduler::{BackendRegistry, SchedulerConfig};
24use crate::state::{AgentResultCache, RunCheckpoint, RunStore};
25use blake3::Hasher;
26use chrono::Utc;
27use serde::{Deserialize, Serialize};
28use std::collections::HashMap;
29use std::path::Path;
30use std::sync::{Arc, RwLock};
31use std::time::{Duration, SystemTime, UNIX_EPOCH};
32use thiserror::Error;
33
34#[derive(Error, Debug)]
39pub enum JournalError {
40 #[error("run not found: {0}")]
41 RunNotFound(RunId),
42 #[error("run is not resumable (status: {status:?})")]
43 NotResumable { status: String },
44 #[error("I/O error: {0}")]
45 Io(#[from] std::io::Error),
46 #[error("serialization error: {0}")]
47 Serde(#[from] serde_json::Error),
48 #[error("journal corrupted: {0}")]
49 Corrupted(String),
50}
51
52#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
59pub struct AgentCacheKey {
60 pub hash: String,
61 pub prompt_preview: String,
63 pub phase_id: PhaseId,
64}
65
66impl AgentCacheKey {
67 pub fn new(prompt: &str, phase_id: PhaseId) -> Self {
70 let normalized = normalize_prompt(prompt);
71 let preview = if normalized.chars().count() > 80 {
72 format!("{}...", normalized.chars().take(80).collect::<String>())
73 } else {
74 normalized.clone()
75 };
76
77 let mut h = Hasher::new();
78 h.update(normalized.as_bytes());
79 h.update(b"\0");
80 h.update(&phase_id.to_le_bytes());
81
82 Self {
83 hash: h.finalize().to_hex().to_string(),
84 prompt_preview: preview,
85 phase_id,
86 }
87 }
88}
89
90fn normalize_prompt(prompt: &str) -> String {
91 prompt
92 .replace("\r\n", "\n")
93 .replace('\r', "\n")
94 .split_whitespace()
95 .collect::<Vec<_>>()
96 .join(" ")
97}
98
99pub struct JournalStore {
113 inner: Arc<RunStore>,
115 cache_index: RwLock<HashMap<String, AgentResultCache>>,
118 event_tx: Option<EventSender>,
120}
121
122impl std::fmt::Debug for JournalStore {
123 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
124 f.debug_struct("JournalStore")
125 .field("inner", &self.inner)
126 .field("cache_index_size", &self.cache_index.read().unwrap().len())
127 .field("has_event_tx", &self.event_tx.is_some())
128 .finish()
129 }
130}
131
132impl JournalStore {
133 pub fn new(run_dir: &Path) -> Result<Self, JournalError> {
136 tracing::debug!(path = %run_dir.display(), "creating journal store");
137 let inner = RunStore::new(run_dir)?;
138 Ok(Self {
139 inner,
140 cache_index: RwLock::new(HashMap::new()),
141 event_tx: None,
142 })
143 }
144
145 pub fn init_run(&self, run_id: RunId, task: &str) -> Result<(), JournalError> {
147 tracing::info!(%run_id, %task, "initializing run in journal");
148 self.inner.init_run(run_id, task)?;
149 Ok(())
150 }
151
152 pub fn init_run_with_meta(
154 &self,
155 run_id: RunId,
156 task: &str,
157 workflow_meta: serde_json::Value,
158 ) -> Result<(), JournalError> {
159 tracing::info!(
160 %run_id, %task,
161 "initializing run in journal with meta"
162 );
163 self.inner.init_run_with_meta(run_id, task, workflow_meta)?;
164 Ok(())
165 }
166
167 pub fn open(&self, run_id: RunId) -> Result<RunCheckpoint, JournalError> {
174 tracing::info!(%run_id, "opening journal for resume");
175 let checkpoint = self
176 .inner
177 .open_run(run_id)?
178 .ok_or(JournalError::RunNotFound(run_id))?;
179
180 if matches!(
181 checkpoint.status,
182 crate::state::CheckpointStatus::Completed | crate::state::CheckpointStatus::Cancelled
183 ) {
184 return Err(JournalError::NotResumable {
185 status: format!("{:?}", checkpoint.status),
186 });
187 }
188
189 let mut index = HashMap::new();
192 for (agent_id, cache) in &checkpoint.agent_results {
193 index.insert(agent_id.to_string(), cache.clone());
194 if let Some(ref hash) = cache.cache_key_hash {
195 index.insert(hash.clone(), cache.clone());
196 }
197 }
198 *self.cache_index.write().unwrap() = index;
199
200 Ok(checkpoint)
201 }
202
203 #[allow(clippy::too_many_arguments)]
209 pub fn cache_agent(
210 &self,
211 cache_key: &AgentCacheKey,
212 agent_id: AgentId,
213 phase_id: PhaseId,
214 status: AgentStatus,
215 output: serde_json::Value,
216 findings: Vec<Finding>,
217 tokens: TokenUsage,
218 ) -> Result<AgentCacheKey, JournalError> {
219 let ts = current_timestamp();
220 let cache = AgentResultCache {
221 agent_id,
222 phase_id,
223 status: status.as_str().to_string(),
224 output,
225 findings,
226 tokens: tokens.total(),
227 completed_at: ts,
228 cache_key_hash: Some(cache_key.hash.clone()),
229 description: None,
230 role: None,
231 };
232
233 {
235 let mut index = self.cache_index.write().unwrap();
236 index.insert(cache_key.hash.clone(), cache.clone());
237 index.insert(agent_id.to_string(), cache.clone());
239 }
240
241 if let Err(e) = self.inner.upsert_agent_result(&cache) {
243 tracing::warn!(%agent_id, error = %e, "failed to persist agent cache");
244 }
245
246 let event = AgentEvent::AgentDone {
248 run_id: self
249 .inner
250 .get_checkpoint()
251 .map(|c| c.run_id)
252 .unwrap_or_else(uuid::Uuid::nil),
253 agent_id,
254 status,
255 tokens,
256 elapsed_ms: 0,
257 name: None,
258 agent_seq: 0,
259 output: serde_json::Value::Null,
260 findings: Vec::new(),
261 prompt: String::new(),
262 retry_count: 0,
263 ts: Utc::now(),
264 };
265 self.inner.append_event(&event)?;
266
267 if let Some(ref tx) = self.event_tx {
269 let _ = tx.send(event);
270 }
271
272 Ok(cache_key.clone())
273 }
274
275 #[allow(clippy::too_many_arguments)]
283 pub fn record_result(
284 &self,
285 cache_key: &AgentCacheKey,
286 agent_id: AgentId,
287 phase_id: PhaseId,
288 status: AgentStatus,
289 output: serde_json::Value,
290 findings: Vec<Finding>,
291 tokens: TokenUsage,
292 ) {
293 let cache = AgentResultCache {
294 agent_id,
295 phase_id,
296 status: status.as_str().to_string(),
297 output,
298 findings,
299 tokens: tokens.total(),
300 completed_at: current_timestamp(),
301 cache_key_hash: Some(cache_key.hash.clone()),
302 description: None,
303 role: None,
304 };
305
306 {
307 let mut index = self.cache_index.write().unwrap();
308 index.insert(cache_key.hash.clone(), cache.clone());
309 index.insert(agent_id.to_string(), cache.clone());
310 }
311
312 if let Err(e) = self.inner.upsert_agent_result(&cache) {
313 tracing::warn!(%agent_id, error = %e, "failed to persist agent result");
314 }
315 }
316
317 pub fn store(&self) -> Arc<RunStore> {
321 self.inner.clone()
322 }
323
324 pub fn append_event(&self, event: &AgentEvent) -> Result<(), JournalError> {
326 self.inner.append_event(event)?;
327 Ok(())
328 }
329
330 pub fn has_completed(&self, cache_key: &AgentCacheKey) -> bool {
333 let index = self.cache_index.read().unwrap();
334 index.contains_key(&cache_key.hash)
335 }
336
337 pub fn get_cached(&self, cache_key: &AgentCacheKey) -> Option<AgentResultCache> {
340 let index = self.cache_index.read().unwrap();
341 index.get(&cache_key.hash).cloned()
342 }
343
344 pub fn completed_keys(&self) -> Vec<AgentCacheKey> {
347 let index = self.cache_index.read().unwrap();
348 index
349 .keys()
350 .map(|k| AgentCacheKey {
351 hash: k.clone(),
352 prompt_preview: String::new(),
353 phase_id: 0,
354 })
355 .collect()
356 }
357
358 pub fn get_checkpoint(&self) -> Option<RunCheckpoint> {
360 self.inner.get_checkpoint()
361 }
362
363 pub fn flush(&self) -> Result<(), JournalError> {
365 Ok(())
367 }
368
369 pub fn cancel(&self) -> Result<(), JournalError> {
371 self.inner.cancel()?;
372 Ok(())
373 }
374}
375
376pub struct CompositeJournalCallback {
382 callbacks: Vec<Arc<dyn crate::scheduler::JournalCallback>>,
383}
384
385impl CompositeJournalCallback {
386 pub fn new(callbacks: Vec<Arc<dyn crate::scheduler::JournalCallback>>) -> Self {
387 Self { callbacks }
388 }
389}
390
391#[async_trait::async_trait]
392impl crate::scheduler::JournalCallback for CompositeJournalCallback {
393 async fn on_agent_done(
394 &self,
395 agent_id: AgentId,
396 phase_id: PhaseId,
397 status: AgentStatus,
398 output: serde_json::Value,
399 tokens: TokenUsage,
400 ) {
401 for cb in &self.callbacks {
402 cb.on_agent_done(agent_id, phase_id, status.clone(), output.clone(), tokens)
403 .await;
404 }
405 }
406}
407
408#[async_trait::async_trait]
409impl crate::scheduler::JournalCallback for JournalStore {
410 async fn on_agent_done(
411 &self,
412 agent_id: AgentId,
413 phase_id: PhaseId,
414 status: AgentStatus,
415 output: serde_json::Value,
416 tokens: TokenUsage,
417 ) {
418 let ts = current_timestamp();
419
420 let existing = {
425 let index = self.cache_index.read().unwrap();
426 index.get(&agent_id.to_string()).cloned()
427 };
428
429 let cache = AgentResultCache {
430 agent_id,
431 phase_id: existing.as_ref().map(|c| c.phase_id).unwrap_or(phase_id),
432 status: status.as_str().to_string(),
433 output: existing
434 .as_ref()
435 .filter(|c| !c.output.is_null())
436 .map(|c| c.output.clone())
437 .unwrap_or(output),
438 findings: existing
439 .as_ref()
440 .filter(|c| !c.findings.is_empty())
441 .map(|c| c.findings.clone())
442 .unwrap_or_default(),
443 tokens: tokens.total(),
444 completed_at: ts,
445 cache_key_hash: existing.as_ref().and_then(|c| c.cache_key_hash.clone()),
446 description: existing.as_ref().and_then(|c| c.description.clone()),
447 role: existing.as_ref().and_then(|c| c.role.clone()),
448 };
449
450 {
453 let mut index = self.cache_index.write().unwrap();
454 index.insert(agent_id.to_string(), cache.clone());
455 if let Some(ref hash) = cache.cache_key_hash {
456 index.insert(hash.clone(), cache.clone());
457 }
458 }
459
460 if let Err(e) = self.inner.upsert_agent_result(&cache) {
462 tracing::warn!(%agent_id, error = %e, "failed to persist agent result from callback");
463 }
464 }
465}
466
467#[derive(Debug)]
473pub struct ResumeContext {
474 pub run_id: RunId,
475 pub checkpoint: RunCheckpoint,
476 pub journal: Arc<JournalStore>,
477 pub scheduler_config: SchedulerConfig,
478 pub backend_registry: BackendRegistry,
479}
480
481#[derive(Debug, Clone)]
483pub enum RunCreationMode {
484 New { task: String },
486 Resume { run_id: RunId, run_dir_name: String },
488 Auto { task: String },
490}
491
492impl RunCreationMode {
493 pub fn resolve(
496 self,
497 journal_dir: &Path,
498 ) -> Result<(RunId, Option<RunCheckpoint>), JournalError> {
499 match self {
500 RunCreationMode::New { task: _ } => {
501 let run_id = uuid::Uuid::now_v7();
502 Ok((run_id, None))
503 }
504 RunCreationMode::Resume {
505 run_id,
506 run_dir_name,
507 } => {
508 let store = JournalStore::new(&journal_dir.join(&run_dir_name))?;
509 let checkpoint = store.open(run_id)?;
510 Ok((run_id, Some(checkpoint)))
511 }
512 RunCreationMode::Auto { task: _ } => {
513 let run_dirs = crate::state::list_runs(journal_dir)?;
515 for dir_name in run_dirs.iter().rev() {
516 let checkpoint_path = journal_dir.join(dir_name).join("checkpoint.json");
517 if let Ok(content) = std::fs::read_to_string(&checkpoint_path) {
518 if let Ok(checkpoint) = serde_json::from_str::<RunCheckpoint>(&content) {
519 if matches!(checkpoint.status, crate::state::CheckpointStatus::Running)
520 {
521 let run_id = checkpoint.run_id;
522 return Ok((run_id, Some(checkpoint)));
523 }
524 }
525 }
526 }
527 let run_id = uuid::Uuid::now_v7();
529 Ok((run_id, None))
530 }
531 }
532 }
533}
534
535pub fn gc_runs(journal_dir: &Path, older_than: Duration) -> Result<usize, JournalError> {
547 let run_dirs = crate::state::list_runs(journal_dir)?;
548 let cutoff = current_timestamp().saturating_sub(older_than.as_secs());
549
550 tracing::debug!("GC: scanning {} runs", run_dirs.len());
551 let mut cleaned = 0;
552 for dir_name in &run_dirs {
553 let run_dir = journal_dir.join(dir_name);
554 let checkpoint_path = run_dir.join("checkpoint.json");
556 if !checkpoint_path.exists() {
557 continue;
558 }
559
560 let content = std::fs::read_to_string(&checkpoint_path)?;
561 let checkpoint: RunCheckpoint = serde_json::from_str(&content)?;
562
563 let is_old = checkpoint.updated_at < cutoff;
564 let is_terminal = matches!(
565 checkpoint.status,
566 crate::state::CheckpointStatus::Completed
567 | crate::state::CheckpointStatus::Cancelled
568 | crate::state::CheckpointStatus::Failed
569 );
570
571 if is_old && is_terminal {
572 tracing::info!(dir = %dir_name, "GC: removing old terminal run");
573 std::fs::remove_dir_all(&run_dir)?;
574 cleaned += 1;
575 }
576 }
577
578 Ok(cleaned)
579}
580
581fn current_timestamp() -> u64 {
582 SystemTime::now()
583 .duration_since(UNIX_EPOCH)
584 .map(|d| d.as_secs())
585 .unwrap_or(0)
586}
587
588#[cfg(test)]
593mod tests {
594 use super::*;
595 use tempfile::tempdir;
596
597 #[test]
599 fn test_journal_lifecycle() {
600 let dir = tempdir().unwrap();
601 let run_id = uuid::Uuid::now_v7();
602 let journal = JournalStore::new(dir.path()).unwrap();
603
604 journal.init_run(run_id, "Test task").unwrap();
606 let cp = journal.get_checkpoint().unwrap();
607 assert_eq!(cp.status, crate::state::CheckpointStatus::Running);
608 assert_eq!(cp.task, "Test task");
609
610 let agent_id = uuid::Uuid::now_v7();
612 let key = AgentCacheKey::new("test prompt", 1);
613 journal
614 .cache_agent(
615 &key,
616 agent_id,
617 1,
618 AgentStatus::Ok,
619 serde_json::json!({"result": "ok"}),
620 vec![],
621 TokenUsage {
622 input: 100,
623 output: 50,
624 cache_read: 0,
625 cache_write: 0,
626 },
627 )
628 .unwrap();
629
630 assert!(journal.has_completed(&key));
632 let cached = journal.get_cached(&key).unwrap();
633 assert_eq!(cached.output, serde_json::json!({"result": "ok"}));
634 assert_eq!(cached.tokens, 150);
635
636 journal.cancel().unwrap();
638 let cp = journal.get_checkpoint().unwrap();
639 assert_eq!(cp.status, crate::state::CheckpointStatus::Cancelled);
640 }
641
642 #[test]
644 fn test_cache_key_uniqueness() {
645 let k1 = AgentCacheKey::new("prompt A", 1);
646 let k2 = AgentCacheKey::new("prompt B", 1);
647 assert_ne!(k1.hash, k2.hash);
648
649 let k4 = AgentCacheKey::new("prompt A", 2);
651 assert_ne!(k1.hash, k4.hash);
652
653 let k5 = AgentCacheKey::new(" prompt \r\nA ", 1);
655 assert_eq!(k1.hash, k5.hash);
656 }
657
658 #[test]
660 fn test_resume_skip_cached() {
661 let dir = tempdir().unwrap();
662 let run_id = uuid::Uuid::now_v7();
663 let journal = JournalStore::new(dir.path()).unwrap();
664 journal.init_run(run_id, "Three agent test").unwrap();
665
666 let k1 = AgentCacheKey::new("task 1", 1);
668 let k2 = AgentCacheKey::new("task 2", 1);
669 let k3 = AgentCacheKey::new("task 3", 1);
670
671 journal
672 .cache_agent(
673 &k1,
674 uuid::Uuid::now_v7(),
675 1,
676 AgentStatus::Ok,
677 serde_json::json!({"done": 1}),
678 vec![],
679 TokenUsage {
680 input: 10,
681 output: 5,
682 cache_read: 0,
683 cache_write: 0,
684 },
685 )
686 .unwrap();
687 journal
688 .cache_agent(
689 &k2,
690 uuid::Uuid::now_v7(),
691 1,
692 AgentStatus::Ok,
693 serde_json::json!({"done": 2}),
694 vec![],
695 TokenUsage {
696 input: 10,
697 output: 5,
698 cache_read: 0,
699 cache_write: 0,
700 },
701 )
702 .unwrap();
703
704 assert!(journal.has_completed(&k1));
706 assert!(journal.has_completed(&k2));
707 assert!(!journal.has_completed(&k3));
708
709 assert!(journal.get_cached(&k3).is_none());
711 }
712
713 #[test]
715 fn test_journal_crash_recovery() {
716 let dir = tempdir().unwrap();
717 let run_id = uuid::Uuid::now_v7();
718
719 {
721 let j = JournalStore::new(dir.path()).unwrap();
722 j.init_run(run_id, "Crash test").unwrap();
723 let key = AgentCacheKey::new("important work", 0);
724 j.cache_agent(
725 &key,
726 uuid::Uuid::now_v7(),
727 0,
728 AgentStatus::Ok,
729 serde_json::json!({"survived": true}),
730 vec![],
731 TokenUsage {
732 input: 1,
733 output: 1,
734 cache_read: 0,
735 cache_write: 0,
736 },
737 )
738 .unwrap();
739 } {
743 let j2 = JournalStore::new(dir.path()).unwrap();
744 let cp = j2.open(run_id).unwrap();
745 assert_eq!(cp.status, crate::state::CheckpointStatus::Running);
746 assert!(!cp.agent_results.is_empty());
747
748 let key = AgentCacheKey::new("important work", 0);
749 let cached = j2.get_cached(&key).unwrap();
750 assert_eq!(cached.output, serde_json::json!({"survived": true}));
751 }
752 }
753
754 #[test]
756 fn test_gc_older_than() {
757 let dir = tempdir().unwrap();
758 let run_dir = dir.path().join("runs");
759 std::fs::create_dir_all(&run_dir).unwrap();
760
761 let run_id = uuid::Uuid::now_v7();
763 let journal = JournalStore::new(&run_dir.join(run_id.to_string())).unwrap();
764 journal.init_run(run_id, "GC me").unwrap();
765
766 if let Some(mut cp) = journal.get_checkpoint() {
768 cp.status = crate::state::CheckpointStatus::Completed;
769 cp.updated_at = 1000; let _ = journal.inner.save_checkpoint(&cp);
771 }
772
773 let cleaned = gc_runs(&run_dir, Duration::from_secs(3600)).unwrap();
775 assert_eq!(cleaned, 1);
776 }
777
778 fn read_checkpoint_status_for(run_dir: &std::path::Path, agent_id: AgentId) -> Option<String> {
790 let cp_path = run_dir.join("checkpoint.json");
791 let content = std::fs::read_to_string(&cp_path).ok()?;
792 let raw: serde_json::Value = serde_json::from_str(&content).ok()?;
793 let ar = raw.get("agent_results")?.as_object()?;
794 for (_k, v) in ar {
795 if v.get("agent_id").and_then(|id| id.as_str()) == Some(&agent_id.to_string()) {
796 return v.get("status").and_then(|s| s.as_str()).map(String::from);
797 }
798 }
799 None
800 }
801
802 fn sample_token_usage(input: u64, output: u64) -> TokenUsage {
803 TokenUsage {
804 input,
805 output,
806 cache_read: 0,
807 cache_write: 0,
808 }
809 }
810
811 #[test]
812 fn cache_agent_persists_snake_case_status_for_each_variant() {
813 let dir = tempdir().unwrap();
818 let run_id = uuid::Uuid::now_v7();
819 let journal = JournalStore::new(dir.path()).unwrap();
820 journal.init_run(run_id, "cache_agent F5").unwrap();
821
822 let cases: Vec<(AgentStatus, &str)> = vec![
823 (AgentStatus::Ok, "ok"),
824 (AgentStatus::Error, "error"),
825 (AgentStatus::Cancelled, "cancelled"),
826 (AgentStatus::TimedOut, "timed_out"),
827 ];
828 for (status, expected) in &cases {
829 let agent_id = uuid::Uuid::now_v7();
830 let key = AgentCacheKey::new("prompt", 1);
831 journal
832 .cache_agent(
833 &key,
834 agent_id,
835 1,
836 status.clone(),
837 serde_json::json!({"v": 1}),
838 vec![],
839 sample_token_usage(10, 5),
840 )
841 .unwrap();
842
843 let persisted = read_checkpoint_status_for(dir.path(), agent_id)
844 .unwrap_or_else(|| panic!("status missing on disk for {status:?}"));
845 assert_eq!(
846 persisted, *expected,
847 "cache_agent({status:?}) must persist status={expected:?} (snake_case); \
848 got {persisted:?}. Reverting to Debug formatting would yield \"timedout\" \
849 for TimedOut and break the on-disk contract."
850 );
851 }
852 }
853
854 #[test]
855 fn cache_agent_timed_out_persists_with_underscore_not_collapsed() {
856 let dir = tempdir().unwrap();
860 let run_id = uuid::Uuid::now_v7();
861 let journal = JournalStore::new(dir.path()).unwrap();
862 journal.init_run(run_id, "timed-out guard").unwrap();
863
864 let agent_id = uuid::Uuid::now_v7();
865 let key = AgentCacheKey::new("p", 0);
866 journal
867 .cache_agent(
868 &key,
869 agent_id,
870 0,
871 AgentStatus::TimedOut,
872 serde_json::json!(null),
873 vec![],
874 sample_token_usage(1, 2),
875 )
876 .unwrap();
877
878 let persisted = read_checkpoint_status_for(dir.path(), agent_id).expect("status on disk");
879 assert_eq!(
880 persisted, "timed_out",
881 "cache_agent(TimedOut) must persist \"timed_out\"; got {persisted:?}"
882 );
883 assert_ne!(
884 persisted, "timedout",
885 "cache_agent(TimedOut) must NOT collapse to Debug-lowercased \"timedout\""
886 );
887 }
888
889 #[test]
890 fn record_result_persists_snake_case_status_for_each_variant() {
891 let dir = tempdir().unwrap();
894 let run_id = uuid::Uuid::now_v7();
895 let journal = JournalStore::new(dir.path()).unwrap();
896 journal.init_run(run_id, "record_result F5").unwrap();
897
898 let cases: Vec<(AgentStatus, &str)> = vec![
899 (AgentStatus::Ok, "ok"),
900 (AgentStatus::Error, "error"),
901 (AgentStatus::Cancelled, "cancelled"),
902 (AgentStatus::TimedOut, "timed_out"),
903 ];
904 for (status, expected) in &cases {
905 let agent_id = uuid::Uuid::now_v7();
906 let key = AgentCacheKey::new("p", 1);
907 journal.record_result(
908 &key,
909 agent_id,
910 1,
911 status.clone(),
912 serde_json::json!({"r": 1}),
913 vec![],
914 sample_token_usage(2, 3),
915 );
916
917 let persisted = read_checkpoint_status_for(dir.path(), agent_id)
918 .unwrap_or_else(|| panic!("status missing on disk for {status:?}"));
919 assert_eq!(
920 persisted, *expected,
921 "record_result({status:?}) must persist status={expected:?}; got {persisted:?}"
922 );
923 }
924 }
925
926 #[test]
927 fn record_result_timed_out_persists_with_underscore() {
928 let dir = tempdir().unwrap();
930 let run_id = uuid::Uuid::now_v7();
931 let journal = JournalStore::new(dir.path()).unwrap();
932 journal.init_run(run_id, "record_result timed-out").unwrap();
933
934 let agent_id = uuid::Uuid::now_v7();
935 let key = AgentCacheKey::new("p", 0);
936 journal.record_result(
937 &key,
938 agent_id,
939 0,
940 AgentStatus::TimedOut,
941 serde_json::json!(null),
942 vec![],
943 sample_token_usage(0, 0),
944 );
945
946 let persisted = read_checkpoint_status_for(dir.path(), agent_id).expect("status on disk");
947 assert_eq!(persisted, "timed_out");
948 assert_ne!(persisted, "timedout");
949 }
950
951 #[tokio::test]
952 async fn journal_callback_on_agent_done_persists_snake_case_status() {
953 let dir = tempdir().unwrap();
957 let run_id = uuid::Uuid::now_v7();
958 let journal = std::sync::Arc::new(JournalStore::new(dir.path()).unwrap());
959 journal.init_run(run_id, "callback F5").unwrap();
960
961 let cases: Vec<(AgentStatus, &str)> = vec![
962 (AgentStatus::Ok, "ok"),
963 (AgentStatus::Error, "error"),
964 (AgentStatus::Cancelled, "cancelled"),
965 (AgentStatus::TimedOut, "timed_out"),
966 ];
967 for (status, expected) in &cases {
968 let agent_id = uuid::Uuid::now_v7();
969 use crate::scheduler::JournalCallback;
970 journal
971 .on_agent_done(
972 agent_id,
973 1,
974 status.clone(),
975 serde_json::json!({}),
976 sample_token_usage(4, 6),
977 )
978 .await;
979
980 let persisted = read_checkpoint_status_for(dir.path(), agent_id)
981 .unwrap_or_else(|| panic!("status missing on disk for {status:?}"));
982 assert_eq!(
983 persisted, *expected,
984 "JournalCallback::on_agent_done({status:?}) must persist status={expected:?}; \
985 got {persisted:?}"
986 );
987 }
988 }
989
990 #[test]
991 fn record_result_then_reopen_uses_snake_case_status() {
992 let dir = tempdir().unwrap();
995 let run_id = uuid::Uuid::now_v7();
996 let journal = JournalStore::new(dir.path()).unwrap();
997 journal.init_run(run_id, "reopen F5").unwrap();
998
999 let agent_id = uuid::Uuid::now_v7();
1000 let key = AgentCacheKey::new("reopen prompt", 1);
1001 journal.record_result(
1002 &key,
1003 agent_id,
1004 1,
1005 AgentStatus::Cancelled,
1006 serde_json::json!({"result": "ok"}),
1007 vec![],
1008 sample_token_usage(7, 11),
1009 );
1010 drop(journal);
1011
1012 let j2 = JournalStore::new(dir.path()).unwrap();
1013 let cp = j2.open(run_id).expect("open after drop");
1014 let cached = cp
1015 .agent_results
1016 .get(&agent_id)
1017 .expect("entry survives reopen");
1018 assert_eq!(
1019 cached.status, "cancelled",
1020 "snake_case status must round-trip through close+reopen"
1021 );
1022 assert_eq!(cached.tokens, 18);
1023 }
1024
1025 #[test]
1026 fn cache_agent_persists_snake_case_status_to_event_log() {
1027 let dir = tempdir().unwrap();
1031 let run_id = uuid::Uuid::now_v7();
1032 let journal = JournalStore::new(dir.path()).unwrap();
1033 journal.init_run(run_id, "event log F5").unwrap();
1034
1035 let agent_id = uuid::Uuid::now_v7();
1036 let key = AgentCacheKey::new("p", 1);
1037 journal
1038 .cache_agent(
1039 &key,
1040 agent_id,
1041 1,
1042 AgentStatus::TimedOut,
1043 serde_json::json!(null),
1044 vec![],
1045 sample_token_usage(1, 1),
1046 )
1047 .unwrap();
1048
1049 let log = journal.store().get_event_log().expect("read events.jsonl");
1053 let agent_done = log
1054 .iter()
1055 .find_map(|e| match e {
1056 AgentEvent::AgentDone {
1057 agent_id: id,
1058 status,
1059 ..
1060 } if id == &agent_id => Some(status.clone()),
1061 _ => None,
1062 })
1063 .expect("AgentDone event in log");
1064 assert!(matches!(agent_done, AgentStatus::TimedOut));
1068 }
1069
1070 #[tokio::test]
1075 async fn on_agent_done_preserves_cache_key_hash_from_record_result() {
1076 let dir = tempdir().unwrap();
1080 let run_id = uuid::Uuid::now_v7();
1081 let journal = std::sync::Arc::new(JournalStore::new(dir.path()).unwrap());
1082 journal.init_run(run_id, "hash preservation").unwrap();
1083
1084 let agent_id = uuid::Uuid::now_v7();
1085 let key = AgentCacheKey::new("preserve me", 1);
1086
1087 journal.record_result(
1089 &key,
1090 agent_id,
1091 1,
1092 AgentStatus::Ok,
1093 serde_json::json!({"answer": 42}),
1094 vec![],
1095 sample_token_usage(10, 5),
1096 );
1097
1098 use crate::scheduler::JournalCallback;
1100 journal
1101 .on_agent_done(
1102 agent_id,
1103 1,
1104 AgentStatus::Ok,
1105 serde_json::json!({}),
1106 sample_token_usage(10, 5),
1107 )
1108 .await;
1109
1110 assert!(
1112 journal.has_completed(&key),
1113 "cache_key_hash must survive on_agent_done"
1114 );
1115
1116 drop(journal);
1118 let j2 = JournalStore::new(dir.path()).unwrap();
1119 j2.open(run_id).expect("reopen");
1120 assert!(
1121 j2.has_completed(&key),
1122 "cache_key_hash must survive reopen after on_agent_done"
1123 );
1124 }
1125
1126 #[tokio::test]
1127 async fn on_agent_done_preserves_cache_key_hash_from_cache_agent() {
1128 let dir = tempdir().unwrap();
1130 let run_id = uuid::Uuid::now_v7();
1131 let journal = std::sync::Arc::new(JournalStore::new(dir.path()).unwrap());
1132 journal.init_run(run_id, "hash preservation 2").unwrap();
1133
1134 let agent_id = uuid::Uuid::now_v7();
1135 let key = AgentCacheKey::new("preserve me 2", 0);
1136
1137 journal
1138 .cache_agent(
1139 &key,
1140 agent_id,
1141 0,
1142 AgentStatus::Ok,
1143 serde_json::json!({"r": 1}),
1144 vec![],
1145 sample_token_usage(1, 1),
1146 )
1147 .unwrap();
1148
1149 use crate::scheduler::JournalCallback;
1150 journal
1151 .on_agent_done(
1152 agent_id,
1153 0,
1154 AgentStatus::Ok,
1155 serde_json::json!({}),
1156 sample_token_usage(1, 1),
1157 )
1158 .await;
1159
1160 assert!(
1161 journal.has_completed(&key),
1162 "cache_key_hash must survive on_agent_done after cache_agent"
1163 );
1164 }
1165}