@temporalio/core-bridge 1.5.2 → 1.7.0

package/sdk-core/core/src/worker/workflow/history_update.rs

@@ -1,57 +1,99 @@
 use crate::{
-    replay::{HistoryInfo, TestHistoryBuilder},
-    worker::client::WorkerClient,
+    protosext::ValidPollWFTQResponse,
+    worker::{
+        client::WorkerClient,
+        workflow::{CacheMissFetchReq, PermittedWFT, PreparedWFT},
+    },
 };
-use futures::{future::BoxFuture, stream, stream::BoxStream, FutureExt, Stream, StreamExt};
+use futures::{future::BoxFuture, FutureExt, Stream};
+use itertools::Itertools;
 use std::{
     collections::VecDeque,
     fmt::Debug,
     future::Future,
+    mem,
+    mem::transmute,
     pin::Pin,
     sync::Arc,
     task::{Context, Poll},
 };
 use temporal_sdk_core_protos::temporal::api::{
     enums::v1::EventType,
-    history::v1::{History, HistoryEvent},
-    workflowservice::v1::GetWorkflowExecutionHistoryResponse,
+    history::v1::{history_event, History, HistoryEvent, WorkflowTaskCompletedEventAttributes},
 };
 use tracing::Instrument;
 
-/// A slimmed down version of a poll workflow task response which includes just the info needed
-/// by [WorkflowManager]. History events are expected to be consumed from it and applied to the
-/// state machines.
+lazy_static::lazy_static! {
+    static ref EMPTY_FETCH_ERR: tonic::Status
+        = tonic::Status::unknown("Fetched empty history page");
+    static ref EMPTY_TASK_ERR: tonic::Status
+        = tonic::Status::unknown("Received an empty workflow task with no queries or history");
+}
+
+/// Represents one or more complete WFT sequences. History events are expected to be consumed from
+/// it and applied to the state machines via [HistoryUpdate::take_next_wft_sequence]
+#[cfg_attr(
+    feature = "save_wf_inputs",
+    derive(serde::Serialize, serde::Deserialize)
+)]
 pub struct HistoryUpdate {
-    events: BoxStream<'static, Result<HistoryEvent, tonic::Status>>,
-    /// It is useful to be able to look ahead up to one workflow task beyond the currently
-    /// requested one. The initial (possibly only) motivation for this being to be able to
-    /// pre-emptively notify lang about patch markers so that calls to `changed` do not need to
-    /// be async.
-    buffered: VecDeque<HistoryEvent>,
-    pub previous_started_event_id: i64,
+    events: Vec<HistoryEvent>,
+    /// The event ID of the last started WFT, as according to the WFT which this update was
+    /// extracted from. Hence, while processing multiple logical WFTs during replay which were part
+    /// of one large history fetched from server, multiple updates may have the same value here.
+    pub previous_wft_started_id: i64,
+    /// The `started_event_id` field from the WFT which this update is tied to. Multiple updates
+    /// may have the same value if they're associated with the same WFT.
+    pub wft_started_id: i64,
+    /// True if this update contains the final WFT in history, and no more attempts to extract
+    /// additional updates should be made.
+    has_last_wft: bool,
 }
 impl Debug for HistoryUpdate {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(
-            f,
-            "HistoryUpdate(previous_started_event_id: {})",
-            self.previous_started_event_id
-        )
+        if self.is_real() {
+            write!(
+                f,
+                "HistoryUpdate(previous_started_event_id: {}, length: {}, first_event_id: {:?})",
+                self.previous_wft_started_id,
+                self.events.len(),
+                self.events.first().map(|e| e.event_id)
+            )
+        } else {
+            write!(f, "DummyHistoryUpdate")
+        }
     }
 }
 
+#[derive(Debug)]
+pub enum NextWFT {
+    ReplayOver,
+    WFT(Vec<HistoryEvent>, bool),
+    NeedFetch,
+}
+
+#[derive(derive_more::DebugCustom)]
+#[debug(fmt = "HistoryPaginator(run_id: {run_id})")]
+#[cfg_attr(
+    feature = "save_wf_inputs",
+    derive(serde::Serialize, serde::Deserialize),
+    serde(default = "HistoryPaginator::fake_deserialized")
+)]
 pub struct HistoryPaginator {
-    // Potentially this could actually be a ref w/ lifetime here
+    pub(crate) wf_id: String,
+    pub(crate) run_id: String,
+    pub(crate) previous_wft_started_id: i64,
+    pub(crate) wft_started_event_id: i64,
+
+    #[cfg_attr(feature = "save_wf_inputs", serde(skip))]
     client: Arc<dyn WorkerClient>,
+    #[cfg_attr(feature = "save_wf_inputs", serde(skip))]
     event_queue: VecDeque<HistoryEvent>,
-    wf_id: String,
-    run_id: String,
+    #[cfg_attr(feature = "save_wf_inputs", serde(skip))]
     next_page_token: NextPageToken,
-    open_history_request:
-        Option<BoxFuture<'static, Result<GetWorkflowExecutionHistoryResponse, tonic::Status>>>,
     /// These are events that should be returned once pagination has finished. This only happens
     /// during cache misses, where we got a partial task but need to fetch history from the start.
-    /// We use this to apply any
+    #[cfg_attr(feature = "save_wf_inputs", serde(skip))]
     final_events: Vec<HistoryEvent>,
 }
 
