Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
71 changes: 44 additions & 27 deletions internal/llmloop/compression.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import (
"context"
"fmt"
"strings"
"sync"
"sync/atomic"
"time"

Expand Down Expand Up @@ -41,6 +42,16 @@ type compressionJob struct {
snapshotLen int // message count when the snapshot was taken
}

// compressionState is the async-compression bookkeeping for a single
// conversation (one RunPerFile call). The Runner is shared by concurrent
// per-file goroutines, so this state must not live on the Runner: a shared
// slot lets one file apply, cancel, or replace another file's compression
// job (#384).
type compressionState struct {
mu sync.Mutex
pendingJob *compressionJob
}

// CountMessagesTokens returns the rough token count of msgs by summing the
// per-message text token count. Exported because both review and scan top
// layers may want it for pre-flight checks.
Expand Down Expand Up @@ -252,31 +263,36 @@ func (r *Runner) runCompression(ctx context.Context, msgs []llm.Message, filePat
return rebuilt, nil
}

// triggerAsyncCompression kicks off a background compression job.
func (r *Runner) triggerAsyncCompression(ctx context.Context, messages []llm.Message, filePath string) {
// triggerAsyncCompression kicks off a background compression job for the
// conversation owning st. A no-op when a job is already pending — the
// check-and-set happens under st.mu so concurrent callers cannot replace
// (and thereby leak) an in-flight job.
func (r *Runner) triggerAsyncCompression(ctx context.Context, st *compressionState, messages []llm.Message, filePath string) {
st.mu.Lock()
if st.pendingJob != nil {
st.mu.Unlock()
return
}
msgSnapshot := copyMessages(messages)

asyncCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Minute)

job := &compressionJob{done: make(chan struct{}), cancel: cancel, snapshotLen: len(messages)}
r.compressionMu.Lock()
r.pendingJob = job
r.compressionMu.Unlock()
st.pendingJob = job
st.mu.Unlock()

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

r.compressionMu.Lock()
defer r.compressionMu.Unlock()
st.mu.Lock()
defer st.mu.Unlock()

if r.pendingJob != job {
if st.pendingJob != job {
return // cancelled or superseded
}
if err != nil {
// Compression failed — abandon the job rather than applying a
// truncated/unmodified snapshot over live messages.
r.pendingJob = nil
st.pendingJob = nil
close(job.done)
return
}
Expand All @@ -288,10 +304,10 @@ func (r *Runner) triggerAsyncCompression(ctx context.Context, messages []llm.Mes
// tryApplyPendingCompression checks whether a background compression has
// completed and swaps the rebuilt messages into place. Returns true if
// applied.
func (r *Runner) tryApplyPendingCompression(messages *[]llm.Message) bool {
r.compressionMu.Lock()
job := r.pendingJob
r.compressionMu.Unlock()
func (r *Runner) tryApplyPendingCompression(st *compressionState, messages *[]llm.Message) bool {
st.mu.Lock()
job := st.pendingJob
st.mu.Unlock()

if job == nil {
return false
Expand All @@ -300,8 +316,8 @@ func (r *Runner) tryApplyPendingCompression(messages *[]llm.Message) bool {
select {
case <-job.done:
applied := false
r.compressionMu.Lock()
if r.pendingJob == job && job.rebuilt != nil {
st.mu.Lock()
if st.pendingJob == job && job.rebuilt != nil {

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

To ensure robust defensive programming, we should add a guard to verify that job.snapshotLen <= len(*messages) before attempting to apply the rebuilt messages. Although messages is currently only expected to grow during the active conversation loop, if any future refactoring or error-handling path shrinks the message history while a background compression job is in flight, job.snapshotLen > len(*messages) could cause a slice bounds out-of-range panic when executing (*messages)[job.snapshotLen:], or lead to inconsistent history state by restoring previously deleted messages.

Suggested change
if st.pendingJob == job && job.rebuilt != nil {
if st.pendingJob == job && job.rebuilt != nil && job.snapshotLen <= len(*messages) {

rebuilt := job.rebuilt
// Preserve any messages appended after the snapshot was taken —
// the background job only compressed messages[:snapshotLen].
Expand All @@ -311,23 +327,24 @@ func (r *Runner) tryApplyPendingCompression(messages *[]llm.Message) bool {
*messages = rebuilt
applied = true
}
if r.pendingJob == job {
r.pendingJob = nil
if st.pendingJob == job {
st.pendingJob = nil
}
r.compressionMu.Unlock()
st.mu.Unlock()
return applied
default:
return false
}
}

// cancelPendingCompression aborts any in-flight background compression.
func (r *Runner) cancelPendingCompression() {
r.compressionMu.Lock()
defer r.compressionMu.Unlock()
// cancelPendingCompression aborts the conversation's in-flight background
// compression, if any.
func (r *Runner) cancelPendingCompression(st *compressionState) {
st.mu.Lock()
defer st.mu.Unlock()

if r.pendingJob != nil {
r.pendingJob.cancel()
r.pendingJob = nil
if st.pendingJob != nil {
st.pendingJob.cancel()
st.pendingJob = nil
}
}
43 changes: 25 additions & 18 deletions internal/llmloop/loop.go
Original file line number Diff line number Diff line change
Expand Up @@ -39,8 +39,9 @@ type Deps struct {
}

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

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

// Async compression is owned by this conversation alone; the deferred
// cancel aborts any job still in flight when the conversation ends.
st := &compressionState{}
defer r.cancelPendingCompression(st)

for toolReqCount > 0 {
select {
case <-ctx.Done():
Expand Down Expand Up @@ -250,7 +254,7 @@ func (r *Runner) RunPerFile(ctx context.Context, messages []llm.Message, newPath
consecutiveEmptyRounds = 0
}

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

r.tryApplyPendingCompression(messages)

tokenCount := CountMessagesTokens(*messages)
r.tryApplyPendingCompression(st, messages)

if tokenCount > warnLimit {
r.cancelPendingCompression()
// A conversation can already be over the warning threshold before this
// round's messages are appended (e.g. an oversized initial prompt).
if CountMessagesTokens(*messages) > warnLimit {
r.cancelPendingCompression(st)
*messages, _ = r.runCompression(ctx, *messages, filePath)
Comment on lines +446 to 448

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

Compress at the exact warning threshold before stopping.

When finalCount == warnLimit, synchronous compression is skipped, async compression is skipped, and the function returns false. This contradicts the contract that stopping occurs only after compression cannot bring the conversation below the threshold.

Proposed boundary fix
-	if CountMessagesTokens(*messages) > warnLimit {
+	if CountMessagesTokens(*messages) >= warnLimit {
 		r.cancelPendingCompression(st)
 		*messages, _ = r.runCompression(ctx, *messages, filePath)
 	}
...
-	if finalCount > warnLimit {
+	if finalCount >= warnLimit {
 		r.cancelPendingCompression(st)
 		*messages, _ = r.runCompression(ctx, *messages, filePath)

Please add an exact-warnLimit regression case.

Also applies to: 461-473

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@internal/llmloop/loop.go` around lines 444 - 446, Update the compression
boundary logic around CountMessagesTokens and the related finalCount handling so
an exact warnLimit value triggers synchronous compression before the stop
decision, while preserving asynchronous compression for values above the
threshold. Add a regression test covering finalCount == warnLimit and verify
compression is attempted before the function returns false.

tokenCount = CountMessagesTokens(*messages)
}

if tokenCount > softLimit && r.pendingJob == nil {
r.triggerAsyncCompression(ctx, *messages, filePath)
}

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

finalCount := CountMessagesTokens(*messages)
if finalCount > warnLimit {
r.cancelPendingCompression()
r.cancelPendingCompression(st)
*messages, _ = r.runCompression(ctx, *messages, filePath)
finalCount = CountMessagesTokens(*messages)
}

// Trigger async compression only after all appends for this update, so
// a job is never started and then immediately cancelled by the same
// call (#384), and never started when we are about to return false.
if finalCount > softLimit && finalCount < warnLimit {
r.triggerAsyncCompression(ctx, st, *messages, filePath)
}

return CountMessagesTokens(*messages) < warnLimit
return finalCount < warnLimit
}

// lookupTool returns the provider for a given tool from the registry, or
Expand Down
Loading