WIP fix cancellation

Author: nikomatsakis

crates/chat-cli/src/cli/acp/client_connection.rs

@@ -124,12 +124,9 @@ pub(super) enum ClientConnectionMethod {
     // This ensures that the prompt termination is ordered with respect
     // to the other notifications that are routed to that same session.
     Prompt(acp::PromptRequest),
-    
+
     #[allow(dead_code)] // Will be used when client-side cancellation is implemented
-    Cancel(
-        acp::CancelNotification,
-        oneshot::Sender<Result<(), acp::Error>>,
-    ),
+    Cancel(acp::CancelNotification, oneshot::Sender<Result<(), acp::Error>>),
 }
 
 impl AcpClientConnectionHandle {
@@ -159,6 +156,7 @@ impl AcpClientConnectionHandle {
                 acp::ClientSideConnection::new(callbacks, outgoing_bytes, incoming_bytes, |fut| {
                     tokio::task::spawn_local(fut);
                 });
+            let client_conn = Arc::new(client_conn);
 
             // Start the client I/O handler
             tokio::task::spawn_local(async move {
@@ -168,42 +166,55 @@ impl AcpClientConnectionHandle {
             });
 
             while let Some(method) = client_rx.recv().await {
-                tracing::debug!(actor="client_connection", event="message received", ?method);
+                tracing::debug!(actor = "client_connection", event = "message received", ?method);
 
                 match method {
                     ClientConnectionMethod::Initialize(initialize_request, sender) => {
                         let response = client_conn.initialize(initialize_request).await;
-                        tracing::debug!(actor="client_connection", event="sending response", ?response);
+                        tracing::debug!(actor = "client_connection", event = "sending response", ?response);
                         ignore_error(sender.send(response));
                     },
                     ClientConnectionMethod::NewSession(new_session_request, sender) => {
                         match client_conn.new_session(new_session_request).await {
                             Ok(session_info) => {
-                                let result = AcpClientSessionHandle::new(
-                                    session_info,
-                                    &client_dispatch,
-                                    client_tx.clone(),
-                                )
-                                .await
-                                .map_err(|_err| acp::Error::internal_error());
-                                tracing::debug!(actor="client_connection", event="sending response", ?result);
+                                let result =
+                                    AcpClientSessionHandle::new(session_info, &client_dispatch, client_tx.clone())
+                                        .await
+                                        .map_err(|_err| acp::Error::internal_error());
+                                tracing::debug!(actor = "client_connection", event = "sending response", ?result);
                                 ignore_error(sender.send(result));
                             },
                             Err(err) => {
-                                tracing::debug!(actor="client_connection", event="sending response", ?err);
+                                tracing::debug!(actor = "client_connection", event = "sending response", ?err);
                                 ignore_error(sender.send(Err(err)));
                             },
                         }
                     },
                     ClientConnectionMethod::Prompt(prompt_request) => {
                         let session_id = prompt_request.session_id.clone();
-                        let response = client_conn.prompt(prompt_request).await;
-                        tracing::debug!(actor="client_connection", event="sending response", ?session_id, ?response);
-                        client_dispatch.client_callback(ClientCallback::PromptResponse(session_id, response));
+
+                        // Spawn off the call to prompt so it runs concurrently.
+                        //
+                        // This way if the user tries to cancel, that message can be received
+                        // and sent to the server. That will cause the server to cancel this prompt call.
+                        tokio::task::spawn_local({
+                            let client_conn = client_conn.clone();
+                            let client_dispatch = client_dispatch.clone();
+                            async move {
+                                let response = client_conn.prompt(prompt_request).await;
+                                tracing::debug!(
+                                    actor = "client_connection",
+                                    event = "sending response",
+                                    ?session_id,
+                                    ?response
+                                );
+                                client_dispatch.client_callback(ClientCallback::PromptResponse(session_id, response));
+                            }
+                        });
                     },
                     ClientConnectionMethod::Cancel(cancel_notification, sender) => {
                         let response = client_conn.cancel(cancel_notification).await;
-                        tracing::debug!(actor="client_connection", event="sending response", ?response);
+                        tracing::debug!(actor = "client_connection", event = "sending response", ?response);
                         ignore_error(sender.send(response));
                     },
                 }
@@ -229,12 +240,10 @@ impl AcpClientConnectionHandle {
         Ok(rx.await??)
     }
 
-    #[allow(dead_code)] // Will be used when client-side cancellation is implemented
+    #[cfg_attr(not(test), allow(dead_code))] // Will be used when client-side cancellation is implemented
     pub async fn cancel(&self, args: acp::CancelNotification) -> Result<()> {
         let (tx, rx) = tokio::sync::oneshot::channel();
-        self.client_tx
-            .send(ClientConnectionMethod::Cancel(args, tx))
-            .await?;
+        self.client_tx.send(ClientConnectionMethod::Cancel(args, tx)).await?;
         Ok(rx.await??)
     }
 }
@@ -258,7 +267,10 @@ impl acp::Client for AcpClientForward {
         todo!()
     }
 
-    async fn write_text_file(&self, _args: acp::WriteTextFileRequest) -> Result<acp::WriteTextFileResponse, acp::Error> {
+    async fn write_text_file(
+        &self,
+        _args: acp::WriteTextFileRequest,
+    ) -> Result<acp::WriteTextFileResponse, acp::Error> {
         todo!()
     }
 
@@ -267,12 +279,16 @@ impl acp::Client for AcpClientForward {
     }
 
     async fn session_notification(&self, args: acp::SessionNotification) -> Result<(), acp::Error> {
-        tracing::debug!(actor="client_connection", event="session_notification", ?args);
+        tracing::debug!(actor = "client_connection", event = "session_notification", ?args);
         let (tx, rx) = oneshot::channel();
         self.client_dispatch
             .client_callback(ClientCallback::Notification(args, tx));
         let result = rx.await;
-        tracing::debug!(actor="client_connection", event="session_notification complete", ?result);
+        tracing::debug!(
+            actor = "client_connection",
+            event = "session_notification complete",
+            ?result
+        );
         result.map_err(acp::Error::into_internal_error)?
     }
 
@@ -342,4 +358,3 @@ impl ClientCallback {
         }
     }
 }
-

crates/chat-cli/src/cli/acp/client_session.rs

@@ -35,6 +35,11 @@ impl AcpClientSessionHandle {
         })
     }
 
