fix(acp): disarm cancellation for buffered responses

Author: benbrandt

src/agent-client-protocol/src/jsonrpc.rs

@@ -1358,6 +1358,9 @@ enum ReplyMessage {
         method: String,
 
         sender: oneshot::Sender<ResponsePayload>,
+
+        #[cfg(feature = "unstable_cancel_request")]
+        cancellation_disarm: SentRequestCancellationDisarm,
     },
 }
 
@@ -1654,6 +1657,9 @@ enum OutgoingMessage {
 
         /// where to send the response when it arrives (includes ack channel)
         response_tx: oneshot::Sender<ResponsePayload>,
+
+        #[cfg(feature = "unstable_cancel_request")]
+        cancellation_disarm: SentRequestCancellationDisarm,
     },
 
     /// Send a notification to the server.
@@ -2003,6 +2009,8 @@ impl<Counterpart: Role> ConnectionTo<Counterpart> {
                     role_id,
                     untyped,
                     response_tx,
+                    #[cfg(feature = "unstable_cancel_request")]
+                    cancellation_disarm: cancellation.disarm_handle(),
                 };
 
                 match self.message_tx.unbounded_send(message) {
@@ -2455,6 +2463,8 @@ impl ResponseRouter<serde_json::Value> {
         id: jsonrpcmsg::Id,
         role_id: RoleId,
         sender: oneshot::Sender<ResponsePayload>,
+        #[cfg(feature = "unstable_cancel_request")]
+        cancellation_disarm: SentRequestCancellationDisarm,
     ) -> Self {
         let response_method = method.clone();
         let response_id = id.clone();
@@ -2475,6 +2485,9 @@ impl ResponseRouter<serde_json::Value> {
                         id = ?response_id,
                         "dropped response because local receiver was gone"
                     );
+                } else {
+                    #[cfg(feature = "unstable_cancel_request")]
+                    cancellation_disarm.disarm();
                 }
                 Ok(())
             }),
@@ -3184,16 +3197,34 @@ fn jsonrpc_id_to_request_id(id: &jsonrpcmsg::Id) -> Result<crate::schema::Reques
 }
 
 #[cfg(feature = "unstable_cancel_request")]
-#[derive(Clone)]
+#[derive(Clone, Debug)]
+pub(crate) struct SentRequestCancellationDisarm {
+    armed: Arc<AtomicBool>,
+}
+
+#[cfg(feature = "unstable_cancel_request")]
+impl SentRequestCancellationDisarm {
+    fn new() -> Self {
+        Self {
+            armed: Arc::new(AtomicBool::new(true)),
+        }
+    }
+
+    fn disarm(&self) {
+        self.armed.store(false, Ordering::Release);
+    }
+}
+
+#[cfg(feature = "unstable_cancel_request")]
 enum SentRequestCancellation {
     Send {
         message_tx: OutgoingMessageTx,
         notification: UntypedMessage,
-        armed: Arc<AtomicBool>,
+        disarm: SentRequestCancellationDisarm,
     },
     Failed {
         error: String,
-        armed: Arc<AtomicBool>,
+        disarm: SentRequestCancellationDisarm,
     },
 }
 
@@ -3211,25 +3242,25 @@ impl SentRequestCancellation {
                 )
             })
             .map_err(|error| error.to_string());
+        let disarm = SentRequestCancellationDisarm::new();
 
         match notification {
             Ok(notification) => Self::Send {
                 message_tx,
                 notification,
-                armed: Arc::new(AtomicBool::new(true)),
-            },
-            Err(error) => Self::Failed {
-                error,
-                armed: Arc::new(AtomicBool::new(true)),
+                disarm,
             },
+            Err(error) => Self::Failed { error, disarm },
         }
     }
 
     fn disarm(&self) {
+        self.disarm_handle().disarm();
+    }
+
+    fn disarm_handle(&self) -> SentRequestCancellationDisarm {
         match self {
-            Self::Send { armed, .. } | Self::Failed { armed, .. } => {
-                armed.store(false, Ordering::Release);
-            }
+            Self::Send { disarm, .. } | Self::Failed { disarm, .. } => disarm.clone(),
         }
     }
 
