Refactor session and team management in backend

Removed session-related methods and fields from database base and CosmosDB client, and updated data models to eliminate redundant session_id fields. Added and refactored current team deletion logic, and improved team selection handling to return the selected team. Updated API and service layers to use new team management methods and fixed related bugs.

Author: Fr4nc3

src/backend/common/database/cosmosdb.py

@@ -190,29 +190,6 @@ async def delete_item(self, item_id: str, partition_key: str) -> None:
             self.logger.error("Failed to delete item from CosmosDB: %s", str(e))
             raise
 
-    # Session Operations
-    async def add_session(self, session: Session) -> None:
-        """Add a session to CosmosDB."""
-        await self.add_item(session)
-
-    async def get_session(self, session_id: str) -> Optional[Session]:
-        """Retrieve a session by session_id."""
-        query = "SELECT * FROM c WHERE c.id=@id AND c.data_type=@data_type"
-        parameters = [
-            {"name": "@id", "value": session_id},
-            {"name": "@data_type", "value": DataType.session},
-        ]
-        results = await self.query_items(query, parameters, Session)
-        return results[0] if results else None
-
-    async def get_all_sessions(self) -> List[Session]:
-        """Retrieve all sessions for the user."""
-        query = "SELECT * FROM c WHERE c.user_id=@user_id AND c.data_type=@data_type"
-        parameters = [
-            {"name": "@user_id", "value": self.user_id},
-            {"name": "@data_type", "value": DataType.session},
-        ]
-        return await self.query_items(query, parameters, Session)
 
     # Plan Operations
     async def add_plan(self, plan: Plan) -> None:
@@ -470,12 +447,32 @@ async def get_current_team(self, user_id: str) -> Optional[UserCurrentTeam]:
         teams = await self.query_items(query, parameters, UserCurrentTeam)
         return teams[0] if teams else None
 
+
+
+    async def delete_current_team(self, user_id: str) -> bool:
+        """Delete the current team for a user."""
+        query = "SELECT c.id, c.session_id FROM c WHERE c.user_id=@user_id AND c.data_type=@data_type"
+
+        params = [
+            {"name": "@user_id", "value": user_id},
+            {"name": "@data_type", "value": DataType.user_current_team},
+        ]
+        items = self.container.query_items(query=query, parameters=params)
+        print("Items to delete:", items)
+        if items:
+            async for doc in items:
+                try:
+                    await self.container.delete_item(doc["id"], partition_key=doc["session_id"])
+                except Exception as e:
+                    self.logger.warning("Failed deleting current team doc %s: %s", doc.get("id"), e)
+
+        return True
+
     async def set_current_team(self, current_team: UserCurrentTeam) -> None:
         """Set the current team for a user."""
         await self._ensure_initialized()
         await self.add_item(current_team)
 
-
     async def update_current_team(self, current_team: UserCurrentTeam) -> None:
         """Update the current team for a user."""
         await self._ensure_initialized()

src/backend/common/database/database_base.py

