Skip to main content

bamboo_domain/session/
persistence.rs

1use std::io;
2use std::sync::Arc;
3
4use crate::session::types::Session;
5
6/// Merge messages from a live runner snapshot into an already-durable
7/// transcript without ever removing or rewriting a durable message.
8///
9/// Runtime sessions are append-oriented, and every newly-created message has a
10/// stable id.  A runner may nevertheless be holding a snapshot that predates a
11/// concurrent append (for example, an injected child-completion message).  A
12/// terminal/error checkpoint must not full-save that stale snapshot: doing so
13/// would shrink the transcript.  Keep the durable ordering and append only the
14/// live messages whose ids are not durable yet.
15pub fn append_missing_runtime_messages(session: &mut Session, durable: &Session) -> usize {
16    let mut seen = durable
17        .messages
18        .iter()
19        .map(|message| message.id.clone())
20        .collect::<std::collections::HashSet<_>>();
21    let missing = session
22        .messages
23        .iter()
24        .filter(|message| seen.insert(message.id.clone()))
25        .cloned()
26        .collect::<Vec<_>>();
27    let appended = missing.len();
28    session.messages = durable.messages.iter().cloned().chain(missing).collect();
29    appended
30}
31
32/// Port for runtime (non-authoritative) session persistence.
33///
34/// Implementors must:
35/// - Serialize concurrent saves per session ID.
36/// - Merge on-disk authoritative metadata (`title`, `pinned`, `title_version`,
37///   `metadata_version`) before writing, so UI edits are never clobbered.
38#[async_trait::async_trait]
39pub trait RuntimeSessionPersistence: Send + Sync {
40    /// Persist the session, merging any newer authoritative metadata from disk.
41    async fn save_runtime_session(&self, session: &mut Session) -> io::Result<()>;
42
43    /// Append-safe checkpoint used at the shared engine execute boundary.
44    ///
45    /// Unlike [`save_runtime_session`](Self::save_runtime_session), this must
46    /// preserve messages that were appended durably by a concurrent writer
47    /// after the runner loaded its snapshot.  Implementations that can provide
48    /// a per-session transaction should override this method and perform the
49    /// load/merge/save under one lock.  The default still reconciles against a
50    /// latest snapshot for lightweight/custom SDK persisters; the built-in
51    /// storage implementation supplies the atomic variant.
52    async fn checkpoint_runtime_session(&self, session: &mut Session) -> io::Result<()> {
53        if let Some(durable) = self.load_runtime_session(&session.id).await? {
54            append_missing_runtime_messages(session, &durable);
55        }
56        self.save_runtime_session(session).await
57    }
58
59    /// Load the latest runtime-visible session snapshot when the persistence
60    /// implementation can coordinate reads. Tools may update a repository-owned
61    /// clone while an agent loop holds its own live Session; the loop uses this
62    /// hook to merge narrowly-scoped tool side effects before its next save.
63    async fn load_runtime_session(&self, _session_id: &str) -> io::Result<Option<Session>> {
64        Ok(None)
65    }
66
67    /// Append one JSON-line analysis record to the session's append-only
68    /// token-usage log (see [`Storage::append_token_usage_record`]). Defaults to
69    /// a no-op so non-file-backed persisters are unaffected.
70    ///
71    /// [`Storage::append_token_usage_record`]: crate::storage::Storage::append_token_usage_record
72    async fn append_token_usage_record(&self, session_id: &str, json_line: &str) -> io::Result<()> {
73        let _ = (session_id, json_line);
74        Ok(())
75    }
76}
77
78#[async_trait::async_trait]
79impl<T: RuntimeSessionPersistence + ?Sized> RuntimeSessionPersistence for Arc<T> {
80    async fn save_runtime_session(&self, session: &mut Session) -> io::Result<()> {
81        (**self).save_runtime_session(session).await
82    }
83
84    async fn checkpoint_runtime_session(&self, session: &mut Session) -> io::Result<()> {
85        (**self).checkpoint_runtime_session(session).await
86    }
87
88    async fn load_runtime_session(&self, session_id: &str) -> io::Result<Option<Session>> {
89        (**self).load_runtime_session(session_id).await
90    }
91
92    async fn append_token_usage_record(&self, session_id: &str, json_line: &str) -> io::Result<()> {
93        (**self)
94            .append_token_usage_record(session_id, json_line)
95            .await
96    }
97}