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::{SessionMetadata, SessionIndex, MessageSummary};
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 = serde_json::to_string_pretty(&self.index).context("serializing session index")?;
93        let tmp = path.with_extension("json.tmp");
94        std::fs::write(&tmp, json)
95            .with_context(|| format!("writing index tmp file {}", tmp.display()))?;
96        std::fs::rename(&tmp, &path)
97            .with_context(|| format!("renaming index tmp file to {}", path.display()))?;
98        Ok(())
99    }
100
101    pub fn save_index(&mut self) -> Result<()> {
102        self.lock.acquire(5000)?;
103        let result = self.save_index_locked();
104        self.lock.release()?;
105        result
106    }
107
108    pub fn start_new(&mut self, project_path: Option<&Path>) -> Result<&Session> {
109        let session = Session::new(project_path);
110        self.current_session = Some(session);
111        self.save_current()?;
112        self.current_session.as_ref()
113            .ok_or_else(|| anyhow::anyhow!("session not found after creation"))
114    }
115
116    pub fn continue_last(&mut self) -> Result<Option<&Session>> {
117        let last_id = self.index.last_session().map(|m| m.id.clone());
118        if let Some(id) = last_id {
119            self.load_session(&id)?;
120            Ok(self.current_session.as_ref())
121        } else {
122            Ok(None)
123        }
124    }
125
126    pub fn resume(&mut self, query: &str) -> Result<Option<&Session>> {
127        let session_id = self.index.find(query).map(|m| m.id.clone());
128        if let Some(id) = session_id {
129            self.load_session(&id)?;
130            Ok(self.current_session.as_ref())
131        } else {
132            Ok(None)
133        }
134    }
135
136    fn load_session(&mut self, id: &str) -> Result<()> {
137        let path = self.session_path(id);
138        if !path.exists() {
139            anyhow::bail!("session file {} not found", path.display());
140        }
141        let data = std::fs::read_to_string(&path)
142            .with_context(|| format!("reading session file {}", path.display()))?;
143        let mut session: Session = serde_json::from_str(&data)
144            .with_context(|| format!("parsing session file {}", path.display()))?;
145
146        session.migrate_legacy();
147
148        if session.metadata.name.is_none()
149            && let Some(index_meta) = self.index.find(id)
150        {
151            session.metadata.name = index_meta.name.clone();
152        }
153
154        self.current_session = Some(session);
155        Ok(())
156    }
157
158    pub fn save_current(&mut self) -> Result<()> {
159        if let Some(ref session) = self.current_session {
160            let session_clone = session.clone();
161
162            self.lock.acquire(5000)?;
163
164            self.index.upsert(session_clone.metadata.clone());
165            self.save_index_locked()?;
166
167            let path = self.session_path(&session_clone.metadata.id);
168            let json = serde_json::to_string(&session_clone).context("serializing session")?;
169            let tmp = path.with_extension("json.tmp");
170            std::fs::write(&tmp, json)
171                .with_context(|| format!("writing session tmp file {}", tmp.display()))?;
172            std::fs::rename(&tmp, &path)
173                .with_context(|| format!("renaming session tmp file to {}", path.display()))?;
174
175            self.lock.release()?;
176        }
177        Ok(())
178    }
179
180    pub fn update_stats(&mut self, last_input_tokens: u32, total_output_tokens: u64) {
181        if let Some(ref mut session) = self.current_session {
182            session.update_stats(last_input_tokens, total_output_tokens);
183        }
184    }
185
186    pub fn record_compression(&mut self, entry: CompressionHistoryEntry) {
187        if let Some(ref mut session) = self.current_session {
188            session.metadata.add_compression_entry(entry);
189        }
190    }
191
192    pub fn set_messages(&mut self, messages: Vec<Message>) {
193        if let Some(ref mut session) = self.current_session {
194            if session.metadata.name.is_none()
195                && !messages.is_empty()
196                && let Some(name) = Self::generate_name_from_messages(&messages)
197            {
198                session.metadata.name = Some(name);
199            }
200
201            session.full_messages = messages.clone();
202            session.message_summaries = messages
203                .iter()
204                .enumerate()
205                .map(|(i, m)| MessageSummary::from_message(m, i))
206                .collect();
207            session.metadata.message_count = session.full_messages.len();
208            session.metadata.updated_at = Utc::now();
209        }
210    }
211
212    pub fn set_compressed_messages(&mut self, compressed: Vec<Message>) {
213        if let Some(ref mut session) = self.current_session {
214            for summary in &mut session.message_summaries {
215                summary.is_compressed = true;
216            }
217
218            for compressed_msg in &compressed {
219                for (idx, full_msg) in session.full_messages.iter().enumerate() {
220                    if session.message_summaries.get(idx).is_some() {
221                        let same_role = compressed_msg.role == full_msg.role;
222                        if same_role
223                            && let Some(summary) = session.message_summaries.get_mut(idx) {
224                                summary.is_compressed = false;
225                            }
226                    }
227                }
228            }
229
230            session.compressed_messages = compressed;
231        }
232    }
233
234    fn generate_name_from_messages(messages: &[Message]) -> Option<String> {
235        let user_messages: Vec<&Message> = messages.iter().filter(|m| m.role == Role::User).collect();
236
237        for msg in user_messages.iter().take(3) {
238            let text = match &msg.content {
239                MessageContent::Text(t) => t.clone(),
240                MessageContent::Blocks(blocks) => blocks
241                    .iter()
242                    .filter_map(|b| {
243                        if let ContentBlock::Text { text } = b {
244                            Some(text.clone())
245                        } else {
246                            None
247                        }
248                    })
249                    .collect::<Vec<_>>()
250                    .join(" "),
251            };
252
253            let cleaned = text.trim().lines().next().unwrap_or("").trim();
254
255            if cleaned.len() < 5 || is_generic_message(cleaned) {
256                continue;
257            }
258
259            let name = if cleaned.chars().count() > 40 {
260                let truncated: String = cleaned.chars().take(37).collect();
261                format!("{}...", truncated)
262            } else {
263                cleaned.to_string()
264            };
265
266            return Some(name);
267        }
268
269        None
270    }
271
272    pub fn api_messages(&self) -> Option<&[Message]> {
273        self.current_session.as_ref().map(|s| s.api_messages())
274    }
275
276    pub fn display_messages(&self) -> Option<&[Message]> {
277        self.current_session.as_ref().map(|s| s.display_messages())
278    }
279
280    pub fn messages(&self) -> Option<&[Message]> {
281        self.current_session.as_ref().map(|s| s.api_messages())
282    }
283
284    pub fn messages_mut(&mut self) -> Option<&mut Vec<Message>> {
285        self.current_session.as_mut().map(|s| &mut s.full_messages)
286    }
287
288    pub fn full_messages(&self) -> Option<&[Message]> {
289        self.current_session.as_ref().map(|s| s.display_messages())
290    }
291
292    pub fn current_id(&self) -> Option<&str> {
293        self.current_session.as_ref().map(|s| s.metadata.id.as_str())
294    }
295
296    pub fn current_name(&self) -> Option<&str> {
297        self.current_session.as_ref().and_then(|s| s.name())
298    }
299
300    pub fn rename_current(&mut self, new_name: &str) -> Result<()> {
301        if let Some(ref session) = self.current_session {
302            let id = session.metadata.id.clone();
303            self.index.rename(&id, new_name)?;
304            if let Some(ref mut session) = self.current_session {
305                session.metadata.name = Some(new_name.to_string());
306            }
307            self.save_current()?;
308        }
309        Ok(())
310    }
311
312    pub fn clear_current(&mut self) -> Result<()> {
313        if let Some(ref session) = self.current_session {
314            self.lock.acquire(5000)?;
315
316            let path = self.session_path(&session.metadata.id);
317            let _ = std::fs::remove_file(&path);
318            self.index.remove(&session.metadata.id);
319            self.save_index_locked()?;
320
321            self.lock.release()?;
322        }
323        self.current_session = None;
324        Ok(())
325    }
326
327    pub fn list_sessions(&self) -> &[SessionMetadata] {
328        &self.index.sessions
329    }
330
331    pub fn cleanup_old_sessions(&mut self, max_age_days: u64) -> Result<usize> {
332        let now = chrono::Utc::now();
333        let threshold = chrono::Duration::days(max_age_days as i64);
334
335        let mut to_remove: Vec<String> = Vec::new();
336
337        for session in &self.index.sessions {
338            let age = now - session.updated_at;
339            if age > threshold {
340                to_remove.push(session.id.clone());
341            }
342        }
343
344        let removed_count = to_remove.len();
345
346        if removed_count > 0 {
347            self.lock.acquire(5000)?;
348
349            for id in &to_remove {
350                let path = self.session_path(id);
351                let _ = std::fs::remove_file(&path);
352                self.index.remove(id);
353            }
354
355            self.save_index_locked()?;
356            self.lock.release()?;
357        }
358
359        Ok(removed_count)
360    }
361
362    pub fn prune_sessions(&mut self, max_sessions: usize) -> Result<usize> {
363        if self.index.sessions.len() <= max_sessions {
364            return Ok(0);
365        }
366
367        let to_remove = self.index.sessions.len() - max_sessions;
368        let mut ids_to_remove: Vec<String> = Vec::new();
369
370        for session in self.index.sessions.iter().skip(max_sessions) {
371            ids_to_remove.push(session.id.clone());
372        }
373
374        self.lock.acquire(5000)?;
375
376        for id in &ids_to_remove {
377            let path = self.session_path(id);
378            let _ = std::fs::remove_file(&path);
379            self.index.remove(id);
380        }
381
382        self.save_index_locked()?;
383        self.lock.release()?;
384
385        Ok(to_remove)
386    }
387
388    pub fn session_count(&self) -> usize {
389        self.index.sessions.len()
390    }
391
392    pub fn has_current(&self) -> bool {
393        self.current_session.is_some()
394    }
395
396    pub fn current_metadata(&self) -> Option<&SessionMetadata> {
397        self.current_session.as_ref().map(|s| &s.metadata)
398    }
399
400    pub fn history_path(&self) -> PathBuf {
401        self.base_dir.join("history.txt")
402    }
403}