Skip to main content

vtcode_core/tools/
file_tracker.rs

1use anyhow::{Context, Result};
2use serde::Serialize;
3use std::path::PathBuf;
4use std::time::SystemTime;
5
6/// Tracks files generated by code execution
7#[derive(Debug, Clone)]
8pub struct FileTracker {
9    workspace_root: PathBuf,
10    tracked_patterns: Vec<String>,
11}
12
13impl FileTracker {
14    pub fn new(workspace_root: PathBuf) -> Self {
15        Self {
16            workspace_root,
17            tracked_patterns: vec![
18                "*.pdf".to_string(),
19                "*.xlsx".to_string(),
20                "*.csv".to_string(),
21                "*.docx".to_string(),
22                "*.png".to_string(),
23                "*.jpg".to_string(),
24                "*.json".to_string(),
25                "*.xml".to_string(),
26            ],
27        }
28    }
29
30    /// Scan for newly created or modified files matching tracked patterns
31    pub async fn detect_new_files(&self, since: SystemTime) -> Result<Vec<TrackedFile>> {
32        use tokio::task;
33        let workspace = self.workspace_root.clone();
34        let patterns = self.tracked_patterns.clone();
35
36        let entries = task::spawn_blocking(move || {
37            let mut files: Vec<PathBuf> = Vec::new();
38            let glob_patterns: Vec<_> = patterns
39                .iter()
40                .map(|p| glob::Pattern::new(p))
41                .collect::<Result<Vec<_>, _>>()
42                .unwrap_or_default();
43
44            let walker = vtcode_commons::walk::build_walker_single_threaded(&workspace).build();
45            for entry in walker.flatten() {
46                if !entry.file_type().is_some_and(|ft| ft.is_file()) {
47                    continue;
48                }
49                let path = entry.path();
50                let relative = path.strip_prefix(&workspace).unwrap_or(path);
51                let relative_str = relative.to_string_lossy();
52                for pattern in &glob_patterns {
53                    if pattern.matches(&relative_str) {
54                        files.push(path.to_path_buf());
55                        break;
56                    }
57                }
58            }
59            files
60        })
61        .await
62        .context("Failed to spawn file search task")?;
63
64        let mut new_files = Vec::new();
65        for file_path in entries {
66            if let Ok(metadata) = tokio::fs::metadata(&file_path).await
67                && let Ok(modified) = metadata.modified()
68                && modified >= since
69            {
70                new_files.push(TrackedFile {
71                    absolute_path: file_path,
72                    size: metadata.len(),
73                    modified,
74                });
75            }
76        }
77
78        Ok(new_files)
79    }
80
81    /// Generate a verification code snippet
82    pub fn generate_verification_code(&self, filename: &str) -> String {
83        format!(
84            r#"
85import os
86import sys
87
88file_path = os.path.abspath('{filename}')
89exists = os.path.exists(file_path)
90is_file = os.path.isfile(file_path) if exists else False
91size = os.path.getsize(file_path) if exists else 0
92
93print(f"FILE_VERIFICATION: {{file_path}} {{exists}} {{is_file}} {{size}}")
94sys.exit(0 if exists and is_file else 1)
95"#
96        )
97    }
98
99    /// Generate a summary of tracked files
100    pub fn generate_file_summary(&self, files: &[TrackedFile]) -> String {
101        if files.is_empty() {
102            return "No generated files detected.".to_string();
103        }
104
105        let mut summary = String::from("Generated files:\n");
106        for file in files {
107            summary.push_str(&format!("  - {} ({} bytes)\n", file.absolute_path.display(), file.size));
108        }
109        summary
110    }
111}
112
113/// Represents a tracked file with metadata
114#[derive(Debug, Clone, Serialize)]
115pub struct TrackedFile {
116    pub absolute_path: PathBuf,
117    pub size: u64,
118    #[serde(with = "system_time_serde")]
119    pub modified: SystemTime,
120}
121
122mod system_time_serde {
123    use serde::Serializer;
124    use std::time::{SystemTime, UNIX_EPOCH};
125
126    pub fn serialize<S>(time: &SystemTime, serializer: S) -> Result<S::Ok, S::Error>
127    where
128        S: Serializer,
129    {
130        let duration = time
131            .duration_since(UNIX_EPOCH)
132            .unwrap_or_else(|_| std::time::Duration::from_secs(0));
133        serializer.serialize_u64(duration.as_secs())
134    }
135}
136
137#[cfg(test)]
138mod tests {
139    use super::*;
140    use tokio;
141
142    #[tokio::test]
143    async fn test_file_tracker_basic() {
144        let tracker = FileTracker::new(PathBuf::from("/test"));
145        assert!(!tracker.tracked_patterns.is_empty());
146    }
147
148    #[tokio::test]
149    async fn test_verification_code_generation() {
150        let tracker = FileTracker::new(PathBuf::from("/workspace"));
151        let code = tracker.generate_verification_code("test.pdf");
152        assert!(code.contains("test.pdf"));
153        assert!(code.contains("FILE_VERIFICATION"));
154    }
155
156    #[test]
157    fn test_file_summary_generation() {
158        let tracker = FileTracker::new(PathBuf::from("/workspace"));
159        let files = vec![TrackedFile {
160            absolute_path: PathBuf::from("/workspace/test.pdf"),
161            size: 1024,
162            modified: SystemTime::now(),
163        }];
164
165        let summary = tracker.generate_file_summary(&files);
166        assert!(summary.contains("test.pdf"));
167        assert!(summary.contains("1024 bytes"));
168    }
169}