Refactor WebSocket message types and models

Moved WebSocket-related types (StreamMessage, StreamingPlanUpdate, etc.) from the service to the models directory for better separation of concerns. Introduced a WebsocketMessageType enum for consistent message type handling. Updated imports and usages across components, hooks, and services to use the new model definitions. Improved message handling logic in WebSocketService for clarity and maintainability.

Author: Fr4nc3

src/frontend/src/components/content/PlanPanelRight.tsx

@@ -12,11 +12,10 @@ import {
   ClockRegular,
   PersonRegular,
 } from "@fluentui/react-icons";
-import { MPlanData } from "../../models";
+import { MPlanData, StreamingPlanUpdate } from "../../models";
 import { TaskService } from "../../services/TaskService";
-import { Step } from "../../models/plan";
 import { PlanDataService } from "../../services/PlanDataService";
-import webSocketService, { StreamingPlanUpdate } from "../../services/WebSocketService";
+import webSocketService from "../../services/WebSocketService";
 import ContentNotFound from "../NotFound/ContentNotFound";
 
 // Clean interface - only display-related props

src/frontend/src/hooks/useWebSocket.tsx

@@ -1,5 +1,6 @@
 import { useCallback, useEffect, useRef, useState } from 'react';
-import { webSocketService, StreamMessage } from '../services/WebSocketService';
+import { webSocketService } from '@/services';
+import { StreamMessage } from '../models';
 
 export interface WebSocketState {
   isConnected: boolean;
@@ -15,7 +16,7 @@ export const useWebSocket = () => {
     isReconnecting: false,
     error: null
   });
-  
+
   const isConnectedRef = useRef(false);
   const isConnectingRef = useRef(false);
   const lastSessionIdRef = useRef<string | null>(null);
