ricecoder_images/
session_manager.rs

1//! Session manager for image integration with ricecoder-sessions.
2//!
3//! This module provides functionality to:
4//! - Store and retrieve images from session context
5//! - Manage image references in sessions
6//! - Support image sharing between sessions
7//! - Persist image metadata with sessions
8
9use crate::models::{ImageAnalysisResult, ImageMetadata};
10use crate::session_integration::{MessageImageMetadata, SessionImageContext};
11use crate::error::{ImageError, ImageResult};
12use std::collections::HashMap;
13
14/// Manages images within a session
15#[derive(Debug, Clone)]
16pub struct SessionImageManager {
17    /// Session ID
18    session_id: String,
19    /// Image context for the session
20    image_context: SessionImageContext,
21    /// Map of image hash to analysis results for quick lookup
22    analysis_cache: HashMap<String, ImageAnalysisResult>,
23}
24
25impl SessionImageManager {
26    /// Create a new session image manager
27    pub fn new(session_id: String) -> Self {
28        Self {
29            session_id,
30            image_context: SessionImageContext::new(),
31            analysis_cache: HashMap::new(),
32        }
33    }
34
35    /// Create a session image manager from persisted context
36    pub fn from_context(session_id: String, context: SessionImageContext) -> Self {
37        Self {
38            session_id,
39            image_context: context,
40            analysis_cache: HashMap::new(),
41        }
42    }
43
44    /// Get the session ID
45    pub fn session_id(&self) -> &str {
46        &self.session_id
47    }
48
49    /// Add an image to the session context
50    ///
51    /// # Arguments
52    ///
53    /// * `metadata` - The image metadata
54    /// * `analysis` - Optional analysis result
55    ///
56    /// # Returns
57    ///
58    /// The image hash
59    pub fn add_image(
60        &mut self,
61        metadata: ImageMetadata,
62        analysis: Option<ImageAnalysisResult>,
63    ) -> ImageResult<String> {
64        let hash = metadata.hash.clone();
65
66        // Add to context
67        self.image_context.add_image(hash.clone(), metadata);
68
69        // Cache analysis if provided
70        if let Some(analysis_result) = analysis {
71            self.analysis_cache.insert(hash.clone(), analysis_result);
72        }
73
74        Ok(hash)
75    }
76
77    /// Remove an image from the session context
78    ///
79    /// Note: The image is removed from current context but remains in history
80    pub fn remove_image(&mut self, hash: &str) -> ImageResult<()> {
81        self.image_context.remove_image(hash);
82        Ok(())
83    }
84
85    /// Get an image from the session
86    pub fn get_image(&self, hash: &str) -> ImageResult<Option<MessageImageMetadata>> {
87        if let Some(metadata) = self.image_context.get_image_metadata(hash) {
88            let analysis = self.analysis_cache.get(hash).cloned();
89            let was_cached = analysis.is_some();
90
91            Ok(Some(MessageImageMetadata::new(
92                hash.to_string(),
93                metadata.clone(),
94                analysis,
95                was_cached,
96            )))
97        } else {
98            Ok(None)
99        }
100    }
101
102    /// Get all current images in the session
103    pub fn get_current_images(&self) -> Vec<MessageImageMetadata> {
104        self.image_context
105            .get_current_images()
106            .into_iter()
107            .map(|metadata| {
108                let hash = metadata.hash.clone();
109                let analysis = self.analysis_cache.get(&hash).cloned();
110                let was_cached = analysis.is_some();
111
112                MessageImageMetadata::new(hash, metadata.clone(), analysis, was_cached)
113            })
114            .collect()
115    }
116
117    /// Get all images ever included in the session
118    pub fn get_all_images(&self) -> Vec<MessageImageMetadata> {
119        self.image_context
120            .get_all_images()
121            .into_iter()
122            .map(|metadata| {
123                let hash = metadata.hash.clone();
124                let analysis = self.analysis_cache.get(&hash).cloned();
125                let was_cached = analysis.is_some();
126
127                MessageImageMetadata::new(hash, metadata.clone(), analysis, was_cached)
128            })
129            .collect()
130    }
131
132    /// Get the number of current images
133    pub fn current_image_count(&self) -> usize {
134        self.image_context.current_image_count()
135    }
136
137    /// Get the total number of images ever included
138    pub fn total_image_count(&self) -> usize {
139        self.image_context.total_image_count()
140    }
141
142    /// Check if an image is in the current context
143    pub fn has_image(&self, hash: &str) -> bool {
144        self.image_context.has_image(hash)
145    }
146
147    /// Clear current images (but keep history)
148    pub fn clear_current(&mut self) -> ImageResult<()> {
149        self.image_context.clear_current();
150        Ok(())
151    }
152
153    /// Clear all images including history
154    pub fn clear_all(&mut self) -> ImageResult<()> {
155        self.image_context.clear_all();
156        self.analysis_cache.clear();
157        Ok(())
158    }
159
160    /// Get the image context for persistence
161    pub fn get_context_for_persistence(&self) -> SessionImageContext {
162        self.image_context.clone()
163    }
164
165    /// Restore image context from persistence
166    pub fn restore_from_persistence(&mut self, context: SessionImageContext) {
167        self.image_context = context;
168    }
169
170    /// Get image hashes in current context
171    pub fn get_current_image_hashes(&self) -> Vec<String> {
172        self.image_context.current_images.clone()
173    }
174
175    /// Get all image hashes ever included
176    pub fn get_all_image_hashes(&self) -> Vec<String> {
177        self.image_context.all_images.clone()
178    }
179
180    /// Update analysis for an image
181    pub fn update_analysis(
182        &mut self,
183        hash: &str,
184        analysis: ImageAnalysisResult,
185    ) -> ImageResult<()> {
186        if !self.image_context.has_image(hash) {
187            return Err(ImageError::InvalidFile(
188                format!("Image with hash {} not found in session", hash),
189            ));
190        }
191
192        self.analysis_cache.insert(hash.to_string(), analysis);
193        Ok(())
194    }
195
196    /// Get analysis for an image
197    pub fn get_analysis(&self, hash: &str) -> Option<&ImageAnalysisResult> {
198        self.analysis_cache.get(hash)
199    }
200
201    /// Get all analyses
202    pub fn get_all_analyses(&self) -> Vec<&ImageAnalysisResult> {
203        self.analysis_cache.values().collect()
204    }
205}
206
207/// Manages images across multiple sessions
208#[derive(Debug, Clone)]
209pub struct MultiSessionImageManager {
210    /// Map of session ID to session image manager
211    sessions: HashMap<String, SessionImageManager>,
212}
213
214impl MultiSessionImageManager {
215    /// Create a new multi-session image manager
216    pub fn new() -> Self {
217        Self {
218            sessions: HashMap::new(),
219        }
220    }
221
222    /// Create a new session
223    pub fn create_session(&mut self, session_id: String) -> ImageResult<()> {
224        if self.sessions.contains_key(&session_id) {
225            return Err(ImageError::InvalidFile(
226                format!("Session {} already exists", session_id),
227            ));
228        }
229
230        self.sessions
231            .insert(session_id.clone(), SessionImageManager::new(session_id));
232        Ok(())
233    }
234
235    /// Get a session manager
236    pub fn get_session(&self, session_id: &str) -> ImageResult<&SessionImageManager> {
237        self.sessions
238            .get(session_id)
239            .ok_or_else(|| ImageError::InvalidFile(format!("Session {} not found", session_id)))
240    }
241
242    /// Get a mutable session manager
243    pub fn get_session_mut(&mut self, session_id: &str) -> ImageResult<&mut SessionImageManager> {
244        self.sessions
245            .get_mut(session_id)
246            .ok_or_else(|| ImageError::InvalidFile(format!("Session {} not found", session_id)))
247    }
248
249    /// Remove a session
250    pub fn remove_session(&mut self, session_id: &str) -> ImageResult<()> {
251        self.sessions
252            .remove(session_id)
253            .ok_or_else(|| ImageError::InvalidFile(format!("Session {} not found", session_id)))?;
254        Ok(())
255    }
256
257    /// Get all session IDs
258    pub fn get_session_ids(&self) -> Vec<String> {
259        self.sessions.keys().cloned().collect()
260    }
261
262    /// Share an image between sessions
263    ///
264    /// Copies an image from one session to another
265    pub fn share_image(
266        &mut self,
267        from_session: &str,
268        to_session: &str,
269        image_hash: &str,
270    ) -> ImageResult<()> {
271        // Get the image from source session
272        let image = self
273            .get_session(from_session)?
274            .get_image(image_hash)?
275            .ok_or_else(|| {
276                ImageError::InvalidFile(format!("Image {} not found in session {}", image_hash, from_session))
277            })?;
278
279        // Add to target session
280        let target_manager = self.get_session_mut(to_session)?;
281        target_manager.add_image(image.metadata, image.analysis)?;
282
283        Ok(())
284    }
285
286    /// Get the number of sessions
287    pub fn session_count(&self) -> usize {
288        self.sessions.len()
289    }
290}
291
292impl Default for MultiSessionImageManager {
293    fn default() -> Self {
294        Self::new()
295    }
296}
297
298#[cfg(test)]
299mod tests {
300    use super::*;
301    use std::path::PathBuf;
302
303    fn create_test_image_metadata() -> ImageMetadata {
304        ImageMetadata::new(
305            PathBuf::from("/path/to/image.png"),
306            crate::formats::ImageFormat::Png,
307            1024 * 1024,
308            800,
309            600,
310            "test_hash_123".to_string(),
311        )
312    }
313
314    fn create_test_analysis() -> ImageAnalysisResult {
315        ImageAnalysisResult::new(
316            "test_hash_123".to_string(),
317            "This is a test image".to_string(),
318            "openai".to_string(),
319            100,
320        )
321    }
322
323    #[test]
324    fn test_session_image_manager_new() {
325        let manager = SessionImageManager::new("session1".to_string());
326        assert_eq!(manager.session_id(), "session1");
327        assert_eq!(manager.current_image_count(), 0);
328    }
329
330    #[test]
331    fn test_session_image_manager_add_image() {
332        let mut manager = SessionImageManager::new("session1".to_string());
333        let metadata = create_test_image_metadata();
334
335        let hash = manager.add_image(metadata, None).unwrap();
336        assert_eq!(hash, "test_hash_123");
337        assert_eq!(manager.current_image_count(), 1);
338    }
339
340    #[test]
341    fn test_session_image_manager_add_with_analysis() {
342        let mut manager = SessionImageManager::new("session1".to_string());
343        let metadata = create_test_image_metadata();
344        let analysis = create_test_analysis();
345
346        manager.add_image(metadata, Some(analysis)).unwrap();
347        assert_eq!(manager.current_image_count(), 1);
348
349        let image = manager.get_image("test_hash_123").unwrap().unwrap();
350        assert!(image.analysis.is_some());
351    }
352
353    #[test]
354    fn test_session_image_manager_remove_image() {
355        let mut manager = SessionImageManager::new("session1".to_string());
356        let metadata = create_test_image_metadata();
357
358        manager.add_image(metadata, None).unwrap();
359        assert_eq!(manager.current_image_count(), 1);
360
361        manager.remove_image("test_hash_123").unwrap();
362        assert_eq!(manager.current_image_count(), 0);
363        assert_eq!(manager.total_image_count(), 1); // Still in history
364    }
365
366    #[test]
367    fn test_session_image_manager_get_image() {
368        let mut manager = SessionImageManager::new("session1".to_string());
369        let metadata = create_test_image_metadata();
370
371        manager.add_image(metadata, None).unwrap();
372
373        let image = manager.get_image("test_hash_123").unwrap();
374        assert!(image.is_some());
375        assert_eq!(image.unwrap().hash, "test_hash_123");
376    }
377
378    #[test]
379    fn test_session_image_manager_clear_current() {
380        let mut manager = SessionImageManager::new("session1".to_string());
381        let metadata = create_test_image_metadata();
382
383        manager.add_image(metadata, None).unwrap();
384        manager.clear_current().unwrap();
385
386        assert_eq!(manager.current_image_count(), 0);
387        assert_eq!(manager.total_image_count(), 1);
388    }
389
390    #[test]
391    fn test_session_image_manager_clear_all() {
392        let mut manager = SessionImageManager::new("session1".to_string());
393        let metadata = create_test_image_metadata();
394
395        manager.add_image(metadata, None).unwrap();
396        manager.clear_all().unwrap();
397
398        assert_eq!(manager.current_image_count(), 0);
399        assert_eq!(manager.total_image_count(), 0);
400    }
401
402    #[test]
403    fn test_multi_session_image_manager_create_session() {
404        let mut manager = MultiSessionImageManager::new();
405        assert!(manager.create_session("session1".to_string()).is_ok());
406        assert_eq!(manager.session_count(), 1);
407    }
408
409    #[test]
410    fn test_multi_session_image_manager_duplicate_session() {
411        let mut manager = MultiSessionImageManager::new();
412        assert!(manager.create_session("session1".to_string()).is_ok());
413        assert!(manager.create_session("session1".to_string()).is_err());
414    }
415
416    #[test]
417    fn test_multi_session_image_manager_get_session() {
418        let mut manager = MultiSessionImageManager::new();
419        manager.create_session("session1".to_string()).unwrap();
420
421        let session = manager.get_session("session1");
422        assert!(session.is_ok());
423    }
424
425    #[test]
426    fn test_multi_session_image_manager_remove_session() {
427        let mut manager = MultiSessionImageManager::new();
428        manager.create_session("session1".to_string()).unwrap();
429        assert_eq!(manager.session_count(), 1);
430
431        manager.remove_session("session1").unwrap();
432        assert_eq!(manager.session_count(), 0);
433    }
434
435    #[test]
436    fn test_multi_session_image_manager_share_image() {
437        let mut manager = MultiSessionImageManager::new();
438        manager.create_session("session1".to_string()).unwrap();
439        manager.create_session("session2".to_string()).unwrap();
440
441        let metadata = create_test_image_metadata();
442        manager
443            .get_session_mut("session1")
444            .unwrap()
445            .add_image(metadata, None)
446            .unwrap();
447
448        manager
449            .share_image("session1", "session2", "test_hash_123")
450            .unwrap();
451
452        assert_eq!(
453            manager.get_session("session2").unwrap().current_image_count(),
454            1
455        );
456    }
457
458    #[test]
459    fn test_session_image_manager_update_analysis() {
460        let mut manager = SessionImageManager::new("session1".to_string());
461        let metadata = create_test_image_metadata();
462
463        manager.add_image(metadata, None).unwrap();
464
465        let analysis = create_test_analysis();
466        manager
467            .update_analysis("test_hash_123", analysis)
468            .unwrap();
469
470        let image = manager.get_image("test_hash_123").unwrap().unwrap();
471        assert!(image.analysis.is_some());
472    }
473
474    #[test]
475    fn test_session_image_manager_persistence() {
476        let mut manager1 = SessionImageManager::new("session1".to_string());
477        let metadata = create_test_image_metadata();
478
479        manager1.add_image(metadata, None).unwrap();
480
481        let context = manager1.get_context_for_persistence();
482
483        let manager2 = SessionImageManager::from_context("session1".to_string(), context);
484        assert_eq!(manager2.current_image_count(), 1);
485    }
486}
ricecoder_images/session_manager.rs

ricecoder_images/
session_manager.rs
