feat(cag): CAG integration — memory adapter, warmup semaphore, cache-aware routing (#47)

## 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
2026-09-27 11:16:40 +00:00
parent aaa9a5d1cb
commit ef5e736673
16 changed files with 1597 additions and 162 deletions
@@ -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)
+			}
+		})
+	}
+}