From 2ce89e5e5e9be5f06bafc15fe3b5f4f40a148dd3 Mon Sep 17 00:00:00 2001 From: Bryan Helmkamp Date: Wed, 11 Mar 2026 22:13:29 -0400 Subject: [PATCH] Simplify compaction code: extract helper method, named constant, avoid cloning Extract duplicated compaction check+trigger blocks in the session loop into Session::compact_if_needed(). Replace the hardcoded 20_000 token budget with COMPACTION_USER_MESSAGE_TOKEN_BUDGET constant. Change extract_recent_user_messages to take ownership of discarded turns instead of borrowing and cloning. Co-Authored-By: Claude Opus 4.6 (1M context) --- lib/crates/arc-agent/src/history.rs | 29 +++++--- lib/crates/arc-agent/src/session.rs | 101 ++++++++++------------------ 2 files changed, 55 insertions(+), 75 deletions(-) diff --git a/lib/crates/arc-agent/src/history.rs b/lib/crates/arc-agent/src/history.rs index d42575d4d..b54a9ac40 100644 --- a/lib/crates/arc-agent/src/history.rs +++ b/lib/crates/arc-agent/src/history.rs @@ -21,8 +21,9 @@ impl History { return; } let preserved = self.turns.split_off(self.turns.len() - preserve_count); - let extracted_user_messages = extract_recent_user_messages(&self.turns, 20_000); - self.turns.clear(); + let discarded = std::mem::take(&mut self.turns); + let extracted_user_messages = + extract_recent_user_messages(discarded, COMPACTION_USER_MESSAGE_TOKEN_BUDGET); self.turns.push(Turn::System { content: summary, timestamp: std::time::SystemTime::now(), @@ -101,27 +102,35 @@ impl History { } } +/// Maximum token budget for user messages extracted from discarded turns during compaction. +const COMPACTION_USER_MESSAGE_TOKEN_BUDGET: usize = 20_000; + /// Walk discarded turns in reverse, collecting `Turn::User` variants up to /// a token budget (estimated at ~4 chars per token). Returns them in /// chronological order so they can be inserted between the summary and the /// preserved tail. -fn extract_recent_user_messages(discarded: &[Turn], token_budget: usize) -> Vec { +fn extract_recent_user_messages(discarded: Vec, token_budget: usize) -> Vec { let char_budget = token_budget * 4; let mut total_chars = 0; - let mut collected: Vec = Vec::new(); + let mut first_kept_index = discarded.len(); - for turn in discarded.iter().rev() { + // Walk backward to find the earliest user message within budget + for (i, turn) in discarded.iter().enumerate().rev() { if let Turn::User { content, .. } = turn { if total_chars + content.len() > char_budget { break; } total_chars += content.len(); - collected.push(turn.clone()); + first_kept_index = i; } } - collected.reverse(); - collected + // Collect kept user messages in forward (chronological) order + discarded + .into_iter() + .skip(first_kept_index) + .filter(|t| matches!(t, Turn::User { .. })) + .collect() } #[cfg(test)] @@ -580,7 +589,7 @@ mod tests { timestamp: SystemTime::now(), }, ]; - let extracted = extract_recent_user_messages(&turns, 20_000); + let extracted = extract_recent_user_messages(turns, 20_000); assert_eq!(extracted.len(), 2); assert!(matches!(&extracted[0], Turn::User { content, .. } if content == "first")); assert!(matches!(&extracted[1], Turn::User { content, .. } if content == "second")); @@ -599,7 +608,7 @@ mod tests { }, ]; // Budget of 30 tokens = 120 chars; second message (100 chars) fits, first would exceed - let extracted = extract_recent_user_messages(&turns, 30); + let extracted = extract_recent_user_messages(turns, 30); assert_eq!(extracted.len(), 1); assert!(matches!(&extracted[0], Turn::User { content, .. } if content.starts_with('b'))); } diff --git a/lib/crates/arc-agent/src/session.rs b/lib/crates/arc-agent/src/session.rs index 55922357b..beb5e82ed 100644 --- a/lib/crates/arc-agent/src/session.rs +++ b/lib/crates/arc-agent/src/session.rs @@ -564,39 +564,8 @@ impl Session { return Err(self.aborted_error()); } - // Pre-turn compaction check: compact before building the request so the - // LLM call uses a trimmed context when we're already over threshold. - let over_threshold = crate::compaction::check_context_usage( - &self.system_prompt, - &self.history, - self.provider_profile.as_ref(), - self.config.compaction_threshold_percent, - &self.event_emitter, - &self.id, - ); - if over_threshold && self.config.enable_context_compaction { - if let Err(e) = crate::compaction::compact_context( - &mut self.history, - &self.llm_client, - self.provider_profile.as_ref(), - &self.system_prompt, - &self.file_tracker, - self.config.compaction_preserve_turns, - &self.event_emitter, - &self.id, - ) - .await - { - self.event_emitter.emit( - self.id.clone(), - AgentEvent::Error { - error: AgentError::InvalidState(format!( - "Context compaction failed: {e}" - )), - }, - ); - } - } + // Pre-turn compaction: trim context before building the request + self.compact_if_needed().await; // Build request let request = self.build_request(); @@ -744,38 +713,8 @@ impl Session { }, ); - // Check context window usage and compact if needed - let over_threshold = crate::compaction::check_context_usage( - &self.system_prompt, - &self.history, - self.provider_profile.as_ref(), - self.config.compaction_threshold_percent, - &self.event_emitter, - &self.id, - ); - if over_threshold && self.config.enable_context_compaction { - if let Err(e) = crate::compaction::compact_context( - &mut self.history, - &self.llm_client, - self.provider_profile.as_ref(), - &self.system_prompt, - &self.file_tracker, - self.config.compaction_preserve_turns, - &self.event_emitter, - &self.id, - ) - .await - { - self.event_emitter.emit( - self.id.clone(), - AgentEvent::Error { - error: AgentError::InvalidState(format!( - "Context compaction failed: {e}" - )), - }, - ); - } - } + // Post-response compaction: trim context after appending assistant turn + self.compact_if_needed().await; // If no tool calls, natural completion if tool_calls.is_empty() { @@ -838,6 +777,38 @@ impl Session { Ok(()) } + async fn compact_if_needed(&mut self) { + let over_threshold = crate::compaction::check_context_usage( + &self.system_prompt, + &self.history, + self.provider_profile.as_ref(), + self.config.compaction_threshold_percent, + &self.event_emitter, + &self.id, + ); + if over_threshold && self.config.enable_context_compaction { + if let Err(e) = crate::compaction::compact_context( + &mut self.history, + &self.llm_client, + self.provider_profile.as_ref(), + &self.system_prompt, + &self.file_tracker, + self.config.compaction_preserve_turns, + &self.event_emitter, + &self.id, + ) + .await + { + self.event_emitter.emit( + self.id.clone(), + AgentEvent::Error { + error: AgentError::InvalidState(format!("Context compaction failed: {e}")), + }, + ); + } + } + } + fn drain_steering(&mut self) { let messages: Vec = self .steering_queue