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];
32
33const VERSION_STORE_BYTES: &[(&str, &[u8])] = &[
34 ("gh-2026-03-10", include_bytes!("../../mcp_store.db.zst")),
35 (
36 "ghec-2026-03-10",
37 include_bytes!("../../mcp_store_vghec-2026-03-10.db.zst"),
38 ),
39 (
40 "ghes-3.21",
41 include_bytes!("../../mcp_store_vghes-3.21.db.zst"),
42 ),
43 (
44 "ghes-3.20",
45 include_bytes!("../../mcp_store_vghes-3.20.db.zst"),
46 ),
47];
48pub fn resolve_store_path(api_version: &str) -> Result<PathBuf> {
67 static EXTRACTED: OnceLock<Mutex<HashMap<String, PathBuf>>> = OnceLock::new();
79 let extracted = EXTRACTED.get_or_init(|| Mutex::new(HashMap::new()));
80 let mut extracted = extracted.lock().unwrap();
81 if let Some(path) = extracted.get(api_version) {
82 return Ok(path.clone());
83 }
84
85 let file = VERSION_STORE_FILES
86 .iter()
87 .find(|(label, _)| *label == api_version)
88 .map(|(_, file)| *file)
89 .with_context(|| format!("unknown api_version '{api_version}' — run the 'versions' command to see what's available"))?;
90 let bytes = VERSION_STORE_BYTES
91 .iter()
92 .find(|(label, _)| *label == api_version)
93 .map(|(_, bytes)| *bytes)
94 .with_context(|| format!("no embedded store data for api_version '{api_version}'"))?;
95
96 let mut dir = std::env::temp_dir();
97 dir.push(concat!(env!("CARGO_PKG_NAME"), "-store"));
98 std::fs::create_dir_all(&dir)
99 .with_context(|| format!("failed to create temp dir '{}'", dir.display()))?;
100
101 let path = dir.join(file);
102 let decompressed = zstd::stream::decode_all(bytes).with_context(|| {
120 format!("failed to decompress embedded store data for api_version '{api_version}'")
121 })?;
122 static UNIQUE: AtomicU64 = AtomicU64::new(0);
123 let tmp_path = dir.join(format!(
124 "{file}.{}.{}.tmp",
125 std::process::id(),
126 UNIQUE.fetch_add(1, Ordering::Relaxed)
127 ));
128 std::fs::write(&tmp_path, decompressed).with_context(|| {
129 format!(
130 "failed to extract embedded store data to '{}'",
131 tmp_path.display()
132 )
133 })?;
134 let mut rename_attempts = 0u32;
144 loop {
145 match std::fs::rename(&tmp_path, &path) {
146 Ok(()) => break,
147 Err(_) if rename_attempts < 5 => {
148 rename_attempts += 1;
149 std::thread::sleep(Duration::from_millis(50 * u64::from(rename_attempts)));
150 }
151 Err(err) => {
152 return Err(err).with_context(|| {
153 format!(
154 "failed to move extracted store data into place at '{}'",
155 path.display()
156 )
157 });
158 }
159 }
160 }
161
162 extracted.insert(api_version.to_string(), path.clone());
163 Ok(path)
164}
165
166const ENDPOINT_COLUMNS: &str = "operation_id, path, method, summary, description, input_schema, output_schema, auth_scheme_ref";
167
168fn register_vec_extension() {
172 REGISTER_VEC_EXTENSION.call_once(|| unsafe {
173 #[allow(clippy::missing_transmute_annotations)]
174 rusqlite::ffi::sqlite3_auto_extension(Some(std::mem::transmute(
175 sqlite_vec::sqlite3_vec_init as *const (),
176 )));
177 });
178}
179
180pub fn open_store(path: &Path) -> Result<Connection> {
184 register_vec_extension();
185 Connection::open_with_flags(path, OpenFlags::SQLITE_OPEN_READ_ONLY)
186 .with_context(|| format!("failed to open '{}'", path.display()))
187}
188
189pub fn open_store_read_write(path: &Path) -> Result<Connection> {
194 register_vec_extension();
195 Connection::open_with_flags(path, OpenFlags::SQLITE_OPEN_READ_WRITE)
196 .with_context(|| format!("failed to open '{}'", path.display()))
197}
198
199pub fn cached_store_connection(api_version: &str) -> Result<&'static Mutex<Connection>> {
217 let path = resolve_store_path(api_version)?;
218 cached_in_memory_connection(api_version, &path)
219}
220
221fn cached_in_memory_connection(cache_key: &str, path: &Path) -> Result<&'static Mutex<Connection>> {
228 static CACHE: OnceLock<Mutex<HashMap<String, &'static Mutex<Connection>>>> = OnceLock::new();
229 let cache = CACHE.get_or_init(|| Mutex::new(HashMap::new()));
230
231 if let Some(conn) = cache.lock().unwrap().get(cache_key) {
232 return Ok(conn);
233 }
234
235 let disk_conn = open_store(path)?;
236 let mut mem_conn =
237 Connection::open_in_memory().context("failed to open an in-memory SQLite connection")?;
238 Backup::new(&disk_conn, &mut mem_conn)
239 .context("failed to start SQLite backup into memory")?
240 .run_to_completion(i32::MAX, Duration::from_millis(0), None)
241 .with_context(|| format!("failed to back up '{}' into memory", path.display()))?;
242 drop(disk_conn);
243
244 let leaked: &'static Mutex<Connection> = Box::leak(Box::new(Mutex::new(mem_conn)));
245 let mut cache = cache.lock().unwrap();
246 Ok(*cache.entry(cache_key.to_string()).or_insert(leaked))
247}
248
249#[derive(Debug, Clone, Serialize)]
250pub struct EndpointRecord {
251 pub operation_id: String,
252 pub path: String,
253 pub method: String,
254 pub summary: Option<String>,
255 pub description: Option<String>,
256 pub input_schema: serde_json::Value,
257 pub output_schema: serde_json::Value,
258 pub auth_scheme_ref: Option<String>,
259}
260
261fn row_to_endpoint(row: &Row) -> rusqlite::Result<EndpointRecord> {
262 let input_schema: String = row.get(5)?;
263 let output_schema: String = row.get(6)?;
264 Ok(EndpointRecord {
265 operation_id: row.get(0)?,
266 path: row.get(1)?,
267 method: row.get(2)?,
268 summary: row.get(3)?,
269 description: row.get(4)?,
270 input_schema: serde_json::from_str(&input_schema).unwrap_or(serde_json::Value::Null),
271 output_schema: serde_json::from_str(&output_schema).unwrap_or(serde_json::Value::Null),
272 auth_scheme_ref: row.get(7)?,
273 })
274}
275
276pub fn get_endpoint(conn: &Connection, operation_id: &str) -> Result<Option<EndpointRecord>> {
277 let mut stmt = conn.prepare(&format!(
278 "SELECT {ENDPOINT_COLUMNS} FROM endpoints WHERE operation_id = ?1"
279 ))?;
280 match stmt.query_row([operation_id], row_to_endpoint) {
281 Ok(record) => Ok(Some(record)),
282 Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
283 Err(err) => Err(err.into()),
284 }
285}
286
287pub fn list_endpoints(conn: &Connection) -> Result<Vec<EndpointRecord>> {
288 let mut stmt = conn.prepare(&format!("SELECT {ENDPOINT_COLUMNS} FROM endpoints"))?;
289 let rows = stmt.query_map([], row_to_endpoint)?;
290 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
291}
292
293#[derive(Debug, Clone, Serialize)]
294pub struct SearchResult {
295 pub operation_id: String,
296 pub summary: Option<String>,
297 pub similarity: f64,
298}
299
300pub fn search_endpoints(
315 conn: &Connection,
316 query_embedding: &[f32],
317 limit: usize,
318) -> Result<Vec<SearchResult>> {
319 let blob: Vec<u8> = query_embedding
320 .iter()
321 .flat_map(|value| value.to_le_bytes())
322 .collect();
323
324 let mut stmt = conn.prepare(
325 "SELECT e.operation_id, e.summary, s.distance
326 FROM semantic_endpoints s
327 JOIN endpoints e ON e.operation_id = s.operation_id
328 WHERE s.embedding MATCH ?1 AND k = ?2
329 ORDER BY s.distance",
330 )?;
331 let rows = stmt.query_map(rusqlite::params![blob, limit], |row| {
332 let distance: f64 = row.get(2)?;
333 Ok(SearchResult {
334 operation_id: row.get(0)?,
335 summary: row.get(1)?,
336 similarity: 1.0 - distance,
337 })
338 })?;
339 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
340}
341
342#[cfg(test)]
343mod tests {
344 use super::*;
345
346 #[test]
352 fn every_version_store_file_has_embedded_bytes() {
353 let file_labels: std::collections::HashSet<_> = VERSION_STORE_FILES
354 .iter()
355 .map(|(label, _)| *label)
356 .collect();
357 let byte_labels: std::collections::HashSet<_> = VERSION_STORE_BYTES
358 .iter()
359 .map(|(label, _)| *label)
360 .collect();
361 assert_eq!(file_labels, byte_labels);
362 }
363
364 #[test]
372 fn resolve_store_path_survives_concurrent_calls() {
373 let api_version = VERSION_STORE_FILES[0].0.to_string();
374 let handles: Vec<_> = (0..16)
375 .map(|_| {
376 let api_version = api_version.clone();
377 std::thread::spawn(move || {
378 let path = resolve_store_path(&api_version).unwrap();
379 open_store(&path).unwrap();
380 })
381 })
382 .collect();
383 for handle in handles {
384 handle.join().unwrap();
385 }
386 }
387
388 fn seeded_store(path: &Path) -> Connection {
394 unsafe {
395 #[allow(clippy::missing_transmute_annotations)]
396 rusqlite::ffi::sqlite3_auto_extension(Some(std::mem::transmute(
397 sqlite_vec::sqlite3_vec_init as *const (),
398 )));
399 }
400 let conn = Connection::open(path).unwrap();
401 conn.execute(
402 "CREATE TABLE endpoints (
403 operation_id TEXT PRIMARY KEY,
404 path TEXT NOT NULL,
405 method TEXT NOT NULL,
406 summary TEXT,
407 description TEXT,
408 input_schema TEXT NOT NULL,
409 output_schema TEXT NOT NULL,
410 auth_scheme_ref TEXT
411 )",
412 [],
413 )
414 .unwrap();
415 conn.execute(
416 "CREATE VIRTUAL TABLE semantic_endpoints USING vec0(
417 operation_id TEXT PRIMARY KEY,
418 embedding FLOAT[4]
419 )",
420 [],
421 )
422 .unwrap();
423 conn.execute(
424 "INSERT INTO endpoints (operation_id, path, method, summary, description, input_schema, output_schema, auth_scheme_ref)
425 VALUES ('listWidgets', '/widgets', 'GET', 'List widgets', NULL, '{}', '[]', NULL)",
426 [],
427 )
428 .unwrap();
429 let embedding: Vec<u8> = [1.0f32, 0.0, 0.0, 0.0]
430 .iter()
431 .flat_map(|v| v.to_le_bytes())
432 .collect();
433 conn.execute(
434 "INSERT INTO semantic_endpoints (operation_id, embedding) VALUES ('listWidgets', ?1)",
435 rusqlite::params![embedding],
436 )
437 .unwrap();
438 conn
439 }
440
441 #[test]
442 fn get_endpoint_returns_a_seeded_row() {
443 let dir = tempfile::tempdir().unwrap();
444 let path = dir.path().join("mcp_store.db");
445 let _conn = seeded_store(&path);
446
447 let store = open_store(&path).unwrap();
448 let endpoint = get_endpoint(&store, "listWidgets").unwrap().unwrap();
449 assert_eq!(endpoint.path, "/widgets");
450 assert_eq!(endpoint.method, "GET");
451 assert_eq!(endpoint.summary.as_deref(), Some("List widgets"));
452 }
453
454 #[test]
455 fn get_endpoint_returns_none_for_an_unknown_operation() {
456 let dir = tempfile::tempdir().unwrap();
457 let path = dir.path().join("mcp_store.db");
458 let _conn = seeded_store(&path);
459
460 let store = open_store(&path).unwrap();
461 assert!(get_endpoint(&store, "unknownOp").unwrap().is_none());
462 }
463
464 #[test]
465 fn list_endpoints_returns_every_row() {
466 let dir = tempfile::tempdir().unwrap();
467 let path = dir.path().join("mcp_store.db");
468 let _conn = seeded_store(&path);
469
470 let store = open_store(&path).unwrap();
471 let endpoints = list_endpoints(&store).unwrap();
472 assert_eq!(endpoints.len(), 1);
473 assert_eq!(endpoints[0].operation_id, "listWidgets");
474 }
475
476 #[test]
477 fn search_endpoints_finds_the_nearest_neighbor() {
478 let dir = tempfile::tempdir().unwrap();
479 let path = dir.path().join("mcp_store.db");
480 let _conn = seeded_store(&path);
481
482 let store = open_store(&path).unwrap();
483 let results = search_endpoints(&store, &[1.0, 0.0, 0.0, 0.0], 5).unwrap();
484 assert_eq!(results.len(), 1);
485 assert_eq!(results[0].operation_id, "listWidgets");
486 assert!(results[0].similarity > 0.99);
487 }
488
489 #[test]
490 fn cached_in_memory_connection_serves_seeded_data() {
491 let dir = tempfile::tempdir().unwrap();
492 let path = dir.path().join("mcp_store.db");
493 let _conn = seeded_store(&path);
494
495 let cached = cached_in_memory_connection("cached-serves-seeded-data", &path).unwrap();
496 let endpoint = get_endpoint(&cached.lock().unwrap(), "listWidgets")
497 .unwrap()
498 .unwrap();
499 assert_eq!(endpoint.path, "/widgets");
500
501 let results = search_endpoints(&cached.lock().unwrap(), &[1.0, 0.0, 0.0, 0.0], 5).unwrap();
505 assert_eq!(results.len(), 1);
506 assert_eq!(results[0].operation_id, "listWidgets");
507 }
508
509 #[test]
510 fn cached_in_memory_connection_holds_no_lingering_lock_on_the_disk_file() {
511 if cfg!(windows) && std::env::var_os("GITHUB_ACTIONS").is_some() {
522 return;
523 }
524
525 let dir = tempfile::tempdir().unwrap();
526 let path = dir.path().join("mcp_store.db");
527 drop(seeded_store(&path));
533
534 let cached = cached_in_memory_connection("cached-releases-disk-handle", &path).unwrap();
535
536 let mut remove_result = std::fs::remove_file(&path);
540 for _ in 0..10 {
541 if remove_result.is_ok() {
542 break;
543 }
544 std::thread::sleep(std::time::Duration::from_millis(50));
545 remove_result = std::fs::remove_file(&path);
546 }
547 remove_result.unwrap();
548
549 let endpoint = get_endpoint(&cached.lock().unwrap(), "listWidgets")
550 .unwrap()
551 .unwrap();
552 assert_eq!(endpoint.path, "/widgets");
553 }
554
555 #[test]
556 fn cached_in_memory_connection_reuses_the_same_connection_for_the_same_key() {
557 let dir = tempfile::tempdir().unwrap();
558 let path = dir.path().join("mcp_store.db");
559 let _conn = seeded_store(&path);
560
561 let first = cached_in_memory_connection("cached-reuses-same-key", &path).unwrap()
562 as *const Mutex<Connection>;
563 let second = cached_in_memory_connection("cached-reuses-same-key", &path).unwrap()
564 as *const Mutex<Connection>;
565 assert_eq!(
566 first, second,
567 "expected the same cached connection, not a fresh backup"
568 );
569 }
570}