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