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