agent_sdk_store_sqlite/
content.rs

1use std::path::{Path, PathBuf};
2
3use agent_sdk_core::{
4    ContentResolutionError, ContentResolutionErrorKind, ContentResolutionPolicy,
5    ContentResolveRequest, ContentResolver, ContentStore, ResolvedContent,
6    content::{ContentRef, ResolvedContent as CoreResolvedContent},
7};
8use rusqlite::params;
9
10use crate::util::{content_error, decode, encode, open, sha256_hex};
11
12const SCHEMA: &str = "
13CREATE TABLE IF NOT EXISTS content (
14    content_ref TEXT PRIMARY KEY,
15    content_json TEXT NOT NULL,
16    bytes BLOB NOT NULL,
17    byte_len INTEGER NOT NULL
18);
19";
20
21#[derive(Clone, Debug)]
22/// SQLite-backed content resolver and store.
23pub struct SqliteContentStore {
24    path: PathBuf,
25}
26
27impl SqliteContentStore {
28    /// Opens or creates a SQLite content store.
29    pub fn open(path: impl AsRef<Path>) -> Result<Self, agent_sdk_core::AgentError> {
30        crate::util::init(path.as_ref(), SCHEMA)?;
31        Ok(Self {
32            path: path.as_ref().to_path_buf(),
33        })
34    }
35}
36
37impl ContentResolver for SqliteContentStore {
38    fn resolve(
39        &self,
40        request: ContentResolveRequest,
41        policy: ContentResolutionPolicy,
42    ) -> Result<ResolvedContent, ContentResolutionError> {
43        if request.requested_version != request.content_ref.version {
44            return Err(content_error(
45                ContentResolutionErrorKind::VersionMismatch,
46                request.content_ref,
47                policy.policy_refs,
48            ));
49        }
50        let connection = open(&self.path).map_err(|_| {
51            content_error(
52                ContentResolutionErrorKind::StorageUnavailable,
53                request.content_ref.clone(),
54                policy.policy_refs.clone(),
55            )
56        })?;
57        let mut statement = connection
58            .prepare("SELECT content_json, bytes FROM content WHERE content_ref = ?1")
59            .map_err(|_| {
60                content_error(
61                    ContentResolutionErrorKind::StorageUnavailable,
62                    request.content_ref.clone(),
63                    policy.policy_refs.clone(),
64                )
65            })?;
66        let row = statement
67            .query_row(params![request.content_ref.content_id.as_str()], |row| {
68                Ok((row.get::<_, String>(0)?, row.get::<_, Vec<u8>>(1)?))
69            })
70            .optional()
71            .map_err(|_| {
72                content_error(
73                    ContentResolutionErrorKind::StorageUnavailable,
74                    request.content_ref.clone(),
75                    policy.policy_refs.clone(),
76                )
77            })?
78            .ok_or_else(|| {
79                content_error(
80                    ContentResolutionErrorKind::Missing,
81                    request.content_ref.clone(),
82                    policy.policy_refs.clone(),
83                )
84            })?;
85        let content_ref: ContentRef = decode(&row.0).map_err(|_| {
86            content_error(
87                ContentResolutionErrorKind::StorageUnavailable,
88                request.content_ref.clone(),
89                policy.policy_refs.clone(),
90            )
91        })?;
92        if !policy.allow_raw_content {
93            return Ok(CoreResolvedContent::redacted(
94                content_ref,
95                policy.policy_refs,
96            ));
97        }
98        if row.1.len() as u64 > policy.max_bytes {
99            return Err(content_error(
100                ContentResolutionErrorKind::MaxBytesExceeded,
101                request.content_ref,
102                policy.policy_refs,
103            ));
104        }
105        if policy.require_hash_match {
106            if let Some(expected) = &content_ref.content_hash {
107                let actual = format!("sha256:{}", sha256_hex(&row.1));
108                if expected != &actual {
109                    return Err(content_error(
110                        ContentResolutionErrorKind::HashMismatch,
111                        request.content_ref,
112                        policy.policy_refs,
113                    ));
114                }
115            }
116        }
117        Ok(ResolvedContent {
118            mime: content_ref.mime.clone(),
119            redacted_summary: content_ref.redacted_summary.clone(),
120            content_ref,
121            bytes: Some(row.1),
122            policy_refs: policy.policy_refs,
123            raw_content_included: true,
124        })
125    }
126
127    fn store_resolved_content(
128        &self,
129        content_ref: &ContentRef,
130        bytes: Vec<u8>,
131    ) -> Result<(), ContentResolutionError> {
132        self.put_content(content_ref, bytes)
133    }
134}
135
136impl ContentStore for SqliteContentStore {
137    fn put_content(
138        &self,
139        content_ref: &ContentRef,
140        bytes: Vec<u8>,
141    ) -> Result<(), ContentResolutionError> {
142        let connection = open(&self.path).map_err(|_| {
143            content_error(
144                ContentResolutionErrorKind::StorageUnavailable,
145                content_ref.clone(),
146                Vec::new(),
147            )
148        })?;
149        connection
150            .execute(
151                "INSERT OR REPLACE INTO content
152                 (content_ref, content_json, bytes, byte_len)
153                 VALUES (?1, ?2, ?3, ?4)",
154                params![
155                    content_ref.content_id.as_str(),
156                    encode(content_ref).map_err(|_| {
157                        content_error(
158                            ContentResolutionErrorKind::StorageUnavailable,
159                            content_ref.clone(),
160                            Vec::new(),
161                        )
162                    })?,
163                    bytes,
164                    bytes.len() as i64,
165                ],
166            )
167            .map_err(|_| {
168                content_error(
169                    ContentResolutionErrorKind::StorageUnavailable,
170                    content_ref.clone(),
171                    Vec::new(),
172                )
173            })?;
174        Ok(())
175    }
176}
177
178use rusqlite::OptionalExtension;
agent_sdk_store_sqlite/content.rs

agent_sdk_store_sqlite/
content.rs
