1 files changed, 18 insertions, 8 deletions
diff --git a/app/assets/javascripts/batch_comments/stores/modules/batch_comments/actions.js b/app/assets/javascripts/batch_comments/stores/modules/batch_comments/actions.js
index a29409c52ae..bd7909bfa76 100644
--- a/app/assets/javascripts/batch_comments/stores/modules/batch_comments/actions.js
+++ b/app/assets/javascripts/batch_comments/stores/modules/batch_comments/actions.js
@@ -2,8 +2,8 @@ import { deprecatedCreateFlash as flash } from '~/flash';
 import { __ } from '~/locale';
 import { scrollToElement } from '~/lib/utils/common_utils';
 import service from '../../../services/drafts_service';
-import * as types from './mutation_types';
 import { CHANGES_TAB, DISCUSSION_TAB, SHOW_TAB } from '../../../constants';
+import * as types from './mutation_types';
 
 export const saveDraft = ({ dispatch }, draft) =>
   dispatch('saveNote', { ...draft, isDraft: true }, { root: true });
@@ -67,13 +67,23 @@ export const publishReview = ({ commit, dispatch, getters }) => {
     .catch(() => commit(types.RECEIVE_PUBLISH_REVIEW_ERROR));
 };
 
-export const updateDiscussionsAfterPublish = ({ dispatch, getters, rootGetters }) =>
-  dispatch('fetchDiscussions', { path: getters.getNotesData.discussionsPath }, { root: true }).then(
-    () =>
-      dispatch('diffs/assignDiscussionsToDiff', rootGetters.discussionsStructuredByLineCode, {
-        root: true,
-      }),
-  );
+export const updateDiscussionsAfterPublish = async ({ dispatch, getters, rootGetters }) => {
+  if (window.gon?.features?.paginatedNotes) {
+    await dispatch('stopPolling', null, { root: true });
+    await dispatch('fetchData', null, { root: true });
+    await dispatch('restartPolling', null, { root: true });
+  } else {
+    await dispatch(
+      'fetchDiscussions',
+      { path: getters.getNotesData.discussionsPath },
+      { root: true },
+    );
+  }
+
+  dispatch('diffs/assignDiscussionsToDiff', rootGetters.discussionsStructuredByLineCode, {
+    root: true,
+  });
+};
 
 export const updateDraft = (
   { commit, getters },