Douwe Osinga
9a91fccfb5
Co-authored-by: Douwe Osinga <douwe@squareup.com> Co-authored-by: Claude Sonnet 4.5 <noreply@anthropic.com>
750 lines
26 KiB
Rust
750 lines
26 KiB
Rust
use anyhow::Result;
|
|
use async_trait::async_trait;
|
|
use futures::StreamExt;
|
|
use goose::agents::{Agent, AgentEvent, SessionConfig};
|
|
use goose::conversation::message::{Message, MessageContent};
|
|
use goose::conversation::Conversation;
|
|
use goose::model::ModelConfig;
|
|
use goose::providers::base::{
|
|
stream_from_single_message, MessageStream, Provider, ProviderDef, ProviderMetadata,
|
|
ProviderUsage, Usage,
|
|
};
|
|
use goose::providers::errors::ProviderError;
|
|
use goose::session::session_manager::SessionType;
|
|
use goose::session::Session;
|
|
use rmcp::model::Tool;
|
|
use std::sync::atomic::{AtomicBool, Ordering};
|
|
use std::sync::Arc;
|
|
use tempfile::TempDir;
|
|
|
|
struct MockCompactionProvider {
|
|
/// Tracks whether compaction has occurred (for context limit recovery case)
|
|
has_compacted: Arc<AtomicBool>,
|
|
}
|
|
|
|
impl MockCompactionProvider {
|
|
fn new() -> Self {
|
|
Self {
|
|
has_compacted: Arc::new(AtomicBool::new(false)),
|
|
}
|
|
}
|
|
|
|
/// Calculate input tokens based on system prompt and messages
|
|
/// Simulates realistic token counts for different scenarios
|
|
fn calculate_input_tokens(&self, system_prompt: &str, messages: &[Message]) -> i32 {
|
|
// Check if this is a compaction call
|
|
let is_compaction_call = messages.len() == 1
|
|
&& messages[0].content.iter().any(|c| {
|
|
if let MessageContent::Text(text) = c {
|
|
text.text.to_lowercase().contains("summarize")
|
|
} else {
|
|
false
|
|
}
|
|
});
|
|
|
|
if is_compaction_call {
|
|
// For compaction: system prompt length is a good proxy for conversation size
|
|
// Base: 6000 (system) + conversation content embedded in prompt
|
|
6000 + (system_prompt.len() as i32 / 4).max(400)
|
|
} else {
|
|
// Regular call: system prompt + messages
|
|
let system_tokens = if system_prompt.is_empty() { 0 } else { 6000 };
|
|
|
|
let message_tokens: i32 = messages
|
|
.iter()
|
|
.map(|msg| {
|
|
let mut tokens = 100;
|
|
for content in &msg.content {
|
|
if let MessageContent::Text(text) = content {
|
|
if text.text.contains("long_tool_call") {
|
|
tokens += 15000;
|
|
}
|
|
}
|
|
}
|
|
tokens
|
|
})
|
|
.sum();
|
|
|
|
system_tokens + message_tokens
|
|
}
|
|
}
|
|
|
|
/// Calculate output tokens based on response type
|
|
fn calculate_output_tokens(&self, is_compaction: bool, messages: &[Message]) -> i32 {
|
|
if is_compaction {
|
|
// Compaction produces a compact summary
|
|
200
|
|
} else {
|
|
// Regular responses vary by content
|
|
let has_hello = messages.iter().any(|msg| {
|
|
msg.content.iter().any(|c| {
|
|
if let MessageContent::Text(text) = c {
|
|
text.text.to_lowercase().contains("hello")
|
|
} else {
|
|
false
|
|
}
|
|
})
|
|
});
|
|
|
|
if has_hello {
|
|
50 // Simple greeting response
|
|
} else {
|
|
100 // Default response
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl Provider for MockCompactionProvider {
|
|
async fn stream(
|
|
&self,
|
|
_model_config: &ModelConfig,
|
|
_session_id: &str,
|
|
system_prompt: &str,
|
|
messages: &[Message],
|
|
_tools: &[Tool],
|
|
) -> Result<MessageStream, ProviderError> {
|
|
// Check if this is a compaction call (message contains "summarize")
|
|
let is_compaction = messages.iter().any(|msg| {
|
|
msg.content.iter().any(|content| {
|
|
if let MessageContent::Text(text) = content {
|
|
text.text.to_lowercase().contains("summarize")
|
|
} else {
|
|
false
|
|
}
|
|
})
|
|
});
|
|
|
|
// Calculate realistic token counts based on actual content
|
|
let input_tokens = self.calculate_input_tokens(system_prompt, messages);
|
|
let output_tokens = self.calculate_output_tokens(is_compaction, messages);
|
|
|
|
// Simulate context limit: if input > 20k tokens and we haven't compacted yet, fail
|
|
const CONTEXT_LIMIT: i32 = 20000;
|
|
if !is_compaction
|
|
&& input_tokens > CONTEXT_LIMIT
|
|
&& !self.has_compacted.load(Ordering::SeqCst)
|
|
{
|
|
return Err(ProviderError::ContextLengthExceeded(format!(
|
|
"Context limit exceeded: {} > {}",
|
|
input_tokens, CONTEXT_LIMIT
|
|
)));
|
|
}
|
|
|
|
// If this is a compaction call, mark that we've compacted
|
|
if is_compaction {
|
|
self.has_compacted.store(true, Ordering::SeqCst);
|
|
}
|
|
|
|
// Generate response
|
|
let message = if is_compaction {
|
|
Message::assistant().with_text("<mock summary of conversation>")
|
|
} else {
|
|
let response_text = if messages.iter().any(|msg| {
|
|
msg.content.iter().any(|c| {
|
|
if let MessageContent::Text(text) = c {
|
|
text.text.to_lowercase().contains("hello")
|
|
} else {
|
|
false
|
|
}
|
|
})
|
|
}) {
|
|
"Hi there! How can I help you?"
|
|
} else {
|
|
"This is a mock response."
|
|
};
|
|
Message::assistant().with_text(response_text)
|
|
};
|
|
|
|
let usage = ProviderUsage::new(
|
|
"mock-model".to_string(),
|
|
Usage::new(
|
|
Some(input_tokens),
|
|
Some(output_tokens),
|
|
Some(input_tokens + output_tokens),
|
|
),
|
|
);
|
|
|
|
Ok(stream_from_single_message(message, usage))
|
|
}
|
|
|
|
fn get_model_config(&self) -> ModelConfig {
|
|
ModelConfig::new("mock-model").unwrap()
|
|
}
|
|
|
|
fn get_name(&self) -> &str {
|
|
"mock-compaction"
|
|
}
|
|
}
|
|
|
|
impl ProviderDef for MockCompactionProvider {
|
|
type Provider = Self;
|
|
|
|
fn metadata() -> ProviderMetadata {
|
|
ProviderMetadata {
|
|
name: "mock".to_string(),
|
|
display_name: "Mock Compaction Provider".to_string(),
|
|
description: "Mock provider for compaction testing".to_string(),
|
|
default_model: "mock-model".to_string(),
|
|
known_models: vec![],
|
|
model_doc_link: "".to_string(),
|
|
config_keys: vec![],
|
|
allows_unlisted_models: false,
|
|
}
|
|
}
|
|
|
|
fn from_env(
|
|
_model: ModelConfig,
|
|
_extensions: Vec<goose::config::ExtensionConfig>,
|
|
) -> futures::future::BoxFuture<'static, anyhow::Result<Self>> {
|
|
Box::pin(async { Ok(Self::new()) })
|
|
}
|
|
}
|
|
|
|
/// Helper: Set up a test session with initial messages and token counts
|
|
async fn setup_test_session(
|
|
agent: &Agent,
|
|
temp_dir: &TempDir,
|
|
session_name: &str,
|
|
messages: Vec<Message>,
|
|
) -> Result<Session> {
|
|
let session = agent
|
|
.config
|
|
.session_manager
|
|
.create_session(
|
|
temp_dir.path().to_path_buf(),
|
|
session_name.to_string(),
|
|
SessionType::Hidden,
|
|
)
|
|
.await?;
|
|
|
|
let conversation = Conversation::new_unvalidated(messages);
|
|
agent
|
|
.config
|
|
.session_manager
|
|
.replace_conversation(&session.id, &conversation)
|
|
.await?;
|
|
|
|
// Set initial token counts
|
|
agent
|
|
.config
|
|
.session_manager
|
|
.update(&session.id)
|
|
.total_tokens(Some(1000))
|
|
.input_tokens(Some(600))
|
|
.output_tokens(Some(400))
|
|
.accumulated_total_tokens(Some(1000))
|
|
.accumulated_input_tokens(Some(600))
|
|
.accumulated_output_tokens(Some(400))
|
|
.apply()
|
|
.await?;
|
|
|
|
Ok(session)
|
|
}
|
|
|
|
/// Helper: Assert conversation has been compacted with proper message visibility
|
|
fn assert_conversation_compacted(conversation: &Conversation) {
|
|
let messages = conversation.messages();
|
|
assert!(!messages.is_empty(), "Conversation should not be empty");
|
|
|
|
// Find the summary message (contains "mock summary")
|
|
let summary_index = messages
|
|
.iter()
|
|
.position(|msg| {
|
|
msg.content.iter().any(|content| {
|
|
if let MessageContent::Text(text) = content {
|
|
text.text.contains("mock summary")
|
|
} else {
|
|
false
|
|
}
|
|
})
|
|
})
|
|
.expect("Conversation should contain the summary message");
|
|
|
|
let summary_msg = &messages[summary_index];
|
|
|
|
// Assert summary message visibility
|
|
assert!(
|
|
summary_msg.is_agent_visible(),
|
|
"Summary message should be agent visible"
|
|
);
|
|
assert!(
|
|
!summary_msg.is_user_visible(),
|
|
"Summary message should NOT be user visible"
|
|
);
|
|
|
|
// Check messages BEFORE the summary (the compacted original messages)
|
|
// These should be made agent-invisible
|
|
for (idx, msg) in messages.iter().enumerate() {
|
|
if idx < summary_index {
|
|
// Old messages before summary: agent can't see them
|
|
assert!(
|
|
!msg.is_agent_visible(),
|
|
"Message before summary at index {} should be agent-invisible",
|
|
idx
|
|
);
|
|
}
|
|
}
|
|
|
|
// Check for continuation message after summary
|
|
// (Should exist and be agent-only)
|
|
if summary_index + 1 < messages.len() {
|
|
let continuation_msg = &messages[summary_index + 1];
|
|
// Continuation message should contain instructions about not mentioning summary
|
|
let has_continuation_text = continuation_msg.content.iter().any(|content| {
|
|
if let MessageContent::Text(text) = content {
|
|
text.text.contains("previous message contains a summary")
|
|
|| text.text.contains("summarization occurred")
|
|
} else {
|
|
false
|
|
}
|
|
});
|
|
|
|
if has_continuation_text {
|
|
assert!(
|
|
continuation_msg.is_agent_visible(),
|
|
"Continuation message should be agent visible"
|
|
);
|
|
assert!(
|
|
!continuation_msg.is_user_visible(),
|
|
"Continuation message should NOT be user visible"
|
|
);
|
|
}
|
|
}
|
|
|
|
// Any messages AFTER the continuation (e.g., preserved recent user message)
|
|
// should be fully visible to both agent and user
|
|
let continuation_end = summary_index + 2;
|
|
for (idx, msg) in messages.iter().enumerate() {
|
|
if idx >= continuation_end {
|
|
assert!(
|
|
msg.is_agent_visible() && msg.is_user_visible(),
|
|
"Message after compaction at index {} should be fully visible",
|
|
idx
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_manual_compaction_updates_token_counts_and_conversation() -> Result<()> {
|
|
let temp_dir = TempDir::new()?;
|
|
let agent = Agent::new();
|
|
|
|
// Setup session with initial messages
|
|
// Each message ~100 tokens, so 4 messages = ~400 tokens in conversation
|
|
let messages = vec![
|
|
Message::user().with_text("Hello, can you help me with something?"),
|
|
Message::assistant().with_text("Of course! What do you need help with?"),
|
|
Message::user().with_text("I need to understand how compaction works."),
|
|
Message::assistant()
|
|
.with_text("Compaction is a process that summarizes conversation history."),
|
|
];
|
|
|
|
let session = setup_test_session(&agent, &temp_dir, "manual-compact-test", messages).await?;
|
|
|
|
// Setup mock provider
|
|
let provider = Arc::new(MockCompactionProvider::new());
|
|
agent.update_provider(provider, &session.id).await?;
|
|
|
|
// Execute manual compaction
|
|
let result = agent.execute_command("/compact", &session.id).await?;
|
|
assert!(result.is_some(), "Compaction should return a result");
|
|
|
|
// Verify token counts
|
|
let updated_session = agent
|
|
.config
|
|
.session_manager
|
|
.get_session(&session.id, true)
|
|
.await?;
|
|
|
|
// Expected token calculation for compaction:
|
|
// During compaction, the 4 messages are embedded in the system prompt template
|
|
// - Input: system prompt with embedded conversation + "Please summarize" message
|
|
// - Output: summary (200 tokens)
|
|
//
|
|
// From mock provider calculation:
|
|
// - System prompt (with 4 embedded messages): varies based on template + content
|
|
// - Single "summarize" message: 100 tokens
|
|
// - Total input observed: ~6100 tokens
|
|
//
|
|
// After compaction:
|
|
// - current input_tokens = summary output (200) - the new compact context
|
|
// - current output_tokens = None (compaction doesn't produce new output)
|
|
// - current total_tokens = 200
|
|
// - accumulated_total = initial (1000) + compaction cost
|
|
let expected_summary_output = 200; // compact summary
|
|
|
|
// Verify the key invariants after manual compaction:
|
|
// After compaction, the current context is ONLY the summary (200 tokens)
|
|
// This is the new agent-visible input context
|
|
assert_eq!(
|
|
updated_session.input_tokens,
|
|
Some(expected_summary_output),
|
|
"Input tokens should be exactly the summary output (200 tokens)"
|
|
);
|
|
assert_eq!(
|
|
updated_session.output_tokens, None,
|
|
"Output tokens should be None after compaction (no new assistant output)"
|
|
);
|
|
assert_eq!(
|
|
updated_session.total_tokens,
|
|
Some(expected_summary_output),
|
|
"Total should equal input (200 tokens) after compaction"
|
|
);
|
|
|
|
// Accumulated tokens increased by the compaction cost
|
|
// Initial: 1000
|
|
// Compaction input: ~6400 (system 6000 + 4 messages ~400)
|
|
// Compaction output: 200
|
|
// Expected accumulated: 1000 + 6400 + 200 = 7600
|
|
let accumulated = updated_session.accumulated_total_tokens.unwrap();
|
|
assert!(
|
|
(7300..=7900).contains(&accumulated),
|
|
"Accumulated should be ~7600 (1000 initial + 6400 input + 200 output). Got: {}",
|
|
accumulated
|
|
);
|
|
|
|
// Verify conversation has been compacted
|
|
let compacted_conversation = updated_session
|
|
.conversation
|
|
.expect("Session should have conversation");
|
|
|
|
assert_conversation_compacted(&compacted_conversation);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_auto_compaction_during_reply() -> Result<()> {
|
|
let temp_dir = TempDir::new()?;
|
|
let agent = Agent::new();
|
|
|
|
// Setup session with many messages to have substantial context
|
|
// 20 exchanges = 40 messages * 100 tokens = ~4000 tokens in conversation
|
|
let mut messages = vec![];
|
|
for i in 0..20 {
|
|
messages.push(Message::user().with_text(format!("User message {}", i)));
|
|
messages.push(Message::assistant().with_text(format!("Assistant response {}", i)));
|
|
}
|
|
|
|
let session = setup_test_session(&agent, &temp_dir, "auto-compact-test", messages).await?;
|
|
|
|
// Capture initial context size before triggering reply
|
|
// Should be: system (6000) + 40 messages (4000) = ~10000 tokens
|
|
let initial_session = agent
|
|
.config
|
|
.session_manager
|
|
.get_session(&session.id, true)
|
|
.await?;
|
|
let initial_input_tokens = initial_session.input_tokens.unwrap_or(0);
|
|
|
|
// Setup mock provider (no context limit enforcement)
|
|
let provider = Arc::new(MockCompactionProvider::new());
|
|
agent.update_provider(provider, &session.id).await?;
|
|
|
|
// Trigger a reply
|
|
// Expected tokens for reply:
|
|
// - Input: system (6000) + 40 messages (4000) + new user message (100) = 10100 tokens
|
|
// - Output: regular response (100 tokens)
|
|
let user_message = Message::user().with_text("Tell me more about compaction");
|
|
|
|
let session_config = SessionConfig {
|
|
id: session.id.clone(),
|
|
schedule_id: None,
|
|
max_turns: None,
|
|
retry_config: None,
|
|
};
|
|
|
|
let reply_stream = agent.reply(user_message, session_config, None).await?;
|
|
tokio::pin!(reply_stream);
|
|
|
|
// Track compaction and context size changes
|
|
let mut compaction_occurred = false;
|
|
let mut input_tokens_after_compaction: Option<i32> = None;
|
|
|
|
while let Some(event_result) = reply_stream.next().await {
|
|
match event_result {
|
|
Ok(AgentEvent::HistoryReplaced(_)) => {
|
|
compaction_occurred = true;
|
|
|
|
// Capture the input tokens immediately after compaction
|
|
let session_after_compact = agent
|
|
.config
|
|
.session_manager
|
|
.get_session(&session.id, true)
|
|
.await?;
|
|
input_tokens_after_compaction = session_after_compact.input_tokens;
|
|
}
|
|
Ok(_) => {}
|
|
Err(e) => return Err(e),
|
|
}
|
|
}
|
|
|
|
let updated_session = agent
|
|
.config
|
|
.session_manager
|
|
.get_session(&session.id, true)
|
|
.await?;
|
|
|
|
if compaction_occurred {
|
|
// Verify that current input context decreased after compaction
|
|
let tokens_after =
|
|
input_tokens_after_compaction.expect("Should have captured tokens after compaction");
|
|
|
|
// Before compaction: system (6000) + 40 messages (4000) = 10,000 tokens
|
|
// After compaction: only the summary (200 tokens) - this becomes the new input
|
|
assert!(
|
|
tokens_after < initial_input_tokens,
|
|
"Input tokens should decrease after compaction. Before: {}, After: {}",
|
|
initial_input_tokens,
|
|
tokens_after
|
|
);
|
|
|
|
// After compaction, input should be exactly the summary: 200 tokens
|
|
assert_eq!(
|
|
tokens_after, 200,
|
|
"Input tokens after compaction should be exactly 200 (summary). Got: {}",
|
|
tokens_after
|
|
);
|
|
|
|
// After the subsequent reply, the current window includes:
|
|
// - system (6000) + summary (200) + new user message (100) + reply (100) = 6400
|
|
let final_input = updated_session.input_tokens.unwrap();
|
|
let final_output = updated_session.output_tokens.unwrap();
|
|
let final_total = updated_session.total_tokens.unwrap();
|
|
|
|
assert!(
|
|
final_input >= 6000,
|
|
"Final input should include at least system prompt (6000). Got: {}",
|
|
final_input
|
|
);
|
|
assert_eq!(
|
|
final_output, 100,
|
|
"Final output should be 100 tokens (default response). Got: {}",
|
|
final_output
|
|
);
|
|
assert_eq!(
|
|
final_total,
|
|
final_input + final_output,
|
|
"Final total should equal input + output"
|
|
);
|
|
|
|
// Accumulated tokens should include:
|
|
// - Initial: 1000
|
|
// - Compaction: ~10,400 input + 200 output = 10,600
|
|
// - Reply: ~6,300 input + 100 output = 6,400
|
|
// Total: 1000 + 10,600 + 6,400 = 18,000
|
|
let accumulated = updated_session.accumulated_total_tokens.unwrap();
|
|
assert!(
|
|
(17000..=19000).contains(&accumulated),
|
|
"Accumulated should be ~18,000 (initial + compaction + reply). Got: {}",
|
|
accumulated
|
|
);
|
|
} else {
|
|
// If no compaction, accumulated should include reply cost
|
|
// - Initial: 1000
|
|
// - Reply: system (6000) + 40 messages (4000) + new message (100) = 10,100 input
|
|
// - Reply output: 100
|
|
// Total: 1000 + 10,100 + 100 = 11,200
|
|
let accumulated = updated_session.accumulated_total_tokens.unwrap();
|
|
assert!(
|
|
(11000..=11500).contains(&accumulated),
|
|
"Accumulated should be ~11,200 (initial + reply). Got: {}",
|
|
accumulated
|
|
);
|
|
|
|
// Current window should be: 10,100 input + 100 output = 10,200
|
|
let final_input = updated_session.input_tokens.unwrap();
|
|
let final_output = updated_session.output_tokens.unwrap();
|
|
|
|
assert!(
|
|
(10000..=10500).contains(&final_input),
|
|
"Input should be ~10,100. Got: {}",
|
|
final_input
|
|
);
|
|
assert_eq!(
|
|
final_output, 100,
|
|
"Output should be 100. Got: {}",
|
|
final_output
|
|
);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_context_limit_recovery_compaction() -> Result<()> {
|
|
let temp_dir = TempDir::new()?;
|
|
let agent = Agent::new();
|
|
|
|
// Setup session with messages that will push context over the limit
|
|
// Each message = 100 tokens, but we'll add a large one
|
|
let messages = vec![
|
|
Message::user().with_text("Hello"),
|
|
Message::assistant().with_text("Hi there"),
|
|
Message::user().with_text("Can you process this long_tool_call result?"),
|
|
Message::assistant().with_text("Processing..."),
|
|
];
|
|
// Token calculation:
|
|
// - 3 regular messages: 300 tokens
|
|
// - 1 message with "long_tool_call": 100 + 15000 = 15100 tokens
|
|
// - Total conversation: ~15400 tokens
|
|
// - With system prompt (6000): 21400 tokens
|
|
|
|
let session = setup_test_session(&agent, &temp_dir, "context-limit-test", messages).await?;
|
|
|
|
// Setup mock provider with context limit of 20000 tokens
|
|
// Initial context (6000 system + 15400 messages = 21400) exceeds this limit
|
|
let provider = Arc::new(MockCompactionProvider::new());
|
|
agent.update_provider(provider, &session.id).await?;
|
|
|
|
// Try to send a message - should trigger context limit, then recover via compaction
|
|
let session_config = SessionConfig {
|
|
id: session.id.clone(),
|
|
schedule_id: None,
|
|
max_turns: None,
|
|
retry_config: None,
|
|
};
|
|
|
|
let reply_stream = agent
|
|
.reply(
|
|
Message::user().with_text("Tell me more"),
|
|
session_config,
|
|
None,
|
|
)
|
|
.await?;
|
|
tokio::pin!(reply_stream);
|
|
|
|
// Track compaction and context size changes
|
|
let mut compaction_occurred = false;
|
|
let mut got_response = false;
|
|
let mut input_tokens_after_compaction: Option<i32> = None;
|
|
|
|
while let Some(event_result) = reply_stream.next().await {
|
|
match event_result {
|
|
Ok(AgentEvent::HistoryReplaced(_)) => {
|
|
compaction_occurred = true;
|
|
|
|
// Capture the input tokens immediately after compaction
|
|
let session_after_compact = agent
|
|
.config
|
|
.session_manager
|
|
.get_session(&session.id, true)
|
|
.await?;
|
|
input_tokens_after_compaction = session_after_compact.input_tokens;
|
|
}
|
|
Ok(AgentEvent::Message(msg)) => {
|
|
// Check if we got a real response (not just a notification)
|
|
if msg
|
|
.content
|
|
.iter()
|
|
.any(|c| matches!(c, MessageContent::Text(_)))
|
|
{
|
|
got_response = true;
|
|
}
|
|
}
|
|
Ok(_) => {}
|
|
Err(e) => return Err(e),
|
|
}
|
|
}
|
|
|
|
// Verify recovery occurred
|
|
assert!(
|
|
compaction_occurred,
|
|
"Compaction should have occurred due to context limit (>20000 tokens)"
|
|
);
|
|
assert!(
|
|
got_response,
|
|
"Should have received a response after recovery"
|
|
);
|
|
|
|
// Verify token counts
|
|
let updated_session = agent
|
|
.config
|
|
.session_manager
|
|
.get_session(&session.id, true)
|
|
.await?;
|
|
|
|
// Expected token flow:
|
|
// 1. Initial attempt: >20000 tokens -> Context limit exceeded
|
|
// 2. Compaction triggered:
|
|
// - Input: system prompt + messages (including long_tool_call with 15k tokens)
|
|
// - Output: 200 tokens (summary)
|
|
// - New context size: 200 tokens
|
|
// 3. Retry with compacted context:
|
|
// - Input: system prompt + summary (200) + new message
|
|
// - Output: 100 tokens (response)
|
|
|
|
// Verify that current input context is dramatically reduced after compaction
|
|
let tokens_after =
|
|
input_tokens_after_compaction.expect("Should have captured tokens after compaction");
|
|
|
|
// After compaction, the input context should be ONLY the summary: 200 tokens
|
|
// Before: system (6000) + long_tool_call messages (~15,400) = 21,400 (exceeded limit!)
|
|
// After: only summary (200 tokens)
|
|
assert_eq!(
|
|
tokens_after, 200,
|
|
"Input tokens after compaction should be exactly 200 (summary only). Got: {}",
|
|
tokens_after
|
|
);
|
|
|
|
// The compacted context is now well under the 20k limit
|
|
assert!(
|
|
tokens_after < 20000,
|
|
"Compacted context should be under 20k limit. Got: {}",
|
|
tokens_after
|
|
);
|
|
|
|
// Check the final token state after recovery
|
|
// Note: The session state reflects the RETRY call (after compaction),
|
|
// which only sees agent-visible messages (summary + continuation + user message)
|
|
let final_input = updated_session.input_tokens.unwrap();
|
|
let final_output = updated_session.output_tokens;
|
|
let final_total = updated_session.total_tokens.unwrap();
|
|
|
|
// After compaction, the retry only sees agent-visible messages:
|
|
// Input: system (6000) + summary (~100) + continuation (~100) + user message (~100) = ~6300
|
|
// Output: 200 (mock detects "summarized" in continuation as compaction)
|
|
// Total: ~6500
|
|
assert!(
|
|
(6000..=6600).contains(&final_input),
|
|
"Final input should reflect retry with agent-visible messages (~6300). Got: {}",
|
|
final_input
|
|
);
|
|
|
|
assert_eq!(
|
|
final_output,
|
|
Some(200),
|
|
"Final output should be 200 (mock detects continuation as compaction). Got: {:?}",
|
|
final_output
|
|
);
|
|
|
|
assert_eq!(
|
|
final_total,
|
|
final_input + final_output.unwrap(),
|
|
"Final total should equal input + output"
|
|
);
|
|
|
|
// Accumulated tokens should include all operations:
|
|
// - Initial: 1000
|
|
// - Compaction: ~6400 input (mock uses system_prompt.len()/4) + 200 output = ~6600
|
|
// - Reply: ~6500 input + 200 output = ~6700
|
|
// Total: 1000 + 6600 + 6700 = ~14300
|
|
let accumulated = updated_session.accumulated_total_tokens.unwrap();
|
|
assert!(
|
|
(13000..=16000).contains(&accumulated),
|
|
"Accumulated should be ~14300 (initial + compaction + reply). Got: {}",
|
|
accumulated
|
|
);
|
|
|
|
// Verify that the conversation was compacted
|
|
let updated_conversation = updated_session
|
|
.conversation
|
|
.expect("Session should have conversation");
|
|
assert_conversation_compacted(&updated_conversation);
|
|
|
|
Ok(())
|
|
}
|