fix: enforce SSE terminal event boundaries

proxy/sse_event.go

@@ -8,7 +8,9 @@ import (
 
 type sseEventTracker struct {
 	buf           []byte
+	recognized    bool
 	terminal      bool
+	afterTerminal bool
 }
 
 func (t *sseEventTracker) Write(p []byte) {
@@ -20,29 +22,39 @@ func (t *sseEventTracker) Write(p []byte) {
 		}
 		event := t.buf[:end]
 		t.buf = t.buf[end+separator:]
-		if terminalSSEEvent(event) {
+		if t.terminal {
+			t.afterTerminal = true
+		}
+		recognized, terminal := classifySSEEvent(event)
+		t.recognized = t.recognized || recognized
+		if terminal {
 			t.terminal = true
 		}
 	}
 }
 
-func (t *sseEventTracker) Complete() bool { return len(t.buf) == 0 }
+func (t *sseEventTracker) Complete() bool {
+	return len(t.buf) == 0 && !t.afterTerminal && (!t.recognized || t.terminal)
+}
 
 func completeSSEEvent(buf []byte) (int, int) {
 	lf := bytes.Index(buf, []byte("\n\n"))
 	crlf := bytes.Index(buf, []byte("\r\n\r\n"))
-	if crlf >= 0 && (lf < 0 || crlf < lf) {
+	cr := bytes.Index(buf, []byte("\r\r"))
-		return crlf, 4
+	end, separator := lf, 2
+	if crlf >= 0 && (end < 0 || crlf < end) {
+		end, separator = crlf, 4
 	}
-	if lf >= 0 {
+	if cr >= 0 && (end < 0 || cr < end) {
-		return lf, 2
+		end, separator = cr, 2
 	}
-	return -1, 0
+	return end, separator
 }
 
-func terminalSSEEvent(event []byte) bool {
+func classifySSEEvent(event []byte) (recognized, terminal bool) {
 	var data strings.Builder
-	for _, line := range strings.Split(strings.ReplaceAll(string(event), "\r\n", "\n"), "\n") {
+	normalized := strings.ReplaceAll(string(event), "\r\n", "\n")
+	for _, line := range strings.Split(strings.ReplaceAll(normalized, "\r", "\n"), "\n") {
 		if !strings.HasPrefix(line, "data:") {
 			continue
 		}
@@ -53,13 +65,31 @@ func terminalSSEEvent(event []byte) bool {
 	}
 	payload := data.String()
 	if payload == "[DONE]" {
-		return true
+		return true, true
 	}
 	var envelope struct {
 		Type       string          `json:"type"`
+		Choices    json.RawMessage `json:"choices"`
+		Candidates []struct {
+			FinishReason string `json:"finishReason"`
+		} `json:"candidates"`
 	}
 	if json.Unmarshal([]byte(payload), &envelope) != nil {
-		return false
+		return false, false
 	}
-	return envelope.Type == "message_stop" || envelope.Type == "response.completed"
+	if envelope.Type != "" {
+		return true, envelope.Type == "message_stop" || envelope.Type == "response.completed"
+	}
+	if envelope.Choices != nil {
+		return true, false
+	}
+	if envelope.Candidates != nil {
+		for _, candidate := range envelope.Candidates {
+			if candidate.FinishReason != "" {
+				return true, true
+			}
+		}
+		return true, false
+	}
+	return false, false
 }

proxy/sse_event_test.go

@@ -0,0 +1,21 @@
+package proxy
+
+import "testing"
+
+func TestSSEEventTrackerRejectsEventAfterTerminalEvent(t *testing.T) {
+	var tracker sseEventTracker
+	tracker.Write([]byte("data: [DONE]\n\ndata: unexpected\n\n"))
+
+	if tracker.Complete() {
+		t.Fatal("SSE stream with an event after its terminal event must not be complete")
+	}
+}
+
+func TestSSEEventTrackerAcceptsCROnlyEventBoundary(t *testing.T) {
+	var tracker sseEventTracker
+	tracker.Write([]byte("data: [DONE]\r\r"))
+
+	if !tracker.Complete() {
+		t.Fatal("terminal SSE event with CR-only boundary should be complete")
+	}
+}