diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index f5fb339c5..950e86305 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -2576,7 +2576,13 @@ impl Agent { let num_tool_requests = frontend_requests.len() + remaining_requests.len(); if num_tool_requests == 0 { - let text = filtered_response.as_concat_text(); + let text = if response.is_user_visible() { + filtered_response + .user_visible_content() + .as_concat_text() + } else { + String::new() + }; if !text.is_empty() { last_assistant_text.push_str(&text); } @@ -3989,7 +3995,7 @@ mod tests { use crate::recipe::Response; use crate::session::session_manager::SessionType; use goose_providers::conversation::token_usage::{ProviderUsage, Usage}; - use rmcp::model::Tool; + use rmcp::model::{Annotations, Role, TextContent, Tool}; use std::path::PathBuf; use std::sync::atomic::{AtomicUsize, Ordering}; use tempfile::TempDir; @@ -4566,6 +4572,48 @@ echo start >> "$PLUGIN_ROOT/hook.log" } } + struct VisibilityTextProvider; + + #[async_trait::async_trait] + impl crate::providers::base::Provider for VisibilityTextProvider { + async fn stream( + &self, + _model_config: &goose_providers::model::ModelConfig, + _system_prompt: &str, + _messages: &[Message], + _tools: &[Tool], + ) -> Result { + let usage = ProviderUsage::new("mock-model".to_string(), Usage::default()); + let mixed_audience = Message::assistant() + .with_content(MessageContent::Text( + TextContent::new("assistant-only block ").with_annotations( + Annotations::default().with_audience(vec![Role::Assistant]), + ), + )) + .with_content(MessageContent::Text( + TextContent::new("visible last") + .with_annotations(Annotations::default().with_audience(vec![Role::User])), + )); + + Ok(Box::pin(futures::stream::iter(vec![ + Ok((Some(Message::assistant().with_text("visible first ")), None)), + Ok(( + Some( + Message::assistant() + .with_text("internal message ") + .agent_only(), + ), + None, + )), + Ok((Some(mixed_audience), Some(usage))), + ]))) + } + + fn get_name(&self) -> &str { + "visibility-text" + } + } + struct OutputLimitMarkerProvider { include_content: bool, call_count: AtomicUsize, @@ -5063,6 +5111,26 @@ echo start >> "$PLUGIN_ROOT/hook.log" Ok(()) } + #[tokio::test] + async fn stop_hook_payload_excludes_non_user_visible_assistant_content() -> Result<()> { + let env = StopHookTestEnv::new(RECORD_PAYLOAD_SCRIPT)?; + let provider = Arc::new(VisibilityTextProvider); + let (agent, session_id) = + create_test_agent(env.data_dir(), env.hook_manager(), provider).await?; + + run_stop_hook_test_turn(&agent, &session_id, "hello").await?; + + let payload = env.stop_payload()?; + assert_eq!( + payload + .get("last_assistant_message") + .and_then(Value::as_str), + Some("visible first visible last") + ); + + Ok(()) + } + #[tokio::test] async fn test_add_final_output_tool() -> Result<()> { let agent = Agent::new(); diff --git a/crates/goose/src/agents/state_machine/ops_stop_hook.rs b/crates/goose/src/agents/state_machine/ops_stop_hook.rs index 86dfa5977..3ee99c866 100644 --- a/crates/goose/src/agents/state_machine/ops_stop_hook.rs +++ b/crates/goose/src/agents/state_machine/ops_stop_hook.rs @@ -2,6 +2,7 @@ use anyhow::Result; use async_trait::async_trait; +use rmcp::model::Role; use crate::agents::state_machine::effects::GooseEffect; use crate::agents::state_machine::{ @@ -55,6 +56,19 @@ impl StopHookOperation { block_cap, } } + + fn trailing_assistant_text(messages: &[Message]) -> String { + messages + .iter() + .rev() + .take_while(|message| message.role == Role::Assistant) + .collect::>() + .into_iter() + .rev() + .filter(|message| message.is_user_visible()) + .map(|message| message.user_visible_content().as_concat_text()) + .collect() + } } #[async_trait] @@ -73,10 +87,7 @@ impl Operation for StopHookOperation { if !ends_turn(messages) { return not_applicable(); } - let last_assistant_text = conversation - .last() - .map(Message::as_concat_text) - .unwrap_or_default(); + let last_assistant_text = Self::trailing_assistant_text(messages); let context = HookContext::new(HookEvent::Stop, &session.id) .with_last_assistant_message(last_assistant_text) diff --git a/crates/goose/src/agents/state_machine/tests/dummy_api.rs b/crates/goose/src/agents/state_machine/tests/dummy_api.rs index a9b48d6d3..4291230ee 100644 --- a/crates/goose/src/agents/state_machine/tests/dummy_api.rs +++ b/crates/goose/src/agents/state_machine/tests/dummy_api.rs @@ -31,6 +31,7 @@ impl Default for ProviderFeatures { #[derive(Clone)] enum ApiResponse { Reply(String), + ReplyWithDistinctIds(Vec), ToolCall { name: String, arguments: String, @@ -262,6 +263,15 @@ impl<'a> ApiRuleBuilder<'a> { self.configured(ApiResponse::Reply(text.into())) } + pub(super) fn reply_with_distinct_ids( + self, + chunks: [&str; N], + ) -> ConfiguredResponse<'a> { + self.configured(ApiResponse::ReplyWithDistinctIds( + chunks.into_iter().map(str::to_string).collect(), + )) + } + pub(super) fn hold_reply(self, text: impl Into) -> ResponseGate { let gate = ResponseGate::new(); self.api @@ -470,6 +480,13 @@ impl DummyApiState { &text, None, )), + ApiResponse::ReplyWithDistinctIds(chunks) => { + let output_tokens: usize = chunks.iter().map(|chunk| chunk.chars().count()).sum(); + sse_response(reply_events_with_distinct_ids( + &meta(output_tokens as i32), + &chunks, + )) + } ApiResponse::ToolCall { name, arguments, @@ -662,6 +679,43 @@ fn reply_events(meta: &ResponseMeta, text: &str, error: Option<&str>) -> String events } +fn reply_events_with_distinct_ids(meta: &ResponseMeta, chunks: &[String]) -> String { + let mut events = String::new(); + for (index, chunk) in chunks.iter().enumerate() { + push_event( + &mut events, + json!({ + "id": format!("{}-{index}", meta.id), + "object": "chat.completion.chunk", + "model": meta.model, + "choices": [{ + "index": 0, + "delta": { "content": chunk }, + "finish_reason": null + }] + }), + ); + } + push_event( + &mut events, + json!({ + "id": format!("{}-{}", meta.id, chunks.len().saturating_sub(1)), + "object": "chat.completion.chunk", + "model": meta.model, + "choices": [{ + "index": 0, + "delta": {}, + "finish_reason": "stop" + }] + }), + ); + if meta.include_usage { + push_event(&mut events, usage_event(meta)); + } + events.push_str("data: [DONE]\n\n"); + events +} + fn no_choices_events(id: &str, model: &str) -> String { let mut events = String::new(); push_event( diff --git a/crates/goose/src/agents/state_machine/tests/hooks_lifecycle.rs b/crates/goose/src/agents/state_machine/tests/hooks_lifecycle.rs index 53f2c7cfe..fcb944085 100644 --- a/crates/goose/src/agents/state_machine/tests/hooks_lifecycle.rs +++ b/crates/goose/src/agents/state_machine/tests/hooks_lifecycle.rs @@ -1,9 +1,11 @@ use anyhow::Result; +use rmcp::model::{Annotations, Role, TextContent}; +use serde_json::Value; use super::calculator_extension::{value, ADD}; use super::pipeline::{test_pipeline, MessageKind::Agent, MessageKind::ToolResponse, MAX_TURNS}; use crate::agents::state_machine::ops_stop_hook::DENIED; -use crate::conversation::message::{MessageContent, SystemNotificationType}; +use crate::conversation::message::{Message, MessageContent, SystemNotificationType}; struct HookTestEnv { _temp_dir: tempfile::TempDir, @@ -44,7 +46,6 @@ impl HookTestEnv { .lines() .count() } - fn last_context(&self) -> serde_json::Value { serde_json::from_str( &std::fs::read_to_string(self.plugin_dir.join("context.json")) @@ -52,12 +53,30 @@ impl HookTestEnv { ) .expect("hook context is valid JSON") } + + fn payloads(&self) -> Vec { + std::fs::read_to_string(self.plugin_dir.join("hook.log")) + .unwrap_or_default() + .lines() + .map(|line| serde_json::from_str(line).expect("hook payload should be JSON")) + .collect() + } } const LOG_AND_ALLOW_SCRIPT: &str = "#!/bin/sh\necho ran >> \"$PLUGIN_ROOT/hook.log\"\nexit 0\n"; const LOG_AND_BLOCK_SCRIPT: &str = "#!/bin/sh\necho blocked >> \"$PLUGIN_ROOT/hook.log\"\necho \"not done yet\" >&2\nexit 2\n"; const LOG_CONTEXT_AND_BLOCK_SCRIPT: &str = "#!/bin/sh\ncat > \"$PLUGIN_ROOT/context.json\"\necho blocked >> \"$PLUGIN_ROOT/hook.log\"\necho \"not done yet\" >&2\nexit 2\n"; +const RECORD_AND_ALLOW_SCRIPT: &str = + "#!/bin/sh\npayload=$(cat)\nprintf '%s\\n' \"$payload\" >> \"$PLUGIN_ROOT/hook.log\"\nexit 0\n"; +const RECORD_AND_BLOCK_MARKER_SCRIPT: &str = "#!/bin/sh +payload=$(cat) +printf '%s\\n' \"$payload\" >> \"$PLUGIN_ROOT/hook.log\" +case \"$payload\" in + *\"policy marker\"*) echo \"policy marker found\" >&2; exit 2 ;; +esac +exit 0 +"; #[tokio::test] async fn stop_hooks_allow_block_and_skip_non_stop_exits() -> Result<()> { @@ -131,6 +150,133 @@ async fn stop_hooks_allow_block_and_skip_non_stop_exits() -> Result<()> { Ok(()) } +#[tokio::test] +async fn stop_hook_receives_complete_user_visible_assistant_response() -> Result<()> { + let same_id = HookTestEnv::new("Stop", RECORD_AND_ALLOW_SCRIPT); + let (pipeline, api) = test_pipeline().await?; + let pipeline = pipeline.with_hook_manager(same_id.hook_manager()); + api.on("same id").reply("one two three four five"); + + let result = pipeline.run(["same id"]).await?; + let response_ids = result + .conversation() + .messages() + .iter() + .filter(|message| message.role == Role::Assistant && !message.as_concat_text().is_empty()) + .filter_map(|message| message.id.as_deref()) + .collect::>(); + assert_eq!(response_ids.len(), 1); + assert_eq!( + same_id.payloads()[0]["last_assistant_message"], + "one two three four five" + ); + + let distinct_ids = HookTestEnv::new("Stop", RECORD_AND_ALLOW_SCRIPT); + let (pipeline, api) = test_pipeline().await?; + let pipeline = pipeline.with_hook_manager(distinct_ids.hook_manager()); + api.on("distinct ids") + .reply_with_distinct_ids(["one ", "two ", "three"]); + + let result = pipeline.run(["distinct ids"]).await?; + let response_ids = result + .conversation() + .messages() + .iter() + .filter(|message| message.role == Role::Assistant && !message.as_concat_text().is_empty()) + .filter_map(|message| message.id.as_deref()) + .collect::>(); + assert_eq!(response_ids.len(), 3); + assert_eq!( + distinct_ids.payloads()[0]["last_assistant_message"], + "one two three" + ); + + let visibility = HookTestEnv::new("Stop", RECORD_AND_ALLOW_SCRIPT); + let (pipeline, _) = test_pipeline().await?; + let pipeline = pipeline.with_hook_manager(visibility.hook_manager()); + pipeline + .seed([ + Message::user().with_text("visibility"), + Message::assistant() + .with_text("visible first ") + .with_id("visible-first"), + Message::assistant() + .with_text("internal message ") + .agent_only() + .with_id("internal"), + Message::assistant() + .with_content(MessageContent::Text( + TextContent::new("assistant-only block ").with_annotations( + Annotations::default().with_audience(vec![Role::Assistant]), + ), + )) + .with_content(MessageContent::Text( + TextContent::new("visible last") + .with_annotations(Annotations::default().with_audience(vec![Role::User])), + )) + .with_id("visible-last"), + ]) + .await?; + + pipeline.resume().await?; + assert_eq!( + visibility.payloads()[0]["last_assistant_message"], + "visible first visible last" + ); + + Ok(()) +} + +#[tokio::test] +async fn stop_hook_distinct_id_denials_retry_once_then_respect_block_cap() -> Result<()> { + let blocked = HookTestEnv::new("Stop", RECORD_AND_BLOCK_MARKER_SCRIPT); + let (pipeline, api) = test_pipeline().await?; + let pipeline = pipeline + .with_hook_manager(blocked.hook_manager()) + .with_stop_hook_block_cap(1); + api.on("hello") + .reply_with_distinct_ids(["policy marker: initial; ", "tail"]); + api.on("blocked ending this turn") + .reply_with_distinct_ids(["policy marker: retry; ", "tail"]); + + let (_, result, _) = pipeline.run_reconstructing_each_step("hello").await?; + assert_eq!(api.call_count(), 2); + assert_eq!(blocked.invocations(), 2); + let payloads = blocked.payloads(); + assert_eq!( + payloads[0]["last_assistant_message"], + "policy marker: initial; tail" + ); + assert_eq!( + payloads[1]["last_assistant_message"], + "policy marker: retry; tail" + ); + assert_eq!( + result + .conversation() + .messages() + .iter() + .filter(|message| message + .metadata + .operation_note("stop_hook", DENIED) + .is_some()) + .count(), + 1 + ); + assert!(result.conversation().last().is_some_and(|message| { + message.content.iter().any(|content| { + matches!( + content, + MessageContent::SystemNotification(notification) + if notification.notification_type == SystemNotificationType::InlineMessage + && notification.msg.contains("GOOSE_STOP_HOOK_BLOCK_CAP") + ) + }) + })); + + Ok(()) +} + #[tokio::test] async fn session_prompt_and_tool_hooks_fire_at_their_boundaries() -> Result<()> { let session_start = HookTestEnv::new("SessionStart", LOG_AND_ALLOW_SCRIPT);