aion_server/worker/

liminal_task_delivery.rs

//! The liminal arm of the delivery seam: push the dispatch out on the worker's
//! existing connection and block for its correlated reply.
//!
//! # Why this arm is the one that needed a seam
//!
//! Before #52 there was no way to reach it from the worker dispatcher at all: a
//! worker with no gRPC stream sender was deregistered on the premise that its
//! presence could only be a leak. This module is the same delivery sequence the
//! outbox arm has always run, lifted so both callers reach one implementation.
//!
//! # What it does NOT own any more
//!
//! Two responsibilities deliberately stayed with the caller, and moving them
//! here would break the invariants they exist to keep:
//!
//! - **The completion token.** One per pass, minted before the candidate walk
//!   and revoked once when the pass places nothing. Minting per delivery would
//!   put a second authorization beside the one a worker may still be holding for
//!   the same attempt.
//! - **The abandonment condition.** See
//!   [`DeliveryIntent`](super::delivery_intent::DeliveryIntent): one of its terms
//!   is unanswerable from inside a transport.

use std::sync::Arc;

use async_trait::async_trait;

use aion_core::{ActivityId, RunId, WorkflowId};
use aion_proto::ProtoActivityTask;

use super::delivery_intent::SharedDeliveryIntent;
use super::intervention::{AttemptKey, AttemptOwnerIndex};
use super::liminal_transport::{AttemptOwnerGuard, DispatchRequest, LiminalCompletionSource};
use super::registry::{WorkerDelivery, WorkerHandle};
use super::task_delivery::{LivenessTracking, TaskDelivery, WorkerTaskDelivery};

/// Delivers by pushing the dispatch out on the worker's liminal connection and
/// blocking for the correlated reply.
///
/// The reply **is** the activity's completion, so this arm re-enters it through
/// the same completion path the gRPC transport's out-of-band completion uses.
pub struct LiminalTaskDelivery {
    completion: Arc<LiminalCompletionSource>,
    /// NOI-6 `attempt -> owning-worker` back-index. `None` (every non-agent
    /// deployment) skips the binding, exactly as the outbox arm does —
    /// intervention is then simply never offered.
    attempt_owners: Option<AttemptOwnerIndex>,
}

impl std::fmt::Debug for LiminalTaskDelivery {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter
            .debug_struct("LiminalTaskDelivery")
            .field("attempt_owners", &self.attempt_owners.is_some())
            .finish_non_exhaustive()
    }
}

impl LiminalTaskDelivery {
    /// Build the liminal delivery arm over the completion sink both transports
    /// share.
    #[must_use]
    pub fn new(completion: Arc<LiminalCompletionSource>) -> Self {
        Self {
            completion,
            attempt_owners: None,
        }
    }

    /// Install the NOI-6 attempt-owner back-index, so a dispatched attempt binds
    /// its owning worker before the push and the intervention router can resolve
    /// the current owner of a live attempt.
    #[must_use]
    pub fn with_attempt_owners(mut self, attempt_owners: AttemptOwnerIndex) -> Self {
        self.attempt_owners = Some(attempt_owners);
        self
    }
}

/// The identity a delivery needs, resolved from the wire task **before** any
/// owner binding is taken.
///
/// 🔴 This type exists to make the ordering structural rather than positional.
/// The owner binding is keyed on the run, so a binding taken before the run is
/// resolved would let an intervention aimed at one continue-as-new generation
/// resolve another generation's worker.
///
/// # Exactly how strong this is, measured rather than asserted
///
/// The binding is taken by [`Resolved::bind_owner`], a method **on the resolved
/// identity**, so it cannot be called before that identity exists: moving the
/// call above the resolution does not compile. That is the mutation this guards
/// against, and it has been run.
///
/// It is **not** proof against a deliberate bypass. [`AttemptKey::new`] is a
/// public constructor over four plain components, so an edit that calls
/// [`AttemptOwnerGuard::bind`] directly with a placeholder run compiles and
/// binds. An earlier version of this comment claimed the key "can only be built
/// from a `Resolved`", which was false — the mutation that proved it survived
/// both this structure and the test below.
///
/// The bypass is also, for the same reason, invisible at runtime: the guard
/// releases on drop, so a binding taken before a refusal is gone again before
/// any caller can look. What defends the property is that the natural edit does
/// not compile and the deliberate one is a different API call, which review can
/// see.
struct Resolved {
    workflow_id: WorkflowId,
    run_id: RunId,
    activity_id: ActivityId,
    attempt: u32,
}

