lean-ctx-server 3.12.4

// SPDX-License-Identifier: Apache-2.0
use std::path::Path;

use lean_ctx::core::tokens::count_tokens;

const SOURCE_EXTENSIONS: &[&str] = &[
    ".rs", ".ts", ".tsx", ".js", ".jsx", ".py", ".go", ".java", ".c", ".cpp", ".h", ".hpp", ".cs",
    ".rb", ".swift", ".kt", ".scala", ".zig", ".toml", ".yaml", ".yml", ".json", ".lock", ".sum",
    ".mod",
];

pub fn handle(project_root: &str, path: &str) -> String {
    if SOURCE_EXTENSIONS.iter().any(|ext| path.ends_with(ext)) {
        return format!(
            "ERROR: refusing to compress source file {path} — \
             ctx_compress_memory is for memory-bank/doc files only (#1323)"
        );
    }

    if let Err(e) = lean_ctx::core::pathjail::enforce_writable(Path::new(path)) {
        return format!("ERROR: {e}");
    }
    let root = Path::new(project_root);

    // The memory file itself may be a symlink a cloned repository shipped.
    // Refuse it before anything is read or backed up, so the bytes of a file
    // outside the project are never copied into the backup. The rewrite keeps
    // the file's own permission bits.
    let mode = match std::fs::symlink_metadata(path) {
        Ok(meta) if meta.file_type().is_symlink() => {
            return format!("ERROR: refusing to compress {path}: it is a symlink");
        }
        Ok(meta) => file_mode(&meta),
        Err(e) => return format!("ERROR: Cannot read {path}: {e}"),
    };

    // Admission gate (G01). This tool rewrites a file from its ORIGINAL bytes
    // and reports token counts only, so nothing admitted or redacted is ever
    // written: the backup receives the original, the file receives the
    // compaction of the original. What admission decides is whether the
    // request may derive from this source at all — a source the read path
    // withholds for the running request is neither compacted nor copied to a
    // backup. (Outside a request scope the read path is a plain local read,
    // as for every other reader.)
    if let Err(error) = crate::tools::ctx_read::read_file_lossy(path)
        && error.kind() == std::io::ErrorKind::PermissionDenied
    {
        return format!("ERROR: Cannot read {path}: {error}");
    }

    let content = match std::fs::read_to_string(path) {
        Ok(c) => c,
        Err(e) => return format!("ERROR: Cannot read {path}: {e}"),
    };

    let original_tokens = count_tokens(&content);

    // The backup name and the file itself can both be symlinks a cloned
    // repository shipped (a dangling `<name>.original.md`, or `<name>` itself),
    // and the jail leaves a dangling link unresolved; write both through the
    // confined write so a symlinked component is refused instead of following
    // it out of the project.
    let backup_path = build_backup_path(path);
    if !Path::new(&backup_path).exists()
        && let Err(e) = lean_ctx::core::pathjail_write::write_requested_below_with_mode(
            root,
            Path::new(&backup_path),
            Path::new(&backup_path),
            content.as_bytes(),
            mode,
        )
    {
        return format!("ERROR: Cannot create backup {backup_path}: {e}");
    }

    let compressed = compress_memory_file(&content);
    let compressed_tokens = count_tokens(&compressed);

    if let Err(e) = lean_ctx::core::pathjail_write::write_requested_below_with_mode(
        root,
        Path::new(path),
        Path::new(path),
        compressed.as_bytes(),
        mode,
    ) {
        return format!("ERROR: Cannot write compressed file: {e}");
    }

    let saved = original_tokens.saturating_sub(compressed_tokens);
    let pct = if original_tokens > 0 {
        (saved as f64 / original_tokens as f64 * 100.0).round() as usize
    } else {
        0
    };

    format!(
        "Compressed {}: {} → {} tokens ({saved} saved, {pct}%)\n\
         Backup: {backup_path}",
        Path::new(path)
            .file_name()
            .and_then(|f| f.to_str())
            .unwrap_or(path),
        original_tokens,
        compressed_tokens,
    )
}

/// Permission bits for the rewritten file and its backup: the original's own
/// bits, so compaction never widens or narrows who may read the file.
#[cfg(unix)]
fn file_mode(meta: &std::fs::Metadata) -> u32 {
    use std::os::unix::fs::PermissionsExt;
    meta.permissions().mode() & 0o777
}

#[cfg(not(unix))]
fn file_mode(_meta: &std::fs::Metadata) -> u32 {
    0o600
}

fn build_backup_path(path: &str) -> String {
    let p = Path::new(path);
    let stem = p.file_stem().and_then(|s| s.to_str()).unwrap_or("file");
    let parent = p.parent().unwrap_or_else(|| Path::new("."));
    parent
        .join(format!("{stem}.original.md"))
        .to_string_lossy()
        .to_string()
}

