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::{Read, 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, validate_path_in_workspace};
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, &bash_state.workspace_root)?;
66    let memory_dir = resolve_memory_dir()?;
67    let relevant_files_data = read_files_content(&relevant_files, 10_000)?;
68    // Scrub credentials before they're written to the on-disk memory file (it is
69    // re-read into context on resume, so a leaked secret would persist).
70    let memory_data =
71        crate::utils::redact::redact(&format_memory(&context, &relevant_files_data)).into_owned();
72    let safe_id = sanitize_filename(&context.id);
73
74    let (memory_file_path, state_file_path) = context_paths_from_dir(&memory_dir, &safe_id);
75    if let Some(response) = write_memory_file(&memory_file_path, &memory_data, &context)? {
76        return Ok(response);
77    }
78
79    write_bash_state_file(&state_file_path, bash_state);
80
81    Ok(context_save_response(&relevant_files, &warnings, &context, &memory_file_path))
82}
83
84fn normalize_context(context: &mut ContextSave) -> Result<()> {
85    if !context.project_root_path.is_empty() {
86        context.project_root_path = expand_user(&context.project_root_path);
87    }
88
89    if context.id.is_empty() {
90        return Err(WinxError::ArgumentParseError("Task ID cannot be empty".to_string()));
91    }
92
93    Ok(())
94}
95
96fn collect_relevant_files(
97    context: &ContextSave,
98    workspace_root: &Path,
99) -> Result<(Vec<PathBuf>, Vec<String>)> {
100    let mut relevant_files = Vec::new();
101    let mut warnings = Vec::new();
102    let mut skipped_outside = 0usize;
103
104    for glob_pattern in &context.relevant_file_globs {
105        // Expand the glob pattern if it contains a tilde
106        let expanded_glob = expand_user(glob_pattern);
107
108        // If the glob is not absolute and we have a project root, make it relative to the project root
109        let final_glob =
110            if !Path::new(&expanded_glob).is_absolute() && !context.project_root_path.is_empty() {
111                PathBuf::from(&context.project_root_path)
112                    .join(expanded_glob)
113                    .to_string_lossy()
114                    .to_string()
115            } else {
116                expanded_glob
117            };
118
119        debug!("Processing glob pattern: {}", final_glob);
120
121        // Use the glob crate to find matching files
122        let matches = glob(&final_glob).map_err(|e| {
123            WinxError::ArgumentParseError(format!("Invalid glob pattern '{final_glob}': {e}"))
124        })?;
125
126        let mut found_files = false;
127        for entry in matches {
128            match entry {
129                Ok(path) => {
130                    if path.is_file() {
131                        // Confine reads to the workspace. ContextSave inlines file
132                        // contents into the saved memory, so without this an
133                        // absolute glob (`/etc/**`, `~/.ssh/*`) or a `..`-escaping
134                        // pattern would exfiltrate arbitrary files — over HTTP that
135                        // is an authenticated read primitive. The other file tools
136                        // already gate on this; ContextSave was the gap.
137                        if validate_path_in_workspace(&path, workspace_root).is_err() {
138                            skipped_outside += 1;
139                            continue;
140                        }
141                        relevant_files.push(path);
142                        found_files = true;
143                        // Limit to 1000 files per glob to avoid excessive processing
144                        if relevant_files.len() >= 1000 {
145                            warn!("Reached limit of 1000 files for glob '{}'", final_glob);
146                            break;
147                        }
148                    }
149                }
150                Err(e) => {
151                    warn!("Error matching glob '{}': {}", final_glob, e);
152                }
153            }
154        }
155
156        if !found_files {
157            warnings.push(format!("Warning: No files found for the glob: {glob_pattern}"));
158        }
159    }
160
161    if skipped_outside > 0 {
162        warnings.push(format!(
163            "Warning: skipped {skipped_outside} file(s) that resolved outside the workspace root."
164        ));
165    }
166
167    debug!("Found {} relevant files", relevant_files.len());
168    Ok((relevant_files, warnings))
169}
170
171fn resolve_memory_dir() -> Result<PathBuf> {
172    let app_dir = match get_app_dir_xdg() {
173        Ok(dir) => dir,
174        Err(e) => {
175            debug!("Failed to get primary app directory: {:?}", e);
176            // Try using temporary directory directly as a last resort
177            let fallback = std::env::temp_dir().join("winx-memory");
178            debug!("Using fallback directory: {}", fallback.display());
179            fs::create_dir_all(&fallback).map_err(|e2| WinxError::FileAccessError {
180                path: fallback.clone(),
181                message: format!(
182                    "Failed to create fallback directory: {e2} (after previous error: {e:?})"
183                ),
184            })?;
185            fallback
186        }
187    };
188
189    let mut memory_dir = app_dir.join("memory");
190    match fs::create_dir_all(&memory_dir) {
191        Ok(()) => {}
192        Err(e) => {
193            debug!("Failed to create memory directory: {}", e);
194            // If we can't create the memory subdirectory, use the app dir directly as the memory_dir
195            debug!("Using app_dir directly as memory_dir due to failed subdirectory creation");
196            memory_dir = app_dir;
197        }
198    }
199
200    Ok(memory_dir)
201}
202
203pub(crate) fn context_paths(id: &str) -> Result<(PathBuf, PathBuf)> {
204    let safe_id = sanitize_filename(id);
205    let memory_dir = resolve_memory_dir()?;
206    Ok(context_paths_from_dir(&memory_dir, &safe_id))
207}
208
209fn context_paths_from_dir(memory_dir: &Path, safe_id: &str) -> (PathBuf, PathBuf) {
210    (
211        memory_dir.join(format!("{safe_id}.txt")),
212        memory_dir.join(format!("{safe_id}_bash_state.json")),
213    )
214}
215
216pub(crate) fn load_saved_context(id: &str) -> Result<Option<(String, Option<BashStateSnapshot>)>> {
217    if id.is_empty() {
218        return Ok(None);
219    }
220
221    let (memory_file_path, state_file_path) = context_paths(id)?;
222    if !memory_file_path.exists() {
223        tracing::warn!(
224            task_id = %id,
225            path = %memory_file_path.display(),
226            "task_id_to_resume requested but no saved context found; starting fresh"
227        );
228        return Ok(None);
229    }
230
231    let memory_data =
232        fs::read_to_string(&memory_file_path).map_err(|e| WinxError::FileAccessError {
233            path: memory_file_path.clone(),
234            message: format!("Failed to read saved context: {e}"),
235        })?;
236    // Cap the injected context so a large saved memory can't blow the Initialize
237    // response past the model's window. Keep the head (project root + plan).
238    let memory_data = truncate_memory_to_tokens(memory_data, memory_token_budget());
239    let state = load_bash_state_from_path(&state_file_path).map_err(|e| {
240        WinxError::SerializationError(format!("Failed to load saved bash state: {e}"))
241    })?;
242
243    Ok(Some((memory_data, state)))
244}
245
246/// Saved memory is treated as non-code prose; cap it like wcgw's
247/// `noncoding_max_tokens` so resuming a task never floods the context window.
248const MEMORY_MAX_TOKENS: usize = 8_000;
249
250/// Memory token budget, overridable via `WINX_NONCODING_TOKEN_BUDGET` (saved
251/// memory is prose, so it shares the non-coding budget with `ReadFiles`).
252fn memory_token_budget() -> usize {
253    crate::utils::encoder::budget_from_env("WINX_NONCODING_TOKEN_BUDGET", MEMORY_MAX_TOKENS)
254}
255
256/// Floor on each file's per-file token slice, so even a glob of many files still
257/// shows a usable head of each one rather than a useless few tokens.
258const MIN_PER_FILE_SAVE_TOKENS: usize = 500;
259
260/// Byte ceiling for reading a single context file off disk. `cap_file_to_tokens`
261/// trims the result to a few KB of tokens anyway, so there's no point pulling more
262/// than this into memory — and a glob that matches a 50 MB `.log` would otherwise
263/// read the whole thing (blocking + GBs of churn) just to drop almost all of it.
264const MAX_CONTEXT_READ_BYTES: u64 = 2 * 1024 * 1024;
265
266/// Read a file's text for context saving, bounded to `MAX_CONTEXT_READ_BYTES`.
267/// Small files go through `read_to_string` unchanged (exact binary detection via
268/// `InvalidData`); oversized files read only the capped head, keeping the valid
269/// UTF-8 prefix when the byte cap lands mid-codepoint, and still report a
270/// genuinely-binary file (invalid from byte 0) as `InvalidData`.
271fn read_text_for_context(path: &Path) -> std::io::Result<String> {
272    let len = fs::metadata(path)?.len();
273    if len <= MAX_CONTEXT_READ_BYTES {
274        return fs::read_to_string(path);
275    }
276    let mut buf = Vec::new();
277    File::open(path)?.take(MAX_CONTEXT_READ_BYTES).read_to_end(&mut buf)?;
278    match String::from_utf8(buf) {
279        Ok(s) => Ok(s),
280        Err(e) => {
281            let valid = e.utf8_error().valid_up_to();
282            if valid == 0 {
283                return Err(std::io::Error::new(
284                    std::io::ErrorKind::InvalidData,
285                    "stream did not contain valid UTF-8",
286                ));
287            }
288            // The byte cap split a multibyte char; the prefix up to `valid` is
289            // valid UTF-8 by construction, so from_utf8_lossy inserts nothing.
290            let bytes = e.into_bytes();
291            Ok(String::from_utf8_lossy(&bytes[..valid]).into_owned())
292        }
293    }
294}
295
296/// Keep the first `max_tokens` of a file's content and note what was dropped.
297/// Used so a single large file in a `ContextSave` glob doesn't crowd out the rest.
298fn cap_file_to_tokens(content: &str, max_tokens: usize) -> String {
299    match crate::utils::encoder::encode_ids(content) {
300        Some(ids) if ids.len() > max_tokens => {
301            let omitted = ids.len() - max_tokens;
302            let head = crate::utils::encoder::decode_ids(&ids[..max_tokens])
303                .unwrap_or_else(|| content.chars().take(max_tokens.saturating_mul(4)).collect());
304            format!("{head}\n[... {omitted} tokens truncated for context save ...]")
305        }
306        _ => content.to_string(),
307    }
308}
309
310/// Truncate saved memory to `max_tokens`, keeping the head (where the project
311/// root and the plan/status live) and appending a marker. Falls back to the
312/// untouched text if the tokenizer is unavailable.
313fn truncate_memory_to_tokens(text: String, max_tokens: usize) -> String {
314    let Some(ids) = crate::utils::encoder::encode_ids(&text) else {
315        return text;
316    };
317    if ids.len() <= max_tokens {
318        return text;
319    }
320    let keep = max_tokens.saturating_sub(8);
321    match crate::utils::encoder::decode_ids(&ids[..keep]) {
322        Some(decoded) => format!("{decoded}\n(...truncated saved context)"),
323        None => text,
324    }
325}
326
327/// Extract the project root recorded at the top of a saved memory file, if any.
328/// Matches the `Project root path: <path>` header written by `format_memory`.
329pub(crate) fn extract_project_root(memory: &str) -> Option<PathBuf> {
330    memory
331        .lines()
332        .find_map(|line| line.strip_prefix("Project root path:").map(str::trim))
333        .filter(|root| !root.is_empty())
334        .map(PathBuf::from)
335}
336
337fn write_memory_file(
338    memory_file_path: &Path,
339    memory_data: &str,
340    context: &ContextSave,
341) -> Result<Option<String>> {
342    match File::create(memory_file_path) {
343        Ok(mut file) => {
344            if let Err(e) = file.write_all(memory_data.as_bytes()) {
345                warn!("Failed to write memory data: {}", e);
346                return Ok(Some(save_to_temp_file(memory_data, context)?));
347            }
348        }
349        Err(e) => {
350            warn!("Failed to create memory file: {}", e);
351            return Ok(Some(save_to_temp_file(memory_data, context)?));
352        }
353    }
354
355    Ok(None)
356}
357
358fn write_bash_state_file(state_file_path: &Path, bash_state: &BashState) {
359    // Try to create and write state file, but don't fail if it doesn't work
360    if let Err(e) = save_bash_state_to_path(state_file_path, &bash_state.snapshot()) {
361        warn!("Failed to write bash state data: {}", e);
362    }
363}
364
365fn context_save_response(
366    relevant_files: &[PathBuf],
367    warnings: &[String],
368    context: &ContextSave,
369    memory_file_path: &Path,
370) -> String {
371    let memory_file_path_str = memory_file_path.to_string_lossy().to_string();
372    if !relevant_files.is_empty() || context.relevant_file_globs.is_empty() {
373        if warnings.is_empty() {
374            memory_file_path_str
375        } else {
376            format!(
377                "{}\n\nContext file successfully saved at {}",
378                warnings.join("\n"),
379                memory_file_path_str
380            )
381        }
382    } else {
383        format!(
384            "Error: No files found for the given globs. Context file successfully saved at \"{memory_file_path_str}\", but please fix the error."
385        )
386    }
387}
388
389/// Format the memory data for saving
390///
391/// # Arguments
392///
393/// * `context` - The `ContextSave` parameters
394/// * `relevant_files_data` - The content of the relevant files
395///
396/// # Returns
397///
398/// A formatted string containing the memory data
399/// Make a glob portable for resume: if it's absolute and under `project_root`,
400/// store it relative to the root (recorded separately in the memory); otherwise
401/// keep it verbatim. A clone or another machine then still resolves it.
402fn relativize_glob(glob: &str, project_root: &str) -> String {
403    if project_root.is_empty() {
404        return glob.to_string();
405    }
406    Path::new(glob)
407        .strip_prefix(project_root)
408        .map_or_else(|_| glob.to_string(), |rel| rel.to_string_lossy().to_string())
409}
410
411fn format_memory(context: &ContextSave, relevant_files_data: &str) -> String {
412    let mut memory_data = String::new();
413
414    // Add project root path if provided
415    if !context.project_root_path.is_empty() {
416        let _ = write!(memory_data, "Project root path: {}\n\n", context.project_root_path);
417    }
418
419    // Add the description
420    memory_data.push_str(&context.description);
421    memory_data.push_str("\n\n");
422
423    // Add the relevant file globs, made portable: an absolute glob under the
424    // project root is stored relative to it (the root is recorded above), so the
425    // saved memory resumes in a clone — or on another machine — where the
426    // absolute path no longer exists. Already-relative globs are kept as-is.
427    let portable_globs = context
428        .relevant_file_globs
429        .iter()
430        .map(|glob| relativize_glob(glob, &context.project_root_path))
431        .collect::<Vec<_>>();
432    let _ = write!(memory_data, "Relevant file globs: {}\n\n", portable_globs.join(", "));
433
434    // Add the content of the relevant files
435    memory_data.push_str("File contents:\n\n");
436    memory_data.push_str(relevant_files_data);
437
438    memory_data
439}
440
441/// Get the application directory for storing data
442///
443/// This function tries multiple locations in order of preference:
444/// 1. `XDG_DATA_HOME/winx` if `XDG_DATA_HOME` is set
445/// 2. HOME/.local/share/winx if HOME is set
446/// 3. Current directory/.winx-data as a fallback
447/// 4. Temporary directory as a last resort
448///
449/// # Returns
450///
451/// A Result with the path to the app directory
452fn get_app_dir_xdg() -> Result<PathBuf> {
453    // Try multiple locations in order of preference
454    let app_dir = get_primary_app_dir().or_else(|_| get_fallback_app_dir())?;
455
456    debug!("Using app directory: {}", app_dir.display());
457    Ok(app_dir)
458}
459
460/// Try to get the primary application directory
461fn get_primary_app_dir() -> Result<PathBuf> {
462    // Try XDG_DATA_HOME first
463    if let Ok(xdg_path) = std::env::var("XDG_DATA_HOME") {
464        let app_dir = PathBuf::from(xdg_path).join("winx");
465        if let Ok(()) = fs::create_dir_all(&app_dir) {
466            return Ok(app_dir);
467        }
468    }
469
470    // Try HOME/.local/share next
471    if let Ok(home) = std::env::var("HOME") {
472        let app_dir = PathBuf::from(home).join(".local/share/winx");
473        if let Ok(()) = fs::create_dir_all(&app_dir) {
474            return Ok(app_dir);
475        }
476    }
477
478    // No successful primary location
479    Err(WinxError::FileAccessError {
480        path: PathBuf::from("<primary-paths>"),
481        message: "Could not create app directory in primary locations".to_string(),
482    })
483}
484
485/// Get a fallback application directory when primary ones fail
486fn get_fallback_app_dir() -> Result<PathBuf> {
487    // Try current directory
488    let current_dir = std::env::current_dir().unwrap_or_else(|_| PathBuf::from("."));
489
490    let app_dir = current_dir.join(".winx-data");
491    if let Ok(()) = fs::create_dir_all(&app_dir) {
492        return Ok(app_dir);
493    }
494
495    // Try temporary directory as a last resort
496    let temp_dir = std::env::temp_dir().join("winx-data");
497    fs::create_dir_all(&temp_dir).map_err(|e| WinxError::FileAccessError {
498        path: temp_dir.clone(),
499        message: format!("Failed to create app directory in any location: {e}"),
500    })?;
501
502    Ok(temp_dir)
503}
504
505/// Read the content of multiple files
506///
507/// # Arguments
508///
509/// * `file_paths` - List of paths to the files to read
510/// * `max_files` - Maximum number of files to read
511///
512/// # Returns
513///
514/// A Result with the content of the files, or an error
515fn read_files_content(file_paths: &[PathBuf], max_files: usize) -> Result<String> {
516    let mut result = String::new();
517    let mut skipped_binary = Vec::new();
518
519    // Give each file a fair slice of the memory budget so one huge file (a 200KB
520    // log) can't eat the whole window and bury the rest. Before this, files were
521    // inlined whole and the global cap afterward kept only the head — usually
522    // just the first file, losing every later file and each file's tail.
523    let shown = file_paths.len().min(max_files).max(1);
524    let per_file_tokens = (memory_token_budget() / shown).max(MIN_PER_FILE_SAVE_TOKENS);
525
526    for (i, path) in file_paths.iter().take(max_files).enumerate() {
527        // Try to read as UTF-8 text (capped, so a giant file can't OOM/block the
528        // executor), skip binary files.
529        match read_text_for_context(path) {
530            Ok(file_content) => {
531                let _ = writeln!(result, "--- File {}: {} ---", i + 1, path.display());
532                result.push_str(&cap_file_to_tokens(&file_content, per_file_tokens));
533                result.push_str("\n\n");
534            }
535            Err(e) if e.kind() == std::io::ErrorKind::InvalidData => {
536                // Binary file or invalid UTF-8 - skip with note
537                skipped_binary.push(path.display().to_string());
538            }
539            Err(e) => {
540                return Err(WinxError::FileAccessError {
541                    path: path.clone(),
542                    message: format!("Failed to read file: {e}"),
543                });
544            }
545        }
546    }
547
548    // Add note about skipped binary files
549    if !skipped_binary.is_empty() {
550        let _ = write!(
551            result,
552            "Note: Skipped {} binary/non-text file(s): {}\n\n",
553            skipped_binary.len(),
554            skipped_binary.join(", ")
555        );
556    }
557
558    if file_paths.len() > max_files {
559        let _ = writeln!(
560            result,
561            "Note: Only showing the first {} files out of {}.",
562            max_files,
563            file_paths.len()
564        );
565    }
566
567    Ok(result)
568}
569
570/// Try to open a file with the default application
571///
572/// # Arguments
573///
574/// * `file_path` - Path to the file to open
575///
576/// # Returns
577///
578/// A Result indicating success or failure
579fn try_open_file(file_path: &str) -> Result<()> {
580    if std::env::consts::OS != "macos" && std::env::consts::OS != "linux" {
581        // Skip on unsupported platforms
582        return Ok(());
583    }
584
585    // Get the command to use based on the OS
586    let cmd = if std::env::consts::OS == "macos" {
587        "open"
588    } else {
589        // Try to find which command is available on Linux
590        for cmd in &["xdg-open", "gnome-open", "kde-open"] {
591            let status = std::process::Command::new("which")
592                .arg(cmd)
593                .stdout(std::process::Stdio::null())
594                .stderr(std::process::Stdio::null())
595                .status();
596
597            if let Ok(status) = status {
598                if status.success() {
599                    // Found an available command, use it
600                    let _ = std::process::Command::new(cmd)
601                        .arg(file_path)
602                        .stdout(std::process::Stdio::null())
603                        .stderr(std::process::Stdio::null())
604                        .spawn()
605                        .map_err(|e| {
606                            WinxError::CommandExecutionError(format!(
607                                "Failed to spawn open command: {e}"
608                            ))
609                        })?;
610
611                    // We don't wait for the command to complete
612                    return Ok(());
613                }
614            }
615        }
616
617        // If no command is available, just return success
618        return Ok(());
619    };
620
621    // Try to open the file
622    let _ = std::process::Command::new(cmd)
623        .arg(file_path)
624        .stdout(std::process::Stdio::null())
625        .stderr(std::process::Stdio::null())
626        .spawn()
627        .map_err(|e| {
628            WinxError::CommandExecutionError(format!("Failed to spawn open command: {e}"))
629        })?;
630
631    // We don't actually need to wait for the command to complete
632    // Just let it run in the background
633    // (This mimics the Python implementation)
634
635    Ok(())
636}
637
638fn should_open_context_file(file_path: &str) -> bool {
639    std::env::var("WINX_OPEN_CONTEXT").is_ok_and(|value| value == "1" || value == "true")
640        && Path::new(file_path).is_file()
641        && !cfg!(test)
642}
643
644/// Save context data to a temporary file as a last resort
645fn save_to_temp_file(memory_data: &str, context: &ContextSave) -> Result<String> {
646    let temp_dir = std::env::temp_dir();
647    let safe_id = sanitize_filename(&context.id);
648    let temp_file_path = temp_dir.join(format!("winx-{safe_id}.txt"));
649
650    let mut file = File::create(&temp_file_path).map_err(|e| WinxError::FileAccessError {
651        path: temp_file_path.clone(),
652        message: format!("Failed to create temporary file: {e}"),
653    })?;
654
655    file.write_all(memory_data.as_bytes()).map_err(|e| WinxError::FileAccessError {
656        path: temp_file_path.clone(),
657        message: format!("Failed to write to temporary file: {e}"),
658    })?;
659
660    let path_str = temp_file_path.to_string_lossy().to_string();
661
662    Ok(format!(
663        "Context was saved to temporary file at {path_str} due to permission issues with regular locations."
664    ))
665}
666
667/// Sanitize a filename to ensure it's valid on all platforms
668fn sanitize_filename(input: &str) -> String {
669    let invalid_chars = vec!['/', '\\', ':', '*', '?', '"', '<', '>', '|'];
670    let mut result = input.to_string();
671
672    for c in invalid_chars {
673        result = result.replace(c, "_");
674    }
675
676    // Limit length to avoid issues. Slice by chars, not bytes: a multibyte
677    // task id whose 45th byte lands mid-code-point would panic on `&result[0..45]`.
678    if result.chars().count() > 50 {
679        use rand::RngExt;
680        let head: String = result.chars().take(45).collect();
681        result = format!("{head}-{}", rand::rng().random_range(1000..9999));
682    }
683
684    result
685}
686
687#[cfg(test)]
688mod tests {
689    #![allow(clippy::unwrap_used)]
690    use super::{
691        cap_file_to_tokens, read_text_for_context, relativize_glob, MAX_CONTEXT_READ_BYTES,
692    };
693
694    #[test]
695    fn read_text_for_context_small_file_reads_whole() {
696        let dir = tempfile::tempdir().unwrap();
697        let p = dir.path().join("small.txt");
698        std::fs::write(&p, "hello\nworld\n").unwrap();
699        assert_eq!(read_text_for_context(&p).unwrap(), "hello\nworld\n");
700    }
701
702    #[test]
703    fn read_text_for_context_binary_is_invalid_data() {
704        let dir = tempfile::tempdir().unwrap();
705        let p = dir.path().join("bin");
706        std::fs::write(&p, [0xff, 0xfe, 0x00, 0x01]).unwrap();
707        assert_eq!(read_text_for_context(&p).unwrap_err().kind(), std::io::ErrorKind::InvalidData);
708    }
709
710    #[test]
711    fn read_text_for_context_caps_oversized_file() {
712        let dir = tempfile::tempdir().unwrap();
713        let p = dir.path().join("big.txt");
714        std::fs::write(&p, vec![b'a'; MAX_CONTEXT_READ_BYTES as usize + 4096]).unwrap();
715        // Reads only the capped head, not the whole oversized file.
716        assert_eq!(read_text_for_context(&p).unwrap().len() as u64, MAX_CONTEXT_READ_BYTES);
717    }
718
719    #[test]
720    fn relativize_glob_makes_absolute_portable() {
721        // Absolute glob under the root -> relative.
722        assert_eq!(relativize_glob("/home/u/proj/src/*.rs", "/home/u/proj"), "src/*.rs");
723        // Already relative -> untouched.
724        assert_eq!(relativize_glob("src/*.rs", "/home/u/proj"), "src/*.rs");
725        // Outside the root -> kept verbatim.
726        assert_eq!(relativize_glob("/etc/*.conf", "/home/u/proj"), "/etc/*.conf");
727        // No project root -> unchanged.
728        assert_eq!(relativize_glob("/abs/x", ""), "/abs/x");
729    }
730
731    #[test]
732    fn small_file_is_untouched() {
733        let content = "def add(a, b):\n    return a + b\n";
734        assert_eq!(cap_file_to_tokens(content, 4_000), content);
735    }
736
737    #[test]
738    fn large_file_is_capped_with_a_note() {
739        // ~5k tokens of distinct words, capped to 100.
740        use std::fmt::Write as _;
741        let mut big = String::new();
742        for i in 0..5_000 {
743            let _ = write!(big, "word{i} ");
744        }
745        let capped = cap_file_to_tokens(&big, 100);
746        assert!(capped.len() < big.len(), "should shrink");
747        assert!(capped.contains("truncated for context save"), "should note truncation");
748        assert!(capped.starts_with("word0"), "should keep the head");
749    }
750}