From b4dbb85d7b8c7c7d8af2fab27265cd559fb88f43 Mon Sep 17 00:00:00 2001 From: xingshining Date: Wed, 26 Aug 2026 10:46:53 +0800 Subject: [PATCH] fix(agent): preserve unconsumed tool results during microcompact --- crates/aion-agent/src/compact/micro.rs | 42 +++++-- crates/aion-agent/src/compact/micro_test.rs | 116 ++++++++++++++---- crates/aion-agent/src/engine_test.rs | 11 +- .../tests/acceptance/compact_test.rs | 13 +- .../aion-agent/tests/engine_compact_test.rs | 15 +-- crates/aion-agent/tests/microcompact_test.rs | 29 +++-- 6 files changed, 165 insertions(+), 61 deletions(-) diff --git a/crates/aion-agent/src/compact/micro.rs b/crates/aion-agent/src/compact/micro.rs index f01ba28f..03a00619 100644 --- a/crates/aion-agent/src/compact/micro.rs +++ b/crates/aion-agent/src/compact/micro.rs @@ -2,7 +2,7 @@ //! //! This is the lightest compaction level. It walks the conversation, //! identifies tool results from compactable tools, and replaces the -//! content of all but the N most recent with a short placeholder. +//! content of old, model-consumed results with a short placeholder. use std::collections::{HashMap, HashSet}; @@ -29,7 +29,7 @@ pub struct MicrocompactResult { /// Returns `true` if **either** trigger fires: /// - **Time**: the most recent assistant message is older than /// `config.micro_gap_seconds`. -/// - **Count**: total compactable (non-cleared) tool results exceed +/// - **Count**: total compactable, model-consumed tool results exceed /// `config.micro_keep_recent * 2`. pub fn should_microcompact(messages: &[Message], config: &CompactConfig) -> bool { if !config.enabled { @@ -59,7 +59,8 @@ fn count_trigger(messages: &[Message], config: &CompactConfig) -> bool { let tool_names = build_tool_name_map(messages); let compactable_set: HashSet<&str> = config.compactable_tools.iter().map(String::as_str).collect(); - let count = count_compactable_results(messages, &tool_names, &compactable_set); + let protected_ids = unconsumed_tool_result_ids(messages); + let count = count_compactable_results(messages, &tool_names, &compactable_set, &protected_ids); count > config.micro_keep_recent * 2 } @@ -74,11 +75,11 @@ fn count_trigger(messages: &[Message], config: &CompactConfig) -> bool { pub fn microcompact(messages: &mut [Message], config: &CompactConfig) -> MicrocompactResult { let tool_names = build_tool_name_map(messages); let compactable_set: HashSet<&str> = config.compactable_tools.iter().map(String::as_str).collect(); + let protected_ids = unconsumed_tool_result_ids(messages); // Collect (message_index, block_index) of all compactable, non-cleared // tool results, in conversation order. - let targets = collect_compactable_locations(messages, &tool_names, &compactable_set); - + let targets = collect_compactable_locations(messages, &tool_names, &compactable_set, &protected_ids); let keep = config.micro_keep_recent.max(1); if targets.len() <= keep { return MicrocompactResult { @@ -128,11 +129,12 @@ fn count_compactable_results( messages: &[Message], tool_names: &HashMap, compactable_set: &HashSet<&str>, + protected_ids: &HashSet<&str>, ) -> usize { messages .iter() .flat_map(|m| &m.content) - .filter(|b| is_compactable_and_live(b, tool_names, compactable_set)) + .filter(|b| is_compactable_and_live(b, tool_names, compactable_set, protected_ids)) .count() } @@ -142,11 +144,12 @@ fn collect_compactable_locations( messages: &[Message], tool_names: &HashMap, compactable_set: &HashSet<&str>, + protected_ids: &HashSet<&str>, ) -> Vec<(usize, usize)> { let mut locations = Vec::new(); for (mi, msg) in messages.iter().enumerate() { for (bi, block) in msg.content.iter().enumerate() { - if is_compactable_and_live(block, tool_names, compactable_set) { + if is_compactable_and_live(block, tool_names, compactable_set, protected_ids) { locations.push((mi, bi)); } } @@ -158,16 +161,19 @@ fn collect_compactable_locations( /// 1. It is a `ToolResult` variant. /// 2. Its corresponding tool name is in the compactable set. /// 3. Its content has not already been cleared. +/// 4. It is not part of the latest assistant tool batch, which the model has +/// not had an opportunity to consume yet. fn is_compactable_and_live( block: &ContentBlock, tool_names: &HashMap, compactable_set: &HashSet<&str>, + protected_ids: &HashSet<&str>, ) -> bool { if let ContentBlock::ToolResult { tool_use_id, content, .. } = block { - if content == CLEARED_TOOL_RESULT { + if content == CLEARED_TOOL_RESULT || protected_ids.contains(tool_use_id.as_str()) { return false; } if let Some(name) = tool_names.get(tool_use_id) { @@ -177,6 +183,26 @@ fn is_compactable_and_live( false } +/// Return the tool ids emitted by the latest assistant message. +/// +/// The engine appends an assistant tool batch before executing it, then appends +/// the corresponding user tool results before the next provider request. Until +/// that request occurs, the entire batch is unconsumed and must not be compacted. +fn unconsumed_tool_result_ids(messages: &[Message]) -> HashSet<&str> { + let Some(message) = messages.iter().rev().find(|message| message.role == Role::Assistant) else { + return HashSet::new(); + }; + + message + .content + .iter() + .filter_map(|block| match block { + ContentBlock::ToolUse { id, .. } => Some(id.as_str()), + _ => None, + }) + .collect() +} + #[cfg(test)] #[path = "micro_test.rs"] mod micro_test; diff --git a/crates/aion-agent/src/compact/micro_test.rs b/crates/aion-agent/src/compact/micro_test.rs index 0e23abed..7f153851 100644 --- a/crates/aion-agent/src/compact/micro_test.rs +++ b/crates/aion-agent/src/compact/micro_test.rs @@ -79,40 +79,45 @@ mod tests { fn live_compactable_result_returns_true() { let tool_names: HashMap = [("t1".into(), "Read".into())].into_iter().collect(); let set: HashSet<&str> = ["Read"].into_iter().collect(); + let protected = HashSet::new(); let block = tool_result_block("t1", "file content here"); - assert!(is_compactable_and_live(&block, &tool_names, &set)); + assert!(is_compactable_and_live(&block, &tool_names, &set, &protected)); } #[test] fn already_cleared_result_returns_false() { let tool_names: HashMap = [("t1".into(), "Read".into())].into_iter().collect(); let set: HashSet<&str> = ["Read"].into_iter().collect(); + let protected = HashSet::new(); let block = tool_result_block("t1", CLEARED_TOOL_RESULT); - assert!(!is_compactable_and_live(&block, &tool_names, &set)); + assert!(!is_compactable_and_live(&block, &tool_names, &set, &protected)); } #[test] fn non_compactable_tool_returns_false() { let tool_names: HashMap = [("t1".into(), "Skill".into())].into_iter().collect(); let set: HashSet<&str> = ["Read", "ExecCommand"].into_iter().collect(); + let protected = HashSet::new(); let block = tool_result_block("t1", "result"); - assert!(!is_compactable_and_live(&block, &tool_names, &set)); + assert!(!is_compactable_and_live(&block, &tool_names, &set, &protected)); } #[test] fn text_block_returns_false() { let tool_names = HashMap::new(); let set: HashSet<&str> = ["Read"].into_iter().collect(); + let protected = HashSet::new(); let block = text_block("hello"); - assert!(!is_compactable_and_live(&block, &tool_names, &set)); + assert!(!is_compactable_and_live(&block, &tool_names, &set, &protected)); } #[test] fn unknown_tool_use_id_returns_false() { let tool_names = HashMap::new(); // no ToolUse registered let set: HashSet<&str> = ["Read"].into_iter().collect(); + let protected = HashSet::new(); let block = tool_result_block("orphan", "data"); - assert!(!is_compactable_and_live(&block, &tool_names, &set)); + assert!(!is_compactable_and_live(&block, &tool_names, &set, &protected)); } // ── time_trigger ──────────────────────────────────────────────────── @@ -166,9 +171,10 @@ mod tests { #[test] fn count_trigger_fires_above_threshold() { - // keep_recent=3, threshold=6. Create 7 compactable results. + // keep_recent=3, threshold=6. The latest round is protected, so + // create 8 results to leave 7 consumed results eligible for compaction. let mut msgs = Vec::new(); - for i in 0..7 { + for i in 0..8 { let id = format!("t{i}"); msgs.push(assistant_msg(vec![tool_use_block(&id, "Read")])); msgs.push(user_msg(vec![tool_result_block(&id, "data")])); @@ -196,11 +202,30 @@ mod tests { assert!(!count_trigger(&msgs, &config)); } + #[test] + fn count_trigger_ignores_unconsumed_tool_round() { + let mut tool_uses = Vec::new(); + let mut tool_results = Vec::new(); + for i in 0..10 { + let id = format!("current-{i}"); + tool_uses.push(tool_use_block(&id, "ExecCommand")); + tool_results.push(tool_result_block(&id, "current output")); + } + let msgs = vec![assistant_msg(tool_uses), user_msg(tool_results)]; + let config = CompactConfig { + micro_keep_recent: 5, + ..default_config() + }; + + assert!(!count_trigger(&msgs, &config)); + } + // ── microcompact ──────────────────────────────────────────────────── #[test] fn clears_oldest_keeps_recent() { - // 5 tool results, keep_recent=2 → clear 3. + // Five tool results, with the latest round protected. Keep two of the + // four consumed results, so only the two oldest are cleared. let mut msgs = Vec::new(); for i in 0..5 { let id = format!("t{i}"); @@ -213,19 +238,19 @@ mod tests { }; let result = microcompact(&mut msgs, &config); - assert_eq!(result.cleared_count, 3); + assert_eq!(result.cleared_count, 2); assert!(result.estimated_tokens_freed > 0); - // First 3 user msgs (indices 1,3,5) should be cleared. - for idx in [1, 3, 5] { + // First two consumed results (indices 1,3) should be cleared. + for idx in [1, 3] { let content = match &msgs[idx].content[0] { ContentBlock::ToolResult { content, .. } => content.as_str(), _ => panic!("expected ToolResult"), }; assert_eq!(content, CLEARED_TOOL_RESULT); } - // Last 2 user msgs (indices 7,9) should retain original content. - for (idx, expected) in [(7, "data-3"), (9, "data-4")] { + // The remaining consumed results and protected current result survive. + for (idx, expected) in [(5, "data-2"), (7, "data-3"), (9, "data-4")] { let content = match &msgs[idx].content[0] { ContentBlock::ToolResult { content, .. } => content.as_str(), _ => panic!("expected ToolResult"), @@ -249,9 +274,45 @@ mod tests { assert_eq!(result.estimated_tokens_freed, 0); } + #[test] + fn preserves_every_result_in_the_unconsumed_tool_round() { + let mut msgs = Vec::new(); + for i in 0..6 { + let id = format!("history-{i}"); + msgs.push(assistant_msg(vec![tool_use_block(&id, "ExecCommand")])); + msgs.push(user_msg(vec![tool_result_block(&id, &format!("history output {i}"))])); + } + + let mut current_calls = Vec::new(); + let mut current_results = Vec::new(); + for i in 0..10 { + let id = format!("current-{i}"); + current_calls.push(tool_use_block(&id, "ExecCommand")); + current_results.push(tool_result_block(&id, &format!("current output {i}"))); + } + msgs.push(assistant_msg(current_calls)); + msgs.push(user_msg(current_results)); + + let config = CompactConfig { + micro_keep_recent: 5, + ..default_config() + }; + let result = microcompact(&mut msgs, &config); + + assert_eq!(result.cleared_count, 1); + for i in 0..10 { + let ContentBlock::ToolResult { content, .. } = &msgs.last().unwrap().content[i] else { + panic!("expected ToolResult"); + }; + assert_eq!(content, &format!("current output {i}")); + } + } + #[test] fn skips_non_compactable_tools() { let mut msgs = vec![ + assistant_msg(vec![tool_use_block("t0", "Read")]), + user_msg(vec![tool_result_block("t0", "older-file-data")]), assistant_msg(vec![tool_use_block("t1", "Read")]), user_msg(vec![tool_result_block("t1", "file-data")]), assistant_msg(vec![tool_use_block("t2", "Skill")]), @@ -267,11 +328,12 @@ mod tests { }; let result = microcompact(&mut msgs, &config); - // Only Read(t1) should be cleared; ExecCommand(t3) kept as most recent. + // Only the oldest consumed Read(t0) should be cleared. The Skill result + // is non-compactable and ExecCommand(t3) is the protected current round. assert_eq!(result.cleared_count, 1); // Skill result untouched. - match &msgs[3].content[0] { + match &msgs[5].content[0] { ContentBlock::ToolResult { content, .. } => { assert_eq!(content, "skill-output"); } @@ -282,18 +344,20 @@ mod tests { #[test] fn does_not_recleared_already_cleared() { let mut msgs = vec![ + assistant_msg(vec![tool_use_block("t0", "Read")]), + user_msg(vec![tool_result_block("t0", CLEARED_TOOL_RESULT)]), assistant_msg(vec![tool_use_block("t1", "Read")]), - user_msg(vec![tool_result_block("t1", CLEARED_TOOL_RESULT)]), + user_msg(vec![tool_result_block("t1", "live-data")]), assistant_msg(vec![tool_use_block("t2", "Read")]), - user_msg(vec![tool_result_block("t2", "live-data")]), + user_msg(vec![tool_result_block("t2", "current-data")]), ]; let config = CompactConfig { micro_keep_recent: 1, ..default_config() }; let result = microcompact(&mut msgs, &config); - // t1 already cleared → not in compactable list. - // Only t2 is compactable, and it's the most recent → keep it. + // t0 already cleared → not in compactable list. + // t1 is consumed but remains within the keep budget; t2 is protected. assert_eq!(result.cleared_count, 0); } @@ -334,10 +398,12 @@ mod tests { fn token_estimate_proportional_to_content() { let long_content = "x".repeat(400); // ~100 tokens let mut msgs = vec![ + assistant_msg(vec![tool_use_block("t0", "Read")]), + user_msg(vec![tool_result_block("t0", &long_content)]), assistant_msg(vec![tool_use_block("t1", "Read")]), - user_msg(vec![tool_result_block("t1", &long_content)]), + user_msg(vec![tool_result_block("t1", "keep")]), assistant_msg(vec![tool_use_block("t2", "Read")]), - user_msg(vec![tool_result_block("t2", "keep")]), + user_msg(vec![tool_result_block("t2", "current")]), ]; let config = CompactConfig { micro_keep_recent: 1, @@ -366,6 +432,8 @@ mod tests { fn keep_recent_floored_at_one() { // Even with keep_recent=0, we never clear everything. let mut msgs = vec![ + assistant_msg(vec![tool_use_block("t0", "Read")]), + user_msg(vec![tool_result_block("t0", "data-0")]), assistant_msg(vec![tool_use_block("t1", "Read")]), user_msg(vec![tool_result_block("t1", "data-1")]), assistant_msg(vec![tool_use_block("t2", "Read")]), @@ -376,10 +444,10 @@ mod tests { ..default_config() }; let result = microcompact(&mut msgs, &config); - // 2 compactable, keep at least 1 → clear 1. + // Two consumed results, keep at least 1 → clear 1. assert_eq!(result.cleared_count, 1); - // The most recent (t2) must survive. - match &msgs[3].content[0] { + // The protected current result (t2) must survive. + match &msgs[5].content[0] { ContentBlock::ToolResult { content, .. } => { assert_eq!(content, "data-2"); } diff --git a/crates/aion-agent/src/engine_test.rs b/crates/aion-agent/src/engine_test.rs index 005ee81b..7fc257ff 100644 --- a/crates/aion-agent/src/engine_test.rs +++ b/crates/aion-agent/src/engine_test.rs @@ -1310,7 +1310,8 @@ mod tests_compact { #[tokio::test] async fn microcompact_clears_old_results() { - // 12 tool results with keep_recent=3 (threshold=6) → should clear 9 + // The latest tool round is protected. Of the 11 consumed results, + // keep_recent=3 leaves 8 eligible for clearing. let mut messages = Vec::new(); for i in 0..12 { let id = format!("t{i}"); @@ -1330,7 +1331,7 @@ mod tests_compact { engine.sync_compact_watermark(); engine.run_compaction().await.unwrap(); - // Last 3 tool results should be preserved + // The three most recent consumed results and the current round remain. let cleared_count = engine .messages .iter() @@ -1338,7 +1339,7 @@ mod tests_compact { .filter(|b| matches!(b, ContentBlock::ToolResult { content, .. } if content == "[Tool result cleared]")) .count(); - assert_eq!(cleared_count, 9); + assert_eq!(cleared_count, 8); assert_eq!(engine.context_state.microcompact_count, 1); assert_eq!(engine.context_state.context_usage, 1_000); assert_eq!(engine.compact_state.last_input_tokens, 1_000); @@ -1348,7 +1349,9 @@ mod tests_compact { #[tokio::test] async fn microcompact_does_not_lower_the_emergency_watermark() { let mut messages = Vec::new(); - for i in 0..3 { + // The latest round is protected, so use four rounds to leave three + // consumed results and still exceed the count trigger threshold. + for i in 0..4 { let id = format!("t{i}"); messages.push(tool_use_msg(&id, "Read")); messages.push(tool_result_msg(&id, &"x".repeat(4_000))); diff --git a/crates/aion-agent/tests/acceptance/compact_test.rs b/crates/aion-agent/tests/acceptance/compact_test.rs index 57779f1f..4da5fd64 100644 --- a/crates/aion-agent/tests/acceptance/compact_test.rs +++ b/crates/aion-agent/tests/acceptance/compact_test.rs @@ -36,8 +36,9 @@ fn tool_result_block(id: &str, content: &str) -> ContentBlock { /// Construct a message history with more than `micro_keep_recent * 2` /// compactable tool results (each with a matching ToolUse block), run -/// microcompact, and verify that old results are cleared while the most -/// recent `micro_keep_recent` are preserved. +/// microcompact, and verify that old, model-consumed results are cleared while +/// the most recent `micro_keep_recent` consumed results and current tool round +/// are preserved. #[test] fn microcompact_clears_old_tool_results() { let keep_recent: usize = 3; @@ -69,8 +70,10 @@ fn microcompact_clears_old_tool_results() { "microcompact should clear at least one tool result, got cleared_count=0" ); - // Exactly total_results - keep_recent should be cleared - let expected_cleared = total_results - keep_recent; + // The latest assistant tool batch has not been consumed by a model turn, + // so only the preceding results are eligible for clearing. + let consumed_results = total_results - 1; + let expected_cleared = consumed_results - keep_recent; assert_eq!( result.cleared_count, expected_cleared, "expected {expected_cleared} cleared, got {}", @@ -91,7 +94,7 @@ fn microcompact_clears_old_tool_results() { } } - // Verify most recent `keep_recent` results are preserved + // Verify the most recent consumed results and current round are preserved. for i in expected_cleared..total_results { let user_msg_idx = i * 2 + 1; match &messages[user_msg_idx].content[0] { diff --git a/crates/aion-agent/tests/engine_compact_test.rs b/crates/aion-agent/tests/engine_compact_test.rs index d33276e4..1fd962d9 100644 --- a/crates/aion-agent/tests/engine_compact_test.rs +++ b/crates/aion-agent/tests/engine_compact_test.rs @@ -451,16 +451,17 @@ async fn tc_2_6_02_micro_before_auto_execution_order() { v }; - // Turns 0-6: tool use. Turn 6 reports high input_tokens + // Turns 0-7: tool use. Turn 7 reports high input_tokens // so that micro and auto both trigger in the SAME cycle - // (turn 7's run_compaction). - // Turn 7 (after compact): text to end the run. + // (turn 8's run_compaction). + // Turn 8 (after compact): text to end the run. // // micro_keep_recent = 3 → count threshold = 6. - // After 7 tool-use turns: 7 > 6 → micro fires. - // After turn 6: context_tokens = 170k > 167k → auto fires. - let events = if count < 7 { - let input_tokens = if count == 6 { 170_000 } else { 10_000 }; + // The newest of 8 tool rounds is protected, leaving 7 consumed + // results. Therefore 7 > 6 and micro fires. + // After turn 7: context_tokens = 170k > 167k → auto fires. + let events = if count < 8 { + let input_tokens = if count == 7 { 170_000 } else { 10_000 }; vec![ LlmEvent::ToolUse { id: format!("t{count}"), diff --git a/crates/aion-agent/tests/microcompact_test.rs b/crates/aion-agent/tests/microcompact_test.rs index 642115a1..08e1d8c3 100644 --- a/crates/aion-agent/tests/microcompact_test.rs +++ b/crates/aion-agent/tests/microcompact_test.rs @@ -61,8 +61,9 @@ fn get_tool_result_content(msg: &Message, block_idx: usize) -> &str { #[test] fn tc_2_3_01_basic_clearing() { - // 10 messages containing 8 tool results (Read x3, ExecCommand x3, Grep x2). - // keep_recent = 3 → oldest 5 cleared. + // Eight tool results (Read x3, ExecCommand x3, Grep x2), where the last + // tool round is protected. Of the seven consumed results, keep_recent = 3 + // leaves the oldest four eligible for clearing. let tool_specs = [ ("r1", "Read"), ("b1", "ExecCommand"), @@ -85,10 +86,10 @@ fn tc_2_3_01_basic_clearing() { }; let result = microcompact(&mut msgs, &config); - assert_eq!(result.cleared_count, 5); + assert_eq!(result.cleared_count, 4); - // First 5 user messages (indices 1, 3, 5, 7, 9) are cleared. - for i in 0..5 { + // First four user messages (indices 1, 3, 5, 7) are cleared. + for i in 0..4 { let user_msg_idx = i * 2 + 1; assert_eq!( get_tool_result_content(&msgs[user_msg_idx], 0), @@ -96,8 +97,8 @@ fn tc_2_3_01_basic_clearing() { "tool result at msg index {user_msg_idx} should be cleared" ); } - // Last 3 user messages (indices 11, 13, 15) retain original content. - for (idx, &(id, _name)) in tool_specs.iter().enumerate().skip(5) { + // Remaining consumed results and current round retain original content. + for (idx, &(id, _name)) in tool_specs.iter().enumerate().skip(4) { let user_msg_idx = idx * 2 + 1; assert_eq!( get_tool_result_content(&msgs[user_msg_idx], 0), @@ -156,8 +157,9 @@ fn tc_2_3_03_only_compactable_tools_cleared() { }; let result = microcompact(&mut msgs, &config); - // 3 compactable (t1-Read, t2-ExecCommand, t4-Read), keep 1 → clear 2. - assert_eq!(result.cleared_count, 2); + // The latest compactable result (t4) is protected. Of the two consumed + // compactable results, keep 1 leaves only t1 eligible for clearing. + assert_eq!(result.cleared_count, 1); // Skill result (t3) must be untouched. assert_eq!(get_tool_result_content(&msgs[5], 0), "skill-output"); @@ -236,7 +238,8 @@ fn tc_2_3_07_no_timestamp_skips_time_check() { #[test] fn tc_2_3_08_token_estimation() { - // 3 tool results with known content lengths, clear all but 1. + // The final tool result is protected. Of the two consumed results, clear + // the oldest and keep the other one. let content_a = "x".repeat(200); // 50 tokens let content_b = "y".repeat(400); // 100 tokens let content_c = "z".repeat(80); // 20 tokens — kept @@ -254,10 +257,10 @@ fn tc_2_3_08_token_estimation() { }; let result = microcompact(&mut msgs, &config); - assert_eq!(result.cleared_count, 2); + assert_eq!(result.cleared_count, 1); assert!(result.estimated_tokens_freed > 0); - // 200/4 + 400/4 = 50 + 100 = 150 - assert_eq!(result.estimated_tokens_freed, 150); + // 200/4 = 50 + assert_eq!(result.estimated_tokens_freed, 50); } // ── TC-2.3-09: Already cleared content not re-cleared ───────────────────────