diff --git a/web/src/api/buster_socket/chats/chatResponses.ts b/web/src/api/buster_socket/chats/chatResponses.ts
index 32390c731..ea460af71 100644
--- a/web/src/api/buster_socket/chats/chatResponses.ts
+++ b/web/src/api/buster_socket/chats/chatResponses.ts
@@ -10,7 +10,6 @@ export enum ChatsResponses {
   '/chats/list:getChatsList' = '/chats/list:getChatsList',
   '/chats/unsubscribe:unsubscribe' = '/chats/unsubscribe:unsubscribe',
   '/chats/get:getChat' = '/chats/get:getChat',
-  '/chats/get:getChatAsset' = '/chats/get:getChatAsset',
   '/chats/post:initializeChat' = '/chats/post:initializeChat',
   '/chats/post:generatingTitle' = '/chats/post:generatingTitle',
   '/chats/post:generatingResponseMessage' = '/chats/post:generatingResponseMessage',
@@ -43,12 +42,6 @@ export type Chat_unsubscribe = {
   onError?: (d: unknown | RustApiError) => void;
 };
 
-export type Chat_getChatAsset = {
-  route: '/chats/get:getChatAsset';
-  callback: (d: BusterChat) => void;
-  onError?: (d: unknown | RustApiError) => void;
-};
-
 /***** CHAT PROGRESS EVENTS START ******/
 
 export type ChatPost_initializeChat = {
@@ -87,7 +80,6 @@ export type ChatResponseTypes =
   | ChatList_getChatsList
   | Chat_unsubscribe
   | Chat_getChat
-  | Chat_getChatAsset
   | ChatPost_initializeChat
   | ChatPost_generatingTitle
   | ChatPost_generatingResponseMessage
diff --git a/web/src/app/app/_components/NewChatModal/NewChatModal.tsx b/web/src/app/app/_components/NewChatModal/NewChatModal.tsx
index 8d318ec9c..61956e6d8 100644
--- a/web/src/app/app/_components/NewChatModal/NewChatModal.tsx
+++ b/web/src/app/app/_components/NewChatModal/NewChatModal.tsx
@@ -1,5 +1,5 @@
 import React, { useEffect, useMemo, useRef, useState } from 'react';
-import { Button, Modal, Input, InputRef, ConfigProvider, Divider, ThemeConfig } from 'antd';
+import { Button, Modal, Input, InputRef, ConfigProvider, ThemeConfig } from 'antd';
 import { AppMaterialIcons } from '@/components/icons';
 import { useMemoizedFn, useMount, useThrottleFn } from 'ahooks';
 import { useAntToken } from '@/styles/useAntToken';
@@ -71,9 +71,10 @@ export const NewChatModal = React.memo<{
   const { run: debouncedGetSuggestedChatPrompts } = useThrottleFn(
     async (v: string) => {
       try {
-        const prompts = await getSuggestedChatPrompts(v);
-        setSuggestedPrompts(prompts);
-        return prompts;
+        // const prompts = await getSuggestedChatPrompts(v);
+        // setSuggestedPrompts(prompts);
+        // return prompts;
+        return [];
       } catch (e) {
         openErrorNotification(e);
       }
@@ -87,10 +88,6 @@ export const NewChatModal = React.memo<{
     });
   });
 
-  const onCloseOrCancel = useMemoizedFn(() => {
-    onClose();
-  });
-
   useEffect(() => {
     if (open) {
       if (defaultSuggestedPrompts.length === 0) {
@@ -112,9 +109,9 @@ export const NewChatModal = React.memo<{
     <ConfigProvider theme={themeConfig}>
       <Modal
         open={open}
-        onCancel={onCloseOrCancel}
+        onCancel={onClose}
         closable={false}
-        onClose={onCloseOrCancel}
+        onClose={onClose}
         width={hasDatasets ? 725 : 350}
         destroyOnClose={true}
         footer={null}
@@ -136,6 +133,7 @@ export const NewChatModal = React.memo<{
               lastKeyPressed={lastKeyPressed}
               activeItem={activeItem}
               prompt={prompt}
+              selectedChatDataSource={selectedChatDataSource}
               setPrompt={setPrompt}
             />
           </div>
@@ -179,6 +177,10 @@ const NewChatInput: React.FC<{
   activeItem: number | null;
   prompt: string;
   setPrompt: (prompt: string) => void;
+  selectedChatDataSource: {
+    id: string;
+    name: string;
+  } | null;
 }> = React.memo(
   ({
     setSuggestedPrompts,
@@ -187,7 +189,8 @@ const NewChatInput: React.FC<{
     shownPrompts,
     lastKeyPressed,
     prompt,
-    setPrompt
+    setPrompt,
+    selectedChatDataSource
   }) => {
     const token = useAntToken();
     const inputRef = useRef<InputRef>(null);
@@ -195,10 +198,12 @@ const NewChatInput: React.FC<{
     const onSelectSearchAsset = useBusterNewChatContextSelector((x) => x.onSelectSearchAsset);
     const [loadingNewChat, setLoadingNewChat] = useState(false);
 
+    console.log(selectedChatDataSource);
+
     const onStartNewChatPreflight = useMemoizedFn(async () => {
       setLoadingNewChat(true);
-      await onStartNewChat(prompt);
-      await timeout(400);
+      await onStartNewChat({ prompt, datasetId: selectedChatDataSource?.id });
+      await timeout(380);
       setPrompt('');
       setLoadingNewChat(false);
     });
@@ -269,6 +274,8 @@ const NewChatInput: React.FC<{
         <Button
           type="primary"
           size="middle"
+          // color="default"
+          // variant="solid"
           icon={<AppMaterialIcons icon="arrow_forward" size={token.fontSizeLG} />}
           loading={loadingNewChat}
           disabled={!inputHasText(prompt)}
diff --git a/web/src/app/app/_layouts/ChatLayout/ChatContainer/ChatContent/ChatContent.tsx b/web/src/app/app/_layouts/ChatLayout/ChatContainer/ChatContent/ChatContent.tsx
index 97129d237..2ea926152 100644
--- a/web/src/app/app/_layouts/ChatLayout/ChatContainer/ChatContent/ChatContent.tsx
+++ b/web/src/app/app/_layouts/ChatLayout/ChatContainer/ChatContent/ChatContent.tsx
@@ -10,20 +10,24 @@ export const ChatContent: React.FC<{ chatContentRef: React.RefObject<HTMLDivElem
     const chatMessageIds = useChatIndividualContextSelector((state) => state.chatMessageIds);
     // const chatMessages = useBusterChatContextSelector((state) => state.chatsMessages);
 
+    const autoClass = 'mx-auto max-w-[600px] w-full';
+
     return (
       <div className="flex h-full w-full flex-col overflow-hidden">
         <div ref={chatContentRef} className="h-full w-full overflow-y-auto">
           <div className="pb-8">
             {chatMessageIds?.map((messageId) => (
               <div key={messageId} className={styles.messageBlock}>
-                <div className="mx-auto max-w-[600px]">
+                <div className={autoClass}>
                   <ChatMessageBlock key={messageId} messageId={messageId} />
                 </div>
               </div>
             ))}
           </div>
         </div>
-        <ChatInput />
+        <div className={autoClass}>
+          <ChatInput />
+        </div>
       </div>
     );
   });
diff --git a/web/src/app/app/_layouts/ChatLayout/ChatContainer/ChatContent/ChatInput/useChatInputFlow.ts b/web/src/app/app/_layouts/ChatLayout/ChatContainer/ChatContent/ChatInput/useChatInputFlow.ts
index 7d8544936..7f202ff9d 100644
--- a/web/src/app/app/_layouts/ChatLayout/ChatContainer/ChatContent/ChatInput/useChatInputFlow.ts
+++ b/web/src/app/app/_layouts/ChatLayout/ChatContainer/ChatContent/ChatInput/useChatInputFlow.ts
@@ -65,7 +65,7 @@ export const useChatInputFlow = ({
         break;
 
       case 'new':
-        await onStartNewChat(inputValue);
+        await onStartNewChat({ prompt: inputValue });
         break;
 
       default:
diff --git a/web/src/context/BusterWebSocket/helpers.ts b/web/src/context/BusterWebSocket/helpers.ts
index 86dfb6ee8..ab1afe22f 100644
--- a/web/src/context/BusterWebSocket/helpers.ts
+++ b/web/src/context/BusterWebSocket/helpers.ts
@@ -15,6 +15,7 @@ import type {
   BusterSocketResponseBase,
   BusterSocketResponseMessage
 } from '@/api/buster_socket/base_interfaces';
+import { ChatsResponses } from '@/api/buster_socket/chats';
 
 export const createBusterResponse = (
   message: BusterSocketResponseMessage
@@ -45,7 +46,8 @@ const isKnownMessageRoute = (parsedMessage: BusterSocketResponseMessage) => {
     ...PermissionsResponses,
     ...TeamResponses,
     ...SearchResponses,
-    ...OrganizationResponses
+    ...OrganizationResponses,
+    ...ChatsResponses
   };
   const event = parsedMessage?.event;
   const route = parsedMessage?.route;
diff --git a/web/src/context/Chats/ChatProvider/useChatSelectors.ts b/web/src/context/Chats/ChatProvider/useChatSelectors.ts
index 0c18512c7..28a3adebf 100644
--- a/web/src/context/Chats/ChatProvider/useChatSelectors.ts
+++ b/web/src/context/Chats/ChatProvider/useChatSelectors.ts
@@ -11,7 +11,7 @@ export const useChatSelectors = ({
   chatsRef: MutableRefObject<Record<string, IBusterChat>>;
   chatsMessagesRef: MutableRefObject<Record<string, IBusterChatMessage>>;
 }) => {
-  const getChatMemoized = useMemoizedFn((chatId: string) => {
+  const getChatMemoized = useMemoizedFn((chatId: string): IBusterChat | undefined => {
     return chatsRef.current[chatId];
   });
 
@@ -24,7 +24,7 @@ export const useChatSelectors = ({
   );
 
   const getChatMessage = useCallback(
-    (messageId: string): IBusterChatMessage => {
+    (messageId: string): IBusterChatMessage | undefined => {
       return chatsMessagesRef.current[messageId];
     },
     [chatsMessagesRef, isPending]
diff --git a/web/src/context/Chats/NewChatProvider/NewChatProvider.tsx b/web/src/context/Chats/NewChatProvider/NewChatProvider.tsx
index 4364922fe..70e297917 100644
--- a/web/src/context/Chats/NewChatProvider/NewChatProvider.tsx
+++ b/web/src/context/Chats/NewChatProvider/NewChatProvider.tsx
@@ -16,7 +16,9 @@ export const useBusterNewChat = () => {
     completeChatCallback,
     startListeningForChatProgress,
     stopListeningForChatProgress,
-    stopChatCallback
+    stopChatCallback,
+    initializeChatCallback,
+    replaceMessageCallback
   } = useChatUpdateMessage();
 
   const onSelectSearchAsset = useMemoizedFn(async (asset: BusterSearchResult) => {
@@ -24,47 +26,96 @@ export const useBusterNewChat = () => {
     await new Promise((resolve) => setTimeout(resolve, 1000));
   });
 
-  const onStartNewChat = useMemoizedFn(async (prompt: string) => {
-    console.log('start new chat');
-    startListeningForChatProgress();
-    const result = await busterSocket.emitAndOnce({
-      emitEvent: {
-        route: '/chats/post',
-        payload: {
-          dataset_id: null, //TODO: add selected dataset id
-          prompt
+  const onStartNewChat = useMemoizedFn(
+    async ({
+      prompt,
+      datasetId,
+      metricId,
+      dashboardId
+    }: {
+      prompt: string;
+      datasetId?: string;
+      metricId?: string;
+      dashboardId?: string;
+    }) => {
+      startListeningForChatProgress();
+
+      await busterSocket.emitAndOnce({
+        emitEvent: {
+          route: '/chats/post',
+          payload: {
+            dataset_id: datasetId, //TODO: add selected dataset id
+            prompt,
+            metric_id: metricId,
+            dashboard_id: dashboardId
+          }
+        },
+        responseEvent: {
+          route: '/chats/post:initializeChat',
+          callback: initializeChatCallback
         }
-      },
-      responseEvent: {
-        route: '/chats/post:complete',
-        callback: completeChatCallback
-      }
-    });
-    stopListeningForChatProgress();
-  });
+      });
+
+      busterSocket
+        .once({
+          route: '/chats/post:complete',
+          callback: completeChatCallback
+        })
+        .then(() => {
+          stopListeningForChatProgress();
+        });
+    }
+  );
 
   const onStartChatFromFile = useMemoizedFn(
     async ({}: { prompt: string; fileId: string; fileType: FileType }) => {
       console.log('start chat from file');
+
       await new Promise((resolve) => setTimeout(resolve, 1000));
     }
   );
 
   const onReplaceMessageInChat = useMemoizedFn(
-    async ({ prompt, messageId }: { prompt: string; messageId: string }) => {
-      console.log('replace message in chat');
-      await new Promise((resolve) => setTimeout(resolve, 1000));
+    async ({
+      prompt,
+      messageId,
+      chatId
+    }: {
+      prompt: string;
+      messageId: string;
+      chatId: string;
+    }) => {
+      startListeningForChatProgress();
+      replaceMessageCallback({
+        prompt,
+        messageId
+      });
+      await busterSocket.emitAndOnce({
+        emitEvent: {
+          route: '/chats/post',
+          payload: {
+            prompt,
+            message_id: messageId,
+            chat_id: chatId
+          }
+        },
+        responseEvent: {
+          route: '/chats/post:complete',
+          callback: completeChatCallback
+        }
+      });
+      stopListeningForChatProgress();
     }
   );
 
   const onFollowUpChat = useMemoizedFn(
     async ({ prompt, chatId }: { prompt: string; chatId: string }) => {
       startListeningForChatProgress();
-      const result = await busterSocket.emitAndOnce({
+
+      await busterSocket.emitAndOnce({
         emitEvent: {
           route: '/chats/post',
           payload: {
-            dataset_id: null,
             prompt,
             chat_id: chatId
           }
diff --git a/web/src/context/Chats/NewChatProvider/useChatUpdateMessage.ts b/web/src/context/Chats/NewChatProvider/useChatUpdateMessage.ts
index 544e93ed0..ee2e0cde1 100644
--- a/web/src/context/Chats/NewChatProvider/useChatUpdateMessage.ts
+++ b/web/src/context/Chats/NewChatProvider/useChatUpdateMessage.ts
@@ -9,25 +9,29 @@ import {
 } from '@/api/buster_socket/chats';
 import { updateChatToIChat } from '@/utils/chat';
 import { useAutoAppendThought } from './useAutoAppendThought';
-import { useHotkeys } from 'react-hotkeys-hook';
+import { useAppLayoutContextSelector } from '@/context/BusterAppLayout';
+import { BusterRoutes } from '@/routes';
 
 export const useChatUpdateMessage = () => {
   const busterSocket = useBusterWebSocket();
+  const onChangePage = useAppLayoutContextSelector((x) => x.onChangePage);
   const onUpdateChat = useBusterChatContextSelector((x) => x.onUpdateChat);
   const getChatMemoized = useBusterChatContextSelector((x) => x.getChatMemoized);
   const onUpdateChatMessage = useBusterChatContextSelector((x) => x.onUpdateChatMessage);
   const getChatMessageMemoized = useBusterChatContextSelector((x) => x.getChatMessageMemoized);
   const onBulkSetChatMessages = useBusterChatContextSelector((x) => x.onBulkSetChatMessages);
+  const onToggleChatsModal = useAppLayoutContextSelector((s) => s.onToggleChatsModal);
 
   const { autoAppendThought } = useAutoAppendThought();
 
   const _generatingTitleCallback = useMemoizedFn((d: ChatEvent_GeneratingTitle) => {
     const { chat_id, title, title_chunk } = d;
     const isCompleted = d.progress === 'completed';
-    const currentTitle = getChatMemoized(chat_id)?.title;
+    const currentChat = getChatMemoized(chat_id)!;
+    const currentTitle = currentChat?.title;
     const newTitle = isCompleted ? title : currentTitle + title_chunk;
     onUpdateChat({
-      ...getChatMemoized(chat_id),
+      ...currentChat,
       title: newTitle
     });
   });
@@ -49,7 +53,7 @@ export const useChatUpdateMessage = () => {
   const _generatingReasoningMessageCallback = useMemoizedFn(
     (d: ChatEvent_GeneratingReasoningMessage) => {
       const { message_id, reasoning, chat_id } = d;
-      const currentReasoning = getChatMessageMemoized(message_id)?.reasoning;
+      const currentReasoning = getChatMessageMemoized(message_id)!.reasoning;
       const isNewMessage = !currentReasoning?.some(({ id }) => id === message_id);
       const updatedReasoning = isNewMessage
         ? [...currentReasoning, reasoning]
@@ -76,6 +80,35 @@ export const useChatUpdateMessage = () => {
     });
   });
 
+  const initializeChatCallback = useMemoizedFn((d: BusterChat) => {
+    const { iChat, iChatMessages } = updateChatToIChat(d);
+    onBulkSetChatMessages(iChatMessages);
+    onUpdateChat(iChat);
+
+    onChangePage({
+      route: BusterRoutes.APP_CHAT_ID,
+      chatId: iChat.id
+    });
+    onToggleChatsModal(false);
+  });
+
+  const replaceMessageCallback = useMemoizedFn(
+    ({ prompt, messageId }: { prompt: string; messageId: string }) => {
+      const currentMessage = getChatMessageMemoized(messageId);
+      const currentRequestMessage = currentMessage?.request_message!;
+
+      onUpdateChatMessage({
+        id: messageId,
+        request_message: {
+          ...currentRequestMessage,
+          request: prompt
+        },
+        reasoning: [],
+        response_messages: []
+      });
+    }
+  );
+
   const listenForGeneratingTitle = useMemoizedFn(() => {
     busterSocket.on({
       route: '/chats/post:generatingTitle',
@@ -131,9 +164,11 @@ export const useChatUpdateMessage = () => {
   });
 
   return {
+    initializeChatCallback,
     completeChatCallback,
     startListeningForChatProgress,
     stopListeningForChatProgress,
-    stopChatCallback
+    stopChatCallback,
+    replaceMessageCallback
   };
 };
diff --git a/web/temp.md b/web/temp.md
index 98939deaa..6f0e32942 100644
--- a/web/temp.md
+++ b/web/temp.md
@@ -69,7 +69,6 @@ export enum ChatsResponses {
   '/chats/list:getChatsList' = '/chats/list:getChatsList',
   '/chats/unsubscribe:unsubscribe' = '/chats/unsubscribe:unsubscribe',
   '/chats/get:getChat' = '/chats/get:getChat',
-  '/chats/get:getChatAsset' = '/chats/get:getChatAsset',
   '/chats/post:initializeChat' = '/chats/post:initializeChat',
   '/chats/post:generatingTitle' = '/chats/post:generatingTitle'
 }