Kitty support did not have any ways of limiting how much memory we use. (#433)

Based on the spec, this now:

* A default maximum of 320 megs, up to 4 gigs of ram.
* If we hit the limit, the first images to go will be the not-visible ones,
  follow by the oldest ones, and then anything is fair game until we meet the
  quota.
* RIS reset drops all kitty images, and entering the alt buffer clears any
  lingering alt images even when you switch back using ?47l (which preserves the alt buffer state)

Author: migueldeicaza

Sources/SwiftTerm/KittyGraphics.swift

@@ -121,6 +121,8 @@ enum KittyGraphicsPayload {
 
 struct KittyGraphicsImage {
     let payload: KittyGraphicsPayload
+    let byteSize: Int
+    var lastAccessTick: UInt64
 }
 
 struct KittyGraphicsPending {
@@ -135,11 +137,14 @@ final class KittyGraphicsState {
     var nextPlacementId: UInt32 = 1
     var pending: KittyGraphicsPending?
     var placementsByKey: [KittyPlacementKey: KittyPlacementRecord] = [:]
+    var totalImageBytes: Int = 0
+    var nextImageAccessTick: UInt64 = 1
 }
 
 extension Terminal {
     private static let kittyMaxImageBytes = 400 * 1024 * 1024
     private static let kittyMaxImageDimension = 10000
+    private static let kittyMaxImageCacheBytes = 4 * 1024 * 1024 * 1024
 
     func handleKittyGraphics(_ data: ArraySlice<UInt8>) {
         guard let (control, payload) = parseKittyGraphicsControl(data) else {
@@ -333,10 +338,7 @@ extension Terminal {
         }
 
         if let id = resolved.imageId {
-            kittyGraphicsState.imagesById[id] = KittyGraphicsImage(payload: payload)
-            if let number = resolved.imageNumber {
-                kittyGraphicsState.imageNumbers[number] = id
-            }
+            storeKittyImage(payload: payload, imageId: id, imageNumber: resolved.imageNumber)
         }
 
         var displayed = true
@@ -440,10 +442,10 @@ extension Terminal {
     }
 
     private func resolveKittyImageForDisplay(control: KittyGraphicsControl) -> (image: KittyGraphicsImage?, imageId: UInt32?, imageNumber: UInt32?, shouldReply: Bool) {
-        if let number = control.imageNumber, let imageId = kittyGraphicsState.imageNumbers[number], let image = kittyGraphicsState.imagesById[imageId] {
+        if let number = control.imageNumber, let imageId = kittyGraphicsState.imageNumbers[number], let image = updateKittyImageAccess(imageId: imageId) {
             return (image, imageId, number, control.suppressResponses == 0)
         }
-        if let imageId = control.imageId, let image = kittyGraphicsState.imagesById[imageId] {
+        if let imageId = control.imageId, let image = updateKittyImageAccess(imageId: imageId) {
             return (image, imageId, nil, control.suppressResponses == 0)
         }
         return (nil, nil, nil, control.suppressResponses == 0)
@@ -1468,9 +1470,23 @@ extension Terminal {
         kittyGraphicsState.imagesById.removeAll()
         kittyGraphicsState.imageNumbers.removeAll()
         kittyGraphicsState.placementsByKey.removeAll()
+        kittyGraphicsState.totalImageBytes = 0
+        kittyGraphicsState.nextImageAccessTick = 1
         updateRange(startLine: buffer.scrollTop, endLine: buffer.scrollBottom)
     }
 
+    func clearKittyImages(in buffer: Buffer, isAlternateBuffer: Bool) {
+        let removedKeys = removeKittyPlacements(in: buffer, lineRange: 0..<buffer.lines.count) { _ in true }
+        let recordKeys = removePlacementRecords { record in
+            record.isAlternateBuffer == isAlternateBuffer
+        }
+        let extraKeys = recordKeys.subtracting(removedKeys)
+        if !extraKeys.isEmpty {
+            _ = removeKittyPlacementsByKey(extraKeys)
+        }
+        cleanupUnusedKittyImages()
+    }
+
     private func deletePlacementsVisibleOnScreen() {
         let start = buffer.yBase
         let end = min(buffer.yBase + rows, buffer.lines.count)
@@ -1704,20 +1720,113 @@ extension Terminal {
     }
 
     private func cleanupUnusedKittyImages() {
+        let used = collectUsedKittyImageIds()
+        let unusedIds = kittyGraphicsState.imagesById.keys.filter { !used.contains($0) }
+        for id in unusedIds {
+            removeKittyImage(imageId: id)
+        }
+    }
+
+    private func storeKittyImage(payload: KittyGraphicsPayload, imageId: UInt32, imageNumber: UInt32?) {
+        let byteSize = kittyPayloadByteSize(payload)
+        let lastAccessTick = nextKittyImageAccessTick()
+        if let existing = kittyGraphicsState.imagesById[imageId] {
+            kittyGraphicsState.totalImageBytes = max(0, kittyGraphicsState.totalImageBytes - existing.byteSize)
+        }
+        kittyGraphicsState.imagesById[imageId] = KittyGraphicsImage(payload: payload,
+                                                                   byteSize: byteSize,
+                                                                   lastAccessTick: lastAccessTick)
+        kittyGraphicsState.totalImageBytes += byteSize
+        if let number = imageNumber {
+            kittyGraphicsState.imageNumbers[number] = imageId
+        }
+        enforceKittyImageCacheLimit()
+    }
+
+    private func updateKittyImageAccess(imageId: UInt32) -> KittyGraphicsImage? {
+        guard var image = kittyGraphicsState.imagesById[imageId] else {
+            return nil
+        }
+        image.lastAccessTick = nextKittyImageAccessTick()
+        kittyGraphicsState.imagesById[imageId] = image
+        return image
+    }
+
+    private func kittyPayloadByteSize(_ payload: KittyGraphicsPayload) -> Int {
+        switch payload {
+        case .png(let data):
+            return data.count
+        case .rgba(let bytes, _, _):
+            return bytes.count
+        }
+    }
+
+    private func nextKittyImageAccessTick() -> UInt64 {
+        let tick = kittyGraphicsState.nextImageAccessTick
+        kittyGraphicsState.nextImageAccessTick &+= 1
+        return tick
+    }
+
+    private func enforceKittyImageCacheLimit() {
+        let limit = clampedKittyImageCacheLimitBytes()
+        guard kittyGraphicsState.totalImageBytes > limit else {
+            return
+        }
+
+        let used = collectUsedKittyImageIds()
+        let unusedIds = kittyGraphicsState.imagesById
+            .filter { !used.contains($0.key) }
+            .sorted { $0.value.lastAccessTick < $1.value.lastAccessTick }
+            .map { $0.key }
+        for id in unusedIds {
+            removeKittyImage(imageId: id)
+            if kittyGraphicsState.totalImageBytes <= limit {
+                return
+            }
+        }
+
+        let oldestIds = kittyGraphicsState.imagesById
+            .sorted { $0.value.lastAccessTick < $1.value.lastAccessTick }
+            .map { $0.key }
+        for id in oldestIds {
+            removeKittyImage(imageId: id)
+            if kittyGraphicsState.totalImageBytes <= limit {
+                return
+            }
+        }
+    }
+
+    private func clampedKittyImageCacheLimitBytes() -> Int {
+        let configured = options.kittyImageCacheLimitBytes
+        if configured <= 0 {
+            return 0
+        }
+        return min(configured, Terminal.kittyMaxImageCacheBytes)
+    }
+
+    private func removeKittyImage(imageId: UInt32) {
+        guard let removed = kittyGraphicsState.imagesById.removeValue(forKey: imageId) else {
+            return
+        }
+        kittyGraphicsState.totalImageBytes = max(0, kittyGraphicsState.totalImageBytes - removed.byteSize)
+        removeKittyImageNumbers(for: imageId)
+    }
+
+    private func removeKittyImageNumbers(for imageId: UInt32) {
+        let numbers = kittyGraphicsState.imageNumbers.filter { $0.value == imageId }.map { $0.key }
+        for number in numbers {
+            kittyGraphicsState.imageNumbers.removeValue(forKey: number)
+        }
+    }
+
+    private func collectUsedKittyImageIds() -> Set<UInt32> {
         var used = Set<UInt32>()
         collectUsedKittyImageIds(from: normalBuffer, into: &used)
         collectUsedKittyImageIds(from: altBuffer, into: &used)
         for record in kittyGraphicsState.placementsByKey.values {
             used.insert(record.imageId)
         }
-        let unusedIds = kittyGraphicsState.imagesById.keys.filter { !used.contains($0) }
-        for id in unusedIds {
-            kittyGraphicsState.imagesById.removeValue(forKey: id)
-        }
-        let unusedNumbers = kittyGraphicsState.imageNumbers.filter { !used.contains($0.value) }.map { $0.key }
-        for number in unusedNumbers {
-            kittyGraphicsState.imageNumbers.removeValue(forKey: number)
-        }
+        return used
     }
 
     private func collectUsedKittyImageIds(from buffer: Buffer, into set: inout Set<UInt32>) {

Sources/SwiftTerm/Terminal.swift

@@ -699,6 +699,7 @@ open class Terminal {
         // when activated.
 
         if clearAlt {
+            clearKittyImages(in: altBuffer, isAlternateBuffer: true)
             altBuffer.clear ()
         }
         buffer = normalBuffer
@@ -716,6 +717,7 @@ open class Terminal {
 
         altBuffer.fillViewportRows(attribute: fillAttr)
         buffer = altBuffer
+        clearKittyImages(in: altBuffer, isAlternateBuffer: true)
     }
 
     func setupTabStops (index: Int = -1)
@@ -4903,6 +4905,7 @@ open class Terminal {
         options.cols = cols
         let savedCursorHidden = cursorHidden
         setup (isReset: true)
+        clearAllKittyImages()
         cursorHidden = savedCursorHidden
         refresh (startRow: 0, endRow: rows-1)
         syncScrollArea ()

Sources/SwiftTerm/TerminalOptions.swift

@@ -58,6 +58,8 @@ public struct TerminalOptions {
     public var tabStopWidth: Int
     /// Whether to report that sixel support is present
     public var enableSixelReported:Bool
+    /// Maximum total bytes to keep for kitty image data; defaults to 320MB and is clamped to 4GB.
+    public var kittyImageCacheLimitBytes: Int
 
     /// Default options
     public static let `default` = TerminalOptions.init(cols: 80,
@@ -68,10 +70,11 @@ public struct TerminalOptions {
                                                        screenReaderMode: false,
                                                        scrollback: 500,
                                                        tabStopWidth: 8,
-                                                       enableSixelReported: true)
+                                                       enableSixelReported: true,
+                                                       kittyImageCacheLimitBytes: 320 * 1024 * 1024)
 
   public init(cols: Int = Self.default.cols, rows: Int = Self.default.rows, convertEol: Bool = Self.default.convertEol, termName: String = Self.default.termName, cursorStyle: CursorStyle = Self.default.cursorStyle, screenReaderMode: Bool = Self.default.screenReaderMode, scrollback: Int = Self.default.scrollback, tabStopWidth: Int = Self.default.tabStopWidth,
-              enableSixelReported: Bool = Self.default.enableSixelReported) {
+              enableSixelReported: Bool = Self.default.enableSixelReported, kittyImageCacheLimitBytes: Int = Self.default.kittyImageCacheLimitBytes) {
         self.cols = cols
         self.rows = rows
         self.convertEol = convertEol
@@ -81,5 +84,6 @@ public struct TerminalOptions {
         self.scrollback = scrollback
         self.tabStopWidth = tabStopWidth
         self.enableSixelReported = enableSixelReported
+        self.kittyImageCacheLimitBytes = kittyImageCacheLimitBytes
     }
 }

Tests/SwiftTermTests/KittyGraphicsLifecycleTests.swift

@@ -0,0 +1,67 @@
+//
+//  KittyGraphicsLifecycleTests.swift
+//
+#if os(macOS)
+import Foundation
+import Testing
+
+@testable import SwiftTerm
+
+final class KittyGraphicsLifecycleTests {
+    private func makeHeadlessTerminal() -> HeadlessTerminal {
+        HeadlessTerminal(queue: SwiftTermTests.queue, options: TerminalOptions(cols: 10, rows: 5)) { _ in }
+    }
+
+    private func sendKitty(terminal: Terminal, control: String, payload: [UInt8]) {
+        let base64 = Data(payload).base64EncodedString()
+        let sequence = "\u{1b}_G\(control);\(base64)\u{1b}\\"
+        terminal.feed(text: sequence)
+    }
+
+    @Test func testKittyImagesClearedOnReset() {
+        let h = makeHeadlessTerminal()
+        let t = h.terminal!
+
+        sendKitty(terminal: t,
+                  control: "a=T,f=24,s=1,v=1,t=d,c=1,r=1,i=1,U=1",
+                  payload: [1, 2, 3])
+
+        #expect(t.kittyGraphicsState.imagesById[1] != nil)
+        #expect(!t.kittyGraphicsState.placementsByKey.isEmpty)
+
+        t.feed(text: "\u{1b}c")
+
+        #expect(t.kittyGraphicsState.imagesById.isEmpty)
+        #expect(t.kittyGraphicsState.imageNumbers.isEmpty)
+        #expect(t.kittyGraphicsState.placementsByKey.isEmpty)
+    }
+
+    @Test func testKittyImagesClearedWhenEnteringAltBuffer() {
+        let h = makeHeadlessTerminal()
+        let t = h.terminal!
+
+        t.feed(text: "\u{1b}[?1049h")
+        #expect(t.isCurrentBufferAlternate)
+
+        sendKitty(terminal: t,
+                  control: "a=T,f=24,s=1,v=1,t=d,c=1,r=1,i=1,U=1",
+                  payload: [1, 2, 3])
+
+        #expect(t.kittyGraphicsState.imagesById[1] != nil)
+        #expect(!t.kittyGraphicsState.placementsByKey.isEmpty)
+
+        t.feed(text: "\u{1b}[?47l")
+        #expect(!t.isCurrentBufferAlternate)
+
+        #expect(t.kittyGraphicsState.imagesById[1] != nil)
+        #expect(!t.kittyGraphicsState.placementsByKey.isEmpty)
+
+        t.feed(text: "\u{1b}[?1049h")
+        #expect(t.isCurrentBufferAlternate)
+
+        #expect(t.kittyGraphicsState.imagesById.isEmpty)
+        #expect(t.kittyGraphicsState.imageNumbers.isEmpty)
+        #expect(t.kittyGraphicsState.placementsByKey.isEmpty)
+    }
+}
+#endif