diff --git a/frontend/app/components/Kai/KaiChat.tsx b/frontend/app/components/Kai/KaiChat.tsx
index 6f9a42cdc..7c82051fb 100644
--- a/frontend/app/components/Kai/KaiChat.tsx
+++ b/frontend/app/components/Kai/KaiChat.tsx
@@ -6,7 +6,7 @@ import { PANEL_SIZES } from 'App/constants/panelSizes';
 import ChatLog from './components/ChatLog';
 import IntroSection from './components/IntroSection';
 import { useQuery } from '@tanstack/react-query';
-import { aiService } from 'App/services';
+import { kaiService } from 'App/services';
 import { toast } from 'react-toastify';
 import { useStore } from 'App/mstore';
 import { observer } from 'mobx-react-lite';
@@ -121,13 +121,13 @@ function ChatsModal({ onSelect }: { onSelect: (threadId: string, title: string)
     refetch,
   } = useQuery({
     queryKey: ['kai', 'chats'],
-    queryFn: () => aiService.getKaiChats(userId, projectId),
+    queryFn: () => kaiService.getKaiChats(userId, projectId),
     staleTime: 1000 * 60,
   });
 
   const onDelete = async (id: string) => {
     try {
-      await aiService.deleteKaiChat(projectId, userId, id);
+      await kaiService.deleteKaiChat(projectId, userId, id);
     } catch (e) {
       toast.error("Something wen't wrong. Please try again later.");
     }
diff --git a/frontend/app/components/Kai/KaiStore.ts b/frontend/app/components/Kai/KaiStore.ts
index 03612bb0c..01835a9a2 100644
--- a/frontend/app/components/Kai/KaiStore.ts
+++ b/frontend/app/components/Kai/KaiStore.ts
@@ -119,23 +119,31 @@ class KaiStore {
     this.chatManager = new ChatManager(settings);
     this.chatManager.setOnMsgHook({
       msgCallback: (msg) => {
-        if (msg.stage === 'start') {
-          this.setProcessingStage({
-            ...msg,
-            content: 'Processing your request...',
-          });
-        }
-        if (msg.stage === 'chart') {
-          this.setProcessingStage(msg);
-        }
-        if (msg.stage === 'final') {
-          const msgObj = {
-            text: msg.content,
-            isUser: false,
-            messageId: msg.messageId,
+        if ('state' in msg) {
+          if (msg.state === 'running') {
+            this.setProcessingStage({ content: 'Processing your request...', stage: 'chart', messageId: Date.now().toPrecision() })
+          } else {
+            this.setProcessingStage(null)
+          }
+        } else {
+          if (msg.stage === 'start') {
+            this.setProcessingStage({
+              ...msg,
+              content: 'Processing your request...',
+            });
+          }
+          if (msg.stage === 'chart') {
+            this.setProcessingStage(msg);
+          }
+          if (msg.stage === 'final') {
+            const msgObj = {
+              text: msg.content,
+              isUser: false,
+              messageId: msg.messageId,
+            }
+            this.addMessage(msgObj);
+            this.setProcessingStage(null);
           }
-          this.addMessage(msgObj);
-          this.setProcessingStage(null);
         }
       },
       titleCallback: setTitle,
diff --git a/frontend/app/components/Kai/SocketManager.ts b/frontend/app/components/Kai/SocketManager.ts
index f40622612..7014f4f29 100644
--- a/frontend/app/components/Kai/SocketManager.ts
+++ b/frontend/app/components/Kai/SocketManager.ts
@@ -52,7 +52,7 @@ export class ChatManager {
     msgCallback,
     titleCallback,
   }: {
-    msgCallback: (msg: BotChunk) => void;
+    msgCallback: (msg: BotChunk | { state: string, type: 'state' }) => void;
     titleCallback: (title: string) => void;
   }) => {
     this.socket.on('chunk', (msg: BotChunk) => {
@@ -63,6 +63,9 @@ export class ChatManager {
       console.log('Received title:', msg);
       titleCallback(msg.content);
     });
+    this.socket.on('state', (state: { message: 'idle' | 'running' }) => {
+      msgCallback({ state: state.message, type: 'state' })
+    })
   };
 
   disconnect = () => {