agent_base/engine/
session_store.rs1use std::collections::HashMap;
2
3use async_trait::async_trait;
4use tokio::sync::Mutex;
5
6use crate::types::{AgentResult, AgentError, RuntimeEvent, SessionId};
7use super::AgentSession;
8
9#[async_trait]
20pub trait SessionStore: Send + Sync {
21 async fn save(&self, session: &AgentSession) -> AgentResult<()>;
23
24 async fn load(&self, session_id: &SessionId) -> AgentResult<Option<AgentSession>>;
26
27 async fn list(&self) -> AgentResult<Vec<SessionId>>;
29
30 async fn delete(&self, session_id: &SessionId) -> AgentResult<()>;
32
33 async fn append_event(&self, _session_id: &SessionId, _event: &RuntimeEvent) -> AgentResult<()> {
37 Ok(())
38 }
39}
40
41pub struct InMemorySessionStore {
42 sessions: Mutex<HashMap<SessionId, AgentSession>>,
43}
44
45impl InMemorySessionStore {
46 pub fn new() -> Self {
47 Self {
48 sessions: Mutex::new(HashMap::new()),
49 }
50 }
51}
52
53impl Default for InMemorySessionStore {
54 fn default() -> Self {
55 Self::new()
56 }
57}
58
59#[async_trait]
60impl SessionStore for InMemorySessionStore {
61 async fn save(&self, session: &AgentSession) -> AgentResult<()> {
62 let session_id = session
63 .id()
64 .ok_or_else(|| AgentError::internal("session has no id"))?;
65 self.sessions
66 .lock()
67 .await
68 .insert(session_id, session.clone());
69 Ok(())
70 }
71
72 async fn load(&self, session_id: &SessionId) -> AgentResult<Option<AgentSession>> {
73 Ok(self.sessions.lock().await.get(session_id).cloned())
74 }
75
76 async fn list(&self) -> AgentResult<Vec<SessionId>> {
77 Ok(self.sessions.lock().await.keys().cloned().collect())
78 }
79
80 async fn delete(&self, session_id: &SessionId) -> AgentResult<()> {
81 self.sessions.lock().await.remove(session_id);
82 Ok(())
83 }
84}