Skip to main content

oxios_kernel/
state_store.rs

1//! Filesystem-based state store.
2//!
3//! All state is persisted as markdown or JSON files organized
4//! by category. This is the "filesystem" of Oxios.
5
6use anyhow::{Result, bail};
7use chrono::{DateTime, Utc};
8use serde::{Deserialize, Deserializer, Serialize, Serializer, de::DeserializeOwned};
9use std::collections::HashMap;
10use std::path::PathBuf;
11use std::sync::Arc;
12use tokio::fs;
13use tokio::io::AsyncWriteExt;
14
15/// Unique identifier for a session.
16#[derive(Debug, Clone, PartialEq, Eq, Hash)]
17pub struct SessionId(pub String);
18
19impl SessionId {
20    /// Creates a new random session ID.
21    pub fn new() -> Self {
22        Self(uuid::Uuid::new_v4().to_string())
23    }
24}
25
26impl Default for SessionId {
27    fn default() -> Self {
28        Self::new()
29    }
30}
31
32impl std::fmt::Display for SessionId {
33    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
34        write!(f, "{}", self.0)
35    }
36}
37
38impl Serialize for SessionId {
39    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
40    where
41        S: Serializer,
42    {
43        serializer.serialize_str(&self.0)
44    }
45}
46
47impl<'de> Deserialize<'de> for SessionId {
48    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
49    where
50        D: Deserializer<'de>,
51    {
52        let s = String::deserialize(deserializer)?;
53        Ok(Self(s))
54    }
55}
56
57/// A user message in a session.
58#[derive(Debug, Clone, Serialize, Deserialize)]
59pub struct UserMessage {
60    /// Message content.
61    pub content: String,
62    /// Timestamp when the message was sent.
63    pub timestamp: DateTime<Utc>,
64}
65
66/// An agent response in a session.
67#[derive(Debug, Clone, Serialize, Deserialize)]
68pub struct AgentResponse {
69    /// Response content.
70    pub content: String,
71    /// Session ID associated with this response.
72    pub session_id: Option<String>,
73    /// Phase reached during orchestration.
74    pub phase_reached: Option<String>,
75    /// Whether evaluation passed.
76    pub evaluation_passed: Option<bool>,
77    /// Timestamp when the response was generated.
78    pub timestamp: DateTime<Utc>,
79    /// Index range into `Session::trajectory_steps` for tool calls that
80    /// occurred during this response. `None` when no tool calls were made.
81    /// Used by the Web UI to render per-turn execution timelines.
82    #[serde(skip_serializing_if = "Option::is_none")]
83    pub trajectory_range: Option<TrajectoryRange>,
84}
85
86/// Index range (exclusive end) into `Session::trajectory_steps`.
87#[derive(Debug, Clone, Serialize, Deserialize)]
88pub struct TrajectoryRange {
89    /// Start index (inclusive).
90    pub start: usize,
91    /// End index (exclusive).
92    pub end: usize,
93}
94
95/// A single tool execution step recorded in a session (RFC-015).
96///
97/// Persisted alongside the agent response so that the Web UI can render the
98/// execution timeline (tool calls, durations, errors) when the user
99/// re-opens the session later. Mirrors `memory::sona::TrajectoryStep` but
100/// is duplicated here to avoid a kernel-state → memory dependency cycle.
101#[derive(Debug, Clone, Serialize, Deserialize)]
102pub struct TrajectoryStepRecord {
103    /// Name of the tool that was called.
104    pub tool_name: String,
105    /// Tool input arguments (JSON).
106    pub tool_args: serde_json::Value,
107    /// Truncated output (max ~500 chars).
108    pub output_summary: String,
109    /// Wall-clock duration in milliseconds.
110    pub duration_ms: u64,
111    /// Whether the tool returned an error.
112    pub is_error: bool,
113    /// Provider-specific tool call ID (for start/end correlation).
114    pub tool_call_id: String,
115    /// Timestamp when the step started.
116    pub timestamp: DateTime<Utc>,
117}
118
119/// Arbitrary key-value metadata for a session.
120pub type SessionMetadata = std::collections::HashMap<String, serde_json::Value>;
121
122/// A session represents a single user conversation.
123///
124/// Sessions track the full message history and metadata for
125/// a user conversation. They are created per user interaction
126/// and persisted for later retrieval.
127#[derive(Debug, Clone, Serialize, Deserialize)]
128pub struct Session {
129    /// Unique session identifier.
130    pub id: SessionId,
131    /// User ID who owns this session.
132    pub user_id: String,
133    /// All user messages in this session.
134    #[serde(default)]
135    pub user_messages: Vec<UserMessage>,
136    /// All agent responses in this session.
137    #[serde(default)]
138    pub agent_responses: Vec<AgentResponse>,
139    /// RFC-015: tool execution trajectory accumulated for this session.
140    /// Appended on each orchestrator run; consumed by the Web UI to render
141    /// the execution timeline when the session is re-opened.
142    #[serde(default, skip_serializing_if = "Vec::is_empty")]
143    pub trajectory_steps: Vec<TrajectoryStepRecord>,
144    /// Currently active persona ID (for future multi-persona support).
145    #[serde(skip_serializing_if = "Option::is_none")]
146    pub active_persona_id: Option<String>,
147    /// RFC-025: Project this session belongs to (singular, grouping only).
148    /// Set by the sidebar/drag-to-reparent; consumed for Project-tree view.
149    #[serde(default, skip_serializing_if = "Option::is_none")]
150    pub project_id: Option<String>,
151    /// Timestamp when the session was created.
152    pub created_at: DateTime<Utc>,
153    /// Timestamp when the session was last updated.
154    pub updated_at: DateTime<Utc>,
155    /// Arbitrary key-value metadata.
156    #[serde(default)]
157    pub metadata: SessionMetadata,
158}
159
160impl Session {
161    /// Creates a new session for a user.
162    pub fn new(user_id: impl Into<String>) -> Self {
163        let now = Utc::now();
164        Self {
165            id: SessionId::new(),
166            user_id: user_id.into(),
167            user_messages: Vec::new(),
168            agent_responses: Vec::new(),
169            trajectory_steps: Vec::new(),
170            active_persona_id: None,
171            project_id: None,
172            created_at: now,
173            updated_at: now,
174            metadata: SessionMetadata::new(),
175        }
176    }
177
178    /// Creates a session with a specific ID.
179    pub fn with_id(user_id: impl Into<String>, session_id: SessionId) -> Self {
180        let now = Utc::now();
181        Self {
182            id: session_id,
183            user_id: user_id.into(),
184            user_messages: Vec::new(),
185            agent_responses: Vec::new(),
186            trajectory_steps: Vec::new(),
187            active_persona_id: None,
188            project_id: None,
189            created_at: now,
190            updated_at: now,
191            metadata: SessionMetadata::new(),
192        }
193    }
194
195    /// Adds a user message to the session.
196    pub fn add_user_message(&mut self, content: impl Into<String>) {
197        self.user_messages.push(UserMessage {
198            content: content.into(),
199            timestamp: Utc::now(),
200        });
201        self.updated_at = Utc::now();
202    }
203
204    /// Adds an agent response to the session.
205    pub fn add_agent_response(&mut self, response: AgentResponse) {
206        self.agent_responses.push(response);
207        self.updated_at = Utc::now();
208    }
209
210    /// Appends trajectory steps to the session (RFC-015).
211    ///
212    /// Called by the orchestrator after each run so the Web UI can
213    /// re-render the execution timeline when the user re-opens the session.
214    pub fn extend_trajectory(&mut self, steps: Vec<TrajectoryStepRecord>) {
215        if steps.is_empty() {
216            return;
217        }
218        self.trajectory_steps.extend(steps);
219        self.updated_at = Utc::now();
220    }
221
222    /// Returns the trajectory steps recorded in this session.
223    pub fn trajectory(&self) -> &[TrajectoryStepRecord] {
224        &self.trajectory_steps
225    }
226
227    /// Sets the active persona ID.
228    pub fn set_active_persona(&mut self, persona_id: Option<String>) {
229        self.active_persona_id = persona_id;
230        self.updated_at = Utc::now();
231    }
232
233    /// Sets a metadata value.
234    pub fn set_metadata(&mut self, key: impl Into<String>, value: serde_json::Value) {
235        self.metadata.insert(key.into(), value);
236        self.updated_at = Utc::now();
237    }
238
239    /// Gets a metadata value.
240    pub fn get_metadata(&self, key: &str) -> Option<&serde_json::Value> {
241        self.metadata.get(key)
242    }
243
244    /// Returns the total number of exchanges in this session.
245    pub fn exchange_count(&self) -> usize {
246        self.user_messages.len().min(self.agent_responses.len())
247    }
248
249    /// Returns true if the session is empty (no messages).
250    pub fn is_empty(&self) -> bool {
251        self.user_messages.is_empty()
252    }
253}
254/// A filesystem-based persistent state store.
255///
256/// Files are organized as `<base_path>/<category>/<name>.md` or
257/// `<base_path>/<category>/<name>.json`.
258#[derive(Clone)]
259pub struct StateStore {
260    /// Root directory for all state files.
261    pub base_path: PathBuf,
262    /// Per-session write locks used by [`StateStore::update_session_with`]
263    /// to serialize concurrent load-modify-save cycles on the same session
264    /// (state-area F1). Legacy callers that do their own
265    /// `load_session` → `save_session` without this primitive remain racy
266    /// and should migrate.
267    session_locks: Arc<parking_lot::RwLock<HashMap<String, Arc<tokio::sync::Mutex<()>>>>>,
268}
269
270impl StateStore {
271    /// Creates a new state store, initializing the directory if needed.
272    ///
273    /// # Example
274    ///
275    /// ```no_run
276    /// use oxios_kernel::state_store::StateStore;
277    /// use std::path::PathBuf;
278    ///
279    /// let store = StateStore::new(PathBuf::from("/tmp/oxios-state")).unwrap();
280    /// ```
281    pub fn new(base_path: PathBuf) -> Result<Self> {
282        Ok(Self {
283            base_path,
284            session_locks: Arc::new(parking_lot::RwLock::new(HashMap::new())),
285        })
286    }
287
288    /// Validate that a category name does not contain path traversal.
289    fn validate_category(category: &str) -> Result<()> {
290        if category.contains("..") || category.contains('\\') {
291            bail!("invalid category name: '{category}'");
292        }
293        if category.is_empty()
294            || category.starts_with('/')
295            || category.ends_with('/')
296            || category.contains("//")
297        {
298            bail!("invalid category name: '{category}'");
299        }
300        Ok(())
301    }
302
303    /// Validate that a file name does not contain path traversal.
304    fn validate_name(name: &str) -> Result<()> {
305        if name.contains("..") || name.contains('/') || name.contains('\\') {
306            bail!("invalid file name: '{name}'");
307        }
308        Ok(())
309    }
310
311    /// Durable atomic write: temp file → fsync(file) → rename → fsync(dir).
312    ///
313    /// `rename(2)` only guarantees *metadata* atomicity, not data
314    /// durability: if the data pages haven't been flushed before the
315    /// rename's metadata commit, a crash can surface the new filename
316    /// with stale or zero contents. fsyncing the temp file before the
317    /// rename, and the parent directory after, closes that window. This
318    /// matters because StateStore is the source of truth for sessions,
319    /// agents, and project metadata. (state-area F9.)
320    async fn durable_write(
321        dir: &std::path::Path,
322        target: &std::path::Path,
323        content: &[u8],
324    ) -> Result<()> {
325        let file_name = target
326            .file_name()
327            .and_then(|n| n.to_str())
328            .unwrap_or("state");
329        let temp_path = dir.join(format!(
330            "{file_name}.{}.{}.tmp",
331            std::process::id(),
332            uuid::Uuid::new_v4()
333        ));
334        {
335            let mut file = fs::File::create(&temp_path).await?;
336            file.write_all(content).await?;
337            file.sync_all().await?;
338        }
339        tokio::fs::rename(&temp_path, target).await?;
340        // Best-effort directory fsync; ignore errors (not all platforms
341        // support it, and the file fsync + rename already did the work).
342        if let Ok(dir_file) = fs::File::open(dir).await {
343            let _ = dir_file.sync_all().await;
344        }
345        Ok(())
346    }
347
348    pub async fn save_markdown(&self, category: &str, name: &str, content: &str) -> Result<()> {
349        Self::validate_category(category)?;
350        Self::validate_name(name)?;
351        let dir = self.base_path.join(category);
352        fs::create_dir_all(&dir).await?;
353        let path = dir.join(format!("{name}.md"));
354        Self::durable_write(&dir, &path, content.as_bytes()).await?;
355        Ok(())
356    }
357
358    /// Load a markdown file from the given category.
359    pub async fn load_markdown(&self, category: &str, name: &str) -> Result<Option<String>> {
360        Self::validate_category(category)?;
361        Self::validate_name(name)?;
362        let path = self.base_path.join(category).join(format!("{name}.md"));
363        match fs::read_to_string(&path).await {
364            Ok(content) => Ok(Some(content)),
365            Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(None),
366            Err(e) => Err(e.into()),
367        }
368    }
369
370    /// List all markdown files in a category (names without extension).
371    pub async fn list_category(&self, category: &str) -> Result<Vec<String>> {
372        Self::validate_category(category)?;
373        let dir = self.base_path.join(category);
374        if !dir.exists() {
375            return Ok(Vec::new());
376        }
377        let mut entries = fs::read_dir(&dir).await?;
378        let mut names = Vec::new();
379        while let Some(entry) = entries.next_entry().await? {
380            let path = entry.path();
381            if let Some(ext) = path.extension()
382                && (ext == "md" || ext == "json")
383                && let Some(stem) = path.file_stem()
384            {
385                names.push(stem.to_string_lossy().into_owned());
386            }
387        }
388        names.sort();
389        Ok(names)
390    }
391
392    /// Save a serializable value as JSON under the given category.
393    pub async fn save_json<T: Serialize>(
394        &self,
395        category: &str,
396        name: &str,
397        data: &T,
398    ) -> Result<()> {
399        Self::validate_category(category)?;
400        Self::validate_name(name)?;
401        let dir = self.base_path.join(category);
402        fs::create_dir_all(&dir).await?;
403        let path = dir.join(format!("{name}.json"));
404        let content = serde_json::to_string_pretty(data)?;
405        Self::durable_write(&dir, &path, content.as_bytes()).await?;
406        Ok(())
407    }
408
409    /// Load a deserializable value from JSON in the given category.
410    pub async fn load_json<T: DeserializeOwned>(
411        &self,
412        category: &str,
413        name: &str,
414    ) -> Result<Option<T>> {
415        Self::validate_category(category)?;
416        Self::validate_name(name)?;
417        let path = self.base_path.join(category).join(format!("{name}.json"));
418        match fs::read_to_string(&path).await {
419            Ok(content) => Ok(Some(serde_json::from_str(&content)?)),
420            Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(None),
421            Err(e) => Err(e.into()),
422        }
423    }
424
425    /// Delete a file from the given category.
426    pub async fn delete_file(&self, category: &str, name: &str) -> Result<bool> {
427        Self::validate_category(category)?;
428        Self::validate_name(name)?;
429        let path = self.base_path.join(category).join(format!("{name}.json"));
430        if path.exists() {
431            tokio::fs::remove_file(path).await?;
432            Ok(true)
433        } else {
434            let path = self.base_path.join(category).join(format!("{name}.md"));
435            if path.exists() {
436                tokio::fs::remove_file(path).await?;
437                Ok(true)
438            } else {
439                Ok(false)
440            }
441        }
442    }
443}
444
445impl std::fmt::Debug for StateStore {
446    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
447        f.debug_struct("StateStore")
448            .field("base_path", &self.base_path)
449            .finish()
450    }
451}
452
453impl StateStore {
454    /// Saves a session to the sessions category.
455    pub async fn save_session(&self, session: &Session) -> Result<()> {
456        self.save_json("sessions", &session.id.0, session).await
457    }
458
459    /// Atomically load, mutate, and save a session under a per-session lock.
460    ///
461    /// This is the safe primitive for read-modify-write on a session: it
462    /// serializes concurrent modifications to the same session, preventing
463    /// the lost-update race that plain `load_session` → `save_session`
464    /// sequences suffer when two callers read the same starting copy and
465    /// the later save clobbers the earlier one (state-area F1).
466    ///
467    /// Returns `Ok(Some(session))` after the mutation+save, or `Ok(None)`
468    /// if the session did not exist.
469    pub async fn update_session_with<F>(
470        &self,
471        session_id: &SessionId,
472        f: F,
473    ) -> Result<Option<Session>>
474    where
475        F: FnOnce(&mut Session) -> Result<()>,
476    {
477        let lock = Self::session_lock(&self.session_locks, &session_id.0);
478        let _guard = lock.lock().await;
479        let mut session = match self.load_session(session_id).await? {
480            Some(s) => s,
481            None => return Ok(None),
482        };
483        f(&mut session)?;
484        self.save_session(&session).await?;
485        Ok(Some(session))
486    }
487
488    /// Get (or lazily create) the per-session mutex. Double-checked under
489    /// the map's read/write guards so concurrent callers for the same
490    /// session share one `Arc<Mutex<()>>`.
491    fn session_lock(
492        map: &parking_lot::RwLock<HashMap<String, Arc<tokio::sync::Mutex<()>>>>,
493        session_id: &str,
494    ) -> Arc<tokio::sync::Mutex<()>> {
495        // Fast path: shared read.
496        if let Some(lock) = map.read().get(session_id) {
497            return Arc::clone(lock);
498        }
499        // Slow path: exclusive write, insert if absent.
500        Arc::clone(
501            map.write()
502                .entry(session_id.to_string())
503                .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))),
504        )
505    }
506
507    /// Saves a session and then runs pruning if auto_prune is enabled.
508    pub async fn save_session_with_prune(
509        &self,
510        session: &Session,
511        prune_config: &PruneConfig,
512    ) -> Result<()> {
513        self.save_session(session).await?;
514        // Prune in the background — don't block the response
515        let store = self.clone();
516        let config = prune_config.clone();
517        tokio::spawn(async move {
518            if let Err(e) = store.prune_sessions(&config).await {
519                tracing::warn!(error = %e, "Background session pruning failed");
520            }
521        });
522        Ok(())
523    }
524
525    /// Loads a session by ID.
526    pub async fn load_session(&self, session_id: &SessionId) -> Result<Option<Session>> {
527        self.load_json("sessions", &session_id.0).await
528    }
529
530    /// RFC-025 Phase 5: Load all sessions in full (messages + trajectories).
531    ///
532    /// Used by the Mount auto-promotion scanner, which needs trajectory
533    /// tool_args to identify paths the agent actually worked on. Cheaper to
534    /// call once per scan than `load_session` per id.
535    ///
536    /// **Warning:** loads every session on disk. Prefer
537    /// [`load_sessions_for_promotion`](Self::load_sessions_for_promotion)
538    /// for the promotion scanner to bound memory usage.
539    pub async fn load_all_sessions(&self) -> Result<Vec<Session>> {
540        let mut sessions = Vec::new();
541        if let Ok(names) = self.list_category("sessions").await {
542            for name in names {
543                if let Ok(Some(session)) = self.load_json::<Session>("sessions", &name).await {
544                    sessions.push(session);
545                }
546            }
547        }
548        Ok(sessions)
549    }
550
551    /// RFC-025 Phase 5: Load only sessions updated at or after `since`.
552    ///
553    /// Bounds memory usage for the Mount auto-promotion scanner: instead of
554    /// deserializing every session into a single `Vec`, we deserialize each
555    /// one and skip it immediately if its `updated_at` predates the cutoff.
556    /// The promotion window is bounded by `PromotionConfig::window_days`, so
557    /// loading older sessions is pure waste.
558    ///
559    /// Sessions that fail to deserialize are skipped (best effort), mirroring
560    /// [`load_all_sessions`](Self::load_all_sessions).
561    pub async fn load_sessions_for_promotion(&self, since: DateTime<Utc>) -> Result<Vec<Session>> {
562        let mut sessions = Vec::new();
563        if let Ok(names) = self.list_category("sessions").await {
564            for name in names {
565                if let Ok(Some(session)) = self.load_json::<Session>("sessions", &name).await {
566                    // Skip sessions outside the promotion window *before*
567                    // collecting — this is the whole point of the cutoff.
568                    if session.updated_at < since {
569                        continue;
570                    }
571                    sessions.push(session);
572                }
573            }
574        }
575        Ok(sessions)
576    }
577
578    /// Lists all sessions (sorted by updated_at descending).
579    pub async fn list_sessions(&self) -> Result<Vec<SessionSummary>> {
580        let mut sessions = Vec::new();
581
582        if let Ok(names) = self.list_category("sessions").await {
583            for name in names {
584                if let Ok(Some(session)) = self.load_json::<Session>("sessions", &name).await {
585                    sessions.push(SessionSummary {
586                        id: session.id.0.clone(),
587                        user_id: session.user_id.clone(),
588                        message_count: session.user_messages.len(),
589                        title: session
590                            .metadata
591                            .get("title")
592                            .and_then(|v| v.as_str())
593                            .map(String::from)
594                            .or_else(|| {
595                                // Auto-generate from first user message
596                                session.user_messages.first().map(|m| {
597                                    let s = m.content.lines().next().unwrap_or("");
598                                    if s.len() > 60 {
599                                        format!("{}…", &s[..s.ceil_char_boundary(59)])
600                                    } else {
601                                        s.to_string()
602                                    }
603                                })
604                            }),
605                        project_id: session
606                            .project_id
607                            .clone()
608                            // Backward-compat: fall back to legacy metadata keys.
609                            .or_else(|| {
610                                session
611                                    .metadata
612                                    .get("project_id")
613                                    .and_then(|v| v.as_str())
614                                    .map(String::from)
615                            })
616                            .or_else(|| {
617                                session
618                                    .metadata
619                                    .get("project_ids")
620                                    .and_then(|v| v.as_str())
621                                    .and_then(|s| s.split(',').next().map(String::from))
622                            }),
623                        created_at: session.created_at,
624                        updated_at: session.updated_at,
625                    });
626                }
627            }
628        }
629
630        // Sort by updated_at descending (most recent first)
631        sessions.sort_by_key(|b| std::cmp::Reverse(b.updated_at));
632        Ok(sessions)
633    }
634
635    /// Deletes a session by ID.
636    pub async fn delete_session(&self, session_id: &SessionId) -> Result<bool> {
637        let path = self
638            .base_path
639            .join("sessions")
640            .join(format!("{}.json", session_id.0));
641        match fs::remove_file(&path).await {
642            Ok(()) => {
643                tracing::info!(session_id = %session_id, "Session deleted");
644                Ok(true)
645            }
646            Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(false),
647            Err(e) => Err(e.into()),
648        }
649    }
650
651    /// Gets or creates a session for a user, initializing with the given session ID.
652    pub async fn get_or_create_session(
653        &self,
654        user_id: &str,
655        session_id: Option<&SessionId>,
656    ) -> Result<Session> {
657        if let Some(sid) = session_id
658            && let Some(existing) = self.load_session(sid).await?
659        {
660            return Ok(existing);
661        }
662
663        // Create new session
664        let session = match session_id {
665            Some(sid) => Session::with_id(user_id, sid.clone()),
666            None => Session::new(user_id),
667        };
668
669        self.save_session(&session).await?;
670        Ok(session)
671    }
672
673    /// Updates an existing session, saving it to disk.
674    pub async fn update_session(&self, session: &Session) -> Result<()> {
675        self.save_session(session).await
676    }
677
678    /// RFC-025: Move a session to a different Project (drag-to-reparent).
679    ///
680    /// Pass `None` to unassign (move to "unfiled").
681    pub async fn move_session_to_project(
682        &self,
683        session_id: &SessionId,
684        project_id: Option<&str>,
685    ) -> Result<bool> {
686        // Use update_session_with so concurrent modifications to the same
687        // session can't lose this project reassignment (state-area F1).
688        let project_id_owned = project_id.map(String::from);
689        let updated = self
690            .update_session_with(session_id, |session| {
691                session.project_id = project_id_owned;
692                session.updated_at = Utc::now();
693                Ok(())
694            })
695            .await?;
696        Ok(updated.is_some())
697    }
698
699    /// Prune sessions based on configuration.
700    ///
701    /// Removes sessions that exceed TTL or exceed the maximum count.
702    /// Returns the number of sessions pruned.
703    pub async fn prune_sessions(&self, config: &PruneConfig) -> Result<usize> {
704        let mut sessions = self.list_sessions().await?;
705        let mut pruned = 0;
706
707        // TTL-based pruning: remove sessions older than ttl_hours
708        if config.ttl_hours > 0 {
709            let cutoff = Utc::now() - chrono::Duration::hours(config.ttl_hours as i64);
710            let to_prune_ttl: std::collections::HashSet<String> = sessions
711                .iter()
712                .filter(|s| s.updated_at < cutoff)
713                .map(|s| s.id.clone())
714                .collect();
715
716            for id in &to_prune_ttl {
717                let sid = SessionId(id.clone());
718                if self.delete_session(&sid).await.is_ok() {
719                    pruned += 1;
720                }
721            }
722
723            // Remove pruned sessions from the list for count-based pruning
724            sessions.retain(|s| !to_prune_ttl.contains(&s.id));
725        }
726
727        // Count-based pruning: keep only the most recent `max_sessions`
728        if config.max_sessions > 0 && sessions.len() > config.max_sessions {
729            // sessions are already sorted by updated_at descending
730            let excess = sessions.len() - config.max_sessions;
731            for session in sessions.into_iter().rev().take(excess) {
732                let sid = SessionId(session.id);
733                if self.delete_session(&sid).await.is_ok() {
734                    pruned += 1;
735                }
736            }
737        }
738
739        if pruned > 0 {
740            tracing::info!(pruned = pruned, "Session pruning completed");
741        }
742
743        Ok(pruned)
744    }
745
746    /// Prune agent records based on config.
747    ///
748    /// 1. TTL-based: delete agents with created_at older than ttl_hours.
749    /// 2. Count-based: if still over max_entries, delete oldest.
750    pub async fn prune_agents_by_config(
751        &self,
752        max_entries: usize,
753        ttl_hours: u64,
754        batch_size: usize,
755    ) -> Result<usize> {
756        let mut pruned = 0usize;
757
758        let names = self.list_category("agents").await?;
759        if names.is_empty() {
760            return Ok(0);
761        }
762
763        let now = Utc::now();
764
765        // 1. TTL-based pruning
766        let mut remaining: Vec<(String, DateTime<Utc>)> = Vec::with_capacity(names.len());
767
768        if ttl_hours > 0 {
769            let cutoff = now - chrono::Duration::hours(ttl_hours as i64);
770            for name in &names {
771                // Load just enough to check created_at
772                if let Ok(Some(info)) = self
773                    .load_json::<crate::types::AgentInfo>("agents", name)
774                    .await
775                {
776                    if info.created_at < cutoff {
777                        if self.delete_file("agents", name).await.unwrap_or(false) {
778                            pruned += 1;
779                        }
780                    } else {
781                        remaining.push((name.clone(), info.created_at));
782                    }
783                }
784            }
785        } else {
786            // Load all created_at for count-based pruning
787            for name in &names {
788                if let Ok(Some(info)) = self
789                    .load_json::<crate::types::AgentInfo>("agents", name)
790                    .await
791                {
792                    remaining.push((name.clone(), info.created_at));
793                }
794            }
795        }
796
797        // 2. Count-based pruning
798        if max_entries > 0 && remaining.len() > max_entries {
799            // Sort by created_at ascending (oldest first)
800            remaining.sort_by_key(|a| a.1);
801
802            let excess = remaining.len() - max_entries;
803            let to_delete = excess.min(batch_size);
804
805            for (name, _) in remaining.iter().take(to_delete) {
806                if self.delete_file("agents", name).await.unwrap_or(false) {
807                    pruned += 1;
808                }
809            }
810        }
811
812        if pruned > 0 {
813            tracing::info!(pruned = pruned, "Agent filesystem pruning completed");
814        }
815
816        Ok(pruned)
817    }
818}
819
820/// Summary of a session for listing (without full message history).
821#[derive(Debug, Clone, Serialize, Deserialize)]
822pub struct SessionSummary {
823    /// Session ID.
824    pub id: String,
825    /// User ID who owns this session.
826    pub user_id: String,
827    /// Number of messages in this session.
828    pub message_count: usize,
829    /// Auto-generated title for this session. Derived from the first user
830    /// message (truncated to ~60 chars) when not explicitly set.
831    #[serde(skip_serializing_if = "Option::is_none")]
832    pub title: Option<String>,
833    /// Active project ID(s) this session belongs to.
834    #[serde(skip_serializing_if = "Option::is_none")]
835    pub project_id: Option<String>,
836    /// When the session was created.
837    pub created_at: DateTime<Utc>,
838    /// When the session was last updated.
839    pub updated_at: DateTime<Utc>,
840}
841
842/// Configuration for session pruning.
843#[derive(Debug, Clone)]
844pub struct PruneConfig {
845    /// Maximum number of sessions to keep. 0 = unlimited.
846    pub max_sessions: usize,
847    /// TTL in hours. Sessions older than this are pruned. 0 = no TTL.
848    pub ttl_hours: u64,
849}
850
851impl Default for PruneConfig {
852    fn default() -> Self {
853        Self {
854            max_sessions: 100,
855            ttl_hours: 168, // 7 days
856        }
857    }
858}
859
860/// Tracks the last time a prune was performed, enabling cooldown.
861pub struct PruneThrottle {
862    /// Instant of the last prune (monotonic).
863    last_prune: std::sync::Mutex<Option<std::time::Instant>>,
864    /// Minimum seconds between prune runs.
865    cooldown_secs: u64,
866}
867
868impl PruneThrottle {
869    /// Create a new throttle with the given cooldown.
870    pub fn new(cooldown_secs: u64) -> Self {
871        Self {
872            last_prune: std::sync::Mutex::new(None),
873            cooldown_secs,
874        }
875    }
876
877    /// Check if enough time has elapsed since the last prune.
878    /// Returns `true` if prune should proceed.
879    pub fn should_prune(&self) -> bool {
880        // std::sync::Mutex can poison if a holder panics; recover here by
881        // taking the inner value so pruning continues.
882        let mut guard = self.last_prune.lock().unwrap_or_else(|e| {
883            tracing::warn!("PruneThrottle mutex poisoned, recovering: {e}");
884            e.into_inner()
885        });
886        let now = std::time::Instant::now();
887        match *guard {
888            Some(last) => {
889                if now.duration_since(last).as_secs() >= self.cooldown_secs {
890                    *guard = Some(now);
891                    true
892                } else {
893                    false
894                }
895            }
896            None => {
897                *guard = Some(now);
898                true
899            }
900        }
901    }
902}
903
904#[cfg(test)]
905mod tests {
906    use super::*;
907    #[tokio::test]
908    async fn test_session_creation_and_persistence() {
909        let temp_dir = tempfile::tempdir().unwrap();
910        let store = StateStore::new(temp_dir.path().to_path_buf()).unwrap();
911
912        // Create a session
913        let mut session = Session::new("user-123");
914        session.add_user_message("Hello");
915
916        // Save and load
917        store.save_session(&session).await.unwrap();
918        let loaded = store.load_session(&session.id).await.unwrap();
919        assert!(loaded.is_some());
920        let loaded = loaded.unwrap();
921        assert_eq!(loaded.user_id, "user-123");
922        assert_eq!(loaded.user_messages.len(), 1);
923    }
924
925    #[tokio::test]
926    async fn test_session_list_sorts_by_updated() {
927        let temp_dir = tempfile::tempdir().unwrap();
928        let store = StateStore::new(temp_dir.path().to_path_buf()).unwrap();
929
930        // Create multiple sessions
931        for i in 0..3 {
932            let mut session = Session::new(format!("user-{}", i));
933            session.add_user_message(format!("Message {}", i));
934            store.save_session(&session).await.unwrap();
935            tokio::time::sleep(std::time::Duration::from_millis(10)).await;
936        }
937
938        let sessions = store.list_sessions().await.unwrap();
939        assert_eq!(sessions.len(), 3);
940        // Most recently updated should be first
941        assert_eq!(sessions[0].user_id, "user-2");
942    }
943
944    #[tokio::test]
945    async fn test_delete_session() {
946        let temp_dir = tempfile::tempdir().unwrap();
947        let store = StateStore::new(temp_dir.path().to_path_buf()).unwrap();
948
949        let session = Session::new("user-123");
950        store.save_session(&session).await.unwrap();
951
952        // Delete and verify
953        let deleted = store.delete_session(&session.id).await.unwrap();
954        assert!(deleted);
955
956        let loaded = store.load_session(&session.id).await.unwrap();
957        assert!(loaded.is_none());
958    }
959
960    #[tokio::test]
961    async fn test_get_or_create_session_existing() {
962        let temp_dir = tempfile::tempdir().unwrap();
963        let store = StateStore::new(temp_dir.path().to_path_buf()).unwrap();
964
965        let mut existing = Session::new("user-123");
966        existing.add_user_message("Original message");
967        store.save_session(&existing).await.unwrap();
968
969        // Get or create with same ID should return existing
970        let retrieved = store
971            .get_or_create_session("user-123", Some(&existing.id))
972            .await
973            .unwrap();
974        assert_eq!(retrieved.id, existing.id);
975        assert_eq!(retrieved.user_messages.len(), 1);
976    }
977
978    #[tokio::test]
979    async fn test_get_or_create_session_new() {
980        let temp_dir = tempfile::tempdir().unwrap();
981        let store = StateStore::new(temp_dir.path().to_path_buf()).unwrap();
982
983        // Get or create without existing session should create new
984        let session = store.get_or_create_session("user-456", None).await.unwrap();
985        assert_eq!(session.user_id, "user-456");
986        assert!(session.user_messages.is_empty());
987    }
988
989    #[tokio::test]
990    async fn test_prune_sessions_by_count() {
991        let temp_dir = tempfile::tempdir().unwrap();
992        let store = StateStore::new(temp_dir.path().to_path_buf()).unwrap();
993
994        // Create 5 sessions
995        for i in 0..5 {
996            let mut session = Session::new(format!("user-{}", i));
997            session.add_user_message(format!("Message {}", i));
998            store.save_session(&session).await.unwrap();
999            tokio::time::sleep(std::time::Duration::from_millis(10)).await;
1000        }
1001
1002        // Prune to max 3
1003        let config = PruneConfig {
1004            max_sessions: 3,
1005            ttl_hours: 0,
1006        };
1007        let pruned = store.prune_sessions(&config).await.unwrap();
1008        assert_eq!(pruned, 2);
1009
1010        let remaining = store.list_sessions().await.unwrap();
1011        assert_eq!(remaining.len(), 3);
1012        // Oldest sessions (user-0, user-1) should be pruned
1013        let remaining_ids: Vec<&str> = remaining.iter().map(|s| s.user_id.as_str()).collect();
1014        assert!(remaining_ids.contains(&"user-2"));
1015        assert!(remaining_ids.contains(&"user-3"));
1016        assert!(remaining_ids.contains(&"user-4"));
1017    }
1018
1019    #[tokio::test]
1020    async fn test_prune_sessions_by_ttl() {
1021        let temp_dir = tempfile::tempdir().unwrap();
1022        let store = StateStore::new(temp_dir.path().to_path_buf()).unwrap();
1023
1024        // Create a session and manually set updated_at to the past
1025        let mut old_session = Session::new("old-user");
1026        old_session.updated_at = Utc::now() - chrono::Duration::hours(48);
1027        store.save_session(&old_session).await.unwrap();
1028
1029        // Create a recent session
1030        let mut recent_session = Session::new("recent-user");
1031        recent_session.add_user_message("Hello");
1032        store.save_session(&recent_session).await.unwrap();
1033
1034        // Prune with 24-hour TTL
1035        let config = PruneConfig {
1036            max_sessions: 0,
1037            ttl_hours: 24,
1038        };
1039        let pruned = store.prune_sessions(&config).await.unwrap();
1040        assert_eq!(pruned, 1);
1041
1042        let remaining = store.list_sessions().await.unwrap();
1043        assert_eq!(remaining.len(), 1);
1044        assert_eq!(remaining[0].user_id, "recent-user");
1045    }
1046
1047    #[tokio::test]
1048    async fn test_load_sessions_for_promotion_filters_by_cutoff() {
1049        // Promo-1: sessions older than the cutoff must be skipped before
1050        // being collected, bounding memory for the promotion scanner.
1051        let temp_dir = tempfile::tempdir().unwrap();
1052        let store = StateStore::new(temp_dir.path().to_path_buf()).unwrap();
1053
1054        // Old session — updated 48h ago.
1055        let mut old_session = Session::new("old-user");
1056        old_session.updated_at = Utc::now() - chrono::Duration::hours(48);
1057        store.save_session(&old_session).await.unwrap();
1058
1059        // Recent session — updated now.
1060        let recent_session = Session::new("recent-user");
1061        store.save_session(&recent_session).await.unwrap();
1062
1063        // Cutoff 24h ago: only the recent session should be loaded.
1064        let cutoff = Utc::now() - chrono::Duration::hours(24);
1065        let sessions = store.load_sessions_for_promotion(cutoff).await.unwrap();
1066        assert_eq!(sessions.len(), 1, "old session must be filtered out");
1067        assert_eq!(sessions[0].user_id, "recent-user");
1068
1069        // A cutoff far in the future returns everything (boundary check).
1070        let far_cutoff = Utc::now() - chrono::Duration::days(365);
1071        let all = store.load_sessions_for_promotion(far_cutoff).await.unwrap();
1072        assert_eq!(all.len(), 2);
1073    }
1074}