rpi_codegraph/
lib.rs

1use std::ffi::c_void;
2use std::io::{BufRead, BufReader, Read, Write};
3use std::path::{Path, PathBuf};
4use std::process::{Child, ChildStdin, Command, Stdio};
5use std::sync::atomic::{AtomicBool, Ordering};
6use std::sync::{mpsc, Arc, Mutex};
7use std::thread;
8use std::time::{Duration, Instant};
9
10use rpi_plugin_sdk::{
11    register_entrypoint_unified, FreeStringFn, PluginApi, StbString, StbStringRef, StepHandle,
12    StepResult, ToolPartialCb,
13};
14use serde_json::{json, Value};
15
16const SESSION_TIMEOUT: Duration = Duration::from_secs(20);
17const POLL_INTERVAL: Duration = Duration::from_millis(50);
18const MAX_OUTPUT_CHARS: usize = 25_000;
19const MAX_DIAGNOSTIC_CHARS: usize = 1_000;
20
21struct Drive {
22    tool: &'static str,
23    params: Value,
24    cancelled: AtomicBool,
25    completed: bool,
26}
27
28struct ToolSpec {
29    name: &'static str,
30    description: &'static str,
31    parameters: String,
32    execute: extern "C" fn(StbStringRef, StbString, Option<FreeStringFn>) -> StepHandle,
33}
34
35fn project_root(params: &Value) -> Result<PathBuf, String> {
36    let requested = params
37        .get("projectPath")
38        .and_then(Value::as_str)
39        .filter(|value| !value.trim().is_empty());
40    let root = match requested {
41        Some(value) => PathBuf::from(value.trim()),
42        None => std::env::current_dir()
43            .map_err(|err| format!("cannot resolve current directory: {err}"))?,
44    };
45    if !root.is_absolute() {
46        return Err("CodeGraph projectPath must be an absolute path".into());
47    }
48    if !root.is_dir() {
49        return Err(format!(
50            "CodeGraph projectPath is not an accessible directory: {}",
51            root.display()
52        ));
53    }
54    Ok(root)
55}
56
57#[cfg(windows)]
58fn spawn_server(root: &Path) -> Result<Child, String> {
59    // Command::new("codegraph") does not reliably resolve npm/Scoop .cmd shims
60    // on Windows. Let PowerShell discover the native application shim.
61    const SCRIPT: &str = concat!(
62        "& { param([string]$ProjectPath) ",
63        "$ErrorActionPreference = 'Stop'; ",
64        "$cmd = Get-Command codegraph -CommandType Application -ErrorAction Stop | Select-Object -First 1; ",
65        "if (-not $cmd) { throw 'codegraph command not found'; }; ",
66        "& $cmd.Source serve --mcp --path $ProjectPath; exit $LASTEXITCODE; }"
67    );
68    Command::new("powershell.exe")
69        .args([
70            "-NoProfile",
71            "-NonInteractive",
72            "-ExecutionPolicy",
73            "Bypass",
74            "-Command",
75            SCRIPT,
76        ])
77        .arg(root)
78        .current_dir(root)
79        .stdin(Stdio::piped())
80        .stdout(Stdio::piped())
81        .stderr(Stdio::piped())
82        .spawn()
83        .map_err(|err| format!("failed to start CodeGraph MCP server: {err}"))
84}
85
86#[cfg(not(windows))]
87fn spawn_server(root: &Path) -> Result<Child, String> {
88    Command::new("codegraph")
89        .args(["serve", "--mcp", "--path"])
90        .arg(root)
91        .current_dir(root)
92        .stdin(Stdio::piped())
93        .stdout(Stdio::piped())
94        .stderr(Stdio::piped())
95        .spawn()
96        .map_err(|err| {
97            format!("failed to start CodeGraph MCP server: {err}; install @colbymchenry/codegraph")
98        })
99}
100
101fn write_rpc(stdin: &mut ChildStdin, payload: &Value) -> Result<(), String> {
102    serde_json::to_writer(&mut *stdin, payload)
103        .map_err(|err| format!("failed to encode CodeGraph request: {err}"))?;
104    stdin
105        .write_all(b"\n")
106        .and_then(|_| stdin.flush())
107        .map_err(|err| format!("failed to write to CodeGraph MCP server: {err}"))
108}
109
110fn wait_for_response(
111    receiver: &mpsc::Receiver<String>,
112    id: u64,
113    deadline: Instant,
114    cancelled: &AtomicBool,
115) -> Result<Value, String> {
116    loop {
117        if cancelled.load(Ordering::SeqCst) {
118            return Err("codegraph request cancelled".into());
119        }
120        let now = Instant::now();
121        if now >= deadline {
122            return Err(format!(
123                "CodeGraph MCP session timed out after {} seconds; try `codegraph unlock`",
124                SESSION_TIMEOUT.as_secs()
125            ));
126        }
127        let wait = POLL_INTERVAL.min(deadline.saturating_duration_since(now));
128        match receiver.recv_timeout(wait) {
129            Ok(line) => {
130                let Ok(message) = serde_json::from_str::<Value>(&line) else {
131                    continue;
132                };
133                if message.get("id").and_then(Value::as_u64) != Some(id) {
134                    continue;
135                }
136                if let Some(error) = message.get("error") {
137                    let text = error
138                        .get("message")
139                        .and_then(Value::as_str)
140                        .map(str::to_owned)
141                        .unwrap_or_else(|| error.to_string());
142                    return Err(text);
143                }
144                return Ok(message.get("result").cloned().unwrap_or(Value::Null));
145            }
146            Err(mpsc::RecvTimeoutError::Timeout) => continue,
147            Err(mpsc::RecvTimeoutError::Disconnected) => {
148                return Err("CodeGraph MCP process closed before responding".into())
149            }
150        }
151    }
152}
153
154fn response_text(result: &Value) -> Result<String, String> {
155    let text = result
156        .get("content")
157        .and_then(Value::as_array)
158        .into_iter()
159        .flatten()
160        .filter(|part| part.get("type").and_then(Value::as_str) == Some("text"))
161        .filter_map(|part| part.get("text").and_then(Value::as_str))
162        .collect::<Vec<_>>()
163        .join("\n");
164    if result.get("isError").and_then(Value::as_bool) == Some(true) {
165        return Err(if text.is_empty() {
166            "CodeGraph tool failed".into()
167        } else {
168            text
169        });
170    }
171    let output = if text.is_empty() {
172        result.to_string()
173    } else {
174        text
175    };
176    Ok(truncate_chars(&output, MAX_OUTPUT_CHARS))
177}
178
179fn truncate_chars(text: &str, max_chars: usize) -> String {
180    if text.chars().count() <= max_chars {
181        return text.to_owned();
182    }
183    let byte_end = text
184        .char_indices()
185        .nth(max_chars)
186        .map(|(index, _)| index)
187        .unwrap_or(text.len());
188    let prefix = &text[..byte_end];
189    let cut = prefix
190        .rfind('\n')
191        .filter(|index| *index >= byte_end.saturating_mul(4) / 5)
192        .unwrap_or(byte_end);
193    format!(
194        "{}\n\n... (output truncated to {max_chars} characters)",
195        &text[..cut]
196    )
197}
198
199fn sanitize_diagnostic(text: &str) -> String {
200    let mut cleaned = String::with_capacity(text.len().min(MAX_DIAGNOSTIC_CHARS));
201    for raw in text.lines() {
202        let mut line = raw.replace('\u{1b}', "");
203        for marker in ["TOKEN=", "SECRET=", "PASSWORD=", "API_KEY=", "APIKEY="] {
204            if let Some(start) = line.to_ascii_uppercase().find(marker) {
205                let value_start = start + marker.len();
206                let value_end = line[value_start..]
207                    .find(char::is_whitespace)
208                    .map(|offset| value_start + offset)
209                    .unwrap_or(line.len());
210                line.replace_range(value_start..value_end, "[redacted]");
211            }
212        }
213        if let Some(start) = line.to_ascii_lowercase().find("bearer ") {
214            let value_start = start + "bearer ".len();
215            let value_end = line[value_start..]
216                .find(char::is_whitespace)
217                .map(|offset| value_start + offset)
218                .unwrap_or(line.len());
219            line.replace_range(value_start..value_end, "[redacted]");
220        }
221        if !cleaned.is_empty() {
222            cleaned.push('\n');
223        }
224        cleaned.push_str(&line);
225        if cleaned.chars().count() >= MAX_DIAGNOSTIC_CHARS {
226            return truncate_chars(&cleaned, MAX_DIAGNOSTIC_CHARS);
227        }
228    }
229    cleaned
230}
231
232fn call_codegraph(tool: &str, params: &Value, cancelled: &AtomicBool) -> Result<String, String> {
233    let root = project_root(params)?;
234    let mut child = spawn_server(&root)?;
235    let mut stdin = child
236        .stdin
237        .take()
238        .ok_or_else(|| "CodeGraph MCP stdin is unavailable".to_string())?;
239    let stdout = child
240        .stdout
241        .take()
242        .ok_or_else(|| "CodeGraph MCP stdout is unavailable".to_string())?;
243    let stderr = child
244        .stderr
245        .take()
246        .ok_or_else(|| "CodeGraph MCP stderr is unavailable".to_string())?;
247
248    let (sender, receiver) = mpsc::channel();
249    thread::spawn(move || {
250        for line in BufReader::new(stdout).lines().map_while(Result::ok) {
251            if sender.send(line).is_err() {
252                break;
253            }
254        }
255    });
256    let diagnostic = Arc::new(Mutex::new(String::new()));
257    let diagnostic_writer = Arc::clone(&diagnostic);
258    thread::spawn(move || {
259        let mut reader = BufReader::new(stderr).take(16 * 1024);
260        let mut text = String::new();
261        let _ = reader.read_to_string(&mut text);
262        if let Ok(mut destination) = diagnostic_writer.lock() {
263            *destination = text;
264        }
265    });
266
267    let deadline = Instant::now() + SESSION_TIMEOUT;
268    let operation = (|| {
269        let path = root.to_string_lossy().replace('\\', "/");
270        let root_uri = if cfg!(windows) {
271            format!("file:///{}", path.trim_start_matches('/'))
272        } else {
273            format!("file://{path}")
274        };
275        write_rpc(
276            &mut stdin,
277            &json!({
278                "jsonrpc": "2.0",
279                "id": 1,
280                "method": "initialize",
281                "params": {
282                    "protocolVersion": "2024-11-05",
283                    "rootUri": root_uri.clone(),
284                    "workspaceFolders": [{"uri": root_uri, "name": root.file_name().and_then(|v| v.to_str()).unwrap_or("project")}],
285                    "capabilities": {},
286                    "clientInfo": {"name": "rpi-codegraph", "version": env!("CARGO_PKG_VERSION")}
287                }
288            }),
289        )?;
290        wait_for_response(&receiver, 1, deadline, cancelled)?;
291        write_rpc(
292            &mut stdin,
293            &json!({"jsonrpc":"2.0","method":"initialized","params":{}}),
294        )?;
295        write_rpc(
296            &mut stdin,
297            &json!({"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":tool,"arguments":params}}),
298        )?;
299        let result = wait_for_response(&receiver, 2, deadline, cancelled)?;
300        response_text(&result)
301    })();
302
303    drop(stdin);
304    let _ = child.kill();
305    let _ = child.wait();
306
307    operation.map_err(|message| {
308        let stderr = diagnostic
309            .lock()
310            .ok()
311            .map(|value| sanitize_diagnostic(&value))
312            .unwrap_or_default();
313        if stderr.is_empty() || message.contains(&stderr) {
314            message
315        } else {
316            format!("{message}: {stderr}")
317        }
318    })
319}
320
321fn execute_tool(tool: &'static str, params: StbString, free: Option<FreeStringFn>) -> StepHandle {
322    let text = params.to_string_lossy();
323    params.free_with(free);
324    Box::into_raw(Box::new(Drive {
325        tool,
326        params: serde_json::from_str(&text).unwrap_or(Value::Null),
327        cancelled: AtomicBool::new(false),
328        completed: false,
329    })) as StepHandle
330}
331
332macro_rules! tool_execute {
333    ($function:ident, $tool:literal) => {
334        extern "C" fn $function(
335            _: StbStringRef,
336            params: StbString,
337            free: Option<FreeStringFn>,
338        ) -> StepHandle {
339            execute_tool($tool, params, free)
340        }
341    };
342}
343
344tool_execute!(execute_search, "codegraph_search");
345tool_execute!(execute_callers, "codegraph_callers");
346tool_execute!(execute_callees, "codegraph_callees");
347tool_execute!(execute_impact, "codegraph_impact");
348tool_execute!(execute_explore, "codegraph_explore");
349tool_execute!(execute_node, "codegraph_node");
350tool_execute!(execute_status, "codegraph_status");
351tool_execute!(execute_files, "codegraph_files");
352
353extern "C" fn poll(handle: StepHandle, _: Option<ToolPartialCb>, _: *mut c_void) -> StepResult {
354    if handle.is_null() {
355        return StepResult::err(StbString::from_string("null codegraph handle".into()));
356    }
357    let drive = unsafe { &mut *(handle as *mut Drive) };
358    if drive.cancelled.load(Ordering::SeqCst) {
359        return StepResult::err(StbString::from_string("codegraph request cancelled".into()));
360    }
361    if drive.completed {
362        return StepResult::err(StbString::from_string(
363            "codegraph polled after completion".into(),
364        ));
365    }
366    drive.completed = true;
367    match call_codegraph(drive.tool, &drive.params, &drive.cancelled) {
368        Ok(text) => StepResult::done(StbString::from_string(
369            json!({"content":[{"type":"text","text":text}]}).to_string(),
370        )),
371        Err(error) => StepResult::err(StbString::from_string(error)),
372    }
373}
374
375extern "C" fn cancel(handle: StepHandle) {
376    if !handle.is_null() {
377        unsafe { &*(handle as *mut Drive) }
378            .cancelled
379            .store(true, Ordering::SeqCst);
380    }
381}
382
383extern "C" fn destroy(handle: StepHandle) {
384    if !handle.is_null() {
385        unsafe { drop(Box::from_raw(handle as *mut Drive)) };
386    }
387}
388
389extern "C" fn free_string(value: StbString) {
390    if !value.is_empty() && !value.ptr.is_null() {
391        unsafe {
392            let slice = std::slice::from_raw_parts(value.ptr as *const u8, value.len);
393            let _ = Box::from_raw(slice as *const [u8] as *mut [u8]);
394        }
395    }
396}
397
398fn parameters(mut properties: serde_json::Map<String, Value>, required: &[&str]) -> String {
399    properties.insert(
400        "projectPath".into(),
401        json!({
402            "type": "string",
403            "description": "Absolute path to an initialized CodeGraph project. Defaults to the current directory."
404        }),
405    );
406    let mut schema = json!({"type":"object", "properties": properties});
407    if !required.is_empty() {
408        schema["required"] = json!(required);
409    }
410    schema.to_string()
411}
412
413fn tool_specs() -> Vec<ToolSpec> {
414    vec![
415        ToolSpec {
416            name: "codegraph_search",
417            description: "Search indexed code symbols by name. Returns compact locations only.",
418            parameters: parameters(json!({
419                "query":{"type":"string"},
420                "kind":{"type":"string","enum":["function","method","class","interface","type","variable","route","component"]},
421                "limit":{"type":"integer","default":10}
422            }).as_object().unwrap().clone(), &["query"]),
423            execute: execute_search,
424        },
425        ToolSpec {
426            name: "codegraph_callers",
427            description: "Find functions or methods that call an indexed symbol.",
428            parameters: parameters(json!({
429                "symbol":{"type":"string"}, "file":{"type":"string"},
430                "limit":{"type":"integer","default":20}
431            }).as_object().unwrap().clone(), &["symbol"]),
432            execute: execute_callers,
433        },
434        ToolSpec {
435            name: "codegraph_callees",
436            description: "Find functions or methods called by an indexed symbol.",
437            parameters: parameters(json!({
438                "symbol":{"type":"string"}, "file":{"type":"string"},
439                "limit":{"type":"integer","default":20}
440            }).as_object().unwrap().clone(), &["symbol"]),
441            execute: execute_callees,
442        },
443        ToolSpec {
444            name: "codegraph_impact",
445            description: "Analyze the dependency impact radius of changing an indexed symbol.",
446            parameters: parameters(json!({
447                "symbol":{"type":"string"}, "file":{"type":"string"},
448                "depth":{"type":"integer","default":2}
449            }).as_object().unwrap().clone(), &["symbol"]),
450            execute: execute_impact,
451        },
452        ToolSpec {
453            name: "codegraph_explore",
454            description: "Explore a code area by question, symbol, file, or flow; returns bounded relevant source grouped by file.",
455            parameters: parameters(json!({
456                "query":{"type":"string"}, "maxFiles":{"type":"integer","default":8}
457            }).as_object().unwrap().clone(), &["query"]),
458            execute: execute_explore,
459        },
460        ToolSpec {
461            name: "codegraph_node",
462            description: "Get one symbol with source and relationships, or read one indexed file with optional line bounds.",
463            parameters: parameters(json!({
464                "symbol":{"type":"string"}, "includeCode":{"type":"boolean","default":false},
465                "file":{"type":"string"}, "offset":{"type":"integer"}, "limit":{"type":"integer"},
466                "symbolsOnly":{"type":"boolean","default":false}, "line":{"type":"integer"}
467            }).as_object().unwrap().clone(), &[]),
468            execute: execute_node,
469        },
470        ToolSpec {
471            name: "codegraph_status",
472            description: "Report CodeGraph index health, counts, and pending synchronization.",
473            parameters: parameters(serde_json::Map::new(), &[]),
474            execute: execute_status,
475        },
476        ToolSpec {
477            name: "codegraph_files",
478            description: "Get the indexed project file tree with optional path, glob, format, metadata, and depth filters.",
479            parameters: parameters(json!({
480                "path":{"type":"string"}, "pattern":{"type":"string"},
481                "format":{"type":"string","enum":["tree","flat","grouped"],"default":"tree"},
482                "includeMetadata":{"type":"boolean","default":true}, "maxDepth":{"type":"integer"}
483            }).as_object().unwrap().clone(), &[]),
484            execute: execute_files,
485        },
486    ]
487}
488
489#[no_mangle]
490pub extern "C" fn rpi_plugin_register(api: *const PluginApi) -> i32 {
491    unsafe { register_entrypoint_unified(api, |api| {
492        let Some(register) = api.register_tool else {
493            return 1;
494        };
495        for spec in tool_specs() {
496            let schema = Box::new(rpi_plugin_sdk::StableToolSchema {
497                name: StbString::from_string(spec.name.into()),
498                description: StbString::from_string(spec.description.into()),
499                parameters: StbString::from_string(spec.parameters),
500            });
501            let rc = register(&*schema, spec.execute, poll, cancel, destroy, free_string);
502            drop(schema);
503            if rc != 0 {
504                return rc;
505            }
506        }
507        0
508    }) }
509}
510
511#[cfg(test)]
512mod tests {
513    use super::*;
514
515    #[test]
516    fn schemas_are_valid_json() {
517        for spec in tool_specs() {
518            serde_json::from_str::<Value>(&spec.parameters)
519                .unwrap_or_else(|error| panic!("invalid schema for {}: {error}", spec.name));
520        }
521    }
522
523    #[test]
524    fn extracts_text_and_rejects_mcp_errors() {
525        let result = json!({"content":[{"type":"text","text":"one"},{"type":"image","data":"x"},{"type":"text","text":"two"}]});
526        assert_eq!(response_text(&result).unwrap(), "one\ntwo");
527        assert!(
528            response_text(&json!({"isError":true,"content":[{"type":"text","text":"bad"}]}))
529                .is_err()
530        );
531    }
532
533    #[test]
534    fn truncates_on_unicode_boundaries() {
535        let text = "你".repeat(30);
536        let output = truncate_chars(&text, 10);
537        assert!(output.starts_with(&"你".repeat(10)));
538        assert!(output.contains("output truncated"));
539    }
540
541    #[test]
542    fn redacts_common_secrets() {
543        let output = sanitize_diagnostic("failed TOKEN=abc Bearer xyz API_KEY=hidden");
544        assert!(!output.contains("abc"));
545        assert!(!output.contains("xyz"));
546        assert!(!output.contains("hidden"));
547    }
548}
rpi_codegraph/lib.rs

rpi_codegraph/
lib.rs
