semlith 0.5.0

//! A stdio MCP server, so an agent can query the store as a tool call.
//!
//! Deliberately hand-rolled: MCP over stdio is newline-delimited JSON-RPC 2.0,
//! and the three methods a tools-only server needs fit in one file with no
//! extra dependency.
//!
//! Everything written to stdout is protocol. Diagnostics go to stderr.

use crate::filter::Filter;
use crate::fleet::Fleet;
use anyhow::Result;
use serde_json::{Value, json};
use std::io::{BufRead, Write};

const PROTOCOL_VERSION: &str = "2024-11-05";

/// Read requests from `input` until EOF, answering on `output`.
pub fn serve(stores: &mut Fleet, input: impl BufRead, mut output: impl Write) -> Result<()> {
    stores.quiet = true;

    for line in input.lines() {
        let line = line?;
        if line.trim().is_empty() {
            continue;
        }

        let req: Value = match serde_json::from_str(&line) {
            Ok(v) => v,
            Err(e) => {
                respond(&mut output, &json!(null), Err((-32700, e.to_string())))?;
                continue;
            }
        };

        // Notifications carry no id and must not be answered.
        let Some(id) = req.get("id").cloned() else {
            continue;
        };
        let method = req.get("method").and_then(Value::as_str).unwrap_or("");
        let params = req.get("params").cloned().unwrap_or(json!({}));

        let result = dispatch(stores, method, &params);
        respond(&mut output, &id, result)?;
    }
    Ok(())
}

fn dispatch(stores: &mut Fleet, method: &str, params: &Value) -> Result<Value, (i64, String)> {
    match method {
        "initialize" => {
            let version = params
                .get("protocolVersion")
                .and_then(Value::as_str)
                .unwrap_or(PROTOCOL_VERSION);
            Ok(json!({
                "protocolVersion": version,
                "capabilities": { "tools": {} },
                "serverInfo": { "name": "semlith", "version": env!("CARGO_PKG_VERSION") },
            }))
        }
        "ping" => Ok(json!({})),
        "tools/list" => Ok(json!({ "tools": tools(stores) })),
        "tools/call" => call_tool(stores, params),
        other => Err((-32601, format!("unknown method: {other}"))),
    }
}

fn tools(stores: &Fleet) -> Value {
    // An agent cannot narrow to a store whose name it has never seen, so the
    // open stores are part of the tool description rather than something to
    // discover by trial.
    let open = stores.labels().join(", ");
    let store_arg = format!(
        "Restrict the search to these stores by name. Open stores: {open}. \
         Omit to search all of them, which is usually right — narrow only when \
         you already know which corpus holds the answer."
    );

    json!([
        {
            "name": "semlith_search",
            "description":
                "Semantic search over the local semlith store: docs, PDFs, code, and notes \
                 that have been indexed on this machine. Returns the most relevant excerpts \
                 with their file path and line range. Use this instead of reading whole files \
                 when you need to find where something is discussed or implemented. \
                 Optionally narrow the search to one part of the corpus with path, ext or lang.",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "What you are looking for. A question in plain English, or an exact identifier — both are searched."
                    },
                    "k": {
                        "type": "integer",
                        "description": "How many excerpts to return (default 8).",
                        "minimum": 1,
                        "maximum": 50
                    },
                    "path": {
                        "type": "array",
                        "items": { "type": "string" },
                        "description":
                            "Restrict the search to files matching these globs, e.g. \"src/**\". \
                             A pattern that does not start with / matches anywhere in the tree. \
                             `*` crosses directory separators, so `src/*` reaches the whole \
                             subtree. Only filter when you already know which part of the \
                             corpus holds the answer — a wrong guess hides it entirely."
                    },
                    "ext": {
                        "type": "array",
                        "items": { "type": "string" },
                        "description": "Restrict the search to these file extensions, e.g. [\"rs\", \"toml\"]."
                    },
                    "lang": {
                        "type": "array",
                        "items": { "type": "string" },
                        "description":
                            "Restrict the search to these languages, e.g. [\"rust\"]. Accepts: \
                             c, cpp, csharp, css, go, haskell, html, java, javascript, json, \
                             kotlin, lua, markdown, ocaml, php, python, ruby, rust, scala, \
                             shell, sql, swift, toml, typescript, yaml."
                    },
                    "store": {
                        "type": "array",
                        "items": { "type": "string" },
                        "description": store_arg
                    }
                },
                "required": ["query"]
            }
        },
        {
            "name": "semlith_stats",
            "description":
                "Report what the local semlith stores currently contain: file count, chunk \
                 count, indexed bytes, and the embedding model, one line per store. Use this \
                 to check whether a corpus is indexed before searching it, and to learn the \
                 store names that semlith_search's store argument accepts.",
            "inputSchema": { "type": "object", "properties": {} }
        }
    ])
}