impl Resolved {
    /// Resolve the task's identity, refusing a task with no run.
    ///
    /// The refusal is the pre-existing one — an activity without a run cannot be
    /// dispatched at all, because an unfenced external effect has no generation
    /// to belong to.
    fn from_task(task: &ProtoActivityTask) -> Result<Self, &'static str> {
        let workflow_id = task
            .workflow_id
            .clone()
            .ok_or("task carries no workflow id")?
            .try_into()
            .map_err(|_| "task carries a malformed workflow id")?;
        let run_id: RunId = task
            .run_id
            .clone()
            .ok_or("activity run id is missing; refusing unfenced external effect")?
            .try_into()
            .map_err(|_| "task carries a malformed run id")?;
        // Infallible: `ProtoActivityId` is the sequence position, and
        // `ActivityId` is a newtype over it. Only its ABSENCE can be refused.
        let activity_id: ActivityId = task
            .activity_id
            .ok_or("task carries no activity id")?
            .into();
        Ok(Self {
            workflow_id,
            run_id,
            activity_id,
            attempt: task.attempt,
        })
    }

    /// The intervention key for this attempt, carrying the **resolved** run.
    fn attempt_key(&self) -> AttemptKey {
        AttemptKey::new(
            self.workflow_id.clone(),
            self.run_id.clone(),
            self.activity_id.clone(),
            self.attempt,
        )
    }

    /// Bind this attempt's owner, returning the guard that releases it.
    ///
    /// 🔴 A method on the **resolved** identity on purpose. The binding is keyed
    /// on the run, so taking it before the run is resolved is the defect; making
    /// it a method on `Resolved` means the call cannot be moved above the
    /// resolution — there is no receiver for it there, and the lift does not
    /// compile. See this type's docs for what that does and does not prove.
    fn bind_owner(
        &self,
        owners: Option<&AttemptOwnerIndex>,
        worker: super::registry::WorkerId,
    ) -> Option<AttemptOwnerGuard> {
        owners.map(|owners| AttemptOwnerGuard::bind(owners.clone(), self.attempt_key(), worker))
    }

    /// The outbox ordinal this activity id was derived from.
    ///
    /// [`ActivityId`] is a newtype over the scheduling sequence position, so
    /// `from_sequence_position` is exactly invertible and the ordinal the wire
    /// needs round-trips without the row being present.
    fn ordinal(&self) -> u64 {
        self.activity_id.sequence_position()
    }
}

/// Build the liminal wire request from the already-built task.
///
/// Every field is carried from the task except `heartbeat_window_ms`, which the
/// task does not have — see [`LivenessTracking`] for why that assignment is a named
/// value rather than a bare zero.
fn request_for_task(task: &ProtoActivityTask, resolved: &Resolved) -> DispatchRequest {
    DispatchRequest {
        activity_type: task.activity_type.clone(),
        workflow_id: resolved.workflow_id.clone(),
        ordinal: resolved.ordinal(),
        run_id: Some(resolved.run_id.clone()),
        completion_token: task.completion_token.clone(),
        idempotency_key: task.idempotency_key.clone(),
        input: task
            .input
            .as_ref()
            .map(|payload| payload.bytes.clone())
            .unwrap_or_default(),
        attempt: resolved.attempt,
        labels: task
            .labels
            .iter()
            .map(|(k, v)| (k.clone(), v.clone()))
            .collect(),
        heartbeat_window_ms: LivenessTracking::NotTrackedPerTask.heartbeat_window_ms(),
    }
}

