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

Manual compaction counting fix + cli cleanup (#5480)

Browse Source
This commit is contained in:
David Katz
2025-11-03 16:27:12 -05:00
committed by GitHub
parent 723cbce7ee
commit 9466d92f04
6 changed files with 46 additions and 127 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/agents/agent.rs
+7 -5
View File
@@ -65,7 +65,7 @@ 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";
pub const MANUAL_COMPACT_TRIGGER: &str = "Please compact this conversation";
/// Context needed for the reply function
pub struct ReplyContext {
@@ -803,9 +803,10 @@ impl Agent {
);
match crate::context_mgmt::compact_messages(self, &conversation_to_compact, false).await {
Ok((compacted_conversation, _token_counts, _summarization_usage)) => {
Ok((compacted_conversation, summarization_usage)) => {
if let Some(session_to_store) = &session {
SessionManager::replace_conversation(&session_to_store.id, &compacted_conversation).await?;
Self::update_session_metrics(session_to_store, &summarization_usage, true).await?;
}
yield AgentEvent::HistoryReplaced(compacted_conversation.clone());
@@ -991,7 +992,7 @@ impl Agent {
// Record usage for the session
if let Some(ref session_config) = &session {
if let Some(ref usage) = usage {
Self::update_session_metrics(session_config, usage).await?;
Self::update_session_metrics(session_config, usage, false).await?;
}
}
@@ -1166,9 +1167,10 @@ impl Agent {
);
match crate::context_mgmt::compact_messages(self, &conversation, true).await {
Ok((compacted_conversation, _token_counts, _usage)) => {
Ok((compacted_conversation, usage)) => {
if let Some(session_to_store) = &session {
SessionManager::replace_conversation(&session_to_store.id, &compacted_conversation).await?
SessionManager::replace_conversation(&session_to_store.id, &compacted_conversation).await?;
Self::update_session_metrics(session_to_store, &usage, true).await?;
}
conversation = compacted_conversation;
crates/goose/src/agents/mod.rs
+1 -1
View File
@@ -26,7 +26,7 @@ mod tool_route_manager;
mod tool_router_index_manager;
pub mod types;
pub use agent::{Agent, AgentEvent};
pub use agent::{Agent, AgentEvent, MANUAL_COMPACT_TRIGGER};
pub use extension::ExtensionConfig;
pub use extension_manager::ExtensionManager;
pub use prompt_manager::PromptManager;
crates/goose/src/agents/reply_parts.rs
+16 -3
View File
@@ -255,6 +255,7 @@ impl Agent {
pub(crate) async fn update_session_metrics(
session_config: &crate::agents::types::SessionConfig,
usage: &ProviderUsage,
is_compaction_usage: bool,
) -> Result<()> {
let session_id = session_config.id.as_str();
let session = SessionManager::get_session(session_id, false).await?;
@@ -273,11 +274,23 @@ impl Agent {
let accumulated_output =
accumulate(session.accumulated_output_tokens, usage.usage.output_tokens);
let (current_total, current_input, current_output) = if is_compaction_usage {
// After compaction: summary output becomes new input context
let new_input = usage.usage.output_tokens;
(new_input, new_input, None)
} else {
(
usage.usage.total_tokens,
usage.usage.input_tokens,
usage.usage.output_tokens,
)
};
SessionManager::update_session(session_id)
.schedule_id(session_config.schedule_id.clone())
.total_tokens(usage.usage.total_tokens)
.input_tokens(usage.usage.input_tokens)
.output_tokens(usage.usage.output_tokens)
.total_tokens(current_total)
.input_tokens(current_input)
.output_tokens(current_output)
.accumulated_total_tokens(accumulated_total)
.accumulated_input_tokens(accumulated_input)
.accumulated_output_tokens(accumulated_output)
crates/goose/src/context_mgmt/mod.rs
+6 -28
View File
@@ -31,13 +31,12 @@ struct SummarizeContext {
/// # Returns
/// * A tuple containing:
/// - `Conversation`: The compacted messages
/// - `Vec<usize>`: Token counts for each message
/// - `Option<ProviderUsage>`: Provider usage from summarization
/// - `ProviderUsage`: Provider usage from summarization
pub async fn compact_messages(
agent: &Agent,
conversation: &Conversation,
preserve_last_user_message: bool,
) -> Result<(Conversation, Vec<usize>, Option<ProviderUsage>)> {
) -> Result<(Conversation, ProviderUsage)> {
info!("Performing message compaction");
let messages = conversation.messages();
@@ -99,16 +98,8 @@ pub async fn compact_messages(
};
let provider = agent.provider().await?;
let summary = do_compact(provider.clone(), messages_to_compact).await?;
let (summary_message, summarization_usage) = match summary {
Some((summary_message, provider_usage)) => (summary_message, Some(provider_usage)),
None => {
// No summary was generated (empty input)
tracing::warn!("Summarization failed. Returning empty messages.");
return Ok((Conversation::empty(), vec![], None));
}
};
let (summary_message, summarization_usage) =
do_compact(provider.clone(), messages_to_compact).await?;
// Create the final message list with updated visibility metadata:
// 1. Original messages become user_visible but not agent_visible
@@ -116,27 +107,17 @@ pub async fn compact_messages(
// 3. Assistant messages to continue the conversation remain both user_visible and agent_visible
let mut final_messages = Vec::new();
let mut final_token_counts = Vec::new();
// Add all original messages with updated visibility (preserve user_visible, set agent_visible=false)
for msg in messages_to_compact.iter().cloned() {
let updated_metadata = msg.metadata.with_agent_invisible();
let updated_msg = msg.with_metadata(updated_metadata);
final_messages.push(updated_msg);
// Token count doesn't matter for agent_visible=false messages, but we'll use 0
final_token_counts.push(0);
}
// 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)
// since that's what will be in the context going forward
let summary_tokens = summarization_usage
.as_ref()
.and_then(|usage| usage.usage.output_tokens)
.unwrap_or(0) as usize;
final_messages.push(summary_msg);
final_token_counts.push(summary_tokens);
// Add an assistant message to continue the conversation (agent_visible=true, user_visible=false)
let assistant_message = Message::assistant()
@@ -146,9 +127,7 @@ Do not mention that you read a summary or that conversation summarization occurr
Just continue the conversation naturally based on the summarized context"
)
.with_metadata(MessageMetadata::agent_only());
let assistant_message_tokens: usize = 0; // Not counted since it's for agent context only
final_messages.push(assistant_message);
final_token_counts.push(assistant_message_tokens);
// Add back the preserved user message if it exists
if let Some(user_text) = preserved_user_text {
@@ -157,7 +136,6 @@ Just continue the conversation naturally based on the summarized context"
Ok((
Conversation::new_unvalidated(final_messages),
final_token_counts,
summarization_usage,
))
}
@@ -222,7 +200,7 @@ pub async fn check_if_compaction_needed(
async fn do_compact(
provider: Arc<dyn Provider>,
messages: &[Message],
) -> Result<Option<(Message, ProviderUsage)>, anyhow::Error> {
) -> Result<(Message, ProviderUsage), anyhow::Error> {
let agent_visible_messages: Vec<&Message> = messages
.iter()
.filter(|msg| msg.is_agent_visible())
@@ -255,7 +233,7 @@ async fn do_compact(
.await
.map_err(|e| anyhow::anyhow!("Failed to ensure usage tokens: {}", e))?;
Ok(Some((response, provider_usage)))
Ok((response, provider_usage))
}
fn format_message_for_compacting(msg: &Message) -> String {