From 6252b600cd6902634921a46c63c6cfe3bd116774 Mon Sep 17 00:00:00 2001 From: oudecheng <13802883547@139.com> Date: Wed, 5 Aug 2026 21:56:41 +0800 Subject: [PATCH] =?UTF-8?q?refactor(compression):=20=E9=87=8D=E6=9E=84?= =?UTF-8?q?=E4=B8=8A=E4=B8=8B=E6=96=87=E5=8E=8B=E7=BC=A9=E4=B8=BA=E4=B8=A4?= =?UTF-8?q?=E9=98=B6=E6=AE=B5+=E5=8F=8C=E9=98=88=E5=80=BC+=E4=B8=89?= =?UTF-8?q?=E6=AE=B5=E4=BF=9D=E7=95=99?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 压缩时机从 process 整体返回后提前到每轮工具调用完成后，更及时控制上下文大小。 两阶段压缩： - 工程化压缩（70% 阈值触发）：截断非子代理 tool 结果到 100 token，仅改内存不影响 DB - LLM 压缩（50% 阈值触发）：工程化压缩后仍超 50% 才调 LLM，避免不必要的 LLM 调用 LLM 压缩三段保留策略： - 保留最旧 5 个 unit 和最新 5 个 unit 原样 - 中间段用 LLM 生成摘要（heavy prompt） - SystemGuard 永远保留不计入配额 - ToolRound 原子性由 parse_to_units 保证，切分在 unit 边界 工程化压缩： - truncate_tool_results_in_place 截断 role="tool" 且 tool_name!="task" 的消息 - 子代理返回（tool_name="task"）保持原样 - char_indices 确保 UTF-8 字符边界安全，截断后追加 "...[已截断]" 标记 AgentLoop 集成： - 新增 CompactionSink trait，LLM 压缩结果通过 sink 回写 DB - AgentProcessResult 新增 compaction_performed 和 engineering_compaction_applied 标记 - finalize_result 据此决定是否 reload DB 历史（避免重复 append）和跳过兜底压缩 错误恢复： - LLM 压缩失败降级为仅工程化压缩，不中断 agent loop - sink.compact 失败记 error 日志继续使用内存压缩结果，DB 下次 process 会重新触发压缩 清理旧的 two-segment 差异化压缩代码（OLDER_BUDGET_RATIO、find_safe_split_point、 build_light_summary_prompt 等），新增 unit_to_messages 辅助函数。 --- src/agent/agent_loop.rs | 162 +++++++++++- src/agent/context_compressor.rs | 456 ++++++++++++++------------------ src/agent/mod.rs | 2 +- src/gateway/agent_factory.rs | 41 +-- src/gateway/execution.rs | 109 +++++++- src/tools/task/runtime.rs | 4 +- 6 files changed, 484 insertions(+), 290 deletions(-) diff --git a/src/agent/agent_loop.rs b/src/agent/agent_loop.rs index 9f47d01..e314f1d 100644 --- a/src/agent/agent_loop.rs +++ b/src/agent/agent_loop.rs @@ -1,4 +1,5 @@ use crate::agent::AgentRuntimeConfig; +use crate::agent::context_compressor::ContextCompressor; use crate::agent::{SystemPromptContext, SystemPromptProvider}; use crate::bus::ChatMessage; use crate::bus::message::ToolMessageState; @@ -689,12 +690,36 @@ pub struct AgentLoop { /// 包装在 Mutex 中以支持 interior mutability —— /// watch::Receiver::changed() 需要 &mut self，但 process() 持有 &self。 cancel_token: Option>>, + /// 上下文压缩器（可选）。配置后会在每轮工具调用完成后按双阈值执行压缩： + /// 1) 真实 prompt_tokens > 70% context_window → 进入压缩流程 + /// 2) 工程化压缩（截断 tool 结果到 100 token，仅改内存） + /// 3) estimate_tokens > 50% context_window → 调 LLM 三段压缩；否则跳过 + compressor: Option>, } #[derive(Debug, Clone)] pub struct AgentProcessResult { pub final_response: ChatMessage, pub emitted_messages: Vec, + /// 本轮 process 是否触发过 LLM 压缩。 + /// true 表示发生过 LLM 压缩（compaction_sink 已被调用）， + /// 调用方据此决定是否需要刷新 DB 中的会话历史。 + pub compaction_performed: bool, + /// 本轮 process 是否触发过工程化压缩（70% 阈值命中，截断 tool 结果）。 + /// 调用方据此跳过兜底 LLM 压缩——因为 in-loop 已判断工程化压缩足够 + /// （或 LLM 压缩失败已降级），兜底基于未压缩历史的判断会不准确。 + pub engineering_compaction_applied: bool, +} + +/// 压缩结果接收端：当 AgentLoop 在 process 内部完成 LLM 压缩后， +/// 通过此 trait 把压缩后的消息序列回写给外部持久化层（DB）。 +/// +/// 工程化压缩（截断 tool 结果）只改内存，不触发此 sink； +/// 只有 LLM 压缩（三段策略）才会触发。 +#[async_trait] +pub trait CompactionSink: Send + Sync + 'static { + /// 用压缩后的消息替换 DB 中该 topic 的历史。 + async fn compact(&self, compressed: &[ChatMessage]) -> Result<(), AgentError>; } #[async_trait] @@ -820,6 +845,7 @@ impl AgentLoop { observer: None, emitted_message_handler: None, cancel_token: None, + compressor: None, max_iterations, }) } @@ -843,6 +869,7 @@ impl AgentLoop { observer: None, emitted_message_handler: None, cancel_token: None, + compressor: None, max_iterations, }) } @@ -867,6 +894,7 @@ impl AgentLoop { observer: None, emitted_message_handler: None, cancel_token: None, + compressor: None, max_iterations, }) } @@ -895,6 +923,7 @@ impl AgentLoop { observer: None, emitted_message_handler: None, cancel_token: None, + compressor: None, max_iterations, }) } @@ -925,6 +954,12 @@ impl AgentLoop { self } + /// 注入上下文压缩器。配置后，`process` 会在每轮工具调用完成后按双阈值执行压缩。 + pub fn with_compressor(mut self, compressor: Option>) -> Self { + self.compressor = compressor; + self + } + pub fn tools(&self) -> &Arc { &self.tools } @@ -940,10 +975,14 @@ impl AgentLoop { /// # 参数 /// - `messages`: 会话历史消息 /// - `system_prompt_context`: 系统提示词上下文（用于动态注入，可选） + /// - `compaction_sink`: 压缩结果回写端（可选）。配置 `compressor` 后， + /// 当 LLM 压缩被触发时通过此 sink 把压缩结果持久化到 DB。 + /// 传 None 则即使配置了 compressor 也只改内存不回写。 pub async fn process( &self, mut messages: Vec, system_prompt_context: Option<&SystemPromptContext>, + compaction_sink: Option<&dyn CompactionSink>, ) -> Result { #[cfg(debug_assertions)] tracing::debug!( @@ -952,6 +991,14 @@ impl AgentLoop { "Starting agent process" ); + // 跟踪本轮 process 是否触发过 LLM 压缩。 + // 工程化压缩（仅截断内存中的 tool 结果）不算 —— 那不影响 DB 状态。 + let mut compaction_performed = false; + // 跟踪本轮 process 是否触发过工程化压缩（70% 阈值命中）。 + // finalize_result 据此跳过兜底 LLM 压缩——因为 in-loop 已判断工程化压缩足够 + // （或 LLM 压缩失败已降级），兜底基于未压缩历史的判断会不准确。 + let mut engineering_compaction_applied = false; + // Sanitize: remove any trailing incomplete tool call sequences // that may have been persisted before a process interruption. { @@ -1151,6 +1198,8 @@ impl AgentLoop { return Ok(AgentProcessResult { final_response: assistant_message, emitted_messages, + compaction_performed, + engineering_compaction_applied, }); } } @@ -1180,7 +1229,7 @@ impl AgentLoop { // If no tool calls, this is the final response if response.tool_calls.is_empty() { - let result = self + let mut result = self .build_final_response( response, &streaming_message_id, @@ -1188,6 +1237,8 @@ impl AgentLoop { &mut emitted_messages, ) .await; + result.compaction_performed = compaction_performed; + result.engineering_compaction_applied = engineering_compaction_applied; return Ok(result); } @@ -1271,6 +1322,89 @@ impl AgentLoop { ) .await; + // === 两阶段压缩（工具调用完成后） === + // 仅当配置了 compressor 时执行。compaction_sink 控制是否回写 DB。 + if let Some(compressor) = &self.compressor { + // 阶段 1：用最近一次 LLM 调用的真实 prompt_tokens 判断 70% 触发阈值 + // 提取 u32（Copy）避免持有 messages 的不可变借用 + let last_prompt_tokens = messages + .iter() + .rev() + .find(|m| m.role == "assistant") + .and_then(|m| m.usage.as_ref()) + .map(|u| u.prompt_tokens); + + if let Some(prompt_tokens) = last_prompt_tokens { + if compressor.should_compress_by_usage(prompt_tokens) { + // 阶段 1a：工程化压缩（截断非子代理 tool 结果到 100 token，仅改内存） + crate::agent::context_compressor::truncate_tool_results_in_place( + &mut messages, + 100, + ); + engineering_compaction_applied = true; + tracing::info!( + iteration, + prompt_tokens, + threshold = compressor.threshold(), + "Engineering compaction applied (tool results truncated to 100 tokens)" + ); + + // 阶段 1b：重新估算，判断是否需要 LLM 压缩（50% 阈值） + let estimated = + crate::agent::context_compressor::estimate_tokens(&messages); + if estimated > compressor.llm_compaction_threshold() { + tracing::info!( + iteration, + estimated_tokens = estimated, + llm_threshold = compressor.llm_compaction_threshold(), + "LLM compaction triggered (still above 50% after engineering compaction)" + ); + // LLM 压缩失败时降级为仅工程化压缩，不中断 agent loop + match compressor + .compress_two_segment_with_provider( + &messages, + self.provider.as_ref(), + ) + .await + { + Ok(compressed) => { + // sink 失败时记日志但不中断——内存已压缩，DB 未更新 + // 下次 process 从 DB 加载时会重新触发压缩 + if let Some(sink) = compaction_sink { + if let Err(e) = sink.compact(&compressed).await { + tracing::error!( + error = %e, + iteration, + "CompactionSink compact failed; \ + in-memory messages still replaced, DB will be re-compacted next round" + ); + } + } + messages = compressed; + compaction_performed = true; + } + Err(e) => { + tracing::warn!( + error = %e, + iteration, + "LLM compaction failed; \ + falling back to engineering-only compaction (in-memory truncated messages retained)" + ); + // 不设置 compaction_performed，messages 保持工程化压缩后的状态 + } + } + } else { + tracing::info!( + iteration, + estimated_tokens = estimated, + llm_threshold = compressor.llm_compaction_threshold(), + "Engineering compaction sufficient (under 50%), skipping LLM compaction" + ); + } + } + } + } + // Loop continues to next iteration with updated messages // PendingUserAction 工具的结果已在上方加入 messages， // 模型将在下一轮看到完整的终端输出并生成智能回复 @@ -1278,14 +1412,18 @@ impl AgentLoop { tracing::debug!( iteration, message_count = messages.len(), + compaction_performed, "Tool execution complete, continuing to next iteration" ); } // Max iterations reached - request final summary from LLM - Ok(self + let mut result = self .run_final_summary(&mut messages, system_prompt_context, &mut emitted_messages) - .await) + .await; + result.compaction_performed = compaction_performed; + result.engineering_compaction_applied = engineering_compaction_applied; + Ok(result) } /// 等待取消信号。若未配置 cancel_token，永远不返回。 @@ -1399,6 +1537,8 @@ impl AgentLoop { AgentProcessResult { final_response: assistant_message, emitted_messages: std::mem::take(emitted_messages), + compaction_performed: false, + engineering_compaction_applied: false, } } @@ -1512,6 +1652,8 @@ impl AgentLoop { return AgentProcessResult { final_response: assistant_message, emitted_messages: std::mem::take(emitted_messages), + compaction_performed: false, + engineering_compaction_applied: false, }; } Err(e) => { @@ -1557,6 +1699,8 @@ impl AgentLoop { return AgentProcessResult { final_response: final_message, emitted_messages: std::mem::take(emitted_messages), + compaction_performed: false, + engineering_compaction_applied: false, }; } } @@ -1582,6 +1726,8 @@ impl AgentLoop { AgentProcessResult { final_response: assistant_message, emitted_messages, + compaction_performed: false, + engineering_compaction_applied: false, } } @@ -2811,7 +2957,9 @@ mod tests { observer: None, emitted_message_handler: None, cancel_token: None, + compressor: None, max_iterations: 1, + // test helper 不需要 tracking 字段 } } @@ -2826,7 +2974,7 @@ mod tests { ])), 3); let result = loop_instance - .process(vec![ChatMessage::user("hello")], None) + .process(vec![ChatMessage::user("hello")], None, None) .await .unwrap(); @@ -2844,7 +2992,7 @@ mod tests { ])), 3); let result = loop_instance - .process(vec![ChatMessage::user("hello")], None) + .process(vec![ChatMessage::user("hello")], None, None) .await .unwrap(); @@ -2863,7 +3011,7 @@ mod tests { ])), 0); let result = loop_instance - .process(vec![ChatMessage::user("hello")], None) + .process(vec![ChatMessage::user("hello")], None, None) .await .unwrap(); @@ -2881,7 +3029,7 @@ mod tests { ])), 1); let result = loop_instance - .process(vec![ChatMessage::user("hello")], None) + .process(vec![ChatMessage::user("hello")], None, None) .await .unwrap(); diff --git a/src/agent/context_compressor.rs b/src/agent/context_compressor.rs index 0605f80..52506d5 100644 --- a/src/agent/context_compressor.rs +++ b/src/agent/context_compressor.rs @@ -20,9 +20,8 @@ pub const SYSTEM_CONTEXT_HISTORY_COMPACTION_NEWER: &str = "history_compaction_ne /// 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; +/// LLM 压缩阈值比例：工程化压缩后仍超过此比例才调 LLM +const LLM_COMPACTION_THRESHOLD_RATIO: f64 = 0.5; // ============================================================================ // HistoryUnit — atomic message units for compression @@ -46,22 +45,6 @@ enum HistoryUnit { 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) - } - } - } -} - // ============================================================================ // Unit parser — one forward pass, O(n) // ============================================================================ @@ -125,9 +108,19 @@ fn parse_to_units(messages: &[ChatMessage]) -> Vec { 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() +/// Flatten a HistoryUnit back into its constituent ChatMessage(s). +/// ToolRound yields [assistant, results...]; others yield a single message. +fn unit_to_messages(unit: &HistoryUnit) -> Vec { + match unit { + HistoryUnit::ToolRound { assistant, results } => { + let mut msgs = vec![assistant.clone()]; + msgs.extend(results.clone()); + msgs + } + HistoryUnit::AssistantText(msg) + | HistoryUnit::UserMessage(msg) + | HistoryUnit::SystemGuard(msg) => vec![msg.clone()], + } } /// Check if a character is CJK (Chinese, Japanese, Korean) @@ -176,6 +169,72 @@ pub fn estimate_tokens(messages: &[ChatMessage]) -> usize { * TOKEN_ESTIMATE_SAFETY_MULTIPLIER) as usize } +/// 估算纯文本的 token 数（不含消息级 JSON 开销，仅内容）。 +fn estimate_text_tokens(text: &str) -> usize { + let mut cjk = 0usize; + let mut other = 0usize; + for ch in text.chars() { + if is_cjk_char(ch) { + cjk += 1; + } else { + other += 1; + } + } + let content_tokens = + (cjk as f64 / CJK_CHARS_PER_TOKEN) + (other as f64 / OTHER_CHARS_PER_TOKEN); + (content_tokens * TOKEN_ESTIMATE_SAFETY_MULTIPLIER) as usize +} + +/// 将文本截断到约 max_tokens，确保 UTF-8 字符边界安全。 +/// 截断后追加 "\n...[已截断]" 标记。 +fn truncate_to_token_limit(text: &str, max_tokens: usize) -> String { + let mut token_count = 0.0f64; + let mut cutoff_byte = text.len(); + + for (byte_idx, ch) in text.char_indices() { + let char_tokens = if is_cjk_char(ch) { 0.5 } else { 0.25 }; + token_count += char_tokens; + if token_count >= max_tokens as f64 { + cutoff_byte = byte_idx; + break; + } + } + + if cutoff_byte >= text.len() { + return text.to_string(); + } + + let mut result = text[..cutoff_byte].to_string(); + result.push_str("\n...[已截断]"); + result +} + +/// 工程化压缩：将非子代理的 tool 结果截断到约 max_tokens。 +/// +/// 规则： +/// - role="tool" 且 tool_name != "task" 的消息，content 截断到约 max_tokens +/// - role="tool" 且 tool_name == "task"（子代理返回）的消息，保持原样 +/// - 其他 role 的消息不受影响 +/// - 截断时用 char_indices 确保 UTF-8 字符边界安全 +/// - 截断后追加 "\n...[已截断]" 标记 +/// +/// **只修改传入的 messages，不涉及 DB 操作。** +pub fn truncate_tool_results_in_place(messages: &mut [ChatMessage], max_tokens: usize) { + for msg in messages.iter_mut() { + if msg.role != "tool" { + continue; + } + if msg.tool_name.as_deref() == Some("task") { + continue; + } + let estimated = estimate_text_tokens(&msg.content); + if estimated <= max_tokens { + continue; + } + msg.content = truncate_to_token_limit(&msg.content, max_tokens); + } +} + /// Configuration for context compression. #[derive(Debug, Clone)] pub struct ContextCompressionConfig { @@ -377,44 +436,6 @@ 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 ) @@ -497,33 +518,9 @@ RECENT SEGMENT (events from just before the current moment): } // ========================================================================= - // Two-segment compression + // Three-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, @@ -614,206 +611,165 @@ RECENT SEGMENT (events from just before the current moment): Ok(take_prefix_chars(transcript, target)) } - /// Main entry point for two-segment compression. + /// Main entry point for three-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 + /// Preserves the oldest 5 units and newest 5 units in full, then + /// summarizes the middle segment with LLM. The result maintains + /// ToolRound atomicity and the middle summary is a pure system message + /// with no tool_calls — 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 mut truncated: Vec = history.to_vec(); + truncate_tool_results_in_place(&mut truncated, 100); + + let tokens = estimate_tokens(&truncated); + if tokens <= self.threshold() { + return Ok(truncated); + } + + let provider = create_provider(AgentRuntimeConfig::from(provider_config.clone()).provider) + .map_err(|e| AgentError::ProviderCreation(e.to_string()))?; + self.compress_two_segment_inner(&truncated, provider.as_ref()) + .await + } + + /// 复用调用方已有的 provider 实例（用于 AgentLoop 内，避免重复创建 provider）。 + /// + /// **注意**：此方法假设调用方已对 `history` 做过工程化压缩 + /// （`truncate_tool_results_in_place`），因此内部不再重复截断。 + /// 这样可避免二次截断导致 `\n...[已截断]` 标记累积和内容损失。 + pub async fn compress_two_segment_with_provider( + &self, + history: &[ChatMessage], + provider: &dyn LLMProvider, ) -> 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()); } + self.compress_two_segment_inner(history, provider).await + } + + /// 三段压缩核心逻辑：保留最旧5 + 最新5 unit，中间段用 LLM 压缩。 + /// + /// 策略： + /// - SystemGuard 永远保留在头部（不计入 5 条配额） + /// - 可压缩单元（UserMessage / AssistantText / ToolRound）按时间顺序： + /// - 最旧 PRESERVE_COUNT 个 unit 原样保留 + /// - 最新 PRESERVE_COUNT 个 unit 原样保留 + /// - 中间段用 LLM 生成摘要（system 消息，无 tool_calls） + /// - ToolRound 原子性由 parse_to_units 保证，切分在 unit 边界 + /// - 中间段摘要为纯文本 system 消息，符合 API 提交要求 + async fn compress_two_segment_inner( + &self, + history: &[ChatMessage], + provider: &dyn LLMProvider, + ) -> Result, AgentError> { + const PRESERVE_COUNT: usize = 5; + let tokens = estimate_tokens(history); tracing::info!( tokens = tokens, threshold = self.threshold(), msg_count = history.len(), - "Starting two-segment compression" + preserve_count = PRESERVE_COUNT, + "Starting three-segment compression (preserve oldest 5 + newest 5)" ); - // Step 1: Parse into atomic units let units = parse_to_units(history); - // Step 2: Separate system guards + user messages from compressible units + // Step 1: Separate SystemGuard (always preserved, not counted in 5) 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 2: If compressible units are too few, skip LLM compression + if compressible.len() <= PRESERVE_COUNT * 2 { + tracing::info!( + compressible_count = compressible.len(), + preserve_threshold = PRESERVE_COUNT * 2, + "Too few compressible units, skipping LLM compaction" + ); + let mut result = system_guards; + for unit in &compressible { + result.extend(unit_to_messages(unit)); + } + return Ok(result); + } - // 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) + // Step 3: Three-segment split + let split = compressible.len() - PRESERVE_COUNT; + let oldest_units = &compressible[..PRESERVE_COUNT]; + let newest_units = &compressible[split..]; + let middle_units = &compressible[PRESERVE_COUNT..split]; + + // Step 4: Build middle segment messages and transcript + let middle_messages: Vec = middle_units + .iter() + .flat_map(unit_to_messages) + .collect(); + let middle_transcript = Self::build_transcript(&middle_messages); + + // Step 5: Summarize middle segment with LLM (heavy prompt) + let budget = self.config.summary_max_chars; + let middle_summary = if middle_messages.is_empty() { + String::new() } 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.summarize_units_segment( + provider, + &middle_messages, + &middle_transcript, + 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) + ) + .await? }; - let older_summary = older_result?; - let newer_summary = newer_result?; - - // Step 7: Assemble compressed history - let mut compressed: Vec = Vec::with_capacity(4); + // Step 6: Assemble compressed history + // [SystemGuards] + [oldest 5 units raw] + [middle summary] + [newest 5 units raw] + let mut compressed: Vec = Vec::with_capacity( + system_guards.len() + middle_messages.len() + 1 + middle_messages.len(), + ); // System guards first compressed.extend(system_guards); - // Latest user message - if let Some(user_msg) = latest_user_msg { - compressed.push(user_msg); + // Oldest PRESERVE_COUNT units (raw) + for unit in oldest_units { + compressed.extend(unit_to_messages(unit)); } - // Heavy compression summary (older) - if !older_summary.is_empty() { + // Middle segment summary (system message, no tool_calls) + if !middle_summary.is_empty() { compressed.push(ChatMessage::system_with_context( - format!("## 较早的操作记录（已压缩）\n\n{}", older_summary), + format!("## 较早的操作记录（已压缩）\n\n{}", middle_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()), - )); + // Newest PRESERVE_COUNT units (raw) + for unit in newest_units { + compressed.extend(unit_to_messages(unit)); } tracing::info!( original_tokens = tokens, + original_msg_count = history.len(), final_tokens = estimate_tokens(&compressed), final_msg_count = compressed.len(), - older_units = older_units.len(), - newer_units = newer_units.len(), - "Two-segment compression completed" + oldest_units = PRESERVE_COUNT, + newest_units = PRESERVE_COUNT, + middle_units = middle_units.len(), + "Three-segment compression completed" ); Ok(compressed) @@ -856,7 +812,7 @@ RECENT SEGMENT (events from just before the current moment): } /// Get the compression threshold in tokens (70% of context window). - fn threshold(&self) -> usize { + pub fn threshold(&self) -> usize { (self.context_window as f64 * self.threshold_ratio) as usize } @@ -864,6 +820,16 @@ RECENT SEGMENT (events from just before the current moment): estimate_tokens(history) > self.threshold() } + /// 触发阈值（70%）：用真实 prompt_tokens 判断是否进入压缩流程。 + pub fn should_compress_by_usage(&self, prompt_tokens: u32) -> bool { + (prompt_tokens as usize) > self.threshold() + } + + /// LLM 压缩阈值（50%）：工程化压缩后用 estimate_tokens 判断是否需要 LLM 压缩。 + pub fn llm_compaction_threshold(&self) -> usize { + (self.context_window as f64 * LLM_COMPACTION_THRESHOLD_RATIO) as usize + } + fn user_turn_ranges(&self, history: &[ChatMessage]) -> Vec { let user_indices: Vec = history .iter() @@ -1394,26 +1360,6 @@ mod tests { assert!(units.is_empty()); } - #[test] - fn test_find_safe_split_point() { - let compressor = ContextCompressor::new(100_000); - // Build a list of units with known token sizes - let messages: Vec = (0..10) - .map(|i| ChatMessage::assistant(&format!("message content number {}", i))) - .collect(); - let units = parse_to_units(&messages); - - // All units are AssistantText, split at 50% token ratio - let split = compressor.find_safe_split_point(&units, 0.5); - // Should split somewhere in the middle (not 0, not len()) - assert!( - split > 0 && split < units.len(), - "split {} should be between 0 and {}", - split, - units.len() - ); - } - #[test] fn test_compress_two_segment_no_tool_calls_in_output() { // This test verifies the critical invariant: diff --git a/src/agent/mod.rs b/src/agent/mod.rs index ce448bc..853ff20 100644 --- a/src/agent/mod.rs +++ b/src/agent/mod.rs @@ -4,7 +4,7 @@ pub mod runtime_config; pub mod system_prompt; pub use agent_loop::{ - AgentError, AgentLoop, AgentProcessResult, EmittedMessageHandler, + AgentError, AgentLoop, AgentProcessResult, CompactionSink, EmittedMessageHandler, PersistingEmittedMessageHandler, SkillProvider, }; pub use context_compressor::ContextCompressor; diff --git a/src/gateway/agent_factory.rs b/src/gateway/agent_factory.rs index 11b3e30..56282b3 100644 --- a/src/gateway/agent_factory.rs +++ b/src/gateway/agent_factory.rs @@ -1,5 +1,6 @@ use std::sync::Arc; +use crate::agent::context_compressor::ContextCompressor; use crate::agent::{AgentError, AgentLoop, CompositeSystemPromptProvider, SystemPromptProvider}; use crate::config::{LLMProviderConfig, ModelResolver}; use crate::domain::CapabilityPolicy; @@ -200,7 +201,7 @@ impl AgentFactory { }; AgentLoop::with_tools_and_system_prompt_provider( - effective_provider_config, + effective_provider_config.clone(), tools, system_prompt_provider, Some(self.skills.clone()), @@ -210,22 +211,28 @@ impl AgentFactory { let tool_chat_id = request .notification_chat_id .unwrap_or(request.session_chat_id); - let mut agent = agent.with_tool_context(ToolContext { - channel_name: Some(request.channel_name.to_string()), - sender_id: request.sender_id.map(str::to_string), - chat_id: Some(tool_chat_id.to_string()), - session_id: Some(session_id), - topic_id: request.topic_id.clone(), - message_id: request.message_id.map(str::to_string), - message_seq: None, - subagent_description: None, - nesting_depth: 0, - task_id: None, - parent_task_id: None, - tool_call_id: None, - // 注入专家 capability，TaskTool 据此强制校验子代理白/黑名单 - parent_capability: expert_capability.clone(), - }); + // 构建上下文压缩器（基于 effective_provider_config 的 context_window_tokens） + let compressor = Arc::new(ContextCompressor::from_provider_config( + &effective_provider_config, + )); + let mut agent = agent + .with_tool_context(ToolContext { + channel_name: Some(request.channel_name.to_string()), + sender_id: request.sender_id.map(str::to_string), + chat_id: Some(tool_chat_id.to_string()), + session_id: Some(session_id), + topic_id: request.topic_id.clone(), + message_id: request.message_id.map(str::to_string), + message_seq: None, + subagent_description: None, + nesting_depth: 0, + task_id: None, + parent_task_id: None, + tool_call_id: None, + // 注入专家 capability，TaskTool 据此强制校验子代理白/黑名单 + parent_capability: expert_capability.clone(), + }) + .with_compressor(Some(compressor)); // 如果有取消信号接收端，注入 Agent if let Some(token) = request.cancel_token { agent = agent.with_cancel_token(token); diff --git a/src/gateway/execution.rs b/src/gateway/execution.rs index f48d420..6ce20d9 100644 --- a/src/gateway/execution.rs +++ b/src/gateway/execution.rs @@ -2,8 +2,8 @@ use std::collections::HashMap; use std::sync::Arc; use crate::agent::{ - AgentError, AgentProcessResult, EmittedMessageHandler, PersistingEmittedMessageHandler, - SystemPromptContext, + AgentError, AgentProcessResult, CompactionSink, EmittedMessageHandler, + PersistingEmittedMessageHandler, SystemPromptContext, }; use crate::bus::message::ToolMessageState; use crate::bus::{ChatMessage, MediaItem, OutboundMessage, SYSTEM_CONTEXT_SCHEDULED_PROMPT}; @@ -24,6 +24,51 @@ impl EmittedMessageHandler for NoOpEmittedMessageHandler { async fn handle(&self, _message: ChatMessage) {} } +/// CompactionSink 实现：在 AgentLoop 内部触发 LLM 压缩时， +/// 把压缩后的消息写回 DB（标记原消息 is_compacted=1 + 插入摘要）。 +/// +/// 不在此处 reload 内存历史——process() 仍在使用局部 messages 变量， +/// 内存历史的刷新由 finalize_result 在 process 返回后统一处理。 +pub(crate) struct CompactionSinkImpl { + session: Arc>, + chat_id: String, + topic_id: String, +} + +impl CompactionSinkImpl { + pub(crate) fn new(session: Arc>, chat_id: String, topic_id: String) -> Self { + Self { + session, + chat_id, + topic_id, + } + } +} + +#[async_trait] +impl CompactionSink for CompactionSinkImpl { + async fn compact(&self, compressed: &[ChatMessage]) -> Result<(), AgentError> { + let mut session_guard = self.session.lock().await; + session_guard.ensure_persistent_session(&self.chat_id)?; + session_guard.ensure_chat_loaded(&self.chat_id, Some(&self.topic_id))?; + + let store = session_guard.store(); + let session_id = session_guard.persistent_session_id(&self.chat_id); + + store + .compact_topic_history(&session_id, &self.topic_id, compressed) + .map_err(|e| AgentError::Other(format!("compact_topic_history error: {}", e)))?; + + tracing::info!( + chat_id = %self.chat_id, + topic_id = %self.topic_id, + compressed_msg_count = compressed.len(), + "In-loop LLM compaction committed to DB (original messages retained as is_compacted=1)" + ); + Ok(()) + } +} + const SCHEDULED_TASK_EXECUTION_SYSTEM_PROMPT: &str = "系统说明：当前输入来自一次已经触发的定时任务执行。你现在需要执行任务内容本身，而不是创建、修改、恢复、暂停或查询新的定时任务。除非当前任务内容明确要求管理调度器，否则不要调用任何定时任务管理工具；像“每小时”、“每天”、“cron”、“定时”等词，只应视为任务背景，不应再解释为新的建任务请求。"; pub(crate) fn compose_scheduled_task_system_prompt(system_prompt: Option<&str>) -> String { @@ -128,8 +173,26 @@ impl AgentExecutionService { // 始终使用执行开始时捕获的 original_topic_id，避免从共享状态重复读取竞态 let target_topic_id = request.original_topic_id.as_deref(); - // 将结果消息保存到确定的话题 - if let Some(topic_id) = target_topic_id { + // 如果 AgentLoop 内部已触发 LLM 压缩，DB 已被 CompactionSink 更新 + // （原消息标记 is_compacted=1 + 压缩摘要已插入）。 + // 此时 emitted_messages 早已通过 handler 持久化，且作为"最新5个 unit" + // 包含在压缩输出中。直接 append 到内存历史会产生重复，因此从 DB 重新加载。 + if request.result.compaction_performed && is_current_turn { + let reload_topic = target_topic_id.unwrap_or(request.chat_id); + if let Err(err) = session.reload_topic_history(request.chat_id, reload_topic) { + tracing::error!( + error = %err, + chat_id = %request.chat_id, + topic_id = %reload_topic, + "Failed to reload topic history after in-loop compaction" + ); + } + tracing::info!( + chat_id = %request.chat_id, + topic_id = %reload_topic, + "In-loop compaction was performed; reloaded topic history from DB" + ); + } else if let Some(topic_id) = target_topic_id { if is_current_turn { // 话题未切换（current_topic == original_topic_id），安全更新内存历史 if let Err(err) = session @@ -196,8 +259,13 @@ impl AgentExecutionService { Vec::new() }; - // 只有当是最新回合时才触发历史压缩 - let should_schedule_compaction = is_current_turn; + // 只有当是最新回合且未在 loop 内触发过任何压缩（工程化或 LLM）时， + // 才触发兜底历史压缩。in-loop 已做工程化压缩时跳过——因为兜底基于 + // 未压缩历史的 estimate_tokens 判断会不准确，可能冗余触发 LLM 压缩， + // 违背"in-loop 已判断工程化压缩足够则不 LLM 压缩"的意图。 + let should_schedule_compaction = is_current_turn + && !request.result.compaction_performed + && !request.result.engineering_compaction_applied; Ok(FinalizedAgentResult { outbound_messages, @@ -294,7 +362,20 @@ impl AgentExecutionService { user_message_count, }; - let result = agent.process(history, Some(&system_prompt_context)).await?; + // 构建 CompactionSink：在 AgentLoop 内部触发 LLM 压缩时把结果写回 DB。 + // topic_id 退化为 chat_id（与 history_key 一致）。 + let compaction_topic_id = original_topic_id + .clone() + .unwrap_or_else(|| request.chat_id.to_string()); + let compaction_sink = CompactionSinkImpl::new( + request.session.clone(), + request.chat_id.to_string(), + compaction_topic_id, + ); + + let result = agent + .process(history, Some(&system_prompt_context), Some(&compaction_sink)) + .await?; let mut metadata = HashMap::new(); // 把用户消息的 UUID 回传给前端，前端用此更新本地消息 ID，使 todo 点击跳转能匹配 metadata.insert("user_message_id".to_string(), user_message.id.clone()); @@ -432,7 +513,19 @@ impl AgentExecutionService { user_message_count, }; - let result = agent.process(history, Some(&system_prompt_context)).await?; + // 构建 CompactionSink：在 AgentLoop 内部触发 LLM 压缩时把结果写回 DB。 + let compaction_topic_id = original_topic_id + .clone() + .unwrap_or_else(|| request.chat_id.to_string()); + let compaction_sink = CompactionSinkImpl::new( + request.session.clone(), + request.chat_id.to_string(), + compaction_topic_id, + ); + + let result = agent + .process(history, Some(&system_prompt_context), Some(&compaction_sink)) + .await?; let outbound_messages = self .finalize_result_and_schedule_compaction( diff --git a/src/tools/task/runtime.rs b/src/tools/task/runtime.rs index 601f159..cf94f39 100644 --- a/src/tools/task/runtime.rs +++ b/src/tools/task/runtime.rs @@ -585,7 +585,7 @@ impl DefaultSubAgentRuntime { let result = tokio::time::timeout( timeout_duration, - agent.process(history, Some(&system_prompt_context)), + agent.process(history, Some(&system_prompt_context), None), ) .await; @@ -630,7 +630,7 @@ impl DefaultSubAgentRuntime { let result = tokio::time::timeout( timeout_duration, - agent.process(history, Some(&system_prompt_context)), + agent.process(history, Some(&system_prompt_context), None), ) .await;