diff --git a/appinfo/info.xml b/appinfo/info.xml
index cc176b2b..24a2ec6a 100644
--- a/appinfo/info.xml
+++ b/appinfo/info.xml
@@ -62,7 +62,7 @@ Known providers:
 
 More details on how to set this up in the [admin docs](https://docs.nextcloud.com/server/latest/admin_manual/ai/index.html)
 ]]>	</description>
-	<version>2.3.0</version>
+	<version>2.4.0</version>
 	<licence>agpl</licence>
 	<author>Julien Veyssier</author>
 	<namespace>Assistant</namespace>
diff --git a/lib/Controller/ChattyLLMController.php b/lib/Controller/ChattyLLMController.php
index b234c9b0..2288a33f 100644
--- a/lib/Controller/ChattyLLMController.php
+++ b/lib/Controller/ChattyLLMController.php
@@ -420,11 +420,7 @@ public function checkMessageGenerationTask(int $taskId, int $sessionId): JSONRes
 		}
 		if ($task->getStatus() === Task::STATUS_SUCCESSFUL) {
 			try {
-				$message = new Message();
-				$message->setSessionId($sessionId);
-				$message->setRole('assistant');
-				$message->setContent(trim($task->getOutput()['output'] ?? ''));
-				$message->setTimestamp(time());
+				$message = $this->messageMapper->getMessageByTaskId($sessionId, $taskId);
 				$jsonMessage = $message->jsonSerialize();
 				$session = $this->sessionMapper->getUserSession($this->userId, $sessionId);
 				$jsonMessage['sessionAgencyPendingActions'] = $session->getAgencyPendingActions();
diff --git a/lib/Db/ChattyLLM/Message.php b/lib/Db/ChattyLLM/Message.php
index 1dd6bf8b..faad0a0f 100644
--- a/lib/Db/ChattyLLM/Message.php
+++ b/lib/Db/ChattyLLM/Message.php
@@ -21,6 +21,8 @@
  * @method \void setContent(string $content)
  * @method \int getTimestamp()
  * @method \void setTimestamp(int $timestamp)
+ * @method \int getOcpTaskId()
+ * @method \void setOcpTaskId(int $ocpTaskId)
  */
 class Message extends Entity implements \JsonSerializable {
 	/** @var int */
@@ -31,6 +33,8 @@ class Message extends Entity implements \JsonSerializable {
 	protected $content;
 	/** @var int */
 	protected $timestamp;
+	/** @var int */
+	protected $ocpTaskId;
 
 	public static $columns = [
 		'id',
@@ -38,6 +42,7 @@ class Message extends Entity implements \JsonSerializable {
 		'role',
 		'content',
 		'timestamp',
+		'ocp_task_id',
 	];
 	public static $fields = [
 		'id',
@@ -45,6 +50,7 @@ class Message extends Entity implements \JsonSerializable {
 		'role',
 		'content',
 		'timestamp',
+		'ocpTaskId',
 	];
 
 	public function __construct() {
@@ -52,6 +58,7 @@ public function __construct() {
 		$this->addType('role', Types::STRING);
 		$this->addType('content', Types::STRING);
 		$this->addType('timestamp', Types::INTEGER);
+		$this->addType('ocp_task_id', Types::INTEGER);
 	}
 
 	#[\ReturnTypeWillChange]
@@ -62,6 +69,7 @@ public function jsonSerialize() {
 			'role' => $this->role,
 			'content' => $this->content,
 			'timestamp' => $this->timestamp,
+			'ocp_task_id' => $this->ocpTaskId,
 		];
 	}
 }
diff --git a/lib/Db/ChattyLLM/MessageMapper.php b/lib/Db/ChattyLLM/MessageMapper.php
index f3023f50..b2c619a7 100644
--- a/lib/Db/ChattyLLM/MessageMapper.php
+++ b/lib/Db/ChattyLLM/MessageMapper.php
@@ -101,6 +101,24 @@ public function getMessageById(int $messageId): Message {
 		return $this->findEntity($qb);
 	}
 
+	/**
+	 * @param int $sessionId
+	 * @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)));
+
+		return $this->findEntity($qb);
+	}
+
 	/**
 	 * @param int $sessionId
 	 * @throws \OCP\DB\Exception
diff --git a/lib/Listener/ChattyLLMTaskListener.php b/lib/Listener/ChattyLLMTaskListener.php
index 7c39aeb7..4cb46b62 100644
--- a/lib/Listener/ChattyLLMTaskListener.php
+++ b/lib/Listener/ChattyLLMTaskListener.php
@@ -57,6 +57,7 @@ public function handle(Event $event): void {
 
 			$message = new Message();
 			$message->setSessionId($sessionId);
+			$message->setOcpTaskId($task->getId());
 			$message->setRole('assistant');
 			$message->setContent(trim($task->getOutput()['output'] ?? ''));
 			$message->setTimestamp(time());
diff --git a/lib/Migration/Version020400Date20250203125359.php b/lib/Migration/Version020400Date20250203125359.php
new file mode 100644
index 00000000..311b8ceb
--- /dev/null
+++ b/lib/Migration/Version020400Date20250203125359.php
@@ -0,0 +1,45 @@
+<?php
+
+declare(strict_types=1);
+/**
+ * SPDX-FileCopyrightText: 2025 Nextcloud GmbH and Nextcloud contributors
+ * SPDX-License-Identifier: AGPL-3.0-or-later
+ */
+
+namespace OCA\Assistant\Migration;
+
+use Closure;
+use OCP\DB\ISchemaWrapper;
+use OCP\DB\Types;
+use OCP\Migration\IOutput;
+use OCP\Migration\SimpleMigrationStep;
+
+class Version020400Date20250203125359 extends SimpleMigrationStep {
+
+	/**
+	 * @param IOutput $output
+	 * @param Closure(): ISchemaWrapper $schemaClosure
+	 * @param array $options
+	 * @return null|ISchemaWrapper
+	 */
+	public function changeSchema(IOutput $output, Closure $schemaClosure, array $options): ?ISchemaWrapper {
+		/** @var ISchemaWrapper $schema */
+		$schema = $schemaClosure();
+		$schemaChanged = false;
+
+		if ($schema->hasTable('assistant_chat_msgs')) {
+			$table = $schema->getTable('assistant_chat_msgs');
+			if (!$table->hasColumn('ocp_task_id')) {
+				$table->addColumn('ocp_task_id', Types::BIGINT, [
+					'notnull' => true,
+					'default' => 0,
+					'unsigned' => true,
+				]);
+				$table->addIndex(['session_id', 'ocp_task_id'], 'assistant_ch_sid_tid');
+				$schemaChanged = true;
+			}
+		}
+
+		return $schemaChanged ? $schema : null;
+	}
+}