Skip to main content

weavatrix_rust/mcp/
server.rs

1use crate::{Analyzer, Weavatrix, tools};
2use mcport::{ServerIdentity, ToolReply, ToolServer, Value};
3use notify::{EventKind, RecursiveMode, Watcher};
4use std::collections::BTreeSet;
5use std::fmt::{Display, Formatter};
6use std::io;
7use std::path::{Path, PathBuf};
8use std::sync::mpsc::{self, Receiver, TryRecvError};
9
10const DERIVED_DIRECTORIES: &[&str] = &[
11    ".git",
12    ".weavatrix",
13    ".codegraph",
14    ".next",
15    ".nuxt",
16    ".svelte-kit",
17    ".turbo",
18    ".venv",
19    "build",
20    "coverage",
21    "dist",
22    "node_modules",
23    "target",
24    "vendor",
25];
26
27#[derive(Debug)]
28pub enum McpError {
29    Io(io::Error),
30    Repository(crate::Error),
31}
32
33impl Display for McpError {
34    fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
35        match self {
36            Self::Io(error) => write!(formatter, "MCP I/O failed: {error}"),
37            Self::Repository(error) => {
38                write!(formatter, "repository initialization failed: {error}")
39            }
40        }
41    }
42}
43
44impl std::error::Error for McpError {}
45
46impl From<io::Error> for McpError {
47    fn from(value: io::Error) -> Self {
48        Self::Io(value)
49    }
50}
51
52impl From<crate::Error> for McpError {
53    fn from(value: crate::Error) -> Self {
54        Self::Repository(value)
55    }
56}
57
58/// Weavatrix tool surface behind the shared `mcport` stdio runtime.
59///
60/// Graph construction finishes before the handshake so startup does not split
61/// CPU between a protocol thread and an analyzer thread. The first tool call
62/// runs an incremental catch-up scan before using it, then starts the
63/// filesystem watcher in the background. This keeps the full cold boundary
64/// deterministic without returning stale evidence or adding an async executor.
65struct WeavatrixServer {
66    profile: super::McpProfile,
67    identity: ServerIdentity,
68    catalog: Value,
69    tool_names: BTreeSet<String>,
70    root: PathBuf,
71    engine: Option<Weavatrix>,
72    first_tool_call: bool,
73    watcher: WatcherState,
74}
75
76impl WeavatrixServer {
77    fn new(root: PathBuf, profile: super::McpProfile) -> Result<Self, McpError> {
78        let definitions = tools::catalog_for_profile(profile);
79        let tool_names = definitions
80            .iter()
81            .map(|definition| definition.name.to_owned())
82            .collect();
83        let catalog = blazingly_json::to_value(definitions)
84            .map_err(|error| McpError::Io(io::Error::other(error)))?;
85        let identity = ServerIdentity::new(
86            "weavatrix-rust",
87            env!("CARGO_PKG_VERSION"),
88            "Local read-only repository intelligence. Inferred evidence is explicitly labelled.",
89        );
90        let engine = Weavatrix::open(&root)?;
91        Ok(Self {
92            profile,
93            identity,
94            catalog,
95            tool_names,
96            root,
97            engine: Some(engine),
98            first_tool_call: true,
99            watcher: WatcherState::NotStarted,
100        })
101    }
102
103    fn engine(&mut self) -> Result<&mut Weavatrix, crate::Error> {
104        if self.engine.is_none() {
105            let engine = Weavatrix::open(&self.root)?;
106            engine.state().prime_weak_components();
107            self.engine = Some(engine);
108        }
109        let Some(engine) = self.engine.as_mut() else {
110            return Err(crate::Error::InvalidRepository(self.root.clone()));
111        };
112        Ok(engine)
113    }
114
115    fn catch_up_graph(&mut self) -> Result<(), String> {
116        let engine = self
117            .engine
118            .as_mut()
119            .ok_or_else(|| "repository graph is not initialized".to_owned())?;
120        if engine
121            .refresh_if_stale()
122            .map_err(|error| format!("repository refresh failed: {error}"))?
123        {
124            engine.state().prime_weak_components();
125        }
126        Ok(())
127    }
128
129    fn start_watcher(&mut self) -> io::Result<()> {
130        let root = self.root.clone();
131        let (sender, receiver) = mpsc::sync_channel(1);
132        std::thread::Builder::new()
133            .name("weavatrix-watch-init".to_owned())
134            .spawn(move || {
135                let _ = sender.send(RepositoryWatcher::new(&root));
136            })?;
137        self.watcher = WatcherState::Starting(receiver);
138        Ok(())
139    }
140
141    fn refresh_before_call(&mut self) -> Result<(), String> {
142        let state = std::mem::replace(&mut self.watcher, WatcherState::NotStarted);
143        let (watcher, catch_up) = match state {
144            WatcherState::NotStarted => {
145                let watcher = RepositoryWatcher::new(&self.root)
146                    .map_err(|error| format!("repository watcher failed: {error}"))?;
147                (watcher, true)
148            }
149            WatcherState::Starting(receiver) => {
150                let watcher = receiver
151                    .recv()
152                    .map_err(|_| "repository watcher startup disconnected".to_owned())?
153                    .map_err(|error| format!("repository watcher failed: {error}"))?;
154                (watcher, true)
155            }
156            WatcherState::Ready(watcher) => (watcher, false),
157        };
158        self.watcher = WatcherState::Ready(watcher);
159
160        let queued_change = self
161            .ready_watcher()
162            .map_err(|error| format!("repository watcher failed: {error}"))?
163            .changed()
164            .map_err(|error| format!("repository watcher failed: {error}"))?;
165        if !catch_up && !queued_change {
166            return Ok(());
167        }
168
169        let engine = self
170            .engine
171            .as_mut()
172            .ok_or_else(|| "repository graph is not initialized".to_owned())?;
173        if engine
174            .refresh_if_stale()
175            .map_err(|error| format!("repository refresh failed: {error}"))?
176        {
177            engine.state().prime_weak_components();
178        }
179
180        // Changes made while the catch-up scan was running were already
181        // registered by the watcher. Apply one more incremental scan when
182        // needed so the tool never observes the pre-change graph.
183        if self
184            .ready_watcher()
185            .map_err(|error| format!("repository watcher failed: {error}"))?
186            .changed()
187            .map_err(|error| format!("repository watcher failed: {error}"))?
188        {
189            let engine = self
190                .engine
191                .as_mut()
192                .ok_or_else(|| "repository graph is not initialized".to_owned())?;
193            if engine
194                .refresh_if_stale()
195                .map_err(|error| format!("repository refresh failed: {error}"))?
196            {
197                engine.state().prime_weak_components();
198            }
199        }
200        Ok(())
201    }
202
203    fn ready_watcher(&self) -> io::Result<&RepositoryWatcher> {
204        match &self.watcher {
205            WatcherState::Ready(watcher) => Ok(watcher),
206            WatcherState::NotStarted | WatcherState::Starting(_) => Err(io::Error::new(
207                io::ErrorKind::NotConnected,
208                "repository watcher is not ready",
209            )),
210        }
211    }
212}
213
214enum WatcherState {
215    NotStarted,
216    Starting(Receiver<io::Result<RepositoryWatcher>>),
217    Ready(RepositoryWatcher),
218}
219
220struct RepositoryWatcher {
221    root: PathBuf,
222    _watcher: notify::RecommendedWatcher,
223    events: Receiver<notify::Result<notify::Event>>,
224}
225
226impl RepositoryWatcher {
227    fn new(root: &Path) -> io::Result<Self> {
228        let root = root.canonicalize()?;
229        let (sender, events) = mpsc::channel();
230        let mut watcher = notify::recommended_watcher(move |event| {
231            let _ = sender.send(event);
232        })
233        .map_err(io::Error::other)?;
234        watcher
235            .watch(&root, RecursiveMode::Recursive)
236            .map_err(io::Error::other)?;
237        Ok(Self {
238            root,
239            _watcher: watcher,
240            events,
241        })
242    }
243
244    fn changed(&self) -> io::Result<bool> {
245        let mut changed = false;
246        loop {
247            match self.events.try_recv() {
248                Ok(Ok(event)) => {
249                    if !matches!(event.kind, EventKind::Access(_))
250                        && event
251                            .paths
252                            .iter()
253                            .any(|path| analysis_input_changed(&self.root, path))
254                    {
255                        changed = true;
256                    }
257                }
258                Ok(Err(error)) => return Err(io::Error::other(error)),
259                Err(TryRecvError::Empty) => return Ok(changed),
260                Err(TryRecvError::Disconnected) => {
261                    return Err(io::Error::new(
262                        io::ErrorKind::BrokenPipe,
263                        "repository filesystem watcher disconnected",
264                    ));
265                }
266            }
267        }
268    }
269}
270
271fn analysis_input_changed(root: &Path, path: &Path) -> bool {
272    let Ok(relative) = path.strip_prefix(root) else {
273        return false;
274    };
275    let normalized = relative.to_string_lossy().replace('\\', "/");
276    let lower = normalized.to_ascii_lowercase();
277    let file_name = relative
278        .file_name()
279        .and_then(|name| name.to_str())
280        .unwrap_or_default()
281        .to_ascii_lowercase();
282
283    if matches!(
284        file_name.as_str(),
285        ".gitignore" | ".ignore" | ".weavatrixignore"
286    ) || matches!(lower.as_str(), ".git/config" | ".git/info/exclude")
287    {
288        return true;
289    }
290    if lower
291        .split('/')
292        .any(|component| DERIVED_DIRECTORIES.contains(&component))
293    {
294        return false;
295    }
296    Analyzer::default().supports_path(&normalized)
297}
298
299impl ToolServer for WeavatrixServer {
300    fn identity(&self) -> ServerIdentity {
301        self.identity.clone()
302    }
303
304    fn identity_ref(&self) -> Option<&ServerIdentity> {
305        Some(&self.identity)
306    }
307
308    fn catalog(&mut self) -> Value {
309        self.catalog.clone()
310    }
311
312    fn catalog_ref(&mut self) -> Option<&Value> {
313        Some(&self.catalog)
314    }
315
316    fn has_tool(&self, name: &str) -> Option<bool> {
317        Some(self.tool_names.contains(name))
318    }
319
320    fn call(&mut self, name: &str, arguments: Value) -> ToolReply {
321        if !self.profile.allows(name) {
322            return ToolReply::error(format!(
323                "tool {name} is unavailable in the {:?} profile",
324                self.profile
325            ));
326        }
327        let graph_was_loaded = self.engine.is_some();
328        let first_tool_call = self.first_tool_call;
329        if first_tool_call {
330            if let Err(error) = self.catch_up_graph() {
331                return ToolReply::error(error);
332            }
333        } else if graph_was_loaded
334            && !matches!(name, "rebuild_graph" | "open_repo")
335            && let Err(error) = self.refresh_before_call()
336        {
337            return ToolReply::error(error);
338        }
339        let structured = arguments
340            .get("output_format")
341            .and_then(Value::as_str)
342            .is_none_or(|format| format == "json");
343        let (reply, opened_root) = {
344            let engine = match self.engine() {
345                Ok(engine) => engine,
346                Err(error) => {
347                    return ToolReply::error(format!("repository initialization failed: {error}"));
348                }
349            };
350            match tools::call(engine, name, arguments) {
351                Ok(value) => {
352                    let opened_root =
353                        (name == "open_repo").then(|| engine.state().root().to_path_buf());
354                    (ToolReply::Success { value, structured }, opened_root)
355                }
356                Err(error) => (ToolReply::error(error), None),
357            }
358        };
359        let opened_repository = opened_root.is_some();
360        if let Some(root) = opened_root {
361            self.root = root;
362        }
363        if (first_tool_call || !graph_was_loaded || opened_repository)
364            && self.engine.is_some()
365            && let Err(error) = self.start_watcher()
366        {
367            return ToolReply::error(format!("repository watcher startup failed: {error}"));
368        }
369        self.first_tool_call = false;
370        reply
371    }
372}
373
374/// Serves the read-only Weavatrix tool catalog over MCP stdio.
375///
376/// # Errors
377///
378/// Returns stdio failures or a missing repository root. Invalid requests are
379/// returned as JSON-RPC errors and do not terminate the server.
380pub fn serve(root: impl AsRef<Path>) -> Result<(), McpError> {
381    serve_with_profile(root, super::McpProfile::All)
382}
383
384/// Serves one capability profile over the same read-only MCP runtime.
385///
386/// The repository root and graph are validated eagerly so misconfiguration
387/// fails before the protocol handshake. The first tool call performs an
388/// incremental catch-up scan, then later calls use filesystem events.
389///
390/// # Errors
391///
392/// Returns stdio failures or a missing repository root.
393pub fn serve_with_profile(
394    root: impl AsRef<Path>,
395    profile: super::McpProfile,
396) -> Result<(), McpError> {
397    let root = root.as_ref().to_path_buf();
398    if !root.is_dir() {
399        return Err(McpError::Io(io::Error::new(
400            io::ErrorKind::NotFound,
401            format!("repository root {} is not a directory", root.display()),
402        )));
403    }
404    let mut server = WeavatrixServer::new(root, profile)?;
405    mcport::serve(&mut server).map_err(McpError::Io)
406}
407
408#[cfg(test)]
409mod tests {
410    use super::{WatcherState, WeavatrixServer};
411    // The request the runtime dispatches is built with the runtime's own JSON
412    // type, so the test exercises the boundary rather than bypassing it.
413    use mcport::{MODERN_PROTOCOL_VERSION, dispatch, json};
414    use std::path::PathBuf;
415
416    fn server(profile: crate::mcp::McpProfile) -> WeavatrixServer {
417        WeavatrixServer::new(PathBuf::from(env!("CARGO_MANIFEST_DIR")), profile).unwrap()
418    }
419
420    #[test]
421    #[allow(clippy::too_many_lines)]
422    fn negotiates_lists_and_calls_tools() {
423        let mut engine = server(crate::mcp::McpProfile::All);
424        let initialized = dispatch(
425            &mut engine,
426            &json!({
427                "jsonrpc": "2.0",
428                "id": 1,
429                "method": "initialize",
430                "params": {"protocolVersion": "2025-06-18"}
431            }),
432        )
433        .expect("initialize is answered");
434        assert_eq!(initialized["result"]["protocolVersion"], "2025-06-18");
435        assert_eq!(
436            initialized["result"]["serverInfo"]["name"],
437            "weavatrix-rust"
438        );
439        assert!(
440            engine.engine.is_some() && matches!(engine.watcher, WatcherState::NotStarted),
441            "initialize must use the ready graph without starting the watcher"
442        );
443
444        let modern_meta = json!({
445            "io.modelcontextprotocol/protocolVersion": MODERN_PROTOCOL_VERSION,
446            "io.modelcontextprotocol/clientInfo": {
447                "name": "weavatrix-test",
448                "version": env!("CARGO_PKG_VERSION")
449            },
450            "io.modelcontextprotocol/clientCapabilities": {}
451        });
452        let discovered = dispatch(
453            &mut engine,
454            &json!({
455                "jsonrpc": "2.0",
456                "id": "discover",
457                "method": "server/discover",
458                "params": {"_meta": modern_meta.clone()}
459            }),
460        )
461        .expect("modern server/discover is answered");
462        assert_eq!(discovered["result"]["resultType"], "complete");
463        assert_eq!(
464            discovered["result"]["supportedVersions"][0],
465            MODERN_PROTOCOL_VERSION
466        );
467        assert_eq!(
468            discovered["result"]["_meta"]["io.modelcontextprotocol/serverInfo"]["name"],
469            "weavatrix-rust"
470        );
471        assert!(
472            engine.engine.is_some() && matches!(engine.watcher, WatcherState::NotStarted),
473            "server/discover must not start the watcher"
474        );
475
476        let listed = dispatch(
477            &mut engine,
478            &json!({"jsonrpc": "2.0", "id": 2, "method": "tools/list"}),
479        )
480        .expect("tools/list is answered");
481        assert_eq!(
482            listed["result"]["tools"].as_array().map(Vec::len),
483            Some(crate::tools::catalog().len())
484        );
485        assert!(
486            engine.engine.is_some() && matches!(engine.watcher, WatcherState::NotStarted),
487            "tools/list must not start the watcher"
488        );
489
490        let modern_listed = dispatch(
491            &mut engine,
492            &json!({
493                "jsonrpc": "2.0",
494                "id": "modern-list",
495                "method": "tools/list",
496                "params": {"_meta": modern_meta}
497            }),
498        )
499        .expect("modern tools/list is answered");
500        assert_eq!(modern_listed["result"]["resultType"], "complete");
501        assert_eq!(
502            modern_listed["result"]["tools"].as_array().map(Vec::len),
503            Some(crate::tools::catalog().len())
504        );
505
506        let called = dispatch(
507            &mut engine,
508            &json!({
509                "jsonrpc": "2.0",
510                "id": 3,
511                "method": "tools/call",
512                "params": {"name": "graph_stats", "arguments": {}}
513            }),
514        )
515        .expect("tools/call is answered");
516        assert_eq!(called["result"]["isError"], false);
517        assert!(
518            called["result"]["structuredContent"]["nodes"]
519                .as_u64()
520                .unwrap()
521                > 0
522        );
523
524        let mut code = server(crate::mcp::McpProfile::Code);
525        let text = dispatch(
526            &mut code,
527            &json!({
528                "jsonrpc": "2.0",
529                "id": 4,
530                "method": "tools/call",
531                "params": {
532                    "name": "graph_stats",
533                    "arguments": {"output_format": "text"}
534                }
535            }),
536        )
537        .expect("tools/call is answered");
538        assert!(text["result"].get("structuredContent").is_none());
539
540        let denied = dispatch(
541            &mut code,
542            &json!({
543                "jsonrpc": "2.0",
544                "id": 5,
545                "method": "tools/call",
546                "params": {"name": "seo_link_suggestions", "arguments": {}}
547            }),
548        )
549        .expect("tools/call is answered");
550        assert_eq!(denied["error"]["code"], -32_602);
551        assert_eq!(
552            denied["error"]["message"],
553            "unknown tool: seo_link_suggestions"
554        );
555    }
556
557    #[test]
558    fn mcp_refreshes_after_a_real_source_change() {
559        let nonce = std::time::SystemTime::now()
560            .duration_since(std::time::UNIX_EPOCH)
561            .unwrap()
562            .as_nanos();
563        let root = std::env::temp_dir().join(format!(
564            "weavatrix-mcp-watcher-{}-{nonce}",
565            std::process::id()
566        ));
567        std::fs::create_dir_all(&root).unwrap();
568        std::fs::write(root.join("source.rs"), "fn first() {}\n").unwrap();
569        let mut engine = WeavatrixServer::new(root.clone(), crate::mcp::McpProfile::All).unwrap();
570        std::fs::write(root.join("source.rs"), "fn first() {}\nfn second() {}\n").unwrap();
571        let first = dispatch(
572            &mut engine,
573            &json!({
574                "jsonrpc": "2.0",
575                "id": 1,
576                "method": "tools/call",
577                "params": {"name": "graph_stats", "arguments": {}}
578            }),
579        )
580        .unwrap();
581        let first_revision = first["result"]["structuredContent"]["revision"]
582            .as_str()
583            .unwrap()
584            .to_owned();
585        assert_eq!(
586            first["result"]["structuredContent"]["node_kinds"]["function"],
587            2
588        );
589
590        std::fs::write(
591            root.join("source.rs"),
592            "fn first() {}\nfn second() {}\nfn third() {}\n",
593        )
594        .unwrap();
595        let second = dispatch(
596            &mut engine,
597            &json!({
598                "jsonrpc": "2.0",
599                "id": 2,
600                "method": "tools/call",
601                "params": {"name": "graph_stats", "arguments": {}}
602            }),
603        )
604        .unwrap();
605        assert_eq!(
606            second["result"]["structuredContent"]["node_kinds"]["function"],
607            3
608        );
609        assert_ne!(
610            second["result"]["structuredContent"]["revision"].as_str(),
611            Some(first_revision.as_str())
612        );
613
614        std::fs::remove_dir_all(root).unwrap();
615    }
616}