use serde_json::Value;
pub use trusty_mcp_core::{error_codes, initialize_response, JsonRpcError, Request, Response};
#[derive(Clone)]
pub struct McpServer {
base_url: String,
http: reqwest::Client,
}
impl McpServer {
pub fn new(base_url: impl Into<String>) -> Self {
Self {
base_url: base_url.into(),
http: reqwest::Client::new(),
}
}
pub fn with_client(base_url: impl Into<String>, http: reqwest::Client) -> Self {
Self {
base_url: base_url.into(),
http,
}
}
pub fn base_url(&self) -> &str {
&self.base_url
}
pub async fn dispatch(&self, req: Request) -> Response {
let is_notification = req.id.is_none();
let id = req.id.clone();
if req.jsonrpc.as_deref() != Some("2.0") {
if is_notification {
return Response::suppressed();
}
return Response::err(id, error_codes::INVALID_REQUEST, "jsonrpc must be \"2.0\"");
}
match req.method.as_str() {
"initialize" => {
return Response::ok(
id,
initialize_response("trusty-search", env!("CARGO_PKG_VERSION"), None),
);
}
"notifications/initialized" | "initialized" => {
return Response::suppressed();
}
_ => {}
}
let params = req.params.clone().unwrap_or(Value::Null);
let (tool, arguments, via_tools_call) = match req.method.as_str() {
"tools/call" => {
let name = params
.get("name")
.and_then(Value::as_str)
.map(str::to_owned);
let args = params
.get("arguments")
.cloned()
.unwrap_or(Value::Object(Default::default()));
match name {
Some(n) => (n, args, true),
None => {
return Response::err(
id,
error_codes::INVALID_PARAMS,
"tools/call requires a 'name' field",
)
}
}
}
"tools/list" => {
return Response::ok(id, serde_json::json!({ "tools": tool_descriptors() }));
}
other => (other.to_string(), params, false),
};
let outcome = self.call_tool(&tool, &arguments).await;
if via_tools_call {
match outcome {
Ok(value) => Response::ok(id, wrap_tool_result(&value, false)),
Err(DispatchError::UnknownTool) => Response::err(
id,
error_codes::METHOD_NOT_FOUND,
format!("unknown tool: {tool}"),
),
Err(DispatchError::InvalidParams(msg)) => Response::ok(id, wrap_tool_error(&msg)),
Err(DispatchError::Transport(msg)) => Response::ok(id, wrap_tool_error(&msg)),
}
} else {
match outcome {
Ok(value) => Response::ok(id, wrap_text_content(&value)),
Err(DispatchError::UnknownTool) => Response::err(
id,
error_codes::METHOD_NOT_FOUND,
format!("unknown tool: {tool}"),
),
Err(DispatchError::InvalidParams(msg)) => {
Response::err(id, error_codes::INVALID_PARAMS, msg)
}
Err(DispatchError::Transport(msg)) => {
Response::err(id, error_codes::INTERNAL_ERROR, msg)
}
}
}
}
async fn call_tool(&self, tool: &str, args: &Value) -> Result<Value, DispatchError> {
match tool {
"search_all" => {
let query = require_str(args, "query")?;
let top_k = args.get("top_k").and_then(Value::as_u64).unwrap_or(10);
let full_content = args
.get("full_content")
.and_then(Value::as_bool)
.unwrap_or(false);
let body = serde_json::json!({
"query": query,
"top_k": top_k,
"full_content": full_content,
});
self.post("/search", &body).await
}
"search_code" => {
let index_id = require_str(args, "index_id")?;
let body = match args.get("query") {
Some(Value::Object(_)) => args.get("query").cloned().unwrap(),
Some(Value::String(text)) => {
let mut b = serde_json::json!({ "text": text });
if let Some(k) = args.get("top_k").and_then(Value::as_u64) {
b["top_k"] = Value::from(k);
}
b
}
_ => {
return Err(DispatchError::InvalidParams(
"missing or invalid 'query' (expected string or object)".into(),
))
}
};
self.post(&format!("/indexes/{index_id}/search"), &body)
.await
}
"index_file" => {
let index_id = require_str(args, "index_id")?;
let path = require_str(args, "path")?;
let content = require_str(args, "content")?;
self.post(
&format!("/indexes/{index_id}/index-file"),
&serde_json::json!({ "path": path, "content": content }),
)
.await
}
"remove_file" => {
let index_id = require_str(args, "index_id")?;
let path = require_str(args, "path")?;
self.post(
&format!("/indexes/{index_id}/remove-file"),
&serde_json::json!({ "path": path }),
)
.await
}
"list_indexes" => self.get("/indexes").await,
"create_index" => {
let id = require_str(args, "id")?;
let root_path = require_str(args, "root_path")?;
self.post(
"/indexes",
&serde_json::json!({ "id": id, "root_path": root_path }),
)
.await
}
"search_similar" => {
let index_id = args
.get("index")
.and_then(Value::as_str)
.unwrap_or("default");
let file = require_str(args, "file")?;
let mut body = serde_json::json!({ "file": file });
if let Some(func) = args.get("function").and_then(Value::as_str) {
body["function"] = Value::String(func.to_string());
}
if let Some(k) = args.get("top_k").and_then(Value::as_u64) {
body["top_k"] = Value::from(k);
}
self.post(&format!("/indexes/{index_id}/search_similar"), &body)
.await
}
"search_health" => self.get("/health").await,
"delete_index" => {
let index_id = require_str(args, "index_id")?;
self.delete(&format!("/indexes/{index_id}")).await
}
"reindex" => {
let index_id = require_str(args, "index_id")?;
let mut body = serde_json::json!({});
if let Some(rp) = args.get("root_path").and_then(Value::as_str) {
body["root_path"] = Value::String(rp.to_string());
}
self.post(&format!("/indexes/{index_id}/reindex"), &body)
.await
}
"index_status" => {
let index_id = require_str(args, "index_id")?;
self.get(&format!("/indexes/{index_id}/status")).await
}
"chat" => {
let index_id = require_str(args, "index_id")?;
let message = args
.get("message")
.and_then(Value::as_str)
.or_else(|| args.get("question").and_then(Value::as_str))
.ok_or_else(|| {
DispatchError::InvalidParams(
"missing required string field: message (or question)".into(),
)
})?;
let mut body = serde_json::json!({
"index_id": index_id,
"message": message,
});
if let Some(history) = args.get("history") {
body["history"] = history.clone();
}
if let Some(model) = args.get("model").and_then(Value::as_str) {
body["model"] = Value::String(model.to_string());
}
if let Some(top_k) = args.get("top_k").and_then(Value::as_u64) {
body["top_k"] = Value::from(top_k);
}
if let Some(key) = args.get("api_key").and_then(Value::as_str) {
body["api_key"] = Value::String(key.to_string());
}
self.post("/chat", &body).await
}
"list_chunks" => {
let index_id = require_str(args, "index_id")?;
let offset = args.get("offset").and_then(Value::as_u64).unwrap_or(0);
let limit = args.get("limit").and_then(Value::as_u64).unwrap_or(100);
self.get(&format!(
"/indexes/{index_id}/chunks?offset={offset}&limit={limit}"
))
.await
}
_ => Err(DispatchError::UnknownTool),
}
}
async fn get(&self, path: &str) -> Result<Value, DispatchError> {
let url = format!("{}{}", self.base_url, path);
let resp = self
.http
.get(&url)
.send()
.await
.map_err(|e| DispatchError::Transport(format!("GET {url}: {e}")))?;
let status = resp.status();
let body: Value = resp
.json()
.await
.map_err(|e| DispatchError::Transport(format!("decode {url}: {e}")))?;
if !status.is_success() {
return Err(DispatchError::Transport(format!(
"GET {url} returned {status}: {body}"
)));
}
Ok(body)
}
async fn delete(&self, path: &str) -> Result<Value, DispatchError> {
let url = format!("{}{}", self.base_url, path);
let resp = self
.http
.delete(&url)
.send()
.await
.map_err(|e| DispatchError::Transport(format!("DELETE {url}: {e}")))?;
let status = resp.status();
let body: Value = resp
.json()
.await
.map_err(|e| DispatchError::Transport(format!("decode {url}: {e}")))?;
if !status.is_success() {
return Err(DispatchError::Transport(format!(
"DELETE {url} returned {status}: {body}"
)));
}
Ok(body)
}
async fn post(&self, path: &str, body: &Value) -> Result<Value, DispatchError> {
let url = format!("{}{}", self.base_url, path);
let resp = self
.http
.post(&url)
.json(body)
.send()
.await
.map_err(|e| DispatchError::Transport(format!("POST {url}: {e}")))?;
let status = resp.status();
let body: Value = resp
.json()
.await
.map_err(|e| DispatchError::Transport(format!("decode {url}: {e}")))?;
if !status.is_success() {
return Err(DispatchError::Transport(format!(
"POST {url} returned {status}: {body}"
)));
}
Ok(body)
}
}
#[derive(Debug)]
enum DispatchError {
UnknownTool,
InvalidParams(String),
Transport(String),
}
fn require_str<'a>(args: &'a Value, key: &str) -> Result<&'a str, DispatchError> {
args.get(key)
.and_then(Value::as_str)
.ok_or_else(|| DispatchError::InvalidParams(format!("missing or non-string '{key}'")))
}
fn wrap_text_content(value: &Value) -> Value {
serde_json::json!({
"content": [{
"type": "text",
"text": serde_json::to_string_pretty(value).unwrap_or_else(|_| value.to_string()),
}]
})
}
fn wrap_tool_result(value: &Value, is_error: bool) -> Value {
serde_json::json!({
"content": [{
"type": "text",
"text": serde_json::to_string_pretty(value).unwrap_or_else(|_| value.to_string()),
}],
"isError": is_error,
})
}
fn wrap_tool_error(msg: &str) -> Value {
serde_json::json!({
"content": [{
"type": "text",
"text": format!("Error: {msg}"),
}],
"isError": true,
})
}
pub fn tool_descriptors() -> Value {
serde_json::json!([
{
"name": "search_all",
"description": "Cross-project hybrid search: fan out to every registered index, merge results via RRF, tag each chunk with its index_id (issue #10).",
"inputSchema": {
"type": "object",
"required": ["query"],
"properties": {
"query": { "type": "string" },
"top_k": { "type": "integer", "default": 10 },
"full_content": { "type": "boolean", "default": false }
}
}
},
{
"name": "search_code",
"description": "Hybrid code search (BM25+vector+KG)",
"inputSchema": {
"type": "object",
"required": ["index_id", "query"],
"properties": {
"index_id": { "type": "string" },
"query": { "type": "string" },
"top_k": { "type": "integer", "default": 10 }
}
}
},
{
"name": "index_file",
"description": "Add or update one file in an index",
"inputSchema": {
"type": "object",
"required": ["index_id", "path", "content"],
"properties": {
"index_id": { "type": "string" },
"path": { "type": "string" },
"content": { "type": "string" }
}
}
},
{
"name": "remove_file",
"description": "Remove a file's chunks from an index",
"inputSchema": {
"type": "object",
"required": ["index_id", "path"],
"properties": {
"index_id": { "type": "string" },
"path": { "type": "string" }
}
}
},
{
"name": "list_indexes",
"description": "List all registered indexes on this daemon",
"inputSchema": { "type": "object", "properties": {} }
},
{
"name": "create_index",
"description": "Register a new (empty) index",
"inputSchema": {
"type": "object",
"required": ["id", "root_path"],
"properties": {
"id": { "type": "string" },
"root_path": { "type": "string" }
}
}
},
{
"name": "search_similar",
"description": "Find chunks semantically similar to a given file/function via HNSW (issue #31)",
"inputSchema": {
"type": "object",
"required": ["file"],
"properties": {
"file": { "type": "string" },
"function": { "type": "string" },
"top_k": { "type": "number" },
"index": { "type": "string" }
}
}
},
{
"name": "search_health",
"description": "Probe daemon liveness and version",
"inputSchema": { "type": "object", "properties": {} }
},
{
"name": "delete_index",
"description": "Delete a registered index and all its data",
"inputSchema": {
"type": "object",
"required": ["index_id"],
"properties": {
"index_id": { "type": "string" }
}
}
},
{
"name": "reindex",
"description": "Trigger a full reindex of a collection (async, returns immediately)",
"inputSchema": {
"type": "object",
"required": ["index_id"],
"properties": {
"index_id": { "type": "string" },
"root_path": { "type": "string" }
}
}
},
{
"name": "index_status",
"description": "Get stats for an index (chunk count, root path)",
"inputSchema": {
"type": "object",
"required": ["index_id"],
"properties": {
"index_id": { "type": "string" }
}
}
},
{
"name": "list_chunks",
"description": "Paginated enumeration of every chunk in an index (issue #54). Stable order by (file, start_line).",
"inputSchema": {
"type": "object",
"required": ["index_id"],
"properties": {
"index_id": { "type": "string" },
"offset": { "type": "integer", "default": 0 },
"limit": { "type": "integer", "default": 100 }
}
}
},
{
"name": "chat",
"description": "Ask a natural-language question about the indexed codebase. \
Automatically searches for the top_k most relevant chunks and \
sends them as context to an OpenRouter LLM (default model: \
anthropic/claude-haiku-4.5). Returns {answer, sources, model}. \
Requires OPENROUTER_API_KEY env var on the daemon, or an \
`api_key` field in the request.",
"inputSchema": {
"type": "object",
"required": ["index_id"],
"properties": {
"index_id": { "type": "string" },
"message": { "type": "string", "description": "User question (alias: question)" },
"question": { "type": "string", "description": "User question (alias: message)" },
"history": { "type": "array", "items": { "type": "object" } },
"model": { "type": "string", "description": "OpenRouter model id (default: anthropic/claude-haiku-4.5)" },
"top_k": { "type": "integer", "description": "Number of context chunks (default: 5)", "default": 5 },
"api_key": { "type": "string", "description": "Fallback OpenRouter API key when OPENROUTER_API_KEY env is unset" }
}
}
}
])
}
#[cfg(test)]
mod tests {
use super::*;
fn req(method: &str, params: Value) -> Request {
Request {
jsonrpc: Some("2.0".into()),
id: Some(Value::from(1u64)),
method: method.into(),
params: Some(params),
}
}
#[tokio::test]
async fn rejects_wrong_jsonrpc_version() {
let server = McpServer::new("http://127.0.0.1:1");
let r = Request {
jsonrpc: Some("1.0".into()),
id: Some(Value::from(7u64)),
method: "search_health".into(),
params: None,
};
let resp = server.dispatch(r).await;
let err = resp.error.expect("expected error");
assert_eq!(err.code, error_codes::INVALID_REQUEST);
assert_eq!(resp.id, Some(Value::from(7u64)));
}
#[tokio::test]
async fn unknown_tool_returns_method_not_found() {
let server = McpServer::new("http://127.0.0.1:1");
let resp = server.dispatch(req("not_a_tool", Value::Null)).await;
let err = resp.error.expect("expected error");
assert_eq!(err.code, error_codes::METHOD_NOT_FOUND);
}
#[tokio::test]
async fn missing_params_returns_invalid_params() {
let server = McpServer::new("http://127.0.0.1:1");
let resp = server
.dispatch(req("index_file", serde_json::json!({})))
.await;
let err = resp.error.expect("expected error");
assert_eq!(err.code, error_codes::INVALID_PARAMS);
}
#[tokio::test]
async fn tools_list_returns_all_tools() {
let server = McpServer::new("http://127.0.0.1:1");
let resp = server.dispatch(req("tools/list", Value::Null)).await;
let result = resp.result.expect("expected result");
let tools = result
.get("tools")
.and_then(Value::as_array)
.expect("array");
assert!(
tools.len() >= 6,
"expected at least 6 tools, got {}",
tools.len()
);
let names: Vec<&str> = tools
.iter()
.filter_map(|t| t.get("name").and_then(Value::as_str))
.collect();
for required in [
"search_code",
"index_file",
"remove_file",
"list_indexes",
"create_index",
"search_health",
] {
assert!(
names.contains(&required),
"missing required tool: {required}"
);
}
}
#[tokio::test]
async fn test_initialize_response() {
let server = McpServer::new("http://127.0.0.1:1");
let r = Request {
jsonrpc: Some("2.0".into()),
id: Some(Value::from(1u64)),
method: "initialize".into(),
params: Some(serde_json::json!({
"protocolVersion": "2024-11-05",
"capabilities": {},
"clientInfo": { "name": "test", "version": "0.0.0" }
})),
};
let resp = server.dispatch(r).await;
assert!(resp.error.is_none(), "initialize must not error");
let result = resp.result.expect("expected result");
assert_eq!(result["protocolVersion"], "2024-11-05");
assert!(result["capabilities"].get("tools").is_some());
assert_eq!(result["serverInfo"]["name"], "trusty-search");
assert!(result["serverInfo"]["version"].is_string());
}
#[tokio::test]
async fn test_tools_list_response() {
let server = McpServer::new("http://127.0.0.1:1");
let resp = server.dispatch(req("tools/list", Value::Null)).await;
let result = resp.result.expect("expected result");
let tools = result
.get("tools")
.and_then(Value::as_array)
.expect("array");
let names: Vec<&str> = tools
.iter()
.filter_map(|t| t.get("name").and_then(Value::as_str))
.collect();
for required in [
"search_code",
"index_file",
"remove_file",
"list_indexes",
"create_index",
"search_health",
] {
assert!(
names.contains(&required),
"tools/list missing '{required}' (got {names:?})"
);
}
for t in tools {
assert!(t.get("name").is_some());
assert!(t.get("inputSchema").is_some());
}
}
#[tokio::test]
async fn test_unknown_method_returns_error() {
let server = McpServer::new("http://127.0.0.1:1");
let resp = server
.dispatch(req("definitely_not_a_method", Value::Null))
.await;
let err = resp.error.expect("expected error");
assert_eq!(err.code, error_codes::METHOD_NOT_FOUND);
}
#[tokio::test]
async fn notification_initialized_is_suppressed() {
let server = McpServer::new("http://127.0.0.1:1");
let r = Request {
jsonrpc: Some("2.0".into()),
id: None, method: "notifications/initialized".into(),
params: None,
};
let resp = server.dispatch(r).await;
assert!(resp.suppress, "notifications must be suppressed");
}
#[tokio::test]
async fn test_tools_list_complete() {
let server = McpServer::new("http://127.0.0.1:1");
let resp = server.dispatch(req("tools/list", Value::Null)).await;
let result = resp.result.expect("expected result");
let tools = result
.get("tools")
.and_then(Value::as_array)
.expect("array");
let names: Vec<&str> = tools
.iter()
.filter_map(|t| t.get("name").and_then(Value::as_str))
.collect();
for required in [
"search_code",
"index_file",
"remove_file",
"list_indexes",
"create_index",
"search_health",
"delete_index",
"reindex",
"index_status",
"list_chunks",
"chat",
"search_all",
] {
assert!(
names.contains(&required),
"tools/list missing '{required}' (got {names:?})"
);
}
}
#[tokio::test]
async fn search_all_missing_query_returns_invalid_params() {
let server = McpServer::new("http://127.0.0.1:1");
let resp = server
.dispatch(req("search_all", serde_json::json!({})))
.await;
let err = resp.error.expect("expected error");
assert_eq!(err.code, error_codes::INVALID_PARAMS);
}
#[tokio::test]
async fn tools_call_without_name_returns_invalid_params() {
let server = McpServer::new("http://127.0.0.1:1");
let resp = server
.dispatch(req("tools/call", serde_json::json!({})))
.await;
let err = resp.error.expect("expected error");
assert_eq!(err.code, error_codes::INVALID_PARAMS);
}
}