lean_ctx_proxy/

tool_kind.rs

// SPDX-License-Identifier: Apache-2.0
//! Classifies what produced a `tool_result` so the proxy never lossy-compresses
//! a file/source-code read the model still needs (e.g. mid-refactor).
//!
//! The request body only carries the tool *result* plus an id linking it to the
//! originating tool *call*. We resolve that id → tool name from the assistant's
//! `tool_use` / `tool_calls` / `function_call` items, then map the name to a
//! [`ToolResultKind`]. A content heuristic ([`looks_like_source_code`]) is the
//! fallback for unknown/custom tools so a file read through a non-standard tool
//! is still protected.

use std::collections::HashMap;

use serde_json::Value;

/// What kind of tool produced a `tool_result`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ToolResultKind {
    /// A file/source read — must reach the model intact (it is what gets edited).
    FileRead,
    /// Shell/command output — safe to run through the pattern compressors.
    Shell,
    /// Search/listing output — safe to compress.
    Search,
    /// Unknown — fall back to the content heuristic before compressing.
    Other,
}

/// Maps a tool name (from any agent) to a [`ToolResultKind`].
///
/// Matching is case-insensitive and substring-based so vendor prefixes
/// (`mcp__fs__read_file`, `functions.read`) and casing variants are covered.
pub fn classify_tool_name(name: &str) -> ToolResultKind {
    let n = name.to_ascii_lowercase();

    // Order matters: a "read_file" must not be caught by a generic "file".
    const FILE_READ: &[&str] = &[
        "read_file",
        "readfile",
        "file_read",
        "fsread",
        "fs_read",
        "view_file",
        "viewfile",
        "open_file",
        "notebookread",
        "notebook_read",
        "cat_file",
        "get_file",
        "fetch_file",
        "ctx_read",
        "ctx_multi_read",
        "multi_read",
        "multiread",
        "read_many", // Gemini CLI `read_many_files`
        "read_files",
        "str_replace_editor", // view sub-mode returns file content
        "ctx_patch",          // #818: diff previews are code — never abbreviate
        "ctx_refactor",       // symbol edits return code diffs
        "ctx_callgraph",      // callers/callees contain source snippets
    ];
    if FILE_READ.iter().any(|k| n.contains(k)) {
        return ToolResultKind::FileRead;
    }
    // Bare "read"/"view"/"cat" as a whole token (Claude Code `Read`, Pi `read`).
    if matches!(n.as_str(), "read" | "view" | "cat" | "open") {
        return ToolResultKind::FileRead;
    }

    const SEARCH: &[&str] = &[
        "grep",
        "ripgrep",
        "search",
        "find",
        "glob",
        "list_dir",
        "listdir",
        "list_files",
        "listfiles",
        "ls",
        "codebase_search",
        "ctx_search",
        "ctx_tree",
    ];
    if SEARCH.iter().any(|k| n.contains(k)) {
        return ToolResultKind::Search;
    }

    const SHELL: &[&str] = &[
        "bash",
        "shell",
        "terminal",
        "run_command",
        "run_terminal",
        "runterminal",
        "execute_command",
        "exec_command",
        "command_exec",
        "ctx_shell",
    ];
    if SHELL.iter().any(|k| n.contains(k)) {
        return ToolResultKind::Shell;
    }
    if matches!(n.as_str(), "run" | "exec" | "execute" | "command" | "sh") {
        return ToolResultKind::Shell;
    }

    // Vendor-prefix fallback. Foreign harnesses namespace their tools
    // (`forge_read`, `pi.shell`, `fs:grep`), which the substring lists above
    // miss. Matching the name's path-like *segments* as whole words catches
    // those without the false positives a bare substring would cause
    // (`thread`, `research`, `already`). FileRead is checked first so a read is
    // never misclassified as compressible.
    for seg in n.split(|c: char| !c.is_ascii_alphanumeric()) {
        match seg {
            "read" | "view" | "cat" | "open" => return ToolResultKind::FileRead,
            "grep" | "search" | "find" | "glob" | "ls" | "rg" => return ToolResultKind::Search,
            "shell" | "bash" | "exec" | "run" | "terminal" | "cmd" => return ToolResultKind::Shell,
            _ => {}
        }
    }

    ToolResultKind::Other
}

