diff --git a/src/agent-client-protocol/src/jsonrpc/handlers.rs b/src/agent-client-protocol/src/jsonrpc/handlers.rs index efc904b..2c7256a 100644 --- a/src/agent-client-protocol/src/jsonrpc/handlers.rs +++ b/src/agent-client-protocol/src/jsonrpc/handlers.rs @@ -499,6 +499,21 @@ impl ChainedHandler { } } +/// Heap-allocates a handler's dispatch future; `inline(never)` keeps its +/// construction out of the caller's poll frame. +#[inline(never)] +fn boxed_dispatch( + handler: &mut H, + message: Dispatch, + connection: ConnectionTo, +) -> crate::BoxFuture<'_, Result, crate::Error>> +where + Counterpart: Role, + H: HandleDispatchFrom, +{ + Box::pin(handler.handle_dispatch_from(message, connection)) +} + impl HandleDispatchFrom for ChainedHandler where H1: HandleDispatchFrom, @@ -517,20 +532,14 @@ where message: Dispatch, connection: ConnectionTo, ) -> Result, crate::Error> { - match self - .handler1 - .handle_dispatch_from(message, connection.clone()) - .await? - { + // Box each link so deep chains don't build one giant nested poll + // frame per handler and overflow small (e.g. GCD 512 KiB) stacks. + match boxed_dispatch(&mut self.handler1, message, connection.clone()).await? { Handled::Yes => Ok(Handled::Yes), Handled::No { message, retry: retry1, - } => match self - .handler2 - .handle_dispatch_from(message, connection) - .await? - { + } => match boxed_dispatch(&mut self.handler2, message, connection).await? { Handled::Yes => Ok(Handled::Yes), Handled::No { message, diff --git a/src/agent-client-protocol/tests/jsonrpc_deep_chain_stack.rs b/src/agent-client-protocol/tests/jsonrpc_deep_chain_stack.rs new file mode 100644 index 0000000..6cc5616 --- /dev/null +++ b/src/agent-client-protocol/tests/jsonrpc_deep_chain_stack.rs @@ -0,0 +1,206 @@ +//! Regression test for stack overflow with deeply chained handlers. +//! +//! Consumers like Zed register 10-20 typed handlers on a single connection, +//! producing a tower of nested `ChainedHandler` futures. In unoptimized +//! (debug) builds, polling that tower used to consume a large stack frame per +//! chain link, overflowing threads with small fixed stacks — notably macOS +//! GCD dispatch threads, which have 512 KiB and ignore `RUST_MIN_STACK`. +//! See zed-industries/zed#61840. +//! +//! This test drives a realistically deep handler chain on a thread restricted +//! to 512 KiB of stack, in the spirit of that environment. It aborts the +//! process (stack overflow) if dispatch frames regress. + +use agent_client_protocol::role::UntypedRole; +use agent_client_protocol::{ + ConnectionTo, JsonRpcMessage, JsonRpcNotification, JsonRpcRequest, JsonRpcResponse, Responder, + SentRequest, +}; +use serde::{Deserialize, Serialize}; +use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; + +/// The stack size of a macOS GCD (libdispatch) worker thread, the smallest +/// stack the dispatch loop is known to run on in the wild. +const GCD_STACK_SIZE: usize = 512 * 1024; + +/// Test helper to block and wait for a JSON-RPC response. +async fn recv( + response: SentRequest, +) -> Result { + let (tx, rx) = tokio::sync::oneshot::channel(); + response.on_receiving_result(async move |result| { + tx.send(result) + .map_err(|_| agent_client_protocol::Error::internal_error()) + })?; + rx.await + .map_err(|_| agent_client_protocol::Error::internal_error())? +} + +macro_rules! test_message { + ($name:ident, $method:literal) => { + #[derive(Debug, Clone, Serialize, Deserialize)] + struct $name { + message: String, + } + + impl JsonRpcMessage for $name { + fn matches_method(method: &str) -> bool { + method == $method + } + + fn method(&self) -> &'static str { + $method + } + + fn to_untyped_message( + &self, + ) -> Result { + agent_client_protocol::UntypedMessage::new(self.method(), self) + } + + fn parse_message( + method: &str, + params: &impl serde::Serialize, + ) -> Result { + if !Self::matches_method(method) { + return Err(agent_client_protocol::Error::method_not_found()); + } + agent_client_protocol::util::json_cast(params) + } + } + }; +} + +test_message!(PingRequest, "ping"); +test_message!(NopRequest, "nop"); +test_message!(NopNotification, "nop_notification"); + +impl JsonRpcRequest for PingRequest { + type Response = PongResponse; +} + +impl JsonRpcRequest for NopRequest { + type Response = PongResponse; +} + +impl JsonRpcNotification for NopNotification {} + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct PongResponse { + echo: String, +} + +impl JsonRpcResponse for PongResponse { + fn into_json(self, _method: &str) -> Result { + serde_json::to_value(self).map_err(agent_client_protocol::Error::into_internal_error) + } + + fn from_value( + _method: &str, + value: serde_json::Value, + ) -> Result { + agent_client_protocol::util::json_cast(&value) + } +} + +/// Registers `PingRequest` first (innermost link, so a "ping" traverses the +/// entire chain before it is claimed), then a pile of handlers for messages +/// that are never sent — mirroring a consumer that registers a callback per +/// protocol method. +macro_rules! deep_chain_server { + ($($_idx:literal),*) => { + UntypedRole + .builder() + .on_receive_request( + async |request: PingRequest, + responder: Responder, + _connection: ConnectionTo| { + responder.respond(PongResponse { + echo: format!("pong: {}", request.message), + }) + }, + agent_client_protocol::on_receive_request!(), + ) + $( + .on_receive_request( + async |request: NopRequest, + responder: Responder, + _connection: ConnectionTo| { + let _ = $_idx; + responder.respond(PongResponse { + echo: request.message, + }) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_notification( + async |_notification: NopNotification, _cx: ConnectionTo| { + Ok(()) + }, + agent_client_protocol::on_receive_notification!(), + ) + )* + }; +} + +#[test] +fn deep_handler_chain_within_small_stack() { + let (client_io, server_io) = tokio::io::duplex(1024); + + // The server side runs its whole dispatch loop on a thread with the same + // fixed stack budget as a macOS GCD worker. + let server = std::thread::Builder::new() + .name("small-stack-dispatch".to_string()) + .stack_size(GCD_STACK_SIZE) + .spawn(move || { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("failed to build server runtime"); + let local = tokio::task::LocalSet::new(); + local.block_on(&runtime, async move { + let (server_reader, server_writer) = tokio::io::split(server_io); + let transport = agent_client_protocol::ByteStreams::new( + server_writer.compat_write(), + server_reader.compat(), + ); + // 16 request + 16 notification handlers plus the ping handler: + // deeper than a full ACP client/agent registration. Before + // the chain links were boxed, this many handlers did not + // even compile (layout query depth overflow), and half as + // many overflowed the stack below at runtime. + let server = + deep_chain_server!(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16); + server.connect_to(transport).await + }) + }) + .expect("failed to spawn small-stack thread"); + + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("failed to build client runtime"); + let local = tokio::task::LocalSet::new(); + let result = local.block_on(&runtime, async move { + let (client_reader, client_writer) = tokio::io::split(client_io); + let transport = agent_client_protocol::ByteStreams::new( + client_writer.compat_write(), + client_reader.compat(), + ); + UntypedRole + .builder() + .connect_with(transport, async |cx| { + let response = recv(cx.send_request(PingRequest { + message: "hello".to_string(), + })) + .await?; + assert_eq!(response.echo, "pong: hello"); + Ok(()) + }) + .await + }); + assert!(result.is_ok(), "client failed: {result:?}"); + + let server_result = server.join().expect("server thread panicked"); + assert!(server_result.is_ok(), "server failed: {server_result:?}"); +}