diff --git a/frontend/appflowy_flutter/lib/plugins/ai_chat/chat_page.dart b/frontend/appflowy_flutter/lib/plugins/ai_chat/chat_page.dart
index e8d282157756..789cce87706f 100644
--- a/frontend/appflowy_flutter/lib/plugins/ai_chat/chat_page.dart
+++ b/frontend/appflowy_flutter/lib/plugins/ai_chat/chat_page.dart
@@ -91,7 +91,6 @@ class _AIChatPageState extends State<AIChatPage> {
                       state.loadingPreviousStatus.when(
                         loading: () => Log.debug("loading"),
                         finish: () {
-                          Log.debug("loading more messages");
                           blocContext
                               .read<ChatBloc>()
                               .add(const ChatEvent.loadPrevMessage());
@@ -117,41 +116,43 @@ class _AIChatPageState extends State<AIChatPage> {
     );
   }
 
-  ChatPopupMenu buildBubble(Message message, Widget child) {
+  Widget buildBubble(Message message, Widget child) {
     final isAuthor = message.author.id == _user.id;
     const borderRadius = BorderRadius.all(Radius.circular(20));
-    if (isMobile) {}
-
-    return ChatPopupMenu(
-      onAction: (action) {
-        switch (action) {
-          case ChatMessageAction.copy:
-            if (message is TextMessage) {
-              Clipboard.setData(
-                ClipboardData(text: message.text),
-              );
-              showMessageToast(
-                LocaleKeys.grid_row_copyProperty.tr(),
-              );
-            }
-            break;
-        }
-      },
-      builder: (context) {
-        return DecoratedBox(
-          decoration: BoxDecoration(
-            borderRadius: borderRadius,
-            color: !isAuthor || message.type == types.MessageType.image
-                ? AFThemeExtension.of(context).tint1
-                : Theme.of(context).colorScheme.primary,
-          ),
-          child: ClipRRect(
-            borderRadius: borderRadius,
-            child: Container(child: child),
-          ),
-        );
-      },
+    final decoratedChild = DecoratedBox(
+      decoration: BoxDecoration(
+        borderRadius: borderRadius,
+        color: !isAuthor || message.type == types.MessageType.image
+            ? AFThemeExtension.of(context).tint1
+            : Theme.of(context).colorScheme.primary,
+      ),
+      child: ClipRRect(
+        borderRadius: borderRadius,
+        child: Container(child: child),
+      ),
     );
+
+    if (isMobile) {
+      return ChatPopupMenu(
+        onAction: (action) {
+          switch (action) {
+            case ChatMessageAction.copy:
+              if (message is TextMessage) {
+                Clipboard.setData(
+                  ClipboardData(text: message.text),
+                );
+                showMessageToast(
+                  LocaleKeys.grid_row_copyProperty.tr(),
+                );
+              }
+              break;
+          }
+        },
+        builder: (context) => decoratedChild,
+      );
+    } else {
+      return decoratedChild;
+    }
   }
 
   AFDefaultChatTheme buildTheme(BuildContext context) {
diff --git a/frontend/rust-lib/event-integration-test/tests/chat/chat_message_test.rs b/frontend/rust-lib/event-integration-test/tests/chat/chat_message_test.rs
index f9fc6c13e714..ff59e6114f81 100644
--- a/frontend/rust-lib/event-integration-test/tests/chat/chat_message_test.rs
+++ b/frontend/rust-lib/event-integration-test/tests/chat/chat_message_test.rs
@@ -116,7 +116,7 @@ async fn af_cloud_load_remote_system_message_test() {
   // Load first five messages
   let rx = test
     .notification_sender
-    .subscribe::<ChatMessageListPB>(&chat_id, ChatNotification::DidLoadLatestChatMessage);
+    .subscribe::<ChatMessageListPB>(&chat_id, ChatNotification::DidLoadPrevChatMessage);
   test
     .load_prev_message(&chat_id, 5, Some(next_back_five.messages[4].message_id))
     .await;
diff --git a/frontend/rust-lib/flowy-chat/src/chat.rs b/frontend/rust-lib/flowy-chat/src/chat.rs
index ef7368a518dd..c850f71100e6 100644
--- a/frontend/rust-lib/flowy-chat/src/chat.rs
+++ b/frontend/rust-lib/flowy-chat/src/chat.rs
@@ -333,9 +333,6 @@ fn stream_send_chat_messages(
             },
           }
         }
-
-        // Mark chat as finished
-        send_notification(&chat_id, ChatNotification::FinishAnswerQuestion).send();
       },
       Err(err) => {
         error!("Failed to send chat message: {}", err);
@@ -376,6 +373,9 @@ fn stream_send_chat_messages(
         })
         .collect::<Vec<_>>();
       insert_chat_messages(conn, &records)?;
+
+      // Mark chat as finished
+      send_notification(&chat_id, ChatNotification::FinishAnswerQuestion).send();
       Ok(())
     }) {
       error!("Failed to save chat messages: {}", err);