app-server-transport/src/transport/stdio.rs
113 lines
use super::CHANNEL_CAPACITY;use super::ConnectionOrigin;use super::TransportEvent;use super::forward_incoming_message;use super::next_connection_id;use super::serialize_outgoing_message;use crate::outgoing_message::QueuedOutgoingMessage;use codex_app_server_protocol::InitializeParams;use codex_app_server_protocol::JSONRPCMessage;use codex_app_server_protocol::JSONRPCRequest;use std::io::ErrorKind;use std::io::Result as IoResult;use tokio::io;use tokio::io::AsyncBufReadExt;use tokio::io::AsyncWriteExt;use tokio::io::BufReader;use tokio::sync::mpsc;use tokio::sync::oneshot;use tokio::task::JoinHandle;use tracing::debug;use tracing::error;use tracing::info;pub async fn start_stdio_connection( transport_event_tx: mpsc::Sender<TransportEvent>, stdio_handles: &mut Vec<JoinHandle<()>>, initialize_client_name_tx: oneshot::Sender<String>,) -> IoResult<()> { let connection_id = next_connection_id(); let (writer_tx, mut writer_rx) = mpsc::channel::<QueuedOutgoingMessage>(CHANNEL_CAPACITY); let writer_tx_for_reader = writer_tx.clone(); transport_event_tx .send(TransportEvent::ConnectionOpened { connection_id, origin: ConnectionOrigin::Stdio, writer: writer_tx, disconnect_sender: None, }) .await .map_err(|_| std::io::Error::new(ErrorKind::BrokenPipe, "processor unavailable"))?; let transport_event_tx_for_reader = transport_event_tx.clone(); stdio_handles.push(tokio::spawn(async move { let stdin = io::stdin(); let reader = BufReader::new(stdin); let mut lines = reader.lines(); let mut initialize_client_name_tx = Some(initialize_client_name_tx); loop { match lines.next_line().await { Ok(Some(line)) => { if let Some(client_name) = stdio_initialize_client_name(&line) && let Some(initialize_client_name_tx) = initialize_client_name_tx.take() { let _ = initialize_client_name_tx.send(client_name); } if !forward_incoming_message( &transport_event_tx_for_reader, &writer_tx_for_reader, connection_id, &line, ) .await { break; } } Ok(None) => break, Err(err) => { error!("Failed reading stdin: {err}"); break; } } } let _ = transport_event_tx_for_reader .send(TransportEvent::ConnectionClosed { connection_id }) .await; debug!("stdin reader finished (EOF)"); })); stdio_handles.push(tokio::spawn(async move { let mut stdout = io::stdout(); while let Some(queued_message) = writer_rx.recv().await { let Some(mut json) = serialize_outgoing_message(queued_message.message) else { continue; }; json.push('\n'); if let Err(err) = stdout.write_all(json.as_bytes()).await { error!("Failed to write to stdout: {err}"); break; } if let Some(write_complete_tx) = queued_message.write_complete_tx { let _ = write_complete_tx.send(()); } } info!("stdout writer exited (channel closed)"); })); Ok(())}fn stdio_initialize_client_name(line: &str) -> Option<String> { let message = serde_json::from_str::<JSONRPCMessage>(line).ok()?; let JSONRPCMessage::Request(JSONRPCRequest { method, params, .. }) = message else { return None; }; if method != "initialize" { return None; } let params = serde_json::from_value::<InitializeParams>(params?).ok()?; Some(params.client_info.name)}