Skip to main content

vtcode_core/utils/
file_workflow.rs

1//! # File-based Workflow Utilities
2//!
3//! This module provides utilities for handling large inputs by reading content from files
4//! referenced in user input, supporting the Claude Code pattern of using files for large inputs.
5
6use anyhow::Result;
7use std::path::Path;
8use tokio::fs;
9
10/// Include content from text files referenced in user input
11///
12/// This function looks for patterns like `@./path/to/file.txt` and replaces them with
13/// the actual file content, allowing users to reference large files without pasting content.
14///
15/// # Arguments
16///
17/// * `input` - The user input text that may contain @ patterns referencing files
18/// * `base_dir` - The base directory to resolve relative paths from
19///
20/// # Returns
21///
22/// * `String` - The input with file references replaced by their content
23pub async fn include_file_content(input: &str, base_dir: &Path) -> Result<String> {
24    let matches = vtcode_commons::at_pattern::find_at_patterns(input);
25    if matches.is_empty() {
26        return Ok(input.to_string());
27    }
28
29    let mut result = String::new();
30    let mut last_end = 0;
31
32    for m in matches {
33        // Add the text before this match
34        if m.start > last_end {
35            result.push_str(&input[last_end..m.start]);
36        }
37
38        // Security validation: prevent directory traversal and validate path
39        if !vtcode_commons::paths::is_safe_relative_path(m.path) {
40            // Skip invalid paths (likely directory traversal attempts)
41            result.push_str(m.full_match);
42            last_end = m.end;
43            continue;
44        }
45
46        // Try to read the file as text content
47        let file_path = base_dir.join(m.path.trim());
48
49        match fs::read_to_string(&file_path).await {
50            Ok(file_content) => {
51                // Include the file content
52                result.push_str(&file_content);
53            }
54            Err(_) => {
55                // If it's not a valid text file, treat as literal @ usage
56                result.push_str(m.full_match);
57            }
58        }
59
60        last_end = m.end;
61    }
62
63    // Add any remaining text after the last match
64    if last_end < input.len() {
65        result.push_str(&input[last_end..]);
66    }
67
68    Ok(result)
69}
70
71#[cfg(test)]
72mod tests {
73    use super::*;
74    use std::io::Write;
75    use tempfile::TempDir;
76
77    #[tokio::test]
78    async fn test_include_file_content_with_text_file() {
79        let temp_dir = TempDir::new().unwrap();
80        let file_path = temp_dir.path().join("test.txt");
81
82        // Create a test file with content
83        let mut temp_file = std::io::BufWriter::new(std::fs::File::create(&file_path).unwrap());
84        writeln!(temp_file, "This is test file content").unwrap();
85        temp_file.flush().unwrap();
86
87        let input = format!(
88            "Look at this file: @{}",
89            file_path.file_name().unwrap().to_string_lossy()
90        );
91
92        let result = include_file_content(&input, temp_dir.path()).await.unwrap();
93
94        // Should contain the file content instead of the @ reference
95        assert!(result.contains("This is test file content"));
96        assert!(!result.contains('@')); // The @ symbol should be replaced
97    }
98
99    #[tokio::test]
100    async fn test_include_file_content_regular_text() {
101        let temp_dir = TempDir::new().unwrap();
102        let input = "This is just regular text with @ symbol not followed by file";
103
104        let result = include_file_content(input, temp_dir.path()).await.unwrap();
105
106        // Should return original text since there's no valid file
107        assert_eq!(result, input);
108    }
109
110    #[tokio::test]
111    async fn test_include_file_content_invalid_file() {
112        let temp_dir = TempDir::new().unwrap();
113        let input = "Look at @nonexistent.txt which doesn't exist";
114
115        let result = include_file_content(input, temp_dir.path()).await.unwrap();
116
117        // Should return original text since file doesn't exist
118        assert_eq!(result, input);
119    }
120
121    #[tokio::test]
122    async fn test_include_file_content_with_quoted_path() {
123        let temp_dir = TempDir::new().unwrap();
124        let file_path = temp_dir.path().join("file with spaces.txt");
125
126        // Create a test file with content
127        let mut temp_file = std::io::BufWriter::new(std::fs::File::create(&file_path).unwrap());
128        writeln!(temp_file, "Content with spaces in filename").unwrap();
129        temp_file.flush().unwrap();
130
131        let input = format!(
132            "Look at this file: @\"{}\"",
133            file_path.file_name().unwrap().to_string_lossy()
134        );
135
136        let result = include_file_content(&input, temp_dir.path()).await.unwrap();
137
138        // Should contain the file content
139        assert!(result.contains("Content with spaces in filename"));
140    }
141
142    #[test]
143    fn test_is_safe_relative_path() {
144        use vtcode_commons::paths::is_safe_relative_path;
145        assert!(!is_safe_relative_path("../../etc/passwd"));
146        assert!(!is_safe_relative_path("../file.txt"));
147        assert!(is_safe_relative_path("file.txt"));
148        assert!(is_safe_relative_path("./path/file.txt"));
149        assert!(is_safe_relative_path(" path with spaces .txt "));
150    }
151}