use crate::error::McpError;
use crate::McpResult;
use super::{McpMessage, McpTransport};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
pub struct McpStdioTransport {
reader: tokio::io::BufReader<tokio::io::Stdin>,
writer_tx: tokio::sync::mpsc::Sender<String>,
}
impl Default for McpStdioTransport {
fn default() -> Self {
Self::new()
}
}
impl McpStdioTransport {
pub fn new() -> Self {
let (writer_tx, mut writer_rx) = tokio::sync::mpsc::channel::<String>(32);
tokio::spawn(async move {
let mut writer = tokio::io::BufWriter::new(tokio::io::stdout());
while let Some(message) = writer_rx.recv().await {
if let Err(e) = writer.write_all(message.as_bytes()).await {
eprintln!("Error writing to stdout: {}", e);
}
if let Err(e) = writer.write_all(b"\n").await {
eprintln!("Error writing newline to stdout: {}", e);
}
if let Err(e) = writer.flush().await {
eprintln!("Error flushing stdout: {}", e);
}
}
});
Self {
reader: tokio::io::BufReader::new(tokio::io::stdin()),
writer_tx,
}
}
}
impl Clone for McpStdioTransport {
fn clone(&self) -> Self {
Self {
reader: BufReader::new(tokio::io::stdin()),
writer_tx: self.writer_tx.clone(),
}
}
}
#[async_trait::async_trait]
impl McpTransport for McpStdioTransport {
async fn send(&mut self, message: McpMessage) -> McpResult<()> {
let json = serde_json::to_string(&message)
.map_err(|err| McpError::Serialization(err.to_string()))?;
self.writer_tx.send(json).await
.map_err(|err| McpError::Transport(format!("Failed to send message to writer: {}", err)))
}
async fn receive(&mut self) -> McpResult<McpMessage> {
let mut line = String::new();
match self.reader.read_line(&mut line).await {
Ok(_) => match serde_json::from_str(&line) {
Ok(parsed) => Ok(parsed),
Err(err) => Err(McpError::Serialization(err.to_string()))
}
Err(err) => {
Err(McpError::Transport(format!("Failed to read: {}", err)))
}
}
}
async fn close(&mut self) -> McpResult<()> {
Ok(())
}
}