Skip to main content

luft_core/
journal.rs

1//! Journal / Resume — checkpoint persistence with replay semantics (M1).
2//!
3//! Provides:
4//! - `JournalStore` — wraps `RunStore` with cache-key index for O(1) lookups
5//! - `AgentCacheKey` — deterministic blake3-based key for agent invocations
6//! - `JournalCallback` trait — scheduler integration hook
7//! - `ResumeContext` — orchestrates run recovery
8//! - `gc_runs()` — cleanup old completed runs
9//!
10//! Thread safety: All public methods take `&self` (interior mutability via RwLock).
11//! The underlying checkpoint data is protected by a single writer lock.
12//!
13//! Lifecycle:
14//!   new() → init_run() → cache_agent()* → flush()
15//!   或:
16//!   open() → has_completed()/get_cached() → workflow resume logic
17
18use 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// ============================================================================
35// Error Types
36// ============================================================================
37
38#[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// ============================================================================
53// Agent Cache Key
54// ============================================================================
55
56/// Deterministic cache key for an agent invocation.
57/// Normalizes whitespace/unicode to ensure cache hits across formatting differences.
58#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
59pub struct AgentCacheKey {
60    pub hash: String,
61    /// Human-readable for debugging
62    pub prompt_preview: String,
63    pub phase_id: PhaseId,
64}
65
66impl AgentCacheKey {
67    /// Generate a cache key from agent parameters.
68    /// Uses blake3 with null separators to prevent field-concatenation collisions.
69    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
99// ============================================================================
100// JournalStore — the journal abstraction over RunStore
101// ============================================================================
102
103/// JournalStore wraps RunStore with replay semantics.
104///
105/// Thread safety: All public methods take `&self` (interior mutability via RwLock).
106/// The underlying checkpoint data is protected by a single writer lock.
107///
108/// Usage lifecycle:
109///   new() → init_run() → cache_agent()* → flush()
110///   或:
111///   open() → has_completed()/get_cached() → workflow resume logic
112pub struct JournalStore {
113    /// Underlying persistence engine (checkpoint.json + events.jsonl).
114    inner: Arc<RunStore>,
115    /// In-memory index: AgentCacheKey hash → AgentResultCache.
116    /// Populated at open() time from the checkpoint's agent_results map.
117    cache_index: RwLock<HashMap<String, AgentResultCache>>,
118    /// Event sender for broadcasting journal updates.
119    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    /// Create a new journal store at the given directory.
134    /// Initializes the underlying RunStore and creates an empty cache index.
135    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    /// Initialize a new run in the journal.
146    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    /// Initialize a new run with declarative workflow metadata.
153    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    /// Open an existing run and rebuild the cache index from persisted data.
168    ///
169    /// This is the entry point for `--resume`. It:
170    /// 1. Loads the checkpoint from disk
171    /// 2. Rebuilds the in-memory cache_index from agent_results
172    /// 3. Returns the checkpoint for the caller to inspect
173    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        // Rebuild cache index — index by both agent_id and cache_key_hash
190        // so that the Lua SDK's has_completed(key) works after resume.
191        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    /// Cache an agent's result in the journal.
204    ///
205    /// Called by the scheduler after an agent completes successfully or fails
206    /// with a non-retryable error. The result is persisted to disk immediately
207    /// (via append_event → update_from_event → write_checkpoint_to_disk).
208    #[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        // Update in-memory index (instant lookup)
234        {
235            let mut index = self.cache_index.write().unwrap();
236            index.insert(cache_key.hash.clone(), cache.clone());
237            // Also index by agent_id for open() compatibility
238            index.insert(agent_id.to_string(), cache.clone());
239        }
240
241        // Persist the full cache entry directly to checkpoint disk (preserves cache_key_hash)
242        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        // Also append event to log (this triggers update_from_event which finds the existing hash)
247        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        // Broadcast via event bus (non-blocking — uses broadcast channel)
268        if let Some(ref tx) = self.event_tx {
269            let _ = tx.send(event);
270        }
271
272        Ok(cache_key.clone())
273    }
274
275    /// Record an agent's output for resume replay, keyed by `cache_key`.
276    ///
277    /// Unlike [`cache_agent`], this does **not** append an `AgentDone` event,
278    /// so it never double-counts tokens against the event-driven checkpoint
279    /// totals. It only upserts the checkpoint entry (preserving `cache_key_hash`
280    /// and the structured output) and refreshes the in-memory cache index.
281    /// Called by the Lua SDK after an agent completes during a live run.
282    #[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    /// Access the underlying run store (shared persistence engine).
318    /// Allows the CLI to route the scheduler event stream through the same
319    /// `RunStore` instance the journal uses, avoiding split-brain checkpoints.
320    pub fn store(&self) -> Arc<RunStore> {
321        self.inner.clone()
322    }
323
324    /// Append an event to the underlying run store (event log + checkpoint).
325    pub fn append_event(&self, event: &AgentEvent) -> Result<(), JournalError> {
326        self.inner.append_event(event)?;
327        Ok(())
328    }
329
330    /// Check if an agent with the given cache key has already completed.
331    /// Used by the Lua SDK's agent() function before submitting to the scheduler.
332    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    /// Get cached result for an agent.
338    /// Returns None if the agent hasn't completed yet.
339    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    /// Get list of all completed agent cache keys.
345    /// Useful for debugging and progress reporting.
346    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    /// Get the underlying checkpoint (read-only snapshot).
359    pub fn get_checkpoint(&self) -> Option<RunCheckpoint> {
360        self.inner.get_checkpoint()
361    }
362
363    /// Flush all pending data to disk.
364    pub fn flush(&self) -> Result<(), JournalError> {
365        // RunStore auto-flushes on append_event; explicit flush for safety.
366        Ok(())
367    }
368
369    /// Mark the run as cancelled.
370    pub fn cancel(&self) -> Result<(), JournalError> {
371        self.inner.cancel()?;
372        Ok(())
373    }
374
375    /// Reset checkpoint status to `Running`. Used when resuming a
376    /// failed/cancelled run.
377    pub fn reset_status_to_running(&self) -> Result<(), JournalError> {
378        self.inner.reset_status_to_running()?;
379        Ok(())
380    }
381}
382
383// ============================================================================
384// Scheduler Integration — JournalCallback trait
385// ============================================================================
386
387/// Composite callback that chains multiple JournalCallback implementations.
388pub 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        // Preserve cache_key_hash and other enriched fields from a prior
428        // cache_agent() / record_result() call.  Without this, the scheduler
429        // callback would overwrite the hash with None, causing the agent to be
430        // re-executed on resume even though it already completed.
431        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        // Update in-memory index so subsequent on_agent_done calls also see
458        // the preserved hash.
459        {
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        // Persist to checkpoint disk
468        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// ============================================================================
475// Resume Orchestration
476// ============================================================================
477
478/// Context for resuming a run.
479#[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/// Options for creating a run (new or resume).
489#[derive(Debug, Clone)]
490pub enum RunCreationMode {
491    /// Start a fresh run.
492    New { task: String },
493    /// Resume from an existing checkpoint.
494    Resume { run_id: RunId, run_dir_name: String },
495    /// Auto-detect: resume if resumable run exists, else new.
496    Auto { task: String },
497}
498
499impl RunCreationMode {
500    /// Resolve the creation mode to concrete parameters.
501    /// For Auto mode, checks journal directory for resumable runs.
502    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                // List all run dirs, find the most recent Running checkpoint
521                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                // No resumable run — create new
535                let run_id = uuid::Uuid::now_v7();
536                Ok((run_id, None))
537            }
538        }
539    }
540}
541
542// ============================================================================
543// GC (Garbage Collection)
544// ============================================================================
545
546/// Clean up old completed/cancelled runs.
547///
548/// Policy:
549/// - Completed/Cancelled runs older than `older_than` are eligible for deletion.
550/// - Running runs are never cleaned.
551///
552/// Returns the number of runs cleaned.
553pub 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        // Peek at checkpoint without full open
562        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// ============================================================================
596// Tests
597// ============================================================================
598
599#[cfg(test)]
600mod tests {
601    use super::*;
602    use tempfile::tempdir;
603
604    /// Test basic journal lifecycle: init → cache → read → cancel
605    #[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        // 1. Init
612        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        // 2. Cache an agent result
618        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        // 3. Verify cache
638        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        // 4. Cancel
644        journal.cancel().unwrap();
645        let cp = journal.get_checkpoint().unwrap();
646        assert_eq!(cp.status, crate::state::CheckpointStatus::Cancelled);
647    }
648
649    /// Test that different prompts produce different cache keys
650    #[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        // Same prompt, different phase
657        let k4 = AgentCacheKey::new("prompt A", 2);
658        assert_ne!(k1.hash, k4.hash);
659
660        // Whitespace normalization
661        let k5 = AgentCacheKey::new("  prompt  \r\nA  ", 1);
662        assert_eq!(k1.hash, k5.hash);
663    }
664
665    /// Test resume: simulate workflow with 3 agents, 2 cached, 1 new
666    #[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        // Cache agent 1 and 2
674        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        // Verify cache hits
712        assert!(journal.has_completed(&k1));
713        assert!(journal.has_completed(&k2));
714        assert!(!journal.has_completed(&k3));
715
716        // Agent 3 should NOT be cached → would go through scheduler
717        assert!(journal.get_cached(&k3).is_none());
718    }
719
720    /// Test that journal survives crash (simulated by re-opening)
721    #[test]
722    fn test_journal_crash_recovery() {
723        let dir = tempdir().unwrap();
724        let run_id = uuid::Uuid::now_v7();
725
726        // Part 1: Create and cache
727        {
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        } // j dropped — simulates crash
747
748        // Part 2: Re-open and verify data survived
749        {
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 GC reference
762    #[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        // Create a completed run
769        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        // Manually mark as completed with old timestamp
774        if let Some(mut cp) = journal.get_checkpoint() {
775            cp.status = crate::state::CheckpointStatus::Completed;
776            cp.updated_at = 1000; // Very old
777            let _ = journal.inner.save_checkpoint(&cp);
778        }
779
780        // GC with very short duration
781        let cleaned = gc_runs(&run_dir, Duration::from_secs(3600)).unwrap();
782        assert_eq!(cleaned, 1);
783    }
784
785    // ----------------------------------------------------------------------
786    // Tests for the F5 contract — AgentResultCache.status persistence.
787    //
788    // cache_agent, record_result, and the JournalCallback impl for
789    // JournalStore all persist AgentResultCache.status. Before F5 the value
790    // was derived from `format!("{:?}", status).to_lowercase()`, which
791    // silently mis-mapped TimedOut → "timedout" (no underscore). The
792    // implementations must now use `AgentStatus::as_str()` and produce
793    // snake_case strings that match the canonical on-disk mapping.
794    // ----------------------------------------------------------------------
795
796    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        // F5 KEY test for JournalStore::cache_agent: the persisted status
821        // MUST equal AgentStatus::as_str() (snake_case), not Debug lowercased.
822        // Particularly important for TimedOut which would otherwise round-trip
823        // as "timedout" (no underscore) and break cross-process resume.
824        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        // Strongest F5 regression guard for cache_agent: TimedOut MUST persist
864        // as "timed_out" with an underscore. The buggy Debug-lowercased path
865        // would produce "timedout" and silently corrupt the journal.
866        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        // F5 test for JournalStore::record_result: same snake_case contract
899        // applies to the non-event-emitting path used by Lua SDK callbacks.
900        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        // Same regression guard for record_result.
936        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        // F5 test for the JournalCallback impl on JournalStore. The scheduler
961        // calls `on_agent_done` when an agent finishes; the persisted
962        // AgentResultCache.status MUST match AgentStatus::as_str() exactly.
963        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        // Snake_case persistence must survive a close+reopen cycle so a
1000        // resumed process sees the canonical strings (not Debug leftovers).
1001        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        // The AgentDone event itself also travels through the same snake_case
1035        // contract (via update_from_event → as_str()). Read events.jsonl back
1036        // and confirm the event log carries the canonical status.
1037        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        // The persisted AgentResultCache.status must already be verified by
1057        // the test above; this test only confirms the event log still parses
1058        // and carries the AgentDone event with the right status enum.
1059        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        // Status enum round-trip is enforced by serde, but the persisted
1072        // cache status string (verified above) is the part that the on-disk
1073        // contract depends on.
1074        assert!(matches!(agent_done, AgentStatus::TimedOut));
1075    }
1076
1077    // ------------------------------------------------------------------
1078    // Regression: on_agent_done must not clobber cache_key_hash
1079    // ------------------------------------------------------------------
1080
1081    #[tokio::test]
1082    async fn on_agent_done_preserves_cache_key_hash_from_record_result() {
1083        // Simulate the race: record_result() writes Some(hash), then the
1084        // scheduler callback on_agent_done() fires for the same agent.
1085        // The hash must survive — otherwise resume re-executes the agent.
1086        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        // 1. record_result writes Some(hash)
1095        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        // 2. scheduler callback fires later — must NOT overwrite hash with None
1106        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        // 3. In-memory index still has the hash entry
1118        assert!(
1119            journal.has_completed(&key),
1120            "cache_key_hash must survive on_agent_done"
1121        );
1122
1123        // 4. Disk checkpoint also preserves the hash
1124        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        // Same scenario but with cache_agent() as the first writer.
1136        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}