diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index 814d351d64cc..4b53c82c8a4d 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -1308,6 +1308,7 @@ impl Agent { let message = Message::user().with_text(FINAL_OUTPUT_CONTINUATION_MESSAGE); messages_to_add.push(message.clone()); yield AgentEvent::Message(message); + messages.extend(messages_to_add); continue } else { let message = Message::assistant().with_text(final_output_tool.final_output.clone().unwrap()); diff --git a/crates/goose/src/agents/final_output_tool.rs b/crates/goose/src/agents/final_output_tool.rs index 14e5d862cc4b..2975e4504d6f 100644 --- a/crates/goose/src/agents/final_output_tool.rs +++ b/crates/goose/src/agents/final_output_tool.rs @@ -38,17 +38,17 @@ impl FinalOutputTool { pub fn tool(&self) -> Tool { let instructions = formatdoc! {r#" - This tool collects the final output for a user and provides validation for structured JSON final output against a predefined schema. + The final_output tool collects the final output for the user and provides validation for structured JSON final output against a predefined schema. - This tool MUST be used for the final output to the user. + This final_output tool MUST be called with the final output for the user. Purpose: - - Collects the final output for a user + - Collects the final output for the user - Ensures that final outputs conform to the expected JSON structure - Provides clear validation feedback when outputs don't match the schema Usage: - - Call the `final_output` tool with your JSON final output + - Call the `final_output` tool with your JSON final output passed as the argument. The expected JSON schema format is: @@ -83,8 +83,8 @@ impl FinalOutputTool { formatdoc! {r#" # Final Output Instructions - You MUST use the `final_output` tool to collect the final output for a user. - The final output MUST be a valid JSON object that matches the following expected schema: + You MUST use the `final_output` tool to collect the final output for the user rather than providing the output directly in your response. + The final output MUST be a valid JSON object that is provided to the `final_output` tool when called and it must match the following schema: {} diff --git a/crates/goose/tests/agent.rs b/crates/goose/tests/agent.rs index d7b502b085f7..99ae07ad44c2 100644 --- a/crates/goose/tests/agent.rs +++ b/crates/goose/tests/agent.rs @@ -661,6 +661,7 @@ mod final_output_tool_tests { #[tokio::test] async fn test_when_final_output_not_called_in_reply() -> Result<()> { use async_trait::async_trait; + use goose::agents::final_output_tool::FINAL_OUTPUT_CONTINUATION_MESSAGE; use goose::conversation::message::Message; use goose::model::ModelConfig; use goose::providers::base::{Provider, ProviderUsage}; @@ -670,6 +671,8 @@ mod final_output_tool_tests { #[derive(Clone)] struct MockProvider { model_config: ModelConfig, + stream_round: std::sync::Arc>, + got_continuation_message: std::sync::Arc>, } #[async_trait] @@ -692,14 +695,38 @@ mod final_output_tool_tests { _messages: &[Message], _tools: &[Tool], ) -> Result { - let deltas = vec![ - Ok((Some(Message::assistant().with_text("Hello")), None)), - Ok((Some(Message::assistant().with_text("Hi!")), None)), - Ok(( - Some(Message::assistant().with_text("What is the final output?")), + if let Some(last_msg) = _messages.last() { + for content in &last_msg.content { + if let goose::conversation::message::MessageContent::Text(text_content) = + content + { + if text_content.text == FINAL_OUTPUT_CONTINUATION_MESSAGE { + let mut got_continuation = + self.got_continuation_message.lock().unwrap(); + *got_continuation = true; + } + } + } + } + + let mut round = self.stream_round.lock().unwrap(); + *round += 1; + + let deltas = if *round == 1 { + vec![ + Ok((Some(Message::assistant().with_text("Hello")), None)), + Ok((Some(Message::assistant().with_text("Hi!")), None)), + Ok(( + Some(Message::assistant().with_text("What is the final output?")), + None, + )), + ] + } else { + vec![Ok(( + Some(Message::assistant().with_text("Additional random delta")), None, - )), - ]; + ))] + }; let stream = stream::iter(deltas.into_iter()); Ok(Box::pin(stream)) @@ -728,7 +755,12 @@ mod final_output_tool_tests { let agent = Agent::new(); let model_config = ModelConfig::new("test-model").unwrap(); - let mock_provider = Arc::new(MockProvider { model_config }); + let mock_provider = Arc::new(MockProvider { + model_config, + stream_round: std::sync::Arc::new(std::sync::Mutex::new(0)), + got_continuation_message: std::sync::Arc::new(std::sync::Mutex::new(false)), + }); + let mock_provider_clone = mock_provider.clone(); agent.update_provider(mock_provider).await?; let response = Response { @@ -764,7 +796,6 @@ mod final_output_tool_tests { } assert!(!responses.is_empty(), "Should have received responses"); - println!("Responses: {:?}", responses); let last_message = responses.last().unwrap(); // Check that the first 3 messages do not have FINAL_OUTPUT_CONTINUATION_MESSAGE @@ -784,6 +815,27 @@ mod final_output_tool_tests { let message_text = last_message.as_concat_text(); assert_eq!(message_text, FINAL_OUTPUT_CONTINUATION_MESSAGE); + // Continue streaming to consume any remaining content, this lets us verify the provider saw the continuation message + while let Some(response_result) = reply_stream.next().await { + match response_result { + Ok(AgentEvent::Message(_response)) => { + break; // Stop after receiving the next message + } + Ok(_) => {} + Err(e) => { + println!("Error while streaming remaining content: {:?}", e); + break; + } + } + } + + // Assert that the provider received the continuation message + let got_continuation = mock_provider_clone.got_continuation_message.lock().unwrap(); + assert!( + *got_continuation, + "Provider should have received the FINAL_OUTPUT_CONTINUATION_MESSAGE" + ); + Ok(()) } }