Skip to main content

bamboo_engine/
session_repository.rs

1//! Canonical session coordinator owned by the framework.
2//!
3//! [`SessionRepository`] bundles the three tiers a Bamboo session lives in — the
4//! in-memory [`SessionCache`], the durable [`Storage`], and the
5//! merge-on-write [`LockedSessionStore`] — and provides the one canonical
6//! load/save coordination (cache → storage → backfill, and dual-write).
7//!
8//! This is a *framework* capability, not a server one: previously the
9//! coordination lived only as inherent methods on the server's `AppState`,
10//! which meant anything outside the HTTP server (the SDK, in-process embedders)
11//! could not load or persist sessions consistently. `SessionRepository` lets any
12//! caller that holds the three tiers share the exact same behaviour; the
13//! server's `AppState` now delegates to it.
14
15use 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/// Framework-owned coordinator over a session's cache / storage / persistence
39/// tiers. Cheap to clone (all fields are `Arc`).
40#[derive(Clone)]
41pub struct SessionRepository {
42    cache: SessionCache,
43    storage: Arc<dyn Storage>,
44    persistence: Arc<LockedSessionStore>,
45    root_actor_directory: Option<Arc<dyn bamboo_domain::ActorDirectoryPort>>,
46    root_actor_owner: Option<(
47        Arc<dyn bamboo_domain::ActorDirectoryPort>,
48        Arc<bamboo_domain::RootActorRuntimeWrite>,
49    )>,
50    root_inbox_receipt: Option<bamboo_domain::SessionInboxReceipt>,
51    root_actor_writers:
52        Arc<dashmap::DashMap<String, std::sync::Weak<bamboo_domain::RootActorRuntimeWrite>>>,
53    #[cfg(test)]
54    post_durable_hook: Option<PostDurableHook>,
55}
56
57impl SessionRepository {
58    pub fn new(
59        cache: SessionCache,
60        storage: Arc<dyn Storage>,
61        persistence: Arc<LockedSessionStore>,
62    ) -> Self {
63        Self {
64            cache,
65            storage,
66            persistence,
67            root_actor_directory: None,
68            root_actor_owner: None,
69            root_inbox_receipt: None,
70            root_actor_writers: Arc::new(Default::default()),
71            #[cfg(test)]
72            post_durable_hook: None,
73        }
74    }
75
76    async fn bind_root_actor_execution_inner(
77        &self,
78        session: &Session,
79        run_id: &str,
80    ) -> std::io::Result<Option<bamboo_domain::RootActorExecutionBinding>> {
81        if session.kind != bamboo_domain::SessionKind::Root
82            || !session.authority_identity.is_ordinary()
83        {
84            return Ok(None);
85        }
86        let Some(directory) = self.root_actor_directory.clone() else {
87            return Ok(None);
88        };
89        if !self.storage.supports_root_actor_runtime_write() {
90            return Err(std::io::Error::new(
91                std::io::ErrorKind::Unsupported,
92                "Root Actor execution requires a fenced storage writer",
93            ));
94        }
95        let Some(durable) = self.storage.load_root_authority(&session.id).await? else {
96            return Err(root_actor_binding_error(
97                "Root must be persisted before activation",
98            ));
99        };
100        if session.created_at != durable.created_at
101            || session.id != durable.id
102            || session.project_id_meta() != durable.project_id_meta()
103            || session.authority_identity != durable.authority_identity
104        {
105            return Err(root_actor_binding_error(
106                "Root identity changed before activation",
107            ));
108        }
109        session
110            .check_conversation_generation(&durable)
111            .map_err(root_actor_binding_error)?;
112        let lease_duration = std::time::Duration::from_secs(15);
113        let now = chrono::Utc::now();
114        let claimed = directory
115            .claim_activation(&bamboo_domain::ActorActivationClaim {
116                actor_id: session.id.clone(),
117                run_id: run_id.to_string(),
118                lease_owner: format!("root-execution-{}", uuid::Uuid::new_v4()),
119                lease_expires_at: now + chrono::Duration::seconds(15),
120                inbox_generation: session
121                    .session_inbox_admission()
122                    .map_or(0, |a| a.last_admitted_sequence),
123                placement_ref: None,
124                now,
125            })
126            .await
127            .map_err(root_actor_binding_error)?;
128        let owner = bamboo_domain::RootActorRuntimeWrite {
129            fence: claimed.fence(),
130            created_at: durable.created_at,
131            conversation_generation: session.conversation_generation_observation(),
132        };
133        let bound = self.bind_root_response_writer(owner.clone())?;
134        let live_owner = bound
135            .root_actor_owner
136            .as_ref()
137            .expect("bound Root owner")
138            .1
139            .clone();
140        let abandon_directory = directory.clone();
141        let abandon_fence = owner.fence.clone();
142        let binding = bamboo_domain::RootActorExecutionBinding::new(
143            Arc::new(bound),
144            directory.clone(),
145            owner,
146            lease_duration,
147            move || {
148                if let Ok(handle) = tokio::runtime::Handle::try_current() {
149                    handle.spawn(async move {
150                        let _ = abandon_directory
151                            .finish_activation(
152                                &abandon_fence,
153                                chrono::Utc::now(),
154                                bamboo_domain::ActorActivationFinish::Cancelled,
155                            )
156                            .await;
157                    });
158                }
159            },
160        );
161        directory
162            .start_activation(&binding.owner.fence, chrono::Utc::now())
163            .await
164            .map_err(root_actor_binding_error)?;
165        // Keep only a weak route to this exact execution's capability. Host
166        // parent-outcome adapters never reconstruct an owner from observation.
167        self.root_actor_writers
168            .insert(session.id.clone(), Arc::downgrade(&live_owner));
169        Ok(Some(binding))
170    }
171
172    /// Build a response coordinator from the exact reserved execution's
173    /// capability. This does not claim or observe a different Actor owner.
174    pub(crate) fn bind_root_response_writer(
175        &self,
176        owner: bamboo_domain::RootActorRuntimeWrite,
177    ) -> std::io::Result<Self> {
178        let directory = self.root_actor_directory.clone().ok_or_else(|| {
179            root_actor_binding_error("Root response requires its Host Actor directory")
180        })?;
181        let host_cache = self.cache.clone();
182        let publisher: bamboo_domain::RootActorRuntimePublisher = Arc::new(move |saved| {
183            host_cache.insert(
184                saved.id.clone(),
185                Arc::new(crate::SessionSnapshot::new(saved.clone())),
186            );
187        });
188        let bound_store = self
189            .persistence
190            .bind_root_actor_writer(owner.clone(), publisher)?;
191        // Deferred repository callbacks can only update this execution-private
192        // cache. The shared Host cache is published inside V2's physical guard.
193        let mut bound = Self::new(
194            Arc::new(Default::default()),
195            self.storage.clone(),
196            Arc::new(bound_store),
197        );
198        bound.root_actor_owner = Some((directory.clone(), Arc::new(owner)));
199        bound.root_actor_writers = self.root_actor_writers.clone();
200        Ok(bound)
201    }
202
203    /// Enable durable ordinary Root execution ownership in a Host whose
204    /// storage implements the named final writer boundary.
205    pub fn with_root_actor_directory(
206        mut self,
207        directory: Arc<dyn bamboo_domain::ActorDirectoryPort>,
208    ) -> Self {
209        self.root_actor_directory = Some(directory);
210        self
211    }
212
213    #[cfg(test)]
214    fn with_post_durable_hook(mut self, hook: PostDurableHook) -> Self {
215        self.post_durable_hook = Some(hook);
216        self
217    }
218
219    #[cfg(test)]
220    fn run_post_durable_hook(&self, operation: &str, marker: &str) {
221        if let Some(hook) = self.post_durable_hook.as_ref() {
222            hook(operation, marker);
223        }
224    }
225
226    pub fn cache(&self) -> &SessionCache {
227        &self.cache
228    }
229
230    pub fn storage(&self) -> &Arc<dyn Storage> {
231        &self.storage
232    }
233
234    pub fn persistence(&self) -> &Arc<LockedSessionStore> {
235        &self.persistence
236    }
237
238    /// Load a session from the memory cache, falling back to durable storage
239    /// (and back-filling the cache on a storage hit). `None` if absent in both.
240    pub async fn load(&self, session_id: &str) -> Option<Session> {
241        if let Some(session) = read_cached_session(&self.cache, session_id) {
242            return Some(session);
243        }
244
245        let _guard = self.persistence.acquire_lock(session_id).await;
246        if let Some(session) = read_cached_session(&self.cache, session_id) {
247            return Some(session);
248        }
249
250        let loaded = self.storage.load_session(session_id).await;
251        #[cfg(test)]
252        self.run_post_durable_hook("load", session_id);
253        match loaded {
254            Ok(Some(session)) => {
255                self.cache.insert(
256                    session_id.to_string(),
257                    Arc::new(crate::SessionSnapshot::new(session.clone())),
258                );
259                Some(session)
260            }
261            _ => None,
262        }
263    }
264
265    /// Like [`load`](Self::load), but surfaces storage errors instead of
266    /// swallowing them to `None`. Cache hit short-circuits; a storage hit
267    /// back-fills the cache.
268    pub async fn try_load(&self, session_id: &str) -> std::io::Result<Option<Session>> {
269        if let Some(session) = read_cached_session(&self.cache, session_id) {
270            return Ok(Some(session));
271        }
272
273        let _guard = self.persistence.acquire_lock(session_id).await;
274        if let Some(session) = read_cached_session(&self.cache, session_id) {
275            return Ok(Some(session));
276        }
277
278        let loaded = self.storage.load_session(session_id).await?;
279        #[cfg(test)]
280        self.run_post_durable_hook("try_load", session_id);
281        if let Some(ref session) = loaded {
282            self.cache.insert(
283                session_id.to_string(),
284                Arc::new(crate::SessionSnapshot::new(session.clone())),
285            );
286        }
287        Ok(loaded)
288    }
289
290    /// Persist the session (merge-on-write) and refresh the cache, surfacing
291    /// storage errors. Use [`save_and_cache`](Self::save_and_cache) for the
292    /// fire-and-forget variant that logs and continues on failure.
293    pub async fn save(&self, session: &mut Session) -> std::io::Result<()> {
294        self.persistence
295            .merge_save_runtime_and_publish(session, |saved, committed| {
296                if committed {
297                    #[cfg(test)]
298                    self.run_post_durable_hook("save_full", &saved.id);
299                    self.cache.insert(
300                        saved.id.clone(),
301                        Arc::new(crate::SessionSnapshot::new(saved.clone())),
302                    );
303                }
304            })
305            .await
306    }
307
308    /// Persist a blocked approval replay before starting its successor. A
309    /// A save error may arrive after the rewritten result was committed. When
310    /// the save was attempted, evict the old approval from cache under the
311    /// same session lock so the next attempt reloads durable state.
312    pub async fn save_replay_resolution(&self, session: &mut Session) -> std::io::Result<()> {
313        self.persistence
314            .merge_save_runtime_and_publish(session, |saved, committed| {
315                if committed {
316                    self.cache.insert(
317                        saved.id.clone(),
318                        Arc::new(crate::SessionSnapshot::new(saved.clone())),
319                    );
320                } else {
321                    self.cache.remove(&saved.id);
322                }
323            })
324            .await
325    }
326
327    /// Atomically mutate the latest durable runtime session and refresh the
328    /// cache with the saved value. This is the safe path for narrow metadata
329    /// indexes that can be updated concurrently with runner message writes.
330    pub async fn update_runtime_session<F>(
331        &self,
332        session_id: &str,
333        metadata_keys: &[&str],
334        mutate: F,
335    ) -> std::io::Result<Option<Session>>
336    where
337        F: FnOnce(&mut Session),
338    {
339        self.persistence
340            .update_runtime_config_and_publish(session_id, mutate, |saved| {
341                if let Some(cached) = self.cache.get(session_id) {
342                    cached.update(|cached| {
343                        for key in metadata_keys {
344                            if let Some(value) = saved.metadata.get(*key) {
345                                cached.metadata.insert((*key).to_string(), value.clone());
346                            } else {
347                                cached.metadata.remove(*key);
348                            }
349                        }
350                    });
351                }
352            })
353            .await
354    }
355
356    /// Load a session, creating a fresh `Session::new(id, model)` if absent.
357    pub async fn load_or_create(&self, session_id: &str, model: &str) -> Session {
358        if let Some(session) = self.load(session_id).await {
359            return session;
360        }
361        Session::new(session_id.to_string(), model.to_string())
362    }
363
364    /// Load a session, reconciling the memory and storage copies via a
365    /// preference heuristic: a live durable Root birth or mode epoch wins first;
366    /// otherwise storage wins when it is strictly newer, or when it is the same
367    /// age but still carries a pending question memory lost.
368    ///
369    /// The cache is refreshed cache-aside. Timestamps cannot make a stale
370    /// Root birth or mode epoch outrank verified durable authority.
371    pub async fn load_merged_checked(&self, session_id: &str) -> std::io::Result<Option<Session>> {
372        let _guard = self.persistence.acquire_lock(session_id).await;
373        let memory_session = read_cached_session(&self.cache, session_id);
374        let storage_session = self.storage.load_session(session_id).await?;
375        #[cfg(test)]
376        self.run_post_durable_hook("load_merged", session_id);
377
378        Ok(match (memory_session, storage_session) {
379            (Some(memory), Some(storage)) => {
380                let root_authority_preference =
381                    root_authority_storage_preference(&memory, &storage);
382                let prefer_storage = should_prefer_storage(&memory, &storage);
383                let diverged = prefer_storage || memory.messages.len() != storage.messages.len();
384                let chosen_len = if prefer_storage {
385                    storage.messages.len()
386                } else {
387                    memory.messages.len()
388                };
389                macro_rules! merged_log {
390                    ($level:ident) => {
391                        tracing::$level!(
392                            "[{}] load_session_merged: memory={} msgs (updated_at={}), storage={} msgs (updated_at={}), prefer_storage={} -> chose {} msgs",
393                            session_id,
394                            memory.messages.len(),
395                            memory.updated_at,
396                            storage.messages.len(),
397                            storage.updated_at,
398                            prefer_storage,
399                            chosen_len,
400                        )
401                    };
402                }
403                if diverged {
404                    merged_log!(debug);
405                } else {
406                    merged_log!(trace);
407                }
408                let memory_updated_at = memory.updated_at;
409                let chosen = if prefer_storage { storage } else { memory };
410                // Cache-aside refresh with a hard no-regression invariant: only
411                // write back when we actually reconciled *to storage* (a memory
412                // win is already the cached copy; re-inserting it would needlessly
413                // replace a possibly-live Arc) AND the reconciled copy is not
414                // older than what memory already holds, except when durable
415                // storage proved a new Root birth or higher mode epoch. Those
416                // authority changes must survive cross-process clock skew.
417                if prefer_storage
418                    && (root_authority_preference == Some(true)
419                        || chosen.updated_at >= memory_updated_at)
420                {
421                    self.cache.insert(
422                        session_id.to_string(),
423                        Arc::new(crate::SessionSnapshot::new(chosen.clone())),
424                    );
425                }
426                Some(chosen)
427            }
428            (Some(memory), None) => Some(memory),
429            (None, Some(storage)) => {
430                self.cache.insert(
431                    session_id.to_string(),
432                    Arc::new(crate::SessionSnapshot::new(storage.clone())),
433                );
434                Some(storage)
435            }
436            (None, None) => None,
437        })
438    }
439
440    /// Compatibility wrapper for read paths where the historical contract
441    /// treated a storage failure like absence. Mutating/recovery paths should
442    /// use [`Self::load_merged_checked`] so they can preserve retry state.
443    pub async fn load_merged(&self, session_id: &str) -> Option<Session> {
444        match self.load_merged_checked(session_id).await {
445            Ok(session) => session,
446            Err(error) => {
447                tracing::warn!(
448                    "[{}] Failed to load merged session from storage: {}",
449                    session_id,
450                    error
451                );
452                read_cached_session(&self.cache, session_id)
453            }
454        }
455    }
456
457    /// Persist the session (merge-on-write, preserving concurrent UI edits to
458    /// the authoritative metadata group) and refresh the in-memory cache.
459    pub async fn save_and_cache(&self, session: &mut Session) {
460        let result = self
461            .persistence
462            .merge_save_runtime_and_publish(session, |saved, _| {
463                #[cfg(test)]
464                self.run_post_durable_hook("save_and_cache", &saved.id);
465                self.cache.insert(
466                    saved.id.clone(),
467                    Arc::new(crate::SessionSnapshot::new(saved.clone())),
468                );
469            })
470            .await;
471        if let Err(error) = result {
472            tracing::warn!("[{}] Failed to save session: {}", session.id, error);
473        }
474    }
475
476    async fn refresh_cached_task_control_plane(&self, session_id: &str) {
477        match self.storage.load_runtime_control_plane(session_id).await {
478            Ok(Some(durable)) => {
479                if let Some(cached) = self.cache.get(session_id) {
480                    cached.update(|cached| adopt_task_control_plane(cached, &durable));
481                }
482            }
483            Ok(None) => {}
484            Err(error) => tracing::warn!(
485                "[{}] Failed to refresh Task control-plane after write conflict: {}",
486                session_id,
487                error
488            ),
489        }
490    }
491}
492
493fn adopt_task_control_plane(target: &mut Session, durable: &Session) {
494    target.task_list = durable.task_list.clone();
495    target
496        .metadata
497        .remove(bamboo_domain::session::runtime_metadata::keys::TASK_LIST_VERSION);
498    if let Some(runtime_metadata) = target.runtime_metadata.as_mut() {
499        runtime_metadata.task_list_version = None;
500    }
501    if target
502        .runtime_metadata
503        .as_ref()
504        .is_some_and(bamboo_domain::session::SessionRuntimeMetadata::is_empty)
505    {
506        target.runtime_metadata = None;
507    }
508    if let Some(version) = durable.task_list_version_meta() {
509        target.set_task_list_version_meta(version);
510    }
511}
512
513fn should_prefer_storage(memory_session: &Session, storage_session: &Session) -> bool {
514    if let Some(prefer_storage) = root_authority_storage_preference(memory_session, storage_session)
515    {
516        return prefer_storage;
517    }
518    // Never reconcile *backwards* to a strictly-older storage copy: if memory is
519    // newer it is authoritative (e.g. it just answered and cleared a pending
520    // question while storage still holds the stale one). Respecting `updated_at`
521    // here is what stops `load_merged` from returning — and caching — stale data.
522    if storage_session.updated_at < memory_session.updated_at {
523        return false;
524    }
525    // Storage is same-age or newer: prefer it when strictly newer, or when it
526    // still carries a pending question that the (same-age) memory copy lost, so
527    // a genuine clarification is never dropped.
528    storage_session.updated_at > memory_session.updated_at
529        || (memory_session.pending_question.is_none() && storage_session.pending_question.is_some())
530}
531
532/// `load_session` has already checked that the storage copy is live. Its Root
533/// birth therefore outranks any cached incarnation with the same ID, even if
534/// the new lifetime reset the mode epoch or a process clock moved backwards.
535/// Within one birth, only a higher durable epoch is an authority advance.
536fn root_authority_storage_preference(
537    memory_session: &Session,
538    storage_session: &Session,
539) -> Option<bool> {
540    use bamboo_domain::SessionKind;
541
542    if memory_session.kind != SessionKind::Root || storage_session.kind != SessionKind::Root {
543        return None;
544    }
545    if memory_session.created_at != storage_session.created_at
546        || memory_session.authority_identity != storage_session.authority_identity
547    {
548        return Some(true);
549    }
550    (memory_session.root_mode_transition_epoch != storage_session.root_mode_transition_epoch)
551        .then_some(
552            storage_session.root_mode_transition_epoch > memory_session.root_mode_transition_epoch,
553        )
554}
555
556/// `SessionRepository` is the canonical `RuntimeSessionPersistence`: the runtime
557/// can persist a session through the same coordinator (merge-on-write + cache
558/// refresh) instead of a bespoke adapter.
559#[async_trait::async_trait]
560impl bamboo_domain::RuntimeSessionPersistence for SessionRepository {
561    fn inherited_child_wait(&self) -> Option<bamboo_domain::InheritedChildWait> {
562        bamboo_domain::RuntimeSessionPersistence::inherited_child_wait(self.persistence.as_ref())
563    }
564
565    fn bind_inherited_child_wait(
566        &self,
567        inherited: bamboo_domain::InheritedChildWait,
568    ) -> std::io::Result<Arc<dyn bamboo_domain::RuntimeSessionPersistence>> {
569        let mut bound = self.clone();
570        bound.persistence = Arc::new(
571            self.persistence
572                .as_ref()
573                .bind_inherited_child_wait(inherited),
574        );
575        Ok(Arc::new(bound))
576    }
577
578    fn root_actor_writer(&self) -> Option<bamboo_domain::RootActorRuntimeWrite> {
579        self.root_actor_owner
580            .as_ref()
581            .map(|(_, owner)| (**owner).clone())
582    }
583
584    fn bind_root_inbox_receipt(
585        &self,
586        receipt: bamboo_domain::SessionInboxReceipt,
587    ) -> std::io::Result<Arc<dyn bamboo_domain::RuntimeSessionPersistence>> {
588        if self.root_actor_owner.is_none()
589            || receipt.generation == 0
590            || self
591                .root_inbox_receipt
592                .as_ref()
593                .is_some_and(|current| current != &receipt)
594        {
595            return Err(root_actor_binding_error(
596                "Root receipt requires its unchanged execution writer",
597            ));
598        }
599        let mut bound = self.clone();
600        bound.root_inbox_receipt = Some(receipt);
601        Ok(Arc::new(bound))
602    }
603
604    fn root_inbox_receipt(&self) -> Option<&bamboo_domain::SessionInboxReceipt> {
605        self.root_inbox_receipt.as_ref()
606    }
607
608    fn root_actor_execution_required(&self, session: &Session) -> bool {
609        self.root_actor_directory.is_some()
610            && session.kind == bamboo_domain::SessionKind::Root
611            && session.authority_identity.is_ordinary()
612    }
613
614    async fn bind_root_actor_execution(
615        &self,
616        session: &Session,
617        run_id: &str,
618    ) -> std::io::Result<Option<bamboo_domain::RootActorExecutionBinding>> {
619        let repository = self.clone();
620        let session = session.clone();
621        let run_id = run_id.to_string();
622        // The complete claim/start handoff owns its state if the adapter stops
623        // waiting. An unobserved returned binding drops its exact-fence cleanup.
624        tokio::spawn(async move {
625            repository
626                .bind_root_actor_execution_inner(&session, &run_id)
627                .await
628        })
629        .await
630        .map_err(|error| root_actor_binding_error(format!("Root binding task failed: {error}")))?
631    }
632
633    async fn save_runtime_session(&self, session: &mut Session) -> std::io::Result<()> {
634        // Runtime authorization reads through this same cache. Refresh it even
635        // when durable storage fails so a current activation can never observe
636        // a previous run's skill allowlist. The error is still returned to the
637        // caller and durable state remains unchanged.
638        self.persistence
639            .merge_save_runtime_and_publish(session, |saved, _| {
640                #[cfg(test)]
641                self.run_post_durable_hook("save_runtime_session", &saved.id);
642                self.cache.insert(
643                    saved.id.clone(),
644                    Arc::new(crate::SessionSnapshot::new(saved.clone())),
645                );
646            })
647            .await
648    }
649
650    async fn save_finalized_runtime_session(&self, session: &mut Session) -> std::io::Result<()> {
651        self.persistence
652            .merge_save_finalized_runtime_and_publish(session, |saved, _| {
653                #[cfg(test)]
654                self.run_post_durable_hook("save_finalized_runtime_session", &saved.id);
655                self.cache.insert(
656                    saved.id.clone(),
657                    Arc::new(crate::SessionSnapshot::new(saved.clone())),
658                );
659            })
660            .await
661    }
662
663    async fn save_finalized_runtime_with_inherited_child_wait(
664        &self,
665        session: &mut Session,
666        inherited: &bamboo_domain::session::runtime_state::WaitingForChildrenState,
667    ) -> std::io::Result<()> {
668        self.persistence
669            .merge_save_inherited_child_wait_and_publish(session, inherited, |saved, _| {
670                self.cache.insert(
671                    saved.id.clone(),
672                    Arc::new(crate::SessionSnapshot::new(saved.clone())),
673                );
674            })
675            .await
676    }
677
678    async fn seed_runtime_activation(&self, session: &mut Session) -> std::io::Result<()> {
679        self.persistence
680            .seed_runtime_activation_and_publish(session, |saved, committed| {
681                #[cfg(test)]
682                self.run_post_durable_hook("seed_runtime_activation", &saved.id);
683                if committed {
684                    self.cache.insert(
685                        saved.id.clone(),
686                        Arc::new(crate::SessionSnapshot::new(saved.clone())),
687                    );
688                }
689            })
690            .await
691    }
692
693    async fn record_permission_posture_activation(
694        &self,
695        session_id: &str,
696        expected_audit_revision: Option<u64>,
697        seed: &bamboo_domain::PermissionAuditSeed,
698    ) -> std::io::Result<Option<Session>> {
699        self.persistence
700            .record_permission_posture_activation_and_publish(
701                session_id,
702                expected_audit_revision,
703                seed,
704                |saved| {
705                    #[cfg(test)]
706                    self.run_post_durable_hook("permission_posture_activation", &saved.id);
707                    self.cache.insert(
708                        saved.id.clone(),
709                        Arc::new(crate::SessionSnapshot::new(saved.clone())),
710                    );
711                },
712            )
713            .await
714    }
715
716    async fn save_runtime_control_plane(&self, session: &mut Session) -> std::io::Result<()> {
717        self.persistence
718            .save_runtime_only_and_publish(session, |saved| {
719                #[cfg(test)]
720                self.run_post_durable_hook("save", &saved.id);
721
722                // A control-plane snapshot may intentionally carry no messages
723                // (for example, child Task synchronization loads the root's
724                // runtime sidecar). Publish its fresh runtime fields without
725                // replacing a cache-resident transcript with that empty
726                // snapshot. SessionInbox admission is coupled to transcript
727                // persistence and is therefore preserved alongside the cached
728                // messages, matching the V2 sidecar overlay contract.
729                if let Some(cached) = self.cache.get(&saved.id) {
730                    cached.update(|cached| {
731                        let messages = cached.messages.clone();
732                        let provider_transcript = cached.provider_transcript.clone();
733                        let admission = cached
734                            .runtime_metadata
735                            .as_ref()
736                            .and_then(|metadata| metadata.session_inbox_admission.clone());
737                        let mut refreshed = saved.clone();
738                        refreshed.messages = messages;
739                        refreshed.provider_transcript = provider_transcript;
740                        if let Some(admission) = admission {
741                            refreshed
742                                .runtime_metadata
743                                .get_or_insert_with(Default::default)
744                                .session_inbox_admission = Some(admission);
745                        } else if let Some(metadata) = refreshed.runtime_metadata.as_mut() {
746                            metadata.session_inbox_admission = None;
747                        }
748                        *cached = refreshed;
749                    });
750                }
751            })
752            .await
753    }
754
755    async fn load_runtime_control_plane(
756        &self,
757        session_id: &str,
758    ) -> std::io::Result<Option<Session>> {
759        bamboo_domain::RuntimeSessionPersistence::load_runtime_control_plane(
760            self.persistence.as_ref(),
761            session_id,
762        )
763        .await
764    }
765
766    async fn update_task_list_control_plane(
767        &self,
768        session_id: &str,
769        task_list: &bamboo_domain::TaskList,
770        version: &str,
771    ) -> std::io::Result<bool> {
772        let result = self
773            .persistence
774            .update_task_list_control_plane_and_publish(session_id, task_list, version, |_| {
775                #[cfg(test)]
776                self.run_post_durable_hook("task", version);
777
778                // The durable transaction changed only Task-owned fields.
779                // Mirror that same narrow patch into the cache so a
780                // concurrent round/status/child transition already present
781                // in memory cannot be replaced by a stale whole-control-
782                // plane snapshot.
783                if let Some(cached) = self.cache.get(session_id) {
784                    cached.update(|cached| {
785                        cached.set_task_list(task_list.clone());
786                        cached.set_task_list_version_meta(version.to_string());
787                    });
788                }
789            })
790            .await;
791        if result
792            .as_ref()
793            .is_err_and(|error| error.kind() == std::io::ErrorKind::WouldBlock)
794        {
795            // The storage final-CAS rejected this unconditional patch because
796            // an independent writer won. Refresh the winner into the cache but
797            // preserve the conflict result so Taskwrite never mistakes it for
798            // the Ok(false) legacy-persistence fallback signal.
799            self.refresh_cached_task_control_plane(session_id).await;
800        }
801        result
802    }
803
804    async fn update_task_list_control_plane_if_version(
805        &self,
806        session_id: &str,
807        expected_version: &str,
808        expected_task_list: &bamboo_domain::TaskList,
809        task_list: &bamboo_domain::TaskList,
810        version: &str,
811    ) -> std::io::Result<bool> {
812        self.persistence
813            .update_task_list_control_plane_if_version_and_publish(
814                session_id,
815                expected_version,
816                expected_task_list,
817                task_list,
818                version,
819                |_| {
820                    if let Some(cached) = self.cache.get(session_id) {
821                        cached.update(|cached| {
822                            cached.set_task_list(task_list.clone());
823                            cached.set_task_list_version_meta(version.to_string());
824                        });
825                    }
826                },
827            )
828            .await
829    }
830
831    async fn update_task_list_control_planes_if_version(
832        &self,
833        session_id: &str,
834        shared_session_id: &str,
835        expected_version: &str,
836        expected_task_list: &bamboo_domain::TaskList,
837        task_list: &bamboo_domain::TaskList,
838        version: &str,
839    ) -> std::io::Result<bool> {
840        self.persistence
841            .update_task_list_control_planes_if_version_and_publish(
842                session_id,
843                shared_session_id,
844                expected_version,
845                expected_task_list,
846                task_list,
847                version,
848                |_, _| {
849                    for id in [session_id, shared_session_id] {
850                        if let Some(cached) = self.cache.get(id) {
851                            cached.update(|cached| {
852                                cached.set_task_list(task_list.clone());
853                                cached.set_task_list_version_meta(version.to_string());
854                            });
855                        }
856                    }
857                },
858            )
859            .await
860    }
861
862    async fn checkpoint_runtime_session(&self, session: &mut Session) -> std::io::Result<()> {
863        // The execute-boundary checkpoint uses LockedSessionStore's atomic
864        // append-safe transcript merge, then publishes that reconciled snapshot
865        // to the runtime cache.  On failure, leave the existing cache alone: the
866        // checkpoint may have failed to load the latest durable transcript, and
867        // replacing a fresher cache entry with the stale runner snapshot would
868        // set up a later SHRINK write.
869        self.persistence
870            .checkpoint_runtime_session_and_publish(session, |saved, committed| {
871                #[cfg(test)]
872                self.run_post_durable_hook("checkpoint", &saved.id);
873                if committed {
874                    self.cache.insert(
875                        saved.id.clone(),
876                        Arc::new(crate::SessionSnapshot::new(saved.clone())),
877                    );
878                }
879            })
880            .await
881    }
882
883    async fn claim_scoped_loop_controls(
884        &self,
885        session: &Session,
886        inbox: Arc<dyn bamboo_domain::SessionInboxPort>,
887        active_run_id: &str,
888        limit: usize,
889    ) -> std::io::Result<Vec<bamboo_domain::ScopedLoopControlClaim>> {
890        use bamboo_domain::ScopedLoopControlClaim;
891        if let Some((directory, owner)) = self.root_actor_owner.as_ref() {
892            let inbox = self.storage.bind_root_actor_inbox(owner, inbox)?;
893            let consumer = bamboo_domain::SessionInboxConsumerId::new();
894            let request = root_inbox::lease_request(directory.as_ref(), owner, &consumer).await?;
895            return inbox
896                .claim_owned_scoped_loop_controls(&session.id, limit, Some(active_run_id), &request)
897                .await
898                .map(|claims| {
899                    claims
900                        .into_iter()
901                        .map(ScopedLoopControlClaim::Owned)
902                        .collect()
903                })
904                .map_err(std::io::Error::other);
905        }
906        inbox
907            .claim_scoped_loop_controls(&session.id, limit, Some(active_run_id))
908            .await
909            .map(|claims| {
910                claims
911                    .into_iter()
912                    .map(ScopedLoopControlClaim::Legacy)
913                    .collect()
914            })
915            .map_err(std::io::Error::other)
916    }
917
918    async fn checkpoint_scoped_loop_control(
919        &self,
920        session: &mut Session,
921        bindings: &[bamboo_domain::ScopedLoopControlBinding],
922        activation: Arc<dyn bamboo_domain::SessionActivationPort>,
923        input: Option<(
924            Arc<dyn bamboo_domain::SessionInboxPort>,
925            bamboo_domain::ScopedLoopControlClaim,
926        )>,
927    ) -> std::io::Result<()> {
928        let _guard = self.persistence.acquire_lock(&session.id).await;
929        let cache = self.cache.clone();
930        let publish: bamboo_domain::RootActorRuntimePublisher = Arc::new(move |saved| {
931            cache.insert(
932                saved.id.clone(),
933                Arc::new(crate::SessionSnapshot::new(saved.clone())),
934            );
935        });
936        self.storage
937            .checkpoint_scoped_loop(bamboo_domain::ScopedLoopControlCheckpoint {
938                session,
939                bindings,
940                activation,
941                owner: self
942                    .root_actor_owner
943                    .as_ref()
944                    .map(|(_, owner)| owner.as_ref()),
945                input,
946                publish,
947            })
948            .await
949    }
950
951    async fn load_loop_control_checkpoint(
952        &self,
953        session_id: &str,
954    ) -> std::io::Result<Option<Session>> {
955        self.storage
956            .load_execution_admission_snapshot(session_id)
957            .await
958    }
959
960    async fn admit_root_inbox(
961        &self,
962        session: &mut Session,
963        inbox: Arc<dyn bamboo_domain::SessionInboxPort>,
964        active_run_id: Option<&str>,
965    ) -> std::io::Result<Option<bamboo_domain::RootInboxAdmission>> {
966        self.admit_owned_root_inbox(session, inbox, active_run_id)
967            .await
968    }
969
970    async fn checkpoint_retrieval_window(
971        &self,
972        expected_base: &Session,
973        staged: &mut Session,
974    ) -> std::io::Result<RetrievalWindowCheckpointOutcome> {
975        self.persistence
976            .checkpoint_retrieval_window_and_publish(expected_base, staged, |saved| {
977                #[cfg(test)]
978                self.run_post_durable_hook("checkpoint_retrieval_window", &saved.id);
979                self.cache.insert(
980                    saved.id.clone(),
981                    Arc::new(crate::SessionSnapshot::new(saved.clone())),
982                );
983            })
984            .await
985    }
986
987    async fn checkpoint_prompt_rewrite(
988        &self,
989        expected_base: &Session,
990        staged: &mut Session,
991    ) -> std::io::Result<RetrievalWindowCheckpointOutcome> {
992        self.persistence
993            .checkpoint_prompt_rewrite_and_publish(expected_base, staged, |saved| {
994                #[cfg(test)]
995                self.run_post_durable_hook("checkpoint_prompt_rewrite", &saved.id);
996                self.cache.insert(
997                    saved.id.clone(),
998                    Arc::new(crate::SessionSnapshot::new(saved.clone())),
999                );
1000            })
1001            .await
1002    }
1003
1004    async fn checkpoint_manual_archive_rejection(
1005        &self,
1006        expected_base: &Session,
1007        staged: &mut Session,
1008    ) -> std::io::Result<RetrievalWindowCheckpointOutcome> {
1009        self.persistence
1010            .checkpoint_manual_archive_rejection_and_publish(expected_base, staged, |saved| {
1011                #[cfg(test)]
1012                self.run_post_durable_hook("checkpoint_manual_archive_rejection", &saved.id);
1013                self.cache.insert(
1014                    saved.id.clone(),
1015                    Arc::new(crate::SessionSnapshot::new(saved.clone())),
1016                );
1017            })
1018            .await
1019    }
1020
1021    async fn checkpoint_manual_archive_consumption(
1022        &self,
1023        expected_base: &Session,
1024        staged: &mut Session,
1025    ) -> std::io::Result<RetrievalWindowCheckpointOutcome> {
1026        self.persistence
1027            .checkpoint_manual_archive_consumption_and_publish(expected_base, staged, |saved| {
1028                #[cfg(test)]
1029                self.run_post_durable_hook("checkpoint_manual_archive_consumption", &saved.id);
1030                self.cache.insert(
1031                    saved.id.clone(),
1032                    Arc::new(crate::SessionSnapshot::new(saved.clone())),
1033                );
1034            })
1035            .await
1036    }
1037
1038    async fn load_runtime_session(&self, session_id: &str) -> std::io::Result<Option<Session>> {
1039        if let Some((directory, owner)) = self.root_actor_owner.as_ref() {
1040            if session_id != owner.fence.actor_id {
1041                return Err(root_actor_binding_error(
1042                    "bound execution cannot load another Root",
1043                ));
1044            }
1045            directory
1046                .validate_fence(&owner.fence, chrono::Utc::now())
1047                .await
1048                .map_err(root_actor_binding_error)?;
1049            let loaded = self.storage.load_session(session_id).await?;
1050            let Some(snapshot) = loaded.as_ref() else {
1051                return Err(root_actor_binding_error("bound Root no longer exists"));
1052            };
1053            validate_bound_root_snapshot(owner, snapshot)?;
1054            directory
1055                .validate_fence(&owner.fence, chrono::Utc::now())
1056                .await
1057                .map_err(root_actor_binding_error)?;
1058            return Ok(loaded);
1059        }
1060        self.try_load(session_id).await
1061    }
1062
1063    async fn load_runtime_admission_snapshot(
1064        &self,
1065        session_id: &str,
1066    ) -> std::io::Result<Option<Session>> {
1067        if let Some((directory, owner)) = self.root_actor_owner.as_ref() {
1068            if session_id != owner.fence.actor_id {
1069                return Err(root_actor_binding_error(
1070                    "bound execution cannot admit another Root",
1071                ));
1072            }
1073            directory
1074                .validate_fence(&owner.fence, chrono::Utc::now())
1075                .await
1076                .map_err(root_actor_binding_error)?;
1077            let loaded = self
1078                .storage
1079                .load_execution_admission_snapshot(session_id)
1080                .await?;
1081            let Some(snapshot) = loaded.as_ref() else {
1082                return Err(root_actor_binding_error("bound Root no longer exists"));
1083            };
1084            validate_bound_root_snapshot(owner, snapshot)?;
1085            // A takeover or delete/rebirth during the full durable read cannot
1086            // release even a previously materialized snapshot to this owner.
1087            directory
1088                .validate_fence(&owner.fence, chrono::Utc::now())
1089                .await
1090                .map_err(root_actor_binding_error)?;
1091            return Ok(loaded);
1092        }
1093        // Runtime/tool reads remain cache-first. An execution admission must
1094        // observe durable facts without publishing a stale or rejected clone.
1095        bamboo_domain::RuntimeSessionPersistence::load_runtime_admission_snapshot(
1096            self.persistence.as_ref(),
1097            session_id,
1098        )
1099        .await
1100    }
1101
1102    async fn clear_legacy_pending_messages(
1103        &self,
1104        session_id: &str,
1105        expected: &[serde_json::Value],
1106    ) -> std::io::Result<bool> {
1107        // Do not use the trait default here: `load_runtime_session` is
1108        // cache-first, so a stale cache could erase a message concurrently
1109        // appended to the durable legacy queue. Delegate the compare-and-clear
1110        // to LockedSessionStore's single per-session critical section.
1111        self.persistence
1112            .clear_legacy_pending_messages_and_publish(session_id, expected, |latest| {
1113                #[cfg(test)]
1114                self.run_post_durable_hook("clear_legacy", session_id);
1115                self.cache.insert(
1116                    session_id.to_string(),
1117                    Arc::new(crate::SessionSnapshot::new(latest.clone())),
1118                );
1119            })
1120            .await
1121    }
1122
1123    async fn append_token_usage_record(
1124        &self,
1125        session_id: &str,
1126        json_line: &str,
1127    ) -> std::io::Result<()> {
1128        self.storage
1129            .append_token_usage_record(session_id, json_line)
1130            .await
1131    }
1132}
1133
1134#[cfg(test)]
1135mod tests {
1136    use super::*;
1137    use bamboo_agent_core::storage::Storage;
1138    use chrono::Utc;
1139    use std::collections::HashMap;
1140    use std::sync::atomic::{AtomicBool, Ordering};
1141    use std::sync::{Condvar, Mutex};
1142    use std::time::Duration;
1143
1144    #[derive(Default)]
1145    struct MapStorage {
1146        sessions: Mutex<HashMap<String, Session>>,
1147        fail_pair_commit: AtomicBool,
1148    }
1149
1150    struct FailingSaveStorage {
1151        persisted: Mutex<Option<Session>>,
1152    }
1153
1154    struct AmbiguousReplaySaveStorage {
1155        persisted: Mutex<Option<Session>>,
1156    }
1157
1158    #[async_trait::async_trait]
1159    impl Storage for MapStorage {
1160        async fn save_session(&self, session: &Session) -> std::io::Result<()> {
1161            self.sessions
1162                .lock()
1163                .unwrap()
1164                .insert(session.id.clone(), session.clone());
1165            Ok(())
1166        }
1167        async fn load_session(&self, session_id: &str) -> std::io::Result<Option<Session>> {
1168            Ok(self.sessions.lock().unwrap().get(session_id).cloned())
1169        }
1170        async fn delete_session(&self, session_id: &str) -> std::io::Result<bool> {
1171            Ok(self.sessions.lock().unwrap().remove(session_id).is_some())
1172        }
1173
1174        async fn save_task_control_plane_if_matches(
1175            &self,
1176            original: &Session,
1177            updated: &Session,
1178        ) -> std::io::Result<bool> {
1179            let mut sessions = self.sessions.lock().unwrap();
1180            let matches_original = sessions.get(&original.id).is_some_and(|current| {
1181                current.task_list_version_meta() == original.task_list_version_meta()
1182                    && serde_json::to_value(&current.task_list).ok()
1183                        == serde_json::to_value(&original.task_list).ok()
1184            });
1185            if !matches_original {
1186                return Ok(false);
1187            }
1188            let Some(current) = sessions.get(&original.id).cloned() else {
1189                return Ok(false);
1190            };
1191            let mut committed = current;
1192            committed.task_list = updated.task_list.clone();
1193            if let Some(version) = updated.task_list_version_meta() {
1194                committed.set_task_list_version_meta(version);
1195            }
1196            sessions.insert(committed.id.clone(), committed);
1197            Ok(true)
1198        }
1199
1200        async fn save_task_control_planes_atomically(
1201            &self,
1202            first_original: &Session,
1203            first_updated: &Session,
1204            second_original: &Session,
1205            second_updated: &Session,
1206        ) -> std::io::Result<bool> {
1207            if self.fail_pair_commit.swap(false, Ordering::SeqCst) {
1208                return Err(std::io::Error::other(
1209                    "injected paired Task transaction failure",
1210                ));
1211            }
1212            let mut sessions = self.sessions.lock().unwrap();
1213            let matches_original = |current: Option<&Session>, original: &Session| {
1214                current.is_some_and(|current| {
1215                    current.task_list_version_meta() == original.task_list_version_meta()
1216                        && serde_json::to_value(&current.task_list).ok()
1217                            == serde_json::to_value(&original.task_list).ok()
1218                })
1219            };
1220            if !matches_original(sessions.get(&first_original.id), first_original)
1221                || !matches_original(sessions.get(&second_original.id), second_original)
1222            {
1223                return Ok(false);
1224            }
1225            sessions.insert(first_updated.id.clone(), first_updated.clone());
1226            sessions.insert(second_updated.id.clone(), second_updated.clone());
1227            Ok(true)
1228        }
1229    }
1230
1231    #[async_trait::async_trait]
1232    impl Storage for FailingSaveStorage {
1233        async fn save_session(&self, _session: &Session) -> std::io::Result<()> {
1234            Err(std::io::Error::other("injected save failure"))
1235        }
1236
1237        async fn load_session(&self, _session_id: &str) -> std::io::Result<Option<Session>> {
1238            Ok(self.persisted.lock().unwrap().clone())
1239        }
1240
1241        async fn delete_session(&self, _session_id: &str) -> std::io::Result<bool> {
1242            Ok(false)
1243        }
1244    }
1245
1246    #[async_trait::async_trait]
1247    impl Storage for AmbiguousReplaySaveStorage {
1248        async fn save_session(&self, session: &Session) -> std::io::Result<()> {
1249            *self.persisted.lock().unwrap() = Some(session.clone());
1250            Err(std::io::Error::other(
1251                "injected post-commit replay save error",
1252            ))
1253        }
1254
1255        async fn load_session(&self, _session_id: &str) -> std::io::Result<Option<Session>> {
1256            Ok(self.persisted.lock().unwrap().clone())
1257        }
1258
1259        async fn delete_session(&self, _session_id: &str) -> std::io::Result<bool> {
1260            Ok(false)
1261        }
1262    }
1263
1264    fn test_repo(storage: Arc<dyn Storage>) -> SessionRepository {
1265        let cache: SessionCache = Arc::default();
1266        let persistence = Arc::new(LockedSessionStore::new(storage.clone()));
1267        SessionRepository::new(cache, storage, persistence)
1268    }
1269
1270    fn cache_put(repo: &SessionRepository, session: &Session) {
1271        repo.cache().insert(
1272            session.id.clone(),
1273            Arc::new(crate::SessionSnapshot::new(session.clone())),
1274        );
1275    }
1276
1277    #[tokio::test]
1278    async fn merged_root_prefers_live_recreated_birth_even_when_epoch_resets() {
1279        let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
1280        let repo = test_repo(storage.clone());
1281        let id = "root-recreated-same-id";
1282        let mut stale = Session::new(id, "stale");
1283        stale.root_mode_transition_epoch = 3;
1284        stale.updated_at += chrono::Duration::hours(1);
1285        cache_put(&repo, &stale);
1286
1287        let mut recreated = Session::new(id, "recreated");
1288        recreated.created_at = stale.created_at + chrono::Duration::seconds(1);
1289        recreated.updated_at = stale.created_at + chrono::Duration::seconds(1);
1290        storage.save_session(&recreated).await.unwrap();
1291
1292        let loaded = repo.load_merged_checked(id).await.unwrap().unwrap();
1293        assert_eq!(loaded.model, "recreated");
1294        assert_eq!(loaded.root_mode_transition_epoch, 0);
1295        assert_eq!(
1296            loaded.root_mode_birth_token(),
1297            recreated.root_mode_birth_token()
1298        );
1299        let cached = read_cached_session(repo.cache(), id).expect("new birth cached");
1300        assert_eq!(cached.model, "recreated");
1301        assert_eq!(
1302            cached.root_mode_birth_token(),
1303            recreated.root_mode_birth_token()
1304        );
1305    }
1306
1307    #[tokio::test]
1308    async fn merged_root_caches_higher_durable_epoch_despite_clock_skew() {
1309        let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
1310        let repo = test_repo(storage.clone());
1311        let id = "root-mode-clock-skew";
1312        let mut stale = Session::new(id, "model");
1313        stale.updated_at += chrono::Duration::hours(1);
1314        cache_put(&repo, &stale);
1315
1316        let mut durable = stale.clone();
1317        durable.root_mode_transition_epoch = 1;
1318        durable.root_orchestration_only = true;
1319        durable.root_tool_authority_revision = 1;
1320        durable.updated_at = durable.created_at;
1321        storage.save_session(&durable).await.unwrap();
1322
1323        let loaded = repo.load_merged_checked(id).await.unwrap().unwrap();
1324        assert_eq!(loaded.root_mode_transition_epoch, 1);
1325        assert!(loaded.root_orchestration_only);
1326        let cached = read_cached_session(repo.cache(), id).expect("higher epoch cached");
1327        assert_eq!(cached.root_mode_transition_epoch, 1);
1328        assert!(cached.root_orchestration_only);
1329    }
1330
1331    #[tokio::test]
1332    async fn runtime_merge_save_adopts_root_mode_terminal_tuple_even_for_policy_neutral_fence() {
1333        use bamboo_domain::{RootModeOperationAction, RootModeOperationRequest};
1334
1335        for action in [
1336            RootModeOperationAction::Select,
1337            RootModeOperationAction::Recover,
1338        ] {
1339            let home = tempfile::tempdir().unwrap();
1340            let concrete = Arc::new(
1341                bamboo_storage::SessionStoreV2::new(home.path().into())
1342                    .await
1343                    .unwrap(),
1344            );
1345            let storage: Arc<dyn Storage> = concrete.clone();
1346            let repo = test_repo(storage.clone());
1347            let mut running = Session::new("runtime-root-mode-merge", "model");
1348            storage.save_session(&running).await.unwrap();
1349            cache_put(&repo, &running);
1350            let request = RootModeOperationRequest {
1351                session_id: running.id.clone(),
1352                operation_id: format!("0:{}", uuid::Uuid::new_v4()),
1353                birth_token: running.root_mode_birth_token(),
1354                expected_epoch: 0,
1355                requested_enabled: true,
1356                action,
1357            };
1358            storage.root_mode_operation(&request).await.unwrap();
1359
1360            // This save can race a host mode operation before the next runtime
1361            // boundary. The production merge must preserve that durable tuple.
1362            running.add_message(bamboo_agent_core::Message::assistant("run result", None));
1363            repo.save(&mut running).await.unwrap();
1364            assert_eq!(running.root_mode_transition_epoch, 1);
1365            assert_eq!(running.root_mode_operations.len(), 1);
1366            assert_eq!(
1367                running.root_orchestration_only_enabled(),
1368                action == RootModeOperationAction::Select
1369            );
1370            let durable = storage.load_session(&running.id).await.unwrap().unwrap();
1371            assert_eq!(durable.root_mode_transition_epoch, 1);
1372            assert_eq!(durable.root_mode_operations, running.root_mode_operations);
1373            assert_eq!(durable.messages.last().unwrap().content, "run result");
1374            concrete.flush_search_index().await;
1375        }
1376    }
1377
1378    fn task_list(session_id: &str, title: &str) -> bamboo_domain::TaskList {
1379        let now = Utc::now();
1380        bamboo_domain::TaskList {
1381            session_id: session_id.to_string(),
1382            title: title.to_string(),
1383            items: Vec::new(),
1384            created_at: now,
1385            updated_at: now,
1386        }
1387    }
1388
1389    fn staged_retrieval_window_archive(expected: &Session) -> Session {
1390        let mut staged = expected.clone();
1391        let mut event = bamboo_domain::CompressionEvent::new(
1392            1,
1393            1,
1394            80.0,
1395            60.0,
1396            0,
1397            bamboo_domain::CompressionTriggerType::Auto,
1398            0.0,
1399            None,
1400            0,
1401        );
1402        event.kind = bamboo_domain::CompressionEventKind::RetrievalWindow;
1403        staged.messages[0].compressed = true;
1404        staged.messages[0].compressed_by_event_id = Some(event.id.clone());
1405        staged.compression_events.push(event);
1406        staged.reset_model_context_epoch(bamboo_domain::ModelContextResetReason::Compression);
1407        staged
1408    }
1409
1410    #[tokio::test]
1411    async fn paired_task_cas_narrowly_updates_child_and_root_cache() {
1412        let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
1413        let repo = test_repo(storage.clone());
1414        let root_id = "paired-cache-root";
1415        let child_id = "paired-cache-child";
1416        let expected_task_list = task_list(root_id, "old shared");
1417
1418        let mut root = Session::new(root_id, "model");
1419        root.add_message(bamboo_agent_core::Message::user("root transcript"));
1420        root.metadata
1421            .insert("unrelated.root".to_string(), "keep".to_string());
1422        root.set_task_list(expected_task_list.clone());
1423        root.set_task_list_version_meta("1");
1424        storage.save_session(&root).await.unwrap();
1425        cache_put(&repo, &root);
1426
1427        let mut child = Session::new_child(child_id, root_id, "model", "child");
1428        child.add_message(bamboo_agent_core::Message::user("child transcript"));
1429        child
1430            .metadata
1431            .insert("unrelated.child".to_string(), "keep".to_string());
1432        child.set_task_list(expected_task_list.clone());
1433        child.set_task_list_version_meta("1");
1434        storage.save_session(&child).await.unwrap();
1435        cache_put(&repo, &child);
1436
1437        assert!(
1438            bamboo_domain::RuntimeSessionPersistence::update_task_list_control_planes_if_version(
1439                &repo,
1440                child_id,
1441                root_id,
1442                "1",
1443                &expected_task_list,
1444                &task_list(root_id, "evaluated"),
1445                "2",
1446            )
1447            .await
1448            .expect("paired cache CAS succeeds")
1449        );
1450
1451        let cached_root = read_cached_session(repo.cache(), root_id).expect("cached root");
1452        let cached_child = read_cached_session(repo.cache(), child_id).expect("cached child");
1453        for (session, transcript, metadata_key) in [
1454            (&cached_root, "root transcript", "unrelated.root"),
1455            (&cached_child, "child transcript", "unrelated.child"),
1456        ] {
1457            assert_eq!(session.task_list_version_meta().as_deref(), Some("2"));
1458            assert_eq!(
1459                session.task_list.as_ref().map(|list| list.title.as_str()),
1460                Some("evaluated")
1461            );
1462            assert_eq!(session.messages.len(), 1);
1463            assert_eq!(session.messages[0].content, transcript);
1464            assert_eq!(
1465                session.metadata.get(metadata_key).map(String::as_str),
1466                Some("keep")
1467            );
1468        }
1469    }
1470
1471    #[tokio::test]
1472    async fn failed_paired_task_transaction_does_not_publish_child_or_root_cache() {
1473        let concrete = Arc::new(MapStorage::default());
1474        let storage: Arc<dyn Storage> = concrete.clone();
1475        let repo = test_repo(storage.clone());
1476        let root_id = "paired-cache-failure-root";
1477        let child_id = "paired-cache-failure-child";
1478        let expected_task_list = task_list(root_id, "old shared");
1479
1480        let mut root = Session::new(root_id, "model");
1481        root.add_message(bamboo_agent_core::Message::user("root transcript"));
1482        root.metadata
1483            .insert("unrelated.root".to_string(), "keep".to_string());
1484        root.set_task_list(expected_task_list.clone());
1485        root.set_task_list_version_meta("1");
1486        storage.save_session(&root).await.unwrap();
1487        cache_put(&repo, &root);
1488
1489        let mut child = Session::new_child(child_id, root_id, "model", "child");
1490        child.add_message(bamboo_agent_core::Message::user("child transcript"));
1491        child
1492            .metadata
1493            .insert("unrelated.child".to_string(), "keep".to_string());
1494        child.set_task_list(expected_task_list.clone());
1495        child.set_task_list_version_meta("1");
1496        storage.save_session(&child).await.unwrap();
1497        cache_put(&repo, &child);
1498
1499        concrete.fail_pair_commit.store(true, Ordering::SeqCst);
1500        let error =
1501            bamboo_domain::RuntimeSessionPersistence::update_task_list_control_planes_if_version(
1502                &repo,
1503                child_id,
1504                root_id,
1505                "1",
1506                &expected_task_list,
1507                &task_list(root_id, "must not publish"),
1508                "2",
1509            )
1510            .await
1511            .expect_err("paired durable transaction fails");
1512        assert!(error.to_string().contains("injected paired"));
1513
1514        for (id, expected_title, transcript, metadata_key) in [
1515            (root_id, "old shared", "root transcript", "unrelated.root"),
1516            (
1517                child_id,
1518                "old shared",
1519                "child transcript",
1520                "unrelated.child",
1521            ),
1522        ] {
1523            let cached = read_cached_session(repo.cache(), id).expect("cached session remains");
1524            assert_eq!(cached.task_list_version_meta().as_deref(), Some("1"));
1525            assert_eq!(
1526                cached.task_list.as_ref().map(|list| list.title.as_str()),
1527                Some(expected_title)
1528            );
1529            assert_eq!(cached.messages[0].content, transcript);
1530            assert_eq!(
1531                cached.metadata.get(metadata_key).map(String::as_str),
1532                Some("keep")
1533            );
1534
1535            let durable = storage.load_session(id).await.unwrap().unwrap();
1536            assert_eq!(durable.task_list_version_meta().as_deref(), Some("1"));
1537            assert_eq!(
1538                durable.task_list.as_ref().map(|list| list.title.as_str()),
1539                Some(expected_title)
1540            );
1541        }
1542    }
1543
1544    fn durable_cache_fence(
1545        operation: impl Into<String>,
1546        marker: impl Into<String>,
1547    ) -> (
1548        PostDurableHook,
1549        tokio::sync::oneshot::Receiver<()>,
1550        Arc<(Mutex<bool>, Condvar)>,
1551    ) {
1552        let operation = operation.into();
1553        let marker = marker.into();
1554        let (started_tx, started_rx) = tokio::sync::oneshot::channel();
1555        let started_tx = Arc::new(Mutex::new(Some(started_tx)));
1556        let release = Arc::new((Mutex::new(false), Condvar::new()));
1557        let hook_release = release.clone();
1558        let hook: PostDurableHook = Arc::new(move |actual_operation, actual_marker| {
1559            if actual_operation != operation || actual_marker != marker {
1560                return;
1561            }
1562            if let Some(started_tx) = started_tx.lock().unwrap().take() {
1563                started_tx.send(()).expect("fence observer still present");
1564            }
1565            let (released, wake) = &*hook_release;
1566            let mut released = released.lock().unwrap();
1567            while !*released {
1568                released = wake.wait(released).unwrap();
1569            }
1570        });
1571        (hook, started_rx, release)
1572    }
1573
1574    fn release_fence(release: &Arc<(Mutex<bool>, Condvar)>) {
1575        let (released, wake) = &**release;
1576        *released.lock().unwrap() = true;
1577        wake.notify_all();
1578    }
1579
1580    async fn assert_second_write_waits_for_cache_publish<T>(
1581        first: tokio::task::JoinHandle<std::io::Result<T>>,
1582        mut second: tokio::task::JoinHandle<std::io::Result<T>>,
1583        release: Arc<(Mutex<bool>, Condvar)>,
1584    ) -> (T, T) {
1585        let second_before_release =
1586            tokio::time::timeout(Duration::from_millis(100), &mut second).await;
1587        let completed_before_release = second_before_release.is_ok();
1588        release_fence(&release);
1589
1590        let first = first
1591            .await
1592            .expect("first writer joins")
1593            .expect("first writer succeeds");
1594        let second = match second_before_release {
1595            Ok(joined) => joined
1596                .expect("second writer joins")
1597                .expect("second writer succeeds"),
1598            Err(_) => second
1599                .await
1600                .expect("second writer joins")
1601                .expect("second writer succeeds"),
1602        };
1603        assert!(
1604            !completed_before_release,
1605            "the second write must remain behind the first write's durable-to-cache fence"
1606        );
1607        (first, second)
1608    }
1609
1610    #[derive(Clone, Copy, Debug)]
1611    enum FullSaveRoute {
1612        InherentSave,
1613        SaveAndCache,
1614        RuntimePersistence,
1615    }
1616
1617    impl FullSaveRoute {
1618        fn operation(self) -> &'static str {
1619            match self {
1620                Self::InherentSave => "save_full",
1621                Self::SaveAndCache => "save_and_cache",
1622                Self::RuntimePersistence => "save_runtime_session",
1623            }
1624        }
1625
1626        fn name(self) -> &'static str {
1627            match self {
1628                Self::InherentSave => "inherent",
1629                Self::SaveAndCache => "save-and-cache",
1630                Self::RuntimePersistence => "runtime-persistence",
1631            }
1632        }
1633    }
1634
1635    async fn assert_full_save_route_serializes_cache_publish(route: FullSaveRoute) {
1636        let temp = tempfile::tempdir().unwrap();
1637        let concrete_storage = Arc::new(
1638            bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
1639                .await
1640                .expect("SessionStoreV2"),
1641        );
1642        let storage: Arc<dyn Storage> = concrete_storage;
1643        let id = format!("root-full-cache-order-{}", route.name());
1644        let (hook, first_durable, release) = durable_cache_fence(route.operation(), id.clone());
1645        let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
1646
1647        let mut initial = Session::new(&id, "model");
1648        initial.add_message(bamboo_agent_core::Message::user("durable transcript"));
1649        initial.set_task_list(task_list(&id, "initial"));
1650        initial.set_task_list_version_meta("0");
1651        initial
1652            .metadata
1653            .insert("unrelated.runtime".to_string(), "keep".to_string());
1654        storage.save_session(&initial).await.unwrap();
1655        cache_put(&repo, &initial);
1656
1657        let first_repo = repo.clone();
1658        let mut root_snapshot = initial.clone();
1659        root_snapshot.add_message(bamboo_agent_core::Message::assistant(
1660            "full-save transcript suffix",
1661            None,
1662        ));
1663        root_snapshot.set_task_list(task_list(&id, "root"));
1664        root_snapshot.set_task_list_version_meta("1");
1665        let first = tokio::spawn(async move {
1666            match route {
1667                FullSaveRoute::InherentSave => first_repo.save(&mut root_snapshot).await,
1668                FullSaveRoute::SaveAndCache => {
1669                    first_repo.save_and_cache(&mut root_snapshot).await;
1670                    Ok(())
1671                }
1672                FullSaveRoute::RuntimePersistence => {
1673                    bamboo_domain::RuntimeSessionPersistence::save_runtime_session(
1674                        first_repo.as_ref(),
1675                        &mut root_snapshot,
1676                    )
1677                    .await
1678                }
1679            }
1680        });
1681        first_durable
1682            .await
1683            .expect("root full durable write reached");
1684
1685        let second_repo = repo.clone();
1686        let second_id = id.clone();
1687        let child_task_list = task_list(&id, "child");
1688        let second = tokio::spawn(async move {
1689            bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
1690                second_repo.as_ref(),
1691                &second_id,
1692                &child_task_list,
1693                "2",
1694            )
1695            .await
1696            .map(|updated| assert!(updated, "root must exist"))
1697        });
1698        assert_second_write_waits_for_cache_publish(first, second, release).await;
1699
1700        let durable = storage.load_session(&id).await.unwrap().unwrap();
1701        let cached = read_cached_session(repo.cache(), &id).expect("cached root");
1702        for (tier, session) in [("durable", durable), ("cache", cached)] {
1703            assert_eq!(
1704                session.task_list_version_meta().as_deref(),
1705                Some("2"),
1706                "{route:?} {tier} must retain the child transaction"
1707            );
1708            assert_eq!(
1709                session.task_list.as_ref().map(|list| list.title.as_str()),
1710                Some("child"),
1711                "{route:?} {tier} must retain the child transaction"
1712            );
1713            assert_eq!(
1714                session
1715                    .metadata
1716                    .get("unrelated.runtime")
1717                    .map(String::as_str),
1718                Some("keep"),
1719                "{route:?} {tier} must preserve unrelated runtime state"
1720            );
1721            assert_eq!(
1722                session.messages.len(),
1723                2,
1724                "{route:?} {tier} must preserve the full-save transcript"
1725            );
1726        }
1727    }
1728
1729    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1730    async fn root_full_saves_and_child_task_patch_share_publish_order() {
1731        for route in [
1732            FullSaveRoute::InherentSave,
1733            FullSaveRoute::SaveAndCache,
1734            FullSaveRoute::RuntimePersistence,
1735        ] {
1736            assert_full_save_route_serializes_cache_publish(route).await;
1737        }
1738    }
1739
1740    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1741    async fn checkpoint_and_child_task_patch_share_publish_order() {
1742        let temp = tempfile::tempdir().unwrap();
1743        let concrete_storage = Arc::new(
1744            bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
1745                .await
1746                .expect("SessionStoreV2"),
1747        );
1748        let storage: Arc<dyn Storage> = concrete_storage;
1749        let id = "checkpoint-cache-order";
1750        let (hook, checkpoint_durable, release) = durable_cache_fence("checkpoint", id);
1751        let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
1752
1753        let mut initial = Session::new(id, "model");
1754        initial.add_message(bamboo_agent_core::Message::user("durable transcript"));
1755        initial.set_task_list(task_list(id, "initial"));
1756        initial.set_task_list_version_meta("0");
1757        initial
1758            .metadata
1759            .insert("unrelated.runtime".to_string(), "keep".to_string());
1760        storage.save_session(&initial).await.unwrap();
1761        cache_put(&repo, &initial);
1762
1763        let checkpoint_repo = repo.clone();
1764        let mut checkpoint_snapshot = initial.clone();
1765        checkpoint_snapshot.add_message(bamboo_agent_core::Message::assistant(
1766            "checkpoint transcript suffix",
1767            None,
1768        ));
1769        checkpoint_snapshot.set_task_list(task_list(id, "checkpoint"));
1770        checkpoint_snapshot.set_task_list_version_meta("1");
1771        let checkpoint = tokio::spawn(async move {
1772            bamboo_domain::RuntimeSessionPersistence::checkpoint_runtime_session(
1773                checkpoint_repo.as_ref(),
1774                &mut checkpoint_snapshot,
1775            )
1776            .await
1777        });
1778        checkpoint_durable
1779            .await
1780            .expect("checkpoint durable write reached");
1781
1782        let child_repo = repo.clone();
1783        let child_task_list = task_list(id, "child");
1784        let child_patch = tokio::spawn(async move {
1785            bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
1786                child_repo.as_ref(),
1787                id,
1788                &child_task_list,
1789                "2",
1790            )
1791            .await
1792            .map(|updated| assert!(updated, "root must exist"))
1793        });
1794        assert_second_write_waits_for_cache_publish(checkpoint, child_patch, release).await;
1795
1796        let durable = storage.load_session(id).await.unwrap().unwrap();
1797        let cached = read_cached_session(repo.cache(), id).expect("cached root");
1798        for (tier, session) in [("durable", durable), ("cache", cached)] {
1799            assert_eq!(
1800                session.task_list_version_meta().as_deref(),
1801                Some("2"),
1802                "{tier} must retain the child transaction"
1803            );
1804            assert_eq!(
1805                session.task_list.as_ref().map(|list| list.title.as_str()),
1806                Some("child"),
1807                "{tier} must retain the child transaction"
1808            );
1809            assert_eq!(
1810                session
1811                    .metadata
1812                    .get("unrelated.runtime")
1813                    .map(String::as_str),
1814                Some("keep"),
1815                "{tier} must preserve unrelated runtime state"
1816            );
1817            assert_eq!(
1818                session.messages.len(),
1819                2,
1820                "{tier} must preserve the checkpoint transcript"
1821            );
1822        }
1823    }
1824
1825    #[derive(Clone, Copy, Debug)]
1826    enum CacheBackfillRoute {
1827        Load,
1828        TryLoad,
1829    }
1830
1831    impl CacheBackfillRoute {
1832        fn operation(self) -> &'static str {
1833            match self {
1834                Self::Load => "load",
1835                Self::TryLoad => "try_load",
1836            }
1837        }
1838
1839        fn name(self) -> &'static str {
1840            match self {
1841                Self::Load => "load",
1842                Self::TryLoad => "try-load",
1843            }
1844        }
1845    }
1846
1847    async fn assert_cache_backfill_serializes_with_task_patch(route: CacheBackfillRoute) {
1848        let temp = tempfile::tempdir().unwrap();
1849        let concrete_storage = Arc::new(
1850            bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
1851                .await
1852                .expect("SessionStoreV2"),
1853        );
1854        let storage: Arc<dyn Storage> = concrete_storage;
1855        let id = format!("cache-backfill-order-{}", route.name());
1856        let (hook, loaded_old_durable, release) =
1857            durable_cache_fence(route.operation(), id.clone());
1858        let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
1859
1860        let mut initial = Session::new(&id, "model");
1861        initial.set_task_list(task_list(&id, "initial"));
1862        initial.set_task_list_version_meta("0");
1863        storage.save_session(&initial).await.unwrap();
1864        assert!(
1865            read_cached_session(repo.cache(), &id).is_none(),
1866            "the race requires a genuine cache miss"
1867        );
1868
1869        let load_repo = repo.clone();
1870        let load_id = id.clone();
1871        let load = tokio::spawn(async move {
1872            let loaded = match route {
1873                CacheBackfillRoute::Load => load_repo.load(&load_id).await,
1874                CacheBackfillRoute::TryLoad => {
1875                    load_repo.try_load(&load_id).await.expect("storage load")
1876                }
1877            };
1878            assert!(loaded.is_some(), "seeded session must load");
1879            Ok(())
1880        });
1881        loaded_old_durable
1882            .await
1883            .expect("old durable snapshot loaded");
1884
1885        let patch_repo = repo.clone();
1886        let patch_id = id.clone();
1887        let child_task_list = task_list(&id, "child");
1888        let patch = tokio::spawn(async move {
1889            bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
1890                patch_repo.as_ref(),
1891                &patch_id,
1892                &child_task_list,
1893                "1",
1894            )
1895            .await
1896            .map(|updated| assert!(updated, "root must exist"))
1897        });
1898        assert_second_write_waits_for_cache_publish(load, patch, release).await;
1899
1900        let durable = storage.load_session(&id).await.unwrap().unwrap();
1901        let cached = read_cached_session(repo.cache(), &id).expect("backfilled cache");
1902        for (tier, session) in [("durable", durable), ("cache", cached)] {
1903            assert_eq!(
1904                session.task_list_version_meta().as_deref(),
1905                Some("1"),
1906                "{route:?} {tier} must retain the child transaction"
1907            );
1908            assert_eq!(
1909                session.task_list.as_ref().map(|list| list.title.as_str()),
1910                Some("child"),
1911                "{route:?} {tier} must retain the child transaction"
1912            );
1913        }
1914    }
1915
1916    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1917    async fn cache_miss_backfills_and_child_task_patch_share_publish_order() {
1918        for route in [CacheBackfillRoute::Load, CacheBackfillRoute::TryLoad] {
1919            assert_cache_backfill_serializes_with_task_patch(route).await;
1920        }
1921    }
1922
1923    #[tokio::test]
1924    async fn cache_hits_do_not_wait_for_the_persistence_lock() {
1925        let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
1926        let repo = test_repo(storage);
1927        let id = "cache-hit-lock-free";
1928        let cached = Session::new(id, "cached-model");
1929        cache_put(&repo, &cached);
1930        let persistence_guard = repo.persistence().acquire_lock(id).await;
1931
1932        let loaded = tokio::time::timeout(Duration::from_millis(100), repo.load(id)).await;
1933        let try_loaded = tokio::time::timeout(Duration::from_millis(100), repo.try_load(id)).await;
1934        drop(persistence_guard);
1935
1936        assert_eq!(
1937            loaded
1938                .expect("cache hit must not wait for the persistence lock")
1939                .expect("cached session")
1940                .model,
1941            "cached-model"
1942        );
1943        assert_eq!(
1944            try_loaded
1945                .expect("fallible cache hit must not wait for the persistence lock")
1946                .expect("cache read succeeds")
1947                .expect("cached session")
1948                .model,
1949            "cached-model"
1950        );
1951    }
1952
1953    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1954    async fn load_merged_storage_refresh_and_child_task_patch_share_publish_order() {
1955        let temp = tempfile::tempdir().unwrap();
1956        let concrete_storage = Arc::new(
1957            bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
1958                .await
1959                .expect("SessionStoreV2"),
1960        );
1961        let storage: Arc<dyn Storage> = concrete_storage;
1962        let id = "load-merged-cache-order";
1963        let (hook, loaded_old_durable, release) = durable_cache_fence("load_merged", id);
1964        let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
1965
1966        let mut durable = Session::new(id, "model");
1967        durable.updated_at = Utc::now();
1968        durable.set_task_list(task_list(id, "initial"));
1969        durable.set_task_list_version_meta("0");
1970        storage.save_session(&durable).await.unwrap();
1971
1972        let mut memory = durable.clone();
1973        memory.updated_at = durable.updated_at - chrono::Duration::seconds(1);
1974        memory.set_task_list(task_list(id, "memory"));
1975        cache_put(&repo, &memory);
1976
1977        let load_repo = repo.clone();
1978        let load = tokio::spawn(async move {
1979            assert!(
1980                load_repo.load_merged(id).await.is_some(),
1981                "seeded session must load"
1982            );
1983            Ok(())
1984        });
1985        loaded_old_durable
1986            .await
1987            .expect("old durable snapshot loaded");
1988
1989        let patch_repo = repo.clone();
1990        let child_task_list = task_list(id, "child");
1991        let patch = tokio::spawn(async move {
1992            bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
1993                patch_repo.as_ref(),
1994                id,
1995                &child_task_list,
1996                "1",
1997            )
1998            .await
1999            .map(|updated| assert!(updated, "root must exist"))
2000        });
2001        assert_second_write_waits_for_cache_publish(load, patch, release).await;
2002
2003        let durable = storage.load_session(id).await.unwrap().unwrap();
2004        let cached = read_cached_session(repo.cache(), id).expect("refreshed cache");
2005        for (tier, session) in [("durable", durable), ("cache", cached)] {
2006            assert_eq!(
2007                session.task_list_version_meta().as_deref(),
2008                Some("1"),
2009                "{tier} must retain the child transaction"
2010            );
2011            assert_eq!(
2012                session.task_list.as_ref().map(|list| list.title.as_str()),
2013                Some("child"),
2014                "{tier} must retain the child transaction"
2015            );
2016        }
2017    }
2018
2019    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
2020    async fn legacy_clear_refresh_and_child_task_patch_share_publish_order() {
2021        let temp = tempfile::tempdir().unwrap();
2022        let concrete_storage = Arc::new(
2023            bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
2024                .await
2025                .expect("SessionStoreV2"),
2026        );
2027        let storage: Arc<dyn Storage> = concrete_storage;
2028        let id = "legacy-clear-cache-order";
2029        let expected = vec![serde_json::json!({"content": "legacy"})];
2030        let (hook, loaded_post_cas_snapshot, release) = durable_cache_fence("clear_legacy", id);
2031        let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
2032
2033        let mut initial = Session::new(id, "model");
2034        initial.set_pending_injected_messages(expected.clone());
2035        initial.set_task_list(task_list(id, "initial"));
2036        initial.set_task_list_version_meta("0");
2037        storage.save_session(&initial).await.unwrap();
2038        cache_put(&repo, &initial);
2039
2040        let clear_repo = repo.clone();
2041        let clear_expected = expected.clone();
2042        let clear = tokio::spawn(async move {
2043            bamboo_domain::RuntimeSessionPersistence::clear_legacy_pending_messages(
2044                clear_repo.as_ref(),
2045                id,
2046                &clear_expected,
2047            )
2048            .await
2049            .map(|cleared| assert!(cleared, "legacy queue must match"))
2050        });
2051        loaded_post_cas_snapshot
2052            .await
2053            .expect("post-CAS snapshot loaded");
2054
2055        let patch_repo = repo.clone();
2056        let child_task_list = task_list(id, "child");
2057        let patch = tokio::spawn(async move {
2058            bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
2059                patch_repo.as_ref(),
2060                id,
2061                &child_task_list,
2062                "1",
2063            )
2064            .await
2065            .map(|updated| assert!(updated, "root must exist"))
2066        });
2067        assert_second_write_waits_for_cache_publish(clear, patch, release).await;
2068
2069        let durable = storage.load_session(id).await.unwrap().unwrap();
2070        let cached = read_cached_session(repo.cache(), id).expect("refreshed cache");
2071        for (tier, session) in [("durable", durable), ("cache", cached)] {
2072            assert_eq!(
2073                session.task_list_version_meta().as_deref(),
2074                Some("1"),
2075                "{tier} must retain the child transaction"
2076            );
2077            assert_eq!(
2078                session.task_list.as_ref().map(|list| list.title.as_str()),
2079                Some("child"),
2080                "{tier} must retain the child transaction"
2081            );
2082            assert!(
2083                !session.has_pending_injected_messages(),
2084                "{tier} must retain the successful legacy clear"
2085            );
2086        }
2087    }
2088
2089    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
2090    async fn concurrent_task_patches_publish_cache_in_durable_order() {
2091        let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
2092        let (hook, first_durable, release) = durable_cache_fence("task", "1");
2093        let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
2094        let id = "concurrent-task-cache-order";
2095        let mut initial = Session::new(id, "model");
2096        initial.set_task_list(task_list(id, "initial"));
2097        initial.set_task_list_version_meta("0");
2098        storage.save_session(&initial).await.unwrap();
2099        cache_put(&repo, &initial);
2100
2101        let first_repo = repo.clone();
2102        let first_task_list = task_list(id, "first");
2103        let first = tokio::spawn(async move {
2104            bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
2105                first_repo.as_ref(),
2106                id,
2107                &first_task_list,
2108                "1",
2109            )
2110            .await
2111        });
2112        first_durable.await.expect("first durable write reached");
2113
2114        let second_repo = repo.clone();
2115        let second_task_list = task_list(id, "second");
2116        let second = tokio::spawn(async move {
2117            bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
2118                second_repo.as_ref(),
2119                id,
2120                &second_task_list,
2121                "2",
2122            )
2123            .await
2124        });
2125        let (first_updated, second_updated) =
2126            assert_second_write_waits_for_cache_publish(first, second, release).await;
2127        assert!(first_updated && second_updated);
2128
2129        let durable = storage.load_session(id).await.unwrap().unwrap();
2130        let cached = read_cached_session(repo.cache(), id).expect("cached root");
2131        for (tier, session) in [("durable", durable), ("cache", cached)] {
2132            assert_eq!(
2133                session.task_list_version_meta().as_deref(),
2134                Some("2"),
2135                "{tier} must retain the second transaction"
2136            );
2137            assert_eq!(
2138                session.task_list.as_ref().map(|list| list.title.as_str()),
2139                Some("second"),
2140                "{tier} must retain the second transaction"
2141            );
2142        }
2143    }
2144
2145    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
2146    async fn root_control_plane_save_and_child_task_patch_share_publish_order() {
2147        let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
2148        let (hook, root_durable, release) = durable_cache_fence("save", "root-control-plane-order");
2149        let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
2150        let id = "root-control-plane-order";
2151
2152        let mut initial = Session::new(id, "model");
2153        initial.add_message(bamboo_agent_core::Message::user("durable transcript"));
2154        initial.set_task_list(task_list(id, "initial"));
2155        initial.set_task_list_version_meta("0");
2156        initial
2157            .metadata
2158            .insert("unrelated.runtime".to_string(), "keep".to_string());
2159        storage.save_session(&initial).await.unwrap();
2160        cache_put(&repo, &initial);
2161
2162        let root_repo = repo.clone();
2163        let mut root_snapshot = initial.clone();
2164        root_snapshot.set_task_list(task_list(id, "root"));
2165        root_snapshot.set_task_list_version_meta("1");
2166        let root_save = tokio::spawn(async move {
2167            bamboo_domain::RuntimeSessionPersistence::save_runtime_control_plane(
2168                root_repo.as_ref(),
2169                &mut root_snapshot,
2170            )
2171            .await
2172        });
2173        root_durable.await.expect("root durable write reached");
2174
2175        let child_repo = repo.clone();
2176        let child_task_list = task_list(id, "child");
2177        let child_patch = tokio::spawn(async move {
2178            bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
2179                child_repo.as_ref(),
2180                id,
2181                &child_task_list,
2182                "2",
2183            )
2184            .await
2185            .map(|updated| {
2186                assert!(updated, "root must exist");
2187            })
2188        });
2189        assert_second_write_waits_for_cache_publish(root_save, child_patch, release).await;
2190
2191        let durable = storage.load_session(id).await.unwrap().unwrap();
2192        let cached = read_cached_session(repo.cache(), id).expect("cached root");
2193        for (tier, session) in [("durable", durable), ("cache", cached)] {
2194            assert_eq!(
2195                session.task_list_version_meta().as_deref(),
2196                Some("2"),
2197                "{tier} must retain the child transaction"
2198            );
2199            assert_eq!(
2200                session.task_list.as_ref().map(|list| list.title.as_str()),
2201                Some("child"),
2202                "{tier} must retain the child transaction"
2203            );
2204            assert_eq!(
2205                session
2206                    .metadata
2207                    .get("unrelated.runtime")
2208                    .map(String::as_str),
2209                Some("keep"),
2210                "{tier} must preserve unrelated runtime state"
2211            );
2212            assert_eq!(
2213                session.messages.len(),
2214                1,
2215                "{tier} must preserve the transcript"
2216            );
2217        }
2218    }
2219
2220    #[tokio::test]
2221    async fn stale_control_plane_save_keeps_checkpointed_ledger_in_durable_and_cache() {
2222        let temp = tempfile::tempdir().unwrap();
2223        let storage: Arc<dyn Storage> = Arc::new(
2224            bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
2225                .await
2226                .expect("SessionStoreV2"),
2227        );
2228        let repo = test_repo(storage.clone());
2229        let id = "control-plane-ledger-cache";
2230        let mut initial = Session::new(id, "model");
2231        initial.add_message(bamboo_agent_core::Message::user("durable transcript"));
2232        let assistant = bamboo_agent_core::Message::assistant("normalized", None);
2233        let anchor = assistant.id.clone();
2234        initial.add_message(assistant);
2235        let native_item = bamboo_domain::ProviderTranscriptItem::try_from_payload(
2236            bamboo_domain::ProviderFamily::OpenAi,
2237            bamboo_domain::ProviderProtocol::OpenAiResponsesV1,
2238            bamboo_domain::ProviderTranscriptOrigin::Provider,
2239            bamboo_domain::ProviderTranscriptAuthor::Model,
2240            serde_json::json!({
2241                "type":"tool_search_call","id":"tsc_cache_native","execution":"client","call_id":"cache_native",
2242                "status":"completed","arguments":{"query":"CACHE_NATIVE_PAYLOAD_SENTINEL"}
2243            }),
2244        )
2245        .unwrap();
2246        initial
2247            .append_provider_transcript_group(&anchor, None, vec![native_item])
2248            .unwrap();
2249        storage.save_session(&initial).await.unwrap();
2250        cache_put(&repo, &initial);
2251        let mut stale = initial.clone();
2252
2253        let mut runner = initial;
2254        runner.model_context_state = Some(bamboo_domain::ModelContextState {
2255            state_revision: 1,
2256            prefix_epoch: 1,
2257            cache_scope_sha256: Some("scope".to_string()),
2258            transcript_item_sha256: vec!["runner-l1".to_string()],
2259            ..bamboo_domain::ModelContextState::default()
2260        });
2261        bamboo_domain::RuntimeSessionPersistence::checkpoint_runtime_session(&repo, &mut runner)
2262            .await
2263            .unwrap();
2264
2265        stale
2266            .metadata
2267            .insert("runtime.suspend_reason".to_string(), "waiting".to_string());
2268        bamboo_domain::RuntimeSessionPersistence::save_runtime_control_plane(&repo, &mut stale)
2269            .await
2270            .unwrap();
2271
2272        let expected = runner.model_context_state;
2273        let expected_native = runner.provider_transcript;
2274        let durable = storage.load_session(id).await.unwrap().unwrap();
2275        let cached = read_cached_session(repo.cache(), id).expect("cached session");
2276        for (tier, session) in [("durable", durable), ("cache", cached.clone())] {
2277            assert_eq!(session.model_context_state, expected, "tier={tier}");
2278            assert_eq!(
2279                session.provider_transcript, expected_native,
2280                "tier={tier} must retain the message-anchored native transcript"
2281            );
2282            assert_eq!(
2283                session
2284                    .metadata
2285                    .get("runtime.suspend_reason")
2286                    .map(String::as_str),
2287                Some("waiting"),
2288                "tier={tier}"
2289            );
2290            assert_eq!(session.messages.len(), 2, "tier={tier}");
2291        }
2292
2293        let runtime_json =
2294            std::fs::read_to_string(temp.path().join("sessions").join(id).join("runtime.json"))
2295                .unwrap();
2296        assert!(!runtime_json.contains("CACHE_NATIVE_PAYLOAD_SENTINEL"));
2297
2298        // Prove the cache projection cannot turn a runtime-only update into a
2299        // later durable loss when that cached value becomes a full checkpoint.
2300        let mut cache_writer = cached;
2301        bamboo_domain::RuntimeSessionPersistence::checkpoint_runtime_session(
2302            &repo,
2303            &mut cache_writer,
2304        )
2305        .await
2306        .unwrap();
2307        let restarted = storage.load_session(id).await.unwrap().unwrap();
2308        assert_eq!(restarted.provider_transcript, expected_native);
2309    }
2310
2311    #[tokio::test]
2312    async fn narrow_runtime_metadata_transaction_preserves_live_and_durable_non_owned_state() {
2313        let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
2314        let repo = test_repo(storage.clone());
2315        let id = "narrow-metadata";
2316        let mut durable = Session::new(id, "durable-model");
2317        durable.add_message(bamboo_agent_core::Message::user("durable user turn"));
2318        durable
2319            .metadata
2320            .insert("external.durable".to_string(), "keep".to_string());
2321        storage.save_session(&durable).await.expect("seed durable");
2322
2323        let mut live = durable.clone();
2324        live.add_message(bamboo_agent_core::Message::assistant(
2325            "in-flight assistant tool call",
2326            None,
2327        ));
2328        live.model = "live-model".to_string();
2329        live.metadata
2330            .insert("external.live".to_string(), "keep".to_string());
2331        cache_put(&repo, &live);
2332
2333        repo.update_runtime_session(id, &["workflow.owned"], |latest| {
2334            latest
2335                .metadata
2336                .insert("workflow.owned".to_string(), "active".to_string());
2337        })
2338        .await
2339        .expect("transaction")
2340        .expect("session exists");
2341
2342        let saved = storage
2343            .load_session(id)
2344            .await
2345            .expect("load durable")
2346            .expect("durable exists");
2347        assert_eq!(
2348            saved.messages.len(),
2349            1,
2350            "transaction never writes stale live messages"
2351        );
2352        assert_eq!(
2353            saved.metadata.get("external.durable").map(String::as_str),
2354            Some("keep")
2355        );
2356        assert_eq!(
2357            saved.metadata.get("workflow.owned").map(String::as_str),
2358            Some("active")
2359        );
2360
2361        let cached = read_cached_session(repo.cache(), id).expect("live cache");
2362        assert_eq!(
2363            cached.messages.len(),
2364            2,
2365            "cache live tool call is not replaced"
2366        );
2367        assert_eq!(cached.model, "live-model");
2368        assert_eq!(
2369            cached.metadata.get("external.live").map(String::as_str),
2370            Some("keep")
2371        );
2372        assert_eq!(
2373            cached.metadata.get("workflow.owned").map(String::as_str),
2374            Some("active")
2375        );
2376    }
2377
2378    /// Regression guard: a strictly-newer in-memory session (e.g. one that just
2379    /// answered and cleared its pending question) must win over a strictly-older
2380    /// storage copy that still carries the pending question — both in the value
2381    /// returned AND in the cache (no clobber).
2382    #[tokio::test]
2383    async fn load_merged_does_not_regress_to_older_storage() {
2384        let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
2385        let repo = test_repo(storage.clone());
2386        let id = "s1";
2387
2388        let mut stale = Session::new(id.to_string(), "m");
2389        stale.set_pending_question(
2390            "tc1".into(),
2391            "kind".into(),
2392            "q?".into(),
2393            vec!["OK".into()],
2394            true,
2395        );
2396        stale.updated_at = Utc::now() - chrono::Duration::seconds(10);
2397        storage.save_session(&stale).await.unwrap();
2398
2399        // This is a fresher snapshot of the same Root birth, not a recreated
2400        // Session with the same ID (whose live durable birth takes precedence).
2401        let mut fresh = stale.clone();
2402        fresh.clear_pending_question();
2403        fresh.updated_at = Utc::now();
2404        cache_put(&repo, &fresh);
2405
2406        let merged = repo.load_merged(id).await.expect("session exists");
2407        assert!(
2408            merged.pending_question.is_none(),
2409            "must return the newer answered memory copy, not the stale storage one"
2410        );
2411        let cached = read_cached_session(repo.cache(), id).expect("cached");
2412        assert!(
2413            cached.pending_question.is_none(),
2414            "load_merged must never regress the cache to a stale storage copy"
2415        );
2416    }
2417
2418    /// The pending-question recovery still works when storage is the same age:
2419    /// if memory lost a pending question that same-age storage retains, prefer
2420    /// storage so a genuine clarification is not dropped.
2421    #[tokio::test]
2422    async fn load_merged_recovers_pending_question_from_same_age_storage() {
2423        let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
2424        let repo = test_repo(storage.clone());
2425        let id = "s2";
2426        let ts = Utc::now();
2427
2428        let mut with_pending = Session::new(id.to_string(), "m");
2429        with_pending.set_pending_question(
2430            "tc".into(),
2431            "k".into(),
2432            "q".into(),
2433            vec!["OK".into()],
2434            true,
2435        );
2436        with_pending.updated_at = ts;
2437        storage.save_session(&with_pending).await.unwrap();
2438
2439        let mut lost = with_pending.clone();
2440        lost.clear_pending_question();
2441        lost.updated_at = ts;
2442        cache_put(&repo, &lost);
2443
2444        let merged = repo.load_merged(id).await.expect("session exists");
2445        assert!(
2446            merged.pending_question.is_some(),
2447            "same-age storage carrying a pending question must still be recovered"
2448        );
2449    }
2450
2451    #[tokio::test]
2452    async fn runtime_publish_refreshes_cache_even_when_storage_fails() {
2453        let id = "runtime-selection";
2454        let mut previous = Session::new(id.to_string(), "m");
2455        previous.metadata.insert(
2456            "skill_runtime_selected_skill_ids".to_string(),
2457            "[\"plan\"]".to_string(),
2458        );
2459        let storage: Arc<dyn Storage> = Arc::new(FailingSaveStorage {
2460            persisted: Mutex::new(Some(previous.clone())),
2461        });
2462        let repo = test_repo(storage.clone());
2463        cache_put(&repo, &previous);
2464
2465        let mut current = previous.clone();
2466        current.metadata.insert(
2467            "skill_runtime_selected_skill_ids".to_string(),
2468            "[\"review\"]".to_string(),
2469        );
2470        current.updated_at = Utc::now();
2471
2472        let result =
2473            bamboo_domain::RuntimeSessionPersistence::save_runtime_session(&repo, &mut current)
2474                .await;
2475        assert!(result.is_err(), "durable failure must still be surfaced");
2476
2477        let cached = repo.load(id).await.expect("cached current session");
2478        assert_eq!(
2479            cached
2480                .metadata
2481                .get("skill_runtime_selected_skill_ids")
2482                .map(String::as_str),
2483            Some("[\"review\"]")
2484        );
2485        let allowlist = bamboo_skills::access_control::extract_skill_allowlist(&cached.metadata)
2486            .expect("runtime authorization allowlist");
2487        assert!(allowlist.contains("review"));
2488        assert!(!allowlist.contains("plan"));
2489        let durable = storage
2490            .load_session(id)
2491            .await
2492            .expect("load durable state")
2493            .expect("previous durable session");
2494        assert_eq!(
2495            durable
2496                .metadata
2497                .get("skill_runtime_selected_skill_ids")
2498                .map(String::as_str),
2499            Some("[\"plan\"]")
2500        );
2501    }
2502
2503    #[tokio::test]
2504    async fn finalized_runtime_save_reconciles_late_child_completion_in_disk_and_cache() {
2505        use bamboo_domain::session::runtime_state::{
2506            AgentRuntimeState, AgentStatusState, ChildWaitPolicy, WaitingForChildrenState,
2507        };
2508
2509        let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
2510        let repo = test_repo(storage.clone());
2511        let id = "late-child-finalization";
2512        let mut parent = Session::new(id, "model");
2513        let mut runtime = AgentRuntimeState::new("parent-run");
2514        runtime.status = AgentStatusState::Suspended;
2515        let mut wait = WaitingForChildrenState::for_children(
2516            vec!["child-1".to_string()],
2517            ChildWaitPolicy::All,
2518            Utc::now(),
2519        );
2520        wait.registered_by_tool_call_id = Some("tc_wait".to_string());
2521        runtime.waiting_for_children = Some(wait);
2522        parent.agent_runtime_state = Some(runtime);
2523        parent.metadata.insert(
2524            "runtime.suspend_reason".to_string(),
2525            "waiting_for_children".to_string(),
2526        );
2527        parent.set_last_run_status("suspended");
2528        storage.save_session(&parent).await.unwrap();
2529        cache_put(&repo, &parent);
2530
2531        let mut stale_runner = parent.clone();
2532        let mut completed = parent;
2533        let completed_runtime = completed.agent_runtime_state.as_mut().unwrap();
2534        completed_runtime.waiting_for_children = None;
2535        completed_runtime.status = AgentStatusState::Idle;
2536        completed_runtime.suspension = None;
2537        completed.metadata.remove("runtime.suspend_reason");
2538        let mut outcome = bamboo_agent_core::Message::user("completed child result");
2539        outcome.metadata = Some(serde_json::json!({
2540            "runtime_kind": "child_completion_resume"
2541        }));
2542        completed.add_message(outcome);
2543        storage.save_session(&completed).await.unwrap();
2544
2545        bamboo_domain::RuntimeSessionPersistence::save_finalized_runtime_session(
2546            &repo,
2547            &mut stale_runner,
2548        )
2549        .await
2550        .unwrap();
2551
2552        let durable = storage.load_session(id).await.unwrap().unwrap();
2553        let cached = read_cached_session(repo.cache(), id).expect("final snapshot published");
2554        for saved in [&durable, &cached] {
2555            assert!(saved
2556                .agent_runtime_state
2557                .as_ref()
2558                .unwrap()
2559                .waiting_for_children
2560                .is_none());
2561            assert!(!saved.metadata.contains_key("runtime.suspend_reason"));
2562            assert_eq!(saved.last_run_status().as_deref(), Some("completed"));
2563            assert!(saved
2564                .messages
2565                .iter()
2566                .any(|message| message.content == "completed child result"));
2567        }
2568    }
2569
2570    #[tokio::test]
2571    async fn inherent_save_leaves_existing_cache_untouched_when_storage_fails() {
2572        let id = "inherent-save-failure";
2573        let previous = Session::new(id, "previous");
2574        let storage: Arc<dyn Storage> = Arc::new(FailingSaveStorage {
2575            persisted: Mutex::new(Some(previous.clone())),
2576        });
2577        let repo = test_repo(storage);
2578        cache_put(&repo, &previous);
2579
2580        let mut current = previous.clone();
2581        current.model = "current".to_string();
2582        assert!(repo.save(&mut current).await.is_err());
2583        assert_eq!(
2584            read_cached_session(repo.cache(), id)
2585                .expect("existing cache")
2586                .model,
2587            "previous",
2588            "fallible inherent save must publish only after a durable commit"
2589        );
2590    }
2591
2592    #[tokio::test]
2593    async fn blocked_replay_ambiguous_save_evicts_same_age_stale_approval_cache() {
2594        let id = "blocked-replay-ambiguous-save";
2595        let mut approved = Session::new(id, "model");
2596        approved.metadata.insert(
2597            crate::session_app::respond::PERMISSION_REEXECUTE_METADATA_KEY.into(),
2598            "call-1".into(),
2599        );
2600        let storage: Arc<dyn Storage> = Arc::new(AmbiguousReplaySaveStorage {
2601            persisted: Mutex::new(Some(approved.clone())),
2602        });
2603        let repo = test_repo(storage);
2604        cache_put(&repo, &approved);
2605
2606        let mut resolved = approved.clone();
2607        resolved
2608            .metadata
2609            .remove(crate::session_app::respond::PERMISSION_REEXECUTE_METADATA_KEY);
2610        assert!(repo.save_replay_resolution(&mut resolved).await.is_err());
2611        assert!(read_cached_session(repo.cache(), id).is_none());
2612        let reloaded = repo.load_merged_checked(id).await.unwrap().unwrap();
2613        assert_eq!(reloaded.updated_at, approved.updated_at);
2614        assert!(!reloaded
2615            .metadata
2616            .contains_key(crate::session_app::respond::PERMISSION_REEXECUTE_METADATA_KEY));
2617    }
2618
2619    #[tokio::test]
2620    async fn blocked_replay_precommit_save_error_reloads_retryable_approval() {
2621        let id = "blocked-replay-precommit-save";
2622        let mut approved = Session::new(id, "model");
2623        approved.metadata.insert(
2624            crate::session_app::respond::PERMISSION_REEXECUTE_METADATA_KEY.into(),
2625            "call-1".into(),
2626        );
2627        let storage: Arc<dyn Storage> = Arc::new(FailingSaveStorage {
2628            persisted: Mutex::new(Some(approved.clone())),
2629        });
2630        let repo = test_repo(storage);
2631        cache_put(&repo, &approved);
2632
2633        let mut resolved = approved.clone();
2634        resolved
2635            .metadata
2636            .remove(crate::session_app::respond::PERMISSION_REEXECUTE_METADATA_KEY);
2637        assert!(repo.save_replay_resolution(&mut resolved).await.is_err());
2638        assert!(read_cached_session(repo.cache(), id).is_none());
2639        let reloaded = repo.load_merged_checked(id).await.unwrap().unwrap();
2640        assert_eq!(reloaded.updated_at, approved.updated_at);
2641        assert_eq!(
2642            reloaded
2643                .metadata
2644                .get(crate::session_app::respond::PERMISSION_REEXECUTE_METADATA_KEY)
2645                .map(String::as_str),
2646            Some("call-1")
2647        );
2648    }
2649
2650    #[tokio::test]
2651    async fn save_and_cache_still_refreshes_cache_when_storage_fails() {
2652        let id = "save-and-cache-failure";
2653        let previous = Session::new(id, "previous");
2654        let storage: Arc<dyn Storage> = Arc::new(FailingSaveStorage {
2655            persisted: Mutex::new(Some(previous.clone())),
2656        });
2657        let repo = test_repo(storage);
2658        cache_put(&repo, &previous);
2659
2660        let mut current = previous;
2661        current.model = "current".to_string();
2662        repo.save_and_cache(&mut current).await;
2663        assert_eq!(
2664            read_cached_session(repo.cache(), id)
2665                .expect("refreshed cache")
2666                .model,
2667            "current",
2668            "fire-and-forget save must retain its existing cache-on-failure behavior"
2669        );
2670    }
2671
2672    #[tokio::test]
2673    async fn checkpoint_leaves_existing_cache_untouched_when_storage_fails() {
2674        let id = "checkpoint-failure";
2675        let previous = Session::new(id, "previous");
2676        let storage: Arc<dyn Storage> = Arc::new(FailingSaveStorage {
2677            persisted: Mutex::new(Some(previous.clone())),
2678        });
2679        let repo = test_repo(storage);
2680        cache_put(&repo, &previous);
2681
2682        let mut current = previous.clone();
2683        current.model = "current".to_string();
2684        let result = bamboo_domain::RuntimeSessionPersistence::checkpoint_runtime_session(
2685            &repo,
2686            &mut current,
2687        )
2688        .await;
2689        assert!(result.is_err());
2690        assert_eq!(
2691            read_cached_session(repo.cache(), id)
2692                .expect("existing cache")
2693                .model,
2694            "previous",
2695            "checkpoint must publish only after a durable commit"
2696        );
2697    }
2698
2699    #[tokio::test]
2700    async fn retrieval_window_checkpoint_publishes_the_exact_archived_snapshot() {
2701        let temp = tempfile::tempdir().unwrap();
2702        let storage: Arc<dyn Storage> = Arc::new(
2703            bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
2704                .await
2705                .expect("SessionStoreV2"),
2706        );
2707        let repo = test_repo(storage.clone());
2708        let id = "retrieval-window-repository-checkpoint";
2709        let mut expected = Session::new(id, "model");
2710        expected.add_message(bamboo_agent_core::Message::user("archive candidate"));
2711        expected.add_message(bamboo_agent_core::Message::assistant("retain", None));
2712        storage.save_session(&expected).await.unwrap();
2713        cache_put(&repo, &expected);
2714        let mut staged = staged_retrieval_window_archive(&expected);
2715
2716        let outcome = bamboo_domain::RuntimeSessionPersistence::checkpoint_retrieval_window(
2717            &repo,
2718            &expected,
2719            &mut staged,
2720        )
2721        .await
2722        .unwrap();
2723
2724        assert_eq!(outcome, RetrievalWindowCheckpointOutcome::Committed);
2725        let durable = storage.load_session(id).await.unwrap().unwrap();
2726        let cached = read_cached_session(repo.cache(), id).expect("cached archived Session");
2727        for (tier, saved) in [("durable", durable), ("cache", cached)] {
2728            assert!(saved.messages[0].compressed, "tier={tier}");
2729            assert_eq!(
2730                saved.messages[0].compressed_by_event_id, staged.messages[0].compressed_by_event_id,
2731                "tier={tier}"
2732            );
2733            assert_eq!(
2734                serde_json::to_value(&saved).unwrap(),
2735                serde_json::to_value(&staged).unwrap(),
2736                "tier={tier}"
2737            );
2738        }
2739    }
2740
2741    #[tokio::test]
2742    async fn legacy_clear_uses_durable_cas_and_never_erases_concurrent_append_from_stale_cache() {
2743        let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
2744        let repo = test_repo(storage.clone());
2745        let id = "legacy-cas-race";
2746        let expected = vec![serde_json::json!({"content": "first"})];
2747
2748        let mut stale_cache = Session::new(id, "m");
2749        stale_cache.set_pending_injected_messages(expected.clone());
2750        cache_put(&repo, &stale_cache);
2751
2752        let mut durable = stale_cache.clone();
2753        durable.set_pending_injected_messages(vec![
2754            serde_json::json!({"content": "first"}),
2755            serde_json::json!({"content": "concurrent"}),
2756        ]);
2757        storage.save_session(&durable).await.unwrap();
2758
2759        let cleared = bamboo_domain::RuntimeSessionPersistence::clear_legacy_pending_messages(
2760            &repo, id, &expected,
2761        )
2762        .await
2763        .unwrap();
2764        assert!(!cleared, "the durable compare-and-clear must reject drift");
2765        assert_eq!(
2766            storage
2767                .load_session(id)
2768                .await
2769                .unwrap()
2770                .unwrap()
2771                .pending_injected_messages()
2772                .unwrap(),
2773            durable.pending_injected_messages().unwrap(),
2774            "the concurrent durable append must remain intact"
2775        );
2776        assert_eq!(
2777            read_cached_session(repo.cache(), id)
2778                .unwrap()
2779                .pending_injected_messages()
2780                .unwrap(),
2781            expected,
2782            "a failed CAS must not mutate the existing cache"
2783        );
2784    }
2785
2786    #[tokio::test]
2787    async fn successful_legacy_clear_refreshes_stale_cache_from_durable_state() {
2788        let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
2789        let repo = test_repo(storage.clone());
2790        let id = "legacy-cas-success";
2791        let expected = vec![serde_json::json!({"content": "first"})];
2792
2793        let mut stale_cache = Session::new(id, "stale-model");
2794        stale_cache.set_pending_injected_messages(expected.clone());
2795        cache_put(&repo, &stale_cache);
2796
2797        let mut durable = Session::new(id, "durable-model");
2798        durable.set_pending_injected_messages(expected.clone());
2799        durable
2800            .metadata
2801            .insert("durable-only".to_string(), "keep".to_string());
2802        storage.save_session(&durable).await.unwrap();
2803
2804        assert!(
2805            bamboo_domain::RuntimeSessionPersistence::clear_legacy_pending_messages(
2806                &repo, id, &expected,
2807            )
2808            .await
2809            .unwrap()
2810        );
2811        let cached = read_cached_session(repo.cache(), id).unwrap();
2812        assert!(!cached.has_pending_injected_messages());
2813        assert_eq!(cached.model, "durable-model");
2814        assert_eq!(
2815            cached.metadata.get("durable-only").map(String::as_str),
2816            Some("keep")
2817        );
2818    }
2819
2820    #[tokio::test]
2821    async fn repository_final_save_publishes_host_checkpointed_parent_question() {
2822        use bamboo_domain::{
2823            FunctionCall, ParentQuestion, ParentQuestionCheckpointV1, PendingQuestion,
2824            PendingQuestionSource, ToolCall,
2825        };
2826
2827        let temp = tempfile::tempdir().unwrap();
2828        let store = Arc::new(
2829            bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
2830                .await
2831                .unwrap(),
2832        );
2833        let storage: Arc<dyn Storage> = store.clone();
2834        let repo = test_repo(storage);
2835        let parent = Session::new("repository-question-parent", "model");
2836        store.save_session(&parent).await.unwrap();
2837        let mut child =
2838            Session::new_child_of("repository-question-child", &parent, "model", "assignment");
2839        child.add_message(bamboo_agent_core::Message::user("assignment"));
2840        child.set_last_run_status("running");
2841        store.save_session(&child).await.unwrap();
2842        cache_put(&repo, &child);
2843
2844        let assistant = bamboo_agent_core::Message::assistant(
2845            "",
2846            Some(vec![ToolCall {
2847                id: "question-call".into(),
2848                tool_type: "function".into(),
2849                function: FunctionCall {
2850                    name: "AgenticQuestion".into(),
2851                    arguments: "{}".into(),
2852                },
2853            }]),
2854        );
2855        let result = bamboo_agent_core::Message::tool_result_with_status(
2856            "question-call",
2857            "Clarification needed: Which option?",
2858            true,
2859        );
2860        let pending = PendingQuestion {
2861            tool_call_id: "question-call".into(),
2862            tool_name: "AgenticQuestion".into(),
2863            question: "Which option?".into(),
2864            options: vec!["A".into(), "B".into()],
2865            allow_custom: true,
2866            source: PendingQuestionSource::AgenticClarification,
2867        };
2868        let observation = ParentQuestionCheckpointV1 {
2869            version: 1,
2870            prefix_message_count: child.messages.len(),
2871            prefix_digest: ParentQuestion::prefix_digest(&child.messages).unwrap(),
2872            suffix: vec![assistant, result.clone()],
2873            tool_call_id: pending.tool_call_id.clone(),
2874            tool_result_message_id: result.id.clone(),
2875            question_digest: ParentQuestion::question_digest(&pending, &result.id),
2876            pending,
2877        };
2878        let (_, question) = store
2879            .checkpoint_parent_question(&child, &parent, &observation)
2880            .await
2881            .unwrap()
2882            .unwrap();
2883
2884        let mut old_run = child.clone();
2885        old_run.metadata.insert(
2886            "runtime.actor_parent_question_handoff".into(),
2887            "true".into(),
2888        );
2889        old_run.set_last_run_status("suspended");
2890        old_run.add_message(bamboo_agent_core::Message::assistant(
2891            "old terminal text",
2892            None,
2893        ));
2894        cache_put(&repo, &old_run);
2895        bamboo_domain::RuntimeSessionPersistence::save_runtime_session(&repo, &mut old_run)
2896            .await
2897            .unwrap();
2898
2899        let cached = repo.load(&child.id).await.unwrap();
2900        let reopened = bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
2901            .await
2902            .unwrap()
2903            .load_session(&child.id)
2904            .await
2905            .unwrap()
2906            .unwrap();
2907        for saved in [&cached, &reopened] {
2908            assert_eq!(saved.last_run_status().as_deref(), Some("suspended"));
2909            assert_eq!(saved.messages.len(), child.messages.len() + 2);
2910            assert!(!saved
2911                .messages
2912                .iter()
2913                .any(|message| message.content == "old terminal text"));
2914            assert_eq!(
2915                saved.pending_question.as_ref().unwrap().source,
2916                PendingQuestionSource::DirectParent
2917            );
2918            assert_eq!(
2919                ParentQuestion::for_pending(&parent, saved).unwrap().id,
2920                question.id
2921            );
2922        }
2923    }
2924}
2925
2926fn validate_bound_root_snapshot(
2927    owner: &bamboo_domain::RootActorRuntimeWrite,
2928    snapshot: &Session,
2929) -> std::io::Result<()> {
2930    if snapshot.id != owner.fence.actor_id
2931        || snapshot.created_at != owner.created_at
2932        || snapshot.kind != bamboo_domain::SessionKind::Root
2933        || !snapshot.authority_identity.is_ordinary()
2934    {
2935        return Err(root_actor_binding_error("bound Root identity changed"));
2936    }
2937    owner
2938        .conversation_generation
2939        .check_canonical(snapshot)
2940        .map_err(root_actor_binding_error)
2941}
2942
2943fn root_actor_binding_error(reason: impl std::fmt::Display) -> std::io::Error {
2944    std::io::Error::new(
2945        std::io::ErrorKind::WouldBlock,
2946        bamboo_domain::SessionAuthorityConflict(format!("Root Actor execution rejected: {reason}")),
2947    )
2948}