From 51e06c8f734c528b29f76c5bfe60a4e6dc306e5b Mon Sep 17 00:00:00 2001 From: oudecheng <13802883547@139.com> Date: Fri, 3 Jul 2026 09:52:45 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E5=AE=9E=E7=8E=B0=E4=B8=A4=E6=AE=B5?= =?UTF-8?q?=E5=8E=86=E5=8F=B2=E5=8E=8B=E7=BC=A9=E7=9A=84=E5=90=8C=E6=AD=A5?= =?UTF-8?q?=E5=A4=84=E7=90=86=EF=BC=8C=E6=9B=BF=E6=8D=A2=E4=BC=9A=E8=AF=9D?= =?UTF-8?q?=E7=9A=84=E6=B4=BB=E5=8A=A8=E5=8E=86=E5=8F=B2=E8=AE=B0=E5=BD=95?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/agent/context_compressor.rs | 561 +++++++++++++++++++++++++++++++- src/gateway/compaction.rs | 117 +++---- src/gateway/session.rs | 2 + src/gateway/session_history.rs | 2 + src/storage/mod.rs | 53 +++ src/storage/ports.rs | 17 + 6 files changed, 669 insertions(+), 83 deletions(-) diff --git a/src/agent/context_compressor.rs b/src/agent/context_compressor.rs index 6f90e6a..211598d 100644 --- a/src/agent/context_compressor.rs +++ b/src/agent/context_compressor.rs @@ -5,7 +5,6 @@ use crate::bus::{ use crate::config::LLMProviderConfig; use crate::providers::{ChatCompletionRequest, LLMProvider, Message, create_provider}; use crate::text::{char_count, take_prefix_chars}; - use crate::agent::{AgentError, AgentRuntimeConfig}; const TOKEN_ESTIMATE_SAFETY_MULTIPLIER: f64 = 1.2; @@ -13,6 +12,145 @@ const CJK_CHARS_PER_TOKEN: f64 = 2.0; const OTHER_CHARS_PER_TOKEN: f64 = 4.0; const JSON_OVERHEAD_PER_MESSAGE: usize = 50; +/// System context marker for the heavy compression (older segment) summary. +pub const SYSTEM_CONTEXT_HISTORY_COMPACTION_OLDER: &str = "history_compaction_older"; +/// System context marker for the light compression (newer segment) summary. +pub const SYSTEM_CONTEXT_HISTORY_COMPACTION_NEWER: &str = "history_compaction_newer"; + +/// Default threshold ratio: compress when estimated tokens exceed 70% of context window. +const DEFAULT_THRESHOLD_RATIO: f64 = 0.7; + +/// Default budget split: older segment gets 30% of summary budget, newer gets 70%. +const OLDER_BUDGET_RATIO: f64 = 0.3; +const NEWER_BUDGET_RATIO: f64 = 0.7; + +// ============================================================================ +// HistoryUnit — atomic message units for compression +// ============================================================================ + +/// An indivisible unit of conversation history. Compression operates at the +/// unit level, so tool call sequences can never be split across boundaries. +#[derive(Debug, Clone)] +enum HistoryUnit { + /// Always-preserved system messages (agent_prompt, scheduled_prompt). + SystemGuard(ChatMessage), + /// A user message. + UserMessage(ChatMessage), + /// An atomic tool call sequence: assistant(tool_calls) + all following + /// tool-result messages. Must be preserved or compressed as a whole. + ToolRound { + assistant: ChatMessage, + results: Vec, + }, + /// A plain assistant text response (no tool_calls). + AssistantText(ChatMessage), +} + +impl HistoryUnit { + /// Estimate tokens for this unit alone. + fn estimate_tokens(&self) -> usize { + match self { + HistoryUnit::SystemGuard(msg) | HistoryUnit::UserMessage(msg) | HistoryUnit::AssistantText(msg) => { + estimate_tokens(std::slice::from_ref(msg)) + } + HistoryUnit::ToolRound { assistant, results } => { + let mut all = vec![assistant.clone()]; + all.extend(results.clone()); + estimate_tokens(&all) + } + } + } + + /// Expand this unit back into a flat list of ChatMessages. + #[allow(dead_code)] + fn into_messages(self) -> Vec { + match self { + HistoryUnit::SystemGuard(msg) + | HistoryUnit::UserMessage(msg) + | HistoryUnit::AssistantText(msg) => vec![msg], + HistoryUnit::ToolRound { assistant, results } => { + let mut msgs = vec![assistant]; + msgs.extend(results); + msgs + } + } + } + + /// Returns true if this unit is a ToolRound. + #[allow(dead_code)] + fn is_tool_round(&self) -> bool { + matches!(self, HistoryUnit::ToolRound { .. }) + } +} + +// ============================================================================ +// Unit parser — one forward pass, O(n) +// ============================================================================ + +fn is_system_guard(msg: &ChatMessage) -> bool { + msg.role == "system" + && (msg.has_system_context(SYSTEM_CONTEXT_AGENT_PROMPT) + || msg.has_system_context(SYSTEM_CONTEXT_SCHEDULED_PROMPT)) +} + +fn is_assistant_with_tool_calls(msg: &ChatMessage) -> bool { + msg.role == "assistant" + && msg + .tool_calls + .as_ref() + .map_or(false, |calls| !calls.is_empty()) +} + +/// Parse a flat message list into atomic units. Orphaned tool results +/// (without a preceding assistant with tool_calls) are silently dropped. +fn parse_to_units(messages: &[ChatMessage]) -> Vec { + let mut units = Vec::new(); + let mut i = 0; + + while i < messages.len() { + let msg = &messages[i]; + + if is_system_guard(msg) { + units.push(HistoryUnit::SystemGuard(msg.clone())); + i += 1; + } else if msg.role == "user" { + units.push(HistoryUnit::UserMessage(msg.clone())); + i += 1; + } else if is_assistant_with_tool_calls(msg) { + let assistant = msg.clone(); + let mut results = Vec::new(); + i += 1; + while i < messages.len() && messages[i].role == "tool" { + results.push(messages[i].clone()); + i += 1; + } + units.push(HistoryUnit::ToolRound { assistant, results }); + } else if msg.role == "assistant" { + units.push(HistoryUnit::AssistantText(msg.clone())); + i += 1; + } else if msg.role == "tool" { + // Orphaned tool result — drop with a warning + tracing::warn!( + tool_call_id = ?msg.tool_call_id, + message_index = i, + "Dropping orphaned tool result during unit parsing" + ); + i += 1; + } else { + // system messages that aren't guards — keep as plain messages + units.push(HistoryUnit::AssistantText(msg.clone())); + i += 1; + } + } + + units +} + +/// Estimate total tokens from a slice of units. +fn estimate_tokens_from_units(units: &[HistoryUnit]) -> usize { + units.iter().map(|u| u.estimate_tokens()).sum() +} + /// Check if a character is CJK (Chinese, Japanese, Korean) fn is_cjk_char(c: char) -> bool { matches!(c, @@ -97,7 +235,7 @@ impl Default for ContextCompressionConfig { pub struct ContextCompressor { config: ContextCompressionConfig, context_window: usize, - /// Threshold ratio to trigger compression (50% of context window) + /// Threshold ratio to trigger compression (70% of context window). threshold_ratio: f64, } @@ -113,6 +251,10 @@ impl ContextCompressor { .clamp(MIN_SUMMARY_CHARS, MAX_SUMMARY_CHARS) } + // ========================================================================= + // Transcript building (shared helpers, unchanged from original) + // ========================================================================= + fn format_transcript_entry(message: &ChatMessage) -> String { let role = match message.role.as_str() { "assistant" => "Assistant", @@ -214,6 +356,86 @@ Be concise, aim for {} characters or less. {} +"#, + target_chars, transcript + ) + } + + // ========================================================================= + // Two-segment compression: heavy (older) + light (newer) prompts + // ========================================================================= + + /// Prompt for the older segment — aggressive compression, keep only essentials. + fn build_heavy_summary_prompt(transcript: &str, target_chars: usize) -> String { + format!( + r#"You are a conversation compaction engine. Aggressively summarize the following OLDER conversation segment. These events happened earlier in the session and will be marked as "较早的操作记录（已压缩）". + +A newer portion (not shown here) follows this segment. The model can use tools to re-read files for exact data, so you do NOT need to preserve full tool outputs. + +=== MUST PRESERVE (keep verbatim or near-verbatim) === +- All file paths (absolute and relative) +- All URLs +- Key numeric results and conclusions (one sentence per result) +- Error messages (exact text if important) +- User preferences and decisions +- The original user task or goal + +=== SHOULD CONDENSE === +- Each tool call → one line summary: "tool_name: 关键参数 → 结果概要" +- Multiple similar operations → group: "对 ./data/ 目录下的 10 个 CSV 文件执行了统计分析" +- Verbose tool output → drop entirely unless it contains a critical result or error + +=== SHOULD DROP === +- Full tool output (hundreds of lines of raw data) +- Repeated attempts of the same operation +- Greetings, filler text, markdown formatting + +Be concise, aim for {} characters or less. Output the summary in Chinese if the original conversation was in Chinese. + +--- + +OLDER SEGMENT (events from earlier in the session): + +{} + +"#, + target_chars, transcript + ) + } + + /// Prompt for the newer segment — lighter compression, keep more detail. + fn build_light_summary_prompt(transcript: &str, target_chars: usize) -> String { + format!( + r#"You are a conversation compaction engine. Lightly summarize the following RECENT conversation segment. These events happened just before the current moment and will be marked as "近期操作记录（轻度压缩，保留了更多细节）". + +An older segment (already summarized separately) precedes this. Keep enough detail so the model can continue the task seamlessly without re-reading files for data that appears in these recent results. + +=== MUST PRESERVE (keep with more detail than a normal summary) === +- All file paths, URLs, and identifiers +- The exact sequence of recent operations (step by step) +- Tool parameters (especially file paths, search queries, command strings) +- Key outputs and results (shortened but keep the substance) +- Error messages (complete, not truncated) +- Current task status and what should happen next + +=== SHOULD CONDENSE === +- Tool outputs → truncate to the most meaningful parts (key data, conclusions, not raw output) +- Long text → keep the essence but not every word +- Repeated similar outputs → note the pattern but keep a representative example + +=== SHOULD DROP === +- Completely irrelevant debug output +- Boilerplate text with no informational value +- Trivial operations that have no bearing on the task + +Be concise, aim for {} characters or less. Output the summary in Chinese if the original conversation was in Chinese. + +--- + +RECENT SEGMENT (events from just before the current moment): + +{} + "#, target_chars, transcript ) @@ -295,12 +517,339 @@ Be concise, aim for {} characters or less. Ok(take_prefix_chars(transcript, target_chars)) } + // ========================================================================= + // Two-segment compression + // ========================================================================= + + /// Find a safe split point in the unit list. The split ensures: + /// - Accumulated tokens up to `ratio` of total are placed in the older segment. + /// - The split never lands on a SystemGuard (they're always preserved). + /// - The split lands on a unit boundary (ToolRound, AssistantText, or UserMessage). + fn find_safe_split_point(&self, units: &[HistoryUnit], ratio: f64) -> usize { + let total_tokens = estimate_tokens_from_units(units); + let target = (total_tokens as f64 * ratio) as usize; + let mut accumulated = 0; + + for (i, unit) in units.iter().enumerate() { + // Never split on a SystemGuard — they must be preserved + if matches!(unit, HistoryUnit::SystemGuard(_)) { + continue; + } + accumulated += unit.estimate_tokens(); + if accumulated >= target { + return i + 1; + } + } + + // Fallback: put everything in the older segment + units.len() + } + + /// Summarize a transcript with a custom prompt builder function. + async fn summarize_with_prompt( + &self, + provider: &dyn LLMProvider, + transcript: &str, + target_chars: usize, + build_prompt: fn(&str, usize) -> String, + ) -> Result { + let request = ChatCompletionRequest { + messages: vec![ + Message::system("You are a helpful assistant."), + Message::user(build_prompt(transcript, target_chars)), + ], + temperature: Some(0.3), + max_tokens: Some(1000), + tools: None, + }; + + let response = provider + .chat(request) + .await + .map_err(|e| AgentError::LlmError(e.to_string()))?; + Ok(response.content) + } + + /// Summarize a segment of units with chunked fallback. + async fn summarize_units_segment( + &self, + provider: &dyn LLMProvider, + messages: &[ChatMessage], + transcript: &str, + target_chars: usize, + build_prompt: fn(&str, usize) -> String, + ) -> Result { + if messages.is_empty() { + return Ok(String::new()); + } + + if char_count(transcript) <= target_chars { + return self + .summarize_with_prompt(provider, transcript, target_chars, build_prompt) + .await; + } + + // Use chunked summarization for oversized transcripts + let target = target_chars.max(1); + let mut layer = Self::chunk_messages_for_summary(messages, target); + + if layer.is_empty() { + layer.push(transcript.to_string()); + } + + for _ in 0..6 { + if layer.len() == 1 && char_count(&layer[0]) <= target { + return self + .summarize_with_prompt(provider, &layer[0], target, build_prompt) + .await; + } + + let per_chunk_target = (target / layer.len().max(1)).max(500).min(target); + let mut summaries = Vec::with_capacity(layer.len()); + for chunk in &layer { + summaries.push( + self.summarize_with_prompt(provider, chunk, per_chunk_target, build_prompt) + .await?, + ); + } + + if summaries.len() == 1 { + let summary = summaries.pop().unwrap_or_default(); + if char_count(&summary) <= target { + return Ok(summary); + } + layer = Self::split_text_chunks(&summary, target); + continue; + } + + let merged = summaries.join("\n\n"); + if char_count(&merged) <= target { + return self + .summarize_with_prompt(provider, &merged, target, build_prompt) + .await; + } + + layer = Self::split_text_chunks(&merged, target); + } + + Ok(take_prefix_chars(transcript, target)) + } + + /// Main entry point for two-segment compression. + /// + /// Splits history into older and newer segments, then summarizes each + /// with LLM using different prompts and budget allocations. The result + /// contains no tool_calls or tool-result messages — only system summaries + /// and the original user message — eliminating any risk of API 400 errors + /// from orphaned tool call sequences. + pub async fn compress_two_segment( + &self, + history: &[ChatMessage], + provider_config: &LLMProviderConfig, + ) -> Result, AgentError> { + let tokens = estimate_tokens(history); + if tokens <= self.threshold() { + tracing::info!( + tokens = tokens, + threshold = self.threshold(), + msg_count = history.len(), + "Two-segment compression not needed (under threshold)" + ); + return Ok(history.to_vec()); + } + + tracing::info!( + tokens = tokens, + threshold = self.threshold(), + msg_count = history.len(), + "Starting two-segment compression" + ); + + // Step 1: Parse into atomic units + let units = parse_to_units(history); + + // Step 2: Separate system guards + user messages from compressible units + let mut system_guards: Vec = Vec::new(); + let mut user_messages: Vec = Vec::new(); + let mut compressible: Vec = Vec::new(); + + for unit in units { + match unit { + HistoryUnit::SystemGuard(msg) => system_guards.push(msg), + HistoryUnit::UserMessage(msg) => user_messages.push(msg), + other => compressible.push(other), + } + } + + // Keep only the latest user message in full; older ones go into compression + let latest_user_msg = user_messages.pop(); + + // Step 3: Find split point in compressible units + let total_compressible = estimate_tokens_from_units(&compressible); + let split_point = if total_compressible > 0 { + self.find_safe_split_point(&compressible, 0.5) + } else { + 0 + }; + let older_units: Vec<&HistoryUnit> = compressible[..split_point].iter().collect(); + let newer_units: Vec<&HistoryUnit> = compressible[split_point..].iter().collect(); + + // Step 4: Build transcripts + let older_messages: Vec = older_units + .iter() + .flat_map(|u| match u { + HistoryUnit::ToolRound { assistant, results } => { + let mut msgs = vec![assistant.clone()]; + msgs.extend(results.clone()); + msgs + } + HistoryUnit::AssistantText(msg) => vec![msg.clone()], + HistoryUnit::UserMessage(msg) => { + // Older user messages go into the compressible transcript + vec![msg.clone()] + } + _ => vec![], + }) + .collect(); + + let newer_messages: Vec = newer_units + .iter() + .flat_map(|u| match u { + HistoryUnit::ToolRound { assistant, results } => { + let mut msgs = vec![assistant.clone()]; + msgs.extend(results.clone()); + msgs + } + HistoryUnit::AssistantText(msg) => vec![msg.clone()], + _ => vec![], + }) + .collect(); + + // Include older user messages in the older transcript + let older_user_msgs: Vec = user_messages + .iter() + .map(|m| { + let mut msg = m.clone(); + msg.role = "user".to_string(); + msg + }) + .collect(); + let all_older_messages: Vec = older_user_msgs + .iter() + .chain(older_messages.iter()) + .cloned() + .collect(); + + let older_transcript = Self::build_transcript(&all_older_messages); + let newer_transcript = Self::build_transcript(&newer_messages); + + // Step 5: Budget allocation + let total_budget = self.config.summary_max_chars; + let older_budget = (total_budget as f64 * OLDER_BUDGET_RATIO) as usize; + let newer_budget = (total_budget as f64 * NEWER_BUDGET_RATIO) as usize; + + // Step 6: Create provider and run both summaries (can be parallel) + let runtime_config = AgentRuntimeConfig::from(provider_config.clone()); + let provider = create_provider(runtime_config.provider) + .map_err(|e| AgentError::ProviderCreation(e.to_string()))?; + + let (older_result, newer_result) = if older_units.is_empty() && newer_units.is_empty() { + (Ok(String::new()), Ok(String::new())) + } else if older_units.is_empty() { + let result = self + .summarize_units_segment( + provider.as_ref(), + &newer_messages, + &newer_transcript, + newer_budget, + Self::build_light_summary_prompt, + ) + .await; + (Ok(String::new()), result) + } else if newer_units.is_empty() { + let result = self + .summarize_units_segment( + provider.as_ref(), + &all_older_messages, + &older_transcript, + older_budget, + Self::build_heavy_summary_prompt, + ) + .await; + (result, Ok(String::new())) + } else { + let older_fut = self.summarize_units_segment( + provider.as_ref(), + &all_older_messages, + &older_transcript, + older_budget, + Self::build_heavy_summary_prompt, + ); + let newer_fut = self.summarize_units_segment( + provider.as_ref(), + &newer_messages, + &newer_transcript, + newer_budget, + Self::build_light_summary_prompt, + ); + tokio::join!(older_fut, newer_fut) + }; + + let older_summary = older_result?; + let newer_summary = newer_result?; + + // Step 7: Assemble compressed history + let mut compressed: Vec = Vec::with_capacity(4); + + // System guards first + compressed.extend(system_guards); + + // Latest user message + if let Some(user_msg) = latest_user_msg { + compressed.push(user_msg); + } + + // Heavy compression summary (older) + if !older_summary.is_empty() { + compressed.push(ChatMessage::system_with_context( + format!("## 较早的操作记录（已压缩）\n\n{}", older_summary), + Some(SYSTEM_CONTEXT_HISTORY_COMPACTION_OLDER.to_string()), + )); + } + + // Light compression summary (newer) + if !newer_summary.is_empty() { + compressed.push(ChatMessage::system_with_context( + format!( + "## 近期操作记录（轻度压缩，保留了更多细节）\n\n{}\n\n---\n以上为近期操作记录。如需精确数据，可使用工具重新读取相关文件。", + newer_summary + ), + Some(SYSTEM_CONTEXT_HISTORY_COMPACTION_NEWER.to_string()), + )); + } + + tracing::info!( + original_tokens = tokens, + final_tokens = estimate_tokens(&compressed), + final_msg_count = compressed.len(), + older_units = older_units.len(), + newer_units = newer_units.len(), + "Two-segment compression completed" + ); + + Ok(compressed) + } + + // ========================================================================= + // Legacy methods (kept for backward compatibility during migration) + // ========================================================================= + /// Create a new compressor with the given context window size. pub fn new(context_window: usize) -> Self { Self { config: ContextCompressionConfig::default(), context_window, - threshold_ratio: 0.5, + threshold_ratio: DEFAULT_THRESHOLD_RATIO, } } @@ -323,11 +872,11 @@ Be concise, aim for {} characters or less. Self { config, context_window, - threshold_ratio: 0.5, + threshold_ratio: DEFAULT_THRESHOLD_RATIO, } } - /// Get the compression threshold in tokens. + /// Get the compression threshold in tokens (70% of context window). fn threshold(&self) -> usize { (self.context_window as f64 * self.threshold_ratio) as usize } @@ -707,7 +1256,7 @@ mod tests { #[test] fn test_threshold() { let compressor = ContextCompressor::new(128_000); - assert_eq!(compressor.threshold(), 64_000); + assert_eq!(compressor.threshold(), 89_600); // 70% of 128_000 } #[test] diff --git a/src/gateway/compaction.rs b/src/gateway/compaction.rs index 3f753c3..24465de 100644 --- a/src/gateway/compaction.rs +++ b/src/gateway/compaction.rs @@ -6,97 +6,60 @@ use crate::agent::AgentError; use super::session::Session; +/// Run two-segment history compression synchronously. +/// +/// Unlike the previous background approach (tokio::spawn), this holds the +/// session lock during the LLM calls (2–5 seconds). Since the agent loop +/// has already finished by this point there is no response-time impact, and +/// the synchronous guarantee means the next execution always starts with +/// freshly compacted history. pub(crate) async fn schedule_background_history_compaction( session: Arc>, chat_id: impl Into, ) -> Result<(), AgentError> { let chat_id = chat_id.into(); - let snapshot = { - let mut session_guard = session.lock().await; - let session_record = session_guard.ensure_persistent_session(&chat_id)?; - session_guard.ensure_chat_loaded(&chat_id)?; + let mut session_guard = session.lock().await; + session_guard.ensure_persistent_session(&chat_id)?; + session_guard.ensure_chat_loaded(&chat_id)?; - let history = session_guard.get_or_create_history(&chat_id).clone(); - let compressor = session_guard.compressor().clone(); - if !compressor.should_compress(&history) { - return Ok(()); - } + let history = session_guard.get_or_create_history(&chat_id).clone(); + let compressor = session_guard.compressor().clone(); - if !session_guard.try_start_background_compaction(&chat_id) { - return Ok(()); - } + if !compressor.should_compress(&history) { + return Ok(()); + } - ( - session_guard.store(), - session_guard.persistent_session_id(&chat_id), - session_record.message_count, - history, - compressor, - session_guard.provider_config().clone(), - ) - }; + let store = session_guard.store(); + let session_id = session_guard.persistent_session_id(&chat_id); + let provider_config = session_guard.provider_config().clone(); - let ( - store, - session_id, - snapshot_end_seq, - history, - compressor, - provider_config, - ) = snapshot; - let session_for_task = session.clone(); - let chat_id_for_task = chat_id.clone(); + tracing::info!( + chat_id = %chat_id, + msg_count = history.len(), + "Starting synchronous two-segment compression" + ); - tokio::spawn(async move { - tracing::info!(chat_id = %chat_id_for_task, snapshot_end_seq, "Starting background history compaction"); + // Synchronous compression — holds lock during LLM calls. + // compress_two_segment guarantees the result contains no tool_calls, + // so there is no risk of orphaned tool call sequences. + let compressed = compressor + .compress_two_segment(&history, &provider_config) + .await?; - let compaction_result = compressor - .build_compaction_plan(&history, &provider_config) - .await; - let mut committed = false; + // Replace the entire history with the compressed result. + // Since we hold the lock, no concurrent modifications can occur. + store + .replace_active_history(&session_id, &compressed) + .map_err(|e| AgentError::Other(format!("replace_active_history error: {}", e)))?; - match compaction_result { - Ok(Some(plan)) => match store.compact_active_history( - &session_id, - snapshot_end_seq, - &plan.preserved_system_messages, - &plan.summary_message, - &plan.preserved_messages, - ) { - Ok(true) => { - committed = true; - tracing::info!( - chat_id = %chat_id_for_task, - snapshot_end_seq, - compressed_turns = plan.compressed_turns, - preserved_turns = plan.preserved_turns, - "Background history compaction committed" - ); - } - Ok(false) => { - tracing::info!(chat_id = %chat_id_for_task, snapshot_end_seq, "Background history compaction skipped due to stale snapshot"); - } - Err(error) => { - tracing::warn!(chat_id = %chat_id_for_task, error = %error, "Background history compaction commit failed"); - } - }, - Ok(None) => { - tracing::debug!(chat_id = %chat_id_for_task, "Background history compaction not needed after snapshot analysis"); - } - Err(error) => { - tracing::warn!(chat_id = %chat_id_for_task, error = %error, "Background history compaction build failed"); - } - } + tracing::info!( + chat_id = %chat_id, + compressed_msg_count = compressed.len(), + "Two-segment compression committed" + ); - let mut session_guard = session_for_task.lock().await; - if committed { - if let Err(error) = session_guard.reload_chat_history(&chat_id_for_task) { - tracing::warn!(chat_id = %chat_id_for_task, error = %error, "Failed to reload history after background compaction"); - } - } - session_guard.finish_background_compaction(&chat_id_for_task); - }); + session_guard.reload_chat_history(&chat_id)?; Ok(()) } diff --git a/src/gateway/session.rs b/src/gateway/session.rs index 8cdb2a6..f98f2e8 100644 --- a/src/gateway/session.rs +++ b/src/gateway/session.rs @@ -526,10 +526,12 @@ impl Session { &self.compressor } + #[allow(dead_code)] pub(crate) fn try_start_background_compaction(&mut self, chat_id: &str) -> bool { self.history.try_start_background_compaction(chat_id) } + #[allow(dead_code)] pub(crate) fn finish_background_compaction(&mut self, chat_id: &str) { self.history.finish_background_compaction(chat_id); } diff --git a/src/gateway/session_history.rs b/src/gateway/session_history.rs index d367f4f..bebda25 100644 --- a/src/gateway/session_history.rs +++ b/src/gateway/session_history.rs @@ -235,10 +235,12 @@ impl SessionHistory { Ok(()) } + #[allow(dead_code)] pub(crate) fn try_start_background_compaction(&mut self, chat_id: &str) -> bool { self.compression_in_flight.insert(chat_id.to_string()) } + #[allow(dead_code)] pub(crate) fn finish_background_compaction(&mut self, chat_id: &str) { self.compression_in_flight.remove(chat_id); } diff --git a/src/storage/mod.rs b/src/storage/mod.rs index 138a338..efcd456 100644 --- a/src/storage/mod.rs +++ b/src/storage/mod.rs @@ -814,6 +814,59 @@ impl SessionStore { Ok(true) } + /// Replace the entire active history for a session. + /// + /// This is a simpler alternative to `compact_active_history` for when the + /// compressor has already produced a complete, validated message list + /// (e.g. two-segment compression). It replaces all existing messages + /// with the new list in a single transaction. + pub fn replace_active_history( + &self, + session_id: &str, + messages: &[ChatMessage], + ) -> Result<(), StorageError> { + let conn = self.pool.get()?; + let tx = conn.unchecked_transaction()?; + let now = current_timestamp(); + + // Delete all existing messages for this session + tx.execute( + "DELETE FROM messages WHERE session_id = ?1", + params![session_id], + )?; + + // Insert new messages with sequential seq numbers + let mut active_user_turn_count = 0_i64; + for (i, message) in messages.iter().enumerate() { + let seq = (i + 1) as i64; + if message.role == "user" { + active_user_turn_count += 1; + } + insert_message_with_seq(&tx, session_id, seq, message)?; + } + + tx.execute( + " + UPDATE sessions + SET message_count = ?2, + user_turn_count = ?3, + updated_at = ?4, + last_active_at = ?4, + archived_at = NULL + WHERE id = ?1 AND deleted_at IS NULL + ", + params![ + session_id, + messages.len() as i64, + active_user_turn_count, + now, + ], + )?; + + tx.commit()?; + Ok(()) + } + pub fn mark_agent_prompt_reinjected(&self, session_id: &str) -> Result<(), StorageError> { let now = current_timestamp(); let conn = self.pool.get()?; diff --git a/src/storage/ports.rs b/src/storage/ports.rs index 3f03f9a..0a7aeba 100644 --- a/src/storage/ports.rs +++ b/src/storage/ports.rs @@ -54,6 +54,15 @@ pub trait ConversationRepository: Send + Sync + 'static { summary_message: &ChatMessage, preserved_messages: &[ChatMessage], ) -> Result; + + /// Replace the entire active history for a session with new messages. + /// Used when the compressor has already produced a complete, validated + /// message list (e.g. two-segment compression). + fn replace_active_history( + &self, + session_id: &str, + messages: &[ChatMessage], + ) -> Result<(), StorageError>; } pub trait PromptInjectionRepository: Send + Sync + 'static { @@ -235,6 +244,14 @@ impl ConversationRepository for super::SessionStore { preserved_messages, ) } + + fn replace_active_history( + &self, + session_id: &str, + messages: &[ChatMessage], + ) -> Result<(), StorageError> { + super::SessionStore::replace_active_history(self, session_id, messages) + } } impl PromptInjectionRepository for super::SessionStore {