winx_code_agent/tools/

context_save.rs

use glob::glob;
use std::fmt::Write as FmtWrite;
use std::fs::{self, File};
use std::io::Write;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tokio::sync::Mutex;
use tracing::{debug, warn};

use crate::errors::{Result, WinxError};
use crate::state::bash_state::BashState;
use crate::state::persistence::{
    load_bash_state_from_path, save_bash_state_to_path, BashStateSnapshot,
};
use crate::types::ContextSave;
use crate::utils::path::{expand_user, validate_path_in_workspace};

/// Handle a call to the `ContextSave` tool
///
/// This function processes a `ContextSave` request, saves context information about a task,
/// including file contents from specified globs, to a single file.
///
/// # Arguments
///
/// * `bash_state` - Shared reference to the bash state
/// * `args` - Parameters for the `ContextSave` operation
///
/// # Returns
///
/// A Result with the path where the context file was saved, or an error
pub async fn handle_tool_call(
    bash_state: &Arc<Mutex<Option<BashState>>>,
    args: ContextSave,
) -> Result<String> {
    // Ensure bash state is initialized
    let bash_state_guard = bash_state.lock().await;

    let bash_state = bash_state_guard.as_ref().ok_or(WinxError::BashStateNotInitialized)?;

    // Process the ContextSave request
    let result = save_context(bash_state, args)?;

    if should_open_context_file(&result) {
        if let Err(e) = try_open_file(&result) {
            debug!("Failed to open the context file: {}", e);
        }
    }

    Ok(result)
}

/// Save the context information to a file
///
/// # Arguments
///
/// * `bash_state` - Reference to the bash state
/// * `context` - The `ContextSave` parameters
///
/// # Returns
///
/// A Result with the path where the context file was saved, or an error
fn save_context(bash_state: &BashState, mut context: ContextSave) -> Result<String> {
    normalize_context(&mut context)?;

    let (relevant_files, warnings) = collect_relevant_files(&context, &bash_state.workspace_root)?;
    let memory_dir = resolve_memory_dir()?;
    let relevant_files_data = read_files_content(&relevant_files, 10_000)?;
    let memory_data = format_memory(&context, &relevant_files_data);
    let safe_id = sanitize_filename(&context.id);

    let (memory_file_path, state_file_path) = context_paths_from_dir(&memory_dir, &safe_id);
    if let Some(response) = write_memory_file(&memory_file_path, &memory_data, &context)? {
        return Ok(response);
    }

    write_bash_state_file(&state_file_path, bash_state);

    Ok(context_save_response(&relevant_files, &warnings, &context, &memory_file_path))
}

fn normalize_context(context: &mut ContextSave) -> Result<()> {
    if !context.project_root_path.is_empty() {
        context.project_root_path = expand_user(&context.project_root_path);
    }

    if context.id.is_empty() {
        return Err(WinxError::ArgumentParseError("Task ID cannot be empty".to_string()));
    }

    Ok(())
}