fn compress_memory_file(content: &str) -> String {
    let mut output = Vec::new();
    let mut in_code_block = false;
    let mut code_fence = String::new();

    for line in content.lines() {
        let trimmed = line.trim();

        if !in_code_block && is_code_fence_start(trimmed) {
            in_code_block = true;
            code_fence = trimmed
                .chars()
                .take_while(|c| *c == '`' || *c == '~')
                .collect();
            output.push(line.to_string());
            continue;
        }

        if in_code_block {
            output.push(line.to_string());
            if trimmed.starts_with(&code_fence) && trimmed.len() <= code_fence.len() + 1 {
                in_code_block = false;
                code_fence.clear();
            }
            continue;
        }

        if is_protected_line(trimmed) {
            output.push(line.to_string());
            continue;
        }

        if trimmed.is_empty() {
            if output.last().is_some_and(|l| l.trim().is_empty()) {
                continue;
            }
            output.push(String::new());
            continue;
        }

        let compressed = compress_prose_line(line);
        if !compressed.trim().is_empty() {
            output.push(compressed);
        }
    }

    output.join("\n")
}

fn is_code_fence_start(line: &str) -> bool {
    line.starts_with("```") || line.starts_with("~~~")
}

fn is_protected_line(line: &str) -> bool {
    if line.starts_with('#') {
        return true;
    }
    if line.starts_with("- ") || line.starts_with("* ") || line.starts_with("> ") {
        return true;
    }
    if line.starts_with('|') {
        return true;
    }
    if contains_url_or_path(line) && line.split_whitespace().count() <= 3 {
        return true;
    }
    false
}

fn contains_url_or_path(line: &str) -> bool {
    line.contains("http://")
        || line.contains("https://")
        || line.contains("ftp://")
        || (line.contains('/') && line.contains('.') && !line.contains(' '))
}

fn compress_prose_line(line: &str) -> String {
    let leading_ws: String = line.chars().take_while(|c| c.is_whitespace()).collect();
    let trimmed = line.trim();

    let mut words: Vec<&str> = trimmed.split_whitespace().collect();

    words.retain(|w| !is_filler_word(w));

    let mut result: Vec<String> = Vec::new();
    let mut i = 0;
    while i < words.len() {
        if let Some((replacement, skip)) = try_shorten_phrase(&words, i) {
            result.push(replacement.to_string());
            i += skip;
        } else {
            result.push(words[i].to_string());
            i += 1;
        }
    }

    format!("{}{}", leading_ws, result.join(" "))
}

fn is_filler_word(word: &str) -> bool {
    let w = word.to_lowercase();
    let w = w.trim_matches(|c: char| c.is_ascii_punctuation());
    matches!(
        w,
        "just" | "really" | "basically" | "actually" | "simply" | "please" | "very" | "quite"
    )
}

fn try_shorten_phrase(words: &[&str], pos: usize) -> Option<(&'static str, usize)> {
    if pos + 2 < words.len() {
        let three = format!(
            "{} {} {}",
            words[pos].to_lowercase(),
            words[pos + 1].to_lowercase(),
            words[pos + 2].to_lowercase()
        );
        match three.as_str() {
            "in order to" => return Some(("to", 3)),
            "as well as" => return Some(("and", 3)),
            "due to the" => return Some(("because", 3)),
            "make sure to" => return Some(("ensure", 3)),
            "a lot of" => return Some(("many", 3)),
            "on top of" => return Some(("besides", 3)),
            _ => {}
        }
    }

    if pos + 1 < words.len() {
        let two = format!(
            "{} {}",
            words[pos].to_lowercase(),
            words[pos + 1].to_lowercase()
        );
        match two.as_str() {
            "make sure" => return Some(("ensure", 2)),
            "a lot" => return Some(("many", 2)),
            "as well" => return Some(("also", 2)),
            "in order" => return Some(("to", 2)),
            "prior to" => return Some(("before", 2)),
            "due to" => return Some(("because", 2)),
            _ => {}
        }
    }

    None
}

#[cfg(test)]
mod tests {
    use super::*;

