mod common;
use common::message_text;
use open_agent::{
ContentBlock, Message, MessageRole, TextBlock, ToolResultBlock, ToolUseBlock, estimate_tokens,
is_approaching_limit, truncate_messages,
};
use serde_json::json;
#[test]
fn tool_use_blocks_accumulate_name_id_and_input_lengths() {
let input = json!({"q": "rust"});
assert_eq!(
input.to_string().len(),
12,
"{{\"q\":\"rust\"}} is 12 chars"
);
let message = Message::new(
MessageRole::Assistant,
vec![ContentBlock::ToolUse(ToolUseBlock::new(
"call_12", "search", input,
))],
);
assert_eq!(estimate_tokens(&[message]), 13);
}
#[test]
fn tool_result_blocks_accumulate_id_and_content_lengths() {
let content = json!({"hits": 3});
assert_eq!(content.to_string().len(), 10, "{{\"hits\":3}} is 10 chars");
let message = Message::new(
MessageRole::User,
vec![ContentBlock::ToolResult(ToolResultBlock::new(
"call_12", content,
))],
);
assert_eq!(estimate_tokens(&[message]), 11);
}
fn six_hundred_token_message() -> Vec<Message> {
let messages = vec![Message::new(
MessageRole::User,
vec![ContentBlock::Text(TextBlock::new("x".repeat(2376)))],
)];
assert_eq!(estimate_tokens(&messages), 600);
messages
}
#[test]
fn approaching_limit_scales_the_limit_by_the_margin() {
let messages = six_hundred_token_message();
assert!(is_approaching_limit(&messages, 1000, 0.5));
}
#[test]
fn approaching_limit_is_strict_at_the_threshold() {
let messages = six_hundred_token_message();
assert!(!is_approaching_limit(&messages, 1000, 0.6));
assert!(is_approaching_limit(&messages, 999, 0.6));
}
#[test]
fn truncation_with_keep_zero_drops_everything_but_the_system_prompt() {
let messages = vec![
Message::system("system"),
Message::user("first"),
Message::user("second"),
];
let preserved = truncate_messages(&messages, 0, true);
assert_eq!(preserved.len(), 1);
assert_eq!(preserved[0].role, MessageRole::System);
assert!(truncate_messages(&messages, 0, false).is_empty());
}
#[test]
fn truncation_keeps_the_most_recent_messages_after_the_system_prompt() {
let messages = vec![
Message::system("system"),
Message::user("first"),
Message::user("second"),
Message::user("third"),
];
let kept = truncate_messages(&messages, 2, true);
assert_eq!(kept.len(), 3);
assert_eq!(kept[0].role, MessageRole::System);
assert_eq!(message_text(&kept[1]), "second");
assert_eq!(message_text(&kept[2]), "third");
let tail = truncate_messages(&messages, 2, false);
assert_eq!(tail.len(), 2);
assert_eq!(message_text(&tail[0]), "second");
assert_eq!(message_text(&tail[1]), "third");
}
#[test]
fn truncation_returns_everything_when_keep_covers_the_history() {
let messages = vec![Message::system("system"), Message::user("only")];
assert_eq!(truncate_messages(&messages, 2, true).len(), 2);
assert_eq!(truncate_messages(&messages, 5, true).len(), 2);
assert!(truncate_messages(&[], 5, true).is_empty());
}