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