fn collect_relevant_files(
    context: &ContextSave,
    workspace_root: &Path,
) -> Result<(Vec<PathBuf>, Vec<String>)> {
    let mut relevant_files = Vec::new();
    let mut warnings = Vec::new();
    let mut skipped_outside = 0usize;

    for glob_pattern in &context.relevant_file_globs {
        // Expand the glob pattern if it contains a tilde
        let expanded_glob = expand_user(glob_pattern);

        // If the glob is not absolute and we have a project root, make it relative to the project root
        let final_glob =
            if !Path::new(&expanded_glob).is_absolute() && !context.project_root_path.is_empty() {
                PathBuf::from(&context.project_root_path)
                    .join(expanded_glob)
                    .to_string_lossy()
                    .to_string()
            } else {
                expanded_glob
            };

        debug!("Processing glob pattern: {}", final_glob);

        // Use the glob crate to find matching files
        let matches = glob(&final_glob).map_err(|e| {
            WinxError::ArgumentParseError(format!("Invalid glob pattern '{final_glob}': {e}"))
        })?;

        let mut found_files = false;
        for entry in matches {
            match entry {
                Ok(path) => {
                    if path.is_file() {
                        // Confine reads to the workspace. ContextSave inlines file
                        // contents into the saved memory, so without this an
                        // absolute glob (`/etc/**`, `~/.ssh/*`) or a `..`-escaping
                        // pattern would exfiltrate arbitrary files — over HTTP that
                        // is an authenticated read primitive. The other file tools
                        // already gate on this; ContextSave was the gap.
                        if validate_path_in_workspace(&path, workspace_root).is_err() {
                            skipped_outside += 1;
                            continue;
                        }
                        relevant_files.push(path);
                        found_files = true;
                        // Limit to 1000 files per glob to avoid excessive processing
                        if relevant_files.len() >= 1000 {
                            warn!("Reached limit of 1000 files for glob '{}'", final_glob);
                            break;
                        }
                    }
                }
                Err(e) => {
                    warn!("Error matching glob '{}': {}", final_glob, e);
                }
            }
        }

        if !found_files {
            warnings.push(format!("Warning: No files found for the glob: {glob_pattern}"));
        }
    }

    if skipped_outside > 0 {
        warnings.push(format!(
            "Warning: skipped {skipped_outside} file(s) that resolved outside the workspace root."
        ));
    }

    debug!("Found {} relevant files", relevant_files.len());
    Ok((relevant_files, warnings))
}

fn resolve_memory_dir() -> Result<PathBuf> {
    let app_dir = match get_app_dir_xdg() {
        Ok(dir) => dir,
        Err(e) => {
            debug!("Failed to get primary app directory: {:?}", e);
            // Try using temporary directory directly as a last resort
            let fallback = std::env::temp_dir().join("winx-memory");
            debug!("Using fallback directory: {}", fallback.display());
            fs::create_dir_all(&fallback).map_err(|e2| WinxError::FileAccessError {
                path: fallback.clone(),
                message: format!(
                    "Failed to create fallback directory: {e2} (after previous error: {e:?})"
                ),
            })?;
            fallback
        }
    };

    let mut memory_dir = app_dir.join("memory");
    match fs::create_dir_all(&memory_dir) {
        Ok(()) => {}
        Err(e) => {
            debug!("Failed to create memory directory: {}", e);
            // If we can't create the memory subdirectory, use the app dir directly as the memory_dir
            debug!("Using app_dir directly as memory_dir due to failed subdirectory creation");
            memory_dir = app_dir;
        }
    }

    Ok(memory_dir)
}

pub(crate) fn context_paths(id: &str) -> Result<(PathBuf, PathBuf)> {
    let safe_id = sanitize_filename(id);
    let memory_dir = resolve_memory_dir()?;
    Ok(context_paths_from_dir(&memory_dir, &safe_id))
}

fn context_paths_from_dir(memory_dir: &Path, safe_id: &str) -> (PathBuf, PathBuf) {
    (
        memory_dir.join(format!("{safe_id}.txt")),
        memory_dir.join(format!("{safe_id}_bash_state.json")),
    )
}

pub(crate) fn load_saved_context(id: &str) -> Result<Option<(String, Option<BashStateSnapshot>)>> {
    if id.is_empty() {
        return Ok(None);
    }

    let (memory_file_path, state_file_path) = context_paths(id)?;
    if !memory_file_path.exists() {
        tracing::warn!(
            task_id = %id,
            path = %memory_file_path.display(),
            "task_id_to_resume requested but no saved context found; starting fresh"
        );
        return Ok(None);
    }

    let memory_data =
        fs::read_to_string(&memory_file_path).map_err(|e| WinxError::FileAccessError {
            path: memory_file_path.clone(),
            message: format!("Failed to read saved context: {e}"),
        })?;
    // Cap the injected context so a large saved memory can't blow the Initialize
    // response past the model's window. Keep the head (project root + plan).
    let memory_data = truncate_memory_to_tokens(memory_data, memory_token_budget());
    let state = load_bash_state_from_path(&state_file_path).map_err(|e| {
        WinxError::SerializationError(format!("Failed to load saved bash state: {e}"))
    })?;

    Ok(Some((memory_data, state)))
}

