Refactor WebSocket message handling and types

Standardizes WebSocket message types using the WebsocketMessageType enum across multiple modules. Refactors message sending to include explicit message type and data structure, and updates API endpoint for plan retrieval. Fixes and improves consistency in how messages are sent and handled throughout the backend.

Author: Fr4nc3

src/backend/v3/api/router.py

@@ -1036,8 +1036,8 @@ async def get_plans(request: Request):
 
 
 # Get plans is called in the initial side rendering of the frontend
-@app_v3.get("/plan/{plan_id}")
-async def get_plan_by_id(request: Request, plan_id: str):
+@app_v3.get("/plan")
+async def get_plan_by_id(request: Request,   plan_id: str):
     """
     Retrieve plans for the current user.
 
@@ -1133,21 +1133,8 @@ async def get_plan_by_id(request: Request, plan_id: str):
         )
 
         return [plan_with_steps, formatted_messages]
-
-    current_team = await memory_store.get_current_team(user_id=user_id)
-    if not current_team:
-        return []
-
-    all_plans = await memory_store.get_all_plans_by_team_id(team_id=current_team.id)
-    # Fetch steps for all plans concurrently
-    steps_for_all_plans = await asyncio.gather(
-        *[memory_store.get_steps_by_plan(plan_id=plan.id) for plan in all_plans]
-    )
-    # Create list of PlanWithSteps and update step counts
-    list_of_plans_with_steps = []
-    for plan, steps in zip(all_plans, steps_for_all_plans):
-        plan_with_steps = PlanWithSteps(**plan.model_dump(), steps=steps)
-        plan_with_steps.update_step_counts()
-        list_of_plans_with_steps.append(plan_with_steps)
-
-    return []
+    else:
+        track_event_if_configured(
+            "GetPlanId", {"status_code": 400, "detail": "no plan id"}
+        )
+        raise HTTPException(status_code=400, detail="no plan id")

src/backend/v3/callbacks/response_handlers.py

@@ -56,6 +56,6 @@ async def streaming_agent_response_callback(streaming_message: StreamingChatMess
         if user_id:
             try:
                 message = AgentMessageStreaming(agent_name=streaming_message.name or "Unknown Agent", content=streaming_message.content, is_final=is_final)
-                await connection_config.send_status_update_async(message, user_id)
+                await connection_config.send_status_update_async(message, user_id, message_type=WebsocketMessageType.AGENT_MESSAGE_STREAMING)
             except Exception as e:
                 logging.error(f"Response_callback: Error sending streaming WebSocket message: {e}")

src/backend/v3/config/settings.py

@@ -199,10 +199,20 @@ async def send_status_update_async(
             )
             return
 
+        try:
+            if hasattr(message, "data") and hasattr(message, "type"):
+                message = message.data
+        except Exception as e:
+                print(f"Error loading message data: {e}")
+
+        standard_message = {
+            "type": message_type,
+            "data": message
+        }
         connection = self.get_connection(process_id)
         if connection:
             try:
-                str_message = json.dumps(message, default=str)
+                str_message = json.dumps(standard_message, default=str)
                 await connection.send_text(str_message)
                 logger.debug(f"Message sent to user {user_id} via process {process_id}")
             except Exception as e:

src/backend/v3/magentic_agents/proxy_agent.py

@@ -24,7 +24,7 @@
 from v3.config.settings import (connection_config, current_user_id,
                                 orchestration_config)
 from v3.models.messages import (UserClarificationRequest,
-                                UserClarificationResponse)
+                                UserClarificationResponse, WebsocketMessageType)
 
 
 class DummyAgentThread(AgentThread):
@@ -155,10 +155,10 @@ async def invoke(self, message: str,*, thread: AgentThread | None = None,**kwarg
 
         # Send the approval request to the user's WebSocket
         await connection_config.send_status_update_async({
-            "type": "user_clarification_request", 
+            "type": WebsocketMessageType.USER_CLARIFICATION_REQUEST,
             "data": clarification_message
-        })
-        
+        }, user_id=current_user_id.get(), message_type=WebsocketMessageType.USER_CLARIFICATION_REQUEST)
+
         # Get human input
         human_response = await self._wait_for_user_clarification(clarification_message.request_id)
 
@@ -203,10 +203,10 @@ async def invoke_stream(self, messages, thread=None, **kwargs) -> AsyncIterator[
         # Send the approval request to the user's WebSocket
         # The user_id will be automatically retrieved from context
         await connection_config.send_status_update_async({
-            "type": "user_clarification_request", 
+            "type": WebsocketMessageType.USER_CLARIFICATION_REQUEST, 
             "data": clarification_message
-        })
-        
+        }, user_id=current_user_id.get(), message_type=WebsocketMessageType.USER_CLARIFICATION_REQUEST)
+
         # Get human input - replace with websocket call when available
         human_response = await self._wait_for_user_clarification(clarification_message.request_id)
 

src/backend/v3/orchestration/orchestration_manager.py

@@ -136,7 +136,7 @@ async def run_orchestration(self, user_id, input_task) -> None:
 
                 # Send final result via WebSocket
                 await connection_config.send_status_update_async({
-                    "type": "final_result",
+                    "type": WebsocketMessageType.FINAL_RESULT_MESSAGE,
                     "data": {
                         "content": str(value),
                         "status": "completed",