1use chrono::{DateTime, Utc};
5use serde::{Deserialize, Serialize};
6
7use super::Principal;
8
9pub fn generate_session_id() -> String {
10 let random_bytes: [u8; 10] = rand::random();
11 format!(
12 "sess-{}",
13 base32::encode(base32::Alphabet::Rfc4648 { padding: false }, &random_bytes).to_lowercase()
14 )
15}
16
17pub fn generate_segment_id(session_id: &str, segment_number: u32) -> String {
18 format!("{}-seg-{}", session_id, segment_number)
19}
20
21#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
22pub struct Session {
23 pub id: String,
24 pub principal: Principal,
25 pub created_at: DateTime<Utc>,
26 pub ended_at: Option<DateTime<Utc>>,
27 pub segments: Vec<SessionSegment>,
28 pub current_segment_id: Option<String>,
29}
30
31impl Session {
32 pub fn new(
33 id: String,
34 principal: Principal,
35 provider: String,
36 model: String,
37 policy_id: Option<String>,
38 ) -> Self {
39 let segment_id = generate_segment_id(&id, 1);
40 let segment = SessionSegment {
41 id: segment_id.clone(),
42 provider,
43 model,
44 started_at: Utc::now(),
45 policy_id,
46 };
47 Self {
48 id,
49 principal,
50 created_at: Utc::now(),
51 ended_at: None,
52 segments: vec![segment],
53 current_segment_id: Some(segment_id),
54 }
55 }
56
57 pub fn is_active(&self) -> bool {
58 self.ended_at.is_none()
59 }
60
61 pub fn current_segment(&self) -> Option<&SessionSegment> {
62 self.current_segment_id
63 .as_ref()
64 .and_then(|id| self.segments.iter().find(|s| &s.id == id))
65 }
66
67 pub fn add_segment(
68 &mut self,
69 provider: String,
70 model: String,
71 policy_id: Option<String>,
72 ) -> &SessionSegment {
73 let segment_number = self.segments.len() as u32 + 1;
74 let segment_id = generate_segment_id(&self.id, segment_number);
75 let segment = SessionSegment {
76 id: segment_id.clone(),
77 provider,
78 model,
79 started_at: Utc::now(),
80 policy_id,
81 };
82 self.segments.push(segment);
83 self.current_segment_id = Some(segment_id);
84 self.segments.last().expect("segment was just pushed")
85 }
86
87 pub fn end(&mut self) {
88 self.ended_at = Some(Utc::now());
89 self.current_segment_id = None;
90 }
91}
92
93#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
94pub struct SessionSegment {
95 pub id: String,
96 pub provider: String,
97 pub model: String,
98 pub started_at: DateTime<Utc>,
99 pub policy_id: Option<String>,
100}
101
102#[cfg(test)]
103mod tests {
104 use super::*;
105
106 #[test]
107 fn test_session_creation() {
108 let principal = Principal::new("Test User", "test@example.com");
109 let session = Session::new(
110 "sess-test123".to_string(),
111 principal.clone(),
112 "anthropic".to_string(),
113 "claude-opus-4".to_string(),
114 None,
115 );
116
117 assert_eq!(session.id, "sess-test123");
118 assert_eq!(session.principal, principal);
119 assert!(session.is_active());
120 assert!(session.ended_at.is_none());
121 assert_eq!(session.segments.len(), 1);
122 assert!(session.current_segment_id.is_some());
123 }
124
125 #[test]
126 fn test_segment_id_format() {
127 let segment_id = generate_segment_id("sess-test123", 1);
128 assert_eq!(segment_id, "sess-test123-seg-1");
129
130 let segment_id = generate_segment_id("sess-test123", 2);
131 assert_eq!(segment_id, "sess-test123-seg-2");
132 }
133
134 #[test]
135 fn test_current_segment() {
136 let principal = Principal::new("Test User", "test@example.com");
137 let session = Session::new(
138 "sess-test123".to_string(),
139 principal,
140 "anthropic".to_string(),
141 "claude-opus-4".to_string(),
142 None,
143 );
144
145 let segment = session.current_segment().unwrap();
146 assert_eq!(segment.provider, "anthropic");
147 assert_eq!(segment.model, "claude-opus-4");
148 }
149
150 #[test]
151 fn test_add_segment() {
152 let principal = Principal::new("Test User", "test@example.com");
153 let mut session = Session::new(
154 "sess-test123".to_string(),
155 principal,
156 "anthropic".to_string(),
157 "claude-opus-4".to_string(),
158 None,
159 );
160
161 let segment = session.add_segment(
162 "openai".to_string(),
163 "gpt-4".to_string(),
164 Some("policy-123".to_string()),
165 );
166
167 let segment_id = segment.id.clone();
168 let segment_provider = segment.provider.clone();
169 let segment_model = segment.model.clone();
170 let segment_policy_id = segment.policy_id.clone();
171
172 assert_eq!(session.segments.len(), 2);
173 assert_eq!(segment_id, "sess-test123-seg-2");
174 assert_eq!(segment_provider, "openai");
175 assert_eq!(segment_model, "gpt-4");
176 assert_eq!(segment_policy_id, Some("policy-123".to_string()));
177
178 let current = session.current_segment().unwrap();
179 assert_eq!(current.id, "sess-test123-seg-2");
180 }
181
182 #[test]
183 fn test_end_session() {
184 let principal = Principal::new("Test User", "test@example.com");
185 let mut session = Session::new(
186 "sess-test123".to_string(),
187 principal,
188 "anthropic".to_string(),
189 "claude-opus-4".to_string(),
190 None,
191 );
192
193 assert!(session.is_active());
194
195 session.end();
196
197 assert!(!session.is_active());
198 assert!(session.ended_at.is_some());
199 assert!(session.current_segment_id.is_none());
200 }
201
202 #[test]
203 fn test_session_serialization() {
204 let principal = Principal::new("Test User", "test@example.com");
205 let session = Session::new(
206 "sess-test123".to_string(),
207 principal,
208 "anthropic".to_string(),
209 "claude-opus-4".to_string(),
210 Some("policy-abc".to_string()),
211 );
212
213 let json = serde_json::to_string(&session).unwrap();
214 let deserialized: Session = serde_json::from_str(&json).unwrap();
215
216 assert_eq!(session, deserialized);
217 }
218}