llmgw/internal/router/proxy/proxy_test.go

package proxy

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

func testProxy(mut ...func(*Config)) *Proxy {
	c := DefaultConfig()
	c.FirstTokenTimeout = 2 * time.Second
	c.IdleTimeout = 500 * time.Millisecond
	c.TotalTimeout = 10 * time.Second
	for _, f := range mut {
		f(&c)
	}
	return New(c)
}

// sseServer serves a scripted SSE stream.
func sseServer(t *testing.T, chunks []string, gap time.Duration) (*httptest.Server, *atomic.Bool) {
	t.Helper()
	var cancelled atomic.Bool
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		w.WriteHeader(http.StatusOK)
		f := w.(http.Flusher)
		f.Flush()
		for _, c := range chunks {
			select {
			case <-r.Context().Done():
				cancelled.Store(true)
				return
			case <-time.After(gap):
			}
			if _, err := io.WriteString(w, c); err != nil {
				cancelled.Store(true)
				return
			}
			f.Flush()
		}
	}))
	t.Cleanup(srv.Close)
	return srv, &cancelled
}

func dataChunk(s string) string { return "data: " + s + "\n\n" }

func TestBufferedResponseIsRelayed(t *testing.T) {
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		io.WriteString(w, `{"id":"x","usage":{"prompt_tokens":7,"completion_tokens":3}}`)
	}))
	defer up.Close()

	rec := httptest.NewRecorder()
	out, err := testProxy().Do(t.Context(), rec, Request{URL: up.URL, Body: []byte(`{}`)})
	if err != nil {
		t.Fatal(err)
	}
	if out.Status != 200 || !out.Wrote {
		t.Errorf("outcome=%+v", out)
	}
	if out.Usage == nil || out.Usage.PromptTokens != 7 || out.Usage.TotalTokens != 10 {
		t.Errorf("usage not captured: %+v", out.Usage)
	}
	if !strings.Contains(rec.Body.String(), `"id":"x"`) {
		t.Errorf("body not relayed: %q", rec.Body.String())
	}
}

func TestStreamIsRelayedEventByEvent(t *testing.T) {
	up, _ := sseServer(t, []string{
		dataChunk(`{"choices":[{"delta":{"content":"a"}}]}`),
		dataChunk(`{"choices":[{"delta":{"content":"b"}}]}`),
		"data: [DONE]\n\n",
	}, time.Millisecond)

	rec := httptest.NewRecorder()
	out, err := testProxy().Do(t.Context(), rec, Request{URL: up.URL, Body: []byte(`{}`), Stream: true})
	if err != nil {
		t.Fatal(err)
	}
	body := rec.Body.String()
	if n := strings.Count(body, "data: "); n != 3 {
		t.Errorf("expected 3 data lines, got %d: %q", n, body)
	}
	if !strings.HasSuffix(body, "[DONE]\n\n") {
		t.Errorf("stream did not end with [DONE]: %q", body)
	}
	if out.TTFT == 0 {
		t.Error("TTFT was not recorded")
	}
}

func TestUsageCapturedFromStream(t *testing.T) {
	up, _ := sseServer(t, []string{
		dataChunk(`{"choices":[{"delta":{"content":"a"}}]}`),
		dataChunk(`{"choices":[],"usage":{"prompt_tokens":11,"completion_tokens":5,"total_tokens":16}}`),
		"data: [DONE]\n\n",
	}, time.Millisecond)

	rec := httptest.NewRecorder()
	out, err := testProxy().Do(t.Context(), rec, Request{URL: up.URL, Body: []byte(`{}`), Stream: true})
	if err != nil {
		t.Fatal(err)
	}
	if out.Usage == nil || out.Usage.TotalTokens != 16 {
		t.Fatalf("usage not captured from stream: %+v", out.Usage)
	}
	if !strings.Contains(rec.Body.String(), `"usage"`) {
		t.Error("usage chunk should reach a client that asked for it")
	}
}

func TestUsageChunkSuppressed(t *testing.T) {
	// The router asked the engine for usage the client never requested: record
	// it, but do not change what the client receives.
	up, _ := sseServer(t, []string{
		dataChunk(`{"choices":[{"delta":{"content":"a"}}]}`),
		dataChunk(`{"choices":[],"usage":{"prompt_tokens":11,"completion_tokens":5,"total_tokens":16}}`),
		"data: [DONE]\n\n",
	}, time.Millisecond)

	rec := httptest.NewRecorder()
	out, err := testProxy().Do(t.Context(), rec, Request{
		URL: up.URL, Body: []byte(`{}`), Stream: true, SuppressUsageChunk: true,
	})
	if err != nil {
		t.Fatal(err)
	}
	if out.Usage == nil || out.Usage.TotalTokens != 16 {
		t.Fatalf("usage not recorded: %+v", out.Usage)
	}
	if strings.Contains(rec.Body.String(), `"usage"`) {
		t.Errorf("usage chunk leaked to the client: %q", rec.Body.String())
	}
	if !strings.Contains(rec.Body.String(), "[DONE]") {
		t.Error("[DONE] was dropped along with the usage chunk")
	}
}

