use parking_lot::Mutex; use std::collections::HashMap; use std::fs; use std::path::{Path, PathBuf}; use std::sync::Arc; use uuid::Uuid; const TEMP_DIR_NAME: &str = "hikari-uploads"; pub struct TempFileManager { base_dir: PathBuf, files: HashMap>, } impl TempFileManager { pub fn new() -> Result { let base_dir = std::env::temp_dir().join(TEMP_DIR_NAME); if !base_dir.exists() { fs::create_dir_all(&base_dir) .map_err(|e| format!("Failed to create temp directory: {}", e))?; } Ok(TempFileManager { base_dir, files: HashMap::new(), }) } #[allow(dead_code)] pub fn get_base_dir(&self) -> &Path { &self.base_dir } pub fn save_file( &mut self, conversation_id: &str, data: &[u8], original_filename: Option<&str>, ) -> Result { let unique_id = Uuid::new_v4(); let extension = original_filename .and_then(|name| Path::new(name).extension()) .and_then(|ext| ext.to_str()) .unwrap_or("bin"); let filename = format!("{}_{}.{}", conversation_id, unique_id, extension); let file_path = self.base_dir.join(&filename); fs::write(&file_path, data) .map_err(|e| format!("Failed to write temp file: {}", e))?; self.files .entry(conversation_id.to_string()) .or_default() .push(file_path.clone()); Ok(file_path) } pub fn register_file(&mut self, conversation_id: &str, file_path: PathBuf) { self.files .entry(conversation_id.to_string()) .or_default() .push(file_path); } pub fn get_files_for_conversation(&self, conversation_id: &str) -> Vec { self.files .get(conversation_id) .cloned() .unwrap_or_default() } pub fn cleanup_conversation(&mut self, conversation_id: &str) -> Result<(), String> { if let Some(files) = self.files.remove(conversation_id) { for file_path in files { if file_path.exists() { if let Err(e) = fs::remove_file(&file_path) { eprintln!( "Warning: Failed to remove temp file {:?}: {}", file_path, e ); } } } } Ok(()) } pub fn cleanup_all(&mut self) -> Result<(), String> { let conversation_ids: Vec = self.files.keys().cloned().collect(); for conversation_id in conversation_ids { self.cleanup_conversation(&conversation_id)?; } Ok(()) } pub fn cleanup_orphaned_files(&mut self) -> Result { let mut cleaned_count = 0; if !self.base_dir.exists() { return Ok(0); } let tracked_files: std::collections::HashSet = self.files.values().flatten().cloned().collect(); let entries = fs::read_dir(&self.base_dir) .map_err(|e| format!("Failed to read temp directory: {}", e))?; for entry in entries.flatten() { let path = entry.path(); if path.is_file() && !tracked_files.contains(&path) { if let Err(e) = fs::remove_file(&path) { eprintln!("Warning: Failed to remove orphaned file {:?}: {}", path, e); } else { cleaned_count += 1; } } } Ok(cleaned_count) } } impl Default for TempFileManager { fn default() -> Self { Self::new().expect("Failed to create TempFileManager") } } pub type SharedTempFileManager = Arc>; pub fn create_shared_temp_manager() -> Result { Ok(Arc::new(Mutex::new(TempFileManager::new()?))) } #[cfg(test)] mod tests { use super::*; use std::fs; use tempfile::TempDir; // Helper to create a TempFileManager with a custom base directory for testing fn create_test_manager(base_dir: PathBuf) -> TempFileManager { if !base_dir.exists() { fs::create_dir_all(&base_dir).expect("Failed to create test temp dir"); } TempFileManager { base_dir, files: HashMap::new(), } } #[test] fn test_new_creates_base_directory() { let manager = TempFileManager::new().expect("Failed to create TempFileManager"); assert!(manager.base_dir.exists()); } #[test] fn test_get_base_dir_returns_correct_path() { let temp_dir = TempDir::new().expect("Failed to create temp dir"); let base_path = temp_dir.path().join("hikari-test"); let manager = create_test_manager(base_path.clone()); assert_eq!(manager.get_base_dir(), base_path.as_path()); } #[test] fn test_save_file_creates_file_with_content() { let temp_dir = TempDir::new().expect("Failed to create temp dir"); let base_path = temp_dir.path().join("hikari-test"); let mut manager = create_test_manager(base_path); let data = b"Hello, world!"; let result = manager.save_file("conv-1", data, Some("test.txt")); assert!(result.is_ok()); let file_path = result.unwrap(); assert!(file_path.exists()); let content = fs::read(&file_path).expect("Failed to read file"); assert_eq!(content, data); } #[test] fn test_save_file_uses_correct_extension() { let temp_dir = TempDir::new().expect("Failed to create temp dir"); let base_path = temp_dir.path().join("hikari-test"); let mut manager = create_test_manager(base_path); let data = b"test data"; let result = manager.save_file("conv-1", data, Some("document.pdf")); assert!(result.is_ok()); let file_path = result.unwrap(); assert_eq!(file_path.extension().unwrap(), "pdf"); } #[test] fn test_save_file_uses_bin_extension_when_no_filename() { let temp_dir = TempDir::new().expect("Failed to create temp dir"); let base_path = temp_dir.path().join("hikari-test"); let mut manager = create_test_manager(base_path); let data = b"binary data"; let result = manager.save_file("conv-1", data, None); assert!(result.is_ok()); let file_path = result.unwrap(); assert_eq!(file_path.extension().unwrap(), "bin"); } #[test] fn test_register_file_tracks_file_path() { let temp_dir = TempDir::new().expect("Failed to create temp dir"); let base_path = temp_dir.path().join("hikari-test"); let mut manager = create_test_manager(base_path); let file_path = PathBuf::from("/some/path/file.txt"); manager.register_file("conv-1", file_path.clone()); let files = manager.get_files_for_conversation("conv-1"); assert_eq!(files.len(), 1); assert_eq!(files[0], file_path); } #[test] fn test_get_files_for_conversation_returns_empty_for_unknown() { let temp_dir = TempDir::new().expect("Failed to create temp dir"); let base_path = temp_dir.path().join("hikari-test"); let manager = create_test_manager(base_path); let files = manager.get_files_for_conversation("unknown-conv"); assert!(files.is_empty()); } #[test] fn test_get_files_for_conversation_returns_all_files() { let temp_dir = TempDir::new().expect("Failed to create temp dir"); let base_path = temp_dir.path().join("hikari-test"); let mut manager = create_test_manager(base_path); let data = b"test"; manager.save_file("conv-1", data, Some("file1.txt")).unwrap(); manager.save_file("conv-1", data, Some("file2.txt")).unwrap(); manager.save_file("conv-2", data, Some("file3.txt")).unwrap(); let files_conv1 = manager.get_files_for_conversation("conv-1"); let files_conv2 = manager.get_files_for_conversation("conv-2"); assert_eq!(files_conv1.len(), 2); assert_eq!(files_conv2.len(), 1); } #[test] fn test_cleanup_conversation_removes_files() { let temp_dir = TempDir::new().expect("Failed to create temp dir"); let base_path = temp_dir.path().join("hikari-test"); let mut manager = create_test_manager(base_path); let data = b"test"; let file_path = manager.save_file("conv-1", data, Some("test.txt")).unwrap(); assert!(file_path.exists()); let result = manager.cleanup_conversation("conv-1"); assert!(result.is_ok()); assert!(!file_path.exists()); assert!(manager.get_files_for_conversation("conv-1").is_empty()); } #[test] fn test_cleanup_conversation_handles_missing_files() { let temp_dir = TempDir::new().expect("Failed to create temp dir"); let base_path = temp_dir.path().join("hikari-test"); let mut manager = create_test_manager(base_path); // Register a file that doesn't exist manager.register_file("conv-1", PathBuf::from("/nonexistent/file.txt")); // Should not error, just skip missing files let result = manager.cleanup_conversation("conv-1"); assert!(result.is_ok()); } #[test] fn test_cleanup_conversation_for_unknown_returns_ok() { let temp_dir = TempDir::new().expect("Failed to create temp dir"); let base_path = temp_dir.path().join("hikari-test"); let mut manager = create_test_manager(base_path); let result = manager.cleanup_conversation("unknown-conv"); assert!(result.is_ok()); } #[test] fn test_cleanup_all_removes_all_files() { let temp_dir = TempDir::new().expect("Failed to create temp dir"); let base_path = temp_dir.path().join("hikari-test"); let mut manager = create_test_manager(base_path); let data = b"test"; let file1 = manager.save_file("conv-1", data, Some("f1.txt")).unwrap(); let file2 = manager.save_file("conv-2", data, Some("f2.txt")).unwrap(); assert!(file1.exists()); assert!(file2.exists()); let result = manager.cleanup_all(); assert!(result.is_ok()); assert!(!file1.exists()); assert!(!file2.exists()); assert!(manager.files.is_empty()); } #[test] fn test_cleanup_orphaned_files_removes_untracked() { let temp_dir = TempDir::new().expect("Failed to create temp dir"); let base_path = temp_dir.path().join("hikari-test"); let mut manager = create_test_manager(base_path.clone()); // Create a tracked file let data = b"tracked"; let tracked_path = manager.save_file("conv-1", data, Some("tracked.txt")).unwrap(); // Create an untracked (orphaned) file directly in the temp directory let orphan_path = base_path.join("orphan.txt"); fs::write(&orphan_path, b"orphan").expect("Failed to create orphan file"); assert!(tracked_path.exists()); assert!(orphan_path.exists()); let result = manager.cleanup_orphaned_files(); assert!(result.is_ok()); assert_eq!(result.unwrap(), 1); // One orphan removed assert!(tracked_path.exists()); // Tracked file still exists assert!(!orphan_path.exists()); // Orphan removed } #[test] fn test_cleanup_orphaned_returns_zero_when_none() { let temp_dir = TempDir::new().expect("Failed to create temp dir"); let base_path = temp_dir.path().join("hikari-test"); let mut manager = create_test_manager(base_path); let data = b"test"; manager.save_file("conv-1", data, Some("test.txt")).unwrap(); let result = manager.cleanup_orphaned_files(); assert!(result.is_ok()); assert_eq!(result.unwrap(), 0); } #[test] fn test_cleanup_orphaned_returns_zero_when_dir_missing() { let mut manager = TempFileManager { base_dir: PathBuf::from("/nonexistent/dir"), files: HashMap::new(), }; let result = manager.cleanup_orphaned_files(); assert!(result.is_ok()); assert_eq!(result.unwrap(), 0); } #[test] fn test_default_creates_manager() { // Default should work as long as we can create temp directories let manager = TempFileManager::default(); assert!(manager.base_dir.exists()); } #[test] fn test_create_shared_temp_manager() { let result = create_shared_temp_manager(); assert!(result.is_ok()); let shared = result.unwrap(); let manager = shared.lock(); assert!(manager.base_dir.exists()); } #[test] fn test_multiple_files_same_conversation() { let temp_dir = TempDir::new().expect("Failed to create temp dir"); let base_path = temp_dir.path().join("hikari-test"); let mut manager = create_test_manager(base_path); // Save multiple files to same conversation for i in 0..5 { let data = format!("content {}", i); manager .save_file("conv-1", data.as_bytes(), Some(&format!("file{}.txt", i))) .unwrap(); } let files = manager.get_files_for_conversation("conv-1"); assert_eq!(files.len(), 5); // Each file should have unique content for (i, file_path) in files.iter().enumerate() { let content = fs::read_to_string(file_path).expect("Failed to read"); assert_eq!(content, format!("content {}", i)); } } #[test] fn test_file_paths_contain_conversation_id() { let temp_dir = TempDir::new().expect("Failed to create temp dir"); let base_path = temp_dir.path().join("hikari-test"); let mut manager = create_test_manager(base_path); let file_path = manager .save_file("my-conversation-id", b"test", Some("test.txt")) .unwrap(); let filename = file_path.file_name().unwrap().to_str().unwrap(); assert!(filename.starts_with("my-conversation-id_")); } }