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