fix(provider): handle split SSE JSON chunks safely

Preserve partial stream payloads only when JSON is incomplete so chunk boundaries no longer drop content, while malformed payloads are discarded to avoid poisoning subsequent valid chunks. Add regression tests for split-payload reassembly and malformed-chunk recovery.

Made-with: Cursor
This commit is contained in:
imalasong 2026-04-07 22:05:05 +08:00
parent 7bf6cbe1fa
commit 49309d8911
2 changed files with 54 additions and 2 deletions

View file

@ -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

View file

@ -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