diff --git a/libs/mobile/chat/features/chat/src/lib/component.tsx b/libs/mobile/chat/features/chat/src/lib/component.tsx
index 21a6dc7..acf8177 100644
--- a/libs/mobile/chat/features/chat/src/lib/component.tsx
+++ b/libs/mobile/chat/features/chat/src/lib/component.tsx
@@ -62,7 +62,7 @@ export function Chat({ chatId, selectedModelId, isNewChat, resetToChatsList }: C
   } = useEditMessage({ chat, modelId: selectedModelId });
 
   const history = chat?.chat.history;
-  const isResponseGenerating = !history?.messages[history.currentId].done;
+  const isResponseGenerating = history?.lastAssistantMessage?.done !== true;
 
   const shouldHideContent = isLoading || isRefetching || !isMessagesListLoaded || !selectedModelId;
 
diff --git a/libs/mobile/chat/features/chat/src/lib/components/messages-list/component.tsx b/libs/mobile/chat/features/chat/src/lib/components/messages-list/component.tsx
index fb94147..fafe50d 100644
--- a/libs/mobile/chat/features/chat/src/lib/components/messages-list/component.tsx
+++ b/libs/mobile/chat/features/chat/src/lib/components/messages-list/component.tsx
@@ -11,9 +11,11 @@ import { useSetSelectedModel } from '@open-webui-react-native/mobile/shared/feat
 import { View, AppFlashList } from '@open-webui-react-native/mobile/shared/ui/ui-kit';
 import { ChatScreenParams } from '@open-webui-react-native/mobile/shared/utils/navigation';
 import {
+  Chat,
   chatApi,
   History as ChatHistory,
   Message,
+  patchChatQueryData,
   prepareCompleteChatPayload,
 } from '@open-webui-react-native/shared/data-access/api';
 import { Role } from '@open-webui-react-native/shared/data-access/common';
@@ -131,6 +133,19 @@ export default function ChatMessagesList({
   const handleContinueResponsePress = (messageId: string): void => {
     if (!modelId) return;
 
+    // mark message as not completed
+    patchChatQueryData(chatId, {
+      chat: {
+        history: {
+          messages: {
+            [messageId]: {
+              done: false,
+            },
+          },
+        },
+      } as Chat,
+    });
+
     const completePayload = prepareCompleteChatPayload({
       chatId,
       messages,
@@ -146,7 +161,11 @@ export default function ChatMessagesList({
       const message = history?.messages[item.id];
       if (!message) return null;
 
-      const isLast = item.id === history?.lastAssistantMessage?.id;
+      const lastAssistantMessageInUIList = [...messages]
+        .reverse()
+        .find((m) => history?.messages[m.id]?.role === Role.ASSISTANT);
+
+      const isLast = item.id === lastAssistantMessageInUIList?.id;
 
       return item.role === Role.ASSISTANT ? (
         <AiMessageActions
diff --git a/libs/shared/data-access/api/src/lib/chats/utils/patch-completed-message.ts b/libs/shared/data-access/api/src/lib/chats/utils/patch-completed-message.ts
index b2cdb4d..bebab41 100644
--- a/libs/shared/data-access/api/src/lib/chats/utils/patch-completed-message.ts
+++ b/libs/shared/data-access/api/src/lib/chats/utils/patch-completed-message.ts
@@ -1,5 +1,5 @@
 import { Role } from '@open-webui-react-native/shared/data-access/common';
-import { ChatResponse, Message } from '../models';
+import { ChatResponse, Message, History } from '../models';
 
 export function patchCompletedMessage(oldData: ChatResponse | undefined): ChatResponse | undefined {
   if (
@@ -33,11 +33,10 @@ export function patchCompletedMessage(oldData: ChatResponse | undefined): ChatRe
     : undefined;
 
   const updatedHistory = updatedHistoryMessages
-    ? {
-        ...history,
+    ? new History({
         messages: updatedHistoryMessages,
-        lastAssistantMessage: history?.lastAssistantMessage ?? updatedLastMessage,
-      }
+        currentId: history.currentId,
+      })
     : history;
 
   return {