Skip to content

Commit 647b407

Browse files
committed
Merge remote-tracking branch 'origin/rcc/runtime' into dev/rust
2 parents f8bc5cf + 5eeb7be commit 647b407

1 file changed

Lines changed: 127 additions & 3 deletions

File tree

‎rust/crates/runtime/src/mcp_stdio.rs‎

Lines changed: 127 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,17 +2,19 @@ use std::collections::BTreeMap;
22
use std::io;
33
use std::process::Stdio;
44

5-
use tokio::io::{AsyncReadExt, AsyncWriteExt};
5+
use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader};
66
use tokio::process::{Child, ChildStdin, ChildStdout, Command};
77

8+
use serde_json::Value as JsonRpcMessage;
9+
810
use crate::mcp_client::{McpClientBootstrap, McpClientTransport, McpStdioTransport};
911

1012
#[derive(Debug)]
1113
#[allow(dead_code)]
1214
pub struct McpStdioProcess {
1315
child: Child,
1416
stdin: ChildStdin,
15-
stdout: ChildStdout,
17+
stdout: BufReader<ChildStdout>,
1618
}
1719

1820
#[allow(dead_code)]
@@ -39,7 +41,7 @@ impl McpStdioProcess {
3941
Ok(Self {
4042
child,
4143
stdin,
42-
stdout,
44+
stdout: BufReader::new(stdout),
4345
})
4446
}
4547

@@ -58,6 +60,49 @@ impl McpStdioProcess {
5860
Ok(buffer)
5961
}
6062

63+
pub async fn write_jsonrpc_message(&mut self, message: &JsonRpcMessage) -> io::Result<()> {
64+
let encoded = encode_jsonrpc_message(message)?;
65+
self.write_all(&encoded).await?;
66+
self.flush().await
67+
}
68+
69+
pub async fn read_jsonrpc_message(&mut self) -> io::Result<JsonRpcMessage> {
70+
let payload = self.read_frame().await?;
71+
serde_json::from_slice(&payload)
72+
.map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))
73+
}
74+
75+
async fn read_frame(&mut self) -> io::Result<Vec<u8>> {
76+
let mut content_length = None;
77+
loop {
78+
let mut line = String::new();
79+
let bytes_read = self.stdout.read_line(&mut line).await?;
80+
if bytes_read == 0 {
81+
return Err(io::Error::new(
82+
io::ErrorKind::UnexpectedEof,
83+
"MCP stdio stream closed while reading headers",
84+
));
85+
}
86+
if line == "\r\n" {
87+
break;
88+
}
89+
if let Some(value) = line.strip_prefix("Content-Length:") {
90+
let parsed = value
91+
.trim()
92+
.parse::<usize>()
93+
.map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;
94+
content_length = Some(parsed);
95+
}
96+
}
97+
98+
let content_length = content_length.ok_or_else(|| {
99+
io::Error::new(io::ErrorKind::InvalidData, "missing Content-Length header")
100+
})?;
101+
let mut payload = vec![0_u8; content_length];
102+
self.stdout.read_exact(&mut payload).await?;
103+
Ok(payload)
104+
}
105+
61106
pub async fn terminate(&mut self) -> io::Result<()> {
62107
self.child.kill().await
63108
}
@@ -88,6 +133,15 @@ fn apply_env(command: &mut Command, env: &BTreeMap<String, String>) {
88133
}
89134
}
90135

136+
fn encode_jsonrpc_message(message: &JsonRpcMessage) -> io::Result<Vec<u8>> {
137+
let body = serde_json::to_vec(message)
138+
.map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;
139+
let header = format!("Content-Length: {}\r\n\r\n", body.len());
140+
let mut framed = header.into_bytes();
141+
framed.extend(body);
142+
Ok(framed)
143+
}
144+
91145
#[cfg(test)]
92146
mod tests {
93147
use std::collections::BTreeMap;
@@ -129,6 +183,37 @@ mod tests {
129183
script_path
130184
}
131185

186+
fn write_jsonrpc_script() -> PathBuf {
187+
let root = temp_dir();
188+
fs::create_dir_all(&root).expect("temp dir");
189+
let script_path = root.join("jsonrpc-mcp.py");
190+
let script = [
191+
"#!/usr/bin/env python3",
192+
"import json, sys",
193+
"header = b''",
194+
r"while not header.endswith(b'\r\n\r\n'):",
195+
" chunk = sys.stdin.buffer.read(1)",
196+
" if not chunk:",
197+
" raise SystemExit(1)",
198+
" header += chunk",
199+
"length = 0",
200+
r"for line in header.decode().split('\r\n'):",
201+
r" if line.lower().startswith('content-length:'):",
202+
r" length = int(line.split(':', 1)[1].strip())",
203+
"payload = sys.stdin.buffer.read(length)",
204+
"json.loads(payload.decode())",
205+
r"response = json.dumps({'jsonrpc': '2.0', 'id': 1, 'result': {'echo': True}}).encode()",
206+
r"sys.stdout.buffer.write(f'Content-Length: {len(response)}\r\n\r\n'.encode() + response)",
207+
"sys.stdout.buffer.flush()",
208+
"",
209+
]
210+
.join("\n");
211+
fs::write(&script_path, script).expect("write script");
212+
let mut permissions = fs::metadata(&script_path).expect("metadata").permissions();
213+
permissions.set_mode(0o755);
214+
fs::set_permissions(&script_path, permissions).expect("chmod");
215+
script_path
216+
}
132217
fn sample_bootstrap(script_path: &Path) -> McpClientBootstrap {
133218
let config = ScopedMcpServerConfig {
134219
scope: ConfigSource::Local,
@@ -185,6 +270,45 @@ mod tests {
185270
assert_eq!(error.kind(), ErrorKind::InvalidInput);
186271
}
187272

273+
#[test]
274+
fn round_trips_jsonrpc_messages_over_stdio_frames() {
275+
let runtime = Builder::new_current_thread()
276+
.enable_all()
277+
.build()
278+
.expect("runtime");
279+
runtime.block_on(async {
280+
let script_path = write_jsonrpc_script();
281+
let transport = crate::mcp_client::McpStdioTransport {
282+
command: "python3".to_string(),
283+
args: vec![script_path.to_string_lossy().into_owned()],
284+
env: BTreeMap::new(),
285+
};
286+
let mut process = McpStdioProcess::spawn(&transport).expect("spawn transport directly");
287+
process
288+
.write_jsonrpc_message(&serde_json::json!({
289+
"jsonrpc": "2.0",
290+
"id": 1,
291+
"method": "initialize"
292+
}))
293+
.await
294+
.expect("write jsonrpc message");
295+
296+
let response = process
297+
.read_jsonrpc_message()
298+
.await
299+
.expect("read jsonrpc response");
300+
assert_eq!(response["jsonrpc"], serde_json::json!("2.0"));
301+
assert_eq!(response["id"], serde_json::json!(1));
302+
assert_eq!(response["result"]["echo"], serde_json::json!(true));
303+
304+
let status = process.wait().await.expect("wait for exit");
305+
assert!(status.success());
306+
307+
fs::remove_file(&script_path).expect("cleanup script");
308+
fs::remove_dir_all(script_path.parent().expect("script parent")).expect("cleanup dir");
309+
});
310+
}
311+
188312
#[test]
189313
fn direct_spawn_uses_transport_env() {
190314
let runtime = Builder::new_current_thread()

0 commit comments

Comments
 (0)