feat: preserve last assistant message during compaction
When context window compaction occurs, the last assistant message is now preserved in addition to the system prompt, README, and summary. This improves continuity after compaction by keeping the LLM's most recent response, which often contains important context about what was just done or what comes next. New message order after compaction: [System Prompt] -> [README/AGENTS.md] -> [ACD Stub?] -> [Summary] -> [Last Assistant] -> [Latest User?] Changes: - Add last_assistant_message field to PreservedMessages struct - Modify extract_preserved_messages() to find last assistant message - Modify reset_with_summary_and_stub() to include last assistant message - Add comprehensive integration tests using MockProvider Tests cover edge cases: - No assistant message exists - Tool-call-only assistant messages (still preserved) - Multiple assistant messages (only last one preserved) - No trailing user message
This commit is contained in:
Dhanji R. Prasanna
@@ -298,6 +298,11 @@ Format this as a detailed but concise summary that can be used to resume the con
|
||||
format!("Previous conversation summary:\n\n{}", summary),
|
||||
));
|
||||
|
||||
// Add the last assistant message if present (preserves continuity)
|
||||
if let Some(assistant_msg) = preserved.last_assistant_message {
|
||||
self.add_message(assistant_msg);
|
||||
}
|
||||
|
||||
// Add the latest user message if provided
|
||||
if let Some(user_msg) = latest_user_message {
|
||||
self.add_message(Message::new(MessageRole::User, user_msg));
|
||||
@@ -326,9 +331,18 @@ Format this as a detailed but concise summary that can be used to resume the con
|
||||
}
|
||||
});
|
||||
|
||||
// Find the last assistant message in the conversation
|
||||
let last_assistant_message = self
|
||||
.conversation_history
|
||||
.iter()
|
||||
.rev()
|
||||
.find(|m| matches!(m.role, MessageRole::Assistant))
|
||||
.cloned();
|
||||
|
||||
PreservedMessages {
|
||||
system_prompt,
|
||||
readme,
|
||||
last_assistant_message,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -725,6 +739,7 @@ Format this as a detailed but concise summary that can be used to resume the con
|
||||
struct PreservedMessages {
|
||||
system_prompt: Option<Message>,
|
||||
readme: Option<Message>,
|
||||
last_assistant_message: Option<Message>,
|
||||
}
|
||||
|
||||
impl ThinResult {
|
||||
|
||||
396
crates/g3-core/tests/compaction_test.rs
Normal file
396
crates/g3-core/tests/compaction_test.rs
Normal file
@@ -0,0 +1,396 @@
|
||||
//! Integration tests for context compaction
|
||||
//!
|
||||
//! These tests verify that compaction correctly preserves important messages
|
||||
//! when summarizing conversation history.
|
||||
|
||||
use g3_core::ui_writer::NullUiWriter;
|
||||
use g3_core::Agent;
|
||||
use g3_providers::mock::{MockProvider, MockResponse};
|
||||
use g3_providers::{Message, MessageRole, ProviderRegistry};
|
||||
use tempfile::TempDir;
|
||||
|
||||
/// Helper to create an agent with a mock provider
|
||||
async fn create_agent_with_mock(provider: MockProvider) -> (Agent<NullUiWriter>, TempDir) {
|
||||
let temp_dir = TempDir::new().unwrap();
|
||||
|
||||
// Create a provider registry with the mock provider
|
||||
let mut registry = ProviderRegistry::new();
|
||||
registry.register(provider);
|
||||
|
||||
// Create a minimal config
|
||||
let config = g3_config::Config::default();
|
||||
|
||||
let agent = Agent::new_for_test(config, NullUiWriter, registry)
|
||||
.await
|
||||
.expect("Failed to create agent");
|
||||
|
||||
(agent, temp_dir)
|
||||
}
|
||||
|
||||
/// Helper to find the last assistant message in history
|
||||
fn find_last_assistant_message(history: &[Message]) -> Option<&Message> {
|
||||
history
|
||||
.iter()
|
||||
.rev()
|
||||
.find(|m| matches!(m.role, MessageRole::Assistant))
|
||||
}
|
||||
|
||||
/// Helper to check if a message contains a substring
|
||||
fn message_contains(history: &[Message], role: MessageRole, substring: &str) -> bool {
|
||||
history.iter().any(|m| {
|
||||
std::mem::discriminant(&m.role) == std::mem::discriminant(&role)
|
||||
&& m.content.contains(substring)
|
||||
})
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Compaction Tests
|
||||
// =============================================================================
|
||||
|
||||
/// Test: Last assistant message is preserved after compaction
|
||||
///
|
||||
/// This is the main feature test. After compaction:
|
||||
/// 1. System prompt is preserved
|
||||
/// 2. Summary is added
|
||||
/// 3. Last assistant message should be preserved (NEW)
|
||||
/// 4. Latest user message is preserved
|
||||
///
|
||||
/// The order should be:
|
||||
/// [System] -> [Summary] -> [Last Assistant] -> [Latest User]
|
||||
#[tokio::test]
|
||||
async fn test_compaction_preserves_last_assistant_message() {
|
||||
// Create a provider that will:
|
||||
// 1. Respond to initial conversation
|
||||
// 2. Provide a summary when compaction is triggered
|
||||
let provider = MockProvider::new()
|
||||
// Response 1: Initial assistant response
|
||||
.with_response(MockResponse::text(
|
||||
"I understand you want to build a web server. Let me help you with that.",
|
||||
))
|
||||
// Response 2: Second assistant response (this should be preserved after compaction)
|
||||
.with_response(MockResponse::text(
|
||||
"Here's the implementation plan:\n1. Create main.rs\n2. Add dependencies\n3. Implement routes",
|
||||
))
|
||||
// Response 3: This will be the summary response during compaction
|
||||
.with_response(MockResponse::text(
|
||||
"Summary: User wants to build a web server. We discussed implementation plan with 3 steps.",
|
||||
))
|
||||
// Response 4: Post-compaction response
|
||||
.with_response(MockResponse::text(
|
||||
"Continuing from where we left off...",
|
||||
));
|
||||
|
||||
let (mut agent, _temp_dir) = create_agent_with_mock(provider).await;
|
||||
|
||||
// Build up conversation history
|
||||
agent
|
||||
.execute_task("I want to build a web server in Rust", None, false)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
agent
|
||||
.execute_task("What's the implementation plan?", None, false)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Verify the last assistant message before compaction
|
||||
let history_before = agent.get_context_window().conversation_history.clone();
|
||||
let last_assistant_before = find_last_assistant_message(&history_before)
|
||||
.expect("Should have assistant message before compaction");
|
||||
assert!(
|
||||
last_assistant_before.content.contains("implementation plan"),
|
||||
"Last assistant message should contain 'implementation plan', got: {}",
|
||||
last_assistant_before.content
|
||||
);
|
||||
|
||||
// Trigger manual compaction
|
||||
let compaction_result = agent.force_compact().await;
|
||||
assert!(
|
||||
compaction_result.is_ok(),
|
||||
"Compaction should succeed: {:?}",
|
||||
compaction_result.err()
|
||||
);
|
||||
assert!(
|
||||
compaction_result.unwrap(),
|
||||
"Compaction should return true on success"
|
||||
);
|
||||
|
||||
// Verify the context after compaction
|
||||
let history_after = &agent.get_context_window().conversation_history;
|
||||
|
||||
// Debug: Print the history after compaction
|
||||
eprintln!("\n=== History after compaction ===");
|
||||
for (i, msg) in history_after.iter().enumerate() {
|
||||
eprintln!(
|
||||
" {}: {:?} - {}...",
|
||||
i,
|
||||
msg.role,
|
||||
msg.content.chars().take(80).collect::<String>()
|
||||
);
|
||||
}
|
||||
|
||||
// 1. Should have a summary message
|
||||
assert!(
|
||||
message_contains(history_after, MessageRole::System, "Summary:"),
|
||||
"Should have summary message after compaction"
|
||||
);
|
||||
|
||||
// 2. Should preserve the last assistant message content
|
||||
// This is the key assertion - the last assistant message should be preserved
|
||||
assert!(
|
||||
message_contains(history_after, MessageRole::Assistant, "implementation plan"),
|
||||
"Should preserve last assistant message with 'implementation plan' after compaction.\n\
|
||||
History: {:?}",
|
||||
history_after
|
||||
.iter()
|
||||
.map(|m| format!("{:?}: {}...", m.role, m.content.chars().take(50).collect::<String>()))
|
||||
.collect::<Vec<_>>()
|
||||
);
|
||||
|
||||
// 3. Should preserve the latest user message
|
||||
assert!(
|
||||
message_contains(history_after, MessageRole::User, "implementation plan"),
|
||||
"Should preserve latest user message after compaction"
|
||||
);
|
||||
}
|
||||
|
||||
/// Test: Compaction with no assistant messages doesn't crash
|
||||
///
|
||||
/// Edge case: If there are no assistant messages (e.g., fresh session),
|
||||
/// compaction should still work without errors.
|
||||
#[tokio::test]
|
||||
async fn test_compaction_no_assistant_message() {
|
||||
// Provider that returns a summary
|
||||
let provider = MockProvider::new()
|
||||
.with_response(MockResponse::text("Summary: Empty conversation."));
|
||||
|
||||
let (mut agent, _temp_dir) = create_agent_with_mock(provider).await;
|
||||
|
||||
// Add a user message directly to context without getting a response.
|
||||
// This simulates a state where we have user input but no assistant response yet.
|
||||
use g3_providers::Message;
|
||||
agent.add_message_to_context(Message::new(MessageRole::User, "Hello".to_string()));
|
||||
|
||||
// Trigger compaction - should not crash
|
||||
let result = agent.force_compact().await;
|
||||
assert!(result.is_ok(), "Compaction should succeed even with no assistant messages");
|
||||
}
|
||||
|
||||
/// Test: Compaction preserves tool-call-only assistant message
|
||||
///
|
||||
/// Even if the last assistant message is just a tool call (no prose),
|
||||
/// it should still be preserved as it contains important context.
|
||||
#[tokio::test]
|
||||
async fn test_compaction_preserves_tool_call_only_message() {
|
||||
let provider = MockProvider::new()
|
||||
// Response 1: Text response
|
||||
.with_response(MockResponse::text("Let me check that file."))
|
||||
// Response 2: Tool call response (this is the last assistant message)
|
||||
.with_response(MockResponse::text_with_json_tool(
|
||||
"", // No prose, just tool call
|
||||
"read_file",
|
||||
serde_json::json!({"file_path": "important.rs"}),
|
||||
))
|
||||
// Response 3: Summary
|
||||
.with_response(MockResponse::text("Summary: User asked to check a file."))
|
||||
// Response 4: Post-compaction
|
||||
.with_response(MockResponse::text("Continuing..."));
|
||||
|
||||
let (mut agent, _temp_dir) = create_agent_with_mock(provider).await;
|
||||
|
||||
// Build conversation
|
||||
agent.execute_task("Check the file", None, false).await.unwrap();
|
||||
agent.execute_task("Read important.rs", None, false).await.unwrap();
|
||||
|
||||
// Get the last assistant message before compaction
|
||||
let history_before = agent.get_context_window().conversation_history.clone();
|
||||
let last_assistant_before = find_last_assistant_message(&history_before);
|
||||
|
||||
// Verify we have an assistant message (might contain tool call JSON)
|
||||
assert!(
|
||||
last_assistant_before.is_some(),
|
||||
"Should have an assistant message before compaction"
|
||||
);
|
||||
|
||||
// Trigger compaction
|
||||
let result = agent.force_compact().await;
|
||||
assert!(result.is_ok(), "Compaction should succeed");
|
||||
|
||||
// The assistant message should be preserved (even if it's tool-call-only)
|
||||
let history_after = &agent.get_context_window().conversation_history;
|
||||
let has_assistant = history_after
|
||||
.iter()
|
||||
.any(|m| matches!(m.role, MessageRole::Assistant));
|
||||
|
||||
assert!(
|
||||
has_assistant,
|
||||
"Should have at least one assistant message after compaction (the preserved last one)"
|
||||
);
|
||||
}
|
||||
|
||||
/// Test: Compaction with multiple assistant messages preserves only the last one
|
||||
///
|
||||
/// When there are multiple assistant messages, only the most recent one
|
||||
/// should be preserved (in addition to the summary).
|
||||
#[tokio::test]
|
||||
async fn test_compaction_preserves_only_last_assistant() {
|
||||
let provider = MockProvider::new()
|
||||
// Response 1: First assistant response
|
||||
.with_response(MockResponse::text("FIRST_RESPONSE: Hello!"))
|
||||
// Response 2: Second assistant response
|
||||
.with_response(MockResponse::text("SECOND_RESPONSE: How can I help?"))
|
||||
// Response 3: Third assistant response (this should be preserved)
|
||||
.with_response(MockResponse::text("THIRD_RESPONSE: Let me assist you."))
|
||||
// Response 4: Summary
|
||||
.with_response(MockResponse::text("Summary: Greeted user three times."))
|
||||
// Response 5: Post-compaction
|
||||
.with_response(MockResponse::text("Continuing..."));
|
||||
|
||||
let (mut agent, _temp_dir) = create_agent_with_mock(provider).await;
|
||||
|
||||
// Build conversation with multiple exchanges
|
||||
agent.execute_task("Hi", None, false).await.unwrap();
|
||||
agent.execute_task("What can you do?", None, false).await.unwrap();
|
||||
agent.execute_task("Help me", None, false).await.unwrap();
|
||||
|
||||
// Trigger compaction
|
||||
let result = agent.force_compact().await;
|
||||
assert!(result.is_ok(), "Compaction should succeed");
|
||||
|
||||
let history_after = &agent.get_context_window().conversation_history;
|
||||
|
||||
// Debug output
|
||||
eprintln!("\n=== History after compaction ===");
|
||||
for (i, msg) in history_after.iter().enumerate() {
|
||||
eprintln!(
|
||||
" {}: {:?} - {}",
|
||||
i,
|
||||
msg.role,
|
||||
msg.content.chars().take(100).collect::<String>()
|
||||
);
|
||||
}
|
||||
|
||||
// Should have THIRD_RESPONSE (the last one)
|
||||
assert!(
|
||||
message_contains(history_after, MessageRole::Assistant, "THIRD_RESPONSE"),
|
||||
"Should preserve the LAST assistant message (THIRD_RESPONSE)"
|
||||
);
|
||||
|
||||
// Should NOT have FIRST_RESPONSE or SECOND_RESPONSE as separate messages
|
||||
// (they might be mentioned in the summary, but not as standalone assistant messages)
|
||||
let assistant_messages: Vec<_> = history_after
|
||||
.iter()
|
||||
.filter(|m| matches!(m.role, MessageRole::Assistant))
|
||||
.collect();
|
||||
|
||||
// Should have exactly one assistant message (the preserved last one)
|
||||
assert_eq!(
|
||||
assistant_messages.len(),
|
||||
1,
|
||||
"Should have exactly one assistant message after compaction (the last one), got {}: {:?}",
|
||||
assistant_messages.len(),
|
||||
assistant_messages.iter().map(|m| &m.content).collect::<Vec<_>>()
|
||||
);
|
||||
}
|
||||
|
||||
/// Test: Compaction without a trailing user message
|
||||
///
|
||||
/// Edge case: The last message in history is from the assistant (no user follow-up).
|
||||
/// The assistant message should still be preserved.
|
||||
#[tokio::test]
|
||||
async fn test_compaction_no_trailing_user_message() {
|
||||
let provider = MockProvider::new()
|
||||
// Response 1: Assistant response (will be the last message)
|
||||
.with_response(MockResponse::text("LAST_ASSISTANT_MESSAGE: Here's your answer."))
|
||||
// Response 2: Summary
|
||||
.with_response(MockResponse::text("Summary: Provided an answer."));
|
||||
|
||||
let (mut agent, _temp_dir) = create_agent_with_mock(provider).await;
|
||||
|
||||
// Single exchange - assistant response is the last message
|
||||
agent.execute_task("What's the answer?", None, false).await.unwrap();
|
||||
|
||||
// Verify last message is from assistant
|
||||
let history_before = agent.get_context_window().conversation_history.clone();
|
||||
let last_msg = history_before.last().expect("Should have messages");
|
||||
assert!(
|
||||
matches!(last_msg.role, MessageRole::Assistant),
|
||||
"Last message should be from assistant before compaction"
|
||||
);
|
||||
|
||||
// Trigger compaction
|
||||
let result = agent.force_compact().await;
|
||||
assert!(result.is_ok(), "Compaction should succeed");
|
||||
|
||||
let history_after = &agent.get_context_window().conversation_history;
|
||||
|
||||
// Should preserve the last assistant message
|
||||
assert!(
|
||||
message_contains(history_after, MessageRole::Assistant, "LAST_ASSISTANT_MESSAGE"),
|
||||
"Should preserve last assistant message even without trailing user message"
|
||||
);
|
||||
}
|
||||
|
||||
/// Test: Message order after compaction is correct
|
||||
///
|
||||
/// The order should be:
|
||||
/// [System Prompt] -> [README if present] -> [Summary] -> [Last Assistant] -> [Latest User if present]
|
||||
#[tokio::test]
|
||||
async fn test_compaction_message_order() {
|
||||
let provider = MockProvider::new()
|
||||
.with_response(MockResponse::text("ASSISTANT_TO_PRESERVE: I'll help you."))
|
||||
.with_response(MockResponse::text("SUMMARY_CONTENT: User asked for help."));
|
||||
|
||||
let (mut agent, _temp_dir) = create_agent_with_mock(provider).await;
|
||||
|
||||
agent.execute_task("USER_MESSAGE_TO_PRESERVE: Help me", None, false).await.unwrap();
|
||||
|
||||
// Trigger compaction
|
||||
agent.force_compact().await.unwrap();
|
||||
|
||||
let history = &agent.get_context_window().conversation_history;
|
||||
|
||||
// Debug output
|
||||
eprintln!("\n=== Message order after compaction ===");
|
||||
for (i, msg) in history.iter().enumerate() {
|
||||
eprintln!(
|
||||
" {}: {:?} - {}",
|
||||
i,
|
||||
msg.role,
|
||||
msg.content.chars().take(60).collect::<String>()
|
||||
);
|
||||
}
|
||||
|
||||
// Find indices of key messages
|
||||
let summary_idx = history
|
||||
.iter()
|
||||
.position(|m| m.content.contains("SUMMARY_CONTENT"))
|
||||
.expect("Should have summary");
|
||||
|
||||
let assistant_idx = history
|
||||
.iter()
|
||||
.position(|m| matches!(m.role, MessageRole::Assistant) && m.content.contains("ASSISTANT_TO_PRESERVE"))
|
||||
.expect("Should have preserved assistant message");
|
||||
|
||||
let user_idx = history
|
||||
.iter()
|
||||
.position(|m| matches!(m.role, MessageRole::User) && m.content.contains("USER_MESSAGE_TO_PRESERVE"));
|
||||
|
||||
// Summary should come before assistant message
|
||||
assert!(
|
||||
summary_idx < assistant_idx,
|
||||
"Summary (idx {}) should come before assistant message (idx {})",
|
||||
summary_idx,
|
||||
assistant_idx
|
||||
);
|
||||
|
||||
// If there's a user message, it should come after assistant message
|
||||
if let Some(user_idx) = user_idx {
|
||||
assert!(
|
||||
assistant_idx < user_idx,
|
||||
"Assistant message (idx {}) should come before user message (idx {})",
|
||||
assistant_idx,
|
||||
user_idx
|
||||
);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user