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