fix(goose-acp): heap allocations (#7322)

crates/goose-acp/src/server.rs

@@ -1193,18 +1193,23 @@ impl JrMessageHandler for GooseAcpHandler {
         "goose-acp"
     }
 
-    async fn handle_message(
+    fn handle_message(
         &mut self,
         message: MessageCx,
         cx: JrConnectionCx<AgentToClient>,
-    ) -> Result<Handled<MessageCx>, sacp::Error> {
+    ) -> impl std::future::Future<Output = Result<Handled<MessageCx>, sacp::Error>> + Send {
         use sacp::util::MatchMessageFrom;
         use sacp::JrRequestCx;
 
+        let agent = self.agent.clone();
+
+        // The MatchMessageFrom chain produces an ~85KB async state machine.
+        // Box::pin moves it to the heap so it doesn't overflow the tokio worker stack.
+        Box::pin(async move {
             MatchMessageFrom::new(message, &cx)
                 .if_request(
                     |req: InitializeRequest, req_cx: JrRequestCx<InitializeResponse>| async {
-                    req_cx.respond(self.agent.on_initialize(req).await?)
+                        req_cx.respond(agent.on_initialize(req).await?)
                     },
                 )
                 .await
@@ -1216,19 +1221,19 @@ impl JrMessageHandler for GooseAcpHandler {
                 .await
                 .if_request(
                     |req: NewSessionRequest, req_cx: JrRequestCx<NewSessionResponse>| async {
-                    req_cx.respond(self.agent.on_new_session(req).await?)
+                        req_cx.respond(agent.on_new_session(req).await?)
                     },
                 )
                 .await
                 .if_request(
                     |req: LoadSessionRequest, req_cx: JrRequestCx<LoadSessionResponse>| async {
-                    req_cx.respond(self.agent.on_load_session(req, &cx).await?)
+                        req_cx.respond(agent.on_load_session(req, &cx).await?)
                     },
                 )
                 .await
                 .if_request(
                     |req: PromptRequest, req_cx: JrRequestCx<PromptResponse>| async {
-                    let agent = self.agent.clone();
+                        let agent = agent.clone();
                         let cx_clone = cx.clone();
                         cx.spawn(async move {
                             match agent.on_prompt(req, &cx_clone).await {
@@ -1245,27 +1250,28 @@ impl JrMessageHandler for GooseAcpHandler {
                     },
                 )
                 .await
-            .if_notification(|notif: CancelNotification| async {
+                .if_notification(|notif: CancelNotification| async { agent.on_cancel(notif).await })
-                self.agent.on_cancel(notif).await
-            })
                 .await
                 // Handle methods not yet in the sacp typed API.
                 // - session/set_model: typed support pending in sacp
                 // - _<method>: custom requests that will eventually route to goose-server
                 .otherwise({
-                let agent = self.agent.clone();
+                    let agent = agent.clone();
                     |message: MessageCx| async move {
                         match message {
                             MessageCx::Request(req, request_cx)
                                 if req.method == "session/set_model" =>
                             {
-                            let params: SetSessionModelRequest = serde_json::from_value(req.params)
+                                let params: SetSessionModelRequest =
-                                .map_err(|e| sacp::Error::invalid_params().data(e.to_string()))?;
+                                    serde_json::from_value(req.params).map_err(|e| {
+                                        sacp::Error::invalid_params().data(e.to_string())
+                                    })?;
                                 let resp = agent
                                     .on_set_model(&params.session_id.0, &params.model_id.0)
                                     .await?;
-                            let json = serde_json::to_value(resp)
+                                let json = serde_json::to_value(resp).map_err(|e| {
-                                .map_err(|e| sacp::Error::internal_error().data(e.to_string()))?;
+                                    sacp::Error::internal_error().data(e.to_string())
+                                })?;
                                 request_cx.respond(json)?;
                                 Ok(())
                             }
@@ -1282,14 +1288,20 @@ impl JrMessageHandler for GooseAcpHandler {
                 })
                 .await
                 .map(|()| Handled::Yes)
+        })
     }
 }
 
-pub async fn serve<R, W>(agent: Arc<GooseAcpAgent>, read: R, write: W) -> Result<()>
+pub fn serve<R, W>(
+    agent: Arc<GooseAcpAgent>,
+    read: R,
+    write: W,
+) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<()>> + Send>>
 where
     R: futures::AsyncRead + Unpin + Send + 'static,
     W: futures::AsyncWrite + Unpin + Send + 'static,
 {
+    Box::pin(async move {
         let handler = GooseAcpHandler { agent };
 
         AgentToClient::builder()
@@ -1299,6 +1311,7 @@ where
             .await?;
 
         Ok(())
+    })
 }
 
 pub async fn run(builtins: Vec<String>) -> Result<()> {

crates/goose-acp/src/transport/http.rs

@@ -41,13 +41,11 @@ impl HttpState {
 
         let acp_session_id = uuid::Uuid::new_v4().to_string();
 
-        let handle = tokio::spawn(async move {
         let read_stream = ReceiverToAsyncRead::new(to_agent_rx);
         let write_stream = SenderToAsyncWrite::new(from_agent_tx);
-
+        let fut = crate::server::serve(agent, read_stream.compat(), write_stream.compat_write());
-            if let Err(e) =
+        let handle = tokio::spawn(async move {
-                crate::server::serve(agent, read_stream.compat(), write_stream.compat_write()).await
+            if let Err(e) = fut.await {
-            {
                 error!("ACP session error: {}", e);
             }
         });

crates/goose-acp/src/transport/websocket.rs

@@ -36,13 +36,11 @@ impl WsState {
 
         let acp_session_id = uuid::Uuid::new_v4().to_string();
 
-        let handle = tokio::spawn(async move {
         let read_stream = ReceiverToAsyncRead::new(to_agent_rx);
         let write_stream = SenderToAsyncWrite::new(from_agent_tx);
-
+        let fut = crate::server::serve(agent, read_stream.compat(), write_stream.compat_write());
-            if let Err(e) =
+        let handle = tokio::spawn(async move {
-                crate::server::serve(agent, read_stream.compat(), write_stream.compat_write()).await
+            if let Err(e) = fut.await {
-            {
                 error!("ACP WebSocket session error: {}", e);
             }
         });