refactor(compression): 重构上下文压缩为两阶段+双阈值+三段保留

压缩时机从 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 辅助函数。
This commit is contained in:
oudecheng 2026-08-05 21:56:41 +08:00
parent c790ee1609
commit 6252b600cd
6 changed files with 484 additions and 290 deletions

View File

@ -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<tokio::sync::Mutex<tokio::sync::watch::Receiver<()>>>,
/// 上下文压缩器(可选)。配置后会在每轮工具调用完成后按双阈值执行压缩:
/// 1) 真实 prompt_tokens > 70% context_window → 进入压缩流程
/// 2) 工程化压缩(截断 tool 结果到 100 token仅改内存
/// 3) estimate_tokens > 50% context_window → 调 LLM 三段压缩;否则跳过
compressor: Option<Arc<ContextCompressor>>,
}
#[derive(Debug, Clone)]
pub struct AgentProcessResult {
pub final_response: ChatMessage,
pub emitted_messages: Vec<ChatMessage>,
/// 本轮 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<Arc<ContextCompressor>>) -> Self {
self.compressor = compressor;
self
}
pub fn tools(&self) -> &Arc<ToolRegistry> {
&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<ChatMessage>,
system_prompt_context: Option<&SystemPromptContext>,
compaction_sink: Option<&dyn CompactionSink>,
) -> Result<AgentProcessResult, AgentError> {
#[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% 触发阈值
// 提取 u32Copy避免持有 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_performedmessages 保持工程化压缩后的状态
}
}
} 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();

View File

@ -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<HistoryUnit> {
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<ChatMessage> {
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<Vec<ChatMessage>, AgentError> {
let mut truncated: Vec<ChatMessage> = 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<Vec<ChatMessage>, 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<Vec<ChatMessage>, 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<ChatMessage> = Vec::new();
let mut user_messages: Vec<ChatMessage> = Vec::new();
let mut compressible: Vec<HistoryUnit> = 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<ChatMessage> = 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<ChatMessage> = 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<ChatMessage> = 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<ChatMessage> = user_messages
.iter()
.map(|m| {
let mut msg = m.clone();
msg.role = "user".to_string();
msg
})
.collect();
let all_older_messages: Vec<ChatMessage> = 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.summarize_units_segment(
provider,
&middle_messages,
&middle_transcript,
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)
.await?
};
let older_summary = older_result?;
let newer_summary = newer_result?;
// Step 7: Assemble compressed history
let mut compressed: Vec<ChatMessage> = Vec::with_capacity(4);
// Step 6: Assemble compressed history
// [SystemGuards] + [oldest 5 units raw] + [middle summary] + [newest 5 units raw]
let mut compressed: Vec<ChatMessage> = 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<UserTurnRange> {
let user_indices: Vec<usize> = 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<ChatMessage> = (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:

View File

@ -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;

View File

@ -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,7 +211,12 @@ impl AgentFactory {
let tool_chat_id = request
.notification_chat_id
.unwrap_or(request.session_chat_id);
let mut agent = agent.with_tool_context(ToolContext {
// 构建上下文压缩器(基于 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()),
@ -225,7 +231,8 @@ impl AgentFactory {
tool_call_id: None,
// 注入专家 capabilityTaskTool 据此强制校验子代理白/黑名单
parent_capability: expert_capability.clone(),
});
})
.with_compressor(Some(compressor));
// 如果有取消信号接收端,注入 Agent
if let Some(token) = request.cancel_token {
agent = agent.with_cancel_token(token);

View File

@ -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<Mutex<Session>>,
chat_id: String,
topic_id: String,
}
impl CompactionSinkImpl {
pub(crate) fn new(session: Arc<Mutex<Session>>, 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(

View File

@ -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;