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
1134#[cfg(test)]
1135mod tests {
1136 use super::*;
1137 use bamboo_agent_core::storage::Storage;
1138 use chrono::Utc;
1139 use std::collections::HashMap;
1140 use std::sync::atomic::{AtomicBool, Ordering};
1141 use std::sync::{Condvar, Mutex};
1142 use std::time::Duration;
1143
1144 #[derive(Default)]
1145 struct MapStorage {
1146 sessions: Mutex<HashMap<String, Session>>,
1147 fail_pair_commit: AtomicBool,
1148 }
1149
1150 struct FailingSaveStorage {
1151 persisted: Mutex<Option<Session>>,
1152 }
1153
1154 struct AmbiguousReplaySaveStorage {
1155 persisted: Mutex<Option<Session>>,
1156 }
1157
1158 #[async_trait::async_trait]
1159 impl Storage for MapStorage {
1160 async fn save_session(&self, session: &Session) -> std::io::Result<()> {
1161 self.sessions
1162 .lock()
1163 .unwrap()
1164 .insert(session.id.clone(), session.clone());
1165 Ok(())
1166 }
1167 async fn load_session(&self, session_id: &str) -> std::io::Result<Option<Session>> {
1168 Ok(self.sessions.lock().unwrap().get(session_id).cloned())
1169 }
1170 async fn delete_session(&self, session_id: &str) -> std::io::Result<bool> {
1171 Ok(self.sessions.lock().unwrap().remove(session_id).is_some())
1172 }
1173
1174 async fn save_task_control_plane_if_matches(
1175 &self,
1176 original: &Session,
1177 updated: &Session,
1178 ) -> std::io::Result<bool> {
1179 let mut sessions = self.sessions.lock().unwrap();
1180 let matches_original = sessions.get(&original.id).is_some_and(|current| {
1181 current.task_list_version_meta() == original.task_list_version_meta()
1182 && serde_json::to_value(¤t.task_list).ok()
1183 == serde_json::to_value(&original.task_list).ok()
1184 });
1185 if !matches_original {
1186 return Ok(false);
1187 }
1188 let Some(current) = sessions.get(&original.id).cloned() else {
1189 return Ok(false);
1190 };
1191 let mut committed = current;
1192 committed.task_list = updated.task_list.clone();
1193 if let Some(version) = updated.task_list_version_meta() {
1194 committed.set_task_list_version_meta(version);
1195 }
1196 sessions.insert(committed.id.clone(), committed);
1197 Ok(true)
1198 }
1199
1200 async fn save_task_control_planes_atomically(
1201 &self,
1202 first_original: &Session,
1203 first_updated: &Session,
1204 second_original: &Session,
1205 second_updated: &Session,
1206 ) -> std::io::Result<bool> {
1207 if self.fail_pair_commit.swap(false, Ordering::SeqCst) {
1208 return Err(std::io::Error::other(
1209 "injected paired Task transaction failure",
1210 ));
1211 }
1212 let mut sessions = self.sessions.lock().unwrap();
1213 let matches_original = |current: Option<&Session>, original: &Session| {
1214 current.is_some_and(|current| {
1215 current.task_list_version_meta() == original.task_list_version_meta()
1216 && serde_json::to_value(¤t.task_list).ok()
1217 == serde_json::to_value(&original.task_list).ok()
1218 })
1219 };
1220 if !matches_original(sessions.get(&first_original.id), first_original)
1221 || !matches_original(sessions.get(&second_original.id), second_original)
1222 {
1223 return Ok(false);
1224 }
1225 sessions.insert(first_updated.id.clone(), first_updated.clone());
1226 sessions.insert(second_updated.id.clone(), second_updated.clone());
1227 Ok(true)
1228 }
1229 }
1230
1231 #[async_trait::async_trait]
1232 impl Storage for FailingSaveStorage {
1233 async fn save_session(&self, _session: &Session) -> std::io::Result<()> {
1234 Err(std::io::Error::other("injected save failure"))
1235 }
1236
1237 async fn load_session(&self, _session_id: &str) -> std::io::Result<Option<Session>> {
1238 Ok(self.persisted.lock().unwrap().clone())
1239 }
1240
1241 async fn delete_session(&self, _session_id: &str) -> std::io::Result<bool> {
1242 Ok(false)
1243 }
1244 }
1245
1246 #[async_trait::async_trait]
1247 impl Storage for AmbiguousReplaySaveStorage {
1248 async fn save_session(&self, session: &Session) -> std::io::Result<()> {
1249 *self.persisted.lock().unwrap() = Some(session.clone());
1250 Err(std::io::Error::other(
1251 "injected post-commit replay save error",
1252 ))
1253 }
1254
1255 async fn load_session(&self, _session_id: &str) -> std::io::Result<Option<Session>> {
1256 Ok(self.persisted.lock().unwrap().clone())
1257 }
1258
1259 async fn delete_session(&self, _session_id: &str) -> std::io::Result<bool> {
1260 Ok(false)
1261 }
1262 }
1263
1264 fn test_repo(storage: Arc<dyn Storage>) -> SessionRepository {
1265 let cache: SessionCache = Arc::default();
1266 let persistence = Arc::new(LockedSessionStore::new(storage.clone()));
1267 SessionRepository::new(cache, storage, persistence)
1268 }
1269
1270 fn cache_put(repo: &SessionRepository, session: &Session) {
1271 repo.cache().insert(
1272 session.id.clone(),
1273 Arc::new(crate::SessionSnapshot::new(session.clone())),
1274 );
1275 }
1276
1277 #[tokio::test]
1278 async fn merged_root_prefers_live_recreated_birth_even_when_epoch_resets() {
1279 let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
1280 let repo = test_repo(storage.clone());
1281 let id = "root-recreated-same-id";
1282 let mut stale = Session::new(id, "stale");
1283 stale.root_mode_transition_epoch = 3;
1284 stale.updated_at += chrono::Duration::hours(1);
1285 cache_put(&repo, &stale);
1286
1287 let mut recreated = Session::new(id, "recreated");
1288 recreated.created_at = stale.created_at + chrono::Duration::seconds(1);
1289 recreated.updated_at = stale.created_at + chrono::Duration::seconds(1);
1290 storage.save_session(&recreated).await.unwrap();
1291
1292 let loaded = repo.load_merged_checked(id).await.unwrap().unwrap();
1293 assert_eq!(loaded.model, "recreated");
1294 assert_eq!(loaded.root_mode_transition_epoch, 0);
1295 assert_eq!(
1296 loaded.root_mode_birth_token(),
1297 recreated.root_mode_birth_token()
1298 );
1299 let cached = read_cached_session(repo.cache(), id).expect("new birth cached");
1300 assert_eq!(cached.model, "recreated");
1301 assert_eq!(
1302 cached.root_mode_birth_token(),
1303 recreated.root_mode_birth_token()
1304 );
1305 }
1306
1307 #[tokio::test]
1308 async fn merged_root_caches_higher_durable_epoch_despite_clock_skew() {
1309 let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
1310 let repo = test_repo(storage.clone());
1311 let id = "root-mode-clock-skew";
1312 let mut stale = Session::new(id, "model");
1313 stale.updated_at += chrono::Duration::hours(1);
1314 cache_put(&repo, &stale);
1315
1316 let mut durable = stale.clone();
1317 durable.root_mode_transition_epoch = 1;
1318 durable.root_orchestration_only = true;
1319 durable.root_tool_authority_revision = 1;
1320 durable.updated_at = durable.created_at;
1321 storage.save_session(&durable).await.unwrap();
1322
1323 let loaded = repo.load_merged_checked(id).await.unwrap().unwrap();
1324 assert_eq!(loaded.root_mode_transition_epoch, 1);
1325 assert!(loaded.root_orchestration_only);
1326 let cached = read_cached_session(repo.cache(), id).expect("higher epoch cached");
1327 assert_eq!(cached.root_mode_transition_epoch, 1);
1328 assert!(cached.root_orchestration_only);
1329 }
1330
1331 #[tokio::test]
1332 async fn runtime_merge_save_adopts_root_mode_terminal_tuple_even_for_policy_neutral_fence() {
1333 use bamboo_domain::{RootModeOperationAction, RootModeOperationRequest};
1334
1335 for action in [
1336 RootModeOperationAction::Select,
1337 RootModeOperationAction::Recover,
1338 ] {
1339 let home = tempfile::tempdir().unwrap();
1340 let concrete = Arc::new(
1341 bamboo_storage::SessionStoreV2::new(home.path().into())
1342 .await
1343 .unwrap(),
1344 );
1345 let storage: Arc<dyn Storage> = concrete.clone();
1346 let repo = test_repo(storage.clone());
1347 let mut running = Session::new("runtime-root-mode-merge", "model");
1348 storage.save_session(&running).await.unwrap();
1349 cache_put(&repo, &running);
1350 let request = RootModeOperationRequest {
1351 session_id: running.id.clone(),
1352 operation_id: format!("0:{}", uuid::Uuid::new_v4()),
1353 birth_token: running.root_mode_birth_token(),
1354 expected_epoch: 0,
1355 requested_enabled: true,
1356 action,
1357 };
1358 storage.root_mode_operation(&request).await.unwrap();
1359
1360 running.add_message(bamboo_agent_core::Message::assistant("run result", None));
1363 repo.save(&mut running).await.unwrap();
1364 assert_eq!(running.root_mode_transition_epoch, 1);
1365 assert_eq!(running.root_mode_operations.len(), 1);
1366 assert_eq!(
1367 running.root_orchestration_only_enabled(),
1368 action == RootModeOperationAction::Select
1369 );
1370 let durable = storage.load_session(&running.id).await.unwrap().unwrap();
1371 assert_eq!(durable.root_mode_transition_epoch, 1);
1372 assert_eq!(durable.root_mode_operations, running.root_mode_operations);
1373 assert_eq!(durable.messages.last().unwrap().content, "run result");
1374 concrete.flush_search_index().await;
1375 }
1376 }
1377
1378 fn task_list(session_id: &str, title: &str) -> bamboo_domain::TaskList {
1379 let now = Utc::now();
1380 bamboo_domain::TaskList {
1381 session_id: session_id.to_string(),
1382 title: title.to_string(),
1383 items: Vec::new(),
1384 created_at: now,
1385 updated_at: now,
1386 }
1387 }
1388
1389 fn staged_retrieval_window_archive(expected: &Session) -> Session {
1390 let mut staged = expected.clone();
1391 let mut event = bamboo_domain::CompressionEvent::new(
1392 1,
1393 1,
1394 80.0,
1395 60.0,
1396 0,
1397 bamboo_domain::CompressionTriggerType::Auto,
1398 0.0,
1399 None,
1400 0,
1401 );
1402 event.kind = bamboo_domain::CompressionEventKind::RetrievalWindow;
1403 staged.messages[0].compressed = true;
1404 staged.messages[0].compressed_by_event_id = Some(event.id.clone());
1405 staged.compression_events.push(event);
1406 staged.reset_model_context_epoch(bamboo_domain::ModelContextResetReason::Compression);
1407 staged
1408 }
1409
1410 #[tokio::test]
1411 async fn paired_task_cas_narrowly_updates_child_and_root_cache() {
1412 let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
1413 let repo = test_repo(storage.clone());
1414 let root_id = "paired-cache-root";
1415 let child_id = "paired-cache-child";
1416 let expected_task_list = task_list(root_id, "old shared");
1417
1418 let mut root = Session::new(root_id, "model");
1419 root.add_message(bamboo_agent_core::Message::user("root transcript"));
1420 root.metadata
1421 .insert("unrelated.root".to_string(), "keep".to_string());
1422 root.set_task_list(expected_task_list.clone());
1423 root.set_task_list_version_meta("1");
1424 storage.save_session(&root).await.unwrap();
1425 cache_put(&repo, &root);
1426
1427 let mut child = Session::new_child(child_id, root_id, "model", "child");
1428 child.add_message(bamboo_agent_core::Message::user("child transcript"));
1429 child
1430 .metadata
1431 .insert("unrelated.child".to_string(), "keep".to_string());
1432 child.set_task_list(expected_task_list.clone());
1433 child.set_task_list_version_meta("1");
1434 storage.save_session(&child).await.unwrap();
1435 cache_put(&repo, &child);
1436
1437 assert!(
1438 bamboo_domain::RuntimeSessionPersistence::update_task_list_control_planes_if_version(
1439 &repo,
1440 child_id,
1441 root_id,
1442 "1",
1443 &expected_task_list,
1444 &task_list(root_id, "evaluated"),
1445 "2",
1446 )
1447 .await
1448 .expect("paired cache CAS succeeds")
1449 );
1450
1451 let cached_root = read_cached_session(repo.cache(), root_id).expect("cached root");
1452 let cached_child = read_cached_session(repo.cache(), child_id).expect("cached child");
1453 for (session, transcript, metadata_key) in [
1454 (&cached_root, "root transcript", "unrelated.root"),
1455 (&cached_child, "child transcript", "unrelated.child"),
1456 ] {
1457 assert_eq!(session.task_list_version_meta().as_deref(), Some("2"));
1458 assert_eq!(
1459 session.task_list.as_ref().map(|list| list.title.as_str()),
1460 Some("evaluated")
1461 );
1462 assert_eq!(session.messages.len(), 1);
1463 assert_eq!(session.messages[0].content, transcript);
1464 assert_eq!(
1465 session.metadata.get(metadata_key).map(String::as_str),
1466 Some("keep")
1467 );
1468 }
1469 }
1470
1471 #[tokio::test]
1472 async fn failed_paired_task_transaction_does_not_publish_child_or_root_cache() {
1473 let concrete = Arc::new(MapStorage::default());
1474 let storage: Arc<dyn Storage> = concrete.clone();
1475 let repo = test_repo(storage.clone());
1476 let root_id = "paired-cache-failure-root";
1477 let child_id = "paired-cache-failure-child";
1478 let expected_task_list = task_list(root_id, "old shared");
1479
1480 let mut root = Session::new(root_id, "model");
1481 root.add_message(bamboo_agent_core::Message::user("root transcript"));
1482 root.metadata
1483 .insert("unrelated.root".to_string(), "keep".to_string());
1484 root.set_task_list(expected_task_list.clone());
1485 root.set_task_list_version_meta("1");
1486 storage.save_session(&root).await.unwrap();
1487 cache_put(&repo, &root);
1488
1489 let mut child = Session::new_child(child_id, root_id, "model", "child");
1490 child.add_message(bamboo_agent_core::Message::user("child transcript"));
1491 child
1492 .metadata
1493 .insert("unrelated.child".to_string(), "keep".to_string());
1494 child.set_task_list(expected_task_list.clone());
1495 child.set_task_list_version_meta("1");
1496 storage.save_session(&child).await.unwrap();
1497 cache_put(&repo, &child);
1498
1499 concrete.fail_pair_commit.store(true, Ordering::SeqCst);
1500 let error =
1501 bamboo_domain::RuntimeSessionPersistence::update_task_list_control_planes_if_version(
1502 &repo,
1503 child_id,
1504 root_id,
1505 "1",
1506 &expected_task_list,
1507 &task_list(root_id, "must not publish"),
1508 "2",
1509 )
1510 .await
1511 .expect_err("paired durable transaction fails");
1512 assert!(error.to_string().contains("injected paired"));
1513
1514 for (id, expected_title, transcript, metadata_key) in [
1515 (root_id, "old shared", "root transcript", "unrelated.root"),
1516 (
1517 child_id,
1518 "old shared",
1519 "child transcript",
1520 "unrelated.child",
1521 ),
1522 ] {
1523 let cached = read_cached_session(repo.cache(), id).expect("cached session remains");
1524 assert_eq!(cached.task_list_version_meta().as_deref(), Some("1"));
1525 assert_eq!(
1526 cached.task_list.as_ref().map(|list| list.title.as_str()),
1527 Some(expected_title)
1528 );
1529 assert_eq!(cached.messages[0].content, transcript);
1530 assert_eq!(
1531 cached.metadata.get(metadata_key).map(String::as_str),
1532 Some("keep")
1533 );
1534
1535 let durable = storage.load_session(id).await.unwrap().unwrap();
1536 assert_eq!(durable.task_list_version_meta().as_deref(), Some("1"));
1537 assert_eq!(
1538 durable.task_list.as_ref().map(|list| list.title.as_str()),
1539 Some(expected_title)
1540 );
1541 }
1542 }
1543
1544 fn durable_cache_fence(
1545 operation: impl Into<String>,
1546 marker: impl Into<String>,
1547 ) -> (
1548 PostDurableHook,
1549 tokio::sync::oneshot::Receiver<()>,
1550 Arc<(Mutex<bool>, Condvar)>,
1551 ) {
1552 let operation = operation.into();
1553 let marker = marker.into();
1554 let (started_tx, started_rx) = tokio::sync::oneshot::channel();
1555 let started_tx = Arc::new(Mutex::new(Some(started_tx)));
1556 let release = Arc::new((Mutex::new(false), Condvar::new()));
1557 let hook_release = release.clone();
1558 let hook: PostDurableHook = Arc::new(move |actual_operation, actual_marker| {
1559 if actual_operation != operation || actual_marker != marker {
1560 return;
1561 }
1562 if let Some(started_tx) = started_tx.lock().unwrap().take() {
1563 started_tx.send(()).expect("fence observer still present");
1564 }
1565 let (released, wake) = &*hook_release;
1566 let mut released = released.lock().unwrap();
1567 while !*released {
1568 released = wake.wait(released).unwrap();
1569 }
1570 });
1571 (hook, started_rx, release)
1572 }
1573
1574 fn release_fence(release: &Arc<(Mutex<bool>, Condvar)>) {
1575 let (released, wake) = &**release;
1576 *released.lock().unwrap() = true;
1577 wake.notify_all();
1578 }
1579
1580 async fn assert_second_write_waits_for_cache_publish<T>(
1581 first: tokio::task::JoinHandle<std::io::Result<T>>,
1582 mut second: tokio::task::JoinHandle<std::io::Result<T>>,
1583 release: Arc<(Mutex<bool>, Condvar)>,
1584 ) -> (T, T) {
1585 let second_before_release =
1586 tokio::time::timeout(Duration::from_millis(100), &mut second).await;
1587 let completed_before_release = second_before_release.is_ok();
1588 release_fence(&release);
1589
1590 let first = first
1591 .await
1592 .expect("first writer joins")
1593 .expect("first writer succeeds");
1594 let second = match second_before_release {
1595 Ok(joined) => joined
1596 .expect("second writer joins")
1597 .expect("second writer succeeds"),
1598 Err(_) => second
1599 .await
1600 .expect("second writer joins")
1601 .expect("second writer succeeds"),
1602 };
1603 assert!(
1604 !completed_before_release,
1605 "the second write must remain behind the first write's durable-to-cache fence"
1606 );
1607 (first, second)
1608 }
1609
1610 #[derive(Clone, Copy, Debug)]
1611 enum FullSaveRoute {
1612 InherentSave,
1613 SaveAndCache,
1614 RuntimePersistence,
1615 }
1616
1617 impl FullSaveRoute {
1618 fn operation(self) -> &'static str {
1619 match self {
1620 Self::InherentSave => "save_full",
1621 Self::SaveAndCache => "save_and_cache",
1622 Self::RuntimePersistence => "save_runtime_session",
1623 }
1624 }
1625
1626 fn name(self) -> &'static str {
1627 match self {
1628 Self::InherentSave => "inherent",
1629 Self::SaveAndCache => "save-and-cache",
1630 Self::RuntimePersistence => "runtime-persistence",
1631 }
1632 }
1633 }
1634
1635 async fn assert_full_save_route_serializes_cache_publish(route: FullSaveRoute) {
1636 let temp = tempfile::tempdir().unwrap();
1637 let concrete_storage = Arc::new(
1638 bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
1639 .await
1640 .expect("SessionStoreV2"),
1641 );
1642 let storage: Arc<dyn Storage> = concrete_storage;
1643 let id = format!("root-full-cache-order-{}", route.name());
1644 let (hook, first_durable, release) = durable_cache_fence(route.operation(), id.clone());
1645 let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
1646
1647 let mut initial = Session::new(&id, "model");
1648 initial.add_message(bamboo_agent_core::Message::user("durable transcript"));
1649 initial.set_task_list(task_list(&id, "initial"));
1650 initial.set_task_list_version_meta("0");
1651 initial
1652 .metadata
1653 .insert("unrelated.runtime".to_string(), "keep".to_string());
1654 storage.save_session(&initial).await.unwrap();
1655 cache_put(&repo, &initial);
1656
1657 let first_repo = repo.clone();
1658 let mut root_snapshot = initial.clone();
1659 root_snapshot.add_message(bamboo_agent_core::Message::assistant(
1660 "full-save transcript suffix",
1661 None,
1662 ));
1663 root_snapshot.set_task_list(task_list(&id, "root"));
1664 root_snapshot.set_task_list_version_meta("1");
1665 let first = tokio::spawn(async move {
1666 match route {
1667 FullSaveRoute::InherentSave => first_repo.save(&mut root_snapshot).await,
1668 FullSaveRoute::SaveAndCache => {
1669 first_repo.save_and_cache(&mut root_snapshot).await;
1670 Ok(())
1671 }
1672 FullSaveRoute::RuntimePersistence => {
1673 bamboo_domain::RuntimeSessionPersistence::save_runtime_session(
1674 first_repo.as_ref(),
1675 &mut root_snapshot,
1676 )
1677 .await
1678 }
1679 }
1680 });
1681 first_durable
1682 .await
1683 .expect("root full durable write reached");
1684
1685 let second_repo = repo.clone();
1686 let second_id = id.clone();
1687 let child_task_list = task_list(&id, "child");
1688 let second = tokio::spawn(async move {
1689 bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
1690 second_repo.as_ref(),
1691 &second_id,
1692 &child_task_list,
1693 "2",
1694 )
1695 .await
1696 .map(|updated| assert!(updated, "root must exist"))
1697 });
1698 assert_second_write_waits_for_cache_publish(first, second, release).await;
1699
1700 let durable = storage.load_session(&id).await.unwrap().unwrap();
1701 let cached = read_cached_session(repo.cache(), &id).expect("cached root");
1702 for (tier, session) in [("durable", durable), ("cache", cached)] {
1703 assert_eq!(
1704 session.task_list_version_meta().as_deref(),
1705 Some("2"),
1706 "{route:?} {tier} must retain the child transaction"
1707 );
1708 assert_eq!(
1709 session.task_list.as_ref().map(|list| list.title.as_str()),
1710 Some("child"),
1711 "{route:?} {tier} must retain the child transaction"
1712 );
1713 assert_eq!(
1714 session
1715 .metadata
1716 .get("unrelated.runtime")
1717 .map(String::as_str),
1718 Some("keep"),
1719 "{route:?} {tier} must preserve unrelated runtime state"
1720 );
1721 assert_eq!(
1722 session.messages.len(),
1723 2,
1724 "{route:?} {tier} must preserve the full-save transcript"
1725 );
1726 }
1727 }
1728
1729 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1730 async fn root_full_saves_and_child_task_patch_share_publish_order() {
1731 for route in [
1732 FullSaveRoute::InherentSave,
1733 FullSaveRoute::SaveAndCache,
1734 FullSaveRoute::RuntimePersistence,
1735 ] {
1736 assert_full_save_route_serializes_cache_publish(route).await;
1737 }
1738 }
1739
1740 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1741 async fn checkpoint_and_child_task_patch_share_publish_order() {
1742 let temp = tempfile::tempdir().unwrap();
1743 let concrete_storage = Arc::new(
1744 bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
1745 .await
1746 .expect("SessionStoreV2"),
1747 );
1748 let storage: Arc<dyn Storage> = concrete_storage;
1749 let id = "checkpoint-cache-order";
1750 let (hook, checkpoint_durable, release) = durable_cache_fence("checkpoint", id);
1751 let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
1752
1753 let mut initial = Session::new(id, "model");
1754 initial.add_message(bamboo_agent_core::Message::user("durable transcript"));
1755 initial.set_task_list(task_list(id, "initial"));
1756 initial.set_task_list_version_meta("0");
1757 initial
1758 .metadata
1759 .insert("unrelated.runtime".to_string(), "keep".to_string());
1760 storage.save_session(&initial).await.unwrap();
1761 cache_put(&repo, &initial);
1762
1763 let checkpoint_repo = repo.clone();
1764 let mut checkpoint_snapshot = initial.clone();
1765 checkpoint_snapshot.add_message(bamboo_agent_core::Message::assistant(
1766 "checkpoint transcript suffix",
1767 None,
1768 ));
1769 checkpoint_snapshot.set_task_list(task_list(id, "checkpoint"));
1770 checkpoint_snapshot.set_task_list_version_meta("1");
1771 let checkpoint = tokio::spawn(async move {
1772 bamboo_domain::RuntimeSessionPersistence::checkpoint_runtime_session(
1773 checkpoint_repo.as_ref(),
1774 &mut checkpoint_snapshot,
1775 )
1776 .await
1777 });
1778 checkpoint_durable
1779 .await
1780 .expect("checkpoint durable write reached");
1781
1782 let child_repo = repo.clone();
1783 let child_task_list = task_list(id, "child");
1784 let child_patch = tokio::spawn(async move {
1785 bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
1786 child_repo.as_ref(),
1787 id,
1788 &child_task_list,
1789 "2",
1790 )
1791 .await
1792 .map(|updated| assert!(updated, "root must exist"))
1793 });
1794 assert_second_write_waits_for_cache_publish(checkpoint, child_patch, release).await;
1795
1796 let durable = storage.load_session(id).await.unwrap().unwrap();
1797 let cached = read_cached_session(repo.cache(), id).expect("cached root");
1798 for (tier, session) in [("durable", durable), ("cache", cached)] {
1799 assert_eq!(
1800 session.task_list_version_meta().as_deref(),
1801 Some("2"),
1802 "{tier} must retain the child transaction"
1803 );
1804 assert_eq!(
1805 session.task_list.as_ref().map(|list| list.title.as_str()),
1806 Some("child"),
1807 "{tier} must retain the child transaction"
1808 );
1809 assert_eq!(
1810 session
1811 .metadata
1812 .get("unrelated.runtime")
1813 .map(String::as_str),
1814 Some("keep"),
1815 "{tier} must preserve unrelated runtime state"
1816 );
1817 assert_eq!(
1818 session.messages.len(),
1819 2,
1820 "{tier} must preserve the checkpoint transcript"
1821 );
1822 }
1823 }
1824
1825 #[derive(Clone, Copy, Debug)]
1826 enum CacheBackfillRoute {
1827 Load,
1828 TryLoad,
1829 }
1830
1831 impl CacheBackfillRoute {
1832 fn operation(self) -> &'static str {
1833 match self {
1834 Self::Load => "load",
1835 Self::TryLoad => "try_load",
1836 }
1837 }
1838
1839 fn name(self) -> &'static str {
1840 match self {
1841 Self::Load => "load",
1842 Self::TryLoad => "try-load",
1843 }
1844 }
1845 }
1846
1847 async fn assert_cache_backfill_serializes_with_task_patch(route: CacheBackfillRoute) {
1848 let temp = tempfile::tempdir().unwrap();
1849 let concrete_storage = Arc::new(
1850 bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
1851 .await
1852 .expect("SessionStoreV2"),
1853 );
1854 let storage: Arc<dyn Storage> = concrete_storage;
1855 let id = format!("cache-backfill-order-{}", route.name());
1856 let (hook, loaded_old_durable, release) =
1857 durable_cache_fence(route.operation(), id.clone());
1858 let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
1859
1860 let mut initial = Session::new(&id, "model");
1861 initial.set_task_list(task_list(&id, "initial"));
1862 initial.set_task_list_version_meta("0");
1863 storage.save_session(&initial).await.unwrap();
1864 assert!(
1865 read_cached_session(repo.cache(), &id).is_none(),
1866 "the race requires a genuine cache miss"
1867 );
1868
1869 let load_repo = repo.clone();
1870 let load_id = id.clone();
1871 let load = tokio::spawn(async move {
1872 let loaded = match route {
1873 CacheBackfillRoute::Load => load_repo.load(&load_id).await,
1874 CacheBackfillRoute::TryLoad => {
1875 load_repo.try_load(&load_id).await.expect("storage load")
1876 }
1877 };
1878 assert!(loaded.is_some(), "seeded session must load");
1879 Ok(())
1880 });
1881 loaded_old_durable
1882 .await
1883 .expect("old durable snapshot loaded");
1884
1885 let patch_repo = repo.clone();
1886 let patch_id = id.clone();
1887 let child_task_list = task_list(&id, "child");
1888 let patch = tokio::spawn(async move {
1889 bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
1890 patch_repo.as_ref(),
1891 &patch_id,
1892 &child_task_list,
1893 "1",
1894 )
1895 .await
1896 .map(|updated| assert!(updated, "root must exist"))
1897 });
1898 assert_second_write_waits_for_cache_publish(load, patch, release).await;
1899
1900 let durable = storage.load_session(&id).await.unwrap().unwrap();
1901 let cached = read_cached_session(repo.cache(), &id).expect("backfilled cache");
1902 for (tier, session) in [("durable", durable), ("cache", cached)] {
1903 assert_eq!(
1904 session.task_list_version_meta().as_deref(),
1905 Some("1"),
1906 "{route:?} {tier} must retain the child transaction"
1907 );
1908 assert_eq!(
1909 session.task_list.as_ref().map(|list| list.title.as_str()),
1910 Some("child"),
1911 "{route:?} {tier} must retain the child transaction"
1912 );
1913 }
1914 }
1915
1916 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1917 async fn cache_miss_backfills_and_child_task_patch_share_publish_order() {
1918 for route in [CacheBackfillRoute::Load, CacheBackfillRoute::TryLoad] {
1919 assert_cache_backfill_serializes_with_task_patch(route).await;
1920 }
1921 }
1922
1923 #[tokio::test]
1924 async fn cache_hits_do_not_wait_for_the_persistence_lock() {
1925 let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
1926 let repo = test_repo(storage);
1927 let id = "cache-hit-lock-free";
1928 let cached = Session::new(id, "cached-model");
1929 cache_put(&repo, &cached);
1930 let persistence_guard = repo.persistence().acquire_lock(id).await;
1931
1932 let loaded = tokio::time::timeout(Duration::from_millis(100), repo.load(id)).await;
1933 let try_loaded = tokio::time::timeout(Duration::from_millis(100), repo.try_load(id)).await;
1934 drop(persistence_guard);
1935
1936 assert_eq!(
1937 loaded
1938 .expect("cache hit must not wait for the persistence lock")
1939 .expect("cached session")
1940 .model,
1941 "cached-model"
1942 );
1943 assert_eq!(
1944 try_loaded
1945 .expect("fallible cache hit must not wait for the persistence lock")
1946 .expect("cache read succeeds")
1947 .expect("cached session")
1948 .model,
1949 "cached-model"
1950 );
1951 }
1952
1953 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1954 async fn load_merged_storage_refresh_and_child_task_patch_share_publish_order() {
1955 let temp = tempfile::tempdir().unwrap();
1956 let concrete_storage = Arc::new(
1957 bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
1958 .await
1959 .expect("SessionStoreV2"),
1960 );
1961 let storage: Arc<dyn Storage> = concrete_storage;
1962 let id = "load-merged-cache-order";
1963 let (hook, loaded_old_durable, release) = durable_cache_fence("load_merged", id);
1964 let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
1965
1966 let mut durable = Session::new(id, "model");
1967 durable.updated_at = Utc::now();
1968 durable.set_task_list(task_list(id, "initial"));
1969 durable.set_task_list_version_meta("0");
1970 storage.save_session(&durable).await.unwrap();
1971
1972 let mut memory = durable.clone();
1973 memory.updated_at = durable.updated_at - chrono::Duration::seconds(1);
1974 memory.set_task_list(task_list(id, "memory"));
1975 cache_put(&repo, &memory);
1976
1977 let load_repo = repo.clone();
1978 let load = tokio::spawn(async move {
1979 assert!(
1980 load_repo.load_merged(id).await.is_some(),
1981 "seeded session must load"
1982 );
1983 Ok(())
1984 });
1985 loaded_old_durable
1986 .await
1987 .expect("old durable snapshot loaded");
1988
1989 let patch_repo = repo.clone();
1990 let child_task_list = task_list(id, "child");
1991 let patch = tokio::spawn(async move {
1992 bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
1993 patch_repo.as_ref(),
1994 id,
1995 &child_task_list,
1996 "1",
1997 )
1998 .await
1999 .map(|updated| assert!(updated, "root must exist"))
2000 });
2001 assert_second_write_waits_for_cache_publish(load, patch, release).await;
2002
2003 let durable = storage.load_session(id).await.unwrap().unwrap();
2004 let cached = read_cached_session(repo.cache(), id).expect("refreshed cache");
2005 for (tier, session) in [("durable", durable), ("cache", cached)] {
2006 assert_eq!(
2007 session.task_list_version_meta().as_deref(),
2008 Some("1"),
2009 "{tier} must retain the child transaction"
2010 );
2011 assert_eq!(
2012 session.task_list.as_ref().map(|list| list.title.as_str()),
2013 Some("child"),
2014 "{tier} must retain the child transaction"
2015 );
2016 }
2017 }
2018
2019 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
2020 async fn legacy_clear_refresh_and_child_task_patch_share_publish_order() {
2021 let temp = tempfile::tempdir().unwrap();
2022 let concrete_storage = Arc::new(
2023 bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
2024 .await
2025 .expect("SessionStoreV2"),
2026 );
2027 let storage: Arc<dyn Storage> = concrete_storage;
2028 let id = "legacy-clear-cache-order";
2029 let expected = vec![serde_json::json!({"content": "legacy"})];
2030 let (hook, loaded_post_cas_snapshot, release) = durable_cache_fence("clear_legacy", id);
2031 let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
2032
2033 let mut initial = Session::new(id, "model");
2034 initial.set_pending_injected_messages(expected.clone());
2035 initial.set_task_list(task_list(id, "initial"));
2036 initial.set_task_list_version_meta("0");
2037 storage.save_session(&initial).await.unwrap();
2038 cache_put(&repo, &initial);
2039
2040 let clear_repo = repo.clone();
2041 let clear_expected = expected.clone();
2042 let clear = tokio::spawn(async move {
2043 bamboo_domain::RuntimeSessionPersistence::clear_legacy_pending_messages(
2044 clear_repo.as_ref(),
2045 id,
2046 &clear_expected,
2047 )
2048 .await
2049 .map(|cleared| assert!(cleared, "legacy queue must match"))
2050 });
2051 loaded_post_cas_snapshot
2052 .await
2053 .expect("post-CAS snapshot loaded");
2054
2055 let patch_repo = repo.clone();
2056 let child_task_list = task_list(id, "child");
2057 let patch = tokio::spawn(async move {
2058 bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
2059 patch_repo.as_ref(),
2060 id,
2061 &child_task_list,
2062 "1",
2063 )
2064 .await
2065 .map(|updated| assert!(updated, "root must exist"))
2066 });
2067 assert_second_write_waits_for_cache_publish(clear, patch, release).await;
2068
2069 let durable = storage.load_session(id).await.unwrap().unwrap();
2070 let cached = read_cached_session(repo.cache(), id).expect("refreshed cache");
2071 for (tier, session) in [("durable", durable), ("cache", cached)] {
2072 assert_eq!(
2073 session.task_list_version_meta().as_deref(),
2074 Some("1"),
2075 "{tier} must retain the child transaction"
2076 );
2077 assert_eq!(
2078 session.task_list.as_ref().map(|list| list.title.as_str()),
2079 Some("child"),
2080 "{tier} must retain the child transaction"
2081 );
2082 assert!(
2083 !session.has_pending_injected_messages(),
2084 "{tier} must retain the successful legacy clear"
2085 );
2086 }
2087 }
2088
2089 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
2090 async fn concurrent_task_patches_publish_cache_in_durable_order() {
2091 let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
2092 let (hook, first_durable, release) = durable_cache_fence("task", "1");
2093 let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
2094 let id = "concurrent-task-cache-order";
2095 let mut initial = Session::new(id, "model");
2096 initial.set_task_list(task_list(id, "initial"));
2097 initial.set_task_list_version_meta("0");
2098 storage.save_session(&initial).await.unwrap();
2099 cache_put(&repo, &initial);
2100
2101 let first_repo = repo.clone();
2102 let first_task_list = task_list(id, "first");
2103 let first = tokio::spawn(async move {
2104 bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
2105 first_repo.as_ref(),
2106 id,
2107 &first_task_list,
2108 "1",
2109 )
2110 .await
2111 });
2112 first_durable.await.expect("first durable write reached");
2113
2114 let second_repo = repo.clone();
2115 let second_task_list = task_list(id, "second");
2116 let second = tokio::spawn(async move {
2117 bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
2118 second_repo.as_ref(),
2119 id,
2120 &second_task_list,
2121 "2",
2122 )
2123 .await
2124 });
2125 let (first_updated, second_updated) =
2126 assert_second_write_waits_for_cache_publish(first, second, release).await;
2127 assert!(first_updated && second_updated);
2128
2129 let durable = storage.load_session(id).await.unwrap().unwrap();
2130 let cached = read_cached_session(repo.cache(), id).expect("cached root");
2131 for (tier, session) in [("durable", durable), ("cache", cached)] {
2132 assert_eq!(
2133 session.task_list_version_meta().as_deref(),
2134 Some("2"),
2135 "{tier} must retain the second transaction"
2136 );
2137 assert_eq!(
2138 session.task_list.as_ref().map(|list| list.title.as_str()),
2139 Some("second"),
2140 "{tier} must retain the second transaction"
2141 );
2142 }
2143 }
2144
2145 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
2146 async fn root_control_plane_save_and_child_task_patch_share_publish_order() {
2147 let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
2148 let (hook, root_durable, release) = durable_cache_fence("save", "root-control-plane-order");
2149 let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
2150 let id = "root-control-plane-order";
2151
2152 let mut initial = Session::new(id, "model");
2153 initial.add_message(bamboo_agent_core::Message::user("durable transcript"));
2154 initial.set_task_list(task_list(id, "initial"));
2155 initial.set_task_list_version_meta("0");
2156 initial
2157 .metadata
2158 .insert("unrelated.runtime".to_string(), "keep".to_string());
2159 storage.save_session(&initial).await.unwrap();
2160 cache_put(&repo, &initial);
2161
2162 let root_repo = repo.clone();
2163 let mut root_snapshot = initial.clone();
2164 root_snapshot.set_task_list(task_list(id, "root"));
2165 root_snapshot.set_task_list_version_meta("1");
2166 let root_save = tokio::spawn(async move {
2167 bamboo_domain::RuntimeSessionPersistence::save_runtime_control_plane(
2168 root_repo.as_ref(),
2169 &mut root_snapshot,
2170 )
2171 .await
2172 });
2173 root_durable.await.expect("root durable write reached");
2174
2175 let child_repo = repo.clone();
2176 let child_task_list = task_list(id, "child");
2177 let child_patch = tokio::spawn(async move {
2178 bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
2179 child_repo.as_ref(),
2180 id,
2181 &child_task_list,
2182 "2",
2183 )
2184 .await
2185 .map(|updated| {
2186 assert!(updated, "root must exist");
2187 })
2188 });
2189 assert_second_write_waits_for_cache_publish(root_save, child_patch, release).await;
2190
2191 let durable = storage.load_session(id).await.unwrap().unwrap();
2192 let cached = read_cached_session(repo.cache(), id).expect("cached root");
2193 for (tier, session) in [("durable", durable), ("cache", cached)] {
2194 assert_eq!(
2195 session.task_list_version_meta().as_deref(),
2196 Some("2"),
2197 "{tier} must retain the child transaction"
2198 );
2199 assert_eq!(
2200 session.task_list.as_ref().map(|list| list.title.as_str()),
2201 Some("child"),
2202 "{tier} must retain the child transaction"
2203 );
2204 assert_eq!(
2205 session
2206 .metadata
2207 .get("unrelated.runtime")
2208 .map(String::as_str),
2209 Some("keep"),
2210 "{tier} must preserve unrelated runtime state"
2211 );
2212 assert_eq!(
2213 session.messages.len(),
2214 1,
2215 "{tier} must preserve the transcript"
2216 );
2217 }
2218 }
2219
2220 #[tokio::test]
2221 async fn stale_control_plane_save_keeps_checkpointed_ledger_in_durable_and_cache() {
2222 let temp = tempfile::tempdir().unwrap();
2223 let storage: Arc<dyn Storage> = Arc::new(
2224 bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
2225 .await
2226 .expect("SessionStoreV2"),
2227 );
2228 let repo = test_repo(storage.clone());
2229 let id = "control-plane-ledger-cache";
2230 let mut initial = Session::new(id, "model");
2231 initial.add_message(bamboo_agent_core::Message::user("durable transcript"));
2232 let assistant = bamboo_agent_core::Message::assistant("normalized", None);
2233 let anchor = assistant.id.clone();
2234 initial.add_message(assistant);
2235 let native_item = bamboo_domain::ProviderTranscriptItem::try_from_payload(
2236 bamboo_domain::ProviderFamily::OpenAi,
2237 bamboo_domain::ProviderProtocol::OpenAiResponsesV1,
2238 bamboo_domain::ProviderTranscriptOrigin::Provider,
2239 bamboo_domain::ProviderTranscriptAuthor::Model,
2240 serde_json::json!({
2241 "type":"tool_search_call","id":"tsc_cache_native","execution":"client","call_id":"cache_native",
2242 "status":"completed","arguments":{"query":"CACHE_NATIVE_PAYLOAD_SENTINEL"}
2243 }),
2244 )
2245 .unwrap();
2246 initial
2247 .append_provider_transcript_group(&anchor, None, vec![native_item])
2248 .unwrap();
2249 storage.save_session(&initial).await.unwrap();
2250 cache_put(&repo, &initial);
2251 let mut stale = initial.clone();
2252
2253 let mut runner = initial;
2254 runner.model_context_state = Some(bamboo_domain::ModelContextState {
2255 state_revision: 1,
2256 prefix_epoch: 1,
2257 cache_scope_sha256: Some("scope".to_string()),
2258 transcript_item_sha256: vec!["runner-l1".to_string()],
2259 ..bamboo_domain::ModelContextState::default()
2260 });
2261 bamboo_domain::RuntimeSessionPersistence::checkpoint_runtime_session(&repo, &mut runner)
2262 .await
2263 .unwrap();
2264
2265 stale
2266 .metadata
2267 .insert("runtime.suspend_reason".to_string(), "waiting".to_string());
2268 bamboo_domain::RuntimeSessionPersistence::save_runtime_control_plane(&repo, &mut stale)
2269 .await
2270 .unwrap();
2271
2272 let expected = runner.model_context_state;
2273 let expected_native = runner.provider_transcript;
2274 let durable = storage.load_session(id).await.unwrap().unwrap();
2275 let cached = read_cached_session(repo.cache(), id).expect("cached session");
2276 for (tier, session) in [("durable", durable), ("cache", cached.clone())] {
2277 assert_eq!(session.model_context_state, expected, "tier={tier}");
2278 assert_eq!(
2279 session.provider_transcript, expected_native,
2280 "tier={tier} must retain the message-anchored native transcript"
2281 );
2282 assert_eq!(
2283 session
2284 .metadata
2285 .get("runtime.suspend_reason")
2286 .map(String::as_str),
2287 Some("waiting"),
2288 "tier={tier}"
2289 );
2290 assert_eq!(session.messages.len(), 2, "tier={tier}");
2291 }
2292
2293 let runtime_json =
2294 std::fs::read_to_string(temp.path().join("sessions").join(id).join("runtime.json"))
2295 .unwrap();
2296 assert!(!runtime_json.contains("CACHE_NATIVE_PAYLOAD_SENTINEL"));
2297
2298 let mut cache_writer = cached;
2301 bamboo_domain::RuntimeSessionPersistence::checkpoint_runtime_session(
2302 &repo,
2303 &mut cache_writer,
2304 )
2305 .await
2306 .unwrap();
2307 let restarted = storage.load_session(id).await.unwrap().unwrap();
2308 assert_eq!(restarted.provider_transcript, expected_native);
2309 }
2310
2311 #[tokio::test]
2312 async fn narrow_runtime_metadata_transaction_preserves_live_and_durable_non_owned_state() {
2313 let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
2314 let repo = test_repo(storage.clone());
2315 let id = "narrow-metadata";
2316 let mut durable = Session::new(id, "durable-model");
2317 durable.add_message(bamboo_agent_core::Message::user("durable user turn"));
2318 durable
2319 .metadata
2320 .insert("external.durable".to_string(), "keep".to_string());
2321 storage.save_session(&durable).await.expect("seed durable");
2322
2323 let mut live = durable.clone();
2324 live.add_message(bamboo_agent_core::Message::assistant(
2325 "in-flight assistant tool call",
2326 None,
2327 ));
2328 live.model = "live-model".to_string();
2329 live.metadata
2330 .insert("external.live".to_string(), "keep".to_string());
2331 cache_put(&repo, &live);
2332
2333 repo.update_runtime_session(id, &["workflow.owned"], |latest| {
2334 latest
2335 .metadata
2336 .insert("workflow.owned".to_string(), "active".to_string());
2337 })
2338 .await
2339 .expect("transaction")
2340 .expect("session exists");
2341
2342 let saved = storage
2343 .load_session(id)
2344 .await
2345 .expect("load durable")
2346 .expect("durable exists");
2347 assert_eq!(
2348 saved.messages.len(),
2349 1,
2350 "transaction never writes stale live messages"
2351 );
2352 assert_eq!(
2353 saved.metadata.get("external.durable").map(String::as_str),
2354 Some("keep")
2355 );
2356 assert_eq!(
2357 saved.metadata.get("workflow.owned").map(String::as_str),
2358 Some("active")
2359 );
2360
2361 let cached = read_cached_session(repo.cache(), id).expect("live cache");
2362 assert_eq!(
2363 cached.messages.len(),
2364 2,
2365 "cache live tool call is not replaced"
2366 );
2367 assert_eq!(cached.model, "live-model");
2368 assert_eq!(
2369 cached.metadata.get("external.live").map(String::as_str),
2370 Some("keep")
2371 );
2372 assert_eq!(
2373 cached.metadata.get("workflow.owned").map(String::as_str),
2374 Some("active")
2375 );
2376 }
2377
2378 #[tokio::test]
2383 async fn load_merged_does_not_regress_to_older_storage() {
2384 let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
2385 let repo = test_repo(storage.clone());
2386 let id = "s1";
2387
2388 let mut stale = Session::new(id.to_string(), "m");
2389 stale.set_pending_question(
2390 "tc1".into(),
2391 "kind".into(),
2392 "q?".into(),
2393 vec!["OK".into()],
2394 true,
2395 );
2396 stale.updated_at = Utc::now() - chrono::Duration::seconds(10);
2397 storage.save_session(&stale).await.unwrap();
2398
2399 let mut fresh = stale.clone();
2402 fresh.clear_pending_question();
2403 fresh.updated_at = Utc::now();
2404 cache_put(&repo, &fresh);
2405
2406 let merged = repo.load_merged(id).await.expect("session exists");
2407 assert!(
2408 merged.pending_question.is_none(),
2409 "must return the newer answered memory copy, not the stale storage one"
2410 );
2411 let cached = read_cached_session(repo.cache(), id).expect("cached");
2412 assert!(
2413 cached.pending_question.is_none(),
2414 "load_merged must never regress the cache to a stale storage copy"
2415 );
2416 }
2417
2418 #[tokio::test]
2422 async fn load_merged_recovers_pending_question_from_same_age_storage() {
2423 let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
2424 let repo = test_repo(storage.clone());
2425 let id = "s2";
2426 let ts = Utc::now();
2427
2428 let mut with_pending = Session::new(id.to_string(), "m");
2429 with_pending.set_pending_question(
2430 "tc".into(),
2431 "k".into(),
2432 "q".into(),
2433 vec!["OK".into()],
2434 true,
2435 );
2436 with_pending.updated_at = ts;
2437 storage.save_session(&with_pending).await.unwrap();
2438
2439 let mut lost = with_pending.clone();
2440 lost.clear_pending_question();
2441 lost.updated_at = ts;
2442 cache_put(&repo, &lost);
2443
2444 let merged = repo.load_merged(id).await.expect("session exists");
2445 assert!(
2446 merged.pending_question.is_some(),
2447 "same-age storage carrying a pending question must still be recovered"
2448 );
2449 }
2450
2451 #[tokio::test]
2452 async fn runtime_publish_refreshes_cache_even_when_storage_fails() {
2453 let id = "runtime-selection";
2454 let mut previous = Session::new(id.to_string(), "m");
2455 previous.metadata.insert(
2456 "skill_runtime_selected_skill_ids".to_string(),
2457 "[\"plan\"]".to_string(),
2458 );
2459 let storage: Arc<dyn Storage> = Arc::new(FailingSaveStorage {
2460 persisted: Mutex::new(Some(previous.clone())),
2461 });
2462 let repo = test_repo(storage.clone());
2463 cache_put(&repo, &previous);
2464
2465 let mut current = previous.clone();
2466 current.metadata.insert(
2467 "skill_runtime_selected_skill_ids".to_string(),
2468 "[\"review\"]".to_string(),
2469 );
2470 current.updated_at = Utc::now();
2471
2472 let result =
2473 bamboo_domain::RuntimeSessionPersistence::save_runtime_session(&repo, &mut current)
2474 .await;
2475 assert!(result.is_err(), "durable failure must still be surfaced");
2476
2477 let cached = repo.load(id).await.expect("cached current session");
2478 assert_eq!(
2479 cached
2480 .metadata
2481 .get("skill_runtime_selected_skill_ids")
2482 .map(String::as_str),
2483 Some("[\"review\"]")
2484 );
2485 let allowlist = bamboo_skills::access_control::extract_skill_allowlist(&cached.metadata)
2486 .expect("runtime authorization allowlist");
2487 assert!(allowlist.contains("review"));
2488 assert!(!allowlist.contains("plan"));
2489 let durable = storage
2490 .load_session(id)
2491 .await
2492 .expect("load durable state")
2493 .expect("previous durable session");
2494 assert_eq!(
2495 durable
2496 .metadata
2497 .get("skill_runtime_selected_skill_ids")
2498 .map(String::as_str),
2499 Some("[\"plan\"]")
2500 );
2501 }
2502
2503 #[tokio::test]
2504 async fn finalized_runtime_save_reconciles_late_child_completion_in_disk_and_cache() {
2505 use bamboo_domain::session::runtime_state::{
2506 AgentRuntimeState, AgentStatusState, ChildWaitPolicy, WaitingForChildrenState,
2507 };
2508
2509 let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
2510 let repo = test_repo(storage.clone());
2511 let id = "late-child-finalization";
2512 let mut parent = Session::new(id, "model");
2513 let mut runtime = AgentRuntimeState::new("parent-run");
2514 runtime.status = AgentStatusState::Suspended;
2515 let mut wait = WaitingForChildrenState::for_children(
2516 vec!["child-1".to_string()],
2517 ChildWaitPolicy::All,
2518 Utc::now(),
2519 );
2520 wait.registered_by_tool_call_id = Some("tc_wait".to_string());
2521 runtime.waiting_for_children = Some(wait);
2522 parent.agent_runtime_state = Some(runtime);
2523 parent.metadata.insert(
2524 "runtime.suspend_reason".to_string(),
2525 "waiting_for_children".to_string(),
2526 );
2527 parent.set_last_run_status("suspended");
2528 storage.save_session(&parent).await.unwrap();
2529 cache_put(&repo, &parent);
2530
2531 let mut stale_runner = parent.clone();
2532 let mut completed = parent;
2533 let completed_runtime = completed.agent_runtime_state.as_mut().unwrap();
2534 completed_runtime.waiting_for_children = None;
2535 completed_runtime.status = AgentStatusState::Idle;
2536 completed_runtime.suspension = None;
2537 completed.metadata.remove("runtime.suspend_reason");
2538 let mut outcome = bamboo_agent_core::Message::user("completed child result");
2539 outcome.metadata = Some(serde_json::json!({
2540 "runtime_kind": "child_completion_resume"
2541 }));
2542 completed.add_message(outcome);
2543 storage.save_session(&completed).await.unwrap();
2544
2545 bamboo_domain::RuntimeSessionPersistence::save_finalized_runtime_session(
2546 &repo,
2547 &mut stale_runner,
2548 )
2549 .await
2550 .unwrap();
2551
2552 let durable = storage.load_session(id).await.unwrap().unwrap();
2553 let cached = read_cached_session(repo.cache(), id).expect("final snapshot published");
2554 for saved in [&durable, &cached] {
2555 assert!(saved
2556 .agent_runtime_state
2557 .as_ref()
2558 .unwrap()
2559 .waiting_for_children
2560 .is_none());
2561 assert!(!saved.metadata.contains_key("runtime.suspend_reason"));
2562 assert_eq!(saved.last_run_status().as_deref(), Some("completed"));
2563 assert!(saved
2564 .messages
2565 .iter()
2566 .any(|message| message.content == "completed child result"));
2567 }
2568 }
2569
2570 #[tokio::test]
2571 async fn inherent_save_leaves_existing_cache_untouched_when_storage_fails() {
2572 let id = "inherent-save-failure";
2573 let previous = Session::new(id, "previous");
2574 let storage: Arc<dyn Storage> = Arc::new(FailingSaveStorage {
2575 persisted: Mutex::new(Some(previous.clone())),
2576 });
2577 let repo = test_repo(storage);
2578 cache_put(&repo, &previous);
2579
2580 let mut current = previous.clone();
2581 current.model = "current".to_string();
2582 assert!(repo.save(&mut current).await.is_err());
2583 assert_eq!(
2584 read_cached_session(repo.cache(), id)
2585 .expect("existing cache")
2586 .model,
2587 "previous",
2588 "fallible inherent save must publish only after a durable commit"
2589 );
2590 }
2591
2592 #[tokio::test]
2593 async fn blocked_replay_ambiguous_save_evicts_same_age_stale_approval_cache() {
2594 let id = "blocked-replay-ambiguous-save";
2595 let mut approved = Session::new(id, "model");
2596 approved.metadata.insert(
2597 crate::session_app::respond::PERMISSION_REEXECUTE_METADATA_KEY.into(),
2598 "call-1".into(),
2599 );
2600 let storage: Arc<dyn Storage> = Arc::new(AmbiguousReplaySaveStorage {
2601 persisted: Mutex::new(Some(approved.clone())),
2602 });
2603 let repo = test_repo(storage);
2604 cache_put(&repo, &approved);
2605
2606 let mut resolved = approved.clone();
2607 resolved
2608 .metadata
2609 .remove(crate::session_app::respond::PERMISSION_REEXECUTE_METADATA_KEY);
2610 assert!(repo.save_replay_resolution(&mut resolved).await.is_err());
2611 assert!(read_cached_session(repo.cache(), id).is_none());
2612 let reloaded = repo.load_merged_checked(id).await.unwrap().unwrap();
2613 assert_eq!(reloaded.updated_at, approved.updated_at);
2614 assert!(!reloaded
2615 .metadata
2616 .contains_key(crate::session_app::respond::PERMISSION_REEXECUTE_METADATA_KEY));
2617 }
2618
2619 #[tokio::test]
2620 async fn blocked_replay_precommit_save_error_reloads_retryable_approval() {
2621 let id = "blocked-replay-precommit-save";
2622 let mut approved = Session::new(id, "model");
2623 approved.metadata.insert(
2624 crate::session_app::respond::PERMISSION_REEXECUTE_METADATA_KEY.into(),
2625 "call-1".into(),
2626 );
2627 let storage: Arc<dyn Storage> = Arc::new(FailingSaveStorage {
2628 persisted: Mutex::new(Some(approved.clone())),
2629 });
2630 let repo = test_repo(storage);
2631 cache_put(&repo, &approved);
2632
2633 let mut resolved = approved.clone();
2634 resolved
2635 .metadata
2636 .remove(crate::session_app::respond::PERMISSION_REEXECUTE_METADATA_KEY);
2637 assert!(repo.save_replay_resolution(&mut resolved).await.is_err());
2638 assert!(read_cached_session(repo.cache(), id).is_none());
2639 let reloaded = repo.load_merged_checked(id).await.unwrap().unwrap();
2640 assert_eq!(reloaded.updated_at, approved.updated_at);
2641 assert_eq!(
2642 reloaded
2643 .metadata
2644 .get(crate::session_app::respond::PERMISSION_REEXECUTE_METADATA_KEY)
2645 .map(String::as_str),
2646 Some("call-1")
2647 );
2648 }
2649
2650 #[tokio::test]
2651 async fn save_and_cache_still_refreshes_cache_when_storage_fails() {
2652 let id = "save-and-cache-failure";
2653 let previous = Session::new(id, "previous");
2654 let storage: Arc<dyn Storage> = Arc::new(FailingSaveStorage {
2655 persisted: Mutex::new(Some(previous.clone())),
2656 });
2657 let repo = test_repo(storage);
2658 cache_put(&repo, &previous);
2659
2660 let mut current = previous;
2661 current.model = "current".to_string();
2662 repo.save_and_cache(&mut current).await;
2663 assert_eq!(
2664 read_cached_session(repo.cache(), id)
2665 .expect("refreshed cache")
2666 .model,
2667 "current",
2668 "fire-and-forget save must retain its existing cache-on-failure behavior"
2669 );
2670 }
2671
2672 #[tokio::test]
2673 async fn checkpoint_leaves_existing_cache_untouched_when_storage_fails() {
2674 let id = "checkpoint-failure";
2675 let previous = Session::new(id, "previous");
2676 let storage: Arc<dyn Storage> = Arc::new(FailingSaveStorage {
2677 persisted: Mutex::new(Some(previous.clone())),
2678 });
2679 let repo = test_repo(storage);
2680 cache_put(&repo, &previous);
2681
2682 let mut current = previous.clone();
2683 current.model = "current".to_string();
2684 let result = bamboo_domain::RuntimeSessionPersistence::checkpoint_runtime_session(
2685 &repo,
2686 &mut current,
2687 )
2688 .await;
2689 assert!(result.is_err());
2690 assert_eq!(
2691 read_cached_session(repo.cache(), id)
2692 .expect("existing cache")
2693 .model,
2694 "previous",
2695 "checkpoint must publish only after a durable commit"
2696 );
2697 }
2698
2699 #[tokio::test]
2700 async fn retrieval_window_checkpoint_publishes_the_exact_archived_snapshot() {
2701 let temp = tempfile::tempdir().unwrap();
2702 let storage: Arc<dyn Storage> = Arc::new(
2703 bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
2704 .await
2705 .expect("SessionStoreV2"),
2706 );
2707 let repo = test_repo(storage.clone());
2708 let id = "retrieval-window-repository-checkpoint";
2709 let mut expected = Session::new(id, "model");
2710 expected.add_message(bamboo_agent_core::Message::user("archive candidate"));
2711 expected.add_message(bamboo_agent_core::Message::assistant("retain", None));
2712 storage.save_session(&expected).await.unwrap();
2713 cache_put(&repo, &expected);
2714 let mut staged = staged_retrieval_window_archive(&expected);
2715
2716 let outcome = bamboo_domain::RuntimeSessionPersistence::checkpoint_retrieval_window(
2717 &repo,
2718 &expected,
2719 &mut staged,
2720 )
2721 .await
2722 .unwrap();
2723
2724 assert_eq!(outcome, RetrievalWindowCheckpointOutcome::Committed);
2725 let durable = storage.load_session(id).await.unwrap().unwrap();
2726 let cached = read_cached_session(repo.cache(), id).expect("cached archived Session");
2727 for (tier, saved) in [("durable", durable), ("cache", cached)] {
2728 assert!(saved.messages[0].compressed, "tier={tier}");
2729 assert_eq!(
2730 saved.messages[0].compressed_by_event_id, staged.messages[0].compressed_by_event_id,
2731 "tier={tier}"
2732 );
2733 assert_eq!(
2734 serde_json::to_value(&saved).unwrap(),
2735 serde_json::to_value(&staged).unwrap(),
2736 "tier={tier}"
2737 );
2738 }
2739 }
2740
2741 #[tokio::test]
2742 async fn legacy_clear_uses_durable_cas_and_never_erases_concurrent_append_from_stale_cache() {
2743 let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
2744 let repo = test_repo(storage.clone());
2745 let id = "legacy-cas-race";
2746 let expected = vec![serde_json::json!({"content": "first"})];
2747
2748 let mut stale_cache = Session::new(id, "m");
2749 stale_cache.set_pending_injected_messages(expected.clone());
2750 cache_put(&repo, &stale_cache);
2751
2752 let mut durable = stale_cache.clone();
2753 durable.set_pending_injected_messages(vec![
2754 serde_json::json!({"content": "first"}),
2755 serde_json::json!({"content": "concurrent"}),
2756 ]);
2757 storage.save_session(&durable).await.unwrap();
2758
2759 let cleared = bamboo_domain::RuntimeSessionPersistence::clear_legacy_pending_messages(
2760 &repo, id, &expected,
2761 )
2762 .await
2763 .unwrap();
2764 assert!(!cleared, "the durable compare-and-clear must reject drift");
2765 assert_eq!(
2766 storage
2767 .load_session(id)
2768 .await
2769 .unwrap()
2770 .unwrap()
2771 .pending_injected_messages()
2772 .unwrap(),
2773 durable.pending_injected_messages().unwrap(),
2774 "the concurrent durable append must remain intact"
2775 );
2776 assert_eq!(
2777 read_cached_session(repo.cache(), id)
2778 .unwrap()
2779 .pending_injected_messages()
2780 .unwrap(),
2781 expected,
2782 "a failed CAS must not mutate the existing cache"
2783 );
2784 }
2785
2786 #[tokio::test]
2787 async fn successful_legacy_clear_refreshes_stale_cache_from_durable_state() {
2788 let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
2789 let repo = test_repo(storage.clone());
2790 let id = "legacy-cas-success";
2791 let expected = vec![serde_json::json!({"content": "first"})];
2792
2793 let mut stale_cache = Session::new(id, "stale-model");
2794 stale_cache.set_pending_injected_messages(expected.clone());
2795 cache_put(&repo, &stale_cache);
2796
2797 let mut durable = Session::new(id, "durable-model");
2798 durable.set_pending_injected_messages(expected.clone());
2799 durable
2800 .metadata
2801 .insert("durable-only".to_string(), "keep".to_string());
2802 storage.save_session(&durable).await.unwrap();
2803
2804 assert!(
2805 bamboo_domain::RuntimeSessionPersistence::clear_legacy_pending_messages(
2806 &repo, id, &expected,
2807 )
2808 .await
2809 .unwrap()
2810 );
2811 let cached = read_cached_session(repo.cache(), id).unwrap();
2812 assert!(!cached.has_pending_injected_messages());
2813 assert_eq!(cached.model, "durable-model");
2814 assert_eq!(
2815 cached.metadata.get("durable-only").map(String::as_str),
2816 Some("keep")
2817 );
2818 }
2819
2820 #[tokio::test]
2821 async fn repository_final_save_publishes_host_checkpointed_parent_question() {
2822 use bamboo_domain::{
2823 FunctionCall, ParentQuestion, ParentQuestionCheckpointV1, PendingQuestion,
2824 PendingQuestionSource, ToolCall,
2825 };
2826
2827 let temp = tempfile::tempdir().unwrap();
2828 let store = Arc::new(
2829 bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
2830 .await
2831 .unwrap(),
2832 );
2833 let storage: Arc<dyn Storage> = store.clone();
2834 let repo = test_repo(storage);
2835 let parent = Session::new("repository-question-parent", "model");
2836 store.save_session(&parent).await.unwrap();
2837 let mut child =
2838 Session::new_child_of("repository-question-child", &parent, "model", "assignment");
2839 child.add_message(bamboo_agent_core::Message::user("assignment"));
2840 child.set_last_run_status("running");
2841 store.save_session(&child).await.unwrap();
2842 cache_put(&repo, &child);
2843
2844 let assistant = bamboo_agent_core::Message::assistant(
2845 "",
2846 Some(vec![ToolCall {
2847 id: "question-call".into(),
2848 tool_type: "function".into(),
2849 function: FunctionCall {
2850 name: "AgenticQuestion".into(),
2851 arguments: "{}".into(),
2852 },
2853 }]),
2854 );
2855 let result = bamboo_agent_core::Message::tool_result_with_status(
2856 "question-call",
2857 "Clarification needed: Which option?",
2858 true,
2859 );
2860 let pending = PendingQuestion {
2861 tool_call_id: "question-call".into(),
2862 tool_name: "AgenticQuestion".into(),
2863 question: "Which option?".into(),
2864 options: vec!["A".into(), "B".into()],
2865 allow_custom: true,
2866 source: PendingQuestionSource::AgenticClarification,
2867 };
2868 let observation = ParentQuestionCheckpointV1 {
2869 version: 1,
2870 prefix_message_count: child.messages.len(),
2871 prefix_digest: ParentQuestion::prefix_digest(&child.messages).unwrap(),
2872 suffix: vec![assistant, result.clone()],
2873 tool_call_id: pending.tool_call_id.clone(),
2874 tool_result_message_id: result.id.clone(),
2875 question_digest: ParentQuestion::question_digest(&pending, &result.id),
2876 pending,
2877 };
2878 let (_, question) = store
2879 .checkpoint_parent_question(&child, &parent, &observation)
2880 .await
2881 .unwrap()
2882 .unwrap();
2883
2884 let mut old_run = child.clone();
2885 old_run.metadata.insert(
2886 "runtime.actor_parent_question_handoff".into(),
2887 "true".into(),
2888 );
2889 old_run.set_last_run_status("suspended");
2890 old_run.add_message(bamboo_agent_core::Message::assistant(
2891 "old terminal text",
2892 None,
2893 ));
2894 cache_put(&repo, &old_run);
2895 bamboo_domain::RuntimeSessionPersistence::save_runtime_session(&repo, &mut old_run)
2896 .await
2897 .unwrap();
2898
2899 let cached = repo.load(&child.id).await.unwrap();
2900 let reopened = bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
2901 .await
2902 .unwrap()
2903 .load_session(&child.id)
2904 .await
2905 .unwrap()
2906 .unwrap();
2907 for saved in [&cached, &reopened] {
2908 assert_eq!(saved.last_run_status().as_deref(), Some("suspended"));
2909 assert_eq!(saved.messages.len(), child.messages.len() + 2);
2910 assert!(!saved
2911 .messages
2912 .iter()
2913 .any(|message| message.content == "old terminal text"));
2914 assert_eq!(
2915 saved.pending_question.as_ref().unwrap().source,
2916 PendingQuestionSource::DirectParent
2917 );
2918 assert_eq!(
2919 ParentQuestion::for_pending(&parent, saved).unwrap().id,
2920 question.id
2921 );
2922 }
2923 }
2924}
2925
2926fn validate_bound_root_snapshot(
2927 owner: &bamboo_domain::RootActorRuntimeWrite,
2928 snapshot: &Session,
2929) -> std::io::Result<()> {
2930 if snapshot.id != owner.fence.actor_id
2931 || snapshot.created_at != owner.created_at
2932 || snapshot.kind != bamboo_domain::SessionKind::Root
2933 || !snapshot.authority_identity.is_ordinary()
2934 {
2935 return Err(root_actor_binding_error("bound Root identity changed"));
2936 }
2937 owner
2938 .conversation_generation
2939 .check_canonical(snapshot)
2940 .map_err(root_actor_binding_error)
2941}
2942
2943fn root_actor_binding_error(reason: impl std::fmt::Display) -> std::io::Error {
2944 std::io::Error::new(
2945 std::io::ErrorKind::WouldBlock,
2946 bamboo_domain::SessionAuthorityConflict(format!("Root Actor execution rejected: {reason}")),
2947 )
2948}