suna/backend/agent/tools/message_tool.py

import os
from typing import List, Optional, Union
from agentpress.tool import Tool, ToolResult, openapi_schema, xml_schema

class MessageTool(Tool):
    """Tool for user communication and interaction.
    
    This tool provides methods for notifying users and asking questions, with support for
    attachments and user takeover suggestions.
    """
    
    def __init__(self):
        super().__init__()
    
    @openapi_schema({
        "type": "function",
        "function": {
            "name": "notify",
            "description": "Send a message to user without requiring a response. Use for: 1) Progress updates during long-running tasks, 2) Acknowledging receipt of user instructions, 3) Reporting completion of major milestones, 4) Explaining changes in approach or strategy, 5) Summarizing findings or results without requiring input. IMPORTANT: Use this tool for one-way communication only - do not use when you need user input or confirmation. Always include relevant references when sharing analysis results, generated files, or external resources. The message should be informative, actionable, and provide clear context about the current state of the task.",
            "parameters": {
                "type": "object",
                "properties": {
                    "text": {
                        "type": "string",
                        "description": "Message text to display to user - should be informative and actionable. Include: 1) Current status or progress, 2) Key findings or results, 3) Next steps or implications, 4) Any relevant context or background information."
                    },
                    "references": {
                        "anyOf": [
                            {"type": "string"},
                            {"items": {"type": "string"}, "type": "array"}
                        ],
                        "description": "(Optional) List of files or URLs to reference in the message. Include when: 1) Sharing analysis results or generated files, 2) Referencing external resources or documentation, 3) Providing supporting evidence for findings, 4) Sharing configuration files or scripts. Always use relative paths to /workspace directory."
                    }
                },
                "required": ["text"]
            }
        }
    })
    @xml_schema(
        tag_name="notify",
        mappings=[
            {"param_name": "text", "node_type": "content", "path": "."},
            {"param_name": "references", "node_type": "attribute", "path": ".", "required": False}
        ],
        example='''
        <!-- Use notify for one-way communication when you don't need user input -->
        <!-- Examples of when to use notify: -->
        <!-- 1. Progress updates during long-running tasks -->
        <!-- 2. Completion of major milestones -->
        <!-- 3. Status changes or strategy updates -->
        <!-- 4. Sharing analysis results or findings -->
        <!-- 5. Providing context about next steps -->
        
        <notify references="output/analysis_results.csv,output/visualization.png">
            I've completed the data analysis and generated visualizations of the key trends. The analysis shows a 15% increase in engagement metrics over the last quarter, with the most significant growth in mobile users. I've saved the detailed results in the referenced files. Next, I'll proceed with implementing the recommended optimizations for the mobile experience.
        </notify>
        '''
    )
    async def notify(self, text: str, references: Optional[Union[str, List[str]]] = None) -> ToolResult:
        """Send a notification message to the user without requiring a response.
        
        Args:
            text: The message to display to the user
            references: Optional file paths or URLs to reference in the message
            
        Returns:
            ToolResult indicating success or failure of the notification
        """
        try:
            # Convert single reference to list for consistent handling
            if references and isinstance(references, str):
                references = [references]
                
            # Format the response message
            response_text = f"NOTIFICATION: {text}"
            
            # Add references information if present
            if references:
                reference_list = "\n- ".join(references)
                response_text += f"\n\nReferences:\n- {reference_list}"
            
            return self.success_response(response_text)
        except Exception as e:
            return self.fail_response(f"Error sending notification: {str(e)}")
    
    @openapi_schema({
        "type": "function",
        "function": {
            "name": "ask",
            "description": "Ask user a question and wait for response. Use for: 1) Requesting clarification on ambiguous requirements, 2) Seeking confirmation before proceeding with high-impact changes, 3) Gathering additional information needed to complete a task, 4) Offering options and requesting user preference, 5) Validating assumptions when critical to task success. IMPORTANT: Use this tool only when user input is essential to proceed. Always provide clear context and options when applicable. Include relevant references when the question relates to specific files or resources.",
            "parameters": {
                "type": "object",
                "properties": {
                    "text": {
                        "type": "string",
                        "description": "Question text to present to user - should be specific and clearly indicate what information you need. Include: 1) Clear question or request, 2) Context about why the input is needed, 3) Available options if applicable, 4) Impact of different choices, 5) Any relevant constraints or considerations."
                    },
                    "references": {
                        "anyOf": [
                            {"type": "string"},
                            {"items": {"type": "string"}, "type": "array"}
                        ],
                        "description": "(Optional) List of files or URLs to reference in the question. Include when: 1) Question relates to specific files or configurations, 2) User needs to review content before answering, 3) Options or choices are documented in files, 4) Supporting evidence or context is needed. Always use relative paths to /workspace directory."
                    }
                },
                "required": ["text"]
            }
        }
    })
    @xml_schema(
        tag_name="ask",
        mappings=[
            {"param_name": "text", "node_type": "content", "path": "."},
            {"param_name": "references", "node_type": "attribute", "path": ".", "required": False}
        ],
        example='''
        <!-- Use ask when you need user input to proceed -->
        <!-- Examples of when to use ask: -->
        <!-- 1. Clarifying ambiguous requirements -->
        <!-- 2. Confirming high-impact changes -->
        <!-- 3. Choosing between implementation options -->
        <!-- 4. Validating critical assumptions -->
        <!-- 5. Getting missing information -->
        
        <ask references="recipes/chocolate_cake.txt,photos/cake_examples.jpg">
            I'm planning to bake the chocolate cake for your birthday party. The recipe mentions "rich frosting" but doesn't specify what type. Could you clarify your preferences? For example:
            1. Would you prefer buttercream or cream cheese frosting?
            2. Do you want any specific flavor added to the frosting (vanilla, coffee, etc.)?
            3. Should I add any decorative toppings like sprinkles or fruit?
            4. Do you have any dietary restrictions I should be aware of?
            
            This information will help me make sure the cake meets your expectations for the celebration.
        </ask>
        '''
    )
    async def ask(self, text: str, references: Optional[Union[str, List[str]]] = None) -> ToolResult:
        """Ask the user a question and wait for a response.
        
        Args:
            text: The question to present to the user
            references: Optional file paths or URLs to reference in the question
            
        Returns:
            ToolResult indicating the question was successfully sent
        """
        try:
            # Convert single reference to list for consistent handling
            if references and isinstance(references, str):
                references = [references]
                
            # Format the question message
            response_text = f"QUESTION: {text}"
            
            # Add references information if present
            if references:
                reference_list = "\n- ".join(references)
                response_text += f"\n\nReferences:\n- {reference_list}"
            
            return self.success_response(response_text, requires_response=True)
        except Exception as e:
            return self.fail_response(f"Error asking user: {str(e)}")

    @openapi_schema({
        "type": "function",
        "function": {
            "name": "complete",
            "description": "A special tool to indicate you have completed all tasks and are about to enter complete state. Use ONLY when: 1) All tasks in todo.md are marked complete [x], 2) The user's original request has been fully addressed, 3) There are no pending actions or follow-ups required, 4) You've delivered all final outputs and results to the user. IMPORTANT: This is the ONLY way to properly terminate execution. Never use this tool unless ALL tasks are complete and verified. Always ensure you've provided all necessary outputs and references before using this tool.",
            "parameters": {
                "type": "object"
            }
        }
    })
    @xml_schema(
        tag_name="complete",
        mappings=[],
        example='''
        <!-- Use complete ONLY when ALL tasks are finished -->
        <!-- Prerequisites for using complete: -->
        <!-- 1. All todo.md items marked complete [x] -->
        <!-- 2. User's original request fully addressed -->
        <!-- 3. All outputs and results delivered -->
        <!-- 4. No pending actions or follow-ups -->
        <!-- 5. All tasks verified and validated -->
        
        <complete>
        <!-- This tool indicates successful completion of all tasks -->
        <!-- The system will stop execution after this tool is used -->
        </complete>
        '''
    )
    async def complete(self) -> ToolResult:
        """Indicate that the agent has completed all tasks and is entering complete state.
        
        Returns:
            ToolResult indicating successful transition to complete state
        """
        try:
            return self.success_response("Entering complete state")
        except Exception as e:
            return self.fail_response(f"Error entering complete state: {str(e)}")


if __name__ == "__main__":
    import asyncio
    
    async def test_message_tool():
        message_tool = MessageTool()
        
        # Test notification
        notify_result = await message_tool.notify(
            "Processing has completed successfully!",
            references=["results.txt", "output.log"]
        )
        print("Notification result:", notify_result)
        
        # Test question
        ask_result = await message_tool.ask(
            "Would you like to proceed with the next phase?",
            references="summary.pdf"
        )
        print("Question result:", ask_result)
    
    asyncio.run(test_message_tool())