taquba_workflow/

runner.rs

use std::collections::HashMap;
use std::future::Future;
use std::time::Duration;

use taquba::LeaseHandle;
use tokio_util::sync::CancellationToken;

use crate::effects::EffectsHandle;
use crate::kv::KvReadHandle;
use crate::memo::Memo;

/// A single step within a workflow run, handed to [`StepRunner::run_step`].
///
/// Mirrors [`taquba::JobRecord`]: the `payload` is opaque application bytes
/// and `headers` carries user metadata you set at submission (reserved
/// `workflow.*` keys are filtered out before the runner sees them).
#[derive(Debug, Clone)]
pub struct Step {
    /// Caller-visible run identifier (the value passed to or generated by
    /// [`crate::RunSpec`]).
    pub run_id: String,
    /// Zero-based step number. Step 0 is always the first step of a run, with
    /// the original submission input as its `payload`.
    pub step_number: u32,
    /// Application-defined bytes. For step 0 this is the submission `input`;
    /// for later steps it is the bytes returned by the previous step's
    /// [`StepOutcome::Continue`].
    pub payload: Vec<u8>,
    /// Submitter-supplied metadata, threaded through every step of the run.
    /// Reserved `workflow.*` headers are stripped before the runner sees them.
    pub headers: HashMap<String, String>,
    /// The Taquba job ID for this step, useful for tracing.
    pub job_id: String,
    /// How many times Taquba has attempted to deliver this step. `1` on the
    /// first attempt; `>1` after a lease expiry / nack retry.
    pub attempts: u32,
    /// The attempt limit for this step: a transient failure on the attempt
    /// numbered `max_attempts` ends the run.
    pub max_attempts: u32,
    /// Cooperative cancellation signal for the run. The runtime cancels
    /// this token when [`crate::WorkflowRuntime::cancel`] is called while
    /// this step is in flight, so a long-running runner (e.g. an LLM call,
    /// a slow HTTP request) can short-circuit instead of running to
    /// completion. Typical use:
    ///
    /// ```ignore
    /// tokio::select! {
    ///     out = do_slow_work(step) => out,
    ///     _ = step.cancel_token.cancelled() => {
    ///         Ok(StepOutcome::Cancel { reason: "cooperative".into() })
    ///     }
    /// }
    /// ```
    ///
    /// Runners that ignore the token remain correct: the runtime still
    /// discards the outcome of a cancelled step and fires the terminal
    /// hook with [`crate::TerminalStatus::Cancelled`]. Watching the token
    /// only reduces cancellation latency for slow steps; it doesn't
    /// change semantics.
    ///
    /// The token is a child of the delivery's claim token. A re-delivery
    /// of this step observes `is_cancelled() == true` immediately,
    /// because the queue re-fires the claim's token from the job's
    /// persisted cancellation. Cancelling this token leaves the claim's
    /// token uncancelled, so the runtime does not treat the step as
    /// externally cancelled.
    pub cancel_token: CancellationToken,
    /// The lease handle for this step's delivery. A long-running runner
    /// calls [`LeaseHandle::ensure_at_least`] at progress points (or
    /// once, with a slow call's timeout, before issuing it) so the step
    /// is not re-queued while it still runs. Use
    /// [`LeaseHandle::detached`] when constructing a `Step` in tests; a
    /// detached handle's calls succeed without effect.
    pub lease: LeaseHandle,
    /// Per-step durable key-value store, scoped to this step's
    /// `(run_id, step_number)`. Use to memoize expensive within-step
    /// side effects (LLM calls, paid APIs) so an at-least-once retry
    /// of this step doesn't re-pay for work the prior attempt already
    /// did:
    ///
    /// ```ignore
    /// let response = match step.memo.get("llm").await? {
    ///     Some(cached) => deserialize(&cached),
    ///     None => {
    ///         let fresh = llm.complete(&prompt).await?;
    ///         step.memo.put("llm", &serialize(&fresh)).await?;
    ///         fresh
    ///     }
    /// };
    /// ```
    ///
    /// See [`Memo`] for the full API.
    pub memo: Memo,
    /// Run-scoped durable key-value store, shared by every step of the
    /// run. Entries live beside the per-step [`Step::memo`] entries and
    /// are removed with them when the run's retention expires. Use it
    /// for values a later step reads back, such as an accumulating
    /// journal; the durable channel for the next step's input is
    /// [`StepOutcome::Continue`]'s payload.
    pub run_memo: Memo,
    /// Application KV effects for this step. Writes and deletes staged
    /// here are applied in the same transaction as the settlement that
    /// commits the returned outcome, so application state cannot
    /// diverge from the run's transition on a crash:
    ///
    /// ```ignore
    /// step.effects.put(format!("app/runs/{}", step.run_id), b"done".to_vec())?;
    /// Ok(StepOutcome::Succeed { result })
    /// ```
    ///
    /// See [`EffectsHandle`] for the staging rules. Use
    /// [`EffectsHandle::detached`] when constructing a `Step` in tests.
    pub effects: EffectsHandle,
    /// Read access to the caller KV namespace. A committed value (an
    /// earlier step's applied effect, a [`crate::RunSpec::kv_writes`]
    /// entry, a direct [`taquba::Queue::kv_put`]) is readable here;
    /// effects staged by this step become readable only after it
    /// settles. The intended use is a read-then-stage marker check:
    ///
    /// ```ignore
    /// if step.kv.get(b"app/indexed/doc-1").await?.is_none() {
    ///     index_document(&step.payload).await?; // idempotent
    ///     step.effects.put("app/indexed/doc-1", b"1".to_vec())?;
    /// }
    /// ```
    ///
    /// See [`KvReadHandle`] for the read semantics. Use
    /// [`KvReadHandle::detached`] when constructing a `Step` in tests.
    pub kv: KvReadHandle,
    /// The signal payload, when this step was reached through a
    /// [`Trigger::OnSignal`] wait that a signal resolved: the previous
    /// step continued with `OnSignal`, and a
    /// [`crate::WorkflowRuntime::signal`] call for the correlation key
    /// arrived before the timeout. `None` when the timeout elapsed first
    /// and on every step not preceded by an `OnSignal` wait.
    pub signal: Option<Vec<u8>>,
}

