homelab-frontend/internal/proxy/proxy.go

// Package proxy provides reverse proxying to configured upstreams.
package proxy

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/golang-jwt/jwt/v5"
	"io"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"sort"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"forgejo.riotpiao.com/rock/homelab-frontend/internal/auth"
	"forgejo.riotpiao.com/rock/homelab-frontend/internal/config"
	"forgejo.riotpiao.com/rock/homelab-frontend/internal/identity"
	"forgejo.riotpiao.com/rock/homelab-frontend/internal/logging"
	"forgejo.riotpiao.com/rock/homelab-frontend/internal/observability"
	"forgejo.riotpiao.com/rock/homelab-frontend/internal/tracing"
)

// Handler is a reverse proxy that routes requests to configured upstreams.
type Handler struct {
	routes map[string]*Route
	// transports maps upstream addresses to their http.Transport for connection reuse
	transports map[string]*http.Transport
	// config holds the gateway configuration (for model registry, etc.)
	config *config.Config
	// jwtValidator validates JWT tokens for authenticated endpoints
	jwtValidator *auth.Validator
	// authHTTP is the HTTP client for token exchange with the identity provider.
	authHTTP authClient
	// Default timeouts for synthesized routes (model-based dispatch)
	defaultConnectTimeout time.Duration
	defaultReadTimeout    time.Duration
	defaultWriteTimeout   time.Duration
	defaultMaxBodySize    int64
	// CAG: per-model backend schedulers for cache-aware routing (#45/#46)
	schedulers map[string]*Scheduler
	// CAG: delay applied to warmup requests before queue entry (#45)
	warmupDelay time.Duration
	// engineObservers polls each model backend's /metrics for admission (#64)
	engineObservers map[string]*EngineObserver
	// maxBodyBytes caps synthesized model-route request bodies (MAX_BODY_BYTES)
	maxBodyBytes int64
	// transportsMu guards h.transports for goroutine-safe reads/writes.
	transportsMu sync.RWMutex
}

// Route represents a reverse proxy route.
type Route struct {
	Name      string
	ModelName string // non-empty for model-based routes; used for scheduler lookup
	Upstream  *config.Upstream
	Director  func(*http.Request)
	Transport *http.Transport
}

// Error types for model validation
type modelValidationError struct {
	Kind    string // "invalid_json", "missing_model", "unknown_model"
	Message string
	Model   string // only for unknown_model
}

func (e *modelValidationError) Error() string {
	return e.Message
}

// RFC 9457 Problem Details
type problemDetail struct {
	Type        string   `json:"type"`
	Title       string   `json:"title"`
	Status      int      `json:"status"`
	Detail      string   `json:"detail"`
	ValidModels []string `json:"valid_models,omitempty"`
}

