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(crate::SessionSnapshot::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(crate::SessionSnapshot::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(crate::SessionSnapshot::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                    cached.update(|cached| {
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            })
173            .await
174    }
175
176    /// Load a session, creating a fresh `Session::new(id, model)` if absent.
177    pub async fn load_or_create(&self, session_id: &str, model: &str) -> Session {
178        if let Some(session) = self.load(session_id).await {
179            return session;
180        }
181        Session::new(session_id.to_string(), model.to_string())
182    }
183
184    /// Load a session, reconciling the memory and storage copies via a
185    /// preference heuristic: storage wins when it is strictly newer, or when it
186    /// is the same age but still carries a pending question memory lost. Storage
187    /// is **never** preferred when it is strictly older than memory.
188    ///
189    /// The cache is refreshed cache-aside but with a no-regression guarantee:
190    /// `load_merged` never overwrites a newer cached session with an older
191    /// storage copy, so it is safe to call from hot read paths.
192    pub async fn load_merged_checked(&self, session_id: &str) -> std::io::Result<Option<Session>> {
193        let _guard = self.persistence.acquire_lock(session_id).await;
194        let memory_session = read_cached_session(&self.cache, session_id);
195        let storage_session = self.storage.load_session(session_id).await?;
196        #[cfg(test)]
197        self.run_post_durable_hook("load_merged", session_id);
198
199        Ok(match (memory_session, storage_session) {
200            (Some(memory), Some(storage)) => {
201                let prefer_storage = should_prefer_storage(&memory, &storage);
202                let diverged = prefer_storage || memory.messages.len() != storage.messages.len();
203                let chosen_len = if prefer_storage {
204                    storage.messages.len()
205                } else {
206                    memory.messages.len()
207                };
208                macro_rules! merged_log {
209                    ($level:ident) => {
210                        tracing::$level!(
211                            "[{}] load_session_merged: memory={} msgs (updated_at={}), storage={} msgs (updated_at={}), prefer_storage={} -> chose {} msgs",
212                            session_id,
213                            memory.messages.len(),
214                            memory.updated_at,
215                            storage.messages.len(),
216                            storage.updated_at,
217                            prefer_storage,
218                            chosen_len,
219                        )
220                    };
221                }
222                if diverged {
223                    merged_log!(debug);
224                } else {
225                    merged_log!(trace);
226                }
227                let memory_updated_at = memory.updated_at;
228                let chosen = if prefer_storage { storage } else { memory };
229                // Cache-aside refresh with a hard no-regression invariant: only
230                // write back when we actually reconciled *to storage* (a memory
231                // win is already the cached copy; re-inserting it would needlessly
232                // replace a possibly-live Arc) AND the reconciled copy is not
233                // older than what memory already holds. This is what makes
234                // `load_merged` safe on hot read paths — it can never clobber a
235                // freshly-updated session with a stale storage copy.
236                if prefer_storage && chosen.updated_at >= memory_updated_at {
237                    self.cache.insert(
238                        session_id.to_string(),
239                        Arc::new(crate::SessionSnapshot::new(chosen.clone())),
240                    );
241                }
242                Some(chosen)
243            }
244            (Some(memory), None) => Some(memory),
245            (None, Some(storage)) => {
246                self.cache.insert(
247                    session_id.to_string(),
248                    Arc::new(crate::SessionSnapshot::new(storage.clone())),
249                );
250                Some(storage)
251            }
252            (None, None) => None,
253        })
254    }
255
256    /// Compatibility wrapper for read paths where the historical contract
257    /// treated a storage failure like absence. Mutating/recovery paths should
258    /// use [`Self::load_merged_checked`] so they can preserve retry state.
259    pub async fn load_merged(&self, session_id: &str) -> Option<Session> {
260        match self.load_merged_checked(session_id).await {
261            Ok(session) => session,
262            Err(error) => {
263                tracing::warn!(
264                    "[{}] Failed to load merged session from storage: {}",
265                    session_id,
266                    error
267                );
268                read_cached_session(&self.cache, session_id)
269            }
270        }
271    }
272
273    /// Persist the session (merge-on-write, preserving concurrent UI edits to
274    /// the authoritative metadata group) and refresh the in-memory cache.
275    pub async fn save_and_cache(&self, session: &mut Session) {
276        let result = self
277            .persistence
278            .merge_save_runtime_and_publish(session, |saved, _| {
279                #[cfg(test)]
280                self.run_post_durable_hook("save_and_cache", &saved.id);
281                self.cache.insert(
282                    saved.id.clone(),
283                    Arc::new(crate::SessionSnapshot::new(saved.clone())),
284                );
285            })
286            .await;
287        if let Err(error) = result {
288            tracing::warn!("[{}] Failed to save session: {}", session.id, error);
289        }
290    }
291
292    async fn refresh_cached_task_control_plane(&self, session_id: &str) {
293        match self.storage.load_runtime_control_plane(session_id).await {
294            Ok(Some(durable)) => {
295                if let Some(cached) = self.cache.get(session_id) {
296                    cached.update(|cached| adopt_task_control_plane(cached, &durable));
297                }
298            }
299            Ok(None) => {}
300            Err(error) => tracing::warn!(
301                "[{}] Failed to refresh Task control-plane after write conflict: {}",
302                session_id,
303                error
304            ),
305        }
306    }
307}
308
309fn adopt_task_control_plane(target: &mut Session, durable: &Session) {
310    target.task_list = durable.task_list.clone();
311    target
312        .metadata
313        .remove(bamboo_domain::session::runtime_metadata::keys::TASK_LIST_VERSION);
314    if let Some(runtime_metadata) = target.runtime_metadata.as_mut() {
315        runtime_metadata.task_list_version = None;
316    }
317    if target
318        .runtime_metadata
319        .as_ref()
320        .is_some_and(bamboo_domain::session::SessionRuntimeMetadata::is_empty)
321    {
322        target.runtime_metadata = None;
323    }
324    if let Some(version) = durable.task_list_version_meta() {
325        target.set_task_list_version_meta(version);
326    }
327}
328
329fn should_prefer_storage(memory_session: &Session, storage_session: &Session) -> bool {
330    // Never reconcile *backwards* to a strictly-older storage copy: if memory is
331    // newer it is authoritative (e.g. it just answered and cleared a pending
332    // question while storage still holds the stale one). Respecting `updated_at`
333    // here is what stops `load_merged` from returning — and caching — stale data.
334    if storage_session.updated_at < memory_session.updated_at {
335        return false;
336    }
337    // Storage is same-age or newer: prefer it when strictly newer, or when it
338    // still carries a pending question that the (same-age) memory copy lost, so
339    // a genuine clarification is never dropped.
340    storage_session.updated_at > memory_session.updated_at
341        || (memory_session.pending_question.is_none() && storage_session.pending_question.is_some())
342}
343
344/// `SessionRepository` is the canonical `RuntimeSessionPersistence`: the runtime
345/// can persist a session through the same coordinator (merge-on-write + cache
346/// refresh) instead of a bespoke adapter.
347#[async_trait::async_trait]
348impl bamboo_domain::RuntimeSessionPersistence for SessionRepository {
349    async fn save_runtime_session(&self, session: &mut Session) -> std::io::Result<()> {
350        // Runtime authorization reads through this same cache. Refresh it even
351        // when durable storage fails so a current activation can never observe
352        // a previous run's skill allowlist. The error is still returned to the
353        // caller and durable state remains unchanged.
354        self.persistence
355            .merge_save_runtime_and_publish(session, |saved, _| {
356                #[cfg(test)]
357                self.run_post_durable_hook("save_runtime_session", &saved.id);
358                self.cache.insert(
359                    saved.id.clone(),
360                    Arc::new(crate::SessionSnapshot::new(saved.clone())),
361                );
362            })
363            .await
364    }
365
366    async fn seed_runtime_activation(&self, session: &mut Session) -> std::io::Result<()> {
367        self.persistence
368            .seed_runtime_activation_and_publish(session, |saved, committed| {
369                #[cfg(test)]
370                self.run_post_durable_hook("seed_runtime_activation", &saved.id);
371                if committed {
372                    self.cache.insert(
373                        saved.id.clone(),
374                        Arc::new(crate::SessionSnapshot::new(saved.clone())),
375                    );
376                }
377            })
378            .await
379    }
380
381    async fn record_permission_posture_activation(
382        &self,
383        session_id: &str,
384        expected_audit_revision: Option<u64>,
385        seed: &bamboo_domain::PermissionAuditSeed,
386    ) -> std::io::Result<Option<Session>> {
387        self.persistence
388            .record_permission_posture_activation_and_publish(
389                session_id,
390                expected_audit_revision,
391                seed,
392                |saved| {
393                    #[cfg(test)]
394                    self.run_post_durable_hook("permission_posture_activation", &saved.id);
395                    self.cache.insert(
396                        saved.id.clone(),
397                        Arc::new(crate::SessionSnapshot::new(saved.clone())),
398                    );
399                },
400            )
401            .await
402    }
403
404    async fn save_runtime_control_plane(&self, session: &mut Session) -> std::io::Result<()> {
405        self.persistence
406            .save_runtime_only_and_publish(session, |saved| {
407                #[cfg(test)]
408                self.run_post_durable_hook("save", &saved.id);
409
410                // A control-plane snapshot may intentionally carry no messages
411                // (for example, child Task synchronization loads the root's
412                // runtime sidecar). Publish its fresh runtime fields without
413                // replacing a cache-resident transcript with that empty
414                // snapshot. SessionInbox admission is coupled to transcript
415                // persistence and is therefore preserved alongside the cached
416                // messages, matching the V2 sidecar overlay contract.
417                if let Some(cached) = self.cache.get(&saved.id) {
418                    cached.update(|cached| {
419                        let messages = cached.messages.clone();
420                        let provider_transcript = cached.provider_transcript.clone();
421                        let admission = cached
422                            .runtime_metadata
423                            .as_ref()
424                            .and_then(|metadata| metadata.session_inbox_admission.clone());
425                        let mut refreshed = saved.clone();
426                        refreshed.messages = messages;
427                        refreshed.provider_transcript = provider_transcript;
428                        if let Some(admission) = admission {
429                            refreshed
430                                .runtime_metadata
431                                .get_or_insert_with(Default::default)
432                                .session_inbox_admission = Some(admission);
433                        } else if let Some(metadata) = refreshed.runtime_metadata.as_mut() {
434                            metadata.session_inbox_admission = None;
435                        }
436                        *cached = refreshed;
437                    });
438                }
439            })
440            .await
441    }
442
443    async fn load_runtime_control_plane(
444        &self,
445        session_id: &str,
446    ) -> std::io::Result<Option<Session>> {
447        bamboo_domain::RuntimeSessionPersistence::load_runtime_control_plane(
448            self.persistence.as_ref(),
449            session_id,
450        )
451        .await
452    }
453
454    async fn update_task_list_control_plane(
455        &self,
456        session_id: &str,
457        task_list: &bamboo_domain::TaskList,
458        version: &str,
459    ) -> std::io::Result<bool> {
460        let result = self
461            .persistence
462            .update_task_list_control_plane_and_publish(session_id, task_list, version, |_| {
463                #[cfg(test)]
464                self.run_post_durable_hook("task", version);
465
466                // The durable transaction changed only Task-owned fields.
467                // Mirror that same narrow patch into the cache so a
468                // concurrent round/status/child transition already present
469                // in memory cannot be replaced by a stale whole-control-
470                // plane snapshot.
471                if let Some(cached) = self.cache.get(session_id) {
472                    cached.update(|cached| {
473                        cached.set_task_list(task_list.clone());
474                        cached.set_task_list_version_meta(version.to_string());
475                    });
476                }
477            })
478            .await;
479        if result
480            .as_ref()
481            .is_err_and(|error| error.kind() == std::io::ErrorKind::WouldBlock)
482        {
483            // The storage final-CAS rejected this unconditional patch because
484            // an independent writer won. Refresh the winner into the cache but
485            // preserve the conflict result so Taskwrite never mistakes it for
486            // the Ok(false) legacy-persistence fallback signal.
487            self.refresh_cached_task_control_plane(session_id).await;
488        }
489        result
490    }
491
492    async fn update_task_list_control_plane_if_version(
493        &self,
494        session_id: &str,
495        expected_version: &str,
496        expected_task_list: &bamboo_domain::TaskList,
497        task_list: &bamboo_domain::TaskList,
498        version: &str,
499    ) -> std::io::Result<bool> {
500        self.persistence
501            .update_task_list_control_plane_if_version_and_publish(
502                session_id,
503                expected_version,
504                expected_task_list,
505                task_list,
506                version,
507                |_| {
508                    if let Some(cached) = self.cache.get(session_id) {
509                        cached.update(|cached| {
510                            cached.set_task_list(task_list.clone());
511                            cached.set_task_list_version_meta(version.to_string());
512                        });
513                    }
514                },
515            )
516            .await
517    }
518
519    async fn update_task_list_control_planes_if_version(
520        &self,
521        session_id: &str,
522        shared_session_id: &str,
523        expected_version: &str,
524        expected_task_list: &bamboo_domain::TaskList,
525        task_list: &bamboo_domain::TaskList,
526        version: &str,
527    ) -> std::io::Result<bool> {
528        self.persistence
529            .update_task_list_control_planes_if_version_and_publish(
530                session_id,
531                shared_session_id,
532                expected_version,
533                expected_task_list,
534                task_list,
535                version,
536                |_, _| {
537                    for id in [session_id, shared_session_id] {
538                        if let Some(cached) = self.cache.get(id) {
539                            cached.update(|cached| {
540                                cached.set_task_list(task_list.clone());
541                                cached.set_task_list_version_meta(version.to_string());
542                            });
543                        }
544                    }
545                },
546            )
547            .await
548    }
549
550    async fn checkpoint_runtime_session(&self, session: &mut Session) -> std::io::Result<()> {
551        // The execute-boundary checkpoint uses LockedSessionStore's atomic
552        // append-safe transcript merge, then publishes that reconciled snapshot
553        // to the runtime cache.  On failure, leave the existing cache alone: the
554        // checkpoint may have failed to load the latest durable transcript, and
555        // replacing a fresher cache entry with the stale runner snapshot would
556        // set up a later SHRINK write.
557        self.persistence
558            .checkpoint_runtime_session_and_publish(session, |saved, committed| {
559                #[cfg(test)]
560                self.run_post_durable_hook("checkpoint", &saved.id);
561                if committed {
562                    self.cache.insert(
563                        saved.id.clone(),
564                        Arc::new(crate::SessionSnapshot::new(saved.clone())),
565                    );
566                }
567            })
568            .await
569    }
570
571    async fn load_runtime_session(&self, session_id: &str) -> std::io::Result<Option<Session>> {
572        self.try_load(session_id).await
573    }
574
575    async fn clear_legacy_pending_messages(
576        &self,
577        session_id: &str,
578        expected: &[serde_json::Value],
579    ) -> std::io::Result<bool> {
580        // Do not use the trait default here: `load_runtime_session` is
581        // cache-first, so a stale cache could erase a message concurrently
582        // appended to the durable legacy queue. Delegate the compare-and-clear
583        // to LockedSessionStore's single per-session critical section.
584        self.persistence
585            .clear_legacy_pending_messages_and_publish(session_id, expected, |latest| {
586                #[cfg(test)]
587                self.run_post_durable_hook("clear_legacy", session_id);
588                self.cache.insert(
589                    session_id.to_string(),
590                    Arc::new(crate::SessionSnapshot::new(latest.clone())),
591                );
592            })
593            .await
594    }
595
596    async fn append_token_usage_record(
597        &self,
598        session_id: &str,
599        json_line: &str,
600    ) -> std::io::Result<()> {
601        self.storage
602            .append_token_usage_record(session_id, json_line)
603            .await
604    }
605}
606
607#[cfg(test)]
608mod tests {
609    use super::*;
610    use bamboo_agent_core::storage::Storage;
611    use chrono::Utc;
612    use std::collections::HashMap;
613    use std::sync::atomic::{AtomicBool, Ordering};
614    use std::sync::{Condvar, Mutex};
615    use std::time::Duration;
616
617    #[derive(Default)]
618    struct MapStorage {
619        sessions: Mutex<HashMap<String, Session>>,
620        fail_pair_commit: AtomicBool,
621    }
622
623    struct FailingSaveStorage {
624        persisted: Mutex<Option<Session>>,
625    }
626
627    #[async_trait::async_trait]
628    impl Storage for MapStorage {
629        async fn save_session(&self, session: &Session) -> std::io::Result<()> {
630            self.sessions
631                .lock()
632                .unwrap()
633                .insert(session.id.clone(), session.clone());
634            Ok(())
635        }
636        async fn load_session(&self, session_id: &str) -> std::io::Result<Option<Session>> {
637            Ok(self.sessions.lock().unwrap().get(session_id).cloned())
638        }
639        async fn delete_session(&self, session_id: &str) -> std::io::Result<bool> {
640            Ok(self.sessions.lock().unwrap().remove(session_id).is_some())
641        }
642
643        async fn save_task_control_plane_if_matches(
644            &self,
645            original: &Session,
646            updated: &Session,
647        ) -> std::io::Result<bool> {
648            let mut sessions = self.sessions.lock().unwrap();
649            let matches_original = sessions.get(&original.id).is_some_and(|current| {
650                current.task_list_version_meta() == original.task_list_version_meta()
651                    && serde_json::to_value(&current.task_list).ok()
652                        == serde_json::to_value(&original.task_list).ok()
653            });
654            if !matches_original {
655                return Ok(false);
656            }
657            let Some(current) = sessions.get(&original.id).cloned() else {
658                return Ok(false);
659            };
660            let mut committed = current;
661            committed.task_list = updated.task_list.clone();
662            if let Some(version) = updated.task_list_version_meta() {
663                committed.set_task_list_version_meta(version);
664            }
665            sessions.insert(committed.id.clone(), committed);
666            Ok(true)
667        }
668
669        async fn save_task_control_planes_atomically(
670            &self,
671            first_original: &Session,
672            first_updated: &Session,
673            second_original: &Session,
674            second_updated: &Session,
675        ) -> std::io::Result<bool> {
676            if self.fail_pair_commit.swap(false, Ordering::SeqCst) {
677                return Err(std::io::Error::other(
678                    "injected paired Task transaction failure",
679                ));
680            }
681            let mut sessions = self.sessions.lock().unwrap();
682            let matches_original = |current: Option<&Session>, original: &Session| {
683                current.is_some_and(|current| {
684                    current.task_list_version_meta() == original.task_list_version_meta()
685                        && serde_json::to_value(&current.task_list).ok()
686                            == serde_json::to_value(&original.task_list).ok()
687                })
688            };
689            if !matches_original(sessions.get(&first_original.id), first_original)
690                || !matches_original(sessions.get(&second_original.id), second_original)
691            {
692                return Ok(false);
693            }
694            sessions.insert(first_updated.id.clone(), first_updated.clone());
695            sessions.insert(second_updated.id.clone(), second_updated.clone());
696            Ok(true)
697        }
698    }
699
700    #[async_trait::async_trait]
701    impl Storage for FailingSaveStorage {
702        async fn save_session(&self, _session: &Session) -> std::io::Result<()> {
703            Err(std::io::Error::other("injected save failure"))
704        }
705
706        async fn load_session(&self, _session_id: &str) -> std::io::Result<Option<Session>> {
707            Ok(self.persisted.lock().unwrap().clone())
708        }
709
710        async fn delete_session(&self, _session_id: &str) -> std::io::Result<bool> {
711            Ok(false)
712        }
713    }
714
715    fn test_repo(storage: Arc<dyn Storage>) -> SessionRepository {
716        let cache: SessionCache = Arc::default();
717        let persistence = Arc::new(LockedSessionStore::new(storage.clone()));
718        SessionRepository::new(cache, storage, persistence)
719    }
720
721    fn cache_put(repo: &SessionRepository, session: &Session) {
722        repo.cache().insert(
723            session.id.clone(),
724            Arc::new(crate::SessionSnapshot::new(session.clone())),
725        );
726    }
727
728    fn task_list(session_id: &str, title: &str) -> bamboo_domain::TaskList {
729        let now = Utc::now();
730        bamboo_domain::TaskList {
731            session_id: session_id.to_string(),
732            title: title.to_string(),
733            items: Vec::new(),
734            created_at: now,
735            updated_at: now,
736        }
737    }
738
739    #[tokio::test]
740    async fn paired_task_cas_narrowly_updates_child_and_root_cache() {
741        let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
742        let repo = test_repo(storage.clone());
743        let root_id = "paired-cache-root";
744        let child_id = "paired-cache-child";
745        let expected_task_list = task_list(root_id, "old shared");
746
747        let mut root = Session::new(root_id, "model");
748        root.add_message(bamboo_agent_core::Message::user("root transcript"));
749        root.metadata
750            .insert("unrelated.root".to_string(), "keep".to_string());
751        root.set_task_list(expected_task_list.clone());
752        root.set_task_list_version_meta("1");
753        storage.save_session(&root).await.unwrap();
754        cache_put(&repo, &root);
755
756        let mut child = Session::new_child(child_id, root_id, "model", "child");
757        child.add_message(bamboo_agent_core::Message::user("child transcript"));
758        child
759            .metadata
760            .insert("unrelated.child".to_string(), "keep".to_string());
761        child.set_task_list(expected_task_list.clone());
762        child.set_task_list_version_meta("1");
763        storage.save_session(&child).await.unwrap();
764        cache_put(&repo, &child);
765
766        assert!(
767            bamboo_domain::RuntimeSessionPersistence::update_task_list_control_planes_if_version(
768                &repo,
769                child_id,
770                root_id,
771                "1",
772                &expected_task_list,
773                &task_list(root_id, "evaluated"),
774                "2",
775            )
776            .await
777            .expect("paired cache CAS succeeds")
778        );
779
780        let cached_root = read_cached_session(repo.cache(), root_id).expect("cached root");
781        let cached_child = read_cached_session(repo.cache(), child_id).expect("cached child");
782        for (session, transcript, metadata_key) in [
783            (&cached_root, "root transcript", "unrelated.root"),
784            (&cached_child, "child transcript", "unrelated.child"),
785        ] {
786            assert_eq!(session.task_list_version_meta().as_deref(), Some("2"));
787            assert_eq!(
788                session.task_list.as_ref().map(|list| list.title.as_str()),
789                Some("evaluated")
790            );
791            assert_eq!(session.messages.len(), 1);
792            assert_eq!(session.messages[0].content, transcript);
793            assert_eq!(
794                session.metadata.get(metadata_key).map(String::as_str),
795                Some("keep")
796            );
797        }
798    }
799
800    #[tokio::test]
801    async fn failed_paired_task_transaction_does_not_publish_child_or_root_cache() {
802        let concrete = Arc::new(MapStorage::default());
803        let storage: Arc<dyn Storage> = concrete.clone();
804        let repo = test_repo(storage.clone());
805        let root_id = "paired-cache-failure-root";
806        let child_id = "paired-cache-failure-child";
807        let expected_task_list = task_list(root_id, "old shared");
808
809        let mut root = Session::new(root_id, "model");
810        root.add_message(bamboo_agent_core::Message::user("root transcript"));
811        root.metadata
812            .insert("unrelated.root".to_string(), "keep".to_string());
813        root.set_task_list(expected_task_list.clone());
814        root.set_task_list_version_meta("1");
815        storage.save_session(&root).await.unwrap();
816        cache_put(&repo, &root);
817
818        let mut child = Session::new_child(child_id, root_id, "model", "child");
819        child.add_message(bamboo_agent_core::Message::user("child transcript"));
820        child
821            .metadata
822            .insert("unrelated.child".to_string(), "keep".to_string());
823        child.set_task_list(expected_task_list.clone());
824        child.set_task_list_version_meta("1");
825        storage.save_session(&child).await.unwrap();
826        cache_put(&repo, &child);
827
828        concrete.fail_pair_commit.store(true, Ordering::SeqCst);
829        let error =
830            bamboo_domain::RuntimeSessionPersistence::update_task_list_control_planes_if_version(
831                &repo,
832                child_id,
833                root_id,
834                "1",
835                &expected_task_list,
836                &task_list(root_id, "must not publish"),
837                "2",
838            )
839            .await
840            .expect_err("paired durable transaction fails");
841        assert!(error.to_string().contains("injected paired"));
842
843        for (id, expected_title, transcript, metadata_key) in [
844            (root_id, "old shared", "root transcript", "unrelated.root"),
845            (
846                child_id,
847                "old shared",
848                "child transcript",
849                "unrelated.child",
850            ),
851        ] {
852            let cached = read_cached_session(repo.cache(), id).expect("cached session remains");
853            assert_eq!(cached.task_list_version_meta().as_deref(), Some("1"));
854            assert_eq!(
855                cached.task_list.as_ref().map(|list| list.title.as_str()),
856                Some(expected_title)
857            );
858            assert_eq!(cached.messages[0].content, transcript);
859            assert_eq!(
860                cached.metadata.get(metadata_key).map(String::as_str),
861                Some("keep")
862            );
863
864            let durable = storage.load_session(id).await.unwrap().unwrap();
865            assert_eq!(durable.task_list_version_meta().as_deref(), Some("1"));
866            assert_eq!(
867                durable.task_list.as_ref().map(|list| list.title.as_str()),
868                Some(expected_title)
869            );
870        }
871    }
872
873    fn durable_cache_fence(
874        operation: impl Into<String>,
875        marker: impl Into<String>,
876    ) -> (
877        PostDurableHook,
878        tokio::sync::oneshot::Receiver<()>,
879        Arc<(Mutex<bool>, Condvar)>,
880    ) {
881        let operation = operation.into();
882        let marker = marker.into();
883        let (started_tx, started_rx) = tokio::sync::oneshot::channel();
884        let started_tx = Arc::new(Mutex::new(Some(started_tx)));
885        let release = Arc::new((Mutex::new(false), Condvar::new()));
886        let hook_release = release.clone();
887        let hook: PostDurableHook = Arc::new(move |actual_operation, actual_marker| {
888            if actual_operation != operation || actual_marker != marker {
889                return;
890            }
891            if let Some(started_tx) = started_tx.lock().unwrap().take() {
892                started_tx.send(()).expect("fence observer still present");
893            }
894            let (released, wake) = &*hook_release;
895            let mut released = released.lock().unwrap();
896            while !*released {
897                released = wake.wait(released).unwrap();
898            }
899        });
900        (hook, started_rx, release)
901    }
902
903    fn release_fence(release: &Arc<(Mutex<bool>, Condvar)>) {
904        let (released, wake) = &**release;
905        *released.lock().unwrap() = true;
906        wake.notify_all();
907    }
908
909    async fn assert_second_write_waits_for_cache_publish<T>(
910        first: tokio::task::JoinHandle<std::io::Result<T>>,
911        mut second: tokio::task::JoinHandle<std::io::Result<T>>,
912        release: Arc<(Mutex<bool>, Condvar)>,
913    ) -> (T, T) {
914        let second_before_release =
915            tokio::time::timeout(Duration::from_millis(100), &mut second).await;
916        let completed_before_release = second_before_release.is_ok();
917        release_fence(&release);
918
919        let first = first
920            .await
921            .expect("first writer joins")
922            .expect("first writer succeeds");
923        let second = match second_before_release {
924            Ok(joined) => joined
925                .expect("second writer joins")
926                .expect("second writer succeeds"),
927            Err(_) => second
928                .await
929                .expect("second writer joins")
930                .expect("second writer succeeds"),
931        };
932        assert!(
933            !completed_before_release,
934            "the second write must remain behind the first write's durable-to-cache fence"
935        );
936        (first, second)
937    }
938
939    #[derive(Clone, Copy, Debug)]
940    enum FullSaveRoute {
941        InherentSave,
942        SaveAndCache,
943        RuntimePersistence,
944    }
945
946    impl FullSaveRoute {
947        fn operation(self) -> &'static str {
948            match self {
949                Self::InherentSave => "save_full",
950                Self::SaveAndCache => "save_and_cache",
951                Self::RuntimePersistence => "save_runtime_session",
952            }
953        }
954
955        fn name(self) -> &'static str {
956            match self {
957                Self::InherentSave => "inherent",
958                Self::SaveAndCache => "save-and-cache",
959                Self::RuntimePersistence => "runtime-persistence",
960            }
961        }
962    }
963
964    async fn assert_full_save_route_serializes_cache_publish(route: FullSaveRoute) {
965        let temp = tempfile::tempdir().unwrap();
966        let concrete_storage = Arc::new(
967            bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
968                .await
969                .expect("SessionStoreV2"),
970        );
971        let storage: Arc<dyn Storage> = concrete_storage;
972        let id = format!("root-full-cache-order-{}", route.name());
973        let (hook, first_durable, release) = durable_cache_fence(route.operation(), id.clone());
974        let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
975
976        let mut initial = Session::new(&id, "model");
977        initial.add_message(bamboo_agent_core::Message::user("durable transcript"));
978        initial.set_task_list(task_list(&id, "initial"));
979        initial.set_task_list_version_meta("0");
980        initial
981            .metadata
982            .insert("unrelated.runtime".to_string(), "keep".to_string());
983        storage.save_session(&initial).await.unwrap();
984        cache_put(&repo, &initial);
985
986        let first_repo = repo.clone();
987        let mut root_snapshot = initial.clone();
988        root_snapshot.add_message(bamboo_agent_core::Message::assistant(
989            "full-save transcript suffix",
990            None,
991        ));
992        root_snapshot.set_task_list(task_list(&id, "root"));
993        root_snapshot.set_task_list_version_meta("1");
994        let first = tokio::spawn(async move {
995            match route {
996                FullSaveRoute::InherentSave => first_repo.save(&mut root_snapshot).await,
997                FullSaveRoute::SaveAndCache => {
998                    first_repo.save_and_cache(&mut root_snapshot).await;
999                    Ok(())
1000                }
1001                FullSaveRoute::RuntimePersistence => {
1002                    bamboo_domain::RuntimeSessionPersistence::save_runtime_session(
1003                        first_repo.as_ref(),
1004                        &mut root_snapshot,
1005                    )
1006                    .await
1007                }
1008            }
1009        });
1010        first_durable
1011            .await
1012            .expect("root full durable write reached");
1013
1014        let second_repo = repo.clone();
1015        let second_id = id.clone();
1016        let child_task_list = task_list(&id, "child");
1017        let second = tokio::spawn(async move {
1018            bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
1019                second_repo.as_ref(),
1020                &second_id,
1021                &child_task_list,
1022                "2",
1023            )
1024            .await
1025            .map(|updated| assert!(updated, "root must exist"))
1026        });
1027        assert_second_write_waits_for_cache_publish(first, second, release).await;
1028
1029        let durable = storage.load_session(&id).await.unwrap().unwrap();
1030        let cached = read_cached_session(repo.cache(), &id).expect("cached root");
1031        for (tier, session) in [("durable", durable), ("cache", cached)] {
1032            assert_eq!(
1033                session.task_list_version_meta().as_deref(),
1034                Some("2"),
1035                "{route:?} {tier} must retain the child transaction"
1036            );
1037            assert_eq!(
1038                session.task_list.as_ref().map(|list| list.title.as_str()),
1039                Some("child"),
1040                "{route:?} {tier} must retain the child transaction"
1041            );
1042            assert_eq!(
1043                session
1044                    .metadata
1045                    .get("unrelated.runtime")
1046                    .map(String::as_str),
1047                Some("keep"),
1048                "{route:?} {tier} must preserve unrelated runtime state"
1049            );
1050            assert_eq!(
1051                session.messages.len(),
1052                2,
1053                "{route:?} {tier} must preserve the full-save transcript"
1054            );
1055        }
1056    }
1057
1058    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1059    async fn root_full_saves_and_child_task_patch_share_publish_order() {
1060        for route in [
1061            FullSaveRoute::InherentSave,
1062            FullSaveRoute::SaveAndCache,
1063            FullSaveRoute::RuntimePersistence,
1064        ] {
1065            assert_full_save_route_serializes_cache_publish(route).await;
1066        }
1067    }
1068
1069    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1070    async fn checkpoint_and_child_task_patch_share_publish_order() {
1071        let temp = tempfile::tempdir().unwrap();
1072        let concrete_storage = Arc::new(
1073            bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
1074                .await
1075                .expect("SessionStoreV2"),
1076        );
1077        let storage: Arc<dyn Storage> = concrete_storage;
1078        let id = "checkpoint-cache-order";
1079        let (hook, checkpoint_durable, release) = durable_cache_fence("checkpoint", id);
1080        let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
1081
1082        let mut initial = Session::new(id, "model");
1083        initial.add_message(bamboo_agent_core::Message::user("durable transcript"));
1084        initial.set_task_list(task_list(id, "initial"));
1085        initial.set_task_list_version_meta("0");
1086        initial
1087            .metadata
1088            .insert("unrelated.runtime".to_string(), "keep".to_string());
1089        storage.save_session(&initial).await.unwrap();
1090        cache_put(&repo, &initial);
1091
1092        let checkpoint_repo = repo.clone();
1093        let mut checkpoint_snapshot = initial.clone();
1094        checkpoint_snapshot.add_message(bamboo_agent_core::Message::assistant(
1095            "checkpoint transcript suffix",
1096            None,
1097        ));
1098        checkpoint_snapshot.set_task_list(task_list(id, "checkpoint"));
1099        checkpoint_snapshot.set_task_list_version_meta("1");
1100        let checkpoint = tokio::spawn(async move {
1101            bamboo_domain::RuntimeSessionPersistence::checkpoint_runtime_session(
1102                checkpoint_repo.as_ref(),
1103                &mut checkpoint_snapshot,
1104            )
1105            .await
1106        });
1107        checkpoint_durable
1108            .await
1109            .expect("checkpoint durable write reached");
1110
1111        let child_repo = repo.clone();
1112        let child_task_list = task_list(id, "child");
1113        let child_patch = tokio::spawn(async move {
1114            bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
1115                child_repo.as_ref(),
1116                id,
1117                &child_task_list,
1118                "2",
1119            )
1120            .await
1121            .map(|updated| assert!(updated, "root must exist"))
1122        });
1123        assert_second_write_waits_for_cache_publish(checkpoint, child_patch, release).await;
1124
1125        let durable = storage.load_session(id).await.unwrap().unwrap();
1126        let cached = read_cached_session(repo.cache(), id).expect("cached root");
1127        for (tier, session) in [("durable", durable), ("cache", cached)] {
1128            assert_eq!(
1129                session.task_list_version_meta().as_deref(),
1130                Some("2"),
1131                "{tier} must retain the child transaction"
1132            );
1133            assert_eq!(
1134                session.task_list.as_ref().map(|list| list.title.as_str()),
1135                Some("child"),
1136                "{tier} must retain the child transaction"
1137            );
1138            assert_eq!(
1139                session
1140                    .metadata
1141                    .get("unrelated.runtime")
1142                    .map(String::as_str),
1143                Some("keep"),
1144                "{tier} must preserve unrelated runtime state"
1145            );
1146            assert_eq!(
1147                session.messages.len(),
1148                2,
1149                "{tier} must preserve the checkpoint transcript"
1150            );
1151        }
1152    }
1153
1154    #[derive(Clone, Copy, Debug)]
1155    enum CacheBackfillRoute {
1156        Load,
1157        TryLoad,
1158    }
1159
1160    impl CacheBackfillRoute {
1161        fn operation(self) -> &'static str {
1162            match self {
1163                Self::Load => "load",
1164                Self::TryLoad => "try_load",
1165            }
1166        }
1167
1168        fn name(self) -> &'static str {
1169            match self {
1170                Self::Load => "load",
1171                Self::TryLoad => "try-load",
1172            }
1173        }
1174    }
1175
1176    async fn assert_cache_backfill_serializes_with_task_patch(route: CacheBackfillRoute) {
1177        let temp = tempfile::tempdir().unwrap();
1178        let concrete_storage = Arc::new(
1179            bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
1180                .await
1181                .expect("SessionStoreV2"),
1182        );
1183        let storage: Arc<dyn Storage> = concrete_storage;
1184        let id = format!("cache-backfill-order-{}", route.name());
1185        let (hook, loaded_old_durable, release) =
1186            durable_cache_fence(route.operation(), id.clone());
1187        let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
1188
1189        let mut initial = Session::new(&id, "model");
1190        initial.set_task_list(task_list(&id, "initial"));
1191        initial.set_task_list_version_meta("0");
1192        storage.save_session(&initial).await.unwrap();
1193        assert!(
1194            read_cached_session(repo.cache(), &id).is_none(),
1195            "the race requires a genuine cache miss"
1196        );
1197
1198        let load_repo = repo.clone();
1199        let load_id = id.clone();
1200        let load = tokio::spawn(async move {
1201            let loaded = match route {
1202                CacheBackfillRoute::Load => load_repo.load(&load_id).await,
1203                CacheBackfillRoute::TryLoad => {
1204                    load_repo.try_load(&load_id).await.expect("storage load")
1205                }
1206            };
1207            assert!(loaded.is_some(), "seeded session must load");
1208            Ok(())
1209        });
1210        loaded_old_durable
1211            .await
1212            .expect("old durable snapshot loaded");
1213
1214        let patch_repo = repo.clone();
1215        let patch_id = id.clone();
1216        let child_task_list = task_list(&id, "child");
1217        let patch = tokio::spawn(async move {
1218            bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
1219                patch_repo.as_ref(),
1220                &patch_id,
1221                &child_task_list,
1222                "1",
1223            )
1224            .await
1225            .map(|updated| assert!(updated, "root must exist"))
1226        });
1227        assert_second_write_waits_for_cache_publish(load, patch, release).await;
1228
1229        let durable = storage.load_session(&id).await.unwrap().unwrap();
1230        let cached = read_cached_session(repo.cache(), &id).expect("backfilled cache");
1231        for (tier, session) in [("durable", durable), ("cache", cached)] {
1232            assert_eq!(
1233                session.task_list_version_meta().as_deref(),
1234                Some("1"),
1235                "{route:?} {tier} must retain the child transaction"
1236            );
1237            assert_eq!(
1238                session.task_list.as_ref().map(|list| list.title.as_str()),
1239                Some("child"),
1240                "{route:?} {tier} must retain the child transaction"
1241            );
1242        }
1243    }
1244
1245    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1246    async fn cache_miss_backfills_and_child_task_patch_share_publish_order() {
1247        for route in [CacheBackfillRoute::Load, CacheBackfillRoute::TryLoad] {
1248            assert_cache_backfill_serializes_with_task_patch(route).await;
1249        }
1250    }
1251
1252    #[tokio::test]
1253    async fn cache_hits_do_not_wait_for_the_persistence_lock() {
1254        let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
1255        let repo = test_repo(storage);
1256        let id = "cache-hit-lock-free";
1257        let cached = Session::new(id, "cached-model");
1258        cache_put(&repo, &cached);
1259        let persistence_guard = repo.persistence().acquire_lock(id).await;
1260
1261        let loaded = tokio::time::timeout(Duration::from_millis(100), repo.load(id)).await;
1262        let try_loaded = tokio::time::timeout(Duration::from_millis(100), repo.try_load(id)).await;
1263        drop(persistence_guard);
1264
1265        assert_eq!(
1266            loaded
1267                .expect("cache hit must not wait for the persistence lock")
1268                .expect("cached session")
1269                .model,
1270            "cached-model"
1271        );
1272        assert_eq!(
1273            try_loaded
1274                .expect("fallible cache hit must not wait for the persistence lock")
1275                .expect("cache read succeeds")
1276                .expect("cached session")
1277                .model,
1278            "cached-model"
1279        );
1280    }
1281
1282    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1283    async fn load_merged_storage_refresh_and_child_task_patch_share_publish_order() {
1284        let temp = tempfile::tempdir().unwrap();
1285        let concrete_storage = Arc::new(
1286            bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
1287                .await
1288                .expect("SessionStoreV2"),
1289        );
1290        let storage: Arc<dyn Storage> = concrete_storage;
1291        let id = "load-merged-cache-order";
1292        let (hook, loaded_old_durable, release) = durable_cache_fence("load_merged", id);
1293        let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
1294
1295        let mut durable = Session::new(id, "model");
1296        durable.updated_at = Utc::now();
1297        durable.set_task_list(task_list(id, "initial"));
1298        durable.set_task_list_version_meta("0");
1299        storage.save_session(&durable).await.unwrap();
1300
1301        let mut memory = durable.clone();
1302        memory.updated_at = durable.updated_at - chrono::Duration::seconds(1);
1303        memory.set_task_list(task_list(id, "memory"));
1304        cache_put(&repo, &memory);
1305
1306        let load_repo = repo.clone();
1307        let load = tokio::spawn(async move {
1308            assert!(
1309                load_repo.load_merged(id).await.is_some(),
1310                "seeded session must load"
1311            );
1312            Ok(())
1313        });
1314        loaded_old_durable
1315            .await
1316            .expect("old durable snapshot loaded");
1317
1318        let patch_repo = repo.clone();
1319        let child_task_list = task_list(id, "child");
1320        let patch = tokio::spawn(async move {
1321            bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
1322                patch_repo.as_ref(),
1323                id,
1324                &child_task_list,
1325                "1",
1326            )
1327            .await
1328            .map(|updated| assert!(updated, "root must exist"))
1329        });
1330        assert_second_write_waits_for_cache_publish(load, patch, release).await;
1331
1332        let durable = storage.load_session(id).await.unwrap().unwrap();
1333        let cached = read_cached_session(repo.cache(), id).expect("refreshed cache");
1334        for (tier, session) in [("durable", durable), ("cache", cached)] {
1335            assert_eq!(
1336                session.task_list_version_meta().as_deref(),
1337                Some("1"),
1338                "{tier} must retain the child transaction"
1339            );
1340            assert_eq!(
1341                session.task_list.as_ref().map(|list| list.title.as_str()),
1342                Some("child"),
1343                "{tier} must retain the child transaction"
1344            );
1345        }
1346    }
1347
1348    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1349    async fn legacy_clear_refresh_and_child_task_patch_share_publish_order() {
1350        let temp = tempfile::tempdir().unwrap();
1351        let concrete_storage = Arc::new(
1352            bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
1353                .await
1354                .expect("SessionStoreV2"),
1355        );
1356        let storage: Arc<dyn Storage> = concrete_storage;
1357        let id = "legacy-clear-cache-order";
1358        let expected = vec![serde_json::json!({"content": "legacy"})];
1359        let (hook, loaded_post_cas_snapshot, release) = durable_cache_fence("clear_legacy", id);
1360        let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
1361
1362        let mut initial = Session::new(id, "model");
1363        initial.set_pending_injected_messages(expected.clone());
1364        initial.set_task_list(task_list(id, "initial"));
1365        initial.set_task_list_version_meta("0");
1366        storage.save_session(&initial).await.unwrap();
1367        cache_put(&repo, &initial);
1368
1369        let clear_repo = repo.clone();
1370        let clear_expected = expected.clone();
1371        let clear = tokio::spawn(async move {
1372            bamboo_domain::RuntimeSessionPersistence::clear_legacy_pending_messages(
1373                clear_repo.as_ref(),
1374                id,
1375                &clear_expected,
1376            )
1377            .await
1378            .map(|cleared| assert!(cleared, "legacy queue must match"))
1379        });
1380        loaded_post_cas_snapshot
1381            .await
1382            .expect("post-CAS snapshot loaded");
1383
1384        let patch_repo = repo.clone();
1385        let child_task_list = task_list(id, "child");
1386        let patch = tokio::spawn(async move {
1387            bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
1388                patch_repo.as_ref(),
1389                id,
1390                &child_task_list,
1391                "1",
1392            )
1393            .await
1394            .map(|updated| assert!(updated, "root must exist"))
1395        });
1396        assert_second_write_waits_for_cache_publish(clear, patch, release).await;
1397
1398        let durable = storage.load_session(id).await.unwrap().unwrap();
1399        let cached = read_cached_session(repo.cache(), id).expect("refreshed cache");
1400        for (tier, session) in [("durable", durable), ("cache", cached)] {
1401            assert_eq!(
1402                session.task_list_version_meta().as_deref(),
1403                Some("1"),
1404                "{tier} must retain the child transaction"
1405            );
1406            assert_eq!(
1407                session.task_list.as_ref().map(|list| list.title.as_str()),
1408                Some("child"),
1409                "{tier} must retain the child transaction"
1410            );
1411            assert!(
1412                !session.has_pending_injected_messages(),
1413                "{tier} must retain the successful legacy clear"
1414            );
1415        }
1416    }
1417
1418    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1419    async fn concurrent_task_patches_publish_cache_in_durable_order() {
1420        let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
1421        let (hook, first_durable, release) = durable_cache_fence("task", "1");
1422        let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
1423        let id = "concurrent-task-cache-order";
1424        let mut initial = Session::new(id, "model");
1425        initial.set_task_list(task_list(id, "initial"));
1426        initial.set_task_list_version_meta("0");
1427        storage.save_session(&initial).await.unwrap();
1428        cache_put(&repo, &initial);
1429
1430        let first_repo = repo.clone();
1431        let first_task_list = task_list(id, "first");
1432        let first = tokio::spawn(async move {
1433            bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
1434                first_repo.as_ref(),
1435                id,
1436                &first_task_list,
1437                "1",
1438            )
1439            .await
1440        });
1441        first_durable.await.expect("first durable write reached");
1442
1443        let second_repo = repo.clone();
1444        let second_task_list = task_list(id, "second");
1445        let second = tokio::spawn(async move {
1446            bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
1447                second_repo.as_ref(),
1448                id,
1449                &second_task_list,
1450                "2",
1451            )
1452            .await
1453        });
1454        let (first_updated, second_updated) =
1455            assert_second_write_waits_for_cache_publish(first, second, release).await;
1456        assert!(first_updated && second_updated);
1457
1458        let durable = storage.load_session(id).await.unwrap().unwrap();
1459        let cached = read_cached_session(repo.cache(), id).expect("cached root");
1460        for (tier, session) in [("durable", durable), ("cache", cached)] {
1461            assert_eq!(
1462                session.task_list_version_meta().as_deref(),
1463                Some("2"),
1464                "{tier} must retain the second transaction"
1465            );
1466            assert_eq!(
1467                session.task_list.as_ref().map(|list| list.title.as_str()),
1468                Some("second"),
1469                "{tier} must retain the second transaction"
1470            );
1471        }
1472    }
1473
1474    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1475    async fn root_control_plane_save_and_child_task_patch_share_publish_order() {
1476        let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
1477        let (hook, root_durable, release) = durable_cache_fence("save", "root-control-plane-order");
1478        let repo = Arc::new(test_repo(storage.clone()).with_post_durable_hook(hook));
1479        let id = "root-control-plane-order";
1480
1481        let mut initial = Session::new(id, "model");
1482        initial.add_message(bamboo_agent_core::Message::user("durable transcript"));
1483        initial.set_task_list(task_list(id, "initial"));
1484        initial.set_task_list_version_meta("0");
1485        initial
1486            .metadata
1487            .insert("unrelated.runtime".to_string(), "keep".to_string());
1488        storage.save_session(&initial).await.unwrap();
1489        cache_put(&repo, &initial);
1490
1491        let root_repo = repo.clone();
1492        let mut root_snapshot = initial.clone();
1493        root_snapshot.set_task_list(task_list(id, "root"));
1494        root_snapshot.set_task_list_version_meta("1");
1495        let root_save = tokio::spawn(async move {
1496            bamboo_domain::RuntimeSessionPersistence::save_runtime_control_plane(
1497                root_repo.as_ref(),
1498                &mut root_snapshot,
1499            )
1500            .await
1501        });
1502        root_durable.await.expect("root durable write reached");
1503
1504        let child_repo = repo.clone();
1505        let child_task_list = task_list(id, "child");
1506        let child_patch = tokio::spawn(async move {
1507            bamboo_domain::RuntimeSessionPersistence::update_task_list_control_plane(
1508                child_repo.as_ref(),
1509                id,
1510                &child_task_list,
1511                "2",
1512            )
1513            .await
1514            .map(|updated| {
1515                assert!(updated, "root must exist");
1516            })
1517        });
1518        assert_second_write_waits_for_cache_publish(root_save, child_patch, release).await;
1519
1520        let durable = storage.load_session(id).await.unwrap().unwrap();
1521        let cached = read_cached_session(repo.cache(), id).expect("cached root");
1522        for (tier, session) in [("durable", durable), ("cache", cached)] {
1523            assert_eq!(
1524                session.task_list_version_meta().as_deref(),
1525                Some("2"),
1526                "{tier} must retain the child transaction"
1527            );
1528            assert_eq!(
1529                session.task_list.as_ref().map(|list| list.title.as_str()),
1530                Some("child"),
1531                "{tier} must retain the child transaction"
1532            );
1533            assert_eq!(
1534                session
1535                    .metadata
1536                    .get("unrelated.runtime")
1537                    .map(String::as_str),
1538                Some("keep"),
1539                "{tier} must preserve unrelated runtime state"
1540            );
1541            assert_eq!(
1542                session.messages.len(),
1543                1,
1544                "{tier} must preserve the transcript"
1545            );
1546        }
1547    }
1548
1549    #[tokio::test]
1550    async fn stale_control_plane_save_keeps_checkpointed_ledger_in_durable_and_cache() {
1551        let temp = tempfile::tempdir().unwrap();
1552        let storage: Arc<dyn Storage> = Arc::new(
1553            bamboo_storage::SessionStoreV2::new(temp.path().to_path_buf())
1554                .await
1555                .expect("SessionStoreV2"),
1556        );
1557        let repo = test_repo(storage.clone());
1558        let id = "control-plane-ledger-cache";
1559        let mut initial = Session::new(id, "model");
1560        initial.add_message(bamboo_agent_core::Message::user("durable transcript"));
1561        let assistant = bamboo_agent_core::Message::assistant("normalized", None);
1562        let anchor = assistant.id.clone();
1563        initial.add_message(assistant);
1564        let native_item = bamboo_domain::ProviderTranscriptItem::try_from_payload(
1565            bamboo_domain::ProviderFamily::OpenAi,
1566            bamboo_domain::ProviderProtocol::OpenAiResponsesV1,
1567            bamboo_domain::ProviderTranscriptOrigin::Provider,
1568            bamboo_domain::ProviderTranscriptAuthor::Model,
1569            serde_json::json!({
1570                "type":"tool_search_call","id":"tsc_cache_native","execution":"client","call_id":"cache_native",
1571                "status":"completed","arguments":{"query":"CACHE_NATIVE_PAYLOAD_SENTINEL"}
1572            }),
1573        )
1574        .unwrap();
1575        initial
1576            .append_provider_transcript_group(&anchor, None, vec![native_item])
1577            .unwrap();
1578        storage.save_session(&initial).await.unwrap();
1579        cache_put(&repo, &initial);
1580        let mut stale = initial.clone();
1581
1582        let mut runner = initial;
1583        runner.model_context_state = Some(bamboo_domain::ModelContextState {
1584            state_revision: 1,
1585            prefix_epoch: 1,
1586            cache_scope_sha256: Some("scope".to_string()),
1587            transcript_item_sha256: vec!["runner-l1".to_string()],
1588            ..bamboo_domain::ModelContextState::default()
1589        });
1590        bamboo_domain::RuntimeSessionPersistence::checkpoint_runtime_session(&repo, &mut runner)
1591            .await
1592            .unwrap();
1593
1594        stale
1595            .metadata
1596            .insert("runtime.suspend_reason".to_string(), "waiting".to_string());
1597        bamboo_domain::RuntimeSessionPersistence::save_runtime_control_plane(&repo, &mut stale)
1598            .await
1599            .unwrap();
1600
1601        let expected = runner.model_context_state;
1602        let expected_native = runner.provider_transcript;
1603        let durable = storage.load_session(id).await.unwrap().unwrap();
1604        let cached = read_cached_session(repo.cache(), id).expect("cached session");
1605        for (tier, session) in [("durable", durable), ("cache", cached.clone())] {
1606            assert_eq!(session.model_context_state, expected, "tier={tier}");
1607            assert_eq!(
1608                session.provider_transcript, expected_native,
1609                "tier={tier} must retain the message-anchored native transcript"
1610            );
1611            assert_eq!(
1612                session
1613                    .metadata
1614                    .get("runtime.suspend_reason")
1615                    .map(String::as_str),
1616                Some("waiting"),
1617                "tier={tier}"
1618            );
1619            assert_eq!(session.messages.len(), 2, "tier={tier}");
1620        }
1621
1622        let runtime_json =
1623            std::fs::read_to_string(temp.path().join("sessions").join(id).join("runtime.json"))
1624                .unwrap();
1625        assert!(!runtime_json.contains("CACHE_NATIVE_PAYLOAD_SENTINEL"));
1626
1627        // Prove the cache projection cannot turn a runtime-only update into a
1628        // later durable loss when that cached value becomes a full checkpoint.
1629        let mut cache_writer = cached;
1630        bamboo_domain::RuntimeSessionPersistence::checkpoint_runtime_session(
1631            &repo,
1632            &mut cache_writer,
1633        )
1634        .await
1635        .unwrap();
1636        let restarted = storage.load_session(id).await.unwrap().unwrap();
1637        assert_eq!(restarted.provider_transcript, expected_native);
1638    }
1639
1640    #[tokio::test]
1641    async fn narrow_runtime_metadata_transaction_preserves_live_and_durable_non_owned_state() {
1642        let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
1643        let repo = test_repo(storage.clone());
1644        let id = "narrow-metadata";
1645        let mut durable = Session::new(id, "durable-model");
1646        durable.add_message(bamboo_agent_core::Message::user("durable user turn"));
1647        durable
1648            .metadata
1649            .insert("external.durable".to_string(), "keep".to_string());
1650        storage.save_session(&durable).await.expect("seed durable");
1651
1652        let mut live = durable.clone();
1653        live.add_message(bamboo_agent_core::Message::assistant(
1654            "in-flight assistant tool call",
1655            None,
1656        ));
1657        live.model = "live-model".to_string();
1658        live.metadata
1659            .insert("external.live".to_string(), "keep".to_string());
1660        cache_put(&repo, &live);
1661
1662        repo.update_runtime_session(id, &["workflow.owned"], |latest| {
1663            latest
1664                .metadata
1665                .insert("workflow.owned".to_string(), "active".to_string());
1666        })
1667        .await
1668        .expect("transaction")
1669        .expect("session exists");
1670
1671        let saved = storage
1672            .load_session(id)
1673            .await
1674            .expect("load durable")
1675            .expect("durable exists");
1676        assert_eq!(
1677            saved.messages.len(),
1678            1,
1679            "transaction never writes stale live messages"
1680        );
1681        assert_eq!(
1682            saved.metadata.get("external.durable").map(String::as_str),
1683            Some("keep")
1684        );
1685        assert_eq!(
1686            saved.metadata.get("workflow.owned").map(String::as_str),
1687            Some("active")
1688        );
1689
1690        let cached = read_cached_session(repo.cache(), id).expect("live cache");
1691        assert_eq!(
1692            cached.messages.len(),
1693            2,
1694            "cache live tool call is not replaced"
1695        );
1696        assert_eq!(cached.model, "live-model");
1697        assert_eq!(
1698            cached.metadata.get("external.live").map(String::as_str),
1699            Some("keep")
1700        );
1701        assert_eq!(
1702            cached.metadata.get("workflow.owned").map(String::as_str),
1703            Some("active")
1704        );
1705    }
1706
1707    /// Regression guard: a strictly-newer in-memory session (e.g. one that just
1708    /// answered and cleared its pending question) must win over a strictly-older
1709    /// storage copy that still carries the pending question — both in the value
1710    /// returned AND in the cache (no clobber).
1711    #[tokio::test]
1712    async fn load_merged_does_not_regress_to_older_storage() {
1713        let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
1714        let repo = test_repo(storage.clone());
1715        let id = "s1";
1716
1717        let mut stale = Session::new(id.to_string(), "m");
1718        stale.set_pending_question(
1719            "tc1".into(),
1720            "kind".into(),
1721            "q?".into(),
1722            vec!["OK".into()],
1723            true,
1724        );
1725        stale.updated_at = Utc::now() - chrono::Duration::seconds(10);
1726        storage.save_session(&stale).await.unwrap();
1727
1728        let mut fresh = Session::new(id.to_string(), "m");
1729        fresh.updated_at = Utc::now();
1730        cache_put(&repo, &fresh);
1731
1732        let merged = repo.load_merged(id).await.expect("session exists");
1733        assert!(
1734            merged.pending_question.is_none(),
1735            "must return the newer answered memory copy, not the stale storage one"
1736        );
1737        let cached = read_cached_session(repo.cache(), id).expect("cached");
1738        assert!(
1739            cached.pending_question.is_none(),
1740            "load_merged must never regress the cache to a stale storage copy"
1741        );
1742    }
1743
1744    /// The pending-question recovery still works when storage is the same age:
1745    /// if memory lost a pending question that same-age storage retains, prefer
1746    /// storage so a genuine clarification is not dropped.
1747    #[tokio::test]
1748    async fn load_merged_recovers_pending_question_from_same_age_storage() {
1749        let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
1750        let repo = test_repo(storage.clone());
1751        let id = "s2";
1752        let ts = Utc::now();
1753
1754        let mut with_pending = Session::new(id.to_string(), "m");
1755        with_pending.set_pending_question(
1756            "tc".into(),
1757            "k".into(),
1758            "q".into(),
1759            vec!["OK".into()],
1760            true,
1761        );
1762        with_pending.updated_at = ts;
1763        storage.save_session(&with_pending).await.unwrap();
1764
1765        let mut lost = with_pending.clone();
1766        lost.clear_pending_question();
1767        lost.updated_at = ts;
1768        cache_put(&repo, &lost);
1769
1770        let merged = repo.load_merged(id).await.expect("session exists");
1771        assert!(
1772            merged.pending_question.is_some(),
1773            "same-age storage carrying a pending question must still be recovered"
1774        );
1775    }
1776
1777    #[tokio::test]
1778    async fn runtime_publish_refreshes_cache_even_when_storage_fails() {
1779        let id = "runtime-selection";
1780        let mut previous = Session::new(id.to_string(), "m");
1781        previous.metadata.insert(
1782            "skill_runtime_selected_skill_ids".to_string(),
1783            "[\"plan\"]".to_string(),
1784        );
1785        let storage: Arc<dyn Storage> = Arc::new(FailingSaveStorage {
1786            persisted: Mutex::new(Some(previous.clone())),
1787        });
1788        let repo = test_repo(storage.clone());
1789        cache_put(&repo, &previous);
1790
1791        let mut current = previous.clone();
1792        current.metadata.insert(
1793            "skill_runtime_selected_skill_ids".to_string(),
1794            "[\"review\"]".to_string(),
1795        );
1796        current.updated_at = Utc::now();
1797
1798        let result =
1799            bamboo_domain::RuntimeSessionPersistence::save_runtime_session(&repo, &mut current)
1800                .await;
1801        assert!(result.is_err(), "durable failure must still be surfaced");
1802
1803        let cached = repo.load(id).await.expect("cached current session");
1804        assert_eq!(
1805            cached
1806                .metadata
1807                .get("skill_runtime_selected_skill_ids")
1808                .map(String::as_str),
1809            Some("[\"review\"]")
1810        );
1811        let allowlist = bamboo_skills::access_control::extract_skill_allowlist(&cached.metadata)
1812            .expect("runtime authorization allowlist");
1813        assert!(allowlist.contains("review"));
1814        assert!(!allowlist.contains("plan"));
1815        let durable = storage
1816            .load_session(id)
1817            .await
1818            .expect("load durable state")
1819            .expect("previous durable session");
1820        assert_eq!(
1821            durable
1822                .metadata
1823                .get("skill_runtime_selected_skill_ids")
1824                .map(String::as_str),
1825            Some("[\"plan\"]")
1826        );
1827    }
1828
1829    #[tokio::test]
1830    async fn inherent_save_leaves_existing_cache_untouched_when_storage_fails() {
1831        let id = "inherent-save-failure";
1832        let previous = Session::new(id, "previous");
1833        let storage: Arc<dyn Storage> = Arc::new(FailingSaveStorage {
1834            persisted: Mutex::new(Some(previous.clone())),
1835        });
1836        let repo = test_repo(storage);
1837        cache_put(&repo, &previous);
1838
1839        let mut current = previous.clone();
1840        current.model = "current".to_string();
1841        assert!(repo.save(&mut current).await.is_err());
1842        assert_eq!(
1843            read_cached_session(repo.cache(), id)
1844                .expect("existing cache")
1845                .model,
1846            "previous",
1847            "fallible inherent save must publish only after a durable commit"
1848        );
1849    }
1850
1851    #[tokio::test]
1852    async fn save_and_cache_still_refreshes_cache_when_storage_fails() {
1853        let id = "save-and-cache-failure";
1854        let previous = Session::new(id, "previous");
1855        let storage: Arc<dyn Storage> = Arc::new(FailingSaveStorage {
1856            persisted: Mutex::new(Some(previous.clone())),
1857        });
1858        let repo = test_repo(storage);
1859        cache_put(&repo, &previous);
1860
1861        let mut current = previous;
1862        current.model = "current".to_string();
1863        repo.save_and_cache(&mut current).await;
1864        assert_eq!(
1865            read_cached_session(repo.cache(), id)
1866                .expect("refreshed cache")
1867                .model,
1868            "current",
1869            "fire-and-forget save must retain its existing cache-on-failure behavior"
1870        );
1871    }
1872
1873    #[tokio::test]
1874    async fn checkpoint_leaves_existing_cache_untouched_when_storage_fails() {
1875        let id = "checkpoint-failure";
1876        let previous = Session::new(id, "previous");
1877        let storage: Arc<dyn Storage> = Arc::new(FailingSaveStorage {
1878            persisted: Mutex::new(Some(previous.clone())),
1879        });
1880        let repo = test_repo(storage);
1881        cache_put(&repo, &previous);
1882
1883        let mut current = previous.clone();
1884        current.model = "current".to_string();
1885        let result = bamboo_domain::RuntimeSessionPersistence::checkpoint_runtime_session(
1886            &repo,
1887            &mut current,
1888        )
1889        .await;
1890        assert!(result.is_err());
1891        assert_eq!(
1892            read_cached_session(repo.cache(), id)
1893                .expect("existing cache")
1894                .model,
1895            "previous",
1896            "checkpoint must publish only after a durable commit"
1897        );
1898    }
1899
1900    #[tokio::test]
1901    async fn legacy_clear_uses_durable_cas_and_never_erases_concurrent_append_from_stale_cache() {
1902        let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
1903        let repo = test_repo(storage.clone());
1904        let id = "legacy-cas-race";
1905        let expected = vec![serde_json::json!({"content": "first"})];
1906
1907        let mut stale_cache = Session::new(id, "m");
1908        stale_cache.set_pending_injected_messages(expected.clone());
1909        cache_put(&repo, &stale_cache);
1910
1911        let mut durable = stale_cache.clone();
1912        durable.set_pending_injected_messages(vec![
1913            serde_json::json!({"content": "first"}),
1914            serde_json::json!({"content": "concurrent"}),
1915        ]);
1916        storage.save_session(&durable).await.unwrap();
1917
1918        let cleared = bamboo_domain::RuntimeSessionPersistence::clear_legacy_pending_messages(
1919            &repo, id, &expected,
1920        )
1921        .await
1922        .unwrap();
1923        assert!(!cleared, "the durable compare-and-clear must reject drift");
1924        assert_eq!(
1925            storage
1926                .load_session(id)
1927                .await
1928                .unwrap()
1929                .unwrap()
1930                .pending_injected_messages()
1931                .unwrap(),
1932            durable.pending_injected_messages().unwrap(),
1933            "the concurrent durable append must remain intact"
1934        );
1935        assert_eq!(
1936            read_cached_session(repo.cache(), id)
1937                .unwrap()
1938                .pending_injected_messages()
1939                .unwrap(),
1940            expected,
1941            "a failed CAS must not mutate the existing cache"
1942        );
1943    }
1944
1945    #[tokio::test]
1946    async fn successful_legacy_clear_refreshes_stale_cache_from_durable_state() {
1947        let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
1948        let repo = test_repo(storage.clone());
1949        let id = "legacy-cas-success";
1950        let expected = vec![serde_json::json!({"content": "first"})];
1951
1952        let mut stale_cache = Session::new(id, "stale-model");
1953        stale_cache.set_pending_injected_messages(expected.clone());
1954        cache_put(&repo, &stale_cache);
1955
1956        let mut durable = Session::new(id, "durable-model");
1957        durable.set_pending_injected_messages(expected.clone());
1958        durable
1959            .metadata
1960            .insert("durable-only".to_string(), "keep".to_string());
1961        storage.save_session(&durable).await.unwrap();
1962
1963        assert!(
1964            bamboo_domain::RuntimeSessionPersistence::clear_legacy_pending_messages(
1965                &repo, id, &expected,
1966            )
1967            .await
1968            .unwrap()
1969        );
1970        let cached = read_cached_session(repo.cache(), id).unwrap();
1971        assert!(!cached.has_pending_injected_messages());
1972        assert_eq!(cached.model, "durable-model");
1973        assert_eq!(
1974            cached.metadata.get("durable-only").map(String::as_str),
1975            Some("keep")
1976        );
1977    }
1978}