From 5b9454a21cf6ddc4b3edee994f115dd604adfd70 Mon Sep 17 00:00:00 2001 From: Felix S Klock II Date: Wed, 29 Jul 2026 22:33:10 +0100 Subject: [PATCH] fix(acp): box handler chain links to bound dispatch stack use Consumers that register many typed handlers on one connection (a full ACP client registers 10-20) build a deep tower of nested ChainedHandler futures. In unoptimized builds each nesting level inlined a large poll frame, so dispatching a message on a small fixed stack -- such as a macOS Grand Central Dispatch worker's 512 KiB -- overflowed and aborted the process, and very deep chains failed to compile with a layout query depth overflow. Box each chained handler's future so its state lives on the heap, constructing it in an inline(never) helper so the temporary never stages in the caller's poll frame. Per-handler stack cost drops by roughly an order of magnitude and no longer inlines with depth. Adds a regression test that dispatches through a 33-handler chain on a 512 KiB thread. See zed-industries/zed#61840. --- .../src/jsonrpc/handlers.rs | 29 ++- .../tests/jsonrpc_deep_chain_stack.rs | 206 ++++++++++++++++++ 2 files changed, 225 insertions(+), 10 deletions(-) create mode 100644 src/agent-client-protocol/tests/jsonrpc_deep_chain_stack.rs 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:?}"); +}