spreadsheet_mcp/repository/
virtual_workspace.rs

1use super::{ResolvedWorkbookRef, WorkbookRepository, WorkbookSource};
2use crate::caps::BackendCaps;
3use crate::config::ServerConfig;
4use crate::model::{WorkbookDescriptor, WorkbookId, WorkbookListResponse};
5use crate::tools::filters::WorkbookFilter;
6use crate::utils::{hash_bytes_sha256_hex, hash_path_identity, make_short_workbook_id};
7use crate::workbook::WorkbookContext;
8use anyhow::{Result, anyhow};
9use parking_lot::RwLock;
10use std::collections::HashMap;
11use std::path::Path;
12use std::sync::Arc;
13
14#[derive(Debug, Clone)]
15pub struct VirtualWorkbookInput {
16    pub key: String,
17    pub slug: Option<String>,
18    pub bytes: Vec<u8>,
19}
20
21#[derive(Debug, Clone)]
22struct VirtualWorkbook {
23    key: String,
24    slug: String,
25    workbook_id: WorkbookId,
26    short_id: String,
27    revision_id: String,
28    bytes: Arc<Vec<u8>>,
29}
30
31pub struct VirtualWorkspaceRepository {
32    config: Arc<ServerConfig>,
33    entries: RwLock<HashMap<WorkbookId, VirtualWorkbook>>,
34    alias_index: RwLock<HashMap<String, WorkbookId>>,
35}
36
37impl VirtualWorkspaceRepository {
38    pub fn new(config: Arc<ServerConfig>) -> Self {
39        Self {
40            config,
41            entries: RwLock::new(HashMap::new()),
42            alias_index: RwLock::new(HashMap::new()),
43        }
44    }
45
46    pub fn register(&self, input: VirtualWorkbookInput) -> WorkbookId {
47        let key = input.key;
48        let slug = input.slug.unwrap_or_else(|| sanitize_slug(&key));
49        let workbook_id = WorkbookId(hash_path_identity(Path::new(&format!("virtual/{key}"))));
50        let short_id = make_short_workbook_id(&slug, workbook_id.as_str());
51        let revision_id = hash_bytes_sha256_hex(&input.bytes);
52        let entry = VirtualWorkbook {
53            key: key.clone(),
54            slug,
55            workbook_id: workbook_id.clone(),
56            short_id: short_id.clone(),
57            revision_id,
58            bytes: Arc::new(input.bytes),
59        };
60
61        self.entries.write().insert(workbook_id.clone(), entry);
62        let mut aliases = self.alias_index.write();
63        aliases.insert(key.to_ascii_lowercase(), workbook_id.clone());
64        aliases.insert(short_id.to_ascii_lowercase(), workbook_id.clone());
65        aliases.insert(
66            workbook_id.as_str().to_ascii_lowercase(),
67            workbook_id.clone(),
68        );
69        workbook_id
70    }
71
72    fn lookup(&self, id_or_alias: &WorkbookId) -> Option<VirtualWorkbook> {
73        if let Some(entry) = self.entries.read().get(id_or_alias) {
74            return Some(entry.clone());
75        }
76
77        let lowered = id_or_alias.as_str().to_ascii_lowercase();
78        let id = self.alias_index.read().get(&lowered).cloned()?;
79        self.entries.read().get(&id).cloned()
80    }
81}
82
83impl WorkbookRepository for VirtualWorkspaceRepository {
84    fn list(&self, filter: &WorkbookFilter) -> Result<WorkbookListResponse> {
85        let mut workbooks = Vec::new();
86        for entry in self.entries.read().values() {
87            let virtual_path = Path::new(&entry.key);
88            if !filter.matches(&entry.slug, None, virtual_path) {
89                continue;
90            }
91
92            workbooks.push(WorkbookDescriptor {
93                workbook_id: entry.workbook_id.clone(),
94                short_id: entry.short_id.clone(),
95                slug: entry.slug.clone(),
96                folder: None,
97                path: Some(format!("virtual/{}", entry.key)),
98                client_path: None,
99                bytes: entry.bytes.len() as u64,
100                last_modified: None,
101                revision_id: Some(entry.revision_id.clone()),
102                caps: Some(BackendCaps::xlsx()),
103            });
104        }
105
106        workbooks.sort_by(|a, b| a.slug.cmp(&b.slug));
107
108        Ok(WorkbookListResponse {
109            workbooks,
110            next_offset: None,
111        })
112    }
113
114    fn resolve(&self, id_or_alias: &WorkbookId) -> Result<ResolvedWorkbookRef> {
115        let Some(entry) = self.lookup(id_or_alias) else {
116            return Err(anyhow!("workbook id {} not found", id_or_alias.as_str()));
117        };
118
119        Ok(ResolvedWorkbookRef {
120            workbook_id: entry.workbook_id,
121            short_id: entry.short_id,
122            revision_id: Some(entry.revision_id),
123            source: WorkbookSource::Virtual(entry.key),
124        })
125    }
126
127    fn load_context(&self, resolved: &ResolvedWorkbookRef) -> Result<WorkbookContext> {
128        let WorkbookSource::Virtual(_) = &resolved.source else {
129            return Err(anyhow!(
130                "virtual repository cannot load non-virtual workbook"
131            ));
132        };
133
134        let entry = self
135            .entries
136            .read()
137            .get(&resolved.workbook_id)
138            .cloned()
139            .ok_or_else(|| anyhow!("virtual workbook {} not found", resolved.workbook_id.0))?;
140
141        WorkbookContext::load_from_bytes(
142            &self.config,
143            &entry.key,
144            entry.bytes.as_slice(),
145            resolved.workbook_id.clone(),
146            resolved.short_id.clone(),
147            resolved.revision_id.clone(),
148        )
149    }
150}
151
152fn sanitize_slug(value: &str) -> String {
153    let mut out = value
154        .chars()
155        .map(|ch| {
156            if ch.is_ascii_alphanumeric() || ch == '-' || ch == '_' {
157                ch
158            } else {
159                '-'
160            }
161        })
162        .collect::<String>();
163    out.make_ascii_lowercase();
164    while out.contains("--") {
165        out = out.replace("--", "-");
166    }
167    out.trim_matches('-').to_string()
168}
spreadsheet_mcp/repository/virtual_workspace.rs

spreadsheet_mcp/repository/
virtual_workspace.rs
