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 := bufio.NewScanner(reader)
scanner.Buffer(make([]byte, 0, 1024*1024), 10*1024*1024) // 1MB initial, 10MB max scanner.Buffer(make([]byte, 0, 1024*1024), 10*1024*1024) // 1MB initial, 10MB max
var pendingData string
for scanner.Scan() { for scanner.Scan() {
// Check for context cancellation between chunks // Check for context cancellation between chunks
if err := ctx.Err(); err != nil { if err := ctx.Err(); err != nil {
@ -332,6 +333,13 @@ func parseStreamResponse(
break 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 { var chunk struct {
Choices []struct { Choices []struct {
Delta struct { Delta struct {
@ -350,9 +358,17 @@ func parseStreamResponse(
Usage *UsageInfo `json:"usage"` Usage *UsageInfo `json:"usage"`
} }
if err := json.Unmarshal([]byte(data), &chunk); err != nil { if err := json.Unmarshal([]byte(payload), &chunk); err != nil {
continue // skip malformed chunks // 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 { if chunk.Usage != nil {
usage = chunk.Usage usage = chunk.Usage

View file

@ -2,6 +2,7 @@ package openai_compat
import ( import (
"bytes" "bytes"
"context"
"encoding/json" "encoding/json"
"fmt" "fmt"
"io" "io"
@ -16,6 +17,41 @@ import (
"github.com/sipeed/picoclaw/pkg/providers/protocoltypes" "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) { func TestProviderChat_UsesMaxCompletionTokensForGLM(t *testing.T) {
var requestBody map[string]any var requestBody map[string]any