Stream events already in hand go to the client in one write, not one each

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

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)
+	}
+}

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
 }