Skip to main content

aion_server/api/
worker_grpc.rs

1//! tonic `WorkerProtocol` service — bidirectional stream handler.
2
3use aion_proto::{
4    ProtoActivityResult, ProtoRegisterWorker,
5    generated::{
6        self,
7        worker_protocol_server::{WorkerProtocol, WorkerProtocolServer},
8    },
9};
10use tokio::sync::mpsc;
11use tokio_stream::wrappers::ReceiverStream;
12use tonic::{Request, Response, Status, Streaming};
13
14use crate::worker::PendingActivities;
15use crate::worker::dispatch::{ActivityCompletion, ActivityCompletionSink};
16use crate::worker::registry::{WorkerId, WorkerMessage};
17use crate::{CallerIdentity, ServerState};
18
19/// Cloneable tonic implementation for the worker bidirectional stream.
20#[derive(Clone)]
21pub struct WorkerGrpcService {
22    state: ServerState,
23}
24
25impl WorkerGrpcService {
26    /// Build a tonic worker service from shared server state.
27    #[must_use]
28    pub const fn new(state: ServerState) -> Self {
29        Self { state }
30    }
31}
32
33/// Construct the generated tonic server wrapper for the worker protocol.
34#[must_use]
35pub fn worker_service(state: ServerState) -> WorkerProtocolServer<WorkerGrpcService> {
36    WorkerProtocolServer::new(WorkerGrpcService::new(state))
37}
38
39#[tonic::async_trait]
40impl WorkerProtocol for WorkerGrpcService {
41    type StreamWorkerStream = ReceiverStream<Result<generated::ServerToWorker, Status>>;
42
43    async fn stream_worker(
44        &self,
45        request: Request<Streaming<generated::WorkerToServer>>,
46    ) -> Result<Response<Self::StreamWorkerStream>, Status> {
47        let metadata = request.metadata().clone();
48        let caller = worker_caller_from_metadata(&metadata, &self.state).await?;
49        let token_expires_at = token_expiration_from_metadata(&metadata, &self.state).await?;
50        let heartbeat_grace = self.state.runtime_config().worker.heartbeat_window;
51        let mut inbound = request.into_inner();
52
53        let first = inbound
54            .message()
55            .await?
56            .and_then(|msg| msg.message)
57            .ok_or_else(|| Status::invalid_argument("first message must be RegisterWorker"))?;
58
59        let register = match first {
60            generated::worker_to_server::Message::Register(r) => decode_register(r),
61            _ => {
62                return Err(Status::invalid_argument(
63                    "first message must be RegisterWorker",
64                ));
65            }
66        };
67
68        let (task_tx, task_rx) = mpsc::channel::<Result<generated::ServerToWorker, Status>>(32);
69        let (worker_tx, worker_rx) = mpsc::channel(32);
70
71        let registration = self
72            .state
73            .worker_registry()
74            .accept_registration(self.state.namespace_guard(), &caller, &register, worker_tx)
75            .await
76            .map_err(|error| status_from_server_error(&error))?;
77
78        let pending = self.state.pending_activities().clone();
79        let heartbeat = self.state.heartbeat_tracker().clone();
80        let drain = self.state.drain_state().clone();
81        let registry = self.state.worker_registry().clone();
82        let worker_id = registration
83            .worker_id()
84            .ok_or_else(|| Status::internal("worker registration missing id"))?;
85        // A worker serves a SET of namespaces; the ack echoes them joined in
86        // stable order purely for the worker's logs (the RegisterAck namespace
87        // field is informational, not a routing input).
88        let authorized_namespace = registration
89            .namespaces()
90            .filter(|namespaces| !namespaces.is_empty())
91            .ok_or_else(|| Status::internal("worker registration missing namespace"))?
92            .iter()
93            .cloned()
94            .collect::<Vec<_>>()
95            .join(",");
96
97        // RegisterAck ordering guarantee: the ack is enqueued on `task_tx`
98        // BEFORE the write forwarder that copies dispatched tasks onto the
99        // same channel is spawned, so no task frame can precede it on the
100        // wire. This is a structural ordering proof, not a timing hope.
101        task_tx
102            .try_send(Ok(register_ack_frame(
103                worker_id,
104                &authorized_namespace,
105                heartbeat_grace,
106            )))
107            .map_err(|_| Status::internal("worker response channel closed before RegisterAck"))?;
108
109        tokio::spawn(async move {
110            let write_handle = spawn_write_forwarder(worker_rx, task_tx.clone());
111
112            // Armed BEFORE the inbound loop runs: the sweep in its `Drop`
113            // fires on every exit from this task — clean stream end, stream
114            // error, token expiry, even a panic unwinding `process_inbound`.
115            // The unbounded dispatch wait depends on it.
116            let teardown = StreamTeardown {
117                worker_id,
118                heartbeat: &heartbeat,
119                registry: &registry,
120                pending: &pending,
121                drain: &drain,
122            };
123            let session = WorkerSession {
124                worker_id,
125                pending: &pending,
126                heartbeat: &heartbeat,
127                drain: &drain,
128                token_expires_at,
129                heartbeat_grace,
130                task_tx: task_tx.clone(),
131            };
132            if let Err(status) = process_inbound(inbound, session).await {
133                tracing::info!(
134                    worker_id = ?worker_id,
135                    %status,
136                    "worker stream closed with status"
137                );
138            }
139
140            write_handle.abort();
141            drop(task_tx);
142            drop(teardown);
143            // The teardown sweep already deregistered the stream; consuming
144            // the registration here is an idempotent no-op that still
145            // surfaces a poisoned-lock error loudly.
146            if let Err(error) = registration.deregister() {
147                tracing::error!(
148                    worker_id = ?worker_id,
149                    %error,
150                    "worker deregistration failed during stream teardown"
151                );
152            }
153        });
154
155        Ok(Response::new(ReceiverStream::new(task_rx)))
156    }
157}
158
159/// Spawn the write forwarder: it copies registry-delivered [`WorkerMessage`]s
160/// onto the worker's response stream, and — the #176 zombie fix — TERMINATES
161/// the RPC when the registry drops this worker's delivery sender.
162///
163/// `recv` returning `None` means every delivery sender is gone: the worker
164/// was DEREGISTERED while its stream stayed open (the heartbeat expiry sweep,
165/// or any future administrative removal). Silently deregistering would leave
166/// the worker a zombie — connected but unroutable, believing it is
167/// registered, its heartbeats rejected as "not in flight", never
168/// re-registering until its own stream happens to end. Ending the RPC with a
169/// retryable `Unavailable` status makes the worker OBSERVE the
170/// deregistration and re-register through its reconnect machinery. (On the
171/// normal teardown path this task is aborted before the registration is
172/// consumed, so the status is never sent to a worker that hung up; a send
173/// failure means the response stream's consumer is already gone, so there is
174/// no one left to signal.)
175fn spawn_write_forwarder(
176    mut worker_rx: mpsc::Receiver<WorkerMessage>,
177    task_tx: mpsc::Sender<Result<generated::ServerToWorker, Status>>,
178) -> tokio::task::JoinHandle<()> {
179    tokio::spawn(async move {
180        while let Some(message) = worker_rx.recv().await {
181            let msg = encode_server_to_worker(message);
182            if task_tx.send(Ok(msg)).await.is_err() {
183                return;
184            }
185        }
186        let _ = task_tx
187            .send(Err(Status::unavailable(
188                "worker was deregistered by the server (heartbeat window expired); \
189                 reconnect and re-register",
190            )))
191            .await;
192    })
193}
194
195/// Drop guard that sweeps a torn-down worker stream's in-flight activities:
196/// failed back to the engine mid-run, or parked for restart recovery under a
197/// graceful drain (#207) — see [`teardown_worker_stream`].
198///
199/// A guard rather than a call site so the sweep cannot be skipped by any
200/// exit from the stream task — including a panic unwinding the inbound
201/// loop, which would otherwise leave every dispatch blocked on that worker
202/// waiting forever.
203struct StreamTeardown<'a> {
204    worker_id: WorkerId,
205    heartbeat: &'a crate::worker::HeartbeatTracker,
206    registry: &'a crate::worker::ConnectedWorkerRegistry,
207    pending: &'a PendingActivities,
208    drain: &'a crate::shutdown::DrainState,
209}
210
211impl Drop for StreamTeardown<'_> {
212    fn drop(&mut self) {
213        teardown_worker_stream(
214            self.worker_id,
215            self.heartbeat,
216            self.registry,
217            self.pending,
218            self.drain,
219        );
220    }
221}
222
223/// Sweep a torn-down worker stream's in-flight activities.
224///
225/// The stream is the worker's liveness. When it ends mid-run — process death,
226/// network disconnect, expired token — every activity still assigned to this
227/// worker must be failed back through the completion sink as a retryable
228/// lost-worker error. The activity dispatch wait is unbounded by design (the
229/// engine imposes no activity timeout), so this sweep is what unblocks
230/// dispatches whose worker died mid-activity; the engine's retry policy then
231/// decides re-dispatch.
232///
233/// Under a graceful drain (#207) the stream ending is the EXPECTED worker
234/// response to the drain request, not a death: the worker's in-flight tasks
235/// are PARKED for restart recovery instead — nothing is recorded, nothing is
236/// delivered, the durable log keeps its dangling scheduled/started trail
237/// exactly as a kill -9 would, and post-restart replay re-dispatches it. The
238/// park-vs-fail branch keys on
239/// [`DrainState::is_draining`](crate::shutdown::DrainState::is_draining): a
240/// worker lost while the server is NOT draining still fails-and-retries
241/// byte-identically to before.
242fn teardown_worker_stream(
243    worker_id: WorkerId,
244    heartbeat: &crate::worker::HeartbeatTracker,
245    registry: &crate::worker::ConnectedWorkerRegistry,
246    pending: &PendingActivities,
247    drain: &crate::shutdown::DrainState,
248) {
249    if drain.is_draining() {
250        match heartbeat.park_disconnected_worker(worker_id, registry, pending) {
251            Ok(report) if report.tasks.is_empty() => {}
252            Ok(report) => {
253                tracing::info!(
254                    worker_id = ?worker_id,
255                    parked_tasks = report.tasks.len(),
256                    "worker stream ended during drain; in-flight activities \
257                     parked for restart recovery"
258                );
259            }
260            Err(error) => {
261                tracing::error!(
262                    worker_id = ?worker_id,
263                    %error,
264                    "failed to park draining worker's in-flight activities"
265                );
266            }
267        }
268    } else {
269        match heartbeat.fail_disconnected_worker(worker_id, registry, pending) {
270            Ok(report) if report.tasks.is_empty() => {}
271            Ok(report) => {
272                tracing::warn!(
273                    worker_id = ?worker_id,
274                    failed_tasks = report.tasks.len(),
275                    "worker disconnected with in-flight activities; \
276                     surfaced as retryable lost-worker failures"
277                );
278            }
279            Err(error) => {
280                tracing::error!(
281                    worker_id = ?worker_id,
282                    %error,
283                    "failed to sweep disconnected worker's in-flight activities"
284                );
285            }
286        }
287    }
288    // In-flight accounting may have just reached zero; wake any drain
289    // waiter so shutdown does not sit out its full timeout.
290    drain.notify_activity_drained();
291}
292
293struct WorkerSession<'a> {
294    worker_id: WorkerId,
295    pending: &'a PendingActivities,
296    heartbeat: &'a crate::worker::HeartbeatTracker,
297    drain: &'a crate::shutdown::DrainState,
298    token_expires_at: Option<u64>,
299    heartbeat_grace: std::time::Duration,
300    task_tx: mpsc::Sender<Result<generated::ServerToWorker, Status>>,
301}
302
303async fn process_inbound(
304    mut inbound: Streaming<generated::WorkerToServer>,
305    session: WorkerSession<'_>,
306) -> Result<(), Status> {
307    let mut expired_since: Option<std::time::Instant> = None;
308    while let Some(msg) = inbound.message().await? {
309        let Some(inner) = msg.message else {
310            continue;
311        };
312        match inner {
313            generated::worker_to_server::Message::Result(result) => {
314                let proto_result = decode_activity_result(result);
315                match ActivityCompletion::try_from(proto_result) {
316                    Ok(completion) => {
317                        let workflow_id = completion.workflow_id.clone();
318                        let activity_id = completion.activity_id.clone();
319                        if let Err(error) = session.heartbeat.complete_task(
320                            session.worker_id,
321                            &workflow_id,
322                            &activity_id,
323                        ) {
324                            // A poisoned liveness tracker would also break
325                            // the lost-worker sweep the unbounded dispatch
326                            // wait relies on — never swallow it.
327                            tracing::error!(
328                                worker_id = ?session.worker_id,
329                                workflow_id = %workflow_id,
330                                activity_id = %activity_id,
331                                %error,
332                                "failed to clear in-flight tracking for completed activity"
333                            );
334                        }
335                        session.drain.notify_activity_drained();
336                        if let Err(error) = session.pending.complete_activity(completion) {
337                            tracing::error!(
338                                worker_id = ?session.worker_id,
339                                workflow_id = %workflow_id,
340                                activity_id = %activity_id,
341                                %error,
342                                "activity completion handoff failed"
343                            );
344                        }
345                        // Ack every well-formed result frame — including
346                        // duplicates with no pending waiter; their re-report
347                        // obligation is equally discharged. `try_send`: a
348                        // worker that stopped draining its receive side must
349                        // not wedge the inbound loop; a dropped ack is
350                        // recovered by the next-session re-report.
351                        let ack = result_ack_frame(&workflow_id, &activity_id);
352                        if let Err(error) = session.task_tx.try_send(Ok(ack)) {
353                            tracing::warn!(
354                                worker_id = ?session.worker_id,
355                                workflow_id = %workflow_id,
356                                activity_id = %activity_id,
357                                %error,
358                                "result ack dropped: worker stream channel unavailable"
359                            );
360                        }
361                    }
362                    Err(error) => {
363                        // Malformed result: no ids to ack with. Loud, never
364                        // silent — the worker's entry will re-report and
365                        // re-fail visibly each session.
366                        tracing::error!(
367                            worker_id = ?session.worker_id,
368                            %error,
369                            "malformed activity result frame; no ack sent"
370                        );
371                    }
372                }
373            }
374            generated::worker_to_server::Message::Register(_) => {
375                tracing::warn!(
376                    worker_id = ?session.worker_id,
377                    "ignoring subsequent RegisterWorker message; \
378                     only the first registration is accepted per stream"
379                );
380            }
381            generated::worker_to_server::Message::Heartbeat(heartbeat_msg) => {
382                if let Err(error) = session.heartbeat.record_heartbeat(
383                    session.worker_id,
384                    decode_heartbeat(heartbeat_msg),
385                    std::time::Instant::now(),
386                ) {
387                    // Malformed frames and heartbeats for untracked tasks
388                    // are worker-side defects worth surfacing; a poisoned
389                    // tracker lock is a server-side corruption signal that
390                    // must never vanish silently.
391                    if matches!(error, crate::ServerError::LockPoisoned { .. }) {
392                        tracing::error!(
393                            worker_id = ?session.worker_id,
394                            %error,
395                            "heartbeat tracker lock poisoned; liveness state untrustworthy"
396                        );
397                    } else {
398                        tracing::warn!(
399                            worker_id = ?session.worker_id,
400                            %error,
401                            "worker heartbeat rejected"
402                        );
403                    }
404                }
405                if token_expired(session.token_expires_at) {
406                    let first_expired = *expired_since.get_or_insert_with(std::time::Instant::now);
407                    let _ = session
408                        .task_tx
409                        .send(Err(Status::unauthenticated(
410                            "worker token expired; re-authentication required",
411                        )))
412                        .await;
413                    if first_expired.elapsed() >= session.heartbeat_grace {
414                        return Err(Status::unauthenticated("worker token expired"));
415                    }
416                }
417            }
418        }
419    }
420    Ok(())
421}
422
423async fn worker_caller_from_metadata(
424    metadata: &tonic::metadata::MetadataMap,
425    state: &ServerState,
426) -> Result<CallerIdentity, Status> {
427    crate::api::grpc::caller_from_metadata(metadata, state).await
428}
429
430async fn token_expiration_from_metadata(
431    metadata: &tonic::metadata::MetadataMap,
432    state: &ServerState,
433) -> Result<Option<u64>, Status> {
434    if !state.runtime_config().auth.enabled {
435        return Ok(None);
436    }
437    #[cfg(feature = "auth")]
438    {
439        let bearer = metadata
440            .get("authorization")
441            .and_then(|value| value.to_str().ok())
442            .and_then(parse_bearer)
443            .ok_or_else(|| Status::unauthenticated("missing bearer token"))?;
444        let Some(cache) = state.jwks_cache() else {
445            return Err(Status::unauthenticated("invalid bearer token"));
446        };
447        return cache
448            .validate(&bearer)
449            .await
450            .map(|claims| Some(claims.expires_at()))
451            .map_err(|_error| Status::unauthenticated("invalid bearer token"));
452    }
453    #[cfg(not(feature = "auth"))]
454    {
455        let _ = metadata;
456        // Yield to preserve the async signature required by the auth-feature branch.
457        tokio::task::yield_now().await;
458        Ok(None)
459    }
460}
461
462#[cfg(feature = "auth")]
463fn parse_bearer(value: &str) -> Option<String> {
464    let token = value.strip_prefix("Bearer ")?.trim();
465    if token.is_empty() {
466        return None;
467    }
468    Some(token.to_owned())
469}
470
471fn token_expired(expires_at: Option<u64>) -> bool {
472    expires_at.is_some_and(|expires_at| {
473        #[cfg(feature = "auth")]
474        {
475            crate::auth::jwks::is_expired(expires_at)
476        }
477        #[cfg(not(feature = "auth"))]
478        {
479            let _ = expires_at;
480            false
481        }
482    })
483}
484
485fn status_from_server_error(error: &crate::ServerError) -> Status {
486    let wire = error.to_wire_error();
487    if wire.code == aion_proto::WireErrorCode::NamespaceDenied {
488        Status::permission_denied(wire.message)
489    } else {
490        Status::internal(wire.message)
491    }
492}
493
494/// Build the positive registration acknowledgement frame — the guaranteed
495/// first frame on every successful worker response stream.
496fn register_ack_frame(
497    worker_id: WorkerId,
498    namespace: &str,
499    heartbeat_window: std::time::Duration,
500) -> generated::ServerToWorker {
501    generated::ServerToWorker {
502        message: Some(generated::server_to_worker::Message::RegisterAck(
503            generated::RegisterAck {
504                worker_id: worker_id.value(),
505                namespace: namespace.to_owned(),
506                heartbeat_window_ms: u64::try_from(heartbeat_window.as_millis())
507                    .unwrap_or(u64::MAX),
508            },
509        )),
510    }
511}
512
513/// Build the per-result acknowledgement frame for a consumed `ActivityResult`.
514fn result_ack_frame(
515    workflow_id: &aion_core::WorkflowId,
516    activity_id: &aion_core::ActivityId,
517) -> generated::ServerToWorker {
518    generated::ServerToWorker {
519        message: Some(generated::server_to_worker::Message::ResultAck(
520            generated::ResultAck {
521                workflow_id: Some(generated::WorkflowId {
522                    uuid: workflow_id.to_string(),
523                }),
524                activity_id: Some(generated::ActivityId {
525                    sequence_position: activity_id.sequence_position(),
526                }),
527            },
528        )),
529    }
530}
531
532fn decode_register(r: generated::RegisterWorker) -> ProtoRegisterWorker {
533    ProtoRegisterWorker {
534        namespaces: r.namespaces,
535        activity_types: r.activity_types,
536        task_queue: r.task_queue,
537        node: r.node,
538    }
539}
540
541fn encode_server_to_worker(message: WorkerMessage) -> generated::ServerToWorker {
542    let message = match message {
543        WorkerMessage::ActivityTask(task) => {
544            generated::server_to_worker::Message::Task(encode_task(task))
545        }
546        WorkerMessage::DrainRequest => {
547            generated::server_to_worker::Message::Drain(generated::DrainRequest {})
548        }
549    };
550    generated::ServerToWorker {
551        message: Some(message),
552    }
553}
554
555fn encode_task(task: aion_proto::ProtoActivityTask) -> generated::ActivityTask {
556    generated::ActivityTask {
557        workflow_id: task
558            .workflow_id
559            .map(|id| generated::WorkflowId { uuid: id.uuid }),
560        activity_id: task.activity_id.map(|id| generated::ActivityId {
561            sequence_position: id.sequence_position,
562        }),
563        activity_type: task.activity_type,
564        input: task.input.map(|p| generated::Payload {
565            content_type: p.content_type,
566            bytes: p.bytes,
567        }),
568        attempt: task.attempt,
569        labels: task.labels,
570        run_id: task.run_id.map(|id| generated::RunId { uuid: id.uuid }),
571    }
572}
573
574fn decode_activity_result(r: generated::ActivityResult) -> ProtoActivityResult {
575    ProtoActivityResult {
576        workflow_id: r
577            .workflow_id
578            .map(|id| aion_proto::ProtoWorkflowId { uuid: id.uuid }),
579        activity_id: r.activity_id.map(|id| aion_proto::ProtoActivityId {
580            sequence_position: id.sequence_position,
581        }),
582        outcome: r.outcome.map(decode_outcome),
583        run_id: r.run_id.map(|id| aion_proto::ProtoRunId { uuid: id.uuid }),
584    }
585}
586
587fn decode_heartbeat(r: generated::Heartbeat) -> aion_proto::ProtoHeartbeat {
588    aion_proto::ProtoHeartbeat {
589        workflow_id: r
590            .workflow_id
591            .map(|id| aion_proto::ProtoWorkflowId { uuid: id.uuid }),
592        activity_id: r.activity_id.map(|id| aion_proto::ProtoActivityId {
593            sequence_position: id.sequence_position,
594        }),
595        progress: r.progress.map(|p| aion_proto::ProtoPayload {
596            content_type: p.content_type,
597            bytes: p.bytes,
598        }),
599    }
600}
601
602fn decode_outcome(
603    outcome: generated::activity_result::Outcome,
604) -> aion_proto::proto_activity_result::Outcome {
605    match outcome {
606        generated::activity_result::Outcome::Result(p) => {
607            aion_proto::proto_activity_result::Outcome::Result(aion_proto::ProtoPayload {
608                content_type: p.content_type,
609                bytes: p.bytes,
610            })
611        }
612        generated::activity_result::Outcome::Error(e) => {
613            aion_proto::proto_activity_result::Outcome::Error(aion_proto::ProtoActivityError {
614                kind: e.kind,
615                message: e.message,
616                details: e.details.map(|p| aion_proto::ProtoPayload {
617                    content_type: p.content_type,
618                    bytes: p.bytes,
619                }),
620            })
621        }
622    }
623}
624
625#[cfg(test)]
626mod tests {
627    use std::time::{Duration, Instant};
628
629    use aion_core::{ActivityId, WorkflowId};
630
631    use crate::shutdown::DrainState;
632    use crate::worker::heartbeat::InFlightActivity;
633    use crate::worker::registry::ConnectedWorkerRegistry;
634    use crate::worker::{HeartbeatTracker, PendingActivities};
635
636    use super::teardown_worker_stream;
637
638    type TestError = Box<dyn std::error::Error>;
639
640    /// One tracked in-flight dispatch with a live pending waiter, ready for a
641    /// stream teardown: the registered worker, the shared tracker/pending/drain
642    /// state, and the waiter's receiver.
643    struct TeardownFixture {
644        registry: ConnectedWorkerRegistry,
645        tracker: HeartbeatTracker,
646        pending: PendingActivities,
647        drain: DrainState,
648        worker_id: crate::worker::registry::WorkerId,
649        workflow_id: WorkflowId,
650        activity_id: ActivityId,
651        rx: std::sync::mpsc::Receiver<Result<String, String>>,
652        /// Held so the registered worker stays routable until the teardown
653        /// under test deregisters it (dropping the guard would race that).
654        _registration: crate::worker::registry::WorkerRegistration,
655    }
656
657    fn fixture() -> Result<TeardownFixture, TestError> {
658        let registry = ConnectedWorkerRegistry::default();
659        let (tx, _rx) = tokio::sync::mpsc::channel(1);
660        let activity_types = [String::from("greet")];
661        let registration = registry.register("default", activity_types.iter(), tx)?;
662        let worker_id = registration
663            .worker_id()
664            .ok_or("test worker registration missing id")?;
665        let tracker = HeartbeatTracker::new(Duration::from_secs(5));
666        let pending = PendingActivities::default();
667        let workflow_id = WorkflowId::new_v4();
668        let activity_id = ActivityId::from_sequence_position(0);
669        tracker.track_task(
670            worker_id,
671            InFlightActivity {
672                workflow_id: workflow_id.clone(),
673                activity_id: activity_id.clone(),
674            },
675            Instant::now(),
676        )?;
677        let rx = pending.insert_for_test(workflow_id.clone(), activity_id.clone());
678        Ok(TeardownFixture {
679            registry,
680            tracker,
681            pending,
682            drain: DrainState::default(),
683            worker_id,
684            workflow_id,
685            activity_id,
686            rx,
687            _registration: registration,
688        })
689    }
690
691    /// #207: with drain begun, a stream teardown PARKS the in-flight dispatch —
692    /// the waiter resolves with the ephemeral parked sentinel, no lost-worker
693    /// failure is synthesized, and the tracker empties for drain accounting.
694    #[test]
695    fn teardown_under_drain_parks_instead_of_failing() -> Result<(), TestError> {
696        let fixture = fixture()?;
697        assert!(fixture.drain.begin());
698
699        teardown_worker_stream(
700            fixture.worker_id,
701            &fixture.tracker,
702            &fixture.registry,
703            &fixture.pending,
704            &fixture.drain,
705        );
706
707        let resolved = fixture.rx.recv_timeout(Duration::from_millis(200))?;
708        assert_eq!(
709            resolved,
710            Err(aion::PARKED_ACTIVITY_REASON.to_owned()),
711            "a drain teardown must resolve the waiter with the parked sentinel"
712        );
713        assert_eq!(fixture.tracker.in_flight_count()?, 0);
714        assert!(
715            !fixture.tracker.is_tracked(
716                fixture.worker_id,
717                &fixture.workflow_id,
718                &fixture.activity_id
719            )?,
720            "parking must retire the tracked entry"
721        );
722        Ok(())
723    }
724
725    /// Regression pin: WITHOUT drain, the teardown path is byte-identical to
726    /// before #207 — the waiter resolves with the retryable lost-worker
727    /// failure, so the engine's retry policy still governs genuine worker death.
728    #[test]
729    fn teardown_without_drain_still_fails_with_retryable_lost_worker() -> Result<(), TestError> {
730        let fixture = fixture()?;
731
732        teardown_worker_stream(
733            fixture.worker_id,
734            &fixture.tracker,
735            &fixture.registry,
736            &fixture.pending,
737            &fixture.drain,
738        );
739
740        let resolved = fixture.rx.recv_timeout(Duration::from_millis(200))?;
741        let reason = resolved.err().ok_or("expected a lost-worker failure")?;
742        assert!(
743            reason.starts_with("retryable:"),
744            "a mid-run teardown must stay a retryable failure: {reason}"
745        );
746        assert!(
747            reason.contains("lost before reporting activity result"),
748            "the failure must name worker loss: {reason}"
749        );
750        assert_eq!(fixture.tracker.in_flight_count()?, 0);
751        Ok(())
752    }
753}