1use std::collections::HashMap;
8use std::fmt;
9use std::future::Future;
10use std::pin::Pin;
11use std::sync::{Arc, RwLock as StdRwLock};
12
13use async_trait::async_trait;
14use bamboo_domain::{
15 SessionActivationDisposition, SessionActivationError, SessionActivationPort, SessionInboxPort,
16 SessionInboxWakeCandidate,
17};
18use chrono::{DateTime, Duration, Utc};
19use tokio::sync::{mpsc, watch, Mutex, RwLock};
20
21type AsyncRollback =
22 Box<dyn FnOnce() -> Pin<Box<dyn Future<Output = ()> + Send + 'static>> + Send + 'static>;
23
24pub struct SessionActivationLaunch {
29 pub run_id: String,
30 launch: Option<Box<dyn FnOnce() + Send + 'static>>,
31 rollback: Option<AsyncRollback>,
32 rollback_completion: Option<Arc<RollbackCompletion>>,
33}
34
35struct RollbackCompletion {
36 completed: watch::Sender<bool>,
37}
38
39impl RollbackCompletion {
40 fn new() -> Arc<Self> {
41 let (completed, _receiver) = watch::channel(false);
42 Arc::new(Self { completed })
43 }
44
45 fn complete(&self) {
46 self.completed.send_replace(true);
47 }
48
49 async fn wait(&self) {
50 let mut completed = self.completed.subscribe();
51 while !*completed.borrow() {
52 if completed.changed().await.is_err() {
53 break;
54 }
55 }
56 }
57}
58
59impl SessionActivationLaunch {
60 pub fn new(run_id: impl Into<String>, launch: impl FnOnce() + Send + 'static) -> Self {
61 Self {
62 run_id: run_id.into(),
63 launch: Some(Box::new(launch)),
64 rollback: None,
65 rollback_completion: None,
66 }
67 }
68
69 pub fn new_with_rollback(
75 run_id: impl Into<String>,
76 launch: impl FnOnce() + Send + 'static,
77 rollback: impl FnOnce() + Send + 'static,
78 ) -> Self {
79 Self::new_with_async_rollback(run_id, launch, move || async move {
80 rollback();
81 })
82 }
83
84 pub fn new_with_async_rollback<F, Fut>(
89 run_id: impl Into<String>,
90 launch: impl FnOnce() + Send + 'static,
91 rollback: F,
92 ) -> Self
93 where
94 F: FnOnce() -> Fut + Send + 'static,
95 Fut: Future<Output = ()> + Send + 'static,
96 {
97 let rollback_completion = RollbackCompletion::new();
98 Self {
99 run_id: run_id.into(),
100 launch: Some(Box::new(launch)),
101 rollback: Some(Box::new(move || Box::pin(rollback()))),
102 rollback_completion: Some(rollback_completion),
103 }
104 }
105
106 fn rollback_completion(&self) -> Option<Arc<RollbackCompletion>> {
107 self.rollback_completion.clone()
108 }
109
110 fn launch(mut self) {
111 if let Some(launch) = self.launch.take() {
112 self.rollback = None;
113 self.rollback_completion = None;
114 launch();
115 }
116 }
117}
118
119impl Drop for SessionActivationLaunch {
120 fn drop(&mut self) {
121 if self.launch.is_some() {
122 if let Some(rollback) = self.rollback.take() {
123 let completion = self.rollback_completion.take();
124 if let Ok(runtime) = tokio::runtime::Handle::try_current() {
125 runtime.spawn(async move {
126 rollback().await;
127 if let Some(completion) = completion {
128 completion.complete();
129 }
130 });
131 }
132 }
133 }
134 }
135}
136
137impl fmt::Debug for SessionActivationLaunch {
138 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
139 formatter
140 .debug_struct("SessionActivationLaunch")
141 .field("run_id", &self.run_id)
142 .finish_non_exhaustive()
143 }
144}
145
146#[derive(Debug)]
147pub enum SessionActivationReserveOutcome {
148 Reserved(SessionActivationLaunch),
151 AlreadyRunning { run_id: String },
153 NotFound,
155 NoWork,
157}
158
159#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
165pub enum SessionRunRegistrationError {
166 #[error(
167 "session activation owner collision for {target_session_id}: \
168 existing run {existing_run_id}, attempted run {attempted_run_id}"
169 )]
170 OwnerCollision {
171 target_session_id: String,
172 existing_run_id: String,
173 attempted_run_id: String,
174 },
175}
176
177impl SessionRunRegistrationError {
178 pub fn existing_run_id(&self) -> &str {
180 match self {
181 Self::OwnerCollision {
182 existing_run_id, ..
183 } => existing_run_id,
184 }
185 }
186}
187
188#[async_trait]
194pub trait SessionActivationSpawner: Send + Sync {
195 async fn reserve_activation(
196 &self,
197 target_session_id: &str,
198 inbox_generation: u64,
199 ) -> Result<SessionActivationReserveOutcome, SessionActivationError>;
200}
201
202#[derive(Debug, Clone)]
203struct ActiveOwner {
204 run_id: String,
205 finalizing: bool,
206 registrations: usize,
207 delivery_sink: Option<mpsc::UnboundedSender<u64>>,
211}
212
213#[derive(Debug)]
214struct TargetActivationState {
215 latest_generation: u64,
216 last_dispatched_generation: u64,
223 last_dispatched_coordinator_generation: u64,
227 owner_completion_epoch: u64,
230 last_recovery_attempt: Option<RecoveryAttempt>,
231 owner: Option<ActiveOwner>,
232 activation_reserved: bool,
233 activation_token: u64,
236 activation_epoch: watch::Sender<u64>,
241 notify: watch::Sender<u64>,
242}
243
244impl Default for TargetActivationState {
245 fn default() -> Self {
246 let (activation_epoch, _activation_receiver) = watch::channel(0);
247 let (notify, _receiver) = watch::channel(0);
248 Self {
249 latest_generation: 0,
250 last_dispatched_generation: 0,
251 last_dispatched_coordinator_generation: 0,
252 owner_completion_epoch: 0,
253 last_recovery_attempt: None,
254 owner: None,
255 activation_reserved: false,
256 activation_token: 0,
257 activation_epoch,
258 notify,
259 }
260 }
261}
262
263const INITIAL_RECOVERY_BACKOFF: Duration = Duration::seconds(30);
264const MAX_RECOVERY_BACKOFF: Duration = Duration::minutes(30);
265
266#[derive(Debug, Clone)]
267struct RecoveryAttempt {
268 candidate: SessionInboxWakeCandidate,
269 completion_epoch: u64,
270 backoff: Duration,
271 next_retry_at: DateTime<Utc>,
272 requires_completion: bool,
275}
276
277fn recovery_retry_at(
278 state: &TargetActivationState,
279 candidate: &SessionInboxWakeCandidate,
280 now: DateTime<Utc>,
281) -> Option<Option<DateTime<Utc>>> {
282 let previous = state
283 .last_recovery_attempt
284 .as_ref()
285 .filter(|attempt| attempt.candidate == *candidate)?;
286 if previous.requires_completion && state.owner_completion_epoch <= previous.completion_epoch {
287 return Some(None);
288 }
289 (now < previous.next_retry_at).then_some(Some(previous.next_retry_at))
290}
291
292fn record_recovery_attempt(
293 state: &mut TargetActivationState,
294 candidate: SessionInboxWakeCandidate,
295 now: DateTime<Utc>,
296 requires_completion: bool,
297) {
298 let backoff = state
299 .last_recovery_attempt
300 .as_ref()
301 .filter(|attempt| attempt.candidate == candidate)
302 .map_or(INITIAL_RECOVERY_BACKOFF, |attempt| {
303 (attempt.backoff * 2).min(MAX_RECOVERY_BACKOFF)
304 });
305 state.last_recovery_attempt = Some(RecoveryAttempt {
306 candidate,
307 completion_epoch: state.owner_completion_epoch,
308 backoff,
309 next_retry_at: now + backoff,
310 requires_completion,
311 });
312}
313
314#[derive(Debug, Clone, PartialEq, Eq)]
317pub struct SessionWakeReconcileResult {
318 pub disposition: Option<SessionActivationDisposition>,
319 pub next_due_at: Option<DateTime<Utc>>,
320}
321
322fn reserve_activation_token(state: &mut TargetActivationState) -> u64 {
323 state.activation_reserved = true;
324 state.activation_token = state.activation_token.wrapping_add(1).max(1);
325 state.activation_token
326}
327
328fn release_activation_token(state: &mut TargetActivationState, token: u64) -> bool {
329 if !state.activation_reserved || state.activation_token != token {
330 return false;
331 }
332 state.activation_reserved = false;
333 let next_epoch = (*state.activation_epoch.borrow()).wrapping_add(1);
334 state.activation_epoch.send_replace(next_epoch);
335 true
336}
337
338fn inbox_activation_error(error: bamboo_domain::SessionInboxError) -> SessionActivationError {
339 match error {
340 bamboo_domain::SessionInboxError::TargetNotFound(target) => {
341 SessionActivationError::TargetNotFound(target)
342 }
343 error => SessionActivationError::Internal(error.to_string()),
344 }
345}
346
347struct ActivationReservationLease {
353 router: SessionActivationRouter,
354 target_session_id: String,
355 token: u64,
356 recover_on_drop: bool,
357 armed: bool,
358 rollback_completion: Option<Arc<RollbackCompletion>>,
359}
360
361impl ActivationReservationLease {
362 fn new(
363 router: SessionActivationRouter,
364 target_session_id: &str,
365 token: u64,
366 recover_on_drop: bool,
367 ) -> Self {
368 Self {
369 router,
370 target_session_id: target_session_id.to_string(),
371 token,
372 recover_on_drop,
373 armed: true,
374 rollback_completion: None,
375 }
376 }
377
378 fn wait_for_rollback(&mut self, completion: Option<Arc<RollbackCompletion>>) {
379 self.rollback_completion = completion;
380 }
381
382 fn disarm(&mut self) {
383 self.armed = false;
384 }
385}
386
387impl Drop for ActivationReservationLease {
388 fn drop(&mut self) {
389 if !self.armed {
390 return;
391 }
392 let router = self.router.clone();
393 let target_session_id = self.target_session_id.clone();
394 let token = self.token;
395 let recover_on_drop = self.recover_on_drop;
396 let rollback_completion = self.rollback_completion.take();
397 if let Ok(runtime) = tokio::runtime::Handle::try_current() {
398 runtime.spawn(async move {
399 if let Some(completion) = rollback_completion {
400 completion.wait().await;
401 }
402 router
403 .recover_cancelled_reservation(&target_session_id, token, recover_on_drop)
404 .await;
405 });
406 }
407 }
408}
409
410#[derive(Clone, Default)]
412pub struct SessionActivationRouter {
413 states: Arc<Mutex<HashMap<String, TargetActivationState>>>,
414 spawner: Arc<RwLock<Option<Arc<dyn SessionActivationSpawner>>>>,
415 inbox: Arc<StdRwLock<Option<Arc<dyn SessionInboxPort>>>>,
416 #[cfg(test)]
417 terminal_inspection_pause: Arc<Mutex<Option<TerminalInspectionPause>>>,
418}
419
420#[cfg(test)]
421type TerminalInspectionPause = (Arc<tokio::sync::Notify>, Arc<tokio::sync::Notify>);
422
423type AbortCleanup =
424 Box<dyn FnOnce() -> Pin<Box<dyn Future<Output = ()> + Send + 'static>> + Send + 'static>;
425
426pub struct SessionRunRegistration {
434 router: Arc<SessionActivationRouter>,
435 target_session_id: String,
436 run_id: String,
437 notifications: Option<watch::Receiver<u64>>,
438 abort_cleanup: Option<AbortCleanup>,
439 quarantine_on_drop: bool,
440 selected_input_wakes: Option<watch::Receiver<u64>>,
441 armed: bool,
442}
443
444#[derive(Clone)]
447pub(crate) struct SessionRunOwnership {
448 router: Arc<SessionActivationRouter>,
449 target_session_id: String,
450 run_id: String,
451}
452
453impl SessionRunOwnership {
454 pub(crate) fn matches_router(&self, router: &Arc<SessionActivationRouter>) -> bool {
455 Arc::ptr_eq(&self.router, router)
456 }
457
458 pub(crate) async fn is_current(&self, session_id: &str) -> bool {
459 self.target_session_id == session_id && self.router.owns_run(session_id, &self.run_id).await
460 }
461}
462
463impl SessionRunRegistration {
464 pub(crate) fn ownership(&self) -> SessionRunOwnership {
465 SessionRunOwnership {
466 router: self.router.clone(),
467 target_session_id: self.target_session_id.clone(),
468 run_id: self.run_id.clone(),
469 }
470 }
471
472 pub(crate) fn restrict_to_selected_input(&mut self) {
477 if self.selected_input_wakes.is_none() {
478 self.selected_input_wakes = self.notifications.clone();
479 }
480 }
481
482 pub fn notifications_mut(&mut self) -> &mut watch::Receiver<u64> {
483 self.notifications
484 .as_mut()
485 .expect("run registration notifications are live before finalization")
486 }
487
488 pub fn set_abort_cleanup<F, Fut>(&mut self, cleanup: F)
491 where
492 F: FnOnce() -> Fut + Send + 'static,
493 Fut: Future<Output = ()> + Send + 'static,
494 {
495 self.abort_cleanup = Some(Box::new(move || Box::pin(cleanup())));
496 }
497
498 pub async fn begin_finalization(&mut self) {
499 self.router
500 .begin_finalization(&self.target_session_id, &self.run_id)
501 .await;
502 }
503
504 pub(crate) fn quarantine_on_drop(&mut self) {
508 self.quarantine_on_drop = true;
509 }
510
511 pub async fn quarantine(mut self) {
519 self.armed = false;
520 self.abort_cleanup = None;
521 self.router
522 .begin_finalization(&self.target_session_id, &self.run_id)
523 .await;
524 self.notifications.take();
525 }
526
527 pub async fn abandon(mut self) {
536 self.armed = false;
537 if self.selected_input_wakes.is_some() {
538 drop(self.notifications.take());
539 }
540 let abort_cleanup = self.abort_cleanup.take();
541 let router = self.router.clone();
542 let target_session_id = self.target_session_id.clone();
543 let run_id = self.run_id.clone();
544 let cleanup_target_session_id = target_session_id.clone();
545 let cleanup_run_id = run_id.clone();
546 let selected_input_wakes = self.selected_input_wakes.take();
547 let cleanup = tokio::spawn(async move {
548 if let Some(cleanup) = abort_cleanup {
549 cleanup().await;
550 }
551 router
552 .cleanup_abandoned_registration(
553 &cleanup_target_session_id,
554 &cleanup_run_id,
555 selected_input_wakes,
556 )
557 .await;
558 });
559 if let Err(error) = cleanup.await {
560 tracing::error!(
561 %target_session_id,
562 %run_id,
563 %error,
564 "detached SessionInbox registration cleanup failed"
565 );
566 }
567 }
568
569 pub async fn finish(
572 mut self,
573 admitted_generation: u64,
574 ) -> Result<Option<SessionActivationDisposition>, SessionActivationError> {
575 drop(self.notifications.take());
578 let released = match self.selected_input_wakes.as_ref() {
579 Some(wakes) => {
580 self.router
581 .release_selected_if_unnotified(&self.target_session_id, &self.run_id, wakes)
582 .await
583 }
584 None => false,
585 };
586 let result = if released {
587 Ok(None)
588 } else {
589 self.router
590 .finish_finalization(&self.target_session_id, &self.run_id, admitted_generation)
591 .await
592 };
593 if result.is_ok() {
594 self.armed = false;
595 self.abort_cleanup = None;
596 }
597 result
598 }
599}
600
601impl Drop for SessionRunRegistration {
602 fn drop(&mut self) {
603 if !self.armed {
604 return;
605 }
606 let router = self.router.clone();
607 let target_session_id = self.target_session_id.clone();
608 let run_id = self.run_id.clone();
609 let abort_cleanup = self.abort_cleanup.take();
610 let quarantine = self.quarantine_on_drop;
611 if self.selected_input_wakes.is_some() {
612 drop(self.notifications.take());
613 }
614 let selected_input_wakes = self.selected_input_wakes.take();
615 if let Ok(runtime) = tokio::runtime::Handle::try_current() {
616 runtime.spawn(async move {
617 if quarantine {
618 router.begin_finalization(&target_session_id, &run_id).await;
619 return;
620 }
621 if let Some(cleanup) = abort_cleanup {
622 cleanup().await;
623 }
624 router
625 .cleanup_abandoned_registration(
626 &target_session_id,
627 &run_id,
628 selected_input_wakes,
629 )
630 .await;
631 });
632 }
633 }
634}
635
636impl SessionActivationRouter {
637 async fn coordinator_generation(&self, target: &str) -> Result<u64, SessionActivationError> {
638 let inbox = self
639 .inbox
640 .read()
641 .unwrap_or_else(std::sync::PoisonError::into_inner)
642 .clone();
643 match inbox {
644 Some(inbox) => inbox
645 .coordinator_activation_generation(target)
646 .await
647 .map_err(inbox_activation_error),
648 None => Ok(0),
649 }
650 }
651
652 async fn released_pending_generation(
653 &self,
654 target: &str,
655 after: u64,
656 ) -> Result<u64, SessionActivationError> {
657 if self.coordinator_generation(target).await? <= after {
658 return Ok(0);
659 }
660 let inbox = self
661 .inbox
662 .read()
663 .unwrap_or_else(std::sync::PoisonError::into_inner)
664 .clone();
665 match inbox {
666 Some(inbox) => {
667 let backlog = inbox
668 .inspect(target)
669 .await
670 .map_err(inbox_activation_error)?;
671 Ok(
672 if backlog.activation_pending() && backlog.coordinator_generation > after {
673 backlog.coordinator_generation
674 } else {
675 0
676 },
677 )
678 }
679 None => Ok(0),
680 }
681 }
682
683 pub fn new() -> Arc<Self> {
684 Arc::new(Self::default())
685 }
686
687 pub async fn set_spawner(&self, spawner: Arc<dyn SessionActivationSpawner>) {
690 *self.spawner.write().await = Some(spawner);
691 }
692
693 pub async fn current_run_id(&self, target_session_id: &str) -> Option<String> {
695 self.states
696 .lock()
697 .await
698 .get(target_session_id)
699 .and_then(|state| state.owner.as_ref().map(|owner| owner.run_id.clone()))
700 }
701
702 pub fn set_inbox(&self, inbox: Arc<dyn SessionInboxPort>) {
704 *self
705 .inbox
706 .write()
707 .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(inbox);
708 }
709
710 pub async fn reconcile_wake(
714 &self,
715 target_session_id: &str,
716 now: DateTime<Utc>,
717 ) -> Result<SessionWakeReconcileResult, SessionActivationError> {
718 let inbox = self
719 .inbox
720 .read()
721 .unwrap_or_else(std::sync::PoisonError::into_inner)
722 .clone()
723 .ok_or_else(|| {
724 SessionActivationError::Internal("session inbox is not configured".into())
725 })?;
726 let readiness = inbox
727 .inspect_wake_readiness(target_session_id, now)
728 .await
729 .map_err(inbox_activation_error)?;
730 let Some(candidate) = readiness.ready else {
731 return Ok(SessionWakeReconcileResult {
732 disposition: None,
733 next_due_at: readiness.next_due_at,
734 });
735 };
736 let mut result = self
737 .request_recovery_activation(target_session_id, candidate, now)
738 .await?;
739 result.next_due_at = match (result.next_due_at, readiness.next_due_at) {
740 (Some(left), Some(right)) => Some(left.min(right)),
741 (left, right) => left.or(right),
742 };
743 Ok(result)
744 }
745
746 async fn reconcile_older_message(
750 &self,
751 target: &str,
752 ) -> Result<Option<SessionActivationDisposition>, SessionActivationError> {
753 let high = self
754 .states
755 .lock()
756 .await
757 .get(target)
758 .map_or(0, |state| state.last_dispatched_generation);
759 self.reconcile_message_below(target, high).await
760 }
761
762 async fn reconcile_message_below(
763 &self,
764 target: &str,
765 high: u64,
766 ) -> Result<Option<SessionActivationDisposition>, SessionActivationError> {
767 if high == 0 {
768 return Ok(None);
769 }
770 let inbox = self
771 .inbox
772 .read()
773 .unwrap_or_else(std::sync::PoisonError::into_inner)
774 .clone();
775 let Some(inbox) = inbox else {
776 return Ok(None);
777 };
778 let now = Utc::now();
779 let candidate = inbox
780 .inspect_wake_readiness(target, now)
781 .await
782 .map_err(inbox_activation_error)?
783 .ready;
784 match candidate.filter(|candidate| {
785 candidate.generation < high
786 && candidate.coordinator_message_policy
789 == Some(bamboo_domain::SessionActivationPolicy::InterruptSpecificWait)
790 }) {
791 Some(candidate) => self
792 .request_recovery_activation(target, candidate, now)
793 .await
794 .map(|result| result.disposition),
795 None => Ok(None),
796 }
797 }
798
799 #[cfg(test)]
800 async fn pause_after_terminal_inspection(&self) {
801 let pause = self.terminal_inspection_pause.lock().await.take();
802 if let Some((entered, release)) = pause {
803 entered.notify_one();
804 release.notified().await;
805 }
806 }
807
808 async fn request_recovery_activation(
809 &self,
810 target_session_id: &str,
811 candidate: SessionInboxWakeCandidate,
812 now: DateTime<Utc>,
813 ) -> Result<SessionWakeReconcileResult, SessionActivationError> {
814 loop {
815 let (reservation_wait, reservation_token) = {
816 let mut states = self.states.lock().await;
817 let state = states.entry(target_session_id.to_string()).or_default();
818 state.latest_generation = state.latest_generation.max(candidate.generation);
819 if let Some(owner) = state.owner.as_mut() {
820 if !owner.finalizing {
821 state.notify.send_replace(candidate.generation);
822 if owner
823 .delivery_sink
824 .as_ref()
825 .is_some_and(|sink| sink.send(candidate.generation).is_err())
826 {
827 owner.delivery_sink = None;
828 }
829 return Ok(SessionWakeReconcileResult {
830 disposition: Some(SessionActivationDisposition::ActiveNotified),
831 next_due_at: None,
832 });
833 }
834 return Ok(SessionWakeReconcileResult {
835 disposition: Some(SessionActivationDisposition::ActivationCoalesced),
836 next_due_at: None,
837 });
838 }
839 if let Some(next_due_at) = recovery_retry_at(state, &candidate, now) {
840 return Ok(SessionWakeReconcileResult {
841 disposition: None,
842 next_due_at,
843 });
844 } else if state.activation_reserved {
845 (Some(state.activation_epoch.subscribe()), None)
846 } else {
847 let token = reserve_activation_token(state);
848 (None, Some(token))
849 }
850 };
851
852 if let Some(mut reservation_wait) = reservation_wait {
853 if reservation_wait.changed().await.is_err() {
854 return Err(SessionActivationError::TargetNotFound(
855 target_session_id.to_string(),
856 ));
857 }
858 continue;
859 }
860 let token = reservation_token.expect("recovery activation owns reservation token");
861 let disposition = self
862 .dispatch_reserved_with_recovery(
863 target_session_id,
864 candidate.generation,
865 token,
866 false,
867 Some((candidate, now)),
868 )
869 .await?;
870 return Ok(SessionWakeReconcileResult {
871 disposition: Some(disposition),
872 next_due_at: None,
873 });
874 }
875 }
876
877 #[cfg(test)]
878 pub(crate) async fn hold_state_lock_for_test(
879 &self,
880 entered: Arc<tokio::sync::Notify>,
881 release: Arc<tokio::sync::Notify>,
882 ) {
883 let _states = self.states.lock().await;
884 entered.notify_one();
885 release.notified().await;
886 }
887
888 #[cfg(test)]
889 pub(crate) async fn install_owner_placeholder_for_test(
890 &self,
891 target_session_id: &str,
892 run_id: &str,
893 ) {
894 let mut states = self.states.lock().await;
895 states
896 .entry(target_session_id.to_string())
897 .or_default()
898 .owner = Some(ActiveOwner {
899 run_id: run_id.to_string(),
900 finalizing: false,
901 registrations: 0,
902 delivery_sink: None,
903 });
904 }
905
906 pub async fn register_run(
910 self: &Arc<Self>,
911 target_session_id: &str,
912 run_id: &str,
913 ) -> Result<SessionRunRegistration, SessionRunRegistrationError> {
914 loop {
915 let (notifications, reservation_wait) = {
916 let mut states = self.states.lock().await;
917 let state = states.entry(target_session_id.to_string()).or_default();
918 match state.owner.as_mut() {
919 Some(owner)
920 if owner.run_id == run_id
921 && owner.registrations == 0
922 && !owner.finalizing =>
923 {
924 owner.registrations = 1;
928 (Some(state.notify.subscribe()), None)
929 }
930 Some(owner) => {
931 return Err(SessionRunRegistrationError::OwnerCollision {
932 target_session_id: target_session_id.to_string(),
933 existing_run_id: owner.run_id.clone(),
934 attempted_run_id: run_id.to_string(),
935 });
936 }
937 None if state.activation_reserved => {
938 (None, Some(state.activation_epoch.subscribe()))
944 }
945 None => {
946 state.owner = Some(ActiveOwner {
947 run_id: run_id.to_string(),
948 finalizing: false,
949 registrations: 1,
950 delivery_sink: None,
951 });
952 (Some(state.notify.subscribe()), None)
953 }
954 }
955 };
956
957 if let Some(mut reservation_wait) = reservation_wait {
958 let _ = reservation_wait.changed().await;
959 continue;
960 }
961 return Ok(SessionRunRegistration {
962 router: self.clone(),
963 target_session_id: target_session_id.to_string(),
964 run_id: run_id.to_string(),
965 notifications,
966 abort_cleanup: None,
967 quarantine_on_drop: false,
968 selected_input_wakes: None,
969 armed: true,
970 });
971 }
972 }
973
974 pub async fn attach_delivery_sink(
980 &self,
981 target_session_id: &str,
982 sink: mpsc::UnboundedSender<u64>,
983 ) -> Option<String> {
984 let mut states = self.states.lock().await;
985 let state = states.get_mut(target_session_id)?;
986 let owner = state.owner.as_mut()?;
987 if owner.finalizing {
988 return None;
989 }
990 owner.delivery_sink = Some(sink.clone());
991 let run_id = owner.run_id.clone();
992 if state.latest_generation > 0 && sink.send(state.latest_generation).is_err() {
993 owner.delivery_sink = None;
994 return None;
995 }
996 Some(run_id)
997 }
998
999 pub async fn detach_delivery_sink(&self, target_session_id: &str, run_id: &str) {
1002 let mut states = self.states.lock().await;
1003 let Some(state) = states.get_mut(target_session_id) else {
1004 return;
1005 };
1006 let Some(owner) = state.owner.as_mut() else {
1007 return;
1008 };
1009 if owner.run_id == run_id {
1010 owner.delivery_sink = None;
1011 }
1012 }
1013
1014 pub async fn owns_run(&self, target_session_id: &str, run_id: &str) -> bool {
1016 let states = self.states.lock().await;
1017 states
1018 .get(target_session_id)
1019 .and_then(|state| state.owner.as_ref())
1020 .is_some_and(|owner| owner.run_id == run_id && !owner.finalizing)
1021 }
1022
1023 pub async fn begin_finalization(&self, target_session_id: &str, run_id: &str) {
1028 let mut states = self.states.lock().await;
1029 let state = states.entry(target_session_id.to_string()).or_default();
1030 match state.owner.as_mut() {
1031 Some(owner) if owner.run_id == run_id => owner.finalizing = true,
1032 Some(_) => {}
1033 None => {
1034 state.owner = Some(ActiveOwner {
1035 run_id: run_id.to_string(),
1036 finalizing: true,
1037 registrations: 0,
1038 delivery_sink: None,
1039 });
1040 }
1041 }
1042 }
1043
1044 async fn release_selected_if_unnotified(
1049 &self,
1050 target_session_id: &str,
1051 run_id: &str,
1052 wakes: &watch::Receiver<u64>,
1053 ) -> bool {
1054 let mut states = self.states.lock().await;
1055 let Some(state) = states.get_mut(target_session_id) else {
1056 return true;
1057 };
1058 if !state
1059 .owner
1060 .as_ref()
1061 .is_some_and(|owner| owner.run_id == run_id)
1062 {
1063 return true;
1064 }
1065 if wakes.has_changed().unwrap_or(false) {
1066 return false;
1067 }
1068 state.owner = None;
1069 state.owner_completion_epoch = state.owner_completion_epoch.wrapping_add(1);
1070 if state.latest_generation == 0
1071 && !state.activation_reserved
1072 && state.notify.receiver_count() <= 1
1073 && state.activation_epoch.receiver_count() == 0
1074 {
1075 states.remove(target_session_id);
1076 }
1077 true
1078 }
1079
1080 async fn cleanup_abandoned_registration(
1081 self: &Arc<Self>,
1082 target_session_id: &str,
1083 run_id: &str,
1084 selected_input_wakes: Option<watch::Receiver<u64>>,
1085 ) {
1086 if let Some(wakes) = selected_input_wakes {
1087 if self
1088 .release_selected_if_unnotified(target_session_id, run_id, &wakes)
1089 .await
1090 {
1091 return;
1092 }
1093 }
1094 {
1095 let mut states = self.states.lock().await;
1096 let Some(state) = states.get_mut(target_session_id) else {
1097 return;
1098 };
1099 let Some(owner) = state.owner.as_mut() else {
1100 return;
1101 };
1102 if owner.run_id != run_id {
1103 return;
1104 }
1105 owner.registrations = 0;
1106 owner.finalizing = true;
1107 }
1108
1109 let inbox = self
1114 .inbox
1115 .read()
1116 .unwrap_or_else(std::sync::PoisonError::into_inner)
1117 .clone();
1118 let (durable_generation, released_prefix) = match inbox {
1119 Some(inbox) => match inbox.inspect(target_session_id).await {
1120 Ok(backlog) if backlog.activation_pending() => (
1121 backlog.activation_generation,
1122 backlog.coordinator_generation,
1123 ),
1124 Ok(_) => (0, 0),
1125 Err(error) => {
1126 tracing::warn!(
1127 %target_session_id,
1128 %run_id,
1129 %error,
1130 "failed to reconcile durable SessionInbox generation for abandoned run"
1131 );
1132 (0, 0)
1133 }
1134 },
1135 None => (0, 0),
1136 };
1137
1138 #[cfg(test)]
1139 self.pause_after_terminal_inspection().await;
1140 let (reservation_to_dispatch, previous_high) = {
1141 let mut states = self.states.lock().await;
1142 let Some(state) = states.get_mut(target_session_id) else {
1143 return;
1144 };
1145 if !state.owner.as_ref().is_some_and(|owner| {
1146 owner.run_id == run_id && owner.finalizing && owner.registrations == 0
1147 }) {
1148 return;
1149 }
1150 state.latest_generation = state.latest_generation.max(durable_generation);
1151 let previous_high = state.last_dispatched_generation;
1152 state.owner = None;
1153 state.owner_completion_epoch = state.owner_completion_epoch.wrapping_add(1);
1154 if (state.latest_generation > state.last_dispatched_generation
1155 || released_prefix > state.last_dispatched_coordinator_generation)
1156 && !state.activation_reserved
1157 {
1158 let generation = state.latest_generation;
1159 let token = reserve_activation_token(state);
1160 (Some((generation, token)), previous_high)
1161 } else {
1162 if state.latest_generation == 0
1163 && !state.activation_reserved
1164 && state.notify.receiver_count() == 0
1165 && state.activation_epoch.receiver_count() == 0
1166 {
1167 states.remove(target_session_id);
1168 }
1169 (None, previous_high)
1170 }
1171 };
1172
1173 if let Some((generation, token)) = reservation_to_dispatch {
1174 if let Err(error) = self
1175 .dispatch_reserved(target_session_id, generation, token)
1176 .await
1177 {
1178 tracing::error!(
1179 %target_session_id,
1180 %run_id,
1181 %error,
1182 "failed to reserve successor for abandoned SessionInbox owner"
1183 );
1184 }
1185 } else {
1186 if let Err(error) = self
1190 .reconcile_message_below(target_session_id, previous_high)
1191 .await
1192 {
1193 tracing::warn!(%target_session_id, %error,
1194 "failed to reconcile older eligible message after abandoned owner");
1195 }
1196 }
1197 }
1198
1199 pub async fn finish_finalization(
1203 &self,
1204 target_session_id: &str,
1205 run_id: &str,
1206 admitted_generation: u64,
1207 ) -> Result<Option<SessionActivationDisposition>, SessionActivationError> {
1208 let inbox = self
1211 .inbox
1212 .read()
1213 .unwrap_or_else(std::sync::PoisonError::into_inner)
1214 .clone();
1215 let durable_pending = match inbox {
1216 Some(inbox) => Some(
1217 inbox
1218 .inspect(target_session_id)
1219 .await
1220 .map_err(|error| SessionActivationError::Internal(error.to_string()))?,
1221 ),
1222 None => None,
1223 };
1224 #[cfg(test)]
1225 self.pause_after_terminal_inspection().await;
1226 let (reservation_to_dispatch, previous_high) = {
1227 let mut states = self.states.lock().await;
1228 let state = states.entry(target_session_id.to_string()).or_default();
1229 let previous_high = state.last_dispatched_generation;
1230 if state
1231 .owner
1232 .as_ref()
1233 .is_some_and(|owner| owner.run_id != run_id)
1234 {
1235 return Ok(None);
1236 }
1237 if state.owner.take().is_some() {
1238 state.owner_completion_epoch = state.owner_completion_epoch.wrapping_add(1);
1239 }
1240 let pending = durable_pending
1241 .as_ref()
1242 .is_some_and(|backlog| backlog.activation_pending());
1243 let released_prefix = durable_pending
1244 .as_ref()
1245 .filter(|_| pending)
1246 .map_or(0, |backlog| backlog.coordinator_generation);
1247 if let Some(backlog) = durable_pending.as_ref() {
1248 state.latest_generation =
1249 state.latest_generation.max(backlog.activation_generation);
1250 }
1251 if (pending || state.latest_generation > admitted_generation)
1252 && (state.latest_generation > state.last_dispatched_generation
1253 || released_prefix > state.last_dispatched_coordinator_generation)
1254 && !state.activation_reserved
1255 {
1256 let generation = state.latest_generation;
1257 let token = reserve_activation_token(state);
1258 (Some((generation, token)), previous_high)
1259 } else {
1260 if state.latest_generation <= admitted_generation
1265 && !pending
1266 && state.owner.is_none()
1267 && !state.activation_reserved
1268 && state.notify.receiver_count() == 0
1269 && state.activation_epoch.receiver_count() == 0
1270 {
1271 states.remove(target_session_id);
1272 }
1273 (None, previous_high)
1274 }
1275 };
1276
1277 match reservation_to_dispatch {
1278 Some((generation, token)) => self
1279 .dispatch_reserved(target_session_id, generation, token)
1280 .await
1281 .map(Some),
1282 None => {
1285 self.reconcile_message_below(target_session_id, previous_high)
1286 .await
1287 }
1288 }
1289 }
1290
1291 pub async fn subscribe(&self, target_session_id: &str) -> watch::Receiver<u64> {
1292 let mut states = self.states.lock().await;
1293 states
1294 .entry(target_session_id.to_string())
1295 .or_default()
1296 .notify
1297 .subscribe()
1298 }
1299
1300 async fn recover_cancelled_reservation(
1309 &self,
1310 target_session_id: &str,
1311 reservation_token: u64,
1312 recover_on_drop: bool,
1313 ) {
1314 let previous_prefix = {
1315 let mut states = self.states.lock().await;
1316 let Some(state) = states.get_mut(target_session_id) else {
1317 return;
1318 };
1319 if !release_activation_token(state, reservation_token) || !recover_on_drop {
1320 return;
1321 }
1322 state.last_dispatched_coordinator_generation
1323 };
1324 let released_prefix = self
1325 .released_pending_generation(target_session_id, previous_prefix)
1326 .await
1327 .unwrap_or(0);
1328 let reservation_to_dispatch = {
1329 let mut states = self.states.lock().await;
1330 let Some(state) = states.get_mut(target_session_id) else {
1331 return;
1332 };
1333 if state.owner.is_none()
1334 && (state.latest_generation > state.last_dispatched_generation
1335 || released_prefix > state.last_dispatched_coordinator_generation)
1336 && !state.activation_reserved
1337 {
1338 let generation = state.latest_generation;
1339 let token = reserve_activation_token(state);
1340 Some((generation, token))
1341 } else {
1342 None
1343 }
1344 };
1345
1346 if let Some((generation, token)) = reservation_to_dispatch {
1347 if let Err(error) = self
1348 .dispatch_reserved_with_recovery(target_session_id, generation, token, false, None)
1349 .await
1350 {
1351 tracing::error!(
1352 %target_session_id,
1353 %error,
1354 "failed bounded retry after cancelled SessionInbox activation reservation"
1355 );
1356 }
1357 }
1358 }
1359
1360 async fn dispatch_reserved(
1361 &self,
1362 target_session_id: &str,
1363 generation: u64,
1364 reservation_token: u64,
1365 ) -> Result<SessionActivationDisposition, SessionActivationError> {
1366 self.dispatch_reserved_with_recovery(
1367 target_session_id,
1368 generation,
1369 reservation_token,
1370 true,
1371 None,
1372 )
1373 .await
1374 }
1375
1376 async fn dispatch_reserved_with_recovery(
1377 &self,
1378 target_session_id: &str,
1379 generation: u64,
1380 reservation_token: u64,
1381 recover_on_drop: bool,
1382 recovery: Option<(SessionInboxWakeCandidate, DateTime<Utc>)>,
1383 ) -> Result<SessionActivationDisposition, SessionActivationError> {
1384 let mut lease = ActivationReservationLease::new(
1385 self.clone(),
1386 target_session_id,
1387 reservation_token,
1388 recover_on_drop,
1389 );
1390 let Some(spawner) = self.spawner.read().await.clone() else {
1391 let mut states = self.states.lock().await;
1392 if let Some(state) = states.get_mut(target_session_id) {
1393 release_activation_token(state, reservation_token);
1394 }
1395 lease.disarm();
1396 return Err(SessionActivationError::Internal(
1397 "session activation spawner is not configured".to_string(),
1398 ));
1399 };
1400
1401 let coordinator_generation = self.coordinator_generation(target_session_id).await?;
1404
1405 let outcome = spawner
1406 .reserve_activation(target_session_id, generation)
1407 .await;
1408 if let Ok(SessionActivationReserveOutcome::Reserved(launch)) = &outcome {
1409 lease.wait_for_rollback(launch.rollback_completion());
1410 }
1411 match outcome {
1412 Ok(SessionActivationReserveOutcome::Reserved(launch)) => {
1413 let run_id = launch.run_id.clone();
1414 {
1415 let mut states = self.states.lock().await;
1416 let state = states.entry(target_session_id.to_string()).or_default();
1417 if !release_activation_token(state, reservation_token) {
1418 lease.disarm();
1426 return if state.owner.is_some() {
1427 Ok(SessionActivationDisposition::ActiveNotified)
1428 } else {
1429 Err(SessionActivationError::Internal(
1430 "session activation reservation was superseded".to_string(),
1431 ))
1432 };
1433 }
1434 state.owner = Some(ActiveOwner {
1435 run_id,
1436 finalizing: false,
1437 registrations: 0,
1438 delivery_sink: None,
1439 });
1440 state.last_dispatched_generation =
1441 state.last_dispatched_generation.max(generation);
1442 state.last_dispatched_coordinator_generation = state
1443 .last_dispatched_coordinator_generation
1444 .max(coordinator_generation);
1445 if let Some((candidate, now)) = recovery {
1446 record_recovery_attempt(state, candidate, now, true);
1447 }
1448 }
1449 lease.disarm();
1450 launch.launch();
1453 Ok(SessionActivationDisposition::ActivationReserved)
1454 }
1455 Ok(SessionActivationReserveOutcome::AlreadyRunning { run_id }) => {
1456 let mut states = self.states.lock().await;
1457 let state = states.entry(target_session_id.to_string()).or_default();
1458 if !release_activation_token(state, reservation_token) {
1459 lease.disarm();
1460 return Ok(SessionActivationDisposition::ActiveNotified);
1461 }
1462 state.owner = Some(ActiveOwner {
1463 run_id,
1464 finalizing: false,
1465 registrations: 0,
1466 delivery_sink: None,
1467 });
1468 state.notify.send_replace(generation);
1469 if let Some((candidate, now)) = recovery {
1470 record_recovery_attempt(state, candidate, now, true);
1471 }
1472 lease.disarm();
1473 Ok(SessionActivationDisposition::ActiveNotified)
1474 }
1475 Ok(SessionActivationReserveOutcome::NoWork) => {
1476 let mut states = self.states.lock().await;
1477 if let Some(state) = states.get_mut(target_session_id) {
1478 if release_activation_token(state, reservation_token) {
1479 if let Some((candidate, now)) = recovery {
1480 record_recovery_attempt(state, candidate, now, false);
1481 }
1482 }
1483 }
1484 lease.disarm();
1485 Ok(SessionActivationDisposition::ActivationCoalesced)
1486 }
1487 Ok(SessionActivationReserveOutcome::NotFound) => {
1488 let mut states = self.states.lock().await;
1489 if states
1490 .get(target_session_id)
1491 .is_some_and(|state| state.activation_token == reservation_token)
1492 {
1493 states.remove(target_session_id);
1494 }
1495 lease.disarm();
1496 Err(SessionActivationError::TargetNotFound(
1497 target_session_id.to_string(),
1498 ))
1499 }
1500 Err(error) => {
1501 let mut states = self.states.lock().await;
1502 if let Some(state) = states.get_mut(target_session_id) {
1503 release_activation_token(state, reservation_token);
1504 }
1505 lease.disarm();
1506 Err(error)
1507 }
1508 }
1509 }
1510}
1511
1512#[async_trait]
1513impl SessionActivationPort for SessionActivationRouter {
1514 async fn request_activation(
1515 &self,
1516 target_session_id: &str,
1517 inbox_generation: u64,
1518 ) -> Result<SessionActivationDisposition, SessionActivationError> {
1519 let previous_prefix = {
1520 let states = self.states.lock().await;
1521 states
1522 .get(target_session_id)
1523 .filter(|state| inbox_generation <= state.last_dispatched_generation)
1524 .map(|state| state.last_dispatched_coordinator_generation)
1525 };
1526 let released_prefix = match previous_prefix {
1527 Some(prefix) => {
1528 self.released_pending_generation(target_session_id, prefix)
1529 .await?
1530 }
1531 None => 0,
1532 };
1533 loop {
1534 let (reservation_wait, reservation_token) = {
1535 let mut states = self.states.lock().await;
1536 let state = states.entry(target_session_id.to_string()).or_default();
1537 state.latest_generation = state.latest_generation.max(inbox_generation);
1538 if let Some(owner) = state.owner.as_mut() {
1539 if !owner.finalizing {
1540 state.notify.send_replace(inbox_generation);
1541 if owner
1542 .delivery_sink
1543 .as_ref()
1544 .is_some_and(|sink| sink.send(inbox_generation).is_err())
1545 {
1546 owner.delivery_sink = None;
1547 }
1548 return Ok(SessionActivationDisposition::ActiveNotified);
1549 }
1550 return Ok(SessionActivationDisposition::ActivationCoalesced);
1551 }
1552 if inbox_generation <= state.last_dispatched_generation
1553 && released_prefix <= state.last_dispatched_coordinator_generation
1554 {
1555 drop(states);
1556 return self
1557 .reconcile_older_message(target_session_id)
1558 .await
1559 .map(|result| {
1560 result.unwrap_or(SessionActivationDisposition::ActivationCoalesced)
1561 });
1562 }
1563 if state.activation_reserved {
1564 (Some(state.activation_epoch.subscribe()), None)
1565 } else {
1566 let token = reserve_activation_token(state);
1567 (None, Some(token))
1568 }
1569 };
1570
1571 if let Some(mut reservation_wait) = reservation_wait {
1572 if reservation_wait.changed().await.is_err() {
1578 return Err(SessionActivationError::TargetNotFound(
1579 target_session_id.to_string(),
1580 ));
1581 }
1582 continue;
1583 }
1584
1585 let reservation_token =
1586 reservation_token.expect("non-waiting activation owns a reservation token");
1587 return self
1588 .dispatch_reserved_with_recovery(
1589 target_session_id,
1590 inbox_generation,
1591 reservation_token,
1592 false,
1593 None,
1594 )
1595 .await;
1596 }
1597 }
1598}
1599
1600#[cfg(test)]
1601mod tests {
1602 use super::*;
1603 use std::sync::atomic::{AtomicUsize, Ordering};
1604 use tokio::sync::{Barrier, Notify};
1605
1606 #[tokio::test]
1607 async fn wake_recovery_retries_same_candidate_only_after_owner_completion_and_backoff() {
1608 use bamboo_domain::{
1609 Session, SessionActivationPolicy, SessionInboxLimits, SessionMessageEnvelope, Storage,
1610 };
1611 use bamboo_storage::{FileSessionInbox, SessionStoreV2};
1612
1613 let temp = tempfile::TempDir::new().unwrap();
1614 let store = Arc::new(SessionStoreV2::new(temp.path().into()).await.unwrap());
1615 store
1616 .save_session(&Session::new("target", "model"))
1617 .await
1618 .unwrap();
1619 let inbox = Arc::new(FileSessionInbox::new(store, SessionInboxLimits::default()));
1620 inbox
1621 .deliver_with_activation_intent(
1622 &SessionMessageEnvelope::user_input("target", "still pending"),
1623 SessionActivationPolicy::RespectSpecificWait,
1624 None,
1625 )
1626 .await
1627 .unwrap();
1628 let router = SessionActivationRouter::new();
1629 router.set_inbox(inbox);
1630 let spawner = spawner();
1631 router.set_spawner(spawner.clone()).await;
1632 let now = Utc::now();
1633 assert_eq!(
1634 router
1635 .reconcile_wake("target", now)
1636 .await
1637 .unwrap()
1638 .disposition,
1639 Some(SessionActivationDisposition::ActivationReserved)
1640 );
1641 assert_eq!(
1642 router
1643 .reconcile_wake("target", now)
1644 .await
1645 .unwrap()
1646 .disposition,
1647 Some(SessionActivationDisposition::ActiveNotified)
1648 );
1649 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 1);
1650 let mut registration = router.register_run("target", "run-1").await.unwrap();
1651 registration.begin_finalization().await;
1652 assert_eq!(registration.finish(0).await.unwrap(), None);
1653
1654 let deferred = router
1655 .reconcile_wake("target", now + Duration::seconds(29))
1656 .await
1657 .unwrap();
1658 assert_eq!(deferred.disposition, None);
1659 assert_eq!(deferred.next_due_at, Some(now + INITIAL_RECOVERY_BACKOFF));
1660 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 1);
1661 assert_eq!(
1662 router
1663 .reconcile_wake("target", now + INITIAL_RECOVERY_BACKOFF)
1664 .await
1665 .unwrap()
1666 .disposition,
1667 Some(SessionActivationDisposition::ActivationReserved)
1668 );
1669 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
1670 }
1671
1672 #[tokio::test]
1673 async fn wake_recovery_waits_for_exact_owned_lease_expiry() {
1674 use bamboo_domain::{
1675 Session, SessionActivationPolicy, SessionInboxLeaseRequest, SessionInboxLimits,
1676 SessionMessageEnvelope, Storage,
1677 };
1678 use bamboo_storage::{FileSessionInbox, SessionStoreV2};
1679
1680 let temp = tempfile::TempDir::new().unwrap();
1681 let store = Arc::new(SessionStoreV2::new(temp.path().into()).await.unwrap());
1682 store
1683 .save_session(&Session::new("target", "model"))
1684 .await
1685 .unwrap();
1686 let inbox = Arc::new(FileSessionInbox::new(store, SessionInboxLimits::default()));
1687 inbox
1688 .deliver_with_activation_intent(
1689 &SessionMessageEnvelope::user_input("target", "unacknowledged"),
1690 SessionActivationPolicy::RespectSpecificWait,
1691 None,
1692 )
1693 .await
1694 .unwrap();
1695 let router = SessionActivationRouter::new();
1696 router.set_inbox(inbox.clone());
1697 let spawner = spawner();
1698 router.set_spawner(spawner.clone()).await;
1699 let now = Utc::now();
1700 assert_eq!(
1701 router
1702 .reconcile_wake("target", now)
1703 .await
1704 .unwrap()
1705 .disposition,
1706 Some(SessionActivationDisposition::ActivationReserved)
1707 );
1708 let claim = inbox
1709 .claim_owned(
1710 "target",
1711 1,
1712 None,
1713 &SessionInboxLeaseRequest {
1714 consumer: bamboo_domain::SessionInboxConsumerId::new(),
1715 now,
1716 duration: Duration::seconds(10),
1717 },
1718 )
1719 .await
1720 .unwrap()
1721 .remove(0);
1722 let mut registration = router.register_run("target", "run-1").await.unwrap();
1723 registration.begin_finalization().await;
1724 assert_eq!(registration.finish(0).await.unwrap(), None);
1725 let deferred = router
1726 .reconcile_wake("target", claim.lease.expires_at - Duration::seconds(1))
1727 .await
1728 .unwrap();
1729 assert_eq!(deferred.disposition, None);
1730 assert_eq!(deferred.next_due_at, Some(claim.lease.expires_at));
1731 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 1);
1732 assert_eq!(
1733 router
1734 .reconcile_wake("target", claim.lease.expires_at)
1735 .await
1736 .unwrap()
1737 .disposition,
1738 Some(SessionActivationDisposition::ActivationReserved)
1739 );
1740 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
1741 }
1742
1743 #[tokio::test]
1744 async fn older_staged_release_after_immediate_ack_dispatches_once_without_poison_hot_loop() {
1745 use bamboo_domain::{
1746 Session, SessionActivationPolicy, SessionInboxLimits, SessionMessageEnvelope, Storage,
1747 };
1748 use bamboo_storage::{FileSessionInbox, SessionStoreV2};
1749 async fn release(
1750 inbox: &FileSessionInbox,
1751 receipt: &bamboo_domain::SessionInboxReceipt,
1752 precise: bool,
1753 ) {
1754 if precise {
1755 inbox
1756 .mark_message_activation_eligible(
1757 "parent",
1758 receipt,
1759 SessionActivationPolicy::InterruptSpecificWait,
1760 )
1761 .await
1762 .unwrap();
1763 assert_eq!(
1764 inbox
1765 .coordinator_activation_generation("parent")
1766 .await
1767 .unwrap(),
1768 0
1769 );
1770 } else {
1771 inbox
1772 .mark_activation_eligible(
1773 "parent",
1774 receipt.generation,
1775 SessionActivationPolicy::RespectSpecificWait,
1776 )
1777 .await
1778 .unwrap();
1779 }
1780 }
1781 for (release_during_finalization, precise) in
1782 [(false, false), (true, false), (false, true), (true, true)]
1783 {
1784 let temp = tempfile::TempDir::new().unwrap();
1785 let store = Arc::new(SessionStoreV2::new(temp.path().into()).await.unwrap());
1786 store
1787 .save_session(&Session::new("parent", "model"))
1788 .await
1789 .unwrap();
1790 let inbox = Arc::new(FileSessionInbox::new(
1791 store.clone(),
1792 SessionInboxLimits::default(),
1793 ));
1794 let staged = SessionMessageEnvelope::user_input("parent", "staged result");
1795 let first = inbox.deliver(&staged).await.unwrap();
1796 let immediate = SessionMessageEnvelope::user_input("parent", "immediate steer");
1797 let second = inbox
1798 .deliver_with_activation_intent(
1799 &immediate,
1800 SessionActivationPolicy::RespectSpecificWait,
1801 None,
1802 )
1803 .await
1804 .unwrap();
1805 let reopened_store = Arc::new(SessionStoreV2::new(temp.path().into()).await.unwrap());
1806 let restarted = Arc::new(FileSessionInbox::new(
1807 reopened_store,
1808 SessionInboxLimits::default(),
1809 ));
1810 let router = SessionActivationRouter::new();
1811 router.set_inbox(restarted.clone());
1812 let spawner = spawner();
1813 router.set_spawner(spawner.clone()).await;
1814 let _observer = router.subscribe("parent").await;
1817 let backlog = restarted.inspect("parent").await.unwrap();
1818 assert_eq!(
1819 router
1820 .request_activation("parent", backlog.activation_generation)
1821 .await
1822 .unwrap(),
1823 SessionActivationDisposition::ActivationReserved
1824 );
1825 let claims = restarted.claim("parent", 128).await.unwrap();
1826 assert_eq!(claims.len(), 1);
1827 assert_eq!(claims[0].generation, second.generation);
1828 restarted.ack("parent", &claims[0]).await.unwrap();
1829 let mut registration = router.register_run("parent", "run-1").await.unwrap();
1830 registration.begin_finalization().await;
1831 if release_during_finalization {
1832 release(&restarted, &first, precise).await;
1833 assert_eq!(
1834 router
1835 .request_activation("parent", first.generation)
1836 .await
1837 .unwrap(),
1838 SessionActivationDisposition::ActivationCoalesced
1839 );
1840 assert_eq!(
1841 registration.finish(second.generation).await.unwrap(),
1842 Some(SessionActivationDisposition::ActivationReserved)
1843 );
1844 } else {
1845 assert_eq!(registration.finish(second.generation).await.unwrap(), None);
1846 release(&restarted, &first, precise).await;
1847 assert_eq!(
1848 router
1849 .request_activation("parent", first.generation)
1850 .await
1851 .unwrap(),
1852 SessionActivationDisposition::ActivationReserved
1853 );
1854 }
1855 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
1856 let claims = restarted.claim("parent", 128).await.unwrap();
1857 assert_eq!(claims.len(), 1);
1858 assert_eq!(claims[0].envelope.id, staged.id);
1859 let mut registration = router.register_run("parent", "run-2").await.unwrap();
1862 registration.begin_finalization().await;
1863 assert_eq!(registration.finish(second.generation).await.unwrap(), None);
1864 assert_eq!(
1865 router
1866 .request_activation("parent", first.generation)
1867 .await
1868 .unwrap(),
1869 SessionActivationDisposition::ActivationCoalesced
1870 );
1871 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
1872 restarted.ack("parent", &claims[0]).await.unwrap();
1873 }
1874 }
1875
1876 #[tokio::test]
1877 async fn precise_older_interrupt_below_released_prefix_reaches_one_successor() {
1878 use bamboo_domain::{
1879 Session, SessionActivationPolicy, SessionInboxLimits, SessionMessageEnvelope, Storage,
1880 };
1881 use bamboo_storage::{FileSessionInbox, SessionStoreV2};
1882 for with_older_sibling in [false, true] {
1883 let temp = tempfile::TempDir::new().unwrap();
1884 let store = Arc::new(SessionStoreV2::new(temp.path().into()).await.unwrap());
1885 store
1886 .save_session(&Session::new("parent", "model"))
1887 .await
1888 .unwrap();
1889 let inbox = Arc::new(FileSessionInbox::new(store, SessionInboxLimits::default()));
1890 let sibling = if with_older_sibling {
1891 Some(
1892 inbox
1893 .deliver(&SessionMessageEnvelope::user_input(
1894 "parent",
1895 "older Respect sibling",
1896 ))
1897 .await
1898 .unwrap(),
1899 )
1900 } else {
1901 None
1902 };
1903 let older = inbox
1904 .deliver(&SessionMessageEnvelope::user_input("parent", "older child"))
1905 .await
1906 .unwrap();
1907 let newer = inbox
1908 .deliver(&SessionMessageEnvelope::user_input("parent", "newer Bash"))
1909 .await
1910 .unwrap();
1911 inbox
1912 .mark_activation_eligible(
1913 "parent",
1914 newer.generation,
1915 SessionActivationPolicy::RespectSpecificWait,
1916 )
1917 .await
1918 .unwrap();
1919 let router = SessionActivationRouter::new();
1920 router.set_inbox(inbox.clone());
1921 let spawner = spawner();
1922 router.set_spawner(spawner.clone()).await;
1923 router
1924 .request_activation("parent", newer.generation)
1925 .await
1926 .unwrap();
1927 let claims = inbox.claim("parent", 128).await.unwrap();
1928 assert_eq!(claims.len(), 2 + usize::from(with_older_sibling));
1929 let older_claim = claims
1930 .iter()
1931 .find(|claim| claim.envelope.id == older.id)
1932 .unwrap();
1933 assert_eq!(older_claim.generation, older.generation);
1934 assert_eq!(
1935 older_claim.activation_policy,
1936 SessionActivationPolicy::RespectSpecificWait
1937 );
1938 let newer_claim = claims
1939 .iter()
1940 .find(|claim| claim.envelope.id == newer.id)
1941 .unwrap();
1942 inbox.ack("parent", newer_claim).await.unwrap();
1943 let mut owner = router.register_run("parent", "run-1").await.unwrap();
1944 owner.begin_finalization().await;
1945 assert_eq!(owner.finish(newer.generation).await.unwrap(), None);
1946 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 1);
1947
1948 inbox
1949 .mark_message_activation_eligible(
1950 "parent",
1951 &older,
1952 SessionActivationPolicy::InterruptSpecificWait,
1953 )
1954 .await
1955 .unwrap();
1956 let candidate = inbox
1957 .inspect_wake_readiness("parent", Utc::now())
1958 .await
1959 .unwrap()
1960 .ready
1961 .unwrap();
1962 assert_eq!(candidate.coordinator_generation, newer.generation);
1963 assert_eq!(
1964 candidate.id, older.id,
1965 "a Respect sibling cannot obscure the precise wake"
1966 );
1967 assert_eq!(
1968 candidate.coordinator_message_policy,
1969 Some(SessionActivationPolicy::InterruptSpecificWait)
1970 );
1971 assert_eq!(
1972 router
1973 .request_activation("parent", older.generation)
1974 .await
1975 .unwrap(),
1976 SessionActivationDisposition::ActivationReserved
1977 );
1978 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
1979 let mut successor = router.register_run("parent", "run-2").await.unwrap();
1981 successor.begin_finalization().await;
1982 assert_eq!(successor.finish(newer.generation).await.unwrap(), None);
1983 assert_eq!(
1984 router
1985 .request_activation("parent", older.generation)
1986 .await
1987 .unwrap(),
1988 SessionActivationDisposition::ActivationCoalesced
1989 );
1990 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
1991 inbox.ack("parent", older_claim).await.unwrap();
1992 if let Some(sibling) = sibling {
1993 let claim = claims
1994 .iter()
1995 .find(|claim| claim.envelope.id == sibling.id)
1996 .unwrap();
1997 inbox.ack("parent", claim).await.unwrap();
1998 }
1999 }
2000 }
2001
2002 #[tokio::test]
2003 async fn precise_older_grant_after_terminal_inspection_reaches_one_successor() {
2004 use bamboo_domain::{
2005 Session, SessionActivationPolicy, SessionInboxLimits, SessionMessageEnvelope, Storage,
2006 };
2007 use bamboo_storage::{FileSessionInbox, SessionStoreV2};
2008 for abandon in [false, true] {
2009 let temp = tempfile::TempDir::new().unwrap();
2010 let store = Arc::new(SessionStoreV2::new(temp.path().into()).await.unwrap());
2011 store
2012 .save_session(&Session::new("parent", "model"))
2013 .await
2014 .unwrap();
2015 let inbox = Arc::new(FileSessionInbox::new(store, SessionInboxLimits::default()));
2016 let older = inbox
2017 .deliver(&SessionMessageEnvelope::user_input(
2018 "parent",
2019 "staged child",
2020 ))
2021 .await
2022 .unwrap();
2023 let newer = inbox
2024 .deliver_with_activation_intent(
2025 &SessionMessageEnvelope::user_input("parent", "newer user input"),
2026 SessionActivationPolicy::InterruptSpecificWait,
2027 None,
2028 )
2029 .await
2030 .unwrap();
2031 let router = SessionActivationRouter::new();
2032 router.set_inbox(inbox.clone());
2033 let spawner = spawner();
2034 router.set_spawner(spawner.clone()).await;
2035 router
2036 .request_activation("parent", newer.generation)
2037 .await
2038 .unwrap();
2039 let claims = inbox.claim("parent", 128).await.unwrap();
2040 assert_eq!(claims.len(), 1);
2041 inbox.ack("parent", &claims[0]).await.unwrap();
2042
2043 let mut owner = router.register_run("parent", "run-1").await.unwrap();
2044 owner.begin_finalization().await;
2045 let entered = Arc::new(Notify::new());
2046 let release = Arc::new(Notify::new());
2047 *router.terminal_inspection_pause.lock().await =
2048 Some((entered.clone(), release.clone()));
2049 let finalizing = tokio::spawn(async move {
2050 if abandon {
2051 owner.abandon().await;
2052 None
2053 } else {
2054 owner.finish(newer.generation).await.unwrap()
2055 }
2056 });
2057 entered.notified().await;
2058 assert!(!inbox.inspect("parent").await.unwrap().activation_pending());
2059 inbox
2060 .mark_message_activation_eligible(
2061 "parent",
2062 &older,
2063 SessionActivationPolicy::InterruptSpecificWait,
2064 )
2065 .await
2066 .unwrap();
2067 assert_eq!(
2068 router
2069 .request_activation("parent", older.generation)
2070 .await
2071 .unwrap(),
2072 SessionActivationDisposition::ActivationCoalesced
2073 );
2074 release.notify_one();
2075 let outcome = tokio::time::timeout(std::time::Duration::from_secs(2), finalizing)
2076 .await
2077 .unwrap()
2078 .unwrap();
2079 assert_eq!(
2080 outcome,
2081 (!abandon).then_some(SessionActivationDisposition::ActivationReserved)
2082 );
2083 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
2084 assert!(router.owns_run("parent", "run-2").await);
2085
2086 let claims = inbox.claim("parent", 128).await.unwrap();
2089 assert_eq!(claims.len(), 1);
2090 let mut successor = router.register_run("parent", "run-2").await.unwrap();
2091 successor.begin_finalization().await;
2092 assert_eq!(successor.finish(2).await.unwrap(), None);
2093 assert_eq!(
2094 router
2095 .request_activation("parent", older.generation)
2096 .await
2097 .unwrap(),
2098 SessionActivationDisposition::ActivationCoalesced
2099 );
2100 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
2101 inbox.ack("parent", &claims[0]).await.unwrap();
2102 }
2103 }
2104
2105 struct RecordingSpawner {
2106 reservations: AtomicUsize,
2107 launches: Arc<AtomicUsize>,
2108 entered: Option<Arc<Barrier>>,
2109 release: Option<Arc<Notify>>,
2110 }
2111
2112 #[async_trait]
2113 impl SessionActivationSpawner for RecordingSpawner {
2114 async fn reserve_activation(
2115 &self,
2116 _target_session_id: &str,
2117 _inbox_generation: u64,
2118 ) -> Result<SessionActivationReserveOutcome, SessionActivationError> {
2119 let ordinal = self.reservations.fetch_add(1, Ordering::SeqCst) + 1;
2120 if let Some(entered) = &self.entered {
2121 entered.wait().await;
2122 }
2123 if let Some(release) = &self.release {
2124 release.notified().await;
2125 }
2126 let launches = self.launches.clone();
2127 Ok(SessionActivationReserveOutcome::Reserved(
2128 SessionActivationLaunch::new(format!("run-{ordinal}"), move || {
2129 launches.fetch_add(1, Ordering::SeqCst);
2130 }),
2131 ))
2132 }
2133 }
2134
2135 fn spawner() -> Arc<RecordingSpawner> {
2136 Arc::new(RecordingSpawner {
2137 reservations: AtomicUsize::new(0),
2138 launches: Arc::new(AtomicUsize::new(0)),
2139 entered: None,
2140 release: None,
2141 })
2142 }
2143
2144 struct BlockingInspectInbox {
2145 entered: Arc<Notify>,
2146 release: Arc<Notify>,
2147 drained: Arc<std::sync::atomic::AtomicBool>,
2148 fail_once: Option<Arc<std::sync::atomic::AtomicBool>>,
2149 }
2150
2151 #[async_trait]
2152 impl SessionInboxPort for BlockingInspectInbox {
2153 async fn inspect_wake_readiness(
2154 &self,
2155 _target_session_id: &str,
2156 _now: DateTime<Utc>,
2157 ) -> Result<bamboo_domain::SessionInboxWakeReadiness, bamboo_domain::SessionInboxError>
2158 {
2159 Ok(bamboo_domain::SessionInboxWakeReadiness::default())
2162 }
2163
2164 async fn coordinator_activation_generation(
2165 &self,
2166 _target_session_id: &str,
2167 ) -> Result<u64, bamboo_domain::SessionInboxError> {
2168 Ok(1)
2169 }
2170
2171 async fn deliver(
2172 &self,
2173 _envelope: &bamboo_domain::SessionMessageEnvelope,
2174 ) -> Result<bamboo_domain::SessionInboxReceipt, bamboo_domain::SessionInboxError> {
2175 unreachable!("router cleanup only inspects the durable inbox")
2176 }
2177
2178 async fn mark_activation_eligible(
2179 &self,
2180 _target_session_id: &str,
2181 _generation: u64,
2182 _policy: bamboo_domain::SessionActivationPolicy,
2183 ) -> Result<(), bamboo_domain::SessionInboxError> {
2184 unreachable!("router cleanup only inspects the durable inbox")
2185 }
2186
2187 async fn claim(
2188 &self,
2189 _target_session_id: &str,
2190 _limit: usize,
2191 ) -> Result<Vec<bamboo_domain::SessionInboxClaim>, bamboo_domain::SessionInboxError>
2192 {
2193 unreachable!("router cleanup only inspects the durable inbox")
2194 }
2195
2196 async fn was_admitted(
2197 &self,
2198 _target_session_id: &str,
2199 _id: &bamboo_domain::SessionMessageId,
2200 ) -> Result<bool, bamboo_domain::SessionInboxError> {
2201 unreachable!("router cleanup only inspects the durable inbox")
2202 }
2203
2204 async fn ack(
2205 &self,
2206 _target_session_id: &str,
2207 _claim: &bamboo_domain::SessionInboxClaim,
2208 ) -> Result<(), bamboo_domain::SessionInboxError> {
2209 unreachable!("router cleanup only inspects the durable inbox")
2210 }
2211
2212 async fn inspect(
2213 &self,
2214 _target_session_id: &str,
2215 ) -> Result<bamboo_domain::SessionInboxBacklog, bamboo_domain::SessionInboxError> {
2216 if self
2217 .fail_once
2218 .as_ref()
2219 .is_some_and(|flag| flag.swap(false, Ordering::SeqCst))
2220 {
2221 return Err(bamboo_domain::SessionInboxError::Storage(
2222 "injected inspect failure".into(),
2223 ));
2224 }
2225 if self.drained.load(Ordering::SeqCst) {
2226 return Ok(bamboo_domain::SessionInboxBacklog {
2227 generation: 1,
2228 activation_generation: 1,
2229 ..Default::default()
2230 });
2231 }
2232 self.entered.notify_one();
2233 self.release.notified().await;
2234 Ok(bamboo_domain::SessionInboxBacklog {
2235 pending: 1,
2236 claimed: 0,
2237 generation: 1,
2238 activation_generation: 1,
2239 coordinator_generation: 1,
2240 interrupt_generation: 1,
2241 oldest_generation: Some(1),
2242 })
2243 }
2244 }
2245
2246 #[tokio::test]
2247 async fn rollback_completion_is_retained_before_waiter_subscribes() {
2248 let completion = RollbackCompletion::new();
2249 completion.complete();
2250 tokio::time::timeout(std::time::Duration::from_millis(100), completion.wait())
2251 .await
2252 .expect("completion-before-wait must not lose its wake");
2253 }
2254
2255 #[derive(Clone, Copy)]
2256 enum InjectedFirstOutcome {
2257 NoWork,
2258 Error,
2259 }
2260
2261 struct RetrySpawner {
2262 reservations: AtomicUsize,
2263 launches: Arc<AtomicUsize>,
2264 first_outcome: InjectedFirstOutcome,
2265 first_entered: Arc<Barrier>,
2266 release_first: Arc<Notify>,
2267 }
2268
2269 struct CancelFirstSpawner {
2270 reservations: AtomicUsize,
2271 launches: Arc<AtomicUsize>,
2272 first_entered: Arc<Notify>,
2273 }
2274
2275 struct PanicFirstSpawner {
2276 reservations: AtomicUsize,
2277 launches: Arc<AtomicUsize>,
2278 }
2279
2280 struct AlreadyRunningThenReserveSpawner {
2281 reservations: AtomicUsize,
2282 launches: Arc<AtomicUsize>,
2283 }
2284
2285 #[async_trait]
2286 impl SessionActivationSpawner for AlreadyRunningThenReserveSpawner {
2287 async fn reserve_activation(
2288 &self,
2289 _target_session_id: &str,
2290 _inbox_generation: u64,
2291 ) -> Result<SessionActivationReserveOutcome, SessionActivationError> {
2292 let ordinal = self.reservations.fetch_add(1, Ordering::SeqCst) + 1;
2293 if ordinal == 1 {
2294 return Ok(SessionActivationReserveOutcome::AlreadyRunning {
2295 run_id: "adopted-run".to_string(),
2296 });
2297 }
2298 let launches = self.launches.clone();
2299 Ok(SessionActivationReserveOutcome::Reserved(
2300 SessionActivationLaunch::new("successor-run", move || {
2301 launches.fetch_add(1, Ordering::SeqCst);
2302 }),
2303 ))
2304 }
2305 }
2306
2307 #[async_trait]
2308 impl SessionActivationSpawner for CancelFirstSpawner {
2309 async fn reserve_activation(
2310 &self,
2311 _target_session_id: &str,
2312 _inbox_generation: u64,
2313 ) -> Result<SessionActivationReserveOutcome, SessionActivationError> {
2314 let ordinal = self.reservations.fetch_add(1, Ordering::SeqCst) + 1;
2315 if ordinal == 1 {
2316 self.first_entered.notify_one();
2317 std::future::pending::<()>().await;
2318 }
2319 let launches = self.launches.clone();
2320 Ok(SessionActivationReserveOutcome::Reserved(
2321 SessionActivationLaunch::new(format!("run-{ordinal}"), move || {
2322 launches.fetch_add(1, Ordering::SeqCst);
2323 }),
2324 ))
2325 }
2326 }
2327
2328 #[async_trait]
2329 impl SessionActivationSpawner for PanicFirstSpawner {
2330 async fn reserve_activation(
2331 &self,
2332 _target_session_id: &str,
2333 _inbox_generation: u64,
2334 ) -> Result<SessionActivationReserveOutcome, SessionActivationError> {
2335 let ordinal = self.reservations.fetch_add(1, Ordering::SeqCst) + 1;
2336 if ordinal == 1 {
2337 panic!("injected first activation-spawner panic");
2338 }
2339 let launches = self.launches.clone();
2340 Ok(SessionActivationReserveOutcome::Reserved(
2341 SessionActivationLaunch::new(format!("run-{ordinal}"), move || {
2342 launches.fetch_add(1, Ordering::SeqCst);
2343 }),
2344 ))
2345 }
2346 }
2347
2348 struct RollbackSpawner {
2349 entered: Arc<Barrier>,
2350 release: Arc<Notify>,
2351 launch_ready: Arc<Notify>,
2352 rollbacks: Arc<AtomicUsize>,
2353 }
2354
2355 struct RealRegistryRollbackSpawner {
2356 runners: Arc<
2357 tokio::sync::RwLock<std::collections::HashMap<String, crate::execution::AgentRunner>>,
2358 >,
2359 senders: Arc<
2360 tokio::sync::RwLock<
2361 std::collections::HashMap<
2362 String,
2363 tokio::sync::broadcast::Sender<bamboo_agent_core::AgentEvent>,
2364 >,
2365 >,
2366 >,
2367 sender: tokio::sync::broadcast::Sender<bamboo_agent_core::AgentEvent>,
2368 reservations: AtomicUsize,
2369 first_reserved: Arc<tokio::sync::Notify>,
2370 allow_first_return: Arc<tokio::sync::Notify>,
2371 first_returning: Arc<tokio::sync::Notify>,
2372 rollback_started: Arc<tokio::sync::Notify>,
2373 allow_rollback: Arc<tokio::sync::Notify>,
2374 launches: Arc<AtomicUsize>,
2375 }
2376
2377 #[async_trait]
2378 impl SessionActivationSpawner for RealRegistryRollbackSpawner {
2379 async fn reserve_activation(
2380 &self,
2381 target_session_id: &str,
2382 _inbox_generation: u64,
2383 ) -> Result<SessionActivationReserveOutcome, SessionActivationError> {
2384 let ordinal = self.reservations.fetch_add(1, Ordering::SeqCst) + 1;
2385 match crate::execution::reserve_runner_core(
2386 &self.runners,
2387 &self.senders,
2388 target_session_id,
2389 &self.sender,
2390 )
2391 .await
2392 {
2393 crate::execution::ReserveOutcome::AlreadyRunning(run_id) => {
2394 Ok(SessionActivationReserveOutcome::AlreadyRunning { run_id })
2395 }
2396 crate::execution::ReserveOutcome::Reserved(reservation) => {
2397 let run_id = reservation.run_id.clone();
2398 if ordinal == 1 {
2399 self.first_reserved.notify_one();
2400 self.allow_first_return.notified().await;
2401 self.first_returning.notify_one();
2402 }
2403 let rollback_runners = self.runners.clone();
2404 let rollback_session_id = target_session_id.to_string();
2405 let rollback_run_id = run_id.clone();
2406 let rollback_started = self.rollback_started.clone();
2407 let allow_rollback = self.allow_rollback.clone();
2408 let launches = self.launches.clone();
2409 Ok(SessionActivationReserveOutcome::Reserved(
2410 SessionActivationLaunch::new_with_async_rollback(
2411 run_id,
2412 move || {
2413 launches.fetch_add(1, Ordering::SeqCst);
2414 },
2415 move || async move {
2416 rollback_started.notify_one();
2417 allow_rollback.notified().await;
2418 let mut runners = rollback_runners.write().await;
2419 if runners
2420 .get(&rollback_session_id)
2421 .is_some_and(|runner| runner.run_id == rollback_run_id)
2422 {
2423 runners.remove(&rollback_session_id);
2424 }
2425 },
2426 ),
2427 ))
2428 }
2429 }
2430 }
2431 }
2432
2433 #[async_trait]
2434 impl SessionActivationSpawner for RollbackSpawner {
2435 async fn reserve_activation(
2436 &self,
2437 _target_session_id: &str,
2438 _inbox_generation: u64,
2439 ) -> Result<SessionActivationReserveOutcome, SessionActivationError> {
2440 self.entered.wait().await;
2441 self.release.notified().await;
2442 let rollbacks = self.rollbacks.clone();
2443 let outcome = SessionActivationReserveOutcome::Reserved(
2444 SessionActivationLaunch::new_with_rollback(
2445 "reserved-run",
2446 || {},
2447 move || {
2448 rollbacks.fetch_add(1, Ordering::SeqCst);
2449 },
2450 ),
2451 );
2452 self.launch_ready.notify_one();
2453 Ok(outcome)
2454 }
2455 }
2456
2457 #[async_trait]
2458 impl SessionActivationSpawner for RetrySpawner {
2459 async fn reserve_activation(
2460 &self,
2461 _target_session_id: &str,
2462 _inbox_generation: u64,
2463 ) -> Result<SessionActivationReserveOutcome, SessionActivationError> {
2464 let ordinal = self.reservations.fetch_add(1, Ordering::SeqCst) + 1;
2465 if ordinal == 1 {
2466 self.first_entered.wait().await;
2467 self.release_first.notified().await;
2468 return match self.first_outcome {
2469 InjectedFirstOutcome::NoWork => Ok(SessionActivationReserveOutcome::NoWork),
2470 InjectedFirstOutcome::Error => Err(SessionActivationError::Internal(
2471 "injected first reservation failure".to_string(),
2472 )),
2473 };
2474 }
2475 let launches = self.launches.clone();
2476 Ok(SessionActivationReserveOutcome::Reserved(
2477 SessionActivationLaunch::new(format!("run-{ordinal}"), move || {
2478 launches.fetch_add(1, Ordering::SeqCst);
2479 }),
2480 ))
2481 }
2482 }
2483
2484 async fn assert_newer_generation_retries_after(
2485 first_outcome: InjectedFirstOutcome,
2486 ) -> Result<SessionActivationDisposition, SessionActivationError> {
2487 let router = SessionActivationRouter::new();
2488 let first_entered = Arc::new(Barrier::new(2));
2489 let release_first = Arc::new(Notify::new());
2490 let spawner = Arc::new(RetrySpawner {
2491 reservations: AtomicUsize::new(0),
2492 launches: Arc::new(AtomicUsize::new(0)),
2493 first_outcome,
2494 first_entered: first_entered.clone(),
2495 release_first: release_first.clone(),
2496 });
2497 router.set_spawner(spawner.clone()).await;
2498
2499 let first_router = router.clone();
2500 let first =
2501 tokio::spawn(async move { first_router.request_activation("session", 1).await });
2502 first_entered.wait().await;
2503
2504 let second_router = router.clone();
2505 let second =
2506 tokio::spawn(async move { second_router.request_activation("session", 2).await });
2507 tokio::task::yield_now().await;
2508 assert_eq!(
2509 spawner.reservations.load(Ordering::SeqCst),
2510 1,
2511 "generation 2 must coalesce while generation 1 owns the reservation"
2512 );
2513
2514 release_first.notify_one();
2515 let first_result = first.await.unwrap();
2516 assert_eq!(
2517 second.await.unwrap().unwrap(),
2518 SessionActivationDisposition::ActivationReserved,
2519 "the coalesced newer delivery must take the released reservation"
2520 );
2521 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
2522 assert_eq!(spawner.launches.load(Ordering::SeqCst), 1);
2523 first_result
2524 }
2525
2526 #[tokio::test]
2527 async fn active_owner_is_notified_without_second_reservation() {
2528 let router = SessionActivationRouter::new();
2529 let spawner = spawner();
2530 router.set_spawner(spawner.clone()).await;
2531 let mut registration = router.register_run("session", "run-live").await.unwrap();
2532
2533 assert_eq!(
2534 router.request_activation("session", 7).await.unwrap(),
2535 SessionActivationDisposition::ActiveNotified
2536 );
2537 registration.notifications_mut().changed().await.unwrap();
2538 assert_eq!(*registration.notifications_mut().borrow(), 7);
2539 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 0);
2540 registration.begin_finalization().await;
2541 assert_eq!(registration.finish(7).await.unwrap(), None);
2542 }
2543
2544 #[tokio::test]
2545 async fn different_run_registration_never_overwrites_live_owner() {
2546 let router = SessionActivationRouter::new();
2547 let mut first = router
2548 .register_run("session", "run-1")
2549 .await
2550 .expect("first run owns the logical session");
2551
2552 let error = router
2553 .register_run("session", "run-2")
2554 .await
2555 .err()
2556 .expect("a second exact run must collide");
2557 assert_eq!(
2558 error,
2559 SessionRunRegistrationError::OwnerCollision {
2560 target_session_id: "session".to_string(),
2561 existing_run_id: "run-1".to_string(),
2562 attempted_run_id: "run-2".to_string(),
2563 }
2564 );
2565 assert!(router.owns_run("session", "run-1").await);
2566 assert!(!router.owns_run("session", "run-2").await);
2567
2568 first.begin_finalization().await;
2569 assert_eq!(first.finish(0).await.unwrap(), None);
2570 }
2571
2572 #[tokio::test]
2573 async fn direct_registration_waits_for_activation_to_publish_its_exact_owner() {
2574 let router = SessionActivationRouter::new();
2575 let reservation_entered = Arc::new(Barrier::new(2));
2576 let allow_reservation = Arc::new(Notify::new());
2577 let launches = Arc::new(AtomicUsize::new(0));
2578 let spawner = Arc::new(RecordingSpawner {
2579 reservations: AtomicUsize::new(0),
2580 launches: launches.clone(),
2581 entered: Some(reservation_entered.clone()),
2582 release: Some(allow_reservation.clone()),
2583 });
2584 router.set_spawner(spawner.clone()).await;
2585
2586 let activation_router = router.clone();
2587 let activation =
2588 tokio::spawn(async move { activation_router.request_activation("session", 1).await });
2589 reservation_entered.wait().await;
2590
2591 let direct_router = router.clone();
2592 let direct =
2593 tokio::spawn(async move { direct_router.register_run("session", "manual-run").await });
2594 tokio::task::yield_now().await;
2595 assert!(
2596 !direct.is_finished(),
2597 "manual registration must not supersede an in-flight activation token"
2598 );
2599 assert!(!router.owns_run("session", "manual-run").await);
2600
2601 allow_reservation.notify_one();
2602 assert_eq!(
2603 activation.await.unwrap().unwrap(),
2604 SessionActivationDisposition::ActivationReserved
2605 );
2606 let collision = match direct.await.unwrap() {
2607 Ok(_) => panic!("manual registration must collide with the published activation"),
2608 Err(error) => error,
2609 };
2610 assert_eq!(
2611 collision,
2612 SessionRunRegistrationError::OwnerCollision {
2613 target_session_id: "session".to_string(),
2614 existing_run_id: "run-1".to_string(),
2615 attempted_run_id: "manual-run".to_string(),
2616 }
2617 );
2618 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 1);
2619 assert_eq!(launches.load(Ordering::SeqCst), 1);
2620 assert!(router.owns_run("session", "run-1").await);
2621 assert!(!router.owns_run("session", "manual-run").await);
2622
2623 let mut registration = router.register_run("session", "run-1").await.unwrap();
2624 registration.begin_finalization().await;
2625 assert_eq!(registration.finish(1).await.unwrap(), None);
2626 }
2627
2628 #[tokio::test]
2629 async fn cancelled_explicit_abandon_still_finishes_owned_registration_cleanup() {
2630 let router = SessionActivationRouter::new();
2631 let cleanup_entered = Arc::new(Notify::new());
2632 let cleanup_release = Arc::new(Notify::new());
2633 let cleanup_completed = Arc::new(AtomicUsize::new(0));
2634 let mut registration = router.register_run("session", "run-abandon").await.unwrap();
2635 let entered = cleanup_entered.clone();
2636 let release = cleanup_release.clone();
2637 let completed = cleanup_completed.clone();
2638 registration.set_abort_cleanup(move || async move {
2639 entered.notify_one();
2640 release.notified().await;
2641 completed.fetch_add(1, Ordering::SeqCst);
2642 });
2643
2644 let abandon = tokio::spawn(registration.abandon());
2645 tokio::time::timeout(
2646 std::time::Duration::from_secs(1),
2647 cleanup_entered.notified(),
2648 )
2649 .await
2650 .expect("explicit abandon must start its owned cleanup");
2651 abandon.abort();
2652 assert!(abandon.await.unwrap_err().is_cancelled());
2653 cleanup_release.notify_one();
2654
2655 tokio::time::timeout(std::time::Duration::from_secs(1), async {
2656 loop {
2657 if cleanup_completed.load(Ordering::SeqCst) == 1
2658 && !router.owns_run("session", "run-abandon").await
2659 {
2660 break;
2661 }
2662 tokio::task::yield_now().await;
2663 }
2664 })
2665 .await
2666 .expect("caller cancellation must not cancel detached registration cleanup");
2667 }
2668
2669 #[tokio::test]
2670 async fn delayed_same_run_cannot_adopt_an_owner_during_abandoned_cleanup() {
2671 let router = SessionActivationRouter::new();
2672 let spawner = spawner();
2673 router.set_spawner(spawner.clone()).await;
2674 let inspect_entered = Arc::new(Notify::new());
2675 let inspect_release = Arc::new(Notify::new());
2676 let drained = Arc::new(std::sync::atomic::AtomicBool::new(false));
2677 router.set_inbox(Arc::new(BlockingInspectInbox {
2678 entered: inspect_entered.clone(),
2679 release: inspect_release.clone(),
2680 drained: drained.clone(),
2681 fail_once: None,
2682 }));
2683
2684 let registration = router
2685 .register_run("session", "run-abandoned")
2686 .await
2687 .unwrap();
2688 drop(registration);
2689 tokio::time::timeout(
2690 std::time::Duration::from_secs(1),
2691 inspect_entered.notified(),
2692 )
2693 .await
2694 .expect("abandoned cleanup must reach durable inbox reconciliation");
2695
2696 let error = router
2697 .register_run("session", "run-abandoned")
2698 .await
2699 .err()
2700 .expect("a finalizing abandoned owner is not an adoptable launch placeholder");
2701 assert_eq!(error.existing_run_id(), "run-abandoned");
2702
2703 inspect_release.notify_one();
2704 tokio::time::timeout(std::time::Duration::from_secs(1), async {
2705 while spawner.launches.load(Ordering::SeqCst) == 0 {
2706 tokio::task::yield_now().await;
2707 }
2708 })
2709 .await
2710 .expect("durable pending work must launch one successor");
2711 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 1);
2712 assert_eq!(spawner.launches.load(Ordering::SeqCst), 1);
2713 assert!(!router.owns_run("session", "run-abandoned").await);
2714
2715 let mut successor = router.register_run("session", "run-1").await.unwrap();
2716 drained.store(true, Ordering::SeqCst);
2717 successor.begin_finalization().await;
2718 assert_eq!(successor.finish(1).await.unwrap(), None);
2719 }
2720
2721 #[tokio::test]
2722 async fn delayed_registry_owner_collision_rolls_back_only_its_exact_runner() {
2723 let router = SessionActivationRouter::new();
2724 let runners = Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new()));
2725 let senders = Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new()));
2726 let (sender, _receiver) = tokio::sync::broadcast::channel(8);
2727 let reserved =
2728 match crate::execution::reserve_runner_core(&runners, &senders, "session", &sender)
2729 .await
2730 {
2731 crate::execution::ReserveOutcome::Reserved(reservation) => reservation,
2732 crate::execution::ReserveOutcome::AlreadyRunning(_) => {
2733 panic!("fixture must reserve a fresh server runner")
2734 }
2735 };
2736
2737 let mut direct = router.register_run("session", "sdk-direct").await.unwrap();
2741 let collision = router
2742 .register_run("session", &reserved.run_id)
2743 .await
2744 .err()
2745 .expect("delayed server registration must collide");
2746 let collision_result = Err(bamboo_agent_core::AgentError::LLM(collision.to_string()));
2747 assert!(
2748 crate::execution::finalize_runner_exact(
2749 &runners,
2750 "session",
2751 &reserved.run_id,
2752 &collision_result,
2753 )
2754 .await
2755 );
2756 assert!(router.owns_run("session", "sdk-direct").await);
2757 assert!(matches!(
2758 runners
2759 .read()
2760 .await
2761 .get("session")
2762 .map(|runner| &runner.status),
2763 Some(crate::execution::AgentStatus::Error(_))
2764 ));
2765
2766 direct.begin_finalization().await;
2767 assert_eq!(direct.finish(0).await.unwrap(), None);
2768 }
2769
2770 #[tokio::test]
2771 async fn concurrent_idle_deliveries_coalesce_into_one_launch() {
2772 let router = SessionActivationRouter::new();
2773 let entered = Arc::new(Barrier::new(2));
2774 let release = Arc::new(Notify::new());
2775 let spawner = Arc::new(RecordingSpawner {
2776 reservations: AtomicUsize::new(0),
2777 launches: Arc::new(AtomicUsize::new(0)),
2778 entered: Some(entered.clone()),
2779 release: Some(release.clone()),
2780 });
2781 router.set_spawner(spawner.clone()).await;
2782
2783 let first_router = router.clone();
2784 let first =
2785 tokio::spawn(async move { first_router.request_activation("session", 1).await });
2786 entered.wait().await;
2787 let second_router = router.clone();
2788 let second =
2789 tokio::spawn(async move { second_router.request_activation("session", 2).await });
2790 release.notify_one();
2791 assert_eq!(
2792 first.await.unwrap().unwrap(),
2793 SessionActivationDisposition::ActivationReserved
2794 );
2795 assert_eq!(
2796 second.await.unwrap().unwrap(),
2797 SessionActivationDisposition::ActiveNotified
2798 );
2799 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 1);
2800 assert_eq!(spawner.launches.load(Ordering::SeqCst), 1);
2801 }
2802
2803 #[tokio::test]
2804 async fn newer_generation_re_reserves_after_coalesced_no_work() {
2805 assert_eq!(
2806 assert_newer_generation_retries_after(InjectedFirstOutcome::NoWork)
2807 .await
2808 .unwrap(),
2809 SessionActivationDisposition::ActivationCoalesced
2810 );
2811 }
2812
2813 #[tokio::test]
2814 async fn newer_generation_re_reserves_after_coalesced_error() {
2815 let error = assert_newer_generation_retries_after(InjectedFirstOutcome::Error)
2816 .await
2817 .unwrap_err();
2818 assert!(matches!(error, SessionActivationError::Internal(_)));
2819 }
2820
2821 async fn assert_cancelled_finalization_recovers_without_redelivery(registered_owner: bool) {
2822 let router = SessionActivationRouter::new();
2823 let first_entered = Arc::new(Notify::new());
2824 let spawner = Arc::new(CancelFirstSpawner {
2825 reservations: AtomicUsize::new(0),
2826 launches: Arc::new(AtomicUsize::new(0)),
2827 first_entered: first_entered.clone(),
2828 });
2829 router.set_spawner(spawner.clone()).await;
2830
2831 let finish = if registered_owner {
2832 let mut registration = router.register_run("session", "run-old").await.unwrap();
2833 registration.begin_finalization().await;
2834 assert_eq!(
2835 router.request_activation("session", 1).await.unwrap(),
2836 SessionActivationDisposition::ActivationCoalesced
2837 );
2838 tokio::spawn(async move { registration.finish(0).await })
2839 } else {
2840 router.begin_finalization("session", "run-old").await;
2844 assert_eq!(
2845 router.request_activation("session", 1).await.unwrap(),
2846 SessionActivationDisposition::ActivationCoalesced
2847 );
2848 let finish_router = router.clone();
2849 tokio::spawn(async move {
2850 finish_router
2851 .finish_finalization("session", "run-old", 0)
2852 .await
2853 })
2854 };
2855 first_entered.notified().await;
2856 finish.abort();
2857 assert!(finish.await.unwrap_err().is_cancelled());
2858
2859 tokio::time::timeout(std::time::Duration::from_secs(1), async {
2860 while spawner.launches.load(Ordering::SeqCst) == 0 {
2861 tokio::task::yield_now().await;
2862 }
2863 })
2864 .await
2865 .expect("cancelled finalization must retry pending work without redelivery");
2866 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
2867 assert_eq!(spawner.launches.load(Ordering::SeqCst), 1);
2868 assert!(router.owns_run("session", "run-2").await);
2869
2870 let mut successor = router.register_run("session", "run-2").await.unwrap();
2871 successor.begin_finalization().await;
2872 assert_eq!(successor.finish(1).await.unwrap(), None);
2873 }
2874
2875 #[tokio::test]
2876 async fn cancelled_registered_finalization_retries_without_another_delivery() {
2877 assert_cancelled_finalization_recovers_without_redelivery(true).await;
2878 }
2879
2880 #[tokio::test]
2881 async fn cancelled_reserved_child_finalization_retries_without_another_delivery() {
2882 assert_cancelled_finalization_recovers_without_redelivery(false).await;
2883 }
2884
2885 #[tokio::test]
2886 async fn panicking_finalization_spawner_gets_one_bounded_retry() {
2887 let router = SessionActivationRouter::new();
2888 let spawner = Arc::new(PanicFirstSpawner {
2889 reservations: AtomicUsize::new(0),
2890 launches: Arc::new(AtomicUsize::new(0)),
2891 });
2892 router.set_spawner(spawner.clone()).await;
2893 router.begin_finalization("session", "run-old").await;
2894 assert_eq!(
2895 router.request_activation("session", 1).await.unwrap(),
2896 SessionActivationDisposition::ActivationCoalesced
2897 );
2898
2899 let finish_router = router.clone();
2900 let finish = tokio::spawn(async move {
2901 finish_router
2902 .finish_finalization("session", "run-old", 0)
2903 .await
2904 });
2905 assert!(finish.await.unwrap_err().is_panic());
2906
2907 tokio::time::timeout(std::time::Duration::from_secs(1), async {
2908 while spawner.launches.load(Ordering::SeqCst) == 0 {
2909 tokio::task::yield_now().await;
2910 }
2911 })
2912 .await
2913 .expect("router Drop recovery must survive one spawner panic");
2914 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
2915 assert_eq!(spawner.launches.load(Ordering::SeqCst), 1);
2916
2917 let mut successor = router.register_run("session", "run-2").await.unwrap();
2918 successor.begin_finalization().await;
2919 assert_eq!(successor.finish(1).await.unwrap(), None);
2920 }
2921
2922 #[tokio::test]
2923 async fn cancelled_reservation_releases_token_and_newer_generation_launches() {
2924 let router = SessionActivationRouter::new();
2925 let first_entered = Arc::new(Notify::new());
2926 let spawner = Arc::new(CancelFirstSpawner {
2927 reservations: AtomicUsize::new(0),
2928 launches: Arc::new(AtomicUsize::new(0)),
2929 first_entered: first_entered.clone(),
2930 });
2931 router.set_spawner(spawner.clone()).await;
2932
2933 let first_router = router.clone();
2934 let first =
2935 tokio::spawn(async move { first_router.request_activation("session", 1).await });
2936 first_entered.notified().await;
2937 let second_router = router.clone();
2938 let second =
2939 tokio::spawn(async move { second_router.request_activation("session", 2).await });
2940 tokio::task::yield_now().await;
2941 first.abort();
2942 assert!(first.await.unwrap_err().is_cancelled());
2943
2944 assert_eq!(
2945 tokio::time::timeout(std::time::Duration::from_secs(1), second)
2946 .await
2947 .expect("coalesced caller must not hang behind a cancelled reservation")
2948 .unwrap()
2949 .unwrap(),
2950 SessionActivationDisposition::ActivationReserved
2951 );
2952 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
2953 assert_eq!(spawner.launches.load(Ordering::SeqCst), 1);
2954 }
2955
2956 #[tokio::test]
2957 async fn cancellation_after_external_reservation_rolls_back_unlaunched_slot() {
2958 let router = SessionActivationRouter::new();
2959 let entered = Arc::new(Barrier::new(2));
2960 let release = Arc::new(Notify::new());
2961 let launch_ready = Arc::new(Notify::new());
2962 let rollbacks = Arc::new(AtomicUsize::new(0));
2963 let spawner = Arc::new(RollbackSpawner {
2964 entered: entered.clone(),
2965 release: release.clone(),
2966 launch_ready: launch_ready.clone(),
2967 rollbacks: rollbacks.clone(),
2968 });
2969 router.set_spawner(spawner).await;
2970
2971 let task_router = router.clone();
2972 let task = tokio::spawn(async move { task_router.request_activation("session", 1).await });
2973 entered.wait().await;
2974 let states = router.states.lock().await;
2978 release.notify_one();
2979 launch_ready.notified().await;
2980 task.abort();
2981 assert!(task.await.unwrap_err().is_cancelled());
2982 drop(states);
2983
2984 tokio::time::timeout(std::time::Duration::from_secs(1), async {
2985 while rollbacks.load(Ordering::SeqCst) == 0 {
2986 tokio::task::yield_now().await;
2987 }
2988 })
2989 .await
2990 .expect("unlaunched reservation rollback");
2991 assert_eq!(rollbacks.load(Ordering::SeqCst), 1);
2992 }
2993
2994 #[tokio::test]
2995 async fn redelivery_waits_for_real_runner_rollback_before_reserving() {
2996 let router = SessionActivationRouter::new();
2997 let runners = Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new()));
2998 let senders = Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new()));
2999 let (sender, _receiver) = tokio::sync::broadcast::channel(8);
3000 let first_reserved = Arc::new(tokio::sync::Notify::new());
3001 let allow_first_return = Arc::new(tokio::sync::Notify::new());
3002 let rollback_started = Arc::new(tokio::sync::Notify::new());
3003 let first_returning = Arc::new(tokio::sync::Notify::new());
3004 let allow_rollback = Arc::new(tokio::sync::Notify::new());
3005 let launches = Arc::new(AtomicUsize::new(0));
3006 let spawner = Arc::new(RealRegistryRollbackSpawner {
3007 runners: runners.clone(),
3008 senders,
3009 sender,
3010 reservations: AtomicUsize::new(0),
3011 first_reserved: first_reserved.clone(),
3012 allow_first_return: allow_first_return.clone(),
3013 first_returning: first_returning.clone(),
3014 rollback_started: rollback_started.clone(),
3015 allow_rollback: allow_rollback.clone(),
3016 launches: launches.clone(),
3017 });
3018 router.set_spawner(spawner.clone()).await;
3019
3020 let first_router = router.clone();
3021 let first =
3022 tokio::spawn(async move { first_router.request_activation("session", 1).await });
3023 first_reserved.notified().await;
3024
3025 let state_guard = router.states.lock().await;
3028 allow_first_return.notify_one();
3029 first_returning.notified().await;
3030 first.abort();
3031 assert!(first.await.unwrap_err().is_cancelled());
3032 drop(state_guard);
3033 rollback_started.notified().await;
3034
3035 let second_router = router.clone();
3036 let second =
3037 tokio::spawn(async move { second_router.request_activation("session", 2).await });
3038 tokio::task::yield_now().await;
3039 assert_eq!(
3040 spawner.reservations.load(Ordering::SeqCst),
3041 1,
3042 "redelivery must remain coalesced until exact external rollback completes"
3043 );
3044
3045 allow_rollback.notify_one();
3046 assert_eq!(
3047 tokio::time::timeout(std::time::Duration::from_secs(1), second)
3048 .await
3049 .expect("redelivery must reserve after rollback")
3050 .unwrap()
3051 .unwrap(),
3052 SessionActivationDisposition::ActivationReserved
3053 );
3054 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
3055 assert_eq!(launches.load(Ordering::SeqCst), 1);
3056 let current = runners
3057 .read()
3058 .await
3059 .get("session")
3060 .map(|runner| runner.run_id.clone())
3061 .expect("second real registry reservation remains");
3062 assert_ne!(
3063 current, "",
3064 "the redelivery must never adopt an empty/stale run id"
3065 );
3066 }
3067
3068 #[tokio::test]
3069 async fn delivery_during_finalization_launches_one_successor() {
3070 let router = SessionActivationRouter::new();
3071 let spawner = spawner();
3072 router.set_spawner(spawner.clone()).await;
3073 let mut registration = router.register_run("session", "run-old").await.unwrap();
3074 registration.begin_finalization().await;
3075
3076 assert_eq!(
3077 router.request_activation("session", 11).await.unwrap(),
3078 SessionActivationDisposition::ActivationCoalesced
3079 );
3080 let disposition = registration.finish(10).await.unwrap();
3081 assert_eq!(
3082 disposition,
3083 Some(SessionActivationDisposition::ActivationReserved)
3084 );
3085 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 1);
3086 assert_eq!(spawner.launches.load(Ordering::SeqCst), 1);
3087 }
3088
3089 #[tokio::test]
3090 async fn finalization_does_not_restart_when_cursor_caught_up() {
3091 let router = SessionActivationRouter::new();
3092 let spawner = spawner();
3093 router.set_spawner(spawner.clone()).await;
3094 let mut registration = router.register_run("session", "run-old").await.unwrap();
3095 assert_eq!(
3096 router.request_activation("session", 4).await.unwrap(),
3097 SessionActivationDisposition::ActiveNotified
3098 );
3099 registration.begin_finalization().await;
3100 assert_eq!(registration.finish(4).await.unwrap(), None);
3101 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 0);
3102 assert!(
3103 router.states.lock().await.get("session").is_none(),
3104 "a caught-up completed session must not leak routing state"
3105 );
3106 }
3107
3108 #[tokio::test]
3109 async fn successor_finalization_removes_state_after_raced_generation_catches_up() {
3110 let router = SessionActivationRouter::new();
3111 let spawner = spawner();
3112 router.set_spawner(spawner).await;
3113 let mut registration = router.register_run("session", "run-old").await.unwrap();
3114 registration.begin_finalization().await;
3115
3116 assert_eq!(
3117 router.request_activation("session", 11).await.unwrap(),
3118 SessionActivationDisposition::ActivationCoalesced
3119 );
3120 assert_eq!(
3121 registration.finish(10).await.unwrap(),
3122 Some(SessionActivationDisposition::ActivationReserved)
3123 );
3124 let mut successor_registration = router.register_run("session", "run-1").await.unwrap();
3125 successor_registration.begin_finalization().await;
3126 assert_eq!(successor_registration.finish(11).await.unwrap(), None);
3127 assert!(
3128 router.states.lock().await.get("session").is_none(),
3129 "the successor's caught-up terminal state must be compacted"
3130 );
3131 }
3132
3133 #[tokio::test]
3134 async fn failed_terminal_inspection_keeps_the_exact_owner_cleanup_armed() {
3135 let router = SessionActivationRouter::new();
3136 let spawner = spawner();
3137 router.set_spawner(spawner.clone()).await;
3138 let inspect_release = Arc::new(Notify::new());
3139 inspect_release.notify_one();
3140 router.set_inbox(Arc::new(BlockingInspectInbox {
3141 entered: Arc::new(Notify::new()),
3142 release: inspect_release,
3143 drained: Arc::new(std::sync::atomic::AtomicBool::new(false)),
3144 fail_once: Some(Arc::new(std::sync::atomic::AtomicBool::new(true))),
3145 }));
3146 let mut registration = router
3147 .register_run("session", "failed-inspect-run")
3148 .await
3149 .unwrap();
3150 router.request_activation("session", 1).await.unwrap();
3151 registration.begin_finalization().await;
3152 assert!(registration.finish(1).await.is_err());
3153 tokio::time::timeout(std::time::Duration::from_secs(1), async {
3154 while spawner.launches.load(Ordering::SeqCst) == 0 {
3155 tokio::task::yield_now().await;
3156 }
3157 })
3158 .await
3159 .expect("failed inspection must not leave the old run owning the inbox");
3160 assert!(!router.owns_run("session", "failed-inspect-run").await);
3161 assert_eq!(spawner.launches.load(Ordering::SeqCst), 1);
3162 }
3163
3164 #[tokio::test]
3165 async fn poison_generation_launches_only_one_successor_until_new_work_arrives() {
3166 let router = SessionActivationRouter::new();
3167 let spawner = spawner();
3168 router.set_spawner(spawner.clone()).await;
3169 let inspect_release = Arc::new(Notify::new());
3170 router.set_inbox(Arc::new(BlockingInspectInbox {
3171 entered: Arc::new(Notify::new()),
3172 release: inspect_release.clone(),
3173 drained: Arc::new(std::sync::atomic::AtomicBool::new(false)),
3174 fail_once: None,
3175 }));
3176 let mut registration = router
3177 .register_run("session", "run-original")
3178 .await
3179 .unwrap();
3180 assert_eq!(
3181 router.request_activation("session", 7).await.unwrap(),
3182 SessionActivationDisposition::ActiveNotified
3183 );
3184 registration.begin_finalization().await;
3185 inspect_release.notify_one();
3186 assert_eq!(
3187 registration.finish(0).await.unwrap(),
3188 Some(SessionActivationDisposition::ActivationReserved)
3189 );
3190 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 1);
3191
3192 let mut successor_registration = router.register_run("session", "run-1").await.unwrap();
3196 successor_registration.begin_finalization().await;
3197 inspect_release.notify_one();
3198 assert_eq!(successor_registration.finish(0).await.unwrap(), None);
3199 assert_eq!(
3200 router.request_activation("session", 7).await.unwrap(),
3201 SessionActivationDisposition::ActivationCoalesced
3202 );
3203 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 1);
3204
3205 assert_eq!(
3207 router.request_activation("session", 8).await.unwrap(),
3208 SessionActivationDisposition::ActivationReserved
3209 );
3210 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
3211 }
3212
3213 #[tokio::test]
3214 async fn adopted_terminal_run_preserves_generation_for_one_real_successor() {
3215 let router = SessionActivationRouter::new();
3216 let spawner = Arc::new(AlreadyRunningThenReserveSpawner {
3217 reservations: AtomicUsize::new(0),
3218 launches: Arc::new(AtomicUsize::new(0)),
3219 });
3220 router.set_spawner(spawner.clone()).await;
3221
3222 assert_eq!(
3223 router.request_activation("session", 12).await.unwrap(),
3224 SessionActivationDisposition::ActiveNotified
3225 );
3226 router.begin_finalization("session", "adopted-run").await;
3227 assert_eq!(
3228 router
3229 .finish_finalization("session", "adopted-run", 0)
3230 .await
3231 .unwrap(),
3232 Some(SessionActivationDisposition::ActivationReserved)
3233 );
3234 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
3235 assert_eq!(spawner.launches.load(Ordering::SeqCst), 1);
3236 }
3237 #[tokio::test]
3238 async fn quarantined_registration_never_releases_owner_or_launches_successor() {
3239 let router = SessionActivationRouter::new();
3240 let spawner = spawner();
3241 router.set_spawner(spawner.clone()).await;
3242 let mut registration = router
3243 .register_run("session", "quarantined-run")
3244 .await
3245 .unwrap();
3246 let cleaned = Arc::new(std::sync::atomic::AtomicBool::new(false));
3247 let cleanup_flag = cleaned.clone();
3248 registration.set_abort_cleanup(move || async move {
3249 cleanup_flag.store(true, Ordering::SeqCst);
3250 });
3251 registration.quarantine().await;
3252 assert_eq!(
3253 router.request_activation("session", 19).await.unwrap(),
3254 SessionActivationDisposition::ActivationCoalesced
3255 );
3256 assert!(router.register_run("session", "successor").await.is_err());
3257 assert!(router
3258 .register_run("session", "quarantined-run")
3259 .await
3260 .is_err());
3261 tokio::task::yield_now().await;
3262 assert!(!cleaned.load(Ordering::SeqCst));
3263 assert_eq!(spawner.launches.load(Ordering::SeqCst), 0);
3264 let states = router.states.lock().await;
3265 let owner = states["session"].owner.as_ref().unwrap();
3266 assert_eq!(owner.run_id, "quarantined-run");
3267 assert!(owner.finalizing);
3268 }
3269
3270 #[tokio::test]
3271 async fn physical_execution_registration_drop_is_quarantined() {
3272 let router = SessionActivationRouter::new();
3273 let spawner = spawner();
3274 router.set_spawner(spawner.clone()).await;
3275 let mut registration = router
3276 .register_run("session", "physical-run")
3277 .await
3278 .unwrap();
3279 let cleaned = Arc::new(std::sync::atomic::AtomicBool::new(false));
3280 let cleanup_flag = cleaned.clone();
3281 registration.set_abort_cleanup(move || async move {
3282 cleanup_flag.store(true, Ordering::SeqCst);
3283 });
3284 registration.quarantine_on_drop();
3285 drop(registration);
3286 assert!(matches!(
3290 router.request_activation("session", 7).await.unwrap(),
3291 SessionActivationDisposition::ActiveNotified
3292 | SessionActivationDisposition::ActivationCoalesced
3293 ));
3294 tokio::time::timeout(std::time::Duration::from_secs(1), async {
3295 loop {
3296 if router.states.lock().await["session"]
3297 .owner
3298 .as_ref()
3299 .unwrap()
3300 .finalizing
3301 {
3302 break;
3303 }
3304 tokio::task::yield_now().await;
3305 }
3306 })
3307 .await
3308 .unwrap();
3309 assert_eq!(
3310 router.request_activation("session", 8).await.unwrap(),
3311 SessionActivationDisposition::ActivationCoalesced
3312 );
3313 assert!(!cleaned.load(Ordering::SeqCst));
3314 assert_eq!(spawner.launches.load(Ordering::SeqCst), 0);
3315 assert!(router.register_run("session", "successor").await.is_err());
3316 let states = router.states.lock().await;
3317 assert!(states["session"].owner.as_ref().unwrap().finalizing);
3318 }
3319
3320 #[tokio::test]
3321 async fn cancelling_quarantine_wait_does_not_arm_abandoned_cleanup() {
3322 let router = SessionActivationRouter::new();
3323 let registration = router
3324 .register_run("session", "quarantined-run")
3325 .await
3326 .unwrap();
3327 let states_guard = router.states.lock().await;
3328 let mut quarantine = Box::pin(registration.quarantine());
3329 assert!(futures::poll!(quarantine.as_mut()).is_pending());
3330 drop(quarantine);
3331 drop(states_guard);
3332 tokio::task::yield_now().await;
3333 assert!(router.register_run("session", "successor").await.is_err());
3334 assert_eq!(
3335 router.current_run_id("session").await.as_deref(),
3336 Some("quarantined-run")
3337 );
3338 }
3339 async fn selected_fixture() -> (
3340 tempfile::TempDir,
3341 Arc<SessionActivationRouter>,
3342 Arc<RecordingSpawner>,
3343 Arc<bamboo_storage::FileSessionInbox>,
3344 ) {
3345 use bamboo_domain::{
3346 Session, SessionActivationPolicy, SessionInboxLimits, SessionMessageEnvelope, Storage,
3347 };
3348 let temp = tempfile::tempdir().unwrap();
3349 let store = Arc::new(
3350 bamboo_storage::SessionStoreV2::new(temp.path().into())
3351 .await
3352 .unwrap(),
3353 );
3354 store
3355 .save_session(&Session::new("selected", "model"))
3356 .await
3357 .unwrap();
3358 let inbox = Arc::new(bamboo_storage::FileSessionInbox::new(
3359 store,
3360 SessionInboxLimits::default(),
3361 ));
3362 inbox
3363 .deliver_with_activation_intent(
3364 &SessionMessageEnvelope::user_input("selected", "other eligible input"),
3365 SessionActivationPolicy::RespectSpecificWait,
3366 None,
3367 )
3368 .await
3369 .unwrap();
3370 let router = SessionActivationRouter::new();
3371 router.set_inbox(inbox.clone());
3372 let spawner = spawner();
3373 router.set_spawner(spawner.clone()).await;
3374 (temp, router, spawner, inbox)
3375 }
3376
3377 #[tokio::test]
3378 async fn receipt_run_selected_release_never_inherits_unrequested_durable_backlog() {
3379 for mode in ["finish", "abandon", "drop"] {
3380 let (_temp, router, spawner, inbox) = selected_fixture().await;
3381 let mut selected = router
3382 .register_run("selected", "manual-selected")
3383 .await
3384 .unwrap();
3385 selected.restrict_to_selected_input();
3386 match mode {
3387 "finish" => {
3388 selected.begin_finalization().await;
3389 assert_eq!(selected.finish(0).await.unwrap(), None);
3390 }
3391 "abandon" => selected.abandon().await,
3392 _ => drop(selected),
3393 }
3394 tokio::time::timeout(std::time::Duration::from_secs(2), async {
3395 while router.current_run_id("selected").await.is_some() {
3396 tokio::task::yield_now().await;
3397 }
3398 })
3399 .await
3400 .unwrap();
3401 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 0, "{mode}");
3402 assert_eq!(spawner.launches.load(Ordering::SeqCst), 0);
3403 assert_eq!(inbox.inspect("selected").await.unwrap().pending, 1);
3404 assert_eq!(
3406 router.request_activation("selected", 1).await.unwrap(),
3407 SessionActivationDisposition::ActivationReserved
3408 );
3409 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 1);
3410 }
3411 }
3412
3413 #[tokio::test]
3414 async fn receipt_run_selected_release_preserves_explicit_coalesced_coordinator_wake() {
3415 for mode in ["finish", "abandon", "drop"] {
3416 for notify_before_scope in [false, true] {
3417 let (_temp, router, spawner, _inbox) = selected_fixture().await;
3418 let mut selected = router
3419 .register_run("selected", "manual-selected")
3420 .await
3421 .unwrap();
3422 if !notify_before_scope {
3423 selected.restrict_to_selected_input();
3424 }
3425 assert_eq!(
3426 router.request_activation("selected", 1).await.unwrap(),
3427 SessionActivationDisposition::ActiveNotified
3428 );
3429 if notify_before_scope {
3430 selected.restrict_to_selected_input();
3431 }
3432 match mode {
3433 "finish" => {
3434 selected.begin_finalization().await;
3435 assert_eq!(
3436 selected.finish(0).await.unwrap(),
3437 Some(SessionActivationDisposition::ActivationReserved)
3438 );
3439 }
3440 "abandon" => selected.abandon().await,
3441 _ => drop(selected),
3442 }
3443 tokio::time::timeout(std::time::Duration::from_secs(2), async {
3444 while spawner.reservations.load(Ordering::SeqCst) == 0 {
3445 tokio::task::yield_now().await;
3446 }
3447 })
3448 .await
3449 .unwrap();
3450 assert_eq!(
3451 spawner.reservations.load(Ordering::SeqCst),
3452 1,
3453 "{mode}/{notify_before_scope}"
3454 );
3455 assert_eq!(spawner.launches.load(Ordering::SeqCst), 1);
3456 assert!(router.owns_run("selected", "run-1").await);
3457 }
3458 }
3459 }
3460
3461 #[tokio::test]
3462 async fn receipt_run_cancelled_selected_finish_retains_restricted_drop_cleanup() {
3463 let (_temp, router, spawner, inbox) = selected_fixture().await;
3464 let mut selected = router
3465 .register_run("selected", "manual-selected")
3466 .await
3467 .unwrap();
3468 selected.restrict_to_selected_input();
3469 let lock = router.states.lock().await;
3470 let mut finish = Box::pin(selected.finish(0));
3471 assert!(futures::poll!(finish.as_mut()).is_pending());
3472 drop(finish);
3473 drop(lock);
3474 tokio::time::timeout(std::time::Duration::from_secs(2), async {
3475 while router.current_run_id("selected").await.is_some() {
3476 tokio::task::yield_now().await;
3477 }
3478 })
3479 .await
3480 .unwrap();
3481 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 0);
3482 assert_eq!(inbox.inspect("selected").await.unwrap().pending, 1);
3483 }
3484}