fix(agent): last-iteration false failure, cached schemas, usage in tool rounds

- Run() reported "max iterations reached" even when a valid final answer
  arrived exactly on the last allowed iteration, throwing the response
  away; a completed flag now distinguishes success from budget exhaustion.
- Tool schemas are marshaled once per Run/RunStream instead of once per
  loop iteration — they never change between iterations.
- RunStream's content gate (which hides raw deltas during a tool-call
  round) also swallowed that round's token usage, so callers only ever saw
  the final round's count and context tracking lagged exactly when the
  context grew fastest. Usage is now forwarded in its own chunk.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Victor Hugo Vargas Servin 2026-07-12 16:14:15 -07:00
parent e6830161bc
commit 2e23216932

View file

@ -82,17 +82,21 @@ func (l *Loop) Run(ctx context.Context, input string, history ...llm.Message) (R
start := time.Now() start := time.Now()
messages := l.buildInitialMessages(input, history) messages := l.buildInitialMessages(input, history)
// Tool schemas don't change between iterations, so build the JSON once
// per Run instead of re-marshaling every tool on every loop pass.
toolSchemas := l.getToolSchemas()
var finalContent string var finalContent string
var allToolCalls []llm.ToolCall var allToolCalls []llm.ToolCall
var totalUsage llm.TokenUsage var totalUsage llm.TokenUsage
iterations := 0 iterations := 0
completed := false
for iterations < l.cfg.MaxIters { for iterations < l.cfg.MaxIters {
iterations++ iterations++
resp, err := l.cfg.LLM.Generate(ctx, llm.CompletionRequest{ resp, err := l.cfg.LLM.Generate(ctx, llm.CompletionRequest{
Messages: messages, Messages: messages,
Tools: l.getToolSchemas(), Tools: toolSchemas,
ChatTemplateKwargs: l.cfg.ChatTemplateKwargs, ChatTemplateKwargs: l.cfg.ChatTemplateKwargs,
}) })
if err != nil { if err != nil {
@ -105,6 +109,7 @@ func (l *Loop) Run(ctx context.Context, input string, history ...llm.Message) (R
if len(resp.ToolCalls) == 0 { if len(resp.ToolCalls) == 0 {
finalContent = resp.Content finalContent = resp.Content
completed = true
break break
} }
@ -138,7 +143,11 @@ func (l *Loop) Run(ctx context.Context, input string, history ...llm.Message) (R
} }
duration := time.Since(start) duration := time.Since(start)
if iterations >= l.cfg.MaxIters { // Only report the max-iterations failure when the loop actually ran out
// of budget without producing a final answer — an answer that arrives
// exactly on the last allowed iteration is still a success (the old
// `iterations >= MaxIters` check threw that valid response away).
if !completed {
return Response{}, fmt.Errorf("max iterations (%d) reached", l.cfg.MaxIters) return Response{}, fmt.Errorf("max iterations (%d) reached", l.cfg.MaxIters)
} }
@ -157,6 +166,8 @@ func (l *Loop) Run(ctx context.Context, input string, history ...llm.Message) (R
func (l *Loop) RunStream(ctx context.Context, input string, history ...llm.Message) iter.Seq2[llm.StreamChunk, error] { func (l *Loop) RunStream(ctx context.Context, input string, history ...llm.Message) iter.Seq2[llm.StreamChunk, error] {
return func(yield func(llm.StreamChunk, error) bool) { return func(yield func(llm.StreamChunk, error) bool) {
messages := l.buildInitialMessages(input, history) messages := l.buildInitialMessages(input, history)
// Same as Run: the schemas are identical on every iteration.
toolSchemas := l.getToolSchemas()
iterations := 0 iterations := 0
for iterations < l.cfg.MaxIters { for iterations < l.cfg.MaxIters {
@ -164,7 +175,7 @@ func (l *Loop) RunStream(ctx context.Context, input string, history ...llm.Messa
stream := l.cfg.LLM.Stream(ctx, llm.CompletionRequest{ stream := l.cfg.LLM.Stream(ctx, llm.CompletionRequest{
Messages: messages, Messages: messages,
Tools: l.getToolSchemas(), Tools: toolSchemas,
ChatTemplateKwargs: l.cfg.ChatTemplateKwargs, ChatTemplateKwargs: l.cfg.ChatTemplateKwargs,
}) })
@ -226,11 +237,23 @@ func (l *Loop) RunStream(ctx context.Context, input string, history ...llm.Messa
// real token counts. // real token counts.
hasContent := chunk.Delta != "" || chunk.ReasoningDelta != "" hasContent := chunk.Delta != "" || chunk.ReasoningDelta != ""
hasUsage := chunk.Usage.TotalTokens > 0 hasUsage := chunk.Usage.TotalTokens > 0
if !hasToolCalls && (hasContent || hasUsage) { switch {
case !hasToolCalls && (hasContent || hasUsage):
responseBuilder.WriteString(chunk.Delta) responseBuilder.WriteString(chunk.Delta)
if !yield(chunk, nil) { if !yield(chunk, nil) {
return return
} }
case hasToolCalls && hasUsage:
// The content gate above exists to hide raw provider
// deltas during a tool-call round, but it also swallowed
// that round's token usage — so callers tracking context
// occupancy (e.g. a UI's context bar deciding when to
// compact) only ever saw the usage of the final,
// tool-free round. Forward the usage on its own,
// without the content.
if !yield(llm.StreamChunk{Usage: chunk.Usage}, nil) {
return
}
} }
} }