/// Builds a `tool_use_id → tool_name` map from Anthropic `messages`.
///
/// Scans every assistant content block of `type:"tool_use"`.
pub fn anthropic_tool_names(messages: &[Value]) -> HashMap<String, String> {
    let mut map = HashMap::new();
    for msg in messages {
        let Some(blocks) = msg.get("content").and_then(|c| c.as_array()) else {
            continue;
        };
        for block in blocks {
            if block.get("type").and_then(|t| t.as_str()) != Some("tool_use") {
                continue;
            }
            if let (Some(id), Some(name)) = (
                block.get("id").and_then(|v| v.as_str()),
                block.get("name").and_then(|v| v.as_str()),
            ) {
                map.insert(id.to_string(), name.to_string());
            }
        }
    }
    map
}

/// Builds a `tool_call_id → function_name` map from OpenAI Chat Completions
/// `messages` (assistant `tool_calls[]`).
pub fn openai_tool_names(messages: &[Value]) -> HashMap<String, String> {
    let mut map = HashMap::new();
    for msg in messages {
        let Some(calls) = msg.get("tool_calls").and_then(|c| c.as_array()) else {
            continue;
        };
        for call in calls {
            let id = call.get("id").and_then(|v| v.as_str());
            let name = call
                .get("function")
                .and_then(|f| f.get("name"))
                .and_then(|v| v.as_str());
            if let (Some(id), Some(name)) = (id, name) {
                map.insert(id.to_string(), name.to_string());
            }
        }
    }
    map
}

/// Builds a `call_id → name` map from OpenAI Responses `input` items
/// (`type:"function_call"`).
pub fn responses_tool_names(input: &[Value]) -> HashMap<String, String> {
    let mut map = HashMap::new();
    for item in input {
        if item.get("type").and_then(|t| t.as_str()) != Some("function_call") {
            continue;
        }
        if let (Some(id), Some(name)) = (
            item.get("call_id").and_then(|v| v.as_str()),
            item.get("name").and_then(|v| v.as_str()),
        ) {
            map.insert(id.to_string(), name.to_string());
        }
    }
    map
}

/// Whether a `tool_result` with the given resolved kind and content must be
/// preserved intact (never lossy-compressed) by the proxy.
///
/// File reads are always protected; unknown tools are protected only when the
/// content heuristically looks like source code. Shell/search output is never
/// protected here — it flows through the normal pattern compressors.
pub fn should_protect(kind: ToolResultKind, content: &str) -> bool {
    match kind {
        ToolResultKind::FileRead => true,
        ToolResultKind::Other => looks_like_source_code(content),
        ToolResultKind::Shell | ToolResultKind::Search => false,
    }
}