fn call_tool(stores: &mut Fleet, params: &Value) -> Result<Value, (i64, String)> {
    let name = params.get("name").and_then(Value::as_str).unwrap_or("");
    let args = params.get("arguments").cloned().unwrap_or(json!({}));

    let body = match name {
        "semlith_search" => {
            let Some(query) = args.get("query").and_then(Value::as_str) else {
                return Err((-32602, "missing required argument: query".into()));
            };
            let k = args.get("k").and_then(Value::as_u64).unwrap_or(8) as usize;

            let filter = match Filter::new(
                &strings(&args, "path"),
                &strings(&args, "ext"),
                &strings(&args, "lang"),
            ) {
                Ok(f) => f,
                // An unknown language is the agent's mistake to correct, so it
                // goes back in-band with the list rather than as a protocol error.
                Err(e) => return Ok(tool_error(&e.to_string())),
            };

            // A store name the agent guessed is the agent's mistake to correct,
            // so it comes back in-band with the names that exist.
            let only = strings(&args, "store");

            // Told apart because an agent that scoped to the wrong subsystem
            // should widen the filter, not conclude the corpus is empty. A
            // failure here is a broken store, not an empty selection, and must
            // not be reported as one.
            let selected = if filter.is_empty() {
                1
            } else {
                match stores.matching_files(&filter) {
                    Ok(n) => n,
                    Err(e) => return Ok(tool_error(&format!("search failed: {e}"))),
                }
            };

            if selected == 0 {
                "No indexed file matches that path/ext/lang filter. Try again without it."
                    .to_string()
            } else {
                match stores.search_in(Some(&only), query, k.clamp(1, 50), &filter) {
                    Ok(hits) if hits.is_empty() => "No matches in the semlith store.".to_string(),
                    Ok(hits) => render(&hits),
                    // Tool failures are reported in-band so the agent can react,
                    // rather than as a protocol-level error.
                    Err(e) => return Ok(tool_error(&e.to_string())),
                }
            }
        }
        "semlith_stats" => {
            let many = stores.len() > 1;
            let mut lines = Vec::new();
            for (label, store) in stores.each() {
                let (files, chunks, bytes) = match store.stats() {
                    Ok(s) => s,
                    Err(e) => return Ok(tool_error(&format!("stats failed: {e}"))),
                };
                let body = format!(
                    "{files} files, {chunks} chunks, {} indexed, model {} ({} dim)",
                    crate::human_bytes(bytes),
                    store.model(),
                    store.dim(),
                );
                // One store answers exactly as it did before stores could be
                // combined; a name in front of it would only cost tokens.
                lines.push(if many {
                    format!("{label}: {body}")
                } else {
                    body
                });
            }
            lines.join("\n")
        }
        other => return Err((-32602, format!("unknown tool: {other}"))),
    };

    Ok(json!({ "content": [{ "type": "text", "text": body }] }))
}

/// One filter argument as strings. A bare string is accepted alongside an
/// array, because that is what an agent produces about half the time.
fn strings(args: &Value, key: &str) -> Vec<String> {
    match args.get(key) {
        Some(Value::String(s)) => vec![s.clone()],
        Some(Value::Array(items)) => items
            .iter()
            .filter_map(|i| i.as_str().map(str::to_string))
            .collect(),
        _ => Vec::new(),
    }
}

fn tool_error(message: &str) -> Value {
    json!({ "content": [{ "type": "text", "text": message }], "isError": true })
}

/// Compact, agent-readable: a locator line then the excerpt. Cheap to parse,
/// cheap in tokens, and the locator is enough to go read the real file.
fn render(hits: &[crate::Hit]) -> String {
    let mut out = String::new();
    for (i, h) in hits.iter().enumerate() {
        // The store, when there is more than one, goes in front of the path:
        // one short word, and without it an excerpt from the client is
        // indistinguishable from one from the service.
        let from = match &h.store {
            Some(label) => format!("{label} "),
            None => String::new(),
        };
        out.push_str(&format!(
            "[{}] {from}{}:{}-{} (score {:.3})\n{}\n\n",
            i + 1,
            h.path,
            h.start_line,
            h.end_line,
            h.score,
            h.text.trim_end()
        ));
    }
    out.trim_end().to_string()
}

fn respond(
    out: &mut impl Write,
    id: &Value,
    result: Result<Value, (i64, String)>,
) -> std::io::Result<()> {
    let msg = match result {
        Ok(value) => json!({ "jsonrpc": "2.0", "id": id, "result": value }),
        Err((code, message)) => {
            json!({ "jsonrpc": "2.0", "id": id, "error": { "code": code, "message": message } })
        }
    };
    writeln!(out, "{msg}")?;
    out.flush()
}