swink_agent_memory/
store_async.rs1use 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
36pub type SessionStoreFuture<'a, T> = Pin<Box<dyn Future<Output = io::Result<T>> + Send + 'a>>;
38
39pub(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
52pub 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 pub fn new(store: S) -> Self {
74 Self {
75 inner: Arc::new(store),
76 registry: None,
77 }
78 }
79
80 #[must_use]
85 pub fn with_registry(mut self, registry: Arc<CustomMessageRegistry>) -> Self {
86 self.registry = Some(registry);
87 self
88 }
89}
90
91fn 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 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 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 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 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 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 pub fn list(&self) -> SessionStoreFuture<'_, Vec<SessionMeta>> {
166 let inner = Arc::clone(&self.inner);
167 spawn_store_call(move || inner.list())
168 }
169
170 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 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 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 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 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 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 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 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;