fix(llmloop): finalize exhausted reviews

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 {

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) {