tapir/internal/adapters/llm/bodycap_test.go

package llm

// Observes how much of a gateway response the client will read. The body is read before
// the status is looked at, so a misbehaving gateway or proxy that streams forever was
// bounded only by the client's overall timeout (#53).

import (
	"context"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"
	"time"
)

// endlessGateway answers with the given status and streams filler until the client hangs
// up; wrote reports how many bytes it managed to send.
func endlessGateway(t *testing.T, status int, wrote *atomic.Int64) *httptest.Server {
	t.Helper()
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(status)
		chunk := []byte(strings.Repeat("a", 4096))
		for {
			n, err := w.Write(chunk)
			wrote.Add(int64(n))
			if err != nil {
				return
			}
			if f, ok := w.(http.Flusher); ok {
				f.Flush()
			}
		}
	}))
	t.Cleanup(srv.Close)
	return srv
}

func TestClient_ASuccessBodyLargerThanTheCapIsAnErrorAndTheClientStopsReading(t *testing.T) {
	var wrote atomic.Int64
	srv := endlessGateway(t, http.StatusOK, &wrote)
	c := New(srv.URL, "", "m", 30*time.Second, WithMaxResponseBytes(8192))

	start := time.Now()
	_, err := c.Complete(context.Background(), "s", "u")

	if err == nil || !strings.Contains(err.Error(), "exceeds") {
		t.Fatalf("want a response-too-large error, got %v", err)
	}
	if time.Since(start) > 5*time.Second {
		t.Fatalf("the client kept reading: took %s", time.Since(start))
	}
	// Hanging up lets the server's writes fail; allow what the transport had already buffered.
	time.Sleep(100 * time.Millisecond)
	if got := wrote.Load(); got > 8192+1<<20 {
		t.Fatalf("server managed to send %d bytes after a cap of 8192", got)
	}
}

func TestClient_AFailingBodyIsReadOnlyUpToTheErrorBound(t *testing.T) {
	var wrote atomic.Int64
	srv := endlessGateway(t, http.StatusBadGateway, &wrote)
	c := New(srv.URL, "", "m", 30*time.Second)

	start := time.Now()
	_, err := c.Complete(context.Background(), "s", "u")

	he, ok := err.(*HTTPError)
	if !ok {
		t.Fatalf("want *HTTPError, got %T: %v", err, err)
	}
	if he.Code != http.StatusBadGateway {
		t.Fatalf("code %d", he.Code)
	}
	if len(he.Body) > maxErrorBody+1 {
		t.Fatalf("kept %d bytes of an error body, bound is %d", len(he.Body), maxErrorBody)
	}
	if time.Since(start) > 5*time.Second {
		t.Fatalf("the client kept reading an endless error body: took %s", time.Since(start))
	}
}

// sizedGateway answers 200 with a valid completion padded with trailing whitespace to
// exactly size bytes, so the boundary of the cap can be probed to the byte.
func sizedGateway(t *testing.T, size int) *httptest.Server {
	t.Helper()
	const base = `{"choices":[{"message":{"content":"ok"}}]}`
	payload := base + strings.Repeat(" ", size-len(base))
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(payload))
	}))
	t.Cleanup(srv.Close)
	return srv
}

func TestClient_ABodyOfExactlyTheCapIsReadAndOneByteMoreIsRefused(t *testing.T) {
	const capBytes = 4096
	for _, tt := range []struct {
		name    string
		size    int
		wantErr bool
	}{
		{"one under", capBytes - 1, false},
		{"exactly the cap", capBytes, false},
		{"one over", capBytes + 1, true},
	} {
		t.Run(tt.name, func(t *testing.T) {
			c := New(sizedGateway(t, tt.size).URL, "", "m", 5*time.Second, WithMaxResponseBytes(capBytes))

			got, err := c.Complete(context.Background(), "s", "u")

			if tt.wantErr {
				if err == nil || !strings.Contains(err.Error(), "exceeds") {
					t.Fatalf("want a response-too-large error, got %q, %v", got, err)
				}
				return
			}
			if err != nil || got != "ok" {
				t.Fatalf("want ok, got %q, %v", got, err)
			}
		})
	}
}

// New is what production calls: no option. The default is therefore the live bound.
func TestClient_TheDefaultCapIsEightMiBWhenNoOptionIsGiven(t *testing.T) {
	for _, tt := range []struct {
		name    string
		size    int
		wantErr bool
	}{
		{"exactly 8 MiB", 8 << 20, false},
		{"8 MiB plus one", 8<<20 + 1, true},
	} {
		t.Run(tt.name, func(t *testing.T) {
			c := New(sizedGateway(t, tt.size).URL, "", "m", 10*time.Second)

			_, err := c.Complete(context.Background(), "s", "u")

			if (err != nil) != tt.wantErr {
				t.Fatalf("size %d: err = %v, wantErr %v", tt.size, err, tt.wantErr)
			}
		})
	}
}

func TestClient_ANonPositiveCapOptionKeepsTheDefault(t *testing.T) {
	for _, n := range []int64{0, -1} {
		c := New(sizedGateway(t, 1024).URL, "", "m", 5*time.Second, WithMaxResponseBytes(n))

		if got, err := c.Complete(context.Background(), "s", "u"); err != nil || got != "ok" {
			t.Fatalf("WithMaxResponseBytes(%d): want ok from the default cap, got %q, %v", n, got, err)
		}
	}
}