rpi_codegraph/
lib.rs

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

rpi_codegraph/
lib.rs
