diff --git a/Cargo.lock b/Cargo.lock index bf6e20a..218407f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9,6 +9,7 @@ dependencies = [ "agent-client-protocol-derive", "agent-client-protocol-schema", "agent-client-protocol-test", + "async-io", "async-process", "blocking", "clap", diff --git a/Cargo.toml b/Cargo.toml index e4e9cde..85edc25 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -71,6 +71,7 @@ axum = "0.8" reqwest = { version = "0.13", default-features = false, features = ["rustls", "json"] } eventsource-stream = "0.2" url = "2.5" +async-io = "2" async-process = "2" async-stream = "0.3.6" blocking = "1" diff --git a/md/design.md b/md/design.md index 05268d0..8a1a862 100644 --- a/md/design.md +++ b/md/design.md @@ -144,9 +144,20 @@ stateDiagram-v2 ### Two Connection Modes -**Reactive mode** (`connect_to`): The connection runs handlers until the transport closes. Used for agents and proxies. +**Reactive mode** (`connect_to`): The connection runs handlers until the incoming transport reaches clean EOF, drains responses and notifications already accepted by the outgoing queue through the transport sink, then returns `Ok(())`, including when the builder has long-running `with_spawned` work. Used for agents and proxies. -**Active mode** (`connect_with`): Runs a closure with access to the connection, then closes. Used for clients that drive the interaction. +**Active mode** (`connect_with`): Runs a closure with access to the connection, then closes. Used for clients that drive the interaction. Incoming EOF fails requests that still need responses, but it does not automatically cancel unrelated work in the closure. + +### Clean Incoming EOF + +Incoming EOF is a connection event and a request-liveness boundary: + +- Every pending request is completed with an internal error whose data identifies `incoming_transport_closed` and the request method; `is_incoming_transport_closed()` detects it. +- A request created after EOF fails immediately with the same error. +- `ConnectionTo::incoming_closed()` waits for the close event, and `is_incoming_closed()` reports whether it has completed. +- `Builder::on_close()` runs cleanup callbacks in registration order. Returning an error terminates a still-running `connect_with` foreground; returning `Ok(())` leaves its lifetime under application control. + +This keeps request correctness separate from async cancellation policy. Applications can finish cleanup or notify a central dispatcher without having an arbitrary foreground future dropped at an await point. Pending requests are failed before close callbacks begin; the close signal is published after callbacks finish, so a callback must not await `incoming_closed()` itself. ## Key Source Files diff --git a/src/agent-client-protocol-http/src/client.rs b/src/agent-client-protocol-http/src/client.rs index c531bb5..2c19922 100644 --- a/src/agent-client-protocol-http/src/client.rs +++ b/src/agent-client-protocol-http/src/client.rs @@ -9,7 +9,7 @@ use agent_client_protocol::{ }; use async_tungstenite::tungstenite::Message as WsMessage; use futures::{ - StreamExt, + Stream, StreamExt, channel::mpsc::{self, UnboundedSender}, future::{BoxFuture, FutureExt}, pin_mut, @@ -102,14 +102,23 @@ impl HttpClient { impl ConnectTo for HttpClient { async fn connect_to(self, client: impl ConnectTo) -> Result<(), AcpError> { let (channel, transport) = ConnectTo::::into_channel_and_future(self); + let shutdown_tx = channel.tx.clone(); match futures::future::select( std::pin::pin!(client.connect_to(channel)), std::pin::pin!(transport), ) .await { - futures::future::Either::Left((result, _)) - | futures::future::Either::Right((result, _)) => result, + futures::future::Either::Left((result, transport)) => { + result?; + + // Reject sends from escaped client handles while preserving + // messages already accepted into the channel, then let the + // physical transport finish those messages. + shutdown_tx.close_channel(); + transport.await + } + futures::future::Either::Right((result, _)) => result, } } @@ -139,10 +148,22 @@ async fn run(client: HttpClient, channel: Channel) -> Result<(), AcpError> { let mut lifecycle = HttpTransportLifecycle::new(connection); let mut ordered_posts = PostQueue::default(); let mut response_posts = PostQueue::default(); + let mut outgoing_closed = false; let result = loop { + if outgoing_closed && ordered_posts.is_empty() && response_posts.is_empty() { + break Ok(()); + } + let event = { - let outgoing_next = outgoing.next().fuse(); + let outgoing_next = async { + if outgoing_closed { + futures::future::pending().await + } else { + outgoing.next().await + } + } + .fuse(); let sse_event_next = sse_event_rx.next().fuse(); let sse_failure_next = lifecycle.next_sse_failure().fuse(); let ordered_post_next = ordered_posts.next_completion().fuse(); @@ -171,7 +192,10 @@ async fn run(client: HttpClient, channel: Channel) -> Result<(), AcpError> { error!("upstream channel produced error: {e}"); break Err(e); } - None => break Ok(()), + None => { + outgoing_closed = true; + continue; + } }, HttpLoopEvent::SseEvent(event) => { let Some(event) = event else { @@ -491,6 +515,10 @@ impl PostQueue { self.in_flight = Some(post.into_completion()); } } + + fn is_empty(&self) -> bool { + self.queued.is_empty() && self.in_flight.is_none() + } } #[derive(Clone, Copy, Debug, Eq, PartialEq)] @@ -691,10 +719,6 @@ fn pending_request_key(id: &RequestId) -> Option { async fn run_ws(client: HttpClient, channel: Channel) -> Result<(), AcpError> { let HttpClient { endpoint, .. } = client; - let Channel { - rx: mut outgoing, - tx: incoming, - } = channel; let (ws_stream, response) = async_tungstenite::tokio::connect_async(endpoint.as_str()) .await @@ -703,43 +727,85 @@ async fn run_ws(client: HttpClient, channel: Channel) -> Result<(), AcpError> { status = %response.status(), "WebSocket connection established" ); - let (mut ws_tx, mut ws_rx) = ws_stream.split(); + let (ws_tx, ws_rx) = ws_stream.split(); - loop { - let outgoing_next = outgoing.next().fuse(); - let frame_next = ws_rx.next().fuse(); - pin_mut!(outgoing_next, frame_next); + drive_ws(ws_tx, ws_rx, channel).await +} - futures::select! { - msg = outgoing_next => match msg { - Some(Ok(msg)) => { +trait WsSink { + fn send( + &mut self, + message: WsMessage, + ) -> impl std::future::Future> + Send; +} + +impl WsSink for async_tungstenite::WebSocketSender +where + S: futures::AsyncRead + futures::AsyncWrite + Unpin + Send, +{ + async fn send(&mut self, message: WsMessage) -> Result<(), String> { + async_tungstenite::WebSocketSender::send(self, message) + .await + .map_err(|error| error.to_string()) + } +} + +async fn drive_ws( + mut ws_tx: Tx, + mut ws_rx: Rx, + channel: Channel, +) -> Result<(), AcpError> +where + Tx: WsSink, + Rx: Stream> + Unpin, + RxError: std::fmt::Display, +{ + let Channel { + rx: mut outgoing, + tx: incoming, + } = channel; + let writer = async move { + while let Some(msg) = outgoing.next().await { + match msg { + Ok(msg) => { let text = match serde_json::to_string(&msg) { Ok(t) => t, Err(e) => { error!("failed to serialize outbound message: {e}"); - return Err(AcpError::internal_error() - .data(format!("serialize: {e}"))); + return Err(AcpError::internal_error().data(format!("serialize: {e}"))); } }; if let Err(e) = ws_tx.send(WsMessage::Text(text.into())).await { error!("WebSocket send failed: {e}"); - return Err(AcpError::internal_error() - .data(format!("ws send: {e}"))); + return Err(AcpError::internal_error().data(format!("ws send: {e}"))); } } - Some(Err(e)) => { + Err(e) => { error!("upstream channel produced error: {e}"); return Err(e); } - None => break, - }, - frame = frame_next => match frame { + } + } + + drop(ws_tx.send(WsMessage::Close(None)).await); + Ok(()) + }; + + let reader = async move { + let mut discard_incoming = false; + loop { + match ws_rx.next().await { Some(Ok(WsMessage::Text(text))) => { + if discard_incoming { + continue; + } match serde_json::from_str::(text.as_str()) { Ok(parsed) => { if incoming.unbounded_send(Ok(parsed)).is_err() { - debug!("upstream channel closed; stopping WS reader"); - break; + debug!( + "upstream channel closed; discarding WS input while draining output" + ); + discard_incoming = true; } } Err(e) => { @@ -749,8 +815,10 @@ async fn run_ws(client: HttpClient, channel: Channel) -> Result<(), AcpError> { .unbounded_send(Err(AcpError::parse_error().data(message))) .is_err() { - debug!("upstream channel closed; stopping WS reader"); - break; + debug!( + "upstream channel closed; discarding WS input while draining output" + ); + discard_incoming = true; } } } @@ -758,9 +826,7 @@ async fn run_ws(client: HttpClient, channel: Channel) -> Result<(), AcpError> { Some(Ok(WsMessage::Binary(_))) => { warn!("ignoring binary WebSocket frame (ACP uses text)"); } - Some(Ok( - WsMessage::Ping(_) | WsMessage::Pong(_) | WsMessage::Frame(_), - )) => {} + Some(Ok(WsMessage::Ping(_) | WsMessage::Pong(_) | WsMessage::Frame(_))) => {} Some(Ok(WsMessage::Close(frame))) => { debug!("server closed WebSocket: {frame:?}"); return Err(AcpError::internal_error() @@ -768,18 +834,20 @@ async fn run_ws(client: HttpClient, channel: Channel) -> Result<(), AcpError> { } Some(Err(e)) => { error!("WebSocket receive error: {e}"); - return Err(AcpError::internal_error() - .data(format!("ws recv: {e}"))); + return Err(AcpError::internal_error().data(format!("ws recv: {e}"))); } None => { return Err(AcpError::internal_error().data("WebSocket stream ended")); } - }, + } } - } + }; - drop(ws_tx.send(WsMessage::Close(None)).await); - Ok(()) + pin_mut!(writer, reader); + match futures::future::select(writer, reader).await { + futures::future::Either::Left((result, _)) + | futures::future::Either::Right((result, _)) => result, + } } #[cfg(test)] @@ -810,6 +878,152 @@ mod tests { use super::*; + struct PostsThenExitClient { + finish: Arc, + finished: Arc, + escaped_tx: futures::channel::oneshot::Sender< + futures::channel::mpsc::UnboundedSender>, + >, + } + + struct QueueOutgoingThenText { + text: Option, + outgoing: Option>>, + } + + struct RecordingWsSink(mpsc::UnboundedSender); + + struct BackpressuredWsSink { + output: mpsc::UnboundedSender, + started: mpsc::UnboundedSender<()>, + release: Option>, + } + + struct ReleaseBackpressureOnPoll { + started: mpsc::UnboundedReceiver<()>, + release: Option>, + } + + impl WsSink for RecordingWsSink { + async fn send(&mut self, message: WsMessage) -> Result<(), String> { + self.0 + .unbounded_send(message) + .map_err(|error| error.to_string()) + } + } + + impl WsSink for BackpressuredWsSink { + async fn send(&mut self, message: WsMessage) -> Result<(), String> { + self.output + .unbounded_send(message) + .map_err(|error| error.to_string())?; + if let Some(release) = self.release.take() { + self.started + .unbounded_send(()) + .map_err(|error| error.to_string())?; + release + .await + .map_err(|_| "mock WebSocket reader did not release send".to_string())?; + } + Ok(()) + } + } + + impl Stream for QueueOutgoingThenText { + type Item = Result; + + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + _cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + // Make input ready immediately after queueing output. If the output + // branch was polled first it was still empty, so either poll order + // deterministically selects this input frame first. + if let Some(outgoing) = self.outgoing.take() { + for method in ["custom/first", "custom/second"] { + outgoing + .unbounded_send(Ok(RawJsonRpcMessage::notification( + method.to_string(), + json!({}), + ) + .unwrap())) + .unwrap(); + } + } + if let Some(text) = self.text.take() { + return std::task::Poll::Ready(Some(Ok(text))); + } + std::task::Poll::Pending + } + } + + impl Stream for ReleaseBackpressureOnPoll { + type Item = Result; + + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + if let std::task::Poll::Ready(Some(())) = + std::pin::Pin::new(&mut self.started).poll_next(cx) + && let Some(release) = self.release.take() + { + let _result = release.send(()); + } + std::task::Poll::Pending + } + } + + impl ConnectTo for PostsThenExitClient { + async fn connect_to(self, agent: impl ConnectTo) -> Result<(), AcpError> { + let Self { + finish, + finished, + escaped_tx, + } = self; + let (mut channel, transport) = agent.into_channel_and_future(); + let client = async move { + escaped_tx.send(channel.tx.clone()).map_err(|_| { + AcpError::internal_error().data("escaped sender observer dropped") + })?; + channel + .tx + .unbounded_send(Ok(RawJsonRpcMessage::request( + "initialize".to_string(), + json!({}), + RequestId::Number(1), + ) + .unwrap())) + .map_err(|e| { + AcpError::internal_error().data(format!("send initialize: {e}")) + })?; + channel.rx.next().await.ok_or_else(|| { + AcpError::internal_error().data("initialize response channel closed") + })??; + + for method in ["custom/first", "custom/second"] { + channel + .tx + .unbounded_send(Ok(RawJsonRpcMessage::notification( + method.to_string(), + json!({}), + ) + .unwrap())) + .map_err(|e| { + AcpError::internal_error().data(format!("send {method}: {e}")) + })?; + } + + finish.notified().await; + finished.notify_one(); + Ok(()) + }; + + let ((), ()) = futures::try_join!(transport, client)?; + Ok(()) + } + } + #[test] fn new_targets_standard_acp_endpoint() { assert_eq!( @@ -1194,10 +1408,13 @@ mod tests { } #[tokio::test] - async fn outbound_posts_are_sent_in_order() { + async fn client_completion_drains_ordered_posts_in_order() { let first_started = Arc::new(Notify::new()); let release_first = Arc::new(Notify::new()); let second_seen = Arc::new(Notify::new()); + let finish_client = Arc::new(Notify::new()); + let client_finished = Arc::new(Notify::new()); + let (escaped_tx, escaped_rx) = futures::channel::oneshot::channel(); let app = Router::new().route( "/acp", post({ @@ -1236,40 +1453,14 @@ mod tests { axum::serve(listener, app).await.unwrap(); }); let client = HttpClient::new(format!("http://{addr}")).unwrap(); - let (mut caller, transport) = Channel::duplex(); - let transport = tokio::spawn(run(client, transport)); - - caller - .tx - .unbounded_send(Ok(RawJsonRpcMessage::request( - "initialize".to_string(), - json!({}), - RequestId::Number(1), - ) - .unwrap())) - .unwrap(); - let init_response = timeout(Duration::from_secs(1), caller.rx.next()) + let mut connection = tokio::spawn(client.connect_to(PostsThenExitClient { + finish: finish_client.clone(), + finished: client_finished.clone(), + escaped_tx, + })); + let escaped = timeout(Duration::from_secs(1), escaped_rx) .await .unwrap() - .unwrap() - .unwrap(); - assert!(matches!(init_response, RawJsonRpcMessage::Response(_))); - - caller - .tx - .unbounded_send(Ok(RawJsonRpcMessage::notification( - "custom/first".to_string(), - json!({}), - ) - .unwrap())) - .unwrap(); - caller - .tx - .unbounded_send(Ok(RawJsonRpcMessage::notification( - "custom/second".to_string(), - json!({}), - ) - .unwrap())) .unwrap(); timeout(Duration::from_secs(1), first_started.notified()) @@ -1282,13 +1473,33 @@ mod tests { "second POST must not be sent while the first POST is pending" ); + finish_client.notify_one(); + timeout(Duration::from_secs(1), client_finished.notified()) + .await + .unwrap(); + assert!( + timeout(Duration::from_millis(100), &mut connection) + .await + .is_err(), + "HTTP transport returned before its accepted POSTs completed" + ); + assert!( + escaped + .unbounded_send(Ok(RawJsonRpcMessage::notification( + "custom/too-late".to_string(), + json!({}), + ) + .unwrap())) + .is_err(), + "escaped client sender remained open after client completion" + ); + release_first.notify_one(); timeout(Duration::from_secs(1), second_seen.notified()) .await .unwrap(); - drop(caller); - timeout(Duration::from_secs(1), transport) + timeout(Duration::from_secs(1), connection) .await .unwrap() .unwrap() @@ -1650,6 +1861,97 @@ mod tests { server.abort(); } + #[tokio::test] + async fn websocket_drain_discards_incoming_after_receiver_closes() { + let (caller, transport) = Channel::duplex(); + let Channel { + tx: outgoing, + rx: incoming, + } = caller; + drop(incoming); + + let inbound = + RawJsonRpcMessage::notification("custom/inbound".to_string(), json!({})).unwrap(); + let inbound = WsMessage::Text(serde_json::to_string(&inbound).unwrap().into()); + let ws_rx = QueueOutgoingThenText { + text: Some(inbound), + outgoing: Some(outgoing), + }; + let (ws_output_tx, mut ws_output) = mpsc::unbounded(); + timeout( + Duration::from_secs(1), + drive_ws(RecordingWsSink(ws_output_tx), ws_rx, transport), + ) + .await + .unwrap() + .unwrap(); + let mut frames = Vec::new(); + while let Some(frame) = ws_output.next().await { + frames.push(frame); + } + + let messages = frames + .iter() + .filter_map(|frame| match frame { + WsMessage::Text(text) => { + Some(serde_json::from_str::(text.as_str()).unwrap()) + } + _ => None, + }) + .collect::>(); + let methods = messages + .iter() + .filter_map(method_for_message) + .collect::>(); + assert_eq!(methods, ["custom/first", "custom/second"]); + assert!(matches!(frames.last(), Some(WsMessage::Close(None)))); + } + + #[tokio::test] + async fn websocket_reader_runs_while_send_is_backpressured() { + let (caller, transport) = Channel::duplex(); + let Channel { + tx: outgoing, + rx: incoming, + } = caller; + drop(incoming); + outgoing + .unbounded_send(Ok(RawJsonRpcMessage::notification( + "custom/queued".to_string(), + json!({}), + ) + .unwrap())) + .unwrap(); + drop(outgoing); + + let (started_tx, started_rx) = mpsc::unbounded(); + let (release_tx, release_rx) = futures::channel::oneshot::channel(); + let (ws_output_tx, mut ws_output) = mpsc::unbounded(); + let ws_tx = BackpressuredWsSink { + output: ws_output_tx, + started: started_tx, + release: Some(release_rx), + }; + let ws_rx = ReleaseBackpressureOnPoll { + started: started_rx, + release: Some(release_tx), + }; + + timeout(Duration::from_secs(1), drive_ws(ws_tx, ws_rx, transport)) + .await + .expect("WebSocket reader was not polled while its writer was backpressured") + .unwrap(); + let frames = ws_output.by_ref().collect::>().await; + + let WsMessage::Text(text) = &frames[0] else { + panic!("queued message was not sent as WebSocket text"); + }; + let message = serde_json::from_str::(text.as_str()).unwrap(); + assert_eq!(method_for_message(&message), Some("custom/queued")); + assert!(matches!(frames.get(1), Some(WsMessage::Close(None)))); + assert_eq!(frames.len(), 2); + } + #[tokio::test] async fn peer_ws_close_fails_transport() { let app = Router::new().route("/acp", get(close_ws)); diff --git a/src/agent-client-protocol/Cargo.toml b/src/agent-client-protocol/Cargo.toml index f2f8a8f..72e752c 100644 --- a/src/agent-client-protocol/Cargo.toml +++ b/src/agent-client-protocol/Cargo.toml @@ -45,6 +45,7 @@ serde.workspace = true serde_json.workspace = true tracing.workspace = true uuid.workspace = true +async-io.workspace = true async-process.workspace = true blocking.workspace = true shell-words.workspace = true diff --git a/src/agent-client-protocol/src/acp_agent.rs b/src/agent-client-protocol/src/acp_agent.rs index 53357c2..cba309e 100644 --- a/src/agent-client-protocol/src/acp_agent.rs +++ b/src/agent-client-protocol/src/acp_agent.rs @@ -7,6 +7,7 @@ use std::collections::VecDeque; use std::path::PathBuf; use std::str::FromStr; use std::sync::Arc; +use std::time::Duration; use async_process::Child; use std::pin::pin; @@ -19,6 +20,7 @@ type DebugCallback = Arc; const STDERR_CAPTURE_LIMIT: usize = 64 * 1024; const STDERR_READ_BUFFER_SIZE: usize = 8 * 1024; const STDERR_LINE_TRUNCATION_MARKER: &str = "… [stderr line truncated]"; +const SHUTDOWN_GRACE_PERIOD: Duration = Duration::from_secs(1); /// Direction of a line being sent or received. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -426,19 +428,39 @@ async fn drain_stderr( } } -/// Waits for a child process and returns an error if it exits with non-zero status. -/// -/// The error message includes a bounded tail of stderr collected concurrently. -/// When dropped, the child process is killed. -async fn monitor_child( +struct ExitedChild { + guard: ChildGuard, + status: std::process::ExitStatus, + stderr_rx: futures::channel::oneshot::Receiver, +} + +/// Waits for the direct child process while retaining its process-group guard +/// and stderr receiver for exit reporting. +async fn wait_for_child( mut guard: ChildGuard, stderr_rx: futures::channel::oneshot::Receiver, -) -> Result<(), crate::Error> { +) -> Result { let status = guard .wait() .await .map_err(|e| crate::util::internal_error(format!("Failed to wait for process: {e}")))?; + Ok(ExitedChild { + guard, + status, + stderr_rx, + }) +} + +/// Reports an observed child exit, including a bounded stderr tail for a +/// nonzero status. +async fn finish_child_exit(child: ExitedChild) -> Result<(), crate::Error> { + let ExitedChild { + mut guard, + status, + stderr_rx, + } = child; + // A launcher may exit while a descendant remains alive holding inherited // stdio. Terminate the rest of the group before waiting for stderr EOF. guard.terminate(); @@ -446,7 +468,20 @@ async fn monitor_child( if status.success() { Ok(()) } else { - let stderr = stderr_rx.await.unwrap_or_default(); + let stderr = + match futures::future::select(stderr_rx, async_io::Timer::after(SHUTDOWN_GRACE_PERIOD)) + .await + { + futures::future::Either::Left((stderr, _)) => stderr.unwrap_or_default(), + futures::future::Either::Right((_, stderr_rx)) => { + tracing::debug!( + grace = ?SHUTDOWN_GRACE_PERIOD, + "Agent stderr remained open after process exit; reporting status without it" + ); + drop(stderr_rx); + String::new() + } + }; let message = if stderr.is_empty() { format!("Process exited with {status}") @@ -458,6 +493,84 @@ async fn monitor_child( } } +async fn await_protocol_shutdown_after_successful_child_exit( + protocol_future: F, + grace: Duration, +) -> Result<(), crate::Error> +where + F: std::future::Future> + Unpin, +{ + match futures::future::select(protocol_future, async_io::Timer::after(grace)).await { + futures::future::Either::Left((result, _)) => result, + futures::future::Either::Right((_, protocol_future)) => { + tracing::debug!( + ?grace, + "Protocol transport remained open after successful agent process exit; stopping it" + ); + drop(protocol_future); + Ok(()) + } + } +} + +async fn write_line_with_shutdown_timeout( + writer: &mut W, + line: String, + stdout_eof_rx: &mut Option>, + stdout_eof_seen: &mut bool, + grace: Duration, +) -> std::io::Result<()> +where + W: futures::AsyncWrite + Unpin + ?Sized, +{ + let write = Box::pin(crate::jsonrpc::write_line(writer, line)); + + if *stdout_eof_seen { + return await_write_during_shutdown(write, grace).await; + } + + let Some(stdout_eof) = stdout_eof_rx.as_mut() else { + return write.await; + }; + + match futures::future::select(write, stdout_eof).await { + futures::future::Either::Left((result, _)) => result, + futures::future::Either::Right((stdout_eof, write)) => { + *stdout_eof_rx = None; + if stdout_eof.is_err() { + // Dropping the incoming stream cancels the signal. Only an + // explicit send represents a clean EOF. + return write.await; + } + + *stdout_eof_seen = true; + await_write_during_shutdown(write, grace).await + } + } +} + +async fn await_write_during_shutdown(write: F, grace: Duration) -> std::io::Result<()> +where + F: std::future::Future> + Unpin, +{ + match futures::future::select(write, async_io::Timer::after(grace)).await { + futures::future::Either::Left((result, _)) => result, + futures::future::Either::Right((_, write)) => { + tracing::debug!( + ?grace, + "Pending protocol output did not drain after agent stdout closed" + ); + drop(write); + Err(std::io::Error::new( + std::io::ErrorKind::TimedOut, + format!( + "Agent closed its protocol output but pending protocol output did not drain within {grace:?}" + ), + )) + } + } +} + /// Roles that an ACP agent executable can potentially serve. pub trait AcpAgentCounterpartRole: Role {} @@ -471,7 +584,7 @@ impl crate::ConnectTo for Acp client: impl crate::ConnectTo, ) -> Result<(), crate::Error> { use futures::io::BufReader; - use futures::{AsyncBufReadExt, AsyncWriteExt, StreamExt}; + use futures::{AsyncBufReadExt, StreamExt}; let (child_stdin, child_stdout, child_stderr, child) = self.spawn_process()?; @@ -499,9 +612,9 @@ impl crate::ConnectTo for Acp // Create the guard eagerly so cancelling this connection before the // monitor is first polled still terminates the whole process group. - let child_monitor = monitor_child(ChildGuard(child), stderr_rx); + let child_wait = wait_for_child(ChildGuard(child), stderr_rx); - // Convert stdio to line streams with optional debug inspection + // Convert stdio to line streams with optional debug inspection. let incoming_lines: std::pin::Pin< Box> + Send>, > = if let Some(callback) = self.debug_callback.clone() { @@ -514,31 +627,49 @@ impl crate::ConnectTo for Acp Box::pin(BufReader::new(child_stdout).lines()) }; + // The JSON-RPC transport keeps polling stdout while it drains stdin. + // Signal physical EOF so a child that half-closes stdout and stops + // reading cannot hold a final write open forever. Dropping this stream + // merely cancels the signal and is not treated as EOF. + let (stdout_eof_tx, stdout_eof_rx) = futures::channel::oneshot::channel(); + let mut stdout_eof_tx = Some(stdout_eof_tx); + let mut incoming_lines = incoming_lines; + let incoming_lines = Box::pin(futures::stream::poll_fn(move |cx| { + let next = incoming_lines.as_mut().poll_next(cx); + if matches!(next, std::task::Poll::Ready(None)) + && let Some(stdout_eof_tx) = stdout_eof_tx.take() + { + let _ = stdout_eof_tx.send(()); + } + next + })); + // Create a sink that writes lines (with newlines) to stdin with optional debug logging let outgoing_sink: std::pin::Pin< Box + Send>, - > = if let Some(callback) = self.debug_callback.clone() { - Box::pin(futures::sink::unfold( - (child_stdin, callback), - async move |(mut writer, callback), line: String| { - callback(&line, LineDirection::Stdin); - let mut bytes = line.into_bytes(); - bytes.push(b'\n'); - writer.write_all(&bytes).await?; - Ok::<_, std::io::Error>((writer, callback)) - }, - )) - } else { - Box::pin(futures::sink::unfold( + > = Box::pin(futures::sink::unfold( + ( child_stdin, - async move |mut writer, line: String| { - let mut bytes = line.into_bytes(); - bytes.push(b'\n'); - writer.write_all(&bytes).await?; - Ok::<_, std::io::Error>(writer) - }, - )) - }; + self.debug_callback.clone(), + Some(stdout_eof_rx), + false, + ), + async move |(mut writer, callback, mut stdout_eof_rx, mut stdout_eof_seen), + line: String| { + if let Some(callback) = callback.as_ref() { + callback(&line, LineDirection::Stdin); + } + write_line_with_shutdown_timeout( + &mut writer, + line, + &mut stdout_eof_rx, + &mut stdout_eof_seen, + SHUTDOWN_GRACE_PERIOD, + ) + .await?; + Ok::<_, std::io::Error>((writer, callback, stdout_eof_rx, stdout_eof_seen)) + }, + )); // Race the protocol against child process exit. // Also run stderr collection concurrently. @@ -548,14 +679,44 @@ impl crate::ConnectTo for Acp ); let stderr_future = pin!(stderr_future); - let protocol_future = pin!(protocol_future); - let child_monitor = pin!(child_monitor); + let protocol_future = Box::pin(protocol_future); + let child_wait = Box::pin(child_wait); - // Run stderr reader alongside the main race + // Run stderr reader alongside the main race. Errors still stop the + // connection immediately. After protocol shutdown succeeds, give the + // child a bounded grace period so delayed failures remain observable + // without letting a non-exiting launcher hang shutdown forever. let main_race = async { - match futures::future::select(protocol_future, child_monitor).await { - futures::future::Either::Left((result, _)) - | futures::future::Either::Right((result, _)) => result, + match futures::future::select(protocol_future, child_wait).await { + futures::future::Either::Left((result, child_wait)) => { + result?; + match futures::future::select( + child_wait, + async_io::Timer::after(SHUTDOWN_GRACE_PERIOD), + ) + .await + { + futures::future::Either::Left((child, _)) => { + finish_child_exit(child?).await + } + futures::future::Either::Right((_, child_wait)) => { + tracing::debug!( + grace = ?SHUTDOWN_GRACE_PERIOD, + "Agent process did not exit after protocol shutdown; terminating it" + ); + drop(child_wait); + Ok(()) + } + } + } + futures::future::Either::Right((child, protocol_future)) => { + finish_child_exit(child?).await?; + await_protocol_shutdown_after_successful_child_exit( + protocol_future, + SHUTDOWN_GRACE_PERIOD, + ) + .await + } } }; @@ -797,6 +958,43 @@ mod tests { assert_eq!(*recorded.lock().unwrap(), ["partial"]); } + #[tokio::test] + async fn successful_child_exit_bounds_protocol_shutdown_cleanly() { + let grace = std::time::Duration::from_millis(10); + tokio::time::timeout( + std::time::Duration::from_secs(1), + await_protocol_shutdown_after_successful_child_exit( + futures::future::pending::>(), + grace, + ), + ) + .await + .expect("protocol shutdown wait should be bounded") + .expect("a successful child exit should stop the pending protocol cleanly"); + } + + #[tokio::test] + async fn successful_child_exit_preserves_ready_protocol_error() { + let error = await_protocol_shutdown_after_successful_child_exit( + futures::future::ready(Err(crate::util::internal_error( + "protocol failed during shutdown", + ))), + std::time::Duration::from_secs(1), + ) + .await + .expect_err("a ready protocol error should remain authoritative"); + let detail = error + .data + .as_ref() + .and_then(serde_json::Value::as_str) + .unwrap_or_default(); + + assert!( + detail.contains("protocol failed during shutdown"), + "unexpected protocol error: {error:?}" + ); + } + #[cfg(unix)] #[tokio::test] async fn large_unterminated_stderr_is_fully_drained() { @@ -828,6 +1026,97 @@ mod tests { assert!(tail.ends_with("ACP_END")); } + #[cfg(unix)] + #[tokio::test] + async fn protocol_eof_still_reports_nonzero_child_exit() { + let agent = AcpAgent::from_args([ + "/bin/sh", + "-c", + "exec 1>&-; cat >/dev/null; printf ACP_TEST_FAILURE_AFTER_STDOUT_EOF >&2; exit 17", + ]) + .unwrap(); + + let error = tokio::time::timeout( + std::time::Duration::from_secs(5), + Client.builder().connect_to(agent), + ) + .await + .expect("connection should finish after the child exits") + .expect_err("nonzero child exit after protocol EOF should be reported"); + let detail = error + .data + .as_ref() + .map(serde_json::Value::to_string) + .unwrap_or_default(); + + assert!( + detail.contains("exit status: 17"), + "child exit status should be preserved: {error:?}" + ); + assert!( + detail.contains("ACP_TEST_FAILURE_AFTER_STDOUT_EOF"), + "child stderr should be preserved: {error:?}" + ); + } + + #[cfg(unix)] + #[tokio::test] + async fn successful_child_exit_does_not_cancel_active_foreground() { + let agent = AcpAgent::from_args(["/bin/sh", "-c", "exit 0"]).unwrap(); + let (started_tx, started_rx) = futures::channel::oneshot::channel(); + let (closed_tx, closed_rx) = futures::channel::oneshot::channel(); + let (close_release_tx, close_release_rx) = futures::channel::oneshot::channel(); + let (release_tx, release_rx) = futures::channel::oneshot::channel(); + let connection = tokio::spawn( + Client + .builder() + .on_close(async move |_cx| { + closed_tx.send(()).map_err(|()| { + crate::Error::internal_error().data("close observer dropped") + })?; + close_release_rx.await.map_err(|_| { + crate::Error::internal_error().data("close callback release dropped") + }) + }) + .connect_with(agent, async move |_cx| { + started_tx.send(()).map_err(|()| { + crate::Error::internal_error().data("foreground observer dropped") + })?; + release_rx.await.map_err(|_| { + crate::Error::internal_error().data("foreground release dropped") + }) + }), + ); + + tokio::time::timeout(std::time::Duration::from_secs(5), started_rx) + .await + .expect("foreground should start") + .expect("foreground should report that it started"); + + tokio::time::timeout(std::time::Duration::from_secs(5), closed_rx) + .await + .expect("successful child exit should close the protocol transport") + .expect("successful child exit should invoke close callbacks"); + + tokio::time::sleep(SHUTDOWN_GRACE_PERIOD + std::time::Duration::from_millis(250)).await; + assert!( + !connection.is_finished(), + "successful child exit canceled active cleanup" + ); + + close_release_tx + .send(()) + .expect("clean child exit should preserve close callbacks"); + release_tx + .send(()) + .expect("clean child exit should preserve the foreground"); + tokio::time::timeout(std::time::Duration::from_secs(5), connection) + .await + .expect("released foreground should finish") + .expect("connection task should not panic") + .expect("successful child exit should remain a clean EOF"); + } + #[cfg(unix)] struct KillOnDrop(Option); @@ -920,6 +1209,67 @@ mod tests { assert!(exited, "descendant process {pid} remained alive"); } + #[cfg(unix)] + #[tokio::test] + async fn protocol_eof_terminates_a_child_that_does_not_exit() { + let (agent, mut pid_rx) = + wrapper_agent("echo ACP_TEST_CHILD_PID=$$ >&2; exec 1>&-; while :; do sleep 30; done"); + let mut connection: futures::future::BoxFuture<'static, Result<(), crate::Error>> = + Box::pin(Client.builder().connect_to(agent)); + let child_pid = reported_descendant_pid(&mut connection, &mut pid_rx).await; + let mut cleanup = KillOnDrop(Some(child_pid)); + + assert!(process_is_running(child_pid)); + tokio::time::timeout(std::time::Duration::from_secs(5), &mut connection) + .await + .expect("protocol shutdown should bound its child-exit wait") + .expect("clean protocol shutdown should terminate a non-exiting child"); + assert_process_exits(child_pid).await; + cleanup.disarm(); + } + + #[cfg(unix)] + #[tokio::test] + async fn protocol_eof_bounds_a_blocked_outgoing_drain() { + let (agent, mut pid_rx) = wrapper_agent( + "echo ACP_TEST_CHILD_PID=$$ >&2; exec 1>&-; sleep 30 & child=$!; wait \"$child\"", + ); + let (channel, mut connection) = crate::ConnectTo::::into_channel_and_future(agent); + let crate::Channel { + rx: _incoming, + tx: outgoing, + } = channel; + + let response = crate::RawJsonRpcMessage::response( + crate::schema::v1::RequestId::Number(1), + Ok(serde_json::json!({ "payload": "x".repeat(4 * 1024 * 1024) })), + ); + outgoing + .unbounded_send(Ok(response)) + .expect("response should be accepted before the connection starts"); + outgoing.close_channel(); + + let child_pid = reported_descendant_pid(&mut connection, &mut pid_rx).await; + let mut cleanup = KillOnDrop(Some(child_pid)); + + let error = tokio::time::timeout(std::time::Duration::from_secs(5), &mut connection) + .await + .expect("stdout EOF should bound a blocked outgoing drain") + .expect_err("an undelivered accepted response must not report success"); + let detail = error + .data + .as_ref() + .and_then(serde_json::Value::as_str) + .unwrap_or_default(); + assert!( + detail.contains("pending protocol output did not drain"), + "the error should identify the blocked outgoing drain: {error:?}" + ); + + assert_process_exits(child_pid).await; + cleanup.disarm(); + } + #[cfg(unix)] #[tokio::test] async fn test_connection_drop_kills_wrapper_descendant() { diff --git a/src/agent-client-protocol/src/component.rs b/src/agent-client-protocol/src/component.rs index c1dc283..1222863 100644 --- a/src/agent-client-protocol/src/component.rs +++ b/src/agent-client-protocol/src/component.rs @@ -120,6 +120,12 @@ pub trait ConnectTo: Send + 'static { /// /// A future that resolves when the component stops serving, either successfully /// or with an error. The future must be `Send`. + /// + /// A component that buffers outbound messages should not return `Ok(())` + /// merely because its client completed: it should first finish messages the + /// client already transferred to it. This lets wrappers preserve graceful + /// drain guarantees through to the physical transport sink. Errors may + /// still terminate the connection immediately. fn connect_to( self, client: impl ConnectTo, diff --git a/src/agent-client-protocol/src/concepts/connections.rs b/src/agent-client-protocol/src/concepts/connections.rs index 359afcb..44b3b06 100644 --- a/src/agent-client-protocol/src/concepts/connections.rs +++ b/src/agent-client-protocol/src/concepts/connections.rs @@ -58,6 +58,53 @@ //! # } //! ``` //! +//! # Clean Incoming EOF +//! +//! [`Builder::connect_to`](crate::Builder::connect_to) is reactive: it returns +//! `Ok(())` when the incoming transport reaches clean EOF, after draining +//! responses and notifications already accepted by its outgoing queue through +//! the transport sink. +//! [`Builder::connect_with`](crate::Builder::connect_with) is foreground-owned: +//! EOF fails pending requests, but does not cancel unrelated work in its +//! closure. This avoids dropping application futures at an arbitrary await +//! point. +//! +//! Use [`ConnectionTo::incoming_closed`](crate::ConnectionTo::incoming_closed) +//! to await EOF directly, or [`Builder::on_close`](crate::Builder::on_close) +//! for cleanup and application-specific shutdown policy: +//! +//! ``` +//! # use agent_client_protocol::{Client, ConnectTo, Error}; +//! # async fn example(transport: impl ConnectTo) -> Result<(), Error> { +//! Client.builder() +//! .on_close(async |_cx| { +//! // Notify application-owned work here. Returning an error also +//! // terminates a still-running connect_with foreground. +//! Ok(()) +//! }) +//! .connect_with(transport, async |cx| { +//! cx.incoming_closed().await; +//! Ok(()) +//! }) +//! .await?; +//! # Ok(()) +//! # } +//! ``` +//! +//! Every request still waiting for a response at EOF is completed with an +//! internal error whose data contains +//! `{"reason":"incoming_transport_closed","method":"..."}`. Requests made +//! after EOF fail the same way; use +//! [`is_incoming_transport_closed`](crate::is_incoming_transport_closed) to +//! identify this error. In `connect_with`, notification and response sends +//! remain available so applications can choose their own half-close policy; +//! reactive `connect_to` stops accepting them when its final drain begins. +//! +//! Pending requests are failed before close callbacks begin. The close signal +//! is published after callbacks finish, so a callback must not await +//! [`ConnectionTo::incoming_closed`](crate::ConnectionTo::incoming_closed) +//! itself. +//! //! # Sending Requests //! //! When you call `send_request()`, you get back a [`SentRequest`] that represents diff --git a/src/agent-client-protocol/src/jsonrpc.rs b/src/agent-client-protocol/src/jsonrpc.rs index 99d12e9..784d485 100644 --- a/src/agent-client-protocol/src/jsonrpc.rs +++ b/src/agent-client-protocol/src/jsonrpc.rs @@ -12,9 +12,8 @@ use std::collections::HashMap; use std::fmt::Debug; use std::panic::Location; use std::pin::pin; -use std::sync::Arc; use std::sync::{ - Mutex, + Arc, Mutex, Weak, atomic::{AtomicBool, Ordering}, }; use uuid::Uuid; @@ -743,6 +742,38 @@ where /// Protocol version mode for the public API and wire compatibility layer. protocol_mode: ProtocolMode, + + /// Callbacks run when the incoming transport reaches clean EOF. + on_close: Vec>, +} + +struct OnClose { + callback: Box< + dyn FnOnce(ConnectionTo) -> BoxFuture<'static, Result<(), crate::Error>> + + Send, + >, +} + +impl OnClose { + fn new(callback: F) -> Self + where + F: FnOnce(ConnectionTo) -> Fut + Send + 'static, + Fut: Future> + Send + 'static, + { + Self { + callback: Box::new(move |connection| callback(connection).boxed()), + } + } + + async fn run(self, connection: ConnectionTo) -> Result<(), crate::Error> { + (self.callback)(connection).await + } +} + +impl Debug for OnClose { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.debug_struct("OnClose").finish_non_exhaustive() + } } fn default_protocol_mode() -> ProtocolMode { @@ -768,6 +799,7 @@ impl Builder { handler: NullHandler, responder: NullRun, protocol_mode: default_protocol_mode::(), + on_close: Vec::new(), } } } @@ -784,6 +816,7 @@ where handler, responder: NullRun, protocol_mode: default_protocol_mode::(), + on_close: Vec::new(), } } } @@ -843,8 +876,11 @@ impl< handler: other_handler, responder: other_responder, protocol_mode: other_protocol_mode, + on_close: other_on_close, host: _, } = other; + let mut on_close = self.on_close; + on_close.extend(other_on_close); Builder { host: self.host, name: self.name, @@ -854,6 +890,7 @@ impl< ), responder: ChainRun::new(self.responder, other_responder), protocol_mode: self.protocol_mode.merge(other_protocol_mode), + on_close, } } @@ -871,6 +908,7 @@ impl< handler: ChainedHandler::new(self.handler, handler), responder: self.responder, protocol_mode: self.protocol_mode, + on_close: self.on_close, } } @@ -888,6 +926,7 @@ impl< handler: self.handler, responder: ChainRun::new(self.responder, responder), protocol_mode: self.protocol_mode, + on_close: self.on_close, } } @@ -905,6 +944,47 @@ impl< self.with_responder(SpawnedRun::new(location, task)) } + /// Run a callback when the incoming transport reaches clean EOF. + /// + /// Each callback runs at most once and receives the connection context. A + /// successful callback observes the close without otherwise changing the + /// lifetime of [`connect_with`](Self::connect_with). Returning an error + /// shuts down the connection and cancels a still-running `connect_with` + /// future. + /// + /// Multiple callbacks run sequentially in registration order. All of them + /// run even if an earlier callback fails, after which the first error is + /// returned. Pending requests are failed before callbacks begin, while + /// [`ConnectionTo::incoming_closed`] completes only after they finish. A + /// callback must therefore not await that close future itself. + /// + /// This separation lets applications choose their cancellation policy. A + /// callback can notify application-owned tasks and return `Ok(())` for + /// graceful cleanup, or return an error to stop them immediately. + /// + /// ``` + /// # use agent_client_protocol::{Client, ConnectTo, Error}; + /// # async fn example(transport: impl ConnectTo) -> Result<(), Error> { + /// Client.builder() + /// .on_close(async |_cx| { + /// Err(Error::internal_error().data("agent transport closed")) + /// }) + /// .connect_with(transport, async |_cx| { + /// std::future::pending().await + /// }) + /// .await?; + /// # Ok(()) + /// # } + /// ``` + pub fn on_close(mut self, callback: F) -> Self + where + F: FnOnce(ConnectionTo) -> Fut + Send + 'static, + Fut: Future> + Send + 'static, + { + self.on_close.push(OnClose::new(callback)); + self + } + /// Register a handler for messages that can be either requests OR notifications. /// /// Use this when you want to handle an enum type that contains both request and @@ -1295,6 +1375,10 @@ impl< /// - An error occurs /// - One of your handlers returns an error /// + /// On clean EOF, messages already accepted by the outgoing queue—including + /// handler responses and close-callback notifications—are drained through + /// the transport sink before this returns `Ok(())`. + /// /// The transport is responsible for serializing and deserializing [`RawJsonRpcMessage`] /// values to/from the underlying I/O mechanism (byte streams, channels, etc.). /// @@ -1325,8 +1409,11 @@ impl< self, transport: impl ConnectTo + 'static, ) -> Result<(), crate::Error> { - self.connect_with(transport, async move |_cx| future::pending().await) - .await + self.connect_with(transport, async move |cx| { + cx.incoming_closed().await; + cx.drain_outgoing().await + }) + .await } /// Run the connection until the provided closure completes. @@ -1335,8 +1422,12 @@ impl< /// 1. Running your registered handlers in the background to process incoming messages /// 2. Executing your `main_fn` closure with a [`ConnectionTo`] for sending requests/notifications /// - /// The connection stays active until your `main_fn` returns, then shuts down gracefully. - /// If the connection closes unexpectedly before `main_fn` completes, this returns an error. + /// The connection stays active until your `main_fn` returns, then shuts down. + /// Clean incoming EOF fails every pending request and makes future + /// requests fail immediately. It does not cancel unrelated work in + /// `main_fn`: that future may observe [`ConnectionTo::incoming_closed`], or + /// the builder can use [`on_close`](Self::on_close) to notify it or return + /// an error and stop it. /// /// Use this mode when you need to initiate communication (send requests/notifications) /// in addition to responding to incoming messages. For server-only mode where you just @@ -1384,7 +1475,10 @@ impl< /// /// # Errors /// - /// Returns an error if the connection closes before `main_fn` completes. + /// Returns an error if a handler, background task, transport, or close + /// callback fails, or if `main_fn` returns an error. Clean incoming EOF is + /// observable through [`ConnectionTo::incoming_closed`] and is not itself + /// an error in this mode. pub async fn connect_with( self, transport: impl ConnectTo + 'static, @@ -1409,23 +1503,44 @@ impl< responder, host: me, protocol_mode, + on_close, } = self; let (outgoing_tx, outgoing_rx) = mpsc::unbounded(); let (new_task_tx, new_task_rx) = mpsc::unbounded(); let (dynamic_handler_tx, dynamic_handler_rx) = mpsc::unbounded(); + let pending_replies = PendingReplies::default(); + + // Convert transport into server - this returns a channel for us to use + // and a future that runs the transport. + let transport_component = crate::DynConnectTo::new(transport); + let (transport_channel, transport_future) = transport_component.into_channel_and_future(); + let (transport_completion_tx, transport_completion_rx) = oneshot::channel(); + let transport_completion = transport_completion_rx + .map(|result| { + result.unwrap_or_else(|error| { + Err(crate::util::internal_error(format!( + "transport task dropped before reporting completion: {error}" + ))) + }) + }) + .boxed() + .shared(); + let connection = ConnectionTo::new( me.counterpart(), outgoing_tx, new_task_tx, dynamic_handler_tx, + transport_completion, + pending_replies.registrar(), + on_close, ); - - // Convert transport into server - this returns a channel for us to use - // and a future that runs the transport - let transport_component = crate::DynConnectTo::new(transport); - let (transport_channel, transport_future) = transport_component.into_channel_and_future(); - let spawn_result = connection.spawn(transport_future); + let spawn_result = connection.spawn(async move { + let result = transport_future.await; + drop(transport_completion_tx.send(result.clone())); + result + }); // Destructure the channel endpoints let Channel { @@ -1433,7 +1548,6 @@ impl< tx: transport_outgoing_tx, } = transport_channel; - let (reply_tx, reply_rx) = mpsc::unbounded(); let protocol_compat = ProtocolCompat::new(protocol_mode); let future = crate::util::instrument_with_connection_name(name, { @@ -1442,32 +1556,48 @@ impl< let () = spawn_result?; let background = async { - futures::try_join!( - // Protocol layer: OutgoingMessage -> RawJsonRpcMessage - outgoing_actor::outgoing_protocol_actor( - outgoing_rx, - reply_tx.clone(), - transport_outgoing_tx, - protocol_compat.clone(), - ), - // Protocol layer: RawJsonRpcMessage -> handler/reply routing - incoming_actor::incoming_protocol_actor( - me.counterpart(), - &connection, - transport_incoming_rx, - dynamic_handler_rx, - reply_rx, - handler, - protocol_compat, - ), - task_actor::task_actor(new_task_rx, &connection), - responder.run_with_connection_to(connection.clone()), - )?; - Ok(()) - }; + let incoming = incoming_actor::incoming_protocol_actor( + me.counterpart(), + &connection, + transport_incoming_rx, + dynamic_handler_rx, + pending_replies.clone(), + handler, + protocol_compat.clone(), + ); + let other_actors = async { + futures::try_join!( + // Protocol layer: OutgoingMessage -> RawJsonRpcMessage + outgoing_actor::outgoing_protocol_actor( + outgoing_rx, + pending_replies, + transport_outgoing_tx, + protocol_compat, + ), + task_actor::task_actor(new_task_rx, &connection), + responder.run_with_connection_to(connection.clone()), + )?; + Ok(()) + }; - crate::util::run_until(Box::pin(background), Box::pin(main_fn(connection.clone()))) + // EOF can wake a pending request consumer, which may make + // the task actor fail while close callbacks are running. + // Keep the incoming actor alive until those callbacks have + // all finished, just as we do when the foreground wakes. + run_until_connection_close( + incoming, + other_actors, + connection.incoming_closed.clone(), + ) .await + }; + + run_until_connection_close( + background, + main_fn(connection.clone()), + connection.incoming_closed.clone(), + ) + .await } }); @@ -1515,37 +1645,156 @@ impl std::fmt::Debug for ResponsePayload { } } -/// Message sent to the incoming actor for reply subscription management. -enum ReplyMessage { - /// Subscribe to receive a response for the given request id. - /// When a response with this id arrives, it will be sent through the oneshot - /// along with an ack channel that must be signaled when processing is complete. - /// The method name is stored to allow routing responses through typed handlers. - Subscribe { - id: RequestId, +struct PendingReply { + method: String, + role_id: RoleId, + sender: oneshot::Sender, + cancellation_disarm: SentRequestCancellationDisarm, +} - /// id of the peer this request was sent to - role_id: RoleId, +impl PendingReply { + fn fail(self, error: crate::Error) { + self.cancellation_disarm.disarm(); + if self + .sender + .send(ResponsePayload { + result: Err(error), + ack_tx: None, + }) + .is_err() + { + tracing::trace!(method = %self.method, "Pending request was already dropped"); + } + } - /// (original) method of the request -- the actual request may have been transformed - /// to a successor method, but this will reflect the method of the wrapped request - method: String, + fn fail_incoming_closed(self) { + let error = incoming_transport_closed_error(&self.method); + self.fail(error); + } +} - sender: oneshot::Sender, +#[derive(Default)] +struct PendingRepliesInner { + incoming_closed: bool, + replies: HashMap, +} - cancellation_disarm: SentRequestCancellationDisarm, - }, +#[derive(Clone, Default)] +struct PendingReplies { + inner: Arc>, } -impl std::fmt::Debug for ReplyMessage { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - ReplyMessage::Subscribe { id, method, .. } => f - .debug_struct("Subscribe") - .field("id", id) - .field("method", method) - .finish(), +impl PendingReplies { + fn registrar(&self) -> PendingRepliesRegistrar { + PendingRepliesRegistrar { + inner: Arc::downgrade(&self.inner), + } + } + + fn contains(&self, id: &RequestId) -> bool { + self.inner + .lock() + .expect("pending replies mutex poisoned") + .replies + .contains_key(id) + } + + fn remove(&self, id: &RequestId) -> Option { + self.inner + .lock() + .expect("pending replies mutex poisoned") + .replies + .remove(id) + } + + /// Atomically reject new subscriptions and fail every existing one. + fn close_incoming(&self) -> usize { + let replies = { + let mut inner = self.inner.lock().expect("pending replies mutex poisoned"); + inner.incoming_closed = true; + std::mem::take(&mut inner.replies) + }; + let count = replies.len(); + for (_, reply) in replies { + reply.fail_incoming_closed(); } + count + } +} + +/// A non-owning handle used to register a request before it enters the +/// outgoing queue. Keeping this weak prevents escaped [`ConnectionTo`] clones +/// from extending the lifetime of response senders after the driver stops. +#[derive(Clone)] +struct PendingRepliesRegistrar { + inner: Weak>, +} + +impl PendingRepliesRegistrar { + /// Register a response destination before the request becomes observable. + /// + /// Returns `false` after failing `reply` when EOF has already made a + /// response impossible or the connection driver is no longer running. + fn subscribe( + &self, + id: RequestId, + reply: PendingReply, + incoming_closed: &IncomingClosed, + ) -> bool { + let Some(inner) = self.inner.upgrade() else { + if incoming_closed.is_closing() { + reply.fail_incoming_closed(); + } else { + let method = reply.method.clone(); + reply.fail(crate::util::internal_error(format!( + "failed to send outgoing request `{method}`: connection is no longer running" + ))); + } + return false; + }; + + let result = { + let mut inner = inner.lock().expect("pending replies mutex poisoned"); + if inner.incoming_closed { + Err(reply) + } else { + Ok(inner.replies.insert(id, reply)) + } + }; + + match result { + Err(rejected) => { + rejected.fail_incoming_closed(); + false + } + Ok(replaced) => { + if let Some(replaced) = replaced { + replaced.fail( + crate::Error::internal_error() + .data("outgoing request ID was reused before its response arrived"), + ); + } + true + } + } + } + + fn remove(&self, id: &RequestId) -> Option { + self.inner + .upgrade()? + .lock() + .expect("pending replies mutex poisoned") + .replies + .remove(id) + } +} + +impl Debug for PendingRepliesRegistrar { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("PendingRepliesRegistrar") + .field("is_connected", &(self.inner.strong_count() > 0)) + .finish() } } @@ -1929,6 +2178,10 @@ pub fn is_cancel_request_notification(notification: &N) /// Messages send to be serialized over the transport. #[derive(Debug)] enum OutgoingMessage { + /// Close the outgoing application queue and acknowledge after every + /// already-accepted message has entered the raw transport queue. + CloseAfterDraining { done: oneshot::Sender<()> }, + /// Send a request to the server. Request { /// id assigned to this request (generated by sender) @@ -1937,17 +2190,9 @@ enum OutgoingMessage { /// the original method method: String, - /// the peer we sent this to - role_id: RoleId, - /// the message to send; this may have a distinct method /// depending on the peer untyped: UntypedMessage, - - /// where to send the response when it arrives (includes ack channel) - response_tx: oneshot::Sender, - - cancellation_disarm: SentRequestCancellationDisarm, }, /// Send a notification to the server. @@ -2047,6 +2292,145 @@ pub struct ConnectionTo { message_tx: OutgoingMessageTx, task_tx: TaskTx, dynamic_handler_tx: mpsc::UnboundedSender>, + transport_completion: SharedTransportCompletion, + pending_replies: PendingRepliesRegistrar, + incoming_closed: IncomingClosed, + on_close: Arc>>>>, +} + +type SharedTransportCompletion = future::Shared>>; + +#[derive(Clone)] +struct IncomingClosed { + state: Arc, +} + +struct IncomingClosedState { + closing: AtomicBool, + closed: AtomicBool, + signal_tx: Mutex>>, + signal_rx: future::Shared>, +} + +impl IncomingClosed { + fn new() -> Self { + let (signal_tx, signal_rx) = oneshot::channel(); + Self { + state: Arc::new(IncomingClosedState { + closing: AtomicBool::new(false), + closed: AtomicBool::new(false), + signal_tx: Mutex::new(Some(signal_tx)), + signal_rx: signal_rx.map(|_| ()).boxed().shared(), + }), + } + } + + fn begin_close(&self) { + self.state.closing.store(true, Ordering::Release); + } + + fn finish_close(&self) { + self.state.closed.store(true, Ordering::Release); + let signal_tx = self + .state + .signal_tx + .lock() + .expect("incoming-close signal mutex poisoned") + .take(); + + if let Some(signal_tx) = signal_tx { + let _ = signal_tx.send(()); + } + } + + async fn closed(&self) { + self.state.signal_rx.clone().await; + } + + fn is_closed(&self) -> bool { + self.state.closed.load(Ordering::Acquire) + } + + fn is_closing(&self) -> bool { + self.state.closing.load(Ordering::Acquire) + } +} + +impl Debug for IncomingClosed { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("IncomingClosed") + .field("is_closing", &self.is_closing()) + .field("is_closed", &self.is_closed()) + .finish_non_exhaustive() + } +} + +/// Stable discriminator stored in the `data.reason` field of errors produced +/// when the incoming transport reaches clean EOF before a request receives its +/// response. +pub const INCOMING_TRANSPORT_CLOSED_REASON: &str = "incoming_transport_closed"; + +/// Return whether `error` reports that the incoming transport reached clean +/// EOF before a request received its response. +#[must_use] +pub fn is_incoming_transport_closed(error: &crate::Error) -> bool { + error + .data + .as_ref() + .and_then(|data| data.get("reason")) + .and_then(serde_json::Value::as_str) + == Some(INCOMING_TRANSPORT_CLOSED_REASON) +} + +fn incoming_transport_closed_error(method: &str) -> crate::Error { + let mut error = crate::Error::internal_error(); + error.message = "Incoming transport closed".to_string(); + error.data(serde_json::json!({ + "reason": INCOMING_TRANSPORT_CLOSED_REASON, + "method": method, + })) +} + +/// Run the connection background alongside its foreground while ensuring that +/// a foreground woken by incoming EOF cannot cancel close callbacks midway. +fn run_until_connection_close( + background: impl Future>, + foreground: impl Future>, + incoming_closed: IncomingClosed, +) -> impl Future> { + // Box these before constructing the returned future. Keeping the generic + // connection actors directly in this async state would substantially grow + // every `connect_*` future. + let background = Box::pin(background); + let foreground = Box::pin(foreground); + + async move { + match future::select(background, foreground).await { + Either::Left((background_result, foreground)) => { + background_result?; + foreground.await + } + Either::Right((foreground_result, background)) => { + if !incoming_closed.is_closing() { + return foreground_result; + } + + match future::select(background, Box::pin(incoming_closed.closed())).await { + Either::Left((background_result, _)) => { + background_result?; + foreground_result + } + Either::Right(((), background)) => { + // Poll the background first once more so an error returned + // by the just-finished close callback wins over the ready + // foreground result. + crate::util::run_until(background, future::ready(foreground_result)).await + } + } + } + } + } } impl ConnectionTo { @@ -2055,12 +2439,19 @@ impl ConnectionTo { message_tx: mpsc::UnboundedSender, task_tx: mpsc::UnboundedSender, dynamic_handler_tx: mpsc::UnboundedSender>, + transport_completion: SharedTransportCompletion, + pending_replies: PendingRepliesRegistrar, + on_close: Vec>, ) -> Self { Self { counterpart, message_tx, task_tx, dynamic_handler_tx, + transport_completion, + pending_replies, + incoming_closed: IncomingClosed::new(), + on_close: Arc::new(Mutex::new(Some(on_close))), } } @@ -2069,6 +2460,70 @@ impl ConnectionTo { self.counterpart.clone() } + /// Wait until the incoming transport reaches clean EOF. + /// + /// Transport closure means that no more messages or responses can arrive. + /// Pending requests are failed first; this completes after registered + /// [`Builder::on_close`] callbacks finish. + /// It does not automatically cancel the future passed to + /// [`Builder::connect_with`]; use [`Builder::on_close`] when the connection + /// should run application-specific cleanup or terminate that future. + pub async fn incoming_closed(&self) { + self.incoming_closed.closed().await; + } + + /// Return whether clean incoming-EOF processing has completed. + /// + /// This remains `false` while [`Builder::on_close`] callbacks are running. + #[must_use] + pub fn is_incoming_closed(&self) -> bool { + self.incoming_closed.is_closed() + } + + /// Stop accepting outgoing messages, drain those already accepted through + /// the protocol actor, and wait for the transport sink to finish them. + async fn drain_outgoing(&self) -> Result<(), crate::Error> { + let (done_tx, done_rx) = oneshot::channel(); + let marker_result = send_raw_message( + &self.message_tx, + OutgoingMessage::CloseAfterDraining { done: done_tx }, + ); + let marker_result = match marker_result { + Ok(()) => done_rx.await.map_err(|error| { + crate::util::internal_error(format!( + "outgoing drain marker was dropped before completion: {error}" + )) + }), + Err(error) => Err(error), + }; + + // The marker only proves that all accepted protocol messages entered + // the raw transport queue. Transport completion is the sink-level + // barrier that proves a backpressured writer finished them. + self.transport_completion.clone().await?; + marker_result + } + + fn is_incoming_closing(&self) -> bool { + self.incoming_closed.is_closing() + } + + pub(super) fn begin_incoming_close(&self) { + self.incoming_closed.begin_close(); + } + + pub(super) fn finish_incoming_close(&self) { + self.incoming_closed.finish_close(); + } + + fn take_on_close(&self) -> Vec> { + self.on_close + .lock() + .expect("connection close callback mutex poisoned") + .take() + .unwrap_or_default() + } + /// Spawns a task that will run so long as the JSON-RPC connection is being served. /// /// This is the primary mechanism for offloading expensive work from handler callbacks @@ -2309,40 +2764,61 @@ impl ConnectionTo { let remote_style = self.counterpart.remote_style(peer); let cancellation = SentRequestCancellation::new(self.message_tx.clone(), remote_style, id.clone()); + if self.is_incoming_closing() { + cancellation.disarm(); + drop(response_tx.send(ResponsePayload { + result: Err(incoming_transport_closed_error(&method)), + ack_tx: None, + })); + return SentRequest::new( + id, + method.clone(), + self.task_tx.clone(), + response_rx, + cancellation, + ) + .map(move |json| ::from_value(&method, json)); + } + match remote_style.transform_outgoing_message(request) { Ok(untyped) => { - // Transform the message for the target role - let message = OutgoingMessage::Request { - id: id.clone(), + // Register before enqueueing so incoming EOF can fail every + // observable request before close callbacks begin. The + // outgoing actor checks that the registration still exists + // before sending the request. + let pending_reply = PendingReply { method: method.clone(), role_id, - untyped, - response_tx, + sender: response_tx, cancellation_disarm: cancellation.disarm_handle(), }; - match self.message_tx.unbounded_send(message) { - Ok(()) => (), - Err(error) => { + if self + .pending_replies + .subscribe(id.clone(), pending_reply, &self.incoming_closed) + { + let message = OutgoingMessage::Request { + id: id.clone(), + method: method.clone(), + untyped, + }; + + if let Err(error) = self.message_tx.unbounded_send(message) { cancellation.disarm(); - let OutgoingMessage::Request { - method, - response_tx, - .. - } = error.into_inner() - else { + let OutgoingMessage::Request { id, method, .. } = error.into_inner() else { unreachable!(); }; - response_tx - .send(ResponsePayload { - result: Err(crate::util::internal_error(format!( - "failed to send outgoing request `{method}" - ))), - ack_tx: None, - }) - .unwrap(); + if let Some(pending_reply) = self.pending_replies.remove(&id) { + if self.is_incoming_closing() { + pending_reply.fail_incoming_closed(); + } else { + pending_reply.fail(crate::util::internal_error(format!( + "failed to send outgoing request `{method}`" + ))); + } + } } } } @@ -3504,6 +3980,15 @@ impl JsonRpcNotification for UntypedMessage {} /// the request, then discards the response when it arrives. Requests whose /// eventual response should be ignored, but which should keep running on the /// peer, should use [`detach`](Self::detach) instead. +/// +/// # Incoming Transport EOF +/// +/// If the incoming transport reaches clean EOF before the response arrives, every +/// consumption mode receives an error with the message `Incoming transport +/// closed` and data containing +/// `{"reason":"incoming_transport_closed","method":"..."}`. Requests made +/// after incoming EOF fail immediately with the same error. Use +/// [`is_incoming_transport_closed`] to identify it. #[must_use = "dropping a SentRequest asks the peer to cancel the request and \ discards the response; consume it with `block_task`, \ `on_receiving_result`, `forward_response_to`, or `detach`"] @@ -3842,9 +4327,11 @@ impl SentRequest { /// This is equivalent to calling `on_receiving_result` and manually forwarding /// the result, with two proxy-specific additions: /// - /// - If the pending response is dropped without ever being delivered (for - /// example, the downstream connection closed), the incoming request is + /// - If the pending response cannot be delivered, the incoming request is /// answered with an internal error instead of being left unanswered. + /// Known clean incoming EOF is delivered like any other response + /// error; an unexpected response-channel loss is forwarded as an outer + /// consumption error. /// - When the peer cancels the incoming request, the cancellation is /// forwarded to the outgoing request, and the downstream response /// (normal data or a cancellation error) is still forwarded back. This is @@ -3858,9 +4345,8 @@ impl SentRequest { let this = self.forward_cancellation_from(responder.cancellation()); this.consume_with(async move |response| { - // A response that was never delivered (outer `Err`, e.g. the - // downstream connection closed) is forwarded as an error: the - // incoming request must not be left unanswered. + // An unexpected response-channel loss (outer `Err`) is forwarded + // as an error: the incoming request must not be left unanswered. responder.respond_with_result(response.unwrap_or_else(Err)) }) } @@ -3874,9 +4360,10 @@ impl SentRequest { /// the typed result (`Ok(Result)`). The dispatch loop's ack, if any, /// is sent after `handle` completes. /// - /// If the pending response is dropped without ever being delivered (for - /// example, the connection closed), `handle` receives the outer `Err` - /// describing the loss; there is no ack in that case. + /// Clean incoming EOF is delivered as `Ok(Err(error))`, just like + /// a peer response error, so callback-style consumers still run. If the + /// response channel disappears for another reason, `handle` receives an + /// outer `Err` describing that unexpected loss; there is no ack then. #[track_caller] fn consume_with( self, @@ -4241,9 +4728,59 @@ where IncomingStream: futures::Stream> + Send + 'static, { async fn connect_to(self, client: impl ConnectTo) -> Result<(), crate::Error> { - let (channel, serve_self) = ConnectTo::::into_channel_and_future(self); + let Self { outgoing, incoming } = self; + let (channel, transport_channel) = Channel::duplex(); + let shutdown_tx = channel.tx.clone(); + let Channel { rx, tx } = transport_channel; + + // Once the client completes successfully, its incoming channel is + // gone. Keep consuming successful messages from the physical read + // half without forwarding them so a full-duplex peer cannot block our + // outgoing sink while it is being drained. Transport errors must still + // fail the connection. + let discard_incoming = Arc::new(AtomicBool::new(false)); + let incoming = incoming.filter_map({ + let discard_incoming = discard_incoming.clone(); + move |item| { + let discard_incoming = discard_incoming.load(Ordering::Acquire); + future::ready((!discard_incoming || item.is_err()).then_some(item)) + } + }); + + let outgoing = transport_actor::transport_outgoing_lines_actor(rx, outgoing) + .boxed() + .shared(); + let serve_self = Box::pin({ + let outgoing = outgoing.clone(); + async move { + futures::try_join!( + outgoing, + transport_actor::transport_incoming_lines_actor(incoming, tx), + )?; + Ok(()) + } + }); + match futures::future::select(Box::pin(client.connect_to(channel)), serve_self).await { - Either::Left((result, _)) | Either::Right((result, _)) => result, + Either::Left((result, serve_self)) => { + result?; + shutdown_tx.close_channel(); + discard_incoming.store(true, Ordering::Release); + + // Drive the read half while waiting for the write half, but do + // not require the peer's independent incoming stream to reach + // EOF. If incoming processing finishes successfully first, + // the shared outgoing future still owns and drains the sink. + // A successful `serve_self` result includes its shared + // outgoing clone, while any error must remain authoritative + // instead of being hidden behind the other handle. Poll it + // first so a ready read error wins over clean outgoing + // completion. + match future::select(serve_self, outgoing).await { + Either::Left((result, _)) | Either::Right((result, _)) => result, + } + } + Either::Right((result, _)) => result, } } @@ -4326,6 +4863,38 @@ where pub fn new(outgoing: OB, incoming: IB) -> Self { Self { outgoing, incoming } } + + fn into_lines( + self, + ) -> Lines< + impl futures::Sink + Send + 'static, + impl futures::Stream> + Send + 'static, + > { + use futures::AsyncBufReadExt; + use futures::io::BufReader; + let Self { outgoing, incoming } = self; + + let incoming_lines = Box::pin(BufReader::new(incoming).lines()); + let outgoing_lines = + futures::sink::unfold(Box::pin(outgoing), async move |mut writer, line: String| { + write_line(&mut writer, line).await?; + Ok::<_, std::io::Error>(writer) + }); + + Lines::new(outgoing_lines, incoming_lines) + } +} + +pub(crate) async fn write_line(writer: &mut W, line: String) -> std::io::Result<()> +where + W: AsyncWrite + Unpin + ?Sized, +{ + use futures::AsyncWriteExt as _; + + let mut bytes = line.into_bytes(); + bytes.push(b'\n'); + writer.write_all(&bytes).await?; + writer.flush().await } impl ConnectTo for ByteStreams @@ -4334,34 +4903,11 @@ where IB: AsyncRead + Send + 'static, { async fn connect_to(self, client: impl ConnectTo) -> Result<(), crate::Error> { - let (channel, serve_self) = ConnectTo::::into_channel_and_future(self); - match futures::future::select(pin!(client.connect_to(channel)), serve_self).await { - Either::Left((result, _)) | Either::Right((result, _)) => result, - } + ConnectTo::::connect_to(self.into_lines(), client).await } fn into_channel_and_future(self) -> (Channel, BoxFuture<'static, Result<(), crate::Error>>) { - use futures::AsyncBufReadExt; - use futures::AsyncWriteExt; - use futures::io::BufReader; - let Self { outgoing, incoming } = self; - - // Convert byte streams to line streams - // Box both streams to satisfy Unpin requirements - let incoming_lines = Box::pin(BufReader::new(incoming).lines()); - - // Create a sink that writes lines (with newlines) to the outgoing byte stream - // We need to Box the writer since it may not be Unpin - let outgoing_sink = - futures::sink::unfold(Box::pin(outgoing), async move |mut writer, line: String| { - let mut bytes = line.into_bytes(); - bytes.push(b'\n'); - writer.write_all(&bytes).await?; - Ok::<_, std::io::Error>(writer) - }); - - // Delegate to Lines component - ConnectTo::::into_channel_and_future(Lines::new(outgoing_sink, incoming_lines)) + ConnectTo::::into_channel_and_future(self.into_lines()) } } @@ -4465,6 +5011,16 @@ impl ConnectTo for Channel { mod tests { use super::*; + #[tokio::test] + async fn write_line_flushes_buffered_writers() { + let mut writer = + futures::io::BufWriter::with_capacity(4096, futures::io::Cursor::new(Vec::new())); + + write_line(&mut writer, "message".into()).await.unwrap(); + + assert_eq!(writer.into_inner().into_inner(), b"message\n"); + } + #[test] fn peel_successor_envelopes_returns_plain_messages_unchanged() { let params = serde_json::json!({ "key": "value" }); diff --git a/src/agent-client-protocol/src/jsonrpc/incoming_actor.rs b/src/agent-client-protocol/src/jsonrpc/incoming_actor.rs index a19c923..f94792b 100644 --- a/src/agent-client-protocol/src/jsonrpc/incoming_actor.rs +++ b/src/agent-client-protocol/src/jsonrpc/incoming_actor.rs @@ -1,22 +1,20 @@ // Types re-exported from crate root -use std::collections::HashMap; - use futures::StreamExt as _; use futures::channel::mpsc; -use futures::channel::oneshot; +use futures::stream; use futures_concurrency::stream::StreamExt as _; use rustc_hash::FxHashMap; use uuid::Uuid; use crate::Dispatch; -use crate::RoleId; use crate::UntypedMessage; use crate::jsonrpc::ConnectionTo; use crate::jsonrpc::HandleDispatchFrom; use crate::jsonrpc::OutgoingMessage; +use crate::jsonrpc::PendingReplies; +use crate::jsonrpc::PendingReply; use crate::jsonrpc::RawJsonRpcMessage; use crate::jsonrpc::RawJsonRpcParams; -use crate::jsonrpc::ReplyMessage; use crate::jsonrpc::Responder; use crate::jsonrpc::ResponseRouter; use crate::jsonrpc::dynamic_handler::DynHandleDispatchFrom; @@ -29,20 +27,12 @@ use crate::schema::v1::{RequestId, Response}; use super::Handled; -struct PendingReply { - method: String, - role_id: RoleId, - sender: oneshot::Sender, - cancellation_disarm: super::SentRequestCancellationDisarm, -} - /// Incoming protocol actor: The central dispatch loop for a connection. /// /// This actor handles JSON-RPC protocol semantics: /// - Routes responses to pending request awaiters /// - Routes requests/notifications to registered handlers /// - Converts RawJsonRpcMessage requests/notifications to UntypedMessage for handlers -/// - Manages reply subscriptions from outgoing requests /// /// This is the protocol layer - it has no knowledge of how messages arrived. /// @@ -53,14 +43,19 @@ pub(super) async fn incoming_protocol_actor( connection: &ConnectionTo, transport_rx: mpsc::UnboundedReceiver>, dynamic_handler_rx: mpsc::UnboundedReceiver>, - reply_rx: mpsc::UnboundedReceiver, + pending_replies: PendingReplies, mut handler: impl HandleDispatchFrom, protocol_compat: ProtocolCompat, ) -> Result<(), crate::Error> { - let mut my_rx = transport_rx - .map(IncomingProtocolMsg::Transport) - .merge(dynamic_handler_rx.map(IncomingProtocolMsg::DynamicHandler)) - .merge(reply_rx.map(IncomingProtocolMsg::Reply)); + // `merge` does not expose when one of its source streams ends. Preserve + // transport EOF as an explicit event so the other, connection-internal + // streams cannot hide it. + let transport_with_close = futures::StreamExt::chain( + transport_rx.map(IncomingProtocolMsg::Transport), + stream::iter([IncomingProtocolMsg::TransportClosed]), + ); + let mut my_rx = + transport_with_close.merge(dynamic_handler_rx.map(IncomingProtocolMsg::DynamicHandler)); let mut dynamic_handlers: FxHashMap>> = FxHashMap::default(); @@ -68,33 +63,30 @@ pub(super) async fn incoming_protocol_actor( let request_cancellations = super::RequestCancellationRegistry::new(); - // Map from request ID to (method, sender) for response dispatch. - // The method is stored to allow routing responses through typed handlers. - let mut pending_replies: HashMap = HashMap::new(); - while let Some(message_result) = my_rx.next().await { tracing::trace!(message = ?message_result, actor = "incoming_protocol_actor"); match message_result { - IncomingProtocolMsg::Reply(message) => match message { - ReplyMessage::Subscribe { - id, - role_id, - method, - sender, - cancellation_disarm, - } => { - tracing::trace!(?id, %method, "incoming_actor: subscribing to response"); - pending_replies.insert( - id, - PendingReply { - method, - role_id, - sender, - cancellation_disarm, - }, - ); + IncomingProtocolMsg::TransportClosed => { + connection.begin_incoming_close(); + let pending_reply_count = pending_replies.close_incoming(); + tracing::debug!(pending_reply_count, "Incoming transport closed"); + + let mut callback_error = None; + for callback in connection.take_on_close() { + if let Err(error) = callback.run(connection.clone()).await { + tracing::warn!(?error, "Connection close callback failed"); + if callback_error.is_none() { + callback_error = Some(error); + } + } } - }, + + connection.finish_incoming_close(); + + if let Some(error) = callback_error { + return Err(error); + } + } IncomingProtocolMsg::DynamicHandler(message) => match message { DynamicHandlerMessage::AddDynamicHandler(uuid, mut handler) => { @@ -250,8 +242,8 @@ pub(super) async fn incoming_protocol_actor( #[derive(Debug)] enum IncomingProtocolMsg { Transport(Result), + TransportClosed, DynamicHandler(DynamicHandlerMessage), - Reply(ReplyMessage), } /// Dispatches a JSON-RPC request to the handler. diff --git a/src/agent-client-protocol/src/jsonrpc/outgoing_actor.rs b/src/agent-client-protocol/src/jsonrpc/outgoing_actor.rs index ae7a208..6175960 100644 --- a/src/agent-client-protocol/src/jsonrpc/outgoing_actor.rs +++ b/src/agent-client-protocol/src/jsonrpc/outgoing_actor.rs @@ -2,9 +2,8 @@ use futures::StreamExt as _; use futures::channel::mpsc; -use crate::jsonrpc::ReplyMessage; use crate::jsonrpc::protocol_compat::ProtocolCompat; -use crate::jsonrpc::{OutgoingMessage, RawJsonRpcMessage}; +use crate::jsonrpc::{OutgoingMessage, PendingReplies, RawJsonRpcMessage}; use crate::schema::v1::RequestId; pub type OutgoingMessageTx = mpsc::UnboundedSender; @@ -21,29 +20,42 @@ pub(crate) fn send_raw_message( /// Outgoing protocol actor: Converts application-level OutgoingMessage to protocol-level RawJsonRpcMessage. /// /// This actor handles JSON-RPC protocol semantics: -/// - Subscribes to reply_actor for response correlation +/// - Verifies that outgoing requests still have pending response registrations /// - Converts OutgoingMessage variants to RawJsonRpcMessage /// /// This is the protocol layer - it has no knowledge of how messages are transported. pub(super) async fn outgoing_protocol_actor( mut outgoing_rx: mpsc::UnboundedReceiver, - reply_tx: mpsc::UnboundedSender, + pending_replies: PendingReplies, transport_tx: mpsc::UnboundedSender>, protocol_compat: ProtocolCompat, ) -> Result<(), crate::Error> { + let mut drain_waiters = Vec::new(); + while let Some(message) = outgoing_rx.next().await { tracing::debug!(?message, "outgoing_protocol_actor"); // Create the message to be sent over the transport let json_rpc_message = match message { + OutgoingMessage::CloseAfterDraining { done } => { + // Reject later sends while preserving every message that was + // already accepted into this receiver's buffer. + outgoing_rx.close(); + drain_waiters.push(done); + continue; + } OutgoingMessage::Request { id, - role_id, method, untyped, - response_tx, - cancellation_disarm, } => { + // Requests register their response destination synchronously + // before entering this queue. EOF removes that registration, + // so skip work that can no longer receive a response. + if !pending_replies.contains(&id) { + continue; + } + let request = match protocol_compat .outgoing_message(untyped) .and_then(|untyped| untyped.into_raw_jsonrpc_message(Some(id.clone()))) @@ -51,24 +63,25 @@ pub(super) async fn outgoing_protocol_actor( Ok(request) => request, Err(error) => { tracing::warn!(?id, %method, ?error, "Failed to prepare outgoing request"); - cancellation_disarm.disarm(); - complete_request_with_error(response_tx, error); + if let Some(pending_reply) = pending_replies.remove(&id) { + pending_reply.fail(error); + } continue; } }; - // Record where the reply should be sent once it arrives. - reply_tx - .unbounded_send(ReplyMessage::Subscribe { - id: id.clone(), - role_id, - method, - sender: response_tx, - cancellation_disarm, - }) - .map_err(crate::Error::into_internal_error)?; + if !pending_replies.contains(&id) { + continue; + } - request + if let Err(error) = transport_tx.unbounded_send(Ok(request)) { + let error = crate::Error::into_internal_error(error); + if let Some(pending_reply) = pending_replies.remove(&id) { + pending_reply.fail(error.clone()); + } + return Err(error); + } + continue; } OutgoingMessage::Notification { untyped } => { let messages = match protocol_compat.outgoing_notification(untyped) { @@ -125,20 +138,13 @@ pub(super) async fn outgoing_protocol_actor( .unbounded_send(Ok(json_rpc_message)) .map_err(crate::Error::into_internal_error)?; } - Ok(()) -} -fn complete_request_with_error( - response_tx: futures::channel::oneshot::Sender, - error: crate::Error, -) { - if response_tx - .send(crate::jsonrpc::ResponsePayload { - result: Err(error), - ack_tx: None, - }) - .is_err() - { - tracing::debug!("Dropped failed outgoing request because receiver was gone"); + // Closing the raw queue lets the transport actor finish all buffered + // writes. The caller separately awaits that transport future before + // treating the drain as complete. + drop(transport_tx); + for done in drain_waiters { + let _ = done.send(()); } + Ok(()) } diff --git a/src/agent-client-protocol/src/lib.rs b/src/agent-client-protocol/src/lib.rs index 7d9453a..29dde62 100644 --- a/src/agent-client-protocol/src/lib.rs +++ b/src/agent-client-protocol/src/lib.rs @@ -96,9 +96,9 @@ pub use capabilities::*; pub use jsonrpc::{ Builder, ByteStreams, Channel, ConnectionTo, Dispatch, HandleDispatchFrom, Handled, - IntoHandled, JsonRpcMessage, JsonRpcNotification, JsonRpcRequest, JsonRpcResponse, Lines, - NullHandler, RawJsonRpcMessage, RawJsonRpcParams, Responder, ResponseRouter, SentRequest, - UntypedMessage, + INCOMING_TRANSPORT_CLOSED_REASON, IntoHandled, JsonRpcMessage, JsonRpcNotification, + JsonRpcRequest, JsonRpcResponse, Lines, NullHandler, RawJsonRpcMessage, RawJsonRpcParams, + Responder, ResponseRouter, SentRequest, UntypedMessage, is_incoming_transport_closed, run::{ChainRun, NullRun, RunWithConnectionTo}, }; pub use jsonrpc::{RequestCancellation, is_cancel_request_notification}; diff --git a/src/agent-client-protocol/src/stdio.rs b/src/agent-client-protocol/src/stdio.rs index 07f21f9..d402ac2 100644 --- a/src/agent-client-protocol/src/stdio.rs +++ b/src/agent-client-protocol/src/stdio.rs @@ -54,7 +54,7 @@ impl ConnectTo for Stdio { if let Some(callback) = self.debug_callback { use futures::io::BufReader; - use futures::{AsyncBufReadExt, AsyncWriteExt, StreamExt}; + use futures::{AsyncBufReadExt, StreamExt}; let incoming_callback = callback.clone(); let incoming_lines = Box::pin(BufReader::new(stdin).lines().inspect(move |result| { @@ -68,9 +68,7 @@ impl ConnectTo for Stdio { (stdout, callback), async move |(mut writer, callback), line: String| { callback(&line, LineDirection::Stdout); - let mut bytes = line.into_bytes(); - bytes.push(b'\n'); - writer.write_all(&bytes).await?; + crate::jsonrpc::write_line(&mut writer, line).await?; Ok::<_, std::io::Error>((writer, callback)) }, )) diff --git a/src/agent-client-protocol/tests/jsonrpc_error_handling.rs b/src/agent-client-protocol/tests/jsonrpc_error_handling.rs index 3ae5c05..2168d52 100644 --- a/src/agent-client-protocol/tests/jsonrpc_error_handling.rs +++ b/src/agent-client-protocol/tests/jsonrpc_error_handling.rs @@ -204,7 +204,6 @@ async fn test_invalid_json() { // ============================================================================ #[tokio::test] -#[ignore = "hangs indefinitely - see https://github.com/agentclientprotocol/rust-sdk/issues/64"] async fn test_incomplete_line() { use futures::io::Cursor; diff --git a/src/agent-client-protocol/tests/jsonrpc_transport_close.rs b/src/agent-client-protocol/tests/jsonrpc_transport_close.rs new file mode 100644 index 0000000..8a01b52 --- /dev/null +++ b/src/agent-client-protocol/tests/jsonrpc_transport_close.rs @@ -0,0 +1,1155 @@ +//! Regression tests for incoming transport closure. + +use std::{ + future, io, + sync::{ + Arc, Mutex, + atomic::{AtomicBool, Ordering}, + }, + time::Duration, +}; + +use agent_client_protocol::{ + ByteStreams, Channel, ConnectTo, ConnectionTo, Dispatch, Error, Handled, JsonRpcMessage, + JsonRpcRequest, Lines, RawJsonRpcMessage, UntypedMessage, is_incoming_transport_closed, + role::{Role, UntypedRole}, + schema::v1::{RequestId, Response}, +}; +use agent_client_protocol_test::{MyRequest, MyResponse}; +use futures::{FutureExt as _, SinkExt as _, StreamExt as _, future::join, stream}; +use tokio::io::{AsyncBufReadExt as _, AsyncWriteExt as _}; +use tokio_util::compat::{TokioAsyncReadCompatExt as _, TokioAsyncWriteCompatExt as _}; + +const TIMEOUT: Duration = Duration::from_secs(2); + +struct DelegatingTransport(T); +struct ImmediateClient; + +struct PendingTransport { + started: Option>, + dropped: Arc, +} + +impl Drop for PendingTransport { + fn drop(&mut self) { + self.dropped.store(true, Ordering::Release); + } +} + +impl ConnectTo for PendingTransport { + async fn connect_to(mut self, client: impl ConnectTo) -> Result<(), Error> { + if let Some(started) = self.started.take() { + let _ = started.send(()); + } + + future::pending::<()>().await; + drop((self, client)); + Ok(()) + } +} + +struct QueuedClient { + started: futures::channel::oneshot::Sender<()>, + escaped: futures::channel::oneshot::Sender< + futures::channel::mpsc::UnboundedSender>, + >, +} + +impl ConnectTo for QueuedClient { + async fn connect_to(self, transport: impl ConnectTo) -> Result<(), Error> { + let (channel, transport_future) = transport.into_channel_and_future(); + let message = RawJsonRpcMessage::notification( + "queued".into(), + serde_json::json!({ "payload": "x".repeat(1024) }), + )?; + channel + .tx + .unbounded_send(Ok(message)) + .map_err(Error::into_internal_error)?; + drop(self.escaped.send(channel.tx.clone())); + let _ = self.started.send(()); + drop(channel); + transport_future.await + } +} + +#[derive(Clone)] +struct BlockingRequest { + entered: Arc>>>, + release: Arc>>, +} + +impl std::fmt::Debug for BlockingRequest { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.debug_struct("BlockingRequest").finish() + } +} + +impl JsonRpcMessage for BlockingRequest { + fn matches_method(method: &str) -> bool { + method == "blockingRequest" + } + + fn method(&self) -> &'static str { + "blockingRequest" + } + + fn to_untyped_message(&self) -> Result { + if let Some(entered) = self.entered.lock().unwrap().take() { + let _ = entered.send(()); + } + let _ = self.release.lock().unwrap().recv(); + UntypedMessage::new(self.method(), serde_json::json!({})) + } + + fn parse_message(_method: &str, _params: &impl serde::Serialize) -> Result { + Err(Error::method_not_found().data("test-only outgoing request")) + } +} + +impl JsonRpcRequest for BlockingRequest { + type Response = serde_json::Value; +} + +impl ConnectTo for DelegatingTransport +where + R: Role, + T: ConnectTo, +{ + async fn connect_to(self, client: impl ConnectTo) -> Result<(), Error> { + self.0.connect_to(client).await + } +} + +impl ConnectTo for ImmediateClient { + async fn connect_to(self, _client: impl ConnectTo) -> Result<(), Error> { + Ok(()) + } +} + +fn assert_connection_closed(error: &Error, method: &str) { + assert!(is_incoming_transport_closed(error)); + assert_eq!(error.message, "Incoming transport closed"); + let data = error.data.as_ref().expect("connection-close error data"); + assert_eq!(data["reason"], "incoming_transport_closed"); + assert_eq!(data["method"], method); +} + +async fn receive_requests_then_close(mut peer: Channel, count: usize) { + for _ in 0..count { + assert!(matches!( + peer.rx.next().await, + Some(Ok(RawJsonRpcMessage::Request(_))) + )); + } + drop(peer); +} + +async fn respond_then_close(mut peer: Channel) { + let Some(Ok(RawJsonRpcMessage::Request(request))) = peer.rx.next().await else { + panic!("expected outgoing request"); + }; + peer.tx + .send(Ok(RawJsonRpcMessage::response( + request.id, + Ok(serde_json::to_value(MyResponse { + status: "received".into(), + }) + .unwrap()), + ))) + .await + .expect("send response before EOF"); + drop(peer); +} + +fn assert_response_channel_lost(error: &Error) { + assert!( + !is_incoming_transport_closed(error), + "a response discarded locally must not be attributed to EOF" + ); + assert!( + error.data.as_ref().is_some_and(|data| data + .to_string() + .contains("response to `myRequest` never received")), + "unexpected response-channel error: {error:?}" + ); +} + +async fn assert_connect_to_flushes_final_response( + transport: impl ConnectTo, + mut peer_outgoing: tokio::io::DuplexStream, + peer_incoming: tokio::io::DuplexStream, +) { + let connection = UntypedRole + .builder() + .on_receive_request( + async |_request: MyRequest, responder, _cx: ConnectionTo| { + responder.respond(MyResponse { + status: "received".into(), + }) + }, + agent_client_protocol::on_receive_request!(), + ) + .connect_to(transport); + + let peer = async move { + let request = RawJsonRpcMessage::request( + "myRequest".into(), + serde_json::json!({}), + RequestId::Number(1), + ) + .unwrap(); + let mut request_bytes = serde_json::to_vec(&request).unwrap(); + request_bytes.push(b'\n'); + peer_outgoing.write_all(&request_bytes).await.unwrap(); + + // Half-close the peer's write direction while keeping its read + // direction open for the final response. + drop(peer_outgoing); + + let response_line = tokio::io::BufReader::new(peer_incoming) + .lines() + .next_line() + .await + .unwrap() + .expect("peer should receive a complete response before EOF"); + let RawJsonRpcMessage::Response(Response::Result { id, result }) = + serde_json::from_str(&response_line).expect("response must not be truncated") + else { + panic!("expected successful JSON-RPC response"); + }; + assert_eq!(id, RequestId::Number(1)); + assert_eq!( + serde_json::from_value::(result).unwrap().status, + "received" + ); + }; + + tokio::time::timeout(TIMEOUT, async move { + let (result, ()) = join(connection, peer).await; + result + }) + .await + .expect("connection hung while draining its final response") + .expect("clean EOF should succeed after flushing the response"); +} + +#[tokio::test] +async fn connect_to_returns_cleanly_on_incoming_eof_with_spawned_work() { + let (transport, peer) = Channel::duplex(); + drop(peer); + + let result = tokio::time::timeout( + TIMEOUT, + UntypedRole + .builder() + .with_spawned(|_cx| future::pending::>()) + .connect_to(transport), + ) + .await + .expect("connect_to hung after transport EOF"); + + assert!( + result.is_ok(), + "clean EOF should stop connect_to: {result:?}" + ); +} + +#[tokio::test] +async fn connect_to_flushes_response_queued_before_incoming_eof() { + // A one-byte SDK output buffer forces the JSON response write to yield + // repeatedly. Returning at EOF without a real transport drain truncates it. + let (sdk_outgoing, peer_incoming) = tokio::io::duplex(1); + let (peer_outgoing, sdk_incoming) = tokio::io::duplex(1024); + let transport = ByteStreams::new(sdk_outgoing.compat_write(), sdk_incoming.compat()); + + assert_connect_to_flushes_final_response(transport, peer_outgoing, peer_incoming).await; +} + +#[tokio::test] +async fn connect_to_flushes_a_buffered_byte_writer() { + let (sdk_outgoing, peer_incoming) = tokio::io::duplex(1); + let (peer_outgoing, sdk_incoming) = tokio::io::duplex(1024); + let buffered_outgoing = + futures::io::BufWriter::with_capacity(4096, sdk_outgoing.compat_write()); + let transport = ByteStreams::new(buffered_outgoing, sdk_incoming.compat()); + + assert_connect_to_flushes_final_response(transport, peer_outgoing, peer_incoming).await; +} + +#[tokio::test] +async fn wrapped_byte_streams_flush_response_queued_before_incoming_eof() { + let (sdk_outgoing, peer_incoming) = tokio::io::duplex(1); + let (peer_outgoing, sdk_incoming) = tokio::io::duplex(1024); + let transport = DelegatingTransport(ByteStreams::new( + sdk_outgoing.compat_write(), + sdk_incoming.compat(), + )); + + assert_connect_to_flushes_final_response(transport, peer_outgoing, peer_incoming).await; +} + +#[tokio::test] +async fn wrapped_byte_streams_client_success_does_not_wait_for_peer_eof() { + let (sdk_outgoing, peer_incoming) = tokio::io::duplex(1); + let (peer_outgoing, sdk_incoming) = tokio::io::duplex(1024); + let transport = DelegatingTransport(ByteStreams::new( + sdk_outgoing.compat_write(), + sdk_incoming.compat(), + )); + + tokio::time::timeout( + TIMEOUT, + ConnectTo::::connect_to(transport, ImmediateClient), + ) + .await + .expect("successful client shutdown waited for unrelated peer EOF") + .expect("successful client shutdown should drain outgoing and return"); + + // Keep both peer halves open until after the assertion so the result + // proves the transport's incoming actor was canceled, not completed. + drop(peer_outgoing); + drop(peer_incoming); +} + +#[tokio::test] +async fn outgoing_drain_keeps_the_full_duplex_read_half_moving() { + let (sdk_stream, peer_stream) = tokio::io::duplex(1); + let (sdk_incoming, sdk_outgoing) = tokio::io::split(sdk_stream); + let (peer_incoming, mut peer_outgoing) = tokio::io::split(peer_stream); + let transport = ByteStreams::new(sdk_outgoing.compat_write(), sdk_incoming.compat()); + let (started_tx, started_rx) = futures::channel::oneshot::channel(); + let (escaped_tx, escaped_rx) = futures::channel::oneshot::channel(); + + let connection = ConnectTo::::connect_to( + transport, + QueuedClient { + started: started_tx, + escaped: escaped_tx, + }, + ); + let peer = async move { + started_rx.await.expect("client should queue its message"); + let escaped = escaped_rx.await.expect("client should expose its sender"); + + // Write two frames before reading. The first frame proves shutdown no + // longer tries to forward into the client's closed channel; the second + // proves the read half keeps draining after that point. + for (method, payload_len) in [("first", 16), ("second", 1024)] { + let message = RawJsonRpcMessage::notification( + method.into(), + serde_json::json!({ "payload": "x".repeat(payload_len) }), + ) + .unwrap(); + let mut bytes = serde_json::to_vec(&message).unwrap(); + bytes.push(b'\n'); + peer_outgoing.write_all(&bytes).await.unwrap(); + } + + let line = tokio::io::BufReader::new(peer_incoming) + .lines() + .next_line() + .await + .unwrap() + .expect("queued output should be drained"); + assert!(matches!( + serde_json::from_str::(&line).unwrap(), + RawJsonRpcMessage::Notification(_) + )); + escaped + }; + + let escaped = tokio::time::timeout(TIMEOUT, async move { + let (result, escaped) = join(connection, peer).await; + result.map(|()| escaped) + }) + .await + .expect("full-duplex transport deadlocked while draining output") + .expect("queued output should drain successfully"); + + assert!( + escaped + .unbounded_send(Ok(RawJsonRpcMessage::notification( + "too-late".into(), + serde_json::json!({}), + ) + .unwrap())) + .is_err(), + "escaped sender accepted a message after client completion" + ); +} + +#[tokio::test] +async fn outgoing_drain_propagates_incoming_read_error() { + let (write_started_tx, write_started_rx) = futures::channel::oneshot::channel(); + let outgoing = futures::sink::unfold( + Some(write_started_tx), + |mut write_started, _line: String| async move { + if let Some(write_started) = write_started.take() { + let _ = write_started.send(()); + } + future::pending::>, io::Error>>() + .await + }, + ); + let (incoming_tx, incoming) = futures::channel::mpsc::unbounded::>(); + let (client_started_tx, _client_started_rx) = futures::channel::oneshot::channel(); + let (escaped_tx, escaped_rx) = futures::channel::oneshot::channel(); + + let connection = ConnectTo::::connect_to( + Lines::new(outgoing, incoming), + QueuedClient { + started: client_started_tx, + escaped: escaped_tx, + }, + ); + let inject_error = async move { + let escaped = escaped_rx + .await + .expect("client should expose its outgoing sender"); + write_started_rx + .await + .expect("queued output should reach the backpressured sink"); + assert!( + escaped.is_closed(), + "read error must be injected after the client enters drain mode" + ); + incoming_tx + .unbounded_send(Err(io::Error::other("read failed during outgoing drain"))) + .expect("incoming transport should still be polled during the drain"); + }; + + let error = tokio::time::timeout(TIMEOUT, async move { + let (result, ()) = join(connection, inject_error).await; + result + }) + .await + .expect("incoming read error was swallowed behind the blocked outgoing drain") + .expect_err("incoming read error should fail the connection"); + let detail = error + .data + .as_ref() + .map(serde_json::Value::to_string) + .unwrap_or_default(); + assert!( + detail.contains("read failed during outgoing drain"), + "unexpected transport error: {error:?}" + ); +} + +#[tokio::test] +async fn clean_outgoing_drain_does_not_hide_ready_incoming_read_error() { + let outgoing = futures::sink::unfold((), |(), _line: String| async { Ok::<_, io::Error>(()) }); + let incoming = stream::iter([Err(io::Error::other( + "ready read failed during outgoing drain", + ))]); + + let error = tokio::time::timeout( + TIMEOUT, + ConnectTo::::connect_to(Lines::new(outgoing, incoming), ImmediateClient), + ) + .await + .expect("connection should resolve when both drain sides are ready") + .expect_err("ready incoming read error must win over a clean outgoing drain"); + let detail = error + .data + .as_ref() + .map(serde_json::Value::to_string) + .unwrap_or_default(); + assert!( + detail.contains("ready read failed during outgoing drain"), + "unexpected transport error: {error:?}" + ); +} + +#[tokio::test] +async fn escaped_connection_handle_does_not_retain_the_transport() { + let dropped = Arc::new(AtomicBool::new(false)); + let (started_tx, started_rx) = futures::channel::oneshot::channel(); + let (escaped_tx, escaped_rx) = futures::channel::oneshot::channel(); + + let connection = UntypedRole.builder().connect_with( + PendingTransport { + started: Some(started_tx), + dropped: dropped.clone(), + }, + async move |cx| { + started_rx + .await + .map_err(|_| Error::internal_error().data("transport never started"))?; + let request = cx.send_request(MyRequest {}); + escaped_tx + .send((cx, request)) + .map_err(|_| Error::internal_error().data("escaped handle receiver dropped"))?; + Ok(()) + }, + ); + + let (result, escaped) = tokio::time::timeout(TIMEOUT, join(connection, escaped_rx)) + .await + .expect("connection did not stop after its foreground returned"); + result.expect("foreground shutdown should succeed"); + let (escaped, request) = escaped.expect("foreground should return connection state"); + + assert!( + dropped.load(Ordering::Acquire), + "an escaped handle kept the transport future alive" + ); + let error = tokio::time::timeout(TIMEOUT, request.block_task()) + .await + .expect("escaped handle retained the pending-reply registry") + .expect_err("request should fail when its connection driver stops"); + assert_response_channel_lost(&error); + drop(escaped); +} + +#[tokio::test] +async fn connect_with_can_await_incoming_close() { + let (transport, peer) = Channel::duplex(); + drop(peer); + + tokio::time::timeout( + TIMEOUT, + UntypedRole.builder().connect_with(transport, async |cx| { + cx.incoming_closed().await; + assert!(cx.is_incoming_closed()); + Ok(()) + }), + ) + .await + .expect("incoming close signal was not delivered") + .expect("observing a clean close should succeed"); +} + +#[tokio::test] +async fn on_close_error_stops_a_pending_connect_with_foreground() { + let (transport, peer) = Channel::duplex(); + drop(peer); + + let result = tokio::time::timeout( + TIMEOUT, + UntypedRole + .builder() + .on_close(async |_cx| Err(Error::internal_error().data("stop on transport close"))) + .connect_with(transport, async |_cx| { + future::pending::>().await + }), + ) + .await + .expect("on_close did not stop connect_with"); + + assert!(result.is_err(), "on_close error should stop the connection"); +} + +#[tokio::test] +async fn on_close_error_wins_over_foreground_close_waiter() { + let (transport, peer) = Channel::duplex(); + drop(peer); + + let result = tokio::time::timeout( + TIMEOUT, + UntypedRole + .builder() + .on_close(async |_cx| Err(Error::internal_error().data("close callback failed"))) + .connect_with(transport, async |cx| { + cx.incoming_closed().await; + Ok(()) + }), + ) + .await + .expect("connection hung after close callback error"); + + assert!( + result.is_err(), + "close callback error was swallowed by foreground close waiter" + ); +} + +#[tokio::test] +async fn successful_on_close_does_not_cancel_unrelated_foreground_work() { + let (transport, peer) = Channel::duplex(); + let (closed_tx, closed_rx) = futures::channel::oneshot::channel(); + let (release_tx, release_rx) = futures::channel::oneshot::channel(); + + let connection = UntypedRole + .builder() + .on_close(async move |_cx| { + closed_tx + .send(()) + .map_err(|()| Error::internal_error().data("close observer dropped")) + }) + .connect_with(transport, async move |_cx| { + release_rx + .await + .map_err(|_| Error::internal_error().data("foreground release dropped")) + }); + let connection = tokio::spawn(connection); + + drop(peer); + tokio::time::timeout(TIMEOUT, closed_rx) + .await + .expect("on_close did not run") + .expect("on_close sender dropped"); + tokio::task::yield_now().await; + assert!( + !connection.is_finished(), + "clean EOF should not cancel unrelated connect_with work" + ); + + release_tx + .send(()) + .expect("connection still receives release"); + connection + .await + .expect("connection task panicked") + .expect("foreground release should end the connection cleanly"); +} + +#[tokio::test] +async fn all_on_close_callbacks_run_in_registration_order_after_an_error() { + let (transport, peer) = Channel::duplex(); + let order = Arc::new(Mutex::new(Vec::new())); + let first_order = order.clone(); + let second_order = order.clone(); + + let connection = UntypedRole + .builder() + .on_close(async move |_cx| { + first_order.lock().unwrap().push(1); + Err(Error::internal_error().data("first close callback failed")) + }) + .on_close(async move |_cx| { + second_order.lock().unwrap().push(2); + Ok(()) + }) + .connect_with(transport, async |_cx| { + future::pending::>().await + }); + + drop(peer); + let result = tokio::time::timeout(TIMEOUT, connection) + .await + .expect("close callback error did not stop connection"); + assert!(result.is_err()); + assert_eq!(*order.lock().unwrap(), vec![1, 2]); +} + +#[tokio::test] +async fn pending_requests_fail_before_on_close_finishes() { + let (transport, peer) = Channel::duplex(); + let (release_callback_tx, release_callback_rx) = futures::channel::oneshot::channel(); + let (request_failed_tx, request_failed_rx) = futures::channel::oneshot::channel(); + + let connection = UntypedRole + .builder() + .on_close(async move |_cx| { + release_callback_rx + .await + .map_err(|_| Error::internal_error().data("close callback release dropped")) + }) + .connect_with(transport, async move |cx| { + let error = cx + .send_request(MyRequest {}) + .block_task() + .await + .expect_err("request should fail at EOF"); + assert_connection_closed(&error, "myRequest"); + request_failed_tx + .send(()) + .map_err(|()| Error::internal_error().data("request failure observer dropped"))?; + Ok(()) + }); + let release_callback = async move { + request_failed_rx + .await + .expect("request should fail while close callback is still blocked"); + release_callback_tx + .send(()) + .expect("close callback should still be running"); + }; + + let scenario = async move { + let (result, ((), ())) = join( + connection, + join(receive_requests_then_close(peer, 1), release_callback), + ) + .await; + result + }; + tokio::time::timeout(TIMEOUT, scenario) + .await + .expect("pending request was withheld by close callback") + .expect("connection foreground should handle the close error"); +} + +#[tokio::test] +async fn request_consumer_error_does_not_cancel_on_close_callbacks() { + let (transport, peer) = Channel::duplex(); + let (request_failed_tx, request_failed_rx) = futures::channel::oneshot::channel(); + let callback_order = Arc::new(Mutex::new(Vec::new())); + let first_callback_order = callback_order.clone(); + let second_callback_order = callback_order.clone(); + + let connection = UntypedRole + .builder() + .on_close(async move |_cx| { + // Wait until the EOF error has made the sibling task actor fail. + // Close processing must keep this callback alive regardless. + request_failed_rx + .await + .map_err(|_| Error::internal_error().data("request failure observer dropped"))?; + first_callback_order.lock().unwrap().push(1); + Err(Error::internal_error().data("close callback failed")) + }) + .on_close(async move |_cx| { + second_callback_order.lock().unwrap().push(2); + Ok(()) + }) + .connect_with(transport, async move |cx| { + cx.send_request(MyRequest {}) + .on_receiving_result(async move |result| { + let error = result.expect_err("request should fail at EOF"); + assert_connection_closed(&error, "myRequest"); + request_failed_tx + .send(()) + .map_err(|()| Error::internal_error().data("close callback was dropped"))?; + Err(error) + })?; + + cx.incoming_closed().await; + Ok(()) + }); + + let scenario = async move { + let (result, ()) = join(connection, receive_requests_then_close(peer, 1)).await; + result + }; + let error = tokio::time::timeout(TIMEOUT, scenario) + .await + .expect("connection hung after request consumer failed") + .expect_err("request consumer error should stop the connection"); + assert_eq!( + error.data, + Some(serde_json::json!("close callback failed")), + "the configured close callback error should win over the sibling task error" + ); + assert_eq!(*callback_order.lock().unwrap(), vec![1, 2]); +} + +#[tokio::test] +async fn queued_request_is_failed_as_eof_before_outgoing_actor_runs() { + let (transport, peer) = Channel::duplex(); + let (request_tx, request_rx) = + futures::channel::oneshot::channel::>(); + + let connection = UntypedRole + .builder() + .on_close(async |_cx| Err(Error::internal_error().data("stop after EOF"))) + .connect_with(transport, async move |cx| { + request_tx + .send(cx.send_request(MyRequest {})) + .map_err(|_| Error::internal_error().data("request observer dropped"))?; + future::pending::>().await + }); + let close_after_request_is_queued = async move { + let request = request_rx.await.expect("main should create its request"); + // The outgoing actor has not been repolled since main_fn queued the + // request. Make incoming EOF and its failing callback win that poll. + drop(peer); + request + }; + + let (connection_result, request) = join(connection, close_after_request_is_queued).await; + connection_result.expect_err("close callback should stop the connection"); + let error = request + .block_task() + .await + .expect_err("queued request should be failed by incoming EOF"); + assert_connection_closed(&error, "myRequest"); +} + +#[tokio::test] +async fn on_close_observes_queued_requests_already_failed() { + let (transport, peer) = Channel::duplex(); + let (request_tx, request_rx) = + futures::channel::oneshot::channel::>(); + let (request_queued_tx, request_queued_rx) = futures::channel::oneshot::channel(); + let (callback_done_tx, callback_done_rx) = futures::channel::oneshot::channel(); + + let connection = UntypedRole + .builder() + .on_close(async move |_cx| { + let request = request_rx + .await + .map_err(|_| Error::internal_error().data("queued request sender dropped"))?; + let error = request + .block_task() + .now_or_never() + .expect("queued request was still pending when close callback began") + .expect_err("queued request should fail at EOF"); + assert_connection_closed(&error, "myRequest"); + callback_done_tx + .send(()) + .map_err(|()| Error::internal_error().data("callback observer dropped")) + }) + .connect_with(transport, async move |cx| { + request_tx + .send(cx.send_request(MyRequest {})) + .map_err(|_| Error::internal_error().data("close callback receiver dropped"))?; + request_queued_tx + .send(()) + .map_err(|()| Error::internal_error().data("request queue observer dropped"))?; + callback_done_rx + .await + .map_err(|_| Error::internal_error().data("close callback did not finish"))?; + Ok(()) + }); + let close_after_request_is_queued = async move { + request_queued_rx + .await + .expect("foreground should queue its request"); + drop(peer); + }; + + let scenario = async move { + let (result, ()) = join(connection, close_after_request_is_queued).await; + result + }; + tokio::time::timeout(TIMEOUT, scenario) + .await + .expect("connection hung while closing a queued request") + .expect("close callback and foreground should finish successfully"); +} + +#[tokio::test] +async fn request_finishing_conversion_after_eof_keeps_the_eof_cause() { + let (transport, mut peer) = Channel::duplex(); + let (connection_tx, connection_rx) = + futures::channel::oneshot::channel::>(); + let connection_tx = Arc::new(Mutex::new(Some(connection_tx))); + let connection_observer = connection_tx.clone(); + + let connection = UntypedRole + .builder() + .on_receive_request( + async move |_request: MyRequest, responder, cx: ConnectionTo| { + connection_observer + .lock() + .unwrap() + .take() + .expect("connection observer should run once") + .send(cx) + .map_err(|_| { + Error::internal_error().data("connection handle receiver dropped") + })?; + responder.respond(MyResponse { + status: "ready".into(), + }) + }, + agent_client_protocol::on_receive_request!(), + ) + .connect_to(transport); + let connection = tokio::spawn(connection); + + peer.tx + .send(Ok(RawJsonRpcMessage::request( + "myRequest".into(), + serde_json::json!({}), + RequestId::Number(1), + ) + .unwrap())) + .await + .expect("send request that exposes the connection handle"); + let cx = tokio::time::timeout(TIMEOUT, connection_rx) + .await + .expect("handler did not expose its connection handle") + .expect("connection handle sender dropped"); + assert!(matches!( + tokio::time::timeout(TIMEOUT, peer.rx.next()) + .await + .expect("handler response was not sent"), + Some(Ok(RawJsonRpcMessage::Response(_))) + )); + + let (entered_tx, entered_rx) = futures::channel::oneshot::channel(); + let (release_tx, release_rx) = std::sync::mpsc::channel(); + let request = BlockingRequest { + entered: Arc::new(Mutex::new(Some(entered_tx))), + release: Arc::new(Mutex::new(release_rx)), + }; + let request_thread = std::thread::spawn(move || cx.send_request(request)); + + tokio::time::timeout(TIMEOUT, entered_rx) + .await + .expect("request conversion did not start") + .expect("request conversion entry sender dropped"); + + // EOF starts the reactive drain and closes the outgoing receiver while + // request conversion is still blocked before enqueueing. + drop(peer); + tokio::time::timeout(TIMEOUT, connection) + .await + .expect("reactive connection did not finish its drain") + .expect("connection task panicked") + .expect("clean EOF should finish successfully"); + + release_tx + .send(()) + .expect("release blocked request conversion"); + let request = request_thread.join().expect("request thread panicked"); + let error = request + .block_task() + .await + .expect_err("request should retain the EOF failure cause"); + assert_connection_closed(&error, "blockingRequest"); +} + +#[tokio::test] +async fn incoming_eof_fails_every_pending_request() { + let (transport, peer) = Channel::duplex(); + + let connection = UntypedRole.builder().connect_with(transport, async |cx| { + let first = cx.send_request(MyRequest {}).block_task(); + let second = cx.send_request(MyRequest {}).block_task(); + let (first, second) = join(first, second).await; + + assert_connection_closed(&first.expect_err("first request should fail"), "myRequest"); + assert_connection_closed( + &second.expect_err("second request should fail"), + "myRequest", + ); + Ok(()) + }); + + let scenario = async move { + let (result, ()) = join(connection, receive_requests_then_close(peer, 2)).await; + result + }; + tokio::time::timeout(TIMEOUT, scenario) + .await + .expect("pending requests hung after transport EOF") + .expect("connection foreground should handle both close errors"); +} + +#[tokio::test] +async fn byte_stream_eof_fails_an_in_flight_request() { + let (client_outgoing, peer_incoming) = tokio::io::duplex(1024); + let (peer_outgoing, client_incoming) = tokio::io::duplex(1024); + let transport = ByteStreams::new(client_outgoing.compat_write(), client_incoming.compat()); + + let connection = UntypedRole.builder().connect_with(transport, async |cx| { + let error = cx + .send_request(MyRequest {}) + .block_task() + .await + .expect_err("request should fail when byte stream reaches EOF"); + assert_connection_closed(&error, "myRequest"); + Ok(()) + }); + let peer = async move { + let mut lines = tokio::io::BufReader::new(peer_incoming).lines(); + assert!( + lines.next_line().await.unwrap().is_some(), + "peer should receive the request before closing" + ); + drop(peer_outgoing); + }; + + tokio::time::timeout(TIMEOUT, async move { + let (result, ()) = join(connection, peer).await; + result + }) + .await + .expect("request hung after byte stream EOF") + .expect("connection foreground should handle the close error"); +} + +#[tokio::test] +async fn incoming_eof_delivers_error_to_callback_style_request_consumer() { + let (transport, peer) = Channel::duplex(); + + let connection = UntypedRole.builder().connect_with(transport, async |cx| { + let (callback_tx, callback_rx) = futures::channel::oneshot::channel(); + cx.send_request(MyRequest {}) + .on_receiving_result(async move |result| { + callback_tx + .send(result) + .map_err(|_| Error::internal_error().data("request callback receiver dropped")) + })?; + + let error = callback_rx + .await + .map_err(|_| Error::internal_error().data("request callback did not run"))? + .expect_err("callback should receive the incoming-close error"); + assert_connection_closed(&error, "myRequest"); + Ok(()) + }); + + let scenario = async move { + let (result, ()) = join(connection, receive_requests_then_close(peer, 1)).await; + result + }; + tokio::time::timeout(TIMEOUT, scenario) + .await + .expect("callback-style request consumer hung") + .expect("connection foreground should handle callback error"); +} + +#[tokio::test] +async fn response_buffered_before_eof_is_delivered() { + let (transport, mut peer) = Channel::duplex(); + + let connection = UntypedRole.builder().connect_with(transport, async |cx| { + let response = cx.send_request(MyRequest {}).block_task().await?; + assert_eq!(response.status, "received"); + Ok(()) + }); + let respond_then_close = async move { + let Some(Ok(RawJsonRpcMessage::Request(request))) = peer.rx.next().await else { + panic!("expected outgoing request"); + }; + peer.tx + .send(Ok(RawJsonRpcMessage::response( + request.id, + Ok(serde_json::to_value(MyResponse { + status: "received".into(), + }) + .unwrap()), + ))) + .await + .expect("send response before closing transport"); + drop(peer); + }; + + tokio::time::timeout(TIMEOUT, async move { + let (result, ()) = join(connection, respond_then_close).await; + result + }) + .await + .expect("buffered response was lost at EOF") + .expect("buffered response should win over later EOF"); +} + +#[tokio::test] +async fn response_router_drop_before_eof_is_not_reported_as_eof() { + let (transport, peer) = Channel::duplex(); + + let connection = UntypedRole + .builder() + .on_receive_dispatch( + async |dispatch: Dispatch, _cx: ConnectionTo| match dispatch { + // Claim and drop the router without delivering the response to + // its SentRequest consumer. + Dispatch::Response(..) => Ok(Handled::Yes), + message => Ok(Handled::No { + message, + retry: false, + }), + }, + agent_client_protocol::on_receive_dispatch!(), + ) + .connect_with(transport, async |cx| { + let request = cx.send_request(MyRequest {}); + cx.incoming_closed().await; + + let error = request + .block_task() + .await + .expect_err("dropped response router should cancel its local channel"); + assert_response_channel_lost(&error); + Ok(()) + }); + + let scenario = async move { + let (result, ()) = join(connection, respond_then_close(peer)).await; + result + }; + tokio::time::timeout(TIMEOUT, scenario) + .await + .expect("connection hung after the response router was dropped") + .expect("foreground should handle the local response-channel loss"); +} + +#[tokio::test] +async fn response_router_drop_after_eof_does_not_invoke_result_callback() { + let (transport, peer) = Channel::duplex(); + let callback_ran = Arc::new(AtomicBool::new(false)); + let callback_observer = callback_ran.clone(); + + let connection = UntypedRole + .builder() + .on_receive_dispatch( + async |dispatch: Dispatch, _cx: ConnectionTo| match dispatch { + Dispatch::Response(..) => Ok(Handled::Yes), + message => Ok(Handled::No { + message, + retry: false, + }), + }, + agent_client_protocol::on_receive_dispatch!(), + ) + .connect_with(transport, async move |cx| { + let request = cx.send_request(MyRequest {}); + cx.incoming_closed().await; + + request.on_receiving_result(async move |_result| { + callback_observer.store(true, Ordering::Release); + Ok(()) + })?; + future::pending::>().await + }); + + let scenario = async move { + let (result, ()) = join(connection, respond_then_close(peer)).await; + result + }; + let error = tokio::time::timeout(TIMEOUT, scenario) + .await + .expect("response consumer should fail instead of hanging") + .expect_err("local response-channel loss should stop its consuming task"); + assert_response_channel_lost(&error); + assert!( + !callback_ran.load(Ordering::Acquire), + "on_receiving_result must not treat local channel loss as a peer EOF error" + ); +} + +#[tokio::test] +async fn request_created_after_incoming_eof_fails_immediately() { + let (transport, peer) = Channel::duplex(); + drop(peer); + + tokio::time::timeout( + TIMEOUT, + UntypedRole.builder().connect_with(transport, async |cx| { + cx.incoming_closed().await; + let error = cx + .send_request(MyRequest {}) + .block_task() + .await + .expect_err("post-close request should fail"); + assert_connection_closed(&error, "myRequest"); + Ok(()) + }), + ) + .await + .expect("post-close request hung") + .expect("connection foreground should handle the close error"); +} + +#[tokio::test] +async fn incoming_read_error_is_not_reported_as_clean_eof() { + let outgoing = futures::sink::unfold((), |(), _line: String| async { Ok::<_, io::Error>(()) }); + let incoming = stream::iter([Err(io::Error::other("read failed"))]); + + let result = tokio::time::timeout( + TIMEOUT, + UntypedRole + .builder() + .connect_to(Lines::new(outgoing, incoming)), + ) + .await + .expect("connection hung after transport read error"); + + assert!(result.is_err(), "transport read error was swallowed as EOF"); +}