diff --git a/app/client/src/ce/constants/messages.ts b/app/client/src/ce/constants/messages.ts
index 50f34544a62..98905fa5126 100644
--- a/app/client/src/ce/constants/messages.ts
+++ b/app/client/src/ce/constants/messages.ts
@@ -2520,3 +2520,10 @@ export const JS_EDITOR_SETTINGS = {
   TITLE: () => "Settings",
   ON_LOAD_TITLE: () => "Choose the functions to run on page load",
 };
+
+export const CUSTOM_WIDGET_BUILDER_TAB_TITLE = {
+  AI: () => "AI",
+  HTML: () => "HTML",
+  STYLE: () => "Style",
+  JS: () => "Javascript",
+};
diff --git a/app/client/src/ce/entities/FeatureFlag.ts b/app/client/src/ce/entities/FeatureFlag.ts
index fbfc8e0a2ca..4196cf40f46 100644
--- a/app/client/src/ce/entities/FeatureFlag.ts
+++ b/app/client/src/ce/entities/FeatureFlag.ts
@@ -41,6 +41,7 @@ export const FEATURE_FLAG = {
     "release_ide_datasource_selector_enabled",
   release_table_custom_loading_state_enabled:
     "release_table_custom_loading_state_enabled",
+  release_custom_widget_ai_builder: "release_custom_widget_ai_builder",
 } as const;
 
 export type FeatureFlag = keyof typeof FEATURE_FLAG;
