Skip to main content

aion_server/worker/
bridge.rs

1//! NIF bridge dispatcher that routes `run_activity` calls to connected workers.
2//!
3//! `WorkerActivityDispatcher` implements `aion::ActivityDispatcher` so the
4//! engine's activity NIFs can synchronously dispatch to a remote worker and
5//! block until the result comes back.
6//!
7//! # Threading contract
8//!
9//! The engine invokes [`aion::ActivityDispatcher::dispatch`] from two kinds of
10//! threads: beamr scheduler threads (concurrency combinators) and spawned
11//! tokio tasks (the two-phase `dispatch_activity` completion task). The task
12//! send uses `try_send()` (non-blocking channel push) and the response wait
13//! blocks on `std::sync::mpsc::Receiver::recv`.
14//!
15//! Blocking is harmless on a beamr thread, but on a tokio runtime worker it
16//! must be wrapped in `tokio::task::block_in_place`: the `try_send` wakes the
17//! per-worker gRPC stream forwarder task, and tokio schedules a task woken
18//! from task context into the *current* worker's LIFO slot, which no other
19//! runtime worker can steal. Without the `block_in_place` core handoff the
20//! forwarder sits trapped in that slot while this thread blocks, so the queued
21//! `ActivityTask` is never flushed to the worker even though the worker is
22//! healthy. `block_in_place` moves the worker's scheduler core (LIFO slot
23//! included) to another thread before the wait begins, so dispatch-to-delivery
24//! stays in the millisecond range and the runtime keeps full parallelism.
25//!
26//! # Wait termination
27//!
28//! The engine imposes no activity timeout of its own: agent-style activities
29//! legitimately run for over an hour, so the completion wait is unbounded.
30//! The blocking `recv` terminates on exactly one of:
31//!
32//! - **Completion** — the worker reports a result and the stream handler
33//!   delivers it through [`ActivityCompletionSink::complete_activity`].
34//! - **Worker loss** — the worker's gRPC stream ends (process death,
35//!   disconnect, expired token); the stream teardown sweeps the worker's
36//!   in-flight tasks through the same sink as every other TRANSPORT loss
37//!   ([`HeartbeatTracker::fail_disconnected_worker`]). A liminal-delivered
38//!   worker's loss is observed by its reply router thread instead (the
39//!   correlated-reply awaiter wakes the moment the connection closes) and
40//!   resolves the dispatch with the same transport-loss class. Both are
41//!   classified by [`transport_loss`](crate::worker::transport_loss): `lost:`
42//!   while the transport still has budget to deliver the activity (the engine
43//!   re-dispatches the SAME attempt, recording nothing), `transport-exhausted:`
44//!   once that budget is spent. Neither ever wears the action's retry
45//!   vocabulary — the activity never ran.
46//! - **Graceful-drain park (#207)** — during a drain, a worker stream ending
47//!   (or the drain-timeout backstop) PARKS the worker's in-flight tasks
48//!   through [`ActivityCompletionSink::park_activity`]: the waiter resolves
49//!   with the ephemeral parked sentinel
50//!   ([`aion::PARKED_ACTIVITY_REASON`]), nothing is recorded or delivered,
51//!   and restart recovery re-dispatches the dangling ordinal — kill -9
52//!   convergence.
53//! - **Drain timeout at shutdown** — the shutdown coordinator parks all
54//!   remaining in-flight tasks through the sink
55//!   (`HeartbeatTracker::park_all_in_flight_workers`).
56//! - **Channel teardown** — every sender for the pending entry is dropped
57//!   (a cleanup path removed the entry without completing it); surfaced as
58//!   a channel-closed dispatch error, never a hang.
59//!
60//! An activity's duration is bounded only by the workflow's own
61//! `timeout_seconds` and by worker liveness — never by an engine constant.
62
63use std::collections::BTreeMap;
64use std::sync::{Arc, OnceLock};
65use std::time::{Duration, Instant};
66
67use aion::{ActivityDispatch, ActivityDispatcher};
68use aion_core::{ActivityErrorKind, ActivityId, ContentType, Payload, RunId, WorkflowId};
69use aion_proto::{ProtoActivityId, ProtoActivityTask, ProtoPayload, ProtoWorkflowId};
70use dashmap::DashMap;
71use dashmap::mapref::entry::Entry;
72
73use super::dispatch::{ActivityCompletion, ActivityCompletionOutcome, ActivityCompletionSink};
74use super::envelope::{CompletionFences, CompletionToken, idempotency_key};
75use super::heartbeat::{HeartbeatTracker, InFlightActivity};
76use super::queue_service::{
77    DeliveryRefusal, ExpiredClock, PARK_POLL_INTERVAL, PoolCensus, QueueDeclarationSource,
78    QueueServiceConfig, QueueServiceReason, QueueServiceState, SelectionRefusal, ServiceAddress,
79    ServiceWait, WorkerUnavailable, deliver_within_schedule_to_start, select_worker_or_refuse,
80};
81use super::registry::{
82    ConnectedWorkerRegistry, WorkerArrival, WorkerDelivery, WorkerHandle, WorkerId, WorkerMessage,
83};
84use crate::error::ServerError;
85use crate::shutdown::DrainState;
86use tracing::info_span;
87
88type SyncSender = std::sync::mpsc::SyncSender<Result<String, String>>;
89type SyncReceiver = std::sync::mpsc::Receiver<Result<String, String>>;
90
91/// Execution-scoped key for an in-flight activity dispatch.
92///
93/// The engine seam ([`ActivityDispatch`]) carries the *real* workflow id and
94/// the *real* per-workflow activity ordinal recorded in history, so this pair
95/// uniquely and stably identifies one execution. Keying by bare [`ActivityId`]
96/// would be unsafe across server restarts — a stale result re-reported from a
97/// worker's previous session could complete a *different* post-restart
98/// dispatch reusing the same ordinal — but pairing it with the real workflow
99/// id closes that race: two different workflow executions never share a
100/// workflow id, so a stale `(workflow_id, activity_id)` from a previous server
101/// life can only ever match the exact execution it belongs to.
102///
103/// The wire (`ActivityResult`) carries both ids, plus an attempt discriminator
104/// (`ActivityTask.attempt`). The pending key stays attempt-free for now: a
105/// retry re-dispatches under the same `(workflow_id, activity_id)` and the
106/// outstanding entry is the one awaiting completion. Redelivery bookkeeping
107/// can widen this key with the wire attempt later — no protocol change needed.
108type PendingActivityKey = (WorkflowId, ActivityId);
109
110/// Routes an unmatched durable-outbox completion into the live workflow.
111///
112/// When the outbox is ON a worker completion can arrive at the sink with no
113/// pending oneshot (the dispatch was non-blocking fan-out, or the original
114/// waiter was lost). Rather than dropping it, [`PendingActivities::complete`]
115/// hands it to this callback, which resolves the workflow to its live engine
116/// process and delivers the terminal into its mailbox. The callback is only
117/// installed when the outbox is enabled, so flag-off the unmatched branch
118/// stays a silent drop.
119pub trait OutboxDeliveryCallback: Send + Sync {
120    /// Deliver a successful completion to the live workflow.
121    ///
122    /// Returns `Ok(true)` when delivered to a live workflow and `Ok(false)`
123    /// when no run is currently live (the expected stale-completion case that
124    /// recovery re-arms).
125    ///
126    /// # Errors
127    ///
128    /// Returns [`ServerError`] when the engine rejects the delivery.
129    fn deliver_completion(
130        &self,
131        workflow_id: &WorkflowId,
132        activity_id: &ActivityId,
133        run_id: Option<&RunId>,
134        result: String,
135    ) -> Result<bool, ServerError>;
136
137    /// Deliver a failure to the live workflow. Same `bool`/error contract as
138    /// [`Self::deliver_completion`].
139    ///
140    /// # Errors
141    ///
142    /// Returns [`ServerError`] when the engine rejects the delivery.
143    fn deliver_failure(
144        &self,
145        workflow_id: &WorkflowId,
146        activity_id: &ActivityId,
147        run_id: Option<&RunId>,
148        reason: String,
149    ) -> Result<bool, ServerError>;
150}
151
152/// Tracks in-flight activity dispatches waiting for worker results.
153///
154/// When the server's worker stream handler receives an `ActivityResult`, it
155/// calls [`complete_activity`](ActivityCompletionSink::complete_activity) to
156/// deliver the result to the blocked NIF thread. Entries are keyed by
157/// [`PendingActivityKey`] so a stale result from a previous server life can
158/// never be matched to a different execution (#59).
159///
160/// Clones share both the pending map and the outbox-delivery callback through
161/// `Arc`, so [`set_outbox_delivery`](Self::set_outbox_delivery) called once on
162/// any clone after construction is visible to the clone the dispatcher holds.
163#[derive(Clone)]
164pub struct PendingActivities {
165    pending: Arc<DashMap<PendingActivityKey, SyncSender>>,
166    completion_fences: CompletionFences,
167    outbox_delivery: Arc<OnceLock<Arc<dyn OutboxDeliveryCallback>>>,
168    /// The transport's OWN re-dispatch budget for activities whose worker died
169    /// before reporting (see [`transport_loss`](crate::worker::transport_loss)).
170    /// Shared across clones, so every loss for one execution site lands in one
171    /// budget. There is no default: the window is required at construction, so
172    /// a zero budget cannot be inherited by omission.
173    transport_losses: super::transport_loss::TransportLossLedger,
174}
175
176impl std::fmt::Debug for PendingActivities {
177    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
178        formatter
179            .debug_struct("PendingActivities")
180            .field("pending", &self.pending.len())
181            .field("completion_fences", &self.completion_fences)
182            .field(
183                "outbox_delivery_installed",
184                &self.outbox_delivery.get().is_some(),
185            )
186            .field("transport_losses", &self.transport_losses)
187            .finish()
188    }
189}
190
191impl PendingActivities {
192    fn insert(
193        &self,
194        workflow_id: WorkflowId,
195        run_id: &RunId,
196        activity_id: ActivityId,
197        attempt: u32,
198    ) -> Result<(CompletionToken, SyncReceiver), ServerError> {
199        let completion_token =
200            self.completion_fences
201                .issue(&workflow_id, run_id, &activity_id, attempt)?;
202        let (tx, rx) = std::sync::mpsc::sync_channel(1);
203        let key = (workflow_id, activity_id);
204        match self.pending.entry(key) {
205            Entry::Vacant(entry) => {
206                entry.insert(tx);
207                Ok((completion_token, rx))
208            }
209            Entry::Occupied(entry) => {
210                let (workflow_id, activity_id) = entry.key();
211                Err(ServerError::PendingActivityCollision {
212                    workflow_id: workflow_id.clone(),
213                    activity_id: activity_id.clone(),
214                })
215            }
216        }
217    }
218
219    /// Share the generation registry with non-blocking outbox dispatch.
220    #[must_use]
221    pub fn completion_fences(&self) -> CompletionFences {
222        self.completion_fences.clone()
223    }
224
225    /// Test seam: register a pending waiter exactly as a live dispatch does,
226    /// so crate-internal tests outside this module (the stream-teardown
227    /// park-vs-fail pins) can observe how a sweep resolves it.
228    #[cfg(test)]
229    pub(crate) fn insert_for_test(
230        &self,
231        workflow_id: WorkflowId,
232        run_id: &RunId,
233        activity_id: ActivityId,
234        attempt: u32,
235    ) -> Result<(CompletionToken, SyncReceiver), ServerError> {
236        self.insert(workflow_id, run_id, activity_id, attempt)
237    }
238
239    /// Install the unmatched-completion delivery callback (idempotent).
240    ///
241    /// Set once, after construction, when the durable outbox is enabled. A
242    /// second set is ignored and logged: the callback is process-wide and must
243    /// not silently change identity.
244    pub fn set_outbox_delivery(&self, callback: Arc<dyn OutboxDeliveryCallback>) {
245        if self.outbox_delivery.set(callback).is_err() {
246            tracing::warn!("outbox delivery callback already installed; ignoring duplicate set");
247        }
248    }
249
250    /// Complete a pending dispatch, or route an unmatched completion to the
251    /// outbox delivery callback when one is installed.
252    ///
253    /// A matched entry delivers to its waiting oneshot exactly as before. An
254    /// unmatched completion is dropped silently when no callback is installed
255    /// (outbox OFF — byte-identical to the prior behaviour); with a callback
256    /// installed (outbox ON) it is routed into the live workflow's mailbox.
257    fn complete(
258        &self,
259        workflow_id: &WorkflowId,
260        activity_id: &ActivityId,
261        run_id: Option<&RunId>,
262        result: Result<String, String>,
263    ) -> bool {
264        // Take and drop the DashMap guard before any callback runs: the engine
265        // delivery the callback invokes must never execute under a shard lock.
266        let matched = self
267            .pending
268            .remove(&(workflow_id.clone(), activity_id.clone()));
269        if let Some((_, sender)) = matched {
270            return sender.send(result).is_ok();
271        }
272        let Some(callback) = self.outbox_delivery.get() else {
273            // Outbox OFF: silent drop, byte-identical to the prior behaviour.
274            return false;
275        };
276        let outcome = match result {
277            Ok(payload) => callback.deliver_completion(workflow_id, activity_id, run_id, payload),
278            Err(reason) => callback.deliver_failure(workflow_id, activity_id, run_id, reason),
279        };
280        match outcome {
281            Ok(true) => true,
282            Ok(false) => {
283                // Not live: the expected stale-completion case recovery re-arms.
284                tracing::debug!(
285                    workflow_id = %workflow_id,
286                    activity_id = %activity_id,
287                    "unmatched outbox completion for a workflow that is not currently live; \
288                     recovery will re-arm it"
289                );
290                false
291            }
292            Err(error) => {
293                tracing::warn!(
294                    workflow_id = %workflow_id,
295                    activity_id = %activity_id,
296                    %error,
297                    "failed to deliver unmatched outbox completion to the live workflow"
298                );
299                false
300            }
301        }
302    }
303
304    fn complete_fenced_after_accept(
305        &self,
306        workflow_id: &WorkflowId,
307        activity_id: &ActivityId,
308        run_id: Option<&RunId>,
309        completion_token: &CompletionToken,
310        result: Result<String, String>,
311        after_accept: impl FnOnce() -> Result<bool, ServerError>,
312    ) -> Result<bool, ServerError> {
313        // The acceptance hands back the generation it consumed — every token
314        // outstanding for this attempt, not just the one submitted — so the two
315        // restore arms below put back exactly what was taken rather than
316        // re-deriving it and inventing a narrower generation.
317        let accepted = self
318            .completion_fences
319            .accept(workflow_id, activity_id, completion_token)
320            .inspect_err(|error| {
321                tracing::warn!(
322                    workflow_id = %workflow_id,
323                    activity_id = %activity_id,
324                    %error,
325                    "activity completion rejected by execution-generation fence"
326                );
327            })?;
328        let should_publish = match after_accept() {
329            Ok(should_publish) => should_publish,
330            Err(error) => {
331                if let Err(restore_error) =
332                    self.completion_fences
333                        .restore_if_absent(workflow_id, activity_id, &accepted)
334                {
335                    tracing::error!(%restore_error, "failed to restore completion generation after settlement failure");
336                }
337                return Err(error);
338            }
339        };
340        if !should_publish {
341            // A non-publishing settlement means ANOTHER path owns this
342            // dispatch's resolution (the lost-worker sweep, typically). That
343            // path will present the same token to `accept` — so the
344            // generation this call consumed must be restored, exactly as on
345            // the failure arm, or the true resolver is refused with
346            // `NoCurrentGeneration` and the waiter is never resolved.
347            if let Err(restore_error) =
348                self.completion_fences
349                    .restore_if_absent(workflow_id, activity_id, &accepted)
350            {
351                tracing::error!(%restore_error, "failed to restore completion generation after a non-publishing settlement");
352            }
353            return Ok(false);
354        }
355        // A resolution OUTSIDE the transport domain — a real result, or a real
356        // action failure — retires this execution site's transport-loss budget,
357        // so a run that survived one blip does not carry it forward. A
358        // transport-domain resolution deliberately does not: the running budget
359        // is exactly what bounds a flapping link.
360        let transport_domain = result
361            .as_ref()
362            .err()
363            .is_some_and(|reason| super::transport_loss::is_transport_domain_reason(reason));
364        if !transport_domain
365            && let Err(error) = self.transport_losses.clear(workflow_id, activity_id)
366        {
367            tracing::warn!(
368                workflow_id = %workflow_id,
369                activity_id = %activity_id,
370                %error,
371                "failed to retire the transport-loss budget for a resolved activity"
372            );
373        }
374        Ok(self.complete(workflow_id, activity_id, run_id, result))
375    }
376
377    /// Build a sink whose transport-loss ledger derives its budget from the
378    /// operator's heartbeat window (see
379    /// [`transport_loss`](crate::worker::transport_loss)).
380    ///
381    /// The window is a CONSTRUCTOR parameter rather than a builder step, and
382    /// that is the point: there is no way to obtain a `PendingActivities`
383    /// without naming the window, so a wiring cannot supply everything else and
384    /// silently inherit a zero budget. The previous shape —
385    /// `default().with_heartbeat_window(w)` — was two calls where the second was
386    /// omissible, and nothing in the test suite could tell the wired world from
387    /// the unwired one.
388    #[must_use]
389    pub fn new(heartbeat_window: std::time::Duration) -> Self {
390        Self {
391            pending: Arc::default(),
392            completion_fences: CompletionFences::default(),
393            outbox_delivery: Arc::default(),
394            transport_losses: super::transport_loss::TransportLossLedger::new(heartbeat_window),
395        }
396    }
397
398    /// The transport-loss ledger this sink classifies worker deaths through.
399    #[must_use]
400    pub const fn transport_losses(&self) -> &super::transport_loss::TransportLossLedger {
401        &self.transport_losses
402    }
403
404    /// Classify one worker loss for `(workflow_id, activity_id)` into the
405    /// transport-domain reason the engine seam consumes.
406    ///
407    /// A ledger failure (poisoned lock) is reported as transport exhaustion
408    /// rather than as a re-dispatchable loss: with no trustworthy budget the
409    /// only safe answer is the one that terminates, because an unbounded
410    /// re-dispatch is the failure mode the budget exists to prevent.
411    fn classify_worker_loss(
412        &self,
413        workflow_id: &WorkflowId,
414        activity_id: &ActivityId,
415        worker_id: crate::worker::registry::WorkerId,
416    ) -> String {
417        let detail = super::transport_loss::worker_lost_detail(worker_id);
418        match self
419            .transport_losses
420            .record_loss(workflow_id, activity_id, &detail)
421        {
422            Ok(verdict) => {
423                if verdict.exhausted {
424                    tracing::error!(
425                        operation = "activity_complete",
426                        workflow_id = %workflow_id,
427                        activity_id = %activity_id,
428                        worker_id = ?worker_id,
429                        error_type = "TransportExhausted",
430                        losses = verdict.losses,
431                        budget_ms = self.transport_losses.budget().as_millis(),
432                        "activity abandoned: the transport kept losing its worker past the \
433                         transport-loss budget"
434                    );
435                } else {
436                    tracing::warn!(
437                        operation = "activity_complete",
438                        workflow_id = %workflow_id,
439                        activity_id = %activity_id,
440                        worker_id = ?worker_id,
441                        error_type = "WorkerLost",
442                        losses = verdict.losses,
443                        budget_ms = self.transport_losses.budget().as_millis(),
444                        "worker lost before reporting an activity result; the activity never ran \
445                         and will be re-dispatched attempt-neutrally"
446                    );
447                }
448                verdict.reason
449            }
450            Err(error) => {
451                tracing::error!(
452                    workflow_id = %workflow_id,
453                    activity_id = %activity_id,
454                    %error,
455                    "transport-loss ledger is unreadable; abandoning the activity rather than \
456                     re-dispatching it without a budget"
457                );
458                format!(
459                    "{}{detail} (transport-loss budget unreadable: {error})",
460                    super::transport_loss::TRANSPORT_EXHAUSTED_REASON_PREFIX
461                )
462            }
463        }
464    }
465
466    pub(crate) fn complete_activity_after_accept(
467        &self,
468        completion: ActivityCompletion,
469        after_accept: impl FnOnce() -> Result<(), ServerError>,
470    ) -> Result<(), ServerError> {
471        let result = match completion.outcome {
472            ActivityCompletionOutcome::Succeeded(payload) => {
473                payload_to_string(&payload).map_err(|reason| {
474                    tracing::error!(
475                        operation = "activity_complete",
476                        workflow_id = %completion.workflow_id,
477                        activity_id = %completion.activity_id,
478                        error_type = "ActivityResultDecode",
479                        %reason,
480                        "activity completion failed"
481                    );
482                    ServerError::worker_dispatch("", "", format!("payload decode: {reason}"))
483                })?
484            }
485            ActivityCompletionOutcome::Failed(error) => {
486                let prefix = match error.kind {
487                    ActivityErrorKind::Retryable => "retryable",
488                    ActivityErrorKind::PolicyRefused => "policy_refused",
489                    ActivityErrorKind::Terminal => "terminal",
490                };
491                tracing::error!(
492                    operation = "activity_complete",
493                    workflow_id = %completion.workflow_id,
494                    activity_id = %completion.activity_id,
495                    error_type = "ActivityFailed",
496                    error_kind = prefix,
497                    reason = %error.message,
498                    "activity completion failed"
499                );
500                Err(format!("{prefix}:{}", error.message))
501            }
502            ActivityCompletionOutcome::WorkerLost { worker_id } => Err(self.classify_worker_loss(
503                &completion.workflow_id,
504                &completion.activity_id,
505                worker_id,
506            )),
507        };
508        let accepted_settlement = || after_accept().map(|()| true);
509        self.complete_fenced_after_accept(
510            &completion.workflow_id,
511            &completion.activity_id,
512            completion.run_id.as_ref(),
513            &completion.completion_token,
514            result,
515            accepted_settlement,
516        )?;
517        Ok(())
518    }
519}
520
521impl ActivityCompletionSink for PendingActivities {
522    /// Park one in-flight dispatch for restart recovery (#207): resolve the
523    /// matched waiter with the ephemeral parked sentinel, and nothing else.
524    ///
525    /// This resolution is MANDATORY, not an optimization: the default
526    /// `ActivityDispatcher::dispatch_async` runs the dispatcher's blocking
527    /// `std::sync::mpsc::recv()` on tokio's blocking pool, and tokio `Runtime`
528    /// drop joins blocking threads — an unresolved waiter would wedge process
529    /// exit indefinitely. An unmatched dispatch (already resolved by another
530    /// path) is a no-op: a park is NEVER routed to the outbox delivery
531    /// callback, because it is not a failure and must never reach a workflow.
532    fn park_activity(
533        &self,
534        workflow_id: &WorkflowId,
535        activity_id: &ActivityId,
536    ) -> Result<(), ServerError> {
537        // A park retires the execution site, so it takes the whole current
538        // generation — including a redelivery's sibling token, whose worker is
539        // being parked for the same restart.
540        self.completion_fences
541            .revoke_current(workflow_id, activity_id)?;
542        let matched = self
543            .pending
544            .remove(&(workflow_id.clone(), activity_id.clone()));
545        if let Some((_, sender)) = matched {
546            // A send failure means the waiter side already dropped its
547            // receiver (the dispatch is being cleaned up concurrently) —
548            // benign: there is no thread left to unblock.
549            let _ = sender.send(Err(aion::PARKED_ACTIVITY_REASON.to_owned()));
550        }
551        Ok(())
552    }
553
554    fn complete_activity(&self, completion: ActivityCompletion) -> Result<(), ServerError> {
555        self.complete_activity_after_accept(completion, || Ok(()))
556    }
557}
558
559fn payload_to_string(payload: &Payload) -> Result<Result<String, String>, String> {
560    match payload.content_type() {
561        ContentType::Json => String::from_utf8(payload.bytes().to_vec())
562            .map(Ok)
563            .map_err(|_| "activity result payload is not valid UTF-8".to_owned()),
564    }
565}
566
567/// Dispatcher that routes `run_activity` NIF calls to connected workers.
568///
569/// Synchronous interface — uses `try_send` for the task channel and
570/// `std::sync::mpsc::Receiver::recv` for the response. Callers on a
571/// multi-thread tokio runtime are detected and moved into
572/// `tokio::task::block_in_place` so the blocking wait never starves the
573/// runtime tasks that flush the worker stream (see the module docs).
574pub struct WorkerActivityDispatcher {
575    registry: ConnectedWorkerRegistry,
576    namespace: String,
577    pending: PendingActivities,
578    heartbeat_tracker: HeartbeatTracker,
579    drain_state: DrainState,
580    tokio_handle: Option<tokio::runtime::Handle>,
581    /// NOI-6 attempt→owner back-index. When installed, each liminal-delivered
582    /// dispatch binds its `(workflow, activity, attempt)` to the owning worker
583    /// for the dispatch's lifetime, so the intervention router (and the ops
584    /// console's live-attempts enumeration) can see and target it. `None`
585    /// (isolated tests) binds nothing.
586    attempt_owners: Option<super::intervention::AttemptOwnerIndex>,
587    /// R1 service policies and the two service clocks. Default: `strict` with
588    /// no clocks — refuse the structurally unservable, wait (loudly) for
589    /// everything else, and invent no deadline the operator never wrote.
590    queue_service: QueueServiceConfig,
591    /// Deployed queue declarations, installed once the engine exists. Answers
592    /// `Unknown` until then, and `Unknown` never refuses anything.
593    queue_declarations: QueueDeclarationSource,
594    /// Live unserved-queue state a parked dispatch publishes itself into.
595    queue_state: QueueServiceState,
596    /// Cluster-event publisher an unbounded park announces itself on (#266
597    /// T4). `None` (isolated tests) loses only the pushed echo; the WARN and
598    /// the queryable state above cannot be switched off by wiring.
599    cluster_publisher: Option<crate::cluster_publisher::ClusterEventPublisher>,
600}
601
602impl std::fmt::Debug for WorkerActivityDispatcher {
603    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
604        f.debug_struct("WorkerActivityDispatcher")
605            .field("namespace", &self.namespace)
606            .finish_non_exhaustive()
607    }
608}
609
610impl WorkerActivityDispatcher {
611    /// Build a dispatcher for the given namespace, worker registry, and
612    /// liveness tracker.
613    ///
614    /// The tracker must be the same instance the worker stream handler and
615    /// shutdown coordinator share: the unbounded completion wait relies on
616    /// stream teardown sweeping this tracker's in-flight entries to fail
617    /// dispatches whose worker was lost.
618    #[must_use]
619    pub fn new(
620        registry: ConnectedWorkerRegistry,
621        namespace: impl Into<String>,
622        heartbeat_tracker: HeartbeatTracker,
623    ) -> Self {
624        Self {
625            registry,
626            namespace: namespace.into(),
627            // Derived from the tracker this constructor is already handed,
628            // which already carries the operator's window: the dispatcher
629            // cannot be given a heartbeat window and a mismatched budget.
630            pending: PendingActivities::new(heartbeat_tracker.heartbeat_window()),
631            heartbeat_tracker,
632            drain_state: DrainState::default(),
633            tokio_handle: None,
634            attempt_owners: None,
635            queue_service: QueueServiceConfig::default(),
636            queue_declarations: QueueDeclarationSource::default(),
637            queue_state: QueueServiceState::default(),
638            cluster_publisher: None,
639        }
640    }
641
642    /// Share the deployment-global cluster-event publisher so a dispatch
643    /// parked with no availability deadline is announced on the operator's
644    /// real-time channel, not only in the log (#266 T4).
645    #[must_use]
646    pub fn with_cluster_publisher(
647        mut self,
648        cluster_publisher: crate::cluster_publisher::ClusterEventPublisher,
649    ) -> Self {
650        self.cluster_publisher = Some(cluster_publisher);
651        self
652    }
653
654    /// Share the operator's R1 queue-service settings: the default policy, the
655    /// written per-queue `durable_pending` opt-ins, and the two clocks.
656    #[must_use]
657    pub fn with_queue_service(mut self, queue_service: QueueServiceConfig) -> Self {
658        self.queue_service = queue_service;
659        self
660    }
661
662    /// Share the queue-declaration source the boot path fills in once the
663    /// engine exists (the handle is cloneable; installing on any clone is
664    /// visible here).
665    #[must_use]
666    pub fn with_queue_declarations(mut self, queue_declarations: QueueDeclarationSource) -> Self {
667        self.queue_declarations = queue_declarations;
668        self
669    }
670
671    /// Share the queue-service state so the server can read which addresses are
672    /// unserved and which runs are parked on them.
673    #[must_use]
674    pub fn with_queue_state(mut self, queue_state: QueueServiceState) -> Self {
675        self.queue_state = queue_state;
676        self
677    }
678
679    /// Share the server's NOI-6 attempt→owner back-index so liminal-delivered
680    /// dispatches are visible (and targetable) to the intervention router for
681    /// exactly as long as they are in flight. The production boot passes
682    /// `ServerState`'s index — the SAME instance `intervenable_attempts` and
683    /// `intervene` read — or the console's live-attempts list stays empty for
684    /// every bridge-dispatched agent step.
685    #[must_use]
686    pub fn with_attempt_owners(
687        mut self,
688        attempt_owners: super::intervention::AttemptOwnerIndex,
689    ) -> Self {
690        self.attempt_owners = Some(attempt_owners);
691        self
692    }
693
694    /// The completion sink this dispatcher classifies worker deaths through.
695    ///
696    /// Exposed so the transport-loss budget a wiring actually ends up with is
697    /// observable, rather than being a property only prose asserts.
698    #[must_use]
699    pub const fn pending(&self) -> &PendingActivities {
700        &self.pending
701    }
702
703    /// Share a caller-supplied pending-activities tracker.
704    #[must_use]
705    pub fn with_pending(mut self, pending: PendingActivities) -> Self {
706        self.pending = pending;
707        self
708    }
709
710    /// Share the server drain gate.
711    #[must_use]
712    pub fn with_drain_state(mut self, drain_state: DrainState) -> Self {
713        self.drain_state = drain_state;
714        self
715    }
716
717    /// Share the server runtime handle for sync history writes from dirty NIF threads.
718    #[must_use]
719    pub fn with_tokio_handle(mut self, tokio_handle: tokio::runtime::Handle) -> Self {
720        self.tokio_handle = Some(tokio_handle);
721        self
722    }
723}
724
725impl WorkerActivityDispatcher {
726    /// The drain gate for one dispatch, and the single place a draining server's
727    /// refusal is turned into a reason the engine acts on.
728    ///
729    /// The reason handed back is the #207 parked sentinel
730    /// ([`aion::PARKED_ACTIVITY_REASON`]), NOT the gate's own message. A
731    /// dispatch the drain refused never reached a worker and never ran, which is
732    /// precisely the state that sentinel was minted for: record nothing, deliver
733    /// nothing, and let post-restart replay re-dispatch the dangling ordinal —
734    /// the kill -9 convergence `shutdown::ShutdownOutcome::Parked` documents.
735    ///
736    /// Both alternatives are wrong here, and neither is a style preference:
737    ///
738    /// - the *unprefixed* gate message this used to return is classified as a
739    ///   terminal action failure, so a routine deploy records `ActivityFailed`
740    ///   (and, with no authored retry, `WorkflowFailed`) for work nobody
741    ///   attempted — the synthesized failure #207 exists to abolish;
742    /// - the TRANSPORT-loss class (`lost:`) is attempt-neutral but re-dispatches
743    ///   the SAME attempt live and in-process, which on a draining server means
744    ///   refuse → re-dispatch → refuse, spinning against the very gate that
745    ///   refused it while the process is trying to exit.
746    ///
747    /// The operator-facing detail is not lost: the gate's own message is what is
748    /// logged here. Only the engine-facing classification is the sentinel.
749    fn ensure_accepting(
750        &self,
751        namespace: &str,
752        activity_type: &str,
753        workflow_id: &WorkflowId,
754        activity_id: &ActivityId,
755        worker_id: Option<WorkerId>,
756    ) -> Result<(), String> {
757        self.drain_state
758            .ensure_accepting(namespace, activity_type)
759            .map_err(|error| {
760                log_worker_error(
761                    "WorkerDispatch",
762                    namespace,
763                    activity_type,
764                    workflow_id,
765                    activity_id,
766                    worker_id,
767                    &error.to_string(),
768                );
769                aion::PARKED_ACTIVITY_REASON.to_owned()
770            })
771    }
772
773    /// Select a worker for the address, or refuse with a typed R1 reason.
774    ///
775    /// The wait itself lives in [`super::queue_service::wait`]: every selection
776    /// miss is classified against the deployed contract records and the live
777    /// poller census, published to the queue-service state, and stated at WARN.
778    /// This method owns only the two things the seam cannot delegate — the
779    /// drain gate and the runtime plumbing of one bounded park.
780    fn select_worker_or_wait(
781        &self,
782        address: &ServiceAddress,
783        workflow_id: &WorkflowId,
784        activity_id: &ActivityId,
785    ) -> Result<WorkerHandle, String> {
786        let wait = ServiceWait {
787            registry: &self.registry,
788            declarations: &self.queue_declarations,
789            config: &self.queue_service,
790            state: &self.queue_state,
791            address,
792            workflow_id,
793            activity_id,
794            publisher: self.cluster_publisher.as_ref(),
795        };
796        let mut accepting = || {
797            self.ensure_accepting(
798                &address.namespace,
799                &address.activity_type,
800                workflow_id,
801                activity_id,
802                None,
803            )
804        };
805        let mut park = |arrival: WorkerArrival, budget: Option<Duration>| {
806            self.park_for_worker(arrival, budget);
807        };
808        select_worker_or_refuse(&wait, &mut accepting, &mut park).map_err(|refusal| {
809            let reason = refusal.reason_string();
810            // The drain gate already logged its own refusal at its own site;
811            // logging it twice would double-count the incident.
812            if !matches!(refusal, SelectionRefusal::NotAccepting { .. }) {
813                let error_type = match refusal {
814                    SelectionRefusal::Unavailable(_) => "WorkerUnavailable",
815                    _ => "WorkerRegistry",
816                };
817                log_worker_error(
818                    error_type,
819                    &address.namespace,
820                    &address.activity_type,
821                    workflow_id,
822                    activity_id,
823                    None,
824                    &reason,
825                );
826            }
827            reason
828        })
829    }
830
831    /// Wait once on the `arrival` subscription — bounded by `budget` when one is
832    /// given, and bounded in BOTH arms by the server's own shutdown.
833    ///
834    /// `arrival` was taken by [`select_worker_or_refuse`] BEFORE the selection
835    /// that missed, which is why this method is handed one rather than taking
836    /// its own. Subscribing here would be too late: the registry's wake sources
837    /// store no permit, so a worker that registered during that selection would
838    /// already have fired into an empty waiter list and this park would sleep
839    /// through it.
840    ///
841    /// The arrival wait races the drain latch, and that race is the whole point.
842    /// Without it a `None` budget — the default, because no schedule-to-start
843    /// deadline is configured unless an operator writes one — parked this
844    /// dispatch on a `spawn_blocking` thread with nothing left that could ever
845    /// wake it: the drain gate is consulted once per selection iteration, and
846    /// the iteration was stuck inside this wait. Tokio's runtime `Drop` joins
847    /// the blocking pool, so `aion server` could not exit even after `main`
848    /// returned an exit code — it kept the store's writer lock and blocked the
849    /// next deploy into the same data directory. The captured stack of a real
850    /// one is `docs/evidence/samples/aion-server-hang-after-drain-15202.sample`.
851    /// The `Some(budget)` arm carried the same defect at the scale of the
852    /// budget: exit was delayed by up to the whole availability deadline.
853    ///
854    /// Racing the latch is the fix rather than capping the wait, because the
855    /// wait is legitimately unbounded while the server runs: a queue no worker
856    /// serves yet is a fleet condition an operator resolves by starting a
857    /// worker, and a deadline invented here would refuse work nobody asked to
858    /// have refused.
859    ///
860    /// Waking on the latch decides nothing by itself. The selection loop
861    /// re-selects, misses, and consults the SAME drain gate it always
862    /// consulted — [`Self::ensure_accepting`] — which is where the refusal and
863    /// its classification are produced. This seam gains no second opinion about
864    /// draining; it only stops being unwakeable.
865    fn park_for_worker(&self, arrival: WorkerArrival, budget: Option<Duration>) {
866        let handle = self
867            .tokio_handle
868            .clone()
869            .or_else(|| tokio::runtime::Handle::try_current().ok());
870        let Some(handle) = handle else {
871            // No runtime in reach (a plain OS thread in an isolated test):
872            // sleep the seam's park interval, clamped to the remaining budget.
873            // Bounded by construction, so the loop re-consults the drain gate
874            // within one interval and needs no signal to race.
875            //
876            // `arrival` is DROPPED unawaited here, deliberately: there is no
877            // executor to await it on. That is safe only because this arm is
878            // bounded by construction — the loop re-selects within one interval
879            // whether or not any wake ever fires — and it is the one place the
880            // `park` contract permits discarding the subscription. Its `Drop`
881            // unlinks the waiter it enabled.
882            std::thread::sleep(
883                budget.map_or(PARK_POLL_INTERVAL, |budget| budget.min(PARK_POLL_INTERVAL)),
884            );
885            return;
886        };
887        handle.block_on(async {
888            let raced = async {
889                tokio::select! {
890                    () = arrival => {}
891                    () = self.drain_state.wait_for_drain() => {}
892                }
893            };
894            match budget {
895                None => raced.await,
896                Some(budget) => {
897                    // Timeout is not failure here: the caller re-classifies and
898                    // decides whether the clock has run out.
899                    drop(tokio::time::timeout(budget, raced).await);
900                }
901            }
902        });
903    }
904
905    fn track_worker_task(
906        &self,
907        worker_id: WorkerId,
908        activity_type: &str,
909        workflow_id: &WorkflowId,
910        activity_id: &ActivityId,
911        attempt: u32,
912        completion_token: CompletionToken,
913    ) -> Result<(), String> {
914        self.heartbeat_tracker
915            .track_task(
916                worker_id,
917                InFlightActivity {
918                    workflow_id: workflow_id.clone(),
919                    activity_id: activity_id.clone(),
920                    // The engine-provided attempt this delivery carries — the
921                    // same one stamped on the wire task and bound into the
922                    // NOI-6 attempt→owner index, so the tracker, the wire, and
923                    // the owner index all name one attempt.
924                    attempt,
925                    completion_token,
926                },
927                Instant::now(),
928            )
929            .map_err(|error| {
930                let reason = error.to_string();
931                log_worker_error(
932                    "WorkerHeartbeatTracker",
933                    &self.namespace,
934                    activity_type,
935                    workflow_id,
936                    activity_id,
937                    Some(worker_id),
938                    &reason,
939                );
940                reason
941            })
942    }
943
944    fn cleanup_activity(
945        &self,
946        worker_id: WorkerId,
947        workflow_id: &WorkflowId,
948        activity_id: &ActivityId,
949        completion_token: &CompletionToken,
950    ) {
951        self.pending
952            .pending
953            .remove(&(workflow_id.clone(), activity_id.clone()));
954        // Own-token-only: an undelivered task withdraws the authorization it
955        // was given and never the sibling a concurrently redelivered dispatch
956        // handed to a worker that is still executing.
957        if let Err(error) =
958            self.pending
959                .completion_fences
960                .revoke(workflow_id, activity_id, completion_token)
961        {
962            tracing::error!(
963                workflow_id = %workflow_id,
964                activity_id = %activity_id,
965                %error,
966                "failed to revoke undelivered activity generation"
967            );
968        }
969        let _ = self
970            .heartbeat_tracker
971            .complete_task(worker_id, workflow_id, activity_id);
972        self.drain_state.notify_activity_drained();
973    }
974
975    /// Deliver one dispatched task to the selected worker over ITS transport.
976    ///
977    /// The bridge is transport-agnostic at this seam: selection
978    /// ([`Self::select_worker_or_wait`]) already treats every registry member
979    /// identically, and this match delivers on whichever [`WorkerDelivery`] leg
980    /// the worker registered with — the gRPC stream `mpsc` push, or the liminal
981    /// server-push on the worker's existing connection. Both legs resolve
982    /// through the SAME pending map, so `await_activity_result` is oblivious to
983    /// the transport.
984    fn send_activity_task(
985        &self,
986        worker: &WorkerHandle,
987        task: ProtoActivityTask,
988        address: &ServiceAddress,
989        workflow_id: &WorkflowId,
990        activity_id: &ActivityId,
991        completion_token: &CompletionToken,
992    ) -> Result<(), String> {
993        match worker.delivery() {
994            WorkerDelivery::Grpc(sender) => {
995                let worker_id = worker.id();
996                let mut accepting = || {
997                    self.ensure_accepting(
998                        &address.namespace,
999                        &address.activity_type,
1000                        workflow_id,
1001                        activity_id,
1002                        Some(worker_id),
1003                    )
1004                };
1005                let handed_over = deliver_within_schedule_to_start(
1006                    sender,
1007                    WorkerMessage::ActivityTask(Box::new(task)),
1008                    self.queue_service.schedule_to_start_timeout,
1009                    &mut accepting,
1010                );
1011                let Err(refusal) = handed_over else {
1012                    return Ok(());
1013                };
1014                self.cleanup_activity(worker_id, workflow_id, activity_id, completion_token);
1015                let (error_type, reason) = self.hand_off_failure(&refusal, address);
1016                if !matches!(refusal, DeliveryRefusal::NotAccepting { .. }) {
1017                    log_worker_error(
1018                        error_type,
1019                        &address.namespace,
1020                        &address.activity_type,
1021                        workflow_id,
1022                        activity_id,
1023                        Some(worker_id),
1024                        &reason,
1025                    );
1026                }
1027                Err(reason)
1028            }
1029            // The liminal leg has no bounded intake queue to saturate: the push
1030            // either enqueues on the worker's live connection or fails
1031            // synchronously because that connection is already gone. There is
1032            // therefore no schedule-to-start window to apply here, and reporting
1033            // `SATURATED` for a push failure would name a condition that did not
1034            // happen.
1035            #[cfg(feature = "liminal-transport")]
1036            WorkerDelivery::Liminal(delivery) => self.send_liminal_activity_task(
1037                worker.id(),
1038                delivery,
1039                task,
1040                &address.activity_type,
1041                workflow_id,
1042                activity_id,
1043            ),
1044        }
1045    }
1046
1047    /// Render one hand-off refusal into its log class and failure reason.
1048    ///
1049    /// `SATURATED` is the only one that becomes a typed [`WorkerUnavailable`]:
1050    /// a compatible worker WAS live and would not take the task inside the
1051    /// schedule-to-start clock. A full intake with no clock configured, and a
1052    /// closed transport, keep the failure they always had.
1053    fn hand_off_failure(
1054        &self,
1055        refusal: &DeliveryRefusal,
1056        address: &ServiceAddress,
1057    ) -> (&'static str, String) {
1058        match refusal {
1059            DeliveryRefusal::Saturated { waited } => {
1060                // A census failure must not swallow the refusal: report the
1061                // saturation with an empty census and say why it is empty.
1062                let census = self
1063                    .registry
1064                    .pool_census(
1065                        &address.namespace,
1066                        &address.task_queue,
1067                        &address.activity_type,
1068                        address.node.as_deref(),
1069                    )
1070                    .unwrap_or_else(|error| {
1071                        tracing::error!(
1072                            namespace = %address.namespace,
1073                            task_queue = %address.task_queue,
1074                            activity_type = %address.activity_type,
1075                            %error,
1076                            "poller census failed while reporting a saturated queue; \
1077                             the refusal carries an empty census"
1078                        );
1079                        PoolCensus::default()
1080                    });
1081                let unavailable = WorkerUnavailable {
1082                    reason: QueueServiceReason::Saturated,
1083                    clock: Some(ExpiredClock::ScheduleToStart),
1084                    waited: *waited,
1085                    address: address.clone(),
1086                    census,
1087                };
1088                ("WorkerUnavailable", unavailable.reason_string())
1089            }
1090            DeliveryRefusal::Full => (
1091                "WorkerChannelClosed",
1092                "worker task channel full or closed: no available capacity".to_owned(),
1093            ),
1094            DeliveryRefusal::Closed => (
1095                "WorkerChannelClosed",
1096                "worker task channel full or closed: channel closed".to_owned(),
1097            ),
1098            DeliveryRefusal::NotAccepting { reason } => ("WorkerDispatch", reason.clone()),
1099        }
1100    }
1101
1102    /// Deliver one dispatched task to a liminal-connected worker: push the SAME
1103    /// wire frame the outbox liminal path pushes (a
1104    /// [`DispatchRequest`](super::liminal_transport::DispatchRequest) — the
1105    /// worker's serve loop cannot tell a bridge dispatch from an outbox row) and
1106    /// hand the correlated-reply awaiter to a dedicated router thread that
1107    /// resolves this dispatch's pending entry exactly like a gRPC completion.
1108    ///
1109    /// The wire carries the SAME engine-provided `attempt` and `labels` the gRPC
1110    /// arm's `ActivityTask` carries (a retry over liminal executes with the real
1111    /// attempt, not a re-stamped first delivery), plus the server's heartbeat
1112    /// window so the worker's automatic liveness pump keeps this TRACKED
1113    /// dispatch alive under the #176 expiry sweeper. It carries the engine's
1114    /// concrete run identity, generation proof, and stable effect key exactly
1115    /// like the gRPC bridge task.
1116    ///
1117    /// A successful push also binds the attempt into the NOI-6 attempt→owner
1118    /// back-index (when installed) with the SAME `(workflow, activity, attempt)`
1119    /// key the worker stamps its intervention session with, exactly as the
1120    /// outbox liminal arm binds each row dispatch — so the ops console can
1121    /// enumerate this live attempt and route interventions to its worker. The
1122    /// binding is released when the reply router exits (reply, abandonment, or
1123    /// disconnect — every path). The gRPC arm carries no bind because the agent
1124    /// harness seam exists only on the liminal worker transport.
1125    #[cfg(feature = "liminal-transport")]
1126    fn send_liminal_activity_task(
1127        &self,
1128        worker_id: WorkerId,
1129        delivery: &super::liminal_transport::LiminalWorkerDelivery,
1130        task: ProtoActivityTask,
1131        activity_type: &str,
1132        workflow_id: &WorkflowId,
1133        activity_id: &ActivityId,
1134    ) -> Result<(), String> {
1135        let completion_token =
1136            CompletionToken::from_wire(workflow_id, activity_id, task.completion_token.clone())
1137                .map_err(|error| error.to_string())?;
1138        let heartbeat_window_ms =
1139            u64::try_from(self.heartbeat_tracker.heartbeat_window().as_millis())
1140                .unwrap_or(u64::MAX);
1141        let attempt = task.attempt;
1142        // The run is resolved BEFORE the dispatch is built, because everything
1143        // downstream is keyed on it: the worker refuses a run-less dispatch
1144        // envelope outright, the attempt-owner binding below needs it to name a
1145        // generation, and the attempt's transcript is written under it. Refusing
1146        // here reports the fault against the dispatch that has it, rather than
1147        // shipping a frame whose only possible outcome is the worker's rejection.
1148        let run_id = task
1149            .run_id
1150            .map(RunId::try_from)
1151            .transpose()
1152            .map_err(|error| error.to_string())?
1153            .ok_or_else(|| {
1154                "activity task run id is missing; refusing to dispatch an unidentified run"
1155                    .to_owned()
1156            })?;
1157        let request = super::liminal_transport::DispatchRequest {
1158            activity_type: activity_type.to_owned(),
1159            workflow_id: workflow_id.clone(),
1160            ordinal: activity_id.sequence_position(),
1161            run_id: Some(run_id.clone()),
1162            attempt,
1163            completion_token: task.completion_token,
1164            idempotency_key: task.idempotency_key,
1165            labels: task.labels.into_iter().collect(),
1166            heartbeat_window_ms,
1167            input: task.input.map(|payload| payload.bytes).unwrap_or_default(),
1168        };
1169        // A push-enqueue failure means the worker's connection was already gone
1170        // at push time — the same synchronous-failure contract as a closed gRPC
1171        // stream channel above.
1172        let awaiter = match delivery.push_dispatch(&request) {
1173            Ok(awaiter) => awaiter,
1174            Err(error) => {
1175                let reason = format!("worker liminal push failed: {error}");
1176                self.cleanup_activity(worker_id, workflow_id, activity_id, &completion_token);
1177                log_worker_error(
1178                    "WorkerChannelClosed",
1179                    &self.namespace,
1180                    activity_type,
1181                    workflow_id,
1182                    activity_id,
1183                    Some(worker_id),
1184                    &reason,
1185                );
1186                return Err(reason);
1187            }
1188        };
1189        // NOI-6: the attempt is live on `worker_id` from this push until the
1190        // router resolves it — bind it for exactly that window (the guard is
1191        // dropped when the router thread exits).
1192        let owner_binding = self.attempt_owners.as_ref().map(|owners| {
1193            super::liminal_transport::AttemptOwnerGuard::bind(
1194                owners.clone(),
1195                super::intervention::AttemptKey::new(
1196                    workflow_id.clone(),
1197                    run_id.clone(),
1198                    activity_id.clone(),
1199                    attempt,
1200                ),
1201                worker_id,
1202            )
1203        });
1204        self.spawn_liminal_reply_router(
1205            worker_id,
1206            awaiter,
1207            workflow_id,
1208            activity_id,
1209            &completion_token,
1210            owner_binding,
1211        );
1212        Ok(())
1213    }
1214
1215    /// Waits (on a dedicated router thread, bounded by the dispatch's own
1216    /// lifetime) for the worker's correlated
1217    /// [`DispatchResponse`](super::liminal_transport::DispatchResponse) and
1218    /// re-enters it through the SAME completion bookkeeping the gRPC inbound
1219    /// stream applies (`process_inbound` in `worker_grpc.rs`): clear the
1220    /// in-flight liveness entry, wake any drain waiter, then resolve the
1221    /// bridge's pending map — result, failure, and retryable classification
1222    /// identical (the worker encodes the `retryable:`/`terminal:` reason
1223    /// vocabulary on the wire). An unmatched (already-resolved) REAL reply
1224    /// routes through the outbox delivery callback exactly like a late gRPC
1225    /// result.
1226    ///
1227    /// The #176 heartbeat sweeper covers this dispatch exactly as it covers a
1228    /// gRPC one: the dispatch is tracked in the shared [`HeartbeatTracker`] and
1229    /// the worker's runtime pumps automatic liveness beats over the reserved
1230    /// liminal channel (`WORKER_LIVENESS_CHANNEL`), so a healthy worker running
1231    /// an over-window activity is never falsely expired while a wedged one
1232    /// still is. Prompt worker-DEATH detection additionally rides the
1233    /// connection itself — the awaiter wakes with the typed Disconnected error
1234    /// the moment the connection closes, resolving the SAME retryable
1235    /// lost-worker failure the gRPC stream-teardown sweep reports.
1236    ///
1237    /// Two structural guards mirror the gRPC arm's tracker gating:
1238    ///
1239    /// - A SYNTHESIZED failure (disconnect / receive fault) is delivered only
1240    ///   when this router's own `complete_task` actually retired the tracked
1241    ///   entry — the same "fail only still-tracked tasks" gate
1242    ///   `remove_worker_tasks` gives the gRPC sweeps — so a dispatch already
1243    ///   resolved elsewhere (expiry sweep, shutdown drain, deregistered
1244    ///   fast path) never has a spurious failure injected for an ordinal whose
1245    ///   retry may be live on another worker.
1246    /// - The wait itself ends one reply-poll after the tracked entry
1247    ///   disappears, so an abandoned dispatch never parks this thread for the
1248    ///   remaining life of the worker's connection. A real reply arriving
1249    ///   AFTER that exit is dropped (the resolving path owns the ordinal — its
1250    ///   retry re-executes); this is the one deliberate divergence from the
1251    ///   gRPC arm, whose shared stream task routes any late result to the
1252    ///   outbox callback, and it is the safer half of the trade because a
1253    ///   stale attempt's result can never resolve a newer attempt's entry.
1254    #[cfg(feature = "liminal-transport")]
1255    fn spawn_liminal_reply_router(
1256        &self,
1257        worker_id: WorkerId,
1258        awaiter: liminal_server::server::connection::PushReplyAwaiter,
1259        workflow_id: &WorkflowId,
1260        activity_id: &ActivityId,
1261        completion_token: &CompletionToken,
1262        owner_binding: Option<super::liminal_transport::AttemptOwnerGuard>,
1263    ) {
1264        let pending = self.pending.clone();
1265        let heartbeat_tracker = self.heartbeat_tracker.clone();
1266        let drain_state = self.drain_state.clone();
1267        let workflow_id = workflow_id.clone();
1268        let activity_id = activity_id.clone();
1269        let completion_token = completion_token.clone();
1270        std::thread::spawn(move || {
1271            // Owns the NOI-6 attempt binding for the dispatch's lifetime: it
1272            // drops (releasing the back-index entry) when this router exits,
1273            // on every path — reply, abandonment, disconnect, or panic.
1274            let _owner_binding = owner_binding;
1275            route_liminal_reply(
1276                &pending,
1277                &heartbeat_tracker,
1278                &drain_state,
1279                &awaiter,
1280                (worker_id, &workflow_id, &activity_id, &completion_token),
1281            );
1282        });
1283    }
1284
1285    /// Block until the dispatch terminates (see the module docs for the
1286    /// exhaustive termination list). The wait is deliberately unbounded:
1287    /// the engine imposes no activity timeout of its own.
1288    fn await_activity_result(
1289        &self,
1290        context: &ActivityDispatchContext<'_>,
1291        rx: &SyncReceiver,
1292    ) -> Result<String, String> {
1293        // Close the dispatch/disconnect race before blocking. A worker whose
1294        // stream tore down *before* this dispatch tracked its task was swept
1295        // without this entry, so nothing would ever deliver through `rx`.
1296        // `fail_lost_worker` deregisters before it collects tasks, and this
1297        // dispatch tracked its task before sending, so: if the worker is
1298        // still registered here, any later sweep is guaranteed to include
1299        // this task and unblock the `recv` below.
1300        match self.registry.is_registered(context.worker_id) {
1301            Ok(true) => {}
1302            Ok(false) => {
1303                // A sweep that did include this task may have delivered
1304                // already; prefer its verdict (or a genuine result that
1305                // raced the disconnect) over fabricating one.
1306                if let Ok(result) = rx.try_recv() {
1307                    return self.deliver_result(context, result);
1308                }
1309                self.cleanup_activity(
1310                    context.worker_id,
1311                    context.workflow_id,
1312                    context.activity_id,
1313                    &context.completion_token,
1314                );
1315                // TRANSPORT domain, classified through the shared ledger: the
1316                // activity never executed, so the failure never wears the
1317                // action's retry vocabulary and is re-dispatched
1318                // attempt-neutrally until the transport's own budget is spent.
1319                let reason = self.pending.classify_worker_loss(
1320                    context.workflow_id,
1321                    context.activity_id,
1322                    context.worker_id,
1323                );
1324                log_worker_error(
1325                    "WorkerLost",
1326                    &self.namespace,
1327                    context.activity_type,
1328                    context.workflow_id,
1329                    context.activity_id,
1330                    Some(context.worker_id),
1331                    &reason,
1332                );
1333                return Err(reason);
1334            }
1335            Err(error) => {
1336                self.cleanup_activity(
1337                    context.worker_id,
1338                    context.workflow_id,
1339                    context.activity_id,
1340                    &context.completion_token,
1341                );
1342                let reason = format!("worker registry inspection failed: {error}");
1343                log_worker_error(
1344                    "WorkerRegistry",
1345                    &self.namespace,
1346                    context.activity_type,
1347                    context.workflow_id,
1348                    context.activity_id,
1349                    Some(context.worker_id),
1350                    &reason,
1351                );
1352                return Err(reason);
1353            }
1354        }
1355        if let Ok(result) = rx.recv() {
1356            return self.deliver_result(context, result);
1357        }
1358        // Every sender was dropped without completing: a cleanup path
1359        // removed the pending entry. Surface it instead of hanging.
1360        self.cleanup_activity(
1361            context.worker_id,
1362            context.workflow_id,
1363            context.activity_id,
1364            &context.completion_token,
1365        );
1366        let reason = "activity response channel dropped".to_owned();
1367        log_worker_error(
1368            "WorkerChannelClosed",
1369            &self.namespace,
1370            context.activity_type,
1371            context.workflow_id,
1372            context.activity_id,
1373            Some(context.worker_id),
1374            &reason,
1375        );
1376        Err(reason)
1377    }
1378
1379    fn deliver_result(
1380        &self,
1381        context: &ActivityDispatchContext<'_>,
1382        result: Result<String, String>,
1383    ) -> Result<String, String> {
1384        self.pending
1385            .pending
1386            .remove(&(context.workflow_id.clone(), context.activity_id.clone()));
1387        // A parked dispatch (#207) is not a failure: the server is draining and
1388        // restart recovery re-dispatches the ordinal. Info, never error — a
1389        // routine deploy must not emit an ActivityFailed log per in-flight
1390        // dispatch (the incident's alarm noise).
1391        if let Err(reason) = &result
1392            && aion::is_parked_reason(reason)
1393        {
1394            tracing::info!(
1395                operation = "activity_dispatch",
1396                namespace = %self.namespace,
1397                workflow_id = %context.workflow_id,
1398                activity_id = %context.activity_id,
1399                activity_type = context.activity_type,
1400                worker_id = ?context.worker_id,
1401                "activity parked for restart recovery"
1402            );
1403            return result;
1404        }
1405        log_activity_completion(context, result.is_ok());
1406        result.inspect_err(|reason| {
1407            log_worker_error(
1408                "ActivityFailed",
1409                &self.namespace,
1410                context.activity_type,
1411                context.workflow_id,
1412                context.activity_id,
1413                Some(context.worker_id),
1414                reason,
1415            );
1416        })
1417    }
1418}
1419
1420impl ActivityDispatcher for WorkerActivityDispatcher {
1421    fn dispatch(&self, request: ActivityDispatch) -> Result<String, String> {
1422        match tokio::runtime::Handle::try_current() {
1423            Ok(handle) => match handle.runtime_flavor() {
1424                tokio::runtime::RuntimeFlavor::MultiThread => {
1425                    // We are inside a tokio runtime (the engine spawns the
1426                    // sync dispatch onto its handle). Hand this worker's
1427                    // scheduler core to another thread before blocking so the
1428                    // stream forwarder woken by our `try_send` can actually
1429                    // run — otherwise it is trapped in this worker's
1430                    // non-stealable LIFO slot for as long as we block.
1431                    tokio::task::block_in_place(|| self.dispatch_blocking(request))
1432                }
1433                flavor => Err(format!(
1434                    "activity dispatch blocks the calling thread until the worker responds; \
1435                     a {flavor:?} tokio runtime cannot host that wait because the worker \
1436                     stream forwarder shares its only executor thread and the task could \
1437                     never be delivered — run the engine on a multi-thread tokio runtime"
1438                )),
1439            },
1440            // No tokio context: a beamr scheduler thread or other plain OS
1441            // thread. Blocking here is the designed contract and cannot starve
1442            // the server runtime.
1443            Err(_) => self.dispatch_blocking(request),
1444        }
1445    }
1446}
1447
1448impl WorkerActivityDispatcher {
1449    /// Dispatch the activity and block the calling thread until the worker
1450    /// responds, the worker is declared lost, or the server drains (see the
1451    /// module docs for the exhaustive termination list).
1452    ///
1453    /// The request carries the *real* workflow and activity ids the engine
1454    /// recorded in history, so the worker logs, the pending-completion key,
1455    /// and the heartbeat tracker all correlate directly against the event
1456    /// store. `config` is forwarded by the engine seam but not yet consumed
1457    /// here (the retry executor that reads it is unbuilt).
1458    ///
1459    /// Must never run while the calling thread still owns a tokio scheduler
1460    /// core: the response can only arrive after the runtime's stream
1461    /// forwarder flushes the queued [`WorkerMessage::ActivityTask`] to the
1462    /// worker, so the thread blocking here must not be the one responsible
1463    /// for polling that forwarder. [`ActivityDispatcher::dispatch`] enforces
1464    /// this with `tokio::task::block_in_place`.
1465    fn dispatch_blocking(&self, request: ActivityDispatch) -> Result<String, String> {
1466        let ActivityDispatch {
1467            namespace,
1468            task_queue,
1469            // OPTIONAL within-pool node affinity (NODE-4): `Some(n)` pins this
1470            // dispatch to workers advertising node `n` (require semantics);
1471            // `None` is unpinned and reaches any worker in the pool.
1472            node,
1473            workflow_id,
1474            run_id,
1475            activity_id,
1476            name,
1477            input,
1478            config: _,
1479            attempt,
1480            labels,
1481            // R5 advisory class: engine-side only. The wire carries no such
1482            // field and a worker's behaviour is identical either way — the
1483            // class governs how the ENGINE treats the failure, never how the
1484            // work is done.
1485            advisory: _,
1486        } = request;
1487        let started_at = Instant::now();
1488        self.ensure_accepting(&namespace, &name, &workflow_id, &activity_id, None)?;
1489        let address = ServiceAddress {
1490            namespace: namespace.clone(),
1491            task_queue: task_queue.clone(),
1492            activity_type: name.clone(),
1493            node: node.clone(),
1494        };
1495        let worker = self.select_worker_or_wait(&address, &workflow_id, &activity_id)?;
1496        let worker_id = worker.id();
1497        let span = info_span!(
1498            "activity_dispatch",
1499            operation = "activity_dispatch",
1500            namespace = %namespace,
1501            task_queue = %task_queue,
1502            node = node.as_deref(),
1503            workflow_id = %workflow_id,
1504            activity_id = %activity_id,
1505            activity_type = %name,
1506            worker_id = ?worker_id,
1507        );
1508        let _span_guard = span.enter();
1509        self.ensure_accepting(
1510            &namespace,
1511            &name,
1512            &workflow_id,
1513            &activity_id,
1514            Some(worker_id),
1515        )?;
1516
1517        let (completion_token, rx) = self
1518            .pending
1519            .insert(workflow_id.clone(), &run_id, activity_id.clone(), attempt)
1520            .map_err(|error| error.to_string())?;
1521        let task = activity_task(
1522            &name,
1523            &input,
1524            (&workflow_id, &run_id, &activity_id),
1525            attempt,
1526            labels,
1527            &completion_token,
1528        );
1529        if let Err(error) = self.track_worker_task(
1530            worker_id,
1531            &name,
1532            &workflow_id,
1533            &activity_id,
1534            attempt,
1535            completion_token.clone(),
1536        ) {
1537            self.cleanup_activity(worker_id, &workflow_id, &activity_id, &completion_token);
1538            return Err(error);
1539        }
1540        self.send_activity_task(
1541            &worker,
1542            task,
1543            &address,
1544            &workflow_id,
1545            &activity_id,
1546            &completion_token,
1547        )?;
1548        let context = ActivityDispatchContext {
1549            namespace: &namespace,
1550            activity_type: &name,
1551            worker_id,
1552            workflow_id: &workflow_id,
1553            activity_id: &activity_id,
1554            completion_token,
1555            started_at,
1556        };
1557        self.await_activity_result(&context, &rx)
1558    }
1559}
1560
1561/// Body of one liminal reply-router thread (see
1562/// [`WorkerActivityDispatcher::spawn_liminal_reply_router`] for the contract).
1563///
1564/// Resolves by the key THIS push dispatched (the awaiter is already
1565/// correlation-scoped to it), never by the reply's echoed ids: a buggy echo
1566/// must not cross executions.
1567#[cfg(feature = "liminal-transport")]
1568fn route_liminal_reply(
1569    pending: &PendingActivities,
1570    heartbeat_tracker: &HeartbeatTracker,
1571    drain_state: &DrainState,
1572    awaiter: &liminal_server::server::connection::PushReplyAwaiter,
1573    execution: (WorkerId, &WorkflowId, &ActivityId, &CompletionToken),
1574) {
1575    let (worker_id, workflow_id, activity_id, current_token) = execution;
1576    // The wait re-arms only while this dispatch is still tracked in-flight, so
1577    // a dispatch resolved elsewhere (expiry sweep, shutdown drain, cleanup)
1578    // releases this thread within one reply poll instead of parking it for the
1579    // remaining life of the worker's connection.
1580    let waited = super::liminal_transport::receive_bridge_reply(awaiter, || {
1581        heartbeat_tracker
1582            .is_tracked(worker_id, workflow_id, activity_id)
1583            .unwrap_or(false)
1584    });
1585    // `synthesized` marks a failure this router FABRICATED (disconnect or
1586    // receive fault) as opposed to a real worker reply: only fabricated
1587    // failures are gated on the tracker below.
1588    let (run_id, submitted_token, outcome, synthesized) = match waited {
1589        Ok(Some(response)) => {
1590            let submitted_token = match CompletionToken::from_wire(
1591                workflow_id,
1592                activity_id,
1593                response.completion_token,
1594            ) {
1595                Ok(token) => token,
1596                Err(error) => {
1597                    tracing::warn!(
1598                        worker_id = ?worker_id,
1599                        workflow_id = %workflow_id,
1600                        activity_id = %activity_id,
1601                        %error,
1602                        "liminal activity completion omitted its generation proof"
1603                    );
1604                    return;
1605                }
1606            };
1607            (response.run_id, submitted_token, response.outcome, false)
1608        }
1609        Ok(None) => {
1610            tracing::debug!(
1611                worker_id = ?worker_id,
1612                workflow_id = %workflow_id,
1613                activity_id = %activity_id,
1614                "liminal dispatch resolved by another path; abandoning reply wait"
1615            );
1616            return;
1617        }
1618        Err(error) if error.is_worker_connection_lost() => (
1619            None,
1620            current_token.clone(),
1621            // TRANSPORT domain (the worker's connection closed before it
1622            // replied): classified through the shared ledger, never as an
1623            // action failure.
1624            Err(pending.classify_worker_loss(workflow_id, activity_id, worker_id)),
1625            true,
1626        ),
1627        Err(error) => (
1628            None,
1629            current_token.clone(),
1630            Err(format!("retryable:worker liminal reply failed: {error}")),
1631            true,
1632        ),
1633    };
1634    // The gRPC sweeps fail only still-tracked tasks (`remove_worker_tasks`);
1635    // this is the same structural gate: `complete_task` reports whether THIS
1636    // call retired the tracked entry. A poisoned tracker fails open (deliver)
1637    // so the blocked dispatch thread is never left hanging on a broken lock.
1638    let after_accept = || {
1639        if synthesized {
1640            let was_tracked = clear_completed_task_tracking(
1641                heartbeat_tracker,
1642                worker_id,
1643                workflow_id,
1644                activity_id,
1645            );
1646            if !was_tracked {
1647                // Another path already resolved this dispatch (and notified drain):
1648                // injecting the fabricated lost-worker failure now could reach a retry.
1649                tracing::debug!(
1650                    worker_id = ?worker_id,
1651                    workflow_id = %workflow_id,
1652                    activity_id = %activity_id,
1653                    "liminal dispatch already resolved; dropping synthesized lost-worker failure"
1654                );
1655                return Ok(false);
1656            }
1657        } else {
1658            // A real worker result: the tracking clear is bookkeeping, never
1659            // a gate. Fail open — withholding a received result over a
1660            // poisoned tracker would leave the dispatch waiter parked on an
1661            // unbounded receive with no recovery path, because the lost-worker
1662            // sweep reads the same poisoned lock.
1663            let _ = clear_completed_task_tracking(
1664                heartbeat_tracker,
1665                worker_id,
1666                workflow_id,
1667                activity_id,
1668            );
1669        }
1670        drain_state.notify_activity_drained();
1671        Ok(true)
1672    };
1673    if let Err(error) = pending.complete_fenced_after_accept(
1674        workflow_id,
1675        activity_id,
1676        run_id.as_ref(),
1677        &submitted_token,
1678        outcome,
1679        after_accept,
1680    ) {
1681        tracing::warn!(
1682            worker_id = ?worker_id,
1683            workflow_id = %workflow_id,
1684            activity_id = %activity_id,
1685            %error,
1686            "liminal activity completion handoff rejected"
1687        );
1688    }
1689}
1690
1691/// Clears a completed activity's in-flight tracking entry, reporting whether
1692/// THIS call retired it. Fails open on a poisoned tracker — logged loudly and
1693/// treated as "was tracked" — because the clear is bookkeeping, never a gate:
1694/// withholding a completion over the poisoned lock would strand the dispatch
1695/// waiter on an unbounded receive, and the lost-worker sweep cannot recover it
1696/// (it reads the same lock).
1697pub(crate) fn clear_completed_task_tracking(
1698    heartbeat_tracker: &HeartbeatTracker,
1699    worker_id: WorkerId,
1700    workflow_id: &WorkflowId,
1701    activity_id: &ActivityId,
1702) -> bool {
1703    heartbeat_tracker
1704        .complete_task(worker_id, workflow_id, activity_id)
1705        .unwrap_or_else(|error| {
1706            tracing::error!(
1707                worker_id = ?worker_id,
1708                workflow_id = %workflow_id,
1709                activity_id = %activity_id,
1710                %error,
1711                "failed to clear in-flight tracking for a completed activity; delivering anyway — a poisoned tracker must not withhold a result"
1712            );
1713            true
1714        })
1715}
1716
1717struct ActivityDispatchContext<'a> {
1718    namespace: &'a str,
1719    activity_type: &'a str,
1720    worker_id: WorkerId,
1721    workflow_id: &'a WorkflowId,
1722    activity_id: &'a ActivityId,
1723    completion_token: CompletionToken,
1724    started_at: Instant,
1725}
1726
1727fn activity_task(
1728    activity_type: &str,
1729    input: &str,
1730    execution: (&WorkflowId, &RunId, &ActivityId),
1731    attempt: u32,
1732    labels: BTreeMap<String, String>,
1733    completion_token: &CompletionToken,
1734) -> ProtoActivityTask {
1735    let (workflow_id, run_id, activity_id) = execution;
1736    ProtoActivityTask {
1737        workflow_id: Some(ProtoWorkflowId::from(workflow_id.clone())),
1738        activity_id: Some(ProtoActivityId::from(activity_id.clone())),
1739        activity_type: activity_type.to_owned(),
1740        input: Some(ProtoPayload {
1741            content_type: String::from("application/json"),
1742            bytes: input.as_bytes().to_vec(),
1743        }),
1744        attempt,
1745        labels: labels.into_iter().collect(),
1746        run_id: Some(run_id.clone().into()),
1747        completion_token: completion_token.as_str().to_owned(),
1748        idempotency_key: idempotency_key(workflow_id, run_id, activity_id),
1749    }
1750}
1751
1752fn log_activity_completion(context: &ActivityDispatchContext<'_>, succeeded: bool) {
1753    let duration_ms = duration_ms(context.started_at.elapsed());
1754    tracing::info!(
1755        operation = "activity_complete",
1756        namespace = context.namespace,
1757        workflow_id = %context.workflow_id,
1758        activity_id = %context.activity_id,
1759        activity_type = context.activity_type,
1760        worker_id = ?context.worker_id,
1761        duration_ms,
1762        outcome = if succeeded { "succeeded" } else { "failed" },
1763        "activity completed"
1764    );
1765}
1766
1767fn duration_ms(duration: Duration) -> u64 {
1768    u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
1769}
1770
1771fn log_worker_error(
1772    error_type: &'static str,
1773    namespace: &str,
1774    activity_type: &str,
1775    workflow_id: &WorkflowId,
1776    activity_id: &ActivityId,
1777    worker_id: Option<super::registry::WorkerId>,
1778    reason: &str,
1779) {
1780    tracing::error!(
1781        operation = "activity_dispatch",
1782        namespace,
1783        workflow_id = %workflow_id,
1784        activity_id = %activity_id,
1785        activity_type,
1786        worker_id = ?worker_id,
1787        error_type,
1788        reason,
1789        "worker interaction failed"
1790    );
1791}
1792
1793#[cfg(test)]
1794mod tests {
1795    use std::sync::{
1796        Mutex,
1797        atomic::{AtomicBool, Ordering},
1798    };
1799
1800    use aion_core::{ActivityError, ActivityErrorKind, ContentType, Payload};
1801
1802    use super::*;
1803
1804    fn activity_id(pos: u64) -> ActivityId {
1805        ActivityId::from_sequence_position(pos)
1806    }
1807
1808    /// One delivery of attempt 1 for `(workflow_id, id)` in its own run — what
1809    /// every single-dispatch unit test below is. The run and the attempt are the
1810    /// two discriminators the fence separates a REDELIVERY from a retry and from
1811    /// a new execution generation with, so they are stated here rather than
1812    /// implied.
1813    fn insert_attempt_one(
1814        pending: &PendingActivities,
1815        workflow_id: &WorkflowId,
1816        id: &ActivityId,
1817    ) -> Result<(CompletionToken, SyncReceiver), ServerError> {
1818        pending.insert(workflow_id.clone(), &RunId::new_v4(), id.clone(), 1)
1819    }
1820
1821    #[test]
1822    fn pending_insert_and_complete_delivers_result() -> Result<(), ServerError> {
1823        let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
1824        let workflow_id = WorkflowId::new_v4();
1825        let id = activity_id(1);
1826        let rx = insert_attempt_one(&pending, &workflow_id, &id)?.1;
1827
1828        assert!(pending.complete(&workflow_id, &id, None, Ok("done".to_owned())));
1829        assert_eq!(
1830            rx.recv_timeout(Duration::from_millis(50)),
1831            Ok(Ok("done".to_owned()))
1832        );
1833        Ok(())
1834    }
1835
1836    /// #246 / #254 red: retrying an activity while the same worker still holds
1837    /// the prior attempt must never replace its responder silently. Today's
1838    /// attempt-free `DashMap::insert` drops the first sender, producing the
1839    /// terminal "activity response channel dropped" specimen.
1840    #[test]
1841    fn same_execution_collision_is_a_named_typed_refusal() -> Result<(), ServerError> {
1842        let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
1843        let workflow_id = WorkflowId::new_v4();
1844        let id = activity_id(2);
1845        // The SAME execution redelivered: same run, same attempt. The fence
1846        // treats that as a sibling authorization; the pending map still refuses
1847        // to replace a live responder.
1848        let run_id = RunId::new_v4();
1849        let _first = pending.insert(workflow_id.clone(), &run_id, id.clone(), 1)?;
1850
1851        let error = pending
1852            .insert(workflow_id, &run_id, id, 1)
1853            .err()
1854            .ok_or_else(|| {
1855                ServerError::worker_dispatch(
1856                    "test",
1857                    "same-worker-collision",
1858                    "a live responder collision was silently replaced",
1859                )
1860            })?;
1861
1862        assert!(
1863            error.to_string().contains("pending activity collision"),
1864            "{error}"
1865        );
1866        Ok(())
1867    }
1868
1869    #[test]
1870    fn pending_complete_unknown_returns_false() {
1871        let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
1872        assert!(!pending.complete(
1873            &WorkflowId::new_v4(),
1874            &activity_id(99),
1875            None,
1876            Ok("orphan".to_owned())
1877        ));
1878    }
1879
1880    #[derive(Default)]
1881    struct RecordingOutboxCallback {
1882        completions: Mutex<Vec<(WorkflowId, ActivityId, String)>>,
1883        failures: Mutex<Vec<(WorkflowId, ActivityId, String)>>,
1884        live: bool,
1885    }
1886
1887    impl OutboxDeliveryCallback for RecordingOutboxCallback {
1888        fn deliver_completion(
1889            &self,
1890            workflow_id: &WorkflowId,
1891            activity_id: &ActivityId,
1892            run_id: Option<&RunId>,
1893            result: String,
1894        ) -> Result<bool, ServerError> {
1895            let _ = run_id;
1896            self.completions
1897                .lock()
1898                .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
1899                .push((workflow_id.clone(), activity_id.clone(), result));
1900            Ok(self.live)
1901        }
1902
1903        fn deliver_failure(
1904            &self,
1905            workflow_id: &WorkflowId,
1906            activity_id: &ActivityId,
1907            run_id: Option<&RunId>,
1908            reason: String,
1909        ) -> Result<bool, ServerError> {
1910            let _ = run_id;
1911            self.failures
1912                .lock()
1913                .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
1914                .push((workflow_id.clone(), activity_id.clone(), reason));
1915            Ok(self.live)
1916        }
1917    }
1918
1919    #[test]
1920    fn unmatched_completion_routes_to_outbox_callback_when_installed() -> Result<(), ServerError> {
1921        let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
1922        let callback = Arc::new(RecordingOutboxCallback {
1923            live: true,
1924            ..RecordingOutboxCallback::default()
1925        });
1926        // Install on one clone; the wiring must be visible to every clone.
1927        pending.clone().set_outbox_delivery(callback.clone());
1928
1929        let workflow_id = WorkflowId::new_v4();
1930        let id = activity_id(7);
1931
1932        // No pending entry: the completion is unmatched and must route to the
1933        // callback rather than being dropped. A live workflow reports true.
1934        assert!(pending.complete(&workflow_id, &id, None, Ok("done".to_owned())));
1935        let completions = callback
1936            .completions
1937            .lock()
1938            .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?;
1939        assert_eq!(completions.len(), 1);
1940        assert_eq!(completions[0].0, workflow_id);
1941        assert_eq!(completions[0].1, id);
1942        assert_eq!(completions[0].2, "done");
1943        Ok(())
1944    }
1945
1946    #[test]
1947    fn unmatched_failure_routes_to_outbox_callback_and_not_live_reports_false()
1948    -> Result<(), ServerError> {
1949        let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
1950        // live = false models the expected stale-completion case.
1951        let callback = Arc::new(RecordingOutboxCallback::default());
1952        pending.set_outbox_delivery(callback.clone());
1953
1954        let workflow_id = WorkflowId::new_v4();
1955        let id = activity_id(8);
1956
1957        assert!(!pending.complete(&workflow_id, &id, None, Err("retryable:boom".to_owned())));
1958        let failures = callback
1959            .failures
1960            .lock()
1961            .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?;
1962        assert_eq!(failures.len(), 1);
1963        assert_eq!(failures[0].2, "retryable:boom");
1964        Ok(())
1965    }
1966
1967    #[test]
1968    fn unmatched_completion_is_silent_drop_when_no_callback_installed() {
1969        // Flag-off byte-identical behaviour: no callback, unmatched returns
1970        // false (silent drop) exactly as before.
1971        let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
1972        assert!(!pending.complete(
1973            &WorkflowId::new_v4(),
1974            &activity_id(9),
1975            None,
1976            Ok("x".to_owned())
1977        ));
1978    }
1979
1980    #[test]
1981    fn matched_completion_never_reaches_outbox_callback() -> Result<(), ServerError> {
1982        let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
1983        let callback = Arc::new(RecordingOutboxCallback {
1984            live: true,
1985            ..RecordingOutboxCallback::default()
1986        });
1987        pending.set_outbox_delivery(callback.clone());
1988
1989        let workflow_id = WorkflowId::new_v4();
1990        let id = activity_id(10);
1991        let rx = insert_attempt_one(&pending, &workflow_id, &id)?.1;
1992
1993        assert!(pending.complete(&workflow_id, &id, None, Ok("matched".to_owned())));
1994        assert_eq!(
1995            rx.recv_timeout(Duration::from_millis(50)),
1996            Ok(Ok("matched".to_owned()))
1997        );
1998        assert!(
1999            callback
2000                .completions
2001                .lock()
2002                .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
2003                .is_empty(),
2004            "a matched completion must deliver to its waiter, not the outbox callback"
2005        );
2006        Ok(())
2007    }
2008
2009    /// #207: parking resolves the matched waiter with the ephemeral parked
2010    /// sentinel — the exact string the engine's retry loop classifies as
2011    /// `Parked` — and nothing else.
2012    #[test]
2013    fn park_activity_resolves_matched_waiter_with_the_parked_sentinel() -> Result<(), ServerError> {
2014        let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2015        let workflow_id = WorkflowId::new_v4();
2016        let id = activity_id(11);
2017        let rx = insert_attempt_one(&pending, &workflow_id, &id)?.1;
2018
2019        pending.park_activity(&workflow_id, &id)?;
2020        let result = rx
2021            .recv_timeout(Duration::from_millis(50))
2022            .map_err(|e| ServerError::worker_dispatch("", "", format!("channel: {e}")))?;
2023        assert_eq!(result, Err(aion::PARKED_ACTIVITY_REASON.to_owned()));
2024        Ok(())
2025    }
2026
2027    /// #207: an unmatched park is a no-op and is NEVER routed to the outbox
2028    /// delivery callback — a park is not a failure and must never reach a
2029    /// workflow.
2030    #[test]
2031    fn unmatched_park_is_a_noop_and_never_reaches_the_outbox_callback() -> Result<(), ServerError> {
2032        let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2033        let callback = Arc::new(RecordingOutboxCallback {
2034            live: true,
2035            ..RecordingOutboxCallback::default()
2036        });
2037        pending.set_outbox_delivery(callback.clone());
2038
2039        pending.park_activity(&WorkflowId::new_v4(), &activity_id(12))?;
2040
2041        assert!(
2042            callback
2043                .failures
2044                .lock()
2045                .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
2046                .is_empty(),
2047            "a park must never be delivered as an outbox failure"
2048        );
2049        assert!(
2050            callback
2051                .completions
2052                .lock()
2053                .map_err(|_| ServerError::lock_poisoned("recording outbox callback"))?
2054                .is_empty(),
2055            "a park must never be delivered as an outbox completion"
2056        );
2057        Ok(())
2058    }
2059
2060    #[test]
2061    fn completion_sink_routes_success() -> Result<(), ServerError> {
2062        let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2063        let workflow_id = WorkflowId::new_v4();
2064        let id = activity_id(2);
2065        let (completion_token, rx) = insert_attempt_one(&pending, &workflow_id, &id)?;
2066        let payload = Payload::new(ContentType::Json, br#"{"greeting":"hi"}"#.to_vec());
2067
2068        pending.complete_activity(ActivityCompletion {
2069            workflow_id,
2070            activity_id: id,
2071            run_id: None,
2072            completion_token,
2073            outcome: ActivityCompletionOutcome::Succeeded(payload),
2074        })?;
2075
2076        let result = rx
2077            .recv_timeout(Duration::from_millis(50))
2078            .map_err(|e| ServerError::worker_dispatch("", "", format!("channel: {e}")))?;
2079        assert_eq!(result, Ok(r#"{"greeting":"hi"}"#.to_owned()));
2080        Ok(())
2081    }
2082
2083    #[test]
2084    fn accepted_hook_runs_before_the_result_is_observable() -> Result<(), ServerError> {
2085        let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2086        let workflow_id = WorkflowId::new_v4();
2087        let id = activity_id(13);
2088        let (completion_token, rx) = insert_attempt_one(&pending, &workflow_id, &id)?;
2089        let accepted = AtomicBool::new(false);
2090
2091        pending.complete_activity_after_accept(
2092            ActivityCompletion {
2093                workflow_id,
2094                activity_id: id,
2095                run_id: None,
2096                completion_token,
2097                outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2098                    ContentType::Json,
2099                    br#""ordered""#.to_vec(),
2100                )),
2101            },
2102            || {
2103                accepted.store(true, Ordering::SeqCst);
2104                Ok(())
2105            },
2106        )?;
2107
2108        let result = rx
2109            .recv_timeout(Duration::from_millis(50))
2110            .map_err(|error| ServerError::worker_dispatch("", "", format!("channel: {error}")))?;
2111        assert_eq!(result, Ok(r#""ordered""#.to_owned()));
2112        assert!(
2113            accepted.load(Ordering::SeqCst),
2114            "the accepted-path cleanup hook must finish before the waiter can observe the result"
2115        );
2116        Ok(())
2117    }
2118
2119    #[test]
2120    fn a_non_publishing_settlement_restores_the_generation_for_the_true_resolver()
2121    -> Result<(), ServerError> {
2122        let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2123        let workflow_id = WorkflowId::new_v4();
2124        let id = activity_id(16);
2125        let (completion_token, rx) = insert_attempt_one(&pending, &workflow_id, &id)?;
2126
2127        // The synthesized-failure path finding its dispatch already resolved
2128        // by another path: settlement reports "do not publish". The generation
2129        // it consumed on accept must come back, because the OTHER path (the
2130        // lost-worker sweep) presents the same token next.
2131        let settled = pending.complete_fenced_after_accept(
2132            &workflow_id,
2133            &id,
2134            None,
2135            &completion_token,
2136            Err("retryable:worker lost".to_owned()),
2137            || Ok(false),
2138        )?;
2139        assert!(!settled, "a non-publishing settlement must report false");
2140        assert!(
2141            rx.try_recv().is_err(),
2142            "a non-publishing settlement must not publish"
2143        );
2144
2145        // The true resolver arrives with the same token: it must be accepted,
2146        // not refused with NoCurrentGeneration.
2147        let published = pending.complete_fenced_after_accept(
2148            &workflow_id,
2149            &id,
2150            None,
2151            &completion_token,
2152            Err("retryable:worker lost".to_owned()),
2153            || Ok(true),
2154        )?;
2155        assert!(
2156            published,
2157            "the true resolver must be accepted after a non-publishing settlement"
2158        );
2159        assert_eq!(
2160            rx.recv_timeout(Duration::from_millis(50))
2161                .map_err(|error| ServerError::worker_dispatch(
2162                    "",
2163                    "",
2164                    format!("channel: {error}")
2165                ))?,
2166            Err("retryable:worker lost".to_owned())
2167        );
2168        Ok(())
2169    }
2170
2171    #[test]
2172    fn a_poisoned_tracker_fails_open_and_never_withholds_a_completion() {
2173        let tracker = HeartbeatTracker::new(TEST_HEARTBEAT_WINDOW);
2174        tracker.poison_for_tests();
2175        let workflow_id = WorkflowId::new_v4();
2176        let id = activity_id(17);
2177        let was_tracked =
2178            clear_completed_task_tracking(&tracker, WorkerId::from_value(1), &workflow_id, &id);
2179        assert!(
2180            was_tracked,
2181            "a poisoned tracker must fail open (deliver), never withhold"
2182        );
2183    }
2184
2185    #[test]
2186    fn accepted_hook_failure_blocks_publication_and_restores_generation() -> Result<(), ServerError>
2187    {
2188        let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2189        let workflow_id = WorkflowId::new_v4();
2190        let id = activity_id(15);
2191        let (completion_token, rx) = insert_attempt_one(&pending, &workflow_id, &id)?;
2192
2193        let rejected = pending.complete_activity_after_accept(
2194            ActivityCompletion {
2195                workflow_id: workflow_id.clone(),
2196                activity_id: id.clone(),
2197                run_id: None,
2198                completion_token: completion_token.clone(),
2199                outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2200                    ContentType::Json,
2201                    br#""ordered""#.to_vec(),
2202                )),
2203            },
2204            || Err(ServerError::lock_poisoned("accepted-path settlement")),
2205        );
2206
2207        assert!(matches!(rejected, Err(ServerError::LockPoisoned { .. })));
2208        assert!(
2209            rx.try_recv().is_err(),
2210            "failed accepted-path settlement must not publish the completion"
2211        );
2212
2213        pending.complete_activity_after_accept(
2214            ActivityCompletion {
2215                workflow_id,
2216                activity_id: id,
2217                run_id: None,
2218                completion_token,
2219                outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2220                    ContentType::Json,
2221                    br#""ordered""#.to_vec(),
2222                )),
2223            },
2224            || Ok(()),
2225        )?;
2226        assert_eq!(
2227            rx.recv_timeout(Duration::from_millis(50))
2228                .map_err(|error| ServerError::worker_dispatch(
2229                    "",
2230                    "",
2231                    format!("channel: {error}")
2232                ))?,
2233            Ok(r#""ordered""#.to_owned())
2234        );
2235        Ok(())
2236    }
2237
2238    #[test]
2239    fn refused_generation_never_runs_the_accepted_hook() -> Result<(), ServerError> {
2240        let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2241        let workflow_id = WorkflowId::new_v4();
2242        let id = activity_id(14);
2243        let (_completion_token, rx) = insert_attempt_one(&pending, &workflow_id, &id)?;
2244        let accepted = AtomicBool::new(false);
2245
2246        let rejected = pending.complete_activity_after_accept(
2247            ActivityCompletion {
2248                workflow_id,
2249                activity_id: id,
2250                run_id: None,
2251                completion_token: CompletionToken::for_test(),
2252                outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2253                    ContentType::Json,
2254                    br#""stale""#.to_vec(),
2255                )),
2256            },
2257            || {
2258                accepted.store(true, Ordering::SeqCst);
2259                Ok(())
2260            },
2261        );
2262
2263        assert!(matches!(
2264            rejected,
2265            Err(ServerError::ActivityCompletionRejected { .. })
2266        ));
2267        assert!(
2268            !accepted.load(Ordering::SeqCst),
2269            "a refused generation must preserve liveness by skipping accepted-path cleanup"
2270        );
2271        assert!(
2272            rx.try_recv().is_err(),
2273            "a refused generation must leave the current waiter unresolved"
2274        );
2275        Ok(())
2276    }
2277
2278    #[test]
2279    fn malformed_payload_does_not_consume_the_current_generation() -> Result<(), ServerError> {
2280        let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2281        let workflow_id = WorkflowId::new_v4();
2282        let id = activity_id(12);
2283        let (completion_token, rx) = insert_attempt_one(&pending, &workflow_id, &id)?;
2284
2285        let malformed = pending.complete_activity(ActivityCompletion {
2286            workflow_id: workflow_id.clone(),
2287            activity_id: id.clone(),
2288            run_id: None,
2289            completion_token: completion_token.clone(),
2290            outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2291                ContentType::Json,
2292                vec![0xff],
2293            )),
2294        });
2295        assert!(matches!(malformed, Err(ServerError::WorkerDispatch { .. })));
2296        assert!(
2297            rx.try_recv().is_err(),
2298            "an invalid result must leave the waiter unresolved"
2299        );
2300
2301        pending.complete_activity(ActivityCompletion {
2302            workflow_id,
2303            activity_id: id,
2304            run_id: None,
2305            completion_token,
2306            outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2307                ContentType::Json,
2308                br#""valid""#.to_vec(),
2309            )),
2310        })?;
2311        let result = rx
2312            .recv_timeout(Duration::from_millis(50))
2313            .map_err(|error| ServerError::worker_dispatch("", "", format!("channel: {error}")))?;
2314        assert_eq!(result, Ok(r#""valid""#.to_owned()));
2315        Ok(())
2316    }
2317
2318    #[test]
2319    fn completion_sink_routes_retryable_error() -> Result<(), ServerError> {
2320        let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2321        let workflow_id = WorkflowId::new_v4();
2322        let id = activity_id(3);
2323        let (completion_token, rx) = insert_attempt_one(&pending, &workflow_id, &id)?;
2324
2325        pending.complete_activity(ActivityCompletion {
2326            workflow_id,
2327            activity_id: id,
2328            run_id: None,
2329            completion_token,
2330            outcome: ActivityCompletionOutcome::Failed(ActivityError {
2331                kind: ActivityErrorKind::Retryable,
2332                message: "temporary".to_owned(),
2333                details: None,
2334            }),
2335        })?;
2336
2337        let result = rx
2338            .recv_timeout(Duration::from_millis(50))
2339            .map_err(|e| ServerError::worker_dispatch("", "", format!("channel: {e}")))?;
2340        assert_eq!(result, Err("retryable:temporary".to_owned()));
2341        Ok(())
2342    }
2343
2344    #[test]
2345    fn completion_sink_routes_policy_refusal_with_its_own_prefix() -> Result<(), ServerError> {
2346        let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2347        let workflow_id = WorkflowId::new_v4();
2348        let id = activity_id(4);
2349        let (completion_token, rx) = insert_attempt_one(&pending, &workflow_id, &id)?;
2350
2351        pending.complete_activity(ActivityCompletion {
2352            workflow_id,
2353            activity_id: id,
2354            run_id: None,
2355            completion_token,
2356            outcome: ActivityCompletionOutcome::Failed(ActivityError {
2357                kind: ActivityErrorKind::PolicyRefused,
2358                message: "provider safety policy".to_owned(),
2359                details: None,
2360            }),
2361        })?;
2362
2363        let result = rx
2364            .recv_timeout(Duration::from_millis(50))
2365            .map_err(|error| ServerError::worker_dispatch("", "", format!("channel: {error}")))?;
2366        assert_eq!(
2367            result,
2368            Err("policy_refused:provider safety policy".to_owned())
2369        );
2370        Ok(())
2371    }
2372
2373    /// Regression test (#59, brief D12): pending tracking must be keyed by
2374    /// the full `(WorkflowId, ActivityId)` pair. The dispatcher fabricates
2375    /// activity ids from a process-local counter that resets on server
2376    /// restart, so a stale result re-reported from a worker's previous
2377    /// session carries the same bare `ActivityId` as a fresh post-restart
2378    /// dispatch. Under bare-`ActivityId` keying the stale result completed
2379    /// the wrong execution; with pair keying it is dropped and the genuine
2380    /// result still completes.
2381    #[test]
2382    fn stale_result_for_other_workflow_does_not_complete_pending_dispatch()
2383    -> Result<(), ServerError> {
2384        let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2385        let post_restart_workflow = WorkflowId::new_v4();
2386        let pre_restart_workflow = WorkflowId::new_v4();
2387        // Counter resets to the same sequence position after restart.
2388        let id = activity_id(1);
2389        let (completion_token, rx) = insert_attempt_one(&pending, &post_restart_workflow, &id)?;
2390
2391        // Stale pre-restart result: same activity id, different workflow.
2392        let rejected = pending.complete_activity(ActivityCompletion {
2393            workflow_id: pre_restart_workflow,
2394            activity_id: id.clone(),
2395            run_id: None,
2396            completion_token: CompletionToken::for_test(),
2397            outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2398                ContentType::Json,
2399                br#""stale""#.to_vec(),
2400            )),
2401        });
2402        assert!(matches!(
2403            rejected,
2404            Err(ServerError::ActivityCompletionRejected { .. })
2405        ));
2406        assert!(
2407            rx.try_recv().is_err(),
2408            "stale result for a different workflow must not complete this dispatch"
2409        );
2410
2411        // The genuine result for the pending execution still completes.
2412        pending.complete_activity(ActivityCompletion {
2413            workflow_id: post_restart_workflow,
2414            activity_id: id,
2415            run_id: None,
2416            completion_token,
2417            outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2418                ContentType::Json,
2419                br#""fresh""#.to_vec(),
2420            )),
2421        })?;
2422        let result = rx
2423            .recv_timeout(Duration::from_millis(50))
2424            .map_err(|e| ServerError::worker_dispatch("", "", format!("channel: {e}")))?;
2425        assert_eq!(result, Ok(r#""fresh""#.to_owned()));
2426        Ok(())
2427    }
2428
2429    /// Liveness tracker for dispatcher unit tests; the window only matters
2430    /// to expiry checks, which nothing in these tests drives.
2431    /// The heartbeat window these unit tests wire, named once so the
2432    /// transport-loss budget they run under is explicit rather than inherited.
2433    /// It is the same window [`test_tracker`] uses, so a sink and a tracker
2434    /// built side by side in a test agree exactly as they do in production.
2435    const TEST_HEARTBEAT_WINDOW: Duration = Duration::from_secs(5);
2436
2437    /// The budget a sink ends up with is DERIVED from the operator's window,
2438    /// and this is the arithmetic — not prose about it.
2439    ///
2440    /// Before #100 the wiring was `default().with_heartbeat_window(w)`: two
2441    /// calls where the second was omissible, and deleting it from all five
2442    /// production sites reddened nothing, because every test ran the zero-budget
2443    /// default. This pin and its sibling below are what make that deletion
2444    /// impossible to perform silently.
2445    #[test]
2446    fn a_sink_derives_its_transport_budget_from_the_operator_window() {
2447        let window = Duration::from_secs(7);
2448        let sink = PendingActivities::new(window);
2449
2450        assert_eq!(
2451            sink.transport_losses().budget(),
2452            window * super::super::transport_loss::TRANSPORT_LOSS_BUDGET_WINDOWS,
2453            "the transport-loss budget must be the operator's heartbeat window times \
2454             TRANSPORT_LOSS_BUDGET_WINDOWS; a sink carrying any other budget is one no \
2455             operator declared"
2456        );
2457    }
2458
2459    /// The wiring itself, which is where the defect actually lived: a dispatcher
2460    /// is handed a `HeartbeatTracker` and must not be able to carry a transport
2461    /// budget that disagrees with it.
2462    ///
2463    /// This asserts on the dispatcher a caller really builds, so it sits on the
2464    /// path a forgetful wiring would take rather than on the constructor in
2465    /// isolation.
2466    #[test]
2467    fn a_dispatcher_carries_the_transport_budget_of_the_tracker_it_was_given() {
2468        let window = Duration::from_secs(11);
2469        let registry = ConnectedWorkerRegistry::default();
2470        let dispatcher =
2471            WorkerActivityDispatcher::new(registry, "default", HeartbeatTracker::new(window));
2472
2473        assert_eq!(
2474            dispatcher.pending().transport_losses().budget(),
2475            window * super::super::transport_loss::TRANSPORT_LOSS_BUDGET_WINDOWS,
2476            "the dispatcher's transport budget must be derived from the heartbeat window it \
2477             was constructed with; a zero or mismatched budget here means a worker loss is \
2478             declared transport-exhausted after ONE re-dispatch instead of after the \
2479             operator's window"
2480        );
2481        assert_ne!(
2482            dispatcher.pending().transport_losses().budget(),
2483            Duration::ZERO,
2484            "a ZERO budget is the specific regression this pin exists for: it grants exactly \
2485             one re-dispatchable loss and then declares the infrastructure flapping"
2486        );
2487    }
2488
2489    fn test_tracker() -> HeartbeatTracker {
2490        HeartbeatTracker::new(TEST_HEARTBEAT_WINDOW)
2491    }
2492
2493    /// A `greet` dispatch request carrying real (test-synthesized) ids, the
2494    /// engine-seam shape `WorkerActivityDispatcher::dispatch` now consumes.
2495    fn greet_request() -> ActivityDispatch {
2496        ActivityDispatch {
2497            namespace: "default".to_owned(),
2498            task_queue: "default".to_owned(),
2499            node: None,
2500            workflow_id: WorkflowId::new_v4(),
2501            run_id: RunId::new_v4(),
2502            activity_id: ActivityId::from_sequence_position(0),
2503            name: "greet".to_owned(),
2504            input: "{}".to_owned(),
2505            config: "{}".to_owned(),
2506            attempt: 1,
2507            labels: std::collections::BTreeMap::new(),
2508            advisory: false,
2509        }
2510    }
2511
2512    #[test]
2513    fn dispatcher_fails_immediately_when_draining_without_workers() {
2514        let registry = ConnectedWorkerRegistry::default();
2515        let drain = DrainState::default();
2516        let dispatcher = WorkerActivityDispatcher::new(registry, "default", test_tracker())
2517            .with_drain_state(drain.clone());
2518
2519        let _ = drain.begin();
2520
2521        let result = dispatcher.dispatch(greet_request());
2522
2523        assert!(result.is_err());
2524        let err = result.err().unwrap_or_default();
2525        assert!(
2526            err.contains("drain"),
2527            "expected drain rejection, got: {err}"
2528        );
2529    }
2530
2531    /// Regression test for the production stall where every remote activity
2532    /// timed out: the engine invoked the sync `dispatch` from inside a
2533    /// spawned tokio task (`futures::future::lazy` polled on a runtime
2534    /// worker), and the woken stream-consumer task landed in that blocked
2535    /// worker's non-stealable LIFO slot, so the queued `ActivityTask` was
2536    /// only delivered when the then-extant 30s dispatch timeout fired (the
2537    /// dispatch wait is unbounded today; the stall would now be a hang).
2538    ///
2539    /// Mirrors the real wiring minus tonic: the real registry channel that
2540    /// the gRPC stream forwarder drains, a worker task awaiting that channel
2541    /// on the same runtime, completion through the production
2542    /// `ActivityCompletionSink`, and the sync dispatch invoked from a
2543    /// runtime worker task — the worst case the `block_in_place` guard in
2544    /// `dispatch` defends against (the engine itself now routes through
2545    /// `dispatch_async`, off the async workers).
2546    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2547    async fn dispatch_inside_runtime_task_delivers_promptly_and_round_trips()
2548    -> Result<(), Box<dyn std::error::Error>> {
2549        let registry = ConnectedWorkerRegistry::default();
2550        let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2551        let (worker_tx, mut worker_rx) = tokio::sync::mpsc::channel(32);
2552        let activity_types = [String::from("greet")];
2553        let registration = registry.register("default", activity_types.iter(), worker_tx)?;
2554
2555        let sink = pending.clone();
2556        let echo_worker = tokio::spawn(async move {
2557            let Some(WorkerMessage::ActivityTask(task)) = worker_rx.recv().await else {
2558                return Err("expected an activity task on the worker channel".to_owned());
2559            };
2560            let workflow_id = task
2561                .workflow_id
2562                .ok_or("task missing workflow id")
2563                .and_then(|id| WorkflowId::try_from(id).map_err(|_| "bad workflow id"))?;
2564            let activity_id = task
2565                .activity_id
2566                .map(ActivityId::from)
2567                .ok_or("task missing activity id")?;
2568            let completion_token =
2569                CompletionToken::from_wire(&workflow_id, &activity_id, task.completion_token)
2570                    .map_err(|error| error.to_string())?;
2571            sink.complete_activity(ActivityCompletion {
2572                workflow_id,
2573                activity_id,
2574                run_id: None,
2575                completion_token,
2576                outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2577                    ContentType::Json,
2578                    br#"{"greeting":"hello"}"#.to_vec(),
2579                )),
2580            })
2581            .map_err(|error| error.to_string())
2582        });
2583
2584        let dispatcher = Arc::new(
2585            WorkerActivityDispatcher::new(registry, "default", test_tracker())
2586                .with_pending(pending),
2587        );
2588        let started = Instant::now();
2589        // Invoke the sync dispatch inside the first poll of a spawned task:
2590        // the worst-case calling context for the `block_in_place` guard.
2591        let dispatch_task = tokio::spawn(futures::future::lazy(move |_| {
2592            dispatcher.dispatch(greet_request())
2593        }));
2594        let result = dispatch_task.await.map_err(|error| error.to_string())?;
2595        let elapsed = started.elapsed();
2596
2597        assert_eq!(result, Ok(r#"{"greeting":"hello"}"#.to_owned()));
2598        assert!(
2599            elapsed < Duration::from_secs(5),
2600            "dispatch round trip took {elapsed:?}; task delivery must not \
2601             depend on the blocked dispatch thread"
2602        );
2603        echo_worker.await.map_err(|error| error.to_string())??;
2604        registration.deregister()?;
2605        Ok(())
2606    }
2607
2608    /// A current-thread runtime cannot host the blocking wait (the stream
2609    /// forwarder would share its only executor thread), so dispatch must
2610    /// fail fast with a precise error instead of blocking forever.
2611    #[tokio::test]
2612    async fn dispatch_on_current_thread_runtime_fails_fast()
2613    -> Result<(), Box<dyn std::error::Error>> {
2614        let registry = ConnectedWorkerRegistry::default();
2615        let (worker_tx, _worker_rx) = tokio::sync::mpsc::channel(32);
2616        let activity_types = [String::from("greet")];
2617        let registration = registry.register("default", activity_types.iter(), worker_tx)?;
2618        let dispatcher = WorkerActivityDispatcher::new(registry, "default", test_tracker());
2619
2620        let started = Instant::now();
2621        let result = dispatcher.dispatch(greet_request());
2622        let elapsed = started.elapsed();
2623
2624        let err = result.err().ok_or("expected dispatch to fail")?;
2625        assert!(
2626            err.contains("multi-thread tokio runtime"),
2627            "unexpected error: {err}"
2628        );
2629        assert!(
2630            elapsed < Duration::from_secs(5),
2631            "fail-fast path took {elapsed:?}"
2632        );
2633        registration.deregister()?;
2634        Ok(())
2635    }
2636
2637    /// Bridge-level mirror of the e2e node-pin proof: two workers share the
2638    /// `(namespace, task_queue)` pool but advertise different nodes; an
2639    /// `ActivityDispatch` pinned to one node must reach ONLY the worker on that
2640    /// node through the live engine-seam `WorkerActivityDispatcher`. This is the
2641    /// regression guard for the bridge discarding the dispatch's `node`.
2642    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2643    async fn dispatch_pinned_to_node_reaches_only_that_node()
2644    -> Result<(), Box<dyn std::error::Error>> {
2645        let registry = ConnectedWorkerRegistry::default();
2646        let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2647        let activity_types = [String::from("greet")];
2648        let (n1_tx, mut n1_rx) = tokio::sync::mpsc::channel(32);
2649        let (n2_tx, mut n2_rx) = tokio::sync::mpsc::channel(32);
2650        // Register the DECOY (n2) FIRST so it owns the lowest worker id. The
2651        // bridge's `select_worker` orders candidates by worker id and starts at
2652        // the pool's rotation cursor, which is 0 for a pool nothing has selected
2653        // from yet — so the FIRST unfiltered selection here would name n2 (the
2654        // decoy), and a bridge that DISCARDED the node would route to it: the n1
2655        // echo would never fire and the round trip would time out. With the node
2656        // threaded through, selection is filtered to n1 before the cursor is
2657        // consulted, so the rotation cannot rescue a bridge that dropped it.
2658        let on_n2 = registry.register_namespaces(
2659            [String::from("default")],
2660            "default",
2661            Some(String::from("n2")),
2662            activity_types.iter(),
2663            n2_tx,
2664        )?;
2665        let on_n1 = registry.register_namespaces(
2666            [String::from("default")],
2667            "default",
2668            Some(String::from("n1")),
2669            activity_types.iter(),
2670            n1_tx,
2671        )?;
2672
2673        // Echo only on the n1 channel: the dispatch can only complete if the
2674        // task was routed to n1. If it leaked to n2, the n1 wait would stall and
2675        // the round trip below would time out instead.
2676        let sink = pending.clone();
2677        let echo_n1 = tokio::spawn(async move {
2678            let Some(WorkerMessage::ActivityTask(task)) = n1_rx.recv().await else {
2679                return Err("expected an activity task on the n1 worker channel".to_owned());
2680            };
2681            let workflow_id = task
2682                .workflow_id
2683                .ok_or("task missing workflow id")
2684                .and_then(|id| WorkflowId::try_from(id).map_err(|_| "bad workflow id"))?;
2685            let activity_id = task
2686                .activity_id
2687                .map(ActivityId::from)
2688                .ok_or("task missing activity id")?;
2689            let completion_token =
2690                CompletionToken::from_wire(&workflow_id, &activity_id, task.completion_token)
2691                    .map_err(|error| error.to_string())?;
2692            sink.complete_activity(ActivityCompletion {
2693                workflow_id,
2694                activity_id,
2695                run_id: None,
2696                completion_token,
2697                outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2698                    ContentType::Json,
2699                    br#"{"greeting":"hello"}"#.to_vec(),
2700                )),
2701            })
2702            .map_err(|error| error.to_string())
2703        });
2704
2705        let dispatcher = Arc::new(
2706            WorkerActivityDispatcher::new(registry.clone(), "default", test_tracker())
2707                .with_pending(pending),
2708        );
2709
2710        let pinned = ActivityDispatch {
2711            node: Some(String::from("n1")),
2712            ..greet_request()
2713        };
2714        let started = Instant::now();
2715        let result = tokio::spawn(futures::future::lazy(move |_| dispatcher.dispatch(pinned)))
2716            .await
2717            .map_err(|error| error.to_string())?;
2718        let elapsed = started.elapsed();
2719
2720        assert_eq!(result, Ok(r#"{"greeting":"hello"}"#.to_owned()));
2721        assert!(
2722            elapsed < Duration::from_secs(5),
2723            "pinned dispatch round trip took {elapsed:?}; the task must route to n1"
2724        );
2725        echo_n1.await.map_err(|error| error.to_string())??;
2726
2727        // The n2 worker (wrong node) must never have been handed the task.
2728        assert!(
2729            n2_rx.try_recv().is_err(),
2730            "node=Some(\"n1\") dispatch must not reach the n2 worker"
2731        );
2732
2733        on_n1.deregister()?;
2734        on_n2.deregister()?;
2735        Ok(())
2736    }
2737
2738    /// An unpinned (`node = None`) dispatch is byte-identical to today: it
2739    /// reaches a worker in the pool regardless of the worker's advertised node.
2740    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2741    async fn unpinned_dispatch_reaches_a_pooled_worker_regardless_of_node()
2742    -> Result<(), Box<dyn std::error::Error>> {
2743        let registry = ConnectedWorkerRegistry::default();
2744        let pending = PendingActivities::new(TEST_HEARTBEAT_WINDOW);
2745        let activity_types = [String::from("greet")];
2746        let (n1_tx, mut n1_rx) = tokio::sync::mpsc::channel(32);
2747        let on_n1 = registry.register_namespaces(
2748            [String::from("default")],
2749            "default",
2750            Some(String::from("n1")),
2751            activity_types.iter(),
2752            n1_tx,
2753        )?;
2754
2755        let sink = pending.clone();
2756        let echo = tokio::spawn(async move {
2757            let Some(WorkerMessage::ActivityTask(task)) = n1_rx.recv().await else {
2758                return Err("expected an activity task on the worker channel".to_owned());
2759            };
2760            let workflow_id = task
2761                .workflow_id
2762                .ok_or("task missing workflow id")
2763                .and_then(|id| WorkflowId::try_from(id).map_err(|_| "bad workflow id"))?;
2764            let activity_id = task
2765                .activity_id
2766                .map(ActivityId::from)
2767                .ok_or("task missing activity id")?;
2768            let completion_token =
2769                CompletionToken::from_wire(&workflow_id, &activity_id, task.completion_token)
2770                    .map_err(|error| error.to_string())?;
2771            sink.complete_activity(ActivityCompletion {
2772                workflow_id,
2773                activity_id,
2774                run_id: None,
2775                completion_token,
2776                outcome: ActivityCompletionOutcome::Succeeded(Payload::new(
2777                    ContentType::Json,
2778                    br#"{"greeting":"hello"}"#.to_vec(),
2779                )),
2780            })
2781            .map_err(|error| error.to_string())
2782        });
2783
2784        let dispatcher = Arc::new(
2785            WorkerActivityDispatcher::new(registry.clone(), "default", test_tracker())
2786                .with_pending(pending),
2787        );
2788
2789        // greet_request() carries node: None — the unpinned path.
2790        let result = tokio::spawn(futures::future::lazy(move |_| {
2791            dispatcher.dispatch(greet_request())
2792        }))
2793        .await
2794        .map_err(|error| error.to_string())?;
2795
2796        assert_eq!(result, Ok(r#"{"greeting":"hello"}"#.to_owned()));
2797        echo.await.map_err(|error| error.to_string())??;
2798        on_n1.deregister()?;
2799        Ok(())
2800    }
2801}