Skip to main content

matrixcode_core/session/
manager.rs

1//! Session manager for storage operations
2
3use anyhow::{Context, Result};
4use chrono::Utc;
5use std::path::{Path, PathBuf};
6
7use super::metadata::{MessageSummary, SessionIndex, SessionMetadata};
8use super::session::{Session, SessionFileLock};
9use crate::compress::CompressionHistoryEntry;
10use crate::providers::{ContentBlock, Message, MessageContent, Role};
11
12/// Check if a message is too generic to be a good session name.
13fn is_generic_message(msg: &str) -> bool {
14    let generic = [
15        "继续", "好的", "ok", "yes", "no", "是", "否", "嗯", "对", "行", "可以", "好", "谢谢",
16        "thanks", "hi", "hello", "你好", "开始", "start",
17    ];
18    generic.iter().any(|g| msg.eq_ignore_ascii_case(g))
19}
20
21/// Manager for session storage.
22pub struct SessionManager {
23    base_dir: PathBuf,
24    current_session: Option<Session>,
25    index: SessionIndex,
26    lock: SessionFileLock,
27}
28
29impl SessionManager {
30    /// Create a new session manager.
31    pub fn new() -> Result<Self> {
32        let base_dir = Self::get_base_dir()?;
33        let lock = SessionFileLock::new(&base_dir);
34        let manager = Self {
35            base_dir,
36            current_session: None,
37            index: SessionIndex::default(),
38            lock,
39        };
40        manager.ensure_dirs()?;
41        let mut manager = manager;
42        manager.load_index()?;
43        Ok(manager)
44    }
45
46    fn get_base_dir() -> Result<PathBuf> {
47        let home = std::env::var_os("HOME")
48            .or_else(|| std::env::var_os("USERPROFILE"))
49            .ok_or_else(|| anyhow::anyhow!("HOME or USERPROFILE environment variable not set"))?;
50        let mut p = PathBuf::from(home);
51        p.push(".matrix");
52        Ok(p)
53    }
54
55    fn sessions_dir(&self) -> PathBuf {
56        self.base_dir.join("sessions")
57    }
58
59    fn index_path(&self) -> PathBuf {
60        self.sessions_dir().join("index.json")
61    }
62
63    fn session_path(&self, id: &str) -> PathBuf {
64        self.sessions_dir().join(format!("{}.json", id))
65    }
66
67    fn ensure_dirs(&self) -> Result<()> {
68        std::fs::create_dir_all(&self.base_dir)
69            .with_context(|| format!("creating base dir {}", self.base_dir.display()))?;
70        std::fs::create_dir_all(self.sessions_dir())
71            .with_context(|| format!("creating sessions dir {}", self.sessions_dir().display()))?;
72        Ok(())
73    }
74
75    fn load_index(&mut self) -> Result<()> {
76        let path = self.index_path();
77        if !path.exists() {
78            return Ok(());
79        }
80        let data = std::fs::read_to_string(&path)
81            .with_context(|| format!("reading index file {}", path.display()))?;
82        if data.trim().is_empty() {
83            return Ok(());
84        }
85        self.index = serde_json::from_str(&data)
86            .with_context(|| format!("parsing index file {}", path.display()))?;
87        Ok(())
88    }
89
90    fn save_index_locked(&mut self) -> Result<()> {
91        let path = self.index_path();
92        let json =
93            serde_json::to_string_pretty(&self.index).context("serializing session index")?;
94        let tmp = path.with_extension("json.tmp");
95        std::fs::write(&tmp, json)
96            .with_context(|| format!("writing index tmp file {}", tmp.display()))?;
97        std::fs::rename(&tmp, &path)
98            .with_context(|| format!("renaming index tmp file to {}", path.display()))?;
99        Ok(())
100    }
101
102    pub fn save_index(&mut self) -> Result<()> {
103        self.lock.acquire(5000)?;
104        let result = self.save_index_locked();
105        self.lock.release()?;
106        result
107    }
108
109    pub fn start_new(&mut self, project_path: Option<&Path>) -> Result<&Session> {
110        let session = Session::new(project_path);
111        self.current_session = Some(session);
112        self.save_current()?;
113        self.current_session
114            .as_ref()
115            .ok_or_else(|| anyhow::anyhow!("session not found after creation"))
116    }
117
118    pub fn continue_last(&mut self) -> Result<Option<&Session>> {
119        let last_id = self.index.last_session().map(|m| m.id.clone());
120        if let Some(id) = last_id {
121            self.load_session(&id)?;
122            Ok(self.current_session.as_ref())
123        } else {
124            Ok(None)
125        }
126    }
127
128    pub fn resume(&mut self, query: &str) -> Result<Option<&Session>> {
129        let session_id = self.index.find(query).map(|m| m.id.clone());
130        if let Some(id) = session_id {
131            self.load_session(&id)?;
132            Ok(self.current_session.as_ref())
133        } else {
134            Ok(None)
135        }
136    }
137
138    fn load_session(&mut self, id: &str) -> Result<()> {
139        let path = self.session_path(id);
140        if !path.exists() {
141            anyhow::bail!("session file {} not found", path.display());
142        }
143        let data = std::fs::read_to_string(&path)
144            .with_context(|| format!("reading session file {}", path.display()))?;
145        let mut session: Session = serde_json::from_str(&data)
146            .with_context(|| format!("parsing session file {}", path.display()))?;
147
148        session.migrate_legacy();
149
150        if session.metadata.name.is_none()
151            && let Some(index_meta) = self.index.find(id)
152        {
153            session.metadata.name = index_meta.name.clone();
154        }
155
156        self.current_session = Some(session);
157        Ok(())
158    }
159
160    pub fn save_current(&mut self) -> Result<()> {
161        if let Some(ref session) = self.current_session {
162            let session_clone = session.clone();
163
164            self.lock.acquire(5000)?;
165
166            self.index.upsert(session_clone.metadata.clone());
167            self.save_index_locked()?;
168
169            let path = self.session_path(&session_clone.metadata.id);
170            let json = serde_json::to_string(&session_clone).context("serializing session")?;
171            let tmp = path.with_extension("json.tmp");
172            std::fs::write(&tmp, json)
173                .with_context(|| format!("writing session tmp file {}", tmp.display()))?;
174            std::fs::rename(&tmp, &path)
175                .with_context(|| format!("renaming session tmp file to {}", path.display()))?;
176
177            self.lock.release()?;
178        }
179        Ok(())
180    }
181
182    pub fn update_stats(&mut self, last_input_tokens: u32, total_output_tokens: u64) {
183        if let Some(ref mut session) = self.current_session {
184            session.update_stats(last_input_tokens, total_output_tokens);
185        }
186    }
187
188    pub fn record_compression(&mut self, entry: CompressionHistoryEntry) {
189        if let Some(ref mut session) = self.current_session {
190            session.metadata.add_compression_entry(entry);
191        }
192    }
193
194    pub fn set_messages(&mut self, messages: Vec<Message>) {
195        if let Some(ref mut session) = self.current_session {
196            if session.metadata.name.is_none()
197                && !messages.is_empty()
198                && let Some(name) = Self::generate_name_from_messages(&messages)
199            {
200                session.metadata.name = Some(name);
201            }
202
203            session.full_messages = messages.clone();
204            session.message_summaries = messages
205                .iter()
206                .enumerate()
207                .map(|(i, m)| MessageSummary::from_message(m, i))
208                .collect();
209            session.metadata.message_count = session.full_messages.len();
210            session.metadata.updated_at = Utc::now();
211        }
212    }
213
214    pub fn set_compressed_messages(&mut self, compressed: Vec<Message>) {
215        if let Some(ref mut session) = self.current_session {
216            for summary in &mut session.message_summaries {
217                summary.is_compressed = true;
218            }
219
220            for compressed_msg in &compressed {
221                for (idx, full_msg) in session.full_messages.iter().enumerate() {
222                    if session.message_summaries.get(idx).is_some() {
223                        let same_role = compressed_msg.role == full_msg.role;
224                        if same_role && let Some(summary) = session.message_summaries.get_mut(idx) {
225                            summary.is_compressed = false;
226                        }
227                    }
228                }
229            }
230
231            session.compressed_messages = compressed;
232        }
233    }
234
235    fn generate_name_from_messages(messages: &[Message]) -> Option<String> {
236        let user_messages: Vec<&Message> =
237            messages.iter().filter(|m| m.role == Role::User).collect();
238
239        for msg in user_messages.iter().take(3) {
240            let text = match &msg.content {
241                MessageContent::Text(t) => t.clone(),
242                MessageContent::Blocks(blocks) => blocks
243                    .iter()
244                    .filter_map(|b| {
245                        if let ContentBlock::Text { text } = b {
246                            Some(text.clone())
247                        } else {
248                            None
249                        }
250                    })
251                    .collect::<Vec<_>>()
252                    .join(" "),
253            };
254
255            let cleaned = text.trim().lines().next().unwrap_or("").trim();
256
257            if cleaned.len() < 5 || is_generic_message(cleaned) {
258                continue;
259            }
260
261            let name = if cleaned.chars().count() > 40 {
262                let truncated: String = cleaned.chars().take(37).collect();
263                format!("{}...", truncated)
264            } else {
265                cleaned.to_string()
266            };
267
268            return Some(name);
269        }
270
271        None
272    }
273
274    pub fn api_messages(&self) -> Option<&[Message]> {
275        self.current_session.as_ref().map(|s| s.api_messages())
276    }
277
278    pub fn display_messages(&self) -> Option<&[Message]> {
279        self.current_session.as_ref().map(|s| s.display_messages())
280    }
281
282    pub fn messages(&self) -> Option<&[Message]> {
283        self.current_session.as_ref().map(|s| s.api_messages())
284    }
285
286    pub fn messages_mut(&mut self) -> Option<&mut Vec<Message>> {
287        self.current_session.as_mut().map(|s| &mut s.full_messages)
288    }
289
290    pub fn full_messages(&self) -> Option<&[Message]> {
291        self.current_session.as_ref().map(|s| s.display_messages())
292    }
293
294    pub fn current_id(&self) -> Option<&str> {
295        self.current_session
296            .as_ref()
297            .map(|s| s.metadata.id.as_str())
298    }
299
300    pub fn current_name(&self) -> Option<&str> {
301        self.current_session.as_ref().and_then(|s| s.name())
302    }
303
304    pub fn rename_current(&mut self, new_name: &str) -> Result<()> {
305        if let Some(ref session) = self.current_session {
306            let id = session.metadata.id.clone();
307            self.index.rename(&id, new_name)?;
308            if let Some(ref mut session) = self.current_session {
309                session.metadata.name = Some(new_name.to_string());
310            }
311            self.save_current()?;
312        }
313        Ok(())
314    }
315
316    pub fn clear_current(&mut self) -> Result<()> {
317        if let Some(ref session) = self.current_session {
318            self.lock.acquire(5000)?;
319
320            let path = self.session_path(&session.metadata.id);
321            let _ = std::fs::remove_file(&path);
322            self.index.remove(&session.metadata.id);
323            self.save_index_locked()?;
324
325            self.lock.release()?;
326        }
327        self.current_session = None;
328        Ok(())
329    }
330
331    pub fn list_sessions(&self) -> &[SessionMetadata] {
332        &self.index.sessions
333    }
334
335    pub fn cleanup_old_sessions(&mut self, max_age_days: u64) -> Result<usize> {
336        let now = chrono::Utc::now();
337        let threshold = chrono::Duration::days(max_age_days as i64);
338
339        let mut to_remove: Vec<String> = Vec::new();
340
341        for session in &self.index.sessions {
342            let age = now - session.updated_at;
343            if age > threshold {
344                to_remove.push(session.id.clone());
345            }
346        }
347
348        let removed_count = to_remove.len();
349
350        if removed_count > 0 {
351            self.lock.acquire(5000)?;
352
353            for id in &to_remove {
354                let path = self.session_path(id);
355                let _ = std::fs::remove_file(&path);
356                self.index.remove(id);
357            }
358
359            self.save_index_locked()?;
360            self.lock.release()?;
361        }
362
363        Ok(removed_count)
364    }
365
366    pub fn prune_sessions(&mut self, max_sessions: usize) -> Result<usize> {
367        if self.index.sessions.len() <= max_sessions {
368            return Ok(0);
369        }
370
371        let to_remove = self.index.sessions.len() - max_sessions;
372        let mut ids_to_remove: Vec<String> = Vec::new();
373
374        for session in self.index.sessions.iter().skip(max_sessions) {
375            ids_to_remove.push(session.id.clone());
376        }
377
378        self.lock.acquire(5000)?;
379
380        for id in &ids_to_remove {
381            let path = self.session_path(id);
382            let _ = std::fs::remove_file(&path);
383            self.index.remove(id);
384        }
385
386        self.save_index_locked()?;
387        self.lock.release()?;
388
389        Ok(to_remove)
390    }
391
392    pub fn session_count(&self) -> usize {
393        self.index.sessions.len()
394    }
395
396    pub fn has_current(&self) -> bool {
397        self.current_session.is_some()
398    }
399
400    pub fn current_metadata(&self) -> Option<&SessionMetadata> {
401        self.current_session.as_ref().map(|s| &s.metadata)
402    }
403
404    pub fn history_path(&self) -> PathBuf {
405        self.base_dir.join("history.txt")
406    }
407}