diff --git a/crates/goose-cli/src/commands/web.rs b/crates/goose-cli/src/commands/web.rs index f11e688c1c..f92a4205cf 100644 --- a/crates/goose-cli/src/commands/web.rs +++ b/crates/goose-cli/src/commands/web.rs @@ -567,6 +567,7 @@ async fn process_message_streaming( )) .await; } + // TODO(douwe): delete this MessageContent::ContextLengthExceeded(msg) => { let mut sender = sender.lock().await; let _ = sender @@ -582,7 +583,7 @@ async fn process_message_streaming( .await; let (summarized_messages, _, _) = - agent.summarize_context(messages.messages()).await?; + agent.compact_messages(messages.messages()).await?; SessionManager::replace_conversation( &session_id, &summarized_messages, diff --git a/crates/goose-cli/src/session/mod.rs b/crates/goose-cli/src/session/mod.rs index 0547282c67..3122ba243c 100644 --- a/crates/goose-cli/src/session/mod.rs +++ b/crates/goose-cli/src/session/mod.rs @@ -163,7 +163,7 @@ impl CliSession { agent: &Agent, message_suffix: &str, ) -> Result<()> { - let (summarized_messages, _, _) = agent.summarize_context(messages.messages()).await?; + let (summarized_messages, _, _) = agent.compact_messages(messages.messages()).await?; let msg = format!("Context maxed out\n{}\n{}", "-".repeat(50), message_suffix); output::render_text(&msg, Some(Color::Yellow), true); *messages = summarized_messages; @@ -677,7 +677,7 @@ impl CliSession { // Call the summarize_context method let (summarized_messages, _token_counts, summarization_usage) = self .agent - .summarize_context(self.messages.messages()) + .compact_messages(self.messages.messages()) .await?; // Update the session messages with the summarized ones @@ -947,7 +947,6 @@ impl CliSession { let selected = match context_strategy.as_str() { "clear" => "clear", - "truncate" => "truncate", "summarize" => "summarize", _ => { if interactive { @@ -976,18 +975,6 @@ impl CliSession { output::render_text(&msg, Some(Color::Yellow), true); break; // exit the loop to hand back control to the user } - "truncate" => { - // Truncate messages to fit within context length - let (truncated_messages, _) = self.agent.truncate_context(self.messages.messages()).await?; - let msg = if context_strategy == "truncate" { - format!("Context maxed out - automatically truncated messages.\n{}\nGoose tried its best to truncate messages for you.", "-".repeat(50)) - } else { - format!("Context maxed out\n{}\nGoose tried its best to truncate messages for you.", "-".repeat(50)) - }; - output::render_text("", Some(Color::Yellow), true); - output::render_text(&msg, Some(Color::Yellow), true); - self.messages = truncated_messages; - } "summarize" => { // Use the helper function to summarize context let message_suffix = if context_strategy == "summarize" { @@ -1171,10 +1158,10 @@ impl CliSession { _ => (), } } - Some(Ok(AgentEvent::HistoryReplaced(new_messages))) => { - self.messages = Conversation::new_unvalidated(new_messages.clone()); - } - Some(Ok(AgentEvent::ModelChange { model, mode })) => { + Some(Ok(AgentEvent::HistoryReplaced(updated_conversation))) => { + self.messages = updated_conversation; + } + Some(Ok(AgentEvent::ModelChange { model, mode })) => { // Log model change if in debug mode if self.debug { eprintln!("Model changed to {} in {} mode", model, mode); @@ -1182,6 +1169,7 @@ impl CliSession { } Some(Err(e)) => { + // TODO(Douwe): Delete this // Check if it's a ProviderError::ContextLengthExceeded if e.downcast_ref::() .map(|provider_error| matches!(provider_error, goose::providers::errors::ProviderError::ContextLengthExceeded(_))) diff --git a/crates/goose-server/src/routes/context.rs b/crates/goose-server/src/routes/context.rs index 6f8764b074..943fa38c3b 100644 --- a/crates/goose-server/src/routes/context.rs +++ b/crates/goose-server/src/routes/context.rs @@ -11,8 +11,6 @@ use utoipa::ToSchema; pub struct ContextManageRequest { /// Collection of messages to be managed pub messages: Vec, - /// Operation to perform: "truncation" or "summarize" - pub manage_action: String, /// Optional session ID for session-specific agent pub session_id: String, } @@ -48,28 +46,13 @@ async fn manage_context( ) -> Result, StatusCode> { let agent = state.get_agent_for_route(request.session_id).await?; - let mut processed_messages = Conversation::new_unvalidated(vec![]); - let mut token_counts: Vec = vec![]; - - if request.manage_action == "truncation" { - (processed_messages, token_counts) = agent - .truncate_context(&request.messages) - .await - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; - } else if request.manage_action == "summarize" { - (processed_messages, token_counts, _) = agent - .summarize_context(&request.messages) - .await - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; - } + let (processed_messages, token_counts, _) = agent + .compact_messages(&request.messages) + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; Ok(Json(ContextManageResponse { - messages: processed_messages - .messages() - .iter() - .filter(|m| m.is_user_visible()) - .cloned() - .collect(), + messages: processed_messages.messages().iter().cloned().collect(), token_counts, })) } diff --git a/crates/goose-server/src/routes/reply.rs b/crates/goose-server/src/routes/reply.rs index c8db3b47ce..a419e75769 100644 --- a/crates/goose-server/src/routes/reply.rs +++ b/crates/goose-server/src/routes/reply.rs @@ -121,9 +121,9 @@ impl IntoResponse for SseResponse { } } -#[derive(Debug, Serialize)] +#[derive(Debug, Serialize, utoipa::ToSchema)] #[serde(tag = "type")] -enum MessageEvent { +pub enum MessageEvent { Message { message: Message, }, @@ -141,6 +141,9 @@ enum MessageEvent { request_id: String, message: ServerNotification, }, + UpdateConversation { + conversation: Conversation, + }, Ping, } @@ -306,10 +309,9 @@ pub async fn reply( } } Ok(Some(Ok(AgentEvent::HistoryReplaced(new_messages)))) => { - // Replace the message history with the compacted messages - all_messages = Conversation::new_unvalidated(new_messages); - // Note: We don't send this as a stream event since it's an internal operation - // The client will see the compaction notification message that was sent before this event + all_messages = new_messages.clone(); + stream_event(MessageEvent::UpdateConversation {conversation: new_messages}, &tx, &cancel_token).await; + } Ok(Some(Ok(AgentEvent::ModelChange { model, mode }))) => { stream_event(MessageEvent::ModelChange { model, mode }, &tx, &cancel_token).await; diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index 06da2531fc..3334a67192 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -112,7 +112,7 @@ pub enum AgentEvent { Message(Message), McpNotification((String, ServerNotification)), ModelChange { model: String, mode: String }, - HistoryReplaced(Vec), + HistoryReplaced(Conversation), } impl Default for Agent { @@ -972,7 +972,7 @@ impl Agent { yield AgentEvent::Message( Message::assistant().with_summarization_requested(compaction_message) ); - yield AgentEvent::HistoryReplaced(conversation.messages().clone()); + yield AgentEvent::HistoryReplaced(conversation.clone()); if let Some(session_to_store) = &session { SessionManager::replace_conversation(&session_to_store.id, &conversation).await? } @@ -1306,7 +1306,7 @@ impl Agent { messages_to_add.push(final_message_tool_resp); } } - Err(ProviderError::ContextLengthExceeded(error_msg)) => { + Err(ProviderError::ContextLengthExceeded(_error_msg)) => { info!("Context length exceeded, attempting compaction"); match auto_compact::perform_compaction(self, conversation.messages()).await { @@ -1318,16 +1318,17 @@ impl Agent { "Context limit reached. Conversation has been automatically compacted to continue." ) ); - yield AgentEvent::HistoryReplaced(conversation.messages().to_vec()); + yield AgentEvent::HistoryReplaced(conversation.clone()); if let Some(session_to_store) = &session { SessionManager::replace_conversation(&session_to_store.id, &conversation).await? } continue; } - Err(_) => { - yield AgentEvent::Message(Message::assistant().with_context_length_exceeded( - format!("Context length exceeded and cannot summarize: {}. Unable to continue.", error_msg) - )); + Err(e) => { + error!("Error: {}", e); + yield AgentEvent::Message(Message::assistant().with_text( + format!("Ran into this error trying to compact: {e}.\n\nPlease retry if you think this is a transient or recoverable error.") + )); break; } } diff --git a/crates/goose/src/agents/mod.rs b/crates/goose/src/agents/mod.rs index 39bf903a35..23dabbb8b0 100644 --- a/crates/goose/src/agents/mod.rs +++ b/crates/goose/src/agents/mod.rs @@ -1,5 +1,4 @@ mod agent; -mod context; pub mod extension; pub mod extension_malware_check; pub mod extension_manager; diff --git a/crates/goose/src/agents/subagent_handler.rs b/crates/goose/src/agents/subagent_handler.rs index 25e5d3f532..ad14f282f1 100644 --- a/crates/goose/src/agents/subagent_handler.rs +++ b/crates/goose/src/agents/subagent_handler.rs @@ -127,7 +127,7 @@ fn get_agent_messages( } } - let mut session_messages = + let mut conversation = Conversation::new_unvalidated( vec![Message::user().with_text(text_instruction.clone())], ); @@ -141,15 +141,16 @@ fn get_agent_messages( }; let mut stream = agent - .reply(session_messages.clone(), Some(session_config), None) + .reply(conversation.clone(), Some(session_config), None) .await .map_err(|e| anyhow!("Failed to get reply from agent: {}", e))?; while let Some(message_result) = stream.next().await { match message_result { - Ok(AgentEvent::Message(msg)) => session_messages.push(msg), - Ok(AgentEvent::McpNotification(_)) - | Ok(AgentEvent::ModelChange { .. }) - | Ok(AgentEvent::HistoryReplaced(_)) => {} + Ok(AgentEvent::Message(msg)) => conversation.push(msg), + Ok(AgentEvent::McpNotification(_)) | Ok(AgentEvent::ModelChange { .. }) => {} + Ok(AgentEvent::HistoryReplaced(updated_conversation)) => { + conversation = updated_conversation; + } Err(e) => { tracing::error!("Error receiving message from subagent: {}", e); break; @@ -157,6 +158,6 @@ fn get_agent_messages( } } - Ok(session_messages) + Ok(conversation) }) } diff --git a/crates/goose/src/context_mgmt/auto_compact.rs b/crates/goose/src/context_mgmt/auto_compact.rs index 542d252f31..778b069b45 100644 --- a/crates/goose/src/context_mgmt/auto_compact.rs +++ b/crates/goose/src/context_mgmt/auto_compact.rs @@ -150,7 +150,7 @@ pub async fn perform_compaction(agent: &Agent, messages: &[Message]) -> Result) -> usize { - let model_context_limit = provider.get_model_config().context_limit(); - - // Our conservative estimate of the **target** context limit - // Our token count is an estimate since model providers often don't provide the tokenizer (eg. Claude) - let target_limit = (model_context_limit as f32 * ESTIMATE_FACTOR) as usize; - - // subtract out overhead for system prompt and tools, but ensure we don't go negative - let overhead = SYSTEM_PROMPT_TOKEN_OVERHEAD + TOOLS_TOKEN_OVERHEAD; - if target_limit > overhead { - target_limit - overhead - } else { - // If overhead is larger than target limit, return a minimal usable limit - std::cmp::max(target_limit / 2, 1000) - } -} - -pub fn get_messages_token_counts(token_counter: &TokenCounter, messages: &[Message]) -> Vec { - // Calculate current token count of each message, use count_chat_tokens to ensure we - // capture the full content of the message, include ToolRequests and ToolResponses - messages - .iter() - .map(|msg| token_counter.count_chat_tokens("", std::slice::from_ref(msg), &[])) - .collect() -} - /// Async version of get_messages_token_counts for better performance pub fn get_messages_token_counts_async( token_counter: &AsyncTokenCounter, @@ -49,51 +15,3 @@ pub fn get_messages_token_counts_async( .map(|msg| token_counter.count_chat_tokens("", std::slice::from_ref(msg), &[])) .collect() } - -// These are not being used now but could be useful in the future - -#[allow(dead_code)] -pub struct ChatTokenCounts { - pub system: usize, - pub tools: usize, - pub messages: Vec, -} - -#[allow(dead_code)] -pub fn get_token_counts( - token_counter: &TokenCounter, - messages: &mut [Message], - system_prompt: &str, - tools: &mut Vec, -) -> ChatTokenCounts { - // Take into account the system prompt (includes goosehints), and our tools input - let system_prompt_token_count = token_counter.count_tokens(system_prompt); - let tools_token_count = token_counter.count_tokens_for_tools(tools.as_slice()); - let messages_token_count = get_messages_token_counts(token_counter, messages); - - ChatTokenCounts { - system: system_prompt_token_count, - tools: tools_token_count, - messages: messages_token_count, - } -} - -/// Async version of get_token_counts for better performance -#[allow(dead_code)] -pub fn get_token_counts_async( - token_counter: &AsyncTokenCounter, - messages: &mut [Message], - system_prompt: &str, - tools: &mut Vec, -) -> ChatTokenCounts { - // Take into account the system prompt (includes goosehints), and our tools input - let system_prompt_token_count = token_counter.count_tokens(system_prompt); - let tools_token_count = token_counter.count_tokens_for_tools(tools.as_slice()); - let messages_token_count = get_messages_token_counts_async(token_counter, messages); - - ChatTokenCounts { - system: system_prompt_token_count, - tools: tools_token_count, - messages: messages_token_count, - } -} diff --git a/crates/goose/src/agents/context.rs b/crates/goose/src/context_mgmt/compaction.rs similarity index 60% rename from crates/goose/src/agents/context.rs rename to crates/goose/src/context_mgmt/compaction.rs index 37edcfb565..491fd36b57 100644 --- a/crates/goose/src/agents/context.rs +++ b/crates/goose/src/context_mgmt/compaction.rs @@ -1,68 +1,74 @@ -use anyhow::Ok; - use crate::conversation::message::{Message, MessageMetadata}; use crate::conversation::Conversation; -use crate::token_counter::create_async_token_counter; - -use crate::context_mgmt::summarize::summarize_messages; -use crate::context_mgmt::truncate::{truncate_messages, OldestFirstTruncation}; -use crate::context_mgmt::{estimate_target_context_limit, get_messages_token_counts_async}; +use anyhow::Ok; +use rmcp::model::Role; +use serde::Serialize; +use std::sync::Arc; use super::super::agents::Agent; +use crate::prompt_template::render_global_file; -impl Agent { - /// Public API to truncate oldest messages so that the conversation's token count is within the allowed context limit. - pub async fn truncate_context( - &self, - messages: &[Message], // last message is a user msg that led to assistant message with_context_length_exceeded - ) -> Result<(Conversation, Vec), anyhow::Error> { - let provider = self.provider().await?; - let token_counter = create_async_token_counter() - .await - .map_err(|e| anyhow::anyhow!("Failed to create token counter: {}", e))?; - let target_context_limit = estimate_target_context_limit(provider); - let token_counts = get_messages_token_counts_async(&token_counter, messages); +#[derive(Serialize)] +struct SummarizeContext { + messages: String, +} - let (mut new_messages, mut new_token_counts) = truncate_messages( - messages, - &token_counts, - target_context_limit, - &OldestFirstTruncation, - )?; +use crate::providers::base::{Provider, ProviderUsage}; - // Only add an assistant message if we have room for it and it won't cause another overflow - let assistant_message = Message::assistant().with_text("I had run into a context length exceeded error so I truncated some of the oldest messages in our conversation."); - let assistant_tokens = - token_counter.count_chat_tokens("", std::slice::from_ref(&assistant_message), &[]); - - let current_total: usize = new_token_counts.iter().sum(); - if current_total + assistant_tokens <= target_context_limit { - new_messages.push(assistant_message); - new_token_counts.push(assistant_tokens); - } else { - // If we can't fit the assistant message, at least log what happened - tracing::warn!("Cannot add truncation notice message due to context limits. Current: {}, Assistant: {}, Limit: {}", - current_total, assistant_tokens, target_context_limit); - } - - Ok((new_messages, new_token_counts)) +/// Summarization function that uses the detailed prompt from the markdown template +async fn do_compact_messages( + provider: Arc, + messages: &[Message], +) -> anyhow::Result, anyhow::Error> { + if messages.is_empty() { + return std::prelude::rust_2015::Ok(None); } + // Format all messages as a single string for the summarization prompt + let messages_text = messages + .iter() + .map(|msg| format!("{:?}", msg)) + .collect::>() + .join("\n\n"); + + let context = SummarizeContext { + messages: messages_text, + }; + + // Render the one-shot summarization prompt + let system_prompt = render_global_file("summarize_oneshot.md", &context)?; + + // Create a simple user message requesting summarization + let user_message = Message::user() + .with_text("Please summarize the conversation history provided in the system prompt."); + let summarization_request = vec![user_message]; + + // Send the request to the provider and fetch the response + let (mut response, mut provider_usage) = provider + .complete_fast(&system_prompt, &summarization_request, &[]) + .await?; + + // Set role to user as it will be used in following conversation as user content + response.role = Role::User; + + // Ensure we have token counts, estimating if necessary + provider_usage + .ensure_tokens(&system_prompt, &summarization_request, &response, &[]) + .await + .map_err(|e| anyhow::anyhow!("Failed to ensure usage tokens: {}", e))?; + + std::prelude::rust_2015::Ok(Some((response, provider_usage))) +} + +impl Agent { /// Public API to summarize the conversation so that its token count is within the allowed context limit. /// Returns the summarized messages, token counts, and the ProviderUsage from summarization - pub async fn summarize_context( + pub async fn compact_messages( &self, messages: &[Message], // last message is a user msg that led to assistant message with_context_length_exceeded - ) -> Result< - ( - Conversation, - Vec, - Option, - ), - anyhow::Error, - > { + ) -> Result<(Conversation, Vec, Option), anyhow::Error> { let provider = self.provider().await?; - let summary_result = summarize_messages(provider.clone(), messages).await?; + let summary_result = do_compact_messages(provider.clone(), messages).await?; let (summary_message, summarization_usage) = match summary_result { Some((summary_message, provider_usage)) => (summary_message, Some(provider_usage)), diff --git a/crates/goose/src/context_mgmt/mod.rs b/crates/goose/src/context_mgmt/mod.rs index 00d11d6b87..31bdbd2f3c 100644 --- a/crates/goose/src/context_mgmt/mod.rs +++ b/crates/goose/src/context_mgmt/mod.rs @@ -1,6 +1,5 @@ pub mod auto_compact; mod common; -pub mod summarize; -pub mod truncate; +mod compaction; pub use common::*; diff --git a/crates/goose/src/context_mgmt/summarize.rs b/crates/goose/src/context_mgmt/summarize.rs deleted file mode 100644 index 35e5aef8e2..0000000000 --- a/crates/goose/src/context_mgmt/summarize.rs +++ /dev/null @@ -1,187 +0,0 @@ -use crate::conversation::message::Message; -use crate::prompt_template::render_global_file; -use crate::providers::base::Provider; - -use anyhow::Result; -use rmcp::model::Role; -use serde::Serialize; -use std::sync::Arc; - -#[derive(Serialize)] -struct SummarizeContext { - messages: String, -} - -use crate::providers::base::ProviderUsage; - -/// Summarization function that uses the detailed prompt from the markdown template -pub async fn summarize_messages( - provider: Arc, - messages: &[Message], -) -> Result, anyhow::Error> { - if messages.is_empty() { - return Ok(None); - } - - // Format all messages as a single string for the summarization prompt - let messages_text = messages - .iter() - .map(|msg| format!("{:?}", msg)) - .collect::>() - .join("\n\n"); - - let context = SummarizeContext { - messages: messages_text, - }; - - // Render the one-shot summarization prompt - let system_prompt = render_global_file("summarize_oneshot.md", &context)?; - - // Create a simple user message requesting summarization - let user_message = Message::user() - .with_text("Please summarize the conversation history provided in the system prompt."); - let summarization_request = vec![user_message]; - - // Send the request to the provider and fetch the response - let (mut response, mut provider_usage) = provider - .complete_fast(&system_prompt, &summarization_request, &[]) - .await?; - - // Set role to user as it will be used in following conversation as user content - response.role = Role::User; - - // Ensure we have token counts, estimating if necessary - provider_usage - .ensure_tokens(&system_prompt, &summarization_request, &response, &[]) - .await - .map_err(|e| anyhow::anyhow!("Failed to ensure usage tokens: {}", e))?; - - Ok(Some((response, provider_usage))) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::conversation::message::{Message, MessageContent}; - use crate::model::ModelConfig; - use crate::providers::base::{ProviderMetadata, ProviderUsage, Usage}; - use crate::providers::errors::ProviderError; - use chrono::Utc; - use rmcp::model::Role; - use rmcp::model::Tool; - use rmcp::model::{AnnotateAble, RawTextContent}; - use std::sync::Arc; - - #[derive(Clone)] - struct MockProvider { - model_config: ModelConfig, - } - - #[async_trait::async_trait] - impl Provider for MockProvider { - fn metadata() -> ProviderMetadata { - ProviderMetadata::empty() - } - - fn get_model_config(&self) -> ModelConfig { - self.model_config.clone() - } - - async fn complete_with_model( - &self, - _model_config: &ModelConfig, - _system: &str, - _messages: &[Message], - _tools: &[Tool], - ) -> Result<(Message, ProviderUsage), ProviderError> { - Ok(( - Message::new( - Role::Assistant, - Utc::now().timestamp(), - vec![MessageContent::Text( - RawTextContent { - text: "Summarized content".to_string(), - meta: None, - } - .no_annotation(), - )], - ), - ProviderUsage::new( - "mock".to_string(), - Usage { - input_tokens: Some(100), - output_tokens: Some(50), - total_tokens: Some(150), - }, - ), - )) - } - } - - fn create_mock_provider() -> Result> { - let mock_model_config = ModelConfig::new("test-model")?.with_context_limit(Some(200_000)); - - Ok(Arc::new(MockProvider { - model_config: mock_model_config, - })) - } - - fn create_test_messages() -> Vec { - vec![ - set_up_text_message("Message 1", Role::User), - set_up_text_message("Message 2", Role::Assistant), - set_up_text_message("Message 3", Role::User), - ] - } - - fn set_up_text_message(text: &str, role: Role) -> Message { - Message::new(role, 0, vec![MessageContent::text(text.to_string())]) - } - - #[tokio::test] - async fn test_summarize_messages_basic() { - let provider = create_mock_provider().expect("failed to create mock provider"); - let messages = create_test_messages(); - - let result = summarize_messages(Arc::clone(&provider), &messages).await; - - assert!(result.is_ok(), "The function should return Ok."); - let summary_result = result.unwrap(); - - assert!( - summary_result.is_some(), - "The summary should contain a result." - ); - let (summarized_message, provider_usage) = summary_result.unwrap(); - - assert_eq!( - summarized_message.role, - Role::User, - "The summarized message should be from the user." - ); - assert!( - provider_usage.usage.input_tokens.unwrap_or(0) > 0, - "Should have input token count" - ); - assert!( - provider_usage.usage.output_tokens.unwrap_or(0) > 0, - "Should have output token count" - ); - } - - #[tokio::test] - async fn test_summarize_messages_empty_input() { - let provider = create_mock_provider().expect("failed to create mock provider"); - let messages: Vec = Vec::new(); - - let result = summarize_messages(Arc::clone(&provider), &messages).await; - - assert!(result.is_ok(), "The function should return Ok."); - let summary_result = result.unwrap(); - - assert!( - summary_result.is_none(), - "The summary should be None for empty input." - ); - } -} diff --git a/crates/goose/src/context_mgmt/truncate.rs b/crates/goose/src/context_mgmt/truncate.rs deleted file mode 100644 index b6f74c5972..0000000000 --- a/crates/goose/src/context_mgmt/truncate.rs +++ /dev/null @@ -1,749 +0,0 @@ -use crate::conversation::message::{Message, MessageContent}; -use crate::conversation::Conversation; -use crate::utils::safe_truncate; -use anyhow::{anyhow, Result}; -use rmcp::model::{RawContent, ResourceContents, Role}; -use std::collections::HashSet; -use std::ops::DerefMut; -use tracing::{debug, warn}; - -/// Maximum size for truncated content in characters -const MAX_TRUNCATED_CONTENT_SIZE: usize = 5000; - -/// Handles messages that are individually larger than the context limit -/// by truncating their content rather than removing them entirely -fn handle_oversized_messages( - messages: &[Message], - token_counts: &[usize], - context_limit: usize, - strategy: &dyn TruncationStrategy, -) -> Result<(Conversation, Vec), anyhow::Error> { - let mut truncated_messages = Vec::new(); - let mut truncated_token_counts = Vec::new(); - let mut any_truncated = false; - - // Create a basic token counter for re-estimating truncated content - // Note: This is a rough approximation since we don't have access to the actual tokenizer here - let estimate_tokens = |text: &str| -> usize { - // Rough approximation: 1 token per 4 characters for English text - (text.len() / 4).max(1) - }; - - for (i, (message, &original_tokens)) in messages.iter().zip(token_counts.iter()).enumerate() { - if original_tokens > context_limit { - warn!( - "Message {} has {} tokens, exceeding context limit of {}", - i, original_tokens, context_limit - ); - - // Try to truncate the message content - let truncated_message = truncate_message_content(message, MAX_TRUNCATED_CONTENT_SIZE)?; - let estimated_new_tokens = - estimate_message_tokens(&truncated_message, &estimate_tokens); - - if estimated_new_tokens > context_limit { - // Even truncated message is too large, skip it entirely - warn!("Skipping message {} as even truncated version ({} tokens) exceeds context limit", i, estimated_new_tokens); - any_truncated = true; - continue; - } - - truncated_messages.push(truncated_message); - truncated_token_counts.push(estimated_new_tokens); - any_truncated = true; - } else { - truncated_messages.push(message.clone()); - truncated_token_counts.push(original_tokens); - } - } - - if any_truncated { - debug!("Truncated large message content, now attempting normal truncation"); - // After content truncation, try normal truncation if still needed - return truncate_messages( - &truncated_messages, - &truncated_token_counts, - context_limit, - strategy, - ); - } - - Ok(( - Conversation::new_unvalidated(truncated_messages), - truncated_token_counts, - )) -} - -/// Truncates the content within a message while preserving its structure -fn truncate_message_content(message: &Message, max_content_size: usize) -> Result { - let mut new_message = message.clone(); - - for content in &mut new_message.content { - match content { - MessageContent::Text(text_content) => { - if text_content.text.chars().count() > max_content_size { - let truncated = format!( - "{}\n\n[... content truncated from {} to {} characters ...]", - safe_truncate(&text_content.text, max_content_size), - text_content.text.chars().count(), - max_content_size - ); - text_content.text = truncated; - } - } - MessageContent::ToolResponse(tool_response) => { - if let Ok(ref mut result) = tool_response.tool_result { - for content_item in result { - if let RawContent::Text(ref mut text_content) = content_item.deref_mut() { - if text_content.text.chars().count() > max_content_size { - let truncated = format!( - "{}\n\n[... tool response truncated from {} to {} characters ...]", - safe_truncate(&text_content.text, max_content_size), - text_content.text.chars().count(), - max_content_size - ); - text_content.text = truncated; - } - } - // Handle Resource content which might contain large text - else if let RawContent::Resource(ref mut resource_content) = - content_item.deref_mut() - { - if let ResourceContents::TextResourceContents { text, .. } = - &mut resource_content.resource - { - if text.chars().count() > max_content_size { - let truncated = format!( - "{}\n\n[... resource content truncated from {} to {} characters ...]", - safe_truncate(text, max_content_size), - text.chars().count(), - max_content_size - ); - *text = truncated; - } - } - } - } - } - } - // Other content types are typically smaller, but we could extend this if needed - _ => {} - } - } - - Ok(new_message) -} - -/// Estimates token count for a message using a simple heuristic -fn estimate_message_tokens(message: &Message, estimate_fn: &dyn Fn(&str) -> usize) -> usize { - let mut total_tokens = 10; // Base overhead for message structure - - for content in &message.content { - match content { - MessageContent::Text(text_content) => { - total_tokens += estimate_fn(&text_content.text); - } - MessageContent::ToolResponse(tool_response) => { - if let Ok(ref result) = tool_response.tool_result { - for content_item in result { - match &content_item.raw { - RawContent::Text(text_content) => { - total_tokens += estimate_fn(&text_content.text); - } - RawContent::Resource(resource) => { - match &resource.resource { - ResourceContents::TextResourceContents { text, .. } => { - total_tokens += estimate_fn(text); - } - _ => total_tokens += 5, // Small overhead for other resource types - } - } - _ => { - total_tokens += 5; // Small overhead for other content types - } - } - } - } - } - _ => total_tokens += 5, // Small overhead for other content types - } - } - - total_tokens -} - -/// Truncates the messages to fit within the model's context window. -/// Mutates the input messages and token counts in place. -/// Returns an error if it's impossible to truncate the messages within the context limit. -/// - messages: The vector of messages in the conversation. -/// - token_counts: A parallel vector containing the token count for each message. -/// - context_limit: The maximum allowed context length in tokens. -/// - strategy: The truncation strategy to use. Only option is OldestFirstTruncation. -pub fn truncate_messages( - messages: &[Message], - token_counts: &[usize], - context_limit: usize, - strategy: &dyn TruncationStrategy, -) -> Result<(Conversation, Vec), anyhow::Error> { - let mut messages = messages.to_owned(); - let mut token_counts = token_counts.to_owned(); - - if messages.len() != token_counts.len() { - return Err(anyhow!( - "The vector for messages and token_counts must have same length" - )); - } - - // Step 1: Calculate total tokens - let mut total_tokens: usize = token_counts.iter().sum(); - debug!("Total tokens before truncation: {}", total_tokens); - - // Check if any individual message is larger than the context limit - // First, check for any message that's too large - let max_message_tokens = token_counts.iter().max().copied().unwrap_or(0); - if max_message_tokens > context_limit { - // Try to handle large messages by truncating their content - debug!( - "Found oversized message with {} tokens, attempting content truncation", - max_message_tokens - ); - return handle_oversized_messages(&messages, &token_counts, context_limit, strategy); - } - - let min_user_msg_tokens = messages - .iter() - .zip(token_counts.iter()) - .filter(|(msg, _)| msg.role == Role::User && msg.has_only_text_content()) - .map(|(_, &tokens)| tokens) - .min(); - - // If there are no valid user messages, or the smallest one is too big for the context - if min_user_msg_tokens.is_none() || min_user_msg_tokens.unwrap() > context_limit { - return Err(anyhow!( - "Not possible to truncate messages within context limit: no suitable user messages found" - )); - } - - if total_tokens <= context_limit { - return Ok(( - Conversation::new_unvalidated(messages.to_vec()), - token_counts.to_vec(), - )); // No truncation needed - } - - // Step 2: Determine indices to remove based on strategy - let indices_to_remove = - strategy.determine_indices_to_remove(&messages, &token_counts, context_limit)?; - - // Circuit breaker: if we can't remove enough messages, fail gracefully - let tokens_to_remove: usize = indices_to_remove - .iter() - .map(|&i| token_counts.get(i).copied().unwrap_or(0)) - .sum(); - - if total_tokens - tokens_to_remove > context_limit && !indices_to_remove.is_empty() { - debug!( - "Standard truncation insufficient: {} tokens remain after removing {} tokens", - total_tokens - tokens_to_remove, - tokens_to_remove - ); - // Try more aggressive truncation or content truncation - return handle_oversized_messages(&messages, &token_counts, context_limit, strategy); - } - - if indices_to_remove.is_empty() && total_tokens > context_limit { - return Err(anyhow!( - "Cannot truncate any messages: all messages may be essential or too large individually" - )); - } - - // Step 3: Remove the marked messages - // Vectorize the set and sort in reverse order to avoid shifting indices when removing - let mut indices_to_remove = indices_to_remove.iter().cloned().collect::>(); - indices_to_remove.sort_unstable_by(|a, b| b.cmp(a)); - - for &index in &indices_to_remove { - if index < messages.len() { - let _ = messages.remove(index); - let removed_tokens = token_counts.remove(index); - total_tokens -= removed_tokens; - } - } - - // Step 4: Ensure the last message is a user message with TextContent only - while let Some(last_msg) = messages.last() { - if last_msg.role != Role::User || !last_msg.has_only_text_content() { - let _ = messages.pop().ok_or(anyhow!("Failed to pop message"))?; - let removed_tokens = token_counts - .pop() - .ok_or(anyhow!("Failed to pop token count"))?; - total_tokens -= removed_tokens; - } else { - break; - } - } - - // Step 5: Check first msg is a User message with TextContent only - while let Some(first_msg) = messages.first() { - if first_msg.role != Role::User || !first_msg.has_only_text_content() { - let _ = messages.remove(0); - let removed_tokens = token_counts.remove(0); - total_tokens -= removed_tokens; - } else { - break; - } - } - - debug!("Total tokens after truncation: {}", total_tokens); - - // Ensure we have at least one message remaining and it's within context limit - if messages.is_empty() { - return Err(anyhow!( - "Unable to preserve any messages within context limit" - )); - } - - if total_tokens > context_limit { - return Err(anyhow!( - "Unable to truncate messages within context window." - )); - } - - debug!("Truncation complete. Total tokens: {}", total_tokens); - Ok(( - Conversation::new_unvalidated(messages.to_vec()), - token_counts.to_vec(), - )) -} - -/// Trait representing a truncation strategy -pub trait TruncationStrategy { - /// Determines the indices of messages to remove to fit within the context limit. - /// - /// - `messages`: The list of messages in the conversation. - /// - `token_counts`: A parallel array containing the token count for each message. - /// - `context_limit`: The maximum allowed context length in tokens. - /// - /// Returns a vector of indices to remove. - fn determine_indices_to_remove( - &self, - messages: &[Message], - token_counts: &[usize], - context_limit: usize, - ) -> Result>; -} - -/// Strategy to truncate messages by removing the oldest first -pub struct OldestFirstTruncation; - -impl TruncationStrategy for OldestFirstTruncation { - fn determine_indices_to_remove( - &self, - messages: &[Message], - token_counts: &[usize], - context_limit: usize, - ) -> Result> { - let mut indices_to_remove = HashSet::new(); - let mut total_tokens: usize = token_counts.iter().sum(); - let mut tool_ids_to_remove = HashSet::new(); - - for (i, message) in messages.iter().enumerate() { - if total_tokens <= context_limit { - break; - } - - // Remove the message - indices_to_remove.insert(i); - total_tokens -= token_counts[i]; - debug!( - "OldestFirst: Removing message at index {}. Tokens removed: {}", - i, token_counts[i] - ); - - // If it's a ToolRequest or ToolResponse, mark its pair for removal - if message.is_tool_call() || message.is_tool_response() { - message.get_tool_ids().iter().for_each(|id| { - tool_ids_to_remove.insert((i, id.to_string())); - }); - } - } - - // Now, find and remove paired ToolResponses or ToolRequests - for (i, message) in messages.iter().enumerate() { - let message_tool_ids = message.get_tool_ids(); - // Find the other part of the pair - same tool_id but different message index - for (message_idx, tool_id) in &tool_ids_to_remove { - if message_idx != &i && message_tool_ids.contains(tool_id.as_str()) { - indices_to_remove.insert(i); - // No need to check other tool_ids for this message since it's already marked - break; - } - } - } - - Ok(indices_to_remove) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::conversation::message::Message; - use anyhow::Result; - use rmcp::model::{CallToolRequestParam, Content}; - use rmcp::object; - - // Helper function to create a user text message with a specified token count - fn user_text(index: usize, tokens: usize) -> (Message, usize) { - let content = format!("User message {}", index); - (Message::user().with_text(content), tokens) - } - - // Helper function to create an assistant text message with a specified token count - fn assistant_text(index: usize, tokens: usize) -> (Message, usize) { - let content = format!("Assistant message {}", index); - (Message::assistant().with_text(content), tokens) - } - - // Helper function to create a tool request message with a specified token count - fn assistant_tool_request( - id: &str, - tool_call: CallToolRequestParam, - tokens: usize, - ) -> (Message, usize) { - ( - Message::assistant().with_tool_request(id, Ok(tool_call)), - tokens, - ) - } - - // Helper function to create a tool response message with a specified token count - fn user_tool_response(id: &str, result: Vec, tokens: usize) -> (Message, usize) { - (Message::user().with_tool_response(id, Ok(result)), tokens) - } - - // Helper function to create a large tool response with massive content - fn large_tool_response(id: &str, large_text: String, tokens: usize) -> (Message, usize) { - ( - Message::user().with_tool_response(id, Ok(vec![Content::text(large_text)])), - tokens, - ) - } - - // Helper function to create messages with alternating user and assistant - // text messages of a fixed token count - fn create_messages_with_counts( - num_pairs: usize, - tokens: usize, - remove_last: bool, - ) -> (Conversation, Vec) { - let mut messages = Conversation::new_unvalidated((0..num_pairs).flat_map(|i| { - vec![ - user_text(i * 2, tokens).0, - assistant_text((i * 2) + 1, tokens).0, - ] - })); - - if remove_last { - messages.pop(); - } - - let token_counts = vec![tokens; messages.len()]; - - (messages, token_counts) - } - - #[test] - fn test_handle_oversized_single_message() -> Result<()> { - // Create a scenario similar to the real issue: one very large tool response - let large_content = "A".repeat(50000); // Very large content - let messages = vec![ - user_text(1, 10).0, - assistant_tool_request( - "tool1", - CallToolRequestParam { - name: "read_file".into(), - arguments: Some(object!({"path": "large_file.txt"})), - }, - 20, - ) - .0, - large_tool_response("tool1", large_content, 100000).0, // Massive tool response - user_text(2, 10).0, - ]; - let token_counts = vec![10, 20, 100000, 10]; // One message is huge - let context_limit = 5000; // Much smaller than the large message - - let result = truncate_messages( - &messages, - &token_counts, - context_limit, - &OldestFirstTruncation, - ); - - // Should succeed by truncating the large content - assert!( - result.is_ok(), - "Should handle oversized message by content truncation" - ); - let (truncated_messages, truncated_counts) = result.unwrap(); - - // Should have some messages remaining - assert!( - !truncated_messages.is_empty(), - "Should have some messages left" - ); - - // Total should be within limit - let total_tokens: usize = truncated_counts.iter().sum(); - assert!( - total_tokens <= context_limit, - "Total tokens {} should be <= context limit {}", - total_tokens, - context_limit - ); - - Ok(()) - } - - #[test] - fn test_oldest_first_no_truncation() -> Result<()> { - let (messages, token_counts) = create_messages_with_counts(1, 10, false); - let context_limit = 25; - - let result = truncate_messages( - messages.messages(), - &token_counts, - context_limit, - &OldestFirstTruncation, - )?; - - assert_eq!(result.0.messages(), messages.messages()); - assert_eq!(result.1, token_counts); - Ok(()) - } - - #[test] - fn test_complex_conversation_with_tools() -> Result<()> { - // Simulating a real conversation with multiple tool interactions - let tool_call1 = CallToolRequestParam { - name: "file_read".into(), - arguments: Some(object!({"path": "/tmp/test.txt"})), - }; - let tool_call2 = CallToolRequestParam { - name: "database_query".into(), - arguments: Some(object!({"query": "SELECT * FROM users"})), - }; - - let messages = vec![ - user_text(1, 15).0, // Initial user query - assistant_tool_request("tool1", tool_call1.clone(), 20).0, - user_tool_response( - "tool1", - vec![Content::text("File contents".to_string())], - 10, - ) - .0, - assistant_text(2, 25).0, // Assistant processes file contents - user_text(3, 10).0, // User follow-up - assistant_tool_request("tool2", tool_call2.clone(), 30).0, - user_tool_response( - "tool2", - vec![Content::text("Query results".to_string())], - 20, - ) - .0, - assistant_text(4, 35).0, // Assistant analyzes query results - user_text(5, 5).0, // Final user confirmation - ]; - - let token_counts = vec![15, 20, 10, 25, 10, 30, 20, 35, 5]; - let context_limit = 100; // Force truncation while preserving some tool interactions - - let result = truncate_messages( - &messages, - &token_counts, - context_limit, - &OldestFirstTruncation, - )?; - let (truncated_messages, truncated_counts) = result; - - // Verify that tool pairs are kept together and the conversation remains coherent - assert!(truncated_messages.len() >= 3); // At least one complete interaction should remain - assert!(truncated_messages.last().unwrap().role == Role::User); // Last message should be from user - - // Verify tool pairs are either both present or both removed - let tool_ids: HashSet<_> = truncated_messages - .iter() - .flat_map(|m| m.get_tool_ids()) - .collect(); - - // Each tool ID should appear 0 or 2 times (request + response) - for id in tool_ids { - let count = truncated_messages - .iter() - .flat_map(|m| m.get_tool_ids().into_iter()) - .filter(|&tool_id| tool_id == id) - .count(); - assert!(count == 0 || count == 2, "Tool pair was split: {}", id); - } - - // Total should be within limit - let total_tokens: usize = truncated_counts.iter().sum(); - assert!(total_tokens <= context_limit); - - Ok(()) - } - - #[test] - fn test_edge_case_context_window() -> Result<()> { - // Test case where we're exactly at the context limit - let (messages, token_counts) = create_messages_with_counts(2, 25, false); - let context_limit = 100; // Exactly matches total tokens - - let result = truncate_messages( - messages.messages(), - &token_counts, - context_limit, - &OldestFirstTruncation, - )?; - let (mut messages, mut token_counts) = result; - - assert_eq!(messages.len(), 4); // No truncation needed - assert_eq!(token_counts.iter().sum::(), 100); - - // Now add one more token to force truncation - messages.push(user_text(5, 1).0); - token_counts.push(1); - - let result = truncate_messages( - messages.messages(), - &token_counts, - context_limit, - &OldestFirstTruncation, - )?; - let (messages, token_counts) = result; - - assert!(token_counts.iter().sum::() <= context_limit); - assert!(messages.last().unwrap().role == Role::User); - - Ok(()) - } - - #[test] - fn test_multi_tool_chain() -> Result<()> { - // Simulate a chain of dependent tool calls - let tool_calls = vec![ - CallToolRequestParam { - name: "git_status".into(), - arguments: Some(object!({})), - }, - CallToolRequestParam { - name: "git_diff".into(), - arguments: Some(object!({"file": "main.rs"})), - }, - CallToolRequestParam { - name: "git_commit".into(), - arguments: Some(object!({"message": "Update"})), - }, - ]; - - let mut messages = Vec::new(); - let mut token_counts = Vec::new(); - - // Build a chain of related tool calls - // 30 tokens each round - for (i, tool_call) in tool_calls.into_iter().enumerate() { - let id = format!("git_{}", i); - messages.push(user_text(i, 10).0); - token_counts.push(10); - - messages.push(assistant_tool_request(&id, tool_call, 15).0); - token_counts.push(20); - } - - let context_limit = 50; // Force partial truncation - - let result = truncate_messages( - &messages, - &token_counts, - context_limit, - &OldestFirstTruncation, - )?; - let (truncated_messages, _) = result; - - // Verify that remaining tool chains are complete - let remaining_tool_ids: HashSet<_> = truncated_messages - .iter() - .flat_map(|m| m.get_tool_ids()) - .collect(); - - for _id in remaining_tool_ids { - // Count request/response pairs - let requests = truncated_messages - .iter() - .flat_map(|m| m.get_tool_request_ids().into_iter()) - .count(); - - let responses = truncated_messages - .iter() - .flat_map(|m| m.get_tool_response_ids().into_iter()) - .count(); - - assert_eq!(requests, 1, "Each remaining tool should have one request"); - assert_eq!(responses, 1, "Each remaining tool should have one response"); - } - - Ok(()) - } - - #[test] - fn test_truncation_with_image_content() -> Result<()> { - // Create a conversation with image content mixed in - let messages = vec![ - Message::user().with_image("base64_data", "image/png"), // 50 tokens - Message::assistant().with_text("I see the image"), // 10 tokens - Message::user().with_text("Can you describe it?"), // 10 tokens - Message::assistant().with_text("It shows..."), // 20 tokens - Message::user().with_text("Thanks!"), // 5 tokens - ]; - let token_counts = vec![50, 10, 10, 20, 5]; - let context_limit = 45; // Force truncation - - let result = truncate_messages( - &messages, - &token_counts, - context_limit, - &OldestFirstTruncation, - )?; - let (messages, token_counts) = result; - - // Verify the conversation still makes sense - assert!(!messages.is_empty()); - assert!(messages.last().unwrap().role == Role::User); - assert!(token_counts.iter().sum::() <= context_limit); - - Ok(()) - } - - #[test] - fn test_error_cases() -> Result<()> { - // Test impossibly small context window - let (messages, token_counts) = create_messages_with_counts(1, 10, false); - let result = truncate_messages( - messages.messages(), - &token_counts, - 5, // Impossibly small context - &OldestFirstTruncation, - ); - assert!(result.is_err()); - - // Test unmatched token counts - let messages = vec![user_text(1, 10).0]; - let token_counts = vec![10, 10]; // Mismatched length - let result = truncate_messages(&messages, &token_counts, 100, &OldestFirstTruncation); - assert!(result.is_err()); - - Ok(()) - } -} diff --git a/crates/goose/src/scheduler.rs b/crates/goose/src/scheduler.rs index 42273229e3..a339731498 100644 --- a/crates/goose/src/scheduler.rs +++ b/crates/goose/src/scheduler.rs @@ -1200,7 +1200,7 @@ async fn run_scheduled_job_internal( } if let Some(ref prompt_text) = recipe.prompt { - let mut all_session_messages = + let mut conversation = Conversation::new_unvalidated(vec![Message::user().with_text(prompt_text.clone())]); let session_config = SessionConfig { @@ -1213,11 +1213,7 @@ async fn run_scheduled_job_internal( }; match agent - .reply( - all_session_messages.clone(), - Some(session_config.clone()), - None, - ) + .reply(conversation.clone(), Some(session_config.clone()), None) .await { Ok(mut stream) => { @@ -1231,11 +1227,13 @@ async fn run_scheduled_job_internal( if msg.role == rmcp::model::Role::Assistant { tracing::info!("[Job {}] Assistant: {:?}", job.id, msg.content); } - all_session_messages.push(msg); + conversation.push(msg); } Ok(AgentEvent::McpNotification(_)) => {} Ok(AgentEvent::ModelChange { .. }) => {} - Ok(AgentEvent::HistoryReplaced(_)) => {} + Ok(AgentEvent::HistoryReplaced(updated_conversation)) => { + conversation = updated_conversation; + } Err(e) => { tracing::error!( "[Job {}] Error receiving message from agent: {}", diff --git a/crates/goose/tests/agent.rs b/crates/goose/tests/agent.rs index 157c588817..06b549cfcf 100644 --- a/crates/goose/tests/agent.rs +++ b/crates/goose/tests/agent.rs @@ -120,7 +120,7 @@ async fn run_truncate_test( agent.update_provider(provider).await?; let repeat_count = context_window + 10_000; let large_message_content = "hello ".repeat(repeat_count); - let messages = Conversation::new(vec![ + let mut conversation = Conversation::new(vec![ Message::user().with_text("hi there. what is 2 + 2?"), Message::assistant().with_text("hey! I think it's 4."), Message::user().with_text(&large_message_content), @@ -133,7 +133,7 @@ async fn run_truncate_test( ]) .unwrap(); - let reply_stream = agent.reply(messages, None, None).await?; + let reply_stream = agent.reply(conversation, None, None).await?; tokio::pin!(reply_stream); let mut responses = Vec::new(); @@ -146,8 +146,8 @@ async fn run_truncate_test( Ok(AgentEvent::ModelChange { .. }) => { // Model change events are informational, just continue } - Ok(AgentEvent::HistoryReplaced(_)) => { - // Handle history replacement events if needed + Ok(AgentEvent::HistoryReplaced(updated_conversation)) => { + conversation = updated_conversation; } Err(e) => { println!("Error: {:?}", e); @@ -1108,7 +1108,7 @@ mod max_turns_tests { let provider = Arc::new(MockToolProvider::new()); agent.update_provider(provider).await?; // The mock provider will call a non-existent tool, which will fail and allow the loop to continue - let conversation = Conversation::new(vec![Message::user().with_text("Hello")]).unwrap(); + let mut conversation = Conversation::new(vec![Message::user().with_text("Hello")]).unwrap(); let reply_stream = agent.reply(conversation, None, None).await?; tokio::pin!(reply_stream); @@ -1132,7 +1132,9 @@ mod max_turns_tests { } Ok(AgentEvent::McpNotification(_)) => {} Ok(AgentEvent::ModelChange { .. }) => {} - Ok(AgentEvent::HistoryReplaced(_)) => {} + Ok(AgentEvent::HistoryReplaced(updated_conversation)) => { + conversation = updated_conversation + } Err(e) => { return Err(e); } diff --git a/ui/desktop/openapi.json b/ui/desktop/openapi.json index a53946407b..3db73a4617 100644 --- a/ui/desktop/openapi.json +++ b/ui/desktop/openapi.json @@ -2151,14 +2151,9 @@ "description": "Request payload for context management operations", "required": [ "messages", - "manageAction", "sessionId" ], "properties": { - "manageAction": { - "type": "string", - "description": "Operation to perform: \"truncation\" or \"summarize\"" - }, "messages": { "type": "array", "items": { diff --git a/ui/desktop/src/api/types.gen.ts b/ui/desktop/src/api/types.gen.ts index e4b3ccbeed..3df131bc6a 100644 --- a/ui/desktop/src/api/types.gen.ts +++ b/ui/desktop/src/api/types.gen.ts @@ -80,10 +80,6 @@ export type ContextLengthExceeded = { * Request payload for context management operations */ export type ContextManageRequest = { - /** - * Operation to perform: "truncation" or "summarize" - */ - manageAction: string; /** * Collection of messages to be managed */ diff --git a/ui/desktop/src/hooks/useMessageStream.ts b/ui/desktop/src/hooks/useMessageStream.ts index 1b3781839e..4ca25c69b1 100644 --- a/ui/desktop/src/hooks/useMessageStream.ts +++ b/ui/desktop/src/hooks/useMessageStream.ts @@ -1,7 +1,7 @@ import { useCallback, useEffect, useId, useReducer, useRef, useState } from 'react'; import useSWR from 'swr'; import { createUserMessage, hasCompletedToolCalls } from '../types/message'; -import { Message, Role } from '../api'; +import { Conversation, Message, Role } from '../api'; import { getSession, Session } from '../api'; import { ChatState } from '../types/chatState'; @@ -34,6 +34,7 @@ type MessageEvent = | { type: 'Error'; error: string } | { type: 'Finish'; reason: string } | { type: 'ModelChange'; model: string; mode: string } + | { type: 'UpdateConversation'; conversation: Conversation } | NotificationEvent; export interface UseMessageStreamOptions { @@ -328,6 +329,11 @@ export function useMessageStream({ break; } + case 'UpdateConversation': { + setMessages(parsedEvent.conversation); + break; + } + case 'Error': { // Always throw the error so it gets caught and sets the error state // This ensures the retry UI appears for ALL errors