/// Saved memory is treated as non-code prose; cap it like wcgw's
/// `noncoding_max_tokens` so resuming a task never floods the context window.
const MEMORY_MAX_TOKENS: usize = 8_000;

/// Memory token budget, overridable via `WINX_NONCODING_TOKEN_BUDGET` (saved
/// memory is prose, so it shares the non-coding budget with `ReadFiles`).
fn memory_token_budget() -> usize {
    crate::utils::encoder::budget_from_env("WINX_NONCODING_TOKEN_BUDGET", MEMORY_MAX_TOKENS)
}

/// Floor on each file's per-file token slice, so even a glob of many files still
/// shows a usable head of each one rather than a useless few tokens.
const MIN_PER_FILE_SAVE_TOKENS: usize = 500;

/// Keep the first `max_tokens` of a file's content and note what was dropped.
/// Used so a single large file in a `ContextSave` glob doesn't crowd out the rest.
fn cap_file_to_tokens(content: &str, max_tokens: usize) -> String {
    match crate::utils::encoder::encode_ids(content) {
        Some(ids) if ids.len() > max_tokens => {
            let omitted = ids.len() - max_tokens;
            let head = crate::utils::encoder::decode_ids(&ids[..max_tokens])
                .unwrap_or_else(|| content.chars().take(max_tokens.saturating_mul(4)).collect());
            format!("{head}\n[... {omitted} tokens truncated for context save ...]")
        }
        _ => content.to_string(),
    }
}

/// Truncate saved memory to `max_tokens`, keeping the head (where the project
/// root and the plan/status live) and appending a marker. Falls back to the
/// untouched text if the tokenizer is unavailable.
fn truncate_memory_to_tokens(text: String, max_tokens: usize) -> String {
    let Some(ids) = crate::utils::encoder::encode_ids(&text) else {
        return text;
    };
    if ids.len() <= max_tokens {
        return text;
    }
    let keep = max_tokens.saturating_sub(8);
    match crate::utils::encoder::decode_ids(&ids[..keep]) {
        Some(decoded) => format!("{decoded}\n(...truncated saved context)"),
        None => text,
    }
}

/// Extract the project root recorded at the top of a saved memory file, if any.
/// Matches the `Project root path: <path>` header written by `format_memory`.
pub(crate) fn extract_project_root(memory: &str) -> Option<PathBuf> {
    memory
        .lines()
        .find_map(|line| line.strip_prefix("Project root path:").map(str::trim))
        .filter(|root| !root.is_empty())
        .map(PathBuf::from)
}

fn write_memory_file(
    memory_file_path: &Path,
    memory_data: &str,
    context: &ContextSave,
) -> Result<Option<String>> {
    match File::create(memory_file_path) {
        Ok(mut file) => {
            if let Err(e) = file.write_all(memory_data.as_bytes()) {
                warn!("Failed to write memory data: {}", e);
                return Ok(Some(save_to_temp_file(memory_data, context)?));
            }
        }
        Err(e) => {
            warn!("Failed to create memory file: {}", e);
            return Ok(Some(save_to_temp_file(memory_data, context)?));
        }
    }

    Ok(None)
}

fn write_bash_state_file(state_file_path: &Path, bash_state: &BashState) {
    // Try to create and write state file, but don't fail if it doesn't work
    if let Err(e) = save_bash_state_to_path(state_file_path, &bash_state.snapshot()) {
        warn!("Failed to write bash state data: {}", e);
    }
}

