aion_worker/

context.rs

//! `ActivityContext` heartbeat, cancellation, attempt, and identifier support.

use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};

use aion_core::{ActivityEvent, ActivityId, Payload, RunId, WorkflowId};
use tokio::sync::{Notify, mpsc};

use crate::error::WorkerError;

/// Handler-facing context for one activity execution.
#[derive(Clone, Debug)]
pub struct ActivityContext {
    /// The workflow this activity belongs to.
    workflow_id: WorkflowId,
    /// The concrete run this activity was dispatched by — the generation axis.
    ///
    /// REQUIRED, never optional: [`ActivityContext::run_id`] is the only source
    /// a handler has for the run axis its transcript events must carry, and
    /// [`aion_core::ActivityEvent::run_id`] is itself a required field. An
    /// optional accessor would leave a handler with no honest way to build a
    /// required field, so the run is supplied at construction or the context is
    /// never built — the refusal lives at the dispatch boundary, not here.
    run_id: RunId,
    activity_id: ActivityId,
    attempt: u32,
    idempotency_key: Option<String>,
    cancellation: Arc<CancellationState>,
    heartbeat_sender: Option<mpsc::UnboundedSender<HeartbeatRequest>>,
    /// NOI-5b agent-observability event seam (additive, OPTIONAL). A running
    /// activity (or the harness adapter driving it) emits neutral
    /// [`ActivityEvent`]s here; the worker runtime drains them and forwards them
    /// to the server's transcript sequencer over the same transport activity
    /// results take. A context created WITHOUT this seam — every isolated unit
    /// test and every activity that emits nothing — is a no-op, byte-identical to
    /// today, exactly as the `heartbeat_sender` seam is.
    event_sender: Option<mpsc::UnboundedSender<ActivityEvent>>,
}

/// Internal handle used by the worker runtime to signal cooperative cancellation.
#[derive(Clone, Debug)]
pub struct ActivityCancellationHandle {
    cancellation: Arc<CancellationState>,
}

/// Heartbeat request emitted by [`ActivityContext::heartbeat`].
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct HeartbeatRequest {
    /// Workflow owning the activity whose progress is being reported.
    pub workflow_id: WorkflowId,
    /// Activity whose progress is being reported.
    pub activity_id: ActivityId,
    /// Opaque progress detail supplied by the handler.
    pub detail: Option<Payload>,
}

#[derive(Debug)]
struct CancellationState {
    cancelled: AtomicBool,
    notify: Notify,
}

impl ActivityContext {
    /// Creates a context and the internal handle that can signal cancellation.
    ///
    /// The full dispatch identity is required: an activity execution always
    /// belongs to one `(workflow, run, activity, attempt)`, and a handler reads
    /// the workflow and run back to stamp the transcript events it emits.
    #[must_use]
    pub fn new(
        workflow_id: WorkflowId,
        run_id: RunId,
        activity_id: ActivityId,
        attempt: u32,
    ) -> (Self, ActivityCancellationHandle) {
        Self::for_workflow(workflow_id, run_id, activity_id, attempt, None)
    }

    /// Creates a context whose transcript seam is live, for a host that owns
    /// the receiving end of `events`.
    ///
    /// This is the seam [`Self::emit_event`] publishes on: every event a handler
    /// emits carries this context's `(workflow_id, run_id, activity_id, attempt)`
    /// identity, which is exactly the key the server's transcript sequencer
    /// files it under. A host that executes an activity IN PROCESS (the server's
    /// declared-command path) uses this to hand its own publisher the same
    /// stream a remote worker's drain would have delivered.
    #[must_use]
    pub fn with_transcript(
        workflow_id: WorkflowId,
        run_id: RunId,
        activity_id: ActivityId,
        attempt: u32,
        events: mpsc::UnboundedSender<ActivityEvent>,
    ) -> (Self, ActivityCancellationHandle) {
        Self::for_workflow_with_events(
            workflow_id,
            run_id,
            activity_id,
            attempt,
            None,
            None,
            Some(events),
        )
    }

    /// Returns this activity's identifier.
    #[must_use]
    pub const fn activity_id(&self) -> &ActivityId {
        &self.activity_id
    }

