1489 lines
53 KiB
Rust
1489 lines
53 KiB
Rust
use crate::bus::{
|
|
ChatMessage, SYSTEM_CONTEXT_AGENT_PROMPT, SYSTEM_CONTEXT_HISTORY_COMPACTION,
|
|
SYSTEM_CONTEXT_SCHEDULED_PROMPT,
|
|
};
|
|
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;
|
|
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<ChatMessage>,
|
|
},
|
|
/// 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<ChatMessage> {
|
|
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<HistoryUnit> {
|
|
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,
|
|
'\u{4E00}'..='\u{9FFF}' | // CJK Unified Ideographs
|
|
'\u{3040}'..='\u{309F}' | // Hiragana
|
|
'\u{30A0}'..='\u{30FF}' | // Katakana
|
|
'\u{AC00}'..='\u{D7AF}' | // Korean Hangul
|
|
'\u{3400}'..='\u{4DBF}' | // CJK Extension A
|
|
'\u{20000}'..='\u{2A6DF}' // CJK Extension B
|
|
)
|
|
}
|
|
|
|
/// Token estimation using weighted character counting based on language
|
|
pub fn estimate_tokens(messages: &[ChatMessage]) -> usize {
|
|
let mut cjk_count = 0usize;
|
|
let mut other_count = 0usize;
|
|
let mut media_refs_count = 0usize;
|
|
|
|
for msg in messages {
|
|
// Count content characters
|
|
for c in msg.content.chars() {
|
|
if is_cjk_char(c) {
|
|
cjk_count += 1;
|
|
} else {
|
|
other_count += 1;
|
|
}
|
|
}
|
|
// Count media references
|
|
media_refs_count += msg.media_refs.len();
|
|
}
|
|
|
|
// Weighted token calculation: CJK chars need more tokens per character
|
|
let content_tokens = (cjk_count as f64 / CJK_CHARS_PER_TOKEN)
|
|
+ (other_count as f64 / OTHER_CHARS_PER_TOKEN);
|
|
|
|
// JSON serialization overhead for message structure (fields, brackets, etc.)
|
|
let json_overhead = messages.len() * JSON_OVERHEAD_PER_MESSAGE;
|
|
|
|
// Media references add to JSON size (each path is a string in the array)
|
|
let media_overhead = media_refs_count * 20; // Each media ref adds ~20 chars to JSON
|
|
|
|
// Apply safety multiplier
|
|
((content_tokens + json_overhead as f64 + media_overhead as f64)
|
|
* TOKEN_ESTIMATE_SAFETY_MULTIPLIER) as usize
|
|
}
|
|
|
|
/// Configuration for context compression.
|
|
#[derive(Debug, Clone)]
|
|
pub struct ContextCompressionConfig {
|
|
/// Preserve the latest N complete user turns in full.
|
|
pub retain_last_user_turns: usize,
|
|
/// Maximum characters in summary
|
|
pub summary_max_chars: usize,
|
|
}
|
|
|
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
|
struct UserTurnRange {
|
|
start: usize,
|
|
end_exclusive: usize,
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct HistoryCompactionPlan {
|
|
pub preserved_system_messages: Vec<ChatMessage>,
|
|
pub summary_message: ChatMessage,
|
|
pub preserved_messages: Vec<ChatMessage>,
|
|
pub compressed_turns: usize,
|
|
pub preserved_turns: usize,
|
|
}
|
|
|
|
impl Default for ContextCompressionConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
retain_last_user_turns: 3,
|
|
summary_max_chars: 20_000,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Context compressor that reduces message history when it exceeds token limits.
|
|
#[derive(Clone)]
|
|
pub struct ContextCompressor {
|
|
config: ContextCompressionConfig,
|
|
context_window: usize,
|
|
/// Threshold ratio to trigger compression (70% of context window).
|
|
threshold_ratio: f64,
|
|
}
|
|
|
|
impl ContextCompressor {
|
|
#[cfg(test)]
|
|
fn summary_char_budget_for_context_window(context_window: usize) -> usize {
|
|
const SUMMARY_RATIO: f64 = 0.1;
|
|
const CHARS_PER_TOKEN: f64 = 2.5;
|
|
const MIN_SUMMARY_CHARS: usize = 1_500;
|
|
const MAX_SUMMARY_CHARS: usize = 50_000;
|
|
|
|
((context_window as f64 * SUMMARY_RATIO * CHARS_PER_TOKEN) as usize)
|
|
.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",
|
|
"tool" => "Tool",
|
|
_ => message.role.as_str(),
|
|
};
|
|
let name = message
|
|
.tool_name
|
|
.as_ref()
|
|
.map(|n| format!(" ({})", n))
|
|
.unwrap_or_default();
|
|
format!("{}: {}{}", role, message.content, name)
|
|
}
|
|
|
|
fn build_transcript(messages: &[ChatMessage]) -> String {
|
|
messages
|
|
.iter()
|
|
.map(Self::format_transcript_entry)
|
|
.collect::<Vec<_>>()
|
|
.join("\n\n")
|
|
}
|
|
|
|
fn split_text_chunks(text: &str, max_chars: usize) -> Vec<String> {
|
|
if text.is_empty() {
|
|
return Vec::new();
|
|
}
|
|
|
|
let chunk_size = max_chars.max(1);
|
|
let chars: Vec<char> = text.chars().collect();
|
|
chars
|
|
.chunks(chunk_size)
|
|
.map(|chunk| chunk.iter().collect())
|
|
.collect()
|
|
}
|
|
|
|
fn chunk_messages_for_summary(messages: &[ChatMessage], max_chars: usize) -> Vec<String> {
|
|
if messages.is_empty() {
|
|
return Vec::new();
|
|
}
|
|
|
|
let chunk_limit = max_chars.max(1);
|
|
let mut chunks = Vec::new();
|
|
let mut current = String::new();
|
|
|
|
for entry in messages.iter().map(Self::format_transcript_entry) {
|
|
let separator = if current.is_empty() { "" } else { "\n\n" };
|
|
let candidate = format!("{}{}{}", current, separator, entry);
|
|
|
|
if !current.is_empty() && char_count(&candidate) > chunk_limit {
|
|
chunks.push(current);
|
|
current = String::new();
|
|
}
|
|
|
|
if char_count(&entry) > chunk_limit {
|
|
if !current.is_empty() {
|
|
chunks.push(current);
|
|
current = String::new();
|
|
}
|
|
chunks.extend(Self::split_text_chunks(&entry, chunk_limit));
|
|
continue;
|
|
}
|
|
|
|
if current.is_empty() {
|
|
current = entry;
|
|
} else {
|
|
current.push_str("\n\n");
|
|
current.push_str(&entry);
|
|
}
|
|
}
|
|
|
|
if !current.is_empty() {
|
|
chunks.push(current);
|
|
}
|
|
|
|
chunks
|
|
}
|
|
|
|
fn build_summary_prompt(transcript: &str, target_chars: usize) -> String {
|
|
format!(
|
|
r#"You are a conversation compaction engine. Summarize the following conversation segment.
|
|
|
|
PRESERVE:
|
|
- Each user question or request in full or as a near-verbatim restatement
|
|
- All identifiers (UUIDs, hashes, file paths, URLs)
|
|
- Actions taken (tool calls, file operations, commands)
|
|
- Key information obtained (results, data, errors)
|
|
- Decisions and user preferences
|
|
- Current task status
|
|
|
|
OMIT:
|
|
- Reproducing full tool output verbatim unless it is essential
|
|
- Repeated greetings or filler
|
|
|
|
Do not assume tool content was pre-trimmed. You may receive long tool outputs; keep the important results, errors, and artifacts.
|
|
|
|
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
|
|
)
|
|
}
|
|
|
|
async fn summarize_transcript(
|
|
&self,
|
|
provider: &dyn LLMProvider,
|
|
transcript: &str,
|
|
target_chars: usize,
|
|
) -> Result<String, AgentError> {
|
|
let request = ChatCompletionRequest {
|
|
messages: vec![
|
|
Message::system("You are a helpful assistant."),
|
|
Message::user(Self::build_summary_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)
|
|
}
|
|
|
|
async fn summarize_chunked_transcript(
|
|
&self,
|
|
provider: &dyn LLMProvider,
|
|
messages: &[ChatMessage],
|
|
transcript: &str,
|
|
) -> Result<String, AgentError> {
|
|
let target_chars = self.config.summary_max_chars.max(1);
|
|
let mut layer = Self::chunk_messages_for_summary(messages, target_chars);
|
|
|
|
if layer.is_empty() {
|
|
layer.push(transcript.to_string());
|
|
}
|
|
|
|
for _ in 0..6 {
|
|
if layer.len() == 1 && char_count(&layer[0]) <= target_chars {
|
|
return self
|
|
.summarize_transcript(provider, &layer[0], target_chars)
|
|
.await;
|
|
}
|
|
|
|
let per_chunk_target = (target_chars / layer.len().max(1))
|
|
.max(500)
|
|
.min(target_chars);
|
|
let mut summaries = Vec::with_capacity(layer.len());
|
|
for chunk in &layer {
|
|
summaries.push(
|
|
self.summarize_transcript(provider, chunk, per_chunk_target)
|
|
.await?,
|
|
);
|
|
}
|
|
|
|
if summaries.len() == 1 {
|
|
let summary = summaries.pop().unwrap_or_default();
|
|
if char_count(&summary) <= target_chars {
|
|
return Ok(summary);
|
|
}
|
|
layer = Self::split_text_chunks(&summary, target_chars);
|
|
continue;
|
|
}
|
|
|
|
let merged = summaries.join("\n\n");
|
|
if char_count(&merged) <= target_chars {
|
|
return self
|
|
.summarize_transcript(provider, &merged, target_chars)
|
|
.await;
|
|
}
|
|
|
|
layer = Self::split_text_chunks(&merged, target_chars);
|
|
}
|
|
|
|
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<String, AgentError> {
|
|
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<String, AgentError> {
|
|
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<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());
|
|
}
|
|
|
|
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<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 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<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::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<ChatMessage> = 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: DEFAULT_THRESHOLD_RATIO,
|
|
}
|
|
}
|
|
|
|
pub fn from_provider_config(provider_config: &LLMProviderConfig) -> Self {
|
|
Self::from_runtime_config(&AgentRuntimeConfig::from(provider_config.clone()))
|
|
}
|
|
|
|
pub fn from_runtime_config(config: &AgentRuntimeConfig) -> Self {
|
|
Self::with_config(
|
|
config.context_window_tokens,
|
|
ContextCompressionConfig {
|
|
summary_max_chars: config.context_summary_char_budget,
|
|
..ContextCompressionConfig::default()
|
|
},
|
|
)
|
|
}
|
|
|
|
/// Create with custom configuration.
|
|
pub fn with_config(context_window: usize, config: ContextCompressionConfig) -> Self {
|
|
Self {
|
|
config,
|
|
context_window,
|
|
threshold_ratio: DEFAULT_THRESHOLD_RATIO,
|
|
}
|
|
}
|
|
|
|
/// Get the compression threshold in tokens (70% of context window).
|
|
fn threshold(&self) -> usize {
|
|
(self.context_window as f64 * self.threshold_ratio) as usize
|
|
}
|
|
|
|
pub fn should_compress(&self, history: &[ChatMessage]) -> bool {
|
|
estimate_tokens(history) > self.threshold()
|
|
}
|
|
|
|
fn user_turn_ranges(&self, history: &[ChatMessage]) -> Vec<UserTurnRange> {
|
|
let user_indices: Vec<usize> = history
|
|
.iter()
|
|
.enumerate()
|
|
.filter(|(_, message)| message.role == "user")
|
|
.map(|(index, _)| index)
|
|
.collect();
|
|
|
|
user_indices
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(turn_index, start)| UserTurnRange {
|
|
start: *start,
|
|
end_exclusive: user_indices
|
|
.get(turn_index + 1)
|
|
.copied()
|
|
.unwrap_or(history.len()),
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
fn should_preserve_system_message(&self, message: &ChatMessage) -> bool {
|
|
message.role == "system"
|
|
&& (message.has_system_context(SYSTEM_CONTEXT_AGENT_PROMPT)
|
|
|| message.has_system_context(SYSTEM_CONTEXT_SCHEDULED_PROMPT))
|
|
}
|
|
|
|
fn split_prefix_messages(
|
|
&self,
|
|
history: &[ChatMessage],
|
|
) -> (Vec<ChatMessage>, Vec<ChatMessage>) {
|
|
let preserved_system_messages = history
|
|
.iter()
|
|
.filter(|message| self.should_preserve_system_message(message))
|
|
.cloned()
|
|
.collect();
|
|
let summary_source = history
|
|
.iter()
|
|
.filter(|message| !self.should_preserve_system_message(message))
|
|
.cloned()
|
|
.collect();
|
|
|
|
(preserved_system_messages, summary_source)
|
|
}
|
|
|
|
pub async fn build_compaction_plan(
|
|
&self,
|
|
history: &[ChatMessage],
|
|
provider_config: &LLMProviderConfig,
|
|
) -> Result<Option<HistoryCompactionPlan>, AgentError> {
|
|
if !self.should_compress(history) {
|
|
return Ok(None);
|
|
}
|
|
|
|
let turn_ranges = self.user_turn_ranges(history);
|
|
if turn_ranges.len() <= self.config.retain_last_user_turns {
|
|
return Ok(None);
|
|
}
|
|
|
|
let preserved_turn_start =
|
|
turn_ranges[turn_ranges.len() - self.config.retain_last_user_turns].start;
|
|
if preserved_turn_start == 0 {
|
|
return Ok(None);
|
|
}
|
|
|
|
let (preserved_system_messages, summary_source) =
|
|
self.split_prefix_messages(&history[..preserved_turn_start]);
|
|
|
|
let summary = self
|
|
.summarize_segment(&summary_source, provider_config)
|
|
.await?;
|
|
|
|
// Sanitize preserved messages: the boundary between the summarized
|
|
// and preserved sections can split an assistant tool_calls message
|
|
// from its tool result messages, creating orphaned sequences that
|
|
// would cause API 400 errors.
|
|
let mut preserved_messages = history[preserved_turn_start..].to_vec();
|
|
let removed =
|
|
crate::bus::message::sanitize_incomplete_tool_call_sequences(&mut preserved_messages);
|
|
if removed > 0 {
|
|
tracing::warn!(
|
|
removed_count = removed,
|
|
preserved_turn_start,
|
|
"Compaction plan: removed incomplete tool call sequences from preserved messages"
|
|
);
|
|
}
|
|
|
|
Ok(Some(HistoryCompactionPlan {
|
|
preserved_system_messages,
|
|
summary_message: ChatMessage::system_with_context(
|
|
format!("[Compressed History]\n\n{}", summary),
|
|
Some(SYSTEM_CONTEXT_HISTORY_COMPACTION.to_string()),
|
|
),
|
|
preserved_messages,
|
|
compressed_turns: turn_ranges.len() - self.config.retain_last_user_turns,
|
|
preserved_turns: self.config.retain_last_user_turns,
|
|
}))
|
|
}
|
|
|
|
/// Main entry point - compresses history if over threshold.
|
|
pub async fn compress_if_needed(
|
|
&self,
|
|
history: Vec<ChatMessage>,
|
|
provider_config: &LLMProviderConfig,
|
|
) -> Result<Vec<ChatMessage>, AgentError> {
|
|
let tokens = estimate_tokens(&history);
|
|
if tokens <= self.threshold() {
|
|
#[cfg(debug_assertions)]
|
|
tracing::info!(
|
|
tokens = tokens,
|
|
threshold = self.threshold(),
|
|
msg_count = history.len(),
|
|
"Context compression not needed"
|
|
);
|
|
return Ok(history);
|
|
}
|
|
|
|
tracing::info!(
|
|
tokens = tokens,
|
|
threshold = self.threshold(),
|
|
msg_count = history.len(),
|
|
"Starting context compression"
|
|
);
|
|
|
|
let mut current_history = match self
|
|
.build_compaction_plan(&history, provider_config)
|
|
.await?
|
|
{
|
|
Some(plan) => {
|
|
let mut compressed = Vec::with_capacity(
|
|
plan.preserved_system_messages.len() + plan.preserved_messages.len() + 1,
|
|
);
|
|
compressed.extend(plan.preserved_system_messages);
|
|
compressed.push(plan.summary_message);
|
|
compressed.extend(plan.preserved_messages);
|
|
compressed
|
|
}
|
|
None => history,
|
|
};
|
|
|
|
// Post-compression sanitization: compression can split an assistant
|
|
// tool_calls message from its tool result messages at the boundary
|
|
// between the summarized and preserved sections. This pass removes
|
|
// any orphaned tool_calls or tool results to ensure the message
|
|
// sequence is always valid for the LLM API.
|
|
let removed =
|
|
crate::bus::message::sanitize_incomplete_tool_call_sequences(&mut current_history);
|
|
if removed > 0 {
|
|
tracing::warn!(
|
|
removed_count = removed,
|
|
remaining_messages = current_history.len(),
|
|
"Post-compression sanitization removed incomplete tool call sequences"
|
|
);
|
|
}
|
|
|
|
tracing::info!(
|
|
final_tokens = estimate_tokens(¤t_history),
|
|
final_msg_count = current_history.len(),
|
|
"Context compression completed"
|
|
);
|
|
|
|
Ok(current_history)
|
|
}
|
|
|
|
/// Summarize a segment of messages using LLM.
|
|
async fn summarize_segment(
|
|
&self,
|
|
messages: &[ChatMessage],
|
|
provider_config: &LLMProviderConfig,
|
|
) -> Result<String, AgentError> {
|
|
if messages.is_empty() {
|
|
return Ok(String::new());
|
|
}
|
|
|
|
let runtime_config = AgentRuntimeConfig::from(provider_config.clone());
|
|
let provider = create_provider(runtime_config.provider)
|
|
.map_err(|e| AgentError::ProviderCreation(e.to_string()))?;
|
|
let transcript = Self::build_transcript(messages);
|
|
|
|
let result = if char_count(&transcript) <= self.config.summary_max_chars {
|
|
self.summarize_transcript(
|
|
provider.as_ref(),
|
|
&transcript,
|
|
self.config.summary_max_chars,
|
|
)
|
|
.await
|
|
} else {
|
|
self.summarize_chunked_transcript(provider.as_ref(), messages, &transcript)
|
|
.await
|
|
};
|
|
|
|
match result {
|
|
Ok(summary) => Ok(summary),
|
|
Err(e) => {
|
|
tracing::warn!(error = %e, "LLM summarization failed, using truncated transcript");
|
|
Ok(take_prefix_chars(
|
|
&transcript,
|
|
self.config.summary_max_chars,
|
|
))
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_estimate_tokens() {
|
|
let messages = vec![
|
|
ChatMessage::user("Hello"),
|
|
ChatMessage::assistant("Hi there!"),
|
|
ChatMessage::user("How are you?"),
|
|
];
|
|
|
|
let tokens = estimate_tokens(&messages);
|
|
// Content: "Hello" (5) + "Hi there!" (9) + "How are you?" (12) = 26 chars
|
|
// English: 26 / 4 = 6.5 tokens for content
|
|
// JSON overhead: 3 * 50 = 150
|
|
// Total before multiplier: 156.5
|
|
// After 1.2x: ~188 tokens
|
|
assert!(
|
|
tokens > 100 && tokens < 300,
|
|
"Expected ~150-250 tokens for English content, got {}",
|
|
tokens
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_estimate_tokens_chinese_content() {
|
|
let messages = vec![
|
|
ChatMessage::user("你好，这是一个中文测试"),
|
|
ChatMessage::assistant("这是一个中文回复消息"),
|
|
];
|
|
|
|
let tokens = estimate_tokens(&messages);
|
|
// Content: ~20 CJK chars, CJK uses 2 chars/token = ~10 tokens for content
|
|
// JSON overhead: 2 * 50 = 100
|
|
// Total before multiplier: ~110
|
|
// After 1.2x: ~132 tokens
|
|
assert!(
|
|
tokens > 80 && tokens < 200,
|
|
"Expected ~100-180 tokens for Chinese content, got {}",
|
|
tokens
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_estimate_tokens_mixed_content() {
|
|
let messages = vec![
|
|
ChatMessage::user("Hello 世界 this is 测试"),
|
|
];
|
|
|
|
let tokens = estimate_tokens(&messages);
|
|
// Content: 18 English chars + 4 CJK chars
|
|
// English: 18 / 4 = 4.5, CJK: 4 / 2 = 2, content tokens = 6.5
|
|
// JSON overhead: 1 * 50 = 50
|
|
// Total before multiplier: 56.5
|
|
// After 1.2x: ~68 tokens
|
|
assert!(
|
|
tokens > 40 && tokens < 120,
|
|
"Expected ~50-100 tokens for mixed content, got {}",
|
|
tokens
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_chinese_tokens_higher_than_english() {
|
|
// Use more characters to make the content difference significant
|
|
// compared to JSON overhead (50 tokens per message)
|
|
let english = vec![ChatMessage::user(&"abcdefghij".repeat(20))]; // 200 English chars
|
|
let chinese = vec![ChatMessage::user(&"这是一个测试消息字".repeat(20))]; // 200 CJK chars (10 chars * 20)
|
|
|
|
let english_tokens = estimate_tokens(&english);
|
|
let chinese_tokens = estimate_tokens(&chinese);
|
|
|
|
// 200 English chars: 200/4 = 50 content tokens
|
|
// 200 CJK chars: 200/2 = 100 content tokens
|
|
// With same JSON overhead, Chinese should use ~1.5x tokens
|
|
assert!(
|
|
chinese_tokens > english_tokens * 130 / 100, // At least 1.3x
|
|
"Chinese (200 chars) should use more tokens than English (200 chars): EN={} CN={}",
|
|
english_tokens,
|
|
chinese_tokens
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_estimate_tokens_includes_image_media_refs() {
|
|
let temp_dir = tempfile::tempdir().unwrap();
|
|
let image_path = temp_dir.path().join("demo.jpg");
|
|
std::fs::write(&image_path, vec![0_u8; 12_000]).unwrap();
|
|
|
|
let plain = vec![ChatMessage::user("hello")];
|
|
let with_image = vec![ChatMessage::user_with_media(
|
|
"hello",
|
|
vec![image_path.to_string_lossy().to_string()],
|
|
)];
|
|
|
|
assert!(estimate_tokens(&with_image) > estimate_tokens(&plain));
|
|
}
|
|
|
|
#[test]
|
|
fn test_should_compress() {
|
|
let compressor = ContextCompressor::new(20);
|
|
// Need more content to trigger compression with new weighted calculation
|
|
// 200 English chars / 4 = 50 tokens, plus overhead
|
|
let messages = vec![ChatMessage::user(&"x".repeat(400))];
|
|
assert!(compressor.should_compress(&messages));
|
|
}
|
|
|
|
#[test]
|
|
fn test_user_turn_ranges_follow_user_boundaries() {
|
|
let compressor = ContextCompressor::new(100_000);
|
|
let history = vec![
|
|
ChatMessage::system("system"),
|
|
ChatMessage::user("u1"),
|
|
ChatMessage::assistant("a1"),
|
|
ChatMessage::tool("call-1", "bash", "t1"),
|
|
ChatMessage::user("u2"),
|
|
ChatMessage::assistant("a2"),
|
|
ChatMessage::user("u3"),
|
|
];
|
|
|
|
let turns = compressor.user_turn_ranges(&history);
|
|
assert_eq!(
|
|
turns,
|
|
vec![
|
|
UserTurnRange {
|
|
start: 1,
|
|
end_exclusive: 4
|
|
},
|
|
UserTurnRange {
|
|
start: 4,
|
|
end_exclusive: 6
|
|
},
|
|
UserTurnRange {
|
|
start: 6,
|
|
end_exclusive: 7
|
|
},
|
|
]
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_split_prefix_messages_preserves_key_system_messages() {
|
|
let compressor = ContextCompressor::new(50);
|
|
let prefix = vec![
|
|
ChatMessage::system_with_context(
|
|
"agent prompt",
|
|
Some(SYSTEM_CONTEXT_AGENT_PROMPT.to_string()),
|
|
),
|
|
ChatMessage::user("u1"),
|
|
ChatMessage::assistant("a1"),
|
|
ChatMessage::system_with_context(
|
|
"scheduled prompt",
|
|
Some(SYSTEM_CONTEXT_SCHEDULED_PROMPT.to_string()),
|
|
),
|
|
];
|
|
|
|
let (preserved_system_messages, summary_source) = compressor.split_prefix_messages(&prefix);
|
|
assert_eq!(preserved_system_messages.len(), 2);
|
|
assert_eq!(summary_source.len(), 2);
|
|
assert!(preserved_system_messages[0].has_system_context(SYSTEM_CONTEXT_AGENT_PROMPT));
|
|
assert!(preserved_system_messages[1].has_system_context(SYSTEM_CONTEXT_SCHEDULED_PROMPT));
|
|
}
|
|
|
|
#[test]
|
|
fn test_threshold() {
|
|
let compressor = ContextCompressor::new(128_000);
|
|
assert_eq!(compressor.threshold(), 89_600); // 70% of 128_000
|
|
}
|
|
|
|
#[test]
|
|
fn test_summary_char_budget_for_context_window_scales_and_clamps() {
|
|
assert_eq!(
|
|
ContextCompressor::summary_char_budget_for_context_window(4_096),
|
|
1_500
|
|
);
|
|
assert_eq!(
|
|
ContextCompressor::summary_char_budget_for_context_window(65_536),
|
|
16_384
|
|
);
|
|
assert_eq!(
|
|
ContextCompressor::summary_char_budget_for_context_window(128_000),
|
|
32_000
|
|
);
|
|
assert_eq!(
|
|
ContextCompressor::summary_char_budget_for_context_window(400_000),
|
|
50_000
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_chunk_messages_for_summary_keeps_message_boundaries_when_possible() {
|
|
let messages = vec![
|
|
ChatMessage::user("alpha"),
|
|
ChatMessage::assistant("beta"),
|
|
ChatMessage::user("gamma"),
|
|
];
|
|
|
|
let chunks = ContextCompressor::chunk_messages_for_summary(&messages, 30);
|
|
|
|
assert_eq!(chunks.len(), 2);
|
|
assert!(chunks.iter().all(|chunk| char_count(chunk) <= 30));
|
|
assert_eq!(chunks[0], "user: alpha\n\nAssistant: beta");
|
|
assert_eq!(chunks[1], "user: gamma");
|
|
}
|
|
|
|
#[test]
|
|
fn test_chunk_messages_for_summary_splits_oversized_message() {
|
|
let messages = vec![ChatMessage::user(&"x".repeat(25))];
|
|
|
|
let chunks = ContextCompressor::chunk_messages_for_summary(&messages, 10);
|
|
|
|
assert!(chunks.len() > 1);
|
|
assert!(chunks.iter().all(|chunk| char_count(chunk) <= 10));
|
|
assert_eq!(chunks.concat(), "user: xxxxxxxxxxxxxxxxxxxxxxxxx");
|
|
}
|
|
|
|
// =========================================================================
|
|
// HistoryUnit / parse_to_units tests
|
|
// =========================================================================
|
|
|
|
#[test]
|
|
fn test_parse_to_units_system_guards_preserved() {
|
|
let messages = vec![
|
|
ChatMessage::system_with_context(
|
|
"agent prompt",
|
|
Some(SYSTEM_CONTEXT_AGENT_PROMPT.to_string()),
|
|
),
|
|
ChatMessage::user("hello"),
|
|
ChatMessage::assistant("hi"),
|
|
];
|
|
|
|
let units = parse_to_units(&messages);
|
|
// SystemGuard + UserMessage + AssistantText = 3
|
|
assert_eq!(units.len(), 3);
|
|
assert!(matches!(units[0], HistoryUnit::SystemGuard(_)));
|
|
assert!(matches!(units[1], HistoryUnit::UserMessage(_)));
|
|
assert!(matches!(units[2], HistoryUnit::AssistantText(_)));
|
|
}
|
|
|
|
#[test]
|
|
fn test_parse_to_units_tool_round_atomic() {
|
|
let messages = vec![
|
|
ChatMessage::user("read a file"),
|
|
ChatMessage::assistant_with_tool_calls(
|
|
"let me read",
|
|
vec![crate::domain::messages::ToolCall {
|
|
id: "call_1".to_string(),
|
|
name: "file_read".to_string(),
|
|
arguments: serde_json::json!({"path": "/data/test.txt"}),
|
|
}],
|
|
),
|
|
ChatMessage::tool("call_1", "file_read", "file contents here"),
|
|
ChatMessage::assistant("the file says hello"),
|
|
];
|
|
|
|
let units = parse_to_units(&messages);
|
|
assert_eq!(units.len(), 3); // UserMessage + ToolRound + AssistantText
|
|
assert!(matches!(units[0], HistoryUnit::UserMessage(_)));
|
|
assert!(matches!(units[1], HistoryUnit::ToolRound { .. }));
|
|
assert!(matches!(units[2], HistoryUnit::AssistantText(_)));
|
|
|
|
// Verify ToolRound contains both assistant and results
|
|
if let HistoryUnit::ToolRound { assistant, results } = &units[1] {
|
|
assert_eq!(assistant.role, "assistant");
|
|
assert!(assistant.tool_calls.is_some());
|
|
assert_eq!(results.len(), 1);
|
|
assert_eq!(results[0].role, "tool");
|
|
assert_eq!(results[0].tool_call_id, Some("call_1".to_string()));
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_parse_to_units_orphaned_tool_result_dropped() {
|
|
let messages = vec![
|
|
ChatMessage::user("hello"),
|
|
// Orphaned tool result — no preceding assistant with tool_calls
|
|
ChatMessage::tool("orphan_1", "bash", "some output"),
|
|
ChatMessage::assistant("done"),
|
|
];
|
|
|
|
let units = parse_to_units(&messages);
|
|
// Orphaned tool result should be dropped, leaving UserMessage + AssistantText
|
|
assert_eq!(units.len(), 2);
|
|
assert!(matches!(units[0], HistoryUnit::UserMessage(_)));
|
|
assert!(matches!(units[1], HistoryUnit::AssistantText(_)));
|
|
}
|
|
|
|
#[test]
|
|
fn test_parse_to_units_multiple_tool_results_same_round() {
|
|
let messages = vec![
|
|
ChatMessage::user("do multiple things"),
|
|
ChatMessage::assistant_with_tool_calls(
|
|
"doing things",
|
|
vec![
|
|
crate::domain::messages::ToolCall {
|
|
id: "call_a".to_string(),
|
|
name: "bash".to_string(),
|
|
arguments: serde_json::json!({"command": "ls"}),
|
|
},
|
|
crate::domain::messages::ToolCall {
|
|
id: "call_b".to_string(),
|
|
name: "file_read".to_string(),
|
|
arguments: serde_json::json!({"path": "/data/test.txt"}),
|
|
},
|
|
],
|
|
),
|
|
ChatMessage::tool("call_a", "bash", "file1.txt\nfile2.txt"),
|
|
ChatMessage::tool("call_b", "file_read", "hello world"),
|
|
];
|
|
|
|
let units = parse_to_units(&messages);
|
|
assert_eq!(units.len(), 2); // UserMessage + ToolRound
|
|
if let HistoryUnit::ToolRound { results, .. } = &units[1] {
|
|
assert_eq!(results.len(), 2);
|
|
assert_eq!(results[0].tool_call_id, Some("call_a".to_string()));
|
|
assert_eq!(results[1].tool_call_id, Some("call_b".to_string()));
|
|
} else {
|
|
panic!("Expected ToolRound");
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_parse_to_units_empty_input() {
|
|
let units = parse_to_units(&[]);
|
|
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:
|
|
// compress_two_segment output MUST NOT contain any tool_calls or tool messages.
|
|
let compressor = ContextCompressor::new(128_000);
|
|
|
|
// Build a simple history that's well under the threshold
|
|
// so compress_two_segment returns it unchanged (no LLM call needed).
|
|
let history = vec![
|
|
ChatMessage::system_with_context(
|
|
"You are a helpful assistant.",
|
|
Some(SYSTEM_CONTEXT_AGENT_PROMPT.to_string()),
|
|
),
|
|
ChatMessage::user("hello"),
|
|
ChatMessage::assistant("hi there"),
|
|
];
|
|
|
|
// Dummy config — won't be used because history is under threshold.
|
|
let config = LLMProviderConfig {
|
|
provider_type: "openai".to_string(),
|
|
name: "test".to_string(),
|
|
base_url: "http://localhost".to_string(),
|
|
api_key: "sk-test".to_string(),
|
|
extra_headers: std::collections::HashMap::new(),
|
|
llm_timeout_secs: 120,
|
|
memory_maintenance_timeout_secs: 300,
|
|
model_id: "test-model".to_string(),
|
|
temperature: None,
|
|
max_tokens: None,
|
|
context_window_tokens: Some(128_000),
|
|
model_extra: std::collections::HashMap::new(),
|
|
max_tool_iterations: 100,
|
|
tool_result_max_chars: 100_000,
|
|
context_tool_result_trim_chars: 2_000,
|
|
max_images_in_context: 10,
|
|
max_image_age_rounds: 50,
|
|
};
|
|
|
|
let runtime = tokio::runtime::Runtime::new().unwrap();
|
|
let result = runtime.block_on(compressor.compress_two_segment(&history, &config));
|
|
assert!(result.is_ok());
|
|
let compressed = result.unwrap();
|
|
// Under threshold: should return unchanged (3 messages)
|
|
assert_eq!(compressed.len(), 3);
|
|
// Critical invariant: NO tool_calls or tool_call_id anywhere
|
|
for msg in &compressed {
|
|
assert!(msg.tool_calls.is_none(),
|
|
"compress_two_segment output should never contain tool_calls");
|
|
assert!(msg.tool_call_id.is_none(),
|
|
"compress_two_segment output should never contain tool_call_id");
|
|
}
|
|
}
|
|
}
|