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