fn context_save_response(
    relevant_files: &[PathBuf],
    warnings: &[String],
    context: &ContextSave,
    memory_file_path: &Path,
) -> String {
    let memory_file_path_str = memory_file_path.to_string_lossy().to_string();
    if !relevant_files.is_empty() || context.relevant_file_globs.is_empty() {
        if warnings.is_empty() {
            memory_file_path_str
        } else {
            format!(
                "{}\n\nContext file successfully saved at {}",
                warnings.join("\n"),
                memory_file_path_str
            )
        }
    } else {
        format!(
            "Error: No files found for the given globs. Context file successfully saved at \"{memory_file_path_str}\", but please fix the error."
        )
    }
}

/// Format the memory data for saving
///
/// # Arguments
///
/// * `context` - The `ContextSave` parameters
/// * `relevant_files_data` - The content of the relevant files
///
/// # Returns
///
/// A formatted string containing the memory data
/// Make a glob portable for resume: if it's absolute and under `project_root`,
/// store it relative to the root (recorded separately in the memory); otherwise
/// keep it verbatim. A clone or another machine then still resolves it.
fn relativize_glob(glob: &str, project_root: &str) -> String {
    if project_root.is_empty() {
        return glob.to_string();
    }
    Path::new(glob)
        .strip_prefix(project_root)
        .map_or_else(|_| glob.to_string(), |rel| rel.to_string_lossy().to_string())
}

fn format_memory(context: &ContextSave, relevant_files_data: &str) -> String {
    let mut memory_data = String::new();

    // Add project root path if provided
    if !context.project_root_path.is_empty() {
        let _ = write!(memory_data, "Project root path: {}\n\n", context.project_root_path);
    }

    // Add the description
    memory_data.push_str(&context.description);
    memory_data.push_str("\n\n");

    // Add the relevant file globs, made portable: an absolute glob under the
    // project root is stored relative to it (the root is recorded above), so the
    // saved memory resumes in a clone — or on another machine — where the
    // absolute path no longer exists. Already-relative globs are kept as-is.
    let portable_globs = context
        .relevant_file_globs
        .iter()
        .map(|glob| relativize_glob(glob, &context.project_root_path))
        .collect::<Vec<_>>();
    let _ = write!(memory_data, "Relevant file globs: {}\n\n", portable_globs.join(", "));

    // Add the content of the relevant files
    memory_data.push_str("File contents:\n\n");
    memory_data.push_str(relevant_files_data);

    memory_data
}

/// Get the application directory for storing data
///
/// This function tries multiple locations in order of preference:
/// 1. `XDG_DATA_HOME/winx` if `XDG_DATA_HOME` is set
/// 2. HOME/.local/share/winx if HOME is set
/// 3. Current directory/.winx-data as a fallback
/// 4. Temporary directory as a last resort
///
/// # Returns
///
/// A Result with the path to the app directory
fn get_app_dir_xdg() -> Result<PathBuf> {
    // Try multiple locations in order of preference
    let app_dir = get_primary_app_dir().or_else(|_| get_fallback_app_dir())?;

    debug!("Using app directory: {}", app_dir.display());
    Ok(app_dir)
}

/// Try to get the primary application directory
fn get_primary_app_dir() -> Result<PathBuf> {
    // Try XDG_DATA_HOME first
    if let Ok(xdg_path) = std::env::var("XDG_DATA_HOME") {
        let app_dir = PathBuf::from(xdg_path).join("winx");
        if let Ok(()) = fs::create_dir_all(&app_dir) {
            return Ok(app_dir);
        }
    }

    // Try HOME/.local/share next
    if let Ok(home) = std::env::var("HOME") {
        let app_dir = PathBuf::from(home).join(".local/share/winx");
        if let Ok(()) = fs::create_dir_all(&app_dir) {
            return Ok(app_dir);
        }
    }

    // No successful primary location
    Err(WinxError::FileAccessError {
        path: PathBuf::from("<primary-paths>"),
        message: "Could not create app directory in primary locations".to_string(),
    })
}