// New creates a new reverse proxy handler from configuration.
// It sets up connection pooling and rewriting rules for each route.
func New(cfg *config.Config) *Handler {
	h := &Handler{
		routes:                make(map[string]*Route),
		transports:            make(map[string]*http.Transport),
		config:                cfg,
		defaultConnectTimeout: 10 * time.Second,
		defaultReadTimeout:    1 * time.Hour,
		defaultWriteTimeout:   1 * time.Hour,
		defaultMaxBodySize:    100 * 1024 * 1024,
	}

	// Initialize JWT validator if auth is enabled.
	//
	// Two paths, in order of preference:
	//  1. ZitadelDomain set -> official ZITADEL Go SDK (discovery + remote JWKS
	//     + op.AccessTokenVerifier). This is the production path.
	//  2. JWKSURL set -> legacy raw-JWKS validation, kept for tests and for
	//     any deployment that has not moved to ZITADEL yet.
	if cfg.Auth.Enabled && cfg.Auth.ZitadelDomain != "" {
		ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
		defer cancel()
		v, err := auth.NewZitadelValidator(ctx, cfg.Auth.Issuer, cfg.Auth.ZitadelClientIDs, auth.ZitadelConfig{
			Domain:       cfg.Auth.ZitadelDomain,
			InsecurePort: cfg.Auth.ZitadelInsecurePort,
			ClientIDs:    cfg.Auth.ZitadelClientIDs,
		})
		if err != nil {
			logging.Errorf("zitadel: verifier init failed; auth disabled until restart", err, map[string]string{
				"domain": cfg.Auth.ZitadelDomain,
			})
		} else {
			h.jwtValidator = v
		}
	} else if cfg.Auth.Enabled && cfg.Auth.JWKSURL != "" {
		h.jwtValidator = auth.NewValidator(
			cfg.Auth.Issuer,
			cfg.Auth.Audience,
			cfg.Auth.JWKSURL,
		)
	}

	if cfg.Auth.TokenURL != "" {
		h.authHTTP = newAuthClient()
	}

	// CAG: warmup delay from env (#45)
	warmupMs := int64(500)
	if v := os.Getenv("WARMUP_DELAY_MS"); v != "" {
		if n, err := strconv.ParseInt(v, 10, 64); err == nil && n >= 0 {
			warmupMs = n
		}
	}
	h.warmupDelay = time.Duration(warmupMs) * time.Millisecond

	// CAG: per-model backend schedulers (#45/#46)
	maxSlots := 4
	if v := os.Getenv("VLLM_MAX_NUM_SEQS"); v != "" {
		if n, err := strconv.Atoi(v); err == nil && n > 0 {
			maxSlots = n
		}
	}
	maxQueue := 8
	if v := os.Getenv("VLLM_MAX_QUEUE_DEPTH"); v != "" {
		if n, err := strconv.Atoi(v); err == nil && n > 0 {
			maxQueue = n
		}
	}

	// Issue #64: engine-wide admission thresholds. MaxInFlight defaults to the
	// same value as VLLM_MAX_NUM_SEQS; the observers below read the engine's
	// own gauges, which every replica shares.
	admissionMaxInFlight := maxSlots
	if v := os.Getenv("ADMISSION_MAX_INFLIGHT"); v != "" {
		if n, err := strconv.Atoi(v); err == nil && n > 0 {
			admissionMaxInFlight = n
		}
	}
	admissionMaxKVUsage := 0.90
	if v := os.Getenv("ADMISSION_MAX_KV_USAGE"); v != "" {
		if f, err := strconv.ParseFloat(v, 64); err == nil && f > 0 {
			admissionMaxKVUsage = f
		}
	}
	maxParkSeconds := 120
	if v := os.Getenv("ADMISSION_MAX_PARK_SECONDS"); v != "" {
		if n, err := strconv.Atoi(v); err == nil && n > 0 {
			maxParkSeconds = n
		}
	}
	h.maxBodyBytes = 1048576
	if v := os.Getenv("MAX_BODY_BYTES"); v != "" {
		if n, err := strconv.ParseInt(v, 10, 64); err == nil && n > 0 {
			h.maxBodyBytes = n
		}
	}

	limits := AdmissionLimits{
		MaxInFlight: admissionMaxInFlight,
		MaxKVUsage:  admissionMaxKVUsage,
		MaxPark:     time.Duration(maxParkSeconds) * time.Second,
	}

	// Issue #64: only backends that actually publish the vLLM scheduler gauges
	// may be observed. `reasoning` is the only vLLM engine here; ornith and
	// qwen-cpu are Ollama and embeddings/rerank are TEI, none of which export
	// vllm:* series. Observing them would trip the startup self-check and hold
	// observer_degraded at 1 forever. Unset (the default) observes nothing and
	// leaves admission exactly as it was before this change.
	observeModels := make(map[string]bool)
	for _, name := range strings.Split(os.Getenv("ADMISSION_OBSERVE_MODELS"), ",") {
		if name = strings.TrimSpace(name); name != "" {
			observeModels[name] = true
		}
	}

	h.engineObservers = make(map[string]*EngineObserver)
	h.schedulers = make(map[string]*Scheduler)
	for modelName, model := range cfg.Models {
		addrs := model.AllAddresses()
		if len(addrs) == 0 {
			continue
		}
		if observeModels[modelName] {
			for _, addr := range addrs {
				if _, ok := h.engineObservers[addr]; !ok {
					h.engineObservers[addr] = NewEngineObserver(addr, enginePollInterval)
				}
			}
		}
		h.schedulers[modelName] = NewScheduler(addrs, maxSlots, maxQueue, h.engineObservers, limits)
	}
	// The handler has no request-path lifecycle to hang cancellation on, so the
	// observers poll for the lifetime of the process.
	for _, o := range h.engineObservers {
		o.Start(context.Background())
	}
	// Pre-populate transports for all scheduler backends at startup (single-threaded).
	// getOrCreateTransport is not goroutine-safe; pre-populating avoids concurrent writes.
	for _, model := range cfg.Models {
		for _, addr := range model.AllAddresses() {
			h.getOrCreateTransport(addr, &config.Upstream{
				ConnectTimeout: h.defaultConnectTimeout,
				ReadTimeout:    h.defaultReadTimeout,
				WriteTimeout:   h.defaultWriteTimeout,
			})
		}
	}

	for name, route := range cfg.Routes {
		// Create a transport per unique upstream address for connection reuse
		transport := h.getOrCreateTransport(route.Upstream.Address, &route.Upstream)

		upstreamURL := &url.URL{Scheme: "http", Host: route.Upstream.Address}

		r := &Route{
			Name:      name,
			Upstream:  &route.Upstream,
			Transport: transport,
			Director: func(req *http.Request) {
				directorFunc(req, upstreamURL, &route.Upstream)
			},
		}
		h.routes[name] = r
	}

	return h
}

