excel_mcp_server/
store.rs1use std::collections::HashMap;
2use std::time::{Duration, Instant};
3
4use uuid::Uuid;
5
6use crate::error::ExcelMcpError;
7
8pub struct WorkbookEntry {
10 pub id: String,
11 pub data: zavora_xlsx::Workbook,
12 pub read_only: bool,
13 pub last_access: Instant,
14}
15
16pub 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}