func TestRetryableStatusWritesNothing(t *testing.T) {
	// The caller must be free to try another endpoint, which is only safe if
	// the client has seen nothing.
	for _, code := range []int{429, 502, 503, 504} {
		up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
			http.Error(w, "busy", code)
		}))

		rec := httptest.NewRecorder()
		out, err := testProxy().Do(t.Context(), rec, Request{URL: up.URL, Body: []byte(`{}`)})
		up.Close()

		if err == nil {
			t.Errorf("%d: expected an error", code)
		}
		if !out.Retryable {
			t.Errorf("%d: not marked retryable", code)
		}
		if out.Wrote || rec.Body.Len() > 0 {
			t.Errorf("%d: wrote %q to the client", code, rec.Body.String())
		}
	}
}

func TestNonRetryableStatusIsRelayed(t *testing.T) {
	// A 400 is the client's fault; retrying elsewhere would produce the same
	// answer, so relay it.
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusBadRequest)
		io.WriteString(w, `{"error":{"message":"bad"}}`)
	}))
	defer up.Close()

	rec := httptest.NewRecorder()
	out, err := testProxy().Do(t.Context(), rec, Request{URL: up.URL, Body: []byte(`{}`)})
	if err != nil {
		t.Fatal(err)
	}
	if out.Retryable {
		t.Error("400 must not be retryable")
	}
	if rec.Code != http.StatusBadRequest || !strings.Contains(rec.Body.String(), "bad") {
		t.Errorf("error response not relayed: %d %q", rec.Code, rec.Body.String())
	}
}

func TestConnectionFailureIsRetryable(t *testing.T) {
	rec := httptest.NewRecorder()
	out, err := testProxy(func(c *Config) { c.DialTimeout = 200 * time.Millisecond }).
		Do(t.Context(), rec, Request{URL: "http://127.0.0.1:1/x", Body: []byte(`{}`)})
	if err == nil {
		t.Fatal("expected an error")
	}
	if !out.Retryable || out.Wrote {
		t.Errorf("outcome=%+v", out)
	}
}

func TestClientDisconnectCancelsUpstream(t *testing.T) {
	// Phase 1's exit criterion at the proxy level: the GPU must stop working
	// when the client goes away.
	chunks := make([]string, 200)
	for i := range chunks {
		chunks[i] = dataChunk(fmt.Sprintf(`{"choices":[{"delta":{"content":"%d"}}]}`, i))
	}
	up, cancelled := sseServer(t, chunks, 10*time.Millisecond)

	ctx, cancel := context.WithCancel(context.Background())
	rec := httptest.NewRecorder()

	done := make(chan Outcome, 1)
	go func() {
		out, _ := testProxy().Do(ctx, rec, Request{URL: up.URL, Body: []byte(`{}`), Stream: true})
		done <- out
	}()

	time.Sleep(80 * time.Millisecond) // let a few events through
	cancel()

	select {
	case <-done:
	case <-time.After(3 * time.Second):
		t.Fatal("proxy did not return after client cancellation")
	}

	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
		if cancelled.Load() {
			return
		}
		time.Sleep(10 * time.Millisecond)
	}
	t.Fatal("upstream kept generating after the client disconnected")
}

func TestStalledUpstreamHitsIdleTimeout(t *testing.T) {
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		w.WriteHeader(http.StatusOK)
		f := w.(http.Flusher)
		io.WriteString(w, dataChunk(`{"choices":[{"delta":{"content":"a"}}]}`))
		f.Flush()
		<-r.Context().Done() // then go silent forever
	}))
	defer up.Close()

	rec := httptest.NewRecorder()
	p := testProxy(func(c *Config) { c.IdleTimeout = 200 * time.Millisecond })

	start := time.Now()
	out, err := p.Do(t.Context(), rec, Request{URL: up.URL, Body: []byte(`{}`), Stream: true})
	elapsed := time.Since(start)

	if !errors.Is(err, ErrUpstreamStalled) {
		t.Fatalf("expected ErrUpstreamStalled, got %v", err)
	}
	if !out.Wrote {
		t.Error("the events delivered before the stall should be reported as written")
	}
	if elapsed > 3*time.Second {
		t.Errorf("idle timeout took too long: %v", elapsed)
	}
}

func TestHopByHopHeadersNotForwarded(t *testing.T) {
	var got http.Header
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		got = r.Header.Clone()
		io.WriteString(w, `{}`)
	}))
	defer up.Close()

	h := http.Header{}
	h.Set("Authorization", "Bearer sk-llmgw-secret")
	h.Set("Accept", "text/event-stream")
	h.Set("Traceparent", "00-trace-span-01")

	rec := httptest.NewRecorder()
	if _, err := testProxy().Do(t.Context(), rec, Request{URL: up.URL, Body: []byte(`{}`), Header: h}); err != nil {
		t.Fatal(err)
	}

	if got.Get("Authorization") != "" {
		t.Error("the client's Authorization header was forwarded upstream")
	}
	if got.Get("Accept") != "text/event-stream" {
		t.Error("Accept was not forwarded")
	}
	if got.Get("Traceparent") != "00-trace-span-01" {
		t.Error("Traceparent was not forwarded; tracing would break")
	}
}