fix(tools): stop retrying broken tool-call arguments with identical requests

clients/localai_client.go

@@ -12,6 +12,7 @@ import (
 	"sync"
 
 	"github.com/mudler/cogito"
+	"github.com/mudler/xlog"
 	"github.com/sashabaranov/go-openai"
 )
 
@@ -20,6 +21,10 @@ import (
 // the model's context window once input grows past a few turns.
 const defaultMaxTokens = 16384
 
+// maxSSELineBytes bounds one SSE line. A backend may send a tool call's whole
+// arguments in a single chunk (a file body, a patch), so the bound is generous.
+const maxSSELineBytes = 16 << 20
+
 // Ensure LocalAIClient implements cogito.LLM and cogito.StreamingLLM at compile time.
 var _ cogito.LLM = (*LocalAIClient)(nil)
 var _ cogito.StreamingLLM = (*LocalAIClient)(nil)
@@ -391,6 +396,7 @@ func (llm *LocalAIClient) CreateChatCompletion(ctx context.Context, request open
 	return cogito.LLMReply{
 		ChatCompletionResponse: response,
 		ReasoningContent:       reasoning,
+		MaxTokens:              max(request.MaxTokens, request.MaxCompletionTokens),
 	}, usage, nil
 }
 
@@ -516,8 +522,34 @@ func (llm *LocalAIClient) CreateChatCompletionStream(ctx context.Context, reques
 
 		var lastFinishReason string
 		var streamUsage *openai.Usage
+		maxTokens := request.MaxTokens
+		if request.MaxCompletionTokens > 0 {
+			maxTokens = request.MaxCompletionTokens
+		}
+		done := func() cogito.StreamEvent {
+			return cogito.StreamEvent{Type: cogito.StreamEventDone, FinishReason: lastFinishReason, Usage: usageFromOpenAI(streamUsage), MaxTokens: maxTokens}
+		}
+		// dropped counts data lines that did not decode, and sawToolCall
+		// records whether the stream produced a tool call. A dropped line may
+		// have carried a slice of a tool call's arguments, which would reach
+		// the decision incomplete but looking finished. The stream cannot know
+		// which call the line belonged to, so when both happened it ends with
+		// an ErrStreamInterrupted error instead of Done, and the decision
+		// retries the request. Content-only replies keep their Done: a lost
+		// text slice is not worth a retry.
+		dropped, sawToolCall := 0, false
+		finish := func() {
+			if dropped > 0 && sawToolCall {
+				ch <- cogito.StreamEvent{Type: cogito.StreamEventError, Error: fmt.Errorf("localai stream: dropped %d malformed chunks of a tool call: %w", dropped, cogito.ErrStreamInterrupted)}
+				return
+			}
+			ch <- done()
+		}
 
 		scanner := bufio.NewScanner(resp.Body)
+		// Some tool parsers send a call's whole arguments in one chunk, well
+		// past the scanner's 64KB default token.
+		scanner.Buffer(make([]byte, 0, 64*1024), maxSSELineBytes)
 		for scanner.Scan() {
 			line := scanner.Text()
 
@@ -527,12 +559,18 @@ func (llm *LocalAIClient) CreateChatCompletionStream(ctx context.Context, reques
 			data := strings.TrimPrefix(line, "data: ")
 
 			if data == "[DONE]" {
-				ch <- cogito.StreamEvent{Type: cogito.StreamEventDone, FinishReason: lastFinishReason, Usage: usageFromOpenAI(streamUsage)}
+				finish()
 				return
 			}
 
 			var chunk localAIStreamChunk
 			if err := json.Unmarshal([]byte(data), &chunk); err != nil {
+				dropped++
+				head := data
+				if len(head) > 200 {
+					head = head[:200]
+				}
+				xlog.Warn("localai stream: dropping a chunk that is not valid JSON", "len", len(data), "head", head, "error", err)
 				continue
 			}
 			// End the stream on an error chunk. Without this the chunk has no
@@ -566,6 +604,7 @@ func (llm *LocalAIClient) CreateChatCompletionStream(ctx context.Context, reques
 
 			// Tool call deltas
 			for _, tc := range delta.ToolCalls {
+				sawToolCall = true
 				idx := 0
 				if tc.Index != nil {
 					idx = *tc.Index
@@ -586,11 +625,18 @@ func (llm *LocalAIClient) CreateChatCompletionStream(ctx context.Context, reques
 		}
 
 		if err := scanner.Err(); err != nil {
-			ch <- cogito.StreamEvent{Type: cogito.StreamEventError, Error: err}
+			ch <- cogito.StreamEvent{Type: cogito.StreamEventError, Error: fmt.Errorf("localai stream: %w", err)}
+			return
+		}
+		// The body ended without [DONE]. Some servers omit it, so a stream
+		// that reported a finish_reason is complete. Without one the
+		// connection was cut (or a proxy timed out) mid-reply, and the partial
+		// reply must not look finished.
+		if lastFinishReason == "" {
+			ch <- cogito.StreamEvent{Type: cogito.StreamEventError, Error: fmt.Errorf("localai stream: body ended without [DONE] or a finish_reason: %w", cogito.ErrStreamInterrupted)}
 			return
 		}
-		// If we reach here without [DONE], still emit done
-		ch <- cogito.StreamEvent{Type: cogito.StreamEventDone, FinishReason: lastFinishReason, Usage: usageFromOpenAI(streamUsage)}
+		finish()
 	}()
 
 	return ch, nil

clients/localai_stream_integrity_test.go

@@ -0,0 +1,126 @@
+package clients
+
+import (
+	"context"
+	"errors"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"github.com/mudler/cogito"
+	"github.com/sashabaranov/go-openai"
+)
+
+// rawSSEServer writes the given body verbatim as an SSE stream and closes the
+// connection, so a test controls exactly how the stream ends.
+func rawSSEServer(body string) *httptest.Server {
+	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "text/event-stream")
+		_, _ = w.Write([]byte(body))
+	}))
+}
+
+// collectStream runs one streaming request against body and returns the
+// events the client emitted.
+func collectStream(t *testing.T, body string) []cogito.StreamEvent {
+	t.Helper()
+	srv := rawSSEServer(body)
+	defer srv.Close()
+	llm := NewLocalAILLM("m", "", srv.URL)
+	ch, err := llm.CreateChatCompletionStream(context.Background(), openai.ChatCompletionRequest{
+		Messages: []openai.ChatCompletionMessage{{Role: "user", Content: "hi"}},
+	})
+	if err != nil {
+		t.Fatalf("CreateChatCompletionStream: %v", err)
+	}
+	var events []cogito.StreamEvent
+	for ev := range ch {
+		events = append(events, ev)
+	}
+	if len(events) == 0 {
+		t.Fatal("stream emitted no events")
+	}
+	return events
+}
+
+func lastEvent(events []cogito.StreamEvent) cogito.StreamEvent { return events[len(events)-1] }
+
+func toolArgs(events []cogito.StreamEvent) string {
+	var b strings.Builder
+	for _, ev := range events {
+		if ev.Type == cogito.StreamEventToolCall {
+			b.WriteString(ev.ToolArgs)
+		}
+	}
+	return b.String()
+}
+
+// A connection cut mid tool call (no finish_reason, no [DONE]) must not look
+// like a finished stream: the decision would parse the partial arguments.
+func TestStreamEndedWithoutDoneOrFinishReasonIsInterrupted(t *testing.T) {
+	events := collectStream(t,
+		`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"c1","function":{"name":"search","arguments":"{\"query\":"}}]}}]}`+"\n\n")
+
+	last := lastEvent(events)
+	if last.Type != cogito.StreamEventError {
+		t.Fatalf("want an error event at the end of an interrupted stream, got %q", last.Type)
+	}
+	if !errors.Is(last.Error, cogito.ErrStreamInterrupted) {
+		t.Fatalf("want errors.Is(err, ErrStreamInterrupted), got %v", last.Error)
+	}
+}
+
+// Some servers omit [DONE]; a stream that reported a finish_reason is complete.
+func TestStreamEndedWithoutDoneButWithFinishReasonIsDone(t *testing.T) {
+	events := collectStream(t,
+		`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"c1","function":{"name":"search","arguments":"{\"query\":\"x\"}"}}]}}]}`+"\n\n"+
+			`data: {"choices":[{"delta":{},"finish_reason":"tool_calls"}]}`+"\n\n")
+
+	last := lastEvent(events)
+	if last.Type != cogito.StreamEventDone {
+		t.Fatalf("want a done event, got %q (err %v)", last.Type, last.Error)
+	}
+	if last.FinishReason != "tool_calls" {
+		t.Fatalf("want finish_reason tool_calls, got %q", last.FinishReason)
+	}
+	if last.MaxTokens != defaultMaxTokens {
+		t.Fatalf("want the done event to report the cap sent (%d), got %d", defaultMaxTokens, last.MaxTokens)
+	}
+}
+
+// Some tool parsers send a call's whole arguments in one chunk, which can be
+// far larger than bufio.Scanner's 64KB default token.
+func TestStreamAcceptsLargeSingleChunk(t *testing.T) {
+	big := strings.Repeat("a", 200*1024)
+	events := collectStream(t,
+		`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"c1","function":{"name":"write","arguments":"{\"text\":\"`+big+`\"}"}}]},"finish_reason":"tool_calls"}]}`+"\n\n"+
+			"data: [DONE]\n\n")
+
+	last := lastEvent(events)
+	if last.Type != cogito.StreamEventDone {
+		t.Fatalf("want a done event, got %q (err %v)", last.Type, last.Error)
+	}
+	if got, want := len(toolArgs(events)), len(big)+len(`{"text":""}`); got != want {
+		t.Fatalf("want %d bytes of tool arguments, got %d", want, got)
+	}
+}
+
+// A chunk that does not decode may have carried a slice of the arguments; the
+// stream must say it is lossy rather than finish with incomplete arguments.
+func TestStreamMalformedToolCallChunkIsNotSilentlyDropped(t *testing.T) {
+	events := collectStream(t,
+		`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"c1","function":{"name":"search","arguments":"{\"query\":"}}]}}]}`+"\n\n"+
+			`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\"x`+"\n\n"+
+			`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"}"}}]}}]}`+"\n\n"+
+			`data: {"choices":[{"delta":{},"finish_reason":"tool_calls"}]}`+"\n\n"+
+			"data: [DONE]\n\n")
+
+	last := lastEvent(events)
+	if last.Type != cogito.StreamEventError {
+		t.Fatalf("want an error event after a dropped tool-call chunk, got %q", last.Type)
+	}
+	if !errors.Is(last.Error, cogito.ErrStreamInterrupted) {
+		t.Fatalf("want errors.Is(err, ErrStreamInterrupted), got %v", last.Error)
+	}
+}

clients/openai_client.go

@@ -138,6 +138,7 @@ func (llm *OpenAIClient) CreateChatCompletion(ctx context.Context, request opena
 	return cogito.LLMReply{
 		ChatCompletionResponse: response,
 		ReasoningContent:       response.Choices[0].Message.ReasoningContent,
+		MaxTokens:              max(request.MaxTokens, request.MaxCompletionTokens),
 	}, usage, nil
 }
 
@@ -181,8 +182,12 @@ func (llm *OpenAIClient) CreateChatCompletionStream(ctx context.Context, request
 
 		for {
 			resp, err := stream.Recv()
+			// go-openai returns io.EOF both after [DONE] and when the body
+			// ends without it, so a cut connection cannot be told from a
+			// finished stream here, and a server that sends [DONE] without a
+			// finish_reason is valid. This path is left as it was.
 			if errors.Is(err, io.EOF) {
-				ch <- cogito.StreamEvent{Type: cogito.StreamEventDone, FinishReason: lastFinishReason, Usage: usageFromOpenAI(streamUsage)}
+				ch <- cogito.StreamEvent{Type: cogito.StreamEventDone, FinishReason: lastFinishReason, Usage: usageFromOpenAI(streamUsage), MaxTokens: max(request.MaxTokens, request.MaxCompletionTokens)}
 				return
 			}
 			if err != nil {

decision_length_retry_test.go

@@ -191,9 +191,10 @@ func TestStreamingDecisionLengthRetryCountsFailedAttemptUsage(t *testing.T) {
 
 // scriptedChatLLM is the non-streaming counterpart of scriptedStreamLLM.
 type scriptedChatLLM struct {
-	replies  []openai.ChatCompletionChoice
-	usages   []LLMUsage
-	requests []openai.ChatCompletionRequest
+	replies   []openai.ChatCompletionChoice
+	usages    []LLMUsage
+	reasoning string // ReasoningContent of every reply
+	requests  []openai.ChatCompletionRequest
 }
 
 func (m *scriptedChatLLM) Ask(ctx context.Context, f Fragment) (Fragment, error) { return f, nil }
@@ -204,7 +205,7 @@ func (m *scriptedChatLLM) CreateChatCompletion(ctx context.Context, request open
 	if i >= len(m.replies) {
 		i = len(m.replies) - 1
 	}
-	return LLMReply{ChatCompletionResponse: openai.ChatCompletionResponse{Choices: []openai.ChatCompletionChoice{m.replies[i]}}}, m.usages[i], nil
+	return LLMReply{ChatCompletionResponse: openai.ChatCompletionResponse{Choices: []openai.ChatCompletionChoice{m.replies[i]}}, ReasoningContent: m.reasoning}, m.usages[i], nil
 }
 
 func lengthChoice() openai.ChatCompletionChoice {