1use std::collections::HashMap;
13use std::path::{Path, PathBuf};
14use std::sync::atomic::{AtomicU64, Ordering};
15use std::sync::{Mutex, Once, OnceLock};
16use std::time::Duration;
17
18use anyhow::{Context, Result};
19use rusqlite::backup::Backup;
20use rusqlite::{Connection, OpenFlags, Row};
21use serde::Serialize;
22
23static REGISTER_VEC_EXTENSION: Once = Once::new();
24
25pub const VERSION_STORE_FILES: &[(&str, &str)] = &[
27 ("2025", "mcp_store.db"),
28 ("2022", "mcp_store_v2022.db"),
29 ("2019", "mcp_store_v2019.db"),
30 ("2017", "mcp_store_v2017.db"),
31];
32
33const VERSION_STORE_BYTES: &[(&str, &[u8])] = &[
34 ("2025", include_bytes!("../../mcp_store.db.zst")),
35 ("2022", include_bytes!("../../mcp_store_v2022.db.zst")),
36 ("2019", include_bytes!("../../mcp_store_v2019.db.zst")),
37 ("2017", include_bytes!("../../mcp_store_v2017.db.zst")),
38];
39pub fn resolve_store_path(api_version: &str) -> Result<PathBuf> {
56 let file = VERSION_STORE_FILES
57 .iter()
58 .find(|(label, _)| *label == api_version)
59 .map(|(_, file)| *file)
60 .with_context(|| format!("unknown api_version '{api_version}' — run the 'versions' command to see what's available"))?;
61 let bytes = VERSION_STORE_BYTES
62 .iter()
63 .find(|(label, _)| *label == api_version)
64 .map(|(_, bytes)| *bytes)
65 .with_context(|| format!("no embedded store data for api_version '{api_version}'"))?;
66
67 let mut dir = std::env::temp_dir();
68 dir.push(concat!(env!("CARGO_PKG_NAME"), "-store"));
69 std::fs::create_dir_all(&dir)
70 .with_context(|| format!("failed to create temp dir '{}'", dir.display()))?;
71
72 let path = dir.join(file);
73 let decompressed = zstd::stream::decode_all(bytes).with_context(|| {
95 format!("failed to decompress embedded store data for api_version '{api_version}'")
96 })?;
97 static UNIQUE: AtomicU64 = AtomicU64::new(0);
98 let tmp_path = dir.join(format!(
99 "{file}.{}.{}.tmp",
100 std::process::id(),
101 UNIQUE.fetch_add(1, Ordering::Relaxed)
102 ));
103 std::fs::write(&tmp_path, decompressed).with_context(|| {
104 format!(
105 "failed to extract embedded store data to '{}'",
106 tmp_path.display()
107 )
108 })?;
109 std::fs::rename(&tmp_path, &path).with_context(|| {
110 format!(
111 "failed to move extracted store data into place at '{}'",
112 path.display()
113 )
114 })?;
115
116 Ok(path)
117}
118
119const ENDPOINT_COLUMNS: &str = "operation_id, path, method, summary, description, input_schema, output_schema, auth_scheme_ref";
120
121fn register_vec_extension() {
125 REGISTER_VEC_EXTENSION.call_once(|| unsafe {
126 #[allow(clippy::missing_transmute_annotations)]
127 rusqlite::ffi::sqlite3_auto_extension(Some(std::mem::transmute(
128 sqlite_vec::sqlite3_vec_init as *const (),
129 )));
130 });
131}
132
133pub fn open_store(path: &Path) -> Result<Connection> {
137 register_vec_extension();
138 Connection::open_with_flags(path, OpenFlags::SQLITE_OPEN_READ_ONLY)
139 .with_context(|| format!("failed to open '{}'", path.display()))
140}
141
142pub fn open_store_read_write(path: &Path) -> Result<Connection> {
147 register_vec_extension();
148 Connection::open_with_flags(path, OpenFlags::SQLITE_OPEN_READ_WRITE)
149 .with_context(|| format!("failed to open '{}'", path.display()))
150}
151
152pub fn cached_store_connection(api_version: &str) -> Result<&'static Mutex<Connection>> {
170 let path = resolve_store_path(api_version)?;
171 cached_in_memory_connection(api_version, &path)
172}
173
174fn cached_in_memory_connection(cache_key: &str, path: &Path) -> Result<&'static Mutex<Connection>> {
181 static CACHE: OnceLock<Mutex<HashMap<String, &'static Mutex<Connection>>>> = OnceLock::new();
182 let cache = CACHE.get_or_init(|| Mutex::new(HashMap::new()));
183
184 if let Some(conn) = cache.lock().unwrap().get(cache_key) {
185 return Ok(conn);
186 }
187
188 let disk_conn = open_store(path)?;
189 let mut mem_conn =
190 Connection::open_in_memory().context("failed to open an in-memory SQLite connection")?;
191 Backup::new(&disk_conn, &mut mem_conn)
192 .context("failed to start SQLite backup into memory")?
193 .run_to_completion(i32::MAX, Duration::from_millis(0), None)
194 .with_context(|| format!("failed to back up '{}' into memory", path.display()))?;
195 drop(disk_conn);
196
197 let leaked: &'static Mutex<Connection> = Box::leak(Box::new(Mutex::new(mem_conn)));
198 let mut cache = cache.lock().unwrap();
199 Ok(*cache.entry(cache_key.to_string()).or_insert(leaked))
200}
201
202#[derive(Debug, Clone, Serialize)]
203pub struct EndpointRecord {
204 pub operation_id: String,
205 pub path: String,
206 pub method: String,
207 pub summary: Option<String>,
208 pub description: Option<String>,
209 pub input_schema: serde_json::Value,
210 pub output_schema: serde_json::Value,
211 pub auth_scheme_ref: Option<String>,
212}
213
214fn row_to_endpoint(row: &Row) -> rusqlite::Result<EndpointRecord> {
215 let input_schema: String = row.get(5)?;
216 let output_schema: String = row.get(6)?;
217 Ok(EndpointRecord {
218 operation_id: row.get(0)?,
219 path: row.get(1)?,
220 method: row.get(2)?,
221 summary: row.get(3)?,
222 description: row.get(4)?,
223 input_schema: serde_json::from_str(&input_schema).unwrap_or(serde_json::Value::Null),
224 output_schema: serde_json::from_str(&output_schema).unwrap_or(serde_json::Value::Null),
225 auth_scheme_ref: row.get(7)?,
226 })
227}
228
229pub fn get_endpoint(conn: &Connection, operation_id: &str) -> Result<Option<EndpointRecord>> {
230 let mut stmt = conn.prepare(&format!(
231 "SELECT {ENDPOINT_COLUMNS} FROM endpoints WHERE operation_id = ?1"
232 ))?;
233 match stmt.query_row([operation_id], row_to_endpoint) {
234 Ok(record) => Ok(Some(record)),
235 Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
236 Err(err) => Err(err.into()),
237 }
238}
239
240pub fn list_endpoints(conn: &Connection) -> Result<Vec<EndpointRecord>> {
241 let mut stmt = conn.prepare(&format!("SELECT {ENDPOINT_COLUMNS} FROM endpoints"))?;
242 let rows = stmt.query_map([], row_to_endpoint)?;
243 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
244}
245
246#[derive(Debug, Clone, Serialize)]
247pub struct SearchResult {
248 pub operation_id: String,
249 pub summary: Option<String>,
250 pub similarity: f64,
251}
252
253pub fn search_endpoints(
268 conn: &Connection,
269 query_embedding: &[f32],
270 limit: usize,
271) -> Result<Vec<SearchResult>> {
272 let blob: Vec<u8> = query_embedding
273 .iter()
274 .flat_map(|value| value.to_le_bytes())
275 .collect();
276
277 let mut stmt = conn.prepare(
278 "SELECT e.operation_id, e.summary, s.distance
279 FROM semantic_endpoints s
280 JOIN endpoints e ON e.operation_id = s.operation_id
281 WHERE s.embedding MATCH ?1 AND k = ?2
282 ORDER BY s.distance",
283 )?;
284 let rows = stmt.query_map(rusqlite::params![blob, limit], |row| {
285 let distance: f64 = row.get(2)?;
286 Ok(SearchResult {
287 operation_id: row.get(0)?,
288 summary: row.get(1)?,
289 similarity: 1.0 - distance,
290 })
291 })?;
292 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
293}
294
295#[cfg(test)]
296mod tests {
297 use super::*;
298
299 #[test]
305 fn every_version_store_file_has_embedded_bytes() {
306 let file_labels: std::collections::HashSet<_> = VERSION_STORE_FILES
307 .iter()
308 .map(|(label, _)| *label)
309 .collect();
310 let byte_labels: std::collections::HashSet<_> = VERSION_STORE_BYTES
311 .iter()
312 .map(|(label, _)| *label)
313 .collect();
314 assert_eq!(file_labels, byte_labels);
315 }
316
317 #[test]
325 fn resolve_store_path_survives_concurrent_calls() {
326 let api_version = VERSION_STORE_FILES[0].0.to_string();
327 let handles: Vec<_> = (0..16)
328 .map(|_| {
329 let api_version = api_version.clone();
330 std::thread::spawn(move || {
331 let path = resolve_store_path(&api_version).unwrap();
332 open_store(&path).unwrap();
333 })
334 })
335 .collect();
336 for handle in handles {
337 handle.join().unwrap();
338 }
339 }
340
341 fn seeded_store(path: &Path) -> Connection {
347 unsafe {
348 #[allow(clippy::missing_transmute_annotations)]
349 rusqlite::ffi::sqlite3_auto_extension(Some(std::mem::transmute(
350 sqlite_vec::sqlite3_vec_init as *const (),
351 )));
352 }
353 let conn = Connection::open(path).unwrap();
354 conn.execute(
355 "CREATE TABLE endpoints (
356 operation_id TEXT PRIMARY KEY,
357 path TEXT NOT NULL,
358 method TEXT NOT NULL,
359 summary TEXT,
360 description TEXT,
361 input_schema TEXT NOT NULL,
362 output_schema TEXT NOT NULL,
363 auth_scheme_ref TEXT
364 )",
365 [],
366 )
367 .unwrap();
368 conn.execute(
369 "CREATE VIRTUAL TABLE semantic_endpoints USING vec0(
370 operation_id TEXT PRIMARY KEY,
371 embedding FLOAT[4]
372 )",
373 [],
374 )
375 .unwrap();
376 conn.execute(
377 "INSERT INTO endpoints (operation_id, path, method, summary, description, input_schema, output_schema, auth_scheme_ref)
378 VALUES ('listWidgets', '/widgets', 'GET', 'List widgets', NULL, '{}', '[]', NULL)",
379 [],
380 )
381 .unwrap();
382 let embedding: Vec<u8> = [1.0f32, 0.0, 0.0, 0.0]
383 .iter()
384 .flat_map(|v| v.to_le_bytes())
385 .collect();
386 conn.execute(
387 "INSERT INTO semantic_endpoints (operation_id, embedding) VALUES ('listWidgets', ?1)",
388 rusqlite::params![embedding],
389 )
390 .unwrap();
391 conn
392 }
393
394 #[test]
395 fn get_endpoint_returns_a_seeded_row() {
396 let dir = tempfile::tempdir().unwrap();
397 let path = dir.path().join("mcp_store.db");
398 let _conn = seeded_store(&path);
399
400 let store = open_store(&path).unwrap();
401 let endpoint = get_endpoint(&store, "listWidgets").unwrap().unwrap();
402 assert_eq!(endpoint.path, "/widgets");
403 assert_eq!(endpoint.method, "GET");
404 assert_eq!(endpoint.summary.as_deref(), Some("List widgets"));
405 }
406
407 #[test]
408 fn get_endpoint_returns_none_for_an_unknown_operation() {
409 let dir = tempfile::tempdir().unwrap();
410 let path = dir.path().join("mcp_store.db");
411 let _conn = seeded_store(&path);
412
413 let store = open_store(&path).unwrap();
414 assert!(get_endpoint(&store, "unknownOp").unwrap().is_none());
415 }
416
417 #[test]
418 fn list_endpoints_returns_every_row() {
419 let dir = tempfile::tempdir().unwrap();
420 let path = dir.path().join("mcp_store.db");
421 let _conn = seeded_store(&path);
422
423 let store = open_store(&path).unwrap();
424 let endpoints = list_endpoints(&store).unwrap();
425 assert_eq!(endpoints.len(), 1);
426 assert_eq!(endpoints[0].operation_id, "listWidgets");
427 }
428
429 #[test]
430 fn search_endpoints_finds_the_nearest_neighbor() {
431 let dir = tempfile::tempdir().unwrap();
432 let path = dir.path().join("mcp_store.db");
433 let _conn = seeded_store(&path);
434
435 let store = open_store(&path).unwrap();
436 let results = search_endpoints(&store, &[1.0, 0.0, 0.0, 0.0], 5).unwrap();
437 assert_eq!(results.len(), 1);
438 assert_eq!(results[0].operation_id, "listWidgets");
439 assert!(results[0].similarity > 0.99);
440 }
441
442 #[test]
443 fn cached_in_memory_connection_serves_seeded_data() {
444 let dir = tempfile::tempdir().unwrap();
445 let path = dir.path().join("mcp_store.db");
446 let _conn = seeded_store(&path);
447
448 let cached = cached_in_memory_connection("cached-serves-seeded-data", &path).unwrap();
449 let endpoint = get_endpoint(&cached.lock().unwrap(), "listWidgets")
450 .unwrap()
451 .unwrap();
452 assert_eq!(endpoint.path, "/widgets");
453
454 let results = search_endpoints(&cached.lock().unwrap(), &[1.0, 0.0, 0.0, 0.0], 5).unwrap();
458 assert_eq!(results.len(), 1);
459 assert_eq!(results[0].operation_id, "listWidgets");
460 }
461
462 #[test]
463 fn cached_in_memory_connection_holds_no_lingering_lock_on_the_disk_file() {
464 let dir = tempfile::tempdir().unwrap();
465 let path = dir.path().join("mcp_store.db");
466 let _conn = seeded_store(&path);
467
468 let cached = cached_in_memory_connection("cached-releases-disk-handle", &path).unwrap();
469
470 std::fs::remove_file(&path).unwrap();
474
475 let endpoint = get_endpoint(&cached.lock().unwrap(), "listWidgets")
476 .unwrap()
477 .unwrap();
478 assert_eq!(endpoint.path, "/widgets");
479 }
480
481 #[test]
482 fn cached_in_memory_connection_reuses_the_same_connection_for_the_same_key() {
483 let dir = tempfile::tempdir().unwrap();
484 let path = dir.path().join("mcp_store.db");
485 let _conn = seeded_store(&path);
486
487 let first = cached_in_memory_connection("cached-reuses-same-key", &path).unwrap()
488 as *const Mutex<Connection>;
489 let second = cached_in_memory_connection("cached-reuses-same-key", &path).unwrap()
490 as *const Mutex<Connection>;
491 assert_eq!(
492 first, second,
493 "expected the same cached connection, not a fresh backup"
494 );
495 }
496}