Preserve context estimate after compaction
- commit
- cad135d4231f26abb3c964a893ea77634dbf7953
- author
- patriceckhart · · 2026-08-27 17:40 UTC
- parent
- 8e8b1da
- signature
- unsigned
- refs
- tag: v0.3.52
2 files · +55 -1
M packages/agent/modes/compact_queue_test.go +48 -0
| @@ -94,6 +94,54 @@ func TestPromptSubmittedDuringCompactionStartsFollowUpTurn(t *testing.T) { | |||
| 94 | 94 | } | |
| 95 | 95 | } | |
| 96 | 96 | ||
| 97 | + | func TestCompactionRetainsContextEstimate(t *testing.T) { | |
| 98 | + | client := &compactQueueClient{ | |
| 99 | + | compactionStarted: make(chan struct{}), | |
| 100 | + | releaseCompaction: make(chan struct{}), | |
| 101 | + | followUpRequest: make(chan provider.Request, 1), | |
| 102 | + | } | |
| 103 | + | agent := core.NewAgent(client, "test-model", "", nil) | |
| 104 | + | agent.SetMessages([]provider.Message{ | |
| 105 | + | {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "one"}}}, | |
| 106 | + | {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "two"}}}, | |
| 107 | + | {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "three"}}}, | |
| 108 | + | {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "four"}}}, | |
| 109 | + | {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "five"}}}, | |
| 110 | + | }) | |
| 111 | + | interactive := NewInteractive(InteractiveConfig{Agent: agent}) | |
| 112 | + | ||
| 113 | + | interactive.runCompact(context.Background(), false) | |
| 114 | + | select { | |
| 115 | + | case <-client.compactionStarted: | |
| 116 | + | case <-time.After(2 * time.Second): | |
| 117 | + | t.Fatal("compaction did not start") | |
| 118 | + | } | |
| 119 | + | close(client.releaseCompaction) | |
| 120 | + | ||
| 121 | + | deadline := time.NewTimer(2 * time.Second) | |
| 122 | + | defer deadline.Stop() | |
| 123 | + | poll := time.NewTicker(time.Millisecond) | |
| 124 | + | defer poll.Stop() | |
| 125 | + | for { | |
| 126 | + | interactive.mu.Lock() | |
| 127 | + | busy := interactive.busy | |
| 128 | + | got := interactive.lastCtxInput | |
| 129 | + | interactive.mu.Unlock() | |
| 130 | + | if !busy { | |
| 131 | + | want := estimateTimelineMessageTokens(agent.Messages()) | |
| 132 | + | if got != want || got == 0 { | |
| 133 | + | t.Fatalf("context estimate after compaction = %d, want %d", got, want) | |
| 134 | + | } | |
| 135 | + | break | |
| 136 | + | } | |
| 137 | + | select { | |
| 138 | + | case <-deadline.C: | |
| 139 | + | t.Fatal("compaction did not finish") | |
| 140 | + | case <-poll.C: | |
| 141 | + | } | |
| 142 | + | } | |
| 143 | + | } | |
| 144 | + | ||
| 97 | 145 | func TestPreTurnCompactionPreservesPromptImages(t *testing.T) { | |
| 98 | 146 | client := &compactQueueClient{ | |
| 99 | 147 | compactionStarted: make(chan struct{}), | |
| @@ -94,6 +94,54 @@ func TestPromptSubmittedDuringCompactionStartsFollowUpTurn(t *testing.T) { | |||||
| 94 | } | 94 | } | ||
| 95 | } | 95 | } | ||
| 96 | 96 | ||||
| 97 | + | func TestCompactionRetainsContextEstimate(t *testing.T) { | |||
| 98 | + | client := &compactQueueClient{ | |||
| 99 | + | compactionStarted: make(chan struct{}), | |||
| 100 | + | releaseCompaction: make(chan struct{}), | |||
| 101 | + | followUpRequest: make(chan provider.Request, 1), | |||
| 102 | + | } | |||
| 103 | + | agent := core.NewAgent(client, "test-model", "", nil) | |||
| 104 | + | agent.SetMessages([]provider.Message{ | |||
| 105 | + | {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "one"}}}, | |||
| 106 | + | {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "two"}}}, | |||
| 107 | + | {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "three"}}}, | |||
| 108 | + | {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "four"}}}, | |||
| 109 | + | {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "five"}}}, | |||
| 110 | + | }) | |||
| 111 | + | interactive := NewInteractive(InteractiveConfig{Agent: agent}) | |||
| 112 | + | ||||
| 113 | + | interactive.runCompact(context.Background(), false) | |||
| 114 | + | select { | |||
| 115 | + | case <-client.compactionStarted: | |||
| 116 | + | case <-time.After(2 * time.Second): | |||
| 117 | + | t.Fatal("compaction did not start") | |||
| 118 | + | } | |||
| 119 | + | close(client.releaseCompaction) | |||
| 120 | + | ||||
| 121 | + | deadline := time.NewTimer(2 * time.Second) | |||
| 122 | + | defer deadline.Stop() | |||
| 123 | + | poll := time.NewTicker(time.Millisecond) | |||
| 124 | + | defer poll.Stop() | |||
| 125 | + | for { | |||
| 126 | + | interactive.mu.Lock() | |||
| 127 | + | busy := interactive.busy | |||
| 128 | + | got := interactive.lastCtxInput | |||
| 129 | + | interactive.mu.Unlock() | |||
| 130 | + | if !busy { | |||
| 131 | + | want := estimateTimelineMessageTokens(agent.Messages()) | |||
| 132 | + | if got != want || got == 0 { | |||
| 133 | + | t.Fatalf("context estimate after compaction = %d, want %d", got, want) | |||
| 134 | + | } | |||
| 135 | + | break | |||
| 136 | + | } | |||
| 137 | + | select { | |||
| 138 | + | case <-deadline.C: | |||
| 139 | + | t.Fatal("compaction did not finish") | |||
| 140 | + | case <-poll.C: | |||
| 141 | + | } | |||
| 142 | + | } | |||
| 143 | + | } | |||
| 144 | + | ||||
| 97 | func TestPreTurnCompactionPreservesPromptImages(t *testing.T) { | 145 | func TestPreTurnCompactionPreservesPromptImages(t *testing.T) { | ||
| 98 | client := &compactQueueClient{ | 146 | client := &compactQueueClient{ | ||
| 99 | compactionStarted: make(chan struct{}), | 147 | compactionStarted: make(chan struct{}), | ||
M packages/agent/modes/interactive.go +7 -1
| @@ -5504,7 +5504,7 @@ func (i *Interactive) runCompact(parent context.Context, auto bool) { | |||
| 5504 | 5504 | } | |
| 5505 | 5505 | i.pendingPostCompactNote = "" | |
| 5506 | 5506 | i.extNotes = stripAutoCompactNotes(i.extNotes) | |
| 5507 | - | i.lastCtxInput = 0 | |
| 5507 | + | i.lastCtxInput = estimateTimelineMessageTokens(msgs) | |
| 5508 | 5508 | i.toolCalls = map[string]*tui.ToolCallView{} | |
| 5509 | 5509 | i.toolOrder = nil | |
| 5510 | 5510 | i.toolGate = map[string]int{} | |
| @@ -5524,6 +5524,12 @@ func (i *Interactive) runCompact(parent context.Context, auto bool) { | |||
| 5524 | 5524 | next, i.queued = i.queued[0], i.queued[1:] | |
| 5525 | 5525 | hasNext = true | |
| 5526 | 5526 | } | |
| 5527 | + | // A prompt resumed immediately after compaction must not trigger | |
| 5528 | + | // the pre-turn threshold guard again before fresh provider usage | |
| 5529 | + | // replaces this estimate. | |
| 5530 | + | if hasNext { | |
| 5531 | + | i.lastCtxInput = 0 | |
| 5532 | + | } | |
| 5527 | 5533 | } | |
| 5528 | 5534 | i.mu.Unlock() | |
| 5529 | 5535 | i.invalidate() | |
| @@ -5504,7 +5504,7 @@ func (i *Interactive) runCompact(parent context.Context, auto bool) { | |||||
| 5504 | } | 5504 | } | ||
| 5505 | i.pendingPostCompactNote = "" | 5505 | i.pendingPostCompactNote = "" | ||
| 5506 | i.extNotes = stripAutoCompactNotes(i.extNotes) | 5506 | i.extNotes = stripAutoCompactNotes(i.extNotes) | ||
| 5507 | - | i.lastCtxInput = 0 | 5507 | + | i.lastCtxInput = estimateTimelineMessageTokens(msgs) |
| 5508 | i.toolCalls = map[string]*tui.ToolCallView{} | 5508 | i.toolCalls = map[string]*tui.ToolCallView{} | ||
| 5509 | i.toolOrder = nil | 5509 | i.toolOrder = nil | ||
| 5510 | i.toolGate = map[string]int{} | 5510 | i.toolGate = map[string]int{} | ||
| @@ -5524,6 +5524,12 @@ func (i *Interactive) runCompact(parent context.Context, auto bool) { | |||||
| 5524 | next, i.queued = i.queued[0], i.queued[1:] | 5524 | next, i.queued = i.queued[0], i.queued[1:] | ||
| 5525 | hasNext = true | 5525 | hasNext = true | ||
| 5526 | } | 5526 | } | ||
| 5527 | + | // A prompt resumed immediately after compaction must not trigger | |||
| 5528 | + | // the pre-turn threshold guard again before fresh provider usage | |||
| 5529 | + | // replaces this estimate. | |||
| 5530 | + | if hasNext { | |||
| 5531 | + | i.lastCtxInput = 0 | |||
| 5532 | + | } | |||
| 5527 | } | 5533 | } | ||
| 5528 | i.mu.Unlock() | 5534 | i.mu.Unlock() | ||
| 5529 | i.invalidate() | 5535 | i.invalidate() | ||