Refactor team deletion and partition key handling

Simplifies team deletion logic in CosmosDBClient by removing redundant methods and partition key strategies. Adds session_id as a required field in TeamConfiguration and ensures session_id is generated and used consistently in team creation and deletion. Updates TeamService to use the new deletion method.

Author: Fr4nc3

src/backend/common/database/cosmosdb.py

@@ -129,8 +129,9 @@ async def update_item(self, item: BaseDataModel) -> None:
         try:
             # Convert to dictionary and handle datetime serialization
             document = item.model_dump()
-            document = json.loads(json.dumps(document, cls=DateTimeEncoder))
-
+            for key, value in list(document.items()):
+                if isinstance(value, datetime.datetime):
+                    document[key] = value.isoformat()
             await self.container.upsert_item(body=document)
         except Exception as e:
             self.logger.error("Failed to update item in CosmosDB: %s", str(e))
@@ -164,7 +165,7 @@ async def query_items(
             items = self.container.query_items(query=query, parameters=parameters)
             result_list = []
             async for item in items:
-                item["ts"] = item["_ts"]
+                # item["ts"] = item["_ts"]
                 try:
                     result_list.append(model_class.model_validate(item))
                 except Exception as validation_error:
@@ -352,88 +353,10 @@ async def delete_team(self, team_id: str) -> bool:
         try:
             # First find the team to get its document id and partition key
             team = await self.get_team(team_id)
+            print(team)
             if team:
-                # Use the session_id as partition key, or fall back to user_id if no session_id
-                partition_key = (
-                    team.session_id
-                    if hasattr(team, "session_id") and team.session_id
-                    else team.user_id
-                )
-                await self.container.delete_item(
-                    item=team.id, partition_key=partition_key
-                )
-                return True
-            return False
-        except Exception as e:
-            logging.exception(f"Failed to delete team from Cosmos DB: {e}")
-            return False
-
-    async def delete_team_by_id(self, id: str) -> bool:
-        """Delete a team configuration by its document id.
-
-        Args:
-            id: The document id of the team configuration to delete
-
-        Returns:
-            True if team was found and deleted, False otherwise
-        """
-        await self._ensure_initialized()
-
-        try:
-            # First find the team to get its partition key
-            team = await self.get_team_by_id(id)
-            if team:
-                # Debug logging
-                logging.info(f"Attempting to delete team {id}")
-                logging.info(f"Team user_id: {team.user_id}")
-                logging.info(
-                    f"Team session_id: {getattr(team, 'session_id', 'NOT_SET')}"
-                )
-
-                # Try different partition key strategies
-                partition_keys_to_try = []
-
-                # Strategy 1: Use session_id if it exists
-                if hasattr(team, "session_id") and team.session_id:
-                    partition_keys_to_try.append(team.session_id)
-                    logging.info(f"Will try session_id: {team.session_id}")
-
-                # Strategy 2: Use user_id as fallback
-                if team.user_id:
-                    partition_keys_to_try.append(team.user_id)
-                    logging.info(f"Will try user_id: {team.user_id}")
-
-                # Strategy 3: Use current context session_id
-                if self.session_id:
-                    partition_keys_to_try.append(self.session_id)
-                    logging.info(f"Will try context session_id: {self.session_id}")
-
-                # Strategy 4: Try null/empty partition key (for documents created without session_id)
-                partition_keys_to_try.extend([None, ""])
-                logging.info("Will try null and empty partition keys")
-
-                # Try each partition key until one works
-                for partition_key in partition_keys_to_try:
-                    try:
-                        logging.info(
-                            f"Attempting delete with partition key: {partition_key}"
-                        )
-                        await self.container.delete_item(
-                            item=id, partition_key=partition_key
-                        )
-                        logging.info(
-                            f"Successfully deleted team {id} with partition key: {partition_key}"
-                        )
-                        return True
-                    except Exception as e:
-                        logging.warning(
-                            f"Delete failed with partition key {partition_key}: {str(e)}"
-                        )
-                        continue
-
-                logging.error(f"All partition key strategies failed for team {id}")
-                return False
-            return False
+                await self.delete_item(item_id=team.id, partition_key=team.session_id)
+            return True
         except Exception as e:
             logging.exception(f"Failed to delete team from Cosmos DB: {e}")
             return False
@@ -487,154 +410,3 @@ async def update_team(self, team: TeamConfiguration) -> None:
             team: The TeamConfiguration to update
         """
         await self.update_item(team)
-
-    async def get_team(self, team_id: str) -> Optional[TeamConfiguration]:
-        """Retrieve a specific team configuration by team_id.
-
-        Args:
-            team_id: The team_id of the team configuration to retrieve
-
-        Returns:
-            TeamConfiguration object or None if not found
-        """
-        query = "SELECT * FROM c WHERE c.team_id=@team_id AND c.data_type=@data_type"
-        parameters = [
-            {"name": "@team_id", "value": team_id},
-            {"name": "@data_type", "value": "team_config"},
-        ]
-        teams = await self.query_items(query, parameters, TeamConfiguration)
-        return teams[0] if teams else None
-
-    async def get_team_by_id(self, id: str) -> Optional[TeamConfiguration]:
-        """Retrieve a specific team configuration by its document id.
-
-        Args:
-            id: The document id of the team configuration to retrieve
-
-        Returns:
-            TeamConfiguration object or None if not found
-        """
-        query = "SELECT * FROM c WHERE c.id=@id AND c.data_type=@data_type"
-        parameters = [
-            {"name": "@id", "value": id},
-            {"name": "@data_type", "value": "team_config"},
-        ]
-        teams = await self.query_items(query, parameters, TeamConfiguration)
-        return teams[0] if teams else None
-
-    async def get_all_teams_by_user(self, user_id: str) -> List[TeamConfiguration]:
-        """Retrieve all team configurations for a specific user.
-
-        Args:
-            user_id: The user_id to get team configurations for
-
-        Returns:
-            List of TeamConfiguration objects
-        """
-        query = "SELECT * FROM c WHERE c.user_id=@user_id AND c.data_type=@data_type ORDER BY c.created DESC"
-        parameters = [
-            {"name": "@user_id", "value": user_id},
-            {"name": "@data_type", "value": "team_config"},
-        ]
-        teams = await self.query_items(query, parameters, TeamConfiguration)
-        return teams
-
-    async def delete_team(self, team_id: str) -> bool:
-        """Delete a team configuration by team_id.
-
-        Args:
-            team_id: The team_id of the team configuration to delete
-
-        Returns:
-            True if team was found and deleted, False otherwise
-        """
-        await self.ensure_initialized()
-
-        try:
-            # First find the team to get its document id and partition key
-            team = await self.get_team(team_id)
-            if team:
-                # Use the session_id as partition key, or fall back to user_id if no session_id
-                partition_key = (
-                    team.session_id
-                    if hasattr(team, "session_id") and team.session_id
-                    else team.user_id
-                )
-                await self._container.delete_item(
-                    item=team.id, partition_key=partition_key
-                )
-                return True
-            return False
-        except Exception as e:
-            logging.exception(f"Failed to delete team from Cosmos DB: {e}")
-            return False
-
-    async def delete_team_by_id(self, id: str) -> bool:
-        """Delete a team configuration by its document id.
-
-        Args:
-            id: The document id of the team configuration to delete
-
-        Returns:
-            True if team was found and deleted, False otherwise
-        """
-        await self.ensure_initialized()
-
-        try:
-            # First find the team to get its partition key
-            team = await self.get_team_by_id(id)
-            if team:
-                # Debug logging
-                logging.info(f"Attempting to delete team {id}")
-                logging.info(f"Team user_id: {team.user_id}")
-                logging.info(
-                    f"Team session_id: {getattr(team, 'session_id', 'NOT_SET')}"
-                )
-
-                # Try different partition key strategies
-                partition_keys_to_try = []
-
-                # Strategy 1: Use session_id if it exists
-                if hasattr(team, "session_id") and team.session_id:
-                    partition_keys_to_try.append(team.session_id)
-                    logging.info(f"Will try session_id: {team.session_id}")
-
-                # Strategy 2: Use user_id as fallback
-                if team.user_id:
-                    partition_keys_to_try.append(team.user_id)
-                    logging.info(f"Will try user_id: {team.user_id}")
-
-                # Strategy 3: Use current context session_id
-                if self.session_id:
-                    partition_keys_to_try.append(self.session_id)
-                    logging.info(f"Will try context session_id: {self.session_id}")
-
-                # Strategy 4: Try null/empty partition key (for documents created without session_id)
-                partition_keys_to_try.extend([None, ""])
-                logging.info("Will try null and empty partition keys")
-
-                # Try each partition key until one works
-                for partition_key in partition_keys_to_try:
-                    try:
-                        logging.info(
-                            f"Attempting delete with partition key: {partition_key}"
-                        )
-                        await self._container.delete_item(
-                            item=id, partition_key=partition_key
-                        )
-                        logging.info(
-                            f"Successfully deleted team {id} with partition key: {partition_key}"
-                        )
-                        return True
-                    except Exception as e:
-                        logging.warning(
-                            f"Delete failed with partition key {partition_key}: {str(e)}"
-                        )
-                        continue
-
-                logging.error(f"All partition key strategies failed for team {id}")
-                return False
-            return False
-        except Exception as e:
-            logging.exception(f"Failed to delete team from Cosmos DB: {e}")
-            return False

src/backend/common/database/database_base.py

@@ -157,11 +157,6 @@ async def delete_team(self, team_id: str) -> bool:
         """Delete a team configuration by team_id and return True if deleted."""
         pass
 
-    @abstractmethod
-    async def delete_team_by_id(self, id: str) -> bool:
-        """Delete a team configuration by internal id and return True if deleted."""
-        pass
-
     # Data Management Operations
     @abstractmethod
     async def get_data_by_type(self, data_type: str) -> List[BaseDataModel]:

src/backend/common/models/messages_kernel.py

@@ -178,6 +178,7 @@ class TeamConfiguration(BaseDataModel):
 
     team_id: str
     data_type: Literal["team_config"] = Field("team_config", Literal=True)
+    session_id: str  # Partition key
     name: str
     status: str
     created: str

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

@@ -67,6 +67,7 @@ async def validate_and_parse_team_config(
 
             # Generate unique IDs and timestamps
             unique_team_id = str(uuid.uuid4())
+            session_id = str(uuid.uuid4())
             current_timestamp = datetime.now(timezone.utc).isoformat()
 
             # Validate agents array exists and is not empty
@@ -104,6 +105,7 @@ async def validate_and_parse_team_config(
             # Create team configuration
             team_config = TeamConfiguration(
                 id=unique_team_id,  # Use generated GUID
+                session_id=session_id,
                 team_id=unique_team_id,  # Use generated GUID
                 name=json_data["name"],
                 status=json_data["status"],
@@ -253,17 +255,7 @@ async def delete_team_configuration(self, team_id: str, user_id: str) -> bool:
         """
         try:
             # First, verify the configuration exists and belongs to the user
-            team_config = await self.memory_context.get_team_by_id(team_id)
-
-            if team_config is None:
-                self.logger.warning(
-                    "Team configuration not found for deletion: %s", team_id
-                )
-                return False
-
-            # Delete the configuration using the specific delete_team_by_id method
-            success = await self.memory_context.delete_team_by_id(team_id)
-
+            success = await self.memory_context.delete_team(team_id)
             if success:
                 self.logger.info("Successfully deleted team configuration: %s", team_id)