    /// Returns the workflow this activity belongs to.
    #[must_use]
    pub const fn workflow_id(&self) -> &WorkflowId {
        &self.workflow_id
    }

    /// Returns this activity's attempt number.
    #[must_use]
    pub const fn attempt(&self) -> u32 {
        self.attempt
    }

    /// Returns the concrete run this activity was dispatched by.
    ///
    /// A handler that emits transcript events through [`Self::emit_event`]
    /// stamps this onto every [`ActivityEvent`] it builds: the transcript
    /// keyspace is keyed on `(workflow, run, activity, attempt)`, and without
    /// the run two generations of one continue-as-new chain write to the same
    /// stream. Always present — the run is part of the dispatch identity, so
    /// the handler is never handed an absence it cannot resolve.
    #[must_use]
    pub const fn run_id(&self) -> &RunId {
        &self.run_id
    }

    /// Returns the stable external-effect key delivered with this task.
    ///
    /// Live worker tasks always return `Some`; manually constructed unit-test
    /// contexts return `None` because they are not attached to a server task.
    #[must_use]
    pub fn idempotency_key(&self) -> Option<&str> {
        self.idempotency_key.as_deref()
    }

    /// Emits a cooperative heartbeat request for this activity.
    ///
    /// This is the PROGRESS channel: handlers call it to attach a progress
    /// payload to the activity's liveness record. LIVENESS itself is owned by
    /// the worker runtime, which automatically heartbeats every in-flight
    /// activity within the server-assigned heartbeat window — a handler that
    /// never calls this still stays live for as long as it genuinely runs.
    /// Contexts created without a live heartbeat sender remain no-op contexts
    /// for isolated unit tests.
    ///
    /// # Errors
    ///
    /// Returns [`WorkerError`] when an installed heartbeat seam has been closed.
    pub fn heartbeat(&self, detail: Option<Payload>) -> Result<(), WorkerError> {
        if let Some(sender) = &self.heartbeat_sender {
            sender
                .send(HeartbeatRequest {
                    workflow_id: self.workflow_id.clone(),
                    activity_id: self.activity_id.clone(),
                    detail,
                })
                .map_err(|source| WorkerError::registration(HeartbeatSeamClosed { source }))?;
        }
        Ok(())
    }

    /// Emit a neutral agent-observability [`ActivityEvent`] onto the transcript
    /// seam (NOI-5b).
    ///
    /// Additive and OPTIONAL: on a context created without a live event seam
    /// (every isolated unit test, and every activity that does not run an
    /// instrumented agent) this is a no-op returning `Ok(())`, so behaviour is
    /// byte-identical to today. When a seam is installed the worker runtime drains
    /// these events and forwards them to the server's transcript sequencer, which
    /// stamps the commit-allocated `store_seq` — the producer never assigns it.
    ///
    /// Harness-neutral: the payload is a pure `aion-core` [`ActivityEvent`]; the
    /// per-harness mapping lives in the worker-side adapter, never here.
    ///
    /// # Errors
    ///
    /// Returns [`WorkerError`] when an installed event seam has been closed (the
    /// runtime drain end was dropped) — a dropped transcript event is surfaced,
    /// never silently swallowed.
    pub fn emit_event(&self, event: ActivityEvent) -> Result<(), WorkerError> {
        if let Some(sender) = &self.event_sender {
            sender
                .send(event)
                .map_err(|source| WorkerError::registration(EventSeamClosed { source }))?;
        }
        Ok(())
    }

    /// Returns true once cooperative cancellation has been signalled.
    #[must_use]
    pub fn is_cancelled(&self) -> bool {
        self.cancellation.cancelled.load(Ordering::Acquire)
    }

