refactor: replace state file callbacks with result-based flow

Replace StateFileContext callbacks with StateFileSetupResult type.
State file restoration now collects results and completes at end of
import instead of tracking pending counts with callbacks.

Author: PaulHax

src/io/import/common.ts

@@ -38,6 +38,13 @@ export interface IntermediateResult {
   dataSources: DataSource[];
 }
 
+export interface StateFileSetupResult {
+  type: 'stateFileSetup';
+  dataSources: DataSource[];
+  manifest: Manifest;
+  stateFiles: FileEntry[];
+}
+
 export interface ErrorResult {
   type: 'error';
   error: Error;
@@ -48,6 +55,7 @@ export type ImportResult =
   | LoadableResult
   | ConfigResult
   | IntermediateResult
+  | StateFileSetupResult
   | OkayResult
   | ErrorResult;
 
@@ -101,31 +109,14 @@ export const asOkayResult = (dataSource: DataSource): OkayResult => ({
 export type ArchiveContents = Record<string, File>;
 export type ArchiveCache = Map<File, Awaitable<ArchiveContents>>;
 
-export interface StateFileContext {
-  manifest: Manifest;
-  stateFiles: FileEntry[];
-  stateIDToStoreID: Map<string, string>;
-  pendingLeafCount: number;
-  onLeafImported: (stateID: string, storeID: string) => void;
-  onAllLeavesImported: () => Promise<void>;
-}
-
 export interface ImportContext {
-  // Caches URL responses
   fetchFileCache?: FetchCache<File>;
-  // Caches archives. ArchiveFile -> { [archivePath]: File }
   archiveCache?: ArchiveCache;
-  // Records dicom files
   dicomDataSources?: DataSource[];
   onCleanup?: (fn: () => void) => void;
-  /**
-   * A reference to importDataSources for nested imports.
-   */
   importDataSources?: (
     dataSources: DataSource[]
   ) => Promise<ImportDataSourcesResult[]>;
-  // State file restoration context for 3-phase restoration
-  stateFileContext?: StateFileContext;
 }
 
 export type ImportHandler = ChainHandler<

src/io/import/importDataSources.ts

@@ -10,20 +10,20 @@ import {
   ErrorResult,
   ImportDataSourcesResult,
   asIntermediateResult,
+  StateFileSetupResult,
 } from '@/src/io/import/common';
-import {
-  DataSource,
-  ChunkSource,
-  StateFileLeaf,
-} from '@/src/io/import/dataSource';
+import { DataSource, ChunkSource } from '@/src/io/import/dataSource';
 import handleDicomFile from '@/src/io/import/processors/handleDicomFile';
 import extractArchive from '@/src/io/import/processors/extractArchive';
 import extractArchiveTarget from '@/src/io/import/processors/extractArchiveTarget';
 import handleAmazonS3 from '@/src/io/import/processors/handleAmazonS3';
 import handleGoogleCloudStorage from '@/src/io/import/processors/handleGoogleCloudStorage';
 import importSingleFile from '@/src/io/import/processors/importSingleFile';
 import handleRemoteManifest from '@/src/io/import/processors/remoteManifest';
-import { restoreStateFile } from '@/src/io/import/processors/restoreStateFile';
+import {
+  restoreStateFile,
+  completeStateFileRestore,
+} from '@/src/io/import/processors/restoreStateFile';
 import updateFileMimeType from '@/src/io/import/processors/updateFileMimeType';
 import handleConfig from '@/src/io/import/processors/handleConfig';
 import {
@@ -74,39 +74,18 @@ const applyConfigsPostState = (
     }
   });
 
-function findStateFileLeaf(dataSource: DataSource): StateFileLeaf | undefined {
+function findStateFileLeaf(dataSource: DataSource) {
   let current: DataSource | undefined = dataSource;
   while (current) {
     if (current.stateFileLeaf) return current.stateFileLeaf;
     current = current.parent;
   }
-  // For collections (DICOM volumes), check the first source's parent chain
   if (dataSource.type === 'collection' && dataSource.sources.length > 0) {
     return findStateFileLeaf(dataSource.sources[0]);
   }
   return undefined;
 }
 
-async function handleStateFileResult(
-  result: LoadableResult,
-  importContext: ImportContext
-) {
-  const stateLeaf = findStateFileLeaf(result.dataSource);
-  if (stateLeaf && importContext.stateFileContext) {
-    const ctx = importContext.stateFileContext;
-    ctx.stateIDToStoreID.set(stateLeaf.stateID, result.dataID);
-
-    // Phase 2: Immediately bind view to this data so user sees streaming
-    ctx.onLeafImported(stateLeaf.stateID, result.dataID);
-
-    ctx.pendingLeafCount--;
-    if (ctx.pendingLeafCount === 0) {
-      // Phase 3: Restore tools/segments after all data loaded
-      await ctx.onAllLeavesImported();
-    }
-  }
-}
-
 async function importDicomChunkSources(sources: ChunkSource[]) {
   if (sources.length === 0) return [];
 
@@ -180,6 +159,7 @@ export async function importDataSources(
 
   const chunkSources: DataSource[] = [];
   const configResults: ConfigResult[] = [];
+  const stateFileSetups: StateFileSetupResult[] = [];
   const results: ImportDataSourcesResult[] = [];
 
   let queue = dataSources.map((src) => ({
@@ -194,14 +174,16 @@ export async function importDataSources(
     queue = queue.filter((_, i) => i !== index);
 
     switch (result.type) {
+      case 'stateFileSetup':
+        stateFileSetups.push(result);
+      // fallthrough to handle dataSources
       case 'intermediate': {
         const [chunks, otherSources] = partition(
           (ds) => ds.type === 'chunk',
           result.dataSources
         );
         chunkSources.push(...chunks);
 
-        // try loading intermediate results
         queue.push(
           ...otherSources.map((src) => ({
             promise: evaluateChain(src, handlers, importContext),
@@ -220,11 +202,8 @@ export async function importDataSources(
         break;
       case 'ok':
       case 'error':
-        results.push(result);
-        break;
       case 'data':
         results.push(result);
-        await handleStateFileResult(result, importContext);
         break;
       default:
         throw new Error(`Invalid result: ${result}`);
@@ -243,10 +222,6 @@ export async function importDataSources(
   try {
     const dicomResults = await importDicomChunkSources(dicomChunkSources);
     results.push(...dicomResults);
-    // Handle state file results for DICOM volumes
-    for (const dicomResult of dicomResults) {
-      await handleStateFileResult(dicomResult, importContext);
-    }
   } catch (err) {
     const errorSource =
       dicomChunkSources.length === 1
@@ -255,11 +230,27 @@ export async function importDataSources(
     results.push(asErrorResult(ensureError(err), errorSource));
   }
 
-  // save data sources
-  useDatasetStore().addDataSources(
-    results.filter((result): result is LoadableResult => result.type === 'data')
+  const loadableResults = results.filter(
+    (r): r is LoadableResult => r.type === 'data'
   );
 
+  useDatasetStore().addDataSources(loadableResults);
+
+  for (const setup of stateFileSetups) {
+    const stateIDToStoreID: Record<string, string> = {};
+    for (const loadable of loadableResults) {
+      const leaf = findStateFileLeaf(loadable.dataSource);
+      if (leaf) {
+        stateIDToStoreID[leaf.stateID] = loadable.dataID;
+      }
+    }
+    await completeStateFileRestore(
+      setup.manifest,
+      setup.stateFiles,
+      stateIDToStoreID
+    );
+  }
+
   return results;
 }
 

src/io/import/processors/restoreStateFile.ts

@@ -5,9 +5,8 @@ import {
 } from '@/src/io/state-file/schema';
 import {
   asErrorResult,
-  asIntermediateResult,
   ImportHandler,
-  StateFileContext,
+  StateFileSetupResult,
 } from '@/src/io/import/common';
 import { MANIFEST, isStateFile } from '@/src/io/state-file';
 import { partition, getURLBasename } from '@/src/utils';
@@ -113,27 +112,38 @@ function prepareLeafDataSources(manifest: Manifest, datasetFiles: FileEntry[]) {
   });
 }
 
-async function completeStateFileRestore(ctx: StateFileContext) {
-  const { manifest, stateFiles, stateIDToStoreID } = ctx;
-  const stateIDToStoreIDRecord = Object.fromEntries(stateIDToStoreID);
+export async function completeStateFileRestore(
+  manifest: Manifest,
+  stateFiles: FileEntry[],
+  stateIDToStoreID: Record<string, string>
+) {
+  const viewStore = useViewStore();
 
-  useViewConfigStore().deserializeAll(manifest, stateIDToStoreIDRecord);
+  Object.entries(stateIDToStoreID).forEach(([stateID, storeID]) => {
+    viewStore.bindViewsToData(stateID, storeID, manifest);
+  });
+
+  if (!manifest.viewByID) {
+    const firstStoreID = Object.values(stateIDToStoreID)[0];
+    if (firstStoreID) {
+      viewStore.setDataForAllViews(firstStoreID);
+    }
+  }
+
+  useViewConfigStore().deserializeAll(manifest, stateIDToStoreID);
 
   const segmentGroupIDMap = await useSegmentGroupStore().deserialize(
     manifest,
     stateFiles,
-    stateIDToStoreIDRecord
+    stateIDToStoreID
   );
-  useLayersStore().deserialize(manifest, stateIDToStoreIDRecord);
 
-  useToolStore().deserialize(
-    manifest,
-    segmentGroupIDMap,
-    stateIDToStoreIDRecord
-  );
+  useLayersStore().deserialize(manifest, stateIDToStoreID);
+
+  useToolStore().deserialize(manifest, segmentGroupIDMap, stateIDToStoreID);
 }
 
-export const restoreStateFile: ImportHandler = async (dataSource, context) => {
+export const restoreStateFile: ImportHandler = async (dataSource) => {
   if (dataSource.type === 'file' && (await isStateFile(dataSource.file))) {
     const stateFileContents = await extractFilesFromZip(dataSource.file);
 
@@ -158,61 +168,14 @@ export const restoreStateFile: ImportHandler = async (dataSource, context) => {
       );
     }
 
-    // Phase 1: Set up view layout immediately (without data bindings)
-    const viewStore = useViewStore();
-    viewStore.deserializeLayout(manifest);
-
-    // Prepare leaf data sources with state file tags
-    const leafDataSources = prepareLeafDataSources(manifest, restOfStateFile);
-
-    if (leafDataSources.length === 0) {
-      // No datasets to import, complete restoration immediately
-      await completeStateFileRestore({
-        manifest,
-        stateFiles: restOfStateFile,
-        stateIDToStoreID: new Map(),
-        pendingLeafCount: 0,
-        onLeafImported: () => {},
-        onAllLeavesImported: async () => {},
-      });
-
-      // When viewByID is not in manifest, there's no data to assign
-      return asIntermediateResult([]);
-    }
+    useViewStore().deserializeLayout(manifest);
 
-    // Set up state file context for phase 2 and 3 callbacks
-    const stateFileContext: StateFileContext = {
+    return {
+      type: 'stateFileSetup',
+      dataSources: prepareLeafDataSources(manifest, restOfStateFile),
       manifest,
       stateFiles: restOfStateFile,
-      stateIDToStoreID: new Map(),
-      pendingLeafCount: leafDataSources.length,
-      onLeafImported: (stateID: string, storeID: string) => {
-        // Phase 2: Bind view to data as each leaf completes
-        viewStore.bindViewsToData(stateID, storeID, manifest);
-      },
-      onAllLeavesImported: async () => {
-        // Phase 3: Restore segment groups, tools, layers after all data loaded
-        await completeStateFileRestore(stateFileContext);
-
-        // When viewByID is not in manifest, assign first dataset to all views
-        if (!manifest.viewByID) {
-          const firstStoreID = stateFileContext.stateIDToStoreID
-            .values()
-            .next().value;
-          if (firstStoreID) {
-            viewStore.setDataForAllViews(firstStoreID);
-          }
-        }
-      },
-    };
-
-    // Store context for use by main pipeline
-    if (context) {
-      context.stateFileContext = stateFileContext;
-    }
-
-    // Return leaf data sources to be processed by main pipeline
-    return asIntermediateResult(leafDataSources);
+    } as StateFileSetupResult;
   }
   return Skip;
 };