vv_agent/runtime/stores/
sqlite.rs

1//! SQLite checkpoint v3 store.
2
3use std::path::{Path, PathBuf};
4use std::sync::Mutex;
5
6use rusqlite::{params, Connection, OptionalExtension, Transaction, TransactionBehavior};
7use serde_json::{Map, Value};
8
9use crate::checkpoint::{CheckpointError, CheckpointResult, ClaimMode, EventCursor};
10use crate::runtime::checkpoint_codec::{checkpoint_from_value, checkpoint_to_value};
11use crate::runtime::state::{
12    apply_claim, claim_candidate, prepare_ack, prepare_commit, prepare_event_delivery,
13    prepare_finalize, prepare_finalize_claimed, prepare_progress, prepare_suspend, Checkpoint,
14    CheckpointStore,
15};
16
17const MAX_EXTENSION_STATE_BYTES: u64 = crate::checkpoint::MAX_WIRE_INTEGER;
18const CREATE_CHECKPOINTS_TABLE_SQL: &str = r#"
19CREATE TABLE IF NOT EXISTS checkpoints (
20    checkpoint_key TEXT PRIMARY KEY,
21    schema_version TEXT NOT NULL CHECK (schema_version = 'vv-agent.checkpoint.v3'),
22    run_definition_schema TEXT NOT NULL CHECK (run_definition_schema = 'vv-agent.run-definition.v2'),
23    run_definition TEXT NOT NULL,
24    task_id TEXT NOT NULL,
25    root_run_id TEXT NOT NULL,
26    trace_id TEXT NOT NULL,
27    run_definition_digest TEXT NOT NULL,
28    resume_attempt INTEGER NOT NULL CHECK (resume_attempt >= 1),
29    cycle_index INTEGER NOT NULL CHECK (cycle_index >= 0),
30    status TEXT NOT NULL,
31    messages TEXT NOT NULL,
32    cycles TEXT NOT NULL,
33    model_calls TEXT NOT NULL,
34    shared_state TEXT NOT NULL,
35    budget_usage TEXT,
36    event_cursor TEXT,
37    event_outbox TEXT NOT NULL,
38    extension_state TEXT NOT NULL,
39    model_call_journal TEXT NOT NULL,
40    tool_journal TEXT NOT NULL,
41    revision INTEGER NOT NULL DEFAULT 0 CHECK (revision >= 0),
42    claim_token TEXT,
43    claimed_cycle INTEGER,
44    lease_expires_at_ms INTEGER,
45    terminal_result TEXT,
46    terminal_acknowledged INTEGER NOT NULL DEFAULT 0 CHECK (terminal_acknowledged IN (0, 1)),
47    CHECK (
48        (claim_token IS NULL AND claimed_cycle IS NULL AND lease_expires_at_ms IS NULL)
49        OR
50        (claim_token IS NOT NULL AND claimed_cycle IS NOT NULL AND lease_expires_at_ms IS NOT NULL)
51    ),
52    CHECK (claim_token IS NULL OR claimed_cycle = cycle_index + 1),
53    CHECK (terminal_result IS NULL OR claim_token IS NULL)
54)
55"#;
56const CREATE_CHECKPOINTS_STATUS_INDEX_SQL: &str = r#"
57CREATE INDEX IF NOT EXISTS checkpoints_status_idx ON checkpoints(status)
58"#;
59
60pub struct SqliteCheckpointStore {
61    connection: Mutex<Connection>,
62    location: PathBuf,
63}
64
65impl std::fmt::Debug for SqliteCheckpointStore {
66    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
67        formatter
68            .debug_struct("SqliteCheckpointStore")
69            .field("location", &self.location)
70            .finish_non_exhaustive()
71    }
72}
73
74impl SqliteCheckpointStore {
75    pub fn new(path: impl AsRef<Path>) -> CheckpointResult<Self> {
76        let path = path.as_ref().to_path_buf();
77        let connection = Connection::open(&path).map_err(sqlite_error)?;
78        initialize_schema(&connection)?;
79        Ok(Self {
80            connection: Mutex::new(connection),
81            location: path,
82        })
83    }
84
85    pub fn location(&self) -> &Path {
86        &self.location
87    }
88
89    pub fn save_checkpoint(&self, checkpoint: Checkpoint) -> CheckpointResult<()> {
90        checkpoint.validate()?;
91        let values = SqlValues::from_checkpoint(&checkpoint)?;
92        let connection = self.lock()?;
93        connection
94            .execute(
95                r#"
96                INSERT INTO checkpoints (
97                    checkpoint_key, schema_version, run_definition_schema, run_definition,
98                    task_id, root_run_id, trace_id, run_definition_digest, resume_attempt,
99                    cycle_index, status, messages, cycles, model_calls, shared_state,
100                    budget_usage, event_cursor, event_outbox, extension_state,
101                    model_call_journal, tool_journal, revision, claim_token,
102                    claimed_cycle, lease_expires_at_ms, terminal_result,
103                    terminal_acknowledged
104                ) VALUES (
105                    ?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14,
106                    ?15, ?16, ?17, ?18, ?19, ?20, ?21, ?22, ?23, ?24, ?25, ?26, ?27
107                )
108                ON CONFLICT(checkpoint_key) DO UPDATE SET
109                    schema_version = excluded.schema_version,
110                    run_definition_schema = excluded.run_definition_schema,
111                    run_definition = excluded.run_definition,
112                    task_id = excluded.task_id,
113                    root_run_id = excluded.root_run_id,
114                    trace_id = excluded.trace_id,
115                    run_definition_digest = excluded.run_definition_digest,
116                    resume_attempt = excluded.resume_attempt,
117                    cycle_index = excluded.cycle_index,
118                    status = excluded.status,
119                    messages = excluded.messages,
120                    cycles = excluded.cycles,
121                    model_calls = excluded.model_calls,
122                    shared_state = excluded.shared_state,
123                    budget_usage = excluded.budget_usage,
124                    event_cursor = excluded.event_cursor,
125                    event_outbox = excluded.event_outbox,
126                    extension_state = excluded.extension_state,
127                    model_call_journal = excluded.model_call_journal,
128                    tool_journal = excluded.tool_journal,
129                    revision = excluded.revision,
130                    claim_token = excluded.claim_token,
131                    claimed_cycle = excluded.claimed_cycle,
132                    lease_expires_at_ms = excluded.lease_expires_at_ms,
133                    terminal_result = excluded.terminal_result,
134                    terminal_acknowledged = excluded.terminal_acknowledged
135                "#,
136                values.params(),
137            )
138            .map_err(sqlite_error)?;
139        Ok(())
140    }
141
142    fn lock(&self) -> CheckpointResult<std::sync::MutexGuard<'_, Connection>> {
143        self.connection.lock().map_err(|_| {
144            CheckpointError::new(
145                "checkpoint_store_lock_poisoned",
146                "SQLite store lock poisoned",
147            )
148        })
149    }
150}
151
152fn initialize_schema(connection: &Connection) -> CheckpointResult<()> {
153    connection
154        .execute_batch("PRAGMA journal_mode=WAL;")
155        .map_err(sqlite_error)?;
156    match schema_sql(connection, "table", "checkpoints")? {
157        None => {
158            connection
159                .execute_batch(CREATE_CHECKPOINTS_TABLE_SQL)
160                .map_err(sqlite_error)?;
161            connection
162                .execute_batch(CREATE_CHECKPOINTS_STATUS_INDEX_SQL)
163                .map_err(sqlite_error)?;
164        }
165        Some(existing) => {
166            if normalize_schema_sql(&existing) != normalize_schema_sql(CREATE_CHECKPOINTS_TABLE_SQL)
167            {
168                return Err(schema_mismatch(
169                    "existing checkpoints table does not match the current schema; create a new database",
170                ));
171            }
172            let existing_index = schema_sql(connection, "index", "checkpoints_status_idx")?
173                .ok_or_else(|| {
174                    schema_mismatch(
175                        "existing checkpoints index does not match the current schema; create a new database",
176                    )
177                })?;
178            if normalize_schema_sql(&existing_index)
179                != normalize_schema_sql(CREATE_CHECKPOINTS_STATUS_INDEX_SQL)
180            {
181                return Err(schema_mismatch(
182                    "existing checkpoints index does not match the current schema; create a new database",
183                ));
184            }
185        }
186    }
187    Ok(())
188}
189
190fn schema_sql(
191    connection: &Connection,
192    object_type: &str,
193    name: &str,
194) -> CheckpointResult<Option<String>> {
195    connection
196        .query_row(
197            "SELECT sql FROM sqlite_master WHERE type = ?1 AND name = ?2",
198            params![object_type, name],
199            |row| row.get::<_, Option<String>>(0),
200        )
201        .optional()
202        .map(Option::flatten)
203        .map_err(sqlite_error)
204}
205
206fn normalize_schema_sql(sql: &str) -> String {
207    sql.replace("IF NOT EXISTS", "")
208        .split_whitespace()
209        .collect::<Vec<_>>()
210        .join(" ")
211}
212
213fn schema_mismatch(message: &str) -> CheckpointError {
214    CheckpointError::new("checkpoint_store_schema_mismatch", message)
215}
216
217impl CheckpointStore for SqliteCheckpointStore {
218    fn create_checkpoint(&self, checkpoint: Checkpoint) -> CheckpointResult<bool> {
219        checkpoint.validate()?;
220        let values = SqlValues::from_checkpoint(&checkpoint)?;
221        let connection = self.lock()?;
222        let changed = connection
223            .execute(
224                r#"
225                INSERT OR IGNORE INTO checkpoints (
226                    checkpoint_key, schema_version, run_definition_schema, run_definition,
227                    task_id, root_run_id, trace_id, run_definition_digest, resume_attempt,
228                    cycle_index, status, messages, cycles, model_calls, shared_state,
229                    budget_usage, event_cursor, event_outbox, extension_state,
230                    model_call_journal, tool_journal, revision, claim_token,
231                    claimed_cycle, lease_expires_at_ms, terminal_result,
232                    terminal_acknowledged
233                ) VALUES (
234                    ?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14,
235                    ?15, ?16, ?17, ?18, ?19, ?20, ?21, ?22, ?23, ?24, ?25, ?26, ?27
236                )
237                "#,
238                values.params(),
239            )
240            .map_err(sqlite_error)?;
241        Ok(changed == 1)
242    }
243
244    fn load_checkpoint(&self, checkpoint_key: &str) -> CheckpointResult<Option<Checkpoint>> {
245        let connection = self.lock()?;
246        load_row(&connection, checkpoint_key)
247    }
248
249    fn claim_checkpoint(
250        &self,
251        checkpoint_key: &str,
252        cycle_index: u64,
253        claim_token: &str,
254        lease_expires_at_ms: u64,
255        now_ms: u64,
256        claim_mode: ClaimMode,
257    ) -> CheckpointResult<Option<Checkpoint>> {
258        if claim_token.trim().is_empty() || lease_expires_at_ms <= now_ms {
259            return Err(CheckpointError::new(
260                "checkpoint_claim_invalid",
261                "claim token must be non-empty and lease must be in the future",
262            ));
263        }
264        let mut connection = self.lock()?;
265        let transaction = connection
266            .transaction_with_behavior(TransactionBehavior::Immediate)
267            .map_err(sqlite_error)?;
268        let Some(current) = load_row_transaction(&transaction, checkpoint_key)? else {
269            transaction.commit().map_err(sqlite_error)?;
270            return Ok(None);
271        };
272        if !claim_candidate(&current, cycle_index, now_ms, claim_mode)? {
273            transaction.commit().map_err(sqlite_error)?;
274            return Ok(None);
275        }
276        let mut claimed = current;
277        apply_claim(
278            &mut claimed,
279            cycle_index,
280            claim_token,
281            lease_expires_at_ms,
282            claim_mode,
283        )?;
284        let values = SqlValues::from_checkpoint(&claimed)?;
285        let changed = update_row(&transaction, &values, Some(claimed.revision - 1), None)?;
286        transaction.commit().map_err(sqlite_error)?;
287        if changed {
288            Ok(Some(claimed))
289        } else {
290            Ok(None)
291        }
292    }
293
294    fn progress_checkpoint(
295        &self,
296        checkpoint: Checkpoint,
297        claim_token: &str,
298        expected_revision: u64,
299    ) -> CheckpointResult<bool> {
300        self.replace_claimed(
301            checkpoint,
302            claim_token,
303            expected_revision,
304            ReplaceKind::Progress,
305        )
306    }
307
308    fn suspend_checkpoint(
309        &self,
310        checkpoint: Checkpoint,
311        claim_token: &str,
312        expected_revision: u64,
313    ) -> CheckpointResult<bool> {
314        self.replace_claimed(
315            checkpoint,
316            claim_token,
317            expected_revision,
318            ReplaceKind::Suspend,
319        )
320    }
321
322    fn commit_checkpoint(
323        &self,
324        checkpoint: Checkpoint,
325        claim_token: &str,
326        expected_revision: u64,
327    ) -> CheckpointResult<bool> {
328        self.replace_claimed(
329            checkpoint,
330            claim_token,
331            expected_revision,
332            ReplaceKind::Commit,
333        )
334    }
335
336    fn finalize_claimed_checkpoint(
337        &self,
338        checkpoint: Checkpoint,
339        claim_token: &str,
340        expected_revision: u64,
341    ) -> CheckpointResult<bool> {
342        self.replace_claimed(
343            checkpoint,
344            claim_token,
345            expected_revision,
346            ReplaceKind::FinalizeClaimed,
347        )
348    }
349
350    fn finalize_checkpoint(
351        &self,
352        checkpoint: Checkpoint,
353        expected_revision: u64,
354    ) -> CheckpointResult<bool> {
355        let mut connection = self.lock()?;
356        let transaction = connection
357            .transaction_with_behavior(TransactionBehavior::Immediate)
358            .map_err(sqlite_error)?;
359        let Some(current) = load_row_transaction(&transaction, &checkpoint.checkpoint_key)? else {
360            transaction.commit().map_err(sqlite_error)?;
361            return Ok(false);
362        };
363        let Some(updated) = prepare_finalize(&current, checkpoint, expected_revision)? else {
364            transaction.commit().map_err(sqlite_error)?;
365            return Ok(false);
366        };
367        let values = SqlValues::from_checkpoint(&updated)?;
368        let changed = update_row(&transaction, &values, Some(expected_revision), None)?;
369        transaction.commit().map_err(sqlite_error)?;
370        Ok(changed)
371    }
372
373    fn renew_checkpoint_claim(
374        &self,
375        checkpoint_key: &str,
376        claim_token: &str,
377        lease_expires_at_ms: u64,
378        now_ms: u64,
379    ) -> CheckpointResult<bool> {
380        if claim_token.trim().is_empty() || lease_expires_at_ms <= now_ms {
381            return Err(CheckpointError::new(
382                "checkpoint_claim_invalid",
383                "claim token must be non-empty and lease must be in the future",
384            ));
385        }
386        let mut connection = self.lock()?;
387        let transaction = connection
388            .transaction_with_behavior(TransactionBehavior::Immediate)
389            .map_err(sqlite_error)?;
390        let Some(current) = load_row_transaction(&transaction, checkpoint_key)? else {
391            transaction.commit().map_err(sqlite_error)?;
392            return Ok(false);
393        };
394        if current.claim_token.as_deref() != Some(claim_token)
395            || current
396                .lease_expires_at_ms
397                .is_none_or(|expiry| expiry <= now_ms)
398        {
399            transaction.commit().map_err(sqlite_error)?;
400            return Ok(false);
401        }
402        let changed = transaction
403            .execute(
404                "UPDATE checkpoints SET lease_expires_at_ms = ?1 WHERE checkpoint_key = ?2 AND claim_token = ?3 AND lease_expires_at_ms > ?4",
405                params![
406                    to_i64(lease_expires_at_ms, "lease_expires_at_ms")?,
407                    checkpoint_key,
408                    claim_token,
409                    to_i64(now_ms, "now_ms")?
410                ],
411            )
412            .map_err(sqlite_error)?;
413        transaction.commit().map_err(sqlite_error)?;
414        Ok(changed == 1)
415    }
416
417    fn acknowledge_terminal(
418        &self,
419        checkpoint_key: &str,
420        expected_revision: u64,
421    ) -> CheckpointResult<bool> {
422        let mut connection = self.lock()?;
423        let transaction = connection
424            .transaction_with_behavior(TransactionBehavior::Immediate)
425            .map_err(sqlite_error)?;
426        let Some(current) = load_row_transaction(&transaction, checkpoint_key)? else {
427            transaction.commit().map_err(sqlite_error)?;
428            return Ok(false);
429        };
430        let Some(updated) = prepare_ack(&current, expected_revision)? else {
431            transaction.commit().map_err(sqlite_error)?;
432            return Ok(false);
433        };
434        let values = SqlValues::from_checkpoint(&updated)?;
435        let changed = update_row(&transaction, &values, Some(expected_revision), None)?;
436        transaction.commit().map_err(sqlite_error)?;
437        Ok(changed)
438    }
439
440    fn record_event_delivery(
441        &self,
442        checkpoint_key: &str,
443        claim_token: Option<&str>,
444        expected_revision: u64,
445        event_id: &str,
446        payload_digest: &str,
447        cursor: EventCursor,
448    ) -> CheckpointResult<bool> {
449        let mut connection = self.lock()?;
450        let transaction = connection
451            .transaction_with_behavior(TransactionBehavior::Immediate)
452            .map_err(sqlite_error)?;
453        let Some(current) = load_row_transaction(&transaction, checkpoint_key)? else {
454            transaction.commit().map_err(sqlite_error)?;
455            return Ok(false);
456        };
457        let Some(updated) = prepare_event_delivery(
458            &current,
459            claim_token,
460            expected_revision,
461            event_id,
462            payload_digest,
463            cursor,
464        )?
465        else {
466            transaction.commit().map_err(sqlite_error)?;
467            return Ok(false);
468        };
469        let values = SqlValues::from_checkpoint(&updated)?;
470        let changed = update_row(&transaction, &values, Some(expected_revision), claim_token)?;
471        transaction.commit().map_err(sqlite_error)?;
472        Ok(changed)
473    }
474
475    fn delete_checkpoint(&self, checkpoint_key: &str) -> CheckpointResult<()> {
476        self.lock()?
477            .execute(
478                "DELETE FROM checkpoints WHERE checkpoint_key = ?1",
479                params![checkpoint_key],
480            )
481            .map_err(sqlite_error)?;
482        Ok(())
483    }
484
485    fn list_checkpoints(&self) -> CheckpointResult<Vec<String>> {
486        let connection = self.lock()?;
487        let mut statement = connection
488            .prepare("SELECT checkpoint_key FROM checkpoints ORDER BY checkpoint_key")
489            .map_err(sqlite_error)?;
490        let rows = statement
491            .query_map([], |row| row.get::<_, String>(0))
492            .map_err(sqlite_error)?;
493        rows.collect::<rusqlite::Result<Vec<_>>>()
494            .map_err(sqlite_error)
495    }
496}
497
498impl SqliteCheckpointStore {
499    fn replace_claimed(
500        &self,
501        checkpoint: Checkpoint,
502        claim_token: &str,
503        expected_revision: u64,
504        kind: ReplaceKind,
505    ) -> CheckpointResult<bool> {
506        let mut connection = self.lock()?;
507        let transaction = connection
508            .transaction_with_behavior(TransactionBehavior::Immediate)
509            .map_err(sqlite_error)?;
510        let Some(current) = load_row_transaction(&transaction, &checkpoint.checkpoint_key)? else {
511            transaction.commit().map_err(sqlite_error)?;
512            return Ok(false);
513        };
514        let updated = match kind {
515            ReplaceKind::Progress => {
516                prepare_progress(&current, checkpoint, claim_token, expected_revision)?
517            }
518            ReplaceKind::Suspend => {
519                prepare_suspend(&current, checkpoint, claim_token, expected_revision)?
520            }
521            ReplaceKind::Commit => {
522                prepare_commit(&current, checkpoint, claim_token, expected_revision)?
523            }
524            ReplaceKind::FinalizeClaimed => {
525                prepare_finalize_claimed(&current, checkpoint, claim_token, expected_revision)?
526            }
527        };
528        let Some(updated) = updated else {
529            transaction.commit().map_err(sqlite_error)?;
530            return Ok(false);
531        };
532        let values = SqlValues::from_checkpoint(&updated)?;
533        let changed = update_row(
534            &transaction,
535            &values,
536            Some(expected_revision),
537            Some(claim_token),
538        )?;
539        transaction.commit().map_err(sqlite_error)?;
540        Ok(changed)
541    }
542}
543
544#[derive(Clone, Copy)]
545enum ReplaceKind {
546    Progress,
547    Suspend,
548    Commit,
549    FinalizeClaimed,
550}
551
552struct SqlValues {
553    checkpoint_key: String,
554    schema_version: String,
555    run_definition_schema: String,
556    run_definition: String,
557    task_id: String,
558    root_run_id: String,
559    trace_id: String,
560    run_definition_digest: String,
561    resume_attempt: i64,
562    cycle_index: i64,
563    status: String,
564    messages: String,
565    cycles: String,
566    model_calls: String,
567    shared_state: String,
568    budget_usage: Option<String>,
569    event_cursor: Option<String>,
570    event_outbox: String,
571    extension_state: String,
572    model_call_journal: String,
573    tool_journal: String,
574    revision: i64,
575    claim_token: Option<String>,
576    claimed_cycle: Option<i64>,
577    lease_expires_at_ms: Option<i64>,
578    terminal_result: Option<String>,
579    terminal_acknowledged: i64,
580}
581
582impl SqlValues {
583    fn from_checkpoint(checkpoint: &Checkpoint) -> CheckpointResult<Self> {
584        let value = checkpoint_to_value(checkpoint, MAX_EXTENSION_STATE_BYTES)?;
585        let object = value.as_object().expect("codec emits an object");
586        Ok(Self {
587            checkpoint_key: string_field(object, "checkpoint_key")?,
588            schema_version: string_field(object, "schema_version")?,
589            run_definition_schema: string_field(object, "run_definition_schema")?,
590            run_definition: json_field(object, "run_definition")?,
591            task_id: string_field(object, "task_id")?,
592            root_run_id: string_field(object, "root_run_id")?,
593            trace_id: string_field(object, "trace_id")?,
594            run_definition_digest: string_field(object, "run_definition_digest")?,
595            resume_attempt: to_i64(checkpoint.resume_attempt, "resume_attempt")?,
596            cycle_index: to_i64(checkpoint.cycle_index, "cycle_index")?,
597            status: string_field(object, "status")?,
598            messages: json_field(object, "messages")?,
599            cycles: json_field(object, "cycles")?,
600            model_calls: json_field(object, "model_calls")?,
601            shared_state: json_field(object, "shared_state")?,
602            budget_usage: nullable_json_field(object, "budget_usage")?,
603            event_cursor: nullable_json_field(object, "event_cursor")?,
604            event_outbox: json_field(object, "event_outbox")?,
605            extension_state: json_field(object, "extension_state")?,
606            model_call_journal: json_field(object, "model_call_journal")?,
607            tool_journal: json_field(object, "tool_journal")?,
608            revision: to_i64(checkpoint.revision, "revision")?,
609            claim_token: checkpoint.claim_token.clone(),
610            claimed_cycle: checkpoint
611                .claimed_cycle
612                .map(|value| to_i64(value, "claimed_cycle"))
613                .transpose()?,
614            lease_expires_at_ms: checkpoint
615                .lease_expires_at_ms
616                .map(|value| to_i64(value, "lease_expires_at_ms"))
617                .transpose()?,
618            terminal_result: nullable_json_field(object, "terminal_result")?,
619            terminal_acknowledged: i64::from(checkpoint.terminal_acknowledged),
620        })
621    }
622
623    fn params(&self) -> [&(dyn rusqlite::ToSql + Sync); 27] {
624        [
625            &self.checkpoint_key,
626            &self.schema_version,
627            &self.run_definition_schema,
628            &self.run_definition,
629            &self.task_id,
630            &self.root_run_id,
631            &self.trace_id,
632            &self.run_definition_digest,
633            &self.resume_attempt,
634            &self.cycle_index,
635            &self.status,
636            &self.messages,
637            &self.cycles,
638            &self.model_calls,
639            &self.shared_state,
640            &self.budget_usage,
641            &self.event_cursor,
642            &self.event_outbox,
643            &self.extension_state,
644            &self.model_call_journal,
645            &self.tool_journal,
646            &self.revision,
647            &self.claim_token,
648            &self.claimed_cycle,
649            &self.lease_expires_at_ms,
650            &self.terminal_result,
651            &self.terminal_acknowledged,
652        ]
653    }
654}
655
656fn update_row(
657    transaction: &Transaction<'_>,
658    values: &SqlValues,
659    expected_revision: Option<u64>,
660    claim_token: Option<&str>,
661) -> CheckpointResult<bool> {
662    let Some(expected_revision) = expected_revision else {
663        return Err(CheckpointError::new(
664            "checkpoint_revision_conflict",
665            "an expected revision is required for an update",
666        ));
667    };
668    let changed = transaction
669        .execute(
670            r#"
671            UPDATE checkpoints SET
672                schema_version = ?1, run_definition_schema = ?2, run_definition = ?3,
673                task_id = ?4, root_run_id = ?5, trace_id = ?6, run_definition_digest = ?7,
674                resume_attempt = ?8, cycle_index = ?9, status = ?10, messages = ?11,
675                cycles = ?12, model_calls = ?13, shared_state = ?14, budget_usage = ?15,
676                event_cursor = ?16, event_outbox = ?17, extension_state = ?18,
677                model_call_journal = ?19, tool_journal = ?20, revision = ?21,
678                claim_token = ?22, claimed_cycle = ?23, lease_expires_at_ms = ?24,
679                terminal_result = ?25, terminal_acknowledged = ?26
680            WHERE checkpoint_key = ?27 AND revision = ?28
681              AND (?29 IS NULL OR claim_token = ?29)
682            "#,
683            params![
684                values.schema_version,
685                values.run_definition_schema,
686                values.run_definition,
687                values.task_id,
688                values.root_run_id,
689                values.trace_id,
690                values.run_definition_digest,
691                values.resume_attempt,
692                values.cycle_index,
693                values.status,
694                values.messages,
695                values.cycles,
696                values.model_calls,
697                values.shared_state,
698                values.budget_usage,
699                values.event_cursor,
700                values.event_outbox,
701                values.extension_state,
702                values.model_call_journal,
703                values.tool_journal,
704                values.revision,
705                values.claim_token,
706                values.claimed_cycle,
707                values.lease_expires_at_ms,
708                values.terminal_result,
709                values.terminal_acknowledged,
710                values.checkpoint_key,
711                to_i64(expected_revision, "revision")?,
712                claim_token,
713            ],
714        )
715        .map_err(sqlite_error)?;
716    Ok(changed == 1)
717}
718
719fn load_row(connection: &Connection, checkpoint_key: &str) -> CheckpointResult<Option<Checkpoint>> {
720    let mut statement = connection
721        .prepare(
722            r#"
723            SELECT checkpoint_key, schema_version, run_definition_schema, run_definition,
724                   task_id, root_run_id, trace_id, run_definition_digest, resume_attempt,
725                   cycle_index, status, messages, cycles, model_calls, shared_state,
726                   budget_usage, event_cursor, event_outbox, extension_state,
727                   model_call_journal, tool_journal, revision, claim_token, claimed_cycle,
728                   lease_expires_at_ms, terminal_result, terminal_acknowledged
729            FROM checkpoints WHERE checkpoint_key = ?1
730            "#,
731        )
732        .map_err(sqlite_error)?;
733    statement
734        .query_row(params![checkpoint_key], row_to_checkpoint)
735        .optional()
736        .map_err(sqlite_error)?
737        .transpose()
738}
739
740fn load_row_transaction(
741    transaction: &Transaction<'_>,
742    checkpoint_key: &str,
743) -> CheckpointResult<Option<Checkpoint>> {
744    let mut statement = transaction
745        .prepare(
746            r#"
747            SELECT checkpoint_key, schema_version, run_definition_schema, run_definition,
748                   task_id, root_run_id, trace_id, run_definition_digest, resume_attempt,
749                   cycle_index, status, messages, cycles, model_calls, shared_state,
750                   budget_usage, event_cursor, event_outbox, extension_state,
751                   model_call_journal, tool_journal, revision, claim_token, claimed_cycle,
752                   lease_expires_at_ms, terminal_result, terminal_acknowledged
753            FROM checkpoints WHERE checkpoint_key = ?1
754            "#,
755        )
756        .map_err(sqlite_error)?;
757    statement
758        .query_row(params![checkpoint_key], row_to_checkpoint)
759        .optional()
760        .map_err(sqlite_error)?
761        .transpose()
762}
763
764fn row_to_checkpoint(row: &rusqlite::Row<'_>) -> rusqlite::Result<CheckpointResult<Checkpoint>> {
765    let checkpoint_key: String = row.get(0)?;
766    let schema_version: String = row.get(1)?;
767    let run_definition_schema: String = row.get(2)?;
768    let run_definition: String = row.get(3)?;
769    let task_id: String = row.get(4)?;
770    let root_run_id: String = row.get(5)?;
771    let trace_id: String = row.get(6)?;
772    let run_definition_digest: String = row.get(7)?;
773    let resume_attempt: i64 = row.get(8)?;
774    let cycle_index: i64 = row.get(9)?;
775    let status: String = row.get(10)?;
776    let messages: String = row.get(11)?;
777    let cycles: String = row.get(12)?;
778    let model_calls: String = row.get(13)?;
779    let shared_state: String = row.get(14)?;
780    let budget_usage: Option<String> = row.get(15)?;
781    let event_cursor: Option<String> = row.get(16)?;
782    let event_outbox: String = row.get(17)?;
783    let extension_state: String = row.get(18)?;
784    let model_call_journal: String = row.get(19)?;
785    let tool_journal: String = row.get(20)?;
786    let revision: i64 = row.get(21)?;
787    let claim_token: Option<String> = row.get(22)?;
788    let claimed_cycle: Option<i64> = row.get(23)?;
789    let lease_expires_at_ms: Option<i64> = row.get(24)?;
790    let terminal_result: Option<String> = row.get(25)?;
791    let terminal_acknowledged: i64 = row.get(26)?;
792
793    let result = (|| {
794        let mut object = Map::new();
795        object.insert("schema_version".to_string(), Value::String(schema_version));
796        object.insert(
797            "run_definition_schema".to_string(),
798            Value::String(run_definition_schema),
799        );
800        object.insert("run_definition".to_string(), parse_value(&run_definition)?);
801        object.insert("checkpoint_key".to_string(), Value::String(checkpoint_key));
802        object.insert("task_id".to_string(), Value::String(task_id));
803        object.insert("root_run_id".to_string(), Value::String(root_run_id));
804        object.insert("trace_id".to_string(), Value::String(trace_id));
805        object.insert(
806            "run_definition_digest".to_string(),
807            Value::String(run_definition_digest),
808        );
809        object.insert(
810            "resume_attempt".to_string(),
811            Value::from(to_u64(resume_attempt)?),
812        );
813        object.insert("cycle_index".to_string(), Value::from(to_u64(cycle_index)?));
814        object.insert("status".to_string(), Value::String(status));
815        object.insert("messages".to_string(), parse_value(&messages)?);
816        object.insert("cycles".to_string(), parse_value(&cycles)?);
817        object.insert("model_calls".to_string(), parse_value(&model_calls)?);
818        object.insert("shared_state".to_string(), parse_value(&shared_state)?);
819        object.insert(
820            "budget_usage".to_string(),
821            optional_value(budget_usage.as_deref())?,
822        );
823        object.insert(
824            "event_cursor".to_string(),
825            optional_value(event_cursor.as_deref())?,
826        );
827        object.insert("event_outbox".to_string(), parse_value(&event_outbox)?);
828        object.insert(
829            "extension_state".to_string(),
830            parse_value(&extension_state)?,
831        );
832        object.insert(
833            "model_call_journal".to_string(),
834            parse_value(&model_call_journal)?,
835        );
836        object.insert("tool_journal".to_string(), parse_value(&tool_journal)?);
837        object.insert("revision".to_string(), Value::from(to_u64(revision)?));
838        object.insert(
839            "claim_token".to_string(),
840            claim_token.map_or(Value::Null, Value::String),
841        );
842        object.insert(
843            "claimed_cycle".to_string(),
844            claimed_cycle.map_or(Ok(Value::Null), |value| to_u64(value).map(Value::from))?,
845        );
846        object.insert(
847            "lease_expires_at_ms".to_string(),
848            lease_expires_at_ms.map_or(Ok(Value::Null), |value| to_u64(value).map(Value::from))?,
849        );
850        object.insert(
851            "terminal_result".to_string(),
852            optional_value(terminal_result.as_deref())?,
853        );
854        object.insert(
855            "terminal_acknowledged".to_string(),
856            Value::Bool(terminal_acknowledged != 0),
857        );
858        checkpoint_from_value(&Value::Object(object), MAX_EXTENSION_STATE_BYTES)
859    })();
860    Ok(result)
861}
862
863fn string_field(object: &Map<String, Value>, field: &str) -> CheckpointResult<String> {
864    object
865        .get(field)
866        .and_then(Value::as_str)
867        .map(str::to_string)
868        .ok_or_else(|| {
869            CheckpointError::new("checkpoint_row_invalid", format!("{field} is not a string"))
870        })
871}
872
873fn json_field(object: &Map<String, Value>, field: &str) -> CheckpointResult<String> {
874    serde_json::to_string(object.get(field).ok_or_else(|| {
875        CheckpointError::new("checkpoint_row_invalid", format!("{field} is missing"))
876    })?)
877    .map_err(|error| CheckpointError::new("checkpoint_json_invalid", error.to_string()))
878}
879
880fn nullable_json_field(
881    object: &Map<String, Value>,
882    field: &str,
883) -> CheckpointResult<Option<String>> {
884    match object.get(field) {
885        None | Some(Value::Null) => Ok(None),
886        Some(value) => serde_json::to_string(value)
887            .map(Some)
888            .map_err(|error| CheckpointError::new("checkpoint_json_invalid", error.to_string())),
889    }
890}
891
892fn parse_value(raw: &str) -> CheckpointResult<Value> {
893    serde_json::from_str(raw)
894        .map_err(|error| CheckpointError::new("checkpoint_json_invalid", error.to_string()))
895}
896
897fn optional_value(raw: Option<&str>) -> CheckpointResult<Value> {
898    raw.map_or(Ok(Value::Null), parse_value)
899}
900
901fn to_i64(value: u64, field: &str) -> CheckpointResult<i64> {
902    i64::try_from(value).map_err(|_| {
903        CheckpointError::new(
904            "checkpoint_integer_invalid",
905            format!("{field} does not fit SQLite INTEGER"),
906        )
907    })
908}
909
910fn to_u64(value: i64) -> CheckpointResult<u64> {
911    u64::try_from(value).map_err(|_| {
912        CheckpointError::new(
913            "checkpoint_row_invalid",
914            "negative SQLite integer in checkpoint",
915        )
916    })
917}
918
919fn sqlite_error(error: rusqlite::Error) -> CheckpointError {
920    CheckpointError::new("checkpoint_store_sqlite", error.to_string())
921}
vv_agent/runtime/stores/sqlite.rs

vv_agent/runtime/stores/
sqlite.rs
