yenru0/zed
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Enable compaction for subagents (#64135)

Browse Source 
Release Notes:

- Enabled compaction for subagents

---------

Co-authored-by: zed-zippy[bot] <234243425+zed-zippy[bot]@users.noreply.github.com>
This commit is contained in:
Kirill Bulatovandzed-zippy[bot] authored and GitHub committed 2026-09-16 16:03:30 +00:00
1 parent b84803f56f
commit 2433f1d15d
4 files changed
+839 -398
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

crates/agent/src/agent.rs
+77 -86
View File
@@ -3440,14 +3440,6 @@ impl ThreadEnvironment for NativeThreadEnvironment {
}
}
#[derive(Debug, Clone)]
enum SubagentPromptResult {
Completed,
Cancelled,
ContextWindowWarning,
Error(String),
}
pub struct NativeSubagentHandle {
session_id: acp::SessionId,
parent_thread: WeakEntity<Thread>,
@@ -3487,99 +3479,87 @@ impl SubagentHandle for NativeSubagentHandle {
let parent_thread = self.parent_thread.clone();
cx.spawn(async move |cx| {
let (task, _subscription) = cx.update(|cx| {
let ratio_before_prompt = thread
.read(cx)
.latest_token_usage()
.map(|usage| usage.ratio());
let (task, token_limit_rx, _subscription) = cx.update(|cx| {
parent_thread
.update(cx, |parent_thread, _cx| {
parent_thread.register_running_subagent(thread.downgrade())
})
.ok();
let task = acp_thread.update(cx, |acp_thread, cx| {
acp_thread.send(vec![message.into()], cx)
});
let (token_limit_tx, token_limit_rx) = oneshot::channel::<()>();
let ratio_before_prompt = thread
.read(cx)
.latest_token_usage()
.map_or(TokenUsageRatio::Normal, |usage| usage.ratio());
let (token_limit_tx, token_limit_rx) = oneshot::channel();
let mut token_limit_tx = Some(token_limit_tx);
let subscription = cx.subscribe(
&thread,
move |_thread, event: &TokenUsageUpdated, _cx| {
if let Some(usage) = &event.0 {
let old_ratio = ratio_before_prompt
.clone()
.unwrap_or(TokenUsageRatio::Normal);
let new_ratio = usage.ratio();
if old_ratio == TokenUsageRatio::Normal
&& new_ratio == TokenUsageRatio::Warning
{
if let Some(tx) = token_limit_tx.take() {
tx.send(()).ok();
}
}
move |thread, event: &TokenUsageUpdated, cx| {
if !thread.read(cx).auto_compaction_enabled(cx)
&& event.0.as_ref().is_some_and(|usage| usage.ratio() > ratio_before_prompt)
&& let Some(sender) = token_limit_tx.take()
{
sender.send(()).ok();
}
},
);
let wait_for_prompt = cx
.background_spawn(async move {
futures::select! {
response = task.fuse() => match response {
Ok(Some(response)) => {
match response.stop_reason {
acp::StopReason::Cancelled => SubagentPromptResult::Cancelled,
acp::StopReason::MaxTokens => SubagentPromptResult::Error("The agent reached the maximum number of tokens.".into()),
acp::StopReason::MaxTurnRequests => SubagentPromptResult::Error("The agent reached the maximum number of allowed requests between user turns. Try prompting again.".into()),
acp::StopReason::Refusal => SubagentPromptResult::Error("The agent refused to process that prompt. Try again.".into()),
acp::StopReason::EndTurn | _ => SubagentPromptResult::Completed,
}
}
Ok(None) => SubagentPromptResult::Error("No response from the agent. You can try messaging again.".into()),
Err(error) => SubagentPromptResult::Error(error.to_string()),
},
_ = token_limit_rx.fuse() => SubagentPromptResult::ContextWindowWarning,
}
});
(wait_for_prompt, subscription)
let task = acp_thread.update(cx, |acp_thread, cx| {
acp_thread.send(vec![message.into()], cx)
});
(task, token_limit_rx, subscription)
});
let result = match task.await {
SubagentPromptResult::Completed => thread.read_with(cx, |thread, _cx| {
thread
.last_message()
.and_then(|message| {
let content = message.as_agent_message()?
.content
.iter()
.filter_map(|c| match c {
AgentMessageContent::Text(text) => Some(text.as_str()),
_ => None,
})
.join("\n\n");
if content.is_empty() {
None
} else {
Some( content)
}
})
.context("No response from subagent")
}),
SubagentPromptResult::Cancelled => Err(anyhow!("User canceled")),
SubagentPromptResult::Error(message) => Err(anyhow!("{message}")),
SubagentPromptResult::ContextWindowWarning => {
thread.update(cx, |thread, cx| thread.cancel(cx)).await;
Err(anyhow!(
"The agent is nearing the end of its context window and has been \
stopped. You can prompt the thread again to have the agent wrap up \
or hand off its work."
))
let mut task = task.fuse();
let response = futures::select_biased! {
response = task => response,
_ = token_limit_rx.fuse() => {
if thread.read_with(cx, |thread, _| thread.is_turn_complete()) {
task.await
} else {
thread.update(cx, |thread, cx| thread.cancel(cx)).await;
Err(anyhow!(
"The agent is nearing the end of its context window and has been \
stopped. You can prompt the thread again to have the agent wrap up \
or hand off its work."
))
}
}
};
let discard_partial_output = matches!(
&response,
Ok(Some(response)) if response.stop_reason == acp::StopReason::Cancelled
|| response.stop_reason == acp::StopReason::Refusal
);
let result = match response {
Ok(Some(response)) => match response.stop_reason {
acp::StopReason::Cancelled => Err(anyhow!("User canceled")),
acp::StopReason::MaxTokens => Err(anyhow!("The agent reached the maximum number of tokens.")),
acp::StopReason::MaxTurnRequests => Err(anyhow!("The agent reached the maximum number of allowed requests between user turns. Try prompting again.")),
acp::StopReason::Refusal => Err(anyhow!("The agent refused to process that prompt. Try again.")),
_ => thread.read_with(cx, |thread, _cx| {
thread
.last_message()
.and_then(|message| {
let content = message.as_agent_message()?
.content
.iter()
.filter_map(|content| match content {
AgentMessageContent::Text(text) => Some(text.as_str()),
_ => None,
})
.join("\n\n");
if content.is_empty() {
None
} else {
Some(content)
}
})
.context("No response from subagent")
}),
},
Ok(None) => Err(anyhow!("No response from the agent. You can try messaging again.")),
Err(error) => Err(error),
};
parent_thread
.update(cx, |parent_thread, cx| {
@@ -3587,7 +3567,18 @@ impl SubagentHandle for NativeSubagentHandle {
})
.ok();
result
if discard_partial_output {
result
} else {
result.map_err(|error| {
let partial_output = thread.read_with(cx, |thread, _| thread.subagent_partial_output());
if partial_output.is_empty() {
anyhow!("{error:#}")
} else {
anyhow!("{error:#}\n\nPartial subagent output (last 3 messages, up to 4096 characters each):\n\n{partial_output}")
}
})
}
})
}
}
crates/agent/src/tests/mod.rs
+570 -303
View File
@@ -4,7 +4,7 @@ use acp_thread::{
ThreadStatus,
};
use agent_client_protocol::schema::v1 as acp;
use agent_settings::AgentProfileId;
use agent_settings::{AgentProfileId, AgentSettings, AutoCompactThreshold, COMPACTION_PROMPT};
use anyhow::Result;
use client::{Client, RefreshLlmTokenListener, UserStore};
use collections::IndexMap;
@@ -6586,307 +6586,361 @@ async fn test_parent_cancel_stops_subagent(cx: &mut TestAppContext) {
}
#[gpui::test]
async fn test_subagent_context_window_warning(cx: &mut TestAppContext) {
init_test(cx);
cx.update(|cx| {
LanguageModelRegistry::test(cx);
});
cx.update(|cx| {
cx.update_flags(true, vec!["subagents".to_string()]);
});
async fn test_subagent_auto_compaction(cx: &mut TestAppContext) {
let test = SubagentCompactionTest::new(cx).await;
let send = test.send("subagent task prompt", cx);
test.tool_round(900_000, cx);
let fs = FakeFs::new(cx.executor());
fs.insert_tree(
"/",
json!({
"a": {
"b.md": "Lorem"
}
}),
)
.await;
let project = Project::test(fs.clone(), [path!("/a").as_ref()], cx).await;
let thread_store = cx.new(|cx| ThreadStore::new(cx));
let agent =
cx.update(|cx| NativeAgent::new(thread_store.clone(), Templates::new(), fs.clone(), cx));
let connection = Rc::new(NativeAgentConnection(agent.clone()));
let request = test.compaction_request();
assert_eq!(request.thread_id, Some(test.handle.id().to_string()));
assert_eq!(
request.messages.last().unwrap().string_contents(),
COMPACTION_PROMPT
);
assert_eq!(
request
.messages
.iter()
.flat_map(|message| &message.content)
.filter_map(|content| match content {
MessageContent::ToolResult(result) => Some(result.text_contents()),
_ => None,
})
.collect::<Vec<_>>(),
vec!["tool output"],
);
test.fail_compaction(
http_client::StatusCode::SERVICE_UNAVAILABLE,
Some(Duration::from_secs(3)),
cx,
);
test.assert_running(cx);
assert_eq!(test.model.pending_completions(), Vec::new());
cx.executor()
.advance_clock(crate::maximum_retry_delay_with_jitter(Duration::from_secs(
3,
)));
cx.run_until_parked();
let retry_request = test.compaction_request();
assert_eq!(retry_request.messages, request.messages);
let request = retry_request;
let acp_thread = cx
test.model
.send_completion_stream_text_chunk(&request, "subagent summary");
test.model.end_completion_stream(&request);
cx.run_until_parked();
let request = test.model.pending_completions().pop().unwrap();
assert_eq!(request.intent, Some(CompletionIntent::ToolResults));
assert_eq!(request.thread_id, Some(test.handle.id().to_string()));
assert_eq!(
request
.messages
.iter()
.filter(|message| message.role != Role::System)
.map(|message| message.string_contents())
.collect::<Vec<_>>(),
vec![
"subagent task prompt",
"The previous conversation was compacted. Use this summary as context:\n\nsubagent summary",
],
);
test.model
.send_completion_stream_text_chunk(&request, "subagent answer");
test.model.end_completion_stream(&request);
assert_eq!(send.await.unwrap(), "subagent answer");
test.assert_stopped(cx);
let handle = cx
.update(|cx| {
connection
.clone()
.new_session(project.clone(), PathList::new(&[Path::new("")]), cx)
test.environment
.resume_subagent_thread(test.handle.id(), cx)
})
.await
.unwrap();
let session_id = acp_thread.read_with(cx, |thread, _| thread.session_id().clone());
let thread = agent.read_with(cx, |agent, _| {
agent.sessions.get(&session_id).unwrap().thread.clone()
});
let model = Arc::new(FakeLanguageModel::default());
thread.update(cx, |thread, cx| {
thread.set_model(model.clone(), cx);
});
assert_eq!(handle.id(), test.handle.id());
let send = cx.update(|cx| handle.send("follow-up task".to_string(), &cx.to_async()));
cx.run_until_parked();
// Start the parent turn
let send = acp_thread.update(cx, |thread, cx| thread.send_raw("Prompt", cx));
cx.run_until_parked();
model.send_last_completion_stream_text_chunk("spawning subagent");
let subagent_tool_input = SpawnAgentToolInput {
label: "label".to_string(),
message: "subagent task prompt".to_string(),
session_id: None,
};
let subagent_tool_use = LanguageModelToolUse {
id: "subagent_1".into(),
name: SpawnAgentTool::NAME.into(),
raw_input: serde_json::to_string(&subagent_tool_input).unwrap(),
input: language_model::LanguageModelToolUseInput::Json(
serde_json::to_value(&subagent_tool_input).unwrap(),
),
is_input_complete: true,
thought_signature: None,
};
model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
subagent_tool_use,
));
model.end_last_completion_stream();
cx.run_until_parked();
// Verify subagent is running
let subagent_session_id = thread.read_with(cx, |thread, cx| {
thread
.running_subagent_ids(cx)
.get(0)
.expect("subagent thread should be running")
.clone()
});
// Send a usage update that crosses the warning threshold (80% of 1,000,000)
model.send_last_completion_stream_text_chunk("partial work");
model.send_last_completion_stream_event(LanguageModelCompletionEvent::UsageUpdate(
TokenUsage {
input_tokens: 850_000,
output_tokens: 0,
cache_creation_input_tokens: 0,
cache_read_input_tokens: 0,
},
));
cx.run_until_parked();
// The subagent should no longer be running
thread.read_with(cx, |thread, cx| {
assert!(
thread.running_subagent_ids(cx).is_empty(),
"subagent should be stopped after context window warning"
);
});
// The parent model should get a new completion request to respond to the tool error
model.send_last_completion_stream_text_chunk("Response after warning");
model.end_last_completion_stream();
send.await.unwrap();
// Verify the parent thread shows the warning error in the tool call
let markdown = acp_thread.read_with(cx, |thread, cx| thread.to_markdown(cx));
assert!(
markdown.contains("nearing the end of its context window"),
"tool output should contain context window warning message, got:\n{markdown}"
);
assert!(
markdown.contains("Status: Failed"),
"tool call should have Failed status, got:\n{markdown}"
let request = test.model.pending_completions().pop().unwrap();
assert_eq!(request.intent, Some(CompletionIntent::Subagent));
assert_eq!(
request.messages.last().unwrap().string_contents(),
"follow-up task"
);
test.model
.send_completion_stream_text_chunk(&request, "follow-up answer");
test.model.end_completion_stream(&request);
assert_eq!(send.await.unwrap(), "follow-up answer");
test.assert_stopped(cx);
// Verify the subagent session still exists (can be resumed)
agent.read_with(cx, |agent, _cx| {
assert!(
agent.sessions.contains_key(&subagent_session_id),
"subagent session should still exist for potential resume"
for use_tool in [false, true] {
let send = test.send("refused follow-up", cx);
if use_tool {
test.tool_round(0, cx);
}
test.model
.send_last_completion_stream_event(LanguageModelCompletionEvent::Stop(
StopReason::Refusal,
));
test.model.end_last_completion_stream();
assert_eq!(
send.await.unwrap_err().to_string(),
"The agent refused to process that prompt. Try again."
);
});
test.assert_stopped(cx);
}
}
#[gpui::test]
async fn test_subagent_no_context_window_warning_when_already_at_warning(cx: &mut TestAppContext) {
init_test(cx);
cx.update(|cx| {
LanguageModelRegistry::test(cx);
});
cx.update(|cx| {
cx.update_flags(true, vec!["subagents".to_string()]);
});
async fn test_subagent_compaction_respects_settings_and_context_window(cx: &mut TestAppContext) {
for (enabled, max_tokens, max_output_tokens, input_tokens) in [
(true, 1_000_000, None, 899_999),
(false, 1_000_000, None, 799_999),
(true, 79_999, None, 63_999),
(true, 100_000, Some(20_001), 79_999),
] {
let test = SubagentCompactionTest::new(cx).await;
test.configure_compaction(enabled, max_tokens, max_output_tokens, cx);
let send = test.send("subagent task prompt", cx);
test.tool_round(input_tokens, cx);
let request = test.model.pending_completions().pop().unwrap();
assert_eq!(request.intent, Some(CompletionIntent::ToolResults));
test.assert_running(cx);
test.model
.send_completion_stream_text_chunk(&request, "subagent answer");
test.model.end_completion_stream(&request);
assert_eq!(send.await.unwrap(), "subagent answer");
test.assert_stopped(cx);
}
}
let fs = FakeFs::new(cx.executor());
fs.insert_tree(
"/",
json!({
"a": {
"b.md": "Lorem"
}
}),
)
.await;
let project = Project::test(fs.clone(), [path!("/a").as_ref()], cx).await;
let thread_store = cx.new(|cx| ThreadStore::new(cx));
let agent =
cx.update(|cx| NativeAgent::new(thread_store.clone(), Templates::new(), fs.clone(), cx));
let connection = Rc::new(NativeAgentConnection(agent.clone()));
let acp_thread = cx
.update(|cx| {
connection
.clone()
.new_session(project.clone(), PathList::new(&[Path::new("")]), cx)
})
.await
.unwrap();
let session_id = acp_thread.read_with(cx, |thread, _| thread.session_id().clone());
let thread = agent.read_with(cx, |agent, _| {
agent.sessions.get(&session_id).unwrap().thread.clone()
});
let model = Arc::new(FakeLanguageModel::default());
thread.update(cx, |thread, cx| {
thread.set_model(model.clone(), cx);
});
cx.run_until_parked();
// === First turn: create subagent, trigger context window warning ===
let send = acp_thread.update(cx, |thread, cx| thread.send_raw("First prompt", cx));
cx.run_until_parked();
model.send_last_completion_stream_text_chunk("spawning subagent");
let subagent_tool_input = SpawnAgentToolInput {
label: "initial task".to_string(),
message: "do the first task".to_string(),
session_id: None,
};
let subagent_tool_use = LanguageModelToolUse {
id: "subagent_1".into(),
name: SpawnAgentTool::NAME.into(),
raw_input: serde_json::to_string(&subagent_tool_input).unwrap(),
input: language_model::LanguageModelToolUseInput::Json(
serde_json::to_value(&subagent_tool_input).unwrap(),
),
is_input_complete: true,
thought_signature: None,
};
model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
subagent_tool_use,
));
model.end_last_completion_stream();
cx.run_until_parked();
let subagent_session_id = thread.read_with(cx, |thread, cx| {
thread
.running_subagent_ids(cx)
.get(0)
.expect("subagent thread should be running")
.clone()
});
// Subagent sends a usage update that crosses the warning threshold.
// This triggers Normal→Warning, stopping the subagent.
model.send_last_completion_stream_text_chunk("partial work");
model.send_last_completion_stream_event(LanguageModelCompletionEvent::UsageUpdate(
TokenUsage {
input_tokens: 850_000,
output_tokens: 0,
cache_creation_input_tokens: 0,
cache_read_input_tokens: 0,
},
));
cx.run_until_parked();
// Verify the first turn was stopped with a context window warning
thread.read_with(cx, |thread, cx| {
assert!(
thread.running_subagent_ids(cx).is_empty(),
"subagent should be stopped after context window warning"
#[gpui::test]
async fn test_subagent_context_limit_warning_when_compaction_unavailable(cx: &mut TestAppContext) {
for (enabled, max_tokens, max_output_tokens, warning_tokens) in [
(false, 1_000_000, None, 800_000),
(true, 79_999, None, 64_000),
(true, 100_000, Some(20_001), 80_000),
] {
let test = SubagentCompactionTest::new(cx).await;
test.configure_compaction(enabled, max_tokens, max_output_tokens, cx);
let send = test.send("subagent task prompt", cx);
let request = test.model.pending_completions().pop().unwrap();
test.model
.send_completion_stream_text_chunk(&request, "partial work");
test.update_usage(warning_tokens - 1, cx);
test.assert_running(cx);
assert!(!test.model.is_completion_stream_closed(&request));
test.update_usage(warning_tokens, cx);
assert_eq!(
send.await.unwrap_err().to_string(),
format!(
"{SUBAGENT_CONTEXT_LIMIT_WARNING}\n\nPartial subagent output (last 3 messages, up to 4096 characters each):\n\npartial work"
),
);
});
test.assert_stopped(cx);
assert!(test.model.is_completion_stream_closed(&request));
test.model.end_completion_stream(&request);
assert_eq!(test.model.pending_completions(), Vec::new());
test.assert_no_compaction(cx).await;
// Parent model responds to complete first turn
model.send_last_completion_stream_text_chunk("First response");
model.end_last_completion_stream();
let send = test.send("wrap up", cx);
let request = test.model.pending_completions().pop().unwrap();
test.model
.send_completion_stream_text_chunk(&request, "wrapped up");
test.update_usage(warning_tokens + 1, cx);
test.assert_running(cx);
test.model.end_completion_stream(&request);
assert_eq!(send.await.unwrap(), "wrapped up");
test.assert_stopped(cx);
}
}
send.await.unwrap();
let markdown = acp_thread.read_with(cx, |thread, cx| thread.to_markdown(cx));
assert!(
markdown.contains("nearing the end of its context window"),
"first turn should have context window warning, got:\n{markdown}"
#[gpui::test]
async fn test_subagent_context_limit_exceeded_without_partial_output(cx: &mut TestAppContext) {
let test = SubagentCompactionTest::new(cx).await;
test.configure_compaction(false, 1_000_000, None, cx);
let send = test.send("subagent task prompt", cx);
let request = test.model.pending_completions().pop().unwrap();
test.update_usage(1_000_000, cx);
assert_eq!(
send.await.unwrap_err().to_string(),
SUBAGENT_CONTEXT_LIMIT_WARNING,
);
test.assert_stopped(cx);
assert!(test.model.is_completion_stream_closed(&request));
test.model.end_completion_stream(&request);
assert_eq!(test.model.pending_completions(), Vec::new());
}
// === Second turn: resume the same subagent (now at Warning level) ===
let send2 = acp_thread.update(cx, |thread, cx| thread.send_raw("Follow up", cx));
cx.run_until_parked();
model.send_last_completion_stream_text_chunk("resuming subagent");
let resume_tool_input = SpawnAgentToolInput {
label: "follow-up task".to_string(),
message: "do the follow-up task".to_string(),
session_id: Some(subagent_session_id.clone()),
};
let resume_tool_use = LanguageModelToolUse {
id: "subagent_2".into(),
name: SpawnAgentTool::NAME.into(),
raw_input: serde_json::to_string(&resume_tool_input).unwrap(),
input: language_model::LanguageModelToolUseInput::Json(
serde_json::to_value(&resume_tool_input).unwrap(),
#[gpui::test]
async fn test_subagent_context_limit_preserves_terminal_result_during_checkpoint(
cx: &mut TestAppContext,
) {
for provider_error in [false, true] {
let test = SubagentCompactionTest::new_with_files(json!({".git": {}}), cx).await;
test.configure_compaction(false, 1_000_000, None, cx);
let mut send = test.send("subagent task prompt", cx);
let repository = test.thread.read_with(cx, |thread, cx| {
thread
.project()
.read(cx)
.git_store()
.read(cx)
.active_repository()
.unwrap()
});
let (resume_checkpoint, checkpoint_gate) = oneshot::channel::<()>();
let checkpoint_job = repository.update(cx, |repository, _| {
repository.send_job("hold checkpoint", None, move |_, _| checkpoint_gate)
});
test.model
.send_last_completion_stream_text_chunk("partial work");
if provider_error {
test.model.send_last_completion_stream_error(
LanguageModelCompletionError::from_http_status(
LanguageModelProviderName::new("test"),
http_client::StatusCode::BAD_REQUEST,
"terminal provider error".to_string(),
None,
),
);
}
test.model.end_last_completion_stream();
cx.run_until_parked();
test.thread
.read_with(cx, |thread, _| assert!(thread.is_turn_complete()));
assert!((&mut send).now_or_never().is_none());
repository.read_with(cx, |repository, _| {
let queue = repository.job_debug_queue().to_debug_value();
assert_eq!(
queue["entries"]
.as_array()
.unwrap()
.iter()
.filter(|job| job["description"] == "checkpoint" && job["status"] == "Pending")
.count(),
1,
);
});
test.thread.update(cx, |thread, cx| {
assert!(thread.is_turn_complete());
cx.emit(TokenUsageUpdated(Some(acp_thread::TokenUsage {
max_tokens: 1_000_000,
used_tokens: 800_000,
input_tokens: 800_000,
..acp_thread::TokenUsage::default()
})));
});
cx.run_until_parked();
assert!((&mut send).now_or_never().is_none());
test.parent.read_with(cx, |thread, cx| {
assert_eq!(thread.running_subagent_ids(cx), vec![test.handle.id()]);
});
resume_checkpoint.send(()).unwrap();
checkpoint_job.await.unwrap().unwrap();
if provider_error {
assert_eq!(
send.await.unwrap_err().to_string(),
concat!(
"terminal provider error\n\n",
"Partial subagent output (last 3 messages, up to 4096 characters each):\n\n",
"partial work",
),
);
} else {
assert_eq!(send.await.unwrap(), "partial work");
}
test.assert_stopped(cx);
assert_eq!(test.model.pending_completions(), Vec::new());
}
}
#[gpui::test]
async fn test_subagent_compaction_exhausts_transient_retries(cx: &mut TestAppContext) {
let test = SubagentCompactionTest::new(cx).await;
let send = test.send("subagent task prompt", cx);
test.tool_round(950_000, cx);
for request_index in 0..5 {
test.fail_compaction(
http_client::StatusCode::SERVICE_UNAVAILABLE,
Some(Duration::from_secs(3)),
cx,
);
assert_eq!(test.model.pending_completions(), Vec::new());
if request_index < 4 {
test.assert_running(cx);
cx.executor()
.advance_clock(crate::maximum_retry_delay_with_jitter(Duration::from_secs(
3,
)));
cx.run_until_parked();
}
}
assert_eq!(
send.await.unwrap_err().to_string(),
concat!(
"Automatic context compaction failed: compaction provider error\n\n",
"Partial subagent output (last 3 messages, up to 4096 characters each):\n\n",
"partial work",
),
is_input_complete: true,
thought_signature: None,
};
model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(resume_tool_use));
model.end_last_completion_stream();
cx.run_until_parked();
// Subagent responds with tokens still at warning level (no worse).
// Since ratio_before_prompt was already Warning, this should NOT
// trigger the context window warning again.
model.send_last_completion_stream_text_chunk("follow-up task response");
model.send_last_completion_stream_event(LanguageModelCompletionEvent::UsageUpdate(
TokenUsage {
input_tokens: 870_000,
output_tokens: 0,
cache_creation_input_tokens: 0,
cache_read_input_tokens: 0,
},
));
model.end_last_completion_stream();
cx.run_until_parked();
// Parent model responds to complete second turn
model.send_last_completion_stream_text_chunk("Second response");
model.end_last_completion_stream();
send2.await.unwrap();
// The resumed subagent should have completed normally since the ratio
// didn't transition (it was Warning before and stayed at Warning)
let markdown = acp_thread.read_with(cx, |thread, cx| thread.to_markdown(cx));
assert!(
markdown.contains("follow-up task response"),
"resumed subagent should complete normally when already at warning, got:\n{markdown}"
);
// The second tool call should NOT have a context window warning
let second_tool_pos = markdown
.find("follow-up task")
.expect("should find follow-up tool call");
let after_second_tool = &markdown[second_tool_pos..];
assert!(
!after_second_tool.contains("nearing the end of its context window"),
"should NOT contain context window warning for resumed subagent at same level, got:\n{after_second_tool}"
);
test.assert_stopped(cx);
test.assert_no_compaction(cx).await;
cx.executor()
.advance_clock(crate::maximum_retry_delay_with_jitter(Duration::from_secs(
3,
)));
cx.run_until_parked();
assert_eq!(test.model.pending_completions(), Vec::new());
}
#[gpui::test]
async fn test_subagent_compaction_parent_cancellation(cx: &mut TestAppContext) {
let test = SubagentCompactionTest::new(cx).await;
let send = test.send("subagent task prompt", cx);
test.tool_round(950_000, cx);
let request = test.compaction_request();
test.model
.send_completion_stream_text_chunk(&request, "partial summary");
cx.run_until_parked();
assert!(!test.model.is_completion_stream_closed(&request));
test.parent.update(cx, |thread, cx| thread.cancel(cx)).await;
assert_eq!(send.await.unwrap_err().to_string(), "User canceled");
cx.run_until_parked();
assert!(test.model.is_completion_stream_closed(&request));
test.assert_stopped(cx);
test.assert_no_compaction(cx).await;
test.model.end_completion_stream(&request);
cx.run_until_parked();
assert_eq!(test.model.pending_completions(), Vec::new());
}
#[gpui::test]
async fn test_subagent_compaction_error_propagation(cx: &mut TestAppContext) {
for provider_error in [false, true] {
let test = SubagentCompactionTest::new(cx).await;
let send = test.send("subagent task prompt", cx);
test.tool_round(950_000, cx);
let error = if provider_error {
test.fail_compaction(http_client::StatusCode::BAD_REQUEST, None, cx);
"compaction provider error"
} else {
test.model.end_completion_stream(&test.compaction_request());
"Compaction produced an empty summary"
};
assert_eq!(
send.await.unwrap_err().to_string(),
format!(
"Automatic context compaction failed: {error}\n\n\
Partial subagent output (last 3 messages, up to 4096 characters each):\n\n\
partial work"
)
);
test.assert_stopped(cx);
test.assert_no_compaction(cx).await;
assert_eq!(test.model.pending_completions(), Vec::new());
}
}
#[gpui::test]
@@ -6894,6 +6948,9 @@ async fn test_subagent_error_propagation(cx: &mut TestAppContext) {
init_test(cx);
cx.update(|cx| {
LanguageModelRegistry::test(cx);
let mut settings = AgentSettings::get_global(cx).clone();
settings.auto_compact.enabled = false;
AgentSettings::override_global(settings, cx);
});
cx.update(|cx| {
cx.update_flags(true, vec!["subagents".to_string()]);
@@ -6960,44 +7017,63 @@ async fn test_subagent_error_propagation(cx: &mut TestAppContext) {
cx.run_until_parked();
// Verify subagent is running
thread.read_with(cx, |thread, cx| {
assert!(
!thread.running_subagent_ids(cx).is_empty(),
"subagent should be running"
);
let subagent_session_id = thread.read_with(cx, |thread, cx| {
let subagent_ids = thread.running_subagent_ids(cx);
assert_eq!(subagent_ids.len(), 1);
subagent_ids.into_iter().next().unwrap()
});
// The subagent's model returns a non-retryable error
model.send_last_completion_stream_text_chunk("partial work");
model.send_last_completion_stream_error(LanguageModelCompletionError::from_http_status(
LanguageModelProviderName::new("test"),
http_client::StatusCode::PAYLOAD_TOO_LARGE,
"prompt too large".to_string(),
None,
));
model.end_last_completion_stream();
cx.run_until_parked();
// The subagent should no longer be running
thread.read_with(cx, |thread, cx| {
assert!(
thread.running_subagent_ids(cx).is_empty(),
"subagent should not be running after error"
assert_eq!(
thread.running_subagent_ids(cx),
Vec::new(),
"The subagent should no longer be running"
);
});
// The parent model should get a new completion request to respond to the tool error
model.send_last_completion_stream_text_chunk("Response after error");
model.end_last_completion_stream();
let request = model.pending_completions().pop().unwrap();
assert_eq!(request.intent, Some(CompletionIntent::ToolResults));
let tool_results = request
.messages
.iter()
.flat_map(|message| &message.content)
.filter_map(|content| match content {
MessageContent::ToolResult(result) => Some(result),
_ => None,
})
.collect::<Vec<_>>();
assert_eq!(tool_results.len(), 1);
let tool_result = tool_results.first().unwrap();
assert_eq!(tool_result.tool_use_id.to_string(), "subagent_1");
assert_eq!(&*tool_result.tool_name, SpawnAgentTool::NAME);
assert!(tool_result.is_error);
assert_eq!(
tool_result.text_contents(),
json!({
"session_id": subagent_session_id,
"error": concat!(
"prompt too large\n\n",
"Partial subagent output (last 3 messages, up to 4096 characters each):\n\n",
"partial work",
),
})
.to_string(),
);
model.send_completion_stream_text_chunk(&request, "Response after error");
model.end_completion_stream(&request);
send.await.unwrap();
// Verify the parent thread shows the error in the tool call
let markdown = acp_thread.read_with(cx, |thread, cx| thread.to_markdown(cx));
assert!(
markdown.contains("Status: Failed"),
"tool call should have Failed status after model error, got:\n{markdown}"
);
}
#[gpui::test]
@@ -8743,3 +8819,194 @@ async fn test_mid_turn_model_and_settings_refresh(cx: &mut TestAppContext) {
// Thinking should now be enabled.
assert!(model_b_completions[0].thinking_allowed);
}
const SUBAGENT_CONTEXT_LIMIT_WARNING: &str = "The agent is nearing the end of its context window and has been stopped. You can prompt the thread again to have the agent wrap up or hand off its work.";
struct SubagentCompactionTest {
_agent: Entity<NativeAgent>,
_acp_thread: Entity<AcpThread>,
parent: Entity<Thread>,
thread: Entity<Thread>,
environment: NativeThreadEnvironment,
handle: Rc<dyn SubagentHandle>,
model: Arc<FakeLanguageModel>,
}
impl SubagentCompactionTest {
async fn new(cx: &mut TestAppContext) -> Self {
Self::new_with_files(json!({}), cx).await
}
async fn new_with_files(files: serde_json::Value, cx: &mut TestAppContext) -> Self {
init_test(cx);
cx.update(|cx| {
LanguageModelRegistry::test(cx);
let mut settings = AgentSettings::get_global(cx).clone();
settings.auto_compact.enabled = true;
settings.auto_compact.threshold = AutoCompactThreshold::Percentage(0.9);
for profile in settings.profiles.values_mut() {
profile
.tools
.insert(EchoTool::NAME.to_string().into(), true);
}
AgentSettings::override_global(settings, cx);
});
let fs = FakeFs::new(cx.executor());
fs.insert_tree(path!("/test"), files).await;
let project = Project::test(fs.clone(), [path!("/test").as_ref()], cx).await;
let store = cx.new(ThreadStore::new);
let agent = cx.update(|cx| NativeAgent::new(store, Templates::new(), fs, cx));
let connection = Rc::new(NativeAgentConnection(agent.clone()));
let acp_thread = cx
.update(|cx| connection.new_session(project, PathList::new(&[Path::new("")]), cx))
.await
.unwrap();
let session_id = acp_thread.read_with(cx, |thread, _| thread.session_id().clone());
let parent = agent.read_with(cx, |agent, _| {
agent.sessions.get(&session_id).unwrap().thread.clone()
});
let model = Arc::new(FakeLanguageModel::default());
parent.update(cx, |thread, cx| thread.set_model(model.clone(), cx));
let environment = NativeThreadEnvironment {
agent: agent.downgrade(),
thread: parent.downgrade(),
acp_thread: acp_thread.downgrade(),
};
let handle = cx
.update(|cx| environment.create_subagent_thread("subagent".to_string(), cx))
.unwrap();
let thread = agent.read_with(cx, |agent, _| {
agent.sessions.get(&handle.id()).unwrap().thread.clone()
});
thread.update(cx, |thread, _| thread.add_tool(EchoTool));
cx.run_until_parked();
Self {
_agent: agent,
_acp_thread: acp_thread,
parent,
thread,
environment,
handle,
model,
}
}
fn configure_compaction(
&self,
enabled: bool,
max_tokens: u64,
max_output_tokens: Option<u64>,
cx: &mut TestAppContext,
) {
cx.update(|cx| {
let mut settings = AgentSettings::get_global(cx).clone();
settings.auto_compact.enabled = enabled;
AgentSettings::override_global(settings, cx);
});
self.model.set_max_token_count(max_tokens);
self.model.set_max_output_tokens(max_output_tokens);
}
fn send(&self, message: &str, cx: &mut TestAppContext) -> Task<Result<String>> {
let send = cx.update(|cx| self.handle.send(message.to_string(), &cx.to_async()));
cx.run_until_parked();
send
}
fn update_usage(&self, input_tokens: u64, cx: &mut TestAppContext) {
self.model
.send_last_completion_stream_event(LanguageModelCompletionEvent::UsageUpdate(
TokenUsage {
input_tokens,
..TokenUsage::default()
},
));
cx.run_until_parked();
}
fn tool_round(&self, input_tokens: u64, cx: &mut TestAppContext) {
self.model
.send_last_completion_stream_text_chunk("partial work");
self.update_usage(input_tokens, cx);
self.assert_running(cx);
self.model
.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse(
LanguageModelToolUse {
id: "echo_1".into(),
name: EchoTool::NAME.into(),
raw_input: r#"{"text":"tool output"}"#.to_string(),
input: language_model::LanguageModelToolUseInput::Json(
json!({"text": "tool output"}),
),
is_input_complete: true,
thought_signature: None,
},
));
self.model.end_last_completion_stream();
cx.run_until_parked();
}
fn compaction_request(&self) -> LanguageModelRequest {
let mut requests = self.model.pending_completions();
assert_eq!(requests.len(), 1);
let request = requests.pop().unwrap();
assert_eq!(
request.intent,
Some(CompletionIntent::ThreadContextSummarization)
);
request
}
fn fail_compaction(
&self,
status: http_client::StatusCode,
retry_after: Option<Duration>,
cx: &mut TestAppContext,
) {
let request = self.compaction_request();
self.model
.send_completion_stream_text_chunk(&request, "failed summary");
self.model.send_completion_stream_error(
&request,
LanguageModelCompletionError::from_http_status(
LanguageModelProviderName::new("test"),
status,
"compaction provider error".to_string(),
retry_after,
),
);
self.model.end_completion_stream(&request);
cx.run_until_parked();
}
async fn assert_no_compaction(&self, cx: &mut TestAppContext) {
let saved = self
.thread
.read_with(cx, |thread, cx| thread.to_db(cx))
.await;
assert_eq!(
saved
.messages
.iter()
.filter(|message| matches!(&***message, Message::Compaction(_)))
.count(),
0
);
}
fn assert_running(&self, cx: &mut TestAppContext) {
self.parent.read_with(cx, |thread, cx| {
assert_eq!(thread.running_subagent_ids(cx), vec![self.handle.id()])
});
self.thread
.read_with(cx, |thread, _| assert!(!thread.is_turn_complete()));
}
fn assert_stopped(&self, cx: &mut TestAppContext) {
self.parent.read_with(cx, |thread, cx| {
assert_eq!(thread.running_subagent_ids(cx), Vec::new())
});
self.thread
.read_with(cx, |thread, _| assert!(thread.is_turn_complete()));
}
}
crates/agent/src/thread.rs
+184 -9
View File
@@ -2835,7 +2835,9 @@ impl Thread {
Some(error_message),
)
})?;
return Err(retry_error);
return Err(
retry_error.context("Automatic context compaction failed")
);
}
}
}
@@ -2846,7 +2848,7 @@ impl Thread {
Some(error_message),
)
})?;
return Err(error);
return Err(error.context("Automatic context compaction failed"));
}
}
}
@@ -4329,6 +4331,19 @@ impl Thread {
self.running_turn.is_none()
}
pub(crate) fn auto_compaction_enabled(&self, cx: &App) -> bool {
AgentSettings::get_global(cx).auto_compact.enabled
&& self.input_token_capacity().is_some_and(|max_input_tokens| {
// Models with a small context window don't leave enough headroom for a
// compaction pass; the UI warns the user about the token limit instead.
max_input_tokens >= MIN_COMPACTION_CONTEXT_WINDOW
})
}
pub(crate) fn subagent_partial_output(&self) -> String {
subagent_partial_output_from_messages(&self.messages, self.pending_message.as_ref())
}
fn build_request_messages(
&self,
available_tools: Vec<SharedString>,
@@ -4436,17 +4451,12 @@ impl Thread {
}
fn compaction_message_target_ix(&self, cx: &App) -> Option<usize> {
let auto_compact = AgentSettings::get_global(cx).auto_compact;
if !auto_compact.enabled {
if !self.auto_compaction_enabled(cx) {
return None;
}
let auto_compact = AgentSettings::get_global(cx).auto_compact;
let max_input_tokens = self.input_token_capacity()?;
// Models with a small context window don't leave enough headroom for a
// compaction pass; the UI warns the user about the token limit instead.
if max_input_tokens < MIN_COMPACTION_CONTEXT_WINDOW {
return None;
}
let (usage_ix, usage) = {
let this = &self;
this.messages
@@ -4600,6 +4610,50 @@ impl Thread {
}
}
fn subagent_partial_output_from_messages(
messages: &[Arc<Message>],
pending_message: Option<&AgentMessage>,
) -> String {
let Some(user_message_ix) = messages
.iter()
.rposition(|message| matches!(&**message, Message::User(_)))
else {
return String::new();
};
let mut text_messages = pending_message
.into_iter()
.chain(
messages
.iter()
.skip(user_message_ix + 1)
.rev()
.filter_map(|message| message.as_agent_message()),
)
.filter_map(|message| {
let characters = message
.content
.iter()
.rev()
.filter_map(|content| match content {
AgentMessageContent::Text(text) => Some(text),
_ => None,
})
.flat_map(|text| text.chars().rev())
.take(4096)
.collect::<Vec<_>>();
if characters.is_empty() {
None
} else {
Some(characters.into_iter().rev().collect::<String>())
}
})
.take(3)
.collect::<Vec<_>>();
text_messages.reverse();
text_messages.join("\n\n")
}
fn total_input_tokens(usage: language_model::TokenUsage) -> u64 {
usage
.input_tokens
@@ -7044,6 +7098,127 @@ mod tests {
.collect()
}
#[test]
fn test_subagent_partial_output_filters_non_text() {
let tool_use = LanguageModelToolUse {
id: LanguageModelToolUseId::from("tool"),
name: Arc::from("echo"),
raw_input: "{}".to_string(),
input: language_model::LanguageModelToolUseInput::Json(json!({})),
is_input_complete: true,
thought_signature: None,
};
let tool_result = LanguageModelToolResult {
tool_use_id: tool_use.id.clone(),
tool_name: tool_use.name.clone(),
is_error: false,
content: vec![LanguageModelToolResultContent::Text(Arc::from(
"tool result",
))],
output: Some(json!("raw tool result")),
};
let messages = [
user_text_message(ClientUserMessageId::new(), "old task"),
agent_text_message("old output"),
user_text_message(ClientUserMessageId::new(), "current task"),
Arc::new(Message::Agent(AgentMessage {
content: vec![
AgentMessageContent::Text("first".to_string()),
AgentMessageContent::Thinking {
text: "hidden thinking".to_string(),
signature: None,
},
AgentMessageContent::Text(" part".to_string()),
AgentMessageContent::RedactedThinking("redacted thinking".to_string()),
],
..AgentMessage::default()
})),
Arc::new(Message::Resume),
summary_compaction("compaction summary"),
Arc::new(Message::Agent(AgentMessage {
content: vec![AgentMessageContent::ToolUse(tool_use)],
tool_results: IndexMap::from_iter([(tool_result.tool_use_id.clone(), tool_result)]),
..AgentMessage::default()
})),
agent_text_message(""),
Arc::new(Message::Agent(AgentMessage {
content: vec![AgentMessageContent::Thinking {
text: "more thinking".to_string(),
signature: None,
}],
..AgentMessage::default()
})),
agent_text_message("second part"),
];
let pending_message = AgentMessage {
content: vec![
AgentMessageContent::Text("pending".to_string()),
AgentMessageContent::Thinking {
text: "pending thinking".to_string(),
signature: None,
},
AgentMessageContent::Text(" part".to_string()),
],
..AgentMessage::default()
};
assert_eq!(
subagent_partial_output_from_messages(&messages, Some(&pending_message)),
"first part\n\nsecond part\n\npending part"
);
assert_eq!(subagent_partial_output_from_messages(&[], None), "");
assert_eq!(
subagent_partial_output_from_messages(&[], Some(&pending_message)),
""
);
let mut messages = messages.to_vec();
messages.push(user_text_message(ClientUserMessageId::new(), "next task"));
assert_eq!(subagent_partial_output_from_messages(&messages, None), "");
}
#[test]
fn test_subagent_partial_output_bounds_messages_and_unicode_characters() {
let first = "🦀".repeat(4096);
let second_prefix = "🚀".repeat(2048);
let second_suffix = "🦀".repeat(2048);
let third_prefix = "🌍".repeat(4095);
let messages = [
user_text_message(ClientUserMessageId::new(), "current task"),
agent_text_message("discarded message"),
agent_text_message(&format!("discarded prefix{first}")),
Arc::new(Message::Agent(AgentMessage {
content: vec![
AgentMessageContent::Text(format!("discarded prefix{second_prefix}")),
AgentMessageContent::Thinking {
text: "hidden thinking".to_string(),
signature: None,
},
AgentMessageContent::Text(second_suffix.clone()),
],
..AgentMessage::default()
})),
];
let pending_message = AgentMessage {
content: vec![
AgentMessageContent::Text(format!("discarded prefix{third_prefix}")),
AgentMessageContent::Thinking {
text: "pending thinking".to_string(),
signature: None,
},
AgentMessageContent::Text("🦀".to_string()),
],
..AgentMessage::default()
};
let output = subagent_partial_output_from_messages(&messages, Some(&pending_message));
assert_eq!(
output,
format!("{first}\n\n{second_prefix}{second_suffix}\n\n{third_prefix}🦀")
);
assert_eq!(output.chars().count(), 12_292);
assert_eq!(output.len(), 49_156);
}
#[gpui::test]
async fn test_thread_summary_request_uses_compacted_history(cx: &mut TestAppContext) {
let (thread, _event_stream) = setup_thread_for_test(cx).await;
crates/language_model/src/fake_provider.rs
+8
View File
@@ -272,6 +272,14 @@ impl FakeLanguageModel {
.retain(|(req, _)| req != request);
}
pub fn is_completion_stream_closed(&self, request: &LanguageModelRequest) -> bool {
self.current_completion_txs
.lock()
.iter()
.find(|(pending_request, _)| pending_request == request)
.is_none_or(|(_, sender)| sender.is_closed())
}
pub fn send_last_completion_stream_text_chunk(&self, chunk: impl Into<String>) {
self.send_completion_stream_text_chunk(self.pending_completions().last().unwrap(), chunk);
}