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")),
35 ("2022", include_bytes!("../../mcp_store_v2022.db")),
36 ("2019", include_bytes!("../../mcp_store_v2019.db")),
37 ("2017", include_bytes!("../../mcp_store_v2017.db")),
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 static UNIQUE: AtomicU64 = AtomicU64::new(0);
91 let tmp_path = dir.join(format!(
92 "{file}.{}.{}.tmp",
93 std::process::id(),
94 UNIQUE.fetch_add(1, Ordering::Relaxed)
95 ));
96 std::fs::write(&tmp_path, bytes).with_context(|| {
97 format!(
98 "failed to extract embedded store data to '{}'",
99 tmp_path.display()
100 )
101 })?;
102 std::fs::rename(&tmp_path, &path).with_context(|| {
103 format!(
104 "failed to move extracted store data into place at '{}'",
105 path.display()
106 )
107 })?;
108
109 Ok(path)
110}
111
112const ENDPOINT_COLUMNS: &str = "operation_id, path, method, summary, description, input_schema, output_schema, auth_scheme_ref";
113
114fn register_vec_extension() {
118 REGISTER_VEC_EXTENSION.call_once(|| unsafe {
119 #[allow(clippy::missing_transmute_annotations)]
120 rusqlite::ffi::sqlite3_auto_extension(Some(std::mem::transmute(
121 sqlite_vec::sqlite3_vec_init as *const (),
122 )));
123 });
124}
125
126pub fn open_store(path: &Path) -> Result<Connection> {
130 register_vec_extension();
131 Connection::open_with_flags(path, OpenFlags::SQLITE_OPEN_READ_ONLY)
132 .with_context(|| format!("failed to open '{}'", path.display()))
133}
134
135pub fn open_store_read_write(path: &Path) -> Result<Connection> {
140 register_vec_extension();
141 Connection::open_with_flags(path, OpenFlags::SQLITE_OPEN_READ_WRITE)
142 .with_context(|| format!("failed to open '{}'", path.display()))
143}
144
145pub fn cached_store_connection(api_version: &str) -> Result<&'static Mutex<Connection>> {
163 let path = resolve_store_path(api_version)?;
164 cached_in_memory_connection(api_version, &path)
165}
166
167fn cached_in_memory_connection(cache_key: &str, path: &Path) -> Result<&'static Mutex<Connection>> {
174 static CACHE: OnceLock<Mutex<HashMap<String, &'static Mutex<Connection>>>> = OnceLock::new();
175 let cache = CACHE.get_or_init(|| Mutex::new(HashMap::new()));
176
177 if let Some(conn) = cache.lock().unwrap().get(cache_key) {
178 return Ok(conn);
179 }
180
181 let disk_conn = open_store(path)?;
182 let mut mem_conn =
183 Connection::open_in_memory().context("failed to open an in-memory SQLite connection")?;
184 Backup::new(&disk_conn, &mut mem_conn)
185 .context("failed to start SQLite backup into memory")?
186 .run_to_completion(i32::MAX, Duration::from_millis(0), None)
187 .with_context(|| format!("failed to back up '{}' into memory", path.display()))?;
188 drop(disk_conn);
189
190 let leaked: &'static Mutex<Connection> = Box::leak(Box::new(Mutex::new(mem_conn)));
191 let mut cache = cache.lock().unwrap();
192 Ok(*cache.entry(cache_key.to_string()).or_insert(leaked))
193}
194
195#[derive(Debug, Clone, Serialize)]
196pub struct EndpointRecord {
197 pub operation_id: String,
198 pub path: String,
199 pub method: String,
200 pub summary: Option<String>,
201 pub description: Option<String>,
202 pub input_schema: serde_json::Value,
203 pub output_schema: serde_json::Value,
204 pub auth_scheme_ref: Option<String>,
205}
206
207fn row_to_endpoint(row: &Row) -> rusqlite::Result<EndpointRecord> {
208 let input_schema: String = row.get(5)?;
209 let output_schema: String = row.get(6)?;
210 Ok(EndpointRecord {
211 operation_id: row.get(0)?,
212 path: row.get(1)?,
213 method: row.get(2)?,
214 summary: row.get(3)?,
215 description: row.get(4)?,
216 input_schema: serde_json::from_str(&input_schema).unwrap_or(serde_json::Value::Null),
217 output_schema: serde_json::from_str(&output_schema).unwrap_or(serde_json::Value::Null),
218 auth_scheme_ref: row.get(7)?,
219 })
220}
221
222pub fn get_endpoint(conn: &Connection, operation_id: &str) -> Result<Option<EndpointRecord>> {
223 let mut stmt = conn.prepare(&format!(
224 "SELECT {ENDPOINT_COLUMNS} FROM endpoints WHERE operation_id = ?1"
225 ))?;
226 match stmt.query_row([operation_id], row_to_endpoint) {
227 Ok(record) => Ok(Some(record)),
228 Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
229 Err(err) => Err(err.into()),
230 }
231}
232
233pub fn list_endpoints(conn: &Connection) -> Result<Vec<EndpointRecord>> {
234 let mut stmt = conn.prepare(&format!("SELECT {ENDPOINT_COLUMNS} FROM endpoints"))?;
235 let rows = stmt.query_map([], row_to_endpoint)?;
236 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
237}
238
239#[derive(Debug, Clone, Serialize)]
240pub struct SearchResult {
241 pub operation_id: String,
242 pub summary: Option<String>,
243 pub similarity: f64,
244}
245
246pub fn search_endpoints(
261 conn: &Connection,
262 query_embedding: &[f32],
263 limit: usize,
264) -> Result<Vec<SearchResult>> {
265 let blob: Vec<u8> = query_embedding
266 .iter()
267 .flat_map(|value| value.to_le_bytes())
268 .collect();
269
270 let mut stmt = conn.prepare(
271 "SELECT e.operation_id, e.summary, s.distance
272 FROM semantic_endpoints s
273 JOIN endpoints e ON e.operation_id = s.operation_id
274 WHERE s.embedding MATCH ?1 AND k = ?2
275 ORDER BY s.distance",
276 )?;
277 let rows = stmt.query_map(rusqlite::params![blob, limit], |row| {
278 let distance: f64 = row.get(2)?;
279 Ok(SearchResult {
280 operation_id: row.get(0)?,
281 summary: row.get(1)?,
282 similarity: 1.0 - distance,
283 })
284 })?;
285 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
286}
287
288#[cfg(test)]
289mod tests {
290 use super::*;
291
292 #[test]
298 fn every_version_store_file_has_embedded_bytes() {
299 let file_labels: std::collections::HashSet<_> = VERSION_STORE_FILES
300 .iter()
301 .map(|(label, _)| *label)
302 .collect();
303 let byte_labels: std::collections::HashSet<_> = VERSION_STORE_BYTES
304 .iter()
305 .map(|(label, _)| *label)
306 .collect();
307 assert_eq!(file_labels, byte_labels);
308 }
309
310 #[test]
318 fn resolve_store_path_survives_concurrent_calls() {
319 let api_version = VERSION_STORE_FILES[0].0.to_string();
320 let handles: Vec<_> = (0..16)
321 .map(|_| {
322 let api_version = api_version.clone();
323 std::thread::spawn(move || {
324 let path = resolve_store_path(&api_version).unwrap();
325 open_store(&path).unwrap();
326 })
327 })
328 .collect();
329 for handle in handles {
330 handle.join().unwrap();
331 }
332 }
333
334 fn seeded_store(path: &Path) -> Connection {
340 unsafe {
341 #[allow(clippy::missing_transmute_annotations)]
342 rusqlite::ffi::sqlite3_auto_extension(Some(std::mem::transmute(
343 sqlite_vec::sqlite3_vec_init as *const (),
344 )));
345 }
346 let conn = Connection::open(path).unwrap();
347 conn.execute(
348 "CREATE TABLE endpoints (
349 operation_id TEXT PRIMARY KEY,
350 path TEXT NOT NULL,
351 method TEXT NOT NULL,
352 summary TEXT,
353 description TEXT,
354 input_schema TEXT NOT NULL,
355 output_schema TEXT NOT NULL,
356 auth_scheme_ref TEXT
357 )",
358 [],
359 )
360 .unwrap();
361 conn.execute(
362 "CREATE VIRTUAL TABLE semantic_endpoints USING vec0(
363 operation_id TEXT PRIMARY KEY,
364 embedding FLOAT[4]
365 )",
366 [],
367 )
368 .unwrap();
369 conn.execute(
370 "INSERT INTO endpoints (operation_id, path, method, summary, description, input_schema, output_schema, auth_scheme_ref)
371 VALUES ('listWidgets', '/widgets', 'GET', 'List widgets', NULL, '{}', '[]', NULL)",
372 [],
373 )
374 .unwrap();
375 let embedding: Vec<u8> = [1.0f32, 0.0, 0.0, 0.0]
376 .iter()
377 .flat_map(|v| v.to_le_bytes())
378 .collect();
379 conn.execute(
380 "INSERT INTO semantic_endpoints (operation_id, embedding) VALUES ('listWidgets', ?1)",
381 rusqlite::params![embedding],
382 )
383 .unwrap();
384 conn
385 }
386
387 #[test]
388 fn get_endpoint_returns_a_seeded_row() {
389 let dir = tempfile::tempdir().unwrap();
390 let path = dir.path().join("mcp_store.db");
391 let _conn = seeded_store(&path);
392
393 let store = open_store(&path).unwrap();
394 let endpoint = get_endpoint(&store, "listWidgets").unwrap().unwrap();
395 assert_eq!(endpoint.path, "/widgets");
396 assert_eq!(endpoint.method, "GET");
397 assert_eq!(endpoint.summary.as_deref(), Some("List widgets"));
398 }
399
400 #[test]
401 fn get_endpoint_returns_none_for_an_unknown_operation() {
402 let dir = tempfile::tempdir().unwrap();
403 let path = dir.path().join("mcp_store.db");
404 let _conn = seeded_store(&path);
405
406 let store = open_store(&path).unwrap();
407 assert!(get_endpoint(&store, "unknownOp").unwrap().is_none());
408 }
409
410 #[test]
411 fn list_endpoints_returns_every_row() {
412 let dir = tempfile::tempdir().unwrap();
413 let path = dir.path().join("mcp_store.db");
414 let _conn = seeded_store(&path);
415
416 let store = open_store(&path).unwrap();
417 let endpoints = list_endpoints(&store).unwrap();
418 assert_eq!(endpoints.len(), 1);
419 assert_eq!(endpoints[0].operation_id, "listWidgets");
420 }
421
422 #[test]
423 fn search_endpoints_finds_the_nearest_neighbor() {
424 let dir = tempfile::tempdir().unwrap();
425 let path = dir.path().join("mcp_store.db");
426 let _conn = seeded_store(&path);
427
428 let store = open_store(&path).unwrap();
429 let results = search_endpoints(&store, &[1.0, 0.0, 0.0, 0.0], 5).unwrap();
430 assert_eq!(results.len(), 1);
431 assert_eq!(results[0].operation_id, "listWidgets");
432 assert!(results[0].similarity > 0.99);
433 }
434
435 #[test]
436 fn cached_in_memory_connection_serves_seeded_data() {
437 let dir = tempfile::tempdir().unwrap();
438 let path = dir.path().join("mcp_store.db");
439 let _conn = seeded_store(&path);
440
441 let cached = cached_in_memory_connection("cached-serves-seeded-data", &path).unwrap();
442 let endpoint = get_endpoint(&cached.lock().unwrap(), "listWidgets")
443 .unwrap()
444 .unwrap();
445 assert_eq!(endpoint.path, "/widgets");
446
447 let results = search_endpoints(&cached.lock().unwrap(), &[1.0, 0.0, 0.0, 0.0], 5).unwrap();
451 assert_eq!(results.len(), 1);
452 assert_eq!(results[0].operation_id, "listWidgets");
453 }
454
455 #[test]
456 fn cached_in_memory_connection_holds_no_lingering_lock_on_the_disk_file() {
457 let dir = tempfile::tempdir().unwrap();
458 let path = dir.path().join("mcp_store.db");
459 let _conn = seeded_store(&path);
460
461 let cached = cached_in_memory_connection("cached-releases-disk-handle", &path).unwrap();
462
463 std::fs::remove_file(&path).unwrap();
467
468 let endpoint = get_endpoint(&cached.lock().unwrap(), "listWidgets")
469 .unwrap()
470 .unwrap();
471 assert_eq!(endpoint.path, "/widgets");
472 }
473
474 #[test]
475 fn cached_in_memory_connection_reuses_the_same_connection_for_the_same_key() {
476 let dir = tempfile::tempdir().unwrap();
477 let path = dir.path().join("mcp_store.db");
478 let _conn = seeded_store(&path);
479
480 let first = cached_in_memory_connection("cached-reuses-same-key", &path).unwrap()
481 as *const Mutex<Connection>;
482 let second = cached_in_memory_connection("cached-reuses-same-key", &path).unwrap()
483 as *const Mutex<Connection>;
484 assert_eq!(
485 first, second,
486 "expected the same cached connection, not a fresh backup"
487 );
488 }
489}