    /// Resolves when cooperative cancellation is signalled.
    ///
    /// # Why the waiter exists BEFORE the flag is read
    ///
    /// [`Notify::notify_waiters`] reaches the waiters that EXIST at the moment
    /// it is called and stores nothing for a `Notified` created afterwards —
    /// both halves measured in this module's
    /// `a_notification_is_seen_by_a_waiter_that_existed_before_it_and_lost_on_one_that_did_not`.
    /// A waiter that read the flag first and created its `Notified` second
    /// therefore has a window: the flag reads false, the whole of
    /// [`ActivityCancellationHandle::cancel`] runs inside the window, and the
    /// notification is spent before this waiter exists to receive it. It is
    /// then lost FOREVER, because cancellation is signalled exactly once — and
    /// the awaiting side waits on a cancellation that has already happened,
    /// which is a running command that cannot be killed.
    ///
    /// So the order is inverted. The `Notified` future is created and
    /// `enable()`d — which registers this waiter explicitly, rather than
    /// leaning on the creation-time capture alone — and only then is the flag
    /// read. A `cancel` that lands before the read is seen by the read; one
    /// that lands after it is seen by the waiter already in place. There is no
    /// third position for it to land in.
    ///
    /// The loop is not a spin: `Notified` completes only on a notification, and
    /// re-entering it creates and registers a new waiter before re-reading, so
    /// every iteration keeps the same ordering.
    pub async fn cancelled(&self) {
        loop {
            let notified = self.cancellation.notify.notified();
            let mut notified = std::pin::pin!(notified);
            // Registers this waiter. Everything after this line is covered by
            // a notification, including the flag read on the next line.
            notified.as_mut().enable();
            if self.is_cancelled() {
                return;
            }
            notified.await;
        }
    }

    pub(crate) fn for_workflow(
        workflow_id: WorkflowId,
        run_id: RunId,
        activity_id: ActivityId,
        attempt: u32,
        heartbeat_sender: Option<mpsc::UnboundedSender<HeartbeatRequest>>,
    ) -> (Self, ActivityCancellationHandle) {
        Self::for_workflow_with_events(
            workflow_id,
            run_id,
            activity_id,
            attempt,
            None,
            heartbeat_sender,
            None,
        )
    }

    pub(crate) fn for_task(
        workflow_id: WorkflowId,
        run_id: RunId,
        activity_id: ActivityId,
        attempt: u32,
        idempotency_key: String,
        heartbeat_sender: Option<mpsc::UnboundedSender<HeartbeatRequest>>,
    ) -> (Self, ActivityCancellationHandle) {
        Self::for_workflow_with_events(
            workflow_id,
            run_id,
            activity_id,
            attempt,
            Some(idempotency_key),
            heartbeat_sender,
            None,
        )
    }

    /// Build a context with BOTH the heartbeat seam and the NOI-5b transcript
    /// event seam installed. The runtime uses this when an activity is driven with
    /// observability enabled; the pre-existing constructors default the event seam
    /// to `None` so every current call site is unchanged.
    pub(crate) fn for_workflow_with_events(
        workflow_id: WorkflowId,
        run_id: RunId,
        activity_id: ActivityId,
        attempt: u32,
        idempotency_key: Option<String>,
        heartbeat_sender: Option<mpsc::UnboundedSender<HeartbeatRequest>>,
        event_sender: Option<mpsc::UnboundedSender<ActivityEvent>>,
    ) -> (Self, ActivityCancellationHandle) {
        let cancellation = Arc::new(CancellationState {
            cancelled: AtomicBool::new(false),
            notify: Notify::new(),
        });
        let context = Self {
            workflow_id,
            run_id,
            activity_id,
            attempt,
            idempotency_key,
            cancellation: Arc::clone(&cancellation),
            heartbeat_sender,
            event_sender,
        };
        let handle = ActivityCancellationHandle { cancellation };
        (context, handle)
    }
}

impl ActivityCancellationHandle {
    /// Signals cooperative cancellation to the handler-facing context.
    ///
    /// The FLAG IS SET FIRST and the waiters woken second, and that order is
    /// load-bearing in both directions: a waiter woken by this call finds the
    /// flag already true (so [`ActivityContext::cancelled`] never wakes to a
    /// flag that has not landed yet), and a waiter that registered before the
    /// wake is woken by it (so the flag never lands with nobody told). Waking
    /// first and setting second would let a woken waiter read `false`, loop,
    /// and park forever on a notification that has already been spent.
    ///
    /// `notify_waiters` runs only on the 0→1 transition because cancellation
    /// is signalled once: a second `cancel` has nothing new to say, and the
    /// waiters it would wake are the ones already returned by the flag.
    pub fn cancel(&self) {
        let was_cancelled = self.cancellation.cancelled.swap(true, Ordering::AcqRel);
        if !was_cancelled {
            self.cancellation.notify.notify_waiters();
        }
    }
}

