Extract KittyImageStorage, prevent mem leak on eviction

anthonykim1 · anthonykim1 · commit 27b01b22ef79 · 2026-02-14T15:17:29.000-08:00
diff --git a/addons/addon-image/src/ImageAddon.ts b/addons/addon-image/src/ImageAddon.ts
@@ -9,6 +9,7 @@ import { IIPHandler } from './IIPHandler';
 import { ImageRenderer } from './ImageRenderer';
 import { ImageStorage, CELL_SIZE_DEFAULT } from './ImageStorage';
 import { KittyGraphicsHandler } from './kitty/KittyGraphicsHandler';
+import { KittyImageStorage } from './kitty/KittyImageStorage';
 import { SixelHandler } from './SixelHandler';
 import { SixelImageStorage } from './SixelImageStorage';
 import { IIPImageStorage } from './IIPImageStorage';
@@ -154,9 +155,12 @@ export class ImageAddon implements ITerminalAddon, IImageApi {
 
     // Kitty graphics handler
     if (this._opts.kittySupport) {
-      const kittyHandler = new KittyGraphicsHandler(this._opts, this._renderer!, this._storage!, terminal);
+      const kittyStorage = new KittyImageStorage(this._storage!);
+      const kittyHandler = new KittyGraphicsHandler(this._opts, this._renderer!, kittyStorage, terminal);
       this._handlers.set('kitty', kittyHandler);
       this._disposeLater(
+        kittyStorage,
+        kittyHandler,
         terminal._core._inputHandler._parser.registerApcHandler(0x47, kittyHandler)
       );
     }
diff --git a/addons/addon-image/src/ImageStorage.ts b/addons/addon-image/src/ImageStorage.ts
@@ -124,6 +124,7 @@ export class ImageStorage implements IDisposable {
   private _pixelLimit: number = 2500000;
 
   private _viewportMetrics: { cols: number, rows: number };
+  public onImageDeleted: ((storageId: number) => void) | undefined;
 
   constructor(
     private _terminal: ITerminalExt,
@@ -189,11 +190,13 @@ export class ImageStorage implements IDisposable {
 
   private _delImg(id: number): void {
     const spec = this._images.get(id);
+    if (!spec) return;
     this._images.delete(id);
     // FIXME: really ugly workaround to get bitmaps deallocated :(
-    if (spec && window.ImageBitmap && spec.orig instanceof ImageBitmap) {
+    if (window.ImageBitmap && spec.orig instanceof ImageBitmap) {
       spec.orig.close();
     }
+    this.onImageDeleted?.(id);
   }
 
   /**
diff --git a/addons/addon-image/src/kitty/KittyGraphicsHandler.ts b/addons/addon-image/src/kitty/KittyGraphicsHandler.ts
@@ -3,9 +3,11 @@
  * @license MIT
  */
 
+import { IDisposable } from '@xterm/xterm';
 import { IApcHandler, IImageAddonOptions, IResetHandler, ITerminalExt, ImageLayer } from '../Types';
 import { ImageRenderer } from '../ImageRenderer';
-import { ImageStorage, CELL_SIZE_DEFAULT } from '../ImageStorage';
+import { CELL_SIZE_DEFAULT } from '../ImageStorage';
+import { KittyImageStorage } from './KittyImageStorage';
 import Base64Decoder, { type DecodeStatus } from 'xterm-wasm-parts/lib/base64/Base64Decoder.wasm';
 import {
   KittyAction,
@@ -36,7 +38,7 @@ const SEMICOLON = 0x3B;
 /**
  * Kitty graphics protocol handler with streaming base64 decoding.
  */
-export class KittyGraphicsHandler implements IApcHandler, IResetHandler {
+export class KittyGraphicsHandler implements IApcHandler, IResetHandler, IDisposable {
   private _aborted = false;
   private _decodeError = false;
 
@@ -69,21 +71,11 @@ export class KittyGraphicsHandler implements IApcHandler, IResetHandler {
    * When a chunk arrives with no i=, this key is used to find the pending upload.
    */
   private _lastPendingKey: number | undefined;
-  private _nextImageId = 1;
-  /** Maps Kitty protocol image ID → ImageStorage internal ID for deletion/lookup. */
-  private _kittyIdToStorageId: Map<number, number> = new Map();
-  // TODO: Eliminate double storage — raw image data lives here (as Blob) AND rendered
-  // ImageBitmaps live in ImageStorage. Currently we only use ImageStorage.addImage(bitmap)
-  // for tiling + cursor movement + marker-based eviction.
-  //
-
-  // See: https://github.com/xtermjs/xterm.js/pull/5619#issuecomment-3853678815
-  private _images: Map<number, IKittyImageData> = new Map();
 
   constructor(
     private readonly _opts: IImageAddonOptions,
     private readonly _renderer: ImageRenderer,
-    private readonly _storage: ImageStorage,
+    private readonly _kittyStorage: KittyImageStorage,
     private readonly _coreTerminal: ITerminalExt
   ) {
     // Convert decoded size limit -> max encoded bytes.
@@ -102,8 +94,11 @@ export class KittyGraphicsHandler implements IApcHandler, IResetHandler {
       this._activeDecoder.release();
       this._activeDecoder = null;
     }
-    this._images.clear();
-    this._kittyIdToStorageId.clear();
+    this._kittyStorage.reset();
+  }
+
+  public dispose(): void {
+    this.reset();
   }
 
   public start(): void {
@@ -397,16 +392,13 @@ export class KittyGraphicsHandler implements IApcHandler, IResetHandler {
 
     if (decodeError || bytes.length === 0) return true;
 
-    const id = cmd.id ?? this._nextImageId++;
-    const image: IKittyImageData = {
-      id,
+    this._kittyStorage.storeImage(cmd.id, {
       data: new Blob([bytes as BlobPart]),
       width: cmd.width ?? 0,
       height: cmd.height ?? 0,
       format: (cmd.format ?? KittyFormat.RGBA) as 24 | 32 | 100,
       compression: cmd.compression ?? ''
-    };
-    this._images.set(id, image);
+    });
     return true;
   }
 
@@ -426,8 +418,8 @@ export class KittyGraphicsHandler implements IApcHandler, IResetHandler {
     if (this._pendingTransmissions.has(pendingKey)) return true;
 
     // Display the completed image
-    const id = cmd.id ?? this._nextImageId - 1;
-    const image = this._images.get(id);
+    const id = cmd.id ?? this._kittyStorage.lastImageId;
+    const image = this._kittyStorage.getImage(id);
     if (image) {
       const result = this._displayImage(image, cmd);
       if (cmd.id !== undefined) {
@@ -508,7 +500,7 @@ export class KittyGraphicsHandler implements IApcHandler, IResetHandler {
         }
         this._pendingTransmissions.clear();
         this._lastPendingKey = undefined;
-        this._deleteAll();
+        this._kittyStorage.deleteAll();
         break;
       case 'i':
       case 'I':
@@ -522,7 +514,7 @@ export class KittyGraphicsHandler implements IApcHandler, IResetHandler {
               this._lastPendingKey = undefined;
             }
           }
-          this._deleteById(cmd.id);
+          this._kittyStorage.deleteById(cmd.id);
         }
         break;
       default:
@@ -532,23 +524,6 @@ export class KittyGraphicsHandler implements IApcHandler, IResetHandler {
     return true;
   }
 
-  private _deleteById(id: number): void {
-    this._images.delete(id);
-    const storageId = this._kittyIdToStorageId.get(id);
-    if (storageId !== undefined) {
-      this._storage.deleteImage(storageId);
-      this._kittyIdToStorageId.delete(id);
-    }
-  }
-
-  private _deleteAll(): void {
-    this._images.clear();
-    for (const storageId of this._kittyIdToStorageId.values()) {
-      this._storage.deleteImage(storageId);
-    }
-    this._kittyIdToStorageId.clear();
-  }
-
   private _sendResponse(id: number, message: string, quiet: number): void {
     const isOk = message === 'OK';
     if (isOk && quiet === 1) return;
@@ -602,14 +577,12 @@ export class KittyGraphicsHandler implements IApcHandler, IResetHandler {
     const layer: ImageLayer = (wantsBottom && this._coreTerminal.options.allowTransparency) ? 'bottom' : 'top';
 
     const zIndex = cmd.zIndex ?? 0;
-    let storageId: number;
     if (w !== bitmap.width || h !== bitmap.height) {
       const resized = await createImageBitmap(bitmap, { resizeWidth: w, resizeHeight: h });
-      storageId = this._storage.addImage(resized, true, layer, zIndex);
+      this._kittyStorage.addImage(image.id, resized, true, layer, zIndex);
     } else {
-      storageId = this._storage.addImage(bitmap, true, layer, zIndex);
+      this._kittyStorage.addImage(image.id, bitmap, true, layer, zIndex);
     }
-    this._kittyIdToStorageId.set(image.id, storageId);
 
     // Kitty cursor movement
     // Per spec: cursor placed at first column after last image column,
@@ -723,7 +696,11 @@ export class KittyGraphicsHandler implements IApcHandler, IResetHandler {
   }
 
   public get images(): ReadonlyMap<number, IKittyImageData> {
-    return this._images;
+    return this._kittyStorage.images;
+  }
+
+  public get _kittyIdToStorageId(): ReadonlyMap<number, number> {
+    return this._kittyStorage.kittyIdToStorageId;
   }
 
   public get pendingTransmissions(): ReadonlyMap<number, IPendingTransmission> {
diff --git a/addons/addon-image/src/kitty/KittyImageStorage.ts b/addons/addon-image/src/kitty/KittyImageStorage.ts
@@ -0,0 +1,131 @@
+/**
+ * Copyright (c) 2026 The xterm.js authors. All rights reserved.
+ * @license MIT
+ */
+
+import { IDisposable } from '@xterm/xterm';
+import { ImageStorage } from '../ImageStorage';
+import { ImageLayer } from '../Types';
+import { IKittyImageData } from './KittyGraphicsTypes';
+
+/**
+ * Kitty-specific image storage controller.
+ *
+ * Wraps shared ImageStorage with kitty protocol semantics:
+ * - tracks transmitted image payloads by kitty image id
+ * - tracks kitty image id -> shared ImageStorage id mapping for displayed images
+ * - mirrors shared-storage evictions into kitty maps
+ * - applies protocol-level undisplayed-image eviction policy
+ */
+export class KittyImageStorage implements IDisposable {
+  private static readonly _maxStoredImages = 256;
+
+  private _nextImageId = 1;
+  private readonly _images: Map<number, IKittyImageData> = new Map();
+  private readonly _kittyIdToStorageId: 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;
+      }
+    }
+  };
+
+  constructor(
+    private readonly _storage: ImageStorage
+  ) {
+    this._previousOnImageDeleted = this._storage.onImageDeleted;
+    this._wrappedOnImageDeleted = (storageId: number) => {
+      this._previousOnImageDeleted?.(storageId);
+      this._handleStorageImageDeleted(storageId);
+    };
+    this._storage.onImageDeleted = this._wrappedOnImageDeleted;
+  }
+
+  public reset(): void {
+    this._nextImageId = 1;
+    this._images.clear();
+    this._kittyIdToStorageId.clear();
+  }
+
+  public dispose(): void {
+    this.reset();
+    if (this._storage.onImageDeleted === this._wrappedOnImageDeleted) {
+      this._storage.onImageDeleted = this._previousOnImageDeleted;
+    }
+  }
+
+  public storeImage(id: number | undefined, imageData: Omit<IKittyImageData, 'id'>): number {
+    const imageId = id ?? this._nextImageId++;
+
+    const oldStorageId = this._kittyIdToStorageId.get(imageId);
+    if (oldStorageId !== undefined) {
+      this._storage.deleteImage(oldStorageId);
+      this._kittyIdToStorageId.delete(imageId);
+    }
+
+    if (!this._images.has(imageId) && this._images.size >= KittyImageStorage._maxStoredImages) {
+      this._evictUndisplayedImages();
+    }
+
+    this._images.set(imageId, {
+      ...imageData,
+      id: imageId
+    });
+    return imageId;
+  }
+
+  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);
+  }
+
+  public getImage(kittyId: number): IKittyImageData | undefined {
+    return this._images.get(kittyId);
+  }
+
+  public deleteById(kittyId: number): void {
+    this._images.delete(kittyId);
+    const storageId = this._kittyIdToStorageId.get(kittyId);
+    if (storageId !== undefined) {
+      this._storage.deleteImage(storageId);
+      this._kittyIdToStorageId.delete(kittyId);
+    }
+  }
+
+  public deleteAll(): void {
+    this._images.clear();
+    for (const storageId of this._kittyIdToStorageId.values()) {
+      this._storage.deleteImage(storageId);
+    }
+    this._kittyIdToStorageId.clear();
+  }
+
+  public get images(): ReadonlyMap<number, IKittyImageData> {
+    return this._images;
+  }
+
+  public get kittyIdToStorageId(): ReadonlyMap<number, number> {
+    return this._kittyIdToStorageId;
+  }
+
+  public get lastImageId(): number {
+    return this._nextImageId - 1;
+  }
+
+  private _evictUndisplayedImages(): void {
+    for (const [kittyId] of this._images) {
+      if (this._images.size <= KittyImageStorage._maxStoredImages / 2) {
+        break;
+      }
+      if (!this._kittyIdToStorageId.has(kittyId)) {
+        this._images.delete(kittyId);
+      }
+    }
+  }
+}
diff --git a/addons/addon-image/test/KittyGraphics.test.ts b/addons/addon-image/test/KittyGraphics.test.ts
