fix(context): one line budget per tool result, shared by its text blocks

CHANGELOG.md

@@ -4,6 +4,10 @@ All notable changes to `yoagent` are documented here. The format loosely
 follows [Keep a Changelog](https://keepachangelog.com/), and the project
 adheres to [Semantic Versioning](https://semver.org/).
 
+## Unreleased
+
+- **A tool result's line budget covers all its text blocks.** `tool_output_max_lines` was applied to each text block on its own, so since 0.25.9 sent every block (#347; Bedrock, and results with an image, sent them before) a result of N blocks could reach the model as N × the budget — 50 blocks from an MCP server, 10,000 lines where 200 were meant. The budget is now the result's: blocks no longer than an even share stay whole, the others split what is left (the remainder to the earliest), each cut head + tail with its own marker and stash key as before. A block not left whole keeps at least 5 lines (head, marker, tail; a shorter one stays whole), so when the even share would be under 5 lines — more text blocks than a fifth of the budget — the result goes over the budget by up to 5 lines per such block, never cut without a marker. Below a budget of 5 there is no room for a marker and nothing is shared: each block keeps the budget, as before. One-block results are unchanged, and truncation stays idempotent, so history written by this version is never rewritten. A multi-block result stored by an earlier version that is over the budget in total is cut once more, the first time compaction's level 1 runs over it: one prefix-cache rewrite. A stash key its old markers name may be cut out of the context then (the stash entry itself is kept), and that cut's own marker names no key: compaction does not stash.
+
 ## yoagent-rutis Unreleased
 
 yoagent-rutis only; requires yoagent 0.25, rutis 0.12 and Rust 1.88, as 0.7.0. Additive.

docs/concepts/context-management.md

@@ -148,7 +148,7 @@ let config = ContextConfig::from_context_window(1_000_000);
 
 ### Level 1: Truncate Tool Outputs
 
-Replaces long tool outputs with head + tail, fitting the result into exactly `tool_output_max_lines` lines (the `[... N lines truncated ...]` marker is charged against that budget). This is the cheapest level — it preserves conversation structure and typically saves 50-70% in coding sessions.
+Replaces long tool outputs with head + tail, fitting the result into `tool_output_max_lines` lines (exactly, for a single long block; the `[... N lines truncated ...]` marker is charged against that budget). The budget is the whole result's: when a result has several text blocks (an MCP server's, say), blocks no longer than an even share stay whole and the others split the rest, each cut with its own marker. A block not left whole keeps at least 5 lines (head, marker, tail), so when the even share would be under 5 lines the result goes over the budget by up to 5 lines per such block rather than being cut silently. Below a budget of 5 there is no room for a marker, and each block keeps the budget on its own. This is the cheapest level — it preserves conversation structure and typically saves 50-70% in coding sessions.
 
 Truncation is idempotent: re-running it on an already-truncated output returns it byte for byte, so a session that sits above the budget does not re-truncate the same outputs turn after turn.
 

src/context.rs

@@ -167,7 +167,8 @@ pub struct ContextConfig {
     pub keep_recent: usize,
     /// Minimum first messages to always keep
     pub keep_first: usize,
-    /// Max lines to keep per tool output, for tools with no override below.
+    /// Max lines to keep per tool result, shared by its text blocks (see
+    /// [`truncate_tool_output`]), for tools with no override below.
     ///
     /// Tuned for command output — build logs, test runs, `grep` — where
     /// head+tail is a good cut: the first error is at the top, the summary at
@@ -546,8 +547,8 @@ fn fnv1a(bytes: &[u8]) -> u64 {
 
 /// The stash key for one content block of a tool result.
 ///
-/// Block-qualified because a result can carry several text blocks, each
-/// truncated independently and each getting its own marker. Sharing one key
+/// Block-qualified because a result can carry several text blocks, each cut
+/// with its own marker (from one budget shared across them). Sharing one key
 /// across them made every marker resolve to all the blocks concatenated —
 /// never the one the model had just been told to fetch.
 pub fn block_key(base: &str, block: usize) -> String {
@@ -594,6 +595,9 @@ pub fn block_texts(msg: &AgentMessage) -> Vec<(usize, String)> {
 /// what a single shared key produced and which silently dropped any non-text
 /// content sitting between the blocks.
 ///
+/// Line budgets come from the blocks' line counts alone (one budget per result,
+/// shared), so the unkeyed and keyed passes cut the same lines.
+///
 /// Public so a caller implementing their own stash sink can reproduce exactly
 /// what the loop does: truncate with `None` to learn which blocks carry a
 /// marker, stash those, then truncate again with the base key.
@@ -610,16 +614,16 @@ pub fn truncate_tool_output_keyed(
             is_error,
             timestamp,
         }) => {
-            let max_lines = config.max_lines_for(&tool_name);
+            let budgets = block_budgets(&content, config.max_lines_for(&tool_name));
             let mut marked = Vec::new();
             let truncated_content: Vec<Content> = content
                 .into_iter()
+                .zip(budgets)
                 .enumerate()
-                .map(|(i, c)| match c {
+                .map(|(i, (c, budget))| match c {
                     Content::Text { text } => {
                         let key = base_key.map(|b| block_key(b, i));
-                        let (out, emitted) =
-                            truncate_text_head_tail(&text, max_lines, key.as_deref());
+                        let (out, emitted) = truncate_text_head_tail(&text, budget, key.as_deref());
                         if emitted {
                             marked.push(i);
                         }
@@ -644,8 +648,15 @@ pub fn truncate_tool_output_keyed(
     }
 }
 
-/// Cap a single tool result's text at the tool's line budget, keeping head and
-/// tail.
+/// Fit a single tool result's text into the tool's line budget, keeping head
+/// and tail.
+///
+/// The budget is the result's, shared by its text blocks: blocks no longer
+/// than an even share stay whole, and the others split what is left, each cut
+/// head + tail with its own marker. With a budget of 5 or more a cut block
+/// keeps at least 5 lines (head, marker, tail), so a result with many blocks
+/// can go over the budget; below 5, each block keeps the budget on its own and
+/// gets no marker.
 ///
 /// The budget comes from [`ContextConfig::max_lines_for`], so a tool that
 /// tolerates head+tail badly can be exempted. Non-tool-result messages pass
@@ -660,6 +671,76 @@ pub fn truncate_tool_output(msg: AgentMessage, config: &ContextConfig) -> AgentM
 /// Lines the truncation marker occupies: a blank line, the marker, a blank line.
 const TRUNCATION_MARKER_LINES: usize = 3;
 
+/// The fewest lines a cut block keeps: one head line, the marker's three lines,
+/// one tail line — so every cut block says it was cut (when the budget is at
+/// least this large).
+const MIN_BLOCK_LINES: usize = TRUNCATION_MARKER_LINES + 2;
+
+/// Each content block's line budget, from one budget for the whole result.
+///
+/// The text blocks share `max_lines`: when they fit together, every block is
+/// left alone; otherwise blocks no longer than an even share stay whole and
+/// the rest split what is left evenly (the remainder to the earliest), so a
+/// many-block result (an MCP server's) fits the budget one block would. A
+/// block not left whole is given at least [`MIN_BLOCK_LINES`] (so it is cut
+/// with a marker, or kept whole if shorter): when the even share is smaller
+/// than that, the result goes over `max_lines` by at most `MIN_BLOCK_LINES`
+/// per such block, never silently. When `max_lines < MIN_BLOCK_LINES` the
+/// floor is `max_lines` itself, so every block keeps the whole budget, as
+/// before, and gets no marker.
+///
+/// Idempotent with [`truncate_text_head_tail`]: a cut block is at most its
+/// budget long, so the result either fits afterwards (and a second pass
+/// leaves it alone) or the share fell below the floor, leaving every block
+/// not left whole at most the floor long; a second pass gives each such block
+/// at least the smaller of its length and the floor, so nothing is cut again.
+fn block_budgets(content: &[Content], max_lines: usize) -> Vec<usize> {
+    let lines: Vec<Option<usize>> = content
+        .iter()
+        .map(|c| match c {
+            Content::Text { text } => Some(text.lines().count()),
+            _ => None,
+        })
+        .collect();
+    let mut budgets: Vec<usize> = vec![usize::MAX; content.len()];
+    let total = lines
+        .iter()
+        .flatten()
+        .fold(0usize, |a, n| a.saturating_add(*n));
+    if total <= max_lines {
+        return budgets;
+    }
+    // Shortest first: a block that fits an even share of what is left keeps
+    // all its lines and leaves the rest to the others.
+    let mut order: Vec<(usize, usize)> = lines
+        .iter()
+        .enumerate()
+        .filter_map(|(i, n)| n.map(|n| (n, i)))
+        .collect();
+    order.sort_unstable();
+    let mut left = max_lines;
+    let mut rest = order.len();
+    let mut cut = Vec::new();
+    for (n, i) in order {
+        if cut.is_empty() && n <= left / rest {
+            budgets[i] = n;
+            left -= n;
+            rest -= 1;
+        } else {
+            cut.push(i);
+        }
+    }
+    // The cut blocks, in content order, share what is left.
+    cut.sort_unstable();
+    let share = left / cut.len();
+    let extra = left % cut.len();
+    let floor = MIN_BLOCK_LINES.min(max_lines);
+    for (k, i) in cut.into_iter().enumerate() {
+        budgets[i] = (share + usize::from(k < extra)).max(floor);
+    }
+    budgets
+}
+
 /// Truncate text keeping the head and tail, dropping the middle.
 ///
 /// The marker is charged against `max_lines`, so the result is exactly
@@ -1661,6 +1742,172 @@ mod tests {
         assert!(!marker.chars().any(|c| c.is_ascii_digit()));
     }
 
+    fn lines_block(prefix: &str, n: usize) -> Content {
+        Content::Text {
+            text: (0..n)
+                .map(|i| format!("{prefix}{i}"))
+                .collect::<Vec<_>>()
+                .join("\n"),
+        }
+    }
+
+    fn result_with(content: Vec<Content>) -> AgentMessage {
+        AgentMessage::Llm(Message::ToolResult {
+            tool_call_id: "tc-1".into(),
+            tool_name: "mcp_tool".into(),
+            content,
+            is_error: false,
+            timestamp: 0,
+        })
+    }
+
+    fn content_of(msg: &AgentMessage) -> &[Content] {
+        match msg {
+            AgentMessage::Llm(Message::ToolResult { content, .. }) => content,
+            _ => panic!("expected a tool result"),
+        }
+    }
+
+    fn block_lines(msg: &AgentMessage) -> Vec<usize> {
+        block_texts(msg)
+            .into_iter()
+            .map(|(_, t)| t.lines().count())
+            .collect()
+    }
+
+    #[test]
+    fn the_line_budget_is_the_result_s_not_each_block_s() {
+        let config = ContextConfig {
+            tool_output_max_lines: 200,
+            ..Default::default()
+        };
+        // Twenty blocks of 200 lines: a per-block budget would keep 4,000.
+        let many = result_with(
+            (0..20)
+                .map(|b| lines_block(&format!("b{b}-"), 200))
+                .collect(),
+        );
+        let (once, marked) = truncate_tool_output_keyed(many, &config, None);
+        assert_eq!(block_lines(&once).iter().sum::<usize>(), 200);
+        assert_eq!(
+            marked,
+            (0..20).collect::<Vec<_>>(),
+            "every cut block says so"
+        );
+        assert_eq!(
+            truncate_tool_output(once.clone(), &config),
+            once,
+            "idempotent"
+        );
+
+        // Blocks that fit together are left alone, however many.
+        let fits = result_with((0..4).map(|b| lines_block(&format!("b{b}-"), 50)).collect());
+        assert_eq!(truncate_tool_output(fits.clone(), &config), fits);
+
+        // A short block stays whole; the long ones share what is left, the
+        // remainder to the earliest; an image keeps its place and its index.
+        let mixed = result_with(vec![
+            lines_block("a", 300),
+            Content::Image {
+                data: "x".into(),
+                mime_type: "image/png".into(),
+            },
+            lines_block("s", 10),
+            lines_block("c", 300),
+        ]);
+        let (once, marked) = truncate_tool_output_keyed(mixed, &config, Some("k"));
+        assert_eq!(block_lines(&once), [95, 10, 95]);
+        assert_eq!(marked, [0, 3]);
+        let texts = block_texts(&once);
+        assert_eq!(texts[1].1.lines().count(), 10);
+        assert!(
+            texts[2].1.contains("shared_state get \"k-b3\""),
+            "{}",
+            texts[2].1
+        );
+        assert!(matches!(&content_of(&once)[1], Content::Image { .. }));
+        assert_eq!(
+            truncate_tool_output(once.clone(), &config),
+            once,
+            "idempotent"
+        );
+    }
+
+    #[test]
+    fn the_shared_budget_s_remainder_and_freed_slack() {
+        let config = ContextConfig {
+            tool_output_max_lines: 200,
+            ..Default::default()
+        };
+        // 200 over three long blocks: the remainder goes to the earliest.
+        let three = result_with(
+            (0..3)
+                .map(|b| lines_block(&format!("b{b}-"), 300))
+                .collect(),
+        );
+        assert_eq!(
+            block_lines(&truncate_tool_output(three, &config)),
+            [67, 67, 66]
+        );
+        // 60 lines is over an even share of 150 (50) but fits once the
+        // 10-line block has left 140 for two (70): it stays whole.
+        let config = ContextConfig {
+            tool_output_max_lines: 150,
+            ..Default::default()
+        };
+        let slack = result_with(vec![
+            lines_block("s", 10),
+            lines_block("m", 60),
+            lines_block("l", 300),
+        ]);
+        assert_eq!(
+            block_lines(&truncate_tool_output(slack, &config)),
+            [10, 60, 80]
+        );
+        // A tool exempted by an override stays whole, however many blocks.
+        let mut config = ContextConfig::default();
+        config
+            .tool_output_max_lines_overrides
+            .insert("mcp_tool".into(), usize::MAX);
+        let many = result_with(
+            (0..20)
+                .map(|b| lines_block(&format!("b{b}-"), 200))
+                .collect(),
+        );
+        assert_eq!(truncate_tool_output(many.clone(), &config), many);
+    }
+
+    #[test]
+    fn more_long_blocks_than_the_budget_can_mark_keep_a_marker_each() {
+        let config = ContextConfig {
+            tool_output_max_lines: 20,
+            ..Default::default()
+        };
+        // Ten long blocks and room for two lines each: every one keeps the
+        // fewest lines a marker needs rather than being cut silently.
+        let msg = result_with(
+            (0..10)
+                .map(|b| lines_block(&format!("b{b}-"), 100))
+                .collect(),
+        );
+        let (once, marked) = truncate_tool_output_keyed(msg, &config, None);
+        assert_eq!(block_lines(&once), vec![MIN_BLOCK_LINES; 10]);
+        assert_eq!(marked.len(), 10);
+        assert_eq!(
+            truncate_tool_output(once.clone(), &config),
+            once,
+            "idempotent"
+        );
+        // Short blocks past that point stay whole: sixty 4-line blocks keep
+        // 240 lines of a 200-line budget, none cut.
+        let config = ContextConfig {
+            tool_output_max_lines: 200,
+            ..Default::default()
+        };
+        let short = result_with((0..60).map(|b| lines_block(&format!("b{b}-"), 4)).collect());
+        assert_eq!(truncate_tool_output(short.clone(), &config), short);
+    }
+
     #[test]
     fn test_truncate_tool_output_helper() {
         let big = (0..500)

tests/agent_loop_test.rs

@@ -4211,6 +4211,24 @@ async fn each_text_block_gets_its_own_key_and_resolves_to_itself() {
         _ => false,
     });
     assert!(has_image, "non-text content must pass through untouched");
+
+    // The two blocks share the result's 20 lines, rather than taking 20 each.
+    let lines: Vec<usize> = messages
+        .iter()
+        .find_map(|m| match m {
+            AgentMessage::Llm(Message::ToolResult { content, .. }) => Some(
+                content
+                    .iter()
+                    .filter_map(|c| match c {
+                        Content::Text { text } => Some(text.lines().count()),
+                        _ => None,
+                    })
+                    .collect(),
+            ),
+            _ => None,
+        })
+        .expect("a tool result");
+    assert_eq!(lines, [10, 10]);
 }
 
 /// Only the second text block is long enough to truncate, so `marked` is