// getOrCreateTransport returns a shared http.Transport for the given upstream address.
// This ensures connections are pooled and reused across requests to the same upstream.
func (h *Handler) getOrCreateTransport(addr string, up *config.Upstream) *http.Transport {
	h.transportsMu.RLock()
	if t, ok := h.transports[addr]; ok {
		h.transportsMu.RUnlock()
		return t
	}
	h.transportsMu.RUnlock()

	h.transportsMu.Lock()
	defer h.transportsMu.Unlock()
	if t, ok := h.transports[addr]; ok { // double-check after write lock
		return t
	}

	// Create a transport with timeout settings from the upstream config.
	// Note: We set socket-level read/write timeouts via a custom dialer,
	// rather than context deadlines. Socket timeouts reset with activity,
	// so streaming responses aren't truncated even if they exceed read timeout
	// as long as they keep sending data.
	dialer := &net.Dialer{
		Timeout:   up.ConnectTimeout,
		KeepAlive: 30 * time.Second,
		// Issue #31: TCP_NODELAY disables Nagle's algorithm, reducing latency
		// for streaming responses by sending small packets immediately instead of
		// waiting for larger batches. Critical for low-latency LLM token streaming.
		Control: func(network, address string, c syscall.RawConn) error {
			return c.Control(func(fd uintptr) {
				// TCP_NODELAY disables Nagle's algorithm for immediate packet transmission
				_ = syscall.SetsockoptInt(int(fd), syscall.IPPROTO_TCP, syscall.TCP_NODELAY, 1)
			})
		},
	}

	transport := &http.Transport{
		Dial:            dialer.Dial,
		DialContext:     dialer.DialContext,
		MaxIdleConns:    100,
		IdleConnTimeout: 90 * time.Second,
		// Issue #32: Increase per-host connection limit to support HTTP/2 multiplexing.
		// With HTTP/2, we can serve many concurrent streams over fewer connections,
		// but we still allow more connections for better resource utilization.
		MaxConnsPerHost: 10,
		// Allow persistent connections
		DisableKeepAlives: false,
		// Issue #31: Enable HTTP/2 for client connections to support multiplexing.
		// This allows concurrent requests to stream simultaneously with better flow control.
		ForceAttemptHTTP2: true,
	}

	// Store the upstream config for use in the handler
	if transport.TLSClientConfig == nil {
		// We can't directly set socket timeouts on http.Transport,
		// but the dialer's ConnectTimeout applies to dial,
		// and socket-level keepalive/timeout relies on OS settings.
		// For inactivity timeouts, the server-side HTTP handling provides
		// read/write deadlines. Client-side, we rely on TCP keepalive.
	}

	h.transports[addr] = transport
	return transport // written under transportsMu write lock
}

// directorFunc modifies the request to be sent to the upstream.
// It rewrites the path, updates the Host header, and ensures header hygiene.
func directorFunc(req *http.Request, target *url.URL, upstream *config.Upstream) {
	// Apply path rewrite if configured
	if upstream.PathRewrite != "" {
		req.URL.Path = upstream.PathRewrite
	}

	// Set the scheme and host
	req.URL.Scheme = target.Scheme
	req.URL.Host = target.Host

	// Update the Host header to the upstream address
	req.Host = target.Host

	// Strip hop-by-hop headers as defined in RFC 7230 Section 6.1
	// These must not be forwarded to upstream
	hopByHopHeaders := map[string]bool{
		"connection":          true,
		"keep-alive":          true,
		"proxy-authenticate":  true,
		"proxy-authorization": true,
		"te":                  true,
		"trailers":            true,
		"transfer-encoding":   true,
		"upgrade":             true,
	}

	// Also strip any headers listed in the Connection header
	if conn := req.Header.Get("Connection"); conn != "" {
		for _, h := range strings.Split(conn, ",") {
			hopByHopHeaders[strings.ToLower(strings.TrimSpace(h))] = true
		}
	}

	// Remove all hop-by-hop headers
	// The http.Header.Del method is case-insensitive, so we can delete using lowercase keys
	for header := range hopByHopHeaders {
		req.Header.Del(header)
	}

	// Handle X-Forwarded-For: append the immediate peer
	// Get the peer IP from the request RemoteAddr
	peerIP := getPeerIP(req.RemoteAddr)
	if xForwardedFor := req.Header.Get("X-Forwarded-For"); xForwardedFor != "" {
		// Append the peer IP to the existing X-Forwarded-For
		req.Header.Set("X-Forwarded-For", xForwardedFor+", "+peerIP)
	} else {
		// Create a new X-Forwarded-For with just the peer IP
		req.Header.Set("X-Forwarded-For", peerIP)
	}
}

