From bcaefadcdea888a32d29554d68bf7c647ce4935f Mon Sep 17 00:00:00 2001 From: kite Date: Wed, 12 Aug 2026 17:52:09 +0800 Subject: [PATCH 1/3] feat(llmloop): add grace round after tool-request budget exhausted When RunPerFile exits because MaxToolRequestTimes reaches zero, perform one additional LLM call with only code_comment and task_done available. This gives the model a final chance to submit findings it identified but had not yet reported, preventing loss of review comments on budget stop. --- internal/llmloop/loop.go | 60 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 60 insertions(+) diff --git a/internal/llmloop/loop.go b/internal/llmloop/loop.go index cf99fd1f..d2aa687c 100644 --- a/internal/llmloop/loop.go +++ b/internal/llmloop/loop.go @@ -303,10 +303,70 @@ 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.")) + + 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, nil) + 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 From 2c42f8a7398a3c051ca83abd5e038e7d1167a66d Mon Sep 17 00:00:00 2001 From: kite Date: Wed, 12 Aug 2026 18:37:45 +0800 Subject: [PATCH 2/3] fix(llmloop): address review comments on grace round - Check ctx.Err() before making the grace round LLM call to avoid wasted API calls when the context is already cancelled. - Pass messages copy to AppendTaskRecord and call rec.SetResponse so the grace round interaction is visible in session/debug logs. --- internal/llmloop/loop.go | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/internal/llmloop/loop.go b/internal/llmloop/loop.go index d2aa687c..28726622 100644 --- a/internal/llmloop/loop.go +++ b/internal/llmloop/loop.go @@ -323,6 +323,11 @@ func (r *Runner) runGraceRound(ctx context.Context, messages []llm.Message, newP "- 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, @@ -348,7 +353,8 @@ func (r *Runner) runGraceRound(ctx context.Context, messages []llm.Message, newP } fs := r.deps.Session.GetOrCreateFileSession(newPath) - rec := fs.AppendTaskRecord(session.MainTask, nil) + 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) From 305d7f4386d0c103049f4d3fb617c2083c028bc0 Mon Sep 17 00:00:00 2001 From: kite Date: Wed, 12 Aug 2026 18:40:24 +0800 Subject: [PATCH 3/3] test(llmloop): add unit tests for grace round Cover three scenarios: - Grace round fires and collects code_comment on budget exhaustion - Grace round is skipped when context is already cancelled - Grace round is NOT triggered on StopEmptyRounds --- internal/llmloop/loop_test.go | 172 ++++++++++++++++++++++++++++++++++ 1 file changed, 172 insertions(+) diff --git a/internal/llmloop/loop_test.go b/internal/llmloop/loop_test.go index 5a28dc52..46a40ada 100644 --- a/internal/llmloop/loop_test.go +++ b/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) + } +}