@@ -2,17 +2,19 @@ use std::collections::BTreeMap;
22use std:: io;
33use std:: process:: Stdio ;
44
5- use tokio:: io:: { AsyncReadExt , AsyncWriteExt } ;
5+ use tokio:: io:: { AsyncBufReadExt , AsyncReadExt , AsyncWriteExt , BufReader } ;
66use tokio:: process:: { Child , ChildStdin , ChildStdout , Command } ;
77
8+ use serde_json:: Value as JsonRpcMessage ;
9+
810use crate :: mcp_client:: { McpClientBootstrap , McpClientTransport , McpStdioTransport } ;
911
1012#[ derive( Debug ) ]
1113#[ allow( dead_code) ]
1214pub 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) ]
92146mod 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