pub mod token_auth_provider;
use std::{
collections::HashMap,
path::{Path, PathBuf},
sync::Arc,
};
use async_trait::async_trait;
use rust_mcp_sdk::{
McpServer,
macros::{JsonSchema, mcp_tool},
mcp_server::ServerHandler,
schema::{
CallToolRequestParams, CallToolResult, ListToolsResult, PaginatedRequestParams, RpcError,
TextContent, schema_utils::CallToolError,
},
tool_box,
};
use serde::{Deserialize, Serialize};
use crate::{search, storage};
#[mcp_tool(
name = "search_docs",
description = "Fuzzy-search documentation entries across all indexed docsets. Returns matching entries with their IDs, types, and file paths. Use the returned `id` with `read_doc` to retrieve the full content.",
read_only_hint = true,
idempotent_hint = true,
destructive_hint = false,
open_world_hint = false
)]
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
pub struct SearchDocsTool {
pub query: String,
pub docset: Option<String>,
pub limit: Option<u64>,
}
#[mcp_tool(
name = "read_doc",
description = "Return the Markdown content for a documentation entry by its numeric ID (as returned by search_docs).",
read_only_hint = true,
idempotent_hint = true,
destructive_hint = false,
open_world_hint = false
)]
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
pub struct ReadDocTool {
pub entry_id: i64,
}
#[mcp_tool(
name = "list_docsets",
description = "List docsets with their installed versions. Returns a JSON array of objects with `name`, `installed_versions`, and `available_versions` fields. Filter by name (case-insensitive substring) or set `installed_only` to exclude names with no local installation.",
read_only_hint = true,
idempotent_hint = true,
destructive_hint = false,
open_world_hint = false
)]
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
pub struct ListDocsetsTool {
pub name: Option<String>,
pub installed_only: Option<bool>,
}
tool_box!(DocdbTools, [SearchDocsTool, ReadDocTool, ListDocsetsTool]);
pub struct DocdbHandler {
db_path: PathBuf,
compression_level: i32,
}
impl DocdbHandler {
pub fn new(db_path: PathBuf, compression_level: i32) -> Self {
Self {
db_path,
compression_level,
}
}
}
#[async_trait]
impl ServerHandler for DocdbHandler {
async fn handle_list_tools_request(
&self,
_params: Option<PaginatedRequestParams>,
_runtime: Arc<dyn McpServer>,
) -> Result<ListToolsResult, RpcError> {
Ok(ListToolsResult {
meta: None,
next_cursor: None,
tools: DocdbTools::tools(),
})
}
async fn handle_call_tool_request(
&self,
params: CallToolRequestParams,
_runtime: Arc<dyn McpServer>,
) -> Result<CallToolResult, CallToolError> {
let tool: DocdbTools = DocdbTools::try_from(params).map_err(CallToolError::new)?;
let db_path = self.db_path.clone();
let compression_level = self.compression_level;
tokio::task::spawn_blocking(move || -> Result<CallToolResult, String> {
match tool {
DocdbTools::SearchDocsTool(t) => t
.execute(&db_path, compression_level)
.map_err(|e| e.to_string()),
DocdbTools::ReadDocTool(t) => t
.execute(&db_path, compression_level)
.map_err(|e| e.to_string()),
DocdbTools::ListDocsetsTool(t) => t
.execute(&db_path, compression_level)
.map_err(|e| e.to_string()),
}
})
.await
.map_err(|e| CallToolError::from_message(e.to_string()))?
.map_err(CallToolError::from_message)
}
}
impl SearchDocsTool {
fn execute(
&self,
db_path: &Path,
compression_level: i32,
) -> Result<CallToolResult, CallToolError> {
if !db_path.exists() {
return Err(CallToolError::from_message(
"No database found. Run `docterm download <name>` to add a docset.".to_string(),
));
}
let storage = storage::Storage::open(db_path, compression_level)
.map_err(|e| CallToolError::from_message(e.to_string()))?;
let (entries, scope) = if let Some(spec) = &self.docset {
let (name, version) = parse_docset_spec(spec);
let docset = storage
.find_docset_by_name(name, version)
.map_err(|e| CallToolError::from_message(e.to_string()))?
.ok_or_else(|| {
CallToolError::from_message(format!("Docset '{spec}' not found."))
})?;
let label = format!(
"{} {}",
docset.name,
docset.version.as_deref().unwrap_or("(latest)")
);
let entries = storage
.list_entries_for_docset(docset.id)
.map_err(|e| CallToolError::from_message(e.to_string()))?;
(entries, label)
} else {
let entries = storage
.list_all_entries()
.map_err(|e| CallToolError::from_message(e.to_string()))?;
(entries, "all docsets".to_string())
};
if entries.is_empty() {
return Ok(CallToolResult::text_content(vec![TextContent::from(
"No entries indexed. Run `docterm download <name>` to add a docset.".to_string(),
)]));
}
let limit = self.limit.unwrap_or(20) as usize;
let mut searcher = search::Searcher::new(entries);
let results = searcher.search(&self.query, limit);
if results.is_empty() {
return Ok(CallToolResult::text_content(vec![TextContent::from(
format!("No results for '{}' in {}.", self.query, scope),
)]));
}
let docset_names: HashMap<i64, String> = storage
.list_docsets()
.map_err(|e| CallToolError::from_message(e.to_string()))?
.into_iter()
.map(|d| (d.id, d.name))
.collect();
let mut output = format!(
"Results for '{}' in {} ({} found):\n\n",
self.query,
scope,
results.len()
);
for r in &results {
let docset_name = docset_names
.get(&r.entry.docset_id)
.map(String::as_str)
.unwrap_or("?");
output.push_str(&format!(
"- id={} [{}] {} ({}) — {}\n",
r.entry.id, docset_name, r.entry.name, r.entry.entry_type, r.entry.path
));
}
Ok(CallToolResult::text_content(vec![TextContent::from(
output,
)]))
}
}
impl ReadDocTool {
fn execute(
&self,
db_path: &Path,
compression_level: i32,
) -> Result<CallToolResult, CallToolError> {
if !db_path.exists() {
return Err(CallToolError::from_message(
"No database found. Run `docterm download <name>` to add a docset.".to_string(),
));
}
let storage = storage::Storage::open(db_path, compression_level)
.map_err(|e| CallToolError::from_message(e.to_string()))?;
let entry = storage.get_entry(self.entry_id).map_err(|e| {
CallToolError::from_message(format!("Entry {} not found: {e}", self.entry_id))
})?;
let dict = storage
.get_docset_dictionary(entry.docset_id)
.ok()
.flatten();
let page = storage
.get_page_content_dict(self.entry_id, dict.as_deref())
.map_err(|e| {
CallToolError::from_message(format!("Failed to read doc {}: {e}", self.entry_id))
})?;
Ok(CallToolResult::text_content(vec![TextContent::from(
page.markdown,
)]))
}
}
impl ListDocsetsTool {
fn execute(
&self,
db_path: &Path,
compression_level: i32,
) -> Result<CallToolResult, CallToolError> {
if !db_path.exists() {
return Err(CallToolError::from_message(
"No database found. Run `docterm download <name>` to add a docset.".to_string(),
));
}
let storage = storage::Storage::open(db_path, compression_level)
.map_err(|e| CallToolError::from_message(e.to_string()))?;
let all_docsets = storage
.list_docsets()
.map_err(|e| CallToolError::from_message(e.to_string()))?;
let mut seen: Vec<String> = Vec::new();
let mut display_names: HashMap<String, String> = HashMap::new();
let mut by_name: HashMap<String, Vec<String>> = HashMap::new();
for d in all_docsets {
let key = d.name.to_lowercase();
if !by_name.contains_key(&key) {
seen.push(key.clone());
display_names.insert(key.clone(), d.name.clone());
}
if let Some(v) = d.version {
by_name.entry(key).or_default().push(v);
} else {
by_name.entry(key).or_default();
}
}
let name_filter = self.name.as_deref().map(str::to_lowercase);
let installed_only = self.installed_only.unwrap_or(false);
let results: Vec<DocsetInfo> = seen
.into_iter()
.filter(|key| {
if let Some(ref f) = name_filter {
key.contains(f.as_str())
} else {
true
}
})
.filter_map(|key| {
let installed_versions = by_name.remove(&key).unwrap_or_default();
if installed_only && installed_versions.is_empty() {
return None;
}
let name = display_names.remove(&key).unwrap_or(key);
Some(DocsetInfo {
name,
installed_versions,
available_versions: vec![],
})
})
.collect();
let json = serde_json::to_string_pretty(&results)
.map_err(|e| CallToolError::from_message(e.to_string()))?;
Ok(CallToolResult::text_content(vec![TextContent::from(json)]))
}
}
#[derive(Debug, Serialize)]
struct DocsetInfo {
name: String,
installed_versions: Vec<String>,
available_versions: Vec<String>,
}
fn parse_docset_spec(spec: &str) -> (&str, Option<&str>) {
if let Some((name, version)) = spec.split_once(':') {
(name, Some(version))
} else {
(spec, None)
}
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
#[test]
fn parse_spec_without_version() {
let (name, version) = parse_docset_spec("Rust");
assert_eq!(name, "Rust");
assert!(version.is_none());
}
#[test]
fn parse_spec_with_version() {
let (name, version) = parse_docset_spec("Rust:1.75.0");
assert_eq!(name, "Rust");
assert_eq!(version, Some("1.75.0"));
}
#[test]
fn search_tool_missing_db_returns_error() {
let tool = SearchDocsTool {
query: "Vec".into(),
docset: None,
limit: None,
};
let result = tool.execute(&PathBuf::from("/nonexistent/library.db"), 3);
assert!(result.is_err());
let err = result.unwrap_err();
assert!(err.to_string().contains("No database found"));
}
#[test]
fn search_tool_empty_db_returns_informative_message() {
let dir = TempDir::new().unwrap();
let db_path = dir.path().join("library.db");
let _storage = storage::Storage::open(&db_path, 3).unwrap();
let tool = SearchDocsTool {
query: "Vec".into(),
docset: None,
limit: None,
};
let result = tool.execute(&db_path, 3).unwrap();
let text = extract_text(&result);
assert!(text.contains("No entries indexed"));
}
#[test]
fn search_tool_unknown_docset_returns_error() {
let dir = TempDir::new().unwrap();
let db_path = dir.path().join("library.db");
let _storage = storage::Storage::open(&db_path, 3).unwrap();
let tool = SearchDocsTool {
query: "Vec".into(),
docset: Some("NoSuchDocset".into()),
limit: None,
};
let result = tool.execute(&db_path, 3);
assert!(result.is_err());
assert!(result.unwrap_err().to_string().contains("not found"));
}
#[test]
fn search_tool_finds_entries() {
let dir = TempDir::new().unwrap();
let db_path = dir.path().join("library.db");
let storage = storage::Storage::open(&db_path, 3).unwrap();
let docset_id = storage.insert_docset("Rust", Some("1.0"), None).unwrap();
storage
.insert_entry(docset_id, "Vec", "Struct", "std/vec/struct.Vec.html")
.unwrap();
storage
.insert_entry(
docset_id,
"HashMap",
"Struct",
"std/collections/struct.HashMap.html",
)
.unwrap();
let tool = SearchDocsTool {
query: "Vec".into(),
docset: None,
limit: Some(10),
};
let result = tool.execute(&db_path, 3).unwrap();
let text = extract_text(&result);
assert!(text.contains("Vec"));
assert!(text.contains("id="));
}
#[test]
fn read_doc_missing_db_returns_error() {
let tool = ReadDocTool { entry_id: 1 };
let result = tool.execute(&PathBuf::from("/nonexistent/library.db"), 3);
assert!(result.is_err());
assert!(
result
.unwrap_err()
.to_string()
.contains("No database found")
);
}
#[test]
fn read_doc_unknown_entry_returns_error() {
let dir = TempDir::new().unwrap();
let db_path = dir.path().join("library.db");
let _storage = storage::Storage::open(&db_path, 3).unwrap();
let tool = ReadDocTool { entry_id: 9999 };
let result = tool.execute(&db_path, 3);
assert!(result.is_err());
assert!(result.unwrap_err().to_string().contains("9999"));
}
#[test]
fn read_doc_returns_markdown_content() {
let dir = TempDir::new().unwrap();
let db_path = dir.path().join("library.db");
let storage = storage::Storage::open(&db_path, 3).unwrap();
let docset_id = storage.insert_docset("Rust", Some("1.0"), None).unwrap();
let entry_id = storage
.insert_entry(docset_id, "Vec", "Struct", "std/vec/struct.Vec.html")
.unwrap();
let markdown = "# Vec\n\nA growable array.";
storage.insert_page_dict(entry_id, markdown, None).unwrap();
let tool = ReadDocTool { entry_id };
let result = tool.execute(&db_path, 3).unwrap();
let text = extract_text(&result);
assert_eq!(text, markdown);
}
fn setup_db_with_docsets(dir: &TempDir) -> PathBuf {
let db_path = dir.path().join("library.db");
let storage = storage::Storage::open(&db_path, 3).unwrap();
storage.insert_docset("Rust", Some("1.80.0"), None).unwrap();
storage.insert_docset("Rust", Some("1.94.0"), None).unwrap();
storage.insert_docset("Python", Some("3.12"), None).unwrap();
storage.insert_docset("Go", None, None).unwrap();
db_path
}
#[test]
fn list_docsets_missing_db_returns_error() {
let tool = ListDocsetsTool {
name: None,
installed_only: None,
};
let result = tool.execute(&PathBuf::from("/nonexistent/library.db"), 3);
assert!(result.is_err());
assert!(
result
.unwrap_err()
.to_string()
.contains("No database found")
);
}
#[test]
fn list_docsets_returns_all_names() {
let dir = TempDir::new().unwrap();
let db_path = setup_db_with_docsets(&dir);
let tool = ListDocsetsTool {
name: None,
installed_only: None,
};
let result = tool.execute(&db_path, 3).unwrap();
let text = extract_text(&result);
let parsed: serde_json::Value = serde_json::from_str(&text).unwrap();
let arr = parsed.as_array().unwrap();
let names: Vec<&str> = arr.iter().map(|v| v["name"].as_str().unwrap()).collect();
assert!(names.contains(&"Rust"));
assert!(names.contains(&"Python"));
assert!(names.contains(&"Go"));
}
#[test]
fn list_docsets_version_present_in_installed_versions() {
let dir = TempDir::new().unwrap();
let db_path = setup_db_with_docsets(&dir);
let tool = ListDocsetsTool {
name: Some("rust".into()),
installed_only: None,
};
let result = tool.execute(&db_path, 3).unwrap();
let text = extract_text(&result);
let parsed: serde_json::Value = serde_json::from_str(&text).unwrap();
let arr = parsed.as_array().unwrap();
assert_eq!(arr.len(), 1);
let versions = arr[0]["installed_versions"].as_array().unwrap();
assert_eq!(versions.len(), 1);
assert_eq!(versions[0].as_str().unwrap(), "1.80.0");
}
#[test]
fn list_docsets_name_filter_case_insensitive() {
let dir = TempDir::new().unwrap();
let db_path = setup_db_with_docsets(&dir);
let tool = ListDocsetsTool {
name: Some("PY".into()),
installed_only: None,
};
let result = tool.execute(&db_path, 3).unwrap();
let text = extract_text(&result);
let parsed: serde_json::Value = serde_json::from_str(&text).unwrap();
let arr = parsed.as_array().unwrap();
assert_eq!(arr.len(), 1);
assert_eq!(arr[0]["name"].as_str().unwrap(), "Python");
}
#[test]
fn list_docsets_no_match_returns_empty_array() {
let dir = TempDir::new().unwrap();
let db_path = setup_db_with_docsets(&dir);
let tool = ListDocsetsTool {
name: Some("Haskell".into()),
installed_only: None,
};
let result = tool.execute(&db_path, 3).unwrap();
let text = extract_text(&result);
let parsed: serde_json::Value = serde_json::from_str(&text).unwrap();
assert_eq!(parsed.as_array().unwrap().len(), 0);
}
#[test]
fn list_docsets_installed_only_excludes_empty() {
let dir = TempDir::new().unwrap();
let db_path = dir.path().join("library.db");
let storage = storage::Storage::open(&db_path, 3).unwrap();
storage.insert_docset("Rust", Some("1.94"), None).unwrap();
storage.insert_docset("Go", None, None).unwrap();
let tool = ListDocsetsTool {
name: None,
installed_only: Some(true),
};
let result = tool.execute(&db_path, 3).unwrap();
let text = extract_text(&result);
let parsed: serde_json::Value = serde_json::from_str(&text).unwrap();
let arr = parsed.as_array().unwrap();
let names: Vec<&str> = arr.iter().map(|v| v["name"].as_str().unwrap()).collect();
assert!(names.contains(&"Rust"));
assert!(!names.contains(&"Go"));
}
#[test]
fn list_docsets_available_versions_always_empty() {
let dir = TempDir::new().unwrap();
let db_path = setup_db_with_docsets(&dir);
let tool = ListDocsetsTool {
name: None,
installed_only: None,
};
let result = tool.execute(&db_path, 3).unwrap();
let text = extract_text(&result);
let parsed: serde_json::Value = serde_json::from_str(&text).unwrap();
for entry in parsed.as_array().unwrap() {
assert_eq!(entry["available_versions"].as_array().unwrap().len(), 0);
}
}
#[test]
fn tool_box_has_three_tools() {
let tools = DocdbTools::tools();
assert_eq!(tools.len(), 3);
let names: Vec<&str> = tools.iter().map(|t| t.name.as_str()).collect();
assert!(names.contains(&"search_docs"));
assert!(names.contains(&"read_doc"));
assert!(names.contains(&"list_docsets"));
}
fn extract_text(result: &CallToolResult) -> String {
result
.content
.iter()
.filter_map(|c| c.as_text_content().ok())
.map(|t| t.text.as_str())
.collect::<Vec<_>>()
.join("")
}
}