@temporalio/core-bridge 1.8.5 → 1.9.0-rc.0

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

@@ -13,8 +13,12 @@ use super::{
     TemporalStateMachine,
 };
 use crate::{
+    abstractions::dbg_panic,
     internal_flags::InternalFlags,
-    protosext::{HistoryEventExt, ValidScheduleLA},
+    protosext::{
+        protocol_messages::{IncomingProtocolMessage, IncomingProtocolMessageBody},
+        CompleteLocalActivityData, HistoryEventExt, ValidScheduleLA,
+    },
     telemetry::{metrics::MetricsContext, VecDisplayer},
     worker::{
         workflow::{
@@ -23,13 +27,12 @@ use crate::{
                 activity_state_machine::ActivityMachine,
                 child_workflow_state_machine::ChildWorkflowMachine,
                 modify_workflow_properties_state_machine::modify_workflow_properties,
-                patch_state_machine::VERSION_SEARCH_ATTR_KEY,
+                patch_state_machine::VERSION_SEARCH_ATTR_KEY, update_state_machine::UpdateMachine,
                 upsert_search_attributes_state_machine::upsert_search_attrs_internal,
                 HistEventData,
             },
             CommandID, DrivenWorkflow, HistoryUpdate, InternalFlagsRef, LocalResolution,
-            OutgoingJob, RunBasics, WFCommand, WFMachinesError, WorkflowFetcher,
-            WorkflowStartedInfo,
+            OutgoingJob, RunBasics, WFCommand, WFMachinesError, WorkflowStartedInfo,
         },
         ExecutingLAId, LocalActRequest, LocalActivityExecutionResult, LocalActivityResolution,
     },
@@ -37,11 +40,11 @@ use crate::{
 use siphasher::sip::SipHasher13;
 use slotmap::{SlotMap, SparseSecondaryMap};
 use std::{
-    borrow::{Borrow, BorrowMut},
     cell::RefCell,
     collections::{HashMap, VecDeque},
     convert::TryInto,
     hash::{Hash, Hasher},
+    iter::Peekable,
     rc::Rc,
     time::{Duration, Instant, SystemTime},
 };
@@ -59,7 +62,8 @@ use temporal_sdk_core_protos::{
     temporal::api::{
         command::v1::{command::Attributes as ProtoCmdAttrs, Command as ProtoCommand},
         enums::v1::EventType,
-        history::v1::{history_event, history_event::Attributes, HistoryEvent},
+        history::v1::{history_event, HistoryEvent},
+        protocol::v1::{message::SequencingId, Message as ProtocolMessage},
         sdk::v1::WorkflowTaskCompletedMetadata,
     },
 };
@@ -75,6 +79,8 @@ pub(crate) struct WorkflowMachines {
     /// kept because the lang side polls & completes for every workflow task, but we do not need
     /// to poll the server that often during replay.
     last_history_from_server: HistoryUpdate,
+    /// Protocol messages that have yet to be processed for the current WFT.
+    protocol_msgs: Vec<IncomingProtocolMessage>,
     /// EventId of the last handled WorkflowTaskStarted event
     current_started_event_id: i64,
     /// The event id of the next workflow task started event that the machines need to process.
@@ -88,7 +94,7 @@ pub(crate) struct WorkflowMachines {
     /// True if the workflow is replaying from history
     pub replaying: bool,
     /// Namespace this workflow exists in
-    pub namespace: String,
+    namespace: String,
     /// Workflow identifier
     pub workflow_id: String,
     /// Workflow type identifier. (Function name, class, etc)
@@ -96,7 +102,7 @@ pub(crate) struct WorkflowMachines {
     /// Identifies the current run
     pub run_id: String,
     /// The task queue this workflow is operating within
-    pub task_queue: String,
+    task_queue: String,
     /// Is set to true once we've seen the final event in workflow history, to avoid accidentally
     /// re-applying the final workflow task.
     pub have_seen_terminal_event: bool,
@@ -117,16 +123,19 @@ pub(crate) struct WorkflowMachines {
     history_size_bytes: u64,
     /// Set on each WFT started event
     continue_as_new_suggested: bool,
+    /// Set if the current WFT is already complete and that completion event had a build id in it.
+    current_wft_build_id: Option<String>,
 
     all_machines: SlotMap<MachineKey, Machines>,
     /// If a machine key is in this map, that machine was created internally by core, not as a
     /// command from lang.
     machine_is_core_created: SparseSecondaryMap<MachineKey, ()>,
-
     /// A mapping for accessing machines associated to a particular event, where the key is the id
     /// of the initiating event for that machine.
     machines_by_event_id: HashMap<i64, MachineKey>,
-
+    /// A mapping for accessing machines that were created as a result of protocol messages. The
+    /// key is the protocol's instance id.
+    machines_by_protocol_instance_id: HashMap<String, MachineKey>,
     /// Maps command ids as created by workflow authors to their associated machines.
     id_to_machine: HashMap<CommandID, MachineKey>,
 
@@ -136,6 +145,9 @@ pub(crate) struct WorkflowMachines {
     /// Commands generated by the currently processing workflow task, which will eventually be
     /// transferred to `commands` (and hence eventually sent to the server)
     current_wf_task_commands: VecDeque<CommandAndMachine>,
+    /// Messages generated while processing the current workflow task, which will be sent in the
+    /// next WFT response. Keyed by message id.
+    message_outbox: VecDeque<ProtocolMessage>,
 
     /// Information about patch markers we have already seen while replaying history
     encountered_patch_markers: HashMap<String, ChangeInfo>,
@@ -180,6 +192,9 @@ pub(super) enum MachineResponse {
     /// Pushes a new command into the list that will be sent to server once we respond with the
     /// workflow task completion
     IssueNewCommand(ProtoCommand),
+    /// Pushes a new protocol message into the list that will be sent to server once we respond with
+    /// the workflow task completion
+    IssueNewMessage(ProtocolMessage),
     /// The machine requests the creation of another *different* machine. This acts as if lang
     /// had replied to the activation with a command, but we use a special set of IDs to avoid
     /// collisions.
@@ -230,6 +245,7 @@ impl WorkflowMachines {
         };
         Self {
             last_history_from_server: basics.history,
+            protocol_msgs: vec![],
             namespace: basics.namespace,
             workflow_id: basics.workflow_id,
             workflow_type: basics.workflow_type,
@@ -249,12 +265,15 @@ impl WorkflowMachines {
             observed_internal_flags: Rc::new(RefCell::new(observed_internal_flags)),
             history_size_bytes: 0,
             continue_as_new_suggested: false,
+            current_wft_build_id: None,
             all_machines: Default::default(),
             machine_is_core_created: Default::default(),
             machines_by_event_id: Default::default(),
+            machines_by_protocol_instance_id: Default::default(),
             id_to_machine: Default::default(),
             commands: Default::default(),
             current_wf_task_commands: Default::default(),
+            message_outbox: Default::default(),
             encountered_patch_markers: Default::default(),
             local_activity_data: LocalActivityData::default(),
             have_seen_terminal_event: false,
@@ -274,6 +293,20 @@ impl WorkflowMachines {
             .and_then(|(st, et)| et.duration_since(st).ok())
     }
 
+    /// Must be called every time a new WFT is received
+    pub(crate) fn new_work_from_server(
+        &mut self,
+        update: HistoryUpdate,
+        protocol_messages: Vec<IncomingProtocolMessage>,
+    ) -> Result<()> {
+        if !self.protocol_msgs.is_empty() {
+            dbg_panic!("There are unprocessed protocol messages while receiving new work");
+        }
+        self.protocol_msgs = protocol_messages;
+        self.new_history_from_server(update)?;
+        Ok(())
+    }
+
     pub(crate) fn new_history_from_server(&mut self, update: HistoryUpdate) -> Result<()> {
         self.last_history_from_server = update;
         self.replaying = self.last_history_from_server.previous_wft_started_id > 0;
@@ -335,28 +368,40 @@ impl WorkflowMachines {
         self.drive_me.get_started_info()
     }
 
+    pub(crate) fn get_last_wft_started_id(&self) -> i64 {
+        self.current_started_event_id
+    }
+
+    pub(crate) fn prepare_for_wft_response(&mut self) -> MachinesWFTResponseContent {
+        MachinesWFTResponseContent {
+            replaying: self.replaying,
+            has_pending_jobs: self.has_pending_jobs(),
+            have_seen_terminal_event: self.have_seen_terminal_event,
+            have_pending_la_resolutions: self.has_pending_la_resolutions(),
+            last_processed_event: self.last_processed_event,
+            me: self,
+        }
+    }
+
     /// Fetches commands which are ready for processing from the state machines, generally to be
     /// sent off to the server. They are not removed from the internal queue, that happens when
     /// corresponding history events from the server are being handled.
-    pub(crate) fn get_commands(&self) -> Vec<ProtoCommand> {
+    pub(crate) fn get_commands(&self) -> impl Iterator<Item = ProtoCommand> + '_ {
         // Since we're about to write a WFT, record any internal flags we know about which aren't
         // already recorded.
         (*self.observed_internal_flags)
             .borrow_mut()
             .write_all_known();
-        self.commands
-            .iter()
-            .filter_map(|c| {
-                if !self.machine(c.machine).is_final_state() {
-                    match &c.command {
-                        MachineAssociatedCommand::Real(cmd) => Some((**cmd).clone()),
-                        MachineAssociatedCommand::FakeLocalActivityMarker(_) => None,
-                    }
-                } else {
-                    None
+        self.commands.iter().filter_map(|c| {
+            if !self.machine(c.machine).is_final_state() {
+                match &c.command {
+                    MachineAssociatedCommand::Real(cmd) => Some((**cmd).clone()),
+                    MachineAssociatedCommand::FakeLocalActivityMarker(_) => None,
                 }
-            })
-            .collect()
+            } else {
+                None
+            }
+        })
     }
 
     /// Returns the next activation that needs to be performed by the lang sdk. Things like unblock
@@ -387,6 +432,7 @@ impl WorkflowMachines {
                 .collect(),
             history_size_bytes: self.history_size_bytes,
             continue_as_new_suggested: self.continue_as_new_suggested,
+            build_id_for_current_task: self.current_wft_build_id.clone().unwrap_or_default(),
         }
     }
 
@@ -413,6 +459,30 @@ impl WorkflowMachines {
             .add_lang_used(flags);
     }
 
+    /// Undo a speculative workflow task by resetting to a certain WFT Started ID. This can happen
+    /// when an update request is rejected.
+    pub(crate) fn reset_last_started_id(&mut self, id: i64) {
+        debug!("Resetting back to event id {} due to speculative WFT", id);
+        self.current_started_event_id = id;
+        // This is pretty nasty to just + 1 like this, but, we know WFT complete always follows
+        // WFT started, which the id given to us to reset to must be, and we need to avoid
+        // re-applying the WFT Completed event, so we make sure to consider that processed
+        self.last_processed_event = id + 1;
+        // Then, we have to drop any state machines (which should only be one workflow task machine)
+        // we may have created when servicing the speculative task.
+        // Remove when https://github.com/rust-lang/rust/issues/59618 is stable
+        let remove_these: Vec<_> = self
+            .machines_by_event_id
+            .iter()
+            .filter(|(mid, _)| **mid > id)
+            .map(|(mid, mkey)| (*mid, *mkey))
+            .collect();
+        for (mid, mkey) in remove_these {
+            self.machines_by_event_id.remove(&mid);
+            self.all_machines.remove(mkey);
+        }
+    }
+
     /// Iterate the state machines, which consists of grabbing any pending outgoing commands from
     /// the workflow code, handling them, and preparing them to be sent off to the server.
     pub(crate) fn iterate_machines(&mut self) -> Result<()> {
@@ -446,22 +516,45 @@ impl WorkflowMachines {
             return Ok(0);
         }
 
-        fn update_internal_flags(me: &mut WorkflowMachines) {
-            // Update observed patches with any that were used in the task
-            if let Some(next_complete) = me
-                .last_history_from_server
-                .peek_next_wft_completed(me.last_processed_event)
-            {
-                (*me.observed_internal_flags)
-                    .borrow_mut()
-                    .add_from_complete(next_complete);
-            }
+        fn get_processable_messages(
+            me: &mut WorkflowMachines,
+            for_event_id: i64,
+        ) -> Vec<IncomingProtocolMessage> {
+            // Another thing to replace when `drain_filter` exists
+            let mut ret = vec![];
+            me.protocol_msgs = std::mem::take(&mut me.protocol_msgs)
+                .into_iter()
+                .filter_map(|x| {
+                    if x.processable_after_event_id()
+                        .is_some_and(|eid| eid <= for_event_id)
+                    {
+                        ret.push(x);
+                        None
+                    } else {
+                        Some(x)
+                    }
+                })
+                .collect();
+            ret
         }
 
-        // We update the internal flags before applying the current task (peeking to the completion
-        // of this task), and also at the end (peeking to the completion of the task that lang is
-        // about to generate commands for, and for which we will want those flags active).
-        update_internal_flags(self);
+        // Peek to the next WFT complete and update ourselves with data we might need in it.
+        if let Some(next_complete) = self
+            .last_history_from_server
+            .peek_next_wft_completed(self.last_processed_event)
+        {
+            // We update the internal flags before applying the current task
+            (*self.observed_internal_flags)
+                .borrow_mut()
+                .add_from_complete(next_complete);
+            // Save this tasks' Build ID if it had one
+            if let Some(bid) = next_complete.worker_version.as_ref().map(|wv| &wv.build_id) {
+                self.current_wft_build_id = Some(bid.to_string());
+            } else {
+                // Otherwise we do not want to keep anything previously stored
+                self.current_wft_build_id = None;
+            }
+        }
 
         let last_handled_wft_started_id = self.current_started_event_id;
         let (events, has_final_event) = match self
@@ -499,14 +592,14 @@ impl WorkflowMachines {
         let mut do_handle_event = true;
         let mut history = events.into_iter().peekable();
         while let Some(event) = history.next() {
-            if event.event_id != self.last_processed_event + 1 {
+            let eid = event.event_id;
+            if eid != self.last_processed_event + 1 {
                 return Err(WFMachinesError::Fatal(format!(
                     "History is out of order. Last processed event: {}, event id: {}",
-                    self.last_processed_event, event.event_id
+                    self.last_processed_event, eid
                 )));
             }
             let next_event = history.peek();
-            let eid = event.event_id;
 
             // This definition of replaying here is that we are no longer replaying as soon as we
             // see new events that have never been seen or produced by the SDK.
@@ -520,7 +613,7 @@ impl WorkflowMachines {
             // them.
             if self.replaying
                 && has_final_event
-                && event.event_id > self.last_history_from_server.previous_wft_started_id
+                && eid > self.last_history_from_server.previous_wft_started_id
                 && event.event_type() != EventType::WorkflowTaskCompleted
                 && !event.is_command_event()
             {
@@ -528,6 +621,12 @@ impl WorkflowMachines {
                 self.replaying = false;
             }
 
+            // Process any messages that should be processed before the event we're about to handle
+            let processable_msgs = get_processable_messages(self, eid - 1);
+            for msg in processable_msgs {
+                self.handle_protocol_message(msg)?;
+            }
+
             if do_handle_event {
                 let eho = self.handle_event(
                     HistEventData {
@@ -551,9 +650,14 @@ impl WorkflowMachines {
             self.last_processed_event = eid;
         }
 
+        // Needed to delay mutation of self until after we've iterated over peeked events.
+        enum DelayedAction {
+            WakeLa(MachineKey, CompleteLocalActivityData),
+            ProtocolMessage(IncomingProtocolMessage),
+        }
+        let mut delayed_actions = vec![];
         // Scan through to the next WFT, searching for any patch / la markers, so that we can
         // pre-resolve them.
-        let mut wake_las = vec![];
         for e in self
             .last_history_from_server
             .peek_next_wft_sequence(last_handled_wft_started_id)
@@ -575,7 +679,7 @@ impl WorkflowMachines {
                     if let Ok(mk) =
                         self.get_machine_key(CommandID::LocalActivity(la_dat.marker_dat.seq))
                     {
-                        wake_las.push((mk, la_dat));
+                        delayed_actions.push(DelayedAction::WakeLa(mk, la_dat));
                     } else {
                         self.local_activity_data.insert_peeked_marker(la_dat);
                     }
@@ -584,22 +688,51 @@ impl WorkflowMachines {
                         "Local activity marker was unparsable: {e:?}"
                     )));
                 }
+            } else if let Some(
+                history_event::Attributes::WorkflowExecutionUpdateAcceptedEventAttributes(ref atts),
+            ) = e.attributes
+            {
+                // If we see a workflow update accepted event, initialize the machine for it by
+                // pretending we received the message we would've under not-replay.
+                delayed_actions.push(DelayedAction::ProtocolMessage(IncomingProtocolMessage {
+                    id: atts.accepted_request_message_id.clone(),
+                    protocol_instance_id: atts.protocol_instance_id.clone(),
+                    sequencing_id: Some(SequencingId::EventId(
+                        atts.accepted_request_sequencing_event_id,
+                    )),
+                    body: IncomingProtocolMessageBody::UpdateRequest(
+                        atts.accepted_request
+                            .clone()
+                            .ok_or_else(|| {
+                                WFMachinesError::Fatal(
+                                    "Update accepted event must contain accepted request"
+                                        .to_string(),
+                                )
+                            })?
+                            .try_into()?,
+                    ),
+                }));
             }
         }
-        for (mk, la_dat) in wake_las {
-            let mach = self.machine_mut(mk);
-            if let Machines::LocalActivityMachine(ref mut lam) = *mach {
-                if lam.will_accept_resolve_marker() {
-                    let resps = lam.try_resolve_with_dat(la_dat.into())?;
-                    self.process_machine_responses(mk, resps)?;
-                } else {
-                    self.local_activity_data.insert_peeked_marker(la_dat);
+        for action in delayed_actions {
+            match action {
+                DelayedAction::WakeLa(mk, la_dat) => {
+                    let mach = self.machine_mut(mk);
+                    if let Machines::LocalActivityMachine(ref mut lam) = *mach {
+                        if lam.will_accept_resolve_marker() {
+                            let resps = lam.try_resolve_with_dat(la_dat.into())?;
+                            self.process_machine_responses(mk, resps)?;
+                        } else {
+                            self.local_activity_data.insert_peeked_marker(la_dat);
+                        }
+                    }
+                }
+                DelayedAction::ProtocolMessage(pm) => {
+                    self.handle_protocol_message(pm)?;
                 }
             }
         }
 
-        update_internal_flags(self);
-
         if !self.replaying {
             self.metrics.wf_task_replay_latency(replay_start.elapsed());
         }
@@ -665,27 +798,19 @@ impl WorkflowMachines {
         }
 
         if let Some(initial_cmd_id) = event.get_initial_command_event_id() {
-            // We remove the machine while we it handles events, then return it, to avoid
-            // borrowing from ourself mutably.
-            let maybe_machine = self.machines_by_event_id.remove(&initial_cmd_id);
-            match maybe_machine {
-                Some(sm) => {
-                    self.submachine_handle_event(sm, event_dat)?;
-                    // Restore machine if not in it's final state
-                    if !self.machine(sm).is_final_state() {
-                        self.machines_by_event_id.insert(initial_cmd_id, sm);
-                    }
-                }
-                None => {
-                    return Err(WFMachinesError::Nondeterminism(format!(
+            let mkey = self
+                .machines_by_event_id
+                .get(&initial_cmd_id)
+                .ok_or_else(|| {
+                    WFMachinesError::Nondeterminism(format!(
                         "During event handling, this event had an initial command ID but we \
-                            could not find a matching command for it: {event:?}"
-                    )));
-                }
-            }
+                         could not find a matching command for it: {event:?}"
+                    ))
+                })?;
+            self.submachine_handle_event(*mkey, event_dat)?;
         } else {
             self.handle_non_stateful_event(event_dat)?;
-        }
+        };
 
         Ok(EventHandlingOutcome::Normal)
     }
@@ -769,10 +894,7 @@ impl WorkflowMachines {
     }
 
     fn handle_non_stateful_event(&mut self, event_dat: HistEventData) -> Result<()> {
-        trace!(
-            event = %event_dat.event,
-            "handling non-stateful event"
-        );
+        trace!(event = %event_dat.event, "handling non-stateful event");
         let event_id = event_dat.event.event_id;
         match EventType::from_i32(event_dat.event.event_type) {
             Some(EventType::WorkflowExecutionStarted) => {
@@ -855,6 +977,42 @@ impl WorkflowMachines {
         self.process_machine_responses(sm, machine_responses)?;
         Ok(())
     }
+    /// Handle a single protocol message delivered in a workflow task.
+    ///
+    /// This function will attempt to apply the message to a corresponding state machine for the
+    /// appropriate protocol type, creating it if it does not exist.
+    ///
+    /// On replay, protocol messages may be made up by looking ahead in history to see if there is
+    /// already an event corresponding to the result of some protocol message which would have
+    /// existed to create that result.
+    #[instrument(skip(self))]
+    fn handle_protocol_message(&mut self, message: IncomingProtocolMessage) -> Result<()> {
+        static SEQIDERR: &str = "Update request messages must contain an event sequencing id! \
+                                 This is a server bug.";
+
+        match message.body {
+            IncomingProtocolMessageBody::UpdateRequest(ur) => {
+                let seq_id = if let SequencingId::EventId(eid) = message
+                    .sequencing_id
+                    .ok_or_else(|| WFMachinesError::Fatal(SEQIDERR.to_string()))?
+                {
+                    eid
+                } else {
+                    return Err(WFMachinesError::Fatal(SEQIDERR.to_string()));
+                };
+                let um = UpdateMachine::init(
+                    message.id,
+                    message.protocol_instance_id.clone(),
+                    seq_id,
+                    ur,
+                    self.replaying,
+                );
+                let mk = self.add_new_protocol_machine(um.machine, message.protocol_instance_id);
+                self.process_machine_responses(mk, vec![um.response])?;
+            }
+        }
+        Ok(())
+    }
 
     /// Transfer commands from `current_wf_task_commands` to `commands`, so they may be sent off
     /// to the server. While doing so, [TemporalStateMachine::handle_command] is called on the
@@ -940,6 +1098,13 @@ impl WorkflowMachines {
                         machine: smk,
                     })
                 }
+                MachineResponse::IssueNewMessage(pm) => {
+                    // Messages shouldn't be sent back when replaying. This is true for update,
+                    // currently the only user of protocol messages. May eventually change.
+                    if !self.replaying {
+                        self.message_outbox.push_back(pm);
+                    }
+                }
                 MachineResponse::NewCoreOriginatedCommand(attrs) => match attrs {
                     ProtoCmdAttrs::RequestCancelExternalWorkflowExecutionCommandAttributes(
                         attrs,
@@ -1244,6 +1409,20 @@ impl WorkflowMachines {
                         CommandIdKind::NeverResolves,
                     );
                 }
+                WFCommand::UpdateResponse(ur) => {
+                    let m_key = self.get_machine_by_msg(&ur.protocol_instance_id)?;
+                    let mach = self.machine_mut(m_key);
+                    if let Machines::UpdateMachine(m) = mach {
+                        let resps = m.handle_response(ur)?;
+                        self.process_machine_responses(m_key, resps)?;
+                    } else {
+                        return Err(WFMachinesError::Nondeterminism(format!(
+                            "Tried to handle an update response for \
+                             update with instance id {} but it was not found!",
+                            &ur.protocol_instance_id
+                        )));
+                    }
+                }
                 WFCommand::NoCommandsFromLang => (),
             }
         }
@@ -1267,6 +1446,17 @@ impl WorkflowMachines {
         })?)
     }
 
+    fn get_machine_by_msg(&self, protocol_instance_id: &str) -> Result<MachineKey> {
+        Ok(*self
+            .machines_by_protocol_instance_id
+            .get(protocol_instance_id)
+            .ok_or_else(|| {
+                WFMachinesError::Fatal(format!(
+                    "Missing associated machine for protocol message {protocol_instance_id}"
+                ))
+            })?)
+    }
+
     fn add_terminal_command(&mut self, machine: NewMachineWithCommand) {
         let cwfm = self.add_new_command_machine(machine);
         self.workflow_end_time = Some(SystemTime::now());
@@ -1299,18 +1489,18 @@ impl WorkflowMachines {
         }
     }
 
+    fn add_new_protocol_machine(&mut self, machine: Machines, instance_id: String) -> MachineKey {
+        let k = self.all_machines.insert(machine);
+        self.machines_by_protocol_instance_id.insert(instance_id, k);
+        k
+    }
+
     fn machine(&self, m: MachineKey) -> &Machines {
-        self.all_machines
-            .get(m)
-            .expect("Machine must exist")
-            .borrow()
+        self.all_machines.get(m).expect("Machine must exist")
     }
 
     fn machine_mut(&mut self, m: MachineKey) -> &mut Machines {
-        self.all_machines
-            .get_mut(m)
-            .expect("Machine must exist")
-            .borrow_mut()
+        self.all_machines.get_mut(m).expect("Machine must exist")
     }
 
     fn augment_continue_as_new_with_current_values(
@@ -1355,6 +1545,32 @@ impl WorkflowMachines {
     }
 }
 
+/// Contains everything workflow machine internals need to bubble up when we're getting ready to
+/// respond with a WFT completion. Allows for lazy mutation of the machine, since mutation is not
+/// desired unless we are actually going to respond to the WFT, which may not always happen.
+pub struct MachinesWFTResponseContent<'a> {
+    me: &'a mut WorkflowMachines,
+    pub replaying: bool,
+    pub has_pending_jobs: bool,
+    pub have_seen_terminal_event: bool,
+    pub have_pending_la_resolutions: bool,
+    pub last_processed_event: i64,
+}
+impl<'a> MachinesWFTResponseContent<'a> {
+    pub fn commands(&self) -> Peekable<impl Iterator<Item = ProtoCommand> + '_> {
+        self.me.get_commands().peekable()
+    }
+    pub fn has_messages(&self) -> bool {
+        !self.me.message_outbox.is_empty()
+    }
+    pub fn messages(&mut self) -> Vec<ProtocolMessage> {
+        self.me.message_outbox.drain(..).collect()
+    }
+    pub fn metadata_for_complete(&mut self) -> WorkflowTaskCompletedMetadata {
+        self.me.get_metadata_for_wft_complete()
+    }
+}
+
 fn str_to_randomness_seed(run_id: &str) -> u64 {
     // This was originally `DefaultHasher` but that is potentially unstable across Rust releases.
     // This must forever be `SipHasher13` now or we risk breaking history compat.
@@ -1385,8 +1601,9 @@ fn patch_marker_handling(
     fn skip_one_or_two_events(next_event: Option<&HistoryEvent>) -> Result<EventHandlingOutcome> {
         // Also ignore the subsequent upsert event if present
         let mut skip_next_event = false;
-        if let Some(Attributes::UpsertWorkflowSearchAttributesEventAttributes(atts)) =
-            next_event.and_then(|ne| ne.attributes.as_ref())
+        if let Some(history_event::Attributes::UpsertWorkflowSearchAttributesEventAttributes(
+            atts,
+        )) = next_event.and_then(|ne| ne.attributes.as_ref())
         {
             if let Some(ref sa) = atts.search_attributes {
                 skip_next_event = sa.indexed_fields.contains_key(VERSION_SEARCH_ATTR_KEY);