Skip to main content

aster/session/
store.rs

1//! Session 存储抽象层
2//!
3//! 定义 `SessionStore` trait,允许应用层注入自定义存储实现。
4//! 框架层不再直接依赖具体的存储实现(如 SQLite)。
5
6use crate::conversation::message::Message;
7use crate::conversation::Conversation;
8use crate::model::ModelConfig;
9use crate::recipe::Recipe;
10use crate::session::extension_data::ExtensionData;
11use crate::session::session_manager::{Session, SessionInsights, SessionType};
12use anyhow::Result;
13use async_trait::async_trait;
14use std::collections::HashMap;
15use std::path::PathBuf;
16use std::sync::Arc;
17
18/// Session 存储 trait
19///
20/// 应用层可以实现此 trait 来提供自定义的 session 存储。
21/// 框架提供默认的 SQLite 实现 (`SqliteSessionStore`)。
22#[async_trait]
23pub trait SessionStore: Send + Sync {
24    /// 创建新 session
25    async fn create_session(
26        &self,
27        working_dir: PathBuf,
28        name: String,
29        session_type: SessionType,
30    ) -> Result<Session>;
31
32    /// 获取 session
33    async fn get_session(&self, id: &str, include_messages: bool) -> Result<Session>;
34
35    /// 添加消息到 session
36    async fn add_message(&self, session_id: &str, message: &Message) -> Result<()>;
37
38    /// 替换整个对话历史
39    async fn replace_conversation(
40        &self,
41        session_id: &str,
42        conversation: &Conversation,
43    ) -> Result<()>;
44
45    /// 列出所有 session
46    async fn list_sessions(&self) -> Result<Vec<Session>>;
47
48    /// 按类型列出 session
49    async fn list_sessions_by_types(&self, types: &[SessionType]) -> Result<Vec<Session>>;
50
51    /// 删除 session
52    async fn delete_session(&self, id: &str) -> Result<()>;
53
54    /// 获取统计信息
55    async fn get_insights(&self) -> Result<SessionInsights>;
56
57    /// 导出 session 为 JSON
58    async fn export_session(&self, id: &str) -> Result<String>;
59
60    /// 从 JSON 导入 session
61    async fn import_session(&self, json: &str) -> Result<Session>;
62
63    /// 复制 session
64    async fn copy_session(&self, session_id: &str, new_name: String) -> Result<Session>;
65
66    /// 截断对话(删除指定时间戳之后的消息)
67    async fn truncate_conversation(&self, session_id: &str, timestamp: i64) -> Result<()>;
68
69    /// 更新 session 名称
70    async fn update_session_name(
71        &self,
72        session_id: &str,
73        name: String,
74        user_set: bool,
75    ) -> Result<()>;
76
77    /// 更新 session 扩展数据
78    async fn update_extension_data(
79        &self,
80        session_id: &str,
81        extension_data: ExtensionData,
82    ) -> Result<()>;
83
84    /// 更新 session token 统计
85    async fn update_token_stats(&self, session_id: &str, stats: TokenStatsUpdate) -> Result<()>;
86
87    /// 更新 session 的 provider 和 model 配置
88    async fn update_provider_config(
89        &self,
90        session_id: &str,
91        provider_name: Option<String>,
92        model_config: Option<ModelConfig>,
93    ) -> Result<()>;
94
95    /// 更新 session 的 recipe 配置
96    async fn update_recipe(
97        &self,
98        session_id: &str,
99        recipe: Option<Recipe>,
100        user_recipe_values: Option<HashMap<String, String>>,
101    ) -> Result<()>;
102
103    /// 搜索聊天历史
104    async fn search_chat_history(
105        &self,
106        query: &str,
107        limit: Option<usize>,
108        after_date: Option<chrono::DateTime<chrono::Utc>>,
109        before_date: Option<chrono::DateTime<chrono::Utc>>,
110        exclude_session_id: Option<String>,
111    ) -> Result<Vec<ChatHistoryMatch>>;
112}
113
114/// 聊天历史搜索结果
115#[derive(Debug, Clone)]
116pub struct ChatHistoryMatch {
117    pub session_id: String,
118    pub session_name: String,
119    pub message_role: String,
120    pub message_content: String,
121    pub timestamp: chrono::DateTime<chrono::Utc>,
122    pub relevance_score: f32,
123}
124
125/// Token 统计更新参数
126#[derive(Debug, Clone, Default)]
127pub struct TokenStatsUpdate {
128    pub schedule_id: Option<String>,
129    pub total_tokens: Option<i32>,
130    pub input_tokens: Option<i32>,
131    pub output_tokens: Option<i32>,
132    pub accumulated_total: Option<i32>,
133    pub accumulated_input: Option<i32>,
134    pub accumulated_output: Option<i32>,
135}
136
137/// 空存储实现(不保存任何数据)
138///
139/// 用于不需要持久化的场景,如测试或无状态 API 服务。
140pub struct NoopSessionStore;
141
142#[async_trait]
143impl SessionStore for NoopSessionStore {
144    async fn create_session(
145        &self,
146        working_dir: PathBuf,
147        name: String,
148        session_type: SessionType,
149    ) -> Result<Session> {
150        Ok(Session {
151            id: uuid::Uuid::new_v4().to_string(),
152            working_dir,
153            name,
154            user_set_name: false,
155            session_type,
156            created_at: chrono::Utc::now(),
157            updated_at: chrono::Utc::now(),
158            extension_data: ExtensionData::default(),
159            total_tokens: None,
160            input_tokens: None,
161            output_tokens: None,
162            accumulated_total_tokens: None,
163            accumulated_input_tokens: None,
164            accumulated_output_tokens: None,
165            schedule_id: None,
166            recipe: None,
167            user_recipe_values: None,
168            conversation: Some(Conversation::default()),
169            message_count: 0,
170            provider_name: None,
171            model_config: None,
172        })
173    }
174
175    async fn get_session(&self, _id: &str, _include_messages: bool) -> Result<Session> {
176        Err(anyhow::anyhow!("NoopSessionStore: session not found"))
177    }
178
179    async fn add_message(&self, _session_id: &str, _message: &Message) -> Result<()> {
180        Ok(()) // 静默忽略
181    }
182
183    async fn replace_conversation(
184        &self,
185        _session_id: &str,
186        _conversation: &Conversation,
187    ) -> Result<()> {
188        Ok(())
189    }
190
191    async fn list_sessions(&self) -> Result<Vec<Session>> {
192        Ok(vec![])
193    }
194
195    async fn list_sessions_by_types(&self, _types: &[SessionType]) -> Result<Vec<Session>> {
196        Ok(vec![])
197    }
198
199    async fn delete_session(&self, _id: &str) -> Result<()> {
200        Ok(())
201    }
202
203    async fn get_insights(&self) -> Result<SessionInsights> {
204        Ok(SessionInsights {
205            total_sessions: 0,
206            total_tokens: 0,
207        })
208    }
209
210    async fn export_session(&self, _id: &str) -> Result<String> {
211        Err(anyhow::anyhow!("NoopSessionStore: export not supported"))
212    }
213
214    async fn import_session(&self, _json: &str) -> Result<Session> {
215        Err(anyhow::anyhow!("NoopSessionStore: import not supported"))
216    }
217
218    async fn copy_session(&self, _session_id: &str, _new_name: String) -> Result<Session> {
219        Err(anyhow::anyhow!("NoopSessionStore: copy not supported"))
220    }
221
222    async fn truncate_conversation(&self, _session_id: &str, _timestamp: i64) -> Result<()> {
223        Ok(())
224    }
225
226    async fn update_session_name(
227        &self,
228        _session_id: &str,
229        _name: String,
230        _user_set: bool,
231    ) -> Result<()> {
232        Ok(())
233    }
234
235    async fn update_extension_data(
236        &self,
237        _session_id: &str,
238        _extension_data: ExtensionData,
239    ) -> Result<()> {
240        Ok(())
241    }
242
243    async fn update_token_stats(&self, _session_id: &str, _stats: TokenStatsUpdate) -> Result<()> {
244        Ok(())
245    }
246
247    async fn update_provider_config(
248        &self,
249        _session_id: &str,
250        _provider_name: Option<String>,
251        _model_config: Option<ModelConfig>,
252    ) -> Result<()> {
253        Ok(())
254    }
255
256    async fn update_recipe(
257        &self,
258        _session_id: &str,
259        _recipe: Option<Recipe>,
260        _user_recipe_values: Option<HashMap<String, String>>,
261    ) -> Result<()> {
262        Ok(())
263    }
264
265    async fn search_chat_history(
266        &self,
267        _query: &str,
268        _limit: Option<usize>,
269        _after_date: Option<chrono::DateTime<chrono::Utc>>,
270        _before_date: Option<chrono::DateTime<chrono::Utc>>,
271        _exclude_session_id: Option<String>,
272    ) -> Result<Vec<ChatHistoryMatch>> {
273        Ok(vec![])
274    }
275}
276
277/// 全局 session store 实例
278///
279/// 用于向后兼容,允许现有代码继续使用 `SessionManager::` 静态方法。
280/// 新代码应该使用 `Agent::with_session_store()` 注入存储。
281static GLOBAL_SESSION_STORE: tokio::sync::OnceCell<Arc<dyn SessionStore>> =
282    tokio::sync::OnceCell::const_new();
283
284/// 设置全局 session store
285///
286/// 必须在使用 `SessionManager` 静态方法之前调用。
287/// 通常在应用启动时调用一次。
288pub async fn set_global_session_store(store: Arc<dyn SessionStore>) -> Result<()> {
289    GLOBAL_SESSION_STORE
290        .set(store)
291        .map_err(|_| anyhow::anyhow!("Global session store already set"))
292}
293
294/// 获取全局 session store
295///
296/// 如果未设置,返回错误。
297pub fn get_global_session_store() -> Result<Arc<dyn SessionStore>> {
298    GLOBAL_SESSION_STORE
299        .get()
300        .cloned()
301        .ok_or_else(|| anyhow::anyhow!("Global session store not initialized"))
302}
303
304/// 检查全局 session store 是否已设置
305pub fn is_global_session_store_set() -> bool {
306    GLOBAL_SESSION_STORE.get().is_some()
307}