@temporalio/core-bridge 1.13.1 → 1.14.0

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

@@ -6,7 +6,7 @@ use std::{
     collections::{BTreeSet, HashSet},
     iter,
 };
-use temporal_sdk_core_protos::temporal::api::{
+use temporalio_common::protos::temporal::api::{
     history::v1::WorkflowTaskCompletedEventAttributes, sdk::v1::WorkflowTaskCompletedMetadata,
     workflowservice::v1::get_system_info_response,
 };
@@ -18,9 +18,10 @@ use temporal_sdk_core_protos::temporal::api::{
 /// may be removed from the enum. *Importantly*, all variants must be given explicit values, such
 /// that removing older variants does not create any change in existing values. Removed flag
 /// variants must be reserved forever (a-la protobuf), and should be called out in a comment.
+#[allow(unreachable_pub)] // re-exported in test_help::integ_helpers
 #[repr(u32)]
 #[derive(PartialEq, Eq, PartialOrd, Ord, Hash, Copy, Clone, Debug, enum_iterator::Sequence)]
-pub(crate) enum CoreInternalFlags {
+pub enum CoreInternalFlags {
     /// In this flag additional checks were added to a number of state machines to ensure that
     /// the ID and type of activities, local activities, and child workflows match during replay.
     IdAndTypeDeterminismChecks = 1,
@@ -34,7 +35,7 @@ pub(crate) enum CoreInternalFlags {
     /// multiple workflow coroutines generated a terminal command). This has the
     /// consequence that all non-terminal commands are sent to the server, even
     /// if in the sequence delivered by lang they came after a terminal command.
-    /// See https://github.com/temporalio/features/issues/481.
+    /// See <https://github.com/temporalio/features/issues/481>.
     MoveTerminalCommands = 3,
     /// We received a value higher than this code can understand.
     TooHigh = u32::MAX,
@@ -202,6 +203,15 @@ impl InternalFlags {
             Self::Disabled => Either::Right(iter::empty()),
         }
     }
+
+    pub(crate) fn last_sdk_version(&self) -> Option<&str> {
+        match self {
+            InternalFlags::Enabled {
+                last_sdk_version, ..
+            } if !last_sdk_version.is_empty() => Some(last_sdk_version),
+            InternalFlags::Enabled { .. } | InternalFlags::Disabled => None,
+        }
+    }
 }
 
 impl CoreInternalFlags {
@@ -223,8 +233,9 @@ impl CoreInternalFlags {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use temporal_sdk_core_protos::temporal::api::workflowservice::v1::get_system_info_response::Capabilities;
+    use temporalio_common::protos::temporal::api::workflowservice::v1::get_system_info_response::Capabilities;
 
+    #[allow(clippy::derivable_impls)] // Only want this in test
     impl Default for InternalFlags {
         fn default() -> Self {
             Self::Disabled

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

@@ -12,6 +12,8 @@ extern crate tracing;
 extern crate core;
 
 mod abstractions;
+#[cfg(feature = "antithesis_assertions")]
+mod antithesis;
 #[cfg(feature = "debug-plugin")]
 pub mod debug_client;
 #[cfg(feature = "ephemeral-server")]
@@ -26,25 +28,22 @@ mod worker;
 
 #[cfg(test)]
 mod core_tests;
-#[cfg(test)]
+#[cfg(any(feature = "test-utilities", test))]
 #[macro_use]
-mod test_help;
+pub mod test_help;
 
-pub(crate) use temporal_sdk_core_api::errors;
+pub(crate) use temporalio_common::errors;
 
 pub use pollers::{
-    Client, ClientOptions, ClientOptionsBuilder, ClientTlsConfig, RetryClient, RetryConfig,
-    TlsConfig, WorkflowClientTrait,
+    Client, ClientOptions, ClientTlsOptions, RetryClient, RetryOptions, TlsOptions,
+    WorkflowClientTrait,
 };
-pub use temporal_sdk_core_api as api;
-pub use temporal_sdk_core_protos as protos;
-pub use temporal_sdk_core_protos::TaskToken;
+pub use temporalio_common::protos::TaskToken;
 pub use url::Url;
 pub use worker::{
-    FixedSizeSlotSupplier, RealSysInfo, ResourceBasedSlotsOptions,
-    ResourceBasedSlotsOptionsBuilder, ResourceBasedTuner, ResourceSlotOptions, SlotSupplierOptions,
-    TunerBuilder, TunerHolder, TunerHolderOptions, TunerHolderOptionsBuilder, Worker, WorkerConfig,
-    WorkerConfigBuilder,
+    FixedSizeSlotSupplier, ResourceBasedSlotsOptions, ResourceBasedSlotsOptionsBuilder,
+    ResourceBasedTuner, ResourceSlotOptions, SlotSupplierOptions, TunerBuilder, TunerHolder,
+    TunerHolderOptions, TunerHolderOptionsBuilder, Worker, WorkerConfig, WorkerConfigBuilder,
 };
 
 /// Expose [WorkerClient] symbols
@@ -61,14 +60,14 @@ use crate::{
 };
 use anyhow::bail;
 use futures_util::Stream;
-use std::sync::{Arc, OnceLock};
-use temporal_client::{ConfiguredClient, NamespacedClient, TemporalServiceClientWithMetrics};
-use temporal_sdk_core_api::{
+use std::{sync::Arc, time::Duration};
+use temporalio_client::{ConfiguredClient, NamespacedClient, SharedReplaceableClient};
+use temporalio_common::{
     Worker as WorkerTrait,
     errors::{CompleteActivityError, PollError},
+    protos::coresdk::ActivityHeartbeat,
     telemetry::TelemetryOptions,
 };
-use temporal_sdk_core_protos::coresdk::ActivityHeartbeat;
 
 /// Initialize a worker bound to a task queue.
 ///
@@ -89,39 +88,40 @@ pub fn init_worker<CT>(
 where
     CT: Into<sealed::AnyClient>,
 {
-    let client = init_worker_client(&worker_config, *client.into().into_inner());
-    if client.namespace() != worker_config.namespace {
-        bail!("Passed in client is not bound to the same namespace as the worker");
-    }
-    if client.namespace() == "" {
-        bail!("Client namespace cannot be empty");
+    let namespace = worker_config.namespace.clone();
+    if namespace.is_empty() {
+        bail!("Worker namespace cannot be empty");
     }
-    let client_ident = client.get_identity().to_owned();
-    let sticky_q = sticky_q_name_for_worker(&client_ident, &worker_config);
+
+    let client = RetryClient::new(
+        SharedReplaceableClient::new(init_worker_client(
+            worker_config.namespace.clone(),
+            worker_config.client_identity_override.clone(),
+            client,
+        )),
+        RetryOptions::default(),
+    );
+    let client_ident = client.identity();
+    let sticky_q = sticky_q_name_for_worker(&client_ident, worker_config.max_cached_workflows);
 
     if client_ident.is_empty() {
         bail!("Client identity cannot be empty. Either lang or user should be setting this value");
     }
 
-    let heartbeat_fn = worker_config
-        .heartbeat_interval
-        .map(|_| Arc::new(OnceLock::new()));
-
     let client_bag = Arc::new(WorkerClientBag::new(
         client,
-        worker_config.namespace.clone(),
-        client_ident,
+        namespace.clone(),
+        client_ident.clone(),
         worker_config.versioning_strategy.clone(),
-        heartbeat_fn.clone(),
     ));
 
-    Ok(Worker::new(
-        worker_config,
+    Worker::new(
+        worker_config.clone(),
         sticky_q,
-        client_bag,
+        client_bag.clone(),
         Some(&runtime.telemetry),
-        heartbeat_fn,
-    ))
+        runtime.heartbeat_interval,
+    )
 }
 
 /// Create a worker for replaying one or more existing histories. It will auto-shutdown as soon as
@@ -141,24 +141,28 @@ where
     rwi.into_core_worker()
 }
 
-pub(crate) fn init_worker_client(
-    config: &WorkerConfig,
-    client: ConfiguredClient<TemporalServiceClientWithMetrics>,
-) -> RetryClient<Client> {
-    let mut client = Client::new(client, config.namespace.clone());
-    if let Some(ref id_override) = config.client_identity_override {
+pub(crate) fn init_worker_client<CT>(
+    namespace: String,
+    client_identity_override: Option<String>,
+    client: CT,
+) -> Client
+where
+    CT: Into<sealed::AnyClient>,
+{
+    let mut client = Client::new(*client.into().into_inner(), namespace.clone());
+    if let Some(ref id_override) = client_identity_override {
         client.options_mut().identity.clone_from(id_override);
     }
-    RetryClient::new(client, RetryConfig::default())
+    client
 }
 
 /// Creates a unique sticky queue name for a worker, iff the config allows for 1 or more cached
 /// workflows.
 pub(crate) fn sticky_q_name_for_worker(
     process_identity: &str,
-    config: &WorkerConfig,
+    max_cached_workflows: usize,
 ) -> Option<String> {
-    if config.max_cached_workflows > 0 {
+    if max_cached_workflows > 0 {
         Some(format!(
             "{}-{}",
             &process_identity,
@@ -171,44 +175,57 @@ pub(crate) fn sticky_q_name_for_worker(
 
 mod sealed {
     use super::*;
+    use temporalio_client::{SharedReplaceableClient, TemporalServiceClient};
 
     /// Allows passing different kinds of clients into things that want to be flexible. Motivating
     /// use-case was worker initialization.
     ///
     /// Needs to exist in this crate to avoid blanket impl conflicts.
     pub struct AnyClient {
-        pub(crate) inner: Box<ConfiguredClient<TemporalServiceClientWithMetrics>>,
+        pub(crate) inner: Box<ConfiguredClient<TemporalServiceClient>>,
     }
     impl AnyClient {
-        pub(crate) fn into_inner(self) -> Box<ConfiguredClient<TemporalServiceClientWithMetrics>> {
+        pub(crate) fn into_inner(self) -> Box<ConfiguredClient<TemporalServiceClient>> {
             self.inner
         }
     }
 
-    impl From<RetryClient<ConfiguredClient<TemporalServiceClientWithMetrics>>> for AnyClient {
-        fn from(c: RetryClient<ConfiguredClient<TemporalServiceClientWithMetrics>>) -> Self {
-            Self {
-                inner: Box::new(c.into_inner()),
-            }
+    impl From<ConfiguredClient<TemporalServiceClient>> for AnyClient {
+        fn from(c: ConfiguredClient<TemporalServiceClient>) -> Self {
+            Self { inner: Box::new(c) }
         }
     }
-    impl From<RetryClient<Client>> for AnyClient {
-        fn from(c: RetryClient<Client>) -> Self {
-            Self {
-                inner: Box::new(c.into_inner().into_inner()),
-            }
+
+    impl From<Client> for AnyClient {
+        fn from(c: Client) -> Self {
+            c.into_inner().into()
         }
     }
-    impl From<Arc<RetryClient<Client>>> for AnyClient {
-        fn from(c: Arc<RetryClient<Client>>) -> Self {
-            Self {
-                inner: Box::new(c.get_client().inner().clone()),
-            }
+
+    impl<T> From<RetryClient<T>> for AnyClient
+    where
+        T: Into<AnyClient>,
+    {
+        fn from(c: RetryClient<T>) -> Self {
+            c.into_inner().into()
         }
     }
-    impl From<ConfiguredClient<TemporalServiceClientWithMetrics>> for AnyClient {
-        fn from(c: ConfiguredClient<TemporalServiceClientWithMetrics>) -> Self {
-            Self { inner: Box::new(c) }
+
+    impl<T> From<SharedReplaceableClient<T>> for AnyClient
+    where
+        T: Into<AnyClient> + Clone + Send + Sync,
+    {
+        fn from(c: SharedReplaceableClient<T>) -> Self {
+            c.inner_clone().into()
+        }
+    }
+
+    impl<T> From<Arc<T>> for AnyClient
+    where
+        T: Into<AnyClient> + Clone,
+    {
+        fn from(c: Arc<T>) -> Self {
+            Arc::unwrap_or_clone(c).into()
         }
     }
 }
@@ -220,6 +237,37 @@ pub struct CoreRuntime {
     telemetry: TelemetryInstance,
     runtime: Option<tokio::runtime::Runtime>,
     runtime_handle: tokio::runtime::Handle,
+    heartbeat_interval: Option<Duration>,
+}
+
+/// Holds telemetry options, as well as worker heartbeat_interval. Construct with [RuntimeOptionsBuilder]
+#[derive(Default, derive_builder::Builder)]
+#[builder(build_fn(validate = "Self::validate"))]
+#[non_exhaustive]
+pub struct RuntimeOptions {
+    /// Telemetry configuration options.
+    #[builder(default)]
+    telemetry_options: TelemetryOptions,
+    /// Optional worker heartbeat interval - This configures the heartbeat setting of all
+    /// workers created using this runtime.
+    ///
+    /// Interval must be between 1s and 60s, inclusive.
+    #[builder(default = "Some(Duration::from_secs(60))")]
+    heartbeat_interval: Option<Duration>,
+}
+
+impl RuntimeOptionsBuilder {
+    fn validate(&self) -> Result<(), String> {
+        if let Some(Some(interval)) = self.heartbeat_interval
+            && (interval < Duration::from_secs(1) || interval > Duration::from_secs(60))
+        {
+            return Err(format!(
+                "heartbeat_interval ({interval:?}) must be between 1s and 60s",
+            ));
+        }
+
+        Ok(())
+    }
 }
 
 /// Wraps a [tokio::runtime::Builder] to allow layering multiple on_thread_start functions
@@ -254,13 +302,13 @@ impl CoreRuntime {
     /// If a tokio runtime has already been initialized. To re-use an existing runtime, call
     /// [CoreRuntime::new_assume_tokio].
     pub fn new<F>(
-        telemetry_options: TelemetryOptions,
+        runtime_options: RuntimeOptions,
         mut tokio_builder: TokioRuntimeBuilder<F>,
     ) -> Result<Self, anyhow::Error>
     where
         F: Fn() + Send + Sync + 'static,
     {
-        let telemetry = telemetry_init(telemetry_options)?;
+        let telemetry = telemetry_init(runtime_options.telemetry_options)?;
         let subscriber = telemetry.trace_subscriber();
         let runtime = tokio_builder
             .inner
@@ -275,7 +323,8 @@ impl CoreRuntime {
             })
             .build()?;
         let _rg = runtime.enter();
-        let mut me = Self::new_assume_tokio_initialized_telem(telemetry);
+        let mut me =
+            Self::new_assume_tokio_initialized_telem(telemetry, runtime_options.heartbeat_interval);
         me.runtime = Some(runtime);
         Ok(me)
     }
@@ -285,9 +334,12 @@ impl CoreRuntime {
     ///
     /// # Panics
     /// If there is no currently active Tokio runtime
-    pub fn new_assume_tokio(telemetry_options: TelemetryOptions) -> Result<Self, anyhow::Error> {
-        let telemetry = telemetry_init(telemetry_options)?;
-        Ok(Self::new_assume_tokio_initialized_telem(telemetry))
+    pub fn new_assume_tokio(runtime_options: RuntimeOptions) -> Result<Self, anyhow::Error> {
+        let telemetry = telemetry_init(runtime_options.telemetry_options)?;
+        Ok(Self::new_assume_tokio_initialized_telem(
+            telemetry,
+            runtime_options.heartbeat_interval,
+        ))
     }
 
     /// Construct a runtime from an already-initialized telemetry instance, assuming a tokio runtime
@@ -295,7 +347,10 @@ impl CoreRuntime {
     ///
     /// # Panics
     /// If there is no currently active Tokio runtime
-    pub fn new_assume_tokio_initialized_telem(telemetry: TelemetryInstance) -> Self {
+    pub fn new_assume_tokio_initialized_telem(
+        telemetry: TelemetryInstance,
+        heartbeat_interval: Option<Duration>,
+    ) -> Self {
         let runtime_handle = tokio::runtime::Handle::current();
         if let Some(sub) = telemetry.trace_subscriber() {
             set_trace_subscriber_for_current_thread(sub);
@@ -304,6 +359,7 @@ impl CoreRuntime {
             telemetry,
             runtime: None,
             runtime_handle,
+            heartbeat_interval,
         }
     }
 

package/sdk-core/{core → crates/sdk-core}/src/pollers/mod.rs

@@ -3,9 +3,9 @@ mod poll_buffer;
 pub(crate) use poll_buffer::{
     ActivityTaskOptions, LongPollBuffer, WorkflowTaskOptions, WorkflowTaskPoller,
 };
-pub use temporal_client::{
-    Client, ClientOptions, ClientOptionsBuilder, ClientTlsConfig, RetryClient, RetryConfig,
-    TlsConfig, WorkflowClientTrait,
+pub use temporalio_client::{
+    Client, ClientOptions, ClientTlsOptions, RetryClient, RetryOptions, TlsOptions,
+    WorkflowClientTrait,
 };
 
 use crate::{
@@ -15,23 +15,25 @@ use crate::{
 use anyhow::{anyhow, bail};
 use futures_util::{Stream, stream};
 use std::{fmt::Debug, marker::PhantomData};
-use temporal_sdk_core_api::worker::{ActivitySlotKind, NexusSlotKind, SlotKind, WorkflowSlotKind};
-use temporal_sdk_core_protos::temporal::api::workflowservice::v1::{
-    PollActivityTaskQueueResponse, PollNexusTaskQueueResponse, PollWorkflowTaskQueueResponse,
+use temporalio_common::{
+    protos::temporal::api::workflowservice::v1::{
+        PollActivityTaskQueueResponse, PollNexusTaskQueueResponse, PollWorkflowTaskQueueResponse,
+    },
+    worker::{ActivitySlotKind, NexusSlotKind, SlotKind, WorkflowSlotKind},
 };
 use tokio::select;
 use tokio_util::sync::CancellationToken;
 
-#[cfg(test)]
+#[cfg(any(feature = "test-utilities", test))]
 use futures_util::Future;
-#[cfg(test)]
+#[cfg(any(feature = "test-utilities", test))]
 pub(crate) use poll_buffer::MockPermittedPollBuffer;
 
 pub(crate) type Result<T, E = tonic::Status> = std::result::Result<T, E>;
 
 /// A trait for things that long poll the server.
-#[cfg_attr(test, mockall::automock)]
-#[cfg_attr(test, allow(unused))]
+#[cfg_attr(any(feature = "test-utilities", test), mockall::automock)]
+#[cfg_attr(any(feature = "test-utilities", test), allow(unused))]
 #[async_trait::async_trait]
 pub(crate) trait Poller<PollResult>
 where
@@ -79,7 +81,7 @@ where
     }
 }
 
-#[cfg(test)]
+#[cfg(any(feature = "test-utilities", test))]
 mockall::mock! {
     pub ManualPoller<T: Send + Sync + 'static> {}
     #[allow(unused)]

package/sdk-core/{core → crates/sdk-core}/src/pollers/poll_buffer.rs

@@ -6,6 +6,7 @@ use crate::{
         client::{PollActivityOptions, PollOptions, PollWorkflowOptions, WorkerClient},
     },
 };
+use crossbeam_utils::atomic::AtomicCell;
 use futures_util::{FutureExt, StreamExt, future::BoxFuture};
 use governor::{Quota, RateLimiter};
 use std::{
@@ -16,17 +17,20 @@ use std::{
         Arc,
         atomic::{AtomicBool, AtomicUsize, Ordering},
     },
-    time::Duration,
+    time::{Duration, SystemTime},
 };
-use temporal_client::{ERROR_RETURNED_DUE_TO_SHORT_CIRCUIT, NoRetryOnMatching};
-use temporal_sdk_core_api::worker::{
-    ActivitySlotKind, NexusSlotKind, PollerBehavior, SlotKind, WorkflowSlotKind,
+use temporalio_client::{
+    ERROR_RETURNED_DUE_TO_SHORT_CIRCUIT, request_extensions::NoRetryOnMatching,
 };
-use temporal_sdk_core_protos::temporal::api::{
-    taskqueue::v1::PollerScalingDecision,
-    workflowservice::v1::{
-        PollActivityTaskQueueResponse, PollNexusTaskQueueResponse, PollWorkflowTaskQueueResponse,
+use temporalio_common::{
+    protos::temporal::api::{
+        taskqueue::v1::PollerScalingDecision,
+        workflowservice::v1::{
+            PollActivityTaskQueueResponse, PollNexusTaskQueueResponse,
+            PollWorkflowTaskQueueResponse,
+        },
     },
+    worker::{ActivitySlotKind, NexusSlotKind, PollerBehavior, SlotKind, WorkflowSlotKind},
 };
 use tokio::{
     sync::{
@@ -74,9 +78,15 @@ impl LongPollBuffer<PollWorkflowTaskQueueResponse, WorkflowSlotKind> {
         shutdown: CancellationToken,
         num_pollers_handler: Option<impl Fn(usize) + Send + Sync + 'static>,
         options: WorkflowTaskOptions,
+        last_successful_poll_time: Arc<AtomicCell<Option<SystemTime>>>,
     ) -> Self {
         let is_sticky = sticky_queue.is_some();
-        let poll_scaler = PollScaler::new(poller_behavior, num_pollers_handler, shutdown.clone());
+        let poll_scaler = PollScaler::new(
+            poller_behavior,
+            num_pollers_handler,
+            shutdown.clone(),
+            last_successful_poll_time,
+        );
         if let Some(wftps) = options.wft_poller_shared.as_ref() {
             if is_sticky {
                 wftps.set_sticky_active(poll_scaler.active_rx.clone());
@@ -136,6 +146,7 @@ impl LongPollBuffer<PollWorkflowTaskQueueResponse, WorkflowSlotKind> {
 }
 
 impl LongPollBuffer<PollActivityTaskQueueResponse, ActivitySlotKind> {
+    #[allow(clippy::too_many_arguments)]
     pub(crate) fn new_activity_task(
         client: Arc<dyn WorkerClient>,
         task_queue: String,
@@ -144,6 +155,7 @@ impl LongPollBuffer<PollActivityTaskQueueResponse, ActivitySlotKind> {
         shutdown: CancellationToken,
         num_pollers_handler: Option<impl Fn(usize) + Send + Sync + 'static>,
         options: ActivityTaskOptions,
+        last_successful_poll_time: Arc<AtomicCell<Option<SystemTime>>>,
     ) -> Self {
         let pre_permit_delay = options
             .max_worker_acts_per_second
@@ -183,7 +195,12 @@ impl LongPollBuffer<PollActivityTaskQueueResponse, ActivitySlotKind> {
             }
         };
 
-        let poll_scaler = PollScaler::new(poller_behavior, num_pollers_handler, shutdown.clone());
+        let poll_scaler = PollScaler::new(
+            poller_behavior,
+            num_pollers_handler,
+            shutdown.clone(),
+            last_successful_poll_time,
+        );
         Self::new(
             poll_fn,
             permit_dealer,
@@ -196,6 +213,7 @@ impl LongPollBuffer<PollActivityTaskQueueResponse, ActivitySlotKind> {
 }
 
 impl LongPollBuffer<PollNexusTaskQueueResponse, NexusSlotKind> {
+    #[allow(clippy::too_many_arguments)]
     pub(crate) fn new_nexus_task(
         client: Arc<dyn WorkerClient>,
         task_queue: String,
@@ -203,6 +221,8 @@ impl LongPollBuffer<PollNexusTaskQueueResponse, NexusSlotKind> {
         permit_dealer: MeteredPermitDealer<NexusSlotKind>,
         shutdown: CancellationToken,
         num_pollers_handler: Option<impl Fn(usize) + Send + Sync + 'static>,
+        last_successful_poll_time: Arc<AtomicCell<Option<SystemTime>>>,
+        send_heartbeat: bool,
     ) -> Self {
         let no_retry = if matches!(poller_behavior, PollerBehavior::Autoscaling { .. }) {
             Some(NoRetryOnMatching {
@@ -216,11 +236,14 @@ impl LongPollBuffer<PollNexusTaskQueueResponse, NexusSlotKind> {
             let task_queue = task_queue.clone();
             async move {
                 client
-                    .poll_nexus_task(PollOptions {
-                        task_queue,
-                        no_retry,
-                        timeout_override,
-                    })
+                    .poll_nexus_task(
+                        PollOptions {
+                            task_queue,
+                            no_retry,
+                            timeout_override,
+                        },
+                        send_heartbeat,
+                    )
                     .await
             }
         };
@@ -228,7 +251,12 @@ impl LongPollBuffer<PollNexusTaskQueueResponse, NexusSlotKind> {
             poll_fn,
             permit_dealer,
             shutdown.clone(),
-            PollScaler::new(poller_behavior, num_pollers_handler, shutdown),
+            PollScaler::new(
+                poller_behavior,
+                num_pollers_handler,
+                shutdown,
+                last_successful_poll_time,
+            ),
             None::<fn() -> BoxFuture<'static, ()>>,
             None::<fn(&PollNexusTaskQueueResponse)>,
         )
@@ -413,6 +441,7 @@ where
         behavior: PollerBehavior,
         num_pollers_handler: Option<F>,
         shutdown: CancellationToken,
+        last_successful_poll_time: Arc<AtomicCell<Option<SystemTime>>>,
     ) -> Self {
         let (active_tx, active_rx) = watch::channel(0);
         let num_pollers_handler = num_pollers_handler.map(Arc::new);
@@ -433,6 +462,7 @@ where
             ingested_this_period: Default::default(),
             ingested_last_period: Default::default(),
             scale_up_allowed: AtomicBool::new(true),
+            last_successful_poll_time,
         });
         let rhc = report_handle.clone();
         let ingestor_task = if behavior.is_autoscaling() {
@@ -495,6 +525,7 @@ struct PollScalerReportHandle {
     ingested_this_period: AtomicUsize,
     ingested_last_period: AtomicUsize,
     scale_up_allowed: AtomicBool,
+    last_successful_poll_time: Arc<AtomicCell<Option<SystemTime>>>,
 }
 
 impl PollScalerReportHandle {
@@ -502,6 +533,8 @@ impl PollScalerReportHandle {
     fn poll_result(&self, res: &Result<impl TaskPollerResult, tonic::Status>) -> bool {
         match res {
             Ok(res) => {
+                self.last_successful_poll_time
+                    .store(Some(SystemTime::now()));
                 if let PollerBehavior::SimpleMaximum(_) = self.behavior {
                     // We don't do auto-scaling with the simple max
                     return true;
@@ -666,14 +699,14 @@ impl TaskPollerResult for PollNexusTaskQueueResponse {
     }
 }
 
-#[cfg(test)]
+#[cfg(any(feature = "test-utilities", test))]
 #[derive(derive_more::Constructor)]
 pub(crate) struct MockPermittedPollBuffer<PT, SK: SlotKind> {
     sem: Arc<MeteredPermitDealer<SK>>,
     inner: PT,
 }
 
-#[cfg(test)]
+#[cfg(any(feature = "test-utilities", test))]
 #[async_trait::async_trait]
 impl<T, PT, SK> Poller<(T, OwnedMeteredSemPermit<SK>)> for MockPermittedPollBuffer<PT, SK>
 where
@@ -735,6 +768,7 @@ mod tests {
             WorkflowTaskOptions {
                 wft_poller_shared: Some(Arc::new(WFTPollerShared::new(Some(10)))),
             },
+            Arc::new(AtomicCell::new(None)),
         );
 
         // Poll a bunch of times, "interrupting" it each time, we should only actually have polled
@@ -790,6 +824,7 @@ mod tests {
             WorkflowTaskOptions {
                 wft_poller_shared: Some(Arc::new(WFTPollerShared::new(Some(1)))),
             },
+            Arc::new(AtomicCell::new(None)),
         );
 
         // Should not see error, unwraps should get empty response