do the heartbeat handling here

Author: virajmehta

src/context.rs

@@ -7,7 +7,7 @@ use uuid::Uuid;
 
 use crate::Durable;
 use crate::error::{ControlFlow, TaskError, TaskResult};
-use crate::heartbeat::{HeartbeatHandle, Heartbeater};
+use crate::heartbeat::{HeartbeatHandle, Heartbeater, StepState};
 use crate::task::Task;
 use crate::types::DurableEventPayload;
 use crate::types::{
@@ -165,9 +165,9 @@ where
     /// # Example
     ///
     /// ```ignore
-    /// let payment_id = ctx.step("charge-payment", ctx.task_id, |task_id, _state| async {
+    /// let payment_id = ctx.step("charge-payment", ctx.task_id, |task_id, step_state| async {
     ///     let idempotency_key = format!("{}:charge", task_id);
-    ///     stripe::charge(amount, &idempotency_key).await
+    ///     stripe::charge(amount, &idempotency_key, &step_state.state).await
     /// }).await?;
     /// ```
     #[cfg_attr(
@@ -182,7 +182,7 @@ where
         &mut self,
         base_name: &str,
         params: P,
-        f: fn(P, State) -> Fut,
+        f: fn(P, StepState<State>) -> Fut,
     ) -> TaskResult<T>
     where
         P: Serialize,
@@ -206,13 +206,14 @@ where
         span.record("cached", false);
 
         // Execute the step
-        let result =
-            f(params, self.durable.state().clone())
-                .await
-                .map_err(|e| TaskError::Step {
-                    base_name: base_name.to_string(),
-                    error: e,
-                })?;
+        let step_state = StepState {
+            state: self.durable.state().clone(),
+            heartbeater: self.heartbeat_handle.clone(),
+        };
+        let result = f(params, step_state).await.map_err(|e| TaskError::Step {
+            base_name: base_name.to_string(),
+            error: e,
+        })?;
 
         // Persist checkpoint (also extends claim lease)
         #[cfg(feature = "telemetry")]

src/heartbeat.rs

@@ -96,3 +96,28 @@ impl Heartbeater for NoopHeartbeater {
         Ok(())
     }
 }
+
+/// State provided to `step()` closures, wrapping the user's application state
+/// alongside a [`HeartbeatHandle`] for extending the task lease.
+///
+/// This is passed as the second argument to every `step()` closure, making
+/// heartbeating available without the consumer needing to thread it manually.
+///
+/// # Example
+///
+/// ```ignore
+/// ctx.step("long-operation", params, |params, step_state| async move {
+///     for item in &params.items {
+///         process(item, &step_state.state).await?;
+///         // Extend lease during long-running work
+///         let _ = step_state.heartbeater.heartbeat(None).await;
+///     }
+///     Ok(result)
+/// }).await?;
+/// ```
+pub struct StepState<State> {
+    /// The user's application state.
+    pub state: State,
+    /// Handle for extending the task lease during long-running operations.
+    pub heartbeater: HeartbeatHandle,
+}

src/lib.rs

@@ -110,7 +110,7 @@ pub use client::{Durable, DurableBuilder};
 pub use context::TaskContext;
 pub use cron::{ScheduleFilter, ScheduleInfo, ScheduleOptions, setup_pgcron};
 pub use error::{ControlFlow, DurableError, DurableResult, TaskError, TaskResult};
-pub use heartbeat::{HeartbeatHandle, Heartbeater, NoopHeartbeater};
+pub use heartbeat::{HeartbeatHandle, Heartbeater, NoopHeartbeater, StepState};
 pub use task::{ErasedTask, Task, TaskWrapper};
 pub use types::{
     CancellationPolicy, ClaimedTask, DurableEventPayload, RetryStrategy, SpawnDefaults,

tests/execution_test.rs

@@ -745,13 +745,13 @@ impl durable::Task<AppState> for WriteToDbTask {
     ) -> durable::TaskResult<Self::Output> {
         // Use the app state's db pool to write to a table
         let row_id: i64 = ctx
-            .step("insert", params, |params, state| async move {
+            .step("insert", params, |params, step_state| async move {
                 let (id,): (i64,) = sqlx::query_as(
                     "INSERT INTO test_state_table (key, value) VALUES ($1, $2) RETURNING id",
                 )
                 .bind(&params.key)
                 .bind(&params.value)
-                .fetch_one(&state.db_pool)
+                .fetch_one(&step_state.state.db_pool)
                 .await
                 .map_err(|e| anyhow::anyhow!("DB error: {}", e))?;
                 Ok(id)