@temporalio/core-bridge 1.13.2 → 1.14.0

package/sdk-core/{core → crates/sdk-core}/src/telemetry/prometheus_meter.rs

@@ -12,7 +12,7 @@ use std::{
     sync::Arc,
     time::Duration,
 };
-use temporal_sdk_core_api::telemetry::metrics::{
+use temporalio_common::telemetry::metrics::{
     CoreMeter, Counter, CounterBase, Gauge, GaugeBase, GaugeF64, GaugeF64Base, Histogram,
     HistogramBase, HistogramDuration, HistogramDurationBase, HistogramF64, HistogramF64Base,
     MetricAttributable, MetricAttributes, MetricParameters, NewAttributes, OrderedPromLabelSet,
@@ -459,7 +459,7 @@ pub struct CorePrometheusMeter {
     registry: Registry,
     use_seconds_for_durations: bool,
     unit_suffix: bool,
-    bucket_overrides: temporal_sdk_core_api::telemetry::HistogramBucketOverrides,
+    bucket_overrides: temporalio_common::telemetry::HistogramBucketOverrides,
 }
 
 impl CorePrometheusMeter {
@@ -467,7 +467,7 @@ impl CorePrometheusMeter {
         registry: Registry,
         use_seconds_for_durations: bool,
         unit_suffix: bool,
-        bucket_overrides: temporal_sdk_core_api::telemetry::HistogramBucketOverrides,
+        bucket_overrides: temporalio_common::telemetry::HistogramBucketOverrides,
     ) -> Self {
         Self {
             registry,
@@ -638,7 +638,7 @@ mod tests {
     use super::*;
     use crate::telemetry::{TelemetryInstance, metrics::MetricsContext};
     use prometheus::{Encoder, TextEncoder};
-    use temporal_sdk_core_api::telemetry::{
+    use temporalio_common::telemetry::{
         METRIC_PREFIX,
         metrics::{MetricKeyValue, NewAttributes},
     };
@@ -650,7 +650,7 @@ mod tests {
             registry.clone(),
             false,
             false,
-            temporal_sdk_core_api::telemetry::HistogramBucketOverrides::default(),
+            temporalio_common::telemetry::HistogramBucketOverrides::default(),
         );
 
         let counter = meter.counter(MetricParameters {
@@ -690,7 +690,7 @@ mod tests {
             registry.clone(),
             false,
             false,
-            temporal_sdk_core_api::telemetry::HistogramBucketOverrides::default(),
+            temporalio_common::telemetry::HistogramBucketOverrides::default(),
         );
 
         let base_attrs = meter.new_attributes(NewAttributes::new(vec![
@@ -729,7 +729,7 @@ mod tests {
             registry.clone(),
             false,
             false,
-            temporal_sdk_core_api::telemetry::HistogramBucketOverrides::default(),
+            temporalio_common::telemetry::HistogramBucketOverrides::default(),
         );
 
         let histogram_ms = meter_ms.histogram_duration(MetricParameters {
@@ -759,7 +759,7 @@ mod tests {
             registry_s.clone(),
             true,
             false,
-            temporal_sdk_core_api::telemetry::HistogramBucketOverrides::default(),
+            temporalio_common::telemetry::HistogramBucketOverrides::default(),
         );
 
         let histogram_s = meter_s.histogram_duration(MetricParameters {
@@ -791,7 +791,7 @@ mod tests {
             registry.clone(),
             false,
             false,
-            temporal_sdk_core_api::telemetry::HistogramBucketOverrides::default(),
+            temporalio_common::telemetry::HistogramBucketOverrides::default(),
         );
         let counter = meter.counter(MetricParameters {
             name: "no_labels".into(),
@@ -812,7 +812,7 @@ mod tests {
             registry.clone(),
             false,
             false,
-            temporal_sdk_core_api::telemetry::HistogramBucketOverrides::default(),
+            temporalio_common::telemetry::HistogramBucketOverrides::default(),
         );
         let telem_instance = TelemetryInstance::new(
             None,
@@ -820,7 +820,7 @@ mod tests {
             METRIC_PREFIX.to_string(),
             Some(Arc::new(meter)),
             true,
-            temporal_sdk_core_api::telemetry::TaskQueueLabelStrategy::UseNormal,
+            temporalio_common::telemetry::TaskQueueLabelStrategy::UseNormal,
         );
         let mc = MetricsContext::top_level("foo".to_string(), "q".to_string(), &telem_instance);
         mc.worker_registered();
@@ -844,7 +844,7 @@ mod tests {
             registry.clone(),
             false,
             false,
-            temporal_sdk_core_api::telemetry::HistogramBucketOverrides::default(),
+            temporalio_common::telemetry::HistogramBucketOverrides::default(),
         );
         let dashes = meter.counter(MetricParameters {
             name: "dash-in-name".into(),
@@ -865,7 +865,7 @@ mod tests {
             registry.clone(),
             false,
             false,
-            temporal_sdk_core_api::telemetry::HistogramBucketOverrides::default(),
+            temporalio_common::telemetry::HistogramBucketOverrides::default(),
         );
         let dashes = meter.counter(MetricParameters {
             name: "not@permitted:symbols".into(),

package/sdk-core/{core → crates/sdk-core}/src/telemetry/prometheus_server.rs

@@ -10,7 +10,7 @@ use std::{
     net::{SocketAddr, TcpListener},
     sync::Arc,
 };
-use temporal_sdk_core_api::telemetry::PrometheusExporterOptions;
+use temporalio_common::telemetry::PrometheusExporterOptions;
 use tokio::{io, task::AbortHandle};
 
 pub struct StartedPromServer {
@@ -20,7 +20,7 @@ pub struct StartedPromServer {
 }
 
 /// Builds and runs a prometheus endpoint which can be scraped by prom instances for metrics export.
-/// Returns the meter that can be used as a [CoreMeter].
+/// Returns the meter that can be used as a [`CoreMeter`](temporalio_common::telemetry::metrics::CoreMeter).
 ///
 /// Requires a Tokio runtime to exist, and will block briefly while binding the server endpoint.
 pub fn start_prometheus_metric_exporter(

package/sdk-core/{core → crates/sdk-core}/src/test_help/integ_helpers.rs

@@ -36,39 +36,37 @@ use std::{
     task::{Context, Poll},
     time::Duration,
 };
-use temporal_sdk_core_api::{
+use temporalio_common::{
     Worker as WorkerTrait,
     errors::PollError,
-    worker::{PollerBehavior, WorkerVersioningStrategy},
-};
-use temporal_sdk_core_protos::{
-    coresdk::{
-        workflow_activation::{WorkflowActivationJob, workflow_activation_job},
-        workflow_commands::{CompleteWorkflowExecution, StartTimer},
-        workflow_completion::WorkflowActivationCompletion,
-    },
-    temporal::api::{
-        common::v1::WorkflowExecution,
-        enums::v1::WorkflowTaskFailedCause,
-        failure::v1::Failure,
-        protocol::{self, v1::message},
-        update,
-        workflowservice::v1::{
-            PollActivityTaskQueueResponse, PollNexusTaskQueueResponse,
-            PollWorkflowTaskQueueResponse, RespondWorkflowTaskCompletedResponse,
+    protos::{
+        coresdk::{
+            workflow_activation::{WorkflowActivationJob, workflow_activation_job},
+            workflow_commands::{CompleteWorkflowExecution, StartTimer},
+            workflow_completion::WorkflowActivationCompletion,
         },
+        temporal::api::{
+            common::v1::WorkflowExecution,
+            enums::v1::WorkflowTaskFailedCause,
+            failure::v1::Failure,
+            protocol::{self, v1::message},
+            update,
+            workflowservice::v1::{
+                PollActivityTaskQueueResponse, PollNexusTaskQueueResponse,
+                PollWorkflowTaskQueueResponse, RespondWorkflowTaskCompletedResponse,
+            },
+        },
+        utilities::pack_any,
     },
-    utilities::pack_any,
+    worker::{PollerBehavior, WorkerTaskTypes, WorkerVersioningStrategy},
 };
 use tokio::sync::{Notify, mpsc::unbounded_channel};
 use tokio_stream::wrappers::UnboundedReceiverStream;
+use uuid::Uuid;
 
 /// Default namespace for testing
 pub const NAMESPACE: &str = "default";
 
-/// Default task queue for testing
-pub const TEST_Q: &str = "q";
-
 /// Initiate shutdown, drain the pollers (handling evictions), and wait for shutdown to complete.
 pub async fn drain_pollers_and_shutdown(worker: &dyn WorkerTrait) {
     worker.initiate_shutdown();
@@ -102,11 +100,12 @@ pub async fn drain_pollers_and_shutdown(worker: &dyn WorkerTrait) {
 pub fn test_worker_cfg() -> WorkerConfigBuilder {
     let mut wcb = WorkerConfigBuilder::default();
     wcb.namespace(NAMESPACE)
-        .task_queue(TEST_Q)
+        .task_queue(Uuid::new_v4().to_string())
         .versioning_strategy(WorkerVersioningStrategy::None {
             build_id: "test_bin_id".to_string(),
         })
         .ignore_evicts_on_shutdown(true)
+        .task_types(WorkerTaskTypes::all())
         // Serial polling since it makes mocking much easier.
         .workflow_task_poller_behavior(PollerBehavior::SimpleMaximum(1_usize));
     wcb
@@ -185,11 +184,16 @@ pub fn build_fake_worker(
 }
 
 pub fn mock_worker(mocks: MocksHolder) -> Worker {
-    let sticky_q = sticky_q_name_for_worker("unit-test", &mocks.inputs.config);
-    let act_poller = if mocks.inputs.config.no_remote_activities {
+    let sticky_q = sticky_q_name_for_worker("unit-test", mocks.inputs.config.max_cached_workflows);
+    let act_poller = if mocks.inputs.config.task_types.enable_remote_activities {
+        mocks.inputs.act_poller
+    } else {
         None
+    };
+    let nexus_poller = if mocks.inputs.config.task_types.enable_nexus {
+        mocks.inputs.nexus_poller
     } else {
-        mocks.inputs.act_poller
+        None
     };
     Worker::new_with_pollers(
         mocks.inputs.config,
@@ -198,14 +202,13 @@ pub fn mock_worker(mocks: MocksHolder) -> Worker {
         TaskPollers::Mocked {
             wft_stream: mocks.inputs.wft_stream,
             act_poller,
-            nexus_poller: mocks
-                .inputs
-                .nexus_poller
-                .unwrap_or_else(|| mock_poller_from_resps([])),
+            nexus_poller,
         },
         None,
         None,
+        false,
     )
+    .unwrap()
 }
 
 pub struct FakeWfResponses {
@@ -230,15 +233,38 @@ impl MocksHolder {
         self.inputs.act_poller = Some(poller);
     }
 
+    /// Helper to create and set an activity poller from responses
+    #[cfg(test)]
+    pub(crate) fn set_act_poller_from_resps<ACT>(&mut self, act_tasks: ACT)
+    where
+        ACT: IntoIterator<Item = QueueResponse<PollActivityTaskQueueResponse>>,
+        <ACT as IntoIterator>::IntoIter: Send + 'static,
+    {
+        let act_poller = mock_poller_from_resps(act_tasks);
+        self.set_act_poller(act_poller);
+    }
+
+    /// Helper to create and set a nexus poller from responses
+    #[cfg(test)]
+    pub(crate) fn set_nexus_poller_from_resps<NEX>(&mut self, nexus_tasks: NEX)
+    where
+        NEX: IntoIterator<Item = QueueResponse<PollNexusTaskQueueResponse>>,
+        <NEX as IntoIterator>::IntoIter: Send + 'static,
+    {
+        let nexus_poller = mock_poller_from_resps(nexus_tasks);
+        self.inputs.nexus_poller = Some(nexus_poller);
+    }
+
     /// Can be used for tests that need to avoid auto-shutdown due to running out of mock responses
     pub fn make_wft_stream_interminable(&mut self) {
-        let old_stream = std::mem::replace(&mut self.inputs.wft_stream, stream::pending().boxed());
-        self.inputs.wft_stream = old_stream.chain(stream::pending()).boxed();
+        if let Some(old_stream) = self.inputs.wft_stream.take() {
+            self.inputs.wft_stream = Some(old_stream.chain(stream::pending()).boxed());
+        }
     }
 }
 
 pub struct MockWorkerInputs {
-    pub(crate) wft_stream: BoxStream<'static, Result<ValidPollWFTQResponse, tonic::Status>>,
+    pub(crate) wft_stream: Option<BoxStream<'static, Result<ValidPollWFTQResponse, tonic::Status>>>,
     pub(crate) act_poller: Option<BoxedPoller<PollActivityTaskQueueResponse>>,
     pub(crate) nexus_poller: Option<BoxedPoller<PollNexusTaskQueueResponse>>,
     pub(crate) config: WorkerConfig,
@@ -255,7 +281,7 @@ impl MockWorkerInputs {
         wft_stream: BoxStream<'static, Result<ValidPollWFTQResponse, tonic::Status>>,
     ) -> Self {
         Self {
-            wft_stream,
+            wft_stream: Some(wft_stream),
             act_poller: None,
             nexus_poller: None,
             config: test_worker_cfg().build().unwrap(),
@@ -275,7 +301,7 @@ impl MocksHolder {
         }
     }
 
-    /// Uses the provided list of tasks to create a mock poller for the `TEST_Q`
+    /// Uses the provided list of tasks to create a mock poller with a randomly generated task queue
     pub fn from_client_with_activities<ACT>(
         client: impl WorkerClient + 'static,
         act_tasks: ACT,
@@ -284,10 +310,9 @@ impl MocksHolder {
         ACT: IntoIterator<Item = QueueResponse<PollActivityTaskQueueResponse>>,
         <ACT as IntoIterator>::IntoIter: Send + 'static,
     {
-        let wft_stream = stream::pending().boxed();
         let mock_act_poller = mock_poller_from_resps(act_tasks);
         let mock_worker = MockWorkerInputs {
-            wft_stream,
+            wft_stream: None,
             act_poller: Some(mock_act_poller),
             nexus_poller: None,
             config: test_worker_cfg().build().unwrap(),
@@ -299,6 +324,66 @@ impl MocksHolder {
         }
     }
 
+    /// Uses the provided list of tasks to create a mock poller with a randomly generated task queue
+    pub fn from_client_with_nexus<NEX>(
+        client: impl WorkerClient + 'static,
+        nexus_tasks: NEX,
+    ) -> Self
+    where
+        NEX: IntoIterator<Item = QueueResponse<PollNexusTaskQueueResponse>>,
+        <NEX as IntoIterator>::IntoIter: Send + 'static,
+    {
+        let mock_nexus_poller = mock_poller_from_resps(nexus_tasks);
+        let mock_worker = MockWorkerInputs {
+            wft_stream: None,
+            act_poller: None,
+            nexus_poller: Some(mock_nexus_poller),
+            config: test_worker_cfg().build().unwrap(),
+        };
+        Self {
+            client: Arc::new(client),
+            inputs: mock_worker,
+            outstanding_task_map: None,
+        }
+    }
+
+    /// Create a MocksHolder with custom combination of task pollers.
+    /// Allows any combination of workflow, activity, and nexus tasks.
+    pub fn from_client_with_custom<WFT, ACT, NEX>(
+        client: impl WorkerClient + 'static,
+        wft_stream: Option<WFT>,
+        activity_tasks: Option<ACT>,
+        nexus_tasks: Option<NEX>,
+    ) -> Self
+    where
+        WFT: Stream<Item = PollWorkflowTaskQueueResponse> + Send + 'static,
+        ACT: IntoIterator<Item = QueueResponse<PollActivityTaskQueueResponse>>,
+        <ACT as IntoIterator>::IntoIter: Send + 'static,
+        NEX: IntoIterator<Item = QueueResponse<PollNexusTaskQueueResponse>>,
+        <NEX as IntoIterator>::IntoIter: Send + 'static,
+    {
+        let wft_stream = wft_stream.map(|s| {
+            s.map(|r| Ok(r.try_into().expect("Mock responses must be valid work")))
+                .boxed()
+        });
+
+        let act_poller = activity_tasks.map(|tasks| mock_poller_from_resps(tasks));
+        let nexus_poller = nexus_tasks.map(|tasks| mock_poller_from_resps(tasks));
+
+        let mock_worker = MockWorkerInputs {
+            wft_stream,
+            act_poller,
+            nexus_poller,
+            config: test_worker_cfg().build().unwrap(),
+        };
+
+        Self {
+            client: Arc::new(client),
+            inputs: mock_worker,
+            outstanding_task_map: None,
+        }
+    }
+
     /// Uses the provided task responses and delivers them as quickly as possible when polled.
     /// This is only useful to test buffering, as typically you do not want to pretend that
     /// the server is delivering WFTs super fast for the same run.
@@ -306,9 +391,11 @@ impl MocksHolder {
         client: impl WorkerClient + 'static,
         stream: impl Stream<Item = PollWorkflowTaskQueueResponse> + Send + 'static,
     ) -> Self {
-        let wft_stream = stream
-            .map(|r| Ok(r.try_into().expect("Mock responses must be valid work")))
-            .boxed();
+        let wft_stream = Some(
+            stream
+                .map(|r| Ok(r.try_into().expect("Mock responses must be valid work")))
+                .boxed(),
+        );
         let mock_worker = MockWorkerInputs {
             wft_stream,
             act_poller: None,

package/sdk-core/{core → crates/sdk-core}/src/test_help/unit_helpers.rs

@@ -2,10 +2,12 @@
 
 use futures_util::{StreamExt, stream::FuturesUnordered};
 use std::{collections::HashSet, future::Future};
-use temporal_sdk_core_api::Worker as CoreWorker;
-use temporal_sdk_core_protos::coresdk::{
-    workflow_activation::workflow_activation_job,
-    workflow_completion::{WorkflowActivationCompletion, workflow_activation_completion},
+use temporalio_common::{
+    Worker as CoreWorker,
+    protos::coresdk::{
+        workflow_activation::workflow_activation_job,
+        workflow_completion::{WorkflowActivationCompletion, workflow_activation_completion},
+    },
 };
 
 /// Given a desired number of concurrent executions and a provided function that produces a future,
@@ -44,7 +46,7 @@ macro_rules! job_assert {
 }
 
 type AsserterWithReply<'a> = (
-    &'a dyn Fn(&temporal_sdk_core_protos::coresdk::workflow_activation::WorkflowActivation),
+    &'a dyn Fn(&temporalio_common::protos::coresdk::workflow_activation::WorkflowActivation),
     workflow_activation_completion::Status,
 );
 
@@ -175,14 +177,14 @@ pub(crate) async fn poll_and_reply_clears_outstanding_evicts<'a>(
 }
 
 pub(crate) fn gen_assert_and_reply(
-    asserter: &dyn Fn(&temporal_sdk_core_protos::coresdk::workflow_activation::WorkflowActivation),
+    asserter: &dyn Fn(&temporalio_common::protos::coresdk::workflow_activation::WorkflowActivation),
     reply_commands: Vec<
-        temporal_sdk_core_protos::coresdk::workflow_commands::workflow_command::Variant,
+        temporalio_common::protos::coresdk::workflow_commands::workflow_command::Variant,
     >,
 ) -> AsserterWithReply<'_> {
     (
         asserter,
-        temporal_sdk_core_protos::coresdk::workflow_completion::Success::from_variants(
+        temporalio_common::protos::coresdk::workflow_completion::Success::from_variants(
             reply_commands,
         )
         .into(),
@@ -190,13 +192,13 @@ pub(crate) fn gen_assert_and_reply(
 }
 
 pub(crate) fn gen_assert_and_fail(
-    asserter: &dyn Fn(&temporal_sdk_core_protos::coresdk::workflow_activation::WorkflowActivation),
+    asserter: &dyn Fn(&temporalio_common::protos::coresdk::workflow_activation::WorkflowActivation),
 ) -> AsserterWithReply<'_> {
     (
         asserter,
-        temporal_sdk_core_protos::coresdk::workflow_completion::Failure {
+        temporalio_common::protos::coresdk::workflow_completion::Failure {
             failure: Some(
-                temporal_sdk_core_protos::temporal::api::failure::v1::Failure {
+                temporalio_common::protos::temporal::api::failure::v1::Failure {
                     message: "Intentional test failure".to_string(),
                     ..Default::default()
                 },

package/sdk-core/{core → crates/sdk-core}/src/worker/activities/activity_heartbeat_manager.rs

@@ -9,7 +9,7 @@ use std::{
     sync::Arc,
     time::{Duration, Instant},
 };
-use temporal_sdk_core_protos::{
+use temporalio_common::protos::{
     coresdk::{
         ActivityHeartbeat, IntoPayloadsExt,
         activity_task::{ActivityCancelReason, ActivityCancellationDetails, ActivityTask},
@@ -430,7 +430,7 @@ mod test {
 
     use crate::worker::client::mocks::mock_worker_client;
     use std::time::Duration;
-    use temporal_sdk_core_protos::temporal::api::{
+    use temporalio_common::protos::temporal::api::{
         common::v1::Payload, workflowservice::v1::RecordActivityTaskHeartbeatResponse,
     };
     use tokio::time::sleep;

package/sdk-core/{core → crates/sdk-core}/src/worker/activities/local_activities.rs

@@ -17,18 +17,20 @@ use std::{
     task::{Context, Poll},
     time::{Duration, Instant, SystemTime},
 };
-use temporal_sdk_core_api::worker::LocalActivitySlotKind;
-use temporal_sdk_core_protos::{
-    coresdk::{
-        LocalActivitySlotInfo,
-        activity_result::{Cancellation, Failure as ActFail, Success},
-        activity_task::{ActivityCancelReason, ActivityTask, Start, activity_task},
-    },
-    temporal::api::{
-        common::v1::WorkflowExecution,
-        enums::v1::TimeoutType,
-        failure::v1::{Failure as APIFailure, TimeoutFailureInfo, failure},
+use temporalio_common::{
+    protos::{
+        coresdk::{
+            LocalActivitySlotInfo,
+            activity_result::{Cancellation, Failure as ActFail, Success},
+            activity_task::{ActivityCancelReason, ActivityTask, Start, activity_task},
+        },
+        temporal::api::{
+            common::v1::WorkflowExecution,
+            enums::v1::TimeoutType,
+            failure::v1::{Failure as APIFailure, TimeoutFailureInfo, failure},
+        },
     },
+    worker::LocalActivitySlotKind,
 };
 use tokio::{
     sync::{
@@ -978,7 +980,7 @@ mod tests {
     use super::*;
     use crate::{prost_dur, protosext::LACloseTimeouts, retry_logic::ValidatedRetryPolicy};
     use futures_util::FutureExt;
-    use temporal_sdk_core_protos::temporal::api::{
+    use temporalio_common::protos::temporal::api::{
         common::v1::RetryPolicy,
         failure::v1::{ApplicationFailureInfo, Failure, failure::FailureInfo},
     };

package/sdk-core/{core → crates/sdk-core}/src/worker/activities.rs

@@ -35,17 +35,21 @@ use std::{
     },
     time::{Duration, Instant, SystemTime},
 };
-use temporal_sdk_core_api::worker::ActivitySlotKind;
-use temporal_sdk_core_protos::{
-    coresdk::{
-        ActivityHeartbeat, ActivitySlotInfo,
-        activity_result::{self as ar, activity_execution_result as aer},
-        activity_task::{ActivityCancelReason, ActivityCancellationDetails, ActivityTask},
-    },
-    temporal::api::{
-        failure::v1::{ApplicationFailureInfo, CanceledFailureInfo, Failure, failure::FailureInfo},
-        workflowservice::v1::PollActivityTaskQueueResponse,
+use temporalio_common::{
+    protos::{
+        coresdk::{
+            ActivityHeartbeat, ActivitySlotInfo,
+            activity_result::{self as ar, activity_execution_result as aer},
+            activity_task::{ActivityCancelReason, ActivityCancellationDetails, ActivityTask},
+        },
+        temporal::api::{
+            failure::v1::{
+                ApplicationFailureInfo, CanceledFailureInfo, Failure, failure::FailureInfo,
+            },
+            workflowservice::v1::PollActivityTaskQueueResponse,
+        },
     },
+    worker::ActivitySlotKind,
 };
 use tokio::{
     join,
@@ -735,8 +739,10 @@ mod tests {
         prost_dur,
         worker::client::mocks::mock_worker_client,
     };
-    use temporal_sdk_core_api::worker::PollerBehavior;
-    use temporal_sdk_core_protos::coresdk::activity_result::ActivityExecutionResult;
+    use crossbeam_utils::atomic::AtomicCell;
+    use temporalio_common::{
+        protos::coresdk::activity_result::ActivityExecutionResult, worker::PollerBehavior,
+    };
 
     #[tokio::test]
     async fn per_worker_ratelimit() {
@@ -780,6 +786,7 @@ mod tests {
                 max_worker_acts_per_second: Some(2.0),
                 max_tps: None,
             },
+            Arc::new(AtomicCell::new(None)),
         );
         let atm = WorkerActivityTasks::new(
             sem.clone(),
@@ -871,6 +878,7 @@ mod tests {
                 max_worker_acts_per_second: None,
                 max_tps: None,
             },
+            Arc::new(AtomicCell::new(None)),
         );
         let atm = WorkerActivityTasks::new(
             sem.clone(),
@@ -944,6 +952,7 @@ mod tests {
                 max_worker_acts_per_second: None,
                 max_tps: None,
             },
+            Arc::new(AtomicCell::new(None)),
         );
         let atm = WorkerActivityTasks::new(
             sem.clone(),

package/sdk-core/{core → crates/sdk-core}/src/worker/client/mocks.rs

@@ -1,10 +1,10 @@
 use super::*;
 use futures_util::Future;
 use std::sync::{Arc, LazyLock};
-use temporal_client::SlotManager;
+use temporalio_client::worker::ClientWorkerSet;
 
-pub(crate) static DEFAULT_WORKERS_REGISTRY: LazyLock<Arc<SlotManager>> =
-    LazyLock::new(|| Arc::new(SlotManager::new()));
+pub(crate) static DEFAULT_WORKERS_REGISTRY: LazyLock<Arc<ClientWorkerSet>> =
+    LazyLock::new(|| Arc::new(ClientWorkerSet::new()));
 
 pub(crate) static DEFAULT_TEST_CAPABILITIES: &Capabilities = &Capabilities {
     signal_and_query_header: true,
@@ -30,11 +30,18 @@ pub fn mock_worker_client() -> MockWorkerClient {
         .returning(|| DEFAULT_WORKERS_REGISTRY.clone());
     r.expect_is_mock().returning(|| true);
     r.expect_shutdown_worker()
-        .returning(|_| Ok(ShutdownWorkerResponse {}));
+        .returning(|_, _| Ok(ShutdownWorkerResponse {}));
     r.expect_sdk_name_and_version()
         .returning(|| ("test-core".to_string(), "0.0.0".to_string()));
-    r.expect_get_identity()
+    r.expect_identity()
         .returning(|| "test-identity".to_string());
+    r.expect_worker_grouping_key().returning(Uuid::new_v4);
+    r.expect_set_heartbeat_client_fields().returning(|hb| {
+        hb.sdk_name = "test-core".to_string();
+        hb.sdk_version = "0.0.0".to_string();
+        hb.worker_identity = "test-identity".to_string();
+        hb.heartbeat_time = Some(SystemTime::now().into());
+    });
     r
 }
 
@@ -48,7 +55,7 @@ pub(crate) fn mock_manual_worker_client() -> MockManualWorkerClient {
     r.expect_is_mock().returning(|| true);
     r.expect_sdk_name_and_version()
         .returning(|| ("test-core".to_string(), "0.0.0".to_string()));
-    r.expect_get_identity()
+    r.expect_identity()
         .returning(|| "test-identity".to_string());
     r
 }
@@ -68,7 +75,7 @@ mockall::mock! {
             -> impl Future<Output = Result<PollActivityTaskQueueResponse>> + Send + 'b
             where 'a: 'b, Self: 'b;
 
-        fn poll_nexus_task<'a, 'b>(&self, poll_options: PollOptions)
+        fn poll_nexus_task<'a, 'b>(&self, poll_options: PollOptions, send_heartbeat: bool)
             -> impl Future<Output = Result<PollNexusTaskQueueResponse>> + Send + 'b
             where 'a: 'b, Self: 'b;
 
@@ -139,7 +146,7 @@ mockall::mock! {
         fn respond_legacy_query<'a, 'b>(
             &self,
             task_token: TaskToken,
-        query_result: LegacyQueryResult,
+            query_result: LegacyQueryResult,
         ) -> impl Future<Output = Result<RespondQueryTaskCompletedResponse>> + Send + 'b
             where 'a: 'b, Self: 'b;
 
@@ -147,16 +154,22 @@ mockall::mock! {
           impl Future<Output = Result<DescribeNamespaceResponse>> + Send + 'b
           where 'a: 'b, Self: 'b;
 
-        fn shutdown_worker<'a, 'b>(&self, sticky_task_queue: String) -> impl Future<Output = Result<ShutdownWorkerResponse>> + Send + 'b
+        fn shutdown_worker<'a, 'b>(&self, sticky_task_queue: String, worker_heartbeat: Option<WorkerHeartbeat>) -> impl Future<Output = Result<ShutdownWorkerResponse>> + Send + 'b
             where 'a: 'b, Self: 'b;
 
-        fn record_worker_heartbeat<'a, 'b>(&self, heartbeat: WorkerHeartbeat) -> impl Future<Output = Result<RecordWorkerHeartbeatResponse>> + Send + 'b where 'a: 'b, Self: 'b;
+        fn record_worker_heartbeat<'a, 'b>(
+            &self,
+            namespace: String,
+            heartbeat: Vec<WorkerHeartbeat>
+        ) -> impl Future<Output = Result<RecordWorkerHeartbeatResponse>> + Send + 'b where 'a: 'b, Self: 'b;
 
         fn replace_client(&self, new_client: Client);
         fn capabilities(&self) -> Option<Capabilities>;
-        fn workers(&self) -> Arc<SlotManager>;
+        fn workers(&self) -> Arc<ClientWorkerSet>;
         fn is_mock(&self) -> bool;
         fn sdk_name_and_version(&self) -> (String, String);
-        fn get_identity(&self) -> String;
+        fn identity(&self) -> String;
+        fn worker_grouping_key(&self) -> Uuid;
+        fn set_heartbeat_client_fields(&self, heartbeat: &mut WorkerHeartbeat);
     }
 }