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//! # Rollout
27//!
28//! This strategy is intentionally kept out of the default-safe
29//! [`super::apply_all`] path because replacing historical arguments
30//! with back-references can erase still-relevant call details.
31
32use sha2::{Digest, Sha256};
33use std::collections::HashMap;
34
35use super::ExperimentalStats;
36use crate::provider::{ContentPart, Message};
37
38/// Keep tool-call arguments verbatim in this many trailing messages.
39pub const KEEP_LAST_MESSAGES: usize = 8;
40
41/// Tool-call arguments shorter than this are never collapsed.
42pub const MIN_COLLAPSE_BYTES: usize = 64;
43
44/// Collapse duplicate tool-call arguments into short back-references.
45///
46/// # Examples
47///
48/// ```rust
49/// use codetether_agent::provider::{ContentPart, Message, Role};
50/// use codetether_agent::session::helper::experimental::tool_call_dedup::{
51///     collapse_duplicate_calls, KEEP_LAST_MESSAGES,
52/// };
53///
54/// let args = serde_json::json!({
55///     "path": "src/very/deeply/nested/module.rs",
56///     "encoding": "utf-8",
57/// }).to_string();
58///
59/// let call = |id: &str| Message {
60///     role: Role::Assistant,
61///     content: vec![ContentPart::ToolCall {
62///         id: id.into(),
63///         name: "read_file".into(),
64///         arguments: args.clone(),
65///         thought_signature: None,
66///     }],
67/// };
68///
69/// let mut msgs = vec![call("first"), call("second")];
70/// for i in 0..KEEP_LAST_MESSAGES + 1 {
71///     msgs.push(Message {
72///         role: Role::User,
73///         content: vec![ContentPart::Text { text: format!("q{i}") }],
74///     });
75/// }
76///
77/// let stats = collapse_duplicate_calls(&mut msgs);
78/// assert!(stats.dedup_hits >= 1);
79///
80/// // Second call's arguments now reference the first.
81/// let ContentPart::ToolCall { arguments, .. } = &msgs[1].content[0] else {
82///     panic!();
83/// };
84/// assert!(arguments.starts_with("[REPEAT"));
85/// assert!(arguments.contains("first"));
86/// ```
87pub fn collapse_duplicate_calls(messages: &mut [Message]) -> ExperimentalStats {
88    let mut stats = ExperimentalStats::default();
89    let total = messages.len();
90    if total <= KEEP_LAST_MESSAGES {
91        return stats;
92    }
93    let eligible = total - KEEP_LAST_MESSAGES;
94    let mut seen: HashMap<[u8; 32], String> = HashMap::new();
95
96    for msg in messages[..eligible].iter_mut() {
97        for part in msg.content.iter_mut() {
98            let ContentPart::ToolCall {
99                id,
100                name,
101                arguments,
102                ..
103            } = part
104            else {
105                continue;
106            };
107            if arguments.len() < MIN_COLLAPSE_BYTES {
108                continue;
109            }
110            if arguments.starts_with("[REPEAT") {
111                continue;
112            }
113            let mut h = Sha256::new();
114            h.update(name.as_bytes());
115            h.update(b"\0");
116            h.update(arguments.as_bytes());
117            let key: [u8; 32] = h.finalize().into();
118            match seen.get(&key) {
119                Some(first_id) => {
120                    let marker =
121                        format!("[REPEAT of call_id={first_id}, {} bytes]", arguments.len());
122                    let saved = arguments.len().saturating_sub(marker.len());
123                    *arguments = marker;
124                    stats.dedup_hits += 1;
125                    stats.total_bytes_saved += saved;
126                }
127                None => {
128                    seen.insert(key, id.clone());
129                }
130            }
131        }
132    }
133
134    stats
135}
136
137#[cfg(test)]
138mod tests {
139    use super::*;
140    use crate::provider::Role;
141
142    fn call(id: &str, name: &str, args: &str) -> Message {
143        Message {
144            role: Role::Assistant,
145            content: vec![ContentPart::ToolCall {
146                id: id.into(),
147                name: name.into(),
148                arguments: args.into(),
149                thought_signature: None,
150            }],
151        }
152    }
153
154    #[test]
155    fn different_names_not_merged() {
156        let args = "a".repeat(200);
157        let mut msgs = vec![
158            call("1", "read_file", &args),
159            call("2", "write_file", &args),
160        ];
161        for _ in 0..KEEP_LAST_MESSAGES + 1 {
162            msgs.push(Message {
163                role: Role::User,
164                content: vec![ContentPart::Text { text: "q".into() }],
165            });
166        }
167        let stats = collapse_duplicate_calls(&mut msgs);
168        assert_eq!(stats.dedup_hits, 0);
169    }
170
171    #[test]
172    fn small_args_preserved() {
173        let mut msgs = vec![call("1", "noop", "{}"), call("2", "noop", "{}")];
174        for _ in 0..KEEP_LAST_MESSAGES + 1 {
175            msgs.push(Message {
176                role: Role::User,
177                content: vec![ContentPart::Text { text: "q".into() }],
178            });
179        }
180        let stats = collapse_duplicate_calls(&mut msgs);
181        assert_eq!(stats.dedup_hits, 0);
182    }
183
184    #[test]
185    fn recent_calls_preserved() {
186        let args = "z".repeat(200);
187        let mut msgs = vec![call("1", "f", &args)];
188        for _ in 0..KEEP_LAST_MESSAGES - 1 {
189            msgs.push(Message {
190                role: Role::User,
191                content: vec![ContentPart::Text { text: "q".into() }],
192            });
193        }
194        msgs.push(call("2", "f", &args));
195        let stats = collapse_duplicate_calls(&mut msgs);
196        assert_eq!(stats.dedup_hits, 0);
197    }
198}