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_storage::LockedSessionStore;
20
21use crate::{read_cached_session, SessionCache};
22
23#[cfg(test)]
24type PostDurableHook = Arc<dyn Fn(&str, &str) + Send + Sync>;
25
26/// Framework-owned coordinator over a session's cache / storage / persistence
27/// tiers. Cheap to clone (all fields are `Arc`).
28#[derive(Clone)]
29pub struct SessionRepository {
30    cache: SessionCache,
31    storage: Arc<dyn Storage>,
32    persistence: Arc<LockedSessionStore>,
33    #[cfg(test)]
34    post_durable_hook: Option<PostDurableHook>,
35}
36
37impl SessionRepository {
38    pub fn new(
39        cache: SessionCache,
40        storage: Arc<dyn Storage>,
41        persistence: Arc<LockedSessionStore>,
42    ) -> Self {
43        Self {
44            cache,
45            storage,
46            persistence,
47            #[cfg(test)]
48            post_durable_hook: None,
49        }
50    }
51
52    #[cfg(test)]
53    fn with_post_durable_hook(mut self, hook: PostDurableHook) -> Self {
54        self.post_durable_hook = Some(hook);
55        self
56    }
57
58    #[cfg(test)]
59    fn run_post_durable_hook(&self, operation: &str, marker: &str) {
60        if let Some(hook) = self.post_durable_hook.as_ref() {
61            hook(operation, marker);
62        }
63    }
64
65    pub fn cache(&self) -> &SessionCache {
66        &self.cache
67    }
68
69    pub fn storage(&self) -> &Arc<dyn Storage> {
70        &self.storage
71    }
72
73    pub fn persistence(&self) -> &Arc<LockedSessionStore> {
74        &self.persistence
75    }
76
77    /// Load a session from the memory cache, falling back to durable storage
78    /// (and back-filling the cache on a storage hit). `None` if absent in both.
79    pub async fn load(&self, session_id: &str) -> Option<Session> {
80        if let Some(session) = read_cached_session(&self.cache, session_id) {
81            return Some(session);
82        }
83
84        let _guard = self.persistence.acquire_lock(session_id).await;
85        if let Some(session) = read_cached_session(&self.cache, session_id) {
86            return Some(session);
87        }
88
89        let loaded = self.storage.load_session(session_id).await;
90        #[cfg(test)]
91        self.run_post_durable_hook("load", session_id);
92        match loaded {
93            Ok(Some(session)) => {
94                self.cache.insert(
95                    session_id.to_string(),
96                    Arc::new(parking_lot::RwLock::new(session.clone())),
97                );
98                Some(session)
99            }
100            _ => None,
101        }
102    }
103
104    /// Like [`load`](Self::load), but surfaces storage errors instead of
105    /// swallowing them to `None`. Cache hit short-circuits; a storage hit
106    /// back-fills the cache.
107    pub async fn try_load(&self, session_id: &str) -> std::io::Result<Option<Session>> {
108        if let Some(session) = read_cached_session(&self.cache, session_id) {
109            return Ok(Some(session));
110        }
111
112        let _guard = self.persistence.acquire_lock(session_id).await;
113        if let Some(session) = read_cached_session(&self.cache, session_id) {
114            return Ok(Some(session));
115        }
116
117        let loaded = self.storage.load_session(session_id).await?;
118        #[cfg(test)]
119        self.run_post_durable_hook("try_load", session_id);
120        if let Some(ref session) = loaded {
121            self.cache.insert(
122                session_id.to_string(),
123                Arc::new(parking_lot::RwLock::new(session.clone())),
124            );
125        }
126        Ok(loaded)
127    }
128
129    /// Persist the session (merge-on-write) and refresh the cache, surfacing
130    /// storage errors. Use [`save_and_cache`](Self::save_and_cache) for the
131    /// fire-and-forget variant that logs and continues on failure.
132    pub async fn save(&self, session: &mut Session) -> std::io::Result<()> {
133        self.persistence
134            .merge_save_runtime_and_publish(session, |saved, committed| {
135                if committed {
136                    #[cfg(test)]
137                    self.run_post_durable_hook("save_full", &saved.id);
138                    self.cache.insert(
139                        saved.id.clone(),
140                        Arc::new(parking_lot::RwLock::new(saved.clone())),
141                    );
142                }
143            })
144            .await
145    }
146
147    /// Atomically mutate the latest durable runtime session and refresh the
148    /// cache with the saved value. This is the safe path for narrow metadata
149    /// indexes that can be updated concurrently with runner message writes.
150    pub async fn update_runtime_session<F>(
151        &self,
152        session_id: &str,
153        metadata_keys: &[&str],
154        mutate: F,
155    ) -> std::io::Result<Option<Session>>
156    where
157        F: FnOnce(&mut Session),
158    {
159        self.persistence
160            .update_runtime_config_and_publish(session_id, mutate, |saved| {
161                if let Some(cached) = self.cache.get(session_id) {
162                    let mut cached = cached.write();
163                    for key in metadata_keys {
164                        if let Some(value) = saved.metadata.get(*key) {
165                            cached.metadata.insert((*key).to_string(), value.clone());
166                        } else {
167                            cached.metadata.remove(*key);
168                        }
169                    }
170                }
171            })
172            .await
173    }
174
175    /// Load a session, creating a fresh `Session::new(id, model)` if absent.
176    pub async fn load_or_create(&self, session_id: &str, model: &str) -> Session {
177        if let Some(session) = self.load(session_id).await {
178            return session;
179        }
180        Session::new(session_id.to_string(), model.to_string())
181    }
182
183    /// Load a session, reconciling the memory and storage copies via a
184    /// preference heuristic: storage wins when it is strictly newer, or when it
185    /// is the same age but still carries a pending question memory lost. Storage
186    /// is **never** preferred when it is strictly older than memory.
187    ///
188    /// The cache is refreshed cache-aside but with a no-regression guarantee:
189    /// `load_merged` never overwrites a newer cached session with an older
190    /// storage copy, so it is safe to call from hot read paths.
191    pub async fn load_merged(&self, session_id: &str) -> Option<Session> {
192        let _guard = self.persistence.acquire_lock(session_id).await;
193        let memory_session = read_cached_session(&self.cache, session_id);
194        let storage_session = self
195            .storage
196            .load_session(session_id)
197            .await
198            .unwrap_or_default();
199        #[cfg(test)]
200        self.run_post_durable_hook("load_merged", session_id);
201
202        match (memory_session, storage_session) {
203            (Some(memory), Some(storage)) => {
204                let prefer_storage = should_prefer_storage(&memory, &storage);
205                let diverged = prefer_storage || memory.messages.len() != storage.messages.len();
206                let chosen_len = if prefer_storage {
207                    storage.messages.len()
208                } else {
209                    memory.messages.len()
210                };
211                macro_rules! merged_log {
212                    ($level:ident) => {
213                        tracing::$level!(
214                            "[{}] load_session_merged: memory={} msgs (updated_at={}), storage={} msgs (updated_at={}), prefer_storage={} -> chose {} msgs",
215                            session_id,
216                            memory.messages.len(),
217                            memory.updated_at,
218                            storage.messages.len(),
219                            storage.updated_at,
220                            prefer_storage,
221                            chosen_len,
222                        )
223                    };
224                }
225                if diverged {
226                    merged_log!(debug);
227                } else {
228                    merged_log!(trace);
229                }
230                let memory_updated_at = memory.updated_at;
231                let chosen = if prefer_storage { storage } else { memory };
232                // Cache-aside refresh with a hard no-regression invariant: only
233                // write back when we actually reconciled *to storage* (a memory
234                // win is already the cached copy; re-inserting it would needlessly
235                // replace a possibly-live Arc) AND the reconciled copy is not
236                // older than what memory already holds. This is what makes
237                // `load_merged` safe on hot read paths — it can never clobber a
238                // freshly-updated session with a stale storage copy.
239                if prefer_storage && chosen.updated_at >= memory_updated_at {
240                    self.cache.insert(
241                        session_id.to_string(),
242                        Arc::new(parking_lot::RwLock::new(chosen.clone())),
243                    );
244                }
245                Some(chosen)
246            }
247            (Some(memory), None) => Some(memory),
248            (None, Some(storage)) => {
249                self.cache.insert(
250                    session_id.to_string(),
251                    Arc::new(parking_lot::RwLock::new(storage.clone())),
252                );
253                Some(storage)
254            }
255            (None, None) => None,
256        }
257    }
258
259    /// Persist the session (merge-on-write, preserving concurrent UI edits to
260    /// the authoritative metadata group) and refresh the in-memory cache.
261    pub async fn save_and_cache(&self, session: &mut Session) {
262        let result = self
263            .persistence
264            .merge_save_runtime_and_publish(session, |saved, _| {
265                #[cfg(test)]
266                self.run_post_durable_hook("save_and_cache", &saved.id);
267                self.cache.insert(
268                    saved.id.clone(),
269                    Arc::new(parking_lot::RwLock::new(saved.clone())),
270                );
271            })
272            .await;
273        if let Err(error) = result {
274            tracing::warn!("[{}] Failed to save session: {}", session.id, error);
275        }
276    }
277}
278
279fn should_prefer_storage(memory_session: &Session, storage_session: &Session) -> bool {
280    // Never reconcile *backwards* to a strictly-older storage copy: if memory is
281    // newer it is authoritative (e.g. it just answered and cleared a pending
282    // question while storage still holds the stale one). Respecting `updated_at`
283    // here is what stops `load_merged` from returning — and caching — stale data.
284    if storage_session.updated_at < memory_session.updated_at {
285        return false;
286    }
287    // Storage is same-age or newer: prefer it when strictly newer, or when it
288    // still carries a pending question that the (same-age) memory copy lost, so
289    // a genuine clarification is never dropped.
290    storage_session.updated_at > memory_session.updated_at
291        || (memory_session.pending_question.is_none() && storage_session.pending_question.is_some())
292}
293
294/// `SessionRepository` is the canonical `RuntimeSessionPersistence`: the runtime
295/// can persist a session through the same coordinator (merge-on-write + cache
296/// refresh) instead of a bespoke adapter.
297#[async_trait::async_trait]
298impl bamboo_domain::RuntimeSessionPersistence for SessionRepository {
299    async fn save_runtime_session(&self, session: &mut Session) -> std::io::Result<()> {
300        // Runtime authorization reads through this same cache. Refresh it even
301        // when durable storage fails so a current activation can never observe
302        // a previous run's skill allowlist. The error is still returned to the
303        // caller and durable state remains unchanged.
304        self.persistence
305            .merge_save_runtime_and_publish(session, |saved, _| {
306                #[cfg(test)]
307                self.run_post_durable_hook("save_runtime_session", &saved.id);
308                self.cache.insert(
309                    saved.id.clone(),
310                    Arc::new(parking_lot::RwLock::new(saved.clone())),
311                );
312            })
313            .await
314    }
315
316    async fn seed_runtime_activation(&self, session: &mut Session) -> std::io::Result<()> {
317        self.persistence
318            .seed_runtime_activation_and_publish(session, |saved, committed| {
319                #[cfg(test)]
320                self.run_post_durable_hook("seed_runtime_activation", &saved.id);
321                if committed {
322                    self.cache.insert(
323                        saved.id.clone(),
324                        Arc::new(parking_lot::RwLock::new(saved.clone())),
325                    );
326                }
327            })
328            .await
329    }
330
331    async fn record_permission_posture_activation(
332        &self,
333        session_id: &str,
334        expected_audit_revision: Option<u64>,
335        seed: &bamboo_domain::PermissionAuditSeed,
336    ) -> std::io::Result<Option<Session>> {
337        self.persistence
338            .record_permission_posture_activation_and_publish(
339                session_id,
340                expected_audit_revision,
341                seed,
342                |saved| {
343                    #[cfg(test)]
344                    self.run_post_durable_hook("permission_posture_activation", &saved.id);
345                    self.cache.insert(
346                        saved.id.clone(),
347                        Arc::new(parking_lot::RwLock::new(saved.clone())),
348                    );
349                },
350            )
351            .await
352    }
353
354    async fn save_runtime_control_plane(&self, session: &mut Session) -> std::io::Result<()> {
355        self.persistence
356            .save_runtime_only_and_publish(session, |saved| {
357                #[cfg(test)]
358                self.run_post_durable_hook("save", &saved.id);
359
360                // A control-plane snapshot may intentionally carry no messages
361                // (for example, child Task synchronization loads the root's
362                // runtime sidecar). Publish its fresh runtime fields without
363                // replacing a cache-resident transcript with that empty
364                // snapshot. SessionInbox admission is coupled to transcript
365                // persistence and is therefore preserved alongside the cached
366                // messages, matching the V2 sidecar overlay contract.
367                if let Some(cached) = self.cache.get(&saved.id) {
368                    let mut cached = cached.write();
369                    let messages = cached.messages.clone();
370                    let admission = cached
371                        .runtime_metadata
372                        .as_ref()
373                        .and_then(|metadata| metadata.session_inbox_admission.clone());
374                    let mut refreshed = saved.clone();
375                    refreshed.messages = messages;
376                    if let Some(admission) = admission {
377                        refreshed
378                            .runtime_metadata
379                            .get_or_insert_with(Default::default)
380                            .session_inbox_admission = Some(admission);
381                    } else if let Some(metadata) = refreshed.runtime_metadata.as_mut() {
382                        metadata.session_inbox_admission = None;
383                    }
384                    *cached = refreshed;
385                }
386            })
387            .await
388    }
389
390    async fn load_runtime_control_plane(
391        &self,
392        session_id: &str,
393    ) -> std::io::Result<Option<Session>> {
394        bamboo_domain::RuntimeSessionPersistence::load_runtime_control_plane(
395            self.persistence.as_ref(),
396            session_id,
397        )
398        .await
399    }
400
401    async fn update_task_list_control_plane(
402        &self,
403        session_id: &str,
404        task_list: &bamboo_domain::TaskList,
405        version: &str,
406    ) -> std::io::Result<bool> {
407        self.persistence
408            .update_task_list_control_plane_and_publish(session_id, task_list, version, |_| {
409                #[cfg(test)]
410                self.run_post_durable_hook("task", version);
411
412                // The durable transaction changed only Task-owned fields.
413                // Mirror that same narrow patch into the cache so a
414                // concurrent round/status/child transition already present
415                // in memory cannot be replaced by a stale whole-control-
416                // plane snapshot.
417                if let Some(cached) = self.cache.get(session_id) {
418                    let mut cached = cached.write();
419                    cached.set_task_list(task_list.clone());
420                    cached.set_task_list_version_meta(version.to_string());
421                }
422            })
423            .await
424    }
425
426    async fn checkpoint_runtime_session(&self, session: &mut Session) -> std::io::Result<()> {
427        // The execute-boundary checkpoint uses LockedSessionStore's atomic
428        // append-safe transcript merge, then publishes that reconciled snapshot
429        // to the runtime cache.  On failure, leave the existing cache alone: the
430        // checkpoint may have failed to load the latest durable transcript, and
431        // replacing a fresher cache entry with the stale runner snapshot would
432        // set up a later SHRINK write.
433        self.persistence
434            .checkpoint_runtime_session_and_publish(session, |saved, committed| {
435                #[cfg(test)]
436                self.run_post_durable_hook("checkpoint", &saved.id);
437                if committed {
438                    self.cache.insert(
439                        saved.id.clone(),
440                        Arc::new(parking_lot::RwLock::new(saved.clone())),
441                    );
442                }
443            })
444            .await
445    }
446
447    async fn load_runtime_session(&self, session_id: &str) -> std::io::Result<Option<Session>> {
448        self.try_load(session_id).await
449    }
450
451    async fn clear_legacy_pending_messages(
452        &self,
453        session_id: &str,
454        expected: &[serde_json::Value],
455    ) -> std::io::Result<bool> {
456        // Do not use the trait default here: `load_runtime_session` is
457        // cache-first, so a stale cache could erase a message concurrently
458        // appended to the durable legacy queue. Delegate the compare-and-clear
459        // to LockedSessionStore's single per-session critical section.
460        self.persistence
461            .clear_legacy_pending_messages_and_publish(session_id, expected, |latest| {
462                #[cfg(test)]
463                self.run_post_durable_hook("clear_legacy", session_id);
464                self.cache.insert(
465                    session_id.to_string(),
466                    Arc::new(parking_lot::RwLock::new(latest.clone())),
467                );
468            })
469            .await
470    }
471
472    async fn append_token_usage_record(
473        &self,
474        session_id: &str,
475        json_line: &str,
476    ) -> std::io::Result<()> {
477        self.storage
478            .append_token_usage_record(session_id, json_line)
479            .await
480    }
481}
482
483#[cfg(test)]
484mod tests {
485    use super::*;
486    use bamboo_agent_core::storage::Storage;
487    use chrono::Utc;
488    use std::collections::HashMap;
489    use std::sync::{Condvar, Mutex};
490    use std::time::Duration;
491
492    #[derive(Default)]
493    struct MapStorage {
494        sessions: Mutex<HashMap<String, Session>>,
495    }
496
497    struct FailingSaveStorage {
498        persisted: Mutex<Option<Session>>,
499    }
500
501    #[async_trait::async_trait]
502    impl Storage for MapStorage {
503        async fn save_session(&self, session: &Session) -> std::io::Result<()> {
504            self.sessions
505                .lock()
506                .unwrap()
507                .insert(session.id.clone(), session.clone());
508            Ok(())
509        }
510        async fn load_session(&self, session_id: &str) -> std::io::Result<Option<Session>> {
511            Ok(self.sessions.lock().unwrap().get(session_id).cloned())
512        }
513        async fn delete_session(&self, session_id: &str) -> std::io::Result<bool> {
514            Ok(self.sessions.lock().unwrap().remove(session_id).is_some())
515        }
516    }
517
518    #[async_trait::async_trait]
519    impl Storage for FailingSaveStorage {
520        async fn save_session(&self, _session: &Session) -> std::io::Result<()> {
521            Err(std::io::Error::other("injected save failure"))
522        }
523
524        async fn load_session(&self, _session_id: &str) -> std::io::Result<Option<Session>> {
525            Ok(self.persisted.lock().unwrap().clone())
526        }
527
528        async fn delete_session(&self, _session_id: &str) -> std::io::Result<bool> {
529            Ok(false)
530        }
531    }
532
533    fn test_repo(storage: Arc<dyn Storage>) -> SessionRepository {
534        let cache: SessionCache = Arc::new(dashmap::DashMap::new());
535        let persistence = Arc::new(LockedSessionStore::new(storage.clone()));
536        SessionRepository::new(cache, storage, persistence)
537    }
538
539    fn cache_put(repo: &SessionRepository, session: &Session) {
540        repo.cache().insert(
541            session.id.clone(),
542            Arc::new(parking_lot::RwLock::new(session.clone())),
543        );
544    }
545
546    fn task_list(session_id: &str, title: &str) -> bamboo_domain::TaskList {
547        let now = Utc::now();
548        bamboo_domain::TaskList {
549            session_id: session_id.to_string(),
550            title: title.to_string(),
551            items: Vec::new(),
552            created_at: now,
553            updated_at: now,
554        }
555    }
556
557    fn durable_cache_fence(
558        operation: impl Into<String>,
559        marker: impl Into<String>,
560    ) -> (
561        PostDurableHook,
562        tokio::sync::oneshot::Receiver<()>,
563        Arc<(Mutex<bool>, Condvar)>,
564    ) {
565        let operation = operation.into();
566        let marker = marker.into();
567        let (started_tx, started_rx) = tokio::sync::oneshot::channel();
568        let started_tx = Arc::new(Mutex::new(Some(started_tx)));
569        let release = Arc::new((Mutex::new(false), Condvar::new()));
570        let hook_release = release.clone();
571        let hook: PostDurableHook = Arc::new(move |actual_operation, actual_marker| {
572            if actual_operation != operation || actual_marker != marker {
573                return;
574            }
575            if let Some(started_tx) = started_tx.lock().unwrap().take() {
576                started_tx.send(()).expect("fence observer still present");
577            }
578            let (released, wake) = &*hook_release;
579            let mut released = released.lock().unwrap();
580            while !*released {
581                released = wake.wait(released).unwrap();
582            }
583        });
584        (hook, started_rx, release)
585    }
586
587    fn release_fence(release: &Arc<(Mutex<bool>, Condvar)>) {
588        let (released, wake) = &**release;
589        *released.lock().unwrap() = true;
590        wake.notify_all();
591    }
592
593    async fn assert_second_write_waits_for_cache_publish<T>(
594        first: tokio::task::JoinHandle<std::io::Result<T>>,
595        mut second: tokio::task::JoinHandle<std::io::Result<T>>,
596        release: Arc<(Mutex<bool>, Condvar)>,
597    ) -> (T, T) {
598        let second_before_release =
599            tokio::time::timeout(Duration::from_millis(100), &mut second).await;
600        let completed_before_release = second_before_release.is_ok();
601        release_fence(&release);
602
603        let first = first
604            .await
605            .expect("first writer joins")
606            .expect("first writer succeeds");
607        let second = match second_before_release {
608            Ok(joined) => joined
609                .expect("second writer joins")
610                .expect("second writer succeeds"),
611            Err(_) => second
612                .await
613                .expect("second writer joins")
614                .expect("second writer succeeds"),
615        };
616        assert!(
617            !completed_before_release,
618            "the second write must remain behind the first write's durable-to-cache fence"
619        );
620        (first, second)
621    }
622
623    #[derive(Clone, Copy, Debug)]
624    enum FullSaveRoute {
625        InherentSave,
626        SaveAndCache,
627        RuntimePersistence,
628    }
629
630    impl FullSaveRoute {
631        fn operation(self) -> &'static str {
632            match self {
633                Self::InherentSave => "save_full",
634                Self::SaveAndCache => "save_and_cache",
635                Self::RuntimePersistence => "save_runtime_session",
636            }
637        }
638
639        fn name(self) -> &'static str {
640            match self {
641                Self::InherentSave => "inherent",
642                Self::SaveAndCache => "save-and-cache",
643                Self::RuntimePersistence => "runtime-persistence",
644            }
645        }
646    }
647
648    async fn assert_full_save_route_serializes_cache_publish(route: FullSaveRoute) {
649        let temp = tempfile::tempdir().unwrap();
650        let concrete_storage = Arc::new(
651            bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
652                .await
653                .expect("SessionStoreV2"),
654        );
655        let storage: Arc<dyn Storage> = concrete_storage;
656        let id = format!("root-full-cache-order-{}", route.name());
657        let (hook, first_durable, release) = durable_cache_fence(route.operation(), id.clone());
658        let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
659
660        let mut initial = Session::new(&id, "model");
661        initial.add_message(bamboo_agent_core::Message::user("durable transcript"));
662        initial.set_task_list(task_list(&id, "initial"));
663        initial.set_task_list_version_meta("0");
664        initial
665            .metadata
666            .insert("unrelated.runtime".to_string(), "keep".to_string());
667        storage.save_session(&initial).await.unwrap();
668        cache_put(&repo, &initial);
669
670        let first_repo = repo.clone();
671        let mut root_snapshot = initial.clone();
672        root_snapshot.add_message(bamboo_agent_core::Message::assistant(
673            "full-save transcript suffix",
674            None,
675        ));
676        root_snapshot.set_task_list(task_list(&id, "root"));
677        root_snapshot.set_task_list_version_meta("1");
678        let first = tokio::spawn(async move {
679            match route {
680                FullSaveRoute::InherentSave => first_repo.save(&mut root_snapshot).await,
681                FullSaveRoute::SaveAndCache => {
682                    first_repo.save_and_cache(&mut root_snapshot).await;
683                    Ok(())
684                }
685                FullSaveRoute::RuntimePersistence => {
686                    bamboo_domain::RuntimeSessionPersistence::save_runtime_session(
687                        first_repo.as_ref(),
688                        &mut root_snapshot,
689                    )
690                    .await
691                }
692            }
693        });
694        first_durable
695            .await
696            .expect("root full durable write reached");
697
698        let second_repo = repo.clone();
699        let second_id = id.clone();
700        let child_task_list = task_list(&id, "child");
701        let second = tokio::spawn(async move {
702            bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
703                second_repo.as_ref(),
704                &second_id,
705                &child_task_list,
706                "2",
707            )
708            .await
709            .map(|updated| assert!(updated, "root must exist"))
710        });
711        assert_second_write_waits_for_cache_publish(first, second, release).await;
712
713        let durable = storage.load_session(&id).await.unwrap().unwrap();
714        let cached = read_cached_session(repo.cache(), &id).expect("cached root");
715        for (tier, session) in [("durable", durable), ("cache", cached)] {
716            assert_eq!(
717                session.task_list_version_meta().as_deref(),
718                Some("2"),
719                "{route:?} {tier} must retain the child transaction"
720            );
721            assert_eq!(
722                session.task_list.as_ref().map(|list| list.title.as_str()),
723                Some("child"),
724                "{route:?} {tier} must retain the child transaction"
725            );
726            assert_eq!(
727                session
728                    .metadata
729                    .get("unrelated.runtime")
730                    .map(String::as_str),
731                Some("keep"),
732                "{route:?} {tier} must preserve unrelated runtime state"
733            );
734            assert_eq!(
735                session.messages.len(),
736                2,
737                "{route:?} {tier} must preserve the full-save transcript"
738            );
739        }
740    }
741
742    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
743    async fn root_full_saves_and_child_task_patch_share_publish_order() {
744        for route in [
745            FullSaveRoute::InherentSave,
746            FullSaveRoute::SaveAndCache,
747            FullSaveRoute::RuntimePersistence,
748        ] {
749            assert_full_save_route_serializes_cache_publish(route).await;
750        }
751    }
752
753    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
754    async fn checkpoint_and_child_task_patch_share_publish_order() {
755        let temp = tempfile::tempdir().unwrap();
756        let concrete_storage = Arc::new(
757            bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
758                .await
759                .expect("SessionStoreV2"),
760        );
761        let storage: Arc<dyn Storage> = concrete_storage;
762        let id = "checkpoint-cache-order";
763        let (hook, checkpoint_durable, release) = durable_cache_fence("checkpoint", id);
764        let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
765
766        let mut initial = Session::new(id, "model");
767        initial.add_message(bamboo_agent_core::Message::user("durable transcript"));
768        initial.set_task_list(task_list(id, "initial"));
769        initial.set_task_list_version_meta("0");
770        initial
771            .metadata
772            .insert("unrelated.runtime".to_string(), "keep".to_string());
773        storage.save_session(&initial).await.unwrap();
774        cache_put(&repo, &initial);
775
776        let checkpoint_repo = repo.clone();
777        let mut checkpoint_snapshot = initial.clone();
778        checkpoint_snapshot.add_message(bamboo_agent_core::Message::assistant(
779            "checkpoint transcript suffix",
780            None,
781        ));
782        checkpoint_snapshot.set_task_list(task_list(id, "checkpoint"));
783        checkpoint_snapshot.set_task_list_version_meta("1");
784        let checkpoint = tokio::spawn(async move {
785            bamboo_domain::RuntimeSessionPersistence::checkpoint_runtime_session(
786                checkpoint_repo.as_ref(),
787                &mut checkpoint_snapshot,
788            )
789            .await
790        });
791        checkpoint_durable
792            .await
793            .expect("checkpoint durable write reached");
794
795        let child_repo = repo.clone();
796        let child_task_list = task_list(id, "child");
797        let child_patch = tokio::spawn(async move {
798            bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
799                child_repo.as_ref(),
800                id,
801                &child_task_list,
802                "2",
803            )
804            .await
805            .map(|updated| assert!(updated, "root must exist"))
806        });
807        assert_second_write_waits_for_cache_publish(checkpoint, child_patch, release).await;
808
809        let durable = storage.load_session(id).await.unwrap().unwrap();
810        let cached = read_cached_session(repo.cache(), id).expect("cached root");
811        for (tier, session) in [("durable", durable), ("cache", cached)] {
812            assert_eq!(
813                session.task_list_version_meta().as_deref(),
814                Some("2"),
815                "{tier} must retain the child transaction"
816            );
817            assert_eq!(
818                session.task_list.as_ref().map(|list| list.title.as_str()),
819                Some("child"),
820                "{tier} must retain the child transaction"
821            );
822            assert_eq!(
823                session
824                    .metadata
825                    .get("unrelated.runtime")
826                    .map(String::as_str),
827                Some("keep"),
828                "{tier} must preserve unrelated runtime state"
829            );
830            assert_eq!(
831                session.messages.len(),
832                2,
833                "{tier} must preserve the checkpoint transcript"
834            );
835        }
836    }
837
838    #[derive(Clone, Copy, Debug)]
839    enum CacheBackfillRoute {
840        Load,
841        TryLoad,
842    }
843
844    impl CacheBackfillRoute {
845        fn operation(self) -> &'static str {
846            match self {
847                Self::Load => "load",
848                Self::TryLoad => "try_load",
849            }
850        }
851
852        fn name(self) -> &'static str {
853            match self {
854                Self::Load => "load",
855                Self::TryLoad => "try-load",
856            }
857        }
858    }
859
860    async fn assert_cache_backfill_serializes_with_task_patch(route: CacheBackfillRoute) {
861        let temp = tempfile::tempdir().unwrap();
862        let concrete_storage = Arc::new(
863            bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
864                .await
865                .expect("SessionStoreV2"),
866        );
867        let storage: Arc<dyn Storage> = concrete_storage;
868        let id = format!("cache-backfill-order-{}", route.name());
869        let (hook, loaded_old_durable, release) =
870            durable_cache_fence(route.operation(), id.clone());
871        let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
872
873        let mut initial = Session::new(&id, "model");
874        initial.set_task_list(task_list(&id, "initial"));
875        initial.set_task_list_version_meta("0");
876        storage.save_session(&initial).await.unwrap();
877        assert!(
878            read_cached_session(repo.cache(), &id).is_none(),
879            "the race requires a genuine cache miss"
880        );
881
882        let load_repo = repo.clone();
883        let load_id = id.clone();
884        let load = tokio::spawn(async move {
885            let loaded = match route {
886                CacheBackfillRoute::Load => load_repo.load(&load_id).await,
887                CacheBackfillRoute::TryLoad => {
888                    load_repo.try_load(&load_id).await.expect("storage load")
889                }
890            };
891            assert!(loaded.is_some(), "seeded session must load");
892            Ok(())
893        });
894        loaded_old_durable
895            .await
896            .expect("old durable snapshot loaded");
897
898        let patch_repo = repo.clone();
899        let patch_id = id.clone();
900        let child_task_list = task_list(&id, "child");
901        let patch = tokio::spawn(async move {
902            bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
903                patch_repo.as_ref(),
904                &patch_id,
905                &child_task_list,
906                "1",
907            )
908            .await
909            .map(|updated| assert!(updated, "root must exist"))
910        });
911        assert_second_write_waits_for_cache_publish(load, patch, release).await;
912
913        let durable = storage.load_session(&id).await.unwrap().unwrap();
914        let cached = read_cached_session(repo.cache(), &id).expect("backfilled cache");
915        for (tier, session) in [("durable", durable), ("cache", cached)] {
916            assert_eq!(
917                session.task_list_version_meta().as_deref(),
918                Some("1"),
919                "{route:?} {tier} must retain the child transaction"
920            );
921            assert_eq!(
922                session.task_list.as_ref().map(|list| list.title.as_str()),
923                Some("child"),
924                "{route:?} {tier} must retain the child transaction"
925            );
926        }
927    }
928
929    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
930    async fn cache_miss_backfills_and_child_task_patch_share_publish_order() {
931        for route in [CacheBackfillRoute::Load, CacheBackfillRoute::TryLoad] {
932            assert_cache_backfill_serializes_with_task_patch(route).await;
933        }
934    }
935
936    #[tokio::test]
937    async fn cache_hits_do_not_wait_for_the_persistence_lock() {
938        let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
939        let repo = test_repo(storage);
940        let id = "cache-hit-lock-free";
941        let cached = Session::new(id, "cached-model");
942        cache_put(&repo, &cached);
943        let persistence_guard = repo.persistence().acquire_lock(id).await;
944
945        let loaded = tokio::time::timeout(Duration::from_millis(100), repo.load(id)).await;
946        let try_loaded = tokio::time::timeout(Duration::from_millis(100), repo.try_load(id)).await;
947        drop(persistence_guard);
948
949        assert_eq!(
950            loaded
951                .expect("cache hit must not wait for the persistence lock")
952                .expect("cached session")
953                .model,
954            "cached-model"
955        );
956        assert_eq!(
957            try_loaded
958                .expect("fallible cache hit must not wait for the persistence lock")
959                .expect("cache read succeeds")
960                .expect("cached session")
961                .model,
962            "cached-model"
963        );
964    }
965
966    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
967    async fn load_merged_storage_refresh_and_child_task_patch_share_publish_order() {
968        let temp = tempfile::tempdir().unwrap();
969        let concrete_storage = Arc::new(
970            bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
971                .await
972                .expect("SessionStoreV2"),
973        );
974        let storage: Arc<dyn Storage> = concrete_storage;
975        let id = "load-merged-cache-order";
976        let (hook, loaded_old_durable, release) = durable_cache_fence("load_merged", id);
977        let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
978
979        let mut durable = Session::new(id, "model");
980        durable.updated_at = Utc::now();
981        durable.set_task_list(task_list(id, "initial"));
982        durable.set_task_list_version_meta("0");
983        storage.save_session(&durable).await.unwrap();
984
985        let mut memory = durable.clone();
986        memory.updated_at = durable.updated_at - chrono::Duration::seconds(1);
987        memory.set_task_list(task_list(id, "memory"));
988        cache_put(&repo, &memory);
989
990        let load_repo = repo.clone();
991        let load = tokio::spawn(async move {
992            assert!(
993                load_repo.load_merged(id).await.is_some(),
994                "seeded session must load"
995            );
996            Ok(())
997        });
998        loaded_old_durable
999            .await
1000            .expect("old durable snapshot loaded");
1001
1002        let patch_repo = repo.clone();
1003        let child_task_list = task_list(id, "child");
1004        let patch = tokio::spawn(async move {
1005            bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
1006                patch_repo.as_ref(),
1007                id,
1008                &child_task_list,
1009                "1",
1010            )
1011            .await
1012            .map(|updated| assert!(updated, "root must exist"))
1013        });
1014        assert_second_write_waits_for_cache_publish(load, patch, release).await;
1015
1016        let durable = storage.load_session(id).await.unwrap().unwrap();
1017        let cached = read_cached_session(repo.cache(), id).expect("refreshed cache");
1018        for (tier, session) in [("durable", durable), ("cache", cached)] {
1019            assert_eq!(
1020                session.task_list_version_meta().as_deref(),
1021                Some("1"),
1022                "{tier} must retain the child transaction"
1023            );
1024            assert_eq!(
1025                session.task_list.as_ref().map(|list| list.title.as_str()),
1026                Some("child"),
1027                "{tier} must retain the child transaction"
1028            );
1029        }
1030    }
1031
1032    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1033    async fn legacy_clear_refresh_and_child_task_patch_share_publish_order() {
1034        let temp = tempfile::tempdir().unwrap();
1035        let concrete_storage = Arc::new(
1036            bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
1037                .await
1038                .expect("SessionStoreV2"),
1039        );
1040        let storage: Arc<dyn Storage> = concrete_storage;
1041        let id = "legacy-clear-cache-order";
1042        let expected = vec![serde_json::json!({"content": "legacy"})];
1043        let (hook, loaded_post_cas_snapshot, release) = durable_cache_fence("clear_legacy", id);
1044        let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
1045
1046        let mut initial = Session::new(id, "model");
1047        initial.set_pending_injected_messages(expected.clone());
1048        initial.set_task_list(task_list(id, "initial"));
1049        initial.set_task_list_version_meta("0");
1050        storage.save_session(&initial).await.unwrap();
1051        cache_put(&repo, &initial);
1052
1053        let clear_repo = repo.clone();
1054        let clear_expected = expected.clone();
1055        let clear = tokio::spawn(async move {
1056            bamboo_domain::RuntimeSessionPersistence::clear_legacy_pending_messages(
1057                clear_repo.as_ref(),
1058                id,
1059                &clear_expected,
1060            )
1061            .await
1062            .map(|cleared| assert!(cleared, "legacy queue must match"))
1063        });
1064        loaded_post_cas_snapshot
1065            .await
1066            .expect("post-CAS snapshot loaded");
1067
1068        let patch_repo = repo.clone();
1069        let child_task_list = task_list(id, "child");
1070        let patch = tokio::spawn(async move {
1071            bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
1072                patch_repo.as_ref(),
1073                id,
1074                &child_task_list,
1075                "1",
1076            )
1077            .await
1078            .map(|updated| assert!(updated, "root must exist"))
1079        });
1080        assert_second_write_waits_for_cache_publish(clear, patch, release).await;
1081
1082        let durable = storage.load_session(id).await.unwrap().unwrap();
1083        let cached = read_cached_session(repo.cache(), id).expect("refreshed cache");
1084        for (tier, session) in [("durable", durable), ("cache", cached)] {
1085            assert_eq!(
1086                session.task_list_version_meta().as_deref(),
1087                Some("1"),
1088                "{tier} must retain the child transaction"
1089            );
1090            assert_eq!(
1091                session.task_list.as_ref().map(|list| list.title.as_str()),
1092                Some("child"),
1093                "{tier} must retain the child transaction"
1094            );
1095            assert!(
1096                !session.has_pending_injected_messages(),
1097                "{tier} must retain the successful legacy clear"
1098            );
1099        }
1100    }
1101
1102    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1103    async fn concurrent_task_patches_publish_cache_in_durable_order() {
1104        let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
1105        let (hook, first_durable, release) = durable_cache_fence("task", "1");
1106        let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
1107        let id = "concurrent-task-cache-order";
1108        let mut initial = Session::new(id, "model");
1109        initial.set_task_list(task_list(id, "initial"));
1110        initial.set_task_list_version_meta("0");
1111        storage.save_session(&initial).await.unwrap();
1112        cache_put(&repo, &initial);
1113
1114        let first_repo = repo.clone();
1115        let first_task_list = task_list(id, "first");
1116        let first = tokio::spawn(async move {
1117            bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
1118                first_repo.as_ref(),
1119                id,
1120                &first_task_list,
1121                "1",
1122            )
1123            .await
1124        });
1125        first_durable.await.expect("first durable write reached");
1126
1127        let second_repo = repo.clone();
1128        let second_task_list = task_list(id, "second");
1129        let second = tokio::spawn(async move {
1130            bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
1131                second_repo.as_ref(),
1132                id,
1133                &second_task_list,
1134                "2",
1135            )
1136            .await
1137        });
1138        let (first_updated, second_updated) =
1139            assert_second_write_waits_for_cache_publish(first, second, release).await;
1140        assert!(first_updated && second_updated);
1141
1142        let durable = storage.load_session(id).await.unwrap().unwrap();
1143        let cached = read_cached_session(repo.cache(), id).expect("cached root");
1144        for (tier, session) in [("durable", durable), ("cache", cached)] {
1145            assert_eq!(
1146                session.task_list_version_meta().as_deref(),
1147                Some("2"),
1148                "{tier} must retain the second transaction"
1149            );
1150            assert_eq!(
1151                session.task_list.as_ref().map(|list| list.title.as_str()),
1152                Some("second"),
1153                "{tier} must retain the second transaction"
1154            );
1155        }
1156    }
1157
1158    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1159    async fn root_control_plane_save_and_child_task_patch_share_publish_order() {
1160        let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
1161        let (hook, root_durable, release) = durable_cache_fence("save", "root-control-plane-order");
1162        let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
1163        let id = "root-control-plane-order";
1164
1165        let mut initial = Session::new(id, "model");
1166        initial.add_message(bamboo_agent_core::Message::user("durable transcript"));
1167        initial.set_task_list(task_list(id, "initial"));
1168        initial.set_task_list_version_meta("0");
1169        initial
1170            .metadata
1171            .insert("unrelated.runtime".to_string(), "keep".to_string());
1172        storage.save_session(&initial).await.unwrap();
1173        cache_put(&repo, &initial);
1174
1175        let root_repo = repo.clone();
1176        let mut root_snapshot = initial.clone();
1177        root_snapshot.set_task_list(task_list(id, "root"));
1178        root_snapshot.set_task_list_version_meta("1");
1179        let root_save = tokio::spawn(async move {
1180            bamboo_domain::RuntimeSessionPersistence::save_runtime_control_plane(
1181                root_repo.as_ref(),
1182                &mut root_snapshot,
1183            )
1184            .await
1185        });
1186        root_durable.await.expect("root durable write reached");
1187
1188        let child_repo = repo.clone();
1189        let child_task_list = task_list(id, "child");
1190        let child_patch = tokio::spawn(async move {
1191            bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
1192                child_repo.as_ref(),
1193                id,
1194                &child_task_list,
1195                "2",
1196            )
1197            .await
1198            .map(|updated| {
1199                assert!(updated, "root must exist");
1200            })
1201        });
1202        assert_second_write_waits_for_cache_publish(root_save, child_patch, release).await;
1203
1204        let durable = storage.load_session(id).await.unwrap().unwrap();
1205        let cached = read_cached_session(repo.cache(), id).expect("cached root");
1206        for (tier, session) in [("durable", durable), ("cache", cached)] {
1207            assert_eq!(
1208                session.task_list_version_meta().as_deref(),
1209                Some("2"),
1210                "{tier} must retain the child transaction"
1211            );
1212            assert_eq!(
1213                session.task_list.as_ref().map(|list| list.title.as_str()),
1214                Some("child"),
1215                "{tier} must retain the child transaction"
1216            );
1217            assert_eq!(
1218                session
1219                    .metadata
1220                    .get("unrelated.runtime")
1221                    .map(String::as_str),
1222                Some("keep"),
1223                "{tier} must preserve unrelated runtime state"
1224            );
1225            assert_eq!(
1226                session.messages.len(),
1227                1,
1228                "{tier} must preserve the transcript"
1229            );
1230        }
1231    }
1232
1233    #[tokio::test]
1234    async fn narrow_runtime_metadata_transaction_preserves_live_and_durable_non_owned_state() {
1235        let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
1236        let repo = test_repo(storage.clone());
1237        let id = "narrow-metadata";
1238        let mut durable = Session::new(id, "durable-model");
1239        durable.add_message(bamboo_agent_core::Message::user("durable user turn"));
1240        durable
1241            .metadata
1242            .insert("external.durable".to_string(), "keep".to_string());
1243        storage.save_session(&durable).await.expect("seed durable");
1244
1245        let mut live = durable.clone();
1246        live.add_message(bamboo_agent_core::Message::assistant(
1247            "in-flight assistant tool call",
1248            None,
1249        ));
1250        live.model = "live-model".to_string();
1251        live.metadata
1252            .insert("external.live".to_string(), "keep".to_string());
1253        cache_put(&repo, &live);
1254
1255        repo.update_runtime_session(id, &["workflow.owned"], |latest| {
1256            latest
1257                .metadata
1258                .insert("workflow.owned".to_string(), "active".to_string());
1259        })
1260        .await
1261        .expect("transaction")
1262        .expect("session exists");
1263
1264        let saved = storage
1265            .load_session(id)
1266            .await
1267            .expect("load durable")
1268            .expect("durable exists");
1269        assert_eq!(
1270            saved.messages.len(),
1271            1,
1272            "transaction never writes stale live messages"
1273        );
1274        assert_eq!(
1275            saved.metadata.get("external.durable").map(String::as_str),
1276            Some("keep")
1277        );
1278        assert_eq!(
1279            saved.metadata.get("workflow.owned").map(String::as_str),
1280            Some("active")
1281        );
1282
1283        let cached = read_cached_session(repo.cache(), id).expect("live cache");
1284        assert_eq!(
1285            cached.messages.len(),
1286            2,
1287            "cache live tool call is not replaced"
1288        );
1289        assert_eq!(cached.model, "live-model");
1290        assert_eq!(
1291            cached.metadata.get("external.live").map(String::as_str),
1292            Some("keep")
1293        );
1294        assert_eq!(
1295            cached.metadata.get("workflow.owned").map(String::as_str),
1296            Some("active")
1297        );
1298    }
1299
1300    /// Regression guard: a strictly-newer in-memory session (e.g. one that just
1301    /// answered and cleared its pending question) must win over a strictly-older
1302    /// storage copy that still carries the pending question — both in the value
1303    /// returned AND in the cache (no clobber).
1304    #[tokio::test]
1305    async fn load_merged_does_not_regress_to_older_storage() {
1306        let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
1307        let repo = test_repo(storage.clone());
1308        let id = "s1";
1309
1310        let mut stale = Session::new(id.to_string(), "m");
1311        stale.set_pending_question(
1312            "tc1".into(),
1313            "kind".into(),
1314            "q?".into(),
1315            vec!["OK".into()],
1316            true,
1317        );
1318        stale.updated_at = Utc::now() - chrono::Duration::seconds(10);
1319        storage.save_session(&stale).await.unwrap();
1320
1321        let mut fresh = Session::new(id.to_string(), "m");
1322        fresh.updated_at = Utc::now();
1323        cache_put(&repo, &fresh);
1324
1325        let merged = repo.load_merged(id).await.expect("session exists");
1326        assert!(
1327            merged.pending_question.is_none(),
1328            "must return the newer answered memory copy, not the stale storage one"
1329        );
1330        let cached = read_cached_session(repo.cache(), id).expect("cached");
1331        assert!(
1332            cached.pending_question.is_none(),
1333            "load_merged must never regress the cache to a stale storage copy"
1334        );
1335    }
1336
1337    /// The pending-question recovery still works when storage is the same age:
1338    /// if memory lost a pending question that same-age storage retains, prefer
1339    /// storage so a genuine clarification is not dropped.
1340    #[tokio::test]
1341    async fn load_merged_recovers_pending_question_from_same_age_storage() {
1342        let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
1343        let repo = test_repo(storage.clone());
1344        let id = "s2";
1345        let ts = Utc::now();
1346
1347        let mut with_pending = Session::new(id.to_string(), "m");
1348        with_pending.set_pending_question(
1349            "tc".into(),
1350            "k".into(),
1351            "q".into(),
1352            vec!["OK".into()],
1353            true,
1354        );
1355        with_pending.updated_at = ts;
1356        storage.save_session(&with_pending).await.unwrap();
1357
1358        let mut lost = with_pending.clone();
1359        lost.clear_pending_question();
1360        lost.updated_at = ts;
1361        cache_put(&repo, &lost);
1362
1363        let merged = repo.load_merged(id).await.expect("session exists");
1364        assert!(
1365            merged.pending_question.is_some(),
1366            "same-age storage carrying a pending question must still be recovered"
1367        );
1368    }
1369
1370    #[tokio::test]
1371    async fn runtime_publish_refreshes_cache_even_when_storage_fails() {
1372        let id = "runtime-selection";
1373        let mut previous = Session::new(id.to_string(), "m");
1374        previous.metadata.insert(
1375            "skill_runtime_selected_skill_ids".to_string(),
1376            "[\"plan\"]".to_string(),
1377        );
1378        let storage: Arc<dyn Storage> = Arc::new(FailingSaveStorage {
1379            persisted: Mutex::new(Some(previous.clone())),
1380        });
1381        let repo = test_repo(storage.clone());
1382        cache_put(&repo, &previous);
1383
1384        let mut current = previous.clone();
1385        current.metadata.insert(
1386            "skill_runtime_selected_skill_ids".to_string(),
1387            "[\"review\"]".to_string(),
1388        );
1389        current.updated_at = Utc::now();
1390
1391        let result =
1392            bamboo_domain::RuntimeSessionPersistence::save_runtime_session(&repo, &mut current)
1393                .await;
1394        assert!(result.is_err(), "durable failure must still be surfaced");
1395
1396        let cached = repo.load(id).await.expect("cached current session");
1397        assert_eq!(
1398            cached
1399                .metadata
1400                .get("skill_runtime_selected_skill_ids")
1401                .map(String::as_str),
1402            Some("[\"review\"]")
1403        );
1404        let allowlist = bamboo_skills::access_control::extract_skill_allowlist(&cached.metadata)
1405            .expect("runtime authorization allowlist");
1406        assert!(allowlist.contains("review"));
1407        assert!(!allowlist.contains("plan"));
1408        let durable = storage
1409            .load_session(id)
1410            .await
1411            .expect("load durable state")
1412            .expect("previous durable session");
1413        assert_eq!(
1414            durable
1415                .metadata
1416                .get("skill_runtime_selected_skill_ids")
1417                .map(String::as_str),
1418            Some("[\"plan\"]")
1419        );
1420    }
1421
1422    #[tokio::test]
1423    async fn inherent_save_leaves_existing_cache_untouched_when_storage_fails() {
1424        let id = "inherent-save-failure";
1425        let previous = Session::new(id, "previous");
1426        let storage: Arc<dyn Storage> = Arc::new(FailingSaveStorage {
1427            persisted: Mutex::new(Some(previous.clone())),
1428        });
1429        let repo = test_repo(storage);
1430        cache_put(&repo, &previous);
1431
1432        let mut current = previous.clone();
1433        current.model = "current".to_string();
1434        assert!(repo.save(&mut current).await.is_err());
1435        assert_eq!(
1436            read_cached_session(repo.cache(), id)
1437                .expect("existing cache")
1438                .model,
1439            "previous",
1440            "fallible inherent save must publish only after a durable commit"
1441        );
1442    }
1443
1444    #[tokio::test]
1445    async fn save_and_cache_still_refreshes_cache_when_storage_fails() {
1446        let id = "save-and-cache-failure";
1447        let previous = Session::new(id, "previous");
1448        let storage: Arc<dyn Storage> = Arc::new(FailingSaveStorage {
1449            persisted: Mutex::new(Some(previous.clone())),
1450        });
1451        let repo = test_repo(storage);
1452        cache_put(&repo, &previous);
1453
1454        let mut current = previous;
1455        current.model = "current".to_string();
1456        repo.save_and_cache(&mut current).await;
1457        assert_eq!(
1458            read_cached_session(repo.cache(), id)
1459                .expect("refreshed cache")
1460                .model,
1461            "current",
1462            "fire-and-forget save must retain its existing cache-on-failure behavior"
1463        );
1464    }
1465
1466    #[tokio::test]
1467    async fn checkpoint_leaves_existing_cache_untouched_when_storage_fails() {
1468        let id = "checkpoint-failure";
1469        let previous = Session::new(id, "previous");
1470        let storage: Arc<dyn Storage> = Arc::new(FailingSaveStorage {
1471            persisted: Mutex::new(Some(previous.clone())),
1472        });
1473        let repo = test_repo(storage);
1474        cache_put(&repo, &previous);
1475
1476        let mut current = previous.clone();
1477        current.model = "current".to_string();
1478        let result = bamboo_domain::RuntimeSessionPersistence::checkpoint_runtime_session(
1479            &repo,
1480            &mut current,
1481        )
1482        .await;
1483        assert!(result.is_err());
1484        assert_eq!(
1485            read_cached_session(repo.cache(), id)
1486                .expect("existing cache")
1487                .model,
1488            "previous",
1489            "checkpoint must publish only after a durable commit"
1490        );
1491    }
1492
1493    #[tokio::test]
1494    async fn legacy_clear_uses_durable_cas_and_never_erases_concurrent_append_from_stale_cache() {
1495        let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
1496        let repo = test_repo(storage.clone());
1497        let id = "legacy-cas-race";
1498        let expected = vec![serde_json::json!({"content": "first"})];
1499
1500        let mut stale_cache = Session::new(id, "m");
1501        stale_cache.set_pending_injected_messages(expected.clone());
1502        cache_put(&repo, &stale_cache);
1503
1504        let mut durable = stale_cache.clone();
1505        durable.set_pending_injected_messages(vec![
1506            serde_json::json!({"content": "first"}),
1507            serde_json::json!({"content": "concurrent"}),
1508        ]);
1509        storage.save_session(&durable).await.unwrap();
1510
1511        let cleared = bamboo_domain::RuntimeSessionPersistence::clear_legacy_pending_messages(
1512            &repo, id, &expected,
1513        )
1514        .await
1515        .unwrap();
1516        assert!(!cleared, "the durable compare-and-clear must reject drift");
1517        assert_eq!(
1518            storage
1519                .load_session(id)
1520                .await
1521                .unwrap()
1522                .unwrap()
1523                .pending_injected_messages()
1524                .unwrap(),
1525            durable.pending_injected_messages().unwrap(),
1526            "the concurrent durable append must remain intact"
1527        );
1528        assert_eq!(
1529            read_cached_session(repo.cache(), id)
1530                .unwrap()
1531                .pending_injected_messages()
1532                .unwrap(),
1533            expected,
1534            "a failed CAS must not mutate the existing cache"
1535        );
1536    }
1537
1538    #[tokio::test]
1539    async fn successful_legacy_clear_refreshes_stale_cache_from_durable_state() {
1540        let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
1541        let repo = test_repo(storage.clone());
1542        let id = "legacy-cas-success";
1543        let expected = vec![serde_json::json!({"content": "first"})];
1544
1545        let mut stale_cache = Session::new(id, "stale-model");
1546        stale_cache.set_pending_injected_messages(expected.clone());
1547        cache_put(&repo, &stale_cache);
1548
1549        let mut durable = Session::new(id, "durable-model");
1550        durable.set_pending_injected_messages(expected.clone());
1551        durable
1552            .metadata
1553            .insert("durable-only".to_string(), "keep".to_string());
1554        storage.save_session(&durable).await.unwrap();
1555
1556        assert!(
1557            bamboo_domain::RuntimeSessionPersistence::clear_legacy_pending_messages(
1558                &repo, id, &expected,
1559            )
1560            .await
1561            .unwrap()
1562        );
1563        let cached = read_cached_session(repo.cache(), id).unwrap();
1564        assert!(!cached.has_pending_injected_messages());
1565        assert_eq!(cached.model, "durable-model");
1566        assert_eq!(
1567            cached.metadata.get("durable-only").map(String::as_str),
1568            Some("keep")
1569        );
1570    }
1571}