Refactor AgentMessageResponse model and usage

Removed unused fields from AgentMessageResponse in both backend and frontend models for consistency. Updated related logic in PlanPage and PlanDataService to match the new structure and improve message persistence handling.

Author: Fr4nc3

src/backend/common/models/messages_kernel.py

@@ -2,6 +2,7 @@
 from datetime import datetime, timezone
 from enum import Enum
 from typing import Any, Dict, List, Literal, Optional
+from pydantic import BaseModel
 from semantic_kernel.kernel_pydantic import Field, KernelBaseModel
 
 class DataType(str, Enum):
@@ -53,6 +54,7 @@ class PlanStatus(str, Enum):
     failed = "failed"
     canceled = "canceled"
     approved = "approved"
+    created = "created"
 
 
 class HumanFeedbackStatus(str, Enum):
@@ -109,7 +111,21 @@ class UserCurrentTeam(BaseDataModel):
     user_id: str
     team_id: str
 
-
+class MStep(BaseModel):
+    """model of a step in a plan"""
+    agent: str = ""
+    action: str = ""
+class MPlan(BaseModel):
+    """model of a plan"""
+    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
+    user_id: str = ""
+    team_id: str = ""
+    plan_id: str = ""
+    overall_status: PlanStatus = PlanStatus.created
+    user_request: str = ""
+    team: List[str] = []
+    facts: str = ""
+    steps: List[MStep] = []
 class Plan(BaseDataModel):
     """Represents a plan containing multiple steps."""
 
@@ -120,6 +136,7 @@ class Plan(BaseDataModel):
     overall_status: PlanStatus = PlanStatus.in_progress
     approved: bool = False
     source: str = AgentType.PLANNER.value
+    m_plan: Optional[MPlan] = None
     summary: Optional[str] = None
     team_id: Optional[str] = None
     human_clarification_request: Optional[str] = None

src/backend/v3/api/router.py

