Update messages.py

Author: Fr4nc3

src/backend/v3/models/messages.py

@@ -1,5 +1,6 @@
 """Messages from the backend to the frontend via WebSocket."""
 
+from enum import Enum
 import uuid
 from dataclasses import asdict, dataclass, field
 from typing import Any, Dict, List, Literal, Optional
@@ -135,4 +136,36 @@ class ApprovalRequest(KernelBaseModel):
     session_id: str
     user_id: str
     action: str
-    agent_name: str
+    agent_name: str
+
+
+
+class WebsocketMessageType(str, Enum):
+    AGENT_MESSAGE = "agent_message"
+    AGENT_STREAM_START = "agent_stream_start"
+    AGENT_STREAM_END = "agent_stream_end"
+    AGENT_MESSAGE_STREAMING = "agent_message_streaming"
+    AGENT_TOOL_MESSAGE = "agent_tool_message"
+    PLAN_APPROVAL_REQUEST = "plan_approval_request"
+    PLAN_APPROVAL_RESPONSE = "plan_approval_response"
+    REPLAN_APPROVAL_REQUEST = "replan_approval_request"
+    REPLAN_APPROVAL_RESPONSE = "replan_approval_response"
+    USER_CLARIFICATION_REQUEST = "user_clarification_request"
+    USER_CLARIFICATION_RESPONSE = "user_clarification_response"
+    FINAL_RESULT_MESSAGE = "final_result_message"
+
+
+@dataclass(slots=True)
+class WebsocketMessage:
+    """Generic WebSocket message wrapper."""
+    type: WebsocketMessageType
+    message: Any
+    data: Any
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert the WebsocketMessage to a dictionary for JSON serialization."""
+        return {
+            "type": self.type,
+            "data": self.data.to_dict() if hasattr(self.data, 'to_dict') else self.data,
+            "message": self.message.to_dict() if hasattr(self.message, 'to_dict') else self.message
+        }