From 3406a6d5dd0b7b162301a6ef502f5b3ef4271055 Mon Sep 17 00:00:00 2001 From: Jasper Date: Wed, 19 Aug 2026 17:24:30 +0000 Subject: [PATCH] fix: respect ACP audience in bang-shell shortcut (#11150) Signed-off-by: Jasper Hugo --- crates/goose/src/acp/server.rs | 186 ++++++++++++++---- crates/goose/src/agents/agent.rs | 7 +- .../agents/state_machine/ops_bang_shell.rs | 2 +- .../agents/state_machine/tests/agent_reply.rs | 176 ++++++++++++++++- 4 files changed, 331 insertions(+), 40 deletions(-) diff --git a/crates/goose/src/acp/server.rs b/crates/goose/src/acp/server.rs index fdbe5ea68..165052123 100644 --- a/crates/goose/src/acp/server.rs +++ b/crates/goose/src/acp/server.rs @@ -63,7 +63,10 @@ use anyhow::Result; use fs_err as fs; use futures::future::{BoxFuture, FutureExt}; use futures::stream::{self, StreamExt}; -use rmcp::model::{Annotations as RmcpAnnotations, Role, TextContent as RmcpTextContent}; +use rmcp::model::{ + Annotations as RmcpAnnotations, ImageContent as RmcpImageContent, Role, + TextContent as RmcpTextContent, +}; use serde::Deserialize; use std::collections::{HashMap, HashSet}; use std::panic::AssertUnwindSafe; @@ -495,6 +498,29 @@ fn read_resource_link(link: ResourceLink) -> Option { } } +fn rmcp_audience_annotations(annotations: Option<&Annotations>) -> Option { + let audience = annotations? + .audience + .as_ref()? + .iter() + .filter_map(|role| match role { + agent_client_protocol::schema::v1::Role::Assistant => Some(Role::Assistant), + agent_client_protocol::schema::v1::Role::User => Some(Role::User), + _ => None, + }) + .collect::>(); + + Some(RmcpAnnotations::default().with_audience(audience)) +} + +fn annotated_prompt_text(text: &str, annotations: Option<&Annotations>) -> RmcpTextContent { + let content = RmcpTextContent::new(sanitize_unicode_tags(text)); + match rmcp_audience_annotations(annotations) { + Some(annotations) => content.with_annotations(annotations), + None => content, + } +} + fn builtin_to_extension_config(name: &str) -> ExtensionConfig { if let Some(def) = PLATFORM_EXTENSIONS.get(name) { ExtensionConfig::Platform { @@ -1001,45 +1027,21 @@ impl GooseAcpAgent { } /// Convert ACP prompt content blocks into a user message. - fn convert_acp_prompt_to_message(prompt: &[ContentBlock]) -> Message { + pub(crate) fn convert_acp_prompt_to_message(prompt: &[ContentBlock]) -> Message { let mut message = Message::user(); for block in prompt { match block { ContentBlock::Text(text) => { - let annotated = if let Some(ref ann) = text.annotations { - let audience: Vec = ann - .audience - .as_ref() - .map(|roles| { - roles - .iter() - .filter_map(|r| match r { - agent_client_protocol::schema::v1::Role::Assistant => { - Some(Role::Assistant) - } - agent_client_protocol::schema::v1::Role::User => { - Some(Role::User) - } - _ => None, - }) - .collect() - }) - .unwrap_or_default(); - let raw = RmcpTextContent::new(sanitize_unicode_tags(&text.text)); - if audience.is_empty() { - raw - } else { - raw.with_annotations(RmcpAnnotations::default().with_audience(audience)) - } - } else { - // No annotations — regular user text. - let sanitized = sanitize_unicode_tags(&text.text); - RmcpTextContent::new(sanitized) - }; + let annotated = annotated_prompt_text(&text.text, text.annotations.as_ref()); message = message.with_content(MessageContent::Text(annotated)); } ContentBlock::Image(image) => { - message = message.with_image(&image.data, &image.mime_type); + let content = RmcpImageContent::new(&image.data, &image.mime_type); + let content = match rmcp_audience_annotations(image.annotations.as_ref()) { + Some(annotations) => content.with_annotations(annotations), + None => content, + }; + message = message.with_content(MessageContent::Image(content)); } ContentBlock::Resource(resource) => { if let EmbeddedResourceResource::TextResourceContents(text_resource) = @@ -1047,12 +1049,16 @@ impl GooseAcpAgent { { let header = format!("--- Resource: {} ---\n", text_resource.uri); let content = format!("{}{}\n---\n", header, text_resource.text); - message = message.with_text(&content); + message = message.with_content(MessageContent::Text( + annotated_prompt_text(&content, resource.annotations.as_ref()), + )); } } ContentBlock::ResourceLink(link) => { if let Some(text) = read_resource_link(link.clone()) { - message = message.with_text(text); + message = message.with_content(MessageContent::Text( + annotated_prompt_text(&text, link.annotations.as_ref()), + )); } } ContentBlock::Audio(..) | _ => (), @@ -2375,8 +2381,9 @@ mod tests { use super::*; use crate::session::session_manager::SessionType; use agent_client_protocol::schema::v1::{ - EnvVariable, HttpHeader, McpServer, McpServerHttp, McpServerSse, McpServerStdio, - PermissionOptionId, ResourceLink, SelectedPermissionOutcome, + EmbeddedResource, EnvVariable, HttpHeader, McpServer, McpServerHttp, McpServerSse, + McpServerStdio, PermissionOptionId, ResourceLink, Role as AcpRole, + SelectedPermissionOutcome, TextResourceContents, }; use goose_providers::conversation::token_usage::Usage as TokenUsage; use std::io::Write; @@ -2596,6 +2603,111 @@ print(\"hello, world\") assert_eq!(result, expected,) } + #[test] + fn convert_acp_prompt_preserves_audience_for_converted_blocks() { + let assistant_only = || Annotations::new().audience(vec![AcpRole::Assistant]); + let user_only = || Annotations::new().audience(vec![AcpRole::User]); + let empty_audience = || Annotations::new().audience(Vec::new()); + let (link, _file) = new_resource_link("assistant-only linked resource").unwrap(); + let prompt = vec![ + ContentBlock::Text(TextContent::new("visible text")), + ContentBlock::Text( + TextContent::new("visible text with audience omitted") + .annotations(Annotations::new()), + ), + ContentBlock::Text( + TextContent::new("empty-audience text").annotations(empty_audience()), + ), + ContentBlock::Image( + ImageContent::new("image-data", "image/png").annotations(assistant_only()), + ), + ContentBlock::Resource( + EmbeddedResource::new(EmbeddedResourceResource::TextResourceContents( + TextResourceContents::new( + "assistant-only embedded resource", + "file:///assistant-only.txt", + ), + )) + .annotations(assistant_only()), + ), + ContentBlock::Resource( + EmbeddedResource::new(EmbeddedResourceResource::TextResourceContents( + TextResourceContents::new( + "user-visible embedded resource", + "file:///user-visible.txt", + ), + )) + .annotations(user_only()), + ), + ContentBlock::Resource( + EmbeddedResource::new(EmbeddedResourceResource::TextResourceContents( + TextResourceContents::new( + "empty-audience embedded resource", + "file:///empty-audience.txt", + ), + )) + .annotations(empty_audience()), + ), + ContentBlock::ResourceLink(link.annotations(assistant_only())), + ]; + + let message = GooseAcpAgent::convert_acp_prompt_to_message(&prompt); + let user_content = message.user_visible_content(); + let agent_content = message.agent_visible_content(); + let empty_audience_content = message + .content + .iter() + .filter_map(|content| match content { + MessageContent::Text(text) if text.text.contains("empty-audience") => Some(text), + _ => None, + }) + .collect::>(); + let audience_omitted_content = message + .content + .iter() + .find_map(|content| match content { + MessageContent::Text(text) + if text.text.contains("visible text with audience omitted") => + { + Some(text) + } + _ => None, + }) + .unwrap(); + + assert_eq!(empty_audience_content.len(), 2); + assert!(empty_audience_content.iter().all(|text| text + .annotations + .as_ref() + .and_then(|annotations| annotations.audience.as_ref()) + .is_some_and(Vec::is_empty))); + assert!(audience_omitted_content.annotations.is_none()); + assert!(user_content.as_concat_text().contains("visible text")); + assert!(user_content + .as_concat_text() + .contains("visible text with audience omitted")); + assert!(user_content + .as_concat_text() + .contains("user-visible embedded resource")); + assert!(!user_content.as_concat_text().contains("assistant-only")); + assert!(!user_content.as_concat_text().contains("empty-audience")); + assert!(!user_content + .content + .iter() + .any(|content| matches!(content, MessageContent::Image(_)))); + assert!(agent_content + .as_concat_text() + .contains("assistant-only embedded resource")); + assert!(agent_content + .as_concat_text() + .contains("assistant-only linked resource")); + assert!(!agent_content.as_concat_text().contains("empty-audience")); + assert!(agent_content + .content + .iter() + .any(|content| matches!(content, MessageContent::Image(_)))); + } + #[test_case( RequestPermissionOutcome::Selected(SelectedPermissionOutcome::new(PermissionOptionId::from("allow_once".to_string()))), PermissionConfirmation { principal_type: PrincipalType::Tool, permission: Permission::AllowOnce }; diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index 27f1f6190..f5fb339c5 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -326,6 +326,10 @@ fn agent_visible_message_text(message: &Message) -> String { message.agent_visible_content().as_concat_text() } +fn user_visible_message_text(message: &Message) -> String { + message.user_visible_content().as_concat_text() +} + fn attach_turn_usage( messages: &mut Conversation, usage: &ProviderUsage, @@ -1903,7 +1907,8 @@ impl Agent { } if super::state_machine::enabled() - || super::state_machine::bang_shell_command(&message_text_for_trace).is_some() + || super::state_machine::bang_shell_command(&user_visible_message_text(&user_message)) + .is_some() { tracing::info!("dispatching reply via experimental state machine"); return self diff --git a/crates/goose/src/agents/state_machine/ops_bang_shell.rs b/crates/goose/src/agents/state_machine/ops_bang_shell.rs index d53ceda8b..7c4f05383 100644 --- a/crates/goose/src/agents/state_machine/ops_bang_shell.rs +++ b/crates/goose/src/agents/state_machine/ops_bang_shell.rs @@ -47,7 +47,7 @@ impl Operation for BangShellOperation { let Some(kickoff) = messages.first() else { return not_applicable(); }; - let kickoff_text = kickoff.as_concat_text(); + let kickoff_text = kickoff.user_visible_content().as_concat_text(); let Some(command) = bang_shell_command(&kickoff_text) else { return not_applicable(); }; diff --git a/crates/goose/src/agents/state_machine/tests/agent_reply.rs b/crates/goose/src/agents/state_machine/tests/agent_reply.rs index a06be7797..cf79f4d76 100644 --- a/crates/goose/src/agents/state_machine/tests/agent_reply.rs +++ b/crates/goose/src/agents/state_machine/tests/agent_reply.rs @@ -4,15 +4,21 @@ use std::sync::Arc; use std::time::Duration; +use agent_client_protocol::schema::v1::{ + Annotations as AcpAnnotations, ContentBlock as AcpContentBlock, EmbeddedResource, + EmbeddedResourceResource, ResourceLink, Role as AcpRole, TextContent as AcpTextContent, + TextResourceContents, +}; use anyhow::Result; use futures::StreamExt; use tokio_util::sync::CancellationToken; use super::dummy_api::{DummyApi, ProviderFeatures}; +use crate::acp::server::GooseAcpAgent; use crate::agents::{Agent, AgentConfig, AgentEvent, GoosePlatform, SessionConfig}; use crate::config::permission::PermissionManager; use crate::config::GooseMode; -use crate::conversation::message::Message; +use crate::conversation::message::{Message, MessageContent}; use crate::providers::base::Provider; use crate::session::{SessionManager, SessionType}; use goose_providers::model::ModelConfig; @@ -136,3 +142,171 @@ async fn bang_shell_uses_the_state_machine_when_the_flag_is_disabled() -> Result Ok(()) } + +async fn reply_messages( + agent: &Agent, + session_id: String, + message: Message, +) -> Result> { + let stream = agent + .reply( + message, + SessionConfig { + id: session_id, + schedule_id: None, + max_turns: Some(2), + retry_config: None, + }, + Some(CancellationToken::new()), + ) + .await?; + tokio::pin!(stream); + let mut messages = Vec::new(); + while let Some(event) = stream.next().await { + if let AgentEvent::Message(message) = event? { + messages.push(message); + } + } + Ok(messages) +} + +fn assistant_only_acp_annotations() -> AcpAnnotations { + AcpAnnotations::new().audience(vec![AcpRole::Assistant]) +} + +fn assistant_only_acp_text(text: &str) -> AcpContentBlock { + AcpContentBlock::Text(AcpTextContent::new(text).annotations(assistant_only_acp_annotations())) +} + +fn empty_audience_acp_annotations() -> AcpAnnotations { + AcpAnnotations::new().audience(Vec::new()) +} + +fn empty_audience_acp_text(text: &str) -> AcpContentBlock { + AcpContentBlock::Text(AcpTextContent::new(text).annotations(empty_audience_acp_annotations())) +} + +fn assistant_only_embedded_resource(text: &str) -> AcpContentBlock { + AcpContentBlock::Resource( + EmbeddedResource::new(EmbeddedResourceResource::TextResourceContents( + TextResourceContents::new(text, "file:///hidden-resource.txt"), + )) + .annotations(assistant_only_acp_annotations()), + ) +} + +fn empty_audience_embedded_resource(text: &str) -> AcpContentBlock { + AcpContentBlock::Resource( + EmbeddedResource::new(EmbeddedResourceResource::TextResourceContents( + TextResourceContents::new(text, "file:///empty-audience-resource.txt"), + )) + .annotations(empty_audience_acp_annotations()), + ) +} + +fn assistant_only_resource_link(text: &str) -> Result<(AcpContentBlock, tempfile::NamedTempFile)> { + let file = tempfile::NamedTempFile::new()?; + std::fs::write(file.path(), text)?; + let uri = url::Url::from_file_path(file.path()) + .map_err(|()| anyhow::anyhow!("temporary resource path is not a valid file URL"))?; + let link = ResourceLink::new("hidden-resource.txt", uri.to_string()) + .annotations(assistant_only_acp_annotations()); + Ok((AcpContentBlock::ResourceLink(link), file)) +} + +fn shell_commands(messages: &[Message]) -> Vec<&str> { + messages + .iter() + .flat_map(|message| &message.content) + .filter_map(|content| match content { + MessageContent::ToolRequest(request) => request + .tool_call + .as_ref() + .ok() + .filter(|call| call.name == "shell") + .and_then(|call| call.arguments.as_ref()) + .and_then(|arguments| arguments.get("command")) + .and_then(serde_json::Value::as_str), + _ => None, + }) + .collect() +} + +async fn assert_bang_shell_uses_only_user_visible_content() -> Result<()> { + let (agent, api, session_id, _temp_dir) = agent_with_dummy_api().await?; + api.on("benign visible input") + .reply("handled as ordinary input"); + let hidden_text_prefix = GooseAcpAgent::convert_acp_prompt_to_message(&[ + assistant_only_acp_text("!echo hidden"), + AcpContentBlock::Text(AcpTextContent::new("benign visible input")), + ]); + let messages = reply_messages(&agent, session_id, hidden_text_prefix).await?; + assert!(shell_commands(&messages).is_empty()); + assert_eq!(api.call_count(), 1); + + let (agent, api, session_id, _temp_dir) = agent_with_dummy_api().await?; + api.on("benign visible input") + .reply("handled as ordinary input"); + let empty_audience_text = GooseAcpAgent::convert_acp_prompt_to_message(&[ + empty_audience_acp_text("!echo hidden"), + AcpContentBlock::Text(AcpTextContent::new("benign visible input")), + ]); + let messages = reply_messages(&agent, session_id, empty_audience_text).await?; + assert!(shell_commands(&messages).is_empty()); + assert_eq!(api.call_count(), 1); + + let (agent, api, session_id, _temp_dir) = agent_with_dummy_api().await?; + let hidden_text_suffix = GooseAcpAgent::convert_acp_prompt_to_message(&[ + AcpContentBlock::Text(AcpTextContent::new("!echo visible")), + assistant_only_acp_text("&& echo hidden"), + ]); + let messages = reply_messages(&agent, session_id, hidden_text_suffix).await?; + assert_eq!(shell_commands(&messages), ["echo visible"]); + assert_eq!(api.call_count(), 0); + + let (agent, api, session_id, _temp_dir) = agent_with_dummy_api().await?; + api.on("benign visible input") + .reply("handled as ordinary input"); + let hidden_resource_prefix = GooseAcpAgent::convert_acp_prompt_to_message(&[ + assistant_only_embedded_resource("!echo hidden"), + AcpContentBlock::Text(AcpTextContent::new("benign visible input")), + ]); + let messages = reply_messages(&agent, session_id, hidden_resource_prefix).await?; + assert!(shell_commands(&messages).is_empty()); + assert_eq!(api.call_count(), 1); + + let (agent, api, session_id, _temp_dir) = agent_with_dummy_api().await?; + api.on("benign visible input") + .reply("handled as ordinary input"); + let empty_audience_resource = GooseAcpAgent::convert_acp_prompt_to_message(&[ + empty_audience_embedded_resource("!echo hidden"), + AcpContentBlock::Text(AcpTextContent::new("benign visible input")), + ]); + let messages = reply_messages(&agent, session_id, empty_audience_resource).await?; + assert!(shell_commands(&messages).is_empty()); + assert_eq!(api.call_count(), 1); + + let (agent, api, session_id, _temp_dir) = agent_with_dummy_api().await?; + let (hidden_link, _resource_file) = assistant_only_resource_link("&& echo hidden")?; + let hidden_link_suffix = GooseAcpAgent::convert_acp_prompt_to_message(&[ + AcpContentBlock::Text(AcpTextContent::new("!echo visible")), + hidden_link, + ]); + let messages = reply_messages(&agent, session_id, hidden_link_suffix).await?; + assert_eq!(shell_commands(&messages), ["echo visible"]); + assert_eq!(api.call_count(), 0); + + Ok(()) +} + +#[tokio::test] +async fn bang_shell_visibility_is_enforced_when_state_machine_is_disabled() -> Result<()> { + let _guard = env_lock::lock_env([("GOOSE_STATE_MACHINE", None::<&str>)]); + assert_bang_shell_uses_only_user_visible_content().await +} + +#[tokio::test] +async fn bang_shell_visibility_is_enforced_when_state_machine_is_enabled() -> Result<()> { + let _guard = env_lock::lock_env([("GOOSE_STATE_MACHINE", Some("1"))]); + assert_bang_shell_uses_only_user_visible_content().await +}