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
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 pub fn reset_status_to_running(&self) -> Result<(), JournalError> {
378 self.inner.reset_status_to_running()?;
379 Ok(())
380 }
381}
382
383pub struct CompositeJournalCallback {
389 callbacks: Vec<Arc<dyn crate::scheduler::JournalCallback>>,
390}
391
392impl CompositeJournalCallback {
393 pub fn new(callbacks: Vec<Arc<dyn crate::scheduler::JournalCallback>>) -> Self {
394 Self { callbacks }
395 }
396}
397
398#[async_trait::async_trait]
399impl crate::scheduler::JournalCallback for CompositeJournalCallback {
400 async fn on_agent_done(
401 &self,
402 agent_id: AgentId,
403 phase_id: PhaseId,
404 status: AgentStatus,
405 output: serde_json::Value,
406 tokens: TokenUsage,
407 ) {
408 for cb in &self.callbacks {
409 cb.on_agent_done(agent_id, phase_id, status.clone(), output.clone(), tokens)
410 .await;
411 }
412 }
413}
414
415#[async_trait::async_trait]
416impl crate::scheduler::JournalCallback for JournalStore {
417 async fn on_agent_done(
418 &self,
419 agent_id: AgentId,
420 phase_id: PhaseId,
421 status: AgentStatus,
422 output: serde_json::Value,
423 tokens: TokenUsage,
424 ) {
425 let ts = current_timestamp();
426
427 let existing = {
432 let index = self.cache_index.read().unwrap();
433 index.get(&agent_id.to_string()).cloned()
434 };
435
436 let cache = AgentResultCache {
437 agent_id,
438 phase_id: existing.as_ref().map(|c| c.phase_id).unwrap_or(phase_id),
439 status: status.as_str().to_string(),
440 output: existing
441 .as_ref()
442 .filter(|c| !c.output.is_null())
443 .map(|c| c.output.clone())
444 .unwrap_or(output),
445 findings: existing
446 .as_ref()
447 .filter(|c| !c.findings.is_empty())
448 .map(|c| c.findings.clone())
449 .unwrap_or_default(),
450 tokens: tokens.total(),
451 completed_at: ts,
452 cache_key_hash: existing.as_ref().and_then(|c| c.cache_key_hash.clone()),
453 description: existing.as_ref().and_then(|c| c.description.clone()),
454 role: existing.as_ref().and_then(|c| c.role.clone()),
455 };
456
457 {
460 let mut index = self.cache_index.write().unwrap();
461 index.insert(agent_id.to_string(), cache.clone());
462 if let Some(ref hash) = cache.cache_key_hash {
463 index.insert(hash.clone(), cache.clone());
464 }
465 }
466
467 if let Err(e) = self.inner.upsert_agent_result(&cache) {
469 tracing::warn!(%agent_id, error = %e, "failed to persist agent result from callback");
470 }
471 }
472}
473
474#[derive(Debug)]
480pub struct ResumeContext {
481 pub run_id: RunId,
482 pub checkpoint: RunCheckpoint,
483 pub journal: Arc<JournalStore>,
484 pub scheduler_config: SchedulerConfig,
485 pub backend_registry: BackendRegistry,
486}
487
488#[derive(Debug, Clone)]
490pub enum RunCreationMode {
491 New { task: String },
493 Resume { run_id: RunId, run_dir_name: String },
495 Auto { task: String },
497}
498
499impl RunCreationMode {
500 pub fn resolve(
503 self,
504 journal_dir: &Path,
505 ) -> Result<(RunId, Option<RunCheckpoint>), JournalError> {
506 match self {
507 RunCreationMode::New { task: _ } => {
508 let run_id = uuid::Uuid::now_v7();
509 Ok((run_id, None))
510 }
511 RunCreationMode::Resume {
512 run_id,
513 run_dir_name,
514 } => {
515 let store = JournalStore::new(&journal_dir.join(&run_dir_name))?;
516 let checkpoint = store.open(run_id)?;
517 Ok((run_id, Some(checkpoint)))
518 }
519 RunCreationMode::Auto { task: _ } => {
520 let run_dirs = crate::state::list_runs(journal_dir)?;
522 for dir_name in run_dirs.iter().rev() {
523 let checkpoint_path = journal_dir.join(dir_name).join("checkpoint.json");
524 if let Ok(content) = std::fs::read_to_string(&checkpoint_path) {
525 if let Ok(checkpoint) = serde_json::from_str::<RunCheckpoint>(&content) {
526 if matches!(checkpoint.status, crate::state::CheckpointStatus::Running)
527 {
528 let run_id = checkpoint.run_id;
529 return Ok((run_id, Some(checkpoint)));
530 }
531 }
532 }
533 }
534 let run_id = uuid::Uuid::now_v7();
536 Ok((run_id, None))
537 }
538 }
539 }
540}
541
542pub fn gc_runs(journal_dir: &Path, older_than: Duration) -> Result<usize, JournalError> {
554 let run_dirs = crate::state::list_runs(journal_dir)?;
555 let cutoff = current_timestamp().saturating_sub(older_than.as_secs());
556
557 tracing::debug!("GC: scanning {} runs", run_dirs.len());
558 let mut cleaned = 0;
559 for dir_name in &run_dirs {
560 let run_dir = journal_dir.join(dir_name);
561 let checkpoint_path = run_dir.join("checkpoint.json");
563 if !checkpoint_path.exists() {
564 continue;
565 }
566
567 let content = std::fs::read_to_string(&checkpoint_path)?;
568 let checkpoint: RunCheckpoint = serde_json::from_str(&content)?;
569
570 let is_old = checkpoint.updated_at < cutoff;
571 let is_terminal = matches!(
572 checkpoint.status,
573 crate::state::CheckpointStatus::Completed
574 | crate::state::CheckpointStatus::Cancelled
575 | crate::state::CheckpointStatus::Failed
576 );
577
578 if is_old && is_terminal {
579 tracing::info!(dir = %dir_name, "GC: removing old terminal run");
580 std::fs::remove_dir_all(&run_dir)?;
581 cleaned += 1;
582 }
583 }
584
585 Ok(cleaned)
586}
587
588fn current_timestamp() -> u64 {
589 SystemTime::now()
590 .duration_since(UNIX_EPOCH)
591 .map(|d| d.as_secs())
592 .unwrap_or(0)
593}
594
595#[cfg(test)]
600mod tests {
601 use super::*;
602 use tempfile::tempdir;
603
604 #[test]
606 fn test_journal_lifecycle() {
607 let dir = tempdir().unwrap();
608 let run_id = uuid::Uuid::now_v7();
609 let journal = JournalStore::new(dir.path()).unwrap();
610
611 journal.init_run(run_id, "Test task").unwrap();
613 let cp = journal.get_checkpoint().unwrap();
614 assert_eq!(cp.status, crate::state::CheckpointStatus::Running);
615 assert_eq!(cp.task, "Test task");
616
617 let agent_id = uuid::Uuid::now_v7();
619 let key = AgentCacheKey::new("test prompt", 1);
620 journal
621 .cache_agent(
622 &key,
623 agent_id,
624 1,
625 AgentStatus::Ok,
626 serde_json::json!({"result": "ok"}),
627 vec![],
628 TokenUsage {
629 input: 100,
630 output: 50,
631 cache_read: 0,
632 cache_write: 0,
633 },
634 )
635 .unwrap();
636
637 assert!(journal.has_completed(&key));
639 let cached = journal.get_cached(&key).unwrap();
640 assert_eq!(cached.output, serde_json::json!({"result": "ok"}));
641 assert_eq!(cached.tokens, 150);
642
643 journal.cancel().unwrap();
645 let cp = journal.get_checkpoint().unwrap();
646 assert_eq!(cp.status, crate::state::CheckpointStatus::Cancelled);
647 }
648
649 #[test]
651 fn test_cache_key_uniqueness() {
652 let k1 = AgentCacheKey::new("prompt A", 1);
653 let k2 = AgentCacheKey::new("prompt B", 1);
654 assert_ne!(k1.hash, k2.hash);
655
656 let k4 = AgentCacheKey::new("prompt A", 2);
658 assert_ne!(k1.hash, k4.hash);
659
660 let k5 = AgentCacheKey::new(" prompt \r\nA ", 1);
662 assert_eq!(k1.hash, k5.hash);
663 }
664
665 #[test]
667 fn test_resume_skip_cached() {
668 let dir = tempdir().unwrap();
669 let run_id = uuid::Uuid::now_v7();
670 let journal = JournalStore::new(dir.path()).unwrap();
671 journal.init_run(run_id, "Three agent test").unwrap();
672
673 let k1 = AgentCacheKey::new("task 1", 1);
675 let k2 = AgentCacheKey::new("task 2", 1);
676 let k3 = AgentCacheKey::new("task 3", 1);
677
678 journal
679 .cache_agent(
680 &k1,
681 uuid::Uuid::now_v7(),
682 1,
683 AgentStatus::Ok,
684 serde_json::json!({"done": 1}),
685 vec![],
686 TokenUsage {
687 input: 10,
688 output: 5,
689 cache_read: 0,
690 cache_write: 0,
691 },
692 )
693 .unwrap();
694 journal
695 .cache_agent(
696 &k2,
697 uuid::Uuid::now_v7(),
698 1,
699 AgentStatus::Ok,
700 serde_json::json!({"done": 2}),
701 vec![],
702 TokenUsage {
703 input: 10,
704 output: 5,
705 cache_read: 0,
706 cache_write: 0,
707 },
708 )
709 .unwrap();
710
711 assert!(journal.has_completed(&k1));
713 assert!(journal.has_completed(&k2));
714 assert!(!journal.has_completed(&k3));
715
716 assert!(journal.get_cached(&k3).is_none());
718 }
719
720 #[test]
722 fn test_journal_crash_recovery() {
723 let dir = tempdir().unwrap();
724 let run_id = uuid::Uuid::now_v7();
725
726 {
728 let j = JournalStore::new(dir.path()).unwrap();
729 j.init_run(run_id, "Crash test").unwrap();
730 let key = AgentCacheKey::new("important work", 0);
731 j.cache_agent(
732 &key,
733 uuid::Uuid::now_v7(),
734 0,
735 AgentStatus::Ok,
736 serde_json::json!({"survived": true}),
737 vec![],
738 TokenUsage {
739 input: 1,
740 output: 1,
741 cache_read: 0,
742 cache_write: 0,
743 },
744 )
745 .unwrap();
746 } {
750 let j2 = JournalStore::new(dir.path()).unwrap();
751 let cp = j2.open(run_id).unwrap();
752 assert_eq!(cp.status, crate::state::CheckpointStatus::Running);
753 assert!(!cp.agent_results.is_empty());
754
755 let key = AgentCacheKey::new("important work", 0);
756 let cached = j2.get_cached(&key).unwrap();
757 assert_eq!(cached.output, serde_json::json!({"survived": true}));
758 }
759 }
760
761 #[test]
763 fn test_gc_older_than() {
764 let dir = tempdir().unwrap();
765 let run_dir = dir.path().join("runs");
766 std::fs::create_dir_all(&run_dir).unwrap();
767
768 let run_id = uuid::Uuid::now_v7();
770 let journal = JournalStore::new(&run_dir.join(run_id.to_string())).unwrap();
771 journal.init_run(run_id, "GC me").unwrap();
772
773 if let Some(mut cp) = journal.get_checkpoint() {
775 cp.status = crate::state::CheckpointStatus::Completed;
776 cp.updated_at = 1000; let _ = journal.inner.save_checkpoint(&cp);
778 }
779
780 let cleaned = gc_runs(&run_dir, Duration::from_secs(3600)).unwrap();
782 assert_eq!(cleaned, 1);
783 }
784
785 fn read_checkpoint_status_for(run_dir: &std::path::Path, agent_id: AgentId) -> Option<String> {
797 let cp_path = run_dir.join("checkpoint.json");
798 let content = std::fs::read_to_string(&cp_path).ok()?;
799 let raw: serde_json::Value = serde_json::from_str(&content).ok()?;
800 let ar = raw.get("agent_results")?.as_object()?;
801 for (_k, v) in ar {
802 if v.get("agent_id").and_then(|id| id.as_str()) == Some(&agent_id.to_string()) {
803 return v.get("status").and_then(|s| s.as_str()).map(String::from);
804 }
805 }
806 None
807 }
808
809 fn sample_token_usage(input: u64, output: u64) -> TokenUsage {
810 TokenUsage {
811 input,
812 output,
813 cache_read: 0,
814 cache_write: 0,
815 }
816 }
817
818 #[test]
819 fn cache_agent_persists_snake_case_status_for_each_variant() {
820 let dir = tempdir().unwrap();
825 let run_id = uuid::Uuid::now_v7();
826 let journal = JournalStore::new(dir.path()).unwrap();
827 journal.init_run(run_id, "cache_agent F5").unwrap();
828
829 let cases: Vec<(AgentStatus, &str)> = vec![
830 (AgentStatus::Ok, "ok"),
831 (AgentStatus::Error, "error"),
832 (AgentStatus::Cancelled, "cancelled"),
833 (AgentStatus::TimedOut, "timed_out"),
834 ];
835 for (status, expected) in &cases {
836 let agent_id = uuid::Uuid::now_v7();
837 let key = AgentCacheKey::new("prompt", 1);
838 journal
839 .cache_agent(
840 &key,
841 agent_id,
842 1,
843 status.clone(),
844 serde_json::json!({"v": 1}),
845 vec![],
846 sample_token_usage(10, 5),
847 )
848 .unwrap();
849
850 let persisted = read_checkpoint_status_for(dir.path(), agent_id)
851 .unwrap_or_else(|| panic!("status missing on disk for {status:?}"));
852 assert_eq!(
853 persisted, *expected,
854 "cache_agent({status:?}) must persist status={expected:?} (snake_case); \
855 got {persisted:?}. Reverting to Debug formatting would yield \"timedout\" \
856 for TimedOut and break the on-disk contract."
857 );
858 }
859 }
860
861 #[test]
862 fn cache_agent_timed_out_persists_with_underscore_not_collapsed() {
863 let dir = tempdir().unwrap();
867 let run_id = uuid::Uuid::now_v7();
868 let journal = JournalStore::new(dir.path()).unwrap();
869 journal.init_run(run_id, "timed-out guard").unwrap();
870
871 let agent_id = uuid::Uuid::now_v7();
872 let key = AgentCacheKey::new("p", 0);
873 journal
874 .cache_agent(
875 &key,
876 agent_id,
877 0,
878 AgentStatus::TimedOut,
879 serde_json::json!(null),
880 vec![],
881 sample_token_usage(1, 2),
882 )
883 .unwrap();
884
885 let persisted = read_checkpoint_status_for(dir.path(), agent_id).expect("status on disk");
886 assert_eq!(
887 persisted, "timed_out",
888 "cache_agent(TimedOut) must persist \"timed_out\"; got {persisted:?}"
889 );
890 assert_ne!(
891 persisted, "timedout",
892 "cache_agent(TimedOut) must NOT collapse to Debug-lowercased \"timedout\""
893 );
894 }
895
896 #[test]
897 fn record_result_persists_snake_case_status_for_each_variant() {
898 let dir = tempdir().unwrap();
901 let run_id = uuid::Uuid::now_v7();
902 let journal = JournalStore::new(dir.path()).unwrap();
903 journal.init_run(run_id, "record_result F5").unwrap();
904
905 let cases: Vec<(AgentStatus, &str)> = vec![
906 (AgentStatus::Ok, "ok"),
907 (AgentStatus::Error, "error"),
908 (AgentStatus::Cancelled, "cancelled"),
909 (AgentStatus::TimedOut, "timed_out"),
910 ];
911 for (status, expected) in &cases {
912 let agent_id = uuid::Uuid::now_v7();
913 let key = AgentCacheKey::new("p", 1);
914 journal.record_result(
915 &key,
916 agent_id,
917 1,
918 status.clone(),
919 serde_json::json!({"r": 1}),
920 vec![],
921 sample_token_usage(2, 3),
922 );
923
924 let persisted = read_checkpoint_status_for(dir.path(), agent_id)
925 .unwrap_or_else(|| panic!("status missing on disk for {status:?}"));
926 assert_eq!(
927 persisted, *expected,
928 "record_result({status:?}) must persist status={expected:?}; got {persisted:?}"
929 );
930 }
931 }
932
933 #[test]
934 fn record_result_timed_out_persists_with_underscore() {
935 let dir = tempdir().unwrap();
937 let run_id = uuid::Uuid::now_v7();
938 let journal = JournalStore::new(dir.path()).unwrap();
939 journal.init_run(run_id, "record_result timed-out").unwrap();
940
941 let agent_id = uuid::Uuid::now_v7();
942 let key = AgentCacheKey::new("p", 0);
943 journal.record_result(
944 &key,
945 agent_id,
946 0,
947 AgentStatus::TimedOut,
948 serde_json::json!(null),
949 vec![],
950 sample_token_usage(0, 0),
951 );
952
953 let persisted = read_checkpoint_status_for(dir.path(), agent_id).expect("status on disk");
954 assert_eq!(persisted, "timed_out");
955 assert_ne!(persisted, "timedout");
956 }
957
958 #[tokio::test]
959 async fn journal_callback_on_agent_done_persists_snake_case_status() {
960 let dir = tempdir().unwrap();
964 let run_id = uuid::Uuid::now_v7();
965 let journal = std::sync::Arc::new(JournalStore::new(dir.path()).unwrap());
966 journal.init_run(run_id, "callback F5").unwrap();
967
968 let cases: Vec<(AgentStatus, &str)> = vec![
969 (AgentStatus::Ok, "ok"),
970 (AgentStatus::Error, "error"),
971 (AgentStatus::Cancelled, "cancelled"),
972 (AgentStatus::TimedOut, "timed_out"),
973 ];
974 for (status, expected) in &cases {
975 let agent_id = uuid::Uuid::now_v7();
976 use crate::scheduler::JournalCallback;
977 journal
978 .on_agent_done(
979 agent_id,
980 1,
981 status.clone(),
982 serde_json::json!({}),
983 sample_token_usage(4, 6),
984 )
985 .await;
986
987 let persisted = read_checkpoint_status_for(dir.path(), agent_id)
988 .unwrap_or_else(|| panic!("status missing on disk for {status:?}"));
989 assert_eq!(
990 persisted, *expected,
991 "JournalCallback::on_agent_done({status:?}) must persist status={expected:?}; \
992 got {persisted:?}"
993 );
994 }
995 }
996
997 #[test]
998 fn record_result_then_reopen_uses_snake_case_status() {
999 let dir = tempdir().unwrap();
1002 let run_id = uuid::Uuid::now_v7();
1003 let journal = JournalStore::new(dir.path()).unwrap();
1004 journal.init_run(run_id, "reopen F5").unwrap();
1005
1006 let agent_id = uuid::Uuid::now_v7();
1007 let key = AgentCacheKey::new("reopen prompt", 1);
1008 journal.record_result(
1009 &key,
1010 agent_id,
1011 1,
1012 AgentStatus::Cancelled,
1013 serde_json::json!({"result": "ok"}),
1014 vec![],
1015 sample_token_usage(7, 11),
1016 );
1017 drop(journal);
1018
1019 let j2 = JournalStore::new(dir.path()).unwrap();
1020 let cp = j2.open(run_id).expect("open after drop");
1021 let cached = cp
1022 .agent_results
1023 .get(&agent_id)
1024 .expect("entry survives reopen");
1025 assert_eq!(
1026 cached.status, "cancelled",
1027 "snake_case status must round-trip through close+reopen"
1028 );
1029 assert_eq!(cached.tokens, 18);
1030 }
1031
1032 #[test]
1033 fn cache_agent_persists_snake_case_status_to_event_log() {
1034 let dir = tempdir().unwrap();
1038 let run_id = uuid::Uuid::now_v7();
1039 let journal = JournalStore::new(dir.path()).unwrap();
1040 journal.init_run(run_id, "event log F5").unwrap();
1041
1042 let agent_id = uuid::Uuid::now_v7();
1043 let key = AgentCacheKey::new("p", 1);
1044 journal
1045 .cache_agent(
1046 &key,
1047 agent_id,
1048 1,
1049 AgentStatus::TimedOut,
1050 serde_json::json!(null),
1051 vec![],
1052 sample_token_usage(1, 1),
1053 )
1054 .unwrap();
1055
1056 let log = journal.store().get_event_log().expect("read events.jsonl");
1060 let agent_done = log
1061 .iter()
1062 .find_map(|e| match e {
1063 AgentEvent::AgentDone {
1064 agent_id: id,
1065 status,
1066 ..
1067 } if id == &agent_id => Some(status.clone()),
1068 _ => None,
1069 })
1070 .expect("AgentDone event in log");
1071 assert!(matches!(agent_done, AgentStatus::TimedOut));
1075 }
1076
1077 #[tokio::test]
1082 async fn on_agent_done_preserves_cache_key_hash_from_record_result() {
1083 let dir = tempdir().unwrap();
1087 let run_id = uuid::Uuid::now_v7();
1088 let journal = std::sync::Arc::new(JournalStore::new(dir.path()).unwrap());
1089 journal.init_run(run_id, "hash preservation").unwrap();
1090
1091 let agent_id = uuid::Uuid::now_v7();
1092 let key = AgentCacheKey::new("preserve me", 1);
1093
1094 journal.record_result(
1096 &key,
1097 agent_id,
1098 1,
1099 AgentStatus::Ok,
1100 serde_json::json!({"answer": 42}),
1101 vec![],
1102 sample_token_usage(10, 5),
1103 );
1104
1105 use crate::scheduler::JournalCallback;
1107 journal
1108 .on_agent_done(
1109 agent_id,
1110 1,
1111 AgentStatus::Ok,
1112 serde_json::json!({}),
1113 sample_token_usage(10, 5),
1114 )
1115 .await;
1116
1117 assert!(
1119 journal.has_completed(&key),
1120 "cache_key_hash must survive on_agent_done"
1121 );
1122
1123 drop(journal);
1125 let j2 = JournalStore::new(dir.path()).unwrap();
1126 j2.open(run_id).expect("reopen");
1127 assert!(
1128 j2.has_completed(&key),
1129 "cache_key_hash must survive reopen after on_agent_done"
1130 );
1131 }
1132
1133 #[tokio::test]
1134 async fn on_agent_done_preserves_cache_key_hash_from_cache_agent() {
1135 let dir = tempdir().unwrap();
1137 let run_id = uuid::Uuid::now_v7();
1138 let journal = std::sync::Arc::new(JournalStore::new(dir.path()).unwrap());
1139 journal.init_run(run_id, "hash preservation 2").unwrap();
1140
1141 let agent_id = uuid::Uuid::now_v7();
1142 let key = AgentCacheKey::new("preserve me 2", 0);
1143
1144 journal
1145 .cache_agent(
1146 &key,
1147 agent_id,
1148 0,
1149 AgentStatus::Ok,
1150 serde_json::json!({"r": 1}),
1151 vec![],
1152 sample_token_usage(1, 1),
1153 )
1154 .unwrap();
1155
1156 use crate::scheduler::JournalCallback;
1157 journal
1158 .on_agent_done(
1159 agent_id,
1160 0,
1161 AgentStatus::Ok,
1162 serde_json::json!({}),
1163 sample_token_usage(1, 1),
1164 )
1165 .await;
1166
1167 assert!(
1168 journal.has_completed(&key),
1169 "cache_key_hash must survive on_agent_done after cache_agent"
1170 );
1171 }
1172}