Skip to main content

aion_server/worker/
bridge.rs

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