+    /// Get the session ID for this session
+    pub fn session_id(&self) -> &acp::SessionId {
+        &self.session_info.session_id
+    }
+
     /// Send a message to the agent and read the complete response
     pub async fn prompt(&mut self, message: impl IntoPrompt) -> Result<String> {
         // Construct the prompt
@@ -53,6 +58,7 @@ impl AcpClientSessionHandle {
         // Read notifications until we get the prompt response, then we can return.
         let mut response_text = String::new();
         while let Some(client_callback) = self.callback_rx.recv().await {
+            tracing::debug!(actor="client_session", event="callback received", "session_id"=?self.session_info.session_id, ?client_callback);
             match client_callback {
                 ClientCallback::Notification(notification, tx) => {
                     self.handle_notification(notification, tx, &mut response_text)

crates/chat-cli/src/cli/acp/server.rs

@@ -7,7 +7,7 @@ use agent_client_protocol as acp;
 use serde_json::value::RawValue;
 use tokio::sync::{mpsc, oneshot};
 
-use crate::os::Os;
+use crate::{cli::acp::util::ignore_error, os::Os};
 use super::{server_session::AcpServerSessionHandle, server_connection::AcpServerConnectionHandle};
 
 /// Convert channel errors to ACP errors
@@ -30,6 +30,7 @@ pub struct AcpServerHandle {
 /// Each variant contains:
 /// - Request parameters (the input)
 /// - oneshot::Sender (the "return address" where the actor sends the response back)
+#[derive(Debug)]
 enum ServerMethod {
     Initialize(acp::InitializeRequest, oneshot::Sender<Result<acp::InitializeResponse, acp::Error>>),
     Authenticate(acp::AuthenticateRequest, oneshot::Sender<Result<acp::AuthenticateResponse, acp::Error>>),
@@ -50,6 +51,7 @@ impl AcpServerHandle {
             let mut sessions: HashMap<String, AcpServerSessionHandle> = HashMap::new();
 
             while let Some(method) = server_rx.recv().await {
+                tracing::debug!(actor="server", event="method call received", ?method);
                 match method {
                     ServerMethod::Initialize(args, tx) => {
                         let response = Self::handle_initialize(args).await;
@@ -87,11 +89,7 @@ impl AcpServerHandle {
                         }
                     }
                     ServerMethod::Prompt(args, tx) => {
-                        let response = Self::handle_prompt(args, &sessions).await;
-                        if tx.send(response).is_err() {
-                            tracing::debug!(actor="server", event="response receiver dropped", method="prompt");
-                            break;
-                        }
+                        Self::handle_prompt(args, tx, &sessions).await;
                     }
                     ServerMethod::Cancel(args, tx) => {
                         let response = Self::handle_cancel(args, &sessions).await;
@@ -267,17 +265,20 @@ impl AcpServerHandle {
 
     async fn handle_prompt(
         args: acp::PromptRequest,
+        prompt_tx: oneshot::Sender<Result<acp::PromptResponse, acp::Error>>,
         sessions: &HashMap<String, AcpServerSessionHandle>,
-    ) -> Result<acp::PromptResponse, acp::Error> {
+    ) {
         let session_id = args.session_id.0.as_ref();
 
         // Find the session actor
         if let Some(session_handle) = sessions.get(session_id) {
-            // Forward to session actor
-            session_handle.prompt(args).await
+            // Forward to session actor. Importantly, this actor is responsible
+            // for sending the final result from the prompt to `prompt_tx` -- we just
+            // return immediately.
+            session_handle.prompt(args, prompt_tx).await;
         } else {
             tracing::warn!("Session not found for prompt: {}", session_id);
-            Err(acp::Error::invalid_params())
+            ignore_error(prompt_tx.send(Err(acp::Error::invalid_params())))
         }
     }
 

crates/chat-cli/src/cli/acp/server_connection.rs

@@ -55,16 +55,16 @@ impl AcpServerConnectionHandle {
 
         // Launch the "transport actor", which owns the connection.
         tokio::task::spawn_local(async move {
-            tracing::debug!(actor="transport", event="started");
+            tracing::debug!(actor="server_connection", event="started");
 
             while let Some(method) = transport_rx.recv().await {
-                tracing::debug!(actor="transport", event="message received", ?method);
+                tracing::debug!(actor="server_connection", event="message received", ?method);
                 match method {
                     TransportMethod::SessionNotification(notification, tx) => {
                         let result = connection.session_notification(notification).await;
-                        tracing::debug!(actor="transport", event="notification delivered");
+                        tracing::debug!(actor="server_connection", event="notification delivered");
                         if tx.send(result).is_err() {
-                            tracing::debug!(actor="transport", event="response receiver dropped");
+                            tracing::debug!(actor="server_connection", event="response receiver dropped");
                         }
                     },
                 }