Skip to main content

vtcode_core/tools/
file_tracker.rs

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