vv_agent/
sessions.rs

1use std::collections::{BTreeMap, HashMap};
2use std::path::Path;
3use std::sync::{Arc, Mutex};
4
5use rusqlite::{params, Connection, OptionalExtension, TransactionBehavior};
6use serde::ser::{SerializeMap, SerializeSeq};
7use serde::{Deserialize, Deserializer, Serialize, Serializer};
8use serde_json::Value;
9use sha2::{Digest, Sha256};
10
11use crate::types::{Message, MessageRole, ToolCall};
12
13mod conformance;
14mod redis_store;
15
16pub use conformance::session_store_conformance;
17pub use redis_store::RedisSessionStore;
18
19pub trait Session: Send + Sync {
20    fn session_id(&self) -> &str;
21    fn get_items(&self, limit: Option<usize>) -> SessionFuture<Vec<SessionItem>>;
22    fn add_items(&self, items: Vec<SessionItem>) -> SessionFuture<()>;
23    fn pop_item(&self) -> SessionFuture<Option<SessionItem>>;
24    fn clear(&self) -> SessionFuture<()>;
25
26    fn supports_add_items_once(&self) -> bool {
27        false
28    }
29
30    fn add_items_once(
31        &self,
32        _commit_id: String,
33        _payload_digest: String,
34        _items: Vec<SessionItem>,
35    ) -> SessionFuture<SessionAppendOutcome> {
36        Box::pin(async {
37            Err("checkpoint_session_idempotency_unsupported: session does not support add_items_once"
38                .to_string())
39        })
40    }
41
42    fn clear_session(&self) -> SessionFuture<()> {
43        self.clear()
44    }
45}
46pub type SessionFuture<T> =
47    std::pin::Pin<Box<dyn std::future::Future<Output = Result<T, String>> + Send>>;
48
49#[derive(Debug, Clone, Copy, PartialEq, Eq)]
50pub enum SessionAppendOutcome {
51    Committed,
52    Replayed,
53}
54
55pub fn checkpoint_session_commit_id(checkpoint_key: &str) -> String {
56    let digest = Sha256::digest(checkpoint_key.as_bytes());
57    format!("vv-agent:checkpoint-v2:session:{digest:x}")
58}
59
60pub fn session_commit_payload_digest(items: &[SessionItem]) -> Result<String, String> {
61    let payload = serde_json::json!({
62        "schema_version": "vv-agent.session-commit.v1",
63        "items": items,
64    });
65    let bytes = crate::checkpoint::canonical_json_bytes(&payload, "session commit payload")
66        .map_err(|error| error.to_string())?;
67    Ok(format!("{:x}", Sha256::digest(bytes)))
68}
69
70fn validate_session_commit(
71    commit_id: &str,
72    payload_digest: &str,
73    items: &[SessionItem],
74) -> Result<(), String> {
75    if commit_id.trim().is_empty() {
76        return Err("session_commit_identity_invalid: commit_id must be non-empty".to_string());
77    }
78    if payload_digest.len() != 64
79        || !payload_digest
80            .bytes()
81            .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
82    {
83        return Err(
84            "session_commit_payload_digest_invalid: payload_digest must be lowercase SHA-256"
85                .to_string(),
86        );
87    }
88    if session_commit_payload_digest(items)? != payload_digest {
89        return Err(
90            "session_commit_payload_digest_mismatch: payload_digest does not match items"
91                .to_string(),
92        );
93    }
94    Ok(())
95}
96
97#[derive(Debug, Clone, PartialEq)]
98pub enum SessionItem {
99    User {
100        content: String,
101    },
102    Assistant {
103        content: String,
104    },
105    System {
106        content: String,
107    },
108    Tool {
109        content: String,
110        tool_call_id: String,
111    },
112    Message {
113        message: Message,
114    },
115}
116
117impl SessionItem {
118    pub fn to_message(&self) -> Message {
119        match self {
120            Self::User { content } => Message::user(content.clone()),
121            Self::Assistant { content } => Message::assistant(content.clone()),
122            Self::System { content } => Message::system(content.clone()),
123            Self::Tool {
124                content,
125                tool_call_id,
126            } => Message::tool(content.clone(), tool_call_id.clone()),
127            Self::Message { message } => message.clone(),
128        }
129    }
130
131    pub fn from_message(message: &Message) -> Option<Self> {
132        let has_unrepresentable_tool_call_id = match message.role {
133            MessageRole::Tool => message.tool_call_id.is_none(),
134            _ => message.tool_call_id.is_some(),
135        };
136        if has_unrepresentable_tool_call_id
137            || message.name.is_some()
138            || !message.tool_calls.is_empty()
139            || message.reasoning_content.is_some()
140            || message.image_url.is_some()
141            || !message.metadata.is_empty()
142        {
143            return Some(Self::Message {
144                message: message.clone(),
145            });
146        }
147        match message.role {
148            MessageRole::System => Some(Self::System {
149                content: message.content.clone(),
150            }),
151            MessageRole::User => Some(Self::User {
152                content: message.content.clone(),
153            }),
154            MessageRole::Assistant => Some(Self::Assistant {
155                content: message.content.clone(),
156            }),
157            MessageRole::Tool => Some(Self::Tool {
158                content: message.content.clone(),
159                tool_call_id: message.tool_call_id.clone().unwrap_or_default(),
160            }),
161        }
162    }
163}
164
165impl Serialize for SessionItem {
166    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
167    where
168        S: Serializer,
169    {
170        let message = self.to_message();
171        SessionMessageWire(&message).serialize(serializer)
172    }
173}
174
175impl<'de> Deserialize<'de> for SessionItem {
176    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
177    where
178        D: Deserializer<'de>,
179    {
180        let value = Value::deserialize(deserializer)?;
181        let wire = serde_json::from_value::<SessionMessageInput>(value)
182            .map_err(serde::de::Error::custom)?;
183        let message = wire.into_message().map_err(serde::de::Error::custom)?;
184        SessionItem::from_message(&message)
185            .ok_or_else(|| serde::de::Error::custom("unsupported session message role"))
186    }
187}
188
189#[derive(Deserialize)]
190#[serde(deny_unknown_fields)]
191struct SessionMessageInput {
192    role: String,
193    content: String,
194    #[serde(default, deserialize_with = "deserialize_present_string")]
195    name: Option<String>,
196    #[serde(default, deserialize_with = "deserialize_present_string")]
197    tool_call_id: Option<String>,
198    #[serde(default)]
199    tool_calls: Vec<SessionToolCallInput>,
200    #[serde(default, deserialize_with = "deserialize_present_string")]
201    reasoning_content: Option<String>,
202    #[serde(default, deserialize_with = "deserialize_present_string")]
203    image_url: Option<String>,
204    #[serde(default)]
205    metadata: BTreeMap<String, Value>,
206}
207
208impl SessionMessageInput {
209    fn into_message(self) -> Result<Message, String> {
210        let role = match self.role.as_str() {
211            "system" => MessageRole::System,
212            "user" => MessageRole::User,
213            "assistant" => MessageRole::Assistant,
214            "tool" => MessageRole::Tool,
215            value => return Err(format!("unknown message role: {value}")),
216        };
217        let tool_calls = self
218            .tool_calls
219            .into_iter()
220            .map(SessionToolCallInput::into_tool_call)
221            .collect::<Result<Vec<_>, _>>()?;
222        Ok(Message {
223            role,
224            content: self.content,
225            name: self.name,
226            tool_call_id: self.tool_call_id,
227            tool_calls,
228            reasoning_content: self.reasoning_content,
229            image_url: self.image_url,
230            metadata: self.metadata,
231        })
232    }
233}
234
235#[derive(Deserialize)]
236#[serde(deny_unknown_fields)]
237struct SessionToolCallInput {
238    id: String,
239    #[serde(rename = "type")]
240    kind: String,
241    function: SessionToolFunctionInput,
242    #[serde(default, deserialize_with = "deserialize_present_object")]
243    extra_content: Option<BTreeMap<String, Value>>,
244}
245
246impl SessionToolCallInput {
247    fn into_tool_call(self) -> Result<ToolCall, String> {
248        if self.id.trim().is_empty() {
249            return Err("session tool call id must be non-empty".to_string());
250        }
251        if self.kind != "function" {
252            return Err(format!("unknown session tool call type: {}", self.kind));
253        }
254        if self.function.name.trim().is_empty() {
255            return Err("session tool function name must be non-empty".to_string());
256        }
257        let arguments = serde_json::from_str::<Value>(&self.function.arguments)
258            .map_err(|error| format!("invalid session tool arguments JSON: {error}"))?;
259        let Value::Object(arguments) = arguments else {
260            return Err("session tool arguments must decode to an object".to_string());
261        };
262        Ok(ToolCall {
263            id: self.id,
264            name: self.function.name,
265            arguments: arguments.into_iter().collect(),
266            extra_content: self
267                .extra_content
268                .map(|value| Value::Object(value.into_iter().collect())),
269        })
270    }
271}
272
273#[derive(Deserialize)]
274#[serde(deny_unknown_fields)]
275struct SessionToolFunctionInput {
276    name: String,
277    arguments: String,
278}
279
280fn deserialize_present_string<'de, D>(deserializer: D) -> Result<Option<String>, D::Error>
281where
282    D: Deserializer<'de>,
283{
284    String::deserialize(deserializer).map(Some)
285}
286
287fn deserialize_present_object<'de, D>(
288    deserializer: D,
289) -> Result<Option<BTreeMap<String, Value>>, D::Error>
290where
291    D: Deserializer<'de>,
292{
293    BTreeMap::<String, Value>::deserialize(deserializer).map(Some)
294}
295
296struct SessionMessageWire<'a>(&'a Message);
297
298impl Serialize for SessionMessageWire<'_> {
299    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
300    where
301        S: Serializer,
302    {
303        let message = self.0;
304        let mut field_count = 2;
305        field_count += usize::from(message.name.is_some());
306        field_count += usize::from(message.tool_call_id.is_some());
307        field_count += usize::from(!message.tool_calls.is_empty());
308        field_count += usize::from(message.reasoning_content.is_some());
309        field_count += usize::from(message.image_url.is_some());
310        field_count += usize::from(!message.metadata.is_empty());
311
312        let mut state = serializer.serialize_map(Some(field_count))?;
313        state.serialize_entry("role", &message.role)?;
314        state.serialize_entry("content", &message.content)?;
315        if let Some(name) = &message.name {
316            state.serialize_entry("name", name)?;
317        }
318        if let Some(tool_call_id) = &message.tool_call_id {
319            state.serialize_entry("tool_call_id", tool_call_id)?;
320        }
321        if !message.tool_calls.is_empty() {
322            state.serialize_entry("tool_calls", &SessionToolCallsWire(&message.tool_calls))?;
323        }
324        if let Some(reasoning_content) = &message.reasoning_content {
325            state.serialize_entry("reasoning_content", reasoning_content)?;
326        }
327        if let Some(image_url) = &message.image_url {
328            state.serialize_entry("image_url", image_url)?;
329        }
330        if !message.metadata.is_empty() {
331            state.serialize_entry("metadata", &message.metadata)?;
332        }
333        state.end()
334    }
335}
336
337struct SessionToolCallsWire<'a>(&'a [ToolCall]);
338
339impl Serialize for SessionToolCallsWire<'_> {
340    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
341    where
342        S: Serializer,
343    {
344        let mut sequence = serializer.serialize_seq(Some(self.0.len()))?;
345        for tool_call in self.0 {
346            sequence.serialize_element(&SessionToolCallWire(tool_call))?;
347        }
348        sequence.end()
349    }
350}
351
352struct SessionToolCallWire<'a>(&'a ToolCall);
353
354impl Serialize for SessionToolCallWire<'_> {
355    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
356    where
357        S: Serializer,
358    {
359        let tool_call = self.0;
360        let field_count = 3 + usize::from(tool_call.extra_content.is_some());
361        let mut state = serializer.serialize_map(Some(field_count))?;
362        state.serialize_entry("id", &tool_call.id)?;
363        state.serialize_entry("type", "function")?;
364        state.serialize_entry("function", &SessionToolFunctionWire(tool_call))?;
365        if let Some(extra_content) = &tool_call.extra_content {
366            state.serialize_entry("extra_content", extra_content)?;
367        }
368        state.end()
369    }
370}
371
372struct SessionToolFunctionWire<'a>(&'a ToolCall);
373
374impl Serialize for SessionToolFunctionWire<'_> {
375    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
376    where
377        S: Serializer,
378    {
379        let mut state = serializer.serialize_map(Some(2))?;
380        state.serialize_entry("name", &self.0.name)?;
381        let arguments =
382            serde_json::to_string(&self.0.arguments).map_err(serde::ser::Error::custom)?;
383        state.serialize_entry("arguments", &arguments)?;
384        state.end()
385    }
386}
387
388#[derive(Clone)]
389pub struct MemorySession {
390    session_id: Arc<String>,
391    items: Arc<Mutex<Vec<SessionItem>>>,
392    commits: Arc<Mutex<HashMap<String, String>>>,
393}
394
395impl MemorySession {
396    pub fn new(session_id: impl Into<String>) -> Self {
397        Self {
398            session_id: Arc::new(session_id.into()),
399            items: Arc::new(Mutex::new(Vec::new())),
400            commits: Arc::new(Mutex::new(HashMap::new())),
401        }
402    }
403}
404
405impl Session for MemorySession {
406    fn session_id(&self) -> &str {
407        self.session_id.as_str()
408    }
409
410    fn get_items(&self, limit: Option<usize>) -> SessionFuture<Vec<SessionItem>> {
411        let items = self.items.clone();
412        Box::pin(async move {
413            let items = items
414                .lock()
415                .map_err(|_| "session lock poisoned".to_string())?;
416            let values = match limit {
417                Some(limit) => items
418                    .iter()
419                    .rev()
420                    .take(limit)
421                    .cloned()
422                    .collect::<Vec<_>>()
423                    .into_iter()
424                    .rev()
425                    .collect(),
426                None => items.clone(),
427            };
428            Ok(values)
429        })
430    }
431
432    fn add_items(&self, new_items: Vec<SessionItem>) -> SessionFuture<()> {
433        let items = self.items.clone();
434        Box::pin(async move {
435            items
436                .lock()
437                .map_err(|_| "session lock poisoned".to_string())?
438                .extend(new_items);
439            Ok(())
440        })
441    }
442
443    fn supports_add_items_once(&self) -> bool {
444        true
445    }
446
447    fn add_items_once(
448        &self,
449        commit_id: String,
450        payload_digest: String,
451        new_items: Vec<SessionItem>,
452    ) -> SessionFuture<SessionAppendOutcome> {
453        let items = self.items.clone();
454        let commits = self.commits.clone();
455        Box::pin(async move {
456            validate_session_commit(&commit_id, &payload_digest, &new_items)?;
457            let mut commits = commits
458                .lock()
459                .map_err(|_| "session commit lock poisoned".to_string())?;
460            if let Some(existing) = commits.get(&commit_id) {
461                if existing != &payload_digest {
462                    return Err(
463                        "session_commit_identity_conflict: commit_id has a different payload"
464                            .to_string(),
465                    );
466                }
467                return Ok(SessionAppendOutcome::Replayed);
468            }
469            items
470                .lock()
471                .map_err(|_| "session lock poisoned".to_string())?
472                .extend(new_items);
473            commits.insert(commit_id, payload_digest);
474            Ok(SessionAppendOutcome::Committed)
475        })
476    }
477
478    fn pop_item(&self) -> SessionFuture<Option<SessionItem>> {
479        let items = self.items.clone();
480        Box::pin(async move {
481            Ok(items
482                .lock()
483                .map_err(|_| "session lock poisoned".to_string())?
484                .pop())
485        })
486    }
487
488    fn clear(&self) -> SessionFuture<()> {
489        let items = self.items.clone();
490        let commits = self.commits.clone();
491        Box::pin(async move {
492            items
493                .lock()
494                .map_err(|_| "session lock poisoned".to_string())?
495                .clear();
496            commits
497                .lock()
498                .map_err(|_| "session commit lock poisoned".to_string())?
499                .clear();
500            Ok(())
501        })
502    }
503}
504
505pub trait SessionStore: Send + Sync {
506    fn session(&self, session_id: &str) -> Arc<dyn Session>;
507}
508
509#[derive(Clone, Default)]
510pub struct MemorySessionStore {
511    sessions: Arc<Mutex<HashMap<String, Arc<dyn Session>>>>,
512}
513
514impl MemorySessionStore {
515    pub fn new() -> Self {
516        Self::default()
517    }
518
519    pub fn session(&self, session_id: &str) -> Arc<dyn Session> {
520        <Self as SessionStore>::session(self, session_id)
521    }
522}
523
524impl SessionStore for MemorySessionStore {
525    fn session(&self, session_id: &str) -> Arc<dyn Session> {
526        let mut sessions = self
527            .sessions
528            .lock()
529            .expect("memory session store lock poisoned");
530        sessions
531            .entry(session_id.to_string())
532            .or_insert_with(|| Arc::new(MemorySession::new(session_id)))
533            .clone()
534    }
535}
536
537#[derive(Clone)]
538pub struct SqliteSessionStore {
539    connection: Arc<Mutex<Connection>>,
540}
541
542const SQLITE_SESSION_SCHEMA_VERSION: i64 = 1;
543const CANONICAL_SESSION_COLUMNS: [&str; 3] = ["session_id", "item_index", "payload"];
544const CANONICAL_SESSION_COMMIT_COLUMNS: [&str; 3] = ["session_id", "commit_id", "payload_digest"];
545const CREATE_SESSION_ITEMS_TABLE: &str = r#"
546    CREATE TABLE IF NOT EXISTS session_items (
547        session_id TEXT NOT NULL,
548        item_index INTEGER PRIMARY KEY AUTOINCREMENT,
549        payload TEXT NOT NULL
550    )
551"#;
552const CREATE_SESSION_ITEMS_INDEX: &str = r#"
553    CREATE INDEX IF NOT EXISTS idx_session_items_session_id_item_index
554        ON session_items (session_id, item_index)
555"#;
556const CREATE_SESSION_COMMITS_TABLE: &str = r#"
557    CREATE TABLE IF NOT EXISTS session_commits (
558        session_id TEXT NOT NULL,
559        commit_id TEXT NOT NULL,
560        payload_digest TEXT NOT NULL,
561        PRIMARY KEY (session_id, commit_id)
562    )
563"#;
564
565impl SqliteSessionStore {
566    pub fn open_memory() -> Result<Self, String> {
567        Self::open(":memory:")
568    }
569
570    pub fn open(path: impl AsRef<Path>) -> Result<Self, String> {
571        let mut connection = Connection::open(path).map_err(sqlite_error)?;
572        connection
573            .execute_batch(
574                r#"
575                PRAGMA busy_timeout = 5000;
576                PRAGMA journal_mode=WAL;
577                "#,
578            )
579            .map_err(sqlite_error)?;
580        initialize_sqlite_session_schema(&mut connection)?;
581        Ok(Self {
582            connection: Arc::new(Mutex::new(connection)),
583        })
584    }
585
586    pub fn session(&self, session_id: &str) -> Arc<dyn Session> {
587        <Self as SessionStore>::session(self, session_id)
588    }
589}
590
591impl SessionStore for SqliteSessionStore {
592    fn session(&self, session_id: &str) -> Arc<dyn Session> {
593        Arc::new(SqliteSession {
594            session_id: Arc::new(session_id.to_string()),
595            connection: self.connection.clone(),
596        })
597    }
598}
599
600#[derive(Clone)]
601struct SqliteSession {
602    session_id: Arc<String>,
603    connection: Arc<Mutex<Connection>>,
604}
605
606impl Session for SqliteSession {
607    fn session_id(&self) -> &str {
608        self.session_id.as_str()
609    }
610
611    fn get_items(&self, limit: Option<usize>) -> SessionFuture<Vec<SessionItem>> {
612        let session_id = self.session_id.to_string();
613        let connection = self.connection.clone();
614        Box::pin(async move {
615            let connection = connection
616                .lock()
617                .map_err(|_| "sqlite session store lock poisoned".to_string())?;
618            let mut statement = if limit.is_some() {
619                connection
620                    .prepare(
621                        r#"
622                        SELECT payload
623                        FROM (
624                            SELECT item_index, payload
625                            FROM session_items
626                            WHERE session_id = ?1
627                            ORDER BY item_index DESC
628                            LIMIT ?2
629                        )
630                        ORDER BY item_index ASC
631                        "#,
632                    )
633                    .map_err(sqlite_error)?
634            } else {
635                connection
636                    .prepare(
637                        r#"
638                        SELECT payload
639                        FROM session_items
640                        WHERE session_id = ?1
641                        ORDER BY item_index ASC
642                        "#,
643                    )
644                    .map_err(sqlite_error)?
645            };
646            let mut rows = if let Some(limit) = limit {
647                statement
648                    .query(params![
649                        session_id,
650                        i64::try_from(limit).unwrap_or(i64::MAX)
651                    ])
652                    .map_err(sqlite_error)?
653            } else {
654                statement.query(params![session_id]).map_err(sqlite_error)?
655            };
656            let mut items = Vec::new();
657            while let Some(row) = rows.next().map_err(sqlite_error)? {
658                let payload: String = row.get(0).map_err(sqlite_error)?;
659                items.push(serde_json::from_str(&payload).map_err(json_error)?);
660            }
661            Ok(items)
662        })
663    }
664
665    fn add_items(&self, items: Vec<SessionItem>) -> SessionFuture<()> {
666        let session_id = self.session_id.to_string();
667        let connection = self.connection.clone();
668        Box::pin(async move {
669            if items.is_empty() {
670                return Ok(());
671            }
672            let mut connection = connection
673                .lock()
674                .map_err(|_| "sqlite session store lock poisoned".to_string())?;
675            let transaction = connection.transaction().map_err(sqlite_error)?;
676            for item in items {
677                let payload = serde_json::to_string(&item).map_err(json_error)?;
678                transaction
679                    .execute(
680                        "INSERT INTO session_items (session_id, payload) VALUES (?1, ?2)",
681                        params![session_id, payload],
682                    )
683                    .map_err(sqlite_error)?;
684            }
685            transaction.commit().map_err(sqlite_error)?;
686            Ok(())
687        })
688    }
689
690    fn supports_add_items_once(&self) -> bool {
691        true
692    }
693
694    fn add_items_once(
695        &self,
696        commit_id: String,
697        payload_digest: String,
698        items: Vec<SessionItem>,
699    ) -> SessionFuture<SessionAppendOutcome> {
700        let session_id = self.session_id.to_string();
701        let connection = self.connection.clone();
702        Box::pin(async move {
703            validate_session_commit(&commit_id, &payload_digest, &items)?;
704            let payloads = items
705                .iter()
706                .map(serde_json::to_string)
707                .collect::<Result<Vec<_>, _>>()
708                .map_err(json_error)?;
709            let mut connection = connection
710                .lock()
711                .map_err(|_| "sqlite session store lock poisoned".to_string())?;
712            let transaction = connection
713                .transaction_with_behavior(TransactionBehavior::Immediate)
714                .map_err(sqlite_error)?;
715            let existing = transaction
716                .query_row(
717                    "SELECT payload_digest FROM session_commits WHERE session_id = ?1 AND commit_id = ?2",
718                    params![session_id, commit_id],
719                    |row| row.get::<_, String>(0),
720                )
721                .optional()
722                .map_err(sqlite_error)?;
723            if let Some(existing) = existing {
724                if existing != payload_digest {
725                    return Err(
726                        "session_commit_identity_conflict: commit_id has a different payload"
727                            .to_string(),
728                    );
729                }
730                transaction.commit().map_err(sqlite_error)?;
731                return Ok(SessionAppendOutcome::Replayed);
732            }
733            for payload in payloads {
734                transaction
735                    .execute(
736                        "INSERT INTO session_items (session_id, payload) VALUES (?1, ?2)",
737                        params![session_id, payload],
738                    )
739                    .map_err(sqlite_error)?;
740            }
741            transaction
742                .execute(
743                    "INSERT INTO session_commits (session_id, commit_id, payload_digest) VALUES (?1, ?2, ?3)",
744                    params![session_id, commit_id, payload_digest],
745                )
746                .map_err(sqlite_error)?;
747            transaction.commit().map_err(sqlite_error)?;
748            Ok(SessionAppendOutcome::Committed)
749        })
750    }
751
752    fn pop_item(&self) -> SessionFuture<Option<SessionItem>> {
753        let session_id = self.session_id.to_string();
754        let connection = self.connection.clone();
755        Box::pin(async move {
756            let mut connection = connection
757                .lock()
758                .map_err(|_| "sqlite session store lock poisoned".to_string())?;
759            let transaction = connection.transaction().map_err(sqlite_error)?;
760            let row = transaction
761                .query_row(
762                    r#"
763                    SELECT item_index, payload
764                    FROM session_items
765                    WHERE session_id = ?1
766                    ORDER BY item_index DESC
767                    LIMIT 1
768                    "#,
769                    params![session_id],
770                    |row| Ok((row.get::<_, i64>(0)?, row.get::<_, String>(1)?)),
771                )
772                .optional()
773                .map_err(sqlite_error)?;
774            let Some((item_index, payload)) = row else {
775                transaction.commit().map_err(sqlite_error)?;
776                return Ok(None);
777            };
778            let item = serde_json::from_str(&payload).map_err(json_error)?;
779            transaction
780                .execute(
781                    "DELETE FROM session_items WHERE item_index = ?1",
782                    params![item_index],
783                )
784                .map_err(sqlite_error)?;
785            transaction.commit().map_err(sqlite_error)?;
786            Ok(Some(item))
787        })
788    }
789
790    fn clear(&self) -> SessionFuture<()> {
791        let session_id = self.session_id.to_string();
792        let connection = self.connection.clone();
793        Box::pin(async move {
794            let mut connection = connection
795                .lock()
796                .map_err(|_| "sqlite session store lock poisoned".to_string())?;
797            let transaction = connection.transaction().map_err(sqlite_error)?;
798            transaction
799                .execute(
800                    "DELETE FROM session_items WHERE session_id = ?1",
801                    params![session_id],
802                )
803                .map_err(sqlite_error)?;
804            transaction
805                .execute(
806                    "DELETE FROM session_commits WHERE session_id = ?1",
807                    params![session_id],
808                )
809                .map_err(sqlite_error)?;
810            transaction.commit().map_err(sqlite_error)?;
811            Ok(())
812        })
813    }
814}
815
816fn initialize_sqlite_session_schema(connection: &mut Connection) -> Result<(), String> {
817    let transaction = connection
818        .transaction_with_behavior(TransactionBehavior::Immediate)
819        .map_err(sqlite_error)?;
820    let version = transaction
821        .query_row("PRAGMA user_version", [], |row| row.get::<_, i64>(0))
822        .map_err(sqlite_error)?;
823    let tables = session_table_names(&transaction)?;
824    if tables.is_empty() {
825        if version != 0 {
826            return Err(format!(
827                "empty session database has unexpected schema version {version}"
828            ));
829        }
830        transaction
831            .execute_batch(CREATE_SESSION_ITEMS_TABLE)
832            .map_err(sqlite_error)?;
833        transaction
834            .execute_batch(CREATE_SESSION_ITEMS_INDEX)
835            .map_err(sqlite_error)?;
836        transaction
837            .execute_batch(CREATE_SESSION_COMMITS_TABLE)
838            .map_err(sqlite_error)?;
839        transaction
840            .execute_batch("PRAGMA user_version = 1;")
841            .map_err(sqlite_error)?;
842    } else {
843        if version != SQLITE_SESSION_SCHEMA_VERSION {
844            return Err(format!(
845                "session schema version {version} does not match required version \
846                 {SQLITE_SESSION_SCHEMA_VERSION}"
847            ));
848        }
849        if tables != ["session_commits".to_string(), "session_items".to_string()] {
850            return Err(format!("unsupported session schema tables: {tables:?}"));
851        }
852        let item_columns = session_table_columns(&transaction, "session_items")?;
853        if item_columns != CANONICAL_SESSION_COLUMNS {
854            return Err(format!(
855                "unsupported session_items schema columns: {item_columns:?}"
856            ));
857        }
858        let commit_columns = session_table_columns(&transaction, "session_commits")?;
859        if commit_columns != CANONICAL_SESSION_COMMIT_COLUMNS {
860            return Err(format!(
861                "unsupported session_commits schema columns: {commit_columns:?}"
862            ));
863        }
864        let index_exists = transaction
865            .query_row(
866                "SELECT EXISTS(SELECT 1 FROM sqlite_master \
867                 WHERE type = 'index' AND name = 'idx_session_items_session_id_item_index')",
868                [],
869                |row| row.get::<_, i64>(0),
870            )
871            .map_err(sqlite_error)?
872            != 0;
873        if !index_exists {
874            return Err("session schema is missing the canonical session_items index".to_string());
875        }
876    }
877    transaction.commit().map_err(sqlite_error)
878}
879
880fn session_table_names(connection: &Connection) -> Result<Vec<String>, String> {
881    let mut statement = connection
882        .prepare(
883            "SELECT name FROM sqlite_master \
884             WHERE type = 'table' AND name NOT LIKE 'sqlite_%' ORDER BY name",
885        )
886        .map_err(sqlite_error)?;
887    let rows = statement
888        .query_map([], |row| row.get::<_, String>(0))
889        .map_err(sqlite_error)?;
890    rows.collect::<rusqlite::Result<Vec<_>>>()
891        .map_err(sqlite_error)
892}
893
894fn session_table_columns(connection: &Connection, table: &str) -> Result<Vec<String>, String> {
895    let mut statement = connection
896        .prepare(&format!("PRAGMA table_info({table})"))
897        .map_err(sqlite_error)?;
898    let rows = statement
899        .query_map([], |row| row.get::<_, String>(1))
900        .map_err(sqlite_error)?;
901    rows.collect::<rusqlite::Result<Vec<_>>>()
902        .map_err(sqlite_error)
903}
904
905fn sqlite_error(error: rusqlite::Error) -> String {
906    error.to_string()
907}
908
909fn json_error(error: serde_json::Error) -> String {
910    error.to_string()
911}
vv_agent/sessions.rs

vv_agent/
sessions.rs
