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