From 5550fefaada45912894d754d68397b31bad7e474 Mon Sep 17 00:00:00 2001 From: Ilker Manap Date: Sat, 31 Oct 2026 15:30:19 +0100 Subject: [PATCH] Stream events already in hand go to the client in one write, not one each Co-Authored-By: Claude Opus 5.5 (1M context) --- internal/router/proxy/flush_test.go | 50 +++++++++++++++++++++++++++++ internal/router/proxy/proxy.go | 23 ++++++++++--- 2 files changed, 68 insertions(+), 5 deletions(-) create mode 100644 internal/router/proxy/flush_test.go diff --git a/internal/router/proxy/flush_test.go b/internal/router/proxy/flush_test.go new file mode 100644 index 0000000..a474459 --- /dev/null +++ b/internal/router/proxy/flush_test.go @@ -0,0 +1,50 @@ +package proxy + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +// flushCounter is a response writer that counts its flushes. +type flushCounter struct { + *httptest.ResponseRecorder + flushes int +} + +func (f *flushCounter) Flush() { f.flushes++; f.ResponseRecorder.Flush() } + +// Events the engine sent together go to the client in one write, not one +// each: a write a token was a third of a busy router's CPU. +func TestEventsAlreadyInHandAreFlushedTogether(t *testing.T) { + var b strings.Builder + for i := range 10 { + b.WriteString(`data: {"choices":[{"index":0,"delta":{"content":"t` + string(rune('0'+i)) + `"}}]}` + "\n\n") + } + b.WriteString("data: [DONE]\n\n") + resp := &http.Response{StatusCode: 200, Header: http.Header{"Content-Type": {"text/event-stream"}}, + Body: io.NopCloser(strings.NewReader(b.String()))} + w := &flushCounter{ResponseRecorder: httptest.NewRecorder()} + p := New(Config{}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + if _, err := p.relayStream(ctx, cancel, w, resp, Request{}); err != nil { + t.Fatal(err) + } + got := w.Body.String() + for i := range 10 { + if !strings.Contains(got, `"t`+string(rune('0'+i))+`"`) { + t.Fatalf("event %d missing:\n%s", i, got) + } + } + if !strings.HasSuffix(strings.TrimSpace(got), "data: [DONE]") { + t.Errorf("not ended with [DONE]: %q", got[len(got)-40:]) + } + // One for the headers, and the events: far fewer than one each. + if w.flushes > 3 { + t.Errorf("%d flushes for 11 events arriving together", w.flushes) + } +} diff --git a/internal/router/proxy/proxy.go b/internal/router/proxy/proxy.go index 7cf681b..da86234 100644 --- a/internal/router/proxy/proxy.go +++ b/internal/router/proxy/proxy.go @@ -418,7 +418,7 @@ func (p *Proxy) relayStream(ctx context.Context, abort context.CancelFunc, w htt emit = append(emit, "[DONE]") } } - if err := p.emit(w, flusher, &out, &reply, r, emit); err != nil { + if err := p.emit(w, flusher, &out, &reply, r, emit, br.Buffered() > 0 && !stop); err != nil { return out, err } if stop { @@ -434,7 +434,9 @@ func (p *Proxy) relayStream(ctx context.Context, abort context.CancelFunc, w htt if werr != nil { return out, ErrClientGone } - flusher.Flush() + if br.Buffered() == 0 { + flusher.Flush() + } } } @@ -442,7 +444,7 @@ func (p *Proxy) relayStream(ctx context.Context, abort context.CancelFunc, w htt if errors.Is(err, io.EOF) { if r.StreamFilter != nil { // The engine ended without [DONE]: what was held goes out. - if err := p.emit(w, flusher, &out, &reply, r, r.StreamFilter.OnEnd()); err != nil { + if err := p.emit(w, flusher, &out, &reply, r, r.StreamFilter.OnEnd(), false); err != nil { return out, err } } @@ -479,8 +481,17 @@ func usageOnly(data string) bool { // emit writes events a StreamFilter chose, as data events, and captures the // reply from what the client is actually sent. -func (p *Proxy) emit(w io.Writer, f http.Flusher, out *Outcome, reply *replyBuilder, r Request, data []string) error { +// emit writes events to the client, and flushes them unless more has +// already arrived from the engine (more): those go out with the next write. +// One write a token was a third of a busy router's CPU in system calls; this +// joins only what is already in hand, so no token waits for another. +func (p *Proxy) emit(w io.Writer, f http.Flusher, out *Outcome, reply *replyBuilder, r Request, data []string, more bool) error { if len(data) == 0 { + // What an earlier event left unflushed goes out now that nothing + // more is in hand; with nothing written, a flush writes nothing. + if !more { + f.Flush() + } return nil } for _, d := range data { @@ -493,7 +504,9 @@ func (p *Proxy) emit(w io.Writer, f http.Flusher, out *Outcome, reply *replyBuil return ErrClientGone } } - f.Flush() + if !more { + f.Flush() + } return nil }