From 2f3ed410095cc27826ec714924b728a90d78ff38 Mon Sep 17 00:00:00 2001 From: kite Date: Fri, 19 Jun 2026 20:16:23 +0800 Subject: [PATCH] refactor(agent): split agent.go into util.go and compression.go, fix compression bugs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extract utility functions (stripEmptyPlanBlock, stripMarkdownFences, buildMessageXML, copyMessages, countMessagesTokens, reviewModeString, detectGitBranch) into util.go and all compression logic into compression.go to improve readability (agent.go: 1489 → 1114 lines). Fix pre-existing compression bugs from commit 14adf6f: - Return unmodified messages on compression failure instead of truncating to frozen zone (which discarded all conversation context) - Check error in async compression goroutine to avoid applying failed results - Track snapshot length to append post-snapshot messages when applying async compression, preventing silent message loss - Strip old before appending new one to prevent cumulative frozen zone bloat across multiple compressions - Guard pendingJob cleanup to only clear when it matches the completed job, preventing accidental clearing of a newly scheduled job - Add context.Context and 5s timeout to detectGitBranch - Deep copy ToolCalls in copyMessages to match session/history.go - Reuse finalCount in addNextMessage to avoid redundant token counting --- internal/agent/agent.go | 376 +------------------------------- internal/agent/compression.go | 316 +++++++++++++++++++++++++++ internal/agent/template_test.go | 82 ------- internal/agent/util.go | 112 ++++++++++ internal/agent/util_test.go | 198 +++++++++++++++++ 5 files changed, 627 insertions(+), 457 deletions(-) create mode 100644 internal/agent/compression.go delete mode 100644 internal/agent/template_test.go create mode 100644 internal/agent/util.go create mode 100644 internal/agent/util_test.go diff --git a/internal/agent/agent.go b/internal/agent/agent.go index a71bc07..90d5275 100644 --- a/internal/agent/agent.go +++ b/internal/agent/agent.go @@ -4,8 +4,6 @@ import ( "context" "encoding/json" "fmt" - "os/exec" - "regexp" "strings" "sync" "sync/atomic" @@ -24,27 +22,6 @@ import ( "github.com/open-code-review/open-code-review/internal/tool" ) -// planBlockPattern matches the optional "Review Plan" section in a MAIN_TASK -// template user message: a header line beginning with "### " whose text -// contains "Review Plan" or "审查计划" (with optional ASCII "(Optional)" / -// Chinese "（可选）" suffix), the {{plan_guidance}} placeholder on its own -// line, and one trailing blank line. The ASCII and Chinese header forms -// are matched separately because Go's regexp engine does not define \b -// around CJK ideographs. -var planBlockPattern = regexp.MustCompile( - `(?m)^### [^\n]*(?:Review Plan|审查计划)[^\n]*\n\{\{plan_guidance\}\}\n\n?`) - -// stripEmptyPlanBlock removes the "### Review Plan …\n{{plan_guidance}}\n\n" -// wrapper from a MAIN_TASK user message when the plan phase produced no -// guidance. The previous implementation hard-coded a single Chinese literal, -// which did not match the actual English template shipped in -// task_template.json, so the literal token "{{plan_guidance}}" leaked into -// the rendered prompt on every review where the plan phase was skipped or -// failed. Strip is a no-op when the wrapper is absent. -func stripEmptyPlanBlock(content string) string { - return planBlockPattern.ReplaceAllString(content, "") -} - // Args holds all dependencies and configuration needed to run a review session. type Args struct { // RepoDir is the root of the git repository. @@ -128,34 +105,6 @@ type AgentWarning struct { Type string `json:"type"` } -// compression thresholds as fractions of MaxTokens. -const ( - tokenSoftThreshold = 0.60 // async background compression - tokenWarningThreshold = 0.80 // immediate sync compression -) - -// round groups consecutive messages starting with an assistant message -// followed by zero or more tool result messages. -type round struct { - assistantIdx int - toolIdxs []int -} - -// partitionResult describes how messages should be split for compression. -type partitionResult struct { - frozenEnd int - compressEnd int - rounds []round - activeCount int -} - -// compressionJob tracks an in-flight background compression operation. -type compressionJob struct { - done chan struct{} - rebuilt []llm.Message - cancel context.CancelFunc -} - // Agent orchestrates the AI-powered code review. type Agent struct { args Args @@ -235,7 +184,7 @@ func New(args Args) *Agent { args.CommentCollector = tool.NewCommentCollector() } if args.Session == nil { - gitBranch := detectGitBranch(args.RepoDir) + gitBranch := detectGitBranch(context.Background(), args.RepoDir) mode := args.ReviewMode if mode == "" { mode = reviewModeString(args.From, args.To, args.Commit) @@ -1142,294 +1091,6 @@ func (a *Agent) collectPendingComments() []model.LlmComment { return a.args.CommentCollector.Comments() } -// addNextMessage adds assistant + tool response messages to the conversation history. -// Implements dual-threshold compression: -// - 60% of MaxTokens: trigger async background compression (non-blocking) -// - 80% of MaxTokens: perform synchronous compression immediately -func (a *Agent) addNextMessage(ctx context.Context, assistantContent string, toolCalls []llm.ToolCall, results []tool.ToolCallResult, messages *[]llm.Message, filePath string) bool { - maxAllowed := a.args.Template.MaxTokens - softLimit := int(float64(maxAllowed) * tokenSoftThreshold) - warnLimit := int(float64(maxAllowed) * tokenWarningThreshold) - - // Try to apply any completed async compression from a previous iteration. - a.tryApplyPendingCompression(messages) - - tokenCount := countMessagesTokens(*messages) - - // Hard threshold: synchronous compression. - if tokenCount > warnLimit { - a.cancelPendingCompression() - *messages, _ = a.runCompression(ctx, *messages, filePath) - tokenCount = countMessagesTokens(*messages) - } - - // Soft threshold: async compression. - if tokenCount > softLimit && a.pendingJob == nil { - a.triggerAsyncCompression(ctx, *messages, filePath) - } - - // Add assistant message with tool_calls when present. - if len(toolCalls) > 0 { - *messages = append(*messages, llm.NewToolCallMessage(assistantContent, toolCalls)) - } else if assistantContent != "" { - *messages = append(*messages, llm.NewTextMessage("assistant", assistantContent)) - } - - // Add tool response messages using Claude's tool_result format. - for _, r := range results { - *messages = append(*messages, llm.NewToolResultMessage(r.ToolCallID, r.Result)) - } - - // Final check: compress synchronously if still over warning limit. - finalCount := countMessagesTokens(*messages) - if finalCount > warnLimit { - a.cancelPendingCompression() - *messages, _ = a.runCompression(ctx, *messages, filePath) - } - - return countMessagesTokens(*messages) < warnLimit -} - -func countMessagesTokens(msgs []llm.Message) int { - var total int - for _, m := range msgs { - total += llm.CountTokens(m.ExtractText()) - } - return total -} - -// groupIntoRounds parses messages[start:] into logical (assistant + tool_results) pairs. -func groupIntoRounds(messages []llm.Message, start int) []round { - var rounds []round - i := start - for i < len(messages) { - if messages[i].Role == "assistant" { - r := round{assistantIdx: i} - i++ - for i < len(messages) && messages[i].Role == "tool" { - r.toolIdxs = append(r.toolIdxs, i) - i++ - } - rounds = append(rounds, r) - } else { - i++ - } - } - return rounds -} - -// computeActiveZoneSize returns how many trailing rounds fit within the remaining -// token budget after accounting for frozen zone and the compressed summary. -func computeActiveZoneSize(rounds []round, messages []llm.Message, maxTokens int, reservedTokens int) int { - budget := int(float64(maxTokens)*tokenWarningThreshold) - reservedTokens - if budget <= 0 { - return 0 - } - - count := 0 - tokensUsed := 0 - for i := len(rounds) - 1; i >= 0; i-- { - roundTokens := llm.CountTokens(messages[rounds[i].assistantIdx].ExtractText()) - for _, ti := range rounds[i].toolIdxs { - roundTokens += llm.CountTokens(messages[ti].ExtractText()) - } - if tokensUsed+roundTokens > budget { - break - } - tokensUsed += roundTokens - count++ - } - return count -} - -// partitionMessages divides messages into frozen, compress, and active zones. -// Frozen zone is always messages[0:2]. Active zone preserves the K most recent -// complete rounds based on available token budget. -func partitionMessages(messages []llm.Message, maxTokens int, prevSummaryTokenEstimate int) partitionResult { - result := partitionResult{frozenEnd: 2} - if len(messages) <= 2 { - result.compressEnd = len(messages) - return result - } - - result.rounds = groupIntoRounds(messages, 2) - if len(result.rounds) == 0 { - result.compressEnd = len(messages) - return result - } - - result.activeCount = computeActiveZoneSize(result.rounds, messages, maxTokens, prevSummaryTokenEstimate) - if result.activeCount >= len(result.rounds) { - // Everything fits — no compression needed. - result.compressEnd = len(messages) - result.activeCount = 0 - return result - } - - // compressEnd = index after the last round NOT in active zone. - activeStartIdx := len(result.rounds) - result.activeCount - lastCompressRound := result.rounds[activeStartIdx-1] - if len(lastCompressRound.toolIdxs) > 0 { - result.compressEnd = lastCompressRound.toolIdxs[len(lastCompressRound.toolIdxs)-1] + 1 - } else { - result.compressEnd = lastCompressRound.assistantIdx + 1 - } - - return result -} - -// stripMarkdownFences removes ```json and ``` wrappers that some models -// add around structured outputs. -func stripMarkdownFences(s string) string { - s = strings.TrimSpace(s) - if strings.HasPrefix(s, "```") { - if nl := strings.IndexByte(s, '\n'); nl >= 0 { - s = s[nl+1:] - } else { - s = strings.TrimPrefix(s, "```json") - s = strings.TrimPrefix(s, "```") - } - } - s = strings.TrimSpace(s) - if strings.HasSuffix(s, "```") { - s = strings.TrimSuffix(s, "```") - s = strings.TrimSpace(s) - } - return s -} - -// runCompression performs three-zone memory compression on the given messages. -// It summarizes the compress zone while preserving the active zone intact. -// Returns the rebuilt messages slice: [frozen] + [compressed_summary] + [active]. -func (a *Agent) runCompression(ctx context.Context, msgs []llm.Message, filePath string) ([]llm.Message, error) { - if len(a.args.Template.MemoryCompressionTask.Messages) == 0 || len(msgs) <= 2 { - return msgs[:min(len(msgs), 2)], nil - } - - part := partitionMessages(msgs, a.args.Template.MaxTokens, 0) - if part.compressEnd <= part.frozenEnd { - return msgs, nil - } - - contextXML := buildMessageXML(msgs[part.frozenEnd:part.compressEnd]) - - compressionMsgs := make([]llm.Message, 0, len(a.args.Template.MemoryCompressionTask.Messages)) - for _, m := range a.args.Template.MemoryCompressionTask.Messages { - content := strings.ReplaceAll(m.Content, "{{context}}", contextXML) - compressionMsgs = append(compressionMsgs, llm.NewTextMessage(m.Role, content)) - } - - startTime := time.Now() - resp, err := a.args.LLMClient.CompletionsWithCtx(ctx, llm.ChatRequest{ - Model: a.args.Model, - Messages: compressionMsgs, - MaxTokens: a.args.Template.MaxTokens, - }) - duration := time.Since(startTime) - - fs := a.session.GetOrCreateFileSession(filePath) - rec := fs.AppendTaskRecord(session.MemoryCompressionTask, compressionMsgs) - if err != nil { - rec.SetError(err, duration) - fmt.Fprintf(stdout.Writer(), "[ocr] Memory compression failed: %v\n", err) - return msgs[:part.frozenEnd], fmt.Errorf("memory compression: %w", err) - } - rec.SetResponse(resp, duration) - if resp.Usage != nil { - atomic.AddInt64(&a.totalInputTokens, resp.Usage.PromptTokens) - atomic.AddInt64(&a.totalOutputTokens, resp.Usage.CompletionTokens) - atomic.AddInt64(&a.totalCacheReadTokens, resp.Usage.CacheReadTokens) - atomic.AddInt64(&a.totalCacheWriteTokens, resp.Usage.CacheWriteTokens) - } - - rawSummary := stripMarkdownFences(resp.Content()) - if rawSummary == "" { - return msgs[:part.frozenEnd], nil - } - - rebuilt := make([]llm.Message, 2) - copy(rebuilt, msgs[:2]) - - userMsg := rebuilt[1] - currentText := userMsg.ExtractText() - rebuilt[1] = llm.NewTextMessage(userMsg.Role, currentText+"\n\n\n"+rawSummary+"\n") - - for i := part.compressEnd; i < len(msgs); i++ { - rebuilt = append(rebuilt, msgs[i]) - } - - return rebuilt, nil -} - -// triggerAsyncCompression kicks off a background compression job. -func (a *Agent) triggerAsyncCompression(ctx context.Context, messages []llm.Message, filePath string) { - msgSnapshot := copyMessages(messages) - - asyncCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Minute) - - job := &compressionJob{done: make(chan struct{}), cancel: cancel} - a.compressionMu.Lock() - a.pendingJob = job - a.compressionMu.Unlock() - - go func() { - defer cancel() - rebuilt, _ := a.runCompression(asyncCtx, msgSnapshot, filePath) - - a.compressionMu.Lock() - defer a.compressionMu.Unlock() - - if a.pendingJob != job { - return // cancelled or superseded - } - job.rebuilt = rebuilt - close(job.done) - }() -} - -// tryApplyPendingCompression checks if a background compression completed -// and swaps the rebuilt messages into place. Returns true if applied. -func (a *Agent) tryApplyPendingCompression(messages *[]llm.Message) bool { - a.compressionMu.Lock() - job := a.pendingJob - a.compressionMu.Unlock() - - if job == nil { - return false - } - - select { - case <-job.done: - a.compressionMu.Lock() - if a.pendingJob == job && job.rebuilt != nil { - *messages = job.rebuilt - a.pendingJob = nil - } - a.compressionMu.Unlock() - return true - default: - return false - } -} - -// cancelPendingCompression aborts any in-flight background compression. -func (a *Agent) cancelPendingCompression() { - a.compressionMu.Lock() - defer a.compressionMu.Unlock() - - if a.pendingJob != nil { - a.pendingJob.cancel() - a.pendingJob = nil - } -} - -// copyMessages creates a shallow copy of a message slice. -func copyMessages(msgs []llm.Message) []llm.Message { - out := make([]llm.Message, len(msgs)) - copy(out, msgs) - return out -} - // BuildToolDefs converts toolsconfig.ToolConfigEntry slice into []llm.ToolDef, // filtering by phase (planOnly=true for plan_task, false for main_task). func BuildToolDefs(entries []toolsconfig.ToolConfigEntry, planOnly bool) []llm.ToolDef { @@ -1451,38 +1112,3 @@ func BuildToolDefs(entries []toolsconfig.ToolConfigEntry, planOnly bool) []llm.T } return defs } - -func buildMessageXML(msgs []llm.Message) string { - var sb strings.Builder - for i, m := range msgs { - sb.WriteString(fmt.Sprintf("\n", i, m.Role)) - sb.WriteString(" \n") - sb.WriteString(fmt.Sprintf(" %s\n", m.ExtractText())) - sb.WriteString(" \n") - sb.WriteString("") - if i < len(msgs)-1 { - sb.WriteString("\n") - } - } - return sb.String() -} - -func reviewModeString(from, to, commit string) string { - if commit != "" { - return session.ReviewModeCommit - } - if from != "" && to != "" { - return session.ReviewModeRange - } - return session.ReviewModeWorkspace -} - -// detectGitBranch returns the current git branch name for the given repo, or empty string on failure. -func detectGitBranch(repoDir string) string { - cmd := exec.Command("git", "-C", repoDir, "rev-parse", "--abbrev-ref", "HEAD") - out, err := cmd.Output() - if err != nil || len(out) == 0 { - return "" - } - return strings.TrimSpace(string(out)) -} diff --git a/internal/agent/compression.go b/internal/agent/compression.go new file mode 100644 index 0000000..4656fae --- /dev/null +++ b/internal/agent/compression.go @@ -0,0 +1,316 @@ +package agent + +import ( + "context" + "fmt" + "strings" + "sync/atomic" + "time" + + "github.com/open-code-review/open-code-review/internal/llm" + "github.com/open-code-review/open-code-review/internal/session" + "github.com/open-code-review/open-code-review/internal/stdout" + "github.com/open-code-review/open-code-review/internal/tool" +) + +// compression thresholds as fractions of MaxTokens. +const ( + tokenSoftThreshold = 0.60 // async background compression + tokenWarningThreshold = 0.80 // immediate sync compression +) + +// round groups consecutive messages starting with an assistant message +// followed by zero or more tool result messages. +type round struct { + assistantIdx int + toolIdxs []int +} + +// partitionResult describes how messages should be split for compression. +type partitionResult struct { + frozenEnd int + compressEnd int + rounds []round + activeCount int +} + +// compressionJob tracks an in-flight background compression operation. +type compressionJob struct { + done chan struct{} + rebuilt []llm.Message + cancel context.CancelFunc + snapshotLen int // message count when the snapshot was taken +} + +// groupIntoRounds parses messages[start:] into logical (assistant + tool_results) pairs. +func groupIntoRounds(messages []llm.Message, start int) []round { + var rounds []round + i := start + for i < len(messages) { + if messages[i].Role == "assistant" { + r := round{assistantIdx: i} + i++ + for i < len(messages) && messages[i].Role == "tool" { + r.toolIdxs = append(r.toolIdxs, i) + i++ + } + rounds = append(rounds, r) + } else { + i++ + } + } + return rounds +} + +// computeActiveZoneSize returns how many trailing rounds fit within the remaining +// token budget after accounting for frozen zone and the compressed summary. +func computeActiveZoneSize(rounds []round, messages []llm.Message, maxTokens int, reservedTokens int) int { + budget := int(float64(maxTokens)*tokenWarningThreshold) - reservedTokens + if budget <= 0 { + return 0 + } + + count := 0 + tokensUsed := 0 + for i := len(rounds) - 1; i >= 0; i-- { + roundTokens := llm.CountTokens(messages[rounds[i].assistantIdx].ExtractText()) + for _, ti := range rounds[i].toolIdxs { + roundTokens += llm.CountTokens(messages[ti].ExtractText()) + } + if tokensUsed+roundTokens > budget { + break + } + tokensUsed += roundTokens + count++ + } + return count +} + +// partitionMessages divides messages into frozen, compress, and active zones. +// Frozen zone is always messages[0:2]. Active zone preserves the K most recent +// complete rounds based on available token budget. +func partitionMessages(messages []llm.Message, maxTokens int, prevSummaryTokenEstimate int) partitionResult { + result := partitionResult{frozenEnd: 2} + if len(messages) <= 2 { + result.compressEnd = len(messages) + return result + } + + result.rounds = groupIntoRounds(messages, 2) + if len(result.rounds) == 0 { + result.compressEnd = len(messages) + return result + } + + result.activeCount = computeActiveZoneSize(result.rounds, messages, maxTokens, prevSummaryTokenEstimate) + if result.activeCount >= len(result.rounds) { + // Everything fits — no compression needed. + result.compressEnd = len(messages) + result.activeCount = 0 + return result + } + + // compressEnd = index after the last round NOT in active zone. + activeStartIdx := len(result.rounds) - result.activeCount + lastCompressRound := result.rounds[activeStartIdx-1] + if len(lastCompressRound.toolIdxs) > 0 { + result.compressEnd = lastCompressRound.toolIdxs[len(lastCompressRound.toolIdxs)-1] + 1 + } else { + result.compressEnd = lastCompressRound.assistantIdx + 1 + } + + return result +} + +// addNextMessage adds assistant + tool response messages to the conversation history. +// Implements dual-threshold compression: +// - 60% of MaxTokens: trigger async background compression (non-blocking) +// - 80% of MaxTokens: perform synchronous compression immediately +func (a *Agent) addNextMessage(ctx context.Context, assistantContent string, toolCalls []llm.ToolCall, results []tool.ToolCallResult, messages *[]llm.Message, filePath string) bool { + maxAllowed := a.args.Template.MaxTokens + softLimit := int(float64(maxAllowed) * tokenSoftThreshold) + warnLimit := int(float64(maxAllowed) * tokenWarningThreshold) + + // Try to apply any completed async compression from a previous iteration. + a.tryApplyPendingCompression(messages) + + tokenCount := countMessagesTokens(*messages) + + // Hard threshold: synchronous compression. + if tokenCount > warnLimit { + a.cancelPendingCompression() + *messages, _ = a.runCompression(ctx, *messages, filePath) + tokenCount = countMessagesTokens(*messages) + } + + // Soft threshold: async compression. + if tokenCount > softLimit && a.pendingJob == nil { + a.triggerAsyncCompression(ctx, *messages, filePath) + } + + // Add assistant message with tool_calls when present. + if len(toolCalls) > 0 { + *messages = append(*messages, llm.NewToolCallMessage(assistantContent, toolCalls)) + } else if assistantContent != "" { + *messages = append(*messages, llm.NewTextMessage("assistant", assistantContent)) + } + + // Add tool response messages using Claude's tool_result format. + for _, r := range results { + *messages = append(*messages, llm.NewToolResultMessage(r.ToolCallID, r.Result)) + } + + // Final check: compress synchronously if still over warning limit. + finalCount := countMessagesTokens(*messages) + if finalCount > warnLimit { + a.cancelPendingCompression() + *messages, _ = a.runCompression(ctx, *messages, filePath) + finalCount = countMessagesTokens(*messages) + } + + return finalCount < warnLimit +} + +// runCompression performs three-zone memory compression on the given messages. +// It summarizes the compress zone while preserving the active zone intact. +// Returns the rebuilt messages slice: [frozen] + [compressed_summary] + [active]. +func (a *Agent) runCompression(ctx context.Context, msgs []llm.Message, filePath string) ([]llm.Message, error) { + if len(a.args.Template.MemoryCompressionTask.Messages) == 0 || len(msgs) <= 2 { + return msgs[:min(len(msgs), 2)], nil + } + + part := partitionMessages(msgs, a.args.Template.MaxTokens, 0) + if part.compressEnd <= part.frozenEnd { + return msgs, nil + } + + contextXML := buildMessageXML(msgs[part.frozenEnd:part.compressEnd]) + + compressionMsgs := make([]llm.Message, 0, len(a.args.Template.MemoryCompressionTask.Messages)) + for _, m := range a.args.Template.MemoryCompressionTask.Messages { + content := strings.ReplaceAll(m.Content, "{{context}}", contextXML) + compressionMsgs = append(compressionMsgs, llm.NewTextMessage(m.Role, content)) + } + + startTime := time.Now() + resp, err := a.args.LLMClient.CompletionsWithCtx(ctx, llm.ChatRequest{ + Model: a.args.Model, + Messages: compressionMsgs, + MaxTokens: a.args.Template.MaxTokens, + }) + duration := time.Since(startTime) + + fs := a.session.GetOrCreateFileSession(filePath) + rec := fs.AppendTaskRecord(session.MemoryCompressionTask, compressionMsgs) + if err != nil { + rec.SetError(err, duration) + fmt.Fprintf(stdout.Writer(), "[ocr] Memory compression failed: %v\n", err) + // Intentionally return unmodified msgs: truncating to frozenEnd would + // discard all conversation context, which is worse than staying over + // the token limit temporarily. + return msgs, fmt.Errorf("memory compression: %w", err) + } + rec.SetResponse(resp, duration) + if resp.Usage != nil { + atomic.AddInt64(&a.totalInputTokens, resp.Usage.PromptTokens) + atomic.AddInt64(&a.totalOutputTokens, resp.Usage.CompletionTokens) + atomic.AddInt64(&a.totalCacheReadTokens, resp.Usage.CacheReadTokens) + atomic.AddInt64(&a.totalCacheWriteTokens, resp.Usage.CacheWriteTokens) + } + + rawSummary := stripMarkdownFences(resp.Content()) + if rawSummary == "" { + return msgs, nil + } + + rebuilt := make([]llm.Message, 2) + copy(rebuilt, msgs[:2]) + + userMsg := rebuilt[1] + currentText := userMsg.ExtractText() + if idx := strings.Index(currentText, "\n\n"); idx >= 0 { + currentText = currentText[:idx] + } + rebuilt[1] = llm.NewTextMessage(userMsg.Role, currentText+"\n\n\n"+rawSummary+"\n") + + for i := part.compressEnd; i < len(msgs); i++ { + rebuilt = append(rebuilt, msgs[i]) + } + + return rebuilt, nil +} + +// triggerAsyncCompression kicks off a background compression job. +func (a *Agent) triggerAsyncCompression(ctx context.Context, messages []llm.Message, filePath string) { + msgSnapshot := copyMessages(messages) + + asyncCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Minute) + + job := &compressionJob{done: make(chan struct{}), cancel: cancel, snapshotLen: len(messages)} + a.compressionMu.Lock() + a.pendingJob = job + a.compressionMu.Unlock() + + go func() { + defer cancel() + rebuilt, err := a.runCompression(asyncCtx, msgSnapshot, filePath) + + a.compressionMu.Lock() + defer a.compressionMu.Unlock() + + if a.pendingJob != job { + return // cancelled or superseded + } + if err != nil { + a.pendingJob = nil + close(job.done) + return + } + job.rebuilt = rebuilt + close(job.done) + }() +} + +// tryApplyPendingCompression checks if a background compression completed +// and swaps the rebuilt messages into place. Returns true if applied. +func (a *Agent) tryApplyPendingCompression(messages *[]llm.Message) bool { + a.compressionMu.Lock() + job := a.pendingJob + a.compressionMu.Unlock() + + if job == nil { + return false + } + + select { + case <-job.done: + applied := false + a.compressionMu.Lock() + if a.pendingJob == job && job.rebuilt != nil { + rebuilt := job.rebuilt + if job.snapshotLen < len(*messages) { + rebuilt = append(rebuilt, (*messages)[job.snapshotLen:]...) + } + *messages = rebuilt + applied = true + } + if a.pendingJob == job { + a.pendingJob = nil + } + a.compressionMu.Unlock() + return applied + default: + return false + } +} + +// cancelPendingCompression aborts any in-flight background compression. +func (a *Agent) cancelPendingCompression() { + a.compressionMu.Lock() + defer a.compressionMu.Unlock() + + if a.pendingJob != nil { + a.pendingJob.cancel() + a.pendingJob = nil + } +} diff --git a/internal/agent/template_test.go b/internal/agent/template_test.go deleted file mode 100644 index 55cb1bd..0000000 --- a/internal/agent/template_test.go +++ /dev/null @@ -1,82 +0,0 @@ -package agent - -import ( - "strings" - "testing" -) - -func TestStripEmptyPlanBlock(t *testing.T) { - tests := []struct { - name string - input string - want string - }{ - { - name: "english template wrapper is removed", - input: "header\n### Review Plan (Optional)\n{{plan_guidance}}\n\ntail", - want: "header\ntail", - }, - { - name: "english template wrapper without trailing blank line is removed", - input: "header\n### Review Plan (Optional)\n{{plan_guidance}}\ntail", - want: "header\ntail", - }, - { - name: "chinese template wrapper is removed", - input: "header\n### 审查计划\n{{plan_guidance}}\n\ntail", - want: "header\ntail", - }, - { - name: "chinese optional wrapper is removed", - input: "header\n### 审查计划（可选）\n{{plan_guidance}}\n\ntail", - want: "header\ntail", - }, - { - name: "no wrapper present is a no-op", - input: "no plan block here\njust text", - want: "no plan block here\njust text", - }, - { - name: "multiple wrappers all removed", - input: "### Review Plan (Optional)\n{{plan_guidance}}\n\nmiddle\n### 审查计划\n{{plan_guidance}}\n\nend", - want: "middle\nend", - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - got := stripEmptyPlanBlock(tt.input) - if got != tt.want { - t.Errorf("stripEmptyPlanBlock() = %q, want %q", got, tt.want) - } - if strings.Contains(got, "{{plan_guidance}}") { - t.Errorf("stripEmptyPlanBlock() left literal {{plan_guidance}} in output: %q", got) - } - }) - } -} - -func TestStripEmptyPlanBlock_IntegrationWithReplaceAll(t *testing.T) { - // Regression: validate the executeSubtask call-site flow — the wrapper - // passes through stripEmptyPlanBlock first (token still present), then - // through ReplaceAll. If strip runs after the replace, the token is - // already gone and the wrapper header leaks into the prompt. - template := "header\n### Review Plan (Optional)\n{{plan_guidance}}\n\ntail" - - // Phase 1: strip (token must be present for regex to match) - stripped := stripEmptyPlanBlock(template) - - // Phase 2: replace (no-op since token should already be consumed) - final := strings.ReplaceAll(stripped, "{{plan_guidance}}", "") - - want := "header\ntail" - if final != want { - t.Errorf("stripEmptyPlanBlock integration:\n got: %q\n want: %q", final, want) - } - if strings.Contains(final, "{{plan_guidance}}") { - t.Errorf("literal {{plan_guidance}} leaked: %q", final) - } - if strings.Contains(final, "Review Plan") { - t.Errorf("dangling Review Plan header retained: %q", final) - } -} diff --git a/internal/agent/util.go b/internal/agent/util.go new file mode 100644 index 0000000..7651180 --- /dev/null +++ b/internal/agent/util.go @@ -0,0 +1,112 @@ +package agent + +import ( + "context" + "fmt" + "os/exec" + "regexp" + "strings" + "time" + + "github.com/open-code-review/open-code-review/internal/llm" + "github.com/open-code-review/open-code-review/internal/session" +) + +// planBlockPattern matches the optional "Review Plan" section in a MAIN_TASK +// template user message: a header line beginning with "### " whose text +// contains "Review Plan" or "审查计划" (with optional ASCII "(Optional)" / +// Chinese "（可选）" suffix), the {{plan_guidance}} placeholder on its own +// line, and one trailing blank line. The ASCII and Chinese header forms +// are matched separately because Go's regexp engine does not define \b +// around CJK ideographs. +var planBlockPattern = regexp.MustCompile( + `(?m)^### [^\n]*(?:Review Plan|审查计划)[^\n]*\n\{\{plan_guidance\}\}\n\n?`) + +// stripEmptyPlanBlock removes the "### Review Plan …\n{{plan_guidance}}\n\n" +// wrapper from a MAIN_TASK user message when the plan phase produced no +// guidance. The previous implementation hard-coded a single Chinese literal, +// which did not match the actual English template shipped in +// task_template.json, so the literal token "{{plan_guidance}}" leaked into +// the rendered prompt on every review where the plan phase was skipped or +// failed. Strip is a no-op when the wrapper is absent. +func stripEmptyPlanBlock(content string) string { + return planBlockPattern.ReplaceAllString(content, "") +} + +// stripMarkdownFences removes ```json and ``` wrappers that some models +// add around structured outputs. +func stripMarkdownFences(s string) string { + s = strings.TrimSpace(s) + if strings.HasPrefix(s, "```") { + if nl := strings.IndexByte(s, '\n'); nl >= 0 { + s = s[nl+1:] + } else { + s = strings.TrimPrefix(s, "```json") + s = strings.TrimPrefix(s, "```") + } + } + s = strings.TrimSpace(s) + if strings.HasSuffix(s, "```") { + s = strings.TrimSuffix(s, "```") + s = strings.TrimSpace(s) + } + return s +} + +func buildMessageXML(msgs []llm.Message) string { + var sb strings.Builder + for i, m := range msgs { + sb.WriteString(fmt.Sprintf("\n", i, m.Role)) + sb.WriteString(" \n") + sb.WriteString(fmt.Sprintf(" %s\n", m.ExtractText())) + sb.WriteString(" \n") + sb.WriteString("") + if i < len(msgs)-1 { + sb.WriteString("\n") + } + } + return sb.String() +} + +func copyMessages(msgs []llm.Message) []llm.Message { + out := make([]llm.Message, len(msgs)) + for i, m := range msgs { + out[i] = llm.Message{ + Role: m.Role, + Content: m.Content, + ToolCallID: m.ToolCallID, + ToolCalls: append([]llm.ToolCall(nil), m.ToolCalls...), + } + } + return out +} + +func countMessagesTokens(msgs []llm.Message) int { + var total int + for _, m := range msgs { + total += llm.CountTokens(m.ExtractText()) + } + return total +} + +func reviewModeString(from, to, commit string) string { + if commit != "" { + return session.ReviewModeCommit + } + if from != "" && to != "" { + return session.ReviewModeRange + } + return session.ReviewModeWorkspace +} + +// detectGitBranch returns the current git branch name for the given repo, or empty string on failure. +func detectGitBranch(ctx context.Context, repoDir string) string { + ctx, cancel := context.WithTimeout(ctx, 5*time.Second) + defer cancel() + cmd := exec.CommandContext(ctx, "git", "-C", repoDir, "rev-parse", "--abbrev-ref", "HEAD") + out, err := cmd.Output() + if err != nil || len(out) == 0 { + return "" + } + return strings.TrimSpace(string(out)) +} diff --git a/internal/agent/util_test.go b/internal/agent/util_test.go new file mode 100644 index 0000000..af3b8ff --- /dev/null +++ b/internal/agent/util_test.go @@ -0,0 +1,198 @@ +package agent + +import ( + "strings" + "testing" + + "github.com/open-code-review/open-code-review/internal/llm" + "github.com/open-code-review/open-code-review/internal/session" +) + +func TestStripEmptyPlanBlock(t *testing.T) { + tests := []struct { + name string + input string + want string + }{ + { + name: "english template wrapper is removed", + input: "header\n### Review Plan (Optional)\n{{plan_guidance}}\n\ntail", + want: "header\ntail", + }, + { + name: "english template wrapper without trailing blank line is removed", + input: "header\n### Review Plan (Optional)\n{{plan_guidance}}\ntail", + want: "header\ntail", + }, + { + name: "chinese template wrapper is removed", + input: "header\n### 审查计划\n{{plan_guidance}}\n\ntail", + want: "header\ntail", + }, + { + name: "chinese optional wrapper is removed", + input: "header\n### 审查计划（可选）\n{{plan_guidance}}\n\ntail", + want: "header\ntail", + }, + { + name: "no wrapper present is a no-op", + input: "no plan block here\njust text", + want: "no plan block here\njust text", + }, + { + name: "multiple wrappers all removed", + input: "### Review Plan (Optional)\n{{plan_guidance}}\n\nmiddle\n### 审查计划\n{{plan_guidance}}\n\nend", + want: "middle\nend", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := stripEmptyPlanBlock(tt.input) + if got != tt.want { + t.Errorf("stripEmptyPlanBlock() = %q, want %q", got, tt.want) + } + if strings.Contains(got, "{{plan_guidance}}") { + t.Errorf("stripEmptyPlanBlock() left literal {{plan_guidance}} in output: %q", got) + } + }) + } +} + +func TestStripEmptyPlanBlock_IntegrationWithReplaceAll(t *testing.T) { + template := "header\n### Review Plan (Optional)\n{{plan_guidance}}\n\ntail" + + stripped := stripEmptyPlanBlock(template) + final := strings.ReplaceAll(stripped, "{{plan_guidance}}", "") + + want := "header\ntail" + if final != want { + t.Errorf("stripEmptyPlanBlock integration:\n got: %q\n want: %q", final, want) + } + if strings.Contains(final, "{{plan_guidance}}") { + t.Errorf("literal {{plan_guidance}} leaked: %q", final) + } + if strings.Contains(final, "Review Plan") { + t.Errorf("dangling Review Plan header retained: %q", final) + } +} + +func TestStripMarkdownFences(t *testing.T) { + tests := []struct { + name string + input string + want string + }{ + { + name: "no fences", + input: `["c-0","c-2"]`, + want: `["c-0","c-2"]`, + }, + { + name: "json fenced block", + input: "```json\n[\"c-0\"]\n```", + want: `["c-0"]`, + }, + { + name: "plain fenced block", + input: "```\nhello\n```", + want: "hello", + }, + { + name: "surrounding whitespace", + input: " \n```json\ncontent\n```\n ", + want: "content", + }, + { + name: "empty string", + input: "", + want: "", + }, + { + name: "only opening fence no newline", + input: "```json{}```", + want: "{}", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := stripMarkdownFences(tt.input) + if got != tt.want { + t.Errorf("stripMarkdownFences() = %q, want %q", got, tt.want) + } + }) + } +} + +func TestBuildMessageXML(t *testing.T) { + msgs := []llm.Message{ + llm.NewTextMessage("user", "hello"), + llm.NewTextMessage("assistant", "world"), + } + + got := buildMessageXML(msgs) + + if !strings.Contains(got, ``) { + t.Errorf("missing user message tag in output:\n%s", got) + } + if !strings.Contains(got, ``) { + t.Errorf("missing assistant message tag in output:\n%s", got) + } + if !strings.Contains(got, "hello") || !strings.Contains(got, "world") { + t.Errorf("missing message content in output:\n%s", got) + } +} + +func TestCopyMessages(t *testing.T) { + orig := []llm.Message{ + llm.NewTextMessage("user", "a"), + llm.NewTextMessage("assistant", "b"), + } + + cp := copyMessages(orig) + + if len(cp) != len(orig) { + t.Fatalf("copyMessages length = %d, want %d", len(cp), len(orig)) + } + + cp = append(cp, llm.NewTextMessage("user", "c")) + if len(orig) != 2 { + t.Error("copyMessages: appending to copy modified original slice") + } +} + +func TestCountMessagesTokens(t *testing.T) { + msgs := []llm.Message{ + llm.NewTextMessage("user", "hello world"), + } + + count := countMessagesTokens(msgs) + if count <= 0 { + t.Errorf("countMessagesTokens() = %d, want > 0", count) + } + + empty := countMessagesTokens(nil) + if empty != 0 { + t.Errorf("countMessagesTokens(nil) = %d, want 0", empty) + } +} + +func TestReviewModeString(t *testing.T) { + tests := []struct { + from, to, commit string + want string + }{ + {"", "", "abc123", session.ReviewModeCommit}, + {"main", "feature", "", session.ReviewModeRange}, + {"", "", "", session.ReviewModeWorkspace}, + {"main", "feature", "abc123", session.ReviewModeCommit}, + } + + for _, tt := range tests { + got := reviewModeString(tt.from, tt.to, tt.commit) + if got != tt.want { + t.Errorf("reviewModeString(%q, %q, %q) = %q, want %q", tt.from, tt.to, tt.commit, got, tt.want) + } + } +}