#[cfg(test)]
#[allow(clippy::module_inception)]
mod tests {
use crate::agent::context::compaction::*;
use crate::agent::context::*;
use crate::api::content::Content;
use crate::api::models::Message;
fn score_turn(
turn: &Turn,
messages: &[Message],
query_tokens: &std::collections::HashSet<String>,
) -> f32 {
score_turn_with_idf(turn, messages, query_tokens, None)
}
fn msg(role: &str, content: &str) -> Message {
Message {
role: role.to_string(),
content: Some(Content::text(content)),
reasoning_content: None,
tool_calls: None,
tool_call_id: None,
}
}
#[test]
fn test_new_context() {
let ctx = ConversationContext::new("system prompt".to_string());
assert_eq!(ctx.system_prompt(), "system prompt");
assert!(ctx.messages().is_empty());
}
#[test]
fn test_push_messages() {
let mut ctx = ConversationContext::new("sys".to_string());
ctx.push(msg("user", "hello"));
ctx.push(msg("assistant", "hi"));
assert_eq!(ctx.messages().len(), 2);
}
#[test]
fn test_build_messages_includes_system() {
let mut ctx = ConversationContext::new("system prompt".to_string());
ctx.push(msg("user", "hello"));
let built = ctx.build_messages();
assert_eq!(built.len(), 2);
assert_eq!(built[0].role, "system");
assert_eq!(
built[0].content.as_ref().map(|c| c.text_content()),
Some("system prompt".to_string())
);
assert_eq!(built[1].role, "user");
}
#[test]
fn test_build_messages_with_reasoning() {
let mut ctx = ConversationContext::new("sys".to_string());
ctx.set_last_reasoning("thinking...".to_string());
let built = ctx.build_messages();
let system = built[0].content.as_ref().unwrap();
let text = system.text_content();
assert!(text.contains("<reasoning>"));
assert!(text.contains("thinking..."));
}
#[test]
fn test_clear() {
let mut ctx = ConversationContext::new("sys".to_string());
ctx.push(msg("user", "hello"));
ctx.push(msg("assistant", "hi"));
ctx.clear();
assert!(ctx.messages().is_empty());
}
#[test]
fn test_update_system_prompt() {
let mut ctx = ConversationContext::new("old".to_string());
ctx.update_system_prompt("new".to_string());
assert_eq!(ctx.system_prompt(), "new");
}
#[test]
fn test_budget_remaining() {
let ctx = ConversationContext::new("short".to_string());
assert!(ctx.budget_remaining() > 100_000);
}
#[test]
fn test_restore() {
let messages = vec![msg("user", "hello"), msg("assistant", "hi")];
let ctx = ConversationContext::restore(
"sys".to_string(),
messages,
Some("reasoning".to_string()),
);
assert_eq!(ctx.messages().len(), 2);
assert_eq!(ctx.system_prompt(), "sys");
}
#[test]
fn test_force_compact_with_few_messages() {
let mut ctx = ConversationContext::new("sys".to_string());
ctx.push(msg("user", "hello"));
ctx.push(msg("assistant", "hi"));
ctx.force_compact();
assert_eq!(ctx.messages().len(), 2);
}
#[test]
fn test_auto_compaction_during_push() {
let mut ctx = ConversationContext::new("sys".to_string());
let filler = "x".repeat(10_000);
for i in 0..60 {
ctx.push(msg("user", &format!("message {i}: {filler}")));
ctx.push(msg("assistant", &format!("response {i}: {filler}")));
}
assert!(
ctx.messages().len() < 120,
"Expected compaction, got {} messages",
ctx.messages().len()
);
assert!(ctx.budget_remaining() > 0);
}
#[test]
fn test_compaction_produces_structured_summary() {
let mut ctx = ConversationContext::with_budget("sys".to_string(), 5_000, 0.5);
for i in 0..20 {
ctx.push(msg(
"user",
&format!("Fix the bug in src/main.rs iteration {i}"),
));
ctx.push(msg(
"assistant",
&format!(
"I'll fix the issue by editing src/main.rs {}",
"x".repeat(500)
),
));
}
assert!(ctx.compaction_count() > 0);
let summary = ctx
.messages()
.iter()
.find(|m| {
m.role == "system"
&& m.content
.as_ref()
.is_some_and(|c| c.text_content().contains("[Compacted"))
})
.expect("Should have a compaction summary");
let content = summary.content.as_ref().unwrap().text_content();
assert!(content.contains("Compacted"));
}
#[test]
fn test_compaction_log() {
let mut ctx = ConversationContext::with_budget("sys".to_string(), 5_000, 0.5);
for i in 0..20 {
ctx.push(msg("user", &format!("msg {i} {}", "x".repeat(500))));
ctx.push(msg("assistant", &format!("rsp {i} {}", "x".repeat(500))));
}
assert!(ctx.compaction_count() > 0);
let log = ctx.compaction_log();
assert!(log[0].before_messages > log[0].after_messages);
assert!(log[0].before_tokens > log[0].after_tokens);
}
#[test]
fn test_structured_summary_preserves_decisions() {
let mut ctx = ConversationContext::with_budget("sys".to_string(), 3_000, 0.3);
ctx.push(msg("user", "Fix the lint errors"));
ctx.push(msg(
"assistant",
"I'll fix the lint errors by updating the imports in src/lib.rs",
));
ctx.push(msg("user", "Now add tests"));
ctx.push(msg(
"assistant",
"Let me add unit tests for the new function",
));
for _ in 0..15 {
ctx.push(msg("user", &"x".repeat(500)));
ctx.push(msg("assistant", &"y".repeat(500)));
}
if let Some(s) = ctx.messages().iter().find(|m| {
m.role == "system"
&& m.content
.as_ref()
.is_some_and(|c| c.text_content().contains("[Compacted"))
}) {
let content = s.content.as_ref().unwrap().text_content();
if content.contains("Goal") {
assert!(content.contains("Fix") || content.contains("add"));
}
}
}
#[test]
fn test_extract_file_paths() {
let modified = Vec::new();
let mut read = Vec::new();
extract_file_paths("Check src/main.rs and tests/foo.py", &modified, &mut read);
assert!(read.contains(&"src/main.rs".to_string()));
assert!(read.contains(&"tests/foo.py".to_string()));
}
#[test]
fn test_extract_file_paths_ignores_urls() {
let modified = Vec::new();
let mut read = Vec::new();
extract_file_paths(
"Visit https://example.com/page.html for docs",
&modified,
&mut read,
);
assert!(!read.iter().any(|f| f.contains("example.com")));
}
#[test]
fn test_truncate_str() {
assert_eq!(truncate_str("short", 100), "short");
assert_eq!(truncate_str("hello world", 5), "hello...");
}
#[test]
fn test_multi_pass_compaction() {
let mut ctx = ConversationContext::with_budget("sys".to_string(), 2_000, 0.3);
for i in 0..30 {
ctx.push(msg("user", &format!("msg {i} {}", "x".repeat(300))));
ctx.push(msg("assistant", &format!("rsp {i} {}", "y".repeat(300))));
}
assert!(ctx.compaction_count() >= 1);
assert!(ctx.used_tokens() < ctx.max_context_tokens());
}
#[test]
fn test_simhash_near_duplicate() {
let tokens_a = tokenize_for_scoring("fn main() { println!(\"hello world\"); let x = 42; }");
let tokens_b = tokenize_for_scoring("fn main() { println!(\"hello world\"); let y = 43; }");
let tokens_c =
tokenize_for_scoring("completely different content about database queries sql schema");
let ha = simhash(&tokens_a);
let hb = simhash(&tokens_b);
let hc = simhash(&tokens_c);
assert!(
hamming_distance(ha, hb) < 8,
"Similar content should have low hamming distance, got {}",
hamming_distance(ha, hb)
);
assert!(
hamming_distance(ha, hc) >= 8,
"Different content should have high hamming distance, got {}",
hamming_distance(ha, hc)
);
}
#[test]
fn test_score_turn_recency() {
let messages = vec![
msg("user", "old question about authentication login"),
msg("assistant", "old answer about login implementation"),
msg("user", "new question about database schema migrations"),
msg("assistant", "new answer about migrations and rollback"),
];
let turns = group_into_turns(&messages);
let query: std::collections::HashSet<String> = Default::default();
assert_eq!(turns.len(), 2);
let score_old = score_turn(&turns[0], &messages, &query);
let score_new = score_turn(&turns[1], &messages, &query);
assert!(
score_new > score_old,
"Newer turn should score higher due to recency: {score_new} vs {score_old}"
);
}
#[test]
fn test_smart_compaction_reduces_messages() {
let mut ctx = ConversationContext::with_budget("sys".to_string(), 3_000, 0.3);
for i in 0..20 {
ctx.push(msg("user", &format!("question {i} {}", "x".repeat(300))));
ctx.push(msg("assistant", &format!("answer {i} {}", "y".repeat(300))));
}
assert!(
ctx.compaction_count() >= 1,
"Expected at least one compaction"
);
assert!(
ctx.messages().len() < 40,
"Expected message count to be reduced, got {}",
ctx.messages().len()
);
}
#[test]
fn test_adaptive_threshold_base() {
let ctx = ConversationContext::with_budget("sys".to_string(), 100_000, 0.75);
assert!((ctx.adaptive_threshold() - 0.75).abs() < 0.01);
}
#[test]
fn test_adaptive_threshold_tool_heavy() {
let mut ctx = ConversationContext::with_budget("sys".to_string(), 100_000, 0.75);
for _ in 0..3 {
ctx.messages.push(msg("user", "fix something"));
ctx.messages.push(Message {
role: "assistant".to_string(),
content: Some(Content::text("I'll fix it")),
reasoning_content: None,
tool_calls: Some(vec![crate::api::models::ToolCall {
id: "tc1".to_string(),
call_type: "function".to_string(),
function: crate::api::models::FunctionCall {
name: "file_edit".to_string(),
arguments: "{}".to_string(),
},
}]),
tool_call_id: None,
});
ctx.messages.push(msg("tool", "Success"));
}
ctx.cached_msg_tokens = ctx
.messages
.iter()
.map(ConversationContext::estimate_message_tokens)
.sum();
ctx.cached_tool_count = ctx
.messages
.iter()
.filter(|m| m.role == "tool" || m.tool_calls.is_some())
.count();
assert!(ctx.adaptive_threshold() < 0.75);
}
#[test]
fn test_build_idf_weights() {
let messages = vec![
msg("user", "fix the authentication bug in login"),
msg("assistant", "I'll fix the authentication issue"),
msg("user", "now optimize the database query"),
msg("assistant", "optimizing the query performance"),
];
let weights = build_idf_weights(&messages);
let auth_idf = weights.get("authentication").copied().unwrap_or(0.0);
let db_idf = weights.get("database").copied().unwrap_or(0.0);
assert!(
db_idf > auth_idf,
"Rare term 'database' should have higher IDF than 'authentication': {db_idf} vs {auth_idf}"
);
}
#[test]
fn test_idf_weighted_scoring_prefers_rare_terms() {
let messages = vec![
msg("user", "fix the authentication bug"),
msg("assistant", "fixing authentication now"),
msg("user", "optimize database migration schema"),
msg("assistant", "optimizing the database schema migration"),
];
let idf_weights = build_idf_weights(&messages);
let turns = group_into_turns(&messages);
let query: std::collections::HashSet<String> = ["database", "migration"]
.iter()
.map(|s| s.to_string())
.collect();
let score_auth = score_turn_with_idf(&turns[0], &messages, &query, Some(&idf_weights));
let score_db = score_turn_with_idf(&turns[1], &messages, &query, Some(&idf_weights));
assert!(
score_db > score_auth,
"Database turn should score higher for database query: {score_db} vs {score_auth}"
);
}
#[test]
fn test_observation_masking() {
let mut ctx = ConversationContext::new("sys".to_string());
let big_output = "x".repeat(1000);
ctx.messages.push(msg("user", "read the file"));
ctx.messages.push(msg("tool", &big_output));
ctx.messages.push(msg("user", "now fix it"));
ctx.messages.push(msg("assistant", "fixing"));
ctx.mask_old_tool_outputs(2);
let masked = ctx.messages[1].content.as_ref().unwrap();
let masked_text = masked.text_content();
assert!(masked_text.starts_with("[Tool output masked"));
assert!(masked_text.contains("1000 chars"));
}
#[test]
fn test_observation_masking_preserves_errors() {
let mut ctx = ConversationContext::new("sys".to_string());
let error_output = format!("Error: file not found\n{}", "x".repeat(500));
ctx.messages.push(msg("user", "read the file"));
ctx.messages.push(msg("tool", &error_output));
ctx.messages.push(msg("user", "try again"));
ctx.mask_old_tool_outputs(2);
let masked = ctx.messages[1].content.as_ref().unwrap();
let masked_text = masked.text_content();
assert!(masked_text.contains("error preserved"));
assert!(masked_text.contains("Error: file not found"));
}
#[test]
fn test_adaptive_target_tool_heavy() {
let mut ctx = ConversationContext::new("sys".to_string());
for _ in 0..3 {
ctx.messages.push(msg("user", "do something"));
ctx.messages.push(Message {
role: "assistant".to_string(),
content: None,
reasoning_content: None,
tool_calls: Some(vec![crate::api::models::ToolCall {
id: "tc".to_string(),
call_type: "function".to_string(),
function: crate::api::models::FunctionCall {
name: "bash".to_string(),
arguments: "{}".to_string(),
},
}]),
tool_call_id: None,
});
ctx.messages.push(msg("tool", "output"));
ctx.messages.push(msg("tool", "more output"));
ctx.messages.push(msg("tool", "even more"));
}
ctx.cached_tool_count = ctx
.messages
.iter()
.filter(|m| m.role == "tool" || m.tool_calls.is_some())
.count();
assert!((ctx.adaptive_target_ratio() - 0.20).abs() < 0.01);
}
#[test]
fn test_adaptive_target_conversational() {
let mut ctx = ConversationContext::new("sys".to_string());
for _ in 0..5 {
ctx.messages.push(msg("user", "question"));
ctx.messages.push(msg("assistant", "answer"));
}
assert!((ctx.adaptive_target_ratio() - 0.30).abs() < 0.01);
}
}