ai-space/goose
3
0
Fork 0
mirror of https://github.com/aaif-goose/goose.git synced 2026-07-03 14:10:03 +02:00
Code Issues Packages Projects Releases Wiki Activity

First

Browse Source
This commit is contained in:
Douwe Osinga
2025-10-14 15:25:42 -04:00
parent 05882bc948
commit 7748467765
18 changed files with 120 additions and 1161 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-cli/src/commands/web.rs
+2 -1
View File
@@ -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,
crates/goose-cli/src/session/mod.rs
+7 -19
View File
@@ -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::<goose::providers::errors::ProviderError>()
.map(|provider_error| matches!(provider_error, goose::providers::errors::ProviderError::ContextLengthExceeded(_)))
crates/goose-server/src/routes/context.rs
+5 -22
View File
@@ -11,8 +11,6 @@ use utoipa::ToSchema;
pub struct ContextManageRequest {
/// Collection of messages to be managed
pub messages: Vec<Message>,
/// 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<Json<ContextManageResponse>, 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<usize> = 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,
}))
}
crates/goose-server/src/routes/reply.rs
+8 -6
View File
@@ -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;
crates/goose/src/agents/agent.rs
+9 -8
View File
@@ -112,7 +112,7 @@ pub enum AgentEvent {
Message(Message),
McpNotification((String, ServerNotification)),
ModelChange { model: String, mode: String },
HistoryReplaced(Vec<Message>),
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;
}
}
crates/goose/src/agents/mod.rs
-1
View File
@@ -1,5 +1,4 @@
mod agent;
mod context;
pub mod extension;
pub mod extension_malware_check;
pub mod extension_manager;
crates/goose/src/agents/subagent_handler.rs
+8 -7
View File
@@ -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)
})
}
crates/goose/src/context_mgmt/auto_compact.rs
+1 -1
View File
@@ -150,7 +150,7 @@ pub async fn perform_compaction(agent: &Agent, messages: &[Message]) -> Result<A
// Perform the compaction on messages excluding the preserved user message
let (mut compacted_messages, _, summarization_usage) =
agent.summarize_context(messages_to_compact).await?;
agent.compact_messages(messages_to_compact).await?;
// Add back the preserved user message if it exists
if let Some(user_message) = preserved_user_message {
crates/goose/src/context_mgmt/common.rs
+1 -83
View File
@@ -1,43 +1,9 @@
use std::sync::Arc;
use rmcp::model::Tool;
use crate::conversation::message::Message;
use crate::{
providers::base::Provider,
token_counter::{AsyncTokenCounter, TokenCounter},
};
use crate::token_counter::AsyncTokenCounter;
const ESTIMATE_FACTOR: f32 = 0.7;
pub const SYSTEM_PROMPT_TOKEN_OVERHEAD: usize = 3_000;
pub const TOOLS_TOKEN_OVERHEAD: usize = 5_000;
pub fn estimate_target_context_limit(provider: Arc<dyn Provider>) -> 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<usize> {
// 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<usize>,
}
#[allow(dead_code)]
pub fn get_token_counts(
token_counter: &TokenCounter,
messages: &mut [Message],
system_prompt: &str,
tools: &mut Vec<Tool>,
) -> 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<Tool>,
) -> 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,
}
}
crates/goose/src/agents/context.rs → crates/goose/src/context_mgmt/compaction.rs
+57 -51
View File
@@ -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<usize>), 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<dyn Provider>,
messages: &[Message],
) -> anyhow::Result<Option<(Message, ProviderUsage)>, 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::<Vec<_>>()
.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<usize>,
Option<crate::providers::base::ProviderUsage>,
),
anyhow::Error,
> {
) -> Result<(Conversation, Vec<usize>, Option<ProviderUsage>), 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)),
crates/goose/src/context_mgmt/mod.rs
+1 -2
View File
@@ -1,6 +1,5 @@
pub mod auto_compact;
mod common;
pub mod summarize;
pub mod truncate;
mod compaction;
pub use common::*;
crates/goose/src/context_mgmt/summarize.rs
-187
View File
@@ -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<dyn Provider>,
messages: &[Message],
) -> Result<Option<(Message, ProviderUsage)>, 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::<Vec<_>>()
.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<Arc<dyn Provider>> {
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<Message> {
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<Message> = 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."
);
}
}
crates/goose/src/context_mgmt/truncate.rs
-749
View File
@@ -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<usize>), 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<Message> {
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<usize>), 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::<Vec<usize>>();
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<HashSet<usize>>;
}
/// 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<HashSet<usize>> {
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<Content>, 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<usize>) {
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::<usize>(), 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::<usize>() <= 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::<usize>() <= 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(())
}
}
crates/goose/src/scheduler.rs
+6 -8
View File
@@ -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: {}",
crates/goose/tests/agent.rs
+8 -6
View File
@@ -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);
}
ui/desktop/openapi.json
-5
View File
@@ -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": {
ui/desktop/src/api/types.gen.ts
-4
View File
@@ -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
*/
ui/desktop/src/hooks/useMessageStream.ts
+7 -1
View File
@@ -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