// getPeerIP extracts the IP address from a RemoteAddr string (format: "IP:port")
func getPeerIP(remoteAddr string) string {
	if remoteAddr == "" {
		return ""
	}
	// RemoteAddr is "IP:port", extract just the IP
	if idx := strings.LastIndex(remoteAddr, ":"); idx != -1 {
		return remoteAddr[:idx]
	}
	return remoteAddr
}

// writeProblemDetail writes an RFC 9457 problem detail response.
func writeProblemDetail(w http.ResponseWriter, status int, problemType, title, detail string, validModels []string) {
	w.Header().Set("Content-Type", "application/problem+json")
	w.WriteHeader(status)

	problem := problemDetail{
		Type:        problemType,
		Title:       title,
		Status:      status,
		Detail:      detail,
		ValidModels: validModels,
	}

	json.NewEncoder(w).Encode(problem)
}

// ServeHTTP implements http.Handler.
func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	// Captured before any work so TTFT measures from the moment the gateway
	// accepted the request, not from when the upstream response is copied.
	requestStart := time.Now()

	if h.handleAuthEndpoint(w, r) {
		return
	}

	// Handle /v1/models endpoint (no routing needed, derived from config)
	if r.URL.Path == "/v1/models" && r.Method == "GET" {
		h.handleModelsEndpoint(w, r)
		return
	}

	// Resolve route; handles validation errors and 404.
	route, ok := h.routeOrError(w, r)
	if !ok {
		return
	}

	// Strip spoofed identity headers from all inbound requests.
	// Must happen before any routing — even unauthenticated paths.
	identity.StripIncoming(r)

	// JWT Authentication for /v1/* endpoints.
	if !h.authenticateV1(w, r) {
		return
	}

	// Enforce request body size cap for non-chat routes.
	if !h.enforceBodyLimit(w, r, route) {
		return
	}

	// CAG: warmup deprioritisation + per-backend semaphore (issues #45/#46).
	cacheKey := r.Header.Get("X-Poimen-Cache-Key")
	sched := h.schedulers[route.ModelName]
	proxy, backend, ok := h.setupCAGProxy(r.Context(), w, r, route, sched, cacheKey)
	if !ok {
		return
	}
	if backend != nil {
		defer backend.Release()
	}

	// For streaming responses (SSE, chunked), disable buffering. Issue #33.
	w.Header().Set("X-Accel-Buffering", "no")

	// Capture per-request LLM latency metrics (TTFT/ITL/tokens) while the
	// response is copied from the upstream. Inert for non-SSE responses.
	var respWriter http.ResponseWriter = w
	if route.ModelName != "" {
		respWriter = NewResponseWriterWrapper(w, route.ModelName, observability.Global(), requestStart)
	}

	proxy.ServeHTTP(respWriter, r)
}

// applyWarmup strips the X-Poimen-Warmup header, delays the request while
// honouring context cancellation, and records the metric.
// Returns false if the context was cancelled during the delay (caller must abort).
func (h *Handler) applyWarmup(ctx context.Context, r *http.Request) bool {
	if r.Header.Get("X-Poimen-Warmup") != "true" {
		return true
	}
	r.Header.Del("X-Poimen-Warmup")
	r.Header.Set("X-Priority", "low")
	select {
	case <-time.After(h.warmupDelay):
	case <-ctx.Done():
		return false
	}
	observability.Global().IncrWarmupDelayed()
	return true
}

