Fix issue where cancellation may cancel also compensation calls in sagas

src/tests/implicit_cancellation.rs

@@ -426,3 +426,203 @@ fn replay_while_cancelling() {
     );
     assert_eq!(output.next(), None);
 }
+
+// Reproduces the is_none_or in is_handle_completed bug:
+// a handler that calls a downstream service, awaits it, then sleeps.
+// While sleeping, the invocation is cancelled. The handler catches the cancellation and runs a *compensation*:
+// a NEW downstream call (the "revert"/undo). The compensation call must run to completion and should not be canceled again.
+#[test]
+fn saga_compensation_call_is_not_cancelled() {
+    let mut output = VMTestCase::new()
+        .input(start_message(1))
+        .input(input_entry_message(b"my-data"))
+        // Forward "reserve" call resolves normally, *before* the cancellation:
+        //  - its invocation id (needed later to propagate the cancel to it)
+        //  - its result
+        .input(CallInvocationIdCompletionNotificationMessage {
+            completion_id: 1,
+            invocation_id: "reserve-invocation-id".to_string(),
+        })
+        .input(CallCompletionNotificationMessage {
+            completion_id: 2,
+            result: Some(call_completion_notification_message::Result::Value(
+                Bytes::from_static(b"reserved").into(),
+            )),
+        })
+        .run_without_closing_input(|vm, encoder| {
+            vm.sys_input().unwrap();
+
+            // ── Forward step: reserve via a downstream service, and await it.
+            let reserve = vm
+                .sys_call(
+                    Target {
+                        service: "Inventory".to_string(),
+                        handler: "reserve".to_string(),
+                        key: None,
+                        idempotency_key: None,
+                        scope: None,
+                        limit_key: None,
+                        headers: Vec::new(),
+                    },
+                    Bytes::from_static(b"flight"),
+                    None,
+                    PayloadOptions::default(),
+                )
+                .unwrap();
+            assert_eq!(
+                vm.do_await(UnresolvedFuture::Single(reserve.call_notification_handle))
+                    .unwrap(),
+                AwaitResponse::AnyCompleted
+            );
+            assert2::assert!(
+                let Some(Value::Success(reserved)) = vm
+                    .take_notification(reserve.call_notification_handle)
+                    .unwrap()
+            );
+            assert_eq!(reserved, Bytes::from_static(b"reserved"));
+
+            // ── Long sleep — the window in which we get cancelled.
+            let sleep = vm
+                .sys_sleep(String::default(), Duration::from_secs(60), None)
+                .unwrap();
+
+            // The cancellation arrives now, while we're sleeping.
+            vm.notify_input(encoder.encode(&cancel_signal_notification()));
+
+            // Awaiting the sleep surfaces the cancellation — this is where the
+            // saga would `catch` and start compensating.
+            assert_eq!(
+                vm.do_await(UnresolvedFuture::Single(sleep)).unwrap(),
+                AwaitResponse::CancelSignalReceived
+            );
+
+            // ── Compensation step: unreserve via the downstream service.
+            // A NEW call, made AFTER the cancellation was caught. It must run to
+            // completion and MUST NOT be cancelled by the runtime.
+            let unreserve = vm
+                .sys_call(
+                    Target {
+                        service: "Inventory".to_string(),
+                        handler: "unreserve".to_string(),
+                        key: None,
+                        idempotency_key: None,
+                        scope: None,
+                        limit_key: None,
+                        headers: Vec::new(),
+                    },
+                    Bytes::from_static(b"flight"),
+                    None,
+                    PayloadOptions::default(),
+                )
+                .unwrap();
+
+            // Before the fix, this returned CancelSignalReceived.
+            // With the fix, it completes normally.
+            assert_eq!(
+                vm.do_await(UnresolvedFuture::Single(unreserve.call_notification_handle))
+                    .unwrap(),
+                AwaitResponse::WaitingExternalProgress {
+                    waiting_input: true,
+                    waiting_run_proposal: false
+                }
+            );
+
+            // Send notifications for unreserve
+            vm.notify_input(
+                encoder.encode(&CallInvocationIdCompletionNotificationMessage {
+                    completion_id: 4,
+                    invocation_id: "unreserve-invocation-id".to_string(),
+                }),
+            );
+            vm.notify_input(encoder.encode(&CallCompletionNotificationMessage {
+                completion_id: 5,
+                result: Some(call_completion_notification_message::Result::Value(
+                    Bytes::from_static(b"unreserved").into(),
+                )),
+            }));
+            vm.notify_input_closed();
+
+            // Before the fix, this returned CancelSignalReceived.
+            // With the fix, it completes normally.
+            assert_eq!(
+                vm.do_await(UnresolvedFuture::Single(unreserve.call_notification_handle))
+                    .unwrap(),
+                AwaitResponse::AnyCompleted
+            );
+            assert2::assert!(
+                let Some(Value::Success(unreserved)) = vm
+                    .take_notification(unreserve.call_notification_handle)
+                    .unwrap()
+            );
+            assert_eq!(unreserved, Bytes::from_static(b"unreserved"));
+
+            vm.sys_write_output(
+                NonEmptyValue::Success(unreserved),
+                PayloadOptions::default(),
+            )
+            .unwrap();
+            vm.sys_end().unwrap();
+        });
+
+    // Forward reserve call.
+    assert_that!(
+        output.next_decoded::<CallCommandMessage>().unwrap(),
+        pat!(CallCommandMessage {
+            service_name: eq("Inventory"),
+            handler_name: eq("reserve"),
+            invocation_id_notification_idx: eq(1),
+            result_completion_id: eq(2),
+        })
+    );
+    // The sleep.
+    assert_that!(
+        output.next_decoded::<SleepCommandMessage>().unwrap(),
+        pat!(SleepCommandMessage {
+            result_completion_id: eq(3)
+        })
+    );
+    // Cancellation propagates to the in-flight forward call.
+    assert_eq!(
+        output.next_decoded::<SendSignalCommandMessage>().unwrap(),
+        SendSignalCommandMessage {
+            target_invocation_id: "reserve-invocation-id".to_string(),
+            signal_id: Some(send_signal_command_message::SignalId::Idx(CANCEL_SIGNAL_ID)),
+            result: Some(send_signal_command_message::Result::Void(Default::default())),
+            ..Default::default()
+        }
+    );
+    // The compensation call is journaled...
+    assert_that!(
+        output.next_decoded::<CallCommandMessage>().unwrap(),
+        pat!(CallCommandMessage {
+            service_name: eq("Inventory"),
+            handler_name: eq("unreserve"),
+            invocation_id_notification_idx: eq(4),
+            result_completion_id: eq(5),
+        })
+    );
+    // The awaiting on the new call
+    assert_that!(
+        output.next_decoded::<AwaitingOnMessage>().unwrap(),
+        pat!(AwaitingOnMessage {
+            awaiting_on: some(pat!(Future {
+                waiting_signals: unordered_elements_are![eq(1)],
+                waiting_completions: unordered_elements_are![eq(5)],
+                waiting_named_signals: empty(),
+                nested_futures: empty(),
+                combinator_type: eq(CombinatorType::FirstCompleted as i32)
+            })),
+            executing_side_effects: eq(false)
+        })
+    );
+    // Finally, the output. No SendSignalCommandMessage targeting "unreserve-invocation-id".
+    assert_that!(
+        output.next_decoded::<OutputCommandMessage>().unwrap(),
+        is_output_with_success(b"unreserved")
+    );
+    assert_eq!(
+        output.next_decoded::<EndMessage>().unwrap(),
+        EndMessage::default()
+    );
+    assert_eq!(output.next(), None);
+}

src/vm/async_results_state.rs

@@ -308,7 +308,7 @@ impl AsyncResultsState {
     pub(crate) fn is_handle_completed(&self, handle: NotificationHandle) -> bool {
         self.handle_mapping
             .get(&handle)
-            .is_none_or(|id| self.ready.contains_key(id))
+            .is_some_and(|id| self.ready.contains_key(id))
     }
 
     fn resolve_handle_state(&self, handle: NotificationHandle) -> HandleState {
@@ -417,7 +417,11 @@ impl AsyncResultsState {
     pub(crate) fn take_handle(&mut self, handle: NotificationHandle) -> Option<NotificationResult> {
         let id = self.handle_mapping.get(&handle)?;
         if let Some(res) = self.ready.remove(id) {
-            self.handle_mapping.remove(&handle);
+            if handle != CANCEL_NOTIFICATION_HANDLE {
+                self.handle_mapping.remove(&handle);
+                // Don't remove the CANCEL handle mapping
+                // TODO improve the code here around handling built-in signals handles
+            }
             Some(res)
         } else {
             None