Pass spawn in

Co-authored-by: Max Brunsfeld <maxbrunsfeld@gmail.com>
Co-authored-by: Conrad Irwin <conrad.irwin@gmail.com>
Co-authored-by: Nathan Sobo <nathan@zed.dev>

Author: 4 people

rust/acp.rs

@@ -40,6 +40,7 @@ impl AgentConnection {
         handler: H,
         outgoing_bytes: impl Unpin + AsyncWrite,
         incoming_bytes: impl Unpin + AsyncRead,
+        spawn: impl Fn(LocalBoxFuture<'static, ()>),
     ) -> (
         Self,
         impl Future<Output = ()>,
@@ -53,6 +54,7 @@ impl AgentConnection {
             }),
             outgoing_bytes,
             incoming_bytes,
+            spawn,
         );
         (Self(connection), handler_task, io_task)
     }
@@ -81,6 +83,7 @@ impl ClientConnection {
         handler: H,
         outgoing_bytes: impl Unpin + AsyncWrite,
         incoming_bytes: impl Unpin + AsyncRead,
+        spawn: impl Fn(LocalBoxFuture<'static, ()>),
     ) -> (
         Self,
         impl Future<Output = ()>,
@@ -94,6 +97,7 @@ impl ClientConnection {
             }),
             outgoing_bytes,
             incoming_bytes,
+            spawn,
         );
         (Self(connection), handler_task, io_task)
     }
@@ -189,6 +193,7 @@ where
         request_handler: Box<dyn 'static + Fn(In) -> LocalBoxFuture<'static, Result<In::Response>>>,
         outgoing_bytes: impl Unpin + AsyncWrite,
         incoming_bytes: impl Unpin + AsyncRead,
+        spawn: impl Fn(LocalBoxFuture<'static, ()>),
     ) -> (
         Self,
         impl Future<Output = ()>,
@@ -201,7 +206,7 @@ where
             outgoing_tx: outgoing_tx.clone(),
             next_id: AtomicI32::new(0),
         };
-        let handler_task = Self::handle_incoming(outgoing_tx, incoming_rx, request_handler);
+        let handler_task = Self::handle_incoming(outgoing_tx, incoming_rx, request_handler, spawn);
         let io_task = Self::handle_io(
             outgoing_rx,
             incoming_tx,
@@ -309,27 +314,35 @@ where
         incoming_handler: Box<
             dyn 'static + Fn(In) -> LocalBoxFuture<'static, Result<In::Response>>,
         >,
+        spawn: impl Fn(LocalBoxFuture<'static, ()>),
     ) {
         while let Some((id, params)) = incoming_rx.next().await {
-            let result = incoming_handler(params).await;
-            match result {
-                Ok(result) => {
-                    outgoing_tx
-                        .unbounded_send(OutgoingMessage::OkResponse { id, result })
-                        .ok();
-                }
-                Err(error) => {
-                    outgoing_tx
-                        .unbounded_send(OutgoingMessage::ErrorResponse {
-                            id,
-                            error: Error {
-                                code: -32603,
-                                message: error.to_string(),
-                            },
-                        })
-                        .ok();
+            let result = incoming_handler(params);
+            let outgoing_tx = outgoing_tx.clone();
+            spawn(
+                async move {
+                    let result = result.await;
+                    match result {
+                        Ok(result) => {
+                            outgoing_tx
+                                .unbounded_send(OutgoingMessage::OkResponse { id, result })
+                                .ok();
+                        }
+                        Err(error) => {
+                            outgoing_tx
+                                .unbounded_send(OutgoingMessage::ErrorResponse {
+                                    id,
+                                    error: Error {
+                                        code: -32603,
+                                        message: error.to_string(),
+                                    },
+                                })
+                                .ok();
+                        }
+                    }
                 }
-            }
+                .boxed_local(),
+            )
         }
     }
 }

rust/acp_tests.rs

@@ -67,9 +67,23 @@ async fn test_client_agent_communication() {
             let (agent_to_client_tx, agent_to_client_rx) = async_pipe::pipe();
 
             let (client_connection, client_handle_task, client_io_task) =
-                AgentConnection::connect_to_agent(client, client_to_agent_tx, agent_to_client_rx);
+                AgentConnection::connect_to_agent(
+                    client,
+                    client_to_agent_tx,
+                    agent_to_client_rx,
+                    |fut| {
+                        tokio::task::spawn_local(fut);
+                    },
+                );
             let (agent_connection, agent_handle_task, agent_io_task) =
-                ClientConnection::connect_to_client(agent, agent_to_client_tx, client_to_agent_rx);
+                ClientConnection::connect_to_client(
+                    agent,
+                    agent_to_client_tx,
+                    client_to_agent_rx,
+                    |fut| {
+                        tokio::task::spawn_local(fut);
+                    },
+                );
 
             let _task = tokio::task::spawn_local(client_handle_task);
             let _task = tokio::task::spawn_local(agent_handle_task);

rust/schema.rs

@@ -16,8 +16,8 @@ pub struct Method {
     pub response_payload: bool,
 }
 
-pub trait AnyRequest: Serialize + Sized {
-    type Response: Serialize;
+pub trait AnyRequest: Serialize + Sized + 'static {
+    type Response: Serialize + 'static;
     fn from_method_and_params(method: &str, params: &RawValue) -> Result<Self>;
     fn response_from_method_and_result(method: &str, params: &RawValue) -> Result<Self::Response>;
 }