diff --git a/internal/llmloop/loop.go b/internal/llmloop/loop.go index a4e5651a..cf03a7cb 100644 --- a/internal/llmloop/loop.go +++ b/internal/llmloop/loop.go @@ -166,6 +166,23 @@ const ( StopCompression ) +const ( + maxFinalizationRounds = 2 + finalizationPrompt = "Finish the review using the evidence already gathered. Submit all confirmed findings with code_comment, then call task_done with state DONE. Do not request more context." +) + +// finalizationToolDefs prevents a model that has already finished reasoning, +// or exhausted its context rounds, from opening another investigation branch. +func finalizationToolDefs(defs []llm.ToolDef) []llm.ToolDef { + final := make([]llm.ToolDef, 0, 2) + for _, def := range defs { + if def.Function.Name == tool.CodeComment.Name() || def.Function.Name == tool.TaskDone.Name() { + final = append(final, def) + } + } + return final +} + // RunPerFile drives the main LLM conversation loop for a single file. // It sends messages with the configured tool definitions, executes any // tool calls returned by the model, and collects review comments until @@ -179,6 +196,9 @@ func (r *Runner) RunPerFile(ctx context.Context, messages []llm.Message, newPath const maxConsecutiveEmptyRounds = 3 consecutiveEmptyRounds := 0 sessionID := uuid.NewString() + finalizationRounds := 0 + finalizationStop := StopNone + finalTools := finalizationToolDefs(r.deps.MainToolDefs) // Async compression is owned by this conversation alone; the deferred // cancel aborts any job still in flight when the conversation ends. @@ -189,14 +209,21 @@ func (r *Runner) RunPerFile(ctx context.Context, messages []llm.Message, newPath // reached zero, the run stopped on the round budget. The empty-round and // compression breaks overwrite it at their trigger points. stop := StopMaxRounds - for toolReqCount > 0 { + for toolReqCount > 0 || finalizationRounds > 0 { select { case <-ctx.Done(): return false, StopNone, ctx.Err() default: } - toolReqCount-- + finalizing := finalizationRounds > 0 + toolDefs := r.deps.MainToolDefs + if finalizing { + finalizationRounds-- + toolDefs = finalTools + } else { + toolReqCount-- + } fs := r.deps.Session.GetOrCreateFileSession(newPath) rec := fs.AppendTaskRecord(session.MainTask, append([]llm.Message(nil), messages...)) @@ -206,7 +233,7 @@ func (r *Runner) RunPerFile(ctx context.Context, messages []llm.Message, newPath resp, err := r.deps.LLMClient.CompletionsWithCtx(ctx, llm.ChatRequest{ Model: r.deps.Model, Messages: messages, - Tools: r.deps.MainToolDefs, + Tools: toolDefs, MaxTokens: r.deps.Template.CompletionTokenLimit(), SessionID: sessionID, }) @@ -235,11 +262,24 @@ func (r *Runner) RunPerFile(ctx context.Context, messages []llm.Message, newPath calls := resp.ToolCalls() if len(calls) == 0 { + if finalizing { + if content != "" { + messages = append(messages, llm.NewTextMessage("assistant", content)) + } + if finalizationRounds > 0 { + messages = append(messages, llm.NewTextMessage("user", finalizationPrompt)) + continue + } + stop = finalizationStop + break + } fmt.Fprintf(stdout.Writer(), "[ocr] No tool calls parsed for %s, retrying...\n", newPath) - messages = append(messages, llm.NewTextMessage("user", "You did not successfully call any tools. Please try again or use task_done if finished.")) if content != "" { - messages = append(messages[:len(messages)-1], llm.NewTextMessage("assistant", content), messages[len(messages)-1]) + messages = append(messages, llm.NewTextMessage("assistant", content)) } + messages = append(messages, llm.NewTextMessage("user", finalizationPrompt)) + finalizationRounds = maxFinalizationRounds + finalizationStop = StopEmptyRounds continue } @@ -295,6 +335,19 @@ func (r *Runner) RunPerFile(ctx context.Context, messages []llm.Message, newPath stop = StopCompression break } + if finalizing { + if finalizationRounds > 0 { + messages = append(messages, llm.NewTextMessage("user", finalizationPrompt)) + continue + } + stop = finalizationStop + break + } + if toolReqCount == 0 { + messages = append(messages, llm.NewTextMessage("user", finalizationPrompt)) + finalizationRounds = maxFinalizationRounds + finalizationStop = StopMaxRounds + } } if stop == StopMaxRounds { diff --git a/internal/llmloop/loop_test.go b/internal/llmloop/loop_test.go index 42358e10..f15debdb 100644 --- a/internal/llmloop/loop_test.go +++ b/internal/llmloop/loop_test.go @@ -81,6 +81,33 @@ func fileReadToolCallResponse(callID, args string) *llm.ChatResponse { } } +func textResponse(content string) *llm.ChatResponse { + return &llm.ChatResponse{ + Choices: []llm.Choice{{Message: llm.ResponseMessage{Content: &content}}}, + Model: "fake", + Usage: &llm.UsageInfo{PromptTokens: 10, CompletionTokens: 5}, + } +} + +func codeCommentResponse() *llm.ChatResponse { + content := "" + return &llm.ChatResponse{ + Choices: []llm.Choice{{Message: llm.ResponseMessage{ + Content: &content, + ToolCalls: []llm.ToolCall{{ + ID: "comment_1", + Type: "function", + Function: llm.FunctionCall{ + Name: "code_comment", + Arguments: `{"comments":[{"content":"issue","existing_code":"package main","category":"bug","severity":"medium"}]}`, + }, + }}, + }}}, + Model: "fake", + Usage: &llm.UsageInfo{PromptTokens: 10, CompletionTokens: 5}, + } +} + type fakeFileReadProvider struct { result string } @@ -91,14 +118,21 @@ func (f *fakeFileReadProvider) Execute(_ context.Context, _ map[string]any) (str } func newTestDeps(client llm.LLMClient) Deps { + collector := tool.NewCommentCollector() reg := tool.NewRegistry() reg.Register(&fakeFileReadProvider{result: "package main\n"}) + reg.Register(&tool.CodeCommentProvider{Collector: collector}) return Deps{ - LLMClient: client, - Model: "fake", - Template: template.Template{MaxTokens: 100000, MaxToolRequestTimes: 10}, - Tools: reg, - CommentCollector: tool.NewCommentCollector(), + LLMClient: client, + Model: "fake", + Template: template.Template{MaxTokens: 100000, MaxToolRequestTimes: 10}, + Tools: reg, + MainToolDefs: []llm.ToolDef{ + {Type: "function", Function: llm.FunctionDef{Name: "task_done"}}, + {Type: "function", Function: llm.FunctionDef{Name: "code_comment"}}, + {Type: "function", Function: llm.FunctionDef{Name: "file_read"}}, + }, + CommentCollector: collector, Session: session.New("/tmp/test-repo", "main", "fake", session.SessionOptions{}), } } @@ -309,12 +343,11 @@ func TestRunPerFile_UnknownTool(t *testing.T) { } func TestRunPerFile_MaxToolRequestsWithoutTaskDoneDoesNotComplete(t *testing.T) { - content := "" - client := &fakeClient{responses: []*llm.ChatResponse{{ - Choices: []llm.Choice{{Message: llm.ResponseMessage{Content: &content}}}, - Model: "fake", - Usage: &llm.UsageInfo{PromptTokens: 5, CompletionTokens: 5}, - }}} + client := &fakeClient{responses: []*llm.ChatResponse{ + fileReadToolCallResponse("call_1", `{"path":"main.go"}`), + textResponse("Unable to finalize."), + textResponse("Still unable to finalize."), + }} deps := newTestDeps(client) deps.Template.MaxToolRequestTimes = 1 runner := NewRunner(deps) @@ -330,6 +363,99 @@ func TestRunPerFile_MaxToolRequestsWithoutTaskDoneDoesNotComplete(t *testing.T) if stop != StopMaxRounds { t.Fatalf("expected StopMaxRounds, got %v", stop) } + if client.calls != 3 { + t.Fatalf("LLM calls = %d, want one context round and two failed finalization rounds", client.calls) + } +} + +func TestRunPerFile_FinalizesAfterToolRoundBudget(t *testing.T) { + client := &fakeClient{responses: []*llm.ChatResponse{ + fileReadToolCallResponse("call_1", `{"path":"main.go"}`), + taskDoneResponse(), + }} + deps := newTestDeps(client) + deps.Template.MaxToolRequestTimes = 1 + runner := NewRunner(deps) + + completed, stop, err := runner.RunPerFile( + context.Background(), + []llm.Message{llm.NewTextMessage("user", "review")}, + "main.go", + ) + if err != nil { + t.Fatalf("RunPerFile: %v", err) + } + if !completed || stop != StopNone { + t.Fatalf("completed = %v, stop = %v; want completed finalization", completed, stop) + } + if client.calls != 2 { + t.Fatalf("LLM calls = %d, want one context round and one finalization round", client.calls) + } + assertFinalizationTools(t, client.requests[1].Tools) +} + +func TestRunPerFile_TextOnlyResponseEntersFinalization(t *testing.T) { + client := &fakeClient{responses: []*llm.ChatResponse{ + textResponse("Analysis complete; no issues found."), + taskDoneResponse(), + }} + runner := NewRunner(newTestDeps(client)) + + completed, stop, err := runner.RunPerFile( + context.Background(), + []llm.Message{llm.NewTextMessage("user", "review")}, + "main.go", + ) + if err != nil { + t.Fatalf("RunPerFile: %v", err) + } + if !completed || stop != StopNone { + t.Fatalf("completed = %v, stop = %v; want completed finalization", completed, stop) + } + assertFinalizationTools(t, client.requests[1].Tools) +} + +func TestRunPerFile_FinalizationCompletesAfterCommentsAndTaskDone(t *testing.T) { + client := &fakeClient{responses: []*llm.ChatResponse{ + fileReadToolCallResponse("call_1", `{"path":"main.go"}`), + codeCommentResponse(), + taskDoneResponse(), + }} + deps := newTestDeps(client) + deps.Template.MaxToolRequestTimes = 1 + runner := NewRunner(deps) + + completed, stop, err := runner.RunPerFile( + context.Background(), + []llm.Message{llm.NewTextMessage("user", "review")}, + "main.go", + ) + if err != nil { + t.Fatalf("RunPerFile: %v", err) + } + if !completed || stop != StopNone { + t.Fatalf("completed = %v, stop = %v; want task_done after submitted comments", completed, stop) + } + if len(runner.CollectPendingComments()) != 1 { + t.Fatal("expected the finalization comment to be collected") + } +} + +func assertFinalizationTools(t *testing.T, defs []llm.ToolDef) { + t.Helper() + if len(defs) != 2 { + t.Fatalf("finalization tools = %d, want task_done and code_comment", len(defs)) + } + want := map[string]bool{"task_done": true, "code_comment": true} + for _, def := range defs { + if !want[def.Function.Name] { + t.Fatalf("unexpected finalization tool %q", def.Function.Name) + } + delete(want, def.Function.Name) + } + if len(want) != 0 { + t.Fatalf("missing finalization tools: %v", want) + } } func TestRunPerFile_EmptyToolResultsStopWithEmptyRounds(t *testing.T) {