diff --git a/internal/adapters/llm/bodycap_test.go b/internal/adapters/llm/bodycap_test.go new file mode 100644 index 0000000..293381f --- /dev/null +++ b/internal/adapters/llm/bodycap_test.go @@ -0,0 +1,81 @@ +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)) + } +}