Skip to main content

a3s_code_core/store/
mod.rs

1//! Session persistence layer
2//!
3//! Provides pluggable session storage via the `SessionStore` trait.
4//!
5//! ## Default Implementation
6//!
7//! `FileSessionStore` stores each session as a JSON file:
8//! - Session metadata (id, name, timestamps)
9//! - Configuration (system prompt, policies)
10//! - Conversation history (messages)
11//! - Context usage statistics
12//!
13//! ## Custom Backends
14//!
15//! Implement `SessionStore` trait for custom backends (Redis, PostgreSQL, etc.):
16//!
17//! ```ignore
18//! use a3s_code::store::{
19//!     SessionData, SessionSnapshotV1, SessionStore, SessionStoreCapabilities,
20//! };
21//!
22//! struct RedisStore { /* ... */ }
23//!
24//! #[async_trait::async_trait]
25//! impl SessionStore for RedisStore {
26//!     // Required by AgentSession::save: commit the entire value in one
27//!     // backend transaction / atomic replacement.
28//!     async fn save_snapshot(&self, snapshot: &SessionSnapshotV1) -> Result<()> { /* ... */ }
29//!     async fn load_snapshot(&self, id: &str) -> Result<Option<SessionSnapshotV1>> { /* ... */ }
30//!     fn capabilities(&self) -> SessionStoreCapabilities {
31//!         SessionStoreCapabilities { atomic_session_snapshots: true }
32//!     }
33//!
34//!     // Legacy fragment APIs remain available for migration compatibility.
35//!     async fn save(&self, session: &SessionData) -> Result<()> { /* ... */ }
36//!     async fn load(&self, id: &str) -> Result<Option<SessionData>> { /* ... */ }
37//!     async fn delete(&self, id: &str) -> Result<()> { /* ... */ }
38//!     async fn list(&self) -> Result<Vec<String>> { /* ... */ }
39//!     async fn exists(&self, id: &str) -> Result<bool> { /* ... */ }
40//! }
41//! ```
42
43mod encryption;
44mod file_store;
45mod lease;
46mod memory_store;
47mod session_data;
48mod session_snapshot;
49mod wal;
50mod watch;
51
52#[cfg(test)]
53mod tests;
54
55pub use encryption::SessionStoreAtRestCipher;
56pub use file_store::FileSessionStore;
57pub use lease::{SessionStoreWriterLeaseV1, SESSION_STORE_WRITER_LEASE_SCHEMA_V1};
58pub use memory_store::MemorySessionStore;
59pub use session_data::{
60    ContextUsage, LlmConfigData, SessionConfig, SessionData, SessionState,
61    DEFAULT_AUTO_COMPACT_THRESHOLD,
62};
63pub use session_snapshot::{SessionSnapshotV1, SESSION_SNAPSHOT_SCHEMA_VERSION};
64pub use wal::{
65    snapshot_content_digest, FileSessionStoreWal, SessionStoreWalEntryV1, SessionStoreWalPhaseV1,
66    SESSION_STORE_WAL_ENTRY_SCHEMA_V1,
67};
68pub use watch::{
69    SessionStoreCommitEventV1, SessionStoreCommitWatch, SESSION_STORE_COMMIT_EVENT_SCHEMA_V1,
70};
71
72use crate::loop_checkpoint::LoopCheckpoint;
73use crate::run::RunRecord;
74use crate::subagent_task_tracker::SubagentTaskSnapshot;
75use crate::tools::ArtifactStore;
76use crate::trace::TraceEvent;
77use crate::verification::VerificationReport;
78use anyhow::{bail, Result};
79
80/// Persistence guarantees advertised by a session store implementation.
81///
82/// Hosts inspect these flags before relying on a durability semantics; a
83/// missing guarantee means the caller must arrange it above the store. The
84/// flags describe what an implementation already proves with its own tests,
85/// not aspirations (KRN-6).
86#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
87pub struct SessionStoreCapabilities {
88    /// A complete [`SessionSnapshotV1`] is committed as one atomic generation.
89    pub atomic_session_snapshots: bool,
90    /// Saving a generation whose identity already exists either replaces it
91    /// under one compare-and-swap decision or fails without a partial write;
92    /// concurrent writers never interleave generations.
93    pub aggregate_cas: bool,
94    /// Events are appended once with monotonically increasing sequence and
95    /// can be replayed after reopen; duplicates fail closed.
96    pub append_only_event_log: bool,
97    /// Cross-process writers are fenced by a lease so a stale process cannot
98    /// overwrite a newer generation after a takeover.
99    pub lease_fencing: bool,
100    /// Persisted bytes are encrypted at rest by the backend itself.
101    pub encrypted_at_rest: bool,
102    /// The store can notify watchers of committed generations.
103    pub watch: bool,
104    /// Artifact garbage collection is reference-aware: content reachable
105    /// from a retained provenance receipt, review finding, checkpoint, or
106    /// publication is never removed.
107    pub reference_aware_artifact_gc: bool,
108}
109
110// ============================================================================
111// Session Store Trait
112// ============================================================================
113
114/// Session storage trait
115#[async_trait::async_trait]
116pub trait SessionStore: Send + Sync {
117    /// Save session data
118    async fn save(&self, session: &SessionData) -> Result<()>;
119
120    /// Load session data by ID
121    async fn load(&self, id: &str) -> Result<Option<SessionData>>;
122
123    /// Delete session data
124    async fn delete(&self, id: &str) -> Result<()>;
125
126    /// List all session IDs
127    async fn list(&self) -> Result<Vec<String>>;
128
129    /// Check if session exists
130    async fn exists(&self, id: &str) -> Result<bool>;
131
132    /// Save a complete session generation.
133    ///
134    /// There is deliberately no fragmented default write. Silently mapping an
135    /// aggregate save to several independent writes would acknowledge a
136    /// generation that readers can observe only partially. Backends must make
137    /// their write semantics explicit by overriding this method.
138    async fn save_snapshot(&self, _snapshot: &SessionSnapshotV1) -> Result<()> {
139        bail!(
140            "session store '{}' does not support aggregate session snapshots",
141            self.backend_name()
142        )
143    }
144
145    /// Compare-and-swap save of one complete generation (STORE-CAS1).
146    ///
147    /// When `expected_current_digest` is `None`, the write is unconditional.
148    /// When `Some(digest)`, the currently durable snapshot for the same
149    /// session id must exist and match that digest; otherwise the method
150    /// returns `Ok(false)` without writing. Successful commits return
151    /// `Ok(true)`.
152    async fn save_snapshot_cas(
153        &self,
154        snapshot: &SessionSnapshotV1,
155        expected_current_digest: Option<&str>,
156    ) -> Result<bool> {
157        if expected_current_digest.is_some() {
158            bail!(
159                "session store '{}' does not support aggregate snapshot CAS",
160                self.backend_name()
161            );
162        }
163        self.save_snapshot(snapshot).await?;
164        Ok(true)
165    }
166
167    /// Acquire (or take over) the store-wide writer lease (STORE-LEASE1).
168    ///
169    /// Successful callers receive the new durable epoch. After a takeover,
170    /// any previously held epoch is stale and must fail closed on the next
171    /// fenced snapshot commit.
172    async fn acquire_writer_lease(&self, _holder_id: &str) -> Result<SessionStoreWriterLeaseV1> {
173        bail!(
174            "session store '{}' does not support writer lease fencing",
175            self.backend_name()
176        )
177    }
178
179    /// Read the currently durable writer lease, if any.
180    async fn writer_lease(&self) -> Result<Option<SessionStoreWriterLeaseV1>> {
181        Ok(None)
182    }
183
184    /// Subscribe to durable snapshot commit notifications (STORE-WATCH1).
185    ///
186    /// Events are published only after a complete generation is durable.
187    /// Lagged subscribers may skip intermediate commits.
188    async fn watch_commits(&self) -> Result<SessionStoreCommitWatch> {
189        bail!(
190            "session store '{}' does not support commit watch notifications",
191            self.backend_name()
192        )
193    }
194
195    /// Load one complete session generation.
196    ///
197    /// Legacy backends are assembled through the fragment APIs. This path is
198    /// best-effort and may observe concurrent fragment updates; callers can
199    /// inspect [`Self::capabilities`] before relying on atomicity.
200    async fn load_snapshot(&self, id: &str) -> Result<Option<SessionSnapshotV1>> {
201        let Some(session) = self.load(id).await? else {
202            return Ok(None);
203        };
204        let artifacts = self.load_artifacts(id).await?.unwrap_or_default();
205        Ok(Some(SessionSnapshotV1::new(
206            session,
207            &artifacts,
208            self.load_trace_events(id).await?.unwrap_or_default(),
209            self.load_run_records(id).await?.unwrap_or_default(),
210            self.load_verification_reports(id)
211                .await?
212                .unwrap_or_default(),
213            self.load_subagent_tasks(id).await?.unwrap_or_default(),
214        )))
215    }
216
217    /// Report persistence guarantees without requiring a write probe.
218    fn capabilities(&self) -> SessionStoreCapabilities {
219        SessionStoreCapabilities::default()
220    }
221
222    /// Save artifacts associated with a session.
223    async fn save_artifacts(&self, _id: &str, artifacts: &ArtifactStore) -> Result<()> {
224        if !artifacts.is_empty() {
225            bail!(
226                "session store '{}' does not support artifacts",
227                self.backend_name()
228            );
229        }
230        Ok(())
231    }
232
233    /// Load artifacts associated with a session.
234    async fn load_artifacts(&self, _id: &str) -> Result<Option<ArtifactStore>> {
235        Ok(None)
236    }
237
238    /// Save compact trace events associated with a session.
239    async fn save_trace_events(&self, _id: &str, events: &[TraceEvent]) -> Result<()> {
240        if !events.is_empty() {
241            bail!(
242                "session store '{}' does not support trace events",
243                self.backend_name()
244            );
245        }
246        Ok(())
247    }
248
249    /// Load compact trace events associated with a session.
250    async fn load_trace_events(&self, _id: &str) -> Result<Option<Vec<TraceEvent>>> {
251        Ok(None)
252    }
253
254    /// Save run snapshots and replayable runtime events associated with a session.
255    async fn save_run_records(&self, _id: &str, records: &[RunRecord]) -> Result<()> {
256        if !records.is_empty() {
257            bail!(
258                "session store '{}' does not support run records",
259                self.backend_name()
260            );
261        }
262        Ok(())
263    }
264
265    /// Load run snapshots and replayable runtime events associated with a session.
266    async fn load_run_records(&self, _id: &str) -> Result<Option<Vec<RunRecord>>> {
267        Ok(None)
268    }
269
270    /// Save structured verification reports associated with a session.
271    async fn save_verification_reports(
272        &self,
273        _id: &str,
274        reports: &[VerificationReport],
275    ) -> Result<()> {
276        if !reports.is_empty() {
277            bail!(
278                "session store '{}' does not support verification reports",
279                self.backend_name()
280            );
281        }
282        Ok(())
283    }
284
285    /// Load structured verification reports associated with a session.
286    async fn load_verification_reports(
287        &self,
288        _id: &str,
289    ) -> Result<Option<Vec<VerificationReport>>> {
290        Ok(None)
291    }
292
293    /// Save the session's delegated subagent task tracker snapshots.
294    ///
295    /// Cluster-grade hosts need this so a migrated session keeps a
296    /// queryable history of its delegated child runs. Cancellers are
297    /// **not** persisted — they are runtime-only and re-attaching them
298    /// is the executor's job at task respawn time.
299    async fn save_subagent_tasks(&self, _id: &str, tasks: &[SubagentTaskSnapshot]) -> Result<()> {
300        if !tasks.is_empty() {
301            bail!(
302                "session store '{}' does not support subagent tasks",
303                self.backend_name()
304            );
305        }
306        Ok(())
307    }
308
309    /// Load the session's delegated subagent task tracker snapshots.
310    async fn load_subagent_tasks(&self, _id: &str) -> Result<Option<Vec<SubagentTaskSnapshot>>> {
311        Ok(None)
312    }
313
314    /// Save the latest per-tool-round loop checkpoint for `run_id`.
315    ///
316    /// The agent loop calls this through the
317    /// [`SessionStoreCheckpointSink`](crate::loop_checkpoint::SessionStoreCheckpointSink)
318    /// adapter after each completed tool round. Implementations should
319    /// **overwrite** any earlier checkpoint for the same `run_id` — the
320    /// loop only ever needs the most recent boundary.
321    async fn save_loop_checkpoint(
322        &self,
323        _run_id: &str,
324        _checkpoint: &LoopCheckpoint,
325    ) -> Result<()> {
326        Ok(())
327    }
328
329    /// Load the latest loop checkpoint for `run_id`.
330    async fn load_loop_checkpoint(&self, _run_id: &str) -> Result<Option<LoopCheckpoint>> {
331        Ok(None)
332    }
333
334    /// Delete the loop checkpoint for `run_id`, if present.
335    ///
336    /// Called by the run lifecycle when a run reaches a terminal state
337    /// **in-process** (completed, failed, or cancelled) — at that point
338    /// the checkpoint is dead weight. Only a process crash (the agent
339    /// loop never returns) should leave a checkpoint behind for
340    /// crash-recovery resume. Without this, every tool-using run would
341    /// leak a checkpoint forever — the dominant unbounded-growth source
342    /// for long-running cluster deployments.
343    ///
344    /// Deleting a non-existent checkpoint is a no-op success.
345    async fn delete_loop_checkpoint(&self, _run_id: &str) -> Result<()> {
346        Ok(())
347    }
348
349    /// Persist a workflow checkpoint, overwriting any earlier one for the same
350    /// `workflow_id`. The resumable orchestration combinators call this at each
351    /// step boundary so an interrupted workflow resumes from the last
352    /// completed step (here or, after migration, on another node).
353    async fn save_workflow_checkpoint(
354        &self,
355        _workflow_id: &str,
356        _checkpoint: &crate::orchestration::WorkflowCheckpoint,
357    ) -> Result<()> {
358        Ok(())
359    }
360
361    /// Load the latest workflow checkpoint for `workflow_id`.
362    async fn load_workflow_checkpoint(
363        &self,
364        _workflow_id: &str,
365    ) -> Result<Option<crate::orchestration::WorkflowCheckpoint>> {
366        Ok(None)
367    }
368
369    /// Delete the workflow checkpoint for `workflow_id`, if present. Called
370    /// when a workflow reaches a terminal state in-process; only a crash should
371    /// leave one behind for resume. Deleting a non-existent checkpoint is a
372    /// no-op success.
373    async fn delete_workflow_checkpoint(&self, _workflow_id: &str) -> Result<()> {
374        Ok(())
375    }
376
377    /// Health check — verify the store backend is reachable and operational
378    async fn health_check(&self) -> Result<()> {
379        Ok(())
380    }
381
382    /// Backend name for diagnostics
383    fn backend_name(&self) -> &str {
384        "unknown"
385    }
386}