@@ -77,8 +119,77 @@ impl From<Vec<u8>> for NextPageToken {
 }
 
 impl HistoryPaginator {
-    pub(crate) fn new(
+    /// Use a new poll response to create a new [WFTPaginator], returning it and the
+    /// [PreparedWFT] extracted from it that can be fed into workflow state.
+    pub(super) async fn from_poll(
+        wft: ValidPollWFTQResponse,
+        client: Arc<dyn WorkerClient>,
+    ) -> Result<(Self, PreparedWFT), tonic::Status> {
+        let empty_hist = wft.history.events.is_empty();
+        let npt = if empty_hist {
+            NextPageToken::FetchFromStart
+        } else {
+            wft.next_page_token.into()
+        };
+        let mut paginator = HistoryPaginator::new(
+            wft.history,
+            wft.previous_started_event_id,
+            wft.started_event_id,
+            wft.workflow_execution.workflow_id.clone(),
+            wft.workflow_execution.run_id.clone(),
+            npt,
+            client,
+        );
+        if empty_hist && wft.legacy_query.is_none() && wft.query_requests.is_empty() {
+            return Err(EMPTY_TASK_ERR.clone());
+        }
+        let update = if empty_hist {
+            HistoryUpdate::from_events(
+                [],
+                wft.previous_started_event_id,
+                wft.started_event_id,
+                true,
+            )
+            .0
+        } else {
+            paginator.extract_next_update().await?
+        };
+        let prepared = PreparedWFT {
+            task_token: wft.task_token,
+            attempt: wft.attempt,
+            execution: wft.workflow_execution,
+            workflow_type: wft.workflow_type,
+            legacy_query: wft.legacy_query,
+            query_requests: wft.query_requests,
+            update,
+        };
+        Ok((paginator, prepared))
+    }
+
+    pub(super) async fn from_fetchreq(
+        mut req: CacheMissFetchReq,
+        client: Arc<dyn WorkerClient>,
+    ) -> Result<PermittedWFT, tonic::Status> {
+        let mut paginator = Self {
+            wf_id: req.original_wft.work.execution.workflow_id.clone(),
+            run_id: req.original_wft.work.execution.run_id.clone(),
+            previous_wft_started_id: req.original_wft.work.update.previous_wft_started_id,
+            wft_started_event_id: req.original_wft.work.update.wft_started_id,
+            client,
+            event_queue: Default::default(),
+            next_page_token: NextPageToken::FetchFromStart,
+            final_events: vec![],
+        };
+        let first_update = paginator.extract_next_update().await?;
+        req.original_wft.work.update = first_update;
+        req.original_wft.paginator = paginator;
+        Ok(req.original_wft)
+    }
+
+    fn new(
         initial_history: History,
+        previous_wft_started_id: i64,
+        wft_started_event_id: i64,
         wf_id: String,
         run_id: String,
         next_page_token: impl Into<NextPageToken>,
@@ -97,20 +208,120 @@ impl HistoryPaginator {
             wf_id,
             run_id,
             next_page_token,
-            open_history_request: None,
             final_events,
+            previous_wft_started_id,
+            wft_started_event_id,
+        }
+    }
+
+    #[cfg(feature = "save_wf_inputs")]
+    pub(super) fn fake_deserialized() -> HistoryPaginator {
+        use crate::worker::client::mocks::mock_manual_workflow_client;
+        HistoryPaginator {
+            client: Arc::new(mock_manual_workflow_client()),
+            event_queue: Default::default(),
+            wf_id: "".to_string(),
+            run_id: "".to_string(),
+            next_page_token: NextPageToken::FetchFromStart,
+            final_events: vec![],
+            previous_wft_started_id: -2,
+            wft_started_event_id: -2,
+        }
+    }
+
+    /// Return at least the next two WFT sequences (as determined by the passed-in ID) as a
+    /// [HistoryUpdate]. Two sequences supports the required peek-ahead during replay without
+    /// unnecessary back-and-forth.
+    ///
+    /// If there are already enough events buffered in memory, they will all be returned. Including
+    /// possibly (likely, during replay) more than just the next two WFTs.
+    ///
+    /// If there are insufficient events to constitute two WFTs, then we will fetch pages until
+    /// we have two, or until we are at the end of history.
+    pub(crate) async fn extract_next_update(&mut self) -> Result<HistoryUpdate, tonic::Status> {
+        loop {
+            let no_next_page = !self.get_next_page().await?;
+            let current_events = mem::take(&mut self.event_queue);
+            let seen_enough_events = current_events
+                .back()
+                .map(|e| e.event_id)
+                .unwrap_or_default()
+                >= self.wft_started_event_id;
+            if current_events.is_empty() || (no_next_page && !seen_enough_events) {
+                // If next page fetching happened, and we still ended up with no or insufficient
+                // events, something is wrong. We're expecting there to be more events to be able to
+                // extract this update, but server isn't giving us any. We have no choice except to
+                // give up and evict.
+                error!(
+                    "We expected to be able to fetch more events but server says there are none"
+                );
+                return Err(EMPTY_FETCH_ERR.clone());
+            }
+            let first_event_id = current_events.front().unwrap().event_id;
+            // We only *really* have the last WFT if the events go all the way up to at least the
+            // WFT started event id. Otherwise we somehow still have partial history.
+            let no_more = matches!(self.next_page_token, NextPageToken::Done) && seen_enough_events;
+            let (update, extra) = HistoryUpdate::from_events(
+                current_events,
+                self.previous_wft_started_id,
+                self.wft_started_event_id,
+                no_more,
+            );
+            let extra_eid_same = extra
+                .first()
+                .map(|e| e.event_id == first_event_id)
+                .unwrap_or_default();
+            // If there are some events at the end of the fetched events which represent only a
+            // portion of a complete WFT, retain them to be used in the next extraction.
+            self.event_queue = extra.into();
+            if !no_more && extra_eid_same {
+                // There was not a meaningful WFT in the whole page. We must fetch more.
+                continue;
+            }
+            return Ok(update);
         }
     }
 
-    fn extend_queue_with_new_page(&mut self, resp: GetWorkflowExecutionHistoryResponse) {
-        self.next_page_token = resp.next_page_token.into();
+    /// Fetches the next page and adds it to the internal queue. Returns true if a fetch was
+    /// performed, false if there is no next page.
+    async fn get_next_page(&mut self) -> Result<bool, tonic::Status> {
+        let history = loop {
+            let npt = match mem::replace(&mut self.next_page_token, NextPageToken::Done) {
+                // If the last page token we got was empty, we're done.
+                NextPageToken::Done => return Ok(false),
+                NextPageToken::FetchFromStart => vec![],
+                NextPageToken::Next(v) => v,
+            };
+            debug!(run_id=%self.run_id, "Fetching new history page");
+            let fetch_res = self
+                .client
+                .get_workflow_execution_history(self.wf_id.clone(), Some(self.run_id.clone()), npt)
+                .instrument(span!(tracing::Level::TRACE, "fetch_history_in_paginator"))
+                .await?;
+
+            self.next_page_token = fetch_res.next_page_token.into();
+
+            let history_is_empty = fetch_res
+                .history
+                .as_ref()
+                .map(|h| h.events.is_empty())
+                .unwrap_or(true);
+            if history_is_empty && matches!(&self.next_page_token, NextPageToken::Next(_)) {
+                // If the fetch returned an empty history, but there *was* a next page token,
+                // immediately try to get that.
+                continue;
+            }
+            // Async doesn't love recursion so we do this instead.
+            break fetch_res.history;
+        };
+
         self.event_queue
-            .extend(resp.history.map(|h| h.events).unwrap_or_default());
+            .extend(history.map(|h| h.events).unwrap_or_default());
         if matches!(&self.next_page_token, NextPageToken::Done) {
             // If finished, we need to extend the queue with the final events, skipping any
             // which are already present.
             if let Some(last_event_id) = self.event_queue.back().map(|e| e.event_id) {
-                let final_events = std::mem::take(&mut self.final_events);
+                let final_events = mem::take(&mut self.final_events);
                 self.event_queue.extend(
                     final_events
                         .into_iter()
@@ -118,374 +329,576 @@ impl HistoryPaginator {
                 );
             }
         };
+        Ok(true)
     }
 }
 
-impl Stream for HistoryPaginator {
+#[pin_project::pin_project]
+struct StreamingHistoryPaginator {
+    inner: HistoryPaginator,
+    #[pin]
+    open_history_request: Option<BoxFuture<'static, Result<(), tonic::Status>>>,
+}
+
+impl StreamingHistoryPaginator {
+    // Kept since can be used for history downloading
+    #[cfg(test)]
+    pub fn new(inner: HistoryPaginator) -> Self {
+        Self {
+            inner,
+            open_history_request: None,
+        }
+    }
+}
+
+impl Stream for StreamingHistoryPaginator {
     type Item = Result<HistoryEvent, tonic::Status>;
 
-    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-        if let Some(e) = self.event_queue.pop_front() {
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let mut this = self.project();
+
+        if let Some(e) = this.inner.event_queue.pop_front() {
             return Poll::Ready(Some(Ok(e)));
         }
-        let history_req = if let Some(req) = self.open_history_request.as_mut() {
-            req
-        } else {
-            let npt = match std::mem::replace(&mut self.next_page_token, NextPageToken::Done) {
-                // If there's no open request and the last page token we got was empty, we're done.
-                NextPageToken::Done => return Poll::Ready(None),
-                NextPageToken::FetchFromStart => vec![],
-                NextPageToken::Next(v) => v,
-            };
-            debug!(run_id=%self.run_id, "Fetching new history page");
-            let gw = self.client.clone();
-            let wid = self.wf_id.clone();
-            let rid = self.run_id.clone();
-            let resp_fut = async move {
-                gw.get_workflow_execution_history(wid, Some(rid), npt)
-                    .instrument(span!(tracing::Level::TRACE, "fetch_history_in_paginator"))
-                    .await
-            };
-            self.open_history_request.insert(resp_fut.boxed())
-        };
+        if this.open_history_request.is_none() {
+            // SAFETY: This is safe because the inner paginator cannot be dropped before the future,
+            //   and the future won't be moved from out of this struct.
+            this.open_history_request.set(Some(unsafe {
+                transmute(HistoryPaginator::get_next_page(this.inner).boxed())
+            }));
+        }
+        let history_req = this.open_history_request.as_mut().as_pin_mut().unwrap();
 
-        return match Future::poll(history_req.as_mut(), cx) {
+        match Future::poll(history_req, cx) {
             Poll::Ready(resp) => {
-                self.open_history_request = None;
+                this.open_history_request.set(None);
                 match resp {
                     Err(neterr) => Poll::Ready(Some(Err(neterr))),
-                    Ok(resp) => {
-                        self.extend_queue_with_new_page(resp);
-                        Poll::Ready(self.event_queue.pop_front().map(Ok))
-                    }
+                    Ok(_) => Poll::Ready(this.inner.event_queue.pop_front().map(Ok)),
                 }
             }
             Poll::Pending => Poll::Pending,
-        };
+        }
     }
 }
 
 impl HistoryUpdate {
-    pub fn new(history_iterator: HistoryPaginator, previous_wft_started_id: i64) -> Self {
+    /// Sometimes it's useful to take an update out of something without needing to use an option
+    /// field. Use this to replace the field with an empty update.
+    pub fn dummy() -> Self {
         Self {
-            events: history_iterator.fuse().boxed(),
-            buffered: VecDeque::new(),
-            previous_started_event_id: previous_wft_started_id,
+            events: vec![],
+            previous_wft_started_id: -1,
+            wft_started_id: -1,
+            has_last_wft: false,
+        }
+    }
+    pub fn is_real(&self) -> bool {
+        self.previous_wft_started_id >= 0
+    }
+    pub fn first_event_id(&self) -> Option<i64> {
+        self.events.get(0).map(|e| e.event_id)
+    }
+
+    /// Create an instance of an update directly from events. If the passed in event iterator has a
+    /// partial WFT sequence at the end, all events after the last complete WFT sequence (ending
+    /// with WFT started) are returned back to the caller, since the history update only works in
+    /// terms of complete WFT sequences.
+    pub fn from_events<I: IntoIterator<Item = HistoryEvent>>(
+        events: I,
+        previous_wft_started_id: i64,
+        wft_started_id: i64,
+        has_last_wft: bool,
+    ) -> (Self, Vec<HistoryEvent>)
+    where
+        <I as IntoIterator>::IntoIter: Send + 'static,
+    {
+        let mut all_events: Vec<_> = events.into_iter().collect();
+        let mut last_end =
+            find_end_index_of_next_wft_seq(all_events.as_slice(), previous_wft_started_id);
+        if matches!(last_end, NextWFTSeqEndIndex::Incomplete(_)) {
+            return if has_last_wft {
+                (
+                    Self {
+                        events: all_events,
+                        previous_wft_started_id,
+                        wft_started_id,
+                        has_last_wft,
+                    },
+                    vec![],
+                )
+            } else {
+                (
+                    Self {
+                        events: vec![],
+                        previous_wft_started_id,
+                        wft_started_id,
+                        has_last_wft,
+                    },
+                    all_events,
+                )
+            };
+        }
+        while let NextWFTSeqEndIndex::Complete(next_end_ix) = last_end {
+            let next_end_eid = all_events[next_end_ix].event_id;
+            // To save skipping all events at the front of this slice, only pass the relevant
+            // portion, but that means the returned index must be adjusted, hence the addition.
+            let next_end = find_end_index_of_next_wft_seq(&all_events[next_end_ix..], next_end_eid)
+                .add(next_end_ix);
+            if matches!(next_end, NextWFTSeqEndIndex::Incomplete(_)) {
+                break;
+            }
+            last_end = next_end;
         }
+        let remaining_events = if all_events.is_empty() {
+            vec![]
+        } else {
+            all_events.split_off(last_end.index() + 1)
+        };
+
+        (
+            Self {
+                events: all_events,
+                previous_wft_started_id,
+                wft_started_id,
+                has_last_wft,
+            },
+            remaining_events,
+        )
     }
 
-    /// Create an instance of an update directly from events - should only be used for replaying.
+    /// Create an instance of an update directly from events. The passed in events *must* consist
+    /// of one or more complete WFT sequences. IE: The event iterator must not end in the middle
+    /// of a WFT sequence.
+    #[cfg(test)]
     pub fn new_from_events<I: IntoIterator<Item = HistoryEvent>>(
         events: I,
         previous_wft_started_id: i64,
+        wft_started_id: i64,
     ) -> Self
     where
         <I as IntoIterator>::IntoIter: Send + 'static,
     {
         Self {
-            events: stream::iter(events.into_iter().map(Ok)).boxed(),
-            buffered: VecDeque::new(),
-            previous_started_event_id: previous_wft_started_id,
+            events: events.into_iter().collect(),
+            previous_wft_started_id,
+            wft_started_id,
+            has_last_wft: true,
         }
     }
 
-    /// Given a workflow task started id, return all events starting at that number (inclusive) to
-    /// the next WFT started event (inclusive). If there is no subsequent WFT started event,
-    /// remaining history is returned.
-    ///
-    /// Events are *consumed* by this process, to keep things efficient in workflow machines, and
-    /// the function may call out to server to fetch more pages if they are known to exist and
-    /// needed to complete the WFT sequence.
+    /// Given a workflow task started id, return all events starting at that number (exclusive) to
+    /// the next WFT started event (inclusive).
     ///
-    /// Always buffers the WFT sequence *after* the returned one as well, if it is available.
+    /// Events are *consumed* by this process, to keep things efficient in workflow machines.
     ///
-    /// Can return a tonic error in the event that fetching additional history was needed and failed
-    pub async fn take_next_wft_sequence(
-        &mut self,
-        from_wft_started_id: i64,
-    ) -> Result<Vec<HistoryEvent>, tonic::Status> {
-        let (next_wft_events, maybe_bonus_events) = self
-            .take_next_wft_sequence_impl(from_wft_started_id)
-            .await?;
-        if !maybe_bonus_events.is_empty() {
-            self.buffered.extend(maybe_bonus_events);
+    /// If we are out of WFT sequences that can be yielded by this update, it will return an empty
+    /// vec, indicating more pages will need to be fetched.
+    pub fn take_next_wft_sequence(&mut self, from_wft_started_id: i64) -> NextWFT {
+        // First, drop any events from the queue which are earlier than the passed-in id.
+        if let Some(ix_first_relevant) = self.starting_index_after_skipping(from_wft_started_id) {
+            self.events.drain(0..ix_first_relevant);
         }
-
-        if let Some(last_event_id) = next_wft_events.last().map(|he| he.event_id) {
-            // Always attempt to fetch the *next* WFT sequence as well, to buffer it for lookahead
-            let (buffer_these_events, maybe_bonus_events) =
-                self.take_next_wft_sequence_impl(last_event_id).await?;
-            self.buffered.extend(buffer_these_events);
-            if !maybe_bonus_events.is_empty() {
-                self.buffered.extend(maybe_bonus_events);
+        let next_wft_ix = find_end_index_of_next_wft_seq(&self.events, from_wft_started_id);
+        match next_wft_ix {
+            NextWFTSeqEndIndex::Incomplete(siz) => {
+                if self.has_last_wft {
+                    if siz == 0 {
+                        NextWFT::ReplayOver
+                    } else {
+                        self.build_next_wft(siz)
+                    }
+                } else {
+                    if siz != 0 {
+                        panic!(
+                            "HistoryUpdate was created with an incomplete WFT. This is an SDK bug."
+                        );
+                    }
+                    NextWFT::NeedFetch
+                }
             }
+            NextWFTSeqEndIndex::Complete(next_wft_ix) => self.build_next_wft(next_wft_ix),
         }
+    }
 
-        Ok(next_wft_events)
+    fn build_next_wft(&mut self, drain_this_much: usize) -> NextWFT {
+        NextWFT::WFT(
+            self.events.drain(0..=drain_this_much).collect(),
+            self.events.is_empty() && self.has_last_wft,
+        )
     }
 
     /// Lets the caller peek ahead at the next WFT sequence that will be returned by
-    /// [take_next_wft_sequence]. Will always return an empty iterator if that has not been called
-    /// first. May also return an empty iterator or incomplete sequence if we are at the end of
-    /// history.
-    pub fn peek_next_wft_sequence(&self) -> impl Iterator<Item = &HistoryEvent> {
-        self.buffered.iter()
-    }
-
-    /// Retrieve the next WFT sequence, first from buffered events and then from the real stream.
-    /// Returns (events up to the next logical wft sequence, extra events that were taken but
-    /// should be re-appended to the end of the buffer).
-    async fn take_next_wft_sequence_impl(
-        &mut self,
-        from_event_id: i64,
-    ) -> Result<(Vec<HistoryEvent>, Vec<HistoryEvent>), tonic::Status> {
-        let mut events_to_next_wft_started: Vec<HistoryEvent> = vec![];
-
-        // This flag tracks if, while determining events to be returned, we have seen the next
-        // logically significant WFT started event which follows the one that was passed in as a
-        // parameter. If a WFT fails, times out, or is devoid of commands (ie: a heartbeat) it is
-        // not significant. So we will stop returning events (exclusive) as soon as we see an event
-        // following a WFT started that is *not* failed, timed out, or completed with a command.
-        let mut next_wft_state = NextWftState::NotSeen;
-        let mut should_pop = |e: &HistoryEvent| {
-            if e.event_id <= from_event_id {
-                return true;
-            } else if e.event_type() == EventType::WorkflowTaskStarted {
-                next_wft_state = NextWftState::Seen;
-                return true;
+    /// [take_next_wft_sequence]. Will always return the first available WFT sequence if that has
+    /// not been called first. May also return an empty iterator or incomplete sequence if we are at
+    /// the end of history.
+    pub fn peek_next_wft_sequence(&self, from_wft_started_id: i64) -> &[HistoryEvent] {
+        let ix_first_relevant = self
+            .starting_index_after_skipping(from_wft_started_id)
+            .unwrap_or_default();
+        let relevant_events = &self.events[ix_first_relevant..];
+        if relevant_events.is_empty() {
+            return relevant_events;
+        }
+        let ix_end = find_end_index_of_next_wft_seq(relevant_events, from_wft_started_id).index();
+        &relevant_events[0..=ix_end]
+    }
+
+    /// Returns true if this update has the next needed WFT sequence, false if events will need to
+    /// be fetched in order to create a complete update with the entire next WFT sequence.
+    pub fn can_take_next_wft_sequence(&self, from_wft_started_id: i64) -> bool {
+        let next_wft_ix = find_end_index_of_next_wft_seq(&self.events, from_wft_started_id);
+        if let NextWFTSeqEndIndex::Incomplete(_) = next_wft_ix {
+            if !self.has_last_wft {
+                return false;
             }
+        }
+        true
+    }
 
-            match next_wft_state {
-                NextWftState::Seen => {
-                    // Must ignore failures and timeouts
-                    if e.event_type() == EventType::WorkflowTaskFailed
-                        || e.event_type() == EventType::WorkflowTaskTimedOut
-                    {
-                        next_wft_state = NextWftState::NotSeen;
-                        return true;
-                    } else if e.event_type() == EventType::WorkflowTaskCompleted {
-                        next_wft_state = NextWftState::SeenCompleted;
-                        return true;
-                    }
-                    false
+    /// Returns the next WFT completed event attributes, if any, starting at (inclusive) the
+    /// `from_id`
+    pub fn peek_next_wft_completed(
+        &self,
+        from_id: i64,
+    ) -> Option<&WorkflowTaskCompletedEventAttributes> {
+        self.events
+            .iter()
+            .skip_while(|e| e.event_id < from_id)
+            .find_map(|e| match &e.attributes {
+                Some(history_event::Attributes::WorkflowTaskCompletedEventAttributes(ref a)) => {
+                    Some(a)
                 }
-                NextWftState::SeenCompleted => {
-                    // If we've seen the WFT be completed, and this event is another scheduled, then
-                    // this was an empty heartbeat we should ignore.
-                    if e.event_type() == EventType::WorkflowTaskScheduled {
-                        next_wft_state = NextWftState::NotSeen;
-                        return true;
-                    }
-                    // Otherwise, we're done here
-                    false
-                }
-                NextWftState::NotSeen => true,
-            }
-        };
+                _ => None,
+            })
+    }
+
+    fn starting_index_after_skipping(&self, from_wft_started_id: i64) -> Option<usize> {
+        self.events
+            .iter()
+            .find_position(|e| e.event_id > from_wft_started_id)
+            .map(|(ix, _)| ix)
+    }
+}
 
-        // Fetch events from the buffer first, then from the network
-        let mut event_q = stream::iter(self.buffered.drain(..).map(Ok)).chain(&mut self.events);
+#[derive(Debug, Copy, Clone)]
+enum NextWFTSeqEndIndex {
+    /// The next WFT sequence is completely contained within the passed-in iterator
+    Complete(usize),
+    /// The next WFT sequence is not found within the passed-in iterator, and the contained
+    /// value is the last index of the iterator.
+    Incomplete(usize),
+}
+impl NextWFTSeqEndIndex {
+    fn index(self) -> usize {
+        match self {
+            NextWFTSeqEndIndex::Complete(ix) | NextWFTSeqEndIndex::Incomplete(ix) => ix,
+        }
+    }
+    fn add(self, val: usize) -> Self {
+        match self {
+            NextWFTSeqEndIndex::Complete(ix) => NextWFTSeqEndIndex::Complete(ix + val),
+            NextWFTSeqEndIndex::Incomplete(ix) => NextWFTSeqEndIndex::Incomplete(ix + val),
+        }
+    }
+}
 
-        let mut extra_e = vec![];
-        let mut last_seen_id = None;
-        while let Some(e) = event_q.next().await {
-            let e = e?;
+/// Discovers the index of the last event in next WFT sequence within the passed-in slice
+fn find_end_index_of_next_wft_seq(
+    events: &[HistoryEvent],
+    from_event_id: i64,
+) -> NextWFTSeqEndIndex {
+    if events.is_empty() {
+        return NextWFTSeqEndIndex::Incomplete(0);
+    }
+    let mut last_index = 0;
+    let mut saw_any_non_wft_event = false;
+    for (ix, e) in events.iter().enumerate() {
+        last_index = ix;
 
-            // This little block prevents us from infinitely fetching work from the server in the
-            // event that, for whatever reason, it keeps returning stuff we've already seen.
-            if let Some(last_id) = last_seen_id {
-                if e.event_id <= last_id {
-                    error!("Server returned history event IDs that went backwards!");
-                    break;
+        // It's possible to have gotten a new history update without eviction (ex: unhandled
+        // command on completion), where we may need to skip events we already handled.
+        if e.event_id <= from_event_id {
+            continue;
+        }
+
+        if !matches!(
+            e.event_type(),
+            EventType::WorkflowTaskFailed
+                | EventType::WorkflowTaskTimedOut
+                | EventType::WorkflowTaskScheduled
+                | EventType::WorkflowTaskStarted
+                | EventType::WorkflowTaskCompleted
+        ) {
+            saw_any_non_wft_event = true;
+        }
+        if e.is_final_wf_execution_event() {
+            return NextWFTSeqEndIndex::Complete(last_index);
+        }
+
+        if e.event_type() == EventType::WorkflowTaskStarted {
+            if let Some(next_event) = events.get(ix + 1) {
+                let et = next_event.event_type();
+                // If the next event is WFT timeout or fail, or abrupt WF execution end, that
+                // doesn't conclude a WFT sequence.
+                if matches!(
+                    et,
+                    EventType::WorkflowTaskFailed
+                        | EventType::WorkflowTaskTimedOut
+                        | EventType::WorkflowExecutionTimedOut
+                        | EventType::WorkflowExecutionTerminated
+                        | EventType::WorkflowExecutionCanceled
+                ) {
+                    continue;
                 }
-            }
-            last_seen_id = Some(e.event_id);
-
-            // It's possible to have gotten a new history update without eviction (ex: unhandled
-            // command on completion), where we may need to skip events we already handled.
-            if e.event_id > from_event_id {
-                if !should_pop(&e) {
-                    if next_wft_state == NextWftState::SeenCompleted {
-                        // We have seen the wft completed event, but decided to exit. We don't
-                        // want to return that event as part of this sequence, so include it for
-                        // re-buffering along with the event we're currently on.
-                        extra_e.push(
-                            events_to_next_wft_started
-                                .pop()
-                                .expect("There is an element here by definition"),
-                        );
+                // If we've never seen an interesting event and the next two events are a completion
+                // followed immediately again by scheduled, then this is a WFT heartbeat and also
+                // doesn't conclude the sequence.
+                else if et == EventType::WorkflowTaskCompleted {
+                    if let Some(next_next_event) = events.get(ix + 2) {
+                        if next_next_event.event_type() == EventType::WorkflowTaskScheduled {
+                            continue;
+                        } else {
+                            saw_any_non_wft_event = true;
+                        }
                     }
-                    extra_e.push(e);
-                    break;
                 }
-                events_to_next_wft_started.push(e);
+            }
+            if saw_any_non_wft_event {
+                return NextWFTSeqEndIndex::Complete(ix);
             }
         }
-
-        Ok((events_to_next_wft_started, extra_e))
     }
-}
 
-#[derive(Eq, PartialEq, Debug)]
-enum NextWftState {
-    NotSeen,
-    Seen,
-    SeenCompleted,
+    NextWFTSeqEndIndex::Incomplete(last_index)
 }
 
-impl From<HistoryInfo> for HistoryUpdate {
-    fn from(v: HistoryInfo) -> Self {
-        Self::new_from_events(v.events().to_vec(), v.previous_started_event_id())
+#[cfg(test)]
+pub mod tests {
+    use super::*;
+    use crate::{
+        replay::{HistoryInfo, TestHistoryBuilder},
+        test_help::canned_histories,
+        worker::client::mocks::mock_workflow_client,
+    };
+    use futures_util::TryStreamExt;
+    use temporal_sdk_core_protos::temporal::api::workflowservice::v1::GetWorkflowExecutionHistoryResponse;
+
+    impl From<HistoryInfo> for HistoryUpdate {
+        fn from(v: HistoryInfo) -> Self {
+            Self::new_from_events(
+                v.events().to_vec(),
+                v.previous_started_event_id(),
+                v.workflow_task_started_event_id(),
+            )
+        }
     }
-}
 
-pub trait TestHBExt {
-    fn as_history_update(&self) -> HistoryUpdate;
-}
+    pub trait TestHBExt {
+        fn as_history_update(&self) -> HistoryUpdate;
+    }
 
-impl TestHBExt for TestHistoryBuilder {
-    fn as_history_update(&self) -> HistoryUpdate {
-        self.get_full_history_info().unwrap().into()
+    impl TestHBExt for TestHistoryBuilder {
+        fn as_history_update(&self) -> HistoryUpdate {
+            self.get_full_history_info().unwrap().into()
+        }
     }
-}
 
-#[cfg(test)]
-pub mod tests {
-    use super::*;
-    use crate::{test_help::canned_histories, worker::client::mocks::mock_workflow_client};
+    impl NextWFT {
+        fn unwrap_events(self) -> Vec<HistoryEvent> {
+            match self {
+                NextWFT::WFT(e, _) => e,
+                o => panic!("Must be complete WFT: {o:?}"),
+            }
+        }
+    }
 
-    #[tokio::test]
-    async fn consumes_standard_wft_sequence() {
+    fn next_check_peek(update: &mut HistoryUpdate, from_id: i64) -> Vec<HistoryEvent> {
+        let seq_peeked = update.peek_next_wft_sequence(from_id).to_vec();
+        let seq = update.take_next_wft_sequence(from_id).unwrap_events();
+        assert_eq!(seq, seq_peeked);
+        seq
+    }
+
+    #[test]
+    fn consumes_standard_wft_sequence() {
         let timer_hist = canned_histories::single_timer("t");
         let mut update = timer_hist.as_history_update();
-        let seq_1 = update.take_next_wft_sequence(0).await.unwrap();
+        let seq_1 = next_check_peek(&mut update, 0);
         assert_eq!(seq_1.len(), 3);
         assert_eq!(seq_1.last().unwrap().event_id, 3);
-        let seq_2 = update.take_next_wft_sequence(3).await.unwrap();
+        let seq_2_peeked = update.peek_next_wft_sequence(0).to_vec();
+        let seq_2 = next_check_peek(&mut update, 3);
+        assert_eq!(seq_2, seq_2_peeked);
         assert_eq!(seq_2.len(), 5);
         assert_eq!(seq_2.last().unwrap().event_id, 8);
     }
 
-    #[tokio::test]
-    async fn skips_wft_failed() {
+    #[test]
+    fn skips_wft_failed() {
         let failed_hist = canned_histories::workflow_fails_with_reset_after_timer("t", "runid");
         let mut update = failed_hist.as_history_update();
-        let seq_1 = update.take_next_wft_sequence(0).await.unwrap();
+        let seq_1 = next_check_peek(&mut update, 0);
         assert_eq!(seq_1.len(), 3);
         assert_eq!(seq_1.last().unwrap().event_id, 3);
-        let seq_2 = update.take_next_wft_sequence(3).await.unwrap();
+        let seq_2 = next_check_peek(&mut update, 3);
         assert_eq!(seq_2.len(), 8);
         assert_eq!(seq_2.last().unwrap().event_id, 11);
     }
 
-    #[tokio::test]
-    async fn skips_wft_timeout() {
+    #[test]
+    fn skips_wft_timeout() {
         let failed_hist = canned_histories::wft_timeout_repro();
         let mut update = failed_hist.as_history_update();
-        let seq_1 = update.take_next_wft_sequence(0).await.unwrap();
+        let seq_1 = next_check_peek(&mut update, 0);
         assert_eq!(seq_1.len(), 3);
         assert_eq!(seq_1.last().unwrap().event_id, 3);
-        let seq_2 = update.take_next_wft_sequence(3).await.unwrap();
+        let seq_2 = next_check_peek(&mut update, 3);
         assert_eq!(seq_2.len(), 11);
         assert_eq!(seq_2.last().unwrap().event_id, 14);
     }
 
-    #[tokio::test]
-    async fn skips_events_before_desired_wft() {
+    #[test]
+    fn skips_events_before_desired_wft() {
         let timer_hist = canned_histories::single_timer("t");
         let mut update = timer_hist.as_history_update();
         // We haven't processed the first 3 events, but we should still only get the second sequence
-        let seq_2 = update.take_next_wft_sequence(3).await.unwrap();
+        let seq_2 = update.take_next_wft_sequence(3).unwrap_events();
         assert_eq!(seq_2.len(), 5);
         assert_eq!(seq_2.last().unwrap().event_id, 8);
     }
 
-    #[tokio::test]
-    async fn history_ends_abruptly() {
+    #[test]
+    fn history_ends_abruptly() {
         let mut timer_hist = canned_histories::single_timer("t");
         timer_hist.add_workflow_execution_terminated();
         let mut update = timer_hist.as_history_update();
-        let seq_2 = update.take_next_wft_sequence(3).await.unwrap();
-        assert_eq!(seq_2.len(), 5);
-        assert_eq!(seq_2.last().unwrap().event_id, 8);
+        let seq_2 = update.take_next_wft_sequence(3).unwrap_events();
+        assert_eq!(seq_2.len(), 6);
+        assert_eq!(seq_2.last().unwrap().event_id, 9);
     }
 
-    #[tokio::test]
-    async fn heartbeats_skipped() {
+    #[test]
+    fn heartbeats_skipped() {
         let mut t = TestHistoryBuilder::default();
         t.add_by_type(EventType::WorkflowExecutionStarted);
         t.add_full_wf_task();
-        t.add_full_wf_task();
-        t.add_get_event_id(EventType::TimerStarted, None);
+        t.add_full_wf_task(); // wft started 6
+        t.add_by_type(EventType::TimerStarted);
+        t.add_full_wf_task(); // wft started 10
         t.add_full_wf_task();
         t.add_full_wf_task();
-        t.add_full_wf_task();
-        t.add_full_wf_task();
-        t.add_get_event_id(EventType::TimerStarted, None);
-        t.add_full_wf_task();
+        t.add_full_wf_task(); // wft started 19
+        t.add_by_type(EventType::TimerStarted);
+        t.add_full_wf_task(); // wft started 23
         t.add_we_signaled("whee", vec![]);
         t.add_full_wf_task();
         t.add_workflow_execution_completed();
 
         let mut update = t.as_history_update();
-        let seq = update.take_next_wft_sequence(0).await.unwrap();
+        let seq = next_check_peek(&mut update, 0);
         assert_eq!(seq.len(), 6);
-        let seq = update.take_next_wft_sequence(6).await.unwrap();
+        let seq = next_check_peek(&mut update, 6);
         assert_eq!(seq.len(), 13);
-        let seq = update.take_next_wft_sequence(19).await.unwrap();
+        let seq = next_check_peek(&mut update, 19);
         assert_eq!(seq.len(), 4);
-        let seq = update.take_next_wft_sequence(23).await.unwrap();
+        let seq = next_check_peek(&mut update, 23);
         assert_eq!(seq.len(), 4);
-        let seq = update.take_next_wft_sequence(27).await.unwrap();
+        let seq = next_check_peek(&mut update, 27);
         assert_eq!(seq.len(), 2);
     }
 
-    #[tokio::test]
-    async fn paginator_fetches_new_pages() {
-        // Note that this test triggers the "event ids that went backwards" error, acceptably.
-        // Can be fixed by having mock not return earlier events.
-        let wft_count = 500;
-        let long_hist = canned_histories::long_sequential_timers(wft_count);
-        let initial_hist = long_hist.get_history_info(10).unwrap();
-        let prev_started = initial_hist.previous_started_event_id();
+    #[test]
+    fn heartbeat_marker_end() {
+        let mut t = TestHistoryBuilder::default();
+        t.add_by_type(EventType::WorkflowExecutionStarted);
+        t.add_full_wf_task();
+        t.add_full_wf_task();
+        t.add_local_activity_result_marker(1, "1", "done".into());
+        t.add_workflow_execution_completed();
+
+        let mut update = t.as_history_update();
+        let seq = next_check_peek(&mut update, 3);
+        // completed, sched, started
+        assert_eq!(seq.len(), 3);
+        let seq = next_check_peek(&mut update, 6);
+        assert_eq!(seq.len(), 3);
+    }
+
+    fn paginator_setup(history: TestHistoryBuilder, chunk_size: usize) -> HistoryPaginator {
+        let hinfo = history.get_full_history_info().unwrap();
+        let wft_started = hinfo.workflow_task_started_event_id();
+        let full_hist = hinfo.into_events();
+        let initial_hist = full_hist.chunks(chunk_size).next().unwrap().to_vec();
         let mut mock_client = mock_workflow_client();
 
-        let mut npt = 2;
+        let mut npt = 1;
         mock_client
             .expect_get_workflow_execution_history()
             .returning(move |_, _, passed_npt| {
                 assert_eq!(passed_npt, vec![npt]);
-                let history = long_hist.get_history_info(10 * npt as usize).unwrap();
+                let mut hist_chunks = full_hist.chunks(chunk_size).peekable();
+                let next_chunks = hist_chunks.nth(npt.into()).unwrap_or_default();
                 npt += 1;
+                let next_page_token = if hist_chunks.peek().is_none() {
+                    vec![]
+                } else {
+                    vec![npt]
+                };
                 Ok(GetWorkflowExecutionHistoryResponse {
-                    history: Some(history.into()),
+                    history: Some(History {
+                        events: next_chunks.into(),
+                    }),
                     raw_history: vec![],
-                    next_page_token: vec![npt],
+                    next_page_token,
                     archived: false,
                 })
             });
 
-        let mut update = HistoryUpdate::new(
-            HistoryPaginator::new(
-                initial_hist.into(),
-                "wfid".to_string(),
-                "runid".to_string(),
-                vec![2], // Start at page "2"
-                Arc::new(mock_client),
-            ),
-            prev_started,
+        HistoryPaginator::new(
+            History {
+                events: initial_hist,
+            },
+            0,
+            wft_started,
+            "wfid".to_string(),
+            "runid".to_string(),
+            vec![1],
+            Arc::new(mock_client),
+        )
+    }
+
+    #[rstest::rstest]
+    #[tokio::test]
+    async fn paginator_extracts_updates(#[values(10, 11, 12, 13, 14)] chunk_size: usize) {
+        let wft_count = 100;
+        let mut paginator = paginator_setup(
+            canned_histories::long_sequential_timers(wft_count),
+            chunk_size,
         );
+        let mut update = paginator.extract_next_update().await.unwrap();
 
-        let seq = update.take_next_wft_sequence(0).await.unwrap();
+        let seq = update.take_next_wft_sequence(0).unwrap_events();
         assert_eq!(seq.len(), 3);
 
         let mut last_event_id = 3;
         let mut last_started_id = 3;
-        for _ in 1..wft_count {
-            let seq = update
-                .take_next_wft_sequence(last_started_id)
-                .await
-                .unwrap();
+        for i in 1..wft_count {
+            let seq = {
+                match update.take_next_wft_sequence(last_started_id) {
+                    NextWFT::WFT(seq, _) => seq,
+                    NextWFT::NeedFetch => {
+                        update = paginator.extract_next_update().await.unwrap();
+                        update
+                            .take_next_wft_sequence(last_started_id)
+                            .unwrap_events()
+                    }
+                    NextWFT::ReplayOver => {
+                        assert_eq!(i, wft_count - 1);
+                        break;
+                    }
+                }
+            };
             for e in &seq {
                 last_event_id += 1;
                 assert_eq!(e.event_id, last_event_id);
@@ -495,10 +908,139 @@ pub mod tests {
         }
     }
 
+    #[tokio::test]
+    async fn paginator_streams() {
+        let wft_count = 10;
+        let paginator = StreamingHistoryPaginator::new(paginator_setup(
+            canned_histories::long_sequential_timers(wft_count),
+            10,
+        ));
+        let everything: Vec<_> = paginator.try_collect().await.unwrap();
+        assert_eq!(everything.len(), (wft_count + 1) * 5);
+        everything.iter().fold(1, |event_id, e| {
+            assert_eq!(event_id, e.event_id);
+            e.event_id + 1
+        });
+    }
+
+    fn three_wfts_then_heartbeats() -> TestHistoryBuilder {
+        let mut t = TestHistoryBuilder::default();
+        // Start with two complete normal WFTs
+        t.add_by_type(EventType::WorkflowExecutionStarted);
+        t.add_full_wf_task(); // wft start - 3
+        t.add_by_type(EventType::TimerStarted);
+        t.add_full_wf_task(); // wft start - 7
+        t.add_by_type(EventType::TimerStarted);
+        t.add_full_wf_task(); // wft start - 11
+        for _ in 1..50 {
+            // Add a bunch of heartbeats with no commands, which count as one task
+            t.add_full_wf_task();
+        }
+        t.add_workflow_execution_completed();
+        t
+    }
+
+    #[tokio::test]
+    async fn needs_fetch_if_ending_in_middle_of_wft_seq() {
+        let t = three_wfts_then_heartbeats();
+        let mut ends_in_middle_of_seq = t.as_history_update().events;
+        ends_in_middle_of_seq.truncate(19);
+        // The update should contain the first two complete WFTs, ending on the 8th event which
+        // is WFT started. The remaining events should be returned. False flags means the creator
+        // knows there are more events, so we should return need fetch
+        let (mut update, remaining) = HistoryUpdate::from_events(
+            ends_in_middle_of_seq,
+            0,
+            t.get_full_history_info()
+                .unwrap()
+                .workflow_task_started_event_id(),
+            false,
+        );
+        assert_eq!(remaining[0].event_id, 8);
+        assert_eq!(remaining.last().unwrap().event_id, 19);
+        let seq = update.take_next_wft_sequence(0).unwrap_events();
+        assert_eq!(seq.last().unwrap().event_id, 3);
+        let seq = update.take_next_wft_sequence(3).unwrap_events();
+        assert_eq!(seq.last().unwrap().event_id, 7);
+        let next = update.take_next_wft_sequence(7);
+        assert_matches!(next, NextWFT::NeedFetch);
+    }
+
+    // Like the above, but if the history happens to be cut off at a wft boundary, (even though
+    // there may have been many heartbeats after we have no way of knowing about), it's going to
+    // count events 7-20 as a WFT since there is started, completed, timer command, ..heartbeats..
+    #[tokio::test]
+    async fn needs_fetch_after_complete_seq_with_heartbeats() {
+        let t = three_wfts_then_heartbeats();
+        let mut ends_in_middle_of_seq = t.as_history_update().events;
+        ends_in_middle_of_seq.truncate(20);
+        let (mut update, remaining) = HistoryUpdate::from_events(
+            ends_in_middle_of_seq,
+            0,
+            t.get_full_history_info()
+                .unwrap()
+                .workflow_task_started_event_id(),
+            false,
+        );
+        assert!(remaining.is_empty());
+        let seq = update.take_next_wft_sequence(0).unwrap_events();
+        assert_eq!(seq.last().unwrap().event_id, 3);
+        let seq = update.take_next_wft_sequence(3).unwrap_events();
+        assert_eq!(seq.last().unwrap().event_id, 7);
+        let seq = update.take_next_wft_sequence(7).unwrap_events();
+        assert_eq!(seq.last().unwrap().event_id, 20);
+        let next = update.take_next_wft_sequence(20);
+        assert_matches!(next, NextWFT::NeedFetch);
+    }
+
+    #[rstest::rstest]
+    #[tokio::test]
+    async fn paginator_works_with_wft_over_multiple_pages(
+        #[values(10, 11, 12, 13, 14)] chunk_size: usize,
+    ) {
+        let t = three_wfts_then_heartbeats();
+        let mut paginator = paginator_setup(t, chunk_size);
+        let mut update = paginator.extract_next_update().await.unwrap();
+        let mut last_id = 0;
+        loop {
+            let seq = update.take_next_wft_sequence(last_id);
+            match seq {
+                NextWFT::WFT(seq, _) => {
+                    last_id = seq.last().unwrap().event_id;
+                }
+                NextWFT::NeedFetch => {
+                    update = paginator.extract_next_update().await.unwrap();
+                }
+                NextWFT::ReplayOver => break,
+            }
+        }
+        assert_eq!(last_id, 160);
+    }
+
+    #[tokio::test]
+    async fn task_just_before_heartbeat_chain_is_taken() {
+        let t = three_wfts_then_heartbeats();
+        let mut update = t.as_history_update();
+        let seq = update.take_next_wft_sequence(0).unwrap_events();
+        assert_eq!(seq.last().unwrap().event_id, 3);
+        let seq = update.take_next_wft_sequence(3).unwrap_events();
+        assert_eq!(seq.last().unwrap().event_id, 7);
+        let seq = update.take_next_wft_sequence(7).unwrap_events();
+        assert_eq!(seq.last().unwrap().event_id, 158);
+        let seq = update.take_next_wft_sequence(158).unwrap_events();
+        assert_eq!(seq.last().unwrap().event_id, 160);
+        assert_eq!(
+            seq.last().unwrap().event_type(),
+            EventType::WorkflowExecutionCompleted
+        );
+    }
+
     #[tokio::test]
     async fn handles_cache_misses() {
         let timer_hist = canned_histories::single_timer("t");
         let partial_task = timer_hist.get_one_wft(2).unwrap();
+        let prev_started_wft_id = partial_task.previous_started_event_id();
+        let wft_started_id = partial_task.workflow_task_started_event_id();
         let mut history_from_get: GetWorkflowExecutionHistoryResponse =
             timer_hist.get_history_info(2).unwrap().into();
         // Chop off the last event, which is WFT started, which server doesn't return in get
@@ -509,24 +1051,121 @@ pub mod tests {
             .expect_get_workflow_execution_history()
             .returning(move |_, _, _| Ok(history_from_get.clone()));
 
-        let mut update = HistoryUpdate::new(
-            HistoryPaginator::new(
-                partial_task.into(),
-                "wfid".to_string(),
-                "runid".to_string(),
-                // A cache miss means we'll try to fetch from start
-                NextPageToken::FetchFromStart,
-                Arc::new(mock_client),
-            ),
-            1,
+        let mut paginator = HistoryPaginator::new(
+            partial_task.into(),
+            prev_started_wft_id,
+            wft_started_id,
+            "wfid".to_string(),
+            "runid".to_string(),
+            // A cache miss means we'll try to fetch from start
+            NextPageToken::FetchFromStart,
+            Arc::new(mock_client),
         );
+        let mut update = paginator.extract_next_update().await.unwrap();
         // We expect if we try to take the first task sequence that the first event is the first
         // event in the sequence.
-        let seq = update.take_next_wft_sequence(0).await.unwrap();
+        let seq = update.take_next_wft_sequence(0).unwrap_events();
         assert_eq!(seq[0].event_id, 1);
-        let seq = update.take_next_wft_sequence(3).await.unwrap();
+        let seq = update.take_next_wft_sequence(3).unwrap_events();
         // Verify anything extra (which should only ever be WFT started) was re-appended to the
         // end of the event iteration after fetching the old history.
         assert_eq!(seq.last().unwrap().event_id, 8);
     }
+
+    #[test]
+    fn la_marker_chunking() {
+        let mut t = TestHistoryBuilder::default();
+        t.add_by_type(EventType::WorkflowExecutionStarted);
+        t.add_full_wf_task();
+        t.add_we_signaled("whatever", vec![]);
+        t.add_full_wf_task(); // started - 7
+        t.add_local_activity_result_marker(1, "hi", Default::default());
+        let act_s = t.add_activity_task_scheduled("1");
+        let act_st = t.add_activity_task_started(act_s);
+        t.add_activity_task_completed(act_s, act_st, Default::default());
+        t.add_workflow_task_scheduled_and_started();
+        t.add_workflow_task_timed_out();
+        t.add_workflow_task_scheduled_and_started();
+        t.add_workflow_task_timed_out();
+        t.add_workflow_task_scheduled_and_started();
+
+        let mut update = t.as_history_update();
+        let seq = next_check_peek(&mut update, 0);
+        assert_eq!(seq.len(), 3);
+        let seq = next_check_peek(&mut update, 3);
+        assert_eq!(seq.len(), 4);
+        let seq = next_check_peek(&mut update, 7);
+        assert_eq!(seq.len(), 13);
+    }
+
+    #[tokio::test]
+    async fn handles_blank_fetch_response() {
+        let timer_hist = canned_histories::single_timer("t");
+        let partial_task = timer_hist.get_one_wft(2).unwrap();
+        let prev_started_wft_id = partial_task.previous_started_event_id();
+        let wft_started_id = partial_task.workflow_task_started_event_id();
+        let mut mock_client = mock_workflow_client();
+        mock_client
+            .expect_get_workflow_execution_history()
+            .returning(move |_, _, _| Ok(Default::default()));
+
+        let mut paginator = HistoryPaginator::new(
+            partial_task.into(),
+            prev_started_wft_id,
+            wft_started_id,
+            "wfid".to_string(),
+            "runid".to_string(),
+            // A cache miss means we'll try to fetch from start
+            NextPageToken::FetchFromStart,
+            Arc::new(mock_client),
+        );
+        let err = paginator.extract_next_update().await.unwrap_err();
+        assert_matches!(err.code(), tonic::Code::Unknown);
+    }
+
+    #[tokio::test]
+    async fn handles_empty_page_with_next_token() {
+        let timer_hist = canned_histories::single_timer("t");
+        let partial_task = timer_hist.get_one_wft(2).unwrap();
+        let prev_started_wft_id = partial_task.previous_started_event_id();
+        let wft_started_id = partial_task.workflow_task_started_event_id();
+        let full_resp: GetWorkflowExecutionHistoryResponse =
+            timer_hist.get_full_history_info().unwrap().into();
+        let mut mock_client = mock_workflow_client();
+        mock_client
+            .expect_get_workflow_execution_history()
+            .returning(move |_, _, _| {
+                Ok(GetWorkflowExecutionHistoryResponse {
+                    history: Some(History { events: vec![] }),
+                    raw_history: vec![],
+                    next_page_token: vec![2],
+                    archived: false,
+                })
+            })
+            .times(1);
+        mock_client
+            .expect_get_workflow_execution_history()
+            .returning(move |_, _, _| Ok(full_resp.clone()))
+            .times(1);
+
+        let mut paginator = HistoryPaginator::new(
+            partial_task.into(),
+            prev_started_wft_id,
+            wft_started_id,
+            "wfid".to_string(),
+            "runid".to_string(),
+            // A cache miss means we'll try to fetch from start
+            NextPageToken::FetchFromStart,
+            Arc::new(mock_client),
+        );
+        let mut update = paginator.extract_next_update().await.unwrap();
+        let seq = update.take_next_wft_sequence(0).unwrap_events();
+        assert_eq!(seq.last().unwrap().event_id, 3);
+        let seq = update.take_next_wft_sequence(3).unwrap_events();
+        assert_eq!(seq.last().unwrap().event_id, 8);
+        assert_matches!(update.take_next_wft_sequence(8), NextWFT::ReplayOver);
+    }
+
+    // TODO: Test we dont re-feed pointless updates if fetching returns <= events we already
+    //   processed
 }