fix(ui): clear image cache when FilePicker closes to prevent unbounded memory growth (#2158)

M1xA created 3 days ago

* fix(ui): clear image cache when FilePicker closes to prevent unbounded memory growth

* fix(ui): only reset image cache on FilePicker close, not on new session

* fix(ui): reset image cache after Dialog close

* fix(ui): rename image.Reset to image.ResetCache for clarity

Change summary

internal/ui/image/image.go      |  7 +++++
internal/ui/image/image_test.go | 46 +++++++++++++++++++++++++++++++++++
internal/ui/model/ui.go         |  9 ++++++
3 files changed, 62 insertions(+)

Detailed changes

internal/ui/image/image.go

@@ -68,6 +68,13 @@ var (
 	cachedMutex  sync.RWMutex
 )
 
+// ResetCache clears the image cache, freeing all cached decoded images.
+func ResetCache() {
+	cachedMutex.Lock()
+	clear(cachedImages)
+	cachedMutex.Unlock()
+}
+
 // fitImage resizes the image to fit within the specified dimensions in
 // terminal cells, maintaining the aspect ratio.
 func fitImage(id string, img image.Image, cs CellSize, cols, rows int) image.Image {

internal/ui/image/image_test.go

@@ -0,0 +1,46 @@
+package image
+
+import (
+	"image"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestResetCache(t *testing.T) {
+	t.Parallel()
+
+	cachedMutex.Lock()
+	cachedImages[imageKey{id: "a", cols: 10, rows: 10}] = cachedImage{
+		img:  image.NewRGBA(image.Rect(0, 0, 1, 1)),
+		cols: 10,
+		rows: 10,
+	}
+	cachedImages[imageKey{id: "b", cols: 20, rows: 20}] = cachedImage{
+		img:  image.NewRGBA(image.Rect(0, 0, 1, 1)),
+		cols: 20,
+		rows: 20,
+	}
+	cachedMutex.Unlock()
+
+	ResetCache()
+
+	cachedMutex.RLock()
+	length := len(cachedImages)
+	cachedMutex.RUnlock()
+
+	require.Equal(t, 0, length)
+}
+
+func TestResetIdempotent(t *testing.T) {
+	t.Parallel()
+
+	// Calling Reset on an empty cache should not panic.
+	ResetCache()
+
+	cachedMutex.RLock()
+	length := len(cachedImages)
+	cachedMutex.RUnlock()
+
+	require.Equal(t, 0, length)
+}

internal/ui/model/ui.go

@@ -41,6 +41,7 @@ import (
 	"github.com/charmbracelet/crush/internal/ui/common"
 	"github.com/charmbracelet/crush/internal/ui/completions"
 	"github.com/charmbracelet/crush/internal/ui/dialog"
+	fimage "github.com/charmbracelet/crush/internal/ui/image"
 	"github.com/charmbracelet/crush/internal/ui/logo"
 	"github.com/charmbracelet/crush/internal/ui/styles"
 	"github.com/charmbracelet/crush/internal/ui/util"
@@ -1137,6 +1138,10 @@ func (m *UI) handleDialogMsg(msg tea.Msg) tea.Cmd {
 			break
 		}
 
+		if m.dialog.ContainsDialog(dialog.FilePickerID) {
+			defer fimage.ResetCache()
+		}
+
 		m.dialog.CloseFrontDialog()
 
 		if isOnboarding {
@@ -1351,6 +1356,10 @@ func (m *UI) handleDialogMsg(msg tea.Msg) tea.Cmd {
 				m.dialog.CloseDialog(dialog.FilePickerID)
 				return nil
 			},
+			func() tea.Msg {
+				fimage.ResetCache()
+				return nil
+			},
 		))
 
 	case dialog.ActionRunCustomCommand: