1use crate::error::SpecError;
4use crate::models::{
5 ApprovalGate, ConversationMessage, MessageRole, SpecPhase, SpecWritingSession,
6};
7use chrono::Utc;
8use std::collections::HashMap;
9
10#[derive(Debug, Clone)]
12pub struct ConversationManager {
13 sessions: HashMap<String, SpecWritingSession>,
15}
16
17impl ConversationManager {
18 pub fn new() -> Self {
20 ConversationManager {
21 sessions: HashMap::new(),
22 }
23 }
24
25 pub fn create_session(
27 &mut self,
28 session_id: String,
29 spec_id: String,
30 ) -> Result<SpecWritingSession, SpecError> {
31 if self.sessions.contains_key(&session_id) {
32 return Err(SpecError::ConversationError(format!(
33 "Session {} already exists",
34 session_id
35 )));
36 }
37
38 let now = Utc::now();
39 let session = SpecWritingSession {
40 id: session_id.clone(),
41 spec_id,
42 phase: SpecPhase::Discovery,
43 conversation_history: vec![],
44 approval_gates: vec![],
45 created_at: now,
46 updated_at: now,
47 };
48
49 self.sessions.insert(session_id, session.clone());
50 Ok(session)
51 }
52
53 pub fn get_session(&self, session_id: &str) -> Result<SpecWritingSession, SpecError> {
55 self.sessions.get(session_id).cloned().ok_or_else(|| {
56 SpecError::ConversationError(format!("Session {} not found", session_id))
57 })
58 }
59
60 pub fn add_message(
62 &mut self,
63 session_id: &str,
64 message_id: String,
65 role: MessageRole,
66 content: String,
67 ) -> Result<ConversationMessage, SpecError> {
68 let session = self.sessions.get_mut(session_id).ok_or_else(|| {
69 SpecError::ConversationError(format!("Session {} not found", session_id))
70 })?;
71
72 let message = ConversationMessage {
73 id: message_id,
74 spec_id: session.spec_id.clone(),
75 role,
76 content,
77 timestamp: Utc::now(),
78 };
79
80 session.conversation_history.push(message.clone());
81 session.updated_at = Utc::now();
82
83 Ok(message)
84 }
85
86 pub fn get_conversation_history(
88 &self,
89 session_id: &str,
90 ) -> Result<Vec<ConversationMessage>, SpecError> {
91 let session = self.get_session(session_id)?;
92 Ok(session.conversation_history)
93 }
94
95 pub fn get_message(
97 &self,
98 session_id: &str,
99 message_id: &str,
100 ) -> Result<ConversationMessage, SpecError> {
101 let session = self.get_session(session_id)?;
102 session
103 .conversation_history
104 .iter()
105 .find(|m| m.id == message_id)
106 .cloned()
107 .ok_or_else(|| {
108 SpecError::ConversationError(format!("Message {} not found", message_id))
109 })
110 }
111
112 pub fn update_phase(
114 &mut self,
115 session_id: &str,
116 new_phase: SpecPhase,
117 ) -> Result<SpecWritingSession, SpecError> {
118 let session = self.sessions.get_mut(session_id).ok_or_else(|| {
119 SpecError::ConversationError(format!("Session {} not found", session_id))
120 })?;
121
122 session.phase = new_phase;
123 session.updated_at = Utc::now();
124
125 Ok(session.clone())
126 }
127
128 pub fn add_approval_gate(
130 &mut self,
131 session_id: &str,
132 gate: ApprovalGate,
133 ) -> Result<(), SpecError> {
134 let session = self.sessions.get_mut(session_id).ok_or_else(|| {
135 SpecError::ConversationError(format!("Session {} not found", session_id))
136 })?;
137
138 session.approval_gates.push(gate);
139 session.updated_at = Utc::now();
140
141 Ok(())
142 }
143
144 pub fn approve_phase(
146 &mut self,
147 session_id: &str,
148 phase: SpecPhase,
149 approved_by: Option<String>,
150 feedback: Option<String>,
151 ) -> Result<(), SpecError> {
152 let session = self.sessions.get_mut(session_id).ok_or_else(|| {
153 SpecError::ConversationError(format!("Session {} not found", session_id))
154 })?;
155
156 if let Some(gate) = session.approval_gates.iter_mut().find(|g| g.phase == phase) {
158 gate.approved = true;
159 gate.approved_at = Some(Utc::now());
160 gate.approved_by = approved_by;
161 gate.feedback = feedback;
162 } else {
163 let gate = ApprovalGate {
164 phase,
165 approved: true,
166 approved_at: Some(Utc::now()),
167 approved_by,
168 feedback,
169 };
170 session.approval_gates.push(gate);
171 }
172
173 session.updated_at = Utc::now();
174 Ok(())
175 }
176
177 pub fn get_approval_status(
179 &self,
180 session_id: &str,
181 phase: SpecPhase,
182 ) -> Result<bool, SpecError> {
183 let session = self.get_session(session_id)?;
184 Ok(session
185 .approval_gates
186 .iter()
187 .find(|g| g.phase == phase)
188 .map(|g| g.approved)
189 .unwrap_or(false))
190 }
191
192 pub fn delete_session(&mut self, session_id: &str) -> Result<(), SpecError> {
194 self.sessions.remove(session_id).ok_or_else(|| {
195 SpecError::ConversationError(format!("Session {} not found", session_id))
196 })?;
197 Ok(())
198 }
199
200 pub fn list_sessions(&self) -> Vec<String> {
202 self.sessions.keys().cloned().collect()
203 }
204
205 pub fn message_count(&self, session_id: &str) -> Result<usize, SpecError> {
207 let session = self.get_session(session_id)?;
208 Ok(session.conversation_history.len())
209 }
210
211 pub fn clear_history(&mut self, session_id: &str) -> Result<(), SpecError> {
213 let session = self.sessions.get_mut(session_id).ok_or_else(|| {
214 SpecError::ConversationError(format!("Session {} not found", session_id))
215 })?;
216
217 session.conversation_history.clear();
218 session.updated_at = Utc::now();
219
220 Ok(())
221 }
222}
223
224impl Default for ConversationManager {
225 fn default() -> Self {
226 Self::new()
227 }
228}
229
230#[cfg(test)]
231mod tests {
232 use super::*;
233
234 #[test]
235 fn test_create_session() {
236 let mut manager = ConversationManager::new();
237 let session = manager
238 .create_session("session-1".to_string(), "spec-1".to_string())
239 .unwrap();
240
241 assert_eq!(session.id, "session-1");
242 assert_eq!(session.spec_id, "spec-1");
243 assert_eq!(session.phase, SpecPhase::Discovery);
244 assert!(session.conversation_history.is_empty());
245 assert!(session.approval_gates.is_empty());
246 }
247
248 #[test]
249 fn test_create_duplicate_session_fails() {
250 let mut manager = ConversationManager::new();
251 manager
252 .create_session("session-1".to_string(), "spec-1".to_string())
253 .unwrap();
254
255 let result = manager.create_session("session-1".to_string(), "spec-2".to_string());
256 assert!(result.is_err());
257 }
258
259 #[test]
260 fn test_get_session() {
261 let mut manager = ConversationManager::new();
262 manager
263 .create_session("session-1".to_string(), "spec-1".to_string())
264 .unwrap();
265
266 let session = manager.get_session("session-1").unwrap();
267 assert_eq!(session.id, "session-1");
268 assert_eq!(session.spec_id, "spec-1");
269 }
270
271 #[test]
272 fn test_get_nonexistent_session_fails() {
273 let manager = ConversationManager::new();
274 let result = manager.get_session("nonexistent");
275 assert!(result.is_err());
276 }
277
278 #[test]
279 fn test_add_message() {
280 let mut manager = ConversationManager::new();
281 manager
282 .create_session("session-1".to_string(), "spec-1".to_string())
283 .unwrap();
284
285 let message = manager
286 .add_message(
287 "session-1",
288 "msg-1".to_string(),
289 MessageRole::User,
290 "Hello".to_string(),
291 )
292 .unwrap();
293
294 assert_eq!(message.id, "msg-1");
295 assert_eq!(message.role, MessageRole::User);
296 assert_eq!(message.content, "Hello");
297 }
298
299 #[test]
300 fn test_add_message_to_nonexistent_session_fails() {
301 let mut manager = ConversationManager::new();
302 let result = manager.add_message(
303 "nonexistent",
304 "msg-1".to_string(),
305 MessageRole::User,
306 "Hello".to_string(),
307 );
308 assert!(result.is_err());
309 }
310
311 #[test]
312 fn test_get_conversation_history() {
313 let mut manager = ConversationManager::new();
314 manager
315 .create_session("session-1".to_string(), "spec-1".to_string())
316 .unwrap();
317
318 manager
319 .add_message(
320 "session-1",
321 "msg-1".to_string(),
322 MessageRole::User,
323 "Hello".to_string(),
324 )
325 .unwrap();
326
327 manager
328 .add_message(
329 "session-1",
330 "msg-2".to_string(),
331 MessageRole::Assistant,
332 "Hi there".to_string(),
333 )
334 .unwrap();
335
336 let history = manager.get_conversation_history("session-1").unwrap();
337 assert_eq!(history.len(), 2);
338 assert_eq!(history[0].role, MessageRole::User);
339 assert_eq!(history[1].role, MessageRole::Assistant);
340 }
341
342 #[test]
343 fn test_get_message() {
344 let mut manager = ConversationManager::new();
345 manager
346 .create_session("session-1".to_string(), "spec-1".to_string())
347 .unwrap();
348
349 manager
350 .add_message(
351 "session-1",
352 "msg-1".to_string(),
353 MessageRole::User,
354 "Hello".to_string(),
355 )
356 .unwrap();
357
358 let message = manager.get_message("session-1", "msg-1").unwrap();
359 assert_eq!(message.id, "msg-1");
360 assert_eq!(message.content, "Hello");
361 }
362
363 #[test]
364 fn test_get_nonexistent_message_fails() {
365 let mut manager = ConversationManager::new();
366 manager
367 .create_session("session-1".to_string(), "spec-1".to_string())
368 .unwrap();
369
370 let result = manager.get_message("session-1", "nonexistent");
371 assert!(result.is_err());
372 }
373
374 #[test]
375 fn test_update_phase() {
376 let mut manager = ConversationManager::new();
377 manager
378 .create_session("session-1".to_string(), "spec-1".to_string())
379 .unwrap();
380
381 let session = manager
382 .update_phase("session-1", SpecPhase::Requirements)
383 .unwrap();
384
385 assert_eq!(session.phase, SpecPhase::Requirements);
386 }
387
388 #[test]
389 fn test_add_approval_gate() {
390 let mut manager = ConversationManager::new();
391 manager
392 .create_session("session-1".to_string(), "spec-1".to_string())
393 .unwrap();
394
395 let gate = ApprovalGate {
396 phase: SpecPhase::Requirements,
397 approved: false,
398 approved_at: None,
399 approved_by: None,
400 feedback: None,
401 };
402
403 manager.add_approval_gate("session-1", gate).unwrap();
404
405 let session = manager.get_session("session-1").unwrap();
406 assert_eq!(session.approval_gates.len(), 1);
407 }
408
409 #[test]
410 fn test_approve_phase() {
411 let mut manager = ConversationManager::new();
412 manager
413 .create_session("session-1".to_string(), "spec-1".to_string())
414 .unwrap();
415
416 manager
417 .approve_phase(
418 "session-1",
419 SpecPhase::Requirements,
420 Some("reviewer".to_string()),
421 Some("Looks good".to_string()),
422 )
423 .unwrap();
424
425 let approved = manager
426 .get_approval_status("session-1", SpecPhase::Requirements)
427 .unwrap();
428
429 assert!(approved);
430 }
431
432 #[test]
433 fn test_get_approval_status() {
434 let mut manager = ConversationManager::new();
435 manager
436 .create_session("session-1".to_string(), "spec-1".to_string())
437 .unwrap();
438
439 let approved = manager
440 .get_approval_status("session-1", SpecPhase::Requirements)
441 .unwrap();
442
443 assert!(!approved);
444
445 manager
446 .approve_phase("session-1", SpecPhase::Requirements, None, None)
447 .unwrap();
448
449 let approved = manager
450 .get_approval_status("session-1", SpecPhase::Requirements)
451 .unwrap();
452
453 assert!(approved);
454 }
455
456 #[test]
457 fn test_delete_session() {
458 let mut manager = ConversationManager::new();
459 manager
460 .create_session("session-1".to_string(), "spec-1".to_string())
461 .unwrap();
462
463 manager.delete_session("session-1").unwrap();
464
465 let result = manager.get_session("session-1");
466 assert!(result.is_err());
467 }
468
469 #[test]
470 fn test_list_sessions() {
471 let mut manager = ConversationManager::new();
472 manager
473 .create_session("session-1".to_string(), "spec-1".to_string())
474 .unwrap();
475 manager
476 .create_session("session-2".to_string(), "spec-2".to_string())
477 .unwrap();
478
479 let sessions = manager.list_sessions();
480 assert_eq!(sessions.len(), 2);
481 assert!(sessions.contains(&"session-1".to_string()));
482 assert!(sessions.contains(&"session-2".to_string()));
483 }
484
485 #[test]
486 fn test_message_count() {
487 let mut manager = ConversationManager::new();
488 manager
489 .create_session("session-1".to_string(), "spec-1".to_string())
490 .unwrap();
491
492 assert_eq!(manager.message_count("session-1").unwrap(), 0);
493
494 manager
495 .add_message(
496 "session-1",
497 "msg-1".to_string(),
498 MessageRole::User,
499 "Hello".to_string(),
500 )
501 .unwrap();
502
503 assert_eq!(manager.message_count("session-1").unwrap(), 1);
504 }
505
506 #[test]
507 fn test_clear_history() {
508 let mut manager = ConversationManager::new();
509 manager
510 .create_session("session-1".to_string(), "spec-1".to_string())
511 .unwrap();
512
513 manager
514 .add_message(
515 "session-1",
516 "msg-1".to_string(),
517 MessageRole::User,
518 "Hello".to_string(),
519 )
520 .unwrap();
521
522 assert_eq!(manager.message_count("session-1").unwrap(), 1);
523
524 manager.clear_history("session-1").unwrap();
525
526 assert_eq!(manager.message_count("session-1").unwrap(), 0);
527 }
528
529 #[test]
530 fn test_session_lifecycle() {
531 let mut manager = ConversationManager::new();
532
533 let session = manager
535 .create_session("session-1".to_string(), "spec-1".to_string())
536 .unwrap();
537 assert_eq!(session.phase, SpecPhase::Discovery);
538
539 manager
541 .add_message(
542 "session-1",
543 "msg-1".to_string(),
544 MessageRole::User,
545 "Create a task system".to_string(),
546 )
547 .unwrap();
548
549 manager
551 .update_phase("session-1", SpecPhase::Requirements)
552 .unwrap();
553
554 manager
556 .approve_phase("session-1", SpecPhase::Requirements, None, None)
557 .unwrap();
558
559 let session = manager.get_session("session-1").unwrap();
561 assert_eq!(session.phase, SpecPhase::Requirements);
562 assert_eq!(session.conversation_history.len(), 1);
563 assert_eq!(session.approval_gates.len(), 1);
564 assert!(session.approval_gates[0].approved);
565 }
566
567 #[test]
568 fn test_conversation_history_preservation() {
569 let mut manager = ConversationManager::new();
570 manager
571 .create_session("session-1".to_string(), "spec-1".to_string())
572 .unwrap();
573
574 let messages = vec![
576 ("msg-1", MessageRole::User, "Hello"),
577 ("msg-2", MessageRole::Assistant, "Hi"),
578 ("msg-3", MessageRole::User, "How are you?"),
579 ("msg-4", MessageRole::Assistant, "I'm good"),
580 ];
581
582 for (id, role, content) in messages {
583 manager
584 .add_message("session-1", id.to_string(), role, content.to_string())
585 .unwrap();
586 }
587
588 let history = manager.get_conversation_history("session-1").unwrap();
590 assert_eq!(history.len(), 4);
591
592 assert_eq!(history[0].id, "msg-1");
594 assert_eq!(history[1].id, "msg-2");
595 assert_eq!(history[2].id, "msg-3");
596 assert_eq!(history[3].id, "msg-4");
597
598 assert_eq!(history[0].role, MessageRole::User);
600 assert_eq!(history[1].role, MessageRole::Assistant);
601 assert_eq!(history[2].role, MessageRole::User);
602 assert_eq!(history[3].role, MessageRole::Assistant);
603 }
604}