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 ("gh-2026-03-10", "mcp_store.db"),
28 ("ghec-2026-03-10", "mcp_store_vghec-2026-03-10.db"),
29 ("ghes-3.21", "mcp_store_vghes-3.21.db"),
30 ("ghes-3.20", "mcp_store_vghes-3.20.db"),
31 ("ghes-3.19", "mcp_store_vghes-3.19.db"),
32];
33
34const VERSION_STORE_BYTES: &[(&str, &[u8])] = &[
35 ("gh-2026-03-10", include_bytes!("../../mcp_store.db.zst")),
36 (
37 "ghec-2026-03-10",
38 include_bytes!("../../mcp_store_vghec-2026-03-10.db.zst"),
39 ),
40 (
41 "ghes-3.21",
42 include_bytes!("../../mcp_store_vghes-3.21.db.zst"),
43 ),
44 (
45 "ghes-3.20",
46 include_bytes!("../../mcp_store_vghes-3.20.db.zst"),
47 ),
48 (
49 "ghes-3.19",
50 include_bytes!("../../mcp_store_vghes-3.19.db.zst"),
51 ),
52];
53pub fn resolve_store_path(api_version: &str) -> Result<PathBuf> {
72 static EXTRACTED: OnceLock<Mutex<HashMap<String, PathBuf>>> = OnceLock::new();
84 let extracted = EXTRACTED.get_or_init(|| Mutex::new(HashMap::new()));
85 let mut extracted = extracted.lock().unwrap();
86 if let Some(path) = extracted.get(api_version) {
87 return Ok(path.clone());
88 }
89
90 let file = VERSION_STORE_FILES
91 .iter()
92 .find(|(label, _)| *label == api_version)
93 .map(|(_, file)| *file)
94 .with_context(|| format!("unknown api_version '{api_version}' — run the 'versions' command to see what's available"))?;
95 let bytes = VERSION_STORE_BYTES
96 .iter()
97 .find(|(label, _)| *label == api_version)
98 .map(|(_, bytes)| *bytes)
99 .with_context(|| format!("no embedded store data for api_version '{api_version}'"))?;
100
101 let mut dir = std::env::temp_dir();
102 dir.push(concat!(env!("CARGO_PKG_NAME"), "-store"));
103 std::fs::create_dir_all(&dir)
104 .with_context(|| format!("failed to create temp dir '{}'", dir.display()))?;
105
106 let path = dir.join(file);
107 let decompressed = zstd::stream::decode_all(bytes).with_context(|| {
124 format!("failed to decompress embedded store data for api_version '{api_version}'")
125 })?;
126 static UNIQUE: AtomicU64 = AtomicU64::new(0);
127 let tmp_path = dir.join(format!(
128 "{file}.{}.{}.tmp",
129 std::process::id(),
130 UNIQUE.fetch_add(1, Ordering::Relaxed)
131 ));
132 std::fs::write(&tmp_path, decompressed).with_context(|| {
133 format!(
134 "failed to extract embedded store data to '{}'",
135 tmp_path.display()
136 )
137 })?;
138 std::fs::rename(&tmp_path, &path).with_context(|| {
139 format!(
140 "failed to move extracted store data into place at '{}'",
141 path.display()
142 )
143 })?;
144
145 extracted.insert(api_version.to_string(), path.clone());
146 Ok(path)
147}
148
149const ENDPOINT_COLUMNS: &str = "operation_id, path, method, summary, description, input_schema, output_schema, auth_scheme_ref";
150
151fn register_vec_extension() {
155 REGISTER_VEC_EXTENSION.call_once(|| unsafe {
156 #[allow(clippy::missing_transmute_annotations)]
157 rusqlite::ffi::sqlite3_auto_extension(Some(std::mem::transmute(
158 sqlite_vec::sqlite3_vec_init as *const (),
159 )));
160 });
161}
162
163pub fn open_store(path: &Path) -> Result<Connection> {
167 register_vec_extension();
168 Connection::open_with_flags(path, OpenFlags::SQLITE_OPEN_READ_ONLY)
169 .with_context(|| format!("failed to open '{}'", path.display()))
170}
171
172pub fn open_store_read_write(path: &Path) -> Result<Connection> {
177 register_vec_extension();
178 Connection::open_with_flags(path, OpenFlags::SQLITE_OPEN_READ_WRITE)
179 .with_context(|| format!("failed to open '{}'", path.display()))
180}
181
182pub fn cached_store_connection(api_version: &str) -> Result<&'static Mutex<Connection>> {
200 let path = resolve_store_path(api_version)?;
201 cached_in_memory_connection(api_version, &path)
202}
203
204fn cached_in_memory_connection(cache_key: &str, path: &Path) -> Result<&'static Mutex<Connection>> {
211 static CACHE: OnceLock<Mutex<HashMap<String, &'static Mutex<Connection>>>> = OnceLock::new();
212 let cache = CACHE.get_or_init(|| Mutex::new(HashMap::new()));
213
214 if let Some(conn) = cache.lock().unwrap().get(cache_key) {
215 return Ok(conn);
216 }
217
218 let disk_conn = open_store(path)?;
219 let mut mem_conn =
220 Connection::open_in_memory().context("failed to open an in-memory SQLite connection")?;
221 Backup::new(&disk_conn, &mut mem_conn)
222 .context("failed to start SQLite backup into memory")?
223 .run_to_completion(i32::MAX, Duration::from_millis(0), None)
224 .with_context(|| format!("failed to back up '{}' into memory", path.display()))?;
225 drop(disk_conn);
226
227 let leaked: &'static Mutex<Connection> = Box::leak(Box::new(Mutex::new(mem_conn)));
228 let mut cache = cache.lock().unwrap();
229 Ok(*cache.entry(cache_key.to_string()).or_insert(leaked))
230}
231
232#[derive(Debug, Clone, Serialize)]
233pub struct EndpointRecord {
234 pub operation_id: String,
235 pub path: String,
236 pub method: String,
237 pub summary: Option<String>,
238 pub description: Option<String>,
239 pub input_schema: serde_json::Value,
240 pub output_schema: serde_json::Value,
241 pub auth_scheme_ref: Option<String>,
242}
243
244fn row_to_endpoint(row: &Row) -> rusqlite::Result<EndpointRecord> {
245 let input_schema: String = row.get(5)?;
246 let output_schema: String = row.get(6)?;
247 Ok(EndpointRecord {
248 operation_id: row.get(0)?,
249 path: row.get(1)?,
250 method: row.get(2)?,
251 summary: row.get(3)?,
252 description: row.get(4)?,
253 input_schema: serde_json::from_str(&input_schema).unwrap_or(serde_json::Value::Null),
254 output_schema: serde_json::from_str(&output_schema).unwrap_or(serde_json::Value::Null),
255 auth_scheme_ref: row.get(7)?,
256 })
257}
258
259pub fn get_endpoint(conn: &Connection, operation_id: &str) -> Result<Option<EndpointRecord>> {
260 let mut stmt = conn.prepare(&format!(
261 "SELECT {ENDPOINT_COLUMNS} FROM endpoints WHERE operation_id = ?1"
262 ))?;
263 match stmt.query_row([operation_id], row_to_endpoint) {
264 Ok(record) => Ok(Some(record)),
265 Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
266 Err(err) => Err(err.into()),
267 }
268}
269
270pub fn list_endpoints(conn: &Connection) -> Result<Vec<EndpointRecord>> {
271 let mut stmt = conn.prepare(&format!("SELECT {ENDPOINT_COLUMNS} FROM endpoints"))?;
272 let rows = stmt.query_map([], row_to_endpoint)?;
273 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
274}
275
276#[derive(Debug, Clone, Serialize)]
277pub struct SearchResult {
278 pub operation_id: String,
279 pub summary: Option<String>,
280 pub similarity: f64,
281}
282
283pub fn search_endpoints(
298 conn: &Connection,
299 query_embedding: &[f32],
300 limit: usize,
301) -> Result<Vec<SearchResult>> {
302 let blob: Vec<u8> = query_embedding
303 .iter()
304 .flat_map(|value| value.to_le_bytes())
305 .collect();
306
307 let mut stmt = conn.prepare(
308 "SELECT e.operation_id, e.summary, s.distance
309 FROM semantic_endpoints s
310 JOIN endpoints e ON e.operation_id = s.operation_id
311 WHERE s.embedding MATCH ?1 AND k = ?2
312 ORDER BY s.distance",
313 )?;
314 let rows = stmt.query_map(rusqlite::params![blob, limit], |row| {
315 let distance: f64 = row.get(2)?;
316 Ok(SearchResult {
317 operation_id: row.get(0)?,
318 summary: row.get(1)?,
319 similarity: 1.0 - distance,
320 })
321 })?;
322 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
323}
324
325#[cfg(test)]
326mod tests {
327 use super::*;
328
329 #[test]
335 fn every_version_store_file_has_embedded_bytes() {
336 let file_labels: std::collections::HashSet<_> = VERSION_STORE_FILES
337 .iter()
338 .map(|(label, _)| *label)
339 .collect();
340 let byte_labels: std::collections::HashSet<_> = VERSION_STORE_BYTES
341 .iter()
342 .map(|(label, _)| *label)
343 .collect();
344 assert_eq!(file_labels, byte_labels);
345 }
346
347 #[test]
355 fn resolve_store_path_survives_concurrent_calls() {
356 let api_version = VERSION_STORE_FILES[0].0.to_string();
357 let handles: Vec<_> = (0..16)
358 .map(|_| {
359 let api_version = api_version.clone();
360 std::thread::spawn(move || {
361 let path = resolve_store_path(&api_version).unwrap();
362 open_store(&path).unwrap();
363 })
364 })
365 .collect();
366 for handle in handles {
367 handle.join().unwrap();
368 }
369 }
370
371 fn seeded_store(path: &Path) -> Connection {
377 unsafe {
378 #[allow(clippy::missing_transmute_annotations)]
379 rusqlite::ffi::sqlite3_auto_extension(Some(std::mem::transmute(
380 sqlite_vec::sqlite3_vec_init as *const (),
381 )));
382 }
383 let conn = Connection::open(path).unwrap();
384 conn.execute(
385 "CREATE TABLE endpoints (
386 operation_id TEXT PRIMARY KEY,
387 path TEXT NOT NULL,
388 method TEXT NOT NULL,
389 summary TEXT,
390 description TEXT,
391 input_schema TEXT NOT NULL,
392 output_schema TEXT NOT NULL,
393 auth_scheme_ref TEXT
394 )",
395 [],
396 )
397 .unwrap();
398 conn.execute(
399 "CREATE VIRTUAL TABLE semantic_endpoints USING vec0(
400 operation_id TEXT PRIMARY KEY,
401 embedding FLOAT[4]
402 )",
403 [],
404 )
405 .unwrap();
406 conn.execute(
407 "INSERT INTO endpoints (operation_id, path, method, summary, description, input_schema, output_schema, auth_scheme_ref)
408 VALUES ('listWidgets', '/widgets', 'GET', 'List widgets', NULL, '{}', '[]', NULL)",
409 [],
410 )
411 .unwrap();
412 let embedding: Vec<u8> = [1.0f32, 0.0, 0.0, 0.0]
413 .iter()
414 .flat_map(|v| v.to_le_bytes())
415 .collect();
416 conn.execute(
417 "INSERT INTO semantic_endpoints (operation_id, embedding) VALUES ('listWidgets', ?1)",
418 rusqlite::params![embedding],
419 )
420 .unwrap();
421 conn
422 }
423
424 #[test]
425 fn get_endpoint_returns_a_seeded_row() {
426 let dir = tempfile::tempdir().unwrap();
427 let path = dir.path().join("mcp_store.db");
428 let _conn = seeded_store(&path);
429
430 let store = open_store(&path).unwrap();
431 let endpoint = get_endpoint(&store, "listWidgets").unwrap().unwrap();
432 assert_eq!(endpoint.path, "/widgets");
433 assert_eq!(endpoint.method, "GET");
434 assert_eq!(endpoint.summary.as_deref(), Some("List widgets"));
435 }
436
437 #[test]
438 fn get_endpoint_returns_none_for_an_unknown_operation() {
439 let dir = tempfile::tempdir().unwrap();
440 let path = dir.path().join("mcp_store.db");
441 let _conn = seeded_store(&path);
442
443 let store = open_store(&path).unwrap();
444 assert!(get_endpoint(&store, "unknownOp").unwrap().is_none());
445 }
446
447 #[test]
448 fn list_endpoints_returns_every_row() {
449 let dir = tempfile::tempdir().unwrap();
450 let path = dir.path().join("mcp_store.db");
451 let _conn = seeded_store(&path);
452
453 let store = open_store(&path).unwrap();
454 let endpoints = list_endpoints(&store).unwrap();
455 assert_eq!(endpoints.len(), 1);
456 assert_eq!(endpoints[0].operation_id, "listWidgets");
457 }
458
459 #[test]
460 fn search_endpoints_finds_the_nearest_neighbor() {
461 let dir = tempfile::tempdir().unwrap();
462 let path = dir.path().join("mcp_store.db");
463 let _conn = seeded_store(&path);
464
465 let store = open_store(&path).unwrap();
466 let results = search_endpoints(&store, &[1.0, 0.0, 0.0, 0.0], 5).unwrap();
467 assert_eq!(results.len(), 1);
468 assert_eq!(results[0].operation_id, "listWidgets");
469 assert!(results[0].similarity > 0.99);
470 }
471
472 #[test]
473 fn cached_in_memory_connection_serves_seeded_data() {
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-serves-seeded-data", &path).unwrap();
479 let endpoint = get_endpoint(&cached.lock().unwrap(), "listWidgets")
480 .unwrap()
481 .unwrap();
482 assert_eq!(endpoint.path, "/widgets");
483
484 let results = search_endpoints(&cached.lock().unwrap(), &[1.0, 0.0, 0.0, 0.0], 5).unwrap();
488 assert_eq!(results.len(), 1);
489 assert_eq!(results[0].operation_id, "listWidgets");
490 }
491
492 #[test]
493 fn cached_in_memory_connection_holds_no_lingering_lock_on_the_disk_file() {
494 let dir = tempfile::tempdir().unwrap();
495 let path = dir.path().join("mcp_store.db");
496 drop(seeded_store(&path));
502
503 let cached = cached_in_memory_connection("cached-releases-disk-handle", &path).unwrap();
504
505 std::fs::remove_file(&path).unwrap();
509
510 let endpoint = get_endpoint(&cached.lock().unwrap(), "listWidgets")
511 .unwrap()
512 .unwrap();
513 assert_eq!(endpoint.path, "/widgets");
514 }
515
516 #[test]
517 fn cached_in_memory_connection_reuses_the_same_connection_for_the_same_key() {
518 let dir = tempfile::tempdir().unwrap();
519 let path = dir.path().join("mcp_store.db");
520 let _conn = seeded_store(&path);
521
522 let first = cached_in_memory_connection("cached-reuses-same-key", &path).unwrap()
523 as *const Mutex<Connection>;
524 let second = cached_in_memory_connection("cached-reuses-same-key", &path).unwrap()
525 as *const Mutex<Connection>;
526 assert_eq!(
527 first, second,
528 "expected the same cached connection, not a fresh backup"
529 );
530 }
531}