fix(cag): extract applyWarmup/acquireScheduledBackend helpers, add race-free tests

- proxy.go: extract applyWarmup + acquireScheduledBackend to reduce ServeHTTP complexity
- proxy_test.go: TestProxyQueueFull, TestProxyCacheHitLogging (AC5/AC6/AC7 HTTP-level coverage)
- scheduler_test.go: TestBackendAcquireCancelledContext + runtime.Gosched spin-wait (no time.Sleep)
- router_test.go: TestAuthenticate_MethodCapabilityOverrideDirect with real RS256 JWT + JWKS server

internal/proxy/proxy.go

@@ -2,6 +2,7 @@
 package proxy
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -468,56 +469,13 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	var selectedBackend *Backend
 	cacheKey := r.Header.Get("X-Poimen-Cache-Key")
 	if sched := h.schedulers[route.ModelName]; sched != nil {
-		selectedBackend = sched.Select(cacheKey)
-
-		// Warmup: sleep BEFORE queue entry so agents arriving in that window win slots first.
-		if r.Header.Get("X-Poimen-Warmup") == "true" {
-			r.Header.Del("X-Poimen-Warmup")
-			r.Header.Set("X-Priority", "low")
-			time.Sleep(h.warmupDelay)
-			observability.Global().IncrWarmupDelayed()
+		var ok bool
+		selectedBackend, route, ok = h.acquireScheduledBackend(r.Context(), w, r, sched, route, cacheKey)
+		if !ok {
+			return
 		}
-
 		if selectedBackend != nil {
-			if err := selectedBackend.Acquire(r.Context()); err != nil {
-				if errors.Is(err, ErrQueueFull) {
-					observability.Global().IncrQueueFull()
-					logging.Warnf("backend queue full", map[string]string{
-						"backend": selectedBackend.Address,
-						"path":    r.URL.Path,
-						"model":   route.ModelName,
-					})
-					w.WriteHeader(http.StatusServiceUnavailable)
-					fmt.Fprintf(w, "queue full")
-				} else {
-					logging.Warnf("backend acquire failed", map[string]string{
-						"backend": selectedBackend.Address,
-						"path":    r.URL.Path,
-						"model":   route.ModelName,
-						"err":     err.Error(),
-					})
-					w.WriteHeader(http.StatusServiceUnavailable)
-					fmt.Fprintf(w, "upstream busy")
-				}
-				return
-			}
 			defer selectedBackend.Release()
-
-			// Patch route to selected backend if it differs from the config default.
-			if selectedBackend.Address != route.Upstream.Address {
-				patchedCfg := *route.Upstream
-				patchedCfg.Address = selectedBackend.Address
-				patchedURL, _ := url.Parse("http://" + selectedBackend.Address)
-				route = &Route{
-					Name:      route.Name,
-					ModelName: route.ModelName,
-					Upstream:  &patchedCfg,
-					Transport: h.getOrCreateTransport(selectedBackend.Address, &patchedCfg),
-					Director: func(req *http.Request) {
-						directorFunc(req, patchedURL, &patchedCfg)
-					},
-				}
-			}
 		}
 	}
 
@@ -586,6 +544,68 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 
 
+// applyWarmup strips the X-Poimen-Warmup header, delays the request, and records
+// the metric. It is a no-op when the header is absent or not "true".
+func (h *Handler) applyWarmup(r *http.Request) {
+	if r.Header.Get("X-Poimen-Warmup") != "true" {
+		return
+	}
+	r.Header.Del("X-Poimen-Warmup")
+	r.Header.Set("X-Priority", "low")
+	time.Sleep(h.warmupDelay)
+	observability.Global().IncrWarmupDelayed()
+}
+
+// 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)
+	h.applyWarmup(r)
+	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.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
+}
+
 // getValidModels returns a sorted list of all configured model names.
 func (h *Handler) getValidModels() []string {
 	var models []string

internal/proxy/proxy_test.go

@@ -5,11 +5,13 @@ import (
 	"io"
 	"net/http"
 	"net/http/httptest"
+	"runtime"
 	"strings"
 	"testing"
 	"time"
 
 	"forgejo.riotpiao.com/rock/homelab-frontend/internal/config"
+	"forgejo.riotpiao.com/rock/homelab-frontend/internal/observability"
 )
 
 func TestProxyBasic(t *testing.T) {
@@ -415,6 +417,140 @@ func TestProxyPreservesBody(t *testing.T) {
 	}
 }
 
+// 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.

internal/proxy/scheduler_test.go

@@ -3,6 +3,7 @@ package proxy
 import (
 	"context"
 	"errors"
+	"runtime"
 	"sync"
 	"testing"
 	"time"
@@ -23,8 +24,10 @@ func TestBackendAcquireBlocksThenUnblocksOnRelease(t *testing.T) {
 		g2Done <- b.Acquire(context.Background())
 	}()
 
-	// Give G2 time to block.
-	time.Sleep(20 * time.Millisecond)
+	// Wait until G2 enters the queue channel (proves Acquire is blocking on slots).
+	for len(b.queue) == 0 {
+		runtime.Gosched()
+	}
 
 	select {
 	case err := <-g2Done:
@@ -119,8 +122,10 @@ func TestBackendAcquireCancelledContext(t *testing.T) {
 		g2Done <- b.Acquire(ctx)
 	}()
 
-	// Give G2 time to enter the queue and block waiting for a slot.
-	time.Sleep(20 * time.Millisecond)
+	// 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 {

internal/serviceadapter/router_test.go

@@ -1,10 +1,17 @@
 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"
@@ -321,58 +328,140 @@ func TestDispatch_MethodCapabilityOverridesPOSTDefault(t *testing.T) {
 	}
 }
 
-// TestAuthenticate_MethodCapabilityOverrideDirect calls authenticate() directly
-// (same-package access) with a non-nil jwtValidator so the auth gate is not bypassed.
-// Two sub-tests verify:
-//  1. authenticate() is actually invoked when jwtValidator != nil (returns 401 for
-//     a missing Authorization header — proves the function ran, not the nil-bypass path).
-//  2. The method-level Auth.Capability override selects "memory:read" over the
-//     POST verb-derived default "memory:write" — exercises the same lines that
-//     authenticate() uses internally (lines: required = method.Auth.Capability).
+// 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) {
-	// Validator with a dummy JWKS URL (never fetched — ValidateBearerToken is not
-	// reached in these sub-tests because we hit the auth-header guard first).
-	validator := auth.NewValidator("", "", "http://127.0.0.1:0/jwks")
-
-	// Method with method-level capability override: POST normally derives "memory:write"
-	// but this endpoint is a read-only standing-context retrieval → override to "memory:read".
-	method := &Method{
-		Verb:         "POST",
-		UpstreamPath: "/memory/standing-context",
-		Auth:         &Auth{Capability: "memory:read"},
+	// 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)
 	}
 
-	d := &Dispatcher{jwtValidator: validator}
+	// 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),
+			},
+		},
+	})
 
