refac

Author: tjbck

backend/open_webui/main.py

@@ -1787,11 +1787,12 @@ async def process_chat(request, form_data, user, metadata, model):
             log.info('Chat processing was cancelled')
             try:
                 event_emitter = await get_event_emitter(metadata)
-                await asyncio.shield(
-                    event_emitter(
-                        {'type': 'chat:tasks:cancel'},
+                if event_emitter:
+                    await asyncio.shield(
+                        event_emitter(
+                            {'type': 'chat:tasks:cancel'},
+                        )
                     )
-                )
             except Exception as e:
                 pass
             finally:
@@ -1812,15 +1813,16 @@ async def process_chat(request, form_data, user, metadata, model):
                         )
 
                     event_emitter = await get_event_emitter(metadata)
-                    await event_emitter(
-                        {
-                            'type': 'chat:message:error',
-                            'data': {'error': {'content': str(e)}},
-                        }
-                    )
-                    await event_emitter(
-                        {'type': 'chat:tasks:cancel'},
-                    )
+                    if event_emitter:
+                        await event_emitter(
+                            {
+                                'type': 'chat:message:error',
+                                'data': {'error': {'content': str(e)}},
+                            }
+                        )
+                        await event_emitter(
+                            {'type': 'chat:tasks:cancel'},
+                        )
 
                 except Exception:
                     pass

backend/open_webui/utils/middleware.py

@@ -4692,25 +4692,33 @@ async def restricted_import(name, globals=None, locals=None, fromlist=(), level=
                 await background_tasks_handler(ctx)
             except asyncio.CancelledError:
                 log.warning('Task was cancelled!')
-                await event_emitter({'type': 'chat:tasks:cancel'})
+                try:
+                    await asyncio.shield(event_emitter({'type': 'chat:tasks:cancel'}))
 
-                if not ENABLE_REALTIME_CHAT_SAVE:
-                    # Save message in the database
-                    await Chats.upsert_message_to_chat_by_id_and_message_id(
-                        metadata['chat_id'],
-                        metadata['message_id'],
-                        {
-                            'done': True,
-                            'content': serialize_output(output),
-                            'output': output,
-                        },
-                    )
-                else:
-                    await Chats.upsert_message_to_chat_by_id_and_message_id(
-                        metadata['chat_id'],
-                        metadata['message_id'],
-                        {'done': True},
-                    )
+                    if not ENABLE_REALTIME_CHAT_SAVE:
+                        # Save message in the database
+                        await asyncio.shield(
+                            Chats.upsert_message_to_chat_by_id_and_message_id(
+                                metadata['chat_id'],
+                                metadata['message_id'],
+                                {
+                                    'done': True,
+                                    'content': serialize_output(output),
+                                    'output': output,
+                                },
+                            )
+                        )
+                    else:
+                        await asyncio.shield(
+                            Chats.upsert_message_to_chat_by_id_and_message_id(
+                                metadata['chat_id'],
+                                metadata['message_id'],
+                                {'done': True},
+                            )
+                        )
+                except Exception:
+                    pass
+                raise  # re-raise CancelledError for proper propagation
 
             if response.background is not None:
                 await response.background()