fix(tools): preserve accurate truncation notices

This commit is contained in:
leokun
2026-09-04 16:12:31 +08:00
parent 3039eaff21
commit 0f564c0e7c
@@ -387,7 +387,15 @@ fn gate_mcp(tool: &mut pb::McpToolCall) {
"[truncated: MCP content items exceeded {MCP_CONTENT_ITEM_LIMIT} items; showing {MCP_CONTENT_ITEM_LIMIT} of {original_items} items]"
));
}
let original_text_bytes = success.content.iter().fold(0usize, |total, item| {
let bytes = match item.content.as_ref() {
Some(pb::mcp_tool_result_content_item::Content::Text(text)) => text.text.len(),
_ => 0,
};
total.saturating_add(bytes)
});
let mut remaining_text = MCP_TEXT_LIMIT;
let mut text_truncated = false;
let mut content = Vec::with_capacity(success.content.len() + notices.len());
for mut item in std::mem::take(&mut success.content) {
// MCP images are sent to the client as inline binary data. Truncating
@@ -399,19 +407,23 @@ fn gate_mcp(tool: &mut pb::McpToolCall) {
let original = text.text.clone();
let next = truncate_text("MCP content item", &original, MCP_TEXT_LIMIT);
if remaining_text == 0 {
notices.push(truncation_notice(
"MCP text",
MCP_TEXT_LIMIT,
MCP_TEXT_LIMIT,
MCP_TEXT_LIMIT.saturating_add(original.len()),
));
text_truncated |= !next.is_empty();
continue;
}
text.text = truncate_text("MCP text", &next, remaining_text);
text_truncated |= text.text != next;
remaining_text = remaining_text.saturating_sub(text.text.len());
}
content.push(item);
}
if text_truncated {
notices.push(truncation_notice(
"MCP text",
MCP_TEXT_LIMIT,
MCP_TEXT_LIMIT.saturating_sub(remaining_text),
original_text_bytes,
));
}
content.extend(notices.into_iter().map(mcp_notice));
success.content = content;
}
@@ -644,17 +656,40 @@ fn truncate_text(tool_name: &str, content: &str, limit: usize) -> String {
}
let kept = utf8_prefix(content, limit - notice.len());
// `notice.len()` grows with the digit count of `shown`, so `kept.len()`
// can alternate between two values across a power-of-ten boundary
// instead of reaching a fixed point. Settle on the first repeat; the
// reported count is then off by one at most and the result still fits.
// can alternate between two values across a power-of-ten boundary.
if kept.len() == shown || previous == Some(kept.len()) {
return format!("{}{notice}", kept.trim_end_matches('\n'));
return truncated_prefix_with_notice(tool_name, content, limit, original, kept.len());
}
previous = Some(shown);
shown = kept.len();
}
}
fn truncated_prefix_with_notice(
tool_name: &str,
content: &str,
limit: usize,
original: usize,
initial_cap: usize,
) -> String {
let mut cap = initial_cap;
loop {
let kept = utf8_prefix(content, cap).trim_end_matches('\n');
let notice = format!(
"\n\n[truncated: {tool_name} result exceeded {limit} bytes; showing {} of {original} bytes]",
kept.len()
);
if notice.len() >= limit {
return utf8_prefix(content, limit).to_string();
}
let available = limit - notice.len();
if kept.len() <= available {
return format!("{kept}{notice}");
}
cap = available;
}
}
fn truncate_edges(tool_name: &str, content: &str, limit: usize) -> String {
if content.len() <= limit {
return content.to_string();
@@ -732,6 +767,30 @@ mod tests {
})
}
fn mcp_tool(texts: Vec<String>) -> pb::McpToolCall {
pb::McpToolCall {
args: None,
result: Some(pb::McpToolResult {
result: Some(pb::mcp_tool_result::Result::Success(pb::McpSuccess {
content: texts
.into_iter()
.map(|text| pb::McpToolResultContentItem {
content: Some(pb::mcp_tool_result_content_item::Content::Text(
pb::McpTextContent {
text,
output_location: None,
},
)),
})
.collect(),
is_error: false,
structured_content: None,
})),
}),
description: None,
}
}
#[test]
fn truncate_text_never_exceeds_its_limit() {
let content = "b".repeat(200);
@@ -754,6 +813,32 @@ mod tests {
assert!(truncate_text("MCP text", &"b".repeat(500), 82).len() <= 82);
}
#[test]
fn truncate_text_reports_the_actual_utf8_prefix_size() {
let content = "😀".repeat(1_000);
let output = truncate_text("Grep", &content, 81);
let (kept, notice) = output
.split_once("\n\n[truncated:")
.expect("the truncation notice fits");
assert!(
notice.contains(&format!("showing {} of", kept.len())),
"notice must report the actual UTF-8 prefix size: {output}"
);
assert!(output.len() <= 81);
}
#[test]
fn mcp_total_text_truncation_always_adds_a_notice() {
let mut tool = mcp_tool(vec!["a".repeat(MCP_TEXT_LIMIT - 8), "b".repeat(100)]);
gate_mcp(&mut tool);
let success = match tool.result.unwrap().result.unwrap() {
pb::mcp_tool_result::Result::Success(success) => success,
_ => panic!("expected MCP success"),
};
assert_eq!(success.content.len(), 3);
assert!(is_mcp_notice(success.content.last().unwrap()));
}
#[test]
fn grep_content_gate_survives_a_nearly_exhausted_byte_budget() {
// 16 matches leave 16 bytes of the 32 KiB content budget, which is less