// acquireScheduledBackend selects a backend from sched, applies warmup
// deprioritisation, acquires a concurrency slot, and patches route when the
// selected backend differs from the config default. It returns the backend
// to Release (caller must defer), the (possibly patched) route, whether SSE
// headers were already written while parking, and ok=true.
// If the slot cannot be acquired it writes a 503 and returns ok=false.
func (h *Handler) acquireScheduledBackend(ctx context.Context, w http.ResponseWriter, r *http.Request, sched *Scheduler, route *Route, cacheKey string) (*Backend, *Route, bool, bool) {
	backend := sched.Select(cacheKey)
	if !h.applyWarmup(ctx, r) {
		return nil, nil, false, false
	}
	if backend == nil {
		return nil, route, false, true
	}

	// Decision Q22: a request that actually parks on the engine gate and wants
	// an SSE stream gets its headers pre-written plus a 15 s keep-alive, so no
	// intermediary times the connection out before the engine admits it.
	var (
		prewritten bool
		hbStop     chan struct{}
		hbDone     chan struct{}
	)
	onPark := func() {
		if !requestWantsStream(r) {
			return
		}
		flusher, ok := w.(http.Flusher)
		if !ok {
			return
		}
		prewritten = true
		hbStop = make(chan struct{})
		hbDone = make(chan struct{})
		w.Header().Set("Content-Type", "text/event-stream")
		w.Header().Set("Cache-Control", "no-cache")
		w.Header().Set("Connection", "keep-alive")
		w.Header().Set("X-Accel-Buffering", "no")
		w.WriteHeader(http.StatusOK)
		fmt.Fprint(w, ": keep-alive\n\n")
		flusher.Flush()
		go func() {
			defer close(hbDone)
			ticker := time.NewTicker(15 * time.Second)
			defer ticker.Stop()
			for {
				select {
				case <-hbStop:
					return
				case <-ctx.Done():
					return
				case <-ticker.C:
					fmt.Fprint(w, ": keep-alive\n\n")
					flusher.Flush()
				}
			}
		}()
	}

	err := backend.acquire(ctx, onPark)
	if prewritten {
		close(hbStop)
		<-hbDone
	}
	if err != nil {
		// Once the SSE status line is spent, report the rejection as an event
		// instead of a second WriteHeader.
		reject := func(body string) {
			if prewritten {
				fmt.Fprintf(w, "event: error\ndata: {\"error\":%q}\n\n", body)
				if f, ok := w.(http.Flusher); ok {
					f.Flush()
				}
				return
			}
			w.WriteHeader(http.StatusServiceUnavailable)
			fmt.Fprintf(w, "%s", body)
		}
		if errors.Is(err, ErrAdmissionTimeout) {
			observability.Global().IncrAdmissionRejected("park_timeout")
			logging.Warnf("admission timeout", map[string]string{
				"backend": backend.Address,
				"path":    r.URL.Path,
				"model":   route.ModelName,
			})
			if !prewritten {
				w.Header().Set("Retry-After", "5")
			}
			reject("admission timeout")
		} else if errors.Is(err, ErrQueueFull) {
			observability.Global().IncrQueueFull()
			observability.Global().IncrAdmissionRejected("queue_full")
			logging.Warnf("backend queue full", map[string]string{
				"backend": backend.Address,
				"path":    r.URL.Path,
				"model":   route.ModelName,
			})
			if !prewritten {
				w.Header().Set("Retry-After", "5")
			}
			reject("queue full")
		} else {
			logging.Warnf("backend acquire failed", map[string]string{
				"backend": backend.Address,
				"path":    r.URL.Path,
				"model":   route.ModelName,
				"err":     err.Error(),
			})
			reject("upstream busy")
		}
		return nil, nil, prewritten, false
	}
	if backend.Address != route.Upstream.Address {
		patchedCfg := *route.Upstream
		patchedCfg.Address = backend.Address
		patchedURL := &url.URL{Scheme: "http", Host: backend.Address}
		route = &Route{
			Name:      route.Name,
			ModelName: route.ModelName,
			Upstream:  &patchedCfg,
			Transport: h.getOrCreateTransport(backend.Address, &patchedCfg),
			Director: func(req *http.Request) {
				directorFunc(req, patchedURL, &patchedCfg)
			},
		}
	}
	return backend, route, prewritten, true
}

// requestWantsStream reports whether the JSON request body asks for an SSE
// stream ("stream": true). The body is restored for the upstream.
func requestWantsStream(r *http.Request) bool {
	if r.Body == nil {
		return false
	}
	body, err := io.ReadAll(r.Body)
	if err != nil {
		return false
	}
	r.Body = io.NopCloser(bytes.NewReader(body))
	var probe struct {
		Stream bool `json:"stream"`
	}
	if err := json.Unmarshal(body, &probe); err != nil {
		return false
	}
	return probe.Stream
}

// handleAuthEndpoint routes /auth/* requests and returns true if handled.
func (h *Handler) handleAuthEndpoint(w http.ResponseWriter, r *http.Request) bool {
	switch r.URL.Path {
	case "/auth/token":
		h.handleAuthToken(w, r)
		return true
	case "/auth/refresh":
		h.handleAuthRefresh(w, r)
		return true
	case "/auth/exchange":
		h.handleAuthExchange(w, r)
		return true
	}
	return false
}