@@ -32,54 +33,54 @@ export const useWebSocket = () => {
 
   const connectWebSocket = useCallback(async (sessionId: string, processId?: string) => {
     if (isConnectedRef.current || isConnectingRef.current) return;
-    
+
     setIsConnecting(true);
     lastSessionIdRef.current = sessionId;
     lastProcessIdRef.current = processId;
-      
+
     try {
       await webSocketService.connect(sessionId, processId);
       isConnectedRef.current = true;
-      setState(prev => ({ 
-        ...prev, 
-        isConnected: true, 
-        isConnecting: false, 
-        error: null 
+      setState(prev => ({
+        ...prev,
+        isConnected: true,
+        isConnecting: false,
+        error: null
       }));
     } catch (error) {
       console.error('Failed to connect to WebSocket:', error);
       isConnectedRef.current = false;
       isConnectingRef.current = false;
-      setState(prev => ({ 
-        ...prev, 
-        isConnected: false, 
-        isConnecting: false, 
-        error: 'Failed to connect to server' 
+      setState(prev => ({
+        ...prev,
+        isConnected: false,
+        isConnecting: false,
+        error: 'Failed to connect to server'
       }));
     }
   }, [setIsConnecting]);
 
   const reconnect = useCallback(async () => {
     if (!lastSessionIdRef.current) return;
-    
+
     setIsReconnecting(true);
     try {
       await webSocketService.connect(lastSessionIdRef.current, lastProcessIdRef.current);
       isConnectedRef.current = true;
-      setState(prev => ({ 
-        ...prev, 
-        isConnected: true, 
-        isReconnecting: false, 
-        error: null 
+      setState(prev => ({
+        ...prev,
+        isConnected: true,
+        isReconnecting: false,
+        error: null
       }));
     } catch (error) {
       console.error('Failed to reconnect to WebSocket:', error);
       isConnectedRef.current = false;
-      setState(prev => ({ 
-        ...prev, 
-        isConnected: false, 
-        isReconnecting: false, 
-        error: 'Failed to reconnect to server' 
+      setState(prev => ({
+        ...prev,
+        isConnected: false,
+        isReconnecting: false,
+        error: 'Failed to reconnect to server'
       }));
     }
   }, [setIsReconnecting]);
@@ -88,11 +89,11 @@ export const useWebSocket = () => {
     webSocketService.disconnect();
     isConnectedRef.current = false;
     isConnectingRef.current = false;
-    setState(prev => ({ 
-      ...prev, 
-      isConnected: false, 
-      isConnecting: false, 
-      isReconnecting: false 
+    setState(prev => ({
+      ...prev,
+      isConnected: false,
+      isConnecting: false,
+      isReconnecting: false
     }));
   }, []);
 
@@ -102,8 +103,8 @@ export const useWebSocket = () => {
       if (message.data?.connected !== undefined) {
         const connected = message.data.connected;
         isConnectedRef.current = connected;
-        setState(prev => ({ 
-          ...prev, 
+        setState(prev => ({
+          ...prev,
           isConnected: connected,
           isConnecting: false,
           isReconnecting: false,
@@ -115,10 +116,10 @@ export const useWebSocket = () => {
     // Set up error listener
     const unsubscribeError = webSocketService.on('error', (message: StreamMessage) => {
       isConnectedRef.current = false;
-      setState(prev => ({ 
-        ...prev, 
+      setState(prev => ({
+        ...prev,
         isConnected: false,
-        error: message.data?.error || 'WebSocket error occurred' 
+        error: message.data?.error || 'WebSocket error occurred'
       }));
     });
 

src/frontend/src/models/enums.tsx

@@ -235,3 +235,19 @@ export enum HumanFeedbackStatus {
     ACCEPTED = "accepted",
     REJECTED = "rejected"
 }
+
+export enum WebsocketMessageType {
+    SYSTEM_MESSAGE = "system_message",
+    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"
+}

src/frontend/src/models/index.tsx

@@ -22,11 +22,5 @@ export type { Agent as TaskAgent } from './taskDetails';
 // Export Team models (Agent interface takes precedence)
 export * from './Team';
 
-// Export WebSocket service types that are needed by components
-export type { 
-    StreamingPlanUpdate, 
-    StreamMessage, 
-    ParsedPlanApprovalRequest 
-} from '../services/WebSocketService';
 
 // Add other model exports as needed

src/frontend/src/models/messages.tsx

@@ -1,4 +1,5 @@
-import { AgentType, StepStatus, PlanStatus } from './enums';
+import { AgentType, StepStatus, PlanStatus, WebsocketMessageType } from './enums';
+import { MPlanData } from './plan';
 
 /**
  * Message roles compatible with Semantic Kernel
@@ -114,4 +115,57 @@ export interface PlanStateUpdate {
     session_id: string;
     /** Overall status of the plan */
     overall_status: PlanStatus;
+}
+
+
+
+export interface StreamMessage {
+    type: WebsocketMessageType
+    plan_id?: string;
+    session_id?: string;
+    data?: any;
+    timestamp?: string | number;
+}
+
+export interface StreamingPlanUpdate {
+    plan_id: string;
+    session_id?: string;
+    step_id?: string;
+    agent_name?: string;
+    content?: string;
+    status?: 'in_progress' | 'completed' | 'error' | 'creating_plan' | 'pending_approval';
+    message_type?: 'thinking' | 'action' | 'result' | 'clarification_needed' | 'plan_approval_request';
+    timestamp?: number;
+    is_final?: boolean;
+}
+
+export interface PlanApprovalRequestData {
+    plan_id: string;
+    session_id: string;
+    plan: {
+        steps: Array<{
+            id: string;
+            description: string;
+            agent: string;
+            estimated_duration?: string;
+        }>;
+        total_steps: number;
+        estimated_completion?: string;
+    };
+    status: 'PENDING_APPROVAL';
+}
+
+export interface PlanApprovalResponseData {
+    plan_id: string;
+    session_id: string;
+    approved: boolean;
+    feedback?: string;
+}
+
+// Structured plan approval request
+export interface ParsedPlanApprovalRequest {
+    type: 'parsed_plan_approval_request';
+    plan_id: string;
+    parsedData: MPlanData;
+    rawData: string;
 }

src/frontend/src/models/plan.tsx

@@ -1,5 +1,5 @@
 import { AgentType, PlanStatus, StepStatus, HumanFeedbackStatus } from './enums';
-import { StreamingPlanUpdate } from '../services/WebSocketService';
+import { StreamingPlanUpdate } from './messages';
 
 /**
  * Base interface with common fields

src/frontend/src/pages/PlanPage.tsx

@@ -20,9 +20,9 @@ import LoadingMessage, { loadingMessages } from "../coral/components/LoadingMess
 import { RAIErrorCard, RAIErrorData } from "../components/errors";
 import { TeamConfig } from "../models/Team";
 import { TeamService } from "../services/TeamService";
-import webSocketService, { StreamMessage, StreamingPlanUpdate } from "../services/WebSocketService";
+import webSocketService from "../services/WebSocketService";
 import { APIService } from "../api/apiService";
-import { Step } from "../models/plan";
+import { StreamMessage, StreamingPlanUpdate } from "../models";
 
 import "../styles/PlanPage.css"