1use std::collections::HashMap;
13use std::path::{Path, PathBuf};
14use std::sync::{Mutex, Once, OnceLock};
15use std::time::Duration;
16
17use anyhow::{Context, Result};
18use rusqlite::backup::Backup;
19use rusqlite::{Connection, OpenFlags, Row};
20use serde::Serialize;
21
22static REGISTER_VEC_EXTENSION: Once = Once::new();
23
24pub const VERSION_STORE_FILES: &[(&str, &str)] = &[
26 ("gh-2026-03-10", "mcp_store.db"),
27 ("ghec-2026-03-10", "mcp_store_vghec-2026-03-10.db"),
28 ("ghes-3.21", "mcp_store_vghes-3.21.db"),
29 ("ghes-3.20", "mcp_store_vghes-3.20.db"),
30 ("ghes-3.19", "mcp_store_vghes-3.19.db"),
31];
32
33const VERSION_STORE_BYTES: &[(&str, &[u8])] = &[
39 ("gh-2026-03-10", include_bytes!("../../mcp_store.db.zst")),
40 (
41 "ghec-2026-03-10",
42 include_bytes!("../../mcp_store_vghec-2026-03-10.db.zst"),
43 ),
44 (
45 "ghes-3.21",
46 include_bytes!("../../mcp_store_vghes-3.21.db.zst"),
47 ),
48 (
49 "ghes-3.20",
50 include_bytes!("../../mcp_store_vghes-3.20.db.zst"),
51 ),
52 (
53 "ghes-3.19",
54 include_bytes!("../../mcp_store_vghes-3.19.db.zst"),
55 ),
56];
57pub fn resolve_store_path(api_version: &str) -> Result<PathBuf> {
74 let file = VERSION_STORE_FILES
75 .iter()
76 .find(|(label, _)| *label == api_version)
77 .map(|(_, file)| *file)
78 .with_context(|| format!("unknown api_version '{api_version}' — run the 'versions' command to see what's available"))?;
79 let bytes = VERSION_STORE_BYTES
80 .iter()
81 .find(|(label, _)| *label == api_version)
82 .map(|(_, bytes)| *bytes)
83 .with_context(|| format!("no embedded store data for api_version '{api_version}'"))?;
84
85 let mut dir = std::env::temp_dir();
86 dir.push(concat!(env!("CARGO_PKG_NAME"), "-store"));
87 std::fs::create_dir_all(&dir)
88 .with_context(|| format!("failed to create temp dir '{}'", dir.display()))?;
89
90 let path = dir.join(file);
91 let decompressed = zstd::stream::decode_all(bytes).with_context(|| {
101 format!("failed to decompress embedded store data for api_version '{api_version}'")
102 })?;
103 std::fs::write(&path, decompressed).with_context(|| {
104 format!(
105 "failed to extract embedded store data to '{}'",
106 path.display()
107 )
108 })?;
109
110 Ok(path)
111}
112
113const ENDPOINT_COLUMNS: &str = "operation_id, path, method, summary, description, input_schema, output_schema, auth_scheme_ref";
114
115fn register_vec_extension() {
119 REGISTER_VEC_EXTENSION.call_once(|| unsafe {
120 #[allow(clippy::missing_transmute_annotations)]
121 rusqlite::ffi::sqlite3_auto_extension(Some(std::mem::transmute(
122 sqlite_vec::sqlite3_vec_init as *const (),
123 )));
124 });
125}
126
127pub fn open_store(path: &Path) -> Result<Connection> {
131 register_vec_extension();
132 Connection::open_with_flags(path, OpenFlags::SQLITE_OPEN_READ_ONLY)
133 .with_context(|| format!("failed to open '{}'", path.display()))
134}
135
136pub fn open_store_read_write(path: &Path) -> Result<Connection> {
141 register_vec_extension();
142 Connection::open_with_flags(path, OpenFlags::SQLITE_OPEN_READ_WRITE)
143 .with_context(|| format!("failed to open '{}'", path.display()))
144}
145
146pub fn cached_store_connection(api_version: &str) -> Result<&'static Mutex<Connection>> {
164 let path = resolve_store_path(api_version)?;
165 cached_in_memory_connection(api_version, &path)
166}
167
168fn cached_in_memory_connection(cache_key: &str, path: &Path) -> Result<&'static Mutex<Connection>> {
175 static CACHE: OnceLock<Mutex<HashMap<String, &'static Mutex<Connection>>>> = OnceLock::new();
176 let cache = CACHE.get_or_init(|| Mutex::new(HashMap::new()));
177
178 if let Some(conn) = cache.lock().unwrap().get(cache_key) {
179 return Ok(conn);
180 }
181
182 let disk_conn = open_store(path)?;
183 let mut mem_conn =
184 Connection::open_in_memory().context("failed to open an in-memory SQLite connection")?;
185 Backup::new(&disk_conn, &mut mem_conn)
186 .context("failed to start SQLite backup into memory")?
187 .run_to_completion(i32::MAX, Duration::from_millis(0), None)
188 .with_context(|| format!("failed to back up '{}' into memory", path.display()))?;
189 drop(disk_conn);
190
191 let leaked: &'static Mutex<Connection> = Box::leak(Box::new(Mutex::new(mem_conn)));
192 let mut cache = cache.lock().unwrap();
193 Ok(*cache.entry(cache_key.to_string()).or_insert(leaked))
194}
195
196#[derive(Debug, Clone, Serialize)]
197pub struct EndpointRecord {
198 pub operation_id: String,
199 pub path: String,
200 pub method: String,
201 pub summary: Option<String>,
202 pub description: Option<String>,
203 pub input_schema: serde_json::Value,
204 pub output_schema: serde_json::Value,
205 pub auth_scheme_ref: Option<String>,
206}
207
208fn row_to_endpoint(row: &Row) -> rusqlite::Result<EndpointRecord> {
209 let input_schema: String = row.get(5)?;
210 let output_schema: String = row.get(6)?;
211 Ok(EndpointRecord {
212 operation_id: row.get(0)?,
213 path: row.get(1)?,
214 method: row.get(2)?,
215 summary: row.get(3)?,
216 description: row.get(4)?,
217 input_schema: serde_json::from_str(&input_schema).unwrap_or(serde_json::Value::Null),
218 output_schema: serde_json::from_str(&output_schema).unwrap_or(serde_json::Value::Null),
219 auth_scheme_ref: row.get(7)?,
220 })
221}
222
223pub fn get_endpoint(conn: &Connection, operation_id: &str) -> Result<Option<EndpointRecord>> {
224 let mut stmt = conn.prepare(&format!(
225 "SELECT {ENDPOINT_COLUMNS} FROM endpoints WHERE operation_id = ?1"
226 ))?;
227 match stmt.query_row([operation_id], row_to_endpoint) {
228 Ok(record) => Ok(Some(record)),
229 Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
230 Err(err) => Err(err.into()),
231 }
232}
233
234pub fn list_endpoints(conn: &Connection) -> Result<Vec<EndpointRecord>> {
235 let mut stmt = conn.prepare(&format!("SELECT {ENDPOINT_COLUMNS} FROM endpoints"))?;
236 let rows = stmt.query_map([], row_to_endpoint)?;
237 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
238}
239
240#[derive(Debug, Clone, Serialize)]
241pub struct SearchResult {
242 pub operation_id: String,
243 pub summary: Option<String>,
244 pub similarity: f64,
245}
246
247pub fn search_endpoints(
262 conn: &Connection,
263 query_embedding: &[f32],
264 limit: usize,
265) -> Result<Vec<SearchResult>> {
266 let blob: Vec<u8> = query_embedding
267 .iter()
268 .flat_map(|value| value.to_le_bytes())
269 .collect();
270
271 let mut stmt = conn.prepare(
272 "SELECT e.operation_id, e.summary, s.distance
273 FROM semantic_endpoints s
274 JOIN endpoints e ON e.operation_id = s.operation_id
275 WHERE s.embedding MATCH ?1 AND k = ?2
276 ORDER BY s.distance",
277 )?;
278 let rows = stmt.query_map(rusqlite::params![blob, limit], |row| {
279 let distance: f64 = row.get(2)?;
280 Ok(SearchResult {
281 operation_id: row.get(0)?,
282 summary: row.get(1)?,
283 similarity: 1.0 - distance,
284 })
285 })?;
286 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
287}
288
289#[cfg(test)]
290mod tests {
291 use super::*;
292
293 #[test]
299 fn every_version_store_file_has_embedded_bytes() {
300 let file_labels: std::collections::HashSet<_> = VERSION_STORE_FILES
301 .iter()
302 .map(|(label, _)| *label)
303 .collect();
304 let byte_labels: std::collections::HashSet<_> = VERSION_STORE_BYTES
305 .iter()
306 .map(|(label, _)| *label)
307 .collect();
308 assert_eq!(file_labels, byte_labels);
309 }
310
311 fn seeded_store(path: &Path) -> Connection {
317 unsafe {
318 #[allow(clippy::missing_transmute_annotations)]
319 rusqlite::ffi::sqlite3_auto_extension(Some(std::mem::transmute(
320 sqlite_vec::sqlite3_vec_init as *const (),
321 )));
322 }
323 let conn = Connection::open(path).unwrap();
324 conn.execute(
325 "CREATE TABLE endpoints (
326 operation_id TEXT PRIMARY KEY,
327 path TEXT NOT NULL,
328 method TEXT NOT NULL,
329 summary TEXT,
330 description TEXT,
331 input_schema TEXT NOT NULL,
332 output_schema TEXT NOT NULL,
333 auth_scheme_ref TEXT
334 )",
335 [],
336 )
337 .unwrap();
338 conn.execute(
339 "CREATE VIRTUAL TABLE semantic_endpoints USING vec0(
340 operation_id TEXT PRIMARY KEY,
341 embedding FLOAT[4]
342 )",
343 [],
344 )
345 .unwrap();
346 conn.execute(
347 "INSERT INTO endpoints (operation_id, path, method, summary, description, input_schema, output_schema, auth_scheme_ref)
348 VALUES ('listWidgets', '/widgets', 'GET', 'List widgets', NULL, '{}', '[]', NULL)",
349 [],
350 )
351 .unwrap();
352 let embedding: Vec<u8> = [1.0f32, 0.0, 0.0, 0.0]
353 .iter()
354 .flat_map(|v| v.to_le_bytes())
355 .collect();
356 conn.execute(
357 "INSERT INTO semantic_endpoints (operation_id, embedding) VALUES ('listWidgets', ?1)",
358 rusqlite::params![embedding],
359 )
360 .unwrap();
361 conn
362 }
363
364 #[test]
365 fn get_endpoint_returns_a_seeded_row() {
366 let dir = tempfile::tempdir().unwrap();
367 let path = dir.path().join("mcp_store.db");
368 let _conn = seeded_store(&path);
369
370 let store = open_store(&path).unwrap();
371 let endpoint = get_endpoint(&store, "listWidgets").unwrap().unwrap();
372 assert_eq!(endpoint.path, "/widgets");
373 assert_eq!(endpoint.method, "GET");
374 assert_eq!(endpoint.summary.as_deref(), Some("List widgets"));
375 }
376
377 #[test]
378 fn get_endpoint_returns_none_for_an_unknown_operation() {
379 let dir = tempfile::tempdir().unwrap();
380 let path = dir.path().join("mcp_store.db");
381 let _conn = seeded_store(&path);
382
383 let store = open_store(&path).unwrap();
384 assert!(get_endpoint(&store, "unknownOp").unwrap().is_none());
385 }
386
387 #[test]
388 fn list_endpoints_returns_every_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 endpoints = list_endpoints(&store).unwrap();
395 assert_eq!(endpoints.len(), 1);
396 assert_eq!(endpoints[0].operation_id, "listWidgets");
397 }
398
399 #[test]
400 fn search_endpoints_finds_the_nearest_neighbor() {
401 let dir = tempfile::tempdir().unwrap();
402 let path = dir.path().join("mcp_store.db");
403 let _conn = seeded_store(&path);
404
405 let store = open_store(&path).unwrap();
406 let results = search_endpoints(&store, &[1.0, 0.0, 0.0, 0.0], 5).unwrap();
407 assert_eq!(results.len(), 1);
408 assert_eq!(results[0].operation_id, "listWidgets");
409 assert!(results[0].similarity > 0.99);
410 }
411
412 #[test]
413 fn cached_in_memory_connection_serves_seeded_data() {
414 let dir = tempfile::tempdir().unwrap();
415 let path = dir.path().join("mcp_store.db");
416 let _conn = seeded_store(&path);
417
418 let cached = cached_in_memory_connection("cached-serves-seeded-data", &path).unwrap();
419 let endpoint = get_endpoint(&cached.lock().unwrap(), "listWidgets")
420 .unwrap()
421 .unwrap();
422 assert_eq!(endpoint.path, "/widgets");
423
424 let results = search_endpoints(&cached.lock().unwrap(), &[1.0, 0.0, 0.0, 0.0], 5).unwrap();
428 assert_eq!(results.len(), 1);
429 assert_eq!(results[0].operation_id, "listWidgets");
430 }
431
432 #[test]
433 fn cached_in_memory_connection_holds_no_lingering_lock_on_the_disk_file() {
434 let dir = tempfile::tempdir().unwrap();
435 let path = dir.path().join("mcp_store.db");
436 let _conn = seeded_store(&path);
437
438 let cached = cached_in_memory_connection("cached-releases-disk-handle", &path).unwrap();
439
440 std::fs::remove_file(&path).unwrap();
444
445 let endpoint = get_endpoint(&cached.lock().unwrap(), "listWidgets")
446 .unwrap()
447 .unwrap();
448 assert_eq!(endpoint.path, "/widgets");
449 }
450
451 #[test]
452 fn cached_in_memory_connection_reuses_the_same_connection_for_the_same_key() {
453 let dir = tempfile::tempdir().unwrap();
454 let path = dir.path().join("mcp_store.db");
455 let _conn = seeded_store(&path);
456
457 let first = cached_in_memory_connection("cached-reuses-same-key", &path).unwrap()
458 as *const Mutex<Connection>;
459 let second = cached_in_memory_connection("cached-reuses-same-key", &path).unwrap()
460 as *const Mutex<Connection>;
461 assert_eq!(
462 first, second,
463 "expected the same cached connection, not a fresh backup"
464 );
465 }
466}