From 0c6b53d296a5cde4e1d780aab93a83237e03a719 Mon Sep 17 00:00:00 2001 From: Lukasz Raczylo Date: Thu, 20 Aug 2026 02:39:21 +0100 Subject: [PATCH] feat: SSE reader and flushing writers --- sse.go | 167 +++++++++++++++++++++++++++++ sse_test.go | 301 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 468 insertions(+) create mode 100644 sse.go create mode 100644 sse_test.go diff --git a/sse.go b/sse.go new file mode 100644 index 0000000..b612488 --- /dev/null +++ b/sse.go @@ -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 +} diff --git a/sse_test.go b/sse_test.go new file mode 100644 index 0000000..10b3046 --- /dev/null +++ b/sse_test.go @@ -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") + } +}