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])] = &[
44 ("2025", include_bytes!("../../mcp_store.db.zst")),
45 ("2022", include_bytes!("../../mcp_store_v2022.db.zst")),
46 ("2019", include_bytes!("../../mcp_store_v2019.db.zst")),
47 ("2017", include_bytes!("../../mcp_store_v2017.db.zst")),
48];
49pub fn resolve_store_path(api_version: &str) -> Result<PathBuf> {
66 let file = VERSION_STORE_FILES
67 .iter()
68 .find(|(label, _)| *label == api_version)
69 .map(|(_, file)| *file)
70 .with_context(|| format!("unknown api_version '{api_version}' — run the 'versions' command to see what's available"))?;
71 let bytes = VERSION_STORE_BYTES
72 .iter()
73 .find(|(label, _)| *label == api_version)
74 .map(|(_, bytes)| *bytes)
75 .with_context(|| format!("no embedded store data for api_version '{api_version}'"))?;
76
77 let mut dir = std::env::temp_dir();
78 dir.push(concat!(env!("CARGO_PKG_NAME"), "-store"));
79 std::fs::create_dir_all(&dir)
80 .with_context(|| format!("failed to create temp dir '{}'", dir.display()))?;
81
82 let path = dir.join(file);
83 let decompressed = zstd::stream::decode_all(bytes).with_context(|| {
105 format!("failed to decompress embedded store data for api_version '{api_version}'")
106 })?;
107 static UNIQUE: AtomicU64 = AtomicU64::new(0);
108 let tmp_path = dir.join(format!(
109 "{file}.{}.{}.tmp",
110 std::process::id(),
111 UNIQUE.fetch_add(1, Ordering::Relaxed)
112 ));
113 std::fs::write(&tmp_path, decompressed).with_context(|| {
114 format!(
115 "failed to extract embedded store data to '{}'",
116 tmp_path.display()
117 )
118 })?;
119 std::fs::rename(&tmp_path, &path).with_context(|| {
120 format!(
121 "failed to move extracted store data into place at '{}'",
122 path.display()
123 )
124 })?;
125
126 Ok(path)
127}
128
129const ENDPOINT_COLUMNS: &str = "operation_id, path, method, summary, description, input_schema, output_schema, auth_scheme_ref";
130
131fn register_vec_extension() {
135 REGISTER_VEC_EXTENSION.call_once(|| unsafe {
136 #[allow(clippy::missing_transmute_annotations)]
137 rusqlite::ffi::sqlite3_auto_extension(Some(std::mem::transmute(
138 sqlite_vec::sqlite3_vec_init as *const (),
139 )));
140 });
141}
142
143pub fn open_store(path: &Path) -> Result<Connection> {
147 register_vec_extension();
148 Connection::open_with_flags(path, OpenFlags::SQLITE_OPEN_READ_ONLY)
149 .with_context(|| format!("failed to open '{}'", path.display()))
150}
151
152pub fn open_store_read_write(path: &Path) -> Result<Connection> {
157 register_vec_extension();
158 Connection::open_with_flags(path, OpenFlags::SQLITE_OPEN_READ_WRITE)
159 .with_context(|| format!("failed to open '{}'", path.display()))
160}
161
162pub fn cached_store_connection(api_version: &str) -> Result<&'static Mutex<Connection>> {
180 let path = resolve_store_path(api_version)?;
181 cached_in_memory_connection(api_version, &path)
182}
183
184fn cached_in_memory_connection(cache_key: &str, path: &Path) -> Result<&'static Mutex<Connection>> {
191 static CACHE: OnceLock<Mutex<HashMap<String, &'static Mutex<Connection>>>> = OnceLock::new();
192 let cache = CACHE.get_or_init(|| Mutex::new(HashMap::new()));
193
194 if let Some(conn) = cache.lock().unwrap().get(cache_key) {
195 return Ok(conn);
196 }
197
198 let disk_conn = open_store(path)?;
199 let mut mem_conn =
200 Connection::open_in_memory().context("failed to open an in-memory SQLite connection")?;
201 Backup::new(&disk_conn, &mut mem_conn)
202 .context("failed to start SQLite backup into memory")?
203 .run_to_completion(i32::MAX, Duration::from_millis(0), None)
204 .with_context(|| format!("failed to back up '{}' into memory", path.display()))?;
205 drop(disk_conn);
206
207 let leaked: &'static Mutex<Connection> = Box::leak(Box::new(Mutex::new(mem_conn)));
208 let mut cache = cache.lock().unwrap();
209 Ok(*cache.entry(cache_key.to_string()).or_insert(leaked))
210}
211
212#[derive(Debug, Clone, Serialize)]
213pub struct EndpointRecord {
214 pub operation_id: String,
215 pub path: String,
216 pub method: String,
217 pub summary: Option<String>,
218 pub description: Option<String>,
219 pub input_schema: serde_json::Value,
220 pub output_schema: serde_json::Value,
221 pub auth_scheme_ref: Option<String>,
222}
223
224fn row_to_endpoint(row: &Row) -> rusqlite::Result<EndpointRecord> {
225 let input_schema: String = row.get(5)?;
226 let output_schema: String = row.get(6)?;
227 Ok(EndpointRecord {
228 operation_id: row.get(0)?,
229 path: row.get(1)?,
230 method: row.get(2)?,
231 summary: row.get(3)?,
232 description: row.get(4)?,
233 input_schema: serde_json::from_str(&input_schema).unwrap_or(serde_json::Value::Null),
234 output_schema: serde_json::from_str(&output_schema).unwrap_or(serde_json::Value::Null),
235 auth_scheme_ref: row.get(7)?,
236 })
237}
238
239pub fn get_endpoint(conn: &Connection, operation_id: &str) -> Result<Option<EndpointRecord>> {
240 let mut stmt = conn.prepare(&format!(
241 "SELECT {ENDPOINT_COLUMNS} FROM endpoints WHERE operation_id = ?1"
242 ))?;
243 match stmt.query_row([operation_id], row_to_endpoint) {
244 Ok(record) => Ok(Some(record)),
245 Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
246 Err(err) => Err(err.into()),
247 }
248}
249
250pub fn list_endpoints(conn: &Connection) -> Result<Vec<EndpointRecord>> {
251 let mut stmt = conn.prepare(&format!("SELECT {ENDPOINT_COLUMNS} FROM endpoints"))?;
252 let rows = stmt.query_map([], row_to_endpoint)?;
253 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
254}
255
256#[derive(Debug, Clone, Serialize)]
257pub struct SearchResult {
258 pub operation_id: String,
259 pub summary: Option<String>,
260 pub similarity: f64,
261}
262
263pub fn search_endpoints(
278 conn: &Connection,
279 query_embedding: &[f32],
280 limit: usize,
281) -> Result<Vec<SearchResult>> {
282 let blob: Vec<u8> = query_embedding
283 .iter()
284 .flat_map(|value| value.to_le_bytes())
285 .collect();
286
287 let mut stmt = conn.prepare(
288 "SELECT e.operation_id, e.summary, s.distance
289 FROM semantic_endpoints s
290 JOIN endpoints e ON e.operation_id = s.operation_id
291 WHERE s.embedding MATCH ?1 AND k = ?2
292 ORDER BY s.distance",
293 )?;
294 let rows = stmt.query_map(rusqlite::params![blob, limit], |row| {
295 let distance: f64 = row.get(2)?;
296 Ok(SearchResult {
297 operation_id: row.get(0)?,
298 summary: row.get(1)?,
299 similarity: 1.0 - distance,
300 })
301 })?;
302 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
303}
304
305#[cfg(test)]
306mod tests {
307 use super::*;
308
309 #[test]
315 fn every_version_store_file_has_embedded_bytes() {
316 let file_labels: std::collections::HashSet<_> = VERSION_STORE_FILES
317 .iter()
318 .map(|(label, _)| *label)
319 .collect();
320 let byte_labels: std::collections::HashSet<_> = VERSION_STORE_BYTES
321 .iter()
322 .map(|(label, _)| *label)
323 .collect();
324 assert_eq!(file_labels, byte_labels);
325 }
326
327 #[test]
335 fn resolve_store_path_survives_concurrent_calls() {
336 let api_version = VERSION_STORE_FILES[0].0.to_string();
337 let handles: Vec<_> = (0..16)
338 .map(|_| {
339 let api_version = api_version.clone();
340 std::thread::spawn(move || {
341 let path = resolve_store_path(&api_version).unwrap();
342 open_store(&path).unwrap();
343 })
344 })
345 .collect();
346 for handle in handles {
347 handle.join().unwrap();
348 }
349 }
350
351 fn seeded_store(path: &Path) -> Connection {
357 unsafe {
358 #[allow(clippy::missing_transmute_annotations)]
359 rusqlite::ffi::sqlite3_auto_extension(Some(std::mem::transmute(
360 sqlite_vec::sqlite3_vec_init as *const (),
361 )));
362 }
363 let conn = Connection::open(path).unwrap();
364 conn.execute(
365 "CREATE TABLE endpoints (
366 operation_id TEXT PRIMARY KEY,
367 path TEXT NOT NULL,
368 method TEXT NOT NULL,
369 summary TEXT,
370 description TEXT,
371 input_schema TEXT NOT NULL,
372 output_schema TEXT NOT NULL,
373 auth_scheme_ref TEXT
374 )",
375 [],
376 )
377 .unwrap();
378 conn.execute(
379 "CREATE VIRTUAL TABLE semantic_endpoints USING vec0(
380 operation_id TEXT PRIMARY KEY,
381 embedding FLOAT[4]
382 )",
383 [],
384 )
385 .unwrap();
386 conn.execute(
387 "INSERT INTO endpoints (operation_id, path, method, summary, description, input_schema, output_schema, auth_scheme_ref)
388 VALUES ('listWidgets', '/widgets', 'GET', 'List widgets', NULL, '{}', '[]', NULL)",
389 [],
390 )
391 .unwrap();
392 let embedding: Vec<u8> = [1.0f32, 0.0, 0.0, 0.0]
393 .iter()
394 .flat_map(|v| v.to_le_bytes())
395 .collect();
396 conn.execute(
397 "INSERT INTO semantic_endpoints (operation_id, embedding) VALUES ('listWidgets', ?1)",
398 rusqlite::params![embedding],
399 )
400 .unwrap();
401 conn
402 }
403
404 #[test]
405 fn get_endpoint_returns_a_seeded_row() {
406 let dir = tempfile::tempdir().unwrap();
407 let path = dir.path().join("mcp_store.db");
408 let _conn = seeded_store(&path);
409
410 let store = open_store(&path).unwrap();
411 let endpoint = get_endpoint(&store, "listWidgets").unwrap().unwrap();
412 assert_eq!(endpoint.path, "/widgets");
413 assert_eq!(endpoint.method, "GET");
414 assert_eq!(endpoint.summary.as_deref(), Some("List widgets"));
415 }
416
417 #[test]
418 fn get_endpoint_returns_none_for_an_unknown_operation() {
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 assert!(get_endpoint(&store, "unknownOp").unwrap().is_none());
425 }
426
427 #[test]
428 fn list_endpoints_returns_every_row() {
429 let dir = tempfile::tempdir().unwrap();
430 let path = dir.path().join("mcp_store.db");
431 let _conn = seeded_store(&path);
432
433 let store = open_store(&path).unwrap();
434 let endpoints = list_endpoints(&store).unwrap();
435 assert_eq!(endpoints.len(), 1);
436 assert_eq!(endpoints[0].operation_id, "listWidgets");
437 }
438
439 #[test]
440 fn search_endpoints_finds_the_nearest_neighbor() {
441 let dir = tempfile::tempdir().unwrap();
442 let path = dir.path().join("mcp_store.db");
443 let _conn = seeded_store(&path);
444
445 let store = open_store(&path).unwrap();
446 let results = search_endpoints(&store, &[1.0, 0.0, 0.0, 0.0], 5).unwrap();
447 assert_eq!(results.len(), 1);
448 assert_eq!(results[0].operation_id, "listWidgets");
449 assert!(results[0].similarity > 0.99);
450 }
451
452 #[test]
453 fn cached_in_memory_connection_serves_seeded_data() {
454 let dir = tempfile::tempdir().unwrap();
455 let path = dir.path().join("mcp_store.db");
456 let _conn = seeded_store(&path);
457
458 let cached = cached_in_memory_connection("cached-serves-seeded-data", &path).unwrap();
459 let endpoint = get_endpoint(&cached.lock().unwrap(), "listWidgets")
460 .unwrap()
461 .unwrap();
462 assert_eq!(endpoint.path, "/widgets");
463
464 let results = search_endpoints(&cached.lock().unwrap(), &[1.0, 0.0, 0.0, 0.0], 5).unwrap();
468 assert_eq!(results.len(), 1);
469 assert_eq!(results[0].operation_id, "listWidgets");
470 }
471
472 #[test]
473 fn cached_in_memory_connection_holds_no_lingering_lock_on_the_disk_file() {
474 let dir = tempfile::tempdir().unwrap();
475 let path = dir.path().join("mcp_store.db");
476 let _conn = seeded_store(&path);
477
478 let cached = cached_in_memory_connection("cached-releases-disk-handle", &path).unwrap();
479
480 std::fs::remove_file(&path).unwrap();
484
485 let endpoint = get_endpoint(&cached.lock().unwrap(), "listWidgets")
486 .unwrap()
487 .unwrap();
488 assert_eq!(endpoint.path, "/widgets");
489 }
490
491 #[test]
492 fn cached_in_memory_connection_reuses_the_same_connection_for_the_same_key() {
493 let dir = tempfile::tempdir().unwrap();
494 let path = dir.path().join("mcp_store.db");
495 let _conn = seeded_store(&path);
496
497 let first = cached_in_memory_connection("cached-reuses-same-key", &path).unwrap()
498 as *const Mutex<Connection>;
499 let second = cached_in_memory_connection("cached-reuses-same-key", &path).unwrap()
500 as *const Mutex<Connection>;
501 assert_eq!(
502 first, second,
503 "expected the same cached connection, not a fresh backup"
504 );
505 }
506}