/// When the next step of a run becomes claimable. Set on the `when`
/// field of [`StepOutcome::Continue`].
#[derive(Debug, Clone)]
#[non_exhaustive]
pub enum Trigger {
    /// The next step is claimable immediately.
    Immediate,
    /// The next step is claimable `Duration` from now.
    After(Duration),
    /// The next step is claimable when a signal for `correlation_key`
    /// arrives via [`crate::WorkflowRuntime::signal`], or after `timeout`,
    /// whichever comes first. The next step reads [`Step::signal`] to
    /// distinguish the two: `Some(payload)` when a signal arrived, `None`
    /// when the timeout elapsed first. A signal that arrived before this
    /// step settled is consumed at settlement and the next step runs
    /// immediately.
    ///
    /// One waiter per correlation key: registering a second waiter while
    /// one is already waiting fails the step permanently. Choose keys
    /// that are unique per waiter (e.g. include the run id).
    OnSignal {
        /// Caller-chosen key the signaller addresses.
        correlation_key: String,
        /// Upper bound on the wait; the next step runs with
        /// [`Step::signal`] `None` when it elapses first.
        timeout: Duration,
    },
}

/// What the runner wants the runtime to do after this step.
#[derive(Debug, Clone)]
#[non_exhaustive]
pub enum StepOutcome {
    /// Run is not finished. Enqueue the next step with `payload` as its
    /// bytes; `when` decides when it becomes claimable. The runtime
    /// advances `step_number` by 1. The constructors
    /// [`Self::continue_now`] and [`Self::continue_after`] build the
    /// common forms.
    Continue {
        /// Bytes to hand to the next step's [`Step::payload`].
        payload: Vec<u8>,
        /// When the next step becomes claimable.
        when: Trigger,
    },
    /// The run is finished successfully. The runtime acks the step and fires
    /// the configured terminal hook with
    /// [`crate::TerminalStatus::Succeeded`] and `result` as the body.
    Succeed {
        /// Final result bytes handed to the terminal hook.
        result: Vec<u8>,
    },
    /// The run is finished as failed by the runner's verdict; the runner
    /// ran to completion but the workflow's logical outcome is "no" (e.g.
    /// a validation rule rejected the input, a policy check denied the
    /// request, an agent decided the task can't be fulfilled). The runtime
    /// acks the step and fires the terminal hook with
    /// [`crate::TerminalStatus::Failed`] and `reason` as the error.
    ///
    /// Use this for *workflow-level* failures. For *infrastructure*
    /// failures (network outage, downstream service down, etc.) return
    /// `Err(StepError::transient)` or `Err(StepError::permanent)` instead;
    /// those dead-letter the step so an operator can find it via
    /// [`taquba::Queue::dead_jobs`]. `Fail` is a successful execution with
    /// a negative outcome and does not dead-letter.
    Fail {
        /// Human-readable reason recorded on [`crate::RunOutcome::error`].
        reason: String,
    },
    /// The run is finished as cancelled by the runner. Use this when the
    /// runner decides on its own that the workflow should stop early
    /// without it being a logical failure (e.g. a downstream cancellation
    /// signal arrived mid-step, the user-supplied input is now obsolete).
    /// The runtime acks the step and fires the terminal hook with
    /// [`crate::TerminalStatus::Cancelled`] and `reason` as the error.
    ///
    /// For *external* cancellation requested by another component in the
    /// process, call [`crate::WorkflowRuntime::cancel`] instead; the
    /// runtime translates that into the same `Cancelled` terminal state.
    Cancel {
        /// Human-readable reason recorded on [`crate::RunOutcome::error`].
        reason: String,
    },
}

