clarify message mapper, use sessionId in deleteMessageById and getMessageById

Signed-off-by: Julien Veyssier <julien-nc@posteo.net>

Author: julien-nc

lib/Controller/ChattyLLMController.php

@@ -402,10 +402,7 @@ public function getMessage(int $sessionId, int $messageId): JSONResponse {
 				return new JSONResponse(['error' => $this->l10n->t('Session not found')], Http::STATUS_NOT_FOUND);
 			}
 
-			$message = $this->messageMapper->getMessageById($messageId);
-			if ($message->getSessionId() !== $sessionId) {
-				return new JSONResponse(['error' => $this->l10n->t('Message not found')], Http::STATUS_NOT_FOUND);
-			}
+			$message = $this->messageMapper->getMessageById($sessionId, $messageId);
 
 			return new JSONResponse($message->jsonSerialize());
 		} catch (\OCP\DB\Exception $e) {
@@ -439,23 +436,19 @@ public function deleteMessage(int $messageId, int $sessionId): JSONResponse {
 			if (!$sessionExists) {
 				return new JSONResponse(['error' => $this->l10n->t('Session not found')], Http::STATUS_NOT_FOUND);
 			}
-			$message = $this->messageMapper->getMessageById($messageId);
-			// otherwise one can delete a message from other sessions
-			if ($message->getSessionId() !== $sessionId) {
-				return new JSONResponse(['error' => $this->l10n->t('Message not found')], Http::STATUS_NOT_FOUND);
-			}
+			$message = $this->messageMapper->getMessageById($sessionId, $messageId);
+			$ocpTaskId = $message->getOcpTaskId();
+
+			$this->messageMapper->deleteMessageById($sessionId, $messageId);
 
 			// delete the related task
-			$ocpTaskId = $message->getOcpTaskId();
 			if ($ocpTaskId !== 0) {
 				try {
 					$task = $this->taskProcessingManager->getTask($ocpTaskId);
 					$this->taskProcessingManager->deleteTask($task);
 				} catch (\OCP\TaskProcessing\Exception\Exception) {
 				}
 			}
-
-			$this->messageMapper->deleteMessageById($messageId);
 			return new JSONResponse();
 		} catch (\OCP\DB\Exception|\RuntimeException $e) {
 			$this->logger->warning('Failed to delete a chat message', ['exception' => $e]);
@@ -633,7 +626,7 @@ public function regenerateForSession(int $sessionId, int $messageId): JSONRespon
 		}
 
 		try {
-			$this->messageMapper->deleteMessageById($messageId);
+			$this->messageMapper->deleteMessageById($sessionId, $messageId);
 		} catch (\OCP\DB\Exception|\RuntimeException $e) {
 			$this->logger->warning('Failed to delete the last message', ['exception' => $e]);
 			return new JSONResponse(['error' => $this->l10n->t('Failed to delete the last message')], Http::STATUS_INTERNAL_SERVER_ERROR);

lib/Db/ChattyLLM/MessageMapper.php

@@ -12,6 +12,7 @@
 use OCP\AppFramework\Db\DoesNotExistException;
 use OCP\AppFramework\Db\MultipleObjectsReturnedException;
 use OCP\AppFramework\Db\QBMapper;
+use OCP\DB\Exception;
 use OCP\DB\QueryBuilder\IQueryBuilder;
 use OCP\IDBConnection;
 
@@ -99,17 +100,19 @@ public function getMessages(int $sessionId, int $cursor, int $limit): array {
 	}
 
 	/**
+	 * @param int $sessionId
 	 * @param integer $messageId
 	 * @return Message
-	 * @throws \OCP\DB\Exception
-	 * @throws MultipleObjectsReturnedException
 	 * @throws DoesNotExistException
+	 * @throws Exception
+	 * @throws MultipleObjectsReturnedException
 	 */
-	public function getMessageById(int $messageId): Message {
+	public function getMessageById(int $sessionId, int $messageId): Message {
 		$qb = $this->db->getQueryBuilder();
 		$qb->select(Message::$columns)
 			->from($this->getTableName())
-			->where($qb->expr()->eq('id', $qb->createPositionalParameter($messageId, IQueryBuilder::PARAM_INT)));
+			->where($qb->expr()->eq('id', $qb->createPositionalParameter($messageId, IQueryBuilder::PARAM_INT)))
+			->andWhere($qb->expr()->eq('session_id', $qb->createPositionalParameter($sessionId, IQueryBuilder::PARAM_INT)));
 
 		return $this->findEntity($qb);
 	}
@@ -147,15 +150,16 @@ public function deleteMessagesBySession(int $sessionId): void {
 	}
 
 	/**
+	 * @param int $sessionId
 	 * @param integer $messageId
-	 * @throws \OCP\DB\Exception
-	 * @throws \RuntimeException
 	 * @return void
+	 * @throws Exception
 	 */
-	public function deleteMessageById(int $messageId): void {
+	public function deleteMessageById(int $sessionId, int $messageId): void {
 		$qb = $this->db->getQueryBuilder();
 		$qb->delete($this->getTableName())
-			->where($qb->expr()->eq('id', $qb->createPositionalParameter($messageId, IQueryBuilder::PARAM_INT)));
+			->where($qb->expr()->eq('id', $qb->createPositionalParameter($messageId, IQueryBuilder::PARAM_INT)))
+			->andWhere($qb->expr()->eq('session_id', $qb->createPositionalParameter($sessionId, IQueryBuilder::PARAM_INT)));
 
 		$qb->executeStatement();
 	}

lib/Listener/ChattyLLMTaskListener.php

@@ -92,7 +92,7 @@ public function handle(Event $event): void {
 				// now we have the transcription of the user audio input
 				if (preg_match('/^chatty-llm:\d+:(\d+)$/', $customId, $matches)) {
 					$queryMessageId = (int)$matches[1];
-					$queryMessage = $this->messageMapper->getMessageById($queryMessageId);
+					$queryMessage = $this->messageMapper->getMessageById($sessionId, $queryMessageId);
 					$queryMessage->setContent(trim($taskOutput['input_transcript'] ?? ''));
 					$this->messageMapper->update($queryMessage);
 					// TODO update session title if it's the first message