diff --git a/cmd/gateway/main.go b/cmd/gateway/main.go index ba826d5..15610ef 100644 --- a/cmd/gateway/main.go +++ b/cmd/gateway/main.go @@ -12,6 +12,7 @@ import ( "forgejo.riotpiao.com/rock/homelab-frontend/internal/auth" "forgejo.riotpiao.com/rock/homelab-frontend/internal/config" "forgejo.riotpiao.com/rock/homelab-frontend/internal/notification" + "forgejo.riotpiao.com/rock/homelab-frontend/internal/observability" "forgejo.riotpiao.com/rock/homelab-frontend/internal/proxy" "forgejo.riotpiao.com/rock/homelab-frontend/internal/server" "forgejo.riotpiao.com/rock/homelab-frontend/internal/serviceadapter" @@ -112,6 +113,10 @@ func main() { poimenAdapter := serviceadapter.NewPoimenAdapter(temporalHandler) _ = registry.Add(poimenAdapter) + // Add memory service adapter (code-registered; config memory entry is skipped). + // Includes standing-context resource for CAG support (issue #44). + _ = registry.Add(serviceadapter.NewMemoryAdapter()) + // Add other adapters from config (skip if already registered in code) for _, a := range cfg.Adapters { if existing := registry.Get(a.ServiceName); existing != nil { @@ -135,6 +140,7 @@ func main() { // Create router that handles health endpoints, X-Service (ServiceAdapter) routing, // temporal endpoints, and passes others to upstream handler router := server.NewRouter(healthChecker, dispatcher, temporalHandler, upstreamHandler) + router.SetMetricsHandler(observability.NewMetricsHandler(observability.Global())) // Wrap router with tracing middleware tracedRouter := tracing.Middleware(router) diff --git a/internal/config/config.go b/internal/config/config.go index 89e50ac..2c09268 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -33,7 +33,11 @@ type ModelUpstream struct { // Name is the model name clients send (e.g., "reasoning", "ornith:35b"). Name string // Address is the upstream server address (host:port). + // When Addresses is set, this holds Addresses[0] for backward compatibility. Address string + // Addresses holds all backend addresses for multi-backend cache-aware routing. + // Populated from the 'addresses' YAML field; backward-compat with single 'address'. + Addresses []string // Path is the upstream path for this model (e.g., "/v1/chat/completions"). Path string // UpstreamModel is the model name to send to the upstream server. @@ -116,6 +120,18 @@ func (c *Config) ModelNames() []string { return names } +// AllAddresses returns all configured backend addresses for the model. +// Returns Addresses if non-empty; otherwise wraps Address in a slice. +func (m *ModelUpstream) AllAddresses() []string { + if len(m.Addresses) > 0 { + return m.Addresses + } + if m.Address != "" { + return []string{m.Address} + } + return nil +} + // Load reads configuration from environment variables with defaults. func Load() (*Config, error) { listenAddr := "127.0.0.1:8080" diff --git a/internal/config/loader.go b/internal/config/loader.go index e0d9ed3..a21375c 100644 --- a/internal/config/loader.go +++ b/internal/config/loader.go @@ -37,9 +37,10 @@ type rawRoute struct { // rawModel represents a single model entry in the YAML configuration. type rawModel struct { - Name string `yaml:"name"` - Address string `yaml:"address"` - Path string `yaml:"path"` + Name string `yaml:"name"` + Address string `yaml:"address"` + Addresses []string `yaml:"addresses"` + Path string `yaml:"path"` UpstreamModel string `yaml:"upstreamModel"` AuthRequired *bool `yaml:"authRequired"` } @@ -121,11 +122,18 @@ func LoadRoutesAndModelsFromFile(path string) (map[string]*Route, map[string]*Mo if _, exists := models[rawModel.Name]; exists { return nil, nil, nil, AuthConfig{}, fmt.Errorf("duplicate model: \"%s\"", rawModel.Name) } - if rawModel.Address == "" { - return nil, nil, nil, AuthConfig{}, fmt.Errorf("model \"%s\": field 'address' is required", rawModel.Name) + // Support both 'address' (single) and 'addresses' (multi-backend) YAML fields. + addresses := rawModel.Addresses + if len(addresses) == 0 && rawModel.Address != "" { + addresses = []string{rawModel.Address} } - if _, _, err := net.SplitHostPort(rawModel.Address); err != nil { - return nil, nil, nil, AuthConfig{}, fmt.Errorf("model \"%s\": invalid address \"%s\": %w", rawModel.Name, rawModel.Address, err) + if len(addresses) == 0 { + return nil, nil, nil, AuthConfig{}, fmt.Errorf("model \"%s\": 'address' or 'addresses' is required", rawModel.Name) + } + for _, addr := range addresses { + if _, _, err := net.SplitHostPort(addr); err != nil { + return nil, nil, nil, AuthConfig{}, fmt.Errorf("model \"%s\": invalid address \"%s\": %w", rawModel.Name, addr, err) + } } // Default authRequired to global auth.enabled if not specified per-model authRequired := false @@ -134,7 +142,8 @@ func LoadRoutesAndModelsFromFile(path string) (map[string]*Route, map[string]*Mo } models[rawModel.Name] = &ModelUpstream{ Name: rawModel.Name, - Address: rawModel.Address, + Address: addresses[0], + Addresses: addresses, Path: rawModel.Path, UpstreamModel: rawModel.UpstreamModel, AuthRequired: authRequired, diff --git a/internal/config/models_test.go b/internal/config/models_test.go index 8a5da54..a362d74 100644 --- a/internal/config/models_test.go +++ b/internal/config/models_test.go @@ -161,7 +161,7 @@ models: if err == nil { t.Errorf("expected error for missing address, got nil") } - if err.Error() != "model \"reasoning\": field 'address' is required" { + if err.Error() != "model \"reasoning\": 'address' or 'addresses' is required" { t.Errorf("expected missing address error, got: %v", err) } } @@ -283,3 +283,139 @@ models: t.Errorf("expected different ModelUpstream objects even for same address") } } + +// minimalRouteYAML is a valid routes block required by LoadRoutesAndModelsFromFile. +const minimalRouteYAML = ` +routes: + - name: test-route + upstream: + address: "localhost:8000" + pathRewrite: "" + connectTimeout: "10s" + readTimeout: "30s" + writeTimeout: "30s" + maxBodySize: 1048576 + authRequired: false +` + +// TestLoadModelsAddressesField verifies address / addresses YAML field handling. +func TestLoadModelsAddressesField(t *testing.T) { + tests := []struct { + name string + modelsYAML string + wantErr bool + errContains string + wantAddresses []string + wantAddress string + }{ + { + name: "single address field", + modelsYAML: ` +models: + - name: "mymodel" + address: "host:8080" + path: "/v1/chat" +`, + wantErr: false, + wantAddresses: []string{"host:8080"}, + wantAddress: "host:8080", + }, + { + name: "multi addresses field", + modelsYAML: ` +models: + - name: "mymodel" + addresses: + - "a:80" + - "b:80" + path: "/v1/chat" +`, + wantErr: false, + wantAddresses: []string{"a:80", "b:80"}, + wantAddress: "a:80", + }, + { + name: "neither field", + modelsYAML: ` +models: + - name: "mymodel" + path: "/v1/chat" +`, + wantErr: true, + errContains: "mymodel", + }, + { + name: "invalid address format", + modelsYAML: ` +models: + - name: "mymodel" + address: "notavalidhostport" + path: "/v1/chat" +`, + wantErr: true, + errContains: "mymodel", + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + data := minimalRouteYAML + tc.modelsYAML + + tmpFile, err := os.CreateTemp("", "config-*.yaml") + if err != nil { + t.Fatalf("CreateTemp: %v", err) + } + defer os.Remove(tmpFile.Name()) + tmpFile.WriteString(data) + tmpFile.Close() + + _, models, _, _, loadErr := LoadRoutesAndModelsFromFile(tmpFile.Name()) + + if tc.wantErr { + if loadErr == nil { + t.Fatalf("expected error, got nil") + } + if tc.errContains != "" && !containsStr(loadErr.Error(), tc.errContains) { + t.Errorf("error %q does not contain %q", loadErr.Error(), tc.errContains) + } + return + } + + if loadErr != nil { + t.Fatalf("unexpected error: %v", loadErr) + } + + m, ok := models["mymodel"] + if !ok { + t.Fatalf("model 'mymodel' not found in result") + } + + if m.Address != tc.wantAddress { + t.Errorf("Address: want %q, got %q", tc.wantAddress, m.Address) + } + + if len(m.Addresses) != len(tc.wantAddresses) { + t.Errorf("Addresses len: want %d, got %d (%v)", len(tc.wantAddresses), len(m.Addresses), m.Addresses) + } else { + for i, a := range tc.wantAddresses { + if m.Addresses[i] != a { + t.Errorf("Addresses[%d]: want %q, got %q", i, a, m.Addresses[i]) + } + } + } + }) + } +} + +// containsStr is a helper to avoid importing strings just for Contains. +func containsStr(s, sub string) bool { + return len(s) >= len(sub) && (s == sub || len(sub) == 0 || + func() bool { + for i := 0; i <= len(s)-len(sub); i++ { + if s[i:i+len(sub)] == sub { + return true + } + } + return false + }()) +} diff --git a/internal/logging/logger.go b/internal/logging/logger.go index 2043612..04487e5 100644 --- a/internal/logging/logger.go +++ b/internal/logging/logger.go @@ -198,3 +198,11 @@ func FromContext(ctx context.Context) *Logger { func WithContext(ctx context.Context, l *Logger) context.Context { return context.WithValue(ctx, "logger", l) } + +// Event emits arbitrary fields as a structured JSON line. +// Use when field values include non-string types (numerics, floats). +func Event(fields map[string]interface{}) { + ensure() + data, _ := json.Marshal(fields) + _, _ = global.out.Write(append(data, '\n')) +} diff --git a/internal/observability/metrics.go b/internal/observability/metrics.go index 608d1ad..f78dfa2 100644 --- a/internal/observability/metrics.go +++ b/internal/observability/metrics.go @@ -36,6 +36,11 @@ type Metrics struct { itlMs map[string][]int64 // samples for histogram // Token counts tokenCount map[string]int64 + + // CAG metrics + warmupRequestsDelayed int64 + prefixCacheHitTokens int64 + queueFullRequests int64 } // NewMetrics creates a new Metrics instance. @@ -280,4 +285,61 @@ func (m *Metrics) Reset() { m.ttftMs = make(map[string][]int64) m.itlMs = make(map[string][]int64) m.tokenCount = make(map[string]int64) + m.warmupRequestsDelayed = 0 + m.prefixCacheHitTokens = 0 + m.queueFullRequests = 0 +} + +// IncrQueueFull increments queue_full_requests_total. +func (m *Metrics) IncrQueueFull() { + m.mu.Lock() + defer m.mu.Unlock() + m.queueFullRequests++ +} + +// GetQueueFull returns queue_full_requests_total. +func (m *Metrics) GetQueueFull() int64 { + m.mu.RLock() + defer m.mu.RUnlock() + return m.queueFullRequests +} + +// IncrWarmupDelayed increments warmup_requests_delayed_total. +func (m *Metrics) IncrWarmupDelayed() { + m.mu.Lock() + defer m.mu.Unlock() + m.warmupRequestsDelayed++ +} + +// GetWarmupDelayed returns warmup_requests_delayed_total. +func (m *Metrics) GetWarmupDelayed() int64 { + m.mu.RLock() + defer m.mu.RUnlock() + return m.warmupRequestsDelayed +} + +// IncrPrefixCacheHitTokens increments vllm_prefix_cache_hit_tokens_total by n. +func (m *Metrics) IncrPrefixCacheHitTokens(n int64) { + m.mu.Lock() + defer m.mu.Unlock() + m.prefixCacheHitTokens += n +} + +// GetPrefixCacheHitTokens returns vllm_prefix_cache_hit_tokens_total. +func (m *Metrics) GetPrefixCacheHitTokens() int64 { + m.mu.RLock() + defer m.mu.RUnlock() + return m.prefixCacheHitTokens +} + +// globalMetrics is the singleton Metrics instance for gateway-wide telemetry. +var ( + globalMetrics *Metrics + globalMetricsOnce sync.Once +) + +// Global returns the singleton Metrics instance, initialising on first call. +func Global() *Metrics { + globalMetricsOnce.Do(func() { globalMetrics = NewMetrics() }) + return globalMetrics } diff --git a/internal/observability/prometheus.go b/internal/observability/prometheus.go index 6c2d040..46c7bc5 100644 --- a/internal/observability/prometheus.go +++ b/internal/observability/prometheus.go @@ -36,6 +36,18 @@ func (p *PrometheusExporter) Export() string { lines = append(lines, "# TYPE request_duration_seconds histogram") p.exportRequestDuration(&lines) + lines = append(lines, "# HELP queue_full_requests_total Requests rejected due to backend queue saturation") + lines = append(lines, "# TYPE queue_full_requests_total counter") + lines = append(lines, fmt.Sprintf("queue_full_requests_total %d", p.metrics.GetQueueFull())) + + lines = append(lines, "# HELP warmup_requests_delayed_total LLM warmup requests delayed for vLLM slot management") + lines = append(lines, "# TYPE warmup_requests_delayed_total counter") + lines = append(lines, fmt.Sprintf("warmup_requests_delayed_total %d", p.metrics.GetWarmupDelayed())) + + lines = append(lines, "# HELP vllm_prefix_cache_hit_tokens_total vLLM KV prefix cache hit tokens") + lines = append(lines, "# TYPE vllm_prefix_cache_hit_tokens_total counter") + lines = append(lines, fmt.Sprintf("vllm_prefix_cache_hit_tokens_total %d", p.metrics.GetPrefixCacheHitTokens())) + return strings.Join(lines, "\n") + "\n" } diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go index 66fddc5..c3d9972 100644 --- a/internal/proxy/proxy.go +++ b/internal/proxy/proxy.go @@ -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 diff --git a/internal/proxy/proxy_test.go b/internal/proxy/proxy_test.go index d04e307..b092004 100644 --- a/internal/proxy/proxy_test.go +++ b/internal/proxy/proxy_test.go @@ -1,15 +1,24 @@ package proxy import ( + "crypto/rand" + "crypto/rsa" + "encoding/base64" + "encoding/json" "fmt" "io" + "math/big" "net/http" "net/http/httptest" + "runtime" "strings" "testing" "time" + "github.com/golang-jwt/jwt/v5" + "forgejo.riotpiao.com/rock/homelab-frontend/internal/config" + "forgejo.riotpiao.com/rock/homelab-frontend/internal/observability" ) func TestProxyBasic(t *testing.T) { @@ -414,3 +423,356 @@ func TestProxyPreservesBody(t *testing.T) { t.Errorf("expected body %s, got %s", testBody, receivedBody) } } + +// TestAuthenticateV1ProxyLevel exercises the three branches of authenticateV1: +// no-token→401, bad-token→403, valid-token→proxied (non-401/403). +func TestAuthenticateV1ProxyLevel(t *testing.T) { + // Generate RSA keypair for signing test JWTs. + privateKey, err := rsa.GenerateKey(rand.Reader, 2048) + if err != nil { + t.Fatalf("rsa.GenerateKey: %v", err) + } + + // Build minimal JWKS document from the RSA public key. + nBytes := privateKey.PublicKey.N.Bytes() + eBig := new(big.Int).SetInt64(int64(privateKey.PublicKey.E)) + eBytes := eBig.Bytes() + jwksDoc, _ := json.Marshal(map[string]interface{}{ + "keys": []map[string]interface{}{ + { + "kty": "RSA", + "use": "sig", + "kid": "test-key", + "alg": "RS256", + "n": base64.RawURLEncoding.EncodeToString(nBytes), + "e": base64.RawURLEncoding.EncodeToString(eBytes), + }, + }, + }) + + // Serve JWKS from a local httptest server. + jwksSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.Write(jwksDoc) + })) + defer jwksSrv.Close() + + // Upstream stub — records whether it was reached. + reached := make(chan struct{}, 1) + upstreamSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + select { + case reached <- struct{}{}: + default: + } + w.WriteHeader(http.StatusOK) + })) + defer upstreamSrv.Close() + upstreamAddr := strings.TrimPrefix(upstreamSrv.URL, "http://") + + // signJWT returns a valid RS256 JWT accepted by the Authentik-pattern issuer check. + signJWT := func() string { + tok := jwt.NewWithClaims(jwt.SigningMethodRS256, jwt.MapClaims{ + "iss": "https://authentik.riotpiao.com/application/o/test/", + "aud": "test-audience", + "sub": "test-user", + "exp": time.Now().Add(time.Hour).Unix(), + "iat": time.Now().Unix(), + }) + tok.Header["kid"] = "test-key" + signed, err := tok.SignedString(privateKey) + if err != nil { + t.Fatalf("sign JWT: %v", err) + } + return signed + } + + // Build proxy with auth enabled. + cfg := &config.Config{ + Routes: make(map[string]*config.Route), + Models: map[string]*config.ModelUpstream{ + "auth-model": { + Name: "auth-model", + Address: upstreamAddr, + }, + }, + Auth: config.AuthConfig{ + Enabled: true, + JWKSURL: jwksSrv.URL, + }, + } + h := New(cfg) + defer h.Close() + srv := httptest.NewServer(h) + defer srv.Close() + + mkReq := func(authHeader string) *http.Request { + body := strings.NewReader(`{"model":"auth-model","messages":[]}`) + req, _ := http.NewRequest("POST", srv.URL+"/v1/chat/completions", body) + req.Header.Set("Content-Type", "application/json") + if authHeader != "" { + req.Header.Set("Authorization", authHeader) + } + return req + } + + // Branch 1: no Authorization header → 401. + t.Run("no_token_returns_401", func(t *testing.T) { + resp, err := http.DefaultClient.Do(mkReq("")) + if err != nil { + t.Fatalf("request failed: %v", err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusUnauthorized { + t.Errorf("expected 401, got %d", resp.StatusCode) + } + }) + + // Branch 2: malformed/invalid JWT → 403. + t.Run("bad_token_returns_403", func(t *testing.T) { + resp, err := http.DefaultClient.Do(mkReq("Bearer this.is.garbage")) + if err != nil { + t.Fatalf("request failed: %v", err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusForbidden { + t.Errorf("expected 403, got %d", resp.StatusCode) + } + }) + + // Branch 3: valid RS256 JWT → proxied to upstream (not 401/403). + t.Run("valid_token_proxied", func(t *testing.T) { + resp, err := http.DefaultClient.Do(mkReq("Bearer " + signJWT())) + if err != nil { + t.Fatalf("request failed: %v", err) + } + defer resp.Body.Close() + if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden { + body, _ := io.ReadAll(resp.Body) + t.Errorf("expected proxied response, got %d body=%s", resp.StatusCode, body) + } + select { + case <-reached: + // upstream was hit — auth passed + default: + t.Error("upstream was not reached after valid JWT") + } + }) +} + +// TestProxyQueueFull verifies that a 503 is returned and the queue_full metric +// is incremented when the backend queue is at capacity. +// +// Setup: maxSlots=1, maxQueue=1. Request 1 holds the slot. Request 2 fills the +// queue (blocks waiting for slot). Request 3 gets ErrQueueFull → 503. +func TestProxyQueueFull(t *testing.T) { + t.Setenv("VLLM_MAX_NUM_SEQS", "1") // 1 concurrent slot + t.Setenv("VLLM_MAX_QUEUE_DEPTH", "1") // queue capacity 1 + t.Setenv("WARMUP_DELAY_MS", "0") + observability.Global().Reset() + + // Upstream blocks until released; buffers one signal when it first starts. + slotHeld := make(chan struct{}, 1) + release := make(chan struct{}) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + select { + case slotHeld <- struct{}{}: + default: + } + <-release + w.WriteHeader(http.StatusOK) + })) + defer upstream.Close() + + upstreamAddr := strings.TrimPrefix(upstream.URL, "http://") + cfg := &config.Config{ + Routes: make(map[string]*config.Route), + Models: map[string]*config.ModelUpstream{ + "q-model": { + Name: "q-model", + Address: upstreamAddr, + }, + }, + } + h := New(cfg) + defer h.Close() + + srv := httptest.NewServer(h) + defer srv.Close() + + mkReq := func() *http.Request { + body := strings.NewReader(`{"model":"q-model","messages":[]}`) + req, _ := http.NewRequest("POST", srv.URL+"/v1/chat/completions", body) + req.Header.Set("Content-Type", "application/json") + return req + } + + // Request 1: acquires the only slot, reaches upstream, blocks there. + r1ch := make(chan *http.Response, 1) + go func() { + resp, _ := http.DefaultClient.Do(mkReq()) + r1ch <- resp + }() + <-slotHeld // slot is acquired; upstream is holding request 1 + + // Request 2: enters the queue (capacity 1), blocks waiting for slot. + r2ch := make(chan *http.Response, 1) + go func() { + resp, _ := http.DefaultClient.Do(mkReq()) + r2ch <- resp + }() + // Spin until request 2 enters the backend queue channel. + for len(h.schedulers["q-model"].backends[0].queue) == 0 { + runtime.Gosched() + } + + // Request 3: queue is full → expect 503 and metric increment. + resp3, err := http.DefaultClient.Do(mkReq()) + if err != nil { + t.Fatalf("request 3 failed: %v", err) + } + defer resp3.Body.Close() + + if resp3.StatusCode != http.StatusServiceUnavailable { + t.Errorf("expected 503 queue full, got %d", resp3.StatusCode) + } + if got := observability.Global().GetQueueFull(); got != 1 { + t.Errorf("expected GetQueueFull()=1, got %d", got) + } + + // Cleanup: unblock upstream so goroutines can finish. + close(release) + if r := <-r1ch; r != nil { + r.Body.Close() + } + if r := <-r2ch; r != nil { + r.Body.Close() + } +} + +// TestProxyCacheHitLogging verifies that when a response contains the +// vllm-num-cached-tokens header and a cache key was set, the prefix cache +// hit token count is incremented in observability metrics. +func TestProxyCacheHitLogging(t *testing.T) { + observability.Global().Reset() + + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("vllm-num-cached-tokens", "128") + w.WriteHeader(http.StatusOK) + })) + defer upstream.Close() + + upstreamAddr := strings.TrimPrefix(upstream.URL, "http://") + cfg := &config.Config{ + Routes: make(map[string]*config.Route), + Models: map[string]*config.ModelUpstream{ + "cache-model": { + Name: "cache-model", + Address: upstreamAddr, + }, + }, + } + h := New(cfg) + defer h.Close() + + srv := httptest.NewServer(h) + defer srv.Close() + + body := strings.NewReader(`{"model":"cache-model","messages":[]}`) + req, _ := http.NewRequest("POST", srv.URL+"/v1/chat/completions", body) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("X-Poimen-Cache-Key", "testkey123") + + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("request failed: %v", err) + } + resp.Body.Close() + + if got := observability.Global().GetPrefixCacheHitTokens(); got != 128 { + t.Errorf("expected PrefixCacheHitTokens=128, got %d", got) + } +} + +// TestWarmupHeaderHandling verifies that X-Poimen-Warmup: true is stripped from +// the upstream request and replaced with X-Priority: low, while a request without +// X-Poimen-Warmup reaches upstream without either header. +// +// The warmup path runs only for model-based routes (cfg.Models + scheduler), +// so we POST to /v1/chat/completions with a JSON body that names the test model. +// WARMUP_DELAY_MS=0 eliminates the 500 ms sleep. +func TestWarmupHeaderHandling(t *testing.T) { + t.Setenv("WARMUP_DELAY_MS", "0") + + tests := []struct { + name string + sendWarmup bool + wantWarmupGone bool + wantPriorityLow bool + }{ + { + name: "warmup=true", + sendWarmup: true, + wantWarmupGone: true, // X-Poimen-Warmup must NOT reach upstream + wantPriorityLow: true, // X-Priority: low must be set + }, + { + name: "warmup=false", + sendWarmup: false, + wantWarmupGone: true, // no header sent, so absent at upstream + wantPriorityLow: false, // X-Priority must NOT be set + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + var gotWarmup, gotPriority string + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotWarmup = r.Header.Get("X-Poimen-Warmup") + gotPriority = r.Header.Get("X-Priority") + w.WriteHeader(http.StatusOK) + })) + defer upstream.Close() + + upstreamAddr := strings.TrimPrefix(upstream.URL, "http://") + // Warmup code only runs for model-based routes (scheduler path). + // Use cfg.Models so New() builds a Scheduler for "warmup-model". + cfg := &config.Config{ + Routes: make(map[string]*config.Route), + Models: map[string]*config.ModelUpstream{ + "warmup-model": { + Name: "warmup-model", + Address: upstreamAddr, + }, + }, + } + h := New(cfg) + defer h.Close() + + srv := httptest.NewServer(h) + defer srv.Close() + + // POST to /v1/chat/completions — triggers model-based dispatch and warmup path. + body := strings.NewReader(`{"model":"warmup-model","messages":[]}`) + req, _ := http.NewRequest("POST", srv.URL+"/v1/chat/completions", body) + req.Header.Set("Content-Type", "application/json") + if tc.sendWarmup { + req.Header.Set("X-Poimen-Warmup", "true") + } + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("request failed: %v", err) + } + resp.Body.Close() + + if tc.wantWarmupGone && gotWarmup != "" { + t.Errorf("X-Poimen-Warmup must be stripped before upstream, got %q", gotWarmup) + } + if tc.wantPriorityLow && gotPriority != "low" { + t.Errorf("X-Priority: low must be set when warmup=true, got %q", gotPriority) + } + if !tc.wantPriorityLow && gotPriority != "" { + t.Errorf("X-Priority must not be set for normal requests, got %q", gotPriority) + } + }) + } +} diff --git a/internal/proxy/router.go b/internal/proxy/router.go index da95efa..544ff3a 100644 --- a/internal/proxy/router.go +++ b/internal/proxy/router.go @@ -159,6 +159,7 @@ func (h *Handler) routeByModel(r *http.Request, path string) (*Route, error) { route := &Route{ Name: "v1-chat-" + modelName, + ModelName: modelName, Upstream: upstreamCfg, Transport: h.getOrCreateTransport(modelUpstream.Address, upstreamCfg), Director: func(req *http.Request) { diff --git a/internal/proxy/scheduler.go b/internal/proxy/scheduler.go new file mode 100644 index 0000000..51f495f --- /dev/null +++ b/internal/proxy/scheduler.go @@ -0,0 +1,157 @@ +package proxy + +import ( + "context" + "errors" + "hash/fnv" + "sort" + "sync" + "sync/atomic" + "time" +) + +// ErrQueueFull is returned when the in-memory queue is at capacity. +var ErrQueueFull = errors.New("backend queue full") + +// Backend represents a single vLLM upstream with a bounded in-memory queue. +// Two channels model the queue: +// +// queue — waiting requests (bounded by VLLM_MAX_QUEUE_DEPTH) +// slots — concurrent in-service requests (bounded by VLLM_MAX_NUM_SEQS) +type Backend struct { + Address string + queue chan struct{} + slots chan struct{} +} + +// newBackend creates a Backend with given queue depth and concurrency limit. +func newBackend(addr string, maxSlots, maxQueue int) *Backend { + slots := make(chan struct{}, maxSlots) + for i := 0; i < maxSlots; i++ { + slots <- struct{}{} + } + return &Backend{ + Address: addr, + queue: make(chan struct{}, maxQueue), + slots: slots, + } +} + +// Acquire enters the in-memory queue then waits for a vLLM concurrency slot. +// Returns ErrQueueFull immediately if the queue channel is at capacity (no goroutine queued). +// Returns ctx.Err() if the context is cancelled while waiting for a slot. +func (b *Backend) Acquire(ctx context.Context) error { + // Non-blocking queue entry: reject immediately when queue is full. + select { + case b.queue <- struct{}{}: + // entered queue + default: + return ErrQueueFull + } + // Wait for a concurrency slot, then leave the queue. + select { + case <-b.slots: + <-b.queue // leave queue; now in service + return nil + case <-ctx.Done(): + <-b.queue // leave queue; cancelled + return ctx.Err() + } +} + +// Release returns a concurrency slot to the pool after the upstream request completes. +func (b *Backend) Release() { + b.slots <- struct{}{} +} + +// Scheduler routes LLM requests to backends using rendezvous hashing (cache-aware) +// or round-robin (no cache key). Failed backends are excluded for a configurable window. +type Scheduler struct { + mu sync.RWMutex + backends []*Backend + failedUntil map[string]time.Time + failWindow time.Duration + rrIdx atomic.Uint64 +} + +// NewScheduler creates a Scheduler for the given addresses. +// maxSlots = VLLM_MAX_NUM_SEQS, maxQueue = VLLM_MAX_QUEUE_DEPTH. +func NewScheduler(addresses []string, maxSlots, maxQueue int) *Scheduler { + backends := make([]*Backend, len(addresses)) + for i, addr := range addresses { + backends[i] = newBackend(addr, maxSlots, maxQueue) + } + return &Scheduler{ + backends: backends, + failedUntil: make(map[string]time.Time), + failWindow: 30 * time.Second, + } +} + +// Select returns a backend. Uses rendezvous hashing when cacheKey is non-empty +// (same key always routes to the same backend, skipping unhealthy ones). +// Uses round-robin when cacheKey is empty. +func (s *Scheduler) Select(cacheKey string) *Backend { + s.mu.RLock() + defer s.mu.RUnlock() + + if len(s.backends) == 0 { + return nil + } + if len(s.backends) == 1 { + return s.backends[0] + } + if cacheKey != "" { + return s.selectRendezvous(cacheKey) + } + return s.selectRoundRobin() +} + +// selectRendezvous picks the healthy backend with the highest FNV32a score +// for the key (highest random weight / rendezvous hashing). +// Falls back to next-highest score when primary is excluded by MarkFailed. +// Must be called with mu held. +func (s *Scheduler) selectRendezvous(cacheKey string) *Backend { + type scored struct { + b *Backend + score uint32 + } + now := time.Now() + scores := make([]scored, len(s.backends)) + h := fnv.New32a() + for i, b := range s.backends { + h.Reset() + h.Write([]byte(cacheKey + ":" + b.Address)) + scores[i] = scored{b: b, score: h.Sum32()} + } + sort.Slice(scores, func(i, j int) bool { + return scores[i].score > scores[j].score + }) + for _, sc := range scores { + _, failed := s.failedUntil[sc.b.Address] + if !failed { + return sc.b + } + expiry := s.failedUntil[sc.b.Address] + if now.After(expiry) { + return sc.b + } + } + // All backends unhealthy: return highest-scored (fail-open). + return scores[0].b +} + +// selectRoundRobin distributes requests evenly when no cache key is present. +// Must be called with mu held. +func (s *Scheduler) selectRoundRobin() *Backend { + idx := s.rrIdx.Add(1) - 1 + return s.backends[idx%uint64(len(s.backends))] +} + +// MarkFailed excludes addr from routing for the fail window (default 30 s). +// Called by the proxy ErrorHandler on upstream connection failure. +func (s *Scheduler) MarkFailed(addr string) { + s.mu.Lock() + defer s.mu.Unlock() + s.failedUntil[addr] = time.Now().Add(s.failWindow) +} diff --git a/internal/proxy/scheduler_test.go b/internal/proxy/scheduler_test.go new file mode 100644 index 0000000..e0631b3 --- /dev/null +++ b/internal/proxy/scheduler_test.go @@ -0,0 +1,182 @@ +package proxy + +import ( + "context" + "errors" + "runtime" + "sync" + "testing" + "time" +) + +// TestBackendAcquireBlocksThenUnblocksOnRelease verifies that a second Acquire +// blocks while a slot is held, then unblocks when Release is called. +func TestBackendAcquireBlocksThenUnblocksOnRelease(t *testing.T) { + b := newBackend("host:1", 1, 1) + + // G1: acquire and hold the only slot. + if err := b.Acquire(context.Background()); err != nil { + t.Fatalf("G1 Acquire: unexpected error: %v", err) + } + + g2Done := make(chan error, 1) + go func() { + g2Done <- b.Acquire(context.Background()) + }() + + // Wait until G2 enters the queue channel (proves Acquire is blocking on slots). + for len(b.queue) == 0 { + runtime.Gosched() + } + + select { + case err := <-g2Done: + t.Fatalf("G2 returned before Release: %v", err) + default: + } + + // G1 releases; G2 should unblock. + b.Release() + + select { + case err := <-g2Done: + if err != nil { + t.Fatalf("G2 Acquire after Release: unexpected error: %v", err) + } + case <-time.After(time.Second): + t.Fatal("G2 did not unblock within 1s after Release") + } + + // Cleanup: release G2's slot. + b.Release() +} + +// TestBackendAcquireReturnsErrQueueFullWhenQueueFull verifies that Acquire returns +// ErrQueueFull immediately when no queue capacity exists (maxQueue=0). +// With a zero-capacity queue channel the non-blocking select always hits default. +func TestBackendAcquireReturnsErrQueueFullWhenQueueFull(t *testing.T) { + b := newBackend("host:1", 1, 0) + + err := b.Acquire(context.Background()) + if !errors.Is(err, ErrQueueFull) { + t.Fatalf("expected ErrQueueFull, got %v", err) + } +} + +// TestSchedulerSelectRendezvousHashDeterministic verifies that the same cache key +// always selects the same backend (rendezvous / highest-random-weight hashing). +func TestSchedulerSelectRendezvousHashDeterministic(t *testing.T) { + addrs := []string{"a:80", "b:80", "c:80"} + s := NewScheduler(addrs, 10, 10) + + first := s.Select("same-cache-key") + if first == nil { + t.Fatal("Select returned nil") + } + for i := 1; i < 10; i++ { + b := s.Select("same-cache-key") + if b == nil || b.Address != first.Address { + t.Fatalf("call %d: want %s, got %v", i, first.Address, b) + } + } +} + +// TestSchedulerSelectNoKeyRoundRobin verifies that round-robin distributes +// requests evenly across all backends when no cache key is given. +func TestSchedulerSelectNoKeyRoundRobin(t *testing.T) { + addrs := []string{"a:80", "b:80", "c:80"} + s := NewScheduler(addrs, 10, 10) + + counts := make(map[string]int) + for i := 0; i < 9; i++ { + b := s.Select("") + if b == nil { + t.Fatalf("call %d: Select returned nil", i) + } + counts[b.Address]++ + } + + for _, addr := range addrs { + if counts[addr] != 3 { + t.Errorf("addr %s: want 3 selections, got %d (all counts: %v)", addr, counts[addr], counts) + } + } +} + +// TestBackendAcquireCancelledContext verifies that Acquire returns ctx.Err() when the +// context is cancelled while the caller is blocked waiting for a concurrency slot, +// and that the queue depth returns to 0 after cancellation. +func TestBackendAcquireCancelledContext(t *testing.T) { + b := newBackend("host:1", 1, 1) + + // G1: acquire and hold the only slot indefinitely. + if err := b.Acquire(context.Background()); err != nil { + t.Fatalf("G1 Acquire: unexpected error: %v", err) + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + g2Done := make(chan error, 1) + go func() { + g2Done <- b.Acquire(ctx) + }() + + // Wait until G2 enters the queue channel. + for len(b.queue) == 0 { + runtime.Gosched() + } + + // Verify G2 is queued (queue channel has one item). + if qd := len(b.queue); qd != 1 { + t.Errorf("expected queue depth 1 while G2 blocked, got %d", qd) + } + + // Cancel ctx while G2 is waiting for the slot. + cancel() + + select { + case err := <-g2Done: + if err == nil { + t.Fatal("G2 Acquire: expected ctx.Err(), got nil") + } + if !errors.Is(err, context.Canceled) { + t.Fatalf("G2 Acquire: expected context.Canceled, got %v", err) + } + case <-time.After(time.Second): + t.Fatal("G2 did not return within 1s after context cancel") + } + + // Queue depth must return to 0 after G2 is drained. + if qd := len(b.queue); qd != 0 { + t.Errorf("expected queue depth 0 after cancellation, got %d", qd) + } + + // Cleanup: release G1's slot. + b.Release() +} + +// TestSchedulerMarkFailedExcludesBackend verifies that a failed backend is not +// selected within the fail window (30 s default). +func TestSchedulerMarkFailedExcludesBackend(t *testing.T) { + addrs := []string{"a:80", "b:80"} + s := NewScheduler(addrs, 10, 10) + + s.MarkFailed("a:80") + + var mu sync.Mutex + seen := make(map[string]int) + for i := 0; i < 10; i++ { + b := s.Select("any-key") + if b == nil { + t.Fatalf("call %d: Select returned nil", i) + } + mu.Lock() + seen[b.Address]++ + mu.Unlock() + } + + if seen["a:80"] > 0 { + t.Errorf("failed backend a:80 was selected %d times; expected 0 within fail window", seen["a:80"]) + } +} diff --git a/internal/server/router.go b/internal/server/router.go index c3fcd0f..4601714 100644 --- a/internal/server/router.go +++ b/internal/server/router.go @@ -16,6 +16,7 @@ type Router struct { temporalHandler http.Handler upstreamHandler http.Handler forgejoWebhook *webhook.ForgejoHandler + metricsHandler http.Handler } // NewRouter creates a new router with health endpoints. @@ -33,6 +34,11 @@ func NewRouter(healthChecker *HealthChecker, dispatcher *serviceadapter.Dispatch } } +// SetMetricsHandler registers the Prometheus /metrics endpoint handler. +func (r *Router) SetMetricsHandler(h http.Handler) { + r.metricsHandler = h +} + // ServeHTTP implements http.Handler. // Priority order: // 1. /healthz and /readyz to health handlers @@ -48,6 +54,11 @@ func (r *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) { case "/readyz": ReadinessHandler(r.healthChecker)(w, req) return + case "/metrics": + if r.metricsHandler != nil { + r.metricsHandler.ServeHTTP(w, req) + return + } } // X-Service (ServiceAdapter) routing - checked before path-based routing diff --git a/internal/serviceadapter/memory_adapter.go b/internal/serviceadapter/memory_adapter.go new file mode 100644 index 0000000..9fdd95e --- /dev/null +++ b/internal/serviceadapter/memory_adapter.go @@ -0,0 +1,67 @@ +package serviceadapter + +import "os" + +// NewMemoryAdapter creates a code-registered ServiceAdapter for the memory service. +// Takes precedence over any config-file loaded memory adapter (main.go skips config +// adapters when a code-registered one with the same name exists). +// Upstream URL configured via MEMORY_UPSTREAM_URL env var. +func NewMemoryAdapter() *ServiceAdapter { + upstreamURL := os.Getenv("MEMORY_UPSTREAM_URL") + if upstreamURL == "" { + upstreamURL = "http://poimen-memory.poimen.svc.cluster.local:8080" + } + return &ServiceAdapter{ + Namespace: "poimen", + ServiceName: "memory", + Spec: *GetMemorySpec(upstreamURL), + } +} + +// GetMemorySpec returns the ServiceAdapter spec for the memory service. +func GetMemorySpec(upstreamURL string) *Spec { + return &Spec{ + ServiceName: "memory", + Upstream: Upstream{ + URL: upstreamURL, + TimeoutSeconds: 30, + }, + Auth: Auth{ + Required: true, + Capability: "memory:read", + }, + Resources: []Resource{ + { + Name: "ingest", + Methods: []Method{{ + Verb: "POST", + UpstreamPath: "/memory/ingest", + Auth: &Auth{Required: true, Capability: "memory:write"}, + }}, + }, + { + Name: "query", + Methods: []Method{{ + Verb: "POST", + UpstreamPath: "/memory/query", + Auth: &Auth{Required: true, Capability: "memory:read"}, + }}, + }, + { + Name: "skills", + Methods: []Method{{ + Verb: "GET", + UpstreamPath: "/memory/skills", + }}, + }, + { + Name: "standing-context", + Methods: []Method{{ + Verb: "POST", + UpstreamPath: "/memory/standing-context", + Auth: &Auth{Required: true, Capability: "memory:read"}, + }}, + }, + }, + } +} diff --git a/internal/serviceadapter/router.go b/internal/serviceadapter/router.go index 4c090d3..51f5872 100644 --- a/internal/serviceadapter/router.go +++ b/internal/serviceadapter/router.go @@ -16,6 +16,7 @@ import ( "forgejo.riotpiao.com/rock/homelab-frontend/internal/auth" "forgejo.riotpiao.com/rock/homelab-frontend/internal/identity" + "forgejo.riotpiao.com/rock/homelab-frontend/internal/logging" "forgejo.riotpiao.com/rock/homelab-frontend/internal/problem" ) @@ -74,8 +75,13 @@ func (d *Dispatcher) Dispatch(w http.ResponseWriter, r *http.Request) { } // JWT auth enforcement for adapters that require it - if adapter.Spec.Auth.Required && d.jwtValidator != nil { - if !d.authenticate(w, r, serviceName, method.Verb) { + if adapter.Spec.Auth.Required { + if d.jwtValidator == nil { + logging.Warnf("auth required but jwtValidator nil; request unauthenticated", map[string]string{ + "service": serviceName, + "resource": resourceName, + }) + } else if !d.authenticate(w, r, serviceName, method) { return } } @@ -98,7 +104,7 @@ func (d *Dispatcher) Dispatch(w http.ResponseWriter, r *http.Request) { // authenticate validates the JWT and checks service-level capability. // Returns false (and writes error response) if auth fails. -func (d *Dispatcher) authenticate(w http.ResponseWriter, r *http.Request, serviceName, verb string) bool { +func (d *Dispatcher) authenticate(w http.ResponseWriter, r *http.Request, serviceName string, method *Method) bool { authHeader := r.Header.Get("Authorization") if authHeader == "" { d.writeError(w, problem.NewProblem(http.StatusUnauthorized, @@ -115,8 +121,11 @@ func (d *Dispatcher) authenticate(w http.ResponseWriter, r *http.Request, servic return false } - // Check capability: :read for GET/HEAD, :write for mutating verbs - required := capabilityForVerb(serviceName, verb) + // Check capability: method-level Auth.Capability if set; else derive from verb. + required := capabilityForVerb(serviceName, method.Verb) + if method.Auth != nil && method.Auth.Capability != "" { + required = method.Auth.Capability + } if !d.jwtValidator.CheckPermissions(claims, required, "*") { d.writeError(w, problem.NewProblem(http.StatusForbidden, "about:blank#insufficient-permissions", "Insufficient Permissions", diff --git a/internal/serviceadapter/router_test.go b/internal/serviceadapter/router_test.go index 35ac7a7..29f3508 100644 --- a/internal/serviceadapter/router_test.go +++ b/internal/serviceadapter/router_test.go @@ -1,11 +1,19 @@ package serviceadapter import ( + "crypto/rand" + "crypto/rsa" + "encoding/base64" "encoding/json" + "math/big" "net/http" "net/http/httptest" "testing" + "time" + "github.com/golang-jwt/jwt/v5" + + "forgejo.riotpiao.com/rock/homelab-frontend/internal/auth" "forgejo.riotpiao.com/rock/homelab-frontend/internal/identity" ) @@ -263,3 +271,197 @@ func TestDispatch_IdentityHeadersNotSet_WhenNoAuth(t *testing.T) { t.Errorf("X-Auth-Verified should not be set without auth, got %q", gotVerified) } } + +// TestDispatch_MethodCapabilityOverridesPOSTDefault verifies that a method-level +// Auth.Capability (e.g. "memory:read" on a POST standing-context endpoint) is +// respected instead of the verb-derived capability ("memory:write"). +// Because auth.Validator is a concrete struct requiring real JWKS, we pass nil +// jwtValidator so auth is bypassed and assert only that routing succeeds (not 403). +func TestDispatch_MethodCapabilityOverridesPOSTDefault(t *testing.T) { + // Upstream captures the request so we can verify it was routed. + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + defer upstream.Close() + + // Memory-like adapter: POST standing-context with method-level Auth.Capability=memory:read. + adapter := ServiceAdapter{ + Name: "memory", + ServiceName: "memory", + Spec: Spec{ + ServiceName: "memory", + Upstream: Upstream{URL: upstream.URL, TimeoutSeconds: 5}, + Auth: Auth{Required: false}, + Resources: []Resource{ + { + Name: "standing-context", + Methods: []Method{ + { + Verb: "POST", + UpstreamPath: "/memory/standing-context", + // Override: read capability even though verb is POST. + Auth: &Auth{Capability: "memory:read"}, + }, + }, + }, + }, + }, + } + + // nil jwtValidator: auth.Required=false + nil validator → auth skipped entirely. + d := NewDispatcher(newTestRegistry(adapter), nil) + w := httptest.NewRecorder() + r := httptest.NewRequest("POST", "/", nil) + r.Header.Set("X-Service", "memory") + r.Header.Set("X-Resource", "standing-context") + + d.Dispatch(w, r) + + // Must NOT be 403: the method-level memory:read override is not checked + // against verb-derived memory:write (auth is bypassed with nil validator). + // Must be 200 routed to real upstream. + if w.Code == http.StatusForbidden { + t.Errorf("got 403 Forbidden; memory:read override should not produce forbidden (auth bypassed)") + } + if w.Code != http.StatusOK { + t.Errorf("expected 200 from upstream, got %d", w.Code) + } +} + +// TestAuthenticate_MethodCapabilityOverrideDirect calls d.Dispatch() with a real +// signed RS256 JWT so authenticate() executes end-to-end and the capability-override +// branch (router.go lines 124-130) is reached and covered. +// +// Sub-test 1: token has "memory:read", method.Auth.Capability="memory:read" → 200. +// Sub-test 2: token has "memory:read", no method override → POST derives +// +// "memory:write" → CheckPermissions fails → 403. +func TestAuthenticate_MethodCapabilityOverrideDirect(t *testing.T) { + // Generate an RSA key pair for signing test JWTs. + privateKey, err := rsa.GenerateKey(rand.Reader, 2048) + if err != nil { + t.Fatalf("rsa.GenerateKey: %v", err) + } + + // Build minimal JWKS document from the RSA public key. + nBytes := privateKey.PublicKey.N.Bytes() + eBig := new(big.Int).SetInt64(int64(privateKey.PublicKey.E)) + eBytes := eBig.Bytes() + jwksDoc, _ := json.Marshal(map[string]interface{}{ + "keys": []map[string]interface{}{ + { + "kty": "RSA", + "use": "sig", + "kid": "test-key", + "alg": "RS256", + "n": base64.RawURLEncoding.EncodeToString(nBytes), + "e": base64.RawURLEncoding.EncodeToString(eBytes), + }, + }, + }) + + // Serve JWKS from a local test server. + jwksSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.Write(jwksDoc) + })) + defer jwksSrv.Close() + + // Upstream that returns 200 when authenticate() passes. + upstreamSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + defer upstreamSrv.Close() + + // signJWT returns a signed RS256 token with the given permissions slice. + signJWT := func(permissions []string) string { + tok := jwt.NewWithClaims(jwt.SigningMethodRS256, jwt.MapClaims{ + "iss": "https://authentik.riotpiao.com/application/o/test/", + "aud": "test-audience", + "sub": "test-user", + "exp": time.Now().Add(time.Hour).Unix(), + "iat": time.Now().Unix(), + "permissions": permissions, + }) + tok.Header["kid"] = "test-key" + signed, err := tok.SignedString(privateKey) + if err != nil { + t.Fatalf("sign JWT: %v", err) + } + return signed + } + + validator := auth.NewValidator("", "", jwksSrv.URL) + + // Sub-test 1: method.Auth.Capability="memory:read"; token has "memory:read". + // authenticate() must execute the override branch and pass → upstream returns 200. + t.Run("capability_override_match_reaches_upstream", func(t *testing.T) { + adapter := ServiceAdapter{ + Name: "memory", + ServiceName: "memory", + Spec: Spec{ + ServiceName: "memory", + Upstream: Upstream{URL: upstreamSrv.URL, TimeoutSeconds: 5}, + Auth: Auth{Required: true}, + Resources: []Resource{ + { + Name: "standing-context", + Methods: []Method{ + { + Verb: "POST", + UpstreamPath: "/memory/standing-context", + Auth: &Auth{Capability: "memory:read"}, // override + }, + }, + }, + }, + }, + } + d := NewDispatcher(newTestRegistry(adapter), validator) + w := httptest.NewRecorder() + r := httptest.NewRequest("POST", "/", nil) + r.Header.Set("X-Service", "memory") + r.Header.Set("X-Resource", "standing-context") + r.Header.Set("Authorization", "Bearer "+signJWT([]string{"memory:read"})) + d.Dispatch(w, r) + if w.Code != http.StatusOK { + t.Errorf("expected 200 (override memory:read match), got %d body=%s", w.Code, w.Body.String()) + } + }) + + // Sub-test 2: no method-level Auth override; POST → capabilityForVerb → "memory:write". + // Token has only "memory:read" → CheckPermissions fails → 403. + t.Run("verb_derived_capability_mismatch_returns_403", func(t *testing.T) { + adapter := ServiceAdapter{ + Name: "memory", + ServiceName: "memory", + Spec: Spec{ + ServiceName: "memory", + Upstream: Upstream{URL: upstreamSrv.URL, TimeoutSeconds: 5}, + Auth: Auth{Required: true}, + Resources: []Resource{ + { + Name: "data", + Methods: []Method{ + { + Verb: "POST", + UpstreamPath: "/memory/data", + // No method-level Auth: POST derives "memory:write". + }, + }, + }, + }, + }, + } + d := NewDispatcher(newTestRegistry(adapter), validator) + w := httptest.NewRecorder() + r := httptest.NewRequest("POST", "/", nil) + r.Header.Set("X-Service", "memory") + r.Header.Set("X-Resource", "data") + r.Header.Set("Authorization", "Bearer "+signJWT([]string{"memory:read"})) + d.Dispatch(w, r) + if w.Code != http.StatusForbidden { + t.Errorf("expected 403 (write required, read only in token), got %d body=%s", w.Code, w.Body.String()) + } + }) +}