Merge pull request #8 from VictorVargas/fix/providers-agent-loop

Fix providers (OpenAI/llamacpp/Anthropic) and agent loop correctness

pkg/agent/loop.go

@@ -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,7 +143,11 @@ 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)
 	}
 
@@ -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
+					}
 				}
 			}
 

pkg/llm/providers/anthropic/client.go

@@ -7,8 +7,8 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
-	"net/http"
 	"iter"
+	"net/http"
 	"strings"
 
 	"github.com/VictorVargas/rony-llm-agent/pkg/llm"
@@ -19,6 +19,13 @@ const (
 	defaultModel     = "claude-opus-4-8"
 	defaultMaxTokens = 8192
 	anthropicVersion = "2023-06-01"
+	// defaultContextWindow is what Capabilities() reports when
+	// Config.ContextWindow is unset: 200k tokens, the standard window for
+	// Claude models. Callers use this number to decide when to compact
+	// their conversation, so over-reporting it (the old code assumed 1M
+	// for anything that wasn't haiku) meant compaction fired far too late
+	// and requests started overflowing the real window.
+	defaultContextWindow = 200000
 )
 
 // Config holds the settings needed to create an Anthropic client.
@@ -27,6 +34,7 @@ type Config struct {
 	Model         string // defaults to claude-opus-4-8
 	BaseURL       string // defaults to https://api.anthropic.com/v1
 	MaxTokens     int    // default max_tokens sent on every request (Anthropic requires one); 0 = defaultMaxTokens
+	ContextWindow int    // model's context window in tokens (0 = defaultContextWindow); raise it only for models/plans with an extended window
 	Temperature   *float32
 	TopP          *float32
 }
