use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::{Mutex, Once, OnceLock};
use std::time::Duration;
use anyhow::{Context, Result};
use rusqlite::backup::Backup;
use rusqlite::{Connection, OpenFlags, Row};
use serde::Serialize;
static REGISTER_VEC_EXTENSION: Once = Once::new();
pub const VERSION_STORE_FILES: &[(&str, &str)] = &[
("gh-2026-03-10", "mcp_store.db"),
("ghec-2026-03-10", "mcp_store_vghec-2026-03-10.db"),
("ghes-3.21", "mcp_store_vghes-3.21.db"),
("ghes-3.20", "mcp_store_vghes-3.20.db"),
("ghes-3.19", "mcp_store_vghes-3.19.db"),
("ghes-2.22", "mcp_store_vghes-2.22.db"),
];
const VERSION_STORE_BYTES: &[(&str, &[u8])] = &[
("gh-2026-03-10", include_bytes!("../../mcp_store.db")),
("ghec-2026-03-10", include_bytes!("../../mcp_store_vghec-2026-03-10.db")),
("ghes-3.21", include_bytes!("../../mcp_store_vghes-3.21.db")),
("ghes-3.20", include_bytes!("../../mcp_store_vghes-3.20.db")),
("ghes-3.19", include_bytes!("../../mcp_store_vghes-3.19.db")),
("ghes-2.22", include_bytes!("../../mcp_store_vghes-2.22.db")),
];
pub fn resolve_store_path(api_version: &str) -> Result<PathBuf> {
let file = VERSION_STORE_FILES
.iter()
.find(|(label, _)| *label == api_version)
.map(|(_, file)| *file)
.with_context(|| format!("unknown api_version '{api_version}' — run the 'versions' command to see what's available"))?;
let bytes = VERSION_STORE_BYTES
.iter()
.find(|(label, _)| *label == api_version)
.map(|(_, bytes)| *bytes)
.with_context(|| format!("no embedded store data for api_version '{api_version}'"))?;
let mut dir = std::env::temp_dir();
dir.push(concat!(env!("CARGO_PKG_NAME"), "-store"));
std::fs::create_dir_all(&dir)
.with_context(|| format!("failed to create temp dir '{}'", dir.display()))?;
let path = dir.join(file);
std::fs::write(&path, bytes).with_context(|| {
format!(
"failed to extract embedded store data to '{}'",
path.display()
)
})?;
Ok(path)
}
const ENDPOINT_COLUMNS: &str = "operation_id, path, method, summary, description, input_schema, output_schema, auth_scheme_ref";
fn register_vec_extension() {
REGISTER_VEC_EXTENSION.call_once(|| unsafe {
#[allow(clippy::missing_transmute_annotations)]
rusqlite::ffi::sqlite3_auto_extension(Some(std::mem::transmute(
sqlite_vec::sqlite3_vec_init as *const (),
)));
});
}
pub fn open_store(path: &Path) -> Result<Connection> {
register_vec_extension();
Connection::open_with_flags(path, OpenFlags::SQLITE_OPEN_READ_ONLY)
.with_context(|| format!("failed to open '{}'", path.display()))
}
pub fn open_store_read_write(path: &Path) -> Result<Connection> {
register_vec_extension();
Connection::open_with_flags(path, OpenFlags::SQLITE_OPEN_READ_WRITE)
.with_context(|| format!("failed to open '{}'", path.display()))
}
pub fn cached_store_connection(api_version: &str) -> Result<&'static Mutex<Connection>> {
let path = resolve_store_path(api_version)?;
cached_in_memory_connection(api_version, &path)
}
fn cached_in_memory_connection(cache_key: &str, path: &Path) -> Result<&'static Mutex<Connection>> {
static CACHE: OnceLock<Mutex<HashMap<String, &'static Mutex<Connection>>>> = OnceLock::new();
let cache = CACHE.get_or_init(|| Mutex::new(HashMap::new()));
if let Some(conn) = cache.lock().unwrap().get(cache_key) {
return Ok(conn);
}
let disk_conn = open_store(path)?;
let mut mem_conn =
Connection::open_in_memory().context("failed to open an in-memory SQLite connection")?;
Backup::new(&disk_conn, &mut mem_conn)
.context("failed to start SQLite backup into memory")?
.run_to_completion(i32::MAX, Duration::from_millis(0), None)
.with_context(|| format!("failed to back up '{}' into memory", path.display()))?;
drop(disk_conn);
let leaked: &'static Mutex<Connection> = Box::leak(Box::new(Mutex::new(mem_conn)));
let mut cache = cache.lock().unwrap();
Ok(*cache.entry(cache_key.to_string()).or_insert(leaked))
}
#[derive(Debug, Clone, Serialize)]
pub struct EndpointRecord {
pub operation_id: String,
pub path: String,
pub method: String,
pub summary: Option<String>,
pub description: Option<String>,
pub input_schema: serde_json::Value,
pub output_schema: serde_json::Value,
pub auth_scheme_ref: Option<String>,
}
fn row_to_endpoint(row: &Row) -> rusqlite::Result<EndpointRecord> {
let input_schema: String = row.get(5)?;
let output_schema: String = row.get(6)?;
Ok(EndpointRecord {
operation_id: row.get(0)?,
path: row.get(1)?,
method: row.get(2)?,
summary: row.get(3)?,
description: row.get(4)?,
input_schema: serde_json::from_str(&input_schema).unwrap_or(serde_json::Value::Null),
output_schema: serde_json::from_str(&output_schema).unwrap_or(serde_json::Value::Null),
auth_scheme_ref: row.get(7)?,
})
}
pub fn get_endpoint(conn: &Connection, operation_id: &str) -> Result<Option<EndpointRecord>> {
let mut stmt = conn.prepare(&format!(
"SELECT {ENDPOINT_COLUMNS} FROM endpoints WHERE operation_id = ?1"
))?;
match stmt.query_row([operation_id], row_to_endpoint) {
Ok(record) => Ok(Some(record)),
Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
Err(err) => Err(err.into()),
}
}
pub fn list_endpoints(conn: &Connection) -> Result<Vec<EndpointRecord>> {
let mut stmt = conn.prepare(&format!("SELECT {ENDPOINT_COLUMNS} FROM endpoints"))?;
let rows = stmt.query_map([], row_to_endpoint)?;
Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
}
#[derive(Debug, Clone, Serialize)]
pub struct SearchResult {
pub operation_id: String,
pub summary: Option<String>,
pub similarity: f64,
}
pub fn search_endpoints(
conn: &Connection,
query_embedding: &[f32],
limit: usize,
) -> Result<Vec<SearchResult>> {
let blob: Vec<u8> = query_embedding
.iter()
.flat_map(|value| value.to_le_bytes())
.collect();
let mut stmt = conn.prepare(
"SELECT e.operation_id, e.summary, s.distance
FROM semantic_endpoints s
JOIN endpoints e ON e.operation_id = s.operation_id
WHERE s.embedding MATCH ?1 AND k = ?2
ORDER BY s.distance",
)?;
let rows = stmt.query_map(rusqlite::params![blob, limit], |row| {
let distance: f64 = row.get(2)?;
Ok(SearchResult {
operation_id: row.get(0)?,
summary: row.get(1)?,
similarity: 1.0 - distance,
})
})?;
Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn every_version_store_file_has_embedded_bytes() {
let file_labels: std::collections::HashSet<_> = VERSION_STORE_FILES
.iter()
.map(|(label, _)| *label)
.collect();
let byte_labels: std::collections::HashSet<_> = VERSION_STORE_BYTES
.iter()
.map(|(label, _)| *label)
.collect();
assert_eq!(file_labels, byte_labels);
}
fn seeded_store(path: &Path) -> Connection {
unsafe {
#[allow(clippy::missing_transmute_annotations)]
rusqlite::ffi::sqlite3_auto_extension(Some(std::mem::transmute(
sqlite_vec::sqlite3_vec_init as *const (),
)));
}
let conn = Connection::open(path).unwrap();
conn.execute(
"CREATE TABLE endpoints (
operation_id TEXT PRIMARY KEY,
path TEXT NOT NULL,
method TEXT NOT NULL,
summary TEXT,
description TEXT,
input_schema TEXT NOT NULL,
output_schema TEXT NOT NULL,
auth_scheme_ref TEXT
)",
[],
)
.unwrap();
conn.execute(
"CREATE VIRTUAL TABLE semantic_endpoints USING vec0(
operation_id TEXT PRIMARY KEY,
embedding FLOAT[4]
)",
[],
)
.unwrap();
conn.execute(
"INSERT INTO endpoints (operation_id, path, method, summary, description, input_schema, output_schema, auth_scheme_ref)
VALUES ('listWidgets', '/widgets', 'GET', 'List widgets', NULL, '{}', '[]', NULL)",
[],
)
.unwrap();
let embedding: Vec<u8> = [1.0f32, 0.0, 0.0, 0.0]
.iter()
.flat_map(|v| v.to_le_bytes())
.collect();
conn.execute(
"INSERT INTO semantic_endpoints (operation_id, embedding) VALUES ('listWidgets', ?1)",
rusqlite::params![embedding],
)
.unwrap();
conn
}
#[test]
fn get_endpoint_returns_a_seeded_row() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("mcp_store.db");
let _conn = seeded_store(&path);
let store = open_store(&path).unwrap();
let endpoint = get_endpoint(&store, "listWidgets").unwrap().unwrap();
assert_eq!(endpoint.path, "/widgets");
assert_eq!(endpoint.method, "GET");
assert_eq!(endpoint.summary.as_deref(), Some("List widgets"));
}
#[test]
fn get_endpoint_returns_none_for_an_unknown_operation() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("mcp_store.db");
let _conn = seeded_store(&path);
let store = open_store(&path).unwrap();
assert!(get_endpoint(&store, "unknownOp").unwrap().is_none());
}
#[test]
fn list_endpoints_returns_every_row() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("mcp_store.db");
let _conn = seeded_store(&path);
let store = open_store(&path).unwrap();
let endpoints = list_endpoints(&store).unwrap();
assert_eq!(endpoints.len(), 1);
assert_eq!(endpoints[0].operation_id, "listWidgets");
}
#[test]
fn search_endpoints_finds_the_nearest_neighbor() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("mcp_store.db");
let _conn = seeded_store(&path);
let store = open_store(&path).unwrap();
let results = search_endpoints(&store, &[1.0, 0.0, 0.0, 0.0], 5).unwrap();
assert_eq!(results.len(), 1);
assert_eq!(results[0].operation_id, "listWidgets");
assert!(results[0].similarity > 0.99);
}
#[test]
fn cached_in_memory_connection_serves_seeded_data() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("mcp_store.db");
let _conn = seeded_store(&path);
let cached = cached_in_memory_connection("cached-serves-seeded-data", &path).unwrap();
let endpoint = get_endpoint(&cached.lock().unwrap(), "listWidgets")
.unwrap()
.unwrap();
assert_eq!(endpoint.path, "/widgets");
let results = search_endpoints(&cached.lock().unwrap(), &[1.0, 0.0, 0.0, 0.0], 5).unwrap();
assert_eq!(results.len(), 1);
assert_eq!(results[0].operation_id, "listWidgets");
}
#[test]
fn cached_in_memory_connection_holds_no_lingering_lock_on_the_disk_file() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("mcp_store.db");
let _conn = seeded_store(&path);
let cached = cached_in_memory_connection("cached-releases-disk-handle", &path).unwrap();
std::fs::remove_file(&path).unwrap();
let endpoint = get_endpoint(&cached.lock().unwrap(), "listWidgets")
.unwrap()
.unwrap();
assert_eq!(endpoint.path, "/widgets");
}
#[test]
fn cached_in_memory_connection_reuses_the_same_connection_for_the_same_key() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("mcp_store.db");
let _conn = seeded_store(&path);
let first = cached_in_memory_connection("cached-reuses-same-key", &path).unwrap()
as *const Mutex<Connection>;
let second = cached_in_memory_connection("cached-reuses-same-key", &path).unwrap()
as *const Mutex<Connection>;
assert_eq!(
first, second,
"expected the same cached connection, not a fresh backup"
);
}
}