Skip to main content

luft_core/
state.rs

1//! Progress persistence and resume for long-running workflows.
2//!
3//! This module defines the data types for checkpointing and the
4//! `CheckpointBackend` trait that persistence engines (e.g. SQLite) implement.
5//!
6//! Key features:
7//! - Event log persistence
8//! - Agent result caching
9//! - Resume from last checkpoint
10//! - Run state management
11
12use crate::contract::event::AgentEvent;
13use crate::contract::finding::Finding;
14use crate::contract::ids::{AgentId, PhaseId, RunId};
15use serde::{Deserialize, Serialize};
16use std::collections::HashMap;
17use std::path::Path;
18
19
20// ============================================================================
21// Data Types (frozen contracts)
22// ============================================================================
23
24/// Run state persisted to the backend.
25#[derive(Debug, Clone, Serialize, Deserialize)]
26pub struct RunCheckpoint {
27    pub run_id: RunId,
28    pub task: String,
29    pub status: CheckpointStatus,
30    pub current_phase: u32,
31    pub completed_phases: Vec<PhaseSummary>,
32    pub agent_results: HashMap<AgentId, AgentResultCache>,
33    #[serde(default)]
34    pub agent_sessions: HashMap<AgentId, AgentSessionCheckpoint>,
35    pub findings: Vec<Finding>,
36    pub total_tokens: u64,
37    pub created_at: u64,
38    pub updated_at: u64,
39    #[serde(default)]
40    pub workflow_meta: Option<serde_json::Value>,
41    #[serde(default)]
42    pub started_agent_ids: Vec<AgentId>,
43}
44
45#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
46#[serde(rename_all = "lowercase")]
47pub enum CheckpointStatus {
48    Running,
49    Completed,
50    Failed,
51    Cancelled,
52}
53
54impl std::fmt::Display for CheckpointStatus {
55    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
56        let s = match self {
57            CheckpointStatus::Running => "Running",
58            CheckpointStatus::Completed => "Completed",
59            CheckpointStatus::Failed => "Failed",
60            CheckpointStatus::Cancelled => "Cancelled",
61        };
62        f.write_str(s)
63    }
64}
65
66impl CheckpointStatus {
67    pub fn as_str(&self) -> &'static str {
68        match self {
69            CheckpointStatus::Running => "running",
70            CheckpointStatus::Completed => "completed",
71            CheckpointStatus::Failed => "failed",
72            CheckpointStatus::Cancelled => "cancelled",
73        }
74    }
75
76    pub fn parse_str(s: &str) -> Self {
77        match s.to_lowercase().as_str() {
78            "completed" => CheckpointStatus::Completed,
79            "failed" => CheckpointStatus::Failed,
80            "cancelled" => CheckpointStatus::Cancelled,
81            _ => CheckpointStatus::Running,
82        }
83    }
84}
85
86#[derive(Debug, Clone, Serialize, Deserialize)]
87pub struct PhaseSummary {
88    pub phase_id: PhaseId,
89    pub label: String,
90    pub planned: usize,
91    pub ok: usize,
92    pub failed: usize,
93    #[serde(default)]
94    pub description: Option<String>,
95    #[serde(default)]
96    pub role: Option<String>,
97}
98
99#[derive(Debug, Clone, Serialize, Deserialize)]
100pub struct AgentResultCache {
101    pub agent_id: AgentId,
102    pub phase_id: PhaseId,
103    pub status: String,
104    pub output: serde_json::Value,
105    pub findings: Vec<Finding>,
106    pub tokens: u64,
107    pub completed_at: u64,
108    #[serde(default)]
109    pub cache_key_hash: Option<String>,
110    #[serde(default)]
111    pub description: Option<String>,
112    #[serde(default)]
113    pub role: Option<String>,
114}
115
116#[derive(Debug, Clone, Serialize, Deserialize)]
117pub struct AgentSessionCheckpoint {
118    pub agent_id: AgentId,
119    #[serde(default)]
120    pub backend_id: Option<String>,
121    #[serde(default)]
122    pub protocol_session_id: Option<String>,
123    pub session_id: String,
124    pub status: String,
125    pub updated_at: u64,
126    #[serde(default)]
127    pub resumable: bool,
128}
129
130// ============================================================================
131// CheckpointBackend Trait
132// ============================================================================
133
134/// Persistence backend for a single run.
135///
136/// Implementations (e.g. `SqliteCheckpointBackend` in `luft-storage`)
137/// provide checkpoint + event log storage. All methods are synchronous;
138/// async backends bridge internally via `block_in_place`.
139pub trait CheckpointBackend: Send + Sync + std::fmt::Debug {
140    /// Initialize a new run.
141    fn init_run(&self, run_id: RunId, task: &str, run_dir: &str) -> anyhow::Result<()>;
142
143    /// Initialize a new run with declarative workflow metadata.
144    fn init_run_with_meta(
145        &self,
146        run_id: RunId,
147        task: &str,
148        run_dir: &str,
149        workflow_meta: serde_json::Value,
150    ) -> anyhow::Result<()>;
151
152    /// Open an existing run for resume. Returns None if not found.
153    fn open_run(&self, run_id: RunId) -> anyhow::Result<Option<RunCheckpoint>>;
154
155    /// Append an event to the log and update checkpoint state.
156    fn append_event(&self, event: &AgentEvent) -> anyhow::Result<()>;
157
158    /// Insert or update an agent result.
159    fn upsert_agent_result(&self, cache: &AgentResultCache) -> anyhow::Result<()>;
160
161    /// Insert or update an agent session.
162    fn upsert_agent_session(&self, session: &AgentSessionCheckpoint) -> anyhow::Result<()>;
163
164    /// Get current checkpoint (from in-memory cache).
165    fn get_checkpoint(&self) -> Option<RunCheckpoint>;
166
167    /// Get all findings.
168    fn get_findings(&self) -> Vec<Finding>;
169
170    /// Get event log as a vector.
171    fn get_event_log(&self) -> anyhow::Result<Vec<AgentEvent>>;
172
173    /// Check if a run can be resumed.
174    fn can_resume(&self) -> bool;
175
176    /// Reset checkpoint status to Running.
177    fn reset_status_to_running(&self) -> anyhow::Result<()>;
178
179    /// Mark run as cancelled.
180    fn cancel(&self) -> anyhow::Result<()>;
181
182    /// Save checkpoint (full overwrite).
183    fn save_checkpoint(&self, checkpoint: &RunCheckpoint) -> anyhow::Result<()>;
184}
185
186/// Helper: current unix timestamp.
187pub fn current_timestamp() -> u64 {
188    std::time::SystemTime::now()
189        .duration_since(std::time::UNIX_EPOCH)
190        .map(|d| d.as_secs())
191        .unwrap_or(0)
192}
193
194// ============================================================================
195// Factory — callers provide a backend at construction time.
196// ============================================================================
197
198/// List all run directory names under the base dir.
199/// For SQLite backends, this is derived from the `runs` table.
200/// This helper remains for filesystem-based discovery.
201pub fn list_run_dirs(base_dir: &Path) -> anyhow::Result<Vec<String>> {
202    if !base_dir.exists() {
203        return Ok(vec![]);
204    }
205    let mut run_dirs = Vec::new();
206    for entry in std::fs::read_dir(base_dir)? {
207        let entry = entry?;
208        let path = entry.path();
209        if path.is_dir() {
210            if let Some(name) = path.file_name().and_then(|n| n.to_str()) {
211                run_dirs.push(name.to_string());
212            }
213        }
214    }
215    run_dirs.sort();
216    Ok(run_dirs)
217}