@@ -1340,12 +1340,12 @@ async def get_plan_by_id(request: Request,  plan_id: Optional[str] = Query(None)
 
             team = await memory_store.get_team_by_id(team_id=plan.team_id)
             agent_messages = await memory_store.get_agent_messages(plan_id=plan.plan_id)
-            m_plan = await memory_store.get_mplan(plan_id=plan.plan_id)
+
             return {
                 "plan": plan,
                 "team": team if team else None,
                 "messages": agent_messages,
-                "m_plan": m_plan if m_plan else None,
+                "m_plan": plan.m_plan,
             }
         else:
             track_event_if_configured(

src/backend/v3/common/services/plan_service.py

@@ -158,8 +158,8 @@ async def handle_plan_approval(
                     orchestration_config.plans[human_feedback.m_plan_id] = mplan
                     if plan:
                         plan.overall_status = PlanStatus.approved
+                        plan.m_plan = mplan
                         await memory_store.update_plan(plan)
-                        await memory_store.add_mplan(mplan)
                         track_event_if_configured(
                             "PlanApproved",
                             {

src/backend/v3/models/messages.py

@@ -139,18 +139,14 @@ class ApprovalRequest(KernelBaseModel):
 
 @dataclass(slots=True)
 class AgentMessageResponse:
-    """Message sent to HumanAgent to request approval for a step."""
+    """Response message representing an agent's message."""
     plan_id: str
     agent: str
     content: str
     agent_type: AgentMessageType
     is_final: bool = False
-    m_plan_id: Optional[str] = None
-    user_id: Optional[str] = None
-    timestamp: Optional[str] = field(default_factory=lambda: str(time.time()))
-    raw_data: Optional[str] = None
-    steps: List[Any] = Field(default_factory=list) 
-    next_steps: List[Any] = Field(default_factory=list)
+    raw_data: str = None
+
 
 class WebsocketMessageType(str, Enum):
     """Types of WebSocket messages."""

src/backend/v3/models/models.py

@@ -24,8 +24,6 @@ class MStep(BaseModel):
 class MPlan(BaseModel):
     """model of a plan"""
     id: str = Field(default_factory=lambda: str(uuid.uuid4()))
-    data_type: Literal[DataType.m_plan] = Field(DataType.m_plan, Literal=True)
-    session_id: str = Field(default_factory=lambda: str(uuid.uuid4())) # db partition key
     user_id: str = ""
     team_id: str = ""
     plan_id: str = ""

src/frontend/src/models/agentMessage.tsx

@@ -35,7 +35,7 @@ export interface AgentMessageData {
  * Corresponds to the Python AgentMessageResponse class.
  */
 export interface AgentMessageResponse {
-    is_final?: boolean;
+
     /** Plan identifier */
     plan_id: string;
     /** Agent name or identifier */
@@ -44,18 +44,10 @@ export interface AgentMessageResponse {
     content: string;
     /** Type of agent (Human or AI) */
     agent_type: AgentMessageType;
-    /** Associated m_plan identifier */
-    m_plan_id?: string;
-    /** User identifier */
-    user_id?: string;
-    /** Timestamp when the message was created */
-    timestamp?: string;
+    is_final: boolean;
     /** Raw data associated with the message */
-    raw_data?: string;
-    /** Steps associated with the message */
-    steps?: any[];
-    /** Next steps associated with the message */
-    next_steps?: any[];
+    raw_data: string;
+
 }
 
 export interface FinalMessage {

src/frontend/src/pages/PlanPage.tsx

@@ -69,6 +69,7 @@ const PlanPage: React.FC = () => {
         console.log(is_final)
         console.log(agentMessageData)
         const agentMessageResponse = PlanDataService.createAgentMessageResponse(agentMessageData, planData, is_final);
+        console.log('📤 Persisting agent message:', agentMessageResponse);
         void apiService.sendAgentMessage(agentMessageResponse)
             .then(saved => {
                 console.log('[agent_message][persisted]', {
@@ -215,24 +216,28 @@ const PlanPage: React.FC = () => {
                 timestamp: Date.now(),
                 steps: [],   // intentionally always empty
                 next_steps: [],  // intentionally always empty
-                content: "🎉🎉 " + (finalMessage.content || ''),
+                content: "🎉🎉 " + (finalMessage.data?.content || ''),
                 raw_data: finalMessage || '',
             } as AgentMessageData;
 
 
             console.log('✅ Parsed final result message:', agentMessageData);
-            setStreamingMessageBuffer("");
-            setShowProcessingPlanSpinner(false);
-            setAgentMessages(prev => [...prev, agentMessageData]);
-            scrollToBottom();
-            // Persist the agent message
-            const is_final = finalMessage?.status === 'completed' || false;
-            if (finalMessage?.status === 'completed' && planData?.plan) {
-                planData.plan.overall_status = PlanStatus.COMPLETED;
-                setPlanData({ ...planData });
+            // we ignore the terminated message 
+            if (finalMessage?.data?.status === PlanStatus.COMPLETED) {
+                setStreamingMessageBuffer("");
+                setShowProcessingPlanSpinner(false);
+                setAgentMessages(prev => [...prev, agentMessageData]);
+                scrollToBottom();
+                // Persist the agent message
+                const is_final = true;
+                if (planData?.plan) {
+                    planData.plan.overall_status = PlanStatus.COMPLETED;
+                    setPlanData({ ...planData });
+                }
+
+                processAgentMessage(agentMessageData, planData, is_final);
             }
 
-            processAgentMessage(agentMessageData, planData, is_final);
 
         });
 
@@ -344,6 +349,9 @@ const PlanPage: React.FC = () => {
                 console.log("Plan data fetched:", planResult);
                 if (planResult?.plan?.overall_status === PlanStatus.IN_PROGRESS) {
                     setShowApprovalButtons(true);
+
+                } else {
+                    setWaitingForPlan(false);
                 }
                 if (planResult?.plan?.overall_status !== PlanStatus.COMPLETED) {
                     setContinueWithWebsocketFlow(true);

src/frontend/src/services/PlanDataService.tsx

@@ -197,17 +197,14 @@ export class PlanDataService {
       console.log("Invalid plan data provided to createAgentMessageResponse");
     }
     return {
-      is_final: is_final,
+
       plan_id: planData.plan.plan_id,
       agent: agentMessage.agent,
       content: agentMessage.content,
       agent_type: agentMessage.agent_type,
-      m_plan_id: planData.mplan?.id || undefined,
-      user_id: planData.plan.user_id,
-      timestamp: agentMessage.timestamp ? agentMessage.timestamp.toString() : new Date().getTime().toString(),
-      raw_data: agentMessage.raw_data,
-      steps: agentMessage.steps,
-      next_steps: agentMessage.next_steps
+      is_final: is_final,
+      raw_data: JSON.stringify(agentMessage.raw_data),
+
     };
   }