gestura_core/checkpoints/
session.rs1use serde::{Deserialize, Serialize};
8
9use crate::agent_sessions::{AgentSession, AgentSessionStore};
10use crate::config::{AppConfig, GlobalPermissionSettings, PipelineSettings};
11use crate::hooks::HooksSettings;
12use crate::tasks::{TaskList, TaskManager};
13
14use super::CheckpointManager;
15use super::types::{CheckpointError, CheckpointId, CheckpointMetadata, CheckpointSnapshot};
16
17pub const SESSION_CHECKPOINT_SCHEMA_V1: u32 = 1;
19
20#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
24pub struct SessionCheckpointConfig {
25 pub pipeline: PipelineSettings,
27
28 pub permissions: GlobalPermissionSettings,
30
31 pub llm_primary_provider: String,
33
34 #[serde(default, skip_serializing_if = "Option::is_none")]
36 pub llm_fallback_provider: Option<String>,
37
38 #[serde(default, skip_serializing_if = "Option::is_none")]
40 pub openai_model: Option<String>,
41
42 #[serde(default, skip_serializing_if = "Option::is_none")]
44 pub anthropic_model: Option<String>,
45
46 #[serde(default, skip_serializing_if = "Option::is_none")]
48 pub grok_model: Option<String>,
49
50 #[serde(default, skip_serializing_if = "Option::is_none")]
52 pub ollama_model: Option<String>,
53
54 #[serde(default)]
59 pub hooks: HooksSettings,
60}
61
62impl SessionCheckpointConfig {
63 pub fn from_app_config(config: &AppConfig) -> Self {
67 Self {
68 pipeline: config.pipeline.clone(),
69 permissions: config.permissions.clone(),
70 llm_primary_provider: config.llm.primary.clone(),
71 llm_fallback_provider: config.llm.fallback.clone(),
72 openai_model: config.llm.openai.as_ref().map(|c| c.model.clone()),
73 anthropic_model: config.llm.anthropic.as_ref().map(|c| c.model.clone()),
74 grok_model: config.llm.grok.as_ref().map(|c| c.model.clone()),
75 ollama_model: config.llm.ollama.as_ref().map(|c| c.model.clone()),
76 hooks: config.hooks.clone(),
77 }
78 }
79}
80
81#[derive(Debug, Clone, Serialize, Deserialize)]
83pub struct SessionCheckpointPayload {
84 pub schema_version: u32,
86
87 pub session: AgentSession,
89
90 pub tasks: TaskList,
92
93 pub config: SessionCheckpointConfig,
95}
96
97impl SessionCheckpointPayload {
98 pub fn validate(&self) -> Result<(), CheckpointError> {
100 if self.schema_version != SESSION_CHECKPOINT_SCHEMA_V1 {
101 return Err(CheckpointError::UnsupportedSchema {
102 expected: SESSION_CHECKPOINT_SCHEMA_V1,
103 found: self.schema_version,
104 });
105 }
106
107 if self.session.id != self.tasks.session_id {
108 return Err(CheckpointError::InvalidInput(format!(
109 "checkpoint payload session/task mismatch: session.id='{}' tasks.session_id='{}'",
110 self.session.id, self.tasks.session_id
111 )));
112 }
113
114 Ok(())
115 }
116}
117
118impl CheckpointManager {
119 pub fn create_session_checkpoint(
126 &self,
127 session_id: &str,
128 session_store: &dyn AgentSessionStore,
129 task_manager: &TaskManager,
130 config: &AppConfig,
131 label: Option<String>,
132 ) -> Result<CheckpointMetadata, CheckpointError> {
133 let session = session_store
134 .load(session_id)
135 .map_err(|e| CheckpointError::AgentSession(e.to_string()))?;
136
137 let tasks = task_manager
138 .load_task_list(session_id)
139 .map_err(|e| CheckpointError::Tasks(e.to_string()))?;
140
141 let payload = SessionCheckpointPayload {
142 schema_version: SESSION_CHECKPOINT_SCHEMA_V1,
143 session,
144 tasks,
145 config: SessionCheckpointConfig::from_app_config(config),
146 };
147 payload.validate()?;
148
149 let snapshot = CheckpointSnapshot {
150 session_id: Some(session_id.to_string()),
151 payload: serde_json::to_value(payload)?,
152 };
153
154 self.create_checkpoint(snapshot, label)
155 }
156
157 pub fn restore_session_checkpoint(
159 &self,
160 id: &CheckpointId,
161 ) -> Result<SessionCheckpointPayload, CheckpointError> {
162 let snapshot = self.restore_checkpoint(id)?;
163 let payload: SessionCheckpointPayload = serde_json::from_value(snapshot.payload)?;
164 payload.validate()?;
165 Ok(payload)
166 }
167
168 pub fn apply_session_checkpoint(
172 &self,
173 id: &CheckpointId,
174 session_store: &dyn AgentSessionStore,
175 task_manager: &TaskManager,
176 ) -> Result<SessionCheckpointPayload, CheckpointError> {
177 let payload = self.restore_session_checkpoint(id)?;
178
179 session_store
180 .save(&payload.session)
181 .map_err(|e| CheckpointError::AgentSession(e.to_string()))?;
182 task_manager
183 .replace_task_list(payload.tasks.clone())
184 .map_err(|e| CheckpointError::Tasks(e.to_string()))?;
185
186 Ok(payload)
187 }
188
189 pub fn list_session_checkpoints(
191 &self,
192 session_id: &str,
193 ) -> Result<Vec<CheckpointMetadata>, CheckpointError> {
194 let all = self.list_checkpoints()?;
195 Ok(all
196 .into_iter()
197 .filter(|m| m.session_id.as_deref() == Some(session_id))
198 .collect())
199 }
200}
201
202#[cfg(test)]
203mod tests {
204 use super::*;
205
206 use crate::agent_sessions::{FileAgentSessionStore, MessageSource};
207 use crate::checkpoints::{CheckpointRetentionPolicy, FileCheckpointStore};
208 use tempfile::tempdir;
209
210 #[test]
211 fn create_and_apply_session_checkpoint_restores_session_and_tasks() {
212 let temp = tempdir().unwrap();
213
214 let sessions_dir = temp.path().join("sessions");
215 let session_store = FileAgentSessionStore::new(sessions_dir);
216
217 let checkpoint_store = FileCheckpointStore::new(temp.path().join("checkpoints"));
218 let manager =
219 CheckpointManager::new(checkpoint_store, CheckpointRetentionPolicy::default());
220
221 let task_manager = TaskManager::new(temp.path());
222 let config = AppConfig::default();
223
224 let workspace_dir = temp.path().join("workspace");
226 std::fs::create_dir_all(&workspace_dir).unwrap();
227 let mut session =
228 AgentSession::new_with_workspace(workspace_dir, Some("m".to_string())).unwrap();
229 session.add_user_message("hello", MessageSource::Text);
230 session_store.save(&session).unwrap();
231
232 let t = task_manager
234 .create_task(&session.id, "Task", "Desc", None)
235 .unwrap();
236 task_manager
237 .set_current_task_id(&session.id, Some(t.id.clone()))
238 .unwrap();
239
240 let meta = manager
242 .create_session_checkpoint(
243 &session.id,
244 &session_store,
245 &task_manager,
246 &config,
247 Some("before-change".to_string()),
248 )
249 .unwrap();
250
251 let mut mutated = session_store.load(&session.id).unwrap();
253 mutated.add_user_message("later", MessageSource::Text);
254 session_store.save(&mutated).unwrap();
255
256 let t2 = task_manager
257 .create_task(&session.id, "Task2", "Desc2", None)
258 .unwrap();
259 task_manager
260 .set_current_task_id(&session.id, Some(t2.id.clone()))
261 .unwrap();
262
263 let applied = manager
265 .apply_session_checkpoint(&meta.id, &session_store, &task_manager)
266 .unwrap();
267
268 let rewound = session_store.load(&session.id).unwrap();
270 assert_eq!(rewound.message_count(), 1);
271 assert_eq!(rewound.state.messages[0].content, "hello");
272
273 let loaded_tasks = task_manager.load_task_list(&session.id).unwrap();
275 assert_eq!(loaded_tasks.tasks.len(), 1);
276 assert_eq!(loaded_tasks.current_task_id(), Some(t.id.as_str()));
277
278 assert_eq!(applied.schema_version, SESSION_CHECKPOINT_SCHEMA_V1);
280 assert_eq!(applied.config.pipeline, config.pipeline);
281 assert_eq!(applied.config.permissions, config.permissions);
282 }
283
284 #[test]
285 fn retention_deletes_oldest_files() {
286 let temp = tempdir().unwrap();
287
288 let session_store = FileAgentSessionStore::new(temp.path().join("sessions"));
289 let checkpoint_dir = temp.path().join("checkpoints");
290 let checkpoint_store = FileCheckpointStore::new(checkpoint_dir.clone());
291 let manager = CheckpointManager::new(
292 checkpoint_store,
293 CheckpointRetentionPolicy { max_checkpoints: 2 },
294 );
295
296 let task_manager = TaskManager::new(temp.path());
297 let config = AppConfig::default();
298
299 let workspace_dir = temp.path().join("workspace");
300 std::fs::create_dir_all(&workspace_dir).unwrap();
301 let session = AgentSession::new_with_workspace(workspace_dir, None).unwrap();
302 session_store.save(&session).unwrap();
303
304 manager
305 .create_session_checkpoint(&session.id, &session_store, &task_manager, &config, None)
306 .unwrap();
307 manager
308 .create_session_checkpoint(&session.id, &session_store, &task_manager, &config, None)
309 .unwrap();
310 manager
311 .create_session_checkpoint(&session.id, &session_store, &task_manager, &config, None)
312 .unwrap();
313
314 let file_count = std::fs::read_dir(&checkpoint_dir)
315 .unwrap()
316 .filter(|e| {
317 e.as_ref()
318 .ok()
319 .and_then(|e| {
320 e.path()
321 .extension()
322 .map(|ext| ext == std::ffi::OsStr::new("json"))
323 })
324 .unwrap_or(false)
325 })
326 .count();
327 assert_eq!(file_count, 2);
328 }
329}