Skip to main content

runifold_store_sqlite/store/
conversation.rs

1//! Conversation persistence and atomic terminal checkpoint commit.
2
3use std::sync::Arc;
4
5use futures_executor::block_on;
6use runifold_agent::{
7    ConversationAppend, ConversationCreateOutcome, ConversationId, ConversationSequence,
8    ConversationStore, ConversationStoreError, ConversationStoreErrorKind, ConversationStoreFuture,
9    ConversationSummary, ConversationSummaryBatch, ConversationSummaryCommit,
10    ConversationTranscriptEntry, ConversationVersion, ConversationView, ConversationWindow,
11    DurableConversationCommit, DurableConversationStore, InMemoryConversationStore,
12    MemoryNamespace, SemanticMemory, SemanticMemoryQuery, SemanticMemoryUpsert,
13};
14use rusqlite::{OptionalExtension, Transaction, TransactionBehavior, params};
15
16use super::{SqliteStore, current_revision, sqlite_revision};
17
18const SNAPSHOT_FORMAT_VERSION: i64 = 1;
19
20impl SqliteStore {
21    fn execute_conversation<T, F>(
22        &self,
23        operation: F,
24    ) -> ConversationStoreFuture<'_, Result<T, ConversationStoreError>>
25    where
26        T: Send + 'static,
27        F: FnOnce(&InMemoryConversationStore) -> Result<T, ConversationStoreError> + Send + 'static,
28    {
29        let connection = Arc::clone(&self.connection);
30        Box::pin(async move {
31            let runtime = tokio::runtime::Handle::try_current().map_err(|_| {
32                storage_error("SQLite conversation operations require a Tokio runtime")
33            })?;
34            runtime
35                .spawn_blocking(move || {
36                    let mut connection = connection
37                        .lock()
38                        .unwrap_or_else(std::sync::PoisonError::into_inner);
39                    let transaction = connection
40                        .transaction_with_behavior(TransactionBehavior::Immediate)
41                        .map_err(|error| database_error(&error))?;
42                    let state = load_state(&transaction)?;
43                    let output = operation(&state)?;
44                    save_state(&transaction, &state)?;
45                    transaction
46                        .commit()
47                        .map_err(|error| database_error(&error))?;
48                    Ok(output)
49                })
50                .await
51                .map_err(|error| {
52                    storage_error(format!("SQLite conversation task failed: {error}"))
53                })?
54        })
55    }
56}
57
58impl ConversationStore for SqliteStore {
59    fn create(
60        &self,
61        conversation_id: ConversationId,
62        namespace: MemoryNamespace,
63    ) -> ConversationStoreFuture<'_, Result<ConversationCreateOutcome, ConversationStoreError>>
64    {
65        self.execute_conversation(move |store| block_on(store.create(conversation_id, namespace)))
66    }
67
68    fn load_view(
69        &self,
70        conversation_id: ConversationId,
71        namespace: MemoryNamespace,
72        window: ConversationWindow,
73        summary_batch: ConversationSummaryBatch,
74    ) -> ConversationStoreFuture<'_, Result<ConversationView, ConversationStoreError>> {
75        self.execute_conversation(move |store| {
76            block_on(store.load_view(conversation_id, namespace, window, summary_batch))
77        })
78    }
79
80    fn list_transcript(
81        &self,
82        conversation_id: ConversationId,
83        namespace: MemoryNamespace,
84        after: Option<ConversationSequence>,
85        limit: ConversationWindow,
86    ) -> ConversationStoreFuture<'_, Result<Vec<ConversationTranscriptEntry>, ConversationStoreError>>
87    {
88        self.execute_conversation(move |store| {
89            block_on(store.list_transcript(conversation_id, namespace, after, limit))
90        })
91    }
92
93    fn append(
94        &self,
95        namespace: MemoryNamespace,
96        command: ConversationAppend,
97    ) -> ConversationStoreFuture<'_, Result<ConversationVersion, ConversationStoreError>> {
98        self.execute_conversation(move |store| block_on(store.append(namespace, command)))
99    }
100
101    fn commit_summary(
102        &self,
103        namespace: MemoryNamespace,
104        command: ConversationSummaryCommit,
105    ) -> ConversationStoreFuture<'_, Result<ConversationSummary, ConversationStoreError>> {
106        self.execute_conversation(move |store| block_on(store.commit_summary(namespace, command)))
107    }
108
109    fn upsert_memory(
110        &self,
111        command: SemanticMemoryUpsert,
112    ) -> ConversationStoreFuture<'_, Result<SemanticMemory, ConversationStoreError>> {
113        self.execute_conversation(move |store| block_on(store.upsert_memory(command)))
114    }
115
116    fn search_memory(
117        &self,
118        query: SemanticMemoryQuery,
119    ) -> ConversationStoreFuture<'_, Result<Vec<SemanticMemory>, ConversationStoreError>> {
120        self.execute_conversation(move |store| block_on(store.search_memory(query)))
121    }
122}
123
124impl DurableConversationStore for SqliteStore {
125    fn commit_durable_turn(
126        &self,
127        command: DurableConversationCommit,
128    ) -> ConversationStoreFuture<'_, Result<ConversationVersion, ConversationStoreError>> {
129        let connection = Arc::clone(&self.connection);
130        Box::pin(async move {
131            let runtime = tokio::runtime::Handle::try_current().map_err(|_| {
132                storage_error("SQLite durable conversation commit requires a Tokio runtime")
133            })?;
134            runtime
135                .spawn_blocking(move || {
136                    let mut connection = connection
137                        .lock()
138                        .unwrap_or_else(std::sync::PoisonError::into_inner);
139                    let transaction = connection
140                        .transaction_with_behavior(TransactionBehavior::Immediate)
141                        .map_err(|error| database_error(&error))?;
142                    let state = load_state(&transaction)?;
143                    let version = block_on(state.append(command.namespace, command.append))?;
144                    compare_and_swap_checkpoint(
145                        &transaction,
146                        &command.checkpoint,
147                        command.expected_checkpoint_revision,
148                    )?;
149                    save_state(&transaction, &state)?;
150                    transaction
151                        .commit()
152                        .map_err(|error| database_error(&error))?;
153                    Ok(version)
154                })
155                .await
156                .map_err(|error| {
157                    storage_error(format!("SQLite durable commit task failed: {error}"))
158                })?
159        })
160    }
161}
162
163fn load_state(
164    transaction: &Transaction<'_>,
165) -> Result<InMemoryConversationStore, ConversationStoreError> {
166    let stored = transaction
167        .query_row(
168            "SELECT format_version, state_blob
169             FROM runifold_conversation_state
170             WHERE singleton_id = 1",
171            [],
172            |row| Ok((row.get::<_, i64>(0)?, row.get::<_, Vec<u8>>(1)?)),
173        )
174        .optional()
175        .map_err(|error| database_error(&error))?;
176    match stored {
177        Some((SNAPSHOT_FORMAT_VERSION, encoded)) => {
178            InMemoryConversationStore::from_persistent_snapshot(&encoded)
179        }
180        Some((format_version, _)) => Err(storage_error(format!(
181            "unsupported SQLite conversation state format version {format_version}"
182        ))),
183        None => Ok(InMemoryConversationStore::new()),
184    }
185}
186
187fn save_state(
188    transaction: &Transaction<'_>,
189    state: &InMemoryConversationStore,
190) -> Result<(), ConversationStoreError> {
191    let encoded = state.export_persistent_snapshot()?;
192    transaction
193        .execute(
194            "INSERT INTO runifold_conversation_state (
195                 singleton_id, format_version, state_blob, updated_at_ms
196             ) VALUES (1, ?1, ?2,
197                 CAST(strftime('%s', 'now') AS INTEGER) * 1000
198                 + CAST(substr(strftime('%f', 'now'), 4, 3) AS INTEGER))
199             ON CONFLICT(singleton_id) DO UPDATE SET
200                 format_version = excluded.format_version,
201                 state_blob = excluded.state_blob,
202                 updated_at_ms = excluded.updated_at_ms",
203            params![SNAPSHOT_FORMAT_VERSION, encoded],
204        )
205        .map_err(|error| database_error(&error))?;
206    Ok(())
207}
208
209fn compare_and_swap_checkpoint(
210    transaction: &Transaction<'_>,
211    checkpoint: &runifold_core::Checkpoint,
212    expected_revision: u64,
213) -> Result<(), ConversationStoreError> {
214    let revision = sqlite_revision(checkpoint.revision).map_err(storage_error)?;
215    let expected = sqlite_revision(expected_revision).map_err(storage_error)?;
216    if revision
217        != expected.checked_add(1).ok_or_else(|| {
218            conflict_error("checkpoint revision overflow during durable conversation commit")
219        })?
220    {
221        return Err(conflict_error(
222            "durable conversation checkpoint revision is not the expected successor",
223        ));
224    }
225    let current = current_revision(
226        transaction,
227        "runifold_checkpoints",
228        "checkpoint_id",
229        &checkpoint.id.to_string(),
230    )
231    .map_err(|error| database_error(&error))?;
232    if current != Some(expected) {
233        return Err(conflict_error(
234            "durable conversation checkpoint revision precondition failed",
235        ));
236    }
237    let record = serde_json::to_string(checkpoint)
238        .map_err(|error| storage_error(format!("checkpoint encoding failed: {error}")))?;
239    let changed = transaction
240        .execute(
241            "UPDATE runifold_checkpoints
242             SET revision = ?1, record_json = ?2
243             WHERE checkpoint_id = ?3 AND revision = ?4",
244            params![revision, record, checkpoint.id.to_string(), expected],
245        )
246        .map_err(|error| database_error(&error))?;
247    if changed != 1 {
248        return Err(conflict_error(
249            "durable conversation checkpoint compare-and-swap failed",
250        ));
251    }
252    Ok(())
253}
254
255fn database_error(error: &rusqlite::Error) -> ConversationStoreError {
256    storage_error(format!("SQLite conversation operation failed: {error}"))
257}
258
259fn storage_error(message: impl Into<String>) -> ConversationStoreError {
260    ConversationStoreError::new(ConversationStoreErrorKind::Storage, message)
261}
262
263fn conflict_error(message: impl Into<String>) -> ConversationStoreError {
264    ConversationStoreError::new(ConversationStoreErrorKind::Conflict, message)
265}