/// Get a fallback application directory when primary ones fail
fn get_fallback_app_dir() -> Result<PathBuf> {
    // Try current directory
    let current_dir = std::env::current_dir().unwrap_or_else(|_| PathBuf::from("."));

    let app_dir = current_dir.join(".winx-data");
    if let Ok(()) = fs::create_dir_all(&app_dir) {
        return Ok(app_dir);
    }

    // Try temporary directory as a last resort
    let temp_dir = std::env::temp_dir().join("winx-data");
    fs::create_dir_all(&temp_dir).map_err(|e| WinxError::FileAccessError {
        path: temp_dir.clone(),
        message: format!("Failed to create app directory in any location: {e}"),
    })?;

    Ok(temp_dir)
}

/// Read the content of multiple files
///
/// # Arguments
///
/// * `file_paths` - List of paths to the files to read
/// * `max_files` - Maximum number of files to read
///
/// # Returns
///
/// A Result with the content of the files, or an error
fn read_files_content(file_paths: &[PathBuf], max_files: usize) -> Result<String> {
    let mut result = String::new();
    let mut skipped_binary = Vec::new();

    // Give each file a fair slice of the memory budget so one huge file (a 200KB
    // log) can't eat the whole window and bury the rest. Before this, files were
    // inlined whole and the global cap afterward kept only the head — usually
    // just the first file, losing every later file and each file's tail.
    let shown = file_paths.len().min(max_files).max(1);
    let per_file_tokens = (memory_token_budget() / shown).max(MIN_PER_FILE_SAVE_TOKENS);

    for (i, path) in file_paths.iter().take(max_files).enumerate() {
        // Try to read as UTF-8 text, skip binary files
        match fs::read_to_string(path) {
            Ok(file_content) => {
                let _ = writeln!(result, "--- File {}: {} ---", i + 1, path.display());
                result.push_str(&cap_file_to_tokens(&file_content, per_file_tokens));
                result.push_str("\n\n");
            }
            Err(e) if e.kind() == std::io::ErrorKind::InvalidData => {
                // Binary file or invalid UTF-8 - skip with note
                skipped_binary.push(path.display().to_string());
            }
            Err(e) => {
                return Err(WinxError::FileAccessError {
                    path: path.clone(),
                    message: format!("Failed to read file: {e}"),
                });
            }
        }
    }

    // Add note about skipped binary files
    if !skipped_binary.is_empty() {
        let _ = write!(
            result,
            "Note: Skipped {} binary/non-text file(s): {}\n\n",
            skipped_binary.len(),
            skipped_binary.join(", ")
        );
    }

    if file_paths.len() > max_files {
        let _ = writeln!(
            result,
            "Note: Only showing the first {} files out of {}.",
            max_files,
            file_paths.len()
        );
    }

    Ok(result)
}

/// Try to open a file with the default application
///
/// # Arguments
///
/// * `file_path` - Path to the file to open
///
/// # Returns
///
/// A Result indicating success or failure
fn try_open_file(file_path: &str) -> Result<()> {
    if std::env::consts::OS != "macos" && std::env::consts::OS != "linux" {
        // Skip on unsupported platforms
        return Ok(());
    }

    // Get the command to use based on the OS
    let cmd = if std::env::consts::OS == "macos" {
        "open"
    } else {
        // Try to find which command is available on Linux
        for cmd in &["xdg-open", "gnome-open", "kde-open"] {
            let status = std::process::Command::new("which")
                .arg(cmd)
                .stdout(std::process::Stdio::null())
                .stderr(std::process::Stdio::null())
                .status();

            if let Ok(status) = status {
                if status.success() {
                    // Found an available command, use it
                    let _ = std::process::Command::new(cmd)
                        .arg(file_path)
                        .stdout(std::process::Stdio::null())
                        .stderr(std::process::Stdio::null())
                        .spawn()
                        .map_err(|e| {
                            WinxError::CommandExecutionError(format!(
                                "Failed to spawn open command: {e}"
                            ))
                        })?;

                    // We don't wait for the command to complete
                    return Ok(());
                }
            }
        }

        // If no command is available, just return success
        return Ok(());
    };

    // Try to open the file
    let _ = std::process::Command::new(cmd)
        .arg(file_path)
        .stdout(std::process::Stdio::null())
        .stderr(std::process::Stdio::null())
        .spawn()
        .map_err(|e| {
            WinxError::CommandExecutionError(format!("Failed to spawn open command: {e}"))
        })?;

    // We don't actually need to wait for the command to complete
    // Just let it run in the background
    // (This mimics the Python implementation)

    Ok(())
}

