diff --git a/src/views/components/StopButton/index.tsx b/src/views/components/StopButton/index.tsx
index 6cf840e..01b9d5a 100644
--- a/src/views/components/StopButton/index.tsx
+++ b/src/views/components/StopButton/index.tsx
@@ -25,7 +25,7 @@ const StopButton = observer(() => {
                 }
             }}
             onClick={() => {
-                chat.stopGenerating(false, null);
+                chat.stopGenerating(false, '', chat.currentMessage);
                 messageUtil.sendMessage({
                     command: 'stopDevChat'
                 });
diff --git a/src/views/pages/ChatPanel.tsx b/src/views/pages/ChatPanel.tsx
index 2ea1ed6..565b129 100644
--- a/src/views/pages/ChatPanel.tsx
+++ b/src/views/pages/ChatPanel.tsx
@@ -88,8 +88,7 @@ const chatPanel = observer(() => {
             timer.start();
         });
         messageUtil.registerHandler('receiveMessage', (message: { text: string; isError: boolean, hash }) => {
-            const messageItem = Message.create({ type: 'bot', message: message.text, hash: message.hash });
-            chat.stopGenerating(true, messageItem);
+            chat.stopGenerating(true, message.hash, message.text);
             if (message.isError) {
                 chat.happendError(message.text);
             }
diff --git a/src/views/stores/ChatStore.ts b/src/views/stores/ChatStore.ts
index 16d2864..dfb4283 100644
--- a/src/views/stores/ChatStore.ts
+++ b/src/views/stores/ChatStore.ts
@@ -170,20 +170,20 @@ You can configure DevChat from [Settings](#settings).`;
                     command: 'regeneration'
                 });
             },
-            stopGenerating: (hasDone: boolean, message: Instance<typeof Message> | Instance<typeof types.null>) => {
+            stopGenerating: (hasDone: boolean, hash: string = '', message: string = '') => {
                 self.generating = false;
                 self.responsed = false;
                 self.hasDone = hasDone;
+                const messagesLength = self.messages.length;
                 if (hasDone) {
-                    const { hash } = message ? message : { hash: '' };
-                    const messagesLength = self.messages.length;
-
                     if (messagesLength > 1) {
                         self.messages[messagesLength - 2].hash = hash;
                         self.messages[messagesLength - 1].hash = hash;
                     } else if (messagesLength > 0) {
                         self.messages[messagesLength - 1].hash = hash;
                     }
+                } else {
+                    self.messages[messagesLength - 1].message = message;
                 }
             },
             startResponsing: (message: string) => {