    /// A cloned repository can ship the backup name as a dangling symlink, or
    /// the memory file itself as a symlink; neither may carry the write out of
    /// the project.
    #[cfg(unix)]
    #[test]
    fn handle_refuses_symlinked_backup_and_target() {
        use std::os::unix::fs::symlink;
        let _iso = lean_ctx::core::data_dir::isolated_data_dir();
        let tmp = tempfile::tempdir().unwrap();
        let root = tmp.path().join("proj");
        std::fs::create_dir_all(&root).unwrap();
        let outside = tmp.path().join("outside");
        std::fs::create_dir_all(&outside).unwrap();

        // The backup name is a dangling link out of the project.
        let file = root.join("CLAUDE.md");
        std::fs::write(&file, "# T\n\n\n\nverbose prose here.\n").unwrap();
        symlink(outside.join("stolen"), root.join("CLAUDE.original.md")).unwrap();
        let out = handle(
            root.to_string_lossy().as_ref(),
            file.to_string_lossy().as_ref(),
        );
        assert!(out.starts_with("ERROR") && out.contains("symlink"), "{out}");
        assert!(!outside.join("stolen").exists());
        assert_eq!(
            std::fs::read_to_string(&file).unwrap(),
            "# T\n\n\n\nverbose prose here.\n"
        );

        // The memory file itself is a link to a file outside the project.
        let victim = outside.join("victim");
        std::fs::write(&victim, "keep").unwrap();
        let linked = root.join("NOTES.md");
        symlink(&victim, &linked).unwrap();
        let out = handle(
            root.to_string_lossy().as_ref(),
            linked.to_string_lossy().as_ref(),
        );
        assert!(out.starts_with("ERROR"), "{out}");
        assert_eq!(std::fs::read_to_string(&victim).unwrap(), "keep");
        // Nothing of the linked file is copied into a backup either.
        assert!(!root.join("NOTES.original.md").exists());

        // A plain memory file still compacts and gets its backup in the project.
        let plain = root.join("MEM.md");
        std::fs::write(&plain, "# M\n\n\n\nverbose prose here.\n").unwrap();
        let out = handle(
            root.to_string_lossy().as_ref(),
            plain.to_string_lossy().as_ref(),
        );
        assert!(out.starts_with("Compressed MEM.md"), "{out}");
        assert!(root.join("MEM.original.md").exists());
    }

    /// The confined write replaces the file; the result and its backup keep
    /// the original's permission bits (as a create applies them, through the
    /// umask) instead of a fixed 0600.
    #[cfg(unix)]
    #[test]
    fn handle_keeps_the_file_mode() {
        use std::os::unix::fs::{OpenOptionsExt, PermissionsExt};
        let _iso = lean_ctx::core::data_dir::isolated_data_dir();
        let tmp = tempfile::tempdir().unwrap();
        let root = tmp.path().join("proj");
        std::fs::create_dir_all(&root).unwrap();
        let mode = |p: &Path| std::fs::metadata(p).unwrap().permissions().mode() & 0o777;

        // 0o700 survives any umask that keeps the owner's bits, so a fixed
        // 0600 fails here whatever umask the test runs under.
        for (name, original) in [("AGENTS", 0o644), ("NOTES", 0o700)] {
            let probe = root.join(format!("{name}.probe"));
            std::fs::OpenOptions::new()
                .write(true)
                .create_new(true)
                .mode(original)
                .open(&probe)
                .unwrap();
            let expected = mode(&probe);

            let file = root.join(format!("{name}.md"));
            std::fs::write(&file, "# A\n\n\n\nverbose prose here.\n").unwrap();
            std::fs::set_permissions(&file, std::fs::Permissions::from_mode(original)).unwrap();
            let out = handle(
                root.to_string_lossy().as_ref(),
                file.to_string_lossy().as_ref(),
            );
            assert!(out.starts_with(&format!("Compressed {name}.md")), "{out}");
            assert_eq!(mode(&file), expected, "{name}.md");
            assert_eq!(
                mode(&root.join(format!("{name}.original.md"))),
                expected,
                "{name}.original.md"
            );
        }
    }

    #[test]
    fn preserves_code_blocks() {
        let input = "Some text just really here.\n\n```rust\nfn main() {\n    println!(\"hello\");\n}\n```\n\nMore text.";
        let result = compress_memory_file(input);
        assert!(result.contains("fn main()"));
        assert!(result.contains("println!"));
    }

    #[test]
    fn preserves_headings() {
        let input = "# Main Heading\n\nJust some filler text here.\n\n## Sub Heading";
        let result = compress_memory_file(input);
        assert!(result.contains("# Main Heading"));
        assert!(result.contains("## Sub Heading"));
    }

    #[test]
    fn preserves_urls() {
        let input = "Visit https://example.com for details.\nJust some really basic text.";
        let result = compress_memory_file(input);
        assert!(result.contains("https://example.com"));
    }

    #[test]
    fn removes_filler_words() {
        let input = "You should just really basically make sure to check this.";
        let result = compress_prose_line(input);
        assert!(!result.contains("just"));
        assert!(!result.contains("really"));
        assert!(!result.contains("basically"));
        assert!(result.contains("ensure"));
    }

