llmgw/internal/router/proxy/proxy.go

// Package proxy streams a request to one upstream engine and relays the
// response to the client (plan §7.8).
//
// It handles exactly one attempt. Choosing the endpoint, retrying and tripping
// breakers belong to the caller, because those decisions need the scheduler and
// the endpoint registry. What the proxy owes the caller is an honest answer to
// "did anything reach the client yet?" — once the first byte is out, retrying
// elsewhere would corrupt the response.
package proxy

import (
	"bufio"
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

// Config holds proxy settings.
type Config struct {
	// DialTimeout bounds establishing the TCP connection.
	DialTimeout time.Duration
	// ResponseHeaderTimeout bounds the wait for upstream response headers on a
	// streaming request, where headers arrive before generation begins and so
	// are a liveness check.
	ResponseHeaderTimeout time.Duration
	// GenerationTimeout bounds a non-streaming request, where the engine sends
	// no headers until the completion is finished.
	//
	// The two were one setting, and that made ResponseHeaderTimeout a total
	// generation budget in disguise for every non-streaming request. At 60s it
	// read as a sensible liveness guard and behaved as a cap of roughly a
	// hundred tokens on a model generating 1.7 a second — a request for 120
	// tokens failed at exactly one minute, and nothing in the name or the
	// message said why. Zero falls back to TotalTimeout.
	GenerationTimeout time.Duration
	// FirstTokenTimeout bounds the wait for the first streamed event.
	FirstTokenTimeout time.Duration
	// IdleTimeout bounds the gap between streamed events. A stalled engine
	// looks identical to a slow one over a single connection, so without this
	// a wedged request holds a client forever.
	IdleTimeout time.Duration
	// TotalTimeout bounds the whole request. Generous by default: a long
	// generation is legitimately slow.
	TotalTimeout time.Duration
	// MaxIdleConnsPerHost keeps connections warm; engines are few and busy.
	MaxIdleConnsPerHost int
	// MaxResponseBytes bounds a buffered (non-streaming) response.
	MaxResponseBytes int64
}

// DefaultConfig returns sensible defaults.
func DefaultConfig() Config {
	return Config{
		DialTimeout:           5 * time.Second,
		ResponseHeaderTimeout: 60 * time.Second,
		GenerationTimeout:     10 * time.Minute,
		FirstTokenTimeout:     120 * time.Second,
		IdleTimeout:           60 * time.Second,
		TotalTimeout:          30 * time.Minute,
		MaxIdleConnsPerHost:   64,
		MaxResponseBytes:      64 << 20,
	}
}

func (c *Config) normalize() {
	d := DefaultConfig()
	if c.DialTimeout <= 0 {
		c.DialTimeout = d.DialTimeout
	}
	if c.ResponseHeaderTimeout <= 0 {
		c.ResponseHeaderTimeout = d.ResponseHeaderTimeout
	}
	if c.FirstTokenTimeout <= 0 {
		c.FirstTokenTimeout = d.FirstTokenTimeout
	}
	if c.IdleTimeout <= 0 {
		c.IdleTimeout = d.IdleTimeout
	}
	if c.TotalTimeout <= 0 {
		c.TotalTimeout = d.TotalTimeout
	}
	if c.GenerationTimeout <= 0 {
		c.GenerationTimeout = d.GenerationTimeout
	}
	if c.MaxIdleConnsPerHost <= 0 {
		c.MaxIdleConnsPerHost = d.MaxIdleConnsPerHost
	}
	if c.MaxResponseBytes <= 0 {
		c.MaxResponseBytes = d.MaxResponseBytes
	}
}

// Proxy forwards requests to engine endpoints.
type Proxy struct {
	cfg Config
	// client is for streaming requests, where headers arrive before generation
	// and a short header timeout is a genuine liveness check.
	client *http.Client
	// buffered is for non-streaming requests, where the engine sends nothing
	// until the completion is done, so the header timeout is the generation
	// budget. Two transports rather than one, because the same number cannot
	// mean both things.
	buffered *http.Client
}

// New builds a proxy.
func New(cfg Config) *Proxy {
	cfg.normalize()
	transport := func(headerTimeout time.Duration) *http.Transport {
		return &http.Transport{
			DialContext:           (&net.Dialer{Timeout: cfg.DialTimeout}).DialContext,
			ResponseHeaderTimeout: headerTimeout,
			MaxIdleConnsPerHost:   cfg.MaxIdleConnsPerHost,
			MaxIdleConns:          cfg.MaxIdleConnsPerHost * 8,
			IdleConnTimeout:       90 * time.Second,
			// Engines rarely speak h2c, and HTTP/2 over cleartext would
			// need explicit setup anyway.
			ForceAttemptHTTP2:  false,
			DisableCompression: true,
		}
	}
	return &Proxy{
		cfg: cfg,
		buffered: &http.Client{
			Transport: transport(cfg.GenerationTimeout),
		},
		client: &http.Client{
			// No client timeout: streaming responses run for minutes and the
			// per-phase watchdogs below bound them more precisely.
			Transport: transport(cfg.ResponseHeaderTimeout),
		},
	}
}

// Request describes one upstream attempt.
type Request struct {
	// URL is the fully qualified upstream URL.
	URL string
	// Body is the request body to send.
	Body []byte
	// Header carries headers to forward (authorization is deliberately not).
	Header http.Header
	// APIKey is the engine's own key, sent in place of the client's.
	APIKey string
	// Stream says whether an SSE response is expected.
	Stream bool
	// SuppressUsageChunk drops the final usage-only chunk from the client while
	// still recording it. Set when the router asked for usage the client did
	// not request.
	SuppressUsageChunk bool
	// CaptureReply asks for the assistant's text back in the Outcome, up to
	// this many bytes. Zero means capture nothing, which is every request that
	// is not continuing a stored conversation.
	//
	// Opt-in and bounded on purpose. This is the hottest path in the system,
	// and a gateway that accumulated every reply in memory would be one long
	// answer away from trouble; a request that stores nothing must pay nothing.
	CaptureReply int
}

// Usage holds the token counts reported by the engine.
type Usage struct {
	PromptTokens     int
	CompletionTokens int
	TotalTokens      int
	CachedTokens     int
}

// Outcome reports what happened during one attempt.
type Outcome struct {
	// Wrote is true once any byte reached the client. Retrying is then unsafe.
	Wrote bool
	// Status is the upstream status code, 0 if the request never completed.
	Status int
	// Retryable means the attempt failed in a way another endpoint might survive.
	Retryable bool
	// TTFT is the time until the first streamed event, zero if none arrived.
	TTFT time.Duration
	// Usage is the engine-reported token usage, if it was present.
	Usage *Usage
	// Bytes is how many bytes were relayed to the client.
	Bytes int64
	// Reply is the assistant's text, when CaptureReply asked for it. Truncated
	// at the cap rather than dropped: a shortened turn replays as a shortened
	// turn, while nothing at all would replay as the assistant having said
	// nothing, which is a different and wrong conversation.
	Reply string
	// ReplyTruncated says the cap was reached, so nothing downstream records a
	// half answer as a whole one.
	ReplyTruncated bool
}

var (
	// ErrClientGone means the client disconnected; the upstream was cancelled.
	ErrClientGone = errors.New("client disconnected")
	// ErrUpstreamStalled means no event arrived within the idle or first-token timeout.
	ErrUpstreamStalled = errors.New("upstream stalled")
)

// Do performs one attempt against one endpoint.
//
// On a retryable failure it writes nothing to w, leaving the caller free to try
// another endpoint. Once it has written anything, Outcome.Wrote is true and the
// caller must not retry.
func (p *Proxy) Do(ctx context.Context, w http.ResponseWriter, r Request) (Outcome, error) {
	ctx, cancel := context.WithTimeout(ctx, p.cfg.TotalTimeout)
	defer cancel()

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, r.URL, bytes.NewReader(r.Body))
	if err != nil {
		return Outcome{Retryable: false}, err
	}
	req.Header = cloneForUpstream(r.Header)
	req.Header.Set("Content-Type", "application/json")
	if r.APIKey != "" {
		req.Header.Set("Authorization", "Bearer "+r.APIKey)
	}
	req.ContentLength = int64(len(r.Body))

	// A non-streaming request waits for the whole completion before a single
	// header arrives, so it is bounded by the generation budget rather than the
	// liveness timeout.
	client := p.client
	if !r.Stream {
		client = p.buffered
	}
	resp, err := client.Do(req)
	if err != nil {
		// The client going away is not the endpoint's fault; never blame it.
		if ctx.Err() != nil && ctxCauseIsClient(ctx) {
			return Outcome{Retryable: false}, ErrClientGone
		}
		return Outcome{Retryable: true}, fmt.Errorf("upstream request failed: %w", err)
	}
	defer resp.Body.Close()

	if retryableStatus(resp.StatusCode) {
		// Drain a little so the connection can be reused, then let the caller
		// pick somewhere else. Nothing has been written to the client.
		_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 4<<10))
		return Outcome{Status: resp.StatusCode, Retryable: true},
			fmt.Errorf("upstream returned %s", resp.Status)
	}

	if r.Stream && isEventStream(resp.Header) {
		return p.relayStream(ctx, cancel, w, resp, r)
	}
	return p.relayBuffered(w, resp, r)
}

