Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
lizhengfeng101 merged 3 commits into from
Aug 13, 2026
Merged

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

Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
bcaefad
feat(llmloop): add grace round after tool-request budget exhausted
lizhengfeng101 Aug 12, 2026
2c42f8a
fix(llmloop): address review comments on grace round
lizhengfeng101 Aug 12, 2026
305d7f4
test(llmloop): add unit tests for grace round
lizhengfeng101 Aug 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
60 changes: 60 additions & 0 deletions internal/llmloop/loop.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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{
Comment thread
lizhengfeng101 marked this conversation as resolved.
Model: r.deps.Model,
Messages: messages,
Tools: graceDefs,

@stay-foolish-forever stay-foolish-forever Aug 13, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Restricting the tool list for this final call will likely cost us the whole KV cache prefix.

Placing the tool list in or next to the system prompt is the mainstream practice among current LLMs, precisely so that the prefix cache keeps hitting across a session. Swapping r.deps.MainToolDefs for the two-tool subset rewrites that part of the prefix, so everything behind it is invalidated.

The timing makes it expensive: runGraceRound only fires on StopMaxRounds, i.e. when the conversation has accumulated a full budget's worth of tool_result file contents and is at its longest. We re-read that prompt at full price and write a fresh cache entry that nothing downstream will ever read. The new totalCacheReadTokens/totalCacheWriteTokens counters should show it directly: read ≈ 0, write ≈ the whole prompt.

That said, there's a genuine trade-off here between effectiveness and cost. Narrowing the tool list is a hard constraint, whereas leaving MainToolDefs intact and relying on the prose instruction (optionally filtering resp.ToolCalls() on the response side) is only a soft one — the model could spend its last turn on a read tool. If the temporary override does measurably improve how often the grace round lands a real finding, then this may simply be a price we have to pay.

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)
}
Comment thread
lizhengfeng101 marked this conversation as resolved.
}

// 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
Expand Down
Loading