diff --git a/frontend/.yarnrc.yml b/frontend/.yarnrc.yml
index f3bd10e60..da3ec2ca7 100644
--- a/frontend/.yarnrc.yml
+++ b/frontend/.yarnrc.yml
@@ -2,6 +2,6 @@ compressionLevel: 1
 
 enableGlobalCache: true
 
-nodeLinker: pnpm
+nodeLinker: node-modules
 
 yarnPath: .yarn/releases/yarn-4.7.0.cjs
diff --git a/frontend/app/components/Kai/KaiStore.ts b/frontend/app/components/Kai/KaiStore.ts
new file mode 100644
index 000000000..9d6b92fa8
--- /dev/null
+++ b/frontend/app/components/Kai/KaiStore.ts
@@ -0,0 +1,207 @@
+import { makeAutoObservable, runInAction } from 'mobx';
+import { BotChunk, ChatManager, Message } from './SocketManager';
+import { aiService } from 'App/services';
+import { toast } from 'react-toastify';
+
+class KaiStore {
+  chatManager: ChatManager | null = null;
+  processingStage: BotChunk | null = null;
+  messages: Message[] = [];
+  queryText = '';
+  loadingChat = false;
+  replacing = false;
+
+  constructor() {
+    makeAutoObservable(this);
+  }
+
+  get lastHumanMessage() {
+    let msg = null;
+    let index = null;
+    for (let i = this.messages.length - 1; i >= 0; i--) {
+      const message = this.messages[i];
+      if (message.isUser) {
+        msg = message;
+        index = i;
+        break;
+      }
+    }
+    return { msg, index };
+  }
+
+  get lastKaiMessage() {
+    let msg = null;
+    let index = null;
+    for (let i = this.messages.length - 1; i >= 0; i--) {
+      const message = this.messages[i];
+      if (!message.isUser) {
+        msg = message;
+        index = i;
+        break;
+      }
+    }
+    return { msg, index };
+  }
+
+  setQueryText = (text: string) => {
+    this.queryText = text;
+  };
+
+  setLoadingChat = (loading: boolean) => {
+    this.loadingChat = loading;
+  };
+
+  setChatManager = (chatManager: ChatManager) => {
+    this.chatManager = chatManager;
+  };
+
+  setProcessingStage = (stage: BotChunk | null) => {
+    this.processingStage = stage;
+  };
+
+  setMessages = (messages: Message[]) => {
+    this.messages = messages;
+  };
+
+  addMessage = (message: Message) => {
+    this.messages.push(message);
+  };
+
+  editMessage = (text: string) => {
+    this.setQueryText(text);
+    this.setReplacing(true);
+  };
+
+  replaceAtIndex = (message: Message, index: number) => {
+    const messages = [...this.messages];
+    messages[index] = message;
+    this.setMessages(messages);
+  };
+
+  deleteAtIndex = (indexes: number[]) => {
+    if (!indexes.length) return;
+    const messages = this.messages.filter((_, i) => !indexes.includes(i));
+    console.log(messages, indexes)
+    runInAction(() => {
+      this.messages = messages;
+    })
+  }
+
+  getChat = async (projectId: string, userId: string, threadId: string) => {
+    this.setLoadingChat(true);
+    try {
+      const res = await aiService.getKaiChat(projectId, userId, threadId);
+      if (res && res.length) {
+        this.setMessages(
+          res.map((m) => {
+            const isUser = m.role === 'human';
+            return {
+              text: m.content,
+              isUser: isUser,
+              messageId: m.message_id,
+            };
+          }),
+        );
+      }
+    } catch (e) {
+      console.error(e);
+      toast.error("Couldn't load chat history. Please try again later.");
+    } finally {
+      this.setLoadingChat(false);
+    }
+  };
+
+  createChatManager = (
+    settings: { projectId: string; userId: string; threadId: string },
+    setTitle: (title: string) => void,
+    initialMsg: string | null,
+  ) => {
+    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,
+          }
+          this.addMessage(msgObj);
+          this.setProcessingStage(null);
+        }
+      },
+      titleCallback: setTitle,
+    });
+
+    if (initialMsg) {
+      this.sendMessage(initialMsg);
+    }
+  };
+
+  setReplacing = (replacing: boolean) => {
+    this.replacing = replacing;
+  };
+
+  sendMessage = (message: string) => {
+    console.log('send')
+    if (this.chatManager) {
+      this.chatManager.sendMessage(message, this.replacing);
+    }
+    if (this.replacing) {
+      console.log(this.lastHumanMessage, this.lastKaiMessage, 'replacing these two')
+      const deleting = []
+      if (this.lastHumanMessage.index !== null) {
+        deleting.push(this.lastHumanMessage.index);
+      }
+      if (this.lastKaiMessage.index !== null) {
+        deleting.push(this.lastKaiMessage.index)
+      }
+      this.deleteAtIndex(deleting);
+      this.setReplacing(false)
+    }
+    this.addMessage({
+      text: message,
+      isUser: true,
+      messageId: Date.now().toString(),
+    });
+  };
+
+  sendMsgFeedback = (feedback: string, messageId: string) => {
+    const settings = { projectId: '2325', userId: '0' };
+    aiService
+      .feedback(
+        feedback === 'like',
+        messageId,
+        settings.projectId,
+        settings.userId,
+      )
+      .then(() => {
+        toast.success('Feedback saved.');
+      })
+      .catch((e) => {
+        console.error(e);
+        toast.error('Failed to send feedback. Please try again later.');
+      });
+  };
+
+  clearChat = () => {
+    this.setMessages([]);
+    this.setProcessingStage(null);
+    this.setLoadingChat(false);
+    this.setQueryText('');
+    if (this.chatManager) {
+      this.chatManager.disconnect();
+      this.chatManager = null;
+    }
+  };
+}
+
+export const kaiStore = new KaiStore();
diff --git a/frontend/app/components/Kai/SocketManager.ts b/frontend/app/components/Kai/SocketManager.ts
index 020cf18c8..c2e755f55 100644
--- a/frontend/app/components/Kai/SocketManager.ts
+++ b/frontend/app/components/Kai/SocketManager.ts
@@ -36,12 +36,13 @@ export class ChatManager {
     this.socket = socket;
   }
 