#[derive(Debug, thiserror::Error)]
#[error("activity heartbeat seam is closed: {source}")]
struct HeartbeatSeamClosed {
    source: mpsc::error::SendError<HeartbeatRequest>,
}

#[derive(Debug, thiserror::Error)]
#[error("activity transcript event seam is closed: {source}")]
struct EventSeamClosed {
    source: mpsc::error::SendError<ActivityEvent>,
}

#[cfg(test)]
mod tests {
    use std::future::Future as _;
    use std::sync::Arc;
    use std::sync::atomic::{AtomicBool, Ordering};
    use std::task::{Context as TaskContext, Poll, Wake, Waker};
    use std::time::Duration;

    use aion_core::ActivityId;

    use super::ActivityContext;

    /// What a test returns. Every fallible step is carried rather than
    /// unwrapped, because the workspace denies panicking accessors in test code
    /// as firmly as in library code.
    type TestResult = Result<(), Box<dyn std::error::Error>>;

    /// NOI-5b: a context WITHOUT an event seam is a no-op — `emit_event` returns
    /// `Ok(())` and drops the event, byte-identical to a context that predates the
    /// seam. This is the additive guarantee: an activity that emits nothing (and
    /// every isolated unit test) is unaffected.
    #[tokio::test]
    async fn emit_event_is_a_no_op_without_an_installed_seam() {
        use aion_core::{ActivityEvent, ActivityEventKind, MessageRole, RunId, WorkflowId};
        use chrono::Utc;
        use uuid::Uuid;

        let workflow_id = WorkflowId::new(Uuid::from_u128(1));
        let run_id = RunId::new(Uuid::from_u128(0x11));
        let (context, _cancellation) = ActivityContext::new(
            workflow_id.clone(),
            run_id.clone(),
            ActivityId::from_sequence_position(1),
            0,
        );
        // The dispatch identity is always readable — a handler stamping a
        // transcript event never has to invent the run axis.
        assert_eq!(context.run_id(), &run_id);
        assert_eq!(context.workflow_id(), &workflow_id);
        let event = ActivityEvent {
            workflow_id,
            run_id,
            activity_id: ActivityId::from_sequence_position(1),
            attempt: 0,
            agent_id: Uuid::from_u128(2),
            agent_role: "orchestrator".to_owned(),
            emitted_at: Utc::now(),
            worker_seq: 1,
            store_seq: None,
            ephemeral: false,
            kind: ActivityEventKind::Message {
                role: MessageRole::Assistant,
                text: "hello".to_owned(),
            },
        };
        assert!(context.emit_event(event).is_ok());
    }

    /// With an event seam installed, `emit_event` forwards the neutral event to
    /// the runtime drain end — the additive worker->server ingestion seam.
    #[tokio::test]
    async fn emit_event_forwards_to_installed_seam() -> Result<(), Box<dyn std::error::Error>> {
        use aion_core::{ActivityEvent, ActivityEventKind, MessageRole, RunId, WorkflowId};
        use chrono::Utc;
        use uuid::Uuid;

        let run_id = RunId::new(Uuid::from_u128(0x11));
        let (sender, mut drain) = super::mpsc::unbounded_channel();
        let (context, _cancellation) = ActivityContext::for_workflow_with_events(
            WorkflowId::new(Uuid::from_u128(1)),
            run_id.clone(),
            ActivityId::from_sequence_position(1),
            0,
            None,
            None,
            Some(sender),
        );
        // A dispatched context exposes the run its events must be stamped with.
        assert_eq!(context.run_id(), &run_id);
        let event = ActivityEvent {
            workflow_id: WorkflowId::new(Uuid::from_u128(1)),
            run_id: run_id.clone(),
            activity_id: ActivityId::from_sequence_position(1),
            attempt: 0,
            agent_id: Uuid::from_u128(2),
            agent_role: "orchestrator".to_owned(),
            emitted_at: Utc::now(),
            worker_seq: 7,
            store_seq: None,
            ephemeral: false,
            kind: ActivityEventKind::Message {
                role: MessageRole::Assistant,
                text: "steer".to_owned(),
            },
        };
        context.emit_event(event.clone())?;
        let delivered = drain.recv().await.ok_or("event must be delivered")?;
        assert_eq!(delivered.worker_seq, 7);
        assert_eq!(delivered, event);
        Ok(())
    }

