diff --git a/backend/open_webui/tools/builtin.py b/backend/open_webui/tools/builtin.py
index 7427592380..fe20555f9d 100644
--- a/backend/open_webui/tools/builtin.py
+++ b/backend/open_webui/tools/builtin.py
@@ -249,11 +249,13 @@ async def generate_image(
 
         # Persist files to DB if chat context is available
         if __chat_id__ and __message_id__ and images:
-            image_files = Chats.add_message_files_by_id_and_message_id(
+            db_files = Chats.add_message_files_by_id_and_message_id(
                 __chat_id__,
                 __message_id__,
                 image_files,
             )
+            if db_files is not None:
+                image_files = db_files
 
         # Emit the images to the UI if event emitter is available
         if __event_emitter__ and image_files:
@@ -314,11 +316,13 @@ async def edit_image(
 
         # Persist files to DB if chat context is available
         if __chat_id__ and __message_id__ and images:
-            image_files = Chats.add_message_files_by_id_and_message_id(
+            db_files = Chats.add_message_files_by_id_and_message_id(
                 __chat_id__,
                 __message_id__,
                 image_files,
             )
+            if db_files is not None:
+                image_files = db_files
 
         # Emit the images to the UI if event emitter is available
         if __event_emitter__ and image_files:
diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py
index 27323742a4..f3d5d6821c 100644
--- a/backend/open_webui/utils/middleware.py
+++ b/backend/open_webui/utils/middleware.py
@@ -3777,14 +3777,17 @@ async def streaming_chat_response_handler(response, ctx):
                                         delta.get("images", []), request, metadata, user
                                     )
                                     if image_urls:
+                                        image_file_list = [
+                                            {"type": "image", "url": url}
+                                            for url in image_urls
+                                        ]
                                         message_files = Chats.add_message_files_by_id_and_message_id(
                                             metadata["chat_id"],
                                             metadata["message_id"],
-                                            [
-                                                {"type": "image", "url": url}
-                                                for url in image_urls
-                                            ],
+                                            image_file_list,
                                         )
+                                        if message_files is None:
+                                            message_files = image_file_list
 
                                         await event_emitter(
                                             {