-  sendMesage = (message: string) => {
+  sendMessage = (message: string, isReplace = false) => {
     this.socket.emit(
       'message',
       JSON.stringify({
         message,
         threadId: this.threadId,
+        replace: isReplace,
       }),
     );
   };
@@ -78,3 +79,7 @@ export interface Message {
   isUser: boolean;
   messageId: string;
 }
+
+export interface SentMessage extends Message {
+  replace: boolean
+}
diff --git a/frontend/app/components/Kai/components/ChatInput.tsx b/frontend/app/components/Kai/components/ChatInput.tsx
index 522228b5d..e3b671c23 100644
--- a/frontend/app/components/Kai/components/ChatInput.tsx
+++ b/frontend/app/components/Kai/components/ChatInput.tsx
@@ -1,25 +1,37 @@
 import React from 'react'
 import { Button, Input } from "antd";
 import { SendHorizonal } from "lucide-react";
+import { kaiStore } from "../KaiStore";
+import { observer } from "mobx-react-lite";
 
 function ChatInput({ isLoading, onSubmit }: { isLoading?: boolean, onSubmit: (str: string) => void }) {
-  const [inputValue, setInputValue] = React.useState<string>('');
+  const inputRef = React.useRef<Input>(null);
+  const inputValue = kaiStore.queryText;
+  const setInputValue = (text: string) => {
+    kaiStore.setQueryText(text)
+  }
 
   const submit = () => {
-    onSubmit(inputValue)
-    setInputValue('')
+    if (inputValue.length > 0) {
+      onSubmit(inputValue)
+      setInputValue('')
+    }
   }
+
+  React.useEffect(() => {
+    if (inputRef.current) {
+      inputRef.current.focus()
+    }
+  }, [inputValue])
+
   return (
     <Input
+      onPressEnter={submit}
+      ref={inputRef}
       placeholder={'Ask anything about your product and users...'}
       size={'large'}
       value={inputValue}
       onChange={(e) => setInputValue(e.target.value)}
-      onKeyDown={(e) => {
-        if (e.key === 'Enter' && inputValue) {
-          submit()
-        }
-      }}
       suffix={
         <Button
           loading={isLoading}
@@ -34,4 +46,4 @@ function ChatInput({ isLoading, onSubmit }: { isLoading?: boolean, onSubmit: (st
   )
 }
 
-export default ChatInput
+export default observer(ChatInput)
diff --git a/frontend/app/components/Kai/components/ChatLog.tsx b/frontend/app/components/Kai/components/ChatLog.tsx
index 3904c9fc2..1c6024006 100644
--- a/frontend/app/components/Kai/components/ChatLog.tsx
+++ b/frontend/app/components/Kai/components/ChatLog.tsx
@@ -1,11 +1,9 @@
 import React from 'react';
 import ChatInput from './ChatInput';
 import { ChatMsg, ChatNotice } from './ChatMsg';
-import { ChatManager } from '../SocketManager';
-import type { BotChunk, Message } from '../SocketManager';
-import { aiService } from 'App/services';
-import { toast } from 'react-toastify';
 import { Loader } from 'UI';
+import { kaiStore } from '../KaiStore'
+import { observer } from 'mobx-react-lite';
 
 function ChatLog({
   projectId,
@@ -24,92 +22,28 @@ function ChatLog({
   initialMsg: string | null;
   setInitialMsg: (msg: string | null) => void;
 }) {
-  const chatManager = React.useRef<ChatManager | null>(null);
+  const messages = kaiStore.messages;
+  const loading = kaiStore.loadingChat;
   const chatRef = React.useRef<HTMLDivElement>(null);
-  const [messages, setMessages] = React.useState<Message[]>(
-    initialMsg ? [{ text: initialMsg, isUser: true, messageId: '123' }] : [],
-  );
-  const [processingStage, setProcessing] = React.useState<BotChunk | null>(
-    null,
-  );
-  const [isLoading, setLoading] = React.useState(false);
+  const processingStage = kaiStore.processingStage;
 
   React.useEffect(() => {
     //const settings = { projectId: projectId ?? 2325, userId: userId ?? 65 };
     const settings = { projectId: '2325', userId: '0', threadId, };
     if (threadId && !initialMsg) {
-      setLoading(true);
-      aiService
-        .getKaiChat(settings.projectId, settings.userId, threadId)
-        .then((res) => {
-          if (res && res.length) {
-            setMessages(
-              res.map((m) => {
-                const isUser = m.role === 'human';
-                return {
-                  text: m.content,
-                  isUser: isUser,
-                  messageId: m.message_id,
-                };
-              }),
-            );
-          }
-        })
-        .catch((e) => {
-          console.error(e);
-          toast.error("Couldn't load chat history. Please try again later.");
-        })
-        .finally(() => {
-          setLoading(false);
-        });
+      void kaiStore.getChat(settings.projectId, settings.userId, threadId)
     }
     if (threadId) {
-      chatManager.current = new ChatManager(settings);
-      chatManager.current.setOnMsgHook({
-        msgCallback: (msg) => {
-          if (msg.stage === 'chart') {
-            setProcessing(msg);
-          }
-          if (msg.stage === 'start') {
-            setProcessing({ ...msg, content: 'Processing your request...' });
-          }
-          if (msg.stage === 'final') {
-            setMessages((prev) => [
-              ...prev,
-              {
-                text: msg.content,
-                isUser: false,
-                userName: 'Kai',
-                messageId: msg.messageId,
-              },
-            ]);
-            setProcessing(null);
-          }
-        },
-        titleCallback: (title) => onTitleChange(title),
-      });
+      kaiStore.createChatManager(settings, onTitleChange, initialMsg)
     }
     return () => {
-      chatManager.current?.disconnect();
+      kaiStore.clearChat();
       setInitialMsg(null);
     };
   }, [threadId]);
 
-  React.useEffect(() => {
-    if (initialMsg) {
-      chatManager.current?.sendMesage(initialMsg);
-    }
-  }, [initialMsg]);
-
   const onSubmit = (text: string) => {
-    chatManager.current?.sendMesage(text);
-    setMessages((prev) => [
-      ...prev,
-      {
-        text,
-        isUser: true,
-      },
-    ]);
+    kaiStore.sendMessage(text)
   };
 
   React.useEffect(() => {
@@ -119,8 +53,10 @@ function ChatLog({
     });
   }, [messages.length, processingStage]);
 
+  const lastHumanMsgInd: null | number = kaiStore.lastHumanMessage.index;
+  console.log('messages', messages)
   return (
-    <Loader loading={isLoading} className={'w-full h-full'}>
+    <Loader loading={loading} className={'w-full h-full'}>
       <div
         ref={chatRef}
         className={
@@ -135,6 +71,7 @@ function ChatLog({
               isUser={msg.isUser}
               userName={userLetter}
               messageId={msg.messageId}
+              isLast={index === lastHumanMsgInd}
             />
           ))}
           {processingStage ? (
@@ -149,4 +86,4 @@ function ChatLog({
   );
 }
 
-export default ChatLog;
+export default observer(ChatLog);
diff --git a/frontend/app/components/Kai/components/ChatMsg.tsx b/frontend/app/components/Kai/components/ChatMsg.tsx
index 5186ee76e..d35bbc21c 100644
--- a/frontend/app/components/Kai/components/ChatMsg.tsx
+++ b/frontend/app/components/Kai/components/ChatMsg.tsx
@@ -4,34 +4,27 @@ import cn from 'classnames';
 import Markdown from 'react-markdown';
 import { Loader, ThumbsUp, ThumbsDown, ListRestart } from 'lucide-react';
 import { toast } from 'react-toastify';
-import { aiService } from 'App/services';
+import { kaiStore } from '../KaiStore';
 
 export function ChatMsg({
   text,
   isUser,
   userName,
   messageId,
+  isLast,
 }: {
   text: string;
   isUser: boolean;
   messageId: string;
   userName?: string;
+  isLast?: boolean;
 }) {
-  const onClick = () => {
-    toast.info('I do nothing!');
-  };
-  const onFeedback = (feedback: 'like' | 'dislike', messageId: string) => {
-    const settings = { projectId: '2325', userId: '0' };
-    aiService
-      .feedback(feedback === 'like', messageId, settings.projectId, settings.userId)
-      .then(() => {
-        toast.success('Feedback saved.');
-      })
-      .catch((e) => {
-        console.error(e);
-        toast.error('Failed to send feedback. Please try again later.');
-      });
+  const onRetry = () => {
+    kaiStore.editMessage(text)
   }
+  const onFeedback = (feedback: 'like' | 'dislike', messageId: string) => {
+    kaiStore.sendMsgFeedback(feedback, messageId);
+  };
   return (
     <div
       className={cn(
@@ -58,7 +51,19 @@ export function ChatMsg({
       )}
       <div className={'mt-1'}>
         <Markdown>{text}</Markdown>
-        {isUser ? null : (
+        {isUser ? (
+          isLast ? (
+            <div
+              onClick={onRetry}
+              className={
+                'flex items-center gap-2 px-2 rounded-lg border border-gray-medium text-sm cursor-pointer hover:border-main hover:text-main w-fit'
+              }
+            >
+              <ListRestart size={16} />
+              <div>Edit</div>
+            </div>
+          ) : null
+        ) : (
           <div className={'flex items-center gap-2'}>
             <IconButton onClick={() => onFeedback('like', messageId)}>
               <ThumbsUp size={16} />
@@ -66,16 +71,6 @@ export function ChatMsg({
             <IconButton onClick={() => onFeedback('dislike', messageId)}>
               <ThumbsDown size={16} />
             </IconButton>
-
-            <div
-              onClick={onClick}
-              className={
-                'flex items-center gap-2 px-2 rounded-lg border border-gray-medium text-sm cursor-pointer hover:border-main hover:text-main'
-              }
-            >
-              <ListRestart size={16} />
-              <div>Retry</div>
-            </div>
           </div>
         )}
       </div>
diff --git a/frontend/app/mstore/index.tsx b/frontend/app/mstore/index.tsx
index f380c727f..ea9bd5631 100644
--- a/frontend/app/mstore/index.tsx
+++ b/frontend/app/mstore/index.tsx
@@ -33,8 +33,6 @@ import UiPlayerStore from './uiPlayerStore';
 import userStore from './userStore';
 import UxtestingStore from './uxtestingStore';
 import WeeklyReportStore from './weeklyReportConfigStore';
-import KaiStore from './kaiStore';
-
 import logger from '@/logger';
 
 const projectStore = new ProjectsStore();
@@ -112,7 +110,6 @@ export class RootStore {
   searchStoreLive: SearchStoreLive;
   integrationsStore: IntegrationsStore;
   projectsStore: ProjectsStore;
-  kaiStore: KaiStore;
 
   constructor() {
     this.dashboardStore = new DashboardStore();
@@ -145,7 +142,6 @@ export class RootStore {
     this.searchStore = searchStore;
     this.searchStoreLive = searchStoreLive;
     this.integrationsStore = new IntegrationsStore();
-    this.kaiStore = new KaiStore();
   }
 
   initClient() {
diff --git a/frontend/app/mstore/kaiStore.ts b/frontend/app/mstore/kaiStore.ts
deleted file mode 100644
index 75eadfc0f..000000000
--- a/frontend/app/mstore/kaiStore.ts
+++ /dev/null
@@ -1,75 +0,0 @@
-import { makeAutoObservable } from 'mobx';
-
-type threadId = number | string;
-
-export interface RecentChat {
-  title: string;
-  thread_id: threadId;
-}
-
-export interface ChartData {
-  timestamp: number,
-  [key: string]: number
-}
-
-export interface Message {
-  role: 'human' | 'kai';
-  content: string;
-  chart?: { type: string, data?: ChartData[] }
-  message_id: number;
-}
-
-export default class KaiStore {
-  recentChats: RecentChat[] = [];
-  chats: Record<threadId, Message[]> = {};
-  activeThreadId: threadId = 0;
-  loadingHistory = false;
-  loadingAnswer = false;
-  status = 'thinking';
-
-  constructor() {
-    makeAutoObservable(this);
-  }
-
-  setActiveThreadId = (threadId: threadId) => {
-    this.activeThreadId = threadId;
-  }
-
-  setStatus = (status: string) => {
-    this.status = status;
-  }
-
-  setLoadingHistory = (loading: boolean) => {
-    this.loadingHistory = loading;
-  }
-
-  setLoadingAnswer = (loading: boolean) => {
-    this.loadingAnswer = loading;
-  }
-
-  addRecentChat = (chat: RecentChat) => {
-    this.recentChats.push(chat);
-  }
-
-  setRecentChats = (chats: RecentChat[]) => {
-    this.recentChats = chats;
-  }
-
-  addMessage = (threadId: number, message: Message) => {
-    const id = `_${threadId}` as threadId;
-    if (!this.chats[id]) {
-      this.chats[id] = [];
-    }
-    this.chats[id].push(message);
-  }
-
-  setMessages = (threadId: number, messages: Message[]) => {
-    const id = `_${threadId}` as threadId;
-    this.chats[id] = messages;
-  }
-
-  getMessages = (threadId: number) => {
-    const id = `_${threadId}` as threadId;
-    return this.chats[id] || [];
-  }
-}
\ No newline at end of file