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!("Look at this file: @{}", file_path.file_name().unwrap().to_string_lossy());
88
89        let result = include_file_content(&input, temp_dir.path()).await.unwrap();
90
91        // Should contain the file content instead of the @ reference
92        assert!(result.contains("This is test file content"));
93        assert!(!result.contains('@')); // The @ symbol should be replaced
94    }
95
96    #[tokio::test]
97    async fn test_include_file_content_regular_text() {
98        let temp_dir = TempDir::new().unwrap();
99        let input = "This is just regular text with @ symbol not followed by file";
100
101        let result = include_file_content(input, temp_dir.path()).await.unwrap();
102
103        // Should return original text since there's no valid file
104        assert_eq!(result, input);
105    }
106
107    #[tokio::test]
108    async fn test_include_file_content_invalid_file() {
109        let temp_dir = TempDir::new().unwrap();
110        let input = "Look at @nonexistent.txt which doesn't exist";
111
112        let result = include_file_content(input, temp_dir.path()).await.unwrap();
113
114        // Should return original text since file doesn't exist
115        assert_eq!(result, input);
116    }
117
118    #[tokio::test]
119    async fn test_include_file_content_with_quoted_path() {
120        let temp_dir = TempDir::new().unwrap();
121        let file_path = temp_dir.path().join("file with spaces.txt");
122
123        // Create a test file with content
124        let mut temp_file = std::io::BufWriter::new(std::fs::File::create(&file_path).unwrap());
125        writeln!(temp_file, "Content with spaces in filename").unwrap();
126        temp_file.flush().unwrap();
127
128        let input = format!("Look at this file: @\"{}\"", file_path.file_name().unwrap().to_string_lossy());
129
130        let result = include_file_content(&input, temp_dir.path()).await.unwrap();
131
132        // Should contain the file content
133        assert!(result.contains("Content with spaces in filename"));
134    }
135
136    #[test]
137    fn test_is_safe_relative_path() {
138        use vtcode_commons::paths::is_safe_relative_path;
139        assert!(!is_safe_relative_path("../../etc/passwd"));
140        assert!(!is_safe_relative_path("../file.txt"));
141        assert!(is_safe_relative_path("file.txt"));
142        assert!(is_safe_relative_path("./path/file.txt"));
143        assert!(is_safe_relative_path(" path with spaces .txt "));
144    }
145}