Skip to main content

agent_sdk_store_supabase/
content.rs

1use agent_sdk_core::{
2    ContentResolutionError, ContentResolutionErrorKind, ContentResolutionPolicy,
3    ContentResolveRequest, ContentResolver, ContentStore, ResolvedContent, content::ContentRef,
4};
5use base64::{Engine as _, engine::general_purpose::STANDARD};
6use serde_json::json;
7
8use crate::client::SupabaseClient;
9
10#[derive(Clone)]
11/// Supabase-backed content resolver and store.
12pub struct SupabaseContentStore {
13    client: SupabaseClient,
14}
15
16impl SupabaseContentStore {
17    pub fn new(client: SupabaseClient) -> Self {
18        Self { client }
19    }
20}
21
22impl ContentResolver for SupabaseContentStore {
23    fn resolve(
24        &self,
25        request: ContentResolveRequest,
26        policy: ContentResolutionPolicy,
27    ) -> Result<ResolvedContent, ContentResolutionError> {
28        if !policy.allow_raw_content {
29            return Ok(ResolvedContent::redacted(
30                request.content_ref,
31                policy.policy_refs,
32            ));
33        }
34        let query = format!(
35            "store_scope=eq.{}&content_id=eq.{}&select=content_ref,bytes_base64&limit=1",
36            self.client.config().store_scope(),
37            request.content_ref.content_id.as_str()
38        );
39        let response = self
40            .client
41            .select("agent_sdk_content", &query)
42            .map_err(|_| {
43                content_error(
44                    ContentResolutionErrorKind::StorageUnavailable,
45                    request.content_ref.clone(),
46                )
47            })?;
48        if !(200..300).contains(&response.status) {
49            return Err(content_error(
50                ContentResolutionErrorKind::StorageUnavailable,
51                request.content_ref,
52            ));
53        }
54        let rows =
55            serde_json::from_slice::<Vec<serde_json::Value>>(&response.body).map_err(|_| {
56                content_error(
57                    ContentResolutionErrorKind::StorageUnavailable,
58                    request.content_ref.clone(),
59                )
60            })?;
61        let Some(row) = rows.into_iter().next() else {
62            return Err(content_error(
63                ContentResolutionErrorKind::Missing,
64                request.content_ref,
65            ));
66        };
67        let content_ref = serde_json::from_value::<ContentRef>(row["content_ref"].clone())
68            .unwrap_or(request.content_ref);
69        let encoded = row["bytes_base64"].as_str().ok_or_else(|| {
70            content_error(
71                ContentResolutionErrorKind::StorageUnavailable,
72                content_ref.clone(),
73            )
74        })?;
75        let bytes = STANDARD.decode(encoded).map_err(|_| {
76            content_error(
77                ContentResolutionErrorKind::StorageUnavailable,
78                content_ref.clone(),
79            )
80        })?;
81        if bytes.len() as u64 > policy.max_bytes {
82            return Err(content_error(
83                ContentResolutionErrorKind::MaxBytesExceeded,
84                content_ref,
85            ));
86        }
87        Ok(ResolvedContent {
88            mime: content_ref.mime.clone(),
89            redacted_summary: content_ref.redacted_summary.clone(),
90            content_ref,
91            bytes: Some(bytes),
92            policy_refs: policy.policy_refs,
93            raw_content_included: true,
94        })
95    }
96
97    fn store_resolved_content(
98        &self,
99        content_ref: &ContentRef,
100        bytes: Vec<u8>,
101    ) -> Result<(), ContentResolutionError> {
102        self.put_content(content_ref, bytes)
103    }
104}
105
106impl ContentStore for SupabaseContentStore {
107    fn put_content(
108        &self,
109        content_ref: &ContentRef,
110        bytes: Vec<u8>,
111    ) -> Result<(), ContentResolutionError> {
112        let response = self
113            .client
114            .insert(
115                "agent_sdk_content",
116                &json!({
117                    "store_scope": self.client.config().store_scope(),
118                    "content_id": content_ref.content_id.as_str(),
119                    "content_ref": content_ref,
120                    "bytes_base64": STANDARD.encode(&bytes),
121                    "byte_len": bytes.len(),
122                }),
123            )
124            .map_err(|_| {
125                content_error(
126                    ContentResolutionErrorKind::StorageUnavailable,
127                    content_ref.clone(),
128                )
129            })?;
130        if !(200..300).contains(&response.status) {
131            return Err(content_error(
132                ContentResolutionErrorKind::StorageUnavailable,
133                content_ref.clone(),
134            ));
135        }
136        Ok(())
137    }
138}
139
140fn content_error(
141    kind: ContentResolutionErrorKind,
142    content_ref: ContentRef,
143) -> ContentResolutionError {
144    ContentResolutionError {
145        kind,
146        redacted_summary: content_ref.redacted_summary.clone(),
147        content_ref: Box::new(content_ref),
148        policy_refs: Vec::new(),
149    }
150}