From 0e79e29828fb8311013402091cdd766fb051c5b4 Mon Sep 17 00:00:00 2001
From: Marco Beretta <81851188+berry-13@users.noreply.github.com>
Date: Sun, 5 Jul 2026 03:11:50 +0200
Subject: [PATCH] fix: collect tool-output attachment ids when capping files

collectConversationFileIds only read Message.files, but agent/tool outputs
(code artifacts, generated images) are stored on Message.attachments with their
own file_id references, and those File rows also commonly lack a
conversationId. A forced-retention conversion or migration could therefore
convert the conversation and messages while leaving attachment file metadata at
expiredAt: null, so storage cleanup never swept it.

Match messages carrying either field and collect file ids from both arrays.
---
 .../data-schemas/src/methods/message.spec.ts  |  6 +++
 packages/data-schemas/src/utils/retention.ts  | 40 +++++++++++++------
 2 files changed, 33 insertions(+), 13 deletions(-)

diff --git a/packages/data-schemas/src/methods/message.spec.ts b/packages/data-schemas/src/methods/message.spec.ts
index ca0694cf7b..ca7803143d 100644
--- a/packages/data-schemas/src/methods/message.spec.ts
+++ b/packages/data-schemas/src/methods/message.spec.ts
@@ -1864,6 +1864,7 @@ describe('Message Operations', () => {
       const forcedExpiredAt = new Date(Date.now() + 24 * 60 * 60 * 1000);
       const conversationId = uuidv4();
       const messageFileId = uuidv4();
+      const toolOutputFileId = uuidv4();
       const convoFileId = uuidv4();
 
       await Conversation().create({
@@ -1879,9 +1880,11 @@ describe('Message Operations', () => {
         user: 'user123',
         text: 'with attachment',
         files: [{ file_id: messageFileId, filename: 'doc.pdf' }],
+        attachments: [{ file_id: toolOutputFileId, filename: 'chart.png' }],
       });
       await File().collection.insertMany([
         { file_id: messageFileId, user: new mongoose.Types.ObjectId(), expiredAt: null },
+        { file_id: toolOutputFileId, user: new mongoose.Types.ObjectId(), expiredAt: null },
         { file_id: convoFileId, user: new mongoose.Types.ObjectId(), expiredAt: null },
       ]);
 
@@ -1898,6 +1901,9 @@ describe('Message Operations', () => {
       const messageFile = await File().findOne({ file_id: messageFileId }).lean();
       expect(messageFile?.expiredAt?.getTime()).toBe(forcedExpiredAt.getTime());
 
+      const toolOutputFile = await File().findOne({ file_id: toolOutputFileId }).lean();
+      expect(toolOutputFile?.expiredAt?.getTime()).toBe(forcedExpiredAt.getTime());
+
       const convoFile = await File().findOne({ file_id: convoFileId }).lean();
       expect(convoFile?.expiredAt?.getTime()).toBe(forcedExpiredAt.getTime());
     });
diff --git a/packages/data-schemas/src/utils/retention.ts b/packages/data-schemas/src/utils/retention.ts
index 3c9a6615bd..5d2705a7f0 100644
--- a/packages/data-schemas/src/utils/retention.ts
+++ b/packages/data-schemas/src/utils/retention.ts
@@ -163,10 +163,11 @@ export const capConversationSharedLinks = async (
 
 /**
  * Collects the file ids a set of conversations references. Message-attachment uploads create
- * File rows without a `conversationId` — they are referenced only from `Message.files[].file_id`
- * and the conversation's own `files` array — so conversation-scoped file caps must also target
- * these ids. `seedFileIds` takes the conversations' `files` arrays; message references are read
- * in one pass over the conversations' messages.
+ * File rows without a `conversationId` — they are referenced only from `Message.files[].file_id`,
+ * `Message.attachments[].file_id` (tool/agent outputs), and the conversation's own `files`
+ * array — so conversation-scoped file caps must also target these ids. `seedFileIds` takes the
+ * conversations' `files` arrays; message references are read in one pass over the conversations'
+ * messages.
  */
 export const collectConversationFileIds = async (
   Message: Model<IMessage>,
@@ -175,6 +176,15 @@ export const collectConversationFileIds = async (
   seedFileIds?: Iterable<string | null | undefined>,
 ): Promise<string[]> => {
   const fileIds = new Set<string>();
+  const addFileIds = (references?: Array<{ file_id?: unknown } | null>) => {
+    for (const reference of references ?? []) {
+      const fileId = reference?.file_id;
+      if (typeof fileId === 'string' && fileId.length > 0) {
+        fileIds.add(fileId);
+      }
+    }
+  };
+
   for (const fileId of seedFileIds ?? []) {
     if (typeof fileId === 'string' && fileId.length > 0) {
       fileIds.add(fileId);
@@ -185,17 +195,21 @@ export const collectConversationFileIds = async (
       {
         user: userId,
         conversationId: { $in: conversationIds },
-        files: { $exists: true, $ne: null },
+        $or: [
+          { files: { $exists: true, $ne: null } },
+          { attachments: { $exists: true, $ne: null } },
+        ],
       },
-      'files',
-    ).lean<Array<{ files?: Array<{ file_id?: unknown } | null> }>>();
+      'files attachments',
+    ).lean<
+      Array<{
+        files?: Array<{ file_id?: unknown } | null>;
+        attachments?: Array<{ file_id?: unknown } | null>;
+      }>
+    >();
     for (const message of messages) {
-      for (const file of message.files ?? []) {
-        const fileId = file?.file_id;
-        if (typeof fileId === 'string' && fileId.length > 0) {
-          fileIds.add(fileId);
-        }
-      }
+      addFileIds(message.files);
+      addFileIds(message.attachments);
     }
   }
   return [...fileIds];