spreadsheet_mcp/
state.rs

1use crate::config::ServerConfig;
2#[cfg(feature = "recalc")]
3use crate::fork::{ForkConfig, ForkRegistry};
4use crate::model::{WorkbookId, WorkbookListResponse};
5#[cfg(feature = "recalc")]
6use crate::recalc::{
7    GlobalRecalcLock, LibreOfficeBackend, RecalcBackend, RecalcConfig, create_executor,
8};
9use crate::tools::filters::WorkbookFilter;
10use crate::utils::{hash_path_metadata, make_short_workbook_id};
11use crate::workbook::{WorkbookContext, build_workbook_list};
12use anyhow::{Result, anyhow};
13use lru::LruCache;
14use parking_lot::RwLock;
15use std::collections::HashMap;
16use std::fs;
17use std::num::NonZeroUsize;
18use std::path::{Path, PathBuf};
19use std::sync::Arc;
20use tokio::task;
21use walkdir::WalkDir;
22
23pub struct AppState {
24    config: Arc<ServerConfig>,
25    cache: RwLock<LruCache<WorkbookId, Arc<WorkbookContext>>>,
26    index: RwLock<HashMap<WorkbookId, PathBuf>>,
27    alias_index: RwLock<HashMap<String, WorkbookId>>,
28    #[cfg(feature = "recalc")]
29    fork_registry: Option<Arc<ForkRegistry>>,
30    #[cfg(feature = "recalc")]
31    recalc_backend: Option<Arc<dyn RecalcBackend>>,
32    #[cfg(feature = "recalc")]
33    recalc_semaphore: Option<GlobalRecalcLock>,
34}
35
36impl AppState {
37    pub fn new(config: Arc<ServerConfig>) -> Self {
38        let capacity = NonZeroUsize::new(config.cache_capacity.max(1)).unwrap();
39
40        #[cfg(feature = "recalc")]
41        let (fork_registry, recalc_backend, recalc_semaphore) = if config.recalc_enabled {
42            let fork_config = ForkConfig::default();
43            let registry = ForkRegistry::new(fork_config)
44                .map(Arc::new)
45                .map_err(|e| tracing::warn!("failed to init fork registry: {}", e))
46                .ok();
47
48            if let Some(registry) = &registry {
49                registry.clone().start_cleanup_task();
50            }
51
52            let executor = create_executor(&RecalcConfig::default());
53            let backend: Arc<dyn RecalcBackend> = Arc::new(LibreOfficeBackend::new(executor));
54            let backend = if backend.is_available() {
55                Some(backend)
56            } else {
57                tracing::warn!("recalc backend not available (soffice not found)");
58                None
59            };
60
61            let semaphore = GlobalRecalcLock::new(config.max_concurrent_recalcs);
62
63            (registry, backend, Some(semaphore))
64        } else {
65            (None, None, None)
66        };
67
68        Self {
69            config,
70            cache: RwLock::new(LruCache::new(capacity)),
71            index: RwLock::new(HashMap::new()),
72            alias_index: RwLock::new(HashMap::new()),
73            #[cfg(feature = "recalc")]
74            fork_registry,
75            #[cfg(feature = "recalc")]
76            recalc_backend,
77            #[cfg(feature = "recalc")]
78            recalc_semaphore,
79        }
80    }
81
82    pub fn config(&self) -> Arc<ServerConfig> {
83        self.config.clone()
84    }
85
86    #[cfg(feature = "recalc")]
87    pub fn fork_registry(&self) -> Option<&Arc<ForkRegistry>> {
88        self.fork_registry.as_ref()
89    }
90
91    #[cfg(feature = "recalc")]
92    pub fn recalc_backend(&self) -> Option<&Arc<dyn RecalcBackend>> {
93        self.recalc_backend.as_ref()
94    }
95
96    #[cfg(feature = "recalc")]
97    pub fn recalc_semaphore(&self) -> Option<&GlobalRecalcLock> {
98        self.recalc_semaphore.as_ref()
99    }
100
101    pub fn list_workbooks(&self, filter: WorkbookFilter) -> Result<WorkbookListResponse> {
102        let response = build_workbook_list(&self.config, &filter)?;
103        {
104            let mut index = self.index.write();
105            let mut aliases = self.alias_index.write();
106            for descriptor in &response.workbooks {
107                let abs_path = self.config.resolve_path(PathBuf::from(&descriptor.path));
108                index.insert(descriptor.workbook_id.clone(), abs_path);
109                aliases.insert(
110                    descriptor.short_id.to_ascii_lowercase(),
111                    descriptor.workbook_id.clone(),
112                );
113            }
114        }
115        Ok(response)
116    }
117
118    pub async fn open_workbook(&self, workbook_id: &WorkbookId) -> Result<Arc<WorkbookContext>> {
119        let canonical = self.canonicalize_workbook_id(workbook_id)?;
120        {
121            let mut cache = self.cache.write();
122            if let Some(entry) = cache.get(&canonical) {
123                return Ok(entry.clone());
124            }
125        }
126
127        let path = self.resolve_workbook_path(&canonical)?;
128        let config = self.config.clone();
129        let path_buf = path.clone();
130        let workbook_id_clone = canonical.clone();
131        let workbook =
132            task::spawn_blocking(move || WorkbookContext::load(&config, &path_buf)).await??;
133        let workbook = Arc::new(workbook);
134
135        {
136            let mut aliases = self.alias_index.write();
137            aliases.insert(
138                workbook.short_id.to_ascii_lowercase(),
139                workbook_id_clone.clone(),
140            );
141        }
142
143        let mut cache = self.cache.write();
144        cache.put(workbook_id_clone, workbook.clone());
145        Ok(workbook)
146    }
147
148    pub fn close_workbook(&self, workbook_id: &WorkbookId) -> Result<()> {
149        let canonical = self.canonicalize_workbook_id(workbook_id)?;
150        let mut cache = self.cache.write();
151        cache.pop(&canonical);
152        Ok(())
153    }
154
155    fn resolve_workbook_path(&self, workbook_id: &WorkbookId) -> Result<PathBuf> {
156        if let Some(path) = self.index.read().get(workbook_id).cloned() {
157            return Ok(path);
158        }
159
160        let located = self.scan_for_workbook(workbook_id.as_str())?;
161        self.register_location(&located);
162        Ok(located.path)
163    }
164
165    fn canonicalize_workbook_id(&self, workbook_id: &WorkbookId) -> Result<WorkbookId> {
166        if self.index.read().contains_key(workbook_id) {
167            return Ok(workbook_id.clone());
168        }
169        let aliases = self.alias_index.read();
170        if let Some(mapped) = aliases.get(workbook_id.as_str()).cloned() {
171            return Ok(mapped);
172        }
173        let lowered = workbook_id.as_str().to_ascii_lowercase();
174        if lowered != workbook_id.as_str()
175            && let Some(mapped) = aliases.get(&lowered).cloned()
176        {
177            return Ok(mapped);
178        }
179
180        let located = self.scan_for_workbook(workbook_id.as_str())?;
181        let canonical = located.workbook_id.clone();
182        self.register_location(&located);
183        Ok(canonical)
184    }
185
186    fn scan_for_workbook(&self, candidate: &str) -> Result<LocatedWorkbook> {
187        let candidate_lower = candidate.to_ascii_lowercase();
188
189        if let Some(single) = self.config.single_workbook() {
190            let metadata = fs::metadata(single)?;
191            let canonical = WorkbookId(hash_path_metadata(single, &metadata));
192            let slug = single
193                .file_stem()
194                .map(|s| s.to_string_lossy().to_string())
195                .unwrap_or_else(|| "workbook".to_string());
196            let short_id = make_short_workbook_id(&slug, canonical.as_str());
197            if candidate == canonical.as_str() || candidate_lower == short_id {
198                return Ok(LocatedWorkbook {
199                    workbook_id: canonical,
200                    short_id,
201                    path: single.to_path_buf(),
202                });
203            }
204            return Err(anyhow!(
205                "workbook id {} not found in single-workbook mode (expected {} or {})",
206                candidate,
207                canonical.as_str(),
208                short_id
209            ));
210        }
211
212        for entry in WalkDir::new(&self.config.workspace_root) {
213            let entry = entry?;
214            if !entry.file_type().is_file() {
215                continue;
216            }
217            let path = entry.path();
218            if !has_supported_extension(&self.config.supported_extensions, path) {
219                continue;
220            }
221            let metadata = entry.metadata()?;
222            let canonical = WorkbookId(hash_path_metadata(path, &metadata));
223            let slug = path
224                .file_stem()
225                .map(|s| s.to_string_lossy().to_string())
226                .unwrap_or_else(|| "workbook".to_string());
227            let short_id = make_short_workbook_id(&slug, canonical.as_str());
228            if candidate == canonical.as_str() || candidate_lower == short_id {
229                return Ok(LocatedWorkbook {
230                    workbook_id: canonical,
231                    short_id,
232                    path: path.to_path_buf(),
233                });
234            }
235        }
236        Err(anyhow!("workbook id {} not found", candidate))
237    }
238
239    fn register_location(&self, located: &LocatedWorkbook) {
240        self.index
241            .write()
242            .insert(located.workbook_id.clone(), located.path.clone());
243        self.alias_index.write().insert(
244            located.short_id.to_ascii_lowercase(),
245            located.workbook_id.clone(),
246        );
247    }
248}
249
250struct LocatedWorkbook {
251    workbook_id: WorkbookId,
252    short_id: String,
253    path: PathBuf,
254}
255
256fn has_supported_extension(allowed: &[String], path: &Path) -> bool {
257    path.extension()
258        .and_then(|ext| ext.to_str())
259        .map(|ext| {
260            let lower = ext.to_ascii_lowercase();
261            allowed.iter().any(|candidate| candidate == &lower)
262        })
263        .unwrap_or(false)
264}