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(Box<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 =
344 SessionExecutionActivationOwnership::Registered(Box::new(registration));
345 Ok(())
346 }
347 Err(error) => {
348 self.disarm_after_registration_rejection(error.existing_run_id());
353 Err(error)
354 }
355 }
356 }
357
358 pub async fn abandon(mut self) {
362 self.armed = false;
363 self.cancel_token.cancel();
364 if let Some(owner) = self.root_actor.take() {
365 (*owner)
366 .finish(bamboo_domain::ActorActivationFinish::Cancelled)
367 .await;
368 }
369 let activation = std::mem::replace(
370 &mut self.activation,
371 SessionExecutionActivationOwnership::Unrouted,
372 );
373 let runners = self.runners.clone();
374 let session_id = self.session_id.clone();
375 let run_id = self.run_id.clone();
376 let cleanup_session_id = session_id.clone();
377 let cleanup_run_id = run_id.clone();
378 let selected_input = self.selected_input;
379 let cleanup = tokio::spawn(async move {
380 cleanup_execution_reservation(
381 activation,
382 runners,
383 cleanup_session_id,
384 cleanup_run_id,
385 selected_input,
386 )
387 .await;
388 });
389 if let Err(error) = cleanup.await {
390 tracing::error!(
391 %session_id,
392 %run_id,
393 %error,
394 "detached execution-reservation cleanup failed"
395 );
396 }
397 }
398
399 fn disarm_after_registration_rejection(&mut self, existing_run_id: &str) {
400 if existing_run_id != self.run_id {
401 self.cancel_token.cancel();
402 }
403 self.armed = false;
404 }
405
406 pub(crate) fn matches_execution_target(
407 &self,
408 session_id: &str,
409 domain_session_id: &str,
410 runners: &Arc<RwLock<HashMap<String, AgentRunner>>>,
411 ) -> bool {
412 self.session_id == session_id
413 && domain_session_id == session_id
414 && Arc::ptr_eq(&self.runners, runners)
415 }
416
417 pub(crate) fn disarm_for_execution(
418 &mut self,
419 ) -> (CancellationToken, Option<SessionRunRegistration>) {
420 self.armed = false;
421 let registration = match std::mem::replace(
422 &mut self.activation,
423 SessionExecutionActivationOwnership::Unrouted,
424 ) {
425 SessionExecutionActivationOwnership::Registered(registration) => Some(*registration),
426 SessionExecutionActivationOwnership::Unrouted => None,
427 SessionExecutionActivationOwnership::UnpublishedActivation(_) => {
428 unreachable!("unpublished activation cannot transfer to execution")
429 }
430 SessionExecutionActivationOwnership::RegistrationPending(_) => {
431 unreachable!("ensure_registered adopts every pending router registration")
432 }
433 };
434 (self.cancel_token.clone(), registration)
435 }
436}
437
438impl Drop for SessionExecutionReservation {
439 fn drop(&mut self) {
440 if !self.armed {
441 return;
442 }
443 self.armed = false;
444 self.cancel_token.cancel();
445 let activation = std::mem::replace(
446 &mut self.activation,
447 SessionExecutionActivationOwnership::Unrouted,
448 );
449 let runners = self.runners.clone();
450 let session_id = self.session_id.clone();
451 let run_id = self.run_id.clone();
452 let selected_input = self.selected_input;
453 if let Ok(runtime) = tokio::runtime::Handle::try_current() {
454 runtime.spawn(async move {
455 cleanup_execution_reservation(
456 activation,
457 runners,
458 session_id,
459 run_id,
460 selected_input,
461 )
462 .await;
463 });
464 }
465 }
466}
467
468pub enum SessionExecutionReserveOutcome {
470 Reserved(SessionExecutionReservation),
471 AlreadyRunning { run_id: String },
472}
473
474async fn register_reserved_activation(
475 router: Arc<SessionActivationRouter>,
476 runners: &Arc<RwLock<HashMap<String, AgentRunner>>>,
477 session_id: &str,
478 run_id: &str,
479 selected_input: bool,
480) -> Result<SessionRunRegistration, SessionRunRegistrationError> {
481 let mut registration = match router.register_run(session_id, run_id).await {
482 Ok(registration) => registration,
483 Err(error) => {
484 let collision = Err(AgentError::LLM(error.to_string()));
485 finalize_rejected_runner_if_distinct(
486 runners,
487 session_id,
488 error.existing_run_id(),
489 run_id,
490 &collision,
491 )
492 .await;
493 return Err(error);
494 }
495 };
496 if selected_input {
497 registration.restrict_to_selected_input();
498 }
499 let cleanup_runners = runners.clone();
500 let cleanup_session_id = session_id.to_string();
501 let cleanup_run_id = run_id.to_string();
502 registration.set_abort_cleanup(move || async move {
503 let abandoned = Err(AgentError::Cancelled);
504 finalize_runner_exact(
505 &cleanup_runners,
506 &cleanup_session_id,
507 &cleanup_run_id,
508 &abandoned,
509 )
510 .await;
511 });
512 Ok(registration)
513}
514
515async fn cleanup_execution_reservation(
516 activation: SessionExecutionActivationOwnership,
517 runners: Arc<RwLock<HashMap<String, AgentRunner>>>,
518 session_id: String,
519 run_id: String,
520 selected_input: bool,
521) {
522 match activation {
523 SessionExecutionActivationOwnership::Registered(registration) => {
524 (*registration).abandon().await;
525 }
526 SessionExecutionActivationOwnership::RegistrationPending(router) => {
527 match register_reserved_activation(
528 router,
529 &runners,
530 &session_id,
531 &run_id,
532 selected_input,
533 )
534 .await
535 {
536 Ok(registration) => registration.abandon().await,
537 Err(error) => {
538 tracing::debug!(
539 %session_id,
540 %run_id,
541 existing_run_id = %error.existing_run_id(),
542 "abandoned activation placeholder was already adopted or superseded"
543 );
544 }
545 }
546 }
547 SessionExecutionActivationOwnership::UnpublishedActivation(_) => {
548 remove_runner_exact(&runners, &session_id, &run_id).await;
549 }
550 SessionExecutionActivationOwnership::Unrouted => {
551 let abandoned = Err(AgentError::Cancelled);
552 finalize_runner_exact(&runners, &session_id, &run_id, &abandoned).await;
553 }
554 }
555}
556
557async fn remove_runner_exact(
558 runners: &Arc<RwLock<HashMap<String, AgentRunner>>>,
559 session_id: &str,
560 run_id: &str,
561) -> bool {
562 let mut runners = runners.write().await;
563 if runners
564 .get(session_id)
565 .is_some_and(|runner| runner.run_id == run_id)
566 {
567 super::runner_lifecycle::remove_runner_entry(&mut runners, session_id).await;
568 true
569 } else {
570 false
571 }
572}
573
574pub async fn reserve_session_execution(
582 agent: &Arc<Agent>,
583 runners: &Arc<RwLock<HashMap<String, AgentRunner>>>,
584 senders: &Arc<RwLock<HashMap<String, broadcast::Sender<AgentEvent>>>>,
585 session_id: &str,
586 event_sender: &broadcast::Sender<AgentEvent>,
587) -> SessionExecutionReserveOutcome {
588 reserve_session_execution_inner(agent, runners, senders, session_id, event_sender, false).await
589}
590
591pub async fn reserve_selected_input_execution(
594 agent: &Arc<Agent>,
595 runners: &Arc<RwLock<HashMap<String, AgentRunner>>>,
596 senders: &Arc<RwLock<HashMap<String, broadcast::Sender<AgentEvent>>>>,
597 session_id: &str,
598 event_sender: &broadcast::Sender<AgentEvent>,
599) -> SessionExecutionReserveOutcome {
600 reserve_session_execution_inner(agent, runners, senders, session_id, event_sender, true).await
601}
602
603async fn reserve_session_execution_inner(
604 agent: &Arc<Agent>,
605 runners: &Arc<RwLock<HashMap<String, AgentRunner>>>,
606 senders: &Arc<RwLock<HashMap<String, broadcast::Sender<AgentEvent>>>>,
607 session_id: &str,
608 event_sender: &broadcast::Sender<AgentEvent>,
609 selected_input: bool,
610) -> SessionExecutionReserveOutcome {
611 let reservation = match reserve_runner_core(runners, senders, session_id, event_sender).await {
612 ReserveOutcome::AlreadyRunning(run_id) => {
613 return SessionExecutionReserveOutcome::AlreadyRunning { run_id };
614 }
615 ReserveOutcome::Reserved(reservation) => reservation,
616 };
617
618 let mut execution_reservation = SessionExecutionReservation {
624 session_id: session_id.to_string(),
625 run_id: reservation.run_id,
626 cancel_token: reservation.cancel_token,
627 runners: runners.clone(),
628 activation: SessionExecutionActivationOwnership::Unrouted,
629 root_actor: None,
630 execution_persistence: None,
631 persistence_bound: false,
632 selected_input,
633 armed: true,
634 };
635
636 if let Some(router) = agent.activation_router().cloned() {
637 execution_reservation.activation =
638 SessionExecutionActivationOwnership::RegistrationPending(router);
639 if let Err(error) = execution_reservation.ensure_registered().await {
640 tracing::warn!(
641 %session_id,
642 attempted_run_id = %execution_reservation.run_id(),
643 existing_run_id = %error.existing_run_id(),
644 %error,
645 "runner reservation collided with an existing logical-session owner"
646 );
647 return SessionExecutionReserveOutcome::AlreadyRunning {
648 run_id: error.existing_run_id().to_string(),
649 };
650 }
651 }
652 SessionExecutionReserveOutcome::Reserved(execution_reservation)
653}
654
655pub fn read_cached_session(cache: &SessionCache, id: &str) -> Option<bamboo_agent_core::Session> {
663 cache
664 .get(id)
665 .map(|e| e.value().clone())
666 .map(|a| a.read().clone())
667}
668
669const SKILL_CONTEXT_START_MARKER: &str = "<!-- BAMBOO_SKILL_CONTEXT_START -->";
670const TOOL_GUIDE_START_MARKER: &str = "<!-- BAMBOO_TOOL_GUIDE_START -->";
671const EXTERNAL_MEMORY_START_MARKER: &str = "<!-- BAMBOO_EXTERNAL_MEMORY_START -->";
672const TASK_LIST_START_MARKER: &str = "<!-- BAMBOO_TASK_LIST_START -->";
673
674pub struct SessionExecutionOutcome {
681 pub success: bool,
683 pub cancelled: bool,
685 pub error: Option<String>,
687}
688
689impl SessionExecutionOutcome {
690 fn from_result(result: &Result<(), AgentError>) -> Self {
691 match result {
692 Ok(()) => Self {
693 success: true,
694 cancelled: false,
695 error: None,
696 },
697 Err(error) => Self {
698 success: false,
699 cancelled: error.is_cancelled(),
700 error: Some(error.to_string()),
701 },
702 }
703 }
704}
705
706pub type SessionCompletionHook = Box<
716 dyn for<'a> FnOnce(
717 SessionExecutionOutcome,
718 &'a mut Session,
719 ) -> Pin<Box<dyn Future<Output = ()> + Send + 'a>>
720 + Send,
721>;
722
723pub struct SessionExecutionArgs {
729 pub agent: Arc<Agent>,
731 pub session_id: String,
732 pub session: Session,
733 pub execution_reservation: SessionExecutionReservation,
736
737 pub tools_override: Option<Arc<dyn ToolExecutor>>,
739 pub provider_override: Option<Arc<dyn LLMProvider>>,
740 pub model_roster: ModelRoster,
744 pub reasoning_effort: Option<ReasoningEffort>,
745 pub reasoning_effort_source: String,
746 pub auxiliary_model_resolver:
747 Option<Arc<dyn Fn() -> crate::runtime::config::AuxiliaryModelConfig + Send + Sync>>,
748 pub disabled_filter_resolver: Option<DisabledFilterResolver>,
752 pub disabled_tools: Option<BTreeSet<String>>,
753 pub disabled_skill_ids: Option<BTreeSet<String>>,
754 pub selected_skill_ids: Option<Vec<String>>,
755 pub selected_skill_mode: Option<String>,
756 pub mpsc_tx: mpsc::Sender<AgentEvent>,
757 pub history_commit_barrier: HistoryCommitBarrier,
760 pub image_fallback: Option<ImageFallbackConfig>,
761 pub gold_config: Option<GoldConfig>,
762 pub guardian_config: Option<GuardianConfig>,
764 pub guardian_spawner: Option<Arc<dyn GuardianSpawner>>,
767 pub bash_resume_hook: Option<Arc<dyn BashResumeHook>>,
769 pub bash_completion_sink: Option<Arc<dyn BashCompletionSink>>,
771 pub app_data_dir: Option<std::path::PathBuf>,
772 pub run_budget: Option<bamboo_config::RunBudgetConfig>,
775
776 pub runners: Arc<RwLock<HashMap<String, AgentRunner>>>,
778 pub sessions_cache: SessionCache,
779
780 pub on_complete: Option<SessionCompletionHook>,
783
784 pub child_completion_handler: Option<Arc<dyn super::ChildCompletionHandler>>,
796}
797
798struct ExecuteRequestParams {
808 tools: Option<Arc<dyn ToolExecutor>>,
809 provider_override: Option<Arc<dyn LLMProvider>>,
810 model_roster: ModelRoster,
811 reasoning_effort: Option<ReasoningEffort>,
812 auxiliary_model_resolver: Option<Arc<dyn Fn() -> AuxiliaryModelConfig + Send + Sync>>,
813 disabled_filter_resolver: Option<DisabledFilterResolver>,
814 disabled_tools: Option<BTreeSet<String>>,
815 disabled_skill_ids: Option<BTreeSet<String>>,
816 selected_skill_ids: Option<Vec<String>>,
817 selected_skill_mode: Option<String>,
818 image_fallback: Option<ImageFallbackConfig>,
819 gold_config: Option<GoldConfig>,
820 guardian_config: Option<GuardianConfig>,
821 guardian_spawner: Option<Arc<dyn GuardianSpawner>>,
822 bash_resume_hook: Option<Arc<dyn BashResumeHook>>,
823 bash_completion_sink: Option<Arc<dyn BashCompletionSink>>,
824 app_data_dir: Option<std::path::PathBuf>,
825 run_budget: Option<bamboo_config::RunBudgetConfig>,
826}
827
828fn build_execute_request(
835 initial_message: String,
836 event_tx: mpsc::Sender<AgentEvent>,
837 cancel_token: CancellationToken,
838 params: ExecuteRequestParams,
839) -> ExecuteRequest {
840 let ExecuteRequestParams {
841 tools,
842 provider_override,
843 model_roster,
844 reasoning_effort,
845 auxiliary_model_resolver,
846 disabled_filter_resolver,
847 disabled_tools,
848 disabled_skill_ids,
849 selected_skill_ids,
850 selected_skill_mode,
851 image_fallback,
852 gold_config,
853 guardian_config,
854 guardian_spawner,
855 bash_resume_hook,
856 bash_completion_sink,
857 app_data_dir,
858 run_budget,
859 } = params;
860
861 let mut builder = ExecuteRequestBuilder::new(initial_message, event_tx, cancel_token)
862 .model_roster(model_roster)
863 .gold_config(gold_config)
864 .guardian_config(guardian_config)
865 .guardian_spawner(guardian_spawner)
866 .bash_resume_hook(bash_resume_hook)
867 .bash_completion_sink(bash_completion_sink);
868
869 if let Some(run_budget) = run_budget {
870 builder = builder.run_budget(run_budget);
871 }
872
873 if let Some(tools) = tools {
874 builder = builder.tools(tools);
875 }
876 if let Some(provider_override) = provider_override {
877 builder = builder.provider_override(provider_override);
878 }
879 if let Some(reasoning_effort) = reasoning_effort {
880 builder = builder.reasoning_effort(reasoning_effort);
881 }
882 if let Some(disabled_filter_resolver) = disabled_filter_resolver {
883 builder = builder.disabled_filter_resolver(disabled_filter_resolver);
884 }
885 if let Some(auxiliary_model_resolver) = auxiliary_model_resolver {
886 builder = builder.auxiliary_model_resolver(auxiliary_model_resolver);
887 }
888 if let Some(disabled_tools) = disabled_tools {
889 builder = builder.disabled_tools(disabled_tools);
890 }
891 if let Some(disabled_skill_ids) = disabled_skill_ids {
892 builder = builder.disabled_skill_ids(disabled_skill_ids);
893 }
894 if let Some(selected_skill_ids) = selected_skill_ids {
895 builder = builder.selected_skill_ids(selected_skill_ids);
896 }
897 if let Some(selected_skill_mode) = selected_skill_mode {
898 builder = builder.selected_skill_mode(selected_skill_mode);
899 }
900 if let Some(image_fallback) = image_fallback {
901 builder = builder.image_fallback(image_fallback);
902 }
903 if let Some(app_data_dir) = app_data_dir {
904 builder = builder.app_data_dir(app_data_dir);
905 }
906
907 builder.build()
908}
909
910pub fn spawn_session_execution(args: SessionExecutionArgs) {
919 spawn_session_execution_with_inputs(args, None);
920}
921
922pub fn spawn_session_execution_with_inputs(
925 args: SessionExecutionArgs,
926 inputs: Option<crate::runtime::config::UntrustedExecutionInputs>,
927) {
928 let span_session_id = args.session_id.clone();
929 let session_span = tracing::info_span!("agent_execution", session_id = %span_session_id);
930
931 tokio::spawn(
932 async move {
933 let SessionExecutionArgs {
934 agent,
935 session_id,
936 mut session,
937 mut execution_reservation,
938 tools_override,
939 provider_override,
940 model_roster,
941 reasoning_effort,
942 reasoning_effort_source,
943 auxiliary_model_resolver,
944 disabled_filter_resolver,
945 disabled_tools,
946 disabled_skill_ids,
947 selected_skill_ids,
948 selected_skill_mode,
949 mpsc_tx,
950 mut history_commit_barrier,
951 image_fallback,
952 gold_config,
953 guardian_config,
954 guardian_spawner,
955 bash_resume_hook,
956 bash_completion_sink,
957 app_data_dir,
958 run_budget,
959 runners,
960 sessions_cache,
961 on_complete,
962 child_completion_handler,
963 } = args;
964
965 if !execution_reservation.matches_execution_target(&session_id, &session.id, &runners) {
966 tracing::error!(
967 %session_id,
968 domain_session_id = %session.id,
969 reservation_session_id = %execution_reservation.session_id(),
970 run_id = %execution_reservation.run_id(),
971 same_runner_registry = Arc::ptr_eq(&execution_reservation.runners, &runners),
972 "refusing mismatched session execution reservation"
973 );
974 execution_reservation.abandon().await;
975 return;
976 }
977 if let Err(error) = execution_reservation.ensure_registered().await {
978 tracing::warn!(
979 %session_id,
980 run_id = %execution_reservation.run_id(),
981 %error,
982 "session execution could not adopt its router activation owner"
983 );
984 return;
985 }
986 if execution_reservation.root_actor.is_none()
987 && agent.persistence().root_actor_execution_required(&session)
988 {
989 let message =
990 "Root execution route requires a bound Actor writer and event handoff";
991 tracing::warn!(%session_id, "{message}");
992 let _ = mpsc_tx
993 .send(AgentEvent::Error {
994 message: message.into(),
995 })
996 .await;
997 execution_reservation.abandon().await;
998 return;
999 }
1000 if let Err(error) = execution_reservation.bind_execution_persistence(&agent, &session) {
1003 let _ = mpsc_tx
1004 .send(AgentEvent::Error {
1005 message: error.to_string(),
1006 })
1007 .await;
1008 execution_reservation.abandon().await;
1009 return;
1010 }
1011 let ownership = match &execution_reservation.activation {
1012 SessionExecutionActivationOwnership::Registered(run) => Some(run.ownership()),
1013 _ => None,
1014 };
1015 let admission_persistence = execution_reservation
1016 .execution_persistence()
1017 .unwrap_or_else(|| agent.persistence().clone());
1018 if let (Some(router), Some(ownership)) = (agent.activation_router(), ownership.as_ref())
1019 {
1020 if !router.bind_run_lifetime(ownership, &session).await {
1021 let _ = mpsc_tx
1022 .send(AgentEvent::Error {
1023 message: "registered run Session lifetime differs".into(),
1024 })
1025 .await;
1026 execution_reservation.abandon().await;
1027 return;
1028 }
1029 }
1030 if let Err(error) = crate::runtime::counted_loop::admit_registered(
1031 &mut session,
1032 ownership.as_ref(),
1033 agent.activation_router(),
1034 admission_persistence.as_ref(),
1035 )
1036 .await
1037 {
1038 let _ = mpsc_tx
1039 .send(AgentEvent::Error {
1040 message: error.to_string(),
1041 })
1042 .await;
1043 execution_reservation.abandon().await;
1044 return;
1045 }
1046 let mut root_actor = execution_reservation.root_actor.take().map(|owner| *owner);
1047 let root_actor_bound = root_actor.is_some();
1048 let persistence = execution_reservation
1049 .execution_persistence()
1050 .unwrap_or_else(|| agent.persistence().clone());
1051 let agent = Arc::new(agent.with_execution_persistence(persistence));
1052 let activation_run_id = execution_reservation.run_id().to_string();
1053 let prior_message_ids = session.messages.iter().map(|m| m.id.clone()).collect();
1054 let (cancel_token, mut activation_registration) =
1055 execution_reservation.disarm_for_execution();
1056
1057 let model = model_roster.model.clone().unwrap_or_default();
1061
1062 let initial_message = initial_user_message_for_session(&session);
1063 let selected_skill_ids =
1064 selected_skill_ids.or_else(|| selected_skill_ids_for_session(&session));
1065 let selected_skill_mode =
1066 selected_skill_mode.or_else(|| selected_skill_mode_for_session(&session));
1067
1068 tracing::info!(
1069 "[{}] Using resolved session model: {}, reasoning_effort={}, reasoning_source={}",
1070 session_id,
1071 model,
1072 reasoning_effort
1073 .map(ReasoningEffort::as_str)
1074 .unwrap_or("none"),
1075 reasoning_effort_source
1076 );
1077
1078 crate::session_app::execution_prep::prepare_session_for_execution(
1085 &mut session,
1086 None,
1087 Some(&model),
1088 );
1089
1090 let system_prompt = system_prompt_for_debug_log(&mut session);
1091 if let Some(prompt) = system_prompt.as_ref() {
1092 log_base_system_prompt_snapshot(&session_id, prompt);
1093 }
1094
1095 let inherited_child_wait = agent.persistence().inherited_child_wait();
1099 let execute_request = build_execute_request(
1100 initial_message,
1101 mpsc_tx.clone(),
1102 cancel_token,
1103 ExecuteRequestParams {
1104 tools: tools_override,
1105 provider_override,
1106 model_roster,
1107 reasoning_effort,
1108 auxiliary_model_resolver,
1109 disabled_filter_resolver,
1110 disabled_tools,
1111 disabled_skill_ids,
1112 selected_skill_ids,
1113 selected_skill_mode,
1114 image_fallback,
1115 gold_config,
1116 guardian_config,
1117 guardian_spawner,
1118 bash_resume_hook,
1119 bash_completion_sink,
1120 app_data_dir,
1121 run_budget,
1122 },
1123 );
1124
1125 let result = {
1132 use futures::FutureExt;
1133 match std::panic::AssertUnwindSafe(agent.execute_with_inputs(
1134 &mut session,
1135 execute_request,
1136 inputs,
1137 ))
1138 .catch_unwind()
1139 .await
1140 {
1141 Ok(result) => result,
1142 Err(panic) => {
1143 let message = panic
1144 .downcast_ref::<&str>()
1145 .map(|s| (*s).to_string())
1146 .or_else(|| panic.downcast_ref::<String>().cloned())
1147 .unwrap_or_else(|| "non-string panic payload".to_string());
1148 tracing::error!(
1149 "[{}] agent execution panicked; finalizing as terminal error: {}",
1150 session_id,
1151 message
1152 );
1153 Err(AgentError::LLM(format!(
1154 "agent execution panicked: {message}"
1155 )))
1156 }
1157 }
1158 };
1159
1160 if let Some(error_event) = terminal_error_event_for_result(&result) {
1162 let _ = mpsc_tx.send(error_event).await;
1163 }
1164
1165 let suspended_non_terminal = result.is_ok()
1182 && session
1183 .metadata
1184 .get("runtime.suspend_reason")
1185 .is_some_and(|reason| !reason.trim().is_empty());
1186 match &result {
1187 Ok(()) if suspended_non_terminal => {
1188 session.set_last_run_status("suspended");
1189 session.clear_last_run_error();
1190 }
1191 Ok(()) => {
1192 session.set_last_run_status("completed");
1193 session.clear_last_run_error();
1194 }
1195 Err(error) if error.is_cancelled() => {
1196 session.set_last_run_status("cancelled");
1197 session.set_last_run_error(error.to_string());
1198 }
1199 Err(error) => {
1200 session.set_last_run_status("error");
1201 session.set_last_run_error(error.to_string());
1202 }
1203 }
1204
1205 if let Some(on_complete) = on_complete {
1209 on_complete(SessionExecutionOutcome::from_result(&result), &mut session).await;
1210 }
1211
1212 let executed_admitted_generation = session
1216 .session_inbox_admission()
1217 .map_or(0, |state| state.last_admitted_sequence);
1218 let legacy_migration =
1219 crate::runtime::runner::state_bridge::migrate_legacy_pending_only(
1220 &mut session,
1221 Some(agent.storage()),
1222 Some(agent.persistence()),
1223 agent.session_inbox(),
1224 )
1225 .await;
1226 let pending_boundary_generation = session
1227 .session_inbox_admission()
1228 .and_then(|state| state.pending_activation_generation());
1229 let pending_generation = match (
1230 pending_boundary_generation,
1231 legacy_migration.highest_generation,
1232 ) {
1233 (Some(left), Some(right)) => Some(left.max(right)),
1234 (left, right) => left.or(right),
1235 };
1236 if let (Some(generation), Some(router)) =
1237 (pending_generation, agent.activation_router())
1238 {
1239 let activation_ready = if let Some(inbox) = agent.session_inbox() {
1240 match inbox
1241 .mark_activation_eligible(
1242 &session_id,
1243 generation,
1244 bamboo_domain::SessionActivationPolicy::InterruptSpecificWait,
1245 )
1246 .await
1247 {
1248 Ok(()) => true,
1249 Err(error) => {
1250 tracing::error!(
1251 session_id = %session_id,
1252 %error,
1253 "failed to persist unadmitted SessionInbox activation watermark"
1254 );
1255 false
1256 }
1257 }
1258 } else {
1259 false
1260 };
1261 if activation_ready {
1262 if let Err(error) = bamboo_domain::SessionActivationPort::request_activation(
1263 router.as_ref(),
1264 &session_id,
1265 generation,
1266 )
1267 .await
1268 {
1269 tracing::error!(
1270 session_id = %session_id,
1271 %error,
1272 "failed to hand unadmitted SessionInbox generation to activation router"
1273 );
1274 }
1275 }
1276 }
1277
1278 if let Some(registration) = activation_registration.as_mut() {
1282 registration.begin_finalization().await;
1283 }
1284
1285 super::child_completion::ChildCompletionSource::prepare(
1289 &mut session,
1290 &activation_run_id,
1291 &prior_message_ids,
1292 );
1293 let saved = save_finalized_runtime_with_inherited_child_wait(
1294 agent.persistence().as_ref(),
1295 &mut session,
1296 inherited_child_wait
1297 .as_ref()
1298 .map(|inherited| inherited.wait()),
1299 )
1300 .await;
1301 let history_committed = saved.is_ok();
1302 let completion_source =
1303 super::child_completion::ChildCompletionSource::after_final_save(
1304 &session,
1305 history_committed,
1306 );
1307 let completed_at = chrono::Utc::now();
1308 if let Err(error) = &saved {
1309 tracing::warn!("[{}] Failed to save session: {}", session_id, error);
1310 }
1311
1312 if history_committed
1318 && !history_commit_barrier
1319 .send_and_wait(&mpsc_tx, session_id.clone())
1320 .await
1321 {
1322 tracing::warn!(
1323 session_id = %session_id,
1324 "history commit barrier could not be published before runner finalization"
1325 );
1326 }
1327
1328 finalize_runner_exact(&runners, &session_id, &activation_run_id, &result).await;
1334
1335 if let Some(owner) = root_actor.take() {
1336 let outcome = if result.as_ref().is_err_and(|error| error.is_cancelled()) {
1337 bamboo_domain::ActorActivationFinish::Cancelled
1338 } else if result.is_err() || !history_committed {
1339 bamboo_domain::ActorActivationFinish::Failed
1340 } else {
1341 bamboo_domain::ActorActivationFinish::Succeeded
1342 };
1343 owner.finish(outcome).await;
1344 }
1345
1346 let finalization = if let Some(registration) = activation_registration.take() {
1347 registration.finish(executed_admitted_generation).await
1348 } else {
1349 Ok(None)
1350 };
1351 if let Err(error) = finalization {
1352 tracing::error!(
1356 session_id = %session_id,
1357 %error,
1358 "failed to activate successor for finalization-racing SessionInbox delivery"
1359 );
1360 }
1361
1362 let child_completion = child_completion_handler.filter(|_| {
1371 session.kind == bamboo_agent_core::SessionKind::Child
1372 && session.parent_session_id.is_some()
1373 });
1374 let parent_session_id = session.parent_session_id.clone();
1375 let child_status = session.last_run_status();
1376 let child_error = session.last_run_error();
1377
1378 publish_final_session_cache(&sessions_cache, session, root_actor_bound, &saved);
1379
1380 if let (Some(handler), Some(parent_session_id), Some(status)) =
1381 (child_completion, parent_session_id, child_status)
1382 {
1383 use futures::FutureExt;
1384 let completion = ChildCompletion {
1385 parent_session_id: parent_session_id.clone(),
1386 child_session_id: session_id.clone(),
1387 status,
1388 error: child_error,
1389 completed_at,
1390 source: completion_source,
1391 };
1392 if std::panic::AssertUnwindSafe(handler.on_child_completed(completion))
1393 .catch_unwind()
1394 .await
1395 .is_err()
1396 {
1397 tracing::error!(
1398 %parent_session_id,
1399 child_session_id = %session_id,
1400 "child completion handler panicked on resumed-child terminal"
1401 );
1402 }
1403 }
1404
1405 tracing::info!("[{}] Agent execution completed", session_id);
1406 }
1407 .instrument(session_span),
1408 );
1409}
1410
1411async fn save_finalized_runtime_with_inherited_child_wait(
1413 persistence: &dyn bamboo_domain::RuntimeSessionPersistence,
1414 session: &mut Session,
1415 inherited_child_wait: Option<&bamboo_domain::session::runtime_state::WaitingForChildrenState>,
1416) -> std::io::Result<()> {
1417 if let Some(inherited) = inherited_child_wait.filter(|wait| {
1418 session
1419 .agent_runtime_state
1420 .as_ref()
1421 .and_then(|runtime| runtime.waiting_for_children.as_ref())
1422 == Some(*wait)
1423 }) {
1424 persistence
1425 .save_finalized_runtime_with_inherited_child_wait(session, inherited)
1426 .await
1427 } else {
1428 persistence.save_finalized_runtime_session(session).await
1429 }
1430}
1431
1432fn publish_final_session_cache(
1433 cache: &SessionCache,
1434 session: Session,
1435 root_actor_bound: bool,
1436 saved: &std::io::Result<()>,
1437) {
1438 let authority_conflict = saved.as_ref().err().is_some_and(|error| {
1439 error
1440 .get_ref()
1441 .is_some_and(|cause| cause.is::<bamboo_domain::SessionAuthorityConflict>())
1442 });
1443 if !authority_conflict && !root_actor_bound {
1446 cache.insert(
1447 session.id.clone(),
1448 Arc::new(crate::SessionSnapshot::new(session)),
1449 );
1450 }
1451}
1452
1453pub fn log_base_system_prompt_snapshot(session_id: &str, prompt: &str) {
1455 tracing::info!(
1456 "[{}] Base system prompt snapshot: len={} chars, has_skill={}, has_tool_guide={}, has_external_memory={}, has_task_list={}",
1457 session_id,
1458 prompt.len(),
1459 prompt.contains(SKILL_CONTEXT_START_MARKER),
1460 prompt.contains(TOOL_GUIDE_START_MARKER),
1461 prompt.contains(EXTERNAL_MEMORY_START_MARKER),
1462 prompt.contains(TASK_LIST_START_MARKER),
1463 );
1464
1465 tracing::debug!(
1466 "[{}] ========== BASE SYSTEM PROMPT SNAPSHOT ==========",
1467 session_id
1468 );
1469 tracing::debug!("[{}] Snapshot length: {} chars", session_id, prompt.len());
1470 tracing::debug!("[{}] -----------------------------------", session_id);
1471 tracing::debug!("[{}] {}", session_id, prompt);
1472 tracing::debug!(
1473 "[{}] ========== END BASE SYSTEM PROMPT SNAPSHOT ==========",
1474 session_id
1475 );
1476}
1477
1478pub fn terminal_error_event_for_result(result: &Result<(), AgentError>) -> Option<AgentEvent> {
1480 match result {
1481 Ok(_) => None,
1482 Err(error) if error.is_cancelled() => Some(AgentEvent::Error {
1483 message: "Agent execution cancelled by user".to_string(),
1484 }),
1485 Err(error) => Some(AgentEvent::Error {
1486 message: error.to_string(),
1487 }),
1488 }
1489}
1490
1491fn system_prompt_for_session(session: &Session) -> Option<String> {
1494 session
1495 .messages
1496 .iter()
1497 .find(|message| matches!(message.role, bamboo_agent_core::Role::System))
1498 .map(|message| message.content.clone())
1499}
1500
1501fn system_prompt_for_debug_log(session: &mut Session) -> Option<String> {
1505 crate::runtime::runner::session_setup::migrate_legacy_workspace_prompt(session);
1506 system_prompt_for_session(session)
1507}
1508
1509fn initial_user_message_for_session(session: &Session) -> String {
1510 session
1511 .messages
1512 .last()
1513 .filter(|message| matches!(message.role, bamboo_agent_core::Role::User))
1514 .map(|message| message.content.clone())
1515 .unwrap_or_default()
1516}
1517
1518fn selected_skill_ids_for_session(session: &Session) -> Option<Vec<String>> {
1519 session
1520 .metadata
1521 .get("selected_skill_ids")
1522 .and_then(|raw| bamboo_skills::selection::parse_selected_skill_ids_metadata(raw))
1523}
1524
1525fn selected_skill_mode_for_session(session: &Session) -> Option<String> {
1526 let value = session
1527 .metadata
1528 .get("skill_mode")
1529 .or_else(|| session.metadata.get("mode"))?;
1530 let trimmed = value.trim();
1531 if trimmed.is_empty() {
1532 None
1533 } else {
1534 Some(trimmed.to_string())
1535 }
1536}
1537
1538pub(super) struct RootActorExecution {
1539 binding: bamboo_domain::RootActorExecutionBinding,
1540 stop_renewal: CancellationToken,
1541 renewal: Option<tokio::task::JoinHandle<()>>,
1542 response_events: Vec<AgentEvent>,
1543}
1544
1545impl RootActorExecution {
1546 pub(super) fn start(binding: bamboo_domain::RootActorExecutionBinding) -> Self {
1547 let stop_renewal = CancellationToken::new();
1548 let stop = stop_renewal.clone();
1549 let directory = binding.directory.clone();
1550 let fence = binding.owner.fence.clone();
1551 let duration = binding.lease_duration;
1552 let renewal = tokio::spawn(async move {
1553 let mut interval = tokio::time::interval(duration / 3);
1554 interval.tick().await;
1555 loop {
1556 tokio::select! {
1557 _ = stop.cancelled() => break,
1558 _ = interval.tick() => {
1559 let now = chrono::Utc::now();
1560 let expires = now + chrono::Duration::from_std(duration).expect("bounded Root lease");
1561 if let Err(error) = directory.renew_activation(&fence, now, expires).await {
1562 tracing::warn!(actor_id = %fence.actor_id, run_id = %fence.run_id,
1566 %error, "Root Actor renewal lost its execution authority");
1567 break;
1568 }
1569 }
1570 }
1571 }
1572 });
1573 Self {
1574 binding,
1575 stop_renewal,
1576 renewal: Some(renewal),
1577 response_events: Vec::new(),
1578 }
1579 }
1580
1581 pub(super) fn persistence(
1582 &self,
1583 ) -> std::sync::Arc<dyn bamboo_domain::RuntimeSessionPersistence> {
1584 self.binding.persistence.clone()
1585 }
1586
1587 pub(super) async fn finish(mut self, outcome: bamboo_domain::ActorActivationFinish) {
1588 self.stop_renewal.cancel();
1589 if let Some(renewal) = self.renewal.take() {
1590 let _ = renewal.await;
1591 }
1592 let directory = self.binding.directory.clone();
1593 let fence = self.binding.owner.fence.clone();
1594 let finish = tokio::spawn(async move {
1595 if let Err(error) = directory
1596 .finish_activation(&fence, chrono::Utc::now(), outcome)
1597 .await
1598 {
1599 tracing::warn!(actor_id = %fence.actor_id, run_id = %fence.run_id,
1600 %error, "Root Actor finish rejected its obsolete execution fence");
1601 }
1602 });
1603 self.binding.disarm_abandonment();
1604 let _ = finish.await;
1605 }
1606}
1607
1608impl Drop for RootActorExecution {
1609 fn drop(&mut self) {
1610 self.stop_renewal.cancel();
1611 }
1614}
1615
1616#[cfg(test)]
1617#[path = "agent_spawn_peer_privacy_tests.rs"]
1618mod peer_privacy_tests;
1619
1620#[cfg(test)]
1621mod reservation_tests {
1622 use super::*;
1623 use crate::runtime::execution::runner_state::AgentStatus;
1624
1625 struct CompletionAtFinalSave {
1626 directory: std::path::PathBuf,
1627 persistence: bamboo_storage::LockedSessionStore,
1628 cleared: std::sync::atomic::AtomicBool,
1629 }
1630
1631 impl CompletionAtFinalSave {
1632 async fn clear_wait(&self, id: &str) {
1633 let mut process = std::process::Command::new(std::env::current_exe().unwrap())
1634 .args(["--exact", "runtime::execution::agent_spawn::reservation_tests::inherited_child_wait_adopts_independent_completion_at_final_save", "--nocapture"])
1635 .env("BAMBOO_FINAL_WAIT_COMPLETION_DIR", &self.directory)
1636 .env("BAMBOO_FINAL_WAIT_COMPLETION_ID", id)
1637 .spawn().unwrap();
1638 let deadline = std::time::Instant::now() + std::time::Duration::from_secs(15);
1639 loop {
1640 if let Some(status) = process.try_wait().unwrap() {
1641 assert!(status.success(), "completion process failed");
1642 break;
1643 }
1644 if std::time::Instant::now() >= deadline {
1645 let _ = process.kill();
1646 let _ = process.wait();
1647 panic!("bounded completion process deadline");
1648 }
1649 std::thread::sleep(std::time::Duration::from_millis(5));
1650 }
1651 self.cleared
1652 .store(true, std::sync::atomic::Ordering::SeqCst);
1653 }
1654 }
1655
1656 #[async_trait::async_trait]
1657 impl bamboo_domain::RuntimeSessionPersistence for CompletionAtFinalSave {
1658 async fn save_runtime_session(&self, session: &mut Session) -> std::io::Result<()> {
1659 self.clear_wait(&session.id).await;
1660 self.persistence.merge_save_finalized_runtime(session).await
1661 }
1662 async fn save_finalized_runtime_with_inherited_child_wait(
1663 &self,
1664 session: &mut Session,
1665 inherited: &bamboo_domain::session::runtime_state::WaitingForChildrenState,
1666 ) -> std::io::Result<()> {
1667 self.clear_wait(&session.id).await;
1668 self.persistence
1669 .merge_save_inherited_child_wait_and_publish(session, inherited, |_, _| {})
1670 .await
1671 }
1672 }
1673
1674 #[tokio::test]
1675 async fn inherited_child_wait_adopts_independent_completion_at_final_save() {
1676 use bamboo_agent_core::storage::Storage;
1677 use bamboo_domain::session::runtime_state::{
1678 AgentRuntimeState, AgentStatusState, ChildWaitPolicy, WaitingForChildrenState,
1679 };
1680 if let Some(directory) = std::env::var_os("BAMBOO_FINAL_WAIT_COMPLETION_DIR") {
1681 let remote = bamboo_storage::SessionStoreV2::new(directory.into())
1682 .await
1683 .unwrap();
1684 let id = std::env::var("BAMBOO_FINAL_WAIT_COMPLETION_ID").unwrap();
1685 let mut latest = remote.load_session(&id).await.unwrap().unwrap();
1686 let runtime = latest.agent_runtime_state.as_mut().unwrap();
1687 runtime.waiting_for_children = None;
1688 runtime.status = AgentStatusState::Idle;
1689 runtime.suspension = None;
1690 latest.metadata.remove("runtime.suspend_reason");
1691 remote.save_runtime_state(&latest).await.unwrap();
1692 return;
1693 }
1694 let temp = tempfile::tempdir().unwrap();
1695 let local = Arc::new(
1696 bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
1697 .await
1698 .unwrap(),
1699 );
1700 let mut runner = Session::new("cross-host-final-wait", "model");
1704 let inherited = WaitingForChildrenState::for_children(
1705 vec!["child".into()],
1706 ChildWaitPolicy::All,
1707 chrono::Utc::now(),
1708 );
1709 let mut runtime = AgentRuntimeState::new("run");
1710 runtime.waiting_for_children = Some(inherited.clone());
1711 runtime.status = AgentStatusState::Suspended;
1712 runner.agent_runtime_state = Some(runtime);
1713 runner.metadata.insert(
1714 "runtime.suspend_reason".into(),
1715 "waiting_for_children".into(),
1716 );
1717 runner.set_last_run_status("suspended");
1718 local.save_session(&runner).await.unwrap();
1719 let persistence = CompletionAtFinalSave {
1720 directory: temp.path().to_path_buf(),
1721 persistence: bamboo_storage::LockedSessionStore::new(local.clone()),
1722 cleared: std::sync::atomic::AtomicBool::new(false),
1723 };
1724 save_finalized_runtime_with_inherited_child_wait(
1725 &persistence,
1726 &mut runner,
1727 Some(&inherited),
1728 )
1729 .await
1730 .unwrap();
1731 assert!(persistence
1732 .cleared
1733 .load(std::sync::atomic::Ordering::SeqCst));
1734 let durable = local.load_session(&runner.id).await.unwrap().unwrap();
1735 for session in [&runner, &durable] {
1736 let runtime = session.agent_runtime_state.as_ref().unwrap();
1737 assert!(
1738 runtime.waiting_for_children.is_none(),
1739 "independent completion must not be resurrected"
1740 );
1741 assert_eq!(runtime.status, AgentStatusState::Idle);
1742 assert!(runtime.suspension.is_none());
1743 assert!(!session.metadata.contains_key("runtime.suspend_reason"));
1744 assert_eq!(session.last_run_status().as_deref(), Some("completed"));
1745 }
1746 }
1747
1748 #[tokio::test]
1749 async fn inherited_child_wait_rejects_aba_and_missing_birth_without_cache_publication() {
1750 use bamboo_agent_core::storage::Storage;
1751 use bamboo_domain::session::runtime_state::{
1752 AgentRuntimeState, ChildWaitPolicy, WaitingForChildrenState,
1753 };
1754 use bamboo_storage::{LockedSessionStore, SessionStoreV2};
1755
1756 for missing in [false, true] {
1757 let temp = tempfile::tempdir().unwrap();
1758 let storage = Arc::new(
1759 SessionStoreV2::new(temp.path().to_path_buf())
1760 .await
1761 .unwrap(),
1762 );
1763 let persistence = LockedSessionStore::new(storage.clone());
1764 let cache: SessionCache = Arc::default();
1765 storage
1766 .save_session(&Session::new("root", "model"))
1767 .await
1768 .unwrap();
1769 let mut stale = Session::new("nested-child-parent", "model");
1770 stale.kind = bamboo_domain::SessionKind::Child;
1771 stale.parent_session_id = Some("root".into());
1772 stale.root_session_id = "root".into();
1773 stale.spawn_depth = 1;
1774 let wait = WaitingForChildrenState::for_children(
1775 vec!["grandchild".into()],
1776 ChildWaitPolicy::All,
1777 chrono::Utc::now(),
1778 );
1779 let mut runtime = AgentRuntimeState::new("old-run");
1780 runtime.waiting_for_children = Some(wait.clone());
1781 stale.agent_runtime_state = Some(runtime);
1782 stale.add_message(bamboo_agent_core::Message::assistant("STALE", None));
1783 storage.save_session(&stale).await.unwrap();
1784
1785 let mut replacement = stale.clone();
1786 replacement.created_at += chrono::Duration::seconds(1);
1787 replacement.messages.clear();
1788 replacement.add_message(bamboo_agent_core::Message::assistant("REPLACEMENT", None));
1789 replacement
1790 .agent_runtime_state
1791 .as_mut()
1792 .unwrap()
1793 .waiting_for_children = None;
1794 storage.delete_session(&stale.id).await.unwrap();
1795 if !missing {
1796 storage.save_session(&replacement).await.unwrap();
1797 }
1798 if !missing {
1799 cache.insert(
1800 replacement.id.clone(),
1801 Arc::new(crate::SessionSnapshot::new(replacement.clone())),
1802 );
1803 }
1804
1805 let saved = save_finalized_runtime_with_inherited_child_wait(
1806 &persistence,
1807 &mut stale,
1808 Some(&wait),
1809 )
1810 .await;
1811 assert!(saved.is_err(), "stale incarnation must not be saved");
1812 publish_final_session_cache(&cache, stale, false, &saved);
1815 if missing {
1816 assert!(
1817 read_cached_session(&cache, &replacement.id).is_none(),
1818 "deleted incarnation must not be resurrected in cache"
1819 );
1820 } else {
1821 let cached = read_cached_session(&cache, &replacement.id).unwrap();
1822 assert_eq!(
1823 cached.created_at, replacement.created_at,
1824 "ABA rejection must retain replacement cache birth"
1825 );
1826 assert_eq!(
1827 serde_json::to_value(&cached.messages).unwrap(),
1828 serde_json::to_value(&replacement.messages).unwrap(),
1829 "ABA rejection must retain replacement cache transcript"
1830 );
1831 }
1832 assert!(saved
1833 .as_ref()
1834 .unwrap_err()
1835 .get_ref()
1836 .is_some_and(|cause| cause.is::<bamboo_domain::SessionAuthorityConflict>()));
1837 let durable = storage.load_session(&replacement.id).await.unwrap();
1838 if missing {
1839 assert!(durable.is_none());
1840 } else {
1841 let durable = durable.unwrap();
1842 assert_eq!(durable.created_at, replacement.created_at);
1843 assert_eq!(
1844 serde_json::to_value(&durable.messages).unwrap(),
1845 serde_json::to_value(&replacement.messages).unwrap()
1846 );
1847 }
1848 }
1849 }
1850
1851 #[test]
1852 fn debug_prompt_snapshot_migrates_legacy_host_paths_before_logging() {
1853 let legacy_project = format!(
1854 "{}\nProject ID: legacy-project\nProject path: /private/legacy-project\nProject home: /private/legacy-home\n{}",
1855 crate::runtime::context::PROJECT_CONTEXT_START_MARKER,
1856 crate::runtime::context::PROJECT_CONTEXT_END_MARKER,
1857 );
1858 let legacy_workspace =
1859 crate::runtime::context::build_workspace_prompt_context("/private/legacy-workspace")
1860 .expect("legacy workspace context");
1861 let legacy_instruction = format!(
1862 "{}\nSource: /private/legacy-workspace/AGENTS.md\nlegacy policy\n{}",
1863 crate::runtime::context::instruction::INSTRUCTION_CONTEXT_START_MARKER,
1864 crate::runtime::context::instruction::INSTRUCTION_CONTEXT_END_MARKER,
1865 );
1866 let mut session = Session::new("legacy-debug-log", "model");
1867 session.add_message(bamboo_agent_core::Message::system(format!(
1868 "Base prompt\n\n{legacy_project}\n\n{legacy_workspace}\n\n{legacy_instruction}"
1869 )));
1870
1871 let prompt = system_prompt_for_debug_log(&mut session).expect("normalized System");
1872
1873 assert_eq!(prompt, "Base prompt");
1874 assert_eq!(
1875 session.workspace_path_meta().as_deref(),
1876 Some("/private/legacy-workspace")
1877 );
1878 for private_path in [
1879 "/private/legacy-project",
1880 "/private/legacy-home",
1881 "/private/legacy-workspace",
1882 ] {
1883 assert!(!prompt.contains(private_path));
1884 }
1885 assert!(!prompt.contains(crate::runtime::context::PROJECT_CONTEXT_START_MARKER));
1886 assert!(!prompt.contains(crate::runtime::context::WORKSPACE_CONTEXT_START_MARKER));
1887 assert!(!prompt
1888 .contains(crate::runtime::context::instruction::INSTRUCTION_CONTEXT_START_MARKER));
1889 }
1890
1891 #[test]
1892 fn reservation_target_requires_domain_id_and_exact_runner_registry() {
1893 let runners = Arc::new(RwLock::new(HashMap::new()));
1894 let other_runners = Arc::new(RwLock::new(HashMap::new()));
1895 let mut reservation = SessionExecutionReservation {
1896 session_id: "session-a".to_string(),
1897 run_id: "run-a".to_string(),
1898 cancel_token: CancellationToken::new(),
1899 runners: runners.clone(),
1900 activation: SessionExecutionActivationOwnership::Unrouted,
1901 root_actor: None,
1902 execution_persistence: None,
1903 persistence_bound: false,
1904 selected_input: false,
1905 armed: true,
1906 };
1907
1908 assert!(reservation.matches_execution_target("session-a", "session-a", &runners));
1909 assert!(!reservation.matches_execution_target("session-b", "session-a", &runners));
1910 assert!(!reservation.matches_execution_target("session-a", "session-b", &runners));
1911 assert!(!reservation.matches_execution_target("session-a", "session-a", &other_runners));
1912 reservation.armed = false;
1913 }
1914
1915 #[tokio::test]
1916 async fn forced_same_run_registration_rejection_never_cancels_live_owner() {
1917 let runners = Arc::new(RwLock::new(HashMap::new()));
1918 let mut runner = AgentRunner::new();
1919 runner.status = AgentStatus::Running;
1920 let run_id = runner.run_id.clone();
1921 let live_cancel_token = runner.cancel_token.clone();
1922 runners.write().await.insert("same-run".to_string(), runner);
1923
1924 let router = SessionActivationRouter::new();
1925 let mut live_registration = router
1926 .register_run("same-run", &run_id)
1927 .await
1928 .expect("first registration owns the run");
1929 let mut rejected = SessionExecutionReservation {
1930 session_id: "same-run".to_string(),
1931 run_id: run_id.clone(),
1932 cancel_token: live_cancel_token.clone(),
1933 runners: runners.clone(),
1934 activation: SessionExecutionActivationOwnership::RegistrationPending(router.clone()),
1935 root_actor: None,
1936 execution_persistence: None,
1937 persistence_bound: false,
1938 selected_input: false,
1939 armed: true,
1940 };
1941
1942 let error = match rejected.ensure_registered().await {
1943 Ok(()) => panic!("a duplicate registration for the same run must be rejected"),
1944 Err(error) => error,
1945 };
1946 assert_eq!(error.existing_run_id(), run_id);
1947 drop(rejected);
1948
1949 assert!(!live_cancel_token.is_cancelled());
1950 assert!(matches!(
1951 runners
1952 .read()
1953 .await
1954 .get("same-run")
1955 .map(|runner| &runner.status),
1956 Some(AgentStatus::Running)
1957 ));
1958 assert!(router.owns_run("same-run", &run_id).await);
1959 live_registration.begin_finalization().await;
1960 assert_eq!(live_registration.finish(0).await.unwrap(), None);
1961 }
1962}