Skip to main content

bamboo_engine/
session_repository.rs

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