/// Heuristic fallback: does this text look like source code (vs command output)?
///
/// Deliberately conservative — it only returns `true` when code signals clearly
/// dominate and shell/log signals are essentially absent, so genuine logs and
/// build output are still compressed. Used only when the tool name is unknown.
///
/// GH #628: the previous version under-counted real source — a decorative
/// separator comment (`// ————`, `// ====`) and the call-shaped scaffolding of a
/// test file (`describe(…) {`, `});`) scored as non-code, so a genuine source
/// read routed through an unrecognized tool was lossy-compressed and silently
/// lost those separator lines, breaking the model's subsequent exact-match edit.
/// The fix: comment lines are *neutral* (never dilute the ratio) and top-level
/// call/closer shapes count as code even at column 0. The shell-signal veto is
/// untouched, so genuine logs and build output are still compressed.
pub fn looks_like_source_code(content: &str) -> bool {
    let mut code_signals = 0usize;
    let mut shell_signals = 0usize;
    let mut considered = 0usize;

    for raw in content.lines().take(200) {
        let line = raw.trim_end();
        let trimmed = line.trim_start();
        if trimmed.is_empty() {
            continue;
        }

        // Comment lines are part of source but carry no compress-vs-keep signal
        // on their own, and a decorative separator (`// ————`, `// ====`) or a
        // doc-block continuation (` * @param`) must never *dilute* the code ratio
        // — that false-negative is what let the proxy strip those lines (#628).
        // Treat C-style comments as neutral: skip without counting. `#` is
        // deliberately excluded (ambiguous with shell prompts / log levels /
        // Python) so genuine shell output is still detected and compressed.
        if trimmed.starts_with("//")
            || trimmed.starts_with("/*")
            || trimmed.starts_with("*/")
            || trimmed.starts_with("* ")
        {
            continue;
        }

        considered += 1;

        // Command/log markers — strong evidence this is NOT a file read.
        if trimmed.starts_with("$ ")
            || trimmed.starts_with("% ")
            || trimmed.starts_with(">>> ")
            || trimmed.starts_with("warning:")
            || trimmed.starts_with("error:")
            || trimmed.starts_with("error[")
            || trimmed.starts_with("INFO ")
            || trimmed.starts_with("WARN ")
            || trimmed.starts_with("DEBUG ")
            || trimmed.starts_with("ERROR ")
            || trimmed.starts_with("Compiling ")
            || trimmed.starts_with("Downloaded ")
            || trimmed.starts_with("test result:")
        {
            shell_signals += 1;
            continue;
        }

        // Code markers.
        let is_indented = line.len() != trimmed.len();
        let has_code_punct = trimmed.ends_with('{')
            || trimmed.ends_with('}')
            || trimmed.ends_with(';')
            || trimmed.ends_with("=>")
            || trimmed.ends_with("->")
            || trimmed.ends_with(':');
        // Top-level declarations, call statements and block closers carry code
        // punctuation even at column 0 (`describe("x", () => {`, `});`), so the
        // bare `is_indented && has_code_punct` test missed them — the exact
        // test-DSL shape (`describe`/`it`/`expect`) behind the #628 false-negative.
        // A call/closer *shape* with code punctuation is a strong code signal
        // regardless of indentation.
        let is_call_or_closer = (trimmed.contains('(') && trimmed.contains(')'))
            || trimmed.starts_with('}')
            || trimmed.starts_with(')');
        let has_keyword = [
            "fn ",
            "def ",
            "class ",
            "import ",
            "from ",
            "function ",
            "func ",
            "pub ",
            "const ",
            "let ",
            "var ",
            "package ",
            "public ",
            "private ",
            "struct ",
            "enum ",
            "impl ",
            "#include",
            "return ",
            "async ",
            "export ",
        ]
        .iter()
        .any(|k| trimmed.starts_with(k) || trimmed.contains(k));

        // Code punctuation counts when the line is indented (a statement in a
        // block) OR is a top-level call/closer shape (`describe(…) {`, `});`).
        let has_code_shape = has_code_punct && (is_indented || is_call_or_closer);
        if has_code_shape || has_keyword {
            code_signals += 1;
        }
    }

    if considered < 5 || shell_signals > 0 {
        return false;
    }
    // Require a clear majority of code-shaped lines.
    code_signals * 2 >= considered
}

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

    #[test]
    fn classifies_file_read_tools() {
        for name in [
            "Read",
            "read_file",
            "view_file",
            "ctx_read",
            "mcp__fs__readFile",
            // Multi-file reads return file content and must be protected too.
            "ctx_multi_read",
            "read_many_files",
        ] {
            assert_eq!(
                classify_tool_name(name),
                ToolResultKind::FileRead,
                "{name} should be FileRead"
            );
        }
    }

    #[test]
    fn classifies_shell_and_search() {
        assert_eq!(classify_tool_name("Bash"), ToolResultKind::Shell);
        assert_eq!(
            classify_tool_name("run_terminal_cmd"),
            ToolResultKind::Shell
        );
        assert_eq!(classify_tool_name("Grep"), ToolResultKind::Search);
        assert_eq!(
            classify_tool_name("codebase_search"),
            ToolResultKind::Search
        );
    }

    #[test]
    fn unknown_tool_is_other() {
        assert_eq!(classify_tool_name("submit_pr"), ToolResultKind::Other);
    }

    #[test]
    fn classifies_vendor_prefixed_foreign_tools() {
        // Foreign harnesses (forge / pi) namespace their tools; the segment
        // fallback must still route them so source reads stay protected and
        // shell/search output stays compressible.
        assert_eq!(classify_tool_name("forge_read"), ToolResultKind::FileRead);
        assert_eq!(classify_tool_name("pi.read"), ToolResultKind::FileRead);
        assert_eq!(classify_tool_name("forge_shell"), ToolResultKind::Shell);
        assert_eq!(classify_tool_name("forge_exec"), ToolResultKind::Shell);
        assert_eq!(classify_tool_name("fs:grep"), ToolResultKind::Search);
    }

    #[test]
    fn segment_fallback_has_no_substring_false_positives() {
        // Whole-word segments only: "thread" contains "read", "spread" contains
        // "read" — neither may be misclassified as a file read.
        assert_eq!(classify_tool_name("thread_create"), ToolResultKind::Other);
        assert_eq!(classify_tool_name("spread_values"), ToolResultKind::Other);
        assert_eq!(
            classify_tool_name("readme_generator"),
            ToolResultKind::Other
        );
        assert_eq!(classify_tool_name("submit_pull"), ToolResultKind::Other);
    }

    #[test]
    fn anthropic_names_resolve_from_tool_use() {
        let messages = vec![
            serde_json::json!({
                "role": "assistant",
                "content": [
                    {"type": "text", "text": "reading"},
                    {"type": "tool_use", "id": "toolu_1", "name": "Read", "input": {}}
                ]
            }),
            serde_json::json!({
                "role": "user",
                "content": [{"type": "tool_result", "tool_use_id": "toolu_1", "content": "x"}]
            }),
        ];
        let names = anthropic_tool_names(&messages);
        assert_eq!(names.get("toolu_1").map(String::as_str), Some("Read"));
    }

    #[test]
    fn openai_names_resolve_from_tool_calls() {
        let messages = vec![serde_json::json!({
            "role": "assistant",
            "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "read_file"}}]
        })];
        let names = openai_tool_names(&messages);
        assert_eq!(names.get("call_1").map(String::as_str), Some("read_file"));
    }

    #[test]
    fn responses_names_resolve_from_function_call() {
        let input = vec![serde_json::json!({
            "type": "function_call", "call_id": "call_1", "name": "Read", "arguments": "{}"
        })];
        let names = responses_tool_names(&input);
        assert_eq!(names.get("call_1").map(String::as_str), Some("Read"));
    }

    #[test]
    fn source_code_detected() {
        let code = "pub fn build(cfg: &Config) -> Result<App> {\n    let mut app = App::new();\n    app.configure(cfg);\n    for route in cfg.routes() {\n        app.register(route);\n    }\n    Ok(app)\n}";
        assert!(looks_like_source_code(code));
    }

    #[test]
    fn command_output_not_code() {
        let log = "$ cargo build\n   Compiling foo v0.1.0\n   Compiling bar v0.2.0\nwarning: unused variable\n    Finished dev target\nerror: could not compile";
        assert!(!looks_like_source_code(log));
    }

    #[test]
    fn plain_prose_not_code() {
        let prose = "This is a normal paragraph of text.\nIt has several sentences.\nNone of them are code.\nThey are just words on lines.\nMore words follow here.";
        assert!(!looks_like_source_code(prose));
    }

    /// Regression for GH #628: a real test file whose only "non-code-shaped"
    /// lines are decorative separator comments must be recognized as source, so
    /// the proxy protects it (when routed through an unrecognized tool) instead
    /// of lossy-compressing it and silently dropping the `// ————` separators —
    /// the exact divergence that made the model's `ctx_edit` fail on a whitespace
    /// mismatch.
    #[test]
    fn test_file_with_separator_comments_is_source() {
        let code = "import { describe, it, expect } from \"vitest\";\n\
            \n\
            // ————————————————————————————————————————————————————————\n\
            // Section: arithmetic\n\
            // ————————————————————————————————————————————————————————\n\
            describe(\"add\", () => {\n\
            \x20 it(\"adds\", () => {\n\
            \x20   expect(1 + 1).toBe(2);\n\
            \x20 });\n\
            });\n\
            \n\
            // ----------------------------------------------------------\n\
            // Section: strings\n\
            // ----------------------------------------------------------\n\
            describe(\"concat\", () => {\n\
            \x20 it(\"joins\", () => {\n\
            \x20   expect(\"a\" + \"b\").toBe(\"ab\");\n\
            \x20 });\n\
            });\n";
        assert!(
            looks_like_source_code(code),
            "a .test.ts with decorative separator comments must read as source"
        );
        assert!(
            should_protect(ToolResultKind::Other, code),
            "an unrecognized tool returning this source must still be protected"
        );
    }

    /// A comment-heavy source file (license header, doc block) is still source:
    /// the neutral-comment rule must not let comments dilute the code ratio.
    #[test]
    fn comment_heavy_source_still_detected() {
        let code = "/*\n\
            \x20* Copyright (c) 2026. All rights reserved.\n\
            \x20* This module wires the request pipeline.\n\
            \x20*/\n\
            export function build(cfg) {\n\
            \x20 const app = create();\n\
            \x20 app.use(cfg);\n\
            \x20 return app;\n\
            }\n";
        assert!(looks_like_source_code(code));
    }

    /// The loosened call/closer signal must NOT start treating parenthesized log
    /// output as code — the shell-signal veto and the punctuation requirement keep
    /// genuine command output compressible.
    #[test]
    fn parenthesized_log_output_still_not_code() {
        let log = "INFO  starting worker (pid=4211)\n\
            processing batch (size=128) ok\n\
            processing batch (size=64) ok\n\
            WARN  slow response (842ms) from upstream\n\
            done in 3.2s (0 errors)\n";
        assert!(!looks_like_source_code(log));
    }

    /// #818: ctx_patch and ctx_refactor must be classified as FileRead
    /// so their diff previews are never abbreviation-compressed.
    #[test]
    fn ctx_patch_and_refactor_classified_as_file_read() {
        assert_eq!(classify_tool_name("ctx_patch"), ToolResultKind::FileRead);
        assert_eq!(classify_tool_name("ctx_refactor"), ToolResultKind::FileRead);
        assert_eq!(
            classify_tool_name("ctx_callgraph"),
            ToolResultKind::FileRead
        );
    }

    /// #818: diff preview with - /+ prefixes must be protected.
    #[test]
    fn diff_preview_is_protected_when_kind_is_file_read() {
        let diff = "--- src/main.rs\n\
            - fn check_all_segments(command: &str, allowlist: &[String]) -> Result<(), ShellError> {\n\
            -     if allowlist.is_empty() {\n\
            -         return Ok(());\n\
            + fn check_all_segments(cmd: &str, list: &[String]) -> Result<(), ShellError> {\n\
            +     if list.is_empty() {\n\
            +         return Ok(());\n";
        assert!(
            should_protect(ToolResultKind::FileRead, diff),
            "diff preview must be protected when kind is FileRead"
        );
    }
}