vv_agent/runtime/stores/
sqlite_v2.rs

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

vv_agent/runtime/stores/
sqlite_v2.rs
