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 ("ghes-2.22", "mcp_store_vghes-2.22.db"),
32];
33
34const VERSION_STORE_BYTES: &[(&str, &[u8])] = &[
35 ("gh-2026-03-10", include_bytes!("../../mcp_store.db")),
36 (
37 "ghec-2026-03-10",
38 include_bytes!("../../mcp_store_vghec-2026-03-10.db"),
39 ),
40 ("ghes-3.21", include_bytes!("../../mcp_store_vghes-3.21.db")),
41 ("ghes-3.20", include_bytes!("../../mcp_store_vghes-3.20.db")),
42 ("ghes-3.19", include_bytes!("../../mcp_store_vghes-3.19.db")),
43 ("ghes-2.22", include_bytes!("../../mcp_store_vghes-2.22.db")),
44];
45pub fn resolve_store_path(api_version: &str) -> Result<PathBuf> {
62 let file = VERSION_STORE_FILES
63 .iter()
64 .find(|(label, _)| *label == api_version)
65 .map(|(_, file)| *file)
66 .with_context(|| format!("unknown api_version '{api_version}' — run the 'versions' command to see what's available"))?;
67 let bytes = VERSION_STORE_BYTES
68 .iter()
69 .find(|(label, _)| *label == api_version)
70 .map(|(_, bytes)| *bytes)
71 .with_context(|| format!("no embedded store data for api_version '{api_version}'"))?;
72
73 let mut dir = std::env::temp_dir();
74 dir.push(concat!(env!("CARGO_PKG_NAME"), "-store"));
75 std::fs::create_dir_all(&dir)
76 .with_context(|| format!("failed to create temp dir '{}'", dir.display()))?;
77
78 let path = dir.join(file);
79 std::fs::write(&path, bytes).with_context(|| {
86 format!(
87 "failed to extract embedded store data to '{}'",
88 path.display()
89 )
90 })?;
91
92 Ok(path)
93}
94
95const ENDPOINT_COLUMNS: &str = "operation_id, path, method, summary, description, input_schema, output_schema, auth_scheme_ref";
96
97fn register_vec_extension() {
101 REGISTER_VEC_EXTENSION.call_once(|| unsafe {
102 #[allow(clippy::missing_transmute_annotations)]
103 rusqlite::ffi::sqlite3_auto_extension(Some(std::mem::transmute(
104 sqlite_vec::sqlite3_vec_init as *const (),
105 )));
106 });
107}
108
109pub fn open_store(path: &Path) -> Result<Connection> {
113 register_vec_extension();
114 Connection::open_with_flags(path, OpenFlags::SQLITE_OPEN_READ_ONLY)
115 .with_context(|| format!("failed to open '{}'", path.display()))
116}
117
118pub fn open_store_read_write(path: &Path) -> Result<Connection> {
123 register_vec_extension();
124 Connection::open_with_flags(path, OpenFlags::SQLITE_OPEN_READ_WRITE)
125 .with_context(|| format!("failed to open '{}'", path.display()))
126}
127
128pub fn cached_store_connection(api_version: &str) -> Result<&'static Mutex<Connection>> {
146 let path = resolve_store_path(api_version)?;
147 cached_in_memory_connection(api_version, &path)
148}
149
150fn cached_in_memory_connection(cache_key: &str, path: &Path) -> Result<&'static Mutex<Connection>> {
157 static CACHE: OnceLock<Mutex<HashMap<String, &'static Mutex<Connection>>>> = OnceLock::new();
158 let cache = CACHE.get_or_init(|| Mutex::new(HashMap::new()));
159
160 if let Some(conn) = cache.lock().unwrap().get(cache_key) {
161 return Ok(conn);
162 }
163
164 let disk_conn = open_store(path)?;
165 let mut mem_conn =
166 Connection::open_in_memory().context("failed to open an in-memory SQLite connection")?;
167 Backup::new(&disk_conn, &mut mem_conn)
168 .context("failed to start SQLite backup into memory")?
169 .run_to_completion(i32::MAX, Duration::from_millis(0), None)
170 .with_context(|| format!("failed to back up '{}' into memory", path.display()))?;
171 drop(disk_conn);
172
173 let leaked: &'static Mutex<Connection> = Box::leak(Box::new(Mutex::new(mem_conn)));
174 let mut cache = cache.lock().unwrap();
175 Ok(*cache.entry(cache_key.to_string()).or_insert(leaked))
176}
177
178#[derive(Debug, Clone, Serialize)]
179pub struct EndpointRecord {
180 pub operation_id: String,
181 pub path: String,
182 pub method: String,
183 pub summary: Option<String>,
184 pub description: Option<String>,
185 pub input_schema: serde_json::Value,
186 pub output_schema: serde_json::Value,
187 pub auth_scheme_ref: Option<String>,
188}
189
190fn row_to_endpoint(row: &Row) -> rusqlite::Result<EndpointRecord> {
191 let input_schema: String = row.get(5)?;
192 let output_schema: String = row.get(6)?;
193 Ok(EndpointRecord {
194 operation_id: row.get(0)?,
195 path: row.get(1)?,
196 method: row.get(2)?,
197 summary: row.get(3)?,
198 description: row.get(4)?,
199 input_schema: serde_json::from_str(&input_schema).unwrap_or(serde_json::Value::Null),
200 output_schema: serde_json::from_str(&output_schema).unwrap_or(serde_json::Value::Null),
201 auth_scheme_ref: row.get(7)?,
202 })
203}
204
205pub fn get_endpoint(conn: &Connection, operation_id: &str) -> Result<Option<EndpointRecord>> {
206 let mut stmt = conn.prepare(&format!(
207 "SELECT {ENDPOINT_COLUMNS} FROM endpoints WHERE operation_id = ?1"
208 ))?;
209 match stmt.query_row([operation_id], row_to_endpoint) {
210 Ok(record) => Ok(Some(record)),
211 Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
212 Err(err) => Err(err.into()),
213 }
214}
215
216pub fn list_endpoints(conn: &Connection) -> Result<Vec<EndpointRecord>> {
217 let mut stmt = conn.prepare(&format!("SELECT {ENDPOINT_COLUMNS} FROM endpoints"))?;
218 let rows = stmt.query_map([], row_to_endpoint)?;
219 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
220}
221
222#[derive(Debug, Clone, Serialize)]
223pub struct SearchResult {
224 pub operation_id: String,
225 pub summary: Option<String>,
226 pub similarity: f64,
227}
228
229pub fn search_endpoints(
244 conn: &Connection,
245 query_embedding: &[f32],
246 limit: usize,
247) -> Result<Vec<SearchResult>> {
248 let blob: Vec<u8> = query_embedding
249 .iter()
250 .flat_map(|value| value.to_le_bytes())
251 .collect();
252
253 let mut stmt = conn.prepare(
254 "SELECT e.operation_id, e.summary, s.distance
255 FROM semantic_endpoints s
256 JOIN endpoints e ON e.operation_id = s.operation_id
257 WHERE s.embedding MATCH ?1 AND k = ?2
258 ORDER BY s.distance",
259 )?;
260 let rows = stmt.query_map(rusqlite::params![blob, limit], |row| {
261 let distance: f64 = row.get(2)?;
262 Ok(SearchResult {
263 operation_id: row.get(0)?,
264 summary: row.get(1)?,
265 similarity: 1.0 - distance,
266 })
267 })?;
268 Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
269}
270
271#[cfg(test)]
272mod tests {
273 use super::*;
274
275 #[test]
281 fn every_version_store_file_has_embedded_bytes() {
282 let file_labels: std::collections::HashSet<_> = VERSION_STORE_FILES
283 .iter()
284 .map(|(label, _)| *label)
285 .collect();
286 let byte_labels: std::collections::HashSet<_> = VERSION_STORE_BYTES
287 .iter()
288 .map(|(label, _)| *label)
289 .collect();
290 assert_eq!(file_labels, byte_labels);
291 }
292
293 fn seeded_store(path: &Path) -> Connection {
299 unsafe {
300 #[allow(clippy::missing_transmute_annotations)]
301 rusqlite::ffi::sqlite3_auto_extension(Some(std::mem::transmute(
302 sqlite_vec::sqlite3_vec_init as *const (),
303 )));
304 }
305 let conn = Connection::open(path).unwrap();
306 conn.execute(
307 "CREATE TABLE endpoints (
308 operation_id TEXT PRIMARY KEY,
309 path TEXT NOT NULL,
310 method TEXT NOT NULL,
311 summary TEXT,
312 description TEXT,
313 input_schema TEXT NOT NULL,
314 output_schema TEXT NOT NULL,
315 auth_scheme_ref TEXT
316 )",
317 [],
318 )
319 .unwrap();
320 conn.execute(
321 "CREATE VIRTUAL TABLE semantic_endpoints USING vec0(
322 operation_id TEXT PRIMARY KEY,
323 embedding FLOAT[4]
324 )",
325 [],
326 )
327 .unwrap();
328 conn.execute(
329 "INSERT INTO endpoints (operation_id, path, method, summary, description, input_schema, output_schema, auth_scheme_ref)
330 VALUES ('listWidgets', '/widgets', 'GET', 'List widgets', NULL, '{}', '[]', NULL)",
331 [],
332 )
333 .unwrap();
334 let embedding: Vec<u8> = [1.0f32, 0.0, 0.0, 0.0]
335 .iter()
336 .flat_map(|v| v.to_le_bytes())
337 .collect();
338 conn.execute(
339 "INSERT INTO semantic_endpoints (operation_id, embedding) VALUES ('listWidgets', ?1)",
340 rusqlite::params![embedding],
341 )
342 .unwrap();
343 conn
344 }
345
346 #[test]
347 fn get_endpoint_returns_a_seeded_row() {
348 let dir = tempfile::tempdir().unwrap();
349 let path = dir.path().join("mcp_store.db");
350 let _conn = seeded_store(&path);
351
352 let store = open_store(&path).unwrap();
353 let endpoint = get_endpoint(&store, "listWidgets").unwrap().unwrap();
354 assert_eq!(endpoint.path, "/widgets");
355 assert_eq!(endpoint.method, "GET");
356 assert_eq!(endpoint.summary.as_deref(), Some("List widgets"));
357 }
358
359 #[test]
360 fn get_endpoint_returns_none_for_an_unknown_operation() {
361 let dir = tempfile::tempdir().unwrap();
362 let path = dir.path().join("mcp_store.db");
363 let _conn = seeded_store(&path);
364
365 let store = open_store(&path).unwrap();
366 assert!(get_endpoint(&store, "unknownOp").unwrap().is_none());
367 }
368
369 #[test]
370 fn list_endpoints_returns_every_row() {
371 let dir = tempfile::tempdir().unwrap();
372 let path = dir.path().join("mcp_store.db");
373 let _conn = seeded_store(&path);
374
375 let store = open_store(&path).unwrap();
376 let endpoints = list_endpoints(&store).unwrap();
377 assert_eq!(endpoints.len(), 1);
378 assert_eq!(endpoints[0].operation_id, "listWidgets");
379 }
380
381 #[test]
382 fn search_endpoints_finds_the_nearest_neighbor() {
383 let dir = tempfile::tempdir().unwrap();
384 let path = dir.path().join("mcp_store.db");
385 let _conn = seeded_store(&path);
386
387 let store = open_store(&path).unwrap();
388 let results = search_endpoints(&store, &[1.0, 0.0, 0.0, 0.0], 5).unwrap();
389 assert_eq!(results.len(), 1);
390 assert_eq!(results[0].operation_id, "listWidgets");
391 assert!(results[0].similarity > 0.99);
392 }
393
394 #[test]
395 fn cached_in_memory_connection_serves_seeded_data() {
396 let dir = tempfile::tempdir().unwrap();
397 let path = dir.path().join("mcp_store.db");
398 let _conn = seeded_store(&path);
399
400 let cached = cached_in_memory_connection("cached-serves-seeded-data", &path).unwrap();
401 let endpoint = get_endpoint(&cached.lock().unwrap(), "listWidgets")
402 .unwrap()
403 .unwrap();
404 assert_eq!(endpoint.path, "/widgets");
405
406 let results = search_endpoints(&cached.lock().unwrap(), &[1.0, 0.0, 0.0, 0.0], 5).unwrap();
410 assert_eq!(results.len(), 1);
411 assert_eq!(results[0].operation_id, "listWidgets");
412 }
413
414 #[test]
415 fn cached_in_memory_connection_holds_no_lingering_lock_on_the_disk_file() {
416 let dir = tempfile::tempdir().unwrap();
417 let path = dir.path().join("mcp_store.db");
418 let _conn = seeded_store(&path);
419
420 let cached = cached_in_memory_connection("cached-releases-disk-handle", &path).unwrap();
421
422 std::fs::remove_file(&path).unwrap();
426
427 let endpoint = get_endpoint(&cached.lock().unwrap(), "listWidgets")
428 .unwrap()
429 .unwrap();
430 assert_eq!(endpoint.path, "/widgets");
431 }
432
433 #[test]
434 fn cached_in_memory_connection_reuses_the_same_connection_for_the_same_key() {
435 let dir = tempfile::tempdir().unwrap();
436 let path = dir.path().join("mcp_store.db");
437 let _conn = seeded_store(&path);
438
439 let first = cached_in_memory_connection("cached-reuses-same-key", &path).unwrap()
440 as *const Mutex<Connection>;
441 let second = cached_in_memory_connection("cached-reuses-same-key", &path).unwrap()
442 as *const Mutex<Connection>;
443 assert_eq!(
444 first, second,
445 "expected the same cached connection, not a fresh backup"
446 );
447 }
448}