Skip to main content

excel_mcp_server/
store.rs

1use std::collections::HashMap;
2use std::time::{Duration, Instant};
3
4use uuid::Uuid;
5
6use crate::error::ExcelMcpError;
7
8/// A single workbook in the store — always backed by zavora-xlsx.
9pub struct WorkbookEntry {
10    pub id: String,
11    pub data: zavora_xlsx::Workbook,
12    pub read_only: bool,
13    pub last_access: Instant,
14}
15
16/// Manages the lifecycle of all open workbooks.
17pub struct WorkbookStore {
18    workbooks: HashMap<String, WorkbookEntry>,
19    max_capacity: usize,
20    ttl: Duration,
21}
22
23impl std::fmt::Debug for WorkbookStore {
24    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
25        f.debug_struct("WorkbookStore")
26            .field("open_count", &self.workbooks.len())
27            .field("max_capacity", &self.max_capacity)
28            .field("ttl", &self.ttl)
29            .finish()
30    }
31}
32
33impl WorkbookStore {
34    pub fn new() -> Self {
35        Self {
36            workbooks: HashMap::new(),
37            max_capacity: 10,
38            ttl: Duration::from_secs(30 * 60),
39        }
40    }
41
42    pub fn with_config(max_capacity: usize, ttl: Duration) -> Self {
43        Self {
44            workbooks: HashMap::new(),
45            max_capacity,
46            ttl,
47        }
48    }
49
50    pub fn insert(&mut self, mut entry: WorkbookEntry) -> Result<String, ExcelMcpError> {
51        self.evict_expired();
52        if self.workbooks.len() >= self.max_capacity {
53            return Err(ExcelMcpError::CapacityExceeded(format!(
54                "Workbook store is at maximum capacity ({}). Save and close an existing workbook first.",
55                self.max_capacity
56            )));
57        }
58        let id = Uuid::new_v4().to_string();
59        entry.id = id.clone();
60        entry.last_access = Instant::now();
61        self.workbooks.insert(id.clone(), entry);
62        Ok(id)
63    }
64
65    pub fn get(&mut self, id: &str) -> Option<&WorkbookEntry> {
66        self.evict_expired();
67        if let Some(entry) = self.workbooks.get_mut(id) {
68            entry.last_access = Instant::now();
69        }
70        self.workbooks.get(id)
71    }
72
73    pub fn get_mut(&mut self, id: &str) -> Option<&mut WorkbookEntry> {
74        self.evict_expired();
75        if let Some(entry) = self.workbooks.get_mut(id) {
76            entry.last_access = Instant::now();
77        }
78        self.workbooks.get_mut(id)
79    }
80
81    pub fn remove(&mut self, id: &str) -> Option<WorkbookEntry> {
82        self.workbooks.remove(id)
83    }
84
85    pub fn evict_expired(&mut self) -> Vec<String> {
86        let now = Instant::now();
87        let ttl = self.ttl;
88        let expired: Vec<String> = self
89            .workbooks
90            .iter()
91            .filter(|(_, e)| now.duration_since(e.last_access) > ttl)
92            .map(|(id, _)| id.clone())
93            .collect();
94        for id in &expired {
95            self.workbooks.remove(id);
96        }
97        expired
98    }
99
100    pub fn open_ids(&self) -> Vec<String> {
101        self.workbooks.keys().cloned().collect()
102    }
103    pub fn is_full(&self) -> bool {
104        self.workbooks.len() >= self.max_capacity
105    }
106}
107
108impl Default for WorkbookStore {
109    fn default() -> Self {
110        Self::new()
111    }
112}
113
114#[cfg(test)]
115mod tests {
116    use super::*;
117
118    fn make_entry() -> WorkbookEntry {
119        WorkbookEntry {
120            id: String::new(),
121            data: zavora_xlsx::Workbook::new(),
122            read_only: false,
123            last_access: Instant::now(),
124        }
125    }
126
127    #[test]
128    fn test_capacity_enforcement() {
129        let mut store = WorkbookStore::with_config(2, Duration::from_secs(600));
130        let _id1 = store.insert(make_entry()).unwrap();
131        let _id2 = store.insert(make_entry()).unwrap();
132        assert!(store.insert(make_entry()).is_err());
133    }
134
135    #[test]
136    fn test_ttl_eviction() {
137        let mut store = WorkbookStore::with_config(10, Duration::from_millis(1));
138        let id = store.insert(make_entry()).unwrap();
139        std::thread::sleep(Duration::from_millis(10));
140        assert!(store.get(&id).is_none());
141    }
142
143    #[test]
144    fn test_remove() {
145        let mut store = WorkbookStore::with_config(10, Duration::from_secs(600));
146        let id = store.insert(make_entry()).unwrap();
147        assert!(store.remove(&id).is_some());
148        assert!(store.get(&id).is_none());
149    }
150}