llmgw/internal/router/proxy/flush_test.go

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