1use std::sync::Arc;
16
17use bamboo_agent_core::storage::Storage;
18use bamboo_agent_core::Session;
19use bamboo_domain::RetrievalWindowCheckpointOutcome;
20use bamboo_storage::LockedSessionStore;
21
22use crate::{read_cached_session, SessionCache};
23
24#[path = "session_repository_root_inbox.rs"]
25mod root_inbox;
26
27#[path = "session_repository_parent_outcome.rs"]
28mod parent_outcome;
29pub use parent_outcome::{ParentOutcomeRoute, ParentOutcomeWriter};
30
31#[cfg(test)]
32type PostDurableHook = Arc<dyn Fn(&str, &str) + Send + Sync>;
33
34#[derive(Clone)]
37pub struct SessionRepository {
38 cache: SessionCache,
39 storage: Arc<dyn Storage>,
40 persistence: Arc<LockedSessionStore>,
41 root_actor_directory: Option<Arc<dyn bamboo_domain::ActorDirectoryPort>>,
42 root_actor_owner: Option<(
43 Arc<dyn bamboo_domain::ActorDirectoryPort>,
44 Arc<bamboo_domain::RootActorRuntimeWrite>,
45 )>,
46 root_actor_writers:
47 Arc<dashmap::DashMap<String, std::sync::Weak<bamboo_domain::RootActorRuntimeWrite>>>,
48 #[cfg(test)]
49 post_durable_hook: Option<PostDurableHook>,
50}
51
52impl SessionRepository {
53 pub fn new(
54 cache: SessionCache,
55 storage: Arc<dyn Storage>,
56 persistence: Arc<LockedSessionStore>,
57 ) -> Self {
58 Self {
59 cache,
60 storage,
61 persistence,
62 root_actor_directory: None,
63 root_actor_owner: None,
64 root_actor_writers: Arc::new(Default::default()),
65 #[cfg(test)]
66 post_durable_hook: None,
67 }
68 }
69
70 async fn bind_root_actor_execution_inner(
71 &self,
72 session: &Session,
73 run_id: &str,
74 ) -> std::io::Result<Option<bamboo_domain::RootActorExecutionBinding>> {
75 if session.kind != bamboo_domain::SessionKind::Root
76 || !session.authority_identity.is_ordinary()
77 {
78 return Ok(None);
79 }
80 let Some(directory) = self.root_actor_directory.clone() else {
81 return Ok(None);
82 };
83 if !self.storage.supports_root_actor_runtime_write() {
84 return Err(std::io::Error::new(
85 std::io::ErrorKind::Unsupported,
86 "Root Actor execution requires a fenced storage writer",
87 ));
88 }
89 let Some(durable) = self.storage.load_root_authority(&session.id).await? else {
90 return Err(root_actor_binding_error(
91 "Root must be persisted before activation",
92 ));
93 };
94 if session.created_at != durable.created_at
95 || session.id != durable.id
96 || session.project_id_meta() != durable.project_id_meta()
97 || session.authority_identity != durable.authority_identity
98 {
99 return Err(root_actor_binding_error(
100 "Root identity changed before activation",
101 ));
102 }
103 let lease_duration = std::time::Duration::from_secs(15);
104 let now = chrono::Utc::now();
105 let claimed = directory
106 .claim_activation(&bamboo_domain::ActorActivationClaim {
107 actor_id: session.id.clone(),
108 run_id: run_id.to_string(),
109 lease_owner: format!("root-execution-{}", uuid::Uuid::new_v4()),
110 lease_expires_at: now + chrono::Duration::seconds(15),
111 inbox_generation: session
112 .session_inbox_admission()
113 .map_or(0, |a| a.last_admitted_sequence),
114 placement_ref: None,
115 now,
116 })
117 .await
118 .map_err(root_actor_binding_error)?;
119 let owner = bamboo_domain::RootActorRuntimeWrite {
120 fence: claimed.fence(),
121 created_at: durable.created_at,
122 };
123 let bound = self.bind_root_response_writer(owner.clone())?;
124 let live_owner = bound
125 .root_actor_owner
126 .as_ref()
127 .expect("bound Root owner")
128 .1
129 .clone();
130 let abandon_directory = directory.clone();
131 let abandon_fence = owner.fence.clone();
132 let binding = bamboo_domain::RootActorExecutionBinding::new(
133 Arc::new(bound),
134 directory.clone(),
135 owner,
136 lease_duration,
137 move || {
138 if let Ok(handle) = tokio::runtime::Handle::try_current() {
139 handle.spawn(async move {
140 let _ = abandon_directory
141 .finish_activation(
142 &abandon_fence,
143 chrono::Utc::now(),
144 bamboo_domain::ActorActivationFinish::Cancelled,
145 )
146 .await;
147 });
148 }
149 },
150 );
151 directory
152 .start_activation(&binding.owner.fence, chrono::Utc::now())
153 .await
154 .map_err(root_actor_binding_error)?;
155 self.root_actor_writers
158 .insert(session.id.clone(), Arc::downgrade(&live_owner));
159 Ok(Some(binding))
160 }
161
162 pub(crate) fn bind_root_response_writer(
165 &self,
166 owner: bamboo_domain::RootActorRuntimeWrite,
167 ) -> std::io::Result<Self> {
168 let directory = self.root_actor_directory.clone().ok_or_else(|| {
169 root_actor_binding_error("Root response requires its Host Actor directory")
170 })?;
171 let host_cache = self.cache.clone();
172 let publisher: bamboo_domain::RootActorRuntimePublisher = Arc::new(move |saved| {
173 host_cache.insert(
174 saved.id.clone(),
175 Arc::new(crate::SessionSnapshot::new(saved.clone())),
176 );
177 });
178 let bound_store = self
179 .persistence
180 .bind_root_actor_writer(owner.clone(), publisher)?;
181 let mut bound = Self::new(
184 Arc::new(Default::default()),
185 self.storage.clone(),
186 Arc::new(bound_store),
187 );
188 bound.root_actor_owner = Some((directory.clone(), Arc::new(owner)));
189 bound.root_actor_writers = self.root_actor_writers.clone();
190 Ok(bound)
191 }
192
193 pub fn with_root_actor_directory(
196 mut self,
197 directory: Arc<dyn bamboo_domain::ActorDirectoryPort>,
198 ) -> Self {
199 self.root_actor_directory = Some(directory);
200 self
201 }
202
203 #[cfg(test)]
204 fn with_post_durable_hook(mut self, hook: PostDurableHook) -> Self {
205 self.post_durable_hook = Some(hook);
206 self
207 }
208
209 #[cfg(test)]
210 fn run_post_durable_hook(&self, operation: &str, marker: &str) {
211 if let Some(hook) = self.post_durable_hook.as_ref() {
212 hook(operation, marker);
213 }
214 }
215
216 pub fn cache(&self) -> &SessionCache {
217 &self.cache
218 }
219
220 pub fn storage(&self) -> &Arc<dyn Storage> {
221 &self.storage
222 }
223
224 pub fn persistence(&self) -> &Arc<LockedSessionStore> {
225 &self.persistence
226 }
227
228 pub async fn load(&self, session_id: &str) -> Option<Session> {
231 if let Some(session) = read_cached_session(&self.cache, session_id) {
232 return Some(session);
233 }
234
235 let _guard = self.persistence.acquire_lock(session_id).await;
236 if let Some(session) = read_cached_session(&self.cache, session_id) {
237 return Some(session);
238 }
239
240 let loaded = self.storage.load_session(session_id).await;
241 #[cfg(test)]
242 self.run_post_durable_hook("load", session_id);
243 match loaded {
244 Ok(Some(session)) => {
245 self.cache.insert(
246 session_id.to_string(),
247 Arc::new(crate::SessionSnapshot::new(session.clone())),
248 );
249 Some(session)
250 }
251 _ => None,
252 }
253 }
254
255 pub async fn try_load(&self, session_id: &str) -> std::io::Result<Option<Session>> {
259 if let Some(session) = read_cached_session(&self.cache, session_id) {
260 return Ok(Some(session));
261 }
262
263 let _guard = self.persistence.acquire_lock(session_id).await;
264 if let Some(session) = read_cached_session(&self.cache, session_id) {
265 return Ok(Some(session));
266 }
267
268 let loaded = self.storage.load_session(session_id).await?;
269 #[cfg(test)]
270 self.run_post_durable_hook("try_load", session_id);
271 if let Some(ref session) = loaded {
272 self.cache.insert(
273 session_id.to_string(),
274 Arc::new(crate::SessionSnapshot::new(session.clone())),
275 );
276 }
277 Ok(loaded)
278 }
279
280 pub async fn save(&self, session: &mut Session) -> std::io::Result<()> {
284 self.persistence
285 .merge_save_runtime_and_publish(session, |saved, committed| {
286 if committed {
287 #[cfg(test)]
288 self.run_post_durable_hook("save_full", &saved.id);
289 self.cache.insert(
290 saved.id.clone(),
291 Arc::new(crate::SessionSnapshot::new(saved.clone())),
292 );
293 }
294 })
295 .await
296 }
297
298 pub async fn save_replay_resolution(&self, session: &mut Session) -> std::io::Result<()> {
303 self.persistence
304 .merge_save_runtime_and_publish(session, |saved, committed| {
305 if committed {
306 self.cache.insert(
307 saved.id.clone(),
308 Arc::new(crate::SessionSnapshot::new(saved.clone())),
309 );
310 } else {
311 self.cache.remove(&saved.id);
312 }
313 })
314 .await
315 }
316
317 pub async fn update_runtime_session<F>(
321 &self,
322 session_id: &str,
323 metadata_keys: &[&str],
324 mutate: F,
325 ) -> std::io::Result<Option<Session>>
326 where
327 F: FnOnce(&mut Session),
328 {
329 self.persistence
330 .update_runtime_config_and_publish(session_id, mutate, |saved| {
331 if let Some(cached) = self.cache.get(session_id) {
332 cached.update(|cached| {
333 for key in metadata_keys {
334 if let Some(value) = saved.metadata.get(*key) {
335 cached.metadata.insert((*key).to_string(), value.clone());
336 } else {
337 cached.metadata.remove(*key);
338 }
339 }
340 });
341 }
342 })
343 .await
344 }
345
346 pub async fn load_or_create(&self, session_id: &str, model: &str) -> Session {
348 if let Some(session) = self.load(session_id).await {
349 return session;
350 }
351 Session::new(session_id.to_string(), model.to_string())
352 }
353
354 pub async fn load_merged_checked(&self, session_id: &str) -> std::io::Result<Option<Session>> {
362 let _guard = self.persistence.acquire_lock(session_id).await;
363 let memory_session = read_cached_session(&self.cache, session_id);
364 let storage_session = self.storage.load_session(session_id).await?;
365 #[cfg(test)]
366 self.run_post_durable_hook("load_merged", session_id);
367
368 Ok(match (memory_session, storage_session) {
369 (Some(memory), Some(storage)) => {
370 let root_authority_preference =
371 root_authority_storage_preference(&memory, &storage);
372 let prefer_storage = should_prefer_storage(&memory, &storage);
373 let diverged = prefer_storage || memory.messages.len() != storage.messages.len();
374 let chosen_len = if prefer_storage {
375 storage.messages.len()
376 } else {
377 memory.messages.len()
378 };
379 macro_rules! merged_log {
380 ($level:ident) => {
381 tracing::$level!(
382 "[{}] load_session_merged: memory={} msgs (updated_at={}), storage={} msgs (updated_at={}), prefer_storage={} -> chose {} msgs",
383 session_id,
384 memory.messages.len(),
385 memory.updated_at,
386 storage.messages.len(),
387 storage.updated_at,
388 prefer_storage,
389 chosen_len,
390 )
391 };
392 }
393 if diverged {
394 merged_log!(debug);
395 } else {
396 merged_log!(trace);
397 }
398 let memory_updated_at = memory.updated_at;
399 let chosen = if prefer_storage { storage } else { memory };
400 if prefer_storage
408 && (root_authority_preference == Some(true)
409 || chosen.updated_at >= memory_updated_at)
410 {
411 self.cache.insert(
412 session_id.to_string(),
413 Arc::new(crate::SessionSnapshot::new(chosen.clone())),
414 );
415 }
416 Some(chosen)
417 }
418 (Some(memory), None) => Some(memory),
419 (None, Some(storage)) => {
420 self.cache.insert(
421 session_id.to_string(),
422 Arc::new(crate::SessionSnapshot::new(storage.clone())),
423 );
424 Some(storage)
425 }
426 (None, None) => None,
427 })
428 }
429
430 pub async fn load_merged(&self, session_id: &str) -> Option<Session> {
434 match self.load_merged_checked(session_id).await {
435 Ok(session) => session,
436 Err(error) => {
437 tracing::warn!(
438 "[{}] Failed to load merged session from storage: {}",
439 session_id,
440 error
441 );
442 read_cached_session(&self.cache, session_id)
443 }
444 }
445 }
446
447 pub async fn save_and_cache(&self, session: &mut Session) {
450 let result = self
451 .persistence
452 .merge_save_runtime_and_publish(session, |saved, _| {
453 #[cfg(test)]
454 self.run_post_durable_hook("save_and_cache", &saved.id);
455 self.cache.insert(
456 saved.id.clone(),
457 Arc::new(crate::SessionSnapshot::new(saved.clone())),
458 );
459 })
460 .await;
461 if let Err(error) = result {
462 tracing::warn!("[{}] Failed to save session: {}", session.id, error);
463 }
464 }
465
466 async fn refresh_cached_task_control_plane(&self, session_id: &str) {
467 match self.storage.load_runtime_control_plane(session_id).await {
468 Ok(Some(durable)) => {
469 if let Some(cached) = self.cache.get(session_id) {
470 cached.update(|cached| adopt_task_control_plane(cached, &durable));
471 }
472 }
473 Ok(None) => {}
474 Err(error) => tracing::warn!(
475 "[{}] Failed to refresh Task control-plane after write conflict: {}",
476 session_id,
477 error
478 ),
479 }
480 }
481}
482
483fn adopt_task_control_plane(target: &mut Session, durable: &Session) {
484 target.task_list = durable.task_list.clone();
485 target
486 .metadata
487 .remove(bamboo_domain::session::runtime_metadata::keys::TASK_LIST_VERSION);
488 if let Some(runtime_metadata) = target.runtime_metadata.as_mut() {
489 runtime_metadata.task_list_version = None;
490 }
491 if target
492 .runtime_metadata
493 .as_ref()
494 .is_some_and(bamboo_domain::session::SessionRuntimeMetadata::is_empty)
495 {
496 target.runtime_metadata = None;
497 }
498 if let Some(version) = durable.task_list_version_meta() {
499 target.set_task_list_version_meta(version);
500 }
501}
502
503fn should_prefer_storage(memory_session: &Session, storage_session: &Session) -> bool {
504 if let Some(prefer_storage) = root_authority_storage_preference(memory_session, storage_session)
505 {
506 return prefer_storage;
507 }
508 if storage_session.updated_at < memory_session.updated_at {
513 return false;
514 }
515 storage_session.updated_at > memory_session.updated_at
519 || (memory_session.pending_question.is_none() && storage_session.pending_question.is_some())
520}
521
522fn root_authority_storage_preference(
527 memory_session: &Session,
528 storage_session: &Session,
529) -> Option<bool> {
530 use bamboo_domain::SessionKind;
531
532 if memory_session.kind != SessionKind::Root || storage_session.kind != SessionKind::Root {
533 return None;
534 }
535 if memory_session.created_at != storage_session.created_at
536 || memory_session.authority_identity != storage_session.authority_identity
537 {
538 return Some(true);
539 }
540 (memory_session.root_mode_transition_epoch != storage_session.root_mode_transition_epoch)
541 .then_some(
542 storage_session.root_mode_transition_epoch > memory_session.root_mode_transition_epoch,
543 )
544}
545
546#[async_trait::async_trait]
550impl bamboo_domain::RuntimeSessionPersistence for SessionRepository {
551 fn inherited_child_wait(&self) -> Option<bamboo_domain::InheritedChildWait> {
552 bamboo_domain::RuntimeSessionPersistence::inherited_child_wait(self.persistence.as_ref())
553 }
554
555 fn bind_inherited_child_wait(
556 &self,
557 inherited: bamboo_domain::InheritedChildWait,
558 ) -> std::io::Result<Arc<dyn bamboo_domain::RuntimeSessionPersistence>> {
559 let mut bound = self.clone();
560 bound.persistence = Arc::new(
561 self.persistence
562 .as_ref()
563 .bind_inherited_child_wait(inherited),
564 );
565 Ok(Arc::new(bound))
566 }
567
568 fn root_actor_writer(&self) -> Option<bamboo_domain::RootActorRuntimeWrite> {
569 self.root_actor_owner
570 .as_ref()
571 .map(|(_, owner)| (**owner).clone())
572 }
573
574 fn root_actor_execution_required(&self, session: &Session) -> bool {
575 self.root_actor_directory.is_some()
576 && session.kind == bamboo_domain::SessionKind::Root
577 && session.authority_identity.is_ordinary()
578 }
579
580 async fn bind_root_actor_execution(
581 &self,
582 session: &Session,
583 run_id: &str,
584 ) -> std::io::Result<Option<bamboo_domain::RootActorExecutionBinding>> {
585 let repository = self.clone();
586 let session = session.clone();
587 let run_id = run_id.to_string();
588 tokio::spawn(async move {
591 repository
592 .bind_root_actor_execution_inner(&session, &run_id)
593 .await
594 })
595 .await
596 .map_err(|error| root_actor_binding_error(format!("Root binding task failed: {error}")))?
597 }
598
599 async fn save_runtime_session(&self, session: &mut Session) -> std::io::Result<()> {
600 self.persistence
605 .merge_save_runtime_and_publish(session, |saved, _| {
606 #[cfg(test)]
607 self.run_post_durable_hook("save_runtime_session", &saved.id);
608 self.cache.insert(
609 saved.id.clone(),
610 Arc::new(crate::SessionSnapshot::new(saved.clone())),
611 );
612 })
613 .await
614 }
615
616 async fn save_finalized_runtime_session(&self, session: &mut Session) -> std::io::Result<()> {
617 self.persistence
618 .merge_save_finalized_runtime_and_publish(session, |saved, _| {
619 #[cfg(test)]
620 self.run_post_durable_hook("save_finalized_runtime_session", &saved.id);
621 self.cache.insert(
622 saved.id.clone(),
623 Arc::new(crate::SessionSnapshot::new(saved.clone())),
624 );
625 })
626 .await
627 }
628
629 async fn save_finalized_runtime_with_inherited_child_wait(
630 &self,
631 session: &mut Session,
632 inherited: &bamboo_domain::session::runtime_state::WaitingForChildrenState,
633 ) -> std::io::Result<()> {
634 self.persistence
635 .merge_save_inherited_child_wait_and_publish(session, inherited, |saved, _| {
636 self.cache.insert(
637 saved.id.clone(),
638 Arc::new(crate::SessionSnapshot::new(saved.clone())),
639 );
640 })
641 .await
642 }
643
644 async fn seed_runtime_activation(&self, session: &mut Session) -> std::io::Result<()> {
645 self.persistence
646 .seed_runtime_activation_and_publish(session, |saved, committed| {
647 #[cfg(test)]
648 self.run_post_durable_hook("seed_runtime_activation", &saved.id);
649 if committed {
650 self.cache.insert(
651 saved.id.clone(),
652 Arc::new(crate::SessionSnapshot::new(saved.clone())),
653 );
654 }
655 })
656 .await
657 }
658
659 async fn record_permission_posture_activation(
660 &self,
661 session_id: &str,
662 expected_audit_revision: Option<u64>,
663 seed: &bamboo_domain::PermissionAuditSeed,
664 ) -> std::io::Result<Option<Session>> {
665 self.persistence
666 .record_permission_posture_activation_and_publish(
667 session_id,
668 expected_audit_revision,
669 seed,
670 |saved| {
671 #[cfg(test)]
672 self.run_post_durable_hook("permission_posture_activation", &saved.id);
673 self.cache.insert(
674 saved.id.clone(),
675 Arc::new(crate::SessionSnapshot::new(saved.clone())),
676 );
677 },
678 )
679 .await
680 }
681
682 async fn save_runtime_control_plane(&self, session: &mut Session) -> std::io::Result<()> {
683 self.persistence
684 .save_runtime_only_and_publish(session, |saved| {
685 #[cfg(test)]
686 self.run_post_durable_hook("save", &saved.id);
687
688 if let Some(cached) = self.cache.get(&saved.id) {
696 cached.update(|cached| {
697 let messages = cached.messages.clone();
698 let provider_transcript = cached.provider_transcript.clone();
699 let admission = cached
700 .runtime_metadata
701 .as_ref()
702 .and_then(|metadata| metadata.session_inbox_admission.clone());
703 let mut refreshed = saved.clone();
704 refreshed.messages = messages;
705 refreshed.provider_transcript = provider_transcript;
706 if let Some(admission) = admission {
707 refreshed
708 .runtime_metadata
709 .get_or_insert_with(Default::default)
710 .session_inbox_admission = Some(admission);
711 } else if let Some(metadata) = refreshed.runtime_metadata.as_mut() {
712 metadata.session_inbox_admission = None;
713 }
714 *cached = refreshed;
715 });
716 }
717 })
718 .await
719 }
720
721 async fn load_runtime_control_plane(
722 &self,
723 session_id: &str,
724 ) -> std::io::Result<Option<Session>> {
725 bamboo_domain::RuntimeSessionPersistence::load_runtime_control_plane(
726 self.persistence.as_ref(),
727 session_id,
728 )
729 .await
730 }
731
732 async fn update_task_list_control_plane(
733 &self,
734 session_id: &str,
735 task_list: &bamboo_domain::TaskList,
736 version: &str,
737 ) -> std::io::Result<bool> {
738 let result = self
739 .persistence
740 .update_task_list_control_plane_and_publish(session_id, task_list, version, |_| {
741 #[cfg(test)]
742 self.run_post_durable_hook("task", version);
743
744 if let Some(cached) = self.cache.get(session_id) {
750 cached.update(|cached| {
751 cached.set_task_list(task_list.clone());
752 cached.set_task_list_version_meta(version.to_string());
753 });
754 }
755 })
756 .await;
757 if result
758 .as_ref()
759 .is_err_and(|error| error.kind() == std::io::ErrorKind::WouldBlock)
760 {
761 self.refresh_cached_task_control_plane(session_id).await;
766 }
767 result
768 }
769
770 async fn update_task_list_control_plane_if_version(
771 &self,
772 session_id: &str,
773 expected_version: &str,
774 expected_task_list: &bamboo_domain::TaskList,
775 task_list: &bamboo_domain::TaskList,
776 version: &str,
777 ) -> std::io::Result<bool> {
778 self.persistence
779 .update_task_list_control_plane_if_version_and_publish(
780 session_id,
781 expected_version,
782 expected_task_list,
783 task_list,
784 version,
785 |_| {
786 if let Some(cached) = self.cache.get(session_id) {
787 cached.update(|cached| {
788 cached.set_task_list(task_list.clone());
789 cached.set_task_list_version_meta(version.to_string());
790 });
791 }
792 },
793 )
794 .await
795 }
796
797 async fn update_task_list_control_planes_if_version(
798 &self,
799 session_id: &str,
800 shared_session_id: &str,
801 expected_version: &str,
802 expected_task_list: &bamboo_domain::TaskList,
803 task_list: &bamboo_domain::TaskList,
804 version: &str,
805 ) -> std::io::Result<bool> {
806 self.persistence
807 .update_task_list_control_planes_if_version_and_publish(
808 session_id,
809 shared_session_id,
810 expected_version,
811 expected_task_list,
812 task_list,
813 version,
814 |_, _| {
815 for id in [session_id, shared_session_id] {
816 if let Some(cached) = self.cache.get(id) {
817 cached.update(|cached| {
818 cached.set_task_list(task_list.clone());
819 cached.set_task_list_version_meta(version.to_string());
820 });
821 }
822 }
823 },
824 )
825 .await
826 }
827
828 async fn checkpoint_runtime_session(&self, session: &mut Session) -> std::io::Result<()> {
829 self.persistence
836 .checkpoint_runtime_session_and_publish(session, |saved, committed| {
837 #[cfg(test)]
838 self.run_post_durable_hook("checkpoint", &saved.id);
839 if committed {
840 self.cache.insert(
841 saved.id.clone(),
842 Arc::new(crate::SessionSnapshot::new(saved.clone())),
843 );
844 }
845 })
846 .await
847 }
848
849 async fn admit_root_inbox(
850 &self,
851 session: &mut Session,
852 inbox: Arc<dyn bamboo_domain::SessionInboxPort>,
853 active_run_id: Option<&str>,
854 ) -> std::io::Result<Option<bamboo_domain::RootInboxAdmission>> {
855 self.admit_owned_root_inbox(session, inbox, active_run_id)
856 .await
857 }
858
859 async fn checkpoint_retrieval_window(
860 &self,
861 expected_base: &Session,
862 staged: &mut Session,
863 ) -> std::io::Result<RetrievalWindowCheckpointOutcome> {
864 self.persistence
865 .checkpoint_retrieval_window_and_publish(expected_base, staged, |saved| {
866 #[cfg(test)]
867 self.run_post_durable_hook("checkpoint_retrieval_window", &saved.id);
868 self.cache.insert(
869 saved.id.clone(),
870 Arc::new(crate::SessionSnapshot::new(saved.clone())),
871 );
872 })
873 .await
874 }
875
876 async fn checkpoint_prompt_rewrite(
877 &self,
878 expected_base: &Session,
879 staged: &mut Session,
880 ) -> std::io::Result<RetrievalWindowCheckpointOutcome> {
881 self.persistence
882 .checkpoint_prompt_rewrite_and_publish(expected_base, staged, |saved| {
883 #[cfg(test)]
884 self.run_post_durable_hook("checkpoint_prompt_rewrite", &saved.id);
885 self.cache.insert(
886 saved.id.clone(),
887 Arc::new(crate::SessionSnapshot::new(saved.clone())),
888 );
889 })
890 .await
891 }
892
893 async fn checkpoint_manual_archive_rejection(
894 &self,
895 expected_base: &Session,
896 staged: &mut Session,
897 ) -> std::io::Result<RetrievalWindowCheckpointOutcome> {
898 self.persistence
899 .checkpoint_manual_archive_rejection_and_publish(expected_base, staged, |saved| {
900 #[cfg(test)]
901 self.run_post_durable_hook("checkpoint_manual_archive_rejection", &saved.id);
902 self.cache.insert(
903 saved.id.clone(),
904 Arc::new(crate::SessionSnapshot::new(saved.clone())),
905 );
906 })
907 .await
908 }
909
910 async fn checkpoint_manual_archive_consumption(
911 &self,
912 expected_base: &Session,
913 staged: &mut Session,
914 ) -> std::io::Result<RetrievalWindowCheckpointOutcome> {
915 self.persistence
916 .checkpoint_manual_archive_consumption_and_publish(expected_base, staged, |saved| {
917 #[cfg(test)]
918 self.run_post_durable_hook("checkpoint_manual_archive_consumption", &saved.id);
919 self.cache.insert(
920 saved.id.clone(),
921 Arc::new(crate::SessionSnapshot::new(saved.clone())),
922 );
923 })
924 .await
925 }
926
927 async fn load_runtime_session(&self, session_id: &str) -> std::io::Result<Option<Session>> {
928 if let Some((directory, owner)) = self.root_actor_owner.as_ref() {
929 if session_id != owner.fence.actor_id {
930 return Err(root_actor_binding_error(
931 "bound execution cannot load another Root",
932 ));
933 }
934 directory
935 .validate_fence(&owner.fence, chrono::Utc::now())
936 .await
937 .map_err(root_actor_binding_error)?;
938 return self.storage.load_session(session_id).await;
939 }
940 self.try_load(session_id).await
941 }
942
943 async fn clear_legacy_pending_messages(
944 &self,
945 session_id: &str,
946 expected: &[serde_json::Value],
947 ) -> std::io::Result<bool> {
948 self.persistence
953 .clear_legacy_pending_messages_and_publish(session_id, expected, |latest| {
954 #[cfg(test)]
955 self.run_post_durable_hook("clear_legacy", session_id);
956 self.cache.insert(
957 session_id.to_string(),
958 Arc::new(crate::SessionSnapshot::new(latest.clone())),
959 );
960 })
961 .await
962 }
963
964 async fn append_token_usage_record(
965 &self,
966 session_id: &str,
967 json_line: &str,
968 ) -> std::io::Result<()> {
969 self.storage
970 .append_token_usage_record(session_id, json_line)
971 .await
972 }
973}
974
975#[cfg(test)]
976mod tests {
977 use super::*;
978 use bamboo_agent_core::storage::Storage;
979 use chrono::Utc;
980 use std::collections::HashMap;
981 use std::sync::atomic::{AtomicBool, Ordering};
982 use std::sync::{Condvar, Mutex};
983 use std::time::Duration;
984
985 #[derive(Default)]
986 struct MapStorage {
987 sessions: Mutex<HashMap<String, Session>>,
988 fail_pair_commit: AtomicBool,
989 }
990
991 struct FailingSaveStorage {
992 persisted: Mutex<Option<Session>>,
993 }
994
995 struct AmbiguousReplaySaveStorage {
996 persisted: Mutex<Option<Session>>,
997 }
998
999 #[async_trait::async_trait]
1000 impl Storage for MapStorage {
1001 async fn save_session(&self, session: &Session) -> std::io::Result<()> {
1002 self.sessions
1003 .lock()
1004 .unwrap()
1005 .insert(session.id.clone(), session.clone());
1006 Ok(())
1007 }
1008 async fn load_session(&self, session_id: &str) -> std::io::Result<Option<Session>> {
1009 Ok(self.sessions.lock().unwrap().get(session_id).cloned())
1010 }
1011 async fn delete_session(&self, session_id: &str) -> std::io::Result<bool> {
1012 Ok(self.sessions.lock().unwrap().remove(session_id).is_some())
1013 }
1014
1015 async fn save_task_control_plane_if_matches(
1016 &self,
1017 original: &Session,
1018 updated: &Session,
1019 ) -> std::io::Result<bool> {
1020 let mut sessions = self.sessions.lock().unwrap();
1021 let matches_original = sessions.get(&original.id).is_some_and(|current| {
1022 current.task_list_version_meta() == original.task_list_version_meta()
1023 && serde_json::to_value(¤t.task_list).ok()
1024 == serde_json::to_value(&original.task_list).ok()
1025 });
1026 if !matches_original {
1027 return Ok(false);
1028 }
1029 let Some(current) = sessions.get(&original.id).cloned() else {
1030 return Ok(false);
1031 };
1032 let mut committed = current;
1033 committed.task_list = updated.task_list.clone();
1034 if let Some(version) = updated.task_list_version_meta() {
1035 committed.set_task_list_version_meta(version);
1036 }
1037 sessions.insert(committed.id.clone(), committed);
1038 Ok(true)
1039 }
1040
1041 async fn save_task_control_planes_atomically(
1042 &self,
1043 first_original: &Session,
1044 first_updated: &Session,
1045 second_original: &Session,
1046 second_updated: &Session,
1047 ) -> std::io::Result<bool> {
1048 if self.fail_pair_commit.swap(false, Ordering::SeqCst) {
1049 return Err(std::io::Error::other(
1050 "injected paired Task transaction failure",
1051 ));
1052 }
1053 let mut sessions = self.sessions.lock().unwrap();
1054 let matches_original = |current: Option<&Session>, original: &Session| {
1055 current.is_some_and(|current| {
1056 current.task_list_version_meta() == original.task_list_version_meta()
1057 && serde_json::to_value(¤t.task_list).ok()
1058 == serde_json::to_value(&original.task_list).ok()
1059 })
1060 };
1061 if !matches_original(sessions.get(&first_original.id), first_original)
1062 || !matches_original(sessions.get(&second_original.id), second_original)
1063 {
1064 return Ok(false);
1065 }
1066 sessions.insert(first_updated.id.clone(), first_updated.clone());
1067 sessions.insert(second_updated.id.clone(), second_updated.clone());
1068 Ok(true)
1069 }
1070 }
1071
1072 #[async_trait::async_trait]
1073 impl Storage for FailingSaveStorage {
1074 async fn save_session(&self, _session: &Session) -> std::io::Result<()> {
1075 Err(std::io::Error::other("injected save failure"))
1076 }
1077
1078 async fn load_session(&self, _session_id: &str) -> std::io::Result<Option<Session>> {
1079 Ok(self.persisted.lock().unwrap().clone())
1080 }
1081
1082 async fn delete_session(&self, _session_id: &str) -> std::io::Result<bool> {
1083 Ok(false)
1084 }
1085 }
1086
1087 #[async_trait::async_trait]
1088 impl Storage for AmbiguousReplaySaveStorage {
1089 async fn save_session(&self, session: &Session) -> std::io::Result<()> {
1090 *self.persisted.lock().unwrap() = Some(session.clone());
1091 Err(std::io::Error::other(
1092 "injected post-commit replay save error",
1093 ))
1094 }
1095
1096 async fn load_session(&self, _session_id: &str) -> std::io::Result<Option<Session>> {
1097 Ok(self.persisted.lock().unwrap().clone())
1098 }
1099
1100 async fn delete_session(&self, _session_id: &str) -> std::io::Result<bool> {
1101 Ok(false)
1102 }
1103 }
1104
1105 fn test_repo(storage: Arc<dyn Storage>) -> SessionRepository {
1106 let cache: SessionCache = Arc::default();
1107 let persistence = Arc::new(LockedSessionStore::new(storage.clone()));
1108 SessionRepository::new(cache, storage, persistence)
1109 }
1110
1111 fn cache_put(repo: &SessionRepository, session: &Session) {
1112 repo.cache().insert(
1113 session.id.clone(),
1114 Arc::new(crate::SessionSnapshot::new(session.clone())),
1115 );
1116 }
1117
1118 #[tokio::test]
1119 async fn merged_root_prefers_live_recreated_birth_even_when_epoch_resets() {
1120 let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
1121 let repo = test_repo(storage.clone());
1122 let id = "root-recreated-same-id";
1123 let mut stale = Session::new(id, "stale");
1124 stale.root_mode_transition_epoch = 3;
1125 stale.updated_at += chrono::Duration::hours(1);
1126 cache_put(&repo, &stale);
1127
1128 let mut recreated = Session::new(id, "recreated");
1129 recreated.created_at = stale.created_at + chrono::Duration::seconds(1);
1130 recreated.updated_at = stale.created_at + chrono::Duration::seconds(1);
1131 storage.save_session(&recreated).await.unwrap();
1132
1133 let loaded = repo.load_merged_checked(id).await.unwrap().unwrap();
1134 assert_eq!(loaded.model, "recreated");
1135 assert_eq!(loaded.root_mode_transition_epoch, 0);
1136 assert_eq!(
1137 loaded.root_mode_birth_token(),
1138 recreated.root_mode_birth_token()
1139 );
1140 let cached = read_cached_session(repo.cache(), id).expect("new birth cached");
1141 assert_eq!(cached.model, "recreated");
1142 assert_eq!(
1143 cached.root_mode_birth_token(),
1144 recreated.root_mode_birth_token()
1145 );
1146 }
1147
1148 #[tokio::test]
1149 async fn merged_root_caches_higher_durable_epoch_despite_clock_skew() {
1150 let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
1151 let repo = test_repo(storage.clone());
1152 let id = "root-mode-clock-skew";
1153 let mut stale = Session::new(id, "model");
1154 stale.updated_at += chrono::Duration::hours(1);
1155 cache_put(&repo, &stale);
1156
1157 let mut durable = stale.clone();
1158 durable.root_mode_transition_epoch = 1;
1159 durable.root_orchestration_only = true;
1160 durable.root_tool_authority_revision = 1;
1161 durable.updated_at = durable.created_at;
1162 storage.save_session(&durable).await.unwrap();
1163
1164 let loaded = repo.load_merged_checked(id).await.unwrap().unwrap();
1165 assert_eq!(loaded.root_mode_transition_epoch, 1);
1166 assert!(loaded.root_orchestration_only);
1167 let cached = read_cached_session(repo.cache(), id).expect("higher epoch cached");
1168 assert_eq!(cached.root_mode_transition_epoch, 1);
1169 assert!(cached.root_orchestration_only);
1170 }
1171
1172 #[tokio::test]
1173 async fn runtime_merge_save_adopts_root_mode_terminal_tuple_even_for_policy_neutral_fence() {
1174 use bamboo_domain::{RootModeOperationAction, RootModeOperationRequest};
1175
1176 for action in [
1177 RootModeOperationAction::Select,
1178 RootModeOperationAction::Recover,
1179 ] {
1180 let home = tempfile::tempdir().unwrap();
1181 let concrete = Arc::new(
1182 bamboo_storage::SessionStoreV2::new(home.path().into())
1183 .await
1184 .unwrap(),
1185 );
1186 let storage: Arc<dyn Storage> = concrete.clone();
1187 let repo = test_repo(storage.clone());
1188 let mut running = Session::new("runtime-root-mode-merge", "model");
1189 storage.save_session(&running).await.unwrap();
1190 cache_put(&repo, &running);
1191 let request = RootModeOperationRequest {
1192 session_id: running.id.clone(),
1193 operation_id: format!("0:{}", uuid::Uuid::new_v4()),
1194 birth_token: running.root_mode_birth_token(),
1195 expected_epoch: 0,
1196 requested_enabled: true,
1197 action,
1198 };
1199 storage.root_mode_operation(&request).await.unwrap();
1200
1201 running.add_message(bamboo_agent_core::Message::assistant("run result", None));
1204 repo.save(&mut running).await.unwrap();
1205 assert_eq!(running.root_mode_transition_epoch, 1);
1206 assert_eq!(running.root_mode_operations.len(), 1);
1207 assert_eq!(
1208 running.root_orchestration_only_enabled(),
1209 action == RootModeOperationAction::Select
1210 );
1211 let durable = storage.load_session(&running.id).await.unwrap().unwrap();
1212 assert_eq!(durable.root_mode_transition_epoch, 1);
1213 assert_eq!(durable.root_mode_operations, running.root_mode_operations);
1214 assert_eq!(durable.messages.last().unwrap().content, "run result");
1215 concrete.flush_search_index().await;
1216 }
1217 }
1218
1219 fn task_list(session_id: &str, title: &str) -> bamboo_domain::TaskList {
1220 let now = Utc::now();
1221 bamboo_domain::TaskList {
1222 session_id: session_id.to_string(),
1223 title: title.to_string(),
1224 items: Vec::new(),
1225 created_at: now,
1226 updated_at: now,
1227 }
1228 }
1229
1230 fn staged_retrieval_window_archive(expected: &Session) -> Session {
1231 let mut staged = expected.clone();
1232 let mut event = bamboo_domain::CompressionEvent::new(
1233 1,
1234 1,
1235 80.0,
1236 60.0,
1237 0,
1238 bamboo_domain::CompressionTriggerType::Auto,
1239 0.0,
1240 None,
1241 0,
1242 );
1243 event.kind = bamboo_domain::CompressionEventKind::RetrievalWindow;
1244 staged.messages[0].compressed = true;
1245 staged.messages[0].compressed_by_event_id = Some(event.id.clone());
1246 staged.compression_events.push(event);
1247 staged.reset_model_context_epoch(bamboo_domain::ModelContextResetReason::Compression);
1248 staged
1249 }
1250
1251 #[tokio::test]
1252 async fn paired_task_cas_narrowly_updates_child_and_root_cache() {
1253 let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
1254 let repo = test_repo(storage.clone());
1255 let root_id = "paired-cache-root";
1256 let child_id = "paired-cache-child";
1257 let expected_task_list = task_list(root_id, "old shared");
1258
1259 let mut root = Session::new(root_id, "model");
1260 root.add_message(bamboo_agent_core::Message::user("root transcript"));
1261 root.metadata
1262 .insert("unrelated.root".to_string(), "keep".to_string());
1263 root.set_task_list(expected_task_list.clone());
1264 root.set_task_list_version_meta("1");
1265 storage.save_session(&root).await.unwrap();
1266 cache_put(&repo, &root);
1267
1268 let mut child = Session::new_child(child_id, root_id, "model", "child");
1269 child.add_message(bamboo_agent_core::Message::user("child transcript"));
1270 child
1271 .metadata
1272 .insert("unrelated.child".to_string(), "keep".to_string());
1273 child.set_task_list(expected_task_list.clone());
1274 child.set_task_list_version_meta("1");
1275 storage.save_session(&child).await.unwrap();
1276 cache_put(&repo, &child);
1277
1278 assert!(
1279 bamboo_domain::RuntimeSessionPersistence::update_task_list_control_planes_if_version(
1280 &repo,
1281 child_id,
1282 root_id,
1283 "1",
1284 &expected_task_list,
1285 &task_list(root_id, "evaluated"),
1286 "2",
1287 )
1288 .await
1289 .expect("paired cache CAS succeeds")
1290 );
1291
1292 let cached_root = read_cached_session(repo.cache(), root_id).expect("cached root");
1293 let cached_child = read_cached_session(repo.cache(), child_id).expect("cached child");
1294 for (session, transcript, metadata_key) in [
1295 (&cached_root, "root transcript", "unrelated.root"),
1296 (&cached_child, "child transcript", "unrelated.child"),
1297 ] {
1298 assert_eq!(session.task_list_version_meta().as_deref(), Some("2"));
1299 assert_eq!(
1300 session.task_list.as_ref().map(|list| list.title.as_str()),
1301 Some("evaluated")
1302 );
1303 assert_eq!(session.messages.len(), 1);
1304 assert_eq!(session.messages[0].content, transcript);
1305 assert_eq!(
1306 session.metadata.get(metadata_key).map(String::as_str),
1307 Some("keep")
1308 );
1309 }
1310 }
1311
1312 #[tokio::test]
1313 async fn failed_paired_task_transaction_does_not_publish_child_or_root_cache() {
1314 let concrete = Arc::new(MapStorage::default());
1315 let storage: Arc<dyn Storage> = concrete.clone();
1316 let repo = test_repo(storage.clone());
1317 let root_id = "paired-cache-failure-root";
1318 let child_id = "paired-cache-failure-child";
1319 let expected_task_list = task_list(root_id, "old shared");
1320
1321 let mut root = Session::new(root_id, "model");
1322 root.add_message(bamboo_agent_core::Message::user("root transcript"));
1323 root.metadata
1324 .insert("unrelated.root".to_string(), "keep".to_string());
1325 root.set_task_list(expected_task_list.clone());
1326 root.set_task_list_version_meta("1");
1327 storage.save_session(&root).await.unwrap();
1328 cache_put(&repo, &root);
1329
1330 let mut child = Session::new_child(child_id, root_id, "model", "child");
1331 child.add_message(bamboo_agent_core::Message::user("child transcript"));
1332 child
1333 .metadata
1334 .insert("unrelated.child".to_string(), "keep".to_string());
1335 child.set_task_list(expected_task_list.clone());
1336 child.set_task_list_version_meta("1");
1337 storage.save_session(&child).await.unwrap();
1338 cache_put(&repo, &child);
1339
1340 concrete.fail_pair_commit.store(true, Ordering::SeqCst);
1341 let error =
1342 bamboo_domain::RuntimeSessionPersistence::update_task_list_control_planes_if_version(
1343 &repo,
1344 child_id,
1345 root_id,
1346 "1",
1347 &expected_task_list,
1348 &task_list(root_id, "must not publish"),
1349 "2",
1350 )
1351 .await
1352 .expect_err("paired durable transaction fails");
1353 assert!(error.to_string().contains("injected paired"));
1354
1355 for (id, expected_title, transcript, metadata_key) in [
1356 (root_id, "old shared", "root transcript", "unrelated.root"),
1357 (
1358 child_id,
1359 "old shared",
1360 "child transcript",
1361 "unrelated.child",
1362 ),
1363 ] {
1364 let cached = read_cached_session(repo.cache(), id).expect("cached session remains");
1365 assert_eq!(cached.task_list_version_meta().as_deref(), Some("1"));
1366 assert_eq!(
1367 cached.task_list.as_ref().map(|list| list.title.as_str()),
1368 Some(expected_title)
1369 );
1370 assert_eq!(cached.messages[0].content, transcript);
1371 assert_eq!(
1372 cached.metadata.get(metadata_key).map(String::as_str),
1373 Some("keep")
1374 );
1375
1376 let durable = storage.load_session(id).await.unwrap().unwrap();
1377 assert_eq!(durable.task_list_version_meta().as_deref(), Some("1"));
1378 assert_eq!(
1379 durable.task_list.as_ref().map(|list| list.title.as_str()),
1380 Some(expected_title)
1381 );
1382 }
1383 }
1384
1385 fn durable_cache_fence(
1386 operation: impl Into<String>,
1387 marker: impl Into<String>,
1388 ) -> (
1389 PostDurableHook,
1390 tokio::sync::oneshot::Receiver<()>,
1391 Arc<(Mutex<bool>, Condvar)>,
1392 ) {
1393 let operation = operation.into();
1394 let marker = marker.into();
1395 let (started_tx, started_rx) = tokio::sync::oneshot::channel();
1396 let started_tx = Arc::new(Mutex::new(Some(started_tx)));
1397 let release = Arc::new((Mutex::new(false), Condvar::new()));
1398 let hook_release = release.clone();
1399 let hook: PostDurableHook = Arc::new(move |actual_operation, actual_marker| {
1400 if actual_operation != operation || actual_marker != marker {
1401 return;
1402 }
1403 if let Some(started_tx) = started_tx.lock().unwrap().take() {
1404 started_tx.send(()).expect("fence observer still present");
1405 }
1406 let (released, wake) = &*hook_release;
1407 let mut released = released.lock().unwrap();
1408 while !*released {
1409 released = wake.wait(released).unwrap();
1410 }
1411 });
1412 (hook, started_rx, release)
1413 }
1414
1415 fn release_fence(release: &Arc<(Mutex<bool>, Condvar)>) {
1416 let (released, wake) = &**release;
1417 *released.lock().unwrap() = true;
1418 wake.notify_all();
1419 }
1420
1421 async fn assert_second_write_waits_for_cache_publish<T>(
1422 first: tokio::task::JoinHandle<std::io::Result<T>>,
1423 mut second: tokio::task::JoinHandle<std::io::Result<T>>,
1424 release: Arc<(Mutex<bool>, Condvar)>,
1425 ) -> (T, T) {
1426 let second_before_release =
1427 tokio::time::timeout(Duration::from_millis(100), &mut second).await;
1428 let completed_before_release = second_before_release.is_ok();
1429 release_fence(&release);
1430
1431 let first = first
1432 .await
1433 .expect("first writer joins")
1434 .expect("first writer succeeds");
1435 let second = match second_before_release {
1436 Ok(joined) => joined
1437 .expect("second writer joins")
1438 .expect("second writer succeeds"),
1439 Err(_) => second
1440 .await
1441 .expect("second writer joins")
1442 .expect("second writer succeeds"),
1443 };
1444 assert!(
1445 !completed_before_release,
1446 "the second write must remain behind the first write's durable-to-cache fence"
1447 );
1448 (first, second)
1449 }
1450
1451 #[derive(Clone, Copy, Debug)]
1452 enum FullSaveRoute {
1453 InherentSave,
1454 SaveAndCache,
1455 RuntimePersistence,
1456 }
1457
1458 impl FullSaveRoute {
1459 fn operation(self) -> &'static str {
1460 match self {
1461 Self::InherentSave => "save_full",
1462 Self::SaveAndCache => "save_and_cache",
1463 Self::RuntimePersistence => "save_runtime_session",
1464 }
1465 }
1466
1467 fn name(self) -> &'static str {
1468 match self {
1469 Self::InherentSave => "inherent",
1470 Self::SaveAndCache => "save-and-cache",
1471 Self::RuntimePersistence => "runtime-persistence",
1472 }
1473 }
1474 }
1475
1476 async fn assert_full_save_route_serializes_cache_publish(route: FullSaveRoute) {
1477 let temp = tempfile::tempdir().unwrap();
1478 let concrete_storage = Arc::new(
1479 bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
1480 .await
1481 .expect("SessionStoreV2"),
1482 );
1483 let storage: Arc<dyn Storage> = concrete_storage;
1484 let id = format!("root-full-cache-order-{}", route.name());
1485 let (hook, first_durable, release) = durable_cache_fence(route.operation(), id.clone());
1486 let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
1487
1488 let mut initial = Session::new(&id, "model");
1489 initial.add_message(bamboo_agent_core::Message::user("durable transcript"));
1490 initial.set_task_list(task_list(&id, "initial"));
1491 initial.set_task_list_version_meta("0");
1492 initial
1493 .metadata
1494 .insert("unrelated.runtime".to_string(), "keep".to_string());
1495 storage.save_session(&initial).await.unwrap();
1496 cache_put(&repo, &initial);
1497
1498 let first_repo = repo.clone();
1499 let mut root_snapshot = initial.clone();
1500 root_snapshot.add_message(bamboo_agent_core::Message::assistant(
1501 "full-save transcript suffix",
1502 None,
1503 ));
1504 root_snapshot.set_task_list(task_list(&id, "root"));
1505 root_snapshot.set_task_list_version_meta("1");
1506 let first = tokio::spawn(async move {
1507 match route {
1508 FullSaveRoute::InherentSave => first_repo.save(&mut root_snapshot).await,
1509 FullSaveRoute::SaveAndCache => {
1510 first_repo.save_and_cache(&mut root_snapshot).await;
1511 Ok(())
1512 }
1513 FullSaveRoute::RuntimePersistence => {
1514 bamboo_domain::RuntimeSessionPersistence::save_runtime_session(
1515 first_repo.as_ref(),
1516 &mut root_snapshot,
1517 )
1518 .await
1519 }
1520 }
1521 });
1522 first_durable
1523 .await
1524 .expect("root full durable write reached");
1525
1526 let second_repo = repo.clone();
1527 let second_id = id.clone();
1528 let child_task_list = task_list(&id, "child");
1529 let second = tokio::spawn(async move {
1530 bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
1531 second_repo.as_ref(),
1532 &second_id,
1533 &child_task_list,
1534 "2",
1535 )
1536 .await
1537 .map(|updated| assert!(updated, "root must exist"))
1538 });
1539 assert_second_write_waits_for_cache_publish(first, second, release).await;
1540
1541 let durable = storage.load_session(&id).await.unwrap().unwrap();
1542 let cached = read_cached_session(repo.cache(), &id).expect("cached root");
1543 for (tier, session) in [("durable", durable), ("cache", cached)] {
1544 assert_eq!(
1545 session.task_list_version_meta().as_deref(),
1546 Some("2"),
1547 "{route:?} {tier} must retain the child transaction"
1548 );
1549 assert_eq!(
1550 session.task_list.as_ref().map(|list| list.title.as_str()),
1551 Some("child"),
1552 "{route:?} {tier} must retain the child transaction"
1553 );
1554 assert_eq!(
1555 session
1556 .metadata
1557 .get("unrelated.runtime")
1558 .map(String::as_str),
1559 Some("keep"),
1560 "{route:?} {tier} must preserve unrelated runtime state"
1561 );
1562 assert_eq!(
1563 session.messages.len(),
1564 2,
1565 "{route:?} {tier} must preserve the full-save transcript"
1566 );
1567 }
1568 }
1569
1570 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1571 async fn root_full_saves_and_child_task_patch_share_publish_order() {
1572 for route in [
1573 FullSaveRoute::InherentSave,
1574 FullSaveRoute::SaveAndCache,
1575 FullSaveRoute::RuntimePersistence,
1576 ] {
1577 assert_full_save_route_serializes_cache_publish(route).await;
1578 }
1579 }
1580
1581 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1582 async fn checkpoint_and_child_task_patch_share_publish_order() {
1583 let temp = tempfile::tempdir().unwrap();
1584 let concrete_storage = Arc::new(
1585 bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
1586 .await
1587 .expect("SessionStoreV2"),
1588 );
1589 let storage: Arc<dyn Storage> = concrete_storage;
1590 let id = "checkpoint-cache-order";
1591 let (hook, checkpoint_durable, release) = durable_cache_fence("checkpoint", id);
1592 let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
1593
1594 let mut initial = Session::new(id, "model");
1595 initial.add_message(bamboo_agent_core::Message::user("durable transcript"));
1596 initial.set_task_list(task_list(id, "initial"));
1597 initial.set_task_list_version_meta("0");
1598 initial
1599 .metadata
1600 .insert("unrelated.runtime".to_string(), "keep".to_string());
1601 storage.save_session(&initial).await.unwrap();
1602 cache_put(&repo, &initial);
1603
1604 let checkpoint_repo = repo.clone();
1605 let mut checkpoint_snapshot = initial.clone();
1606 checkpoint_snapshot.add_message(bamboo_agent_core::Message::assistant(
1607 "checkpoint transcript suffix",
1608 None,
1609 ));
1610 checkpoint_snapshot.set_task_list(task_list(id, "checkpoint"));
1611 checkpoint_snapshot.set_task_list_version_meta("1");
1612 let checkpoint = tokio::spawn(async move {
1613 bamboo_domain::RuntimeSessionPersistence::checkpoint_runtime_session(
1614 checkpoint_repo.as_ref(),
1615 &mut checkpoint_snapshot,
1616 )
1617 .await
1618 });
1619 checkpoint_durable
1620 .await
1621 .expect("checkpoint durable write reached");
1622
1623 let child_repo = repo.clone();
1624 let child_task_list = task_list(id, "child");
1625 let child_patch = tokio::spawn(async move {
1626 bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
1627 child_repo.as_ref(),
1628 id,
1629 &child_task_list,
1630 "2",
1631 )
1632 .await
1633 .map(|updated| assert!(updated, "root must exist"))
1634 });
1635 assert_second_write_waits_for_cache_publish(checkpoint, child_patch, release).await;
1636
1637 let durable = storage.load_session(id).await.unwrap().unwrap();
1638 let cached = read_cached_session(repo.cache(), id).expect("cached root");
1639 for (tier, session) in [("durable", durable), ("cache", cached)] {
1640 assert_eq!(
1641 session.task_list_version_meta().as_deref(),
1642 Some("2"),
1643 "{tier} must retain the child transaction"
1644 );
1645 assert_eq!(
1646 session.task_list.as_ref().map(|list| list.title.as_str()),
1647 Some("child"),
1648 "{tier} must retain the child transaction"
1649 );
1650 assert_eq!(
1651 session
1652 .metadata
1653 .get("unrelated.runtime")
1654 .map(String::as_str),
1655 Some("keep"),
1656 "{tier} must preserve unrelated runtime state"
1657 );
1658 assert_eq!(
1659 session.messages.len(),
1660 2,
1661 "{tier} must preserve the checkpoint transcript"
1662 );
1663 }
1664 }
1665
1666 #[derive(Clone, Copy, Debug)]
1667 enum CacheBackfillRoute {
1668 Load,
1669 TryLoad,
1670 }
1671
1672 impl CacheBackfillRoute {
1673 fn operation(self) -> &'static str {
1674 match self {
1675 Self::Load => "load",
1676 Self::TryLoad => "try_load",
1677 }
1678 }
1679
1680 fn name(self) -> &'static str {
1681 match self {
1682 Self::Load => "load",
1683 Self::TryLoad => "try-load",
1684 }
1685 }
1686 }
1687
1688 async fn assert_cache_backfill_serializes_with_task_patch(route: CacheBackfillRoute) {
1689 let temp = tempfile::tempdir().unwrap();
1690 let concrete_storage = Arc::new(
1691 bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
1692 .await
1693 .expect("SessionStoreV2"),
1694 );
1695 let storage: Arc<dyn Storage> = concrete_storage;
1696 let id = format!("cache-backfill-order-{}", route.name());
1697 let (hook, loaded_old_durable, release) =
1698 durable_cache_fence(route.operation(), id.clone());
1699 let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
1700
1701 let mut initial = Session::new(&id, "model");
1702 initial.set_task_list(task_list(&id, "initial"));
1703 initial.set_task_list_version_meta("0");
1704 storage.save_session(&initial).await.unwrap();
1705 assert!(
1706 read_cached_session(repo.cache(), &id).is_none(),
1707 "the race requires a genuine cache miss"
1708 );
1709
1710 let load_repo = repo.clone();
1711 let load_id = id.clone();
1712 let load = tokio::spawn(async move {
1713 let loaded = match route {
1714 CacheBackfillRoute::Load => load_repo.load(&load_id).await,
1715 CacheBackfillRoute::TryLoad => {
1716 load_repo.try_load(&load_id).await.expect("storage load")
1717 }
1718 };
1719 assert!(loaded.is_some(), "seeded session must load");
1720 Ok(())
1721 });
1722 loaded_old_durable
1723 .await
1724 .expect("old durable snapshot loaded");
1725
1726 let patch_repo = repo.clone();
1727 let patch_id = id.clone();
1728 let child_task_list = task_list(&id, "child");
1729 let patch = tokio::spawn(async move {
1730 bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
1731 patch_repo.as_ref(),
1732 &patch_id,
1733 &child_task_list,
1734 "1",
1735 )
1736 .await
1737 .map(|updated| assert!(updated, "root must exist"))
1738 });
1739 assert_second_write_waits_for_cache_publish(load, patch, release).await;
1740
1741 let durable = storage.load_session(&id).await.unwrap().unwrap();
1742 let cached = read_cached_session(repo.cache(), &id).expect("backfilled cache");
1743 for (tier, session) in [("durable", durable), ("cache", cached)] {
1744 assert_eq!(
1745 session.task_list_version_meta().as_deref(),
1746 Some("1"),
1747 "{route:?} {tier} must retain the child transaction"
1748 );
1749 assert_eq!(
1750 session.task_list.as_ref().map(|list| list.title.as_str()),
1751 Some("child"),
1752 "{route:?} {tier} must retain the child transaction"
1753 );
1754 }
1755 }
1756
1757 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1758 async fn cache_miss_backfills_and_child_task_patch_share_publish_order() {
1759 for route in [CacheBackfillRoute::Load, CacheBackfillRoute::TryLoad] {
1760 assert_cache_backfill_serializes_with_task_patch(route).await;
1761 }
1762 }
1763
1764 #[tokio::test]
1765 async fn cache_hits_do_not_wait_for_the_persistence_lock() {
1766 let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
1767 let repo = test_repo(storage);
1768 let id = "cache-hit-lock-free";
1769 let cached = Session::new(id, "cached-model");
1770 cache_put(&repo, &cached);
1771 let persistence_guard = repo.persistence().acquire_lock(id).await;
1772
1773 let loaded = tokio::time::timeout(Duration::from_millis(100), repo.load(id)).await;
1774 let try_loaded = tokio::time::timeout(Duration::from_millis(100), repo.try_load(id)).await;
1775 drop(persistence_guard);
1776
1777 assert_eq!(
1778 loaded
1779 .expect("cache hit must not wait for the persistence lock")
1780 .expect("cached session")
1781 .model,
1782 "cached-model"
1783 );
1784 assert_eq!(
1785 try_loaded
1786 .expect("fallible cache hit must not wait for the persistence lock")
1787 .expect("cache read succeeds")
1788 .expect("cached session")
1789 .model,
1790 "cached-model"
1791 );
1792 }
1793
1794 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1795 async fn load_merged_storage_refresh_and_child_task_patch_share_publish_order() {
1796 let temp = tempfile::tempdir().unwrap();
1797 let concrete_storage = Arc::new(
1798 bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
1799 .await
1800 .expect("SessionStoreV2"),
1801 );
1802 let storage: Arc<dyn Storage> = concrete_storage;
1803 let id = "load-merged-cache-order";
1804 let (hook, loaded_old_durable, release) = durable_cache_fence("load_merged", id);
1805 let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
1806
1807 let mut durable = Session::new(id, "model");
1808 durable.updated_at = Utc::now();
1809 durable.set_task_list(task_list(id, "initial"));
1810 durable.set_task_list_version_meta("0");
1811 storage.save_session(&durable).await.unwrap();
1812
1813 let mut memory = durable.clone();
1814 memory.updated_at = durable.updated_at - chrono::Duration::seconds(1);
1815 memory.set_task_list(task_list(id, "memory"));
1816 cache_put(&repo, &memory);
1817
1818 let load_repo = repo.clone();
1819 let load = tokio::spawn(async move {
1820 assert!(
1821 load_repo.load_merged(id).await.is_some(),
1822 "seeded session must load"
1823 );
1824 Ok(())
1825 });
1826 loaded_old_durable
1827 .await
1828 .expect("old durable snapshot loaded");
1829
1830 let patch_repo = repo.clone();
1831 let child_task_list = task_list(id, "child");
1832 let patch = tokio::spawn(async move {
1833 bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
1834 patch_repo.as_ref(),
1835 id,
1836 &child_task_list,
1837 "1",
1838 )
1839 .await
1840 .map(|updated| assert!(updated, "root must exist"))
1841 });
1842 assert_second_write_waits_for_cache_publish(load, patch, release).await;
1843
1844 let durable = storage.load_session(id).await.unwrap().unwrap();
1845 let cached = read_cached_session(repo.cache(), id).expect("refreshed cache");
1846 for (tier, session) in [("durable", durable), ("cache", cached)] {
1847 assert_eq!(
1848 session.task_list_version_meta().as_deref(),
1849 Some("1"),
1850 "{tier} must retain the child transaction"
1851 );
1852 assert_eq!(
1853 session.task_list.as_ref().map(|list| list.title.as_str()),
1854 Some("child"),
1855 "{tier} must retain the child transaction"
1856 );
1857 }
1858 }
1859
1860 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1861 async fn legacy_clear_refresh_and_child_task_patch_share_publish_order() {
1862 let temp = tempfile::tempdir().unwrap();
1863 let concrete_storage = Arc::new(
1864 bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
1865 .await
1866 .expect("SessionStoreV2"),
1867 );
1868 let storage: Arc<dyn Storage> = concrete_storage;
1869 let id = "legacy-clear-cache-order";
1870 let expected = vec![serde_json::json!({"content": "legacy"})];
1871 let (hook, loaded_post_cas_snapshot, release) = durable_cache_fence("clear_legacy", id);
1872 let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
1873
1874 let mut initial = Session::new(id, "model");
1875 initial.set_pending_injected_messages(expected.clone());
1876 initial.set_task_list(task_list(id, "initial"));
1877 initial.set_task_list_version_meta("0");
1878 storage.save_session(&initial).await.unwrap();
1879 cache_put(&repo, &initial);
1880
1881 let clear_repo = repo.clone();
1882 let clear_expected = expected.clone();
1883 let clear = tokio::spawn(async move {
1884 bamboo_domain::RuntimeSessionPersistence::clear_legacy_pending_messages(
1885 clear_repo.as_ref(),
1886 id,
1887 &clear_expected,
1888 )
1889 .await
1890 .map(|cleared| assert!(cleared, "legacy queue must match"))
1891 });
1892 loaded_post_cas_snapshot
1893 .await
1894 .expect("post-CAS snapshot loaded");
1895
1896 let patch_repo = repo.clone();
1897 let child_task_list = task_list(id, "child");
1898 let patch = tokio::spawn(async move {
1899 bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
1900 patch_repo.as_ref(),
1901 id,
1902 &child_task_list,
1903 "1",
1904 )
1905 .await
1906 .map(|updated| assert!(updated, "root must exist"))
1907 });
1908 assert_second_write_waits_for_cache_publish(clear, patch, release).await;
1909
1910 let durable = storage.load_session(id).await.unwrap().unwrap();
1911 let cached = read_cached_session(repo.cache(), id).expect("refreshed cache");
1912 for (tier, session) in [("durable", durable), ("cache", cached)] {
1913 assert_eq!(
1914 session.task_list_version_meta().as_deref(),
1915 Some("1"),
1916 "{tier} must retain the child transaction"
1917 );
1918 assert_eq!(
1919 session.task_list.as_ref().map(|list| list.title.as_str()),
1920 Some("child"),
1921 "{tier} must retain the child transaction"
1922 );
1923 assert!(
1924 !session.has_pending_injected_messages(),
1925 "{tier} must retain the successful legacy clear"
1926 );
1927 }
1928 }
1929
1930 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1931 async fn concurrent_task_patches_publish_cache_in_durable_order() {
1932 let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
1933 let (hook, first_durable, release) = durable_cache_fence("task", "1");
1934 let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
1935 let id = "concurrent-task-cache-order";
1936 let mut initial = Session::new(id, "model");
1937 initial.set_task_list(task_list(id, "initial"));
1938 initial.set_task_list_version_meta("0");
1939 storage.save_session(&initial).await.unwrap();
1940 cache_put(&repo, &initial);
1941
1942 let first_repo = repo.clone();
1943 let first_task_list = task_list(id, "first");
1944 let first = tokio::spawn(async move {
1945 bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
1946 first_repo.as_ref(),
1947 id,
1948 &first_task_list,
1949 "1",
1950 )
1951 .await
1952 });
1953 first_durable.await.expect("first durable write reached");
1954
1955 let second_repo = repo.clone();
1956 let second_task_list = task_list(id, "second");
1957 let second = tokio::spawn(async move {
1958 bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
1959 second_repo.as_ref(),
1960 id,
1961 &second_task_list,
1962 "2",
1963 )
1964 .await
1965 });
1966 let (first_updated, second_updated) =
1967 assert_second_write_waits_for_cache_publish(first, second, release).await;
1968 assert!(first_updated && second_updated);
1969
1970 let durable = storage.load_session(id).await.unwrap().unwrap();
1971 let cached = read_cached_session(repo.cache(), id).expect("cached root");
1972 for (tier, session) in [("durable", durable), ("cache", cached)] {
1973 assert_eq!(
1974 session.task_list_version_meta().as_deref(),
1975 Some("2"),
1976 "{tier} must retain the second transaction"
1977 );
1978 assert_eq!(
1979 session.task_list.as_ref().map(|list| list.title.as_str()),
1980 Some("second"),
1981 "{tier} must retain the second transaction"
1982 );
1983 }
1984 }
1985
1986 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1987 async fn root_control_plane_save_and_child_task_patch_share_publish_order() {
1988 let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
1989 let (hook, root_durable, release) = durable_cache_fence("save", "root-control-plane-order");
1990 let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
1991 let id = "root-control-plane-order";
1992
1993 let mut initial = Session::new(id, "model");
1994 initial.add_message(bamboo_agent_core::Message::user("durable transcript"));
1995 initial.set_task_list(task_list(id, "initial"));
1996 initial.set_task_list_version_meta("0");
1997 initial
1998 .metadata
1999 .insert("unrelated.runtime".to_string(), "keep".to_string());
2000 storage.save_session(&initial).await.unwrap();
2001 cache_put(&repo, &initial);
2002
2003 let root_repo = repo.clone();
2004 let mut root_snapshot = initial.clone();
2005 root_snapshot.set_task_list(task_list(id, "root"));
2006 root_snapshot.set_task_list_version_meta("1");
2007 let root_save = tokio::spawn(async move {
2008 bamboo_domain::RuntimeSessionPersistence::save_runtime_control_plane(
2009 root_repo.as_ref(),
2010 &mut root_snapshot,
2011 )
2012 .await
2013 });
2014 root_durable.await.expect("root durable write reached");
2015
2016 let child_repo = repo.clone();
2017 let child_task_list = task_list(id, "child");
2018 let child_patch = tokio::spawn(async move {
2019 bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
2020 child_repo.as_ref(),
2021 id,
2022 &child_task_list,
2023 "2",
2024 )
2025 .await
2026 .map(|updated| {
2027 assert!(updated, "root must exist");
2028 })
2029 });
2030 assert_second_write_waits_for_cache_publish(root_save, child_patch, release).await;
2031
2032 let durable = storage.load_session(id).await.unwrap().unwrap();
2033 let cached = read_cached_session(repo.cache(), id).expect("cached root");
2034 for (tier, session) in [("durable", durable), ("cache", cached)] {
2035 assert_eq!(
2036 session.task_list_version_meta().as_deref(),
2037 Some("2"),
2038 "{tier} must retain the child transaction"
2039 );
2040 assert_eq!(
2041 session.task_list.as_ref().map(|list| list.title.as_str()),
2042 Some("child"),
2043 "{tier} must retain the child transaction"
2044 );
2045 assert_eq!(
2046 session
2047 .metadata
2048 .get("unrelated.runtime")
2049 .map(String::as_str),
2050 Some("keep"),
2051 "{tier} must preserve unrelated runtime state"
2052 );
2053 assert_eq!(
2054 session.messages.len(),
2055 1,
2056 "{tier} must preserve the transcript"
2057 );
2058 }
2059 }
2060
2061 #[tokio::test]
2062 async fn stale_control_plane_save_keeps_checkpointed_ledger_in_durable_and_cache() {
2063 let temp = tempfile::tempdir().unwrap();
2064 let storage: Arc<dyn Storage> = Arc::new(
2065 bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
2066 .await
2067 .expect("SessionStoreV2"),
2068 );
2069 let repo = test_repo(storage.clone());
2070 let id = "control-plane-ledger-cache";
2071 let mut initial = Session::new(id, "model");
2072 initial.add_message(bamboo_agent_core::Message::user("durable transcript"));
2073 let assistant = bamboo_agent_core::Message::assistant("normalized", None);
2074 let anchor = assistant.id.clone();
2075 initial.add_message(assistant);
2076 let native_item = bamboo_domain::ProviderTranscriptItem::try_from_payload(
2077 bamboo_domain::ProviderFamily::OpenAi,
2078 bamboo_domain::ProviderProtocol::OpenAiResponsesV1,
2079 bamboo_domain::ProviderTranscriptOrigin::Provider,
2080 bamboo_domain::ProviderTranscriptAuthor::Model,
2081 serde_json::json!({
2082 "type":"tool_search_call","id":"tsc_cache_native","execution":"client","call_id":"cache_native",
2083 "status":"completed","arguments":{"query":"CACHE_NATIVE_PAYLOAD_SENTINEL"}
2084 }),
2085 )
2086 .unwrap();
2087 initial
2088 .append_provider_transcript_group(&anchor, None, vec![native_item])
2089 .unwrap();
2090 storage.save_session(&initial).await.unwrap();
2091 cache_put(&repo, &initial);
2092 let mut stale = initial.clone();
2093
2094 let mut runner = initial;
2095 runner.model_context_state = Some(bamboo_domain::ModelContextState {
2096 state_revision: 1,
2097 prefix_epoch: 1,
2098 cache_scope_sha256: Some("scope".to_string()),
2099 transcript_item_sha256: vec!["runner-l1".to_string()],
2100 ..bamboo_domain::ModelContextState::default()
2101 });
2102 bamboo_domain::RuntimeSessionPersistence::checkpoint_runtime_session(&repo, &mut runner)
2103 .await
2104 .unwrap();
2105
2106 stale
2107 .metadata
2108 .insert("runtime.suspend_reason".to_string(), "waiting".to_string());
2109 bamboo_domain::RuntimeSessionPersistence::save_runtime_control_plane(&repo, &mut stale)
2110 .await
2111 .unwrap();
2112
2113 let expected = runner.model_context_state;
2114 let expected_native = runner.provider_transcript;
2115 let durable = storage.load_session(id).await.unwrap().unwrap();
2116 let cached = read_cached_session(repo.cache(), id).expect("cached session");
2117 for (tier, session) in [("durable", durable), ("cache", cached.clone())] {
2118 assert_eq!(session.model_context_state, expected, "tier={tier}");
2119 assert_eq!(
2120 session.provider_transcript, expected_native,
2121 "tier={tier} must retain the message-anchored native transcript"
2122 );
2123 assert_eq!(
2124 session
2125 .metadata
2126 .get("runtime.suspend_reason")
2127 .map(String::as_str),
2128 Some("waiting"),
2129 "tier={tier}"
2130 );
2131 assert_eq!(session.messages.len(), 2, "tier={tier}");
2132 }
2133
2134 let runtime_json =
2135 std::fs::read_to_string(temp.path().join("sessions").join(id).join("runtime.json"))
2136 .unwrap();
2137 assert!(!runtime_json.contains("CACHE_NATIVE_PAYLOAD_SENTINEL"));
2138
2139 let mut cache_writer = cached;
2142 bamboo_domain::RuntimeSessionPersistence::checkpoint_runtime_session(
2143 &repo,
2144 &mut cache_writer,
2145 )
2146 .await
2147 .unwrap();
2148 let restarted = storage.load_session(id).await.unwrap().unwrap();
2149 assert_eq!(restarted.provider_transcript, expected_native);
2150 }
2151
2152 #[tokio::test]
2153 async fn narrow_runtime_metadata_transaction_preserves_live_and_durable_non_owned_state() {
2154 let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
2155 let repo = test_repo(storage.clone());
2156 let id = "narrow-metadata";
2157 let mut durable = Session::new(id, "durable-model");
2158 durable.add_message(bamboo_agent_core::Message::user("durable user turn"));
2159 durable
2160 .metadata
2161 .insert("external.durable".to_string(), "keep".to_string());
2162 storage.save_session(&durable).await.expect("seed durable");
2163
2164 let mut live = durable.clone();
2165 live.add_message(bamboo_agent_core::Message::assistant(
2166 "in-flight assistant tool call",
2167 None,
2168 ));
2169 live.model = "live-model".to_string();
2170 live.metadata
2171 .insert("external.live".to_string(), "keep".to_string());
2172 cache_put(&repo, &live);
2173
2174 repo.update_runtime_session(id, &["workflow.owned"], |latest| {
2175 latest
2176 .metadata
2177 .insert("workflow.owned".to_string(), "active".to_string());
2178 })
2179 .await
2180 .expect("transaction")
2181 .expect("session exists");
2182
2183 let saved = storage
2184 .load_session(id)
2185 .await
2186 .expect("load durable")
2187 .expect("durable exists");
2188 assert_eq!(
2189 saved.messages.len(),
2190 1,
2191 "transaction never writes stale live messages"
2192 );
2193 assert_eq!(
2194 saved.metadata.get("external.durable").map(String::as_str),
2195 Some("keep")
2196 );
2197 assert_eq!(
2198 saved.metadata.get("workflow.owned").map(String::as_str),
2199 Some("active")
2200 );
2201
2202 let cached = read_cached_session(repo.cache(), id).expect("live cache");
2203 assert_eq!(
2204 cached.messages.len(),
2205 2,
2206 "cache live tool call is not replaced"
2207 );
2208 assert_eq!(cached.model, "live-model");
2209 assert_eq!(
2210 cached.metadata.get("external.live").map(String::as_str),
2211 Some("keep")
2212 );
2213 assert_eq!(
2214 cached.metadata.get("workflow.owned").map(String::as_str),
2215 Some("active")
2216 );
2217 }
2218
2219 #[tokio::test]
2224 async fn load_merged_does_not_regress_to_older_storage() {
2225 let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
2226 let repo = test_repo(storage.clone());
2227 let id = "s1";
2228
2229 let mut stale = Session::new(id.to_string(), "m");
2230 stale.set_pending_question(
2231 "tc1".into(),
2232 "kind".into(),
2233 "q?".into(),
2234 vec!["OK".into()],
2235 true,
2236 );
2237 stale.updated_at = Utc::now() - chrono::Duration::seconds(10);
2238 storage.save_session(&stale).await.unwrap();
2239
2240 let mut fresh = stale.clone();
2243 fresh.clear_pending_question();
2244 fresh.updated_at = Utc::now();
2245 cache_put(&repo, &fresh);
2246
2247 let merged = repo.load_merged(id).await.expect("session exists");
2248 assert!(
2249 merged.pending_question.is_none(),
2250 "must return the newer answered memory copy, not the stale storage one"
2251 );
2252 let cached = read_cached_session(repo.cache(), id).expect("cached");
2253 assert!(
2254 cached.pending_question.is_none(),
2255 "load_merged must never regress the cache to a stale storage copy"
2256 );
2257 }
2258
2259 #[tokio::test]
2263 async fn load_merged_recovers_pending_question_from_same_age_storage() {
2264 let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
2265 let repo = test_repo(storage.clone());
2266 let id = "s2";
2267 let ts = Utc::now();
2268
2269 let mut with_pending = Session::new(id.to_string(), "m");
2270 with_pending.set_pending_question(
2271 "tc".into(),
2272 "k".into(),
2273 "q".into(),
2274 vec!["OK".into()],
2275 true,
2276 );
2277 with_pending.updated_at = ts;
2278 storage.save_session(&with_pending).await.unwrap();
2279
2280 let mut lost = with_pending.clone();
2281 lost.clear_pending_question();
2282 lost.updated_at = ts;
2283 cache_put(&repo, &lost);
2284
2285 let merged = repo.load_merged(id).await.expect("session exists");
2286 assert!(
2287 merged.pending_question.is_some(),
2288 "same-age storage carrying a pending question must still be recovered"
2289 );
2290 }
2291
2292 #[tokio::test]
2293 async fn runtime_publish_refreshes_cache_even_when_storage_fails() {
2294 let id = "runtime-selection";
2295 let mut previous = Session::new(id.to_string(), "m");
2296 previous.metadata.insert(
2297 "skill_runtime_selected_skill_ids".to_string(),
2298 "[\"plan\"]".to_string(),
2299 );
2300 let storage: Arc<dyn Storage> = Arc::new(FailingSaveStorage {
2301 persisted: Mutex::new(Some(previous.clone())),
2302 });
2303 let repo = test_repo(storage.clone());
2304 cache_put(&repo, &previous);
2305
2306 let mut current = previous.clone();
2307 current.metadata.insert(
2308 "skill_runtime_selected_skill_ids".to_string(),
2309 "[\"review\"]".to_string(),
2310 );
2311 current.updated_at = Utc::now();
2312
2313 let result =
2314 bamboo_domain::RuntimeSessionPersistence::save_runtime_session(&repo, &mut current)
2315 .await;
2316 assert!(result.is_err(), "durable failure must still be surfaced");
2317
2318 let cached = repo.load(id).await.expect("cached current session");
2319 assert_eq!(
2320 cached
2321 .metadata
2322 .get("skill_runtime_selected_skill_ids")
2323 .map(String::as_str),
2324 Some("[\"review\"]")
2325 );
2326 let allowlist = bamboo_skills::access_control::extract_skill_allowlist(&cached.metadata)
2327 .expect("runtime authorization allowlist");
2328 assert!(allowlist.contains("review"));
2329 assert!(!allowlist.contains("plan"));
2330 let durable = storage
2331 .load_session(id)
2332 .await
2333 .expect("load durable state")
2334 .expect("previous durable session");
2335 assert_eq!(
2336 durable
2337 .metadata
2338 .get("skill_runtime_selected_skill_ids")
2339 .map(String::as_str),
2340 Some("[\"plan\"]")
2341 );
2342 }
2343
2344 #[tokio::test]
2345 async fn finalized_runtime_save_reconciles_late_child_completion_in_disk_and_cache() {
2346 use bamboo_domain::session::runtime_state::{
2347 AgentRuntimeState, AgentStatusState, ChildWaitPolicy, WaitingForChildrenState,
2348 };
2349
2350 let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
2351 let repo = test_repo(storage.clone());
2352 let id = "late-child-finalization";
2353 let mut parent = Session::new(id, "model");
2354 let mut runtime = AgentRuntimeState::new("parent-run");
2355 runtime.status = AgentStatusState::Suspended;
2356 let mut wait = WaitingForChildrenState::for_children(
2357 vec!["child-1".to_string()],
2358 ChildWaitPolicy::All,
2359 Utc::now(),
2360 );
2361 wait.registered_by_tool_call_id = Some("tc_wait".to_string());
2362 runtime.waiting_for_children = Some(wait);
2363 parent.agent_runtime_state = Some(runtime);
2364 parent.metadata.insert(
2365 "runtime.suspend_reason".to_string(),
2366 "waiting_for_children".to_string(),
2367 );
2368 parent.set_last_run_status("suspended");
2369 storage.save_session(&parent).await.unwrap();
2370 cache_put(&repo, &parent);
2371
2372 let mut stale_runner = parent.clone();
2373 let mut completed = parent;
2374 let completed_runtime = completed.agent_runtime_state.as_mut().unwrap();
2375 completed_runtime.waiting_for_children = None;
2376 completed_runtime.status = AgentStatusState::Idle;
2377 completed_runtime.suspension = None;
2378 completed.metadata.remove("runtime.suspend_reason");
2379 let mut outcome = bamboo_agent_core::Message::user("completed child result");
2380 outcome.metadata = Some(serde_json::json!({
2381 "runtime_kind": "child_completion_resume"
2382 }));
2383 completed.add_message(outcome);
2384 storage.save_session(&completed).await.unwrap();
2385
2386 bamboo_domain::RuntimeSessionPersistence::save_finalized_runtime_session(
2387 &repo,
2388 &mut stale_runner,
2389 )
2390 .await
2391 .unwrap();
2392
2393 let durable = storage.load_session(id).await.unwrap().unwrap();
2394 let cached = read_cached_session(repo.cache(), id).expect("final snapshot published");
2395 for saved in [&durable, &cached] {
2396 assert!(saved
2397 .agent_runtime_state
2398 .as_ref()
2399 .unwrap()
2400 .waiting_for_children
2401 .is_none());
2402 assert!(!saved.metadata.contains_key("runtime.suspend_reason"));
2403 assert_eq!(saved.last_run_status().as_deref(), Some("completed"));
2404 assert!(saved
2405 .messages
2406 .iter()
2407 .any(|message| message.content == "completed child result"));
2408 }
2409 }
2410
2411 #[tokio::test]
2412 async fn inherent_save_leaves_existing_cache_untouched_when_storage_fails() {
2413 let id = "inherent-save-failure";
2414 let previous = Session::new(id, "previous");
2415 let storage: Arc<dyn Storage> = Arc::new(FailingSaveStorage {
2416 persisted: Mutex::new(Some(previous.clone())),
2417 });
2418 let repo = test_repo(storage);
2419 cache_put(&repo, &previous);
2420
2421 let mut current = previous.clone();
2422 current.model = "current".to_string();
2423 assert!(repo.save(&mut current).await.is_err());
2424 assert_eq!(
2425 read_cached_session(repo.cache(), id)
2426 .expect("existing cache")
2427 .model,
2428 "previous",
2429 "fallible inherent save must publish only after a durable commit"
2430 );
2431 }
2432
2433 #[tokio::test]
2434 async fn blocked_replay_ambiguous_save_evicts_same_age_stale_approval_cache() {
2435 let id = "blocked-replay-ambiguous-save";
2436 let mut approved = Session::new(id, "model");
2437 approved.metadata.insert(
2438 crate::session_app::respond::PERMISSION_REEXECUTE_METADATA_KEY.into(),
2439 "call-1".into(),
2440 );
2441 let storage: Arc<dyn Storage> = Arc::new(AmbiguousReplaySaveStorage {
2442 persisted: Mutex::new(Some(approved.clone())),
2443 });
2444 let repo = test_repo(storage);
2445 cache_put(&repo, &approved);
2446
2447 let mut resolved = approved.clone();
2448 resolved
2449 .metadata
2450 .remove(crate::session_app::respond::PERMISSION_REEXECUTE_METADATA_KEY);
2451 assert!(repo.save_replay_resolution(&mut resolved).await.is_err());
2452 assert!(read_cached_session(repo.cache(), id).is_none());
2453 let reloaded = repo.load_merged_checked(id).await.unwrap().unwrap();
2454 assert_eq!(reloaded.updated_at, approved.updated_at);
2455 assert!(!reloaded
2456 .metadata
2457 .contains_key(crate::session_app::respond::PERMISSION_REEXECUTE_METADATA_KEY));
2458 }
2459
2460 #[tokio::test]
2461 async fn blocked_replay_precommit_save_error_reloads_retryable_approval() {
2462 let id = "blocked-replay-precommit-save";
2463 let mut approved = Session::new(id, "model");
2464 approved.metadata.insert(
2465 crate::session_app::respond::PERMISSION_REEXECUTE_METADATA_KEY.into(),
2466 "call-1".into(),
2467 );
2468 let storage: Arc<dyn Storage> = Arc::new(FailingSaveStorage {
2469 persisted: Mutex::new(Some(approved.clone())),
2470 });
2471 let repo = test_repo(storage);
2472 cache_put(&repo, &approved);
2473
2474 let mut resolved = approved.clone();
2475 resolved
2476 .metadata
2477 .remove(crate::session_app::respond::PERMISSION_REEXECUTE_METADATA_KEY);
2478 assert!(repo.save_replay_resolution(&mut resolved).await.is_err());
2479 assert!(read_cached_session(repo.cache(), id).is_none());
2480 let reloaded = repo.load_merged_checked(id).await.unwrap().unwrap();
2481 assert_eq!(reloaded.updated_at, approved.updated_at);
2482 assert_eq!(
2483 reloaded
2484 .metadata
2485 .get(crate::session_app::respond::PERMISSION_REEXECUTE_METADATA_KEY)
2486 .map(String::as_str),
2487 Some("call-1")
2488 );
2489 }
2490
2491 #[tokio::test]
2492 async fn save_and_cache_still_refreshes_cache_when_storage_fails() {
2493 let id = "save-and-cache-failure";
2494 let previous = Session::new(id, "previous");
2495 let storage: Arc<dyn Storage> = Arc::new(FailingSaveStorage {
2496 persisted: Mutex::new(Some(previous.clone())),
2497 });
2498 let repo = test_repo(storage);
2499 cache_put(&repo, &previous);
2500
2501 let mut current = previous;
2502 current.model = "current".to_string();
2503 repo.save_and_cache(&mut current).await;
2504 assert_eq!(
2505 read_cached_session(repo.cache(), id)
2506 .expect("refreshed cache")
2507 .model,
2508 "current",
2509 "fire-and-forget save must retain its existing cache-on-failure behavior"
2510 );
2511 }
2512
2513 #[tokio::test]
2514 async fn checkpoint_leaves_existing_cache_untouched_when_storage_fails() {
2515 let id = "checkpoint-failure";
2516 let previous = Session::new(id, "previous");
2517 let storage: Arc<dyn Storage> = Arc::new(FailingSaveStorage {
2518 persisted: Mutex::new(Some(previous.clone())),
2519 });
2520 let repo = test_repo(storage);
2521 cache_put(&repo, &previous);
2522
2523 let mut current = previous.clone();
2524 current.model = "current".to_string();
2525 let result = bamboo_domain::RuntimeSessionPersistence::checkpoint_runtime_session(
2526 &repo,
2527 &mut current,
2528 )
2529 .await;
2530 assert!(result.is_err());
2531 assert_eq!(
2532 read_cached_session(repo.cache(), id)
2533 .expect("existing cache")
2534 .model,
2535 "previous",
2536 "checkpoint must publish only after a durable commit"
2537 );
2538 }
2539
2540 #[tokio::test]
2541 async fn retrieval_window_checkpoint_publishes_the_exact_archived_snapshot() {
2542 let temp = tempfile::tempdir().unwrap();
2543 let storage: Arc<dyn Storage> = Arc::new(
2544 bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
2545 .await
2546 .expect("SessionStoreV2"),
2547 );
2548 let repo = test_repo(storage.clone());
2549 let id = "retrieval-window-repository-checkpoint";
2550 let mut expected = Session::new(id, "model");
2551 expected.add_message(bamboo_agent_core::Message::user("archive candidate"));
2552 expected.add_message(bamboo_agent_core::Message::assistant("retain", None));
2553 storage.save_session(&expected).await.unwrap();
2554 cache_put(&repo, &expected);
2555 let mut staged = staged_retrieval_window_archive(&expected);
2556
2557 let outcome = bamboo_domain::RuntimeSessionPersistence::checkpoint_retrieval_window(
2558 &repo,
2559 &expected,
2560 &mut staged,
2561 )
2562 .await
2563 .unwrap();
2564
2565 assert_eq!(outcome, RetrievalWindowCheckpointOutcome::Committed);
2566 let durable = storage.load_session(id).await.unwrap().unwrap();
2567 let cached = read_cached_session(repo.cache(), id).expect("cached archived Session");
2568 for (tier, saved) in [("durable", durable), ("cache", cached)] {
2569 assert!(saved.messages[0].compressed, "tier={tier}");
2570 assert_eq!(
2571 saved.messages[0].compressed_by_event_id, staged.messages[0].compressed_by_event_id,
2572 "tier={tier}"
2573 );
2574 assert_eq!(
2575 serde_json::to_value(&saved).unwrap(),
2576 serde_json::to_value(&staged).unwrap(),
2577 "tier={tier}"
2578 );
2579 }
2580 }
2581
2582 #[tokio::test]
2583 async fn legacy_clear_uses_durable_cas_and_never_erases_concurrent_append_from_stale_cache() {
2584 let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
2585 let repo = test_repo(storage.clone());
2586 let id = "legacy-cas-race";
2587 let expected = vec![serde_json::json!({"content": "first"})];
2588
2589 let mut stale_cache = Session::new(id, "m");
2590 stale_cache.set_pending_injected_messages(expected.clone());
2591 cache_put(&repo, &stale_cache);
2592
2593 let mut durable = stale_cache.clone();
2594 durable.set_pending_injected_messages(vec![
2595 serde_json::json!({"content": "first"}),
2596 serde_json::json!({"content": "concurrent"}),
2597 ]);
2598 storage.save_session(&durable).await.unwrap();
2599
2600 let cleared = bamboo_domain::RuntimeSessionPersistence::clear_legacy_pending_messages(
2601 &repo, id, &expected,
2602 )
2603 .await
2604 .unwrap();
2605 assert!(!cleared, "the durable compare-and-clear must reject drift");
2606 assert_eq!(
2607 storage
2608 .load_session(id)
2609 .await
2610 .unwrap()
2611 .unwrap()
2612 .pending_injected_messages()
2613 .unwrap(),
2614 durable.pending_injected_messages().unwrap(),
2615 "the concurrent durable append must remain intact"
2616 );
2617 assert_eq!(
2618 read_cached_session(repo.cache(), id)
2619 .unwrap()
2620 .pending_injected_messages()
2621 .unwrap(),
2622 expected,
2623 "a failed CAS must not mutate the existing cache"
2624 );
2625 }
2626
2627 #[tokio::test]
2628 async fn successful_legacy_clear_refreshes_stale_cache_from_durable_state() {
2629 let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
2630 let repo = test_repo(storage.clone());
2631 let id = "legacy-cas-success";
2632 let expected = vec![serde_json::json!({"content": "first"})];
2633
2634 let mut stale_cache = Session::new(id, "stale-model");
2635 stale_cache.set_pending_injected_messages(expected.clone());
2636 cache_put(&repo, &stale_cache);
2637
2638 let mut durable = Session::new(id, "durable-model");
2639 durable.set_pending_injected_messages(expected.clone());
2640 durable
2641 .metadata
2642 .insert("durable-only".to_string(), "keep".to_string());
2643 storage.save_session(&durable).await.unwrap();
2644
2645 assert!(
2646 bamboo_domain::RuntimeSessionPersistence::clear_legacy_pending_messages(
2647 &repo, id, &expected,
2648 )
2649 .await
2650 .unwrap()
2651 );
2652 let cached = read_cached_session(repo.cache(), id).unwrap();
2653 assert!(!cached.has_pending_injected_messages());
2654 assert_eq!(cached.model, "durable-model");
2655 assert_eq!(
2656 cached.metadata.get("durable-only").map(String::as_str),
2657 Some("keep")
2658 );
2659 }
2660
2661 #[tokio::test]
2662 async fn repository_final_save_publishes_host_checkpointed_parent_question() {
2663 use bamboo_domain::{
2664 FunctionCall, ParentQuestion, ParentQuestionCheckpointV1, PendingQuestion,
2665 PendingQuestionSource, ToolCall,
2666 };
2667
2668 let temp = tempfile::tempdir().unwrap();
2669 let store = Arc::new(
2670 bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
2671 .await
2672 .unwrap(),
2673 );
2674 let storage: Arc<dyn Storage> = store.clone();
2675 let repo = test_repo(storage);
2676 let parent = Session::new("repository-question-parent", "model");
2677 store.save_session(&parent).await.unwrap();
2678 let mut child =
2679 Session::new_child_of("repository-question-child", &parent, "model", "assignment");
2680 child.add_message(bamboo_agent_core::Message::user("assignment"));
2681 child.set_last_run_status("running");
2682 store.save_session(&child).await.unwrap();
2683 cache_put(&repo, &child);
2684
2685 let assistant = bamboo_agent_core::Message::assistant(
2686 "",
2687 Some(vec![ToolCall {
2688 id: "question-call".into(),
2689 tool_type: "function".into(),
2690 function: FunctionCall {
2691 name: "AgenticQuestion".into(),
2692 arguments: "{}".into(),
2693 },
2694 }]),
2695 );
2696 let result = bamboo_agent_core::Message::tool_result_with_status(
2697 "question-call",
2698 "Clarification needed: Which option?",
2699 true,
2700 );
2701 let pending = PendingQuestion {
2702 tool_call_id: "question-call".into(),
2703 tool_name: "AgenticQuestion".into(),
2704 question: "Which option?".into(),
2705 options: vec!["A".into(), "B".into()],
2706 allow_custom: true,
2707 source: PendingQuestionSource::AgenticClarification,
2708 };
2709 let observation = ParentQuestionCheckpointV1 {
2710 version: 1,
2711 prefix_message_count: child.messages.len(),
2712 prefix_digest: ParentQuestion::prefix_digest(&child.messages).unwrap(),
2713 suffix: vec![assistant, result.clone()],
2714 tool_call_id: pending.tool_call_id.clone(),
2715 tool_result_message_id: result.id.clone(),
2716 question_digest: ParentQuestion::question_digest(&pending, &result.id),
2717 pending,
2718 };
2719 let (_, question) = store
2720 .checkpoint_parent_question(&child, &parent, &observation)
2721 .await
2722 .unwrap()
2723 .unwrap();
2724
2725 let mut old_run = child.clone();
2726 old_run.metadata.insert(
2727 "runtime.actor_parent_question_handoff".into(),
2728 "true".into(),
2729 );
2730 old_run.set_last_run_status("suspended");
2731 old_run.add_message(bamboo_agent_core::Message::assistant(
2732 "old terminal text",
2733 None,
2734 ));
2735 cache_put(&repo, &old_run);
2736 bamboo_domain::RuntimeSessionPersistence::save_runtime_session(&repo, &mut old_run)
2737 .await
2738 .unwrap();
2739
2740 let cached = repo.load(&child.id).await.unwrap();
2741 let reopened = bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
2742 .await
2743 .unwrap()
2744 .load_session(&child.id)
2745 .await
2746 .unwrap()
2747 .unwrap();
2748 for saved in [&cached, &reopened] {
2749 assert_eq!(saved.last_run_status().as_deref(), Some("suspended"));
2750 assert_eq!(saved.messages.len(), child.messages.len() + 2);
2751 assert!(!saved
2752 .messages
2753 .iter()
2754 .any(|message| message.content == "old terminal text"));
2755 assert_eq!(
2756 saved.pending_question.as_ref().unwrap().source,
2757 PendingQuestionSource::DirectParent
2758 );
2759 assert_eq!(
2760 ParentQuestion::for_pending(&parent, saved).unwrap().id,
2761 question.id
2762 );
2763 }
2764 }
2765}
2766
2767fn root_actor_binding_error(reason: impl std::fmt::Display) -> std::io::Error {
2768 std::io::Error::new(
2769 std::io::ErrorKind::WouldBlock,
2770 bamboo_domain::SessionAuthorityConflict(format!("Root Actor execution rejected: {reason}")),
2771 )
2772}