// relayBuffered copies a non-streaming response through, capturing usage.
func (p *Proxy) relayBuffered(w http.ResponseWriter, resp *http.Response, r Request) (Outcome, error) {
	body, err := io.ReadAll(io.LimitReader(resp.Body, p.cfg.MaxResponseBytes))
	if err != nil {
		return Outcome{Status: resp.StatusCode, Retryable: true}, err
	}

	copyHeaders(w.Header(), resp.Header)
	w.WriteHeader(resp.StatusCode)
	n, werr := w.Write(body)

	out := Outcome{
		Wrote:  n > 0,
		Status: resp.StatusCode,
		Bytes:  int64(n),
		Usage:  usageFromJSON(body),
	}
	if r.CaptureReply > 0 && resp.StatusCode == http.StatusOK {
		out.Reply, out.ReplyTruncated = capReply(replyFromJSON(body), r.CaptureReply)
	}
	if werr != nil {
		return out, ErrClientGone
	}
	return out, nil
}

// relayStream relays an SSE response event by event, flushing as it goes.
//
// abort must cancel the context the upstream request was built with: the idle
// watchdog can only stop a silent engine by aborting that read.
func (p *Proxy) relayStream(ctx context.Context, abort context.CancelFunc, w http.ResponseWriter, resp *http.Response, r Request) (result Outcome, rerr error) {
	flusher, ok := w.(http.Flusher)
	if !ok {
		return Outcome{Status: resp.StatusCode, Retryable: false},
			errors.New("response writer does not support streaming")
	}

	copyHeaders(w.Header(), resp.Header)
	w.WriteHeader(resp.StatusCode)
	flusher.Flush()

	out := Outcome{Wrote: true, Status: resp.StatusCode}
	start := time.Now()
	// Nothing is accumulated unless a thread asked for it, so a request that
	// stores nothing pays nothing on the hottest path in the system.
	reply := replyBuilder{cap: r.CaptureReply}
	// On every exit rather than at the one that looks like the end: a stream
	// can also end with the client gone or the engine stalled, and what the
	// assistant said up to that point is still what it said.
	defer func() {
		result.Reply, result.ReplyTruncated = reply.result()
	}()

	// One watchdog covers both phases: until the first event it enforces
	// FirstTokenTimeout, afterwards IdleTimeout.
	wd := newWatchdog(p.cfg.FirstTokenTimeout)
	defer wd.run(ctx, abort)()

	br := bufio.NewReaderSize(resp.Body, 16<<10)
	var ev eventBuffer

	for {
		line, err := br.ReadString('\n')
		if len(line) > 0 {
			wd.beat()

			ev.add(line)
			if ev.complete() {
				if out.TTFT == 0 {
					out.TTFT = time.Since(start)
					wd.setInterval(p.cfg.IdleTimeout)
				}
				if r.CaptureReply > 0 {
					reply.add(replyDelta(ev.data.String()))
				}
				if u := ev.usage(); u != nil {
					out.Usage = u
					if r.SuppressUsageChunk {
						// The client never asked for usage; record it and drop
						// the chunk rather than changing what they receive.
						ev.reset()
						continue
					}
				}
				n, werr := io.WriteString(w, ev.string())
				out.Bytes += int64(n)
				ev.reset()
				if werr != nil {
					return out, ErrClientGone
				}
				flusher.Flush()
			}
		}

		if err != nil {
			if errors.Is(err, io.EOF) {
				// Flush a trailing event that never got its blank line.
				if s := ev.string(); s != "" {
					n, _ := io.WriteString(w, s)
					out.Bytes += int64(n)
					flusher.Flush()
				}
				return out, nil
			}
			if wd.fired() {
				return out, ErrUpstreamStalled
			}
			if ctx.Err() != nil {
				return out, ErrClientGone
			}
			return out, err
		}
	}
}

