perf(chat): cache the prefixed render of chat messages

Christian Rocha and Charm Crush created

Avoid rebuilding the focus and selection prefix for every line of every
message on every frame. Each item now keeps the prefixed output cached
and only rebuilds it when something actually changes.

Co-Authored-By: Charm Crush <crush@charm.land>

Change summary

internal/ui/chat/assistant.go         |  22 +++
internal/ui/chat/messages.go          |  44 ++++++
internal/ui/chat/prefix_cache_test.go | 196 +++++++++++++++++++++++++++++
internal/ui/chat/tools.go             |  24 +++
internal/ui/chat/user.go              |  20 ++
5 files changed, 302 insertions(+), 4 deletions(-)

Detailed changes

internal/ui/chat/assistant.go 🔗

@@ -114,6 +114,22 @@ func (a *AssistantMessageItem) Render(width int) string {
 	// it's wrapping logic.
 	// We already know that the content is wrapped to the correct width in
 	// RawRender, so we can just apply the styles directly to each line.
+	//
+	// The split + per-line prefix loop is O(L); cache the result keyed
+	// by (width, focused) so steady-state Render becomes a pointer
+	// return. Bypass the cache while spinning (RawRender's spinner
+	// suffix changes every animation frame) or while a highlight range
+	// is active (selection drag).
+	useCache := !a.isSpinning() && !a.isHighlighted()
+	var key uint64
+	if a.focused {
+		key = 1
+	}
+	if useCache {
+		if cached, ok := a.getCachedPrefixedRender(width, key); ok {
+			return cached
+		}
+	}
 	focused := a.sty.Messages.AssistantFocused.Render()
 	blurred := a.sty.Messages.AssistantBlurred.Render()
 	rendered := a.RawRender(width)
@@ -125,7 +141,11 @@ func (a *AssistantMessageItem) Render(width int) string {
 			lines[i] = blurred + line
 		}
 	}
-	return strings.Join(lines, "\n")
+	out := strings.Join(lines, "\n")
+	if useCache {
+		a.setCachedPrefixedRender(out, width, key)
+	}
+	return out
 }
 
 // renderMessageContent renders the message content including thinking, main content, and finish reason.

internal/ui/chat/messages.go 🔗

@@ -156,6 +156,20 @@ type cachedMessageItem struct {
 	// width and height are the dimensions of the cached render
 	width  int
 	height int
+
+	// prefixedRendered caches the per-line-prefixed Render output (the
+	// result of splitting RawRender by newlines and prepending a focus
+	// or selection prefix to every line). Items rebuild this every
+	// frame today; caching it keyed by (prefixedWidth, prefixedKey)
+	// turns Render into a pointer return when item state is stable.
+	//
+	// Invalidation lives in clearCache; callers must additionally
+	// bypass this cache whenever the prefixed output would not be
+	// stable (spinner ticks, active highlight ranges) by not calling
+	// setCachedPrefixedRender for those frames.
+	prefixedRendered string
+	prefixedWidth    int
+	prefixedKey      uint64
 }
 
 // getCachedRender returns the cached render if it exists for the given width.
@@ -173,11 +187,31 @@ func (c *cachedMessageItem) setCachedRender(rendered string, width, height int)
 	c.height = height
 }
 
+// getCachedPrefixedRender returns the cached prefixed render if it exists
+// for the given (width, key). The key encodes any state that changes the
+// per-line prefix (focused/blurred, compact, ...).
+func (c *cachedMessageItem) getCachedPrefixedRender(width int, key uint64) (string, bool) {
+	if c.prefixedRendered != "" && c.prefixedWidth == width && c.prefixedKey == key {
+		return c.prefixedRendered, true
+	}
+	return "", false
+}
+
+// setCachedPrefixedRender stores the cached prefixed render.
+func (c *cachedMessageItem) setCachedPrefixedRender(rendered string, width int, key uint64) {
+	c.prefixedRendered = rendered
+	c.prefixedWidth = width
+	c.prefixedKey = key
+}
+
 // clearCache clears the cached render.
 func (c *cachedMessageItem) clearCache() {
 	c.rendered = ""
 	c.width = 0
 	c.height = 0
+	c.prefixedRendered = ""
+	c.prefixedWidth = 0
+	c.prefixedKey = 0
 }
 
 // focusableMessageItem is a base struct for message items that can be focused.
