diff --git a/proxy/sse_event.go b/proxy/sse_event.go index 59c01b6..f572905 100644 --- a/proxy/sse_event.go +++ b/proxy/sse_event.go @@ -7,8 +7,10 @@ import ( ) type sseEventTracker struct { - buf []byte - terminal bool + 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) { - return crlf, 4 + cr := bytes.Index(buf, []byte("\r\r")) + end, separator := lf, 2 + if crlf >= 0 && (end < 0 || crlf < end) { + end, separator = crlf, 4 } - if lf >= 0 { - return lf, 2 + if cr >= 0 && (end < 0 || cr < end) { + 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"` + 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 } diff --git a/proxy/sse_event_test.go b/proxy/sse_event_test.go new file mode 100644 index 0000000..ea96855 --- /dev/null +++ b/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") + } +}