Remove orphaned tool calls before compaction

- Add fix_conversation to handle agent_visible vs user_visible messages separately
- Update fix_lead_trail to have more specific error message for trailing assistant messages
- Add two new tests for handling pending tool requests in last assistant message
- Conversation fixes now only apply to agent_visible messages, preserving non-visible ones
This commit is contained in:
David Katz
2025-10-07 19:05:05 -04:00
parent 00302f5c89
commit 24fe252b73
2 changed files with 113 additions and 35 deletions
+2 -31
View File
@@ -167,8 +167,7 @@ pub async fn perform_compaction(agent: &Agent, messages: &[Message]) -> Result<A
/// Check if messages need compaction and compact them if necessary
///
/// This is a convenience wrapper function that combines checking and compaction.
/// If the most recent message is a user message, it will be preserved by removing it
/// before compaction and adding it back afterwards.
/// Uses perform_compaction internally to handle the actual compaction process.
///
/// # Arguments
/// * `agent` - The agent to use for context management
@@ -207,35 +206,7 @@ pub async fn check_and_compact_messages(
check_result.usage_ratio * 100.0
);
// Check if the most recent message is a user message
let (messages_to_compact, preserved_user_message) = if let Some(last_message) = messages.last()
{
if matches!(last_message.role, rmcp::model::Role::User) {
// Remove the last user message before auto-compaction
(&messages[..messages.len() - 1], Some(last_message.clone()))
} else {
(messages, None)
}
} else {
(messages, None)
};
// Perform the compaction on messages excluding the preserved user message
// The summarize_context method already handles the visibility properly
let (mut summary_messages, _, summarization_usage) =
agent.summarize_context(messages_to_compact).await?;
// Add back the preserved user message if it exists
// (keeps default visibility: both true)
if let Some(user_message) = preserved_user_message {
summary_messages.push(user_message);
}
Ok(AutoCompactResult {
compacted: true,
messages: summary_messages,
summarization_usage,
})
perform_compaction(agent, messages).await
}
#[cfg(test)]
+111 -4
View File
@@ -161,10 +161,35 @@ impl<'a> IntoIterator for &'a Conversation {
/// Fix a conversation that we're about to send to an LLM. So the last and first
/// messages should always be from the user.
///
/// This function handles agent_visible and non-visible messages separately to ensure
/// that only agent-visible messages are fixed, while preserving non-visible messages.
pub fn fix_conversation(conversation: Conversation) -> (Conversation, Vec<String>) {
let messages = conversation.messages().clone();
let (messages, issues) = fix_messages(messages);
(Conversation::new_unvalidated(messages), issues)
let all_messages = conversation.messages().clone();
// Separate agent-visible and non-visible messages
let mut agent_visible_messages = Vec::new();
let mut non_visible_messages = Vec::new();
for msg in all_messages {
if msg.metadata.agent_visible {
agent_visible_messages.push(msg);
} else {
non_visible_messages.push(msg);
}
}
// Only fix agent-visible messages
let (fixed_visible, issues) = fix_messages(agent_visible_messages);
// Combine back: non-visible messages stay unchanged, visible messages are fixed
let mut final_messages = non_visible_messages;
final_messages.extend(fixed_visible);
// Sort by timestamp to maintain order
final_messages.sort_by_key(|m| m.created);
(Conversation::new_unvalidated(final_messages), issues)
}
fn fix_messages(messages: Vec<Message>) -> (Vec<Message>, Vec<String>) {
@@ -389,7 +414,8 @@ fn fix_lead_trail(mut messages: Vec<Message>) -> (Vec<Message>, Vec<String>) {
if let Some(last) = messages.last() {
if last.role == Role::Assistant {
messages.pop();
issues.push("Removed trailing assistant message".to_string());
issues
.push("Removed trailing assistant message with pending tool requests".to_string());
}
}
@@ -752,4 +778,85 @@ mod tests {
panic!("Expected second item to be an image");
}
}
#[test]
fn test_last_assistant_message_with_pending_tool_request() {
// Test the scenario where the last assistant message has a tool request
// This should be removed to prevent orphaned tool responses after compaction
let messages = vec![
Message::user().with_text("Help me with something"),
Message::assistant()
.with_text("I'll help you with that")
.with_tool_request(
"tool_1",
Ok(CallToolRequestParam {
name: "some_tool".into(),
arguments: Some(object!({})),
}),
),
];
let (fixed, issues) = run_verify(messages);
// The entire last assistant message should be removed
assert_eq!(fixed.len(), 1);
assert_eq!(issues.len(), 2);
// Should have removed the orphaned tool request and the trailing assistant message
assert!(issues
.iter()
.any(|i| i.contains("Removed orphaned tool request 'tool_1'")));
assert!(issues
.iter()
.any(|i| i.contains("Removed trailing assistant message")));
// Only the user message should remain
assert_eq!(fixed[0].role, Role::User);
assert!(fixed[0].as_concat_text().contains("Help me with something"));
}
#[test]
fn test_last_assistant_message_with_multiple_pending_tool_requests() {
// Test with multiple tool requests in the last assistant message
let messages = vec![
Message::user().with_text("Do multiple things"),
Message::assistant()
.with_text("I'll do multiple things")
.with_tool_request(
"tool_1",
Ok(CallToolRequestParam {
name: "tool_a".into(),
arguments: Some(object!({})),
}),
)
.with_tool_request(
"tool_2",
Ok(CallToolRequestParam {
name: "tool_b".into(),
arguments: Some(object!({})),
}),
),
];
let (fixed, issues) = run_verify(messages);
// The entire last assistant message should be removed
assert_eq!(fixed.len(), 1);
assert_eq!(issues.len(), 3);
// Should have removed both orphaned tool requests and the trailing assistant message
assert!(issues
.iter()
.any(|i| i.contains("Removed orphaned tool request 'tool_1'")));
assert!(issues
.iter()
.any(|i| i.contains("Removed orphaned tool request 'tool_2'")));
assert!(issues
.iter()
.any(|i| i.contains("Removed trailing assistant message")));
// Only the user message should remain
assert_eq!(fixed[0].role, Role::User);
assert!(fixed[0].as_concat_text().contains("Do multiple things"));
}
}