@@ -237,12 +271,20 @@ func (a *AssistantInfoItem) RawRender(width int) string {
 
 // Render implements MessageItem.
 func (a *AssistantInfoItem) Render(width int) string {
+	// AssistantInfoItem uses a single, state-independent prefix; key 0
+	// is sufficient. The cache is invalidated whenever the underlying
+	// cachedMessageItem render is cleared.
+	if cached, ok := a.getCachedPrefixedRender(width, 0); ok {
+		return cached
+	}
 	prefix := a.sty.Messages.SectionHeader.Render()
 	lines := strings.Split(a.RawRender(width), "\n")
 	for i, line := range lines {
 		lines[i] = prefix + line
 	}
-	return strings.Join(lines, "\n")
+	out := strings.Join(lines, "\n")
+	a.setCachedPrefixedRender(out, width, 0)
+	return out
 }
 
 func (a *AssistantInfoItem) renderContent(width int) string {

internal/ui/chat/prefix_cache_test.go 🔗

@@ -0,0 +1,196 @@
+package chat
+
+import (
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/charmbracelet/crush/internal/message"
+	"github.com/charmbracelet/crush/internal/ui/attachments"
+	"github.com/charmbracelet/crush/internal/ui/styles"
+	"github.com/stretchr/testify/require"
+)
+
+// finishedAssistantMessage builds an assistant message with text content and a
+// finish part so AssistantMessageItem.isSpinning returns false and the
+// prefix cache is exercised.
+func finishedAssistantMessage(id, text string) *message.Message {
+	return &message.Message{
+		ID:   id,
+		Role: message.Assistant,
+		Parts: []message.ContentPart{
+			message.TextContent{Text: text},
+			message.Finish{Reason: message.FinishReasonEndTurn, Time: time.Now().Unix()},
+		},
+	}
+}
+
+// TestAssistantMessageItemRender_PrefixCacheFocusBlur covers the F3 invariant
+// that focus → blur → focus produces the correct prefix every time and never
+// leaks the previous focus state out of the cache.
+func TestAssistantMessageItemRender_PrefixCacheFocusBlur(t *testing.T) {
+	t.Parallel()
+
+	sty := styles.CharmtonePantera()
+	msg := finishedAssistantMessage("m1", "Hello world from the cache test.")
+	item := NewAssistantMessageItem(&sty, msg).(*AssistantMessageItem)
+
+	const width = 60
+
+	item.SetFocused(true)
+	focused1 := item.Render(width)
+	focused2 := item.Render(width)
+	require.Equal(t, focused1, focused2, "second render must hit the cache and match the first")
+
+	item.SetFocused(false)
+	blurred1 := item.Render(width)
+	require.NotEqual(t, focused1, blurred1, "blur must produce a different prefixed render than focus")
+
+	item.SetFocused(true)
+	focused3 := item.Render(width)
+	require.Equal(t, focused1, focused3, "re-focus must produce identical output to the original focused render")
+}
+
+// TestAssistantMessageItemRender_PrefixCacheWidthInvalidates asserts that a
+// width change does not return the previous width's cached output.
+func TestAssistantMessageItemRender_PrefixCacheWidthInvalidates(t *testing.T) {
+	t.Parallel()
+
+	sty := styles.CharmtonePantera()
+	msg := finishedAssistantMessage("m2", "Some content that wraps differently at different widths so the rendered output diverges.")
+	item := NewAssistantMessageItem(&sty, msg).(*AssistantMessageItem)
+	item.SetFocused(true)
+
+	narrow := item.Render(40)
+	wide := item.Render(100)
+	require.NotEqual(t, narrow, wide, "different widths must produce different rendered output")
+
+	narrowAgain := item.Render(40)
+	require.Equal(t, narrow, narrowAgain, "returning to the original width must hit (or repopulate) the cache with the same output")
+}
+
+// TestAssistantMessageItemRender_PrefixCacheHighlightOnTop guarantees that
+// activating a highlight range bypasses the prefix cache so selection drags
+// reflect immediately, and that clearing the highlight returns to the cached
+// prefixed output unchanged.
+func TestAssistantMessageItemRender_PrefixCacheHighlightOnTop(t *testing.T) {
+	t.Parallel()
+
+	sty := styles.CharmtonePantera()
+	msg := finishedAssistantMessage("m3", "Hello world from the highlight test.")
+	item := NewAssistantMessageItem(&sty, msg).(*AssistantMessageItem)
+	item.SetFocused(true)
+
+	const width = 60
+	plain := item.Render(width)
+
+	// Activating a highlight must change the rendered output (selection
+	// painted on top) without poisoning the cache for the un-highlighted
+	// state that follows.
+	item.SetHighlight(0, 0, 0, 5)
+	highlighted := item.Render(width)
+	require.NotEqual(t, plain, highlighted, "active highlight must change Render output")
+
+	// Clear the highlight; the cached un-highlighted prefix render must
+	// be returned unchanged.
+	item.SetHighlight(-1, -1, -1, -1)
+	plainAfter := item.Render(width)
+	require.Equal(t, plain, plainAfter, "clearing the highlight must restore the cached prefixed output exactly")
+}
+
+// TestUserMessageItemRender_PrefixCacheFocusBlur is the user-message
+// counterpart of the assistant focus/blur cache test.
+func TestUserMessageItemRender_PrefixCacheFocusBlur(t *testing.T) {
+	t.Parallel()
+
+	sty := styles.CharmtonePantera()
+	msg := &message.Message{
+		ID:   "u1",
+		Role: message.User,
+		Parts: []message.ContentPart{
+			message.TextContent{Text: "Hello from the user."},
+		},
+	}
+	r := attachments.NewRenderer(
+		sty.Attachments.Normal,
+		sty.Attachments.Deleting,
+		sty.Attachments.Image,
+		sty.Attachments.Text,
+	)
+	item := NewUserMessageItem(&sty, msg, r).(*UserMessageItem)
+
+	const width = 60
+
+	item.SetFocused(true)
+	focused1 := item.Render(width)
+	focused2 := item.Render(width)
+	require.Equal(t, focused1, focused2)
+
+	item.SetFocused(false)
+	blurred := item.Render(width)
+	require.NotEqual(t, focused1, blurred)
+
+	item.SetFocused(true)
+	focused3 := item.Render(width)
+	require.Equal(t, focused1, focused3)
+}
+
+// TestCachedMessageItem_PrefixCacheSemantics covers the constant-prefix
+// path used by AssistantInfoItem and the (width, key) keying used by every
+// item, against the underlying cachedMessageItem helper directly. This
+// avoids constructing a full *config.Config with an initialized provider
+// map just to exercise cache plumbing that is identical for all callers.
+func TestCachedMessageItem_PrefixCacheSemantics(t *testing.T) {
+	t.Parallel()
+
+	c := &cachedMessageItem{}
+
+	// Empty cache: miss.
+	_, ok := c.getCachedPrefixedRender(80, 0)
+	require.False(t, ok)
+
+	// Set then hit at the same (width, key).
+	c.setCachedPrefixedRender("hello", 80, 0)
+	got, ok := c.getCachedPrefixedRender(80, 0)
+	require.True(t, ok)
+	require.Equal(t, "hello", got)
+
+	// Different width: miss.
+	_, ok = c.getCachedPrefixedRender(120, 0)
+	require.False(t, ok)
+
+	// Different key (focused vs blurred): miss.
+	_, ok = c.getCachedPrefixedRender(80, 1)
+	require.False(t, ok)
+
+	// clearCache drops the prefixed cache too.
+	c.setCachedRender("raw", 80, 1)
+	c.setCachedPrefixedRender("hello", 80, 0)
+	c.clearCache()
+	_, ok = c.getCachedPrefixedRender(80, 0)
+	require.False(t, ok, "clearCache must drop the prefixed render cache")
+	_, _, ok = c.getCachedRender(80)
+	require.False(t, ok, "clearCache must also drop the raw render cache")
+}
+
+// TestAssistantMessageItemRender_PrefixCacheNoCacheLeak guards against a
+// regression where the cache returned the prefixed output of the previous
+// width. We verify that the cached output for width=W contains the W-sized
+// prefix and not a stale wider one by checking that line lengths are
+// consistent on cache hit.
+func TestAssistantMessageItemRender_PrefixCacheNoCacheLeak(t *testing.T) {
+	t.Parallel()
+
+	sty := styles.CharmtonePantera()
+	msg := finishedAssistantMessage("m4", strings.Repeat("word ", 40))
+	item := NewAssistantMessageItem(&sty, msg).(*AssistantMessageItem)
+	item.SetFocused(true)
+
+	out80 := item.Render(80)
+	out120 := item.Render(120)
+	require.NotEqual(t, out80, out120)
+
+	// Hit each cached entry again and confirm stability.
+	require.Equal(t, out80, item.Render(80))
+	require.Equal(t, out120, item.Render(120))
+}

internal/ui/chat/tools.go 🔗

@@ -332,6 +332,24 @@ func (t *baseToolMessageItem) RawRender(width int) string {
 
 // Render renders the tool message item at the given width.
 func (t *baseToolMessageItem) Render(width int) string {
+	// Cache the prefixed output keyed by (width, prefix variant).
+	// Bypass the cache while spinning (RawRender output is
+	// frame-dependent) or while a highlight range is active.
+	useCache := !t.isSpinning() && !t.isHighlighted()
+	var key uint64
+	switch {
+	case t.isCompact:
+		key = 2
+	case t.focused:
+		key = 1
+	default:
+		key = 0
+	}
+	if useCache {
+		if cached, ok := t.getCachedPrefixedRender(width, key); ok {
+			return cached
+		}
+	}
 	var prefix string
 	if t.isCompact {
 		prefix = t.sty.Messages.ToolCallCompact.Render()
@@ -344,7 +362,11 @@ func (t *baseToolMessageItem) Render(width int) string {
 	for i, ln := range lines {
 		lines[i] = prefix + ln
 	}
-	return strings.Join(lines, "\n")
+	out := strings.Join(lines, "\n")
+	if useCache {
+		t.setCachedPrefixedRender(out, width, key)
+	}
+	return out
 }
 
 // ToolCall returns the tool call associated with this message item.

internal/ui/chat/user.go 🔗

@@ -70,6 +70,20 @@ func (m *UserMessageItem) RawRender(width int) string {
 
 // Render implements MessageItem.
 func (m *UserMessageItem) Render(width int) string {
+	// Bypass the prefix cache while a highlight range is active so
+	// selection drags reflect immediately without invalidating the
+	// cache. Highlight changes are intentionally applied "above" the
+	// prefix cache.
+	useCache := !m.isHighlighted()
+	var key uint64
+	if m.focused {
+		key = 1
+	}
+	if useCache {
+		if cached, ok := m.getCachedPrefixedRender(width, key); ok {
+			return cached
+		}
+	}
 	var prefix string
 	if m.focused {
 		prefix = m.sty.Messages.UserFocused.Render()
@@ -80,7 +94,11 @@ func (m *UserMessageItem) Render(width int) string {
 	for i, line := range lines {
 		lines[i] = prefix + line
 	}
-	return strings.Join(lines, "\n")
+	out := strings.Join(lines, "\n")
+	if useCache {
+		m.setCachedPrefixedRender(out, width, key)
+	}
+	return out
 }
 
 // ID implements MessageItem.