diff --git a/crates/agent/src/thread.rs b/crates/agent/src/thread.rs index 1134e1f50d0..ba60f2e6292 100644 --- a/crates/agent/src/thread.rs +++ b/crates/agent/src/thread.rs @@ -11,14 +11,16 @@ use acp_thread::{MentionUri, UserMessageId}; use action_log::ActionLog; use agent_settings::UserAgentsMd; use feature_flags::{ - CreateThreadToolFeatureFlag, FeatureFlagAppExt as _, LspToolFeatureFlag, RenameToolFeatureFlag, - UpdatePlanToolFeatureFlag, UpdateTitleToolFeatureFlag, + CreateThreadToolFeatureFlag, FeatureFlagAppExt as _, HandoffFeatureFlag, LspToolFeatureFlag, + RenameToolFeatureFlag, UpdatePlanToolFeatureFlag, UpdateTitleToolFeatureFlag, }; +use zed_env_vars::{EnvVar, env_var}; use crate::sandboxing::{SandboxRequest, ThreadSandboxGrants, sandboxing_enabled}; use agent_client_protocol::schema as acp; use agent_settings::{ - AgentProfileId, AgentSettings, SUMMARIZE_THREAD_DETAILED_PROMPT, SUMMARIZE_THREAD_PROMPT, + AgentProfileId, AgentSettings, COMPACTION_PROMPT, SUMMARIZE_THREAD_DETAILED_PROMPT, + SUMMARIZE_THREAD_PROMPT, }; use anyhow::{Context as _, Result, anyhow}; use chrono::{DateTime, Local, Utc}; @@ -53,8 +55,8 @@ use serde::{Deserialize, Serialize}; use settings::{ LanguageModelSelection, Settings, SettingsStore, ToolPermissionMode, update_settings_file, }; -use std::cell::RefCell; use std::fmt::Write; +use std::{cell::RefCell, ops::ControlFlow}; use std::{ collections::BTreeMap, marker::PhantomData, @@ -71,6 +73,17 @@ const TOOL_CANCELED_MESSAGE: &str = "Tool canceled by user"; pub const MAX_TOOL_NAME_LENGTH: usize = 64; pub const MAX_SUBAGENT_DEPTH: u8 = 1; +const AGENT_COMPACTION_REMAINING_TOKEN_BUDGET: u64 = 40_000; + +/// Auto-compaction is only available for models whose context window is at least +/// this large. For smaller models there isn't enough headroom for a compaction +/// pass to be worthwhile, so we leave the thread uncompacted and let the UI warn +/// the user instead. +pub const MIN_COMPACTION_CONTEXT_WINDOW: u64 = 80_000; + +static AGENT_COMPACTION_REMAINING_TOKEN_BUDGET_ENV_VAR: std::sync::LazyLock = + env_var!("AGENT_COMPACTION_REMAINING_TOKEN_BUDGET"); + // 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; @@ -2223,6 +2236,41 @@ impl Thread { let mut attempt = 0; let mut intent = CompletionIntent::UserPrompt; loop { + if cx.update(|cx| cx.has_flag::()) { + match Self::perform_compaction_if_needed( + this, + event_stream, + cancellation_rx.clone(), + cx, + ) + .await + { + Ok(ControlFlow::Continue(())) => {} + Ok(ControlFlow::Break(())) => return Ok(()), + Err(error) => { + log::error!("Compaction failed: {}", error); + match error.downcast::() { + Ok(error) => { + match Self::retry_completion_error( + this, + event_stream, + &mut cancellation_rx, + error, + attempt, + cx, + ) + .await? + { + ControlFlow::Break(()) => return Ok(()), + ControlFlow::Continue(()) => continue, + } + } + Err(error) => return Err(error), + } + } + } + } + // Re-read the model and refresh tools on each iteration so that // mid-turn changes (e.g. the user switches model, toggles tools, // or changes profile) take effect between tool-call rounds. @@ -2387,20 +2435,18 @@ impl Thread { if let Some(error) = error { attempt += 1; - let retry = this.update(cx, |this, cx| { - let user_store = this.user_store.read(cx); - this.handle_completion_error(error, attempt, user_store.plan()) - })??; - let timer = cx.background_executor().timer(retry.duration); - event_stream.send_retry(retry); - futures::select! { - _ = timer.fuse() => {} - _ = cancellation_rx.changed().fuse() => { - if *cancellation_rx.borrow() { - log::debug!("Turn cancelled during retry delay, exiting"); - return Ok(()); - } - } + match Self::retry_completion_error( + this, + event_stream, + &mut cancellation_rx, + error, + attempt, + cx, + ) + .await? + { + ControlFlow::Break(_) => return Ok(()), + ControlFlow::Continue(_) => {} } this.update(cx, |this, _cx| { if let Some(Message::Agent(message)) = this.last_message() { @@ -2424,6 +2470,126 @@ impl Thread { } } + /// Computes the retry status for a failed completion, notifies listeners, + /// and waits out the backoff delay (or returns early if the turn is + /// cancelled while waiting). Returns an error if the completion is not + /// retryable or retries are exhausted. + async fn retry_completion_error( + this: &WeakEntity, + event_stream: &ThreadEventStream, + cancellation_rx: &mut watch::Receiver, + error: LanguageModelCompletionError, + attempt: u8, + cx: &mut AsyncApp, + ) -> Result> { + let retry = this.update(cx, |this, cx| { + let user_store = this.user_store.read(cx); + this.handle_completion_error(error, attempt, user_store.plan()) + })??; + let timer = cx.background_executor().timer(retry.duration); + event_stream.send_retry(retry); + futures::select! { + _ = timer.fuse() => {} + _ = cancellation_rx.changed().fuse() => { + if *cancellation_rx.borrow() { + log::debug!("Turn cancelled during retry delay, exiting"); + return Ok(ControlFlow::Break(())); + } + } + } + Ok(ControlFlow::Continue(())) + } + + async fn perform_compaction_if_needed( + this: &WeakEntity, + event_stream: &ThreadEventStream, + mut cancellation_rx: watch::Receiver, + cx: &mut AsyncApp, + ) -> Result> { + let Some((model, request, insertion_ix)) = this.update(cx, |this, cx| { + let Some(insertion_ix) = this.compaction_message_target_ix() else { + return None; + }; + let model = this.model.clone()?; + let request = this.build_compaction_request(insertion_ix, &model, cx); + Some((model, request, insertion_ix)) + })? + else { + return Ok(ControlFlow::Continue(())); + }; + + log::debug!("Running compaction"); + let stream = futures::select! { + result = model.stream_completion(request, cx).fuse() => result, + _ = cancellation_rx.changed().fuse() => { + if *cancellation_rx.borrow() { + log::debug!("Compaction cancelled before request started"); + return Ok(ControlFlow::Break(())); + } + return Ok(ControlFlow::Continue(())); + } + }; + let mut stream = stream?; + + let mut summary = String::new(); + loop { + let event = futures::select! { + event = stream.next().fuse() => event, + _ = cancellation_rx.changed().fuse() => { + if *cancellation_rx.borrow() { + log::debug!("Compaction cancelled while summarizing"); + return Ok(ControlFlow::Break(())); + } + continue; + } + }; + + let Some(event) = event else { + break; + }; + + match event? { + LanguageModelCompletionEvent::Text(text) => summary.push_str(&text), + LanguageModelCompletionEvent::Stop(_) + | LanguageModelCompletionEvent::Started + | LanguageModelCompletionEvent::Queued { .. } + | LanguageModelCompletionEvent::UsageUpdate(_) + | LanguageModelCompletionEvent::Thinking { .. } + | LanguageModelCompletionEvent::RedactedThinking { .. } + | LanguageModelCompletionEvent::ReasoningDetails(_) + | LanguageModelCompletionEvent::ToolUse(_) + | LanguageModelCompletionEvent::ToolUseJsonParseError { .. } + | LanguageModelCompletionEvent::StartMessage { .. } => {} + } + } + + if *cancellation_rx.borrow() { + log::debug!("Compaction cancelled after summarizing"); + return Ok(ControlFlow::Break(())); + } + + let summary = summary.trim().to_string(); + if summary.is_empty() { + log::warn!("Compaction produced an empty summary"); + return Err(anyhow::anyhow!("Compaction produced an empty summary")); + } + + log::debug!("Compaction succeeded:\n{summary}"); + + this.update(cx, |this, cx| { + let compaction = Arc::new(Message::Compaction(CompactionInfo::Summary(summary.into()))); + if insertion_ix <= this.messages.len() { + this.messages.insert(insertion_ix, compaction); + } else { + this.messages.push(compaction); + } + event_stream.send_context_compaction(); + cx.notify(); + })?; + + Ok(ControlFlow::Continue(())) + } + fn process_tool_result( this: &WeakEntity, event_stream: &ThreadEventStream, @@ -3353,10 +3519,24 @@ impl Thread { available_tools: Vec, cx: &App, ) -> Vec { - log::trace!( - "Building request messages from {} thread messages", - self.messages.len() - ); + let mut messages = + self.build_request_messages_until(available_tools, self.messages.len(), cx); + + if let Some(message) = self.pending_message.as_ref() { + messages.extend(message.to_request()); + } + + messages + } + + fn build_request_messages_until( + &self, + available_tools: Vec, + end_ix: usize, + cx: &App, + ) -> Vec { + let end_ix = end_ix.min(self.messages.len()); + log::trace!("Building request messages from {} thread messages", end_ix); let user_agents_md = UserAgentsMd::global(cx).and_then(|s| s.content().cloned()); let system_prompt = SystemPromptTemplate { @@ -3376,22 +3556,22 @@ impl Thread { cache: false, reasoning_details: None, }]; - self.extend_request_history(&mut messages); + self.extend_request_history_until(&mut messages, end_ix); if let Some(last_message) = messages.last_mut() { last_message.cache = true; } - if let Some(message) = self.pending_message.as_ref() { - messages.extend(message.to_request()); - } - messages } - fn extend_request_history(&self, messages: &mut Vec) { - let Some(compaction_ix) = self.latest_compaction_message_ix() else { - for message in &self.messages { + fn extend_request_history_until( + &self, + messages: &mut Vec, + end_ix: usize, + ) { + let Some(compaction_ix) = self.latest_compaction_message_ix_before(end_ix) else { + for message in &self.messages[..end_ix] { messages.extend(message.to_request()); } return; @@ -3404,17 +3584,99 @@ impl Thread { messages.extend(self.retained_user_request_messages_before(compaction_ix)); } - for message in &self.messages[compaction_ix..] { + for message in &self.messages[compaction_ix..end_ix] { messages.extend(message.to_request()); } } - fn latest_compaction_message_ix(&self) -> Option { - self.messages + fn latest_compaction_message_ix_before(&self, end_ix: usize) -> Option { + self.messages[..end_ix] .iter() .rposition(|message| matches!(&**message, Message::Compaction(_))) } + fn compaction_message_target_ix(&self) -> Option { + let model = self.model.as_ref()?; + // 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 model.max_token_count() < MIN_COMPACTION_CONTEXT_WINDOW { + return None; + } + let (usage_ix, usage) = { + let this = &self; + this.messages + .iter() + .enumerate() + .rev() + .find_map(|(ix, message)| { + let Message::User(user_message) = &**message else { + return None; + }; + this.request_token_usage + .get(&user_message.id) + .copied() + .map(|usage| (ix, usage)) + }) + }?; + if self + .latest_compaction_message_ix_before(self.messages.len()) + .is_some_and(|compaction_ix| compaction_ix > usage_ix) + { + return None; + } + + let active_tokens = total_input_tokens(usage).saturating_add(usage.output_tokens); + + let remaining_budget = AGENT_COMPACTION_REMAINING_TOKEN_BUDGET_ENV_VAR + .value + .as_ref() + .and_then(|v| v.parse().ok()) + .unwrap_or(AGENT_COMPACTION_REMAINING_TOKEN_BUDGET); + + let compaction_threshold = model.max_token_count().saturating_sub(remaining_budget); + if active_tokens < compaction_threshold { + return None; + } + + let insertion_ix = match self.messages.last() { + Some(message) + if matches!( + &**message, + Message::User(UserMessage { id, .. }) if !self.request_token_usage.contains_key(id) + ) => + { + self.messages.len().saturating_sub(1) + } + _ => self.messages.len(), + }; + Some(insertion_ix) + } + + fn build_compaction_request( + &self, + insertion_ix: usize, + model: &Arc, + cx: &App, + ) -> LanguageModelRequest { + let mut request = LanguageModelRequest { + thread_id: Some(self.id.to_string()), + prompt_id: Some(self.prompt_id.to_string()), + intent: Some(CompletionIntent::ThreadContextSummarization), + temperature: AgentSettings::temperature_for_model(model, cx), + messages: self.build_request_messages_until(Vec::new(), insertion_ix, cx), + ..Default::default() + }; + + request.messages.push(LanguageModelRequestMessage { + role: Role::User, + content: vec![COMPACTION_PROMPT.into()], + cache: false, + reasoning_details: None, + }); + + request + } + fn retained_user_request_messages_before( &self, compaction_ix: usize, @@ -5129,6 +5391,137 @@ mod tests { .collect() } + #[gpui::test] + async fn test_compaction_threshold_uses_latest_reported_usage(cx: &mut TestAppContext) { + let (thread, _event_stream) = setup_thread_for_test(cx).await; + let model = Arc::new(FakeLanguageModel::default()); + let user_message_id = UserMessageId::new(); + + cx.update(|cx| { + thread.update(cx, |thread, cx| { + thread.set_model(model, cx); + thread + .messages + .push(user_text_message(user_message_id.clone(), "near limit")); + thread.request_token_usage.insert( + user_message_id.clone(), + language_model::TokenUsage { + input_tokens: 960_000, + ..Default::default() + }, + ); + + assert_eq!(thread.compaction_message_target_ix(), Some(1)); + }); + }); + } + + #[gpui::test] + async fn test_compaction_unavailable_for_small_context_window(cx: &mut TestAppContext) { + let (thread, _event_stream) = setup_thread_for_test(cx).await; + let model = Arc::new(FakeLanguageModel::default()); + // A context window below the minimum disables auto-compaction. + model.set_max_token_count(MIN_COMPACTION_CONTEXT_WINDOW - 1); + let user_message_id = UserMessageId::new(); + + cx.update(|cx| { + thread.update(cx, |thread, cx| { + thread.set_model(model, cx); + thread + .messages + .push(user_text_message(user_message_id.clone(), "near limit")); + thread.request_token_usage.insert( + user_message_id.clone(), + language_model::TokenUsage { + input_tokens: u64::MAX, + ..Default::default() + }, + ); + + assert_eq!(thread.compaction_message_target_ix(), None); + }); + }); + } + + #[gpui::test] + async fn test_compaction_inserts_before_new_user_and_requests_compacted_window( + cx: &mut TestAppContext, + ) { + let (thread, _event_stream) = setup_thread_for_test(cx).await; + let model = Arc::new(FakeLanguageModel::default()); + let old_user_message_id = UserMessageId::new(); + let new_user_message_id = UserMessageId::new(); + + cx.update(|cx| { + cx.update_flags(true, vec!["handoff".to_string()]); + thread.update(cx, |thread, cx| { + thread.set_model(model.clone(), cx); + thread + .messages + .push(user_text_message(old_user_message_id.clone(), "old user")); + thread.messages.push(agent_text_message("old assistant")); + thread.request_token_usage.insert( + old_user_message_id.clone(), + language_model::TokenUsage { + input_tokens: 960_000, + ..Default::default() + }, + ); + }); + }); + + let _events = cx + .update(|cx| { + thread.update(cx, |thread, cx| { + thread.send(new_user_message_id, vec!["new prompt"], cx) + }) + }) + .unwrap(); + cx.run_until_parked(); + + let compaction_request = model.pending_completions().pop().unwrap(); + assert_eq!( + compaction_request.intent, + Some(CompletionIntent::ThreadContextSummarization) + ); + let compaction_texts = request_texts_after_system(&compaction_request.messages); + assert_eq!(compaction_texts.len(), 3); + assert_eq!(compaction_texts[0], "old user"); + assert_eq!(compaction_texts[1], "old assistant"); + assert_eq!(compaction_texts[2], COMPACTION_PROMPT); + + model.send_completion_stream_text_chunk(&compaction_request, "compacted old context"); + model.end_completion_stream(&compaction_request); + cx.run_until_parked(); + + let final_request = model.pending_completions().pop().unwrap(); + assert_eq!(final_request.intent, Some(CompletionIntent::UserPrompt)); + assert_eq!( + request_texts_after_system(&final_request.messages), + vec![ + "old user".to_string(), + summary_request_text("compacted old context"), + "new prompt".to_string(), + ] + ); + + model.send_completion_stream_text_chunk(&final_request, "answer"); + model.end_completion_stream(&final_request); + cx.run_until_parked(); + + cx.update(|cx| { + thread.read_with(cx, |thread, _cx| { + assert!(matches!(&*thread.messages[0], Message::User(_))); + assert!(matches!(&*thread.messages[1], Message::Agent(_))); + assert!(matches!( + &*thread.messages[2], + Message::Compaction(CompactionInfo::Summary(summary)) if summary.as_ref() == "compacted old context" + )); + assert!(matches!(&*thread.messages[3], Message::User(_))); + }); + }); + } + #[gpui::test] async fn test_replay_emits_context_compaction(cx: &mut TestAppContext) { let (thread, _event_stream) = setup_thread_for_test(cx).await; diff --git a/crates/agent_settings/src/agent_settings.rs b/crates/agent_settings/src/agent_settings.rs index b25081080fb..80a26095046 100644 --- a/crates/agent_settings/src/agent_settings.rs +++ b/crates/agent_settings/src/agent_settings.rs @@ -26,6 +26,7 @@ pub use crate::user_agents_md::{UserAgentsMd, UserAgentsMdState, init as init_us pub const SUMMARIZE_THREAD_PROMPT: &str = include_str!("prompts/summarize_thread_prompt.txt"); pub const SUMMARIZE_THREAD_DETAILED_PROMPT: &str = include_str!("prompts/summarize_thread_detailed_prompt.txt"); +pub const COMPACTION_PROMPT: &str = include_str!("prompts/compaction_prompt.txt"); #[derive(Debug, Clone, Default, PartialEq, Eq)] pub struct PanelLayout { diff --git a/crates/agent_settings/src/prompts/compaction_prompt.txt b/crates/agent_settings/src/prompts/compaction_prompt.txt new file mode 100644 index 00000000000..94a5aa7603e --- /dev/null +++ b/crates/agent_settings/src/prompts/compaction_prompt.txt @@ -0,0 +1,10 @@ +You are compacting this conversation into a handoff for another agent that will resume the work. + +Include: +- Goal: what the user is ultimately trying to achieve +- State: progress so far, current blockers, and decisions made +- Context: constraints, preferences, and critical data/examples/references needed to continue +- Next: the specific steps that remain +- Pitfalls: anything tried that didn't work + +Write it so the next agent can act without re-asking the user. Be concise and well-structured. diff --git a/crates/agent_ui/src/conversation_view.rs b/crates/agent_ui/src/conversation_view.rs index 1caa5089375..3b21b1c0462 100644 --- a/crates/agent_ui/src/conversation_view.rs +++ b/crates/agent_ui/src/conversation_view.rs @@ -24,7 +24,7 @@ use editor::scroll::Autoscroll; use editor::{ Editor, EditorEvent, EditorMode, MultiBuffer, PathKey, SelectionEffects, SizingBehavior, }; -use feature_flags::{AgentSharingFeatureFlag, FeatureFlagAppExt as _}; +use feature_flags::{AgentSharingFeatureFlag, FeatureFlagAppExt as _, HandoffFeatureFlag}; use file_icons::FileIcons; use fs::Fs; use futures::FutureExt as _; diff --git a/crates/agent_ui/src/conversation_view/thread_view.rs b/crates/agent_ui/src/conversation_view/thread_view.rs index c4ae1a2a956..2624836245a 100644 --- a/crates/agent_ui/src/conversation_view/thread_view.rs +++ b/crates/agent_ui/src/conversation_view/thread_view.rs @@ -9650,6 +9650,18 @@ impl ThreadView { } let token_usage = self.thread.read(cx).token_usage()?; + + // When auto-compaction is available (the handoff feature flag is enabled + // and the model's context window is large enough), the thread is + // compacted automatically before it reaches the limit, so there's no + // need to warn the user. Models with a context window that's too small + // can't be auto-compacted, so we fall back to the normal warning. + if cx.has_flag::() + && token_usage.max_tokens >= agent::MIN_COMPACTION_CONTEXT_WINDOW + { + return None; + } + let ratio = token_usage.ratio(); let (severity, icon, title) = match ratio { diff --git a/crates/language_model/src/fake_provider.rs b/crates/language_model/src/fake_provider.rs index dfef78b5fce..d0527d495d7 100644 --- a/crates/language_model/src/fake_provider.rs +++ b/crates/language_model/src/fake_provider.rs @@ -11,7 +11,7 @@ use http_client::Result; use parking_lot::Mutex; use std::sync::{ Arc, - atomic::{AtomicBool, Ordering::SeqCst}, + atomic::{AtomicBool, AtomicU64, Ordering::SeqCst}, }; #[derive(Clone)] @@ -126,6 +126,7 @@ pub struct FakeLanguageModel { supports_thinking: AtomicBool, supports_streaming_tools: AtomicBool, supports_images: AtomicBool, + max_token_count: AtomicU64, } impl Default for FakeLanguageModel { @@ -140,6 +141,7 @@ impl Default for FakeLanguageModel { supports_thinking: AtomicBool::new(false), supports_streaming_tools: AtomicBool::new(false), supports_images: AtomicBool::new(false), + max_token_count: AtomicU64::new(1_000_000), } } } @@ -180,6 +182,10 @@ impl FakeLanguageModel { self.supports_images.store(supports, SeqCst); } + pub fn set_max_token_count(&self, count: u64) { + self.max_token_count.store(count, SeqCst); + } + pub fn pending_completions(&self) -> Vec { self.current_completion_txs .lock() @@ -302,7 +308,7 @@ impl LanguageModel for FakeLanguageModel { } fn max_token_count(&self) -> u64 { - 1000000 + self.max_token_count.load(SeqCst) } fn stream_completion(