diff --git a/x-pack/plugins/lens/public/datasources/text_based/text_based_languages.test.ts b/x-pack/plugins/lens/public/datasources/text_based/text_based_languages.test.ts
index cef1bfa96b8a..25bab766558e 100644
--- a/x-pack/plugins/lens/public/datasources/text_based/text_based_languages.test.ts
+++ b/x-pack/plugins/lens/public/datasources/text_based/text_based_languages.test.ts
@@ -217,21 +217,24 @@ describe('IndexPattern Data Source', () => {
   describe('#removeLayer', () => {
     it('should remove a layer', () => {
       expect(TextBasedDatasource.removeLayer(baseState, 'a')).toEqual({
-        ...baseState,
-        layers: {
-          a: {
-            columns: [],
-            allColumns: [
-              {
-                columnId: 'col1',
-                fieldName: 'Test 1',
-                meta: {
-                  type: 'number',
+        removedLayerIds: ['a'],
+        newState: {
+          ...baseState,
+          layers: {
+            a: {
+              columns: [],
+              allColumns: [
+                {
+                  columnId: 'col1',
+                  fieldName: 'Test 1',
+                  meta: {
+                    type: 'number',
+                  },
                 },
-              },
-            ],
-            query: { sql: 'SELECT * FROM foo' },
-            index: 'foo',
+              ],
+              query: { sql: 'SELECT * FROM foo' },
+              index: 'foo',
+            },
           },
         },
       });
diff --git a/x-pack/plugins/lens/public/editor_frame_service/editor_frame/workspace_panel/chart_switch.test.tsx b/x-pack/plugins/lens/public/editor_frame_service/editor_frame/workspace_panel/chart_switch.test.tsx
index e6b762911d4d..276b56cc69ff 100644
--- a/x-pack/plugins/lens/public/editor_frame_service/editor_frame/workspace_panel/chart_switch.test.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_service/editor_frame/workspace_panel/chart_switch.test.tsx
@@ -494,8 +494,8 @@ describe('chart_switch', () => {
 
     switchTo('visB', instance);
     expect(datasourceMap.testDatasource.removeLayer).toHaveBeenCalledWith({}, 'a');
-    expect(datasourceMap.testDatasource.removeLayer).toHaveBeenCalledWith(undefined, 'b');
-    expect(datasourceMap.testDatasource.removeLayer).toHaveBeenCalledWith(undefined, 'c');
+    expect(datasourceMap.testDatasource.removeLayer).toHaveBeenCalledWith({}, 'b');
+    expect(datasourceMap.testDatasource.removeLayer).toHaveBeenCalledWith({}, 'c');
     expect(visualizationMap.visB.getSuggestions).toHaveBeenCalledWith(
       expect.objectContaining({
         keptLayerIds: ['a'],
diff --git a/x-pack/plugins/lens/public/mocks/datasource_mock.ts b/x-pack/plugins/lens/public/mocks/datasource_mock.ts
index cfb93882559e..986a753a1a4a 100644
--- a/x-pack/plugins/lens/public/mocks/datasource_mock.ts
+++ b/x-pack/plugins/lens/public/mocks/datasource_mock.ts
@@ -26,7 +26,7 @@ export function createMockDatasource(id: string): DatasourceMock {
 
   return {
     id: 'testDatasource',
-    clearLayer: jest.fn((state, _layerId) => state),
+    clearLayer: jest.fn((state, _layerId) => ({ newState: state, removedLayerIds: [] })),
     getDatasourceSuggestionsForField: jest.fn((_state, _item, filterFn, _indexPatterns) => []),
     getDatasourceSuggestionsForVisualizeField: jest.fn(
       (_state, _indexpatternId, _fieldName, _indexPatterns) => []
@@ -44,7 +44,7 @@ export function createMockDatasource(id: string): DatasourceMock {
     renderLayerPanel: jest.fn(),
     toExpression: jest.fn((_frame, _state, _indexPatterns) => null),
     insertLayer: jest.fn((_state, _newLayerId) => ({})),
-    removeLayer: jest.fn((_state, _layerId) => {}),
+    removeLayer: jest.fn((state, layerId) => ({ newState: state, removedLayerIds: [layerId] })),
     cloneLayer: jest.fn((_state, _layerId, _newLayerId, getNewId) => {}),
     removeColumn: jest.fn((props) => {}),
     getLayers: jest.fn((_state) => []),
diff --git a/x-pack/plugins/lens/public/state_management/lens_slice.ts b/x-pack/plugins/lens/public/state_management/lens_slice.ts
index 105cc31d54ee..3f3490906f88 100644
--- a/x-pack/plugins/lens/public/state_management/lens_slice.ts
+++ b/x-pack/plugins/lens/public/state_management/lens_slice.ts
@@ -996,9 +996,12 @@ export const makeLensReducer = (storeDeps: LensStoreDeps) => {
             );
           }) ?? [];
         if (layerDatasourceId) {
-          state.datasourceStates[layerDatasourceId].state = datasourceMap[
-            layerDatasourceId
-          ].removeLayer(current(state).datasourceStates[layerDatasourceId].state, layerId);
+          const { newState } = datasourceMap[layerDatasourceId].removeLayer(
+            current(state).datasourceStates[layerDatasourceId].state,
+            layerId
+          );
+          state.datasourceStates[layerDatasourceId].state = newState;
+          // TODO - call removeLayer for any extra (linked) layers removed by the datasource
         }
       });
     },