fix(explore): preserve shape size across projection reloads

apps/web/src/explore/dataset-controller.eat.test.ts

@@ -5,7 +5,9 @@ import { createEmptyExploreViewRequest } from './url-state';
 const mocks = vi.hoisted(() => ({
   loadData: vi.fn(),
   markLastLoadStatus: vi.fn(),
+  readTooltipAnnotations: vi.fn(() => [] as string[]),
   resolvePendingLoadFinalization: vi.fn(),
+  writeTooltipAnnotations: vi.fn(),
 }));
 
 vi.mock('./data-renderer', () => ({
@@ -28,8 +30,8 @@ vi.mock('./opfs-dataset-store', () => ({
 }));
 
 vi.mock('./tooltip-annotations-store', () => ({
-  readTooltipAnnotations: () => [],
-  writeTooltipAnnotations: vi.fn(),
+  readTooltipAnnotations: mocks.readTooltipAnnotations,
+  writeTooltipAnnotations: mocks.writeTooltipAnnotations,
 }));
 
 import { createDatasetController } from './dataset-controller';
@@ -49,90 +51,205 @@ const data: VisualizationData = {
   annotation_data: { ec: new Int32Array([0]) },
 };
 
-describe('dataset controller EAT settings restore', () => {
+type DatasetControllerOptions = Parameters<typeof createDatasetController>[0];
+
+function createControllerHarness({
+  getLatestViewRequest = () => createEmptyExploreViewRequest(),
+  loadQueue: loadQueueOverrides = {},
+}: {
+  getLatestViewRequest?: () => ReturnType<typeof createEmptyExploreViewRequest>;
+  loadQueue?: Partial<DatasetControllerOptions['loadQueue']>;
+} = {}) {
+  const controlBar = {
+    clearForNewDataset: vi.fn(),
+    hasFileSettings: false,
+  };
+  const legendElement = {
+    clearForNewDataset: vi.fn(),
+    setFileSettings: vi.fn(),
+    applyEatSettings: vi.fn(),
+  };
+  const plotElement = {
+    eatOverlayEnabled: true,
+  };
+  const viewController = {
+    subscribeToViewChanges: vi.fn(() => () => {}),
+    resolveLatestView: vi.fn(),
+    getLatestViewRequest: vi.fn(getLatestViewRequest),
+    applyLatestViewForDatasetLoad: vi.fn(),
+    setRequestedView: vi.fn(),
+  };
+  const loadQueue = {
+    registerFileLoad: vi.fn(),
+    getLoadMetaForFile: vi.fn(),
+    getRunningLoadMeta: () => null,
+    getLatestSequence: () => 0,
+    resolvePendingLoadFinalization: mocks.resolvePendingLoadFinalization,
+    ...loadQueueOverrides,
+  };
+  const options = {
+    controlBar,
+    dataLoader: {},
+    defaultDatasetName: 'default.parquetbundle',
+    getIsDisposed: () => false,
+    interactionController: {},
+    legendElement,
+    loadQueue,
+    overlayController: { update: vi.fn() },
+    plotElement,
+    setCurrentDatasetIsDemo: vi.fn(),
+    setCurrentDatasetName: vi.fn(),
+    structureViewer: {},
+    viewController,
+  } as unknown as DatasetControllerOptions;
+
+  return {
+    controlBar,
+    controller: createDatasetController(options),
+    legendElement,
+    plotElement,
+    viewController,
+  };
+}
+
+describe('dataset controller', () => {
   beforeEach(() => {
     vi.clearAllMocks();
     mocks.loadData.mockResolvedValue(undefined);
     mocks.markLastLoadStatus.mockResolvedValue(undefined);
+    mocks.readTooltipAnnotations.mockReturnValue([]);
   });
 
-  it('applies embedded EAT settings after an OPFS reload while retaining OPFS legend precedence', async () => {
-    const controlBar = {
-      clearForNewDataset: vi.fn(),
-      hasFileSettings: false,
-    };
-    const legendElement = {
-      clearForNewDataset: vi.fn(),
-      setFileSettings: vi.fn(),
-      applyEatSettings: vi.fn(),
-    };
-    const plotElement = {
-      eatOverlayEnabled: true,
-    };
-    const viewController = {
-      subscribeToViewChanges: vi.fn(() => () => {}),
-      resolveLatestView: vi.fn(),
-      getLatestViewRequest: vi.fn(() => createEmptyExploreViewRequest()),
-      applyLatestViewForDatasetLoad: vi.fn(),
-      setRequestedView: vi.fn(),
-    };
-    const options = {
-      controlBar,
-      dataLoader: {},
-      defaultDatasetName: 'default.parquetbundle',
-      getIsDisposed: () => false,
-      interactionController: {},
-      legendElement,
-      loadQueue: {
-        registerFileLoad: vi.fn(),
-        getLoadMetaForFile: vi.fn(),
-        getRunningLoadMeta: () => ({ sequence: 7, kind: 'opfs' as const }),
-        getLatestSequence: () => 7,
-        resolvePendingLoadFinalization: mocks.resolvePendingLoadFinalization,
-      },
-      overlayController: { update: vi.fn() },
-      plotElement,
-      setCurrentDatasetIsDemo: vi.fn(),
-      setCurrentDatasetName: vi.fn(),
-      structureViewer: {},
-      viewController,
-    } as unknown as Parameters<typeof createDatasetController>[0];
-    const controller = createDatasetController(options);
-
-    await controller.handleDataLoaded({
-      detail: {
-        data,
-        settings: {
-          legendSettings: { ec: { categories: {} } },
-          exportOptions: {},
-          eatOverlayEnabled: false,
-          eatConfidenceThreshold: 0.75,
+  describe('EAT settings restore', () => {
+    it('applies embedded EAT settings after an OPFS reload while retaining OPFS legend precedence', async () => {
+      const { controlBar, controller, legendElement } = createControllerHarness({
+        loadQueue: {
+          getRunningLoadMeta: () => ({ sequence: 7, kind: 'opfs' as const }),
+          getLatestSequence: () => 7,
+        },
+      });
+
+      await controller.handleDataLoaded({
+        detail: {
+          data,
+          settings: {
+            legendSettings: { ec: { categories: {} } },
+            exportOptions: {},
+            eatOverlayEnabled: false,
+            eatConfidenceThreshold: 0.75,
+          },
+          source: 'auto',
         },
-        source: 'auto',
-      },
-    } as unknown as Event);
-
-    expect(controlBar.clearForNewDataset).toHaveBeenCalledOnce();
-    expect(legendElement.applyEatSettings).toHaveBeenCalledWith(false, 0.75);
-    expect(controlBar.hasFileSettings).toBe(true);
-    expect(legendElement.setFileSettings).not.toHaveBeenCalled();
-    expect(mocks.markLastLoadStatus).toHaveBeenCalledWith('success');
-    expect(mocks.resolvePendingLoadFinalization).toHaveBeenCalledWith(7);
-
-    await controller.handleDataLoaded({
-      detail: {
+      } as unknown as Event);
+
+      expect(controlBar.clearForNewDataset).toHaveBeenCalledOnce();
+      expect(legendElement.applyEatSettings).toHaveBeenCalledWith(false, 0.75);
+      expect(controlBar.hasFileSettings).toBe(true);
+      expect(legendElement.setFileSettings).not.toHaveBeenCalled();
+      expect(mocks.markLastLoadStatus).toHaveBeenCalledWith('success');
+      expect(mocks.resolvePendingLoadFinalization).toHaveBeenCalledWith(7);
+
+      await controller.handleDataLoaded({
+        detail: {
+          data,
+          settings: {
+            legendSettings: {},
+            exportOptions: {},
+            eatOverlayEnabled: true,
+          },
+          source: 'auto',
+        },
+      } as unknown as Event);
+      expect(legendElement.applyEatSettings).toHaveBeenLastCalledWith(
+        true,
+        DEFAULT_EAT_CONFIDENCE_THRESHOLD,
+      );
+    });
+  });
+
+  describe('legend persistence lifecycle', () => {
+    it('reuses dataset lifecycle across initial load, user import, and explicit reset', async () => {
+      let latestViewRequest = createEmptyExploreViewRequest();
+      const { controlBar, controller, legendElement, viewController } = createControllerHarness({
+        getLatestViewRequest: () => latestViewRequest,
+        loadQueue: {
+          getLoadMetaForFile: vi.fn(() => ({ sequence: 1, kind: 'user' as const })),
+        },
+      });
+      const defaultEventDetail = {
         data,
         settings: {
           legendSettings: {},
           exportOptions: {},
-          eatOverlayEnabled: true,
         },
-        source: 'auto',
-      },
-    } as unknown as Event);
-    expect(legendElement.applyEatSettings).toHaveBeenLastCalledWith(
-      true,
-      DEFAULT_EAT_CONFIDENCE_THRESHOLD,
-    );
+        source: 'auto' as const,
+      };
+      const event = { detail: defaultEventDetail } as unknown as Event;
+
+      await controller.handleDataLoaded(event);
+
+      expect(legendElement.clearForNewDataset).toHaveBeenNthCalledWith(
+        1,
+        expect.any(String),
+        false,
+      );
+      expect(controlBar.clearForNewDataset).toHaveBeenNthCalledWith(1, expect.any(String), false);
+      expect(legendElement.setFileSettings).toHaveBeenNthCalledWith(
+        1,
+        {},
+        expect.any(String),
+        false,
+      );
+
+      latestViewRequest = {
+        requested: { tooltip: ['stale-annotation'] },
+        present: { annotation: false, projection: false, tooltip: true },
+        normalize: { annotation: false, projection: false, tooltip: false },
+      };
+      await controller.handleDataLoaded({
+        detail: {
+          ...defaultEventDetail,
+          source: 'user',
+          file: { name: 'custom.parquetbundle' } as File,
+        },
+      } as unknown as Event);
+
+      expect(legendElement.clearForNewDataset).toHaveBeenNthCalledWith(2, expect.any(String), true);
+      expect(viewController.setRequestedView).toHaveBeenCalledWith({
+        requested: { tooltip: undefined },
+        present: { annotation: false, projection: false, tooltip: false },
+        normalize: { annotation: false, projection: false, tooltip: true },
+      });
+
+      await controller.handleDataLoaded(event);
+
+      expect(legendElement.clearForNewDataset).toHaveBeenNthCalledWith(3, expect.any(String), true);
+      expect(controlBar.clearForNewDataset).toHaveBeenNthCalledWith(3, expect.any(String), true);
+      expect(legendElement.setFileSettings).toHaveBeenNthCalledWith(
+        3,
+        {},
+        expect.any(String),
+        true,
+      );
+    });
+
+    it('restores saved tooltip annotations on an automatic load with a silent URL', async () => {
+      mocks.readTooltipAnnotations.mockReturnValue(['ec']);
+      const { controller, viewController } = createControllerHarness();
+
+      await controller.handleDataLoaded({
+        detail: {
+          data,
+          settings: { legendSettings: {}, exportOptions: {} },
+          source: 'auto',
+        },
+      } as unknown as Event);
+
+      expect(viewController.setRequestedView).toHaveBeenCalledWith({
+        requested: { tooltip: ['ec'] },
+        present: { annotation: false, projection: false, tooltip: true },
+        normalize: { annotation: false, projection: false, tooltip: false },
+      });
+    });
   });
 });

apps/web/src/explore/dataset-controller.ts

@@ -108,7 +108,7 @@
       loadSequence = loadMeta.sequence;
 
       if (runningLoadMeta && loadMeta.sequence !== runningLoadMeta.sequence) {
         console.log('Ignoring stale data load result:', {
           source,
           fileName: file?.name ?? null,
           loadKind: loadMeta.kind,
@@ -132,16 +132,22 @@
       }
 
       const datasetHash = generateDatasetHash(data);
+      const hadPreviousDataset = currentDatasetHash !== null;
       const shouldClearPersistedState =
-        loadMeta.kind === 'default' || (loadMeta.kind === 'user' && settings != null);
+        (loadMeta.kind === 'default' && hadPreviousDataset) ||
+        (loadMeta.kind === 'user' && settings != null);
 
       legendElement.clearForNewDataset(datasetHash, shouldClearPersistedState);
       controlBar.clearForNewDataset(datasetHash, shouldClearPersistedState);
 
       await loadData(data);
 
       if (settings && loadMeta.kind !== 'opfs') {
-        legendElement.setFileSettings(settings.legendSettings, datasetHash, true);
+        legendElement.setFileSettings(
+          settings.legendSettings,
+          datasetHash,
+          shouldClearPersistedState,
+        );
       }
       if (settings) {
         const eatOverlayEnabled = settings.eatOverlayEnabled ?? true;
@@ -167,7 +173,6 @@
       // Must be set before the restore block so that any view-change emitted by
       // setRequestedView below is persisted under the new dataset's key, not the
       // previous dataset's key.
-      const hadPreviousDataset = currentDatasetHash !== null;
       currentDatasetHash = datasetHash;
 
       const latestRequest = viewController.getLatestViewRequest();
@@ -258,7 +263,7 @@
     const loadSequence = runningLoadMeta?.sequence ?? null;
 
     if (customEvent.detail.originalError?.name === 'AbortError') {
       console.log('Data load cancelled by user');
       if (loadSequence !== null) {
         loadQueue.resolvePendingLoadFinalization(loadSequence);
       }
@@ -303,7 +308,7 @@
     loadPersistedOrDefaultDataset: persistedDatasetController.loadPersistedOrDefaultDataset,
     tryLoadPersistedAgain: persistedDatasetController.tryLoadPersistedAgain,
     handleLoadingStart() {
       console.log('Data loading started');
       overlayController.update(true, 5, 'Analyzing file structure...', 'Starting upload...');
     },
     handleLoadingProgress(event: Event) {

apps/web/tests/dataset-reload.spec.ts

@@ -5,7 +5,9 @@ import {
   clickLegendItem,
   dismissTourIfPresent,
   getFirstLegendItemValue,
+  getShapeSizeState,
   isLegendItemHidden,
+  setShapeSize,
   waitForExploreDataLoad,
   waitForExploreInteractionReady,
   waitForPersistedExploreDataset,
@@ -285,7 +287,7 @@ async function hasLegacyNotificationHelperArtifacts(page: Page): Promise<boolean
 // Tests
 // ---------------------------------------------------------------------------
 
-test.describe('Dataset reload resets state (#178)', () => {
+test.describe('Dataset reload preserves legend state (#340)', () => {
   test.beforeEach(async ({ page }) => {
     // Each Playwright test receives a fresh context; shared storage state only
     // seeds the completed product-tour key, so OPFS starts empty here.
@@ -294,9 +296,7 @@ test.describe('Dataset reload resets state (#178)', () => {
     await dismissTourIfPresent(page);
   });
 
-  test('page reload restores default legend state and clears persisted hidden values', async ({
-    page,
-  }) => {
+  test('page reload restores persisted legend state and hidden values', async ({ page }) => {
     const itemValue = await getFirstLegendItemValue(page);
 
     expect(await isLegendItemHidden(page, itemValue)).toBe(false);
@@ -322,8 +322,8 @@ test.describe('Dataset reload resets state (#178)', () => {
     await waitForExploreDataLoad(page);
     await dismissTourIfPresent(page);
 
-    expect(await isLegendItemHidden(page, itemValue)).toBe(false);
-    expect(await itemHiddenInStorage()).toBe(false);
+    expect(await isLegendItemHidden(page, itemValue)).toBe(true);
+    expect(await itemHiddenInStorage()).toBe(true);
   });
 });
 
@@ -369,6 +369,8 @@ test.describe('Persisted custom datasets in OPFS (#176)', () => {
 
   test('reset to demo clears the persisted custom dataset', async ({ page }) => {
     const defaultCount = await getProteinCount(page);
+    const defaultShapeSize = await getShapeSizeState(page);
+    await setShapeSize(page, 42);
 
     await loadCustomDatasetFromImportMenu(page, CUSTOM_5K_BUNDLE_PATH);
     await page.waitForFunction(
@@ -385,6 +387,7 @@ test.describe('Persisted custom datasets in OPFS (#176)', () => {
 
     await loadDemoDatasetFromImportMenu(page);
     await waitForProteinCount(page, defaultCount);
+    await expect.poll(() => getShapeSizeState(page)).toEqual(defaultShapeSize);
 
     await page.reload();
     await waitForExploreDataLoad(page);