From 2e232169329e444b9bf9a0f0e19fc97bf7cf1046 Mon Sep 17 00:00:00 2001
From: Victor Vargas <victorhvargas@outlook.com>
Date: Sun, 12 Jul 2026 16:14:15 -0700
Subject: [PATCH] fix(agent): last-iteration false failure, cached schemas,
 usage in tool rounds
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- 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>
---
 pkg/agent/loop.go | 59 ++++++++++++++++++++++++++++++++---------------
 1 file changed, 41 insertions(+), 18 deletions(-)

diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go
index 59dc93e..f73d882 100644
--- a/pkg/agent/loop.go
+++ b/pkg/agent/loop.go
@@ -47,19 +47,19 @@ type Config struct {
 
 // Iteration represents a single cycle of the agent loop.
 type Iteration struct {
-	Number     int
-	ToolCalls  []llm.ToolCall
-	ToolsUsed  int
-	Duration   time.Duration
+	Number    int
+	ToolCalls []llm.ToolCall
+	ToolsUsed int
+	Duration  time.Duration
 }
 
 // Response is the final output of the agent loop.
 type Response struct {
-	Content     string
-	ToolCalls   []llm.ToolCall
-	Iterations  int
-	Duration    time.Duration
-	TokenUsage  llm.TokenUsage
+	Content    string
+	ToolCalls  []llm.ToolCall
+	Iterations int
+	Duration   time.Duration
+	TokenUsage llm.TokenUsage
 }
 
 // Loop is the main agent loop that orchestrates LLM calls and tool execution.
@@ -82,17 +82,21 @@ func (l *Loop) Run(ctx context.Context, input string, history ...llm.Message) (R
 	start := time.Now()
 
 	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 allToolCalls []llm.ToolCall
 	var totalUsage llm.TokenUsage
 	iterations := 0
+	completed := false
 
 	for iterations < l.cfg.MaxIters {
 		iterations++
 
 		resp, err := l.cfg.LLM.Generate(ctx, llm.CompletionRequest{
 			Messages:           messages,
-			Tools:              l.getToolSchemas(),
+			Tools:              toolSchemas,
 			ChatTemplateKwargs: l.cfg.ChatTemplateKwargs,
 		})
 		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 {
 			finalContent = resp.Content
+			completed = true
 			break
 		}
 
@@ -138,16 +143,20 @@ func (l *Loop) Run(ctx context.Context, input string, history ...llm.Message) (R
 	}
 
 	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{
-		Content:     finalContent,
-		ToolCalls:   allToolCalls,
-		Iterations:  iterations,
-		Duration:    duration,
-		TokenUsage:  totalUsage,
+		Content:    finalContent,
+		ToolCalls:  allToolCalls,
+		Iterations: iterations,
+		Duration:   duration,
+		TokenUsage: totalUsage,
 	}, nil
 }
 
@@ -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] {
 	return func(yield func(llm.StreamChunk, error) bool) {
 		messages := l.buildInitialMessages(input, history)
+		// Same as Run: the schemas are identical on every iteration.
+		toolSchemas := l.getToolSchemas()
 		iterations := 0
 
 		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{
 				Messages:           messages,
-				Tools:              l.getToolSchemas(),
+				Tools:              toolSchemas,
 				ChatTemplateKwargs: l.cfg.ChatTemplateKwargs,
 			})
 
@@ -226,11 +237,23 @@ func (l *Loop) RunStream(ctx context.Context, input string, history ...llm.Messa
 				// real token counts.
 				hasContent := chunk.Delta != "" || chunk.ReasoningDelta != ""
 				hasUsage := chunk.Usage.TotalTokens > 0
-				if !hasToolCalls && (hasContent || hasUsage) {
+				switch {
+				case !hasToolCalls && (hasContent || hasUsage):
 					responseBuilder.WriteString(chunk.Delta)
 					if !yield(chunk, nil) {
 						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
+					}
 				}
 			}