impl StepOutcome {
    /// Continue the run; the next step is claimable immediately.
    pub fn continue_now(payload: Vec<u8>) -> Self {
        Self::Continue {
            payload,
            when: Trigger::Immediate,
        }
    }

    /// Continue the run; the next step is claimable `delay` from now.
    pub fn continue_after(payload: Vec<u8>, delay: Duration) -> Self {
        Self::Continue {
            payload,
            when: Trigger::After(delay),
        }
    }

    /// Continue the run; the next step is claimable when a signal for
    /// `correlation_key` arrives, or after `timeout` at the latest.
    pub fn continue_on_signal(
        payload: Vec<u8>,
        correlation_key: impl Into<String>,
        timeout: Duration,
    ) -> Self {
        Self::Continue {
            payload,
            when: Trigger::OnSignal {
                correlation_key: correlation_key.into(),
                timeout,
            },
        }
    }
}

/// Failure outcomes the runner can return.
#[derive(Debug)]
pub struct StepError {
    /// Human-readable message recorded on the underlying job's `last_error`.
    pub message: String,
    /// Whether to retry the step or fail the run immediately.
    pub kind: StepErrorKind,
}

impl StepError {
    /// Build a transient error: Taquba retries the step per the queue's
    /// backoff/`max_attempts`. Once `max_attempts` is exhausted, the step is
    /// dead-lettered and the run terminates as failed.
    pub fn transient(message: impl Into<String>) -> Self {
        Self {
            message: message.into(),
            kind: StepErrorKind::Transient,
        }
    }

    /// Build a permanent error: the step is dead-lettered immediately and the
    /// run terminates as failed.
    pub fn permanent(message: impl Into<String>) -> Self {
        Self {
            message: message.into(),
            kind: StepErrorKind::Permanent,
        }
    }
}

impl std::fmt::Display for StepError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(&self.message)
    }
}

impl std::error::Error for StepError {}

impl From<crate::Error> for StepError {
    fn from(err: crate::Error) -> Self {
        let permanent = err.is_permanent();
        let message = err.to_string();
        if permanent {
            Self::permanent(message)
        } else {
            Self::transient(message)
        }
    }
}

/// Whether a [`StepError`] should retry or fail the run.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum StepErrorKind {
    /// Retry per the queue's backoff policy until `max_attempts` is reached.
    Transient,
    /// Dead-letter the step immediately; terminate the run as failed.
    Permanent,
}

/// User-implemented logic that advances a single workflow step.
///
/// Implementations must be idempotent for the same `(run_id, step_number)`:
/// Taquba is at-least-once, so a step can be claimed and processed more than
/// once if a lease expires before the worker acks. Returning the same
/// `StepOutcome` for the same input is the easiest way to satisfy this.
pub trait StepRunner: Send + Sync {
    /// Process a single step of a workflow run. Return [`StepOutcome::Continue`]
    /// to enqueue the next step, [`StepOutcome::Succeed`] to finish the run
    /// successfully, [`StepOutcome::Fail`] to terminate the run as Failed by
    /// runner verdict, [`StepOutcome::Cancel`] to terminate the run as
    /// Cancelled by runner verdict, or `Err(StepError)` to retry /
    /// dead-letter on infrastructure errors.
    fn run_step(
        &self,
        step: &Step,
    ) -> impl Future<Output = std::result::Result<StepOutcome, StepError>> + Send;
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn from_workflow_error_maps_via_is_permanent() {
        let permanent: StepError = crate::Error::InputMismatch("run-1".into()).into();
        assert_eq!(permanent.kind, StepErrorKind::Permanent);

        let store_err = taquba::object_store::Error::NotFound {
            path: "x".into(),
            source: "missing".into(),
        };
        let transient: StepError = crate::Error::Store(store_err).into();
        assert_eq!(transient.kind, StepErrorKind::Transient);
    }
}