1pub 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
26pub 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; Ok(())
106 }
107 }
108}
109
110pub 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 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}