Merge branch 'github_main' into 'develop'

feat(backend): 修复chunk无法跨进程广播的问题 (#620)

See merge request weibo_rd/common/wecode/wegent!877

Author: 李泓宇

backend/app/services/chat/ai_trigger.py

@@ -124,15 +124,12 @@ async def _trigger_direct_chat(
     """
     from app.api.ws.events import ServerEvents
 
-    # Emit chat:start event
+    # Emit chat:start event using global emitter for cross-worker broadcasting
     logger.info(f"[ai_trigger] Emitting chat:start event")
-    await namespace.emit(
-        ServerEvents.CHAT_START,
-        {
-            "task_id": task.id,
-            "subtask_id": assistant_subtask.id,
-        },
-        room=task_room,
+    emitter = get_ws_emitter()
+    await emitter.emit_chat_start(
+        task_id=task.id,
+        subtask_id=assistant_subtask.id,
     )
     logger.info(f"[ai_trigger] chat:start emitted")
 
@@ -259,10 +256,11 @@ async def _stream_chat_response(
             # Record error in span
             span_manager.record_error(TelemetryEventNames.BOT_NOT_FOUND, error_msg)
 
-            await namespace.emit(
-                ServerEvents.CHAT_ERROR,
-                {"subtask_id": subtask_id, "error": error_msg},
-                room=task_room,
+            emitter = get_ws_emitter()
+            await emitter.emit_chat_error(
+                task_id=task_id,
+                subtask_id=subtask_id,
+                error=error_msg,
             )
             return
 
@@ -367,6 +365,9 @@ async def _stream_chat_response(
         from app.services.chat.providers import get_provider
         from app.services.chat.providers.base import ChunkType
 
+        # Get global emitter
+        emitter = get_ws_emitter()
+
         # Check if this is a group chat - get history from database with user names
         is_group_chat = payload.is_group_chat
         if is_group_chat:
@@ -405,10 +406,10 @@ async def _stream_chat_response(
                 TelemetryEventNames.PROVIDER_CREATION_FAILED, error_msg
             )
 
-            await namespace.emit(
-                ServerEvents.CHAT_ERROR,
-                {"subtask_id": subtask_id, "error": error_msg},
-                room=task_room,
+            await emitter.emit_chat_error(
+                task_id=task_id,
+                subtask_id=subtask_id,
+                error=error_msg,
             )
             return
 
@@ -429,25 +430,21 @@ async def _stream_chat_response(
         async for chunk in stream_gen:
             if cancel_event.is_set() or await session_manager.is_cancelled(subtask_id):
                 # Cancelled
-                await namespace.emit(
-                    ServerEvents.CHAT_CANCELLED,
-                    {"subtask_id": subtask_id},
-                    room=task_room,
+                await emitter.emit_chat_cancelled(
+                    task_id=task_id,
+                    subtask_id=subtask_id,
                 )
                 break
 
             if chunk.type == ChunkType.CONTENT and chunk.content:
                 full_response += chunk.content
 
                 # Emit chunk
-                await namespace.emit(
-                    ServerEvents.CHAT_CHUNK,
-                    {
-                        "subtask_id": subtask_id,
-                        "content": chunk.content,
-                        "offset": offset,
-                    },
-                    room=task_room,
+                await emitter.emit_chat_chunk(
+                    task_id=task_id,
+                    subtask_id=subtask_id,
+                    content=chunk.content,
+                    offset=offset,
                 )
                 offset += len(chunk.content)
 
@@ -480,13 +477,10 @@ async def _stream_chat_response(
                     TelemetryEventNames.STREAM_CHUNK_ERROR, error_msg, model_config
                 )
 
-                await namespace.emit(
-                    ServerEvents.CHAT_ERROR,
-                    {
-                        "subtask_id": subtask_id,
-                        "error": error_msg,
-                    },
-                    room=task_room,
+                await emitter.emit_chat_error(
+                    task_id=task_id,
+                    subtask_id=subtask_id,
+                    error=error_msg,
                 )
                 await db_handler.update_subtask_status(
                     subtask_id, "FAILED", error=chunk.error
@@ -551,10 +545,12 @@ async def _stream_chat_response(
         # Record error in span
         span_manager.record_exception(e)
 
-        await namespace.emit(
-            ServerEvents.CHAT_ERROR,
-            {"subtask_id": subtask_id, "error": str(e)},
-            room=task_room,
+        # Use global emitter for cross-worker broadcasting
+        error_emitter = get_ws_emitter()
+        await error_emitter.emit_chat_error(
+            task_id=task_id,
+            subtask_id=subtask_id,
+            error=str(e),
         )
     finally:
         # Detach OTEL context first (before exiting span)

backend/app/services/chat/ws_emitter.py

@@ -52,6 +52,8 @@ async def emit_chat_start(
         task_id: int,
         subtask_id: int,
         bot_name: Optional[str] = None,
+        shell_type: str = "Chat",
+        message_id: Optional[int] = None,
     ) -> None:
         """
         Emit chat:start event to task room.
@@ -60,18 +62,24 @@ async def emit_chat_start(
             task_id: Task ID
             subtask_id: Subtask ID
             bot_name: Optional bot name
+            shell_type: Shell type for frontend display logic (default: "Chat")
+            message_id: Optional message ID for ordering
         """
         await self.sio.emit(
             ServerEvents.CHAT_START,
             {
                 "task_id": task_id,
                 "subtask_id": subtask_id,
                 "bot_name": bot_name,
+                "shell_type": shell_type,
+                "message_id": message_id,
             },
             room=f"task:{task_id}",
             namespace=self.namespace,
         )
-        logger.debug(f"[WS] emit chat:start task={task_id} subtask={subtask_id}")
+        logger.debug(
+            f"[WS] emit chat:start task={task_id} subtask={subtask_id} shell_type={shell_type}"
+        )
 
     async def emit_chat_chunk(
         self,

backend/app/services/chat_v2/ai_trigger.py

@@ -118,6 +118,7 @@ async def _trigger_direct_chat(
     Emits chat:start event and starts streaming in background task.
     """
     from app.api.ws.events import ServerEvents
+    from app.services.chat.ws_emitter import get_ws_emitter
 
     # Extract data from ORM objects before starting background task
     # This prevents DetachedInstanceError
@@ -241,10 +242,13 @@ async def _stream_chat_response(
         if not team:
             error_msg = "Team not found"
             span_manager.record_error(TelemetryEventNames.TEAM_NOT_FOUND, error_msg)
-            await namespace.emit(
-                ServerEvents.CHAT_ERROR,
-                {"subtask_id": subtask_id, "error": error_msg},
-                room=task_room,
+            from app.services.chat.ws_emitter import get_ws_emitter
+
+            error_emitter = get_ws_emitter()
+            await error_emitter.emit_chat_error(
+                task_id=task_data["id"],
+                subtask_id=subtask_id,
+                error=error_msg,
             )
             return
 
@@ -269,10 +273,13 @@ async def _stream_chat_response(
             span_manager.record_error(
                 TelemetryEventNames.CONFIG_BUILD_FAILED, error_msg
             )
-            await namespace.emit(
-                ServerEvents.CHAT_ERROR,
-                {"subtask_id": subtask_id, "error": error_msg},
-                room=task_room,
+            from app.services.chat.ws_emitter import get_ws_emitter
+
+            error_emitter = get_ws_emitter()
+            await error_emitter.emit_chat_error(
+                task_id=task_data["id"],
+                subtask_id=subtask_id,
+                error=error_msg,
             )
             return
 
@@ -337,10 +344,14 @@ async def _stream_chat_response(
         logger.exception("[ai_trigger] Stream error subtask=%d: %s", subtask_id, e)
         # Record error in span
         span_manager.record_exception(e)
-        await namespace.emit(
-            ServerEvents.CHAT_ERROR,
-            {"subtask_id": subtask_id, "error": str(e)},
-            room=task_room,
+        # Use global emitter for cross-worker broadcasting
+        from app.services.chat.ws_emitter import get_ws_emitter
+
+        error_emitter = get_ws_emitter()
+        await error_emitter.emit_chat_error(
+            task_id=task_data["id"],
+            subtask_id=subtask_id,
+            error=str(e),
         )
     finally:
         # Detach OTEL context first (before exiting span)

backend/app/services/chat_v2/service.py

@@ -555,7 +555,7 @@ async def stream_to_websocket(
         task_room = config.task_room
 
         # Create WebSocket emitter
-        emitter = WebSocketEmitter(namespace, task_room)
+        emitter = WebSocketEmitter(namespace, task_room, task_id)
 
         # Create streaming state
         state = StreamingState(

backend/app/services/chat_v2/streaming/emitters.py

@@ -153,35 +153,35 @@ def has_events(self) -> bool:
 
 
 class WebSocketEmitter(StreamEmitter):
-    """WebSocket emitter using Socket.IO namespace.
+    """WebSocket emitter using global ws_emitter for cross-worker broadcasting.
 
-    Emits events through a Socket.IO namespace to a specific room.
+    Uses the global ws_emitter (with Redis adapter) to ensure events are
+    broadcast to all backend replicas in multi-instance deployments.
     """
 
-    def __init__(self, namespace: Any, task_room: str):
+    def __init__(self, namespace: Any, task_room: str, task_id: int):
         """Initialize WebSocket emitter.
 
         Args:
-            namespace: Socket.IO namespace instance
+            namespace: Socket.IO namespace instance (kept for compatibility)
             task_room: Room name for broadcasting events
+            task_id: Task ID for emitting events
         """
         self.namespace = namespace
         self.task_room = task_room
+        self.task_id = task_id
 
     async def emit_start(
         self, task_id: int, subtask_id: int, shell_type: str = "Chat"
     ) -> None:
-        """Emit chat:start event with shell_type."""
-        from app.api.ws.events import ServerEvents
-
-        await self.namespace.emit(
-            ServerEvents.CHAT_START,
-            {
-                "task_id": task_id,
-                "subtask_id": subtask_id,
-                "shell_type": shell_type,  # Include shell_type for frontend display logic
-            },
-            room=self.task_room,
+        """Emit chat:start event using global emitter with shell_type."""
+        from app.services.chat.ws_emitter import get_ws_emitter
+
+        emitter = get_ws_emitter()
+        await emitter.emit_chat_start(
+            task_id=task_id,
+            subtask_id=subtask_id,
+            shell_type=shell_type,
         )
         logger.info("[WS_EMITTER] chat:start emitted with shell_type=%s", shell_type)
 
@@ -192,11 +192,11 @@ async def emit_chunk(
         subtask_id: int,
         result: dict[str, Any] | None = None,
     ) -> None:
-        """Emit chat:chunk event with optional result data (thinking, workbench).
+        """Emit chat:chunk event using global emitter with optional result data.
 
         This follows the same pattern as executor tasks to enable thinking/workbench display.
         """
-        from app.api.ws.events import ServerEvents
+        from app.services.chat.ws_emitter import get_ws_emitter
 
         logger.debug(
             "[WS_EMITTER] emit_chunk: subtask_id=%d, offset=%d, content_len=%d, has_result=%s",
@@ -205,19 +205,13 @@ async def emit_chunk(
             len(content),
             result is not None,
         )
-        payload: dict[str, Any] = {
-            "subtask_id": subtask_id,
-            "content": content,
-            "offset": offset,
-        }
-        # Include full result if provided (for chat_v2 thinking/workbench display)
-        if result is not None:
-            payload["result"] = result
-
-        await self.namespace.emit(
-            ServerEvents.CHAT_CHUNK,
-            payload,
-            room=self.task_room,
+        emitter = get_ws_emitter()
+        await emitter.emit_chat_chunk(
+            task_id=self.task_id,
+            subtask_id=subtask_id,
+            content=content,
+            offset=offset,
+            result=result,
         )
 
     async def emit_done(
@@ -228,40 +222,38 @@ async def emit_done(
         result: dict[str, Any],
         message_id: int | None = None,
     ) -> None:
-        """Emit chat:done event."""
-        from app.api.ws.events import ServerEvents
-
-        await self.namespace.emit(
-            ServerEvents.CHAT_DONE,
-            {
-                "task_id": task_id,
-                "subtask_id": subtask_id,
-                "offset": offset,
-                "result": result,
-                "message_id": message_id,
-            },
-            room=self.task_room,
+        """Emit chat:done event using global emitter."""
+        from app.services.chat.ws_emitter import get_ws_emitter
+
+        emitter = get_ws_emitter()
+        await emitter.emit_chat_done(
+            task_id=task_id,
+            subtask_id=subtask_id,
+            offset=offset,
+            result=result,
+            message_id=message_id,
         )
         logger.info("[WS_EMITTER] chat:done emitted message_id=%s", message_id)
 
     async def emit_error(self, subtask_id: int, error: str) -> None:
-        """Emit chat:error event."""
-        from app.api.ws.events import ServerEvents
-
-        await self.namespace.emit(
-            ServerEvents.CHAT_ERROR,
-            {"subtask_id": subtask_id, "error": error},
-            room=self.task_room,
+        """Emit chat:error event using global emitter."""
+        from app.services.chat.ws_emitter import get_ws_emitter
+
+        emitter = get_ws_emitter()
+        await emitter.emit_chat_error(
+            task_id=self.task_id,
+            subtask_id=subtask_id,
+            error=error,
         )
         logger.warning("[WS_EMITTER] chat:error emitted: %s", error)
 
     async def emit_cancelled(self, subtask_id: int) -> None:
-        """Emit chat:cancelled event."""
-        from app.api.ws.events import ServerEvents
+        """Emit chat:cancelled event using global emitter."""
+        from app.services.chat.ws_emitter import get_ws_emitter
 
-        await self.namespace.emit(
-            ServerEvents.CHAT_CANCELLED,
-            {"subtask_id": subtask_id},
-            room=self.task_room,
+        emitter = get_ws_emitter()
+        await emitter.emit_chat_cancelled(
+            task_id=self.task_id,
+            subtask_id=subtask_id,
         )
         logger.info("[WS_EMITTER] chat:cancelled emitted")