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