feat(llmloop): add grace round after tool-request budget exhausted

internal/llmloop/loop.go

@@ -303,10 +303,76 @@ func (r *Runner) RunPerFile(ctx context.Context, messages []llm.Message, newPath
 
 	if stop == StopMaxRounds {
 		fmt.Fprintf(stdout.Writer(), "[ocr] Max tool requests reached for %s.\n", newPath)
+		r.runGraceRound(ctx, messages, newPath, sessionID)
 	}
 	return false, stop, nil
 }
 
+// runGraceRound performs one final LLM call after the tool-request budget is
+// exhausted, giving the model a chance to submit any findings it identified
+// but did not yet report via code_comment.
+func (r *Runner) runGraceRound(ctx context.Context, messages []llm.Message, newPath string, sessionID string) {
+	graceDefs := graceRoundToolDefs(r.deps.MainToolDefs)
+	if len(graceDefs) == 0 {
+		return
+	}
+
+	messages = append(messages, llm.NewTextMessage("user",
+		"Your tool-call budget is exhausted. This is your FINAL round. You may ONLY:\n"+
+			"- Call code_comment to submit any findings you have identified but not yet reported.\n"+
+			"- Call task_done if you have nothing more to report.\n"+
+			"No other tools are available. Do not attempt further analysis."))
+
+	if ctx.Err() != nil {
+		fmt.Fprintf(stdout.Writer(), "[ocr] Grace round skipped for %s: context cancelled\n", newPath)
+		return
+	}
+
+	resp, err := r.deps.LLMClient.CompletionsWithCtx(ctx, llm.ChatRequest{
+		Model:     r.deps.Model,
+		Messages:  messages,
+		Tools:     graceDefs,
+		MaxTokens: r.deps.Template.CompletionTokenLimit(),
+		SessionID: sessionID,
+	})
+	if err != nil {
+		fmt.Fprintf(stdout.Writer(), "[ocr] Grace round LLM error for %s: %v\n", newPath, err)
+		return
+	}
+
+	if resp.Usage != nil {
+		atomic.AddInt64(&r.totalInputTokens, resp.Usage.PromptTokens)
+		atomic.AddInt64(&r.totalOutputTokens, resp.Usage.CompletionTokens)
+		atomic.AddInt64(&r.totalCacheReadTokens, resp.Usage.CacheReadTokens)
+		atomic.AddInt64(&r.totalCacheWriteTokens, resp.Usage.CacheWriteTokens)
+	}
+
+	calls := resp.ToolCalls()
+	if len(calls) == 0 {
+		return
+	}
+
+	fs := r.deps.Session.GetOrCreateFileSession(newPath)
+	rec := fs.AppendTaskRecord(session.MainTask, append([]llm.Message(nil), messages...))
+	rec.SetResponse(resp, 0)
+	thinking := resp.ReasoningContent()
+	for _, call := range calls {
+		r.executeToolCall(ctx, newPath, call, rec, thinking)
+	}
+}
+
+// graceRoundToolDefs returns the subset of tool definitions containing only
+// code_comment and task_done.
+func graceRoundToolDefs(defs []llm.ToolDef) []llm.ToolDef {
+	out := make([]llm.ToolDef, 0, 2)
+	for _, d := range defs {
+		if d.Function.Name == "code_comment" || d.Function.Name == "task_done" {
+			out = append(out, d)
+		}
+	}
+	return out
+}
+
 // executeToolCall dispatches a single tool call from the LLM response and
 // records the result in session history. code_comment handling includes
 // optional async dispatch through CommentWorkerPool plus line-number

internal/llmloop/loop_test.go

@@ -594,3 +594,175 @@ func TestExecuteToolCall_CodeCommentOverridesHallucinatedPath(t *testing.T) {
 		t.Errorf("path override: got %q, want %q", comments[0].Path, "correct.go")
 	}
 }