// routeOrError resolves a route for r, writing error responses on failure.
// Returns (route, true) on success or (nil, false) when a response has been written.
func (h *Handler) routeOrError(w http.ResponseWriter, r *http.Request) (*Route, bool) {
	route, err := h.RouteRequest(r)
	if validationErr, ok := err.(*modelValidationError); ok {
		h.writeModelValidationError(w, validationErr, h.getValidModels())
		logging.Errorf("client error", validationErr, map[string]string{
			"path":   r.URL.Path,
			"method": r.Method,
			"reason": validationErr.Kind,
		})
		return nil, false
	}
	if err != nil || route == nil {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprintf(w, "not found")
		if err != nil {
			logging.Errorf("routing failed", err, map[string]string{
				"path":   r.URL.Path,
				"method": r.Method,
			})
		} else {
			logging.Errorf("no route matches", fmt.Errorf("path=%s method=%s", r.URL.Path, r.Method), nil)
		}
		return nil, false
	}
	return route, true
}

// writeModelValidationError writes a problem-detail response for model validation errors.
func (h *Handler) writeModelValidationError(w http.ResponseWriter, validationErr *modelValidationError, validModels []string) {
	var status int
	var problemType, title, detail string
	var models []string

	switch validationErr.Kind {
	case "invalid_json":
		status = http.StatusBadRequest
		problemType = "https://api.example.com/problems/invalid-request-body"
		title = "Invalid Request Body"
		detail = validationErr.Message
	case "missing_model", "empty_model", "null_model":
		status = http.StatusBadRequest
		problemType = "https://api.example.com/problems/missing-model"
		title = "Missing Model"
		detail = "The 'model' field is required and must be a non-empty string"
		models = validModels
	case "unknown_model":
		status = http.StatusBadRequest
		problemType = "https://api.example.com/problems/unknown-model"
		title = "Unknown Model"
		detail = fmt.Sprintf("Model %q is not available. See valid_models for available options.", validationErr.Model)
		models = validModels
	default:
		status = http.StatusBadRequest
		problemType = "https://api.example.com/problems/invalid-request"
		title = "Invalid Request"
		detail = validationErr.Message
	}

	writeProblemDetail(w, status, problemType, title, detail, models)
}

// authenticateV1 validates JWT for /v1/* paths.
// Returns false when auth fails (response already written).
func (h *Handler) authenticateV1(w http.ResponseWriter, r *http.Request) bool {
	if h.jwtValidator == nil {
		logging.Warnf("JWT auth disabled; /v1/* request unauthenticated", map[string]string{"path": r.URL.Path})
		return true
	}
	if !strings.HasPrefix(r.URL.Path, "/v1/") {
		return true
	}
	authHeader := r.Header.Get("Authorization")
	if authHeader == "" {
		writeProblemDetail(w, http.StatusUnauthorized,
			"https://api.example.com/problems/unauthorized",
			"Unauthorized",
			"Authorization header required",
			nil)
		logging.Errorf("auth failed", fmt.Errorf("missing auth header"), map[string]string{
			"path": r.URL.Path,
		})
		return false
	}
	claims, err := h.jwtValidator.ValidateBearerToken(authHeader)
	if err != nil {
		writeProblemDetail(w, http.StatusForbidden,
			"https://api.example.com/problems/forbidden",
			"Forbidden",
			"JWT validation failed",
			nil)
		logging.Errorf("auth failed", err, map[string]string{
			"path": r.URL.Path,
		})
		return false
	}
	identity.Inject(r, claims)
	if caps := requiredCapabilities(h.config.Auth.RequiredCapability); len(caps) > 0 {
		args := append(caps, "*")
		if !h.jwtValidator.CheckPermissions(claims, args...) && !h.allowListedMachineIdentity(claims) {
			writeProblemDetail(w, http.StatusForbidden,
				"https://api.example.com/problems/insufficient-permissions",
				"Insufficient Permissions",
				fmt.Sprintf("Required capability: %s", strings.Join(caps, " or ")),
				nil)
			logging.Errorf("auth failed", fmt.Errorf("insufficient permissions"), map[string]string{
				"path":     r.URL.Path,
				"required": strings.Join(caps, ","),
			})
			return false
		}
	}
	return true
}

