@temporalio/core-bridge 1.5.2 → 1.7.0

package/sdk-core/core/src/worker/workflow/machines/activity_state_machine.rs

@@ -4,6 +4,10 @@ use super::{
     workflow_machines::MachineResponse, Cancellable, EventInfo, NewMachineWithCommand,
     OnEventWrapper, WFMachinesAdapter, WFMachinesError,
 };
+use crate::{
+    internal_flags::CoreInternalFlags,
+    worker::workflow::{machines::HistEventData, InternalFlagsRef},
+};
 use rustfsm::{fsm, MachineError, StateMachine, TransitionResult};
 use std::convert::{TryFrom, TryInto};
 use temporal_sdk_core_protos::{
@@ -16,10 +20,7 @@ use temporal_sdk_core_protos::{
         command::v1::{command, Command, RequestCancelActivityTaskCommandAttributes},
         common::v1::{ActivityType, Payload, Payloads},
         enums::v1::{CommandType, EventType, RetryState},
-        failure::v1::{
-            self as failure, failure::FailureInfo, ActivityFailureInfo, CanceledFailureInfo,
-            Failure,
-        },
+        failure::v1::{failure::FailureInfo, ActivityFailureInfo, CanceledFailureInfo, Failure},
         history::v1::{
             history_event, ActivityTaskCanceledEventAttributes,
             ActivityTaskCompletedEventAttributes, ActivityTaskFailedEventAttributes,
@@ -37,7 +38,7 @@ fsm! {
     Created --(Schedule, on_schedule)--> ScheduleCommandCreated;
 
     ScheduleCommandCreated --(CommandScheduleActivityTask) --> ScheduleCommandCreated;
-    ScheduleCommandCreated --(ActivityTaskScheduled(i64),
+    ScheduleCommandCreated --(ActivityTaskScheduled(ActTaskScheduledData),
         shared on_activity_task_scheduled) --> ScheduledEventRecorded;
     ScheduleCommandCreated --(Cancel, shared on_canceled) --> Canceled;
 
@@ -45,7 +46,7 @@ fsm! {
     ScheduledEventRecorded --(ActivityTaskTimedOut(ActivityTaskTimedOutEventAttributes),
         shared on_task_timed_out) --> TimedOut;
     ScheduledEventRecorded --(Cancel, shared on_canceled) --> ScheduledActivityCancelCommandCreated;
-    ScheduledEventRecorded --(Abandon, shared on_abandoned) --> Canceled;
+    ScheduledEventRecorded --(Abandon, on_abandoned) --> Canceled;
 
     Started --(ActivityTaskCompleted(ActivityTaskCompletedEventAttributes),
         on_activity_task_completed) --> Completed;
@@ -54,7 +55,7 @@ fsm! {
     Started --(ActivityTaskTimedOut(ActivityTaskTimedOutEventAttributes),
         shared on_activity_task_timed_out) --> TimedOut;
     Started --(Cancel, shared on_canceled) --> StartedActivityCancelCommandCreated;
-    Started --(Abandon, shared on_abandoned) --> Canceled;
+    Started --(Abandon, on_abandoned) --> Canceled;
 
     ScheduledActivityCancelCommandCreated --(CommandRequestCancelActivityTask) --> ScheduledActivityCancelCommandCreated;
     ScheduledActivityCancelCommandCreated --(ActivityTaskCancelRequested) --> ScheduledActivityCancelEventRecorded;
@@ -89,39 +90,46 @@ pub(super) enum ActivityMachineCommand {
     #[display(fmt = "Fail")]
     Fail(Failure),
     #[display(fmt = "Cancel")]
-    Cancel(Option<Payloads>),
+    Cancel(Option<ActivityTaskCanceledEventAttributes>),
     #[display(fmt = "RequestCancellation")]
     RequestCancellation(Command),
 }
 
-/// Creates a new activity state machine and a command to schedule it on the server.
-pub(super) fn new_activity(attribs: ScheduleActivity) -> NewMachineWithCommand {
-    let (activity, add_cmd) = ActivityMachine::new_scheduled(attribs);
-    NewMachineWithCommand {
-        command: add_cmd,
-        machine: activity.into(),
-    }
+pub(super) struct ActTaskScheduledData {
+    event_id: i64,
+    act_type: String,
+    act_id: String,
+    last_task_in_history: bool,
 }
 
 impl ActivityMachine {
     /// Create a new activity and immediately schedule it.
-    fn new_scheduled(attribs: ScheduleActivity) -> (Self, Command) {
-        let mut s = Self {
-            state: Created {}.into(),
-            shared_state: SharedState {
-                cancellation_type: ActivityCancellationType::from_i32(attribs.cancellation_type)
+    pub(super) fn new_scheduled(
+        attrs: ScheduleActivity,
+        internal_flags: InternalFlagsRef,
+    ) -> NewMachineWithCommand {
+        let mut s = Self::from_parts(
+            Created {}.into(),
+            SharedState {
+                cancellation_type: ActivityCancellationType::from_i32(attrs.cancellation_type)
                     .unwrap(),
-                attrs: attribs,
-                ..Default::default()
+                attrs,
+                internal_flags,
+                scheduled_event_id: 0,
+                started_event_id: 0,
+                cancelled_before_sent: false,
             },
-        };
+        );
         OnEventWrapper::on_event_mut(&mut s, ActivityMachineEvents::Schedule)
             .expect("Scheduling activities doesn't fail");
-        let cmd = Command {
+        let command = Command {
             command_type: CommandType::ScheduleActivityTask as i32,
             attributes: Some(s.shared_state().attrs.clone().into()),
         };
-        (s, cmd)
+        NewMachineWithCommand {
+            command,
+            machine: s.into(),
+        }
     }
 
     fn machine_responses_from_cancel_request(&self, cancel_cmd: Command) -> Vec<MachineResponse> {
@@ -136,45 +144,46 @@ impl ActivityMachine {
         r
     }
 
-    fn create_cancelation_resolve(&self, details: Option<Payload>) -> ResolveActivity {
+    fn create_cancelation_resolve(
+        &self,
+        attrs: Option<ActivityTaskCanceledEventAttributes>,
+    ) -> ResolveActivity {
+        let attrs = attrs.unwrap_or_default();
         ResolveActivity {
             seq: self.shared_state.attrs.seq,
             result: Some(ActivityResolution {
                 status: Some(activity_resolution::Status::Cancelled(Cancellation {
-                    failure: Some(Failure {
-                        message: "Activity cancelled".to_string(),
-                        cause: Some(Box::from(Failure {
-                            failure_info: Some(FailureInfo::CanceledFailureInfo(
-                                CanceledFailureInfo {
-                                    details: details.map(Into::into),
-                                },
-                            )),
-                            ..Default::default()
-                        })),
-                        failure_info: Some(FailureInfo::ActivityFailureInfo(ActivityFailureInfo {
-                            scheduled_event_id: self.shared_state.scheduled_event_id,
-                            started_event_id: self.shared_state.started_event_id,
-                            activity_type: Some(ActivityType {
-                                name: self.shared_state.attrs.activity_type.clone(),
-                            }),
-                            activity_id: self.shared_state.attrs.activity_id.clone(),
-                            retry_state: RetryState::CancelRequested as i32,
-                            ..Default::default()
-                        })),
-                        ..Default::default()
-                    }),
+                    failure: Some(new_cancel_failure(&self.shared_state, attrs)),
                 })),
             }),
         }
     }
 }
 
-impl TryFrom<HistoryEvent> for ActivityMachineEvents {
+impl TryFrom<HistEventData> for ActivityMachineEvents {
     type Error = WFMachinesError;
 
-    fn try_from(e: HistoryEvent) -> Result<Self, Self::Error> {
+    fn try_from(e: HistEventData) -> Result<Self, Self::Error> {
+        let last_task_in_history = e.current_task_is_last_in_history;
+        let e = e.event;
         Ok(match e.event_type() {
-            EventType::ActivityTaskScheduled => Self::ActivityTaskScheduled(e.event_id),
+            EventType::ActivityTaskScheduled => {
+                if let Some(history_event::Attributes::ActivityTaskScheduledEventAttributes(
+                    attrs,
+                )) = e.attributes
+                {
+                    Self::ActivityTaskScheduled(ActTaskScheduledData {
+                        event_id: e.event_id,
+                        act_id: attrs.activity_id,
+                        act_type: attrs.activity_type.unwrap_or_default().name,
+                        last_task_in_history,
+                    })
+                } else {
+                    return Err(WFMachinesError::Fatal(format!(
+                        "Activity scheduled attributes were unset: {e}"
+                    )));
+                }
+            }
             EventType::ActivityTaskStarted => Self::ActivityTaskStarted(e.event_id),
             EventType::ActivityTaskCompleted => {
                 if let Some(history_event::Attributes::ActivityTaskCompletedEventAttributes(
@@ -184,8 +193,7 @@ impl TryFrom<HistoryEvent> for ActivityMachineEvents {
                     Self::ActivityTaskCompleted(attrs)
                 } else {
                     return Err(WFMachinesError::Fatal(format!(
-                        "Activity completion attributes were unset: {}",
-                        e
+                        "Activity completion attributes were unset: {e}"
                     )));
                 }
             }
@@ -196,8 +204,7 @@ impl TryFrom<HistoryEvent> for ActivityMachineEvents {
                     Self::ActivityTaskFailed(attrs)
                 } else {
                     return Err(WFMachinesError::Fatal(format!(
-                        "Activity failure attributes were unset: {}",
-                        e
+                        "Activity failure attributes were unset: {e}"
                     )));
                 }
             }
@@ -208,8 +215,7 @@ impl TryFrom<HistoryEvent> for ActivityMachineEvents {
                     Self::ActivityTaskTimedOut(attrs)
                 } else {
                     return Err(WFMachinesError::Fatal(format!(
-                        "Activity timeout attributes were unset: {}",
-                        e
+                        "Activity timeout attributes were unset: {e}"
                     )));
                 }
             }
@@ -221,15 +227,13 @@ impl TryFrom<HistoryEvent> for ActivityMachineEvents {
                     Self::ActivityTaskCanceled(attrs)
                 } else {
                     return Err(WFMachinesError::Fatal(format!(
-                        "Activity cancellation attributes were unset: {}",
-                        e
+                        "Activity cancellation attributes were unset: {e}"
                     )));
                 }
             }
             _ => {
                 return Err(WFMachinesError::Nondeterminism(format!(
-                    "Activity machine does not handle this event: {}",
-                    e
+                    "Activity machine does not handle this event: {e}"
                 )))
             }
         })
@@ -268,10 +272,8 @@ impl WFMachinesAdapter for ActivityMachine {
             ActivityMachineCommand::RequestCancellation(c) => {
                 self.machine_responses_from_cancel_request(c)
             }
-            ActivityMachineCommand::Cancel(details) => {
-                vec![self
-                    .create_cancelation_resolve(convert_payloads(event_info, details)?)
-                    .into()]
+            ActivityMachineCommand::Cancel(attrs) => {
+                vec![self.create_cancelation_resolve(attrs).into()]
             }
         })
     }
@@ -330,16 +332,9 @@ impl Cancellable for ActivityMachine {
                     self.machine_responses_from_cancel_request(cmd)
                 }
                 ActivityMachineCommand::Cancel(details) => {
-                    vec![self
-                        .create_cancelation_resolve(
-                            details
-                                .map(TryInto::try_into)
-                                .transpose()
-                                .unwrap_or_default(),
-                        )
-                        .into()]
+                    vec![self.create_cancelation_resolve(details).into()]
                 }
-                x => panic!("Invalid cancel event response {:?}", x),
+                x => panic!("Invalid cancel event response {x:?}"),
             })
             .collect();
         Ok(res)
@@ -350,13 +345,14 @@ impl Cancellable for ActivityMachine {
     }
 }
 
-#[derive(Default, Clone)]
+#[derive(Clone)]
 pub(super) struct SharedState {
     scheduled_event_id: i64,
     started_event_id: i64,
     attrs: ScheduleActivity,
     cancellation_type: ActivityCancellationType,
     cancelled_before_sent: bool,
+    internal_flags: InternalFlagsRef,
 }
 
 #[derive(Default, Clone)]
@@ -375,28 +371,37 @@ pub(super) struct ScheduleCommandCreated {}
 impl ScheduleCommandCreated {
     pub(super) fn on_activity_task_scheduled(
         self,
-        dat: SharedState,
-        scheduled_event_id: i64,
+        dat: &mut SharedState,
+        sched_dat: ActTaskScheduledData,
     ) -> ActivityMachineTransition<ScheduledEventRecorded> {
-        ActivityMachineTransition::ok_shared(
-            vec![],
-            ScheduledEventRecorded::default(),
-            SharedState {
-                scheduled_event_id,
-                ..dat
-            },
-        )
+        if dat.internal_flags.borrow_mut().try_use(
+            CoreInternalFlags::IdAndTypeDeterminismChecks,
+            sched_dat.last_task_in_history,
+        ) {
+            if sched_dat.act_id != dat.attrs.activity_id {
+                return TransitionResult::Err(WFMachinesError::Nondeterminism(format!(
+                    "Activity id of scheduled event '{}' does not \
+                 match activity id of activity command '{}'",
+                    sched_dat.act_id, dat.attrs.activity_id
+                )));
+            }
+            if sched_dat.act_type != dat.attrs.activity_type {
+                return TransitionResult::Err(WFMachinesError::Nondeterminism(format!(
+                    "Activity type of scheduled event '{}' does not \
+                 match activity type of activity command '{}'",
+                    sched_dat.act_type, dat.attrs.activity_type
+                )));
+            }
+        }
+        dat.scheduled_event_id = sched_dat.event_id;
+        ActivityMachineTransition::default()
     }
-    pub(super) fn on_canceled(self, dat: SharedState) -> ActivityMachineTransition<Canceled> {
-        let canceled_state = SharedState {
-            cancelled_before_sent: true,
-            ..dat
-        };
+
+    pub(super) fn on_canceled(self, dat: &mut SharedState) -> ActivityMachineTransition<Canceled> {
+        dat.cancelled_before_sent = true;
         match dat.cancellation_type {
-            ActivityCancellationType::Abandon => {
-                ActivityMachineTransition::ok_shared(vec![], Canceled::default(), canceled_state)
-            }
-            _ => notify_lang_activity_cancelled(canceled_state, None),
+            ActivityCancellationType::Abandon => ActivityMachineTransition::default(),
+            _ => notify_lang_activity_cancelled(None),
         }
     }
 }
@@ -407,21 +412,15 @@ pub(super) struct ScheduledEventRecorded {}
 impl ScheduledEventRecorded {
     pub(super) fn on_task_started(
         self,
-        dat: SharedState,
+        dat: &mut SharedState,
         started_event_id: i64,
     ) -> ActivityMachineTransition<Started> {
-        ActivityMachineTransition::ok_shared(
-            vec![],
-            Started::default(),
-            SharedState {
-                started_event_id,
-                ..dat
-            },
-        )
+        dat.started_event_id = started_event_id;
+        ActivityMachineTransition::default()
     }
     pub(super) fn on_task_timed_out(
         self,
-        dat: SharedState,
+        dat: &mut SharedState,
         attrs: ActivityTaskTimedOutEventAttributes,
     ) -> ActivityMachineTransition<TimedOut> {
         notify_lang_activity_timed_out(dat, attrs)
@@ -429,15 +428,15 @@ impl ScheduledEventRecorded {
 
     pub(super) fn on_canceled(
         self,
-        dat: SharedState,
+        dat: &mut SharedState,
     ) -> ActivityMachineTransition<ScheduledActivityCancelCommandCreated> {
         create_request_cancel_activity_task_command(
             dat,
             ScheduledActivityCancelCommandCreated::default(),
         )
     }
-    pub(super) fn on_abandoned(self, dat: SharedState) -> ActivityMachineTransition<Canceled> {
-        notify_lang_activity_cancelled(dat, None)
+    pub(super) fn on_abandoned(self) -> ActivityMachineTransition<Canceled> {
+        notify_lang_activity_cancelled(None)
     }
 }
 
@@ -456,7 +455,7 @@ impl Started {
     }
     pub(super) fn on_activity_task_failed(
         self,
-        dat: SharedState,
+        dat: &mut SharedState,
         attrs: ActivityTaskFailedEventAttributes,
     ) -> ActivityMachineTransition<Failed> {
         ActivityMachineTransition::ok(
@@ -467,7 +466,7 @@ impl Started {
 
     pub(super) fn on_activity_task_timed_out(
         self,
-        dat: SharedState,
+        dat: &mut SharedState,
         attrs: ActivityTaskTimedOutEventAttributes,
     ) -> ActivityMachineTransition<TimedOut> {
         notify_lang_activity_timed_out(dat, attrs)
@@ -475,15 +474,15 @@ impl Started {
 
     pub(super) fn on_canceled(
         self,
-        dat: SharedState,
+        dat: &mut SharedState,
     ) -> ActivityMachineTransition<StartedActivityCancelCommandCreated> {
         create_request_cancel_activity_task_command(
             dat,
             StartedActivityCancelCommandCreated::default(),
         )
     }
-    pub(super) fn on_abandoned(self, dat: SharedState) -> ActivityMachineTransition<Canceled> {
-        notify_lang_activity_cancelled(dat, None)
+    pub(super) fn on_abandoned(self) -> ActivityMachineTransition<Canceled> {
+        notify_lang_activity_cancelled(None)
     }
 }
 
@@ -496,15 +495,15 @@ pub(super) struct ScheduledActivityCancelEventRecorded {}
 impl ScheduledActivityCancelEventRecorded {
     pub(super) fn on_activity_task_canceled(
         self,
-        dat: SharedState,
+        dat: &mut SharedState,
         attrs: ActivityTaskCanceledEventAttributes,
     ) -> ActivityMachineTransition<Canceled> {
-        notify_if_not_already_cancelled(dat, |dat| notify_lang_activity_cancelled(dat, Some(attrs)))
+        notify_if_not_already_cancelled(dat, |_| notify_lang_activity_cancelled(Some(attrs)))
     }
 
     pub(super) fn on_activity_task_timed_out(
         self,
-        dat: SharedState,
+        dat: &mut SharedState,
         attrs: ActivityTaskTimedOutEventAttributes,
     ) -> ActivityMachineTransition<TimedOut> {
         notify_if_not_already_cancelled(dat, |dat| notify_lang_activity_timed_out(dat, attrs))
@@ -526,7 +525,7 @@ pub(super) struct StartedActivityCancelEventRecorded {}
 impl StartedActivityCancelEventRecorded {
     pub(super) fn on_activity_task_completed(
         self,
-        dat: SharedState,
+        dat: &mut SharedState,
         attrs: ActivityTaskCompletedEventAttributes,
     ) -> ActivityMachineTransition<Completed> {
         notify_if_not_already_cancelled(dat, |_| {
@@ -535,7 +534,7 @@ impl StartedActivityCancelEventRecorded {
     }
     pub(super) fn on_activity_task_failed(
         self,
-        dat: SharedState,
+        dat: &mut SharedState,
         attrs: ActivityTaskFailedEventAttributes,
     ) -> ActivityMachineTransition<Failed> {
         notify_if_not_already_cancelled(dat, |dat| {
@@ -544,23 +543,23 @@ impl StartedActivityCancelEventRecorded {
     }
     pub(super) fn on_activity_task_timed_out(
         self,
-        dat: SharedState,
+        dat: &mut SharedState,
         attrs: ActivityTaskTimedOutEventAttributes,
     ) -> ActivityMachineTransition<TimedOut> {
         notify_if_not_already_cancelled(dat, |dat| notify_lang_activity_timed_out(dat, attrs))
     }
     pub(super) fn on_activity_task_canceled(
         self,
-        dat: SharedState,
+        dat: &mut SharedState,
         attrs: ActivityTaskCanceledEventAttributes,
     ) -> ActivityMachineTransition<Canceled> {
-        notify_if_not_already_cancelled(dat, |dat| notify_lang_activity_cancelled(dat, Some(attrs)))
+        notify_if_not_already_cancelled(dat, |_| notify_lang_activity_cancelled(Some(attrs)))
     }
 }
 
 fn notify_if_not_already_cancelled<S>(
-    dat: SharedState,
-    notifier: impl FnOnce(SharedState) -> ActivityMachineTransition<S>,
+    dat: &mut SharedState,
+    notifier: impl FnOnce(&mut SharedState) -> ActivityMachineTransition<S>,
 ) -> ActivityMachineTransition<S>
 where
     S: Into<ActivityMachineState> + Default,
@@ -639,7 +638,7 @@ pub(super) struct Canceled {}
 impl Canceled {
     pub(super) fn on_activity_task_started(
         self,
-        dat: SharedState,
+        dat: &mut SharedState,
         seq_num: i64,
     ) -> ActivityMachineTransition<Canceled> {
         // Abandoned activities might start anyway. Ignore the result.
@@ -648,14 +647,13 @@ impl Canceled {
         } else {
             TransitionResult::Err(WFMachinesError::Nondeterminism(format!(
                 "Non-Abandon cancel mode activities cannot be started after being cancelled. \
-                 Seq: {:?}",
-                seq_num
+                 Seq: {seq_num:?}"
             )))
         }
     }
     pub(super) fn on_activity_task_completed(
         self,
-        dat: SharedState,
+        dat: &mut SharedState,
         attrs: ActivityTaskCompletedEventAttributes,
     ) -> ActivityMachineTransition<Canceled> {
         // Abandoned activities might complete anyway. Ignore the result.
@@ -663,15 +661,14 @@ impl Canceled {
             TransitionResult::default()
         } else {
             TransitionResult::Err(WFMachinesError::Nondeterminism(format!(
-                "Non-Abandon cancel mode activities cannot be completed after being cancelled: {:?}",
-                attrs
+                "Non-Abandon cancel mode activities cannot be completed after being cancelled: {attrs:?}"
             )))
         }
     }
 }
 
 fn create_request_cancel_activity_task_command<S>(
-    dat: SharedState,
+    dat: &mut SharedState,
     next_state: S,
 ) -> ActivityMachineTransition<S>
 where
@@ -696,81 +693,102 @@ where
 /// Notifies lang side that activity has timed out by sending a failure with timeout error as a cause.
 /// State machine will transition into the TimedOut state.
 fn notify_lang_activity_timed_out(
-    dat: SharedState,
+    dat: &SharedState,
     attrs: ActivityTaskTimedOutEventAttributes,
 ) -> TransitionResult<ActivityMachine, TimedOut> {
-    ActivityMachineTransition::ok_shared(
-        vec![ActivityMachineCommand::Fail(new_timeout_failure(
-            &dat, attrs,
-        ))],
-        TimedOut::default(),
-        dat,
-    )
+    ActivityMachineTransition::commands(vec![ActivityMachineCommand::Fail(new_timeout_failure(
+        dat, attrs,
+    ))])
 }
 
-/// Notifies lang side that activity has been cancelled by sending a failure with cancelled failure
-/// as a cause. Optional cancelled_event, if passed, is used to supply event IDs. State machine will
-/// transition into the `next_state` provided as a parameter.
+/// Returns the transition to indicate activity has been cancelled by sending a failure with
+/// cancelled failure as a cause. Optional cancelled_event, if passed, is used to supply event IDs.
 fn notify_lang_activity_cancelled(
-    dat: SharedState,
     canceled_event: Option<ActivityTaskCanceledEventAttributes>,
 ) -> ActivityMachineTransition<Canceled> {
-    ActivityMachineTransition::ok_shared(
-        vec![ActivityMachineCommand::Cancel(
-            canceled_event.and_then(|e| e.details),
-        )],
-        Canceled::default(),
-        dat,
-    )
+    ActivityMachineTransition::commands(vec![ActivityMachineCommand::Cancel(canceled_event)])
 }
 
-fn new_failure(dat: SharedState, attrs: ActivityTaskFailedEventAttributes) -> Failure {
+fn new_failure(dat: &SharedState, attrs: ActivityTaskFailedEventAttributes) -> Failure {
+    let rs = attrs.retry_state();
     Failure {
         message: "Activity task failed".to_owned(),
         cause: attrs.failure.map(Box::new),
-        failure_info: Some(FailureInfo::ActivityFailureInfo(
-            failure::ActivityFailureInfo {
-                identity: attrs.identity,
-                activity_type: Some(ActivityType {
-                    name: dat.attrs.activity_type,
-                }),
-                activity_id: dat.attrs.activity_id,
-                retry_state: attrs.retry_state,
-                started_event_id: attrs.started_event_id,
-                scheduled_event_id: attrs.scheduled_event_id,
-            },
+        failure_info: Some(activity_fail_info(
+            dat.attrs.activity_type.clone(),
+            dat.attrs.activity_id.clone(),
+            Some(attrs.identity),
+            rs,
+            attrs.started_event_id,
+            attrs.scheduled_event_id,
         )),
         ..Default::default()
     }
 }
 
 fn new_timeout_failure(dat: &SharedState, attrs: ActivityTaskTimedOutEventAttributes) -> Failure {
-    let failure_info = ActivityFailureInfo {
-        activity_id: dat.attrs.activity_id.to_string(),
-        activity_type: Some(ActivityType {
-            name: dat.attrs.activity_type.to_string(),
-        }),
-        scheduled_event_id: attrs.scheduled_event_id,
-        started_event_id: attrs.started_event_id,
-        retry_state: attrs.retry_state,
-        ..Default::default()
-    };
+    let rs = attrs.retry_state();
     Failure {
         message: "Activity task timed out".to_string(),
         cause: attrs.failure.map(Box::new),
-        failure_info: Some(FailureInfo::ActivityFailureInfo(failure_info)),
+        failure_info: Some(activity_fail_info(
+            dat.attrs.activity_type.clone(),
+            dat.attrs.activity_id.clone(),
+            None,
+            rs,
+            attrs.started_event_id,
+            attrs.scheduled_event_id,
+        )),
         ..Default::default()
     }
 }
 
+fn new_cancel_failure(dat: &SharedState, attrs: ActivityTaskCanceledEventAttributes) -> Failure {
+    Failure {
+        message: "Activity cancelled".to_string(),
+        cause: Some(Box::from(Failure {
+            failure_info: Some(FailureInfo::CanceledFailureInfo(CanceledFailureInfo {
+                details: attrs.details.map(Into::into),
+            })),
+            ..Default::default()
+        })),
+        failure_info: Some(activity_fail_info(
+            dat.attrs.activity_type.clone(),
+            dat.attrs.activity_id.clone(),
+            Some(attrs.identity),
+            RetryState::CancelRequested,
+            attrs.started_event_id,
+            attrs.scheduled_event_id,
+        )),
+        ..Default::default()
+    }
+}
+
+pub fn activity_fail_info(
+    act_type: String,
+    act_id: String,
+    identity: Option<String>,
+    retry_state: RetryState,
+    started_event_id: i64,
+    scheduled_event_id: i64,
+) -> FailureInfo {
+    FailureInfo::ActivityFailureInfo(ActivityFailureInfo {
+        identity: identity.unwrap_or_default(),
+        activity_type: Some(ActivityType { name: act_type }),
+        activity_id: act_id,
+        retry_state: retry_state as i32,
+        started_event_id,
+        scheduled_event_id,
+    })
+}
+
 fn convert_payloads(
     event_info: Option<EventInfo>,
     result: Option<Payloads>,
 ) -> Result<Option<Payload>, WFMachinesError> {
     result.map(TryInto::try_into).transpose().map_err(|pe| {
         WFMachinesError::Fatal(format!(
-            "Not exactly one payload in activity result ({}) for event: {:?}",
-            pe, event_info
+            "Not exactly one payload in activity result ({pe}) for event: {event_info:?}"
         ))
     })
 }
@@ -779,15 +797,17 @@ fn convert_payloads(
 mod test {
     use super::*;
     use crate::{
-        replay::TestHistoryBuilder, test_help::canned_histories, worker::workflow::ManagedWFFunc,
+        internal_flags::InternalFlags, replay::TestHistoryBuilder, test_help::canned_histories,
+        worker::workflow::ManagedWFFunc,
     };
     use rstest::{fixture, rstest};
-    use std::mem::discriminant;
+    use std::{cell::RefCell, mem::discriminant, rc::Rc};
     use temporal_sdk::{
         ActivityOptions, CancellableFuture, WfContext, WorkflowFunction, WorkflowResult,
     };
-    use temporal_sdk_core_protos::coresdk::workflow_activation::{
-        workflow_activation_job, WorkflowActivationJob,
+    use temporal_sdk_core_protos::{
+        coresdk::workflow_activation::{workflow_activation_job, WorkflowActivationJob},
+        DEFAULT_ACTIVITY_TYPE,
     };
 
     #[fixture]
@@ -807,7 +827,13 @@ mod test {
     }
 
     async fn activity_wf(command_sink: WfContext) -> WorkflowResult<()> {
-        command_sink.activity(ActivityOptions::default()).await;
+        command_sink
+            .activity(ActivityOptions {
+                activity_id: Some("activity-id-1".to_string()),
+                activity_type: DEFAULT_ACTIVITY_TYPE.to_string(),
+                ..Default::default()
+            })
+            .await;
         Ok(().into())
     }
 
@@ -895,13 +921,20 @@ mod test {
             TimedOut {}.into(),
             Completed {}.into(),
         ] {
-            let mut s = ActivityMachine {
-                state: state.clone(),
-                shared_state: Default::default(),
-            };
+            let mut s = ActivityMachine::from_parts(
+                state.clone(),
+                SharedState {
+                    scheduled_event_id: 0,
+                    started_event_id: 0,
+                    attrs: Default::default(),
+                    cancellation_type: Default::default(),
+                    cancelled_before_sent: false,
+                    internal_flags: Rc::new(RefCell::new(InternalFlags::new(&Default::default()))),
+                },
+            );
             let cmds = s.cancel().unwrap();
             assert_eq!(cmds.len(), 0);
-            assert_eq!(discriminant(&state), discriminant(&s.state));
+            assert_eq!(discriminant(&state), discriminant(s.state()));
         }
     }
 }