diff --git a/src/legacy/core_plugins/visualizations/public/embeddable/visualize_embeddable.ts b/src/legacy/core_plugins/visualizations/public/embeddable/visualize_embeddable.ts
index 72a0ef72b569..32bbae13b79b 100644
--- a/src/legacy/core_plugins/visualizations/public/embeddable/visualize_embeddable.ts
+++ b/src/legacy/core_plugins/visualizations/public/embeddable/visualize_embeddable.ts
@@ -369,9 +369,7 @@ export class VisualizeEmbeddable extends Embeddable<VisualizeInput, VisualizeOut
         query: this.input.query,
         filters: this.input.filters,
       },
-      extraHandlers: {
-        uiState: this.uiState,
-      },
+      uiState: this.uiState,
     };
     this.expression = await buildPipeline(this.vis, {
       searchSource: this.savedVisualization.searchSource,
diff --git a/src/plugins/expressions/public/loader.ts b/src/plugins/expressions/public/loader.ts
index 4600922e076f..fbe2f37c648d 100644
--- a/src/plugins/expressions/public/loader.ts
+++ b/src/plugins/expressions/public/loader.ts
@@ -167,7 +167,7 @@ export class ExpressionLoader {
   };
 
   private render(data: Data): void {
-    this.renderHandler.render(data, this.params.extraHandlers);
+    this.renderHandler.render(data, this.params.uiState);
   }
 
   private setParams(params?: IExpressionLoaderParams) {
@@ -182,8 +182,8 @@ export class ExpressionLoader {
         this.params.searchContext || {}
       ) as any;
     }
-    if (params.extraHandlers && this.params) {
-      this.params.extraHandlers = params.extraHandlers;
+    if (params.uiState && this.params) {
+      this.params.uiState = params.uiState;
     }
     if (params.variables && this.params) {
       this.params.variables = params.variables;
diff --git a/src/plugins/expressions/public/render.ts b/src/plugins/expressions/public/render.ts
index 86e360f8135e..ad4d16bcd132 100644
--- a/src/plugins/expressions/public/render.ts
+++ b/src/plugins/expressions/public/render.ts
@@ -94,7 +94,7 @@ export class ExpressionRenderHandler {
     };
   }
 
-  render = async (data: any, extraHandlers: IExpressionRendererExtraHandlers = {}) => {
+  render = async (data: any, uiState: any = {}) => {
     if (!data || typeof data !== 'object') {
       return this.handleRenderError(new Error('invalid data provided to the expression renderer'));
     }
@@ -119,7 +119,7 @@ export class ExpressionRenderHandler {
         .get(data.as)!
         .render(this.element, data.value, {
           ...this.handlers,
-          ...extraHandlers,
+          uiState,
         } as any);
     } catch (e) {
       return this.handleRenderError(e);
diff --git a/src/plugins/expressions/public/types/index.ts b/src/plugins/expressions/public/types/index.ts
index c77698d3661c..b5781ef213fd 100644
--- a/src/plugins/expressions/public/types/index.ts
+++ b/src/plugins/expressions/public/types/index.ts
@@ -48,7 +48,7 @@ export interface IExpressionLoaderParams {
   disableCaching?: boolean;
   customFunctions?: [];
   customRenderers?: [];
-  extraHandlers?: Record<string, any>;
+  uiState?: unknown;
   inspectorAdapters?: Adapters;
   onRenderError?: RenderErrorHandlerFnType;
 }