tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 5 Packages Projects Releases Wiki Activity

Replace compaction notifications with system notifications (#5218)

Browse Source
This commit is contained in:
David Katz
2025-10-22 16:09:44 -04:00
committed by GitHub
parent 755e9f893d
commit 23412e270f
31 changed files with 350 additions and 1060 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-cli/src/session/export.rs
+2 -2
View File
@@ -369,8 +369,8 @@ pub fn message_to_markdown(message: &Message, export_all_content: bool) -> Strin
md.push_str("**Thinking:**\n");
md.push_str("> *Thinking was redacted*\n\n");
}
MessageContent::ConversationCompacted(summarization) => {
md.push_str(&format!("*{}*\n\n", summarization.msg));
MessageContent::SystemNotification(notification) => {
md.push_str(&format!("*{}*\n\n", notification.msg));
}
_ => {
md.push_str(
crates/goose-cli/src/session/output.rs
+12 -2
View File
@@ -185,8 +185,18 @@ pub fn render_message(message: &Message, debug: bool) {
println!("\n{}", style("Thinking:").dim().italic());
print_markdown("Thinking was redacted", theme);
}
MessageContent::ConversationCompacted(summarization) => {
println!("\n{}", style(&summarization.msg).yellow());
MessageContent::SystemNotification(notification) => {
use goose::conversation::message::SystemNotificationType;
match notification.notification_type {
SystemNotificationType::ThinkingMessage => {
show_thinking();
set_thinking_message(&notification.msg);
}
SystemNotificationType::InlineMessage => {
println!("\n{}", style(&notification.msg).yellow());
}
}
}
_ => {
println!("WARNING: Message content type could not be rendered");
crates/goose-server/src/openapi.rs
+5 -6
View File
@@ -18,8 +18,9 @@ use goose::config::declarative_providers::{
DeclarativeProviderConfig, LoadedProvider, ProviderEngine,
};
use goose::conversation::message::{
ConversationCompacted, FrontendToolRequest, Message, MessageContent, MessageMetadata,
RedactedThinkingContent, ThinkingContent, ToolConfirmationRequest, ToolRequest, ToolResponse,
FrontendToolRequest, Message, MessageContent, MessageMetadata, RedactedThinkingContent,
SystemNotificationContent, SystemNotificationType, ThinkingContent, ToolConfirmationRequest,
ToolRequest, ToolResponse,
};
use utoipa::openapi::schema::{
@@ -350,7 +351,6 @@ derive_utoipa!(Icon as IconSchema);
super::routes::agent::update_router_tool_selector,
super::routes::reply::confirm_permission,
super::routes::reply::reply,
super::routes::context::manage_context,
super::routes::session::list_sessions,
super::routes::session::get_session,
super::routes::session::get_session_insights,
@@ -393,8 +393,6 @@ derive_utoipa!(Icon as IconSchema);
super::routes::config_management::UpdateCustomProviderRequest,
super::routes::reply::PermissionConfirmationRequest,
super::routes::reply::ChatRequest,
super::routes::context::ContextManageRequest,
super::routes::context::ContextManageResponse,
super::routes::session::ImportSessionRequest,
super::routes::session::SessionListResponse,
super::routes::session::UpdateSessionDescriptionRequest,
@@ -420,7 +418,8 @@ derive_utoipa!(Icon as IconSchema);
RedactedThinkingContent,
FrontendToolRequest,
ResourceContentsSchema,
ConversationCompacted,
SystemNotificationType,
SystemNotificationContent,
JsonObjectSchema,
RoleSchema,
ProviderMetadata,
crates/goose-server/src/routes/context.rs
-68
View File
@@ -1,68 +0,0 @@
use crate::state::AppState;
use axum::{extract::State, http::StatusCode, routing::post, Json, Router};
use goose::conversation::{message::Message, Conversation};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use utoipa::ToSchema;
/// Request payload for context management operations
#[derive(Debug, Deserialize, ToSchema)]
#[serde(rename_all = "camelCase")]
pub struct ContextManageRequest {
/// Collection of messages to be managed
pub messages: Vec<Message>,
/// Optional session ID for session-specific agent
pub session_id: String,
}
/// Response from context management operations
#[derive(Debug, Serialize, ToSchema)]
#[serde(rename_all = "camelCase")]
pub struct ContextManageResponse {
/// Processed messages after the operation
pub messages: Vec<Message>,
/// Token counts for each processed message
pub token_counts: Vec<usize>,
}
#[utoipa::path(
post,
path = "/context/manage",
request_body = ContextManageRequest,
responses(
(status = 200, description = "Context managed successfully", body = ContextManageResponse),
(status = 401, description = "Unauthorized - Invalid or missing API key"),
(status = 412, description = "Precondition failed - Agent not available"),
(status = 500, description = "Internal server error")
),
security(
("api_key" = [])
),
tag = "Context Management"
)]
async fn manage_context(
State(state): State<Arc<AppState>>,
Json(request): Json<ContextManageRequest>,
) -> Result<Json<ContextManageResponse>, StatusCode> {
let agent = state.get_agent_for_route(request.session_id).await?;
let conversation = Conversation::new_unvalidated(request.messages);
let (processed_messages, token_counts, _) =
goose::context_mgmt::compact_messages(&agent, &conversation, false)
.await
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
// TODO(Douwe): store into db
Ok(Json(ContextManageResponse {
messages: processed_messages.messages().to_vec(),
token_counts,
}))
}
// Configure routes for this module
pub fn routes(state: Arc<AppState>) -> Router {
Router::new()
.route("/context/manage", post(manage_context))
.with_state(state)
}
crates/goose-server/src/routes/mod.rs
-2
View File
@@ -1,7 +1,6 @@
pub mod agent;
pub mod audio;
pub mod config_management;
pub mod context;
pub mod errors;
pub mod extension;
pub mod recipe;
@@ -23,7 +22,6 @@ pub fn configure(state: Arc<crate::state::AppState>) -> Router {
.merge(reply::routes(state.clone()))
.merge(agent::routes(state.clone()))
.merge(audio::routes(state.clone()))
.merge(context::routes(state.clone()))
.merge(extension::routes(state.clone()))
.merge(config_management::routes(state.clone()))
.merge(recipe::routes(state.clone()))
crates/goose/src/agents/agent.rs
+104 -74
View File
@@ -59,11 +59,13 @@ use super::model_selector::autopilot::AutoPilot;
use super::platform_tools;
use super::tool_execution::{ToolCallResult, CHAT_MODE_TOOL_SKIPPED_RESPONSE, DECLINED_RESPONSE};
use crate::agents::subagent_task_config::TaskConfig;
use crate::conversation::message::{Message, ToolRequest};
use crate::conversation::message::{Message, MessageContent, SystemNotificationType, ToolRequest};
use crate::session::extension_data::{EnabledExtensionsState, ExtensionState};
use crate::session::SessionManager;
const DEFAULT_MAX_TURNS: u32 = 1000;
const COMPACTION_THINKING_TEXT: &str = "goose is compacting the conversation...";
const MANUAL_COMPACT_TRIGGER: &str = "Please compact this conversation";
/// Context needed for the reply function
pub struct ReplyContext {
@@ -745,77 +747,99 @@ impl Agent {
session: Option<SessionConfig>,
cancel_token: Option<CancellationToken>,
) -> Result<BoxStream<'_, Result<AgentEvent>>> {
// Try to get session metadata for more accurate token counts
let session_metadata = if let Some(session_config) = &session {
SessionManager::get_session(&session_config.id, false)
.await
.ok()
} else {
None
};
let check_result = crate::context_mgmt::check_if_compaction_needed(
self,
&unfixed_conversation,
None,
session_metadata.as_ref(),
)
.await;
let (did_compact, compacted_conversation, compaction_error) = match check_result {
// TODO(dkatz): send a notification that we are starting compaction here.
Ok(true) => {
match crate::context_mgmt::compact_messages(self, &unfixed_conversation, false)
.await
{
Ok((conversation, _token_counts, _summarization_usage)) => {
(true, conversation, None)
}
Err(e) => (false, unfixed_conversation.clone(), Some(e)),
let is_manual_compact = unfixed_conversation.messages().last().is_some_and(|msg| {
msg.content.iter().any(|c| {
if let MessageContent::Text(text) = c {
text.text.trim() == MANUAL_COMPACT_TRIGGER
} else {
false
}
})
});
if !is_manual_compact {
let session_metadata = if let Some(session_config) = &session {
SessionManager::get_session(&session_config.id, false)
.await
.ok()
} else {
None
};
let needs_auto_compact = crate::context_mgmt::check_if_compaction_needed(
self,
&unfixed_conversation,
None,
session_metadata.as_ref(),
)
.await?;
if !needs_auto_compact {
return self
.reply_internal(unfixed_conversation, session, cancel_token)
.await;
}
Ok(false) => (false, unfixed_conversation, None),
Err(e) => (false, unfixed_conversation.clone(), Some(e)),
};
}
if did_compact {
// Get threshold from config to include in message
let config = crate::config::Config::global();
let threshold = config
.get_param::<f64>("GOOSE_AUTO_COMPACT_THRESHOLD")
.unwrap_or(DEFAULT_COMPACTION_THRESHOLD);
let threshold_percentage = (threshold * 100.0) as u32;
let conversation_to_compact = unfixed_conversation.clone();
let compaction_msg = format!(
"Exceeded auto-compact threshold of {}%. Context has been summarized and reduced.\n\n",
threshold_percentage
Ok(Box::pin(async_stream::try_stream! {
if !is_manual_compact {
let config = crate::config::Config::global();
let threshold = config
.get_param::<f64>("GOOSE_AUTO_COMPACT_THRESHOLD")
.unwrap_or(DEFAULT_COMPACTION_THRESHOLD);
let threshold_percentage = (threshold * 100.0) as u32;
let inline_msg = format!(
"Exceeded auto-compact threshold of {}%. Performing auto-compaction...",
threshold_percentage
);
yield AgentEvent::Message(
Message::assistant().with_system_notification(
SystemNotificationType::InlineMessage,
inline_msg,
)
);
}
yield AgentEvent::Message(
Message::assistant().with_system_notification(
SystemNotificationType::ThinkingMessage,
COMPACTION_THINKING_TEXT,
)
);
Ok(Box::pin(async_stream::try_stream! {
// TODO(Douwe): send this before we actually compact:
yield AgentEvent::Message(
Message::assistant().with_conversation_compacted(compaction_msg)
);
yield AgentEvent::HistoryReplaced(compacted_conversation.clone());
if let Some(session_to_store) = &session {
SessionManager::replace_conversation(&session_to_store.id, &compacted_conversation).await?
}
match crate::context_mgmt::compact_messages(self, &conversation_to_compact, false).await {
Ok((compacted_conversation, _token_counts, _summarization_usage)) => {
if let Some(session_to_store) = &session {
SessionManager::replace_conversation(&session_to_store.id, &compacted_conversation).await?;
}
let mut reply_stream = self.reply_internal(compacted_conversation, session, cancel_token).await?;
while let Some(event) = reply_stream.next().await {
yield event?;
yield AgentEvent::HistoryReplaced(compacted_conversation.clone());
yield AgentEvent::Message(
Message::assistant().with_system_notification(
SystemNotificationType::InlineMessage,
"Compaction complete",
)
);
if !is_manual_compact {
let mut reply_stream = self.reply_internal(compacted_conversation, session, cancel_token).await?;
while let Some(event) = reply_stream.next().await {
yield event?;
}
}
}
}))
} else if let Some(error) = compaction_error {
Ok(Box::pin(async_stream::try_stream! {
yield AgentEvent::Message(Message::assistant().with_text(
format!("Ran into this error trying to auto-compact: {error}.\n\nPlease try again or create a new session")
));
}))
} else {
self.reply_internal(compacted_conversation, session, cancel_token)
.await
}
Err(e) => {
yield AgentEvent::Message(Message::assistant().with_text(
format!("Ran into this error trying to compact: {e}.\n\nPlease try again or create a new session")
));
}
}
}))
}
/// Main reply method that handles the actual agent processing
@@ -1138,23 +1162,29 @@ impl Agent {
}
}
Err(ProviderError::ContextLengthExceeded(_error_msg)) => {
info!("Context length exceeded, attempting compaction");
yield AgentEvent::Message(
Message::assistant().with_system_notification(
SystemNotificationType::InlineMessage,
"Context limit reached. Compacting to continue conversation...",
)
);
yield AgentEvent::Message(
Message::assistant().with_system_notification(
SystemNotificationType::ThinkingMessage,
COMPACTION_THINKING_TEXT,
)
);
// TODO(dkatz): send a notification that we are starting compaction here.
match crate::context_mgmt::compact_messages(self, &conversation, true).await {
Ok((compacted_conversation, _token_counts, _usage)) => {
if let Some(session_to_store) = &session {
SessionManager::replace_conversation(&session_to_store.id, &compacted_conversation).await?
}
conversation = compacted_conversation;
did_recovery_compact_this_iteration = true;
yield AgentEvent::Message(
Message::assistant().with_conversation_compacted(
"Context limit reached. Conversation has been automatically compacted to continue."
)
);
yield AgentEvent::HistoryReplaced(conversation.clone());
if let Some(session_to_store) = &session {
SessionManager::replace_conversation(&session_to_store.id, &conversation).await?
}
continue;
}
Err(e) => {
crates/goose/src/context_mgmt/mod.rs
+3 -9
View File
@@ -93,14 +93,6 @@ pub async fn compact_messages(
final_token_counts.push(0);
}
// Add the compaction marker (user_visible=true, agent_visible=false)
let compaction_marker = Message::assistant()
.with_conversation_compacted("Conversation compacted and summarized")
.with_metadata(MessageMetadata::user_only());
let compaction_marker_tokens: usize = 0; // Not counted since agent_visible=false
final_messages.push(compaction_marker);
final_token_counts.push(compaction_marker_tokens);
// Add the summary message (agent_visible=true, user_visible=false)
let summary_msg = summary_message.with_metadata(MessageMetadata::agent_only());
// For token counting purposes, we use the output tokens (the actual summary content)
@@ -281,7 +273,9 @@ fn format_message_for_compacting(msg: &Message) -> String {
}
MessageContent::Thinking(thinking) => format!("thinking: {}", thinking.thinking),
MessageContent::RedactedThinking(_) => "redacted_thinking".to_string(),
MessageContent::ConversationCompacted(compact) => format!("compacted: {}", compact.msg),
MessageContent::SystemNotification(notification) => {
format!("system_notification: {}", notification.msg)
}
})
.collect();
crates/goose/src/conversation/message.rs
+31 -12
View File
@@ -112,7 +112,16 @@ pub struct FrontendToolRequest {
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, ToSchema)]
pub struct ConversationCompacted {
#[serde(rename_all = "camelCase")]
pub enum SystemNotificationType {
ThinkingMessage,
InlineMessage,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, ToSchema)]
#[serde(rename_all = "camelCase")]
pub struct SystemNotificationContent {
pub notification_type: SystemNotificationType,
pub msg: String,
}
@@ -128,7 +137,7 @@ pub enum MessageContent {
FrontendToolRequest(FrontendToolRequest),
Thinking(ThinkingContent),
RedactedThinking(RedactedThinkingContent),
ConversationCompacted(ConversationCompacted),
SystemNotification(SystemNotificationContent),
}
impl fmt::Display for MessageContent {
@@ -156,8 +165,8 @@ impl fmt::Display for MessageContent {
},
MessageContent::Thinking(t) => write!(f, "[Thinking: {}]", t.thinking),
MessageContent::RedactedThinking(_r) => write!(f, "[RedactedThinking]"),
MessageContent::ConversationCompacted(r) => {
write!(f, "[SummarizationRequested: {}]", r.msg)
MessageContent::SystemNotification(r) => {
write!(f, "[SystemNotification: {}]", r.msg)
}
}
}
@@ -237,14 +246,19 @@ impl MessageContent {
})
}
pub fn conversation_compacted<S: Into<String>>(msg: S) -> Self {
MessageContent::ConversationCompacted(ConversationCompacted { msg: msg.into() })
pub fn system_notification<S: Into<String>>(
notification_type: SystemNotificationType,
msg: S,
) -> Self {
MessageContent::SystemNotification(SystemNotificationContent {
notification_type,
msg: msg.into(),
})
}
// Add this new method to check for summarization requested content
pub fn as_summarization_requested(&self) -> Option<&ConversationCompacted> {
if let MessageContent::ConversationCompacted(ref summarization_requested) = self {
Some(summarization_requested)
pub fn as_system_notification(&self) -> Option<&SystemNotificationContent> {
if let MessageContent::SystemNotification(ref notification) = self {
Some(notification)
} else {
None
}
@@ -650,8 +664,13 @@ impl Message {
.all(|c| matches!(c, MessageContent::Text(_)))
}
pub fn with_conversation_compacted<S: Into<String>>(self, msg: S) -> Self {
self.with_content(MessageContent::conversation_compacted(msg))
pub fn with_system_notification<S: Into<String>>(
self,
notification_type: SystemNotificationType,
msg: S,
) -> Self {
self.with_content(MessageContent::system_notification(notification_type, msg))
.with_metadata(MessageMetadata::user_only())
}
/// Set the visibility metadata for the message
crates/goose/src/providers/formats/anthropic.rs
+1 -1
View File
@@ -90,7 +90,7 @@ pub fn format_messages(messages: &[Message]) -> Vec<Value> {
MessageContent::ToolConfirmationRequest(_tool_confirmation_request) => {
// Skip tool confirmation requests
}
MessageContent::ConversationCompacted(_) => {
MessageContent::SystemNotification(_) => {
// Skip
}
MessageContent::Thinking(thinking) => {
crates/goose/src/providers/formats/bedrock.rs
+2 -2
View File
@@ -48,8 +48,8 @@ pub fn to_bedrock_message_content(content: &MessageContent) -> Result<bedrock::C
// Redacted thinking blocks are not supported in Bedrock - skip
bedrock::ContentBlock::Text("".to_string())
}
MessageContent::ConversationCompacted(_) => {
bail!("SummarizationRequested should not get passed to the provider")
MessageContent::SystemNotification(_) => {
bail!("SystemNotification should not get passed to the provider")
}
MessageContent::ToolRequest(tool_req) => {
let tool_use_id = tool_req.id.to_string();
crates/goose/src/providers/formats/databricks.rs
+1 -1
View File
@@ -128,7 +128,7 @@ fn format_messages(messages: &[Message], image_format: &ImageFormat) -> Vec<Data
}
}
}
MessageContent::ConversationCompacted(_) => {
MessageContent::SystemNotification(_) => {
continue;
}
MessageContent::ToolResponse(response) => {
crates/goose/src/providers/formats/openai.rs
+1 -1
View File
@@ -96,7 +96,7 @@ pub fn format_messages(messages: &[Message], image_format: &ImageFormat) -> Vec<
// Redacted thinking blocks are not directly used in OpenAI format
continue;
}
MessageContent::ConversationCompacted(_) => {
MessageContent::SystemNotification(_) => {
continue;
}
MessageContent::ToolRequest(request) => match &request.tool_call {
crates/goose/src/providers/formats/snowflake.rs
+1 -1
View File
@@ -53,7 +53,7 @@ pub fn format_messages(messages: &[Message]) -> Vec<Value> {
MessageContent::ToolConfirmationRequest(_) => {
// Skip tool confirmation requests
}
MessageContent::ConversationCompacted(_) => {
MessageContent::SystemNotification(_) => {
// Skip
}
MessageContent::Thinking(_thinking) => {