Skip to main content

swink_agent_memory/
store_async.rs

1//! Async adapter over synchronous session stores.
2//!
3//! This module provides [`BlockingSessionStore`], which wraps any
4//! [`SessionStore`](crate::store::SessionStore) and exposes async methods by
5//! offloading each call to [`tokio::task::spawn_blocking`].
6//!
7//! # Why no async trait?
8//!
9//! A previous version of this module defined `AsyncSessionStore`, a trait that
10//! mirrored `SessionStore` with async signatures. That trait was removed because:
11//!
12//! - It offered no behaviour beyond bridging sync → async via `spawn_blocking`.
13//! - Its `load` signature accepted a per-call `registry` argument that
14//!   `BlockingSessionStore` could never honour (a `&CustomMessageRegistry`
15//!   reference cannot cross a `spawn_blocking` boundary), creating a silent
16//!   footgun in the API.
17//! - No implementation other than `BlockingSessionStore` existed or was planned.
18//!
19//! Callers that previously held `Box<dyn AsyncSessionStore>` should use
20//! `Arc<BlockingSessionStore<S>>` directly, which is callable from async
21//! contexts and has the same method set.
22
23use std::future::Future;
24use std::io;
25use std::pin::Pin;
26use std::sync::Arc;
27
28use swink_agent::{AgentMessage, CustomMessageRegistry};
29
30use crate::entry::SessionEntry;
31use crate::interrupt::InterruptState;
32use crate::load_options::LoadOptions;
33use crate::meta::SessionMeta;
34use crate::search::{SessionHit, SessionSearchOptions};
35
36/// A boxed future returned by [`BlockingSessionStore`] methods.
37pub type SessionStoreFuture<'a, T> = Pin<Box<dyn Future<Output = io::Result<T>> + Send + 'a>>;
38
39/// Run blocking filesystem work on Tokio's blocking pool.
40///
41/// A panic in `f` surfaces as an [`io::Error`] rather than being swallowed.
42pub(crate) fn spawn_store_call<T: Send + 'static>(
43    f: impl FnOnce() -> io::Result<T> + Send + 'static,
44) -> SessionStoreFuture<'static, T> {
45    Box::pin(async move {
46        tokio::task::spawn_blocking(f)
47            .await
48            .map_err(io::Error::other)?
49    })
50}
51
52/// Adapter that wraps a synchronous [`SessionStore`](crate::store::SessionStore)
53/// and exposes async methods by running each call via `tokio::task::spawn_blocking`.
54///
55/// Custom messages are preserved faithfully: `save`/`append` snapshot custom
56/// messages to their JSON envelope before crossing the thread boundary, and
57/// `load` uses the registry configured at construction time to restore them.
58///
59/// To restore custom messages on load, provide the registry once at construction
60/// via [`BlockingSessionStore::with_registry`].  A `&CustomMessageRegistry`
61/// reference cannot cross `spawn_blocking` boundaries, so a per-call registry
62/// is not supported.
63///
64/// Concurrent writes to the same session may corrupt the file.
65/// Callers are expected to enforce single-writer access.
66pub struct BlockingSessionStore<S: crate::store::SessionStore + 'static> {
67    inner: Arc<S>,
68    registry: Option<Arc<CustomMessageRegistry>>,
69}
70
71impl<S: crate::store::SessionStore + 'static> BlockingSessionStore<S> {
72    /// Create a new blocking adapter wrapping the given session store.
73    pub fn new(store: S) -> Self {
74        Self {
75            inner: Arc::new(store),
76            registry: None,
77        }
78    }
79
80    /// Attach a [`CustomMessageRegistry`] for deserializing custom messages on load.
81    ///
82    /// Because `&CustomMessageRegistry` cannot cross `spawn_blocking` boundaries,
83    /// the registry must be provided once at construction rather than per call.
84    #[must_use]
85    pub fn with_registry(mut self, registry: Arc<CustomMessageRegistry>) -> Self {
86        self.registry = Some(registry);
87        self
88    }
89}
90
91/// Clone messages for transfer across `spawn_blocking`.
92///
93/// Delegates to [`swink_agent::clone_messages_for_send`] which snapshots
94/// `Custom` variants into `SerializedCustomMessage` wrappers so they can
95/// cross thread boundaries faithfully.
96fn clone_messages_for_blocking(messages: &[AgentMessage]) -> Vec<AgentMessage> {
97    swink_agent::clone_messages_for_send(messages)
98}
99
100impl<S: crate::store::SessionStore + 'static> BlockingSessionStore<S> {
101    /// Persist a session asynchronously, including both LLM and custom messages.
102    pub fn save(
103        &self,
104        id: &str,
105        meta: &SessionMeta,
106        messages: &[AgentMessage],
107    ) -> SessionStoreFuture<'_, ()> {
108        let inner = Arc::clone(&self.inner);
109        let id = id.to_string();
110        let meta = meta.clone();
111        let messages = clone_messages_for_blocking(messages);
112        spawn_store_call(move || inner.save(&id, &meta, &messages))
113    }
114
115    /// Persist a session transcript plus its state snapshot asynchronously.
116    pub fn save_full(
117        &self,
118        id: &str,
119        meta: &SessionMeta,
120        messages: &[AgentMessage],
121        state: &serde_json::Value,
122    ) -> SessionStoreFuture<'_, SessionMeta> {
123        let inner = Arc::clone(&self.inner);
124        let id = id.to_string();
125        let meta = meta.clone();
126        let messages = clone_messages_for_blocking(messages);
127        let state = state.clone();
128        spawn_store_call(move || inner.save_full(&id, &meta, &messages, &state))
129    }
130
131    /// Append messages to an existing session asynchronously.
132    pub fn append(&self, id: &str, messages: &[AgentMessage]) -> SessionStoreFuture<'_, ()> {
133        let inner = Arc::clone(&self.inner);
134        let id = id.to_string();
135        let messages = clone_messages_for_blocking(messages);
136        spawn_store_call(move || inner.append(&id, &messages))
137    }
138
139    /// Load a session by ID asynchronously.
140    ///
141    /// Custom messages are restored using the registry supplied to
142    /// [`BlockingSessionStore::with_registry`]. Without a registry, custom
143    /// messages are discarded on load because the blocking adapter has no
144    /// serialized-wrapper fallback.
145    pub fn load(&self, id: &str) -> SessionStoreFuture<'_, (SessionMeta, Vec<AgentMessage>)> {
146        let inner = Arc::clone(&self.inner);
147        let id = id.to_string();
148        let registry = self.registry.clone();
149        spawn_store_call(move || inner.load(&id, registry.as_deref()))
150    }
151
152    /// Load a session transcript plus its state snapshot asynchronously from
153    /// one backend-defined read boundary.
154    pub fn load_full(
155        &self,
156        id: &str,
157    ) -> SessionStoreFuture<'_, (SessionMeta, Vec<AgentMessage>, Option<serde_json::Value>)> {
158        let inner = Arc::clone(&self.inner);
159        let id = id.to_string();
160        let registry = self.registry.clone();
161        spawn_store_call(move || inner.load_full(&id, registry.as_deref()))
162    }
163
164    /// List all saved sessions asynchronously.
165    pub fn list(&self) -> SessionStoreFuture<'_, Vec<SessionMeta>> {
166        let inner = Arc::clone(&self.inner);
167        spawn_store_call(move || inner.list())
168    }
169
170    /// Delete a session by ID asynchronously.
171    pub fn delete(&self, id: &str) -> SessionStoreFuture<'_, ()> {
172        let inner = Arc::clone(&self.inner);
173        let id = id.to_string();
174        spawn_store_call(move || inner.delete(&id))
175    }
176
177    /// Save session state snapshot asynchronously.
178    pub fn save_state(&self, id: &str, state: &serde_json::Value) -> SessionStoreFuture<'_, ()> {
179        let inner = Arc::clone(&self.inner);
180        let id = id.to_string();
181        let state = state.clone();
182        spawn_store_call(move || inner.save_state(&id, &state))
183    }
184
185    /// Load session state snapshot asynchronously. Returns `None` if not set.
186    pub fn load_state(&self, id: &str) -> SessionStoreFuture<'_, Option<serde_json::Value>> {
187        let inner = Arc::clone(&self.inner);
188        let id = id.to_string();
189        spawn_store_call(move || inner.load_state(&id))
190    }
191
192    /// Persist interrupt state for a session asynchronously.
193    pub fn save_interrupt(&self, id: &str, state: &InterruptState) -> SessionStoreFuture<'_, ()> {
194        let inner = Arc::clone(&self.inner);
195        let id = id.to_string();
196        let state = state.clone();
197        spawn_store_call(move || inner.save_interrupt(&id, &state))
198    }
199
200    /// Load interrupt state for a session asynchronously.
201    pub fn load_interrupt(&self, id: &str) -> SessionStoreFuture<'_, Option<InterruptState>> {
202        let inner = Arc::clone(&self.inner);
203        let id = id.to_string();
204        spawn_store_call(move || inner.load_interrupt(&id))
205    }
206
207    /// Clear interrupt state for a session asynchronously.
208    pub fn clear_interrupt(&self, id: &str) -> SessionStoreFuture<'_, ()> {
209        let inner = Arc::clone(&self.inner);
210        let id = id.to_string();
211        spawn_store_call(move || inner.clear_interrupt(&id))
212    }
213
214    /// Load a session with filtering options asynchronously.
215    pub fn load_with_options(
216        &self,
217        id: &str,
218        options: &LoadOptions,
219    ) -> SessionStoreFuture<'_, (SessionMeta, Vec<SessionEntry>)> {
220        let inner = Arc::clone(&self.inner);
221        let id = id.to_string();
222        let options = options.clone();
223        spawn_store_call(move || inner.load_with_options(&id, &options))
224    }
225
226    /// Search persisted sessions asynchronously.
227    pub fn search(
228        &self,
229        query: &str,
230        options: &SessionSearchOptions,
231    ) -> SessionStoreFuture<'_, Vec<SessionHit>> {
232        let inner = Arc::clone(&self.inner);
233        let query = query.to_string();
234        let options = options.clone();
235        spawn_store_call(move || inner.search(&query, &options))
236    }
237}
238
239#[cfg(test)]
240#[path = "store_async_tests.rs"]
241mod tests;