agent: Respect compaction boundaries in request windows (#58241)
Agent request construction now honors persisted compaction markers by rebuilding the model-visible window from the latest boundary, retaining recent user context for summary compactions, and treating provider-native compactions as suffix-only until native replay lands. Release Notes: - N/A Co-authored-by: Bennet Bo Fenner <bennetbo@gmx.de>
This commit is contained in:
Richard FeldmanandBennet Bo Fenner authored and GitHub committed
1 parent
0584c934dd
commit
6d3015d2ba
1 file changed
+301
-16
+301
-16
@@ -40,8 +40,8 @@ use language_model::{
|
||||
LanguageModelId, LanguageModelImage, LanguageModelProviderId, LanguageModelRegistry,
|
||||
LanguageModelRequest, LanguageModelRequestMessage, LanguageModelRequestTool,
|
||||
LanguageModelToolResult, LanguageModelToolResultContent, LanguageModelToolSchemaFormat,
|
||||
LanguageModelToolUse, LanguageModelToolUseId, Role, SelectedModel, Speed, StopReason,
|
||||
TokenUsage, ZED_CLOUD_PROVIDER_ID,
|
||||
LanguageModelToolUse, LanguageModelToolUseId, MessageContent, Role, SelectedModel, Speed,
|
||||
StopReason, TokenUsage, ZED_CLOUD_PROVIDER_ID,
|
||||
};
|
||||
use project::Project;
|
||||
use prompt_store::ProjectContext;
|
||||
@@ -68,6 +68,9 @@ const TOOL_CANCELED_MESSAGE: &str = "Tool canceled by user";
|
||||
pub const MAX_TOOL_NAME_LENGTH: usize = 64;
|
||||
pub const MAX_SUBAGENT_DEPTH: u8 = 1;
|
||||
|
||||
// Using the heuristic that 1 token is about 4 bytes, keep the last 80K bytes of user-message content (~20k tokens).
|
||||
const COMPACTION_RETAINED_USER_MESSAGES_BYTE_BUDGET: usize = 80_000;
|
||||
|
||||
/// Returned when a turn is attempted but no language model has been selected.
|
||||
#[derive(Debug)]
|
||||
pub struct NoModelConfiguredError;
|
||||
@@ -3232,9 +3235,7 @@ impl Thread {
|
||||
cache: false,
|
||||
reasoning_details: None,
|
||||
}];
|
||||
for message in &self.messages {
|
||||
messages.extend(message.to_request());
|
||||
}
|
||||
self.extend_request_history(&mut messages);
|
||||
|
||||
if let Some(last_message) = messages.last_mut() {
|
||||
last_message.cache = true;
|
||||
@@ -3247,6 +3248,67 @@ impl Thread {
|
||||
messages
|
||||
}
|
||||
|
||||
fn extend_request_history(&self, messages: &mut Vec<LanguageModelRequestMessage>) {
|
||||
let Some(compaction_ix) = self.latest_compaction_message_ix() else {
|
||||
for message in &self.messages {
|
||||
messages.extend(message.to_request());
|
||||
}
|
||||
return;
|
||||
};
|
||||
|
||||
if matches!(
|
||||
&*self.messages[compaction_ix],
|
||||
Message::Compaction(CompactionInfo::Summary(_))
|
||||
) {
|
||||
messages.extend(self.retained_user_request_messages_before(compaction_ix));
|
||||
}
|
||||
|
||||
for message in &self.messages[compaction_ix..] {
|
||||
messages.extend(message.to_request());
|
||||
}
|
||||
}
|
||||
|
||||
fn latest_compaction_message_ix(&self) -> Option<usize> {
|
||||
self.messages
|
||||
.iter()
|
||||
.rposition(|message| matches!(&**message, Message::Compaction(_)))
|
||||
}
|
||||
|
||||
fn retained_user_request_messages_before(
|
||||
&self,
|
||||
compaction_ix: usize,
|
||||
) -> Vec<LanguageModelRequestMessage> {
|
||||
let mut remaining_bytes = COMPACTION_RETAINED_USER_MESSAGES_BYTE_BUDGET;
|
||||
let mut retained_messages = Vec::new();
|
||||
|
||||
for message in self.messages[..compaction_ix].iter().rev() {
|
||||
let Message::User(user_message) = &**message else {
|
||||
continue;
|
||||
};
|
||||
if user_message.content.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
let request_message = user_message.to_request();
|
||||
let byte_count = user_message_byte_len(&request_message);
|
||||
if let Some(bytes) = remaining_bytes.checked_sub(byte_count) {
|
||||
remaining_bytes = bytes;
|
||||
retained_messages.push(request_message);
|
||||
} else {
|
||||
if remaining_bytes > 0
|
||||
&& let Some(request_message) =
|
||||
truncate_user_message_to_byte_budget(request_message, remaining_bytes)
|
||||
{
|
||||
retained_messages.push(request_message);
|
||||
}
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
retained_messages.reverse();
|
||||
retained_messages
|
||||
}
|
||||
|
||||
pub fn to_markdown(&self) -> String {
|
||||
let mut markdown = String::new();
|
||||
for (ix, message) in self.messages.iter().enumerate() {
|
||||
@@ -3385,6 +3447,83 @@ fn total_input_tokens(usage: language_model::TokenUsage) -> u64 {
|
||||
.saturating_add(usage.cache_read_input_tokens)
|
||||
}
|
||||
|
||||
fn user_message_byte_len(message: &LanguageModelRequestMessage) -> usize {
|
||||
message
|
||||
.content
|
||||
.iter()
|
||||
.map(|content| match content {
|
||||
MessageContent::Text(text) => text.len(),
|
||||
MessageContent::Image(image) => image.len(),
|
||||
// These can never occur in a user message
|
||||
MessageContent::Thinking { .. }
|
||||
| MessageContent::RedactedThinking(_)
|
||||
| MessageContent::ToolResult(_)
|
||||
| MessageContent::ToolUse(_) => 0,
|
||||
})
|
||||
.sum()
|
||||
}
|
||||
|
||||
fn truncate_user_message_to_byte_budget(
|
||||
mut message: LanguageModelRequestMessage,
|
||||
byte_budget: usize,
|
||||
) -> Option<LanguageModelRequestMessage> {
|
||||
let mut remaining_bytes = byte_budget;
|
||||
let mut content = Vec::with_capacity(message.content.len());
|
||||
|
||||
for item in message.content {
|
||||
match item {
|
||||
MessageContent::Text(text) => {
|
||||
let fits = text.len() <= remaining_bytes;
|
||||
if let Some(text) = take_text_within_byte_budget(text, &mut remaining_bytes) {
|
||||
content.push(MessageContent::Text(text));
|
||||
}
|
||||
if !fits {
|
||||
break;
|
||||
}
|
||||
}
|
||||
MessageContent::Image(image) => {
|
||||
let byte_len = image.len();
|
||||
if let Some(bytes) = remaining_bytes.checked_sub(byte_len) {
|
||||
remaining_bytes = bytes;
|
||||
content.push(MessageContent::Image(image));
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
// These can never occur in a user message
|
||||
MessageContent::Thinking { .. }
|
||||
| MessageContent::RedactedThinking(_)
|
||||
| MessageContent::ToolResult(_)
|
||||
| MessageContent::ToolUse(_) => {}
|
||||
}
|
||||
}
|
||||
|
||||
if content.is_empty() {
|
||||
None
|
||||
} else {
|
||||
message.content = content;
|
||||
Some(message)
|
||||
}
|
||||
}
|
||||
|
||||
fn take_text_within_byte_budget(text: String, remaining_bytes: &mut usize) -> Option<String> {
|
||||
if text.is_empty() || *remaining_bytes == 0 {
|
||||
return None;
|
||||
}
|
||||
|
||||
if let Some(bytes) = remaining_bytes.checked_sub(text.len()) {
|
||||
*remaining_bytes = bytes;
|
||||
return Some(text);
|
||||
}
|
||||
|
||||
let end = text.floor_char_boundary((*remaining_bytes).min(text.len()));
|
||||
*remaining_bytes = 0;
|
||||
|
||||
let text = text[..end].to_string();
|
||||
|
||||
if text.is_empty() { None } else { Some(text) }
|
||||
}
|
||||
|
||||
struct RunningTurn {
|
||||
/// Holds the task that handles agent interaction until the end of the turn.
|
||||
/// Survives across multiple requests as the model performs tool calls and
|
||||
@@ -4602,6 +4741,38 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
fn user_text_message(id: UserMessageId, text: &str) -> Arc<Message> {
|
||||
Arc::new(Message::User(UserMessage {
|
||||
id,
|
||||
content: vec![UserMessageContent::Text(text.to_string())].into(),
|
||||
}))
|
||||
}
|
||||
|
||||
fn agent_text_message(text: &str) -> Arc<Message> {
|
||||
Arc::new(Message::Agent(AgentMessage {
|
||||
content: vec![AgentMessageContent::Text(text.to_string())],
|
||||
..Default::default()
|
||||
}))
|
||||
}
|
||||
|
||||
fn summary_compaction(summary: &str) -> Arc<Message> {
|
||||
Arc::new(Message::Compaction(CompactionInfo::Summary(summary.into())))
|
||||
}
|
||||
|
||||
fn summary_request_text(summary: &str) -> String {
|
||||
format!(
|
||||
"The previous conversation was compacted. Use this summary as context:\n\n{summary}"
|
||||
)
|
||||
}
|
||||
|
||||
fn request_texts_after_system(messages: &[LanguageModelRequestMessage]) -> Vec<String> {
|
||||
messages
|
||||
.iter()
|
||||
.skip(1)
|
||||
.map(LanguageModelRequestMessage::string_contents)
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[gpui::test]
|
||||
async fn test_replay_emits_context_compaction(cx: &mut TestAppContext) {
|
||||
let (thread, _event_stream) = setup_thread_for_test(cx).await;
|
||||
@@ -4609,19 +4780,11 @@ mod tests {
|
||||
|
||||
let mut replay_events = cx.update(|cx| {
|
||||
thread.update(cx, |thread, cx| {
|
||||
thread.messages.push(Arc::new(Message::User(UserMessage {
|
||||
id: user_message_id.clone(),
|
||||
content: vec![UserMessageContent::Text("before".to_string())].into(),
|
||||
})));
|
||||
thread
|
||||
.messages
|
||||
.push(Arc::new(Message::Compaction(CompactionInfo::Summary(
|
||||
"summary".into(),
|
||||
))));
|
||||
thread.messages.push(Arc::new(Message::Agent(AgentMessage {
|
||||
content: vec![AgentMessageContent::Text("after".to_string())],
|
||||
..Default::default()
|
||||
})));
|
||||
.push(user_text_message(user_message_id.clone(), "before"));
|
||||
thread.messages.push(summary_compaction("summary"));
|
||||
thread.messages.push(agent_text_message("after"));
|
||||
|
||||
thread.replay(cx)
|
||||
})
|
||||
@@ -4649,6 +4812,128 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[gpui::test]
|
||||
async fn test_native_compaction_boundary(cx: &mut TestAppContext) {
|
||||
let (thread, _event_stream) = setup_thread_for_test(cx).await;
|
||||
|
||||
let request_messages = cx.update(|cx| {
|
||||
thread.update(cx, |thread, cx| {
|
||||
thread
|
||||
.messages
|
||||
.push(user_text_message(UserMessageId::new(), "before native"));
|
||||
thread.messages.push(Arc::new(Message::Compaction(
|
||||
CompactionInfo::ProviderNative {
|
||||
provider: LanguageModelProviderId::from("openai".to_string()),
|
||||
items: vec![json!({"type": "compaction"})],
|
||||
},
|
||||
)));
|
||||
thread
|
||||
.messages
|
||||
.push(user_text_message(UserMessageId::new(), "after native"));
|
||||
|
||||
thread.build_request_messages(Vec::new(), cx)
|
||||
})
|
||||
});
|
||||
|
||||
assert_eq!(
|
||||
request_texts_after_system(&request_messages),
|
||||
vec!["after native".to_string()]
|
||||
);
|
||||
}
|
||||
|
||||
#[gpui::test]
|
||||
async fn test_retained_users_truncate_oldest(cx: &mut TestAppContext) {
|
||||
let (thread, _event_stream) = setup_thread_for_test(cx).await;
|
||||
let mut long_text = "START".to_string();
|
||||
long_text.push_str(&"x".repeat(COMPACTION_RETAINED_USER_MESSAGES_BYTE_BUDGET));
|
||||
long_text.push_str("END");
|
||||
|
||||
let request_messages = cx.update(|cx| {
|
||||
thread.update(cx, |thread, cx| {
|
||||
thread.messages.push(user_text_message(
|
||||
UserMessageId::new(),
|
||||
"dropped older user",
|
||||
));
|
||||
thread
|
||||
.messages
|
||||
.push(agent_text_message("dropped assistant"));
|
||||
thread
|
||||
.messages
|
||||
.push(user_text_message(UserMessageId::new(), &long_text));
|
||||
thread
|
||||
.messages
|
||||
.push(user_text_message(UserMessageId::new(), "new"));
|
||||
thread.messages.push(summary_compaction("summary context"));
|
||||
thread.messages.push(agent_text_message("after assistant"));
|
||||
thread
|
||||
.messages
|
||||
.push(user_text_message(UserMessageId::new(), "after user"));
|
||||
|
||||
thread.build_request_messages(Vec::new(), cx)
|
||||
})
|
||||
});
|
||||
|
||||
let request_texts = request_texts_after_system(&request_messages);
|
||||
assert_eq!(request_texts.len(), 5);
|
||||
assert_eq!(
|
||||
request_texts[0],
|
||||
format!(
|
||||
"START{}",
|
||||
"x".repeat(
|
||||
COMPACTION_RETAINED_USER_MESSAGES_BYTE_BUDGET - "START".len() - "new".len()
|
||||
)
|
||||
)
|
||||
);
|
||||
assert_eq!(request_texts[1], "new");
|
||||
assert_eq!(request_texts[2], summary_request_text("summary context"));
|
||||
assert_eq!(request_texts[3], "after assistant");
|
||||
assert_eq!(request_texts[4], "after user");
|
||||
assert!(request_texts.iter().all(
|
||||
|text| !text.contains("dropped older user") && !text.contains("dropped assistant")
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_truncate_text_utf8_boundary() {
|
||||
let message = LanguageModelRequestMessage {
|
||||
role: Role::User,
|
||||
content: vec![MessageContent::Text("hello 👋 world".to_string())],
|
||||
cache: false,
|
||||
reasoning_details: None,
|
||||
};
|
||||
|
||||
let truncated = truncate_user_message_to_byte_budget(message, 8).unwrap();
|
||||
assert_eq!(
|
||||
truncated.content,
|
||||
vec![MessageContent::Text("hello ".to_string())]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_truncate_keeps_fitting_images() {
|
||||
let image = LanguageModelImage {
|
||||
source: "image".into(),
|
||||
};
|
||||
let message = LanguageModelRequestMessage {
|
||||
role: Role::User,
|
||||
content: vec![
|
||||
MessageContent::Text("abc".to_string()),
|
||||
MessageContent::Image(image.clone()),
|
||||
],
|
||||
cache: false,
|
||||
reasoning_details: None,
|
||||
};
|
||||
|
||||
let truncated = truncate_user_message_to_byte_budget(message, 8).unwrap();
|
||||
assert_eq!(
|
||||
truncated.content,
|
||||
vec![
|
||||
MessageContent::Text("abc".to_string()),
|
||||
MessageContent::Image(image),
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
fn setup_parent_with_subagents(
|
||||
cx: &mut TestAppContext,
|
||||
parent: &Entity<Thread>,
|
||||
|
||||
Reference in new issue
Block a user