@@ -37,6 +45,7 @@ type Client struct {
 	baseURL       string
 	model         string
 	maxTokens     int
+	contextWindow int
 	temperature   *float32
 	topP          *float32
 	http          *http.Client
@@ -63,11 +72,17 @@ func New(cfg Config) (*Client, error) {
 		maxTokens = defaultMaxTokens
 	}
 
+	contextWindow := cfg.ContextWindow
+	if contextWindow == 0 {
+		contextWindow = defaultContextWindow
+	}
+
 	return &Client{
 		apiKey:        cfg.APIKey,
 		baseURL:       baseURL,
 		model:         model,
 		maxTokens:     maxTokens,
+		contextWindow: contextWindow,
 		temperature:   cfg.Temperature,
 		topP:          cfg.TopP,
 		http:          http.DefaultClient,
@@ -257,15 +272,11 @@ func (c *Client) Name() string {
 }
 
 func (c *Client) Capabilities() llm.ProviderCapabilities {
-	maxContext := 1000000
-	if strings.Contains(c.model, "haiku") {
-		maxContext = 200000
-	}
 	return llm.ProviderCapabilities{
 		SupportsTools:    true,
 		SupportsVision:   true,
 		SupportsJSON:     true,
-		MaxContextWindow: maxContext,
+		MaxContextWindow: c.contextWindow,
 	}
 }
 

pkg/llm/providers/llamacpp/client.go

@@ -2,12 +2,13 @@ package llamacpp
 
 import (
 	"bufio"
+	"bytes"
 	"context"
 	"encoding/json"
 	"fmt"
 	"io"
-	"net/http"
 	"iter"
+	"net/http"
 	"strings"
 	"time"
 
@@ -123,7 +124,7 @@ func (c *Client) Generate(ctx context.Context, req llm.CompletionRequest) (llm.C
 		return llm.CompletionResponse{}, fmt.Errorf("decoding response: %w", err)
 	}
 
-	return c.toResponse(apiResp), nil
+	return c.toResponse(apiResp)
 }
 
 func (c *Client) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] {
@@ -191,6 +192,11 @@ func (c *Client) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq
 		}
 
 		scanner := bufio.NewScanner(resp.Body)
+		// A single SSE line can exceed bufio.Scanner's 64KB default cap
+		// (e.g. a large tool-call arguments delta or a long reasoning
+		// event), which would kill the stream mid-turn with "token too
+		// long" — same headroom the anthropic client already reserves.
+		scanner.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
 		for scanner.Scan() {
 			line := scanner.Text()
 			if !strings.HasPrefix(line, "data: ") {
@@ -370,11 +376,16 @@ func (c *Client) buildRequest(req llm.CompletionRequest, stream bool) (io.Reader
 	if err != nil {
 		return nil, fmt.Errorf("marshaling request: %w", err)
 	}
-	return strings.NewReader(string(data)), nil
+	return bytes.NewReader(data), nil
 }
 
 // toResponse converts a llama.cpp API response to our CompletionResponse.
-func (c *Client) toResponse(resp llamaChatResponse) llm.CompletionResponse {
+func (c *Client) toResponse(resp llamaChatResponse) (llm.CompletionResponse, error) {
+	// Guard against a 200 response with no choices (e.g. a misbehaving
+	// server or proxy) — indexing Choices[0] blindly panics the whole app.
+	if len(resp.Choices) == 0 {
+		return llm.CompletionResponse{}, fmt.Errorf("llamacpp: response contained no choices")
+	}
 	choice := resp.Choices[0]
 	result := llm.CompletionResponse{
 		ID:         resp.ID,
@@ -398,7 +409,7 @@ func (c *Client) toResponse(resp llamaChatResponse) llm.CompletionResponse {
 		TotalTokens:  resp.Usage.TotalTokens,
 	}
 
-	return result
+	return result, nil
 }
 
 // llama.cpp API types

pkg/llm/providers/openai/client.go

@@ -2,12 +2,13 @@ package openai
 
 import (
 	"bufio"
+	"bytes"
 	"context"
 	"encoding/json"
 	"fmt"
 	"io"
-	"net/http"
 	"iter"
+	"net/http"
 	"strings"
 
 	"github.com/VictorVargas/rony-llm-agent/pkg/llm"
@@ -23,6 +24,7 @@ type Config struct {
 // Client implements llm.LLMClient for OpenAI.
 type Client struct {
 	apiKey  string
+	model   string
 	baseURL string
 	http    *http.Client
 }
@@ -40,6 +42,7 @@ func New(cfg Config) (*Client, error) {
 
 	return &Client{
 		apiKey:  cfg.APIKey,
+		model:   cfg.Model,
 		baseURL: baseURL,
 		http:    http.DefaultClient,
 	}, nil
@@ -48,7 +51,7 @@ func New(cfg Config) (*Client, error) {
 func (c *Client) Generate(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) {
 	endpoint := c.baseURL + "/chat/completions"
 
-	payload, err := c.buildRequest(req)
+	payload, err := c.buildRequest(req, false)
 	if err != nil {
 		return llm.CompletionResponse{}, fmt.Errorf("building request: %w", err)
 	}
@@ -76,14 +79,14 @@ func (c *Client) Generate(ctx context.Context, req llm.CompletionRequest) (llm.C
 		return llm.CompletionResponse{}, fmt.Errorf("decoding response: %w", err)
 	}
 
-	return c.toResponse(apiResp), nil
+	return c.toResponse(apiResp)
 }
 
 func (c *Client) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] {
 	return func(yield func(llm.StreamChunk, error) bool) {
 		endpoint := c.baseURL + "/chat/completions"
 
-		payload, err := c.buildRequest(req)
+		payload, err := c.buildRequest(req, true)
 		if err != nil {
 			yield(llm.StreamChunk{}, fmt.Errorf("building request: %w", err))
 			return
@@ -111,7 +114,45 @@ func (c *Client) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq
 			return
 		}
 
+		// toolCallAccum buffers one tool call's fragments as they stream in:
+		// the SSE format sends the id/name in the first delta for a given
+		// tool-call index and the (potentially large) arguments JSON in
+		// pieces across many subsequent deltas, so it can't be handed to a
+		// tool handler until it's fully assembled. Same accumulation the
+		// llamacpp client does — without it, tool calls made over a stream
+		// were silently dropped and the agent loop never executed them.
+		type toolCallAccum struct {
+			id   string
+			name string
+			args strings.Builder
+		}
+		toolCallFrags := map[int]*toolCallAccum{}
+		var toolCallOrder []int
+
+		flushToolCalls := func() []llm.ToolCall {
+			if len(toolCallOrder) == 0 {
+				return nil
+			}
+			calls := make([]llm.ToolCall, 0, len(toolCallOrder))
+			for _, idx := range toolCallOrder {
+				frag := toolCallFrags[idx]
+				calls = append(calls, llm.ToolCall{
+					ID:        frag.id,
+					Name:      frag.name,
+					Arguments: json.RawMessage(frag.args.String()),
+				})
+			}
+			toolCallFrags = map[int]*toolCallAccum{}
+			toolCallOrder = nil
+			return calls
+		}
+
 		scanner := bufio.NewScanner(resp.Body)
+		// A single SSE line can exceed bufio.Scanner's 64KB default cap
+		// (e.g. a large tool-call arguments delta), which would kill the
+		// stream with "token too long" — same headroom the anthropic
+		// client already reserves.
+		scanner.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
 		for scanner.Scan() {
 			line := scanner.Text()
 			if !strings.HasPrefix(line, "data: ") {
@@ -128,13 +169,61 @@ func (c *Client) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq
 				return
 			}
 
+			var usage llm.TokenUsage
+			if event.Usage != nil {
+				usage = llm.TokenUsage{
+					InputTokens:  event.Usage.PromptTokens,
+					OutputTokens: event.Usage.CompletionTokens,
+					TotalTokens:  event.Usage.TotalTokens,
+				}
+			}
+
+			if len(event.Choices) == 0 {
+				// The usage-only event (per stream_options.include_usage)
+				// carries no choices, so it needs its own chunk.
+				if event.Usage != nil {
+					if !yield(llm.StreamChunk{Usage: usage}, nil) {
+						return
+					}
+				}
+				continue
+			}
+
 			for _, choice := range event.Choices {
+				hasFragment := len(choice.Delta.ToolCalls) > 0
+				for _, tc := range choice.Delta.ToolCalls {
+					frag, ok := toolCallFrags[tc.Index]
+					if !ok {
+						frag = &toolCallAccum{}
+						toolCallFrags[tc.Index] = frag
+						toolCallOrder = append(toolCallOrder, tc.Index)
+					}
+					if tc.ID != "" {
+						frag.id = tc.ID
+					}
+					if tc.Function.Name != "" {
+						frag.name = tc.Function.Name
+					}
+					frag.args.WriteString(tc.Function.Arguments)
+				}
+
 				chunk := llm.StreamChunk{
 					Delta: choice.Delta.Content,
+					Usage: usage,
 				}
 				if choice.FinishReason != "" {
 					chunk.FinishReason = choice.FinishReason
+					chunk.ToolCalls = flushToolCalls()
 				}
+
+				// A fragment-only event (a piece of a tool call's streamed
+				// arguments, with nothing else in this delta) has nothing
+				// yet for the agent loop to act on: it was buffered above,
+				// so skip yielding an empty chunk for it.
+				if hasFragment && chunk.Delta == "" && chunk.FinishReason == "" {
+					continue
+				}
+
 				if !yield(chunk, nil) {
 					return
 				}
@@ -160,7 +249,7 @@ func (c *Client) Capabilities() llm.ProviderCapabilities {
 }
 
 // buildRequest converts an llm.CompletionRequest to the OpenAI API format.
-func (c *Client) buildRequest(req llm.CompletionRequest) (io.Reader, error) {
+func (c *Client) buildRequest(req llm.CompletionRequest, stream bool) (io.Reader, error) {
 	// Convert messages to OpenAI format
 	messages := make([]openaiMessage, len(req.Messages))
 	for i, m := range req.Messages {
@@ -195,10 +284,25 @@ func (c *Client) buildRequest(req llm.CompletionRequest) (io.Reader, error) {
 		tools[i] = tool
 	}
 
+	// The per-request model wins when set; otherwise fall back to the
+	// client's configured one (Config.Model used to be discarded entirely,
+	// so every request went out with an empty model — a hard API error on
+	// OpenAI, and the agent loop never sets req.Model).
+	model := req.Model
+	if model == "" {
+		model = c.model
+	}
+
 	openaiReq := openaiChatRequest{
-		Model:    req.Model,
+		Model:    model,
 		Messages: messages,
-		Stream:   false,
+		Stream:   stream,
+	}
+	if stream {
+		// Ask for a final SSE event carrying token usage (OpenAI-style
+		// streaming omits it otherwise), so callers can track real token
+		// counts per turn instead of always seeing zero.
+		openaiReq.StreamOptions = &openaiStreamOptions{IncludeUsage: true}
 	}
 	if len(tools) > 0 {
 		openaiReq.Tools = tools
@@ -222,11 +326,14 @@ func (c *Client) buildRequest(req llm.CompletionRequest) (io.Reader, error) {
 	if err != nil {
 		return nil, fmt.Errorf("marshaling request: %w", err)
 	}
-	return strings.NewReader(string(data)), nil
+	return bytes.NewReader(data), nil
 }
 
 // toResponse converts an OpenAI API response to our CompletionResponse.
-func (c *Client) toResponse(resp openaiChatResponse) llm.CompletionResponse {
+func (c *Client) toResponse(resp openaiChatResponse) (llm.CompletionResponse, error) {
+	if len(resp.Choices) == 0 {
+		return llm.CompletionResponse{}, fmt.Errorf("openai: response contained no choices")
+	}
 	choice := resp.Choices[0]
 	result := llm.CompletionResponse{
 		ID:         resp.ID,
@@ -249,7 +356,7 @@ func (c *Client) toResponse(resp openaiChatResponse) llm.CompletionResponse {
 		TotalTokens:  resp.Usage.TotalTokens,
 	}
 
-	return result
+	return result, nil
 }
 
 // OpenAI API types
@@ -263,6 +370,11 @@ type openaiChatRequest struct {
 	MaxTokens     *int                 `json:"max_tokens,omitempty"`
 	Stop          []string             `json:"stop,omitempty"`
 	Stream        bool                 `json:"stream"`
+	StreamOptions *openaiStreamOptions `json:"stream_options,omitempty"`
+}
+
+type openaiStreamOptions struct {
+	IncludeUsage bool `json:"include_usage"`
 }
 
 type openaiMessage struct {
@@ -319,6 +431,7 @@ type openaiUsage struct {
 type openaiStreamEvent struct {
 	ID      string               `json:"id"`
 	Choices []openaiStreamChoice `json:"choices"`
+	Usage   *openaiUsage         `json:"usage"`
 }
 
 type openaiStreamChoice struct {

pkg/llm/types_test.go

@@ -139,7 +139,5 @@ func TestStreamChunk_JSON(t *testing.T) {
 	}
 }
 
-
-
 func float32Ptr(f float32) *float32 { return &f }
 func intPtr(i int) *int             { return &i }