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