// allowListedMachineIdentity admits a service account whose token carries no roles
// claim, when its client id appears in the explicit allow-list.
//
// ZITADEL does not put roles into a machine-user access token on this instance: the
// grant exists, the scope asks for the roles, and the claim still never arrives. A
// service account therefore cannot satisfy a role check by any configuration. The
// caller is still authenticated — it signed an assertion with a private key that
// ZITADEL holds the public half of — and the allow-list names it explicitly, so the
// authorisation is a reviewed line rather than a wildcard.
func (h *Handler) allowListedMachineIdentity(claims jwt.MapClaims) bool {
	if !h.config.Auth.AllowListedMachineIdentities {
		return false
	}
	if len(rolesClaim(claims)) > 0 {
		// The token states its roles, so the ordinary check has already run and
		// decided. Falling through here would be a second, weaker rule.
		return false
	}
	clientID, _ := claims["client_id"].(string)
	if clientID == "" {
		return false
	}
	for _, allowed := range h.config.Auth.ZitadelClientIDs {
		if allowed == clientID {
			logging.Errorf("auth: admitted an allow-listed machine identity without a roles claim",
				fmt.Errorf("no roles claim"),
				map[string]string{"client_id": clientID})
			return true
		}
	}
	return false
}

// rolesClaim collects the role names from every claim shape the issuers use.
func rolesClaim(claims jwt.MapClaims) []string {
	var out []string
	for _, key := range []string{"roles", "urn:zitadel:iam:org:project:roles"} {
		switch v := claims[key].(type) {
		case []interface{}:
			for _, item := range v {
				if name, ok := item.(string); ok {
					out = append(out, name)
				}
			}
		case map[string]interface{}:
			for name := range v {
				out = append(out, name)
			}
		}
	}
	return out
}

// requiredCapabilities splits the comma-separated requiredCapability setting
// into a list. Several spellings are accepted because the same capability is
// named differently by each identity provider: Authentik used "llm:inference"
// as a permission string, while ZITADEL uses role keys such as
// "llm-inference". A token satisfies the gate when it carries any one of them.
func requiredCapabilities(raw string) []string {
	parts := strings.Split(raw, ",")
	out := make([]string, 0, len(parts))
	for _, p := range parts {
		if p = strings.TrimSpace(p); p != "" {
			out = append(out, p)
		}
	}
	return out
}

// enforceBodyLimit caps request body size for non-chat routes.
// Returns false when the request is rejected (response already written).
func (h *Handler) enforceBodyLimit(w http.ResponseWriter, r *http.Request, route *Route) bool {
	if r.URL.Path == "/v1/chat/completions" {
		return true
	}
	if route.Upstream.MaxBodySize > 0 && r.ContentLength > route.Upstream.MaxBodySize {
		w.WriteHeader(http.StatusRequestEntityTooLarge)
		fmt.Fprintf(w, "request body too large")
		logging.Errorf("request rejected", fmt.Errorf("body_too_large"), map[string]string{
			"reason":         "body_too_large",
			"route":          route.Name,
			"upstream":       route.Upstream.Address,
			"content_length": fmt.Sprintf("%d", r.ContentLength),
			"max_body_size":  fmt.Sprintf("%d", route.Upstream.MaxBodySize),
		})
		return false
	}
	if route.Upstream.MaxBodySize > 0 && r.Body != nil {
		r.Body = io.NopCloser(io.LimitReader(r.Body, route.Upstream.MaxBodySize))
	}
	return true
}

