mirror of
https://github.com/alibaba/open-code-review.git
synced 2026-07-09 17:28:58 +00:00
refactor(agent): split agent.go into util.go and compression.go, fix compression bugs
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 <previous_review_summary> 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
This commit is contained in:
kite
parent
ad870f9549
commit
602e3bbe94
5 changed files with 627 additions and 457 deletions
|
|
@ -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<previous_review_summary>\n"+rawSummary+"\n</previous_review_summary>")
|
||||
|
||||
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("<message id=\"%d\" role=\"%s\">\n", i, m.Role))
|
||||
sb.WriteString(" <content>\n")
|
||||
sb.WriteString(fmt.Sprintf(" %s\n", m.ExtractText()))
|
||||
sb.WriteString(" </content>\n")
|
||||
sb.WriteString("</message>")
|
||||
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))
|
||||
}
|
||||
|
|
|
|||
316
internal/agent/compression.go
Normal file
316
internal/agent/compression.go
Normal file
|
|
@ -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<previous_review_summary>"); idx >= 0 {
|
||||
currentText = currentText[:idx]
|
||||
}
|
||||
rebuilt[1] = llm.NewTextMessage(userMsg.Role, currentText+"\n\n<previous_review_summary>\n"+rawSummary+"\n</previous_review_summary>")
|
||||
|
||||
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
|
||||
}
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
112
internal/agent/util.go
Normal file
112
internal/agent/util.go
Normal file
|
|
@ -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("<message id=\"%d\" role=\"%s\">\n", i, m.Role))
|
||||
sb.WriteString(" <content>\n")
|
||||
sb.WriteString(fmt.Sprintf(" %s\n", m.ExtractText()))
|
||||
sb.WriteString(" </content>\n")
|
||||
sb.WriteString("</message>")
|
||||
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))
|
||||
}
|
||||
198
internal/agent/util_test.go
Normal file
198
internal/agent/util_test.go
Normal file
|
|
@ -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, `<message id="0" role="user">`) {
|
||||
t.Errorf("missing user message tag in output:\n%s", got)
|
||||
}
|
||||
if !strings.Contains(got, `<message id="1" role="assistant">`) {
|
||||
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)
|
||||
}
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue