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