Pre-calculate + fix bug to make sure we dont reset per chunk

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

@@ -65,6 +65,10 @@ export class KittyGraphicsHandler implements IApcHandler, IResetHandler {
   private _decodedChunks: Uint8Array[] = [];
   private _totalDecodedSize = 0;
 
+  /** Pre-calculated encoded size limit */
+  private _encodedSizeLimit = 0;
+  private _totalEncodedSize = 0;
+
   /** Parsed command. These are the control data before semicolon. */
   private _parsedCommand: IKittyCommand | null = null;
 
@@ -100,6 +104,9 @@ export class KittyGraphicsHandler implements IApcHandler, IResetHandler {
     this._decodedChunks = [];
     this._totalDecodedSize = 0;
     this._parsedCommand = null;
+    // Pre-calculate encoded limit once: base64 is 4 bytes encoded → 3 bytes decoded
+    this._encodedSizeLimit = Math.ceil(this._opts.kittySizeLimit * 4 / 3);
+    this._totalEncodedSize = 0;
   }
 
   public put(data: Uint32Array, start: number, end: number): void {
@@ -158,9 +165,14 @@ export class KittyGraphicsHandler implements IApcHandler, IResetHandler {
   private _streamPayload(data: Uint32Array, start: number, end: number): void {
     if (this._aborted) return;
 
-    // Check size limit
-    const estimatedTotal = this._totalDecodedSize + Math.ceil(this._shardBufferPos * 3 / 4) + Math.ceil((end - start) * 3 / 4);
-    if (estimatedTotal > this._opts.kittySizeLimit) {
+    // Check size limit (compare encoded bytes against pre-calculated limit)
+    // Include cumulative size from pending transmission for multi-chunk images
+    const pendingKey = this._parsedCommand?.id ?? 0;
+    const pending = this._pendingTransmissions.get(pendingKey);
+    const previousEncodedSize = pending?.totalEncodedSize ?? 0;
+    this._totalEncodedSize += end - start;
+    const cumulativeEncodedSize = previousEncodedSize + this._totalEncodedSize;
+    if (cumulativeEncodedSize > this._encodedSizeLimit) {
       this._aborted = true;
       return;
     }
@@ -290,6 +302,7 @@ export class KittyGraphicsHandler implements IApcHandler, IResetHandler {
           cmd: { ...cmd },
           chunks: [],
           totalSize: 0,
+          totalEncodedSize: 0,
           leftover: new Uint32Array(4),
           leftoverLength: 0
         };
@@ -466,6 +479,8 @@ export class KittyGraphicsHandler implements IApcHandler, IResetHandler {
         pending.chunks.push(bytes);
         pending.totalSize += bytes.length;
       }
+      // Track cumulative encoded size for size limit enforcement
+      pending.totalEncodedSize += this._totalEncodedSize;
 
       if (isMoreComing) return true;
 
@@ -498,6 +513,7 @@ export class KittyGraphicsHandler implements IApcHandler, IResetHandler {
         cmd: { ...cmd },
         chunks: bytes.length > 0 ? [bytes] : [],
         totalSize: bytes.length,
+        totalEncodedSize: this._totalEncodedSize,
         leftover: new Uint32Array(4),
         leftoverLength: 0
       });

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

@@ -103,8 +103,10 @@ export interface IPendingTransmission {
   cmd: IKittyCommand;
   /** Accumulated decoded payload chunks */
   chunks: Uint8Array[];
-  /** Total size of accumulated chunks */
+  /** Total size of accumulated decoded chunks */
   totalSize: number;
+  /** Total encoded (base64) bytes received across all chunks - for size limit enforcement */
+  totalEncodedSize: number;
   /** Leftover base64 bytes (0-3) that weren't aligned in previous chunk */
   leftover: Uint32Array;
   /** Number of valid bytes in leftover */

addons/addon-image/test/ImageAddon.test.ts

@@ -391,6 +391,39 @@ test.describe('ImageAddon', () => {
       deepStrictEqual(storedData, KITTY_BLACK_1X1_BYTES);
     });
 
+    test('enforces size limit across chunked transmissions', async () => {
+      // Create a custom addon with very small size limit (100 bytes)
+      // The 1x1 PNG is ~164 bytes base64, so 2 chunks should exceed 100
+      await ctx.page.evaluate(() => {
+        (window as any).smallLimitAddon = new ImageAddon({
+          kittySupport: true,
+          kittySizeLimit: 100  // Very small limit
+        });
+        (window as any).term.loadAddon((window as any).smallLimitAddon);
+      });
+
+      // Split the base64 data into two chunks
+      const half = Math.floor(KITTY_BLACK_1X1_BASE64.length / 2);
+      const part1 = KITTY_BLACK_1X1_BASE64.substring(0, half);
+      const part2 = KITTY_BLACK_1X1_BASE64.substring(half);
+
+      // Send chunked data - first chunk (~82 bytes) is under limit
+      await ctx.proxy.write(`\x1b_Ga=t,f=100,i=777,m=1;${part1}\x1b\\`);
+      await timeout(50);
+
+      // Second chunk brings total to ~164 bytes, exceeding 100 byte limit
+      await ctx.proxy.write(`\x1b_Ga=t,f=100,i=777;${part2}\x1b\\`);
+      await timeout(100);
+
+      // Image should NOT be stored due to size limit
+      strictEqual(await ctx.page.evaluate(`window.smallLimitAddon._handlers.get('kitty').images.has(777)`), false);
+
+      // Cleanup
+      await ctx.page.evaluate(() => {
+        (window as any).smallLimitAddon.dispose();
+      });
+    });
+
     test('delete command (a=d) removes specific image by id', async () => {
       await ctx.proxy.write(`\x1b_Ga=t,f=100,i=10;${KITTY_BLACK_1X1_BASE64}\x1b\\`);
       await timeout(50);