1use std::collections::{BTreeSet, HashMap};
7use std::future::Future;
8use std::pin::Pin;
9use std::sync::Arc;
10
11use tokio::sync::{broadcast, mpsc, RwLock};
12use tokio_util::sync::CancellationToken;
13use tracing::Instrument;
14
15use bamboo_agent_core::tools::ToolExecutor;
16use bamboo_agent_core::{AgentError, AgentEvent, Session};
17use bamboo_domain::ReasoningEffort;
18use bamboo_llm::LLMProvider;
19
20use crate::runtime::config::{
21 AuxiliaryModelConfig, BashCompletionSink, BashResumeHook, DisabledFilterResolver, GoldConfig,
22 GuardianConfig, GuardianSpawner, ImageFallbackConfig,
23};
24use crate::runtime::execution::child_completion::ChildCompletion;
25use crate::runtime::execution::event_forwarder::HistoryCommitBarrier;
26use crate::runtime::execution::runner_lifecycle::{
27 finalize_rejected_runner_if_distinct, finalize_runner_exact, reserve_runner_core,
28 reserve_runner_for_registration, ReserveOutcome, RunnerReservation,
29};
30use crate::runtime::execution::runner_state::AgentRunner;
31use crate::runtime::model_roster::ModelRoster;
32use crate::runtime::Agent;
33use crate::runtime::{ExecuteRequest, ExecuteRequestBuilder};
34use crate::session_activation::{
35 SessionActivationRouter, SessionRunRegistration, SessionRunRegistrationError,
36};
37
38pub use crate::session_cache::SessionCache;
39
40enum SessionExecutionActivationOwnership {
41 Unrouted,
43 UnpublishedActivation(Arc<SessionActivationRouter>),
46 RegistrationPending(Arc<SessionActivationRouter>),
50 Registered(Box<SessionRunRegistration>),
53}
54
55pub struct SessionExecutionReservation {
61 session_id: String,
62 run_id: String,
63 cancel_token: CancellationToken,
64 runners: Arc<RwLock<HashMap<String, AgentRunner>>>,
65 activation: SessionExecutionActivationOwnership,
66 root_actor: Option<Box<RootActorExecution>>,
67 child_startup: Option<Box<ChildHostStartup>>,
68 child_closure: Option<tokio::task::JoinHandle<()>>,
69 physical_admission: Option<super::root_execution_lifetime::RootRunAdmission>,
70 execution_persistence: Option<Arc<dyn bamboo_domain::RuntimeSessionPersistence>>,
71 persistence_bound: bool,
72 selected_input: bool,
73 armed: bool,
74}
75
76impl SessionExecutionReservation {
77 pub fn session_id(&self) -> &str {
78 &self.session_id
79 }
80
81 pub fn run_id(&self) -> &str {
82 &self.run_id
83 }
84
85 pub fn cancel_token(&self) -> &CancellationToken {
86 &self.cancel_token
87 }
88
89 pub async fn arm_child_host_startup(
93 &mut self,
94 agent: &Agent,
95 expected: &Session,
96 cache: SessionCache,
97 events: tokio::sync::broadcast::Sender<AgentEvent>,
98 ) -> std::io::Result<()> {
99 self.arm_child_host_startup_inner(agent, expected, cache, Some(events))
100 .await
101 }
102
103 async fn arm_child_host_startup_inner(
104 &mut self,
105 agent: &Agent,
106 expected: &Session,
107 cache: SessionCache,
108 events: Option<tokio::sync::broadcast::Sender<AgentEvent>>,
109 ) -> std::io::Result<()> {
110 if expected.kind != bamboo_domain::SessionKind::Child || expected.id != self.session_id {
111 return Err(std::io::Error::other("Child startup target mismatch"));
112 }
113 if self.child_startup.is_none() {
114 let publication = self
115 .runners
116 .read()
117 .await
118 .get(&self.session_id)
119 .filter(|runner| runner.run_id == self.run_id)
120 .map(|runner| runner.event_publication.clone())
121 .ok_or_else(|| std::io::Error::other("Child startup runner was superseded"))?;
122 let sent = std::sync::atomic::AtomicBool::new(false);
123 self.child_startup = Some(Box::new(ChildHostStartup {
124 storage: agent.storage().clone(),
125 expected: expected.clone(),
126 run_id: self.run_id.clone(),
127 publish: Arc::new(move |saved| {
128 publication.publish(|| {
129 cache.insert(
130 saved.id.clone(),
131 Arc::new(crate::SessionSnapshot::new(saved.clone())),
132 );
133 if !sent.swap(true, std::sync::atomic::Ordering::AcqRel) {
134 if let Some(events) = &events {
135 let _ = events.send(AgentEvent::Error {
136 message: saved.last_run_error().unwrap_or_else(|| {
137 "Child startup failed before execution".into()
138 }),
139 });
140 }
141 }
142 });
143 }),
144 }));
145 }
146 Ok(())
147 }
148
149 pub async fn fail_child_host_startup(&mut self, detail: &str) {
150 if let Some(startup) = self.child_startup.take() {
151 let detail = detail.to_string();
152 self.child_closure = Some(tokio::spawn(async move {
154 startup.close(&detail).await;
155 }));
156 }
157 if let Some(closure) = self.child_closure.as_mut() {
158 let _ = closure.await;
159 }
160 self.child_closure = None;
161 }
162
163 pub async fn bind_root_actor(
166 &mut self,
167 agent: &Agent,
168 session: &Session,
169 ) -> std::io::Result<()> {
170 if self.root_actor.is_none() {
171 if self.session_id != session.id {
172 return Err(std::io::Error::new(
173 std::io::ErrorKind::InvalidInput,
174 "Root binding target mismatch",
175 ));
176 }
177 if let Some(binding) = agent
178 .persistence()
179 .bind_root_actor_execution(session, &self.run_id)
180 .await?
181 {
182 self.root_actor = Some(Box::new(RootActorExecution::start(binding)));
183 }
184 }
185 self.bind_execution_persistence(agent, session)
186 }
187
188 fn bind_execution_persistence(
189 &mut self,
190 agent: &Agent,
191 session: &Session,
192 ) -> std::io::Result<()> {
193 if self.session_id != session.id {
194 return Err(std::io::Error::other(
195 "execution persistence target mismatch",
196 ));
197 }
198 if self.persistence_bound {
199 return Ok(());
200 }
201 let persistence = self
202 .root_actor
203 .as_ref()
204 .map_or_else(|| agent.persistence().clone(), |owner| owner.persistence());
205 if let Some(inherited) = persistence
206 .inherited_child_wait()
207 .or_else(|| bamboo_domain::InheritedChildWait::capture(session))
208 {
209 inherited.validate_session(session)?;
210 self.execution_persistence = Some(persistence.bind_inherited_child_wait(inherited)?);
211 } else if self.root_actor.is_some() {
212 self.execution_persistence = Some(persistence);
213 }
214 self.persistence_bound = true;
215 Ok(())
216 }
217
218 pub fn execution_persistence(
219 &self,
220 ) -> Option<Arc<dyn bamboo_domain::RuntimeSessionPersistence>> {
221 self.execution_persistence.clone()
222 }
223
224 pub fn bind_root_inbox_receipt(
227 &mut self,
228 receipt: bamboo_domain::SessionInboxReceipt,
229 ) -> std::io::Result<()> {
230 if !self.selected_input || self.root_actor.is_none() || self.cancel_token.is_cancelled() {
231 return Err(std::io::Error::other(
232 "Root receipt requires a live bound reservation",
233 ));
234 }
235 let persistence = self
236 .execution_persistence
237 .as_ref()
238 .ok_or_else(|| std::io::Error::other("Root execution persistence unavailable"))?;
239 self.execution_persistence = Some(persistence.bind_root_inbox_receipt(receipt)?);
240 Ok(())
241 }
242
243 pub async fn retain_selected_run_output(&self) -> std::io::Result<bool> {
246 if !self.selected_input {
247 return Ok(false);
248 }
249 let owner = self
250 .root_actor_writer()
251 .ok_or_else(|| std::io::Error::other("Selected output requires its Root writer"))?;
252 let receipt = self
253 .execution_persistence
254 .as_ref()
255 .and_then(|p| p.root_inbox_receipt())
256 .cloned()
257 .ok_or_else(|| std::io::Error::other("Selected output requires its bound receipt"))?;
258 if owner.fence.actor_id != self.session_id
259 || owner.fence.run_id != self.run_id
260 || self.cancel_token.is_cancelled()
261 {
262 return Err(std::io::Error::other(
263 "Selected output reservation is not live",
264 ));
265 }
266 let runners = self.runners.read().await;
267 let runner = runners
268 .get(&self.session_id)
269 .filter(|r| r.run_id == self.run_id && !r.cancel_token.is_cancelled())
270 .ok_or_else(|| std::io::Error::other("Selected output runner is unavailable"))?;
271 runner
272 .visible_messages
273 .retain_selected_output(&owner, receipt)?;
274 Ok(true)
275 }
276
277 pub fn root_actor_writer(&self) -> Option<bamboo_domain::RootActorRuntimeWrite> {
278 self.root_actor
279 .as_ref()
280 .map(|owner| owner.binding.owner.clone())
281 }
282
283 pub fn queue_root_response_event(&mut self, event: AgentEvent) -> Option<AgentEvent> {
286 let Some(owner) = self.root_actor.as_mut() else {
287 return Some(event);
288 };
289 owner.response_events.push(event);
290 None
291 }
292
293 pub fn take_root_response_events(&mut self) -> Vec<AgentEvent> {
294 self.root_actor
295 .as_mut()
296 .map(|owner| std::mem::take(&mut owner.response_events))
297 .unwrap_or_default()
298 }
299
300 pub async fn finish_root_actor(&mut self, outcome: bamboo_domain::ActorActivationFinish) {
303 if let Some(owner) = self.root_actor.take() {
304 (*owner).finish(outcome).await;
305 }
306 }
307
308 pub(crate) fn from_activation_placeholder(
314 session_id: impl Into<String>,
315 reservation: RunnerReservation,
316 router: Arc<SessionActivationRouter>,
317 runners: Arc<RwLock<HashMap<String, AgentRunner>>>,
318 ) -> Self {
319 Self {
320 session_id: session_id.into(),
321 run_id: reservation.run_id,
322 cancel_token: reservation.cancel_token,
323 runners,
324 activation: SessionExecutionActivationOwnership::UnpublishedActivation(router),
325 root_actor: None,
326 physical_admission: reservation.physical_admission,
327 execution_persistence: None,
328 persistence_bound: false,
329 child_startup: None,
330 child_closure: None,
331 selected_input: false,
332 armed: true,
333 }
334 }
335
336 pub(crate) fn from_pending_registration(
340 session_id: impl Into<String>,
341 reservation: RunnerReservation,
342 router: Option<Arc<SessionActivationRouter>>,
343 runners: Arc<RwLock<HashMap<String, AgentRunner>>>,
344 ) -> Self {
345 let activation = match router {
346 Some(router) => SessionExecutionActivationOwnership::RegistrationPending(router),
347 None => SessionExecutionActivationOwnership::Unrouted,
348 };
349 Self {
350 session_id: session_id.into(),
351 run_id: reservation.run_id,
352 cancel_token: reservation.cancel_token,
353 runners,
354 activation,
355 root_actor: None,
356 physical_admission: reservation.physical_admission,
357 execution_persistence: None,
358 persistence_bound: false,
359 child_startup: None,
360 child_closure: None,
361 selected_input: false,
362 armed: true,
363 }
364 }
365
366 pub(crate) fn mark_activation_published(&mut self) {
369 let activation = std::mem::replace(
370 &mut self.activation,
371 SessionExecutionActivationOwnership::Unrouted,
372 );
373 self.activation = match activation {
374 SessionExecutionActivationOwnership::UnpublishedActivation(router) => {
375 SessionExecutionActivationOwnership::RegistrationPending(router)
376 }
377 other => other,
378 };
379 }
380
381 pub(crate) async fn rollback_unpublished_activation(mut self) {
386 self.armed = false;
387 self.cancel_token.cancel();
388 self.finish_root_actor(bamboo_domain::ActorActivationFinish::Cancelled)
389 .await;
390 let activation = std::mem::replace(
391 &mut self.activation,
392 SessionExecutionActivationOwnership::Unrouted,
393 );
394 debug_assert!(matches!(
395 activation,
396 SessionExecutionActivationOwnership::UnpublishedActivation(_)
397 ));
398 remove_runner_exact(&self.runners, &self.session_id, &self.run_id).await;
399 }
400
401 pub async fn ensure_registered(&mut self) -> Result<(), SessionRunRegistrationError> {
407 let router = match &self.activation {
408 SessionExecutionActivationOwnership::RegistrationPending(router) => router.clone(),
409 SessionExecutionActivationOwnership::UnpublishedActivation(_) => {
410 unreachable!("unpublished activation entered an execution adapter")
411 }
412 SessionExecutionActivationOwnership::Unrouted
413 | SessionExecutionActivationOwnership::Registered(_) => return Ok(()),
414 };
415 match register_reserved_activation(
416 router,
417 &self.runners,
418 &self.session_id,
419 &self.run_id,
420 self.selected_input,
421 )
422 .await
423 {
424 Ok(registration) => {
425 self.activation =
426 SessionExecutionActivationOwnership::Registered(Box::new(registration));
427 Ok(())
428 }
429 Err(error) => {
430 self.disarm_after_registration_rejection(error.existing_run_id());
435 Err(error)
436 }
437 }
438 }
439
440 pub async fn abandon(mut self) {
444 self.armed = false;
445 self.cancel_token.cancel();
446 if let Some(owner) = self.root_actor.take() {
447 (*owner)
448 .finish(bamboo_domain::ActorActivationFinish::Cancelled)
449 .await;
450 }
451 let activation = std::mem::replace(
452 &mut self.activation,
453 SessionExecutionActivationOwnership::Unrouted,
454 );
455 let runners = self.runners.clone();
456 let session_id = self.session_id.clone();
457 let run_id = self.run_id.clone();
458 let cleanup_session_id = session_id.clone();
459 let cleanup_run_id = run_id.clone();
460 let selected_input = self.selected_input;
461 let child_startup = self.child_startup.take();
462 let child_closure = self.child_closure.take();
463 let cleanup = tokio::spawn(async move {
464 if let Some(closure) = child_closure {
465 let _ = closure.await;
466 }
467 if let Some(startup) = child_startup {
468 startup
469 .close("startup reservation abandoned before execution")
470 .await;
471 }
472 cleanup_execution_reservation(
473 activation,
474 runners,
475 cleanup_session_id,
476 cleanup_run_id,
477 selected_input,
478 )
479 .await;
480 });
481 if let Err(error) = cleanup.await {
482 tracing::error!(
483 %session_id,
484 %run_id,
485 %error,
486 "detached execution-reservation cleanup failed"
487 );
488 }
489 }
490
491 fn disarm_after_registration_rejection(&mut self, existing_run_id: &str) {
492 if existing_run_id != self.run_id {
493 self.cancel_token.cancel();
494 }
495 self.armed = false;
496 if let Some(startup) = self.child_startup.take() {
497 if existing_run_id != self.run_id {
498 tokio::spawn(async move {
499 startup
500 .close("startup registration rejected before execution")
501 .await;
502 });
503 }
504 }
505 }
506
507 pub(crate) fn matches_execution_target(
508 &self,
509 session_id: &str,
510 domain_session_id: &str,
511 runners: &Arc<RwLock<HashMap<String, AgentRunner>>>,
512 ) -> bool {
513 self.session_id == session_id
514 && domain_session_id == session_id
515 && Arc::ptr_eq(&self.runners, runners)
516 }
517
518 pub(crate) fn disarm_for_execution(
519 &mut self,
520 ) -> (CancellationToken, Option<SessionRunRegistration>) {
521 self.armed = false;
522 self.child_startup = None;
523 let registration = match std::mem::replace(
524 &mut self.activation,
525 SessionExecutionActivationOwnership::Unrouted,
526 ) {
527 SessionExecutionActivationOwnership::Registered(registration) => Some(*registration),
528 SessionExecutionActivationOwnership::Unrouted => None,
529 SessionExecutionActivationOwnership::UnpublishedActivation(_) => {
530 unreachable!("unpublished activation cannot transfer to execution")
531 }
532 SessionExecutionActivationOwnership::RegistrationPending(_) => {
533 unreachable!("ensure_registered adopts every pending router registration")
534 }
535 };
536 (self.cancel_token.clone(), registration)
537 }
538}
539
540impl Drop for SessionExecutionReservation {
541 fn drop(&mut self) {
542 if !self.armed {
543 return;
544 }
545 self.armed = false;
546 self.cancel_token.cancel();
547 let activation = std::mem::replace(
548 &mut self.activation,
549 SessionExecutionActivationOwnership::Unrouted,
550 );
551 let runners = self.runners.clone();
552 let session_id = self.session_id.clone();
553 let run_id = self.run_id.clone();
554 let selected_input = self.selected_input;
555 let child_startup = self.child_startup.take();
556 let child_closure = self.child_closure.take();
557 if let Ok(runtime) = tokio::runtime::Handle::try_current() {
558 runtime.spawn(async move {
559 if let Some(closure) = child_closure {
560 let _ = closure.await;
561 }
562 if let Some(startup) = child_startup {
563 startup
564 .close("startup reservation dropped before execution")
565 .await;
566 }
567 cleanup_execution_reservation(
568 activation,
569 runners,
570 session_id,
571 run_id,
572 selected_input,
573 )
574 .await;
575 });
576 }
577 }
578}
579
580struct ChildHostStartup {
581 storage: Arc<dyn bamboo_domain::Storage>,
582 expected: Session,
583 run_id: String,
584 publish: bamboo_domain::storage::RootActorRuntimePublisher,
585}
586
587impl ChildHostStartup {
588 async fn close(self, detail: &str) {
589 if let Err(error) = self
590 .storage
591 .fail_child_host_startup(&self.expected, &self.run_id, detail, self.publish)
592 .await
593 {
594 tracing::warn!(session_id = %self.expected.id, run_id = %self.run_id, %error, "Child pre-execution closure unconfirmed");
595 }
596 }
597}
598
599pub struct SessionPreparationReservation {
602 reservation: Option<SessionExecutionReservation>,
603}
604
605impl SessionPreparationReservation {
606 pub async fn bind_root_actor(
607 &mut self,
608 agent: &Agent,
609 session: &Session,
610 ) -> std::io::Result<()> {
611 self.reservation
612 .as_mut()
613 .expect("live preparation")
614 .bind_root_actor(agent, session)
615 .await
616 }
617
618 pub fn root_actor_writer(&self) -> Option<bamboo_domain::RootActorRuntimeWrite> {
619 self.reservation
620 .as_ref()
621 .and_then(SessionExecutionReservation::root_actor_writer)
622 }
623
624 pub async fn finish(mut self, outcome: bamboo_domain::ActorActivationFinish) {
625 let reservation = self.reservation.take().expect("live preparation");
626 let _ = Self::release(reservation, outcome).await;
627 }
628
629 pub async fn abandon(self) {
630 self.finish(bamboo_domain::ActorActivationFinish::Cancelled)
631 .await;
632 }
633
634 fn release(
635 mut reservation: SessionExecutionReservation,
636 outcome: bamboo_domain::ActorActivationFinish,
637 ) -> tokio::task::JoinHandle<()> {
638 tokio::spawn(async move {
641 reservation.finish_root_actor(outcome).await;
642 reservation.abandon().await;
643 })
644 }
645}
646
647impl Drop for SessionPreparationReservation {
648 fn drop(&mut self) {
649 if let Some(reservation) = self.reservation.take() {
650 if tokio::runtime::Handle::try_current().is_ok() {
651 drop(Self::release(
652 reservation,
653 bamboo_domain::ActorActivationFinish::Cancelled,
654 ));
655 }
656 }
657 }
658}
659
660pub async fn reserve_session_preparation(
663 agent: &Arc<Agent>,
664 runners: &Arc<RwLock<HashMap<String, AgentRunner>>>,
665 session_id: &str,
666) -> std::io::Result<SessionPreparationReservation> {
667 let router = agent.activation_router().cloned().ok_or_else(|| {
668 std::io::Error::new(
669 std::io::ErrorKind::Unsupported,
670 "Preparation requires its activation router",
671 )
672 })?;
673 let busy = || {
674 std::io::Error::new(
675 std::io::ErrorKind::WouldBlock,
676 "Session has an active runner or logical owner",
677 )
678 };
679 if runners.read().await.get(session_id).is_some_and(|runner| {
680 runner.root_reset_admission_reserved()
681 || matches!(
682 runner.status,
683 super::runner_state::AgentStatus::Pending
684 | super::runner_state::AgentStatus::Running
685 )
686 }) {
687 return Err(busy());
688 }
689 let reservation = reserve_logical_session(router, runners, session_id, true)
692 .await
693 .map_err(|_| busy())?;
694 let preparation = SessionPreparationReservation {
695 reservation: Some(reservation),
696 };
697 if runners.read().await.get(session_id).is_some_and(|runner| {
698 runner.root_reset_admission_reserved()
699 || matches!(
700 runner.status,
701 super::runner_state::AgentStatus::Pending
702 | super::runner_state::AgentStatus::Running
703 )
704 }) {
705 preparation.abandon().await;
706 return Err(busy());
707 }
708 Ok(preparation)
709}
710
711async fn reserve_logical_session(
712 router: Arc<SessionActivationRouter>,
713 runners: &Arc<RwLock<HashMap<String, AgentRunner>>>,
714 session_id: &str,
715 selected_input: bool,
716) -> Result<SessionExecutionReservation, SessionRunRegistrationError> {
717 let mut reservation = SessionExecutionReservation {
718 session_id: session_id.to_owned(),
719 run_id: uuid::Uuid::new_v4().to_string(),
720 cancel_token: CancellationToken::new(),
721 runners: runners.clone(),
722 activation: SessionExecutionActivationOwnership::Unrouted,
723 root_actor: None,
724 physical_admission: None,
725 execution_persistence: None,
726 persistence_bound: false,
727 child_startup: None,
728 child_closure: None,
729 selected_input: false,
730 armed: true,
731 };
732 let registration = register_reserved_activation(
733 router,
734 runners,
735 session_id,
736 &reservation.run_id,
737 selected_input,
738 )
739 .await?;
740 reservation.activation =
741 SessionExecutionActivationOwnership::Registered(Box::new(registration));
742 Ok(reservation)
743}
744
745pub enum SessionExecutionReserveOutcome {
747 Reserved(SessionExecutionReservation),
748 AlreadyRunning { run_id: String },
749}
750
751async fn register_reserved_activation(
752 router: Arc<SessionActivationRouter>,
753 runners: &Arc<RwLock<HashMap<String, AgentRunner>>>,
754 session_id: &str,
755 run_id: &str,
756 selected_input: bool,
757) -> Result<SessionRunRegistration, SessionRunRegistrationError> {
758 let mut registration = match router.register_run(session_id, run_id).await {
759 Ok(registration) => registration,
760 Err(error) => {
761 let collision = Err(AgentError::LLM(error.to_string()));
762 finalize_rejected_runner_if_distinct(
763 runners,
764 session_id,
765 error.existing_run_id(),
766 run_id,
767 &collision,
768 )
769 .await;
770 return Err(error);
771 }
772 };
773 if selected_input {
774 registration.restrict_to_selected_input();
775 }
776 let cleanup_runners = runners.clone();
777 let cleanup_session_id = session_id.to_string();
778 let cleanup_run_id = run_id.to_string();
779 registration.set_abort_cleanup(move || async move {
780 let abandoned = Err(AgentError::Cancelled);
781 finalize_runner_exact(
782 &cleanup_runners,
783 &cleanup_session_id,
784 &cleanup_run_id,
785 &abandoned,
786 )
787 .await;
788 });
789 Ok(registration)
790}
791
792async fn cleanup_execution_reservation(
793 activation: SessionExecutionActivationOwnership,
794 runners: Arc<RwLock<HashMap<String, AgentRunner>>>,
795 session_id: String,
796 run_id: String,
797 selected_input: bool,
798) {
799 match activation {
800 SessionExecutionActivationOwnership::Registered(registration) => {
801 (*registration).abandon().await;
802 }
803 SessionExecutionActivationOwnership::RegistrationPending(router) => {
804 match register_reserved_activation(
805 router,
806 &runners,
807 &session_id,
808 &run_id,
809 selected_input,
810 )
811 .await
812 {
813 Ok(registration) => registration.abandon().await,
814 Err(error) => {
815 tracing::debug!(
816 %session_id,
817 %run_id,
818 existing_run_id = %error.existing_run_id(),
819 "abandoned activation placeholder was already adopted or superseded"
820 );
821 }
822 }
823 }
824 SessionExecutionActivationOwnership::UnpublishedActivation(_) => {
825 remove_runner_exact(&runners, &session_id, &run_id).await;
826 }
827 SessionExecutionActivationOwnership::Unrouted => {
828 let abandoned = Err(AgentError::Cancelled);
829 finalize_runner_exact(&runners, &session_id, &run_id, &abandoned).await;
830 }
831 }
832}
833
834async fn remove_runner_exact(
835 runners: &Arc<RwLock<HashMap<String, AgentRunner>>>,
836 session_id: &str,
837 run_id: &str,
838) -> bool {
839 let mut runners = runners.write().await;
840 if runners
841 .get(session_id)
842 .is_some_and(|runner| runner.run_id == run_id)
843 {
844 super::runner_lifecycle::remove_runner_entry(&mut runners, session_id).await;
845 true
846 } else {
847 false
848 }
849}
850
851pub async fn reserve_session_execution(
859 agent: &Arc<Agent>,
860 runners: &Arc<RwLock<HashMap<String, AgentRunner>>>,
861 senders: &Arc<RwLock<HashMap<String, broadcast::Sender<AgentEvent>>>>,
862 session_id: &str,
863 event_sender: &broadcast::Sender<AgentEvent>,
864) -> SessionExecutionReserveOutcome {
865 reserve_session_execution_inner(agent, runners, senders, session_id, event_sender, false).await
866}
867
868pub async fn reserve_selected_input_execution(
871 agent: &Arc<Agent>,
872 runners: &Arc<RwLock<HashMap<String, AgentRunner>>>,
873 senders: &Arc<RwLock<HashMap<String, broadcast::Sender<AgentEvent>>>>,
874 session_id: &str,
875 event_sender: &broadcast::Sender<AgentEvent>,
876) -> SessionExecutionReserveOutcome {
877 reserve_session_execution_inner(agent, runners, senders, session_id, event_sender, true).await
878}
879
880async fn reserve_session_execution_inner(
881 agent: &Arc<Agent>,
882 runners: &Arc<RwLock<HashMap<String, AgentRunner>>>,
883 senders: &Arc<RwLock<HashMap<String, broadcast::Sender<AgentEvent>>>>,
884 session_id: &str,
885 event_sender: &broadcast::Sender<AgentEvent>,
886 selected_input: bool,
887) -> SessionExecutionReserveOutcome {
888 if let Some(router) = agent.activation_router().cloned() {
889 if let Some(run_id) = runners
890 .read()
891 .await
892 .get(session_id)
893 .filter(|runner| {
894 runner.root_reset_admission_reserved()
895 || matches!(runner.status, super::runner_state::AgentStatus::Running)
896 })
897 .map(|runner| runner.run_id.clone())
898 {
899 return SessionExecutionReserveOutcome::AlreadyRunning { run_id };
900 }
901 let mut reservation =
904 match reserve_logical_session(router, runners, session_id, selected_input).await {
905 Ok(reservation) => reservation,
906 Err(error) => {
907 return SessionExecutionReserveOutcome::AlreadyRunning {
908 run_id: error.existing_run_id().to_owned(),
909 }
910 }
911 };
912 let identity = RunnerReservation {
913 run_id: reservation.run_id.clone(),
914 cancel_token: reservation.cancel_token.clone(),
915 physical_admission: None,
916 };
917 match reserve_runner_for_registration(runners, senders, session_id, event_sender, identity)
918 .await
919 {
920 ReserveOutcome::Reserved(raw) => {
921 reservation.physical_admission = raw.physical_admission;
922 reservation.selected_input = selected_input;
923 return SessionExecutionReserveOutcome::Reserved(reservation);
924 }
925 ReserveOutcome::AlreadyRunning(run_id) => {
926 reservation.abandon().await;
927 return SessionExecutionReserveOutcome::AlreadyRunning { run_id };
928 }
929 }
930 }
931
932 let reservation = match reserve_runner_core(runners, senders, session_id, event_sender).await {
933 ReserveOutcome::AlreadyRunning(run_id) => {
934 return SessionExecutionReserveOutcome::AlreadyRunning { run_id };
935 }
936 ReserveOutcome::Reserved(reservation) => reservation,
937 };
938
939 let execution_reservation = SessionExecutionReservation {
940 session_id: session_id.to_string(),
941 run_id: reservation.run_id,
942 cancel_token: reservation.cancel_token,
943 runners: runners.clone(),
944 activation: SessionExecutionActivationOwnership::Unrouted,
945 root_actor: None,
946 physical_admission: reservation.physical_admission,
947 execution_persistence: None,
948 persistence_bound: false,
949 child_startup: None,
950 child_closure: None,
951 selected_input,
952 armed: true,
953 };
954 SessionExecutionReserveOutcome::Reserved(execution_reservation)
955}
956
957pub fn read_cached_session(cache: &SessionCache, id: &str) -> Option<bamboo_agent_core::Session> {
965 cache
966 .get(id)
967 .map(|e| e.value().clone())
968 .map(|a| a.read().clone())
969}
970
971const SKILL_CONTEXT_START_MARKER: &str = "<!-- BAMBOO_SKILL_CONTEXT_START -->";
972const TOOL_GUIDE_START_MARKER: &str = "<!-- BAMBOO_TOOL_GUIDE_START -->";
973const EXTERNAL_MEMORY_START_MARKER: &str = "<!-- BAMBOO_EXTERNAL_MEMORY_START -->";
974const TASK_LIST_START_MARKER: &str = "<!-- BAMBOO_TASK_LIST_START -->";
975
976pub struct SessionExecutionOutcome {
983 pub success: bool,
985 pub cancelled: bool,
987 pub error: Option<String>,
989}
990
991impl SessionExecutionOutcome {
992 fn from_result(result: &Result<(), AgentError>) -> Self {
993 match result {
994 Ok(()) => Self {
995 success: true,
996 cancelled: false,
997 error: None,
998 },
999 Err(error) => Self {
1000 success: false,
1001 cancelled: error.is_cancelled(),
1002 error: Some(error.to_string()),
1003 },
1004 }
1005 }
1006}
1007
1008pub type SessionCompletionHook = Box<
1018 dyn for<'a> FnOnce(
1019 SessionExecutionOutcome,
1020 &'a mut Session,
1021 ) -> Pin<Box<dyn Future<Output = ()> + Send + 'a>>
1022 + Send,
1023>;
1024
1025pub struct SessionExecutionArgs {
1031 pub agent: Arc<Agent>,
1033 pub session_id: String,
1034 pub session: Session,
1035 pub execution_reservation: SessionExecutionReservation,
1038
1039 pub tools_override: Option<Arc<dyn ToolExecutor>>,
1041 pub provider_override: Option<Arc<dyn LLMProvider>>,
1042 pub model_roster: ModelRoster,
1046 pub reasoning_effort: Option<ReasoningEffort>,
1047 pub reasoning_effort_source: String,
1048 pub auxiliary_model_resolver:
1049 Option<Arc<dyn Fn() -> crate::runtime::config::AuxiliaryModelConfig + Send + Sync>>,
1050 pub disabled_filter_resolver: Option<DisabledFilterResolver>,
1054 pub disabled_tools: Option<BTreeSet<String>>,
1055 pub disabled_skill_ids: Option<BTreeSet<String>>,
1056 pub selected_skill_ids: Option<Vec<String>>,
1057 pub selected_skill_mode: Option<String>,
1058 pub mpsc_tx: mpsc::Sender<AgentEvent>,
1059 pub history_commit_barrier: HistoryCommitBarrier,
1062 pub image_fallback: Option<ImageFallbackConfig>,
1063 pub gold_config: Option<GoldConfig>,
1064 pub guardian_config: Option<GuardianConfig>,
1066 pub guardian_spawner: Option<Arc<dyn GuardianSpawner>>,
1069 pub bash_resume_hook: Option<Arc<dyn BashResumeHook>>,
1071 pub bash_completion_sink: Option<Arc<dyn BashCompletionSink>>,
1073 pub app_data_dir: Option<std::path::PathBuf>,
1074 pub run_budget: Option<bamboo_config::RunBudgetConfig>,
1077
1078 pub runners: Arc<RwLock<HashMap<String, AgentRunner>>>,
1080 pub sessions_cache: SessionCache,
1081
1082 pub on_complete: Option<SessionCompletionHook>,
1085
1086 pub child_completion_handler: Option<Arc<dyn super::ChildCompletionHandler>>,
1098}
1099
1100struct ExecuteRequestParams {
1110 tools: Option<Arc<dyn ToolExecutor>>,
1111 provider_override: Option<Arc<dyn LLMProvider>>,
1112 model_roster: ModelRoster,
1113 reasoning_effort: Option<ReasoningEffort>,
1114 auxiliary_model_resolver: Option<Arc<dyn Fn() -> AuxiliaryModelConfig + Send + Sync>>,
1115 disabled_filter_resolver: Option<DisabledFilterResolver>,
1116 disabled_tools: Option<BTreeSet<String>>,
1117 disabled_skill_ids: Option<BTreeSet<String>>,
1118 selected_skill_ids: Option<Vec<String>>,
1119 selected_skill_mode: Option<String>,
1120 image_fallback: Option<ImageFallbackConfig>,
1121 gold_config: Option<GoldConfig>,
1122 guardian_config: Option<GuardianConfig>,
1123 guardian_spawner: Option<Arc<dyn GuardianSpawner>>,
1124 bash_resume_hook: Option<Arc<dyn BashResumeHook>>,
1125 bash_completion_sink: Option<Arc<dyn BashCompletionSink>>,
1126 app_data_dir: Option<std::path::PathBuf>,
1127 run_budget: Option<bamboo_config::RunBudgetConfig>,
1128}
1129
1130fn build_execute_request(
1137 initial_message: String,
1138 event_tx: mpsc::Sender<AgentEvent>,
1139 cancel_token: CancellationToken,
1140 params: ExecuteRequestParams,
1141) -> ExecuteRequest {
1142 let ExecuteRequestParams {
1143 tools,
1144 provider_override,
1145 model_roster,
1146 reasoning_effort,
1147 auxiliary_model_resolver,
1148 disabled_filter_resolver,
1149 disabled_tools,
1150 disabled_skill_ids,
1151 selected_skill_ids,
1152 selected_skill_mode,
1153 image_fallback,
1154 gold_config,
1155 guardian_config,
1156 guardian_spawner,
1157 bash_resume_hook,
1158 bash_completion_sink,
1159 app_data_dir,
1160 run_budget,
1161 } = params;
1162
1163 let mut builder = ExecuteRequestBuilder::new(initial_message, event_tx, cancel_token)
1164 .model_roster(model_roster)
1165 .gold_config(gold_config)
1166 .guardian_config(guardian_config)
1167 .guardian_spawner(guardian_spawner)
1168 .bash_resume_hook(bash_resume_hook)
1169 .bash_completion_sink(bash_completion_sink);
1170
1171 if let Some(run_budget) = run_budget {
1172 builder = builder.run_budget(run_budget);
1173 }
1174
1175 if let Some(tools) = tools {
1176 builder = builder.tools(tools);
1177 }
1178 if let Some(provider_override) = provider_override {
1179 builder = builder.provider_override(provider_override);
1180 }
1181 if let Some(reasoning_effort) = reasoning_effort {
1182 builder = builder.reasoning_effort(reasoning_effort);
1183 }
1184 if let Some(disabled_filter_resolver) = disabled_filter_resolver {
1185 builder = builder.disabled_filter_resolver(disabled_filter_resolver);
1186 }
1187 if let Some(auxiliary_model_resolver) = auxiliary_model_resolver {
1188 builder = builder.auxiliary_model_resolver(auxiliary_model_resolver);
1189 }
1190 if let Some(disabled_tools) = disabled_tools {
1191 builder = builder.disabled_tools(disabled_tools);
1192 }
1193 if let Some(disabled_skill_ids) = disabled_skill_ids {
1194 builder = builder.disabled_skill_ids(disabled_skill_ids);
1195 }
1196 if let Some(selected_skill_ids) = selected_skill_ids {
1197 builder = builder.selected_skill_ids(selected_skill_ids);
1198 }
1199 if let Some(selected_skill_mode) = selected_skill_mode {
1200 builder = builder.selected_skill_mode(selected_skill_mode);
1201 }
1202 if let Some(image_fallback) = image_fallback {
1203 builder = builder.image_fallback(image_fallback);
1204 }
1205 if let Some(app_data_dir) = app_data_dir {
1206 builder = builder.app_data_dir(app_data_dir);
1207 }
1208
1209 builder.build()
1210}
1211
1212pub fn spawn_session_execution(args: SessionExecutionArgs) {
1221 spawn_session_execution_with_inputs(args, None);
1222}
1223
1224pub fn spawn_session_execution_with_inputs(
1227 args: SessionExecutionArgs,
1228 inputs: Option<crate::runtime::config::UntrustedExecutionInputs>,
1229) {
1230 spawn_session_execution_with_upstream_retry(args, inputs, Default::default());
1231}
1232
1233pub fn spawn_session_execution_with_upstream_retry(
1236 args: SessionExecutionArgs,
1237 inputs: Option<crate::runtime::config::UntrustedExecutionInputs>,
1238 upstream_retry: bamboo_llm::retry::UpstreamRetry,
1239) {
1240 let span_session_id = args.session_id.clone();
1241 let session_span = tracing::info_span!("agent_execution", session_id = %span_session_id);
1242 let is_root = args.session.kind == bamboo_domain::SessionKind::Root;
1243 let outer = if is_root {
1244 match args
1245 .execution_reservation
1246 .physical_admission
1247 .as_ref()
1248 .map(|admission| admission.attach(&args.session))
1249 .transpose()
1250 {
1251 Ok(outer) => outer,
1252 Err(error) => {
1253 tokio::spawn(async move {
1256 let _ = args
1257 .mpsc_tx
1258 .send(AgentEvent::Error {
1259 message: error.to_string(),
1260 })
1261 .await;
1262 args.execution_reservation.abandon().await;
1263 });
1264 return;
1265 }
1266 }
1267 } else {
1268 None
1269 };
1270 let task_outer = outer.clone();
1271
1272 let task = tokio::spawn(
1273 async move {
1274 let SessionExecutionArgs {
1275 agent,
1276 session_id,
1277 mut session,
1278 mut execution_reservation,
1279 tools_override,
1280 provider_override,
1281 model_roster,
1282 reasoning_effort,
1283 reasoning_effort_source,
1284 auxiliary_model_resolver,
1285 disabled_filter_resolver,
1286 disabled_tools,
1287 disabled_skill_ids,
1288 selected_skill_ids,
1289 selected_skill_mode,
1290 mpsc_tx,
1291 mut history_commit_barrier,
1292 image_fallback,
1293 gold_config,
1294 guardian_config,
1295 guardian_spawner,
1296 bash_resume_hook,
1297 bash_completion_sink,
1298 app_data_dir,
1299 run_budget,
1300 runners,
1301 sessions_cache,
1302 on_complete,
1303 child_completion_handler,
1304 } = args;
1305
1306 if !execution_reservation.matches_execution_target(&session_id, &session.id, &runners) {
1307 tracing::error!(
1308 %session_id,
1309 domain_session_id = %session.id,
1310 reservation_session_id = %execution_reservation.session_id(),
1311 run_id = %execution_reservation.run_id(),
1312 same_runner_registry = Arc::ptr_eq(&execution_reservation.runners, &runners),
1313 "refusing mismatched session execution reservation"
1314 );
1315 execution_reservation.abandon().await;
1316 return;
1317 }
1318 if let Err(error) = execution_reservation.ensure_registered().await {
1319 tracing::warn!(
1320 %session_id,
1321 run_id = %execution_reservation.run_id(),
1322 %error,
1323 "session execution could not adopt its router activation owner"
1324 );
1325 return;
1326 }
1327 if execution_reservation.root_actor.is_none()
1328 && agent.persistence().root_actor_execution_required(&session)
1329 {
1330 let message =
1331 "Root execution route requires a bound Actor writer and event handoff";
1332 tracing::warn!(%session_id, "{message}");
1333 let _ = mpsc_tx
1334 .send(AgentEvent::Error {
1335 message: message.into(),
1336 })
1337 .await;
1338 execution_reservation.abandon().await;
1339 return;
1340 }
1341 if let Err(error) = execution_reservation.bind_execution_persistence(&agent, &session) {
1344 let _ = mpsc_tx
1345 .send(AgentEvent::Error {
1346 message: error.to_string(),
1347 })
1348 .await;
1349 execution_reservation.abandon().await;
1350 return;
1351 }
1352 let ownership = match &execution_reservation.activation {
1353 SessionExecutionActivationOwnership::Registered(run) => Some(run.ownership()),
1354 _ => None,
1355 };
1356 let admission_persistence = execution_reservation
1357 .execution_persistence()
1358 .unwrap_or_else(|| agent.persistence().clone());
1359 if let (Some(router), Some(ownership)) = (agent.activation_router(), ownership.as_ref())
1360 {
1361 if !router.bind_run_lifetime(ownership, &session).await {
1362 let _ = mpsc_tx
1363 .send(AgentEvent::Error {
1364 message: "registered run Session lifetime differs".into(),
1365 })
1366 .await;
1367 execution_reservation.abandon().await;
1368 return;
1369 }
1370 }
1371 if let Err(error) = crate::runtime::counted_loop::admit_registered(
1372 &mut session,
1373 ownership.as_ref(),
1374 agent.activation_router(),
1375 admission_persistence.as_ref(),
1376 )
1377 .await
1378 {
1379 let _ = mpsc_tx
1380 .send(AgentEvent::Error {
1381 message: error.to_string(),
1382 })
1383 .await;
1384 execution_reservation.abandon().await;
1385 return;
1386 }
1387 let mut root_actor = execution_reservation.root_actor.take().map(|owner| *owner);
1388 let root_actor_bound = root_actor.is_some();
1389 let persistence = execution_reservation
1390 .execution_persistence()
1391 .unwrap_or_else(|| agent.persistence().clone());
1392 let agent = Arc::new(agent.with_execution_persistence(persistence));
1393 let activation_run_id = execution_reservation.run_id().to_string();
1394 if session.kind == bamboo_agent_core::SessionKind::Child {
1395 if let Err(error) = execution_reservation
1399 .arm_child_host_startup_inner(&agent, &session, sessions_cache.clone(), None)
1400 .await
1401 {
1402 let _ = mpsc_tx
1403 .send(AgentEvent::Error {
1404 message: error.to_string(),
1405 })
1406 .await;
1407 execution_reservation.abandon().await;
1408 return;
1409 }
1410 if let Err(error) = agent
1411 .storage()
1412 .begin_host_run(&session, &activation_run_id)
1413 .await
1414 {
1415 let _ = mpsc_tx
1416 .send(AgentEvent::Error {
1417 message: error.to_string(),
1418 })
1419 .await;
1420 execution_reservation
1421 .fail_child_host_startup(&error.to_string())
1422 .await;
1423 execution_reservation.abandon().await;
1424 return;
1425 }
1426 session.set_host_run_id(&activation_run_id);
1427 session.set_last_run_status("running");
1428 }
1429 let prior_message_ids = session.messages.iter().map(|m| m.id.clone()).collect();
1430 let (cancel_token, mut activation_registration) =
1431 execution_reservation.disarm_for_execution();
1432
1433 let model = model_roster.model.clone().unwrap_or_default();
1437
1438 let initial_message = initial_user_message_for_session(&session);
1439 let selected_skill_ids =
1440 selected_skill_ids.or_else(|| selected_skill_ids_for_session(&session));
1441 let selected_skill_mode =
1442 selected_skill_mode.or_else(|| selected_skill_mode_for_session(&session));
1443
1444 tracing::info!(
1445 "[{}] Using resolved session model: {}, reasoning_effort={}, reasoning_source={}",
1446 session_id,
1447 model,
1448 reasoning_effort
1449 .map(ReasoningEffort::as_str)
1450 .unwrap_or("none"),
1451 reasoning_effort_source
1452 );
1453
1454 crate::session_app::execution_prep::prepare_session_for_execution(
1461 &mut session,
1462 None,
1463 Some(&model),
1464 );
1465
1466 let system_prompt = system_prompt_for_debug_log(&mut session);
1467 if let Some(prompt) = system_prompt.as_ref() {
1468 log_base_system_prompt_snapshot(&session_id, prompt);
1469 }
1470
1471 let inherited_child_wait = agent.persistence().inherited_child_wait();
1475 let execute_request = build_execute_request(
1476 initial_message,
1477 mpsc_tx.clone(),
1478 cancel_token,
1479 ExecuteRequestParams {
1480 tools: tools_override,
1481 provider_override,
1482 model_roster,
1483 reasoning_effort,
1484 auxiliary_model_resolver,
1485 disabled_filter_resolver,
1486 disabled_tools,
1487 disabled_skill_ids,
1488 selected_skill_ids,
1489 selected_skill_mode,
1490 image_fallback,
1491 gold_config,
1492 guardian_config,
1493 guardian_spawner,
1494 bash_resume_hook,
1495 bash_completion_sink,
1496 app_data_dir,
1497 run_budget,
1498 },
1499 );
1500
1501 let result = {
1508 use futures::FutureExt;
1509 match std::panic::AssertUnwindSafe(agent.execute_with_inputs_and_upstream_retry(
1510 &mut session,
1511 execute_request,
1512 inputs,
1513 upstream_retry,
1514 ))
1515 .catch_unwind()
1516 .await
1517 {
1518 Ok(result) => result,
1519 Err(panic) => {
1520 if let Some(outer) = &task_outer {
1521 outer
1522 .unconfirmed("Root loop unwound; physical effects are unconfirmed");
1523 }
1524 let message = panic
1525 .downcast_ref::<&str>()
1526 .map(|s| (*s).to_string())
1527 .or_else(|| panic.downcast_ref::<String>().cloned())
1528 .unwrap_or_else(|| "non-string panic payload".to_string());
1529 tracing::error!(
1530 "[{}] agent execution panicked; finalizing as terminal error: {}",
1531 session_id,
1532 message
1533 );
1534 Err(AgentError::LLM(format!(
1535 "agent execution panicked: {message}"
1536 )))
1537 }
1538 }
1539 };
1540
1541 if let Some(error_event) = terminal_error_event_for_result(&result) {
1543 let _ = mpsc_tx.send(error_event).await;
1544 }
1545
1546 let suspended_non_terminal = result.is_ok() && session.has_runtime_wait();
1563 match &result {
1564 Ok(()) if suspended_non_terminal => {
1565 session.set_last_run_status("suspended");
1566 session.clear_last_run_error();
1567 }
1568 Ok(()) => {
1569 session.set_last_run_status("completed");
1570 session.clear_last_run_error();
1571 }
1572 Err(error) if error.is_cancelled() => {
1573 session.set_last_run_status("cancelled");
1574 session.set_last_run_error(error.to_string());
1575 }
1576 Err(error) => {
1577 session.set_last_run_status("error");
1578 session.set_last_run_error(error.to_string());
1579 }
1580 }
1581
1582 session.synchronize_terminal_runtime_status();
1583
1584 if let Some(on_complete) = on_complete {
1588 on_complete(SessionExecutionOutcome::from_result(&result), &mut session).await;
1589 }
1590
1591 let executed_admitted_generation = session
1595 .session_inbox_admission()
1596 .map_or(0, |state| state.last_admitted_sequence);
1597 let legacy_migration =
1598 crate::runtime::runner::state_bridge::migrate_legacy_pending_only(
1599 &mut session,
1600 Some(agent.storage()),
1601 Some(agent.persistence()),
1602 agent.session_inbox(),
1603 )
1604 .await;
1605 let pending_boundary_generation = session
1606 .session_inbox_admission()
1607 .and_then(|state| state.pending_activation_generation());
1608 let pending_generation = match (
1609 pending_boundary_generation,
1610 legacy_migration.highest_generation,
1611 ) {
1612 (Some(left), Some(right)) => Some(left.max(right)),
1613 (left, right) => left.or(right),
1614 };
1615 if let (Some(generation), Some(router)) =
1616 (pending_generation, agent.activation_router())
1617 {
1618 let activation_ready = if let Some(inbox) = agent.session_inbox() {
1619 match inbox
1620 .mark_activation_eligible(
1621 &session_id,
1622 generation,
1623 bamboo_domain::SessionActivationPolicy::InterruptSpecificWait,
1624 )
1625 .await
1626 {
1627 Ok(()) => true,
1628 Err(error) => {
1629 tracing::error!(
1630 session_id = %session_id,
1631 %error,
1632 "failed to persist unadmitted SessionInbox activation watermark"
1633 );
1634 false
1635 }
1636 }
1637 } else {
1638 false
1639 };
1640 if activation_ready {
1641 if let Err(error) = bamboo_domain::SessionActivationPort::request_activation(
1642 router.as_ref(),
1643 &session_id,
1644 generation,
1645 )
1646 .await
1647 {
1648 tracing::error!(
1649 session_id = %session_id,
1650 %error,
1651 "failed to hand unadmitted SessionInbox generation to activation router"
1652 );
1653 }
1654 }
1655 }
1656
1657 if let Some(registration) = activation_registration.as_mut() {
1661 registration.begin_finalization().await;
1662 }
1663
1664 super::child_completion::ChildCompletionSource::prepare(
1668 &mut session,
1669 &activation_run_id,
1670 &prior_message_ids,
1671 );
1672 let saved = save_finalized_runtime_with_inherited_child_wait(
1673 agent.persistence().as_ref(),
1674 &mut session,
1675 inherited_child_wait
1676 .as_ref()
1677 .map(|inherited| inherited.wait()),
1678 )
1679 .await;
1680 let history_committed = saved.is_ok();
1681 let completion_source =
1682 super::child_completion::ChildCompletionSource::after_final_save(
1683 &session,
1684 history_committed,
1685 );
1686 let completed_at = chrono::Utc::now();
1687 if let Err(error) = &saved {
1688 if let Some(outer) = &task_outer {
1689 outer.unconfirmed(format!(
1690 "Root final persistence outcome unconfirmed: {error}"
1691 ));
1692 }
1693 tracing::warn!("[{}] Failed to save session: {}", session_id, error);
1694 }
1695
1696 if history_committed
1702 && !history_commit_barrier
1703 .send_and_wait(&mpsc_tx, session_id.clone())
1704 .await
1705 {
1706 if let Some(outer) = &task_outer {
1707 outer.unconfirmed("Root history publication barrier unconfirmed");
1708 }
1709 tracing::warn!(
1710 session_id = %session_id,
1711 "history commit barrier could not be published before runner finalization"
1712 );
1713 }
1714
1715 finalize_runner_exact(&runners, &session_id, &activation_run_id, &result).await;
1721
1722 #[cfg(test)]
1723 if let Some(outer) = &task_outer {
1724 outer.pause_before_actor_finish().await;
1725 }
1726
1727 if let Some(owner) = root_actor.take() {
1728 let outcome = if result.as_ref().is_err_and(|error| error.is_cancelled()) {
1729 bamboo_domain::ActorActivationFinish::Cancelled
1730 } else if result.is_err() || !history_committed {
1731 bamboo_domain::ActorActivationFinish::Failed
1732 } else {
1733 bamboo_domain::ActorActivationFinish::Succeeded
1734 };
1735 let finished = owner.finish_observed(outcome).await;
1736 if let Err(error) = &finished {
1737 tracing::warn!(%session_id, %error, "Root Actor closure unconfirmed");
1738 }
1739 if let Some(outer) = &task_outer {
1740 outer.actor_finished(finished);
1741 }
1742 }
1743
1744 let finalization = if let Some(registration) = activation_registration.take() {
1745 registration.finish(executed_admitted_generation).await
1746 } else {
1747 Ok(None)
1748 };
1749 if let Err(error) = finalization {
1750 if let Some(outer) = &task_outer {
1751 outer.unconfirmed(format!(
1752 "Root registration finalization unconfirmed: {error}"
1753 ));
1754 }
1755 tracing::error!(
1759 session_id = %session_id,
1760 %error,
1761 "failed to activate successor for finalization-racing SessionInbox delivery"
1762 );
1763 }
1764
1765 let child_completion = child_completion_handler.filter(|_| {
1774 session.kind == bamboo_agent_core::SessionKind::Child
1775 && session.parent_session_id.is_some()
1776 });
1777 let parent_session_id = session.parent_session_id.clone();
1778 let child_status = session.last_run_status();
1779 let child_error = session.last_run_error();
1780
1781 publish_final_session_cache(&sessions_cache, session, root_actor_bound, &saved);
1782
1783 if let (Some(handler), Some(parent_session_id), Some(status)) =
1784 (child_completion, parent_session_id, child_status)
1785 {
1786 use futures::FutureExt;
1787 let completion = ChildCompletion {
1788 parent_session_id: parent_session_id.clone(),
1789 child_session_id: session_id.clone(),
1790 status,
1791 error: child_error,
1792 completed_at,
1793 source: completion_source,
1794 };
1795 if std::panic::AssertUnwindSafe(handler.on_child_completed(completion))
1796 .catch_unwind()
1797 .await
1798 .is_err()
1799 {
1800 tracing::error!(
1801 %parent_session_id,
1802 child_session_id = %session_id,
1803 "child completion handler panicked on resumed-child terminal"
1804 );
1805 }
1806 }
1807
1808 tracing::info!("[{}] Agent execution completed", session_id);
1809 if let Some(outer) = &task_outer {
1810 outer.tail_completed();
1811 }
1812 }
1813 .instrument(session_span),
1814 );
1815 if let Some(outer) = outer {
1816 outer.retain(task);
1817 } else if is_root {
1818 tokio::spawn(async move {
1821 let _ = task.await;
1822 });
1823 }
1824}
1825
1826async fn save_finalized_runtime_with_inherited_child_wait(
1828 persistence: &dyn bamboo_domain::RuntimeSessionPersistence,
1829 session: &mut Session,
1830 inherited_child_wait: Option<&bamboo_domain::session::runtime_state::WaitingForChildrenState>,
1831) -> std::io::Result<()> {
1832 if let Some(inherited) = inherited_child_wait.filter(|wait| {
1833 session
1834 .agent_runtime_state
1835 .as_ref()
1836 .and_then(|runtime| runtime.waiting_for_children.as_ref())
1837 == Some(*wait)
1838 }) {
1839 persistence
1840 .save_finalized_runtime_with_inherited_child_wait(session, inherited)
1841 .await
1842 } else {
1843 persistence.save_finalized_runtime_session(session).await
1844 }
1845}
1846
1847fn publish_final_session_cache(
1848 cache: &SessionCache,
1849 session: Session,
1850 root_actor_bound: bool,
1851 saved: &std::io::Result<()>,
1852) {
1853 let authority_conflict = saved.as_ref().err().is_some_and(|error| {
1854 error
1855 .get_ref()
1856 .is_some_and(|cause| cause.is::<bamboo_domain::SessionAuthorityConflict>())
1857 });
1858 if !authority_conflict && !root_actor_bound {
1861 cache.insert(
1862 session.id.clone(),
1863 Arc::new(crate::SessionSnapshot::new(session)),
1864 );
1865 }
1866}
1867
1868pub fn log_base_system_prompt_snapshot(session_id: &str, prompt: &str) {
1870 tracing::info!(
1871 "[{}] Base system prompt snapshot: len={} chars, has_skill={}, has_tool_guide={}, has_external_memory={}, has_task_list={}",
1872 session_id,
1873 prompt.len(),
1874 prompt.contains(SKILL_CONTEXT_START_MARKER),
1875 prompt.contains(TOOL_GUIDE_START_MARKER),
1876 prompt.contains(EXTERNAL_MEMORY_START_MARKER),
1877 prompt.contains(TASK_LIST_START_MARKER),
1878 );
1879
1880 tracing::debug!(
1881 "[{}] ========== BASE SYSTEM PROMPT SNAPSHOT ==========",
1882 session_id
1883 );
1884 tracing::debug!("[{}] Snapshot length: {} chars", session_id, prompt.len());
1885 tracing::debug!("[{}] -----------------------------------", session_id);
1886 tracing::debug!("[{}] {}", session_id, prompt);
1887 tracing::debug!(
1888 "[{}] ========== END BASE SYSTEM PROMPT SNAPSHOT ==========",
1889 session_id
1890 );
1891}
1892
1893pub fn terminal_error_event_for_result(result: &Result<(), AgentError>) -> Option<AgentEvent> {
1895 match result {
1896 Ok(_) => None,
1897 Err(error) if error.is_cancelled() => Some(AgentEvent::Error {
1898 message: "Agent execution cancelled by user".to_string(),
1899 }),
1900 Err(error) => Some(AgentEvent::Error {
1901 message: error.to_string(),
1902 }),
1903 }
1904}
1905
1906fn system_prompt_for_session(session: &Session) -> Option<String> {
1909 session
1910 .messages
1911 .iter()
1912 .find(|message| matches!(message.role, bamboo_agent_core::Role::System))
1913 .map(|message| message.content.clone())
1914}
1915
1916fn system_prompt_for_debug_log(session: &mut Session) -> Option<String> {
1920 crate::runtime::runner::session_setup::migrate_legacy_workspace_prompt(session);
1921 system_prompt_for_session(session)
1922}
1923
1924fn initial_user_message_for_session(session: &Session) -> String {
1925 session
1926 .messages
1927 .last()
1928 .filter(|message| matches!(message.role, bamboo_agent_core::Role::User))
1929 .map(|message| message.content.clone())
1930 .unwrap_or_default()
1931}
1932
1933fn selected_skill_ids_for_session(session: &Session) -> Option<Vec<String>> {
1934 session
1935 .metadata
1936 .get("selected_skill_ids")
1937 .and_then(|raw| bamboo_skills::selection::parse_selected_skill_ids_metadata(raw))
1938}
1939
1940fn selected_skill_mode_for_session(session: &Session) -> Option<String> {
1941 let value = session
1942 .metadata
1943 .get("skill_mode")
1944 .or_else(|| session.metadata.get("mode"))?;
1945 let trimmed = value.trim();
1946 if trimmed.is_empty() {
1947 None
1948 } else {
1949 Some(trimmed.to_string())
1950 }
1951}
1952
1953pub(super) struct RootActorExecution {
1954 binding: bamboo_domain::RootActorExecutionBinding,
1955 stop_renewal: CancellationToken,
1956 renewal: Option<tokio::task::JoinHandle<()>>,
1957 response_events: Vec<AgentEvent>,
1958}
1959
1960impl RootActorExecution {
1961 pub(super) fn start(binding: bamboo_domain::RootActorExecutionBinding) -> Self {
1962 let stop_renewal = CancellationToken::new();
1963 let stop = stop_renewal.clone();
1964 let directory = binding.directory.clone();
1965 let fence = binding.owner.fence.clone();
1966 let duration = binding.lease_duration;
1967 let renewal = tokio::spawn(async move {
1968 let mut interval = tokio::time::interval(duration / 3);
1969 interval.tick().await;
1970 loop {
1971 tokio::select! {
1972 _ = stop.cancelled() => break,
1973 _ = interval.tick() => {
1974 let now = chrono::Utc::now();
1975 let expires = now + chrono::Duration::from_std(duration).expect("bounded Root lease");
1976 if let Err(error) = directory.renew_activation(&fence, now, expires).await {
1977 if matches!(error, bamboo_domain::ActorDirectoryError::Storage(_)) {
1985 tracing::warn!(actor_id = %fence.actor_id, run_id = %fence.run_id,
1986 %error, "Root Actor renewal hit transient storage failure; retrying next tick");
1987 continue;
1988 }
1989 tracing::warn!(actor_id = %fence.actor_id, run_id = %fence.run_id,
1990 %error, "Root Actor renewal lost its execution authority");
1991 break;
1992 }
1993 }
1994 }
1995 }
1996 });
1997 Self {
1998 binding,
1999 stop_renewal,
2000 renewal: Some(renewal),
2001 response_events: Vec::new(),
2002 }
2003 }
2004
2005 pub(super) fn persistence(
2006 &self,
2007 ) -> std::sync::Arc<dyn bamboo_domain::RuntimeSessionPersistence> {
2008 self.binding.persistence.clone()
2009 }
2010
2011 pub(super) async fn finish(self, outcome: bamboo_domain::ActorActivationFinish) {
2012 if let Err(error) = self.finish_observed(outcome).await {
2013 tracing::warn!(%error, "Root Actor closure unconfirmed");
2014 }
2015 }
2016
2017 async fn finish_observed(
2018 mut self,
2019 outcome: bamboo_domain::ActorActivationFinish,
2020 ) -> std::io::Result<()> {
2021 self.stop_renewal.cancel();
2022 let renewal = match self.renewal.take() {
2023 Some(renewal) => renewal.await.map_err(std::io::Error::other),
2024 None => Ok(()),
2025 };
2026 let directory = self.binding.directory.clone();
2027 let owner = self.binding.owner.clone();
2028 let publish = self.binding.terminal_publisher.clone();
2029 let finish = tokio::spawn(async move {
2030 directory
2031 .finish_root_activation_observed(&owner, chrono::Utc::now(), outcome, publish)
2032 .await
2033 .map(|_| ())
2034 .map_err(std::io::Error::other)
2035 });
2036 self.binding.disarm_abandonment();
2037 finish.await.map_err(std::io::Error::other)??;
2038 renewal
2039 }
2040}
2041
2042impl Drop for RootActorExecution {
2043 fn drop(&mut self) {
2044 self.stop_renewal.cancel();
2045 }
2048}
2049
2050#[cfg(test)]
2051#[path = "agent_spawn_peer_privacy_tests.rs"]
2052mod peer_privacy_tests;
2053
2054#[cfg(test)]
2055mod reservation_tests {
2056 use super::*;
2057 use crate::runtime::execution::runner_state::AgentStatus;
2058
2059 struct CompletionAtFinalSave {
2060 directory: std::path::PathBuf,
2061 persistence: bamboo_storage::LockedSessionStore,
2062 cleared: std::sync::atomic::AtomicBool,
2063 }
2064
2065 impl CompletionAtFinalSave {
2066 async fn clear_wait(&self, id: &str) {
2067 let mut process = std::process::Command::new(std::env::current_exe().unwrap())
2068 .args(["--exact", "runtime::execution::agent_spawn::reservation_tests::inherited_child_wait_adopts_independent_completion_at_final_save", "--nocapture"])
2069 .env("BAMBOO_FINAL_WAIT_COMPLETION_DIR", &self.directory)
2070 .env("BAMBOO_FINAL_WAIT_COMPLETION_ID", id)
2071 .spawn().unwrap();
2072 let deadline = std::time::Instant::now() + std::time::Duration::from_secs(15);
2073 loop {
2074 if let Some(status) = process.try_wait().unwrap() {
2075 assert!(status.success(), "completion process failed");
2076 break;
2077 }
2078 if std::time::Instant::now() >= deadline {
2079 let _ = process.kill();
2080 let _ = process.wait();
2081 panic!("bounded completion process deadline");
2082 }
2083 std::thread::sleep(std::time::Duration::from_millis(5));
2084 }
2085 self.cleared
2086 .store(true, std::sync::atomic::Ordering::SeqCst);
2087 }
2088 }
2089
2090 #[async_trait::async_trait]
2091 impl bamboo_domain::RuntimeSessionPersistence for CompletionAtFinalSave {
2092 async fn save_runtime_session(&self, session: &mut Session) -> std::io::Result<()> {
2093 self.clear_wait(&session.id).await;
2094 self.persistence.merge_save_finalized_runtime(session).await
2095 }
2096 async fn save_finalized_runtime_with_inherited_child_wait(
2097 &self,
2098 session: &mut Session,
2099 inherited: &bamboo_domain::session::runtime_state::WaitingForChildrenState,
2100 ) -> std::io::Result<()> {
2101 self.clear_wait(&session.id).await;
2102 self.persistence
2103 .merge_save_inherited_child_wait_and_publish(session, inherited, |_, _| {})
2104 .await
2105 }
2106 }
2107
2108 #[tokio::test]
2109 async fn inherited_child_wait_adopts_independent_completion_at_final_save() {
2110 use bamboo_agent_core::storage::Storage;
2111 use bamboo_domain::session::runtime_state::{
2112 AgentRuntimeState, AgentStatusState, ChildWaitPolicy, WaitingForChildrenState,
2113 };
2114 if let Some(directory) = std::env::var_os("BAMBOO_FINAL_WAIT_COMPLETION_DIR") {
2115 let remote = bamboo_storage::SessionStoreV2::new(directory.into())
2116 .await
2117 .unwrap();
2118 let id = std::env::var("BAMBOO_FINAL_WAIT_COMPLETION_ID").unwrap();
2119 let mut latest = remote.load_session(&id).await.unwrap().unwrap();
2120 let runtime = latest.agent_runtime_state.as_mut().unwrap();
2121 runtime.waiting_for_children = None;
2122 runtime.status = AgentStatusState::Idle;
2123 runtime.suspension = None;
2124 latest.metadata.remove("runtime.suspend_reason");
2125 remote.save_runtime_state(&latest).await.unwrap();
2126 return;
2127 }
2128 let temp = tempfile::tempdir().unwrap();
2129 let local = Arc::new(
2130 bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
2131 .await
2132 .unwrap(),
2133 );
2134 let mut runner = Session::new("cross-host-final-wait", "model");
2138 let inherited = WaitingForChildrenState::for_children(
2139 vec!["child".into()],
2140 ChildWaitPolicy::All,
2141 chrono::Utc::now(),
2142 );
2143 let mut runtime = AgentRuntimeState::new("run");
2144 runtime.waiting_for_children = Some(inherited.clone());
2145 runtime.status = AgentStatusState::Suspended;
2146 runner.agent_runtime_state = Some(runtime);
2147 runner.metadata.insert(
2148 "runtime.suspend_reason".into(),
2149 "waiting_for_children".into(),
2150 );
2151 runner.set_last_run_status("suspended");
2152 local.save_session(&runner).await.unwrap();
2153 let persistence = CompletionAtFinalSave {
2154 directory: temp.path().to_path_buf(),
2155 persistence: bamboo_storage::LockedSessionStore::new(local.clone()),
2156 cleared: std::sync::atomic::AtomicBool::new(false),
2157 };
2158 save_finalized_runtime_with_inherited_child_wait(
2159 &persistence,
2160 &mut runner,
2161 Some(&inherited),
2162 )
2163 .await
2164 .unwrap();
2165 assert!(persistence
2166 .cleared
2167 .load(std::sync::atomic::Ordering::SeqCst));
2168 let durable = local.load_session(&runner.id).await.unwrap().unwrap();
2169 for session in [&runner, &durable] {
2170 let runtime = session.agent_runtime_state.as_ref().unwrap();
2171 assert!(
2172 runtime.waiting_for_children.is_none(),
2173 "independent completion must not be resurrected"
2174 );
2175 assert_eq!(runtime.status, AgentStatusState::Idle);
2176 assert!(runtime.suspension.is_none());
2177 assert!(!session.metadata.contains_key("runtime.suspend_reason"));
2178 assert_eq!(session.last_run_status().as_deref(), Some("completed"));
2179 }
2180 }
2181
2182 #[tokio::test]
2183 async fn inherited_child_wait_rejects_aba_and_missing_birth_without_cache_publication() {
2184 use bamboo_agent_core::storage::Storage;
2185 use bamboo_domain::session::runtime_state::{
2186 AgentRuntimeState, ChildWaitPolicy, WaitingForChildrenState,
2187 };
2188 use bamboo_storage::{LockedSessionStore, SessionStoreV2};
2189
2190 for missing in [false, true] {
2191 let temp = tempfile::tempdir().unwrap();
2192 let storage = Arc::new(
2193 SessionStoreV2::new(temp.path().to_path_buf())
2194 .await
2195 .unwrap(),
2196 );
2197 let persistence = LockedSessionStore::new(storage.clone());
2198 let cache: SessionCache = Arc::default();
2199 storage
2200 .save_session(&Session::new("root", "model"))
2201 .await
2202 .unwrap();
2203 let mut stale = Session::new("nested-child-parent", "model");
2204 stale.kind = bamboo_domain::SessionKind::Child;
2205 stale.parent_session_id = Some("root".into());
2206 stale.root_session_id = "root".into();
2207 stale.spawn_depth = 1;
2208 let wait = WaitingForChildrenState::for_children(
2209 vec!["grandchild".into()],
2210 ChildWaitPolicy::All,
2211 chrono::Utc::now(),
2212 );
2213 let mut runtime = AgentRuntimeState::new("old-run");
2214 runtime.waiting_for_children = Some(wait.clone());
2215 stale.agent_runtime_state = Some(runtime);
2216 stale.add_message(bamboo_agent_core::Message::assistant("STALE", None));
2217 storage.save_session(&stale).await.unwrap();
2218
2219 let mut replacement = stale.clone();
2220 replacement.created_at += chrono::Duration::seconds(1);
2221 replacement.messages.clear();
2222 replacement.add_message(bamboo_agent_core::Message::assistant("REPLACEMENT", None));
2223 replacement
2224 .agent_runtime_state
2225 .as_mut()
2226 .unwrap()
2227 .waiting_for_children = None;
2228 storage.delete_session(&stale.id).await.unwrap();
2229 if !missing {
2230 storage.save_session(&replacement).await.unwrap();
2231 }
2232 if !missing {
2233 cache.insert(
2234 replacement.id.clone(),
2235 Arc::new(crate::SessionSnapshot::new(replacement.clone())),
2236 );
2237 }
2238
2239 let saved = save_finalized_runtime_with_inherited_child_wait(
2240 &persistence,
2241 &mut stale,
2242 Some(&wait),
2243 )
2244 .await;
2245 assert!(saved.is_err(), "stale incarnation must not be saved");
2246 publish_final_session_cache(&cache, stale, false, &saved);
2249 if missing {
2250 assert!(
2251 read_cached_session(&cache, &replacement.id).is_none(),
2252 "deleted incarnation must not be resurrected in cache"
2253 );
2254 } else {
2255 let cached = read_cached_session(&cache, &replacement.id).unwrap();
2256 assert_eq!(
2257 cached.created_at, replacement.created_at,
2258 "ABA rejection must retain replacement cache birth"
2259 );
2260 assert_eq!(
2261 serde_json::to_value(&cached.messages).unwrap(),
2262 serde_json::to_value(&replacement.messages).unwrap(),
2263 "ABA rejection must retain replacement cache transcript"
2264 );
2265 }
2266 assert!(saved
2267 .as_ref()
2268 .unwrap_err()
2269 .get_ref()
2270 .is_some_and(|cause| cause.is::<bamboo_domain::SessionAuthorityConflict>()));
2271 let durable = storage.load_session(&replacement.id).await.unwrap();
2272 if missing {
2273 assert!(durable.is_none());
2274 } else {
2275 let durable = durable.unwrap();
2276 assert_eq!(durable.created_at, replacement.created_at);
2277 assert_eq!(
2278 serde_json::to_value(&durable.messages).unwrap(),
2279 serde_json::to_value(&replacement.messages).unwrap()
2280 );
2281 }
2282 }
2283 }
2284
2285 #[test]
2286 fn debug_prompt_snapshot_migrates_legacy_host_paths_before_logging() {
2287 let legacy_project = format!(
2288 "{}\nProject ID: legacy-project\nProject path: /private/legacy-project\nProject home: /private/legacy-home\n{}",
2289 crate::runtime::context::PROJECT_CONTEXT_START_MARKER,
2290 crate::runtime::context::PROJECT_CONTEXT_END_MARKER,
2291 );
2292 let legacy_workspace =
2293 crate::runtime::context::build_workspace_prompt_context("/private/legacy-workspace")
2294 .expect("legacy workspace context");
2295 let legacy_instruction = format!(
2296 "{}\nSource: /private/legacy-workspace/AGENTS.md\nlegacy policy\n{}",
2297 crate::runtime::context::instruction::INSTRUCTION_CONTEXT_START_MARKER,
2298 crate::runtime::context::instruction::INSTRUCTION_CONTEXT_END_MARKER,
2299 );
2300 let mut session = Session::new("legacy-debug-log", "model");
2301 session.add_message(bamboo_agent_core::Message::system(format!(
2302 "Base prompt\n\n{legacy_project}\n\n{legacy_workspace}\n\n{legacy_instruction}"
2303 )));
2304
2305 let prompt = system_prompt_for_debug_log(&mut session).expect("normalized System");
2306
2307 assert_eq!(prompt, "Base prompt");
2308 assert_eq!(
2309 session.workspace_path_meta().as_deref(),
2310 Some("/private/legacy-workspace")
2311 );
2312 for private_path in [
2313 "/private/legacy-project",
2314 "/private/legacy-home",
2315 "/private/legacy-workspace",
2316 ] {
2317 assert!(!prompt.contains(private_path));
2318 }
2319 assert!(!prompt.contains(crate::runtime::context::PROJECT_CONTEXT_START_MARKER));
2320 assert!(!prompt.contains(crate::runtime::context::WORKSPACE_CONTEXT_START_MARKER));
2321 assert!(!prompt
2322 .contains(crate::runtime::context::instruction::INSTRUCTION_CONTEXT_START_MARKER));
2323 }
2324
2325 #[test]
2326 fn reservation_target_requires_domain_id_and_exact_runner_registry() {
2327 let runners = Arc::new(RwLock::new(HashMap::new()));
2328 let other_runners = Arc::new(RwLock::new(HashMap::new()));
2329 let mut reservation = SessionExecutionReservation {
2330 session_id: "session-a".to_string(),
2331 run_id: "run-a".to_string(),
2332 cancel_token: CancellationToken::new(),
2333 runners: runners.clone(),
2334 activation: SessionExecutionActivationOwnership::Unrouted,
2335 root_actor: None,
2336 physical_admission: None,
2337 execution_persistence: None,
2338 persistence_bound: false,
2339 child_startup: None,
2340 child_closure: None,
2341 selected_input: false,
2342 armed: true,
2343 };
2344
2345 assert!(reservation.matches_execution_target("session-a", "session-a", &runners));
2346 assert!(!reservation.matches_execution_target("session-b", "session-a", &runners));
2347 assert!(!reservation.matches_execution_target("session-a", "session-b", &runners));
2348 assert!(!reservation.matches_execution_target("session-a", "session-a", &other_runners));
2349 reservation.armed = false;
2350 }
2351
2352 #[tokio::test]
2353 async fn forced_same_run_registration_rejection_never_cancels_live_owner() {
2354 let runners = Arc::new(RwLock::new(HashMap::new()));
2355 let mut runner = AgentRunner::new();
2356 runner.status = AgentStatus::Running;
2357 let run_id = runner.run_id.clone();
2358 let live_cancel_token = runner.cancel_token.clone();
2359 runners.write().await.insert("same-run".to_string(), runner);
2360
2361 let router = SessionActivationRouter::new();
2362 let mut live_registration = router
2363 .register_run("same-run", &run_id)
2364 .await
2365 .expect("first registration owns the run");
2366 let mut rejected = SessionExecutionReservation {
2367 session_id: "same-run".to_string(),
2368 run_id: run_id.clone(),
2369 cancel_token: live_cancel_token.clone(),
2370 runners: runners.clone(),
2371 activation: SessionExecutionActivationOwnership::RegistrationPending(router.clone()),
2372 root_actor: None,
2373 physical_admission: None,
2374 execution_persistence: None,
2375 persistence_bound: false,
2376 child_startup: None,
2377 child_closure: None,
2378 selected_input: false,
2379 armed: true,
2380 };
2381
2382 let error = match rejected.ensure_registered().await {
2383 Ok(()) => panic!("a duplicate registration for the same run must be rejected"),
2384 Err(error) => error,
2385 };
2386 assert_eq!(error.existing_run_id(), run_id);
2387 drop(rejected);
2388
2389 assert!(!live_cancel_token.is_cancelled());
2390 assert!(matches!(
2391 runners
2392 .read()
2393 .await
2394 .get("same-run")
2395 .map(|runner| &runner.status),
2396 Some(AgentStatus::Running)
2397 ));
2398 assert!(router.owns_run("same-run", &run_id).await);
2399 live_registration.begin_finalization().await;
2400 assert_eq!(live_registration.finish(0).await.unwrap(), None);
2401 }
2402}