tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 7 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> {
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
let messages_for_provider = if config.toolshim {
convert_tool_messages_to_text(messages)
convert_tool_messages_to_text(&filtered_messages)
} else {
Conversation::new_unvalidated(messages.to_vec())
Conversation::new_unvalidated(filtered_messages)
};
// 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>> {
let mut items = Vec::new();
for message in messages.iter().filter(|m| m.is_agent_visible()) {
let filtered = message.agent_visible_content();
let role = match filtered.role {
for message in messages {
let role = match message.role {
Role::User => Some("user"),
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 {
MessageContent::Text(text) => {
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");
// Add conversation history
for message in messages.iter().filter(|m| m.is_agent_visible()) {
let filtered = message.agent_visible_content();
let role_prefix = match filtered.role {
for message in messages {
let role_prefix = match message.role {
Role::User => "Human: ",
Role::Assistant => "Assistant: ",
};
full_prompt.push_str(role_prefix);
for content in &filtered.content {
for content in &message.content {
match content {
MessageContent::Text(text_content) => {
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> {
let mut anthropic_messages = Vec::new();
let filtered_messages: Vec<Message> = messages
.iter()
.filter(|m| m.is_agent_visible())
.map(|m| m.agent_visible_content())
.collect();
for message in &filtered_messages {
for message in messages {
let role = match message.role {
Role::User => USER_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};
pub fn to_bedrock_message(message: &Message) -> Result<bedrock::Message> {
let filtered = message.agent_visible_content();
bedrock::Message::builder()
.role(to_bedrock_role(&filtered.role))
.role(to_bedrock_role(&message.role))
.set_content(Some(
filtered
message
.content
.iter()
.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
fn format_messages(messages: &[Message], image_format: &ImageFormat) -> Vec<DatabricksMessage> {
let mut result = Vec::new();
for message in messages.iter().filter(|m| m.is_agent_visible()) {
let filtered = message.agent_visible_content();
for message in messages {
let mut converted = DatabricksMessage {
content: Value::Null,
role: match filtered.role {
role: match message.role {
Role::User => "user".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_multiple_content = false;
for content in &filtered.content {
for content in &message.content {
match content {
MessageContent::Text(text) => {
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)}));
}
},
MessageContent::ToolResponse(response) => {
match &response.tool_result {
Ok(result) => {
// Send only contents with no audience or with Assistant in the audience
let abridged: Vec<_> = result
.content
.iter()
.filter(|content| {
content.audience().is_none_or(|audience| {
audience.contains(&Role::Assistant)
})
})
.map(|content| content.raw.clone())
.collect();
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());
}
MessageContent::ToolResponse(response) => match &response.tool_result {
Ok(result) => {
let mut tool_content = Vec::new();
for content in result.content.iter().map(|c| c.raw.clone()) {
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 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));
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 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) => {
let mut part = Map::new();
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> {
let mut messages_spec = Vec::new();
for message in messages.iter().filter(|m| m.is_agent_visible()) {
let filtered = message.agent_visible_content();
for message in messages {
let mut converted = json!({
"role": filtered.role
"role": message.role
});
let mut output = Vec::new();
let mut content_array = Vec::new();
let mut text_array = Vec::new();
for content in &filtered.content {
for content in &message.content {
match content {
MessageContent::Text(text) => {
if !text.text.is_empty() {
@@ -136,13 +135,11 @@ pub fn format_messages(messages: &[Message], image_format: &ImageFormat) -> Vec<
MessageContent::ToolResponse(response) => {
match &response.tool_result {
Ok(result) => {
let abridged: Vec<_> = result.content.to_vec();
// Process all content, replacing images with placeholder text
let mut tool_content = Vec::new();
let mut image_messages = Vec::new();
for content in abridged {
for content in result.content.iter() {
match content.deref() {
RawContent::Image(image) => {
// 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);
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]) {
for message in messages.iter().filter(|m| m.is_agent_visible()) {
let filtered = message.agent_visible_content();
for content in &filtered.content {
for message in messages {
for content in &message.content {
if let MessageContent::ToolResponse(response) = content {
match &response.tool_result {
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> {
let mut snowflake_messages = Vec::new();
let filtered_messages: Vec<Message> = messages
.iter()
.filter(|m| m.is_agent_visible())
.map(|m| m.agent_visible_content())
.collect();
for message in &filtered_messages {
for message in messages {
let role = match message.role {
Role::User => "user",
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
.iter()
.map(|message| {
let filtered = message.agent_visible_content();
let mut new_content = Vec::new();
let mut has_tool_content = false;
for content in &filtered.content {
for content in &message.content {
match content {
MessageContent::ToolRequest(req) => {
has_tool_content = true;
@@ -370,9 +369,9 @@ pub fn convert_tool_messages_to_text(messages: &[Message]) -> Conversation {
}
if has_tool_content {
Message::new(filtered.role.clone(), filtered.created, new_content)
Message::new(message.role.clone(), message.created, new_content)
} else {
filtered
message.clone()
}
})
.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
// Note: The current session state reflects the compaction operation,
// as the agent records compaction metrics before retrying
// Note: The session state reflects the RETRY call (after compaction),
// which only sees agent-visible messages (summary + continuation + user message)
let final_input = updated_session.input_tokens.unwrap();
let final_output = updated_session.output_tokens;
let final_total = updated_session.total_tokens.unwrap();
// After compaction during recovery, the session shows the compaction tokens
// Input: system (6000) + long_tool_call messages (~15,400) + new message (100) = ~21,500
// Output: 200 (compaction summary)
// Total: ~21,700
// After compaction, the retry only sees agent-visible messages:
// Input: system (6000) + summary (~100) + continuation (~100) + user message (~100) = ~6300
// Output: 200 (mock detects "summarized" in continuation as compaction)
// Total: ~6500
assert!(
(21000..=22000).contains(&final_input),
"Final input should reflect compaction input (~21,500). Got: {}",
(6000..=6600).contains(&final_input),
"Final input should reflect retry with agent-visible messages (~6300). Got: {}",
final_input
);
assert_eq!(
final_output,
Some(200),
"Final output should be compaction output (200). Got: {:?}",
"Final output should be 200 (mock detects continuation as compaction). Got: {:?}",
final_output
);
@@ -715,13 +715,13 @@ async fn test_context_limit_recovery_compaction() -> Result<()> {
// Accumulated tokens should include all operations:
// - Initial: 1000
// - Compaction: ~21,600 input (system + long messages) + 200 output = 21,800
// - Reply: ~6,300 input + 100 output = 6,400
// Total: 1000 + 21,800 + 6,400 = 29,200
// - Compaction: ~6400 input (mock uses system_prompt.len()/4) + 200 output = ~6600
// - Reply: ~6500 input + 200 output = ~6700
// Total: 1000 + 6600 + 6700 = ~14300
let accumulated = updated_session.accumulated_total_tokens.unwrap();
assert!(
(28000..=30000).contains(&accumulated),
"Accumulated should be ~29,200 (initial + compaction + reply). Got: {}",
(13000..=16000).contains(&accumulated),
"Accumulated should be ~14300 (initial + compaction + reply). Got: {}",
accumulated
);