chasm_cli/
storage.rs

1// Copyright (c) 2024-2026 Nervosys LLC
2// SPDX-License-Identifier: Apache-2.0
3//! VS Code storage (SQLite database) operations
4
5use crate::error::{CsmError, Result};
6use crate::models::{ChatSession, ChatSessionIndex, ChatSessionIndexEntry};
7use crate::workspace::{get_empty_window_sessions_path, get_workspace_storage_path};
8use regex::Regex;
9use rusqlite::Connection;
10use std::path::{Path, PathBuf};
11use sysinfo::System;
12
13/// Sanitize JSON content by replacing lone surrogates with replacement character.
14/// VS Code sometimes writes invalid JSON with lone Unicode surrogates (e.g., \udde0).
15fn sanitize_json_unicode(content: &str) -> String {
16    // Match lone high surrogates (D800-DBFF) not followed by low surrogate (DC00-DFFF)
17    // and lone low surrogates (DC00-DFFF) not preceded by high surrogate
18    let re = Regex::new(r"\\u[dD][89aAbB][0-9a-fA-F]{2}(?!\\u[dD][cCdDeEfF][0-9a-fA-F]{2})|(?<!\\u[dD][89aAbB][0-9a-fA-F]{2})\\u[dD][cCdDeEfF][0-9a-fA-F]{2}")
19        .unwrap();
20    re.replace_all(content, "\\uFFFD").to_string()
21}
22
23/// Try to parse JSON, sanitizing invalid Unicode if needed
24pub fn parse_session_json(content: &str) -> std::result::Result<ChatSession, serde_json::Error> {
25    match serde_json::from_str::<ChatSession>(content) {
26        Ok(session) => Ok(session),
27        Err(e) => {
28            // If parsing fails due to Unicode issue, try sanitizing
29            if e.to_string().contains("surrogate") || e.to_string().contains("escape") {
30                let sanitized = sanitize_json_unicode(content);
31                serde_json::from_str::<ChatSession>(&sanitized)
32            } else {
33                Err(e)
34            }
35        }
36    }
37}
38
39/// Get the path to the workspace storage database
40pub fn get_workspace_storage_db(workspace_id: &str) -> Result<PathBuf> {
41    let storage_path = get_workspace_storage_path()?;
42    Ok(storage_path.join(workspace_id).join("state.vscdb"))
43}
44
45/// Read the chat session index from VS Code storage
46pub fn read_chat_session_index(db_path: &Path) -> Result<ChatSessionIndex> {
47    let conn = Connection::open(db_path)?;
48
49    let result: std::result::Result<String, rusqlite::Error> = conn.query_row(
50        "SELECT value FROM ItemTable WHERE key = ?",
51        ["chat.ChatSessionStore.index"],
52        |row| row.get(0),
53    );
54
55    match result {
56        Ok(json_str) => serde_json::from_str(&json_str)
57            .map_err(|e| CsmError::InvalidSessionFormat(e.to_string())),
58        Err(rusqlite::Error::QueryReturnedNoRows) => Ok(ChatSessionIndex::default()),
59        Err(e) => Err(CsmError::SqliteError(e)),
60    }
61}
62
63/// Write the chat session index to VS Code storage
64pub fn write_chat_session_index(db_path: &Path, index: &ChatSessionIndex) -> Result<()> {
65    let conn = Connection::open(db_path)?;
66    let json_str = serde_json::to_string(index)?;
67
68    // Check if the key exists
69    let exists: bool = conn.query_row(
70        "SELECT COUNT(*) > 0 FROM ItemTable WHERE key = ?",
71        ["chat.ChatSessionStore.index"],
72        |row| row.get(0),
73    )?;
74
75    if exists {
76        conn.execute(
77            "UPDATE ItemTable SET value = ? WHERE key = ?",
78            [&json_str, "chat.ChatSessionStore.index"],
79        )?;
80    } else {
81        conn.execute(
82            "INSERT INTO ItemTable (key, value) VALUES (?, ?)",
83            ["chat.ChatSessionStore.index", &json_str],
84        )?;
85    }
86
87    Ok(())
88}
89
90/// Add a session to the VS Code index
91pub fn add_session_to_index(
92    db_path: &Path,
93    session_id: &str,
94    title: &str,
95    last_message_date_ms: i64,
96    is_imported: bool,
97    initial_location: &str,
98    is_empty: bool,
99) -> Result<()> {
100    let mut index = read_chat_session_index(db_path)?;
101
102    index.entries.insert(
103        session_id.to_string(),
104        ChatSessionIndexEntry {
105            session_id: session_id.to_string(),
106            title: title.to_string(),
107            last_message_date: last_message_date_ms,
108            is_imported,
109            initial_location: initial_location.to_string(),
110            is_empty,
111        },
112    );
113
114    write_chat_session_index(db_path, &index)
115}
116
117/// Remove a session from the VS Code index
118pub fn remove_session_from_index(db_path: &Path, session_id: &str) -> Result<bool> {
119    let mut index = read_chat_session_index(db_path)?;
120    let removed = index.entries.remove(session_id).is_some();
121    if removed {
122        write_chat_session_index(db_path, &index)?;
123    }
124    Ok(removed)
125}
126
127/// Sync the VS Code index with sessions on disk (remove stale entries, add missing ones)
128pub fn sync_session_index(
129    workspace_id: &str,
130    chat_sessions_dir: &Path,
131    force: bool,
132) -> Result<(usize, usize)> {
133    let db_path = get_workspace_storage_db(workspace_id)?;
134
135    if !db_path.exists() {
136        return Err(CsmError::WorkspaceNotFound(format!(
137            "Database not found: {}",
138            db_path.display()
139        )));
140    }
141
142    // Check if VS Code is running
143    if !force && is_vscode_running() {
144        return Err(CsmError::VSCodeRunning);
145    }
146
147    // Get current index
148    let mut index = read_chat_session_index(&db_path)?;
149
150    // Get session files on disk
151    let mut files_on_disk: std::collections::HashSet<String> = std::collections::HashSet::new();
152    if chat_sessions_dir.exists() {
153        for entry in std::fs::read_dir(chat_sessions_dir)? {
154            let entry = entry?;
155            let path = entry.path();
156            if path.extension().map(|e| e == "json").unwrap_or(false) {
157                if let Some(stem) = path.file_stem() {
158                    files_on_disk.insert(stem.to_string_lossy().to_string());
159                }
160            }
161        }
162    }
163
164    // Remove stale entries (in index but not on disk)
165    let stale_ids: Vec<String> = index
166        .entries
167        .keys()
168        .filter(|id| !files_on_disk.contains(*id))
169        .cloned()
170        .collect();
171
172    let removed = stale_ids.len();
173    for id in &stale_ids {
174        index.entries.remove(id);
175    }
176
177    // Add/update sessions from disk
178    let mut added = 0;
179    for entry in std::fs::read_dir(chat_sessions_dir)? {
180        let entry = entry?;
181        let path = entry.path();
182
183        if path.extension().map(|e| e == "json").unwrap_or(false) {
184            if let Ok(content) = std::fs::read_to_string(&path) {
185                if let Ok(session) = parse_session_json(&content) {
186                    let session_id = session.session_id.clone().unwrap_or_else(|| {
187                        path.file_stem()
188                            .map(|s| s.to_string_lossy().to_string())
189                            .unwrap_or_else(|| uuid::Uuid::new_v4().to_string())
190                    });
191
192                    let title = session.title();
193                    let is_empty = session.is_empty();
194                    let last_message_date = session.last_message_date;
195                    let initial_location = session.initial_location.clone();
196
197                    index.entries.insert(
198                        session_id.clone(),
199                        ChatSessionIndexEntry {
200                            session_id,
201                            title,
202                            last_message_date,
203                            is_imported: session.is_imported,
204                            initial_location,
205                            is_empty,
206                        },
207                    );
208                    added += 1;
209                }
210            }
211        }
212    }
213
214    // Write the synced index
215    write_chat_session_index(&db_path, &index)?;
216
217    Ok((added, removed))
218}
219
220/// Register all sessions from a directory into the VS Code index
221pub fn register_all_sessions_from_directory(
222    workspace_id: &str,
223    chat_sessions_dir: &Path,
224    force: bool,
225) -> Result<usize> {
226    let db_path = get_workspace_storage_db(workspace_id)?;
227
228    if !db_path.exists() {
229        return Err(CsmError::WorkspaceNotFound(format!(
230            "Database not found: {}",
231            db_path.display()
232        )));
233    }
234
235    // Check if VS Code is running
236    if !force && is_vscode_running() {
237        return Err(CsmError::VSCodeRunning);
238    }
239
240    // Use sync to ensure index matches disk
241    let (added, removed) = sync_session_index(workspace_id, chat_sessions_dir, force)?;
242
243    // Print individual session info
244    for entry in std::fs::read_dir(chat_sessions_dir)? {
245        let entry = entry?;
246        let path = entry.path();
247
248        if path.extension().map(|e| e == "json").unwrap_or(false) {
249            if let Ok(content) = std::fs::read_to_string(&path) {
250                if let Ok(session) = parse_session_json(&content) {
251                    let session_id = session.session_id.clone().unwrap_or_else(|| {
252                        path.file_stem()
253                            .map(|s| s.to_string_lossy().to_string())
254                            .unwrap_or_else(|| uuid::Uuid::new_v4().to_string())
255                    });
256
257                    let title = session.title();
258
259                    println!(
260                        "[OK] Registered: {} ({}...)",
261                        title,
262                        &session_id[..12.min(session_id.len())]
263                    );
264                }
265            }
266        }
267    }
268
269    if removed > 0 {
270        println!("[OK] Removed {} stale index entries", removed);
271    }
272
273    Ok(added)
274}
275
276/// Check if VS Code is currently running
277pub fn is_vscode_running() -> bool {
278    let mut sys = System::new();
279    sys.refresh_processes();
280
281    for process in sys.processes().values() {
282        let name = process.name().to_lowercase();
283        if name.contains("code") && !name.contains("codec") {
284            return true;
285        }
286    }
287
288    false
289}
290
291/// Backup workspace sessions to a timestamped directory
292pub fn backup_workspace_sessions(workspace_dir: &Path) -> Result<Option<PathBuf>> {
293    let chat_sessions_dir = workspace_dir.join("chatSessions");
294
295    if !chat_sessions_dir.exists() {
296        return Ok(None);
297    }
298
299    let timestamp = std::time::SystemTime::now()
300        .duration_since(std::time::UNIX_EPOCH)
301        .unwrap()
302        .as_secs();
303
304    let backup_dir = workspace_dir.join(format!("chatSessions-backup-{}", timestamp));
305
306    // Copy directory recursively
307    copy_dir_all(&chat_sessions_dir, &backup_dir)?;
308
309    Ok(Some(backup_dir))
310}
311
312/// Recursively copy a directory
313fn copy_dir_all(src: &Path, dst: &Path) -> Result<()> {
314    std::fs::create_dir_all(dst)?;
315
316    for entry in std::fs::read_dir(src)? {
317        let entry = entry?;
318        let src_path = entry.path();
319        let dst_path = dst.join(entry.file_name());
320
321        if src_path.is_dir() {
322            copy_dir_all(&src_path, &dst_path)?;
323        } else {
324            std::fs::copy(&src_path, &dst_path)?;
325        }
326    }
327
328    Ok(())
329}
330
331// =============================================================================
332// Empty Window Sessions (ALL SESSIONS)
333// =============================================================================
334
335/// Read all empty window chat sessions (not tied to any workspace)
336/// These appear in VS Code's "ALL SESSIONS" panel
337pub fn read_empty_window_sessions() -> Result<Vec<ChatSession>> {
338    let sessions_path = get_empty_window_sessions_path()?;
339
340    if !sessions_path.exists() {
341        return Ok(Vec::new());
342    }
343
344    let mut sessions = Vec::new();
345
346    for entry in std::fs::read_dir(&sessions_path)? {
347        let entry = entry?;
348        let path = entry.path();
349
350        if path.extension().is_some_and(|e| e == "json") {
351            if let Ok(content) = std::fs::read_to_string(&path) {
352                if let Ok(session) = parse_session_json(&content) {
353                    sessions.push(session);
354                }
355            }
356        }
357    }
358
359    // Sort by last message date (most recent first)
360    sessions.sort_by(|a, b| b.last_message_date.cmp(&a.last_message_date));
361
362    Ok(sessions)
363}
364
365/// Get a specific empty window session by ID
366#[allow(dead_code)]
367pub fn get_empty_window_session(session_id: &str) -> Result<Option<ChatSession>> {
368    let sessions_path = get_empty_window_sessions_path()?;
369    let session_path = sessions_path.join(format!("{}.json", session_id));
370
371    if !session_path.exists() {
372        return Ok(None);
373    }
374
375    let content = std::fs::read_to_string(&session_path)?;
376    let session: ChatSession = serde_json::from_str(&content)
377        .map_err(|e| CsmError::InvalidSessionFormat(e.to_string()))?;
378
379    Ok(Some(session))
380}
381
382/// Write an empty window session
383#[allow(dead_code)]
384pub fn write_empty_window_session(session: &ChatSession) -> Result<PathBuf> {
385    let sessions_path = get_empty_window_sessions_path()?;
386
387    // Create directory if it doesn't exist
388    std::fs::create_dir_all(&sessions_path)?;
389
390    let session_id = session.session_id.as_deref().unwrap_or("unknown");
391    let session_path = sessions_path.join(format!("{}.json", session_id));
392    let content = serde_json::to_string_pretty(session)?;
393    std::fs::write(&session_path, content)?;
394
395    Ok(session_path)
396}
397
398/// Delete an empty window session
399#[allow(dead_code)]
400pub fn delete_empty_window_session(session_id: &str) -> Result<bool> {
401    let sessions_path = get_empty_window_sessions_path()?;
402    let session_path = sessions_path.join(format!("{}.json", session_id));
403
404    if session_path.exists() {
405        std::fs::remove_file(&session_path)?;
406        Ok(true)
407    } else {
408        Ok(false)
409    }
410}
411
412/// Count empty window sessions
413pub fn count_empty_window_sessions() -> Result<usize> {
414    let sessions_path = get_empty_window_sessions_path()?;
415
416    if !sessions_path.exists() {
417        return Ok(0);
418    }
419
420    let count = std::fs::read_dir(&sessions_path)?
421        .filter_map(|e| e.ok())
422        .filter(|e| e.path().extension().is_some_and(|ext| ext == "json"))
423        .count();
424
425    Ok(count)
426}