//! Python Virtual Environment Manager //! //! Manages isolated Python virtual environments for packs with Python dependencies. //! Each pack gets its own venv to prevent dependency conflicts. use super::dependency::{ DependencyError, DependencyManager, DependencyResult, DependencySpec, EnvironmentInfo, }; use async_trait::async_trait; use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::path::{Path, PathBuf}; use std::process::Stdio; use tokio::fs; use tokio::io::AsyncWriteExt; use tokio::process::Command; use tracing::{debug, info, warn}; /// Python virtual environment manager pub struct PythonVenvManager { /// Base directory for all virtual environments base_dir: PathBuf, /// Python interpreter to use for creating venvs python_path: PathBuf, /// Cache of environment info env_cache: tokio::sync::RwLock>, } /// Metadata stored with each environment #[derive(Debug, Clone, Serialize, Deserialize)] struct VenvMetadata { pack_ref: String, dependencies: Vec, created_at: chrono::DateTime, updated_at: chrono::DateTime, python_version: String, dependency_hash: String, } impl PythonVenvManager { /// Create a new Python venv manager pub fn new(base_dir: PathBuf) -> Self { Self { base_dir, python_path: PathBuf::from("python3"), env_cache: tokio::sync::RwLock::new(HashMap::new()), } } /// Create a new Python venv manager with custom Python path pub fn with_python_path(base_dir: PathBuf, python_path: PathBuf) -> Self { Self { base_dir, python_path, env_cache: tokio::sync::RwLock::new(HashMap::new()), } } /// Get the directory path for a pack's venv fn get_venv_path(&self, pack_ref: &str) -> PathBuf { // Sanitize pack_ref to create a valid directory name let safe_name = pack_ref.replace(['/', '\\', '.'], "_"); self.base_dir.join(safe_name) } /// Get the Python executable path within a venv fn get_venv_python(&self, venv_path: &Path) -> PathBuf { if cfg!(windows) { venv_path.join("Scripts").join("python.exe") } else { venv_path.join("bin").join("python") } } /// Get the pip executable path within a venv fn get_venv_pip(&self, venv_path: &Path) -> PathBuf { if cfg!(windows) { venv_path.join("Scripts").join("pip.exe") } else { venv_path.join("bin").join("pip") } } /// Get the metadata file path for a venv fn get_metadata_path(&self, venv_path: &Path) -> PathBuf { venv_path.join("attune_metadata.json") } /// Calculate a hash of dependencies for change detection fn calculate_dependency_hash(&self, spec: &DependencySpec) -> String { use std::collections::hash_map::DefaultHasher; use std::hash::{Hash, Hasher}; let mut hasher = DefaultHasher::new(); // Sort dependencies for consistent hashing let mut deps = spec.dependencies.clone(); deps.sort(); for dep in &deps { dep.hash(&mut hasher); } if let Some(ref content) = spec.requirements_file_content { content.hash(&mut hasher); } format!("{:x}", hasher.finish()) } /// Create a new virtual environment async fn create_venv(&self, venv_path: &Path) -> DependencyResult<()> { info!( "Creating Python virtual environment at: {}", venv_path.display() ); // Ensure base directory exists if let Some(parent) = venv_path.parent() { fs::create_dir_all(parent).await?; } // Create venv using python -m venv let output = Command::new(&self.python_path) .arg("-m") .arg("venv") .arg(venv_path) .arg("--clear") // Clear if exists .stdout(Stdio::piped()) .stderr(Stdio::piped()) .output() .await .map_err(|e| { DependencyError::CreateEnvironmentFailed(format!( "Failed to spawn venv command: {}", e )) })?; if !output.status.success() { let stderr = String::from_utf8_lossy(&output.stderr); return Err(DependencyError::CreateEnvironmentFailed(format!( "venv creation failed: {}", stderr ))); } // Upgrade pip to latest version let pip_path = self.get_venv_pip(venv_path); let output = Command::new(&pip_path) .arg("install") .arg("--upgrade") .arg("pip") .stdout(Stdio::piped()) .stderr(Stdio::piped()) .output() .await .map_err(|e| DependencyError::InstallFailed(format!("Failed to upgrade pip: {}", e)))?; if !output.status.success() { warn!("Failed to upgrade pip, continuing anyway"); } info!("Virtual environment created successfully"); Ok(()) } /// Install dependencies in a venv async fn install_dependencies( &self, venv_path: &Path, spec: &DependencySpec, ) -> DependencyResult<()> { if !spec.has_dependencies() { debug!("No dependencies to install"); return Ok(()); } info!("Installing dependencies in venv: {}", venv_path.display()); let pip_path = self.get_venv_pip(venv_path); // Install from requirements file content if provided if let Some(ref requirements_content) = spec.requirements_file_content { let req_file = venv_path.join("requirements.txt"); fs::write(&req_file, requirements_content).await?; let output = Command::new(&pip_path) .arg("install") .arg("-r") .arg(&req_file) .stdout(Stdio::piped()) .stderr(Stdio::piped()) .output() .await .map_err(|e| { DependencyError::InstallFailed(format!( "Failed to install from requirements.txt: {}", e )) })?; if !output.status.success() { let stderr = String::from_utf8_lossy(&output.stderr); return Err(DependencyError::InstallFailed(format!( "pip install failed: {}", stderr ))); } info!("Dependencies installed from requirements.txt"); } else if !spec.dependencies.is_empty() { // Install individual dependencies let output = Command::new(&pip_path) .arg("install") .args(&spec.dependencies) .stdout(Stdio::piped()) .stderr(Stdio::piped()) .output() .await .map_err(|e| { DependencyError::InstallFailed(format!("Failed to install dependencies: {}", e)) })?; if !output.status.success() { let stderr = String::from_utf8_lossy(&output.stderr); return Err(DependencyError::InstallFailed(format!( "pip install failed: {}", stderr ))); } info!("Installed {} dependencies", spec.dependencies.len()); } Ok(()) } /// Get Python version from a venv async fn get_python_version(&self, venv_path: &Path) -> DependencyResult { let python_path = self.get_venv_python(venv_path); let output = Command::new(&python_path) .arg("--version") .stdout(Stdio::piped()) .stderr(Stdio::piped()) .output() .await .map_err(|e| { DependencyError::ProcessError(format!("Failed to get Python version: {}", e)) })?; if !output.status.success() { return Err(DependencyError::ProcessError( "Failed to get Python version".to_string(), )); } let version = String::from_utf8_lossy(&output.stdout); Ok(version.trim().to_string()) } /// List installed packages in a venv async fn list_installed_packages(&self, venv_path: &Path) -> DependencyResult> { let pip_path = self.get_venv_pip(venv_path); let output = Command::new(&pip_path) .arg("list") .arg("--format=freeze") .stdout(Stdio::piped()) .stderr(Stdio::piped()) .output() .await .map_err(|e| { DependencyError::ProcessError(format!("Failed to list packages: {}", e)) })?; if !output.status.success() { return Ok(Vec::new()); } let packages = String::from_utf8_lossy(&output.stdout) .lines() .map(|s| s.to_string()) .collect(); Ok(packages) } /// Save metadata for a venv async fn save_metadata( &self, venv_path: &Path, metadata: &VenvMetadata, ) -> DependencyResult<()> { let metadata_path = self.get_metadata_path(venv_path); let json = serde_json::to_string_pretty(metadata).map_err(|e| { DependencyError::LockFileError(format!("Failed to serialize metadata: {}", e)) })?; let mut file = fs::File::create(&metadata_path).await.map_err(|e| { DependencyError::LockFileError(format!("Failed to create metadata file: {}", e)) })?; file.write_all(json.as_bytes()).await.map_err(|e| { DependencyError::LockFileError(format!("Failed to write metadata: {}", e)) })?; Ok(()) } /// Load metadata for a venv async fn load_metadata(&self, venv_path: &Path) -> DependencyResult> { let metadata_path = self.get_metadata_path(venv_path); if !metadata_path.exists() { return Ok(None); } let content = fs::read_to_string(&metadata_path).await.map_err(|e| { DependencyError::LockFileError(format!("Failed to read metadata: {}", e)) })?; let metadata: VenvMetadata = serde_json::from_str(&content).map_err(|e| { DependencyError::LockFileError(format!("Failed to parse metadata: {}", e)) })?; Ok(Some(metadata)) } /// Check if a venv exists and is valid async fn is_valid_venv(&self, venv_path: &Path) -> bool { if !venv_path.exists() { return false; } let python_path = self.get_venv_python(venv_path); if !python_path.exists() { return false; } // Try to run python --version to verify it works let result = Command::new(&python_path) .arg("--version") .stdout(Stdio::null()) .stderr(Stdio::null()) .status() .await; matches!(result, Ok(status) if status.success()) } /// Build environment info from a venv async fn build_env_info( &self, pack_ref: &str, venv_path: &Path, ) -> DependencyResult { let is_valid = self.is_valid_venv(venv_path).await; let python_path = self.get_venv_python(venv_path); let (python_version, installed_deps, created_at, updated_at) = if is_valid { let version = self .get_python_version(venv_path) .await .unwrap_or_else(|_| "Unknown".to_string()); let deps = self .list_installed_packages(venv_path) .await .unwrap_or_default(); let metadata = self.load_metadata(venv_path).await.ok().flatten(); let created = metadata .as_ref() .map(|m| m.created_at) .unwrap_or_else(chrono::Utc::now); let updated = metadata .as_ref() .map(|m| m.updated_at) .unwrap_or_else(chrono::Utc::now); (version, deps, created, updated) } else { ( "Unknown".to_string(), Vec::new(), chrono::Utc::now(), chrono::Utc::now(), ) }; Ok(EnvironmentInfo { id: pack_ref.to_string(), path: venv_path.to_path_buf(), runtime: "python".to_string(), runtime_version: python_version, installed_dependencies: installed_deps, created_at, updated_at, is_valid, executable_path: python_path, }) } } #[async_trait] impl DependencyManager for PythonVenvManager { fn runtime_type(&self) -> &str { "python" } async fn ensure_environment( &self, pack_ref: &str, spec: &DependencySpec, ) -> DependencyResult { info!("Ensuring Python environment for pack: {}", pack_ref); let venv_path = self.get_venv_path(pack_ref); let dependency_hash = self.calculate_dependency_hash(spec); // Check if environment exists and is up to date if venv_path.exists() { if let Some(metadata) = self.load_metadata(&venv_path).await? { if metadata.dependency_hash == dependency_hash && self.is_valid_venv(&venv_path).await { debug!("Using existing venv (dependencies unchanged)"); let env_info = self.build_env_info(pack_ref, &venv_path).await?; // Update cache let mut cache = self.env_cache.write().await; cache.insert(pack_ref.to_string(), env_info.clone()); return Ok(env_info); } info!("Dependencies changed or venv invalid, recreating environment"); } } // Create or recreate the venv self.create_venv(&venv_path).await?; // Install dependencies self.install_dependencies(&venv_path, spec).await?; // Get Python version let python_version = self.get_python_version(&venv_path).await?; // Save metadata let metadata = VenvMetadata { pack_ref: pack_ref.to_string(), dependencies: spec.dependencies.clone(), created_at: chrono::Utc::now(), updated_at: chrono::Utc::now(), python_version: python_version.clone(), dependency_hash, }; self.save_metadata(&venv_path, &metadata).await?; // Build environment info let env_info = self.build_env_info(pack_ref, &venv_path).await?; // Update cache let mut cache = self.env_cache.write().await; cache.insert(pack_ref.to_string(), env_info.clone()); info!("Python environment ready for pack: {}", pack_ref); Ok(env_info) } async fn get_environment(&self, pack_ref: &str) -> DependencyResult> { // Check cache first { let cache = self.env_cache.read().await; if let Some(env_info) = cache.get(pack_ref) { return Ok(Some(env_info.clone())); } } let venv_path = self.get_venv_path(pack_ref); if !venv_path.exists() { return Ok(None); } let env_info = self.build_env_info(pack_ref, &venv_path).await?; // Update cache let mut cache = self.env_cache.write().await; cache.insert(pack_ref.to_string(), env_info.clone()); Ok(Some(env_info)) } async fn remove_environment(&self, pack_ref: &str) -> DependencyResult<()> { info!("Removing Python environment for pack: {}", pack_ref); let venv_path = self.get_venv_path(pack_ref); if venv_path.exists() { fs::remove_dir_all(&venv_path).await?; } // Remove from cache let mut cache = self.env_cache.write().await; cache.remove(pack_ref); info!("Environment removed"); Ok(()) } async fn validate_environment(&self, pack_ref: &str) -> DependencyResult { let venv_path = self.get_venv_path(pack_ref); Ok(self.is_valid_venv(&venv_path).await) } async fn get_executable_path(&self, pack_ref: &str) -> DependencyResult { let venv_path = self.get_venv_path(pack_ref); let python_path = self.get_venv_python(&venv_path); if !python_path.exists() { return Err(DependencyError::EnvironmentNotFound(format!( "Python executable not found for pack: {}", pack_ref ))); } Ok(python_path) } async fn list_environments(&self) -> DependencyResult> { let mut environments = Vec::new(); let mut entries = fs::read_dir(&self.base_dir).await?; while let Some(entry) = entries.next_entry().await? { if entry.file_type().await?.is_dir() { let venv_path = entry.path(); if self.is_valid_venv(&venv_path).await { // Extract pack_ref from directory name if let Some(dir_name) = venv_path.file_name().and_then(|n| n.to_str()) { if let Ok(env_info) = self.build_env_info(dir_name, &venv_path).await { environments.push(env_info); } } } } } Ok(environments) } async fn cleanup(&self, keep_recent: usize) -> DependencyResult> { info!( "Cleaning up Python virtual environments (keeping {} most recent)", keep_recent ); let mut environments = self.list_environments().await?; // Sort by updated_at, newest first environments.sort_by(|a, b| b.updated_at.cmp(&a.updated_at)); let mut removed = Vec::new(); // Remove environments beyond keep_recent threshold for env in environments.iter().skip(keep_recent) { // Also skip if environment is invalid if !env.is_valid { if let Err(e) = self.remove_environment(&env.id).await { warn!("Failed to remove environment {}: {}", env.id, e); } else { removed.push(env.id.clone()); } } } info!("Cleaned up {} environments", removed.len()); Ok(removed) } async fn needs_update(&self, pack_ref: &str, spec: &DependencySpec) -> DependencyResult { let venv_path = self.get_venv_path(pack_ref); if !venv_path.exists() { return Ok(true); } if !self.is_valid_venv(&venv_path).await { return Ok(true); } // Check if dependency hash matches if let Some(metadata) = self.load_metadata(&venv_path).await? { let current_hash = self.calculate_dependency_hash(spec); Ok(metadata.dependency_hash != current_hash) } else { // No metadata, assume needs update Ok(true) } } } #[cfg(test)] mod tests { use super::*; use tempfile::TempDir; #[tokio::test] async fn test_venv_path_sanitization() { let temp_dir = TempDir::new().unwrap(); let manager = PythonVenvManager::new(temp_dir.path().to_path_buf()); let path = manager.get_venv_path("core.http"); assert!(path.to_string_lossy().contains("core_http")); let path = manager.get_venv_path("my/pack"); assert!(path.to_string_lossy().contains("my_pack")); } #[test] fn test_dependency_hash_consistency() { let temp_dir = TempDir::new().unwrap(); let manager = PythonVenvManager::new(temp_dir.path().to_path_buf()); let spec1 = DependencySpec::new("python") .with_dependency("requests==2.28.0") .with_dependency("flask==2.0.0"); let spec2 = DependencySpec::new("python") .with_dependency("flask==2.0.0") .with_dependency("requests==2.28.0"); // Hashes should be the same regardless of order (we sort) let hash1 = manager.calculate_dependency_hash(&spec1); let hash2 = manager.calculate_dependency_hash(&spec2); assert_eq!(hash1, hash2); } #[test] fn test_dependency_hash_different() { let temp_dir = TempDir::new().unwrap(); let manager = PythonVenvManager::new(temp_dir.path().to_path_buf()); let spec1 = DependencySpec::new("python").with_dependency("requests==2.28.0"); let spec2 = DependencySpec::new("python").with_dependency("requests==2.29.0"); let hash1 = manager.calculate_dependency_hash(&spec1); let hash2 = manager.calculate_dependency_hash(&spec2); assert_ne!(hash1, hash2); } }