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

fix: pass complete response to stop hooks (#11366)

Browse Source 
Signed-off-by: Jasper Hugo <jasper@spiral.xyz>
This commit is contained in:
Jasper
2026-08-20 10:50:45 +00:00
committed by GitHub
parent df899654fe
commit ea07f7bb41
4 changed files with 287 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/agents/agent.rs
+70 -2
View File
@@ -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<MessageStream, ProviderError> {
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();
crates/goose/src/agents/state_machine/ops_stop_hook.rs
+15 -4
View File
@@ -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::<Vec<_>>()
.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<Session, GooseEffect> 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)
crates/goose/src/agents/state_machine/tests/dummy_api.rs
+54
View File
@@ -31,6 +31,7 @@ impl Default for ProviderFeatures {
#[derive(Clone)]
enum ApiResponse {
Reply(String),
ReplyWithDistinctIds(Vec<String>),
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<const N: usize>(
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<String>) -> 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(
crates/goose/src/agents/state_machine/tests/hooks_lifecycle.rs
+148 -2
View File
@@ -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<Value> {
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::<std::collections::HashSet<_>>();
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::<std::collections::HashSet<_>>();
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);