-	// Sub-test 1: authenticate() is called and returns 401 for a missing
-	// Authorization header — proves the function runs (not the nil-validator bypass).
-	t.Run("authenticate_called_returns_401_without_token", func(t *testing.T) {
+	// 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)
-		got := d.authenticate(w, r, "memory", method)
-		if got {
-			t.Error("expected authenticate() to return false for missing Authorization header")
-		}
-		if w.Code != http.StatusUnauthorized {
-			t.Errorf("expected 401 Unauthorized, got %d", w.Code)
+		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: prove that the capability override selects "memory:read" (not the
-	// verb-derived "memory:write") using the identical logic from authenticate().
-	t.Run("method_level_capability_overrides_verb_derived_default", func(t *testing.T) {
-		// Replicate the exact lines from authenticate() (router.go:119-121):
-		required := capabilityForVerb("memory", method.Verb) // POST → "memory:write"
-		if method.Auth != nil && method.Auth.Capability != "" {
-			required = method.Auth.Capability // override → "memory:read"
+	// 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".
+							},
+						},
+					},
+				},
+			},
 		}
-		if required != "memory:read" {
-			t.Errorf("method-level Auth.Capability override: got %q, want \"memory:read\"", required)
-		}
-		// Confirm the verb-derived value without override would have been wrong.
-		verbDefault := capabilityForVerb("memory", method.Verb)
-		if verbDefault != "memory:write" {
-			t.Errorf("verb-derived default for POST: got %q, want \"memory:write\"", verbDefault)
+		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())
 		}
 	})
 }