// setupCAGProxy acquires a scheduled backend (if any), builds a reverse proxy,
// wires ErrorHandler and ModifyResponse, and returns (proxy, backend, true).
// Returns (nil, nil, false) when the backend slot cannot be acquired (response written).
func (h *Handler) setupCAGProxy(ctx context.Context, w http.ResponseWriter, r *http.Request, route *Route, sched *Scheduler, cacheKey string) (*httputil.ReverseProxy, *Backend, bool) {
	var selectedBackend *Backend
	var headersPrewritten bool
	if sched != nil {
		var ok bool
		selectedBackend, route, headersPrewritten, ok = h.acquireScheduledBackend(ctx, w, r, sched, route, cacheKey)
		if !ok {
			return nil, nil, false
		}
	}

	proxy := httputil.NewSingleHostReverseProxy(&url.URL{
		Scheme: "http",
		Host:   route.Upstream.Address,
	})
	proxy.Director = route.Director
	proxy.Transport = tracing.NewTransport(route.Transport)

	capturedBackend := selectedBackend
	capturedRoute := route
	if headersPrewritten {
		// The 200 + text/event-stream status is already on the wire, so an
		// upstream status can no longer be forwarded: httputil.ReverseProxy
		// calls WriteHeader again and net/http keeps the first one. Frame a
		// failure as an SSE event rather than letting a raw error body land
		// inside the stream, which a streaming client reads as an empty
		// completion and reports as success.
		proxy.ModifyResponse = func(resp *http.Response) error {
			if resp.StatusCode < http.StatusBadRequest {
				return nil
			}
			body, readErr := io.ReadAll(io.LimitReader(resp.Body, 4096))
			resp.Body.Close()
			if readErr != nil {
				logging.Warnf("reading upstream error body failed", map[string]string{
					"upstream": capturedRoute.Upstream.Address,
					"err":      readErr.Error(),
				})
			}
			logging.Warnf("upstream error after parked stream started", map[string]string{
				"upstream": capturedRoute.Upstream.Address,
				"status":   strconv.Itoa(resp.StatusCode),
			})
			resp.Body = io.NopCloser(strings.NewReader(fmt.Sprintf(
				"event: error\ndata: {\"upstream_status\":%d,\"error\":%q}\n\n",
				resp.StatusCode, strings.TrimSpace(string(body)))))
			resp.ContentLength = -1
			resp.Header.Del("Content-Length")
			resp.Header.Set("Content-Type", "text/event-stream")
			return nil
		}
	}
	proxy.ErrorHandler = func(w http.ResponseWriter, r *http.Request, err error) {
		if capturedBackend != nil {
			if s := h.schedulers[capturedRoute.ModelName]; s != nil {
				s.MarkFailed(capturedBackend.Address)
			}
		}
		if headersPrewritten {
			// SSE headers and a 200 are already on the wire: a 502 here would be
			// an ignored WriteHeader and would tear the event stream. Report the
			// failure as an event instead.
			fmt.Fprint(w, "event: error\ndata: {\"error\":\"upstream error\"}\n\n")
			if f, ok := w.(http.Flusher); ok {
				f.Flush()
			}
			logging.Errorf("upstream error", err, map[string]string{
				"upstream": capturedRoute.Upstream.Address,
				"path":     r.URL.Path,
			})
			return
		}
		w.WriteHeader(http.StatusBadGateway)
		fmt.Fprintf(w, "upstream error")
		logging.Errorf("upstream error", err, map[string]string{
			"upstream": capturedRoute.Upstream.Address,
			"path":     r.URL.Path,
		})
	}

	if cacheKey != "" && r.URL.Path == "/v1/chat/completions" {
		proxy.ModifyResponse = func(resp *http.Response) error {
			raw := resp.Header.Get("vllm-num-cached-tokens")
			if raw == "" {
				return nil
			}
			n, err := strconv.ParseInt(raw, 10, 64)
			if err != nil || n <= 0 {
				return nil
			}
			observability.Global().IncrPrefixCacheHitTokens(n)
			logging.Event(map[string]interface{}{
				"event":         "llm_request",
				"cache_key":     cacheKey,
				"cached_tokens": n,
			})
			return nil
		}
	}

	return proxy, selectedBackend, true
}

// getValidModels returns a sorted list of all configured model names.
func (h *Handler) getValidModels() []string {
	var models []string
	for name := range h.config.Models {
		models = append(models, name)
	}
	sort.Strings(models)
	return models
}

// modelsListResponse represents the response for GET /v1/models
type modelsListResponse struct {
	Object string            `json:"object"`
	Data   []modelsListEntry `json:"data"`
}

// modelsListEntry represents a single model in the list
type modelsListEntry struct {
	ID      string `json:"id"`
	Object  string `json:"object"`
	OwnedBy string `json:"owned_by"`
	Created int64  `json:"created"`
}

// handleModelsEndpoint serves GET /v1/models
// Returns a list of all configured models, derived from config not hardcoded
func (h *Handler) handleModelsEndpoint(w http.ResponseWriter, r *http.Request) {
	// Get all model names from config
	modelNames := h.getValidModels()

	// Build the response
	data := make([]modelsListEntry, len(modelNames))
	for i, name := range modelNames {
		data[i] = modelsListEntry{
			ID:      name,
			Object:  "model",
			OwnedBy: "api.riotpiao.com",
			Created: 1700000000, // Fixed timestamp; can be made configurable if needed
		}
	}

	response := modelsListResponse{
		Object: "list",
		Data:   data,
	}

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)
	json.NewEncoder(w).Encode(response)
}

// Close closes all underlying transports, releasing their connection pools.
func (h *Handler) Close() error {
	for _, transport := range h.transports {
		transport.CloseIdleConnections()
	}
	return nil
}