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