shuyi_cli/mcp/
forward.rs

1//! The projects' own servers (ticket 2546): a call whose project is in
2//! another cargo workspace than the server's goes to a `shuyi mcp` built
3//! from that workspace's engine, a **forwarded server**, so an agent in a
4//! worktree plays and checks the worktree's code, whichever checkout the
5//! client started the server in.
6//!
7//! A forwarded server is the project's `command` and `mcp` ([`words`]),
8//! run in the project's root with pipes to its stdin, stdout and stderr,
9//! and [`FORWARDED`] set, so it answers every call itself. It gets an
10//! `initialize` without roots as it starts, and its build runs while the
11//! server goes on. It ends when its stdin closes: when the server drops it,
12//! or ends, or is killed.
13//!
14//! [`Forwarding`] is the seam the server starts them and looks up
15//! workspaces through ([`Processes`] for real), so its tests run no cargo.
16
17use std::fmt::Write as _;
18use std::fs;
19use std::io::{BufRead, BufReader, Write};
20use std::path::{Path, PathBuf};
21use std::process::{Child, Command, Stdio};
22use std::sync::mpsc::{Receiver, TryRecvError, channel};
23use std::sync::{Arc, Mutex};
24use std::time::Duration;
25
26use serde_json::{Value, json};
27
28use super::child::keep_tail;
29use super::server::PROTOCOL_VERSIONS;
30use crate::project::Project;
31
32/// The environment variable a forwarded server is started with: it hands
33/// nothing on, so a server is never forwarded twice.
34pub const FORWARDED: &str = "SHUYI_MCP_FORWARDED";
35
36/// The id of the `initialize` request a forwarded server gets.
37const INITIALIZE: u64 = 0;
38
39/// A project's own server that the server hands calls to.
40pub trait Forward {
41    /// Whether it has answered `initialize`: none while it builds (or,
42    /// with `wait`, once it has answered or stopped), or why it never will.
43    fn ready(&mut self, wait: bool) -> Option<Result<(), String>>;
44
45    /// Its response to the request `method` with `params`, once it is
46    /// ready, or why there is none: it stopped, with its last output.
47    fn request(&mut self, method: &str, params: &Value) -> Result<Value, String>;
48}
49
50/// Looks up workspaces and starts forwarded servers.
51pub trait Forwarding<'a> {
52    /// The root of the cargo workspace `folder` is in, none outside one.
53    /// Paths of one folder compare equal.
54    fn workspace(&self, folder: &Path) -> Option<PathBuf>;
55
56    /// A forwarded server: `words` run in `root`.
57    fn spawn(&self, words: &[String], root: &Path) -> Box<dyn Forward + 'a>;
58}
59
60/// The words that start the server of the project at `root`, in the
61/// workspace `workspace`, and the folder they run in:
62///
63/// - the project's `command` (`shuyi.ron`) and `mcp`, in its root, if it
64///   names one;
65/// - else the `shuyi` entry of the workspace's `.mcp.json` (the engine's
66///   own repository has one and no `shuyi.ron` at its root), in the
67///   workspace;
68/// - else `shuyi mcp` in the project's root.
69///
70/// None outside a project with no entry: the server answers itself.
71pub fn words(root: &Path, workspace: &Path) -> Option<(Vec<String>, PathBuf)> {
72    let project = Project::find(root).ok();
73    if let Some(project) = &project
74        && let Some(command) = &project.config.command
75    {
76        return Some((with_mcp(command), project.root.clone()));
77    }
78    if let Some(words) = entry(&workspace.join(".mcp.json")) {
79        return Some((words, workspace.to_path_buf()));
80    }
81    project.map(|project| (with_mcp(project.config.command()), project.root))
82}
83
84/// `command`'s words, then `mcp`.
85fn with_mcp(command: &str) -> Vec<String> {
86    command
87        .split_whitespace()
88        .chain(["mcp"])
89        .map(str::to_owned)
90        .collect()
91}
92
93/// The `shuyi` server's command and args in the `.mcp.json` at `path`.
94fn entry(path: &Path) -> Option<Vec<String>> {
95    let file: Value = serde_json::from_str(&fs::read_to_string(path).ok()?).ok()?;
96    let server = &file["mcpServers"]["shuyi"];
97    let command = server["command"].as_str()?;
98    let args = server["args"].as_array().map_or(&[][..], Vec::as_slice);
99    std::iter::once(Some(command))
100        .chain(args.iter().map(Value::as_str))
101        .map(|word| word.map(str::to_owned))
102        .collect()
103}
104
105/// The folder a `file://` URI names (a client's root), or none if it
106/// names none: its percent-escapes decoded, and a Windows drive's slash
107/// dropped (`file:///C:/x` is `C:/x`).
108pub fn path_of_uri(uri: &str) -> Option<PathBuf> {
109    let rest = uri.strip_prefix("file://")?;
110    let path = rest.strip_prefix("localhost").unwrap_or(rest);
111    if !path.starts_with('/') {
112        return None;
113    }
114    let text = String::from_utf8(decoded(path)?).ok()?;
115    let b = text.as_bytes();
116    let text = if b.len() >= 3 && b[1].is_ascii_alphabetic() && b[2] == b':' {
117        &text[1..]
118    } else {
119        &text
120    };
121    if cfg!(windows) {
122        Some(PathBuf::from(text.replace('/', "\\")))
123    } else {
124        Some(PathBuf::from(text))
125    }
126}
127
128/// `text` with each `%XX` made its byte, or none if one is broken.
129fn decoded(text: &str) -> Option<Vec<u8>> {
130    let mut out = Vec::with_capacity(text.len());
131    let mut bytes = text.bytes();
132    while let Some(byte) = bytes.next() {
133        if byte == b'%' {
134            let hex = [bytes.next()?, bytes.next()?];
135            out.push(u8::from_str_radix(std::str::from_utf8(&hex).ok()?, 16).ok()?);
136        } else {
137            out.push(byte);
138        }
139    }
140    Some(out)
141}
142
143/// `path` without the `\\?\` before a Windows drive that canonicalizing
144/// gives, which some programs don't read.
145fn plain(path: PathBuf) -> PathBuf {
146    match path.to_str().and_then(|text| text.strip_prefix(r"\\?\")) {
147        Some(rest) if rest.as_bytes().get(1) == Some(&b':') => PathBuf::from(rest),
148        _ => path,
149    }
150}
151
152/// Looks up workspaces with cargo and starts forwarded servers as
153/// processes.
154#[derive(Debug, Clone, Copy, Default)]
155pub struct Processes;
156
157impl<'a> Forwarding<'a> for Processes {
158    fn workspace(&self, folder: &Path) -> Option<PathBuf> {
159        let cargo = crate::cargo();
160        let output = Command::new(cargo)
161            .args(["locate-project", "--workspace", "--message-format", "plain"])
162            .current_dir(folder)
163            .stdin(Stdio::null())
164            .stderr(Stdio::null())
165            .output()
166            .ok()?;
167        if !output.status.success() {
168            return None;
169        }
170        let manifest = PathBuf::from(String::from_utf8(output.stdout).ok()?.trim());
171        let root = manifest.parent()?;
172        Some(plain(
173            fs::canonicalize(root).unwrap_or_else(|_| root.to_path_buf()),
174        ))
175    }
176
177    fn spawn(&self, words: &[String], root: &Path) -> Box<dyn Forward + 'a> {
178        match Process::start(words, root) {
179            Ok(process) => Box::new(process),
180            Err(why) => Box::new(Failed(why)),
181        }
182    }
183}
184
185/// A forwarded server that never started: each call says why.
186struct Failed(String);
187
188impl Forward for Failed {
189    fn ready(&mut self, _: bool) -> Option<Result<(), String>> {
190        Some(Err(self.0.clone()))
191    }
192
193    fn request(&mut self, _: &str, _: &Value) -> Result<Value, String> {
194        Err(self.0.clone())
195    }
196}
197
198/// A forwarded server: its stdin, the lines of its stdout, and the last
199/// of its stderr.
200pub struct Process {
201    /// What it is, for errors: `the project's own server ("…" in …)`.
202    named: String,
203    child: Option<Child>,
204    stdin: Option<Box<dyn Write + Send>>,
205    lines: Receiver<String>,
206    tail: Arc<Mutex<Vec<u8>>>,
207    /// Hears when its stderr has been copied to the end.
208    copied: Option<Receiver<()>>,
209    /// Whether it answered `initialize`, once known.
210    state: Option<Result<(), String>>,
211    /// The id of its next request.
212    next: u64,
213}
214
215impl Process {
216    /// A forwarded server called `named`, writing to `stdin` and answering
217    /// on `lines`, its stderr's last bytes in `tail`; sent `initialize` now.
218    pub fn new(
219        named: String,
220        stdin: Box<dyn Write + Send>,
221        lines: Receiver<String>,
222        tail: Arc<Mutex<Vec<u8>>>,
223    ) -> Self {
224        let mut process = Self {
225            named,
226            child: None,
227            stdin: Some(stdin),
228            lines,
229            tail,
230            copied: None,
231            state: None,
232            next: INITIALIZE + 1,
233        };
234        process.send(&json!({
235            "jsonrpc": "2.0",
236            "id": INITIALIZE,
237            "method": "initialize",
238            "params": {
239                "protocolVersion": PROTOCOL_VERSIONS[0],
240                "capabilities": {},
241                "clientInfo": { "name": "shuyi", "version": env!("CARGO_PKG_VERSION") },
242            },
243        }));
244        process
245    }
246
247    /// `words` run in `root` as a forwarded server, or why it didn't
248    /// start.
249    fn start(words: &[String], root: &Path) -> Result<Self, String> {
250        let named = format!(
251            "the project's own server (\"{}\" in {})",
252            words.join(" "),
253            root.display()
254        );
255        let Some((program, args)) = words.split_first() else {
256            return Err(format!("{named} has no command"));
257        };
258        eprintln!("mcp: {} is answered by {named}", root.display());
259        let spawned = Command::new(program)
260            .args(args)
261            .current_dir(root)
262            .env(FORWARDED, "1")
263            // So `cargo` picks the toolchain of the project's checkout.
264            .env_remove("RUSTUP_TOOLCHAIN")
265            .stdin(Stdio::piped())
266            .stdout(Stdio::piped())
267            .stderr(Stdio::piped())
268            .spawn();
269        let mut child = match spawned {
270            Ok(child) => child,
271            Err(e) => return Err(format!("{named} didn't start: {e}")),
272        };
273        let (Some(stdin), Some(stdout), Some(stderr)) =
274            (child.stdin.take(), child.stdout.take(), child.stderr.take())
275        else {
276            let _ = child.kill();
277            return Err(format!("{named} started without its pipes"));
278        };
279        let (sent, lines) = channel();
280        std::thread::spawn(move || {
281            for line in BufReader::new(stdout).lines() {
282                let Ok(line) = line else { break };
283                if sent.send(line).is_err() {
284                    break;
285                }
286            }
287        });
288        let tail = Arc::new(Mutex::new(Vec::new()));
289        let kept = Arc::clone(&tail);
290        let (ended, copied) = channel();
291        std::thread::spawn(move || {
292            keep_tail(stderr, &kept);
293            let _ = ended.send(());
294        });
295        let mut process = Self::new(named, Box::new(stdin), lines, tail);
296        process.child = Some(child);
297        process.copied = Some(copied);
298        Ok(process)
299    }
300
301    /// Writes `message` on a line of its stdin; false if it is closed.
302    fn send(&mut self, message: &Value) -> bool {
303        let sent = self
304            .stdin
305            .as_mut()
306            .map(|stdin| writeln!(stdin, "{message}").and_then(|()| stdin.flush()));
307        matches!(sent, Some(Ok(())))
308    }
309
310    /// Why it answers nothing: it stopped, with its exit code and the
311    /// last of its output. It is stopped for good.
312    fn stopped(&mut self) -> String {
313        if let Some(Err(why)) = &self.state {
314            return why.clone();
315        }
316        self.stdin = None;
317        let status = self
318            .child
319            .as_mut()
320            .and_then(|child| child.wait().ok())
321            .and_then(|status| status.code());
322        // Its stderr closes as it ends: let the copy catch up, briefly.
323        if let Some(copied) = &self.copied {
324            let _ = copied.recv_timeout(Duration::from_secs(1));
325        }
326        let tail = self
327            .tail
328            .lock()
329            .map(|tail| String::from_utf8_lossy(&tail).trim().to_owned())
330            .unwrap_or_default();
331        let mut out = format!("{} stopped before it answered", self.named);
332        if let Some(code) = status {
333            let _ = write!(out, " (exit code {code})");
334        }
335        if !tail.is_empty() {
336            out.push_str(": ");
337            out.push_str(&tail);
338        }
339        self.state = Some(Err(out.clone()));
340        out
341    }
342
343    /// Reads its lines until the response to `id`, waiting for each with
344    /// `wait`: `Err` if it stops first, none if `wait` is false and it
345    /// hasn't come yet.
346    fn response(&mut self, id: u64, wait: bool) -> Result<Option<Value>, String> {
347        loop {
348            let next = if wait {
349                self.lines.recv().ok()
350            } else {
351                match self.lines.try_recv() {
352                    Ok(line) => Some(line),
353                    Err(TryRecvError::Empty) => return Ok(None),
354                    Err(TryRecvError::Disconnected) => None,
355                }
356            };
357            // Its stdout closed.
358            let Some(line) = next else {
359                return Err(self.stopped());
360            };
361            // A notification, or a line that isn't JSON: not the answer.
362            let Ok(message) = serde_json::from_str::<Value>(&line) else {
363                continue;
364            };
365            if message.get("id") == Some(&json!(id)) {
366                return Ok(Some(message));
367            }
368        }
369    }
370}
371
372impl Forward for Process {
373    fn ready(&mut self, wait: bool) -> Option<Result<(), String>> {
374        if self.state.is_none() {
375            let state = match self.response(INITIALIZE, wait) {
376                Ok(None) => return None,
377                Err(why) => Err(why),
378                Ok(Some(answer)) if answer.get("result").is_some() => {
379                    self.send(&json!({ "jsonrpc": "2.0", "method": "notifications/initialized" }));
380                    Ok(())
381                }
382                Ok(Some(answer)) => Err(format!(
383                    "{} refused to start: {}",
384                    self.named, answer["error"]
385                )),
386            };
387            self.state = Some(state);
388        }
389        self.state.clone()
390    }
391
392    fn request(&mut self, method: &str, params: &Value) -> Result<Value, String> {
393        if let Some(Err(why)) = self.ready(true) {
394            return Err(why);
395        }
396        let id = self.next;
397        self.next += 1;
398        let request = json!({ "jsonrpc": "2.0", "id": id, "method": method, "params": params });
399        if !self.send(&request) {
400            return Err(self.stopped());
401        }
402        self.response(id, true)?
403            .ok_or_else(|| format!("{} answered nothing", self.named))
404    }
405}
406
407impl Drop for Process {
408    /// Closes its stdin, which ends it, and makes sure it has ended.
409    fn drop(&mut self) {
410        self.stdin = None;
411        if let Some(child) = self.child.as_mut() {
412            // A moment to end by itself, as a server does when its stdin
413            // closes; then it is ended.
414            let _ = self
415                .copied
416                .as_ref()
417                .map(|copied| copied.recv_timeout(Duration::from_millis(500)));
418            let _ = child.kill();
419            let _ = child.wait();
420        }
421    }
422}
423
424#[cfg(test)]
425mod tests;
shuyi_cli/mcp/forward.rs

shuyi_cli/mcp/
forward.rs