    #[test]
    fn live_task_context_exposes_the_server_idempotency_key() {
        let run_id = aion_core::RunId::new_v4();
        let (context, cancellation) = ActivityContext::for_task(
            aion_core::WorkflowId::new_v4(),
            run_id.clone(),
            ActivityId::from_sequence_position(7),
            3,
            String::from("effect-key"),
            None,
        );

        assert_eq!(context.idempotency_key(), Some("effect-key"));
        assert_eq!(context.attempt(), 3);
        assert_eq!(context.run_id(), &run_id);
        drop(cancellation);
    }

    /// Reads the cancellation flag AT THE MOMENT it is woken, which is the only
    /// place from which the ordering inside `cancel()` is observable: `wake` is
    /// called from inside `notify_waiters`, so what the flag says here is what
    /// it said before the wake was sent.
    struct CancelWatcher {
        context: ActivityContext,
        woken: AtomicBool,
        flag_at_wake: AtomicBool,
    }

    impl Wake for CancelWatcher {
        fn wake(self: Arc<Self>) {
            self.wake_by_ref();
        }

        fn wake_by_ref(self: &Arc<Self>) {
            self.flag_at_wake
                .store(self.context.is_cancelled(), Ordering::Release);
            self.woken.store(true, Ordering::Release);
        }
    }

    /// THE ORDERING GUARANTEE, both halves, measured rather than argued.
    ///
    /// A waiter that has polled once is REGISTERED (it is woken by a later
    /// `cancel`), and `cancel` sets the flag BEFORE it wakes anyone (the flag
    /// already reads true from inside the wake). Those two facts are what make
    /// a cancellation impossible to lose: whichever side moves first, the other
    /// sees it.
    #[test]
    fn cancel_sets_the_flag_before_it_wakes_a_registered_waiter() {
        let (context, handle) = ActivityContext::new(
            aion_core::WorkflowId::new_v4(),
            aion_core::RunId::new_v4(),
            ActivityId::from_sequence_position(1),
            1,
        );
        let watcher = Arc::new(CancelWatcher {
            context: context.clone(),
            woken: AtomicBool::new(false),
            flag_at_wake: AtomicBool::new(false),
        });
        let waker = Waker::from(Arc::clone(&watcher));
        let mut task = TaskContext::from_waker(&waker);
        let mut cancelled = std::pin::pin!(context.cancelled());

        assert_eq!(
            cancelled.as_mut().poll(&mut task),
            Poll::Pending,
            "nothing has cancelled yet"
        );
        assert!(
            !watcher.woken.load(Ordering::Acquire),
            "no wake before a cancel"
        );

        handle.cancel();

        assert!(
            watcher.woken.load(Ordering::Acquire),
            "the waiter registered its interest on its first poll, so `cancel` found it"
        );
        assert!(
            watcher.flag_at_wake.load(Ordering::Acquire),
            "`cancel` must set the flag BEFORE it notifies: a waiter woken to a flag that has \
             not landed yet loops and parks forever"
        );
        assert_eq!(
            cancelled.as_mut().poll(&mut task),
            Poll::Ready(()),
            "the woken waiter resolves"
        );
    }

