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