cuenv_core/hooks/
state.rs

1//! State management for hook execution tracking
2
3use crate::hooks::types::{ExecutionStatus, HookResult};
4use crate::{Error, Result};
5use chrono::{DateTime, Utc};
6use fs4::tokio::AsyncFileExt;
7use serde::{Deserialize, Serialize};
8use std::collections::HashMap;
9use std::path::{Path, PathBuf};
10use tokio::fs;
11use tokio::fs::OpenOptions;
12use tokio::io::{AsyncReadExt, AsyncWriteExt};
13use tracing::{debug, error, info, warn};
14
15/// Manages persistent state for hook execution sessions
16#[derive(Debug, Clone)]
17pub struct StateManager {
18    state_dir: PathBuf,
19}
20
21impl StateManager {
22    /// Create a new state manager with the specified state directory
23    pub fn new(state_dir: PathBuf) -> Self {
24        Self { state_dir }
25    }
26
27    /// Get the default state directory (~/.cuenv/state)
28    pub fn default_state_dir() -> Result<PathBuf> {
29        // Check for CUENV_STATE_DIR environment variable first
30        if let Ok(state_dir) = std::env::var("CUENV_STATE_DIR") {
31            return Ok(PathBuf::from(state_dir));
32        }
33
34        let home = dirs::home_dir()
35            .ok_or_else(|| Error::configuration("Could not determine home directory"))?;
36        Ok(home.join(".cuenv").join("state"))
37    }
38
39    /// Create a state manager using the default state directory
40    pub fn with_default_dir() -> Result<Self> {
41        Ok(Self::new(Self::default_state_dir()?))
42    }
43
44    /// Get the state directory path
45    pub fn get_state_dir(&self) -> &Path {
46        &self.state_dir
47    }
48
49    /// Ensure the state directory exists
50    pub async fn ensure_state_dir(&self) -> Result<()> {
51        if !self.state_dir.exists() {
52            fs::create_dir_all(&self.state_dir)
53                .await
54                .map_err(|e| Error::Io {
55                    source: e,
56                    path: Some(self.state_dir.clone().into_boxed_path()),
57                    operation: "create_dir_all".to_string(),
58                })?;
59            debug!("Created state directory: {}", self.state_dir.display());
60        }
61        Ok(())
62    }
63
64    /// Generate a state file path for a given directory hash
65    fn state_file_path(&self, instance_hash: &str) -> PathBuf {
66        self.state_dir.join(format!("{}.json", instance_hash))
67    }
68
69    /// Get the state file path for a given directory hash (public for PID files)
70    pub fn get_state_file_path(&self, instance_hash: &str) -> PathBuf {
71        self.state_dir.join(format!("{}.json", instance_hash))
72    }
73
74    /// Save execution state to disk with atomic write and locking
75    pub async fn save_state(&self, state: &HookExecutionState) -> Result<()> {
76        self.ensure_state_dir().await?;
77
78        let state_file = self.state_file_path(&state.instance_hash);
79        let json = serde_json::to_string_pretty(state)
80            .map_err(|e| Error::configuration(format!("Failed to serialize state: {e}")))?;
81
82        // Write to a temporary file first, then rename atomically
83        let temp_path = state_file.with_extension("tmp");
84
85        // Open temp file with exclusive lock for writing
86        let mut file = OpenOptions::new()
87            .write(true)
88            .create(true)
89            .truncate(true)
90            .open(&temp_path)
91            .await
92            .map_err(|e| Error::Io {
93                source: e,
94                path: Some(temp_path.clone().into_boxed_path()),
95                operation: "open".to_string(),
96            })?;
97
98        // Acquire exclusive lock (only one writer allowed)
99        file.lock_exclusive().map_err(|e| {
100            Error::configuration(format!(
101                "Failed to acquire exclusive lock on state temp file: {}",
102                e
103            ))
104        })?;
105
106        file.write_all(json.as_bytes())
107            .await
108            .map_err(|e| Error::Io {
109                source: e,
110                path: Some(temp_path.clone().into_boxed_path()),
111                operation: "write_all".to_string(),
112            })?;
113
114        file.sync_all().await.map_err(|e| Error::Io {
115            source: e,
116            path: Some(temp_path.clone().into_boxed_path()),
117            operation: "sync_all".to_string(),
118        })?;
119
120        // Unlock happens automatically when file is dropped
121        drop(file);
122
123        // Atomically rename temp file to final location
124        fs::rename(&temp_path, &state_file)
125            .await
126            .map_err(|e| Error::Io {
127                source: e,
128                path: Some(state_file.clone().into_boxed_path()),
129                operation: "rename".to_string(),
130            })?;
131
132        debug!(
133            "Saved execution state for directory hash: {}",
134            state.instance_hash
135        );
136        Ok(())
137    }
138
139    /// Load execution state from disk with shared locking
140    pub async fn load_state(&self, instance_hash: &str) -> Result<Option<HookExecutionState>> {
141        let state_file = self.state_file_path(instance_hash);
142
143        if !state_file.exists() {
144            return Ok(None);
145        }
146
147        // Open file with shared lock for reading
148        let mut file = match OpenOptions::new().read(true).open(&state_file).await {
149            Ok(f) => f,
150            Err(e) => {
151                // File might have been deleted between exists check and open
152                if e.kind() == std::io::ErrorKind::NotFound {
153                    return Ok(None);
154                }
155                return Err(Error::Io {
156                    source: e,
157                    path: Some(state_file.clone().into_boxed_path()),
158                    operation: "open".to_string(),
159                });
160            }
161        };
162
163        // Acquire shared lock (multiple readers allowed)
164        file.lock_shared().map_err(|e| {
165            Error::configuration(format!(
166                "Failed to acquire shared lock on state file: {}",
167                e
168            ))
169        })?;
170
171        let mut contents = String::new();
172        file.read_to_string(&mut contents)
173            .await
174            .map_err(|e| Error::Io {
175                source: e,
176                path: Some(state_file.clone().into_boxed_path()),
177                operation: "read_to_string".to_string(),
178            })?;
179
180        // Unlock happens automatically when file is dropped
181        drop(file);
182
183        let state: HookExecutionState = serde_json::from_str(&contents)
184            .map_err(|e| Error::configuration(format!("Failed to deserialize state: {e}")))?;
185
186        debug!(
187            "Loaded execution state for directory hash: {}",
188            instance_hash
189        );
190        Ok(Some(state))
191    }
192
193    /// Remove state file for a directory
194    pub async fn remove_state(&self, instance_hash: &str) -> Result<()> {
195        let state_file = self.state_file_path(instance_hash);
196
197        if state_file.exists() {
198            fs::remove_file(&state_file).await.map_err(|e| Error::Io {
199                source: e,
200                path: Some(state_file.into_boxed_path()),
201                operation: "remove_file".to_string(),
202            })?;
203            debug!(
204                "Removed execution state for directory hash: {}",
205                instance_hash
206            );
207        }
208
209        Ok(())
210    }
211
212    /// List all active execution states
213    pub async fn list_active_states(&self) -> Result<Vec<HookExecutionState>> {
214        if !self.state_dir.exists() {
215            return Ok(Vec::new());
216        }
217
218        let mut states = Vec::new();
219        let mut dir = fs::read_dir(&self.state_dir).await.map_err(|e| Error::Io {
220            source: e,
221            path: Some(self.state_dir.clone().into_boxed_path()),
222            operation: "read_dir".to_string(),
223        })?;
224
225        while let Some(entry) = dir.next_entry().await.map_err(|e| Error::Io {
226            source: e,
227            path: Some(self.state_dir.clone().into_boxed_path()),
228            operation: "next_entry".to_string(),
229        })? {
230            let path = entry.path();
231            if path.extension().and_then(|s| s.to_str()) == Some("json")
232                && let Some(stem) = path.file_stem().and_then(|s| s.to_str())
233                && let Ok(Some(state)) = self.load_state(stem).await
234            {
235                states.push(state);
236            }
237        }
238
239        Ok(states)
240    }
241
242    /// Clean up the entire state directory
243    pub async fn cleanup_state_directory(&self) -> Result<usize> {
244        if !self.state_dir.exists() {
245            return Ok(0);
246        }
247
248        let mut cleaned_count = 0;
249        let mut dir = fs::read_dir(&self.state_dir).await.map_err(|e| Error::Io {
250            source: e,
251            path: Some(self.state_dir.clone().into_boxed_path()),
252            operation: "read_dir".to_string(),
253        })?;
254
255        while let Some(entry) = dir.next_entry().await.map_err(|e| Error::Io {
256            source: e,
257            path: Some(self.state_dir.clone().into_boxed_path()),
258            operation: "next_entry".to_string(),
259        })? {
260            let path = entry.path();
261
262            // Only clean up JSON state files
263            if path.extension().and_then(|s| s.to_str()) == Some("json") {
264                // Try to load and check if it's a completed state
265                if let Some(stem) = path.file_stem().and_then(|s| s.to_str()) {
266                    match self.load_state(stem).await {
267                        Ok(Some(state)) if state.is_complete() => {
268                            // Remove completed states
269                            if let Err(e) = fs::remove_file(&path).await {
270                                warn!("Failed to remove state file {}: {}", path.display(), e);
271                            } else {
272                                cleaned_count += 1;
273                                debug!("Cleaned up state file: {}", path.display());
274                            }
275                        }
276                        Ok(Some(_)) => {
277                            // Keep running states
278                            debug!("Keeping active state file: {}", path.display());
279                        }
280                        Ok(None) => {}
281                        Err(e) => {
282                            // If we can't parse it, it might be corrupted - remove it
283                            warn!("Failed to parse state file {}: {}", path.display(), e);
284                            if let Err(rm_err) = fs::remove_file(&path).await {
285                                error!(
286                                    "Failed to remove corrupted state file {}: {}",
287                                    path.display(),
288                                    rm_err
289                                );
290                            } else {
291                                cleaned_count += 1;
292                                info!("Removed corrupted state file: {}", path.display());
293                            }
294                        }
295                    }
296                }
297            }
298        }
299
300        if cleaned_count > 0 {
301            info!("Cleaned up {} state files from directory", cleaned_count);
302        }
303
304        Ok(cleaned_count)
305    }
306
307    /// Clean up orphaned state files (states without corresponding processes)
308    pub async fn cleanup_orphaned_states(&self, max_age: chrono::Duration) -> Result<usize> {
309        let cutoff = Utc::now() - max_age;
310        let mut cleaned_count = 0;
311
312        for state in self.list_active_states().await? {
313            // Remove states that are stuck in running but are too old
314            if state.status == ExecutionStatus::Running && state.started_at < cutoff {
315                warn!(
316                    "Found orphaned running state for {} (started {}), removing",
317                    state.directory_path.display(),
318                    state.started_at
319                );
320                self.remove_state(&state.instance_hash).await?;
321                cleaned_count += 1;
322            }
323        }
324
325        if cleaned_count > 0 {
326            info!("Cleaned up {} orphaned state files", cleaned_count);
327        }
328
329        Ok(cleaned_count)
330    }
331}
332
333/// Represents the state of hook execution for a specific directory
334#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
335pub struct HookExecutionState {
336    /// Hash combining directory path and config (instance identifier)
337    pub instance_hash: String,
338    /// Path to the directory being processed
339    pub directory_path: PathBuf,
340    /// Hash of the configuration that was approved
341    pub config_hash: String,
342    /// Current status of execution
343    pub status: ExecutionStatus,
344    /// Total number of hooks to execute
345    pub total_hooks: usize,
346    /// Number of hooks completed so far
347    pub completed_hooks: usize,
348    /// Index of currently executing hook (if any)
349    pub current_hook_index: Option<usize>,
350    /// Results of completed hooks
351    pub hook_results: HashMap<usize, HookResult>,
352    /// Timestamp when execution started
353    pub started_at: DateTime<Utc>,
354    /// Timestamp when execution finished (if completed)
355    pub finished_at: Option<DateTime<Utc>>,
356    /// Error message if execution failed
357    pub error_message: Option<String>,
358    /// Environment variables captured from source hooks
359    pub environment_vars: HashMap<String, String>,
360    /// Previous environment variables (for diff/unset support)
361    pub previous_env: Option<HashMap<String, String>>,
362}
363
364impl HookExecutionState {
365    /// Create a new execution state
366    pub fn new(
367        directory_path: PathBuf,
368        instance_hash: String,
369        config_hash: String,
370        total_hooks: usize,
371    ) -> Self {
372        Self {
373            instance_hash,
374            directory_path,
375            config_hash,
376            status: ExecutionStatus::Running,
377            total_hooks,
378            completed_hooks: 0,
379            current_hook_index: None,
380            hook_results: HashMap::new(),
381            started_at: Utc::now(),
382            finished_at: None,
383            error_message: None,
384            environment_vars: HashMap::new(),
385            previous_env: None,
386        }
387    }
388
389    /// Mark a hook as currently executing
390    pub fn mark_hook_running(&mut self, hook_index: usize) {
391        self.current_hook_index = Some(hook_index);
392        info!(
393            "Started executing hook {} of {}",
394            hook_index + 1,
395            self.total_hooks
396        );
397    }
398
399    /// Record the result of a hook execution
400    pub fn record_hook_result(&mut self, hook_index: usize, result: HookResult) {
401        self.hook_results.insert(hook_index, result.clone());
402        self.completed_hooks += 1;
403        self.current_hook_index = None;
404
405        if result.success {
406            info!(
407                "Hook {} of {} completed successfully",
408                hook_index + 1,
409                self.total_hooks
410            );
411        } else {
412            error!(
413                "Hook {} of {} failed: {:?}",
414                hook_index + 1,
415                self.total_hooks,
416                result.error
417            );
418            self.status = ExecutionStatus::Failed;
419            self.error_message = result.error.clone();
420            self.finished_at = Some(Utc::now());
421            return;
422        }
423
424        // Check if all hooks are complete
425        if self.completed_hooks == self.total_hooks {
426            self.status = ExecutionStatus::Completed;
427            self.finished_at = Some(Utc::now());
428            info!("All {} hooks completed successfully", self.total_hooks);
429        }
430    }
431
432    /// Mark execution as cancelled
433    pub fn mark_cancelled(&mut self, reason: Option<String>) {
434        self.status = ExecutionStatus::Cancelled;
435        self.finished_at = Some(Utc::now());
436        self.error_message = reason;
437        self.current_hook_index = None;
438    }
439
440    /// Check if execution is complete (success, failure, or cancelled)
441    pub fn is_complete(&self) -> bool {
442        matches!(
443            self.status,
444            ExecutionStatus::Completed | ExecutionStatus::Failed | ExecutionStatus::Cancelled
445        )
446    }
447
448    /// Get a human-readable progress display
449    pub fn progress_display(&self) -> String {
450        match &self.status {
451            ExecutionStatus::Running => {
452                if let Some(current) = self.current_hook_index {
453                    format!(
454                        "Executing hook {} of {} ({})",
455                        current + 1,
456                        self.total_hooks,
457                        self.status
458                    )
459                } else {
460                    format!(
461                        "{} of {} hooks completed",
462                        self.completed_hooks, self.total_hooks
463                    )
464                }
465            }
466            ExecutionStatus::Completed => "All hooks completed successfully".to_string(),
467            ExecutionStatus::Failed => {
468                if let Some(error) = &self.error_message {
469                    format!("Hook execution failed: {}", error)
470                } else {
471                    "Hook execution failed".to_string()
472                }
473            }
474            ExecutionStatus::Cancelled => {
475                if let Some(reason) = &self.error_message {
476                    format!("Hook execution cancelled: {}", reason)
477                } else {
478                    "Hook execution cancelled".to_string()
479                }
480            }
481        }
482    }
483
484    /// Get execution duration
485    pub fn duration(&self) -> chrono::Duration {
486        let end = self.finished_at.unwrap_or_else(Utc::now);
487        end - self.started_at
488    }
489}
490
491/// Compute a hash for a unique execution instance (directory + config)
492pub fn compute_instance_hash(path: &Path, config_hash: &str) -> String {
493    use sha2::{Digest, Sha256};
494    let mut hasher = Sha256::new();
495    hasher.update(path.to_string_lossy().as_bytes());
496    hasher.update(b":");
497    hasher.update(config_hash.as_bytes());
498    format!("{:x}", hasher.finalize())[..16].to_string()
499}
500
501#[cfg(test)]
502mod tests {
503    use super::*;
504    use crate::hooks::types::{Hook, HookResult};
505    use std::collections::HashMap;
506    use std::os::unix::process::ExitStatusExt;
507    use std::sync::Arc;
508    use std::time::Duration;
509    use tempfile::TempDir;
510
511    #[test]
512    fn test_compute_instance_hash() {
513        let path = Path::new("/test/path");
514        let config_hash = "test_config";
515        let hash = compute_instance_hash(path, config_hash);
516        assert_eq!(hash.len(), 16);
517
518        // Same path and config should produce same hash
519        let hash2 = compute_instance_hash(path, config_hash);
520        assert_eq!(hash, hash2);
521
522        // Different path should produce different hash
523        let different_path = Path::new("/other/path");
524        let different_hash = compute_instance_hash(different_path, config_hash);
525        assert_ne!(hash, different_hash);
526
527        // Same path but different config should produce different hash
528        let different_config_hash = compute_instance_hash(path, "different_config");
529        assert_ne!(hash, different_config_hash);
530    }
531
532    #[tokio::test]
533    async fn test_state_manager_operations() {
534        let temp_dir = TempDir::new().unwrap();
535        let state_manager = StateManager::new(temp_dir.path().to_path_buf());
536
537        let directory_path = PathBuf::from("/test/dir");
538        let config_hash = "test_config_hash".to_string();
539        let instance_hash = compute_instance_hash(&directory_path, &config_hash);
540
541        let mut state =
542            HookExecutionState::new(directory_path, instance_hash.clone(), config_hash, 2);
543
544        // Save initial state
545        state_manager.save_state(&state).await.unwrap();
546
547        // Load state back
548        let loaded_state = state_manager
549            .load_state(&instance_hash)
550            .await
551            .unwrap()
552            .unwrap();
553        assert_eq!(loaded_state.instance_hash, state.instance_hash);
554        assert_eq!(loaded_state.total_hooks, 2);
555        assert_eq!(loaded_state.status, ExecutionStatus::Running);
556
557        // Update state with hook result
558        let hook = Hook {
559            command: "echo".to_string(),
560            args: vec!["test".to_string()],
561            dir: None,
562            inputs: Vec::new(),
563            source: Some(false),
564        };
565
566        let result = HookResult::success(
567            hook,
568            std::process::ExitStatus::from_raw(0),
569            "test\n".to_string(),
570            "".to_string(),
571            100,
572        );
573
574        state.record_hook_result(0, result);
575        state_manager.save_state(&state).await.unwrap();
576
577        // Load updated state
578        let updated_state = state_manager
579            .load_state(&instance_hash)
580            .await
581            .unwrap()
582            .unwrap();
583        assert_eq!(updated_state.completed_hooks, 1);
584        assert_eq!(updated_state.hook_results.len(), 1);
585
586        // Remove state
587        state_manager.remove_state(&instance_hash).await.unwrap();
588        let removed_state = state_manager.load_state(&instance_hash).await.unwrap();
589        assert!(removed_state.is_none());
590    }
591
592    #[test]
593    fn test_hook_execution_state() {
594        let directory_path = PathBuf::from("/test/dir");
595        let instance_hash = "test_hash".to_string();
596        let config_hash = "config_hash".to_string();
597        let mut state = HookExecutionState::new(directory_path, instance_hash, config_hash, 3);
598
599        // Initial state
600        assert_eq!(state.status, ExecutionStatus::Running);
601        assert_eq!(state.total_hooks, 3);
602        assert_eq!(state.completed_hooks, 0);
603        assert!(!state.is_complete());
604
605        // Mark hook as running
606        state.mark_hook_running(0);
607        assert_eq!(state.current_hook_index, Some(0));
608
609        // Record successful hook result
610        let hook = Hook {
611            command: "echo".to_string(),
612            args: vec![],
613            dir: None,
614            inputs: Vec::new(),
615            source: Some(false),
616        };
617
618        let result = HookResult::success(
619            hook.clone(),
620            std::process::ExitStatus::from_raw(0),
621            "".to_string(),
622            "".to_string(),
623            100,
624        );
625
626        state.record_hook_result(0, result);
627        assert_eq!(state.completed_hooks, 1);
628        assert_eq!(state.current_hook_index, None);
629        assert_eq!(state.status, ExecutionStatus::Running);
630        assert!(!state.is_complete());
631
632        // Record failed hook result
633        let failed_result = HookResult::failure(
634            hook,
635            Some(std::process::ExitStatus::from_raw(256)),
636            "".to_string(),
637            "error".to_string(),
638            50,
639            "Command failed".to_string(),
640        );
641
642        state.record_hook_result(1, failed_result);
643        assert_eq!(state.completed_hooks, 2);
644        assert_eq!(state.status, ExecutionStatus::Failed);
645        assert!(state.is_complete());
646        assert!(state.error_message.is_some());
647
648        // Test cancellation
649        let mut cancelled_state = HookExecutionState::new(
650            PathBuf::from("/test"),
651            "hash".to_string(),
652            "config".to_string(),
653            1,
654        );
655        cancelled_state.mark_cancelled(Some("User cancelled".to_string()));
656        assert_eq!(cancelled_state.status, ExecutionStatus::Cancelled);
657        assert!(cancelled_state.is_complete());
658    }
659
660    #[test]
661    fn test_progress_display() {
662        let directory_path = PathBuf::from("/test/dir");
663        let instance_hash = "test_hash".to_string();
664        let config_hash = "config_hash".to_string();
665        let mut state = HookExecutionState::new(directory_path, instance_hash, config_hash, 2);
666
667        // Running state
668        let display = state.progress_display();
669        assert!(display.contains("0 of 2"));
670
671        // Running with current hook
672        state.mark_hook_running(0);
673        let display = state.progress_display();
674        assert!(display.contains("Executing hook 1 of 2"));
675
676        // Completed state
677        state.status = ExecutionStatus::Completed;
678        state.current_hook_index = None;
679        let display = state.progress_display();
680        assert_eq!(display, "All hooks completed successfully");
681
682        // Failed state
683        state.status = ExecutionStatus::Failed;
684        state.error_message = Some("Test error".to_string());
685        let display = state.progress_display();
686        assert!(display.contains("Hook execution failed: Test error"));
687    }
688
689    #[tokio::test]
690    async fn test_state_directory_cleanup() {
691        let temp_dir = TempDir::new().unwrap();
692        let state_manager = StateManager::new(temp_dir.path().to_path_buf());
693
694        // Create multiple states with different statuses
695        let completed_state = HookExecutionState {
696            instance_hash: "completed_hash".to_string(),
697            directory_path: PathBuf::from("/completed"),
698            config_hash: "config1".to_string(),
699            status: ExecutionStatus::Completed,
700            total_hooks: 1,
701            completed_hooks: 1,
702            current_hook_index: None,
703            hook_results: HashMap::new(),
704            environment_vars: HashMap::new(),
705            started_at: Utc::now() - chrono::Duration::hours(1),
706            finished_at: Some(Utc::now() - chrono::Duration::minutes(30)),
707            error_message: None,
708            previous_env: None,
709        };
710
711        let running_state = HookExecutionState {
712            instance_hash: "running_hash".to_string(),
713            directory_path: PathBuf::from("/running"),
714            config_hash: "config2".to_string(),
715            status: ExecutionStatus::Running,
716            total_hooks: 2,
717            completed_hooks: 1,
718            current_hook_index: Some(1),
719            hook_results: HashMap::new(),
720            environment_vars: HashMap::new(),
721            started_at: Utc::now() - chrono::Duration::minutes(5),
722            finished_at: None,
723            error_message: None,
724            previous_env: None,
725        };
726
727        let failed_state = HookExecutionState {
728            instance_hash: "failed_hash".to_string(),
729            directory_path: PathBuf::from("/failed"),
730            config_hash: "config3".to_string(),
731            status: ExecutionStatus::Failed,
732            total_hooks: 1,
733            completed_hooks: 0,
734            current_hook_index: None,
735            hook_results: HashMap::new(),
736            environment_vars: HashMap::new(),
737            started_at: Utc::now() - chrono::Duration::hours(2),
738            finished_at: Some(Utc::now() - chrono::Duration::hours(1)),
739            error_message: Some("Test failure".to_string()),
740            previous_env: None,
741        };
742
743        // Save all states
744        state_manager.save_state(&completed_state).await.unwrap();
745        state_manager.save_state(&running_state).await.unwrap();
746        state_manager.save_state(&failed_state).await.unwrap();
747
748        // Verify all states exist
749        let states = state_manager.list_active_states().await.unwrap();
750        assert_eq!(states.len(), 3);
751
752        // Clean up completed states
753        let cleaned = state_manager.cleanup_state_directory().await.unwrap();
754        assert_eq!(cleaned, 2); // Should clean up completed and failed states
755
756        // Verify only running state remains
757        let remaining_states = state_manager.list_active_states().await.unwrap();
758        assert_eq!(remaining_states.len(), 1);
759        assert_eq!(remaining_states[0].instance_hash, "running_hash");
760    }
761
762    #[tokio::test]
763    async fn test_cleanup_orphaned_states() {
764        let temp_dir = TempDir::new().unwrap();
765        let state_manager = StateManager::new(temp_dir.path().to_path_buf());
766
767        // Create an old running state (orphaned)
768        let orphaned_state = HookExecutionState {
769            instance_hash: "orphaned_hash".to_string(),
770            directory_path: PathBuf::from("/orphaned"),
771            config_hash: "config".to_string(),
772            status: ExecutionStatus::Running,
773            total_hooks: 1,
774            completed_hooks: 0,
775            current_hook_index: Some(0),
776            hook_results: HashMap::new(),
777            environment_vars: HashMap::new(),
778            started_at: Utc::now() - chrono::Duration::hours(3),
779            finished_at: None,
780            error_message: None,
781            previous_env: None,
782        };
783
784        // Create a recent running state (not orphaned)
785        let recent_state = HookExecutionState {
786            instance_hash: "recent_hash".to_string(),
787            directory_path: PathBuf::from("/recent"),
788            config_hash: "config".to_string(),
789            status: ExecutionStatus::Running,
790            total_hooks: 1,
791            completed_hooks: 0,
792            current_hook_index: Some(0),
793            hook_results: HashMap::new(),
794            environment_vars: HashMap::new(),
795            started_at: Utc::now() - chrono::Duration::minutes(5),
796            finished_at: None,
797            error_message: None,
798            previous_env: None,
799        };
800
801        // Save both states
802        state_manager.save_state(&orphaned_state).await.unwrap();
803        state_manager.save_state(&recent_state).await.unwrap();
804
805        // Clean up orphaned states older than 1 hour
806        let cleaned = state_manager
807            .cleanup_orphaned_states(chrono::Duration::hours(1))
808            .await
809            .unwrap();
810        assert_eq!(cleaned, 1); // Should clean up only the orphaned state
811
812        // Verify only recent state remains
813        let remaining_states = state_manager.list_active_states().await.unwrap();
814        assert_eq!(remaining_states.len(), 1);
815        assert_eq!(remaining_states[0].instance_hash, "recent_hash");
816    }
817
818    #[tokio::test]
819    async fn test_corrupted_state_file_handling() {
820        let temp_dir = TempDir::new().unwrap();
821        let state_dir = temp_dir.path().join("state");
822        let state_manager = StateManager::new(state_dir.clone());
823
824        // Ensure state directory exists
825        state_manager.ensure_state_dir().await.unwrap();
826
827        // Write corrupted JSON to a state file
828        let corrupted_file = state_dir.join("corrupted.json");
829        tokio::fs::write(&corrupted_file, "{invalid json}")
830            .await
831            .unwrap();
832
833        // List active states should handle the corrupted file gracefully
834        let states = state_manager.list_active_states().await.unwrap();
835        assert_eq!(states.len(), 0); // Corrupted file should be skipped
836
837        // Cleanup should remove the corrupted file
838        let cleaned = state_manager.cleanup_state_directory().await.unwrap();
839        assert_eq!(cleaned, 1);
840
841        // Verify the corrupted file is gone
842        assert!(!corrupted_file.exists());
843    }
844
845    #[tokio::test]
846    async fn test_concurrent_state_modifications() {
847        use tokio::task;
848
849        let temp_dir = TempDir::new().unwrap();
850        let state_manager = Arc::new(StateManager::new(temp_dir.path().to_path_buf()));
851
852        // Create initial state
853        let initial_state = HookExecutionState {
854            instance_hash: "concurrent_hash".to_string(),
855            directory_path: PathBuf::from("/concurrent"),
856            config_hash: "config".to_string(),
857            status: ExecutionStatus::Running,
858            total_hooks: 10,
859            completed_hooks: 0,
860            current_hook_index: Some(0),
861            hook_results: HashMap::new(),
862            environment_vars: HashMap::new(),
863            started_at: Utc::now(),
864            finished_at: None,
865            error_message: None,
866            previous_env: None,
867        };
868
869        state_manager.save_state(&initial_state).await.unwrap();
870
871        // Spawn multiple tasks that concurrently modify the state
872        let mut handles = vec![];
873
874        for i in 0..5 {
875            let sm = state_manager.clone();
876            let path = initial_state.directory_path.clone();
877
878            let handle = task::spawn(async move {
879                // Load state - it might have been modified by another task
880                let instance_hash = compute_instance_hash(&path, "concurrent_config");
881
882                // Simulate some work
883                tokio::time::sleep(Duration::from_millis(10)).await;
884
885                // Load state, modify, and save (handle potential concurrent modifications)
886                if let Ok(Some(mut state)) = sm.load_state(&instance_hash).await {
887                    state.completed_hooks += 1;
888                    state.current_hook_index = Some(i + 1);
889
890                    // Save state - ignore errors from concurrent saves
891                    let _ = sm.save_state(&state).await;
892                }
893            });
894
895            handles.push(handle);
896        }
897
898        // Wait for all tasks to complete
899        for handle in handles {
900            handle.await.unwrap();
901        }
902
903        // Verify final state - due to concurrent writes, the exact values may vary
904        // but the state should be loadable and valid
905        let final_state = state_manager
906            .load_state(&initial_state.instance_hash)
907            .await
908            .unwrap();
909
910        // The state might exist or not depending on timing of concurrent operations
911        if let Some(state) = final_state {
912            assert_eq!(state.instance_hash, "concurrent_hash");
913            // Completed hooks will be 0 if all concurrent writes failed, or > 0 if some succeeded
914        }
915    }
916
917    #[tokio::test]
918    async fn test_state_with_unicode_and_special_chars() {
919        let temp_dir = TempDir::new().unwrap();
920        let state_manager = StateManager::new(temp_dir.path().to_path_buf());
921
922        // Create state with unicode and special characters
923        let mut unicode_state = HookExecutionState {
924            instance_hash: "unicode_hash".to_string(),
925            directory_path: PathBuf::from("/测试/目录/🚀"),
926            config_hash: "config_ñ_é_ü".to_string(),
927            status: ExecutionStatus::Failed,
928            total_hooks: 1,
929            completed_hooks: 1,
930            current_hook_index: None,
931            hook_results: HashMap::new(),
932            environment_vars: HashMap::new(),
933            started_at: Utc::now(),
934            finished_at: Some(Utc::now()),
935            error_message: Some("Error: 错误信息 with émojis 🔥💥".to_string()),
936            previous_env: None,
937        };
938
939        // Add hook result with unicode output
940        let unicode_hook = Hook {
941            command: "echo".to_string(),
942            args: vec![],
943            dir: None,
944            inputs: vec![],
945            source: None,
946        };
947        let unicode_result = HookResult {
948            hook: unicode_hook,
949            success: false,
950            exit_status: Some(1),
951            stdout: "输出: Hello 世界! 🌍".to_string(),
952            stderr: "错误: ñoño error ⚠️".to_string(),
953            duration_ms: 100,
954            error: Some("失败了 😢".to_string()),
955        };
956        unicode_state.hook_results.insert(0, unicode_result);
957
958        // Save and load the state
959        state_manager.save_state(&unicode_state).await.unwrap();
960
961        let loaded = state_manager
962            .load_state(&unicode_state.instance_hash)
963            .await
964            .unwrap()
965            .unwrap();
966
967        // Verify all unicode content is preserved
968        assert_eq!(loaded.config_hash, "config_ñ_é_ü");
969        assert_eq!(
970            loaded.error_message,
971            Some("Error: 错误信息 with émojis 🔥💥".to_string())
972        );
973
974        let hook_result = loaded.hook_results.get(&0).unwrap();
975        assert_eq!(hook_result.stdout, "输出: Hello 世界! 🌍");
976        assert_eq!(hook_result.stderr, "错误: ñoño error ⚠️");
977        assert_eq!(hook_result.error, Some("失败了 😢".to_string()));
978    }
979
980    #[tokio::test]
981    async fn test_state_directory_with_many_states() {
982        let temp_dir = TempDir::new().unwrap();
983        let state_manager = StateManager::new(temp_dir.path().to_path_buf());
984
985        // Create many states to test scalability
986        for i in 0..50 {
987            let state = HookExecutionState {
988                instance_hash: format!("hash_{}", i),
989                directory_path: PathBuf::from(format!("/dir/{}", i)),
990                config_hash: format!("config_{}", i),
991                status: if i % 3 == 0 {
992                    ExecutionStatus::Completed
993                } else if i % 3 == 1 {
994                    ExecutionStatus::Running
995                } else {
996                    ExecutionStatus::Failed
997                },
998                total_hooks: 1,
999                completed_hooks: if i % 3 == 0 { 1 } else { 0 },
1000                current_hook_index: if i % 3 == 1 { Some(0) } else { None },
1001                hook_results: HashMap::new(),
1002                environment_vars: HashMap::new(),
1003                started_at: Utc::now() - chrono::Duration::hours(i as i64),
1004                finished_at: if i % 3 != 1 {
1005                    Some(Utc::now() - chrono::Duration::hours(i as i64 - 1))
1006                } else {
1007                    None
1008                },
1009                error_message: if i % 3 == 2 {
1010                    Some(format!("Error {}", i))
1011                } else {
1012                    None
1013                },
1014                previous_env: None,
1015            };
1016            state_manager.save_state(&state).await.unwrap();
1017        }
1018
1019        // List all states
1020        let listed = state_manager.list_active_states().await.unwrap();
1021        assert_eq!(listed.len(), 50);
1022
1023        // Clean up old completed states (older than 24 hours)
1024        let cleaned = state_manager
1025            .cleanup_orphaned_states(chrono::Duration::hours(24))
1026            .await
1027            .unwrap();
1028
1029        // Should clean up states older than 24 hours
1030        assert!(cleaned > 0);
1031    }
1032}