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

refactor: centralize audience filtering before providers receive messages (#6728)

Browse Source 
Signed-off-by: rabi <ramishra@redhat.com>
This commit is contained in:
Rabi Mishra
2026-01-29 22:04:27 +05:30
committed by GitHub
parent e89b7d8159
commit 0a5581723c
12 changed files with 109 additions and 154 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/agents/reply_parts.rs
+8 -2
View File
@@ -186,11 +186,17 @@ impl Agent {
) -> Result<MessageStream, ProviderError> { ) -> Result<MessageStream, ProviderError> {
let config = provider.get_model_config(); let config = provider.get_model_config();
let filtered_messages: Vec<Message> = messages
.iter()
.filter(|m| m.is_agent_visible())
.map(|m| m.agent_visible_content())
.collect();
// Convert tool messages to text if toolshim is enabled // Convert tool messages to text if toolshim is enabled
let messages_for_provider = if config.toolshim { let messages_for_provider = if config.toolshim {
convert_tool_messages_to_text(messages) convert_tool_messages_to_text(&filtered_messages)
} else { } else {
Conversation::new_unvalidated(messages.to_vec()) Conversation::new_unvalidated(filtered_messages)
}; };
// Clone owned data to move into the async stream // Clone owned data to move into the async stream
crates/goose/src/providers/chatgpt_codex.rs
+3 -4
View File
@@ -78,9 +78,8 @@ static CHATGPT_CODEX_AUTH_STATE: LazyLock<Arc<ChatGptCodexAuthState>> =
fn build_input_items(messages: &[Message]) -> Result<Vec<Value>> { fn build_input_items(messages: &[Message]) -> Result<Vec<Value>> {
let mut items = Vec::new(); let mut items = Vec::new();
for message in messages.iter().filter(|m| m.is_agent_visible()) { for message in messages {
let filtered = message.agent_visible_content(); let role = match message.role {
let role = match filtered.role {
Role::User => Some("user"), Role::User => Some("user"),
Role::Assistant => Some("assistant"), Role::Assistant => Some("assistant"),
}; };
@@ -96,7 +95,7 @@ fn build_input_items(messages: &[Message]) -> Result<Vec<Value>> {
} }
}; };
for content in &filtered.content { for content in &message.content {
match content { match content {
MessageContent::Text(text) => { MessageContent::Text(text) => {
if !text.text.is_empty() { if !text.text.is_empty() {
crates/goose/src/providers/cursor_agent.rs
+3 -4
View File
@@ -64,15 +64,14 @@ impl CursorAgentProvider {
full_prompt.push_str("\n\n"); full_prompt.push_str("\n\n");
// Add conversation history // Add conversation history
for message in messages.iter().filter(|m| m.is_agent_visible()) { for message in messages {
let filtered = message.agent_visible_content(); let role_prefix = match message.role {
let role_prefix = match filtered.role {
Role::User => "Human: ", Role::User => "Human: ",
Role::Assistant => "Assistant: ", Role::Assistant => "Assistant: ",
}; };
full_prompt.push_str(role_prefix); full_prompt.push_str(role_prefix);
for content in &filtered.content { for content in &message.content {
match content { match content {
MessageContent::Text(text_content) => { MessageContent::Text(text_content) => {
full_prompt.push_str(&text_content.text); full_prompt.push_str(&text_content.text);
crates/goose/src/providers/formats/anthropic.rs
+1 -7
View File
@@ -34,13 +34,7 @@ const DATA_FIELD: &str = "data";
pub fn format_messages(messages: &[Message]) -> Vec<Value> { pub fn format_messages(messages: &[Message]) -> Vec<Value> {
let mut anthropic_messages = Vec::new(); let mut anthropic_messages = Vec::new();
let filtered_messages: Vec<Message> = messages for message in messages {
.iter()
.filter(|m| m.is_agent_visible())
.map(|m| m.agent_visible_content())
.collect();
for message in &filtered_messages {
let role = match message.role { let role = match message.role {
Role::User => USER_ROLE, Role::User => USER_ROLE,
Role::Assistant => ASSISTANT_ROLE, Role::Assistant => ASSISTANT_ROLE,
crates/goose/src/providers/formats/bedrock.rs
+2 -4
View File
@@ -18,12 +18,10 @@ use super::super::base::Usage;
use crate::conversation::message::{Message, MessageContent}; use crate::conversation::message::{Message, MessageContent};
pub fn to_bedrock_message(message: &Message) -> Result<bedrock::Message> { pub fn to_bedrock_message(message: &Message) -> Result<bedrock::Message> {
let filtered = message.agent_visible_content();
bedrock::Message::builder() bedrock::Message::builder()
.role(to_bedrock_role(&filtered.role)) .role(to_bedrock_role(&message.role))
.set_content(Some( .set_content(Some(
filtered message
.content .content
.iter() .iter()
.map(to_bedrock_message_content) .map(to_bedrock_message_content)
crates/goose/src/providers/formats/databricks.rs
+3 -4
View File
@@ -106,11 +106,10 @@ fn format_tool_response(
/// even though the message structure is otherwise following openai, the enum switches this /// even though the message structure is otherwise following openai, the enum switches this
fn format_messages(messages: &[Message], image_format: &ImageFormat) -> Vec<DatabricksMessage> { fn format_messages(messages: &[Message], image_format: &ImageFormat) -> Vec<DatabricksMessage> {
let mut result = Vec::new(); let mut result = Vec::new();
for message in messages.iter().filter(|m| m.is_agent_visible()) { for message in messages {
let filtered = message.agent_visible_content();
let mut converted = DatabricksMessage { let mut converted = DatabricksMessage {
content: Value::Null, content: Value::Null,
role: match filtered.role { role: match message.role {
Role::User => "user".to_string(), Role::User => "user".to_string(),
Role::Assistant => "assistant".to_string(), Role::Assistant => "assistant".to_string(),
}, },
@@ -122,7 +121,7 @@ fn format_messages(messages: &[Message], image_format: &ImageFormat) -> Vec<Data
let mut has_tool_calls = false; let mut has_tool_calls = false;
let mut has_multiple_content = false; let mut has_multiple_content = false;
for content in &filtered.content { for content in &message.content {
match content { match content {
MessageContent::Text(text) => { MessageContent::Text(text) => {
if !text.text.is_empty() { if !text.text.is_empty() {
crates/goose/src/providers/formats/google.rs
+64 -93
View File
@@ -105,104 +105,75 @@ pub fn format_messages(messages: &[Message]) -> Vec<Value> {
parts.push(json!({"text":format!("Error: {}", e)})); parts.push(json!({"text":format!("Error: {}", e)}));
} }
}, },
MessageContent::ToolResponse(response) => { MessageContent::ToolResponse(response) => match &response.tool_result {
match &response.tool_result { Ok(result) => {
Ok(result) => { let mut tool_content = Vec::new();
// Send only contents with no audience or with Assistant in the audience for content in result.content.iter().map(|c| c.raw.clone()) {
let abridged: Vec<_> = result match content {
.content RawContent::Image(image) => {
.iter() parts.push(json!({
.filter(|content| { "inline_data": {
content.audience().is_none_or(|audience| { "mime_type": image.mime_type,
audience.contains(&Role::Assistant) "data": image.data,
}) }
}) }));
.map(|content| content.raw.clone()) }
.collect(); _ => {
tool_content.push(content.no_annotation());
let mut tool_content = Vec::new();
for content in abridged {
match content {
RawContent::Image(image) => {
parts.push(json!({
"inline_data": {
"mime_type": image.mime_type,
"data": image.data,
}
}));
}
_ => {
tool_content.push(content.no_annotation());
}
} }
} }
let mut text = tool_content
.iter()
.filter_map(|c| match c.deref() {
RawContent::Text(t) => Some(t.text.clone()),
RawContent::Resource(raw_embedded_resource) => Some(
raw_embedded_resource
.clone()
.no_annotation()
.get_text(),
),
_ => None,
})
.collect::<Vec<_>>()
.join("\n");
if text.is_empty() {
text = "Tool call is done.".to_string();
}
let mut part = Map::new();
let mut function_response = Map::new();
function_response.insert("name".to_string(), json!(response.id));
function_response.insert(
"response".to_string(),
json!({"content": {"text": text}}),
);
part.insert(
"functionResponse".to_string(),
json!(function_response),
);
if include_signature {
if let Some(signature) =
get_thought_signature(&response.metadata)
{
part.insert(
THOUGHT_SIGNATURE_KEY.to_string(),
json!(signature),
);
}
}
parts.push(json!(part));
} }
Err(e) => { let mut text = tool_content
let mut part = Map::new(); .iter()
let mut function_response = Map::new(); .filter_map(|c| match c.deref() {
function_response.insert("name".to_string(), json!(response.id)); RawContent::Text(t) => Some(t.text.clone()),
function_response.insert( RawContent::Resource(raw_embedded_resource) => Some(
"response".to_string(), raw_embedded_resource.clone().no_annotation().get_text(),
json!({"content": {"text": format!("Error: {}", e)}}), ),
); _ => None,
part.insert( })
"functionResponse".to_string(), .collect::<Vec<_>>()
json!(function_response), .join("\n");
);
if include_signature { if text.is_empty() {
if let Some(signature) = text = "Tool call is done.".to_string();
get_thought_signature(&response.metadata)
{
part.insert(
THOUGHT_SIGNATURE_KEY.to_string(),
json!(signature),
);
}
}
parts.push(json!(part));
} }
let mut part = Map::new();
let mut function_response = Map::new();
function_response.insert("name".to_string(), json!(response.id));
function_response
.insert("response".to_string(), json!({"content": {"text": text}}));
part.insert("functionResponse".to_string(), json!(function_response));
if include_signature {
if let Some(signature) = get_thought_signature(&response.metadata) {
part.insert(
THOUGHT_SIGNATURE_KEY.to_string(),
json!(signature),
);
}
}
parts.push(json!(part));
} }
} Err(e) => {
let mut part = Map::new();
let mut function_response = Map::new();
function_response.insert("name".to_string(), json!(response.id));
function_response.insert(
"response".to_string(),
json!({"content": {"text": format!("Error: {}", e)}}),
);
part.insert("functionResponse".to_string(), json!(function_response));
if include_signature {
if let Some(signature) = get_thought_signature(&response.metadata) {
part.insert(
THOUGHT_SIGNATURE_KEY.to_string(),
json!(signature),
);
}
}
parts.push(json!(part));
}
},
MessageContent::Thinking(thinking) => { MessageContent::Thinking(thinking) => {
let mut part = Map::new(); let mut part = Map::new();
part.insert("text".to_string(), json!(thinking.thinking)); part.insert("text".to_string(), json!(thinking.thinking));
crates/goose/src/providers/formats/openai.rs
+5 -8
View File
@@ -59,17 +59,16 @@ struct StreamingChunk {
pub fn format_messages(messages: &[Message], image_format: &ImageFormat) -> Vec<Value> { pub fn format_messages(messages: &[Message], image_format: &ImageFormat) -> Vec<Value> {
let mut messages_spec = Vec::new(); let mut messages_spec = Vec::new();
for message in messages.iter().filter(|m| m.is_agent_visible()) { for message in messages {
let filtered = message.agent_visible_content();
let mut converted = json!({ let mut converted = json!({
"role": filtered.role "role": message.role
}); });
let mut output = Vec::new(); let mut output = Vec::new();
let mut content_array = Vec::new(); let mut content_array = Vec::new();
let mut text_array = Vec::new(); let mut text_array = Vec::new();
for content in &filtered.content { for content in &message.content {
match content { match content {
MessageContent::Text(text) => { MessageContent::Text(text) => {
if !text.text.is_empty() { if !text.text.is_empty() {
@@ -136,13 +135,11 @@ pub fn format_messages(messages: &[Message], image_format: &ImageFormat) -> Vec<
MessageContent::ToolResponse(response) => { MessageContent::ToolResponse(response) => {
match &response.tool_result { match &response.tool_result {
Ok(result) => { Ok(result) => {
let abridged: Vec<_> = result.content.to_vec();
// Process all content, replacing images with placeholder text // Process all content, replacing images with placeholder text
let mut tool_content = Vec::new(); let mut tool_content = Vec::new();
let mut image_messages = Vec::new(); let mut image_messages = Vec::new();
for content in abridged { for content in result.content.iter() {
match content.deref() { match content.deref() {
RawContent::Image(image) => { RawContent::Image(image) => {
// Add placeholder text in the tool response // Add placeholder text in the tool response
@@ -164,7 +161,7 @@ pub fn format_messages(messages: &[Message], image_format: &ImageFormat) -> Vec<
tool_content.push(Content::text(text)); tool_content.push(Content::text(text));
} }
_ => { _ => {
tool_content.push(content); tool_content.push(content.clone());
} }
} }
} }
crates/goose/src/providers/formats/openai_responses.rs
+2 -3
View File
@@ -306,9 +306,8 @@ fn add_function_calls(input_items: &mut Vec<Value>, messages: &[Message]) {
} }
fn add_function_call_outputs(input_items: &mut Vec<Value>, messages: &[Message]) { fn add_function_call_outputs(input_items: &mut Vec<Value>, messages: &[Message]) {
for message in messages.iter().filter(|m| m.is_agent_visible()) { for message in messages {
let filtered = message.agent_visible_content(); for content in &message.content {
for content in &filtered.content {
if let MessageContent::ToolResponse(response) = content { if let MessageContent::ToolResponse(response) = content {
match &response.tool_result { match &response.tool_result {
Ok(contents) => { Ok(contents) => {
crates/goose/src/providers/formats/snowflake.rs
+1 -7
View File
@@ -12,13 +12,7 @@ use std::collections::HashSet;
pub fn format_messages(messages: &[Message]) -> Vec<Value> { pub fn format_messages(messages: &[Message]) -> Vec<Value> {
let mut snowflake_messages = Vec::new(); let mut snowflake_messages = Vec::new();
let filtered_messages: Vec<Message> = messages for message in messages {
.iter()
.filter(|m| m.is_agent_visible())
.map(|m| m.agent_visible_content())
.collect();
for message in &filtered_messages {
let role = match message.role { let role = match message.role {
Role::User => "user", Role::User => "user",
Role::Assistant => "assistant", Role::Assistant => "assistant",
crates/goose/src/providers/toolshim.rs
+3 -4
View File
@@ -321,11 +321,10 @@ pub fn convert_tool_messages_to_text(messages: &[Message]) -> Conversation {
let converted_messages: Vec<Message> = messages let converted_messages: Vec<Message> = messages
.iter() .iter()
.map(|message| { .map(|message| {
let filtered = message.agent_visible_content();
let mut new_content = Vec::new(); let mut new_content = Vec::new();
let mut has_tool_content = false; let mut has_tool_content = false;
for content in &filtered.content { for content in &message.content {
match content { match content {
MessageContent::ToolRequest(req) => { MessageContent::ToolRequest(req) => {
has_tool_content = true; has_tool_content = true;
@@ -370,9 +369,9 @@ pub fn convert_tool_messages_to_text(messages: &[Message]) -> Conversation {
} }
if has_tool_content { if has_tool_content {
Message::new(filtered.role.clone(), filtered.created, new_content) Message::new(message.role.clone(), message.created, new_content)
} else { } else {
filtered message.clone()
} }
}) })
.collect(); .collect();
crates/goose/tests/compaction.rs
+14 -14
View File
@@ -684,26 +684,26 @@ async fn test_context_limit_recovery_compaction() -> Result<()> {
); );
// Check the final token state after recovery // Check the final token state after recovery
// Note: The current session state reflects the compaction operation, // Note: The session state reflects the RETRY call (after compaction),
// as the agent records compaction metrics before retrying // which only sees agent-visible messages (summary + continuation + user message)
let final_input = updated_session.input_tokens.unwrap(); let final_input = updated_session.input_tokens.unwrap();
let final_output = updated_session.output_tokens; let final_output = updated_session.output_tokens;
let final_total = updated_session.total_tokens.unwrap(); let final_total = updated_session.total_tokens.unwrap();
// After compaction during recovery, the session shows the compaction tokens // After compaction, the retry only sees agent-visible messages:
// Input: system (6000) + long_tool_call messages (~15,400) + new message (100) = ~21,500 // Input: system (6000) + summary (~100) + continuation (~100) + user message (~100) = ~6300
// Output: 200 (compaction summary) // Output: 200 (mock detects "summarized" in continuation as compaction)
// Total: ~21,700 // Total: ~6500
assert!( assert!(
(21000..=22000).contains(&final_input), (6000..=6600).contains(&final_input),
"Final input should reflect compaction input (~21,500). Got: {}", "Final input should reflect retry with agent-visible messages (~6300). Got: {}",
final_input final_input
); );
assert_eq!( assert_eq!(
final_output, final_output,
Some(200), Some(200),
"Final output should be compaction output (200). Got: {:?}", "Final output should be 200 (mock detects continuation as compaction). Got: {:?}",
final_output final_output
); );
@@ -715,13 +715,13 @@ async fn test_context_limit_recovery_compaction() -> Result<()> {
// Accumulated tokens should include all operations: // Accumulated tokens should include all operations:
// - Initial: 1000 // - Initial: 1000
// - Compaction: ~21,600 input (system + long messages) + 200 output = 21,800 // - Compaction: ~6400 input (mock uses system_prompt.len()/4) + 200 output = ~6600
// - Reply: ~6,300 input + 100 output = 6,400 // - Reply: ~6500 input + 200 output = ~6700
// Total: 1000 + 21,800 + 6,400 = 29,200 // Total: 1000 + 6600 + 6700 = ~14300
let accumulated = updated_session.accumulated_total_tokens.unwrap(); let accumulated = updated_session.accumulated_total_tokens.unwrap();
assert!( assert!(
(28000..=30000).contains(&accumulated), (13000..=16000).contains(&accumulated),
"Accumulated should be ~29,200 (initial + compaction + reply). Got: {}", "Accumulated should be ~14300 (initial + compaction + reply). Got: {}",
accumulated accumulated
); );