1use std::sync::Arc;
16
17use bamboo_agent_core::storage::Storage;
18use bamboo_agent_core::Session;
19use bamboo_storage::LockedSessionStore;
20
21use crate::{read_cached_session, SessionCache};
22
23#[cfg(test)]
24type PostDurableHook = Arc<dyn Fn(&str, &str) + Send + Sync>;
25
26#[derive(Clone)]
29pub struct SessionRepository {
30 cache: SessionCache,
31 storage: Arc<dyn Storage>,
32 persistence: Arc<LockedSessionStore>,
33 #[cfg(test)]
34 post_durable_hook: Option<PostDurableHook>,
35}
36
37impl SessionRepository {
38 pub fn new(
39 cache: SessionCache,
40 storage: Arc<dyn Storage>,
41 persistence: Arc<LockedSessionStore>,
42 ) -> Self {
43 Self {
44 cache,
45 storage,
46 persistence,
47 #[cfg(test)]
48 post_durable_hook: None,
49 }
50 }
51
52 #[cfg(test)]
53 fn with_post_durable_hook(mut self, hook: PostDurableHook) -> Self {
54 self.post_durable_hook = Some(hook);
55 self
56 }
57
58 #[cfg(test)]
59 fn run_post_durable_hook(&self, operation: &str, marker: &str) {
60 if let Some(hook) = self.post_durable_hook.as_ref() {
61 hook(operation, marker);
62 }
63 }
64
65 pub fn cache(&self) -> &SessionCache {
66 &self.cache
67 }
68
69 pub fn storage(&self) -> &Arc<dyn Storage> {
70 &self.storage
71 }
72
73 pub fn persistence(&self) -> &Arc<LockedSessionStore> {
74 &self.persistence
75 }
76
77 pub async fn load(&self, session_id: &str) -> Option<Session> {
80 if let Some(session) = read_cached_session(&self.cache, session_id) {
81 return Some(session);
82 }
83
84 let _guard = self.persistence.acquire_lock(session_id).await;
85 if let Some(session) = read_cached_session(&self.cache, session_id) {
86 return Some(session);
87 }
88
89 let loaded = self.storage.load_session(session_id).await;
90 #[cfg(test)]
91 self.run_post_durable_hook("load", session_id);
92 match loaded {
93 Ok(Some(session)) => {
94 self.cache.insert(
95 session_id.to_string(),
96 Arc::new(parking_lot::RwLock::new(session.clone())),
97 );
98 Some(session)
99 }
100 _ => None,
101 }
102 }
103
104 pub async fn try_load(&self, session_id: &str) -> std::io::Result<Option<Session>> {
108 if let Some(session) = read_cached_session(&self.cache, session_id) {
109 return Ok(Some(session));
110 }
111
112 let _guard = self.persistence.acquire_lock(session_id).await;
113 if let Some(session) = read_cached_session(&self.cache, session_id) {
114 return Ok(Some(session));
115 }
116
117 let loaded = self.storage.load_session(session_id).await?;
118 #[cfg(test)]
119 self.run_post_durable_hook("try_load", session_id);
120 if let Some(ref session) = loaded {
121 self.cache.insert(
122 session_id.to_string(),
123 Arc::new(parking_lot::RwLock::new(session.clone())),
124 );
125 }
126 Ok(loaded)
127 }
128
129 pub async fn save(&self, session: &mut Session) -> std::io::Result<()> {
133 self.persistence
134 .merge_save_runtime_and_publish(session, |saved, committed| {
135 if committed {
136 #[cfg(test)]
137 self.run_post_durable_hook("save_full", &saved.id);
138 self.cache.insert(
139 saved.id.clone(),
140 Arc::new(parking_lot::RwLock::new(saved.clone())),
141 );
142 }
143 })
144 .await
145 }
146
147 pub async fn update_runtime_session<F>(
151 &self,
152 session_id: &str,
153 metadata_keys: &[&str],
154 mutate: F,
155 ) -> std::io::Result<Option<Session>>
156 where
157 F: FnOnce(&mut Session),
158 {
159 self.persistence
160 .update_runtime_config_and_publish(session_id, mutate, |saved| {
161 if let Some(cached) = self.cache.get(session_id) {
162 let mut cached = cached.write();
163 for key in metadata_keys {
164 if let Some(value) = saved.metadata.get(*key) {
165 cached.metadata.insert((*key).to_string(), value.clone());
166 } else {
167 cached.metadata.remove(*key);
168 }
169 }
170 }
171 })
172 .await
173 }
174
175 pub async fn load_or_create(&self, session_id: &str, model: &str) -> Session {
177 if let Some(session) = self.load(session_id).await {
178 return session;
179 }
180 Session::new(session_id.to_string(), model.to_string())
181 }
182
183 pub async fn load_merged_checked(&self, session_id: &str) -> std::io::Result<Option<Session>> {
192 let _guard = self.persistence.acquire_lock(session_id).await;
193 let memory_session = read_cached_session(&self.cache, session_id);
194 let storage_session = self.storage.load_session(session_id).await?;
195 #[cfg(test)]
196 self.run_post_durable_hook("load_merged", session_id);
197
198 Ok(match (memory_session, storage_session) {
199 (Some(memory), Some(storage)) => {
200 let prefer_storage = should_prefer_storage(&memory, &storage);
201 let diverged = prefer_storage || memory.messages.len() != storage.messages.len();
202 let chosen_len = if prefer_storage {
203 storage.messages.len()
204 } else {
205 memory.messages.len()
206 };
207 macro_rules! merged_log {
208 ($level:ident) => {
209 tracing::$level!(
210 "[{}] load_session_merged: memory={} msgs (updated_at={}), storage={} msgs (updated_at={}), prefer_storage={} -> chose {} msgs",
211 session_id,
212 memory.messages.len(),
213 memory.updated_at,
214 storage.messages.len(),
215 storage.updated_at,
216 prefer_storage,
217 chosen_len,
218 )
219 };
220 }
221 if diverged {
222 merged_log!(debug);
223 } else {
224 merged_log!(trace);
225 }
226 let memory_updated_at = memory.updated_at;
227 let chosen = if prefer_storage { storage } else { memory };
228 if prefer_storage && chosen.updated_at >= memory_updated_at {
236 self.cache.insert(
237 session_id.to_string(),
238 Arc::new(parking_lot::RwLock::new(chosen.clone())),
239 );
240 }
241 Some(chosen)
242 }
243 (Some(memory), None) => Some(memory),
244 (None, Some(storage)) => {
245 self.cache.insert(
246 session_id.to_string(),
247 Arc::new(parking_lot::RwLock::new(storage.clone())),
248 );
249 Some(storage)
250 }
251 (None, None) => None,
252 })
253 }
254
255 pub async fn load_merged(&self, session_id: &str) -> Option<Session> {
259 match self.load_merged_checked(session_id).await {
260 Ok(session) => session,
261 Err(error) => {
262 tracing::warn!(
263 "[{}] Failed to load merged session from storage: {}",
264 session_id,
265 error
266 );
267 read_cached_session(&self.cache, session_id)
268 }
269 }
270 }
271
272 pub async fn save_and_cache(&self, session: &mut Session) {
275 let result = self
276 .persistence
277 .merge_save_runtime_and_publish(session, |saved, _| {
278 #[cfg(test)]
279 self.run_post_durable_hook("save_and_cache", &saved.id);
280 self.cache.insert(
281 saved.id.clone(),
282 Arc::new(parking_lot::RwLock::new(saved.clone())),
283 );
284 })
285 .await;
286 if let Err(error) = result {
287 tracing::warn!("[{}] Failed to save session: {}", session.id, error);
288 }
289 }
290
291 async fn refresh_cached_task_control_plane(&self, session_id: &str) {
292 match self.storage.load_runtime_control_plane(session_id).await {
293 Ok(Some(durable)) => {
294 if let Some(cached) = self.cache.get(session_id) {
295 adopt_task_control_plane(&mut cached.write(), &durable);
296 }
297 }
298 Ok(None) => {}
299 Err(error) => tracing::warn!(
300 "[{}] Failed to refresh Task control-plane after write conflict: {}",
301 session_id,
302 error
303 ),
304 }
305 }
306}
307
308fn adopt_task_control_plane(target: &mut Session, durable: &Session) {
309 target.task_list = durable.task_list.clone();
310 target
311 .metadata
312 .remove(bamboo_domain::session::runtime_metadata::keys::TASK_LIST_VERSION);
313 if let Some(runtime_metadata) = target.runtime_metadata.as_mut() {
314 runtime_metadata.task_list_version = None;
315 }
316 if target
317 .runtime_metadata
318 .as_ref()
319 .is_some_and(bamboo_domain::session::SessionRuntimeMetadata::is_empty)
320 {
321 target.runtime_metadata = None;
322 }
323 if let Some(version) = durable.task_list_version_meta() {
324 target.set_task_list_version_meta(version);
325 }
326}
327
328fn should_prefer_storage(memory_session: &Session, storage_session: &Session) -> bool {
329 if storage_session.updated_at < memory_session.updated_at {
334 return false;
335 }
336 storage_session.updated_at > memory_session.updated_at
340 || (memory_session.pending_question.is_none() && storage_session.pending_question.is_some())
341}
342
343#[async_trait::async_trait]
347impl bamboo_domain::RuntimeSessionPersistence for SessionRepository {
348 async fn save_runtime_session(&self, session: &mut Session) -> std::io::Result<()> {
349 self.persistence
354 .merge_save_runtime_and_publish(session, |saved, _| {
355 #[cfg(test)]
356 self.run_post_durable_hook("save_runtime_session", &saved.id);
357 self.cache.insert(
358 saved.id.clone(),
359 Arc::new(parking_lot::RwLock::new(saved.clone())),
360 );
361 })
362 .await
363 }
364
365 async fn seed_runtime_activation(&self, session: &mut Session) -> std::io::Result<()> {
366 self.persistence
367 .seed_runtime_activation_and_publish(session, |saved, committed| {
368 #[cfg(test)]
369 self.run_post_durable_hook("seed_runtime_activation", &saved.id);
370 if committed {
371 self.cache.insert(
372 saved.id.clone(),
373 Arc::new(parking_lot::RwLock::new(saved.clone())),
374 );
375 }
376 })
377 .await
378 }
379
380 async fn record_permission_posture_activation(
381 &self,
382 session_id: &str,
383 expected_audit_revision: Option<u64>,
384 seed: &bamboo_domain::PermissionAuditSeed,
385 ) -> std::io::Result<Option<Session>> {
386 self.persistence
387 .record_permission_posture_activation_and_publish(
388 session_id,
389 expected_audit_revision,
390 seed,
391 |saved| {
392 #[cfg(test)]
393 self.run_post_durable_hook("permission_posture_activation", &saved.id);
394 self.cache.insert(
395 saved.id.clone(),
396 Arc::new(parking_lot::RwLock::new(saved.clone())),
397 );
398 },
399 )
400 .await
401 }
402
403 async fn save_runtime_control_plane(&self, session: &mut Session) -> std::io::Result<()> {
404 self.persistence
405 .save_runtime_only_and_publish(session, |saved| {
406 #[cfg(test)]
407 self.run_post_durable_hook("save", &saved.id);
408
409 if let Some(cached) = self.cache.get(&saved.id) {
417 let mut cached = cached.write();
418 let messages = cached.messages.clone();
419 let admission = cached
420 .runtime_metadata
421 .as_ref()
422 .and_then(|metadata| metadata.session_inbox_admission.clone());
423 let mut refreshed = saved.clone();
424 refreshed.messages = messages;
425 if let Some(admission) = admission {
426 refreshed
427 .runtime_metadata
428 .get_or_insert_with(Default::default)
429 .session_inbox_admission = Some(admission);
430 } else if let Some(metadata) = refreshed.runtime_metadata.as_mut() {
431 metadata.session_inbox_admission = None;
432 }
433 *cached = refreshed;
434 }
435 })
436 .await
437 }
438
439 async fn load_runtime_control_plane(
440 &self,
441 session_id: &str,
442 ) -> std::io::Result<Option<Session>> {
443 bamboo_domain::RuntimeSessionPersistence::load_runtime_control_plane(
444 self.persistence.as_ref(),
445 session_id,
446 )
447 .await
448 }
449
450 async fn update_task_list_control_plane(
451 &self,
452 session_id: &str,
453 task_list: &bamboo_domain::TaskList,
454 version: &str,
455 ) -> std::io::Result<bool> {
456 let result = self
457 .persistence
458 .update_task_list_control_plane_and_publish(session_id, task_list, version, |_| {
459 #[cfg(test)]
460 self.run_post_durable_hook("task", version);
461
462 if let Some(cached) = self.cache.get(session_id) {
468 let mut cached = cached.write();
469 cached.set_task_list(task_list.clone());
470 cached.set_task_list_version_meta(version.to_string());
471 }
472 })
473 .await;
474 if result
475 .as_ref()
476 .is_err_and(|error| error.kind() == std::io::ErrorKind::WouldBlock)
477 {
478 self.refresh_cached_task_control_plane(session_id).await;
483 }
484 result
485 }
486
487 async fn update_task_list_control_plane_if_version(
488 &self,
489 session_id: &str,
490 expected_version: &str,
491 expected_task_list: &bamboo_domain::TaskList,
492 task_list: &bamboo_domain::TaskList,
493 version: &str,
494 ) -> std::io::Result<bool> {
495 self.persistence
496 .update_task_list_control_plane_if_version_and_publish(
497 session_id,
498 expected_version,
499 expected_task_list,
500 task_list,
501 version,
502 |_| {
503 if let Some(cached) = self.cache.get(session_id) {
504 let mut cached = cached.write();
505 cached.set_task_list(task_list.clone());
506 cached.set_task_list_version_meta(version.to_string());
507 }
508 },
509 )
510 .await
511 }
512
513 async fn update_task_list_control_planes_if_version(
514 &self,
515 session_id: &str,
516 shared_session_id: &str,
517 expected_version: &str,
518 expected_task_list: &bamboo_domain::TaskList,
519 task_list: &bamboo_domain::TaskList,
520 version: &str,
521 ) -> std::io::Result<bool> {
522 self.persistence
523 .update_task_list_control_planes_if_version_and_publish(
524 session_id,
525 shared_session_id,
526 expected_version,
527 expected_task_list,
528 task_list,
529 version,
530 |_, _| {
531 for id in [session_id, shared_session_id] {
532 if let Some(cached) = self.cache.get(id) {
533 let mut cached = cached.write();
534 cached.set_task_list(task_list.clone());
535 cached.set_task_list_version_meta(version.to_string());
536 }
537 }
538 },
539 )
540 .await
541 }
542
543 async fn checkpoint_runtime_session(&self, session: &mut Session) -> std::io::Result<()> {
544 self.persistence
551 .checkpoint_runtime_session_and_publish(session, |saved, committed| {
552 #[cfg(test)]
553 self.run_post_durable_hook("checkpoint", &saved.id);
554 if committed {
555 self.cache.insert(
556 saved.id.clone(),
557 Arc::new(parking_lot::RwLock::new(saved.clone())),
558 );
559 }
560 })
561 .await
562 }
563
564 async fn load_runtime_session(&self, session_id: &str) -> std::io::Result<Option<Session>> {
565 self.try_load(session_id).await
566 }
567
568 async fn clear_legacy_pending_messages(
569 &self,
570 session_id: &str,
571 expected: &[serde_json::Value],
572 ) -> std::io::Result<bool> {
573 self.persistence
578 .clear_legacy_pending_messages_and_publish(session_id, expected, |latest| {
579 #[cfg(test)]
580 self.run_post_durable_hook("clear_legacy", session_id);
581 self.cache.insert(
582 session_id.to_string(),
583 Arc::new(parking_lot::RwLock::new(latest.clone())),
584 );
585 })
586 .await
587 }
588
589 async fn append_token_usage_record(
590 &self,
591 session_id: &str,
592 json_line: &str,
593 ) -> std::io::Result<()> {
594 self.storage
595 .append_token_usage_record(session_id, json_line)
596 .await
597 }
598}
599
600#[cfg(test)]
601mod tests {
602 use super::*;
603 use bamboo_agent_core::storage::Storage;
604 use chrono::Utc;
605 use std::collections::HashMap;
606 use std::sync::atomic::{AtomicBool, Ordering};
607 use std::sync::{Condvar, Mutex};
608 use std::time::Duration;
609
610 #[derive(Default)]
611 struct MapStorage {
612 sessions: Mutex<HashMap<String, Session>>,
613 fail_pair_commit: AtomicBool,
614 }
615
616 struct FailingSaveStorage {
617 persisted: Mutex<Option<Session>>,
618 }
619
620 #[async_trait::async_trait]
621 impl Storage for MapStorage {
622 async fn save_session(&self, session: &Session) -> std::io::Result<()> {
623 self.sessions
624 .lock()
625 .unwrap()
626 .insert(session.id.clone(), session.clone());
627 Ok(())
628 }
629 async fn load_session(&self, session_id: &str) -> std::io::Result<Option<Session>> {
630 Ok(self.sessions.lock().unwrap().get(session_id).cloned())
631 }
632 async fn delete_session(&self, session_id: &str) -> std::io::Result<bool> {
633 Ok(self.sessions.lock().unwrap().remove(session_id).is_some())
634 }
635
636 async fn save_task_control_plane_if_matches(
637 &self,
638 original: &Session,
639 updated: &Session,
640 ) -> std::io::Result<bool> {
641 let mut sessions = self.sessions.lock().unwrap();
642 let matches_original = sessions.get(&original.id).is_some_and(|current| {
643 current.task_list_version_meta() == original.task_list_version_meta()
644 && serde_json::to_value(¤t.task_list).ok()
645 == serde_json::to_value(&original.task_list).ok()
646 });
647 if !matches_original {
648 return Ok(false);
649 }
650 let Some(current) = sessions.get(&original.id).cloned() else {
651 return Ok(false);
652 };
653 let mut committed = current;
654 committed.task_list = updated.task_list.clone();
655 if let Some(version) = updated.task_list_version_meta() {
656 committed.set_task_list_version_meta(version);
657 }
658 sessions.insert(committed.id.clone(), committed);
659 Ok(true)
660 }
661
662 async fn save_task_control_planes_atomically(
663 &self,
664 first_original: &Session,
665 first_updated: &Session,
666 second_original: &Session,
667 second_updated: &Session,
668 ) -> std::io::Result<bool> {
669 if self.fail_pair_commit.swap(false, Ordering::SeqCst) {
670 return Err(std::io::Error::other(
671 "injected paired Task transaction failure",
672 ));
673 }
674 let mut sessions = self.sessions.lock().unwrap();
675 let matches_original = |current: Option<&Session>, original: &Session| {
676 current.is_some_and(|current| {
677 current.task_list_version_meta() == original.task_list_version_meta()
678 && serde_json::to_value(¤t.task_list).ok()
679 == serde_json::to_value(&original.task_list).ok()
680 })
681 };
682 if !matches_original(sessions.get(&first_original.id), first_original)
683 || !matches_original(sessions.get(&second_original.id), second_original)
684 {
685 return Ok(false);
686 }
687 sessions.insert(first_updated.id.clone(), first_updated.clone());
688 sessions.insert(second_updated.id.clone(), second_updated.clone());
689 Ok(true)
690 }
691 }
692
693 #[async_trait::async_trait]
694 impl Storage for FailingSaveStorage {
695 async fn save_session(&self, _session: &Session) -> std::io::Result<()> {
696 Err(std::io::Error::other("injected save failure"))
697 }
698
699 async fn load_session(&self, _session_id: &str) -> std::io::Result<Option<Session>> {
700 Ok(self.persisted.lock().unwrap().clone())
701 }
702
703 async fn delete_session(&self, _session_id: &str) -> std::io::Result<bool> {
704 Ok(false)
705 }
706 }
707
708 fn test_repo(storage: Arc<dyn Storage>) -> SessionRepository {
709 let cache: SessionCache = Arc::new(dashmap::DashMap::new());
710 let persistence = Arc::new(LockedSessionStore::new(storage.clone()));
711 SessionRepository::new(cache, storage, persistence)
712 }
713
714 fn cache_put(repo: &SessionRepository, session: &Session) {
715 repo.cache().insert(
716 session.id.clone(),
717 Arc::new(parking_lot::RwLock::new(session.clone())),
718 );
719 }
720
721 fn task_list(session_id: &str, title: &str) -> bamboo_domain::TaskList {
722 let now = Utc::now();
723 bamboo_domain::TaskList {
724 session_id: session_id.to_string(),
725 title: title.to_string(),
726 items: Vec::new(),
727 created_at: now,
728 updated_at: now,
729 }
730 }
731
732 #[tokio::test]
733 async fn paired_task_cas_narrowly_updates_child_and_root_cache() {
734 let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
735 let repo = test_repo(storage.clone());
736 let root_id = "paired-cache-root";
737 let child_id = "paired-cache-child";
738 let expected_task_list = task_list(root_id, "old shared");
739
740 let mut root = Session::new(root_id, "model");
741 root.add_message(bamboo_agent_core::Message::user("root transcript"));
742 root.metadata
743 .insert("unrelated.root".to_string(), "keep".to_string());
744 root.set_task_list(expected_task_list.clone());
745 root.set_task_list_version_meta("1");
746 storage.save_session(&root).await.unwrap();
747 cache_put(&repo, &root);
748
749 let mut child = Session::new_child(child_id, root_id, "model", "child");
750 child.add_message(bamboo_agent_core::Message::user("child transcript"));
751 child
752 .metadata
753 .insert("unrelated.child".to_string(), "keep".to_string());
754 child.set_task_list(expected_task_list.clone());
755 child.set_task_list_version_meta("1");
756 storage.save_session(&child).await.unwrap();
757 cache_put(&repo, &child);
758
759 assert!(
760 bamboo_domain::RuntimeSessionPersistence::update_task_list_control_planes_if_version(
761 &repo,
762 child_id,
763 root_id,
764 "1",
765 &expected_task_list,
766 &task_list(root_id, "evaluated"),
767 "2",
768 )
769 .await
770 .expect("paired cache CAS succeeds")
771 );
772
773 let cached_root = read_cached_session(repo.cache(), root_id).expect("cached root");
774 let cached_child = read_cached_session(repo.cache(), child_id).expect("cached child");
775 for (session, transcript, metadata_key) in [
776 (&cached_root, "root transcript", "unrelated.root"),
777 (&cached_child, "child transcript", "unrelated.child"),
778 ] {
779 assert_eq!(session.task_list_version_meta().as_deref(), Some("2"));
780 assert_eq!(
781 session.task_list.as_ref().map(|list| list.title.as_str()),
782 Some("evaluated")
783 );
784 assert_eq!(session.messages.len(), 1);
785 assert_eq!(session.messages[0].content, transcript);
786 assert_eq!(
787 session.metadata.get(metadata_key).map(String::as_str),
788 Some("keep")
789 );
790 }
791 }
792
793 #[tokio::test]
794 async fn failed_paired_task_transaction_does_not_publish_child_or_root_cache() {
795 let concrete = Arc::new(MapStorage::default());
796 let storage: Arc<dyn Storage> = concrete.clone();
797 let repo = test_repo(storage.clone());
798 let root_id = "paired-cache-failure-root";
799 let child_id = "paired-cache-failure-child";
800 let expected_task_list = task_list(root_id, "old shared");
801
802 let mut root = Session::new(root_id, "model");
803 root.add_message(bamboo_agent_core::Message::user("root transcript"));
804 root.metadata
805 .insert("unrelated.root".to_string(), "keep".to_string());
806 root.set_task_list(expected_task_list.clone());
807 root.set_task_list_version_meta("1");
808 storage.save_session(&root).await.unwrap();
809 cache_put(&repo, &root);
810
811 let mut child = Session::new_child(child_id, root_id, "model", "child");
812 child.add_message(bamboo_agent_core::Message::user("child transcript"));
813 child
814 .metadata
815 .insert("unrelated.child".to_string(), "keep".to_string());
816 child.set_task_list(expected_task_list.clone());
817 child.set_task_list_version_meta("1");
818 storage.save_session(&child).await.unwrap();
819 cache_put(&repo, &child);
820
821 concrete.fail_pair_commit.store(true, Ordering::SeqCst);
822 let error =
823 bamboo_domain::RuntimeSessionPersistence::update_task_list_control_planes_if_version(
824 &repo,
825 child_id,
826 root_id,
827 "1",
828 &expected_task_list,
829 &task_list(root_id, "must not publish"),
830 "2",
831 )
832 .await
833 .expect_err("paired durable transaction fails");
834 assert!(error.to_string().contains("injected paired"));
835
836 for (id, expected_title, transcript, metadata_key) in [
837 (root_id, "old shared", "root transcript", "unrelated.root"),
838 (
839 child_id,
840 "old shared",
841 "child transcript",
842 "unrelated.child",
843 ),
844 ] {
845 let cached = read_cached_session(repo.cache(), id).expect("cached session remains");
846 assert_eq!(cached.task_list_version_meta().as_deref(), Some("1"));
847 assert_eq!(
848 cached.task_list.as_ref().map(|list| list.title.as_str()),
849 Some(expected_title)
850 );
851 assert_eq!(cached.messages[0].content, transcript);
852 assert_eq!(
853 cached.metadata.get(metadata_key).map(String::as_str),
854 Some("keep")
855 );
856
857 let durable = storage.load_session(id).await.unwrap().unwrap();
858 assert_eq!(durable.task_list_version_meta().as_deref(), Some("1"));
859 assert_eq!(
860 durable.task_list.as_ref().map(|list| list.title.as_str()),
861 Some(expected_title)
862 );
863 }
864 }
865
866 fn durable_cache_fence(
867 operation: impl Into<String>,
868 marker: impl Into<String>,
869 ) -> (
870 PostDurableHook,
871 tokio::sync::oneshot::Receiver<()>,
872 Arc<(Mutex<bool>, Condvar)>,
873 ) {
874 let operation = operation.into();
875 let marker = marker.into();
876 let (started_tx, started_rx) = tokio::sync::oneshot::channel();
877 let started_tx = Arc::new(Mutex::new(Some(started_tx)));
878 let release = Arc::new((Mutex::new(false), Condvar::new()));
879 let hook_release = release.clone();
880 let hook: PostDurableHook = Arc::new(move |actual_operation, actual_marker| {
881 if actual_operation != operation || actual_marker != marker {
882 return;
883 }
884 if let Some(started_tx) = started_tx.lock().unwrap().take() {
885 started_tx.send(()).expect("fence observer still present");
886 }
887 let (released, wake) = &*hook_release;
888 let mut released = released.lock().unwrap();
889 while !*released {
890 released = wake.wait(released).unwrap();
891 }
892 });
893 (hook, started_rx, release)
894 }
895
896 fn release_fence(release: &Arc<(Mutex<bool>, Condvar)>) {
897 let (released, wake) = &**release;
898 *released.lock().unwrap() = true;
899 wake.notify_all();
900 }
901
902 async fn assert_second_write_waits_for_cache_publish<T>(
903 first: tokio::task::JoinHandle<std::io::Result<T>>,
904 mut second: tokio::task::JoinHandle<std::io::Result<T>>,
905 release: Arc<(Mutex<bool>, Condvar)>,
906 ) -> (T, T) {
907 let second_before_release =
908 tokio::time::timeout(Duration::from_millis(100), &mut second).await;
909 let completed_before_release = second_before_release.is_ok();
910 release_fence(&release);
911
912 let first = first
913 .await
914 .expect("first writer joins")
915 .expect("first writer succeeds");
916 let second = match second_before_release {
917 Ok(joined) => joined
918 .expect("second writer joins")
919 .expect("second writer succeeds"),
920 Err(_) => second
921 .await
922 .expect("second writer joins")
923 .expect("second writer succeeds"),
924 };
925 assert!(
926 !completed_before_release,
927 "the second write must remain behind the first write's durable-to-cache fence"
928 );
929 (first, second)
930 }
931
932 #[derive(Clone, Copy, Debug)]
933 enum FullSaveRoute {
934 InherentSave,
935 SaveAndCache,
936 RuntimePersistence,
937 }
938
939 impl FullSaveRoute {
940 fn operation(self) -> &'static str {
941 match self {
942 Self::InherentSave => "save_full",
943 Self::SaveAndCache => "save_and_cache",
944 Self::RuntimePersistence => "save_runtime_session",
945 }
946 }
947
948 fn name(self) -> &'static str {
949 match self {
950 Self::InherentSave => "inherent",
951 Self::SaveAndCache => "save-and-cache",
952 Self::RuntimePersistence => "runtime-persistence",
953 }
954 }
955 }
956
957 async fn assert_full_save_route_serializes_cache_publish(route: FullSaveRoute) {
958 let temp = tempfile::tempdir().unwrap();
959 let concrete_storage = Arc::new(
960 bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
961 .await
962 .expect("SessionStoreV2"),
963 );
964 let storage: Arc<dyn Storage> = concrete_storage;
965 let id = format!("root-full-cache-order-{}", route.name());
966 let (hook, first_durable, release) = durable_cache_fence(route.operation(), id.clone());
967 let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
968
969 let mut initial = Session::new(&id, "model");
970 initial.add_message(bamboo_agent_core::Message::user("durable transcript"));
971 initial.set_task_list(task_list(&id, "initial"));
972 initial.set_task_list_version_meta("0");
973 initial
974 .metadata
975 .insert("unrelated.runtime".to_string(), "keep".to_string());
976 storage.save_session(&initial).await.unwrap();
977 cache_put(&repo, &initial);
978
979 let first_repo = repo.clone();
980 let mut root_snapshot = initial.clone();
981 root_snapshot.add_message(bamboo_agent_core::Message::assistant(
982 "full-save transcript suffix",
983 None,
984 ));
985 root_snapshot.set_task_list(task_list(&id, "root"));
986 root_snapshot.set_task_list_version_meta("1");
987 let first = tokio::spawn(async move {
988 match route {
989 FullSaveRoute::InherentSave => first_repo.save(&mut root_snapshot).await,
990 FullSaveRoute::SaveAndCache => {
991 first_repo.save_and_cache(&mut root_snapshot).await;
992 Ok(())
993 }
994 FullSaveRoute::RuntimePersistence => {
995 bamboo_domain::RuntimeSessionPersistence::save_runtime_session(
996 first_repo.as_ref(),
997 &mut root_snapshot,
998 )
999 .await
1000 }
1001 }
1002 });
1003 first_durable
1004 .await
1005 .expect("root full durable write reached");
1006
1007 let second_repo = repo.clone();
1008 let second_id = id.clone();
1009 let child_task_list = task_list(&id, "child");
1010 let second = tokio::spawn(async move {
1011 bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
1012 second_repo.as_ref(),
1013 &second_id,
1014 &child_task_list,
1015 "2",
1016 )
1017 .await
1018 .map(|updated| assert!(updated, "root must exist"))
1019 });
1020 assert_second_write_waits_for_cache_publish(first, second, release).await;
1021
1022 let durable = storage.load_session(&id).await.unwrap().unwrap();
1023 let cached = read_cached_session(repo.cache(), &id).expect("cached root");
1024 for (tier, session) in [("durable", durable), ("cache", cached)] {
1025 assert_eq!(
1026 session.task_list_version_meta().as_deref(),
1027 Some("2"),
1028 "{route:?} {tier} must retain the child transaction"
1029 );
1030 assert_eq!(
1031 session.task_list.as_ref().map(|list| list.title.as_str()),
1032 Some("child"),
1033 "{route:?} {tier} must retain the child transaction"
1034 );
1035 assert_eq!(
1036 session
1037 .metadata
1038 .get("unrelated.runtime")
1039 .map(String::as_str),
1040 Some("keep"),
1041 "{route:?} {tier} must preserve unrelated runtime state"
1042 );
1043 assert_eq!(
1044 session.messages.len(),
1045 2,
1046 "{route:?} {tier} must preserve the full-save transcript"
1047 );
1048 }
1049 }
1050
1051 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1052 async fn root_full_saves_and_child_task_patch_share_publish_order() {
1053 for route in [
1054 FullSaveRoute::InherentSave,
1055 FullSaveRoute::SaveAndCache,
1056 FullSaveRoute::RuntimePersistence,
1057 ] {
1058 assert_full_save_route_serializes_cache_publish(route).await;
1059 }
1060 }
1061
1062 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1063 async fn checkpoint_and_child_task_patch_share_publish_order() {
1064 let temp = tempfile::tempdir().unwrap();
1065 let concrete_storage = Arc::new(
1066 bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
1067 .await
1068 .expect("SessionStoreV2"),
1069 );
1070 let storage: Arc<dyn Storage> = concrete_storage;
1071 let id = "checkpoint-cache-order";
1072 let (hook, checkpoint_durable, release) = durable_cache_fence("checkpoint", id);
1073 let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
1074
1075 let mut initial = Session::new(id, "model");
1076 initial.add_message(bamboo_agent_core::Message::user("durable transcript"));
1077 initial.set_task_list(task_list(id, "initial"));
1078 initial.set_task_list_version_meta("0");
1079 initial
1080 .metadata
1081 .insert("unrelated.runtime".to_string(), "keep".to_string());
1082 storage.save_session(&initial).await.unwrap();
1083 cache_put(&repo, &initial);
1084
1085 let checkpoint_repo = repo.clone();
1086 let mut checkpoint_snapshot = initial.clone();
1087 checkpoint_snapshot.add_message(bamboo_agent_core::Message::assistant(
1088 "checkpoint transcript suffix",
1089 None,
1090 ));
1091 checkpoint_snapshot.set_task_list(task_list(id, "checkpoint"));
1092 checkpoint_snapshot.set_task_list_version_meta("1");
1093 let checkpoint = tokio::spawn(async move {
1094 bamboo_domain::RuntimeSessionPersistence::checkpoint_runtime_session(
1095 checkpoint_repo.as_ref(),
1096 &mut checkpoint_snapshot,
1097 )
1098 .await
1099 });
1100 checkpoint_durable
1101 .await
1102 .expect("checkpoint durable write reached");
1103
1104 let child_repo = repo.clone();
1105 let child_task_list = task_list(id, "child");
1106 let child_patch = tokio::spawn(async move {
1107 bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
1108 child_repo.as_ref(),
1109 id,
1110 &child_task_list,
1111 "2",
1112 )
1113 .await
1114 .map(|updated| assert!(updated, "root must exist"))
1115 });
1116 assert_second_write_waits_for_cache_publish(checkpoint, child_patch, release).await;
1117
1118 let durable = storage.load_session(id).await.unwrap().unwrap();
1119 let cached = read_cached_session(repo.cache(), id).expect("cached root");
1120 for (tier, session) in [("durable", durable), ("cache", cached)] {
1121 assert_eq!(
1122 session.task_list_version_meta().as_deref(),
1123 Some("2"),
1124 "{tier} must retain the child transaction"
1125 );
1126 assert_eq!(
1127 session.task_list.as_ref().map(|list| list.title.as_str()),
1128 Some("child"),
1129 "{tier} must retain the child transaction"
1130 );
1131 assert_eq!(
1132 session
1133 .metadata
1134 .get("unrelated.runtime")
1135 .map(String::as_str),
1136 Some("keep"),
1137 "{tier} must preserve unrelated runtime state"
1138 );
1139 assert_eq!(
1140 session.messages.len(),
1141 2,
1142 "{tier} must preserve the checkpoint transcript"
1143 );
1144 }
1145 }
1146
1147 #[derive(Clone, Copy, Debug)]
1148 enum CacheBackfillRoute {
1149 Load,
1150 TryLoad,
1151 }
1152
1153 impl CacheBackfillRoute {
1154 fn operation(self) -> &'static str {
1155 match self {
1156 Self::Load => "load",
1157 Self::TryLoad => "try_load",
1158 }
1159 }
1160
1161 fn name(self) -> &'static str {
1162 match self {
1163 Self::Load => "load",
1164 Self::TryLoad => "try-load",
1165 }
1166 }
1167 }
1168
1169 async fn assert_cache_backfill_serializes_with_task_patch(route: CacheBackfillRoute) {
1170 let temp = tempfile::tempdir().unwrap();
1171 let concrete_storage = Arc::new(
1172 bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
1173 .await
1174 .expect("SessionStoreV2"),
1175 );
1176 let storage: Arc<dyn Storage> = concrete_storage;
1177 let id = format!("cache-backfill-order-{}", route.name());
1178 let (hook, loaded_old_durable, release) =
1179 durable_cache_fence(route.operation(), id.clone());
1180 let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
1181
1182 let mut initial = Session::new(&id, "model");
1183 initial.set_task_list(task_list(&id, "initial"));
1184 initial.set_task_list_version_meta("0");
1185 storage.save_session(&initial).await.unwrap();
1186 assert!(
1187 read_cached_session(repo.cache(), &id).is_none(),
1188 "the race requires a genuine cache miss"
1189 );
1190
1191 let load_repo = repo.clone();
1192 let load_id = id.clone();
1193 let load = tokio::spawn(async move {
1194 let loaded = match route {
1195 CacheBackfillRoute::Load => load_repo.load(&load_id).await,
1196 CacheBackfillRoute::TryLoad => {
1197 load_repo.try_load(&load_id).await.expect("storage load")
1198 }
1199 };
1200 assert!(loaded.is_some(), "seeded session must load");
1201 Ok(())
1202 });
1203 loaded_old_durable
1204 .await
1205 .expect("old durable snapshot loaded");
1206
1207 let patch_repo = repo.clone();
1208 let patch_id = id.clone();
1209 let child_task_list = task_list(&id, "child");
1210 let patch = tokio::spawn(async move {
1211 bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
1212 patch_repo.as_ref(),
1213 &patch_id,
1214 &child_task_list,
1215 "1",
1216 )
1217 .await
1218 .map(|updated| assert!(updated, "root must exist"))
1219 });
1220 assert_second_write_waits_for_cache_publish(load, patch, release).await;
1221
1222 let durable = storage.load_session(&id).await.unwrap().unwrap();
1223 let cached = read_cached_session(repo.cache(), &id).expect("backfilled cache");
1224 for (tier, session) in [("durable", durable), ("cache", cached)] {
1225 assert_eq!(
1226 session.task_list_version_meta().as_deref(),
1227 Some("1"),
1228 "{route:?} {tier} must retain the child transaction"
1229 );
1230 assert_eq!(
1231 session.task_list.as_ref().map(|list| list.title.as_str()),
1232 Some("child"),
1233 "{route:?} {tier} must retain the child transaction"
1234 );
1235 }
1236 }
1237
1238 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1239 async fn cache_miss_backfills_and_child_task_patch_share_publish_order() {
1240 for route in [CacheBackfillRoute::Load, CacheBackfillRoute::TryLoad] {
1241 assert_cache_backfill_serializes_with_task_patch(route).await;
1242 }
1243 }
1244
1245 #[tokio::test]
1246 async fn cache_hits_do_not_wait_for_the_persistence_lock() {
1247 let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
1248 let repo = test_repo(storage);
1249 let id = "cache-hit-lock-free";
1250 let cached = Session::new(id, "cached-model");
1251 cache_put(&repo, &cached);
1252 let persistence_guard = repo.persistence().acquire_lock(id).await;
1253
1254 let loaded = tokio::time::timeout(Duration::from_millis(100), repo.load(id)).await;
1255 let try_loaded = tokio::time::timeout(Duration::from_millis(100), repo.try_load(id)).await;
1256 drop(persistence_guard);
1257
1258 assert_eq!(
1259 loaded
1260 .expect("cache hit must not wait for the persistence lock")
1261 .expect("cached session")
1262 .model,
1263 "cached-model"
1264 );
1265 assert_eq!(
1266 try_loaded
1267 .expect("fallible cache hit must not wait for the persistence lock")
1268 .expect("cache read succeeds")
1269 .expect("cached session")
1270 .model,
1271 "cached-model"
1272 );
1273 }
1274
1275 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1276 async fn load_merged_storage_refresh_and_child_task_patch_share_publish_order() {
1277 let temp = tempfile::tempdir().unwrap();
1278 let concrete_storage = Arc::new(
1279 bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
1280 .await
1281 .expect("SessionStoreV2"),
1282 );
1283 let storage: Arc<dyn Storage> = concrete_storage;
1284 let id = "load-merged-cache-order";
1285 let (hook, loaded_old_durable, release) = durable_cache_fence("load_merged", id);
1286 let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
1287
1288 let mut durable = Session::new(id, "model");
1289 durable.updated_at = Utc::now();
1290 durable.set_task_list(task_list(id, "initial"));
1291 durable.set_task_list_version_meta("0");
1292 storage.save_session(&durable).await.unwrap();
1293
1294 let mut memory = durable.clone();
1295 memory.updated_at = durable.updated_at - chrono::Duration::seconds(1);
1296 memory.set_task_list(task_list(id, "memory"));
1297 cache_put(&repo, &memory);
1298
1299 let load_repo = repo.clone();
1300 let load = tokio::spawn(async move {
1301 assert!(
1302 load_repo.load_merged(id).await.is_some(),
1303 "seeded session must load"
1304 );
1305 Ok(())
1306 });
1307 loaded_old_durable
1308 .await
1309 .expect("old durable snapshot loaded");
1310
1311 let patch_repo = repo.clone();
1312 let child_task_list = task_list(id, "child");
1313 let patch = tokio::spawn(async move {
1314 bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
1315 patch_repo.as_ref(),
1316 id,
1317 &child_task_list,
1318 "1",
1319 )
1320 .await
1321 .map(|updated| assert!(updated, "root must exist"))
1322 });
1323 assert_second_write_waits_for_cache_publish(load, patch, release).await;
1324
1325 let durable = storage.load_session(id).await.unwrap().unwrap();
1326 let cached = read_cached_session(repo.cache(), id).expect("refreshed cache");
1327 for (tier, session) in [("durable", durable), ("cache", cached)] {
1328 assert_eq!(
1329 session.task_list_version_meta().as_deref(),
1330 Some("1"),
1331 "{tier} must retain the child transaction"
1332 );
1333 assert_eq!(
1334 session.task_list.as_ref().map(|list| list.title.as_str()),
1335 Some("child"),
1336 "{tier} must retain the child transaction"
1337 );
1338 }
1339 }
1340
1341 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1342 async fn legacy_clear_refresh_and_child_task_patch_share_publish_order() {
1343 let temp = tempfile::tempdir().unwrap();
1344 let concrete_storage = Arc::new(
1345 bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
1346 .await
1347 .expect("SessionStoreV2"),
1348 );
1349 let storage: Arc<dyn Storage> = concrete_storage;
1350 let id = "legacy-clear-cache-order";
1351 let expected = vec![serde_json::json!({"content": "legacy"})];
1352 let (hook, loaded_post_cas_snapshot, release) = durable_cache_fence("clear_legacy", id);
1353 let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
1354
1355 let mut initial = Session::new(id, "model");
1356 initial.set_pending_injected_messages(expected.clone());
1357 initial.set_task_list(task_list(id, "initial"));
1358 initial.set_task_list_version_meta("0");
1359 storage.save_session(&initial).await.unwrap();
1360 cache_put(&repo, &initial);
1361
1362 let clear_repo = repo.clone();
1363 let clear_expected = expected.clone();
1364 let clear = tokio::spawn(async move {
1365 bamboo_domain::RuntimeSessionPersistence::clear_legacy_pending_messages(
1366 clear_repo.as_ref(),
1367 id,
1368 &clear_expected,
1369 )
1370 .await
1371 .map(|cleared| assert!(cleared, "legacy queue must match"))
1372 });
1373 loaded_post_cas_snapshot
1374 .await
1375 .expect("post-CAS snapshot loaded");
1376
1377 let patch_repo = repo.clone();
1378 let child_task_list = task_list(id, "child");
1379 let patch = tokio::spawn(async move {
1380 bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
1381 patch_repo.as_ref(),
1382 id,
1383 &child_task_list,
1384 "1",
1385 )
1386 .await
1387 .map(|updated| assert!(updated, "root must exist"))
1388 });
1389 assert_second_write_waits_for_cache_publish(clear, patch, release).await;
1390
1391 let durable = storage.load_session(id).await.unwrap().unwrap();
1392 let cached = read_cached_session(repo.cache(), id).expect("refreshed cache");
1393 for (tier, session) in [("durable", durable), ("cache", cached)] {
1394 assert_eq!(
1395 session.task_list_version_meta().as_deref(),
1396 Some("1"),
1397 "{tier} must retain the child transaction"
1398 );
1399 assert_eq!(
1400 session.task_list.as_ref().map(|list| list.title.as_str()),
1401 Some("child"),
1402 "{tier} must retain the child transaction"
1403 );
1404 assert!(
1405 !session.has_pending_injected_messages(),
1406 "{tier} must retain the successful legacy clear"
1407 );
1408 }
1409 }
1410
1411 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1412 async fn concurrent_task_patches_publish_cache_in_durable_order() {
1413 let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
1414 let (hook, first_durable, release) = durable_cache_fence("task", "1");
1415 let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
1416 let id = "concurrent-task-cache-order";
1417 let mut initial = Session::new(id, "model");
1418 initial.set_task_list(task_list(id, "initial"));
1419 initial.set_task_list_version_meta("0");
1420 storage.save_session(&initial).await.unwrap();
1421 cache_put(&repo, &initial);
1422
1423 let first_repo = repo.clone();
1424 let first_task_list = task_list(id, "first");
1425 let first = tokio::spawn(async move {
1426 bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
1427 first_repo.as_ref(),
1428 id,
1429 &first_task_list,
1430 "1",
1431 )
1432 .await
1433 });
1434 first_durable.await.expect("first durable write reached");
1435
1436 let second_repo = repo.clone();
1437 let second_task_list = task_list(id, "second");
1438 let second = tokio::spawn(async move {
1439 bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
1440 second_repo.as_ref(),
1441 id,
1442 &second_task_list,
1443 "2",
1444 )
1445 .await
1446 });
1447 let (first_updated, second_updated) =
1448 assert_second_write_waits_for_cache_publish(first, second, release).await;
1449 assert!(first_updated && second_updated);
1450
1451 let durable = storage.load_session(id).await.unwrap().unwrap();
1452 let cached = read_cached_session(repo.cache(), id).expect("cached root");
1453 for (tier, session) in [("durable", durable), ("cache", cached)] {
1454 assert_eq!(
1455 session.task_list_version_meta().as_deref(),
1456 Some("2"),
1457 "{tier} must retain the second transaction"
1458 );
1459 assert_eq!(
1460 session.task_list.as_ref().map(|list| list.title.as_str()),
1461 Some("second"),
1462 "{tier} must retain the second transaction"
1463 );
1464 }
1465 }
1466
1467 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1468 async fn root_control_plane_save_and_child_task_patch_share_publish_order() {
1469 let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
1470 let (hook, root_durable, release) = durable_cache_fence("save", "root-control-plane-order");
1471 let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
1472 let id = "root-control-plane-order";
1473
1474 let mut initial = Session::new(id, "model");
1475 initial.add_message(bamboo_agent_core::Message::user("durable transcript"));
1476 initial.set_task_list(task_list(id, "initial"));
1477 initial.set_task_list_version_meta("0");
1478 initial
1479 .metadata
1480 .insert("unrelated.runtime".to_string(), "keep".to_string());
1481 storage.save_session(&initial).await.unwrap();
1482 cache_put(&repo, &initial);
1483
1484 let root_repo = repo.clone();
1485 let mut root_snapshot = initial.clone();
1486 root_snapshot.set_task_list(task_list(id, "root"));
1487 root_snapshot.set_task_list_version_meta("1");
1488 let root_save = tokio::spawn(async move {
1489 bamboo_domain::RuntimeSessionPersistence::save_runtime_control_plane(
1490 root_repo.as_ref(),
1491 &mut root_snapshot,
1492 )
1493 .await
1494 });
1495 root_durable.await.expect("root durable write reached");
1496
1497 let child_repo = repo.clone();
1498 let child_task_list = task_list(id, "child");
1499 let child_patch = tokio::spawn(async move {
1500 bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
1501 child_repo.as_ref(),
1502 id,
1503 &child_task_list,
1504 "2",
1505 )
1506 .await
1507 .map(|updated| {
1508 assert!(updated, "root must exist");
1509 })
1510 });
1511 assert_second_write_waits_for_cache_publish(root_save, child_patch, release).await;
1512
1513 let durable = storage.load_session(id).await.unwrap().unwrap();
1514 let cached = read_cached_session(repo.cache(), id).expect("cached root");
1515 for (tier, session) in [("durable", durable), ("cache", cached)] {
1516 assert_eq!(
1517 session.task_list_version_meta().as_deref(),
1518 Some("2"),
1519 "{tier} must retain the child transaction"
1520 );
1521 assert_eq!(
1522 session.task_list.as_ref().map(|list| list.title.as_str()),
1523 Some("child"),
1524 "{tier} must retain the child transaction"
1525 );
1526 assert_eq!(
1527 session
1528 .metadata
1529 .get("unrelated.runtime")
1530 .map(String::as_str),
1531 Some("keep"),
1532 "{tier} must preserve unrelated runtime state"
1533 );
1534 assert_eq!(
1535 session.messages.len(),
1536 1,
1537 "{tier} must preserve the transcript"
1538 );
1539 }
1540 }
1541
1542 #[tokio::test]
1543 async fn stale_control_plane_save_keeps_checkpointed_ledger_in_durable_and_cache() {
1544 let temp = tempfile::tempdir().unwrap();
1545 let storage: Arc<dyn Storage> = Arc::new(
1546 bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
1547 .await
1548 .expect("SessionStoreV2"),
1549 );
1550 let repo = test_repo(storage.clone());
1551 let id = "control-plane-ledger-cache";
1552 let mut initial = Session::new(id, "model");
1553 initial.add_message(bamboo_agent_core::Message::user("durable transcript"));
1554 storage.save_session(&initial).await.unwrap();
1555 cache_put(&repo, &initial);
1556 let mut stale = initial.clone();
1557
1558 let mut runner = initial;
1559 runner.model_context_state = Some(bamboo_domain::ModelContextState {
1560 state_revision: 1,
1561 prefix_epoch: 1,
1562 cache_scope_sha256: Some("scope".to_string()),
1563 transcript_item_sha256: vec!["runner-l1".to_string()],
1564 ..bamboo_domain::ModelContextState::default()
1565 });
1566 bamboo_domain::RuntimeSessionPersistence::checkpoint_runtime_session(&repo, &mut runner)
1567 .await
1568 .unwrap();
1569
1570 stale
1571 .metadata
1572 .insert("runtime.suspend_reason".to_string(), "waiting".to_string());
1573 bamboo_domain::RuntimeSessionPersistence::save_runtime_control_plane(&repo, &mut stale)
1574 .await
1575 .unwrap();
1576
1577 let expected = runner.model_context_state;
1578 let durable = storage.load_session(id).await.unwrap().unwrap();
1579 let cached = read_cached_session(repo.cache(), id).expect("cached session");
1580 for (tier, session) in [("durable", durable), ("cache", cached)] {
1581 assert_eq!(session.model_context_state, expected, "tier={tier}");
1582 assert_eq!(
1583 session
1584 .metadata
1585 .get("runtime.suspend_reason")
1586 .map(String::as_str),
1587 Some("waiting"),
1588 "tier={tier}"
1589 );
1590 assert_eq!(session.messages.len(), 1, "tier={tier}");
1591 }
1592 }
1593
1594 #[tokio::test]
1595 async fn narrow_runtime_metadata_transaction_preserves_live_and_durable_non_owned_state() {
1596 let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
1597 let repo = test_repo(storage.clone());
1598 let id = "narrow-metadata";
1599 let mut durable = Session::new(id, "durable-model");
1600 durable.add_message(bamboo_agent_core::Message::user("durable user turn"));
1601 durable
1602 .metadata
1603 .insert("external.durable".to_string(), "keep".to_string());
1604 storage.save_session(&durable).await.expect("seed durable");
1605
1606 let mut live = durable.clone();
1607 live.add_message(bamboo_agent_core::Message::assistant(
1608 "in-flight assistant tool call",
1609 None,
1610 ));
1611 live.model = "live-model".to_string();
1612 live.metadata
1613 .insert("external.live".to_string(), "keep".to_string());
1614 cache_put(&repo, &live);
1615
1616 repo.update_runtime_session(id, &["workflow.owned"], |latest| {
1617 latest
1618 .metadata
1619 .insert("workflow.owned".to_string(), "active".to_string());
1620 })
1621 .await
1622 .expect("transaction")
1623 .expect("session exists");
1624
1625 let saved = storage
1626 .load_session(id)
1627 .await
1628 .expect("load durable")
1629 .expect("durable exists");
1630 assert_eq!(
1631 saved.messages.len(),
1632 1,
1633 "transaction never writes stale live messages"
1634 );
1635 assert_eq!(
1636 saved.metadata.get("external.durable").map(String::as_str),
1637 Some("keep")
1638 );
1639 assert_eq!(
1640 saved.metadata.get("workflow.owned").map(String::as_str),
1641 Some("active")
1642 );
1643
1644 let cached = read_cached_session(repo.cache(), id).expect("live cache");
1645 assert_eq!(
1646 cached.messages.len(),
1647 2,
1648 "cache live tool call is not replaced"
1649 );
1650 assert_eq!(cached.model, "live-model");
1651 assert_eq!(
1652 cached.metadata.get("external.live").map(String::as_str),
1653 Some("keep")
1654 );
1655 assert_eq!(
1656 cached.metadata.get("workflow.owned").map(String::as_str),
1657 Some("active")
1658 );
1659 }
1660
1661 #[tokio::test]
1666 async fn load_merged_does_not_regress_to_older_storage() {
1667 let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
1668 let repo = test_repo(storage.clone());
1669 let id = "s1";
1670
1671 let mut stale = Session::new(id.to_string(), "m");
1672 stale.set_pending_question(
1673 "tc1".into(),
1674 "kind".into(),
1675 "q?".into(),
1676 vec!["OK".into()],
1677 true,
1678 );
1679 stale.updated_at = Utc::now() - chrono::Duration::seconds(10);
1680 storage.save_session(&stale).await.unwrap();
1681
1682 let mut fresh = Session::new(id.to_string(), "m");
1683 fresh.updated_at = Utc::now();
1684 cache_put(&repo, &fresh);
1685
1686 let merged = repo.load_merged(id).await.expect("session exists");
1687 assert!(
1688 merged.pending_question.is_none(),
1689 "must return the newer answered memory copy, not the stale storage one"
1690 );
1691 let cached = read_cached_session(repo.cache(), id).expect("cached");
1692 assert!(
1693 cached.pending_question.is_none(),
1694 "load_merged must never regress the cache to a stale storage copy"
1695 );
1696 }
1697
1698 #[tokio::test]
1702 async fn load_merged_recovers_pending_question_from_same_age_storage() {
1703 let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
1704 let repo = test_repo(storage.clone());
1705 let id = "s2";
1706 let ts = Utc::now();
1707
1708 let mut with_pending = Session::new(id.to_string(), "m");
1709 with_pending.set_pending_question(
1710 "tc".into(),
1711 "k".into(),
1712 "q".into(),
1713 vec!["OK".into()],
1714 true,
1715 );
1716 with_pending.updated_at = ts;
1717 storage.save_session(&with_pending).await.unwrap();
1718
1719 let mut lost = with_pending.clone();
1720 lost.clear_pending_question();
1721 lost.updated_at = ts;
1722 cache_put(&repo, &lost);
1723
1724 let merged = repo.load_merged(id).await.expect("session exists");
1725 assert!(
1726 merged.pending_question.is_some(),
1727 "same-age storage carrying a pending question must still be recovered"
1728 );
1729 }
1730
1731 #[tokio::test]
1732 async fn runtime_publish_refreshes_cache_even_when_storage_fails() {
1733 let id = "runtime-selection";
1734 let mut previous = Session::new(id.to_string(), "m");
1735 previous.metadata.insert(
1736 "skill_runtime_selected_skill_ids".to_string(),
1737 "[\"plan\"]".to_string(),
1738 );
1739 let storage: Arc<dyn Storage> = Arc::new(FailingSaveStorage {
1740 persisted: Mutex::new(Some(previous.clone())),
1741 });
1742 let repo = test_repo(storage.clone());
1743 cache_put(&repo, &previous);
1744
1745 let mut current = previous.clone();
1746 current.metadata.insert(
1747 "skill_runtime_selected_skill_ids".to_string(),
1748 "[\"review\"]".to_string(),
1749 );
1750 current.updated_at = Utc::now();
1751
1752 let result =
1753 bamboo_domain::RuntimeSessionPersistence::save_runtime_session(&repo, &mut current)
1754 .await;
1755 assert!(result.is_err(), "durable failure must still be surfaced");
1756
1757 let cached = repo.load(id).await.expect("cached current session");
1758 assert_eq!(
1759 cached
1760 .metadata
1761 .get("skill_runtime_selected_skill_ids")
1762 .map(String::as_str),
1763 Some("[\"review\"]")
1764 );
1765 let allowlist = bamboo_skills::access_control::extract_skill_allowlist(&cached.metadata)
1766 .expect("runtime authorization allowlist");
1767 assert!(allowlist.contains("review"));
1768 assert!(!allowlist.contains("plan"));
1769 let durable = storage
1770 .load_session(id)
1771 .await
1772 .expect("load durable state")
1773 .expect("previous durable session");
1774 assert_eq!(
1775 durable
1776 .metadata
1777 .get("skill_runtime_selected_skill_ids")
1778 .map(String::as_str),
1779 Some("[\"plan\"]")
1780 );
1781 }
1782
1783 #[tokio::test]
1784 async fn inherent_save_leaves_existing_cache_untouched_when_storage_fails() {
1785 let id = "inherent-save-failure";
1786 let previous = Session::new(id, "previous");
1787 let storage: Arc<dyn Storage> = Arc::new(FailingSaveStorage {
1788 persisted: Mutex::new(Some(previous.clone())),
1789 });
1790 let repo = test_repo(storage);
1791 cache_put(&repo, &previous);
1792
1793 let mut current = previous.clone();
1794 current.model = "current".to_string();
1795 assert!(repo.save(&mut current).await.is_err());
1796 assert_eq!(
1797 read_cached_session(repo.cache(), id)
1798 .expect("existing cache")
1799 .model,
1800 "previous",
1801 "fallible inherent save must publish only after a durable commit"
1802 );
1803 }
1804
1805 #[tokio::test]
1806 async fn save_and_cache_still_refreshes_cache_when_storage_fails() {
1807 let id = "save-and-cache-failure";
1808 let previous = Session::new(id, "previous");
1809 let storage: Arc<dyn Storage> = Arc::new(FailingSaveStorage {
1810 persisted: Mutex::new(Some(previous.clone())),
1811 });
1812 let repo = test_repo(storage);
1813 cache_put(&repo, &previous);
1814
1815 let mut current = previous;
1816 current.model = "current".to_string();
1817 repo.save_and_cache(&mut current).await;
1818 assert_eq!(
1819 read_cached_session(repo.cache(), id)
1820 .expect("refreshed cache")
1821 .model,
1822 "current",
1823 "fire-and-forget save must retain its existing cache-on-failure behavior"
1824 );
1825 }
1826
1827 #[tokio::test]
1828 async fn checkpoint_leaves_existing_cache_untouched_when_storage_fails() {
1829 let id = "checkpoint-failure";
1830 let previous = Session::new(id, "previous");
1831 let storage: Arc<dyn Storage> = Arc::new(FailingSaveStorage {
1832 persisted: Mutex::new(Some(previous.clone())),
1833 });
1834 let repo = test_repo(storage);
1835 cache_put(&repo, &previous);
1836
1837 let mut current = previous.clone();
1838 current.model = "current".to_string();
1839 let result = bamboo_domain::RuntimeSessionPersistence::checkpoint_runtime_session(
1840 &repo,
1841 &mut current,
1842 )
1843 .await;
1844 assert!(result.is_err());
1845 assert_eq!(
1846 read_cached_session(repo.cache(), id)
1847 .expect("existing cache")
1848 .model,
1849 "previous",
1850 "checkpoint must publish only after a durable commit"
1851 );
1852 }
1853
1854 #[tokio::test]
1855 async fn legacy_clear_uses_durable_cas_and_never_erases_concurrent_append_from_stale_cache() {
1856 let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
1857 let repo = test_repo(storage.clone());
1858 let id = "legacy-cas-race";
1859 let expected = vec![serde_json::json!({"content": "first"})];
1860
1861 let mut stale_cache = Session::new(id, "m");
1862 stale_cache.set_pending_injected_messages(expected.clone());
1863 cache_put(&repo, &stale_cache);
1864
1865 let mut durable = stale_cache.clone();
1866 durable.set_pending_injected_messages(vec![
1867 serde_json::json!({"content": "first"}),
1868 serde_json::json!({"content": "concurrent"}),
1869 ]);
1870 storage.save_session(&durable).await.unwrap();
1871
1872 let cleared = bamboo_domain::RuntimeSessionPersistence::clear_legacy_pending_messages(
1873 &repo, id, &expected,
1874 )
1875 .await
1876 .unwrap();
1877 assert!(!cleared, "the durable compare-and-clear must reject drift");
1878 assert_eq!(
1879 storage
1880 .load_session(id)
1881 .await
1882 .unwrap()
1883 .unwrap()
1884 .pending_injected_messages()
1885 .unwrap(),
1886 durable.pending_injected_messages().unwrap(),
1887 "the concurrent durable append must remain intact"
1888 );
1889 assert_eq!(
1890 read_cached_session(repo.cache(), id)
1891 .unwrap()
1892 .pending_injected_messages()
1893 .unwrap(),
1894 expected,
1895 "a failed CAS must not mutate the existing cache"
1896 );
1897 }
1898
1899 #[tokio::test]
1900 async fn successful_legacy_clear_refreshes_stale_cache_from_durable_state() {
1901 let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
1902 let repo = test_repo(storage.clone());
1903 let id = "legacy-cas-success";
1904 let expected = vec![serde_json::json!({"content": "first"})];
1905
1906 let mut stale_cache = Session::new(id, "stale-model");
1907 stale_cache.set_pending_injected_messages(expected.clone());
1908 cache_put(&repo, &stale_cache);
1909
1910 let mut durable = Session::new(id, "durable-model");
1911 durable.set_pending_injected_messages(expected.clone());
1912 durable
1913 .metadata
1914 .insert("durable-only".to_string(), "keep".to_string());
1915 storage.save_session(&durable).await.unwrap();
1916
1917 assert!(
1918 bamboo_domain::RuntimeSessionPersistence::clear_legacy_pending_messages(
1919 &repo, id, &expected,
1920 )
1921 .await
1922 .unwrap()
1923 );
1924 let cached = read_cached_session(repo.cache(), id).unwrap();
1925 assert!(!cached.has_pending_injected_messages());
1926 assert_eq!(cached.model, "durable-model");
1927 assert_eq!(
1928 cached.metadata.get("durable-only").map(String::as_str),
1929 Some("keep")
1930 );
1931 }
1932}