From de446952c2db1168dca5c6e3da6d961e6f55a8eb Mon Sep 17 00:00:00 2001
From: Timothy Jaeryang Baek <tim@openwebui.com>
Date: Wed, 14 Oct 2026 23:58:58 +0400
Subject: [PATCH] refac

---
 backend/open_webui/main.py              | 14 +++++++++-----
 backend/open_webui/utils/automations.py | 17 +++++++++++++----
 2 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/backend/open_webui/main.py b/backend/open_webui/main.py
index 863d8be4a0..5db490fbf2 100644
--- a/backend/open_webui/main.py
+++ b/backend/open_webui/main.py
@@ -1813,7 +1813,7 @@ async def chat_completion(
 
     # Fan out: one task per model
     if metadata.get('session_id') and metadata.get('chat_id'):
-        task_ids = []
+        chat_tasks = {}
         subagent_results = []
         is_internal = getattr(request.state, 'internal', False) is True
         chat_id = metadata['chat_id']
@@ -1869,13 +1869,13 @@ async def chat_completion(
                 subagent_results.append(await process)
                 continue
 
-            task_id, _ = await create_task(
+            task_id, task = await create_task(
                 request.app.state.redis,
                 process,
                 id=chat_id,
                 task_id=per_model_metadata['task_id'],
             )
-            task_ids.append(task_id)
+            chat_tasks[task_id] = task
 
         if is_internal:
             return {
@@ -1886,7 +1886,7 @@ async def chat_completion(
             }
 
         # Emit chat:active=true
-        if task_ids:
+        if chat_tasks:
             event_emitter = await get_event_emitter(
                 {**metadata, 'message_id': message_ids[0]['message_id']},
                 update_db=False,
@@ -1895,9 +1895,13 @@ async def chat_completion(
                 folder_id = metadata.get('folder_id') or await Chats.get_chat_folder_id(chat_id, user.id)
                 await event_emitter({'type': 'chat:active', 'data': {'active': True, 'folder_id': folder_id}})
 
+        # Automation chat runs must finish before their caller records success.
+        if metadata.get('automation_id') and is_saved_chat_id(chat_id):
+            await asyncio.gather(*chat_tasks.values())
+
         return {
             'status': True,
-            'task_ids': task_ids,
+            'task_ids': list(chat_tasks),
             'chat_id': chat_id,
         }
     else:
diff --git a/backend/open_webui/utils/automations.py b/backend/open_webui/utils/automations.py
index 0bda3a3d1b..ce4c9af3ec 100644
--- a/backend/open_webui/utils/automations.py
+++ b/backend/open_webui/utils/automations.py
@@ -357,6 +357,7 @@ async def execute_automation(app, automation: AutomationModel) -> None:
     session_id + chat_id + message_id → async task → pipeline handles everything
     (filters, model params, knowledge/RAG, tools, DB saves, webhooks).
     """
+    chat = None
     try:
         user = await Users.get_user_by_id(automation.user_id)
         if not user:
@@ -521,6 +522,12 @@ async def execute_automation(app, automation: AutomationModel) -> None:
         request = _build_request(app, token=token)
         await app.state.CHAT_COMPLETION_HANDLER(request, form_data, user=user)
 
+        reply = await Chats.get_message_by_id_and_message_id(chat.id, assistant_msg_id)
+        if reply and reply.get('error'):
+            raise ValueError(str(reply['error']['content']))
+        if not reply or not reply.get('done'):
+            raise ValueError('Automation reply did not complete')
+
         # Notify user
         from open_webui.socket.main import sio
 
@@ -544,16 +551,18 @@ async def execute_automation(app, automation: AutomationModel) -> None:
             data={'name': automation.name, 'chat_id': chat.id},
         )
 
-    except Exception as e:
+    except (Exception, asyncio.CancelledError) as e:
         log.exception(f'Automation {automation.id} failed')
-        error = str(e)[:4000]
-        await _record_run(automation.id, 'error', error=error)
+        error = 'Automation reply was cancelled' if isinstance(e, asyncio.CancelledError) else str(e)[:4000]
+        await _record_run(automation.id, 'error', chat_id=chat.id if chat else None, error=error)
         await publish_event(
             app,
             EVENTS.AUTOMATION_RUN_FAILED,
             subject_id=automation.id,
-            data={'name': automation.name, 'error': error},
+            data={'name': automation.name, 'error': error, **({'chat_id': chat.id} if chat else {})},
         )
+        if isinstance(e, asyncio.CancelledError):
+            raise
 
 
 ####################