Skip to main content

codetether_agent/session/helper/experimental/
tool_call_dedup.rs

1//! Collapse redundant identical tool calls.
2//!
3//! When an agent re-issues the same tool call with the same arguments
4//! (e.g. `read_file { path: "src/lib.rs" }` called at step 3 and again
5//! at step 17), the *call* itself is wasted tokens — the earlier
6//! response is already in history (and by this point usually already
7//! deduplicated by [`super::dedup`]).
8//!
9//! This pass walks the buffer, hashes each `ContentPart::ToolCall` by
10//! `(name, arguments)`, and replaces the *arguments* of later
11//! duplicates with a short `[REPEAT of call_X]` marker. The tool's
12//! result block for the repeat is left untouched so the model still
13//! sees a consistent call→result pairing; dedup will have collapsed
14//! the result body if the output was identical.
15//!
16//! # Safety
17//!
18//! * Only the **arguments** string is rewritten. Name, id, and
19//!   thought_signature are preserved so every provider adapter still
20//!   round-trips.
21//! * Tool calls in the most recent [`KEEP_LAST_MESSAGES`] messages are
22//!   never touched — the model is likely actively reasoning over them.
23//! * Only triggers when arguments exceed [`MIN_COLLAPSE_BYTES`] so
24//!   micro-arguments like `{}` stay readable.
25//!
26//! # Always-on
27//!
28//! No config flag.
29
30use sha2::{Digest, Sha256};
31use std::collections::HashMap;
32
33use super::ExperimentalStats;
34use crate::provider::{ContentPart, Message};
35
36/// Keep tool-call arguments verbatim in this many trailing messages.
37pub const KEEP_LAST_MESSAGES: usize = 8;
38
39/// Tool-call arguments shorter than this are never collapsed.
40pub const MIN_COLLAPSE_BYTES: usize = 64;
41
42/// Collapse duplicate tool-call arguments into short back-references.
43///
44/// # Examples
45///
46/// ```rust
47/// use codetether_agent::provider::{ContentPart, Message, Role};
48/// use codetether_agent::session::helper::experimental::tool_call_dedup::{
49///     collapse_duplicate_calls, KEEP_LAST_MESSAGES,
50/// };
51///
52/// let args = serde_json::json!({
53///     "path": "src/very/deeply/nested/module.rs",
54///     "encoding": "utf-8",
55/// }).to_string();
56///
57/// let call = |id: &str| Message {
58///     role: Role::Assistant,
59///     content: vec![ContentPart::ToolCall {
60///         id: id.into(),
61///         name: "read_file".into(),
62///         arguments: args.clone(),
63///         thought_signature: None,
64///     }],
65/// };
66///
67/// let mut msgs = vec![call("first"), call("second")];
68/// for i in 0..KEEP_LAST_MESSAGES + 1 {
69///     msgs.push(Message {
70///         role: Role::User,
71///         content: vec![ContentPart::Text { text: format!("q{i}") }],
72///     });
73/// }
74///
75/// let stats = collapse_duplicate_calls(&mut msgs);
76/// assert!(stats.dedup_hits >= 1);
77///
78/// // Second call's arguments now reference the first.
79/// let ContentPart::ToolCall { arguments, .. } = &msgs[1].content[0] else {
80///     panic!();
81/// };
82/// assert!(arguments.starts_with("[REPEAT"));
83/// assert!(arguments.contains("first"));
84/// ```
85pub fn collapse_duplicate_calls(messages: &mut [Message]) -> ExperimentalStats {
86    let mut stats = ExperimentalStats::default();
87    let total = messages.len();
88    if total <= KEEP_LAST_MESSAGES {
89        return stats;
90    }
91    let eligible = total - KEEP_LAST_MESSAGES;
92    let mut seen: HashMap<[u8; 32], String> = HashMap::new();
93
94    for msg in messages[..eligible].iter_mut() {
95        for part in msg.content.iter_mut() {
96            let ContentPart::ToolCall {
97                id,
98                name,
99                arguments,
100                ..
101            } = part
102            else {
103                continue;
104            };
105            if arguments.len() < MIN_COLLAPSE_BYTES {
106                continue;
107            }
108            if arguments.starts_with("[REPEAT") {
109                continue;
110            }
111            let mut h = Sha256::new();
112            h.update(name.as_bytes());
113            h.update(b"\0");
114            h.update(arguments.as_bytes());
115            let key: [u8; 32] = h.finalize().into();
116            match seen.get(&key) {
117                Some(first_id) => {
118                    let marker =
119                        format!("[REPEAT of call_id={first_id}, {} bytes]", arguments.len());
120                    let saved = arguments.len().saturating_sub(marker.len());
121                    *arguments = marker;
122                    stats.dedup_hits += 1;
123                    stats.total_bytes_saved += saved;
124                }
125                None => {
126                    seen.insert(key, id.clone());
127                }
128            }
129        }
130    }
131
132    stats
133}
134
135#[cfg(test)]
136mod tests {
137    use super::*;
138    use crate::provider::Role;
139
140    fn call(id: &str, name: &str, args: &str) -> Message {
141        Message {
142            role: Role::Assistant,
143            content: vec![ContentPart::ToolCall {
144                id: id.into(),
145                name: name.into(),
146                arguments: args.into(),
147                thought_signature: None,
148            }],
149        }
150    }
151
152    #[test]
153    fn different_names_not_merged() {
154        let args = "a".repeat(200);
155        let mut msgs = vec![
156            call("1", "read_file", &args),
157            call("2", "write_file", &args),
158        ];
159        for _ in 0..KEEP_LAST_MESSAGES + 1 {
160            msgs.push(Message {
161                role: Role::User,
162                content: vec![ContentPart::Text { text: "q".into() }],
163            });
164        }
165        let stats = collapse_duplicate_calls(&mut msgs);
166        assert_eq!(stats.dedup_hits, 0);
167    }
168
169    #[test]
170    fn small_args_preserved() {
171        let mut msgs = vec![call("1", "noop", "{}"), call("2", "noop", "{}")];
172        for _ in 0..KEEP_LAST_MESSAGES + 1 {
173            msgs.push(Message {
174                role: Role::User,
175                content: vec![ContentPart::Text { text: "q".into() }],
176            });
177        }
178        let stats = collapse_duplicate_calls(&mut msgs);
179        assert_eq!(stats.dedup_hits, 0);
180    }
181
182    #[test]
183    fn recent_calls_preserved() {
184        let args = "z".repeat(200);
185        let mut msgs = vec![call("1", "f", &args)];
186        for _ in 0..KEEP_LAST_MESSAGES - 1 {
187            msgs.push(Message {
188                role: Role::User,
189                content: vec![ContentPart::Text { text: "q".into() }],
190            });
191        }
192        msgs.push(call("2", "f", &args));
193        let stats = collapse_duplicate_calls(&mut msgs);
194        assert_eq!(stats.dedup_hits, 0);
195    }
196}