fix more mem leak + cleanups

Author: anthonykim1

addons/addon-image/src/kitty/KittyGraphicsHandler.ts

@@ -81,11 +81,7 @@ export class KittyGraphicsHandler implements IApcHandler, IResetHandler, IDispos
   }
 
   public reset(): void {
-    for (const pending of this._pendingTransmissions.values()) {
-      pending.decoder.release();
-    }
-    this._pendingTransmissions.clear();
-    this._lastPendingKey = undefined;
+    this._cleanupAllPending();
     if (this._activeDecoder) {
       this._activeDecoder.release();
       this._activeDecoder = null;
@@ -97,6 +93,21 @@ export class KittyGraphicsHandler implements IApcHandler, IResetHandler, IDispos
     this.reset();
   }
 
+  private _removePendingEntry(key: number): void {
+    this._pendingTransmissions.delete(key);
+    if (this._lastPendingKey === key) {
+      this._lastPendingKey = undefined;
+    }
+  }
+
+  private _cleanupAllPending(): void {
+    for (const pending of this._pendingTransmissions.values()) {
+      pending.decoder.release();
+    }
+    this._pendingTransmissions.clear();
+    this._lastPendingKey = undefined;
+  }
+
   public start(): void {
     this._aborted = false;
     this._decodeError = false;
@@ -178,10 +189,7 @@ export class KittyGraphicsHandler implements IApcHandler, IResetHandler, IDispos
       }
       this._activeDecoder = null;
       if (pending) {
-        this._pendingTransmissions.delete(pendingKey);
-        if (this._lastPendingKey === pendingKey) {
-          this._lastPendingKey = undefined;
-        }
+        this._removePendingEntry(pendingKey);
       }
       this._aborted = true;
       return;
@@ -202,10 +210,7 @@ export class KittyGraphicsHandler implements IApcHandler, IResetHandler, IDispos
       this._activeDecoder = null;
       this._decodeError = true;
       if (pending) {
-        this._pendingTransmissions.delete(pendingKey);
-        if (this._lastPendingKey === pendingKey) {
-          this._lastPendingKey = undefined;
-        }
+        this._removePendingEntry(pendingKey);
       }
     }
   }
@@ -404,14 +409,8 @@ export class KittyGraphicsHandler implements IApcHandler, IResetHandler, IDispos
       return true;
     }
 
-    const pendingKey = cmd.id ?? 0;
-
     this._handleTransmit(cmd, bytes, decodeError);
 
-    // If still accumulating chunks, don't display yet
-    if (this._pendingTransmissions.has(pendingKey)) return true;
-
-    // Display the completed image
     const id = cmd.id ?? this._kittyStorage.lastImageId;
     const image = this._kittyStorage.getImage(id);
     if (image) {
@@ -488,26 +487,17 @@ export class KittyGraphicsHandler implements IApcHandler, IResetHandler, IDispos
     switch (selector) {
       case 'a':
       case 'A':
-        // Delete all — also abort all in-flight uploads
-        for (const pending of this._pendingTransmissions.values()) {
-          pending.decoder.release();
-        }
-        this._pendingTransmissions.clear();
-        this._lastPendingKey = undefined;
+        this._cleanupAllPending();
         this._kittyStorage.deleteAll();
         break;
       case 'i':
       case 'I':
-        // Delete by image ID — only abort the targeted upload
         if (cmd.id !== undefined) {
           const pending = this._pendingTransmissions.get(cmd.id);
           if (pending) {
             pending.decoder.release();
-            this._pendingTransmissions.delete(cmd.id);
-            if (this._lastPendingKey === cmd.id) {
-              this._lastPendingKey = undefined;
-            }
           }
+          this._removePendingEntry(cmd.id);
           this._kittyStorage.deleteById(cmd.id);
         }
         break;
@@ -574,6 +564,7 @@ export class KittyGraphicsHandler implements IApcHandler, IResetHandler, IDispos
     const zIndex = cmd.zIndex ?? 0;
     if (w !== bitmap.width || h !== bitmap.height) {
       const resized = await createImageBitmap(bitmap, { resizeWidth: w, resizeHeight: h });
+      bitmap.close();
       this._kittyStorage.addImage(image.id, resized, true, layer, zIndex);
     } else {
       this._kittyStorage.addImage(image.id, bitmap, true, layer, zIndex);
@@ -615,7 +606,10 @@ export class KittyGraphicsHandler implements IApcHandler, IResetHandler, IDispos
             canvas.getContext('2d')?.drawImage(img, 0, 0);
             createImageBitmap(canvas).then(resolve).catch(reject);
           });
-          img.addEventListener('error', reject);
+          img.addEventListener('error', () => {
+            URL.revokeObjectURL(url);
+            reject(new Error('Failed to load image'));
+          });
           img.src = url;
         });
       }

