From 1a92111f1a37792f837ae45f56b9a4549c07ed1e Mon Sep 17 00:00:00 2001 From: David Katz Date: Mon, 29 Sep 2025 17:20:57 -0400 Subject: [PATCH] Add filtering for agentVisible: false messages on streaming providers (#4847) --- crates/goose/src/context_mgmt/common.rs | 3 +-- crates/goose/src/conversation/mod.rs | 21 ++++++++++++++++- crates/goose/src/providers/base.rs | 23 +++---------------- crates/goose/src/providers/bedrock.rs | 1 + crates/goose/src/providers/claude_code.rs | 2 +- crates/goose/src/providers/cursor_agent.rs | 2 +- .../goose/src/providers/formats/anthropic.rs | 3 +-- .../goose/src/providers/formats/databricks.rs | 2 +- crates/goose/src/providers/formats/google.rs | 1 + crates/goose/src/providers/formats/openai.rs | 2 +- .../goose/src/providers/formats/snowflake.rs | 2 +- crates/goose/src/providers/gemini_cli.rs | 2 +- 12 files changed, 33 insertions(+), 31 deletions(-) diff --git a/crates/goose/src/context_mgmt/common.rs b/crates/goose/src/context_mgmt/common.rs index 8bb5ff00..106f8a28 100644 --- a/crates/goose/src/context_mgmt/common.rs +++ b/crates/goose/src/context_mgmt/common.rs @@ -43,10 +43,9 @@ pub fn get_messages_token_counts_async( token_counter: &AsyncTokenCounter, 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() + .filter(|m| m.is_agent_visible()) .map(|msg| token_counter.count_chat_tokens("", std::slice::from_ref(msg), &[])) .collect() } diff --git a/crates/goose/src/conversation/mod.rs b/crates/goose/src/conversation/mod.rs index 927bfea7..ba956fd7 100644 --- a/crates/goose/src/conversation/mod.rs +++ b/crates/goose/src/conversation/mod.rs @@ -1,4 +1,4 @@ -use crate::conversation::message::{Message, MessageContent}; +use crate::conversation::message::{Message, MessageContent, MessageMetadata}; use rmcp::model::Role; use serde::{Deserialize, Serialize}; use std::collections::HashSet; @@ -103,6 +103,25 @@ impl Conversation { self.0.clear(); } + pub fn filtered_messages(&self, filter: F) -> Vec + where + F: Fn(&MessageMetadata) -> bool, + { + self.0 + .iter() + .filter(|msg| filter(&msg.metadata)) + .cloned() + .collect() + } + + pub fn agent_visible_messages(&self) -> Vec { + self.filtered_messages(|meta| meta.agent_visible) + } + + pub fn user_visible_messages(&self) -> Vec { + self.filtered_messages(|meta| meta.user_visible) + } + fn validate(self) -> Result { let (_messages, issues) = fix_messages(self.0.clone()); if !issues.is_empty() { diff --git a/crates/goose/src/providers/base.rs b/crates/goose/src/providers/base.rs index 26a1716a..44792309 100644 --- a/crates/goose/src/providers/base.rs +++ b/crates/goose/src/providers/base.rs @@ -330,7 +330,6 @@ pub trait Provider: Send + Sync { ) -> Result<(Message, ProviderUsage), ProviderError>; // Default implementation: use the provider's configured model - // This method filters messages to only include agent_visible ones async fn complete( &self, system: &str, @@ -338,20 +337,11 @@ pub trait Provider: Send + Sync { tools: &[Tool], ) -> Result<(Message, ProviderUsage), ProviderError> { let model_config = self.get_model_config(); - - // Filter messages to only include agent_visible ones - let agent_visible_messages: Vec = messages - .iter() - .filter(|m| m.is_agent_visible()) - .cloned() - .collect(); - - self.complete_with_model(&model_config, system, &agent_visible_messages, tools) + self.complete_with_model(&model_config, system, messages, tools) .await } // Check if a fast model is configured, otherwise fall back to regular model - // This method filters messages to only include agent_visible ones async fn complete_fast( &self, system: &str, @@ -361,15 +351,8 @@ pub trait Provider: Send + Sync { let model_config = self.get_model_config(); let fast_config = model_config.use_fast_model(); - // Filter messages to only include agent_visible ones - let agent_visible_messages: Vec = messages - .iter() - .filter(|m| m.is_agent_visible()) - .cloned() - .collect(); - match self - .complete_with_model(&fast_config, system, &agent_visible_messages, tools) + .complete_with_model(&fast_config, system, messages, tools) .await { Ok(result) => Ok(result), @@ -381,7 +364,7 @@ pub trait Provider: Send + Sync { e, model_config.model_name ); - self.complete_with_model(&model_config, system, &agent_visible_messages, tools) + self.complete_with_model(&model_config, system, messages, tools) .await } else { Err(e) diff --git a/crates/goose/src/providers/bedrock.rs b/crates/goose/src/providers/bedrock.rs index ce048a05..e7246aa9 100644 --- a/crates/goose/src/providers/bedrock.rs +++ b/crates/goose/src/providers/bedrock.rs @@ -124,6 +124,7 @@ impl BedrockProvider { .set_messages(Some( messages .iter() + .filter(|m| m.is_agent_visible()) .map(to_bedrock_message) .collect::>()?, )); diff --git a/crates/goose/src/providers/claude_code.rs b/crates/goose/src/providers/claude_code.rs index 3907d9e6..16c36c23 100644 --- a/crates/goose/src/providers/claude_code.rs +++ b/crates/goose/src/providers/claude_code.rs @@ -129,7 +129,7 @@ impl ClaudeCodeProvider { fn messages_to_claude_format(&self, _system: &str, messages: &[Message]) -> Result { let mut claude_messages = Vec::new(); - for message in messages { + for message in messages.iter().filter(|m| m.is_agent_visible()) { let role = match message.role { Role::User => "user", Role::Assistant => "assistant", diff --git a/crates/goose/src/providers/cursor_agent.rs b/crates/goose/src/providers/cursor_agent.rs index d1c767e5..44cd6232 100644 --- a/crates/goose/src/providers/cursor_agent.rs +++ b/crates/goose/src/providers/cursor_agent.rs @@ -133,7 +133,7 @@ impl CursorAgentProvider { full_prompt.push_str("\n\n"); // Add conversation history - for message in messages { + for message in messages.iter().filter(|m| m.is_agent_visible()) { let role_prefix = match message.role { Role::User => "Human: ", Role::Assistant => "Assistant: ", diff --git a/crates/goose/src/providers/formats/anthropic.rs b/crates/goose/src/providers/formats/anthropic.rs index ac9458ec..e3767d8e 100644 --- a/crates/goose/src/providers/formats/anthropic.rs +++ b/crates/goose/src/providers/formats/anthropic.rs @@ -31,8 +31,7 @@ const DATA_FIELD: &str = "data"; pub fn format_messages(messages: &[Message]) -> Vec { let mut anthropic_messages = Vec::new(); - // Convert messages to Anthropic format - for message in messages { + for message in messages.iter().filter(|m| m.is_agent_visible()) { let role = match message.role { Role::User => USER_ROLE, Role::Assistant => ASSISTANT_ROLE, diff --git a/crates/goose/src/providers/formats/databricks.rs b/crates/goose/src/providers/formats/databricks.rs index b19cb509..be125b4c 100644 --- a/crates/goose/src/providers/formats/databricks.rs +++ b/crates/goose/src/providers/formats/databricks.rs @@ -29,7 +29,7 @@ struct DatabricksMessage { /// even though the message structure is otherwise following openai, the enum switches this fn format_messages(messages: &[Message], image_format: &ImageFormat) -> Vec { let mut result = Vec::new(); - for message in messages { + for message in messages.iter().filter(|m| m.is_agent_visible()) { let mut converted = DatabricksMessage { content: Value::Null, role: match message.role { diff --git a/crates/goose/src/providers/formats/google.rs b/crates/goose/src/providers/formats/google.rs index 5c608148..a59b917d 100644 --- a/crates/goose/src/providers/formats/google.rs +++ b/crates/goose/src/providers/formats/google.rs @@ -17,6 +17,7 @@ use std::ops::Deref; pub fn format_messages(messages: &[Message]) -> Vec { messages .iter() + .filter(|m| m.is_agent_visible()) .filter(|message| { message .content diff --git a/crates/goose/src/providers/formats/openai.rs b/crates/goose/src/providers/formats/openai.rs index d71ba14a..ce9f3eb8 100644 --- a/crates/goose/src/providers/formats/openai.rs +++ b/crates/goose/src/providers/formats/openai.rs @@ -59,7 +59,7 @@ struct StreamingChunk { /// even though the message structure is otherwise following openai, the enum switches this pub fn format_messages(messages: &[Message], image_format: &ImageFormat) -> Vec { let mut messages_spec = Vec::new(); - for message in messages { + for message in messages.iter().filter(|m| m.is_agent_visible()) { let mut converted = json!({ "role": message.role }); diff --git a/crates/goose/src/providers/formats/snowflake.rs b/crates/goose/src/providers/formats/snowflake.rs index b5faabbb..660f2838 100644 --- a/crates/goose/src/providers/formats/snowflake.rs +++ b/crates/goose/src/providers/formats/snowflake.rs @@ -13,7 +13,7 @@ pub fn format_messages(messages: &[Message]) -> Vec { let mut snowflake_messages = Vec::new(); // Convert messages to Snowflake format - for message in messages { + for message in messages.iter().filter(|m| m.is_agent_visible()) { let role = match message.role { Role::User => "user", Role::Assistant => "assistant", diff --git a/crates/goose/src/providers/gemini_cli.rs b/crates/goose/src/providers/gemini_cli.rs index 4b142702..35d59fb8 100644 --- a/crates/goose/src/providers/gemini_cli.rs +++ b/crates/goose/src/providers/gemini_cli.rs @@ -140,7 +140,7 @@ impl GeminiCliProvider { full_prompt.push_str("\n\n"); // Add conversation history - for message in messages { + for message in messages.iter().filter(|m| m.is_agent_visible()) { let role_prefix = match message.role { Role::User => "Human: ", Role::Assistant => "Assistant: ",