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