    /// THE WINDOW THE FIX CLOSES, measured directly on the primitive rather
    /// than asserted about it.
    ///
    /// Three arms, and the middle one is the whole reason the order in
    /// `cancelled()` is what it is:
    ///
    /// * a `Notified` created AFTER a `notify_waiters` never sees it — the
    ///   notification is gone, and cancellation is signalled exactly once, so a
    ///   lost one is lost for good;
    /// * a `Notified` created BEFORE it does see it, even though it had not yet
    ///   been polled;
    /// * and one that was additionally `enable`d sees it too.
    ///
    /// So the guarantee `cancelled()` needs is that the `Notified` EXISTS
    /// before the flag is read. It is enabled as well, which registers it
    /// explicitly at a point this code chooses, so the guarantee does not rest
    /// on the creation-time capture alone.
    #[tokio::test]
    async fn a_notification_is_seen_by_a_waiter_that_existed_before_it_and_lost_on_one_that_did_not()
    -> TestResult {
        let notify = tokio::sync::Notify::new();
        notify.notify_waiters();
        let mut late = std::pin::pin!(notify.notified());
        let missed = tokio::time::timeout(Duration::from_millis(200), late.as_mut()).await;
        assert!(
            missed.is_err(),
            "a notification sent before the waiter existed is LOST — this is the window"
        );

        let notify = tokio::sync::Notify::new();
        let mut created = std::pin::pin!(notify.notified());
        notify.notify_waiters();
        let landed = tokio::time::timeout(Duration::from_millis(200), created.as_mut()).await;
        assert!(
            landed.is_ok(),
            "a waiter that existed when the notification was sent sees it"
        );

        let notify = tokio::sync::Notify::new();
        let mut registered = std::pin::pin!(notify.notified());
        assert!(
            !registered.as_mut().enable(),
            "nothing has been notified yet, so enabling only registers"
        );
        notify.notify_waiters();
        let landed = tokio::time::timeout(Duration::from_millis(200), registered.as_mut()).await;
        assert!(
            landed.is_ok(),
            "a waiter that registered before the notification was sent sees it"
        );
        Ok(())
    }

    /// THE RACE, RUN: a cancel landing while the waiter is on its way into the
    /// wait must never be lost. Bounded — a fixed number of rounds, each with a
    /// patience wide enough that a healthy round never approaches it — and it
    /// FAILS rather than hangs, naming the round that lost its cancellation.
    ///
    /// The defect this guards was reproducible at roughly one round in tens of
    /// thousands, so a green run of this test is a guard and not a proof; the
    /// proof is the ordering test above and the `enable`-before-read structure
    /// it pins. What this catches is a future change that quietly reintroduces
    /// the read-then-register order.
    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn a_cancel_racing_the_waiter_into_the_wait_is_never_lost() -> TestResult {
        const ROUNDS: usize = 20_000;
        const PATIENCE: Duration = Duration::from_secs(30);

        for round in 0..ROUNDS {
            let (context, handle) = ActivityContext::new(
                aion_core::WorkflowId::new_v4(),
                aion_core::RunId::new_v4(),
                ActivityId::from_sequence_position(1),
                1,
            );
            let waiter = tokio::spawn(async move { context.cancelled().await });
            // The canceller runs on a BLOCKING pool thread, which is a real
            // thread of its own rather than a task queued behind the waiter:
            // the two sides have to be genuinely parallel for the window to be
            // reachable at all, and a canceller that can only run after the
            // waiter has parked would make every round pass vacuously.
            let canceller = tokio::task::spawn_blocking(move || handle.cancel());
            match tokio::time::timeout(PATIENCE, waiter).await {
                Ok(joined) => joined?,
                Err(elapsed) => {
                    return Err(format!(
                        "round {round}: the waiter never woke ({elapsed}) — a cancellation was \
                         lost, which is the read-then-register window"
                    )
                    .into());
                }
            }
            canceller.await?;
        }
        Ok(())
    }

    #[tokio::test]
    async fn context_exposes_identity_attempt_and_cancellation_signal() {
        let activity_id = ActivityId::from_sequence_position(42);
        let (context, cancellation) = ActivityContext::new(
            aion_core::WorkflowId::new_v4(),
            aion_core::RunId::new_v4(),
            activity_id.clone(),
            3,
        );

        assert_eq!(context.activity_id(), &activity_id);
        assert_eq!(context.attempt(), 3);
        assert!(!context.is_cancelled());

        cancellation.cancel();

        assert!(context.is_cancelled());
        let cancelled = tokio::time::timeout(Duration::from_millis(50), context.cancelled()).await;
        assert!(cancelled.is_ok());
    }
}