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

rpi_codegraph/
lib.rs
