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 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>),
51 Registered(SessionRunRegistration),
54}
55
56pub struct SessionExecutionReservation {
62 session_id: String,
63 run_id: String,
64 cancel_token: CancellationToken,
65 runners: Arc<RwLock<HashMap<String, AgentRunner>>>,
66 activation: SessionExecutionActivationOwnership,
67 root_actor: Option<Box<RootActorExecution>>,
68 execution_persistence: Option<Arc<dyn bamboo_domain::RuntimeSessionPersistence>>,
69 persistence_bound: bool,
70 selected_input: bool,
71 armed: bool,
72}
73
74impl SessionExecutionReservation {
75 pub fn session_id(&self) -> &str {
76 &self.session_id
77 }
78
79 pub fn run_id(&self) -> &str {
80 &self.run_id
81 }
82
83 pub fn cancel_token(&self) -> &CancellationToken {
84 &self.cancel_token
85 }
86
87 pub async fn bind_root_actor(
90 &mut self,
91 agent: &Agent,
92 session: &Session,
93 ) -> std::io::Result<()> {
94 if self.root_actor.is_none() {
95 if self.session_id != session.id {
96 return Err(std::io::Error::new(
97 std::io::ErrorKind::InvalidInput,
98 "Root binding target mismatch",
99 ));
100 }
101 if let Some(binding) = agent
102 .persistence()
103 .bind_root_actor_execution(session, &self.run_id)
104 .await?
105 {
106 self.root_actor = Some(Box::new(RootActorExecution::start(binding)));
107 }
108 }
109 self.bind_execution_persistence(agent, session)
110 }
111
112 fn bind_execution_persistence(
113 &mut self,
114 agent: &Agent,
115 session: &Session,
116 ) -> std::io::Result<()> {
117 if self.session_id != session.id {
118 return Err(std::io::Error::other(
119 "execution persistence target mismatch",
120 ));
121 }
122 if self.persistence_bound {
123 return Ok(());
124 }
125 let persistence = self
126 .root_actor
127 .as_ref()
128 .map_or_else(|| agent.persistence().clone(), |owner| owner.persistence());
129 if let Some(inherited) = persistence
130 .inherited_child_wait()
131 .or_else(|| bamboo_domain::InheritedChildWait::capture(session))
132 {
133 inherited.validate_session(session)?;
134 self.execution_persistence = Some(persistence.bind_inherited_child_wait(inherited)?);
135 } else if self.root_actor.is_some() {
136 self.execution_persistence = Some(persistence);
137 }
138 self.persistence_bound = true;
139 Ok(())
140 }
141
142 pub fn execution_persistence(
143 &self,
144 ) -> Option<Arc<dyn bamboo_domain::RuntimeSessionPersistence>> {
145 self.execution_persistence.clone()
146 }
147
148 pub fn bind_root_inbox_receipt(
151 &mut self,
152 receipt: bamboo_domain::SessionInboxReceipt,
153 ) -> std::io::Result<()> {
154 if !self.selected_input || self.root_actor.is_none() || self.cancel_token.is_cancelled() {
155 return Err(std::io::Error::other(
156 "Root receipt requires a live bound reservation",
157 ));
158 }
159 let persistence = self
160 .execution_persistence
161 .as_ref()
162 .ok_or_else(|| std::io::Error::other("Root execution persistence unavailable"))?;
163 self.execution_persistence = Some(persistence.bind_root_inbox_receipt(receipt)?);
164 Ok(())
165 }
166
167 pub async fn retain_selected_run_output(&self) -> std::io::Result<bool> {
170 if !self.selected_input {
171 return Ok(false);
172 }
173 let owner = self
174 .root_actor_writer()
175 .ok_or_else(|| std::io::Error::other("Selected output requires its Root writer"))?;
176 let receipt = self
177 .execution_persistence
178 .as_ref()
179 .and_then(|p| p.root_inbox_receipt())
180 .cloned()
181 .ok_or_else(|| std::io::Error::other("Selected output requires its bound receipt"))?;
182 if owner.fence.actor_id != self.session_id
183 || owner.fence.run_id != self.run_id
184 || self.cancel_token.is_cancelled()
185 {
186 return Err(std::io::Error::other(
187 "Selected output reservation is not live",
188 ));
189 }
190 let runners = self.runners.read().await;
191 let runner = runners
192 .get(&self.session_id)
193 .filter(|r| r.run_id == self.run_id && !r.cancel_token.is_cancelled())
194 .ok_or_else(|| std::io::Error::other("Selected output runner is unavailable"))?;
195 runner
196 .visible_messages
197 .retain_selected_output(&owner, receipt)?;
198 Ok(true)
199 }
200
201 pub fn root_actor_writer(&self) -> Option<bamboo_domain::RootActorRuntimeWrite> {
202 self.root_actor
203 .as_ref()
204 .map(|owner| owner.binding.owner.clone())
205 }
206
207 pub fn queue_root_response_event(&mut self, event: AgentEvent) -> Option<AgentEvent> {
210 let Some(owner) = self.root_actor.as_mut() else {
211 return Some(event);
212 };
213 owner.response_events.push(event);
214 None
215 }
216
217 pub fn take_root_response_events(&mut self) -> Vec<AgentEvent> {
218 self.root_actor
219 .as_mut()
220 .map(|owner| std::mem::take(&mut owner.response_events))
221 .unwrap_or_default()
222 }
223
224 pub async fn finish_root_actor(&mut self, outcome: bamboo_domain::ActorActivationFinish) {
227 if let Some(owner) = self.root_actor.take() {
228 (*owner).finish(outcome).await;
229 }
230 }
231
232 pub(crate) fn from_activation_placeholder(
238 session_id: impl Into<String>,
239 reservation: RunnerReservation,
240 router: Arc<SessionActivationRouter>,
241 runners: Arc<RwLock<HashMap<String, AgentRunner>>>,
242 ) -> Self {
243 Self {
244 session_id: session_id.into(),
245 run_id: reservation.run_id,
246 cancel_token: reservation.cancel_token,
247 runners,
248 activation: SessionExecutionActivationOwnership::UnpublishedActivation(router),
249 root_actor: None,
250 execution_persistence: None,
251 persistence_bound: false,
252 selected_input: false,
253 armed: true,
254 }
255 }
256
257 pub(crate) fn from_pending_registration(
261 session_id: impl Into<String>,
262 reservation: RunnerReservation,
263 router: Option<Arc<SessionActivationRouter>>,
264 runners: Arc<RwLock<HashMap<String, AgentRunner>>>,
265 ) -> Self {
266 let activation = match router {
267 Some(router) => SessionExecutionActivationOwnership::RegistrationPending(router),
268 None => SessionExecutionActivationOwnership::Unrouted,
269 };
270 Self {
271 session_id: session_id.into(),
272 run_id: reservation.run_id,
273 cancel_token: reservation.cancel_token,
274 runners,
275 activation,
276 root_actor: None,
277 execution_persistence: None,
278 persistence_bound: false,
279 selected_input: false,
280 armed: true,
281 }
282 }
283
284 pub(crate) fn mark_activation_published(&mut self) {
287 let activation = std::mem::replace(
288 &mut self.activation,
289 SessionExecutionActivationOwnership::Unrouted,
290 );
291 self.activation = match activation {
292 SessionExecutionActivationOwnership::UnpublishedActivation(router) => {
293 SessionExecutionActivationOwnership::RegistrationPending(router)
294 }
295 other => other,
296 };
297 }
298
299 pub(crate) async fn rollback_unpublished_activation(mut self) {
304 self.armed = false;
305 self.cancel_token.cancel();
306 self.finish_root_actor(bamboo_domain::ActorActivationFinish::Cancelled)
307 .await;
308 let activation = std::mem::replace(
309 &mut self.activation,
310 SessionExecutionActivationOwnership::Unrouted,
311 );
312 debug_assert!(matches!(
313 activation,
314 SessionExecutionActivationOwnership::UnpublishedActivation(_)
315 ));
316 remove_runner_exact(&self.runners, &self.session_id, &self.run_id).await;
317 }
318
319 pub async fn ensure_registered(&mut self) -> Result<(), SessionRunRegistrationError> {
325 let router = match &self.activation {
326 SessionExecutionActivationOwnership::RegistrationPending(router) => router.clone(),
327 SessionExecutionActivationOwnership::UnpublishedActivation(_) => {
328 unreachable!("unpublished activation entered an execution adapter")
329 }
330 SessionExecutionActivationOwnership::Unrouted
331 | SessionExecutionActivationOwnership::Registered(_) => return Ok(()),
332 };
333 match register_reserved_activation(
334 router,
335 &self.runners,
336 &self.session_id,
337 &self.run_id,
338 self.selected_input,
339 )
340 .await
341 {
342 Ok(registration) => {
343 self.activation = SessionExecutionActivationOwnership::Registered(registration);
344 Ok(())
345 }
346 Err(error) => {
347 self.disarm_after_registration_rejection(error.existing_run_id());
352 Err(error)
353 }
354 }
355 }
356
357 pub async fn abandon(mut self) {
361 self.armed = false;
362 self.cancel_token.cancel();
363 if let Some(owner) = self.root_actor.take() {
364 (*owner)
365 .finish(bamboo_domain::ActorActivationFinish::Cancelled)
366 .await;
367 }
368 let activation = std::mem::replace(
369 &mut self.activation,
370 SessionExecutionActivationOwnership::Unrouted,
371 );
372 let runners = self.runners.clone();
373 let session_id = self.session_id.clone();
374 let run_id = self.run_id.clone();
375 let cleanup_session_id = session_id.clone();
376 let cleanup_run_id = run_id.clone();
377 let selected_input = self.selected_input;
378 let cleanup = tokio::spawn(async move {
379 cleanup_execution_reservation(
380 activation,
381 runners,
382 cleanup_session_id,
383 cleanup_run_id,
384 selected_input,
385 )
386 .await;
387 });
388 if let Err(error) = cleanup.await {
389 tracing::error!(
390 %session_id,
391 %run_id,
392 %error,
393 "detached execution-reservation cleanup failed"
394 );
395 }
396 }
397
398 fn disarm_after_registration_rejection(&mut self, existing_run_id: &str) {
399 if existing_run_id != self.run_id {
400 self.cancel_token.cancel();
401 }
402 self.armed = false;
403 }
404
405 pub(crate) fn matches_execution_target(
406 &self,
407 session_id: &str,
408 domain_session_id: &str,
409 runners: &Arc<RwLock<HashMap<String, AgentRunner>>>,
410 ) -> bool {
411 self.session_id == session_id
412 && domain_session_id == session_id
413 && Arc::ptr_eq(&self.runners, runners)
414 }
415
416 pub(crate) fn disarm_for_execution(
417 &mut self,
418 ) -> (CancellationToken, Option<SessionRunRegistration>) {
419 self.armed = false;
420 let registration = match std::mem::replace(
421 &mut self.activation,
422 SessionExecutionActivationOwnership::Unrouted,
423 ) {
424 SessionExecutionActivationOwnership::Registered(registration) => Some(registration),
425 SessionExecutionActivationOwnership::Unrouted => None,
426 SessionExecutionActivationOwnership::UnpublishedActivation(_) => {
427 unreachable!("unpublished activation cannot transfer to execution")
428 }
429 SessionExecutionActivationOwnership::RegistrationPending(_) => {
430 unreachable!("ensure_registered adopts every pending router registration")
431 }
432 };
433 (self.cancel_token.clone(), registration)
434 }
435}
436
437impl Drop for SessionExecutionReservation {
438 fn drop(&mut self) {
439 if !self.armed {
440 return;
441 }
442 self.armed = false;
443 self.cancel_token.cancel();
444 let activation = std::mem::replace(
445 &mut self.activation,
446 SessionExecutionActivationOwnership::Unrouted,
447 );
448 let runners = self.runners.clone();
449 let session_id = self.session_id.clone();
450 let run_id = self.run_id.clone();
451 let selected_input = self.selected_input;
452 if let Ok(runtime) = tokio::runtime::Handle::try_current() {
453 runtime.spawn(async move {
454 cleanup_execution_reservation(
455 activation,
456 runners,
457 session_id,
458 run_id,
459 selected_input,
460 )
461 .await;
462 });
463 }
464 }
465}
466
467pub enum SessionExecutionReserveOutcome {
469 Reserved(SessionExecutionReservation),
470 AlreadyRunning { run_id: String },
471}
472
473async fn register_reserved_activation(
474 router: Arc<SessionActivationRouter>,
475 runners: &Arc<RwLock<HashMap<String, AgentRunner>>>,
476 session_id: &str,
477 run_id: &str,
478 selected_input: bool,
479) -> Result<SessionRunRegistration, SessionRunRegistrationError> {
480 let mut registration = match router.register_run(session_id, run_id).await {
481 Ok(registration) => registration,
482 Err(error) => {
483 let collision = Err(AgentError::LLM(error.to_string()));
484 finalize_rejected_runner_if_distinct(
485 runners,
486 session_id,
487 error.existing_run_id(),
488 run_id,
489 &collision,
490 )
491 .await;
492 return Err(error);
493 }
494 };
495 if selected_input {
496 registration.restrict_to_selected_input();
497 }
498 let cleanup_runners = runners.clone();
499 let cleanup_session_id = session_id.to_string();
500 let cleanup_run_id = run_id.to_string();
501 registration.set_abort_cleanup(move || async move {
502 let abandoned = Err(AgentError::Cancelled);
503 finalize_runner_exact(
504 &cleanup_runners,
505 &cleanup_session_id,
506 &cleanup_run_id,
507 &abandoned,
508 )
509 .await;
510 });
511 Ok(registration)
512}
513
514async fn cleanup_execution_reservation(
515 activation: SessionExecutionActivationOwnership,
516 runners: Arc<RwLock<HashMap<String, AgentRunner>>>,
517 session_id: String,
518 run_id: String,
519 selected_input: bool,
520) {
521 match activation {
522 SessionExecutionActivationOwnership::Registered(registration) => {
523 registration.abandon().await;
524 }
525 SessionExecutionActivationOwnership::RegistrationPending(router) => {
526 match register_reserved_activation(
527 router,
528 &runners,
529 &session_id,
530 &run_id,
531 selected_input,
532 )
533 .await
534 {
535 Ok(registration) => registration.abandon().await,
536 Err(error) => {
537 tracing::debug!(
538 %session_id,
539 %run_id,
540 existing_run_id = %error.existing_run_id(),
541 "abandoned activation placeholder was already adopted or superseded"
542 );
543 }
544 }
545 }
546 SessionExecutionActivationOwnership::UnpublishedActivation(_) => {
547 remove_runner_exact(&runners, &session_id, &run_id).await;
548 }
549 SessionExecutionActivationOwnership::Unrouted => {
550 let abandoned = Err(AgentError::Cancelled);
551 finalize_runner_exact(&runners, &session_id, &run_id, &abandoned).await;
552 }
553 }
554}
555
556async fn remove_runner_exact(
557 runners: &Arc<RwLock<HashMap<String, AgentRunner>>>,
558 session_id: &str,
559 run_id: &str,
560) -> bool {
561 let mut runners = runners.write().await;
562 if runners
563 .get(session_id)
564 .is_some_and(|runner| runner.run_id == run_id)
565 {
566 super::runner_lifecycle::remove_runner_entry(&mut runners, session_id).await;
567 true
568 } else {
569 false
570 }
571}
572
573pub async fn reserve_session_execution(
581 agent: &Arc<Agent>,
582 runners: &Arc<RwLock<HashMap<String, AgentRunner>>>,
583 senders: &Arc<RwLock<HashMap<String, broadcast::Sender<AgentEvent>>>>,
584 session_id: &str,
585 event_sender: &broadcast::Sender<AgentEvent>,
586) -> SessionExecutionReserveOutcome {
587 reserve_session_execution_inner(agent, runners, senders, session_id, event_sender, false).await
588}
589
590pub async fn reserve_selected_input_execution(
593 agent: &Arc<Agent>,
594 runners: &Arc<RwLock<HashMap<String, AgentRunner>>>,
595 senders: &Arc<RwLock<HashMap<String, broadcast::Sender<AgentEvent>>>>,
596 session_id: &str,
597 event_sender: &broadcast::Sender<AgentEvent>,
598) -> SessionExecutionReserveOutcome {
599 reserve_session_execution_inner(agent, runners, senders, session_id, event_sender, true).await
600}
601
602async fn reserve_session_execution_inner(
603 agent: &Arc<Agent>,
604 runners: &Arc<RwLock<HashMap<String, AgentRunner>>>,
605 senders: &Arc<RwLock<HashMap<String, broadcast::Sender<AgentEvent>>>>,
606 session_id: &str,
607 event_sender: &broadcast::Sender<AgentEvent>,
608 selected_input: bool,
609) -> SessionExecutionReserveOutcome {
610 let reservation = match reserve_runner_core(runners, senders, session_id, event_sender).await {
611 ReserveOutcome::AlreadyRunning(run_id) => {
612 return SessionExecutionReserveOutcome::AlreadyRunning { run_id };
613 }
614 ReserveOutcome::Reserved(reservation) => reservation,
615 };
616
617 let mut execution_reservation = SessionExecutionReservation {
623 session_id: session_id.to_string(),
624 run_id: reservation.run_id,
625 cancel_token: reservation.cancel_token,
626 runners: runners.clone(),
627 activation: SessionExecutionActivationOwnership::Unrouted,
628 root_actor: None,
629 execution_persistence: None,
630 persistence_bound: false,
631 selected_input,
632 armed: true,
633 };
634
635 if let Some(router) = agent.activation_router().cloned() {
636 execution_reservation.activation =
637 SessionExecutionActivationOwnership::RegistrationPending(router);
638 if let Err(error) = execution_reservation.ensure_registered().await {
639 tracing::warn!(
640 %session_id,
641 attempted_run_id = %execution_reservation.run_id(),
642 existing_run_id = %error.existing_run_id(),
643 %error,
644 "runner reservation collided with an existing logical-session owner"
645 );
646 return SessionExecutionReserveOutcome::AlreadyRunning {
647 run_id: error.existing_run_id().to_string(),
648 };
649 }
650 }
651 SessionExecutionReserveOutcome::Reserved(execution_reservation)
652}
653
654pub fn read_cached_session(cache: &SessionCache, id: &str) -> Option<bamboo_agent_core::Session> {
662 cache
663 .get(id)
664 .map(|e| e.value().clone())
665 .map(|a| a.read().clone())
666}
667
668const SKILL_CONTEXT_START_MARKER: &str = "<!-- BAMBOO_SKILL_CONTEXT_START -->";
669const TOOL_GUIDE_START_MARKER: &str = "<!-- BAMBOO_TOOL_GUIDE_START -->";
670const EXTERNAL_MEMORY_START_MARKER: &str = "<!-- BAMBOO_EXTERNAL_MEMORY_START -->";
671const TASK_LIST_START_MARKER: &str = "<!-- BAMBOO_TASK_LIST_START -->";
672
673pub struct SessionExecutionOutcome {
680 pub success: bool,
682 pub cancelled: bool,
684 pub error: Option<String>,
686}
687
688impl SessionExecutionOutcome {
689 fn from_result(result: &Result<(), AgentError>) -> Self {
690 match result {
691 Ok(()) => Self {
692 success: true,
693 cancelled: false,
694 error: None,
695 },
696 Err(error) => Self {
697 success: false,
698 cancelled: error.is_cancelled(),
699 error: Some(error.to_string()),
700 },
701 }
702 }
703}
704
705pub type SessionCompletionHook = Box<
715 dyn for<'a> FnOnce(
716 SessionExecutionOutcome,
717 &'a mut Session,
718 ) -> Pin<Box<dyn Future<Output = ()> + Send + 'a>>
719 + Send,
720>;
721
722pub struct SessionExecutionArgs {
728 pub agent: Arc<Agent>,
730 pub session_id: String,
731 pub session: Session,
732 pub execution_reservation: SessionExecutionReservation,
735
736 pub tools_override: Option<Arc<dyn ToolExecutor>>,
738 pub provider_override: Option<Arc<dyn LLMProvider>>,
739 pub model_roster: ModelRoster,
743 pub reasoning_effort: Option<ReasoningEffort>,
744 pub reasoning_effort_source: String,
745 pub auxiliary_model_resolver:
746 Option<Arc<dyn Fn() -> crate::runtime::config::AuxiliaryModelConfig + Send + Sync>>,
747 pub disabled_filter_resolver: Option<DisabledFilterResolver>,
751 pub disabled_tools: Option<BTreeSet<String>>,
752 pub disabled_skill_ids: Option<BTreeSet<String>>,
753 pub selected_skill_ids: Option<Vec<String>>,
754 pub selected_skill_mode: Option<String>,
755 pub mpsc_tx: mpsc::Sender<AgentEvent>,
756 pub history_commit_barrier: HistoryCommitBarrier,
759 pub image_fallback: Option<ImageFallbackConfig>,
760 pub gold_config: Option<GoldConfig>,
761 pub guardian_config: Option<GuardianConfig>,
763 pub guardian_spawner: Option<Arc<dyn GuardianSpawner>>,
766 pub bash_resume_hook: Option<Arc<dyn BashResumeHook>>,
768 pub bash_completion_sink: Option<Arc<dyn BashCompletionSink>>,
770 pub app_data_dir: Option<std::path::PathBuf>,
771 pub run_budget: Option<bamboo_config::RunBudgetConfig>,
774
775 pub runners: Arc<RwLock<HashMap<String, AgentRunner>>>,
777 pub sessions_cache: SessionCache,
778
779 pub on_complete: Option<SessionCompletionHook>,
782
783 pub child_completion_handler: Option<Arc<dyn super::ChildCompletionHandler>>,
795}
796
797struct ExecuteRequestParams {
807 tools: Option<Arc<dyn ToolExecutor>>,
808 provider_override: Option<Arc<dyn LLMProvider>>,
809 model_roster: ModelRoster,
810 reasoning_effort: Option<ReasoningEffort>,
811 auxiliary_model_resolver: Option<Arc<dyn Fn() -> AuxiliaryModelConfig + Send + Sync>>,
812 disabled_filter_resolver: Option<DisabledFilterResolver>,
813 disabled_tools: Option<BTreeSet<String>>,
814 disabled_skill_ids: Option<BTreeSet<String>>,
815 selected_skill_ids: Option<Vec<String>>,
816 selected_skill_mode: Option<String>,
817 image_fallback: Option<ImageFallbackConfig>,
818 gold_config: Option<GoldConfig>,
819 guardian_config: Option<GuardianConfig>,
820 guardian_spawner: Option<Arc<dyn GuardianSpawner>>,
821 bash_resume_hook: Option<Arc<dyn BashResumeHook>>,
822 bash_completion_sink: Option<Arc<dyn BashCompletionSink>>,
823 app_data_dir: Option<std::path::PathBuf>,
824 run_budget: Option<bamboo_config::RunBudgetConfig>,
825}
826
827fn build_execute_request(
834 initial_message: String,
835 event_tx: mpsc::Sender<AgentEvent>,
836 cancel_token: CancellationToken,
837 params: ExecuteRequestParams,
838) -> ExecuteRequest {
839 let ExecuteRequestParams {
840 tools,
841 provider_override,
842 model_roster,
843 reasoning_effort,
844 auxiliary_model_resolver,
845 disabled_filter_resolver,
846 disabled_tools,
847 disabled_skill_ids,
848 selected_skill_ids,
849 selected_skill_mode,
850 image_fallback,
851 gold_config,
852 guardian_config,
853 guardian_spawner,
854 bash_resume_hook,
855 bash_completion_sink,
856 app_data_dir,
857 run_budget,
858 } = params;
859
860 let mut builder = ExecuteRequestBuilder::new(initial_message, event_tx, cancel_token)
861 .model_roster(model_roster)
862 .gold_config(gold_config)
863 .guardian_config(guardian_config)
864 .guardian_spawner(guardian_spawner)
865 .bash_resume_hook(bash_resume_hook)
866 .bash_completion_sink(bash_completion_sink);
867
868 if let Some(run_budget) = run_budget {
869 builder = builder.run_budget(run_budget);
870 }
871
872 if let Some(tools) = tools {
873 builder = builder.tools(tools);
874 }
875 if let Some(provider_override) = provider_override {
876 builder = builder.provider_override(provider_override);
877 }
878 if let Some(reasoning_effort) = reasoning_effort {
879 builder = builder.reasoning_effort(reasoning_effort);
880 }
881 if let Some(disabled_filter_resolver) = disabled_filter_resolver {
882 builder = builder.disabled_filter_resolver(disabled_filter_resolver);
883 }
884 if let Some(auxiliary_model_resolver) = auxiliary_model_resolver {
885 builder = builder.auxiliary_model_resolver(auxiliary_model_resolver);
886 }
887 if let Some(disabled_tools) = disabled_tools {
888 builder = builder.disabled_tools(disabled_tools);
889 }
890 if let Some(disabled_skill_ids) = disabled_skill_ids {
891 builder = builder.disabled_skill_ids(disabled_skill_ids);
892 }
893 if let Some(selected_skill_ids) = selected_skill_ids {
894 builder = builder.selected_skill_ids(selected_skill_ids);
895 }
896 if let Some(selected_skill_mode) = selected_skill_mode {
897 builder = builder.selected_skill_mode(selected_skill_mode);
898 }
899 if let Some(image_fallback) = image_fallback {
900 builder = builder.image_fallback(image_fallback);
901 }
902 if let Some(app_data_dir) = app_data_dir {
903 builder = builder.app_data_dir(app_data_dir);
904 }
905
906 builder.build()
907}
908
909pub fn spawn_session_execution(args: SessionExecutionArgs) {
918 spawn_session_execution_with_inputs(args, None);
919}
920
921pub fn spawn_session_execution_with_inputs(
924 args: SessionExecutionArgs,
925 inputs: Option<crate::runtime::config::UntrustedExecutionInputs>,
926) {
927 let span_session_id = args.session_id.clone();
928 let session_span = tracing::info_span!("agent_execution", session_id = %span_session_id);
929
930 tokio::spawn(
931 async move {
932 let SessionExecutionArgs {
933 agent,
934 session_id,
935 mut session,
936 mut execution_reservation,
937 tools_override,
938 provider_override,
939 model_roster,
940 reasoning_effort,
941 reasoning_effort_source,
942 auxiliary_model_resolver,
943 disabled_filter_resolver,
944 disabled_tools,
945 disabled_skill_ids,
946 selected_skill_ids,
947 selected_skill_mode,
948 mpsc_tx,
949 mut history_commit_barrier,
950 image_fallback,
951 gold_config,
952 guardian_config,
953 guardian_spawner,
954 bash_resume_hook,
955 bash_completion_sink,
956 app_data_dir,
957 run_budget,
958 runners,
959 sessions_cache,
960 on_complete,
961 child_completion_handler,
962 } = args;
963
964 if !execution_reservation.matches_execution_target(&session_id, &session.id, &runners) {
965 tracing::error!(
966 %session_id,
967 domain_session_id = %session.id,
968 reservation_session_id = %execution_reservation.session_id(),
969 run_id = %execution_reservation.run_id(),
970 same_runner_registry = Arc::ptr_eq(&execution_reservation.runners, &runners),
971 "refusing mismatched session execution reservation"
972 );
973 execution_reservation.abandon().await;
974 return;
975 }
976 if let Err(error) = execution_reservation.ensure_registered().await {
977 tracing::warn!(
978 %session_id,
979 run_id = %execution_reservation.run_id(),
980 %error,
981 "session execution could not adopt its router activation owner"
982 );
983 return;
984 }
985 if execution_reservation.root_actor.is_none()
986 && agent.persistence().root_actor_execution_required(&session)
987 {
988 let message =
989 "Root execution route requires a bound Actor writer and event handoff";
990 tracing::warn!(%session_id, "{message}");
991 let _ = mpsc_tx
992 .send(AgentEvent::Error {
993 message: message.into(),
994 })
995 .await;
996 execution_reservation.abandon().await;
997 return;
998 }
999 if let Err(error) = execution_reservation.bind_execution_persistence(&agent, &session) {
1002 let _ = mpsc_tx
1003 .send(AgentEvent::Error {
1004 message: error.to_string(),
1005 })
1006 .await;
1007 execution_reservation.abandon().await;
1008 return;
1009 }
1010 let ownership = match &execution_reservation.activation {
1011 SessionExecutionActivationOwnership::Registered(run) => Some(run.ownership()),
1012 _ => None,
1013 };
1014 let admission_persistence = execution_reservation
1015 .execution_persistence()
1016 .unwrap_or_else(|| agent.persistence().clone());
1017 if let Err(error) = crate::runtime::counted_loop::admit_registered(
1018 &mut session,
1019 ownership.as_ref(),
1020 agent.activation_router(),
1021 admission_persistence.as_ref(),
1022 )
1023 .await
1024 {
1025 let _ = mpsc_tx
1026 .send(AgentEvent::Error {
1027 message: error.to_string(),
1028 })
1029 .await;
1030 execution_reservation.abandon().await;
1031 return;
1032 }
1033 let mut root_actor = execution_reservation.root_actor.take().map(|owner| *owner);
1034 let root_actor_bound = root_actor.is_some();
1035 let persistence = execution_reservation
1036 .execution_persistence()
1037 .unwrap_or_else(|| agent.persistence().clone());
1038 let agent = Arc::new(agent.with_execution_persistence(persistence));
1039 let activation_run_id = execution_reservation.run_id().to_string();
1040 let prior_message_ids = session.messages.iter().map(|m| m.id.clone()).collect();
1041 let (cancel_token, mut activation_registration) =
1042 execution_reservation.disarm_for_execution();
1043
1044 let model = model_roster.model.clone().unwrap_or_default();
1048
1049 let initial_message = initial_user_message_for_session(&session);
1050 let selected_skill_ids =
1051 selected_skill_ids.or_else(|| selected_skill_ids_for_session(&session));
1052 let selected_skill_mode =
1053 selected_skill_mode.or_else(|| selected_skill_mode_for_session(&session));
1054
1055 tracing::info!(
1056 "[{}] Using resolved session model: {}, reasoning_effort={}, reasoning_source={}",
1057 session_id,
1058 model,
1059 reasoning_effort
1060 .map(ReasoningEffort::as_str)
1061 .unwrap_or("none"),
1062 reasoning_effort_source
1063 );
1064
1065 crate::session_app::execution_prep::prepare_session_for_execution(
1072 &mut session,
1073 None,
1074 Some(&model),
1075 );
1076
1077 let system_prompt = system_prompt_for_debug_log(&mut session);
1078 if let Some(prompt) = system_prompt.as_ref() {
1079 log_base_system_prompt_snapshot(&session_id, prompt);
1080 }
1081
1082 let inherited_child_wait = agent.persistence().inherited_child_wait();
1086 let execute_request = build_execute_request(
1087 initial_message,
1088 mpsc_tx.clone(),
1089 cancel_token,
1090 ExecuteRequestParams {
1091 tools: tools_override,
1092 provider_override,
1093 model_roster,
1094 reasoning_effort,
1095 auxiliary_model_resolver,
1096 disabled_filter_resolver,
1097 disabled_tools,
1098 disabled_skill_ids,
1099 selected_skill_ids,
1100 selected_skill_mode,
1101 image_fallback,
1102 gold_config,
1103 guardian_config,
1104 guardian_spawner,
1105 bash_resume_hook,
1106 bash_completion_sink,
1107 app_data_dir,
1108 run_budget,
1109 },
1110 );
1111
1112 let result = {
1119 use futures::FutureExt;
1120 match std::panic::AssertUnwindSafe(agent.execute_with_inputs(
1121 &mut session,
1122 execute_request,
1123 inputs,
1124 ))
1125 .catch_unwind()
1126 .await
1127 {
1128 Ok(result) => result,
1129 Err(panic) => {
1130 let message = panic
1131 .downcast_ref::<&str>()
1132 .map(|s| (*s).to_string())
1133 .or_else(|| panic.downcast_ref::<String>().cloned())
1134 .unwrap_or_else(|| "non-string panic payload".to_string());
1135 tracing::error!(
1136 "[{}] agent execution panicked; finalizing as terminal error: {}",
1137 session_id,
1138 message
1139 );
1140 Err(AgentError::LLM(format!(
1141 "agent execution panicked: {message}"
1142 )))
1143 }
1144 }
1145 };
1146
1147 if let Some(error_event) = terminal_error_event_for_result(&result) {
1149 let _ = mpsc_tx.send(error_event).await;
1150 }
1151
1152 let suspended_non_terminal = result.is_ok()
1169 && session
1170 .metadata
1171 .get("runtime.suspend_reason")
1172 .is_some_and(|reason| !reason.trim().is_empty());
1173 match &result {
1174 Ok(()) if suspended_non_terminal => {
1175 session.set_last_run_status("suspended");
1176 session.clear_last_run_error();
1177 }
1178 Ok(()) => {
1179 session.set_last_run_status("completed");
1180 session.clear_last_run_error();
1181 }
1182 Err(error) if error.is_cancelled() => {
1183 session.set_last_run_status("cancelled");
1184 session.set_last_run_error(error.to_string());
1185 }
1186 Err(error) => {
1187 session.set_last_run_status("error");
1188 session.set_last_run_error(error.to_string());
1189 }
1190 }
1191
1192 if let Some(on_complete) = on_complete {
1196 on_complete(SessionExecutionOutcome::from_result(&result), &mut session).await;
1197 }
1198
1199 let executed_admitted_generation = session
1203 .session_inbox_admission()
1204 .map_or(0, |state| state.last_admitted_sequence);
1205 let legacy_migration =
1206 crate::runtime::runner::state_bridge::migrate_legacy_pending_only(
1207 &mut session,
1208 Some(agent.storage()),
1209 Some(agent.persistence()),
1210 agent.session_inbox(),
1211 )
1212 .await;
1213 let pending_boundary_generation = session
1214 .session_inbox_admission()
1215 .and_then(|state| state.pending_activation_generation());
1216 let pending_generation = match (
1217 pending_boundary_generation,
1218 legacy_migration.highest_generation,
1219 ) {
1220 (Some(left), Some(right)) => Some(left.max(right)),
1221 (left, right) => left.or(right),
1222 };
1223 if let (Some(generation), Some(router)) =
1224 (pending_generation, agent.activation_router())
1225 {
1226 let activation_ready = if let Some(inbox) = agent.session_inbox() {
1227 match inbox
1228 .mark_activation_eligible(
1229 &session_id,
1230 generation,
1231 bamboo_domain::SessionActivationPolicy::InterruptSpecificWait,
1232 )
1233 .await
1234 {
1235 Ok(()) => true,
1236 Err(error) => {
1237 tracing::error!(
1238 session_id = %session_id,
1239 %error,
1240 "failed to persist unadmitted SessionInbox activation watermark"
1241 );
1242 false
1243 }
1244 }
1245 } else {
1246 false
1247 };
1248 if activation_ready {
1249 if let Err(error) = bamboo_domain::SessionActivationPort::request_activation(
1250 router.as_ref(),
1251 &session_id,
1252 generation,
1253 )
1254 .await
1255 {
1256 tracing::error!(
1257 session_id = %session_id,
1258 %error,
1259 "failed to hand unadmitted SessionInbox generation to activation router"
1260 );
1261 }
1262 }
1263 }
1264
1265 if let Some(registration) = activation_registration.as_mut() {
1269 registration.begin_finalization().await;
1270 }
1271
1272 super::child_completion::ChildCompletionSource::prepare(
1276 &mut session,
1277 &activation_run_id,
1278 &prior_message_ids,
1279 );
1280 let saved = save_finalized_runtime_with_inherited_child_wait(
1281 agent.persistence().as_ref(),
1282 &mut session,
1283 inherited_child_wait
1284 .as_ref()
1285 .map(|inherited| inherited.wait()),
1286 )
1287 .await;
1288 let history_committed = saved.is_ok();
1289 let completion_source =
1290 super::child_completion::ChildCompletionSource::after_final_save(
1291 &session,
1292 history_committed,
1293 );
1294 let completed_at = chrono::Utc::now();
1295 if let Err(error) = &saved {
1296 tracing::warn!("[{}] Failed to save session: {}", session_id, error);
1297 }
1298
1299 if history_committed
1305 && !history_commit_barrier
1306 .send_and_wait(&mpsc_tx, session_id.clone())
1307 .await
1308 {
1309 tracing::warn!(
1310 session_id = %session_id,
1311 "history commit barrier could not be published before runner finalization"
1312 );
1313 }
1314
1315 finalize_runner_exact(&runners, &session_id, &activation_run_id, &result).await;
1321
1322 if let Some(owner) = root_actor.take() {
1323 let outcome = if result.as_ref().is_err_and(|error| error.is_cancelled()) {
1324 bamboo_domain::ActorActivationFinish::Cancelled
1325 } else if result.is_err() || !history_committed {
1326 bamboo_domain::ActorActivationFinish::Failed
1327 } else {
1328 bamboo_domain::ActorActivationFinish::Succeeded
1329 };
1330 owner.finish(outcome).await;
1331 }
1332
1333 let finalization = if let Some(registration) = activation_registration.take() {
1334 registration.finish(executed_admitted_generation).await
1335 } else {
1336 Ok(None)
1337 };
1338 if let Err(error) = finalization {
1339 tracing::error!(
1343 session_id = %session_id,
1344 %error,
1345 "failed to activate successor for finalization-racing SessionInbox delivery"
1346 );
1347 }
1348
1349 let child_completion = child_completion_handler.filter(|_| {
1358 session.kind == bamboo_agent_core::SessionKind::Child
1359 && session.parent_session_id.is_some()
1360 });
1361 let parent_session_id = session.parent_session_id.clone();
1362 let child_status = session.last_run_status();
1363 let child_error = session.last_run_error();
1364
1365 publish_final_session_cache(&sessions_cache, session, root_actor_bound, &saved);
1366
1367 if let (Some(handler), Some(parent_session_id), Some(status)) =
1368 (child_completion, parent_session_id, child_status)
1369 {
1370 use futures::FutureExt;
1371 let completion = ChildCompletion {
1372 parent_session_id: parent_session_id.clone(),
1373 child_session_id: session_id.clone(),
1374 status,
1375 error: child_error,
1376 completed_at,
1377 source: completion_source,
1378 };
1379 if std::panic::AssertUnwindSafe(handler.on_child_completed(completion))
1380 .catch_unwind()
1381 .await
1382 .is_err()
1383 {
1384 tracing::error!(
1385 %parent_session_id,
1386 child_session_id = %session_id,
1387 "child completion handler panicked on resumed-child terminal"
1388 );
1389 }
1390 }
1391
1392 tracing::info!("[{}] Agent execution completed", session_id);
1393 }
1394 .instrument(session_span),
1395 );
1396}
1397
1398async fn save_finalized_runtime_with_inherited_child_wait(
1400 persistence: &dyn bamboo_domain::RuntimeSessionPersistence,
1401 session: &mut Session,
1402 inherited_child_wait: Option<&bamboo_domain::session::runtime_state::WaitingForChildrenState>,
1403) -> std::io::Result<()> {
1404 if let Some(inherited) = inherited_child_wait.filter(|wait| {
1405 session
1406 .agent_runtime_state
1407 .as_ref()
1408 .and_then(|runtime| runtime.waiting_for_children.as_ref())
1409 == Some(*wait)
1410 }) {
1411 persistence
1412 .save_finalized_runtime_with_inherited_child_wait(session, inherited)
1413 .await
1414 } else {
1415 persistence.save_finalized_runtime_session(session).await
1416 }
1417}
1418
1419fn publish_final_session_cache(
1420 cache: &SessionCache,
1421 session: Session,
1422 root_actor_bound: bool,
1423 saved: &std::io::Result<()>,
1424) {
1425 let authority_conflict = saved.as_ref().err().is_some_and(|error| {
1426 error
1427 .get_ref()
1428 .is_some_and(|cause| cause.is::<bamboo_domain::SessionAuthorityConflict>())
1429 });
1430 if !authority_conflict && !root_actor_bound {
1433 cache.insert(
1434 session.id.clone(),
1435 Arc::new(crate::SessionSnapshot::new(session)),
1436 );
1437 }
1438}
1439
1440pub fn log_base_system_prompt_snapshot(session_id: &str, prompt: &str) {
1442 tracing::info!(
1443 "[{}] Base system prompt snapshot: len={} chars, has_skill={}, has_tool_guide={}, has_external_memory={}, has_task_list={}",
1444 session_id,
1445 prompt.len(),
1446 prompt.contains(SKILL_CONTEXT_START_MARKER),
1447 prompt.contains(TOOL_GUIDE_START_MARKER),
1448 prompt.contains(EXTERNAL_MEMORY_START_MARKER),
1449 prompt.contains(TASK_LIST_START_MARKER),
1450 );
1451
1452 tracing::debug!(
1453 "[{}] ========== BASE SYSTEM PROMPT SNAPSHOT ==========",
1454 session_id
1455 );
1456 tracing::debug!("[{}] Snapshot length: {} chars", session_id, prompt.len());
1457 tracing::debug!("[{}] -----------------------------------", session_id);
1458 tracing::debug!("[{}] {}", session_id, prompt);
1459 tracing::debug!(
1460 "[{}] ========== END BASE SYSTEM PROMPT SNAPSHOT ==========",
1461 session_id
1462 );
1463}
1464
1465pub fn terminal_error_event_for_result(result: &Result<(), AgentError>) -> Option<AgentEvent> {
1467 match result {
1468 Ok(_) => None,
1469 Err(error) if error.is_cancelled() => Some(AgentEvent::Error {
1470 message: "Agent execution cancelled by user".to_string(),
1471 }),
1472 Err(error) => Some(AgentEvent::Error {
1473 message: error.to_string(),
1474 }),
1475 }
1476}
1477
1478fn system_prompt_for_session(session: &Session) -> Option<String> {
1481 session
1482 .messages
1483 .iter()
1484 .find(|message| matches!(message.role, bamboo_agent_core::Role::System))
1485 .map(|message| message.content.clone())
1486}
1487
1488fn system_prompt_for_debug_log(session: &mut Session) -> Option<String> {
1492 crate::runtime::runner::session_setup::migrate_legacy_workspace_prompt(session);
1493 system_prompt_for_session(session)
1494}
1495
1496fn initial_user_message_for_session(session: &Session) -> String {
1497 session
1498 .messages
1499 .last()
1500 .filter(|message| matches!(message.role, bamboo_agent_core::Role::User))
1501 .map(|message| message.content.clone())
1502 .unwrap_or_default()
1503}
1504
1505fn selected_skill_ids_for_session(session: &Session) -> Option<Vec<String>> {
1506 session
1507 .metadata
1508 .get("selected_skill_ids")
1509 .and_then(|raw| bamboo_skills::selection::parse_selected_skill_ids_metadata(raw))
1510}
1511
1512fn selected_skill_mode_for_session(session: &Session) -> Option<String> {
1513 let value = session
1514 .metadata
1515 .get("skill_mode")
1516 .or_else(|| session.metadata.get("mode"))?;
1517 let trimmed = value.trim();
1518 if trimmed.is_empty() {
1519 None
1520 } else {
1521 Some(trimmed.to_string())
1522 }
1523}
1524
1525pub(super) struct RootActorExecution {
1526 binding: bamboo_domain::RootActorExecutionBinding,
1527 stop_renewal: CancellationToken,
1528 renewal: Option<tokio::task::JoinHandle<()>>,
1529 response_events: Vec<AgentEvent>,
1530}
1531
1532impl RootActorExecution {
1533 pub(super) fn start(binding: bamboo_domain::RootActorExecutionBinding) -> Self {
1534 let stop_renewal = CancellationToken::new();
1535 let stop = stop_renewal.clone();
1536 let directory = binding.directory.clone();
1537 let fence = binding.owner.fence.clone();
1538 let duration = binding.lease_duration;
1539 let renewal = tokio::spawn(async move {
1540 let mut interval = tokio::time::interval(duration / 3);
1541 interval.tick().await;
1542 loop {
1543 tokio::select! {
1544 _ = stop.cancelled() => break,
1545 _ = interval.tick() => {
1546 let now = chrono::Utc::now();
1547 let expires = now + chrono::Duration::from_std(duration).expect("bounded Root lease");
1548 if let Err(error) = directory.renew_activation(&fence, now, expires).await {
1549 tracing::warn!(actor_id = %fence.actor_id, run_id = %fence.run_id,
1553 %error, "Root Actor renewal lost its execution authority");
1554 break;
1555 }
1556 }
1557 }
1558 }
1559 });
1560 Self {
1561 binding,
1562 stop_renewal,
1563 renewal: Some(renewal),
1564 response_events: Vec::new(),
1565 }
1566 }
1567
1568 pub(super) fn persistence(
1569 &self,
1570 ) -> std::sync::Arc<dyn bamboo_domain::RuntimeSessionPersistence> {
1571 self.binding.persistence.clone()
1572 }
1573
1574 pub(super) async fn finish(mut self, outcome: bamboo_domain::ActorActivationFinish) {
1575 self.stop_renewal.cancel();
1576 if let Some(renewal) = self.renewal.take() {
1577 let _ = renewal.await;
1578 }
1579 let directory = self.binding.directory.clone();
1580 let fence = self.binding.owner.fence.clone();
1581 let finish = tokio::spawn(async move {
1582 if let Err(error) = directory
1583 .finish_activation(&fence, chrono::Utc::now(), outcome)
1584 .await
1585 {
1586 tracing::warn!(actor_id = %fence.actor_id, run_id = %fence.run_id,
1587 %error, "Root Actor finish rejected its obsolete execution fence");
1588 }
1589 });
1590 self.binding.disarm_abandonment();
1591 let _ = finish.await;
1592 }
1593}
1594
1595impl Drop for RootActorExecution {
1596 fn drop(&mut self) {
1597 self.stop_renewal.cancel();
1598 }
1601}
1602
1603#[cfg(test)]
1604#[path = "agent_spawn_peer_privacy_tests.rs"]
1605mod peer_privacy_tests;
1606
1607#[cfg(test)]
1608mod reservation_tests {
1609 use super::*;
1610 use crate::runtime::execution::runner_state::AgentStatus;
1611
1612 struct CompletionAtFinalSave {
1613 directory: std::path::PathBuf,
1614 persistence: bamboo_storage::LockedSessionStore,
1615 cleared: std::sync::atomic::AtomicBool,
1616 }
1617
1618 impl CompletionAtFinalSave {
1619 async fn clear_wait(&self, id: &str) {
1620 let mut process = std::process::Command::new(std::env::current_exe().unwrap())
1621 .args(["--exact", "runtime::execution::agent_spawn::reservation_tests::inherited_child_wait_adopts_independent_completion_at_final_save", "--nocapture"])
1622 .env("BAMBOO_FINAL_WAIT_COMPLETION_DIR", &self.directory)
1623 .env("BAMBOO_FINAL_WAIT_COMPLETION_ID", id)
1624 .spawn().unwrap();
1625 let deadline = std::time::Instant::now() + std::time::Duration::from_secs(15);
1626 loop {
1627 if let Some(status) = process.try_wait().unwrap() {
1628 assert!(status.success(), "completion process failed");
1629 break;
1630 }
1631 if std::time::Instant::now() >= deadline {
1632 let _ = process.kill();
1633 let _ = process.wait();
1634 panic!("bounded completion process deadline");
1635 }
1636 std::thread::sleep(std::time::Duration::from_millis(5));
1637 }
1638 self.cleared
1639 .store(true, std::sync::atomic::Ordering::SeqCst);
1640 }
1641 }
1642
1643 #[async_trait::async_trait]
1644 impl bamboo_domain::RuntimeSessionPersistence for CompletionAtFinalSave {
1645 async fn save_runtime_session(&self, session: &mut Session) -> std::io::Result<()> {
1646 self.clear_wait(&session.id).await;
1647 self.persistence.merge_save_finalized_runtime(session).await
1648 }
1649 async fn save_finalized_runtime_with_inherited_child_wait(
1650 &self,
1651 session: &mut Session,
1652 inherited: &bamboo_domain::session::runtime_state::WaitingForChildrenState,
1653 ) -> std::io::Result<()> {
1654 self.clear_wait(&session.id).await;
1655 self.persistence
1656 .merge_save_inherited_child_wait_and_publish(session, inherited, |_, _| {})
1657 .await
1658 }
1659 }
1660
1661 #[tokio::test]
1662 async fn inherited_child_wait_adopts_independent_completion_at_final_save() {
1663 use bamboo_agent_core::storage::Storage;
1664 use bamboo_domain::session::runtime_state::{
1665 AgentRuntimeState, AgentStatusState, ChildWaitPolicy, WaitingForChildrenState,
1666 };
1667 if let Some(directory) = std::env::var_os("BAMBOO_FINAL_WAIT_COMPLETION_DIR") {
1668 let remote = bamboo_storage::SessionStoreV2::new(directory.into())
1669 .await
1670 .unwrap();
1671 let id = std::env::var("BAMBOO_FINAL_WAIT_COMPLETION_ID").unwrap();
1672 let mut latest = remote.load_session(&id).await.unwrap().unwrap();
1673 let runtime = latest.agent_runtime_state.as_mut().unwrap();
1674 runtime.waiting_for_children = None;
1675 runtime.status = AgentStatusState::Idle;
1676 runtime.suspension = None;
1677 latest.metadata.remove("runtime.suspend_reason");
1678 remote.save_runtime_state(&latest).await.unwrap();
1679 return;
1680 }
1681 let temp = tempfile::tempdir().unwrap();
1682 let local = Arc::new(
1683 bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
1684 .await
1685 .unwrap(),
1686 );
1687 let mut runner = Session::new("cross-host-final-wait", "model");
1691 let inherited = WaitingForChildrenState::for_children(
1692 vec!["child".into()],
1693 ChildWaitPolicy::All,
1694 chrono::Utc::now(),
1695 );
1696 let mut runtime = AgentRuntimeState::new("run");
1697 runtime.waiting_for_children = Some(inherited.clone());
1698 runtime.status = AgentStatusState::Suspended;
1699 runner.agent_runtime_state = Some(runtime);
1700 runner.metadata.insert(
1701 "runtime.suspend_reason".into(),
1702 "waiting_for_children".into(),
1703 );
1704 runner.set_last_run_status("suspended");
1705 local.save_session(&runner).await.unwrap();
1706 let persistence = CompletionAtFinalSave {
1707 directory: temp.path().to_path_buf(),
1708 persistence: bamboo_storage::LockedSessionStore::new(local.clone()),
1709 cleared: std::sync::atomic::AtomicBool::new(false),
1710 };
1711 save_finalized_runtime_with_inherited_child_wait(
1712 &persistence,
1713 &mut runner,
1714 Some(&inherited),
1715 )
1716 .await
1717 .unwrap();
1718 assert!(persistence
1719 .cleared
1720 .load(std::sync::atomic::Ordering::SeqCst));
1721 let durable = local.load_session(&runner.id).await.unwrap().unwrap();
1722 for session in [&runner, &durable] {
1723 let runtime = session.agent_runtime_state.as_ref().unwrap();
1724 assert!(
1725 runtime.waiting_for_children.is_none(),
1726 "independent completion must not be resurrected"
1727 );
1728 assert_eq!(runtime.status, AgentStatusState::Idle);
1729 assert!(runtime.suspension.is_none());
1730 assert!(!session.metadata.contains_key("runtime.suspend_reason"));
1731 assert_eq!(session.last_run_status().as_deref(), Some("completed"));
1732 }
1733 }
1734
1735 #[tokio::test]
1736 async fn inherited_child_wait_rejects_aba_and_missing_birth_without_cache_publication() {
1737 use bamboo_agent_core::storage::Storage;
1738 use bamboo_domain::session::runtime_state::{
1739 AgentRuntimeState, ChildWaitPolicy, WaitingForChildrenState,
1740 };
1741 use bamboo_storage::{LockedSessionStore, SessionStoreV2};
1742
1743 for missing in [false, true] {
1744 let temp = tempfile::tempdir().unwrap();
1745 let storage = Arc::new(
1746 SessionStoreV2::new(temp.path().to_path_buf())
1747 .await
1748 .unwrap(),
1749 );
1750 let persistence = LockedSessionStore::new(storage.clone());
1751 let cache: SessionCache = Arc::default();
1752 storage
1753 .save_session(&Session::new("root", "model"))
1754 .await
1755 .unwrap();
1756 let mut stale = Session::new("nested-child-parent", "model");
1757 stale.kind = bamboo_domain::SessionKind::Child;
1758 stale.parent_session_id = Some("root".into());
1759 stale.root_session_id = "root".into();
1760 stale.spawn_depth = 1;
1761 let wait = WaitingForChildrenState::for_children(
1762 vec!["grandchild".into()],
1763 ChildWaitPolicy::All,
1764 chrono::Utc::now(),
1765 );
1766 let mut runtime = AgentRuntimeState::new("old-run");
1767 runtime.waiting_for_children = Some(wait.clone());
1768 stale.agent_runtime_state = Some(runtime);
1769 stale.add_message(bamboo_agent_core::Message::assistant("STALE", None));
1770 storage.save_session(&stale).await.unwrap();
1771
1772 let mut replacement = stale.clone();
1773 replacement.created_at += chrono::Duration::seconds(1);
1774 replacement.messages.clear();
1775 replacement.add_message(bamboo_agent_core::Message::assistant("REPLACEMENT", None));
1776 replacement
1777 .agent_runtime_state
1778 .as_mut()
1779 .unwrap()
1780 .waiting_for_children = None;
1781 storage.delete_session(&stale.id).await.unwrap();
1782 if !missing {
1783 storage.save_session(&replacement).await.unwrap();
1784 }
1785 if !missing {
1786 cache.insert(
1787 replacement.id.clone(),
1788 Arc::new(crate::SessionSnapshot::new(replacement.clone())),
1789 );
1790 }
1791
1792 let saved = save_finalized_runtime_with_inherited_child_wait(
1793 &persistence,
1794 &mut stale,
1795 Some(&wait),
1796 )
1797 .await;
1798 assert!(saved.is_err(), "stale incarnation must not be saved");
1799 publish_final_session_cache(&cache, stale, false, &saved);
1802 if missing {
1803 assert!(
1804 read_cached_session(&cache, &replacement.id).is_none(),
1805 "deleted incarnation must not be resurrected in cache"
1806 );
1807 } else {
1808 let cached = read_cached_session(&cache, &replacement.id).unwrap();
1809 assert_eq!(
1810 cached.created_at, replacement.created_at,
1811 "ABA rejection must retain replacement cache birth"
1812 );
1813 assert_eq!(
1814 serde_json::to_value(&cached.messages).unwrap(),
1815 serde_json::to_value(&replacement.messages).unwrap(),
1816 "ABA rejection must retain replacement cache transcript"
1817 );
1818 }
1819 assert!(saved
1820 .as_ref()
1821 .unwrap_err()
1822 .get_ref()
1823 .is_some_and(|cause| cause.is::<bamboo_domain::SessionAuthorityConflict>()));
1824 let durable = storage.load_session(&replacement.id).await.unwrap();
1825 if missing {
1826 assert!(durable.is_none());
1827 } else {
1828 let durable = durable.unwrap();
1829 assert_eq!(durable.created_at, replacement.created_at);
1830 assert_eq!(
1831 serde_json::to_value(&durable.messages).unwrap(),
1832 serde_json::to_value(&replacement.messages).unwrap()
1833 );
1834 }
1835 }
1836 }
1837
1838 #[test]
1839 fn debug_prompt_snapshot_migrates_legacy_host_paths_before_logging() {
1840 let legacy_project = format!(
1841 "{}\nProject ID: legacy-project\nProject path: /private/legacy-project\nProject home: /private/legacy-home\n{}",
1842 crate::runtime::context::PROJECT_CONTEXT_START_MARKER,
1843 crate::runtime::context::PROJECT_CONTEXT_END_MARKER,
1844 );
1845 let legacy_workspace =
1846 crate::runtime::context::build_workspace_prompt_context("/private/legacy-workspace")
1847 .expect("legacy workspace context");
1848 let legacy_instruction = format!(
1849 "{}\nSource: /private/legacy-workspace/AGENTS.md\nlegacy policy\n{}",
1850 crate::runtime::context::instruction::INSTRUCTION_CONTEXT_START_MARKER,
1851 crate::runtime::context::instruction::INSTRUCTION_CONTEXT_END_MARKER,
1852 );
1853 let mut session = Session::new("legacy-debug-log", "model");
1854 session.add_message(bamboo_agent_core::Message::system(format!(
1855 "Base prompt\n\n{legacy_project}\n\n{legacy_workspace}\n\n{legacy_instruction}"
1856 )));
1857
1858 let prompt = system_prompt_for_debug_log(&mut session).expect("normalized System");
1859
1860 assert_eq!(prompt, "Base prompt");
1861 assert_eq!(
1862 session.workspace_path_meta().as_deref(),
1863 Some("/private/legacy-workspace")
1864 );
1865 for private_path in [
1866 "/private/legacy-project",
1867 "/private/legacy-home",
1868 "/private/legacy-workspace",
1869 ] {
1870 assert!(!prompt.contains(private_path));
1871 }
1872 assert!(!prompt.contains(crate::runtime::context::PROJECT_CONTEXT_START_MARKER));
1873 assert!(!prompt.contains(crate::runtime::context::WORKSPACE_CONTEXT_START_MARKER));
1874 assert!(!prompt
1875 .contains(crate::runtime::context::instruction::INSTRUCTION_CONTEXT_START_MARKER));
1876 }
1877
1878 #[test]
1879 fn reservation_target_requires_domain_id_and_exact_runner_registry() {
1880 let runners = Arc::new(RwLock::new(HashMap::new()));
1881 let other_runners = Arc::new(RwLock::new(HashMap::new()));
1882 let mut reservation = SessionExecutionReservation {
1883 session_id: "session-a".to_string(),
1884 run_id: "run-a".to_string(),
1885 cancel_token: CancellationToken::new(),
1886 runners: runners.clone(),
1887 activation: SessionExecutionActivationOwnership::Unrouted,
1888 root_actor: None,
1889 execution_persistence: None,
1890 persistence_bound: false,
1891 selected_input: false,
1892 armed: true,
1893 };
1894
1895 assert!(reservation.matches_execution_target("session-a", "session-a", &runners));
1896 assert!(!reservation.matches_execution_target("session-b", "session-a", &runners));
1897 assert!(!reservation.matches_execution_target("session-a", "session-b", &runners));
1898 assert!(!reservation.matches_execution_target("session-a", "session-a", &other_runners));
1899 reservation.armed = false;
1900 }
1901
1902 #[tokio::test]
1903 async fn forced_same_run_registration_rejection_never_cancels_live_owner() {
1904 let runners = Arc::new(RwLock::new(HashMap::new()));
1905 let mut runner = AgentRunner::new();
1906 runner.status = AgentStatus::Running;
1907 let run_id = runner.run_id.clone();
1908 let live_cancel_token = runner.cancel_token.clone();
1909 runners.write().await.insert("same-run".to_string(), runner);
1910
1911 let router = SessionActivationRouter::new();
1912 let mut live_registration = router
1913 .register_run("same-run", &run_id)
1914 .await
1915 .expect("first registration owns the run");
1916 let mut rejected = SessionExecutionReservation {
1917 session_id: "same-run".to_string(),
1918 run_id: run_id.clone(),
1919 cancel_token: live_cancel_token.clone(),
1920 runners: runners.clone(),
1921 activation: SessionExecutionActivationOwnership::RegistrationPending(router.clone()),
1922 root_actor: None,
1923 execution_persistence: None,
1924 persistence_bound: false,
1925 selected_input: false,
1926 armed: true,
1927 };
1928
1929 let error = match rejected.ensure_registered().await {
1930 Ok(()) => panic!("a duplicate registration for the same run must be rejected"),
1931 Err(error) => error,
1932 };
1933 assert_eq!(error.existing_run_id(), run_id);
1934 drop(rejected);
1935
1936 assert!(!live_cancel_token.is_cancelled());
1937 assert!(matches!(
1938 runners
1939 .read()
1940 .await
1941 .get("same-run")
1942 .map(|runner| &runner.status),
1943 Some(AgentStatus::Running)
1944 ));
1945 assert!(router.owns_run("same-run", &run_id).await);
1946 live_registration.begin_finalization().await;
1947 assert_eq!(live_registration.finish(0).await.unwrap(), None);
1948 }
1949}