1use serde::{Deserialize, Serialize};
2use std::collections::BTreeMap;
3
4pub type AttributeBag = BTreeMap<String, serde_json::Value>;
5
6#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
7pub struct RuntimeContext {
8 pub partition_key: String,
9 #[serde(default)]
10 pub attributes: AttributeBag,
11}
12
13impl RuntimeContext {
14 pub fn new(partition_key: impl Into<String>) -> Self {
15 Self {
16 partition_key: partition_key.into(),
17 attributes: AttributeBag::new(),
18 }
19 }
20
21 pub fn legacy_default() -> Self {
22 Self::new("default")
23 }
24
25 pub fn with_attribute(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
26 self.attributes.insert(key.into(), value);
27 self
28 }
29
30 pub fn validate(&self) -> Result<(), String> {
31 let partition_key = self.partition_key.trim();
32 if partition_key.is_empty() {
33 return Err("partition_key is required".to_string());
34 }
35 if partition_key.len() > 256 {
36 return Err("partition_key is too long".to_string());
37 }
38 for (key, value) in &self.attributes {
39 validate_attribute_key(key)?;
40 validate_attribute_value(value)?;
41 }
42 Ok(())
43 }
44}
45
46pub fn empty_attribute_bag() -> AttributeBag {
47 AttributeBag::new()
48}
49
50pub fn is_attribute_bag_empty(value: &AttributeBag) -> bool {
51 value.is_empty()
52}
53
54fn validate_attribute_key(key: &str) -> Result<(), String> {
55 if key.is_empty() || key.len() > 64 {
56 return Err(format!("invalid attribute key: {key}"));
57 }
58 if !key
59 .chars()
60 .all(|ch| ch == '_' || ch.is_ascii_lowercase() || ch.is_ascii_digit())
61 {
62 return Err(format!("invalid attribute key: {key}"));
63 }
64 Ok(())
65}
66
67fn validate_attribute_value(value: &serde_json::Value) -> Result<(), String> {
68 match value {
69 serde_json::Value::Null | serde_json::Value::Bool(_) | serde_json::Value::Number(_) => {
70 Ok(())
71 }
72 serde_json::Value::String(value) => {
73 if value.len() > 2048 {
74 Err("attribute string value is too long".to_string())
75 } else {
76 Ok(())
77 }
78 }
79 serde_json::Value::Array(values) => {
80 if values.len() > 32 {
81 return Err("attribute array value is too long".to_string());
82 }
83 for value in values {
84 validate_attribute_value(value)?;
85 }
86 Ok(())
87 }
88 serde_json::Value::Object(values) => {
89 if values.len() > 64 {
90 return Err("attribute object value is too large".to_string());
91 }
92 for (key, value) in values {
93 validate_attribute_key(key)?;
94 validate_attribute_value(value)?;
95 }
96 Ok(())
97 }
98 }
99}
100
101#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
102#[serde(rename_all = "lowercase")]
103pub enum StorageType {
104 LocalFile,
105 Database,
106 Redis,
107 Mysql,
108 Mongodb,
109}
110
111#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
112pub enum DatabaseType {
113 Postgres,
114 Mysql,
115}
116
117#[derive(Debug, Clone, Serialize, Deserialize)]
118pub struct DatabaseConfig {
119 pub host: String,
120 pub port: u16,
121 pub database: String,
122 pub user: String,
123 pub password: String,
124 pub max: Option<u32>,
125 pub idle_timeout_millis: Option<u64>,
126 pub connection_timeout_millis: Option<u64>,
127}
128
129#[derive(Debug, Clone)]
130pub struct WriteEntryOptions {
131 pub storage_type: StorageType,
132 pub session_id: String,
133 pub role: String,
134 pub content: String,
135 pub tools: Option<serde_json::Value>,
136 pub token_consumption: Option<u32>,
137 pub create_at: Option<String>,
138 pub database_config: Option<DatabaseConfig>,
139 pub runtime_context: Option<RuntimeContext>,
140 pub metadata: AttributeBag,
141}
142
143#[derive(Debug, Clone)]
144pub struct WriteCompactedEntryOptions {
145 pub storage_type: StorageType,
146 pub session_id: String,
147 pub summary: String,
148 pub trigger_entry_id: String,
149 pub create_at: Option<String>,
150 pub database_config: Option<DatabaseConfig>,
151 pub runtime_context: Option<RuntimeContext>,
152 pub metadata: AttributeBag,
153}
154
155#[derive(Debug, Clone, Serialize, Deserialize)]
156pub struct Entry {
157 #[serde(skip_serializing_if = "Option::is_none")]
158 pub partition_key: Option<String>,
159 pub entry_id: String,
160 pub session_id: String,
161 pub content: String,
162 pub role: String,
163 #[serde(skip_serializing_if = "Option::is_none")]
164 pub token_consumption: Option<u32>,
165 pub status: i16,
166 #[serde(skip_serializing_if = "Option::is_none")]
167 pub tools: Option<serde_json::Value>,
168 pub create_at: String,
169 pub is_compaction: i16,
170 #[serde(default, skip_serializing_if = "is_attribute_bag_empty")]
171 pub attributes: AttributeBag,
172 #[serde(default, skip_serializing_if = "is_attribute_bag_empty")]
173 pub metadata: AttributeBag,
174}
175
176#[derive(Debug, Clone, Serialize, Deserialize)]
177pub struct Session {
178 #[serde(skip_serializing_if = "Option::is_none")]
179 pub partition_key: Option<String>,
180 pub session_id: String,
181 pub status: i16,
182 pub create_at: String,
183 #[serde(default, skip_serializing_if = "is_attribute_bag_empty")]
184 pub attributes: AttributeBag,
185 #[serde(default, skip_serializing_if = "is_attribute_bag_empty")]
186 pub metadata: AttributeBag,
187}
188
189#[derive(Debug, Clone, Serialize, Deserialize)]
190pub struct SessionData {
191 pub session: Option<Session>,
192 pub entries: Vec<Entry>,
193 pub compacted_entries: Vec<CompactedEntry>,
194}
195
196#[derive(Debug, Clone, Serialize, Deserialize)]
197pub struct CompactedEntry {
198 #[serde(skip_serializing_if = "Option::is_none")]
199 pub partition_key: Option<String>,
200 pub entry_id: String,
201 pub session_id: String,
202 pub trigger_entry_id: String,
203 pub summary: String,
204 pub create_at: String,
205 pub status: i16,
206 #[serde(default, skip_serializing_if = "is_attribute_bag_empty")]
207 pub attributes: AttributeBag,
208 #[serde(default, skip_serializing_if = "is_attribute_bag_empty")]
209 pub metadata: AttributeBag,
210}
211
212#[derive(Debug, Clone, Serialize, Deserialize)]
213pub struct StorageFile {
214 pub sessions: Vec<StorageSession>,
215}
216
217#[derive(Debug, Clone, Serialize, Deserialize)]
218pub struct StorageSession {
219 #[serde(skip_serializing_if = "Option::is_none")]
220 pub partition_key: Option<String>,
221 pub session_id: String,
222 pub status: i16,
223 pub create_at: String,
224 #[serde(default, skip_serializing_if = "is_attribute_bag_empty")]
225 pub attributes: AttributeBag,
226 #[serde(default, skip_serializing_if = "is_attribute_bag_empty")]
227 pub metadata: AttributeBag,
228 #[serde(default)]
229 pub entries: Vec<Entry>,
230 #[serde(default)]
231 pub compacted_entries: Vec<CompactedEntry>,
232}
233
234#[cfg(test)]
235mod tests {
236 use super::*;
237
238 #[test]
239 fn runtime_context_accepts_arbitrary_business_attributes() {
240 let context = RuntimeContext::new("tenant:t1:workspace:w9")
241 .with_attribute("tenant_id", serde_json::json!("t1"))
242 .with_attribute("workspace_id", serde_json::json!("w9"))
243 .with_attribute("project_id", serde_json::json!("p3"));
244
245 assert!(context.validate().is_ok());
246 assert_eq!(
247 context.attributes.get("project_id"),
248 Some(&serde_json::json!("p3"))
249 );
250 }
251
252 #[test]
253 fn runtime_context_rejects_invalid_attribute_keys() {
254 let context =
255 RuntimeContext::new("user:u1").with_attribute("User Id", serde_json::json!("u1"));
256
257 assert!(context.validate().is_err());
258 }
259}