Skip to main content

otherone_storage/
lib.rs

1// 作用:存储模块统一入口
2// 关联:被 otherone-agent、otherone-context、用户调用
3// 预期结果:根据存储类型分发到不同后端实现
4
5pub mod database;
6pub mod error;
7pub mod localfile;
8pub mod redis;
9pub mod types;
10
11use error::StorageError;
12use types::{RuntimeContext, StorageType, WriteCompactedEntryOptions, WriteEntryOptions};
13
14fn required_runtime_context(
15    runtime_context: &Option<RuntimeContext>,
16) -> Result<&RuntimeContext, StorageError> {
17    let runtime_context = runtime_context.as_ref().ok_or_else(|| {
18        StorageError::ConfigError("runtime_context is required for shared storage".to_string())
19    })?;
20    runtime_context
21        .validate()
22        .map_err(StorageError::ConfigError)?;
23    Ok(runtime_context)
24}
25
26/// 写入 entry 数据的统一入口
27pub async fn write_entry(options: &WriteEntryOptions) -> Result<(), StorageError> {
28    match options.storage_type {
29        StorageType::LocalFile => {
30            localfile::writer::write_entry_to_file(
31                &options.session_id,
32                &options.role,
33                &options.content,
34                options.tools.as_ref(),
35                options.token_consumption,
36                options.create_at.as_deref(),
37                &options.metadata,
38            )?;
39            Ok(())
40        }
41        StorageType::Database | StorageType::Mysql => {
42            let db_config = options.database_config.as_ref().ok_or_else(|| {
43                StorageError::ConfigError("database_config is required".to_string())
44            })?;
45            let runtime_context = required_runtime_context(&options.runtime_context)?;
46            if matches!(options.storage_type, StorageType::Mysql) {
47                database::mysql::write_entry_mysql(
48                    db_config,
49                    runtime_context,
50                    &options.session_id,
51                    &options.role,
52                    &options.content,
53                    options.tools.as_ref(),
54                    options.token_consumption,
55                    &options.metadata,
56                )
57                .await?;
58            } else {
59                database::writer::write_entry_to_database(
60                    db_config,
61                    runtime_context,
62                    &options.session_id,
63                    &options.role,
64                    &options.content,
65                    options.tools.as_ref(),
66                    options.token_consumption,
67                    options.create_at.as_deref(),
68                    &options.metadata,
69                )
70                .await?;
71            }
72            Ok(())
73        }
74        StorageType::Mongodb => {
75            let db_config = options.database_config.as_ref().ok_or_else(|| {
76                StorageError::ConfigError("database_config is required".to_string())
77            })?;
78            let runtime_context = required_runtime_context(&options.runtime_context)?;
79            let mongo = database::mongodb::MongoClient::new(
80                &format!(
81                    "mongodb://{}:{}@{}:{}/{}",
82                    db_config.user,
83                    db_config.password,
84                    db_config.host,
85                    db_config.port,
86                    db_config.database
87                ),
88                &db_config.database,
89            )
90            .await?;
91            mongo
92                .write_entry(
93                    runtime_context,
94                    &options.session_id,
95                    &options.role,
96                    &options.content,
97                    options.tools.as_ref(),
98                    options.token_consumption,
99                )
100                .await?;
101            Ok(())
102        }
103        StorageType::Redis => {
104            let _ = options; // Redis is a cache layer, not primary storage
105            Ok(())
106        }
107    }
108}
109
110/// 写入压缩记录的统一入口
111pub async fn write_compacted_entry(
112    options: &WriteCompactedEntryOptions,
113) -> Result<(), StorageError> {
114    match options.storage_type {
115        StorageType::LocalFile => {
116            localfile::writer::write_compacted_entry_to_file(
117                &options.session_id,
118                &options.summary,
119                &options.trigger_entry_id,
120                options.create_at.as_deref(),
121                &options.metadata,
122            )?;
123            Ok(())
124        }
125        StorageType::Database | StorageType::Mysql | StorageType::Mongodb | StorageType::Redis => {
126            let db_config = options.database_config.as_ref().ok_or_else(|| {
127                StorageError::ConfigError("database_config is required".to_string())
128            })?;
129            let runtime_context = required_runtime_context(&options.runtime_context)?;
130            if matches!(options.storage_type, StorageType::Database) {
131                database::writer::write_compacted_entry_to_database(
132                    db_config,
133                    runtime_context,
134                    &options.session_id,
135                    &options.summary,
136                    &options.trigger_entry_id,
137                    options.create_at.as_deref(),
138                    &options.metadata,
139                )
140                .await?;
141            } else if matches!(options.storage_type, StorageType::Mysql) {
142                database::mysql::write_compacted_entry_mysql(
143                    db_config,
144                    runtime_context,
145                    &options.session_id,
146                    &options.summary,
147                    &options.trigger_entry_id,
148                    &options.metadata,
149                )
150                .await?;
151            }
152            // MySQL/MongoDB/Redis compacted entries: basic support via the same interface
153            Ok(())
154        }
155    }
156}
157
158#[cfg(test)]
159mod tests {
160    use super::*;
161    use crate::types::{DatabaseConfig, RuntimeContext};
162
163    fn test_database_config() -> DatabaseConfig {
164        DatabaseConfig {
165            host: "127.0.0.1".to_string(),
166            port: 5432,
167            database: "otherone".to_string(),
168            user: "otherone".to_string(),
169            password: "password".to_string(),
170            max: None,
171            idle_timeout_millis: None,
172            connection_timeout_millis: None,
173        }
174    }
175
176    #[tokio::test]
177    async fn database_write_requires_runtime_context() {
178        let result = write_entry(&WriteEntryOptions {
179            storage_type: StorageType::Database,
180            session_id: "session-1".to_string(),
181            role: "user".to_string(),
182            content: "hello".to_string(),
183            tools: None,
184            token_consumption: None,
185            create_at: None,
186            database_config: Some(test_database_config()),
187            runtime_context: None,
188            metadata: Default::default(),
189        })
190        .await;
191
192        assert!(
193            matches!(result, Err(StorageError::ConfigError(message)) if message.contains("runtime_context"))
194        );
195    }
196
197    #[tokio::test]
198    async fn database_write_validates_runtime_context_before_connecting() {
199        let result = write_entry(&WriteEntryOptions {
200            storage_type: StorageType::Database,
201            session_id: "session-1".to_string(),
202            role: "user".to_string(),
203            content: "hello".to_string(),
204            tools: None,
205            token_consumption: None,
206            create_at: None,
207            database_config: Some(test_database_config()),
208            runtime_context: Some(
209                RuntimeContext::new("tenant:t1")
210                    .with_attribute("Bad Key", serde_json::json!("value")),
211            ),
212            metadata: Default::default(),
213        })
214        .await;
215
216        assert!(
217            matches!(result, Err(StorageError::ConfigError(message)) if message.contains("invalid attribute key"))
218        );
219    }
220}