Skip to main content

litellm_rs/core/models/user/
session.rs

1//! User session management
2
3use crate::core::models::Metadata;
4use serde::{Deserialize, Serialize};
5use std::collections::HashMap;
6use uuid::Uuid;
7
8/// User session
9#[derive(Debug, Clone, Serialize, Deserialize)]
10pub struct UserSession {
11    /// Session metadata
12    #[serde(flatten)]
13    pub metadata: Metadata,
14    /// User ID
15    pub user_id: Uuid,
16    /// Session token
17    #[serde(skip_serializing)]
18    pub token: String,
19    /// Session type
20    pub session_type: SessionType,
21    /// IP address
22    pub ip_address: Option<String>,
23    /// User agent
24    pub user_agent: Option<String>,
25    /// Expires at
26    pub expires_at: chrono::DateTime<chrono::Utc>,
27    /// Last activity
28    pub last_activity: chrono::DateTime<chrono::Utc>,
29    /// Session data
30    pub data: HashMap<String, serde_json::Value>,
31}
32
33/// Session type
34#[derive(Debug, Clone, Serialize, Deserialize)]
35#[serde(rename_all = "snake_case")]
36pub enum SessionType {
37    /// Web session
38    Web,
39    /// API session
40    Api,
41    /// Mobile session
42    Mobile,
43    /// CLI session
44    Cli,
45}
46
47impl UserSession {
48    /// Create a new session
49    pub fn new(
50        user_id: Uuid,
51        token: String,
52        session_type: SessionType,
53        expires_at: chrono::DateTime<chrono::Utc>,
54    ) -> Self {
55        Self {
56            metadata: Metadata::new(),
57            user_id,
58            token,
59            session_type,
60            ip_address: None,
61            user_agent: None,
62            expires_at,
63            last_activity: chrono::Utc::now(),
64            data: HashMap::new(),
65        }
66    }
67
68    /// Check if session is expired
69    pub fn is_expired(&self) -> bool {
70        chrono::Utc::now() > self.expires_at
71    }
72
73    /// Update last activity
74    pub fn update_activity(&mut self) {
75        self.last_activity = chrono::Utc::now();
76    }
77
78    /// Set session data
79    pub fn set_data<K: Into<String>, V: Into<serde_json::Value>>(&mut self, key: K, value: V) {
80        self.data.insert(key.into(), value.into());
81    }
82
83    /// Get session data
84    pub fn get_data(&self, key: &str) -> Option<&serde_json::Value> {
85        self.data.get(key)
86    }
87}
88
89#[cfg(test)]
90mod tests {
91    use super::*;
92    use chrono::{Duration, Utc};
93
94    // ==================== SessionType Tests ====================
95
96    #[test]
97    fn test_session_type_web() {
98        let session_type = SessionType::Web;
99        let json = serde_json::to_string(&session_type).unwrap();
100        assert_eq!(json, "\"web\"");
101    }
102
103    #[test]
104    fn test_session_type_api() {
105        let session_type = SessionType::Api;
106        let json = serde_json::to_string(&session_type).unwrap();
107        assert_eq!(json, "\"api\"");
108    }
109
110    #[test]
111    fn test_session_type_mobile() {
112        let session_type = SessionType::Mobile;
113        let json = serde_json::to_string(&session_type).unwrap();
114        assert_eq!(json, "\"mobile\"");
115    }
116
117    #[test]
118    fn test_session_type_cli() {
119        let session_type = SessionType::Cli;
120        let json = serde_json::to_string(&session_type).unwrap();
121        assert_eq!(json, "\"cli\"");
122    }
123
124    #[test]
125    fn test_session_type_deserialize() {
126        let web: SessionType = serde_json::from_str("\"web\"").unwrap();
127        assert!(matches!(web, SessionType::Web));
128
129        let api: SessionType = serde_json::from_str("\"api\"").unwrap();
130        assert!(matches!(api, SessionType::Api));
131    }
132
133    #[test]
134    fn test_session_type_clone() {
135        let original = SessionType::Mobile;
136        let cloned = original.clone();
137        let json1 = serde_json::to_string(&original).unwrap();
138        let json2 = serde_json::to_string(&cloned).unwrap();
139        assert_eq!(json1, json2);
140    }
141
142    #[test]
143    fn test_session_type_debug() {
144        let session_type = SessionType::Cli;
145        let debug_str = format!("{:?}", session_type);
146        assert!(debug_str.contains("Cli"));
147    }
148
149    // ==================== UserSession Creation Tests ====================
150
151    #[test]
152    fn test_user_session_new() {
153        let user_id = Uuid::new_v4();
154        let token = "test_token_123".to_string();
155        let expires_at = Utc::now() + Duration::hours(24);
156
157        let session = UserSession::new(user_id, token.clone(), SessionType::Web, expires_at);
158
159        assert_eq!(session.user_id, user_id);
160        assert_eq!(session.token, token);
161        assert!(matches!(session.session_type, SessionType::Web));
162        assert_eq!(session.expires_at, expires_at);
163        assert!(session.ip_address.is_none());
164        assert!(session.user_agent.is_none());
165        assert!(session.data.is_empty());
166    }
167
168    #[test]
169    fn test_user_session_new_with_api_type() {
170        let user_id = Uuid::new_v4();
171        let expires_at = Utc::now() + Duration::days(30);
172
173        let session = UserSession::new(
174            user_id,
175            "api_token".to_string(),
176            SessionType::Api,
177            expires_at,
178        );
179
180        assert!(matches!(session.session_type, SessionType::Api));
181    }
182
183    #[test]
184    fn test_user_session_new_with_mobile_type() {
185        let user_id = Uuid::new_v4();
186        let expires_at = Utc::now() + Duration::days(7);
187
188        let session = UserSession::new(
189            user_id,
190            "mobile_token".to_string(),
191            SessionType::Mobile,
192            expires_at,
193        );
194
195        assert!(matches!(session.session_type, SessionType::Mobile));
196    }
197
198    #[test]
199    fn test_user_session_new_with_cli_type() {
200        let user_id = Uuid::new_v4();
201        let expires_at = Utc::now() + Duration::hours(1);
202
203        let session = UserSession::new(
204            user_id,
205            "cli_token".to_string(),
206            SessionType::Cli,
207            expires_at,
208        );
209
210        assert!(matches!(session.session_type, SessionType::Cli));
211    }
212
213    // ==================== UserSession Expiration Tests ====================
214
215    #[test]
216    fn test_user_session_not_expired() {
217        let user_id = Uuid::new_v4();
218        let expires_at = Utc::now() + Duration::hours(24);
219
220        let session = UserSession::new(user_id, "token".to_string(), SessionType::Web, expires_at);
221
222        assert!(!session.is_expired());
223    }
224
225    #[test]
226    fn test_user_session_expired() {
227        let user_id = Uuid::new_v4();
228        let expires_at = Utc::now() - Duration::hours(1);
229
230        let session = UserSession::new(user_id, "token".to_string(), SessionType::Web, expires_at);
231
232        assert!(session.is_expired());
233    }
234
235    #[test]
236    fn test_user_session_just_expired() {
237        let user_id = Uuid::new_v4();
238        let expires_at = Utc::now() - Duration::seconds(1);
239
240        let session = UserSession::new(user_id, "token".to_string(), SessionType::Web, expires_at);
241
242        assert!(session.is_expired());
243    }
244
245    // ==================== UserSession Activity Tests ====================
246
247    #[test]
248    fn test_user_session_update_activity() {
249        let user_id = Uuid::new_v4();
250        let expires_at = Utc::now() + Duration::hours(24);
251
252        let mut session =
253            UserSession::new(user_id, "token".to_string(), SessionType::Web, expires_at);
254
255        let initial_activity = session.last_activity;
256        std::thread::sleep(std::time::Duration::from_millis(10));
257        session.update_activity();
258
259        assert!(session.last_activity >= initial_activity);
260    }
261
262    // ==================== UserSession Data Tests ====================
263
264    #[test]
265    fn test_user_session_set_data() {
266        let user_id = Uuid::new_v4();
267        let expires_at = Utc::now() + Duration::hours(24);
268
269        let mut session =
270            UserSession::new(user_id, "token".to_string(), SessionType::Web, expires_at);
271
272        session.set_data("key1", "value1");
273        session.set_data("key2", 42);
274        session.set_data("key3", true);
275
276        assert_eq!(session.data.len(), 3);
277    }
278
279    #[test]
280    fn test_user_session_get_data() {
281        let user_id = Uuid::new_v4();
282        let expires_at = Utc::now() + Duration::hours(24);
283
284        let mut session =
285            UserSession::new(user_id, "token".to_string(), SessionType::Web, expires_at);
286
287        session.set_data("test_key", "test_value");
288
289        let value = session.get_data("test_key");
290        assert!(value.is_some());
291        assert_eq!(value.unwrap(), "test_value");
292    }
293
294    #[test]
295    fn test_user_session_get_data_missing() {
296        let user_id = Uuid::new_v4();
297        let expires_at = Utc::now() + Duration::hours(24);
298
299        let session = UserSession::new(user_id, "token".to_string(), SessionType::Web, expires_at);
300
301        assert!(session.get_data("nonexistent").is_none());
302    }
303
304    #[test]
305    fn test_user_session_set_data_overwrite() {
306        let user_id = Uuid::new_v4();
307        let expires_at = Utc::now() + Duration::hours(24);
308
309        let mut session =
310            UserSession::new(user_id, "token".to_string(), SessionType::Web, expires_at);
311
312        session.set_data("key", "original");
313        session.set_data("key", "updated");
314
315        let value = session.get_data("key");
316        assert_eq!(value.unwrap(), "updated");
317    }
318
319    #[test]
320    fn test_user_session_set_data_various_types() {
321        let user_id = Uuid::new_v4();
322        let expires_at = Utc::now() + Duration::hours(24);
323
324        let mut session =
325            UserSession::new(user_id, "token".to_string(), SessionType::Web, expires_at);
326
327        session.set_data("string", "hello");
328        session.set_data("number", 123);
329        session.set_data("float", 1.234);
330        session.set_data("bool", true);
331        session.set_data("null", serde_json::Value::Null);
332
333        assert_eq!(session.data.len(), 5);
334    }
335
336    // ==================== UserSession Serialization Tests ====================
337
338    #[test]
339    fn test_user_session_serialize() {
340        let user_id = Uuid::new_v4();
341        let expires_at = Utc::now() + Duration::hours(24);
342
343        let session = UserSession::new(
344            user_id,
345            "secret_token".to_string(),
346            SessionType::Web,
347            expires_at,
348        );
349
350        let json = serde_json::to_string(&session).unwrap();
351
352        // Token should NOT be serialized (skip_serializing)
353        assert!(!json.contains("secret_token"));
354        // User ID should be serialized
355        assert!(json.contains(&user_id.to_string()));
356        // Session type should be serialized
357        assert!(json.contains("\"session_type\":\"web\""));
358    }
359
360    #[test]
361    fn test_user_session_serialize_with_optional_fields() {
362        let user_id = Uuid::new_v4();
363        let expires_at = Utc::now() + Duration::hours(24);
364
365        let mut session =
366            UserSession::new(user_id, "token".to_string(), SessionType::Api, expires_at);
367
368        session.ip_address = Some("192.168.1.1".to_string());
369        session.user_agent = Some("Mozilla/5.0".to_string());
370
371        let json = serde_json::to_string(&session).unwrap();
372
373        assert!(json.contains("192.168.1.1"));
374        assert!(json.contains("Mozilla/5.0"));
375    }
376
377    // ==================== UserSession Clone Tests ====================
378
379    #[test]
380    fn test_user_session_clone() {
381        let user_id = Uuid::new_v4();
382        let expires_at = Utc::now() + Duration::hours(24);
383
384        let mut session = UserSession::new(
385            user_id,
386            "token".to_string(),
387            SessionType::Mobile,
388            expires_at,
389        );
390
391        session.ip_address = Some("10.0.0.1".to_string());
392        session.set_data("key", "value");
393
394        let cloned = session.clone();
395
396        assert_eq!(session.user_id, cloned.user_id);
397        assert_eq!(session.token, cloned.token);
398        assert_eq!(session.ip_address, cloned.ip_address);
399        assert_eq!(session.data.len(), cloned.data.len());
400    }
401
402    // ==================== UserSession Debug Tests ====================
403
404    #[test]
405    fn test_user_session_debug() {
406        let user_id = Uuid::new_v4();
407        let expires_at = Utc::now() + Duration::hours(24);
408
409        let session = UserSession::new(user_id, "token".to_string(), SessionType::Cli, expires_at);
410
411        let debug_str = format!("{:?}", session);
412        assert!(debug_str.contains("UserSession"));
413        assert!(debug_str.contains("Cli"));
414    }
415
416    // ==================== UserSession Edge Cases ====================
417
418    #[test]
419    fn test_user_session_empty_token() {
420        let user_id = Uuid::new_v4();
421        let expires_at = Utc::now() + Duration::hours(24);
422
423        let session = UserSession::new(user_id, "".to_string(), SessionType::Web, expires_at);
424
425        assert!(session.token.is_empty());
426    }
427
428    #[test]
429    fn test_user_session_long_token() {
430        let user_id = Uuid::new_v4();
431        let expires_at = Utc::now() + Duration::hours(24);
432        let long_token = "a".repeat(1000);
433
434        let session = UserSession::new(user_id, long_token.clone(), SessionType::Web, expires_at);
435
436        assert_eq!(session.token.len(), 1000);
437    }
438
439    #[test]
440    fn test_user_session_far_future_expiry() {
441        let user_id = Uuid::new_v4();
442        let expires_at = Utc::now() + Duration::days(365 * 10);
443
444        let session = UserSession::new(user_id, "token".to_string(), SessionType::Api, expires_at);
445
446        assert!(!session.is_expired());
447    }
448}