#[async_trait]
impl WorkerTaskDelivery for LiminalTaskDelivery {
    async fn deliver(
        &self,
        worker: &WorkerHandle,
        task: &ProtoActivityTask,
        intent: &SharedDeliveryIntent,
    ) -> TaskDelivery {
        // Identity is resolved FIRST — before the transport is even consulted —
        // and the binding below is a method on what this produces, so it cannot
        // be moved above this line.
        //
        // Ahead of the transport check deliberately. A task with no run is
        // malformed **regardless of which worker it was aimed at**, so refusing
        // it with "not delivered over liminal" would name the wrong defect and
        // send an operator after the wrong remedy.
        let resolved = match Resolved::from_task(task) {
            Ok(resolved) => resolved,
            Err(reason) => return TaskDelivery::failed(reason),
        };

        let delivery = match worker.delivery() {
            WorkerDelivery::Liminal(delivery) => delivery.clone(),
            WorkerDelivery::Grpc(_) => {
                // A caller chose the wrong transport for the worker it selected.
                // The worker is alive and correctly registered, so its
                // registration stands — this is #52's defect refused at its
                // mirror site.
                return TaskDelivery::failed(
                    "selected worker is not delivered over liminal; the liminal transport cannot \
                     reach it",
                );
            }
        };

        // NOI-6: bind this attempt's owner BEFORE the push, so an intervention
        // that races the dispatch resolves the worker. The guard releases on
        // every exit path (reply, error, panic) so the index never keeps a
        // finished attempt.
        let _owner_guard = resolved.bind_owner(self.attempt_owners.as_ref(), worker.id());

        let request = request_for_task(task, &resolved);
        // The push is a blocking, thread-based liminal call; run it off the
        // async runtime so a long-running activity cannot starve a runtime
        // worker. The caller's intent is re-asked at every poll boundary of the
        // wait, which is why it arrives behind an `Arc`.
        let waiting_intent = Arc::clone(intent);
        let dispatched = tokio::task::spawn_blocking(move || {
            delivery.dispatch_held(&request, || waiting_intent.still_wanted())
        })
        .await;

        let response = match dispatched {
            Ok(Ok(Some(response))) => response,
            Ok(Ok(None)) => {
                // The caller withdrew while the delivery waited. The worker is
                // alive and correctly registered; only this pass gave up, and a
                // late reply is discarded by the fences.
                return TaskDelivery::failed("delivery wait abandoned before worker reply");
            }
            Ok(Err(error)) => {
                return TaskDelivery::failed(format!("liminal dispatch failed: {error}"));
            }
            Err(error) => {
                return TaskDelivery::failed(format!("dispatch task join failed: {error}"));
            }
        };

        // Re-enter the worker's result through the SAME completion path the gRPC
        // transport uses; terminal dedup applies unchanged.
        //
        // 🔴 A failure HERE is the one outcome whose name reads narrower than its
        // meaning: the worker took the task, executed it, and replied — and the
        // reply did not record. `Delivered` would be wrong (the caller would
        // settle the work without its result) and `WorkerUnreachable` would be
        // wrong twice over (the worker is demonstrably alive, and deregistering
        // it would destroy a healthy registration). `DeliveryFailed` carries the
        // right obligations — registration stands, the caller withdraws its
        // token — which is what the type encodes and what the outbox arm has
        // always done here.
        if let Err(error) = self.completion.deliver(&response) {
            return TaskDelivery::failed(format!(
                "worker replied but the completion could not be recorded: {error}"
            ));
        }
        TaskDelivery::Delivered
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use aion_core::{ActivityId, InterventionCapabilities, RunId, WorkflowId};
    use aion_proto::{ProtoActivityId, ProtoActivityTask, ProtoWorkflowId};
    use uuid::Uuid;

    use crate::error::ServerError;
    use crate::worker::bridge::OutboxDeliveryCallback;
    use crate::worker::delivery_intent::{AlwaysWanted, SharedDeliveryIntent};
    use crate::worker::intervention::AttemptOwnerIndex;
    use crate::worker::liminal_transport::LiminalCompletionSource;
    use crate::worker::registry::{ConnectedWorkerRegistry, WorkerDelivery};
    use crate::worker::task_delivery::{TaskDelivery, WorkerTaskDelivery};

    use super::{LiminalTaskDelivery, Resolved};

    /// The completion sink is never reached by these tests — every one of them
    /// refuses before a worker is pushed to — so both methods report "no live
    /// run" if they are ever called, rather than pretending to succeed.
    struct NoopCallback;

    impl OutboxDeliveryCallback for NoopCallback {
        fn deliver_completion(
            &self,
            _workflow_id: &WorkflowId,
            _activity_id: &ActivityId,
            _run_id: Option<&RunId>,
            _result: String,
        ) -> Result<bool, ServerError> {
            Ok(false)
        }
        fn deliver_failure(
            &self,
            _workflow_id: &WorkflowId,
            _activity_id: &ActivityId,
            _run_id: Option<&RunId>,
            _reason: String,
        ) -> Result<bool, ServerError> {
            Ok(false)
        }
    }

    const WORKFLOW: u128 = 0x51;

    /// A task carrying everything a delivery needs EXCEPT a run id.
    fn task_without_a_run() -> ProtoActivityTask {
        ProtoActivityTask {
            workflow_id: Some(ProtoWorkflowId::from(WorkflowId::new(Uuid::from_u128(
                WORKFLOW,
            )))),
            activity_id: Some(ProtoActivityId::from(ActivityId::from_sequence_position(3))),
            activity_type: String::from("agent"),
            input: None,
            attempt: 1,
            labels: std::collections::HashMap::new(),
            // The whole point of this fixture.
            run_id: None,
            completion_token: String::from("token"),
            idempotency_key: String::from("key"),
        }
    }

    fn liminal_delivery(owners: &AttemptOwnerIndex) -> LiminalTaskDelivery {
        LiminalTaskDelivery::new(Arc::new(LiminalCompletionSource::new(Arc::new(
            NoopCallback,
        ))))
        .with_attempt_owners(owners.clone())
    }

    /// 🔴 NAMED PROPERTY: the run is resolved BEFORE any attempt owner is bound.
    ///
    /// # What breaks without it
    ///
    /// The owner binding is keyed on the run. A binding taken before the run is
    /// resolved would let an intervention aimed at one continue-as-new
    /// generation resolve a DIFFERENT generation's worker — an operator
    /// stopping run A reaching the worker executing run B.
    ///
    /// # 🔴 What this test does and does NOT witness — measured by mutation
    ///
    /// It witnesses the **refusal**: a run-less task is rejected, is diagnosed
    /// by its missing run rather than by the transport, does not deregister the
    /// worker, and leaves no owner binding behind.
    ///
    /// It does **not** witness the ordering, and an earlier version of this
    /// comment claimed it did. The mutation — lifting the bind above the
    /// resolution with a placeholder run — was run, and **this test stayed
    /// green**, because [`AttemptOwnerGuard`] releases on drop: a binding taken
    /// before a refusal is already gone by the time any caller can look. There
    /// is no observation window on this path, and there cannot be one.
    ///
    /// The ordering is defended structurally instead: the bind is a method on
    /// [`Resolved`], so the natural lift does not compile. That claim was also
    /// checked by mutation rather than asserted — see [`Resolved`] for exactly
    /// how far it goes, including the bypass it does not stop.
    #[tokio::test]
    async fn run_resolution_precedes_attempt_owner_binding()
    -> Result<(), Box<dyn std::error::Error>> {
        let owners = AttemptOwnerIndex::new();
        let delivery = liminal_delivery(&owners);

        // A real registration yields a real WorkerId; no id here is fabricated.
        let registry = ConnectedWorkerRegistry::default();
        let (sender, _receiver) = tokio::sync::mpsc::channel(1);
        let types = [String::from("agent")];
        let registration = registry.register_delivery_with_capabilities(
            [String::from("default")],
            String::from("default"),
            None,
            types.iter(),
            WorkerDelivery::Grpc(sender),
            InterventionCapabilities::none(),
        )?;
        let worker_id = registration
            .worker_id()
            .ok_or("a registration must assign a worker id")?;
        let worker = registry
            .worker_by_id(worker_id)?
            .ok_or("the worker just registered must be readable")?;

        let intent: SharedDeliveryIntent = Arc::new(AlwaysWanted);
        let outcome = delivery
            .deliver(&worker, &task_without_a_run(), &intent)
            .await;

        match outcome {
            TaskDelivery::Delivered => {
                return Err("a task with no run must never be delivered".into());
            }
            TaskDelivery::Undeliverable(undeliverable) => {
                // The refusal names the RUN, not the transport: identity is
                // resolved before the transport is consulted, so a run-less task
                // is diagnosed as run-less whichever worker it was aimed at.
                assert!(
                    undeliverable.reason().contains("run id is missing"),
                    "a run-less task must be refused BY ITS MISSING RUN, not by the transport; \
                     got: {}",
                    undeliverable.reason()
                );
                assert!(
                    !undeliverable.deregisters_worker(),
                    "a malformed task is not evidence that the worker is gone"
                );
            }
        }

        // 🔴 THE PROPERTY, as far as this test can witness it: nothing is bound
        // when the run refuses. Note what this does NOT say — an earlier
        // version of this comment claimed nothing COULD be bound because an
        // `AttemptKey` is only constructible from a resolved identity, and that
        // was false. `AttemptKey::new` is public over four plain components.
        // See this test's docs and [`Resolved`] for what actually holds.
        let bound = owners.attempts_for_workflow(&WorkflowId::new(Uuid::from_u128(WORKFLOW)));
        assert!(
            bound.is_empty(),
            "the attempt-owner index must be untouched when the run refused; a binding here \
             means the bind was lifted above the run resolution, and an intervention could \
             resolve the wrong generation's worker. Found: {bound:?}"
        );
        Ok(())
    }

    /// The property's SECOND witness, and the one that outlives a refactor of
    /// the call site: identity resolution refuses a run-less task, so no
    /// `Resolved` is produced — and [`Resolved::bind_owner`], the only binding
    /// path the delivery takes, is a method on it.
    ///
    /// 🔴 That is a claim about the CALL SITE, not about the key type. An
    /// earlier version said an `AttemptKey` "cannot be built without" a
    /// `Resolved`; it can, and the mutation that proved it survived. What this
    /// witnesses is that resolution refuses first.
    ///
    /// Kept beside the behavioural test deliberately. The refusal and the
    /// structure are different claims, and a change that removes one should
    /// still meet the other.
    #[test]
    fn identity_resolution_refuses_a_task_with_no_run() -> Result<(), Box<dyn std::error::Error>> {
        let Err(refusal) = Resolved::from_task(&task_without_a_run()) else {
            return Err("a task with no run must not resolve".into());
        };
        assert!(
            refusal.contains("run id is missing"),
            "the refusal must name the run so an operator is sent to the right remedy: {refusal}"
        );
        Ok(())
    }
}