fix: Missing final structured output should provide guidance to model (#4654)

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());

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:
 
             {}
 

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<std::sync::Mutex<i32>>,
+            got_continuation_message: std::sync::Arc<std::sync::Mutex<bool>>,
         }
 
         #[async_trait]
@@ -692,14 +695,38 @@ mod final_output_tool_tests {
                 _messages: &[Message],
                 _tools: &[Tool],
             ) -> Result<MessageStream, ProviderError> {
-                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(())
     }
 }