feat: Structured output for recipes (#3188)

2025-07-02 12:16:57 +10:00
parent 620474b76e
commit 0a00b0f588
13 changed files with 754 additions and 7 deletions
@@ -10,6 +10,7 @@ use futures_util::stream;
 use futures_util::stream::StreamExt;
 use mcp_core::protocol::JsonRpcMessage;

+use crate::agents::final_output_tool::{FINAL_OUTPUT_CONTINUATION_MESSAGE, FINAL_OUTPUT_TOOL_NAME};
 use crate::agents::sub_recipe_manager::SubRecipeManager;
 use crate::config::{Config, ExtensionConfigManager, PermissionManager};
 use crate::message::Message;
@@ -17,7 +18,7 @@ use crate::permission::permission_judge::check_tool_permissions;
 use crate::permission::PermissionConfirmation;
 use crate::providers::base::Provider;
 use crate::providers::errors::ProviderError;
-use crate::recipe::{Author, Recipe, Settings, SubRecipe};
+use crate::recipe::{Author, Recipe, Response, Settings, SubRecipe};
 use crate::scheduler_trait::SchedulerTrait;
 use crate::tool_monitor::{ToolCall, ToolMonitor};
 use regex::Regex;
@@ -47,6 +48,7 @@ use mcp_core::{

 use crate::agents::subagent_tools::SUBAGENT_RUN_TASK_TOOL_NAME;

+use super::final_output_tool::FinalOutputTool;
 use super::platform_tools;
 use super::router_tools;
 use super::subagent_manager::SubAgentManager;
@@ -58,6 +60,7 @@ pub struct Agent {
    pub(super) provider: Mutex<Option<Arc<dyn Provider>>>,
    pub(super) extension_manager: RwLock<ExtensionManager>,
    pub(super) sub_recipe_manager: Mutex<SubRecipeManager>,
+    pub(super) final_output_tool: Mutex<Option<FinalOutputTool>>,
    pub(super) frontend_tools: Mutex<HashMap<String, FrontendTool>>,
    pub(super) frontend_instructions: Mutex<Option<String>>,
    pub(super) prompt_manager: Mutex<PromptManager>,
@@ -131,6 +134,7 @@ impl Agent {
            provider: Mutex::new(None),
            extension_manager: RwLock::new(ExtensionManager::new()),
            sub_recipe_manager: Mutex::new(SubRecipeManager::new()),
+            final_output_tool: Mutex::new(None),
            frontend_tools: Mutex::new(HashMap::new()),
            frontend_instructions: Mutex::new(None),
            prompt_manager: Mutex::new(PromptManager::new()),
@@ -205,6 +209,14 @@ impl Agent {
        Ok(tools)
    }

+    pub async fn add_final_output_tool(&self, response: Response) {
+        let mut final_output_tool = self.final_output_tool.lock().await;
+        let created_final_output_tool = FinalOutputTool::new(response);
+        let final_output_system_prompt = created_final_output_tool.system_prompt();
+        *final_output_tool = Some(created_final_output_tool);
+        self.extend_system_prompt(final_output_system_prompt).await;
+    }
+
    pub async fn add_sub_recipes(&self, sub_recipes: Vec<SubRecipe>) {
        let mut sub_recipe_manager = self.sub_recipe_manager.lock().await;
        sub_recipe_manager.add_sub_recipe_tools(sub_recipes);
@@ -258,6 +270,20 @@ impl Agent {
            return (request_id, Ok(ToolCallResult::from(result)));
        }

+        if tool_call.name == FINAL_OUTPUT_TOOL_NAME {
+            if let Some(final_output_tool) = self.final_output_tool.lock().await.as_mut() {
+                let result = final_output_tool.execute_tool_call(tool_call.clone()).await;
+                return (request_id, Ok(result));
+            } else {
+                return (
+                    request_id,
+                    Err(ToolError::ExecutionError(
+                        "Final output tool not defined".to_string(),
+                    )),
+                );
+            }
+        }
+
        let extension_manager = self.extension_manager.read().await;
        let sub_recipe_manager = self.sub_recipe_manager.lock().await;

@@ -544,6 +570,10 @@ impl Agent {
        if extension_name.is_none() {
            let sub_recipe_manager = self.sub_recipe_manager.lock().await;
            prefixed_tools.extend(sub_recipe_manager.sub_recipe_tools.values().cloned());
+
+            if let Some(final_output_tool) = self.final_output_tool.lock().await.as_ref() {
+                prefixed_tools.push(final_output_tool.tool());
+            }
        }

        prefixed_tools
@@ -766,6 +796,15 @@ impl Agent {

                        let num_tool_requests = frontend_requests.len() + remaining_requests.len();
                        if num_tool_requests == 0 {
+                            if let Some(final_output_tool) = self.final_output_tool.lock().await.as_ref() {
+                                if final_output_tool.final_output.is_none() {
+                                    tracing::warn!("Final output tool has not been called yet. Continuing agent loop.");
+                                    yield AgentEvent::Message(Message::user().with_text(FINAL_OUTPUT_CONTINUATION_MESSAGE));
+                                    continue;
+                                } else {
+                                    yield AgentEvent::Message(Message::assistant().with_text(final_output_tool.final_output.clone().unwrap()));
+                                }
+                            }
                            break;
                        }

@@ -1260,3 +1299,43 @@ impl Agent {
        Ok(recipe)
    }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::recipe::Response;
+
+    #[tokio::test]
+    async fn test_add_final_output_tool() -> Result<()> {
+        let agent = Agent::new();
+
+        let response = Response {
+            json_schema: Some(serde_json::json!({
+                "type": "object",
+                "properties": {
+                    "result": {"type": "string"}
+                }
+            })),
+        };
+
+        agent.add_final_output_tool(response).await;
+
+        let tools = agent.list_tools(None).await;
+        let final_output_tool = tools.iter().find(|tool| tool.name == "final_output");
+
+        assert!(
+            final_output_tool.is_some(),
+            "Final output tool should be present after adding"
+        );
+
+        let prompt_manager = agent.prompt_manager.lock().await;
+        let system_prompt =
+            prompt_manager.build_system_prompt(vec![], None, serde_json::Value::Null, None, None);
+
+        let final_output_tool_ref = agent.final_output_tool.lock().await;
+        let final_output_tool_system_prompt =
+            final_output_tool_ref.as_ref().unwrap().system_prompt();
+        assert!(system_prompt.contains(&final_output_tool_system_prompt));
+        Ok(())
+    }
+}
@@ -0,0 +1,261 @@
+use crate::agents::tool_execution::ToolCallResult;
+use crate::recipe::Response;
+use indoc::formatdoc;
+use mcp_core::{
+    tool::{Tool, ToolAnnotations},
+    Content, ToolCall, ToolError,
+};
+use serde_json::Value;
+
+pub const FINAL_OUTPUT_TOOL_NAME: &str = "final_output";
+pub const FINAL_OUTPUT_CONTINUATION_MESSAGE: &str =
+    "You MUST call the `final_output` tool with your final output for the user.";
+
+pub struct FinalOutputTool {
+    pub response: Response,
+    /// The final output collected for the user. It will be a single line string for easy script extraction from output.
+    pub final_output: Option<String>,
+}
+
+impl FinalOutputTool {
+    pub fn new(response: Response) -> Self {
+        if response.json_schema.is_none() {
+            panic!("Cannot create FinalOutputTool: json_schema is required");
+        }
+        let schema = response.json_schema.as_ref().unwrap();
+
+        if let Some(obj) = schema.as_object() {
+            if obj.is_empty() {
+                panic!("Cannot create FinalOutputTool: empty json_schema is not allowed");
+            }
+        }
+
+        jsonschema::meta::validate(schema).unwrap();
+        Self {
+            response,
+            final_output: None,
+        }
+    }
+
+    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.
+
+            This tool MUST be used for the final output to the user.
+            
+            Purpose:
+            - Collects the final output for a 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
+            
+            The expected JSON schema format is:
+
+            {}
+            
+            When validation fails, you'll receive:
+            - Specific validation errors
+            - The expected format
+        "#, serde_json::to_string_pretty(self.response.json_schema.as_ref().unwrap()).unwrap()};
+
+        Tool::new(
+            FINAL_OUTPUT_TOOL_NAME.to_string(),
+            instructions,
+            self.response.json_schema.as_ref().unwrap().clone(),
+            Some(ToolAnnotations {
+                title: Some("Final Output".to_string()),
+                read_only_hint: false,
+                destructive_hint: false,
+                idempotent_hint: true,
+                open_world_hint: false,
+            }),
+        )
+    }
+
+    pub fn system_prompt(&self) -> String {
+        formatdoc! {r#"
+            # Final Ouptut 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:
+
+            {}
+
+            ----
+        "#, serde_json::to_string_pretty(self.response.json_schema.as_ref().unwrap()).unwrap()}
+    }
+
+    async fn validate_json_output(&self, output: &Value) -> Result<Value, String> {
+        let compiled_schema =
+            match jsonschema::validator_for(self.response.json_schema.as_ref().unwrap()) {
+                Ok(schema) => schema,
+                Err(e) => {
+                    return Err(format!("Internal error: Failed to compile schema: {}", e));
+                }
+            };
+
+        let validation_errors: Vec<String> = compiled_schema
+            .iter_errors(output)
+            .map(|error| format!("- {}: {}", error.instance_path, error))
+            .collect();
+
+        if validation_errors.is_empty() {
+            Ok(output.clone())
+        } else {
+            Err(format!(
+                "Validation failed:\n{}\n\nExpected format:\n{}\n\nPlease correct your output to match the expected JSON schema and try again.",
+                validation_errors.join("\n"),
+                serde_json::to_string_pretty(self.response.json_schema.as_ref().unwrap()).unwrap_or_else(|_| "Invalid schema".to_string())
+            ))
+        }
+    }
+
+    pub async fn execute_tool_call(&mut self, tool_call: ToolCall) -> ToolCallResult {
+        match tool_call.name.as_str() {
+            FINAL_OUTPUT_TOOL_NAME => {
+                let result = self.validate_json_output(&tool_call.arguments).await;
+                match result {
+                    Ok(parsed_value) => {
+                        self.final_output = Some(Self::parsed_final_output_string(parsed_value));
+                        ToolCallResult::from(Ok(vec![Content::text(
+                            "Final output successfully collected.".to_string(),
+                        )]))
+                    }
+                    Err(error) => ToolCallResult::from(Err(ToolError::InvalidParameters(error))),
+                }
+            }
+            _ => ToolCallResult::from(Err(ToolError::NotFound(format!(
+                "Unknown tool: {}",
+                tool_call.name
+            )))),
+        }
+    }
+
+    // Formats the parsed JSON as a single line string so its easy to extract from the output
+    fn parsed_final_output_string(parsed_json: Value) -> String {
+        serde_json::to_string(&parsed_json).unwrap()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::recipe::Response;
+    use serde_json::json;
+
+    fn create_complex_test_schema() -> Value {
+        json!({
+            "type": "object",
+            "properties": {
+                "user": {
+                    "type": "object",
+                    "properties": {
+                        "name": {"type": "string"},
+                        "age": {"type": "number"}
+                    },
+                    "required": ["name", "age"]
+                },
+                "tags": {
+                    "type": "array",
+                    "items": {"type": "string"}
+                }
+            },
+            "required": ["user", "tags"]
+        })
+    }
+
+    #[test]
+    #[should_panic(expected = "Cannot create FinalOutputTool: json_schema is required")]
+    fn test_new_with_missing_schema() {
+        let response = Response { json_schema: None };
+        FinalOutputTool::new(response);
+    }
+
+    #[test]
+    #[should_panic(expected = "Cannot create FinalOutputTool: empty json_schema is not allowed")]
+    fn test_new_with_empty_schema() {
+        let response = Response {
+            json_schema: Some(json!({})),
+        };
+        FinalOutputTool::new(response);
+    }
+
+    #[test]
+    #[should_panic]
+    fn test_new_with_invalid_schema() {
+        let response = Response {
+            json_schema: Some(json!({
+                "type": "invalid_type",
+                "properties": {
+                    "message": {
+                        "type": "unknown_type"
+                    }
+                }
+            })),
+        };
+        FinalOutputTool::new(response);
+    }
+
+    #[tokio::test]
+    async fn test_execute_tool_call_schema_validation_failure() {
+        let response = Response {
+            json_schema: Some(json!({
+                "type": "object",
+                "properties": {
+                    "message": {
+                        "type": "string"
+                    },
+                    "count": {
+                        "type": "number"
+                    }
+                },
+                "required": ["message", "count"]
+            })),
+        };
+
+        let mut tool = FinalOutputTool::new(response);
+        let tool_call = ToolCall {
+            name: FINAL_OUTPUT_TOOL_NAME.to_string(),
+            arguments: json!({
+                "message": "Hello"  // Missing required "count" field
+            }),
+        };
+
+        let result = tool.execute_tool_call(tool_call).await;
+        let tool_result = result.result.await;
+        assert!(tool_result.is_err());
+        if let Err(error) = tool_result {
+            assert!(error.to_string().contains("Validation failed"));
+        }
+    }
+
+    #[tokio::test]
+    async fn test_execute_tool_call_complex_valid_json() {
+        let response = Response {
+            json_schema: Some(create_complex_test_schema()),
+        };
+
+        let mut tool = FinalOutputTool::new(response);
+        let tool_call = ToolCall {
+            name: FINAL_OUTPUT_TOOL_NAME.to_string(),
+            arguments: json!({
+                "user": {
+                    "name": "John",
+                    "age": 30
+                },
+                "tags": ["developer", "rust"]
+            }),
+        };
+
+        let result = tool.execute_tool_call(tool_call).await;
+        let tool_result = result.result.await;
+        assert!(tool_result.is_ok());
+        assert!(tool.final_output.is_some());
+
+        let final_output = tool.final_output.unwrap();
+        assert!(serde_json::from_str::<Value>(&final_output).is_ok());
+        assert!(!final_output.contains('\n'));
+    }
+}
@@ -2,6 +2,7 @@ mod agent;
 mod context;
 pub mod extension;
 pub mod extension_manager;
+pub mod final_output_tool;
 mod large_response_handler;
 pub mod platform_tools;
 pub mod prompt_manager;