@@ -77,6 +78,7 @@ export const DEFAULT_FEATURE_FLAG_VALUE: FeatureFlags = {
   release_ide_animations_enabled: false,
   release_ide_datasource_selector_enabled: false,
   release_table_custom_loading_state_enabled: false,
+  release_custom_widget_ai_builder: false,
 };
 
 export const AB_TESTING_EVENT_KEYS = {
diff --git a/app/client/src/pages/Editor/CustomWidgetBuilder/Editor/ChatBot/ChatBot.tsx b/app/client/src/pages/Editor/CustomWidgetBuilder/Editor/ChatBot/ChatBot.tsx
new file mode 100644
index 00000000000..ae3c98e9c82
--- /dev/null
+++ b/app/client/src/pages/Editor/CustomWidgetBuilder/Editor/ChatBot/ChatBot.tsx
@@ -0,0 +1,103 @@
+import React, {
+  useCallback,
+  useContext,
+  useEffect,
+  useMemo,
+  useRef,
+} from "react";
+import type { ContentProps } from "../CodeEditors/types";
+import { CustomWidgetBuilderContext } from "../..";
+import {
+  CUSTOM_WIDGET_AI_BOT_MESSAGE_RESPONSE_DEBOUNCE_TIMEOUT,
+  CUSTOM_WIDGET_AI_BOT_URL,
+  CUSTOM_WIDGET_AI_CHAT_TYPE,
+  CUSTOM_WIDGET_AI_INITIALISED_MESSAGE,
+} from "../../constants";
+import { isObject } from "lodash";
+
+export const ChatBot = (props: ContentProps) => {
+  const ref = useRef<HTMLIFrameElement>(null);
+  const lastUpdateFromBot = useRef<number>(0);
+  const { bulkUpdate, parentEntityId, uncompiledSrcDoc, widgetId } = useContext(
+    CustomWidgetBuilderContext,
+  );
+
+  const handleSrcDocUpdates = useCallback(() => {
+    // Don't send updates back to bot if the last update came from the bot within the last 100ms
+    if (
+      Date.now() - lastUpdateFromBot.current <
+      CUSTOM_WIDGET_AI_BOT_MESSAGE_RESPONSE_DEBOUNCE_TIMEOUT
+    ) {
+      return;
+    }
+
+    // Send src doc to the chatbot iframe
+    if (ref.current && ref.current.contentWindow && uncompiledSrcDoc) {
+      ref.current.contentWindow.postMessage(
+        {
+          html_code: uncompiledSrcDoc.html,
+          css_code: uncompiledSrcDoc.css,
+          js_code: uncompiledSrcDoc.js,
+          chatType: CUSTOM_WIDGET_AI_CHAT_TYPE,
+        },
+        "*",
+      );
+    }
+  }, [uncompiledSrcDoc]);
+
+  const updateContents = useCallback(
+    (
+      event: MessageEvent<
+        string | { html_code?: string; css_code?: string; js_code?: string }
+      >,
+    ) => {
+      const iframeWindow =
+        ref.current?.contentWindow || ref.current?.contentDocument?.defaultView;
+
+      // Accept messages only from the current iframe
+      if (event.source !== iframeWindow) return;
+
+      if (event.data === CUSTOM_WIDGET_AI_INITIALISED_MESSAGE) {
+        handleSrcDocUpdates();
+
+        return;
+      }
+
+      if (!bulkUpdate) return;
+
+      if (isObject(event.data)) {
+        lastUpdateFromBot.current = Date.now();
+
+        bulkUpdate({
+          html: event.data.html_code || "",
+          css: event.data.css_code || "",
+          js: event.data.js_code || "",
+        });
+      }
+    },
+    [bulkUpdate, handleSrcDocUpdates],
+  );
+
+  useEffect(
+    function addEventListenerForBotUpdates() {
+      // add a listener to update the contents
+      window.addEventListener("message", updateContents, false);
+
+      // clean up
+      return () => window.removeEventListener("message", updateContents, false);
+    },
+    [updateContents],
+  );
+
+  useEffect(handleSrcDocUpdates, [handleSrcDocUpdates]);
+
+  const instanceId = `${widgetId}-${parentEntityId}`;
+
+  const srcUrl = useMemo(() => {
+    return CUSTOM_WIDGET_AI_BOT_URL(instanceId);
+  }, [instanceId]);
+
+  return (
+    <iframe height={`${props.height}px`} ref={ref} src={srcUrl} width="100%" />
+  );
+};
diff --git a/app/client/src/pages/Editor/CustomWidgetBuilder/Editor/Header/index.tsx b/app/client/src/pages/Editor/CustomWidgetBuilder/Editor/Header/index.tsx
index 8b78e09196f..4003d2eae0c 100644
--- a/app/client/src/pages/Editor/CustomWidgetBuilder/Editor/Header/index.tsx
+++ b/app/client/src/pages/Editor/CustomWidgetBuilder/Editor/Header/index.tsx
@@ -1,7 +1,6 @@
 import React from "react";
 import styles from "./styles.module.css";
 import WidgetName from "./widgetName";
-import LayoutControls from "./layoutControls";
 import ReferenceTrigger from "./referenceTrigger";
 import { CodeTemplates } from "./CodeTemplates";
 
@@ -11,7 +10,6 @@ export default function Header() {
       <div className={styles.headerControlsLeft}>
         <WidgetName />
         <CodeTemplates />
-        <LayoutControls />
       </div>
       <div className={styles.headerControlsRight}>
         <ReferenceTrigger />
diff --git a/app/client/src/pages/Editor/CustomWidgetBuilder/Editor/Header/layoutControls.tsx b/app/client/src/pages/Editor/CustomWidgetBuilder/Editor/Header/layoutControls.tsx
deleted file mode 100644
index 1b748d6c58e..00000000000
--- a/app/client/src/pages/Editor/CustomWidgetBuilder/Editor/Header/layoutControls.tsx
+++ /dev/null
@@ -1,49 +0,0 @@
-import { SegmentedControl } from "@appsmith/ads";
-import React, { useContext } from "react";
-import { CustomWidgetBuilderContext } from "../..";
-import styles from "./styles.module.css";
-import styled from "styled-components";
-import { CUSTOM_WIDGET_FEATURE, createMessage } from "ee/constants/messages";
-import AnalyticsUtil from "ee/utils/AnalyticsUtil";
-
-const StyledSegmentedControl = styled(SegmentedControl)`
-  & .ads-v2-icon {
-    transform: rotate(90deg);
-  }
-`;
-
-export default function LayoutControls() {
-  const { selectedLayout, selectLayout, widgetId } = useContext(
-    CustomWidgetBuilderContext,
-  );
-
-  const onChange = (value: string) => {
-    selectLayout?.(value);
-
-    AnalyticsUtil.logEvent("CUSTOM_WIDGET_BUILDER_LAYOUT_CHANGED", {
-      widgetId: widgetId,
-      layoutName: value,
-    });
-  };
-
-  return (
-    <div className={styles.layoutControl}>
-      <StyledSegmentedControl
-        onChange={onChange}
-        options={[
-          {
-            label: createMessage(CUSTOM_WIDGET_FEATURE.layout.split),
-            startIcon: "layout-column-line",
-            value: "split",
-          },
-          {
-            label: createMessage(CUSTOM_WIDGET_FEATURE.layout.tab),
-            startIcon: "layout-left-2-line",
-            value: "tabs",
-          },
-        ]}
-        value={selectedLayout}
-      />
-    </div>
-  );
-}
diff --git a/app/client/src/pages/Editor/CustomWidgetBuilder/Editor/Layouts/SplitLayout/index.tsx b/app/client/src/pages/Editor/CustomWidgetBuilder/Editor/Layouts/SplitLayout/index.tsx
deleted file mode 100644
index 9f0b9d6f2f3..00000000000
--- a/app/client/src/pages/Editor/CustomWidgetBuilder/Editor/Layouts/SplitLayout/index.tsx
+++ /dev/null
@@ -1,112 +0,0 @@
-import React, { useEffect, useMemo, useRef, useState } from "react";
-import styles from "./styles.module.css";
-import type { ContentProps } from "../../CodeEditors/types";
-import { Icon } from "@appsmith/ads";
-import { set } from "lodash";
-import classNames from "classnames";
-import { useLocalStorage } from "utils/hooks/localstorage";
-
-interface Row {
-  title: string;
-  titleControls?: React.ReactNode;
-  children: (props: ContentProps) => React.ReactNode;
-}
-
-interface Props {
-  rows: Array<Row>;
-}
-
-const HEADER_HEIGHT = 38;
-
-const CUSTOM_WIDGET_SPLIT_VIEW_VISIBILITY =
-  "CUSTOM_WIDGET_SPLIT_VIEW_VISIBILITY";
-
-export default function SplitLayout(props: Props) {
-  const { rows } = props;
-  const containerRef = useRef<HTMLDivElement>(null);
-  const [height, setHeight] = React.useState(0);
-  const [loading, setLoading] = useState(true);
-  const [isRowCollapsed, setIsRowCollapsed] = useLocalStorage(
-    CUSTOM_WIDGET_SPLIT_VIEW_VISIBILITY,
-    Array(rows.length).fill(false),
-  );
-
-  useEffect(() => {
-    if (containerRef.current) {
-      setHeight(
-        window.innerHeight - containerRef.current.getBoundingClientRect().top,
-      );
-
-      setLoading(false);
-    }
-
-    const handler = () => {
-      if (containerRef.current) {
-        setHeight(
-          window.innerHeight - containerRef.current.getBoundingClientRect().top,
-        );
-      }
-    };
-
-    window.addEventListener("resize", handler);
-
-    return () => {
-      window.removeEventListener("resize", handler);
-    };
-  }, []);
-
-  const computedHeight = useMemo(() => {
-    const expandedRows = isRowCollapsed.filter((row: boolean) => !row).length;
-
-    return (height - HEADER_HEIGHT * rows.length) / expandedRows;
-  }, [isRowCollapsed, height]);
-
-  return (
-    <div className={styles.wrapper} ref={containerRef}>
-      {!loading &&
-        rows.map((row, index) => (
-          <div className={classNames(styles.editor)} key={row.title}>
-            <div
-              className={classNames(
-                styles.editorHeader,
-                isRowCollapsed[index] && index === rows.length - 1
-                  ? styles.hasBorder
-                  : "",
-              )}
-              onClick={() =>
-                setIsRowCollapsed(
-                  set([...isRowCollapsed], index, !isRowCollapsed[index]),
-                )
-              }
-            >
-              <div>{row.title}</div>
-              <div className={styles.editorHeaderControls}>
-                {row.titleControls}
-                <Icon
-                  name={
-                    !isRowCollapsed[index]
-                      ? "arrow-down-s-line"
-                      : "arrow-up-s-line"
-                  }
-                  size="lg"
-                  style={{ cursor: "pointer" }}
-                />
-              </div>
-            </div>
-
-            {!isRowCollapsed[index] && (
-              <div
-                className={styles.editorBody}
-                style={{ height: computedHeight }}
-              >
-                {row.children({
-                  height: computedHeight,
-                  width: "100%",
-                })}
-              </div>
-            )}
-          </div>
-        ))}
-    </div>
-  );
-}
diff --git a/app/client/src/pages/Editor/CustomWidgetBuilder/Editor/Layouts/SplitLayout/styles.module.css b/app/client/src/pages/Editor/CustomWidgetBuilder/Editor/Layouts/SplitLayout/styles.module.css
deleted file mode 100644
index 4d8bad08700..00000000000
--- a/app/client/src/pages/Editor/CustomWidgetBuilder/Editor/Layouts/SplitLayout/styles.module.css
+++ /dev/null
@@ -1,25 +0,0 @@
-.editorHeader {
-  height: 38px;
-  display: flex;
-  align-items: center;
-  padding: 0 10px;
-  border-top: 1px solid var(--ads-v2-color-gray-300);
-  font-weight: 500;
-  justify-content: space-between;
-  cursor: pointer;
-}
-
-.editorHeaderControls {
-  display: flex;
-  gap: 1rem;
-  justify-content: right;
-}
-
-.editorBody {
-  border-top: 1px solid var(--ads-v2-color-gray-300);
-  overflow: hidden;
-}
-
-.hasBorder {
-  border-bottom: 1px solid var(--ads-v2-color-gray-300);
-}
diff --git a/app/client/src/pages/Editor/CustomWidgetBuilder/Editor/Layouts/TabsLayout/index.tsx b/app/client/src/pages/Editor/CustomWidgetBuilder/Editor/Layouts/TabsLayout/index.tsx
index 42bab0b4974..fbd8fc6b982 100644
--- a/app/client/src/pages/Editor/CustomWidgetBuilder/Editor/Layouts/TabsLayout/index.tsx
+++ b/app/client/src/pages/Editor/CustomWidgetBuilder/Editor/Layouts/TabsLayout/index.tsx
@@ -3,6 +3,10 @@ import styles from "./styles.module.css";
 import { Tab, TabPanel, Tabs, TabsList } from "@appsmith/ads";
 import type { ContentProps } from "../../CodeEditors/types";
 import useLocalStorageState from "utils/hooks/useLocalStorageState";
+import classNames from "classnames";
+import { useFeatureFlag } from "utils/hooks/useFeatureFlag";
+import { FEATURE_FLAG } from "ee/entities/FeatureFlag";
+import { CUSTOM_WIDGET_BUILDER_TABS } from "../../../constants";
 
 interface Props {
   tabs: Array<{
@@ -17,9 +21,15 @@ const LOCAL_STORAGE_KEYS = "custom-widget-layout-tabs-state";
 export default function TabLayout(props: Props) {
   const { tabs } = props;
 
+  const isDefaultAITab = useFeatureFlag(
+    FEATURE_FLAG.release_custom_widget_ai_builder,
+  );
+
   const [selectedTab, setSelectedTab] = useLocalStorageState<string>(
     LOCAL_STORAGE_KEYS,
-    tabs[0].title,
+    isDefaultAITab
+      ? CUSTOM_WIDGET_BUILDER_TABS.AI
+      : CUSTOM_WIDGET_BUILDER_TABS.JS,
   );
 
   useEffect(() => {
@@ -88,7 +98,10 @@ export default function TabLayout(props: Props) {
           </TabsList>
           {tabs.map((tab) => (
             <TabPanel
-              className={styles.tabPanel}
+              className={classNames(styles.tabPanel, {
+                "data-[state=inactive]:hidden": selectedTab !== tab.title,
+              })}
+              forceMount
               key={tab.title}
               value={tab.title}
             >
diff --git a/app/client/src/pages/Editor/CustomWidgetBuilder/Editor/Layouts/index.tsx b/app/client/src/pages/Editor/CustomWidgetBuilder/Editor/Layouts/index.tsx
index 78456951a96..3f911235736 100644
--- a/app/client/src/pages/Editor/CustomWidgetBuilder/Editor/Layouts/index.tsx
+++ b/app/client/src/pages/Editor/CustomWidgetBuilder/Editor/Layouts/index.tsx
@@ -1,8 +1,7 @@
-import React, { useContext, useMemo } from "react";
+import React, { useContext } from "react";
 import styles from "./styles.module.css";
 import { CustomWidgetBuilderContext } from "../..";
 import TabLayout from "./TabsLayout";
-import SplitLayout from "./SplitLayout";
 import type { ContentProps } from "../CodeEditors/types";
 
 interface Props {
@@ -18,17 +17,6 @@ export default function Layout(props: Props) {
 
   const context = useContext(CustomWidgetBuilderContext);
 
-  const layout = useMemo(() => {
-    switch (context.selectedLayout) {
-      case "tabs":
-        return <TabLayout tabs={content} />;
-      case "split":
-        return <SplitLayout rows={content} />;
-      default:
-        return <div>Tab Layout</div>;
-    }
-  }, [context.selectedLayout]);
-
   return (
     <div
       className={
@@ -37,7 +25,7 @@ export default function Layout(props: Props) {
           : styles.editorLayoutFullWidth
       }
     >
-      {layout}
+      <TabLayout tabs={content} />
     </div>
   );
 }
diff --git a/app/client/src/pages/Editor/CustomWidgetBuilder/Editor/index.tsx b/app/client/src/pages/Editor/CustomWidgetBuilder/Editor/index.tsx
index d72b4381a72..1b8da78be36 100644
--- a/app/client/src/pages/Editor/CustomWidgetBuilder/Editor/index.tsx
+++ b/app/client/src/pages/Editor/CustomWidgetBuilder/Editor/index.tsx
@@ -1,4 +1,4 @@
-import React, { useContext } from "react";
+import React, { useContext, useMemo } from "react";
 import styles from "./styles.module.css";
 import Layout from "./Layouts";
 import Header from "./Header";
@@ -8,9 +8,43 @@ import StyleEditor, { TitleControls } from "./CodeEditors/StyleEditor";
 import JSEditor from "./CodeEditors/JSEditor";
 import type { ContentProps } from "./CodeEditors/types";
 import References from "./References";
+import { ChatBot } from "./ChatBot/ChatBot";
+import { useFeatureFlag } from "utils/hooks/useFeatureFlag";
+import { FEATURE_FLAG } from "ee/entities/FeatureFlag";
+import { CUSTOM_WIDGET_BUILDER_TABS } from "../constants";
 
 export default function Editor() {
   const { isReferenceOpen } = useContext(CustomWidgetBuilderContext);
+  const isAIEnabled = useFeatureFlag(
+    FEATURE_FLAG.release_custom_widget_ai_builder,
+  );
+
+  const tabs = useMemo(() => {
+    const defaultTabs = [
+      {
+        title: CUSTOM_WIDGET_BUILDER_TABS.HTML,
+        children: (props: ContentProps) => <HTMLEditor {...props} />,
+      },
+      {
+        title: CUSTOM_WIDGET_BUILDER_TABS.STYLE,
+        titleControls: <TitleControls />,
+        children: (props: ContentProps) => <StyleEditor {...props} />,
+      },
+      {
+        title: CUSTOM_WIDGET_BUILDER_TABS.JS,
+        children: (props: ContentProps) => <JSEditor {...props} />,
+      },
+    ];
+
+    if (isAIEnabled) {
+      defaultTabs.push({
+        title: CUSTOM_WIDGET_BUILDER_TABS.AI,
+        children: (props: ContentProps) => <ChatBot {...props} />,
+      });
+    }
+
+    return defaultTabs;
+  }, [isAIEnabled]);
 
   return (
     <div>
@@ -18,23 +52,7 @@ export default function Editor() {
         <Header />
       </div>
       <div className={styles.contentRightBody}>
-        <Layout
-          content={[
-            {
-              title: "HTML",
-              children: (props: ContentProps) => <HTMLEditor {...props} />,
-            },
-            {
-              title: "Style",
-              titleControls: <TitleControls />,
-              children: (props: ContentProps) => <StyleEditor {...props} />,
-            },
-            {
-              title: "Javascript",
-              children: (props: ContentProps) => <JSEditor {...props} />,
-            },
-          ]}
-        />
+        <Layout content={tabs} />
         {isReferenceOpen && <References />}
       </div>
     </div>
diff --git a/app/client/src/pages/Editor/CustomWidgetBuilder/constants.ts b/app/client/src/pages/Editor/CustomWidgetBuilder/constants.ts
index 77e18f5633f..a98b116c8df 100644
--- a/app/client/src/pages/Editor/CustomWidgetBuilder/constants.ts
+++ b/app/client/src/pages/Editor/CustomWidgetBuilder/constants.ts
@@ -1,3 +1,8 @@
+import {
+  createMessage,
+  CUSTOM_WIDGET_BUILDER_TAB_TITLE,
+} from "ee/constants/messages";
+
 export const CUSTOM_WIDGET_BUILDER_EVENTS = {
   READY: "READY",
   READY_ACK: "READY_ACK",
@@ -13,12 +18,10 @@ export const CUSTOM_WIDGET_BUILDER_EVENTS = {
 export const LOCAL_STORAGE_KEYS_IS_REFERENCE_OPEN =
   "custom-widget-builder-context-state-is-reference-open";
 
-export const LOCAL_STORAGE_KEYS_SELECTED_LAYOUT =
-  "custom-widget-builder-context-state-selected-layout";
-
 export const DEFAULT_CONTEXT_VALUE = {
   name: "",
   widgetId: "",
+  parentEntityId: "",
   srcDoc: {
     html: "",
     js: "",
@@ -61,3 +64,19 @@ export const CUSTOM_WIDGET_DEFAULT_MODEL_DOC_URL = `${CUSTOM_WIDGET_DOC_URL}#def
 export const CUSTOM_WIDGET_ONREADY_DOC_URL = `${CUSTOM_WIDGET_DOC_URL}#onready`;
 
 export const CUSTOM_WIDGET_HEIGHT_DOC_URL = `${CUSTOM_WIDGET_DOC_URL}#height`;
+
+export const CUSTOM_WIDGET_AI_CHAT_TYPE = "CUSTOM_WIDGET";
+
+export const CUSTOM_WIDGET_AI_INITIALISED_MESSAGE = "CHAT_INITIALISED";
+
+export const CUSTOM_WIDGET_AI_BOT_URL = (instanceId: string) =>
+  `https://internal.appsmith.com/app/app-builder-bot/custom-widget-bot-672b2020d37b7d0b29dcfa71?embed=true&chatType=${CUSTOM_WIDGET_AI_CHAT_TYPE}&chatInstance=${instanceId}&url=${encodeURIComponent(window.location.origin)}`;
+
+export const CUSTOM_WIDGET_AI_BOT_MESSAGE_RESPONSE_DEBOUNCE_TIMEOUT = 100;
+
+export const CUSTOM_WIDGET_BUILDER_TABS = {
+  AI: createMessage(CUSTOM_WIDGET_BUILDER_TAB_TITLE.AI),
+  HTML: createMessage(CUSTOM_WIDGET_BUILDER_TAB_TITLE.HTML),
+  STYLE: createMessage(CUSTOM_WIDGET_BUILDER_TAB_TITLE.STYLE),
+  JS: createMessage(CUSTOM_WIDGET_BUILDER_TAB_TITLE.JS),
+};
diff --git a/app/client/src/pages/Editor/CustomWidgetBuilder/types.ts b/app/client/src/pages/Editor/CustomWidgetBuilder/types.ts
index 9cc3d5e76e8..6aec321cc3c 100644
--- a/app/client/src/pages/Editor/CustomWidgetBuilder/types.ts
+++ b/app/client/src/pages/Editor/CustomWidgetBuilder/types.ts
@@ -28,9 +28,9 @@ export interface CustomWidgetBuilderContextValueType {
   //Custom widget name
   name: string;
   widgetId: string;
+  parentEntityId: string;
 
   isReferenceOpen: boolean;
-  selectedLayout: string;
 
   //Compiled src doc
   srcDoc: SrcDoc;
@@ -57,7 +57,6 @@ export interface CustomWidgetBuilderContextValueType {
 
 export interface CustomWidgetBuilderContextFunctionType {
   toggleReference: () => void;
-  selectLayout: (layout: string) => void;
   close: () => void;
   update: (editor: string, value: string) => void;
   updateModel: (model: Record<string, unknown>) => void;
diff --git a/app/client/src/pages/Editor/CustomWidgetBuilder/useCustomBuilder.tsx b/app/client/src/pages/Editor/CustomWidgetBuilder/useCustomBuilder.tsx
index 583fb787239..a1fdfc81bb7 100644
--- a/app/client/src/pages/Editor/CustomWidgetBuilder/useCustomBuilder.tsx
+++ b/app/client/src/pages/Editor/CustomWidgetBuilder/useCustomBuilder.tsx
@@ -3,7 +3,6 @@ import {
   CUSTOM_WIDGET_BUILDER_EVENTS,
   DEFAULT_CONTEXT_VALUE,
   LOCAL_STORAGE_KEYS_IS_REFERENCE_OPEN,
-  LOCAL_STORAGE_KEYS_SELECTED_LAYOUT,
 } from "./constants";
 import history from "utils/history";
 import useLocalStorageState from "utils/hooks/useLocalStorageState";
@@ -17,6 +16,8 @@ import {
 } from "./types";
 import { compileSrcDoc } from "./utility";
 import AnalyticsUtil from "ee/utils/AnalyticsUtil";
+import { useEditorType } from "ee/hooks";
+import { useParentEntityInfo } from "ee/hooks/datasourceEditorHooks";
 
 let connectionTimeout: number;
 
@@ -28,10 +29,8 @@ export function useCustomBuilder(): [CustomWidgetBuilderContextType, boolean] {
     true,
   );
 
-  const [selectedLayout, setSelectedLayout] = useLocalStorageState<string>(
-    LOCAL_STORAGE_KEYS_SELECTED_LAYOUT,
-    "tabs",
-  );
+  const editorType = useEditorType(location.pathname);
+  const { parentEntityId } = useParentEntityInfo(editorType);
 
   const [contextValue, setContextValue] =
     useState<CustomWidgetBuilderContextValueType>(DEFAULT_CONTEXT_VALUE);
@@ -93,9 +92,6 @@ export function useCustomBuilder(): [CustomWidgetBuilderContextType, boolean] {
       toggleReference: () => {
         setIsReferenceOpen(!isReferenceOpen);
       },
-      selectLayout: (layout) => {
-        setSelectedLayout(layout);
-      },
       close: () => {
         window.opener?.focus();
         window.close();
@@ -154,22 +150,16 @@ export function useCustomBuilder(): [CustomWidgetBuilderContextType, boolean] {
         });
       },
     }),
-    [
-      contextValue.uncompiledSrcDoc,
-      setIsReferenceOpen,
-      isReferenceOpen,
-      setSelectedLayout,
-    ],
+    [contextValue.uncompiledSrcDoc, setIsReferenceOpen, isReferenceOpen],
   );
 
   const context = useMemo(
     () => ({
       ...contextValue,
       isReferenceOpen,
-      selectedLayout,
       ...contextFunctions,
     }),
-    [contextValue, contextFunctions, isReferenceOpen, selectedLayout],
+    [contextValue, contextFunctions, isReferenceOpen],
   );
 
   useEffect(replay, [contextValue.srcDoc]);
@@ -186,6 +176,7 @@ export function useCustomBuilder(): [CustomWidgetBuilderContextType, boolean] {
               ...prev,
               name: event.data.name,
               widgetId: event.data.widgetId,
+              parentEntityId,
               srcDoc: event.data.srcDoc,
               uncompiledSrcDoc: event.data.uncompiledSrcDoc,
               initialSrcDoc: event.data.uncompiledSrcDoc,
@@ -223,6 +214,7 @@ export function useCustomBuilder(): [CustomWidgetBuilderContextType, boolean] {
               showConnectionLostMessage: false,
               name: event.data.name,
               widgetId: event.data.widgetId,
+              parentEntityId,
               srcDoc: event.data.srcDoc,
               uncompiledSrcDoc: event.data.uncompiledSrcDoc,
               initialSrcDoc: event.data.uncompiledSrcDoc,