fn should_open_context_file(file_path: &str) -> bool {
    std::env::var("WINX_OPEN_CONTEXT").is_ok_and(|value| value == "1" || value == "true")
        && Path::new(file_path).is_file()
        && !cfg!(test)
}

/// Save context data to a temporary file as a last resort
fn save_to_temp_file(memory_data: &str, context: &ContextSave) -> Result<String> {
    let temp_dir = std::env::temp_dir();
    let safe_id = sanitize_filename(&context.id);
    let temp_file_path = temp_dir.join(format!("winx-{safe_id}.txt"));

    let mut file = File::create(&temp_file_path).map_err(|e| WinxError::FileAccessError {
        path: temp_file_path.clone(),
        message: format!("Failed to create temporary file: {e}"),
    })?;

    file.write_all(memory_data.as_bytes()).map_err(|e| WinxError::FileAccessError {
        path: temp_file_path.clone(),
        message: format!("Failed to write to temporary file: {e}"),
    })?;

    let path_str = temp_file_path.to_string_lossy().to_string();

    Ok(format!(
        "Context was saved to temporary file at {path_str} due to permission issues with regular locations."
    ))
}

/// Sanitize a filename to ensure it's valid on all platforms
fn sanitize_filename(input: &str) -> String {
    let invalid_chars = vec!['/', '\\', ':', '*', '?', '"', '<', '>', '|'];
    let mut result = input.to_string();

    for c in invalid_chars {
        result = result.replace(c, "_");
    }

    // Limit length to avoid issues. Slice by chars, not bytes: a multibyte
    // task id whose 45th byte lands mid-code-point would panic on `&result[0..45]`.
    if result.chars().count() > 50 {
        use rand::RngExt;
        let head: String = result.chars().take(45).collect();
        result = format!("{head}-{}", rand::rng().random_range(1000..9999));
    }

    result
}

#[cfg(test)]
mod tests {
    use super::{cap_file_to_tokens, relativize_glob};

    #[test]
    fn relativize_glob_makes_absolute_portable() {
        // Absolute glob under the root -> relative.
        assert_eq!(relativize_glob("/home/u/proj/src/*.rs", "/home/u/proj"), "src/*.rs");
        // Already relative -> untouched.
        assert_eq!(relativize_glob("src/*.rs", "/home/u/proj"), "src/*.rs");
        // Outside the root -> kept verbatim.
        assert_eq!(relativize_glob("/etc/*.conf", "/home/u/proj"), "/etc/*.conf");
        // No project root -> unchanged.
        assert_eq!(relativize_glob("/abs/x", ""), "/abs/x");
    }

    #[test]
    fn small_file_is_untouched() {
        let content = "def add(a, b):\n    return a + b\n";
        assert_eq!(cap_file_to_tokens(content, 4_000), content);
    }

    #[test]
    fn large_file_is_capped_with_a_note() {
        // ~5k tokens of distinct words, capped to 100.
        use std::fmt::Write as _;
        let mut big = String::new();
        for i in 0..5_000 {
            let _ = write!(big, "word{i} ");
        }
        let capped = cap_file_to_tokens(&big, 100);
        assert!(capped.len() < big.len(), "should shrink");
        assert!(capped.contains("truncated for context save"), "should note truncation");
        assert!(capped.starts_with("word0"), "should keep the head");
    }
}