@@ -3238,9 +3269,9 @@ impl SentRequestCancellation {
             Self::Send {
                 message_tx,
                 notification,
-                armed,
+                disarm,
             } => {
-                if !armed.swap(false, Ordering::AcqRel) {
+                if !disarm.armed.swap(false, Ordering::AcqRel) {
                     return Ok(());
                 }
 
@@ -3251,8 +3282,8 @@ impl SentRequestCancellation {
                     },
                 )
             }
-            Self::Failed { error, armed } => {
-                if !armed.swap(false, Ordering::AcqRel) {
+            Self::Failed { error, disarm } => {
+                if !disarm.armed.swap(false, Ordering::AcqRel) {
                     return Ok(());
                 }
 
@@ -3279,17 +3310,17 @@ impl Debug for SentRequestCancellation {
         match self {
             Self::Send {
                 notification,
-                armed,
+                disarm,
                 ..
             } => f
                 .debug_struct("SentRequestCancellation")
                 .field("notification", notification)
-                .field("armed", &armed.load(Ordering::Acquire))
+                .field("armed", &disarm.armed.load(Ordering::Acquire))
                 .finish(),
-            Self::Failed { error, armed } => f
+            Self::Failed { error, disarm } => f
                 .debug_struct("SentRequestCancellation")
                 .field("error", error)
-                .field("armed", &armed.load(Ordering::Acquire))
+                .field("armed", &disarm.armed.load(Ordering::Acquire))
                 .finish(),
         }
     }

src/agent-client-protocol/src/jsonrpc/incoming_actor.rs

@@ -30,6 +30,8 @@ struct PendingReply {
     method: String,
     role_id: RoleId,
     sender: oneshot::Sender<crate::jsonrpc::ResponsePayload>,
+    #[cfg(feature = "unstable_cancel_request")]
+    cancellation_disarm: super::SentRequestCancellationDisarm,
 }
 
 /// Incoming protocol actor: The central dispatch loop for a connection.
@@ -78,6 +80,8 @@ pub(super) async fn incoming_protocol_actor<Counterpart: Role>(
                     role_id,
                     method,
                     sender,
+                    #[cfg(feature = "unstable_cancel_request")]
+                    cancellation_disarm,
                 } => {
                     tracing::trace!(?id, %method, "incoming_actor: subscribing to response");
                     let id = serde_json::to_value(&id).unwrap();
@@ -87,6 +91,8 @@ pub(super) async fn incoming_protocol_actor<Counterpart: Role>(
                             method,
                             role_id,
                             sender,
+                            #[cfg(feature = "unstable_cancel_request")]
+                            cancellation_disarm,
                         },
                     );
                 }
@@ -260,10 +266,19 @@ fn dispatch_from_response(
         method,
         role_id,
         sender,
+        #[cfg(feature = "unstable_cancel_request")]
+        cancellation_disarm,
     } = pending_reply;
 
     // Create a Dispatch::Response with a ResponseRouter that routes to the oneshot
-    let router = ResponseRouter::new(method.clone(), id.clone(), role_id, sender);
+    let router = ResponseRouter::new(
+        method.clone(),
+        id.clone(),
+        role_id,
+        sender,
+        #[cfg(feature = "unstable_cancel_request")]
+        cancellation_disarm,
+    );
     Dispatch::Response(result, router)
 }
 

src/agent-client-protocol/src/jsonrpc/outgoing_actor.rs

@@ -41,6 +41,8 @@ pub(super) async fn outgoing_protocol_actor(
                 method,
                 untyped,
                 response_tx,
+                #[cfg(feature = "unstable_cancel_request")]
+                cancellation_disarm,
             } => {
                 let request = match protocol_compat
                     .outgoing_message(untyped)
@@ -49,6 +51,8 @@ pub(super) async fn outgoing_protocol_actor(
                     Ok(request) => request,
                     Err(error) => {
                         tracing::warn!(?id, %method, ?error, "Failed to convert outgoing request");
+                        #[cfg(feature = "unstable_cancel_request")]
+                        cancellation_disarm.disarm();
                         complete_request_with_error(response_tx, error);
                         continue;
                     }
@@ -61,6 +65,8 @@ pub(super) async fn outgoing_protocol_actor(
                         role_id,
                         method,
                         sender: response_tx,
+                        #[cfg(feature = "unstable_cancel_request")]
+                        cancellation_disarm,
                     })
                     .map_err(crate::Error::into_internal_error)?;
 
@@ -167,6 +173,8 @@ mod tests {
                 method: "session/new".into(),
                 untyped: malformed_v2_known_method()?,
                 response_tx,
+                #[cfg(feature = "unstable_cancel_request")]
+                cancellation_disarm: crate::jsonrpc::SentRequestCancellationDisarm::new(),
             })
             .map_err(crate::Error::into_internal_error)?;
         drop(outgoing_tx);

src/agent-client-protocol/tests/jsonrpc_request_cancellation.rs

@@ -628,6 +628,78 @@ async fn late_response_after_dropped_sent_request_does_not_close_connection() {
         .await;
 }
 
+#[tokio::test(flavor = "current_thread")]
+async fn response_buffered_before_drop_disarms_auto_cancellation() {
+    use tokio::task::LocalSet;
+
+    let local = LocalSet::new();
+
+    local
+        .run_until(async {
+            let received = Arc::new(Mutex::new(Vec::new()));
+            let received_for_handler = received.clone();
+
+            let (server_reader, server_writer, client_reader, client_writer) = setup_test_streams();
+            let server_transport =
+                agent_client_protocol::ByteStreams::new(server_writer, server_reader);
+            let server = UntypedRole
+                .builder()
+                .on_receive_request(
+                    async |request: SimpleRequest,
+                           responder: Responder<SimpleResponse>,
+                           _connection: ConnectionTo<UntypedRole>| {
+                        responder.respond(SimpleResponse {
+                            result: format!("echo: {}", request.message),
+                        })
+                    },
+                    agent_client_protocol::on_receive_request!(),
+                )
+                .on_receive_notification(
+                    async move |notification: CancelRequestNotification,
+                                _connection: ConnectionTo<UntypedRole>| {
+                        received_for_handler
+                            .lock()
+                            .unwrap()
+                            .push(notification.request_id);
+                        Ok(())
+                    },
+                    agent_client_protocol::on_receive_notification!(),
+                );
+
+            tokio::task::spawn_local(async move {
+                if let Err(error) = server.connect_to(server_transport).await {
+                    panic!("server should stay alive: {error:?}");
+                }
+            });
+
+            let client_transport =
+                agent_client_protocol::ByteStreams::new(client_writer, client_reader);
+            let response = UntypedRole
+                .builder()
+                .connect_with(client_transport, async |cx| {
+                    let request: SentRequest<SimpleResponse> = cx.send_request(SimpleRequest {
+                        message: "buffered".into(),
+                    });
+
+                    tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
+                    drop(request);
+                    tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
+
+                    cx.send_request(SimpleRequest {
+                        message: "after buffered".into(),
+                    })
+                    .block_task()
+                    .await
+                })
+                .await
+                .unwrap();
+
+            assert_eq!(response.result, "echo: after buffered");
+            assert!(received.lock().unwrap().is_empty());
+        })
+        .await;
+}
+
 #[tokio::test(flavor = "current_thread")]
 async fn completed_sent_request_does_not_send_cancellation_on_drop() {
     use tokio::task::LocalSet;