feat: SSE reader and flushing writers

2026-10-13 04:58:50 +00:00 · 2026-08-20 02:39:21 +01:00
parent 2acc0a2a84
commit 0c6b53d296
2 changed files with 468 additions and 0 deletions
@@ -0,0 +1,167 @@
+package traefikllmgateway
+
+import (
+	"bufio"
+	"io"
+	"net/http"
+	"strings"
+)
+
+// sseScannerInitialBufSize is bufio.Scanner's starting token buffer for
+// readSSE. It grows on demand up to sseScannerMaxLineSize.
+const sseScannerInitialBufSize = 64 * 1024
+
+// sseScannerMaxLineSize caps a single text/event-stream field line at
+// 4MiB. Upstream LLM SSE events can carry large data lines — a
+// base64-encoded image, for instance — so the cap is generous, but a line
+// beyond it makes bufio.Scanner return bufio.ErrTooLong instead of
+// growing its buffer without bound.
+const sseScannerMaxLineSize = 4 * 1024 * 1024
+
+// sseDataPrefix is the "data:" field prefix sseWriter writes before every
+// event payload, per the text/event-stream wire format.
+const sseDataPrefix = "data: "
+
+// sseEvent is one parsed text/event-stream event: its "event:" field
+// (empty when the stream omits it) and its "data:" field, with multi-line
+// data joined by "\n" per the SSE spec.
+type sseEvent struct {
+	event string
+	data  []byte
+}
+
+// readSSE parses r as a text/event-stream body and calls fn once per
+// event. It accumulates "event:" and "data:" field lines — multiple
+// "data:" lines join with "\n" — and dispatches fn on a blank line, then
+// resets the accumulated fields. Comment lines (starting with ":") and
+// "id:"/"retry:" fields are recognized and silently ignored. Line endings
+// may be "\n" or "\r\n": bufio.ScanLines, the scanner's split function,
+// strips an optional trailing "\r" itself.
+//
+// readSSE stops and returns fn's error the first time fn returns one. On
+// a clean EOF it returns nil, first dispatching a final un-terminated
+// event (fields accumulated with no trailing blank line) if the stream
+// ended with one pending. A field line longer than sseScannerMaxLineSize
+// makes the underlying scanner fail; that error (bufio.ErrTooLong) is
+// returned as-is.
+func readSSE(r io.Reader, fn func(sseEvent) error) error {
+	scanner := bufio.NewScanner(r)
+	scanner.Buffer(make([]byte, sseScannerInitialBufSize), sseScannerMaxLineSize)
+
+	var (
+		event     string
+		dataLines []string
+		pending   bool
+	)
+
+	dispatch := func() error {
+		if !pending {
+			return nil
+		}
+		ev := sseEvent{event: event, data: []byte(strings.Join(dataLines, "\n"))}
+		event, dataLines, pending = "", nil, false
+		return fn(ev)
+	}
+
+	for scanner.Scan() {
+		line := scanner.Text()
+		if line == "" {
+			if err := dispatch(); err != nil {
+				return err
+			}
+			continue
+		}
+		if strings.HasPrefix(line, ":") {
+			continue // comment line, per the SSE spec
+		}
+
+		field, value, _ := strings.Cut(line, ":")
+		value = strings.TrimPrefix(value, " ")
+		switch field {
+		case "event":
+			event = value
+			pending = true
+		case "data":
+			dataLines = append(dataLines, value)
+			pending = true
+		case "id", "retry":
+			// recognized fields, intentionally ignored
+		}
+	}
+	if err := scanner.Err(); err != nil {
+		return err
+	}
+	return dispatch()
+}
+
+// sseWriter writes Server-Sent Events over an http.ResponseWriter,
+// flushing after each event so a proxied client sees each chunk as it is
+// produced instead of buffered until the response closes.
+type sseWriter struct {
+	w http.ResponseWriter
+	f http.Flusher
+}
+
+// newSSEWriter returns an sseWriter over w, setting the SSE response
+// headers before anything is written: Content-Type: text/event-stream,
+// Cache-Control: no-cache, and X-Accel-Buffering: no (disables buffering
+// in a reverse proxy such as nginx sitting in front of the gateway). w
+// need not implement http.Flusher — the Gateway's statusTrackingWriter
+// always does, since it delegates — but writeData/writeDone tolerate its
+// absence by skipping the flush.
+func newSSEWriter(w http.ResponseWriter) *sseWriter {
+	h := w.Header()
+	h.Set("Content-Type", "text/event-stream")
+	h.Set("Cache-Control", "no-cache")
+	h.Set("X-Accel-Buffering", "no")
+
+	f, _ := w.(http.Flusher)
+	return &sseWriter{w: w, f: f}
+}
+
+// writeData writes one SSE data event — "data: " + b + "\n\n" — as a
+// single Write call, then flushes when the underlying writer supports it.
+func (s *sseWriter) writeData(b []byte) error {
+	buf := make([]byte, 0, len(sseDataPrefix)+len(b)+2)
+	buf = append(buf, sseDataPrefix...)
+	buf = append(buf, b...)
+	buf = append(buf, '\n', '\n')
+
+	if _, err := s.w.Write(buf); err != nil {
+		return err
+	}
+	if s.f != nil {
+		s.f.Flush()
+	}
+	return nil
+}
+
+// writeDone writes the SSE stream's terminal "data: [DONE]\n\n" sentinel
+// (the OpenAI streaming convention) and flushes. Write errors are not
+// returned — writeDone is the last thing a handler does before returning,
+// and the connection is about to close either way.
+func (s *sseWriter) writeDone() {
+	_, _ = s.w.Write([]byte(sseDataPrefix + "[DONE]\n\n"))
+	if s.f != nil {
+		s.f.Flush()
+	}
+}
+
+// flushWriter wraps an io.Writer, flushing after every Write. Passthrough
+// and MCP proxying use it as the destination for io.Copy, so a streamed
+// reverse-proxy body reaches the client incrementally instead of sitting
+// in a buffer until the copy finishes.
+type flushWriter struct {
+	w io.Writer
+	f http.Flusher
+}
+
+// Write delegates to the wrapped io.Writer, then flushes when f is
+// non-nil.
+func (fw *flushWriter) Write(b []byte) (int, error) {
+	n, err := fw.w.Write(b)
+	if fw.f != nil {
+		fw.f.Flush()
+	}
+	return n, err
+}
@@ -0,0 +1,301 @@
+package traefikllmgateway
+
+import (
+	"bufio"
+	"bytes"
+	"errors"
+	"net/http"
+	"reflect"
+	"strings"
+	"testing"
+)
+
+// recordingResponseWriter is a minimal http.ResponseWriter that also
+// implements http.Flusher, recording every write and flush so tests can
+// pin sseWriter's exact wire output and its one-flush-per-event contract.
+type recordingResponseWriter struct {
+	header  http.Header
+	body    bytes.Buffer
+	flushed int
+}
+
+func newRecordingResponseWriter() *recordingResponseWriter {
+	return &recordingResponseWriter{header: make(http.Header)}
+}
+
+func (w *recordingResponseWriter) Header() http.Header         { return w.header }
+func (w *recordingResponseWriter) Write(b []byte) (int, error) { return w.body.Write(b) }
+func (w *recordingResponseWriter) WriteHeader(int)             {}
+func (w *recordingResponseWriter) Flush()                      { w.flushed++ }
+
+// noFlushResponseWriter is an http.ResponseWriter that deliberately does
+// NOT implement http.Flusher, proving sseWriter tolerates its absence.
+type noFlushResponseWriter struct {
+	header http.Header
+	body   bytes.Buffer
+}
+
+func newNoFlushResponseWriter() *noFlushResponseWriter {
+	return &noFlushResponseWriter{header: make(http.Header)}
+}
+
+func (w *noFlushResponseWriter) Header() http.Header         { return w.header }
+func (w *noFlushResponseWriter) Write(b []byte) (int, error) { return w.body.Write(b) }
+func (w *noFlushResponseWriter) WriteHeader(int)             {}
+
+// recordingFlusher is a bare http.Flusher used to pin flushWriter's
+// flush-after-every-write contract independent of any io.Writer.
+type recordingFlusher struct{ flushed int }
+
+func (f *recordingFlusher) Flush() { f.flushed++ }
+
+// TestReadSSE_Fixture parses a fixture mixing a comment line, CRLF and LF
+// line endings, and a multi-line data field, and pins the exact two events
+// it must produce: multi-line data joins with "\n", and the field order in
+// the stream does not matter to the result.
+func TestReadSSE_Fixture(t *testing.T) {
+	raw := ": keep-alive\r\n" +
+		"event: message\r\n" +
+		"data: line one\r\n" +
+		"data: line two\r\n" +
+		"\r\n" +
+		"data: second event\n" +
+		"\n"
+
+	var got []sseEvent
+	err := readSSE(strings.NewReader(raw), func(ev sseEvent) error {
+		got = append(got, ev)
+		return nil
+	})
+	if err != nil {
+		t.Fatalf("readSSE: %v", err)
+	}
+
+	want := []sseEvent{
+		{event: "message", data: []byte("line one\nline two")},
+		{event: "", data: []byte("second event")},
+	}
+	if !reflect.DeepEqual(got, want) {
+		t.Errorf("got %#v, want %#v", got, want)
+	}
+}
+
+// TestReadSSE_IDAndRetryIgnored proves "id:" and "retry:" fields are
+// recognized and dropped: they neither leak into the dispatched event nor
+// trigger a dispatch of their own.
+func TestReadSSE_IDAndRetryIgnored(t *testing.T) {
+	raw := "id: 42\n" +
+		"retry: 3000\n" +
+		"data: hello\n" +
+		"\n"
+
+	var got []sseEvent
+	err := readSSE(strings.NewReader(raw), func(ev sseEvent) error {
+		got = append(got, ev)
+		return nil
+	})
+	if err != nil {
+		t.Fatalf("readSSE: %v", err)
+	}
+
+	want := []sseEvent{{event: "", data: []byte("hello")}}
+	if !reflect.DeepEqual(got, want) {
+		t.Errorf("got %#v, want %#v", got, want)
+	}
+}
+
+// TestReadSSE_FinalEventFlushedAtEOF proves an event with no trailing
+// blank line before EOF still gets dispatched once, not silently dropped.
+func TestReadSSE_FinalEventFlushedAtEOF(t *testing.T) {
+	raw := "event: partial\ndata: leftover"
+
+	var got []sseEvent
+	err := readSSE(strings.NewReader(raw), func(ev sseEvent) error {
+		got = append(got, ev)
+		return nil
+	})
+	if err != nil {
+		t.Fatalf("readSSE: %v", err)
+	}
+
+	want := []sseEvent{{event: "partial", data: []byte("leftover")}}
+	if !reflect.DeepEqual(got, want) {
+		t.Errorf("got %#v, want %#v", got, want)
+	}
+}
+
+// TestReadSSE_StopsOnFnError proves readSSE returns fn's error verbatim
+// (checkable with errors.Is) and stops after the first event instead of
+// continuing to parse and dispatch the rest of the stream.
+func TestReadSSE_StopsOnFnError(t *testing.T) {
+	raw := "data: first\n\ndata: second\n\n"
+	sentinel := errors.New("boom")
+
+	var calls int
+	err := readSSE(strings.NewReader(raw), func(sseEvent) error {
+		calls++
+		return sentinel
+	})
+
+	if !errors.Is(err, sentinel) {
+		t.Fatalf("err = %v, want sentinel %v", err, sentinel)
+	}
+	if calls != 1 {
+		t.Errorf("calls = %d, want 1 (must stop after fn's first error)", calls)
+	}
+}
+
+// TestReadSSE_LineTooLong proves a single field line beyond
+// sseScannerMaxLineSize surfaces the scanner's own error instead of
+// growing the buffer without bound.
+func TestReadSSE_LineTooLong(t *testing.T) {
+	huge := strings.Repeat("a", sseScannerMaxLineSize+10)
+	raw := "data: " + huge + "\n\n"
+
+	err := readSSE(strings.NewReader(raw), func(sseEvent) error { return nil })
+	if !errors.Is(err, bufio.ErrTooLong) {
+		t.Fatalf("err = %v, want bufio.ErrTooLong", err)
+	}
+}
+
+// TestReadSSE_BlankLinesWithoutFieldsDoNotDispatch proves blank lines and
+// comment-only content, with no "event:"/"data:" field accumulated,
+// produce zero dispatched events rather than empty ones.
+func TestReadSSE_BlankLinesWithoutFieldsDoNotDispatch(t *testing.T) {
+	raw := "\n\n: just a comment\n\n"
+
+	var calls int
+	err := readSSE(strings.NewReader(raw), func(sseEvent) error {
+		calls++
+		return nil
+	})
+	if err != nil {
+		t.Fatalf("readSSE: %v", err)
+	}
+	if calls != 0 {
+		t.Errorf("calls = %d, want 0", calls)
+	}
+}
+
+// TestNewSSEWriter_SetsHeaders proves newSSEWriter sets the three SSE
+// response headers on construction, before any body write.
+func TestNewSSEWriter_SetsHeaders(t *testing.T) {
+	w := newRecordingResponseWriter()
+	newSSEWriter(w)
+
+	cases := map[string]string{
+		"Content-Type":      "text/event-stream",
+		"Cache-Control":     "no-cache",
+		"X-Accel-Buffering": "no",
+	}
+	for header, want := range cases {
+		if got := w.Header().Get(header); got != want {
+			t.Errorf("header %q = %q, want %q", header, got, want)
+		}
+	}
+	if w.body.Len() != 0 {
+		t.Errorf("body = %q, want empty (headers must be set before any write)", w.body.String())
+	}
+}
+
+// TestSSEWriter_WriteData pins the exact wire bytes writeData produces —
+// "data: " + payload + "\n\n" — and proves it flushes once per call.
+func TestSSEWriter_WriteData(t *testing.T) {
+	w := newRecordingResponseWriter()
+	s := newSSEWriter(w)
+
+	if err := s.writeData([]byte(`{"a":1}`)); err != nil {
+		t.Fatalf("writeData: %v", err)
+	}
+	if err := s.writeData([]byte("second")); err != nil {
+		t.Fatalf("writeData: %v", err)
+	}
+
+	want := "data: {\"a\":1}\n\ndata: second\n\n"
+	if got := w.body.String(); got != want {
+		t.Errorf("body = %q, want %q", got, want)
+	}
+	if w.flushed != 2 {
+		t.Errorf("flushed = %d, want 2 (one Flush per writeData call)", w.flushed)
+	}
+}
+
+// TestSSEWriter_WriteDone pins the "[DONE]" sentinel's exact wire bytes and
+// proves it flushes.
+func TestSSEWriter_WriteDone(t *testing.T) {
+	w := newRecordingResponseWriter()
+	s := newSSEWriter(w)
+
+	s.writeDone()
+
+	want := "data: [DONE]\n\n"
+	if got := w.body.String(); got != want {
+		t.Errorf("body = %q, want %q", got, want)
+	}
+	if w.flushed != 1 {
+		t.Errorf("flushed = %d, want 1", w.flushed)
+	}
+}
+
+// TestSSEWriter_ToleratesMissingFlusher proves writeData/writeDone still
+// write the correct bytes, without panicking, when the underlying
+// http.ResponseWriter does not implement http.Flusher.
+func TestSSEWriter_ToleratesMissingFlusher(t *testing.T) {
+	w := newNoFlushResponseWriter()
+	s := newSSEWriter(w)
+
+	if err := s.writeData([]byte("x")); err != nil {
+		t.Fatalf("writeData: %v", err)
+	}
+	s.writeDone()
+
+	want := "data: x\n\ndata: [DONE]\n\n"
+	if got := w.body.String(); got != want {
+		t.Errorf("body = %q, want %q", got, want)
+	}
+}
+
+// TestFlushWriter_WritesAndFlushes proves Write delegates to the wrapped
+// io.Writer and calls Flush once per Write when f is non-nil.
+func TestFlushWriter_WritesAndFlushes(t *testing.T) {
+	var buf bytes.Buffer
+	fl := &recordingFlusher{}
+	fw := &flushWriter{w: &buf, f: fl}
+
+	n, err := fw.Write([]byte("chunk"))
+	if err != nil {
+		t.Fatalf("Write: %v", err)
+	}
+	if n != len("chunk") {
+		t.Errorf("n = %d, want %d", n, len("chunk"))
+	}
+
+	if _, err := fw.Write([]byte("more")); err != nil {
+		t.Fatalf("Write: %v", err)
+	}
+
+	if got := buf.String(); got != "chunkmore" {
+		t.Errorf("buf = %q, want %q", got, "chunkmore")
+	}
+	if fl.flushed != 2 {
+		t.Errorf("flushed = %d, want 2", fl.flushed)
+	}
+}
+
+// TestFlushWriter_NilFlusher proves Write still delegates correctly, and
+// does not panic, when f is nil.
+func TestFlushWriter_NilFlusher(t *testing.T) {
+	var buf bytes.Buffer
+	fw := &flushWriter{w: &buf}
+
+	n, err := fw.Write([]byte("ok"))
+	if err != nil {
+		t.Fatalf("Write: %v", err)
+	}
+	if n != len("ok") {
+		t.Errorf("n = %d, want %d", n, len("ok"))
+	}
+	if got := buf.String(); got != "ok" {
+		t.Errorf("buf = %q, want %q", got, "ok")
+	}
+}