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

package/sdk-core/client/src/worker_registry/mod.rs

@@ -0,0 +1,264 @@
+//! This module enables the tracking of workers that are associated with a client instance.
+//! This is needed to implement Eager Workflow Start, a latency optimization in which the client,
+//!  after reserving a slot, directly forwards a WFT to a local worker.
+
+use parking_lot::RwLock;
+use slotmap::SlotMap;
+use std::collections::{hash_map::Entry::Vacant, HashMap};
+
+use temporal_sdk_core_protos::temporal::api::workflowservice::v1::PollWorkflowTaskQueueResponse;
+
+slotmap::new_key_type! {
+    /// Registration key for a worker
+    pub struct WorkerKey;
+}
+
+/// This trait is implemented by an object associated with a worker, which provides WFT processing slots.
+#[cfg_attr(test, mockall::automock)]
+pub trait SlotProvider: std::fmt::Debug {
+    /// The namespace for the WFTs that it can process.
+    fn namespace(&self) -> &str;
+    /// The task queue this provider listens to.
+    fn task_queue(&self) -> &str;
+    /// Try to reserve a slot on this worker.
+    fn try_reserve_wft_slot(&self) -> Option<Box<dyn Slot + Send>>;
+}
+
+/// This trait represents a slot reserved for processing a WFT by a worker.
+#[cfg_attr(test, mockall::automock)]
+pub trait Slot {
+    /// Consumes this slot by dispatching a WFT to its worker. This can only be called once.
+    fn schedule_wft(
+        self: Box<Self>,
+        task: PollWorkflowTaskQueueResponse,
+    ) -> Result<(), anyhow::Error>;
+}
+
+#[derive(PartialEq, Eq, Hash, Debug, Clone)]
+struct SlotKey {
+    namespace: String,
+    task_queue: String,
+}
+
+impl SlotKey {
+    fn new(namespace: String, task_queue: String) -> SlotKey {
+        SlotKey {
+            namespace,
+            task_queue,
+        }
+    }
+}
+
+/// This is an inner class for [SlotManager] needed to hide the mutex.
+#[derive(Default, Debug)]
+struct SlotManagerImpl {
+    /// Maps keys, i.e., namespace#task_queue, to provider.
+    providers: HashMap<SlotKey, Box<dyn SlotProvider + Send + Sync>>,
+    /// Maps ids to keys in `providers`.
+    index: SlotMap<WorkerKey, SlotKey>,
+}
+
+impl SlotManagerImpl {
+    /// Factory method.
+    fn new() -> Self {
+        Self {
+            index: Default::default(),
+            providers: Default::default(),
+        }
+    }
+
+    fn try_reserve_wft_slot(
+        &self,
+        namespace: String,
+        task_queue: String,
+    ) -> Option<Box<dyn Slot + Send>> {
+        let key = SlotKey::new(namespace, task_queue);
+        if let Some(p) = self.providers.get(&key) {
+            if let Some(slot) = p.try_reserve_wft_slot() {
+                return Some(slot);
+            }
+        }
+        None
+    }
+
+    fn register(&mut self, provider: Box<dyn SlotProvider + Send + Sync>) -> Option<WorkerKey> {
+        let key = SlotKey::new(
+            provider.namespace().to_string(),
+            provider.task_queue().to_string(),
+        );
+        if let Vacant(p) = self.providers.entry(key.clone()) {
+            p.insert(provider);
+            Some(self.index.insert(key))
+        } else {
+            warn!("Ignoring registration for worker: {key:?}.");
+            None
+        }
+    }
+
+    fn unregister(&mut self, id: WorkerKey) {
+        if let Some(key) = self.index.remove(id) {
+            self.providers.remove(&key);
+        }
+    }
+
+    #[cfg(test)]
+    fn num_providers(&self) -> (usize, usize) {
+        (self.index.len(), self.providers.len())
+    }
+}
+
+/// Enables local workers to made themselves visible to a shared client instance.
+/// There can only be one worker registered per namespace+queue_name+client, others will get ignored.
+/// It also provides a convenient method to find compatible slots within the collection.
+#[derive(Default, Debug)]
+pub struct SlotManager {
+    manager: RwLock<SlotManagerImpl>,
+}
+
+impl SlotManager {
+    /// Factory method.
+    pub fn new() -> Self {
+        Self {
+            manager: RwLock::new(SlotManagerImpl::new()),
+        }
+    }
+
+    /// Try to reserve a compatible processing slot in any of the registered workers.
+    pub(crate) fn try_reserve_wft_slot(
+        &self,
+        namespace: String,
+        task_queue: String,
+    ) -> Option<Box<dyn Slot + Send>> {
+        self.manager
+            .read()
+            .try_reserve_wft_slot(namespace, task_queue)
+    }
+
+    /// Register a local worker that can provide WFT processing slots.
+    pub fn register(&self, provider: Box<dyn SlotProvider + Send + Sync>) -> Option<WorkerKey> {
+        self.manager.write().register(provider)
+    }
+
+    /// Unregister a provider, typically when its worker starts shutdown.
+    pub fn unregister(&self, id: WorkerKey) {
+        self.manager.write().unregister(id)
+    }
+
+    #[cfg(test)]
+    /// Returns (num_providers, num_buckets), where a bucket key is namespace+task_queue.
+    /// There is only one provider per bucket so `num_providers` should be equal to `num_buckets`.
+    pub fn num_providers(&self) -> (usize, usize) {
+        self.manager.read().num_providers()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn new_mock_slot(with_error: bool) -> Box<MockSlot> {
+        let mut mock_slot = MockSlot::new();
+        if with_error {
+            mock_slot
+                .expect_schedule_wft()
+                .returning(|_| Err(anyhow::anyhow!("Changed my mind")));
+        } else {
+            mock_slot.expect_schedule_wft().returning(|_| Ok(()));
+        }
+        Box::new(mock_slot)
+    }
+
+    fn new_mock_provider(
+        namespace: String,
+        task_queue: String,
+        with_error: bool,
+        no_slots: bool,
+    ) -> MockSlotProvider {
+        let mut mock_provider = MockSlotProvider::new();
+        mock_provider
+            .expect_try_reserve_wft_slot()
+            .returning(move || {
+                if no_slots {
+                    None
+                } else {
+                    Some(new_mock_slot(with_error))
+                }
+            });
+        mock_provider.expect_namespace().return_const(namespace);
+        mock_provider.expect_task_queue().return_const(task_queue);
+        mock_provider
+    }
+
+    #[test]
+    fn registry_respects_registration_order() {
+        let mock_provider1 =
+            new_mock_provider("foo".to_string(), "bar_q".to_string(), false, false);
+        let mock_provider2 = new_mock_provider("foo".to_string(), "bar_q".to_string(), false, true);
+
+        let manager = SlotManager::new();
+        let some_slots = manager.register(Box::new(mock_provider1));
+        let no_slots = manager.register(Box::new(mock_provider2));
+        assert!(no_slots.is_none());
+
+        let mut found = 0;
+        for _ in 0..10 {
+            if manager
+                .try_reserve_wft_slot("foo".to_string(), "bar_q".to_string())
+                .is_some()
+            {
+                found += 1;
+            }
+        }
+        assert_eq!(found, 10);
+        assert_eq!((1, 1), manager.num_providers());
+
+        manager.unregister(some_slots.unwrap());
+        assert_eq!((0, 0), manager.num_providers());
+
+        let mock_provider1 =
+            new_mock_provider("foo".to_string(), "bar_q".to_string(), false, false);
+        let mock_provider2 = new_mock_provider("foo".to_string(), "bar_q".to_string(), false, true);
+
+        let no_slots = manager.register(Box::new(mock_provider2));
+        let some_slots = manager.register(Box::new(mock_provider1));
+        assert!(some_slots.is_none());
+
+        let mut not_found = 0;
+        for _ in 0..10 {
+            if manager
+                .try_reserve_wft_slot("foo".to_string(), "bar_q".to_string())
+                .is_none()
+            {
+                not_found += 1;
+            }
+        }
+        assert_eq!(not_found, 10);
+        assert_eq!((1, 1), manager.num_providers());
+        manager.unregister(no_slots.unwrap());
+        assert_eq!((0, 0), manager.num_providers());
+    }
+
+    #[test]
+    fn registry_keeps_one_provider_per_namespace() {
+        let manager = SlotManager::new();
+        let mut worker_keys = vec![];
+        for i in 0..10 {
+            let namespace = format!("myId{}", i % 3);
+            let mock_provider = new_mock_provider(namespace, "bar_q".to_string(), false, false);
+            worker_keys.push(manager.register(Box::new(mock_provider)));
+        }
+        assert_eq!((3, 3), manager.num_providers());
+
+        let count = worker_keys
+            .iter()
+            .filter(|key| key.is_some())
+            .fold(0, |count, key| {
+                manager.unregister(key.unwrap());
+                // Should be idempotent
+                manager.unregister(key.unwrap());
+                count + 1
+            });
+        assert_eq!(3, count);
+        assert_eq!((0, 0), manager.num_providers());
+    }
+}

package/sdk-core/client/src/workflow_handle/mod.rs

@@ -54,6 +54,7 @@ pub struct WorkflowHandle<ClientT, ResultT> {
 }
 
 /// Holds needed information to refer to a specific workflow run, or workflow execution chain
+#[derive(Debug)]
 pub struct WorkflowExecutionInfo {
     /// Namespace the workflow lives in
     pub namespace: String,
@@ -100,7 +101,7 @@ where
             let server_res = self
                 .client
                 .clone()
-                .workflow_client()
+                .workflow_client_mut()
                 .get_workflow_execution_history(GetWorkflowExecutionHistoryRequest {
                     namespace: self.info.namespace.to_string(),
                     execution: Some(WorkflowExecution {

package/sdk-core/core/Cargo.toml

@@ -28,7 +28,7 @@ async-trait = "0.1"
 base64 = "0.21"
 console-subscriber = { version = "0.1", optional = true }
 crossbeam = "0.8"
-dashmap = "5.0"
+dashmap = "5.5"
 derive_builder = "0.12"
 derive_more = "0.99"
 enum_dispatch = "0.3"
@@ -36,18 +36,19 @@ enum-iterator = "1.4"
 flate2 = { version = "1.0", optional = true }
 futures = "0.3"
 futures-util = "0.3"
-governor = "0.5"
+governor = "0.6"
 http = "0.2"
 hyper = "0.14"
-itertools = "0.10"
+itertools = "0.11"
 lazy_static = "1.4"
-lru = "0.10"
+lru = "0.11"
 mockall = "0.11"
-nix = { version = "0.26", optional = true }
+nix = { version = "0.27", optional = true, features = ["process", "signal"] }
 once_cell = "1.5"
-opentelemetry = { version = "0.18", features = ["rt-tokio"] }
-opentelemetry-otlp = { version = "0.11", features = ["tokio", "metrics"] }
-opentelemetry-prometheus = "0.11"
+opentelemetry = { workspace = true, features = ["metrics"] }
+opentelemetry_sdk = { version = "0.21", features = ["rt-tokio", "metrics"] }
+opentelemetry-otlp = { version = "0.14", features = ["tokio", "metrics"] }
+opentelemetry-prometheus = "0.14"
 parking_lot = { version = "0.12", features = ["send_guard"] }
 pin-project = "1.0"
 prometheus = "0.13"
@@ -59,51 +60,48 @@ ringbuf = "0.3"
 rmp-serde = { version = "1.1", optional = true }
 serde = "1.0"
 serde_json = "1.0"
-siphasher = "0.3"
+siphasher = "1.0"
 slotmap = "1.0"
 tar = { version = "0.4", optional = true }
 thiserror = "1.0"
 tokio = { version = "1.26", features = ["rt", "rt-multi-thread", "parking_lot", "time", "fs", "process"] }
 tokio-util = { version = "0.7", features = ["io", "io-util"] }
 tokio-stream = "0.1"
-tonic = { version = "0.8", features = ["tls", "tls-roots"] }
+tonic = { workspace = true, features = ["tls", "tls-roots"] }
 tracing = "0.1"
 tracing-futures = "0.2"
-tracing-opentelemetry = "0.18"
 tracing-subscriber = { version = "0.3", features = ["parking_lot", "env-filter", "registry"] }
 url = "2.2"
 uuid = { version = "1.1", features = ["v4"] }
 zip = { version = "0.6.3", optional = true }
+log = "0.4.20"
 
 # 1st party local deps
 [dependencies.temporal-sdk-core-api]
 path = "../core-api"
-version = "0.1"
+features = ["otel_impls"]
 
 [dependencies.temporal-sdk-core-protos]
 path = "../sdk-core-protos"
-version = "0.1"
 features = ["history_builders"]
 
 [dependencies.temporal-client]
 path = "../client"
-version = "0.1"
 
 [dependencies.rustfsm]
 path = "../fsm"
-version = "0.1"
 
 [dev-dependencies]
 assert_matches = "1.4"
 bimap = "0.6.1"
 clap = { version = "4.0", features = ["derive"] }
-criterion = "0.4"
-rstest = "0.17"
+criterion = "0.5"
+rstest = "0.18"
 temporal-sdk-core-test-utils = { path = "../test-utils" }
 temporal-sdk = { path = "../sdk" }
 
 [build-dependencies]
-tonic-build = "0.8"
+tonic-build = { workspace = true }
 
 [[test]]
 name = "integ_tests"

package/sdk-core/core/src/core_tests/child_workflows.rs

@@ -1,12 +1,13 @@
 use crate::{
     replay::DEFAULT_WORKFLOW_TYPE,
     test_help::{
-        canned_histories, mock_sdk, mock_worker, single_hist_mock_sg, MockPollCfg, ResponseType,
+        build_fake_sdk, canned_histories, mock_sdk, mock_worker, single_hist_mock_sg, MockPollCfg,
+        ResponseType,
     },
-    worker::{client::mocks::mock_workflow_client, ManagedWFFunc},
+    worker::client::mocks::mock_workflow_client,
 };
 use temporal_client::WorkflowOptions;
-use temporal_sdk::{ChildWorkflowOptions, Signal, WfContext, WorkflowFunction, WorkflowResult};
+use temporal_sdk::{ChildWorkflowOptions, Signal, WfContext, WorkflowResult};
 use temporal_sdk_core_api::Worker;
 use temporal_sdk_core_protos::coresdk::{
     child_workflow::{child_workflow_result, ChildWorkflowCancellationType},
@@ -100,11 +101,10 @@ async fn parent_cancels_child_wf(ctx: WfContext) -> WorkflowResult<()> {
 
 #[tokio::test]
 async fn cancel_child_workflow() {
-    let func = WorkflowFunction::new(parent_cancels_child_wf);
     let t = canned_histories::single_child_workflow_cancelled("child-id-1");
-    let mut wfm = ManagedWFFunc::new(t, func, vec![]);
-    wfm.process_all_activations().await.unwrap();
-    wfm.shutdown().await.unwrap();
+    let mut worker = build_fake_sdk(MockPollCfg::from_resps(t, [ResponseType::AllHistory]));
+    worker.register_wf(DEFAULT_WORKFLOW_TYPE, parent_cancels_child_wf);
+    worker.run().await.unwrap();
 }
 
 #[rstest::rstest]

package/sdk-core/core/src/core_tests/mod.rs

@@ -4,6 +4,7 @@ mod determinism;
 mod local_activities;
 mod queries;
 mod replay_flag;
+mod updates;
 mod workers;
 mod workflow_cancels;
 mod workflow_tasks;

package/sdk-core/core/src/core_tests/replay_flag.rs

@@ -1,26 +1,25 @@
 use crate::{
     test_help::{
-        build_mock_pollers, canned_histories, hist_to_poll_resp, mock_worker, MockPollCfg,
-        ResponseType,
+        build_fake_sdk, build_mock_pollers, canned_histories, hist_to_poll_resp, mock_worker,
+        MockPollCfg, ResponseType,
     },
-    worker::{client::mocks::mock_workflow_client, ManagedWFFunc, LEGACY_QUERY_ID},
+    worker::{client::mocks::mock_workflow_client, LEGACY_QUERY_ID},
 };
 use rstest::{fixture, rstest};
 use std::{collections::VecDeque, time::Duration};
-use temporal_sdk::{WfContext, WorkflowFunction};
-use temporal_sdk_core_api::Worker;
+use temporal_sdk::{WfContext, Worker, WorkflowFunction};
+use temporal_sdk_core_api::Worker as CoreWorker;
 use temporal_sdk_core_protos::{
     coresdk::{
         workflow_activation::{workflow_activation_job, WorkflowActivationJob},
         workflow_completion::WorkflowActivationCompletion,
     },
-    temporal::api::{
-        enums::v1::{CommandType, EventType},
-        query::v1::WorkflowQuery,
-    },
-    TestHistoryBuilder,
+    temporal::api::{enums::v1::EventType, query::v1::WorkflowQuery},
+    TestHistoryBuilder, DEFAULT_WORKFLOW_TYPE,
+};
+use temporal_sdk_core_test_utils::{
+    interceptors::ActivationAssertionsInterceptor, query_ok, start_timer_cmd,
 };
-use temporal_sdk_core_test_utils::{query_ok, start_timer_cmd};
 
 fn timers_wf(num_timers: u32) -> WorkflowFunction {
     WorkflowFunction::new(move |command_sink: WfContext| async move {
@@ -32,38 +31,31 @@ fn timers_wf(num_timers: u32) -> WorkflowFunction {
 }
 
 #[fixture(num_timers = 1)]
-fn fire_happy_hist(num_timers: u32) -> ManagedWFFunc {
+fn fire_happy_hist(num_timers: u32) -> Worker {
     let func = timers_wf(num_timers);
     // Add 1 b/c history takes # wf tasks, not timers
     let t = canned_histories::long_sequential_timers(num_timers as usize);
-    ManagedWFFunc::new(t, func, vec![])
+    let mut worker = build_fake_sdk(MockPollCfg::from_resps(t, [ResponseType::AllHistory]));
+    worker.register_wf(DEFAULT_WORKFLOW_TYPE, func);
+    worker
 }
 
 #[rstest]
 #[case::one_timer(fire_happy_hist(1), 1)]
 #[case::five_timers(fire_happy_hist(5), 5)]
 #[tokio::test]
-async fn replay_flag_is_correct(#[case] mut wfm: ManagedWFFunc, #[case] num_timers: usize) {
+async fn replay_flag_is_correct(#[case] mut worker: Worker, #[case] num_timers: usize) {
     // Verify replay flag is correct by constructing a workflow manager that already has a complete
     // history fed into it. It should always be replaying, because history is complete.
 
-    for _ in 1..=num_timers {
-        let act = wfm.get_next_activation().await.unwrap();
-        assert!(act.is_replaying);
-        let commands = wfm.get_server_commands().commands;
-        assert_eq!(commands.len(), 1);
-        assert_eq!(commands[0].command_type, CommandType::StartTimer as i32);
+    let mut aai = ActivationAssertionsInterceptor::default();
+
+    for _ in 1..=num_timers + 1 {
+        aai.then(|a| assert!(a.is_replaying));
     }
 
-    let act = wfm.get_next_activation().await.unwrap();
-    assert!(act.is_replaying);
-    let commands = wfm.get_server_commands().commands;
-    assert_eq!(commands.len(), 1);
-    assert_eq!(
-        commands[0].command_type,
-        CommandType::CompleteWorkflowExecution as i32
-    );
-    wfm.shutdown().await.unwrap();
+    worker.set_worker_interceptor(aai);
+    worker.run().await.unwrap();
 }
 
 #[tokio::test(flavor = "multi_thread")]
@@ -71,15 +63,14 @@ async fn replay_flag_is_correct_partial_history() {
     let func = timers_wf(1);
     // Add 1 b/c history takes # wf tasks, not timers
     let t = canned_histories::long_sequential_timers(2);
-    let mut wfm =
-        ManagedWFFunc::new_from_update(t.get_history_info(1).unwrap().into(), func, vec![]);
-
-    let act = wfm.get_next_activation().await.unwrap();
-    assert!(!act.is_replaying);
-    let commands = wfm.get_server_commands().commands;
-    assert_eq!(commands.len(), 1);
-    assert_eq!(commands[0].command_type, CommandType::StartTimer as i32);
-    wfm.shutdown().await.unwrap();
+    let mut worker = build_fake_sdk(MockPollCfg::from_resps(t, [1]));
+    worker.register_wf(DEFAULT_WORKFLOW_TYPE, func);
+
+    let mut aai = ActivationAssertionsInterceptor::default();
+    aai.then(|a| assert!(!a.is_replaying));
+
+    worker.set_worker_interceptor(aai);
+    worker.run().await.unwrap();
 }
 
 #[tokio::test]
@@ -99,7 +90,6 @@ async fn replay_flag_correct_with_query() {
             let h = pr.history.as_mut().unwrap();
             h.events.truncate(5);
             pr.started_event_id = 3;
-            dbg!(&pr.resp);
             pr
         },
         hist_to_poll_resp(&t, wfid.to_owned(), 2.into()),

package/sdk-core/core/src/core_tests/updates.rs

@@ -0,0 +1,73 @@
+use crate::test_help::{build_mock_pollers, mock_worker, MockPollCfg, ResponseType};
+use temporal_sdk_core_api::Worker;
+use temporal_sdk_core_protos::{
+    coresdk::{
+        workflow_activation::{workflow_activation_job, WorkflowActivationJob},
+        workflow_commands::{update_response::Response, CompleteWorkflowExecution, UpdateResponse},
+        workflow_completion::WorkflowActivationCompletion,
+    },
+    temporal::api::{common::v1::Payload, enums::v1::EventType},
+    TestHistoryBuilder,
+};
+
+#[tokio::test]
+async fn replay_with_empty_first_task() {
+    let mut t = TestHistoryBuilder::default();
+    t.add_by_type(EventType::WorkflowExecutionStarted);
+    t.add_full_wf_task();
+    t.add_full_wf_task();
+    let accept_id = t.add_update_accepted("upd1", "update");
+    t.add_we_signaled("hi", vec![]);
+    t.add_full_wf_task();
+    t.add_update_completed(accept_id);
+    t.add_workflow_execution_completed();
+
+    let mock = MockPollCfg::from_resps(t, [ResponseType::AllHistory]);
+    let mut mock = build_mock_pollers(mock);
+    mock.worker_cfg(|wc| wc.max_cached_workflows = 1);
+    let core = mock_worker(mock);
+
+    let task = core.poll_workflow_activation().await.unwrap();
+    assert_matches!(
+        task.jobs.as_slice(),
+        [
+            WorkflowActivationJob {
+                variant: Some(workflow_activation_job::Variant::DoUpdate(_)),
+            },
+            WorkflowActivationJob {
+                variant: Some(workflow_activation_job::Variant::StartWorkflow(_)),
+            },
+        ]
+    );
+    core.complete_workflow_activation(WorkflowActivationCompletion::from_cmd(
+        task.run_id,
+        UpdateResponse {
+            protocol_instance_id: "upd1".to_string(),
+            response: Some(Response::Accepted(())),
+        }
+        .into(),
+    ))
+    .await
+    .unwrap();
+
+    let task = core.poll_workflow_activation().await.unwrap();
+    assert_matches!(
+        task.jobs.as_slice(),
+        [WorkflowActivationJob {
+            variant: Some(workflow_activation_job::Variant::SignalWorkflow(_)),
+        }]
+    );
+    core.complete_workflow_activation(WorkflowActivationCompletion::from_cmds(
+        task.run_id,
+        vec![
+            UpdateResponse {
+                protocol_instance_id: "upd1".to_string(),
+                response: Some(Response::Completed(Payload::default())),
+            }
+            .into(),
+            CompleteWorkflowExecution { result: None }.into(),
+        ],
+    ))
+    .await
+    .unwrap();
+}

package/sdk-core/core/src/core_tests/workflow_tasks.rs

@@ -45,7 +45,7 @@ use temporal_sdk_core_protos::{
     default_act_sched, default_wes_attribs,
     temporal::api::{
         command::v1::command::Attributes,
-        common::v1::{Payload, RetryPolicy},
+        common::v1::{Payload, RetryPolicy, WorkerVersionStamp},
         enums::v1::{EventType, WorkflowTaskFailedCause},
         failure::v1::Failure,
         history::v1::{
@@ -2807,3 +2807,54 @@ async fn use_compatible_version_flag(
         .await
         .unwrap();
 }
+
+#[tokio::test]
+async fn sets_build_id_from_wft_complete() {
+    let wfid = "fake_wf_id";
+
+    let mut t = TestHistoryBuilder::default();
+    t.add_by_type(EventType::WorkflowExecutionStarted);
+    t.add_full_wf_task();
+    let timer_started_event_id = t.add_by_type(EventType::TimerStarted);
+    t.add_timer_fired(timer_started_event_id, "1".to_string());
+    t.add_full_wf_task();
+    t.modify_event(t.current_event_id(), |he| {
+        if let history_event::Attributes::WorkflowTaskCompletedEventAttributes(a) =
+            he.attributes.as_mut().unwrap()
+        {
+            a.worker_version = Some(WorkerVersionStamp {
+                build_id: "enchi-cat".to_string(),
+                ..Default::default()
+            });
+        }
+    });
+    let timer_started_event_id = t.add_by_type(EventType::TimerStarted);
+    t.add_timer_fired(timer_started_event_id, "2".to_string());
+    t.add_full_wf_task();
+
+    let mock = mock_workflow_client();
+    let mut worker = mock_sdk_cfg(
+        MockPollCfg::from_resp_batches(wfid, t, [ResponseType::AllHistory], mock),
+        |cfg| {
+            cfg.worker_build_id = "fierce-predator".to_string();
+            cfg.max_cached_workflows = 1;
+        },
+    );
+
+    worker.register_wf(DEFAULT_WORKFLOW_TYPE, |ctx: WfContext| async move {
+        // First task, it should be empty, since replaying and nothing in first WFT completed
+        ctx.timer(Duration::from_secs(1)).await;
+        assert_eq!(ctx.current_build_id(), None);
+        ctx.timer(Duration::from_secs(1)).await;
+        assert_eq!(ctx.current_build_id(), Some("enchi-cat".to_string()));
+        ctx.timer(Duration::from_secs(1)).await;
+        // Not replaying at this point, so we should see the worker's build id
+        assert_eq!(ctx.current_build_id(), Some("fierce-predator".to_string()));
+        Ok(().into())
+    });
+    worker
+        .submit_wf(wfid, DEFAULT_WORKFLOW_TYPE, vec![], Default::default())
+        .await
+        .unwrap();
+    worker.run_until_done().await.unwrap();
+}