// --- SSE event assembly ---

// eventBuffer accumulates the lines of one SSE event. Events are relayed whole
// rather than line by line so a usage-only chunk can be dropped atomically.
type eventBuffer struct {
	lines []string
	data  strings.Builder
}

func (e *eventBuffer) add(line string) {
	e.lines = append(e.lines, line)
	if payload, ok := strings.CutPrefix(line, "data: "); ok {
		e.data.WriteString(strings.TrimRight(payload, "\r\n"))
	}
}

// complete reports whether the blank line terminating an event has arrived.
func (e *eventBuffer) complete() bool {
	if len(e.lines) == 0 {
		return false
	}
	last := e.lines[len(e.lines)-1]
	return last == "\n" || last == "\r\n"
}

func (e *eventBuffer) string() string {
	if len(e.lines) == 0 {
		return ""
	}
	return strings.Join(e.lines, "")
}

func (e *eventBuffer) reset() {
	e.lines = e.lines[:0]
	e.data.Reset()
}

// usage returns the token usage carried by this event, if any.
func (e *eventBuffer) usage() *Usage {
	s := e.data.String()
	if s == "" || s == "[DONE]" || !strings.Contains(s, "\"usage\"") {
		return nil
	}
	return usageFromJSON([]byte(s))
}

// --- watchdog ---

