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 + } } }