diff --git a/pkg/providers/openai_compat/provider.go b/pkg/providers/openai_compat/provider.go index d25a0fce4..915f209bb 100644 --- a/pkg/providers/openai_compat/provider.go +++ b/pkg/providers/openai_compat/provider.go @@ -316,6 +316,7 @@ func parseStreamResponse( scanner := bufio.NewScanner(reader) scanner.Buffer(make([]byte, 0, 1024*1024), 10*1024*1024) // 1MB initial, 10MB max + var pendingData string for scanner.Scan() { // Check for context cancellation between chunks if err := ctx.Err(); err != nil { @@ -332,6 +333,13 @@ func parseStreamResponse( break } + // Some SSE implementations can split a JSON payload across multiple data lines. + // Keep incomplete fragments and retry parse when the next fragment arrives. + payload := data + if pendingData != "" { + payload = pendingData + payload + } + var chunk struct { Choices []struct { Delta struct { @@ -350,9 +358,17 @@ func parseStreamResponse( Usage *UsageInfo `json:"usage"` } - if err := json.Unmarshal([]byte(data), &chunk); err != nil { - continue // skip malformed chunks + if err := json.Unmarshal([]byte(payload), &chunk); err != nil { + // Keep buffer only for likely truncated JSON. + // For other malformed payloads, drop it to avoid poisoning later valid chunks. + if strings.Contains(err.Error(), "unexpected end of JSON input") { + pendingData = payload + } else { + pendingData = "" + } + continue } + pendingData = "" if chunk.Usage != nil { usage = chunk.Usage diff --git a/pkg/providers/openai_compat/provider_test.go b/pkg/providers/openai_compat/provider_test.go index d140d63d6..df8e2e42b 100644 --- a/pkg/providers/openai_compat/provider_test.go +++ b/pkg/providers/openai_compat/provider_test.go @@ -2,6 +2,7 @@ package openai_compat import ( "bytes" + "context" "encoding/json" "fmt" "io" @@ -16,6 +17,41 @@ import ( "github.com/sipeed/picoclaw/pkg/providers/protocoltypes" ) +func TestParseStreamResponse_ReassemblesSplitJSONAcrossDataLines(t *testing.T) { + stream := strings.Join([]string{ + `data: {"choices":[{"delta":{"content":"hel`, + `data: lo"},"finish_reason":"stop"}]}`, + `data: [DONE]`, + "", + }, "\n") + + out, err := parseStreamResponse(context.Background(), strings.NewReader(stream), nil) + if err != nil { + t.Fatalf("parseStreamResponse() error = %v", err) + } + if out.Content != "hello" { + t.Fatalf("Content = %q, want %q", out.Content, "hello") + } +} + +func TestParseStreamResponse_DropsMalformedChunkAndRecovers(t *testing.T) { + stream := strings.Join([]string{ + `data: {"choices":[{"delta":{"content":"bad"}}]} garbage`, + `data: {"choices":[{"delta":{"content":"ok"},"finish_reason":"stop"}]}`, + `data: [DONE]`, + "", + }, "\n") + + out, err := parseStreamResponse(context.Background(), strings.NewReader(stream), nil) + if err != nil { + t.Fatalf("parseStreamResponse() error = %v", err) + } + // The malformed chunk should be skipped, and the next valid chunk should still parse. + if out.Content != "ok" { + t.Fatalf("Content = %q, want %q", out.Content, "ok") + } +} + func TestProviderChat_UsesMaxCompletionTokensForGLM(t *testing.T) { var requestBody map[string]any