telos_agent/agent/compaction/
message_truncation.rs

1//! Message-level truncation: cap oversized text, thinking, and tool-result fields.
2//!
3//! This layer shrinks individual messages without changing conversation shape.
4
5use std::{fmt::Debug, sync::Arc};
6
7use serde_json::Value;
8
9use crate::model::message::{ContentBlock, Message, Role};
10
11// ── Compression strategy ─────────────────────────────────────────────────────
12
13/// Pluggable strategy for compressing a single content string.
14pub trait ContentCompressor: Debug + Send + Sync {
15    /// Compress `content` so its JSON-serialised form fits within `target_bytes`.
16    fn compress(&self, content: &str, target_bytes: usize) -> String;
17}
18
19/// Brute-force Unicode truncation — keeps the longest prefix whose serialised
20/// JSON form fits in `target_bytes`.
21#[derive(Debug, Default)]
22pub struct TruncationCompressor;
23
24impl ContentCompressor for TruncationCompressor {
25    fn compress(&self, input: &str, target_bytes: usize) -> String {
26        if serde_json::to_string(input).unwrap().len() <= target_bytes {
27            return input.to_string();
28        }
29
30        let lo = binary_search_largest_fitting(input.chars().count(), |mid| {
31            let candidate: String = input.chars().take(mid).collect();
32            serde_json::to_string(&candidate).unwrap().len() <= target_bytes
33        });
34        input.chars().take(lo).collect()
35    }
36}
37
38// ── Config ───────────────────────────────────────────────────────────────────
39
40/// Knobs for message-level truncation. `None` disables the corresponding budget.
41#[derive(Debug, Clone, Default)]
42pub struct MessageTruncationConfig {
43    /// Maximum serialised bytes for a single content field (text / tool result / thinking).
44    pub max_block_content_bytes: Option<usize>,
45    /// Maximum serialised bytes for the entire message.
46    pub max_message_bytes: Option<usize>,
47    /// Compression strategy — defaults to [`TruncationCompressor`] when `None`.
48    pub compressor: Option<Arc<dyn ContentCompressor>>,
49}
50
51// ── Result ───────────────────────────────────────────────────────────────────
52
53/// Outcome of [`truncate_message`].
54#[derive(Debug, Clone)]
55pub struct MessageTruncationResult {
56    pub message: Message,
57    pub compacted: bool,
58}
59
60// ── Public API ───────────────────────────────────────────────────────────────
61
62/// Compress content fields inside message blocks so the total serialised size
63/// fits within the budget. Delegates to the configured [`ContentCompressor`].
64pub fn truncate_message(
65    message: Message,
66    config: &MessageTruncationConfig,
67) -> MessageTruncationResult {
68    if config.max_block_content_bytes.is_none() && config.max_message_bytes.is_none() {
69        return MessageTruncationResult { message, compacted: false };
70    }
71
72    let compressor: &dyn ContentCompressor =
73        config.compressor.as_deref().unwrap_or(&TruncationCompressor);
74    let Message { role, mut blocks } = message;
75    let mut compacted = false;
76
77    if let Some(max_block_content_bytes) = config.max_block_content_bytes {
78        compacted |= compress_blocks_over_cap(&mut blocks, max_block_content_bytes, compressor);
79    }
80
81    if let Some(max_message_bytes) = config.max_message_bytes
82        && message_serialized_len(role, &blocks) > max_message_bytes
83    {
84        let cap = largest_block_cap_for_message(role, &blocks, max_message_bytes, compressor);
85        compacted |= compress_blocks_over_cap(&mut blocks, cap, compressor);
86    }
87
88    MessageTruncationResult { message: Message { role, blocks }, compacted }
89}
90
91// ── Helpers ──────────────────────────────────────────────────────────────────
92
93/// Serialised byte length of the truncatable content inside a block.
94pub(crate) fn content_serialized_len(block: &ContentBlock) -> Option<usize> {
95    match block {
96        ContentBlock::Text(t) => Some(serde_json::to_string(&t.text).unwrap().len()),
97        ContentBlock::ToolResult(r) => Some(r.content.to_string().len()),
98        ContentBlock::Thinking(t) => Some(serde_json::to_string(&t.text).unwrap().len()),
99        ContentBlock::ToolCall(_) => None,
100    }
101}
102
103fn largest_block_cap_for_message(
104    role: Role,
105    blocks: &[ContentBlock],
106    max_message_bytes: usize,
107    compressor: &dyn ContentCompressor,
108) -> usize {
109    let upper_bound = blocks.iter().filter_map(content_serialized_len).max().unwrap_or(0);
110
111    binary_search_largest_fitting(upper_bound, |cap| {
112        let mut capped_blocks = blocks.to_vec();
113        compress_blocks_over_cap(&mut capped_blocks, cap, compressor);
114        message_serialized_len(role, &capped_blocks) <= max_message_bytes
115    })
116}
117
118/// Return the largest value in `0..=upper_bound` accepted by a monotonic predicate.
119fn binary_search_largest_fitting(
120    mut upper_bound: usize,
121    mut fits: impl FnMut(usize) -> bool,
122) -> usize {
123    let mut lo = 0usize;
124
125    while lo < upper_bound {
126        let mid = (lo + upper_bound).div_ceil(2);
127        if fits(mid) {
128            lo = mid;
129        } else {
130            upper_bound = mid - 1;
131        }
132    }
133
134    lo
135}
136
137fn message_serialized_len(role: Role, blocks: &[ContentBlock]) -> usize {
138    serde_json::to_string(&Message { role, blocks: blocks.to_vec() }).unwrap().len()
139}
140
141fn compress_blocks_over_cap(
142    blocks: &mut [ContentBlock],
143    max_content_bytes: usize,
144    compressor: &dyn ContentCompressor,
145) -> bool {
146    let mut compacted = false;
147
148    for block in blocks {
149        if content_serialized_len(block).is_some_and(|bytes| bytes > max_content_bytes) {
150            truncate_block_content(block, max_content_bytes, compressor);
151            compacted = true;
152        }
153    }
154
155    compacted
156}
157
158const TRUNCATION_MARKER: &str = "...[truncated]";
159
160/// Compress the block's content field to fit within `target_bytes`.
161pub(crate) fn truncate_block_content(
162    block: &mut ContentBlock,
163    target_bytes: usize,
164    compressor: &dyn ContentCompressor,
165) {
166    match block {
167        ContentBlock::Text(t) => {
168            t.text = compress_with_marker(&t.text, target_bytes, compressor);
169        }
170        ContentBlock::ToolResult(r) => {
171            let s = r.content.to_string();
172            r.content = Value::String(compress_with_marker(&s, target_bytes, compressor));
173        }
174        ContentBlock::Thinking(t) => {
175            t.text = compress_with_marker(&t.text, target_bytes, compressor);
176        }
177        ContentBlock::ToolCall(_) => {}
178    }
179}
180
181fn compress_with_marker(
182    content: &str,
183    target_bytes: usize,
184    compressor: &dyn ContentCompressor,
185) -> String {
186    let marker_budget = serde_json::to_string(TRUNCATION_MARKER).unwrap().len() - 2;
187    if target_bytes <= marker_budget + 2 {
188        return compressor.compress(content, target_bytes);
189    }
190
191    let compressed = compressor.compress(content, target_bytes - marker_budget);
192    format!("{compressed}{TRUNCATION_MARKER}")
193}
194
195// ── Tests ────────────────────────────────────────────────────────────────────
196
197#[cfg(test)]
198mod tests {
199    use super::*;
200    use crate::model::message::*;
201
202    fn tool_result(text: &str) -> ToolResult {
203        ToolResult {
204            tool_call_id: "1".into(),
205            name: "t".into(),
206            is_error: false,
207            content: serde_json::json!({"data": text}),
208        }
209    }
210
211    fn config(block_cap: Option<usize>, msg_cap: Option<usize>) -> MessageTruncationConfig {
212        MessageTruncationConfig {
213            max_block_content_bytes: block_cap,
214            max_message_bytes: msg_cap,
215            compressor: None,
216        }
217    }
218
219    #[test]
220    fn no_op_when_budget_disabled() {
221        let msg = Message::tool_results(vec![tool_result(&"x".repeat(10_000))]);
222        let r = truncate_message(msg.clone(), &MessageTruncationConfig::default());
223        assert!(!r.compacted);
224        assert_eq!(r.message, msg);
225    }
226
227    #[test]
228    fn per_block_cap_truncates_oversized_tool_result() {
229        let msg = Message::tool_results(vec![tool_result(&"x".repeat(10_000))]);
230        let r = truncate_message(msg, &config(Some(100), None));
231        assert!(r.compacted);
232        let block_len = content_serialized_len(&r.message.blocks[0]).unwrap();
233        assert!(block_len <= 100, "len={}", block_len);
234    }
235
236    #[test]
237    fn default_compressor_respects_json_serialized_length() {
238        let input = "\\".repeat(100);
239        let compressed = TruncationCompressor.compress(&input, 20);
240        let serialized_len = serde_json::to_string(&compressed).unwrap().len();
241        assert!(serialized_len <= 20, "len={}", serialized_len);
242    }
243
244    #[test]
245    fn aggregate_budget_truncates_multiple_blocks() {
246        let msg = Message::tool_results(vec![
247            tool_result(&"x".repeat(500)),
248            tool_result(&"y".repeat(500)),
249        ]);
250        let r = truncate_message(msg, &config(None, Some(300)));
251        assert!(r.compacted);
252        let s = serde_json::to_string(&r.message).unwrap();
253        assert!(s.len() < 600, "len={}", s.len());
254    }
255
256    #[test]
257    fn both_budgets_fit_message_budget() {
258        let msg = Message::tool_results(vec![
259            tool_result(&"x".repeat(10_000)),
260            tool_result(&"y".repeat(10_000)),
261        ]);
262        let r = truncate_message(msg, &config(Some(200), Some(300)));
263        assert!(r.compacted);
264        let s = serde_json::to_string(&r.message).unwrap();
265        assert!(s.len() < 500, "len={}", s.len());
266    }
267
268    #[test]
269    fn leaves_tool_calls_untouched() {
270        let msg = Message {
271            role: Role::Assistant,
272            blocks: vec![
273                ContentBlock::Text(TextBlock { text: "x".repeat(10_000) }),
274                ContentBlock::ToolCall(ToolCall {
275                    id: "call_1".into(),
276                    name: "read".into(),
277                    arguments: serde_json::json!({"path": "/etc/passwd"}),
278                }),
279            ],
280        };
281        let r = truncate_message(msg, &config(Some(50), None));
282        assert!(r.compacted);
283        assert_eq!(r.message.tool_calls().count(), 1);
284    }
285
286    #[test]
287    fn truncated_content_has_marker() {
288        let msg = Message::user("hello world this is a long message".repeat(10));
289        let r = truncate_message(msg, &config(Some(30), None));
290        assert!(r.compacted);
291        assert!(r.message.text_content().contains("[truncated]"));
292    }
293}
telos_agent/agent/compaction/message_truncation.rs

telos_agent/agent/compaction/
message_truncation.rs
