vv_agent/
sessions.rs

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

vv_agent/
sessions.rs
