From 49309d8911e9386847ea49074388ab70bc08d957 Mon Sep 17 00:00:00 2001
From: imalasong <imalasong@qq.com>
Date: Tue, 7 Apr 2026 22:05:05 +0800
Subject: [PATCH] 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
---
 pkg/providers/openai_compat/provider.go      | 20 +++++++++--
 pkg/providers/openai_compat/provider_test.go | 36 ++++++++++++++++++++
 2 files changed, 54 insertions(+), 2 deletions(-)

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