Add agent message support to database and services

Introduces methods for adding, updating, and retrieving agent messages in the database layer (CosmosDBClient and DatabaseBase). Implements conversion and persistence logic for agent messages in PlanService. Also adds a session_id field to MPlan for partitioning.

Author: Fr4nc3

src/backend/common/database/cosmosdb.py

@@ -23,6 +23,7 @@
 
 from .database_base import DatabaseBase
 from ..models.messages_kernel import (
+    AgentMessageData,
     BaseDataModel,
     Plan,
     Step,
@@ -479,4 +480,23 @@ async def get_mplan(self, plan_id: str) -> Optional[messages.MPlan]:
             {"name": "@data_type", "value": DataType.m_plan},
         ]
         results = await self.query_items(query, parameters, messages.MPlan)
-        return results[0] if results else None
+        return results[0] if results else None
+    
+
+    async def add_agent_message(self, message: messages.AgentMessageResponse) -> None:
+        """Add an agent message to the database."""
+        await self.add_item(message)
+
+    async def update_agent_message(self, message: messages.AgentMessageResponse) -> None:
+        """Update an agent message in the database."""
+        await self.update_item(message)
+
+    async def get_agent_messages(self, plan_id: str) -> List[messages.AgentMessageResponse]:
+        """Retrieve an agent message by message_id."""
+        query = "SELECT * FROM c WHERE c.plan_id=@plan_id AND c.data_type=@data_type"
+        parameters = [
+            {"name": "@plan_id", "value": plan_id},
+            {"name": "@data_type", "value": DataType.m_plan_message},
+        ]
+
+        return await self.query_items(query, parameters, messages.AgentMessageResponse)

src/backend/common/database/database_base.py

@@ -4,6 +4,7 @@
 from typing import Any, Dict, List, Optional, Type
 import v3.models.messages as messages
 from ..models.messages_kernel import (
+    AgentMessageData,
     BaseDataModel,
     Plan,
     Step,
@@ -213,5 +214,19 @@ async def update_mplan(self, mplan: messages.MPlan) -> None:
 
     @abstractmethod
     async def get_mplan(self, plan_id: str) -> Optional[messages.MPlan]:
-        """Retrieve a mplan configuration by mplan_id."""
+        """Retrieve a mplan configuration by plan_id."""
+        pass
+    
+    @abstractmethod
+    async def add_agent_message(self, message: AgentMessageData) -> None:
+        pass
+
+    @abstractmethod
+    async def update_agent_message(self, message: AgentMessageResponse) -> None:
+        """Update an agent message in the database."""
+        pass
+
+    @abstractmethod
+    async def get_agent_messages(self, plan_id: str) -> Optional[AgentMessageResponse]:
+        """Retrieve an agent message by message_id."""
         pass

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

@@ -13,8 +13,6 @@
 import uuid
 from semantic_kernel.kernel_pydantic import Field
 
-class MPlanExpanded(MPlan):
-    session_id: str = Field(default_factory=lambda: str(uuid.uuid4()))
 
 logger = logging.getLogger(__name__)
 def build_agent_message_from_user_clarification(
@@ -39,6 +37,70 @@ def build_agent_message_from_user_clarification(
         next_steps = []      # intentionally empty
     )
 
+def build_agent_message_from_agent_message_response(
+    agent_response: messages.AgentMessageResponse,
+    user_id: str,
+) -> AgentMessageData:
+    """
+    Convert a messages.AgentMessageResponse into common.models.messages_kernel.AgentMessageData.
+    This is defensive: it tolerates missing fields and different timestamp formats.
+    """
+    # Robust timestamp parsing (accepts seconds or ms or missing)
+
+
+    # Raw data serialization
+    raw = getattr(agent_response, "raw_data", None)
+    try:
+        if raw is None:
+            # try asdict if it's a dataclass-like
+            try:
+                raw_str = json.dumps(asdict(agent_response))
+            except Exception:
+                raw_str = json.dumps({k: getattr(agent_response, k) for k in dir(agent_response) if not k.startswith("_")})
+        elif isinstance(raw, (dict, list)):
+            raw_str = json.dumps(raw)
+        else:
+            raw_str = str(raw)
+    except Exception:
+        raw_str = json.dumps({"raw": str(raw)})
+
+    # Steps / next_steps defaulting
+    steps = getattr(agent_response, "steps", []) or []
+    next_steps = getattr(agent_response, "next_steps", []) or []
+
+    # Agent name and type
+    agent_name = getattr(agent_response, "agent", "") or getattr(agent_response, "agent_name", "") or getattr(agent_response, "source", "")
+    # Try to infer agent_type, fallback to AI_AGENT
+    agent_type_raw = getattr(agent_response, "agent_type", None)
+    if isinstance(agent_type_raw, AgentMessageType):
+        agent_type = agent_type_raw
+    else:
+        # Normalize common strings
+        agent_type_str = str(agent_type_raw or "").lower()
+        if "human" in agent_type_str:
+            agent_type = AgentMessageType.HUMAN_AGENT
+        else:
+            agent_type = AgentMessageType.AI_AGENT
+
+    # Content
+    content = getattr(agent_response, "content", "") or getattr(agent_response, "text", "") or ""
+
+    # plan_id / user_id fallback
+    plan_id_val = getattr(agent_response, "plan_id", "") or ""
+    user_id_val = getattr(agent_response, "user_id", "") or user_id
+
+    return AgentMessageData(
+        plan_id=plan_id_val,
+        user_id=user_id_val,
+        m_plan_id=getattr(agent_response, "m_plan_id", ""),
+        agent=agent_name,
+        agent_type=agent_type,
+        content=content,
+        raw_data=raw_str,
+        steps=list(steps),
+        next_steps=list(next_steps),
+    )
+
 
 class PlanService:
 
@@ -118,7 +180,18 @@ async def handle_agent_messages(agent_message: messages.AgentMessageResponse, us
         Raises:
             ValueError on invalid state
         """
-        return True
+        try:
+            agent_msg = build_agent_message_from_agent_message_response(agent_message, user_id)
+
+            # Persist if your database layer supports it.
+            # Look for or implement something like: memory_store.add_agent_message(agent_msg)
+            memory_store = await DatabaseFactory.get_database(user_id=user_id)
+            await memory_store.add_agent_message(agent_msg)
+            return True
+        except Exception as e:
+            logger.exception("Failed to handle human clarification -> agent message: %s", e)
+            return False
+
 
     @staticmethod
     async def handle_human_clarification(human_feedback: messages.UserClarificationResponse, user_id: str) -> bool:
@@ -141,14 +214,7 @@ async def handle_human_clarification(human_feedback: messages.UserClarificationR
             # Persist if your database layer supports it.
             # Look for or implement something like: memory_store.add_agent_message(agent_msg)
             memory_store = await DatabaseFactory.get_database(user_id=user_id)
-            if hasattr(memory_store, "add_agent_message"):
-                await memory_store.add_agent_message(agent_msg)
-            else:
-                # Fallback: log or ignore if persistence not yet implemented
-                logging.debug("add_agent_message not implemented; skipping persistence")
-
-            # Optionally emit over websocket if you have a broadcaster:
-            # await websocket_manager.broadcast(WebsocketMessageType.AGENT_MESSAGE, agent_msg.model_dump())
+            await memory_store.add_agent_message(agent_msg)
 
             return True
         except Exception as e:

src/backend/v3/models/models.py

@@ -25,6 +25,7 @@ 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 = ""