addons/addon-image/src/kitty/KittyImageStorage.ts

@@ -21,16 +21,16 @@ export class KittyImageStorage implements IDisposable {
   private _nextImageId = 1;
   private readonly _images: Map<number, IKittyImageData> = new Map();
   private readonly _kittyIdToStorageId: Map<number, number> = new Map();
+  private readonly _storageIdToKittyId: Map<number, number> = new Map();
 
   private readonly _previousOnImageDeleted: ((storageId: number) => void) | undefined;
   private readonly _wrappedOnImageDeleted: (storageId: number) => void;
   private readonly _handleStorageImageDeleted = (storageId: number): void => {
-    for (const [kittyId, mappedStorageId] of this._kittyIdToStorageId) {
-      if (mappedStorageId === storageId) {
-        this._kittyIdToStorageId.delete(kittyId);
-        this._images.delete(kittyId);
-        break;
-      }
+    const kittyId = this._storageIdToKittyId.get(storageId);
+    if (kittyId !== undefined) {
+      this._kittyIdToStorageId.delete(kittyId);
+      this._storageIdToKittyId.delete(storageId);
+      this._images.delete(kittyId);
     }
   };
 
@@ -49,6 +49,7 @@ export class KittyImageStorage implements IDisposable {
     this._nextImageId = 1;
     this._images.clear();
     this._kittyIdToStorageId.clear();
+    this._storageIdToKittyId.clear();
   }
 
   public dispose(): void {
@@ -65,6 +66,7 @@ export class KittyImageStorage implements IDisposable {
     if (oldStorageId !== undefined) {
       this._storage.deleteImage(oldStorageId);
       this._kittyIdToStorageId.delete(imageId);
+      this._storageIdToKittyId.delete(oldStorageId);
     }
 
     if (!this._images.has(imageId) && this._images.size >= KittyImageStorage._maxStoredImages) {
@@ -81,6 +83,7 @@ export class KittyImageStorage implements IDisposable {
   public addImage(kittyId: number, image: HTMLCanvasElement | ImageBitmap, scrolling: boolean, layer: ImageLayer, zIndex: number): void {
     const storageId = this._storage.addImage(image, scrolling, layer, zIndex);
     this._kittyIdToStorageId.set(kittyId, storageId);
+    this._storageIdToKittyId.set(storageId, kittyId);
   }
 
   public getImage(kittyId: number): IKittyImageData | undefined {
@@ -93,6 +96,7 @@ export class KittyImageStorage implements IDisposable {
     if (storageId !== undefined) {
       this._storage.deleteImage(storageId);
       this._kittyIdToStorageId.delete(kittyId);
+      this._storageIdToKittyId.delete(storageId);
     }
   }
 
@@ -102,6 +106,7 @@ export class KittyImageStorage implements IDisposable {
       this._storage.deleteImage(storageId);
     }
     this._kittyIdToStorageId.clear();
+    this._storageIdToKittyId.clear();
   }
 
   public get images(): ReadonlyMap<number, IKittyImageData> {