ai_agents_runtime/spec/
storage.rs

1//! Storage configuration types
2
3use serde::{Deserialize, Serialize};
4
5/// Storage configuration using tagged enum for type safety and extensibility
6#[derive(Debug, Clone, Serialize, Deserialize)]
7#[serde(tag = "type")]
8pub enum StorageConfig {
9    #[serde(rename = "none")]
10    None,
11
12    #[serde(rename = "file")]
13    File(FileStorageConfig),
14
15    #[serde(rename = "sqlite")]
16    Sqlite(SqliteStorageConfig),
17
18    #[serde(rename = "redis")]
19    Redis(RedisStorageConfig),
20}
21
22#[derive(Debug, Clone, Serialize, Deserialize)]
23pub struct FileStorageConfig {
24    pub path: String,
25}
26
27#[derive(Debug, Clone, Serialize, Deserialize)]
28pub struct SqliteStorageConfig {
29    pub path: String,
30
31    #[serde(default)]
32    pub table: Option<String>,
33}
34
35#[derive(Debug, Clone, Serialize, Deserialize)]
36pub struct RedisStorageConfig {
37    pub url: String,
38
39    #[serde(default)]
40    pub prefix: Option<String>,
41
42    #[serde(default)]
43    pub ttl_seconds: Option<u64>,
44}
45
46impl Default for StorageConfig {
47    fn default() -> Self {
48        StorageConfig::None
49    }
50}
51
52impl StorageConfig {
53    pub fn none() -> Self {
54        StorageConfig::None
55    }
56
57    pub fn file(path: impl Into<String>) -> Self {
58        StorageConfig::File(FileStorageConfig { path: path.into() })
59    }
60
61    pub fn sqlite(path: impl Into<String>) -> Self {
62        StorageConfig::Sqlite(SqliteStorageConfig {
63            path: path.into(),
64            table: None,
65        })
66    }
67
68    pub fn redis(url: impl Into<String>) -> Self {
69        StorageConfig::Redis(RedisStorageConfig {
70            url: url.into(),
71            prefix: None,
72            ttl_seconds: None,
73        })
74    }
75
76    pub fn is_none(&self) -> bool {
77        matches!(self, StorageConfig::None)
78    }
79
80    pub fn is_file(&self) -> bool {
81        matches!(self, StorageConfig::File(_))
82    }
83
84    pub fn is_sqlite(&self) -> bool {
85        matches!(self, StorageConfig::Sqlite(_))
86    }
87
88    pub fn is_redis(&self) -> bool {
89        matches!(self, StorageConfig::Redis(_))
90    }
91
92    pub fn storage_type(&self) -> &'static str {
93        match self {
94            StorageConfig::None => "none",
95            StorageConfig::File(_) => "file",
96            StorageConfig::Sqlite(_) => "sqlite",
97            StorageConfig::Redis(_) => "redis",
98        }
99    }
100
101    pub fn get_path(&self) -> Option<&str> {
102        match self {
103            StorageConfig::File(c) => Some(&c.path),
104            StorageConfig::Sqlite(c) => Some(&c.path),
105            _ => None,
106        }
107    }
108
109    pub fn get_url(&self) -> Option<&str> {
110        match self {
111            StorageConfig::Redis(c) => Some(&c.url),
112            _ => None,
113        }
114    }
115
116    pub fn get_prefix(&self) -> &str {
117        match self {
118            StorageConfig::Redis(c) => c.prefix.as_deref().unwrap_or("agent:"),
119            _ => "agent:",
120        }
121    }
122
123    pub fn get_ttl(&self) -> Option<u64> {
124        match self {
125            StorageConfig::Redis(c) => c.ttl_seconds,
126            _ => None,
127        }
128    }
129
130    pub fn get_table(&self) -> Option<&str> {
131        match self {
132            StorageConfig::Sqlite(c) => c.table.as_deref(),
133            _ => None,
134        }
135    }
136
137    pub fn as_file(&self) -> Option<&FileStorageConfig> {
138        match self {
139            StorageConfig::File(c) => Some(c),
140            _ => None,
141        }
142    }
143
144    pub fn as_sqlite(&self) -> Option<&SqliteStorageConfig> {
145        match self {
146            StorageConfig::Sqlite(c) => Some(c),
147            _ => None,
148        }
149    }
150
151    pub fn as_redis(&self) -> Option<&RedisStorageConfig> {
152        match self {
153            StorageConfig::Redis(c) => Some(c),
154            _ => None,
155        }
156    }
157}
158
159impl RedisStorageConfig {
160    pub fn with_prefix(mut self, prefix: impl Into<String>) -> Self {
161        self.prefix = Some(prefix.into());
162        self
163    }
164
165    pub fn with_ttl(mut self, ttl_seconds: u64) -> Self {
166        self.ttl_seconds = Some(ttl_seconds);
167        self
168    }
169}
170
171impl SqliteStorageConfig {
172    pub fn with_table(mut self, table: impl Into<String>) -> Self {
173        self.table = Some(table.into());
174        self
175    }
176}
177
178use ai_agents_storage::StorageConfig as StorageStorageConfig;
179
180/// Convert spec StorageConfig to storage crate StorageConfig for backend instantiation.
181pub fn to_storage_config(config: &StorageConfig) -> StorageStorageConfig {
182    match config {
183        StorageConfig::None => StorageStorageConfig::None,
184        StorageConfig::File(fc) => StorageStorageConfig::File {
185            path: fc.path.clone(),
186        },
187        StorageConfig::Sqlite(sc) => StorageStorageConfig::Sqlite {
188            path: sc.path.clone(),
189        },
190        StorageConfig::Redis(rc) => StorageStorageConfig::Redis {
191            url: rc.url.clone(),
192            prefix: rc.prefix.clone(),
193            ttl_seconds: rc.ttl_seconds,
194        },
195    }
196}
197
198#[cfg(test)]
199mod tests {
200    use super::*;
201
202    #[test]
203    fn test_storage_config_default() {
204        let config = StorageConfig::default();
205        assert!(config.is_none());
206        assert!(!config.is_file());
207        assert!(!config.is_sqlite());
208        assert!(!config.is_redis());
209        assert_eq!(config.storage_type(), "none");
210    }
211
212    #[test]
213    fn test_storage_config_none_yaml() {
214        let yaml = "type: none\n";
215        let config: StorageConfig = serde_yaml::from_str(yaml).unwrap();
216        assert!(config.is_none());
217    }
218
219    #[test]
220    fn test_storage_config_file() {
221        let yaml = r#"
222type: file
223path: "./data/sessions"
224"#;
225        let config: StorageConfig = serde_yaml::from_str(yaml).unwrap();
226        assert!(config.is_file());
227        assert_eq!(config.get_path(), Some("./data/sessions"));
228        assert_eq!(config.storage_type(), "file");
229    }
230
231    #[test]
232    fn test_storage_config_file_builder() {
233        let config = StorageConfig::file("./data/sessions");
234        assert!(config.is_file());
235        assert_eq!(config.get_path(), Some("./data/sessions"));
236    }
237
238    #[test]
239    fn test_storage_config_sqlite() {
240        let yaml = r#"
241type: sqlite
242path: "./data/sessions.db"
243"#;
244        let config: StorageConfig = serde_yaml::from_str(yaml).unwrap();
245        assert!(config.is_sqlite());
246        assert_eq!(config.get_path(), Some("./data/sessions.db"));
247        assert_eq!(config.storage_type(), "sqlite");
248    }
249
250    #[test]
251    fn test_storage_config_sqlite_with_table() {
252        let yaml = r#"
253type: sqlite
254path: "./data/sessions.db"
255table: "custom_sessions"
256"#;
257        let config: StorageConfig = serde_yaml::from_str(yaml).unwrap();
258        assert!(config.is_sqlite());
259        assert_eq!(config.get_table(), Some("custom_sessions"));
260    }
261
262    #[test]
263    fn test_storage_config_sqlite_builder() {
264        let config = StorageConfig::sqlite("./data/sessions.db");
265        assert!(config.is_sqlite());
266        assert_eq!(config.get_path(), Some("./data/sessions.db"));
267    }
268
269    #[test]
270    fn test_storage_config_redis() {
271        let yaml = r#"
272type: redis
273url: "redis://localhost:6379"
274prefix: "myagent:"
275ttl_seconds: 86400
276"#;
277        let config: StorageConfig = serde_yaml::from_str(yaml).unwrap();
278        assert!(config.is_redis());
279        assert_eq!(config.get_url(), Some("redis://localhost:6379"));
280        assert_eq!(config.get_prefix(), "myagent:");
281        assert_eq!(config.get_ttl(), Some(86400));
282        assert_eq!(config.storage_type(), "redis");
283    }
284
285    #[test]
286    fn test_storage_config_redis_builder() {
287        let config = StorageConfig::redis("redis://localhost:6379");
288        assert!(config.is_redis());
289        assert_eq!(config.get_url(), Some("redis://localhost:6379"));
290        assert_eq!(config.get_prefix(), "agent:");
291        assert_eq!(config.get_ttl(), None);
292    }
293
294    #[test]
295    fn test_storage_config_default_prefix() {
296        let config = StorageConfig::default();
297        assert_eq!(config.get_prefix(), "agent:");
298
299        let config = StorageConfig::file("./data");
300        assert_eq!(config.get_prefix(), "agent:");
301
302        let yaml = r#"
303type: redis
304url: "redis://localhost:6379"
305"#;
306        let config: StorageConfig = serde_yaml::from_str(yaml).unwrap();
307        assert_eq!(config.get_prefix(), "agent:");
308    }
309
310    #[test]
311    fn test_storage_config_accessors() {
312        let file_config = StorageConfig::file("./data");
313        assert!(file_config.as_file().is_some());
314        assert!(file_config.as_sqlite().is_none());
315        assert!(file_config.as_redis().is_none());
316
317        let sqlite_config = StorageConfig::sqlite("./data.db");
318        assert!(sqlite_config.as_file().is_none());
319        assert!(sqlite_config.as_sqlite().is_some());
320        assert!(sqlite_config.as_redis().is_none());
321
322        let redis_config = StorageConfig::redis("redis://localhost");
323        assert!(redis_config.as_file().is_none());
324        assert!(redis_config.as_sqlite().is_none());
325        assert!(redis_config.as_redis().is_some());
326    }
327
328    #[test]
329    fn test_redis_config_builder_methods() {
330        let config = RedisStorageConfig {
331            url: "redis://localhost:6379".to_string(),
332            prefix: None,
333            ttl_seconds: None,
334        }
335        .with_prefix("test:")
336        .with_ttl(3600);
337
338        assert_eq!(config.prefix, Some("test:".to_string()));
339        assert_eq!(config.ttl_seconds, Some(3600));
340    }
341
342    #[test]
343    fn test_sqlite_config_builder_methods() {
344        let config = SqliteStorageConfig {
345            path: "./data.db".to_string(),
346            table: None,
347        }
348        .with_table("custom_table");
349
350        assert_eq!(config.table, Some("custom_table".to_string()));
351    }
352
353    #[test]
354    fn test_storage_config_serialization() {
355        let config = StorageConfig::redis("redis://localhost:6379");
356        let yaml = serde_yaml::to_string(&config).unwrap();
357        assert!(yaml.contains("type: redis"));
358        assert!(yaml.contains("url: redis://localhost:6379"));
359    }
360
361    #[test]
362    fn test_to_storage_config_none() {
363        use ai_agents_storage::StorageConfig as SC;
364        let result = to_storage_config(&StorageConfig::None);
365        assert!(matches!(result, SC::None));
366    }
367
368    #[test]
369    fn test_to_storage_config_file() {
370        use ai_agents_storage::StorageConfig as SC;
371        let config = StorageConfig::file("./data/sessions");
372        let result = to_storage_config(&config);
373        match result {
374            SC::File { path } => assert_eq!(path, "./data/sessions"),
375            other => panic!("expected File, got {:?}", other),
376        }
377    }
378
379    #[test]
380    fn test_to_storage_config_sqlite() {
381        use ai_agents_storage::StorageConfig as SC;
382        let config = StorageConfig::sqlite("./data/db.sqlite");
383        let result = to_storage_config(&config);
384        match result {
385            SC::Sqlite { path } => assert_eq!(path, "./data/db.sqlite"),
386            other => panic!("expected Sqlite, got {:?}", other),
387        }
388    }
389}
ai_agents_runtime/spec/storage.rs

ai_agents_runtime/spec/
storage.rs
