Skip to content

Commit 47a9222

Browse files
committed
fix(llmloop): scope async memory compression to each RunPerFile conversation
The Runner is shared by all concurrent per-file review goroutines, but it held a single compressionMu/pendingJob slot for async memory compression. With concurrency > 1 that shared slot caused four defects (#384): 1. Cross-file apply: tryApplyPendingCompression had no owner check, so file B could splice file A's rebuilt history into its own messages. 2. Cross-file cancel/replace: the warning-threshold paths canceled whichever file's job happened to be pending, surfacing spurious "context canceled" errors at the gateway; triggerAsyncCompression overwrote the slot unconditionally, wasting the superseded request. 3. Same-call start-then-cancel: the soft-threshold trigger fired before the new messages were appended, so an append that crossed the warning threshold canceled the job started microseconds earlier. 4. The pendingJob == nil fast-path read in addNextMessage was unlocked. Fix: move the bookkeeping into a compressionState owned by each RunPerFile call and thread it through trigger/apply/cancel. The nil-pending gate is now an atomic check-and-set inside triggerAsyncCompression under st.mu, and the async trigger moved to the end of addNextMessage, gated on the post-append count sitting strictly between the soft and warning thresholds. RunPerFile defers a cancel so no job outlives its conversation. Aggregate token counters and warnings stay Runner-level. Intentional behavior deltas, all strictly safer: the async snapshot now includes the just-appended round; an in-flight job is canceled when RunPerFile returns instead of running up to 5 minutes orphaned; no async job starts when the call is about to return false. Verified: new regression tests (cross-file isolation via a channel-gated fake client, owner-only summary apply with post-snapshot suffix preserved, no start-then-cancel in one update, 4 concurrent RunPerFile calls under -race); all existing compression tests updated to the per-conversation API; full suite green with -race; coverage 81.1%. E2E: 2-file concurrent review against a local OpenAI-compatible stub with compression exercised — zero "context canceled", both files done. Fixes #384
1 parent 0b226fa commit 47a9222

3 files changed

Lines changed: 337 additions & 62 deletions

File tree

internal/llmloop/compression.go

Lines changed: 44 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import (
44
"context"
55
"fmt"
66
"strings"
7+
"sync"
78
"sync/atomic"
89
"time"
910

@@ -41,6 +42,16 @@ type compressionJob struct {
4142
snapshotLen int // message count when the snapshot was taken
4243
}
4344

45+
// compressionState is the async-compression bookkeeping for a single
46+
// conversation (one RunPerFile call). The Runner is shared by concurrent
47+
// per-file goroutines, so this state must not live on the Runner: a shared
48+
// slot lets one file apply, cancel, or replace another file's compression
49+
// job (#384).
50+
type compressionState struct {
51+
mu sync.Mutex
52+
pendingJob *compressionJob
53+
}
54+
4455
// CountMessagesTokens returns the rough token count of msgs by summing the
4556
// per-message text token count. Exported because both review and scan top
4657
// layers may want it for pre-flight checks.
@@ -252,31 +263,36 @@ func (r *Runner) runCompression(ctx context.Context, msgs []llm.Message, filePat
252263
return rebuilt, nil
253264
}
254265

255-
// triggerAsyncCompression kicks off a background compression job.
256-
func (r *Runner) triggerAsyncCompression(ctx context.Context, messages []llm.Message, filePath string) {
266+
// triggerAsyncCompression kicks off a background compression job for the
267+
// conversation owning st. A no-op when a job is already pending — the
268+
// check-and-set happens under st.mu so concurrent callers cannot replace
269+
// (and thereby leak) an in-flight job.
270+
func (r *Runner) triggerAsyncCompression(ctx context.Context, st *compressionState, messages []llm.Message, filePath string) {
271+
st.mu.Lock()
272+
if st.pendingJob != nil {
273+
st.mu.Unlock()
274+
return
275+
}
257276
msgSnapshot := copyMessages(messages)
258-
259277
asyncCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Minute)
260-
261278
job := &compressionJob{done: make(chan struct{}), cancel: cancel, snapshotLen: len(messages)}
262-
r.compressionMu.Lock()
263-
r.pendingJob = job
264-
r.compressionMu.Unlock()
279+
st.pendingJob = job
280+
st.mu.Unlock()
265281

266282
go func() {
267283
defer cancel()
268284
rebuilt, err := r.runCompression(asyncCtx, msgSnapshot, filePath)
269285

270-
r.compressionMu.Lock()
271-
defer r.compressionMu.Unlock()
286+
st.mu.Lock()
287+
defer st.mu.Unlock()
272288

273-
if r.pendingJob != job {
289+
if st.pendingJob != job {
274290
return // cancelled or superseded
275291
}
276292
if err != nil {
277293
// Compression failed — abandon the job rather than applying a
278294
// truncated/unmodified snapshot over live messages.
279-
r.pendingJob = nil
295+
st.pendingJob = nil
280296
close(job.done)
281297
return
282298
}
@@ -288,10 +304,10 @@ func (r *Runner) triggerAsyncCompression(ctx context.Context, messages []llm.Mes
288304
// tryApplyPendingCompression checks whether a background compression has
289305
// completed and swaps the rebuilt messages into place. Returns true if
290306
// applied.
291-
func (r *Runner) tryApplyPendingCompression(messages *[]llm.Message) bool {
292-
r.compressionMu.Lock()
293-
job := r.pendingJob
294-
r.compressionMu.Unlock()
307+
func (r *Runner) tryApplyPendingCompression(st *compressionState, messages *[]llm.Message) bool {
308+
st.mu.Lock()
309+
job := st.pendingJob
310+
st.mu.Unlock()
295311

296312
if job == nil {
297313
return false
@@ -300,8 +316,8 @@ func (r *Runner) tryApplyPendingCompression(messages *[]llm.Message) bool {
300316
select {
301317
case <-job.done:
302318
applied := false
303-
r.compressionMu.Lock()
304-
if r.pendingJob == job && job.rebuilt != nil {
319+
st.mu.Lock()
320+
if st.pendingJob == job && job.rebuilt != nil {
305321
rebuilt := job.rebuilt
306322
// Preserve any messages appended after the snapshot was taken —
307323
// the background job only compressed messages[:snapshotLen].
@@ -311,23 +327,24 @@ func (r *Runner) tryApplyPendingCompression(messages *[]llm.Message) bool {
311327
*messages = rebuilt
312328
applied = true
313329
}
314-
if r.pendingJob == job {
315-
r.pendingJob = nil
330+
if st.pendingJob == job {
331+
st.pendingJob = nil
316332
}
317-
r.compressionMu.Unlock()
333+
st.mu.Unlock()
318334
return applied
319335
default:
320336
return false
321337
}
322338
}
323339

324-
// cancelPendingCompression aborts any in-flight background compression.
325-
func (r *Runner) cancelPendingCompression() {
326-
r.compressionMu.Lock()
327-
defer r.compressionMu.Unlock()
340+
// cancelPendingCompression aborts the conversation's in-flight background
341+
// compression, if any.
342+
func (r *Runner) cancelPendingCompression(st *compressionState) {
343+
st.mu.Lock()
344+
defer st.mu.Unlock()
328345

329-
if r.pendingJob != nil {
330-
r.pendingJob.cancel()
331-
r.pendingJob = nil
346+
if st.pendingJob != nil {
347+
st.pendingJob.cancel()
348+
st.pendingJob = nil
332349
}
333350
}

internal/llmloop/loop.go

Lines changed: 25 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -39,8 +39,9 @@ type Deps struct {
3939
}
4040

4141
// Runner is a per-session (across files) executor of the LLM tool-use
42-
// loop. Token counters, warnings, and the optional background compression
43-
// job are aggregated across every RunPerFile call.
42+
// loop. Token counters and warnings are aggregated across every RunPerFile
43+
// call; background memory compression is scoped to each RunPerFile
44+
// conversation (see compressionState).
4445
type Runner struct {
4546
deps Deps
4647
totalInputTokens int64 // atomically updated
@@ -51,8 +52,6 @@ type Runner struct {
5152
warnings []AgentWarning
5253
toolCallsMu sync.Mutex
5354
toolCalls map[string]int64
54-
compressionMu sync.Mutex
55-
pendingJob *compressionJob
5655
}
5756

5857
// NewRunner returns a Runner bound to the given dependencies.
@@ -153,6 +152,11 @@ func (r *Runner) RunPerFile(ctx context.Context, messages []llm.Message, newPath
153152
consecutiveEmptyRounds := 0
154153
sessionID := uuid.NewString()
155154

155+
// Async compression is owned by this conversation alone; the deferred
156+
// cancel aborts any job still in flight when the conversation ends.
157+
st := &compressionState{}
158+
defer r.cancelPendingCompression(st)
159+
156160
for toolReqCount > 0 {
157161
select {
158162
case <-ctx.Done():
@@ -250,7 +254,7 @@ func (r *Runner) RunPerFile(ctx context.Context, messages []llm.Message, newPath
250254
consecutiveEmptyRounds = 0
251255
}
252256

253-
succeed := r.addNextMessage(ctx, content, calls, results, &messages, newPath)
257+
succeed := r.addNextMessage(ctx, content, calls, results, &messages, newPath, st)
254258
if !succeed {
255259
fmt.Fprintf(stdout.Writer(), "[ocr] Context compression exceeded threshold for %s, stopping.\n", newPath)
256260
break
@@ -430,23 +434,18 @@ func (r *Runner) executeToolCall(ctx context.Context, newPath string, call llm.T
430434
// warning (80%) MaxTokens thresholds. Returns false when even after
431435
// synchronous compression the conversation is still over the warning
432436
// threshold — caller should stop the loop in that case.
433-
func (r *Runner) addNextMessage(ctx context.Context, assistantContent string, toolCalls []llm.ToolCall, results []tool.ToolCallResult, messages *[]llm.Message, filePath string) bool {
437+
func (r *Runner) addNextMessage(ctx context.Context, assistantContent string, toolCalls []llm.ToolCall, results []tool.ToolCallResult, messages *[]llm.Message, filePath string, st *compressionState) bool {
434438
maxAllowed := r.deps.Template.MaxTokens
435439
softLimit := int(float64(maxAllowed) * tokenSoftThreshold)
436440
warnLimit := int(float64(maxAllowed) * tokenWarningThreshold)
437441

438-
r.tryApplyPendingCompression(messages)
439-
440-
tokenCount := CountMessagesTokens(*messages)
442+
r.tryApplyPendingCompression(st, messages)
441443

442-
if tokenCount > warnLimit {
443-
r.cancelPendingCompression()
444+
// A conversation can already be over the warning threshold before this
445+
// round's messages are appended (e.g. an oversized initial prompt).
446+
if CountMessagesTokens(*messages) > warnLimit {
447+
r.cancelPendingCompression(st)
444448
*messages, _ = r.runCompression(ctx, *messages, filePath)
445-
tokenCount = CountMessagesTokens(*messages)
446-
}
447-
448-
if tokenCount > softLimit && r.pendingJob == nil {
449-
r.triggerAsyncCompression(ctx, *messages, filePath)
450449
}
451450

452451
if len(toolCalls) > 0 {
@@ -461,11 +460,19 @@ func (r *Runner) addNextMessage(ctx context.Context, assistantContent string, to
461460

462461
finalCount := CountMessagesTokens(*messages)
463462
if finalCount > warnLimit {
464-
r.cancelPendingCompression()
463+
r.cancelPendingCompression(st)
465464
*messages, _ = r.runCompression(ctx, *messages, filePath)
465+
finalCount = CountMessagesTokens(*messages)
466+
}
467+
468+
// Trigger async compression only after all appends for this update, so
469+
// a job is never started and then immediately cancelled by the same
470+
// call (#384), and never started when we are about to return false.
471+
if finalCount > softLimit && finalCount < warnLimit {
472+
r.triggerAsyncCompression(ctx, st, *messages, filePath)
466473
}
467474

468-
return CountMessagesTokens(*messages) < warnLimit
475+
return finalCount < warnLimit
469476
}
470477

471478
// lookupTool returns the provider for a given tool from the registry, or

0 commit comments

Comments
 (0)