1use 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#[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
130pub trait CheckpointBackend: Send + Sync + std::fmt::Debug {
140 fn init_run(&self, run_id: RunId, task: &str, run_dir: &str) -> anyhow::Result<()>;
142
143 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 fn open_run(&self, run_id: RunId) -> anyhow::Result<Option<RunCheckpoint>>;
154
155 fn append_event(&self, event: &AgentEvent) -> anyhow::Result<()>;
157
158 fn upsert_agent_result(&self, cache: &AgentResultCache) -> anyhow::Result<()>;
160
161 fn upsert_agent_session(&self, session: &AgentSessionCheckpoint) -> anyhow::Result<()>;
163
164 fn get_checkpoint(&self) -> Option<RunCheckpoint>;
166
167 fn get_findings(&self) -> Vec<Finding>;
169
170 fn get_event_log(&self) -> anyhow::Result<Vec<AgentEvent>>;
172
173 fn can_resume(&self) -> bool;
175
176 fn reset_status_to_running(&self) -> anyhow::Result<()>;
178
179 fn cancel(&self) -> anyhow::Result<()>;
181
182 fn save_checkpoint(&self, checkpoint: &RunCheckpoint) -> anyhow::Result<()>;
184}
185
186pub 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
194pub 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}