feat(cag): CAG integration — memory adapter, warmup semaphore, cache-aware routing (#47)
CI / CI (push) Successful in 3m20s

## Summary

Three issues in one PR. Implements Cache-Augmented Generation (CAG) support in the api-gw.

---

### #44 — Memory ServiceAdapter: standing-context resource

- Code-registers `memory` adapter with `standing-context` resource (`POST /memory/standing-context`)
- Fixes dispatcher bug: method-level `Auth.Capability` override now respected (was always using verb-derived capability)
- Auth: `memory:read` required via method-level override; `memory:write` for ingest
- `MEMORY_UPSTREAM_URL` env configures upstream

### #45 — Continuous batching: warmup deprioritisation + per-backend semaphore

- In-memory bounded queue per backend: `VLLM_MAX_QUEUE_DEPTH` slots (default 20)
- Concurrency limiter: `VLLM_MAX_NUM_SEQS` slots (default 4)
- `X-Poimen-Warmup: true` → sleep `WARMUP_DELAY_MS` (default 500ms) before queue entry, strip header, add `X-Priority: low`
- Queue full → 503 immediately, no goroutine pile-up
- Metrics: `warmup_requests_delayed_total`, `queue_full_requests_total` at `/metrics`

### #46 — Cache-aware routing: rendezvous hash + prefix cache hit logging

- `addresses: [...]` YAML field for multi-backend models (backward-compat with `address`)
- `X-Poimen-Cache-Key` header → rendezvous-hash backend selection (same key = same GPU)
- No cache key → round-robin
- Failed backend excluded 30s; next-highest-score fallback
- `vllm-num-cached-tokens` response header captured → `vllm_prefix_cache_hit_tokens_total` counter + structured log

---

## Verification

```
go build ./...    ✅
go vet ./...      ✅
go test -race ./... ✅ (all packages, 0 races)
```

## Reviewer verdicts

- reviewer-sre: SIGN-OFF
- reviewer-quality: SIGN-OFF

## Non-blocking findings (🟡🔵 only)

- `getOrCreateTransport` not goroutine-safe by design; pre-population invariant documented
- `time.Sleep(warmupDelay)` holds goroutine; context-cancel alternative left as TODO
- OTel trace_id not propagated into CAG log events (existing gap)

Closes #44 Closes #45 Closes #46

---------

Co-authored-by: Admin Bot <[email protected]>
Reviewed-on: #47
This commit was merged in pull request #47.
This commit is contained in:
2026-09-27 11:16:40 +00:00
co-authored by Admin Bot
parent aaa9a5d1cb
commit ef5e736673
16 changed files with 1597 additions and 162 deletions
+343 -148
View File
@@ -2,15 +2,20 @@
package proxy
import (
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net"
"net/http"
"net/http/httputil"
"net/url"
"os"
"sort"
"strconv"
"strings"
"sync"
"syscall"
"time"
@@ -18,6 +23,7 @@ import (
"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"
)
@@ -37,13 +43,20 @@ type Handler struct {
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
// transportsMu guards h.transports for goroutine-safe reads/writes.
transportsMu sync.RWMutex
}
// Route represents a reverse proxy route.
type Route struct {
Name string
Upstream *config.Upstream
Director func(*http.Request)
Name string
ModelName string // non-empty for model-based routes; used for scheduler lookup
Upstream *config.Upstream
Director func(*http.Request)
Transport *http.Transport
}
@@ -93,6 +106,47 @@ func New(cfg *config.Config) *Handler {
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 := 20
if v := os.Getenv("VLLM_MAX_QUEUE_DEPTH"); v != "" {
if n, err := strconv.Atoi(v); err == nil && n > 0 {
maxQueue = n
}
}
h.schedulers = make(map[string]*Scheduler)
for modelName, model := range cfg.Models {
addrs := model.AllAddresses()
if len(addrs) > 0 {
h.schedulers[modelName] = NewScheduler(addrs, maxSlots, maxQueue)
}
}
// 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)
@@ -116,7 +170,16 @@ func New(cfg *config.Config) *Handler {
// 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
}
@@ -165,7 +228,7 @@ func (h *Handler) getOrCreateTransport(addr string, up *config.Upstream) *http.T
}
h.transports[addr] = transport
return transport
return transport // written under transportsMu write lock
}
// directorFunc modifies the request to be sent to the upstream.
@@ -251,17 +314,7 @@ func writeProblemDetail(w http.ResponseWriter, status int, problemType, title, d
// ServeHTTP implements http.Handler.
func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
// Auth endpoints — no JWT required (they issue tokens)
if r.URL.Path == "/auth/token" {
h.handleAuthToken(w, r)
return
}
if r.URL.Path == "/auth/refresh" {
h.handleAuthRefresh(w, r)
return
}
if r.URL.Path == "/auth/exchange" {
h.handleAuthExchange(w, r)
if h.handleAuthEndpoint(w, r) {
return
}
@@ -271,58 +324,146 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
return
}
// Try to find a matching route (including body-based dispatch for /v1/chat/completions).
// Note: /workflows endpoint is deprecated. Use X-Service: workflow + X-Resource headers instead.
// 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")
proxy.ServeHTTP(w, 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, and ok=true.
// If the slot cannot be acquired it writes a 503 and returns (nil, nil, false).
func (h *Handler) acquireScheduledBackend(ctx context.Context, w http.ResponseWriter, r *http.Request, sched *Scheduler, route *Route, cacheKey string) (*Backend, *Route, bool) {
backend := sched.Select(cacheKey)
if !h.applyWarmup(ctx, r) {
return nil, nil, false
}
if backend == nil {
return nil, route, true
}
if err := backend.Acquire(ctx); err != nil {
if errors.Is(err, ErrQueueFull) {
observability.Global().IncrQueueFull()
logging.Warnf("backend queue full", map[string]string{
"backend": backend.Address,
"path": r.URL.Path,
"model": route.ModelName,
})
w.Header().Set("Retry-After", "5")
w.WriteHeader(http.StatusServiceUnavailable)
fmt.Fprintf(w, "queue full")
} else {
logging.Warnf("backend acquire failed", map[string]string{
"backend": backend.Address,
"path": r.URL.Path,
"model": route.ModelName,
"err": err.Error(),
})
w.WriteHeader(http.StatusServiceUnavailable)
fmt.Fprintf(w, "upstream busy")
}
return nil, nil, false
}
if backend.Address != route.Upstream.Address {
patchedCfg := *route.Upstream
patchedCfg.Address = backend.Address
patchedURL, _ := url.Parse("http://" + 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, true
}
// 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)
// Check if this is a model validation error (from body-based dispatch)
if validationErr, ok := err.(*modelValidationError); ok {
// This is a client error, not a routing error
var status int
var problemType string
var title string
var detail 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"
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)
default:
status = http.StatusBadRequest
problemType = "https://api.example.com/problems/invalid-request"
title = "Invalid Request"
detail = validationErr.Message
}
// Get list of valid models (only for model-related errors)
var validModels []string
if validationErr.Kind == "unknown_model" || validationErr.Kind == "missing_model" || validationErr.Kind == "empty_model" || validationErr.Kind == "null_model" {
validModels = h.getValidModels()
}
writeProblemDetail(w, status, problemType, title, detail, validModels)
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
return nil, false
}
if err != nil || route == nil {
// Route not found or error determining route
w.WriteHeader(http.StatusNotFound)
fmt.Fprintf(w, "not found")
if err != nil {
@@ -333,126 +474,180 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
} else {
logging.Errorf("no route matches", fmt.Errorf("path=%s method=%s", r.URL.Path, r.Method), nil)
}
return
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
}
// Strip spoofed identity headers from all inbound requests.
// Must happen before any routing — even unauthenticated paths.
identity.StripIncoming(r)
writeProblemDetail(w, status, problemType, title, detail, models)
}
// JWT Authentication for /v1/* endpoints
if h.jwtValidator != nil && strings.HasPrefix(r.URL.Path, "/v1/") {
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
}
claims, err := h.jwtValidator.ValidateBearerToken(authHeader)
if err != nil {
// 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 h.config.Auth.RequiredCapability != "" {
if !h.jwtValidator.CheckPermissions(claims, h.config.Auth.RequiredCapability, "*") {
writeProblemDetail(w, http.StatusForbidden,
"https://api.example.com/problems/forbidden",
"Forbidden",
"JWT validation failed",
"https://api.example.com/problems/insufficient-permissions",
"Insufficient Permissions",
fmt.Sprintf("Required capability: %s", h.config.Auth.RequiredCapability),
nil)
logging.Errorf("auth failed", err, map[string]string{
"path": r.URL.Path,
logging.Errorf("auth failed", fmt.Errorf("insufficient permissions"), map[string]string{
"path": r.URL.Path,
"required": h.config.Auth.RequiredCapability,
})
return
return false
}
}
return true
}
// Inject identity headers for downstream services
identity.Inject(r, claims)
// 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
}
// Check required capability if configured
if h.config.Auth.RequiredCapability != "" {
if !h.jwtValidator.CheckPermissions(claims, h.config.Auth.RequiredCapability, "*") {
writeProblemDetail(w, http.StatusForbidden,
"https://api.example.com/problems/insufficient-permissions",
"Insufficient Permissions",
fmt.Sprintf("Required capability: %s", h.config.Auth.RequiredCapability),
nil)
logging.Errorf("auth failed", fmt.Errorf("insufficient permissions"), map[string]string{
"path": r.URL.Path,
"required": h.config.Auth.RequiredCapability,
})
return
}
// 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
if sched != nil {
var ok bool
selectedBackend, route, ok = h.acquireScheduledBackend(ctx, w, r, sched, route, cacheKey)
if !ok {
return nil, nil, false
}
}
// Note: Body size checking already happened in RouteRequest (body was read for model dispatch).
// For other paths, we still need to enforce the cap.
// For /v1/chat/completions, the body was already read and validated.
// Enforce request body size cap for non-chat routes
if r.URL.Path != "/v1/chat/completions" {
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
}
// Wrap request body with size limiter
// This enforces the cap at read time, not after buffering
if route.Upstream.MaxBodySize > 0 && r.Body != nil {
r.Body = io.NopCloser(io.LimitReader(r.Body, route.Upstream.MaxBodySize))
}
}
// Create the reverse proxy
proxy := httputil.NewSingleHostReverseProxy(&url.URL{
Scheme: "http",
Host: route.Upstream.Address,
})
// Set the director to apply path rewriting
proxy.Director = route.Director
// Use the connection-pooled transport wrapped with tracing
proxy.Transport = tracing.NewTransport(route.Transport)
// Set error handler to log upstream errors
capturedBackend := selectedBackend
capturedRoute := route
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)
}
}
w.WriteHeader(http.StatusBadGateway)
fmt.Fprintf(w, "upstream error")
logging.Errorf("upstream error", err, map[string]string{
"upstream": route.Upstream.Address,
"upstream": capturedRoute.Upstream.Address,
"path": r.URL.Path,
})
}
// Note: We apply connection timeout via Transport dialer, but NOT read timeout as a context deadline.
// Read timeout should apply to inactivity (socket read timeout), not total request duration.
// A streaming response that's continuously sending should not be cut off.
// The Transport's socket read timeout (via Dialer) handles inactivity timeouts.
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
}
}
// For streaming responses (SSE, chunked), disable buffering to ensure events
// reach clients immediately. Issue #33: X-Accel-Buffering:no tells nginx/Ingress
// to stream instead of buffer. ResponseController.Flush() in upstream handler
// pairs with this to deliver unbuffered chunks.
w.Header().Set("X-Accel-Buffering", "no")
// Serve the request through the proxy
proxy.ServeHTTP(w, r)
return proxy, selectedBackend, true
}
// getValidModels returns a sorted list of all configured model names.
func (h *Handler) getValidModels() []string {
var models []string