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