From 3f9e8cca2a0dd2db3cb23a6f61dfe29d6f0816d4 Mon Sep 17 00:00:00 2001 From: Cian Johnston Date: Thu, 13 Aug 2026 08:42:23 +0100 Subject: [PATCH] chore: add test coverage for chatd compaction (#28053) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Adds test coverage for the three compaction-decision functions in chatd that had zero tests: `latestPromptUsage`, `shouldCompactPromptUsage`, and `contextTokensFromUsage`. AIGOV-585 hypothesized that chatd's token counting logic was incorrect — that it compared a cumulative sum of prompt tokens across all agentic-loop steps against the context window. The tests disprove this: `latestPromptUsage` returns the last persisted assistant message's usage, not a sum. The actual bug was in the aibridge streaming interceptor, which summed usage across SSE chunks and persisted inflated values (fixed in `ad100452d4`). ## What's tested - `TestLatestPromptUsage` — pins that the compaction path reads the last step's usage (5,400), not a cumulative sum across steps (15,600). If someone wires `TotalUsage` into the compaction path as the issue suggested, this fails. - `TestShouldCompactPromptUsage` — covers the threshold decision with the inflated value from the issue (417,012 → compacts), the correct value (6,000 → doesn't compact), cache token counting, and both disable guards (threshold=100, contextLimit=0).
Plan / investigation notes - Traced the full flow: `chatloop.go:993` sets `result.usage = part.Usage` from the per-step `StreamPartTypeFinish` event, not the accumulated `TotalUsage` from `agent.go:544`. chatd never calls fantasy's `Agent` interface. - The `TotalUsage` accumulation in `agent.go:544` is only used for cost attribution, not context occupancy. - Commit `ad100452d4` fixed the real bug in `aibridge/intercept/chatcompletions/streaming.go` (cross-chunk usage summation for vLLM-style backends). - Tests reuse existing `dbMessage` and `withUsage` helpers from `message_conversion_test.go` (same package).
Generated by [Coder Agents](https://coder.com) --------- Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- .../generation_preparer_internal_test.go | 98 +++++++++++++++++++ 1 file changed, 98 insertions(+) diff --git a/coderd/x/chatd/generation_preparer_internal_test.go b/coderd/x/chatd/generation_preparer_internal_test.go index 38a39cbd2e..83fd496739 100644 --- a/coderd/x/chatd/generation_preparer_internal_test.go +++ b/coderd/x/chatd/generation_preparer_internal_test.go @@ -9,6 +9,7 @@ import ( fantasyopenai "charm.land/fantasy/providers/openai" "github.com/google/uuid" "github.com/sqlc-dev/pqtype" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "cdr.dev/slog/v3/sloggers/slogtest" @@ -523,3 +524,100 @@ func TestDeriveFinalTurnRunResult(t *testing.T) { require.Empty(t, result.FallbackModel) }) } + +// TestLatestPromptUsage verifies the compaction path reads the last +// persisted assistant usage, not a cumulative sum across steps. +// This is the chatd side of AIGOV-585: the issue blamed chatd for +// summing usage across steps, but latestPromptUsage returns the +// single most-recent non-zero usage. The real bug was in the +// aibridge streaming interceptor (fixed in ad100452d4). +func TestLatestPromptUsage(t *testing.T) { + t.Parallel() + + t.Run("returns last assistant step not a sum", func(t *testing.T) { + t.Parallel() + // Three-step agentic turn: each step re-sends the whole + // conversation, so InputTokens are 5000, 5200, 5400. + // A sum would be 15600; the correct answer is 5400. + messages := []database.ChatMessage{ + dbMessage(t, 1, database.ChatMessageRoleUser, false, codersdk.ChatMessageText("hi")), + withUsage(dbMessage(t, 2, database.ChatMessageRoleAssistant, false, codersdk.ChatMessageText("step1")), 5000, 0), + dbMessage(t, 3, database.ChatMessageRoleTool, false, codersdk.ChatMessageText("result")), + withUsage(dbMessage(t, 4, database.ChatMessageRoleAssistant, false, codersdk.ChatMessageText("step2")), 5200, 0), + dbMessage(t, 5, database.ChatMessageRoleTool, false, codersdk.ChatMessageText("result")), + withUsage(dbMessage(t, 6, database.ChatMessageRoleAssistant, false, codersdk.ChatMessageText("step3")), 5400, 0), + } + usage := latestPromptUsage(messages) + assert.Equal(t, int64(5400), usage.InputTokens, + "must return the last step's usage, not the sum across steps") + }) + + t.Run("returns zero when no messages have usage", func(t *testing.T) { + t.Parallel() + messages := []database.ChatMessage{ + dbMessage(t, 1, database.ChatMessageRoleUser, false, codersdk.ChatMessageText("hi")), + dbMessage(t, 2, database.ChatMessageRoleAssistant, false, codersdk.ChatMessageText("hi")), + } + assert.Equal(t, fantasy.Usage{}, latestPromptUsage(messages)) + }) + + t.Run("returns zero for empty message list", func(t *testing.T) { + t.Parallel() + assert.Equal(t, fantasy.Usage{}, latestPromptUsage(nil)) + }) +} + +// TestShouldCompactPromptUsage verifies the compaction threshold decision +// is correct for both the inflated values the aibridge bug produced and +// accurate per-step values. +func TestShouldCompactPromptUsage(t *testing.T) { + t.Parallel() + + const contextLimit = int64(262144) // 256K, as in the poolside report + + t.Run("inflated cumulative usage triggers compaction", func(t *testing.T) { + t.Parallel() + // 417,012 tokens: what the aibridge cross-chunk sum bug + // produced for a ~6,000-token conversation. + assert.True(t, shouldCompactPromptUsage( + fantasy.Usage{InputTokens: 417012, TotalTokens: 418846}, + contextLimit, 80)) + }) + + t.Run("correct per-step usage does not trigger", func(t *testing.T) { + t.Parallel() + assert.False(t, shouldCompactPromptUsage( + fantasy.Usage{InputTokens: 6000, TotalTokens: 6030}, + contextLimit, 80)) + }) + + t.Run("threshold 100 disables compaction", func(t *testing.T) { + t.Parallel() + assert.False(t, shouldCompactPromptUsage( + fantasy.Usage{InputTokens: 500000}, contextLimit, 100)) + }) + + t.Run("zero context limit disables compaction", func(t *testing.T) { + t.Parallel() + assert.False(t, shouldCompactPromptUsage( + fantasy.Usage{InputTokens: 6000}, 0, 80)) + }) + + t.Run("counts cache read and creation tokens", func(t *testing.T) { + t.Parallel() + usage := fantasy.Usage{ + InputTokens: 6000, + CacheReadTokens: 200000, + CacheCreationTokens: 5000, + } + // 211,000 / 262,144 = ~80.5% + assert.True(t, shouldCompactPromptUsage(usage, contextLimit, 80)) + }) + + t.Run("falls back to TotalTokens when granular fields are missing", func(t *testing.T) { + t.Parallel() + assert.True(t, shouldCompactPromptUsage( + fantasy.Usage{TotalTokens: 211000}, + contextLimit, 80)) + }) +}