// watchdog aborts the upstream read once an interval passes with no beat, and
// records that it did so, so the resulting read error can be attributed to the
// engine stalling rather than to the client going away.
type watchdog struct {
	interval atomic.Int64
	last     atomic.Int64
	tripped  atomic.Bool
}

func newWatchdog(interval time.Duration) *watchdog {
	w := &watchdog{}
	w.interval.Store(int64(interval))
	w.last.Store(time.Now().UnixNano())
	return w
}

func (w *watchdog) beat() { w.last.Store(time.Now().UnixNano()) }

func (w *watchdog) setInterval(d time.Duration) { w.interval.Store(int64(d)) }

func (w *watchdog) fired() bool { return w.tripped.Load() }

// run starts the watchdog and returns a stop function. abort is called if an
// interval elapses without a beat.
func (w *watchdog) run(ctx context.Context, abort context.CancelFunc) func() {
	done := make(chan struct{})

	go func() {
		// Tick well below the shortest interval so the overshoot stays small.
		// The interval shrinks once the first token arrives, so derive the tick
		// from the smaller of the two rather than from whatever is set now.
		tick := max(time.Duration(w.interval.Load())/4, 10*time.Millisecond)
		t := time.NewTicker(tick)
		defer t.Stop()
		for {
			select {
			case <-ctx.Done():
				return
			case <-done:
				return
			case <-t.C:
				interval := time.Duration(w.interval.Load())
				if time.Since(time.Unix(0, w.last.Load())) > interval {
					w.tripped.Store(true)
					abort()
					return
				}
				// Keep the tick proportional as the interval changes.
				if want := max(interval/4, 10*time.Millisecond); want != tick {
					tick = want
					t.Reset(tick)
				}
			}
		}
	}()

	var once sync.Once
	return func() { once.Do(func() { close(done) }) }
}

// --- helpers ---

func retryableStatus(code int) bool {
	switch code {
	case http.StatusTooManyRequests,
		http.StatusBadGateway,
		http.StatusServiceUnavailable,
		http.StatusGatewayTimeout:
		return true
	default:
		return false
	}
}

func isEventStream(h http.Header) bool {
	return strings.HasPrefix(h.Get("Content-Type"), "text/event-stream")
}

// hopByHop headers must not be forwarded (RFC 9110 §7.6.1).
var hopByHop = map[string]bool{
	"Connection":          true,
	"Proxy-Connection":    true,
	"Keep-Alive":          true,
	"Proxy-Authenticate":  true,
	"Proxy-Authorization": true,
	"Te":                  true,
	"Trailer":             true,
	"Transfer-Encoding":   true,
	"Upgrade":             true,
}

func copyHeaders(dst, src http.Header) {
	for k, vs := range src {
		if hopByHop[http.CanonicalHeaderKey(k)] {
			continue
		}
		for _, v := range vs {
			dst.Add(k, v)
		}
	}
}

// cloneForUpstream forwards only what the engine needs. The client's
// Authorization header is deliberately dropped: it authenticates the caller to
// llmgw, not llmgw to the engine. llmgw's own key for the engine is set by Do.
func cloneForUpstream(h http.Header) http.Header {
	out := make(http.Header, 4)
	for _, k := range []string{"Accept", "User-Agent", "Traceparent", "Tracestate", "X-Request-Id"} {
		if v := h.Get(k); v != "" {
			out.Set(k, v)
		}
	}
	return out
}

func ctxCauseIsClient(ctx context.Context) bool {
	return errors.Is(context.Cause(ctx), context.Canceled)
}