ricecoder_specs/
conversation.rs

1//! Conversation history storage and retrieval for spec writing sessions
2
3use crate::error::SpecError;
4use crate::models::{
5    ApprovalGate, ConversationMessage, MessageRole, SpecPhase, SpecWritingSession,
6};
7use chrono::Utc;
8use std::collections::HashMap;
9
10/// Manages conversation history and session lifecycle
11#[derive(Debug, Clone)]
12pub struct ConversationManager {
13    /// In-memory storage of sessions (session_id -> session)
14    sessions: HashMap<String, SpecWritingSession>,
15}
16
17impl ConversationManager {
18    /// Create a new conversation manager
19    pub fn new() -> Self {
20        ConversationManager {
21            sessions: HashMap::new(),
22        }
23    }
24
25    /// Create a new spec writing session
26    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    /// Retrieve a session by ID
54    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    /// Add a message to a session's conversation history
61    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    /// Get conversation history for a session
87    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    /// Get a specific message from a session
96    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    /// Update a session's phase
113    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    /// Add an approval gate to a session
129    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    /// Approve a phase in a session
145    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        // Find or create the approval gate for this phase
157        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    /// Get approval status for a phase
178    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    /// Delete a session
193    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    /// List all session IDs
201    pub fn list_sessions(&self) -> Vec<String> {
202        self.sessions.keys().cloned().collect()
203    }
204
205    /// Get the number of messages in a session
206    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    /// Clear conversation history for a session (but keep the session)
212    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        // Create session
534        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        // Add messages
540        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        // Update phase
550        manager
551            .update_phase("session-1", SpecPhase::Requirements)
552            .unwrap();
553
554        // Approve phase
555        manager
556            .approve_phase("session-1", SpecPhase::Requirements, None, None)
557            .unwrap();
558
559        // Verify final state
560        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        // Add multiple messages
575        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        // Retrieve and verify
589        let history = manager.get_conversation_history("session-1").unwrap();
590        assert_eq!(history.len(), 4);
591
592        // Verify order is preserved
593        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        // Verify roles are preserved
599        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}