1use 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#[derive(Debug, Clone, PartialEq, Eq, Hash)]
17pub struct SessionId(pub String);
18
19impl SessionId {
20 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#[derive(Debug, Clone, Serialize, Deserialize)]
59pub struct UserMessage {
60 pub content: String,
62 pub timestamp: DateTime<Utc>,
64}
65
66#[derive(Debug, Clone, Serialize, Deserialize)]
68pub struct AgentResponse {
69 pub content: String,
71 pub session_id: Option<String>,
73 pub phase_reached: Option<String>,
75 pub evaluation_passed: Option<bool>,
77 pub timestamp: DateTime<Utc>,
79 #[serde(skip_serializing_if = "Option::is_none")]
83 pub trajectory_range: Option<TrajectoryRange>,
84}
85
86#[derive(Debug, Clone, Serialize, Deserialize)]
88pub struct TrajectoryRange {
89 pub start: usize,
91 pub end: usize,
93}
94
95#[derive(Debug, Clone, Serialize, Deserialize)]
102pub struct TrajectoryStepRecord {
103 pub tool_name: String,
105 pub tool_args: serde_json::Value,
107 pub output_summary: String,
109 pub duration_ms: u64,
111 pub is_error: bool,
113 pub tool_call_id: String,
115 pub timestamp: DateTime<Utc>,
117}
118
119pub type SessionMetadata = std::collections::HashMap<String, serde_json::Value>;
121
122#[derive(Debug, Clone, Serialize, Deserialize)]
128pub struct Session {
129 pub id: SessionId,
131 pub user_id: String,
133 #[serde(default)]
135 pub user_messages: Vec<UserMessage>,
136 #[serde(default)]
138 pub agent_responses: Vec<AgentResponse>,
139 #[serde(default, skip_serializing_if = "Vec::is_empty")]
143 pub trajectory_steps: Vec<TrajectoryStepRecord>,
144 #[serde(skip_serializing_if = "Option::is_none")]
146 pub active_persona_id: Option<String>,
147 #[serde(default, skip_serializing_if = "Option::is_none")]
150 pub project_id: Option<String>,
151 pub created_at: DateTime<Utc>,
153 pub updated_at: DateTime<Utc>,
155 #[serde(default)]
157 pub metadata: SessionMetadata,
158}
159
160impl Session {
161 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 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 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 pub fn add_agent_response(&mut self, response: AgentResponse) {
206 self.agent_responses.push(response);
207 self.updated_at = Utc::now();
208 }
209
210 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 pub fn trajectory(&self) -> &[TrajectoryStepRecord] {
224 &self.trajectory_steps
225 }
226
227 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 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 pub fn get_metadata(&self, key: &str) -> Option<&serde_json::Value> {
241 self.metadata.get(key)
242 }
243
244 pub fn exchange_count(&self) -> usize {
246 self.user_messages.len().min(self.agent_responses.len())
247 }
248
249 pub fn is_empty(&self) -> bool {
251 self.user_messages.is_empty()
252 }
253}
254#[derive(Clone)]
259pub struct StateStore {
260 pub base_path: PathBuf,
262 session_locks: Arc<parking_lot::RwLock<HashMap<String, Arc<tokio::sync::Mutex<()>>>>>,
268}
269
270impl StateStore {
271 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 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 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 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 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 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 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 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 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 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 pub async fn save_session(&self, session: &Session) -> Result<()> {
456 self.save_json("sessions", &session.id.0, session).await
457 }
458
459 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 fn session_lock(
492 map: &parking_lot::RwLock<HashMap<String, Arc<tokio::sync::Mutex<()>>>>,
493 session_id: &str,
494 ) -> Arc<tokio::sync::Mutex<()>> {
495 if let Some(lock) = map.read().get(session_id) {
497 return Arc::clone(lock);
498 }
499 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 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 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 pub async fn load_session(&self, session_id: &SessionId) -> Result<Option<Session>> {
527 self.load_json("sessions", &session_id.0).await
528 }
529
530 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 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 if session.updated_at < since {
569 continue;
570 }
571 sessions.push(session);
572 }
573 }
574 }
575 Ok(sessions)
576 }
577
578 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 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 .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 sessions.sort_by_key(|b| std::cmp::Reverse(b.updated_at));
632 Ok(sessions)
633 }
634
635 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 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 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 pub async fn update_session(&self, session: &Session) -> Result<()> {
675 self.save_session(session).await
676 }
677
678 pub async fn move_session_to_project(
682 &self,
683 session_id: &SessionId,
684 project_id: Option<&str>,
685 ) -> Result<bool> {
686 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 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 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 sessions.retain(|s| !to_prune_ttl.contains(&s.id));
725 }
726
727 if config.max_sessions > 0 && sessions.len() > config.max_sessions {
729 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 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 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 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 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 if max_entries > 0 && remaining.len() > max_entries {
799 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#[derive(Debug, Clone, Serialize, Deserialize)]
822pub struct SessionSummary {
823 pub id: String,
825 pub user_id: String,
827 pub message_count: usize,
829 #[serde(skip_serializing_if = "Option::is_none")]
832 pub title: Option<String>,
833 #[serde(skip_serializing_if = "Option::is_none")]
835 pub project_id: Option<String>,
836 pub created_at: DateTime<Utc>,
838 pub updated_at: DateTime<Utc>,
840}
841
842#[derive(Debug, Clone)]
844pub struct PruneConfig {
845 pub max_sessions: usize,
847 pub ttl_hours: u64,
849}
850
851impl Default for PruneConfig {
852 fn default() -> Self {
853 Self {
854 max_sessions: 100,
855 ttl_hours: 168, }
857 }
858}
859
860pub struct PruneThrottle {
862 last_prune: std::sync::Mutex<Option<std::time::Instant>>,
864 cooldown_secs: u64,
866}
867
868impl PruneThrottle {
869 pub fn new(cooldown_secs: u64) -> Self {
871 Self {
872 last_prune: std::sync::Mutex::new(None),
873 cooldown_secs,
874 }
875 }
876
877 pub fn should_prune(&self) -> bool {
880 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 let mut session = Session::new("user-123");
914 session.add_user_message("Hello");
915
916 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 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 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 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 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 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 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 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 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 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 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 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 let temp_dir = tempfile::tempdir().unwrap();
1052 let store = StateStore::new(temp_dir.path().to_path_buf()).unwrap();
1053
1054 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 let recent_session = Session::new("recent-user");
1061 store.save_session(&recent_session).await.unwrap();
1062
1063 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 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}