    #[test]
    fn shortens_phrases() {
        let input = "In order to fix this, make sure to check the config.";
        let result = compress_prose_line(input);
        assert!(!result.contains("In order to"));
        assert!(result.contains("to"));
        assert!(result.contains("ensure"));
    }

    #[test]
    fn collapses_blank_lines() {
        let input = "Line 1\n\n\n\nLine 2\n\n\nLine 3";
        let result = compress_memory_file(input);
        assert!(!result.contains("\n\n\n"));
    }

    #[test]
    fn preserves_tables() {
        let input = "| Col A | Col B |\n|-------|-------|\n| val1  | val2  |";
        let result = compress_memory_file(input);
        assert!(result.contains("| Col A | Col B |"));
    }

    #[test]
    fn backup_path_computed_correctly() {
        assert_eq!(
            Path::new(&build_backup_path("/home/user/.cursorrules")),
            Path::new("/home/user")
                .join(".cursorrules.original.md")
                .as_path()
        );
        assert_eq!(
            Path::new(&build_backup_path("/project/CLAUDE.md")),
            Path::new("/project").join("CLAUDE.original.md").as_path()
        );
    }

    /// #475: in-place memory compaction must refuse to rewrite a file inside a
    /// read-only root and must not drop a `*.original.md` backup there either.
    #[cfg(not(feature = "no-jail"))]
    #[test]
    fn handle_denies_write_into_read_only_root() {
        let _iso = lean_ctx::core::data_dir::isolated_data_dir();
        let dir = tempfile::tempdir().unwrap();
        let ro = dir.path().join("refrepo");
        std::fs::create_dir_all(&ro).unwrap();
        let file = ro.join("CLAUDE.md");
        let original = "# Title\n\nJust some really verbose prose here.\n";
        std::fs::write(&file, original).unwrap();

        let ro_canon = lean_ctx::core::pathjail::canonicalize_or_self(&ro);
        lean_ctx::test_env::set_var(
            "LEAN_CTX_READ_ONLY_ROOTS",
            ro_canon.to_string_lossy().as_ref(),
        );
        let out = handle(
            ro.to_string_lossy().as_ref(),
            file.to_string_lossy().as_ref(),
        );
        lean_ctx::test_env::remove_var("LEAN_CTX_READ_ONLY_ROOTS");

        assert!(
            out.starts_with("ERROR") && out.contains("read-only"),
            "compaction into a read-only root must be refused: {out}"
        );
        assert_eq!(
            std::fs::read_to_string(&file).unwrap(),
            original,
            "the file must be untouched"
        );
        assert!(
            !ro.join("CLAUDE.original.md").exists(),
            "no backup may be written into a read-only root"
        );
    }

    /// G01: a memory file the request may not read was compacted in place and
    /// copied to a `*.original.md` backup. It is now left untouched; a file
    /// the request may read is compacted from its original bytes as before.
    #[test]
    fn a_withheld_memory_file_is_neither_compacted_nor_backed_up() {
        let _iso = lean_ctx::core::data_dir::isolated_data_dir();
        let dir = tempfile::tempdir().unwrap();
        let root = dir.path().to_string_lossy().into_owned();
        let in_request = |path: &std::path::Path| {
            crate::tools::admitted_sources::as_agent_acquisition(&root, || {
                handle(&root, path.to_string_lossy().as_ref())
            })
        };
        let open = dir.path().join("NOTES.md");
        std::fs::write(&open, "# Title\n\n\n\nJust some prose here.\n").unwrap();
        {
            let _community = lean_ctx::core::policy::runtime::TestPolicyOverride::set(None);
            let out = in_request(&open);
            assert!(out.starts_with("Compressed NOTES.md: "), "{out}");
            assert!(dir.path().join("NOTES.original.md").exists());
        }
        let policy = lean_ctx::core::policy::load(
            "name = \"ctx-compress-memory-test\"\nversion = \"1.0.0\"\ndescription = \"test\"\n\
             [filters]\nclassification = \"block\"\nblocked_labels = [\"CONFIDENTIAL\"]\n",
        )
        .unwrap();
        let _policy = lean_ctx::core::policy::runtime::TestPolicyOverride::set(Some(policy));
        let plan = dir.path().join("PLAN.md");
        let original = "CONFIDENTIAL\n\n\n\nroadmap-marker for the launch\n";
        std::fs::write(&plan, original).unwrap();
        let out = in_request(&plan);
        assert!(out.starts_with("ERROR: Cannot read "), "{out}");
        assert!(!out.contains("roadmap-marker"), "{out}");
        assert_eq!(std::fs::read_to_string(&plan).unwrap(), original);
        assert!(!dir.path().join("PLAN.original.md").exists());
    }
}