@temporalio/core-bridge 0.21.1 → 0.22.0

package/Cargo.lock

@@ -1568,6 +1568,7 @@ dependencies = [
 name = "temporal-client"
 version = "0.1.0"
 dependencies = [
+ "anyhow",
  "async-trait",
  "backoff",
  "derive_builder",

package/index.d.ts

@@ -66,9 +66,13 @@ export interface TelemetryOptions {
    * Which determines what tracing data is collected in the Core SDK
    */
   tracingFilter?: string;
-  /** What level, if any, logs should be forwarded from core at */
+  /**
+   * What level, if any, logs should be forwarded from core at
+   *
+   * @default OFF
+   */
   // These strings should match the log::LevelFilter enum in rust
-  logForwardingLevel: 'OFF' | LogLevel;
+  logForwardingLevel?: 'OFF' | LogLevel;
   /** If set, metrics will be exposed on an http server in this process for direct scraping by
    *  prometheus. If used in conjunction with the OTel collector, metrics will *not* be exported
    *  to the collector, but traces will be.

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@temporalio/core-bridge",
-  "version": "0.21.1",
+  "version": "0.22.0",
   "description": "Temporal.io SDK Core<>Node bridge",
   "main": "index.js",
   "types": "index.d.ts",
@@ -20,7 +20,7 @@
   "license": "MIT",
   "dependencies": {
     "@opentelemetry/api": "^1.0.3",
-    "@temporalio/internal-non-workflow-common": "^0.21.0",
+    "@temporalio/internal-non-workflow-common": "^0.22.0",
     "arg": "^5.0.1",
     "cargo-cp-artifact": "^0.1.4",
     "which": "^2.0.2"
@@ -43,5 +43,5 @@
   "publishConfig": {
     "access": "public"
   },
-  "gitHead": "c241ab2f8e92966555dbd6d01dc072c8965ee6f4"
+  "gitHead": "3aa1f14982bd170d21b728cbf016dc4f1b595a76"
 }

package/sdk-core/.buildkite/docker/docker-compose.yaml

@@ -9,7 +9,7 @@ services:
   #      - '9042:9042'
 
   temporal:
-    image: temporalio/auto-setup:1.15.0
+    image: temporalio/auto-setup:1.16.0
     ports:
       - "7233:7233"
       - "7234:7234"
@@ -26,12 +26,13 @@ services:
       - cassandra
 
   temporal-web:
-    image: temporalio/web:1.13.0
+    image: temporalio/ui:0.10.2
     logging:
       driver: none
     ports:
-      - "8088:8088"
+      - "8080:8080"
     environment:
-      - "TEMPORAL_GRPC_ENDPOINT=temporal:7233"
+      - TEMPORAL_ADDRESS=temporal:7233
+      - TEMPORAL_CORS_ORIGINS=http://localhost:3000
     depends_on:
       - temporal

package/sdk-core/client/Cargo.toml

@@ -11,6 +11,7 @@ keywords = ["temporal", "workflow"]
 categories = ["development-tools"]
 
 [dependencies]
+anyhow = "1.0"
 async-trait = "0.1"
 backoff = "0.4"
 derive_builder = "0.11"

package/sdk-core/client/src/lib.rs

@@ -10,13 +10,17 @@ extern crate tracing;
 mod metrics;
 mod raw;
 mod retry;
+mod workflow_handle;
 
 pub use crate::retry::{CallType, RetryClient};
 pub use raw::WorkflowService;
+pub use workflow_handle::{WorkflowExecutionInfo, WorkflowExecutionResult};
 
 use crate::{
     metrics::{GrpcMetricSvc, MetricsContext},
     raw::{sealed::RawClientLike, AttachMetricLabels},
+    sealed::{RawClientLikeUser, WfHandleClient},
+    workflow_handle::UntypedWorkflowHandle,
 };
 use backoff::{ExponentialBackoff, SystemClock};
 use http::uri::InvalidUri;
@@ -37,7 +41,7 @@ use temporal_sdk_core_protos::{
         enums::v1::{TaskQueueKind, WorkflowTaskFailedCause},
         failure::v1::Failure,
         query::v1::{WorkflowQuery, WorkflowQueryResult},
-        taskqueue::v1::{StickyExecutionAttributes, TaskQueue},
+        taskqueue::v1::{StickyExecutionAttributes, TaskQueue, TaskQueueMetadata},
         workflowservice::v1::{workflow_service_client::WorkflowServiceClient, *},
     },
     TaskToken,
@@ -53,8 +57,6 @@ use tower::ServiceBuilder;
 use url::Url;
 use uuid::Uuid;
 
-use temporal_sdk_core_protos::temporal::api::taskqueue::v1::TaskQueueMetadata;
-
 static LONG_POLL_METHOD_NAMES: [&str; 2] = ["PollWorkflowTaskQueue", "PollActivityTaskQueue"];
 /// The server times out polls after 60 seconds. Set our timeout to be slightly beyond that.
 const LONG_POLL_TIMEOUT: Duration = Duration::from_secs(70);
@@ -431,7 +433,7 @@ pub type WorkflowServiceClientWithMetrics = WorkflowServiceClient<InterceptedMet
 type InterceptedMetricsSvc = InterceptedService<GrpcMetricSvc, ServiceCallInterceptor>;
 
 /// Contains an instance of a namespace-bound client for interacting with the Temporal server
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct Client {
     /// Client for interacting with workflow service
     inner: ConfiguredClient<WorkflowServiceClientWithMetrics>,
@@ -475,11 +477,6 @@ impl Client {
     pub fn options(&self) -> &ClientOptions {
         &self.inner.options
     }
-
-    /// Used to access the client as a [WorkflowService] implementor rather than the raw struct
-    fn wf_svc(&self) -> impl RawClientLike<SvcType = InterceptedMetricsSvc> {
-        self.raw_client().clone()
-    }
 }
 
 /// This trait provides higher-level friendlier interaction with the server.
@@ -1036,3 +1033,49 @@ impl WorkflowClientTrait for Client {
         &self.namespace
     }
 }
+
+mod sealed {
+    use crate::{InterceptedMetricsSvc, RawClientLike, WorkflowClientTrait};
+
+    pub trait RawClientLikeUser {
+        type RawClientT: RawClientLike<SvcType = InterceptedMetricsSvc>;
+        /// Used to access the client as a [WorkflowService] implementor rather than the raw struct
+        fn wf_svc(&self) -> Self::RawClientT;
+    }
+
+    pub trait WfHandleClient: WorkflowClientTrait + RawClientLikeUser {}
+    impl<T> WfHandleClient for T where T: WorkflowClientTrait + RawClientLikeUser {}
+}
+
+impl RawClientLikeUser for Client {
+    type RawClientT = WorkflowServiceClientWithMetrics;
+
+    fn wf_svc(&self) -> Self::RawClientT {
+        self.raw_client().clone()
+    }
+}
+
+/// Additional methods for workflow clients
+pub trait WfClientExt: WfHandleClient + Sized {
+    /// Create an untyped handle for a workflow execution, which can be used to do things like
+    /// wait for that workflow's result. `run_id` may be left blank to target the latest run.
+    fn get_untyped_workflow_handle(
+        &self,
+        workflow_id: impl Into<String>,
+        run_id: impl Into<String>,
+    ) -> UntypedWorkflowHandle<Self::RawClientT>
+    where
+        Self::RawClientT: Clone,
+    {
+        let rid = run_id.into();
+        UntypedWorkflowHandle::new(
+            self.wf_svc(),
+            WorkflowExecutionInfo {
+                namespace: self.namespace().to_string(),
+                workflow_id: workflow_id.into(),
+                run_id: if rid.is_empty() { None } else { Some(rid) },
+            },
+        )
+    }
+}
+impl<T> WfClientExt for T where T: WfHandleClient + Sized {}

package/sdk-core/client/src/raw.rs

@@ -15,7 +15,7 @@ use tonic::{body::BoxBody, client::GrpcService, metadata::KeyAndValueRef};
 
 pub(super) mod sealed {
     use super::*;
-    use crate::{ConfiguredClient, RetryClient};
+    use crate::{Client, ConfiguredClient, InterceptedMetricsSvc, RetryClient};
     use futures::TryFutureExt;
     use tonic::{Request, Response, Status};
 
@@ -104,6 +104,14 @@ pub(super) mod sealed {
             &mut self.client
         }
     }
+
+    impl RawClientLike for Client {
+        type SvcType = InterceptedMetricsSvc;
+
+        fn client(&mut self) -> &mut WorkflowServiceClient<Self::SvcType> {
+            &mut self.inner
+        }
+    }
 }
 
 /// Helper for cloning a tonic request as long as the inner message may be cloned.

package/sdk-core/client/src/retry.rs

@@ -1,5 +1,5 @@
 use crate::{
-    ClientOptions, Result, RetryConfig, WorkflowClientTrait, WorkflowOptions,
+    ClientOptions, RawClientLikeUser, Result, RetryConfig, WorkflowClientTrait, WorkflowOptions,
     WorkflowTaskCompletion,
 };
 use backoff::{backoff::Backoff, ExponentialBackoff};
@@ -592,3 +592,14 @@ mod tests {
         }
     }
 }
+
+impl<C> RawClientLikeUser for RetryClient<C>
+where
+    C: RawClientLikeUser,
+{
+    type RawClientT = C::RawClientT;
+
+    fn wf_svc(&self) -> Self::RawClientT {
+        self.client.wf_svc()
+    }
+}

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

@@ -0,0 +1,183 @@
+use crate::{InterceptedMetricsSvc, RawClientLike};
+use anyhow::{anyhow, bail};
+use std::marker::PhantomData;
+use temporal_sdk_core_protos::{
+    coresdk::{common::Payload, FromPayloadsExt},
+    temporal::api::{
+        common::v1::WorkflowExecution, enums::v1::HistoryEventFilterType, failure::v1::Failure,
+        history::v1::history_event::Attributes,
+        workflowservice::v1::GetWorkflowExecutionHistoryRequest,
+    },
+};
+
+/// Enumerates terminal states for a particular workflow execution
+// TODO: Add non-proto failure types, flesh out details, etc.
+#[derive(Debug)]
+#[allow(clippy::large_enum_variant)]
+pub enum WorkflowExecutionResult<T> {
+    /// The workflow finished successfully
+    Succeeded(T),
+    /// The workflow finished in failure
+    Failed(Failure),
+    /// The workflow was cancelled
+    Cancelled(Vec<Payload>),
+    /// The workflow was terminated
+    Terminated(Vec<Payload>),
+    /// The workflow timed out
+    TimedOut,
+    /// The workflow continued as new
+    ContinuedAsNew,
+}
+
+/// Options for fetching workflow results
+#[derive(Debug, Clone, Copy)]
+pub struct GetWorkflowResultOpts {
+    /// If true (the default), follows to the next workflow run in the execution chain while
+    /// retrieving results.
+    pub follow_runs: bool,
+}
+impl Default for GetWorkflowResultOpts {
+    fn default() -> Self {
+        Self { follow_runs: true }
+    }
+}
+
+/// A workflow handle which can refer to a specific workflow run, or a chain of workflow runs with
+/// the same workflow id.
+pub struct WorkflowHandle<ClientT, ResultT> {
+    client: ClientT,
+    info: WorkflowExecutionInfo,
+
+    _res_type: PhantomData<ResultT>,
+}
+
+/// Holds needed information to refer to a specific workflow run, or workflow execution chain
+pub struct WorkflowExecutionInfo {
+    /// Namespace the workflow lives in
+    pub namespace: String,
+    /// The workflow's id
+    pub workflow_id: String,
+    /// If set, target this specific run of the workflow
+    pub run_id: Option<String>,
+}
+
+impl WorkflowExecutionInfo {
+    /// Bind the workflow info to a specific client, turning it into a workflow handle
+    pub fn bind_untyped<CT>(self, client: CT) -> UntypedWorkflowHandle<CT>
+    where
+        CT: RawClientLike<SvcType = InterceptedMetricsSvc> + Clone,
+    {
+        UntypedWorkflowHandle::new(client, self)
+    }
+}
+
+/// A workflow handle to a workflow with unknown types. Uses raw payloads.
+pub type UntypedWorkflowHandle<CT> = WorkflowHandle<CT, Vec<Payload>>;
+
+impl<CT, RT> WorkflowHandle<CT, RT>
+where
+    CT: RawClientLike<SvcType = InterceptedMetricsSvc> + Clone,
+    // TODO: Make more generic, capable of (de)serialization w/ serde
+    RT: FromPayloadsExt,
+{
+    pub(crate) fn new(client: CT, info: WorkflowExecutionInfo) -> Self {
+        Self {
+            client,
+            info,
+            _res_type: PhantomData::<RT>,
+        }
+    }
+
+    pub async fn get_workflow_result(
+        &self,
+        opts: GetWorkflowResultOpts,
+    ) -> Result<WorkflowExecutionResult<RT>, anyhow::Error> {
+        let mut next_page_tok = vec![];
+        let mut run_id = self.info.run_id.clone().unwrap_or_default();
+        loop {
+            let server_res = self
+                .client
+                .clone()
+                .client()
+                .get_workflow_execution_history(GetWorkflowExecutionHistoryRequest {
+                    namespace: self.info.namespace.to_string(),
+                    execution: Some(WorkflowExecution {
+                        workflow_id: self.info.workflow_id.clone(),
+                        run_id: run_id.clone(),
+                    }),
+                    skip_archival: true,
+                    wait_new_event: true,
+                    history_event_filter_type: HistoryEventFilterType::CloseEvent as i32,
+                    next_page_token: next_page_tok.clone(),
+                    ..Default::default()
+                })
+                .await?
+                .into_inner();
+
+            let mut history = server_res
+                .history
+                .ok_or_else(|| anyhow!("Server returned an empty history!"))?;
+
+            if history.events.is_empty() {
+                next_page_tok = server_res.next_page_token;
+                continue;
+            }
+            // If page token was previously set, clear it.
+            next_page_tok = vec![];
+
+            let event_attrs = history.events.pop().and_then(|ev| ev.attributes);
+
+            macro_rules! follow {
+                ($attrs:ident) => {
+                    if opts.follow_runs && $attrs.new_execution_run_id != "" {
+                        run_id = $attrs.new_execution_run_id;
+                        continue;
+                    }
+                };
+            }
+
+            break match event_attrs {
+                Some(Attributes::WorkflowExecutionCompletedEventAttributes(attrs)) => {
+                    follow!(attrs);
+                    Ok(WorkflowExecutionResult::Succeeded(RT::from_payloads(
+                        attrs.result,
+                    )))
+                }
+                Some(Attributes::WorkflowExecutionFailedEventAttributes(attrs)) => {
+                    follow!(attrs);
+                    Ok(WorkflowExecutionResult::Failed(
+                        attrs.failure.unwrap_or_default(),
+                    ))
+                }
+                Some(Attributes::WorkflowExecutionCanceledEventAttributes(attrs)) => Ok(
+                    WorkflowExecutionResult::Cancelled(Vec::from_payloads(attrs.details)),
+                ),
+                Some(Attributes::WorkflowExecutionTimedOutEventAttributes(attrs)) => {
+                    follow!(attrs);
+                    Ok(WorkflowExecutionResult::TimedOut)
+                }
+                Some(Attributes::WorkflowExecutionTerminatedEventAttributes(attrs)) => Ok(
+                    WorkflowExecutionResult::Terminated(Vec::from_payloads(attrs.details)),
+                ),
+                Some(Attributes::WorkflowExecutionContinuedAsNewEventAttributes(attrs)) => {
+                    if opts.follow_runs {
+                        if !attrs.new_execution_run_id.is_empty() {
+                            run_id = attrs.new_execution_run_id;
+                            continue;
+                        } else {
+                            bail!("New execution run id was empty in continue as new event!");
+                        }
+                    } else {
+                        Ok(WorkflowExecutionResult::ContinuedAsNew)
+                    }
+                }
+                o => Err(anyhow!(
+                    "Server returned an event that didn't match the CloseEvent filter. \
+                     This is either a server bug or a new event the SDK does not understand. \
+                     Event details: {:?}",
+                    o
+                )),
+            };
+        }
+    }
+}

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

@@ -1,18 +1,14 @@
 use crate::{
     replay::DEFAULT_WORKFLOW_TYPE,
-    test_help::{
-        build_mock_pollers, canned_histories, mock_worker, MockPollCfg, ResponseType, TEST_Q,
-    },
+    test_help::{canned_histories, mock_sdk, MockPollCfg, ResponseType},
     worker::client::mocks::mock_workflow_client,
     workflow::managed_wf::ManagedWFFunc,
 };
-use std::sync::Arc;
 use temporal_client::WorkflowOptions;
 use temporal_sdk::{ChildWorkflowOptions, Signal, WfContext, WorkflowFunction, WorkflowResult};
 use temporal_sdk_core_protos::coresdk::child_workflow::{
     child_workflow_result, ChildWorkflowCancellationType,
 };
-use temporal_sdk_core_test_utils::TestWorker;
 use tokio::join;
 
 const SIGNAME: &str = "SIGNAME";
@@ -26,10 +22,12 @@ async fn signal_child_workflow(#[case] serial: bool) {
     let wf_type = DEFAULT_WORKFLOW_TYPE;
     let t = canned_histories::single_child_workflow_signaled("child-id-1", SIGNAME);
     let mock = mock_workflow_client();
-    let mh = MockPollCfg::from_resp_batches(wf_id, t, [ResponseType::AllHistory], mock);
-    let mock = build_mock_pollers(mh);
-    let core = mock_worker(mock);
-    let mut worker = TestWorker::new(Arc::new(core), TEST_Q.to_string());
+    let mut worker = mock_sdk(MockPollCfg::from_resp_batches(
+        wf_id,
+        t,
+        [ResponseType::AllHistory],
+        mock,
+    ));
 
     let wf = move |ctx: WfContext| async move {
         let child = ctx.child_workflow(ChildWorkflowOptions {

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

@@ -1,21 +1,15 @@
 use crate::{
     replay::DEFAULT_WORKFLOW_TYPE,
-    test_help::{
-        build_mock_pollers, canned_histories, mock_worker, MockPollCfg, ResponseType, TEST_Q,
-    },
+    test_help::{canned_histories, mock_sdk, mock_sdk_cfg, MockPollCfg, ResponseType},
     worker::client::mocks::mock_workflow_client,
 };
 use std::{
-    sync::{
-        atomic::{AtomicBool, AtomicUsize, Ordering},
-        Arc,
-    },
+    sync::atomic::{AtomicBool, AtomicUsize, Ordering},
     time::Duration,
 };
 use temporal_client::WorkflowOptions;
 use temporal_sdk::{WfContext, WorkflowResult};
 use temporal_sdk_core_protos::temporal::api::enums::v1::WorkflowTaskFailedCause;
-use temporal_sdk_core_test_utils::TestWorker;
 
 static DID_FAIL: AtomicBool = AtomicBool::new(false);
 pub async fn timer_wf_fails_once(ctx: WfContext) -> WorkflowResult<()> {
@@ -32,7 +26,7 @@ pub async fn timer_wf_fails_once(ctx: WfContext) -> WorkflowResult<()> {
 /// Verifies that workflow panics (which in this case the Rust SDK turns into workflow activation
 /// failures) are turned into unspecified WFT failures.
 #[tokio::test]
-async fn test_wf_task_rejected_properly() {
+async fn test_panic_wf_task_rejected_properly() {
     let wf_id = "fakeid";
     let wf_type = DEFAULT_WORKFLOW_TYPE;
     let t = canned_histories::workflow_fails_with_failure_after_timer("1");
@@ -43,9 +37,7 @@ async fn test_wf_task_rejected_properly() {
     mh.num_expected_fails = Some(1);
     mh.expect_fail_wft_matcher =
         Box::new(|_, cause, _| matches!(cause, WorkflowTaskFailedCause::Unspecified));
-    let mock = build_mock_pollers(mh);
-    let core = mock_worker(mock);
-    let mut worker = TestWorker::new(Arc::new(core), TEST_Q.to_string());
+    let mut worker = mock_sdk(mh);
 
     worker.register_wf(wf_type.to_owned(), timer_wf_fails_once);
     worker
@@ -82,14 +74,11 @@ async fn test_wf_task_rejected_properly_due_to_nondeterminism(#[case] use_cache:
     mh.num_expected_fails = Some(1);
     mh.expect_fail_wft_matcher =
         Box::new(|_, cause, _| matches!(cause, WorkflowTaskFailedCause::NonDeterministicError));
-    let mut mock = build_mock_pollers(mh);
-    if use_cache {
-        mock.worker_cfg(|cfg| {
+    let mut worker = mock_sdk_cfg(mh, |cfg| {
+        if use_cache {
             cfg.max_cached_workflows = 2;
-        });
-    }
-    let core = mock_worker(mock);
-    let mut worker = TestWorker::new(Arc::new(core), TEST_Q.to_string());
+        }
+    });
 
     let started_count: &'static _ = Box::leak(Box::new(AtomicUsize::new(0)));
     worker.register_wf(wf_type.to_owned(), move |ctx: WfContext| async move {

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

@@ -1,27 +1,23 @@
 use crate::{
     replay::{default_wes_attribs, TestHistoryBuilder, DEFAULT_WORKFLOW_TYPE},
-    test_help::{build_mock_pollers, mock_worker, MockPollCfg, ResponseType, TEST_Q},
+    test_help::{mock_sdk, mock_sdk_cfg, MockPollCfg, ResponseType},
     worker::client::mocks::mock_workflow_client,
 };
 use anyhow::anyhow;
 use futures::future::join_all;
 use std::{
-    sync::{
-        atomic::{AtomicUsize, Ordering},
-        Arc,
-    },
+    sync::atomic::{AtomicUsize, Ordering},
     time::Duration,
 };
 use temporal_client::WorkflowOptions;
-use temporal_sdk::{LocalActivityOptions, WfContext, WorkflowResult};
+use temporal_sdk::{ActContext, LocalActivityOptions, WfContext, WorkflowResult};
 use temporal_sdk_core_protos::{
     coresdk::{common::RetryPolicy, AsJsonPayloadExt},
     temporal::api::{enums::v1::EventType, failure::v1::Failure},
 };
-use temporal_sdk_core_test_utils::TestWorker;
 use tokio::sync::Barrier;
 
-async fn echo(e: String) -> anyhow::Result<String> {
+async fn echo(_ctx: ActContext, e: String) -> anyhow::Result<String> {
     Ok(e)
 }
 
@@ -52,12 +48,11 @@ async fn local_act_two_wfts_before_marker(#[case] replay: bool, #[case] cached:
         vec![1.into(), 2.into(), ResponseType::AllHistory]
     };
     let mh = MockPollCfg::from_resp_batches(wf_id, t, resps, mock);
-    let mut mock = build_mock_pollers(mh);
-    if cached {
-        mock.worker_cfg(|cfg| cfg.max_cached_workflows = 1);
-    }
-    let core = mock_worker(mock);
-    let mut worker = TestWorker::new(Arc::new(core), TEST_Q.to_string());
+    let mut worker = mock_sdk_cfg(mh, |cfg| {
+        if cached {
+            cfg.max_cached_workflows = 1;
+        }
+    });
 
     worker.register_wf(
         DEFAULT_WORKFLOW_TYPE.to_owned(),
@@ -119,12 +114,13 @@ async fn local_act_many_concurrent() {
     let wf_id = "fakeid";
     let mock = mock_workflow_client();
     let mh = MockPollCfg::from_resp_batches(wf_id, t, [1, 2, 3], mock);
-    let mock = build_mock_pollers(mh);
-    let core = mock_worker(mock);
-    let mut worker = TestWorker::new(Arc::new(core), TEST_Q.to_string());
+    let mut worker = mock_sdk(mh);
 
     worker.register_wf(DEFAULT_WORKFLOW_TYPE.to_owned(), local_act_fanout_wf);
-    worker.register_activity("echo", |str: String| async move { Ok(str) });
+    worker.register_activity(
+        "echo",
+        |_ctx: ActContext, str: String| async move { Ok(str) },
+    );
     worker
         .submit_wf(
             wf_id.to_owned(),
@@ -167,10 +163,9 @@ async fn local_act_heartbeat(#[case] shutdown_middle: bool) {
     // and might poll an extra time
     let mut mh = MockPollCfg::from_resp_batches(wf_id, t, [1, 2, 2, 2, 2], mock);
     mh.enforce_correct_number_of_polls = false;
-    let mut mock = build_mock_pollers(mh);
-    mock.worker_cfg(|wc| wc.max_cached_workflows = 1);
-    let core = Arc::new(mock_worker(mock));
-    let mut worker = TestWorker::new(core.clone(), TEST_Q.to_string());
+    let mut worker = mock_sdk_cfg(mh, |wc| wc.max_cached_workflows = 1);
+    let core = worker.orig_core_worker.clone();
+
     let shutdown_barr: &'static Barrier = Box::leak(Box::new(Barrier::new(2)));
 
     worker.register_wf(
@@ -185,7 +180,7 @@ async fn local_act_heartbeat(#[case] shutdown_middle: bool) {
             Ok(().into())
         },
     );
-    worker.register_activity("echo", move |str: String| async move {
+    worker.register_activity("echo", move |_ctx: ActContext, str: String| async move {
         if shutdown_middle {
             shutdown_barr.wait().await;
         }
@@ -226,9 +221,7 @@ async fn local_act_fail_and_retry(#[case] eventually_pass: bool) {
     let wf_id = "fakeid";
     let mock = mock_workflow_client();
     let mh = MockPollCfg::from_resp_batches(wf_id, t, [1], mock);
-    let mock = build_mock_pollers(mh);
-    let core = mock_worker(mock);
-    let mut worker = TestWorker::new(Arc::new(core), TEST_Q.to_string());
+    let mut worker = mock_sdk(mh);
 
     worker.register_wf(
         DEFAULT_WORKFLOW_TYPE.to_owned(),
@@ -256,7 +249,7 @@ async fn local_act_fail_and_retry(#[case] eventually_pass: bool) {
         },
     );
     let attempts: &'static _ = Box::leak(Box::new(AtomicUsize::new(0)));
-    worker.register_activity("echo", move |_: String| async move {
+    worker.register_activity("echo", move |_ctx: ActContext, _: String| async move {
         // Succeed on 3rd attempt (which is ==2 since fetch_add returns prev val)
         if 2 == attempts.fetch_add(1, Ordering::Relaxed) && eventually_pass {
             Ok(())
@@ -309,10 +302,7 @@ async fn local_act_retry_long_backoff_uses_timer() {
         [1.into(), 2.into(), ResponseType::AllHistory],
         mock,
     );
-    let mut mock = build_mock_pollers(mh);
-    mock.worker_cfg(|w| w.max_cached_workflows = 1);
-    let core = mock_worker(mock);
-    let mut worker = TestWorker::new(Arc::new(core), TEST_Q.to_string());
+    let mut worker = mock_sdk_cfg(mh, |w| w.max_cached_workflows = 1);
 
     worker.register_wf(
         DEFAULT_WORKFLOW_TYPE.to_owned(),
@@ -338,7 +328,7 @@ async fn local_act_retry_long_backoff_uses_timer() {
             Ok(().into())
         },
     );
-    worker.register_activity("echo", move |_: String| async move {
+    worker.register_activity("echo", move |_ctx: ActContext, _: String| async move {
         Result::<(), _>::Err(anyhow!("Oh no I failed!"))
     });
     worker