mirror of
https://wget.la/https://github.com/leookun/cursor-byok
synced 2026-10-08 07:21:13 +08:00
fix(tools): preserve accurate truncation notices
This commit is contained in:
@@ -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]"
|
"[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 remaining_text = MCP_TEXT_LIMIT;
|
||||||
|
let mut text_truncated = false;
|
||||||
let mut content = Vec::with_capacity(success.content.len() + notices.len());
|
let mut content = Vec::with_capacity(success.content.len() + notices.len());
|
||||||
for mut item in std::mem::take(&mut success.content) {
|
for mut item in std::mem::take(&mut success.content) {
|
||||||
// MCP images are sent to the client as inline binary data. Truncating
|
// 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 original = text.text.clone();
|
||||||
let next = truncate_text("MCP content item", &original, MCP_TEXT_LIMIT);
|
let next = truncate_text("MCP content item", &original, MCP_TEXT_LIMIT);
|
||||||
if remaining_text == 0 {
|
if remaining_text == 0 {
|
||||||
notices.push(truncation_notice(
|
text_truncated |= !next.is_empty();
|
||||||
"MCP text",
|
|
||||||
MCP_TEXT_LIMIT,
|
|
||||||
MCP_TEXT_LIMIT,
|
|
||||||
MCP_TEXT_LIMIT.saturating_add(original.len()),
|
|
||||||
));
|
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
text.text = truncate_text("MCP text", &next, remaining_text);
|
text.text = truncate_text("MCP text", &next, remaining_text);
|
||||||
|
text_truncated |= text.text != next;
|
||||||
remaining_text = remaining_text.saturating_sub(text.text.len());
|
remaining_text = remaining_text.saturating_sub(text.text.len());
|
||||||
}
|
}
|
||||||
content.push(item);
|
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));
|
content.extend(notices.into_iter().map(mcp_notice));
|
||||||
success.content = content;
|
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());
|
let kept = utf8_prefix(content, limit - notice.len());
|
||||||
// `notice.len()` grows with the digit count of `shown`, so `kept.len()`
|
// `notice.len()` grows with the digit count of `shown`, so `kept.len()`
|
||||||
// can alternate between two values across a power-of-ten boundary
|
// 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.
|
|
||||||
if kept.len() == shown || previous == Some(kept.len()) {
|
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);
|
previous = Some(shown);
|
||||||
shown = kept.len();
|
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 {
|
fn truncate_edges(tool_name: &str, content: &str, limit: usize) -> String {
|
||||||
if content.len() <= limit {
|
if content.len() <= limit {
|
||||||
return content.to_string();
|
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]
|
#[test]
|
||||||
fn truncate_text_never_exceeds_its_limit() {
|
fn truncate_text_never_exceeds_its_limit() {
|
||||||
let content = "b".repeat(200);
|
let content = "b".repeat(200);
|
||||||
@@ -754,6 +813,32 @@ mod tests {
|
|||||||
assert!(truncate_text("MCP text", &"b".repeat(500), 82).len() <= 82);
|
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]
|
#[test]
|
||||||
fn grep_content_gate_survives_a_nearly_exhausted_byte_budget() {
|
fn grep_content_gate_survives_a_nearly_exhausted_byte_budget() {
|
||||||
// 16 matches leave 16 bytes of the 32 KiB content budget, which is less
|
// 16 matches leave 16 bytes of the 32 KiB content budget, which is less
|
||||||
|
|||||||
Reference in New Issue
Block a user