Skip to main content

aion_server/worker/
bridge.rs

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