+
+func graceRoundCommentResponse() *llm.ChatResponse {
+	content := ""
+	return &llm.ChatResponse{
+		Choices: []llm.Choice{{
+			Message: llm.ResponseMessage{
+				Content: &content,
+				ToolCalls: []llm.ToolCall{{
+					ID:   "call_grace",
+					Type: "function",
+					Function: llm.FunctionCall{
+						Name:      "code_comment",
+						Arguments: `{"comments":[{"content":"found a bug","existing_code":"x := 1"}]}`,
+					},
+				}},
+			},
+		}},
+		Model: "fake",
+		Usage: &llm.UsageInfo{PromptTokens: 50, CompletionTokens: 20},
+	}
+}
+
+func TestRunPerFile_GraceRoundSubmitsComment(t *testing.T) {
+	// Round 1: file_read (exhausts budget with MaxToolRequestTimes=1)
+	// Grace round: model calls code_comment
+	client := &fakeClient{responses: []*llm.ChatResponse{
+		fileReadToolCallResponse("call_1", `{"path":"main.go"}`),
+		graceRoundCommentResponse(),
+	}}
+	collector := tool.NewCommentCollector()
+	reg := tool.NewRegistry()
+	reg.Register(&fakeFileReadProvider{result: "package main\n"})
+	reg.Register(&tool.CodeCommentProvider{Collector: collector})
+	reg.Freeze()
+
+	deps := Deps{
+		LLMClient:        client,
+		Model:            "fake",
+		Template:         template.Template{MaxTokens: 100000, MaxToolRequestTimes: 1},
+		Tools:            reg,
+		CommentCollector: collector,
+		MainToolDefs: []llm.ToolDef{
+			{Type: "function", Function: llm.FunctionDef{Name: "code_comment"}},
+			{Type: "function", Function: llm.FunctionDef{Name: "task_done"}},
+			{Type: "function", Function: llm.FunctionDef{Name: "file_read"}},
+		},
+		Session: session.New("/tmp/test-repo", "main", "fake", session.SessionOptions{}),
+	}
+	runner := NewRunner(deps)
+
+	msgs := []llm.Message{llm.NewTextMessage("user", "review")}
+	completed, stop, err := runner.RunPerFile(context.Background(), msgs, "main.go")
+	if err != nil {
+		t.Fatalf("RunPerFile: %v", err)
+	}
+	if completed {
+		t.Fatal("expected not completed (budget exhausted)")
+	}
+	if stop != StopMaxRounds {
+		t.Fatalf("stop = %v, want StopMaxRounds", stop)
+	}
+	// Grace round should have been called (call 2)
+	if client.calls != 2 {
+		t.Fatalf("LLM calls = %d, want 2 (1 main + 1 grace)", client.calls)
+	}
+	// The grace round should only have code_comment + task_done tools
+	graceReq := client.requests[1]
+	if len(graceReq.Tools) != 2 {
+		t.Fatalf("grace round tools = %d, want 2", len(graceReq.Tools))
+	}
+	// Comment should have been collected
+	comments := collector.Comments()
+	if len(comments) != 1 {
+		t.Fatalf("comments = %d, want 1", len(comments))
+	}
+	if comments[0].Path != "main.go" {
+		t.Errorf("comment path = %q, want main.go", comments[0].Path)
+	}
+	// Token usage should include grace round
+	if runner.TotalInputTokens() != 70 {
+		t.Errorf("TotalInputTokens = %d, want 70 (20+50)", runner.TotalInputTokens())
+	}
+}
+
+func TestRunPerFile_GraceRoundSkippedWhenContextCancelled(t *testing.T) {
+	client := &fakeClient{responses: []*llm.ChatResponse{
+		fileReadToolCallResponse("call_1", `{"path":"main.go"}`),
+	}}
+	deps := newTestDeps(client)
+	deps.Template.MaxToolRequestTimes = 1
+	deps.MainToolDefs = []llm.ToolDef{
+		{Type: "function", Function: llm.FunctionDef{Name: "code_comment"}},
+		{Type: "function", Function: llm.FunctionDef{Name: "task_done"}},
+	}
+	runner := NewRunner(deps)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	// Cancel after the main loop exits but before grace round runs.
+	// We simulate this by using a client that cancels ctx after the first call.
+	origClient := client
+	client.responses = []*llm.ChatResponse{
+		fileReadToolCallResponse("call_1", `{"path":"main.go"}`),
+	}
+	_ = origClient
+
+	// Use a wrapper that cancels after first call
+	cancelClient := &cancelAfterNClient{inner: client, cancelAt: 1, cancel: cancel}
+	deps.LLMClient = cancelClient
+	runner = NewRunner(deps)
+
+	msgs := []llm.Message{llm.NewTextMessage("user", "review")}
+	_, stop, _ := runner.RunPerFile(ctx, msgs, "main.go")
+	if stop != StopMaxRounds {
+		t.Fatalf("stop = %v, want StopMaxRounds", stop)
+	}
+	// Grace round should have been skipped (only 1 LLM call total)
+	if cancelClient.calls != 1 {
+		t.Fatalf("LLM calls = %d, want 1 (grace skipped due to ctx cancel)", cancelClient.calls)
+	}
+}
+
+type cancelAfterNClient struct {
+	inner    *fakeClient
+	cancelAt int
+	cancel   context.CancelFunc
+	calls    int
+}
+
+func (c *cancelAfterNClient) CompletionsWithCtx(ctx context.Context, req llm.ChatRequest) (*llm.ChatResponse, error) {
+	c.calls++
+	resp, err := c.inner.CompletionsWithCtx(ctx, req)
+	if c.calls >= c.cancelAt {
+		c.cancel()
+	}
+	return resp, err
+}
+
+func TestRunPerFile_GraceRoundNotTriggeredOnEmptyRoundsStop(t *testing.T) {
+	client := &fakeClient{responses: []*llm.ChatResponse{
+		fileReadToolCallResponse("call_1", `{"path":"main.go"}`),
+		fileReadToolCallResponse("call_2", `{"path":"main.go"}`),
+		fileReadToolCallResponse("call_3", `{"path":"main.go"}`),
+	}}
+	reg := tool.NewRegistry()
+	reg.Register(&fakeFileReadProvider{result: ""})
+	deps := Deps{
+		LLMClient:        client,
+		Model:            "fake",
+		Template:         template.Template{MaxTokens: 100000, MaxToolRequestTimes: 10},
+		Tools:            reg,
+		CommentCollector: tool.NewCommentCollector(),
+		MainToolDefs: []llm.ToolDef{
+			{Type: "function", Function: llm.FunctionDef{Name: "code_comment"}},
+			{Type: "function", Function: llm.FunctionDef{Name: "task_done"}},
+		},
+		Session: session.New("/tmp/test-repo", "main", "fake", session.SessionOptions{}),
+	}
+	runner := NewRunner(deps)
+
+	msgs := []llm.Message{llm.NewTextMessage("user", "review")}
+	_, stop, err := runner.RunPerFile(context.Background(), msgs, "main.go")
+	if err != nil {
+		t.Fatalf("RunPerFile: %v", err)
+	}
+	if stop != StopEmptyRounds {
+		t.Fatalf("stop = %v, want StopEmptyRounds", stop)
+	}
+	// Should NOT trigger grace round — only 3 LLM calls (empty rounds)
+	if client.calls != 3 {
+		t.Fatalf("LLM calls = %d, want 3 (no grace round on empty-rounds stop)", client.calls)
+	}
+}