Merge pull request #392 from kevin9327/fix/tool-result-truncation-budget

server/src/cursor/tools/tool_call_result/gate.rs

@@ -258,7 +258,9 @@ fn gate_grep_content(content: &mut pb::GrepContentResult, budget: &mut GrepBudge
                 truncated = true;
                 break;
             }
-            budget.content_bytes -= next_match.content.len();
+            budget.content_bytes = budget
+                .content_bytes
+                .saturating_sub(next_match.content.len());
             budget.matches -= 1;
             next.matches.push(next_match);
         }
@@ -630,15 +632,25 @@ fn truncate_text(tool_name: &str, content: &str, limit: usize) -> String {
     }
     let original = content.len();
     let mut shown = limit;
+    let mut previous = None;
     loop {
         let notice = format!(
             "\n\n[truncated: {tool_name} result exceeded {limit} bytes; showing {shown} of {original} bytes]"
         );
-        let available = limit.saturating_sub(notice.len());
-        let kept = utf8_prefix(content, available);
-        if kept.len() == shown {
+        if notice.len() >= limit {
+            // The notice alone would blow the budget; keep a plain prefix so the
+            // result never costs more than `limit` bytes.
+            return utf8_prefix(content, limit).to_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.
+        if kept.len() == shown || previous == Some(kept.len()) {
             return format!("{}{notice}", kept.trim_end_matches('\n'));
         }
+        previous = Some(shown);
         shown = kept.len();
     }
 }
@@ -685,3 +697,83 @@ fn utf8_suffix(value: &str, limit: usize) -> &str {
     }
     &value[start..]
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn grep_tool(matches: Vec<String>) -> pb::tool_call::Tool {
+        pb::tool_call::Tool::GrepToolCall(pb::GrepToolCall {
+            args: None,
+            result: Some(pb::GrepResult {
+                result: Some(pb::grep_result::Result::Success(pb::GrepSuccess {
+                    active_editor_result: Some(pb::GrepUnionResult {
+                        result: Some(pb::grep_union_result::Result::Content(
+                            pb::GrepContentResult {
+                                matches: vec![pb::GrepFileMatch {
+                                    file: "src/lib.rs".into(),
+                                    matches: matches
+                                        .into_iter()
+                                        .enumerate()
+                                        .map(|(index, content)| pb::GrepContentMatch {
+                                            line_number: index as i32 + 1,
+                                            content,
+                                            ..Default::default()
+                                        })
+                                        .collect(),
+                                }],
+                                ..Default::default()
+                            },
+                        )),
+                    }),
+                    ..Default::default()
+                })),
+            }),
+        })
+    }
+
+    #[test]
+    fn truncate_text_never_exceeds_its_limit() {
+        let content = "b".repeat(200);
+        for limit in 1..=250 {
+            let output = truncate_text("Grep", &content, limit);
+            assert!(
+                output.len() <= limit,
+                "limit {limit} produced {} bytes",
+                output.len()
+            );
+        }
+    }
+
+    #[test]
+    fn truncate_text_terminates_when_the_notice_length_oscillates() {
+        // `limit` values where the notice grows and shrinks with the digit count
+        // of the reported byte count, so the fixed point is never reached.
+        assert!(truncate_text("Grep", &"b".repeat(200), 78).len() <= 78);
+        assert!(truncate_text("Grep", &"b".repeat(200), 170).len() <= 170);
+        assert!(truncate_text("MCP text", &"b".repeat(500), 82).len() <= 82);
+    }
+
+    #[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
+        // than the truncation notice for the 17th match.
+        let mut matches = vec!["a".repeat(2047); 16];
+        matches.push("b".repeat(100));
+        let mut tool = grep_tool(matches);
+        let mut content = String::new();
+        tool_completion("Grep", &mut tool, &mut content);
+    }
+
+    #[test]
+    fn grep_content_gate_terminates_on_an_oscillating_remaining_budget() {
+        // The same path, tuned so the remaining budget lands on a `limit` where
+        // the truncation notice length oscillates.
+        let mut matches = vec!["a".repeat(2043); 15];
+        matches.push("a".repeat(2045));
+        matches.push("b".repeat(200));
+        let mut tool = grep_tool(matches);
+        let mut content = String::new();
+        tool_completion("Grep", &mut tool, &mut content);
+    }
+}