diff --git a/mcp-stdio/src/lib.rs b/mcp-stdio/src/lib.rs index 8697a29..00b59f7 100644 --- a/mcp-stdio/src/lib.rs +++ b/mcp-stdio/src/lib.rs @@ -22,9 +22,8 @@ pub async fn read_mcp_message( break; } - let lower_line = line.to_lowercase(); - if let Some(len_str) = lower_line.strip_prefix("content-length:") { - length = len_str.trim().parse().unwrap_or(0); + if line.len() >= 15 && line[..15].eq_ignore_ascii_case("content-length:") { + length = line[15..].trim().parse().unwrap_or(0); } } @@ -39,3 +38,4 @@ pub async fn read_mcp_message( String::from_utf8(buffer).ok() } + diff --git a/server/src/fix_stub.py b/server/src/fix_stub.py new file mode 100644 index 0000000..9a8bc8d --- /dev/null +++ b/server/src/fix_stub.py @@ -0,0 +1,66 @@ +import os + +def fix_stub_leaks(): + filepath = 'stub/src/main.rs' + with open(filepath, 'r', encoding='utf-8') as f: + content = f.read() + + # 1. Fix connect_async to handle shutdown and timeout + old_connect = """ match tokio_tungstenite::connect_async(request).await {""" + new_connect = """ let connect_result = tokio::select! { + _ = shutdown_rx.recv() => { + tracing::info!("Shutdown received during connect"); + break; + } + res = tokio::time::timeout( + tokio::time::Duration::from_secs(5), + tokio_tungstenite::connect_async(request) + ) => res, + }; + + match connect_result { + Ok(Ok((ws_stream, _))) => {""" + + if old_connect in content: + content = content.replace(old_connect, new_connect) + print("Replaced connect_async") + + # Fix Err block to match the new match structure + old_err = """ Err(e) => { + tracing::error!("Failed to connect via WSS: {}", e); + tokio::time::sleep(tokio::time::Duration::from_secs(1)).await; + }""" + new_err = """ Ok(Err(e)) => { + tracing::error!("Failed to connect via WSS: {}", e); + tokio::time::sleep(tokio::time::Duration::from_secs(1)).await; + } + Err(_) => { + tracing::error!("Connection attempt timed out"); + tokio::time::sleep(tokio::time::Duration::from_secs(1)).await; + }""" + + if old_err in content: + content = content.replace(old_err, new_err) + print("Replaced Err branch") + + # 2. Fix the break in send_task that exits the stub instead of reconnecting + old_select_send = """ _ = &mut send_task => { + tracing::error!("Send task exited"); + tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; + recv_task.abort(); + break; + }""" + new_select_send = """ _ = &mut send_task => { + tracing::error!("Send task exited"); + tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; + recv_task.abort(); + }""" + + if old_select_send in content: + content = content.replace(old_select_send, new_select_send) + print("Replaced select send_task") + + with open(filepath, 'w', encoding='utf-8') as f: + f.write(content) + +fix_stub_leaks() diff --git a/stub/src/main.rs b/stub/src/main.rs index d5565e1..028193b 100644 --- a/stub/src/main.rs +++ b/stub/src/main.rs @@ -74,8 +74,19 @@ fn main() -> Result<(), Box> { } }; - match tokio_tungstenite::connect_async(request).await { - Ok((ws_stream, _)) => { + let connect_result = tokio::select! { + _ = shutdown_rx.recv() => { + tracing::info!("Shutdown received during connect"); + break; + } + res = tokio::time::timeout( + tokio::time::Duration::from_secs(5), + tokio_tungstenite::connect_async(request) + ) => res, + }; + + match connect_result { + Ok(Ok((ws_stream, _))) => { tracing::info!("Successfully connected to target server"); let (mut write, mut read) = ws_stream.split(); @@ -138,7 +149,6 @@ fn main() -> Result<(), Box> { tracing::error!("Send task exited"); tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; recv_task.abort(); - break; } _ = &mut recv_task => { tracing::error!("Recv task exited"); @@ -147,12 +157,18 @@ fn main() -> Result<(), Box> { } tokio::time::sleep(tokio::time::Duration::from_secs(1)).await; } - Err(e) => { + Ok(Err(e)) => { tracing::error!("Failed to connect via WSS: {}", e); tokio::time::sleep(tokio::time::Duration::from_secs(1)).await; } + Err(_) => { + tracing::error!("Connection attempt timed out"); + tokio::time::sleep(tokio::time::Duration::from_secs(1)).await; + } } } Ok(()) }) } + +