@@ -5,7 +5,6 @@
 
 from ..models.messages_kernel import (
     BaseDataModel,
-    Session,
     Plan,
     Step,
     TeamConfiguration,
@@ -59,21 +58,6 @@ async def delete_item(self, item_id: str, partition_key: str) -> None:
         """Delete an item from the database."""
         pass
 
-    # Session Operations
-    @abstractmethod
-    async def add_session(self, session: Session) -> None:
-        """Add a session to the database."""
-        pass
-
-    @abstractmethod
-    async def get_session(self, session_id: str) -> Optional[Session]:
-        """Retrieve a session by session_id."""
-        pass
-
-    @abstractmethod
-    async def get_all_sessions(self) -> List[Session]:
-        """Retrieve all sessions for the user."""
-        pass
 
     # Plan Operations
     @abstractmethod
@@ -86,11 +70,6 @@ async def update_plan(self, plan: Plan) -> None:
         """Update a plan in the database."""
         pass
 
-    @abstractmethod
-    async def get_plan_by_session(self, session_id: str) -> Optional[Plan]:
-        """Retrieve a plan by session_id."""
-        pass
-
     @abstractmethod
     async def get_plan_by_plan_id(self, plan_id: str) -> Optional[Plan]:
         """Retrieve a plan by plan_id."""
@@ -207,6 +186,11 @@ async def get_current_team(self, user_id: str) -> Optional[UserCurrentTeam]:
         """Retrieve the current team for a user."""
         pass
 
+    @abstractmethod
+    async def delete_current_team(self, user_id: str) -> Optional[UserCurrentTeam]:
+        """Retrieve the current team for a user."""
+        pass
+
     @abstractmethod
     async def set_current_team(self, current_team: UserCurrentTeam) -> None:
         pass

src/backend/common/models/messages_kernel.py

@@ -78,6 +78,7 @@ class BaseDataModel(KernelBaseModel):
     """Base data model with common fields."""
 
     id: str = Field(default_factory=lambda: str(uuid.uuid4()))
+    session_id: str = Field(default_factory=lambda: str(uuid.uuid4()))
     timestamp: Optional[datetime] = Field(
         default_factory=lambda: datetime.now(timezone.utc)
     )
@@ -87,8 +88,6 @@ class AgentMessage(BaseDataModel):
     """Base class for messages sent between agents."""
 
     data_type: Literal[DataType.agent_message] = Field(DataType.agent_message, Literal=True)
-    session_id: str
-    user_id: str
     plan_id: str
     content: str
     source: str
@@ -118,7 +117,6 @@ class Plan(BaseDataModel):
 
     data_type: Literal[DataType.plan] = Field(DataType.plan, Literal=True)
     plan_id: str
-    session_id: str
     user_id: str
     initial_goal: str
     overall_status: PlanStatus = PlanStatus.in_progress
@@ -135,7 +133,6 @@ class Step(BaseDataModel):
 
     data_type: Literal[DataType.step] = Field(DataType.step, Literal=True)
     plan_id: str
-    session_id: str  # Partition key
     user_id: str
     action: str
     agent: AgentType
@@ -146,7 +143,7 @@ class Step(BaseDataModel):
     updated_action: Optional[str] = None
 
 
-class TeamSelectionRequest(KernelBaseModel):
+class TeamSelectionRequest(BaseDataModel):
     """Request model for team selection."""
 
     team_id: str

src/backend/v3/api/router.py

@@ -117,9 +117,12 @@ async def init_team(
         team_service = TeamService(memory_store)
         user_current_team = await memory_store.get_current_team(user_id=user_id)
         if not user_current_team:
-            await team_service.handle_team_selection(
+            print("User has no current team, setting to default:", init_team_id)
+            user_current_team =  await team_service.handle_team_selection(
                 user_id=user_id, team_id=init_team_id
             )
+            if user_current_team:
+                init_team_id = user_current_team.team_id
         else:
             init_team_id = user_current_team.team_id
         # Verify the team exists and user has access to it
@@ -896,7 +899,7 @@ async def select_team(selection: TeamSelectionRequest, request: Request):
         team_configuration = await team_service.get_team_configuration(
             selection.team_id, user_id
         )
-        if team_config is None:
+        if team_configuration is None: # ensure that id is valid
             raise HTTPException(
                 status_code=404,
                 detail=f"Team configuration '{selection.team_id}' not found or access denied",

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

@@ -235,15 +235,15 @@ async def delete_user_current_team(self, user_id: str) -> bool:
             True if successful, False otherwise
         """
         try:
-            await self.memory_context.delete_user_current_team(user_id)
+            await self.memory_context.delete_current_team(user_id)
             self.logger.info("Successfully deleted current team for user %s", user_id)
             return True
 
         except Exception as e:
             self.logger.error("Error deleting current team: %s", str(e))
             return False
 
-    async def handle_team_selection(self, user_id: str, team_id: str) -> bool:
+    async def handle_team_selection(self, user_id: str, team_id: str) -> UserCurrentTeam:
         """
         Set a default team for a user.
 
@@ -254,25 +254,21 @@ async def handle_team_selection(self, user_id: str, team_id: str) -> bool:
         Returns:
             True if successful, False otherwise
         """
+        print("Handling team selection for user:", user_id, "team:", team_id)
         try:
-            current_team = await self.memory_context.get_current_team(user_id)
-
-            if current_team is None:
-                current_team = UserCurrentTeam(user_id=user_id, team_id=team_id)
-                await self.memory_context.set_current_team(current_team)
-                return True
-            else:
-                current_team.team_id = team_id
-                await self.memory_context.update_current_team(current_team)
-                return True
+            await self.memory_context.delete_current_team(user_id)
+            current_team = UserCurrentTeam(
+                user_id=user_id,
+                team_id=team_id,
+            )
+            await self.memory_context.set_current_team(current_team)
+            return current_team
 
         except Exception as e:
             self.logger.error("Error setting default team: %s", str(e))
-            return False
+            return None
 
-    async def get_all_team_configurations(
-        self
-    ) -> List[TeamConfiguration]:
+    async def get_all_team_configurations(self) -> List[TeamConfiguration]:
         """
         Retrieve all team configurations for a user.