From f417732279e8ca1a284e8c2aed0096bc7a9e85a7 Mon Sep 17 00:00:00 2001 From: ParthSareen Date: Thu, 11 Jun 2026 21:02:40 -0700 Subject: [PATCH] agent: bound markdown and compaction caches --- agent/compactor.go | 22 +++++++++++++++++- agent/compactor_test.go | 44 +++++++++++++++++++++++++++++++++++ cmd/tui/chat_markdown.go | 41 +++++++++++++++++++++++++++----- cmd/tui/chat_markdown_test.go | 17 ++++++++++++++ 4 files changed, 117 insertions(+), 7 deletions(-) diff --git a/agent/compactor.go b/agent/compactor.go index ea09df6d6..9ebddacde 100644 --- a/agent/compactor.go +++ b/agent/compactor.go @@ -15,6 +15,8 @@ const ( defaultCompactionThreshold = 0.8 compactionSummaryMessagePrefix = "Conversation summary:\n" + maxCompactionSummaryBytes = 16 * 1024 + compactionSummaryTruncated = "\n\n[summary truncated]" ) type Compactor interface { @@ -91,7 +93,7 @@ func (c *SimpleCompactor) MaybeCompact(ctx context.Context, req CompactionReques result.Reason = err.Error() return result, err } - summary = strings.TrimSpace(summary) + summary = truncateCompactionSummary(strings.TrimSpace(summary)) if summary == "" { result.Reason = "summary was empty" return result, nil @@ -201,6 +203,24 @@ func (c *SimpleCompactor) summarize(ctx context.Context, req CompactionRequest, return summary.String(), nil } +func truncateCompactionSummary(summary string) string { + if len(summary) <= maxCompactionSummaryBytes { + return summary + } + limit := maxCompactionSummaryBytes - len(compactionSummaryTruncated) + if limit < 0 { + limit = 0 + } + var b strings.Builder + for _, r := range summary { + if b.Len()+len(string(r)) > limit { + break + } + b.WriteRune(r) + } + return strings.TrimSpace(b.String()) + compactionSummaryTruncated +} + func estimateCompactionTokens(text string) int { text = strings.TrimSpace(text) if text == "" { diff --git a/agent/compactor_test.go b/agent/compactor_test.go index 77840bd40..e8960844c 100644 --- a/agent/compactor_test.go +++ b/agent/compactor_test.go @@ -85,6 +85,50 @@ func TestSimpleCompactorSummarizesOldMessages(t *testing.T) { } } +func TestSimpleCompactorTruncatesOversizedSummary(t *testing.T) { + longSummary := strings.Repeat("x", maxCompactionSummaryBytes+1024) + client := &fakeClient{ + responses: [][]api.ChatResponse{{ + {Message: api.Message{Role: "assistant", Content: longSummary}}, + }}, + } + store := &compactionStore{} + compactor := NewSimpleCompactor(client, store, CompactionOptions{ + ContextWindowTokens: 100, + KeepUserTurns: 1, + Threshold: 0.5, + }) + + result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ + ChatID: "chat-1", + Model: "model", + Messages: []api.Message{ + {Role: "user", Content: "old one"}, + {Role: "assistant", Content: "old answer"}, + {Role: "user", Content: "recent one"}, + }, + Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 75}}, + }) + if err != nil { + t.Fatal(err) + } + if !result.Compacted { + t.Fatal("expected compaction") + } + if len(result.Summary) > maxCompactionSummaryBytes { + t.Fatalf("summary bytes = %d, want <= %d", len(result.Summary), maxCompactionSummaryBytes) + } + if !strings.HasSuffix(result.Summary, compactionSummaryTruncated) { + t.Fatalf("summary missing truncation marker") + } + if store.summary != result.Summary { + t.Fatalf("stored summary mismatch") + } + if !strings.Contains(result.Messages[0].Content, compactionSummaryTruncated) { + t.Fatalf("compacted message missing truncation marker: %#v", result.Messages) + } +} + func TestSimpleCompactorSkipsBelowThreshold(t *testing.T) { client := &fakeClient{} compactor := NewSimpleCompactor(client, nil, CompactionOptions{ diff --git a/cmd/tui/chat_markdown.go b/cmd/tui/chat_markdown.go index 75614f663..c9c02b56a 100644 --- a/cmd/tui/chat_markdown.go +++ b/cmd/tui/chat_markdown.go @@ -17,7 +17,19 @@ var markdownLinkPattern = regexp.MustCompile(`\[([^\]]+)\]\((https?://[^)\s]+)\) var markdownTableSeparatorPattern = regexp.MustCompile(`^:?-{3,}:?$`) -var chatMarkdownRenderers sync.Map +const maxMarkdownRendererCacheEntries = 8 + +var chatMarkdownRenderers = newMarkdownRendererCache() + +type markdownRendererCache struct { + mu sync.Mutex + renderers map[int]*cachedMarkdownRenderer + order []int +} + +func newMarkdownRendererCache() *markdownRendererCache { + return &markdownRendererCache{renderers: make(map[int]*cachedMarkdownRenderer)} +} type cachedMarkdownRenderer struct { renderer *glamour.TermRenderer @@ -461,10 +473,15 @@ func widestColumn(widths []int) int { } func markdownRendererForWidth(width int) (*cachedMarkdownRenderer, error) { - if cached, ok := chatMarkdownRenderers.Load(width); ok { - return cached.(*cachedMarkdownRenderer), nil - } + return chatMarkdownRenderers.renderer(width) +} +func (c *markdownRendererCache) renderer(width int) (*cachedMarkdownRenderer, error) { + c.mu.Lock() + defer c.mu.Unlock() + if cached, ok := c.renderers[width]; ok { + return cached, nil + } renderer, err := glamour.NewTermRenderer( glamour.WithStyles(compactMarkdownStyle()), glamour.WithWordWrap(width), @@ -475,8 +492,20 @@ func markdownRendererForWidth(width int) (*cachedMarkdownRenderer, error) { return nil, err } cached := &cachedMarkdownRenderer{renderer: renderer} - actual, _ := chatMarkdownRenderers.LoadOrStore(width, cached) - return actual.(*cachedMarkdownRenderer), nil + if len(c.order) >= maxMarkdownRendererCacheEntries { + evict := c.order[0] + c.order = c.order[1:] + delete(c.renderers, evict) + } + c.renderers[width] = cached + c.order = append(c.order, width) + return cached, nil +} + +func (c *markdownRendererCache) len() int { + c.mu.Lock() + defer c.mu.Unlock() + return len(c.renderers) } func renderMarkdownDiffFences(markdown string, width int) (string, bool) { diff --git a/cmd/tui/chat_markdown_test.go b/cmd/tui/chat_markdown_test.go index 4a13374d0..ff14bdc45 100644 --- a/cmd/tui/chat_markdown_test.go +++ b/cmd/tui/chat_markdown_test.go @@ -21,6 +21,23 @@ func TestChatMarkdownRendersAssistantAndSystemOutput(t *testing.T) { } } +func TestMarkdownRendererCacheIsBounded(t *testing.T) { + previous := chatMarkdownRenderers + chatMarkdownRenderers = newMarkdownRendererCache() + defer func() { + chatMarkdownRenderers = previous + }() + + for i := 0; i < maxMarkdownRendererCacheEntries+4; i++ { + if _, err := markdownRendererForWidth(40 + i); err != nil { + t.Fatal(err) + } + } + if got := chatMarkdownRenderers.len(); got != maxMarkdownRendererCacheEntries { + t.Fatalf("cache size = %d, want %d", got, maxMarkdownRendererCacheEntries) + } +} + func TestChatMarkdownExposesLinks(t *testing.T) { m := chatModel{width: 100, height: 30} m.entries = []chatEntry{{