zesdex/src/app/runtime/context/dedup.rs

#![allow(dead_code)]
//! Cross-call tool-result deduplication: when a read-only tool is called
//! again with identical arguments, the earlier result is replaced with a
//! placeholder so only the latest copy occupies context.
//!
//! Flow: pair each `Role::Tool` message to its originating `ToolCall` via
//! `tool_call_id` -> key on `(function.name, sha256(canonical_json(args)))`
//! -> for read-only tools, keep only the last occurrence of each key in
//! full, placeholder the rest.
//!
//! Why: reading the same file (or re-running the same grep) twice in a
//! session otherwise keeps both full copies in context until compaction
//! eventually drops the older one wholesale, along with everything else
//! from that period. Mutating tools (`write`, `edit`, `bash`, `delete`,
//! `git_operator`, ...) are never touched, even with identical
//! arguments, because call order and repetition can be semantically
//! meaningful (e.g. retrying a flaky `bash` command until it passes).
use crate::app::subagent::division::tool_scope::READ_TOOLS;
use crate::dto::chat::message::{ChatMessage, Role};
use sha2::Digest;
use std::collections::HashMap;

const DUPLICATE_PLACEHOLDER: &str =
    "[duplicate result — superseded by a later identical call, see below]";

/// Replace superseded read-only tool results with a placeholder.
///
/// Return: a `Vec<ChatMessage>` the same length as `messages`, and
/// `true` iff at least one entry was replaced. The caller uses the
/// `bool` to decide whether the result is worth persisting/announcing,
/// without `ChatMessage` needing to implement `PartialEq`.
pub fn collapse(messages: &[ChatMessage]) -> (Vec<ChatMessage>, bool) {
    // tool_call_id -> (tool name, canonical JSON of its arguments)
    let mut call_info: HashMap<String, (String, String)> = HashMap::new();
    for m in messages {
        if let Some(calls) = &m.tool_calls {
            for call in calls {
                let canonical = serde_json::to_string(&call.function.arguments).unwrap_or_default();
                call_info.insert(call.id.clone(), (call.function.name.clone(), canonical));
            }
        }
    }

    // For each (tool, args-hash) key among read-only tools, find the
    // index of its LAST occurrence — that's the one kept in full.
    let mut last_index_for_key: HashMap<String, usize> = HashMap::new();
    for (idx, m) in messages.iter().enumerate() {
        if m.role != Role::Tool {
            continue;
        }
        let Some(id) = &m.tool_call_id else { continue };
        let Some((name, args)) = call_info.get(id) else {
            continue;
        };
        if !READ_TOOLS.contains(&name.as_str()) {
            continue;
        }
        last_index_for_key.insert(dedup_key(name, args), idx);
    }

    let mut changed = false;
    let result = messages
        .iter()
        .enumerate()
        .map(|(idx, m)| {
            if m.role != Role::Tool {
                return m.clone();
            }
            let Some(id) = &m.tool_call_id else {
                return m.clone();
            };
            let Some((name, args)) = call_info.get(id) else {
                return m.clone();
            };
            if !READ_TOOLS.contains(&name.as_str()) {
                return m.clone();
            }
            let key = dedup_key(name, args);
            if last_index_for_key.get(&key) == Some(&idx) {
                return m.clone();
            }
            changed = true;
            ChatMessage::tool_result(id.clone(), DUPLICATE_PLACEHOLDER.to_string())
        })
        .collect();

    (result, changed)
}

/// Build the dedup key for a tool call.
///
/// Why hash the arguments: keeps the key a fixed, short size regardless
/// of argument payload size. `serde_json::to_string` is already
/// canonical here — this codebase doesn't enable `serde_json`'s
/// `preserve_order` feature, so `Value::Object` is backed by a
/// `BTreeMap` and always serializes keys in sorted order.
fn dedup_key(tool_name: &str, canonical_args: &str) -> String {
    let hash = hex::encode(sha2::Sha256::digest(canonical_args.as_bytes()));
    format!("{tool_name}:{hash}")
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::dto::chat::message::ChatMessage;
    use crate::dto::chat::tool::{ToolCall, ToolFunction};
    use serde_json::json;

    fn assistant_with_call(id: &str, name: &str, args: serde_json::Value) -> ChatMessage {
        let mut m = ChatMessage::assistant(None);
        m.tool_calls = Some(vec![ToolCall {
            id: id.to_string(),
            type_: "function".to_string(),
            function: ToolFunction {
                name: name.to_string(),
                arguments: args,
            },
        }]);
        m
    }

    #[test]
    fn older_result_of_same_read_tool_and_args_is_replaced() {
        let messages = vec![
            assistant_with_call("call-1", "read", json!({"path": "a.rs"})),
            ChatMessage::tool_result("call-1".to_string(), "first read of a.rs".to_string()),
            assistant_with_call("call-2", "read", json!({"path": "a.rs"})),
            ChatMessage::tool_result("call-2".to_string(), "second read of a.rs".to_string()),
        ];

        let (result, changed) = collapse(&messages);

        assert!(changed);
        assert_eq!(result[1].content.as_deref(), Some(DUPLICATE_PLACEHOLDER));
        assert_eq!(result[3].content.as_deref(), Some("second read of a.rs"));
    }

    #[test]
    fn different_arguments_are_not_deduplicated() {
        let messages = vec![
            assistant_with_call("call-1", "read", json!({"path": "a.rs"})),
            ChatMessage::tool_result("call-1".to_string(), "read of a.rs".to_string()),
            assistant_with_call("call-2", "read", json!({"path": "b.rs"})),
            ChatMessage::tool_result("call-2".to_string(), "read of b.rs".to_string()),
        ];

        let (result, changed) = collapse(&messages);

        assert!(!changed);
        assert_eq!(result[1].content.as_deref(), Some("read of a.rs"));
        assert_eq!(result[3].content.as_deref(), Some("read of b.rs"));
    }

    #[test]
    fn key_order_in_arguments_does_not_prevent_dedup() {
        let messages = vec![
            assistant_with_call("call-1", "grep", json!({"pattern": "foo", "path": "."})),
            ChatMessage::tool_result("call-1".to_string(), "first grep".to_string()),
            assistant_with_call("call-2", "grep", json!({"path": ".", "pattern": "foo"})),
            ChatMessage::tool_result("call-2".to_string(), "second grep".to_string()),
        ];

        let (result, changed) = collapse(&messages);

        assert!(changed);
        assert_eq!(result[1].content.as_deref(), Some(DUPLICATE_PLACEHOLDER));
    }

    #[test]
    fn mutating_tool_with_identical_args_is_never_deduplicated() {
        let messages = vec![
            assistant_with_call("call-1", "bash", json!({"command": "cargo test"})),
            ChatMessage::tool_result("call-1".to_string(), "first run: 3 failed".to_string()),
            assistant_with_call("call-2", "bash", json!({"command": "cargo test"})),
            ChatMessage::tool_result("call-2".to_string(), "second run: 0 failed".to_string()),
        ];

        let (result, changed) = collapse(&messages);

        assert!(!changed);
        assert_eq!(result[1].content.as_deref(), Some("first run: 3 failed"));
        assert_eq!(result[3].content.as_deref(), Some("second run: 0 failed"));
    }

    #[test]
    fn tool_result_with_no_matching_call_is_left_untouched() {
        let messages = vec![ChatMessage::tool_result(
            "orphan-id".to_string(),
            "some result".to_string(),
        )];

        let (result, changed) = collapse(&messages);

        assert!(!changed);
        assert_eq!(result[0].content.as_deref(), Some("some result"));
    }
}