gestura_core/checkpoints/
session.rs

1//! Session-specific checkpoint (rewind) helpers.
2//!
3//! This layer builds on the generic [`crate::checkpoints`] storage primitives to
4//! snapshot and restore an agent session together with related state (tasks + a
5//! **redacted** subset of configuration).
6
7use 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
17/// Schema version used for [`SessionCheckpointPayload`].
18pub const SESSION_CHECKPOINT_SCHEMA_V1: u32 = 1;
19
20/// A **redacted** subset of configuration captured in a session checkpoint.
21///
22/// This intentionally excludes secrets such as API keys.
23#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
24pub struct SessionCheckpointConfig {
25    /// Pipeline behavior relevant to prompt/history management.
26    pub pipeline: PipelineSettings,
27
28    /// Default permission behavior for new sessions.
29    pub permissions: GlobalPermissionSettings,
30
31    /// Global primary LLM provider id (e.g., "openai", "anthropic", "ollama").
32    pub llm_primary_provider: String,
33
34    /// Optional global fallback LLM provider id.
35    #[serde(default, skip_serializing_if = "Option::is_none")]
36    pub llm_fallback_provider: Option<String>,
37
38    /// Optional global model selector for OpenAI (no secrets).
39    #[serde(default, skip_serializing_if = "Option::is_none")]
40    pub openai_model: Option<String>,
41
42    /// Optional global model selector for Anthropic (no secrets).
43    #[serde(default, skip_serializing_if = "Option::is_none")]
44    pub anthropic_model: Option<String>,
45
46    /// Optional global model selector for Grok (no secrets).
47    #[serde(default, skip_serializing_if = "Option::is_none")]
48    pub grok_model: Option<String>,
49
50    /// Optional global model selector for Ollama (no secrets).
51    #[serde(default, skip_serializing_if = "Option::is_none")]
52    pub ollama_model: Option<String>,
53
54    /// Hooks configuration (no secrets).
55    ///
56    /// This is safe to capture because hook execution is still governed by
57    /// explicit allow-listing of programs and hooks are disabled by default.
58    #[serde(default)]
59    pub hooks: HooksSettings,
60}
61
62impl SessionCheckpointConfig {
63    /// Build a checkpoint-safe config snapshot from the full application config.
64    ///
65    /// This method **must not** include secrets (API keys).
66    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/// Typed payload stored inside a session checkpoint.
82#[derive(Debug, Clone, Serialize, Deserialize)]
83pub struct SessionCheckpointPayload {
84    /// Payload schema version.
85    pub schema_version: u32,
86
87    /// The agent session state (messages, tool calls, session overrides).
88    pub session: AgentSession,
89
90    /// The task list for this session, including the persisted `current_task_id` pointer.
91    pub tasks: TaskList,
92
93    /// Redacted configuration snapshot.
94    pub config: SessionCheckpointConfig,
95}
96
97impl SessionCheckpointPayload {
98    /// Validate basic invariants for this payload.
99    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    /// Create a checkpoint for a specific session.
120    ///
121    /// This snapshots:
122    /// - agent session state (history + session overrides)
123    /// - task list state (including `current_task_id`)
124    /// - a redacted subset of global configuration
125    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    /// Restore a session checkpoint payload without applying it.
158    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    /// Apply a session checkpoint by writing the restored state back to persistence.
169    ///
170    /// Returns the restored payload (useful for callers that want to update UI state).
171    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    /// List checkpoints belonging to a particular session.
190    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        // Create a session.
225        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        // Create tasks + set current pointer.
233        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        // Snapshot.
241        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        // Mutate persisted state.
252        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        // Apply checkpoint.
264        let applied = manager
265            .apply_session_checkpoint(&meta.id, &session_store, &task_manager)
266            .unwrap();
267
268        // Verify session rewound.
269        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        // Verify tasks rewound.
274        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        // Verify config included and schema is valid.
279        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}