From 771117cafeedffc8a9b5859503fc055bcd525d3e Mon Sep 17 00:00:00 2001 From: Christian Rocha Date: Tue, 12 May 2026 09:39:46 -0400 Subject: [PATCH] perf(chat): cache the prefixed render of chat messages 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 --- 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(-) create mode 100644 internal/ui/chat/prefix_cache_test.go diff --git a/internal/ui/chat/assistant.go b/internal/ui/chat/assistant.go index f49cd1d6ceca12b15f705640400da6941972342e..2977e7000c2945250bb327c1c79dc152b8c51f92 100644 --- a/internal/ui/chat/assistant.go +++ b/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. diff --git a/internal/ui/chat/messages.go b/internal/ui/chat/messages.go index 6c86cb74951e0bf8f20fe7af4fa420b4527936ca..3c1c5a13b7ba13a5e6990b4cc546c8813c0d5252 100644 --- a/internal/ui/chat/messages.go +++ b/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 { diff --git a/internal/ui/chat/prefix_cache_test.go b/internal/ui/chat/prefix_cache_test.go new file mode 100644 index 0000000000000000000000000000000000000000..443f5a68dab73d8e518b37d867fc9aa8c758ca25 --- /dev/null +++ b/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)) +} diff --git a/internal/ui/chat/tools.go b/internal/ui/chat/tools.go index 4715cfafa6a25f01b93ee56caf41aa55aee3196e..560fd2fc5312852b5cc91a228b4587f4accb0d2a 100644 --- a/internal/ui/chat/tools.go +++ b/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. diff --git a/internal/ui/chat/user.go b/internal/ui/chat/user.go index b1160d2a1531cb6c2915e8df67cb1962079620e0..8194955a6fa99e20a7d949651d6d1400d8f3fb72 100644 --- a/internal/ui/chat/user.go +++ b/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.