ai_agents_runtime/spec/
storage.rs

1//! Storage configuration types
2
3use serde::{Deserialize, Deserializer, Serialize};
4
5/// Storage configuration using tagged enum for type safety and extensibility
6#[derive(Debug, Clone, Serialize, Default)]
7#[serde(tag = "type")]
8pub enum StorageConfig {
9    #[default]
10    #[serde(rename = "none")]
11    None,
12
13    #[serde(rename = "file")]
14    File(FileStorageConfig),
15
16    #[serde(rename = "sqlite")]
17    Sqlite(SqliteStorageConfig),
18
19    #[serde(rename = "redis")]
20    Redis(RedisStorageConfig),
21}
22
23#[derive(Deserialize)]
24#[serde(tag = "type", rename_all = "lowercase", deny_unknown_fields)]
25enum StorageConfigWire {
26    None {},
27    File {
28        path: String,
29    },
30    Sqlite {
31        path: String,
32        #[serde(default)]
33        table: Option<String>,
34    },
35    Redis {
36        url: String,
37        #[serde(default)]
38        prefix: Option<String>,
39        #[serde(default)]
40        ttl_seconds: Option<u64>,
41    },
42}
43
44impl<'de> Deserialize<'de> for StorageConfig {
45    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
46    where
47        D: Deserializer<'de>,
48    {
49        Ok(match StorageConfigWire::deserialize(deserializer)? {
50            StorageConfigWire::None {} => Self::None,
51            StorageConfigWire::File { path } => Self::File(FileStorageConfig { path }),
52            StorageConfigWire::Sqlite { path, table } => {
53                Self::Sqlite(SqliteStorageConfig { path, table })
54            }
55            StorageConfigWire::Redis {
56                url,
57                prefix,
58                ttl_seconds,
59            } => Self::Redis(RedisStorageConfig {
60                url,
61                prefix,
62                ttl_seconds,
63            }),
64        })
65    }
66}
67
68#[derive(Debug, Clone, Serialize, Deserialize)]
69pub struct FileStorageConfig {
70    pub path: String,
71}
72
73#[derive(Debug, Clone, Serialize, Deserialize)]
74pub struct SqliteStorageConfig {
75    pub path: String,
76
77    #[serde(default)]
78    pub table: Option<String>,
79}
80
81#[derive(Debug, Clone, Serialize, Deserialize)]
82pub struct RedisStorageConfig {
83    pub url: String,
84
85    #[serde(default)]
86    pub prefix: Option<String>,
87
88    #[serde(default)]
89    pub ttl_seconds: Option<u64>,
90}
91
92impl StorageConfig {
93    pub fn none() -> Self {
94        StorageConfig::None
95    }
96
97    pub fn file(path: impl Into<String>) -> Self {
98        StorageConfig::File(FileStorageConfig { path: path.into() })
99    }
100
101    pub fn sqlite(path: impl Into<String>) -> Self {
102        StorageConfig::Sqlite(SqliteStorageConfig {
103            path: path.into(),
104            table: None,
105        })
106    }
107
108    pub fn redis(url: impl Into<String>) -> Self {
109        StorageConfig::Redis(RedisStorageConfig {
110            url: url.into(),
111            prefix: None,
112            ttl_seconds: None,
113        })
114    }
115
116    pub fn is_none(&self) -> bool {
117        matches!(self, StorageConfig::None)
118    }
119
120    pub fn is_file(&self) -> bool {
121        matches!(self, StorageConfig::File(_))
122    }
123
124    pub fn is_sqlite(&self) -> bool {
125        matches!(self, StorageConfig::Sqlite(_))
126    }
127
128    pub fn is_redis(&self) -> bool {
129        matches!(self, StorageConfig::Redis(_))
130    }
131
132    pub fn storage_type(&self) -> &'static str {
133        match self {
134            StorageConfig::None => "none",
135            StorageConfig::File(_) => "file",
136            StorageConfig::Sqlite(_) => "sqlite",
137            StorageConfig::Redis(_) => "redis",
138        }
139    }
140
141    pub fn get_path(&self) -> Option<&str> {
142        match self {
143            StorageConfig::File(c) => Some(&c.path),
144            StorageConfig::Sqlite(c) => Some(&c.path),
145            _ => None,
146        }
147    }
148
149    pub fn get_url(&self) -> Option<&str> {
150        match self {
151            StorageConfig::Redis(c) => Some(&c.url),
152            _ => None,
153        }
154    }
155
156    pub fn get_prefix(&self) -> &str {
157        match self {
158            StorageConfig::Redis(c) => c.prefix.as_deref().unwrap_or("agent:"),
159            _ => "agent:",
160        }
161    }
162
163    pub fn get_ttl(&self) -> Option<u64> {
164        match self {
165            StorageConfig::Redis(c) => c.ttl_seconds,
166            _ => None,
167        }
168    }
169
170    pub fn get_table(&self) -> Option<&str> {
171        match self {
172            StorageConfig::Sqlite(c) => c.table.as_deref(),
173            _ => None,
174        }
175    }
176
177    pub fn as_file(&self) -> Option<&FileStorageConfig> {
178        match self {
179            StorageConfig::File(c) => Some(c),
180            _ => None,
181        }
182    }
183
184    pub fn as_sqlite(&self) -> Option<&SqliteStorageConfig> {
185        match self {
186            StorageConfig::Sqlite(c) => Some(c),
187            _ => None,
188        }
189    }
190
191    pub fn as_redis(&self) -> Option<&RedisStorageConfig> {
192        match self {
193            StorageConfig::Redis(c) => Some(c),
194            _ => None,
195        }
196    }
197}
198
199impl RedisStorageConfig {
200    pub fn with_prefix(mut self, prefix: impl Into<String>) -> Self {
201        self.prefix = Some(prefix.into());
202        self
203    }
204
205    pub fn with_ttl(mut self, ttl_seconds: u64) -> Self {
206        self.ttl_seconds = Some(ttl_seconds);
207        self
208    }
209}
210
211impl SqliteStorageConfig {
212    pub fn with_table(mut self, table: impl Into<String>) -> Self {
213        self.table = Some(table.into());
214        self
215    }
216}
217
218use ai_agents_storage::StorageConfig as StorageStorageConfig;
219
220/// Convert spec StorageConfig to storage crate StorageConfig for backend instantiation.
221pub fn to_storage_config(config: &StorageConfig) -> StorageStorageConfig {
222    match config {
223        StorageConfig::None => StorageStorageConfig::None,
224        StorageConfig::File(fc) => StorageStorageConfig::File {
225            path: fc.path.clone(),
226        },
227        StorageConfig::Sqlite(sc) => StorageStorageConfig::Sqlite {
228            path: sc.path.clone(),
229        },
230        StorageConfig::Redis(rc) => StorageStorageConfig::Redis {
231            url: rc.url.clone(),
232            prefix: rc.prefix.clone(),
233            ttl_seconds: rc.ttl_seconds,
234        },
235    }
236}
237
238#[cfg(test)]
239mod tests {
240    use super::*;
241
242    #[test]
243    fn test_storage_config_default() {
244        let config = StorageConfig::default();
245        assert!(config.is_none());
246        assert!(!config.is_file());
247        assert!(!config.is_sqlite());
248        assert!(!config.is_redis());
249        assert_eq!(config.storage_type(), "none");
250    }
251
252    #[test]
253    fn test_storage_config_none_yaml() {
254        let yaml = "type: none\n";
255        let config: StorageConfig = serde_yaml::from_str(yaml).unwrap();
256        assert!(config.is_none());
257    }
258
259    fn assert_unknown_field(yaml: &str, field: &str) {
260        let error = serde_yaml::from_str::<StorageConfig>(yaml)
261            .unwrap_err()
262            .to_string();
263        assert!(
264            error.contains(&format!("unknown field `{field}`")),
265            "{error}"
266        );
267    }
268
269    #[test]
270    fn test_storage_config_rejects_unknown_none_field() {
271        assert_unknown_field("type: none\nunexpected: true\n", "unexpected");
272    }
273
274    #[test]
275    fn test_storage_config_rejects_unknown_file_field() {
276        assert_unknown_field("type: file\npath: ./data\nread_only: true\n", "read_only");
277    }
278
279    #[test]
280    fn test_storage_config_rejects_unknown_sqlite_field() {
281        assert_unknown_field("type: sqlite\npath: ./data.db\ntabl: sessions\n", "tabl");
282    }
283
284    #[test]
285    fn test_storage_config_rejects_unknown_redis_field() {
286        assert_unknown_field(
287            "type: redis\nurl: redis://localhost:6379\nttl_second: 60\n",
288            "ttl_second",
289        );
290    }
291
292    #[test]
293    fn test_storage_config_file() {
294        let yaml = r#"
295type: file
296path: "./data/sessions"
297"#;
298        let config: StorageConfig = serde_yaml::from_str(yaml).unwrap();
299        assert!(config.is_file());
300        assert_eq!(config.get_path(), Some("./data/sessions"));
301        assert_eq!(config.storage_type(), "file");
302    }
303
304    #[test]
305    fn test_storage_config_file_builder() {
306        let config = StorageConfig::file("./data/sessions");
307        assert!(config.is_file());
308        assert_eq!(config.get_path(), Some("./data/sessions"));
309    }
310
311    #[test]
312    fn test_storage_config_sqlite() {
313        let yaml = r#"
314type: sqlite
315path: "./data/sessions.db"
316"#;
317        let config: StorageConfig = serde_yaml::from_str(yaml).unwrap();
318        assert!(config.is_sqlite());
319        assert_eq!(config.get_path(), Some("./data/sessions.db"));
320        assert_eq!(config.storage_type(), "sqlite");
321    }
322
323    #[test]
324    fn test_storage_config_sqlite_with_table() {
325        let yaml = r#"
326type: sqlite
327path: "./data/sessions.db"
328table: "custom_sessions"
329"#;
330        let config: StorageConfig = serde_yaml::from_str(yaml).unwrap();
331        assert!(config.is_sqlite());
332        assert_eq!(config.get_table(), Some("custom_sessions"));
333    }
334
335    #[test]
336    fn test_storage_config_sqlite_builder() {
337        let config = StorageConfig::sqlite("./data/sessions.db");
338        assert!(config.is_sqlite());
339        assert_eq!(config.get_path(), Some("./data/sessions.db"));
340    }
341
342    #[test]
343    fn test_storage_config_redis() {
344        let yaml = r#"
345type: redis
346url: "redis://localhost:6379"
347prefix: "myagent:"
348ttl_seconds: 86400
349"#;
350        let config: StorageConfig = serde_yaml::from_str(yaml).unwrap();
351        assert!(config.is_redis());
352        assert_eq!(config.get_url(), Some("redis://localhost:6379"));
353        assert_eq!(config.get_prefix(), "myagent:");
354        assert_eq!(config.get_ttl(), Some(86400));
355        assert_eq!(config.storage_type(), "redis");
356    }
357
358    #[test]
359    fn test_storage_config_redis_builder() {
360        let config = StorageConfig::redis("redis://localhost:6379");
361        assert!(config.is_redis());
362        assert_eq!(config.get_url(), Some("redis://localhost:6379"));
363        assert_eq!(config.get_prefix(), "agent:");
364        assert_eq!(config.get_ttl(), None);
365    }
366
367    #[test]
368    fn test_storage_config_default_prefix() {
369        let config = StorageConfig::default();
370        assert_eq!(config.get_prefix(), "agent:");
371
372        let config = StorageConfig::file("./data");
373        assert_eq!(config.get_prefix(), "agent:");
374
375        let yaml = r#"
376type: redis
377url: "redis://localhost:6379"
378"#;
379        let config: StorageConfig = serde_yaml::from_str(yaml).unwrap();
380        assert_eq!(config.get_prefix(), "agent:");
381    }
382
383    #[test]
384    fn test_storage_config_accessors() {
385        let file_config = StorageConfig::file("./data");
386        assert!(file_config.as_file().is_some());
387        assert!(file_config.as_sqlite().is_none());
388        assert!(file_config.as_redis().is_none());
389
390        let sqlite_config = StorageConfig::sqlite("./data.db");
391        assert!(sqlite_config.as_file().is_none());
392        assert!(sqlite_config.as_sqlite().is_some());
393        assert!(sqlite_config.as_redis().is_none());
394
395        let redis_config = StorageConfig::redis("redis://localhost");
396        assert!(redis_config.as_file().is_none());
397        assert!(redis_config.as_sqlite().is_none());
398        assert!(redis_config.as_redis().is_some());
399    }
400
401    #[test]
402    fn test_redis_config_builder_methods() {
403        let config = RedisStorageConfig {
404            url: "redis://localhost:6379".to_string(),
405            prefix: None,
406            ttl_seconds: None,
407        }
408        .with_prefix("test:")
409        .with_ttl(3600);
410
411        assert_eq!(config.prefix, Some("test:".to_string()));
412        assert_eq!(config.ttl_seconds, Some(3600));
413    }
414
415    #[test]
416    fn test_sqlite_config_builder_methods() {
417        let config = SqliteStorageConfig {
418            path: "./data.db".to_string(),
419            table: None,
420        }
421        .with_table("custom_table");
422
423        assert_eq!(config.table, Some("custom_table".to_string()));
424    }
425
426    #[test]
427    fn test_storage_config_serialization() {
428        let config = StorageConfig::redis("redis://localhost:6379");
429        let yaml = serde_yaml::to_string(&config).unwrap();
430        assert!(yaml.contains("type: redis"));
431        assert!(yaml.contains("url: redis://localhost:6379"));
432    }
433
434    #[test]
435    fn test_storage_config_valid_round_trips() {
436        let configs = [
437            StorageConfig::none(),
438            StorageConfig::file("./data/sessions"),
439            StorageConfig::sqlite("./data/sessions.db"),
440            StorageConfig::Sqlite(
441                SqliteStorageConfig {
442                    path: "./data/custom.db".to_string(),
443                    table: None,
444                }
445                .with_table("custom_sessions"),
446            ),
447            StorageConfig::Redis(
448                RedisStorageConfig {
449                    url: "redis://localhost:6379".to_string(),
450                    prefix: None,
451                    ttl_seconds: None,
452                }
453                .with_prefix("custom:")
454                .with_ttl(3600),
455            ),
456        ];
457
458        for config in configs {
459            let yaml = serde_yaml::to_string(&config).unwrap();
460            let restored: StorageConfig = serde_yaml::from_str(&yaml).unwrap();
461            assert_eq!(
462                serde_yaml::to_value(restored).unwrap(),
463                serde_yaml::to_value(config).unwrap()
464            );
465        }
466    }
467
468    #[test]
469    fn test_to_storage_config_none() {
470        use ai_agents_storage::StorageConfig as SC;
471        let result = to_storage_config(&StorageConfig::None);
472        assert!(matches!(result, SC::None));
473    }
474
475    #[test]
476    fn test_to_storage_config_file() {
477        use ai_agents_storage::StorageConfig as SC;
478        let config = StorageConfig::file("./data/sessions");
479        let result = to_storage_config(&config);
480        match result {
481            SC::File { path } => assert_eq!(path, "./data/sessions"),
482            other => panic!("expected File, got {:?}", other),
483        }
484    }
485
486    #[test]
487    fn test_to_storage_config_sqlite() {
488        use ai_agents_storage::StorageConfig as SC;
489        let config = StorageConfig::sqlite("./data/db.sqlite");
490        let result = to_storage_config(&config);
491        match result {
492            SC::Sqlite { path } => assert_eq!(path, "./data/db.sqlite"),
493            other => panic!("expected Sqlite, got {:?}", other),
494        }
495    }
496}
ai_agents_runtime/spec/storage.rs

ai_agents_runtime/spec/
storage.rs
