Skip to main content

winx_code_agent/tools/
context_save.rs

1use glob::glob;
2use std::fmt::Write as FmtWrite;
3use std::fs::{self, File};
4use std::io::Write;
5use std::path::{Path, PathBuf};
6use std::sync::Arc;
7use tokio::sync::Mutex;
8use tracing::{debug, warn};
9
10use crate::errors::{Result, WinxError};
11use crate::state::bash_state::BashState;
12use crate::state::persistence::{
13    load_bash_state_from_path, save_bash_state_to_path, BashStateSnapshot,
14};
15use crate::types::ContextSave;
16use crate::utils::path::expand_user;
17
18/// Handle a call to the `ContextSave` tool
19///
20/// This function processes a `ContextSave` request, saves context information about a task,
21/// including file contents from specified globs, to a single file.
22///
23/// # Arguments
24///
25/// * `bash_state` - Shared reference to the bash state
26/// * `args` - Parameters for the `ContextSave` operation
27///
28/// # Returns
29///
30/// A Result with the path where the context file was saved, or an error
31pub async fn handle_tool_call(
32    bash_state: &Arc<Mutex<Option<BashState>>>,
33    args: ContextSave,
34) -> Result<String> {
35    // Ensure bash state is initialized
36    let bash_state_guard = bash_state.lock().await;
37
38    let bash_state = bash_state_guard.as_ref().ok_or(WinxError::BashStateNotInitialized)?;
39
40    // Process the ContextSave request
41    let result = save_context(bash_state, args)?;
42
43    if should_open_context_file(&result) {
44        if let Err(e) = try_open_file(&result) {
45            debug!("Failed to open the context file: {}", e);
46        }
47    }
48
49    Ok(result)
50}
51
52/// Save the context information to a file
53///
54/// # Arguments
55///
56/// * `bash_state` - Reference to the bash state
57/// * `context` - The `ContextSave` parameters
58///
59/// # Returns
60///
61/// A Result with the path where the context file was saved, or an error
62fn save_context(bash_state: &BashState, mut context: ContextSave) -> Result<String> {
63    normalize_context(&mut context)?;
64
65    let (relevant_files, warnings) = collect_relevant_files(&context)?;
66    let memory_dir = resolve_memory_dir()?;
67    let relevant_files_data = read_files_content(&relevant_files, 10_000)?;
68    let memory_data = format_memory(&context, &relevant_files_data);
69    let safe_id = sanitize_filename(&context.id);
70
71    let (memory_file_path, state_file_path) = context_paths_from_dir(&memory_dir, &safe_id);
72    if let Some(response) = write_memory_file(&memory_file_path, &memory_data, &context)? {
73        return Ok(response);
74    }
75
76    write_bash_state_file(&state_file_path, bash_state);
77
78    Ok(context_save_response(&relevant_files, &warnings, &context, &memory_file_path))
79}
80
81fn normalize_context(context: &mut ContextSave) -> Result<()> {
82    if !context.project_root_path.is_empty() {
83        context.project_root_path = expand_user(&context.project_root_path);
84    }
85
86    if context.id.is_empty() {
87        return Err(WinxError::ArgumentParseError("Task ID cannot be empty".to_string()));
88    }
89
90    Ok(())
91}
92
93fn collect_relevant_files(context: &ContextSave) -> Result<(Vec<PathBuf>, Vec<String>)> {
94    let mut relevant_files = Vec::new();
95    let mut warnings = Vec::new();
96
97    for glob_pattern in &context.relevant_file_globs {
98        // Expand the glob pattern if it contains a tilde
99        let expanded_glob = expand_user(glob_pattern);
100
101        // If the glob is not absolute and we have a project root, make it relative to the project root
102        let final_glob =
103            if !Path::new(&expanded_glob).is_absolute() && !context.project_root_path.is_empty() {
104                PathBuf::from(&context.project_root_path)
105                    .join(expanded_glob)
106                    .to_string_lossy()
107                    .to_string()
108            } else {
109                expanded_glob
110            };
111
112        debug!("Processing glob pattern: {}", final_glob);
113
114        // Use the glob crate to find matching files
115        let matches = glob(&final_glob).map_err(|e| {
116            WinxError::ArgumentParseError(format!("Invalid glob pattern '{final_glob}': {e}"))
117        })?;
118
119        let mut found_files = false;
120        for entry in matches {
121            match entry {
122                Ok(path) => {
123                    if path.is_file() {
124                        relevant_files.push(path);
125                        found_files = true;
126                        // Limit to 1000 files per glob to avoid excessive processing
127                        if relevant_files.len() >= 1000 {
128                            warn!("Reached limit of 1000 files for glob '{}'", final_glob);
129                            break;
130                        }
131                    }
132                }
133                Err(e) => {
134                    warn!("Error matching glob '{}': {}", final_glob, e);
135                }
136            }
137        }
138
139        if !found_files {
140            warnings.push(format!("Warning: No files found for the glob: {glob_pattern}"));
141        }
142    }
143
144    debug!("Found {} relevant files", relevant_files.len());
145    Ok((relevant_files, warnings))
146}
147
148fn resolve_memory_dir() -> Result<PathBuf> {
149    let app_dir = match get_app_dir_xdg() {
150        Ok(dir) => dir,
151        Err(e) => {
152            debug!("Failed to get primary app directory: {:?}", e);
153            // Try using temporary directory directly as a last resort
154            let fallback = std::env::temp_dir().join("winx-memory");
155            debug!("Using fallback directory: {}", fallback.display());
156            fs::create_dir_all(&fallback).map_err(|e2| WinxError::FileAccessError {
157                path: fallback.clone(),
158                message: format!(
159                    "Failed to create fallback directory: {e2} (after previous error: {e:?})"
160                ),
161            })?;
162            fallback
163        }
164    };
165
166    let mut memory_dir = app_dir.join("memory");
167    match fs::create_dir_all(&memory_dir) {
168        Ok(()) => {}
169        Err(e) => {
170            debug!("Failed to create memory directory: {}", e);
171            // If we can't create the memory subdirectory, use the app dir directly as the memory_dir
172            debug!("Using app_dir directly as memory_dir due to failed subdirectory creation");
173            memory_dir = app_dir;
174        }
175    }
176
177    Ok(memory_dir)
178}
179
180pub(crate) fn context_paths(id: &str) -> Result<(PathBuf, PathBuf)> {
181    let safe_id = sanitize_filename(id);
182    let memory_dir = resolve_memory_dir()?;
183    Ok(context_paths_from_dir(&memory_dir, &safe_id))
184}
185
186fn context_paths_from_dir(memory_dir: &Path, safe_id: &str) -> (PathBuf, PathBuf) {
187    (
188        memory_dir.join(format!("{safe_id}.txt")),
189        memory_dir.join(format!("{safe_id}_bash_state.json")),
190    )
191}
192
193pub(crate) fn load_saved_context(id: &str) -> Result<Option<(String, Option<BashStateSnapshot>)>> {
194    if id.is_empty() {
195        return Ok(None);
196    }
197
198    let (memory_file_path, state_file_path) = context_paths(id)?;
199    if !memory_file_path.exists() {
200        return Ok(None);
201    }
202
203    let memory_data =
204        fs::read_to_string(&memory_file_path).map_err(|e| WinxError::FileAccessError {
205            path: memory_file_path.clone(),
206            message: format!("Failed to read saved context: {e}"),
207        })?;
208    let state = load_bash_state_from_path(&state_file_path).map_err(|e| {
209        WinxError::SerializationError(format!("Failed to load saved bash state: {e}"))
210    })?;
211
212    Ok(Some((memory_data, state)))
213}
214
215fn write_memory_file(
216    memory_file_path: &Path,
217    memory_data: &str,
218    context: &ContextSave,
219) -> Result<Option<String>> {
220    match File::create(memory_file_path) {
221        Ok(mut file) => {
222            if let Err(e) = file.write_all(memory_data.as_bytes()) {
223                warn!("Failed to write memory data: {}", e);
224                return Ok(Some(save_to_temp_file(memory_data, context)?));
225            }
226        }
227        Err(e) => {
228            warn!("Failed to create memory file: {}", e);
229            return Ok(Some(save_to_temp_file(memory_data, context)?));
230        }
231    }
232
233    Ok(None)
234}
235
236fn write_bash_state_file(state_file_path: &Path, bash_state: &BashState) {
237    // Try to create and write state file, but don't fail if it doesn't work
238    if let Err(e) = save_bash_state_to_path(state_file_path, &bash_state.snapshot()) {
239        warn!("Failed to write bash state data: {}", e);
240    }
241}
242
243fn context_save_response(
244    relevant_files: &[PathBuf],
245    warnings: &[String],
246    context: &ContextSave,
247    memory_file_path: &Path,
248) -> String {
249    let memory_file_path_str = memory_file_path.to_string_lossy().to_string();
250    if !relevant_files.is_empty() || context.relevant_file_globs.is_empty() {
251        if warnings.is_empty() {
252            memory_file_path_str
253        } else {
254            format!(
255                "{}\n\nContext file successfully saved at {}",
256                warnings.join("\n"),
257                memory_file_path_str
258            )
259        }
260    } else {
261        format!(
262            "Error: No files found for the given globs. Context file successfully saved at \"{memory_file_path_str}\", but please fix the error."
263        )
264    }
265}
266
267/// Format the memory data for saving
268///
269/// # Arguments
270///
271/// * `context` - The `ContextSave` parameters
272/// * `relevant_files_data` - The content of the relevant files
273///
274/// # Returns
275///
276/// A formatted string containing the memory data
277fn format_memory(context: &ContextSave, relevant_files_data: &str) -> String {
278    let mut memory_data = String::new();
279
280    // Add project root path if provided
281    if !context.project_root_path.is_empty() {
282        let _ = write!(memory_data, "Project root path: {}\n\n", context.project_root_path);
283    }
284
285    // Add the description
286    memory_data.push_str(&context.description);
287    memory_data.push_str("\n\n");
288
289    // Add the relevant file globs
290    let _ =
291        write!(memory_data, "Relevant file globs: {}\n\n", context.relevant_file_globs.join(", "));
292
293    // Add the content of the relevant files
294    memory_data.push_str("File contents:\n\n");
295    memory_data.push_str(relevant_files_data);
296
297    memory_data
298}
299
300/// Get the application directory for storing data
301///
302/// This function tries multiple locations in order of preference:
303/// 1. `XDG_DATA_HOME/winx` if `XDG_DATA_HOME` is set
304/// 2. HOME/.local/share/winx if HOME is set
305/// 3. Current directory/.winx-data as a fallback
306/// 4. Temporary directory as a last resort
307///
308/// # Returns
309///
310/// A Result with the path to the app directory
311fn get_app_dir_xdg() -> Result<PathBuf> {
312    // Try multiple locations in order of preference
313    let app_dir = get_primary_app_dir().or_else(|_| get_fallback_app_dir())?;
314
315    debug!("Using app directory: {}", app_dir.display());
316    Ok(app_dir)
317}
318
319/// Try to get the primary application directory
320fn get_primary_app_dir() -> Result<PathBuf> {
321    // Try XDG_DATA_HOME first
322    if let Ok(xdg_path) = std::env::var("XDG_DATA_HOME") {
323        let app_dir = PathBuf::from(xdg_path).join("winx");
324        if let Ok(()) = fs::create_dir_all(&app_dir) {
325            return Ok(app_dir);
326        }
327    }
328
329    // Try HOME/.local/share next
330    if let Ok(home) = std::env::var("HOME") {
331        let app_dir = PathBuf::from(home).join(".local/share/winx");
332        if let Ok(()) = fs::create_dir_all(&app_dir) {
333            return Ok(app_dir);
334        }
335    }
336
337    // No successful primary location
338    Err(WinxError::FileAccessError {
339        path: PathBuf::from("<primary-paths>"),
340        message: "Could not create app directory in primary locations".to_string(),
341    })
342}
343
344/// Get a fallback application directory when primary ones fail
345fn get_fallback_app_dir() -> Result<PathBuf> {
346    // Try current directory
347    let current_dir = std::env::current_dir().unwrap_or_else(|_| PathBuf::from("."));
348
349    let app_dir = current_dir.join(".winx-data");
350    if let Ok(()) = fs::create_dir_all(&app_dir) {
351        return Ok(app_dir);
352    }
353
354    // Try temporary directory as a last resort
355    let temp_dir = std::env::temp_dir().join("winx-data");
356    fs::create_dir_all(&temp_dir).map_err(|e| WinxError::FileAccessError {
357        path: temp_dir.clone(),
358        message: format!("Failed to create app directory in any location: {e}"),
359    })?;
360
361    Ok(temp_dir)
362}
363
364/// Read the content of multiple files
365///
366/// # Arguments
367///
368/// * `file_paths` - List of paths to the files to read
369/// * `max_files` - Maximum number of files to read
370///
371/// # Returns
372///
373/// A Result with the content of the files, or an error
374fn read_files_content(file_paths: &[PathBuf], max_files: usize) -> Result<String> {
375    let mut result = String::new();
376    let mut skipped_binary = Vec::new();
377
378    for (i, path) in file_paths.iter().take(max_files).enumerate() {
379        // Try to read as UTF-8 text, skip binary files
380        match fs::read_to_string(path) {
381            Ok(file_content) => {
382                let _ = writeln!(result, "--- File {}: {} ---", i + 1, path.display());
383                result.push_str(&file_content);
384                result.push_str("\n\n");
385            }
386            Err(e) if e.kind() == std::io::ErrorKind::InvalidData => {
387                // Binary file or invalid UTF-8 - skip with note
388                skipped_binary.push(path.display().to_string());
389            }
390            Err(e) => {
391                return Err(WinxError::FileAccessError {
392                    path: path.clone(),
393                    message: format!("Failed to read file: {e}"),
394                });
395            }
396        }
397    }
398
399    // Add note about skipped binary files
400    if !skipped_binary.is_empty() {
401        let _ = write!(
402            result,
403            "Note: Skipped {} binary/non-text file(s): {}\n\n",
404            skipped_binary.len(),
405            skipped_binary.join(", ")
406        );
407    }
408
409    if file_paths.len() > max_files {
410        let _ = writeln!(
411            result,
412            "Note: Only showing the first {} files out of {}.",
413            max_files,
414            file_paths.len()
415        );
416    }
417
418    Ok(result)
419}
420
421/// Try to open a file with the default application
422///
423/// # Arguments
424///
425/// * `file_path` - Path to the file to open
426///
427/// # Returns
428///
429/// A Result indicating success or failure
430fn try_open_file(file_path: &str) -> Result<()> {
431    if std::env::consts::OS != "macos" && std::env::consts::OS != "linux" {
432        // Skip on unsupported platforms
433        return Ok(());
434    }
435
436    // Get the command to use based on the OS
437    let cmd = if std::env::consts::OS == "macos" {
438        "open"
439    } else {
440        // Try to find which command is available on Linux
441        for cmd in &["xdg-open", "gnome-open", "kde-open"] {
442            let status = std::process::Command::new("which")
443                .arg(cmd)
444                .stdout(std::process::Stdio::null())
445                .stderr(std::process::Stdio::null())
446                .status();
447
448            if let Ok(status) = status {
449                if status.success() {
450                    // Found an available command, use it
451                    let _ = std::process::Command::new(cmd)
452                        .arg(file_path)
453                        .stdout(std::process::Stdio::null())
454                        .stderr(std::process::Stdio::null())
455                        .spawn()
456                        .map_err(|e| {
457                            WinxError::CommandExecutionError(format!(
458                                "Failed to spawn open command: {e}"
459                            ))
460                        })?;
461
462                    // We don't wait for the command to complete
463                    return Ok(());
464                }
465            }
466        }
467
468        // If no command is available, just return success
469        return Ok(());
470    };
471
472    // Try to open the file
473    let _ = std::process::Command::new(cmd)
474        .arg(file_path)
475        .stdout(std::process::Stdio::null())
476        .stderr(std::process::Stdio::null())
477        .spawn()
478        .map_err(|e| {
479            WinxError::CommandExecutionError(format!("Failed to spawn open command: {e}"))
480        })?;
481
482    // We don't actually need to wait for the command to complete
483    // Just let it run in the background
484    // (This mimics the Python implementation)
485
486    Ok(())
487}
488
489fn should_open_context_file(file_path: &str) -> bool {
490    std::env::var("WINX_OPEN_CONTEXT").is_ok_and(|value| value == "1" || value == "true")
491        && Path::new(file_path).is_file()
492        && !cfg!(test)
493}
494
495/// Save context data to a temporary file as a last resort
496fn save_to_temp_file(memory_data: &str, context: &ContextSave) -> Result<String> {
497    let temp_dir = std::env::temp_dir();
498    let safe_id = sanitize_filename(&context.id);
499    let temp_file_path = temp_dir.join(format!("winx-{safe_id}.txt"));
500
501    let mut file = File::create(&temp_file_path).map_err(|e| WinxError::FileAccessError {
502        path: temp_file_path.clone(),
503        message: format!("Failed to create temporary file: {e}"),
504    })?;
505
506    file.write_all(memory_data.as_bytes()).map_err(|e| WinxError::FileAccessError {
507        path: temp_file_path.clone(),
508        message: format!("Failed to write to temporary file: {e}"),
509    })?;
510
511    let path_str = temp_file_path.to_string_lossy().to_string();
512
513    Ok(format!(
514        "Context was saved to temporary file at {path_str} due to permission issues with regular locations."
515    ))
516}
517
518/// Sanitize a filename to ensure it's valid on all platforms
519fn sanitize_filename(input: &str) -> String {
520    let invalid_chars = vec!['/', '\\', ':', '*', '?', '"', '<', '>', '|'];
521    let mut result = input.to_string();
522
523    for c in invalid_chars {
524        result = result.replace(c, "_");
525    }
526
527    // Limit length to avoid issues
528    if result.len() > 50 {
529        use rand::RngExt;
530        result = format!("{}-{}", &result[0..45], rand::rng().random_range(1000..9999));
531    }
532
533    result
534}