Fix little bug with run cancellation

Author: slinkydeveloper

src/endpoint/context.rs

@@ -16,7 +16,7 @@ use futures::{FutureExt, TryFutureExt};
 use pin_project_lite::pin_project;
 use restate_sdk_shared_core::{
     CoreVM, DoProgressResponse, Error as CoreError, NonEmptyValue, NotificationHandle, RetryPolicy,
-    RunExitResult, TakeOutputResult, Target, Value, VM,
+    RunExitResult, TakeOutputResult, Target, TerminalFailure, Value, VM,
 };
 use std::borrow::Cow;
 use std::collections::HashMap;
@@ -859,6 +859,19 @@ where
                                 closure_fut: closure.run(),
                             });
                         }
+                        Ok(DoProgressResponse::CancelSignalReceived) => {
+                            drop(inner_ctx);
+                            // Got cancellation!
+                            this.state.set(RunState::WaitingResultFut {
+                                result_fut: async {
+                                    Ok(Err(TerminalError::from(TerminalFailure {
+                                        code: 409,
+                                        message: "cancelled".to_string(),
+                                    })))
+                                }
+                                .boxed(),
+                            })
+                        }
                         _ => {
                             drop(inner_ctx);
                             // In all the other cases, just move on waiting the result,