From 2f7dc60337940bc9d960692c24be2ac6d1b1483f Mon Sep 17 00:00:00 2001
From: euvre <93761161+euvre@users.noreply.github.com>
Date: Wed, 22 Jul 2026 19:21:17 +0800
Subject: [PATCH] fix: restore download button when file type is selected in
 agent message (#17184)

---
 agent/canvas.py                               |  3 +++
 .../components/next-message-item/index.tsx    | 26 -------------------
 web/src/hooks/use-send-message.ts             |  9 +++++--
 .../agent/chat/use-send-agent-message.ts      | 13 ++++++++--
 4 files changed, 21 insertions(+), 30 deletions(-)

diff --git a/agent/canvas.py b/agent/canvas.py
index bbe91e089d..87de8a5394 100644
--- a/agent/canvas.py
+++ b/agent/canvas.py
@@ -981,6 +981,9 @@ class Canvas(Graph):
             message_end["status"] = cpn_obj.get_param("status")
         if isinstance(cpn_obj.output("attachment"), dict):
             message_end["attachment"] = cpn_obj.output("attachment")
+        downloads = cpn_obj.output("downloads")
+        if isinstance(downloads, list) and downloads:
+            message_end["downloads"] = downloads
         if self._has_reference():
             message_end["reference"] = self.get_reference()
         # NOTE: aggregated run token usage is intentionally NOT attached here.
diff --git a/web/src/components/next-message-item/index.tsx b/web/src/components/next-message-item/index.tsx
index 9c5f3d5ce0..a294474043 100644
--- a/web/src/components/next-message-item/index.tsx
+++ b/web/src/components/next-message-item/index.tsx
@@ -315,32 +315,6 @@ function MessageItem({
                 ))}
               </div>
             )}
-            {/* {isAssistant && item.attachment && item.attachment.doc_id && (
-              <div className="w-full flex items-center justify-end">
-                <Button
-                  variant="link"
-                  className="p-1 m-0 h-auto text-text-sub-title-invert"
-                  onClick={async () => {
-                    if (item.attachment?.doc_id) {
-                      try {
-                        const response = await downloadFile({
-                          docId: item.attachment.doc_id,
-                          ext: item.attachment.format,
-                        });
-                        const blob = new Blob([response.data], {
-                          type: response.data.type,
-                        });
-                        downloadFileFromBlob(blob, item.attachment.file_name);
-                      } catch (error) {
-                        console.error('Download failed:', error);
-                      }
-                    }
-                  }}
-                >
-                  <Download size={16} />
-                </Button>
-              </div>
-            )} */}
           </section>
         </div>
       </section>
diff --git a/web/src/hooks/use-send-message.ts b/web/src/hooks/use-send-message.ts
index e418ce79de..49d1e31b85 100644
--- a/web/src/hooks/use-send-message.ts
+++ b/web/src/hooks/use-send-message.ts
@@ -1,7 +1,10 @@
 import message from '@/components/ui/message';
 import { Authorization } from '@/constants/authorization';
 import { ResponseType } from '@/interfaces/database/base';
-import { IReferenceObject } from '@/interfaces/database/chat';
+import {
+  IDocumentDownloadInfo,
+  IReferenceObject,
+} from '@/interfaces/database/chat';
 import { BeginQuery } from '@/pages/agent/interface';
 import { getAuthorization } from '@/utils/authorization-util';
 import { EventSourceParserStream } from 'eventsource-parser/stream';
@@ -59,7 +62,9 @@ export interface IMessageData {
 }
 
 export interface IMessageEndData {
-  reference: IReferenceObject;
+  reference?: IReferenceObject;
+  attachment?: IAttachment;
+  downloads?: IDocumentDownloadInfo[];
 }
 
 export interface ILogData extends INodeData {
diff --git a/web/src/pages/agent/chat/use-send-agent-message.ts b/web/src/pages/agent/chat/use-send-agent-message.ts
index 855e819fe2..6af27139e0 100644
--- a/web/src/pages/agent/chat/use-send-agent-message.ts
+++ b/web/src/pages/agent/chat/use-send-agent-message.ts
@@ -83,12 +83,21 @@ export function findMessageFromList(eventList: IEventList) {
   const workflowFinished = eventList.find(
     (x) => x.event === MessageEventType.WorkflowFinished,
   ) as IMessageEvent;
+  const messageEndEvent = [...eventList]
+    .reverse()
+    .find((x) => x.event === MessageEventType.MessageEnd) as IMessageEndEvent;
   return {
     id: eventList[0]?.message_id,
     content: nextContent,
     audio_binary: audioBinary,
-    attachment: workflowFinished?.data?.outputs?.attachment || {},
-    downloads: workflowFinished?.data?.outputs?.downloads || [],
+    attachment:
+      workflowFinished?.data?.outputs?.attachment ||
+      messageEndEvent?.data?.attachment ||
+      {},
+    downloads:
+      workflowFinished?.data?.outputs?.downloads ||
+      messageEndEvent?.data?.downloads ||
+      [],
   };
 }