diff --git a/src/plugins/vis_builder/common/vis_builder_saved_object_attributes.ts b/src/plugins/vis_builder/common/vis_builder_saved_object_attributes.ts
index 243e455d7157..3e157f9bf6df 100644
--- a/src/plugins/vis_builder/common/vis_builder_saved_object_attributes.ts
+++ b/src/plugins/vis_builder/common/vis_builder_saved_object_attributes.ts
@@ -13,6 +13,7 @@ export interface VisBuilderSavedObjectAttributes extends SavedObjectAttributes {
   visualizationState?: string;
   updated_at?: string;
   styleState?: string;
+  uiState?: string;
   version: number;
   searchSourceFields?: {
     index?: string;
diff --git a/src/plugins/vis_builder/public/application/components/right_nav.tsx b/src/plugins/vis_builder/public/application/components/right_nav.tsx
index 5d3f298c6bf3..fde4f3110d1c 100644
--- a/src/plugins/vis_builder/public/application/components/right_nav.tsx
+++ b/src/plugins/vis_builder/public/application/components/right_nav.tsx
@@ -122,6 +122,6 @@ const OptionItem = ({ icon, title }: { icon: IconType; title: string }) => (
   </>
 );
 
-// The app uses EuiResizableContainer that triggers a rerender for ever mouseover action.
+// The app uses EuiResizableContainer that triggers a rerender for every mouseover action.
 // To prevent this child component from unnecessarily rerendering in that instance, it needs to be memoized
 export const RightNav = React.memo(RightNavUI);
diff --git a/src/plugins/vis_builder/public/application/components/workspace.tsx b/src/plugins/vis_builder/public/application/components/workspace.tsx
index 214a735f6ba8..43f3ad04fe0f 100644
--- a/src/plugins/vis_builder/public/application/components/workspace.tsx
+++ b/src/plugins/vis_builder/public/application/components/workspace.tsx
@@ -10,7 +10,7 @@ import { useOpenSearchDashboards } from '../../../../opensearch_dashboards_react
 import { IExpressionLoaderParams } from '../../../../expressions/public';
 import { VisBuilderServices } from '../../types';
 import { validateSchemaState, validateAggregations } from '../utils/validations';
-import { useTypedSelector } from '../utils/state_management';
+import { useTypedDispatch, useTypedSelector, setUIStateState } from '../utils/state_management';
 import { useAggs, useVisualizationType } from '../utils/use';
 import { PersistedState } from '../../../../visualizations/public';
 
@@ -39,8 +39,25 @@ export const WorkspaceUI = () => {
     timeRange: data.query.timefilter.timefilter.getTime(),
   });
   const rootState = useTypedSelector((state) => state);
-  // Visualizations require the uiState to persist even when the expression changes
-  const uiState = useMemo(() => new PersistedState(), []);
+  const dispatch = useTypedDispatch();
+  // Visualizations require the uiState object to persist even when the expression changes
+  // eslint-disable-next-line react-hooks/exhaustive-deps
+  const uiState = useMemo(() => new PersistedState(rootState.ui), []);
+
+  useEffect(() => {
+    if (rootState.metadata.editor.state === 'loaded') {
+      uiState.setSilent(rootState.ui);
+    }
+    // To update uiState once saved object data is loaded
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [rootState.metadata.editor.state, uiState]);
+
+  useEffect(() => {
+    uiState.on('change', (args) => {
+      // Store changes to UI state
+      dispatch(setUIStateState(uiState.toJSON()));
+    });
+  }, [dispatch, uiState]);
 
   useEffect(() => {
     async function loadExpression() {
@@ -133,6 +150,6 @@ export const WorkspaceUI = () => {
   );
 };
 
-// The app uses EuiResizableContainer that triggers a rerender for ever mouseover action.
+// The app uses EuiResizableContainer that triggers a rerender for every mouseover action.
 // To prevent this child component from unnecessarily rerendering in that instance, it needs to be memoized
 export const Workspace = React.memo(WorkspaceUI);
diff --git a/src/plugins/vis_builder/public/application/utils/schema.json b/src/plugins/vis_builder/public/application/utils/schema.json
index 9effed97b2be..7cf8bbc2534f 100644
--- a/src/plugins/vis_builder/public/application/utils/schema.json
+++ b/src/plugins/vis_builder/public/application/utils/schema.json
@@ -1,28 +1,47 @@
 {
-    "type": "object",
-    "properties": {
-      "styleState": {
-        "type": "object"
-      },
-      "visualizationState": {
-        "type": "object",
-        "properties": {
-          "activeVisualization": {
-            "type": "object",
-            "properties": {
-              "name": { "type": "string" },
-              "aggConfigParams": { "type": "array" }
+  "type": "object",
+  "properties": {
+    "styleState": {
+      "type": "object"
+    },
+    "visualizationState": {
+      "type": "object",
+      "properties": {
+        "activeVisualization": {
+          "type": "object",
+          "properties": {
+            "name": {
+              "type": "string"
             },
-            "required": ["name", "aggConfigParams"],
-            "additionalProperties": false
+            "aggConfigParams": {
+              "type": "array"
+            }
           },
-          "indexPattern": { "type": "string" },
-          "searchField": { "type": "string" }
+          "required": [
+            "name",
+            "aggConfigParams"
+          ],
+          "additionalProperties": false
         },
-        "required": ["searchField"],
-        "additionalProperties": false
-      }
+        "indexPattern": {
+          "type": "string"
+        },
+        "searchField": {
+          "type": "string"
+        }
+      },
+      "required": [
+        "searchField"
+      ],
+      "additionalProperties": false
     },
-    "required": ["styleState", "visualizationState"],
-    "additionalProperties": false
+    "uiState": {
+      "type": "object"
+    }
+  },
+  "required": [
+    "styleState",
+    "visualizationState"
+  ],
+  "additionalProperties": false
 }
\ No newline at end of file
diff --git a/src/plugins/vis_builder/public/application/utils/state_management/metadata_slice.ts b/src/plugins/vis_builder/public/application/utils/state_management/metadata_slice.ts
index 05ceb324aaa1..880c15f3e44a 100644
--- a/src/plugins/vis_builder/public/application/utils/state_management/metadata_slice.ts
+++ b/src/plugins/vis_builder/public/application/utils/state_management/metadata_slice.ts
@@ -11,7 +11,7 @@ import { VisBuilderServices } from '../../../types';
  * Clean state: when viz finished loading and ready to be edited
  * Dirty state: when there are changes applied to the viz after it finished loading
  */
-type EditorState = 'loading' | 'clean' | 'dirty';
+type EditorState = 'loading' | 'loaded' | 'clean' | 'dirty';
 
 export interface MetadataState {
   editor: {
diff --git a/src/plugins/vis_builder/public/application/utils/state_management/preload.ts b/src/plugins/vis_builder/public/application/utils/state_management/preload.ts
index 43aa2e7b8ede..f7a0f6bd7ad3 100644
--- a/src/plugins/vis_builder/public/application/utils/state_management/preload.ts
+++ b/src/plugins/vis_builder/public/application/utils/state_management/preload.ts
@@ -8,6 +8,7 @@ import { VisBuilderServices } from '../../..';
 import { getPreloadedState as getPreloadedStyleState } from './style_slice';
 import { getPreloadedState as getPreloadedVisualizationState } from './visualization_slice';
 import { getPreloadedState as getPreloadedMetadataState } from './metadata_slice';
+import { getPreloadedState as getPreloadedUIState } from './ui_state_slice';
 import { RootState } from './store';
 
 export const getPreloadedState = async (
@@ -16,10 +17,12 @@ export const getPreloadedState = async (
   const styleState = await getPreloadedStyleState(services);
   const visualizationState = await getPreloadedVisualizationState(services);
   const metadataState = await getPreloadedMetadataState(services);
+  const uiState = await getPreloadedUIState(services);
 
   return {
     style: styleState,
     visualization: visualizationState,
     metadata: metadataState,
+    ui: uiState,
   };
 };
diff --git a/src/plugins/vis_builder/public/application/utils/state_management/redux_persistence.test.tsx b/src/plugins/vis_builder/public/application/utils/state_management/redux_persistence.test.tsx
index 91f760bbf231..a46d5c027656 100644
--- a/src/plugins/vis_builder/public/application/utils/state_management/redux_persistence.test.tsx
+++ b/src/plugins/vis_builder/public/application/utils/state_management/redux_persistence.test.tsx
@@ -18,6 +18,7 @@ describe('test redux state persistence', () => {
       style: 'style',
       visualization: 'visualization',
       metadata: 'metadata',
+      ui: 'ui',
     };
   });
 
@@ -33,6 +34,7 @@ describe('test redux state persistence', () => {
         editor: { errors: {}, state: 'loading' },
         originatingApp: undefined,
       },
+      ui: {},
     };
 
     const returnStates = await loadReduxState(mockServices);
diff --git a/src/plugins/vis_builder/public/application/utils/state_management/redux_persistence.ts b/src/plugins/vis_builder/public/application/utils/state_management/redux_persistence.ts
index a531986a9ac9..3ebfa47268ec 100644
--- a/src/plugins/vis_builder/public/application/utils/state_management/redux_persistence.ts
+++ b/src/plugins/vis_builder/public/application/utils/state_management/redux_persistence.ts
@@ -12,22 +12,21 @@ export const loadReduxState = async (services: VisBuilderServices) => {
     const serializedState = services.osdUrlStateStorage.get<RootState>('_a');
     if (serializedState !== null) return serializedState;
   } catch (err) {
-    /* eslint-disable no-console */
+    // eslint-disable-next-line no-console
     console.error(err);
-    /* eslint-enable no-console */
   }
 
   return await getPreloadedState(services);
 };
 
 export const persistReduxState = (
-  { style, visualization, metadata },
+  { style, visualization, metadata, ui }: RootState,
   services: VisBuilderServices
 ) => {
   try {
     services.osdUrlStateStorage.set<RootState>(
       '_a',
-      { style, visualization, metadata },
+      { style, visualization, metadata, ui },
       {
         replace: true,
       }
diff --git a/src/plugins/vis_builder/public/application/utils/state_management/store.ts b/src/plugins/vis_builder/public/application/utils/state_management/store.ts
index f1b1c0eeae2a..820631f7568e 100644
--- a/src/plugins/vis_builder/public/application/utils/state_management/store.ts
+++ b/src/plugins/vis_builder/public/application/utils/state_management/store.ts
@@ -8,12 +8,14 @@ import { isEqual } from 'lodash';
 import { reducer as styleReducer } from './style_slice';
 import { reducer as visualizationReducer } from './visualization_slice';
 import { reducer as metadataReducer } from './metadata_slice';
+import { reducer as uiStateReducer } from './ui_state_slice';
 import { VisBuilderServices } from '../../..';
 import { loadReduxState, persistReduxState } from './redux_persistence';
 import { handlerEditorState } from './handlers/editor_state';
 import { handlerParentAggs } from './handlers/parent_aggs';
 
 const rootReducer = combineReducers({
+  ui: uiStateReducer,
   style: styleReducer,
   visualization: visualizationReducer,
   metadata: metadataReducer,
@@ -60,3 +62,5 @@ export type AppDispatch = Store['dispatch'];
 
 export { setState as setStyleState, StyleState } from './style_slice';
 export { setState as setVisualizationState, VisualizationState } from './visualization_slice';
+export { MetadataState } from './metadata_slice';
+export { setState as setUIStateState, UIStateState } from './ui_state_slice';
diff --git a/src/plugins/vis_builder/public/application/utils/state_management/ui_state_slice.ts b/src/plugins/vis_builder/public/application/utils/state_management/ui_state_slice.ts
new file mode 100644
index 000000000000..826fe9d9873d
--- /dev/null
+++ b/src/plugins/vis_builder/public/application/utils/state_management/ui_state_slice.ts
@@ -0,0 +1,42 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { createSlice, PayloadAction } from '@reduxjs/toolkit';
+import { VisBuilderServices } from '../../../types';
+
+export type UIStateState<T = any> = T;
+
+const initialState = {} as UIStateState;
+
+export const getPreloadedState = async ({
+  types,
+  data,
+}: VisBuilderServices): Promise<UIStateState> => {
+  return initialState;
+};
+
+export const uiStateSlice = createSlice({
+  name: 'ui',
+  initialState,
+  reducers: {
+    setState<T>(state: T, action: PayloadAction<UIStateState<T>>) {
+      return action.payload;
+    },
+    updateState<T>(state: T, action: PayloadAction<Partial<UIStateState<T>>>) {
+      state = {
+        ...state,
+        ...action.payload,
+      };
+    },
+  },
+});
+
+// Exposing the state functions as generics
+export const setState = uiStateSlice.actions.setState as <T>(payload: T) => PayloadAction<T>;
+export const updateState = uiStateSlice.actions.updateState as <T>(
+  payload: Partial<T>
+) => PayloadAction<Partial<T>>;
+
+export const { reducer } = uiStateSlice;
diff --git a/src/plugins/vis_builder/public/application/utils/use/use_saved_vis_builder_vis.ts b/src/plugins/vis_builder/public/application/utils/use/use_saved_vis_builder_vis.ts
index 604c90a25ccd..44ffbaf75953 100644
--- a/src/plugins/vis_builder/public/application/utils/use/use_saved_vis_builder_vis.ts
+++ b/src/plugins/vis_builder/public/application/utils/use/use_saved_vis_builder_vis.ts
@@ -14,7 +14,12 @@ import {
 import { EDIT_PATH, PLUGIN_ID } from '../../../../common';
 import { VisBuilderServices } from '../../../types';
 import { getCreateBreadcrumbs, getEditBreadcrumbs } from '../breadcrumbs';
-import { useTypedDispatch, setStyleState, setVisualizationState } from '../state_management';
+import {
+  useTypedDispatch,
+  setStyleState,
+  setVisualizationState,
+  setUIStateState,
+} from '../state_management';
 import { useOpenSearchDashboards } from '../../../../../opensearch_dashboards_react/public';
 import { setEditorState } from '../state_management/metadata_slice';
 import { getStateFromSavedObject } from '../../../saved_visualizations/transforms';
@@ -46,6 +51,7 @@ export const useSavedVisBuilderVis = (visualizationIdFromUrl: string | undefined
 
     const loadSavedVisBuilderVis = async () => {
       try {
+        dispatch(setEditorState({ state: 'loading' }));
         const savedVisBuilderVis = await getSavedVisBuilderVis(
           savedVisBuilderLoader,
           visualizationIdFromUrl
@@ -56,8 +62,10 @@ export const useSavedVisBuilderVis = (visualizationIdFromUrl: string | undefined
           chrome.setBreadcrumbs(getEditBreadcrumbs(title, navigateToApp));
           chrome.docTitle.change(title);
 
+          dispatch(setUIStateState(state.ui));
           dispatch(setStyleState(state.style));
           dispatch(setVisualizationState(state.visualization));
+          dispatch(setEditorState({ state: 'loaded' }));
         } else {
           chrome.setBreadcrumbs(getCreateBreadcrumbs(navigateToApp));
         }
diff --git a/src/plugins/vis_builder/public/embeddable/vis_builder_embeddable.tsx b/src/plugins/vis_builder/public/embeddable/vis_builder_embeddable.tsx
index 7142c9c3c0cf..a931877ffe6d 100644
--- a/src/plugins/vis_builder/public/embeddable/vis_builder_embeddable.tsx
+++ b/src/plugins/vis_builder/public/embeddable/vis_builder_embeddable.tsx
@@ -22,6 +22,7 @@ import {
 } from '../../../expressions/public';
 import {
   Filter,
+  IIndexPattern,
   opensearchFilters,
   Query,
   TimefilterContract,
@@ -37,20 +38,23 @@ import {
 import { PersistedState } from '../../../visualizations/public';
 import { VisBuilderSavedVis } from '../saved_visualizations/transforms';
 import { handleVisEvent } from '../application/utils/handle_vis_event';
+import { VisBuilderEmbeddableFactoryDeps } from './vis_builder_embeddable_factory';
 
 // Apparently this needs to match the saved object type for the clone and replace panel actions to work
 export const VISBUILDER_EMBEDDABLE = VISBUILDER_SAVED_OBJECT;
 
 export interface VisBuilderEmbeddableConfiguration {
   savedVis: VisBuilderSavedVis;
-  // TODO: add indexPatterns as part of configuration
-  // indexPatterns?: IIndexPattern[];
+  indexPatterns?: IIndexPattern[];
   editPath: string;
   editUrl: string;
   editable: boolean;
+  deps: VisBuilderEmbeddableFactoryDeps;
 }
 
-export type VisBuilderInput = SavedObjectEmbeddableInput;
+export interface VisBuilderInput extends SavedObjectEmbeddableInput {
+  uiState?: any;
+}
 
 export interface VisBuilderOutput extends EmbeddableOutput {
   /**
@@ -58,11 +62,12 @@ export interface VisBuilderOutput extends EmbeddableOutput {
    * `input.savedObjectId`. If the id is invalid, this may be undefined.
    */
   savedVis?: VisBuilderSavedVis;
+  indexPatterns?: IIndexPattern[];
 }
 
 type ExpressionLoader = InstanceType<ExpressionsStart['ExpressionLoader']>;
 
-export class VisBuilderEmbeddable extends Embeddable<SavedObjectEmbeddableInput, VisBuilderOutput> {
+export class VisBuilderEmbeddable extends Embeddable<VisBuilderInput, VisBuilderOutput> {
   public readonly type = VISBUILDER_EMBEDDABLE;
   private handler?: ExpressionLoader;
   private timeRange?: TimeRange;
@@ -75,11 +80,19 @@ export class VisBuilderEmbeddable extends Embeddable<SavedObjectEmbeddableInput,
   private node?: HTMLElement;
   private savedVis?: VisBuilderSavedVis;
   private serializedState?: string;
-  private uiState?: PersistedState;
+  private uiState: PersistedState;
+  private readonly deps: VisBuilderEmbeddableFactoryDeps;
 
   constructor(
     timefilter: TimefilterContract,
-    { savedVis, editPath, editUrl, editable }: VisBuilderEmbeddableConfiguration,
+    {
+      savedVis,
+      editPath,
+      editUrl,
+      editable,
+      deps,
+      indexPatterns,
+    }: VisBuilderEmbeddableConfiguration,
     initialInput: SavedObjectEmbeddableInput,
     {
       parent,
@@ -96,12 +109,16 @@ export class VisBuilderEmbeddable extends Embeddable<SavedObjectEmbeddableInput,
         editUrl,
         editable,
         savedVis,
+        indexPatterns,
       },
       parent
     );
 
+    this.deps = deps;
     this.savedVis = savedVis;
-    this.uiState = new PersistedState();
+    this.uiState = new PersistedState(savedVis.state.ui);
+    this.uiState.on('change', this.uiStateChangeHandler);
+    this.uiState.on('reload', this.reload);
 
     this.autoRefreshFetchSubscription = timefilter
       .getAutoRefreshFetch$()
@@ -216,6 +233,8 @@ export class VisBuilderEmbeddable extends Embeddable<SavedObjectEmbeddableInput,
   public destroy() {
     super.destroy();
     this.subscriptions.forEach((s) => s.unsubscribe());
+    this.uiState.off('change', this.uiStateChangeHandler);
+    this.uiState.off('reload', this.reload);
     if (this.node) {
       ReactDOM.unmountComponentAtNode(this.node);
     }
@@ -249,6 +268,7 @@ export class VisBuilderEmbeddable extends Embeddable<SavedObjectEmbeddableInput,
 
   public async handleChanges() {
     // TODO: refactor (here and in visualize) to remove lodash dependency - immer probably a better choice
+    this.transferInputToUiState();
 
     let dirty = false;
 
@@ -303,6 +323,17 @@ export class VisBuilderEmbeddable extends Embeddable<SavedObjectEmbeddableInput,
     this.updateOutput({ loading: false, error });
   };
 
+  private uiStateChangeHandler = () => {
+    this.updateInput({
+      uiState: this.uiState.toJSON(),
+    });
+  };
+
+  private transferInputToUiState = () => {
+    if (JSON.stringify(this.input.uiState) !== this.uiState.toString())
+      this.uiState.set(this.input.uiState);
+  };
+
   // TODO: we may eventually need to add support for visualizations that use triggers like filter or brush, but current VisBuilder vis types don't support triggers
   // public supportedTriggers(): TriggerId[] {
   //   return this.visType.getSupportedTriggers?.() ?? [];
diff --git a/src/plugins/vis_builder/public/embeddable/vis_builder_embeddable_factory.tsx b/src/plugins/vis_builder/public/embeddable/vis_builder_embeddable_factory.tsx
index f80ad18ee363..90048ba91322 100644
--- a/src/plugins/vis_builder/public/embeddable/vis_builder_embeddable_factory.tsx
+++ b/src/plugins/vis_builder/public/embeddable/vis_builder_embeddable_factory.tsx
@@ -5,7 +5,6 @@
 
 import { i18n } from '@osd/i18n';
 import {
-  EmbeddableFactory,
   EmbeddableFactoryDefinition,
   EmbeddableOutput,
   ErrorEmbeddable,
@@ -35,16 +34,14 @@ import {
   getTimeFilter,
   getUISettings,
 } from '../plugin_services';
+import { StartServicesGetter } from '../../../opensearch_dashboards_utils/public';
+import { VisBuilderPluginStartDependencies } from '../types';
 
-// TODO: use or remove?
-export type VisBuilderEmbeddableFactory = EmbeddableFactory<
-  SavedObjectEmbeddableInput,
-  VisBuilderOutput | EmbeddableOutput,
-  VisBuilderEmbeddable | DisabledEmbeddable,
-  VisBuilderSavedObjectAttributes
->;
+export interface VisBuilderEmbeddableFactoryDeps {
+  start: StartServicesGetter<VisBuilderPluginStartDependencies>;
+}
 
-export class VisBuilderEmbeddableFactoryDefinition
+export class VisBuilderEmbeddableFactory
   implements
     EmbeddableFactoryDefinition<
       SavedObjectEmbeddableInput,
@@ -62,7 +59,7 @@ export class VisBuilderEmbeddableFactoryDefinition
   };
 
   // TODO: Would it be better to explicitly declare start service dependencies?
-  constructor() {}
+  constructor(private readonly deps: VisBuilderEmbeddableFactoryDeps) {}
 
   public canCreateNew() {
     // Because VisBuilder creation starts with the visualization modal, no need to have a separate entry for VisBuilder until it's separate
@@ -90,13 +87,22 @@ export class VisBuilderEmbeddableFactoryDefinition
         return new DisabledEmbeddable(PLUGIN_NAME, input);
       }
 
+      const savedVis = getStateFromSavedObject(savedObject);
+      const indexPatternService = this.deps.start().plugins.data.indexPatterns;
+      const indexPattern = await indexPatternService.get(
+        savedVis.state.visualization.indexPattern || ''
+      );
+      const indexPatterns = indexPattern ? [indexPattern] : [];
+
       return new VisBuilderEmbeddable(
         getTimeFilter(),
         {
-          savedVis: getStateFromSavedObject(savedObject),
+          savedVis,
           editUrl,
           editPath,
           editable: true,
+          deps: this.deps,
+          indexPatterns,
         },
         {
           ...input,
diff --git a/src/plugins/vis_builder/public/plugin.ts b/src/plugins/vis_builder/public/plugin.ts
index 0c1d569f6bed..7ed9e68ff7d7 100644
--- a/src/plugins/vis_builder/public/plugin.ts
+++ b/src/plugins/vis_builder/public/plugin.ts
@@ -23,7 +23,7 @@ import {
   VisBuilderSetup,
   VisBuilderStart,
 } from './types';
-import { VisBuilderEmbeddableFactoryDefinition, VISBUILDER_EMBEDDABLE } from './embeddable';
+import { VisBuilderEmbeddableFactory, VISBUILDER_EMBEDDABLE } from './embeddable';
 import visBuilderIconSecondaryFill from './assets/vis_builder_icon_secondary_fill.svg';
 import visBuilderIcon from './assets/vis_builder_icon.svg';
 import {
@@ -54,6 +54,7 @@ import { ConfigSchema } from '../config';
 import {
   createOsdUrlStateStorage,
   createOsdUrlTracker,
+  createStartServicesGetter,
   withNotifyOnErrors,
 } from '../../opensearch_dashboards_utils/public';
 import { opensearchFilters } from '../../data/public';
@@ -176,10 +177,8 @@ export class VisBuilderPlugin
     });
 
     // Register embeddable
-    // TODO: investigate simplification via getter a la visualizations:
-    // const start = createStartServicesGetter(core.getStartServices));
-    // const embeddableFactory = new VisBuilderEmbeddableFactoryDefinition({ start });
-    const embeddableFactory = new VisBuilderEmbeddableFactoryDefinition();
+    const start = createStartServicesGetter(core.getStartServices);
+    const embeddableFactory = new VisBuilderEmbeddableFactory({ start });
     embeddable.registerEmbeddableFactory(VISBUILDER_EMBEDDABLE, embeddableFactory);
 
     // Register the plugin as an alias to create visualization
diff --git a/src/plugins/vis_builder/public/saved_visualizations/_saved_vis.ts b/src/plugins/vis_builder/public/saved_visualizations/_saved_vis.ts
index 53ccc7a9d7d0..021af777df17 100644
--- a/src/plugins/vis_builder/public/saved_visualizations/_saved_vis.ts
+++ b/src/plugins/vis_builder/public/saved_visualizations/_saved_vis.ts
@@ -22,6 +22,7 @@ export function createSavedVisBuilderVisClass(services: SavedObjectOpenSearchDas
       description: 'text',
       visualizationState: 'text',
       styleState: 'text',
+      uiState: 'text',
       version: 'integer',
     };
 
@@ -44,7 +45,8 @@ export function createSavedVisBuilderVisClass(services: SavedObjectOpenSearchDas
           description: '',
           visualizationState: '{}',
           styleState: '{}',
-          version: 2,
+          uiState: '{}',
+          version: 3,
         },
       });
       this.showInRecentlyAccessed = true;
diff --git a/src/plugins/vis_builder/public/saved_visualizations/transforms.test.ts b/src/plugins/vis_builder/public/saved_visualizations/transforms.test.ts
index efbcfd23f799..68c24dfe4af1 100644
--- a/src/plugins/vis_builder/public/saved_visualizations/transforms.test.ts
+++ b/src/plugins/vis_builder/public/saved_visualizations/transforms.test.ts
@@ -25,7 +25,7 @@ describe('transforms', () => {
       savedObject = {} as VisBuilderSavedObject;
       rootState = {
         metadata: { editor: { state: 'loading', errors: {} } },
-        style: '',
+        style: {},
         visualization: {
           searchField: '',
           indexPattern: TEST_INDEX_PATTERN_ID,
@@ -34,6 +34,7 @@ describe('transforms', () => {
             aggConfigParams: [],
           },
         },
+        ui: {},
       };
       indexPattern = getStubIndexPattern(
         TEST_INDEX_PATTERN_ID,
@@ -49,6 +50,7 @@ describe('transforms', () => {
 
       expect(savedObject.visualizationState).not.toContain(TEST_INDEX_PATTERN_ID);
       expect(savedObject.styleState).toEqual(JSON.stringify(rootState.style));
+      expect(savedObject.uiState).toEqual(JSON.stringify(rootState.ui));
       expect(savedObject.searchSourceFields?.index?.id).toEqual(TEST_INDEX_PATTERN_ID);
     });
 
@@ -69,6 +71,7 @@ describe('transforms', () => {
       visualizationState: JSON.stringify({
         searchField: '',
       }),
+      uiState: '{}',
       searchSourceFields: {
         index: 'test-index',
       },
@@ -80,6 +83,7 @@ describe('transforms', () => {
       expect(state).toMatchInlineSnapshot(`
         Object {
           "style": Object {},
+          "ui": Object {},
           "visualization": Object {
             "indexPattern": "test-index",
             "searchField": "",
diff --git a/src/plugins/vis_builder/public/saved_visualizations/transforms.ts b/src/plugins/vis_builder/public/saved_visualizations/transforms.ts
index 0a7a6e529a6b..9f8dd705e3e4 100644
--- a/src/plugins/vis_builder/public/saved_visualizations/transforms.ts
+++ b/src/plugins/vis_builder/public/saved_visualizations/transforms.ts
@@ -27,6 +27,7 @@ export const saveStateToSavedObject = (
   );
   obj.styleState = JSON.stringify(state.style);
   obj.searchSourceFields = { index: indexPattern };
+  obj.uiState = JSON.stringify(state.ui);
 
   return obj;
 };
@@ -40,7 +41,8 @@ export const getStateFromSavedObject = (
   obj: VisBuilderSavedObjectAttributes
 ): VisBuilderSavedVis => {
   const { id, title, description } = obj;
-  const styleState = JSON.parse(obj.styleState || '');
+  const styleState = JSON.parse(obj.styleState || '{}');
+  const uiState = JSON.parse(obj.uiState || '{}');
   const vizStateWithoutIndex = JSON.parse(obj.visualizationState || '');
   const visualizationState: VisualizationState = {
     searchField: '',
@@ -48,7 +50,7 @@ export const getStateFromSavedObject = (
     indexPattern: obj.searchSourceFields?.index,
   };
 
-  const validateResult = validateVisBuilderState({ styleState, visualizationState });
+  const validateResult = validateVisBuilderState({ styleState, visualizationState, uiState });
 
   if (!validateResult.valid) {
     throw new InvalidJSONProperty(
@@ -75,6 +77,7 @@ export const getStateFromSavedObject = (
     state: {
       visualization: visualizationState,
       style: styleState,
+      ui: uiState,
     },
   };
 };
diff --git a/src/plugins/vis_builder/public/types.ts b/src/plugins/vis_builder/public/types.ts
index 5221a1c513ec..1ba8843e016a 100644
--- a/src/plugins/vis_builder/public/types.ts
+++ b/src/plugins/vis_builder/public/types.ts
@@ -62,6 +62,7 @@ export interface ISavedVis {
   description?: string;
   visualizationState?: string;
   styleState?: string;
+  uiState?: string;
   version?: number;
 }
 
diff --git a/src/plugins/vis_builder/server/saved_objects/vis_builder_app.ts b/src/plugins/vis_builder/server/saved_objects/vis_builder_app.ts
index 15d785b3b451..029557010bee 100644
--- a/src/plugins/vis_builder/server/saved_objects/vis_builder_app.ts
+++ b/src/plugins/vis_builder/server/saved_objects/vis_builder_app.ts
@@ -46,6 +46,10 @@ export const visBuilderSavedObjectType: SavedObjectsType = {
         type: 'text',
         index: false,
       },
+      uiState: {
+        type: 'text',
+        index: false,
+      },
       version: { type: 'integer' },
       // Need to add a kibanaSavedObjectMeta attribute here to follow the current saved object flow
       // When we save a saved object, the saved object plugin will extract the search source into two parts