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

fix(acp): preserve content audience boundaries (#10493)

Browse Source
This commit is contained in:
Jasper
2026-07-17 10:50:36 +02:00
committed by GitHub
parent ae83c0a212
commit 8168cb8378
18 changed files with 1937 additions and 194 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/acp/provider.rs
+239 -33
View File
@@ -1,12 +1,13 @@
use agent_client_protocol::schema::v1::{
ClientCapabilities, CloseSessionRequest, ContentBlock, ContentChunk, EnvVariable, HttpHeader,
ImageContent, InitializeRequest, InitializeResponse, McpCapabilities, McpServer, McpServerHttp,
McpServerStdio, NewSessionRequest, NewSessionResponse, PromptRequest, PromptResponse,
RequestPermissionOutcome, RequestPermissionRequest, RequestPermissionResponse,
SessionConfigKind, SessionConfigOption, SessionConfigOptionCategory,
SessionConfigSelectOptions, SessionId, SessionModeState, SessionNotification, SessionUpdate,
SetSessionConfigOptionRequest, SetSessionModeRequest, SetSessionModeResponse, StopReason,
TextContent, ToolCallContent, ToolCallStatus, ToolKind,
Annotations as AcpAnnotations, ClientCapabilities, CloseSessionRequest, ContentBlock,
ContentChunk, EnvVariable, HttpHeader, ImageContent, InitializeRequest, InitializeResponse,
McpCapabilities, McpServer, McpServerHttp, McpServerStdio, NewSessionRequest,
NewSessionResponse, PromptRequest, PromptResponse, RequestPermissionOutcome,
RequestPermissionRequest, RequestPermissionResponse, Role as AcpRole, SessionConfigKind,
SessionConfigOption, SessionConfigOptionCategory, SessionConfigSelectOptions, SessionId,
SessionModeState, SessionNotification, SessionUpdate, SetSessionConfigOptionRequest,
SetSessionModeRequest, SetSessionModeResponse, StopReason, TextContent, ToolCallContent,
ToolCallStatus, ToolKind,
};
use agent_client_protocol::schema::ProtocolVersion;
use agent_client_protocol::{Agent, Client, ConnectionTo};
@@ -35,10 +36,12 @@ use crate::acp::{map_permission_response, PermissionDecision};
use crate::config::{ExtensionConfig, GooseMode};
use crate::context_mgmt::format_message_for_compacting;
use crate::conversation::message::{Message, MessageContent, TOOL_META_EXTERNAL_DISPATCH_KEY};
use crate::conversation::Conversation;
use crate::permission::permission_confirmation::PrincipalType;
use crate::permission::{Permission, PermissionConfirmation};
use crate::providers::base::{MessageStream, PermissionRouting, Provider};
use crate::subprocess::configure_subprocess;
use crate::utils::sanitize_unicode_tags;
use goose_providers::errors::ProviderError;
use goose_providers::model::ModelConfig;
@@ -97,7 +100,7 @@ type ClientLoopFn = Box<
#[derive(Debug)]
enum AcpUpdate {
Text(String),
Text(TextContent),
Thought(String),
ToolCallStart {
id: String,
@@ -484,8 +487,17 @@ impl Provider for AcpProvider {
ProviderError::RequestFailed(format!("Failed to set ACP model option: {e}"))
})?;
let current_prompt_blocks = messages_to_prompt(messages, false);
if current_prompt_blocks.is_empty() {
return Ok(Box::pin(futures::stream::empty()));
}
let claim = self.claim_handoff_context(messages);
let prompt_blocks = messages_to_prompt(messages, claim.include_context);
let prompt_blocks = if claim.include_context {
messages_to_prompt(messages, true)
} else {
current_prompt_blocks
};
// Drop any tool-call buffer state left over from a prior prompt
// (e.g. cancelled or interrupted before its terminal status arrived).
if let Ok(mut buffer) = self.pending_tool_updates.lock() {
@@ -526,9 +538,7 @@ impl Provider for AcpProvider {
let (id, ts) = text_run
.get_or_insert_with(fresh_text_run)
.clone();
let message = Message::new(Role::Assistant, ts, vec![])
.with_text(text)
.with_id(id);
let message = acp_text_update_message(text, id, ts);
yield (Some(message), None);
}
}
@@ -815,7 +825,7 @@ impl AcpClientLoop {
{
match notification.update {
SessionUpdate::AgentMessageChunk(ContentChunk {
content: ContentBlock::Text(TextContent { text, .. }),
content: ContentBlock::Text(text),
..
}) => {
let _ = tx.try_send(AcpUpdate::Text(text));
@@ -1353,26 +1363,19 @@ fn filter_supported_servers(
}
fn messages_to_prompt(messages: &[Message], include_handoff_context: bool) -> Vec<ContentBlock> {
let mut content_blocks = Vec::new();
let Some(last_user_index) = last_user_message_index(messages) else {
return content_blocks;
return Vec::new();
};
if include_handoff_context {
if let Some(memo) = build_handoff_context_memo(&messages[..last_user_index]) {
content_blocks.push(ContentBlock::Text(TextContent::new(memo)));
}
}
let message = &messages[last_user_index];
let message = messages[last_user_index].agent_visible_content();
let mut current_prompt_blocks = Vec::new();
for content in &message.content {
match content {
MessageContent::Text(text) => {
content_blocks.push(ContentBlock::Text(TextContent::new(text.text.clone())));
current_prompt_blocks.push(ContentBlock::Text(TextContent::new(text.text.clone())));
}
MessageContent::Image(image) => {
content_blocks.push(ContentBlock::Image(ImageContent::new(
current_prompt_blocks.push(ContentBlock::Image(ImageContent::new(
&image.data,
&image.mime_type,
)));
@@ -1381,6 +1384,15 @@ fn messages_to_prompt(messages: &[Message], include_handoff_context: bool) -> Ve
}
}
if current_prompt_blocks.is_empty() || !include_handoff_context {
return current_prompt_blocks;
}
let mut content_blocks = Vec::new();
if let Some(memo) = build_handoff_context_memo(&messages[..last_user_index]) {
content_blocks.push(ContentBlock::Text(TextContent::new(memo)));
}
content_blocks.extend(current_prompt_blocks);
content_blocks
}
@@ -1399,11 +1411,12 @@ fn has_handoff_context(messages: &[Message]) -> bool {
}
fn build_handoff_context_memo(prior_messages: &[Message]) -> Option<String> {
let formatted_messages: Vec<String> = prior_messages
.iter()
.filter(|message| message.is_agent_visible())
.map(format_message_for_compacting)
.collect();
let formatted_messages: Vec<String> =
Conversation::new_unvalidated(prior_messages.iter().cloned())
.agent_visible_messages()
.iter()
.map(|message| format_message_for_compacting(&message.agent_visible_content()))
.collect();
if formatted_messages.is_empty() {
return None;
@@ -1419,6 +1432,48 @@ do not treat it as a new task or mention this handoff unless relevant."
))
}
fn acp_audience_to_rmcp(annotations: Option<&AcpAnnotations>) -> Option<Vec<Role>> {
let audience = annotations?.audience.as_ref()?;
let audience = audience
.iter()
.filter_map(|role| match role {
AcpRole::Assistant => Some(Role::Assistant),
AcpRole::User => Some(Role::User),
_ => None,
})
.collect::<Vec<_>>();
if audience.is_empty() {
None
} else {
Some(audience)
}
}
fn acp_text_content_to_rmcp(text: TextContent) -> RmcpContent {
let audience = acp_audience_to_rmcp(text.annotations.as_ref());
let mut content = RmcpContent::text(sanitize_unicode_tags(&text.text));
if let Some(audience) = audience {
content = content.with_audience(audience);
}
content
}
fn acp_image_content_to_rmcp(image: ImageContent) -> RmcpContent {
let audience = acp_audience_to_rmcp(image.annotations.as_ref());
let mut content = RmcpContent::image(image.data, image.mime_type);
if let Some(audience) = audience {
content = content.with_audience(audience);
}
content
}
fn acp_text_update_message(text: TextContent, id: String, created: i64) -> Message {
Message::new(Role::Assistant, created, vec![])
.with_content(acp_text_content_to_rmcp(text).into())
.with_id(id)
}
/// Convert ACP `ToolCallContent` blocks into the rmcp `Content` shape goose's
/// `Message::with_tool_response` consumes. Handles `Content` (text/image/other),
/// `Diff`, and `Terminal` variants; falls back to a JSON serialization of
@@ -1433,10 +1488,10 @@ fn acp_tool_call_content_to_rmcp(
match block {
ToolCallContent::Content(val) => match val.content {
ContentBlock::Text(text) => {
out.push(RmcpContent::text(text.text));
out.push(acp_text_content_to_rmcp(text));
}
ContentBlock::Image(image) => {
out.push(RmcpContent::image(image.data, image.mime_type));
out.push(acp_image_content_to_rmcp(image));
}
other => {
if let Ok(json) = serde_json::to_string(&other) {
@@ -1611,6 +1666,7 @@ mod tests {
use agent_client_protocol::schema::v1::{
SessionConfigSelectOption, SessionMode, SessionModeId,
};
use rmcp::model::AnnotateAble;
use test_case::test_case;
fn prompt_text(block: &ContentBlock) -> &str {
@@ -1690,6 +1746,26 @@ mod tests {
assert_eq!(prompt_text(&blocks[1]), "continue from there");
}
#[test]
fn messages_to_prompt_drops_user_only_acp_rows_from_handoff() {
let user_only = TextContent::new("SECRET_USER_ONLY")
.annotations(AcpAnnotations::new().audience(vec![AcpRole::User]));
let messages = vec![
Message::user().with_text("visible prior"),
acp_text_update_message(user_only, "acp-message".to_string(), 123),
Message::user().with_text("current request"),
];
let blocks = messages_to_prompt(&messages, true);
assert_eq!(blocks.len(), 2);
let memo = prompt_text(&blocks[0]);
assert!(memo.contains("visible prior"));
assert!(!memo.contains("SECRET_USER_ONLY"));
assert!(!memo.contains("<empty message>"));
assert_eq!(prompt_text(&blocks[1]), "current request");
}
#[test]
fn messages_to_prompt_keeps_latest_user_images_after_handoff_memo() {
let messages = vec![
@@ -1713,6 +1789,107 @@ mod tests {
assert_eq!(prompt_text(&blocks[2]), "describe this");
}
#[test]
fn messages_to_prompt_excludes_user_only_current_and_handoff_content() {
use rmcp::model::RawTextContent;
fn user_only_text(text: &str) -> MessageContent {
MessageContent::Text(
RawTextContent {
text: text.to_string(),
meta: None,
}
.no_annotation()
.with_audience(vec![Role::User]),
)
}
let messages = vec![
Message::user()
.with_text("visible prior")
.with_content(user_only_text("SECRET_PRIOR")),
Message::user()
.with_text("visible current")
.with_content(user_only_text("SECRET_CURRENT")),
];
let rendered = messages_to_prompt(&messages, true)
.iter()
.filter_map(|block| match block {
ContentBlock::Text(text) => Some(text.text.as_str()),
_ => None,
})
.collect::<Vec<_>>()
.join("\n");
assert!(rendered.contains("visible prior"));
assert!(rendered.contains("visible current"));
assert!(!rendered.contains("SECRET_PRIOR"));
assert!(!rendered.contains("SECRET_CURRENT"));
}
#[test]
fn messages_to_prompt_drops_handoff_when_current_content_is_user_only() {
use rmcp::model::RawTextContent;
let current = MessageContent::Text(
RawTextContent {
text: "user-only".to_string(),
meta: None,
}
.no_annotation()
.with_audience(vec![Role::User]),
);
let messages = vec![
Message::assistant().with_text("prior context"),
Message::user().with_content(current),
];
assert!(messages_to_prompt(&messages, true).is_empty());
}
#[tokio::test]
async fn stream_skips_user_only_prompt_without_claiming_handoff_context() {
use futures::StreamExt;
use rmcp::model::RawTextContent;
let (tx, mut rx) = mpsc::channel(1);
let (provider, model) = test_provider_with_tx(Some(tx));
let current = MessageContent::Text(
RawTextContent {
text: "user-only".to_string(),
meta: None,
}
.no_annotation()
.with_audience(vec![Role::User]),
);
let messages = vec![
Message::assistant().with_text("prior context"),
Message::user().with_content(current),
];
let mut stream = provider.stream(&model, "", &messages, &[]).await.unwrap();
assert!(stream.next().await.is_none());
assert!(rx.try_recv().is_err());
assert!(!provider.handoff_context_sent.load(Ordering::Acquire));
}
#[test]
fn live_acp_text_update_preserves_assistant_only_audience() {
let text = TextContent::new("assistant-only")
.annotations(AcpAnnotations::new().audience(vec![AcpRole::Assistant]));
let message = acp_text_update_message(text, "message-id".to_string(), 123);
let MessageContent::Text(text) = &message.content[0] else {
panic!("expected text content");
};
let audience = text.audience().expect("audience annotation should survive");
assert!(audience.contains(&Role::Assistant));
assert!(!audience.contains(&Role::User));
}
#[test]
fn handoff_context_is_sent_only_on_first_provider_prompt() {
let (provider, _) = test_provider();
@@ -2351,6 +2528,35 @@ mod tests {
);
}
#[test]
fn acp_tool_call_content_preserves_audience_annotations() {
let text_block = ToolCallContent::Content(agent_client_protocol::schema::v1::Content::new(
ContentBlock::Text(
TextContent::new("user-only")
.annotations(AcpAnnotations::new().audience(vec![AcpRole::User])),
),
));
let image_block = ToolCallContent::Content(
agent_client_protocol::schema::v1::Content::new(ContentBlock::Image(
ImageContent::new("base64data", "image/png")
.annotations(AcpAnnotations::new().audience(vec![AcpRole::Assistant])),
)),
);
let out = acp_tool_call_content_to_rmcp(Some(vec![text_block, image_block]), None);
let text_audience = out[0]
.audience()
.expect("text audience annotation should survive");
assert!(text_audience.contains(&Role::User));
assert!(!text_audience.contains(&Role::Assistant));
let image_audience = out[1]
.audience()
.expect("image audience annotation should survive");
assert!(image_audience.contains(&Role::Assistant));
assert!(!image_audience.contains(&Role::User));
}
#[test]
fn acp_tool_call_content_falls_back_to_raw_output_when_blocks_empty() {
let out =
crates/goose/src/acp/server/load_session.rs
+1 -5
View File
@@ -38,7 +38,7 @@ fn replay_conversation_to_client(
let messages = session
.conversation
.as_ref()
.map(|c| c.messages().to_vec())
.map(|c| c.user_visible_messages())
.unwrap_or_default();
debug!(
target: "perf",
@@ -51,10 +51,6 @@ fn replay_conversation_to_client(
HashMap::<String, crate::conversation::message::ToolRequest>::new();
for message in &messages {
if !message.metadata.user_visible {
continue;
}
for content_item in &message.content {
match content_item {
MessageContent::Text(text) => {
crates/goose/src/agents/agent.rs
+217 -10
View File
@@ -276,6 +276,14 @@ pub enum AgentEvent {
HistoryReplaced(Conversation),
}
fn project_message_for_user_event(message: &Message) -> Message {
message.user_visible_content()
}
fn agent_visible_message_text(message: &Message) -> String {
message.agent_visible_content().as_concat_text()
}
fn attach_turn_usage(
messages: &mut Conversation,
usage: &ProviderUsage,
@@ -285,9 +293,10 @@ fn attach_turn_usage(
.iter_mut()
.rev()
.find(|m| m.role == rmcp::model::Role::Assistant)?;
let has_user_visible_content = !message.user_visible_content().content.is_empty();
let message_usage = MessageUsage::from_provider_usage(usage, false);
message.metadata.usage = Some(Box::new(message_usage.clone()));
Some((message.id.clone(), message_usage))
has_user_visible_content.then(|| (message.id.clone(), message_usage))
}
impl Default for Agent {
@@ -1541,7 +1550,7 @@ impl Agent {
) -> Result<BoxStream<'_, Result<AgentEvent>>> {
let session_manager = self.config.session_manager.clone();
let message_text_for_trace = user_message.as_concat_text();
let message_text_for_trace = agent_visible_message_text(&user_message);
tracing::Span::current().record("user_message", message_text_for_trace.as_str());
tracing::Span::current().record("trace_input", message_text_for_trace.as_str());
@@ -1581,17 +1590,34 @@ impl Agent {
}
}
let message_text = user_message.as_concat_text();
let message_text = message_text_for_trace;
let session = session_manager
.get_session(&session_config.id, true)
.await?;
let is_first_turn = session
let is_first_agent_turn = session
.conversation
.as_ref()
.map(|conversation| conversation.messages().is_empty())
.map(|conversation| {
conversation.messages().iter().all(|message| {
!message.is_agent_visible()
|| message.agent_visible_content().content.is_empty()
})
})
.unwrap_or(true);
if is_first_turn {
if !user_message.is_agent_visible()
|| user_message.agent_visible_content().content.is_empty()
{
let user_visibility = user_message.is_user_visible();
let user_message = user_message.with_visibility(user_visibility, false);
session_manager
.add_message(&session_config.id, &user_message)
.await?;
return Ok(Box::pin(futures::stream::empty()));
}
if is_first_agent_turn {
self.emit_hook(crate::hooks::HookEvent::SessionStart, &session_config.id)
.await;
}
@@ -1921,7 +1947,7 @@ impl Agent {
if can_drain_pending_steers {
for message in self.drain_pending_steers(&session_config.id).await {
let message_text = message.as_concat_text();
let message_text = agent_visible_message_text(&message);
if self
.hook_manager
.has_hooks(crate::hooks::HookEvent::UserPromptSubmit)
@@ -2036,6 +2062,7 @@ impl Agent {
let mut did_recovery_compact_this_iteration = false;
let mut exit_chat = false;
let mut provider_errored = false;
let mut provider_produced_content = false;
let mut pending_final_output: Option<String> = None;
let mut pending_turn_usage: Option<ProviderUsage> = None;
@@ -2071,6 +2098,24 @@ impl Agent {
continue;
}
provider_produced_content |= response.content.iter().any(|content| {
match content {
MessageContent::Text(text) => !text.text.is_empty(),
MessageContent::Image(image) => !image.data.is_empty(),
MessageContent::Thinking(thinking) => {
!thinking.thinking.is_empty()
|| !thinking.signature.is_empty()
}
MessageContent::RedactedThinking(thinking) => {
!thinking.data.is_empty()
}
MessageContent::SystemNotification(notification) => {
!notification.msg.is_empty()
}
_ => true,
}
});
let ToolCategorizeResult {
frontend_requests,
remaining_requests,
@@ -2104,8 +2149,10 @@ impl Agent {
},
);
yield AgentEvent::Message(filtered_response.clone());
tokio::task::yield_now().await;
if !filtered_response.content.is_empty() {
yield AgentEvent::Message(filtered_response.clone());
tokio::task::yield_now().await;
}
let num_tool_requests = frontend_requests.len() + remaining_requests.len();
if num_tool_requests == 0 {
@@ -2467,7 +2514,7 @@ impl Agent {
request_msg.created = final_response.created;
}
messages_to_add.push(request_msg);
yield AgentEvent::Message(final_response.clone());
yield AgentEvent::Message(project_message_for_user_event(&final_response));
messages_to_add.push(final_response);
}
@@ -2633,6 +2680,7 @@ impl Agent {
&& !exit_chat
&& !provider_errored
&& !did_recovery_compact_this_iteration
&& !provider_produced_content
&& last_assistant_text.is_empty();
if empty_response {
@@ -3261,6 +3309,7 @@ impl Agent {
.filter(super::reply_parts::is_tool_visible_to_model)
.collect();
messages = Conversation::new_unvalidated(messages.agent_visible_messages());
messages.push(Message::user().with_text(recipe_prompt));
let (messages, issues) = fix_conversation(messages);
@@ -3472,6 +3521,48 @@ mod tests {
));
}
#[test]
fn user_event_projection_preserves_hidden_tool_response_wrapper() {
use rmcp::model::{Content, Role};
let hidden_only = Message::user().with_tool_response(
"tool-1",
Ok(CallToolResult::success(vec![Content::text(
"provider-only",
)
.with_audience(vec![Role::Assistant])])),
);
let projected = project_message_for_user_event(&hidden_only);
let result = projected.content[0]
.as_tool_response()
.expect("hidden tool response wrapper")
.tool_result
.as_ref()
.expect("successful hidden tool result");
assert!(result.content.is_empty());
}
#[test]
fn agent_visible_message_text_excludes_user_only_blocks() {
use rmcp::model::{AnnotateAble, RawTextContent, Role};
let user_only = RawTextContent {
text: "SECRET_USER_ONLY".to_string(),
meta: None,
}
.no_annotation()
.with_audience(vec![Role::User]);
let message = Message::user()
.with_text("/goal visible objective")
.with_content(MessageContent::Text(user_only));
assert_eq!(
agent_visible_message_text(&message),
"/goal visible objective"
);
}
struct ActionRequiredProvider {
handled: tokio::sync::Mutex<Vec<(String, PermissionConfirmation)>>,
}
@@ -3956,6 +4047,90 @@ echo start >> "$PLUGIN_ROOT/hook.log"
Ok(())
}
#[tokio::test]
async fn skipped_user_message_does_not_enter_empty_response_retry_loop() -> Result<()> {
use rmcp::model::{AnnotateAble, RawTextContent, Role};
let env = SessionStartHookTestEnv::new()?;
let provider = Arc::new(CountingTextProvider::new());
let hook_manager = env.hook_manager();
let (agent, session_id) =
create_test_agent(env.data_dir(), hook_manager, provider.clone()).await?;
let session_config = SessionConfig {
id: session_id.clone(),
schedule_id: None,
max_turns: Some(10),
retry_config: None,
};
let user_only_content = MessageContent::Text(
RawTextContent {
text: "user-only".to_string(),
meta: None,
}
.no_annotation()
.with_audience(vec![Role::User]),
);
let mut stream = agent
.reply(
Message::user().with_content(user_only_content),
session_config,
None,
)
.await?;
assert!(stream.next().await.is_none());
assert_eq!(provider.call_count.load(Ordering::SeqCst), 0);
assert_eq!(env.hook_invocations(), 0);
let session = agent
.config
.session_manager
.get_session(&session_id, true)
.await?;
let conversation = session.conversation.unwrap();
assert_eq!(conversation.messages().len(), 1);
assert!(!conversation.messages()[0].is_agent_visible());
let visible_session_config = SessionConfig {
id: session_id.clone(),
schedule_id: None,
max_turns: Some(10),
retry_config: None,
};
let mut visible_stream = agent
.reply(
Message::user().with_text("agent-visible"),
visible_session_config,
None,
)
.await?;
while let Some(event) = visible_stream.next().await {
event?;
}
assert_eq!(provider.call_count.load(Ordering::SeqCst), 1);
assert_eq!(env.hook_invocations(), 1);
let final_session_config = SessionConfig {
id: session_id,
schedule_id: None,
max_turns: Some(10),
retry_config: None,
};
let mut final_stream = agent
.reply(
Message::user().with_text("second-agent-visible"),
final_session_config,
None,
)
.await?;
while let Some(event) = final_stream.next().await {
event?;
}
assert_eq!(provider.call_count.load(Ordering::SeqCst), 2);
assert_eq!(env.hook_invocations(), 1);
Ok(())
}
#[tokio::test]
async fn stop_hook_block_cap_allows_configured_consecutive_blocks_then_overrides() -> Result<()>
{
@@ -4209,4 +4384,36 @@ echo start >> "$PLUGIN_ROOT/hook.log"
"user message must stay untouched"
);
}
#[test]
fn attach_turn_usage_suppresses_notification_for_assistant_only_message() {
use rmcp::model::{AnnotateAble, RawTextContent, Role};
let usage = ProviderUsage::new(
"test-model".to_string(),
Usage::new(Some(1200), Some(340), None),
);
let assistant_only = RawTextContent {
text: "provider-only state".to_string(),
meta: None,
}
.no_annotation()
.with_audience(vec![Role::Assistant]);
let mut conversation = Conversation::new_unvalidated([
Message::user().with_text("hi"),
Message::assistant()
.with_id("hidden")
.with_content(MessageContent::Text(assistant_only)),
]);
assert!(attach_turn_usage(&mut conversation, &usage).is_none());
let stored = conversation.messages()[1]
.metadata
.usage
.as_deref()
.expect("usage must remain stored on the hidden assistant message");
assert_eq!(stored.input_tokens, Some(1200));
assert_eq!(stored.output_tokens, Some(340));
}
}
crates/goose/src/agents/platform_extensions/chatrecall.rs
+101 -37
View File
@@ -1,6 +1,7 @@
use crate::agents::extension::PlatformExtensionContext;
use crate::agents::mcp_client::{Error, McpClientTrait};
use crate::agents::tool_execution::ToolCallContext;
use crate::conversation::Conversation;
use crate::session::session_manager::SessionType;
use anyhow::Result;
use async_trait::async_trait;
@@ -39,6 +40,46 @@ pub struct ChatRecallClient {
context: PlatformExtensionContext,
}
fn format_agent_visible_excerpt(conversation: &Conversation) -> Option<(usize, String)> {
let messages = conversation.agent_visible_messages();
let total = messages.len();
if total == 0 {
return None;
}
let mut output = String::new();
let first_count = std::cmp::min(3, total);
output.push_str("--- First Few Messages ---\n\n");
for (idx, message) in messages.iter().take(first_count).enumerate() {
output.push_str(&format!("{}. [{:?}] ", idx + 1, message.role));
for content in &message.content {
if let Some(text) = content.as_text() {
output.push_str(text);
output.push('\n');
}
}
output.push('\n');
}
if total > first_count {
output.push_str("--- Last Few Messages ---\n\n");
let last_count = std::cmp::min(3, total);
let skip_count = total.saturating_sub(last_count);
for (idx, message) in messages.iter().skip(skip_count).enumerate() {
output.push_str(&format!("{}. [{:?}] ", skip_count + idx + 1, message.role));
for content in &message.content {
if let Some(text) = content.as_text() {
output.push_str(text);
output.push('\n');
}
}
output.push('\n');
}
}
Some((total, output))
}
impl ChatRecallClient {
pub fn new(context: PlatformExtensionContext) -> Result<Self> {
let info = InitializeResult::new(ServerCapabilities::builder().enable_tools().build())
@@ -92,15 +133,14 @@ impl ChatRecallClient {
))]);
}
let msgs = conversation.unwrap().messages();
let total = msgs.len();
if total == 0 {
let Some((total, excerpt)) =
format_agent_visible_excerpt(conversation.unwrap())
else {
return Ok(vec![Content::text(format!(
"Session {} has no messages.",
sid
))]);
}
};
let mut output = format!(
"Session: {} (ID: {})\nWorking Dir: {}\nTotal Messages: {}\n\n",
@@ -110,38 +150,7 @@ impl ChatRecallClient {
total
);
let first_count = std::cmp::min(3, total);
output.push_str("--- First Few Messages ---\n\n");
for (idx, msg) in msgs.iter().take(first_count).enumerate() {
output.push_str(&format!("{}. [{:?}] ", idx + 1, msg.role));
for content in &msg.content {
if let Some(text) = content.as_text() {
output.push_str(text);
output.push('\n');
}
}
output.push('\n');
}
if total > first_count {
output.push_str("--- Last Few Messages ---\n\n");
let last_count = std::cmp::min(3, total);
let skip_count = total.saturating_sub(last_count);
for (idx, msg) in msgs.iter().skip(skip_count).enumerate() {
output.push_str(&format!(
"{}. [{:?}] ",
skip_count + idx + 1,
msg.role
));
for content in &msg.content {
if let Some(text) = content.as_text() {
output.push_str(text);
output.push('\n');
}
}
output.push('\n');
}
}
output.push_str(&excerpt);
Ok(vec![Content::text(output)])
}
@@ -307,3 +316,58 @@ impl McpClientTrait for ChatRecallClient {
Some(&self.info)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::conversation::message::{Message, MessageContent, MessageMetadata};
use rmcp::model::{AnnotateAble, RawTextContent, Role};
fn annotated_text(text: &str, audience: Vec<Role>) -> MessageContent {
MessageContent::Text(
RawTextContent {
text: text.to_string(),
meta: None,
}
.no_annotation()
.with_audience(audience),
)
}
#[test]
fn loaded_excerpt_projects_audience_before_selecting_endpoints() {
let conversation = Conversation::new_unvalidated([
Message::user()
.with_text("hidden first row")
.with_metadata(MessageMetadata::user_only()),
Message::user()
.with_text("visible first")
.with_content(annotated_text("user-only first secret", vec![Role::User])),
Message::assistant().with_text("visible middle"),
Message::assistant()
.with_text("hidden last row")
.with_metadata(MessageMetadata::user_only()),
Message::user()
.with_text("visible last")
.with_content(annotated_text("user-only last secret", vec![Role::User])),
]);
let (total, excerpt) = format_agent_visible_excerpt(&conversation).unwrap();
assert_eq!(total, 3);
assert!(excerpt.contains("visible first"));
assert!(excerpt.contains("visible last"));
assert!(!excerpt.contains("hidden first row"));
assert!(!excerpt.contains("hidden last row"));
assert!(!excerpt.contains("user-only first secret"));
assert!(!excerpt.contains("user-only last secret"));
let canonical_user_text = conversation
.user_visible_messages()
.iter()
.map(Message::as_concat_text)
.collect::<Vec<_>>()
.join("\n");
assert!(canonical_user_text.contains("user-only first secret"));
assert!(canonical_user_text.contains("user-only last secret"));
}
}
crates/goose/src/agents/platform_extensions/orchestrator.rs
+43 -3
View File
@@ -5,6 +5,7 @@ use crate::agents::{AgentEvent, SessionConfig};
use crate::config::{Config, ExtensionConfig, GooseMode};
use crate::context_mgmt::format_message_for_compacting;
use crate::conversation::message::Message;
use crate::conversation::Conversation;
use crate::execution::manager::AgentManager;
use crate::providers;
use crate::providers::base::Provider;
@@ -282,7 +283,7 @@ impl OrchestratorClient {
match mode {
"first_last" => {
if let Some(conversation) = &session.conversation {
let messages = conversation.messages();
let messages = agent_visible_session_messages(conversation);
if messages.is_empty() {
output.push("No messages in this session.".to_string());
} else {
@@ -335,9 +336,9 @@ impl OrchestratorClient {
) -> Result<String, String> {
let provider = self.get_provider().await?;
let conversation_text = messages
let conversation_text = Conversation::new_unvalidated(messages.iter().cloned())
.agent_visible_messages()
.iter()
.filter(|m| m.is_agent_visible())
.map(format_message_for_compacting)
.collect::<Vec<_>>()
.join("\n");
@@ -576,6 +577,10 @@ impl OrchestratorClient {
}
}
fn agent_visible_session_messages(conversation: &Conversation) -> Vec<Message> {
conversation.agent_visible_messages()
}
#[async_trait]
impl McpClientTrait for OrchestratorClient {
async fn list_tools(
@@ -672,3 +677,38 @@ fn extract_string(args: &JsonObject, key: &str) -> Result<String, String> {
.map(|s| s.to_string())
.ok_or_else(|| format!("Missing or invalid '{}'", key))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::conversation::message::MessageContent;
use rmcp::model::{AnnotateAble, RawTextContent, Role};
#[test]
fn first_last_projection_drops_hidden_endpoints_and_content() {
let user_only = |text: &str| {
MessageContent::Text(
RawTextContent {
text: text.to_string(),
meta: None,
}
.no_annotation()
.with_audience(vec![Role::User]),
)
};
let conversation = Conversation::new_unvalidated([
Message::assistant().with_content(user_only("hidden first")),
Message::user().with_text("visible first"),
Message::assistant()
.with_content(user_only("hidden block"))
.with_text("visible last"),
Message::assistant().with_content(user_only("hidden last")),
]);
let messages = agent_visible_session_messages(&conversation);
assert_eq!(messages.len(), 2);
assert_eq!(messages[0].as_concat_text(), "visible first");
assert_eq!(messages[1].as_concat_text(), "visible last");
}
}
crates/goose/src/agents/reply_parts.rs
+197 -10
View File
@@ -13,7 +13,7 @@ use super::super::agents::Agent;
use crate::agents::platform_extensions::code_execution;
use crate::config::Config;
use crate::conversation::message::{Message, MessageContent, MessageUsage, ToolRequest};
use crate::conversation::Conversation;
use crate::conversation::{fix_conversation, Conversation};
#[cfg(test)]
use crate::providers::base::stream_from_single_message;
use crate::providers::base::{MessageStream, Provider};
@@ -291,17 +291,16 @@ impl Agent {
) -> Result<MessageStream, ProviderError> {
let config = model_config.clone();
let filtered_messages: Vec<Message> = messages
.iter()
.filter(|m| m.is_agent_visible())
.map(|m| m.agent_visible_content())
.collect();
let projected_messages =
Conversation::new_unvalidated(messages.iter().cloned()).agent_visible_messages();
let (filtered_messages, _) =
fix_conversation(Conversation::new_unvalidated(projected_messages));
// Convert tool messages to text if toolshim is enabled
let messages_for_provider = if config.toolshim {
convert_tool_messages_to_text(&filtered_messages)
convert_tool_messages_to_text(filtered_messages.messages())
} else {
Conversation::new_unvalidated(filtered_messages)
filtered_messages
};
// Clone owned data to move into the async stream
@@ -364,7 +363,7 @@ impl Agent {
(
Some(MessageContent::Text(last_text)),
MessageContent::Text(new_text),
) => {
) if last_text.audience() == new_text.audience() => {
last_text.text.push_str(&new_text.text);
}
_ => {
@@ -526,7 +525,9 @@ impl Agent {
MessageContent::Thinking(_) | MessageContent::RedactedThinking(_)
if should_suppress_replayed_thinking => {}
_ => {
filtered_content.push(content.clone());
if let Some(content) = user_visible_provider_content(content) {
filtered_content.push(content);
}
}
}
}
@@ -621,6 +622,10 @@ impl Agent {
}
}
fn user_visible_provider_content(content: &MessageContent) -> Option<MessageContent> {
content.user_visible_content()
}
/// Check whether a tool should be callable by an app based on MCP Apps visibility metadata.
///
/// Per the MCP Apps spec (2026-01-26), if `_meta.ui.visibility` is present and does not
@@ -674,7 +679,9 @@ mod tests {
use async_trait::async_trait;
use goose_providers::conversation::token_usage::{ProviderStats, ProviderUsage, Usage};
use goose_providers::model::ModelConfig;
use rmcp::model::{AnnotateAble, RawTextContent, Role};
use rmcp::object;
use std::sync::Mutex;
use std::time::{Duration, Instant};
#[derive(Clone)]
@@ -699,6 +706,161 @@ mod tests {
}
}
#[derive(Clone)]
struct CapturingProvider {
messages: Arc<Mutex<Vec<Message>>>,
}
#[async_trait]
impl Provider for CapturingProvider {
fn get_name(&self) -> &str {
"capturing"
}
async fn stream(
&self,
_model_config: &ModelConfig,
_system: &str,
messages: &[Message],
_tools: &[Tool],
) -> Result<MessageStream, ProviderError> {
*self.messages.lock().unwrap() = messages.to_vec();
let message = Message::assistant().with_text("ok");
let usage = ProviderUsage::new("capturing".to_string(), Usage::default());
Ok(stream_from_single_message(message, usage))
}
}
#[tokio::test]
async fn provider_input_drops_rows_empty_after_agent_projection() {
let user_only = RawTextContent {
text: "user-only ACP output".to_string(),
meta: None,
}
.no_annotation()
.with_audience(vec![Role::User]);
let messages = vec![
Message::assistant().with_content(MessageContent::Text(user_only)),
Message::user().with_text("current request"),
];
let captured = Arc::new(Mutex::new(Vec::new()));
let provider = Arc::new(CapturingProvider {
messages: captured.clone(),
});
let _stream = crate::agents::Agent::stream_response_from_provider(
provider,
ModelConfig::new("test-model"),
"test-session",
"system",
&messages,
&[],
&[],
)
.await
.unwrap();
let captured = captured.lock().unwrap();
assert_eq!(captured.len(), 1);
assert_eq!(captured[0].role, Role::User);
assert_eq!(captured[0].as_concat_text(), "current request");
}
#[tokio::test]
async fn provider_input_refixes_roles_after_agent_projection() {
let user_only = RawTextContent {
text: "hidden separator".to_string(),
meta: None,
}
.no_annotation()
.with_audience(vec![Role::User]);
let messages = vec![
Message::user().with_text("first request"),
Message::assistant().with_content(MessageContent::Text(user_only)),
Message::user().with_text("second request"),
];
let captured = Arc::new(Mutex::new(Vec::new()));
let provider = Arc::new(CapturingProvider {
messages: captured.clone(),
});
let _stream = crate::agents::Agent::stream_response_from_provider(
provider,
ModelConfig::new("test-model"),
"test-session",
"system",
&messages,
&[],
&[],
)
.await
.unwrap();
let captured = captured.lock().unwrap();
assert_eq!(captured.len(), 1);
assert_eq!(captured[0].role, Role::User);
assert_eq!(
captured[0].as_concat_text(),
"first request\nsecond request"
);
assert!(!captured[0].as_concat_text().contains("hidden separator"));
}
#[tokio::test]
async fn provider_input_refixes_tool_result_emptied_by_agent_projection() {
let user_only_result =
rmcp::model::Content::text("hidden result").with_audience(vec![Role::User]);
let messages = vec![
Message::user().with_text("run the tool"),
Message::assistant().with_tool_request(
"tool-1",
Ok(rmcp::model::CallToolRequestParams::new("test_tool")),
),
Message::user().with_tool_response(
"tool-1",
Ok(rmcp::model::CallToolResult::success(vec![user_only_result])),
),
];
let captured = Arc::new(Mutex::new(Vec::new()));
let provider = Arc::new(CapturingProvider {
messages: captured.clone(),
});
let _stream = crate::agents::Agent::stream_response_from_provider(
provider,
ModelConfig::new("test-model"),
"test-session",
"system",
&messages,
&[],
&[],
)
.await
.unwrap();
let captured = captured.lock().unwrap();
let tool_response = captured
.iter()
.flat_map(|message| &message.content)
.find_map(|content| match content {
MessageContent::ToolResponse(response) => Some(response),
_ => None,
})
.expect("projected tool response should remain paired");
let result = tool_response
.tool_result
.as_ref()
.expect("tool response should remain successful");
assert_eq!(result.content.len(), 1);
assert_eq!(
result.content[0]
.as_text()
.expect("placeholder should be text")
.text,
"(empty result)"
);
}
#[tokio::test]
async fn prepare_tools_returns_sorted_tools_including_frontend() -> anyhow::Result<()> {
let data_dir = tempfile::tempdir()?;
@@ -859,6 +1021,31 @@ mod tests {
));
}
#[tokio::test]
async fn categorize_tool_requests_excludes_assistant_only_text_from_user_events() {
let agent = crate::agents::Agent::new();
let assistant_only = RawTextContent {
text: "assistant-only".to_string(),
meta: None,
}
.no_annotation()
.with_audience(vec![Role::Assistant]);
let response = Message::assistant()
.with_content(MessageContent::Text(assistant_only))
.with_text("user-visible")
.with_thinking("visible reasoning", "");
let (_frontend_requests, _other_requests, filtered_message) =
agent.categorize_tool_requests(&response, &[], false).await;
assert_eq!(response.as_concat_text(), "assistant-only\nuser-visible");
assert_eq!(filtered_message.as_concat_text(), "user-visible");
assert!(filtered_message
.content
.iter()
.any(|content| matches!(content, MessageContent::Thinking(_))));
}
#[tokio::test]
async fn categorize_tool_requests_skips_externally_dispatched_and_preserves_marker() {
// External requests must (1) survive coercion with goose.external_dispatch
crates/goose/src/context_mgmt/mod.rs
+221 -57
View File
@@ -89,37 +89,32 @@ pub async fn compact_messages(
has_text && !has_tool_content
};
let extract_text = |msg: &Message| -> Option<String> {
let text_parts: Vec<String> = msg
.content
.iter()
.filter_map(|c| {
if let MessageContent::Text(text) = c {
Some(text.text.clone())
} else {
None
}
})
.collect();
if text_parts.is_empty() {
None
} else {
Some(text_parts.join("\n"))
}
};
// Find and preserve the most recent user message for non-manual compacts
let (preserved_user_message, is_most_recent) = if !manual_compact {
let found_msg = messages.iter().enumerate().rev().find(|(_, msg)| {
msg.is_agent_visible()
&& matches!(msg.role, rmcp::model::Role::User)
&& has_text_only(msg)
let found_msg = messages.iter().enumerate().rev().find_map(|(idx, msg)| {
if !msg.is_agent_visible() || !matches!(msg.role, rmcp::model::Role::User) {
return None;
}
let projected = msg.agent_visible_content();
if !has_text_only(&projected) {
return None;
}
let preserved = projected
.content
.into_iter()
.filter(|content| matches!(content, MessageContent::Text(_)))
.fold(
Message::user().with_metadata(MessageMetadata::agent_only()),
Message::with_content,
);
Some((idx, preserved))
});
if let Some((idx, msg)) = found_msg {
let is_last = idx == messages.len() - 1;
(Some(msg.clone()), is_last)
(Some(msg), is_last)
} else {
(None, false)
}
@@ -138,16 +133,8 @@ pub async fn compact_messages(
// 3. Assistant messages to continue the conversation are also agent_visible but not user_visible
let mut final_messages = Vec::new();
for (idx, msg) in messages_to_compact.iter().enumerate() {
let updated_metadata = if is_most_recent
&& idx == messages_to_compact.len() - 1
&& preserved_user_message.is_some()
{
// This is the most recent message and we're preserving it by adding a fresh copy
MessageMetadata::invisible()
} else {
msg.metadata.clone().with_agent_invisible()
};
for msg in messages_to_compact {
let updated_metadata = msg.metadata.clone().with_agent_invisible();
let updated_msg = msg.clone().with_metadata(updated_metadata);
final_messages.push(updated_msg);
}
@@ -173,9 +160,7 @@ pub async fn compact_messages(
final_messages.extend(merged_continuation);
if let Some(user_msg) = preserved_user_message {
if let Some(text) = extract_text(&user_msg) {
final_messages.push(Message::user().with_text(&text));
}
final_messages.push(user_msg);
}
Ok((
@@ -295,11 +280,8 @@ async fn do_compact(
session_id: &str,
messages: &[Message],
) -> Result<(Message, ProviderUsage), anyhow::Error> {
let agent_visible_messages: Vec<Message> = messages
.iter()
.filter(|msg| msg.is_agent_visible())
.map(|msg| msg.agent_visible_content())
.collect();
let agent_visible_messages =
Conversation::new_unvalidated(messages.iter().cloned()).agent_visible_messages();
// Try progressively removing more tool response messages from the middle to reduce context length
let removal_percentages = [0, 10, 20, 50, 100];
@@ -490,16 +472,9 @@ pub fn tool_ids_to_summarize(
.collect()
}
pub async fn summarize_tool_call(
provider: &dyn Provider,
model_config: &ModelConfig,
session_id: &str,
conversation: &Conversation,
tool_id: &str,
) -> Result<Message> {
let messages = conversation.messages();
let matching_messages: Vec<&Message> = messages
fn agent_visible_tool_pair(conversation: &Conversation, tool_id: &str) -> Result<Vec<Message>> {
let matching_messages = conversation
.messages()
.iter()
.filter(|m| {
m.content.iter().any(|c| match c {
@@ -508,18 +483,42 @@ pub async fn summarize_tool_call(
_ => false,
})
})
.collect();
.cloned()
.collect::<Vec<_>>();
let matching_messages =
Conversation::new_unvalidated(matching_messages).agent_visible_messages();
if matching_messages.is_empty() {
let has_request = matching_messages.iter().any(|message| {
message.content.iter().any(
|content| matches!(content, MessageContent::ToolRequest(request) if request.id == tool_id),
)
});
let has_response = matching_messages.iter().any(|message| {
message.content.iter().any(
|content| matches!(content, MessageContent::ToolResponse(response) if response.id == tool_id),
)
});
if !has_request || !has_response {
return Err(anyhow::anyhow!(
"No messages found for tool id: {}",
"No agent-visible tool pair found for tool id: {}",
tool_id
));
}
Ok(matching_messages)
}
pub async fn summarize_tool_call(
provider: &dyn Provider,
model_config: &ModelConfig,
session_id: &str,
conversation: &Conversation,
tool_id: &str,
) -> Result<Message> {
let matching_messages = agent_visible_tool_pair(conversation, tool_id)?;
let formatted = matching_messages
.iter()
.map(|msg| format_message_for_compacting(msg))
.map(format_message_for_compacting)
.collect::<Vec<_>>()
.join("\n");
@@ -735,6 +734,171 @@ mod tests {
.expect("compaction should produce a valid conversation");
}
#[tokio::test]
async fn preserved_user_message_keeps_audience_projection_after_compaction() {
use rmcp::model::{RawTextContent, Role};
let annotated_text = |text: &str, audience| {
MessageContent::Text(
RawTextContent {
text: text.to_string(),
meta: None,
}
.no_annotation()
.with_audience(audience),
)
};
let current_request = Message::user()
.with_text("visible current request")
.with_content(annotated_text("user-only secret", vec![Role::User]))
.with_content(annotated_text(
"assistant-only preprompt",
vec![Role::Assistant],
));
let conversation = Conversation::new_unvalidated([
Message::user().with_text("earlier request"),
Message::assistant().with_text("earlier response"),
current_request,
]);
let provider = MockProvider::new(Message::assistant().with_text("summary"), 1000);
let (compacted, _) = compact_messages(
&provider,
&provider.config,
"test-session-id",
&conversation,
false,
)
.await
.unwrap();
let preserved_copies = compacted
.messages()
.iter()
.filter(|message| message.as_concat_text().contains("visible current request"))
.collect::<Vec<_>>();
assert_eq!(preserved_copies.len(), 2);
let archived = preserved_copies
.iter()
.find(|message| message.is_user_visible())
.unwrap();
assert!(!archived.is_agent_visible());
assert!(archived.as_concat_text().contains("user-only secret"));
let replay = preserved_copies
.iter()
.find(|message| message.is_agent_visible())
.unwrap();
assert!(!replay.is_user_visible());
assert!(replay.as_concat_text().contains("assistant-only preprompt"));
assert!(!replay.as_concat_text().contains("user-only secret"));
let agent_text = compacted
.agent_visible_messages()
.iter()
.map(Message::as_concat_text)
.collect::<Vec<_>>()
.join("\n");
assert!(agent_text.contains("visible current request"));
assert!(agent_text.contains("assistant-only preprompt"));
assert!(!agent_text.contains("user-only secret"));
let user_text = compacted
.user_visible_messages()
.iter()
.map(Message::as_concat_text)
.collect::<Vec<_>>()
.join("\n");
assert!(user_text.contains("user-only secret"));
assert!(!user_text.contains("assistant-only preprompt"));
}
#[tokio::test]
async fn tool_pair_summary_projects_nested_audiences_before_provider_input() {
let provider = MockProvider::new(Message::assistant().with_text("summary"), 1000);
let conversation = Conversation::new_unvalidated([
Message::assistant()
.with_tool_request("tool_0", Ok(CallToolRequestParams::new("read_file"))),
Message::user().with_tool_response(
"tool_0",
Ok(rmcp::model::CallToolResult::success(vec![
RawContent::text("visible result").no_annotation(),
RawContent::text("user-only secret")
.no_annotation()
.with_audience(vec![Role::User]),
])),
),
]);
let projected = agent_visible_tool_pair(&conversation, "tool_0").unwrap();
let formatted = projected
.iter()
.map(format_message_for_compacting)
.collect::<Vec<_>>()
.join("\n");
assert!(formatted.contains("visible result"));
assert!(!formatted.contains("user-only secret"));
let user_only_conversation = Conversation::new_unvalidated([
Message::assistant()
.with_tool_request("tool_1", Ok(CallToolRequestParams::new("read_file"))),
Message::user().with_tool_response(
"tool_1",
Ok(rmcp::model::CallToolResult::success(vec![
RawContent::text("user-only secret")
.no_annotation()
.with_audience(vec![Role::User]),
])),
),
]);
let user_only_formatted = agent_visible_tool_pair(&user_only_conversation, "tool_1")
.unwrap()
.iter()
.map(format_message_for_compacting)
.collect::<Vec<_>>()
.join("\n");
assert!(!user_only_formatted.contains("user-only secret"));
summarize_tool_call(
&provider,
&provider.config,
"test-session-id",
&conversation,
"tool_0",
)
.await
.unwrap();
}
#[tokio::test]
async fn tool_pair_summary_rejects_agent_hidden_response() {
let provider = MockProvider::new(Message::assistant().with_text("summary"), 1000);
let conversation = Conversation::new_unvalidated([
Message::assistant()
.with_tool_request("tool_0", Ok(CallToolRequestParams::new("read_file"))),
Message::user()
.with_tool_response(
"tool_0",
Ok(rmcp::model::CallToolResult::success(vec![
RawContent::text("user-only secret").no_annotation(),
])),
)
.with_metadata(MessageMetadata::user_only()),
]);
let error = summarize_tool_call(
&provider,
&provider.config,
"test-session-id",
&conversation,
"tool_0",
)
.await
.unwrap_err();
assert!(error.to_string().contains("No agent-visible tool pair"));
}
#[tokio::test]
async fn test_progressive_removal_on_context_exceeded() {
let response_message = Message::assistant().with_text("<mock summary>");
crates/goose/src/session/chat_history_search.rs
+152 -5
View File
@@ -2,6 +2,7 @@ use crate::conversation::message::MessageContent;
use crate::session::session_manager::SessionType;
use anyhow::Result;
use chrono::{DateTime, Utc};
use rmcp::model::Role;
use serde::Serialize;
use sqlx::{Pool, Sqlite};
use std::collections::HashMap;
@@ -142,9 +143,24 @@ impl<'a> ChatHistorySearch<'a> {
m.timestamp
FROM messages m
INNER JOIN sessions s ON m.session_id = s.id
WHERE EXISTS (
SELECT 1 FROM json_each(m.content_json)
WHERE json_extract(value, '$.type') = 'text'
WHERE COALESCE(
CASE
WHEN json_valid(m.metadata_json)
THEN json_extract(m.metadata_json, '$.agentVisible')
END,
1
) = 1
AND EXISTS (
SELECT 1 FROM json_each(m.content_json) AS content
WHERE json_extract(content.value, '$.type') = 'text'
AND (
json_type(content.value, '$.annotations.audience') IS NULL
OR EXISTS (
SELECT 1
FROM json_each(content.value, '$.annotations.audience') AS audience
WHERE audience.value = 'assistant'
)
)
AND (
"#,
);
@@ -153,7 +169,7 @@ impl<'a> ChatHistorySearch<'a> {
if i > 0 {
sql.push_str(" OR ");
}
sql.push_str("LOWER(json_extract(value, '$.text')) LIKE ?");
sql.push_str("LOWER(json_extract(content.value, '$.text')) LIKE ?");
}
sql.push_str(
@@ -203,7 +219,11 @@ impl<'a> ChatHistorySearch<'a> {
) in rows
{
if let Ok(content_vec) = serde_json::from_str::<Vec<MessageContent>>(&content_json) {
let text_parts = Self::extract_text_content(content_vec);
let agent_visible_content = content_vec
.into_iter()
.filter_map(|content| content.filter_for_audience(Role::Assistant))
.collect();
let text_parts = Self::extract_text_content(agent_visible_content);
if !text_parts.is_empty() {
let entry = session_messages.entry(session_id.clone()).or_insert((
@@ -302,3 +322,130 @@ impl<'a> ChatHistorySearch<'a> {
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::conversation::message::{Message, MessageContent, MessageMetadata};
use rmcp::model::{AnnotateAble, RawTextContent};
use sqlx::sqlite::SqlitePoolOptions;
fn user_only_text(text: &str) -> MessageContent {
MessageContent::Text(
RawTextContent {
text: text.to_string(),
meta: None,
}
.no_annotation()
.with_audience(vec![Role::User]),
)
}
async fn insert_message(pool: &Pool<Sqlite>, message: &Message, timestamp: DateTime<Utc>) {
sqlx::query(
r#"
INSERT INTO messages (session_id, role, content_json, timestamp, metadata_json)
VALUES ('session-1', ?, ?, ?, ?)
"#,
)
.bind(match message.role {
Role::User => "user",
Role::Assistant => "assistant",
})
.bind(serde_json::to_string(&message.content).unwrap())
.bind(timestamp)
.bind(serde_json::to_string(&message.metadata).unwrap())
.execute(pool)
.await
.unwrap();
}
#[tokio::test]
async fn search_projects_audience_before_matching_and_limiting() {
let pool = SqlitePoolOptions::new()
.max_connections(1)
.connect("sqlite::memory:")
.await
.unwrap();
sqlx::query(
r#"
CREATE TABLE sessions (
id TEXT PRIMARY KEY,
description TEXT NOT NULL,
working_dir TEXT NOT NULL,
created_at TIMESTAMP NOT NULL,
session_type TEXT NOT NULL
);
CREATE TABLE messages (
id INTEGER PRIMARY KEY AUTOINCREMENT,
session_id TEXT NOT NULL,
role TEXT NOT NULL,
content_json TEXT NOT NULL,
timestamp TIMESTAMP NOT NULL,
metadata_json TEXT
);
"#,
)
.execute(&pool)
.await
.unwrap();
sqlx::query(
"INSERT INTO sessions (id, description, working_dir, created_at, session_type) VALUES ('session-1', 'test', '/tmp', ?, 'user')",
)
.bind(Utc::now())
.execute(&pool)
.await
.unwrap();
let now = Utc::now();
insert_message(
&pool,
&Message::user().with_text("needle public"),
now - chrono::Duration::seconds(3),
)
.await;
insert_message(
&pool,
&Message::user()
.with_text("haystack visible")
.with_content(user_only_text("needle secret-only")),
now - chrono::Duration::seconds(2),
)
.await;
insert_message(
&pool,
&Message::user()
.with_text("needle hidden row")
.with_metadata(MessageMetadata::user_only()),
now - chrono::Duration::seconds(1),
)
.await;
let needle = ChatHistorySearch::new(&pool, "needle", Some(1), None, None, None, vec![])
.execute()
.await
.unwrap();
assert_eq!(needle.total_matches, 1);
assert_eq!(needle.results[0].messages[0].content, "needle public");
let haystack =
ChatHistorySearch::new(&pool, "haystack", Some(10), None, None, None, vec![])
.execute()
.await
.unwrap();
assert_eq!(haystack.total_matches, 1);
assert!(haystack.results[0].messages[0]
.content
.contains("haystack visible"));
assert!(!haystack.results[0].messages[0]
.content
.contains("needle secret-only"));
let hidden_only =
ChatHistorySearch::new(&pool, "secret-only", Some(10), None, None, None, vec![])
.execute()
.await
.unwrap();
assert_eq!(hidden_only.total_matches, 0);
}
}