Handle multiple messages for one taskId

lib/Controller/ChattyLLMController.php

@@ -581,7 +581,6 @@ public function regenerateForSession(int $sessionId, int $messageId): JSONRespon
 	 * @param int $taskId The message generation task ID
 	 * @param int $sessionId The chat session ID
 	 * @return JSONResponse<Http::STATUS_OK, AssistantChatAgencyMessage, array{}>|JSONResponse<Http::STATUS_EXPECTATION_FAILED, array{task_status: int, slow_pickup: bool, task_output?: array<string, list<numeric|string>|numeric|string>|null}, array{}>|JSONResponse<Http::STATUS_BAD_REQUEST, array{error: string, task_status?: int, userFacingErrorMessage?: ?string}, array{}>|JSONResponse<Http::STATUS_INTERNAL_SERVER_ERROR|Http::STATUS_UNAUTHORIZED|Http::STATUS_NOT_FOUND, array{error: string}, array{}>
-	 * @throws MultipleObjectsReturnedException
 	 * @throws \OCP\DB\Exception
 	 *
 	 * 200: The task was successful, a message has been generated

lib/Db/ChattyLLM/MessageMapper.php

@@ -150,17 +150,24 @@ public function getMessageById(int $sessionId, int $messageId): Message {
 	 * @param int $ocpTaskId
 	 * @return Message
 	 * @throws DoesNotExistException
-	 * @throws MultipleObjectsReturnedException
 	 * @throws \OCP\DB\Exception
 	 */
 	public function getMessageByTaskId(int $sessionId, int $ocpTaskId): Message {
 		$qb = $this->db->getQueryBuilder();
 		$qb->select(Message::$columns)
 			->from($this->getTableName())
 			->where($qb->expr()->eq('session_id', $qb->createPositionalParameter($sessionId, IQueryBuilder::PARAM_INT)))
-			->andWhere($qb->expr()->eq('ocp_task_id', $qb->createPositionalParameter($ocpTaskId, IQueryBuilder::PARAM_INT)));
+			->andWhere($qb->expr()->eq('ocp_task_id', $qb->createPositionalParameter($ocpTaskId, IQueryBuilder::PARAM_INT)))
+			// concurrent task processing executors may run the same task twice,
+			// which results in duplicate rows with the same ocp_task_id
+			->orderBy('id', 'ASC')
+			->setMaxResults(1);
 
-		return $this->findEntity($qb);
+		$entities = $this->findEntities($qb);
+		if ($entities === []) {
+			throw new DoesNotExistException('Assistant chat message with session_id ' . $sessionId . ' and ocp_task_id ' . $ocpTaskId . ' does not exist');
+		}
+		return $entities[0];
 	}
 
 	/**

lib/Listener/ChattyLLMTaskListener.php

@@ -15,6 +15,7 @@
 use OCA\Assistant\Db\ChattyLLM\SessionMapper;
 use OCA\Assistant\Service\NotificationService;
 use OCA\Assistant\Service\TaskProcessingService;
+use OCP\AppFramework\Db\DoesNotExistException;
 use OCP\EventDispatcher\Event;
 use OCP\EventDispatcher\IEventListener;
 use OCP\TaskProcessing\Events\TaskFailedEvent;
@@ -76,6 +77,15 @@ public function handle(Event $event): void {
 		if (preg_match('/^chatty-llm:(\d+)/', $customId, $matches)) {
 			$sessionId = (int)$matches[1];
 
+			// concurrent task processing executors can process the same task twice and
+			// dispatch this event twice, in which case the message was already stored
+			try {
+				$this->messageMapper->getMessageByTaskId($sessionId, $task->getId());
+				return;
+			} catch (DoesNotExistException $e) {
+				// the message was not stored yet, continue with the normal handling
+			}
+
 			$isAgency = class_exists('OCP\\TaskProcessing\\TaskTypes\\ContextAgentInteraction')
 				&& $taskTypeId === \OCP\TaskProcessing\TaskTypes\ContextAgentInteraction::ID;
 			$isAudioChat = class_exists('OCP\\TaskProcessing\\TaskTypes\\AudioToAudioChat')

tests/unit/Listener/ChattyLLMTaskListenerTest.php

@@ -0,0 +1,141 @@
+<?php
+
+declare(strict_types=1);
+
+/**
+ * SPDX-FileCopyrightText: 2026 Nextcloud GmbH and Nextcloud contributors
+ * SPDX-License-Identifier: AGPL-3.0-or-later
+ */
+
+namespace OCA\Assistant\Tests\Unit\Listener;
+
+use OCA\Assistant\Db\ChattyLLM\Message;
+use OCA\Assistant\Db\ChattyLLM\MessageMapper;
+use OCA\Assistant\Db\ChattyLLM\Session;
+use OCA\Assistant\Db\ChattyLLM\SessionMapper;
+use OCA\Assistant\Listener\ChattyLLMTaskListener;
+use OCA\Assistant\Service\NotificationService;
+use OCA\Assistant\Service\TaskProcessingService;
+use OCP\AppFramework\Db\DoesNotExistException;
+use OCP\TaskProcessing\Events\TaskSuccessfulEvent;
+use OCP\TaskProcessing\Task;
+use OCP\TaskProcessing\TaskTypes\TextToTextChat;
+use PHPUnit\Framework\MockObject\MockObject;
+use Psr\Log\LoggerInterface;
+use Test\TestCase;
+
+/**
+ * Concurrent task processing executors can run the same task twice,
+ * which must not store the chat message twice.
+ */
+class ChattyLLMTaskListenerTest extends TestCase {
+
+	private const SESSION_ID = 23;
+	private const TASK_ID = 91;
+
+	private MessageMapper&MockObject $messageMapper;
+	private SessionMapper&MockObject $sessionMapper;
+	private ChattyLLMTaskListener $listener;
+
+	protected function setUp(): void {
+		parent::setUp();
+
+		$this->messageMapper = $this->createMock(MessageMapper::class);
+		$this->sessionMapper = $this->createMock(SessionMapper::class);
+
+		$session = new Session();
+		$session->setAssignmentId(null);
+		$this->sessionMapper->method('getUserSession')->willReturn($session);
+
+		$this->listener = new ChattyLLMTaskListener(
+			$this->messageMapper,
+			$this->sessionMapper,
+			$this->createMock(TaskProcessingService::class),
+			$this->createMock(LoggerInterface::class),
+			$this->createMock(NotificationService::class),
+		);
+	}
+
+	private function makeTask(): Task {
+		$task = new Task(
+			TextToTextChat::ID,
+			[
+				'input' => 'What is a cat?',
+				'system_prompt' => '',
+				'history' => [],
+			],
+			'assistant:chatty-llm',
+			'admin',
+			'chatty-llm:' . self::SESSION_ID,
+		);
+		$task->setId(self::TASK_ID);
+		$task->setOutput(['output' => 'A cat is an animal.']);
+		return $task;
+	}
+
+	private function handleTask(): void {
+		$event = new TaskSuccessfulEvent($this->makeTask());
+		$this->listener->handle($event);
+	}
+
+	public function testMessageIsStored(): void {
+		$this->messageMapper->method('getMessageByTaskId')
+			->with(self::SESSION_ID, self::TASK_ID)
+			->willThrowException(new DoesNotExistException('Message not found'));
+		$storedMessage = null;
+		$this->messageMapper->expects($this->once())
+			->method('insert')
+			->willReturnCallback(static function (Message $message) use (&$storedMessage): Message {
+				$storedMessage = $message;
+				$message->setId(3);
+				return $message;
+			});
+		$this->sessionMapper->expects($this->once())->method('update');
+
+		$this->handleTask();
+
+		$this->assertNotNull($storedMessage);
+		$this->assertSame(self::SESSION_ID, $storedMessage->getSessionId());
+		$this->assertSame(self::TASK_ID, $storedMessage->getOcpTaskId());
+		$this->assertSame(Message::ROLE_ASSISTANT, $storedMessage->getRole());
+		$this->assertSame('A cat is an animal.', $storedMessage->getContent());
+	}
+
+	public function testMessageIsStoredOnlyOnce(): void {
+		$storedMessage = new Message();
+		$storedMessage->setSessionId(self::SESSION_ID);
+		$storedMessage->setOcpTaskId(self::TASK_ID);
+		$this->messageMapper->method('getMessageByTaskId')
+			->with(self::SESSION_ID, self::TASK_ID)
+			->willReturn($storedMessage);
+		$this->messageMapper->expects($this->never())->method('insert');
+		$this->sessionMapper->expects($this->never())->method('update');
+
+		$this->handleTask();
+	}
+
+	public function testTitleTasksAreUntouched(): void {
+		$task = new Task(
+			TextToTextChat::ID,
+			[
+				'input' => 'Generate a title',
+				'system_prompt' => '',
+				'history' => [],
+			],
+			'assistant:chatty-llm',
+			'admin',
+			'chatty-title:' . self::SESSION_ID,
+		);
+		$task->setId(self::TASK_ID);
+		$task->setOutput(['output' => 'Cats']);
+
+		$event = new TaskSuccessfulEvent($task);
+		$this->messageMapper->expects($this->never())->method('getMessageByTaskId');
+		$this->messageMapper->expects($this->never())->method('insert');
+		$this->sessionMapper->expects($this->once())
+			->method('updateSessionTitle')
+			->with('admin', self::SESSION_ID, 'Cats');
+
+		$this->listener->handle($event);
+	}
+}