feat(workspaces): change tab behavior to prevent modified tabs from being closed by accident COMPASS-5022 (#5952)

* feat(workspaces): add workspace helper methods to allow controlling tab closing flow; apply initial logic for preventing tabs from closing by accident

* chore: refactor to canClose / canReplace sync handler response

* fix(workspaces): always open new tab when explicitly required

* chore: separate single handler to two distinct ones for close and replace

* chore: lint fixes

* chore(workspaces): clean-up store and add new unit tests

* chore: add e2e tests for tab behavior

* chore: don't click the query bar to not trigger AI feature by accident

* chore: small text fixes

Author: gribnoysup

packages/compass-aggregations/src/stores/store.ts

@@ -262,6 +262,14 @@ export function activateAggregationsPlugin(
     maxTimeMSChanged(preferences.getPreferences().maxTimeMS || null)
   );
 
+  const onCloseOrReplace = () => {
+    return !store.getState().isModified;
+  };
+
+  addCleanup(workspaces.onTabReplace?.(onCloseOrReplace));
+
+  addCleanup(workspaces.onTabClose?.(onCloseOrReplace));
+
   return {
     store,
     deactivate: cleanup,

packages/compass-e2e-tests/helpers/commands/close-workspace-tabs.ts

@@ -2,7 +2,8 @@ import type { CompassBrowser } from '../compass-browser';
 import * as Selectors from '../selectors';
 
 export async function closeWorkspaceTabs(
-  browser: CompassBrowser
+  browser: CompassBrowser,
+  autoConfirmTabClose = true
 ): Promise<void> {
   const countTabs = async () => {
     return (await browser.$$(Selectors.workspaceTab(null))).length;
@@ -15,6 +16,14 @@ export async function closeWorkspaceTabs(
     await currentActiveTab.click();
     await browser.waitUntil(async () => {
       await currentActiveTab.$(Selectors.CloseWorkspaceTab).click();
+      if (autoConfirmTabClose) {
+        // Tabs in "dirty" state can't be closed without confirmation
+        if (await browser.$(Selectors.ConfirmTabCloseModal).isExisting()) {
+          await browser.clickVisible(
+            browser.$(Selectors.ConfirmTabCloseModal).$('button=Close tab')
+          );
+        }
+      }
       return (await currentActiveTab.isExisting()) === false;
     });
   }

packages/compass-e2e-tests/helpers/commands/collection-workspaces.ts

@@ -5,17 +5,21 @@ async function navigateToCollection(
   browser: CompassBrowser,
   // TODO(COMPASS-8002): take connectionName into account
   dbName: string,
-  collectionName: string
+  collectionName: string,
+
+  // Close all the workspace tabs to get rid of all the state we
+  // might have accumulated. This is the only way to get back to the zero
+  // state of Schema, and Validation tabs without re-connecting.
+  closeExistingTabs = true
 ): Promise<void> {
   const collectionSelector = Selectors.sidebarCollection(
     dbName,
     collectionName
   );
 
-  // Close all the workspace tabs to get rid of all the state we
-  // might have accumulated. This is the only way to get back to the zero
-  // state of Schema, and Validation tabs without re-connecting.
-  await browser.closeWorkspaceTabs();
+  if (closeExistingTabs) {
+    await browser.closeWorkspaceTabs();
+  }
 
   // search for the collection and wait for the collection to be there and visible
   await browser.clickVisible(Selectors.SidebarFilterInput);
@@ -39,9 +43,15 @@ export async function navigateToCollectionTab(
     | 'Aggregations'
     | 'Schema'
     | 'Indexes'
-    | 'Validation' = 'Documents'
+    | 'Validation' = 'Documents',
+  closeExistingTabs = true
 ): Promise<void> {
-  await navigateToCollection(browser, dbName, collectionName);
+  await navigateToCollection(
+    browser,
+    dbName,
+    collectionName,
+    closeExistingTabs
+  );
   await navigateWithinCurrentCollectionTabs(browser, tabName);
 }
 

packages/compass-e2e-tests/helpers/selectors.ts

@@ -1160,6 +1160,8 @@ export const sidebarInstanceNavigationItem = (
   return `${Sidebar} [aria-label="${tabName}"]`;
 };
 export const SidebarMyQueriesTab = `${Sidebar} [aria-label="My Queries"]`;
+export const WorkspaceTab =
+  '[role="tablist"][aria-label="Workspace Tabs"] [role="tab"]';
 export const workspaceTab = (
   title: string | null,
   active: boolean | null = null
@@ -1171,7 +1173,7 @@ export const workspaceTab = (
       : ['My Queries', 'Performance', 'Databases'].includes(title)
       ? `[title="${title}"]`
       : `[data-namespace="${title}"]`;
-  return `[role="tablist"][aria-label="Workspace Tabs"] [role="tab"]${_title}${_active}`;
+  return `${WorkspaceTab}${_title}${_active}`;
 };
 export const instanceWorkspaceTab = (
   tabName: 'Performance' | 'Databases',
@@ -1303,3 +1305,6 @@ export const AgreeAndContinueButton = 'button=Agree and continue';
 
 // Any toast
 export const AnyToastDismissButton = '[data-testid="lg-toast-dismiss-button"]';
+
+// Close tab confirmation
+export const ConfirmTabCloseModal = '[data-testid="confirm-tab-close"]';

packages/compass-e2e-tests/tests/tabs.test.ts

@@ -0,0 +1,119 @@
+import type { CompassBrowser } from '../helpers/compass-browser';
+import { init, cleanup, screenshotIfFailed } from '../helpers/compass';
+import type { Compass } from '../helpers/compass';
+import * as Selectors from '../helpers/selectors';
+import { createNumbersCollection } from '../helpers/insert-data';
+import { expect } from 'chai';
+
+describe('Global Tabs', function () {
+  let compass: Compass;
+  let browser: CompassBrowser;
+
+  const collections = ['a', 'b', 'c'];
+
+  before(async function () {
+    compass = await init(this.test?.fullTitle());
+    browser = compass.browser;
+  });
+
+  beforeEach(async function () {
+    for (const collName of collections) {
+      await createNumbersCollection(collName, 1);
+      await createNumbersCollection(collName, 1);
+      await createNumbersCollection(collName, 1);
+    }
+    await browser.connectWithConnectionString();
+  });
+
+  afterEach(async function () {
+    await screenshotIfFailed(compass, this.currentTest);
+  });
+
+  after(async function () {
+    await cleanup(compass);
+  });
+
+  it('should open tabs over each other when not modified', async function () {
+    for (const collName of collections) {
+      await browser.navigateToCollectionTab(
+        'test',
+        collName,
+        'Documents',
+        false
+      );
+    }
+    expect(await browser.$$(Selectors.workspaceTab(null))).to.have.lengthOf(1);
+  });
+
+  it('should open tabs over each other when not modified', async function () {
+    for (const collName of collections) {
+      await browser.navigateToCollectionTab(
+        'test',
+        collName,
+        'Documents',
+        false
+      );
+      // Click something to make sure we "modified" the tab
+      await browser.clickVisible(
+        Selectors.queryBarApplyFilterButton('Documents')
+      );
+    }
+    expect(await browser.$$(Selectors.workspaceTab(null))).to.have.lengthOf(3);
+  });
+
+  it('should close tabs without warning even when "modified" by interacting with the tab', async function () {
+    for (const collName of collections) {
+      await browser.navigateToCollectionTab(
+        'test',
+        collName,
+        'Documents',
+        false
+      );
+      // Click something to make sure we "modified" the tab
+      await browser.clickVisible(
+        Selectors.queryBarApplyFilterButton('Documents')
+      );
+    }
+    await browser.closeWorkspaceTabs(false);
+    expect(await browser.$$(Selectors.workspaceTab(null))).to.have.lengthOf(0);
+  });
+
+  it('should ask for confirmation when closing modified Aggregations tab', async function () {
+    await browser.navigateToCollectionTab('test', 'a', 'Aggregations');
+
+    await browser.clickVisible(
+      Selectors.aggregationPipelineModeToggle('as-text')
+    );
+
+    await browser.setCodemirrorEditorValue(
+      Selectors.AggregationAsTextEditor,
+      '[{$match: { i: 0 }}]'
+    );
+
+    await browser.clickVisible(Selectors.CloseWorkspaceTab);
+    await browser.$(Selectors.ConfirmTabCloseModal).waitForDisplayed();
+
+    await browser.clickVisible(
+      browser.$(Selectors.ConfirmTabCloseModal).$('button=Cancel')
+    );
+    await browser
+      .$(Selectors.ConfirmTabCloseModal)
+      .waitForExist({ reverse: true });
+
+    // Checking first that cancel leaves the tab on the screen
+    expect(await browser.$$(Selectors.workspaceTab(null))).to.have.lengthOf(1);
+
+    await browser.clickVisible(Selectors.CloseWorkspaceTab);
+    await browser.$(Selectors.ConfirmTabCloseModal).waitForDisplayed();
+
+    await browser.clickVisible(
+      browser.$(Selectors.ConfirmTabCloseModal).$('button=Close tab')
+    );
+    await browser
+      .$(Selectors.ConfirmTabCloseModal)
+      .waitForExist({ reverse: true });
+
+    // When confirmed, should remove the tab
+    expect(await browser.$$(Selectors.workspaceTab(null))).to.have.lengthOf(0);
+  });
+});

packages/compass-workspaces/src/components/workspace-close-handler.tsx

@@ -0,0 +1,156 @@
+import { useEffect, useRef } from 'react';
+import type { WorkspaceTab } from '../stores/workspaces';
+import { useWorkspaceTabId } from './workspace-tab-state-provider';
+
+export type WorkspaceDestroyHandler = () => boolean;
+
+const HANDLERS = {
+  close: new Map<string, WorkspaceDestroyHandler[]>(),
+  replace: new Map<string, WorkspaceDestroyHandler[]>(),
+} as const;
+
+const resolveTabDestroyState = (
+  type: 'close' | 'replace',
+  tab: WorkspaceTab
+): boolean => {
+  const handlers = HANDLERS[type].get(tab.id) ?? [];
+
+  for (const handler of handlers) {
+    if (handler() === false) {
+      return false;
+    }
+  }
+
+  return true;
+};
+
+export const canCloseTab = (tab: WorkspaceTab) => {
+  return resolveTabDestroyState('close', tab);
+};
+
+export const canReplaceTab = (tab: WorkspaceTab) => {
+  return resolveTabDestroyState('replace', tab);
+};
+
+/**
+ * Exported only for testing purposes
+ * @internal
+ */
+export const setTabDestroyHandler = (
+  type: 'close' | 'replace',
+  tabId: string,
+  handler: WorkspaceDestroyHandler
+) => {
+  const handlerMap = HANDLERS[type];
+  const handlers = handlerMap.get(tabId) ?? [];
+  handlerMap.set(tabId, handlers.concat(handler));
+  return () => {
+    cleanupTabDestroyHandler(type, tabId, handler);
+  };
+};
+
+export const cleanupTabDestroyHandler = (
+  type?: 'close' | 'replace',
+  tabId?: string,
+  handler?: WorkspaceDestroyHandler
+) => {
+  if (handler && tabId && type) {
+    HANDLERS[type].set(
+      tabId,
+      (HANDLERS[type].get(tabId) ?? []).filter((fn) => {
+        return fn !== handler;
+      })
+    );
+  } else if (tabId && type) {
+    HANDLERS[type].delete(tabId);
+  } else if (type) {
+    HANDLERS[type].clear();
+  } else {
+    Object.values(HANDLERS).forEach((map) => {
+      map.clear();
+    });
+  }
+};
+
+function useOnTabDestroyHandler(
+  type: 'close' | 'replace',
+  handler: WorkspaceDestroyHandler
+) {
+  const tabId = useWorkspaceTabId();
+  const handlerRef = useRef(handler);
+  handlerRef.current = handler;
+  useEffect(() => {
+    const onClose: WorkspaceDestroyHandler = () => {
+      return handlerRef.current();
+    };
+    return setTabDestroyHandler(type, tabId, onClose);
+  }, [type, tabId]);
+}
+
+/**
+ * A hook that registers a tab close handler. Before closing the tab, this
+ * handler will be called and can return either `true` to allow the tab to be
+ * closed, or `false`, preventing tab to be closed before user confirms the tab
+ * closure.
+ *
+ * Multiple handlers can be registered for one tab, they will be called in order
+ * of registration and any one of them returning `false` will prevent the tab
+ * from closing.
+ *
+ * @example
+ * function TabWithInput() {
+ *   const [value, setValue] = useState('');
+ *   // Will prevent tab from closing if text input is not empty. This will
+ *   // cause a confirmation modal to appear before closing the tab
+ *   useOnTabClose(() => {
+ *     return value !== '';
+ *   });
+ *   return <TextInput value={value} onChange={setValue} />
+ * }
+ */
+export const useOnTabClose = useOnTabDestroyHandler.bind(null, 'close');
+
+/**
+ *
+ * A hook that registers a tab replace handler. By default when opening a new
+ * workspace, it will be opened in the same tab, destroying the content of the
+ * current workspace. In that case, the registered handler can return either
+ * `true` to allow the workspace to be destroyed, or `false` to prevent the tab
+ * from being destroyed and forcing the new workspace to open in the new tab.
+ *
+ * Multiple handlers can be registered for one tab, they will be called in order
+ * of registration and any one of them returning `false` will prevent the tab
+ * from being replaced.
+ *
+ * @example
+ * function TabWithInput() {
+ *   const [value, setValue] = useState('');
+ *   // Will prevent tab from being replaced if text input is not empty. This
+ *   // will cause new workspace to open in a new tab
+ *   useOnTabReplace(() => {
+ *     return value !== '';
+ *   });
+ *   return <TextInput value={value} onChange={setValue} />
+ * }
+ */
+export const useOnTabReplace = useOnTabDestroyHandler.bind(null, 'replace');
+
+export function useRegisterTabDestroyHandler() {
+  let tabId: string | undefined;
+  try {
+    // We are not breaking rules of hooks here, this method will either always
+    // fail or always return a value when rendered in a certain spot in the
+    // application
+    // eslint-disable-next-line react-hooks/rules-of-hooks
+    tabId = useWorkspaceTabId();
+  } catch {
+    // do nothing, we are just outside of the workspace tab tree
+  }
+
+  if (tabId) {
+    return {
+      onTabClose: setTabDestroyHandler.bind(null, 'close', tabId),
+      onTabReplace: setTabDestroyHandler.bind(null, 'replace', tabId),
+    };
+  }
+}

packages/compass-workspaces/src/components/workspace-tab-state-provider.tsx

@@ -98,7 +98,7 @@ export const WorkspaceTabStateProvider = ({
   );
 };
 
-function useWorkspaceTabId() {
+export function useWorkspaceTabId() {
   let tabId = useContext(WorkspaceTabIdContext);
   if (!tabId) {
     if (process.env.NODE_ENV !== 'test') {