feat(workspaces): Save and restore tabs COMPASS-9499

packages/atlas-service/src/atlas-service.ts

@@ -85,6 +85,7 @@ export class AtlasService {
       | 'favoriteQueries'
       | 'recentQueries'
       | 'favoriteAggregations'
+      | 'savedWorkspaces'
       | 'dataModelDescriptions',
     id?: string
   ): string {

packages/compass-preferences-model/src/feature-flags.ts

@@ -216,6 +216,16 @@ export const FEATURE_FLAG_DEFINITIONS = [
     },
   },
 
+  {
+    name: 'enableRestoreWorkspaces',
+    stage: 'development',
+    atlasCloudFeatureFlagName:
+      'DATA_EXPLORER_COMPASS_WEB_ENABLE_RESTORE_WORKSPACES',
+    description: {
+      short: 'Enable restoring previous workspace tabs on startup',
+    },
+  },
+
   {
     name: 'enableAutomaticRelationshipInference',
     stage: 'released',

packages/compass-user-data/src/index.ts

@@ -1,3 +1,3 @@
 export type { ReadAllResult } from './user-data';
-export { type IUserData, FileUserData, AtlasUserData } from './user-data';
+export { IUserData, FileUserData, AtlasUserData } from './user-data';
 export { z } from 'zod';

packages/compass-web/src/entrypoint.tsx

@@ -13,6 +13,7 @@ import type {
 } from '@mongodb-js/compass-workspaces';
 import WorkspacesPlugin, {
   WorkspacesProvider,
+  WorkspacesStorageServiceProviderWeb,
 } from '@mongodb-js/compass-workspaces';
 import {
   CollectionsWorkspaceTab,
@@ -125,7 +126,8 @@ const WithStorageProviders = createServiceProvider(
       const type = pathParts[0] as
         | 'favoriteQueries'
         | 'recentQueries'
-        | 'favoriteAggregations';
+        | 'favoriteAggregations'
+        | 'savedWorkspaces';
       const pathOrgId = pathParts[1];
       const pathProjectId = pathParts[2];
       const id = pathParts[3];
@@ -176,7 +178,14 @@ const WithStorageProviders = createServiceProvider(
       <PipelineStorageProvider value={pipelineStorage.current}>
         <FavoriteQueryStorageProvider value={favoriteQueryStorage.current}>
           <RecentQueryStorageProvider value={recentQueryStorage.current}>
-            {children}
+            <WorkspacesStorageServiceProviderWeb
+              orgId={orgId}
+              projectId={projectId}
+              getResourceUrl={getResourceUrl}
+              authenticatedFetch={authenticatedFetch}
+            >
+              {children}
+            </WorkspacesStorageServiceProviderWeb>
           </RecentQueryStorageProvider>
         </FavoriteQueryStorageProvider>
       </PipelineStorageProvider>

packages/compass-workspaces/src/index.ts

@@ -14,11 +14,13 @@ import workspacesReducer, {
   connectionDisconnected,
   updateDatabaseInfo,
   updateCollectionInfo,
+  loadSavedWorkspaces,
   beforeUnloading,
 } from './stores/workspaces';
 import Workspaces from './components';
 import { applyMiddleware, createStore } from 'redux';
 import thunk from 'redux-thunk';
+import { workspacesStateChangeMiddleware } from './stores/workspaces-middleware';
 import type { MongoDBInstance } from '@mongodb-js/compass-app-stores/provider';
 import { mongoDBInstancesManagerLocator } from '@mongodb-js/compass-app-stores/provider';
 import type Collection from 'mongodb-collection-model';
@@ -37,13 +39,17 @@ import {
 } from '@mongodb-js/compass-app-stores/provider';
 import type { PreferencesAccess } from 'compass-preferences-model/provider';
 import { preferencesLocator } from 'compass-preferences-model/provider';
+import { workspacesStorageServiceLocator } from './services/workspaces-storage';
+import { type WorkspacesStateSchema } from './types';
+import { type IUserData } from '@mongodb-js/compass-user-data';
 
 export type WorkspacesServices = {
   globalAppRegistry: AppRegistry;
   instancesManager: MongoDBInstancesManager;
   connections: ConnectionsService;
   logger: Logger;
   preferences: PreferencesAccess;
+  userData: IUserData<typeof WorkspacesStateSchema>;
 };
 
 export function configureStore(
@@ -65,7 +71,10 @@ export function configureStore(
       collectionInfo: {},
       databaseInfo: {},
     },
-    applyMiddleware(thunk.withExtraArgument(services))
+    applyMiddleware(
+      thunk.withExtraArgument(services),
+      workspacesStateChangeMiddleware(services)
+    )
   );
 
   return store;
@@ -85,6 +94,7 @@ export function activateWorkspacePlugin(
     connections,
     logger,
     preferences,
+    userData,
   }: WorkspacesServices,
   { on, cleanup, addCleanup }: ActivateHelpers
 ) {
@@ -94,8 +104,11 @@ export function activateWorkspacePlugin(
     connections,
     logger,
     preferences,
+    userData,
   });
 
+  void store.dispatch(loadSavedWorkspaces());
+
   addCleanup(cleanupLocalAppRegistries);
 
   const setupInstanceListeners = (
@@ -188,9 +201,11 @@ const WorkspacesPlugin = registerCompassPlugin(
     connections: connectionsLocator,
     logger: createLoggerLocator('COMPASS-WORKSPACES-UI'),
     preferences: preferencesLocator,
+    userData: workspacesStorageServiceLocator,
   }
 );
 
+export { WorkspacesStateSchema } from './types';
 export default WorkspacesPlugin;
 export { WorkspacesProvider } from './components/workspaces-provider';
 export type { OpenWorkspaceOptions, CollectionTabInfo };
@@ -210,3 +225,6 @@ export type {
   CollectionSubtab,
   WorkspacePluginProps,
 } from './types';
+
+export { WorkspacesStorageServiceProviderDesktop } from './services/workspaces-storage-desktop';
+export { WorkspacesStorageServiceProviderWeb } from './services/workspaces-storage-web';

packages/compass-workspaces/src/services/workspaces-storage-desktop.tsx

@@ -0,0 +1,23 @@
+import React, { useRef } from 'react';
+import { FileUserData, type IUserData } from '@mongodb-js/compass-user-data';
+import { WorkspacesStorageServiceContext } from './workspaces-storage';
+import { WorkspacesStateSchema } from '../types';
+import { EJSON } from 'bson';
+
+export const WorkspacesStorageServiceProviderDesktop: React.FunctionComponent =
+  ({ children }) => {
+    const storageRef = useRef<IUserData<typeof WorkspacesStateSchema>>(
+      new FileUserData(WorkspacesStateSchema, 'WorkspacesState', {
+        serialize: (content) =>
+          EJSON.stringify(content, {
+            relaxed: false,
+          }),
+        deserialize: (content: string) => EJSON.parse(content),
+      }) as IUserData<typeof WorkspacesStateSchema>
+    );
+    return (
+      <WorkspacesStorageServiceContext.Provider value={storageRef.current}>
+        {children}
+      </WorkspacesStorageServiceContext.Provider>
+    );
+  };

packages/compass-workspaces/src/services/workspaces-storage-web.tsx

@@ -0,0 +1,34 @@
+import React, { useRef } from 'react';
+import { AtlasUserData, type IUserData } from '@mongodb-js/compass-user-data';
+import { WorkspacesStorageServiceContext } from './workspaces-storage';
+import { WorkspacesStateSchema } from '../types';
+import { EJSON } from 'bson';
+
+export const WorkspacesStorageServiceProviderWeb: React.FunctionComponent<{
+  orgId: string;
+  projectId: string;
+  getResourceUrl: (path?: string) => string;
+  authenticatedFetch: (
+    url: RequestInfo | URL,
+    options?: RequestInit
+  ) => Promise<Response>;
+}> = ({ orgId, projectId, getResourceUrl, authenticatedFetch, children }) => {
+  const storageRef = useRef<IUserData<typeof WorkspacesStateSchema>>(
+    new AtlasUserData(WorkspacesStateSchema, 'WorkspacesState', {
+      orgId,
+      projectId,
+      getResourceUrl,
+      authenticatedFetch,
+      serialize: (content) =>
+        EJSON.stringify(content, {
+          relaxed: false,
+        }),
+      deserialize: (content: string) => EJSON.parse(content),
+    })
+  );
+  return (
+    <WorkspacesStorageServiceContext.Provider value={storageRef.current}>
+      {children}
+    </WorkspacesStorageServiceContext.Provider>
+  );
+};

packages/compass-workspaces/src/services/workspaces-storage.tsx

@@ -0,0 +1,50 @@
+import { createServiceLocator } from '@mongodb-js/compass-app-registry';
+import {
+  type z,
+  IUserData,
+  type ReadAllResult,
+} from '@mongodb-js/compass-user-data';
+import React, { useContext } from 'react';
+import { WorkspacesStateSchema } from '../types';
+
+const throwIfNotTestEnv = () => {
+  if (process.env.NODE_ENV !== 'test') {
+    throw new Error("Can't find Workspaces storage service in React context");
+  }
+};
+
+export class noopUserData<T extends z.Schema> extends IUserData<T> {
+  write(): Promise<boolean> {
+    throwIfNotTestEnv();
+    return Promise.resolve(true);
+  }
+  delete(): Promise<boolean> {
+    throwIfNotTestEnv();
+    return Promise.resolve(true);
+  }
+  readAll(): Promise<ReadAllResult<T>> {
+    throwIfNotTestEnv();
+    return Promise.resolve({ data: [], errors: [] });
+  }
+  readOne(): Promise<z.output<T>> {
+    throwIfNotTestEnv();
+    return Promise.resolve(undefined);
+  }
+  updateAttributes(): Promise<boolean> {
+    throwIfNotTestEnv();
+    return Promise.resolve(true);
+  }
+}
+
+export const noopWorkspacesStorageService: IUserData<
+  typeof WorkspacesStateSchema
+> = new noopUserData(WorkspacesStateSchema, 'WorkspacesState');
+
+export const WorkspacesStorageServiceContext = React.createContext<
+  IUserData<typeof WorkspacesStateSchema>
+>(noopWorkspacesStorageService);
+
+export const workspacesStorageServiceLocator = createServiceLocator(() => {
+  const service = useContext(WorkspacesStorageServiceContext);
+  return service;
+}, 'workspacesStorageServiceLocator');

packages/compass-workspaces/src/stores/workspaces-middleware.ts

@@ -0,0 +1,69 @@
+import type { Middleware, AnyAction } from 'redux';
+import type { WorkspacesState } from './workspaces';
+import type { WorkspacesStateData } from '../types';
+import type { WorkspacesServices } from '..';
+import { mongoLogId } from '@mongodb-js/compass-logging/provider';
+
+/**
+ * Debounced handler to save the workspaces state.
+ */
+const handleWorkspacesStateChange = (() => {
+  let timeoutId: ReturnType<typeof setTimeout>;
+  return (state: WorkspacesState, services: WorkspacesServices) => {
+    clearTimeout(timeoutId);
+    timeoutId = setTimeout(() => {
+      // Fire and forget - don't await to avoid blocking the action
+      void saveWorkspaceStateToUserData(state, services);
+    }, 250);
+  };
+})();
+
+/**
+ * Middleware that runs a callback whenever the workspaces state changes.
+ * This allows you to perform side effects when the state is updated.
+ */
+export function workspacesStateChangeMiddleware(
+  services: WorkspacesServices
+): Middleware<Record<string, never>, WorkspacesState> {
+  return (store) => (next) => (action: AnyAction) => {
+    const prevState = store.getState();
+    const result = next(action);
+    const nextState = store.getState();
+
+    // Only call the callback if the workspaces state actually changed
+    if (prevState !== nextState) {
+      if (services.preferences.getPreferences().enableRestoreWorkspaces) {
+        handleWorkspacesStateChange(nextState, services);
+      }
+    }
+
+    return result;
+  };
+}
+
+/**
+ * Saves the workspace state to persistent storage using UserData
+ */
+async function saveWorkspaceStateToUserData(
+  state: WorkspacesState,
+  services: WorkspacesServices
+) {
+  try {
+    // Transform the state to the format we want to save
+    const stateToSave: WorkspacesStateData = {
+      tabs: state.tabs,
+      activeTabId: state.activeTabId,
+      timestamp: Date.now(),
+    };
+
+    // Save to UserData with a fixed ID
+    await services.userData.write('saved-workspaces', stateToSave);
+  } catch (error) {
+    services.logger.log.error(
+      mongoLogId(1_001_000_229),
+      'Workspaces middleware',
+      'Failed to save workspace state to UserData',
+      { error }
+    );
+  }
+}