1use std::collections::HashMap;
8use std::fmt;
9use std::future::Future;
10use std::pin::Pin;
11use std::sync::{Arc, Mutex as StdMutex, RwLock as StdRwLock, Weak};
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 instance: Arc<()>,
205 run_id: String,
206 session_created_at: Option<DateTime<Utc>>,
208 finalizing: bool,
209 registrations: usize,
210 child_completions_need_successor: bool,
213 delivery_sink: Option<mpsc::UnboundedSender<u64>>,
217 scoped_loop: Option<SupervisorLoopGrant>,
218}
219
220#[derive(Debug, Clone)]
221struct SupervisorLoopGrant {
222 binding: bamboo_domain::ScopedLoopControlBinding,
223 execution_cancel: tokio_util::sync::CancellationToken,
224 host_stop: tokio_util::sync::CancellationToken,
225}
226impl SupervisorLoopGrant {
227 fn stopped(&self) -> bool {
228 self.execution_cancel.is_cancelled() || self.host_stop.is_cancelled()
229 }
230}
231
232struct SupervisorLoopDeliveryGuard {
233 inner: ExpectedRunDeliveryGuard,
234 grant: SupervisorLoopGrant,
235}
236impl bamboo_domain::SessionRunDeliveryGuard for SupervisorLoopDeliveryGuard {
237 fn matches(&self, target: &str, birth: DateTime<Utc>, run: &str) -> bool {
238 !self.grant.stopped()
239 && bamboo_domain::SessionRunDeliveryGuard::matches(&self.inner, target, birth, run)
240 }
241 fn notify_committed_delivery(&self, generation: u64) -> bool {
242 self.inner.notify_if(generation, || !self.grant.stopped())
243 }
244}
245
246#[derive(Debug)]
247struct TargetActivationState {
248 latest_generation: u64,
249 latest_delivery_generation: u64,
251 last_dispatched_generation: u64,
258 last_dispatched_coordinator_generation: u64,
262 owner_completion_epoch: u64,
265 last_recovery_attempt: Option<RecoveryAttempt>,
266 owner: Option<ActiveOwner>,
267 activation_reserved: bool,
268 activation_token: u64,
271 activation_epoch: watch::Sender<u64>,
276 notify: watch::Sender<u64>,
277}
278
279impl Default for TargetActivationState {
280 fn default() -> Self {
281 let (activation_epoch, _activation_receiver) = watch::channel(0);
282 let (notify, _receiver) = watch::channel(0);
283 Self {
284 latest_generation: 0,
285 latest_delivery_generation: 0,
286 last_dispatched_generation: 0,
287 last_dispatched_coordinator_generation: 0,
288 owner_completion_epoch: 0,
289 last_recovery_attempt: None,
290 owner: None,
291 activation_reserved: false,
292 activation_token: 0,
293 activation_epoch,
294 notify,
295 }
296 }
297}
298
299const INITIAL_RECOVERY_BACKOFF: Duration = Duration::seconds(30);
300const MAX_RECOVERY_BACKOFF: Duration = Duration::minutes(30);
301
302#[derive(Debug, Clone)]
303struct RecoveryAttempt {
304 candidate: SessionInboxWakeCandidate,
305 completion_epoch: u64,
306 backoff: Duration,
307 next_retry_at: DateTime<Utc>,
308 requires_completion: bool,
311}
312
313fn recovery_retry_at(
314 state: &TargetActivationState,
315 candidate: &SessionInboxWakeCandidate,
316 now: DateTime<Utc>,
317) -> Option<Option<DateTime<Utc>>> {
318 let previous = state
319 .last_recovery_attempt
320 .as_ref()
321 .filter(|attempt| attempt.candidate == *candidate)?;
322 if previous.requires_completion && state.owner_completion_epoch <= previous.completion_epoch {
323 return Some(None);
324 }
325 (now < previous.next_retry_at).then_some(Some(previous.next_retry_at))
326}
327
328fn record_recovery_attempt(
329 state: &mut TargetActivationState,
330 candidate: SessionInboxWakeCandidate,
331 now: DateTime<Utc>,
332 requires_completion: bool,
333) {
334 let backoff = state
335 .last_recovery_attempt
336 .as_ref()
337 .filter(|attempt| attempt.candidate == candidate)
338 .map_or(INITIAL_RECOVERY_BACKOFF, |attempt| {
339 (attempt.backoff * 2).min(MAX_RECOVERY_BACKOFF)
340 });
341 state.last_recovery_attempt = Some(RecoveryAttempt {
342 candidate,
343 completion_epoch: state.owner_completion_epoch,
344 backoff,
345 next_retry_at: now + backoff,
346 requires_completion,
347 });
348}
349
350#[derive(Debug, Clone, PartialEq, Eq)]
353pub struct SessionWakeReconcileResult {
354 pub disposition: Option<SessionActivationDisposition>,
355 pub next_due_at: Option<DateTime<Utc>>,
356}
357
358fn reserve_activation_token(state: &mut TargetActivationState) -> u64 {
359 state.activation_reserved = true;
360 state.activation_token = state.activation_token.wrapping_add(1).max(1);
361 state.activation_token
362}
363
364fn release_activation_token(state: &mut TargetActivationState, token: u64) -> bool {
365 if !state.activation_reserved || state.activation_token != token {
366 return false;
367 }
368 state.activation_reserved = false;
369 let next_epoch = (*state.activation_epoch.borrow()).wrapping_add(1);
370 state.activation_epoch.send_replace(next_epoch);
371 true
372}
373
374fn inbox_activation_error(error: bamboo_domain::SessionInboxError) -> SessionActivationError {
375 match error {
376 bamboo_domain::SessionInboxError::TargetNotFound(target) => {
377 SessionActivationError::TargetNotFound(target)
378 }
379 error => SessionActivationError::Internal(error.to_string()),
380 }
381}
382
383struct ActivationReservationLease {
389 router: SessionActivationRouter,
390 target_session_id: String,
391 token: u64,
392 recover_on_drop: bool,
393 armed: bool,
394 rollback_completion: Option<Arc<RollbackCompletion>>,
395}
396
397impl ActivationReservationLease {
398 fn new(
399 router: SessionActivationRouter,
400 target_session_id: &str,
401 token: u64,
402 recover_on_drop: bool,
403 ) -> Self {
404 Self {
405 router,
406 target_session_id: target_session_id.to_string(),
407 token,
408 recover_on_drop,
409 armed: true,
410 rollback_completion: None,
411 }
412 }
413
414 fn wait_for_rollback(&mut self, completion: Option<Arc<RollbackCompletion>>) {
415 self.rollback_completion = completion;
416 }
417
418 fn disarm(&mut self) {
419 self.armed = false;
420 }
421}
422
423impl Drop for ActivationReservationLease {
424 fn drop(&mut self) {
425 if !self.armed {
426 return;
427 }
428 let router = self.router.clone();
429 let target_session_id = self.target_session_id.clone();
430 let token = self.token;
431 let recover_on_drop = self.recover_on_drop;
432 let rollback_completion = self.rollback_completion.take();
433 if let Ok(runtime) = tokio::runtime::Handle::try_current() {
434 runtime.spawn(async move {
435 if let Some(completion) = rollback_completion {
436 completion.wait().await;
437 }
438 router
439 .recover_cancelled_reservation(&target_session_id, token, recover_on_drop)
440 .await;
441 });
442 }
443 }
444}
445
446struct ExpectedRunDeliveryGuard {
447 target_session_id: String,
448 session_created_at: DateTime<Utc>,
449 run_id: String,
450 states: Arc<Mutex<HashMap<String, TargetActivationState>>>,
451 _publication: tokio::sync::OwnedMutexGuard<()>,
452}
453
454impl bamboo_domain::SessionRunDeliveryGuard for ExpectedRunDeliveryGuard {
455 fn matches(
456 &self,
457 target_session_id: &str,
458 target_created_at: DateTime<Utc>,
459 expected_run_id: &str,
460 ) -> bool {
461 self.target_session_id == target_session_id
462 && self.session_created_at == target_created_at
463 && self.run_id == expected_run_id
464 }
465
466 fn notify_committed_delivery(&self, inbox_generation: u64) -> bool {
467 self.notify_if(inbox_generation, || true)
468 }
469}
470
471impl ExpectedRunDeliveryGuard {
472 fn notify_if(&self, inbox_generation: u64, allowed: impl FnOnce() -> bool) -> bool {
473 let mut states = self.states.blocking_lock();
476 let Some(state) = states.get_mut(&self.target_session_id) else {
477 return false;
478 };
479 let Some(owner) = state.owner.as_mut().filter(|owner| {
480 owner.run_id == self.run_id
481 && owner.session_created_at == Some(self.session_created_at)
482 && !owner.finalizing
483 && owner.registrations > 0
484 && allowed()
485 }) else {
486 return false;
487 };
488 state.latest_delivery_generation = state.latest_delivery_generation.max(inbox_generation);
490 state.notify.send_replace(inbox_generation);
491 if owner
492 .delivery_sink
493 .as_ref()
494 .is_some_and(|sink| sink.send(inbox_generation).is_err())
495 {
496 owner.delivery_sink = None;
497 }
498 true
499 }
500}
501
502#[derive(Clone, Default)]
504pub struct SessionActivationRouter {
505 runtime_instance: Arc<()>,
506 states: Arc<Mutex<HashMap<String, TargetActivationState>>>,
507 publication_fences: TargetPublicationFences,
510 spawner: Arc<RwLock<Option<Arc<dyn SessionActivationSpawner>>>>,
511 inbox: Arc<StdRwLock<Option<Arc<dyn SessionInboxPort>>>>,
512 #[cfg(test)]
513 terminal_inspection_pause: Arc<Mutex<Option<TerminalInspectionPause>>>,
514}
515
516type TargetPublicationFences = Arc<StdMutex<HashMap<String, Weak<Mutex<()>>>>>;
517
518#[cfg(test)]
519type TerminalInspectionPause = (Arc<tokio::sync::Notify>, Arc<tokio::sync::Notify>);
520
521type AbortCleanup =
522 Box<dyn FnOnce() -> Pin<Box<dyn Future<Output = ()> + Send + 'static>> + Send + 'static>;
523
524type ScopedTransportCleanup =
525 Arc<dyn Fn() -> Pin<Box<dyn Future<Output = std::io::Result<()>> + Send>> + Send + Sync>;
526
527struct ClosedLoopTransportGuard(ExpectedRunDeliveryGuard);
528impl bamboo_domain::SessionRunDeliveryGuard for ClosedLoopTransportGuard {
529 fn matches(&self, target: &str, birth: DateTime<Utc>, run: &str) -> bool {
530 bamboo_domain::SessionRunDeliveryGuard::matches(&self.0, target, birth, run)
531 }
532 fn notify_committed_delivery(&self, _generation: u64) -> bool {
533 false
534 }
535}
536
537pub struct SessionRunRegistration {
545 instance: Arc<()>,
546 router: Arc<SessionActivationRouter>,
547 target_session_id: String,
548 run_id: String,
549 notifications: Option<watch::Receiver<u64>>,
550 abort_cleanup: Option<AbortCleanup>,
551 scoped_transport_cleanup: Option<ScopedTransportCleanup>,
552 quarantine_on_drop: bool,
553 tool_suspension_observation: Option<bamboo_agent_core::tools::SuspensionRestriction>,
556 selected_input_wakes: Option<watch::Receiver<u64>>,
557 armed: bool,
558}
559
560#[derive(Clone)]
563pub(crate) struct SessionRunOwnership {
564 instance: Arc<()>,
565 router: Arc<SessionActivationRouter>,
566 target_session_id: String,
567 run_id: String,
568}
569
570impl SessionRunOwnership {
571 pub(crate) async fn close_scoped_transport(
574 &self,
575 binding: &bamboo_domain::ScopedLoopControlBinding,
576 ) -> std::io::Result<()> {
577 if binding.target_session_id != self.target_session_id
578 || binding.owner_run_id != self.run_id
579 {
580 return Err(std::io::Error::other(
581 "target transport proof identity differs",
582 ));
583 }
584 let _publication = self.router.lock_publication(&self.target_session_id).await;
585 let mut states = self.router.states.lock().await;
586 let owner = states
587 .get_mut(&self.target_session_id)
588 .and_then(|state| state.owner.as_mut())
589 .filter(|owner| {
590 Arc::ptr_eq(&owner.instance, &self.instance)
591 && owner.run_id == self.run_id
592 && owner.session_created_at == Some(binding.target_created_at)
593 && owner.registrations > 0
594 })
595 .ok_or_else(|| {
596 std::io::Error::other("target transport owner is no longer registered")
597 })?;
598 if owner
599 .scoped_loop
600 .as_ref()
601 .is_some_and(|grant| grant.binding.execution_id != binding.execution_id)
602 {
603 return Err(std::io::Error::other("target transport execution differs"));
604 }
605 owner.finalizing = true;
606 Ok(())
607 }
608
609 pub(crate) fn retain_closed_scoped_transport(
610 &self,
611 binding: &bamboo_domain::ScopedLoopControlBinding,
612 ) -> std::io::Result<Arc<dyn bamboo_domain::SessionRunDeliveryGuard>> {
613 if binding.target_session_id != self.target_session_id
614 || binding.owner_run_id != self.run_id
615 {
616 return Err(std::io::Error::other("closed transport identity differs"));
617 }
618 let publication = self
619 .router
620 .publication_fence(&self.target_session_id)
621 .blocking_lock_owned();
622 let states = self.router.states.blocking_lock();
623 if !states
624 .get(&self.target_session_id)
625 .and_then(|state| state.owner.as_ref())
626 .is_some_and(|owner| {
627 Arc::ptr_eq(&owner.instance, &self.instance)
628 && owner.run_id == self.run_id
629 && owner.session_created_at == Some(binding.target_created_at)
630 && owner.finalizing
631 && owner.registrations > 0
632 })
633 {
634 return Err(std::io::Error::other("closed transport owner changed"));
635 }
636 drop(states);
637 Ok(Arc::new(ClosedLoopTransportGuard(
638 ExpectedRunDeliveryGuard {
639 target_session_id: self.target_session_id.clone(),
640 session_created_at: binding.target_created_at,
641 run_id: self.run_id.clone(),
642 states: self.router.states.clone(),
643 _publication: publication,
644 },
645 )))
646 }
647
648 pub(crate) fn run_id(&self) -> &str {
649 &self.run_id
650 }
651
652 pub(crate) fn router(&self) -> &Arc<SessionActivationRouter> {
653 &self.router
654 }
655
656 pub(crate) fn matches_router(&self, router: &Arc<SessionActivationRouter>) -> bool {
657 Arc::ptr_eq(&self.router, router)
658 }
659
660 pub(crate) async fn is_current(&self, session_id: &str) -> bool {
661 self.target_session_id == session_id && self.router.owns_run(session_id, &self.run_id).await
662 }
663}
664
665impl SessionRunRegistration {
666 pub(crate) fn set_scoped_transport_cleanup<F, Fut>(&mut self, cleanup: F)
667 where
668 F: Fn() -> Fut + Send + Sync + 'static,
669 Fut: Future<Output = std::io::Result<()>> + Send + 'static,
670 {
671 self.scoped_transport_cleanup = Some(Arc::new(move || Box::pin(cleanup())));
672 }
673
674 pub(crate) fn bind_tool_suspension_observation(
675 &mut self,
676 token: bamboo_agent_core::tools::SuspensionRestriction,
677 ) {
678 self.tool_suspension_observation = Some(token);
679 }
680
681 pub(crate) fn observed_forbidden_running(&self) -> bool {
682 self.tool_suspension_observation
683 .as_ref()
684 .is_some_and(|token| token.observation().forbidden_running_seen)
685 }
686
687 pub(crate) fn ownership(&self) -> SessionRunOwnership {
688 SessionRunOwnership {
689 instance: self.instance.clone(),
690 router: self.router.clone(),
691 target_session_id: self.target_session_id.clone(),
692 run_id: self.run_id.clone(),
693 }
694 }
695
696 pub(crate) fn restrict_to_selected_input(&mut self) {
701 if self.selected_input_wakes.is_none() {
702 self.selected_input_wakes = self.notifications.clone();
703 }
704 }
705
706 pub fn notifications_mut(&mut self) -> &mut watch::Receiver<u64> {
707 self.notifications
708 .as_mut()
709 .expect("run registration notifications are live before finalization")
710 }
711
712 pub fn set_abort_cleanup<F, Fut>(&mut self, cleanup: F)
715 where
716 F: FnOnce() -> Fut + Send + 'static,
717 Fut: Future<Output = ()> + Send + 'static,
718 {
719 self.abort_cleanup = Some(Box::new(move || Box::pin(cleanup())));
720 }
721
722 pub async fn begin_finalization(&mut self) {
723 self.router
724 .begin_finalization(&self.target_session_id, &self.run_id)
725 .await;
726 }
727
728 pub(crate) fn quarantine_on_drop(&mut self) {
732 self.quarantine_on_drop = true;
733 }
734
735 pub async fn quarantine(mut self) {
743 self.armed = false;
744 self.abort_cleanup = None;
745 self.router
746 .begin_finalization(&self.target_session_id, &self.run_id)
747 .await;
748 if let Some(cleanup) = self.scoped_transport_cleanup.take() {
749 let task = tokio::spawn(async move { cleanup().await });
750 match task.await {
751 Ok(Ok(())) => {}
752 result => tracing::warn!(?result, "quarantined Loop transport cleanup incomplete"),
753 }
754 }
755 self.notifications.take();
756 }
757
758 pub async fn abandon(mut self) {
767 self.armed = false;
768 if self.selected_input_wakes.is_some() {
769 drop(self.notifications.take());
770 }
771 let scoped_transport_cleanup = self.scoped_transport_cleanup.clone();
772 let abort_cleanup = self.abort_cleanup.take();
773 let router = self.router.clone();
774 let target_session_id = self.target_session_id.clone();
775 let run_id = self.run_id.clone();
776 let cleanup_target_session_id = target_session_id.clone();
777 let cleanup_run_id = run_id.clone();
778 let selected_input_wakes = self.selected_input_wakes.take();
779 let cleanup = tokio::spawn(async move {
780 if let Some(cleanup) = scoped_transport_cleanup {
781 if let Err(error) = cleanup().await {
782 tracing::warn!(%error, "closed Loop transport cleanup failed; retaining abandoned owner");
783 return;
784 }
785 }
786 if let Some(cleanup) = abort_cleanup {
787 cleanup().await;
788 }
789 router
790 .cleanup_abandoned_registration(
791 &cleanup_target_session_id,
792 &cleanup_run_id,
793 selected_input_wakes,
794 )
795 .await;
796 });
797 if let Err(error) = cleanup.await {
798 tracing::error!(
799 %target_session_id,
800 %run_id,
801 %error,
802 "detached SessionInbox registration cleanup failed"
803 );
804 }
805 }
806
807 pub async fn finish(
810 mut self,
811 admitted_generation: u64,
812 ) -> Result<Option<SessionActivationDisposition>, SessionActivationError> {
813 if let Some(cleanup) = &self.scoped_transport_cleanup {
814 cleanup()
815 .await
816 .map_err(|error| SessionActivationError::Internal(error.to_string()))?;
817 }
818 drop(self.notifications.take());
821 let released = match self.selected_input_wakes.as_ref() {
822 Some(wakes) => {
823 self.router
824 .release_selected_if_unnotified(&self.target_session_id, &self.run_id, wakes)
825 .await
826 }
827 None => false,
828 };
829 let result = if released {
830 Ok(None)
831 } else {
832 self.router
833 .finish_finalization(&self.target_session_id, &self.run_id, admitted_generation)
834 .await
835 };
836 if result.is_ok() {
837 self.armed = false;
838 self.abort_cleanup = None;
839 }
840 result
841 }
842}
843
844impl Drop for SessionRunRegistration {
845 fn drop(&mut self) {
846 if !self.armed {
847 return;
848 }
849 let router = self.router.clone();
850 let target_session_id = self.target_session_id.clone();
851 let run_id = self.run_id.clone();
852 let scoped_transport_cleanup = self.scoped_transport_cleanup.clone();
853 let abort_cleanup = self.abort_cleanup.take();
854 let quarantine = self.quarantine_on_drop || self.observed_forbidden_running();
855 if self.selected_input_wakes.is_some() {
856 drop(self.notifications.take());
857 }
858 let selected_input_wakes = self.selected_input_wakes.take();
859 if let Ok(runtime) = tokio::runtime::Handle::try_current() {
860 runtime.spawn(async move {
861 if let Some(cleanup) = scoped_transport_cleanup {
862 if let Err(error) = cleanup().await {
863 tracing::warn!(%error, "closed Loop transport cleanup failed; retaining dropped owner");
864 return;
865 }
866 }
867 if quarantine {
868 router.begin_finalization(&target_session_id, &run_id).await;
869 return;
870 }
871 if let Some(cleanup) = abort_cleanup {
872 cleanup().await;
873 }
874 router
875 .cleanup_abandoned_registration(
876 &target_session_id,
877 &run_id,
878 selected_input_wakes,
879 )
880 .await;
881 });
882 }
883 }
884}
885
886impl SessionActivationRouter {
887 fn publication_fence(&self, target: &str) -> Arc<Mutex<()>> {
888 let mut fences = self
889 .publication_fences
890 .lock()
891 .unwrap_or_else(std::sync::PoisonError::into_inner);
892 fences.retain(|_, fence| fence.strong_count() > 0);
893 if let Some(fence) = fences.get(target).and_then(Weak::upgrade) {
894 return fence;
895 }
896 let fence = Arc::new(Mutex::new(()));
897 fences.insert(target.into(), Arc::downgrade(&fence));
898 fence
899 }
900
901 async fn lock_publication(&self, target: &str) -> tokio::sync::OwnedMutexGuard<()> {
902 self.publication_fence(target).lock_owned().await
905 }
906
907 async fn coordinator_generation(&self, target: &str) -> Result<u64, SessionActivationError> {
908 let inbox = self
909 .inbox
910 .read()
911 .unwrap_or_else(std::sync::PoisonError::into_inner)
912 .clone();
913 match inbox {
914 Some(inbox) => inbox
915 .coordinator_activation_generation(target)
916 .await
917 .map_err(inbox_activation_error),
918 None => Ok(0),
919 }
920 }
921
922 async fn released_pending_generation(
923 &self,
924 target: &str,
925 after: u64,
926 ) -> Result<u64, SessionActivationError> {
927 if self.coordinator_generation(target).await? <= after {
928 return Ok(0);
929 }
930 let inbox = self
931 .inbox
932 .read()
933 .unwrap_or_else(std::sync::PoisonError::into_inner)
934 .clone();
935 match inbox {
936 Some(inbox) => {
937 let backlog = inbox
938 .inspect(target)
939 .await
940 .map_err(inbox_activation_error)?;
941 Ok(
942 if backlog.activation_pending() && backlog.coordinator_generation > after {
943 backlog.coordinator_generation
944 } else {
945 0
946 },
947 )
948 }
949 None => Ok(0),
950 }
951 }
952
953 pub fn new() -> Arc<Self> {
954 Arc::new(Self::default())
955 }
956
957 pub async fn set_spawner(&self, spawner: Arc<dyn SessionActivationSpawner>) {
960 *self.spawner.write().await = Some(spawner);
961 }
962
963 pub async fn current_run_id(&self, target_session_id: &str) -> Option<String> {
965 self.states
966 .lock()
967 .await
968 .get(target_session_id)
969 .and_then(|state| state.owner.as_ref().map(|owner| owner.run_id.clone()))
970 }
971
972 pub(crate) async fn child_completion_successor_run(
977 &self,
978 target_session_id: &str,
979 wait_or_guardian_transition: bool,
980 previous_fenced_run: Option<&str>,
981 ) -> Option<String> {
982 let mut states = self.states.lock().await;
983 let owner = states.get_mut(target_session_id)?.owner.as_mut()?;
984 owner.child_completions_need_successor |=
985 wait_or_guardian_transition || previous_fenced_run == Some(owner.run_id.as_str());
986 (owner.finalizing || owner.child_completions_need_successor).then(|| owner.run_id.clone())
987 }
988
989 pub fn set_inbox(&self, inbox: Arc<dyn SessionInboxPort>) {
991 *self
992 .inbox
993 .write()
994 .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(inbox);
995 }
996
997 pub async fn reconcile_wake(
1001 &self,
1002 target_session_id: &str,
1003 now: DateTime<Utc>,
1004 ) -> Result<SessionWakeReconcileResult, SessionActivationError> {
1005 let inbox = self
1006 .inbox
1007 .read()
1008 .unwrap_or_else(std::sync::PoisonError::into_inner)
1009 .clone()
1010 .ok_or_else(|| {
1011 SessionActivationError::Internal("session inbox is not configured".into())
1012 })?;
1013 let readiness = inbox
1014 .inspect_wake_readiness(target_session_id, now)
1015 .await
1016 .map_err(inbox_activation_error)?;
1017 let Some(candidate) = readiness.ready else {
1018 return Ok(SessionWakeReconcileResult {
1019 disposition: None,
1020 next_due_at: readiness.next_due_at,
1021 });
1022 };
1023 let mut result = self
1024 .request_recovery_activation(target_session_id, candidate, now)
1025 .await?;
1026 result.next_due_at = match (result.next_due_at, readiness.next_due_at) {
1027 (Some(left), Some(right)) => Some(left.min(right)),
1028 (left, right) => left.or(right),
1029 };
1030 Ok(result)
1031 }
1032
1033 async fn reconcile_older_message(
1037 &self,
1038 target: &str,
1039 ) -> Result<Option<SessionActivationDisposition>, SessionActivationError> {
1040 let high = self
1041 .states
1042 .lock()
1043 .await
1044 .get(target)
1045 .map_or(0, |state| state.last_dispatched_generation);
1046 self.reconcile_message_below(target, high).await
1047 }
1048
1049 async fn reconcile_message_below(
1050 &self,
1051 target: &str,
1052 high: u64,
1053 ) -> Result<Option<SessionActivationDisposition>, SessionActivationError> {
1054 if high == 0 {
1055 return Ok(None);
1056 }
1057 let inbox = self
1058 .inbox
1059 .read()
1060 .unwrap_or_else(std::sync::PoisonError::into_inner)
1061 .clone();
1062 let Some(inbox) = inbox else {
1063 return Ok(None);
1064 };
1065 let now = Utc::now();
1066 let candidate = inbox
1067 .inspect_wake_readiness(target, now)
1068 .await
1069 .map_err(inbox_activation_error)?
1070 .ready;
1071 match candidate.filter(|candidate| {
1072 candidate.generation < high
1073 && candidate.coordinator_message_policy
1076 == Some(bamboo_domain::SessionActivationPolicy::InterruptSpecificWait)
1077 }) {
1078 Some(candidate) => self
1079 .request_recovery_activation(target, candidate, now)
1080 .await
1081 .map(|result| result.disposition),
1082 None => Ok(None),
1083 }
1084 }
1085
1086 #[cfg(test)]
1087 async fn pause_after_terminal_inspection(&self) {
1088 let pause = self.terminal_inspection_pause.lock().await.take();
1089 if let Some((entered, release)) = pause {
1090 entered.notify_one();
1091 release.notified().await;
1092 }
1093 }
1094
1095 async fn request_recovery_activation(
1096 &self,
1097 target_session_id: &str,
1098 candidate: SessionInboxWakeCandidate,
1099 now: DateTime<Utc>,
1100 ) -> Result<SessionWakeReconcileResult, SessionActivationError> {
1101 loop {
1102 let (reservation_wait, reservation_token) = {
1103 let mut states = self.states.lock().await;
1104 let state = states.entry(target_session_id.to_string()).or_default();
1105 state.latest_generation = state.latest_generation.max(candidate.generation);
1106 if let Some(owner) = state.owner.as_mut() {
1107 if !owner.finalizing {
1108 state.notify.send_replace(candidate.generation);
1109 if owner
1110 .delivery_sink
1111 .as_ref()
1112 .is_some_and(|sink| sink.send(candidate.generation).is_err())
1113 {
1114 owner.delivery_sink = None;
1115 }
1116 return Ok(SessionWakeReconcileResult {
1117 disposition: Some(SessionActivationDisposition::ActiveNotified),
1118 next_due_at: None,
1119 });
1120 }
1121 return Ok(SessionWakeReconcileResult {
1122 disposition: Some(SessionActivationDisposition::ActivationCoalesced),
1123 next_due_at: None,
1124 });
1125 }
1126 if let Some(next_due_at) = recovery_retry_at(state, &candidate, now) {
1127 return Ok(SessionWakeReconcileResult {
1128 disposition: None,
1129 next_due_at,
1130 });
1131 } else if state.activation_reserved {
1132 (Some(state.activation_epoch.subscribe()), None)
1133 } else {
1134 let token = reserve_activation_token(state);
1135 (None, Some(token))
1136 }
1137 };
1138
1139 if let Some(mut reservation_wait) = reservation_wait {
1140 if reservation_wait.changed().await.is_err() {
1141 return Err(SessionActivationError::TargetNotFound(
1142 target_session_id.to_string(),
1143 ));
1144 }
1145 continue;
1146 }
1147 let token = reservation_token.expect("recovery activation owns reservation token");
1148 let disposition = self
1149 .dispatch_reserved_with_recovery(
1150 target_session_id,
1151 candidate.generation,
1152 token,
1153 false,
1154 Some((candidate, now)),
1155 )
1156 .await?;
1157 return Ok(SessionWakeReconcileResult {
1158 disposition: Some(disposition),
1159 next_due_at: None,
1160 });
1161 }
1162 }
1163
1164 #[cfg(test)]
1165 pub(crate) async fn hold_state_lock_for_test(
1166 &self,
1167 entered: Arc<tokio::sync::Notify>,
1168 release: Arc<tokio::sync::Notify>,
1169 ) {
1170 let _states = self.states.lock().await;
1171 entered.notify_one();
1172 release.notified().await;
1173 }
1174
1175 #[cfg(test)]
1176 pub(crate) async fn install_owner_placeholder_for_test(
1177 &self,
1178 target_session_id: &str,
1179 run_id: &str,
1180 ) {
1181 let _publication = self.lock_publication(target_session_id).await;
1182 let mut states = self.states.lock().await;
1183 states
1184 .entry(target_session_id.to_string())
1185 .or_default()
1186 .owner = Some(ActiveOwner {
1187 instance: Arc::new(()),
1188 run_id: run_id.to_string(),
1189 session_created_at: None,
1190 finalizing: false,
1191 registrations: 0,
1192 child_completions_need_successor: false,
1193 delivery_sink: None,
1194 scoped_loop: None,
1195 });
1196 }
1197
1198 pub async fn register_run(
1202 self: &Arc<Self>,
1203 target_session_id: &str,
1204 run_id: &str,
1205 ) -> Result<SessionRunRegistration, SessionRunRegistrationError> {
1206 self.register_run_with_birth(target_session_id, run_id, None)
1207 .await
1208 }
1209
1210 pub async fn register_session_run(
1213 self: &Arc<Self>,
1214 session: &bamboo_domain::Session,
1215 run_id: &str,
1216 ) -> Result<SessionRunRegistration, SessionRunRegistrationError> {
1217 self.register_run_with_birth(&session.id, run_id, Some(session.created_at))
1218 .await
1219 }
1220
1221 pub(crate) async fn bind_run_lifetime(
1222 self: &Arc<Self>,
1223 ownership: &SessionRunOwnership,
1224 session: &bamboo_domain::Session,
1225 ) -> bool {
1226 if !ownership.matches_router(self) || ownership.target_session_id != session.id {
1227 return false;
1228 }
1229 let _publication = self.lock_publication(&session.id).await;
1230 let mut states = self.states.lock().await;
1231 let Some(owner) = states
1232 .get_mut(&session.id)
1233 .and_then(|state| state.owner.as_mut())
1234 else {
1235 return false;
1236 };
1237 if owner.run_id != ownership.run_id
1238 || owner.finalizing
1239 || owner.registrations == 0
1240 || owner
1241 .session_created_at
1242 .is_some_and(|birth| birth != session.created_at)
1243 {
1244 return false;
1245 }
1246 owner.session_created_at = Some(session.created_at);
1247 true
1248 }
1249
1250 pub(crate) async fn bind_supervisor_loop_control(
1254 self: &Arc<Self>,
1255 ownership: &SessionRunOwnership,
1256 binding: bamboo_domain::ScopedLoopControlBinding,
1257 execution_cancel: tokio_util::sync::CancellationToken,
1258 host_stop: tokio_util::sync::CancellationToken,
1259 ) -> Result<(), bamboo_domain::SessionInboxError> {
1260 let inbox = self
1261 .inbox
1262 .read()
1263 .unwrap_or_else(std::sync::PoisonError::into_inner)
1264 .clone()
1265 .ok_or_else(|| {
1266 bamboo_domain::SessionInboxError::InvalidClaim("scoped Loop Inbox missing".into())
1267 })?;
1268 binding
1269 .validate()
1270 .map_err(bamboo_domain::SessionInboxError::InvalidClaim)?;
1271 if execution_cancel.is_cancelled() || host_stop.is_cancelled() {
1272 return Err(bamboo_domain::SessionInboxError::InvalidClaim(
1273 "Loop owner is stopped".into(),
1274 ));
1275 }
1276 if !ownership.matches_router(self)
1277 || ownership.target_session_id != binding.target_session_id
1278 || ownership.run_id != binding.owner_run_id
1279 {
1280 return Err(bamboo_domain::SessionInboxError::InvalidClaim(
1281 "Loop owner proof mismatch".into(),
1282 ));
1283 }
1284 let _publication = self.lock_publication(&binding.target_session_id).await;
1285 let mut states = self.states.lock().await;
1286 let owner = states
1287 .get_mut(&binding.target_session_id)
1288 .and_then(|state| state.owner.as_mut())
1289 .filter(|owner| {
1290 Arc::ptr_eq(&owner.instance, &ownership.instance)
1291 && owner.run_id == ownership.run_id
1292 && !owner.finalizing
1293 && owner.registrations > 0
1294 && owner.session_created_at == Some(binding.target_created_at)
1295 })
1296 .ok_or_else(|| {
1297 bamboo_domain::SessionInboxError::InvalidClaim(
1298 "Loop owner is no longer current".into(),
1299 )
1300 })?;
1301 if execution_cancel.is_cancelled() || host_stop.is_cancelled() {
1303 return Err(bamboo_domain::SessionInboxError::InvalidClaim(
1304 "Loop owner stopped during grant binding".into(),
1305 ));
1306 }
1307 if let Some(old) = &owner.scoped_loop {
1308 if old.binding != binding || old.stopped() {
1309 return Err(bamboo_domain::SessionInboxError::InvalidClaim(
1310 "different or stopped Loop grant already bound".into(),
1311 ));
1312 }
1313 return Ok(());
1315 }
1316 let runtime: Arc<dyn SessionActivationPort> = self.clone();
1317 if !inbox
1318 .bind_scoped_transport_runtime(Arc::downgrade(&runtime), self.runtime_instance.clone())
1319 {
1320 return Err(bamboo_domain::SessionInboxError::InvalidClaim(
1321 "scoped Inbox belongs to another live runtime or is unsupported".into(),
1322 ));
1323 }
1324 owner.scoped_loop = Some(SupervisorLoopGrant {
1325 binding,
1326 execution_cancel,
1327 host_stop,
1328 });
1329 Ok(())
1330 }
1331
1332 pub async fn revoke_supervisor_loop_control(
1335 &self,
1336 target: &str,
1337 grant_id: &uuid::Uuid,
1338 ) -> bool {
1339 let _publication = self.lock_publication(target).await;
1340 let mut states = self.states.lock().await;
1341 let Some(owner) = states
1342 .get_mut(target)
1343 .and_then(|state| state.owner.as_mut())
1344 else {
1345 return false;
1346 };
1347 if owner
1348 .scoped_loop
1349 .as_ref()
1350 .is_some_and(|grant| grant.binding.grant_id == *grant_id)
1351 {
1352 owner.scoped_loop = None;
1353 true
1354 } else {
1355 false
1356 }
1357 }
1358
1359 async fn register_run_with_birth(
1360 self: &Arc<Self>,
1361 target_session_id: &str,
1362 run_id: &str,
1363 birth: Option<DateTime<Utc>>,
1364 ) -> Result<SessionRunRegistration, SessionRunRegistrationError> {
1365 loop {
1366 let (notifications, reservation_wait) = {
1367 let _publication = self.lock_publication(target_session_id).await;
1368 let mut states = self.states.lock().await;
1369 let state = states.entry(target_session_id.to_string()).or_default();
1370 match state.owner.as_mut() {
1371 Some(owner)
1372 if owner.run_id == run_id
1373 && owner.registrations == 0
1374 && !owner.finalizing =>
1375 {
1376 if owner
1380 .session_created_at
1381 .is_some_and(|existing| birth.is_some_and(|birth| existing != birth))
1382 {
1383 return Err(SessionRunRegistrationError::OwnerCollision {
1384 target_session_id: target_session_id.into(),
1385 existing_run_id: owner.run_id.clone(),
1386 attempted_run_id: run_id.into(),
1387 });
1388 }
1389 if birth.is_some() {
1390 owner.session_created_at = birth;
1391 }
1392 owner.registrations = 1;
1393 (
1394 Some((
1395 state.notify.subscribe(),
1396 state
1397 .owner
1398 .as_ref()
1399 .expect("registered owner")
1400 .instance
1401 .clone(),
1402 )),
1403 None,
1404 )
1405 }
1406 Some(owner) => {
1407 return Err(SessionRunRegistrationError::OwnerCollision {
1408 target_session_id: target_session_id.to_string(),
1409 existing_run_id: owner.run_id.clone(),
1410 attempted_run_id: run_id.to_string(),
1411 });
1412 }
1413 None if state.activation_reserved => {
1414 (None, Some(state.activation_epoch.subscribe()))
1420 }
1421 None => {
1422 state.owner = Some(ActiveOwner {
1423 instance: Arc::new(()),
1424 run_id: run_id.to_string(),
1425 session_created_at: birth,
1426 finalizing: false,
1427 registrations: 1,
1428 child_completions_need_successor: false,
1429 delivery_sink: None,
1430 scoped_loop: None,
1431 });
1432 (
1433 Some((
1434 state.notify.subscribe(),
1435 state
1436 .owner
1437 .as_ref()
1438 .expect("registered owner")
1439 .instance
1440 .clone(),
1441 )),
1442 None,
1443 )
1444 }
1445 }
1446 };
1447
1448 if let Some(mut reservation_wait) = reservation_wait {
1449 let _ = reservation_wait.changed().await;
1450 continue;
1451 }
1452 let (notifications, instance) =
1453 notifications.expect("registered owner has a receiver and instance");
1454 return Ok(SessionRunRegistration {
1455 instance,
1456 router: self.clone(),
1457 target_session_id: target_session_id.to_string(),
1458 run_id: run_id.to_string(),
1459 notifications: Some(notifications),
1460 abort_cleanup: None,
1461 scoped_transport_cleanup: None,
1462 quarantine_on_drop: false,
1463 tool_suspension_observation: None,
1464 selected_input_wakes: None,
1465 armed: true,
1466 });
1467 }
1468 }
1469
1470 pub async fn attach_delivery_sink(
1476 &self,
1477 target_session_id: &str,
1478 sink: mpsc::UnboundedSender<u64>,
1479 ) -> Option<String> {
1480 let mut states = self.states.lock().await;
1481 let state = states.get_mut(target_session_id)?;
1482 let owner = state.owner.as_mut()?;
1483 if owner.finalizing {
1484 return None;
1485 }
1486 owner.delivery_sink = Some(sink.clone());
1487 let run_id = owner.run_id.clone();
1488 let delivery_generation = state
1489 .latest_generation
1490 .max(state.latest_delivery_generation);
1491 if delivery_generation > 0 && sink.send(delivery_generation).is_err() {
1492 owner.delivery_sink = None;
1493 return None;
1494 }
1495 Some(run_id)
1496 }
1497
1498 pub async fn detach_delivery_sink(&self, target_session_id: &str, run_id: &str) {
1501 let mut states = self.states.lock().await;
1502 let Some(state) = states.get_mut(target_session_id) else {
1503 return;
1504 };
1505 let Some(owner) = state.owner.as_mut() else {
1506 return;
1507 };
1508 if owner.run_id == run_id {
1509 owner.delivery_sink = None;
1510 }
1511 }
1512
1513 pub async fn owns_run(&self, target_session_id: &str, run_id: &str) -> bool {
1515 let states = self.states.lock().await;
1516 states
1517 .get(target_session_id)
1518 .and_then(|state| state.owner.as_ref())
1519 .is_some_and(|owner| owner.run_id == run_id && !owner.finalizing)
1520 }
1521
1522 pub async fn begin_finalization(&self, target_session_id: &str, run_id: &str) {
1527 let _publication = self.lock_publication(target_session_id).await;
1528 let mut states = self.states.lock().await;
1529 let state = states.entry(target_session_id.to_string()).or_default();
1530 match state.owner.as_mut() {
1531 Some(owner) if owner.run_id == run_id => owner.finalizing = true,
1532 Some(_) => {}
1533 None => {
1534 state.owner = Some(ActiveOwner {
1535 instance: Arc::new(()),
1536 run_id: run_id.to_string(),
1537 session_created_at: None,
1538 finalizing: true,
1539 registrations: 0,
1540 child_completions_need_successor: false,
1541 delivery_sink: None,
1542 scoped_loop: None,
1543 });
1544 }
1545 }
1546 }
1547
1548 async fn release_selected_if_unnotified(
1553 &self,
1554 target_session_id: &str,
1555 run_id: &str,
1556 wakes: &watch::Receiver<u64>,
1557 ) -> bool {
1558 let _publication = self.lock_publication(target_session_id).await;
1559 let mut states = self.states.lock().await;
1560 let Some(state) = states.get_mut(target_session_id) else {
1561 return true;
1562 };
1563 if !state
1564 .owner
1565 .as_ref()
1566 .is_some_and(|owner| owner.run_id == run_id)
1567 {
1568 return true;
1569 }
1570 if wakes.has_changed().unwrap_or(false) {
1571 return false;
1572 }
1573 state.owner = None;
1574 state.owner_completion_epoch = state.owner_completion_epoch.wrapping_add(1);
1575 if state.latest_generation == 0
1576 && !state.activation_reserved
1577 && state.notify.receiver_count() <= 1
1578 && state.activation_epoch.receiver_count() == 0
1579 {
1580 states.remove(target_session_id);
1581 }
1582 true
1583 }
1584
1585 async fn cleanup_abandoned_registration(
1586 self: &Arc<Self>,
1587 target_session_id: &str,
1588 run_id: &str,
1589 selected_input_wakes: Option<watch::Receiver<u64>>,
1590 ) {
1591 if let Some(wakes) = selected_input_wakes {
1592 if self
1593 .release_selected_if_unnotified(target_session_id, run_id, &wakes)
1594 .await
1595 {
1596 return;
1597 }
1598 }
1599 {
1600 let _publication = self.lock_publication(target_session_id).await;
1601 let mut states = self.states.lock().await;
1602 let Some(state) = states.get_mut(target_session_id) else {
1603 return;
1604 };
1605 let Some(owner) = state.owner.as_mut() else {
1606 return;
1607 };
1608 if owner.run_id != run_id {
1609 return;
1610 }
1611 owner.registrations = 0;
1612 owner.finalizing = true;
1613 }
1614
1615 let inbox = self
1620 .inbox
1621 .read()
1622 .unwrap_or_else(std::sync::PoisonError::into_inner)
1623 .clone();
1624 let (durable_generation, released_prefix) = match inbox {
1625 Some(inbox) => match inbox.inspect(target_session_id).await {
1626 Ok(backlog) if backlog.activation_pending() => (
1627 backlog.activation_generation,
1628 backlog.coordinator_generation,
1629 ),
1630 Ok(_) => (0, 0),
1631 Err(error) => {
1632 tracing::warn!(
1633 %target_session_id,
1634 %run_id,
1635 %error,
1636 "failed to reconcile durable SessionInbox generation for abandoned run"
1637 );
1638 (0, 0)
1639 }
1640 },
1641 None => (0, 0),
1642 };
1643
1644 #[cfg(test)]
1645 self.pause_after_terminal_inspection().await;
1646 let (reservation_to_dispatch, previous_high) = {
1647 let _publication = self.lock_publication(target_session_id).await;
1648 let mut states = self.states.lock().await;
1649 let Some(state) = states.get_mut(target_session_id) else {
1650 return;
1651 };
1652 if !state.owner.as_ref().is_some_and(|owner| {
1653 owner.run_id == run_id && owner.finalizing && owner.registrations == 0
1654 }) {
1655 return;
1656 }
1657 state.latest_generation = state.latest_generation.max(durable_generation);
1658 let previous_high = state.last_dispatched_generation;
1659 state.owner = None;
1660 state.owner_completion_epoch = state.owner_completion_epoch.wrapping_add(1);
1661 if (state.latest_generation > state.last_dispatched_generation
1662 || released_prefix > state.last_dispatched_coordinator_generation)
1663 && !state.activation_reserved
1664 {
1665 let generation = state.latest_generation;
1666 let token = reserve_activation_token(state);
1667 (Some((generation, token)), previous_high)
1668 } else {
1669 if state.latest_generation == 0
1670 && !state.activation_reserved
1671 && state.notify.receiver_count() == 0
1672 && state.activation_epoch.receiver_count() == 0
1673 {
1674 states.remove(target_session_id);
1675 }
1676 (None, previous_high)
1677 }
1678 };
1679
1680 if let Some((generation, token)) = reservation_to_dispatch {
1681 if let Err(error) = self
1682 .dispatch_reserved(target_session_id, generation, token)
1683 .await
1684 {
1685 tracing::error!(
1686 %target_session_id,
1687 %run_id,
1688 %error,
1689 "failed to reserve successor for abandoned SessionInbox owner"
1690 );
1691 }
1692 } else {
1693 if let Err(error) = self
1697 .reconcile_message_below(target_session_id, previous_high)
1698 .await
1699 {
1700 tracing::warn!(%target_session_id, %error,
1701 "failed to reconcile older eligible message after abandoned owner");
1702 }
1703 }
1704 }
1705
1706 pub async fn finish_finalization(
1710 &self,
1711 target_session_id: &str,
1712 run_id: &str,
1713 admitted_generation: u64,
1714 ) -> Result<Option<SessionActivationDisposition>, SessionActivationError> {
1715 let inbox = self
1718 .inbox
1719 .read()
1720 .unwrap_or_else(std::sync::PoisonError::into_inner)
1721 .clone();
1722 let durable_pending = match inbox {
1723 Some(inbox) => Some(
1724 inbox
1725 .inspect(target_session_id)
1726 .await
1727 .map_err(|error| SessionActivationError::Internal(error.to_string()))?,
1728 ),
1729 None => None,
1730 };
1731 #[cfg(test)]
1732 self.pause_after_terminal_inspection().await;
1733 let (reservation_to_dispatch, previous_high) = {
1734 let _publication = self.lock_publication(target_session_id).await;
1735 let mut states = self.states.lock().await;
1736 let state = states.entry(target_session_id.to_string()).or_default();
1737 let previous_high = state.last_dispatched_generation;
1738 if state
1739 .owner
1740 .as_ref()
1741 .is_some_and(|owner| owner.run_id != run_id)
1742 {
1743 return Ok(None);
1744 }
1745 if state.owner.take().is_some() {
1746 state.owner_completion_epoch = state.owner_completion_epoch.wrapping_add(1);
1747 }
1748 let pending = durable_pending
1749 .as_ref()
1750 .is_some_and(|backlog| backlog.activation_pending());
1751 let released_prefix = durable_pending
1752 .as_ref()
1753 .filter(|_| pending)
1754 .map_or(0, |backlog| backlog.coordinator_generation);
1755 if let Some(backlog) = durable_pending.as_ref() {
1756 state.latest_generation =
1757 state.latest_generation.max(backlog.activation_generation);
1758 }
1759 if (pending || state.latest_generation > admitted_generation)
1760 && (state.latest_generation > state.last_dispatched_generation
1761 || released_prefix > state.last_dispatched_coordinator_generation)
1762 && !state.activation_reserved
1763 {
1764 let generation = state.latest_generation;
1765 let token = reserve_activation_token(state);
1766 (Some((generation, token)), previous_high)
1767 } else {
1768 if state.latest_generation <= admitted_generation
1773 && !pending
1774 && state.owner.is_none()
1775 && !state.activation_reserved
1776 && state.notify.receiver_count() == 0
1777 && state.activation_epoch.receiver_count() == 0
1778 {
1779 states.remove(target_session_id);
1780 }
1781 (None, previous_high)
1782 }
1783 };
1784
1785 match reservation_to_dispatch {
1786 Some((generation, token)) => self
1787 .dispatch_reserved(target_session_id, generation, token)
1788 .await
1789 .map(Some),
1790 None => {
1793 self.reconcile_message_below(target_session_id, previous_high)
1794 .await
1795 }
1796 }
1797 }
1798
1799 pub async fn subscribe(&self, target_session_id: &str) -> watch::Receiver<u64> {
1800 let mut states = self.states.lock().await;
1801 states
1802 .entry(target_session_id.to_string())
1803 .or_default()
1804 .notify
1805 .subscribe()
1806 }
1807
1808 async fn recover_cancelled_reservation(
1817 &self,
1818 target_session_id: &str,
1819 reservation_token: u64,
1820 recover_on_drop: bool,
1821 ) {
1822 let previous_prefix = {
1823 let mut states = self.states.lock().await;
1824 let Some(state) = states.get_mut(target_session_id) else {
1825 return;
1826 };
1827 if !release_activation_token(state, reservation_token) || !recover_on_drop {
1828 return;
1829 }
1830 state.last_dispatched_coordinator_generation
1831 };
1832 let released_prefix = self
1833 .released_pending_generation(target_session_id, previous_prefix)
1834 .await
1835 .unwrap_or(0);
1836 let reservation_to_dispatch = {
1837 let mut states = self.states.lock().await;
1838 let Some(state) = states.get_mut(target_session_id) else {
1839 return;
1840 };
1841 if state.owner.is_none()
1842 && (state.latest_generation > state.last_dispatched_generation
1843 || released_prefix > state.last_dispatched_coordinator_generation)
1844 && !state.activation_reserved
1845 {
1846 let generation = state.latest_generation;
1847 let token = reserve_activation_token(state);
1848 Some((generation, token))
1849 } else {
1850 None
1851 }
1852 };
1853
1854 if let Some((generation, token)) = reservation_to_dispatch {
1855 if let Err(error) = self
1856 .dispatch_reserved_with_recovery(target_session_id, generation, token, false, None)
1857 .await
1858 {
1859 tracing::error!(
1860 %target_session_id,
1861 %error,
1862 "failed bounded retry after cancelled SessionInbox activation reservation"
1863 );
1864 }
1865 }
1866 }
1867
1868 async fn dispatch_reserved(
1869 &self,
1870 target_session_id: &str,
1871 generation: u64,
1872 reservation_token: u64,
1873 ) -> Result<SessionActivationDisposition, SessionActivationError> {
1874 self.dispatch_reserved_with_recovery(
1875 target_session_id,
1876 generation,
1877 reservation_token,
1878 true,
1879 None,
1880 )
1881 .await
1882 }
1883
1884 async fn dispatch_reserved_with_recovery(
1885 &self,
1886 target_session_id: &str,
1887 generation: u64,
1888 reservation_token: u64,
1889 recover_on_drop: bool,
1890 recovery: Option<(SessionInboxWakeCandidate, DateTime<Utc>)>,
1891 ) -> Result<SessionActivationDisposition, SessionActivationError> {
1892 let mut lease = ActivationReservationLease::new(
1893 self.clone(),
1894 target_session_id,
1895 reservation_token,
1896 recover_on_drop,
1897 );
1898 let Some(spawner) = self.spawner.read().await.clone() else {
1899 let mut states = self.states.lock().await;
1900 if let Some(state) = states.get_mut(target_session_id) {
1901 release_activation_token(state, reservation_token);
1902 }
1903 lease.disarm();
1904 return Err(SessionActivationError::Internal(
1905 "session activation spawner is not configured".to_string(),
1906 ));
1907 };
1908
1909 let coordinator_generation = self.coordinator_generation(target_session_id).await?;
1912
1913 let outcome = spawner
1914 .reserve_activation(target_session_id, generation)
1915 .await;
1916 if let Ok(SessionActivationReserveOutcome::Reserved(launch)) = &outcome {
1917 lease.wait_for_rollback(launch.rollback_completion());
1918 }
1919 match outcome {
1920 Ok(SessionActivationReserveOutcome::Reserved(launch)) => {
1921 let run_id = launch.run_id.clone();
1922 {
1923 let _publication = self.lock_publication(target_session_id).await;
1924 let mut states = self.states.lock().await;
1925 let state = states.entry(target_session_id.to_string()).or_default();
1926 if !release_activation_token(state, reservation_token) {
1927 lease.disarm();
1935 return if state.owner.is_some() {
1936 Ok(SessionActivationDisposition::ActiveNotified)
1937 } else {
1938 Err(SessionActivationError::Internal(
1939 "session activation reservation was superseded".to_string(),
1940 ))
1941 };
1942 }
1943 state.owner = Some(ActiveOwner {
1944 instance: Arc::new(()),
1945 run_id,
1946 session_created_at: None,
1947 finalizing: false,
1948 registrations: 0,
1949 child_completions_need_successor: false,
1950 delivery_sink: None,
1951 scoped_loop: None,
1952 });
1953 state.last_dispatched_generation =
1954 state.last_dispatched_generation.max(generation);
1955 state.last_dispatched_coordinator_generation = state
1956 .last_dispatched_coordinator_generation
1957 .max(coordinator_generation);
1958 if let Some((candidate, now)) = recovery {
1959 record_recovery_attempt(state, candidate, now, true);
1960 }
1961 }
1962 lease.disarm();
1963 launch.launch();
1966 Ok(SessionActivationDisposition::ActivationReserved)
1967 }
1968 Ok(SessionActivationReserveOutcome::AlreadyRunning { run_id }) => {
1969 let _publication = self.lock_publication(target_session_id).await;
1970 let mut states = self.states.lock().await;
1971 let state = states.entry(target_session_id.to_string()).or_default();
1972 if !release_activation_token(state, reservation_token) {
1973 lease.disarm();
1974 return Ok(SessionActivationDisposition::ActiveNotified);
1975 }
1976 state.owner = Some(ActiveOwner {
1977 instance: Arc::new(()),
1978 run_id,
1979 session_created_at: None,
1980 finalizing: false,
1981 registrations: 0,
1982 child_completions_need_successor: false,
1983 delivery_sink: None,
1984 scoped_loop: None,
1985 });
1986 state.notify.send_replace(generation);
1987 if let Some((candidate, now)) = recovery {
1988 record_recovery_attempt(state, candidate, now, true);
1989 }
1990 lease.disarm();
1991 Ok(SessionActivationDisposition::ActiveNotified)
1992 }
1993 Ok(SessionActivationReserveOutcome::NoWork) => {
1994 let mut states = self.states.lock().await;
1995 if let Some(state) = states.get_mut(target_session_id) {
1996 if release_activation_token(state, reservation_token) {
1997 if let Some((candidate, now)) = recovery {
1998 record_recovery_attempt(state, candidate, now, false);
1999 }
2000 }
2001 }
2002 lease.disarm();
2003 Ok(SessionActivationDisposition::ActivationCoalesced)
2004 }
2005 Ok(SessionActivationReserveOutcome::NotFound) => {
2006 let _publication = self.lock_publication(target_session_id).await;
2007 let mut states = self.states.lock().await;
2008 if states
2009 .get(target_session_id)
2010 .is_some_and(|state| state.activation_token == reservation_token)
2011 {
2012 states.remove(target_session_id);
2013 }
2014 lease.disarm();
2015 Err(SessionActivationError::TargetNotFound(
2016 target_session_id.to_string(),
2017 ))
2018 }
2019 Err(error) => {
2020 let mut states = self.states.lock().await;
2021 if let Some(state) = states.get_mut(target_session_id) {
2022 release_activation_token(state, reservation_token);
2023 }
2024 lease.disarm();
2025 Err(error)
2026 }
2027 }
2028 }
2029}
2030
2031#[async_trait]
2032impl SessionActivationPort for SessionActivationRouter {
2033 fn scoped_transport_runtime_instance(&self) -> Option<Arc<()>> {
2034 Some(self.runtime_instance.clone())
2035 }
2036
2037 fn retain_closed_scoped_transport_blocking(
2038 &self,
2039 binding: &bamboo_domain::ScopedLoopControlBinding,
2040 ) -> Result<Arc<dyn bamboo_domain::SessionRunDeliveryGuard>, bamboo_domain::SessionRunConflict>
2041 {
2042 let publication = self
2043 .publication_fence(&binding.target_session_id)
2044 .blocking_lock_owned();
2045 let states = self.states.blocking_lock();
2046 let matches = states
2047 .get(&binding.target_session_id)
2048 .and_then(|state| state.owner.as_ref())
2049 .is_some_and(|owner| {
2050 owner.finalizing
2051 && owner.registrations > 0
2052 && owner.run_id == binding.owner_run_id
2053 && owner.session_created_at == Some(binding.target_created_at)
2054 && owner
2055 .scoped_loop
2056 .as_ref()
2057 .is_none_or(|grant| grant.binding.execution_id == binding.execution_id)
2058 });
2059 if !matches {
2060 return Err(bamboo_domain::SessionRunConflict::unsupported(
2061 &binding.owner_run_id,
2062 ));
2063 }
2064 drop(states);
2065 Ok(Arc::new(ClosedLoopTransportGuard(
2066 ExpectedRunDeliveryGuard {
2067 target_session_id: binding.target_session_id.clone(),
2068 session_created_at: binding.target_created_at,
2069 run_id: binding.owner_run_id.clone(),
2070 states: self.states.clone(),
2071 _publication: publication,
2072 },
2073 )))
2074 }
2075
2076 fn retain_expected_run_blocking(
2077 &self,
2078 target_session_id: &str,
2079 target_created_at: DateTime<Utc>,
2080 expected_run_id: &str,
2081 ) -> Result<Arc<dyn bamboo_domain::SessionRunDeliveryGuard>, bamboo_domain::SessionRunConflict>
2082 {
2083 use bamboo_domain::{SessionRunConflict, SessionRunConflictReason};
2084 let publication = self
2087 .publication_fence(target_session_id)
2088 .blocking_lock_owned();
2089 let states = self.states.blocking_lock();
2090 let owner = states
2091 .get(target_session_id)
2092 .and_then(|state| state.owner.as_ref());
2093 let reason = match owner {
2094 None => Some(SessionRunConflictReason::NoActiveRun),
2095 Some(owner)
2096 if owner
2097 .session_created_at
2098 .is_some_and(|birth| birth != target_created_at) =>
2099 {
2100 Some(SessionRunConflictReason::SessionBirthChanged)
2101 }
2102 Some(owner) if owner.finalizing => Some(SessionRunConflictReason::Finalizing),
2103 Some(owner) if owner.registrations == 0 => {
2104 Some(SessionRunConflictReason::NotRegistered)
2105 }
2106 Some(owner) if owner.session_created_at.is_none() => {
2107 Some(SessionRunConflictReason::UnboundLifetime)
2108 }
2109 Some(owner) if owner.run_id != expected_run_id => {
2110 Some(SessionRunConflictReason::RunChanged)
2111 }
2112 Some(_) => None,
2113 };
2114 if let Some(reason) = reason {
2115 return Err(SessionRunConflict {
2116 reason,
2117 expected_run_id: expected_run_id.into(),
2118 observed_run_id: owner
2119 .filter(|owner| owner.session_created_at == Some(target_created_at))
2120 .map(|owner| owner.run_id.clone()),
2121 });
2122 }
2123 Ok(Arc::new(ExpectedRunDeliveryGuard {
2124 target_session_id: target_session_id.into(),
2125 session_created_at: target_created_at,
2126 run_id: expected_run_id.into(),
2127 states: self.states.clone(),
2128 _publication: publication,
2129 }))
2130 }
2131
2132 async fn observe_supervisor_loop(
2133 &self,
2134 supervisor: &bamboo_domain::SupervisorReference,
2135 target: &str,
2136 expected_link_revision: Option<u64>,
2137 action: bamboo_domain::LoopControlAction,
2138 ) -> Result<bamboo_domain::ScopedLoopControlBinding, bamboo_domain::SessionInboxError> {
2139 let _publication = self.lock_publication(target).await;
2140 let states = self.states.lock().await;
2141 let owner = states
2142 .get(target)
2143 .and_then(|state| state.owner.as_ref())
2144 .filter(|owner| !owner.finalizing && owner.registrations > 0)
2145 .ok_or_else(|| {
2146 bamboo_domain::SessionInboxError::InvalidClaim("no live Loop owner".into())
2147 })?;
2148 let grant = owner
2149 .scoped_loop
2150 .as_ref()
2151 .filter(|grant| !grant.stopped())
2152 .ok_or_else(|| {
2153 bamboo_domain::SessionInboxError::InvalidClaim(
2154 "host did not grant Supervisor Loop control".into(),
2155 )
2156 })?;
2157 let binding = &grant.binding;
2158 binding
2159 .validate()
2160 .map_err(bamboo_domain::SessionInboxError::InvalidClaim)?;
2161 if binding.supervisor != *supervisor
2162 || binding.target_session_id != target
2163 || owner.run_id != binding.owner_run_id
2164 || owner.session_created_at != Some(binding.target_created_at)
2165 || !binding.allowed_actions.contains(&action)
2166 || expected_link_revision.is_some_and(|revision| revision != binding.link_revision)
2167 || (expected_link_revision.is_none()
2168 && action != bamboo_domain::LoopControlAction::Status)
2169 {
2170 return Err(bamboo_domain::SessionInboxError::InvalidClaim(
2171 "live Loop grant observation mismatch".into(),
2172 ));
2173 }
2174 Ok(binding.clone())
2175 }
2176
2177 fn retain_supervisor_loop_blocking(
2178 &self,
2179 binding: &bamboo_domain::ScopedLoopControlBinding,
2180 ) -> Result<Arc<dyn bamboo_domain::SessionRunDeliveryGuard>, bamboo_domain::SessionRunConflict>
2181 {
2182 use bamboo_domain::{SessionRunConflict, SessionRunConflictReason};
2183 let publication = self
2184 .publication_fence(&binding.target_session_id)
2185 .blocking_lock_owned();
2186 let states = self.states.blocking_lock();
2187 let owner = states
2188 .get(&binding.target_session_id)
2189 .and_then(|state| state.owner.as_ref());
2190 if binding.validate().is_err()
2191 || owner.is_none_or(|owner| {
2192 owner.finalizing
2193 || owner.registrations == 0
2194 || owner.run_id != binding.owner_run_id
2195 || owner.session_created_at != Some(binding.target_created_at)
2196 || owner
2197 .scoped_loop
2198 .as_ref()
2199 .is_none_or(|grant| grant.stopped() || grant.binding != *binding)
2200 })
2201 {
2202 return Err(SessionRunConflict {
2203 reason: SessionRunConflictReason::RunChanged,
2204 expected_run_id: binding.owner_run_id.clone(),
2205 observed_run_id: owner
2206 .filter(|owner| owner.session_created_at == Some(binding.target_created_at))
2207 .map(|owner| owner.run_id.clone()),
2208 });
2209 }
2210 let grant = owner
2211 .and_then(|owner| owner.scoped_loop.as_ref())
2212 .expect("validated Loop grant")
2213 .clone();
2214 Ok(Arc::new(SupervisorLoopDeliveryGuard {
2215 inner: ExpectedRunDeliveryGuard {
2216 target_session_id: binding.target_session_id.clone(),
2217 session_created_at: binding.target_created_at,
2218 run_id: binding.owner_run_id.clone(),
2219 states: self.states.clone(),
2220 _publication: publication,
2221 },
2222 grant,
2223 }))
2224 }
2225
2226 async fn notify_active_delivery(&self, target_session_id: &str, inbox_generation: u64) -> bool {
2227 let mut states = self.states.lock().await;
2228 let Some(state) = states.get_mut(target_session_id) else {
2229 return false;
2230 };
2231 let Some(owner) = state.owner.as_mut().filter(|owner| !owner.finalizing) else {
2232 return false;
2233 };
2234 state.latest_delivery_generation = state.latest_delivery_generation.max(inbox_generation);
2235 state.notify.send_replace(inbox_generation);
2236 if owner
2237 .delivery_sink
2238 .as_ref()
2239 .is_some_and(|sink| sink.send(inbox_generation).is_err())
2240 {
2241 owner.delivery_sink = None;
2242 }
2243 true
2244 }
2245
2246 async fn request_activation(
2247 &self,
2248 target_session_id: &str,
2249 inbox_generation: u64,
2250 ) -> Result<SessionActivationDisposition, SessionActivationError> {
2251 let previous_prefix = {
2252 let states = self.states.lock().await;
2253 states
2254 .get(target_session_id)
2255 .filter(|state| inbox_generation <= state.last_dispatched_generation)
2256 .map(|state| state.last_dispatched_coordinator_generation)
2257 };
2258 let released_prefix = match previous_prefix {
2259 Some(prefix) => {
2260 self.released_pending_generation(target_session_id, prefix)
2261 .await?
2262 }
2263 None => 0,
2264 };
2265 loop {
2266 let (reservation_wait, reservation_token) = {
2267 let mut states = self.states.lock().await;
2268 let state = states.entry(target_session_id.to_string()).or_default();
2269 state.latest_generation = state.latest_generation.max(inbox_generation);
2270 if let Some(owner) = state.owner.as_mut() {
2271 if !owner.finalizing {
2272 state.notify.send_replace(inbox_generation);
2273 if owner
2274 .delivery_sink
2275 .as_ref()
2276 .is_some_and(|sink| sink.send(inbox_generation).is_err())
2277 {
2278 owner.delivery_sink = None;
2279 }
2280 return Ok(SessionActivationDisposition::ActiveNotified);
2281 }
2282 return Ok(SessionActivationDisposition::ActivationCoalesced);
2283 }
2284 if inbox_generation <= state.last_dispatched_generation
2285 && released_prefix <= state.last_dispatched_coordinator_generation
2286 {
2287 drop(states);
2288 return self
2289 .reconcile_older_message(target_session_id)
2290 .await
2291 .map(|result| {
2292 result.unwrap_or(SessionActivationDisposition::ActivationCoalesced)
2293 });
2294 }
2295 if state.activation_reserved {
2296 (Some(state.activation_epoch.subscribe()), None)
2297 } else {
2298 let token = reserve_activation_token(state);
2299 (None, Some(token))
2300 }
2301 };
2302
2303 if let Some(mut reservation_wait) = reservation_wait {
2304 if reservation_wait.changed().await.is_err() {
2310 return Err(SessionActivationError::TargetNotFound(
2311 target_session_id.to_string(),
2312 ));
2313 }
2314 continue;
2315 }
2316
2317 let reservation_token =
2318 reservation_token.expect("non-waiting activation owns a reservation token");
2319 return self
2320 .dispatch_reserved_with_recovery(
2321 target_session_id,
2322 inbox_generation,
2323 reservation_token,
2324 false,
2325 None,
2326 )
2327 .await;
2328 }
2329 }
2330}
2331
2332#[cfg(test)]
2333mod tests {
2334 use super::*;
2335 use std::sync::atomic::{AtomicUsize, Ordering};
2336 use tokio::sync::{Barrier, Notify};
2337
2338 #[tokio::test]
2339 async fn wake_recovery_retries_same_candidate_only_after_owner_completion_and_backoff() {
2340 use bamboo_domain::{
2341 Session, SessionActivationPolicy, SessionInboxLimits, SessionMessageEnvelope, Storage,
2342 };
2343 use bamboo_storage::{FileSessionInbox, SessionStoreV2};
2344
2345 let temp = tempfile::TempDir::new().unwrap();
2346 let store = Arc::new(SessionStoreV2::new(temp.path().into()).await.unwrap());
2347 store
2348 .save_session(&Session::new("target", "model"))
2349 .await
2350 .unwrap();
2351 let inbox = Arc::new(FileSessionInbox::new(store, SessionInboxLimits::default()));
2352 inbox
2353 .deliver_with_activation_intent(
2354 &SessionMessageEnvelope::user_input("target", "still pending"),
2355 SessionActivationPolicy::RespectSpecificWait,
2356 None,
2357 )
2358 .await
2359 .unwrap();
2360 let router = SessionActivationRouter::new();
2361 router.set_inbox(inbox);
2362 let spawner = spawner();
2363 router.set_spawner(spawner.clone()).await;
2364 let now = Utc::now();
2365 assert_eq!(
2366 router
2367 .reconcile_wake("target", now)
2368 .await
2369 .unwrap()
2370 .disposition,
2371 Some(SessionActivationDisposition::ActivationReserved)
2372 );
2373 assert_eq!(
2374 router
2375 .reconcile_wake("target", now)
2376 .await
2377 .unwrap()
2378 .disposition,
2379 Some(SessionActivationDisposition::ActiveNotified)
2380 );
2381 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 1);
2382 let mut registration = router.register_run("target", "run-1").await.unwrap();
2383 registration.begin_finalization().await;
2384 assert_eq!(registration.finish(0).await.unwrap(), None);
2385
2386 let deferred = router
2387 .reconcile_wake("target", now + Duration::seconds(29))
2388 .await
2389 .unwrap();
2390 assert_eq!(deferred.disposition, None);
2391 assert_eq!(deferred.next_due_at, Some(now + INITIAL_RECOVERY_BACKOFF));
2392 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 1);
2393 assert_eq!(
2394 router
2395 .reconcile_wake("target", now + INITIAL_RECOVERY_BACKOFF)
2396 .await
2397 .unwrap()
2398 .disposition,
2399 Some(SessionActivationDisposition::ActivationReserved)
2400 );
2401 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
2402 }
2403
2404 #[tokio::test]
2405 async fn wake_recovery_waits_for_exact_owned_lease_expiry() {
2406 use bamboo_domain::{
2407 Session, SessionActivationPolicy, SessionInboxLeaseRequest, SessionInboxLimits,
2408 SessionMessageEnvelope, Storage,
2409 };
2410 use bamboo_storage::{FileSessionInbox, SessionStoreV2};
2411
2412 let temp = tempfile::TempDir::new().unwrap();
2413 let store = Arc::new(SessionStoreV2::new(temp.path().into()).await.unwrap());
2414 store
2415 .save_session(&Session::new("target", "model"))
2416 .await
2417 .unwrap();
2418 let inbox = Arc::new(FileSessionInbox::new(store, SessionInboxLimits::default()));
2419 inbox
2420 .deliver_with_activation_intent(
2421 &SessionMessageEnvelope::user_input("target", "unacknowledged"),
2422 SessionActivationPolicy::RespectSpecificWait,
2423 None,
2424 )
2425 .await
2426 .unwrap();
2427 let router = SessionActivationRouter::new();
2428 router.set_inbox(inbox.clone());
2429 let spawner = spawner();
2430 router.set_spawner(spawner.clone()).await;
2431 let now = Utc::now();
2432 assert_eq!(
2433 router
2434 .reconcile_wake("target", now)
2435 .await
2436 .unwrap()
2437 .disposition,
2438 Some(SessionActivationDisposition::ActivationReserved)
2439 );
2440 let claim = inbox
2441 .claim_owned(
2442 "target",
2443 1,
2444 None,
2445 &SessionInboxLeaseRequest {
2446 consumer: bamboo_domain::SessionInboxConsumerId::new(),
2447 now,
2448 duration: Duration::seconds(10),
2449 },
2450 )
2451 .await
2452 .unwrap()
2453 .remove(0);
2454 let mut registration = router.register_run("target", "run-1").await.unwrap();
2455 registration.begin_finalization().await;
2456 assert_eq!(registration.finish(0).await.unwrap(), None);
2457 let deferred = router
2458 .reconcile_wake("target", claim.lease.expires_at - Duration::seconds(1))
2459 .await
2460 .unwrap();
2461 assert_eq!(deferred.disposition, None);
2462 assert_eq!(deferred.next_due_at, Some(claim.lease.expires_at));
2463 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 1);
2464 assert_eq!(
2465 router
2466 .reconcile_wake("target", claim.lease.expires_at)
2467 .await
2468 .unwrap()
2469 .disposition,
2470 Some(SessionActivationDisposition::ActivationReserved)
2471 );
2472 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
2473 }
2474
2475 #[tokio::test]
2476 async fn older_staged_release_after_immediate_ack_dispatches_once_without_poison_hot_loop() {
2477 use bamboo_domain::{
2478 Session, SessionActivationPolicy, SessionInboxLimits, SessionMessageEnvelope, Storage,
2479 };
2480 use bamboo_storage::{FileSessionInbox, SessionStoreV2};
2481 async fn release(
2482 inbox: &FileSessionInbox,
2483 receipt: &bamboo_domain::SessionInboxReceipt,
2484 precise: bool,
2485 ) {
2486 if precise {
2487 inbox
2488 .mark_message_activation_eligible(
2489 "parent",
2490 receipt,
2491 SessionActivationPolicy::InterruptSpecificWait,
2492 )
2493 .await
2494 .unwrap();
2495 assert_eq!(
2496 inbox
2497 .coordinator_activation_generation("parent")
2498 .await
2499 .unwrap(),
2500 0
2501 );
2502 } else {
2503 inbox
2504 .mark_activation_eligible(
2505 "parent",
2506 receipt.generation,
2507 SessionActivationPolicy::RespectSpecificWait,
2508 )
2509 .await
2510 .unwrap();
2511 }
2512 }
2513 for (release_during_finalization, precise) in
2514 [(false, false), (true, false), (false, true), (true, true)]
2515 {
2516 let temp = tempfile::TempDir::new().unwrap();
2517 let store = Arc::new(SessionStoreV2::new(temp.path().into()).await.unwrap());
2518 store
2519 .save_session(&Session::new("parent", "model"))
2520 .await
2521 .unwrap();
2522 let inbox = Arc::new(FileSessionInbox::new(
2523 store.clone(),
2524 SessionInboxLimits::default(),
2525 ));
2526 let staged = SessionMessageEnvelope::user_input("parent", "staged result");
2527 let first = inbox.deliver(&staged).await.unwrap();
2528 let immediate = SessionMessageEnvelope::user_input("parent", "immediate steer");
2529 let second = inbox
2530 .deliver_with_activation_intent(
2531 &immediate,
2532 SessionActivationPolicy::RespectSpecificWait,
2533 None,
2534 )
2535 .await
2536 .unwrap();
2537 let reopened_store = Arc::new(SessionStoreV2::new(temp.path().into()).await.unwrap());
2538 let restarted = Arc::new(FileSessionInbox::new(
2539 reopened_store,
2540 SessionInboxLimits::default(),
2541 ));
2542 let router = SessionActivationRouter::new();
2543 router.set_inbox(restarted.clone());
2544 let spawner = spawner();
2545 router.set_spawner(spawner.clone()).await;
2546 let _observer = router.subscribe("parent").await;
2549 let backlog = restarted.inspect("parent").await.unwrap();
2550 assert_eq!(
2551 router
2552 .request_activation("parent", backlog.activation_generation)
2553 .await
2554 .unwrap(),
2555 SessionActivationDisposition::ActivationReserved
2556 );
2557 let claims = restarted.claim("parent", 128).await.unwrap();
2558 assert_eq!(claims.len(), 1);
2559 assert_eq!(claims[0].generation, second.generation);
2560 restarted.ack("parent", &claims[0]).await.unwrap();
2561 let mut registration = router.register_run("parent", "run-1").await.unwrap();
2562 registration.begin_finalization().await;
2563 if release_during_finalization {
2564 release(&restarted, &first, precise).await;
2565 assert_eq!(
2566 router
2567 .request_activation("parent", first.generation)
2568 .await
2569 .unwrap(),
2570 SessionActivationDisposition::ActivationCoalesced
2571 );
2572 assert_eq!(
2573 registration.finish(second.generation).await.unwrap(),
2574 Some(SessionActivationDisposition::ActivationReserved)
2575 );
2576 } else {
2577 assert_eq!(registration.finish(second.generation).await.unwrap(), None);
2578 release(&restarted, &first, precise).await;
2579 assert_eq!(
2580 router
2581 .request_activation("parent", first.generation)
2582 .await
2583 .unwrap(),
2584 SessionActivationDisposition::ActivationReserved
2585 );
2586 }
2587 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
2588 let claims = restarted.claim("parent", 128).await.unwrap();
2589 assert_eq!(claims.len(), 1);
2590 assert_eq!(claims[0].envelope.id, staged.id);
2591 let mut registration = router.register_run("parent", "run-2").await.unwrap();
2594 registration.begin_finalization().await;
2595 assert_eq!(registration.finish(second.generation).await.unwrap(), None);
2596 assert_eq!(
2597 router
2598 .request_activation("parent", first.generation)
2599 .await
2600 .unwrap(),
2601 SessionActivationDisposition::ActivationCoalesced
2602 );
2603 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
2604 restarted.ack("parent", &claims[0]).await.unwrap();
2605 }
2606 }
2607
2608 #[tokio::test]
2609 async fn precise_older_interrupt_below_released_prefix_reaches_one_successor() {
2610 use bamboo_domain::{
2611 Session, SessionActivationPolicy, SessionInboxLimits, SessionMessageEnvelope, Storage,
2612 };
2613 use bamboo_storage::{FileSessionInbox, SessionStoreV2};
2614 for with_older_sibling in [false, true] {
2615 let temp = tempfile::TempDir::new().unwrap();
2616 let store = Arc::new(SessionStoreV2::new(temp.path().into()).await.unwrap());
2617 store
2618 .save_session(&Session::new("parent", "model"))
2619 .await
2620 .unwrap();
2621 let inbox = Arc::new(FileSessionInbox::new(store, SessionInboxLimits::default()));
2622 let sibling = if with_older_sibling {
2623 Some(
2624 inbox
2625 .deliver(&SessionMessageEnvelope::user_input(
2626 "parent",
2627 "older Respect sibling",
2628 ))
2629 .await
2630 .unwrap(),
2631 )
2632 } else {
2633 None
2634 };
2635 let older = inbox
2636 .deliver(&SessionMessageEnvelope::user_input("parent", "older child"))
2637 .await
2638 .unwrap();
2639 let newer = inbox
2640 .deliver(&SessionMessageEnvelope::user_input("parent", "newer Bash"))
2641 .await
2642 .unwrap();
2643 inbox
2644 .mark_activation_eligible(
2645 "parent",
2646 newer.generation,
2647 SessionActivationPolicy::RespectSpecificWait,
2648 )
2649 .await
2650 .unwrap();
2651 let router = SessionActivationRouter::new();
2652 router.set_inbox(inbox.clone());
2653 let spawner = spawner();
2654 router.set_spawner(spawner.clone()).await;
2655 router
2656 .request_activation("parent", newer.generation)
2657 .await
2658 .unwrap();
2659 let claims = inbox.claim("parent", 128).await.unwrap();
2660 assert_eq!(claims.len(), 2 + usize::from(with_older_sibling));
2661 let older_claim = claims
2662 .iter()
2663 .find(|claim| claim.envelope.id == older.id)
2664 .unwrap();
2665 assert_eq!(older_claim.generation, older.generation);
2666 assert_eq!(
2667 older_claim.activation_policy,
2668 SessionActivationPolicy::RespectSpecificWait
2669 );
2670 let newer_claim = claims
2671 .iter()
2672 .find(|claim| claim.envelope.id == newer.id)
2673 .unwrap();
2674 inbox.ack("parent", newer_claim).await.unwrap();
2675 let mut owner = router.register_run("parent", "run-1").await.unwrap();
2676 owner.begin_finalization().await;
2677 assert_eq!(owner.finish(newer.generation).await.unwrap(), None);
2678 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 1);
2679
2680 inbox
2681 .mark_message_activation_eligible(
2682 "parent",
2683 &older,
2684 SessionActivationPolicy::InterruptSpecificWait,
2685 )
2686 .await
2687 .unwrap();
2688 let candidate = inbox
2689 .inspect_wake_readiness("parent", Utc::now())
2690 .await
2691 .unwrap()
2692 .ready
2693 .unwrap();
2694 assert_eq!(candidate.coordinator_generation, newer.generation);
2695 assert_eq!(
2696 candidate.id, older.id,
2697 "a Respect sibling cannot obscure the precise wake"
2698 );
2699 assert_eq!(
2700 candidate.coordinator_message_policy,
2701 Some(SessionActivationPolicy::InterruptSpecificWait)
2702 );
2703 assert_eq!(
2704 router
2705 .request_activation("parent", older.generation)
2706 .await
2707 .unwrap(),
2708 SessionActivationDisposition::ActivationReserved
2709 );
2710 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
2711 let mut successor = router.register_run("parent", "run-2").await.unwrap();
2713 successor.begin_finalization().await;
2714 assert_eq!(successor.finish(newer.generation).await.unwrap(), None);
2715 assert_eq!(
2716 router
2717 .request_activation("parent", older.generation)
2718 .await
2719 .unwrap(),
2720 SessionActivationDisposition::ActivationCoalesced
2721 );
2722 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
2723 inbox.ack("parent", older_claim).await.unwrap();
2724 if let Some(sibling) = sibling {
2725 let claim = claims
2726 .iter()
2727 .find(|claim| claim.envelope.id == sibling.id)
2728 .unwrap();
2729 inbox.ack("parent", claim).await.unwrap();
2730 }
2731 }
2732 }
2733
2734 #[tokio::test]
2735 async fn precise_older_grant_after_terminal_inspection_reaches_one_successor() {
2736 use bamboo_domain::{
2737 Session, SessionActivationPolicy, SessionInboxLimits, SessionMessageEnvelope, Storage,
2738 };
2739 use bamboo_storage::{FileSessionInbox, SessionStoreV2};
2740 for abandon in [false, true] {
2741 let temp = tempfile::TempDir::new().unwrap();
2742 let store = Arc::new(SessionStoreV2::new(temp.path().into()).await.unwrap());
2743 store
2744 .save_session(&Session::new("parent", "model"))
2745 .await
2746 .unwrap();
2747 let inbox = Arc::new(FileSessionInbox::new(store, SessionInboxLimits::default()));
2748 let older = inbox
2749 .deliver(&SessionMessageEnvelope::user_input(
2750 "parent",
2751 "staged child",
2752 ))
2753 .await
2754 .unwrap();
2755 let newer = inbox
2756 .deliver_with_activation_intent(
2757 &SessionMessageEnvelope::user_input("parent", "newer user input"),
2758 SessionActivationPolicy::InterruptSpecificWait,
2759 None,
2760 )
2761 .await
2762 .unwrap();
2763 let router = SessionActivationRouter::new();
2764 router.set_inbox(inbox.clone());
2765 let spawner = spawner();
2766 router.set_spawner(spawner.clone()).await;
2767 router
2768 .request_activation("parent", newer.generation)
2769 .await
2770 .unwrap();
2771 let claims = inbox.claim("parent", 128).await.unwrap();
2772 assert_eq!(claims.len(), 1);
2773 inbox.ack("parent", &claims[0]).await.unwrap();
2774
2775 let mut owner = router.register_run("parent", "run-1").await.unwrap();
2776 owner.begin_finalization().await;
2777 let entered = Arc::new(Notify::new());
2778 let release = Arc::new(Notify::new());
2779 *router.terminal_inspection_pause.lock().await =
2780 Some((entered.clone(), release.clone()));
2781 let finalizing = tokio::spawn(async move {
2782 if abandon {
2783 owner.abandon().await;
2784 None
2785 } else {
2786 owner.finish(newer.generation).await.unwrap()
2787 }
2788 });
2789 entered.notified().await;
2790 assert!(!inbox.inspect("parent").await.unwrap().activation_pending());
2791 inbox
2792 .mark_message_activation_eligible(
2793 "parent",
2794 &older,
2795 SessionActivationPolicy::InterruptSpecificWait,
2796 )
2797 .await
2798 .unwrap();
2799 assert_eq!(
2800 router
2801 .request_activation("parent", older.generation)
2802 .await
2803 .unwrap(),
2804 SessionActivationDisposition::ActivationCoalesced
2805 );
2806 release.notify_one();
2807 let outcome = tokio::time::timeout(std::time::Duration::from_secs(2), finalizing)
2808 .await
2809 .unwrap()
2810 .unwrap();
2811 assert_eq!(
2812 outcome,
2813 (!abandon).then_some(SessionActivationDisposition::ActivationReserved)
2814 );
2815 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
2816 assert!(router.owns_run("parent", "run-2").await);
2817
2818 let claims = inbox.claim("parent", 128).await.unwrap();
2821 assert_eq!(claims.len(), 1);
2822 let mut successor = router.register_run("parent", "run-2").await.unwrap();
2823 successor.begin_finalization().await;
2824 assert_eq!(successor.finish(2).await.unwrap(), None);
2825 assert_eq!(
2826 router
2827 .request_activation("parent", older.generation)
2828 .await
2829 .unwrap(),
2830 SessionActivationDisposition::ActivationCoalesced
2831 );
2832 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
2833 inbox.ack("parent", &claims[0]).await.unwrap();
2834 }
2835 }
2836
2837 struct RecordingSpawner {
2838 reservations: AtomicUsize,
2839 launches: Arc<AtomicUsize>,
2840 entered: Option<Arc<Barrier>>,
2841 release: Option<Arc<Notify>>,
2842 }
2843
2844 #[async_trait]
2845 impl SessionActivationSpawner for RecordingSpawner {
2846 async fn reserve_activation(
2847 &self,
2848 _target_session_id: &str,
2849 _inbox_generation: u64,
2850 ) -> Result<SessionActivationReserveOutcome, SessionActivationError> {
2851 let ordinal = self.reservations.fetch_add(1, Ordering::SeqCst) + 1;
2852 if let Some(entered) = &self.entered {
2853 entered.wait().await;
2854 }
2855 if let Some(release) = &self.release {
2856 release.notified().await;
2857 }
2858 let launches = self.launches.clone();
2859 Ok(SessionActivationReserveOutcome::Reserved(
2860 SessionActivationLaunch::new(format!("run-{ordinal}"), move || {
2861 launches.fetch_add(1, Ordering::SeqCst);
2862 }),
2863 ))
2864 }
2865 }
2866
2867 fn spawner() -> Arc<RecordingSpawner> {
2868 Arc::new(RecordingSpawner {
2869 reservations: AtomicUsize::new(0),
2870 launches: Arc::new(AtomicUsize::new(0)),
2871 entered: None,
2872 release: None,
2873 })
2874 }
2875
2876 struct BlockingInspectInbox {
2877 entered: Arc<Notify>,
2878 release: Arc<Notify>,
2879 drained: Arc<std::sync::atomic::AtomicBool>,
2880 fail_once: Option<Arc<std::sync::atomic::AtomicBool>>,
2881 }
2882
2883 #[async_trait]
2884 impl SessionInboxPort for BlockingInspectInbox {
2885 async fn inspect_wake_readiness(
2886 &self,
2887 _target_session_id: &str,
2888 _now: DateTime<Utc>,
2889 ) -> Result<bamboo_domain::SessionInboxWakeReadiness, bamboo_domain::SessionInboxError>
2890 {
2891 Ok(bamboo_domain::SessionInboxWakeReadiness::default())
2894 }
2895
2896 async fn coordinator_activation_generation(
2897 &self,
2898 _target_session_id: &str,
2899 ) -> Result<u64, bamboo_domain::SessionInboxError> {
2900 Ok(1)
2901 }
2902
2903 async fn deliver(
2904 &self,
2905 _envelope: &bamboo_domain::SessionMessageEnvelope,
2906 ) -> Result<bamboo_domain::SessionInboxReceipt, bamboo_domain::SessionInboxError> {
2907 unreachable!("router cleanup only inspects the durable inbox")
2908 }
2909
2910 async fn mark_activation_eligible(
2911 &self,
2912 _target_session_id: &str,
2913 _generation: u64,
2914 _policy: bamboo_domain::SessionActivationPolicy,
2915 ) -> Result<(), bamboo_domain::SessionInboxError> {
2916 unreachable!("router cleanup only inspects the durable inbox")
2917 }
2918
2919 async fn claim(
2920 &self,
2921 _target_session_id: &str,
2922 _limit: usize,
2923 ) -> Result<Vec<bamboo_domain::SessionInboxClaim>, bamboo_domain::SessionInboxError>
2924 {
2925 unreachable!("router cleanup only inspects the durable inbox")
2926 }
2927
2928 async fn was_admitted(
2929 &self,
2930 _target_session_id: &str,
2931 _id: &bamboo_domain::SessionMessageId,
2932 ) -> Result<bool, bamboo_domain::SessionInboxError> {
2933 unreachable!("router cleanup only inspects the durable inbox")
2934 }
2935
2936 async fn ack(
2937 &self,
2938 _target_session_id: &str,
2939 _claim: &bamboo_domain::SessionInboxClaim,
2940 ) -> Result<(), bamboo_domain::SessionInboxError> {
2941 unreachable!("router cleanup only inspects the durable inbox")
2942 }
2943
2944 async fn inspect(
2945 &self,
2946 _target_session_id: &str,
2947 ) -> Result<bamboo_domain::SessionInboxBacklog, bamboo_domain::SessionInboxError> {
2948 if self
2949 .fail_once
2950 .as_ref()
2951 .is_some_and(|flag| flag.swap(false, Ordering::SeqCst))
2952 {
2953 return Err(bamboo_domain::SessionInboxError::Storage(
2954 "injected inspect failure".into(),
2955 ));
2956 }
2957 if self.drained.load(Ordering::SeqCst) {
2958 return Ok(bamboo_domain::SessionInboxBacklog {
2959 generation: 1,
2960 activation_generation: 1,
2961 ..Default::default()
2962 });
2963 }
2964 self.entered.notify_one();
2965 self.release.notified().await;
2966 Ok(bamboo_domain::SessionInboxBacklog {
2967 pending: 1,
2968 claimed: 0,
2969 generation: 1,
2970 activation_generation: 1,
2971 coordinator_generation: 1,
2972 interrupt_generation: 1,
2973 oldest_generation: Some(1),
2974 })
2975 }
2976 }
2977
2978 #[tokio::test]
2979 async fn rollback_completion_is_retained_before_waiter_subscribes() {
2980 let completion = RollbackCompletion::new();
2981 completion.complete();
2982 tokio::time::timeout(std::time::Duration::from_millis(100), completion.wait())
2983 .await
2984 .expect("completion-before-wait must not lose its wake");
2985 }
2986
2987 #[derive(Clone, Copy)]
2988 enum InjectedFirstOutcome {
2989 NoWork,
2990 Error,
2991 }
2992
2993 struct RetrySpawner {
2994 reservations: AtomicUsize,
2995 launches: Arc<AtomicUsize>,
2996 first_outcome: InjectedFirstOutcome,
2997 first_entered: Arc<Barrier>,
2998 release_first: Arc<Notify>,
2999 }
3000
3001 struct CancelFirstSpawner {
3002 reservations: AtomicUsize,
3003 launches: Arc<AtomicUsize>,
3004 first_entered: Arc<Notify>,
3005 }
3006
3007 struct PanicFirstSpawner {
3008 reservations: AtomicUsize,
3009 launches: Arc<AtomicUsize>,
3010 }
3011
3012 struct AlreadyRunningThenReserveSpawner {
3013 reservations: AtomicUsize,
3014 launches: Arc<AtomicUsize>,
3015 }
3016
3017 #[async_trait]
3018 impl SessionActivationSpawner for AlreadyRunningThenReserveSpawner {
3019 async fn reserve_activation(
3020 &self,
3021 _target_session_id: &str,
3022 _inbox_generation: u64,
3023 ) -> Result<SessionActivationReserveOutcome, SessionActivationError> {
3024 let ordinal = self.reservations.fetch_add(1, Ordering::SeqCst) + 1;
3025 if ordinal == 1 {
3026 return Ok(SessionActivationReserveOutcome::AlreadyRunning {
3027 run_id: "adopted-run".to_string(),
3028 });
3029 }
3030 let launches = self.launches.clone();
3031 Ok(SessionActivationReserveOutcome::Reserved(
3032 SessionActivationLaunch::new("successor-run", move || {
3033 launches.fetch_add(1, Ordering::SeqCst);
3034 }),
3035 ))
3036 }
3037 }
3038
3039 #[async_trait]
3040 impl SessionActivationSpawner for CancelFirstSpawner {
3041 async fn reserve_activation(
3042 &self,
3043 _target_session_id: &str,
3044 _inbox_generation: u64,
3045 ) -> Result<SessionActivationReserveOutcome, SessionActivationError> {
3046 let ordinal = self.reservations.fetch_add(1, Ordering::SeqCst) + 1;
3047 if ordinal == 1 {
3048 self.first_entered.notify_one();
3049 std::future::pending::<()>().await;
3050 }
3051 let launches = self.launches.clone();
3052 Ok(SessionActivationReserveOutcome::Reserved(
3053 SessionActivationLaunch::new(format!("run-{ordinal}"), move || {
3054 launches.fetch_add(1, Ordering::SeqCst);
3055 }),
3056 ))
3057 }
3058 }
3059
3060 #[async_trait]
3061 impl SessionActivationSpawner for PanicFirstSpawner {
3062 async fn reserve_activation(
3063 &self,
3064 _target_session_id: &str,
3065 _inbox_generation: u64,
3066 ) -> Result<SessionActivationReserveOutcome, SessionActivationError> {
3067 let ordinal = self.reservations.fetch_add(1, Ordering::SeqCst) + 1;
3068 if ordinal == 1 {
3069 panic!("injected first activation-spawner panic");
3070 }
3071 let launches = self.launches.clone();
3072 Ok(SessionActivationReserveOutcome::Reserved(
3073 SessionActivationLaunch::new(format!("run-{ordinal}"), move || {
3074 launches.fetch_add(1, Ordering::SeqCst);
3075 }),
3076 ))
3077 }
3078 }
3079
3080 struct RollbackSpawner {
3081 entered: Arc<Barrier>,
3082 release: Arc<Notify>,
3083 launch_ready: Arc<Notify>,
3084 rollbacks: Arc<AtomicUsize>,
3085 }
3086
3087 struct RealRegistryRollbackSpawner {
3088 runners: Arc<
3089 tokio::sync::RwLock<std::collections::HashMap<String, crate::execution::AgentRunner>>,
3090 >,
3091 senders: Arc<
3092 tokio::sync::RwLock<
3093 std::collections::HashMap<
3094 String,
3095 tokio::sync::broadcast::Sender<bamboo_agent_core::AgentEvent>,
3096 >,
3097 >,
3098 >,
3099 sender: tokio::sync::broadcast::Sender<bamboo_agent_core::AgentEvent>,
3100 reservations: AtomicUsize,
3101 first_reserved: Arc<tokio::sync::Notify>,
3102 allow_first_return: Arc<tokio::sync::Notify>,
3103 first_returning: Arc<tokio::sync::Notify>,
3104 rollback_started: Arc<tokio::sync::Notify>,
3105 allow_rollback: Arc<tokio::sync::Notify>,
3106 launches: Arc<AtomicUsize>,
3107 }
3108
3109 #[async_trait]
3110 impl SessionActivationSpawner for RealRegistryRollbackSpawner {
3111 async fn reserve_activation(
3112 &self,
3113 target_session_id: &str,
3114 _inbox_generation: u64,
3115 ) -> Result<SessionActivationReserveOutcome, SessionActivationError> {
3116 let ordinal = self.reservations.fetch_add(1, Ordering::SeqCst) + 1;
3117 match crate::execution::reserve_runner_core(
3118 &self.runners,
3119 &self.senders,
3120 target_session_id,
3121 &self.sender,
3122 )
3123 .await
3124 {
3125 crate::execution::ReserveOutcome::AlreadyRunning(run_id) => {
3126 Ok(SessionActivationReserveOutcome::AlreadyRunning { run_id })
3127 }
3128 crate::execution::ReserveOutcome::Reserved(reservation) => {
3129 let run_id = reservation.run_id.clone();
3130 if ordinal == 1 {
3131 self.first_reserved.notify_one();
3132 self.allow_first_return.notified().await;
3133 self.first_returning.notify_one();
3134 }
3135 let rollback_runners = self.runners.clone();
3136 let rollback_session_id = target_session_id.to_string();
3137 let rollback_run_id = run_id.clone();
3138 let rollback_started = self.rollback_started.clone();
3139 let allow_rollback = self.allow_rollback.clone();
3140 let launches = self.launches.clone();
3141 Ok(SessionActivationReserveOutcome::Reserved(
3142 SessionActivationLaunch::new_with_async_rollback(
3143 run_id,
3144 move || {
3145 launches.fetch_add(1, Ordering::SeqCst);
3146 },
3147 move || async move {
3148 rollback_started.notify_one();
3149 allow_rollback.notified().await;
3150 let mut runners = rollback_runners.write().await;
3151 if runners
3152 .get(&rollback_session_id)
3153 .is_some_and(|runner| runner.run_id == rollback_run_id)
3154 {
3155 runners.remove(&rollback_session_id);
3156 }
3157 },
3158 ),
3159 ))
3160 }
3161 }
3162 }
3163 }
3164
3165 #[async_trait]
3166 impl SessionActivationSpawner for RollbackSpawner {
3167 async fn reserve_activation(
3168 &self,
3169 _target_session_id: &str,
3170 _inbox_generation: u64,
3171 ) -> Result<SessionActivationReserveOutcome, SessionActivationError> {
3172 self.entered.wait().await;
3173 self.release.notified().await;
3174 let rollbacks = self.rollbacks.clone();
3175 let outcome = SessionActivationReserveOutcome::Reserved(
3176 SessionActivationLaunch::new_with_rollback(
3177 "reserved-run",
3178 || {},
3179 move || {
3180 rollbacks.fetch_add(1, Ordering::SeqCst);
3181 },
3182 ),
3183 );
3184 self.launch_ready.notify_one();
3185 Ok(outcome)
3186 }
3187 }
3188
3189 #[async_trait]
3190 impl SessionActivationSpawner for RetrySpawner {
3191 async fn reserve_activation(
3192 &self,
3193 _target_session_id: &str,
3194 _inbox_generation: u64,
3195 ) -> Result<SessionActivationReserveOutcome, SessionActivationError> {
3196 let ordinal = self.reservations.fetch_add(1, Ordering::SeqCst) + 1;
3197 if ordinal == 1 {
3198 self.first_entered.wait().await;
3199 self.release_first.notified().await;
3200 return match self.first_outcome {
3201 InjectedFirstOutcome::NoWork => Ok(SessionActivationReserveOutcome::NoWork),
3202 InjectedFirstOutcome::Error => Err(SessionActivationError::Internal(
3203 "injected first reservation failure".to_string(),
3204 )),
3205 };
3206 }
3207 let launches = self.launches.clone();
3208 Ok(SessionActivationReserveOutcome::Reserved(
3209 SessionActivationLaunch::new(format!("run-{ordinal}"), move || {
3210 launches.fetch_add(1, Ordering::SeqCst);
3211 }),
3212 ))
3213 }
3214 }
3215
3216 async fn assert_newer_generation_retries_after(
3217 first_outcome: InjectedFirstOutcome,
3218 ) -> Result<SessionActivationDisposition, SessionActivationError> {
3219 let router = SessionActivationRouter::new();
3220 let first_entered = Arc::new(Barrier::new(2));
3221 let release_first = Arc::new(Notify::new());
3222 let spawner = Arc::new(RetrySpawner {
3223 reservations: AtomicUsize::new(0),
3224 launches: Arc::new(AtomicUsize::new(0)),
3225 first_outcome,
3226 first_entered: first_entered.clone(),
3227 release_first: release_first.clone(),
3228 });
3229 router.set_spawner(spawner.clone()).await;
3230
3231 let first_router = router.clone();
3232 let first =
3233 tokio::spawn(async move { first_router.request_activation("session", 1).await });
3234 first_entered.wait().await;
3235
3236 let second_router = router.clone();
3237 let second =
3238 tokio::spawn(async move { second_router.request_activation("session", 2).await });
3239 tokio::task::yield_now().await;
3240 assert_eq!(
3241 spawner.reservations.load(Ordering::SeqCst),
3242 1,
3243 "generation 2 must coalesce while generation 1 owns the reservation"
3244 );
3245
3246 release_first.notify_one();
3247 let first_result = first.await.unwrap();
3248 assert_eq!(
3249 second.await.unwrap().unwrap(),
3250 SessionActivationDisposition::ActivationReserved,
3251 "the coalesced newer delivery must take the released reservation"
3252 );
3253 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
3254 assert_eq!(spawner.launches.load(Ordering::SeqCst), 1);
3255 first_result
3256 }
3257
3258 #[tokio::test]
3259 async fn active_owner_is_notified_without_second_reservation() {
3260 let router = SessionActivationRouter::new();
3261 let spawner = spawner();
3262 router.set_spawner(spawner.clone()).await;
3263 let mut registration = router.register_run("session", "run-live").await.unwrap();
3264
3265 assert_eq!(
3266 router.request_activation("session", 7).await.unwrap(),
3267 SessionActivationDisposition::ActiveNotified
3268 );
3269 registration.notifications_mut().changed().await.unwrap();
3270 assert_eq!(*registration.notifications_mut().borrow(), 7);
3271 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 0);
3272 registration.begin_finalization().await;
3273 assert_eq!(registration.finish(7).await.unwrap(), None);
3274 }
3275
3276 #[tokio::test]
3277 async fn different_run_registration_never_overwrites_live_owner() {
3278 let router = SessionActivationRouter::new();
3279 let mut first = router
3280 .register_run("session", "run-1")
3281 .await
3282 .expect("first run owns the logical session");
3283
3284 let error = router
3285 .register_run("session", "run-2")
3286 .await
3287 .err()
3288 .expect("a second exact run must collide");
3289 assert_eq!(
3290 error,
3291 SessionRunRegistrationError::OwnerCollision {
3292 target_session_id: "session".to_string(),
3293 existing_run_id: "run-1".to_string(),
3294 attempted_run_id: "run-2".to_string(),
3295 }
3296 );
3297 assert!(router.owns_run("session", "run-1").await);
3298 assert!(!router.owns_run("session", "run-2").await);
3299
3300 first.begin_finalization().await;
3301 assert_eq!(first.finish(0).await.unwrap(), None);
3302 }
3303
3304 #[tokio::test]
3305 async fn direct_registration_waits_for_activation_to_publish_its_exact_owner() {
3306 let router = SessionActivationRouter::new();
3307 let reservation_entered = Arc::new(Barrier::new(2));
3308 let allow_reservation = Arc::new(Notify::new());
3309 let launches = Arc::new(AtomicUsize::new(0));
3310 let spawner = Arc::new(RecordingSpawner {
3311 reservations: AtomicUsize::new(0),
3312 launches: launches.clone(),
3313 entered: Some(reservation_entered.clone()),
3314 release: Some(allow_reservation.clone()),
3315 });
3316 router.set_spawner(spawner.clone()).await;
3317
3318 let activation_router = router.clone();
3319 let activation =
3320 tokio::spawn(async move { activation_router.request_activation("session", 1).await });
3321 reservation_entered.wait().await;
3322
3323 let direct_router = router.clone();
3324 let direct =
3325 tokio::spawn(async move { direct_router.register_run("session", "manual-run").await });
3326 tokio::task::yield_now().await;
3327 assert!(
3328 !direct.is_finished(),
3329 "manual registration must not supersede an in-flight activation token"
3330 );
3331 assert!(!router.owns_run("session", "manual-run").await);
3332
3333 allow_reservation.notify_one();
3334 assert_eq!(
3335 activation.await.unwrap().unwrap(),
3336 SessionActivationDisposition::ActivationReserved
3337 );
3338 let collision = match direct.await.unwrap() {
3339 Ok(_) => panic!("manual registration must collide with the published activation"),
3340 Err(error) => error,
3341 };
3342 assert_eq!(
3343 collision,
3344 SessionRunRegistrationError::OwnerCollision {
3345 target_session_id: "session".to_string(),
3346 existing_run_id: "run-1".to_string(),
3347 attempted_run_id: "manual-run".to_string(),
3348 }
3349 );
3350 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 1);
3351 assert_eq!(launches.load(Ordering::SeqCst), 1);
3352 assert!(router.owns_run("session", "run-1").await);
3353 assert!(!router.owns_run("session", "manual-run").await);
3354
3355 let mut registration = router.register_run("session", "run-1").await.unwrap();
3356 registration.begin_finalization().await;
3357 assert_eq!(registration.finish(1).await.unwrap(), None);
3358 }
3359
3360 #[tokio::test]
3361 async fn cancelled_explicit_abandon_still_finishes_owned_registration_cleanup() {
3362 let router = SessionActivationRouter::new();
3363 let cleanup_entered = Arc::new(Notify::new());
3364 let cleanup_release = Arc::new(Notify::new());
3365 let cleanup_completed = Arc::new(AtomicUsize::new(0));
3366 let mut registration = router.register_run("session", "run-abandon").await.unwrap();
3367 let entered = cleanup_entered.clone();
3368 let release = cleanup_release.clone();
3369 let completed = cleanup_completed.clone();
3370 registration.set_abort_cleanup(move || async move {
3371 entered.notify_one();
3372 release.notified().await;
3373 completed.fetch_add(1, Ordering::SeqCst);
3374 });
3375
3376 let abandon = tokio::spawn(registration.abandon());
3377 tokio::time::timeout(
3378 std::time::Duration::from_secs(1),
3379 cleanup_entered.notified(),
3380 )
3381 .await
3382 .expect("explicit abandon must start its owned cleanup");
3383 abandon.abort();
3384 assert!(abandon.await.unwrap_err().is_cancelled());
3385 cleanup_release.notify_one();
3386
3387 tokio::time::timeout(std::time::Duration::from_secs(1), async {
3388 loop {
3389 if cleanup_completed.load(Ordering::SeqCst) == 1
3390 && !router.owns_run("session", "run-abandon").await
3391 {
3392 break;
3393 }
3394 tokio::task::yield_now().await;
3395 }
3396 })
3397 .await
3398 .expect("caller cancellation must not cancel detached registration cleanup");
3399 }
3400
3401 #[tokio::test]
3402 async fn delayed_same_run_cannot_adopt_an_owner_during_abandoned_cleanup() {
3403 let router = SessionActivationRouter::new();
3404 let spawner = spawner();
3405 router.set_spawner(spawner.clone()).await;
3406 let inspect_entered = Arc::new(Notify::new());
3407 let inspect_release = Arc::new(Notify::new());
3408 let drained = Arc::new(std::sync::atomic::AtomicBool::new(false));
3409 router.set_inbox(Arc::new(BlockingInspectInbox {
3410 entered: inspect_entered.clone(),
3411 release: inspect_release.clone(),
3412 drained: drained.clone(),
3413 fail_once: None,
3414 }));
3415
3416 let registration = router
3417 .register_run("session", "run-abandoned")
3418 .await
3419 .unwrap();
3420 drop(registration);
3421 tokio::time::timeout(
3422 std::time::Duration::from_secs(1),
3423 inspect_entered.notified(),
3424 )
3425 .await
3426 .expect("abandoned cleanup must reach durable inbox reconciliation");
3427
3428 let error = router
3429 .register_run("session", "run-abandoned")
3430 .await
3431 .err()
3432 .expect("a finalizing abandoned owner is not an adoptable launch placeholder");
3433 assert_eq!(error.existing_run_id(), "run-abandoned");
3434
3435 inspect_release.notify_one();
3436 tokio::time::timeout(std::time::Duration::from_secs(1), async {
3437 while spawner.launches.load(Ordering::SeqCst) == 0 {
3438 tokio::task::yield_now().await;
3439 }
3440 })
3441 .await
3442 .expect("durable pending work must launch one successor");
3443 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 1);
3444 assert_eq!(spawner.launches.load(Ordering::SeqCst), 1);
3445 assert!(!router.owns_run("session", "run-abandoned").await);
3446
3447 let mut successor = router.register_run("session", "run-1").await.unwrap();
3448 drained.store(true, Ordering::SeqCst);
3449 successor.begin_finalization().await;
3450 assert_eq!(successor.finish(1).await.unwrap(), None);
3451 }
3452
3453 #[tokio::test]
3454 async fn delayed_registry_owner_collision_rolls_back_only_its_exact_runner() {
3455 let router = SessionActivationRouter::new();
3456 let runners = Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new()));
3457 let senders = Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new()));
3458 let (sender, _receiver) = tokio::sync::broadcast::channel(8);
3459 let reserved =
3460 match crate::execution::reserve_runner_core(&runners, &senders, "session", &sender)
3461 .await
3462 {
3463 crate::execution::ReserveOutcome::Reserved(reservation) => reservation,
3464 crate::execution::ReserveOutcome::AlreadyRunning(_) => {
3465 panic!("fixture must reserve a fresh server runner")
3466 }
3467 };
3468
3469 let mut direct = router.register_run("session", "sdk-direct").await.unwrap();
3473 let collision = router
3474 .register_run("session", &reserved.run_id)
3475 .await
3476 .err()
3477 .expect("delayed server registration must collide");
3478 let collision_result = Err(bamboo_agent_core::AgentError::LLM(collision.to_string()));
3479 assert!(
3480 crate::execution::finalize_runner_exact(
3481 &runners,
3482 "session",
3483 &reserved.run_id,
3484 &collision_result,
3485 )
3486 .await
3487 );
3488 assert!(router.owns_run("session", "sdk-direct").await);
3489 assert!(matches!(
3490 runners
3491 .read()
3492 .await
3493 .get("session")
3494 .map(|runner| &runner.status),
3495 Some(crate::execution::AgentStatus::Error(_))
3496 ));
3497
3498 direct.begin_finalization().await;
3499 assert_eq!(direct.finish(0).await.unwrap(), None);
3500 }
3501
3502 #[tokio::test]
3503 async fn concurrent_idle_deliveries_coalesce_into_one_launch() {
3504 let router = SessionActivationRouter::new();
3505 let entered = Arc::new(Barrier::new(2));
3506 let release = Arc::new(Notify::new());
3507 let spawner = Arc::new(RecordingSpawner {
3508 reservations: AtomicUsize::new(0),
3509 launches: Arc::new(AtomicUsize::new(0)),
3510 entered: Some(entered.clone()),
3511 release: Some(release.clone()),
3512 });
3513 router.set_spawner(spawner.clone()).await;
3514
3515 let first_router = router.clone();
3516 let first =
3517 tokio::spawn(async move { first_router.request_activation("session", 1).await });
3518 entered.wait().await;
3519 let second_router = router.clone();
3520 let second =
3521 tokio::spawn(async move { second_router.request_activation("session", 2).await });
3522 release.notify_one();
3523 assert_eq!(
3524 first.await.unwrap().unwrap(),
3525 SessionActivationDisposition::ActivationReserved
3526 );
3527 assert_eq!(
3528 second.await.unwrap().unwrap(),
3529 SessionActivationDisposition::ActiveNotified
3530 );
3531 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 1);
3532 assert_eq!(spawner.launches.load(Ordering::SeqCst), 1);
3533 }
3534
3535 #[tokio::test]
3536 async fn newer_generation_re_reserves_after_coalesced_no_work() {
3537 assert_eq!(
3538 assert_newer_generation_retries_after(InjectedFirstOutcome::NoWork)
3539 .await
3540 .unwrap(),
3541 SessionActivationDisposition::ActivationCoalesced
3542 );
3543 }
3544
3545 #[tokio::test]
3546 async fn newer_generation_re_reserves_after_coalesced_error() {
3547 let error = assert_newer_generation_retries_after(InjectedFirstOutcome::Error)
3548 .await
3549 .unwrap_err();
3550 assert!(matches!(error, SessionActivationError::Internal(_)));
3551 }
3552
3553 async fn assert_cancelled_finalization_recovers_without_redelivery(registered_owner: bool) {
3554 let router = SessionActivationRouter::new();
3555 let first_entered = Arc::new(Notify::new());
3556 let spawner = Arc::new(CancelFirstSpawner {
3557 reservations: AtomicUsize::new(0),
3558 launches: Arc::new(AtomicUsize::new(0)),
3559 first_entered: first_entered.clone(),
3560 });
3561 router.set_spawner(spawner.clone()).await;
3562
3563 let finish = if registered_owner {
3564 let mut registration = router.register_run("session", "run-old").await.unwrap();
3565 registration.begin_finalization().await;
3566 assert_eq!(
3567 router.request_activation("session", 1).await.unwrap(),
3568 SessionActivationDisposition::ActivationCoalesced
3569 );
3570 tokio::spawn(async move { registration.finish(0).await })
3571 } else {
3572 router.begin_finalization("session", "run-old").await;
3576 assert_eq!(
3577 router.request_activation("session", 1).await.unwrap(),
3578 SessionActivationDisposition::ActivationCoalesced
3579 );
3580 let finish_router = router.clone();
3581 tokio::spawn(async move {
3582 finish_router
3583 .finish_finalization("session", "run-old", 0)
3584 .await
3585 })
3586 };
3587 first_entered.notified().await;
3588 finish.abort();
3589 assert!(finish.await.unwrap_err().is_cancelled());
3590
3591 tokio::time::timeout(std::time::Duration::from_secs(1), async {
3592 while spawner.launches.load(Ordering::SeqCst) == 0 {
3593 tokio::task::yield_now().await;
3594 }
3595 })
3596 .await
3597 .expect("cancelled finalization must retry pending work without redelivery");
3598 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
3599 assert_eq!(spawner.launches.load(Ordering::SeqCst), 1);
3600 assert!(router.owns_run("session", "run-2").await);
3601
3602 let mut successor = router.register_run("session", "run-2").await.unwrap();
3603 successor.begin_finalization().await;
3604 assert_eq!(successor.finish(1).await.unwrap(), None);
3605 }
3606
3607 #[tokio::test]
3608 async fn cancelled_registered_finalization_retries_without_another_delivery() {
3609 assert_cancelled_finalization_recovers_without_redelivery(true).await;
3610 }
3611
3612 #[tokio::test]
3613 async fn cancelled_reserved_child_finalization_retries_without_another_delivery() {
3614 assert_cancelled_finalization_recovers_without_redelivery(false).await;
3615 }
3616
3617 #[tokio::test]
3618 async fn panicking_finalization_spawner_gets_one_bounded_retry() {
3619 let router = SessionActivationRouter::new();
3620 let spawner = Arc::new(PanicFirstSpawner {
3621 reservations: AtomicUsize::new(0),
3622 launches: Arc::new(AtomicUsize::new(0)),
3623 });
3624 router.set_spawner(spawner.clone()).await;
3625 router.begin_finalization("session", "run-old").await;
3626 assert_eq!(
3627 router.request_activation("session", 1).await.unwrap(),
3628 SessionActivationDisposition::ActivationCoalesced
3629 );
3630
3631 let finish_router = router.clone();
3632 let finish = tokio::spawn(async move {
3633 finish_router
3634 .finish_finalization("session", "run-old", 0)
3635 .await
3636 });
3637 assert!(finish.await.unwrap_err().is_panic());
3638
3639 tokio::time::timeout(std::time::Duration::from_secs(1), async {
3640 while spawner.launches.load(Ordering::SeqCst) == 0 {
3641 tokio::task::yield_now().await;
3642 }
3643 })
3644 .await
3645 .expect("router Drop recovery must survive one spawner panic");
3646 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
3647 assert_eq!(spawner.launches.load(Ordering::SeqCst), 1);
3648
3649 let mut successor = router.register_run("session", "run-2").await.unwrap();
3650 successor.begin_finalization().await;
3651 assert_eq!(successor.finish(1).await.unwrap(), None);
3652 }
3653
3654 #[tokio::test]
3655 async fn cancelled_reservation_releases_token_and_newer_generation_launches() {
3656 let router = SessionActivationRouter::new();
3657 let first_entered = Arc::new(Notify::new());
3658 let spawner = Arc::new(CancelFirstSpawner {
3659 reservations: AtomicUsize::new(0),
3660 launches: Arc::new(AtomicUsize::new(0)),
3661 first_entered: first_entered.clone(),
3662 });
3663 router.set_spawner(spawner.clone()).await;
3664
3665 let first_router = router.clone();
3666 let first =
3667 tokio::spawn(async move { first_router.request_activation("session", 1).await });
3668 first_entered.notified().await;
3669 let second_router = router.clone();
3670 let second =
3671 tokio::spawn(async move { second_router.request_activation("session", 2).await });
3672 tokio::task::yield_now().await;
3673 first.abort();
3674 assert!(first.await.unwrap_err().is_cancelled());
3675
3676 assert_eq!(
3677 tokio::time::timeout(std::time::Duration::from_secs(1), second)
3678 .await
3679 .expect("coalesced caller must not hang behind a cancelled reservation")
3680 .unwrap()
3681 .unwrap(),
3682 SessionActivationDisposition::ActivationReserved
3683 );
3684 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
3685 assert_eq!(spawner.launches.load(Ordering::SeqCst), 1);
3686 }
3687
3688 #[tokio::test]
3689 async fn cancellation_after_external_reservation_rolls_back_unlaunched_slot() {
3690 let router = SessionActivationRouter::new();
3691 let entered = Arc::new(Barrier::new(2));
3692 let release = Arc::new(Notify::new());
3693 let launch_ready = Arc::new(Notify::new());
3694 let rollbacks = Arc::new(AtomicUsize::new(0));
3695 let spawner = Arc::new(RollbackSpawner {
3696 entered: entered.clone(),
3697 release: release.clone(),
3698 launch_ready: launch_ready.clone(),
3699 rollbacks: rollbacks.clone(),
3700 });
3701 router.set_spawner(spawner).await;
3702
3703 let task_router = router.clone();
3704 let task = tokio::spawn(async move { task_router.request_activation("session", 1).await });
3705 entered.wait().await;
3706 let states = router.states.lock().await;
3710 release.notify_one();
3711 launch_ready.notified().await;
3712 task.abort();
3713 assert!(task.await.unwrap_err().is_cancelled());
3714 drop(states);
3715
3716 tokio::time::timeout(std::time::Duration::from_secs(1), async {
3717 while rollbacks.load(Ordering::SeqCst) == 0 {
3718 tokio::task::yield_now().await;
3719 }
3720 })
3721 .await
3722 .expect("unlaunched reservation rollback");
3723 assert_eq!(rollbacks.load(Ordering::SeqCst), 1);
3724 }
3725
3726 #[tokio::test]
3727 async fn redelivery_waits_for_real_runner_rollback_before_reserving() {
3728 let router = SessionActivationRouter::new();
3729 let runners = Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new()));
3730 let senders = Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new()));
3731 let (sender, _receiver) = tokio::sync::broadcast::channel(8);
3732 let first_reserved = Arc::new(tokio::sync::Notify::new());
3733 let allow_first_return = Arc::new(tokio::sync::Notify::new());
3734 let rollback_started = Arc::new(tokio::sync::Notify::new());
3735 let first_returning = Arc::new(tokio::sync::Notify::new());
3736 let allow_rollback = Arc::new(tokio::sync::Notify::new());
3737 let launches = Arc::new(AtomicUsize::new(0));
3738 let spawner = Arc::new(RealRegistryRollbackSpawner {
3739 runners: runners.clone(),
3740 senders,
3741 sender,
3742 reservations: AtomicUsize::new(0),
3743 first_reserved: first_reserved.clone(),
3744 allow_first_return: allow_first_return.clone(),
3745 first_returning: first_returning.clone(),
3746 rollback_started: rollback_started.clone(),
3747 allow_rollback: allow_rollback.clone(),
3748 launches: launches.clone(),
3749 });
3750 router.set_spawner(spawner.clone()).await;
3751
3752 let first_router = router.clone();
3753 let first =
3754 tokio::spawn(async move { first_router.request_activation("session", 1).await });
3755 first_reserved.notified().await;
3756
3757 let state_guard = router.states.lock().await;
3760 allow_first_return.notify_one();
3761 first_returning.notified().await;
3762 first.abort();
3763 assert!(first.await.unwrap_err().is_cancelled());
3764 drop(state_guard);
3765 rollback_started.notified().await;
3766
3767 let second_router = router.clone();
3768 let second =
3769 tokio::spawn(async move { second_router.request_activation("session", 2).await });
3770 tokio::task::yield_now().await;
3771 assert_eq!(
3772 spawner.reservations.load(Ordering::SeqCst),
3773 1,
3774 "redelivery must remain coalesced until exact external rollback completes"
3775 );
3776
3777 allow_rollback.notify_one();
3778 assert_eq!(
3779 tokio::time::timeout(std::time::Duration::from_secs(1), second)
3780 .await
3781 .expect("redelivery must reserve after rollback")
3782 .unwrap()
3783 .unwrap(),
3784 SessionActivationDisposition::ActivationReserved
3785 );
3786 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
3787 assert_eq!(launches.load(Ordering::SeqCst), 1);
3788 let current = runners
3789 .read()
3790 .await
3791 .get("session")
3792 .map(|runner| runner.run_id.clone())
3793 .expect("second real registry reservation remains");
3794 assert_ne!(
3795 current, "",
3796 "the redelivery must never adopt an empty/stale run id"
3797 );
3798 }
3799
3800 #[tokio::test]
3801 async fn delivery_during_finalization_launches_one_successor() {
3802 let router = SessionActivationRouter::new();
3803 let spawner = spawner();
3804 router.set_spawner(spawner.clone()).await;
3805 let mut registration = router.register_run("session", "run-old").await.unwrap();
3806 registration.begin_finalization().await;
3807
3808 assert_eq!(
3809 router.request_activation("session", 11).await.unwrap(),
3810 SessionActivationDisposition::ActivationCoalesced
3811 );
3812 let disposition = registration.finish(10).await.unwrap();
3813 assert_eq!(
3814 disposition,
3815 Some(SessionActivationDisposition::ActivationReserved)
3816 );
3817 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 1);
3818 assert_eq!(spawner.launches.load(Ordering::SeqCst), 1);
3819 }
3820
3821 #[tokio::test]
3822 async fn finalization_does_not_restart_when_cursor_caught_up() {
3823 let router = SessionActivationRouter::new();
3824 let spawner = spawner();
3825 router.set_spawner(spawner.clone()).await;
3826 let mut registration = router.register_run("session", "run-old").await.unwrap();
3827 assert_eq!(
3828 router.request_activation("session", 4).await.unwrap(),
3829 SessionActivationDisposition::ActiveNotified
3830 );
3831 registration.begin_finalization().await;
3832 assert_eq!(registration.finish(4).await.unwrap(), None);
3833 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 0);
3834 assert!(
3835 router.states.lock().await.get("session").is_none(),
3836 "a caught-up completed session must not leak routing state"
3837 );
3838 }
3839
3840 #[tokio::test]
3841 async fn successor_finalization_removes_state_after_raced_generation_catches_up() {
3842 let router = SessionActivationRouter::new();
3843 let spawner = spawner();
3844 router.set_spawner(spawner).await;
3845 let mut registration = router.register_run("session", "run-old").await.unwrap();
3846 registration.begin_finalization().await;
3847
3848 assert_eq!(
3849 router.request_activation("session", 11).await.unwrap(),
3850 SessionActivationDisposition::ActivationCoalesced
3851 );
3852 assert_eq!(
3853 registration.finish(10).await.unwrap(),
3854 Some(SessionActivationDisposition::ActivationReserved)
3855 );
3856 let mut successor_registration = router.register_run("session", "run-1").await.unwrap();
3857 successor_registration.begin_finalization().await;
3858 assert_eq!(successor_registration.finish(11).await.unwrap(), None);
3859 assert!(
3860 router.states.lock().await.get("session").is_none(),
3861 "the successor's caught-up terminal state must be compacted"
3862 );
3863 }
3864
3865 #[tokio::test]
3866 async fn failed_terminal_inspection_keeps_the_exact_owner_cleanup_armed() {
3867 let router = SessionActivationRouter::new();
3868 let spawner = spawner();
3869 router.set_spawner(spawner.clone()).await;
3870 let inspect_release = Arc::new(Notify::new());
3871 inspect_release.notify_one();
3872 router.set_inbox(Arc::new(BlockingInspectInbox {
3873 entered: Arc::new(Notify::new()),
3874 release: inspect_release,
3875 drained: Arc::new(std::sync::atomic::AtomicBool::new(false)),
3876 fail_once: Some(Arc::new(std::sync::atomic::AtomicBool::new(true))),
3877 }));
3878 let mut registration = router
3879 .register_run("session", "failed-inspect-run")
3880 .await
3881 .unwrap();
3882 router.request_activation("session", 1).await.unwrap();
3883 registration.begin_finalization().await;
3884 assert!(registration.finish(1).await.is_err());
3885 tokio::time::timeout(std::time::Duration::from_secs(1), async {
3886 while spawner.launches.load(Ordering::SeqCst) == 0 {
3887 tokio::task::yield_now().await;
3888 }
3889 })
3890 .await
3891 .expect("failed inspection must not leave the old run owning the inbox");
3892 assert!(!router.owns_run("session", "failed-inspect-run").await);
3893 assert_eq!(spawner.launches.load(Ordering::SeqCst), 1);
3894 }
3895
3896 #[tokio::test]
3897 async fn poison_generation_launches_only_one_successor_until_new_work_arrives() {
3898 let router = SessionActivationRouter::new();
3899 let spawner = spawner();
3900 router.set_spawner(spawner.clone()).await;
3901 let inspect_release = Arc::new(Notify::new());
3902 router.set_inbox(Arc::new(BlockingInspectInbox {
3903 entered: Arc::new(Notify::new()),
3904 release: inspect_release.clone(),
3905 drained: Arc::new(std::sync::atomic::AtomicBool::new(false)),
3906 fail_once: None,
3907 }));
3908 let mut registration = router
3909 .register_run("session", "run-original")
3910 .await
3911 .unwrap();
3912 assert_eq!(
3913 router.request_activation("session", 7).await.unwrap(),
3914 SessionActivationDisposition::ActiveNotified
3915 );
3916 registration.begin_finalization().await;
3917 inspect_release.notify_one();
3918 assert_eq!(
3919 registration.finish(0).await.unwrap(),
3920 Some(SessionActivationDisposition::ActivationReserved)
3921 );
3922 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 1);
3923
3924 let mut successor_registration = router.register_run("session", "run-1").await.unwrap();
3928 successor_registration.begin_finalization().await;
3929 inspect_release.notify_one();
3930 assert_eq!(successor_registration.finish(0).await.unwrap(), None);
3931 assert_eq!(
3932 router.request_activation("session", 7).await.unwrap(),
3933 SessionActivationDisposition::ActivationCoalesced
3934 );
3935 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 1);
3936
3937 assert_eq!(
3939 router.request_activation("session", 8).await.unwrap(),
3940 SessionActivationDisposition::ActivationReserved
3941 );
3942 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
3943 }
3944
3945 #[tokio::test]
3946 async fn adopted_terminal_run_preserves_generation_for_one_real_successor() {
3947 let router = SessionActivationRouter::new();
3948 let spawner = Arc::new(AlreadyRunningThenReserveSpawner {
3949 reservations: AtomicUsize::new(0),
3950 launches: Arc::new(AtomicUsize::new(0)),
3951 });
3952 router.set_spawner(spawner.clone()).await;
3953
3954 assert_eq!(
3955 router.request_activation("session", 12).await.unwrap(),
3956 SessionActivationDisposition::ActiveNotified
3957 );
3958 router.begin_finalization("session", "adopted-run").await;
3959 assert_eq!(
3960 router
3961 .finish_finalization("session", "adopted-run", 0)
3962 .await
3963 .unwrap(),
3964 Some(SessionActivationDisposition::ActivationReserved)
3965 );
3966 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
3967 assert_eq!(spawner.launches.load(Ordering::SeqCst), 1);
3968 }
3969 #[tokio::test]
3970 async fn quarantined_registration_never_releases_owner_or_launches_successor() {
3971 let router = SessionActivationRouter::new();
3972 let spawner = spawner();
3973 router.set_spawner(spawner.clone()).await;
3974 let mut registration = router
3975 .register_run("session", "quarantined-run")
3976 .await
3977 .unwrap();
3978 let cleaned = Arc::new(std::sync::atomic::AtomicBool::new(false));
3979 let cleanup_flag = cleaned.clone();
3980 registration.set_abort_cleanup(move || async move {
3981 cleanup_flag.store(true, Ordering::SeqCst);
3982 });
3983 registration.quarantine().await;
3984 assert_eq!(
3985 router.request_activation("session", 19).await.unwrap(),
3986 SessionActivationDisposition::ActivationCoalesced
3987 );
3988 assert!(router.register_run("session", "successor").await.is_err());
3989 assert!(router
3990 .register_run("session", "quarantined-run")
3991 .await
3992 .is_err());
3993 tokio::task::yield_now().await;
3994 assert!(!cleaned.load(Ordering::SeqCst));
3995 assert_eq!(spawner.launches.load(Ordering::SeqCst), 0);
3996 let states = router.states.lock().await;
3997 let owner = states["session"].owner.as_ref().unwrap();
3998 assert_eq!(owner.run_id, "quarantined-run");
3999 assert!(owner.finalizing);
4000 }
4001
4002 #[tokio::test]
4003 async fn physical_execution_registration_drop_is_quarantined() {
4004 let router = SessionActivationRouter::new();
4005 let spawner = spawner();
4006 router.set_spawner(spawner.clone()).await;
4007 let mut registration = router
4008 .register_run("session", "physical-run")
4009 .await
4010 .unwrap();
4011 let cleaned = Arc::new(std::sync::atomic::AtomicBool::new(false));
4012 let cleanup_flag = cleaned.clone();
4013 registration.set_abort_cleanup(move || async move {
4014 cleanup_flag.store(true, Ordering::SeqCst);
4015 });
4016 registration.quarantine_on_drop();
4017 drop(registration);
4018 assert!(matches!(
4022 router.request_activation("session", 7).await.unwrap(),
4023 SessionActivationDisposition::ActiveNotified
4024 | SessionActivationDisposition::ActivationCoalesced
4025 ));
4026 tokio::time::timeout(std::time::Duration::from_secs(1), async {
4027 loop {
4028 if router.states.lock().await["session"]
4029 .owner
4030 .as_ref()
4031 .unwrap()
4032 .finalizing
4033 {
4034 break;
4035 }
4036 tokio::task::yield_now().await;
4037 }
4038 })
4039 .await
4040 .unwrap();
4041 assert_eq!(
4042 router.request_activation("session", 8).await.unwrap(),
4043 SessionActivationDisposition::ActivationCoalesced
4044 );
4045 assert!(!cleaned.load(Ordering::SeqCst));
4046 assert_eq!(spawner.launches.load(Ordering::SeqCst), 0);
4047 assert!(router.register_run("session", "successor").await.is_err());
4048 let states = router.states.lock().await;
4049 assert!(states["session"].owner.as_ref().unwrap().finalizing);
4050 }
4051
4052 #[tokio::test]
4053 async fn cancelling_quarantine_wait_does_not_arm_abandoned_cleanup() {
4054 let router = SessionActivationRouter::new();
4055 let registration = router
4056 .register_run("session", "quarantined-run")
4057 .await
4058 .unwrap();
4059 let states_guard = router.states.lock().await;
4060 let mut quarantine = Box::pin(registration.quarantine());
4061 assert!(futures::poll!(quarantine.as_mut()).is_pending());
4062 drop(quarantine);
4063 drop(states_guard);
4064 tokio::task::yield_now().await;
4065 assert!(router.register_run("session", "successor").await.is_err());
4066 assert_eq!(
4067 router.current_run_id("session").await.as_deref(),
4068 Some("quarantined-run")
4069 );
4070 }
4071 async fn selected_fixture() -> (
4072 tempfile::TempDir,
4073 Arc<SessionActivationRouter>,
4074 Arc<RecordingSpawner>,
4075 Arc<bamboo_storage::FileSessionInbox>,
4076 ) {
4077 use bamboo_domain::{
4078 Session, SessionActivationPolicy, SessionInboxLimits, SessionMessageEnvelope, Storage,
4079 };
4080 let temp = tempfile::tempdir().unwrap();
4081 let store = Arc::new(
4082 bamboo_storage::SessionStoreV2::new(temp.path().into())
4083 .await
4084 .unwrap(),
4085 );
4086 store
4087 .save_session(&Session::new("selected", "model"))
4088 .await
4089 .unwrap();
4090 let inbox = Arc::new(bamboo_storage::FileSessionInbox::new(
4091 store,
4092 SessionInboxLimits::default(),
4093 ));
4094 inbox
4095 .deliver_with_activation_intent(
4096 &SessionMessageEnvelope::user_input("selected", "other eligible input"),
4097 SessionActivationPolicy::RespectSpecificWait,
4098 None,
4099 )
4100 .await
4101 .unwrap();
4102 let router = SessionActivationRouter::new();
4103 router.set_inbox(inbox.clone());
4104 let spawner = spawner();
4105 router.set_spawner(spawner.clone()).await;
4106 (temp, router, spawner, inbox)
4107 }
4108
4109 #[tokio::test]
4110 async fn receipt_run_selected_release_never_inherits_unrequested_durable_backlog() {
4111 for mode in ["finish", "abandon", "drop"] {
4112 let (_temp, router, spawner, inbox) = selected_fixture().await;
4113 let mut selected = router
4114 .register_run("selected", "manual-selected")
4115 .await
4116 .unwrap();
4117 selected.restrict_to_selected_input();
4118 match mode {
4119 "finish" => {
4120 selected.begin_finalization().await;
4121 assert_eq!(selected.finish(0).await.unwrap(), None);
4122 }
4123 "abandon" => selected.abandon().await,
4124 _ => drop(selected),
4125 }
4126 tokio::time::timeout(std::time::Duration::from_secs(2), async {
4127 while router.current_run_id("selected").await.is_some() {
4128 tokio::task::yield_now().await;
4129 }
4130 })
4131 .await
4132 .unwrap();
4133 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 0, "{mode}");
4134 assert_eq!(spawner.launches.load(Ordering::SeqCst), 0);
4135 assert_eq!(inbox.inspect("selected").await.unwrap().pending, 1);
4136 assert_eq!(
4138 router.request_activation("selected", 1).await.unwrap(),
4139 SessionActivationDisposition::ActivationReserved
4140 );
4141 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 1);
4142 }
4143 }
4144
4145 #[tokio::test]
4146 async fn receipt_run_selected_release_preserves_explicit_coalesced_coordinator_wake() {
4147 for mode in ["finish", "abandon", "drop"] {
4148 for notify_before_scope in [false, true] {
4149 let (_temp, router, spawner, _inbox) = selected_fixture().await;
4150 let mut selected = router
4151 .register_run("selected", "manual-selected")
4152 .await
4153 .unwrap();
4154 if !notify_before_scope {
4155 selected.restrict_to_selected_input();
4156 }
4157 assert_eq!(
4158 router.request_activation("selected", 1).await.unwrap(),
4159 SessionActivationDisposition::ActiveNotified
4160 );
4161 if notify_before_scope {
4162 selected.restrict_to_selected_input();
4163 }
4164 match mode {
4165 "finish" => {
4166 selected.begin_finalization().await;
4167 assert_eq!(
4168 selected.finish(0).await.unwrap(),
4169 Some(SessionActivationDisposition::ActivationReserved)
4170 );
4171 }
4172 "abandon" => selected.abandon().await,
4173 _ => drop(selected),
4174 }
4175 tokio::time::timeout(std::time::Duration::from_secs(2), async {
4176 while spawner.reservations.load(Ordering::SeqCst) == 0 {
4177 tokio::task::yield_now().await;
4178 }
4179 })
4180 .await
4181 .unwrap();
4182 assert_eq!(
4183 spawner.reservations.load(Ordering::SeqCst),
4184 1,
4185 "{mode}/{notify_before_scope}"
4186 );
4187 assert_eq!(spawner.launches.load(Ordering::SeqCst), 1);
4188 assert!(router.owns_run("selected", "run-1").await);
4189 }
4190 }
4191 }
4192
4193 #[tokio::test]
4194 async fn receipt_run_cancelled_selected_finish_retains_restricted_drop_cleanup() {
4195 let (_temp, router, spawner, inbox) = selected_fixture().await;
4196 let mut selected = router
4197 .register_run("selected", "manual-selected")
4198 .await
4199 .unwrap();
4200 selected.restrict_to_selected_input();
4201 let lock = router.states.lock().await;
4202 let mut finish = Box::pin(selected.finish(0));
4203 assert!(futures::poll!(finish.as_mut()).is_pending());
4204 drop(finish);
4205 drop(lock);
4206 tokio::time::timeout(std::time::Duration::from_secs(2), async {
4207 while router.current_run_id("selected").await.is_some() {
4208 tokio::task::yield_now().await;
4209 }
4210 })
4211 .await
4212 .unwrap();
4213 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 0);
4214 assert_eq!(inbox.inspect("selected").await.unwrap().pending, 1);
4215 }
4216 struct PeerIntentFixture {
4217 _temp: tempfile::TempDir,
4218 store: Arc<bamboo_storage::SessionStoreV2>,
4219 inbox: Arc<bamboo_storage::FileSessionInbox>,
4220 router: Arc<SessionActivationRouter>,
4221 spawner: Arc<RecordingSpawner>,
4222 messenger: crate::SessionMessenger,
4223 }
4224
4225 impl PeerIntentFixture {
4226 async fn new() -> Self {
4227 use bamboo_domain::{Session, SessionInboxLimits, Storage};
4228 let temp = tempfile::tempdir().unwrap();
4229 let store = Arc::new(
4230 bamboo_storage::SessionStoreV2::new(temp.path().into())
4231 .await
4232 .unwrap(),
4233 );
4234 let target = Session::new("target", "model");
4235 let source = Session::new_child_of("sender", &target, "model", "task");
4236 store.save_session(&target).await.unwrap();
4237 store.save_session(&source).await.unwrap();
4238 let inbox = Arc::new(bamboo_storage::FileSessionInbox::new(
4239 store.clone(),
4240 SessionInboxLimits::default(),
4241 ));
4242 let router = SessionActivationRouter::new();
4243 router.set_inbox(inbox.clone());
4244 let spawner = spawner();
4245 router.set_spawner(spawner.clone()).await;
4246 let messenger =
4247 crate::SessionMessenger::new(store.clone(), inbox.clone(), router.clone());
4248 Self {
4249 _temp: temp,
4250 store,
4251 inbox,
4252 router,
4253 spawner,
4254 messenger,
4255 }
4256 }
4257 }
4258
4259 #[tokio::test]
4260 async fn queue_only_delivery_never_schedules_idle_finalizing_or_restarted_target() {
4261 use bamboo_domain::SessionMessageDeliveryIntent::QueueOnly;
4262 for timing in [
4263 "idle",
4264 "active",
4265 "finalizing",
4266 "after_terminal_inspection",
4267 "selected_finish",
4268 "selected_abandon",
4269 ] {
4270 let fixture = PeerIntentFixture::new().await;
4271 let mut registration = if timing == "idle" {
4272 None
4273 } else {
4274 Some(
4275 fixture
4276 .router
4277 .register_run("target", "manual-run")
4278 .await
4279 .unwrap(),
4280 )
4281 };
4282 if timing == "selected_finish" || timing == "selected_abandon" {
4283 registration.as_mut().unwrap().restrict_to_selected_input();
4284 }
4285 if timing == "finalizing" || timing == "after_terminal_inspection" {
4286 registration.as_mut().unwrap().begin_finalization().await;
4287 }
4288 let mut race = None;
4289 let release = Arc::new(Notify::new());
4290 if timing == "after_terminal_inspection" {
4291 let entered = Arc::new(Notify::new());
4292 *fixture.router.terminal_inspection_pause.lock().await =
4293 Some((entered.clone(), release.clone()));
4294 let registration = registration.take().unwrap();
4295 race = Some(tokio::spawn(async move { registration.finish(0).await }));
4296 tokio::time::timeout(std::time::Duration::from_secs(2), entered.notified())
4297 .await
4298 .unwrap();
4299 }
4300 let receipt = fixture
4301 .messenger
4302 .send_peer_with_intent("sender", "target", "queue", "queued", QueueOnly)
4303 .await
4304 .unwrap();
4305 assert_eq!(receipt.activation, None, "{timing}");
4306 assert_eq!(
4307 receipt.active_delivery_notified,
4308 matches!(timing, "active" | "selected_finish" | "selected_abandon")
4309 );
4310 if let Some(mut registration) = registration {
4311 if timing == "selected_abandon" {
4312 registration.abandon().await;
4313 } else {
4314 registration.begin_finalization().await;
4315 assert_eq!(registration.finish(0).await.unwrap(), None);
4316 }
4317 }
4318 if let Some(race) = race {
4319 release.notify_one();
4320 assert_eq!(race.await.unwrap().unwrap(), None);
4321 }
4322 assert_eq!(
4323 fixture.spawner.reservations.load(Ordering::SeqCst),
4324 0,
4325 "{timing}"
4326 );
4327 assert_eq!(fixture.spawner.launches.load(Ordering::SeqCst), 0);
4328 assert_eq!(fixture.inbox.inspect("target").await.unwrap().pending, 1);
4329 let restarted = SessionActivationRouter::new();
4331 restarted.set_inbox(fixture.inbox.clone());
4332 restarted.set_spawner(fixture.spawner.clone()).await;
4333 assert_eq!(
4334 restarted
4335 .reconcile_wake("target", Utc::now())
4336 .await
4337 .unwrap()
4338 .disposition,
4339 None
4340 );
4341 assert_eq!(fixture.spawner.reservations.load(Ordering::SeqCst), 0);
4342 }
4343 }
4344
4345 #[tokio::test]
4346 async fn queue_only_notifies_late_bound_active_driver_and_checkpoints_at_safe_point() {
4347 use bamboo_domain::{SessionMessageDeliveryIntent::QueueOnly, Storage};
4348 let fixture = PeerIntentFixture::new().await;
4349 let mut registration = fixture
4350 .router
4351 .register_run("target", "manual-run")
4352 .await
4353 .unwrap();
4354 let receipt = fixture
4355 .messenger
4356 .send_peer_with_intent("sender", "target", "queue", "queued", QueueOnly)
4357 .await
4358 .unwrap();
4359 assert!(receipt.active_delivery_notified);
4360 assert!(registration.notifications_mut().has_changed().unwrap());
4361 let (tx, mut rx) = mpsc::unbounded_channel();
4362 assert_eq!(
4363 fixture
4364 .router
4365 .attach_delivery_sink("target", tx)
4366 .await
4367 .as_deref(),
4368 Some("manual-run")
4369 );
4370 assert_eq!(rx.try_recv().unwrap(), receipt.delivery.generation);
4371 let claims = fixture
4372 .inbox
4373 .claim_for_turn("target", 8, Some("manual-run"))
4374 .await
4375 .unwrap();
4376 assert_eq!(claims.len(), 1);
4377 let mut saved = fixture.store.load_session("target").await.unwrap().unwrap();
4378 saved.add_message(claims[0].envelope.to_provider_message().unwrap());
4379 fixture.store.save_session(&saved).await.unwrap();
4380 fixture.inbox.ack("target", &claims[0]).await.unwrap();
4381 registration.begin_finalization().await;
4382 assert_eq!(
4383 registration
4384 .finish(receipt.delivery.generation)
4385 .await
4386 .unwrap(),
4387 None
4388 );
4389 assert_eq!(fixture.spawner.reservations.load(Ordering::SeqCst), 0);
4390 let saved = fixture.store.load_session("target").await.unwrap().unwrap();
4391 assert_eq!(
4392 saved
4393 .messages
4394 .iter()
4395 .filter(|message| message.id == receipt.delivery.id.as_str())
4396 .count(),
4397 1
4398 );
4399 }
4400
4401 #[tokio::test]
4402 async fn explicit_trigger_turn_can_start_a_turn_after_queue_only_delivery() {
4403 use bamboo_domain::SessionMessageDeliveryIntent::{QueueOnly, TriggerTurn};
4404 let fixture = PeerIntentFixture::new().await;
4405 fixture
4406 .messenger
4407 .send_peer_with_intent("sender", "target", "queue", "queued", QueueOnly)
4408 .await
4409 .unwrap();
4410 assert_eq!(fixture.spawner.reservations.load(Ordering::SeqCst), 0);
4411 let followup = fixture
4412 .messenger
4413 .send_peer_with_intent("sender", "target", "followup", "followup", TriggerTurn)
4414 .await
4415 .unwrap();
4416 assert_eq!(
4417 followup.activation,
4418 Some(SessionActivationDisposition::ActivationReserved)
4419 );
4420 assert_eq!(fixture.spawner.reservations.load(Ordering::SeqCst), 1);
4421 assert_eq!(fixture.spawner.launches.load(Ordering::SeqCst), 1);
4422 let claims = fixture
4423 .inbox
4424 .claim_for_turn("target", 8, Some("run-1"))
4425 .await
4426 .unwrap();
4427 assert_eq!(claims.len(), 2);
4428 assert!(claims.iter().all(|claim| claim.activation_policy
4429 == bamboo_domain::SessionActivationPolicy::RespectSpecificWait));
4430 }
4431
4432 #[tokio::test]
4433 async fn acknowledged_explicit_peer_retry_cannot_activate_or_notify_a_new_owner() {
4434 use bamboo_domain::{SessionMessageDeliveryIntent, Storage};
4435 for intent in [
4436 SessionMessageDeliveryIntent::TriggerTurn,
4437 SessionMessageDeliveryIntent::QueueOnly,
4438 ] {
4439 for active_owner in [false, true] {
4440 let fixture = PeerIntentFixture::new().await;
4441 let first = fixture
4442 .messenger
4443 .send_peer_with_intent("sender", "target", "ack-operation", "input", intent)
4444 .await
4445 .unwrap();
4446 let claim = fixture
4447 .inbox
4448 .claim_for_turn("target", 1, Some("first-run"))
4449 .await
4450 .unwrap()
4451 .remove(0);
4452 let mut saved = fixture.store.load_session("target").await.unwrap().unwrap();
4453 saved.add_message(claim.envelope.to_provider_message().unwrap());
4454 saved
4455 .session_inbox_admission_mut()
4456 .record(claim.envelope.id.clone(), claim.generation);
4457 fixture.store.save_session(&saved).await.unwrap();
4458 fixture.inbox.ack("target", &claim).await.unwrap();
4459
4460 let restarted = SessionActivationRouter::new();
4463 restarted.set_inbox(fixture.inbox.clone());
4464 let fresh_spawner = spawner();
4465 restarted.set_spawner(fresh_spawner.clone()).await;
4466 let mut registration = if active_owner {
4467 Some(restarted.register_run("target", "new-owner").await.unwrap())
4468 } else {
4469 None
4470 };
4471 let messenger = crate::SessionMessenger::new(
4472 fixture.store.clone(),
4473 fixture.inbox.clone(),
4474 restarted,
4475 );
4476 let retry = messenger
4477 .send_peer_with_intent("sender", "target", "ack-operation", "input", intent)
4478 .await
4479 .unwrap();
4480 assert_eq!(retry.delivery, first.delivery);
4481 assert_eq!(retry.activation, None, "{intent:?} active={active_owner}");
4482 assert!(!retry.active_delivery_notified);
4483 assert_eq!(fresh_spawner.reservations.load(Ordering::SeqCst), 0);
4484 assert_eq!(fresh_spawner.launches.load(Ordering::SeqCst), 0);
4485 if let Some(mut registration) = registration.take() {
4486 assert!(!registration.notifications_mut().has_changed().unwrap());
4487 registration.begin_finalization().await;
4488 assert_eq!(registration.finish(claim.generation).await.unwrap(), None);
4489 }
4490 }
4491 }
4492 }
4493 #[tokio::test]
4494 async fn expected_run_guard_conflicts_and_blocks_finalization_until_publication_finishes() {
4495 use bamboo_domain::{SessionRunConflictReason, SessionRunDeliveryGuard};
4496 async fn retain(
4497 router: Arc<SessionActivationRouter>,
4498 birth: DateTime<Utc>,
4499 expected: &'static str,
4500 ) -> Result<Arc<dyn SessionRunDeliveryGuard>, bamboo_domain::SessionRunConflict> {
4501 tokio::task::spawn_blocking(move || {
4502 router.retain_expected_run_blocking("target", birth, expected)
4503 })
4504 .await
4505 .unwrap()
4506 }
4507 let router = SessionActivationRouter::new();
4508 let target = bamboo_domain::Session::new("target", "model");
4509 let birth = target.created_at;
4510 assert_eq!(
4511 retain(router.clone(), birth, "run-a")
4512 .await
4513 .err()
4514 .unwrap()
4515 .reason,
4516 SessionRunConflictReason::NoActiveRun
4517 );
4518 let mut registration = router.register_session_run(&target, "run-a").await.unwrap();
4519 assert_eq!(
4520 retain(router.clone(), birth, "replacement")
4521 .await
4522 .err()
4523 .unwrap()
4524 .reason,
4525 SessionRunConflictReason::RunChanged
4526 );
4527 let stale_birth = retain(
4528 router.clone(),
4529 birth + chrono::Duration::seconds(1),
4530 "run-a",
4531 )
4532 .await
4533 .err()
4534 .unwrap();
4535 assert_eq!(
4536 stale_birth.reason,
4537 SessionRunConflictReason::SessionBirthChanged
4538 );
4539 assert_eq!(stale_birth.observed_run_id, None);
4540 let guard = retain(router.clone(), birth, "run-a").await.unwrap();
4541 assert!(guard.matches("target", birth, "run-a"));
4542 assert!(!guard.matches("other-target", birth, "run-a"));
4543 let notify = guard.clone();
4544 assert!(
4545 tokio::task::spawn_blocking(move || notify.notify_committed_delivery(7))
4546 .await
4547 .unwrap()
4548 );
4549 let mut finalizing = Box::pin(router.begin_finalization("target", "run-a"));
4550 assert!(futures::poll!(finalizing.as_mut()).is_pending());
4551 drop(guard);
4552 finalizing.await;
4553 assert_eq!(
4554 retain(router.clone(), birth, "run-a")
4555 .await
4556 .err()
4557 .unwrap()
4558 .reason,
4559 SessionRunConflictReason::Finalizing
4560 );
4561 assert_eq!(*registration.notifications_mut().borrow(), 7);
4562 assert_eq!(registration.finish(0).await.unwrap(), None);
4563 assert_eq!(router.current_run_id("target").await, None);
4564 }
4565
4566 #[tokio::test]
4567 async fn expected_run_publication_fence_does_not_block_other_sessions() {
4568 let router = SessionActivationRouter::new();
4569 let target = bamboo_domain::Session::new("target", "model");
4570 let registration = router.register_session_run(&target, "run-a").await.unwrap();
4571 let query = router.clone();
4572 let guard = tokio::task::spawn_blocking(move || {
4573 query.retain_expected_run_blocking(&target.id, target.created_at, "run-a")
4574 })
4575 .await
4576 .unwrap()
4577 .unwrap();
4578 tokio::time::timeout(std::time::Duration::from_secs(1), async {
4581 let other = bamboo_domain::Session::new("independent", "model");
4582 let mut other_run = router
4583 .register_session_run(&other, "independent-run")
4584 .await
4585 .unwrap();
4586 assert!(router.notify_active_delivery(&other.id, 11).await);
4587 assert_eq!(*other_run.notifications_mut().borrow(), 11);
4588 other_run.begin_finalization().await;
4589 assert_eq!(other_run.finish(0).await.unwrap(), None);
4590 assert_eq!(router.current_run_id(&other.id).await, None);
4591 })
4592 .await
4593 .expect("a parked publication must not hold the global router lock");
4594 drop(guard);
4595 assert_eq!(registration.finish(0).await.unwrap(), None);
4596 }
4597
4598 #[tokio::test]
4599 async fn expected_run_requires_real_lifetime_binding_and_cannot_rebind() {
4600 use bamboo_domain::SessionRunConflictReason;
4601 let router = SessionActivationRouter::new();
4602 let target = bamboo_domain::Session::new("target", "model");
4603 let mut registration = router.register_run(&target.id, "run-a").await.unwrap();
4604 let query = router.clone();
4605 let birth = target.created_at;
4606 let conflict = tokio::task::spawn_blocking(move || {
4607 query.retain_expected_run_blocking("target", birth, "run-a")
4608 })
4609 .await
4610 .unwrap()
4611 .err()
4612 .unwrap();
4613 assert_eq!(conflict.reason, SessionRunConflictReason::UnboundLifetime);
4614 assert_eq!(conflict.observed_run_id, None);
4615 let other_router = SessionActivationRouter::new();
4616 assert!(
4617 !other_router
4618 .bind_run_lifetime(®istration.ownership(), &target)
4619 .await
4620 );
4621 assert!(
4622 router
4623 .bind_run_lifetime(®istration.ownership(), &target)
4624 .await
4625 );
4626 let mut replacement = target.clone();
4627 replacement.created_at += chrono::Duration::seconds(1);
4628 assert!(
4629 !router
4630 .bind_run_lifetime(®istration.ownership(), &replacement)
4631 .await
4632 );
4633 let query = router.clone();
4634 let guard = tokio::task::spawn_blocking(move || {
4635 query.retain_expected_run_blocking("target", birth, "run-a")
4636 })
4637 .await
4638 .unwrap()
4639 .unwrap();
4640 assert!(guard.matches("target", birth, "run-a"));
4641 drop(guard);
4642 registration.begin_finalization().await;
4643 assert!(
4644 !router
4645 .bind_run_lifetime(®istration.ownership(), &target)
4646 .await
4647 );
4648 assert_eq!(registration.finish(0).await.unwrap(), None);
4649 }
4650
4651 #[tokio::test]
4652 async fn expected_run_guard_rejects_unregistered_reservations_without_notification() {
4653 let router = SessionActivationRouter::new();
4654 router
4655 .install_owner_placeholder_for_test("target", "run-a")
4656 .await;
4657 let query = router.clone();
4658 let result = tokio::task::spawn_blocking(move || {
4659 query.retain_expected_run_blocking("target", Utc::now(), "run-a")
4660 })
4661 .await
4662 .unwrap();
4663 assert_eq!(
4664 result.err().unwrap().reason,
4665 bamboo_domain::SessionRunConflictReason::NotRegistered
4666 );
4667 assert_eq!(
4668 router.current_run_id("target").await.as_deref(),
4669 Some("run-a")
4670 );
4671 assert_eq!(
4672 router.states.lock().await["target"].latest_delivery_generation,
4673 0
4674 );
4675 }
4676
4677 mod tool_suspension_registration_tests {
4678 use super::*;
4679 use bamboo_agent_core::tools::{
4680 capture_suspension_restriction, observe_tool_suspension, with_suspension_disabled,
4681 AsyncWaitKind, RunningCompletion, RunningHandle, ToolOutcome, ToolResult,
4682 ToolSuspensionObservation,
4683 };
4684
4685 async fn bind_running_probe(
4686 registration: &mut SessionRunRegistration,
4687 deny: bool,
4688 kills: Arc<AtomicUsize>,
4689 ) -> ToolSuspensionObservation {
4690 let (_, report) = observe_tool_suspension(false, async {
4691 registration
4692 .bind_tool_suspension_observation(capture_suspension_restriction().unwrap());
4693 let collapse = async move {
4694 ToolOutcome::Running(RunningHandle {
4695 tool_call_id: "registration-running-probe".into(),
4696 ack: ToolResult::text(true, "ordinary running ack"),
4697 completion: RunningCompletion::Detached,
4698 wait_kind: AsyncWaitKind::AsyncTools,
4699 kill: Box::new(move || {
4700 kills.fetch_add(1, Ordering::SeqCst);
4701 }),
4702 })
4703 .into_tool_result()
4704 };
4705 let result = if deny {
4706 with_suspension_disabled(collapse).await
4707 } else {
4708 collapse.await
4709 };
4710 assert_eq!(result.success, !deny);
4711 })
4712 .await;
4713 report
4714 }
4715
4716 #[tokio::test]
4717 async fn forbidden_observation_retains_exact_owner_after_drop_or_observer_cancel() {
4718 for cancel_observer in [false, true] {
4719 let router = SessionActivationRouter::new();
4720 let spawner = spawner();
4721 router.set_spawner(spawner.clone()).await;
4722 let mut registration = router
4723 .register_run("session", "observed-running-owner")
4724 .await
4725 .unwrap();
4726 let cleaned = Arc::new(AtomicUsize::new(0));
4727 let clean_count = cleaned.clone();
4728 registration.set_abort_cleanup(move || async move {
4729 clean_count.fetch_add(1, Ordering::SeqCst);
4730 });
4731 assert_eq!(
4732 router.request_activation("session", 7).await.unwrap(),
4733 SessionActivationDisposition::ActiveNotified
4734 );
4735 let kills = Arc::new(AtomicUsize::new(0));
4736 let report = if cancel_observer {
4737 let (ready_tx, ready_rx) = tokio::sync::oneshot::channel();
4738 let child_kills = kills.clone();
4739 let task = tokio::spawn(async move {
4740 observe_tool_suspension(false, async move {
4741 let report =
4742 bind_running_probe(&mut registration, true, child_kills).await;
4743 assert!(registration.observed_forbidden_running());
4744 let _ = ready_tx.send(report);
4745 std::future::pending::<()>().await;
4746 drop(registration);
4749 })
4750 .await;
4751 });
4752 let ready =
4753 tokio::time::timeout(std::time::Duration::from_secs(2), ready_rx).await;
4754 task.abort();
4755 let task_exit = task.await;
4756 assert!(task_exit.unwrap_err().is_cancelled());
4757 ready.unwrap().unwrap()
4758 } else {
4759 let report = bind_running_probe(&mut registration, true, kills.clone()).await;
4760 assert!(registration.observed_forbidden_running());
4761 drop(registration);
4764 report
4765 };
4766 assert!(report.running_seen && report.forbidden_running_seen);
4767 assert_eq!(kills.load(Ordering::SeqCst), 1);
4768
4769 tokio::time::timeout(std::time::Duration::from_secs(2), async {
4770 loop {
4771 let finalized = router
4772 .states
4773 .lock()
4774 .await
4775 .get("session")
4776 .and_then(|state| state.owner.as_ref())
4777 .is_some_and(|owner| {
4778 owner.run_id == "observed-running-owner" && owner.finalizing
4779 });
4780 if finalized {
4781 break;
4782 }
4783 tokio::task::yield_now().await;
4784 }
4785 })
4786 .await
4787 .expect("the exact retained owner must enter finalization");
4788 assert_eq!(
4789 router.request_activation("session", 8).await.unwrap(),
4790 SessionActivationDisposition::ActivationCoalesced
4791 );
4792 for attempted in ["successor", "observed-running-owner"] {
4793 let collision = router
4794 .register_run("session", attempted)
4795 .await
4796 .err()
4797 .expect("a quarantined owner cannot be replaced or adopted");
4798 assert_eq!(collision.existing_run_id(), "observed-running-owner");
4799 }
4800 assert_eq!(cleaned.load(Ordering::SeqCst), 0);
4801 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 0);
4802 assert_eq!(spawner.launches.load(Ordering::SeqCst), 0);
4803 let states = router.states.lock().await;
4804 let owner = states["session"].owner.as_ref().unwrap();
4805 assert_eq!(owner.run_id, "observed-running-owner");
4806 assert!(owner.finalizing);
4807 assert_eq!(owner.registrations, 1);
4808 }
4809 }
4810
4811 #[tokio::test]
4812 async fn allowed_observation_preserves_ordinary_drop_and_finish_successors() {
4813 for finish_normally in [false, true] {
4814 let router = SessionActivationRouter::new();
4815 let spawner = spawner();
4816 router.set_spawner(spawner.clone()).await;
4817 let mut registration = router
4818 .register_run("session", "ordinary-observed-owner")
4819 .await
4820 .unwrap();
4821 let cleaned = Arc::new(AtomicUsize::new(0));
4822 let clean_count = cleaned.clone();
4823 registration.set_abort_cleanup(move || async move {
4824 clean_count.fetch_add(1, Ordering::SeqCst);
4825 });
4826 let kills = Arc::new(AtomicUsize::new(0));
4827 let report = bind_running_probe(&mut registration, false, kills.clone()).await;
4828 assert!(report.running_seen);
4829 assert!(!report.forbidden_running_seen);
4830 assert!(!registration.observed_forbidden_running());
4831 assert_eq!(kills.load(Ordering::SeqCst), 0);
4832 assert_eq!(
4833 router.request_activation("session", 7).await.unwrap(),
4834 SessionActivationDisposition::ActiveNotified
4835 );
4836 if finish_normally {
4837 registration.begin_finalization().await;
4838 assert_eq!(
4839 registration.finish(0).await.unwrap(),
4840 Some(SessionActivationDisposition::ActivationReserved)
4841 );
4842 } else {
4843 drop(registration);
4844 }
4845 tokio::time::timeout(std::time::Duration::from_secs(2), async {
4846 while spawner.launches.load(Ordering::SeqCst) == 0 {
4847 tokio::task::yield_now().await;
4848 }
4849 })
4850 .await
4851 .expect("ordinary handoff must still launch its one successor");
4852 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 1);
4853 assert_eq!(spawner.launches.load(Ordering::SeqCst), 1);
4854 assert_eq!(
4855 cleaned.load(Ordering::SeqCst),
4856 usize::from(!finish_normally)
4857 );
4858 assert!(router.owns_run("session", "run-1").await);
4859 let mut successor = router.register_run("session", "run-1").await.unwrap();
4860 successor.begin_finalization().await;
4861 assert_eq!(successor.finish(7).await.unwrap(), None);
4862 assert!(router.current_run_id("session").await.is_none());
4863 assert_eq!(spawner.launches.load(Ordering::SeqCst), 1);
4864 }
4865 }
4866 }
4867}
4868
4869#[cfg(test)]
4870#[path = "session_activation_child_completion_tests.rs"]
4871mod child_completion_fence_tests;
4872
4873#[cfg(test)]
4874#[path = "session_activation/scoped_loop_tests.rs"]
4875mod scoped_loop_tests;