Keep wasm download feedback active through module initialization

src/worker.ts

@@ -4,11 +4,19 @@ import type { IDriveFSOptions } from '@jupyterlite/cockle';
 import { BaseShellWorker } from '@jupyterlite/cockle';
 import { DriveFS } from '@jupyterlite/services';
 
+type InitializeArgs = Parameters<BaseShellWorker['initialize']>;
+type DownloadModuleCallback = InitializeArgs[3];
+
 /**
  * Shell web worker that uses DriveFS via service worker.
  * Note that this is not exported as it is accessed from Shell via the filename.
  */
 class ShellWorker extends BaseShellWorker {
+  override async initialize(...args: InitializeArgs): Promise<void> {
+    args[3] = this._wrapDownloadModuleCallback(args[3]);
+    await super.initialize(...args);
+  }
+
   /**
    * Initialize the DriveFS to mount an external file system, if available.
    */
@@ -36,6 +44,32 @@ class ShellWorker extends BaseShellWorker {
       console.warn('Terminal not connected to shared drive');
     }
   }
+
+  private _wrapDownloadModuleCallback(
+    callback: DownloadModuleCallback
+  ): DownloadModuleCallback {
+    return ((packageName: string, moduleName: string, start: boolean): void => {
+      if (start) {
+        (self as any).Module = undefined;
+        callback(packageName, moduleName, start);
+        return;
+      }
+
+      const module = (self as any).Module;
+      if (typeof module !== 'function') {
+        callback(packageName, moduleName, start);
+        return;
+      }
+
+      (self as any).Module = async (...args: any[]): Promise<any> => {
+        try {
+          return await module(...args);
+        } finally {
+          callback(packageName, moduleName, start);
+        }
+      };
+    }) as DownloadModuleCallback;
+  }
 }
 
 const worker = new ShellWorker();