rgx/
mcp.rs

1//! Minimal MCP stdio server exposing rgx's search to AI agents.
2//!
3//! Speaks JSON-RPC 2.0 over stdio (newline-delimited), implementing the handshake plus three tools:
4//! `content_search`, `file_search`, and `status`. Results come back as ripgrep-style text (the shape
5//! models already know), per `docs/mcp.md`. Parsing uses `serde_json` so UTF-8, escapes, and key
6//! order are handled correctly.
7
8use std::io::{BufRead, Write};
9use std::path::Path;
10use std::path::PathBuf;
11
12use anyhow::Result;
13use serde_json::{Value, json};
14
15use crate::client;
16use crate::compact::{self, CompactOpts};
17use crate::confirm::SearchOptions;
18use crate::cursor::{self, Mode};
19use crate::proto::Request;
20
21const PROTOCOL_VERSION: &str = "2024-11-05";
22
23/// Run the MCP server rooted at `root` until stdin closes.
24pub fn run(root: PathBuf) -> Result<()> {
25    let stdin = std::io::stdin();
26    let mut stdout = std::io::stdout();
27    let mut line = String::new();
28    loop {
29        line.clear();
30        if stdin.lock().read_line(&mut line)? == 0 {
31            break;
32        }
33        if line.trim().is_empty() {
34            continue;
35        }
36        if let Some(resp) = handle_message(line.trim(), &root) {
37            writeln!(stdout, "{resp}")?;
38            stdout.flush()?;
39        }
40    }
41    Ok(())
42}
43
44/// Handle one JSON-RPC message; returns a response line, or `None` for notifications / unparseable
45/// input.
46fn handle_message(msg: &str, root: &Path) -> Option<String> {
47    let v: Value = serde_json::from_str(msg).ok()?;
48    let id = v.get("id").cloned().unwrap_or(Value::Null);
49    let method = v.get("method")?.as_str()?;
50    match method {
51        "initialize" => Some(result(
52            id,
53            json!({
54                "protocolVersion": PROTOCOL_VERSION,
55                "capabilities": {"tools": {}},
56                "serverInfo": {"name": "rgx", "version": env!("CARGO_PKG_VERSION")},
57            }),
58        )),
59        "tools/list" => Some(result(id, tools())),
60        "tools/call" => Some(handle_tool_call(id, &v, root)),
61        m if m.starts_with("notifications/") => None,
62        _ => {
63            if id.is_null() {
64                None
65            } else {
66                Some(error(id, -32601, "method not found"))
67            }
68        }
69    }
70}
71
72fn handle_tool_call(id: Value, msg: &Value, root: &Path) -> String {
73    let params = msg.get("params");
74    let name = params.and_then(|p| p.get("name")).and_then(Value::as_str);
75    let args = params.and_then(|p| p.get("arguments"));
76    match name {
77        Some("content_search") => {
78            // A cursor carries the entire query + resume position, so it supersedes the other args.
79            let query = if let Some(tok) = arg_str(args, "cursor") {
80                let blob = match client::take_cursor(root, tok) {
81                    Ok(Some(blob)) => blob,
82                    Ok(None) => {
83                        return error(id, -32602, "pagination expired — re-run the search");
84                    }
85                    Err(e) => return error(id, -32603, &format!("{e}")),
86                };
87                match cursor::decode(&blob) {
88                    Ok(c) => Query {
89                        start_after: c.last_path.clone().map(|p| (p, c.last_lineno)),
90                        prev: Some((c.prev_total, c.fingerprint)),
91                        pattern: c.pattern,
92                        opts: c.opts,
93                        mode: c.mode,
94                        page_size: c.page_size,
95                    },
96                    Err(e) => return error(id, -32602, &format!("invalid cursor: {e}")),
97                }
98            } else {
99                let Some(pattern) = arg_str(args, "pattern") else {
100                    return error(id, -32602, "missing required argument 'pattern'");
101                };
102                Query {
103                    pattern: pattern.to_string(),
104                    opts: SearchOptions {
105                        case_insensitive: arg_bool(args, "case_insensitive"),
106                        word: arg_bool(args, "word"),
107                        fixed_strings: arg_bool(args, "fixed_strings"),
108                        multi_line: arg_bool(args, "multi_line"),
109                        ..Default::default()
110                    },
111                    mode: if arg_bool(args, "count") {
112                        Mode::Count
113                    } else if arg_bool(args, "files_only") {
114                        Mode::Files
115                    } else {
116                        Mode::Matches
117                    },
118                    start_after: None,
119                    page_size: arg_usize(args, "page_size").unwrap_or(compact::DEFAULT_PAGE_SIZE),
120                    prev: None,
121                }
122            };
123            tool_result(id, &compact_search(root, query))
124        }
125        Some("file_search") => {
126            let Some(query) = arg_str(args, "query") else {
127                return error(id, -32602, "missing required argument 'query'");
128            };
129            let limit = arg_usize(args, "limit").unwrap_or(200) as u32;
130            let after = arg_str(args, "after").map(str::to_string);
131            tool_result(id, &file_search(root, query, after, limit))
132        }
133        Some("status") => tool_result(id, &run_request(root, &Request::Status)),
134        Some(other) => error(id, -32602, &format!("unknown tool {other:?}")),
135        None => error(id, -32602, "missing tool name"),
136    }
137}
138
139/// A resolved content query: from explicit args, or unpacked from a `cursor`.
140struct Query {
141    pattern: String,
142    opts: SearchOptions,
143    mode: Mode,
144    start_after: Option<(String, u64)>,
145    page_size: usize,
146    /// `(total, fingerprint)` when resuming a cursor, for the staleness note.
147    prev: Option<(usize, u32)>,
148}
149
150fn arg_bool(args: Option<&Value>, key: &str) -> bool {
151    args.and_then(|a| a.get(key))
152        .and_then(Value::as_bool)
153        .unwrap_or(false)
154}
155
156fn arg_usize(args: Option<&Value>, key: &str) -> Option<usize> {
157    args.and_then(|a| a.get(key))
158        .and_then(Value::as_u64)
159        .map(|n| n as usize)
160}
161
162fn arg_str<'a>(args: Option<&'a Value>, key: &str) -> Option<&'a str> {
163    args.and_then(|a| a.get(key)).and_then(Value::as_str)
164}
165
166/// Run a content search and return the token-savings view: results grouped by file and paged, with a
167/// cursor to fetch the next page. Matching is identical to `rg`; only presentation differs (see
168/// `compact`). Paging is cheap (warm index), so an agent pulls more on demand rather than dumping all.
169fn compact_search(root: &Path, q: Query) -> String {
170    let raw = match crate::collect_search(root, &q.pattern, q.opts) {
171        Ok(b) => b,
172        Err(e) => return format!("error: {e}"),
173    };
174    let p = compact::format(
175        &raw,
176        &q.pattern,
177        q.opts,
178        CompactOpts {
179            mode: q.mode,
180            start_after: q.start_after,
181            page_size: q.page_size,
182            max_cols: compact::DEFAULT_MAX_COLS,
183        },
184    );
185    let mut text = format!("{}\n{}", p.header, p.body);
186    if let Some(note) = p.staleness_note(q.prev) {
187        text.push_str(&format!("\nnote: {note}"));
188    }
189    // root_hint is None: the MCP server root is authoritative, so the cursor never carries a path.
190    // The blob is parked in this root's daemon; the agent echoes back the short token. On a store
191    // failure, still tell the agent more remains so it can't mistake a partial page for the whole.
192    if let Some(next) = p.next_cursor(q.mode, q.pattern, q.opts, q.page_size, None) {
193        match client::store_cursor(root, cursor::encode(&next)) {
194            Ok(token) => text.push_str(&format!(
195                "\n(more: call content_search with cursor: \"{token}\")"
196            )),
197            Err(e) => text.push_str(&format!(
198                "\nnote: more results exist but the pagination cursor could not be stored ({e}); \
199                 re-run the search"
200            )),
201        }
202    }
203    text
204}
205
206/// File-name search returning the token-savings shape: a `[files X-Y of N]` header, one path per
207/// line, and a hint to fetch more (the daemon reports the true total and the keyset resume key).
208fn file_search(root: &Path, query: &str, after: Option<String>, limit: u32) -> String {
209    let bytes = match client::request(
210        root,
211        &Request::Find {
212            needle: query.to_string(),
213            after,
214            limit,
215        },
216    ) {
217        Ok(b) => b,
218        Err(e) => return format!("error: {e}"),
219    };
220    let (header, body) = crate::proto::parse_find_header(&bytes);
221    let body = String::from_utf8_lossy(body);
222    let Some(h) = header else {
223        return body.into_owned();
224    };
225    let first = if h.returned == 0 { 0 } else { h.start + 1 };
226    let mut text = format!(
227        "[files {first}-{} of {}]\n{body}",
228        h.start + h.returned,
229        h.total
230    );
231    if let Some(next) = h.next_after {
232        text.push_str(&format!(
233            "\n(more: call file_search with after: \"{next}\")"
234        ));
235    }
236    text
237}
238
239fn run_request(root: &Path, req: &Request) -> String {
240    match client::request(root, req) {
241        Ok(bytes) => String::from_utf8_lossy(&bytes).into_owned(),
242        Err(e) => format!("error: {e}"),
243    }
244}
245
246fn tools() -> Value {
247    json!({"tools": [
248        {
249            "name": "content_search",
250            "description": concat!(
251                "Search file contents with a regex (ripgrep semantics, accelerated by an index). ",
252                "Results are grouped by file and paged: the match set is identical to ripgrep, nothing ",
253                "is dropped. The header reports the total match/file count, so you know how much you ",
254                "have NOT seen; when more remains, fetch it by passing the opaque `cursor` from the ",
255                "response (it carries the exact same query, so the next page can't drift). Paging is ",
256                "cheap (the index is warm). For a quick sense of scope, use `files_only` (paths only) ",
257                "or `count` (per-file counts) instead of a page-walk. Long lines are trimmed around ",
258                "the match (read the file for the full line)."
259            ),
260            "inputSchema": {
261                "type": "object",
262                "properties": {
263                    "pattern": {"type": "string", "description": "required for a new search; omit when paging via cursor"},
264                    "case_insensitive": {"type": "boolean"},
265                    "word": {"type": "boolean", "description": "match only whole words (-w)"},
266                    "fixed_strings": {"type": "boolean", "description": "treat pattern as a literal (-F)"},
267                    "multi_line": {"type": "boolean"},
268                    "files_only": {"type": "boolean", "description": "list matching file paths only (-l)"},
269                    "count": {"type": "boolean", "description": "per-file match counts only (-c)"},
270                    "page_size": {"type": "integer", "description": "matches (or files, for -l/-c) per page; default 50"},
271                    "cursor": {"type": "string", "description": "opaque token from a previous response; fetches the next page and supersedes all other args"}
272                }
273            }
274        },
275        {
276            "name": "file_search",
277            "description": concat!(
278                "Find files/directories by name or path substring (fd/find-style). Returns a header ",
279                "with the true total, then one path per line; when more remain than the page holds, ",
280                "the response gives an `after` key to fetch the next page."
281            ),
282            "inputSchema": {
283                "type": "object",
284                "properties": {
285                    "query": {"type": "string"},
286                    "limit": {"type": "integer", "description": "max paths per page; default 200"},
287                    "after": {"type": "string", "description": "resume key from a previous response (keyset paging)"}
288                },
289                "required": ["query"]
290            }
291        },
292        {
293            "name": "status",
294            "description": "Report index health: whether it is ready, file and trigram counts.",
295            "inputSchema": {"type": "object", "properties": {}}
296        }
297    ]})
298}
299
300fn result(id: Value, result: Value) -> String {
301    json!({"jsonrpc": "2.0", "id": id, "result": result}).to_string()
302}
303
304fn error(id: Value, code: i32, message: &str) -> String {
305    json!({"jsonrpc": "2.0", "id": id, "error": {"code": code, "message": message}}).to_string()
306}
307
308fn tool_result(id: Value, text: &str) -> String {
309    json!({"jsonrpc": "2.0", "id": id, "result": {"content": [{"type": "text", "text": text}]}})
310        .to_string()
311}
312
313#[cfg(test)]
314mod tests {
315    use super::*;
316
317    #[test]
318    fn initialize_and_unicode_pattern_parse() {
319        let resp = handle_message(
320            r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}"#,
321            Path::new("."),
322        )
323        .unwrap();
324        assert!(resp.contains("\"protocolVersion\""));
325        // A non-ASCII pattern must round-trip through the parser unmolested.
326        let v: Value = serde_json::from_str(
327            r#"{"params":{"name":"content_search","arguments":{"pattern":"café"}}}"#,
328        )
329        .unwrap();
330        let pat = v["params"]["arguments"]["pattern"].as_str().unwrap();
331        assert_eq!(pat, "café");
332    }
333
334    #[test]
335    fn content_search_advertises_cursor_paging() {
336        let listed = tools().to_string();
337        assert!(listed.contains("content_search"));
338        assert!(listed.contains("\"cursor\""));
339        assert!(listed.contains("\"page_size\""));
340        assert!(!listed.contains("\"page\""));
341    }
342
343    #[test]
344    fn file_search_advertises_keyset_paging() {
345        let listed = tools().to_string();
346        assert!(listed.contains("file_search"));
347        assert!(listed.contains("\"after\""));
348        assert!(listed.contains("\"limit\""));
349    }
350
351    #[test]
352    fn notifications_get_no_response() {
353        assert!(
354            handle_message(
355                r#"{"jsonrpc":"2.0","method":"notifications/initialized"}"#,
356                Path::new(".")
357            )
358            .is_none()
359        );
360    }
361}
rgx/mcp.rs

rgx/
mcp.rs
