agent-sdk-store-supabase 0.1.0-alpha.4

Supabase REST-backed durable store adapters for the Agent SDK.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127

use agent_sdk_core::{
    AgentError, CheckpointPrunePolicy, CheckpointPruneReport, CheckpointSaveOutcome,
    CheckpointStore, JournalCursor, RunCheckpoint, RunId,
};
use serde_json::json;

use crate::{client::SupabaseClient, transport::supabase_error};

#[derive(Clone)]
/// Supabase-backed checkpoint store.
pub struct SupabaseCheckpointStore {
    client: SupabaseClient,
}

impl SupabaseCheckpointStore {
    pub fn new(client: SupabaseClient) -> Self {
        Self { client }
    }
}

impl CheckpointStore for SupabaseCheckpointStore {
    fn save(
        &self,
        checkpoint: RunCheckpoint,
        latest_journal_seq: u64,
    ) -> Result<CheckpointSaveOutcome, AgentError> {
        checkpoint.validate_against_latest_seq(latest_journal_seq)?;
        let checkpoint_ref = checkpoint.checkpoint_id.clone();
        let covers_journal_seq = checkpoint.covers_journal_seq;
        let response = self.client.insert(
            "agent_sdk_checkpoints",
            &json!({
                "store_scope": self.client.config().store_scope(),
                "run_id": checkpoint.run_id.as_str(),
                "checkpoint_id": checkpoint.checkpoint_id,
                "covers_journal_seq": checkpoint.covers_journal_seq,
                "checkpoint": checkpoint,
            }),
        )?;
        if !(200..300).contains(&response.status) {
            return Err(supabase_error(format!(
                "supabase checkpoint save failed with status {}",
                response.status
            )));
        }
        Ok(CheckpointSaveOutcome {
            checkpoint_ref,
            covers_journal_seq,
            terminal_checkpoint: false,
        })
    }

    fn load_latest(&self, run_id: &RunId) -> Result<Option<RunCheckpoint>, AgentError> {
        let query = format!(
            "store_scope=eq.{}&run_id=eq.{}&select=checkpoint&order=covers_journal_seq.desc&limit=1",
            self.client.config().store_scope(),
            run_id.as_str()
        );
        load_checkpoint(self.client.select("agent_sdk_checkpoints", &query)?)
    }

    fn load_at_or_before(
        &self,
        run_id: &RunId,
        cursor: &JournalCursor,
    ) -> Result<Option<RunCheckpoint>, AgentError> {
        let seq = cursor
            .as_str()
            .strip_prefix("journal.")
            .unwrap_or(cursor.as_str());
        let query = format!(
            "store_scope=eq.{}&run_id=eq.{}&covers_journal_seq=lte.{}&select=checkpoint&order=covers_journal_seq.desc&limit=1",
            self.client.config().store_scope(),
            run_id.as_str(),
            seq
        );
        load_checkpoint(self.client.select("agent_sdk_checkpoints", &query)?)
    }

    fn prune(
        &self,
        run_id: &RunId,
        policy: CheckpointPrunePolicy,
    ) -> Result<CheckpointPruneReport, AgentError> {
        let response = self.client.rpc(
            "agent_sdk_prune_checkpoints",
            &json!({
                "p_store_scope": self.client.config().store_scope(),
                "p_run_id": run_id.as_str(),
                "p_prune_covered_before": policy.prune_covered_before,
                "p_preserve_latest_terminal": policy.preserve_latest_terminal,
            }),
        )?;
        if !(200..300).contains(&response.status) {
            return Err(supabase_error(format!(
                "supabase checkpoint prune failed with status {}",
                response.status
            )));
        }
        Ok(CheckpointPruneReport {
            run_id: run_id.clone(),
            pruned_count: 0,
            retained_count: 0,
            preserved_terminal_checkpoint: None,
        })
    }
}

fn load_checkpoint(
    response: crate::transport::SupabaseHttpResponse,
) -> Result<Option<RunCheckpoint>, AgentError> {
    if !(200..300).contains(&response.status) {
        return Err(supabase_error(format!(
            "supabase checkpoint read failed with status {}",
            response.status
        )));
    }
    let rows = serde_json::from_slice::<Vec<serde_json::Value>>(&response.body)
        .map_err(|error| supabase_error(error.to_string()))?;
    rows.into_iter()
        .next()
        .map(|row| {
            serde_json::from_value::<RunCheckpoint>(row["checkpoint"].clone())
                .map_err(|error| supabase_error(error.to_string()))
        })
        .transpose()
}