merge main

proxy/internal/accesslog/logger.go

@@ -0,0 +1,442 @@
+package accesslog
+
+import (
+	"context"
+	"maps"
+	"net/netip"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"github.com/rs/xid"
+	log "github.com/sirupsen/logrus"
+	"google.golang.org/grpc"
+	"google.golang.org/protobuf/types/known/timestamppb"
+
+	"github.com/netbirdio/netbird/proxy/auth"
+	"github.com/netbirdio/netbird/proxy/internal/types"
+	"github.com/netbirdio/netbird/shared/management/proto"
+	"github.com/netbirdio/netbird/trustedproxy"
+)
+
+const (
+	requestThreshold    = 10000              // Log every 10k requests
+	bytesThreshold      = 1024 * 1024 * 1024 // Log every 1GB
+	usageCleanupPeriod  = 1 * time.Hour      // Clean up stale counters every hour
+	usageInactiveWindow = 24 * time.Hour     // Consider domain inactive if no traffic for 24 hours
+	logSendTimeout      = 10 * time.Second
+
+	// denyCooldown is the min interval between deny log entries per service+reason
+	// to prevent flooding from denied connections (e.g. UDP packets from blocked IPs).
+	denyCooldown = 10 * time.Second
+
+	// maxDenyBuckets caps tracked deny rate-limit entries to bound memory under DDoS.
+	maxDenyBuckets = 10000
+
+	// maxLogWorkers caps concurrent gRPC send goroutines.
+	maxLogWorkers = 4096
+)
+
+type domainUsage struct {
+	requestCount     int64
+	requestStartTime time.Time
+
+	bytesTransferred int64
+	bytesStartTime   time.Time
+
+	lastActivity time.Time // Track last activity for cleanup
+}
+
+type gRPCClient interface {
+	SendAccessLog(ctx context.Context, in *proto.SendAccessLogRequest, opts ...grpc.CallOption) (*proto.SendAccessLogResponse, error)
+}
+
+// denyBucketKey identifies a rate-limited deny log stream.
+type denyBucketKey struct {
+	ServiceID types.ServiceID
+	Reason    string
+}
+
+// denyBucket tracks rate-limited deny log entries.
+type denyBucket struct {
+	lastLogged time.Time
+	suppressed int64
+}
+
+// Logger sends access log entries to the management server via gRPC.
+type Logger struct {
+	client         gRPCClient
+	logger         *log.Logger
+	trustedProxies *trustedproxy.List
+
+	usageMux    sync.Mutex
+	domainUsage map[string]*domainUsage
+
+	denyMu      sync.Mutex
+	denyBuckets map[denyBucketKey]*denyBucket
+
+	logSem        chan struct{}
+	cleanupCancel context.CancelFunc
+	dropped       atomic.Int64
+}
+
+// NewLogger creates a new access log Logger. The trustedProxies parameter
+// configures which upstream proxy IP ranges are trusted for extracting
+// the real client IP from X-Forwarded-For headers.
+func NewLogger(client gRPCClient, logger *log.Logger, trustedProxies *trustedproxy.List) *Logger {
+	if logger == nil {
+		logger = log.StandardLogger()
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	l := &Logger{
+		client:         client,
+		logger:         logger,
+		trustedProxies: trustedProxies,
+		domainUsage:    make(map[string]*domainUsage),
+		denyBuckets:    make(map[denyBucketKey]*denyBucket),
+		logSem:         make(chan struct{}, maxLogWorkers),
+		cleanupCancel:  cancel,
+	}
+
+	// Start background cleanup routine
+	go l.cleanupStaleUsage(ctx)
+
+	return l
+}
+
+// Close stops the cleanup routine. Should be called during graceful shutdown.
+func (l *Logger) Close() {
+	if l.cleanupCancel != nil {
+		l.cleanupCancel()
+	}
+}
+
+type logEntry struct {
+	ID            string
+	AccountID     types.AccountID
+	ServiceID     types.ServiceID
+	Host          string
+	Path          string
+	DurationMs    int64
+	Method        string
+	ResponseCode  int32
+	SourceIP      netip.Addr
+	AuthMechanism string
+	UserID        string
+	AuthSuccess   bool
+	BytesUpload   int64
+	BytesDownload int64
+	Protocol      Protocol
+	Metadata      map[string]string
+	AgentNetwork  bool
+}
+
+// Protocol identifies the transport protocol of an access log entry.
+type Protocol string
+
+const (
+	ProtocolHTTP Protocol = "http"
+	ProtocolTCP  Protocol = "tcp"
+	ProtocolUDP  Protocol = "udp"
+	ProtocolTLS  Protocol = "tls"
+)
+
+// L4Entry holds the data for a layer-4 (TCP/UDP) access log entry.
+type L4Entry struct {
+	AccountID     types.AccountID
+	ServiceID     types.ServiceID
+	Protocol      Protocol
+	Host          string // SNI hostname or listen address
+	SourceIP      netip.Addr
+	DurationMs    int64
+	BytesUpload   int64
+	BytesDownload int64
+	// DenyReason, when non-empty, indicates the connection was denied.
+	// Values match the HTTP auth mechanism strings: "ip_restricted",
+	// "country_restricted", "geo_unavailable", "crowdsec_ban", etc.
+	DenyReason string
+	// Metadata carries extra context about the connection (e.g. CrowdSec verdict).
+	Metadata map[string]string
+}
+
+// LogL4 sends an access log entry for a layer-4 connection (TCP or UDP).
+// The call is non-blocking: the gRPC send happens in a background goroutine.
+func (l *Logger) LogL4(entry L4Entry) {
+	le := logEntry{
+		ID:            xid.New().String(),
+		AccountID:     entry.AccountID,
+		ServiceID:     entry.ServiceID,
+		Protocol:      entry.Protocol,
+		Host:          entry.Host,
+		SourceIP:      entry.SourceIP,
+		DurationMs:    entry.DurationMs,
+		BytesUpload:   entry.BytesUpload,
+		BytesDownload: entry.BytesDownload,
+		Metadata:      maps.Clone(entry.Metadata),
+	}
+	if entry.DenyReason != "" {
+		if !l.allowDenyLog(entry.ServiceID, entry.DenyReason) {
+			return
+		}
+		le.AuthMechanism = entry.DenyReason
+		le.AuthSuccess = false
+	}
+	l.log(le)
+	l.trackUsage(entry.Host, entry.BytesUpload+entry.BytesDownload)
+}
+
+// allowDenyLog rate-limits deny log entries per service+reason combination.
+func (l *Logger) allowDenyLog(serviceID types.ServiceID, reason string) bool {
+	key := denyBucketKey{ServiceID: serviceID, Reason: reason}
+	now := time.Now()
+
+	l.denyMu.Lock()
+	defer l.denyMu.Unlock()
+
+	b, ok := l.denyBuckets[key]
+	if !ok {
+		if len(l.denyBuckets) >= maxDenyBuckets {
+			return false
+		}
+		l.denyBuckets[key] = &denyBucket{lastLogged: now}
+		return true
+	}
+
+	if now.Sub(b.lastLogged) >= denyCooldown {
+		if b.suppressed > 0 {
+			l.logger.Debugf("access restriction: suppressed %d deny log entries for %s (%s)", b.suppressed, serviceID, reason)
+		}
+		b.lastLogged = now
+		b.suppressed = 0
+		return true
+	}
+
+	b.suppressed++
+	return false
+}
+
+// usageMetadataKeys is the allowlist of metadata retained on a stripped,
+// usage-only agent-network entry. Mirrors the llm.* / cost.* keys in
+// proxy/internal/middleware/keys.go — only the dimensions management needs to
+// record a usage row (provider / model / tokens / cost / groups).
+var usageMetadataKeys = map[string]struct{}{
+	"llm.provider":              {},
+	"llm.model":                 {},
+	"llm.resolved_provider_id":  {},
+	"llm.input_tokens":          {},
+	"llm.output_tokens":         {},
+	"llm.total_tokens":          {},
+	"llm.cached_input_tokens":   {},
+	"llm.cache_creation_tokens": {},
+	"cost.usd_input":            {},
+	"cost.usd_cached_input":     {},
+	"cost.usd_cache_creation":   {},
+	"cost.usd_output":           {},
+	"cost.usd_total":            {},
+	"cost.usd_cache":            {},
+	"llm.authorising_groups":    {},
+}
+
+// stripAgentNetworkEntryForUsage returns the entry reduced to what's needed to
+// record usage/cost: it drops request detail (host / path / source IP) and any
+// prompt capture, keeping the LLM usage metadata plus the caller identity
+// (user / auth mechanism) needed for attribution. Shipped when an
+// agent-network account has log collection disabled but usage must still be
+// collected. logEntry is passed by value, so mutating it here is safe; Metadata
+// is replaced with a fresh map rather than mutated in place.
+func stripAgentNetworkEntryForUsage(entry logEntry) logEntry {
+	entry.Host = ""
+	entry.Path = ""
+	entry.SourceIP = netip.Addr{}
+	// Drop the rest of the per-request telemetry too — a usage-only entry
+	// must carry the LLM usage metadata and caller identity, nothing that
+	// describes the individual request.
+	entry.Method = ""
+	entry.ResponseCode = 0
+	entry.DurationMs = 0
+	entry.BytesUpload = 0
+	entry.BytesDownload = 0
+	entry.Protocol = ""
+
+	if len(entry.Metadata) > 0 {
+		stripped := make(map[string]string, len(usageMetadataKeys))
+		for k := range usageMetadataKeys {
+			if v, ok := entry.Metadata[k]; ok {
+				stripped[k] = v
+			}
+		}
+		entry.Metadata = stripped
+	}
+	return entry
+}
+
+func (l *Logger) log(entry logEntry) {
+	// Fire off the log request in a separate routine.
+	// This increases the possibility of losing a log message
+	// (although it should still get logged in the event of an error),
+	// but it will reduce latency returning the request in the
+	// middleware.
+	// There is also a chance that log messages will arrive at
+	// the server out of order; however, the timestamp should
+	// allow for resolving that on the server.
+	now := timestamppb.Now()
+	select {
+	case l.logSem <- struct{}{}:
+	default:
+		total := l.dropped.Add(1)
+		l.logger.Debugf("access log send dropped: worker limit reached (total dropped: %d)", total)
+		return
+	}
+	go func() {
+		defer func() { <-l.logSem }()
+		logCtx, cancel := context.WithTimeout(context.Background(), logSendTimeout)
+		defer cancel()
+		// Only OIDC sessions have a meaningful user identity.
+		if entry.AuthMechanism != auth.MethodOIDC.String() {
+			entry.UserID = ""
+		}
+
+		var sourceIP string
+		if entry.SourceIP.IsValid() {
+			sourceIP = entry.SourceIP.String()
+		}
+
+		if _, err := l.client.SendAccessLog(logCtx, &proto.SendAccessLogRequest{
+			Log: &proto.AccessLog{
+				LogId:         entry.ID,
+				AccountId:     string(entry.AccountID),
+				Timestamp:     now,
+				ServiceId:     string(entry.ServiceID),
+				Host:          entry.Host,
+				Path:          entry.Path,
+				DurationMs:    entry.DurationMs,
+				Method:        entry.Method,
+				ResponseCode:  entry.ResponseCode,
+				SourceIp:      sourceIP,
+				AuthMechanism: entry.AuthMechanism,
+				UserId:        entry.UserID,
+				AuthSuccess:   entry.AuthSuccess,
+				BytesUpload:   entry.BytesUpload,
+				BytesDownload: entry.BytesDownload,
+				Protocol:      string(entry.Protocol),
+				Metadata:      entry.Metadata,
+				AgentNetwork:  entry.AgentNetwork,
+			},
+		}); err != nil {
+			l.logger.WithFields(log.Fields{
+				"service_id":     entry.ServiceID,
+				"host":           entry.Host,
+				"path":           entry.Path,
+				"duration":       entry.DurationMs,
+				"method":         entry.Method,
+				"response_code":  entry.ResponseCode,
+				"source_ip":      sourceIP,
+				"auth_mechanism": entry.AuthMechanism,
+				"user_id":        entry.UserID,
+				"auth_success":   entry.AuthSuccess,
+				"error":          err,
+			}).Error("Error sending access log on gRPC connection")
+		}
+	}()
+}
+
+// trackUsage records request and byte counts per domain, logging when thresholds are hit.
+func (l *Logger) trackUsage(domain string, bytesTransferred int64) {
+	if domain == "" {
+		return
+	}
+
+	l.usageMux.Lock()
+	defer l.usageMux.Unlock()
+
+	now := time.Now()
+	usage, exists := l.domainUsage[domain]
+	if !exists {
+		usage = &domainUsage{
+			requestStartTime: now,
+			bytesStartTime:   now,
+			lastActivity:     now,
+		}
+		l.domainUsage[domain] = usage
+	}
+
+	usage.lastActivity = now
+
+	usage.requestCount++
+	if usage.requestCount >= requestThreshold {
+		elapsed := time.Since(usage.requestStartTime)
+		l.logger.WithFields(log.Fields{
+			"domain":   domain,
+			"requests": usage.requestCount,
+			"duration": elapsed.String(),
+		}).Infof("domain %s had %d requests over %s", domain, usage.requestCount, elapsed)
+
+		usage.requestCount = 0
+		usage.requestStartTime = now
+	}
+
+	usage.bytesTransferred += bytesTransferred
+	if usage.bytesTransferred >= bytesThreshold {
+		elapsed := time.Since(usage.bytesStartTime)
+		bytesInGB := float64(usage.bytesTransferred) / (1024 * 1024 * 1024)
+		l.logger.WithFields(log.Fields{
+			"domain":   domain,
+			"bytes":    usage.bytesTransferred,
+			"bytes_gb": bytesInGB,
+			"duration": elapsed.String(),
+		}).Infof("domain %s transferred %.2f GB over %s", domain, bytesInGB, elapsed)
+
+		usage.bytesTransferred = 0
+		usage.bytesStartTime = now
+	}
+}
+
+// cleanupStaleUsage removes usage and deny-rate-limit entries that have been inactive.
+func (l *Logger) cleanupStaleUsage(ctx context.Context) {
+	ticker := time.NewTicker(usageCleanupPeriod)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			now := time.Now()
+			l.cleanupDomainUsage(now)
+			l.cleanupDenyBuckets(now)
+		}
+	}
+}
+
+func (l *Logger) cleanupDomainUsage(now time.Time) {
+	l.usageMux.Lock()
+	defer l.usageMux.Unlock()
+
+	removed := 0
+	for domain, usage := range l.domainUsage {
+		if now.Sub(usage.lastActivity) > usageInactiveWindow {
+			delete(l.domainUsage, domain)
+			removed++
+		}
+	}
+	if removed > 0 {
+		l.logger.Debugf("cleaned up %d stale domain usage entries", removed)
+	}
+}
+
+func (l *Logger) cleanupDenyBuckets(now time.Time) {
+	l.denyMu.Lock()
+	defer l.denyMu.Unlock()
+
+	removed := 0
+	for key, bucket := range l.denyBuckets {
+		if now.Sub(bucket.lastLogged) > usageInactiveWindow {
+			delete(l.denyBuckets, key)
+			removed++
+		}
+	}
+	if removed > 0 {
+		l.logger.Debugf("cleaned up %d stale deny rate-limit entries", removed)
+	}
+}

proxy/internal/accesslog/middleware.go

@@ -0,0 +1,107 @@
+package accesslog
+
+import (
+	"net"
+	"net/http"
+	"strings"
+	"time"
+
+	"github.com/rs/xid"
+
+	"github.com/netbirdio/netbird/proxy/internal/proxy"
+	"github.com/netbirdio/netbird/proxy/internal/responsewriter"
+	"github.com/netbirdio/netbird/proxy/web"
+)
+
+// Middleware wraps an HTTP handler to log access entries and resolve client IPs.
+func (l *Logger) Middleware(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// Skip logging for internal proxy assets (CSS, JS, etc.)
+		if strings.HasPrefix(r.URL.Path, web.PathPrefix+"/") {
+			next.ServeHTTP(w, r)
+			return
+		}
+
+		// Generate request ID early so it can be used by error pages and log correlation.
+		requestID := xid.New().String()
+
+		l.logger.Debugf("request: request_id=%s method=%s host=%s path=%s", requestID, r.Method, r.Host, r.URL.Path)
+
+		// Use a response writer wrapper so we can access the status code later.
+		sw := &statusWriter{
+			PassthroughWriter: responsewriter.New(w),
+			status:            http.StatusOK,
+		}
+
+		var bytesRead int64
+		if r.Body != nil {
+			r.Body = &bodyCounter{
+				ReadCloser: r.Body,
+				bytesRead:  &bytesRead,
+			}
+		}
+
+		// Resolve the source IP using trusted proxy configuration before passing
+		// the request on, as the proxy will modify forwarding headers.
+		sourceIp := extractSourceIP(r, l.trustedProxies)
+
+		// Create a mutable struct to capture data from downstream handlers.
+		// We pass a pointer in the context - the pointer itself flows down immutably,
+		// but the struct it points to can be mutated by inner handlers.
+		capturedData := proxy.NewCapturedData(requestID)
+		capturedData.SetClientIP(sourceIp)
+
+		ctx := proxy.WithCapturedData(r.Context(), capturedData)
+
+		start := time.Now()
+		next.ServeHTTP(sw, r.WithContext(ctx))
+		duration := time.Since(start)
+
+		host, _, err := net.SplitHostPort(r.Host)
+		if err != nil {
+			// Fallback to just using the full host value.
+			host = r.Host
+		}
+
+		bytesUpload := bytesRead
+		bytesDownload := sw.bytesWritten
+
+		entry := logEntry{
+			ID:            requestID,
+			ServiceID:     capturedData.GetServiceID(),
+			AccountID:     capturedData.GetAccountID(),
+			Host:          host,
+			Path:          r.URL.Path,
+			DurationMs:    duration.Milliseconds(),
+			Method:        r.Method,
+			ResponseCode:  int32(sw.status),
+			SourceIP:      sourceIp,
+			AuthMechanism: capturedData.GetAuthMethod(),
+			UserID:        capturedData.GetUserID(),
+			AuthSuccess:   sw.status != http.StatusUnauthorized && sw.status != http.StatusForbidden,
+			BytesUpload:   bytesUpload,
+			BytesDownload: bytesDownload,
+			Protocol:      ProtocolHTTP,
+			Metadata:      capturedData.GetMetadata(),
+			AgentNetwork:  capturedData.GetAgentNetwork(),
+		}
+		l.logger.Debugf("response: request_id=%s method=%s host=%s path=%s status=%d duration=%dms source=%s origin=%s service=%s account=%s",
+			requestID, r.Method, host, r.URL.Path, sw.status, duration.Milliseconds(), sourceIp, capturedData.GetOrigin(), capturedData.GetServiceID(), capturedData.GetAccountID())
+
+		// Emit the access log unless the matched target opted out
+		// (agent-network synth targets do this when the account's
+		// EnableLogCollection toggle is off). For agent-network entries we
+		// still ship a stripped, usage-only record even when suppressed, so
+		// usage/cost is collected regardless of the log-collection toggle;
+		// request detail and prompt capture are dropped before sending.
+		switch {
+		case !capturedData.GetSuppressAccessLog():
+			l.log(entry)
+		case entry.AgentNetwork:
+			l.log(stripAgentNetworkEntryForUsage(entry))
+		}
+
+		// Track usage for cost monitoring (upload + download) by domain
+		l.trackUsage(host, bytesUpload+bytesDownload)
+	})
+}

proxy/internal/accesslog/middleware_test.go

@@ -0,0 +1,185 @@
+package accesslog
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc"
+
+	"github.com/netbirdio/netbird/proxy/internal/proxy"
+	"github.com/netbirdio/netbird/shared/management/proto"
+)
+
+// recorderClient is a minimal stub for the access-log gRPCClient interface. It
+// counts SendAccessLog invocations and signals on every call so tests can
+// deterministically wait for the goroutine inside Logger.log without sleeping.
+type recorderClient struct {
+	mu        sync.Mutex
+	calls     int64
+	lastEntry *proto.AccessLog
+	called    chan struct{}
+}
+
+func newRecorderClient() *recorderClient {
+	return &recorderClient{called: make(chan struct{}, 16)}
+}
+
+func (r *recorderClient) SendAccessLog(_ context.Context, in *proto.SendAccessLogRequest, _ ...grpc.CallOption) (*proto.SendAccessLogResponse, error) {
+	r.mu.Lock()
+	r.calls++
+	r.lastEntry = in.GetLog()
+	r.mu.Unlock()
+	select {
+	case r.called <- struct{}{}:
+	default:
+	}
+	return &proto.SendAccessLogResponse{}, nil
+}
+
+func (r *recorderClient) callCount() int64 {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	return r.calls
+}
+
+// newTestLogger builds a Logger backed by the supplied recorderClient. It is
+// the same constructor production uses, just with a stub gRPC client — no
+// mocks, no interface re-implementations.
+func newTestLogger(t *testing.T, client *recorderClient) *Logger {
+	t.Helper()
+	logger := NewLogger(client, nil, nil)
+	t.Cleanup(logger.Close)
+	return logger
+}
+
+// TestMiddleware_SuppressAccessLog_SkipsLogSink asserts the suppression gate.
+// When the inner handler stamps SuppressAccessLog=true on CapturedData (mirrors
+// what reverseproxy does when the matched target's DisableAccessLog flag is
+// set), the middleware must NOT invoke the access-log sink. Bandwidth telemetry
+// (trackUsage) keeps running — it's the call to SendAccessLog that we gate.
+func TestMiddleware_SuppressAccessLog_SkipsLogSink(t *testing.T) {
+	client := newRecorderClient()
+	l := newTestLogger(t, client)
+
+	inner := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		cd := proxy.CapturedDataFromContext(r.Context())
+		require.NotNil(t, cd, "middleware must inject CapturedData into the request context")
+		cd.SetSuppressAccessLog(true)
+		w.WriteHeader(http.StatusOK)
+	})
+
+	srv := httptest.NewServer(l.Middleware(inner))
+	t.Cleanup(srv.Close)
+
+	resp, err := http.Get(srv.URL + "/agent-network/v1/chat/completions")
+	require.NoError(t, err, "GET against suppressed target must succeed")
+	require.NoError(t, resp.Body.Close())
+	require.Equal(t, http.StatusOK, resp.StatusCode, "inner handler must run normally")
+
+	// Give the goroutine fence a beat (Logger.log dispatches in a goroutine).
+	// The negative assertion needs a small window: if a send is going to
+	// happen, it happens promptly.
+	select {
+	case <-client.called:
+		t.Fatalf("access-log sink must not be invoked when SuppressAccessLog=true (got %d call(s))", client.callCount())
+	case <-time.After(150 * time.Millisecond):
+	}
+
+	assert.Equal(t, int64(0), client.callCount(),
+		"SendAccessLog must not be called for suppressed requests")
+}
+
+// TestMiddleware_SuppressAccessLog_DefaultEmitsLog is the regression sanity:
+// when nothing sets SuppressAccessLog (the universal default for every
+// non-agent-network target), the middleware MUST still emit the access-log
+// entry. This is the guarantee that wires-through to the EnableLogCollection
+// gate without breaking anyone who isn't opted in.
+func TestMiddleware_SuppressAccessLog_DefaultEmitsLog(t *testing.T) {
+	client := newRecorderClient()
+	l := newTestLogger(t, client)
+
+	var innerRan atomic.Bool
+	inner := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		innerRan.Store(true)
+		// Intentionally DO NOT touch SuppressAccessLog — mirrors every
+		// non-agent-network target.
+		w.WriteHeader(http.StatusOK)
+	})
+
+	srv := httptest.NewServer(l.Middleware(inner))
+	t.Cleanup(srv.Close)
+
+	resp, err := http.Get(srv.URL + "/service/healthz")
+	require.NoError(t, err, "GET against default target must succeed")
+	require.NoError(t, resp.Body.Close())
+	require.Equal(t, http.StatusOK, resp.StatusCode, "inner handler must run normally")
+	require.True(t, innerRan.Load(), "inner handler must have run")
+
+	select {
+	case <-client.called:
+	case <-time.After(2 * time.Second):
+		t.Fatalf("SendAccessLog must be invoked for non-suppressed requests, none observed (calls=%d)", client.callCount())
+	}
+
+	assert.Equal(t, int64(1), client.callCount(),
+		"non-suppressed request must produce exactly one access-log send")
+}
+
+// TestMiddleware_SuppressAccessLog_PreservesUsageTracking proves the gate is
+// surgical: with SuppressAccessLog=true the access-log send is skipped, but
+// the per-domain usage tracker still records the bytes transferred. This is
+// the cost-monitoring guarantee called out in the gate's comment.
+func TestMiddleware_SuppressAccessLog_PreservesUsageTracking(t *testing.T) {
+	client := newRecorderClient()
+	l := newTestLogger(t, client)
+
+	payload := []byte("ok")
+	inner := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		cd := proxy.CapturedDataFromContext(r.Context())
+		require.NotNil(t, cd, "middleware must inject CapturedData")
+		cd.SetSuppressAccessLog(true)
+		w.WriteHeader(http.StatusOK)
+		_, _ = w.Write(payload)
+	})
+
+	srv := httptest.NewServer(l.Middleware(inner))
+	t.Cleanup(srv.Close)
+
+	resp, err := http.Get(srv.URL + "/agent-network/v1/chat/completions")
+	require.NoError(t, err)
+	require.NoError(t, resp.Body.Close())
+
+	// Allow trackUsage to land — it runs synchronously after l.log(entry) is
+	// (would have been) called.
+	time.Sleep(50 * time.Millisecond)
+
+	l.usageMux.Lock()
+	usage, present := l.domainUsage[hostNoPort(srv.URL)]
+	l.usageMux.Unlock()
+	require.True(t, present, "domain usage must be tracked even when the access-log is suppressed")
+	assert.Greater(t, usage.bytesTransferred, int64(0), "bytesTransferred must include the response payload")
+	assert.Equal(t, int64(0), client.callCount(),
+		"SendAccessLog must remain suppressed across the response write")
+}
+
+// hostNoPort extracts the host name from an httptest server URL. The
+// middleware strips the port before keying domain usage, so the test mirrors
+// that to look the entry up.
+func hostNoPort(url string) string {
+	// httptest URLs are always "http://127.0.0.1:PORT".
+	const prefix = "http://"
+	host := url[len(prefix):]
+	for i := 0; i < len(host); i++ {
+		if host[i] == ':' || host[i] == '/' {
+			return host[:i]
+		}
+	}
+	return host
+}

proxy/internal/accesslog/requestip.go

@@ -0,0 +1,16 @@
+package accesslog
+
+import (
+	"net/http"
+	"net/netip"
+
+	"github.com/netbirdio/netbird/trustedproxy"
+)
+
+// extractSourceIP resolves the real client IP from the request using trusted
+// proxy configuration. When trustedProxies is non-empty and the direct
+// connection is from a trusted source, it walks X-Forwarded-For right-to-left
+// skipping trusted IPs. Otherwise it returns RemoteAddr directly.
+func extractSourceIP(r *http.Request, trustedProxies *trustedproxy.List) netip.Addr {
+	return trustedProxies.ResolveClientIP(r.RemoteAddr, r.Header.Get("X-Forwarded-For"))
+}

proxy/internal/accesslog/statuswriter.go

@@ -0,0 +1,39 @@
+package accesslog
+
+import (
+	"io"
+
+	"github.com/netbirdio/netbird/proxy/internal/responsewriter"
+)
+
+// statusWriter captures the HTTP status code and bytes written from responses.
+// It embeds responsewriter.PassthroughWriter which handles all the optional
+// interfaces (Hijacker, Flusher, Pusher) automatically.
+type statusWriter struct {
+	*responsewriter.PassthroughWriter
+	status       int
+	bytesWritten int64
+}
+
+func (w *statusWriter) WriteHeader(status int) {
+	w.status = status
+	w.PassthroughWriter.WriteHeader(status)
+}
+
+func (w *statusWriter) Write(b []byte) (int, error) {
+	n, err := w.PassthroughWriter.Write(b)
+	w.bytesWritten += int64(n)
+	return n, err
+}
+
+// bodyCounter wraps an io.ReadCloser and counts bytes read from the request body.
+type bodyCounter struct {
+	io.ReadCloser
+	bytesRead *int64
+}
+
+func (bc *bodyCounter) Read(p []byte) (int, error) {
+	n, err := bc.ReadCloser.Read(p)
+	*bc.bytesRead += int64(n)
+	return n, err
+}

proxy/internal/acme/locker.go

@@ -0,0 +1,102 @@
+package acme
+
+import (
+	"context"
+	"path/filepath"
+
+	log "github.com/sirupsen/logrus"
+
+	"github.com/netbirdio/netbird/proxy/internal/flock"
+	"github.com/netbirdio/netbird/proxy/internal/k8s"
+)
+
+// certLocker provides distributed mutual exclusion for certificate operations.
+// Implementations must be safe for concurrent use from multiple goroutines.
+type certLocker interface {
+	// Lock acquires an exclusive lock for the given domain.
+	// It blocks until the lock is acquired, the context is cancelled, or an
+	// unrecoverable error occurs. The returned function releases the lock;
+	// callers must call it exactly once when the critical section is complete.
+	Lock(ctx context.Context, domain string) (unlock func(), err error)
+}
+
+// CertLockMethod controls how ACME certificate locks are coordinated.
+type CertLockMethod string
+
+const (
+	// CertLockAuto detects the environment and selects k8s-lease if running
+	// in a Kubernetes pod, otherwise flock.
+	CertLockAuto CertLockMethod = "auto"
+	// CertLockFlock uses advisory file locks via flock(2).
+	CertLockFlock CertLockMethod = "flock"
+	// CertLockK8sLease uses Kubernetes coordination Leases.
+	CertLockK8sLease CertLockMethod = "k8s-lease"
+)
+
+func newCertLocker(method CertLockMethod, certDir string, logger *log.Logger) certLocker {
+	if logger == nil {
+		logger = log.StandardLogger()
+	}
+
+	if method == "" || method == CertLockAuto {
+		if k8s.InCluster() {
+			method = CertLockK8sLease
+		} else {
+			method = CertLockFlock
+		}
+		logger.Infof("auto-detected cert lock method: %s", method)
+	}
+
+	switch method {
+	case CertLockK8sLease:
+		locker, err := newK8sLeaseLocker(logger)
+		if err != nil {
+			logger.Warnf("create k8s lease locker, falling back to flock: %v", err)
+			return newFlockLocker(certDir, logger)
+		}
+		logger.Infof("using k8s lease locker in namespace %s", locker.client.Namespace())
+		return locker
+	default:
+		logger.Infof("using flock cert locker in %s", certDir)
+		return newFlockLocker(certDir, logger)
+	}
+}
+
+type flockLocker struct {
+	certDir string
+	logger  *log.Logger
+}
+
+func newFlockLocker(certDir string, logger *log.Logger) *flockLocker {
+	if logger == nil {
+		logger = log.StandardLogger()
+	}
+	return &flockLocker{certDir: certDir, logger: logger}
+}
+
+// Lock acquires an advisory file lock for the given domain.
+func (l *flockLocker) Lock(ctx context.Context, domain string) (func(), error) {
+	lockPath := filepath.Join(l.certDir, domain+".lock")
+	lockFile, err := flock.Lock(ctx, lockPath)
+	if err != nil {
+		return nil, err
+	}
+
+	// nil lockFile means locking is not supported (non-unix).
+	if lockFile == nil {
+		return func() { /* no-op: locking unsupported on this platform */ }, nil
+	}
+
+	return func() {
+		if err := flock.Unlock(lockFile); err != nil {
+			l.logger.Debugf("release cert lock for domain %q: %v", domain, err)
+		}
+	}, nil
+}
+
+type noopLocker struct{}
+
+// Lock is a no-op that always succeeds immediately.
+func (noopLocker) Lock(context.Context, string) (func(), error) {
+	return func() { /* no-op: locker disabled */ }, nil
+}

proxy/internal/acme/locker_k8s.go

@@ -0,0 +1,197 @@
+package acme
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"os"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+
+	"github.com/netbirdio/netbird/proxy/internal/k8s"
+)
+
+const (
+	// leaseDurationSec is the Kubernetes Lease TTL. If the holder crashes without
+	// releasing the lock, other replicas must wait this long before taking over.
+	// This is intentionally generous: in the worst case two replicas may both
+	// issue an ACME request for the same domain, which is harmless (the CA
+	// deduplicates and the cache converges).
+	leaseDurationSec = 300
+	retryBaseBackoff = 500 * time.Millisecond
+	retryMaxBackoff  = 10 * time.Second
+)
+
+type k8sLeaseLocker struct {
+	client   *k8s.LeaseClient
+	identity string
+	logger   *log.Logger
+}
+
+func newK8sLeaseLocker(logger *log.Logger) (*k8sLeaseLocker, error) {
+	client, err := k8s.NewLeaseClient()
+	if err != nil {
+		return nil, fmt.Errorf("create k8s lease client: %w", err)
+	}
+
+	identity, err := os.Hostname()
+	if err != nil {
+		return nil, fmt.Errorf("get hostname: %w", err)
+	}
+
+	return &k8sLeaseLocker{
+		client:   client,
+		identity: identity,
+		logger:   logger,
+	}, nil
+}
+
+// Lock acquires a Kubernetes Lease for the given domain using optimistic
+// concurrency. It retries with exponential backoff until the lease is
+// acquired or the context is cancelled.
+func (l *k8sLeaseLocker) Lock(ctx context.Context, domain string) (func(), error) {
+	leaseName := k8s.LeaseNameForDomain(domain)
+	backoff := retryBaseBackoff
+
+	for {
+		acquired, err := l.tryAcquire(ctx, leaseName, domain)
+		if err != nil {
+			return nil, fmt.Errorf("acquire lease %s for %q: %w", leaseName, domain, err)
+		}
+		if acquired {
+			l.logger.Debugf("k8s lease %s acquired for domain %q", leaseName, domain)
+			return l.unlockFunc(leaseName, domain), nil
+		}
+
+		l.logger.Debugf("k8s lease %s held by another replica, retrying in %s", leaseName, backoff)
+
+		timer := time.NewTimer(backoff)
+		select {
+		case <-ctx.Done():
+			timer.Stop()
+			return nil, ctx.Err()
+		case <-timer.C:
+		}
+
+		backoff *= 2
+		if backoff > retryMaxBackoff {
+			backoff = retryMaxBackoff
+		}
+	}
+}
+
+// tryAcquire attempts to create or take over a Lease. Returns (true, nil)
+// on success, (false, nil) if the lease is held and not stale, or an error.
+func (l *k8sLeaseLocker) tryAcquire(ctx context.Context, name, domain string) (bool, error) {
+	existing, err := l.client.Get(ctx, name)
+	if err != nil {
+		return false, err
+	}
+
+	now := k8s.MicroTime{Time: time.Now().UTC()}
+	dur := int32(leaseDurationSec)
+
+	if existing == nil {
+		lease := &k8s.Lease{
+			Metadata: k8s.LeaseMetadata{
+				Name: name,
+				Annotations: map[string]string{
+					"netbird.io/domain": domain,
+				},
+			},
+			Spec: k8s.LeaseSpec{
+				HolderIdentity:       &l.identity,
+				LeaseDurationSeconds: &dur,
+				AcquireTime:          &now,
+				RenewTime:            &now,
+			},
+		}
+
+		if _, err := l.client.Create(ctx, lease); errors.Is(err, k8s.ErrConflict) {
+			return false, nil
+		} else if err != nil {
+			return false, err
+		}
+		return true, nil
+	}
+
+	if !l.canTakeover(existing) {
+		return false, nil
+	}
+
+	existing.Spec.HolderIdentity = &l.identity
+	existing.Spec.LeaseDurationSeconds = &dur
+	existing.Spec.AcquireTime = &now
+	existing.Spec.RenewTime = &now
+
+	if _, err := l.client.Update(ctx, existing); errors.Is(err, k8s.ErrConflict) {
+		return false, nil
+	} else if err != nil {
+		return false, err
+	}
+	return true, nil
+}
+
+// canTakeover returns true if the lease is free (no holder) or stale
+// (renewTime + leaseDuration has passed).
+func (l *k8sLeaseLocker) canTakeover(lease *k8s.Lease) bool {
+	holder := lease.Spec.HolderIdentity
+	if holder == nil || *holder == "" {
+		return true
+	}
+
+	// We already hold it (e.g. from a previous crashed attempt).
+	if *holder == l.identity {
+		return true
+	}
+
+	if lease.Spec.RenewTime == nil || lease.Spec.LeaseDurationSeconds == nil {
+		return true
+	}
+
+	expiry := lease.Spec.RenewTime.Add(time.Duration(*lease.Spec.LeaseDurationSeconds) * time.Second)
+	if time.Now().After(expiry) {
+		l.logger.Infof("k8s lease %s held by %q is stale (expired %s ago), taking over",
+			lease.Metadata.Name, *holder, time.Since(expiry).Round(time.Second))
+		return true
+	}
+
+	return false
+}
+
+// unlockFunc returns a closure that releases the lease by clearing the holder.
+func (l *k8sLeaseLocker) unlockFunc(name, domain string) func() {
+	return func() {
+		// Use a fresh context: the parent may already be cancelled.
+		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+		defer cancel()
+
+		// Re-GET to get current resourceVersion (ours may be stale if
+		// the lock was held for a long time and something updated it).
+		current, err := l.client.Get(ctx, name)
+		if err != nil {
+			l.logger.Debugf("release k8s lease %s for %q: get: %v", name, domain, err)
+			return
+		}
+		if current == nil {
+			return
+		}
+
+		// Only clear if we're still the holder.
+		if current.Spec.HolderIdentity == nil || *current.Spec.HolderIdentity != l.identity {
+			l.logger.Debugf("k8s lease %s for %q: holder changed to %v, skip release",
+				name, domain, current.Spec.HolderIdentity)
+			return
+		}
+
+		empty := ""
+		current.Spec.HolderIdentity = &empty
+		current.Spec.AcquireTime = nil
+		current.Spec.RenewTime = nil
+
+		if _, err := l.client.Update(ctx, current); err != nil {
+			l.logger.Debugf("release k8s lease %s for %q: update: %v", name, domain, err)
+		}
+	}
+}

proxy/internal/acme/locker_test.go

@@ -0,0 +1,65 @@
+package acme
+
+import (
+	"context"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestFlockLockerRoundTrip(t *testing.T) {
+	dir := t.TempDir()
+	locker := newFlockLocker(dir, nil)
+
+	unlock, err := locker.Lock(context.Background(), "example.com")
+	require.NoError(t, err)
+	require.NotNil(t, unlock)
+
+	// Lock file should exist.
+	assert.FileExists(t, filepath.Join(dir, "example.com.lock"))
+
+	unlock()
+}
+
+func TestNoopLocker(t *testing.T) {
+	locker := noopLocker{}
+	unlock, err := locker.Lock(context.Background(), "example.com")
+	require.NoError(t, err)
+	require.NotNil(t, unlock)
+	unlock()
+}
+
+func TestNewCertLockerDefaultsToFlock(t *testing.T) {
+	dir := t.TempDir()
+
+	// t.Setenv registers cleanup to restore the original value.
+	// os.Unsetenv is needed because the production code uses LookupEnv,
+	// which distinguishes "empty" from "not set".
+	t.Setenv("KUBERNETES_SERVICE_HOST", "")
+	os.Unsetenv("KUBERNETES_SERVICE_HOST")
+	locker := newCertLocker(CertLockAuto, dir, nil)
+
+	_, ok := locker.(*flockLocker)
+	assert.True(t, ok, "auto without k8s env should select flockLocker")
+}
+
+func TestNewCertLockerExplicitFlock(t *testing.T) {
+	dir := t.TempDir()
+	locker := newCertLocker(CertLockFlock, dir, nil)
+
+	_, ok := locker.(*flockLocker)
+	assert.True(t, ok, "explicit flock should select flockLocker")
+}
+
+func TestNewCertLockerK8sFallsBackToFlock(t *testing.T) {
+	dir := t.TempDir()
+
+	// k8s-lease without SA files should fall back to flock.
+	locker := newCertLocker(CertLockK8sLease, dir, nil)
+
+	_, ok := locker.(*flockLocker)
+	assert.True(t, ok, "k8s-lease without SA should fall back to flockLocker")
+}

proxy/internal/acme/manager.go

@@ -0,0 +1,643 @@
+package acme
+
+import (
+	"context"
+	"crypto/tls"
+	"crypto/x509"
+	"encoding/asn1"
+	"encoding/base64"
+	"encoding/binary"
+	"encoding/pem"
+	"fmt"
+	"math/rand/v2"
+	"net"
+	"os"
+	"path/filepath"
+	"slices"
+	"strings"
+	"sync"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+	"golang.org/x/crypto/acme"
+	"golang.org/x/crypto/acme/autocert"
+
+	"github.com/netbirdio/netbird/proxy/internal/certwatch"
+	"github.com/netbirdio/netbird/proxy/internal/types"
+	"github.com/netbirdio/netbird/shared/management/domain"
+)
+
+// OID for the SCT list extension (1.3.6.1.4.1.11129.2.4.2)
+var oidSCTList = asn1.ObjectIdentifier{1, 3, 6, 1, 4, 1, 11129, 2, 4, 2}
+
+type certificateNotifier interface {
+	NotifyCertificateIssued(ctx context.Context, accountID types.AccountID, serviceID types.ServiceID, domain string) error
+}
+
+type domainState int
+
+const (
+	domainPending domainState = iota
+	domainReady
+	domainFailed
+)
+
+type domainInfo struct {
+	accountID types.AccountID
+	serviceID types.ServiceID
+	state     domainState
+	err       string
+}
+
+type metricsRecorder interface {
+	RecordCertificateIssuance(duration time.Duration)
+}
+
+// wildcardEntry maps a domain suffix (e.g. ".example.com") to a certwatch
+// watcher that hot-reloads the corresponding wildcard certificate from disk.
+type wildcardEntry struct {
+	suffix  string // e.g. ".example.com"
+	pattern string // e.g. "*.example.com"
+	watcher *certwatch.Watcher
+}
+
+// ManagerConfig holds the configuration values for the ACME certificate manager.
+type ManagerConfig struct {
+	// CertDir is the directory used for caching ACME certificates.
+	CertDir string
+	// ACMEURL is the ACME directory URL (e.g. Let's Encrypt).
+	ACMEURL string
+	// EABKID and EABHMACKey are optional External Account Binding credentials
+	// required by some CAs (e.g. ZeroSSL). EABHMACKey is the base64
+	// URL-encoded string provided by the CA.
+	EABKID     string
+	EABHMACKey string
+	// LockMethod controls the cross-replica coordination strategy.
+	LockMethod CertLockMethod
+	// WildcardDir is an optional path to a directory containing wildcard
+	// certificate pairs (<name>.crt / <name>.key). Wildcard patterns are
+	// extracted from the certificates' SAN lists. Domains matching a
+	// wildcard are served from disk; all others go through ACME.
+	WildcardDir string
+}
+
+// Manager wraps autocert.Manager with domain tracking and cross-replica
+// coordination via a pluggable locking strategy. The locker prevents
+// duplicate ACME requests when multiple replicas share a certificate cache.
+type Manager struct {
+	*autocert.Manager
+
+	certDir string
+	locker  certLocker
+	mu      sync.RWMutex
+	domains map[domain.Domain]*domainInfo
+
+	// wildcards holds all loaded wildcard certificates, keyed by suffix.
+	wildcards []wildcardEntry
+
+	certNotifier certificateNotifier
+	logger       *log.Logger
+	metrics      metricsRecorder
+}
+
+// NewManager creates a new ACME certificate manager.
+func NewManager(cfg ManagerConfig, notifier certificateNotifier, logger *log.Logger, metrics metricsRecorder) (*Manager, error) {
+	if logger == nil {
+		logger = log.StandardLogger()
+	}
+	mgr := &Manager{
+		certDir:      cfg.CertDir,
+		locker:       newCertLocker(cfg.LockMethod, cfg.CertDir, logger),
+		domains:      make(map[domain.Domain]*domainInfo),
+		certNotifier: notifier,
+		logger:       logger,
+		metrics:      metrics,
+	}
+
+	if cfg.WildcardDir != "" {
+		entries, err := loadWildcardDir(cfg.WildcardDir, logger)
+		if err != nil {
+			return nil, fmt.Errorf("load wildcard certificates from %q: %w", cfg.WildcardDir, err)
+		}
+		mgr.wildcards = entries
+	}
+
+	var eab *acme.ExternalAccountBinding
+	if cfg.EABKID != "" && cfg.EABHMACKey != "" {
+		decodedKey, err := base64.RawURLEncoding.DecodeString(cfg.EABHMACKey)
+		if err != nil {
+			logger.Errorf("failed to decode EAB HMAC key: %v", err)
+		} else {
+			eab = &acme.ExternalAccountBinding{
+				KID: cfg.EABKID,
+				Key: decodedKey,
+			}
+			logger.Infof("configured External Account Binding with KID: %s", cfg.EABKID)
+		}
+	}
+
+	mgr.Manager = &autocert.Manager{
+		Prompt:                 autocert.AcceptTOS,
+		HostPolicy:             mgr.hostPolicy,
+		Cache:                  autocert.DirCache(cfg.CertDir),
+		ExternalAccountBinding: eab,
+		Client: &acme.Client{
+			DirectoryURL: cfg.ACMEURL,
+		},
+	}
+	return mgr, nil
+}
+
+// WatchWildcards starts watching all wildcard certificate files for changes.
+// It blocks until ctx is cancelled. It is a no-op if no wildcards are loaded.
+func (mgr *Manager) WatchWildcards(ctx context.Context) {
+	if len(mgr.wildcards) == 0 {
+		return
+	}
+	seen := make(map[*certwatch.Watcher]struct{})
+	var wg sync.WaitGroup
+	for i := range mgr.wildcards {
+		w := mgr.wildcards[i].watcher
+		if _, ok := seen[w]; ok {
+			continue
+		}
+		seen[w] = struct{}{}
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			w.Watch(ctx)
+		}()
+	}
+	wg.Wait()
+}
+
+// loadWildcardDir scans dir for .crt files, pairs each with a matching .key
+// file, loads them, and extracts wildcard SANs (*.example.com) to build
+// the suffix lookup entries.
+func loadWildcardDir(dir string, logger *log.Logger) ([]wildcardEntry, error) {
+	crtFiles, err := filepath.Glob(filepath.Join(dir, "*.crt"))
+	if err != nil {
+		return nil, fmt.Errorf("glob certificate files: %w", err)
+	}
+
+	if len(crtFiles) == 0 {
+		return nil, fmt.Errorf("no .crt files found in %s", dir)
+	}
+
+	var entries []wildcardEntry
+
+	for _, crtPath := range crtFiles {
+		base := strings.TrimSuffix(filepath.Base(crtPath), ".crt")
+		keyPath := filepath.Join(dir, base+".key")
+		if _, err := os.Stat(keyPath); err != nil {
+			logger.Warnf("skipping %s: no matching key file %s", crtPath, keyPath)
+			continue
+		}
+
+		watcher, err := certwatch.NewWatcher(crtPath, keyPath, logger)
+		if err != nil {
+			logger.Warnf("skipping %s: %v", crtPath, err)
+			continue
+		}
+
+		leaf := watcher.Leaf()
+		if leaf == nil {
+			logger.Warnf("skipping %s: no parsed leaf certificate", crtPath)
+			continue
+		}
+
+		for _, san := range leaf.DNSNames {
+			suffix, ok := parseWildcard(san)
+			if !ok {
+				continue
+			}
+			entries = append(entries, wildcardEntry{
+				suffix:  suffix,
+				pattern: san,
+				watcher: watcher,
+			})
+			logger.Infof("wildcard certificate loaded: %s (from %s)", san, filepath.Base(crtPath))
+		}
+	}
+
+	if len(entries) == 0 {
+		return nil, fmt.Errorf("no wildcard SANs (*.example.com) found in certificates in %s", dir)
+	}
+
+	return entries, nil
+}
+
+// parseWildcard validates a wildcard domain pattern like "*.example.com"
+// and returns the suffix ".example.com" for matching.
+func parseWildcard(pattern string) (suffix string, ok bool) {
+	if !strings.HasPrefix(pattern, "*.") {
+		return "", false
+	}
+	parent := pattern[1:] // ".example.com"
+	if strings.Count(parent, ".") < 1 {
+		return "", false
+	}
+	return strings.ToLower(parent), true
+}
+
+// findWildcardEntry returns the wildcard entry that covers host, or nil.
+func (mgr *Manager) findWildcardEntry(host string) *wildcardEntry {
+	if len(mgr.wildcards) == 0 {
+		return nil
+	}
+	host = strings.ToLower(host)
+	for i := range mgr.wildcards {
+		e := &mgr.wildcards[i]
+		if !strings.HasSuffix(host, e.suffix) {
+			continue
+		}
+		// Single-level match: prefix before suffix must have no dots.
+		prefix := strings.TrimSuffix(host, e.suffix)
+		if len(prefix) > 0 && !strings.Contains(prefix, ".") {
+			return e
+		}
+	}
+	return nil
+}
+
+// WildcardPatterns returns the wildcard patterns that are currently loaded.
+func (mgr *Manager) WildcardPatterns() []string {
+	patterns := make([]string, len(mgr.wildcards))
+	for i, e := range mgr.wildcards {
+		patterns[i] = e.pattern
+	}
+	slices.Sort(patterns)
+	return patterns
+}
+
+func (mgr *Manager) hostPolicy(_ context.Context, host string) error {
+	if h, _, err := net.SplitHostPort(host); err == nil {
+		host = h
+	}
+	mgr.mu.RLock()
+	_, exists := mgr.domains[domain.Domain(host)]
+	mgr.mu.RUnlock()
+	if !exists {
+		return fmt.Errorf("unknown domain %q", host)
+	}
+	return nil
+}
+
+// GetCertificate returns the TLS certificate for the given ClientHello.
+// If the requested domain matches a loaded wildcard, the static wildcard
+// certificate is returned. Otherwise, the ACME autocert manager handles
+// the request.
+func (mgr *Manager) GetCertificate(hello *tls.ClientHelloInfo) (*tls.Certificate, error) {
+	if e := mgr.findWildcardEntry(hello.ServerName); e != nil {
+		return e.watcher.GetCertificate(hello)
+	}
+	return mgr.Manager.GetCertificate(hello)
+}
+
+// AddDomain registers a domain for certificate management. Domains that
+// match a loaded wildcard are marked ready immediately (they use the
+// static wildcard certificate) and the method returns true. All other
+// domains go through ACME prefetch and the method returns false.
+//
+// When AddDomain returns true the caller is responsible for sending any
+// certificate-ready notifications after the surrounding operation (e.g.
+// mapping update) has committed successfully.
+func (mgr *Manager) AddDomain(d domain.Domain, accountID types.AccountID, serviceID types.ServiceID) (wildcardHit bool) {
+	name := d.PunycodeString()
+	if e := mgr.findWildcardEntry(name); e != nil {
+		mgr.mu.Lock()
+		mgr.domains[d] = &domainInfo{
+			accountID: accountID,
+			serviceID: serviceID,
+			state:     domainReady,
+		}
+		mgr.mu.Unlock()
+		mgr.logger.Debugf("domain %q matches wildcard %q, using static certificate", name, e.pattern)
+		return true
+	}
+
+	mgr.mu.Lock()
+	mgr.domains[d] = &domainInfo{
+		accountID: accountID,
+		serviceID: serviceID,
+		state:     domainPending,
+	}
+	mgr.mu.Unlock()
+
+	go mgr.prefetchCertificate(d)
+	return false
+}
+
+// prefetchCertificate proactively triggers certificate generation for a domain.
+// It acquires a distributed lock to prevent multiple replicas from issuing
+// duplicate ACME requests. The second replica will block until the first
+// finishes, then find the certificate in the cache.
+// ACME and periodic disk reads race; whichever produces a valid certificate
+// first wins. This handles cases where locking is unreliable and another
+// replica already wrote the cert to the shared cache.
+func (mgr *Manager) prefetchCertificate(d domain.Domain) {
+	time.Sleep(time.Duration(rand.IntN(200)) * time.Millisecond)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
+	defer cancel()
+
+	name := d.PunycodeString()
+
+	mgr.logger.Infof("acquiring cert lock for domain %q", name)
+	lockStart := time.Now()
+	unlock, err := mgr.locker.Lock(ctx, name)
+	if err != nil {
+		mgr.logger.Warnf("acquire cert lock for domain %q, proceeding without lock: %v", name, err)
+	} else {
+		mgr.logger.Infof("acquired cert lock for domain %q in %s", name, time.Since(lockStart))
+		defer unlock()
+	}
+
+	if cert, err := mgr.readCertFromDisk(ctx, name); err == nil {
+		mgr.logger.Infof("certificate for domain %q already on disk, skipping ACME", name)
+		mgr.recordAndNotify(ctx, d, name, cert, 0)
+		return
+	}
+
+	// Run ACME in a goroutine so we can race it against periodic disk reads.
+	// autocert uses its own internal context and cannot be cancelled externally.
+	type acmeResult struct {
+		cert *tls.Certificate
+		err  error
+	}
+	acmeCh := make(chan acmeResult, 1)
+	hello := &tls.ClientHelloInfo{ServerName: name, Conn: &dummyConn{ctx: ctx}}
+	go func() {
+		cert, err := mgr.GetCertificate(hello)
+		acmeCh <- acmeResult{cert, err}
+	}()
+
+	start := time.Now()
+	diskTicker := time.NewTicker(5 * time.Second)
+	defer diskTicker.Stop()
+
+	for {
+		select {
+		case res := <-acmeCh:
+			elapsed := time.Since(start)
+			if res.err != nil {
+				mgr.logger.Warnf("prefetch certificate for domain %q in %s: %v", name, elapsed.String(), res.err)
+				mgr.setDomainState(d, domainFailed, res.err.Error())
+				return
+			}
+			mgr.recordAndNotify(ctx, d, name, res.cert, elapsed)
+			return
+
+		case <-diskTicker.C:
+			cert, err := mgr.readCertFromDisk(context.Background(), name)
+			if err != nil {
+				continue
+			}
+			mgr.logger.Infof("certificate for domain %q appeared on disk after %s", name, time.Since(start).Round(time.Millisecond))
+			// Drain the ACME goroutine before marking ready — autocert holds
+			// an internal write lock on certState while ACME is in flight.
+			go func() {
+				select {
+				case <-acmeCh:
+				default:
+				}
+				mgr.recordAndNotify(context.Background(), d, name, cert, 0)
+			}()
+			return
+
+		case <-ctx.Done():
+			mgr.logger.Warnf("prefetch certificate for domain %q timed out", name)
+			mgr.setDomainState(d, domainFailed, ctx.Err().Error())
+			return
+		}
+	}
+}
+
+// readCertFromDisk reads and parses a certificate directly from the autocert
+// DirCache, bypassing autocert's internal certState mutex. Safe to call
+// concurrently with an in-flight ACME request for the same domain.
+func (mgr *Manager) readCertFromDisk(ctx context.Context, name string) (*tls.Certificate, error) {
+	if mgr.Cache == nil {
+		return nil, fmt.Errorf("no cache configured")
+	}
+	data, err := mgr.Cache.Get(ctx, name)
+	if err != nil {
+		return nil, err
+	}
+	privBlock, certsPEM := pem.Decode(data)
+	if privBlock == nil || !strings.Contains(privBlock.Type, "PRIVATE") {
+		return nil, fmt.Errorf("no private key in cache for %q", name)
+	}
+	cert, err := tls.X509KeyPair(certsPEM, pem.EncodeToMemory(privBlock))
+	if err != nil {
+		return nil, fmt.Errorf("parse cached certificate for %q: %w", name, err)
+	}
+	if len(cert.Certificate) > 0 {
+		leaf, err := x509.ParseCertificate(cert.Certificate[0])
+		if err != nil {
+			return nil, fmt.Errorf("parse leaf for %q: %w", name, err)
+		}
+		if time.Now().After(leaf.NotAfter) {
+			return nil, fmt.Errorf("cached certificate for %q expired at %s", name, leaf.NotAfter)
+		}
+		cert.Leaf = leaf
+	}
+	return &cert, nil
+}
+
+// recordAndNotify records metrics, marks the domain ready, logs cert details,
+// and notifies the cert notifier.
+func (mgr *Manager) recordAndNotify(ctx context.Context, d domain.Domain, name string, cert *tls.Certificate, elapsed time.Duration) {
+	if elapsed > 0 && mgr.metrics != nil {
+		mgr.metrics.RecordCertificateIssuance(elapsed)
+	}
+	mgr.setDomainState(d, domainReady, "")
+	now := time.Now()
+	if cert != nil && cert.Leaf != nil {
+		leaf := cert.Leaf
+		mgr.logger.Infof("certificate for domain %q ready in %s: serial=%s SANs=%v notBefore=%s, notAfter=%s, now=%s",
+			name, elapsed.Round(time.Millisecond),
+			leaf.SerialNumber.Text(16),
+			leaf.DNSNames,
+			leaf.NotBefore.UTC().Format(time.RFC3339),
+			leaf.NotAfter.UTC().Format(time.RFC3339),
+			now.UTC().Format(time.RFC3339),
+		)
+		mgr.logCertificateDetails(name, leaf, now)
+	} else {
+		mgr.logger.Infof("certificate for domain %q ready in %s", name, elapsed.Round(time.Millisecond))
+	}
+	mgr.mu.RLock()
+	info := mgr.domains[d]
+	mgr.mu.RUnlock()
+	if info != nil && mgr.certNotifier != nil {
+		if err := mgr.certNotifier.NotifyCertificateIssued(ctx, info.accountID, info.serviceID, name); err != nil {
+			mgr.logger.Warnf("notify certificate ready for domain %q: %v", name, err)
+		}
+	}
+}
+
+func (mgr *Manager) setDomainState(d domain.Domain, state domainState, errMsg string) {
+	mgr.mu.Lock()
+	defer mgr.mu.Unlock()
+	if info, ok := mgr.domains[d]; ok {
+		info.state = state
+		info.err = errMsg
+	}
+}
+
+// logCertificateDetails logs certificate validity and SCT timestamps.
+func (mgr *Manager) logCertificateDetails(domain string, cert *x509.Certificate, now time.Time) {
+	if cert.NotBefore.After(now) {
+		mgr.logger.Warnf("certificate for %q NotBefore is in the future by %v", domain, cert.NotBefore.Sub(now))
+	}
+
+	sctTimestamps := mgr.parseSCTTimestamps(cert)
+	if len(sctTimestamps) == 0 {
+		return
+	}
+
+	for i, sctTime := range sctTimestamps {
+		if sctTime.After(now) {
+			mgr.logger.Warnf("certificate for %q SCT[%d] timestamp is in the future: %v (by %v)",
+				domain, i, sctTime.UTC(), sctTime.Sub(now))
+		} else {
+			mgr.logger.Debugf("certificate for %q SCT[%d] timestamp: %v (%v in the past)",
+				domain, i, sctTime.UTC(), now.Sub(sctTime))
+		}
+	}
+}
+
+// parseSCTTimestamps extracts SCT timestamps from a certificate.
+func (mgr *Manager) parseSCTTimestamps(cert *x509.Certificate) []time.Time {
+	var timestamps []time.Time
+
+	for _, ext := range cert.Extensions {
+		if !ext.Id.Equal(oidSCTList) {
+			continue
+		}
+
+		// The extension value is an OCTET STRING containing the SCT list
+		var sctListBytes []byte
+		if _, err := asn1.Unmarshal(ext.Value, &sctListBytes); err != nil {
+			mgr.logger.Debugf("failed to unmarshal SCT list outer wrapper: %v", err)
+			continue
+		}
+
+		// SCT list format: 2-byte length prefix, then concatenated SCTs
+		if len(sctListBytes) < 2 {
+			continue
+		}
+
+		listLen := int(binary.BigEndian.Uint16(sctListBytes[:2]))
+		data := sctListBytes[2:]
+		if len(data) < listLen {
+			continue
+		}
+
+		// Parse individual SCTs
+		offset := 0
+		for offset < listLen {
+			if offset+2 > len(data) {
+				break
+			}
+			sctLen := int(binary.BigEndian.Uint16(data[offset : offset+2]))
+			offset += 2
+
+			if offset+sctLen > len(data) {
+				break
+			}
+			sctData := data[offset : offset+sctLen]
+			offset += sctLen
+
+			// SCT format: version (1) + log_id (32) + timestamp (8) + ...
+			if len(sctData) < 41 {
+				continue
+			}
+
+			// Timestamp is at offset 33 (after version + log_id), 8 bytes, milliseconds since epoch
+			tsMillis := binary.BigEndian.Uint64(sctData[33:41])
+			ts := time.UnixMilli(int64(tsMillis))
+			timestamps = append(timestamps, ts)
+		}
+	}
+
+	return timestamps
+}
+
+// dummyConn implements net.Conn to provide context for certificate fetching.
+type dummyConn struct {
+	ctx context.Context
+}
+
+func (c *dummyConn) Read(b []byte) (n int, err error)   { return 0, nil }
+func (c *dummyConn) Write(b []byte) (n int, err error)  { return len(b), nil }
+func (c *dummyConn) Close() error                       { return nil }
+func (c *dummyConn) LocalAddr() net.Addr                { return nil }
+func (c *dummyConn) RemoteAddr() net.Addr               { return nil }
+func (c *dummyConn) SetDeadline(t time.Time) error      { return nil }
+func (c *dummyConn) SetReadDeadline(t time.Time) error  { return nil }
+func (c *dummyConn) SetWriteDeadline(t time.Time) error { return nil }
+
+// RemoveDomain removes a domain from tracking.
+func (mgr *Manager) RemoveDomain(d domain.Domain) {
+	mgr.mu.Lock()
+	defer mgr.mu.Unlock()
+	delete(mgr.domains, d)
+}
+
+// PendingCerts returns the number of certificates currently being prefetched.
+func (mgr *Manager) PendingCerts() int {
+	mgr.mu.RLock()
+	defer mgr.mu.RUnlock()
+	var n int
+	for _, info := range mgr.domains {
+		if info.state == domainPending {
+			n++
+		}
+	}
+	return n
+}
+
+// TotalDomains returns the total number of registered domains.
+func (mgr *Manager) TotalDomains() int {
+	mgr.mu.RLock()
+	defer mgr.mu.RUnlock()
+	return len(mgr.domains)
+}
+
+// PendingDomains returns the domain names currently being prefetched.
+func (mgr *Manager) PendingDomains() []string {
+	return mgr.domainsByState(domainPending)
+}
+
+// ReadyDomains returns domain names that have successfully obtained certificates.
+func (mgr *Manager) ReadyDomains() []string {
+	return mgr.domainsByState(domainReady)
+}
+
+// FailedDomains returns domain names that failed certificate prefetch, mapped to their error.
+func (mgr *Manager) FailedDomains() map[string]string {
+	mgr.mu.RLock()
+	defer mgr.mu.RUnlock()
+	result := make(map[string]string)
+	for d, info := range mgr.domains {
+		if info.state == domainFailed {
+			result[d.PunycodeString()] = info.err
+		}
+	}
+	return result
+}
+
+func (mgr *Manager) domainsByState(state domainState) []string {
+	mgr.mu.RLock()
+	defer mgr.mu.RUnlock()
+	var domains []string
+	for d, info := range mgr.domains {
+		if info.state == state {
+			domains = append(domains, d.PunycodeString())
+		}
+	}
+	slices.Sort(domains)
+	return domains
+}

proxy/internal/acme/manager_test.go

@@ -0,0 +1,306 @@
+package acme
+
+import (
+	"context"
+	"crypto/ecdsa"
+	"crypto/elliptic"
+	"crypto/rand"
+	"crypto/tls"
+	"crypto/x509"
+	"crypto/x509/pkix"
+	"encoding/pem"
+	"math/big"
+	"os"
+	"path/filepath"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/proxy/internal/types"
+)
+
+func TestHostPolicy(t *testing.T) {
+	mgr, err := NewManager(ManagerConfig{CertDir: t.TempDir(), ACMEURL: "https://acme.example.com/directory"}, nil, nil, nil)
+	require.NoError(t, err)
+	mgr.AddDomain("example.com", types.AccountID("acc1"), types.ServiceID("rp1"))
+
+	// Wait for the background prefetch goroutine to finish so the temp dir
+	// can be cleaned up without a race.
+	t.Cleanup(func() {
+		assert.Eventually(t, func() bool {
+			return mgr.PendingCerts() == 0
+		}, 30*time.Second, 50*time.Millisecond)
+	})
+
+	tests := []struct {
+		name    string
+		host    string
+		wantErr bool
+	}{
+		{
+			name: "exact domain match",
+			host: "example.com",
+		},
+		{
+			name: "domain with port",
+			host: "example.com:443",
+		},
+		{
+			name:    "unknown domain",
+			host:    "unknown.com",
+			wantErr: true,
+		},
+		{
+			name:    "unknown domain with port",
+			host:    "unknown.com:443",
+			wantErr: true,
+		},
+		{
+			name:    "empty host",
+			host:    "",
+			wantErr: true,
+		},
+		{
+			name:    "port only",
+			host:    ":443",
+			wantErr: true,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			err := mgr.hostPolicy(context.Background(), tc.host)
+			if tc.wantErr {
+				require.Error(t, err)
+				assert.Contains(t, err.Error(), "unknown domain")
+			} else {
+				assert.NoError(t, err)
+			}
+		})
+	}
+}
+
+func TestDomainStates(t *testing.T) {
+	mgr, err := NewManager(ManagerConfig{CertDir: t.TempDir(), ACMEURL: "https://acme.example.com/directory"}, nil, nil, nil)
+	require.NoError(t, err)
+
+	assert.Equal(t, 0, mgr.PendingCerts(), "initially zero")
+	assert.Equal(t, 0, mgr.TotalDomains(), "initially zero domains")
+	assert.Empty(t, mgr.PendingDomains())
+	assert.Empty(t, mgr.ReadyDomains())
+	assert.Empty(t, mgr.FailedDomains())
+
+	// AddDomain starts as pending, then the prefetch goroutine will fail
+	// (no real ACME server) and transition to failed.
+	mgr.AddDomain("a.example.com", types.AccountID("acc1"), types.ServiceID("rp1"))
+	mgr.AddDomain("b.example.com", types.AccountID("acc1"), types.ServiceID("rp1"))
+
+	assert.Equal(t, 2, mgr.TotalDomains(), "two domains registered")
+
+	// Pending domains should eventually drain after prefetch goroutines finish.
+	assert.Eventually(t, func() bool {
+		return mgr.PendingCerts() == 0
+	}, 30*time.Second, 100*time.Millisecond, "pending certs should return to zero after prefetch completes")
+
+	assert.Empty(t, mgr.PendingDomains())
+	assert.Equal(t, 2, mgr.TotalDomains(), "total domains unchanged")
+
+	// With a fake ACME URL, both should have failed.
+	failed := mgr.FailedDomains()
+	assert.Len(t, failed, 2, "both domains should have failed")
+	assert.Contains(t, failed, "a.example.com")
+	assert.Contains(t, failed, "b.example.com")
+	assert.Empty(t, mgr.ReadyDomains())
+}
+
+func TestParseWildcard(t *testing.T) {
+	tests := []struct {
+		pattern    string
+		wantSuffix string
+		wantOK     bool
+	}{
+		{"*.example.com", ".example.com", true},
+		{"*.foo.example.com", ".foo.example.com", true},
+		{"*.COM", ".com", true},       // single-label TLD
+		{"example.com", "", false},    // no wildcard prefix
+		{"*example.com", "", false},   // missing dot
+		{"**.example.com", "", false}, // double star
+		{"", "", false},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.pattern, func(t *testing.T) {
+			suffix, ok := parseWildcard(tc.pattern)
+			assert.Equal(t, tc.wantOK, ok)
+			if ok {
+				assert.Equal(t, tc.wantSuffix, suffix)
+			}
+		})
+	}
+}
+
+func TestMatchesWildcard(t *testing.T) {
+	wcDir := t.TempDir()
+	generateSelfSignedCert(t, wcDir, "example", "*.example.com")
+
+	acmeDir := t.TempDir()
+	mgr, err := NewManager(ManagerConfig{CertDir: acmeDir, ACMEURL: "https://acme.example.com/directory", WildcardDir: wcDir}, nil, nil, nil)
+	require.NoError(t, err)
+
+	tests := []struct {
+		host  string
+		match bool
+	}{
+		{"foo.example.com", true},
+		{"bar.example.com", true},
+		{"FOO.Example.COM", true},      // case insensitive
+		{"example.com", false},         // bare parent
+		{"sub.foo.example.com", false}, // multi-level
+		{"notexample.com", false},
+		{"", false},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.host, func(t *testing.T) {
+			assert.Equal(t, tc.match, mgr.findWildcardEntry(tc.host) != nil)
+		})
+	}
+}
+
+// generateSelfSignedCert creates a temporary self-signed certificate and key
+// for testing purposes. The baseName controls the output filenames:
+// <baseName>.crt and <baseName>.key.
+func generateSelfSignedCert(t *testing.T, dir, baseName string, dnsNames ...string) {
+	t.Helper()
+
+	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
+	require.NoError(t, err)
+
+	template := &x509.Certificate{
+		SerialNumber: big.NewInt(1),
+		Subject:      pkix.Name{CommonName: dnsNames[0]},
+		DNSNames:     dnsNames,
+		NotBefore:    time.Now().Add(-time.Hour),
+		NotAfter:     time.Now().Add(24 * time.Hour),
+	}
+
+	certDER, err := x509.CreateCertificate(rand.Reader, template, template, &key.PublicKey, key)
+	require.NoError(t, err)
+
+	certFile, err := os.Create(filepath.Join(dir, baseName+".crt"))
+	require.NoError(t, err)
+	require.NoError(t, pem.Encode(certFile, &pem.Block{Type: "CERTIFICATE", Bytes: certDER}))
+	require.NoError(t, certFile.Close())
+
+	keyDER, err := x509.MarshalECPrivateKey(key)
+	require.NoError(t, err)
+	keyFile, err := os.Create(filepath.Join(dir, baseName+".key"))
+	require.NoError(t, err)
+	require.NoError(t, pem.Encode(keyFile, &pem.Block{Type: "EC PRIVATE KEY", Bytes: keyDER}))
+	require.NoError(t, keyFile.Close())
+}
+
+func TestWildcardAddDomainSkipsACME(t *testing.T) {
+	wcDir := t.TempDir()
+	generateSelfSignedCert(t, wcDir, "example", "*.example.com")
+
+	acmeDir := t.TempDir()
+	mgr, err := NewManager(ManagerConfig{CertDir: acmeDir, ACMEURL: "https://acme.example.com/directory", WildcardDir: wcDir}, nil, nil, nil)
+	require.NoError(t, err)
+
+	// Add a wildcard-matching domain — should be immediately ready.
+	mgr.AddDomain("foo.example.com", types.AccountID("acc1"), types.ServiceID("svc1"))
+	assert.Equal(t, 0, mgr.PendingCerts(), "wildcard domain should not be pending")
+	assert.Equal(t, []string{"foo.example.com"}, mgr.ReadyDomains())
+
+	// Add a non-wildcard domain — should go through ACME (pending then failed).
+	mgr.AddDomain("other.net", types.AccountID("acc2"), types.ServiceID("svc2"))
+	assert.Equal(t, 2, mgr.TotalDomains())
+
+	// Wait for the ACME prefetch to fail.
+	assert.Eventually(t, func() bool {
+		return mgr.PendingCerts() == 0
+	}, 30*time.Second, 100*time.Millisecond)
+
+	assert.Equal(t, []string{"foo.example.com"}, mgr.ReadyDomains())
+	assert.Contains(t, mgr.FailedDomains(), "other.net")
+}
+
+func TestWildcardGetCertificate(t *testing.T) {
+	wcDir := t.TempDir()
+	generateSelfSignedCert(t, wcDir, "example", "*.example.com")
+
+	acmeDir := t.TempDir()
+	mgr, err := NewManager(ManagerConfig{CertDir: acmeDir, ACMEURL: "https://acme.example.com/directory", WildcardDir: wcDir}, nil, nil, nil)
+	require.NoError(t, err)
+
+	mgr.AddDomain("foo.example.com", types.AccountID("acc1"), types.ServiceID("svc1"))
+
+	// GetCertificate for a wildcard-matching domain should return the static cert.
+	cert, err := mgr.GetCertificate(&tls.ClientHelloInfo{ServerName: "foo.example.com"})
+	require.NoError(t, err)
+	require.NotNil(t, cert)
+	assert.Contains(t, cert.Leaf.DNSNames, "*.example.com")
+}
+
+func TestMultipleWildcards(t *testing.T) {
+	wcDir := t.TempDir()
+	generateSelfSignedCert(t, wcDir, "example", "*.example.com")
+	generateSelfSignedCert(t, wcDir, "other", "*.other.org")
+
+	acmeDir := t.TempDir()
+	mgr, err := NewManager(ManagerConfig{CertDir: acmeDir, ACMEURL: "https://acme.example.com/directory", WildcardDir: wcDir}, nil, nil, nil)
+	require.NoError(t, err)
+
+	assert.ElementsMatch(t, []string{"*.example.com", "*.other.org"}, mgr.WildcardPatterns())
+
+	// Both wildcards should resolve.
+	mgr.AddDomain("foo.example.com", types.AccountID("acc1"), types.ServiceID("svc1"))
+	mgr.AddDomain("bar.other.org", types.AccountID("acc2"), types.ServiceID("svc2"))
+
+	assert.Equal(t, 0, mgr.PendingCerts())
+	assert.ElementsMatch(t, []string{"foo.example.com", "bar.other.org"}, mgr.ReadyDomains())
+
+	// GetCertificate routes to the correct cert.
+	cert1, err := mgr.GetCertificate(&tls.ClientHelloInfo{ServerName: "foo.example.com"})
+	require.NoError(t, err)
+	assert.Contains(t, cert1.Leaf.DNSNames, "*.example.com")
+
+	cert2, err := mgr.GetCertificate(&tls.ClientHelloInfo{ServerName: "bar.other.org"})
+	require.NoError(t, err)
+	assert.Contains(t, cert2.Leaf.DNSNames, "*.other.org")
+
+	// Non-matching domain falls through to ACME.
+	mgr.AddDomain("custom.net", types.AccountID("acc3"), types.ServiceID("svc3"))
+	assert.Eventually(t, func() bool {
+		return mgr.PendingCerts() == 0
+	}, 30*time.Second, 100*time.Millisecond)
+	assert.Contains(t, mgr.FailedDomains(), "custom.net")
+}
+
+func TestWildcardDirEmpty(t *testing.T) {
+	wcDir := t.TempDir()
+	// Empty directory — no .crt files.
+	_, err := NewManager(ManagerConfig{CertDir: t.TempDir(), ACMEURL: "https://acme.example.com/directory", WildcardDir: wcDir}, nil, nil, nil)
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "no .crt files found")
+}
+
+func TestWildcardDirNonWildcardCert(t *testing.T) {
+	wcDir := t.TempDir()
+	// Certificate without a wildcard SAN.
+	generateSelfSignedCert(t, wcDir, "plain", "plain.example.com")
+
+	_, err := NewManager(ManagerConfig{CertDir: t.TempDir(), ACMEURL: "https://acme.example.com/directory", WildcardDir: wcDir}, nil, nil, nil)
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "no wildcard SANs")
+}
+
+func TestNoWildcardDir(t *testing.T) {
+	// Empty string means no wildcard dir — pure ACME mode.
+	mgr, err := NewManager(ManagerConfig{CertDir: t.TempDir(), ACMEURL: "https://acme.example.com/directory"}, nil, nil, nil)
+	require.NoError(t, err)
+	assert.Empty(t, mgr.WildcardPatterns())
+}

proxy/internal/auth/auth.gohtml

@@ -0,0 +1,18 @@
+<!doctype html>
+{{ range $method, $value := .Methods }}
+{{ if eq $method "pin" }}
+<form>
+    <label for={{ $value }}>PIN:</label>
+    <input name={{ $value }} id={{ $value }} />
+    <button type=submit>Submit</button>
+</form>
+{{ else if eq $method "password" }}
+<form>
+    <label for={{ $value }}>Password:</label>
+    <input name={{ $value }} id={{ $value }}/>
+    <button type=submit>Submit</button>
+</form>
+{{ else if eq $method "oidc" }}
+<a href={{ $value }}>Click here to log in with SSO</a>
+{{ end }}
+{{ end }}

proxy/internal/auth/header.go

@@ -0,0 +1,69 @@
+package auth
+
+import (
+	"errors"
+	"fmt"
+	"net/http"
+
+	"github.com/netbirdio/netbird/proxy/auth"
+	"github.com/netbirdio/netbird/proxy/internal/types"
+	"github.com/netbirdio/netbird/shared/management/proto"
+)
+
+// ErrHeaderAuthFailed indicates that the header was present but the
+// credential did not validate. Callers should return 401 instead of
+// falling through to other auth schemes.
+var ErrHeaderAuthFailed = errors.New("header authentication failed")
+
+// Header implements header-based authentication. The proxy checks for the
+// configured header in each request and validates its value via gRPC.
+type Header struct {
+	id         types.ServiceID
+	accountId  types.AccountID
+	headerName string
+	client     authenticator
+}
+
+// NewHeader creates a Header authentication scheme for the given header name.
+func NewHeader(client authenticator, id types.ServiceID, accountId types.AccountID, headerName string) Header {
+	return Header{
+		id:         id,
+		accountId:  accountId,
+		headerName: headerName,
+		client:     client,
+	}
+}
+
+// Type returns auth.MethodHeader.
+func (Header) Type() auth.Method {
+	return auth.MethodHeader
+}
+
+// Authenticate checks for the configured header in the request. If absent,
+// returns empty (unauthenticated). If present, validates via gRPC.
+func (h Header) Authenticate(r *http.Request) (string, string, error) {
+	value := r.Header.Get(h.headerName)
+	if value == "" {
+		return "", "", nil
+	}
+
+	res, err := h.client.Authenticate(r.Context(), &proto.AuthenticateRequest{
+		Id:        string(h.id),
+		AccountId: string(h.accountId),
+		Request: &proto.AuthenticateRequest_HeaderAuth{
+			HeaderAuth: &proto.HeaderAuthRequest{
+				HeaderValue: value,
+				HeaderName:  h.headerName,
+			},
+		},
+	})
+	if err != nil {
+		return "", "", fmt.Errorf("authenticate header: %w", err)
+	}
+
+	if res.GetSuccess() {
+		return res.GetSessionToken(), "", nil
+	}
+
+	return "", "", ErrHeaderAuthFailed
+}

proxy/internal/auth/identity.go

@@ -0,0 +1,47 @@
+package auth
+
+import (
+	"context"
+	"net/netip"
+)
+
+// PeerIdentity describes the locally-known facts about a peer reachable on
+// the proxy's per-account WireGuard listener. Phase 3 fills PubKey, TunnelIP
+// and FQDN from the embedded client's peerstore. UserID, Email and Groups
+// stay zero in V1 — full identity still travels through ValidateTunnelPeer.
+// Phase V2 will populate them once RemotePeerConfig carries user identity.
+type PeerIdentity struct {
+	PubKey   string
+	TunnelIP netip.Addr
+	FQDN     string
+
+	// V2 fields (zero in V1).
+	UserID string
+	Email  string
+	Groups []string
+}
+
+// TunnelLookupFunc resolves a tunnel IP to a peer identity using locally
+// available peerstore data. ok=false means the IP is not in the calling
+// account's roster.
+type TunnelLookupFunc func(ip netip.Addr) (PeerIdentity, bool)
+
+type tunnelLookupContextKey struct{}
+
+// WithTunnelLookup attaches a per-account peerstore lookup function to
+// the request context. The auth middleware calls this lookup before
+// hitting management's ValidateTunnelPeer to short-circuit unknown IPs
+// and to skip the RPC for already-cached identities.
+func WithTunnelLookup(ctx context.Context, lookup TunnelLookupFunc) context.Context {
+	if lookup == nil {
+		return ctx
+	}
+	return context.WithValue(ctx, tunnelLookupContextKey{}, lookup)
+}
+
+// TunnelLookupFromContext returns the peerstore lookup attached to ctx,
+// or nil when the request did not arrive on a per-account listener.
+func TunnelLookupFromContext(ctx context.Context) TunnelLookupFunc {
+	v, _ := ctx.Value(tunnelLookupContextKey{}).(TunnelLookupFunc)
+	return v
+}

proxy/internal/auth/middleware.go

@@ -0,0 +1,740 @@
+package auth
+
+import (
+	"context"
+	"crypto/ed25519"
+	"encoding/base64"
+	"errors"
+	"fmt"
+	"html"
+	"net"
+	"net/http"
+	"net/netip"
+	"net/url"
+	"sync"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+	"google.golang.org/grpc"
+
+	"github.com/netbirdio/netbird/proxy/auth"
+	"github.com/netbirdio/netbird/proxy/internal/proxy"
+	"github.com/netbirdio/netbird/proxy/internal/restrict"
+	"github.com/netbirdio/netbird/proxy/internal/types"
+	"github.com/netbirdio/netbird/proxy/web"
+	"github.com/netbirdio/netbird/shared/management/proto"
+)
+
+// errValidationUnavailable indicates that session validation failed due to
+// an infrastructure error (e.g. gRPC unavailable), not an invalid token.
+var errValidationUnavailable = errors.New("session validation unavailable")
+
+type authenticator interface {
+	Authenticate(ctx context.Context, in *proto.AuthenticateRequest, opts ...grpc.CallOption) (*proto.AuthenticateResponse, error)
+}
+
+// SessionValidator validates session tokens and checks user access permissions.
+type SessionValidator interface {
+	ValidateSession(ctx context.Context, in *proto.ValidateSessionRequest, opts ...grpc.CallOption) (*proto.ValidateSessionResponse, error)
+	ValidateTunnelPeer(ctx context.Context, in *proto.ValidateTunnelPeerRequest, opts ...grpc.CallOption) (*proto.ValidateTunnelPeerResponse, error)
+}
+
+// Scheme defines an authentication mechanism for a domain.
+type Scheme interface {
+	Type() auth.Method
+	// Authenticate checks the request and determines whether it represents
+	// an authenticated user. An empty token indicates an unauthenticated
+	// request; optionally, promptData may be returned for the login UI.
+	// An error indicates an infrastructure failure (e.g. gRPC unavailable).
+	Authenticate(*http.Request) (token string, promptData string, err error)
+}
+
+// DomainConfig holds the authentication and restriction settings for a protected domain.
+type DomainConfig struct {
+	Schemes           []Scheme
+	SessionPublicKey  ed25519.PublicKey
+	SessionExpiration time.Duration
+	AccountID         types.AccountID
+	ServiceID         types.ServiceID
+	IPRestrictions    *restrict.Filter
+	// Private routes the domain through ValidateTunnelPeer; failure → 403.
+	Private bool
+}
+
+type validationResult struct {
+	UserID       string
+	UserEmail    string
+	Valid        bool
+	DeniedReason string
+	Groups       []string
+	// GroupNames carries the human-readable display names for Groups,
+	// ordered identically (positional pairing). May be shorter than
+	// Groups for tokens minted before names were embedded; the consumer
+	// falls back to ids for missing positions.
+	GroupNames []string
+}
+
+// Middleware applies per-domain authentication and IP restriction checks.
+type Middleware struct {
+	domainsMux       sync.RWMutex
+	domains          map[string]DomainConfig
+	logger           *log.Logger
+	sessionValidator SessionValidator
+	geo              restrict.GeoResolver
+	tunnelCache      *tunnelValidationCache
+}
+
+// NewMiddleware creates a new authentication middleware. The sessionValidator is
+// optional; if nil, OIDC session tokens are validated locally without group access checks.
+func NewMiddleware(logger *log.Logger, sessionValidator SessionValidator, geo restrict.GeoResolver) *Middleware {
+	if logger == nil {
+		logger = log.StandardLogger()
+	}
+	return &Middleware{
+		domains:          make(map[string]DomainConfig),
+		logger:           logger,
+		sessionValidator: sessionValidator,
+		geo:              geo,
+		tunnelCache:      newTunnelValidationCache(),
+	}
+}
+
+// Protect wraps next with per-domain authentication and IP restriction checks.
+// Requests whose Host is not registered pass through unchanged.
+func (mw *Middleware) Protect(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		host, _, err := net.SplitHostPort(r.Host)
+		if err != nil {
+			host = r.Host
+		}
+
+		config, exists := mw.getDomainConfig(host)
+		mw.logger.Debugf("checking authentication for host: %s, exists: %t", host, exists)
+
+		if !exists {
+			next.ServeHTTP(w, r)
+			return
+		}
+
+		// Set account and service IDs in captured data for access logging.
+		setCapturedIDs(r, config)
+
+		if !mw.checkIPRestrictions(w, r, config) {
+			return
+		}
+
+		// Private services bypass operator schemes and gate on tunnel peer.
+		if config.Private {
+			if mw.forwardWithTunnelPeer(w, r, host, config, next) {
+				return
+			}
+			http.Error(w, "Forbidden", http.StatusForbidden)
+			return
+		}
+
+		// Domains with no authentication schemes pass through after IP checks.
+		if len(config.Schemes) == 0 {
+			next.ServeHTTP(w, r)
+			return
+		}
+
+		if mw.handleOAuthCallbackError(w, r) {
+			return
+		}
+
+		if mw.forwardWithSessionCookie(w, r, host, config, next) {
+			return
+		}
+
+		if mw.forwardWithHeaderAuth(w, r, host, config, next) {
+			return
+		}
+
+		if mw.forwardWithTunnelPeer(w, r, host, config, next) {
+			return
+		}
+
+		if mw.blockOIDCOnPlainHTTP(w, r, config) {
+			return
+		}
+
+		mw.authenticateWithSchemes(w, r, host, config)
+	})
+}
+
+// requestIsPlainHTTP reports whether the request arrived without TLS.
+// Used to gate cookie-on-plain warnings and the OIDC plain-HTTP block.
+func requestIsPlainHTTP(r *http.Request) bool {
+	return r.TLS == nil
+}
+
+// hasOIDCScheme reports whether any of the configured schemes requires
+// TLS to round-trip safely with an external IdP.
+func hasOIDCScheme(schemes []Scheme) bool {
+	for _, s := range schemes {
+		if s.Type() == auth.MethodOIDC {
+			return true
+		}
+	}
+	return false
+}
+
+// blockOIDCOnPlainHTTP fails fast when an OIDC-configured domain is hit
+// over plain HTTP. Most IdPs reject http:// redirect URIs, so surfacing
+// the misconfiguration here yields a clearer error than the IdP's
+// "invalid redirect_uri" round-trip.
+func (mw *Middleware) blockOIDCOnPlainHTTP(w http.ResponseWriter, r *http.Request, config DomainConfig) bool {
+	if !requestIsPlainHTTP(r) {
+		return false
+	}
+	if !hasOIDCScheme(config.Schemes) {
+		return false
+	}
+	mw.logger.WithFields(log.Fields{
+		"host":   r.Host,
+		"remote": r.RemoteAddr,
+	}).Warn("OIDC scheme reached on plain HTTP path; rejecting with 400 — use port 443")
+	http.Error(w, "OIDC requires TLS — use port 443", http.StatusBadRequest)
+	return true
+}
+
+func (mw *Middleware) getDomainConfig(host string) (DomainConfig, bool) {
+	mw.domainsMux.RLock()
+	defer mw.domainsMux.RUnlock()
+	config, exists := mw.domains[host]
+	return config, exists
+}
+
+func setCapturedIDs(r *http.Request, config DomainConfig) {
+	if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
+		cd.SetAccountID(config.AccountID)
+		cd.SetServiceID(config.ServiceID)
+	}
+}
+
+// checkIPRestrictions validates the client IP against the domain's IP restrictions.
+// Uses the resolved client IP from CapturedData (which accounts for trusted proxies)
+// rather than r.RemoteAddr directly.
+func (mw *Middleware) checkIPRestrictions(w http.ResponseWriter, r *http.Request, config DomainConfig) bool {
+	if config.IPRestrictions == nil {
+		return true
+	}
+
+	clientIP := mw.resolveClientIP(r)
+	if !clientIP.IsValid() {
+		mw.logger.Debugf("IP restriction: cannot resolve client address for %q, denying", r.RemoteAddr)
+		http.Error(w, "Forbidden", http.StatusForbidden)
+		return false
+	}
+
+	var verdict restrict.Verdict
+	if types.IsOverlayOrigin(r.Context()) {
+		// Geo/CrowdSec checks don't apply over the WireGuard overlay:
+		// the source address is always inside the NetBird CGNAT range,
+		// which is never in a GeoIP database or a CrowdSec decision
+		// list. Enforcing them here would either no-op (best case) or
+		// fail-closed when the geo database is missing.
+		verdict = config.IPRestrictions.CheckCIDR(clientIP)
+	} else {
+		verdict = config.IPRestrictions.Check(clientIP, mw.geo)
+	}
+	if verdict == restrict.Allow {
+		return true
+	}
+
+	if verdict.IsCrowdSec() {
+		if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
+			cd.SetMetadata("crowdsec_verdict", verdict.String())
+			if config.IPRestrictions.IsObserveOnly(verdict) {
+				cd.SetMetadata("crowdsec_mode", "observe")
+			}
+		}
+	}
+
+	if config.IPRestrictions.IsObserveOnly(verdict) {
+		mw.logger.Debugf("CrowdSec observe: would block %s for %s (%s)", clientIP, r.Host, verdict)
+		return true
+	}
+
+	reason := verdict.String()
+	mw.blockIPRestriction(r, reason)
+	http.Error(w, "Forbidden", http.StatusForbidden)
+	return false
+}
+
+// resolveClientIP extracts the real client IP from CapturedData, falling back to r.RemoteAddr.
+func (mw *Middleware) resolveClientIP(r *http.Request) netip.Addr {
+	if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
+		if ip := cd.GetClientIP(); ip.IsValid() {
+			return ip
+		}
+	}
+
+	clientIPStr, _, _ := net.SplitHostPort(r.RemoteAddr)
+	if clientIPStr == "" {
+		clientIPStr = r.RemoteAddr
+	}
+	addr, err := netip.ParseAddr(clientIPStr)
+	if err != nil {
+		return netip.Addr{}
+	}
+	return addr.Unmap()
+}
+
+// blockIPRestriction sets captured data fields for an IP-restriction block event.
+func (mw *Middleware) blockIPRestriction(r *http.Request, reason string) {
+	if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
+		cd.SetOrigin(proxy.OriginAuth)
+		cd.SetAuthMethod(reason)
+	}
+	mw.logger.Debugf("IP restriction: %s for %s", reason, r.RemoteAddr)
+}
+
+// handleOAuthCallbackError checks for error query parameters from an OAuth
+// callback and renders the access denied page if present.
+func (mw *Middleware) handleOAuthCallbackError(w http.ResponseWriter, r *http.Request) bool {
+	errCode := r.URL.Query().Get("error")
+	if errCode == "" {
+		return false
+	}
+
+	var requestID string
+	if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
+		cd.SetOrigin(proxy.OriginAuth)
+		cd.SetAuthMethod(auth.MethodOIDC.String())
+		requestID = cd.GetRequestID()
+	}
+	errDesc := r.URL.Query().Get("error_description")
+	if errDesc == "" {
+		errDesc = "An error occurred during authentication"
+	} else {
+		errDesc = html.EscapeString(errDesc)
+	}
+	web.ServeAccessDeniedPage(w, r, http.StatusForbidden, "Access Denied", errDesc, requestID)
+	return true
+}
+
+// forwardWithSessionCookie checks for a valid session cookie and, if found,
+// sets the user identity on the request context and forwards to the next handler.
+func (mw *Middleware) forwardWithSessionCookie(w http.ResponseWriter, r *http.Request, host string, config DomainConfig, next http.Handler) bool {
+	cookie, err := r.Cookie(auth.SessionCookieName)
+	if err != nil {
+		return false
+	}
+	userID, email, method, groups, groupNames, err := auth.ValidateSessionJWT(cookie.Value, host, config.SessionPublicKey)
+	if err != nil {
+		return false
+	}
+	if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
+		cd.SetUserID(userID)
+		cd.SetUserEmail(email)
+		cd.SetUserGroups(groups)
+		cd.SetUserGroupNames(groupNames)
+		cd.SetAuthMethod(method)
+	}
+	next.ServeHTTP(w, r)
+	return true
+}
+
+// forwardWithTunnelPeer is the OIDC fast-path for requests originating on the
+// netbird mesh. When the source IP belongs to a private/CGNAT range the proxy
+// asks management to resolve it to a peer/user and to gate by the service's
+// distribution_groups. On success the proxy installs the freshly minted JWT
+// as a session cookie, sets UserID + Method=oidc on the captured data, and
+// forwards directly — operators see the same access-log shape as if the user
+// had completed an OIDC redirect. Any failure (private-range mismatch,
+// management unreachable, peer unknown, user not in group) returns false so
+// the caller falls back to the existing OIDC scheme dispatch.
+//
+// The fast-path is gated on TunnelLookupFromContext(r.Context()) being
+// present — that context value is attached only by the per-account
+// inbound (overlay) listener. The host listener never sets it, so a
+// public client whose source IP happens to fall inside an RFC1918 / ULA
+// / CGNAT range can't impersonate a mesh peer by colliding with a
+// tunnel-IP. Once we know the request arrived over WireGuard the
+// per-account peerstore lookup is consulted: a miss denies fast (no
+// management round-trip), a hit gates the cached ValidateTunnelPeer RPC
+// that mints the session JWT.
+func (mw *Middleware) forwardWithTunnelPeer(w http.ResponseWriter, r *http.Request, host string, config DomainConfig, next http.Handler) bool {
+	if mw.sessionValidator == nil {
+		return false
+	}
+	clientIP := mw.resolveClientIP(r)
+	if !clientIP.IsValid() {
+		return false
+	}
+
+	// Anti-spoof: only honour the tunnel-peer fast-path on requests that
+	// were stamped by an overlay listener. Without that marker an
+	// attacker could send a request from a colliding RFC1918 / CGNAT
+	// source on the public listener and bypass operator auth.
+	lookup := TunnelLookupFromContext(r.Context())
+	if lookup == nil {
+		return false
+	}
+	if !isTunnelSourceIP(clientIP) {
+		return false
+	}
+	if _, ok := lookup(clientIP); !ok {
+		mw.logger.WithFields(log.Fields{
+			"host":   host,
+			"remote": clientIP,
+		}).Debug("local peerstore: tunnel IP not in account roster; denying without RPC")
+		return false
+	}
+
+	resp, _, err := mw.tunnelCache.fetch(r.Context(), tunnelCacheKey{
+		accountID: config.AccountID,
+		tunnelIP:  clientIP,
+		domain:    host,
+	}, mw.validateTunnelPeer)
+	if err != nil {
+		mw.logger.WithError(err).Debug("ValidateTunnelPeer failed; falling back to OIDC")
+		return false
+	}
+	if !resp.GetValid() || resp.GetSessionToken() == "" {
+		return false
+	}
+
+	setSessionCookie(w, resp.GetSessionToken(), config.SessionExpiration)
+	if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
+		cd.SetOrigin(proxy.OriginAuth)
+		cd.SetUserID(resp.GetUserId())
+		cd.SetUserEmail(resp.GetUserEmail())
+		cd.SetUserGroups(resp.GetPeerGroupIds())
+		cd.SetUserGroupNames(resp.GetPeerGroupNames())
+		cd.SetAuthMethod(auth.MethodOIDC.String())
+	}
+	next.ServeHTTP(w, r)
+	return true
+}
+
+// validateTunnelPeer adapts the SessionValidator interface to the cache's
+// validateTunnelPeerFn signature.
+func (mw *Middleware) validateTunnelPeer(ctx context.Context, req *proto.ValidateTunnelPeerRequest) (*proto.ValidateTunnelPeerResponse, error) {
+	return mw.sessionValidator.ValidateTunnelPeer(ctx, req)
+}
+
+// cgnatPrefix covers RFC 6598 100.64.0.0/10, the CGNAT block NetBird
+// allocates tunnel addresses from by default. IsPrivate() doesn't include
+// it, so we check it explicitly.
+var cgnatPrefix = netip.MustParsePrefix("100.64.0.0/10")
+
+// isTunnelSourceIP reports whether ip falls within an address range typical
+// of NetBird tunnels: RFC1918 private space, IPv6 ULA, or CGNAT 100.64/10
+// (NetBird's default range). Loopback and link-local are excluded — the
+// fast-path is meant for peer-to-peer mesh traffic, not localhost.
+func isTunnelSourceIP(ip netip.Addr) bool {
+	if !ip.IsValid() || ip.IsLoopback() || ip.IsLinkLocalUnicast() {
+		return false
+	}
+	if ip.IsPrivate() {
+		return true
+	}
+	return cgnatPrefix.Contains(ip)
+}
+
+// forwardWithHeaderAuth checks for a Header auth scheme. If the header validates,
+// the request is forwarded directly (no redirect), which is important for API clients.
+func (mw *Middleware) forwardWithHeaderAuth(w http.ResponseWriter, r *http.Request, host string, config DomainConfig, next http.Handler) bool {
+	for _, scheme := range config.Schemes {
+		hdr, ok := scheme.(Header)
+		if !ok {
+			continue
+		}
+
+		handled := mw.tryHeaderScheme(w, r, host, config, hdr, next)
+		if handled {
+			return true
+		}
+	}
+	return false
+}
+
+func (mw *Middleware) tryHeaderScheme(w http.ResponseWriter, r *http.Request, host string, config DomainConfig, hdr Header, next http.Handler) bool {
+	token, _, err := hdr.Authenticate(r)
+	if err != nil {
+		return mw.handleHeaderAuthError(w, r, err)
+	}
+	if token == "" {
+		return false
+	}
+
+	result, err := mw.validateSessionToken(r.Context(), host, token, config.SessionPublicKey, auth.MethodHeader)
+	if err != nil {
+		setHeaderCapturedData(r.Context(), "", "", nil, nil)
+		status := http.StatusBadRequest
+		msg := "invalid session token"
+		if errors.Is(err, errValidationUnavailable) {
+			status = http.StatusBadGateway
+			msg = "authentication service unavailable"
+		}
+		http.Error(w, msg, status)
+		return true
+	}
+
+	if !result.Valid {
+		setHeaderCapturedData(r.Context(), result.UserID, result.UserEmail, result.Groups, result.GroupNames)
+		http.Error(w, "Unauthorized", http.StatusUnauthorized)
+		return true
+	}
+
+	setSessionCookie(w, token, config.SessionExpiration)
+	if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
+		cd.SetUserID(result.UserID)
+		cd.SetUserEmail(result.UserEmail)
+		cd.SetUserGroups(result.Groups)
+		cd.SetUserGroupNames(result.GroupNames)
+		cd.SetAuthMethod(auth.MethodHeader.String())
+	}
+
+	next.ServeHTTP(w, r)
+	return true
+}
+
+func (mw *Middleware) handleHeaderAuthError(w http.ResponseWriter, r *http.Request, err error) bool {
+	if errors.Is(err, ErrHeaderAuthFailed) {
+		setHeaderCapturedData(r.Context(), "", "", nil, nil)
+		http.Error(w, "Unauthorized", http.StatusUnauthorized)
+		return true
+	}
+	mw.logger.WithField("scheme", "header").Warnf("header auth infrastructure error: %v", err)
+	if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
+		cd.SetOrigin(proxy.OriginAuth)
+	}
+	http.Error(w, "authentication service unavailable", http.StatusBadGateway)
+	return true
+}
+
+func setHeaderCapturedData(ctx context.Context, userID, userEmail string, groups, groupNames []string) {
+	cd := proxy.CapturedDataFromContext(ctx)
+	if cd == nil {
+		return
+	}
+	cd.SetOrigin(proxy.OriginAuth)
+	cd.SetAuthMethod(auth.MethodHeader.String())
+	cd.SetUserID(userID)
+	cd.SetUserEmail(userEmail)
+	cd.SetUserGroups(groups)
+	cd.SetUserGroupNames(groupNames)
+}
+
+// authenticateWithSchemes tries each configured auth scheme in order.
+// On success it sets a session cookie and redirects; on failure it renders the login page.
+func (mw *Middleware) authenticateWithSchemes(w http.ResponseWriter, r *http.Request, host string, config DomainConfig) {
+	methods := make(map[string]string)
+	var attemptedMethod string
+
+	for _, scheme := range config.Schemes {
+		token, promptData, err := scheme.Authenticate(r)
+		if err != nil {
+			mw.logger.WithField("scheme", scheme.Type().String()).Warnf("authentication infrastructure error: %v", err)
+			if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
+				cd.SetOrigin(proxy.OriginAuth)
+			}
+			http.Error(w, "authentication service unavailable", http.StatusBadGateway)
+			return
+		}
+
+		// Track if credentials were submitted but auth failed
+		if token == "" && wasCredentialSubmitted(r, scheme.Type()) {
+			attemptedMethod = scheme.Type().String()
+		}
+
+		if token != "" {
+			mw.handleAuthenticatedToken(w, r, host, token, config, scheme)
+			return
+		}
+		methods[scheme.Type().String()] = promptData
+	}
+
+	if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
+		cd.SetOrigin(proxy.OriginAuth)
+		if attemptedMethod != "" {
+			cd.SetAuthMethod(attemptedMethod)
+		}
+	}
+
+	if oidcURL, ok := methods[auth.MethodOIDC.String()]; ok && len(methods) == 1 && oidcURL != "" {
+		http.Redirect(w, r, oidcURL, http.StatusFound)
+		return
+	}
+
+	web.ServeHTTP(w, r, map[string]any{"methods": methods}, http.StatusUnauthorized)
+}
+
+// handleAuthenticatedToken validates the token, handles denied access, and on
+// success sets a session cookie and redirects to the original URL.
+func (mw *Middleware) handleAuthenticatedToken(w http.ResponseWriter, r *http.Request, host, token string, config DomainConfig, scheme Scheme) {
+	result, err := mw.validateSessionToken(r.Context(), host, token, config.SessionPublicKey, scheme.Type())
+	if err != nil {
+		if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
+			cd.SetOrigin(proxy.OriginAuth)
+			cd.SetAuthMethod(scheme.Type().String())
+		}
+		status := http.StatusBadRequest
+		msg := "invalid session token"
+		if errors.Is(err, errValidationUnavailable) {
+			status = http.StatusBadGateway
+			msg = "authentication service unavailable"
+		}
+		http.Error(w, msg, status)
+		return
+	}
+
+	if !result.Valid {
+		var requestID string
+		if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
+			cd.SetOrigin(proxy.OriginAuth)
+			cd.SetUserID(result.UserID)
+			cd.SetUserEmail(result.UserEmail)
+			cd.SetUserGroups(result.Groups)
+			cd.SetUserGroupNames(result.GroupNames)
+			cd.SetAuthMethod(scheme.Type().String())
+			requestID = cd.GetRequestID()
+		}
+		web.ServeAccessDeniedPage(w, r, http.StatusForbidden, "Access Denied", "You are not authorized to access this service", requestID)
+		return
+	}
+
+	setSessionCookie(w, token, config.SessionExpiration)
+
+	// Redirect instead of forwarding the auth POST to the backend.
+	// The browser will follow with a GET carrying the new session cookie.
+	if cd := proxy.CapturedDataFromContext(r.Context()); cd != nil {
+		cd.SetOrigin(proxy.OriginAuth)
+		cd.SetUserID(result.UserID)
+		cd.SetUserEmail(result.UserEmail)
+		cd.SetUserGroups(result.Groups)
+		cd.SetUserGroupNames(result.GroupNames)
+		cd.SetAuthMethod(scheme.Type().String())
+	}
+	redirectURL := stripSessionTokenParam(r.URL)
+	http.Redirect(w, r, redirectURL, http.StatusSeeOther)
+}
+
+// setSessionCookie writes a session cookie with secure defaults.
+func setSessionCookie(w http.ResponseWriter, token string, expiration time.Duration) {
+	if expiration == 0 {
+		expiration = auth.DefaultSessionExpiry
+	}
+	http.SetCookie(w, &http.Cookie{
+		Name:     auth.SessionCookieName,
+		Value:    token,
+		Path:     "/",
+		HttpOnly: true,
+		Secure:   true,
+		SameSite: http.SameSiteLaxMode,
+		MaxAge:   int(expiration.Seconds()),
+	})
+}
+
+// wasCredentialSubmitted checks if credentials were submitted for the given auth method.
+func wasCredentialSubmitted(r *http.Request, method auth.Method) bool {
+	switch method {
+	case auth.MethodPIN:
+		return r.FormValue("pin") != ""
+	case auth.MethodPassword:
+		return r.FormValue("password") != ""
+	case auth.MethodOIDC:
+		return r.URL.Query().Get("session_token") != ""
+	}
+	return false
+}
+
+// AddDomain registers authentication schemes for the given domain. With schemes a valid session public key is required.
+// private=true forces ValidateTunnelPeer enforcement (403 on failure) regardless of the schemes list.
+func (mw *Middleware) AddDomain(domain string, schemes []Scheme, publicKeyB64 string, expiration time.Duration, accountID types.AccountID, serviceID types.ServiceID, ipRestrictions *restrict.Filter, private bool) error {
+	if len(schemes) == 0 {
+		mw.domainsMux.Lock()
+		defer mw.domainsMux.Unlock()
+		mw.domains[domain] = DomainConfig{
+			AccountID:      accountID,
+			ServiceID:      serviceID,
+			IPRestrictions: ipRestrictions,
+			Private:        private,
+		}
+		return nil
+	}
+
+	pubKeyBytes, err := base64.StdEncoding.DecodeString(publicKeyB64)
+	if err != nil {
+		return fmt.Errorf("decode session public key for domain %s: %w", domain, err)
+	}
+	if len(pubKeyBytes) != ed25519.PublicKeySize {
+		return fmt.Errorf("invalid session public key size for domain %s: got %d, want %d", domain, len(pubKeyBytes), ed25519.PublicKeySize)
+	}
+
+	mw.domainsMux.Lock()
+	defer mw.domainsMux.Unlock()
+	mw.domains[domain] = DomainConfig{
+		Schemes:           schemes,
+		SessionPublicKey:  pubKeyBytes,
+		SessionExpiration: expiration,
+		AccountID:         accountID,
+		ServiceID:         serviceID,
+		IPRestrictions:    ipRestrictions,
+		Private:           private,
+	}
+	return nil
+}
+
+// RemoveDomain unregisters authentication for the given domain.
+func (mw *Middleware) RemoveDomain(domain string) {
+	mw.domainsMux.Lock()
+	defer mw.domainsMux.Unlock()
+	delete(mw.domains, domain)
+}
+
+// validateSessionToken validates a session token. OIDC tokens with a configured
+// validator go through gRPC for group access checks; other methods validate locally.
+func (mw *Middleware) validateSessionToken(ctx context.Context, host, token string, publicKey ed25519.PublicKey, method auth.Method) (*validationResult, error) {
+	if method == auth.MethodOIDC && mw.sessionValidator != nil {
+		resp, err := mw.sessionValidator.ValidateSession(ctx, &proto.ValidateSessionRequest{
+			Domain:       host,
+			SessionToken: token,
+		})
+		if err != nil {
+			return nil, fmt.Errorf("%w: %w", errValidationUnavailable, err)
+		}
+		if !resp.Valid {
+			mw.logger.WithFields(log.Fields{
+				"domain":        host,
+				"denied_reason": resp.DeniedReason,
+				"user_id":       resp.UserId,
+			}).Debug("Session validation denied")
+			return &validationResult{
+				UserID:       resp.UserId,
+				UserEmail:    resp.GetUserEmail(),
+				Valid:        false,
+				DeniedReason: resp.DeniedReason,
+			}, nil
+		}
+		return &validationResult{
+			UserID:     resp.UserId,
+			UserEmail:  resp.GetUserEmail(),
+			Valid:      true,
+			Groups:     resp.GetPeerGroupIds(),
+			GroupNames: resp.GetPeerGroupNames(),
+		}, nil
+	}
+
+	userID, email, _, groups, groupNames, err := auth.ValidateSessionJWT(token, host, publicKey)
+	if err != nil {
+		return nil, err
+	}
+	return &validationResult{UserID: userID, UserEmail: email, Valid: true, Groups: groups, GroupNames: groupNames}, nil
+}
+
+// stripSessionTokenParam returns the request URI with the session_token query
+// parameter removed so it doesn't linger in the browser's address bar or history.
+func stripSessionTokenParam(u *url.URL) string {
+	q := u.Query()
+	if !q.Has("session_token") {
+		return u.RequestURI()
+	}
+	q.Del("session_token")
+	clean := *u
+	clean.RawQuery = q.Encode()
+	return clean.RequestURI()
+}

proxy/internal/auth/oidc.go

@@ -0,0 +1,66 @@
+package auth
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"net/url"
+
+	"google.golang.org/grpc"
+
+	"github.com/netbirdio/netbird/proxy/auth"
+	"github.com/netbirdio/netbird/proxy/internal/types"
+	"github.com/netbirdio/netbird/shared/management/proto"
+)
+
+type urlGenerator interface {
+	GetOIDCURL(context.Context, *proto.GetOIDCURLRequest, ...grpc.CallOption) (*proto.GetOIDCURLResponse, error)
+}
+
+type OIDC struct {
+	id             types.ServiceID
+	accountId      types.AccountID
+	forwardedProto string
+	client         urlGenerator
+}
+
+// NewOIDC creates a new OIDC authentication scheme
+func NewOIDC(client urlGenerator, id types.ServiceID, accountId types.AccountID, forwardedProto string) OIDC {
+	return OIDC{
+		id:             id,
+		accountId:      accountId,
+		forwardedProto: forwardedProto,
+		client:         client,
+	}
+}
+
+func (OIDC) Type() auth.Method {
+	return auth.MethodOIDC
+}
+
+// Authenticate checks for an OIDC session token or obtains the OIDC redirect URL.
+func (o OIDC) Authenticate(r *http.Request) (string, string, error) {
+	// Check for the session_token query param (from OIDC redirects).
+	// The management server passes the token in the URL because it cannot set
+	// cookies for the proxy's domain (cookies are domain-scoped per RFC 6265).
+	if token := r.URL.Query().Get("session_token"); token != "" {
+		return token, "", nil
+	}
+
+	redirectURL := &url.URL{
+		Scheme: auth.ResolveProto(o.forwardedProto, r.TLS),
+		Host:   r.Host,
+		Path:   r.URL.Path,
+	}
+
+	res, err := o.client.GetOIDCURL(r.Context(), &proto.GetOIDCURLRequest{
+		Id:          string(o.id),
+		AccountId:   string(o.accountId),
+		RedirectUrl: redirectURL.String(),
+	})
+	if err != nil {
+		return "", "", fmt.Errorf("get OIDC URL: %w", err)
+	}
+
+	return "", res.GetUrl(), nil
+}

proxy/internal/auth/password.go

@@ -0,0 +1,63 @@
+package auth
+
+import (
+	"fmt"
+	"net/http"
+
+	"github.com/netbirdio/netbird/proxy/auth"
+	"github.com/netbirdio/netbird/proxy/internal/types"
+	"github.com/netbirdio/netbird/shared/management/proto"
+)
+
+const passwordFormId = "password"
+
+type Password struct {
+	id        types.ServiceID
+	accountId types.AccountID
+	client    authenticator
+}
+
+func NewPassword(client authenticator, id types.ServiceID, accountId types.AccountID) Password {
+	return Password{
+		id:        id,
+		accountId: accountId,
+		client:    client,
+	}
+}
+
+func (Password) Type() auth.Method {
+	return auth.MethodPassword
+}
+
+// Authenticate attempts to authenticate the request using a form
+// value passed in the request.
+// If authentication fails, the required HTTP form ID is returned
+// so that it can be injected into a request from the UI so that
+// authentication may be successful.
+func (p Password) Authenticate(r *http.Request) (string, string, error) {
+	password := r.FormValue(passwordFormId)
+
+	if password == "" {
+		// No password submitted; return the form ID so the UI can prompt the user.
+		return "", passwordFormId, nil
+	}
+
+	res, err := p.client.Authenticate(r.Context(), &proto.AuthenticateRequest{
+		Id:        string(p.id),
+		AccountId: string(p.accountId),
+		Request: &proto.AuthenticateRequest_Password{
+			Password: &proto.PasswordRequest{
+				Password: password,
+			},
+		},
+	})
+	if err != nil {
+		return "", "", fmt.Errorf("authenticate password: %w", err)
+	}
+
+	if res.GetSuccess() {
+		return res.GetSessionToken(), "", nil
+	}
+
+	return "", passwordFormId, nil
+}

proxy/internal/auth/pin.go

@@ -0,0 +1,63 @@
+package auth
+
+import (
+	"fmt"
+	"net/http"
+
+	"github.com/netbirdio/netbird/proxy/auth"
+	"github.com/netbirdio/netbird/proxy/internal/types"
+	"github.com/netbirdio/netbird/shared/management/proto"
+)
+
+const pinFormId = "pin"
+
+type Pin struct {
+	id        types.ServiceID
+	accountId types.AccountID
+	client    authenticator
+}
+
+func NewPin(client authenticator, id types.ServiceID, accountId types.AccountID) Pin {
+	return Pin{
+		id:        id,
+		accountId: accountId,
+		client:    client,
+	}
+}
+
+func (Pin) Type() auth.Method {
+	return auth.MethodPIN
+}
+
+// Authenticate attempts to authenticate the request using a form
+// value passed in the request.
+// If authentication fails, the required HTTP form ID is returned
+// so that it can be injected into a request from the UI so that
+// authentication may be successful.
+func (p Pin) Authenticate(r *http.Request) (string, string, error) {
+	pin := r.FormValue(pinFormId)
+
+	if pin == "" {
+		// No PIN submitted; return the form ID so the UI can prompt the user.
+		return "", pinFormId, nil
+	}
+
+	res, err := p.client.Authenticate(r.Context(), &proto.AuthenticateRequest{
+		Id:        string(p.id),
+		AccountId: string(p.accountId),
+		Request: &proto.AuthenticateRequest_Pin{
+			Pin: &proto.PinRequest{
+				Pin: pin,
+			},
+		},
+	})
+	if err != nil {
+		return "", "", fmt.Errorf("authenticate pin: %w", err)
+	}
+
+	if res.GetSuccess() {
+		return res.GetSessionToken(), "", nil
+	}
+
+	return "", pinFormId, nil
+}

proxy/internal/auth/tunnel_cache.go

@@ -0,0 +1,171 @@
+package auth
+
+import (
+	"context"
+	"net/netip"
+	"sync"
+	"time"
+
+	"golang.org/x/sync/singleflight"
+
+	"github.com/netbirdio/netbird/proxy/internal/types"
+	"github.com/netbirdio/netbird/shared/management/proto"
+)
+
+// tunnelCacheTTL caps how long a positive ValidateTunnelPeer result is
+// reused before re-fetching from management. 5 minutes balances freshness
+// against management load on busy mesh networks.
+const tunnelCacheTTL = 300 * time.Second
+
+// tunnelCachePerAccount caps the number of cached identities per account.
+// Bounded eviction avoids memory growth in pathological cases (huge peer
+// roster, brief request bursts) while staying generous for normal use.
+const tunnelCachePerAccount = 1024
+
+// tunnelCacheKey identifies a cached entry by tunnel IP and originating
+// account. Domain is part of the value, not the key, because the
+// management response is per (account, IP) — domain only gates whether a
+// re-fetch is needed if the operator is accessing a different service.
+type tunnelCacheKey struct {
+	accountID types.AccountID
+	tunnelIP  netip.Addr
+	domain    string
+}
+
+// tunnelCacheEntry stores a positive validation response with the time it
+// was minted. Entries past tunnelCacheTTL are treated as misses.
+type tunnelCacheEntry struct {
+	resp     *proto.ValidateTunnelPeerResponse
+	cachedAt time.Time
+}
+
+// tunnelValidationCache memoizes ValidateTunnelPeer responses keyed by
+// (accountID, tunnelIP, domain). Only successful, valid responses are
+// cached — denials skip the cache so policy changes apply immediately.
+// Single-flight de-duplicates concurrent fetches for the same key so a
+// burst of cold requests collapses into a single RPC.
+type tunnelValidationCache struct {
+	mu      sync.Mutex
+	entries map[types.AccountID]*accountBucket
+	flight  singleflight.Group
+	ttl     time.Duration
+	maxSize int
+	now     func() time.Time
+}
+
+// accountBucket holds the cached entries for a single account, with a
+// FIFO eviction queue used when the bucket exceeds maxSize.
+type accountBucket struct {
+	items map[tunnelCacheKey]tunnelCacheEntry
+	order []tunnelCacheKey
+}
+
+// newTunnelValidationCache constructs a cache with default TTL and bounds.
+func newTunnelValidationCache() *tunnelValidationCache {
+	return &tunnelValidationCache{
+		entries: make(map[types.AccountID]*accountBucket),
+		ttl:     tunnelCacheTTL,
+		maxSize: tunnelCachePerAccount,
+		now:     time.Now,
+	}
+}
+
+// get returns a cached response for the key, or nil when missing or
+// expired. Expired entries are evicted lazily on read.
+func (c *tunnelValidationCache) get(key tunnelCacheKey) *proto.ValidateTunnelPeerResponse {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	bucket, ok := c.entries[key.accountID]
+	if !ok {
+		return nil
+	}
+	entry, ok := bucket.items[key]
+	if !ok {
+		return nil
+	}
+	if c.now().Sub(entry.cachedAt) > c.ttl {
+		delete(bucket.items, key)
+		bucket.order = removeKey(bucket.order, key)
+		return nil
+	}
+	return entry.resp
+}
+
+// put records a positive response under the key. Evicts the oldest entry
+// in the account's bucket when the bound is exceeded.
+func (c *tunnelValidationCache) put(key tunnelCacheKey, resp *proto.ValidateTunnelPeerResponse) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	bucket, ok := c.entries[key.accountID]
+	if !ok {
+		bucket = &accountBucket{items: make(map[tunnelCacheKey]tunnelCacheEntry)}
+		c.entries[key.accountID] = bucket
+	}
+	if _, exists := bucket.items[key]; !exists {
+		bucket.order = append(bucket.order, key)
+	}
+	bucket.items[key] = tunnelCacheEntry{resp: resp, cachedAt: c.now()}
+
+	for len(bucket.order) > c.maxSize {
+		oldest := bucket.order[0]
+		bucket.order = bucket.order[1:]
+		delete(bucket.items, oldest)
+	}
+}
+
+// removeKey drops the first occurrence of needle from order. The cache
+// uses small slices so a linear scan is cheaper than a map+slice combo.
+func removeKey(order []tunnelCacheKey, needle tunnelCacheKey) []tunnelCacheKey {
+	for i, k := range order {
+		if k == needle {
+			return append(order[:i], order[i+1:]...)
+		}
+	}
+	return order
+}
+
+// flightKey turns a cache key into a single-flight string. AccountID and
+// IP isolation by themselves are insufficient because different domains
+// for the same peer/account may have different group access.
+func flightKey(key tunnelCacheKey) string {
+	return string(key.accountID) + "|" + key.tunnelIP.String() + "|" + key.domain
+}
+
+// validateTunnelPeerFn is the RPC entry point the cache wraps. It matches
+// the SessionValidator.ValidateTunnelPeer signature without exposing the
+// gRPC option variadic, since callers don't need it on the cache hot path.
+type validateTunnelPeerFn func(ctx context.Context, req *proto.ValidateTunnelPeerRequest) (*proto.ValidateTunnelPeerResponse, error)
+
+// fetch returns a cached response when present, otherwise calls validate
+// under single-flight and caches the result. Denied responses pass
+// through but are not cached so policy changes apply immediately.
+func (c *tunnelValidationCache) fetch(ctx context.Context, key tunnelCacheKey, validate validateTunnelPeerFn) (*proto.ValidateTunnelPeerResponse, bool, error) {
+	if resp := c.get(key); resp != nil {
+		return resp, true, nil
+	}
+
+	flight := flightKey(key)
+	res, err, _ := c.flight.Do(flight, func() (any, error) {
+		if cached := c.get(key); cached != nil {
+			return cached, nil
+		}
+		resp, err := validate(ctx, &proto.ValidateTunnelPeerRequest{
+			TunnelIp: key.tunnelIP.String(),
+			Domain:   key.domain,
+		})
+		if err != nil {
+			return nil, err
+		}
+		if resp.GetValid() && resp.GetSessionToken() != "" {
+			c.put(key, resp)
+		}
+		return resp, nil
+	})
+	if err != nil {
+		return nil, false, err
+	}
+	resp, _ := res.(*proto.ValidateTunnelPeerResponse)
+	return resp, false, nil
+}

proxy/internal/auth/tunnel_cache_test.go

@@ -0,0 +1,171 @@
+package auth
+
+import (
+	"context"
+	"net/netip"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/proxy/internal/types"
+	"github.com/netbirdio/netbird/shared/management/proto"
+)
+
+func newTestKey(account types.AccountID, ip string, domain string) tunnelCacheKey {
+	return tunnelCacheKey{
+		accountID: account,
+		tunnelIP:  netip.MustParseAddr(ip),
+		domain:    domain,
+	}
+}
+
+func TestTunnelCache_HitSkipsRPC(t *testing.T) {
+	cache := newTunnelValidationCache()
+	key := newTestKey("acct-1", "100.64.0.10", "svc.example")
+
+	var calls int32
+	validate := func(_ context.Context, req *proto.ValidateTunnelPeerRequest) (*proto.ValidateTunnelPeerResponse, error) {
+		atomic.AddInt32(&calls, 1)
+		return &proto.ValidateTunnelPeerResponse{Valid: true, SessionToken: "tok", UserId: "user-1"}, nil
+	}
+
+	resp, fromCache, err := cache.fetch(context.Background(), key, validate)
+	require.NoError(t, err)
+	require.NotNil(t, resp, "first fetch returns RPC response")
+	assert.False(t, fromCache, "first fetch must not be cached")
+
+	resp2, fromCache2, err := cache.fetch(context.Background(), key, validate)
+	require.NoError(t, err)
+	require.NotNil(t, resp2, "second fetch returns cached response")
+	assert.True(t, fromCache2, "second fetch must be served from cache")
+	assert.Equal(t, "user-1", resp2.GetUserId(), "cached response should preserve user identity")
+	assert.Equal(t, int32(1), atomic.LoadInt32(&calls), "validate should run exactly once with one cache hit")
+}
+
+func TestTunnelCache_ExpiredEntryRefetches(t *testing.T) {
+	cache := newTunnelValidationCache()
+	clock := time.Now()
+	cache.now = func() time.Time { return clock }
+
+	key := newTestKey("acct-1", "100.64.0.10", "svc.example")
+	var calls int32
+	validate := func(_ context.Context, _ *proto.ValidateTunnelPeerRequest) (*proto.ValidateTunnelPeerResponse, error) {
+		atomic.AddInt32(&calls, 1)
+		return &proto.ValidateTunnelPeerResponse{Valid: true, SessionToken: "tok"}, nil
+	}
+
+	_, _, err := cache.fetch(context.Background(), key, validate)
+	require.NoError(t, err)
+	assert.Equal(t, int32(1), atomic.LoadInt32(&calls), "first fetch issues one RPC")
+
+	clock = clock.Add(tunnelCacheTTL + time.Second)
+
+	_, fromCache, err := cache.fetch(context.Background(), key, validate)
+	require.NoError(t, err)
+	assert.False(t, fromCache, "expired entry must miss the cache")
+	assert.Equal(t, int32(2), atomic.LoadInt32(&calls), "expired entry forces a re-fetch")
+}
+
+func TestTunnelCache_DeniedResponseNotCached(t *testing.T) {
+	cache := newTunnelValidationCache()
+	key := newTestKey("acct-1", "100.64.0.10", "svc.example")
+
+	var calls int32
+	validate := func(_ context.Context, _ *proto.ValidateTunnelPeerRequest) (*proto.ValidateTunnelPeerResponse, error) {
+		atomic.AddInt32(&calls, 1)
+		return &proto.ValidateTunnelPeerResponse{Valid: false, DeniedReason: "not_in_group"}, nil
+	}
+
+	for i := 0; i < 3; i++ {
+		_, _, err := cache.fetch(context.Background(), key, validate)
+		require.NoError(t, err, "fetch must not error on denied response")
+	}
+	assert.Equal(t, int32(3), atomic.LoadInt32(&calls), "denied responses bypass the cache so policy changes apply immediately")
+}
+
+func TestTunnelCache_ConcurrentColdHitsCoalesce(t *testing.T) {
+	cache := newTunnelValidationCache()
+	key := newTestKey("acct-1", "100.64.0.10", "svc.example")
+
+	gate := make(chan struct{})
+	var calls int32
+	validate := func(_ context.Context, _ *proto.ValidateTunnelPeerRequest) (*proto.ValidateTunnelPeerResponse, error) {
+		atomic.AddInt32(&calls, 1)
+		<-gate
+		return &proto.ValidateTunnelPeerResponse{Valid: true, SessionToken: "tok"}, nil
+	}
+
+	const workers = 16
+	var wg sync.WaitGroup
+	wg.Add(workers)
+	results := make([]bool, workers)
+	for i := 0; i < workers; i++ {
+		go func(idx int) {
+			defer wg.Done()
+			resp, _, err := cache.fetch(context.Background(), key, validate)
+			results[idx] = err == nil && resp.GetValid()
+		}(i)
+	}
+
+	time.Sleep(20 * time.Millisecond)
+	close(gate)
+	wg.Wait()
+
+	for i, ok := range results {
+		assert.Truef(t, ok, "worker %d should observe a successful response", i)
+	}
+	assert.Equal(t, int32(1), atomic.LoadInt32(&calls), "single-flight must collapse concurrent cold fetches into one RPC")
+}
+
+func TestTunnelCache_PerAccountIsolation(t *testing.T) {
+	cache := newTunnelValidationCache()
+	keyA := newTestKey("acct-a", "100.64.0.10", "svc.example")
+	keyB := newTestKey("acct-b", "100.64.0.10", "svc.example")
+
+	var callsA, callsB int32
+	validateA := func(_ context.Context, _ *proto.ValidateTunnelPeerRequest) (*proto.ValidateTunnelPeerResponse, error) {
+		atomic.AddInt32(&callsA, 1)
+		return &proto.ValidateTunnelPeerResponse{Valid: true, SessionToken: "tok-a", UserId: "user-a"}, nil
+	}
+	validateB := func(_ context.Context, _ *proto.ValidateTunnelPeerRequest) (*proto.ValidateTunnelPeerResponse, error) {
+		atomic.AddInt32(&callsB, 1)
+		return &proto.ValidateTunnelPeerResponse{Valid: true, SessionToken: "tok-b", UserId: "user-b"}, nil
+	}
+
+	respA, _, err := cache.fetch(context.Background(), keyA, validateA)
+	require.NoError(t, err)
+	respB, _, err := cache.fetch(context.Background(), keyB, validateB)
+	require.NoError(t, err)
+
+	assert.Equal(t, "user-a", respA.GetUserId(), "account A response should belong to user-a")
+	assert.Equal(t, "user-b", respB.GetUserId(), "account B response must not be served from account A's cache")
+	assert.Equal(t, int32(1), atomic.LoadInt32(&callsA), "validateA called exactly once")
+	assert.Equal(t, int32(1), atomic.LoadInt32(&callsB), "validateB called exactly once")
+}
+
+func TestTunnelCache_BoundedSizeEvictsOldest(t *testing.T) {
+	cache := newTunnelValidationCache()
+	cache.maxSize = 2
+
+	validate := func(_ context.Context, req *proto.ValidateTunnelPeerRequest) (*proto.ValidateTunnelPeerResponse, error) {
+		return &proto.ValidateTunnelPeerResponse{Valid: true, SessionToken: "tok-" + req.GetTunnelIp()}, nil
+	}
+
+	keys := []tunnelCacheKey{
+		newTestKey("acct-1", "100.64.0.10", "svc"),
+		newTestKey("acct-1", "100.64.0.11", "svc"),
+		newTestKey("acct-1", "100.64.0.12", "svc"),
+	}
+	for _, k := range keys {
+		_, _, err := cache.fetch(context.Background(), k, validate)
+		require.NoError(t, err)
+	}
+
+	assert.Nil(t, cache.get(keys[0]), "oldest key should be evicted past maxSize")
+	assert.NotNil(t, cache.get(keys[1]), "second-newest must remain cached")
+	assert.NotNil(t, cache.get(keys[2]), "newest must remain cached")
+}

proxy/internal/auth/tunnel_lookup_test.go

@@ -0,0 +1,355 @@
+package auth
+
+import (
+	"context"
+	"errors"
+	"net/http"
+	"net/http/httptest"
+	"net/netip"
+	"sync/atomic"
+	"testing"
+
+	log "github.com/sirupsen/logrus"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc"
+
+	"github.com/netbirdio/netbird/proxy/internal/proxy"
+	"github.com/netbirdio/netbird/shared/management/proto"
+)
+
+// stubSessionValidator records ValidateTunnelPeer calls and returns the
+// pre-canned response. Counts let tests assert RPC traffic.
+type stubSessionValidator struct {
+	respFn      func(req *proto.ValidateTunnelPeerRequest) *proto.ValidateTunnelPeerResponse
+	respErr     error
+	tunnelCalls atomic.Int32
+}
+
+func (s *stubSessionValidator) ValidateSession(_ context.Context, _ *proto.ValidateSessionRequest, _ ...grpc.CallOption) (*proto.ValidateSessionResponse, error) {
+	return &proto.ValidateSessionResponse{Valid: false}, nil
+}
+
+func (s *stubSessionValidator) ValidateTunnelPeer(_ context.Context, in *proto.ValidateTunnelPeerRequest, _ ...grpc.CallOption) (*proto.ValidateTunnelPeerResponse, error) {
+	s.tunnelCalls.Add(1)
+	if s.respErr != nil {
+		return nil, s.respErr
+	}
+	if s.respFn != nil {
+		return s.respFn(in), nil
+	}
+	return &proto.ValidateTunnelPeerResponse{Valid: false}, nil
+}
+
+func newTunnelMiddleware(t *testing.T, validator SessionValidator) *Middleware {
+	t.Helper()
+	mw := NewMiddleware(log.New(), validator, nil)
+	require.NoError(t, mw.AddDomain("svc.example", nil, "", 0, "acct-1", "svc-1", nil, false))
+	return mw
+}
+
+func newTunnelRequest(remoteAddr string) (*httptest.ResponseRecorder, *http.Request) {
+	w := httptest.NewRecorder()
+	r := httptest.NewRequest(http.MethodGet, "https://svc.example/", nil)
+	r.Host = "svc.example"
+	r.RemoteAddr = remoteAddr
+	return w, r
+}
+
+// TestForwardWithTunnelPeer_LocalLookupUnknownIPDeniesFast verifies the
+// short-circuit: a tunnel IP not in the account's roster never reaches
+// management's ValidateTunnelPeer.
+func TestForwardWithTunnelPeer_LocalLookupUnknownIPDeniesFast(t *testing.T) {
+	validator := &stubSessionValidator{}
+	mw := newTunnelMiddleware(t, validator)
+
+	lookup := TunnelLookupFunc(func(_ netip.Addr) (PeerIdentity, bool) {
+		return PeerIdentity{}, false
+	})
+
+	w, r := newTunnelRequest("100.64.0.99:55555")
+	r = r.WithContext(WithTunnelLookup(r.Context(), lookup))
+
+	called := false
+	next := http.HandlerFunc(func(http.ResponseWriter, *http.Request) { called = true })
+
+	config, _ := mw.getDomainConfig("svc.example")
+	handled := mw.forwardWithTunnelPeer(w, r, "svc.example", config, next)
+
+	assert.False(t, handled, "unknown peer must fall through, not forward")
+	assert.False(t, called, "next handler must not run for unknown peer")
+	assert.Equal(t, int32(0), validator.tunnelCalls.Load(), "ValidateTunnelPeer must be skipped on local-lookup miss")
+}
+
+// TestForwardWithTunnelPeer_GroupsPropagateToCapturedData verifies the proxy
+// surfaces the calling peer's group memberships from ValidateTunnelPeerResponse
+// onto CapturedData so policy-aware middlewares can authorise without an
+// extra management round-trip.
+func TestForwardWithTunnelPeer_GroupsPropagateToCapturedData(t *testing.T) {
+	groups := []string{"engineering", "sre"}
+	validator := &stubSessionValidator{
+		respFn: func(_ *proto.ValidateTunnelPeerRequest) *proto.ValidateTunnelPeerResponse {
+			return &proto.ValidateTunnelPeerResponse{
+				Valid:        true,
+				SessionToken: "tok",
+				UserId:       "user-1",
+				PeerGroupIds: groups,
+			}
+		},
+	}
+	mw := newTunnelMiddleware(t, validator)
+
+	w, r := newTunnelRequest("100.64.0.10:55555")
+	cd := proxy.NewCapturedData("")
+	lookup := TunnelLookupFunc(func(_ netip.Addr) (PeerIdentity, bool) {
+		return PeerIdentity{}, true
+	})
+	r = r.WithContext(proxy.WithCapturedData(WithTunnelLookup(r.Context(), lookup), cd))
+
+	called := false
+	next := http.HandlerFunc(func(http.ResponseWriter, *http.Request) { called = true })
+
+	config, _ := mw.getDomainConfig("svc.example")
+	handled := mw.forwardWithTunnelPeer(w, r, "svc.example", config, next)
+
+	require.True(t, handled, "valid tunnel-peer response must forward")
+	require.True(t, called, "next handler must run")
+	assert.Equal(t, "user-1", cd.GetUserID(), "user id must propagate from tunnel-peer response")
+	assert.Equal(t, groups, cd.GetUserGroups(), "peer group IDs must propagate from tunnel-peer response")
+}
+
+// TestForwardWithTunnelPeer_LocalLookupKnownPeerStillRPCs verifies that a
+// known tunnel IP still triggers ValidateTunnelPeer for the user-identity
+// tail (UserID + group access). Phase 3 only short-circuits the deny path.
+func TestForwardWithTunnelPeer_LocalLookupKnownPeerStillRPCs(t *testing.T) {
+	validator := &stubSessionValidator{
+		respFn: func(_ *proto.ValidateTunnelPeerRequest) *proto.ValidateTunnelPeerResponse {
+			return &proto.ValidateTunnelPeerResponse{Valid: true, SessionToken: "tok", UserId: "user-1"}
+		},
+	}
+	mw := newTunnelMiddleware(t, validator)
+
+	knownIP := netip.MustParseAddr("100.64.0.10")
+	lookup := TunnelLookupFunc(func(ip netip.Addr) (PeerIdentity, bool) {
+		if ip == knownIP {
+			return PeerIdentity{PubKey: "pk", TunnelIP: ip, FQDN: "peer.netbird.cloud"}, true
+		}
+		return PeerIdentity{}, false
+	})
+
+	w, r := newTunnelRequest(knownIP.String() + ":55555")
+	r = r.WithContext(WithTunnelLookup(r.Context(), lookup))
+
+	called := false
+	next := http.HandlerFunc(func(http.ResponseWriter, *http.Request) { called = true })
+
+	config, _ := mw.getDomainConfig("svc.example")
+	handled := mw.forwardWithTunnelPeer(w, r, "svc.example", config, next)
+
+	assert.True(t, handled, "known peer with valid RPC response must forward")
+	assert.True(t, called, "next handler must run on success")
+	assert.Equal(t, int32(1), validator.tunnelCalls.Load(), "RPC must run for the user-identity tail when local lookup confirms the peer")
+}
+
+// TestForwardWithTunnelPeer_NoLookupRefusesFastPath guards the
+// anti-spoof gate: requests that didn't arrive on the per-account
+// inbound listener (no TunnelLookup attached) must never reach
+// management's ValidateTunnelPeer, even when the source IP looks like
+// a tunnel address. A colliding RFC1918 / CGNAT source on the public
+// listener would otherwise impersonate a mesh peer.
+func TestForwardWithTunnelPeer_NoLookupRefusesFastPath(t *testing.T) {
+	validator := &stubSessionValidator{
+		respFn: func(_ *proto.ValidateTunnelPeerRequest) *proto.ValidateTunnelPeerResponse {
+			return &proto.ValidateTunnelPeerResponse{Valid: true, SessionToken: "tok", UserId: "user-1"}
+		},
+	}
+	mw := newTunnelMiddleware(t, validator)
+
+	w, r := newTunnelRequest("100.64.0.10:55555")
+	called := false
+	next := http.HandlerFunc(func(http.ResponseWriter, *http.Request) { called = true })
+
+	config, _ := mw.getDomainConfig("svc.example")
+	handled := mw.forwardWithTunnelPeer(w, r, "svc.example", config, next)
+
+	assert.False(t, handled, "fast-path must refuse without the inbound marker")
+	assert.False(t, called, "next handler must not run")
+	assert.Equal(t, int32(0), validator.tunnelCalls.Load(), "ValidateTunnelPeer must not be invoked without the inbound marker")
+}
+
+// TestForwardWithTunnelPeer_RPCErrorFallsThrough validates that an RPC
+// failure still falls through to the next scheme (no false positive).
+func TestForwardWithTunnelPeer_RPCErrorFallsThrough(t *testing.T) {
+	validator := &stubSessionValidator{respErr: errors.New("management down")}
+	mw := newTunnelMiddleware(t, validator)
+
+	knownIP := netip.MustParseAddr("100.64.0.10")
+	lookup := TunnelLookupFunc(func(ip netip.Addr) (PeerIdentity, bool) {
+		return PeerIdentity{TunnelIP: ip}, true
+	})
+
+	w, r := newTunnelRequest(knownIP.String() + ":55555")
+	r = r.WithContext(WithTunnelLookup(r.Context(), lookup))
+
+	config, _ := mw.getDomainConfig("svc.example")
+	handled := mw.forwardWithTunnelPeer(w, r, "svc.example", config, http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}))
+
+	assert.False(t, handled, "RPC error must let the caller try other schemes")
+	assert.Equal(t, int32(1), validator.tunnelCalls.Load(), "RPC was attempted exactly once")
+}
+
+// TestForwardWithTunnelPeer_CacheReusesPositiveResponse confirms the
+// (account, IP, domain) cache prevents repeated RPCs for the same peer.
+func TestForwardWithTunnelPeer_CacheReusesPositiveResponse(t *testing.T) {
+	validator := &stubSessionValidator{
+		respFn: func(_ *proto.ValidateTunnelPeerRequest) *proto.ValidateTunnelPeerResponse {
+			return &proto.ValidateTunnelPeerResponse{Valid: true, SessionToken: "tok", UserId: "user-1"}
+		},
+	}
+	mw := newTunnelMiddleware(t, validator)
+
+	lookup := TunnelLookupFunc(func(_ netip.Addr) (PeerIdentity, bool) {
+		return PeerIdentity{}, true
+	})
+
+	for i := 0; i < 4; i++ {
+		w, r := newTunnelRequest("100.64.0.10:55555")
+		r = r.WithContext(WithTunnelLookup(r.Context(), lookup))
+		next := http.HandlerFunc(func(http.ResponseWriter, *http.Request) {})
+		config, _ := mw.getDomainConfig("svc.example")
+		handled := mw.forwardWithTunnelPeer(w, r, "svc.example", config, next)
+		require.True(t, handled, "iteration %d should forward", i)
+	}
+
+	assert.Equal(t, int32(1), validator.tunnelCalls.Load(), "subsequent forwards must hit the cache, not management")
+}
+
+// TestForwardWithTunnelPeer_RoutesAccountIDIntoCacheKey ensures cache keys
+// honour account scoping — same tunnel IP on different accounts must not
+// collide.
+func TestForwardWithTunnelPeer_RoutesAccountIDIntoCacheKey(t *testing.T) {
+	validator := &stubSessionValidator{
+		respFn: func(req *proto.ValidateTunnelPeerRequest) *proto.ValidateTunnelPeerResponse {
+			return &proto.ValidateTunnelPeerResponse{Valid: true, SessionToken: "tok", UserId: "user"}
+		},
+	}
+	mw := NewMiddleware(log.New(), validator, nil)
+
+	require.NoError(t, mw.AddDomain("svc-a.example", nil, "", 0, "acct-a", "svc-a", nil, false))
+	require.NoError(t, mw.AddDomain("svc-b.example", nil, "", 0, "acct-b", "svc-b", nil, false))
+
+	// The fast-path requires the inbound-listener marker on the context.
+	// The peerstore lookup itself is account-agnostic at this level
+	// (one TunnelLookupFunc per account is attached by inbound.go); a
+	// trivial "always hit" lookup is enough to exercise the cache-key
+	// branch this test covers.
+	lookup := TunnelLookupFunc(func(_ netip.Addr) (PeerIdentity, bool) {
+		return PeerIdentity{}, true
+	})
+
+	for _, host := range []string{"svc-a.example", "svc-b.example"} {
+		w := httptest.NewRecorder()
+		r := httptest.NewRequest(http.MethodGet, "https://"+host+"/", nil)
+		r.Host = host
+		r.RemoteAddr = "100.64.0.10:55555"
+		r = r.WithContext(WithTunnelLookup(r.Context(), lookup))
+		config, _ := mw.getDomainConfig(host)
+		handled := mw.forwardWithTunnelPeer(w, r, host, config, http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}))
+		require.True(t, handled, "host %s should forward", host)
+	}
+
+	assert.Equal(t, int32(2), validator.tunnelCalls.Load(), "cache must not collide across accounts even when tunnel IPs match")
+}
+
+// TestForwardWithTunnelPeer_LocalLookupShortCircuitDoesNotPopulateCache
+// guarantees that the deny-fast path leaves the cache untouched, so a
+// subsequent request from the same IP after the peerstore catches up
+// goes through the normal RPC flow.
+func TestForwardWithTunnelPeer_LocalLookupShortCircuitDoesNotPopulateCache(t *testing.T) {
+	validator := &stubSessionValidator{
+		respFn: func(_ *proto.ValidateTunnelPeerRequest) *proto.ValidateTunnelPeerResponse {
+			return &proto.ValidateTunnelPeerResponse{Valid: true, SessionToken: "tok"}
+		},
+	}
+	mw := newTunnelMiddleware(t, validator)
+
+	knownIP := netip.MustParseAddr("100.64.0.10")
+	known := false
+	lookup := TunnelLookupFunc(func(ip netip.Addr) (PeerIdentity, bool) {
+		if known && ip == knownIP {
+			return PeerIdentity{TunnelIP: ip}, true
+		}
+		return PeerIdentity{}, false
+	})
+
+	doRequest := func() bool {
+		w, r := newTunnelRequest(knownIP.String() + ":55555")
+		r = r.WithContext(WithTunnelLookup(r.Context(), lookup))
+		config, _ := mw.getDomainConfig("svc.example")
+		return mw.forwardWithTunnelPeer(w, r, "svc.example", config, http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}))
+	}
+
+	require.False(t, doRequest(), "first request must short-circuit")
+	require.Equal(t, int32(0), validator.tunnelCalls.Load(), "short-circuit must not populate the cache")
+
+	known = true
+	require.True(t, doRequest(), "second request with peer in roster must forward via RPC")
+	assert.Equal(t, int32(1), validator.tunnelCalls.Load(), "RPC runs once after peerstore catches up")
+}
+
+func TestPrivateService_FailsClosedOnTunnelPeerFailure(t *testing.T) {
+	mw := NewMiddleware(log.New(), nil, nil)
+	require.NoError(t, mw.AddDomain("private.svc", nil, "", 0, "acct-1", "svc-1", nil, true))
+
+	called := false
+	handler := mw.Protect(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		called = true
+		w.WriteHeader(http.StatusOK)
+	}))
+
+	req := httptest.NewRequest(http.MethodGet, "https://private.svc/", nil)
+	req.Host = "private.svc"
+	req.RemoteAddr = "100.64.0.10:55555"
+	w := httptest.NewRecorder()
+	handler.ServeHTTP(w, req)
+
+	assert.Equal(t, http.StatusForbidden, w.Code)
+	assert.False(t, called)
+}
+
+func TestPrivateService_ForwardsOnTunnelPeerSuccess(t *testing.T) {
+	validator := &stubSessionValidator{
+		respFn: func(_ *proto.ValidateTunnelPeerRequest) *proto.ValidateTunnelPeerResponse {
+			return &proto.ValidateTunnelPeerResponse{
+				Valid:        true,
+				SessionToken: "tok",
+				UserId:       "user-1",
+			}
+		},
+	}
+	mw := NewMiddleware(log.New(), validator, nil)
+	require.NoError(t, mw.AddDomain("private.svc", nil, "", 0, "acct-1", "svc-1", nil, true))
+
+	called := false
+	handler := mw.Protect(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		called = true
+		w.WriteHeader(http.StatusOK)
+	}))
+
+	// Per-account inbound listener attaches WithTunnelLookup; without it
+	// forwardWithTunnelPeer refuses to take the fast-path. Mirror the
+	// real flow so this test exercises the post-gating success branch.
+	lookup := TunnelLookupFunc(func(_ netip.Addr) (PeerIdentity, bool) {
+		return PeerIdentity{}, true
+	})
+
+	req := httptest.NewRequest(http.MethodGet, "https://private.svc/", nil)
+	req.Host = "private.svc"
+	req.RemoteAddr = "100.64.0.10:55555"
+	req = req.WithContext(WithTunnelLookup(req.Context(), lookup))
+	w := httptest.NewRecorder()
+	handler.ServeHTTP(w, req)
+
+	assert.Equal(t, http.StatusOK, w.Code)
+	assert.True(t, called)
+}

proxy/internal/certwatch/watcher.go

@@ -0,0 +1,286 @@
+// Package certwatch watches TLS certificate files on disk and provides
+// a hot-reloading GetCertificate callback for tls.Config.
+package certwatch
+
+import (
+	"context"
+	"crypto/tls"
+	"crypto/x509"
+	"fmt"
+	"path/filepath"
+	"sync"
+	"time"
+
+	"github.com/fsnotify/fsnotify"
+	log "github.com/sirupsen/logrus"
+)
+
+const (
+	defaultPollInterval = 30 * time.Second
+	debounceDelay       = 500 * time.Millisecond
+)
+
+// Watcher monitors TLS certificate files on disk and caches the loaded
+// certificate in memory. It detects changes via fsnotify (with a polling
+// fallback for filesystems like NFS that lack inotify support) and
+// reloads the certificate pair automatically.
+type Watcher struct {
+	certPath string
+	keyPath  string
+
+	mu   sync.RWMutex
+	cert *tls.Certificate
+	leaf *x509.Certificate
+
+	pollInterval time.Duration
+	logger       *log.Logger
+}
+
+// NewWatcher creates a Watcher that monitors the given cert and key files.
+// It performs an initial load of the certificate and returns an error
+// if the initial load fails.
+func NewWatcher(certPath, keyPath string, logger *log.Logger) (*Watcher, error) {
+	if logger == nil {
+		logger = log.StandardLogger()
+	}
+
+	w := &Watcher{
+		certPath:     certPath,
+		keyPath:      keyPath,
+		pollInterval: defaultPollInterval,
+		logger:       logger,
+	}
+
+	if err := w.reload(); err != nil {
+		return nil, fmt.Errorf("initial certificate load: %w", err)
+	}
+
+	return w, nil
+}
+
+// GetCertificate returns the current in-memory certificate.
+// It is safe for concurrent use and compatible with tls.Config.GetCertificate.
+func (w *Watcher) GetCertificate(_ *tls.ClientHelloInfo) (*tls.Certificate, error) {
+	w.mu.RLock()
+	defer w.mu.RUnlock()
+
+	return w.cert, nil
+}
+
+// Leaf returns the parsed leaf certificate, or nil if not yet loaded.
+func (w *Watcher) Leaf() *x509.Certificate {
+	w.mu.RLock()
+	defer w.mu.RUnlock()
+	return w.leaf
+}
+
+// Watch starts watching for certificate file changes. It blocks until
+// ctx is cancelled. It uses fsnotify for immediate detection and falls
+// back to polling if fsnotify is unavailable (e.g. on NFS).
+// Even with fsnotify active, a periodic poll runs as a safety net.
+func (w *Watcher) Watch(ctx context.Context) {
+	// Watch the parent directory rather than individual files. Some volume
+	// mounts use an atomic symlink swap (..data -> timestamped dir), so
+	// watching the parent directory catches the link replacement.
+	certDir := filepath.Dir(w.certPath)
+	keyDir := filepath.Dir(w.keyPath)
+
+	watcher, err := fsnotify.NewWatcher()
+	if err != nil {
+		w.logger.Warnf("fsnotify unavailable, using polling only: %v", err)
+		w.pollLoop(ctx)
+		return
+	}
+	defer func() {
+		if err := watcher.Close(); err != nil {
+			w.logger.Debugf("close fsnotify watcher: %v", err)
+		}
+	}()
+
+	if err := watcher.Add(certDir); err != nil {
+		w.logger.Warnf("fsnotify watch on %s failed, using polling only: %v", certDir, err)
+		w.pollLoop(ctx)
+		return
+	}
+
+	if keyDir != certDir {
+		if err := watcher.Add(keyDir); err != nil {
+			w.logger.Warnf("fsnotify watch on %s failed: %v", keyDir, err)
+		}
+	}
+
+	w.logger.Infof("watching certificate files in %s", certDir)
+	w.fsnotifyLoop(ctx, watcher)
+}
+
+func (w *Watcher) fsnotifyLoop(ctx context.Context, watcher *fsnotify.Watcher) {
+	certBase := filepath.Base(w.certPath)
+	keyBase := filepath.Base(w.keyPath)
+
+	var debounce *time.Timer
+	defer func() {
+		if debounce != nil {
+			debounce.Stop()
+		}
+	}()
+
+	// Periodic poll as a safety net for missed fsnotify events.
+	pollTicker := time.NewTicker(w.pollInterval)
+	defer pollTicker.Stop()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return
+
+		case event, ok := <-watcher.Events:
+			if !ok {
+				return
+			}
+
+			base := filepath.Base(event.Name)
+			if !isRelevantFile(base, certBase, keyBase) {
+				w.logger.Debugf("fsnotify: ignoring event %s on %s", event.Op, event.Name)
+				continue
+			}
+			if !event.Has(fsnotify.Create) && !event.Has(fsnotify.Write) && !event.Has(fsnotify.Rename) {
+				w.logger.Debugf("fsnotify: ignoring op %s on %s", event.Op, base)
+				continue
+			}
+
+			w.logger.Debugf("fsnotify: detected %s on %s, scheduling reload", event.Op, base)
+
+			// Debounce: cert-manager may write cert and key as separate
+			// operations. Wait briefly to load both at once.
+			if debounce != nil {
+				debounce.Stop()
+			}
+			debounce = time.AfterFunc(debounceDelay, func() {
+				if ctx.Err() != nil {
+					return
+				}
+				w.tryReload()
+			})
+
+		case err, ok := <-watcher.Errors:
+			if !ok {
+				return
+			}
+			w.logger.Warnf("fsnotify error: %v", err)
+
+		case <-pollTicker.C:
+			w.tryReload()
+		}
+	}
+}
+
+func (w *Watcher) pollLoop(ctx context.Context) {
+	ticker := time.NewTicker(w.pollInterval)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			w.tryReload()
+		}
+	}
+}
+
+// reload loads the certificate from disk and updates the in-memory cache.
+func (w *Watcher) reload() error {
+	cert, err := tls.LoadX509KeyPair(w.certPath, w.keyPath)
+	if err != nil {
+		return err
+	}
+
+	// Parse the leaf for comparison on subsequent reloads.
+	if cert.Leaf == nil && len(cert.Certificate) > 0 {
+		leaf, err := x509.ParseCertificate(cert.Certificate[0])
+		if err != nil {
+			return fmt.Errorf("parse leaf certificate: %w", err)
+		}
+		cert.Leaf = leaf
+	}
+
+	w.mu.Lock()
+	w.cert = &cert
+	w.leaf = cert.Leaf
+	w.mu.Unlock()
+
+	w.logCertDetails("loaded certificate", cert.Leaf)
+
+	return nil
+}
+
+// tryReload attempts to reload the certificate. It skips the update
+// if the certificate on disk is identical to the one in memory (same
+// serial number and issuer) to avoid redundant log noise.
+func (w *Watcher) tryReload() {
+	cert, err := tls.LoadX509KeyPair(w.certPath, w.keyPath)
+	if err != nil {
+		w.logger.Warnf("reload certificate: %v", err)
+		return
+	}
+
+	if cert.Leaf == nil && len(cert.Certificate) > 0 {
+		leaf, err := x509.ParseCertificate(cert.Certificate[0])
+		if err != nil {
+			w.logger.Warnf("parse reloaded leaf certificate: %v", err)
+			return
+		}
+		cert.Leaf = leaf
+	}
+
+	w.mu.Lock()
+
+	if w.leaf != nil && cert.Leaf != nil &&
+		w.leaf.SerialNumber.Cmp(cert.Leaf.SerialNumber) == 0 &&
+		w.leaf.Issuer.CommonName == cert.Leaf.Issuer.CommonName {
+		w.mu.Unlock()
+		return
+	}
+
+	prev := w.leaf
+	w.cert = &cert
+	w.leaf = cert.Leaf
+	w.mu.Unlock()
+
+	w.logCertChange(prev, cert.Leaf)
+}
+
+func (w *Watcher) logCertDetails(msg string, leaf *x509.Certificate) {
+	if leaf == nil {
+		w.logger.Info(msg)
+		return
+	}
+
+	w.logger.Infof("%s: subject=%q serial=%s SANs=%v notAfter=%s",
+		msg,
+		leaf.Subject.CommonName,
+		leaf.SerialNumber.Text(16),
+		leaf.DNSNames,
+		leaf.NotAfter.UTC().Format(time.RFC3339),
+	)
+}
+
+func (w *Watcher) logCertChange(prev, next *x509.Certificate) {
+	if prev == nil || next == nil {
+		w.logCertDetails("certificate reloaded from disk", next)
+		return
+	}
+
+	w.logger.Infof("certificate reloaded from disk: subject=%q -> %q serial=%s -> %s notAfter=%s -> %s",
+		prev.Subject.CommonName, next.Subject.CommonName,
+		prev.SerialNumber.Text(16), next.SerialNumber.Text(16),
+		prev.NotAfter.UTC().Format(time.RFC3339), next.NotAfter.UTC().Format(time.RFC3339),
+	)
+}
+
+// isRelevantFile returns true if the changed file name is one we care about.
+// This includes the cert/key files themselves and the ..data symlink used
+// by atomic volume mounts.
+func isRelevantFile(changed, certBase, keyBase string) bool {
+	return changed == certBase || changed == keyBase || changed == "..data"
+}

proxy/internal/certwatch/watcher_test.go

@@ -0,0 +1,292 @@
+package certwatch
+
+import (
+	"context"
+	"crypto/ecdsa"
+	"crypto/elliptic"
+	"crypto/rand"
+	"crypto/tls"
+	"crypto/x509"
+	"crypto/x509/pkix"
+	"encoding/pem"
+	"math/big"
+	"os"
+	"path/filepath"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func generateSelfSignedCert(t *testing.T, serial int64) (certPEM, keyPEM []byte) {
+	t.Helper()
+
+	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
+	require.NoError(t, err)
+
+	template := &x509.Certificate{
+		SerialNumber: big.NewInt(serial),
+		Subject:      pkix.Name{CommonName: "test"},
+		NotBefore:    time.Now().Add(-time.Hour),
+		NotAfter:     time.Now().Add(24 * time.Hour),
+	}
+
+	certDER, err := x509.CreateCertificate(rand.Reader, template, template, &key.PublicKey, key)
+	require.NoError(t, err)
+
+	certPEM = pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: certDER})
+
+	keyDER, err := x509.MarshalECPrivateKey(key)
+	require.NoError(t, err)
+	keyPEM = pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: keyDER})
+
+	return certPEM, keyPEM
+}
+
+func writeCert(t *testing.T, dir string, certPEM, keyPEM []byte) {
+	t.Helper()
+
+	require.NoError(t, os.WriteFile(filepath.Join(dir, "tls.crt"), certPEM, 0o600))
+	require.NoError(t, os.WriteFile(filepath.Join(dir, "tls.key"), keyPEM, 0o600))
+}
+
+func TestNewWatcher(t *testing.T) {
+	dir := t.TempDir()
+	certPEM, keyPEM := generateSelfSignedCert(t, 1)
+	writeCert(t, dir, certPEM, keyPEM)
+
+	w, err := NewWatcher(
+		filepath.Join(dir, "tls.crt"),
+		filepath.Join(dir, "tls.key"),
+		nil,
+	)
+	require.NoError(t, err)
+
+	cert, err := w.GetCertificate(nil)
+	require.NoError(t, err)
+	require.NotNil(t, cert)
+	assert.Equal(t, int64(1), cert.Leaf.SerialNumber.Int64())
+}
+
+func TestNewWatcherMissingFiles(t *testing.T) {
+	dir := t.TempDir()
+
+	_, err := NewWatcher(
+		filepath.Join(dir, "tls.crt"),
+		filepath.Join(dir, "tls.key"),
+		nil,
+	)
+	assert.Error(t, err)
+}
+
+func TestReload(t *testing.T) {
+	dir := t.TempDir()
+	certPEM1, keyPEM1 := generateSelfSignedCert(t, 100)
+	writeCert(t, dir, certPEM1, keyPEM1)
+
+	w, err := NewWatcher(
+		filepath.Join(dir, "tls.crt"),
+		filepath.Join(dir, "tls.key"),
+		nil,
+	)
+	require.NoError(t, err)
+
+	cert1, err := w.GetCertificate(nil)
+	require.NoError(t, err)
+	assert.Equal(t, int64(100), cert1.Leaf.SerialNumber.Int64())
+
+	// Write a new cert with a different serial.
+	certPEM2, keyPEM2 := generateSelfSignedCert(t, 200)
+	writeCert(t, dir, certPEM2, keyPEM2)
+
+	// Manually trigger reload.
+	w.tryReload()
+
+	cert2, err := w.GetCertificate(nil)
+	require.NoError(t, err)
+	assert.Equal(t, int64(200), cert2.Leaf.SerialNumber.Int64())
+}
+
+func TestTryReloadSkipsUnchanged(t *testing.T) {
+	dir := t.TempDir()
+	certPEM, keyPEM := generateSelfSignedCert(t, 42)
+	writeCert(t, dir, certPEM, keyPEM)
+
+	w, err := NewWatcher(
+		filepath.Join(dir, "tls.crt"),
+		filepath.Join(dir, "tls.key"),
+		nil,
+	)
+	require.NoError(t, err)
+
+	cert1, err := w.GetCertificate(nil)
+	require.NoError(t, err)
+
+	// Reload with same cert - pointer should remain the same.
+	w.tryReload()
+
+	cert2, err := w.GetCertificate(nil)
+	require.NoError(t, err)
+	assert.Same(t, cert1, cert2, "cert pointer should not change when content is the same")
+}
+
+func TestWatchDetectsChanges(t *testing.T) {
+	dir := t.TempDir()
+	certPEM1, keyPEM1 := generateSelfSignedCert(t, 1)
+	writeCert(t, dir, certPEM1, keyPEM1)
+
+	w, err := NewWatcher(
+		filepath.Join(dir, "tls.crt"),
+		filepath.Join(dir, "tls.key"),
+		nil,
+	)
+	require.NoError(t, err)
+
+	// Use a short poll interval for the test.
+	w.pollInterval = 100 * time.Millisecond
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	go w.Watch(ctx)
+
+	// Write new cert.
+	certPEM2, keyPEM2 := generateSelfSignedCert(t, 999)
+	writeCert(t, dir, certPEM2, keyPEM2)
+
+	// Wait for the watcher to pick it up.
+	require.Eventually(t, func() bool {
+		cert, err := w.GetCertificate(nil)
+		if err != nil {
+			return false
+		}
+		return cert.Leaf.SerialNumber.Int64() == 999
+	}, 5*time.Second, 50*time.Millisecond, "watcher should detect cert change")
+}
+
+func TestIsRelevantFile(t *testing.T) {
+	assert.True(t, isRelevantFile("tls.crt", "tls.crt", "tls.key"))
+	assert.True(t, isRelevantFile("tls.key", "tls.crt", "tls.key"))
+	assert.True(t, isRelevantFile("..data", "tls.crt", "tls.key"))
+	assert.False(t, isRelevantFile("other.txt", "tls.crt", "tls.key"))
+}
+
+// TestWatchSymlinkRotation simulates Kubernetes secret volume updates where
+// the data directory is atomically swapped via a ..data symlink.
+func TestWatchSymlinkRotation(t *testing.T) {
+	base := t.TempDir()
+
+	// Create initial target directory with certs.
+	dir1 := filepath.Join(base, "dir1")
+	require.NoError(t, os.Mkdir(dir1, 0o755))
+	certPEM1, keyPEM1 := generateSelfSignedCert(t, 1)
+	require.NoError(t, os.WriteFile(filepath.Join(dir1, "tls.crt"), certPEM1, 0o600))
+	require.NoError(t, os.WriteFile(filepath.Join(dir1, "tls.key"), keyPEM1, 0o600))
+
+	// Create ..data symlink pointing to dir1.
+	dataLink := filepath.Join(base, "..data")
+	require.NoError(t, os.Symlink(dir1, dataLink))
+
+	// Create tls.crt and tls.key as symlinks to ..data/{file}.
+	certLink := filepath.Join(base, "tls.crt")
+	keyLink := filepath.Join(base, "tls.key")
+	require.NoError(t, os.Symlink(filepath.Join(dataLink, "tls.crt"), certLink))
+	require.NoError(t, os.Symlink(filepath.Join(dataLink, "tls.key"), keyLink))
+
+	w, err := NewWatcher(certLink, keyLink, nil)
+	require.NoError(t, err)
+
+	cert, err := w.GetCertificate(nil)
+	require.NoError(t, err)
+	assert.Equal(t, int64(1), cert.Leaf.SerialNumber.Int64())
+
+	w.pollInterval = 100 * time.Millisecond
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	go w.Watch(ctx)
+
+	// Simulate k8s atomic rotation: create dir2, swap ..data symlink.
+	dir2 := filepath.Join(base, "dir2")
+	require.NoError(t, os.Mkdir(dir2, 0o755))
+	certPEM2, keyPEM2 := generateSelfSignedCert(t, 777)
+	require.NoError(t, os.WriteFile(filepath.Join(dir2, "tls.crt"), certPEM2, 0o600))
+	require.NoError(t, os.WriteFile(filepath.Join(dir2, "tls.key"), keyPEM2, 0o600))
+
+	// Atomic swap: create temp link, then rename over ..data.
+	tmpLink := filepath.Join(base, "..data_tmp")
+	require.NoError(t, os.Symlink(dir2, tmpLink))
+	require.NoError(t, os.Rename(tmpLink, dataLink))
+
+	require.Eventually(t, func() bool {
+		cert, err := w.GetCertificate(nil)
+		if err != nil {
+			return false
+		}
+		return cert.Leaf.SerialNumber.Int64() == 777
+	}, 5*time.Second, 50*time.Millisecond, "watcher should detect symlink rotation")
+}
+
+// TestPollLoopDetectsChanges verifies the poll-only fallback path works.
+func TestPollLoopDetectsChanges(t *testing.T) {
+	dir := t.TempDir()
+	certPEM1, keyPEM1 := generateSelfSignedCert(t, 1)
+	writeCert(t, dir, certPEM1, keyPEM1)
+
+	w, err := NewWatcher(
+		filepath.Join(dir, "tls.crt"),
+		filepath.Join(dir, "tls.key"),
+		nil,
+	)
+	require.NoError(t, err)
+
+	w.pollInterval = 100 * time.Millisecond
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	// Directly use pollLoop to test the fallback path.
+	go w.pollLoop(ctx)
+
+	certPEM2, keyPEM2 := generateSelfSignedCert(t, 555)
+	writeCert(t, dir, certPEM2, keyPEM2)
+
+	require.Eventually(t, func() bool {
+		cert, err := w.GetCertificate(nil)
+		if err != nil {
+			return false
+		}
+		return cert.Leaf.SerialNumber.Int64() == 555
+	}, 5*time.Second, 50*time.Millisecond, "poll loop should detect cert change")
+}
+
+func TestGetCertificateConcurrency(t *testing.T) {
+	dir := t.TempDir()
+	certPEM, keyPEM := generateSelfSignedCert(t, 1)
+	writeCert(t, dir, certPEM, keyPEM)
+
+	w, err := NewWatcher(
+		filepath.Join(dir, "tls.crt"),
+		filepath.Join(dir, "tls.key"),
+		nil,
+	)
+	require.NoError(t, err)
+
+	// Hammer GetCertificate concurrently while reloading.
+	done := make(chan struct{})
+	go func() {
+		for i := 0; i < 100; i++ {
+			w.tryReload()
+		}
+		close(done)
+	}()
+
+	for i := 0; i < 1000; i++ {
+		cert, err := w.GetCertificate(&tls.ClientHelloInfo{})
+		assert.NoError(t, err)
+		assert.NotNil(t, cert)
+	}
+
+	<-done
+}

proxy/internal/conntrack/conn.go

@@ -0,0 +1,51 @@
+package conntrack
+
+import (
+	"bufio"
+	"net"
+	"net/http"
+)
+
+// trackedConn wraps a net.Conn and removes itself from the tracker on Close.
+type trackedConn struct {
+	net.Conn
+	tracker *HijackTracker
+	host    string
+}
+
+func (c *trackedConn) Close() error {
+	c.tracker.remove(c)
+	return c.Conn.Close()
+}
+
+// trackingWriter wraps an http.ResponseWriter and intercepts Hijack calls
+// to replace the raw connection with a trackedConn that auto-deregisters.
+type trackingWriter struct {
+	http.ResponseWriter
+	tracker *HijackTracker
+	host    string
+}
+
+func (w *trackingWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	hijacker, ok := w.ResponseWriter.(http.Hijacker)
+	if !ok {
+		return nil, nil, http.ErrNotSupported
+	}
+	conn, buf, err := hijacker.Hijack()
+	if err != nil {
+		return nil, nil, err
+	}
+	tc := &trackedConn{Conn: conn, tracker: w.tracker, host: w.host}
+	w.tracker.add(tc)
+	return tc, buf, nil
+}
+
+func (w *trackingWriter) Flush() {
+	if flusher, ok := w.ResponseWriter.(http.Flusher); ok {
+		flusher.Flush()
+	}
+}
+
+func (w *trackingWriter) Unwrap() http.ResponseWriter {
+	return w.ResponseWriter
+}

proxy/internal/conntrack/hijacked.go

@@ -0,0 +1,96 @@
+package conntrack
+
+import (
+	"net/http"
+	"sync"
+)
+
+// HijackTracker tracks connections that have been hijacked (e.g. WebSocket
+// upgrades). http.Server.Shutdown does not close hijacked connections, so
+// they must be tracked and closed explicitly during graceful shutdown.
+//
+// Connections are indexed by the request Host so they can be closed
+// per-domain when a service mapping is removed.
+//
+// Use Middleware as the outermost HTTP middleware to ensure hijacked
+// connections are tracked and automatically deregistered when closed.
+type HijackTracker struct {
+	mu    sync.Mutex
+	conns map[*trackedConn]struct{}
+}
+
+// Middleware returns an HTTP middleware that wraps the ResponseWriter so that
+// hijacked connections are tracked and automatically deregistered from the
+// tracker when closed. This should be the outermost middleware in the chain.
+func (t *HijackTracker) Middleware(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		next.ServeHTTP(&trackingWriter{
+			ResponseWriter: w,
+			tracker:        t,
+			host:           hostOnly(r.Host),
+		}, r)
+	})
+}
+
+// CloseAll closes all tracked hijacked connections and returns the count.
+func (t *HijackTracker) CloseAll() int {
+	t.mu.Lock()
+	conns := t.conns
+	t.conns = nil
+	t.mu.Unlock()
+
+	for tc := range conns {
+		_ = tc.Conn.Close()
+	}
+	return len(conns)
+}
+
+// CloseByHost closes all tracked hijacked connections for the given host
+// and returns the number of connections closed.
+func (t *HijackTracker) CloseByHost(host string) int {
+	host = hostOnly(host)
+	t.mu.Lock()
+	var toClose []*trackedConn
+	for tc := range t.conns {
+		if tc.host == host {
+			toClose = append(toClose, tc)
+		}
+	}
+	for _, tc := range toClose {
+		delete(t.conns, tc)
+	}
+	t.mu.Unlock()
+
+	for _, tc := range toClose {
+		_ = tc.Conn.Close()
+	}
+	return len(toClose)
+}
+
+func (t *HijackTracker) add(tc *trackedConn) {
+	t.mu.Lock()
+	if t.conns == nil {
+		t.conns = make(map[*trackedConn]struct{})
+	}
+	t.conns[tc] = struct{}{}
+	t.mu.Unlock()
+}
+
+func (t *HijackTracker) remove(tc *trackedConn) {
+	t.mu.Lock()
+	delete(t.conns, tc)
+	t.mu.Unlock()
+}
+
+// hostOnly strips the port from a host:port string.
+func hostOnly(hostport string) string {
+	for i := len(hostport) - 1; i >= 0; i-- {
+		if hostport[i] == ':' {
+			return hostport[:i]
+		}
+		if hostport[i] < '0' || hostport[i] > '9' {
+			return hostport
+		}
+	}
+	return hostport
+}

proxy/internal/conntrack/hijacked_test.go

@@ -0,0 +1,142 @@
+package conntrack
+
+import (
+	"bufio"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// fakeHijackWriter implements http.ResponseWriter and http.Hijacker for testing.
+type fakeHijackWriter struct {
+	http.ResponseWriter
+	conn net.Conn
+}
+
+func (f *fakeHijackWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	rw := bufio.NewReadWriter(bufio.NewReader(f.conn), bufio.NewWriter(f.conn))
+	return f.conn, rw, nil
+}
+
+func TestCloseByHost(t *testing.T) {
+	var tracker HijackTracker
+
+	// Simulate hijacking two connections for different hosts.
+	connA1, connA2 := net.Pipe()
+	defer connA2.Close()
+	connB1, connB2 := net.Pipe()
+	defer connB2.Close()
+
+	twA := &trackingWriter{
+		ResponseWriter: httptest.NewRecorder(),
+		tracker:        &tracker,
+		host:           "a.example.com",
+	}
+	twB := &trackingWriter{
+		ResponseWriter: httptest.NewRecorder(),
+		tracker:        &tracker,
+		host:           "b.example.com",
+	}
+
+	// Use fakeHijackWriter to provide the Hijack method.
+	twA.ResponseWriter = &fakeHijackWriter{ResponseWriter: twA.ResponseWriter, conn: connA1}
+	twB.ResponseWriter = &fakeHijackWriter{ResponseWriter: twB.ResponseWriter, conn: connB1}
+
+	_, _, err := twA.Hijack()
+	require.NoError(t, err)
+	_, _, err = twB.Hijack()
+	require.NoError(t, err)
+
+	tracker.mu.Lock()
+	assert.Equal(t, 2, len(tracker.conns), "should track 2 connections")
+	tracker.mu.Unlock()
+
+	// Close only host A.
+	n := tracker.CloseByHost("a.example.com")
+	assert.Equal(t, 1, n, "should close 1 connection for host A")
+
+	tracker.mu.Lock()
+	assert.Equal(t, 1, len(tracker.conns), "should have 1 remaining connection")
+	tracker.mu.Unlock()
+
+	// Verify host A's conn is actually closed.
+	buf := make([]byte, 1)
+	_, err = connA2.Read(buf)
+	assert.Error(t, err, "host A pipe should be closed")
+
+	// Host B should still be alive.
+	go func() { _, _ = connB1.Write([]byte("x")) }()
+
+	// Close all remaining.
+	n = tracker.CloseAll()
+	assert.Equal(t, 1, n, "should close remaining 1 connection")
+
+	tracker.mu.Lock()
+	assert.Equal(t, 0, len(tracker.conns), "should have 0 connections after CloseAll")
+	tracker.mu.Unlock()
+}
+
+func TestCloseAll(t *testing.T) {
+	var tracker HijackTracker
+
+	for range 5 {
+		c1, c2 := net.Pipe()
+		defer c2.Close()
+		tc := &trackedConn{Conn: c1, tracker: &tracker, host: "test.com"}
+		tracker.add(tc)
+	}
+
+	tracker.mu.Lock()
+	assert.Equal(t, 5, len(tracker.conns))
+	tracker.mu.Unlock()
+
+	n := tracker.CloseAll()
+	assert.Equal(t, 5, n)
+
+	// Double CloseAll is safe.
+	n = tracker.CloseAll()
+	assert.Equal(t, 0, n)
+}
+
+func TestTrackedConn_AutoDeregister(t *testing.T) {
+	var tracker HijackTracker
+
+	c1, c2 := net.Pipe()
+	defer c2.Close()
+
+	tc := &trackedConn{Conn: c1, tracker: &tracker, host: "auto.com"}
+	tracker.add(tc)
+
+	tracker.mu.Lock()
+	assert.Equal(t, 1, len(tracker.conns))
+	tracker.mu.Unlock()
+
+	// Close the tracked conn: should auto-deregister.
+	require.NoError(t, tc.Close())
+
+	tracker.mu.Lock()
+	assert.Equal(t, 0, len(tracker.conns), "should auto-deregister on close")
+	tracker.mu.Unlock()
+}
+
+func TestHostOnly(t *testing.T) {
+	tests := []struct {
+		input string
+		want  string
+	}{
+		{"example.com:443", "example.com"},
+		{"example.com", "example.com"},
+		{"127.0.0.1:8080", "127.0.0.1"},
+		{"[::1]:443", "[::1]"},
+		{"", ""},
+	}
+	for _, tt := range tests {
+		t.Run(tt.input, func(t *testing.T) {
+			assert.Equal(t, tt.want, hostOnly(tt.input))
+		})
+	}
+}

proxy/internal/crowdsec/bouncer.go

@@ -0,0 +1,251 @@
+// Package crowdsec provides a CrowdSec stream bouncer that maintains a local
+// decision cache for IP reputation checks.
+package crowdsec
+
+import (
+	"context"
+	"errors"
+	"net/netip"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"github.com/crowdsecurity/crowdsec/pkg/models"
+	csbouncer "github.com/crowdsecurity/go-cs-bouncer"
+	log "github.com/sirupsen/logrus"
+
+	"github.com/netbirdio/netbird/proxy/internal/restrict"
+)
+
+// Bouncer wraps a CrowdSec StreamBouncer, maintaining a local cache of
+// active decisions for fast IP lookups. It implements restrict.CrowdSecChecker.
+type Bouncer struct {
+	mu       sync.RWMutex
+	ips      map[netip.Addr]*restrict.CrowdSecDecision
+	prefixes map[netip.Prefix]*restrict.CrowdSecDecision
+	ready    atomic.Bool
+
+	apiURL         string
+	apiKey         string
+	tickerInterval time.Duration
+	logger         *log.Entry
+
+	// lifeMu protects cancel and done from concurrent Start/Stop calls.
+	lifeMu sync.Mutex
+	cancel context.CancelFunc
+	done   chan struct{}
+}
+
+// compile-time check
+var _ restrict.CrowdSecChecker = (*Bouncer)(nil)
+
+// NewBouncer creates a bouncer but does not start the stream.
+func NewBouncer(apiURL, apiKey string, logger *log.Entry) *Bouncer {
+	return &Bouncer{
+		apiURL:   apiURL,
+		apiKey:   apiKey,
+		logger:   logger,
+		ips:      make(map[netip.Addr]*restrict.CrowdSecDecision),
+		prefixes: make(map[netip.Prefix]*restrict.CrowdSecDecision),
+	}
+}
+
+// Start launches the background goroutine that streams decisions from the
+// CrowdSec LAPI. The stream runs until Stop is called or ctx is cancelled.
+func (b *Bouncer) Start(ctx context.Context) error {
+	interval := b.tickerInterval
+	if interval == 0 {
+		interval = 10 * time.Second
+	}
+	stream := &csbouncer.StreamBouncer{
+		APIKey:              b.apiKey,
+		APIUrl:              b.apiURL,
+		TickerInterval:      interval.String(),
+		UserAgent:           "netbird-proxy/1.0",
+		Scopes:              []string{"ip", "range"},
+		RetryInitialConnect: true,
+	}
+
+	b.logger.Infof("connecting to CrowdSec LAPI at %s", b.apiURL)
+
+	if err := stream.Init(); err != nil {
+		return err
+	}
+
+	// Reset state from any previous run.
+	b.mu.Lock()
+	b.ips = make(map[netip.Addr]*restrict.CrowdSecDecision)
+	b.prefixes = make(map[netip.Prefix]*restrict.CrowdSecDecision)
+	b.mu.Unlock()
+	b.ready.Store(false)
+
+	ctx, cancel := context.WithCancel(ctx)
+	done := make(chan struct{})
+
+	b.lifeMu.Lock()
+	if b.cancel != nil {
+		b.lifeMu.Unlock()
+		cancel()
+		return errors.New("bouncer already started")
+	}
+	b.cancel = cancel
+	b.done = done
+	b.lifeMu.Unlock()
+
+	var wg sync.WaitGroup
+	wg.Add(2)
+
+	go func() {
+		defer wg.Done()
+		if err := stream.Run(ctx); err != nil && ctx.Err() == nil {
+			b.logger.Errorf("CrowdSec stream ended: %v", err)
+		}
+	}()
+
+	go func() {
+		defer wg.Done()
+		b.consumeStream(ctx, stream)
+	}()
+
+	go func() {
+		wg.Wait()
+		close(done)
+	}()
+
+	return nil
+}
+
+// Stop cancels the stream and waits for all goroutines to finish.
+func (b *Bouncer) Stop() {
+	b.lifeMu.Lock()
+	cancel := b.cancel
+	done := b.done
+	b.cancel = nil
+	b.lifeMu.Unlock()
+
+	if cancel != nil {
+		cancel()
+		<-done
+	}
+}
+
+// Ready returns true after the first batch of decisions has been processed.
+func (b *Bouncer) Ready() bool {
+	return b.ready.Load()
+}
+
+// CheckIP looks up addr in the local decision cache. Returns nil if no
+// active decision exists for the address.
+//
+// Prefix lookups are O(1): instead of scanning all stored prefixes, we
+// probe the map for every possible containing prefix of the address
+// (at most 33 for IPv4, 129 for IPv6).
+func (b *Bouncer) CheckIP(addr netip.Addr) *restrict.CrowdSecDecision {
+	addr = addr.Unmap()
+
+	b.mu.RLock()
+	defer b.mu.RUnlock()
+
+	if d, ok := b.ips[addr]; ok {
+		return d
+	}
+
+	maxBits := 32
+	if addr.Is6() {
+		maxBits = 128
+	}
+	// Walk from most-specific to least-specific prefix so the narrowest
+	// matching decision wins when ranges overlap.
+	for bits := maxBits; bits >= 0; bits-- {
+		prefix := netip.PrefixFrom(addr, bits).Masked()
+		if d, ok := b.prefixes[prefix]; ok {
+			return d
+		}
+	}
+
+	return nil
+}
+
+func (b *Bouncer) consumeStream(ctx context.Context, stream *csbouncer.StreamBouncer) {
+	first := true
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case resp, ok := <-stream.Stream:
+			if !ok {
+				return
+			}
+			b.mu.Lock()
+			b.applyDeleted(resp.Deleted)
+			b.applyNew(resp.New)
+			b.mu.Unlock()
+
+			if first {
+				b.ready.Store(true)
+				b.logger.Info("CrowdSec bouncer synced initial decisions")
+				first = false
+			}
+		}
+	}
+}
+
+func (b *Bouncer) applyDeleted(decisions []*models.Decision) {
+	for _, d := range decisions {
+		if d.Value == nil || d.Scope == nil {
+			continue
+		}
+		value := *d.Value
+
+		if strings.ToLower(*d.Scope) == "range" || strings.Contains(value, "/") {
+			prefix, err := netip.ParsePrefix(value)
+			if err != nil {
+				b.logger.Debugf("skip unparsable CrowdSec range deletion %q: %v", value, err)
+				continue
+			}
+			prefix = normalizePrefix(prefix)
+			delete(b.prefixes, prefix)
+		} else {
+			addr, err := netip.ParseAddr(value)
+			if err != nil {
+				b.logger.Debugf("skip unparsable CrowdSec IP deletion %q: %v", value, err)
+				continue
+			}
+			delete(b.ips, addr.Unmap())
+		}
+	}
+}
+
+func (b *Bouncer) applyNew(decisions []*models.Decision) {
+	for _, d := range decisions {
+		if d.Value == nil || d.Type == nil || d.Scope == nil {
+			continue
+		}
+		dec := &restrict.CrowdSecDecision{Type: restrict.DecisionType(*d.Type)}
+		value := *d.Value
+
+		if strings.ToLower(*d.Scope) == "range" || strings.Contains(value, "/") {
+			prefix, err := netip.ParsePrefix(value)
+			if err != nil {
+				b.logger.Debugf("skip unparsable CrowdSec range %q: %v", value, err)
+				continue
+			}
+			prefix = normalizePrefix(prefix)
+			b.prefixes[prefix] = dec
+		} else {
+			addr, err := netip.ParseAddr(value)
+			if err != nil {
+				b.logger.Debugf("skip unparsable CrowdSec IP %q: %v", value, err)
+				continue
+			}
+			b.ips[addr.Unmap()] = dec
+		}
+	}
+}
+
+// normalizePrefix unmaps v4-mapped-v6 addresses and zeros host bits so
+// the prefix is a valid map key that matches CheckIP's probe logic.
+func normalizePrefix(p netip.Prefix) netip.Prefix {
+	return netip.PrefixFrom(p.Addr().Unmap(), p.Bits()).Masked()
+}

proxy/internal/crowdsec/bouncer_test.go

@@ -0,0 +1,337 @@
+package crowdsec
+
+import (
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"net/netip"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/crowdsecurity/crowdsec/pkg/models"
+	log "github.com/sirupsen/logrus"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/proxy/internal/restrict"
+)
+
+func TestBouncer_CheckIP_Empty(t *testing.T) {
+	b := newTestBouncer()
+	b.ready.Store(true)
+
+	assert.Nil(t, b.CheckIP(netip.MustParseAddr("1.2.3.4")))
+}
+
+func TestBouncer_CheckIP_ExactMatch(t *testing.T) {
+	b := newTestBouncer()
+	b.ready.Store(true)
+	b.ips[netip.MustParseAddr("10.0.0.1")] = &restrict.CrowdSecDecision{Type: restrict.DecisionBan}
+
+	d := b.CheckIP(netip.MustParseAddr("10.0.0.1"))
+	require.NotNil(t, d)
+	assert.Equal(t, restrict.DecisionBan, d.Type)
+
+	assert.Nil(t, b.CheckIP(netip.MustParseAddr("10.0.0.2")))
+}
+
+func TestBouncer_CheckIP_PrefixMatch(t *testing.T) {
+	b := newTestBouncer()
+	b.ready.Store(true)
+	b.prefixes[netip.MustParsePrefix("192.168.1.0/24")] = &restrict.CrowdSecDecision{Type: restrict.DecisionBan}
+
+	d := b.CheckIP(netip.MustParseAddr("192.168.1.100"))
+	require.NotNil(t, d)
+	assert.Equal(t, restrict.DecisionBan, d.Type)
+
+	assert.Nil(t, b.CheckIP(netip.MustParseAddr("192.168.2.1")))
+}
+
+func TestBouncer_CheckIP_UnmapsV4InV6(t *testing.T) {
+	b := newTestBouncer()
+	b.ready.Store(true)
+	b.ips[netip.MustParseAddr("10.0.0.1")] = &restrict.CrowdSecDecision{Type: restrict.DecisionBan}
+
+	d := b.CheckIP(netip.MustParseAddr("::ffff:10.0.0.1"))
+	require.NotNil(t, d)
+	assert.Equal(t, restrict.DecisionBan, d.Type)
+}
+
+func TestBouncer_Ready(t *testing.T) {
+	b := newTestBouncer()
+	assert.False(t, b.Ready())
+
+	b.ready.Store(true)
+	assert.True(t, b.Ready())
+}
+
+func TestBouncer_CheckIP_ExactBeforePrefix(t *testing.T) {
+	b := newTestBouncer()
+	b.ready.Store(true)
+	b.ips[netip.MustParseAddr("10.0.0.1")] = &restrict.CrowdSecDecision{Type: restrict.DecisionCaptcha}
+	b.prefixes[netip.MustParsePrefix("10.0.0.0/8")] = &restrict.CrowdSecDecision{Type: restrict.DecisionBan}
+
+	d := b.CheckIP(netip.MustParseAddr("10.0.0.1"))
+	require.NotNil(t, d)
+	assert.Equal(t, restrict.DecisionCaptcha, d.Type)
+
+	d2 := b.CheckIP(netip.MustParseAddr("10.0.0.2"))
+	require.NotNil(t, d2)
+	assert.Equal(t, restrict.DecisionBan, d2.Type)
+}
+
+func TestBouncer_ApplyNew_IP(t *testing.T) {
+	b := newTestBouncer()
+
+	b.applyNew(makeDecisions(
+		decision{scope: "ip", value: "1.2.3.4", dtype: "ban", scenario: "test/brute"},
+		decision{scope: "ip", value: "5.6.7.8", dtype: "captcha", scenario: "test/crawl"},
+	))
+
+	require.Len(t, b.ips, 2)
+	assert.Equal(t, restrict.DecisionBan, b.ips[netip.MustParseAddr("1.2.3.4")].Type)
+	assert.Equal(t, restrict.DecisionCaptcha, b.ips[netip.MustParseAddr("5.6.7.8")].Type)
+}
+
+func TestBouncer_ApplyNew_Range(t *testing.T) {
+	b := newTestBouncer()
+
+	b.applyNew(makeDecisions(
+		decision{scope: "range", value: "10.0.0.0/8", dtype: "ban"},
+	))
+
+	require.Len(t, b.prefixes, 1)
+	assert.NotNil(t, b.prefixes[netip.MustParsePrefix("10.0.0.0/8")])
+}
+
+func TestBouncer_ApplyDeleted_IP(t *testing.T) {
+	b := newTestBouncer()
+	b.ips[netip.MustParseAddr("1.2.3.4")] = &restrict.CrowdSecDecision{Type: restrict.DecisionBan}
+	b.ips[netip.MustParseAddr("5.6.7.8")] = &restrict.CrowdSecDecision{Type: restrict.DecisionBan}
+
+	b.applyDeleted(makeDecisions(
+		decision{scope: "ip", value: "1.2.3.4", dtype: "ban"},
+	))
+
+	assert.Len(t, b.ips, 1)
+	assert.Nil(t, b.ips[netip.MustParseAddr("1.2.3.4")])
+	assert.NotNil(t, b.ips[netip.MustParseAddr("5.6.7.8")])
+}
+
+func TestBouncer_ApplyDeleted_Range(t *testing.T) {
+	b := newTestBouncer()
+	b.prefixes[netip.MustParsePrefix("10.0.0.0/8")] = &restrict.CrowdSecDecision{Type: restrict.DecisionBan}
+	b.prefixes[netip.MustParsePrefix("192.168.0.0/16")] = &restrict.CrowdSecDecision{Type: restrict.DecisionBan}
+
+	b.applyDeleted(makeDecisions(
+		decision{scope: "range", value: "10.0.0.0/8", dtype: "ban"},
+	))
+
+	require.Len(t, b.prefixes, 1)
+	assert.NotNil(t, b.prefixes[netip.MustParsePrefix("192.168.0.0/16")])
+}
+
+func TestBouncer_ApplyNew_OverwritesExisting(t *testing.T) {
+	b := newTestBouncer()
+	b.ips[netip.MustParseAddr("1.2.3.4")] = &restrict.CrowdSecDecision{Type: restrict.DecisionBan}
+
+	b.applyNew(makeDecisions(
+		decision{scope: "ip", value: "1.2.3.4", dtype: "captcha"},
+	))
+
+	assert.Equal(t, restrict.DecisionCaptcha, b.ips[netip.MustParseAddr("1.2.3.4")].Type)
+}
+
+func TestBouncer_ApplyNew_SkipsInvalid(t *testing.T) {
+	b := newTestBouncer()
+
+	b.applyNew(makeDecisions(
+		decision{scope: "ip", value: "not-an-ip", dtype: "ban"},
+		decision{scope: "range", value: "also-not-valid", dtype: "ban"},
+	))
+
+	assert.Empty(t, b.ips)
+	assert.Empty(t, b.prefixes)
+}
+
+// TestBouncer_StreamIntegration tests the full flow: fake LAPI → StreamBouncer → Bouncer cache → CheckIP.
+func TestBouncer_StreamIntegration(t *testing.T) {
+	lapi := newFakeLAPI()
+	ts := httptest.NewServer(lapi)
+	defer ts.Close()
+
+	// Seed the LAPI with initial decisions.
+	lapi.setDecisions(
+		decision{scope: "ip", value: "1.2.3.4", dtype: "ban", scenario: "crowdsecurity/ssh-bf"},
+		decision{scope: "range", value: "10.0.0.0/8", dtype: "ban", scenario: "crowdsecurity/http-probing"},
+		decision{scope: "ip", value: "5.5.5.5", dtype: "captcha", scenario: "crowdsecurity/http-crawl"},
+	)
+
+	b := NewBouncer(ts.URL, "test-key", log.NewEntry(log.StandardLogger()))
+	b.tickerInterval = 200 * time.Millisecond
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	require.NoError(t, b.Start(ctx))
+	defer b.Stop()
+
+	// Wait for initial sync.
+	require.Eventually(t, b.Ready, 5*time.Second, 50*time.Millisecond, "bouncer should become ready")
+
+	// Verify decisions are cached.
+	d := b.CheckIP(netip.MustParseAddr("1.2.3.4"))
+	require.NotNil(t, d, "1.2.3.4 should be banned")
+	assert.Equal(t, restrict.DecisionBan, d.Type)
+
+	d2 := b.CheckIP(netip.MustParseAddr("10.1.2.3"))
+	require.NotNil(t, d2, "10.1.2.3 should match range ban")
+	assert.Equal(t, restrict.DecisionBan, d2.Type)
+
+	d3 := b.CheckIP(netip.MustParseAddr("5.5.5.5"))
+	require.NotNil(t, d3, "5.5.5.5 should have captcha")
+	assert.Equal(t, restrict.DecisionCaptcha, d3.Type)
+
+	assert.Nil(t, b.CheckIP(netip.MustParseAddr("9.9.9.9")), "unknown IP should be nil")
+
+	// Simulate a delta update: delete one IP, add a new one.
+	lapi.setDelta(
+		[]decision{{scope: "ip", value: "1.2.3.4", dtype: "ban"}},
+		[]decision{{scope: "ip", value: "2.3.4.5", dtype: "throttle", scenario: "crowdsecurity/http-flood"}},
+	)
+
+	// Wait for the delta to be picked up.
+	require.Eventually(t, func() bool {
+		return b.CheckIP(netip.MustParseAddr("2.3.4.5")) != nil
+	}, 5*time.Second, 50*time.Millisecond, "new decision should appear")
+
+	assert.Nil(t, b.CheckIP(netip.MustParseAddr("1.2.3.4")), "deleted decision should be gone")
+
+	d4 := b.CheckIP(netip.MustParseAddr("2.3.4.5"))
+	require.NotNil(t, d4)
+	assert.Equal(t, restrict.DecisionThrottle, d4.Type)
+
+	// Range ban should still be active.
+	assert.NotNil(t, b.CheckIP(netip.MustParseAddr("10.99.99.99")))
+}
+
+// Helpers
+
+func newTestBouncer() *Bouncer {
+	return &Bouncer{
+		ips:      make(map[netip.Addr]*restrict.CrowdSecDecision),
+		prefixes: make(map[netip.Prefix]*restrict.CrowdSecDecision),
+		logger:   log.NewEntry(log.StandardLogger()),
+	}
+}
+
+type decision struct {
+	scope    string
+	value    string
+	dtype    string
+	scenario string
+}
+
+func makeDecisions(decs ...decision) []*models.Decision {
+	out := make([]*models.Decision, len(decs))
+	for i, d := range decs {
+		out[i] = &models.Decision{
+			Scope:    strPtr(d.scope),
+			Value:    strPtr(d.value),
+			Type:     strPtr(d.dtype),
+			Scenario: strPtr(d.scenario),
+			Duration: strPtr("1h"),
+			Origin:   strPtr("cscli"),
+		}
+	}
+	return out
+}
+
+func strPtr(s string) *string { return &s }
+
+// fakeLAPI is a minimal fake CrowdSec LAPI that serves /v1/decisions/stream.
+type fakeLAPI struct {
+	mu       sync.Mutex
+	initial  []decision
+	newDelta []decision
+	delDelta []decision
+	served   bool // true after the initial snapshot has been served
+}
+
+func newFakeLAPI() *fakeLAPI {
+	return &fakeLAPI{}
+}
+
+func (f *fakeLAPI) setDecisions(decs ...decision) {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	f.initial = decs
+	f.served = false
+}
+
+func (f *fakeLAPI) setDelta(deleted, added []decision) {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	f.delDelta = deleted
+	f.newDelta = added
+}
+
+func (f *fakeLAPI) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	if r.URL.Path != "/v1/decisions/stream" {
+		http.NotFound(w, r)
+		return
+	}
+
+	f.mu.Lock()
+	defer f.mu.Unlock()
+
+	resp := streamResponse{}
+
+	if !f.served {
+		for _, d := range f.initial {
+			resp.New = append(resp.New, toLAPIDecision(d))
+		}
+		f.served = true
+	} else {
+		for _, d := range f.delDelta {
+			resp.Deleted = append(resp.Deleted, toLAPIDecision(d))
+		}
+		for _, d := range f.newDelta {
+			resp.New = append(resp.New, toLAPIDecision(d))
+		}
+		// Clear delta after serving once.
+		f.delDelta = nil
+		f.newDelta = nil
+	}
+
+	w.Header().Set("Content-Type", "application/json")
+	json.NewEncoder(w).Encode(resp) //nolint:errcheck
+}
+
+// streamResponse mirrors the CrowdSec LAPI /v1/decisions/stream JSON structure.
+type streamResponse struct {
+	New     []*lapiDecision `json:"new"`
+	Deleted []*lapiDecision `json:"deleted"`
+}
+
+type lapiDecision struct {
+	Duration *string `json:"duration"`
+	Origin   *string `json:"origin"`
+	Scenario *string `json:"scenario"`
+	Scope    *string `json:"scope"`
+	Type     *string `json:"type"`
+	Value    *string `json:"value"`
+}
+
+func toLAPIDecision(d decision) *lapiDecision {
+	return &lapiDecision{
+		Duration: strPtr("1h"),
+		Origin:   strPtr("cscli"),
+		Scenario: strPtr(d.scenario),
+		Scope:    strPtr(d.scope),
+		Type:     strPtr(d.dtype),
+		Value:    strPtr(d.value),
+	}
+}

proxy/internal/crowdsec/registry.go

@@ -0,0 +1,103 @@
+package crowdsec
+
+import (
+	"context"
+	"sync"
+
+	log "github.com/sirupsen/logrus"
+
+	"github.com/netbirdio/netbird/proxy/internal/types"
+)
+
+// Registry manages a single shared Bouncer instance with reference counting.
+// The bouncer starts when the first service acquires it and stops when the
+// last service releases it.
+type Registry struct {
+	mu      sync.Mutex
+	bouncer *Bouncer
+	refs    map[types.ServiceID]struct{}
+	apiURL  string
+	apiKey  string
+	logger  *log.Entry
+	cancel  context.CancelFunc
+}
+
+// NewRegistry creates a registry. The bouncer is not started until Acquire is called.
+func NewRegistry(apiURL, apiKey string, logger *log.Entry) *Registry {
+	return &Registry{
+		apiURL: apiURL,
+		apiKey: apiKey,
+		logger: logger,
+		refs:   make(map[types.ServiceID]struct{}),
+	}
+}
+
+// Available returns true when the LAPI URL and API key are configured.
+func (r *Registry) Available() bool {
+	return r.apiURL != "" && r.apiKey != ""
+}
+
+// Acquire registers svcID as a consumer and starts the bouncer if this is the
+// first consumer. Returns the shared Bouncer (which implements the restrict
+// package's CrowdSecChecker interface). Returns nil if not Available.
+func (r *Registry) Acquire(svcID types.ServiceID) *Bouncer {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	if !r.Available() {
+		return nil
+	}
+
+	if _, exists := r.refs[svcID]; exists {
+		return r.bouncer
+	}
+
+	if r.bouncer == nil {
+		r.startLocked()
+	}
+
+	// startLocked may fail, leaving r.bouncer nil.
+	if r.bouncer == nil {
+		return nil
+	}
+
+	r.refs[svcID] = struct{}{}
+	return r.bouncer
+}
+
+// Release removes svcID as a consumer. Stops the bouncer when the last
+// consumer releases.
+func (r *Registry) Release(svcID types.ServiceID) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	delete(r.refs, svcID)
+
+	if len(r.refs) == 0 && r.bouncer != nil {
+		r.stopLocked()
+	}
+}
+
+func (r *Registry) startLocked() {
+	b := NewBouncer(r.apiURL, r.apiKey, r.logger)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	r.cancel = cancel
+
+	if err := b.Start(ctx); err != nil {
+		r.logger.Errorf("failed to start CrowdSec bouncer: %v", err)
+		cancel()
+		return
+	}
+
+	r.bouncer = b
+	r.logger.Info("CrowdSec bouncer started")
+}
+
+func (r *Registry) stopLocked() {
+	r.bouncer.Stop()
+	r.cancel()
+	r.bouncer = nil
+	r.cancel = nil
+	r.logger.Info("CrowdSec bouncer stopped")
+}

proxy/internal/crowdsec/registry_test.go

@@ -0,0 +1,66 @@
+package crowdsec
+
+import (
+	"testing"
+
+	log "github.com/sirupsen/logrus"
+	"github.com/stretchr/testify/assert"
+
+	"github.com/netbirdio/netbird/proxy/internal/types"
+)
+
+func TestRegistry_Available(t *testing.T) {
+	r := NewRegistry("http://localhost:8080/", "test-key", log.NewEntry(log.StandardLogger()))
+	assert.True(t, r.Available())
+
+	r2 := NewRegistry("", "", log.NewEntry(log.StandardLogger()))
+	assert.False(t, r2.Available())
+
+	r3 := NewRegistry("http://localhost:8080/", "", log.NewEntry(log.StandardLogger()))
+	assert.False(t, r3.Available())
+}
+
+func TestRegistry_Acquire_NotAvailable(t *testing.T) {
+	r := NewRegistry("", "", log.NewEntry(log.StandardLogger()))
+	b := r.Acquire("svc-1")
+	assert.Nil(t, b)
+}
+
+func TestRegistry_Acquire_Idempotent(t *testing.T) {
+	r := newTestRegistry()
+
+	b1 := r.Acquire("svc-1")
+	// Can't start without a real LAPI, but we can verify the ref tracking.
+	// The bouncer will be nil because Start fails, but the ref is tracked.
+	_ = b1
+
+	assert.Len(t, r.refs, 1)
+
+	// Second acquire of same service should not add another ref.
+	r.Acquire("svc-1")
+	assert.Len(t, r.refs, 1)
+}
+
+func TestRegistry_Release_Removes(t *testing.T) {
+	r := newTestRegistry()
+	r.refs[types.ServiceID("svc-1")] = struct{}{}
+
+	r.Release("svc-1")
+	assert.Empty(t, r.refs)
+}
+
+func TestRegistry_Release_Noop(t *testing.T) {
+	r := newTestRegistry()
+	// Releasing a service that was never acquired should not panic.
+	r.Release("nonexistent")
+	assert.Empty(t, r.refs)
+}
+
+func newTestRegistry() *Registry {
+	return &Registry{
+		apiURL: "http://localhost:8080/",
+		apiKey: "test-key",
+		logger: log.NewEntry(log.StandardLogger()),
+		refs:   make(map[types.ServiceID]struct{}),
+	}
+}

proxy/internal/debug/client.go

@@ -0,0 +1,576 @@
+// Package debug provides HTTP debug endpoints and CLI client for the proxy server.
+package debug
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net"
+	"net/http"
+	"net/url"
+	"strings"
+	"time"
+)
+
+// StatusFilters contains filter options for status queries.
+type StatusFilters struct {
+	IPs            []string
+	Names          []string
+	Status         string
+	ConnectionType string
+}
+
+// Client provides CLI access to debug endpoints.
+type Client struct {
+	baseURL    string
+	jsonOutput bool
+	httpClient *http.Client
+	out        io.Writer
+}
+
+// NewClient creates a new debug client.
+func NewClient(baseURL string, jsonOutput bool, out io.Writer) *Client {
+	if !strings.HasPrefix(baseURL, "http://") && !strings.HasPrefix(baseURL, "https://") {
+		baseURL = "http://" + baseURL
+	}
+	baseURL = strings.TrimSuffix(baseURL, "/")
+
+	return &Client{
+		baseURL:    baseURL,
+		jsonOutput: jsonOutput,
+		out:        out,
+		httpClient: &http.Client{
+			Timeout: 30 * time.Second,
+		},
+	}
+}
+
+// Health fetches the health status.
+func (c *Client) Health(ctx context.Context) error {
+	return c.fetchAndPrint(ctx, "/debug/health", c.printHealth)
+}
+
+func (c *Client) printHealth(data map[string]any) {
+	_, _ = fmt.Fprintf(c.out, "Status: %v\n", data["status"])
+	_, _ = fmt.Fprintf(c.out, "Uptime: %v\n", data["uptime"])
+	_, _ = fmt.Fprintf(c.out, "Management Connected: %s\n", boolIcon(data["management_connected"]))
+	_, _ = fmt.Fprintf(c.out, "All Clients Healthy:  %s\n", boolIcon(data["all_clients_healthy"]))
+
+	total, _ := data["certs_total"].(float64)
+	ready, _ := data["certs_ready"].(float64)
+	pending, _ := data["certs_pending"].(float64)
+	failed, _ := data["certs_failed"].(float64)
+	if total > 0 {
+		_, _ = fmt.Fprintf(c.out, "Certificates:         %d ready, %d pending, %d failed (%d total)\n",
+			int(ready), int(pending), int(failed), int(total))
+	}
+	if domains, ok := data["certs_ready_domains"].([]any); ok && len(domains) > 0 {
+		_, _ = fmt.Fprintf(c.out, "  Ready:\n")
+		for _, d := range domains {
+			_, _ = fmt.Fprintf(c.out, "    %v\n", d)
+		}
+	}
+	if domains, ok := data["certs_pending_domains"].([]any); ok && len(domains) > 0 {
+		_, _ = fmt.Fprintf(c.out, "  Pending:\n")
+		for _, d := range domains {
+			_, _ = fmt.Fprintf(c.out, "    %v\n", d)
+		}
+	}
+	if domains, ok := data["certs_failed_domains"].(map[string]any); ok && len(domains) > 0 {
+		_, _ = fmt.Fprintf(c.out, "  Failed:\n")
+		for d, errMsg := range domains {
+			_, _ = fmt.Fprintf(c.out, "    %s: %v\n", d, errMsg)
+		}
+	}
+
+	c.printHealthClients(data)
+}
+
+func (c *Client) printHealthClients(data map[string]any) {
+	clients, ok := data["clients"].(map[string]any)
+	if !ok || len(clients) == 0 {
+		return
+	}
+
+	_, _ = fmt.Fprintf(c.out, "\n%-38s %-9s %-7s %-8s %-8s %-16s %s\n",
+		"ACCOUNT ID", "HEALTHY", "MGMT", "SIGNAL", "RELAYS", "PEERS (P2P/RLY)", "DEGRADED")
+	_, _ = fmt.Fprintln(c.out, strings.Repeat("-", 110))
+
+	for accountID, v := range clients {
+		ch, ok := v.(map[string]any)
+		if !ok {
+			continue
+		}
+
+		healthy := boolIcon(ch["healthy"])
+		mgmt := boolIcon(ch["management_connected"])
+		signal := boolIcon(ch["signal_connected"])
+
+		relaysConn, _ := ch["relays_connected"].(float64)
+		relaysTotal, _ := ch["relays_total"].(float64)
+		relays := fmt.Sprintf("%d/%d", int(relaysConn), int(relaysTotal))
+
+		peersConnected, _ := ch["peers_connected"].(float64)
+		peersTotal, _ := ch["peers_total"].(float64)
+		peersP2P, _ := ch["peers_p2p"].(float64)
+		peersRelayed, _ := ch["peers_relayed"].(float64)
+		peersDegraded, _ := ch["peers_degraded"].(float64)
+		peers := fmt.Sprintf("%d/%d (%d/%d)", int(peersConnected), int(peersTotal), int(peersP2P), int(peersRelayed))
+		degraded := fmt.Sprintf("%d", int(peersDegraded))
+
+		_, _ = fmt.Fprintf(c.out, "%-38s %-9s %-7s %-8s %-8s %-16s %s", accountID, healthy, mgmt, signal, relays, peers, degraded)
+		if errMsg, ok := ch["error"].(string); ok && errMsg != "" {
+			_, _ = fmt.Fprintf(c.out, "  (%s)", errMsg)
+		}
+		_, _ = fmt.Fprintln(c.out)
+	}
+}
+
+func boolIcon(v any) string {
+	b, ok := v.(bool)
+	if !ok {
+		return "?"
+	}
+	if b {
+		return "yes"
+	}
+	return "no"
+}
+
+// ListClients fetches the list of all clients.
+func (c *Client) ListClients(ctx context.Context) error {
+	return c.fetchAndPrint(ctx, "/debug/clients", c.printClients)
+}
+
+func (c *Client) printClients(data map[string]any) {
+	_, _ = fmt.Fprintf(c.out, "Uptime: %v\n", data["uptime"])
+	_, _ = fmt.Fprintf(c.out, "Clients: %v\n\n", data["client_count"])
+
+	clients, ok := data["clients"].([]any)
+	if !ok || len(clients) == 0 {
+		_, _ = fmt.Fprintln(c.out, "No clients connected.")
+		return
+	}
+
+	_, _ = fmt.Fprintf(c.out, "%-38s %-12s %-40s %s\n", "ACCOUNT ID", "AGE", "SERVICES", "HAS CLIENT")
+	_, _ = fmt.Fprintln(c.out, strings.Repeat("-", 110))
+
+	for _, item := range clients {
+		c.printClientRow(item)
+	}
+
+	c.printInboundListeners(clients)
+}
+
+func (c *Client) printInboundListeners(clients []any) {
+	type row struct {
+		accountID string
+		tunnelIP  string
+		httpsPort int
+		httpPort  int
+	}
+	var rows []row
+	for _, item := range clients {
+		client, ok := item.(map[string]any)
+		if !ok {
+			continue
+		}
+		inbound, ok := client["inbound_listener"].(map[string]any)
+		if !ok {
+			continue
+		}
+		tunnelIP, _ := inbound["tunnel_ip"].(string)
+		httpsPort, _ := inbound["https_port"].(float64)
+		httpPort, _ := inbound["http_port"].(float64)
+		accountID, _ := client["account_id"].(string)
+		rows = append(rows, row{
+			accountID: accountID,
+			tunnelIP:  tunnelIP,
+			httpsPort: int(httpsPort),
+			httpPort:  int(httpPort),
+		})
+	}
+	if len(rows) == 0 {
+		return
+	}
+
+	_, _ = fmt.Fprintln(c.out)
+	_, _ = fmt.Fprintln(c.out, "Inbound listeners (per-account):")
+	_, _ = fmt.Fprintf(c.out, "  %-38s %-20s %-7s %s\n", "ACCOUNT ID", "TUNNEL IP", "HTTPS", "HTTP")
+	_, _ = fmt.Fprintln(c.out, "  "+strings.Repeat("-", 78))
+	for _, r := range rows {
+		_, _ = fmt.Fprintf(c.out, "  %-38s %-20s %-7d %d\n", r.accountID, r.tunnelIP, r.httpsPort, r.httpPort)
+	}
+}
+
+func (c *Client) printClientRow(item any) {
+	client, ok := item.(map[string]any)
+	if !ok {
+		return
+	}
+
+	services := c.extractServiceKeys(client)
+	hasClient := "no"
+	if hc, ok := client["has_client"].(bool); ok && hc {
+		hasClient = "yes"
+	}
+
+	_, _ = fmt.Fprintf(c.out, "%-38s %-12v %s %s\n",
+		client["account_id"],
+		client["age"],
+		services,
+		hasClient,
+	)
+}
+
+func (c *Client) extractServiceKeys(client map[string]any) string {
+	d, ok := client["service_keys"].([]any)
+	if !ok || len(d) == 0 {
+		return "-"
+	}
+
+	parts := make([]string, len(d))
+	for i, key := range d {
+		parts[i] = fmt.Sprint(key)
+	}
+	return strings.Join(parts, ", ")
+}
+
+// ClientStatus fetches the status of a specific client.
+func (c *Client) ClientStatus(ctx context.Context, accountID string, filters StatusFilters) error {
+	params := url.Values{}
+	if len(filters.IPs) > 0 {
+		params.Set("filter-by-ips", strings.Join(filters.IPs, ","))
+	}
+	if len(filters.Names) > 0 {
+		params.Set("filter-by-names", strings.Join(filters.Names, ","))
+	}
+	if filters.Status != "" {
+		params.Set("filter-by-status", filters.Status)
+	}
+	if filters.ConnectionType != "" {
+		params.Set("filter-by-connection-type", filters.ConnectionType)
+	}
+
+	path := "/debug/clients/" + url.PathEscape(accountID)
+	if len(params) > 0 {
+		path += "?" + params.Encode()
+	}
+	return c.fetchAndPrint(ctx, path, c.printClientStatus)
+}
+
+func (c *Client) printClientStatus(data map[string]any) {
+	_, _ = fmt.Fprintf(c.out, "Account: %v\n", data["account_id"])
+	if inbound, ok := data["inbound_listener"].(map[string]any); ok {
+		tunnelIP, _ := inbound["tunnel_ip"].(string)
+		httpsPort, _ := inbound["https_port"].(float64)
+		httpPort, _ := inbound["http_port"].(float64)
+		_, _ = fmt.Fprintf(c.out, "Inbound listener: %s (https=%d, http=%d)\n", tunnelIP, int(httpsPort), int(httpPort))
+	}
+	_, _ = fmt.Fprintln(c.out)
+	if status, ok := data["status"].(string); ok {
+		_, _ = fmt.Fprint(c.out, status)
+	}
+}
+
+// ClientSyncResponse fetches the sync response of a specific client.
+func (c *Client) ClientSyncResponse(ctx context.Context, accountID string) error {
+	path := "/debug/clients/" + url.PathEscape(accountID) + "/syncresponse"
+	return c.fetchAndPrintJSON(ctx, path)
+}
+
+// PingTCP performs a TCP ping through a client.
+// ipVersion may be "4", "6", or "" for automatic.
+func (c *Client) PingTCP(ctx context.Context, accountID, host string, port int, timeout time.Duration, ipVersion string) error {
+	params := url.Values{}
+	params.Set("host", host)
+	params.Set("port", fmt.Sprintf("%d", port))
+	if timeout > 0 {
+		params.Set("timeout", timeout.String())
+	}
+	if ipVersion != "" {
+		params.Set("ip_version", ipVersion)
+	}
+
+	path := fmt.Sprintf("/debug/clients/%s/pingtcp?%s", url.PathEscape(accountID), params.Encode())
+	return c.fetchAndPrint(ctx, path, c.printPingResult)
+}
+
+func (c *Client) printPingResult(data map[string]any) {
+	success, _ := data["success"].(bool)
+	host := net.JoinHostPort(fmt.Sprint(data["host"]), fmt.Sprint(data["port"]))
+	if success {
+		remote, _ := data["remote"].(string)
+		if remote != "" && remote != host {
+			_, _ = fmt.Fprintf(c.out, "Success: %s (via %s)\n", host, remote)
+		} else {
+			_, _ = fmt.Fprintf(c.out, "Success: %s\n", host)
+		}
+		_, _ = fmt.Fprintf(c.out, "Latency: %v\n", data["latency"])
+	} else {
+		_, _ = fmt.Fprintf(c.out, "Failed: %s\n", host)
+		c.printError(data)
+	}
+}
+
+// SetLogLevel sets the log level of a specific client.
+func (c *Client) SetLogLevel(ctx context.Context, accountID, level string) error {
+	params := url.Values{}
+	params.Set("level", level)
+
+	path := fmt.Sprintf("/debug/clients/%s/loglevel?%s", url.PathEscape(accountID), params.Encode())
+	return c.fetchAndPrint(ctx, path, c.printLogLevelResult)
+}
+
+func (c *Client) printLogLevelResult(data map[string]any) {
+	success, _ := data["success"].(bool)
+	if success {
+		_, _ = fmt.Fprintf(c.out, "Log level set to: %v\n", data["level"])
+	} else {
+		_, _ = fmt.Fprintln(c.out, "Failed to set log level")
+		c.printError(data)
+	}
+}
+
+// PerfSet live-retunes the tunnel buffer pool cap on all running embedded
+// clients. Batch size is not live-tunable; configure it at proxy startup.
+func (c *Client) PerfSet(ctx context.Context, value uint32) error {
+	path := fmt.Sprintf("/debug/perf?value=%d", value)
+	return c.fetchAndPrint(ctx, path, c.printPerfSet)
+}
+
+func (c *Client) printPerfSet(data map[string]any) {
+	if errMsg, ok := data["error"].(string); ok && errMsg != "" {
+		c.printError(data)
+		return
+	}
+	val, _ := data["value"].(float64)
+	applied, _ := data["applied"].(float64)
+	_, _ = fmt.Fprintf(c.out, "Pool cap set to: %d\n", uint32(val))
+	_, _ = fmt.Fprintf(c.out, "Applied to %d live clients\n", int(applied))
+	if failed, ok := data["failed"].(map[string]any); ok && len(failed) > 0 {
+		_, _ = fmt.Fprintln(c.out, "Failed:")
+		for k, v := range failed {
+			_, _ = fmt.Fprintf(c.out, "  %s: %v\n", k, v)
+		}
+	}
+}
+
+// Runtime fetches runtime stats (heap, goroutines, RSS).
+func (c *Client) Runtime(ctx context.Context) error {
+	return c.fetchAndPrint(ctx, "/debug/runtime", c.printRuntime)
+}
+
+func (c *Client) printRuntime(data map[string]any) {
+	i := func(k string) uint64 {
+		v, _ := data[k].(float64)
+		return uint64(v)
+	}
+	mb := func(n uint64) string { return fmt.Sprintf("%.1f MB", float64(n)/(1<<20)) }
+
+	_, _ = fmt.Fprintf(c.out, "Uptime:       %v\n", data["uptime"])
+	_, _ = fmt.Fprintf(c.out, "Go:           %v on %d CPU (GOMAXPROCS=%d)\n", data["go_version"], uint32(i("num_cpu")), uint32(i("gomaxprocs")))
+	_, _ = fmt.Fprintf(c.out, "Goroutines:   %d\n", i("goroutines"))
+	_, _ = fmt.Fprintf(c.out, "Live objects: %d\n", i("live_objects"))
+	_, _ = fmt.Fprintf(c.out, "GC:           %d cycles, %v pause total\n", i("num_gc"), time.Duration(i("pause_total_ns")))
+	_, _ = fmt.Fprintln(c.out, "Heap:")
+	_, _ = fmt.Fprintf(c.out, "  alloc:      %s\n", mb(i("heap_alloc")))
+	_, _ = fmt.Fprintf(c.out, "  in-use:     %s\n", mb(i("heap_inuse")))
+	_, _ = fmt.Fprintf(c.out, "  idle:       %s\n", mb(i("heap_idle")))
+	_, _ = fmt.Fprintf(c.out, "  released:   %s\n", mb(i("heap_released")))
+	_, _ = fmt.Fprintf(c.out, "  sys:        %s\n", mb(i("heap_sys")))
+	_, _ = fmt.Fprintf(c.out, "Total sys:    %s\n", mb(i("sys")))
+	if _, ok := data["vm_rss"]; ok {
+		_, _ = fmt.Fprintln(c.out, "Process:")
+		_, _ = fmt.Fprintf(c.out, "  VmRSS:      %s\n", mb(i("vm_rss")))
+		_, _ = fmt.Fprintf(c.out, "  VmSize:     %s\n", mb(i("vm_size")))
+		_, _ = fmt.Fprintf(c.out, "  VmData:     %s\n", mb(i("vm_data")))
+	}
+	_, _ = fmt.Fprintf(c.out, "Clients:      %d (%d started)\n", i("clients"), i("started"))
+}
+
+// StartClient starts a specific client.
+func (c *Client) StartClient(ctx context.Context, accountID string) error {
+	path := "/debug/clients/" + url.PathEscape(accountID) + "/start"
+	return c.fetchAndPrint(ctx, path, c.printStartResult)
+}
+
+func (c *Client) printStartResult(data map[string]any) {
+	success, _ := data["success"].(bool)
+	if success {
+		_, _ = fmt.Fprintln(c.out, "Client started")
+	} else {
+		_, _ = fmt.Fprintln(c.out, "Failed to start client")
+		c.printError(data)
+	}
+}
+
+// StopClient stops a specific client.
+func (c *Client) StopClient(ctx context.Context, accountID string) error {
+	path := "/debug/clients/" + url.PathEscape(accountID) + "/stop"
+	return c.fetchAndPrint(ctx, path, c.printStopResult)
+}
+
+func (c *Client) printStopResult(data map[string]any) {
+	success, _ := data["success"].(bool)
+	if success {
+		_, _ = fmt.Fprintln(c.out, "Client stopped")
+	} else {
+		_, _ = fmt.Fprintln(c.out, "Failed to stop client")
+		c.printError(data)
+	}
+}
+
+func (c *Client) printError(data map[string]any) {
+	if errMsg, ok := data["error"].(string); ok {
+		_, _ = fmt.Fprintf(c.out, "Error: %s\n", errMsg)
+	}
+}
+
+// CaptureOptions configures a capture request.
+type CaptureOptions struct {
+	AccountID  string
+	Duration   string
+	FilterExpr string
+	Text       bool
+	Verbose    bool
+	ASCII      bool
+	Output     io.Writer
+}
+
+// Capture streams a packet capture from the debug endpoint. The response body
+// (pcap or text) is written directly to opts.Output until the server closes the
+// connection or the context is cancelled.
+func (c *Client) Capture(ctx context.Context, opts CaptureOptions) error {
+	if opts.AccountID == "" {
+		return fmt.Errorf("account ID is required")
+	}
+	if opts.Output == nil {
+		return fmt.Errorf("output writer is required")
+	}
+
+	params := url.Values{}
+	if opts.Duration != "" {
+		params.Set("duration", opts.Duration)
+	}
+	if opts.FilterExpr != "" {
+		params.Set("filter", opts.FilterExpr)
+	}
+	if opts.Text {
+		params.Set("format", "text")
+	}
+	if opts.Verbose {
+		params.Set("verbose", "true")
+	}
+	if opts.ASCII {
+		params.Set("ascii", "true")
+	}
+
+	path := fmt.Sprintf("/debug/clients/%s/capture", url.PathEscape(opts.AccountID))
+	if len(params) > 0 {
+		path += "?" + params.Encode()
+	}
+
+	fullURL := c.baseURL + path
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, fullURL, nil)
+	if err != nil {
+		return fmt.Errorf("create request: %w", err)
+	}
+
+	// Use a separate client without timeout since captures stream for their full duration.
+	httpClient := &http.Client{}
+	resp, err := httpClient.Do(req)
+	if err != nil {
+		return fmt.Errorf("request failed: %w", err)
+	}
+	defer func() { _ = resp.Body.Close() }()
+
+	if resp.StatusCode >= 400 {
+		body, _ := io.ReadAll(resp.Body)
+		return fmt.Errorf("server error (%d): %s", resp.StatusCode, strings.TrimSpace(string(body)))
+	}
+
+	_, err = io.Copy(opts.Output, resp.Body)
+	if err != nil && ctx.Err() != nil {
+		return nil
+	}
+	return err
+}
+
+func (c *Client) fetchAndPrint(ctx context.Context, path string, printer func(map[string]any)) error {
+	data, raw, err := c.fetch(ctx, path)
+	if err != nil {
+		return err
+	}
+
+	if c.jsonOutput {
+		return c.writeJSON(data)
+	}
+
+	if data != nil {
+		printer(data)
+		return nil
+	}
+
+	_, _ = fmt.Fprintln(c.out, string(raw))
+	return nil
+}
+
+func (c *Client) fetchAndPrintJSON(ctx context.Context, path string) error {
+	data, raw, err := c.fetch(ctx, path)
+	if err != nil {
+		return err
+	}
+
+	if data != nil {
+		return c.writeJSON(data)
+	}
+
+	_, _ = fmt.Fprintln(c.out, string(raw))
+	return nil
+}
+
+func (c *Client) writeJSON(data map[string]any) error {
+	enc := json.NewEncoder(c.out)
+	enc.SetIndent("", "  ")
+	return enc.Encode(data)
+}
+
+func (c *Client) fetch(ctx context.Context, path string) (map[string]any, []byte, error) {
+	fullURL := c.baseURL + path
+	if !strings.Contains(path, "format=json") {
+		if strings.Contains(path, "?") {
+			fullURL += "&format=json"
+		} else {
+			fullURL += "?format=json"
+		}
+	}
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, fullURL, nil)
+	if err != nil {
+		return nil, nil, fmt.Errorf("create request: %w", err)
+	}
+
+	resp, err := c.httpClient.Do(req)
+	if err != nil {
+		return nil, nil, fmt.Errorf("request failed: %w", err)
+	}
+	defer func() { _ = resp.Body.Close() }()
+
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return nil, nil, fmt.Errorf("read response: %w", err)
+	}
+
+	if resp.StatusCode >= 400 {
+		return nil, nil, fmt.Errorf("server error (%d): %s", resp.StatusCode, strings.TrimSpace(string(body)))
+	}
+
+	var data map[string]any
+	if err := json.Unmarshal(body, &data); err != nil {
+		return nil, body, nil
+	}
+
+	return data, body, nil
+}

proxy/internal/debug/client_test.go

@@ -0,0 +1,71 @@
+package debug
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestPrintHealth_WithCertsAndClients(t *testing.T) {
+	var buf bytes.Buffer
+	c := NewClient("localhost:8444", false, &buf)
+
+	data := map[string]any{
+		"status":                "ok",
+		"uptime":                "1h30m",
+		"management_connected":  true,
+		"all_clients_healthy":   true,
+		"certs_total":           float64(3),
+		"certs_ready":           float64(2),
+		"certs_pending":         float64(1),
+		"certs_failed":          float64(0),
+		"certs_ready_domains":   []any{"a.example.com", "b.example.com"},
+		"certs_pending_domains": []any{"c.example.com"},
+		"clients": map[string]any{
+			"acc-1": map[string]any{
+				"healthy":              true,
+				"management_connected": true,
+				"signal_connected":     true,
+				"relays_connected":     float64(1),
+				"relays_total":         float64(2),
+				"peers_connected":      float64(3),
+				"peers_total":          float64(5),
+				"peers_p2p":            float64(2),
+				"peers_relayed":        float64(1),
+				"peers_degraded":       float64(0),
+			},
+		},
+	}
+
+	c.printHealth(data)
+	out := buf.String()
+
+	assert.Contains(t, out, "Status: ok")
+	assert.Contains(t, out, "Uptime: 1h30m")
+	assert.Contains(t, out, "yes") // management_connected
+	assert.Contains(t, out, "2 ready, 1 pending, 0 failed (3 total)")
+	assert.Contains(t, out, "a.example.com")
+	assert.Contains(t, out, "c.example.com")
+	assert.Contains(t, out, "acc-1")
+}
+
+func TestPrintHealth_Minimal(t *testing.T) {
+	var buf bytes.Buffer
+	c := NewClient("localhost:8444", false, &buf)
+
+	data := map[string]any{
+		"status":               "ok",
+		"uptime":               "5m",
+		"management_connected": false,
+		"all_clients_healthy":  false,
+	}
+
+	c.printHealth(data)
+	out := buf.String()
+
+	assert.Contains(t, out, "Status: ok")
+	assert.Contains(t, out, "Uptime: 5m")
+	assert.NotContains(t, out, "Certificates")
+	assert.NotContains(t, out, "ACCOUNT ID")
+}

proxy/internal/debug/handler.go

@@ -0,0 +1,960 @@
+// Package debug provides HTTP debug endpoints for the proxy server.
+package debug
+
+import (
+	"cmp"
+	"context"
+	"embed"
+	"encoding/json"
+	"fmt"
+	"html/template"
+	"maps"
+	"net"
+	"net/http"
+	"os"
+	"runtime"
+	"slices"
+	"strconv"
+	"strings"
+	"sync"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+	"google.golang.org/protobuf/encoding/protojson"
+
+	nbembed "github.com/netbirdio/netbird/client/embed"
+	nbstatus "github.com/netbirdio/netbird/client/status"
+	"github.com/netbirdio/netbird/proxy/internal/health"
+	"github.com/netbirdio/netbird/proxy/internal/roundtrip"
+	"github.com/netbirdio/netbird/proxy/internal/types"
+	"github.com/netbirdio/netbird/version"
+)
+
+//go:embed templates/*.html
+var templateFS embed.FS
+
+const defaultPingTimeout = 10 * time.Second
+
+// formatDuration formats a duration with 2 decimal places using appropriate units.
+func formatDuration(d time.Duration) string {
+	switch {
+	case d >= time.Hour:
+		return fmt.Sprintf("%.2fh", d.Hours())
+	case d >= time.Minute:
+		return fmt.Sprintf("%.2fm", d.Minutes())
+	case d >= time.Second:
+		return fmt.Sprintf("%.2fs", d.Seconds())
+	case d >= time.Millisecond:
+		return fmt.Sprintf("%.2fms", float64(d.Microseconds())/1000)
+	case d >= time.Microsecond:
+		return fmt.Sprintf("%.2fµs", float64(d.Nanoseconds())/1000)
+	default:
+		return fmt.Sprintf("%dns", d.Nanoseconds())
+	}
+}
+
+func sortedAccountIDs(m map[types.AccountID]roundtrip.ClientDebugInfo) []types.AccountID {
+	return slices.Sorted(maps.Keys(m))
+}
+
+// clientProvider provides access to NetBird clients.
+type clientProvider interface {
+	GetClient(accountID types.AccountID) (*nbembed.Client, bool)
+	ListClientsForDebug() map[types.AccountID]roundtrip.ClientDebugInfo
+	ListClientsForStartup() map[types.AccountID]*nbembed.Client
+}
+
+// InboundListenerInfo describes a per-account inbound listener as
+// surfaced through the debug HTTP handler. Mirrors the proto sub-message
+// emitted with SendStatusUpdate so dashboards and CLI tooling see the
+// same shape.
+type InboundListenerInfo struct {
+	TunnelIP  string `json:"tunnel_ip"`
+	HTTPSPort uint16 `json:"https_port"`
+	HTTPPort  uint16 `json:"http_port"`
+}
+
+// InboundProvider exposes per-account inbound listener state. Optional;
+// when nil the debug endpoint omits the inbound section entirely so the
+// existing JSON shape stays additive.
+type InboundProvider interface {
+	InboundListeners() map[types.AccountID]InboundListenerInfo
+}
+
+// healthChecker provides health probe state.
+type healthChecker interface {
+	ReadinessProbe() bool
+	StartupProbe(ctx context.Context) bool
+	CheckClientsConnected(ctx context.Context) (bool, map[types.AccountID]health.ClientHealth)
+}
+
+type certStatus interface {
+	TotalDomains() int
+	PendingDomains() []string
+	ReadyDomains() []string
+	FailedDomains() map[string]string
+}
+
+// Handler provides HTTP debug endpoints.
+type Handler struct {
+	provider   clientProvider
+	health     healthChecker
+	certStatus certStatus
+	inbound    InboundProvider
+	logger     *log.Logger
+	startTime  time.Time
+	templates  *template.Template
+	templateMu sync.RWMutex
+}
+
+// NewHandler creates a new debug handler.
+func NewHandler(provider clientProvider, healthChecker healthChecker, logger *log.Logger) *Handler {
+	if logger == nil {
+		logger = log.StandardLogger()
+	}
+	h := &Handler{
+		provider:  provider,
+		health:    healthChecker,
+		logger:    logger,
+		startTime: time.Now(),
+	}
+	if err := h.loadTemplates(); err != nil {
+		logger.Errorf("failed to load embedded templates: %v", err)
+	}
+	return h
+}
+
+// SetCertStatus sets the certificate status provider for ACME prefetch observability.
+func (h *Handler) SetCertStatus(cs certStatus) {
+	h.certStatus = cs
+}
+
+// SetInboundProvider wires per-account inbound listener observability.
+// Pass nil (or skip the call) to keep the inbound section out of debug
+// responses on proxies that don't run --private.
+func (h *Handler) SetInboundProvider(p InboundProvider) {
+	h.inbound = p
+}
+
+func (h *Handler) loadTemplates() error {
+	tmpl, err := template.ParseFS(templateFS, "templates/*.html")
+	if err != nil {
+		return fmt.Errorf("parse embedded templates: %w", err)
+	}
+
+	h.templateMu.Lock()
+	h.templates = tmpl
+	h.templateMu.Unlock()
+
+	return nil
+}
+
+func (h *Handler) getTemplates() *template.Template {
+	h.templateMu.RLock()
+	defer h.templateMu.RUnlock()
+	return h.templates
+}
+
+// ServeHTTP handles debug requests.
+func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	path := r.URL.Path
+	wantJSON := r.URL.Query().Get("format") == "json" || strings.HasSuffix(path, "/json")
+	path = strings.TrimSuffix(path, "/json")
+
+	switch path {
+	case "/debug", "/debug/":
+		h.handleIndex(w, r, wantJSON)
+	case "/debug/clients":
+		h.handleListClients(w, r, wantJSON)
+	case "/debug/health":
+		h.handleHealth(w, r, wantJSON)
+	case "/debug/perf":
+		h.handlePerf(w, r)
+	case "/debug/runtime":
+		h.handleRuntime(w, r)
+	default:
+		if h.handleClientRoutes(w, r, path, wantJSON) {
+			return
+		}
+		http.NotFound(w, r)
+	}
+}
+
+func (h *Handler) handleClientRoutes(w http.ResponseWriter, r *http.Request, path string, wantJSON bool) bool {
+	if !strings.HasPrefix(path, "/debug/clients/") {
+		return false
+	}
+
+	rest := strings.TrimPrefix(path, "/debug/clients/")
+	parts := strings.SplitN(rest, "/", 2)
+	accountID := types.AccountID(parts[0])
+
+	if len(parts) == 1 {
+		h.handleClientStatus(w, r, accountID, wantJSON)
+		return true
+	}
+
+	switch parts[1] {
+	case "syncresponse":
+		h.handleClientSyncResponse(w, r, accountID, wantJSON)
+	case "tools":
+		h.handleClientTools(w, r, accountID)
+	case "pingtcp":
+		h.handlePingTCP(w, r, accountID)
+	case "loglevel":
+		h.handleLogLevel(w, r, accountID)
+	case "start":
+		h.handleClientStart(w, r, accountID)
+	case "stop":
+		h.handleClientStop(w, r, accountID)
+	case "capture":
+		h.handleCapture(w, r, accountID)
+	default:
+		return false
+	}
+	return true
+}
+
+type failedDomain struct {
+	Domain string
+	Error  string
+}
+
+type indexData struct {
+	Version             string
+	Uptime              string
+	ClientCount         int
+	TotalServices       int
+	CertsTotal          int
+	CertsReady          int
+	CertsPending        int
+	CertsFailed         int
+	CertsPendingDomains []string
+	CertsReadyDomains   []string
+	CertsFailedDomains  []failedDomain
+	Clients             []clientData
+}
+
+type clientData struct {
+	AccountID string
+	Services  string
+	Age       string
+	Status    string
+}
+
+func (h *Handler) handleIndex(w http.ResponseWriter, _ *http.Request, wantJSON bool) {
+	clients := h.provider.ListClientsForDebug()
+	sortedIDs := sortedAccountIDs(clients)
+
+	totalServices := 0
+	for _, info := range clients {
+		totalServices += info.ServiceCount
+	}
+
+	var certsTotal, certsReady, certsPending, certsFailed int
+	var certsPendingDomains, certsReadyDomains []string
+	var certsFailedDomains map[string]string
+	if h.certStatus != nil {
+		certsTotal = h.certStatus.TotalDomains()
+		certsPendingDomains = h.certStatus.PendingDomains()
+		certsReadyDomains = h.certStatus.ReadyDomains()
+		certsFailedDomains = h.certStatus.FailedDomains()
+		certsReady = len(certsReadyDomains)
+		certsPending = len(certsPendingDomains)
+		certsFailed = len(certsFailedDomains)
+	}
+
+	if wantJSON {
+		clientsJSON := make([]map[string]any, 0, len(clients))
+		for _, id := range sortedIDs {
+			info := clients[id]
+			clientsJSON = append(clientsJSON, map[string]any{
+				"account_id":    info.AccountID,
+				"service_count": info.ServiceCount,
+				"service_keys":  info.ServiceKeys,
+				"has_client":    info.HasClient,
+				"created_at":    info.CreatedAt,
+				"age":           time.Since(info.CreatedAt).Round(time.Second).String(),
+			})
+		}
+		resp := map[string]any{
+			"version":        version.NetbirdVersion(),
+			"uptime":         time.Since(h.startTime).Round(time.Second).String(),
+			"client_count":   len(clients),
+			"total_services": totalServices,
+			"certs_total":    certsTotal,
+			"certs_ready":    certsReady,
+			"certs_pending":  certsPending,
+			"certs_failed":   certsFailed,
+			"clients":        clientsJSON,
+		}
+		if len(certsPendingDomains) > 0 {
+			resp["certs_pending_domains"] = certsPendingDomains
+		}
+		if len(certsReadyDomains) > 0 {
+			resp["certs_ready_domains"] = certsReadyDomains
+		}
+		if len(certsFailedDomains) > 0 {
+			resp["certs_failed_domains"] = certsFailedDomains
+		}
+		h.writeJSON(w, resp)
+		return
+	}
+
+	sortedFailed := make([]failedDomain, 0, len(certsFailedDomains))
+	for d, e := range certsFailedDomains {
+		sortedFailed = append(sortedFailed, failedDomain{Domain: d, Error: e})
+	}
+	slices.SortFunc(sortedFailed, func(a, b failedDomain) int {
+		return cmp.Compare(a.Domain, b.Domain)
+	})
+
+	data := indexData{
+		Version:             version.NetbirdVersion(),
+		Uptime:              time.Since(h.startTime).Round(time.Second).String(),
+		ClientCount:         len(clients),
+		TotalServices:       totalServices,
+		CertsTotal:          certsTotal,
+		CertsReady:          certsReady,
+		CertsPending:        certsPending,
+		CertsFailed:         certsFailed,
+		CertsPendingDomains: certsPendingDomains,
+		CertsReadyDomains:   certsReadyDomains,
+		CertsFailedDomains:  sortedFailed,
+		Clients:             make([]clientData, 0, len(clients)),
+	}
+
+	for _, id := range sortedIDs {
+		info := clients[id]
+		services := strings.Join(info.ServiceKeys, ", ")
+		if services == "" {
+			services = "-"
+		}
+		status := "No client"
+		if info.HasClient {
+			status = "Active"
+		}
+		data.Clients = append(data.Clients, clientData{
+			AccountID: string(info.AccountID),
+			Services:  services,
+			Age:       time.Since(info.CreatedAt).Round(time.Second).String(),
+			Status:    status,
+		})
+	}
+
+	h.renderTemplate(w, "index", data)
+}
+
+type clientsData struct {
+	Uptime  string
+	Clients []clientData
+}
+
+func (h *Handler) handleListClients(w http.ResponseWriter, _ *http.Request, wantJSON bool) {
+	clients := h.provider.ListClientsForDebug()
+	sortedIDs := sortedAccountIDs(clients)
+
+	if wantJSON {
+		var inboundAll map[types.AccountID]InboundListenerInfo
+		if h.inbound != nil {
+			inboundAll = h.inbound.InboundListeners()
+		}
+		clientsJSON := make([]map[string]any, 0, len(clients))
+		for _, id := range sortedIDs {
+			info := clients[id]
+			row := map[string]any{
+				"account_id":    info.AccountID,
+				"service_count": info.ServiceCount,
+				"service_keys":  info.ServiceKeys,
+				"has_client":    info.HasClient,
+				"created_at":    info.CreatedAt,
+				"age":           time.Since(info.CreatedAt).Round(time.Second).String(),
+			}
+			if inb, ok := inboundAll[id]; ok {
+				row["inbound_listener"] = inb
+			}
+			clientsJSON = append(clientsJSON, row)
+		}
+		resp := map[string]any{
+			"uptime":       time.Since(h.startTime).Round(time.Second).String(),
+			"client_count": len(clients),
+			"clients":      clientsJSON,
+		}
+		if len(inboundAll) > 0 {
+			resp["inbound_listener_count"] = len(inboundAll)
+		}
+		h.writeJSON(w, resp)
+		return
+	}
+
+	data := clientsData{
+		Uptime:  time.Since(h.startTime).Round(time.Second).String(),
+		Clients: make([]clientData, 0, len(clients)),
+	}
+
+	for _, id := range sortedIDs {
+		info := clients[id]
+		services := strings.Join(info.ServiceKeys, ", ")
+		if services == "" {
+			services = "-"
+		}
+		status := "No client"
+		if info.HasClient {
+			status = "Active"
+		}
+		data.Clients = append(data.Clients, clientData{
+			AccountID: string(info.AccountID),
+			Services:  services,
+			Age:       time.Since(info.CreatedAt).Round(time.Second).String(),
+			Status:    status,
+		})
+	}
+
+	h.renderTemplate(w, "clients", data)
+}
+
+type clientDetailData struct {
+	AccountID string
+	ActiveTab string
+	Content   string
+}
+
+func (h *Handler) handleClientStatus(w http.ResponseWriter, r *http.Request, accountID types.AccountID, wantJSON bool) {
+	client, ok := h.provider.GetClient(accountID)
+	if !ok {
+		http.Error(w, "Client not found: "+string(accountID), http.StatusNotFound)
+		return
+	}
+
+	fullStatus, err := client.Status()
+	if err != nil {
+		http.Error(w, "Error getting status: "+err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	// Parse filter parameters
+	query := r.URL.Query()
+	statusFilter := query.Get("filter-by-status")
+	connectionTypeFilter := query.Get("filter-by-connection-type")
+
+	var prefixNamesFilter []string
+	var prefixNamesFilterMap map[string]struct{}
+	if names := query.Get("filter-by-names"); names != "" {
+		prefixNamesFilter = strings.Split(names, ",")
+		prefixNamesFilterMap = make(map[string]struct{})
+		for _, name := range prefixNamesFilter {
+			prefixNamesFilterMap[strings.ToLower(strings.TrimSpace(name))] = struct{}{}
+		}
+	}
+
+	var ipsFilterMap map[string]struct{}
+	if ips := query.Get("filter-by-ips"); ips != "" {
+		ipsFilterMap = make(map[string]struct{})
+		for _, ip := range strings.Split(ips, ",") {
+			ipsFilterMap[strings.TrimSpace(ip)] = struct{}{}
+		}
+	}
+
+	pbStatus := nbstatus.ToProtoFullStatus(fullStatus)
+	overview := nbstatus.ConvertToStatusOutputOverview(pbStatus, nbstatus.ConvertOptions{
+		StatusFilter:         statusFilter,
+		PrefixNamesFilter:    prefixNamesFilter,
+		PrefixNamesFilterMap: prefixNamesFilterMap,
+		IPsFilter:            ipsFilterMap,
+		ConnectionTypeFilter: connectionTypeFilter,
+	})
+
+	if wantJSON {
+		resp := map[string]any{
+			"account_id": accountID,
+			"status":     overview.FullDetailSummary(),
+		}
+		if info, ok := h.inboundInfoFor(accountID); ok {
+			resp["inbound_listener"] = info
+		}
+		h.writeJSON(w, resp)
+		return
+	}
+
+	data := clientDetailData{
+		AccountID: string(accountID),
+		ActiveTab: "status",
+		Content:   overview.FullDetailSummary(),
+	}
+
+	h.renderTemplate(w, "clientDetail", data)
+}
+
+// inboundInfoFor returns the inbound listener info for an account, or
+// ok=false when no inbound provider is wired or the account has no live
+// listener.
+func (h *Handler) inboundInfoFor(accountID types.AccountID) (InboundListenerInfo, bool) {
+	if h.inbound == nil {
+		return InboundListenerInfo{}, false
+	}
+	all := h.inbound.InboundListeners()
+	info, ok := all[accountID]
+	return info, ok
+}
+
+func (h *Handler) handleClientSyncResponse(w http.ResponseWriter, _ *http.Request, accountID types.AccountID, wantJSON bool) {
+	client, ok := h.provider.GetClient(accountID)
+	if !ok {
+		http.Error(w, "Client not found: "+string(accountID), http.StatusNotFound)
+		return
+	}
+
+	syncResp, err := client.GetLatestSyncResponse()
+	if err != nil {
+		http.Error(w, "Error getting sync response: "+err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	if syncResp == nil {
+		http.Error(w, "No sync response available for client: "+string(accountID), http.StatusNotFound)
+		return
+	}
+
+	opts := protojson.MarshalOptions{
+		EmitUnpopulated: true,
+		UseProtoNames:   true,
+		Indent:          "  ",
+		AllowPartial:    true,
+	}
+
+	jsonBytes, err := opts.Marshal(syncResp)
+	if err != nil {
+		http.Error(w, "Error marshaling sync response: "+err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	if wantJSON {
+		w.Header().Set("Content-Type", "application/json")
+		_, _ = w.Write(jsonBytes)
+		return
+	}
+
+	data := clientDetailData{
+		AccountID: string(accountID),
+		ActiveTab: "syncresponse",
+		Content:   string(jsonBytes),
+	}
+
+	h.renderTemplate(w, "clientDetail", data)
+}
+
+type toolsData struct {
+	AccountID string
+}
+
+func (h *Handler) handleClientTools(w http.ResponseWriter, _ *http.Request, accountID types.AccountID) {
+	_, ok := h.provider.GetClient(accountID)
+	if !ok {
+		http.Error(w, "Client not found: "+string(accountID), http.StatusNotFound)
+		return
+	}
+
+	data := toolsData{
+		AccountID: string(accountID),
+	}
+
+	h.renderTemplate(w, "tools", data)
+}
+
+func (h *Handler) handlePingTCP(w http.ResponseWriter, r *http.Request, accountID types.AccountID) {
+	client, ok := h.provider.GetClient(accountID)
+	if !ok {
+		h.writeJSON(w, map[string]any{"error": "client not found"})
+		return
+	}
+
+	host := r.URL.Query().Get("host")
+	portStr := r.URL.Query().Get("port")
+	if host == "" || portStr == "" {
+		h.writeJSON(w, map[string]any{"error": "host and port parameters required"})
+		return
+	}
+
+	port, err := strconv.Atoi(portStr)
+	if err != nil || port < 1 || port > 65535 {
+		h.writeJSON(w, map[string]any{"error": "invalid port"})
+		return
+	}
+
+	timeout := defaultPingTimeout
+	if t := r.URL.Query().Get("timeout"); t != "" {
+		if d, err := time.ParseDuration(t); err == nil {
+			timeout = d
+		}
+	}
+
+	network := "tcp"
+	if v := r.URL.Query().Get("ip_version"); v == "4" || v == "6" {
+		network += v
+	}
+
+	ctx, cancel := context.WithTimeout(r.Context(), timeout)
+	defer cancel()
+
+	address := net.JoinHostPort(host, strconv.Itoa(port))
+	start := time.Now()
+
+	conn, err := client.Dial(ctx, network, address)
+	if err != nil {
+		h.writeJSON(w, map[string]any{
+			"success": false,
+			"host":    host,
+			"port":    port,
+			"error":   err.Error(),
+		})
+		return
+	}
+
+	remote := conn.RemoteAddr().String()
+	if err := conn.Close(); err != nil {
+		h.logger.Debugf("close tcp ping connection: %v", err)
+	}
+
+	latency := time.Since(start)
+	h.writeJSON(w, map[string]any{
+		"success":    true,
+		"host":       host,
+		"port":       port,
+		"remote":     remote,
+		"latency_ms": latency.Milliseconds(),
+		"latency":    formatDuration(latency),
+	})
+}
+
+func (h *Handler) handleLogLevel(w http.ResponseWriter, r *http.Request, accountID types.AccountID) {
+	client, ok := h.provider.GetClient(accountID)
+	if !ok {
+		h.writeJSON(w, map[string]any{"error": "client not found"})
+		return
+	}
+
+	level := r.URL.Query().Get("level")
+	if level == "" {
+		h.writeJSON(w, map[string]any{"error": "level parameter required (trace, debug, info, warn, error)"})
+		return
+	}
+
+	if err := client.SetLogLevel(level); err != nil {
+		h.writeJSON(w, map[string]any{
+			"success": false,
+			"error":   err.Error(),
+		})
+		return
+	}
+
+	h.writeJSON(w, map[string]any{
+		"success": true,
+		"level":   level,
+	})
+}
+
+const clientActionTimeout = 30 * time.Second
+
+func (h *Handler) handleClientStart(w http.ResponseWriter, r *http.Request, accountID types.AccountID) {
+	client, ok := h.provider.GetClient(accountID)
+	if !ok {
+		h.writeJSON(w, map[string]any{"error": "client not found"})
+		return
+	}
+
+	ctx, cancel := context.WithTimeout(r.Context(), clientActionTimeout)
+	defer cancel()
+
+	if err := client.Start(ctx); err != nil {
+		h.writeJSON(w, map[string]any{
+			"success": false,
+			"error":   err.Error(),
+		})
+		return
+	}
+
+	h.writeJSON(w, map[string]any{
+		"success": true,
+		"message": "client started",
+	})
+}
+
+func (h *Handler) handleClientStop(w http.ResponseWriter, r *http.Request, accountID types.AccountID) {
+	client, ok := h.provider.GetClient(accountID)
+	if !ok {
+		h.writeJSON(w, map[string]any{"error": "client not found"})
+		return
+	}
+
+	ctx, cancel := context.WithTimeout(r.Context(), clientActionTimeout)
+	defer cancel()
+
+	if err := client.Stop(ctx); err != nil {
+		h.writeJSON(w, map[string]any{
+			"success": false,
+			"error":   err.Error(),
+		})
+		return
+	}
+
+	h.writeJSON(w, map[string]any{
+		"success": true,
+		"message": "client stopped",
+	})
+}
+
+func (h *Handler) handlePerf(w http.ResponseWriter, r *http.Request) {
+	raw := r.URL.Query().Get("value")
+	if raw == "" {
+		http.Error(w, "value parameter is required", http.StatusBadRequest)
+		return
+	}
+	n, err := strconv.ParseUint(raw, 10, 32)
+	if err != nil {
+		http.Error(w, fmt.Sprintf("invalid value %q: %v", raw, err), http.StatusBadRequest)
+		return
+	}
+
+	capN := uint32(n)
+	applied := 0
+	failed := map[string]string{}
+	for accountID, client := range h.provider.ListClientsForStartup() {
+		if err := client.SetPerformance(nbembed.Performance{PreallocatedBuffersPerPool: &capN}); err != nil {
+			failed[string(accountID)] = err.Error()
+			continue
+		}
+		applied++
+	}
+
+	resp := map[string]any{
+		"success": true,
+		"value":   capN,
+		"applied": applied,
+	}
+	if len(failed) > 0 {
+		resp["failed"] = failed
+	}
+	h.writeJSON(w, resp)
+}
+
+// handleRuntime returns cheap runtime and process stats. Safe to hit on a
+// running proxy; does not read pprof profiles.
+func (h *Handler) handleRuntime(w http.ResponseWriter, _ *http.Request) {
+	var m runtime.MemStats
+	runtime.ReadMemStats(&m)
+
+	clients := h.provider.ListClientsForDebug()
+	started := 0
+	for _, c := range clients {
+		if c.HasClient {
+			started++
+		}
+	}
+
+	resp := map[string]any{
+		"uptime":         time.Since(h.startTime).Round(time.Second).String(),
+		"goroutines":     runtime.NumGoroutine(),
+		"num_cpu":        runtime.NumCPU(),
+		"gomaxprocs":     runtime.GOMAXPROCS(0),
+		"go_version":     runtime.Version(),
+		"heap_alloc":     m.HeapAlloc,
+		"heap_inuse":     m.HeapInuse,
+		"heap_idle":      m.HeapIdle,
+		"heap_released":  m.HeapReleased,
+		"heap_sys":       m.HeapSys,
+		"sys":            m.Sys,
+		"live_objects":   m.Mallocs - m.Frees,
+		"num_gc":         m.NumGC,
+		"pause_total_ns": m.PauseTotalNs,
+		"clients":        len(clients),
+		"started":        started,
+	}
+
+	if proc := readProcStatus(); proc != nil {
+		resp["vm_rss"] = proc["VmRSS"]
+		resp["vm_size"] = proc["VmSize"]
+		resp["vm_data"] = proc["VmData"]
+	}
+
+	h.writeJSON(w, resp)
+}
+
+// readProcStatus parses /proc/self/status on Linux and returns size fields
+// in bytes. Returns nil on non-Linux or read failure.
+func readProcStatus() map[string]uint64 {
+	raw, err := os.ReadFile("/proc/self/status")
+	if err != nil {
+		return nil
+	}
+	out := map[string]uint64{}
+	for _, line := range strings.Split(string(raw), "\n") {
+		k, v, ok := strings.Cut(line, ":")
+		if !ok {
+			continue
+		}
+		if k != "VmRSS" && k != "VmSize" && k != "VmData" {
+			continue
+		}
+		fields := strings.Fields(v)
+		if len(fields) < 1 {
+			continue
+		}
+		n, err := strconv.ParseUint(fields[0], 10, 64)
+		if err != nil {
+			continue
+		}
+		// Values are reported in kB.
+		out[k] = n * 1024
+	}
+	return out
+}
+
+const maxCaptureDuration = 30 * time.Minute
+
+// handleCapture streams a pcap or text packet capture for the given client.
+//
+// Query params:
+//
+//	duration: capture duration (0 or absent = max, capped at 30m)
+//	format:   "text" for human-readable output (default: pcap)
+//	filter:   BPF-like filter expression (e.g. "host 10.0.0.1 and tcp port 443")
+func (h *Handler) handleCapture(w http.ResponseWriter, r *http.Request, accountID types.AccountID) {
+	client, ok := h.provider.GetClient(accountID)
+	if !ok {
+		http.Error(w, "client not found", http.StatusNotFound)
+		return
+	}
+
+	duration := maxCaptureDuration
+	if durationStr := r.URL.Query().Get("duration"); durationStr != "" {
+		d, err := time.ParseDuration(durationStr)
+		if err != nil {
+			http.Error(w, "invalid duration: "+err.Error(), http.StatusBadRequest)
+			return
+		}
+		if d < 0 {
+			http.Error(w, "duration must not be negative", http.StatusBadRequest)
+			return
+		}
+		if d > 0 {
+			duration = min(d, maxCaptureDuration)
+		}
+	}
+
+	filter := r.URL.Query().Get("filter")
+	wantText := r.URL.Query().Get("format") == "text"
+	verbose := r.URL.Query().Get("verbose") == "true"
+	ascii := r.URL.Query().Get("ascii") == "true"
+
+	opts := nbembed.CaptureOptions{Filter: filter, Verbose: verbose, ASCII: ascii}
+	if wantText {
+		w.Header().Set("Content-Type", "text/plain; charset=utf-8")
+		opts.TextOutput = w
+	} else {
+		w.Header().Set("Content-Type", "application/vnd.tcpdump.pcap")
+		w.Header().Set("Content-Disposition",
+			fmt.Sprintf("attachment; filename=capture-%s.pcap", accountID))
+		opts.Output = w
+	}
+
+	cs, err := client.StartCapture(opts)
+	if err != nil {
+		http.Error(w, "start capture: "+err.Error(), http.StatusServiceUnavailable)
+		return
+	}
+	defer cs.Stop()
+
+	// Flush headers after setup succeeds so errors above can still set status codes.
+	if f, ok := w.(http.Flusher); ok {
+		f.Flush()
+	}
+
+	timer := time.NewTimer(duration)
+	defer timer.Stop()
+
+	select {
+	case <-r.Context().Done():
+	case <-timer.C:
+	}
+
+	cs.Stop()
+
+	stats := cs.Stats()
+	h.logger.Infof("capture for %s finished: %d packets, %d bytes, %d dropped",
+		accountID, stats.Packets, stats.Bytes, stats.Dropped)
+}
+
+func (h *Handler) handleHealth(w http.ResponseWriter, r *http.Request, wantJSON bool) {
+	if !wantJSON {
+		http.Redirect(w, r, "/debug", http.StatusSeeOther)
+		return
+	}
+
+	uptime := time.Since(h.startTime).Round(10 * time.Millisecond).String()
+
+	ready := h.health.ReadinessProbe()
+	allHealthy, clientHealth := h.health.CheckClientsConnected(r.Context())
+
+	status := "ok"
+	// No clients is not a health issue; only degrade when actual clients are unhealthy
+	if !ready || (!allHealthy && len(clientHealth) > 0) {
+		status = "degraded"
+	}
+
+	var certsTotal, certsReady, certsPending, certsFailed int
+	var certsPendingDomains, certsReadyDomains []string
+	var certsFailedDomains map[string]string
+	if h.certStatus != nil {
+		certsTotal = h.certStatus.TotalDomains()
+		certsPendingDomains = h.certStatus.PendingDomains()
+		certsReadyDomains = h.certStatus.ReadyDomains()
+		certsFailedDomains = h.certStatus.FailedDomains()
+		certsReady = len(certsReadyDomains)
+		certsPending = len(certsPendingDomains)
+		certsFailed = len(certsFailedDomains)
+	}
+
+	resp := map[string]any{
+		"status":               status,
+		"uptime":               uptime,
+		"management_connected": ready,
+		"all_clients_healthy":  allHealthy,
+		"certs_total":          certsTotal,
+		"certs_ready":          certsReady,
+		"certs_pending":        certsPending,
+		"certs_failed":         certsFailed,
+		"clients":              clientHealth,
+	}
+	if len(certsPendingDomains) > 0 {
+		resp["certs_pending_domains"] = certsPendingDomains
+	}
+	if len(certsReadyDomains) > 0 {
+		resp["certs_ready_domains"] = certsReadyDomains
+	}
+	if len(certsFailedDomains) > 0 {
+		resp["certs_failed_domains"] = certsFailedDomains
+	}
+	h.writeJSON(w, resp)
+}
+
+func (h *Handler) renderTemplate(w http.ResponseWriter, name string, data any) {
+	w.Header().Set("Content-Type", "text/html; charset=utf-8")
+	tmpl := h.getTemplates()
+	if tmpl == nil {
+		http.Error(w, "Templates not loaded", http.StatusInternalServerError)
+		return
+	}
+	if err := tmpl.ExecuteTemplate(w, name, data); err != nil {
+		h.logger.Errorf("execute template %s: %v", name, err)
+		http.Error(w, "Template error", http.StatusInternalServerError)
+	}
+}
+
+func (h *Handler) writeJSON(w http.ResponseWriter, v any) {
+	w.Header().Set("Content-Type", "application/json")
+	enc := json.NewEncoder(w)
+	enc.SetIndent("", "  ")
+	if err := enc.Encode(v); err != nil {
+		h.logger.Errorf("encode JSON response: %v", err)
+	}
+}

proxy/internal/debug/templates/base.html

@@ -0,0 +1,101 @@
+{{define "style"}}
+body {
+    font-family: monospace;
+    margin: 20px;
+    background: #1a1a1a;
+    color: #eee;
+}
+a {
+    color: #6cf;
+}
+h1, h2, h3 {
+    color: #fff;
+}
+.info {
+    color: #aaa;
+}
+table {
+    border-collapse: collapse;
+    margin: 10px 0;
+}
+th, td {
+    border: 1px solid #444;
+    padding: 8px;
+    text-align: left;
+}
+th {
+    background: #333;
+}
+.nav {
+    margin-bottom: 20px;
+}
+.nav a {
+    margin-right: 15px;
+    padding: 8px 16px;
+    background: #333;
+    text-decoration: none;
+    border-radius: 4px;
+}
+.nav a.active {
+    background: #6cf;
+    color: #000;
+}
+pre {
+    background: #222;
+    padding: 15px;
+    border-radius: 4px;
+    overflow-x: auto;
+    white-space: pre-wrap;
+}
+input, select, textarea {
+    background: #333;
+    color: #eee;
+    border: 1px solid #555;
+    padding: 8px;
+    border-radius: 4px;
+    font-family: monospace;
+}
+input:focus, select:focus, textarea:focus {
+    outline: none;
+    border-color: #6cf;
+}
+button {
+    background: #6cf;
+    color: #000;
+    border: none;
+    padding: 8px 16px;
+    border-radius: 4px;
+    cursor: pointer;
+    font-family: monospace;
+}
+button:hover {
+    background: #5be;
+}
+button:disabled {
+    background: #555;
+    color: #888;
+    cursor: not-allowed;
+}
+.form-group {
+    margin-bottom: 15px;
+}
+.form-group label {
+    display: block;
+    margin-bottom: 5px;
+    color: #aaa;
+}
+.form-row {
+    display: flex;
+    gap: 10px;
+    align-items: flex-end;
+}
+.result {
+    margin-top: 20px;
+}
+.success {
+    color: #5f5;
+}
+.error {
+    color: #f55;
+}
+{{end}}

proxy/internal/debug/templates/client_detail.html

@@ -0,0 +1,19 @@
+{{define "clientDetail"}}
+<!DOCTYPE html>
+<html lang="en">
+<head>
+  <title>Client {{.AccountID}}</title>
+  <style>{{template "style"}}</style>
+</head>
+<body>
+  <h1>Client: {{.AccountID}}</h1>
+  <div class="nav">
+    <a href="/debug">&larr; Back</a>
+    <a href="/debug/clients/{{.AccountID}}/tools"{{if eq .ActiveTab "tools"}} class="active"{{end}}>Tools</a>
+    <a href="/debug/clients/{{.AccountID}}"{{if eq .ActiveTab "status"}} class="active"{{end}}>Status</a>
+    <a href="/debug/clients/{{.AccountID}}/syncresponse"{{if eq .ActiveTab "syncresponse"}} class="active"{{end}}>Sync Response</a>
+  </div>
+  <pre>{{.Content}}</pre>
+</body>
+</html>
+{{end}}

proxy/internal/debug/templates/clients.html

@@ -0,0 +1,33 @@
+{{define "clients"}}
+<!DOCTYPE html>
+<html lang="en">
+<head>
+  <title>Clients</title>
+  <style>{{template "style"}}</style>
+</head>
+<body>
+  <h1>All Clients</h1>
+  <p class="info">Uptime: {{.Uptime}} | <a href="/debug">&larr; Back</a></p>
+  {{if .Clients}}
+  <table>
+    <tr>
+      <th>Account ID</th>
+      <th>Services</th>
+      <th>Age</th>
+      <th>Status</th>
+    </tr>
+    {{range .Clients}}
+    <tr>
+      <td><a href="/debug/clients/{{.AccountID}}/tools">{{.AccountID}}</a></td>
+      <td>{{.Services}}</td>
+      <td>{{.Age}}</td>
+      <td>{{.Status}}</td>
+    </tr>
+    {{end}}
+  </table>
+  {{else}}
+  <p>No clients connected</p>
+  {{end}}
+</body>
+</html>
+{{end}}

proxy/internal/debug/templates/index.html

@@ -0,0 +1,58 @@
+{{define "index"}}
+<!DOCTYPE html>
+<html lang="en">
+<head>
+  <title>NetBird Proxy Debug</title>
+  <style>{{template "style"}}</style>
+</head>
+<body>
+  <h1>NetBird Proxy Debug</h1>
+  <p class="info">Version: {{.Version}} | Uptime: {{.Uptime}}</p>
+  <h2>Certificates: {{.CertsReady}} ready, {{.CertsPending}} pending, {{.CertsFailed}} failed ({{.CertsTotal}} total)</h2>
+  {{if .CertsReadyDomains}}
+  <details>
+    <summary>Ready domains ({{.CertsReady}})</summary>
+    <ul>{{range .CertsReadyDomains}}<li>{{.}}</li>{{end}}</ul>
+  </details>
+  {{end}}
+  {{if .CertsPendingDomains}}
+  <details open>
+    <summary>Pending domains ({{.CertsPending}})</summary>
+    <ul>{{range .CertsPendingDomains}}<li>{{.}}</li>{{end}}</ul>
+  </details>
+  {{end}}
+  {{if .CertsFailedDomains}}
+  <details open>
+    <summary>Failed domains ({{.CertsFailed}})</summary>
+    <ul>{{range .CertsFailedDomains}}<li>{{.Domain}}: {{.Error}}</li>{{end}}</ul>
+  </details>
+  {{end}}
+  <h2>Clients ({{.ClientCount}}) | Services ({{.TotalServices}})</h2>
+  {{if .Clients}}
+  <table>
+    <tr>
+      <th>Account ID</th>
+      <th>Services</th>
+      <th>Age</th>
+      <th>Status</th>
+    </tr>
+    {{range .Clients}}
+    <tr>
+      <td><a href="/debug/clients/{{.AccountID}}/tools">{{.AccountID}}</a></td>
+      <td>{{.Services}}</td>
+      <td>{{.Age}}</td>
+      <td>{{.Status}}</td>
+    </tr>
+    {{end}}
+  </table>
+  {{else}}
+  <p>No clients connected</p>
+  {{end}}
+  <h2>Endpoints</h2>
+  <ul>
+    <li><a href="/debug/clients">/debug/clients</a> - all clients detail</li>
+  </ul>
+  <p class="info">Add ?format=json or /json suffix for JSON output</p>
+</body>
+</html>
+{{end}}

proxy/internal/debug/templates/tools.html

@@ -0,0 +1,142 @@
+{{define "tools"}}
+<!DOCTYPE html>
+<html lang="en">
+<head>
+  <title>Client {{.AccountID}} - Tools</title>
+  <style>{{template "style"}}</style>
+</head>
+<body>
+  <h1>Client: {{.AccountID}}</h1>
+  <div class="nav">
+    <a href="/debug">&larr; Back</a>
+    <a href="/debug/clients/{{.AccountID}}/tools" class="active">Tools</a>
+    <a href="/debug/clients/{{.AccountID}}">Status</a>
+    <a href="/debug/clients/{{.AccountID}}/syncresponse">Sync Response</a>
+  </div>
+
+  <h2>Client Control</h2>
+  <div class="form-row">
+    <div class="form-group">
+      <span>&nbsp;</span>
+      <button onclick="startClient()">Start</button>
+    </div>
+    <div class="form-group">
+      <span>&nbsp;</span>
+      <button onclick="stopClient()">Stop</button>
+    </div>
+  </div>
+  <div id="client-result" class="result"></div>
+
+  <h2>Log Level</h2>
+  <div class="form-row">
+    <div class="form-group">
+      <label for="log-level">Level</label>
+      <select id="log-level" style="width: 120px;">
+        <option value="trace">trace</option>
+        <option value="debug">debug</option>
+        <option value="info">info</option>
+        <option value="warn" selected>warn</option>
+        <option value="error">error</option>
+      </select>
+    </div>
+    <div class="form-group">
+      <span>&nbsp;</span>
+      <button onclick="setLogLevel()">Set Level</button>
+    </div>
+  </div>
+  <div id="log-result" class="result"></div>
+
+  <h2>TCP Ping</h2>
+  <div class="form-row">
+    <div class="form-group">
+      <label for="tcp-host">Host</label>
+      <input type="text" id="tcp-host" placeholder="100.0.0.1 or hostname.netbird.cloud" style="width: 300px;">
+    </div>
+    <div class="form-group">
+      <label for="tcp-port">Port</label>
+      <input type="number" id="tcp-port" placeholder="80" style="width: 80px;">
+    </div>
+    <div class="form-group">
+      <span>&nbsp;</span>
+      <button onclick="doTcpPing()">Connect</button>
+    </div>
+  </div>
+  <div id="tcp-result" class="result"></div>
+
+  <script>
+    const accountID = "{{.AccountID}}";
+
+    async function startClient() {
+        const resultDiv = document.getElementById('client-result');
+        resultDiv.innerHTML = '<span class="info">Starting client...</span>';
+        try {
+            const resp = await fetch('/debug/clients/' + accountID + '/start');
+            const data = await resp.json();
+            if (data.success) {
+                resultDiv.innerHTML = '<span class="success">✓ ' + data.message + '</span>';
+            } else {
+                resultDiv.innerHTML = '<span class="error">✗ ' + data.error + '</span>';
+            }
+        } catch (e) {
+            resultDiv.innerHTML = '<span class="error">Error: ' + e.message + '</span>';
+        }
+    }
+
+    async function stopClient() {
+        const resultDiv = document.getElementById('client-result');
+        resultDiv.innerHTML = '<span class="info">Stopping client...</span>';
+        try {
+            const resp = await fetch('/debug/clients/' + accountID + '/stop');
+            const data = await resp.json();
+            if (data.success) {
+                resultDiv.innerHTML = '<span class="success">✓ ' + data.message + '</span>';
+            } else {
+                resultDiv.innerHTML = '<span class="error">✗ ' + data.error + '</span>';
+            }
+        } catch (e) {
+            resultDiv.innerHTML = '<span class="error">Error: ' + e.message + '</span>';
+        }
+    }
+
+    async function setLogLevel() {
+        const level = document.getElementById('log-level').value;
+        const resultDiv = document.getElementById('log-result');
+        resultDiv.innerHTML = '<span class="info">Setting log level...</span>';
+        try {
+            const resp = await fetch('/debug/clients/' + accountID + '/loglevel?level=' + level);
+            const data = await resp.json();
+            if (data.success) {
+                resultDiv.innerHTML = '<span class="success">✓ Log level set to: ' + data.level + '</span>';
+            } else {
+                resultDiv.innerHTML = '<span class="error">✗ ' + data.error + '</span>';
+            }
+        } catch (e) {
+            resultDiv.innerHTML = '<span class="error">Error: ' + e.message + '</span>';
+        }
+    }
+
+    async function doTcpPing() {
+        const host = document.getElementById('tcp-host').value;
+        const port = document.getElementById('tcp-port').value;
+        if (!host || !port) {
+            alert('Host and port required');
+            return;
+        }
+        const resultDiv = document.getElementById('tcp-result');
+        resultDiv.innerHTML = '<span class="info">Connecting...</span>';
+        try {
+            const resp = await fetch('/debug/clients/' + accountID + '/pingtcp?host=' + encodeURIComponent(host) + '&port=' + port);
+            const data = await resp.json();
+            if (data.success) {
+                resultDiv.innerHTML = '<span class="success">✓ ' + data.host + ':' + data.port + ' connected in ' + data.latency + '</span>';
+            } else {
+                resultDiv.innerHTML = '<span class="error">✗ ' + data.host + ':' + data.port + ': ' + data.error + '</span>';
+            }
+        } catch (e) {
+            resultDiv.innerHTML = '<span class="error">Error: ' + e.message + '</span>';
+        }
+    }
+  </script>
+</body>
+</html>
+{{end}}

proxy/internal/flock/flock_other.go

@@ -0,0 +1,20 @@
+//go:build !unix
+
+package flock
+
+import (
+	"context"
+	"os"
+)
+
+// Lock is a no-op on non-Unix platforms. Returns (nil, nil) to indicate
+// that no lock was acquired; callers must treat a nil file as "proceed
+// without lock" rather than "lock held by someone else."
+func Lock(_ context.Context, _ string) (*os.File, error) {
+	return nil, nil //nolint:nilnil // intentional: nil file signals locking unsupported on this platform
+}
+
+// Unlock is a no-op on non-Unix platforms.
+func Unlock(_ *os.File) error {
+	return nil
+}

proxy/internal/flock/flock_test.go

@@ -0,0 +1,79 @@
+//go:build unix
+
+package flock
+
+import (
+	"context"
+	"os"
+	"path/filepath"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestLockUnlock(t *testing.T) {
+	lockPath := filepath.Join(t.TempDir(), "test.lock")
+
+	f, err := Lock(context.Background(), lockPath)
+	require.NoError(t, err)
+	require.NotNil(t, f)
+
+	_, err = os.Stat(lockPath)
+	assert.NoError(t, err, "lock file should exist")
+
+	err = Unlock(f)
+	assert.NoError(t, err)
+}
+
+func TestUnlockNil(t *testing.T) {
+	err := Unlock(nil)
+	assert.NoError(t, err, "unlocking nil should be a no-op")
+}
+
+func TestLockRespectsContext(t *testing.T) {
+	lockPath := filepath.Join(t.TempDir(), "test.lock")
+
+	f1, err := Lock(context.Background(), lockPath)
+	require.NoError(t, err)
+	defer func() { require.NoError(t, Unlock(f1)) }()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 300*time.Millisecond)
+	defer cancel()
+
+	_, err = Lock(ctx, lockPath)
+	assert.ErrorIs(t, err, context.DeadlineExceeded)
+}
+
+func TestLockBlocks(t *testing.T) {
+	lockPath := filepath.Join(t.TempDir(), "test.lock")
+
+	f1, err := Lock(context.Background(), lockPath)
+	require.NoError(t, err)
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+
+	start := time.Now()
+	var elapsed time.Duration
+
+	go func() {
+		defer wg.Done()
+		f2, err := Lock(context.Background(), lockPath)
+		elapsed = time.Since(start)
+		assert.NoError(t, err)
+		if f2 != nil {
+			assert.NoError(t, Unlock(f2))
+		}
+	}()
+
+	// Hold the lock for 200ms, then release.
+	time.Sleep(200 * time.Millisecond)
+	require.NoError(t, Unlock(f1))
+
+	wg.Wait()
+	assert.GreaterOrEqual(t, elapsed, 150*time.Millisecond,
+		"Lock should have blocked for at least ~200ms")
+}

proxy/internal/flock/flock_unix.go

@@ -0,0 +1,77 @@
+//go:build unix
+
+// Package flock provides best-effort advisory file locking using flock(2).
+//
+// This is used for cross-replica coordination (e.g. preventing duplicate
+// ACME requests). Note that flock(2) does NOT work reliably on NFS volumes:
+// on NFSv3 it depends on the NLM daemon, on NFSv4 Linux emulates it via
+// fcntl locks with different semantics. Callers must treat lock failures
+// as non-fatal and proceed without the lock.
+package flock
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"os"
+	"syscall"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+)
+
+const retryInterval = 100 * time.Millisecond
+
+// Lock acquires an exclusive advisory lock on the given file path.
+// It creates the lock file if it does not exist. The lock attempt
+// respects context cancellation by using non-blocking flock with polling.
+// The caller must call Unlock with the returned *os.File when done.
+func Lock(ctx context.Context, path string) (*os.File, error) {
+	f, err := os.OpenFile(path, os.O_CREATE|os.O_RDWR, 0o600)
+	if err != nil {
+		return nil, fmt.Errorf("open lock file %s: %w", path, err)
+	}
+
+	timer := time.NewTimer(retryInterval)
+	defer timer.Stop()
+
+	for {
+		if err := syscall.Flock(int(f.Fd()), syscall.LOCK_EX|syscall.LOCK_NB); err == nil {
+			return f, nil
+		} else if !errors.Is(err, syscall.EWOULDBLOCK) {
+			if cerr := f.Close(); cerr != nil {
+				log.Debugf("close lock file %s: %v", path, cerr)
+			}
+			return nil, fmt.Errorf("acquire lock on %s: %w", path, err)
+		}
+
+		select {
+		case <-ctx.Done():
+			if cerr := f.Close(); cerr != nil {
+				log.Debugf("close lock file %s: %v", path, cerr)
+			}
+			return nil, ctx.Err()
+		case <-timer.C:
+			timer.Reset(retryInterval)
+		}
+	}
+}
+
+// Unlock releases the lock and closes the file.
+func Unlock(f *os.File) error {
+	if f == nil {
+		return nil
+	}
+
+	defer func() {
+		if cerr := f.Close(); cerr != nil {
+			log.Debugf("close lock file: %v", cerr)
+		}
+	}()
+
+	if err := syscall.Flock(int(f.Fd()), syscall.LOCK_UN); err != nil {
+		return fmt.Errorf("release lock: %w", err)
+	}
+
+	return nil
+}

proxy/internal/geolocation/download.go

@@ -0,0 +1,264 @@
+package geolocation
+
+import (
+	"archive/tar"
+	"bufio"
+	"compress/gzip"
+	"crypto/sha256"
+	"errors"
+	"fmt"
+	"io"
+	"mime"
+	"net/http"
+	"os"
+	"path/filepath"
+	"strings"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+)
+
+const (
+	mmdbTarGZURL  = "https://pkgs.netbird.io/geolocation-dbs/GeoLite2-City/download?suffix=tar.gz"
+	mmdbSha256URL = "https://pkgs.netbird.io/geolocation-dbs/GeoLite2-City/download?suffix=tar.gz.sha256"
+	mmdbInnerName = "GeoLite2-City.mmdb"
+
+	downloadTimeout = 2 * time.Minute
+	maxMMDBSize     = 256 << 20 // 256 MB
+)
+
+// ensureMMDB checks for an existing MMDB file in dataDir. If none is found,
+// it downloads from pkgs.netbird.io with SHA256 verification.
+func ensureMMDB(logger *log.Logger, dataDir string) (string, error) {
+	if err := os.MkdirAll(dataDir, 0o755); err != nil {
+		return "", fmt.Errorf("create geo data directory %s: %w", dataDir, err)
+	}
+
+	pattern := filepath.Join(dataDir, mmdbGlob)
+	if files, _ := filepath.Glob(pattern); len(files) > 0 {
+		mmdbPath := files[len(files)-1]
+		logger.Debugf("using existing geolocation database: %s", mmdbPath)
+		return mmdbPath, nil
+	}
+
+	logger.Info("geolocation database not found, downloading from pkgs.netbird.io")
+	return downloadMMDB(logger, dataDir)
+}
+
+func downloadMMDB(logger *log.Logger, dataDir string) (string, error) {
+	client := &http.Client{Timeout: downloadTimeout}
+
+	datedName, err := fetchRemoteFilename(client, mmdbTarGZURL)
+	if err != nil {
+		return "", fmt.Errorf("get remote filename: %w", err)
+	}
+
+	mmdbFilename := deriveMMDBFilename(datedName)
+	mmdbPath := filepath.Join(dataDir, mmdbFilename)
+
+	tmp, err := os.MkdirTemp("", "geolite-proxy-*")
+	if err != nil {
+		return "", fmt.Errorf("create temp directory: %w", err)
+	}
+	defer os.RemoveAll(tmp)
+
+	checksumFile := filepath.Join(tmp, "checksum.sha256")
+	if err := downloadToFile(client, mmdbSha256URL, checksumFile); err != nil {
+		return "", fmt.Errorf("download checksum: %w", err)
+	}
+
+	expectedHash, err := readChecksumFile(checksumFile)
+	if err != nil {
+		return "", fmt.Errorf("read checksum: %w", err)
+	}
+
+	tarFile := filepath.Join(tmp, datedName)
+	logger.Debugf("downloading geolocation database (%s)", datedName)
+	if err := downloadToFile(client, mmdbTarGZURL, tarFile); err != nil {
+		return "", fmt.Errorf("download database: %w", err)
+	}
+
+	if err := verifySHA256(tarFile, expectedHash); err != nil {
+		return "", fmt.Errorf("verify database checksum: %w", err)
+	}
+
+	if err := extractMMDBFromTarGZ(tarFile, mmdbPath); err != nil {
+		return "", fmt.Errorf("extract database: %w", err)
+	}
+
+	logger.Infof("geolocation database downloaded: %s", mmdbPath)
+	return mmdbPath, nil
+}
+
+// deriveMMDBFilename converts a tar.gz filename to an MMDB filename.
+// Example: GeoLite2-City_20240101.tar.gz -> GeoLite2-City_20240101.mmdb
+func deriveMMDBFilename(tarName string) string {
+	base, _, _ := strings.Cut(tarName, ".")
+	if !strings.Contains(base, "_") {
+		return "GeoLite2-City.mmdb"
+	}
+	return base + ".mmdb"
+}
+
+func fetchRemoteFilename(client *http.Client, url string) (string, error) {
+	resp, err := client.Head(url)
+	if err != nil {
+		return "", err
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		return "", fmt.Errorf("HEAD request: HTTP %d", resp.StatusCode)
+	}
+
+	cd := resp.Header.Get("Content-Disposition")
+	if cd == "" {
+		return "", errors.New("no Content-Disposition header")
+	}
+
+	_, params, err := mime.ParseMediaType(cd)
+	if err != nil {
+		return "", fmt.Errorf("parse Content-Disposition: %w", err)
+	}
+
+	name := filepath.Base(params["filename"])
+	if name == "" || name == "." {
+		return "", errors.New("no filename in Content-Disposition")
+	}
+	return name, nil
+}
+
+func downloadToFile(client *http.Client, url, dest string) error {
+	resp, err := client.Get(url) //nolint:gosec
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		body, _ := io.ReadAll(io.LimitReader(resp.Body, 1024))
+		return fmt.Errorf("HTTP %d: %s", resp.StatusCode, string(body))
+	}
+
+	f, err := os.Create(dest) //nolint:gosec
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+
+	// Cap download at 256 MB to prevent unbounded reads from a compromised server.
+	if _, err := io.Copy(f, io.LimitReader(resp.Body, maxMMDBSize)); err != nil {
+		return err
+	}
+	return nil
+}
+
+func readChecksumFile(path string) (string, error) {
+	f, err := os.Open(path) //nolint:gosec
+	if err != nil {
+		return "", err
+	}
+	defer f.Close()
+
+	scanner := bufio.NewScanner(f)
+	if scanner.Scan() {
+		parts := strings.Fields(scanner.Text())
+		if len(parts) > 0 {
+			return parts[0], nil
+		}
+	}
+	if err := scanner.Err(); err != nil {
+		return "", err
+	}
+	return "", errors.New("empty checksum file")
+}
+
+func verifySHA256(path, expected string) error {
+	f, err := os.Open(path) //nolint:gosec
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+
+	h := sha256.New()
+	if _, err := io.Copy(h, f); err != nil {
+		return err
+	}
+
+	actual := fmt.Sprintf("%x", h.Sum(nil))
+	if actual != expected {
+		return fmt.Errorf("SHA256 mismatch: expected %s, got %s", expected, actual)
+	}
+	return nil
+}
+
+func extractMMDBFromTarGZ(tarGZPath, destPath string) error {
+	f, err := os.Open(tarGZPath) //nolint:gosec
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+
+	gz, err := gzip.NewReader(f)
+	if err != nil {
+		return err
+	}
+	defer gz.Close()
+
+	tr := tar.NewReader(gz)
+	for {
+		hdr, err := tr.Next()
+		if err != nil {
+			if errors.Is(err, io.EOF) {
+				break
+			}
+			return err
+		}
+
+		if hdr.Typeflag == tar.TypeReg && filepath.Base(hdr.Name) == mmdbInnerName {
+			if hdr.Size < 0 || hdr.Size > maxMMDBSize {
+				return fmt.Errorf("mmdb entry size %d exceeds limit %d", hdr.Size, maxMMDBSize)
+			}
+			if err := extractToFileAtomic(io.LimitReader(tr, hdr.Size), destPath); err != nil {
+				return err
+			}
+			return nil
+		}
+	}
+
+	return fmt.Errorf("%s not found in archive", mmdbInnerName)
+}
+
+// extractToFileAtomic writes r to a temporary file in the same directory as
+// destPath, then renames it into place so a crash never leaves a truncated file.
+func extractToFileAtomic(r io.Reader, destPath string) error {
+	dir := filepath.Dir(destPath)
+	tmp, err := os.CreateTemp(dir, ".mmdb-*.tmp")
+	if err != nil {
+		return fmt.Errorf("create temp file: %w", err)
+	}
+	tmpPath := tmp.Name()
+
+	if _, err := io.Copy(tmp, r); err != nil { //nolint:gosec // G110: caller bounds with LimitReader
+		if closeErr := tmp.Close(); closeErr != nil {
+			log.Debugf("failed to close temp file %s: %v", tmpPath, closeErr)
+		}
+		if removeErr := os.Remove(tmpPath); removeErr != nil {
+			log.Debugf("failed to remove temp file %s: %v", tmpPath, removeErr)
+		}
+		return fmt.Errorf("write mmdb: %w", err)
+	}
+	if err := tmp.Close(); err != nil {
+		if removeErr := os.Remove(tmpPath); removeErr != nil {
+			log.Debugf("failed to remove temp file %s: %v", tmpPath, removeErr)
+		}
+		return fmt.Errorf("close temp file: %w", err)
+	}
+	if err := os.Rename(tmpPath, destPath); err != nil {
+		if removeErr := os.Remove(tmpPath); removeErr != nil {
+			log.Debugf("failed to remove temp file %s: %v", tmpPath, removeErr)
+		}
+		return fmt.Errorf("rename to %s: %w", destPath, err)
+	}
+	return nil
+}

proxy/internal/geolocation/geolocation.go

@@ -0,0 +1,152 @@
+// Package geolocation provides IP-to-country lookups using MaxMind GeoLite2 databases.
+package geolocation
+
+import (
+	"fmt"
+	"net/netip"
+	"os"
+	"strconv"
+	"sync"
+
+	"github.com/oschwald/maxminddb-golang"
+	log "github.com/sirupsen/logrus"
+)
+
+const (
+	// EnvDisable disables geolocation lookups entirely when set to a truthy value.
+	EnvDisable = "NB_PROXY_DISABLE_GEOLOCATION"
+
+	mmdbGlob = "GeoLite2-City_*.mmdb"
+)
+
+type record struct {
+	Country struct {
+		ISOCode string `maxminddb:"iso_code"`
+	} `maxminddb:"country"`
+	City struct {
+		Names struct {
+			En string `maxminddb:"en"`
+		} `maxminddb:"names"`
+	} `maxminddb:"city"`
+	Subdivisions []struct {
+		ISOCode string `maxminddb:"iso_code"`
+		Names   struct {
+			En string `maxminddb:"en"`
+		} `maxminddb:"names"`
+	} `maxminddb:"subdivisions"`
+}
+
+// Result holds the outcome of a geo lookup.
+type Result struct {
+	CountryCode     string
+	CityName        string
+	SubdivisionCode string
+	SubdivisionName string
+}
+
+// Lookup provides IP geolocation lookups.
+type Lookup struct {
+	mu     sync.RWMutex
+	db     *maxminddb.Reader
+	logger *log.Logger
+}
+
+// NewLookup opens or downloads the GeoLite2-City MMDB in dataDir.
+// Returns nil without error if geolocation is disabled via environment
+// variable, no data directory is configured, or the download fails
+// (graceful degradation: country restrictions will deny all requests).
+func NewLookup(logger *log.Logger, dataDir string) (*Lookup, error) {
+	if isDisabledByEnv(logger) {
+		logger.Info("geolocation disabled via environment variable")
+		return nil, nil //nolint:nilnil
+	}
+
+	if dataDir == "" {
+		return nil, nil //nolint:nilnil
+	}
+
+	mmdbPath, err := ensureMMDB(logger, dataDir)
+	if err != nil {
+		logger.Warnf("geolocation database unavailable: %v", err)
+		logger.Warn("country-based access restrictions will deny all requests until a database is available")
+		return nil, nil //nolint:nilnil
+	}
+
+	db, err := maxminddb.Open(mmdbPath)
+	if err != nil {
+		return nil, fmt.Errorf("open GeoLite2 database %s: %w", mmdbPath, err)
+	}
+
+	logger.Infof("geolocation database loaded from %s", mmdbPath)
+	return &Lookup{db: db, logger: logger}, nil
+}
+
+// LookupAddr returns the country ISO code and city name for the given IP.
+// Returns an empty Result if the database is nil or the lookup fails.
+func (l *Lookup) LookupAddr(addr netip.Addr) Result {
+	if l == nil {
+		return Result{}
+	}
+
+	l.mu.RLock()
+	defer l.mu.RUnlock()
+
+	if l.db == nil {
+		return Result{}
+	}
+
+	addr = addr.Unmap()
+
+	var rec record
+	if err := l.db.Lookup(addr.AsSlice(), &rec); err != nil {
+		l.logger.Debugf("geolocation lookup %s: %v", addr, err)
+		return Result{}
+	}
+	r := Result{
+		CountryCode: rec.Country.ISOCode,
+		CityName:    rec.City.Names.En,
+	}
+	if len(rec.Subdivisions) > 0 {
+		r.SubdivisionCode = rec.Subdivisions[0].ISOCode
+		r.SubdivisionName = rec.Subdivisions[0].Names.En
+	}
+	return r
+}
+
+// Available reports whether the lookup has a loaded database.
+func (l *Lookup) Available() bool {
+	if l == nil {
+		return false
+	}
+	l.mu.RLock()
+	defer l.mu.RUnlock()
+	return l.db != nil
+}
+
+// Close releases the database resources.
+func (l *Lookup) Close() error {
+	if l == nil {
+		return nil
+	}
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	if l.db != nil {
+		err := l.db.Close()
+		l.db = nil
+		return err
+	}
+	return nil
+}
+
+func isDisabledByEnv(logger *log.Logger) bool {
+	val := os.Getenv(EnvDisable)
+	if val == "" {
+		return false
+	}
+	disabled, err := strconv.ParseBool(val)
+	if err != nil {
+		logger.Warnf("parse %s=%q: %v", EnvDisable, val, err)
+		return false
+	}
+	return disabled
+}

proxy/internal/grpc/auth.go

@@ -0,0 +1,48 @@
+// Package grpc provides gRPC utilities for the proxy client.
+package grpc
+
+import (
+	"context"
+	"os"
+	"strconv"
+
+	log "github.com/sirupsen/logrus"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials"
+)
+
+// EnvProxyAllowInsecure controls whether the proxy token can be sent over non-TLS connections.
+const EnvProxyAllowInsecure = "NB_PROXY_ALLOW_INSECURE"
+
+var _ credentials.PerRPCCredentials = (*proxyAuthToken)(nil)
+
+type proxyAuthToken struct {
+	token         string
+	allowInsecure bool
+}
+
+func (t proxyAuthToken) GetRequestMetadata(context.Context, ...string) (map[string]string, error) {
+	return map[string]string{
+		"authorization": "Bearer " + t.token,
+	}, nil
+}
+
+// RequireTransportSecurity returns true by default to protect the token in transit.
+// Set NB_PROXY_ALLOW_INSECURE=true to allow non-TLS connections (not recommended for production).
+func (t proxyAuthToken) RequireTransportSecurity() bool {
+	return !t.allowInsecure
+}
+
+// WithProxyToken returns a DialOption that sets the proxy access token on each outbound RPC.
+func WithProxyToken(token string) grpc.DialOption {
+	allowInsecure := false
+	if val := os.Getenv(EnvProxyAllowInsecure); val != "" {
+		parsed, err := strconv.ParseBool(val)
+		if err != nil {
+			log.Warnf("invalid value for %s: %v", EnvProxyAllowInsecure, err)
+		} else {
+			allowInsecure = parsed
+		}
+	}
+	return grpc.WithPerRPCCredentials(proxyAuthToken{token: token, allowInsecure: allowInsecure})
+}

proxy/internal/health/health.go

@@ -0,0 +1,405 @@
+// Package health provides health probes for the proxy server.
+package health
+
+import (
+	"context"
+	"encoding/json"
+	"net"
+	"net/http"
+	"strings"
+	"sync"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+
+	"github.com/netbirdio/netbird/client/embed"
+	"github.com/netbirdio/netbird/proxy/internal/types"
+)
+
+const handshakeStaleThreshold = 5 * time.Minute
+
+const (
+	maxConcurrentChecks   = 3
+	maxClientCheckTimeout = 5 * time.Minute
+)
+
+// clientProvider provides access to NetBird clients for health checks.
+type clientProvider interface {
+	ListClientsForStartup() map[types.AccountID]*embed.Client
+}
+
+// Checker tracks health state and provides probe endpoints.
+type Checker struct {
+	logger   *log.Logger
+	provider clientProvider
+
+	mu                  sync.RWMutex
+	managementConnected bool
+	initialSyncComplete bool
+	shuttingDown        bool
+
+	// checkSem limits concurrent client health checks.
+	checkSem chan struct{}
+
+	// checkHealth checks the health of a single client.
+	// Defaults to checkClientHealth; overridable in tests.
+	checkHealth func(*embed.Client) ClientHealth
+}
+
+// ClientHealth represents the health status of a single NetBird client.
+type ClientHealth struct {
+	Healthy             bool   `json:"healthy"`
+	ManagementConnected bool   `json:"management_connected"`
+	SignalConnected     bool   `json:"signal_connected"`
+	RelaysConnected     int    `json:"relays_connected"`
+	RelaysTotal         int    `json:"relays_total"`
+	PeersTotal          int    `json:"peers_total"`
+	PeersConnected      int    `json:"peers_connected"`
+	PeersP2P            int    `json:"peers_p2p"`
+	PeersRelayed        int    `json:"peers_relayed"`
+	PeersDegraded       int    `json:"peers_degraded"`
+	Error               string `json:"error,omitempty"`
+}
+
+// ProbeResponse represents the JSON response for health probes.
+type ProbeResponse struct {
+	Status  string                           `json:"status"`
+	Checks  map[string]bool                  `json:"checks,omitempty"`
+	Clients map[types.AccountID]ClientHealth `json:"clients,omitempty"`
+}
+
+// Server runs the health probe HTTP server on a dedicated port.
+type Server struct {
+	server  *http.Server
+	logger  *log.Logger
+	checker *Checker
+}
+
+// SetManagementConnected updates the management connection state.
+func (c *Checker) SetManagementConnected(connected bool) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.managementConnected = connected
+}
+
+// SetInitialSyncComplete marks that the initial mapping sync has completed.
+func (c *Checker) SetInitialSyncComplete() {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.initialSyncComplete = true
+}
+
+// SetShuttingDown marks the server as shutting down.
+// This causes ReadinessProbe to return false so load balancers stop routing traffic.
+func (c *Checker) SetShuttingDown() {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.shuttingDown = true
+}
+
+// CheckClientsConnected verifies all clients are connected to management/signal/relay.
+// Uses the provided context for timeout/cancellation, with a maximum bound of maxClientCheckTimeout.
+// Limits concurrent checks via semaphore.
+func (c *Checker) CheckClientsConnected(ctx context.Context) (bool, map[types.AccountID]ClientHealth) {
+	// Apply upper bound timeout in case parent context has no deadline
+	ctx, cancel := context.WithTimeout(ctx, maxClientCheckTimeout)
+	defer cancel()
+
+	clients := c.provider.ListClientsForStartup()
+
+	// No clients is not a health issue
+	if len(clients) == 0 {
+		return true, make(map[types.AccountID]ClientHealth)
+	}
+
+	type result struct {
+		accountID types.AccountID
+		health    ClientHealth
+	}
+
+	resultsCh := make(chan result, len(clients))
+	var wg sync.WaitGroup
+
+	for accountID, client := range clients {
+		wg.Add(1)
+		go func(id types.AccountID, cl *embed.Client) {
+			defer wg.Done()
+
+			// Acquire semaphore
+			select {
+			case c.checkSem <- struct{}{}:
+				defer func() { <-c.checkSem }()
+			case <-ctx.Done():
+				resultsCh <- result{id, ClientHealth{Healthy: false, Error: ctx.Err().Error()}}
+				return
+			}
+
+			resultsCh <- result{id, c.checkHealth(cl)}
+		}(accountID, client)
+	}
+
+	go func() {
+		wg.Wait()
+		close(resultsCh)
+	}()
+
+	results := make(map[types.AccountID]ClientHealth)
+	allHealthy := true
+	for r := range resultsCh {
+		results[r.accountID] = r.health
+		if !r.health.Healthy {
+			allHealthy = false
+		}
+	}
+
+	return allHealthy, results
+}
+
+// LivenessProbe returns true if the process is alive.
+// This should always return true if we can respond.
+func (c *Checker) LivenessProbe() bool {
+	return true
+}
+
+// ReadinessProbe returns true if the server can accept traffic.
+func (c *Checker) ReadinessProbe() bool {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+	if c.shuttingDown {
+		return false
+	}
+	return c.managementConnected
+}
+
+// StartupProbe checks if initial startup is complete.
+// Checks management connection, initial sync, and all client health directly.
+// Uses the provided context for timeout/cancellation.
+func (c *Checker) StartupProbe(ctx context.Context) bool {
+	c.mu.RLock()
+	mgmt := c.managementConnected
+	sync := c.initialSyncComplete
+	c.mu.RUnlock()
+
+	if !mgmt || !sync {
+		return false
+	}
+
+	// Check all clients are connected to management/signal/relay.
+	// Returns true when no clients exist (nothing to check).
+	allHealthy, _ := c.CheckClientsConnected(ctx)
+	return allHealthy
+}
+
+// Handler returns an http.Handler for health probe endpoints.
+func (c *Checker) Handler() http.Handler {
+	mux := http.NewServeMux()
+	mux.HandleFunc("/healthz/live", c.handleLiveness)
+	mux.HandleFunc("/healthz/ready", c.handleReadiness)
+	mux.HandleFunc("/healthz/startup", c.handleStartup)
+	mux.HandleFunc("/healthz", c.handleFull)
+	return mux
+}
+
+func (c *Checker) handleLiveness(w http.ResponseWriter, r *http.Request) {
+	if c.LivenessProbe() {
+		c.writeProbeResponse(w, http.StatusOK, "ok", nil, nil)
+		return
+	}
+	c.writeProbeResponse(w, http.StatusServiceUnavailable, "fail", nil, nil)
+}
+
+func (c *Checker) handleReadiness(w http.ResponseWriter, r *http.Request) {
+	c.mu.RLock()
+	checks := map[string]bool{
+		"management_connected": c.managementConnected,
+	}
+	c.mu.RUnlock()
+
+	if c.ReadinessProbe() {
+		c.writeProbeResponse(w, http.StatusOK, "ok", checks, nil)
+		return
+	}
+	c.writeProbeResponse(w, http.StatusServiceUnavailable, "fail", checks, nil)
+}
+
+func (c *Checker) handleStartup(w http.ResponseWriter, r *http.Request) {
+	c.mu.RLock()
+	mgmt := c.managementConnected
+	syncComplete := c.initialSyncComplete
+	c.mu.RUnlock()
+
+	allClientsHealthy, clientHealth := c.CheckClientsConnected(r.Context())
+
+	checks := map[string]bool{
+		"management_connected":  mgmt,
+		"initial_sync_complete": syncComplete,
+		"all_clients_healthy":   allClientsHealthy,
+	}
+
+	ready := mgmt && syncComplete && allClientsHealthy
+	if ready {
+		c.writeProbeResponse(w, http.StatusOK, "ok", checks, clientHealth)
+		return
+	}
+	c.writeProbeResponse(w, http.StatusServiceUnavailable, "fail", checks, clientHealth)
+}
+
+func (c *Checker) handleFull(w http.ResponseWriter, r *http.Request) {
+	c.mu.RLock()
+	mgmt := c.managementConnected
+	sync := c.initialSyncComplete
+	c.mu.RUnlock()
+
+	allClientsHealthy, clientHealth := c.CheckClientsConnected(r.Context())
+
+	checks := map[string]bool{
+		"management_connected":  mgmt,
+		"initial_sync_complete": sync,
+		"all_clients_healthy":   allClientsHealthy,
+	}
+
+	status := "ok"
+	statusCode := http.StatusOK
+	if !c.ReadinessProbe() {
+		status = "fail"
+		statusCode = http.StatusServiceUnavailable
+	}
+
+	c.writeProbeResponse(w, statusCode, status, checks, clientHealth)
+}
+
+func (c *Checker) writeProbeResponse(w http.ResponseWriter, statusCode int, status string, checks map[string]bool, clients map[types.AccountID]ClientHealth) {
+	w.Header().Set("Content-Type", "application/json")
+	w.WriteHeader(statusCode)
+
+	resp := ProbeResponse{
+		Status:  status,
+		Checks:  checks,
+		Clients: clients,
+	}
+	if err := json.NewEncoder(w).Encode(resp); err != nil {
+		c.logger.Debugf("write health response: %v", err)
+	}
+}
+
+// ListenAndServe starts the health probe server.
+func (s *Server) ListenAndServe() error {
+	s.logger.Infof("starting health probe server on %s", s.server.Addr)
+	return s.server.ListenAndServe()
+}
+
+// Serve starts the health probe server on the given listener.
+func (s *Server) Serve(l net.Listener) error {
+	s.logger.Infof("starting health probe server on %s", l.Addr())
+	return s.server.Serve(l)
+}
+
+// Shutdown gracefully shuts down the health probe server.
+func (s *Server) Shutdown(ctx context.Context) error {
+	return s.server.Shutdown(ctx)
+}
+
+// NewChecker creates a new health checker.
+func NewChecker(logger *log.Logger, provider clientProvider) *Checker {
+	if logger == nil {
+		logger = log.StandardLogger()
+	}
+	return &Checker{
+		logger:      logger,
+		provider:    provider,
+		checkSem:    make(chan struct{}, maxConcurrentChecks),
+		checkHealth: checkClientHealth,
+	}
+}
+
+// NewServer creates a new health probe server.
+// If metricsHandler is non-nil, it is mounted at /metrics on the same port.
+func NewServer(addr string, checker *Checker, logger *log.Logger, metricsHandler http.Handler) *Server {
+	if logger == nil {
+		logger = log.StandardLogger()
+	}
+
+	handler := checker.Handler()
+	if metricsHandler != nil {
+		mux := http.NewServeMux()
+		mux.Handle("/metrics", metricsHandler)
+		mux.Handle("/", handler)
+		handler = mux
+	}
+
+	return &Server{
+		server: &http.Server{
+			Addr:         addr,
+			Handler:      handler,
+			ReadTimeout:  5 * time.Second,
+			WriteTimeout: 5 * time.Second,
+		},
+		logger:  logger,
+		checker: checker,
+	}
+}
+
+func checkClientHealth(client *embed.Client) ClientHealth {
+	if client == nil {
+		return ClientHealth{
+			Healthy: false,
+			Error:   "client not initialized",
+		}
+	}
+
+	status, err := client.Status()
+	if err != nil {
+		return ClientHealth{
+			Healthy: false,
+			Error:   err.Error(),
+		}
+	}
+
+	// Count only rel:// and rels:// relays (not stun/turn)
+	var relayCount, relaysConnected int
+	for _, relay := range status.Relays {
+		if !strings.HasPrefix(relay.URI, "rel://") && !strings.HasPrefix(relay.URI, "rels://") {
+			continue
+		}
+		relayCount++
+		if relay.Err == nil {
+			relaysConnected++
+		}
+	}
+
+	// Count peer connection stats
+	now := time.Now()
+	var peersConnected, peersP2P, peersRelayed, peersDegraded int
+	for _, p := range status.Peers {
+		if p.ConnStatus != embed.PeerStatusConnected {
+			continue
+		}
+		peersConnected++
+		if p.Relayed {
+			peersRelayed++
+		} else {
+			peersP2P++
+		}
+		if p.LastWireguardHandshake.IsZero() || now.Sub(p.LastWireguardHandshake) > handshakeStaleThreshold {
+			peersDegraded++
+		}
+	}
+
+	// Client is healthy if connected to management, signal, and at least one relay (if any are defined)
+	healthy := status.ManagementState.Connected &&
+		status.SignalState.Connected &&
+		(relayCount == 0 || relaysConnected > 0)
+
+	return ClientHealth{
+		Healthy:             healthy,
+		ManagementConnected: status.ManagementState.Connected,
+		SignalConnected:     status.SignalState.Connected,
+		RelaysConnected:     relaysConnected,
+		RelaysTotal:         relayCount,
+		PeersTotal:          len(status.Peers),
+		PeersConnected:      peersConnected,
+		PeersP2P:            peersP2P,
+		PeersRelayed:        peersRelayed,
+		PeersDegraded:       peersDegraded,
+	}
+}

proxy/internal/health/health_test.go

@@ -0,0 +1,473 @@
+package health
+
+import (
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/client/embed"
+	"github.com/netbirdio/netbird/proxy/internal/types"
+)
+
+type mockClientProvider struct {
+	clients map[types.AccountID]*embed.Client
+}
+
+func (m *mockClientProvider) ListClientsForStartup() map[types.AccountID]*embed.Client {
+	return m.clients
+}
+
+// newTestChecker creates a checker with a mock health function for testing.
+// The health function returns the provided ClientHealth for every client.
+func newTestChecker(provider clientProvider, healthResult ClientHealth) *Checker {
+	c := NewChecker(nil, provider)
+	c.checkHealth = func(_ *embed.Client) ClientHealth {
+		return healthResult
+	}
+	return c
+}
+
+func TestChecker_LivenessProbe(t *testing.T) {
+	checker := NewChecker(nil, &mockClientProvider{})
+
+	// Liveness should always return true if we can respond.
+	assert.True(t, checker.LivenessProbe())
+}
+
+func TestChecker_ReadinessProbe(t *testing.T) {
+	checker := NewChecker(nil, &mockClientProvider{})
+
+	// Initially not ready (management not connected).
+	assert.False(t, checker.ReadinessProbe())
+
+	// After management connects, should be ready.
+	checker.SetManagementConnected(true)
+	assert.True(t, checker.ReadinessProbe())
+
+	// If management disconnects, should not be ready.
+	checker.SetManagementConnected(false)
+	assert.False(t, checker.ReadinessProbe())
+}
+
+// TestStartupProbe_EmptyServiceList covers the scenario where management has
+// no services configured for this proxy. The proxy should become ready once
+// management is connected and the initial sync completes, even with zero clients.
+func TestStartupProbe_EmptyServiceList(t *testing.T) {
+	checker := NewChecker(nil, &mockClientProvider{})
+
+	// No management connection = not ready.
+	assert.False(t, checker.StartupProbe(context.Background()))
+
+	// Management connected but no sync = not ready.
+	checker.SetManagementConnected(true)
+	assert.False(t, checker.StartupProbe(context.Background()))
+
+	// Management + sync complete + no clients = ready.
+	checker.SetInitialSyncComplete()
+	assert.True(t, checker.StartupProbe(context.Background()))
+}
+
+// TestStartupProbe_WithUnhealthyClients verifies that when services exist
+// and clients have been created but are not yet fully connected (to mgmt,
+// signal, relays), the startup probe does NOT pass.
+func TestStartupProbe_WithUnhealthyClients(t *testing.T) {
+	provider := &mockClientProvider{
+		clients: map[types.AccountID]*embed.Client{
+			"account-1": nil, // concrete client not needed; checkHealth is mocked
+			"account-2": nil,
+		},
+	}
+	checker := newTestChecker(provider, ClientHealth{Healthy: false, Error: "not connected yet"})
+	checker.SetManagementConnected(true)
+	checker.SetInitialSyncComplete()
+
+	assert.False(t, checker.StartupProbe(context.Background()),
+		"startup probe must not pass when clients are unhealthy")
+}
+
+// TestStartupProbe_WithHealthyClients verifies that once all clients are
+// connected and healthy, the startup probe passes.
+func TestStartupProbe_WithHealthyClients(t *testing.T) {
+	provider := &mockClientProvider{
+		clients: map[types.AccountID]*embed.Client{
+			"account-1": nil,
+			"account-2": nil,
+		},
+	}
+	checker := newTestChecker(provider, ClientHealth{
+		Healthy:             true,
+		ManagementConnected: true,
+		SignalConnected:     true,
+		RelaysConnected:     1,
+		RelaysTotal:         1,
+	})
+	checker.SetManagementConnected(true)
+	checker.SetInitialSyncComplete()
+
+	assert.True(t, checker.StartupProbe(context.Background()),
+		"startup probe must pass when all clients are healthy")
+}
+
+// TestStartupProbe_MixedHealthClients verifies that if any single client is
+// unhealthy, the startup probe fails (all-or-nothing).
+func TestStartupProbe_MixedHealthClients(t *testing.T) {
+	provider := &mockClientProvider{
+		clients: map[types.AccountID]*embed.Client{
+			"healthy-account":   nil,
+			"unhealthy-account": nil,
+		},
+	}
+
+	checker := NewChecker(nil, provider)
+	checker.checkHealth = func(cl *embed.Client) ClientHealth {
+		// We identify accounts by their position in the map iteration; since we
+		// can't control map order, make exactly one unhealthy via counter.
+		return ClientHealth{Healthy: false}
+	}
+	checker.SetManagementConnected(true)
+	checker.SetInitialSyncComplete()
+
+	assert.False(t, checker.StartupProbe(context.Background()),
+		"startup probe must fail if any client is unhealthy")
+}
+
+// TestStartupProbe_RequiresAllConditions ensures that each individual
+// prerequisite (management, sync, clients) is necessary. The probe must not
+// pass if any one is missing.
+func TestStartupProbe_RequiresAllConditions(t *testing.T) {
+	provider := &mockClientProvider{
+		clients: map[types.AccountID]*embed.Client{
+			"account-1": nil,
+		},
+	}
+
+	t.Run("no management", func(t *testing.T) {
+		checker := newTestChecker(provider, ClientHealth{Healthy: true})
+		checker.SetInitialSyncComplete()
+		// management NOT connected
+		assert.False(t, checker.StartupProbe(context.Background()))
+	})
+
+	t.Run("no sync", func(t *testing.T) {
+		checker := newTestChecker(provider, ClientHealth{Healthy: true})
+		checker.SetManagementConnected(true)
+		// sync NOT complete
+		assert.False(t, checker.StartupProbe(context.Background()))
+	})
+
+	t.Run("unhealthy client", func(t *testing.T) {
+		checker := newTestChecker(provider, ClientHealth{Healthy: false})
+		checker.SetManagementConnected(true)
+		checker.SetInitialSyncComplete()
+		assert.False(t, checker.StartupProbe(context.Background()))
+	})
+
+	t.Run("all conditions met", func(t *testing.T) {
+		checker := newTestChecker(provider, ClientHealth{Healthy: true})
+		checker.SetManagementConnected(true)
+		checker.SetInitialSyncComplete()
+		assert.True(t, checker.StartupProbe(context.Background()))
+	})
+}
+
+// TestStartupProbe_ConcurrentAccess runs the startup probe from many
+// goroutines simultaneously to check for races.
+func TestStartupProbe_ConcurrentAccess(t *testing.T) {
+	provider := &mockClientProvider{
+		clients: map[types.AccountID]*embed.Client{
+			"account-1": nil,
+			"account-2": nil,
+		},
+	}
+	checker := newTestChecker(provider, ClientHealth{Healthy: true})
+	checker.SetManagementConnected(true)
+	checker.SetInitialSyncComplete()
+
+	var wg sync.WaitGroup
+	const goroutines = 50
+	results := make([]bool, goroutines)
+
+	for i := range goroutines {
+		wg.Add(1)
+		go func(idx int) {
+			defer wg.Done()
+			results[idx] = checker.StartupProbe(context.Background())
+		}(i)
+	}
+	wg.Wait()
+
+	for i, r := range results {
+		assert.True(t, r, "goroutine %d got unexpected result", i)
+	}
+}
+
+// TestStartupProbe_CancelledContext verifies that a cancelled context causes
+// the probe to report unhealthy when client checks are needed.
+func TestStartupProbe_CancelledContext(t *testing.T) {
+	t.Run("no management bypasses context", func(t *testing.T) {
+		checker := NewChecker(nil, &mockClientProvider{})
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+		// Should be false because management isn't connected, context is irrelevant.
+		assert.False(t, checker.StartupProbe(ctx))
+	})
+
+	t.Run("with clients and cancelled context", func(t *testing.T) {
+		provider := &mockClientProvider{
+			clients: map[types.AccountID]*embed.Client{
+				"account-1": nil,
+			},
+		}
+		checker := NewChecker(nil, provider)
+		// Use the real checkHealth path — a cancelled context should cause
+		// the semaphore acquisition to fail, reporting unhealthy.
+		checker.SetManagementConnected(true)
+		checker.SetInitialSyncComplete()
+
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+		assert.False(t, checker.StartupProbe(ctx),
+			"cancelled context must result in unhealthy when clients exist")
+	})
+}
+
+// TestHandler_Startup_EmptyServiceList verifies the HTTP startup endpoint
+// returns 200 when management is connected, sync is complete, and there are
+// no services/clients.
+func TestHandler_Startup_EmptyServiceList(t *testing.T) {
+	checker := NewChecker(nil, &mockClientProvider{})
+	checker.SetManagementConnected(true)
+	checker.SetInitialSyncComplete()
+	handler := checker.Handler()
+
+	req := httptest.NewRequest(http.MethodGet, "/healthz/startup", nil)
+	rec := httptest.NewRecorder()
+	handler.ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusOK, rec.Code)
+
+	var resp ProbeResponse
+	require.NoError(t, json.NewDecoder(rec.Body).Decode(&resp))
+	assert.Equal(t, "ok", resp.Status)
+	assert.True(t, resp.Checks["management_connected"])
+	assert.True(t, resp.Checks["initial_sync_complete"])
+	assert.True(t, resp.Checks["all_clients_healthy"])
+	assert.Empty(t, resp.Clients)
+}
+
+// TestHandler_Startup_WithUnhealthyClients verifies that the HTTP startup
+// endpoint returns 503 when clients exist but are not yet healthy.
+func TestHandler_Startup_WithUnhealthyClients(t *testing.T) {
+	provider := &mockClientProvider{
+		clients: map[types.AccountID]*embed.Client{
+			"account-1": nil,
+		},
+	}
+	checker := newTestChecker(provider, ClientHealth{Healthy: false, Error: "starting"})
+	checker.SetManagementConnected(true)
+	checker.SetInitialSyncComplete()
+	handler := checker.Handler()
+
+	req := httptest.NewRequest(http.MethodGet, "/healthz/startup", nil)
+	rec := httptest.NewRecorder()
+	handler.ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusServiceUnavailable, rec.Code)
+
+	var resp ProbeResponse
+	require.NoError(t, json.NewDecoder(rec.Body).Decode(&resp))
+	assert.Equal(t, "fail", resp.Status)
+	assert.True(t, resp.Checks["management_connected"])
+	assert.True(t, resp.Checks["initial_sync_complete"])
+	assert.False(t, resp.Checks["all_clients_healthy"])
+	require.Contains(t, resp.Clients, types.AccountID("account-1"))
+	assert.Equal(t, "starting", resp.Clients["account-1"].Error)
+}
+
+// TestHandler_Startup_WithHealthyClients verifies the HTTP startup endpoint
+// returns 200 once clients are healthy.
+func TestHandler_Startup_WithHealthyClients(t *testing.T) {
+	provider := &mockClientProvider{
+		clients: map[types.AccountID]*embed.Client{
+			"account-1": nil,
+		},
+	}
+	checker := newTestChecker(provider, ClientHealth{
+		Healthy:             true,
+		ManagementConnected: true,
+		SignalConnected:     true,
+		RelaysConnected:     1,
+		RelaysTotal:         1,
+	})
+	checker.SetManagementConnected(true)
+	checker.SetInitialSyncComplete()
+	handler := checker.Handler()
+
+	req := httptest.NewRequest(http.MethodGet, "/healthz/startup", nil)
+	rec := httptest.NewRecorder()
+	handler.ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusOK, rec.Code)
+
+	var resp ProbeResponse
+	require.NoError(t, json.NewDecoder(rec.Body).Decode(&resp))
+	assert.Equal(t, "ok", resp.Status)
+	assert.True(t, resp.Checks["all_clients_healthy"])
+}
+
+// TestHandler_Startup_NotComplete verifies the startup handler returns 503
+// when prerequisites aren't met.
+func TestHandler_Startup_NotComplete(t *testing.T) {
+	checker := NewChecker(nil, &mockClientProvider{})
+	handler := checker.Handler()
+
+	req := httptest.NewRequest(http.MethodGet, "/healthz/startup", nil)
+	rec := httptest.NewRecorder()
+	handler.ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusServiceUnavailable, rec.Code)
+
+	var resp ProbeResponse
+	require.NoError(t, json.NewDecoder(rec.Body).Decode(&resp))
+	assert.Equal(t, "fail", resp.Status)
+}
+
+func TestChecker_Handler_Liveness(t *testing.T) {
+	checker := NewChecker(nil, &mockClientProvider{})
+	handler := checker.Handler()
+
+	req := httptest.NewRequest(http.MethodGet, "/healthz/live", nil)
+	rec := httptest.NewRecorder()
+	handler.ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusOK, rec.Code)
+
+	var resp ProbeResponse
+	require.NoError(t, json.NewDecoder(rec.Body).Decode(&resp))
+	assert.Equal(t, "ok", resp.Status)
+}
+
+func TestChecker_Handler_Readiness_NotReady(t *testing.T) {
+	checker := NewChecker(nil, &mockClientProvider{})
+	handler := checker.Handler()
+
+	req := httptest.NewRequest(http.MethodGet, "/healthz/ready", nil)
+	rec := httptest.NewRecorder()
+	handler.ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusServiceUnavailable, rec.Code)
+
+	var resp ProbeResponse
+	require.NoError(t, json.NewDecoder(rec.Body).Decode(&resp))
+	assert.Equal(t, "fail", resp.Status)
+	assert.False(t, resp.Checks["management_connected"])
+}
+
+func TestChecker_Handler_Readiness_Ready(t *testing.T) {
+	checker := NewChecker(nil, &mockClientProvider{})
+	checker.SetManagementConnected(true)
+	handler := checker.Handler()
+
+	req := httptest.NewRequest(http.MethodGet, "/healthz/ready", nil)
+	rec := httptest.NewRecorder()
+	handler.ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusOK, rec.Code)
+
+	var resp ProbeResponse
+	require.NoError(t, json.NewDecoder(rec.Body).Decode(&resp))
+	assert.Equal(t, "ok", resp.Status)
+	assert.True(t, resp.Checks["management_connected"])
+}
+
+func TestChecker_Handler_Full(t *testing.T) {
+	checker := NewChecker(nil, &mockClientProvider{})
+	checker.SetManagementConnected(true)
+	handler := checker.Handler()
+
+	req := httptest.NewRequest(http.MethodGet, "/healthz", nil)
+	rec := httptest.NewRecorder()
+	handler.ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusOK, rec.Code)
+
+	var resp ProbeResponse
+	require.NoError(t, json.NewDecoder(rec.Body).Decode(&resp))
+	assert.Equal(t, "ok", resp.Status)
+	assert.NotNil(t, resp.Checks)
+	// Clients may be empty map when no clients exist.
+	assert.Empty(t, resp.Clients)
+}
+
+func TestChecker_SetShuttingDown(t *testing.T) {
+	checker := NewChecker(nil, &mockClientProvider{})
+	checker.SetManagementConnected(true)
+
+	assert.True(t, checker.ReadinessProbe(), "should be ready before shutdown")
+
+	checker.SetShuttingDown()
+
+	assert.False(t, checker.ReadinessProbe(), "should not be ready after shutdown")
+}
+
+func TestChecker_Handler_Readiness_ShuttingDown(t *testing.T) {
+	checker := NewChecker(nil, &mockClientProvider{})
+	checker.SetManagementConnected(true)
+	checker.SetShuttingDown()
+	handler := checker.Handler()
+
+	req := httptest.NewRequest(http.MethodGet, "/healthz/ready", nil)
+	rec := httptest.NewRecorder()
+	handler.ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusServiceUnavailable, rec.Code)
+
+	var resp ProbeResponse
+	require.NoError(t, json.NewDecoder(rec.Body).Decode(&resp))
+	assert.Equal(t, "fail", resp.Status)
+}
+
+func TestNewServer_WithMetricsHandler(t *testing.T) {
+	checker := NewChecker(nil, &mockClientProvider{})
+	checker.SetManagementConnected(true)
+
+	metricsHandler := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.WriteHeader(http.StatusOK)
+		_, _ = w.Write([]byte("metrics"))
+	})
+
+	srv := NewServer(":0", checker, nil, metricsHandler)
+	require.NotNil(t, srv)
+
+	// Verify health endpoint still works through the mux.
+	req := httptest.NewRequest(http.MethodGet, "/healthz/live", nil)
+	rec := httptest.NewRecorder()
+	srv.server.Handler.ServeHTTP(rec, req)
+	assert.Equal(t, http.StatusOK, rec.Code)
+
+	// Verify metrics endpoint is mounted.
+	req = httptest.NewRequest(http.MethodGet, "/metrics", nil)
+	rec = httptest.NewRecorder()
+	srv.server.Handler.ServeHTTP(rec, req)
+	assert.Equal(t, http.StatusOK, rec.Code)
+	assert.Equal(t, "metrics", rec.Body.String())
+}
+
+func TestNewServer_WithoutMetricsHandler(t *testing.T) {
+	checker := NewChecker(nil, &mockClientProvider{})
+	checker.SetManagementConnected(true)
+
+	srv := NewServer(":0", checker, nil, nil)
+	require.NotNil(t, srv)
+
+	req := httptest.NewRequest(http.MethodGet, "/healthz/live", nil)
+	rec := httptest.NewRecorder()
+	srv.server.Handler.ServeHTTP(rec, req)
+	assert.Equal(t, http.StatusOK, rec.Code)
+}

proxy/internal/k8s/lease.go

@@ -0,0 +1,281 @@
+// Package k8s provides a lightweight Kubernetes API client for coordination
+// Leases. It uses raw HTTP calls against the mounted service account
+// credentials, avoiding a dependency on client-go.
+package k8s
+
+import (
+	"bytes"
+	"context"
+	"crypto/sha256"
+	"crypto/tls"
+	"crypto/x509"
+	"encoding/hex"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"os"
+	"strings"
+	"time"
+)
+
+const (
+	saTokenPath     = "/var/run/secrets/kubernetes.io/serviceaccount/token" //nolint:gosec
+	saNamespacePath = "/var/run/secrets/kubernetes.io/serviceaccount/namespace"
+	saCACertPath    = "/var/run/secrets/kubernetes.io/serviceaccount/ca.crt"
+
+	leaseAPIPath = "/apis/coordination.k8s.io/v1"
+)
+
+// ErrConflict is returned when a Lease update fails due to a
+// resourceVersion mismatch (another writer updated the object first).
+var ErrConflict = errors.New("conflict: resource version mismatch")
+
+// Lease represents a coordination.k8s.io/v1 Lease object with only the
+// fields needed for distributed locking.
+type Lease struct {
+	APIVersion string        `json:"apiVersion"`
+	Kind       string        `json:"kind"`
+	Metadata   LeaseMetadata `json:"metadata"`
+	Spec       LeaseSpec     `json:"spec"`
+}
+
+// LeaseMetadata holds the standard k8s object metadata fields used by Leases.
+type LeaseMetadata struct {
+	Name            string            `json:"name"`
+	Namespace       string            `json:"namespace,omitempty"`
+	ResourceVersion string            `json:"resourceVersion,omitempty"`
+	Annotations     map[string]string `json:"annotations,omitempty"`
+}
+
+// LeaseSpec holds the Lease specification fields.
+type LeaseSpec struct {
+	HolderIdentity       *string    `json:"holderIdentity"`
+	LeaseDurationSeconds *int32     `json:"leaseDurationSeconds,omitempty"`
+	AcquireTime          *MicroTime `json:"acquireTime"`
+	RenewTime            *MicroTime `json:"renewTime"`
+}
+
+// MicroTime wraps time.Time with Kubernetes MicroTime JSON formatting.
+type MicroTime struct {
+	time.Time
+}
+
+const microTimeFormat = "2006-01-02T15:04:05.000000Z"
+
+// MarshalJSON implements json.Marshaler with k8s MicroTime format.
+func (t *MicroTime) MarshalJSON() ([]byte, error) {
+	return json.Marshal(t.UTC().Format(microTimeFormat))
+}
+
+// UnmarshalJSON implements json.Unmarshaler with k8s MicroTime format.
+func (t *MicroTime) UnmarshalJSON(data []byte) error {
+	var s string
+	if err := json.Unmarshal(data, &s); err != nil {
+		return err
+	}
+	if s == "" {
+		t.Time = time.Time{}
+		return nil
+	}
+
+	parsed, err := time.Parse(microTimeFormat, s)
+	if err != nil {
+		return fmt.Errorf("parse MicroTime %q: %w", s, err)
+	}
+	t.Time = parsed
+	return nil
+}
+
+// LeaseClient talks to the Kubernetes coordination API using raw HTTP.
+type LeaseClient struct {
+	baseURL    string
+	namespace  string
+	httpClient *http.Client
+}
+
+// NewLeaseClient creates a client that authenticates via the pod's
+// mounted service account. It reads the namespace and CA certificate
+// at construction time (they don't rotate) but reads the bearer token
+// fresh on each request (tokens rotate).
+func NewLeaseClient() (*LeaseClient, error) {
+	host := os.Getenv("KUBERNETES_SERVICE_HOST")
+	port := os.Getenv("KUBERNETES_SERVICE_PORT")
+	if host == "" || port == "" {
+		return nil, fmt.Errorf("KUBERNETES_SERVICE_HOST/PORT not set")
+	}
+
+	ns, err := os.ReadFile(saNamespacePath)
+	if err != nil {
+		return nil, fmt.Errorf("read namespace from %s: %w", saNamespacePath, err)
+	}
+
+	caCert, err := os.ReadFile(saCACertPath)
+	if err != nil {
+		return nil, fmt.Errorf("read CA cert from %s: %w", saCACertPath, err)
+	}
+
+	pool := x509.NewCertPool()
+	if !pool.AppendCertsFromPEM(caCert) {
+		return nil, fmt.Errorf("parse CA certificate from %s", saCACertPath)
+	}
+
+	return &LeaseClient{
+		baseURL:   fmt.Sprintf("https://%s:%s", host, port),
+		namespace: strings.TrimSpace(string(ns)),
+		httpClient: &http.Client{
+			Timeout: 10 * time.Second,
+			Transport: &http.Transport{
+				TLSClientConfig: &tls.Config{
+					RootCAs: pool,
+				},
+			},
+		},
+	}, nil
+}
+
+// Namespace returns the namespace this client operates in.
+func (c *LeaseClient) Namespace() string {
+	return c.namespace
+}
+
+// Get retrieves a Lease by name. Returns (nil, nil) if the Lease does not exist.
+func (c *LeaseClient) Get(ctx context.Context, name string) (*Lease, error) {
+	url := fmt.Sprintf("%s%s/namespaces/%s/leases/%s", c.baseURL, leaseAPIPath, c.namespace, name)
+
+	resp, err := c.doRequest(ctx, http.MethodGet, url, nil)
+	if err != nil {
+		return nil, err
+	}
+	defer func() { _ = resp.Body.Close() }()
+
+	if resp.StatusCode == http.StatusNotFound {
+		return nil, nil //nolint:nilnil
+	}
+	if resp.StatusCode != http.StatusOK {
+		return nil, c.readError(resp)
+	}
+
+	var lease Lease
+	if err := json.NewDecoder(resp.Body).Decode(&lease); err != nil {
+		return nil, fmt.Errorf("decode lease response: %w", err)
+	}
+	return &lease, nil
+}
+
+// Create creates a new Lease. Returns the created Lease with server-assigned
+// fields like resourceVersion populated.
+func (c *LeaseClient) Create(ctx context.Context, lease *Lease) (*Lease, error) {
+	url := fmt.Sprintf("%s%s/namespaces/%s/leases", c.baseURL, leaseAPIPath, c.namespace)
+
+	lease.APIVersion = "coordination.k8s.io/v1"
+	lease.Kind = "Lease"
+	if lease.Metadata.Namespace == "" {
+		lease.Metadata.Namespace = c.namespace
+	}
+
+	resp, err := c.doRequest(ctx, http.MethodPost, url, lease)
+	if err != nil {
+		return nil, err
+	}
+	defer func() { _ = resp.Body.Close() }()
+
+	if resp.StatusCode == http.StatusConflict {
+		return nil, ErrConflict
+	}
+	if resp.StatusCode != http.StatusCreated {
+		return nil, c.readError(resp)
+	}
+
+	var created Lease
+	if err := json.NewDecoder(resp.Body).Decode(&created); err != nil {
+		return nil, fmt.Errorf("decode created lease: %w", err)
+	}
+	return &created, nil
+}
+
+// Update replaces a Lease. The lease.Metadata.ResourceVersion must match
+// the current server value (optimistic concurrency). Returns ErrConflict
+// on version mismatch.
+func (c *LeaseClient) Update(ctx context.Context, lease *Lease) (*Lease, error) {
+	url := fmt.Sprintf("%s%s/namespaces/%s/leases/%s", c.baseURL, leaseAPIPath, c.namespace, lease.Metadata.Name)
+
+	lease.APIVersion = "coordination.k8s.io/v1"
+	lease.Kind = "Lease"
+
+	resp, err := c.doRequest(ctx, http.MethodPut, url, lease)
+	if err != nil {
+		return nil, err
+	}
+	defer func() { _ = resp.Body.Close() }()
+
+	if resp.StatusCode == http.StatusConflict {
+		return nil, ErrConflict
+	}
+	if resp.StatusCode != http.StatusOK {
+		return nil, c.readError(resp)
+	}
+
+	var updated Lease
+	if err := json.NewDecoder(resp.Body).Decode(&updated); err != nil {
+		return nil, fmt.Errorf("decode updated lease: %w", err)
+	}
+	return &updated, nil
+}
+
+func (c *LeaseClient) doRequest(ctx context.Context, method, url string, body any) (*http.Response, error) {
+	token, err := readToken()
+	if err != nil {
+		return nil, fmt.Errorf("read service account token: %w", err)
+	}
+
+	var bodyReader io.Reader
+	if body != nil {
+		data, err := json.Marshal(body)
+		if err != nil {
+			return nil, fmt.Errorf("marshal request body: %w", err)
+		}
+		bodyReader = bytes.NewReader(data)
+	}
+
+	req, err := http.NewRequestWithContext(ctx, method, url, bodyReader)
+	if err != nil {
+		return nil, fmt.Errorf("create request: %w", err)
+	}
+
+	req.Header.Set("Authorization", "Bearer "+token)
+	req.Header.Set("Accept", "application/json")
+	if body != nil {
+		req.Header.Set("Content-Type", "application/json")
+	}
+
+	return c.httpClient.Do(req)
+}
+
+func readToken() (string, error) {
+	data, err := os.ReadFile(saTokenPath)
+	if err != nil {
+		return "", fmt.Errorf("read %s: %w", saTokenPath, err)
+	}
+	return strings.TrimSpace(string(data)), nil
+}
+
+func (c *LeaseClient) readError(resp *http.Response) error {
+	body, _ := io.ReadAll(io.LimitReader(resp.Body, 1024))
+	return fmt.Errorf("k8s API %s %d: %s", resp.Request.URL.Path, resp.StatusCode, string(body))
+}
+
+// LeaseNameForDomain returns a deterministic, DNS-label-safe Lease name
+// for the given domain. The domain is hashed to avoid dots and length issues.
+func LeaseNameForDomain(domain string) string {
+	h := sha256.Sum256([]byte(domain))
+	return "cert-lock-" + hex.EncodeToString(h[:8])
+}
+
+// InCluster reports whether the process is running inside a Kubernetes pod
+// by checking for the KUBERNETES_SERVICE_HOST environment variable.
+func InCluster() bool {
+	_, exists := os.LookupEnv("KUBERNETES_SERVICE_HOST")
+	return exists
+}

proxy/internal/k8s/lease_test.go

@@ -0,0 +1,102 @@
+package k8s
+
+import (
+	"encoding/json"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestLeaseNameForDomain(t *testing.T) {
+	tests := []struct {
+		domain string
+	}{
+		{"example.com"},
+		{"app.example.com"},
+		{"another.domain.io"},
+	}
+
+	seen := make(map[string]string)
+	for _, tc := range tests {
+		name := LeaseNameForDomain(tc.domain)
+
+		assert.True(t, len(name) <= 63, "must be valid DNS label length")
+		assert.Regexp(t, `^cert-lock-[0-9a-f]{16}$`, name,
+			"must match expected format for domain %q", tc.domain)
+
+		// Same input produces same output.
+		assert.Equal(t, name, LeaseNameForDomain(tc.domain), "must be deterministic")
+
+		// Different domains produce different names.
+		if prev, ok := seen[name]; ok {
+			t.Errorf("collision: %q and %q both map to %s", prev, tc.domain, name)
+		}
+		seen[name] = tc.domain
+	}
+}
+
+func TestMicroTimeJSON(t *testing.T) {
+	ts := time.Date(2024, 6, 15, 10, 30, 0, 0, time.UTC)
+	mt := &MicroTime{Time: ts}
+
+	data, err := json.Marshal(mt)
+	require.NoError(t, err)
+	assert.Equal(t, `"2024-06-15T10:30:00.000000Z"`, string(data))
+
+	var decoded MicroTime
+	require.NoError(t, json.Unmarshal(data, &decoded))
+	assert.True(t, ts.Equal(decoded.Time), "round-trip should preserve time")
+}
+
+func TestMicroTimeNullJSON(t *testing.T) {
+	// Null pointer serializes as JSON null via the Lease struct.
+	spec := LeaseSpec{
+		HolderIdentity: nil,
+		AcquireTime:    nil,
+		RenewTime:      nil,
+	}
+
+	data, err := json.Marshal(spec)
+	require.NoError(t, err)
+	assert.Contains(t, string(data), `"acquireTime":null`)
+	assert.Contains(t, string(data), `"renewTime":null`)
+}
+
+func TestLeaseJSONRoundTrip(t *testing.T) {
+	holder := "pod-abc"
+	dur := int32(300)
+	now := MicroTime{Time: time.Now().UTC().Truncate(time.Microsecond)}
+
+	original := Lease{
+		APIVersion: "coordination.k8s.io/v1",
+		Kind:       "Lease",
+		Metadata: LeaseMetadata{
+			Name:            "cert-lock-abcdef0123456789",
+			Namespace:       "default",
+			ResourceVersion: "12345",
+			Annotations: map[string]string{
+				"netbird.io/domain": "app.example.com",
+			},
+		},
+		Spec: LeaseSpec{
+			HolderIdentity:       &holder,
+			LeaseDurationSeconds: &dur,
+			AcquireTime:          &now,
+			RenewTime:            &now,
+		},
+	}
+
+	data, err := json.Marshal(original)
+	require.NoError(t, err)
+
+	var decoded Lease
+	require.NoError(t, json.Unmarshal(data, &decoded))
+
+	assert.Equal(t, original.Metadata.Name, decoded.Metadata.Name)
+	assert.Equal(t, original.Metadata.ResourceVersion, decoded.Metadata.ResourceVersion)
+	assert.Equal(t, *original.Spec.HolderIdentity, *decoded.Spec.HolderIdentity)
+	assert.Equal(t, *original.Spec.LeaseDurationSeconds, *decoded.Spec.LeaseDurationSeconds)
+	assert.True(t, original.Spec.AcquireTime.Equal(decoded.Spec.AcquireTime.Time))
+}

proxy/internal/llm/anthropic.go

@@ -0,0 +1,196 @@
+package llm
+
+import (
+	"encoding/json"
+	"fmt"
+	"strings"
+)
+
+// AnthropicParser implements the Parser interface for the Anthropic Messages
+// and Completions APIs. Detection is substring-based to tolerate upstream
+// path rewrites.
+type AnthropicParser struct{}
+
+var anthropicPathHints = []string{
+	"/v1/messages",
+	"/v1/complete",
+}
+
+// Provider returns ProviderAnthropic.
+func (AnthropicParser) Provider() Provider { return ProviderAnthropic }
+
+// ProviderName returns the stable label used for metrics and metadata.
+func (AnthropicParser) ProviderName() string { return "anthropic" }
+
+// DetectFromURL reports whether the given request path looks like an
+// Anthropic API endpoint. The match is case-insensitive and substring-based.
+func (AnthropicParser) DetectFromURL(path string) bool {
+	lower := strings.ToLower(path)
+	for _, hint := range anthropicPathHints {
+		if strings.Contains(lower, hint) {
+			return true
+		}
+	}
+	return false
+}
+
+type anthropicRequest struct {
+	Model    string             `json:"model"`
+	Stream   *bool              `json:"stream"`
+	System   json.RawMessage    `json:"system"`
+	Messages []anthropicMessage `json:"messages"`
+	// Legacy /v1/complete endpoint.
+	Prompt string `json:"prompt"`
+}
+
+type anthropicMessage struct {
+	Role    string          `json:"role"`
+	Content json.RawMessage `json:"content"`
+}
+
+// ParseRequest extracts the model name and streaming flag from an Anthropic
+// request body. Unknown or missing fields leave the corresponding struct
+// members zero-valued.
+func (AnthropicParser) ParseRequest(body []byte) (RequestFacts, error) {
+	var req anthropicRequest
+	if err := json.Unmarshal(body, &req); err != nil {
+		return RequestFacts{}, fmt.Errorf("decode anthropic request: %w: %v", ErrMalformedRequest, err)
+	}
+	return RequestFacts{
+		Model:  req.Model,
+		Stream: ptrDeref(req.Stream),
+	}, nil
+}
+
+type anthropicResponse struct {
+	Usage struct {
+		InputTokens  int64 `json:"input_tokens"`
+		OutputTokens int64 `json:"output_tokens"`
+		// CacheReadInputTokens and CacheCreationInputTokens are
+		// ADDITIVE to InputTokens (not subset), each billed at its
+		// own rate by the cost meter. cache_read is the cheaper
+		// read-from-cache rate, cache_creation is the more
+		// expensive write-to-cache rate.
+		CacheReadInputTokens     int64 `json:"cache_read_input_tokens"`
+		CacheCreationInputTokens int64 `json:"cache_creation_input_tokens"`
+	} `json:"usage"`
+}
+
+// ParseResponse decodes the non-streaming Anthropic response envelope. Status
+// codes other than 200 are treated as non-LLM responses so the caller can
+// skip cost accounting without aborting the request.
+func (AnthropicParser) ParseResponse(status int, contentType string, body []byte) (Usage, error) {
+	if status != 200 {
+		return Usage{}, fmt.Errorf("anthropic status %d: %w", status, ErrNotLLMResponse)
+	}
+	if isEventStream(contentType) {
+		return Usage{}, ErrStreamingUnsupported
+	}
+	if !isJSON(contentType) {
+		return Usage{}, fmt.Errorf("anthropic content-type %q: %w", contentType, ErrNotLLMResponse)
+	}
+
+	var resp anthropicResponse
+	if err := json.Unmarshal(body, &resp); err != nil {
+		return Usage{}, fmt.Errorf("decode anthropic response: %w: %v", ErrMalformedResponse, err)
+	}
+	return Usage{
+		InputTokens:         resp.Usage.InputTokens,
+		OutputTokens:        resp.Usage.OutputTokens,
+		TotalTokens:         resp.Usage.InputTokens + resp.Usage.OutputTokens + resp.Usage.CacheReadInputTokens + resp.Usage.CacheCreationInputTokens,
+		CachedInputTokens:   resp.Usage.CacheReadInputTokens,
+		CacheCreationTokens: resp.Usage.CacheCreationInputTokens,
+	}, nil
+}
+
+// ExtractPrompt returns the user-visible prompt text from an Anthropic
+// request body. Handles the Messages API (system + messages[]) and the
+// legacy /v1/complete prompt string. Returns "" on any decode failure.
+func (AnthropicParser) ExtractPrompt(body []byte) string {
+	var req anthropicRequest
+	if err := json.Unmarshal(body, &req); err != nil {
+		return ""
+	}
+	var b strings.Builder
+	if len(req.System) > 0 {
+		if s := decodeStringOrJoin(req.System); s != "" {
+			b.WriteString("system: ")
+			b.WriteString(s)
+		}
+	}
+	for _, m := range req.Messages {
+		if b.Len() > 0 {
+			b.WriteByte('\n')
+		}
+		if m.Role != "" {
+			b.WriteString(m.Role)
+			b.WriteString(": ")
+		}
+		b.WriteString(decodeStringOrJoin(m.Content))
+	}
+	if b.Len() == 0 && req.Prompt != "" {
+		b.WriteString(req.Prompt)
+	}
+	return b.String()
+}
+
+// ExtractSessionID is the body-side fallback for Anthropic. Claude Code's
+// authoritative session marker is the X-Claude-Code-Session-Id request
+// header (handled by the request-parser middleware); this only mines the
+// optional metadata.user_id for an embedded "...session_<uuid>" marker.
+// metadata.user_id on its own is a USER identifier, not a session, so the
+// whole value is deliberately NOT used — returning it would mislabel every
+// request from a user as one session. Returns "" when no session marker is
+// present.
+func (AnthropicParser) ExtractSessionID(body []byte) string {
+	var req struct {
+		Metadata struct {
+			UserID string `json:"user_id"`
+		} `json:"metadata"`
+	}
+	if err := json.Unmarshal(body, &req); err != nil {
+		return ""
+	}
+	if idx := strings.LastIndex(req.Metadata.UserID, "session_"); idx >= 0 {
+		if session := req.Metadata.UserID[idx+len("session_"):]; session != "" {
+			return session
+		}
+	}
+	return ""
+}
+
+type anthropicMessageResponse struct {
+	Content []struct {
+		Type string `json:"type"`
+		Text string `json:"text"`
+	} `json:"content"`
+	// Legacy /v1/complete response.
+	Completion string `json:"completion"`
+}
+
+// ExtractCompletion returns the assistant text from a non-streaming Anthropic
+// Messages or Completions response. Returns "" when status/content-type
+// indicate the body is not parseable or no text part is present.
+func (AnthropicParser) ExtractCompletion(status int, contentType string, body []byte) string {
+	if status != 200 || isEventStream(contentType) || !isJSON(contentType) {
+		return ""
+	}
+	var resp anthropicMessageResponse
+	if err := json.Unmarshal(body, &resp); err != nil {
+		return ""
+	}
+	var b strings.Builder
+	for _, part := range resp.Content {
+		if part.Text == "" {
+			continue
+		}
+		if b.Len() > 0 {
+			b.WriteByte('\n')
+		}
+		b.WriteString(part.Text)
+	}
+	if b.Len() == 0 {
+		return resp.Completion
+	}
+	return b.String()
+}

proxy/internal/llm/anthropic_test.go

@@ -0,0 +1,169 @@
+package llm
+
+import (
+	"errors"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestAnthropicDetectFromURL(t *testing.T) {
+	p := AnthropicParser{}
+
+	cases := map[string]bool{
+		"/v1/messages":         true,
+		"/v1/complete":         true,
+		"/V1/Messages":         true,
+		"/proxy/v1/messages?x": true,
+		"/v1/chat/completions": false,
+		"":                     false,
+	}
+	for path, want := range cases {
+		assert.Equal(t, want, p.DetectFromURL(path), "DetectFromURL(%q)", path)
+	}
+}
+
+func TestAnthropicParseRequest(t *testing.T) {
+	p := AnthropicParser{}
+
+	t.Run("stream true", func(t *testing.T) {
+		facts, err := p.ParseRequest([]byte(`{"model":"claude-sonnet-4-5","stream":true}`))
+		require.NoError(t, err)
+		assert.Equal(t, "claude-sonnet-4-5", facts.Model, "model extracted")
+		assert.True(t, facts.Stream, "stream flag honoured")
+	})
+
+	t.Run("stream default", func(t *testing.T) {
+		facts, err := p.ParseRequest([]byte(`{"model":"claude-sonnet-4-5"}`))
+		require.NoError(t, err)
+		assert.False(t, facts.Stream, "missing stream flag defaults to false")
+	})
+
+	t.Run("malformed", func(t *testing.T) {
+		_, err := p.ParseRequest([]byte(`{"model":`))
+		require.Error(t, err)
+		assert.True(t, errors.Is(err, ErrMalformedRequest), "sentinel wrapped")
+	})
+}
+
+func TestAnthropicParseResponse(t *testing.T) {
+	p := AnthropicParser{}
+
+	t.Run("happy fixture", func(t *testing.T) {
+		body, err := os.ReadFile(filepath.Join("fixtures", "anthropic_messages.json"))
+		require.NoError(t, err, "fixture must be readable")
+
+		usage, err := p.ParseResponse(200, "application/json", body)
+		require.NoError(t, err)
+		assert.Equal(t, int64(123), usage.InputTokens, "input tokens extracted")
+		assert.Equal(t, int64(45), usage.OutputTokens, "output tokens extracted")
+		assert.Equal(t, int64(168), usage.TotalTokens, "total computed as sum")
+	})
+
+	t.Run("streaming rejected", func(t *testing.T) {
+		_, err := p.ParseResponse(200, "text/event-stream", []byte(""))
+		require.ErrorIs(t, err, ErrStreamingUnsupported, "SSE responses must use the scanner")
+	})
+
+	t.Run("non-200", func(t *testing.T) {
+		_, err := p.ParseResponse(429, "application/json", []byte(`{}`))
+		require.ErrorIs(t, err, ErrNotLLMResponse, "non-200 rejected as non-LLM")
+	})
+
+	t.Run("non-json content type", func(t *testing.T) {
+		_, err := p.ParseResponse(200, "text/html", []byte(`{}`))
+		require.ErrorIs(t, err, ErrNotLLMResponse, "text/html treated as non-LLM")
+	})
+
+	t.Run("malformed body", func(t *testing.T) {
+		_, err := p.ParseResponse(200, "application/json", []byte(`{`))
+		require.ErrorIs(t, err, ErrMalformedResponse, "bad JSON yields malformed error")
+	})
+
+	// Anthropic's two cache fields are ADDITIVE to input_tokens (not
+	// subset). The parser must surface them so the cost meter can
+	// bill each bucket at its own configured rate. Total includes
+	// every bucket so downstream attribution sees the full token
+	// volume the request consumed.
+	t.Run("cache_read_input_tokens surfaces as CachedInputTokens (additive)", func(t *testing.T) {
+		body := []byte(`{"usage":{"input_tokens":256,"output_tokens":200,"cache_read_input_tokens":768}}`)
+		usage, err := p.ParseResponse(200, "application/json", body)
+		require.NoError(t, err)
+		assert.Equal(t, int64(256), usage.InputTokens, "regular input remains separate from cache buckets")
+		assert.Equal(t, int64(768), usage.CachedInputTokens, "cache_read maps onto CachedInputTokens — same field carries OpenAI cached subset and Anthropic cache reads")
+		assert.Zero(t, usage.CacheCreationTokens)
+		assert.Equal(t, int64(256+200+768), usage.TotalTokens, "total includes every input bucket plus output — cache reads are billable tokens")
+	})
+
+	t.Run("cache_creation_input_tokens surfaces as CacheCreationTokens (additive)", func(t *testing.T) {
+		body := []byte(`{"usage":{"input_tokens":256,"output_tokens":200,"cache_creation_input_tokens":512}}`)
+		usage, err := p.ParseResponse(200, "application/json", body)
+		require.NoError(t, err)
+		assert.Equal(t, int64(256), usage.InputTokens)
+		assert.Zero(t, usage.CachedInputTokens)
+		assert.Equal(t, int64(512), usage.CacheCreationTokens, "cache_creation surfaces — meter applies the write-rate multiplier")
+		assert.Equal(t, int64(256+200+512), usage.TotalTokens)
+	})
+
+	t.Run("both cache buckets present", func(t *testing.T) {
+		body := []byte(`{"usage":{"input_tokens":256,"output_tokens":200,"cache_read_input_tokens":768,"cache_creation_input_tokens":512}}`)
+		usage, err := p.ParseResponse(200, "application/json", body)
+		require.NoError(t, err)
+		assert.Equal(t, int64(768), usage.CachedInputTokens)
+		assert.Equal(t, int64(512), usage.CacheCreationTokens)
+		assert.Equal(t, int64(256+200+768+512), usage.TotalTokens, "all four buckets sum into total")
+	})
+
+	t.Run("absent cache fields leave counts at zero", func(t *testing.T) {
+		body := []byte(`{"usage":{"input_tokens":100,"output_tokens":50}}`)
+		usage, err := p.ParseResponse(200, "application/json", body)
+		require.NoError(t, err)
+		assert.Zero(t, usage.CachedInputTokens, "no cache_read field = no cached count")
+		assert.Zero(t, usage.CacheCreationTokens, "no cache_creation field = no creation count")
+		assert.Equal(t, int64(150), usage.TotalTokens, "back to the simple in+out total when no cache buckets present")
+	})
+}
+
+func TestAnthropicExtractPrompt_Messages(t *testing.T) {
+	body := []byte(`{"model":"claude-sonnet-4-7","system":"be brief","messages":[{"role":"user","content":"hi"},{"role":"assistant","content":"yes"}]}`)
+	got := AnthropicParser{}.ExtractPrompt(body)
+	require.Contains(t, got, "system: be brief", "system surfaces with role label")
+	require.Contains(t, got, "user: hi", "user message surfaces")
+	require.Contains(t, got, "assistant: yes", "assistant message surfaces")
+}
+
+func TestAnthropicExtractPrompt_LegacyComplete(t *testing.T) {
+	body := []byte(`{"model":"claude-2","prompt":"\n\nHuman: hi\n\nAssistant:"}`)
+	got := AnthropicParser{}.ExtractPrompt(body)
+	require.Contains(t, got, "Human: hi", "legacy prompt string surfaces")
+}
+
+func TestAnthropicExtractSessionID(t *testing.T) {
+	t.Run("claude code session suffix", func(t *testing.T) {
+		body := []byte(`{"model":"claude-opus-4-8","metadata":{"user_id":"user_abc123_account_def456_session_9f8e7d6c"},"messages":[]}`)
+		assert.Equal(t, "9f8e7d6c", AnthropicParser{}.ExtractSessionID(body), "session_<id> suffix must be extracted from metadata.user_id")
+	})
+	t.Run("plain user_id is not treated as a session", func(t *testing.T) {
+		body := []byte(`{"model":"claude-opus-4-8","metadata":{"user_id":"acme-team"},"messages":[]}`)
+		assert.Equal(t, "", AnthropicParser{}.ExtractSessionID(body), "a user identifier without a session marker must NOT be used as a session id")
+	})
+	t.Run("no metadata yields empty", func(t *testing.T) {
+		body := []byte(`{"model":"claude-opus-4-8","messages":[{"role":"user","content":"hi"}]}`)
+		assert.Equal(t, "", AnthropicParser{}.ExtractSessionID(body), "absent metadata.user_id yields no session id")
+	})
+}
+
+func TestAnthropicExtractCompletion_Messages(t *testing.T) {
+	body, err := os.ReadFile(filepath.Join("fixtures", "anthropic_messages.json"))
+	require.NoError(t, err)
+	got := AnthropicParser{}.ExtractCompletion(200, "application/json", body)
+	require.NotEmpty(t, got, "anthropic fixture has assistant text")
+}
+
+func TestAnthropicExtractCompletion_Streaming(t *testing.T) {
+	got := AnthropicParser{}.ExtractCompletion(200, "text/event-stream", []byte(""))
+	require.Empty(t, got, "streaming responses are skipped")
+}

proxy/internal/llm/bedrock.go

@@ -0,0 +1,193 @@
+package llm
+
+import (
+	"encoding/json"
+	"fmt"
+	"strings"
+)
+
+// ProviderNameBedrock is the stable label for the AWS Bedrock parser, used as
+// the llm.provider metadata value and the cost-meter formula selector.
+const ProviderNameBedrock = "bedrock"
+
+// BedrockParser implements the Parser interface for the AWS Bedrock runtime.
+// Bedrock carries the model in the URL path (/model/{id}/{action}); the request
+// middleware extracts it there, so this parser focuses on the response shapes:
+// the vendor-native InvokeModel body (e.g. Anthropic's snake_case usage) and the
+// unified Converse body (camelCase usage).
+type BedrockParser struct{}
+
+var bedrockPathHints = []string{"/invoke", "/converse"}
+
+// Provider returns ProviderBedrock.
+func (BedrockParser) Provider() Provider { return ProviderBedrock }
+
+// ProviderName returns the stable label used for metrics and metadata.
+func (BedrockParser) ProviderName() string { return ProviderNameBedrock }
+
+// DetectFromURL reports whether the path is a Bedrock runtime model endpoint.
+func (BedrockParser) DetectFromURL(path string) bool {
+	lower := strings.ToLower(path)
+	if !strings.HasPrefix(lower, "/model/") {
+		return false
+	}
+	for _, hint := range bedrockPathHints {
+		if strings.Contains(lower, hint) {
+			return true
+		}
+	}
+	return false
+}
+
+// ParseRequest is a no-op for Bedrock: the model lives in the URL path, not the
+// body, and the streaming flag is derived from the path action. The request
+// middleware handles both via parseBedrockPath, so this returns empty facts.
+func (BedrockParser) ParseRequest([]byte) (RequestFacts, error) {
+	return RequestFacts{}, nil
+}
+
+// bedrockResponse captures token usage from both Bedrock response shapes:
+// InvokeModel (vendor-native; Anthropic uses snake_case + additive cache
+// buckets) and Converse (camelCase, with a precomputed total).
+type bedrockResponse struct {
+	Usage struct {
+		// InvokeModel (Anthropic-on-Bedrock) — snake_case.
+		InputTokens              int64 `json:"input_tokens"`
+		OutputTokens             int64 `json:"output_tokens"`
+		CacheReadInputTokens     int64 `json:"cache_read_input_tokens"`
+		CacheCreationInputTokens int64 `json:"cache_creation_input_tokens"`
+		// Converse — camelCase; cache buckets are additive to inputTokens (AWS names the write bucket cacheWriteInputTokens).
+		InputTokensCamel      int64 `json:"inputTokens"`
+		OutputTokensCamel     int64 `json:"outputTokens"`
+		TotalTokensCamel      int64 `json:"totalTokens"`
+		CacheReadTokensCamel  int64 `json:"cacheReadInputTokens"`
+		CacheWriteTokensCamel int64 `json:"cacheWriteInputTokens"`
+	} `json:"usage"`
+}
+
+// ParseResponse decodes the non-streaming Bedrock response envelope, handling
+// both the InvokeModel and Converse usage shapes. Non-200 / non-JSON bodies are
+// treated as non-LLM responses so the caller skips cost accounting.
+func (BedrockParser) ParseResponse(status int, contentType string, body []byte) (Usage, error) {
+	if status != 200 {
+		return Usage{}, fmt.Errorf("bedrock status %d: %w", status, ErrNotLLMResponse)
+	}
+	if isAWSEventStream(contentType) || isEventStream(contentType) {
+		return Usage{}, ErrStreamingUnsupported
+	}
+	if !isJSON(contentType) {
+		return Usage{}, fmt.Errorf("bedrock content-type %q: %w", contentType, ErrNotLLMResponse)
+	}
+
+	var resp bedrockResponse
+	if err := json.Unmarshal(body, &resp); err != nil {
+		return Usage{}, fmt.Errorf("decode bedrock response: %w: %v", ErrMalformedResponse, err)
+	}
+	inTok := firstNonZero(resp.Usage.InputTokens, resp.Usage.InputTokensCamel)
+	outTok := firstNonZero(resp.Usage.OutputTokens, resp.Usage.OutputTokensCamel)
+	cacheRead := firstNonZero(resp.Usage.CacheReadInputTokens, resp.Usage.CacheReadTokensCamel)
+	cacheWrite := firstNonZero(resp.Usage.CacheCreationInputTokens, resp.Usage.CacheWriteTokensCamel)
+	total := resp.Usage.TotalTokensCamel
+	if total == 0 {
+		total = inTok + outTok + cacheRead + cacheWrite
+	}
+	return Usage{
+		InputTokens:         inTok,
+		OutputTokens:        outTok,
+		TotalTokens:         total,
+		CachedInputTokens:   cacheRead,
+		CacheCreationTokens: cacheWrite,
+	}, nil
+}
+
+// ExtractPrompt returns the user-visible prompt from a Bedrock request body,
+// handling both the InvokeModel (Anthropic Messages: system + messages[]) and
+// Converse (messages[].content[].text) shapes. Returns "" on decode failure.
+func (BedrockParser) ExtractPrompt(body []byte) string {
+	var req struct {
+		System   json.RawMessage `json:"system"`
+		Messages []struct {
+			Role    string          `json:"role"`
+			Content json.RawMessage `json:"content"`
+		} `json:"messages"`
+	}
+	if err := json.Unmarshal(body, &req); err != nil {
+		return ""
+	}
+	var b strings.Builder
+	if s := decodeStringOrJoin(req.System); s != "" {
+		b.WriteString("system: ")
+		b.WriteString(s)
+	}
+	for _, m := range req.Messages {
+		if b.Len() > 0 {
+			b.WriteByte('\n')
+		}
+		if m.Role != "" {
+			b.WriteString(m.Role)
+			b.WriteString(": ")
+		}
+		b.WriteString(decodeStringOrJoin(m.Content))
+	}
+	return b.String()
+}
+
+// ExtractCompletion returns the assistant text from a non-streaming Bedrock
+// response, handling InvokeModel (Anthropic content[].text) and Converse
+// (output.message.content[].text).
+func (BedrockParser) ExtractCompletion(status int, contentType string, body []byte) string {
+	if status != 200 || isAWSEventStream(contentType) || !isJSON(contentType) {
+		return ""
+	}
+	var resp struct {
+		Content []struct {
+			Text string `json:"text"`
+		} `json:"content"`
+		Output struct {
+			Message struct {
+				Content []struct {
+					Text string `json:"text"`
+				} `json:"content"`
+			} `json:"message"`
+		} `json:"output"`
+	}
+	if err := json.Unmarshal(body, &resp); err != nil {
+		return ""
+	}
+	var b strings.Builder
+	appendText := func(text string) {
+		if text == "" {
+			return
+		}
+		if b.Len() > 0 {
+			b.WriteByte('\n')
+		}
+		b.WriteString(text)
+	}
+	for _, p := range resp.Content {
+		appendText(p.Text)
+	}
+	for _, p := range resp.Output.Message.Content {
+		appendText(p.Text)
+	}
+	return b.String()
+}
+
+// ExtractSessionID has no Bedrock-native marker; session grouping relies on the
+// request headers handled by the middleware. Returns "".
+func (BedrockParser) ExtractSessionID([]byte) string { return "" }
+
+// firstNonZero returns a when non-zero, else b. Folds the snake_case and
+// camelCase usage variants into a single value.
+func firstNonZero(a, b int64) int64 {
+	if a != 0 {
+		return a
+	}
+	return b
+}
+
+// isAWSEventStream reports whether contentType is the AWS binary event-stream
+// framing used by Bedrock's streaming endpoints.
+func isAWSEventStream(contentType string) bool {
+	return strings.Contains(strings.ToLower(contentType), "application/vnd.amazon.eventstream")
+}

proxy/internal/llm/bedrock_model.go

@@ -0,0 +1,38 @@
+package llm
+
+import (
+	"regexp"
+	"strings"
+)
+
+// bedrockRegionPrefixes are the cross-region inference-profile prefixes that
+// front a Bedrock model id (e.g. "eu.anthropic.claude-...").
+var bedrockRegionPrefixes = []string{"us.", "eu.", "apac.", "global."}
+
+// bedrockVersionSuffix matches the trailing "-vN[:N]" or "-YYYYMMDD-vN[:N]"
+// version/throughput suffix of a Bedrock model id.
+var bedrockVersionSuffix = regexp.MustCompile(`-(\d{8}-)?v\d+(:\d+)?$`)
+
+// NormalizeBedrockModel strips an ARN wrapper, a cross-region inference-profile
+// prefix, and the version/throughput suffix from a Bedrock model id so it
+// matches the catalog/pricing key, e.g.
+// "eu.anthropic.claude-sonnet-4-5-20250929-v1:0" -> "anthropic.claude-sonnet-4-5"
+// and the inference-profile ARN's last segment likewise. It is the single
+// source of truth shared by the request parser (which normalizes the request
+// model from the URL path) and the router (which normalizes the operator's
+// registered Bedrock model ids so both sides compare equal).
+func NormalizeBedrockModel(modelID string) string {
+	m := modelID
+	if strings.HasPrefix(m, "arn:") {
+		if i := strings.LastIndex(m, "/"); i >= 0 {
+			m = m[i+1:]
+		}
+	}
+	for _, p := range bedrockRegionPrefixes {
+		if strings.HasPrefix(m, p) {
+			m = m[len(p):]
+			break
+		}
+	}
+	return bedrockVersionSuffix.ReplaceAllString(m, "")
+}

proxy/internal/llm/bedrock_model_test.go

@@ -0,0 +1,23 @@
+package llm
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestNormalizeBedrockModel(t *testing.T) {
+	cases := map[string]string{
+		"eu.anthropic.claude-sonnet-4-5-20250929-v1:0": "anthropic.claude-sonnet-4-5",
+		"us.anthropic.claude-haiku-4-5":                "anthropic.claude-haiku-4-5",
+		"us.anthropic.claude-opus-4-8-20250101-v1:0":   "anthropic.claude-opus-4-8",
+		"anthropic.claude-sonnet-4-5-20250929-v1:0":    "anthropic.claude-sonnet-4-5",
+		"meta.llama3-3-70b-instruct-v1:0":              "meta.llama3-3-70b-instruct",
+		"amazon.nova-pro-v1:0":                         "amazon.nova-pro",
+		// Inference-profile ARN — model id lives in the last path segment.
+		"arn:aws:bedrock:eu-central-1:123456789012:inference-profile/eu.anthropic.claude-sonnet-4-5-20250929-v1:0": "anthropic.claude-sonnet-4-5",
+	}
+	for in, want := range cases {
+		require.Equal(t, want, NormalizeBedrockModel(in), "normalize %q", in)
+	}
+}

proxy/internal/llm/bedrock_test.go

@@ -0,0 +1,77 @@
+package llm
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestBedrockParser_ParseResponse_Invoke(t *testing.T) {
+	body := []byte(`{"usage":{"input_tokens":13,"output_tokens":5,"cache_read_input_tokens":2,"cache_creation_input_tokens":4}}`)
+	u, err := BedrockParser{}.ParseResponse(200, "application/json", body)
+	require.NoError(t, err)
+	require.Equal(t, int64(13), u.InputTokens, "invoke input tokens")
+	require.Equal(t, int64(5), u.OutputTokens, "invoke output tokens")
+	require.Equal(t, int64(2), u.CachedInputTokens, "invoke cache-read tokens")
+	require.Equal(t, int64(4), u.CacheCreationTokens, "invoke cache-creation tokens")
+	require.Equal(t, int64(13+5+2+4), u.TotalTokens, "invoke total is additive")
+}
+
+func TestBedrockParser_ParseResponse_Converse(t *testing.T) {
+	body := []byte(`{"output":{"message":{"content":[{"text":"pong"}]}},"usage":{"inputTokens":11,"outputTokens":3,"totalTokens":14}}`)
+	u, err := BedrockParser{}.ParseResponse(200, "application/json", body)
+	require.NoError(t, err)
+	require.Equal(t, int64(11), u.InputTokens, "converse camelCase input tokens")
+	require.Equal(t, int64(3), u.OutputTokens, "converse camelCase output tokens")
+	require.Equal(t, int64(14), u.TotalTokens, "converse uses provider total")
+}
+
+// Converse camelCase cache fields must land in the billed Usage buckets, same as the InvokeModel snake_case fields.
+func TestBedrockParser_ParseResponse_ConverseCacheBuckets(t *testing.T) {
+	body := []byte(`{"usage":{"inputTokens":11,"outputTokens":3,"cacheReadInputTokens":7,"cacheWriteInputTokens":9}}`)
+	u, err := BedrockParser{}.ParseResponse(200, "application/json", body)
+	require.NoError(t, err)
+	require.Equal(t, int64(11), u.InputTokens, "converse input tokens")
+	require.Equal(t, int64(3), u.OutputTokens, "converse output tokens")
+	require.Equal(t, int64(7), u.CachedInputTokens, "converse cache-read tokens")
+	require.Equal(t, int64(9), u.CacheCreationTokens, "converse cache-write tokens")
+	require.Equal(t, int64(11+3+7+9), u.TotalTokens, "total backfill is additive when the provider omits totalTokens")
+}
+
+func TestBedrockParser_ParseResponse_StreamingUnsupported(t *testing.T) {
+	_, err := BedrockParser{}.ParseResponse(200, "application/vnd.amazon.eventstream", []byte("binary"))
+	require.ErrorIs(t, err, ErrStreamingUnsupported, "event-stream must route to the streaming accumulator")
+}
+
+func TestBedrockParser_ParseResponse_NonSuccess(t *testing.T) {
+	_, err := BedrockParser{}.ParseResponse(404, "application/json", []byte(`{"message":"gated"}`))
+	require.ErrorIs(t, err, ErrNotLLMResponse, "non-200 is not an LLM response")
+}
+
+func TestBedrockParser_ExtractCompletion(t *testing.T) {
+	invoke := BedrockParser{}.ExtractCompletion(200, "application/json", []byte(`{"content":[{"text":"a"},{"text":"b"}]}`))
+	require.Equal(t, "a\nb", invoke, "invoke completion joins content parts")
+
+	converse := BedrockParser{}.ExtractCompletion(200, "application/json", []byte(`{"output":{"message":{"content":[{"text":"x"}]}}}`))
+	require.Equal(t, "x", converse, "converse completion reads output.message.content")
+}
+
+func TestBedrockParser_ExtractPrompt(t *testing.T) {
+	invoke := BedrockParser{}.ExtractPrompt([]byte(`{"messages":[{"role":"user","content":"hi"}]}`))
+	require.Equal(t, "user: hi", invoke, "invoke prompt reads anthropic content string")
+
+	converse := BedrockParser{}.ExtractPrompt([]byte(`{"messages":[{"role":"user","content":[{"text":"hello"}]}]}`))
+	require.Equal(t, "user: hello", converse, "converse prompt reads content parts")
+}
+
+func TestBedrockParser_DetectFromURL(t *testing.T) {
+	require.True(t, BedrockParser{}.DetectFromURL("/model/eu.anthropic.claude/invoke"), "invoke path")
+	require.True(t, BedrockParser{}.DetectFromURL("/model/x/converse-stream"), "converse-stream path")
+	require.False(t, BedrockParser{}.DetectFromURL("/v1/chat/completions"), "openai path is not bedrock")
+}
+
+func TestBedrockParser_RegisteredByName(t *testing.T) {
+	p, ok := ParserByName(ProviderNameBedrock)
+	require.True(t, ok, "bedrock parser is registered")
+	require.Equal(t, ProviderNameBedrock, p.ProviderName())
+}

proxy/internal/llm/errors.go

@@ -0,0 +1,31 @@
+package llm
+
+import "errors"
+
+// Sentinel errors returned by parsers and the pricing loader. Callers use
+// errors.Is to branch on a condition without coupling to parser internals.
+var (
+	// ErrUnknownProvider indicates no parser claimed the request path.
+	ErrUnknownProvider = errors.New("llmobs: unknown provider")
+
+	// ErrUnsupportedModel indicates the response parsed successfully but the
+	// model is absent from the pricing table. Token counts are still valid.
+	ErrUnsupportedModel = errors.New("llmobs: unsupported model")
+
+	// ErrNotLLMResponse indicates the response is not a JSON success body
+	// that a non-streaming parser can consume (non-200 or wrong content type).
+	ErrNotLLMResponse = errors.New("llmobs: not an LLM response")
+
+	// ErrStreamingUnsupported indicates the caller passed an SSE response to
+	// a non-streaming parser. Streaming is handled separately via the SSE
+	// scanner.
+	ErrStreamingUnsupported = errors.New("llmobs: streaming response requires SSE scanner")
+
+	// ErrMalformedResponse indicates the response body could not be decoded
+	// as the provider-specific JSON schema.
+	ErrMalformedResponse = errors.New("llmobs: malformed response body")
+
+	// ErrMalformedRequest indicates the request body could not be decoded as
+	// the provider-specific JSON schema.
+	ErrMalformedRequest = errors.New("llmobs: malformed request body")
+)

proxy/internal/llm/fixtures/anthropic_messages.json

@@ -0,0 +1,17 @@
+{
+  "id": "msg_abc",
+  "type": "message",
+  "role": "assistant",
+  "model": "claude-sonnet-4-5",
+  "content": [
+    {
+      "type": "text",
+      "text": "Hello, world!"
+    }
+  ],
+  "stop_reason": "end_turn",
+  "usage": {
+    "input_tokens": 123,
+    "output_tokens": 45
+  }
+}

proxy/internal/llm/fixtures/anthropic_stream.txt

@@ -0,0 +1,21 @@
+event: message_start
+data: {"type":"message_start","message":{"id":"msg_abc","type":"message","role":"assistant","model":"claude-sonnet-4-5","content":[],"stop_reason":null,"usage":{"input_tokens":123,"output_tokens":1}}}
+
+event: content_block_start
+data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}
+
+event: content_block_delta
+data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}
+
+event: content_block_delta
+data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":", world!"}}
+
+event: content_block_stop
+data: {"type":"content_block_stop","index":0}
+
+event: message_delta
+data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":45}}
+
+event: message_stop
+data: {"type":"message_stop"}
+

proxy/internal/llm/fixtures/openai_chat_completion.json

@@ -0,0 +1,21 @@
+{
+  "id": "chatcmpl-abc",
+  "object": "chat.completion",
+  "created": 1700000000,
+  "model": "gpt-4o-mini",
+  "choices": [
+    {
+      "index": 0,
+      "message": {
+        "role": "assistant",
+        "content": "Hello, world!"
+      },
+      "finish_reason": "stop"
+    }
+  ],
+  "usage": {
+    "prompt_tokens": 123,
+    "completion_tokens": 45,
+    "total_tokens": 168
+  }
+}

proxy/internal/llm/fixtures/openai_responses.json

@@ -0,0 +1,24 @@
+{
+  "id": "resp_abc",
+  "object": "response",
+  "created_at": 1700000000,
+  "model": "gpt-5.4",
+  "output": [
+    {
+      "type": "message",
+      "role": "assistant",
+      "content": [{"type": "output_text", "text": "ok"}]
+    }
+  ],
+  "usage": {
+    "input_tokens": 15,
+    "input_tokens_details": {
+      "cached_tokens": 0
+    },
+    "output_tokens": 414,
+    "output_tokens_details": {
+      "reasoning_tokens": 0
+    },
+    "total_tokens": 429
+  }
+}

proxy/internal/llm/fixtures/openai_responses_stream.txt

@@ -0,0 +1,24 @@
+event: response.created
+data: {"type":"response.created","response":{"id":"resp_abc","object":"response","model":"gpt-5.5","usage":null}}
+
+event: response.in_progress
+data: {"type":"response.in_progress","response":{"id":"resp_abc","usage":null}}
+
+event: response.output_item.added
+data: {"type":"response.output_item.added","output_index":0,"item":{"type":"message","role":"assistant","content":[]}}
+
+event: response.content_part.added
+data: {"type":"response.content_part.added","item_id":"msg_1","output_index":0,"content_index":0,"part":{"type":"output_text","text":""}}
+
+event: response.output_text.delta
+data: {"type":"response.output_text.delta","item_id":"msg_1","output_index":0,"content_index":0,"delta":"Hello"}
+
+event: response.output_text.delta
+data: {"type":"response.output_text.delta","item_id":"msg_1","output_index":0,"content_index":0,"delta":", world!"}
+
+event: response.output_text.done
+data: {"type":"response.output_text.done","item_id":"msg_1","output_index":0,"content_index":0,"text":"Hello, world!"}
+
+event: response.completed
+data: {"type":"response.completed","response":{"id":"resp_abc","object":"response","model":"gpt-5.5","usage":{"input_tokens":123,"input_tokens_details":{"cached_tokens":40},"output_tokens":45,"output_tokens_details":{"reasoning_tokens":12},"total_tokens":168}}}
+

proxy/internal/llm/fixtures/openai_stream.txt

@@ -0,0 +1,8 @@
+data: {"id":"chatcmpl-abc","object":"chat.completion.chunk","created":1700000000,"model":"gpt-4o-mini","choices":[{"index":0,"delta":{"role":"assistant","content":"Hello"},"finish_reason":null}]}
+
+data: {"id":"chatcmpl-abc","object":"chat.completion.chunk","created":1700000000,"model":"gpt-4o-mini","choices":[{"index":0,"delta":{"content":", world!"},"finish_reason":null}]}
+
+data: {"id":"chatcmpl-abc","object":"chat.completion.chunk","created":1700000000,"model":"gpt-4o-mini","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":123,"completion_tokens":45,"total_tokens":168}}
+
+data: [DONE]
+

proxy/internal/llm/fixtures/pricing.yaml

@@ -0,0 +1,59 @@
+# Realistic-pricing starter for llm_observability. Drop this into the
+# directory you point the proxy at via --plugin-data-dir, then reference it
+# from the target's plugin config:
+#
+#   plugins:
+#     - id: llm_observability
+#       enabled: true
+#       params:
+#         pricing_path: pricing.yaml
+#
+# Values are USD per 1_000 tokens. Public list prices drift; treat this as a
+# starting point and keep your production copy current.
+
+openai:
+  # GPT-5 family
+  gpt-5:
+    input_per_1k: 0.00125
+    output_per_1k: 0.01
+  gpt-5-mini:
+    input_per_1k: 0.00025
+    output_per_1k: 0.002
+  gpt-5-nano:
+    input_per_1k: 0.00005
+    output_per_1k: 0.0004
+  gpt-5.4:
+    input_per_1k: 0.00125
+    output_per_1k: 0.01
+  # GPT-4o family
+  gpt-4o:
+    input_per_1k: 0.0025
+    output_per_1k: 0.01
+  gpt-4o-mini:
+    input_per_1k: 0.00015
+    output_per_1k: 0.0006
+  # Embeddings
+  text-embedding-3-large:
+    input_per_1k: 0.00013
+    output_per_1k: 0
+  text-embedding-3-small:
+    input_per_1k: 0.00002
+    output_per_1k: 0
+
+anthropic:
+  # Claude 4.x family
+  claude-opus-4-7:
+    input_per_1k: 0.015
+    output_per_1k: 0.075
+  claude-sonnet-4-7:
+    input_per_1k: 0.003
+    output_per_1k: 0.015
+  claude-sonnet-4-6:
+    input_per_1k: 0.003
+    output_per_1k: 0.015
+  claude-sonnet-4-5:
+    input_per_1k: 0.003
+    output_per_1k: 0.015
+  claude-haiku-4-5:
+    input_per_1k: 0.0008
+    output_per_1k: 0.004

proxy/internal/llm/openai.go

@@ -0,0 +1,412 @@
+package llm
+
+import (
+	"encoding/json"
+	"fmt"
+	"strings"
+)
+
+// OpenAIParser implements the Parser interface for OpenAI-compatible APIs.
+// It recognizes chat.completions, completions, embeddings, and the newer
+// responses endpoint; any proxy path-prefix stripping is tolerated by the
+// substring match in DetectFromURL.
+type OpenAIParser struct{}
+
+// openAIPathHints are substring patterns that mark a request as
+// OpenAI-shaped. The bare `/chat/completions` is listed alongside
+// `/v1/chat/completions` because gateways like Cloudflare AI
+// Gateway place their own version segment before the provider
+// slug (gateway/v1/{account}/{gateway}/openai/chat/completions) —
+// the canonical `/v1/` ends up nowhere near `/chat/completions`,
+// so the `/v1/chat/completions` hint misses. `/chat/completions`
+// is OpenAI's API contract: any service accepting OpenAI bodies
+// serves at this path, so false-positive risk is negligible.
+// `/completions` (legacy), `/embeddings`, and `/responses` are
+// kept on the canonical-only path because their bare forms are
+// too generic to be safe substrings.
+var openAIPathHints = []string{
+	"/v1/chat/completions",
+	"/v1/completions",
+	"/v1/embeddings",
+	"/v1/responses",
+	"/chat/completions",
+}
+
+// Provider returns ProviderOpenAI.
+func (OpenAIParser) Provider() Provider { return ProviderOpenAI }
+
+// ProviderName returns the stable label used for metrics and metadata.
+func (OpenAIParser) ProviderName() string { return "openai" }
+
+// DetectFromURL reports whether the given request path looks like an OpenAI
+// API endpoint. The match is case-insensitive and substring-based so that a
+// reverse proxy prefix strip or rewrite does not defeat detection.
+func (OpenAIParser) DetectFromURL(path string) bool {
+	lower := strings.ToLower(path)
+	for _, hint := range openAIPathHints {
+		if strings.Contains(lower, hint) {
+			return true
+		}
+	}
+	return false
+}
+
+type openAIRequest struct {
+	Model         string `json:"model"`
+	Stream        *bool  `json:"stream"`
+	StreamOptions *struct {
+		IncludeUsage *bool `json:"include_usage"`
+	} `json:"stream_options"`
+	// Chat Completions / Completions: messages[].content (string or array of
+	// content parts). Responses API: input is either a string or an array of
+	// items with content parts. We use json.RawMessage to defer parsing each
+	// shape independently.
+	Messages []openAIMessage `json:"messages"`
+	Prompt   json.RawMessage `json:"prompt"`
+	Input    json.RawMessage `json:"input"`
+}
+
+type openAIMessage struct {
+	Role    string          `json:"role"`
+	Content json.RawMessage `json:"content"`
+}
+
+// ParseRequest extracts the model name and streaming flag from an OpenAI
+// request body. Unknown or missing fields leave the corresponding struct
+// members zero-valued.
+func (OpenAIParser) ParseRequest(body []byte) (RequestFacts, error) {
+	var req openAIRequest
+	if err := json.Unmarshal(body, &req); err != nil {
+		return RequestFacts{}, fmt.Errorf("decode openai request: %w: %v", ErrMalformedRequest, err)
+	}
+	return RequestFacts{
+		Model:  req.Model,
+		Stream: ptrDeref(req.Stream),
+	}, nil
+}
+
+// openAIResponse accepts both naming conventions in a single struct because
+// OpenAI's older Chat Completions API uses prompt_tokens/completion_tokens
+// while the newer Responses API (/v1/responses) uses input_tokens/output_tokens
+// (aligned with Anthropic). Pointer fields let us tell "absent" from "zero".
+//
+// PromptTokensDetails.CachedTokens (Chat Completions) and
+// InputTokensDetails.CachedTokens (Responses API) carry the SUBSET of
+// prompt/input tokens that hit the prompt cache. Cost-meter applies the
+// discount rate to that subset and the regular rate to the remainder so
+// we never double-bill the cached portion.
+type openAIResponse struct {
+	Usage struct {
+		PromptTokens        *int64 `json:"prompt_tokens"`
+		CompletionTokens    *int64 `json:"completion_tokens"`
+		InputTokens         *int64 `json:"input_tokens"`
+		OutputTokens        *int64 `json:"output_tokens"`
+		TotalTokens         *int64 `json:"total_tokens"`
+		PromptTokensDetails *struct {
+			CachedTokens *int64 `json:"cached_tokens"`
+		} `json:"prompt_tokens_details"`
+		InputTokensDetails *struct {
+			CachedTokens *int64 `json:"cached_tokens"`
+		} `json:"input_tokens_details"`
+	} `json:"usage"`
+}
+
+// ParseResponse decodes the non-streaming OpenAI response envelope. Status
+// codes other than 200 are treated as non-LLM responses so the caller can
+// skip cost accounting without aborting the request.
+func (OpenAIParser) ParseResponse(status int, contentType string, body []byte) (Usage, error) {
+	if status != 200 {
+		return Usage{}, fmt.Errorf("openai status %d: %w", status, ErrNotLLMResponse)
+	}
+	if isEventStream(contentType) {
+		return Usage{}, ErrStreamingUnsupported
+	}
+	if !isJSON(contentType) {
+		return Usage{}, fmt.Errorf("openai content-type %q: %w", contentType, ErrNotLLMResponse)
+	}
+
+	var resp openAIResponse
+	if err := json.Unmarshal(body, &resp); err != nil {
+		return Usage{}, fmt.Errorf("decode openai response: %w: %v", ErrMalformedResponse, err)
+	}
+
+	// Responses-API names take precedence when present; fall back to the older
+	// Chat Completions names. This handles both endpoints transparently
+	// without forcing a per-route configuration.
+	u := Usage{
+		InputTokens:       pickInt64(resp.Usage.InputTokens, resp.Usage.PromptTokens),
+		OutputTokens:      pickInt64(resp.Usage.OutputTokens, resp.Usage.CompletionTokens),
+		TotalTokens:       derefInt64(resp.Usage.TotalTokens),
+		CachedInputTokens: openAICachedTokens(resp),
+	}
+	if u.TotalTokens == 0 && (u.InputTokens > 0 || u.OutputTokens > 0) {
+		u.TotalTokens = u.InputTokens + u.OutputTokens
+	}
+	return u, nil
+}
+
+// openAICachedTokens returns the cached-prompt subset reported by
+// either the Responses-API (input_tokens_details.cached_tokens) or
+// the Chat-Completions API (prompt_tokens_details.cached_tokens).
+// Responses-API takes precedence when both are populated.
+func openAICachedTokens(resp openAIResponse) int64 {
+	// Responses-API details are authoritative when present: an explicit
+	// cached_tokens of 0 must be honored, not treated as missing and
+	// overridden by the Chat-Completions field (which would overstate cache).
+	if resp.Usage.InputTokensDetails != nil && resp.Usage.InputTokensDetails.CachedTokens != nil {
+		return derefInt64(resp.Usage.InputTokensDetails.CachedTokens)
+	}
+	if resp.Usage.PromptTokensDetails != nil {
+		return derefInt64(resp.Usage.PromptTokensDetails.CachedTokens)
+	}
+	return 0
+}
+
+// ExtractPrompt returns the user-visible prompt text from an OpenAI request.
+// Handles chat.completions (messages[].content), legacy completions (prompt
+// string), and the Responses API (input as string or content-part array).
+// Returns "" when nothing extractable is found.
+func (OpenAIParser) ExtractPrompt(body []byte) string {
+	var req openAIRequest
+	if err := json.Unmarshal(body, &req); err != nil {
+		return ""
+	}
+	if len(req.Messages) > 0 {
+		return joinMessages(req.Messages)
+	}
+	if len(req.Input) > 0 {
+		return extractResponsesInput(req.Input)
+	}
+	if len(req.Prompt) > 0 {
+		return decodeStringOrJoin(req.Prompt)
+	}
+	return ""
+}
+
+// extractResponsesInput handles the Responses API `input` field. It is one
+// of three shapes: a plain string, an array of message items
+// ({role, content: string | [parts]}) as sent by Codex and the Responses
+// SDK, or a flat array of content parts ({type, text/input_text}). Message
+// items are flattened to "role: text" lines; items without extractable text
+// (reasoning blocks, tool calls) are skipped.
+func extractResponsesInput(raw json.RawMessage) string {
+	if s, ok := tryDecodeString(raw); ok {
+		return s
+	}
+	var items []struct {
+		Role      string          `json:"role"`
+		Content   json.RawMessage `json:"content"`
+		Text      string          `json:"text"`
+		InputText string          `json:"input_text"`
+	}
+	if err := json.Unmarshal(raw, &items); err != nil {
+		return extractContentParts(raw)
+	}
+	var b strings.Builder
+	for _, it := range items {
+		var text string
+		switch {
+		case len(it.Content) > 0:
+			text = decodeStringOrJoin(it.Content)
+		case it.Text != "":
+			text = it.Text
+		case it.InputText != "":
+			text = it.InputText
+		}
+		if text == "" {
+			continue
+		}
+		if b.Len() > 0 {
+			b.WriteByte('\n')
+		}
+		if it.Role != "" {
+			b.WriteString(it.Role)
+			b.WriteString(": ")
+		}
+		b.WriteString(text)
+	}
+	return b.String()
+}
+
+// ExtractSessionID reads the OpenAI session marker. Codex (the Responses
+// API client) stamps client_metadata.session_id on every request body;
+// plain chat.completions traffic carries no session id and yields "".
+func (OpenAIParser) ExtractSessionID(body []byte) string {
+	var req struct {
+		ClientMetadata struct {
+			SessionID string `json:"session_id"`
+		} `json:"client_metadata"`
+	}
+	if err := json.Unmarshal(body, &req); err != nil {
+		return ""
+	}
+	return req.ClientMetadata.SessionID
+}
+
+type openAIChatChoice struct {
+	Message struct {
+		Role    string          `json:"role"`
+		Content json.RawMessage `json:"content"`
+	} `json:"message"`
+	Text string `json:"text"`
+}
+
+type openAIChatResponse struct {
+	Choices []openAIChatChoice `json:"choices"`
+	// Responses API: output[].content[].text
+	Output []struct {
+		Type    string          `json:"type"`
+		Content json.RawMessage `json:"content"`
+		Text    string          `json:"text"`
+	} `json:"output"`
+	OutputText string `json:"output_text"`
+}
+
+// ExtractCompletion returns the assistant text from a non-streaming OpenAI
+// response. Handles chat.completions (choices[].message.content), legacy
+// completions (choices[].text), and Responses API (output[].content[].text
+// or the convenience output_text field).
+func (OpenAIParser) ExtractCompletion(status int, contentType string, body []byte) string {
+	if status != 200 || isEventStream(contentType) || !isJSON(contentType) {
+		return ""
+	}
+	var resp openAIChatResponse
+	if err := json.Unmarshal(body, &resp); err != nil {
+		return ""
+	}
+	if resp.OutputText != "" {
+		return resp.OutputText
+	}
+	for _, c := range resp.Choices {
+		if len(c.Message.Content) > 0 {
+			if s := decodeStringOrJoin(c.Message.Content); s != "" {
+				return s
+			}
+		}
+		if c.Text != "" {
+			return c.Text
+		}
+	}
+	for _, o := range resp.Output {
+		if o.Text != "" {
+			return o.Text
+		}
+		if len(o.Content) > 0 {
+			if s := extractContentParts(o.Content); s != "" {
+				return s
+			}
+		}
+	}
+	return ""
+}
+
+// joinMessages flattens a chat.completions messages array into a single
+// "role: content" string per message, separated by newlines. Roles surface
+// system/user/assistant context which is useful for log review.
+func joinMessages(msgs []openAIMessage) string {
+	var b strings.Builder
+	for i, m := range msgs {
+		if i > 0 {
+			b.WriteByte('\n')
+		}
+		if m.Role != "" {
+			b.WriteString(m.Role)
+			b.WriteString(": ")
+		}
+		b.WriteString(decodeStringOrJoin(m.Content))
+	}
+	return b.String()
+}
+
+// extractContentParts handles the Responses-API content shape, which is
+// either a single string or an array of {type, text} parts. text and
+// input_text both carry user-facing content.
+func extractContentParts(raw json.RawMessage) string {
+	if s, ok := tryDecodeString(raw); ok {
+		return s
+	}
+	var parts []struct {
+		Type      string `json:"type"`
+		Text      string `json:"text"`
+		InputText string `json:"input_text"`
+	}
+	if err := json.Unmarshal(raw, &parts); err != nil {
+		// Last-ditch: array of strings.
+		var arr []string
+		if json.Unmarshal(raw, &arr) == nil {
+			return strings.Join(arr, "\n")
+		}
+		return ""
+	}
+	var b strings.Builder
+	for _, p := range parts {
+		var text string
+		switch {
+		case p.Text != "":
+			text = p.Text
+		case p.InputText != "":
+			text = p.InputText
+		}
+		if text == "" {
+			continue
+		}
+		if b.Len() > 0 {
+			b.WriteByte('\n')
+		}
+		b.WriteString(text)
+	}
+	return b.String()
+}
+
+// decodeStringOrJoin accepts either a JSON string or a content-parts array
+// (chat.completions multimodal) and returns a flat string. Multimodal parts
+// are separated by newlines; non-text parts are skipped.
+func decodeStringOrJoin(raw json.RawMessage) string {
+	if s, ok := tryDecodeString(raw); ok {
+		return s
+	}
+	return extractContentParts(raw)
+}
+
+func tryDecodeString(raw json.RawMessage) (string, bool) {
+	if len(raw) == 0 {
+		return "", false
+	}
+	var s string
+	if err := json.Unmarshal(raw, &s); err == nil {
+		return s, true
+	}
+	return "", false
+}
+
+// pickInt64 returns the first non-nil pointer's value. Used to prefer one
+// naming convention while transparently falling back to another.
+func pickInt64(preferred, fallback *int64) int64 {
+	if preferred != nil {
+		return *preferred
+	}
+	return derefInt64(fallback)
+}
+
+func derefInt64(v *int64) int64 {
+	if v == nil {
+		return 0
+	}
+	return *v
+}
+
+func ptrDeref(b *bool) bool {
+	if b == nil {
+		return false
+	}
+	return *b
+}
+
+func isEventStream(contentType string) bool {
+	return strings.Contains(strings.ToLower(contentType), "text/event-stream")
+}
+
+func isJSON(contentType string) bool {
+	lower := strings.ToLower(contentType)
+	return strings.Contains(lower, "application/json") || strings.Contains(lower, "+json")
+}

proxy/internal/llm/openai_test.go

@@ -0,0 +1,255 @@
+package llm
+
+import (
+	"errors"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestOpenAIDetectFromURL(t *testing.T) {
+	p := OpenAIParser{}
+
+	cases := map[string]bool{
+		"/v1/chat/completions":                  true,
+		"/v1/completions":                       true,
+		"/v1/embeddings":                        true,
+		"/v1/responses":                         true,
+		"/API/V1/Chat/Completions":              true,
+		"/upstream/v1/chat/completions?trace=1": true,
+		// Cloudflare AI Gateway puts its own /v1/{account}/{gateway}
+		// segment between the canonical /v1/ and the provider slug,
+		// so the /v1/chat/completions substring no longer appears
+		// adjacent in the path. The bare /chat/completions hint
+		// catches Cloudflare's OpenAI direct path
+		// (/v1/{account}/{gateway}/openai/chat/completions) and
+		// compat path (/v1/{account}/{gateway}/compat/chat/completions).
+		"/v1/{account}/{gateway}/openai/chat/completions": true,
+		"/v1/{account}/{gateway}/compat/chat/completions": true,
+		"/chat/completions": true,
+		"/v1/messages":      false,
+		"/healthz":          false,
+		"":                  false,
+	}
+	for path, want := range cases {
+		assert.Equal(t, want, p.DetectFromURL(path), "DetectFromURL(%q)", path)
+	}
+}
+
+func TestOpenAIParseRequest(t *testing.T) {
+	p := OpenAIParser{}
+
+	t.Run("stream true", func(t *testing.T) {
+		facts, err := p.ParseRequest([]byte(`{"model":"gpt-4o","stream":true,"stream_options":{"include_usage":true}}`))
+		require.NoError(t, err)
+		assert.Equal(t, "gpt-4o", facts.Model, "request model extracted")
+		assert.True(t, facts.Stream, "request marked as streaming")
+	})
+
+	t.Run("stream default", func(t *testing.T) {
+		facts, err := p.ParseRequest([]byte(`{"model":"gpt-4o-mini"}`))
+		require.NoError(t, err)
+		assert.Equal(t, "gpt-4o-mini", facts.Model, "request model extracted")
+		assert.False(t, facts.Stream, "missing stream flag defaults to false")
+	})
+
+	t.Run("malformed", func(t *testing.T) {
+		_, err := p.ParseRequest([]byte(`{not json}`))
+		require.Error(t, err)
+		assert.True(t, errors.Is(err, ErrMalformedRequest), "sentinel error wrapped")
+	})
+}
+
+func TestOpenAIParseResponse(t *testing.T) {
+	p := OpenAIParser{}
+
+	t.Run("happy fixture", func(t *testing.T) {
+		body, err := os.ReadFile(filepath.Join("fixtures", "openai_chat_completion.json"))
+		require.NoError(t, err, "fixture must be readable")
+
+		usage, err := p.ParseResponse(200, "application/json", body)
+		require.NoError(t, err)
+		assert.Equal(t, int64(123), usage.InputTokens, "prompt tokens become input")
+		assert.Equal(t, int64(45), usage.OutputTokens, "completion tokens become output")
+		assert.Equal(t, int64(168), usage.TotalTokens, "total_tokens carried through")
+	})
+
+	t.Run("total computed when missing", func(t *testing.T) {
+		body := []byte(`{"usage":{"prompt_tokens":10,"completion_tokens":5}}`)
+		usage, err := p.ParseResponse(200, "application/json", body)
+		require.NoError(t, err)
+		assert.Equal(t, int64(15), usage.TotalTokens, "total computed from in+out")
+	})
+
+	t.Run("streaming rejected", func(t *testing.T) {
+		_, err := p.ParseResponse(200, "text/event-stream", []byte(""))
+		require.ErrorIs(t, err, ErrStreamingUnsupported, "SSE responses must use the scanner")
+	})
+
+	t.Run("non-200", func(t *testing.T) {
+		_, err := p.ParseResponse(500, "application/json", []byte(`{"error":"x"}`))
+		require.ErrorIs(t, err, ErrNotLLMResponse, "non-200 rejected as non-LLM")
+	})
+
+	t.Run("non-json content type", func(t *testing.T) {
+		_, err := p.ParseResponse(200, "text/plain", []byte(`{}`))
+		require.ErrorIs(t, err, ErrNotLLMResponse, "text/plain treated as non-LLM")
+	})
+
+	t.Run("malformed body", func(t *testing.T) {
+		_, err := p.ParseResponse(200, "application/json", []byte(`{not json`))
+		require.ErrorIs(t, err, ErrMalformedResponse, "bad JSON yields malformed error")
+	})
+
+	// Responses-API fixture: /v1/responses returns input_tokens/output_tokens
+	// (Anthropic-style) instead of prompt_tokens/completion_tokens. The parser
+	// must accept both.
+	t.Run("responses api fixture", func(t *testing.T) {
+		body, err := os.ReadFile(filepath.Join("fixtures", "openai_responses.json"))
+		require.NoError(t, err, "fixture must be readable")
+
+		usage, err := p.ParseResponse(200, "application/json", body)
+		require.NoError(t, err)
+		assert.Equal(t, int64(15), usage.InputTokens, "input_tokens should map directly")
+		assert.Equal(t, int64(414), usage.OutputTokens, "output_tokens should map directly")
+		assert.Equal(t, int64(429), usage.TotalTokens, "total_tokens carried through")
+	})
+
+	t.Run("responses api naming preferred over chat-completions when both present", func(t *testing.T) {
+		body := []byte(`{"usage":{"prompt_tokens":1,"completion_tokens":2,"input_tokens":15,"output_tokens":414,"total_tokens":429}}`)
+		usage, err := p.ParseResponse(200, "application/json", body)
+		require.NoError(t, err)
+		assert.Equal(t, int64(15), usage.InputTokens, "responses-api names take precedence")
+		assert.Equal(t, int64(414), usage.OutputTokens, "responses-api names take precedence")
+	})
+
+	t.Run("chat-completions naming still works alone", func(t *testing.T) {
+		body := []byte(`{"usage":{"prompt_tokens":15,"completion_tokens":414,"total_tokens":429}}`)
+		usage, err := p.ParseResponse(200, "application/json", body)
+		require.NoError(t, err)
+		assert.Equal(t, int64(15), usage.InputTokens, "prompt_tokens fallback")
+		assert.Equal(t, int64(414), usage.OutputTokens, "completion_tokens fallback")
+	})
+
+	// Cached-prompt accounting. cached_tokens is a SUBSET of
+	// prompt_tokens — input_tokens carries the full prompt count and
+	// the cached subset is reported separately so the cost meter can
+	// apply the discount rate to that portion.
+	t.Run("chat-completions cached_tokens subset surfaces", func(t *testing.T) {
+		body := []byte(`{"usage":{"prompt_tokens":1024,"completion_tokens":200,"total_tokens":1224,"prompt_tokens_details":{"cached_tokens":768}}}`)
+		usage, err := p.ParseResponse(200, "application/json", body)
+		require.NoError(t, err)
+		assert.Equal(t, int64(1024), usage.InputTokens, "input remains the full prompt count — cached is a subset, not a separate bucket")
+		assert.Equal(t, int64(768), usage.CachedInputTokens, "cached_tokens must surface so cost meter can discount the cached subset")
+		assert.Zero(t, usage.CacheCreationTokens, "OpenAI has no cache_creation analogue")
+	})
+
+	t.Run("responses-api input_tokens_details.cached_tokens surfaces", func(t *testing.T) {
+		body := []byte(`{"usage":{"input_tokens":2048,"output_tokens":100,"total_tokens":2148,"input_tokens_details":{"cached_tokens":1500}}}`)
+		usage, err := p.ParseResponse(200, "application/json", body)
+		require.NoError(t, err)
+		assert.Equal(t, int64(2048), usage.InputTokens)
+		assert.Equal(t, int64(1500), usage.CachedInputTokens, "Responses-API input_tokens_details.cached_tokens path must surface too")
+	})
+
+	t.Run("responses-api cached takes precedence over chat-completions when both present", func(t *testing.T) {
+		body := []byte(`{"usage":{"prompt_tokens":1,"input_tokens":2,"output_tokens":3,"prompt_tokens_details":{"cached_tokens":50},"input_tokens_details":{"cached_tokens":99}}}`)
+		usage, err := p.ParseResponse(200, "application/json", body)
+		require.NoError(t, err)
+		assert.Equal(t, int64(99), usage.CachedInputTokens, "Responses-API field wins when both naming conventions are present")
+	})
+
+	t.Run("absent cached_tokens leaves cached counts at zero", func(t *testing.T) {
+		body := []byte(`{"usage":{"prompt_tokens":15,"completion_tokens":414,"total_tokens":429}}`)
+		usage, err := p.ParseResponse(200, "application/json", body)
+		require.NoError(t, err)
+		assert.Zero(t, usage.CachedInputTokens, "no prompt_tokens_details = no cached subset")
+	})
+}
+
+func TestOpenAIExtractPrompt_ChatCompletions(t *testing.T) {
+	body := []byte(`{"model":"gpt-4o-mini","messages":[{"role":"system","content":"be brief"},{"role":"user","content":"ping"}]}`)
+	got := OpenAIParser{}.ExtractPrompt(body)
+	require.NotEmpty(t, got, "messages array must extract")
+	require.Contains(t, got, "system: be brief", "system role and content surface")
+	require.Contains(t, got, "user: ping", "user role and content surface")
+}
+
+func TestOpenAIExtractPrompt_ResponsesAPIStringInput(t *testing.T) {
+	body := []byte(`{"model":"gpt-5.4","input":"Hello there"}`)
+	got := OpenAIParser{}.ExtractPrompt(body)
+	require.Equal(t, "Hello there", got, "string input field should pass through")
+}
+
+func TestOpenAIExtractPrompt_ResponsesAPIInputParts(t *testing.T) {
+	body := []byte(`{"model":"gpt-5.4","input":[{"type":"input_text","input_text":"first"},{"type":"input_text","input_text":"second"}]}`)
+	got := OpenAIParser{}.ExtractPrompt(body)
+	require.Contains(t, got, "first", "first content part surfaces")
+	require.Contains(t, got, "second", "second content part surfaces")
+}
+
+// TestOpenAIExtractPrompt_ResponsesAPIMessageItems guards the live Codex
+// shape: input is an array of message items whose text is nested under
+// content[].text, not flat content parts. The old code fed the outer array
+// to the content-part decoder and extracted nothing, so the stored prompt
+// was empty.
+func TestOpenAIExtractPrompt_ResponsesAPIMessageItems(t *testing.T) {
+	body := []byte(`{"model":"gpt-5.5","input":[` +
+		`{"type":"message","role":"developer","content":[{"type":"input_text","text":"system rules"}]},` +
+		`{"type":"message","role":"user","content":[{"type":"input_text","text":"hello there"}]},` +
+		`{"type":"reasoning","encrypted_content":"opaque","summary":[]},` +
+		`{"type":"message","role":"assistant","content":[{"type":"output_text","text":"prior reply"}]}` +
+		`]}`)
+	got := OpenAIParser{}.ExtractPrompt(body)
+	require.Contains(t, got, "system rules", "developer message content must surface")
+	require.Contains(t, got, "hello there", "user message content must surface")
+	require.Contains(t, got, "developer:", "role labels must prefix each message")
+	require.NotContains(t, got, "opaque", "reasoning items without text must be skipped")
+}
+
+func TestOpenAIExtractPrompt_LegacyCompletion(t *testing.T) {
+	body := []byte(`{"model":"text-davinci-003","prompt":"once upon a time"}`)
+	got := OpenAIParser{}.ExtractPrompt(body)
+	require.Equal(t, "once upon a time", got, "string prompt field should pass through")
+}
+
+func TestOpenAIExtractSessionID(t *testing.T) {
+	t.Run("codex client_metadata.session_id", func(t *testing.T) {
+		body := []byte(`{"model":"gpt-5.5","client_metadata":{"session_id":"019eeb72-ab7c-7cd2","thread_id":"t1"},"input":[]}`)
+		assert.Equal(t, "019eeb72-ab7c-7cd2", OpenAIParser{}.ExtractSessionID(body), "Codex session id must come from client_metadata.session_id")
+	})
+	t.Run("plain chat has no session", func(t *testing.T) {
+		body := []byte(`{"model":"gpt-4o","messages":[{"role":"user","content":"hi"}]}`)
+		assert.Equal(t, "", OpenAIParser{}.ExtractSessionID(body), "plain chat.completions carries no session id")
+	})
+	t.Run("non-JSON yields empty", func(t *testing.T) {
+		assert.Equal(t, "", OpenAIParser{}.ExtractSessionID([]byte("not json")), "malformed body must not error")
+	})
+}
+
+func TestOpenAIExtractCompletion_ChatCompletions(t *testing.T) {
+	body, err := os.ReadFile(filepath.Join("fixtures", "openai_chat_completion.json"))
+	require.NoError(t, err)
+	got := OpenAIParser{}.ExtractCompletion(200, "application/json", body)
+	require.NotEmpty(t, got, "fixture has assistant content")
+}
+
+func TestOpenAIExtractCompletion_ResponsesAPI(t *testing.T) {
+	body, err := os.ReadFile(filepath.Join("fixtures", "openai_responses.json"))
+	require.NoError(t, err)
+	got := OpenAIParser{}.ExtractCompletion(200, "application/json", body)
+	require.NotEmpty(t, got, "responses-api fixture has output content")
+}
+
+func TestOpenAIExtractCompletion_Streaming(t *testing.T) {
+	got := OpenAIParser{}.ExtractCompletion(200, "text/event-stream", []byte(""))
+	require.Empty(t, got, "streaming responses are skipped")
+}
+
+func TestOpenAIExtractCompletion_NonOK(t *testing.T) {
+	got := OpenAIParser{}.ExtractCompletion(500, "application/json", []byte(`{"choices":[{"message":{"content":"x"}}]}`))
+	require.Empty(t, got, "non-200 returns empty")
+}

proxy/internal/llm/parser.go

@@ -0,0 +1,112 @@
+// Package llm provides the shared LLM request and response parsing
+// library consumed by proxy middleware. It is runtime agnostic: the same
+// package is used by the native built-in executor now and will be reused
+// by the WASM adapter later.
+package llm
+
+// Provider identifies an LLM API provider.
+type Provider int
+
+const (
+	// ProviderUnknown signals that no parser matched the request.
+	ProviderUnknown Provider = 0
+	// ProviderOpenAI identifies the OpenAI API surface.
+	ProviderOpenAI Provider = 1
+	// ProviderAnthropic identifies the Anthropic Messages API surface.
+	ProviderAnthropic Provider = 2
+	// ProviderBedrock identifies the AWS Bedrock runtime surface.
+	ProviderBedrock Provider = 3
+)
+
+// RequestFacts captures the subset of the LLM request body that the
+// middleware annotates as metadata (model, streaming flag). Additional
+// fields are added as parsers grow.
+type RequestFacts struct {
+	Model  string
+	Stream bool
+}
+
+// Usage is the provider-agnostic token accounting emitted to metrics and
+// access logs. Downstream consumers map InputTokens/OutputTokens to the
+// plg.llm.* metadata allowlist entries.
+//
+// CachedInputTokens carries OpenAI's prompt_tokens_details.cached_tokens
+// (a SUBSET of InputTokens) when the response is from OpenAI, or
+// Anthropic's cache_read_input_tokens (ADDITIVE to InputTokens) when from
+// Anthropic. The cost meter switches formula on KeyLLMProvider so the
+// two shapes are billed correctly without double-counting.
+//
+// CacheCreationTokens carries Anthropic's cache_creation_input_tokens
+// (ADDITIVE; not present in the OpenAI shape).
+type Usage struct {
+	InputTokens         int64
+	OutputTokens        int64
+	TotalTokens         int64
+	CachedInputTokens   int64
+	CacheCreationTokens int64
+}
+
+// Parser is the per-provider interface implemented in this package. The
+// dispatcher selects a parser by calling DetectFromURL against the incoming
+// request path; ties break by registration order (see Parsers).
+type Parser interface {
+	Provider() Provider
+	ProviderName() string
+	DetectFromURL(path string) bool
+	ParseRequest(body []byte) (RequestFacts, error)
+	ParseResponse(status int, contentType string, body []byte) (Usage, error)
+	// ExtractPrompt returns the user-facing prompt text from a request body.
+	// Different endpoint shapes (chat.completions, responses, messages) are
+	// handled by the per-provider implementation. Returns "" when no prompt
+	// can be extracted; never returns an error — extraction is best-effort
+	// because callers use the result for observability, not authorization.
+	ExtractPrompt(body []byte) string
+	// ExtractCompletion returns the assistant-facing completion text from a
+	// non-streaming response body. status and contentType match the
+	// ParseResponse arguments so implementations can fast-fail uniformly.
+	ExtractCompletion(status int, contentType string, body []byte) string
+	// ExtractSessionID returns a stable identifier that groups requests of
+	// the same conversation / coding session, read from the per-provider
+	// location clients populate (e.g. OpenAI Codex's client_metadata.session_id,
+	// Claude Code's metadata.user_id). Returns "" when the body carries no
+	// recognised session marker; extraction is best-effort and never errors.
+	ExtractSessionID(body []byte) string
+}
+
+// Parsers returns the built-in parser set in a stable order. The order is
+// deterministic so that DetectFromURL ties produce consistent routing.
+func Parsers() []Parser {
+	return []Parser{
+		OpenAIParser{},
+		AnthropicParser{},
+		BedrockParser{},
+	}
+}
+
+// DetectParser returns the first parser whose DetectFromURL matches the given
+// request path. ok=false means no parser claimed the path.
+func DetectParser(path string) (Parser, bool) {
+	for _, p := range Parsers() {
+		if p.DetectFromURL(path) {
+			return p, true
+		}
+	}
+	return nil, false
+}
+
+// ParserByName returns the parser whose ProviderName matches id. Used by
+// callers that already know which provider surface a request will hit
+// (e.g. the agent-network middleware chain configured per synthesised
+// service) so they can skip URL sniffing. ok=false when no parser is
+// registered under that name.
+func ParserByName(id string) (Parser, bool) {
+	if id == "" {
+		return nil, false
+	}
+	for _, p := range Parsers() {
+		if p.ProviderName() == id {
+			return p, true
+		}
+	}
+	return nil, false
+}

proxy/internal/llm/parser_test.go

@@ -0,0 +1,54 @@
+package llm
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestParsers_ProviderNames(t *testing.T) {
+	parsers := Parsers()
+	require.Len(t, parsers, 3, "three built-in parsers expected")
+
+	names := make([]string, 0, len(parsers))
+	for _, p := range parsers {
+		names = append(names, p.ProviderName())
+	}
+	assert.Contains(t, names, "openai", "OpenAI parser should be registered")
+	assert.Contains(t, names, "anthropic", "Anthropic parser should be registered")
+	assert.Contains(t, names, "bedrock", "Bedrock parser should be registered")
+}
+
+func TestDetectParser(t *testing.T) {
+	cases := []struct {
+		name         string
+		path         string
+		expectedName string
+		expectOK     bool
+	}{
+		{"openai chat", "/v1/chat/completions", "openai", true},
+		{"openai prefixed", "/api/v1/chat/completions", "openai", true},
+		{"openai responses", "/v1/responses", "openai", true},
+		{"anthropic messages", "/v1/messages", "anthropic", true},
+		{"anthropic prefixed", "/proxy/v1/messages?query", "anthropic", true},
+		{"unknown path", "/healthz", "", false},
+		{"empty path", "", "", false},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			p, ok := DetectParser(tc.path)
+			require.Equal(t, tc.expectOK, ok, "detection success mismatch for %q", tc.path)
+			if ok {
+				assert.Equal(t, tc.expectedName, p.ProviderName(), "provider name mismatch")
+			}
+		})
+	}
+}
+
+func TestProviderValues(t *testing.T) {
+	assert.Equal(t, Provider(0), ProviderUnknown, "unknown provider is the zero value")
+	assert.Equal(t, ProviderOpenAI, OpenAIParser{}.Provider(), "OpenAI parser reports its provider enum")
+	assert.Equal(t, ProviderAnthropic, AnthropicParser{}.Provider(), "Anthropic parser reports its provider enum")
+}

proxy/internal/llm/pricing/defaults_coverage_test.go

@@ -0,0 +1,65 @@
+package pricing
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// TestDefaultTable_FirstPartyModelCoverage guards the embedded defaults against
+// silent drift/gaps: every metered first-party model the management catalog
+// enumerates must resolve to a price, and a few rates that previously drifted
+// are pinned to their LiteLLM-validated values. Keep this list in step with the
+// catalog (management/server/agentnetwork/catalog) when adding models.
+func TestDefaultTable_FirstPartyModelCoverage(t *testing.T) {
+	tbl := DefaultTable()
+	require.NotNil(t, tbl, "embedded default pricing table must load")
+
+	mustPrice := map[string][]string{
+		// openai parser covers openai_api, azure_openai_api, and mistral_api.
+		"openai": {
+			"gpt-5.5", "gpt-5.5-pro", "gpt-5.4", "gpt-5.4-mini", "gpt-5.4-nano",
+			"gpt-5.3-codex", "gpt-5.3-chat-latest", "o4-mini",
+			"gpt-4.1", "gpt-4.1-mini", "gpt-4.1-nano", "gpt-4o", "gpt-4o-mini",
+			"gpt-4-turbo", "gpt-3.5-turbo", "gpt-35-turbo",
+			"text-embedding-3-large", "text-embedding-3-small",
+			"mistral-large-latest", "mistral-medium-3-5", "codestral-2508",
+			"ministral-8b-latest", "mistral-embed",
+		},
+		"anthropic": {
+			"claude-fable-5", "claude-opus-5", "claude-opus-4-8", "claude-opus-4-7", "claude-opus-4-6",
+			"claude-opus-4-1", "claude-sonnet-4-6", "claude-sonnet-4-5", "claude-haiku-4-5",
+		},
+		// bedrock keys are the normalized ids the request parser emits.
+		"bedrock": {
+			"anthropic.claude-opus-5", "anthropic.claude-opus-4-8", "anthropic.claude-opus-4-7", "anthropic.claude-opus-4-6",
+			"anthropic.claude-opus-4-1", "anthropic.claude-sonnet-4-6", "anthropic.claude-sonnet-4-5",
+			"anthropic.claude-haiku-4-5", "meta.llama3-3-70b-instruct",
+			"amazon.nova-pro", "amazon.nova-lite", "amazon.nova-micro", "amazon.nova-2-lite",
+		},
+	}
+	for provider, models := range mustPrice {
+		for _, m := range models {
+			_, ok := tbl.Cost(provider, m, 1000, 1000, 0, 0)
+			assert.True(t, ok, "%s/%s must be priced in the embedded defaults", provider, m)
+		}
+	}
+
+	// Pin per-direction rates independently (input-only then output-only) so a
+	// swap or skew of input<->output that preserves the combined total is still
+	// caught — these are rates that previously drifted or are easy to mis-enter.
+	in, ok := tbl.Cost("openai", "gpt-5.4", 1000, 0, 0, 0)
+	require.True(t, ok)
+	assert.InDelta(t, 0.0025, in, 1e-9, "gpt-5.4 input = 0.0025 per 1k")
+	out, ok := tbl.Cost("openai", "gpt-5.4", 0, 1000, 0, 0)
+	require.True(t, ok)
+	assert.InDelta(t, 0.015, out, 1e-9, "gpt-5.4 output = 0.015 per 1k")
+
+	in, ok = tbl.Cost("bedrock", "anthropic.claude-sonnet-4-5", 1000, 0, 0, 0)
+	require.True(t, ok)
+	assert.InDelta(t, 0.003, in, 1e-9, "bedrock sonnet-4-5 input = 0.003 per 1k")
+	out, ok = tbl.Cost("bedrock", "anthropic.claude-sonnet-4-5", 0, 1000, 0, 0)
+	require.True(t, ok)
+	assert.InDelta(t, 0.015, out, 1e-9, "bedrock sonnet-4-5 output = 0.015 per 1k")
+}

proxy/internal/llm/pricing/defaults_pricing.yaml

@@ -0,0 +1,298 @@
+# Embedded default pricing for llm_observability. Compiled into the proxy
+# binary via go:embed in pricing.go; cost annotation works out of the box
+# without any operator action.
+#
+# Operators override entries by dropping a pricing.yaml into --plugin-data-dir
+# (or whichever basename is given via params.pricing_path). The override file
+# only needs entries the operator wants to change; missing entries fall
+# through to these defaults.
+#
+# Values are USD per 1_000 tokens. Public list prices drift; ship a fresh
+# binary or override individual entries via the override file as needed.
+#
+# Optional cache fields:
+#   cached_input_per_1k    OpenAI: rate for prompt_tokens_details.cached_tokens
+#                          (a SUBSET of prompt_tokens). Typically 0.5x input.
+#                          Absent → cached portion bills at input_per_1k.
+#   cache_read_per_1k      Anthropic: rate for cache_read_input_tokens
+#                          (ADDITIVE to input_tokens). Typically 0.1x input.
+#                          Absent → cache reads bill at input_per_1k.
+#   cache_creation_per_1k  Anthropic: rate for cache_creation_input_tokens
+#                          (ADDITIVE to input_tokens). Typically 1.25x input.
+#                          Absent → cache writes bill at input_per_1k.
+
+openai:
+  # OpenAI + OpenAI-compatible providers (openai_api, azure_openai_api,
+  # mistral_api, and the openai-parser gateways) all emit llm.provider="openai",
+  # so their models are priced here. Kept in sync with the management catalog;
+  # rates cross-checked against LiteLLM model_prices_and_context_window.json.
+
+  # GPT-5.x family — cache reads 10% of input (0.1x).
+  gpt-5.5:
+    input_per_1k: 0.005
+    output_per_1k: 0.03
+    cached_input_per_1k: 0.0005
+  gpt-5.5-pro:
+    input_per_1k: 0.03
+    output_per_1k: 0.18
+    cached_input_per_1k: 0.003
+  gpt-5.4:
+    input_per_1k: 0.0025
+    output_per_1k: 0.015
+    cached_input_per_1k: 0.00025
+  gpt-5.4-pro:
+    input_per_1k: 0.03
+    output_per_1k: 0.18
+    cached_input_per_1k: 0.003
+  gpt-5.4-mini:
+    input_per_1k: 0.00075
+    output_per_1k: 0.0045
+    cached_input_per_1k: 0.000075
+  gpt-5.4-nano:
+    input_per_1k: 0.0002
+    output_per_1k: 0.00125
+    cached_input_per_1k: 0.00002
+  gpt-5.3-codex:
+    input_per_1k: 0.00175
+    output_per_1k: 0.014
+    cached_input_per_1k: 0.000175
+  gpt-5.3-chat-latest:
+    input_per_1k: 0.00175
+    output_per_1k: 0.014
+    cached_input_per_1k: 0.000175
+  # GPT-5 (2025) family — kept for gateway requests using the unsuffixed ids.
+  gpt-5:
+    input_per_1k: 0.00125
+    output_per_1k: 0.01
+    cached_input_per_1k: 0.000125
+  gpt-5-mini:
+    input_per_1k: 0.00025
+    output_per_1k: 0.002
+    cached_input_per_1k: 0.000025
+  gpt-5-nano:
+    input_per_1k: 0.00005
+    output_per_1k: 0.0004
+    cached_input_per_1k: 0.000005
+  o4-mini:
+    input_per_1k: 0.0011
+    output_per_1k: 0.0044
+    cached_input_per_1k: 0.000275
+  # GPT-4.1 family — cache reads 25% of input.
+  gpt-4.1:
+    input_per_1k: 0.002
+    output_per_1k: 0.008
+    cached_input_per_1k: 0.0005
+  gpt-4.1-mini:
+    input_per_1k: 0.0004
+    output_per_1k: 0.0016
+    cached_input_per_1k: 0.0001
+  gpt-4.1-nano:
+    input_per_1k: 0.0001
+    output_per_1k: 0.0004
+    cached_input_per_1k: 0.000025
+  # GPT-4o family — cache reads 50% of input (0.5x).
+  gpt-4o:
+    input_per_1k: 0.0025
+    output_per_1k: 0.01
+    cached_input_per_1k: 0.00125
+  gpt-4o-mini:
+    input_per_1k: 0.00015
+    output_per_1k: 0.0006
+    cached_input_per_1k: 0.000075
+  # Older GPT — no prompt caching.
+  gpt-4-turbo:
+    input_per_1k: 0.01
+    output_per_1k: 0.03
+  gpt-3.5-turbo:
+    input_per_1k: 0.0005
+    output_per_1k: 0.0015
+  gpt-35-turbo: # Azure deployment alias of gpt-3.5-turbo
+    input_per_1k: 0.0005
+    output_per_1k: 0.0015
+  # Embeddings — no caching, no output tokens.
+  text-embedding-3-large:
+    input_per_1k: 0.00013
+    output_per_1k: 0
+  text-embedding-3-small:
+    input_per_1k: 0.00002
+    output_per_1k: 0
+
+  # Mistral (mistral_api) — routed via the openai parser; no prompt caching.
+  mistral-large-latest:
+    input_per_1k: 0.0005
+    output_per_1k: 0.0015
+  mistral-medium-latest:
+    input_per_1k: 0.0004
+    output_per_1k: 0.002
+  mistral-medium-3-5:
+    input_per_1k: 0.0015
+    output_per_1k: 0.0075
+  mistral-small-latest:
+    input_per_1k: 0.00006
+    output_per_1k: 0.00018
+  magistral-medium-latest:
+    input_per_1k: 0.002
+    output_per_1k: 0.005
+  magistral-small-latest:
+    input_per_1k: 0.0005
+    output_per_1k: 0.0015
+  devstral-medium-latest:
+    input_per_1k: 0.0004
+    output_per_1k: 0.002
+  devstral-small-latest:
+    input_per_1k: 0.0001
+    output_per_1k: 0.0003
+  codestral-2508:
+    input_per_1k: 0.0003
+    output_per_1k: 0.0009
+  codestral-latest:
+    input_per_1k: 0.001
+    output_per_1k: 0.003
+  ministral-3-14b-2512:
+    input_per_1k: 0.0002
+    output_per_1k: 0.0002
+  ministral-8b-latest:
+    input_per_1k: 0.00015
+    output_per_1k: 0.00015
+  ministral-3-3b-2512:
+    input_per_1k: 0.0001
+    output_per_1k: 0.0001
+  mistral-embed:
+    input_per_1k: 0.0001
+    output_per_1k: 0
+
+  # Kimi / Moonshot AI (kimi_api) — OpenAI-compatible /v1 endpoint. Moonshot
+  # reports cache hits OpenAI-style when present; cached input is 10% of
+  # input ($0.30 vs $3.00 per MTok). kimi-k3 is the only model the platform
+  # serves newer accounts (K2-era ids and kimi-latest 404), matching the
+  # management catalog.
+  kimi-k3:
+    input_per_1k: 0.003
+    output_per_1k: 0.015
+    cached_input_per_1k: 0.0003
+
+anthropic:
+  # Claude 4.x family — cache reads ≈10% of input, cache writes ≈125% of input.
+  # Pricing source: Anthropic's current published rates per million tokens,
+  # divided by 1000 for the per-1k figures stored here.
+  claude-fable-5:
+    input_per_1k: 0.010
+    output_per_1k: 0.050
+    cache_read_per_1k: 0.001
+    cache_creation_per_1k: 0.0125
+  claude-opus-5:
+    input_per_1k: 0.005
+    output_per_1k: 0.025
+    cache_read_per_1k: 0.0005
+    cache_creation_per_1k: 0.00625
+  claude-opus-4-8:
+    input_per_1k: 0.005
+    output_per_1k: 0.025
+    cache_read_per_1k: 0.0005
+    cache_creation_per_1k: 0.00625
+  claude-opus-4-7:
+    input_per_1k: 0.005
+    output_per_1k: 0.025
+    cache_read_per_1k: 0.0005
+    cache_creation_per_1k: 0.00625
+  claude-opus-4-6:
+    input_per_1k: 0.005
+    output_per_1k: 0.025
+    cache_read_per_1k: 0.0005
+    cache_creation_per_1k: 0.00625
+  claude-opus-4-1:
+    input_per_1k: 0.015
+    output_per_1k: 0.075
+    cache_read_per_1k: 0.0015
+    cache_creation_per_1k: 0.01875
+  claude-sonnet-4-6:
+    input_per_1k: 0.003
+    output_per_1k: 0.015
+    cache_read_per_1k: 0.0003
+    cache_creation_per_1k: 0.00375
+  claude-sonnet-4-5:
+    input_per_1k: 0.003
+    output_per_1k: 0.015
+    cache_read_per_1k: 0.0003
+    cache_creation_per_1k: 0.00375
+  claude-haiku-4-5:
+    input_per_1k: 0.001
+    output_per_1k: 0.005
+    cache_read_per_1k: 0.0001
+    cache_creation_per_1k: 0.00125
+
+  # Kimi / Moonshot AI (kimi_api) via the Anthropic-compatible endpoint
+  # (/anthropic/v1/messages — the official Claude Code setup). Same rates
+  # as the OpenAI-shape entry above. "kimi-k3[1m]" is the model id some
+  # Claude Code guides set for the 1M-context alias; priced identically so
+  # cost metering doesn't silently skip those requests.
+  kimi-k3:
+    input_per_1k: 0.003
+    output_per_1k: 0.015
+    cache_read_per_1k: 0.0003
+  "kimi-k3[1m]":
+    input_per_1k: 0.003
+    output_per_1k: 0.015
+    cache_read_per_1k: 0.0003
+
+bedrock:
+  # AWS Bedrock model ids, normalised by the request parser (cross-region
+  # inference-profile prefix + version/throughput suffix stripped), e.g.
+  # eu.anthropic.claude-sonnet-4-5-20250929-v1:0 -> anthropic.claude-sonnet-4-5.
+  # Anthropic-on-Bedrock keeps the additive cache buckets (read ≈0.1x input,
+  # write ≈1.25x input); Nova / Llama report no cache, so cost is input+output.
+  anthropic.claude-opus-5:
+    input_per_1k: 0.005
+    output_per_1k: 0.025
+    cache_read_per_1k: 0.0005
+    cache_creation_per_1k: 0.00625
+  anthropic.claude-opus-4-8:
+    input_per_1k: 0.005
+    output_per_1k: 0.025
+    cache_read_per_1k: 0.0005
+    cache_creation_per_1k: 0.00625
+  anthropic.claude-opus-4-7:
+    input_per_1k: 0.005
+    output_per_1k: 0.025
+    cache_read_per_1k: 0.0005
+    cache_creation_per_1k: 0.00625
+  anthropic.claude-opus-4-6:
+    input_per_1k: 0.005
+    output_per_1k: 0.025
+    cache_read_per_1k: 0.0005
+    cache_creation_per_1k: 0.00625
+  anthropic.claude-opus-4-1:
+    input_per_1k: 0.015
+    output_per_1k: 0.075
+    cache_read_per_1k: 0.0015
+    cache_creation_per_1k: 0.01875
+  anthropic.claude-sonnet-4-6:
+    input_per_1k: 0.003
+    output_per_1k: 0.015
+    cache_read_per_1k: 0.0003
+    cache_creation_per_1k: 0.00375
+  anthropic.claude-sonnet-4-5:
+    input_per_1k: 0.003
+    output_per_1k: 0.015
+    cache_read_per_1k: 0.0003
+    cache_creation_per_1k: 0.00375
+  anthropic.claude-haiku-4-5:
+    input_per_1k: 0.001
+    output_per_1k: 0.005
+    cache_read_per_1k: 0.0001
+    cache_creation_per_1k: 0.00125
+  meta.llama3-3-70b-instruct:
+    input_per_1k: 0.00072
+    output_per_1k: 0.00072
+  amazon.nova-2-lite:
+    input_per_1k: 0.0003
+    output_per_1k: 0.0025
+  amazon.nova-pro:
+    input_per_1k: 0.0008
+    output_per_1k: 0.0032
+  amazon.nova-lite:
+    input_per_1k: 0.00006
+    output_per_1k: 0.00024
+  amazon.nova-micro:
+    input_per_1k: 0.000035
+    output_per_1k: 0.00014

proxy/internal/llm/pricing/pricing.go

@@ -0,0 +1,489 @@
+// Package pricing implements the embedded-default + override pricing table
+// shared by middleware that converts LLM token usage into a USD cost
+// estimate. The table is hot-reloadable from a basename under the proxy
+// data directory; missing override files keep the embedded defaults so
+// cost annotation works without operator action.
+package pricing
+
+import (
+	"bytes"
+	"context"
+	_ "embed"
+	"errors"
+	"fmt"
+	"io"
+	"io/fs"
+	"math"
+	"path/filepath"
+	"regexp"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/metric"
+	"gopkg.in/yaml.v3"
+)
+
+//go:embed defaults_pricing.yaml
+var defaultPricingYAML []byte
+
+var (
+	defaultTableOnce sync.Once
+	defaultTablePtr  *Table
+)
+
+// DefaultTable returns the pricing table embedded in the binary. The result
+// is parsed once and shared; callers must not mutate the returned value.
+// Cost annotation works without any operator action because every loader
+// starts with this table.
+func DefaultTable() *Table {
+	defaultTableOnce.Do(func() {
+		t, err := parsePricingBytes(defaultPricingYAML)
+		if err != nil {
+			panic(fmt.Sprintf("llmobs: embedded default pricing failed to parse: %v", err))
+		}
+		defaultTablePtr = t
+	})
+	return defaultTablePtr
+}
+
+// mergeOver returns a new Table containing every entry from base, with any
+// matching entry from overlay replacing the base value. Either argument may
+// be nil. Result is a fresh allocation so callers can mutate / Store safely.
+func mergeOver(base, overlay *Table) *Table {
+	if overlay == nil || len(overlay.entries) == 0 {
+		return base
+	}
+	if base == nil || len(base.entries) == 0 {
+		return overlay
+	}
+	out := make(map[string]map[string]Entry, len(base.entries))
+	for provider, models := range base.entries {
+		inner := make(map[string]Entry, len(models))
+		for model, e := range models {
+			inner[model] = e
+		}
+		out[provider] = inner
+	}
+	for provider, models := range overlay.entries {
+		inner, ok := out[provider]
+		if !ok {
+			inner = make(map[string]Entry, len(models))
+			out[provider] = inner
+		}
+		for model, e := range models {
+			inner[model] = e
+		}
+	}
+	return &Table{entries: out}
+}
+
+// Entry is a single model's input and output pricing, expressed in USD per
+// 1000 tokens.
+//
+// CachedInputPer1K applies to OpenAI's cached prompt tokens, which are a
+// subset of input_tokens — when set, the cached portion is billed at this
+// rate and the non-cached remainder at InputPer1K. Zero means "no discount
+// configured", and cached tokens are billed at InputPer1K (matches current
+// behaviour where cached counts weren't extracted at all).
+//
+// CacheReadPer1K and CacheCreationPer1K apply to Anthropic's two prompt-
+// cache fields, which are additive to input_tokens: cache_read is the
+// cheaper read-from-cache rate, cache_creation is the more expensive
+// write-to-cache rate. Zero means "no rate configured" and the
+// corresponding token bucket is billed at InputPer1K. This is more
+// accurate than today's behaviour, where Anthropic's cache tokens are
+// ignored and not charged at all.
+type Entry struct {
+	InputPer1K         float64
+	OutputPer1K        float64
+	CachedInputPer1K   float64
+	CacheReadPer1K     float64
+	CacheCreationPer1K float64
+}
+
+// Table is a provider-to-model pricing lookup. Instances are immutable once
+// built and are swapped atomically by Loader.
+type Table struct {
+	entries map[string]map[string]Entry
+}
+
+// Cost returns the estimated USD cost for the given token counts. ok is
+// false when the provider or model is not present in the table; the caller
+// can still emit token metrics with a model=unknown label.
+//
+// Provider-shape semantics for cached / cache-creation counts:
+//
+//   - OpenAI: cachedInput is a SUBSET of inTokens. The cached portion is
+//     billed at CachedInputPer1K (or InputPer1K when no override), and the
+//     non-cached remainder of inTokens at InputPer1K. cacheCreation is
+//     ignored (OpenAI has no analogue).
+//   - Anthropic: cachedInput (cache_read) and cacheCreation are ADDITIVE to
+//     inTokens. The three buckets are billed at CacheReadPer1K,
+//     CacheCreationPer1K, and InputPer1K respectively, each falling back
+//     to InputPer1K when the corresponding rate is zero.
+//   - Other providers: cached and cacheCreation are ignored; cost is
+//     inTokens*InputPer1K + outTokens*OutputPer1K.
+func (t *Table) Cost(provider, model string, inTokens, outTokens, cachedInput, cacheCreation int64) (float64, bool) {
+	c, ok := t.Costs(provider, model, inTokens, outTokens, cachedInput, cacheCreation)
+	return c.TotalUSD, ok
+}
+
+// Costs is a per-request cost split. The four per-bucket fields are the base
+// of the breakdown — one per token bucket the provider bills separately — and
+// the two aggregates are derived from them:
+//
+//	TotalUSD = InputUSD + CachedInputUSD + CacheCreationUSD + OutputUSD
+//	CacheUSD = CachedInputUSD + CacheCreationUSD
+//
+// InputUSD is always the cost of the *non-cached* input bucket, for both
+// provider shapes: on OpenAI the cached subset is carved out of inTokens and
+// billed as CachedInputUSD, so the two never double-count. Buckets a provider
+// doesn't bill are zero, which keeps the identities above true everywhere.
+type Costs struct {
+	InputUSD         float64
+	CachedInputUSD   float64
+	CacheCreationUSD float64
+	OutputUSD        float64
+	TotalUSD         float64
+	CacheUSD         float64
+}
+
+// newCosts assembles a split from its per-bucket parts, deriving the two
+// aggregates so TotalUSD and CacheUSD can never drift from the breakdown.
+func newCosts(input, cachedInput, cacheCreation, output float64) Costs {
+	return Costs{
+		InputUSD:         input,
+		CachedInputUSD:   cachedInput,
+		CacheCreationUSD: cacheCreation,
+		OutputUSD:        output,
+		TotalUSD:         input + cachedInput + cacheCreation + output,
+		CacheUSD:         cachedInput + cacheCreation,
+	}
+}
+
+// Costs returns the estimated USD cost split for the given token counts, with
+// the same semantics as Cost.
+func (t *Table) Costs(provider, model string, inTokens, outTokens, cachedInput, cacheCreation int64) (Costs, bool) {
+	// Clamp negatives to zero before any pricing math so a malformed
+	// upstream count can never produce a negative cost.
+	if inTokens < 0 {
+		inTokens = 0
+	}
+	if outTokens < 0 {
+		outTokens = 0
+	}
+	if cachedInput < 0 {
+		cachedInput = 0
+	}
+	if cacheCreation < 0 {
+		cacheCreation = 0
+	}
+	if t == nil {
+		return Costs{}, false
+	}
+	byModel, ok := t.entries[provider]
+	if !ok {
+		return Costs{}, false
+	}
+	entry, ok := byModel[model]
+	if !ok {
+		return Costs{}, false
+	}
+	output := (float64(outTokens) / 1000.0) * entry.OutputPer1K
+	switch provider {
+	case "openai":
+		// cachedInput is a subset of inTokens; clamp so a malformed
+		// upstream (cached > total) can't produce a negative remainder.
+		clamped := cachedInput
+		if clamped > inTokens {
+			clamped = inTokens
+		}
+		cachedRate := entry.CachedInputPer1K
+		if cachedRate <= 0 {
+			cachedRate = entry.InputPer1K
+		}
+		nonCached := float64(inTokens-clamped) / 1000.0 * entry.InputPer1K
+		cached := float64(clamped) / 1000.0 * cachedRate
+		return newCosts(nonCached, cached, 0, output), true
+	case "anthropic", "bedrock":
+		// Bedrock-Anthropic returns the same additive cache buckets as
+		// first-party Anthropic; non-Anthropic Bedrock models simply report
+		// zero cache tokens, so this formula degrades to input + output.
+		readRate := entry.CacheReadPer1K
+		if readRate <= 0 {
+			readRate = entry.InputPer1K
+		}
+		createRate := entry.CacheCreationPer1K
+		if createRate <= 0 {
+			createRate = entry.InputPer1K
+		}
+		input := float64(inTokens) / 1000.0 * entry.InputPer1K
+		read := float64(cachedInput) / 1000.0 * readRate
+		create := float64(cacheCreation) / 1000.0 * createRate
+		return newCosts(input, read, create, output), true
+	default:
+		input := float64(inTokens) / 1000.0 * entry.InputPer1K
+		return newCosts(input, 0, 0, output), true
+	}
+}
+
+// Has reports whether the provider/model pair is present in the table.
+func (t *Table) Has(provider, model string) bool {
+	if t == nil {
+		return false
+	}
+	byModel, ok := t.entries[provider]
+	if !ok {
+		return false
+	}
+	_, ok = byModel[model]
+	return ok
+}
+
+// pricingFile mirrors the on-disk YAML schema. Keys are provider names; the
+// nested map keys are model names.
+type pricingFile map[string]map[string]struct {
+	InputPer1K         float64 `yaml:"input_per_1k"`
+	OutputPer1K        float64 `yaml:"output_per_1k"`
+	CachedInputPer1K   float64 `yaml:"cached_input_per_1k"`
+	CacheReadPer1K     float64 `yaml:"cache_read_per_1k"`
+	CacheCreationPer1K float64 `yaml:"cache_creation_per_1k"`
+}
+
+const (
+	// ReloadInterval is the mtime-poll cadence for the background reloader.
+	ReloadInterval = 30 * time.Second
+
+	// errorBackoff bounds how often the loader logs a repeated parse error.
+	errorBackoff = 5 * time.Minute
+)
+
+var basenameRegex = regexp.MustCompile(`^[a-zA-Z0-9._-]+$`)
+
+// Loader is a confined, hot-reloadable pricing table reader. Construction
+// must succeed against the target file; subsequent reload failures keep the
+// previously-loaded table so callers never observe a blank price list.
+type Loader struct {
+	baseDir  string
+	fullPath string
+	pluginID string
+	table    atomic.Pointer[Table]
+	mtime    atomic.Int64
+	failures metric.Int64Counter
+	interval time.Duration
+}
+
+// NewLoader returns a pricing loader that overlays an optional file-based
+// table on top of the embedded defaults. Missing override file, baseDir, or
+// relPath is not an error: the loader keeps the embedded defaults so cost
+// metadata is still emitted for known models.
+//
+// Errors:
+//   - bad basename, traversal segment, or absolute relPath are rejected so a
+//     misconfigured target surfaces immediately.
+//   - permission errors and YAML parse errors keep the defaults but log a
+//     warning; cost annotation does not silently break.
+//
+// failures is optional; pass nil in tests that do not care about
+// reload-failure telemetry.
+func NewLoader(baseDir, relPath, pluginID string, failures metric.Int64Counter) (*Loader, error) {
+	defaults := DefaultTable()
+	l := &Loader{
+		baseDir:  baseDir,
+		pluginID: pluginID,
+		failures: failures,
+	}
+	l.table.Store(defaults)
+
+	if strings.TrimSpace(baseDir) == "" || strings.TrimSpace(relPath) == "" {
+		return l, nil
+	}
+
+	full, err := resolveMiddlewareDataPath(baseDir, relPath)
+	if err != nil {
+		return nil, err
+	}
+	l.fullPath = full
+
+	overlay, mtime, err := loadPricing(full)
+	if err != nil {
+		if errors.Is(err, fs.ErrNotExist) {
+			// Override file is optional. Defaults already stored.
+			return l, nil
+		}
+		// Symlink rejection, oversize file, parse failure, permission errors
+		// — surface so a misconfigured operator sees the problem instead of
+		// silently running with stale defaults.
+		return nil, fmt.Errorf("load pricing %s: %w", full, err)
+	}
+	l.table.Store(mergeOver(defaults, overlay))
+	l.mtime.Store(mtime.UnixNano())
+	return l, nil
+}
+
+// Get returns the current pricing table. The returned pointer is immutable;
+// callers must not mutate its contents.
+func (l *Loader) Get() *Table {
+	if l == nil {
+		return nil
+	}
+	return l.table.Load()
+}
+
+// WatchesFile reports whether this loader is bound to an override file on
+// disk. False for defaults-only loaders (no operator override given).
+// Callers use this to decide whether to spawn the mtime-poll goroutine.
+func (l *Loader) WatchesFile() bool {
+	if l == nil {
+		return false
+	}
+	return l.fullPath != ""
+}
+
+// SetReloadInterval overrides the mtime-poll cadence used by Reload. Calls
+// after Reload has started have no effect on the running loop. Intended for
+// tests; production code uses the default ReloadInterval.
+func (l *Loader) SetReloadInterval(d time.Duration) {
+	if l == nil || d <= 0 {
+		return
+	}
+	l.interval = d
+}
+
+// Reload runs a polling loop that checks the pricing file mtime every
+// ReloadInterval (or the value passed to SetReloadInterval). Returns when
+// ctx is cancelled.
+func (l *Loader) Reload(ctx context.Context) {
+	if l == nil {
+		return
+	}
+	interval := l.interval
+	if interval <= 0 {
+		interval = ReloadInterval
+	}
+	t := time.NewTicker(interval)
+	defer t.Stop()
+
+	var lastErrAt time.Time
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-t.C:
+			if err := l.reload(); err != nil {
+				if l.failures != nil {
+					l.failures.Add(ctx, 1, metric.WithAttributes(
+						attribute.String("plugin", l.pluginID),
+					))
+				}
+				now := time.Now()
+				if now.Sub(lastErrAt) >= errorBackoff {
+					log.Warnf("llmobs: pricing reload failed for %s: %v", l.fullPath, err)
+					lastErrAt = now
+				}
+			}
+		}
+	}
+}
+
+// reload performs a single-shot mtime check and reload. The reloaded
+// override file is merged on top of the embedded defaults; missing override
+// (e.g. operator deleted the file) is not an error and reverts to defaults.
+func (l *Loader) reload() error {
+	if l.fullPath == "" {
+		// Defaults-only loader; nothing on disk to reload.
+		return nil
+	}
+	mtime, err := statMtime(l.fullPath)
+	if err != nil {
+		if errors.Is(err, fs.ErrNotExist) {
+			// File was removed since startup. Drop back to defaults and
+			// reset mtime so a future re-creation triggers a reload.
+			l.table.Store(DefaultTable())
+			l.mtime.Store(0)
+			return nil
+		}
+		return err
+	}
+	if mtime.UnixNano() == l.mtime.Load() {
+		return nil
+	}
+
+	overlay, newMtime, err := loadPricing(l.fullPath)
+	if err != nil {
+		return err
+	}
+	l.table.Store(mergeOver(DefaultTable(), overlay))
+	l.mtime.Store(newMtime.UnixNano())
+	return nil
+}
+
+// resolveMiddlewareDataPath validates relPath is a safe basename and resolves
+// it under baseDir. An additional cleaned-prefix check guards against
+// CVE-style edge cases where Join is used with trailing path segments.
+func resolveMiddlewareDataPath(baseDir, relPath string) (string, error) {
+	if strings.TrimSpace(baseDir) == "" {
+		return "", errors.New("middleware-data-dir is not configured")
+	}
+	if relPath == "" {
+		return "", errors.New("pricing path is empty")
+	}
+	if !basenameRegex.MatchString(relPath) {
+		return "", fmt.Errorf("pricing path %q is not a safe basename", relPath)
+	}
+	if filepath.IsAbs(relPath) {
+		return "", fmt.Errorf("pricing path %q must be a basename, not absolute", relPath)
+	}
+
+	cleanBase, err := filepath.Abs(filepath.Clean(baseDir))
+	if err != nil {
+		return "", fmt.Errorf("resolve middleware-data-dir: %w", err)
+	}
+	full := filepath.Join(cleanBase, relPath)
+	cleanedFull := filepath.Clean(full)
+	if !strings.HasPrefix(cleanedFull, cleanBase+string(filepath.Separator)) && cleanedFull != cleanBase {
+		return "", fmt.Errorf("pricing path %q escapes middleware-data-dir", relPath)
+	}
+	return cleanedFull, nil
+}
+
+func parsePricingBytes(data []byte) (*Table, error) {
+	dec := yaml.NewDecoder(bytes.NewReader(data))
+	dec.KnownFields(true)
+
+	var raw pricingFile
+	if err := dec.Decode(&raw); err != nil && !errors.Is(err, io.EOF) {
+		return nil, fmt.Errorf("decode pricing yaml: %w", err)
+	}
+
+	out := make(map[string]map[string]Entry, len(raw))
+	for provider, models := range raw {
+		inner := make(map[string]Entry, len(models))
+		for model, entry := range models {
+			for field, v := range map[string]float64{
+				"input_per_1k":          entry.InputPer1K,
+				"output_per_1k":         entry.OutputPer1K,
+				"cached_input_per_1k":   entry.CachedInputPer1K,
+				"cache_read_per_1k":     entry.CacheReadPer1K,
+				"cache_creation_per_1k": entry.CacheCreationPer1K,
+			} {
+				if v < 0 || math.IsNaN(v) || math.IsInf(v, 0) {
+					return nil, fmt.Errorf("pricing %s/%s: %s must be a finite, non-negative rate, got %v", provider, model, field, v)
+				}
+			}
+			inner[model] = Entry{
+				InputPer1K:         entry.InputPer1K,
+				OutputPer1K:        entry.OutputPer1K,
+				CachedInputPer1K:   entry.CachedInputPer1K,
+				CacheReadPer1K:     entry.CacheReadPer1K,
+				CacheCreationPer1K: entry.CacheCreationPer1K,
+			}
+		}
+		out[provider] = inner
+	}
+	return &Table{entries: out}, nil
+}

proxy/internal/llm/pricing/pricing_other.go

@@ -0,0 +1,20 @@
+//go:build !unix
+
+package pricing
+
+import (
+	"fmt"
+	"time"
+)
+
+// loadPricing is unavailable on non-Unix platforms because O_NOFOLLOW and
+// fstat-from-FD are required to honour the spec's symlink-safety rules. The
+// proxy is only deployed on Linux today; a Windows port would need an
+// equivalent path-as-handle implementation.
+func loadPricing(path string) (*Table, time.Time, error) {
+	return nil, time.Time{}, fmt.Errorf("llmobs pricing loader is not supported on this platform: %s", path)
+}
+
+func statMtime(path string) (time.Time, error) {
+	return time.Time{}, fmt.Errorf("llmobs pricing loader is not supported on this platform: %s", path)
+}

proxy/internal/llm/pricing/pricing_test.go

@@ -0,0 +1,432 @@
+//go:build unix
+
+package pricing
+
+import (
+	"context"
+	"os"
+	"path/filepath"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func copyFixture(t *testing.T, src, dst string) {
+	t.Helper()
+	data, err := os.ReadFile(src)
+	require.NoError(t, err, "read source fixture")
+	require.NoError(t, os.WriteFile(dst, data, 0o600), "write target fixture")
+}
+
+func TestNewLoader_HappyPath(t *testing.T) {
+	base := t.TempDir()
+	copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), filepath.Join(base, "pricing.yaml"))
+
+	l, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
+	require.NoError(t, err, "NewLoader must succeed with a valid fixture")
+	table := l.Get()
+	require.NotNil(t, table, "table populated after load")
+
+	cost, ok := table.Cost("openai", "gpt-4o-mini", 1000, 1000, 0, 0)
+	require.True(t, ok, "known provider/model resolves")
+	assert.InDelta(t, 0.00075, cost, 1e-9, "cost = 0.00015 + 0.0006 per 1k tokens")
+
+	cost, ok = table.Cost("openai", "gpt-4o", 2000, 1000, 0, 0)
+	require.True(t, ok, "second known model resolves")
+	assert.InDelta(t, 0.015, cost, 1e-9, "cost for gpt-4o: 2*0.0025 + 1*0.01")
+
+	cost, ok = table.Cost("anthropic", "claude-sonnet-4-5", 1000, 1000, 0, 0)
+	require.True(t, ok, "anthropic model resolves")
+	assert.InDelta(t, 0.018, cost, 1e-9, "cost for claude-sonnet-4-5: 0.003 + 0.015")
+}
+
+// TestCost_OpenAICachedSubsetDiscount proves OpenAI's cached input
+// tokens are billed at the configured cached_input_per_1k rate while
+// the non-cached remainder of input_tokens is billed at the regular
+// rate. Critical because OpenAI returns cached_tokens as a SUBSET of
+// prompt_tokens — naïvely charging the cached count on top of
+// prompt_tokens would double-bill that portion.
+func TestCost_OpenAICachedSubsetDiscount(t *testing.T) {
+	tbl := &Table{entries: map[string]map[string]Entry{
+		"openai": {"gpt-4o": {
+			InputPer1K:       0.0025, // 0.0025 USD per 1k input tokens
+			OutputPer1K:      0.01,
+			CachedInputPer1K: 0.00125, // 0.5x discount on cached
+		}},
+	}}
+	// 1000 prompt tokens, 750 of which were cached. 250 non-cached
+	// at regular rate, 750 cached at the discount rate, 500 output.
+	cost, ok := tbl.Cost("openai", "gpt-4o", 1000, 500, 750, 0)
+	require.True(t, ok, "known model resolves")
+	want := (250.0/1000.0)*0.0025 + (750.0/1000.0)*0.00125 + (500.0/1000.0)*0.01
+	assert.InDelta(t, want, cost, 1e-12,
+		"cached subset must bill at the discount rate; non-cached remainder at regular rate")
+}
+
+// TestCost_OpenAICachedFallsBackToInputRate covers the operator
+// opt-in contract: when CachedInputPer1K is unset (zero), cached
+// tokens bill at the regular input rate. This matches today's
+// behaviour (cached counts weren't extracted at all so they
+// implicitly billed at the input rate via prompt_tokens).
+func TestCost_OpenAICachedFallsBackToInputRate(t *testing.T) {
+	tbl := &Table{entries: map[string]map[string]Entry{
+		"openai": {"gpt-4o": {InputPer1K: 0.0025, OutputPer1K: 0.01}},
+	}}
+	cost, ok := tbl.Cost("openai", "gpt-4o", 1000, 500, 750, 0)
+	require.True(t, ok)
+	want := 0.0025 + (500.0/1000.0)*0.01
+	assert.InDelta(t, want, cost, 1e-12,
+		"absent cached_input_per_1k rate must fall back to input_per_1k — same as pre-feature behaviour")
+}
+
+// TestCost_OpenAIClampsCachedToInputCount is the defensive guard
+// against malformed upstream responses that report cached_tokens >
+// prompt_tokens. We clamp so the formula never produces a negative
+// "non-cached remainder" multiplied by the input rate.
+func TestCost_OpenAIClampsCachedToInputCount(t *testing.T) {
+	tbl := &Table{entries: map[string]map[string]Entry{
+		"openai": {"gpt-4o": {InputPer1K: 0.0025, OutputPer1K: 0.01, CachedInputPer1K: 0.00125}},
+	}}
+	cost, ok := tbl.Cost("openai", "gpt-4o", 100, 0, 9999, 0)
+	require.True(t, ok)
+	// All 100 cached, 0 non-cached. Output is 0.
+	want := (100.0 / 1000.0) * 0.00125
+	assert.InDelta(t, want, cost, 1e-12,
+		"cached count > input count must clamp to input — never bill negative non-cached tokens")
+}
+
+// TestCost_AnthropicCacheReadAndCreationAreAdditive proves the
+// Anthropic shape: cache_read and cache_creation tokens are
+// ADDITIVE to input_tokens (not subset), each billed at its own
+// configured rate. The two rates pull in opposite directions —
+// cache_read is the cheaper read-from-cache rate (≈0.1× input),
+// cache_creation is the more expensive write-to-cache rate
+// (≈1.25× input).
+func TestCost_AnthropicCacheReadAndCreationAreAdditive(t *testing.T) {
+	tbl := &Table{entries: map[string]map[string]Entry{
+		"anthropic": {"claude-sonnet": {
+			InputPer1K:         0.003,
+			OutputPer1K:        0.015,
+			CacheReadPer1K:     0.0003,  // 0.1x of input
+			CacheCreationPer1K: 0.00375, // 1.25x of input
+		}},
+	}}
+	// 256 regular input + 768 cache_read + 512 cache_creation +
+	// 200 output. Each input bucket bills at its own rate.
+	cost, ok := tbl.Cost("anthropic", "claude-sonnet", 256, 200, 768, 512)
+	require.True(t, ok, "known model resolves")
+	want := (256.0/1000.0)*0.003 +
+		(768.0/1000.0)*0.0003 +
+		(512.0/1000.0)*0.00375 +
+		(200.0/1000.0)*0.015
+	assert.InDelta(t, want, cost, 1e-12,
+		"each Anthropic input bucket must bill at its own configured rate")
+}
+
+// TestCost_AnthropicCacheRatesFallBackToInput covers the no-opt-in
+// path: when neither CacheReadPer1K nor CacheCreationPer1K is set,
+// cache tokens bill at the regular input rate. This is more
+// accurate than today's behaviour (cache tokens ignored entirely)
+// without requiring operators to opt in via YAML.
+func TestCost_AnthropicCacheRatesFallBackToInput(t *testing.T) {
+	tbl := &Table{entries: map[string]map[string]Entry{
+		"anthropic": {"claude-sonnet": {InputPer1K: 0.003, OutputPer1K: 0.015}},
+	}}
+	cost, ok := tbl.Cost("anthropic", "claude-sonnet", 256, 200, 768, 512)
+	require.True(t, ok)
+	// Without overrides: every input bucket at input_per_1k.
+	want := ((256.0+768.0+512.0)/1000.0)*0.003 + (200.0/1000.0)*0.015
+	assert.InDelta(t, want, cost, 1e-12,
+		"absent cache rates must fall back to input_per_1k — Anthropic cache tokens were ignored before this change, billing at input rate is more accurate as a default")
+}
+
+func TestNewLoader_UnknownModel(t *testing.T) {
+	base := t.TempDir()
+	copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), filepath.Join(base, "pricing.yaml"))
+
+	l, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
+	require.NoError(t, err)
+
+	_, ok := l.Get().Cost("openai", "fantasy-model", 10, 10, 0, 0)
+	assert.False(t, ok, "unknown model returns ok=false")
+
+	_, ok = l.Get().Cost("cohere", "anything", 10, 10, 0, 0)
+	assert.False(t, ok, "unknown provider returns ok=false")
+}
+
+func TestNewLoader_InvalidYAMLRejected(t *testing.T) {
+	base := t.TempDir()
+	require.NoError(t, os.WriteFile(filepath.Join(base, "pricing.yaml"), []byte("\t- this is not: valid: yaml: :["), 0o600))
+
+	_, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
+	require.Error(t, err, "invalid YAML must surface as construction error")
+}
+
+func TestLoader_ReloadKeepsPreviousOnParseError(t *testing.T) {
+	base := t.TempDir()
+	target := filepath.Join(base, "pricing.yaml")
+	copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), target)
+
+	l, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
+	require.NoError(t, err)
+	require.NotNil(t, l.Get(), "initial table populated")
+
+	// Overwrite with content that violates the strict schema (extra field)
+	// plus a bumped mtime to trigger reload.
+	require.NoError(t, os.WriteFile(target, []byte("openai:\n  gpt-4o:\n    input_per_1k: 1.0\n    output_per_1k: 2.0\n    bogus_field: nope\n"), 0o600))
+	future := time.Now().Add(time.Hour)
+	require.NoError(t, os.Chtimes(target, future, future))
+
+	err = l.reload()
+	require.Error(t, err, "parse error surfaced by reload()")
+
+	cost, ok := l.Get().Cost("openai", "gpt-4o-mini", 1000, 1000, 0, 0)
+	require.True(t, ok, "previous table still available after parse failure")
+	assert.InDelta(t, 0.00075, cost, 1e-9, "previous cost preserved")
+}
+
+func TestLoader_ReloadNoChangeIsNoOp(t *testing.T) {
+	base := t.TempDir()
+	target := filepath.Join(base, "pricing.yaml")
+	copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), target)
+
+	l, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
+	require.NoError(t, err)
+	ptrBefore := l.Get()
+
+	require.NoError(t, l.reload(), "no-change reload must not error")
+	ptrAfter := l.Get()
+	assert.Same(t, ptrBefore, ptrAfter, "table pointer unchanged when mtime unchanged")
+}
+
+func TestLoader_ReloadDetectsChange(t *testing.T) {
+	base := t.TempDir()
+	target := filepath.Join(base, "pricing.yaml")
+	copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), target)
+
+	l, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
+	require.NoError(t, err)
+
+	updated := []byte("openai:\n  gpt-4o-mini:\n    input_per_1k: 1.00\n    output_per_1k: 2.00\n")
+	require.NoError(t, os.WriteFile(target, updated, 0o600))
+	future := time.Now().Add(time.Hour)
+	require.NoError(t, os.Chtimes(target, future, future))
+
+	require.NoError(t, l.reload(), "reload must succeed on valid new content")
+
+	cost, ok := l.Get().Cost("openai", "gpt-4o-mini", 1000, 1000, 0, 0)
+	require.True(t, ok, "updated model still present")
+	assert.InDelta(t, 3.0, cost, 0.0001, "new prices are applied: 1 + 2 per 1k")
+}
+
+// TestLoader_ReloadGoroutinePicksUpChanges proves the background goroutine
+// started via Reload actually swaps the pricing table when the file changes
+// on disk. Without that goroutine running, pricing edits would never reach
+// requests until a proxy restart.
+func TestLoader_ReloadGoroutinePicksUpChanges(t *testing.T) {
+	base := t.TempDir()
+	target := filepath.Join(base, "pricing.yaml")
+	copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), target)
+
+	l, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
+	require.NoError(t, err)
+	l.SetReloadInterval(20 * time.Millisecond)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	done := make(chan struct{})
+	go func() {
+		l.Reload(ctx)
+		close(done)
+	}()
+
+	// Before any rewrite, the loader holds the fixture's prices.
+	costBefore, ok := l.Get().Cost("openai", "gpt-4o-mini", 1000, 1000, 0, 0)
+	require.True(t, ok, "fixture model must resolve initially")
+	assert.InDelta(t, 0.00075, costBefore, 1e-9, "fixture prices apply before rewrite")
+
+	updated := []byte("openai:\n  gpt-4o-mini:\n    input_per_1k: 1.00\n    output_per_1k: 2.00\n")
+	require.NoError(t, os.WriteFile(target, updated, 0o600))
+	future := time.Now().Add(time.Hour)
+	require.NoError(t, os.Chtimes(target, future, future))
+
+	deadline := time.Now().Add(2 * time.Second)
+	for {
+		cost, ok := l.Get().Cost("openai", "gpt-4o-mini", 1000, 1000, 0, 0)
+		if ok && cost > 2.5 {
+			break
+		}
+		if time.Now().After(deadline) {
+			t.Fatalf("background reloader did not pick up rewrite within deadline")
+		}
+		time.Sleep(10 * time.Millisecond)
+	}
+
+	cancel()
+	select {
+	case <-done:
+	case <-time.After(2 * time.Second):
+		t.Fatal("Reload loop did not exit after cancel")
+	}
+}
+
+func TestLoader_ReloadBackgroundLoopCancellation(t *testing.T) {
+	base := t.TempDir()
+	copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), filepath.Join(base, "pricing.yaml"))
+	l, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
+	require.NoError(t, err)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	done := make(chan struct{})
+	go func() {
+		l.Reload(ctx)
+		close(done)
+	}()
+	cancel()
+
+	select {
+	case <-done:
+	case <-time.After(2 * time.Second):
+		t.Fatal("Reload loop did not exit on context cancel")
+	}
+}
+
+func TestNewLoader_PathValidation(t *testing.T) {
+	base := t.TempDir()
+	copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), filepath.Join(base, "pricing.yaml"))
+
+	cases := []struct {
+		name    string
+		relPath string
+	}{
+		{"traversal", "../../etc/passwd"},
+		{"absolute", "/etc/passwd"},
+		{"slash in basename", "sub/pricing.yaml"},
+		{"control chars", "pricing\x00.yaml"},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			_, err := NewLoader(base, tc.relPath, "llm_observability", nil)
+			require.Error(t, err, "NewLoader must reject %q", tc.relPath)
+		})
+	}
+
+	// Empty relPath is no longer a validation error: the loader treats it
+	// as "no override file, defaults only" so cost metadata is still
+	// emitted for the embedded models out of the box.
+	t.Run("empty falls back to defaults", func(t *testing.T) {
+		l, err := NewLoader(base, "", "llm_observability", nil)
+		require.NoError(t, err, "empty relPath should yield a defaults-only loader")
+		require.NotNil(t, l, "loader must be returned")
+		require.False(t, l.WatchesFile(), "no file watching when no override is given")
+		_, ok := l.Get().Cost("openai", "gpt-4o-mini", 1000, 1000, 0, 0)
+		assert.True(t, ok, "embedded defaults should still resolve gpt-4o-mini")
+	})
+}
+
+// TestNewLoader_PathValidation_Extended covers the remaining attack shapes
+// called out in C2: dot references, embedded traversal segments, and a
+// newline in the basename. The basename regex must reject each one even
+// though filepath.Clean would otherwise collapse them.
+func TestNewLoader_PathValidation_Extended(t *testing.T) {
+	base := t.TempDir()
+	copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), filepath.Join(base, "pricing.yaml"))
+
+	cases := []struct {
+		name    string
+		relPath string
+	}{
+		{"dot", "."},
+		{"dotdot", ".."},
+		{"relative traversal", "../pricing.yaml"},
+		{"embedded slash", "pri/cing.yaml"},
+		{"newline", "pricing\n.yaml"},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			_, err := NewLoader(base, tc.relPath, "llm_observability", nil)
+			require.Error(t, err, "NewLoader must reject %q", tc.relPath)
+		})
+	}
+}
+
+// TestNewLoader_ValidBasenameLoads proves the allowlist is exclusive: a
+// basename containing only safe characters under baseDir loads. Without this
+// a regression that over-tightened the regex would silently break valid
+// deployments.
+func TestNewLoader_ValidBasenameLoads(t *testing.T) {
+	base := t.TempDir()
+	copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), filepath.Join(base, "pricing-v2_prod.yaml"))
+
+	l, err := NewLoader(base, "pricing-v2_prod.yaml", "llm_observability", nil)
+	require.NoError(t, err, "basename with _, -, . must load")
+	require.NotNil(t, l.Get(), "table populated")
+}
+
+// TestNewLoader_SymlinkOutsideBaseDirRejected constructs a symlink under
+// baseDir that points to a file outside it. O_NOFOLLOW must refuse to open
+// the symlink even though the symlink path itself is a valid basename under
+// baseDir.
+func TestNewLoader_SymlinkOutsideBaseDirRejected(t *testing.T) {
+	outside := t.TempDir()
+	target := filepath.Join(outside, "evil.yaml")
+	copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), target)
+
+	base := t.TempDir()
+	link := filepath.Join(base, "pricing.yaml")
+	require.NoError(t, os.Symlink(target, link), "symlink setup")
+
+	_, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
+	require.Error(t, err, "O_NOFOLLOW must reject symlink even when it points outside baseDir")
+}
+
+func TestNewLoader_SymlinkRejected(t *testing.T) {
+	base := t.TempDir()
+	concrete := filepath.Join(base, "real.yaml")
+	copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), concrete)
+
+	link := filepath.Join(base, "pricing.yaml")
+	require.NoError(t, os.Symlink(concrete, link), "symlink setup")
+
+	_, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
+	require.Error(t, err, "O_NOFOLLOW must reject symlinked targets")
+}
+
+func TestTableCost_NilSafe(t *testing.T) {
+	var t1 *Table
+	cost, ok := t1.Cost("x", "y", 1, 1, 0, 0)
+	assert.False(t, ok, "nil table reports unknown")
+	assert.Zero(t, cost, "nil table returns zero cost")
+	assert.False(t, t1.Has("x", "y"), "nil table has nothing")
+}
+
+func TestLoaderGet_NilSafe(t *testing.T) {
+	var l *Loader
+	assert.Nil(t, l.Get(), "nil loader returns nil table")
+}
+
+// TestNewLoader_RejectsOversizedFile_FixesM4 proves the loader bounds reads
+// at maxPricingBytes so a hostile file cannot exhaust process memory.
+func TestNewLoader_RejectsOversizedFile_FixesM4(t *testing.T) {
+	base := t.TempDir()
+	target := filepath.Join(base, "pricing.yaml")
+
+	// Build a YAML payload larger than the cap. We pad with valid YAML
+	// comments so a partial read would still fail the size check rather
+	// than the parser.
+	header := "openai:\n"
+	bigComment := make([]byte, maxPricingBytes+1024)
+	for i := range bigComment {
+		bigComment[i] = ' '
+	}
+	bigComment[0] = '#'
+	bigComment[len(bigComment)-1] = '\n'
+	payload := append([]byte(header), bigComment...)
+	require.NoError(t, os.WriteFile(target, payload, 0o600))
+
+	_, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
+	require.Error(t, err, "oversized pricing file must be rejected")
+	assert.Contains(t, err.Error(), "exceeds", "rejection must reference the byte cap")
+}

proxy/internal/llm/pricing/pricing_unix.go

@@ -0,0 +1,68 @@
+//go:build unix
+
+package pricing
+
+import (
+	"fmt"
+	"io"
+	"os"
+	"syscall"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+)
+
+// maxPricingBytes caps the size of the pricing YAML on read so a hostile or
+// runaway file cannot exhaust process memory during reload. 1 MiB is several
+// orders of magnitude larger than any reasonable pricing table.
+const maxPricingBytes int64 = 1 << 20
+
+// loadPricing opens the file with O_NOFOLLOW, fstats the open descriptor,
+// and parses from that same descriptor. Never re-opens by path so a
+// mid-read rename or symlink swap cannot substitute content. Bytes are
+// capped at maxPricingBytes so the loader cannot be coerced into reading an
+// unbounded file.
+func loadPricing(path string) (*Table, time.Time, error) {
+	f, err := os.OpenFile(path, os.O_RDONLY|syscall.O_NOFOLLOW, 0)
+	if err != nil {
+		return nil, time.Time{}, fmt.Errorf("open %s: %w", path, err)
+	}
+	defer func() {
+		if cerr := f.Close(); cerr != nil {
+			log.Debugf("close pricing file %s: %v", path, cerr)
+		}
+	}()
+
+	info, err := f.Stat()
+	if err != nil {
+		return nil, time.Time{}, fmt.Errorf("fstat %s: %w", path, err)
+	}
+	if !info.Mode().IsRegular() {
+		return nil, time.Time{}, fmt.Errorf("pricing file %s is not a regular file", path)
+	}
+
+	data, err := io.ReadAll(io.LimitReader(f, maxPricingBytes+1))
+	if err != nil {
+		return nil, time.Time{}, fmt.Errorf("read %s: %w", path, err)
+	}
+	if int64(len(data)) > maxPricingBytes {
+		return nil, time.Time{}, fmt.Errorf("pricing file %s exceeds %d bytes", path, maxPricingBytes)
+	}
+
+	table, err := parsePricingBytes(data)
+	if err != nil {
+		return nil, time.Time{}, err
+	}
+	return table, info.ModTime(), nil
+}
+
+// statMtime returns the mtime of the file at path. It uses lstat semantics
+// via os.Lstat so a symlink swap is detected even though O_NOFOLLOW will
+// later reject the open.
+func statMtime(path string) (time.Time, error) {
+	info, err := os.Lstat(path)
+	if err != nil {
+		return time.Time{}, fmt.Errorf("lstat %s: %w", path, err)
+	}
+	return info.ModTime(), nil
+}

proxy/internal/llm/sse.go

@@ -0,0 +1,117 @@
+package llm
+
+import (
+	"bufio"
+	"errors"
+	"fmt"
+	"io"
+	"strings"
+)
+
+// Event represents a single server-sent event. Type is the dispatch name
+// carried on an "event:" line (empty when the stream uses only "data:"
+// lines). Data is the concatenation of every "data:" line that made up the
+// event, joined by a single newline.
+type Event struct {
+	Type string
+	Data string
+}
+
+// Scanner reads SSE events from an underlying byte stream. Events are
+// delimited by a blank line ("\n\n"). CRLF line endings are normalized to LF
+// transparently so fixtures captured from live servers can be replayed.
+//
+// Scanner is not safe for concurrent use.
+type Scanner struct {
+	r       *bufio.Reader
+	maxLine int
+}
+
+// NewScanner wraps the given reader. The default underlying buffer size is
+// large enough for typical provider events (~64 KiB); callers needing
+// larger events can wrap the reader in their own bufio.Reader beforehand.
+func NewScanner(r io.Reader) *Scanner {
+	return &Scanner{
+		r:       bufio.NewReaderSize(r, 64*1024),
+		maxLine: 1 << 20,
+	}
+}
+
+// Next returns the next event. It returns io.EOF after the final event has
+// been consumed. A trailing event that is not terminated by a blank line is
+// still returned before io.EOF so that servers which close the connection
+// without a trailing newline are handled correctly.
+func (s *Scanner) Next() (Event, error) {
+	var (
+		event   Event
+		dataBuf strings.Builder
+		hasData bool
+		hasAny  bool
+	)
+
+	for {
+		line, err := s.readLine()
+		if err != nil {
+			if errors.Is(err, io.EOF) && hasAny {
+				event.Data = dataBuf.String()
+				return event, nil
+			}
+			return Event{}, err
+		}
+
+		if line == "" {
+			if !hasAny {
+				continue
+			}
+			event.Data = dataBuf.String()
+			return event, nil
+		}
+
+		hasAny = true
+		if strings.HasPrefix(line, ":") {
+			continue
+		}
+
+		field, value := splitField(line)
+		switch field {
+		case "event":
+			event.Type = value
+		case "data":
+			if hasData {
+				dataBuf.WriteByte('\n')
+			}
+			dataBuf.WriteString(value)
+			hasData = true
+		}
+	}
+}
+
+func (s *Scanner) readLine() (string, error) {
+	line, err := s.r.ReadString('\n')
+	if err != nil {
+		if errors.Is(err, io.EOF) && line != "" {
+			return trimEOL(line), nil
+		}
+		return "", err
+	}
+	if len(line) > s.maxLine {
+		return "", fmt.Errorf("sse line exceeds %d bytes", s.maxLine)
+	}
+	return trimEOL(line), nil
+}
+
+func trimEOL(line string) string {
+	line = strings.TrimRight(line, "\n")
+	line = strings.TrimRight(line, "\r")
+	return line
+}
+
+func splitField(line string) (string, string) {
+	idx := strings.IndexByte(line, ':')
+	if idx < 0 {
+		return line, ""
+	}
+	field := line[:idx]
+	value := strings.TrimPrefix(line[idx+1:], " ")
+	return field, value
+}

proxy/internal/llm/sse_test.go

@@ -0,0 +1,175 @@
+package llm
+
+import (
+	"errors"
+	"io"
+	"os"
+	"path/filepath"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func collectEvents(t *testing.T, r io.Reader) []Event {
+	t.Helper()
+	s := NewScanner(r)
+	var out []Event
+	for {
+		ev, err := s.Next()
+		if errors.Is(err, io.EOF) {
+			return out
+		}
+		require.NoError(t, err, "unexpected error scanning SSE")
+		out = append(out, ev)
+	}
+}
+
+func TestSSEScanner_OpenAIFixture(t *testing.T) {
+	f, err := os.Open(filepath.Join("fixtures", "openai_stream.txt"))
+	require.NoError(t, err, "fixture must be openable")
+	defer f.Close()
+
+	events := collectEvents(t, f)
+	require.Len(t, events, 4, "expected 4 data frames (3 chunks + [DONE])")
+
+	for _, ev := range events {
+		assert.Empty(t, ev.Type, "OpenAI stream uses data-only frames")
+	}
+	assert.Contains(t, events[2].Data, `"usage"`, "third chunk carries usage block")
+	assert.Equal(t, "[DONE]", events[3].Data, "final frame is the OpenAI DONE sentinel")
+}
+
+func TestSSEScanner_AnthropicFixture(t *testing.T) {
+	f, err := os.Open(filepath.Join("fixtures", "anthropic_stream.txt"))
+	require.NoError(t, err, "fixture must be openable")
+	defer f.Close()
+
+	events := collectEvents(t, f)
+	require.Len(t, events, 7, "expected 7 Anthropic events")
+
+	types := make([]string, 0, len(events))
+	for _, ev := range events {
+		types = append(types, ev.Type)
+	}
+	assert.Equal(t, []string{
+		"message_start",
+		"content_block_start",
+		"content_block_delta",
+		"content_block_delta",
+		"content_block_stop",
+		"message_delta",
+		"message_stop",
+	}, types, "Anthropic event ordering matches fixture")
+
+	var deltaUsage Event
+	for _, ev := range events {
+		if ev.Type == "message_delta" {
+			deltaUsage = ev
+			break
+		}
+	}
+	assert.Contains(t, deltaUsage.Data, `"output_tokens":45`, "message_delta carries partial usage")
+}
+
+func TestSSEScanner_MultilineData(t *testing.T) {
+	raw := "event: ping\ndata: line1\ndata: line2\ndata: line3\n\n"
+	events := collectEvents(t, strings.NewReader(raw))
+
+	require.Len(t, events, 1, "one logical event from three data lines")
+	assert.Equal(t, "ping", events[0].Type, "event name honored")
+	assert.Equal(t, "line1\nline2\nline3", events[0].Data, "data lines joined with newline")
+}
+
+func TestSSEScanner_CRLF(t *testing.T) {
+	raw := "event: foo\r\ndata: bar\r\n\r\ndata: baz\r\n\r\n"
+	events := collectEvents(t, strings.NewReader(raw))
+
+	require.Len(t, events, 2, "CRLF-delimited events recognized")
+	assert.Equal(t, "foo", events[0].Type, "first event type preserved")
+	assert.Equal(t, "bar", events[0].Data, "first event data preserved")
+	assert.Empty(t, events[1].Type, "second event has no event name")
+	assert.Equal(t, "baz", events[1].Data, "second event data preserved")
+}
+
+func TestSSEScanner_EmptyInput(t *testing.T) {
+	s := NewScanner(strings.NewReader(""))
+	_, err := s.Next()
+	require.ErrorIs(t, err, io.EOF, "empty input yields immediate EOF")
+}
+
+func TestSSEScanner_CommentIgnored(t *testing.T) {
+	raw := ": this is a comment\ndata: hi\n\n"
+	events := collectEvents(t, strings.NewReader(raw))
+	require.Len(t, events, 1, "comment line does not emit an event")
+	assert.Equal(t, "hi", events[0].Data, "data line honoured after comment")
+}
+
+func TestSSEScanner_TrailingWithoutBlankLine(t *testing.T) {
+	raw := "event: foo\ndata: bar\n"
+	events := collectEvents(t, strings.NewReader(raw))
+	require.Len(t, events, 1, "trailing event without blank line still emitted")
+	assert.Equal(t, "foo", events[0].Type)
+	assert.Equal(t, "bar", events[0].Data)
+}
+
+// TestSSEScanner_ManyConsecutiveEmptyLines feeds a stream that is nothing
+// but empty lines. The scanner must terminate without panic — empty lines
+// alone do not constitute an event and must yield io.EOF.
+func TestSSEScanner_ManyConsecutiveEmptyLines(t *testing.T) {
+	raw := strings.Repeat("\n", 100)
+	s := NewScanner(strings.NewReader(raw))
+	_, err := s.Next()
+	require.ErrorIs(t, err, io.EOF, "100 empty lines must terminate as EOF without panic")
+}
+
+// TestSSEScanner_InterleavedCRLFAndLF mixes \r\n and \n terminators within
+// the same event. The scanner normalizes both and must still recover a
+// coherent event.
+func TestSSEScanner_InterleavedCRLFAndLF(t *testing.T) {
+	raw := "event: mix\r\ndata: first\ndata: second\r\n\n"
+	events := collectEvents(t, strings.NewReader(raw))
+	require.Len(t, events, 1, "mixed line endings must still produce one event")
+	assert.Equal(t, "mix", events[0].Type)
+	assert.Equal(t, "first\nsecond", events[0].Data, "both data lines joined")
+}
+
+// TestSSEScanner_LongSingleDataLine constructs a single data line that
+// exceeds the default bufio buffer (64 KiB) but stays under the scanner
+// maxLine. The scanner must round-trip the value intact without panicking
+// or truncating silently.
+func TestSSEScanner_LongSingleDataLine(t *testing.T) {
+	big := strings.Repeat("x", 80<<10)
+	raw := "data: " + big + "\n\n"
+	events := collectEvents(t, strings.NewReader(raw))
+	require.Len(t, events, 1, "long single-line event must be emitted")
+	assert.Equal(t, big, events[0].Data, "long data preserved")
+}
+
+// TestSSEScanner_BinaryGarbageInData validates that non-printable bytes
+// inside a data line do not crash the parser. The scanner should either
+// round-trip them or return a well-formed error — never panic.
+func TestSSEScanner_BinaryGarbageInData(t *testing.T) {
+	raw := "data: \x00\x01\x02\xff\xfe\n\n"
+	defer func() {
+		if r := recover(); r != nil {
+			t.Fatalf("scanner panicked on binary garbage: %v", r)
+		}
+	}()
+	s := NewScanner(strings.NewReader(raw))
+	ev, err := s.Next()
+	require.NoError(t, err, "binary bytes in data should not surface as error")
+	assert.Equal(t, "\x00\x01\x02\xff\xfe", ev.Data, "binary payload round-trips")
+}
+
+// TestSSEScanner_UnknownFieldsIgnored stresses the field parser by sending
+// unrecognized field names ("id:", "retry:", "custom:"). They must be
+// silently ignored per the SSE spec; the scanner must not panic or emit
+// spurious events.
+func TestSSEScanner_UnknownFieldsIgnored(t *testing.T) {
+	raw := "id: 1\nretry: 5000\ncustom: value\ndata: payload\n\n"
+	events := collectEvents(t, strings.NewReader(raw))
+	require.Len(t, events, 1, "unknown fields must not spawn extra events")
+	assert.Equal(t, "payload", events[0].Data, "data field survives amid unknown fields")
+}

proxy/internal/metrics/l4_metrics_test.go

@@ -0,0 +1,69 @@
+package metrics_test
+
+import (
+	"context"
+	"reflect"
+	"testing"
+	"time"
+
+	promexporter "go.opentelemetry.io/otel/exporters/prometheus"
+	sdkmetric "go.opentelemetry.io/otel/sdk/metric"
+
+	"github.com/netbirdio/netbird/proxy/internal/metrics"
+	"github.com/netbirdio/netbird/proxy/internal/types"
+)
+
+func newTestMetrics(t *testing.T) *metrics.Metrics {
+	t.Helper()
+
+	exporter, err := promexporter.New()
+	if err != nil {
+		t.Fatalf("create prometheus exporter: %v", err)
+	}
+
+	provider := sdkmetric.NewMeterProvider(sdkmetric.WithReader(exporter))
+	pkg := reflect.TypeOf(metrics.Metrics{}).PkgPath()
+	meter := provider.Meter(pkg)
+
+	m, err := metrics.New(context.Background(), meter)
+	if err != nil {
+		t.Fatalf("create metrics: %v", err)
+	}
+	return m
+}
+
+func TestL4ServiceGauge(t *testing.T) {
+	m := newTestMetrics(t)
+
+	m.L4ServiceAdded(types.ServiceModeTCP)
+	m.L4ServiceAdded(types.ServiceModeTCP)
+	m.L4ServiceAdded(types.ServiceModeUDP)
+	m.L4ServiceRemoved(types.ServiceModeTCP)
+}
+
+func TestTCPRelayMetrics(t *testing.T) {
+	m := newTestMetrics(t)
+
+	acct := types.AccountID("acct-1")
+
+	m.TCPRelayStarted(acct)
+	m.TCPRelayStarted(acct)
+	m.TCPRelayEnded(acct, 10*time.Second, 1000, 500)
+	m.TCPRelayDialError(acct)
+	m.TCPRelayRejected(acct)
+}
+
+func TestUDPSessionMetrics(t *testing.T) {
+	m := newTestMetrics(t)
+
+	acct := types.AccountID("acct-2")
+
+	m.UDPSessionStarted(acct)
+	m.UDPSessionStarted(acct)
+	m.UDPSessionEnded(acct)
+	m.UDPSessionDialError(acct)
+	m.UDPSessionRejected(acct)
+	m.UDPPacketRelayed(types.RelayDirectionClientToBackend, 100)
+	m.UDPPacketRelayed(types.RelayDirectionClientToBackend, 200)
+	m.UDPPacketRelayed(types.RelayDirectionBackendToClient, 150)
+}

proxy/internal/metrics/metrics.go

@@ -0,0 +1,459 @@
+package metrics
+
+import (
+	"context"
+	"net/http"
+	"sync"
+	"time"
+
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/metric"
+
+	"github.com/netbirdio/netbird/proxy/internal/proxy"
+	"github.com/netbirdio/netbird/proxy/internal/responsewriter"
+	"github.com/netbirdio/netbird/proxy/internal/types"
+)
+
+// Metrics collects OpenTelemetry metrics for the proxy.
+type Metrics struct {
+	ctx                      context.Context
+	meter                    metric.Meter
+	requestsTotal            metric.Int64Counter
+	activeRequests           metric.Int64UpDownCounter
+	configuredDomains        metric.Int64UpDownCounter
+	totalPaths               metric.Int64UpDownCounter
+	requestDuration          metric.Int64Histogram
+	backendDuration          metric.Int64Histogram
+	certificateIssueDuration metric.Int64Histogram
+
+	// Management sync metrics.
+	snapshotSyncDuration  metric.Int64Histogram
+	snapshotBatchDuration metric.Int64Histogram
+	addPeerDuration       metric.Int64Histogram
+
+	// L4 service-level metrics.
+	l4Services metric.Int64UpDownCounter
+
+	// L4 TCP connection-level metrics.
+	tcpActiveConns  metric.Int64UpDownCounter
+	tcpConnsTotal   metric.Int64Counter
+	tcpConnDuration metric.Int64Histogram
+	tcpBytesTotal   metric.Int64Counter
+
+	// L4 UDP session-level metrics.
+	udpActiveSess    metric.Int64UpDownCounter
+	udpSessionsTotal metric.Int64Counter
+	udpPacketsTotal  metric.Int64Counter
+	udpBytesTotal    metric.Int64Counter
+
+	mappingsMux  sync.Mutex
+	mappingPaths map[string]int
+}
+
+// Meter returns the OpenTelemetry meter the bundle was built with, so other
+// subsystems (e.g. the middleware manager) register instruments on the same
+// meter.
+func (m *Metrics) Meter() metric.Meter {
+	return m.meter
+}
+
+// New creates a Metrics instance using the given OpenTelemetry meter.
+func New(ctx context.Context, meter metric.Meter) (*Metrics, error) {
+	m := &Metrics{
+		ctx:          ctx,
+		meter:        meter,
+		mappingPaths: make(map[string]int),
+	}
+
+	if err := m.initHTTPMetrics(meter); err != nil {
+		return nil, err
+	}
+	if err := m.initSyncMetrics(meter); err != nil {
+		return nil, err
+	}
+	if err := m.initL4Metrics(meter); err != nil {
+		return nil, err
+	}
+
+	return m, nil
+}
+
+func (m *Metrics) initHTTPMetrics(meter metric.Meter) error {
+	var err error
+
+	m.requestsTotal, err = meter.Int64Counter(
+		"proxy.http.request.counter",
+		metric.WithUnit("1"),
+		metric.WithDescription("Total number of requests made to the netbird proxy"),
+	)
+	if err != nil {
+		return err
+	}
+
+	m.activeRequests, err = meter.Int64UpDownCounter(
+		"proxy.http.active_requests",
+		metric.WithUnit("1"),
+		metric.WithDescription("Current in-flight requests handled by the netbird proxy"),
+	)
+	if err != nil {
+		return err
+	}
+
+	m.configuredDomains, err = meter.Int64UpDownCounter(
+		"proxy.domains.count",
+		metric.WithUnit("1"),
+		metric.WithDescription("Current number of domains configured on the netbird proxy"),
+	)
+	if err != nil {
+		return err
+	}
+
+	m.totalPaths, err = meter.Int64UpDownCounter(
+		"proxy.paths.count",
+		metric.WithUnit("1"),
+		metric.WithDescription("Total number of paths configured on the netbird proxy"),
+	)
+	if err != nil {
+		return err
+	}
+
+	m.requestDuration, err = meter.Int64Histogram(
+		"proxy.http.request.duration.ms",
+		metric.WithUnit("milliseconds"),
+		metric.WithDescription("Duration of requests made to the netbird proxy"),
+	)
+	if err != nil {
+		return err
+	}
+
+	m.backendDuration, err = meter.Int64Histogram(
+		"proxy.backend.duration.ms",
+		metric.WithUnit("milliseconds"),
+		metric.WithDescription("Duration of peer round trip time from the netbird proxy"),
+	)
+	if err != nil {
+		return err
+	}
+
+	m.certificateIssueDuration, err = meter.Int64Histogram(
+		"proxy.certificate.issue.duration.ms",
+		metric.WithUnit("milliseconds"),
+		metric.WithDescription("Duration of ACME certificate issuance"),
+	)
+	return err
+}
+
+func (m *Metrics) initSyncMetrics(meter metric.Meter) error {
+	var err error
+
+	m.snapshotSyncDuration, err = meter.Int64Histogram(
+		"proxy.sync.snapshot.duration.ms",
+		metric.WithUnit("milliseconds"),
+		metric.WithDescription("Duration from management connect until the initial snapshot sync is complete"),
+		metric.WithExplicitBucketBoundaries(100, 250, 500, 1000, 2500, 5000, 10000, 30000, 60000, 120000, 300000),
+	)
+	if err != nil {
+		return err
+	}
+
+	m.snapshotBatchDuration, err = meter.Int64Histogram(
+		"proxy.sync.batch.duration.ms",
+		metric.WithUnit("milliseconds"),
+		metric.WithDescription("Duration to process a single mapping batch during initial snapshot sync"),
+		metric.WithExplicitBucketBoundaries(100, 250, 500, 1000, 2500, 5000, 10000, 30000, 60000, 120000, 300000),
+	)
+	if err != nil {
+		return err
+	}
+
+	m.addPeerDuration, err = meter.Int64Histogram(
+		"proxy.peer.add.duration.ms",
+		metric.WithUnit("milliseconds"),
+		metric.WithDescription("Duration to add a peer for an account (keygen + gRPC CreateProxyPeer + embed.New)"),
+		metric.WithExplicitBucketBoundaries(10, 25, 50, 100, 250, 500, 1000, 2500, 5000, 10000),
+	)
+	return err
+}
+
+// RecordSnapshotSyncDuration records the total time from connect to sync-complete.
+func (m *Metrics) RecordSnapshotSyncDuration(d time.Duration) {
+	m.snapshotSyncDuration.Record(m.ctx, d.Milliseconds())
+}
+
+// RecordSnapshotBatchDuration records the time to process one mapping batch during initial sync.
+func (m *Metrics) RecordSnapshotBatchDuration(d time.Duration) {
+	m.snapshotBatchDuration.Record(m.ctx, d.Milliseconds())
+}
+
+// RecordAddPeerDuration records the time to create a new peer for an account.
+func (m *Metrics) RecordAddPeerDuration(d time.Duration, err error) {
+	result := "success"
+	if err != nil {
+		result = "error"
+	}
+	m.addPeerDuration.Record(m.ctx, d.Milliseconds(), metric.WithAttributes(
+		attribute.String("result", result),
+	))
+}
+
+func (m *Metrics) initL4Metrics(meter metric.Meter) error {
+	var err error
+
+	m.l4Services, err = meter.Int64UpDownCounter(
+		"proxy.l4.services.count",
+		metric.WithUnit("1"),
+		metric.WithDescription("Current number of configured L4 services (TCP/TLS/UDP) by mode"),
+	)
+	if err != nil {
+		return err
+	}
+
+	m.tcpActiveConns, err = meter.Int64UpDownCounter(
+		"proxy.tcp.active_connections",
+		metric.WithUnit("1"),
+		metric.WithDescription("Current number of active TCP/TLS relay connections"),
+	)
+	if err != nil {
+		return err
+	}
+
+	m.tcpConnsTotal, err = meter.Int64Counter(
+		"proxy.tcp.connections.total",
+		metric.WithUnit("1"),
+		metric.WithDescription("Total TCP/TLS relay connections by result and account"),
+	)
+	if err != nil {
+		return err
+	}
+
+	m.tcpConnDuration, err = meter.Int64Histogram(
+		"proxy.tcp.connection.duration.ms",
+		metric.WithUnit("milliseconds"),
+		metric.WithDescription("Duration of TCP/TLS relay connections"),
+	)
+	if err != nil {
+		return err
+	}
+
+	m.tcpBytesTotal, err = meter.Int64Counter(
+		"proxy.tcp.bytes.total",
+		metric.WithUnit("bytes"),
+		metric.WithDescription("Total bytes transferred through TCP/TLS relay by direction"),
+	)
+	if err != nil {
+		return err
+	}
+
+	m.udpActiveSess, err = meter.Int64UpDownCounter(
+		"proxy.udp.active_sessions",
+		metric.WithUnit("1"),
+		metric.WithDescription("Current number of active UDP relay sessions"),
+	)
+	if err != nil {
+		return err
+	}
+
+	m.udpSessionsTotal, err = meter.Int64Counter(
+		"proxy.udp.sessions.total",
+		metric.WithUnit("1"),
+		metric.WithDescription("Total UDP relay sessions by result and account"),
+	)
+	if err != nil {
+		return err
+	}
+
+	m.udpPacketsTotal, err = meter.Int64Counter(
+		"proxy.udp.packets.total",
+		metric.WithUnit("1"),
+		metric.WithDescription("Total UDP packets relayed by direction"),
+	)
+	if err != nil {
+		return err
+	}
+
+	m.udpBytesTotal, err = meter.Int64Counter(
+		"proxy.udp.bytes.total",
+		metric.WithUnit("bytes"),
+		metric.WithDescription("Total bytes transferred through UDP relay by direction"),
+	)
+	return err
+}
+
+type responseInterceptor struct {
+	*responsewriter.PassthroughWriter
+	status int
+	size   int
+}
+
+func (w *responseInterceptor) WriteHeader(status int) {
+	w.status = status
+	w.PassthroughWriter.WriteHeader(status)
+}
+
+func (w *responseInterceptor) Write(b []byte) (int, error) {
+	size, err := w.PassthroughWriter.Write(b)
+	w.size += size
+	return size, err
+}
+
+// Unwrap returns the underlying ResponseWriter so http.ResponseController
+// can reach through to the original writer for Hijack/Flush operations.
+func (w *responseInterceptor) Unwrap() http.ResponseWriter {
+	return w.PassthroughWriter
+}
+
+// Middleware wraps an HTTP handler with request metrics.
+func (m *Metrics) Middleware(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		m.requestsTotal.Add(m.ctx, 1)
+		m.activeRequests.Add(m.ctx, 1)
+
+		interceptor := &responseInterceptor{PassthroughWriter: responsewriter.New(w)}
+
+		start := time.Now()
+		defer func() {
+			duration := time.Since(start)
+			m.activeRequests.Add(m.ctx, -1)
+			m.requestDuration.Record(m.ctx, duration.Milliseconds())
+		}()
+
+		next.ServeHTTP(interceptor, r)
+	})
+}
+
+type roundTripperFunc func(*http.Request) (*http.Response, error)
+
+func (f roundTripperFunc) RoundTrip(r *http.Request) (*http.Response, error) {
+	return f(r)
+}
+
+// RoundTripper wraps an http.RoundTripper with backend duration metrics.
+func (m *Metrics) RoundTripper(next http.RoundTripper) http.RoundTripper {
+	return roundTripperFunc(func(req *http.Request) (*http.Response, error) {
+		start := time.Now()
+		res, err := next.RoundTrip(req)
+		duration := time.Since(start)
+
+		m.backendDuration.Record(m.ctx, duration.Milliseconds())
+
+		return res, err
+	})
+}
+
+// AddMapping records that a domain mapping was added.
+func (m *Metrics) AddMapping(mapping proxy.Mapping) {
+	m.mappingsMux.Lock()
+	defer m.mappingsMux.Unlock()
+
+	newPathCount := len(mapping.Paths)
+	oldPathCount, exists := m.mappingPaths[mapping.Host]
+
+	if !exists {
+		m.configuredDomains.Add(m.ctx, 1)
+	}
+
+	pathDelta := newPathCount - oldPathCount
+	if pathDelta != 0 {
+		m.totalPaths.Add(m.ctx, int64(pathDelta))
+	}
+
+	m.mappingPaths[mapping.Host] = newPathCount
+}
+
+// RemoveMapping records that a domain mapping was removed.
+func (m *Metrics) RemoveMapping(mapping proxy.Mapping) {
+	m.mappingsMux.Lock()
+	defer m.mappingsMux.Unlock()
+
+	oldPathCount, exists := m.mappingPaths[mapping.Host]
+	if !exists {
+		return
+	}
+
+	m.configuredDomains.Add(m.ctx, -1)
+	m.totalPaths.Add(m.ctx, -int64(oldPathCount))
+
+	delete(m.mappingPaths, mapping.Host)
+}
+
+// RecordCertificateIssuance records the duration of a certificate issuance.
+func (m *Metrics) RecordCertificateIssuance(duration time.Duration) {
+	m.certificateIssueDuration.Record(m.ctx, duration.Milliseconds())
+}
+
+// L4ServiceAdded increments the L4 service gauge for the given mode.
+func (m *Metrics) L4ServiceAdded(mode types.ServiceMode) {
+	m.l4Services.Add(m.ctx, 1, metric.WithAttributes(attribute.String("mode", string(mode))))
+}
+
+// L4ServiceRemoved decrements the L4 service gauge for the given mode.
+func (m *Metrics) L4ServiceRemoved(mode types.ServiceMode) {
+	m.l4Services.Add(m.ctx, -1, metric.WithAttributes(attribute.String("mode", string(mode))))
+}
+
+// TCPRelayStarted records a new TCP relay connection starting.
+func (m *Metrics) TCPRelayStarted(accountID types.AccountID) {
+	acct := attribute.String("account_id", string(accountID))
+	m.tcpActiveConns.Add(m.ctx, 1, metric.WithAttributes(acct))
+	m.tcpConnsTotal.Add(m.ctx, 1, metric.WithAttributes(acct, attribute.String("result", "success")))
+}
+
+// TCPRelayEnded records a TCP relay connection ending and accumulates bytes and duration.
+func (m *Metrics) TCPRelayEnded(accountID types.AccountID, duration time.Duration, srcToDst, dstToSrc int64) {
+	acct := attribute.String("account_id", string(accountID))
+	m.tcpActiveConns.Add(m.ctx, -1, metric.WithAttributes(acct))
+	m.tcpConnDuration.Record(m.ctx, duration.Milliseconds(), metric.WithAttributes(acct))
+	m.tcpBytesTotal.Add(m.ctx, srcToDst, metric.WithAttributes(attribute.String("direction", "client_to_backend")))
+	m.tcpBytesTotal.Add(m.ctx, dstToSrc, metric.WithAttributes(attribute.String("direction", "backend_to_client")))
+}
+
+// TCPRelayDialError records a dial failure for a TCP relay.
+func (m *Metrics) TCPRelayDialError(accountID types.AccountID) {
+	m.tcpConnsTotal.Add(m.ctx, 1, metric.WithAttributes(
+		attribute.String("account_id", string(accountID)),
+		attribute.String("result", "dial_error"),
+	))
+}
+
+// TCPRelayRejected records a rejected TCP relay (semaphore full).
+func (m *Metrics) TCPRelayRejected(accountID types.AccountID) {
+	m.tcpConnsTotal.Add(m.ctx, 1, metric.WithAttributes(
+		attribute.String("account_id", string(accountID)),
+		attribute.String("result", "rejected"),
+	))
+}
+
+// UDPSessionStarted records a new UDP session starting.
+func (m *Metrics) UDPSessionStarted(accountID types.AccountID) {
+	acct := attribute.String("account_id", string(accountID))
+	m.udpActiveSess.Add(m.ctx, 1, metric.WithAttributes(acct))
+	m.udpSessionsTotal.Add(m.ctx, 1, metric.WithAttributes(acct, attribute.String("result", "success")))
+}
+
+// UDPSessionEnded records a UDP session ending.
+func (m *Metrics) UDPSessionEnded(accountID types.AccountID) {
+	m.udpActiveSess.Add(m.ctx, -1, metric.WithAttributes(attribute.String("account_id", string(accountID))))
+}
+
+// UDPSessionDialError records a dial failure for a UDP session.
+func (m *Metrics) UDPSessionDialError(accountID types.AccountID) {
+	m.udpSessionsTotal.Add(m.ctx, 1, metric.WithAttributes(
+		attribute.String("account_id", string(accountID)),
+		attribute.String("result", "dial_error"),
+	))
+}
+
+// UDPSessionRejected records a rejected UDP session (limit or rate limited).
+func (m *Metrics) UDPSessionRejected(accountID types.AccountID) {
+	m.udpSessionsTotal.Add(m.ctx, 1, metric.WithAttributes(
+		attribute.String("account_id", string(accountID)),
+		attribute.String("result", "rejected"),
+	))
+}
+
+// UDPPacketRelayed records a packet relayed in the given direction with its size in bytes.
+func (m *Metrics) UDPPacketRelayed(direction types.RelayDirection, bytes int) {
+	dir := attribute.String("direction", string(direction))
+	m.udpPacketsTotal.Add(m.ctx, 1, metric.WithAttributes(dir))
+	m.udpBytesTotal.Add(m.ctx, int64(bytes), metric.WithAttributes(dir))
+}

proxy/internal/metrics/metrics_test.go

@@ -0,0 +1,83 @@
+package metrics_test
+
+import (
+	"context"
+	"net/http"
+	"net/url"
+	"reflect"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"go.opentelemetry.io/otel/exporters/prometheus"
+	"go.opentelemetry.io/otel/sdk/metric"
+
+	"github.com/netbirdio/netbird/proxy/internal/metrics"
+)
+
+type testRoundTripper struct {
+	response *http.Response
+	err      error
+}
+
+func (t *testRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
+	return t.response, t.err
+}
+
+func TestMetrics_RoundTripper(t *testing.T) {
+	testResponse := http.Response{
+		StatusCode: http.StatusOK,
+		Body:       http.NoBody,
+	}
+
+	tests := map[string]struct {
+		roundTripper http.RoundTripper
+		request      *http.Request
+		response     *http.Response
+		err          error
+	}{
+		"ok": {
+			roundTripper: &testRoundTripper{response: &testResponse},
+			request:      &http.Request{Method: "GET", URL: &url.URL{Path: "/foo"}},
+			response:     &testResponse,
+		},
+		"nil url": {
+			roundTripper: &testRoundTripper{response: &testResponse},
+			request:      &http.Request{Method: "GET", URL: nil},
+			response:     &testResponse,
+		},
+		"nil response": {
+			roundTripper: &testRoundTripper{response: nil},
+			request:      &http.Request{Method: "GET", URL: &url.URL{Path: "/foo"}},
+		},
+	}
+
+	exporter, err := prometheus.New()
+	if err != nil {
+		t.Fatalf("create prometheus exporter: %v", err)
+	}
+
+	provider := metric.NewMeterProvider(metric.WithReader(exporter))
+	pkg := reflect.TypeOf(metrics.Metrics{}).PkgPath()
+	meter := provider.Meter(pkg)
+
+	m, err := metrics.New(context.Background(), meter)
+	if err != nil {
+		t.Fatalf("create metrics: %v", err)
+	}
+
+	for name, test := range tests {
+		t.Run(name, func(t *testing.T) {
+			rt := m.RoundTripper(test.roundTripper)
+			res, err := rt.RoundTrip(test.request)
+			if res != nil && res.Body != nil {
+				defer res.Body.Close()
+			}
+			if diff := cmp.Diff(test.err, err); diff != "" {
+				t.Errorf("Incorrect error (-want +got):\n%s", diff)
+			}
+			if diff := cmp.Diff(test.response, res); diff != "" {
+				t.Errorf("Incorrect response (-want +got):\n%s", diff)
+			}
+		})
+	}
+}

proxy/internal/middleware/bodypolicy.go

@@ -0,0 +1,63 @@
+package middleware
+
+import (
+	"bytes"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"strconv"
+	"strings"
+)
+
+// ErrExpectContinue is returned when a middleware attempts to replace
+// the body of a request that advertised Expect: 100-continue.
+var ErrExpectContinue = errors.New("body replace rejected: request has Expect: 100-continue")
+
+// ErrOriginalNotDrained is returned when the original body was not
+// fully consumed before replacement. This prevents the backend from
+// seeing a mix of original bytes and the replacement.
+var ErrOriginalNotDrained = errors.New("body replace rejected: original body not drained")
+
+// ErrContentLengthMismatch is returned when the client-advertised
+// Content-Length disagrees with the number of bytes actually read from
+// the body (short-read).
+var ErrContentLengthMismatch = errors.New("body replace rejected: content-length mismatch (short read)")
+
+// ValidateBodyReplace runs the smuggling-prevention rules before a
+// body replacement is applied. Callers must pass originalDrained=true
+// once they have read r.Body to EOF.
+func ValidateBodyReplace(r *http.Request, newBody []byte, originalDrained bool) error {
+	if r == nil {
+		return errors.New("body replace rejected: nil request")
+	}
+	if strings.EqualFold(r.Header.Get("Expect"), "100-continue") {
+		return ErrExpectContinue
+	}
+	if !originalDrained {
+		return ErrOriginalNotDrained
+	}
+	if cl := r.Header.Get("Content-Length"); cl != "" && r.ContentLength > 0 {
+		parsed, err := strconv.ParseInt(cl, 10, 64)
+		if err == nil && parsed != r.ContentLength {
+			return fmt.Errorf("%w: header=%d actual=%d", ErrContentLengthMismatch, parsed, r.ContentLength)
+		}
+	}
+	return nil
+}
+
+// ApplyBodyReplace swaps r.Body for a reader over newBody, recomputes
+// Content-Length, and strips Transfer-Encoding and Trailer so no stale
+// framing reaches the backend.
+func ApplyBodyReplace(r *http.Request, newBody []byte) {
+	if r == nil {
+		return
+	}
+	r.Body = io.NopCloser(bytes.NewReader(newBody))
+	r.ContentLength = int64(len(newBody))
+	r.Header.Set("Content-Length", strconv.Itoa(len(newBody)))
+	r.Header.Del("Transfer-Encoding")
+	r.Header.Del("Trailer")
+	r.TransferEncoding = nil
+	r.Trailer = nil
+}

proxy/internal/middleware/bodytap/request.go

@@ -0,0 +1,344 @@
+// Package bodytap owns the framework-side body capture used by the
+// middleware chain. Request capture buffers up to N bytes of the
+// request body for middleware inspection while replaying the original
+// stream to the upstream. Response capture tees up to N bytes off the
+// streaming response while every byte continues to flow to the client
+// untouched.
+//
+// The package is the single owner of body access — middlewares never
+// read req.Body or hijack the response writer. All inspection happens
+// against the buffer surfaced by the tap, so streaming remains
+// transparent to the client even when middlewares need access to the
+// payload.
+package bodytap
+
+import (
+	"bytes"
+	"encoding/json"
+	"errors"
+	"io"
+	"net/http"
+	"strconv"
+	"strings"
+	"sync"
+)
+
+// MaxRoutingScanBytes bounds how far ScanRoutingFields will read into a
+// request body to recover routing fields when the normal capture is
+// bypassed for size. Sized to comfortably hold a 1M-token context
+// request (whose `model` field a client may place after a multi-MB
+// `messages` array) while still capping pathological inputs.
+const MaxRoutingScanBytes int64 = 32 << 20
+
+// Request bypass reasons emitted as the `mw.capture.bypass_reason`
+// metadata key by the chain when a request body is not surfaced.
+const (
+	BypassUpgradeHeader    = "upgrade_header"
+	BypassConnectionUpgrd  = "connection_upgrade"
+	BypassContentType      = "content_type_not_allowed"
+	BypassBudget           = "capture_budget_exhausted"
+	BypassNoConfig         = "no_capture_config"
+	BypassNoMiddlewares    = "no_middlewares"
+	BypassCapZero          = "cap_zero"
+	BypassContentLengthCap = "content_length_over_cap"
+)
+
+// DefaultCaptureBudgetBytes is the default global capture-budget size.
+const DefaultCaptureBudgetBytes int64 = 256 << 20
+
+// Config holds per-target body capture limits after clamp validation.
+// A zero MaxRequestBytes / MaxResponseBytes disables capture in that
+// direction.
+type Config struct {
+	MaxRequestBytes  int64
+	MaxResponseBytes int64
+	ContentTypes     []string
+}
+
+// Budget is the global token-bucket semaphore shared across all
+// in-flight captures so a single misbehaving target cannot exhaust the
+// proxy.
+type Budget interface {
+	Acquire(n int64) bool
+	Release(n int64)
+}
+
+// NewBudget returns a Budget with the given total byte cap. A zero or
+// negative total disables the budget check.
+func NewBudget(total int64) Budget {
+	return &budget{total: total}
+}
+
+type budget struct {
+	mu    sync.Mutex
+	used  int64
+	total int64
+}
+
+func (b *budget) Acquire(n int64) bool {
+	if n <= 0 {
+		return true
+	}
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	if b.total <= 0 {
+		return true
+	}
+	if b.used+n > b.total {
+		return false
+	}
+	b.used += n
+	return true
+}
+
+func (b *budget) Release(n int64) {
+	if n <= 0 {
+		return
+	}
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	if b.total <= 0 {
+		return
+	}
+	b.used -= n
+	if b.used < 0 {
+		b.used = 0
+	}
+}
+
+// CaptureRequest reads up to cfg.MaxRequestBytes from r.Body into a
+// buffer suitable for middleware inspection, replacing r.Body with a
+// replay reader so the upstream still sees the original bytes. When
+// bypass != "" no body is read and r.Body is left untouched. The
+// returned release function must be invoked once the request is fully
+// processed; it returns the acquired budget tokens to the shared pool.
+// release is always non-nil and is safe to defer immediately after the
+// call.
+func CaptureRequest(r *http.Request, cfg *Config, b Budget) (body []byte, truncated bool, originalSize int64, bypass string, release func(), err error) {
+	release = func() {}
+	if r == nil {
+		return nil, false, 0, BypassNoConfig, release, nil
+	}
+	if cfg == nil {
+		return nil, false, 0, BypassNoConfig, release, nil
+	}
+	if cfg.MaxRequestBytes <= 0 {
+		return nil, false, 0, BypassCapZero, release, nil
+	}
+	if r.Header.Get("Upgrade") != "" {
+		return nil, false, 0, BypassUpgradeHeader, release, nil
+	}
+	if strings.EqualFold(r.Header.Get("Connection"), "upgrade") {
+		return nil, false, 0, BypassConnectionUpgrd, release, nil
+	}
+	if !contentTypeAllowed(r.Header.Get("Content-Type"), cfg.ContentTypes) {
+		return nil, false, 0, BypassContentType, release, nil
+	}
+
+	originalSize = parseContentLength(r.Header.Get("Content-Length"))
+	if originalSize > cfg.MaxRequestBytes {
+		return nil, true, originalSize, BypassContentLengthCap, release, nil
+	}
+
+	limit := cfg.MaxRequestBytes
+	if b != nil && !b.Acquire(limit) {
+		return nil, false, originalSize, BypassBudget, release, nil
+	}
+	if b != nil {
+		var released sync.Once
+		release = func() {
+			released.Do(func() { b.Release(limit) })
+		}
+	}
+
+	if r.Body == nil || r.Body == http.NoBody {
+		release()
+		release = func() {}
+		return nil, false, originalSize, "", release, nil
+	}
+
+	limited := io.LimitReader(r.Body, limit+1)
+	buf, readErr := io.ReadAll(limited)
+	if readErr != nil && !errors.Is(readErr, io.EOF) {
+		release()
+		release = func() {}
+		return nil, false, originalSize, "", release, readErr
+	}
+
+	truncated = int64(len(buf)) > limit
+	if truncated {
+		replay := append([]byte(nil), buf...)
+		viewable := buf[:limit]
+		r.Body = &replayReadCloser{replay: bytes.NewReader(replay), tail: r.Body}
+		return viewable, true, originalSize, "", release, nil
+	}
+	_ = r.Body.Close()
+	r.Body = io.NopCloser(bytes.NewReader(buf))
+	if originalSize <= 0 {
+		originalSize = int64(len(buf))
+	}
+	return buf, false, originalSize, "", release, nil
+}
+
+// replayReadCloser replays the captured prefix and then forwards the
+// remaining bytes from the original body so the upstream sees the
+// full request stream even when capture truncates.
+type replayReadCloser struct {
+	replay  *bytes.Reader
+	tail    io.ReadCloser
+	drained bool
+}
+
+func (r *replayReadCloser) Read(p []byte) (int, error) {
+	if !r.drained {
+		n, err := r.replay.Read(p)
+		if n > 0 {
+			return n, nil
+		}
+		if errors.Is(err, io.EOF) {
+			r.drained = true
+		} else if err != nil {
+			return 0, err
+		}
+	}
+	return r.tail.Read(p)
+}
+
+func (r *replayReadCloser) Close() error {
+	return r.tail.Close()
+}
+
+// ScanRoutingFields recovers the LLM routing fields ("model" and
+// "stream") from a request whose normal capture was bypassed or
+// truncated for size. It reads up to maxScan bytes of r.Body to locate
+// the top-level keys — clients (e.g. Claude Code) may place `model`
+// after a multi-MB `messages` array — then restores r.Body so the
+// upstream still receives the full, untouched stream. Only the small
+// routing fields are extracted; the prompt is never buffered for
+// capture, keeping memory bounded. Returns ok=false when the body isn't
+// a JSON object, the model field isn't found within maxScan, or on a
+// read error.
+func ScanRoutingFields(r *http.Request, maxScan int64) (model string, stream bool, ok bool) {
+	if r == nil || r.Body == nil || r.Body == http.NoBody || maxScan <= 0 {
+		return "", false, false
+	}
+	limited := io.LimitReader(r.Body, maxScan+1)
+	buf, readErr := io.ReadAll(limited)
+	if readErr != nil && !errors.Is(readErr, io.EOF) {
+		// Mid-stream read error (e.g. client disconnect): restore the bytes
+		// read so far plus the untouched tail and abort, rather than
+		// forwarding only the partial prefix as if it were the whole body.
+		r.Body = &replayReadCloser{replay: bytes.NewReader(append([]byte(nil), buf...)), tail: r.Body}
+		return "", false, false
+	}
+	if int64(len(buf)) > maxScan {
+		// Body exceeds the scan ceiling: restore the read prefix plus the
+		// untouched tail so the upstream still gets every byte.
+		r.Body = &replayReadCloser{replay: bytes.NewReader(append([]byte(nil), buf...)), tail: r.Body}
+	} else {
+		_ = r.Body.Close()
+		r.Body = io.NopCloser(bytes.NewReader(buf))
+	}
+	return scanTopLevelModelStream(buf)
+}
+
+// scanTopLevelModelStream walks the top level of a JSON object via a
+// streaming token reader, extracting the "model" string and "stream"
+// bool without materialising large values (each non-target value is
+// skipped as a RawMessage). Tolerant of truncation: returns whatever was
+// found before a malformed/short tail.
+func scanTopLevelModelStream(body []byte) (model string, stream bool, ok bool) {
+	dec := json.NewDecoder(bytes.NewReader(body))
+	tok, err := dec.Token()
+	if err != nil {
+		return "", false, false
+	}
+	if d, isDelim := tok.(json.Delim); !isDelim || d != '{' {
+		return "", false, false
+	}
+	for dec.More() {
+		keyTok, err := dec.Token()
+		if err != nil {
+			return model, stream, ok
+		}
+		key, _ := keyTok.(string)
+		switch key {
+		case "model":
+			var v string
+			if dec.Decode(&v) == nil {
+				model, ok = v, true
+			}
+		case "stream":
+			var v bool
+			if dec.Decode(&v) == nil {
+				stream = v
+			}
+		default:
+			// Skip the value by walking tokens instead of decoding it into
+			// a json.RawMessage — a multi-MB messages array would otherwise
+			// be materialised in full just to be discarded.
+			if err := skipValue(dec); err != nil {
+				return model, stream, ok
+			}
+		}
+	}
+	return model, stream, ok
+}
+
+// skipValue consumes one JSON value from dec without materialising it.
+// Scalars are a single token; objects/arrays are walked to their matching
+// close delimiter so nested structures are skipped in bounded memory.
+func skipValue(dec *json.Decoder) error {
+	tok, err := dec.Token()
+	if err != nil {
+		return err
+	}
+	d, isDelim := tok.(json.Delim)
+	if !isDelim || (d != '{' && d != '[') {
+		return nil
+	}
+	depth := 1
+	for depth > 0 {
+		tok, err := dec.Token()
+		if err != nil {
+			return err
+		}
+		if d, ok := tok.(json.Delim); ok {
+			switch d {
+			case '{', '[':
+				depth++
+			case '}', ']':
+				depth--
+			}
+		}
+	}
+	return nil
+}
+
+func contentTypeAllowed(ct string, allowed []string) bool {
+	if len(allowed) == 0 {
+		return false
+	}
+	media := ct
+	if idx := strings.Index(ct, ";"); idx >= 0 {
+		media = ct[:idx]
+	}
+	media = strings.TrimSpace(strings.ToLower(media))
+	for _, a := range allowed {
+		if strings.EqualFold(strings.TrimSpace(a), media) {
+			return true
+		}
+	}
+	return false
+}
+
+func parseContentLength(v string) int64 {
+	if v == "" {
+		return 0
+	}
+	parsed, err := strconv.ParseInt(v, 10, 64)
+	if err != nil || parsed < 0 {
+		return 0
+	}
+	return parsed
+}

proxy/internal/middleware/bodytap/response.go

@@ -0,0 +1,189 @@
+package bodytap
+
+import (
+	"bytes"
+	"net/http"
+	"sync"
+
+	"github.com/netbirdio/netbird/proxy/internal/responsewriter"
+)
+
+// CapturingResponseWriter wraps an http.ResponseWriter, forwards bytes
+// immediately to the client, and tees a bounded copy into an internal
+// buffer for middleware inspection. Streaming-aware in the sense that
+// every byte the upstream emits flows to the client without queuing
+// — the tee just sees a bounded prefix. SSE-aware parsing happens in
+// the response middleware against the buffered prefix; this writer
+// makes no attempt to demux event boundaries.
+//
+// Flusher and Hijacker are preserved via responsewriter.PassthroughWriter.
+type CapturingResponseWriter struct {
+	*responsewriter.PassthroughWriter
+	mu          sync.Mutex
+	buf         bytes.Buffer
+	cap         int64
+	status      int
+	statusSet   bool
+	written     int64
+	truncated   bool
+	stopped     bool
+	releaseBuf  func()
+	released    sync.Once
+	bypassed    bool
+	bypassReas  string
+	acquiredCap int64
+}
+
+// NewCapturingResponseWriter returns a writer that tees up to maxBytes
+// into a capped buffer while forwarding bytes to the underlying writer
+// immediately. When budget is non-nil the writer pre-acquires maxBytes
+// from it and the returned wrapper must be released by calling
+// Release() once the response is fully forwarded. If the budget cannot
+// be acquired the writer falls back to forwarding the response
+// unmodified, exposes Bypassed()=true with reason BypassBudget, and
+// releases nothing.
+func NewCapturingResponseWriter(w http.ResponseWriter, maxBytes int64, b Budget) *CapturingResponseWriter {
+	cw := &CapturingResponseWriter{
+		PassthroughWriter: responsewriter.New(w),
+		cap:               maxBytes,
+		status:            http.StatusOK,
+		releaseBuf:        func() {},
+	}
+	if maxBytes <= 0 {
+		// Capture disabled: mark stopped so Write never tees and never
+		// flags truncation (a zero cap means "don't capture", not
+		// "captured nothing").
+		cw.stopped = true
+		return cw
+	}
+	if b == nil {
+		return cw
+	}
+	if !b.Acquire(maxBytes) {
+		cw.bypassed = true
+		cw.bypassReas = BypassBudget
+		cw.cap = 0
+		cw.stopped = true
+		return cw
+	}
+	cw.acquiredCap = maxBytes
+	cw.releaseBuf = func() { b.Release(maxBytes) }
+	return cw
+}
+
+// Release returns the response capture budget acquired at construction
+// back to the shared pool. Idempotent. Safe to call from a defer
+// immediately after construction even when the writer ended up
+// bypassing the budget.
+func (c *CapturingResponseWriter) Release() {
+	if c == nil {
+		return
+	}
+	c.released.Do(func() {
+		if c.releaseBuf != nil {
+			c.releaseBuf()
+		}
+	})
+}
+
+// Bypassed reports whether the writer fell through to a no-tee
+// passthrough because the response capture budget could not be
+// acquired.
+func (c *CapturingResponseWriter) Bypassed() bool {
+	if c == nil {
+		return false
+	}
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.bypassed
+}
+
+// BypassReason returns the bypass code recorded by the budget check.
+// Empty when capture proceeded normally.
+func (c *CapturingResponseWriter) BypassReason() string {
+	if c == nil {
+		return ""
+	}
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.bypassReas
+}
+
+// WriteHeader records the status code and forwards it to the underlying
+// writer. Only the first call commits the status — matching HTTP semantics,
+// where superfluous WriteHeader calls (and any call after the body has
+// started) are ignored — so Status() reflects the code actually sent.
+func (c *CapturingResponseWriter) WriteHeader(status int) {
+	c.mu.Lock()
+	if c.statusSet {
+		c.mu.Unlock()
+		return
+	}
+	c.status = status
+	c.statusSet = true
+	c.mu.Unlock()
+	c.PassthroughWriter.WriteHeader(status)
+}
+
+// Write forwards p to the underlying writer unmodified and copies up
+// to the remaining buffer capacity into the tee buffer.
+func (c *CapturingResponseWriter) Write(p []byte) (int, error) {
+	n, err := c.PassthroughWriter.Write(p)
+	if n > 0 {
+		c.mu.Lock()
+		// The first byte commits the status (implicit 200 if WriteHeader was
+		// never called); a later WriteHeader must not change Status().
+		c.statusSet = true
+		c.written += int64(n)
+		if !c.stopped {
+			remaining := c.cap - int64(c.buf.Len())
+			if remaining <= 0 {
+				c.truncated = true
+				c.stopped = true
+			} else {
+				take := int64(n)
+				if take > remaining {
+					take = remaining
+					c.truncated = true
+					c.stopped = true
+				}
+				c.buf.Write(p[:take])
+			}
+		}
+		c.mu.Unlock()
+	}
+	return n, err
+}
+
+// Status returns the captured status code (defaults to 200 when
+// WriteHeader has not been called).
+func (c *CapturingResponseWriter) Status() int {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.status
+}
+
+// Body returns a copy of the buffered response prefix.
+func (c *CapturingResponseWriter) Body() []byte {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	out := make([]byte, c.buf.Len())
+	copy(out, c.buf.Bytes())
+	return out
+}
+
+// Truncated reports whether the buffered prefix stopped short of the
+// full response stream.
+func (c *CapturingResponseWriter) Truncated() bool {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.truncated
+}
+
+// BytesWritten returns the total number of bytes forwarded to the
+// underlying writer.
+func (c *CapturingResponseWriter) BytesWritten() int64 {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.written
+}

proxy/internal/middleware/bodytap/routing_scan_test.go

@@ -0,0 +1,86 @@
+package bodytap
+
+import (
+	"fmt"
+	"io"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// makeBigAnthropicBody builds a request body shaped like Claude Code's:
+// a multi-MB "messages" array with the routing fields (model, stream)
+// placed AFTER it, which is the ordering that defeats a prefix-only
+// capture.
+func makeBigAnthropicBody(t *testing.T, model string, stream bool, messagesBytes int) string {
+	t.Helper()
+	filler := strings.Repeat("x", messagesBytes)
+	return fmt.Sprintf(
+		`{"max_tokens":64000,"messages":[{"role":"user","content":%q}],"model":%q,"stream":%t}`,
+		filler, model, stream,
+	)
+}
+
+func TestScanRoutingFields_ModelAfterLargeMessages(t *testing.T) {
+	body := makeBigAnthropicBody(t, "claude-opus-4-8", true, 3<<20) // 3 MiB messages
+	req := httptest.NewRequest("POST", "https://x/v1/messages", strings.NewReader(body))
+
+	model, stream, ok := ScanRoutingFields(req, MaxRoutingScanBytes)
+	require.True(t, ok, "model must be recovered even when it follows a multi-MB messages array")
+	assert.Equal(t, "claude-opus-4-8", model, "model field must be extracted")
+	assert.True(t, stream, "stream field must be extracted")
+
+	// Body must be fully restored for the upstream.
+	got, err := io.ReadAll(req.Body)
+	require.NoError(t, err)
+	assert.Equal(t, body, string(got), "the full request body must be replayed to upstream after scanning")
+}
+
+func TestScanRoutingFields_SmallBody(t *testing.T) {
+	body := `{"model":"claude-opus-4-8","stream":false,"messages":[]}`
+	req := httptest.NewRequest("POST", "https://x/v1/messages", strings.NewReader(body))
+
+	model, stream, ok := ScanRoutingFields(req, MaxRoutingScanBytes)
+	require.True(t, ok)
+	assert.Equal(t, "claude-opus-4-8", model)
+	assert.False(t, stream)
+
+	got, _ := io.ReadAll(req.Body)
+	assert.Equal(t, body, string(got), "small bodies must also be restored intact")
+}
+
+func TestScanRoutingFields_NoModel(t *testing.T) {
+	body := `{"stream":true,"messages":[]}`
+	req := httptest.NewRequest("POST", "https://x/v1/messages", strings.NewReader(body))
+
+	_, _, ok := ScanRoutingFields(req, MaxRoutingScanBytes)
+	assert.False(t, ok, "ok must be false when no model field is present")
+
+	got, _ := io.ReadAll(req.Body)
+	assert.Equal(t, body, string(got), "body must be restored even when model is absent")
+}
+
+func TestScanRoutingFields_NotJSON(t *testing.T) {
+	body := "this is not json at all"
+	req := httptest.NewRequest("POST", "https://x/v1/messages", strings.NewReader(body))
+
+	_, _, ok := ScanRoutingFields(req, MaxRoutingScanBytes)
+	assert.False(t, ok, "ok must be false for a non-JSON body")
+}
+
+func TestScanRoutingFields_ModelBeyondScanCeiling(t *testing.T) {
+	// model sits after 4 MiB of messages but the scan ceiling is 1 MiB:
+	// model can't be recovered, yet the full body must still replay.
+	body := makeBigAnthropicBody(t, "claude-opus-4-8", true, 4<<20)
+	req := httptest.NewRequest("POST", "https://x/v1/messages", strings.NewReader(body))
+
+	_, _, ok := ScanRoutingFields(req, 1<<20)
+	assert.False(t, ok, "model beyond the scan ceiling is not recoverable")
+
+	got, err := io.ReadAll(req.Body)
+	require.NoError(t, err)
+	assert.Equal(t, body, string(got), "the full body must still replay to upstream even when the scan gives up")
+}

proxy/internal/middleware/builtin/agentnetwork_chain_integration_test.go

@@ -0,0 +1,318 @@
+package builtin_test
+
+import (
+	"context"
+	"net"
+	"runtime"
+	"strconv"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/test/bufconn"
+
+	"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
+	agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
+	mgmtgrpc "github.com/netbirdio/netbird/management/internals/shared/grpc"
+	"github.com/netbirdio/netbird/management/server/store"
+	nbtypes "github.com/netbirdio/netbird/management/server/types"
+	"github.com/netbirdio/netbird/proxy/internal/middleware"
+	"github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_limit_check"
+	"github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_limit_record"
+	"github.com/netbirdio/netbird/shared/management/proto"
+)
+
+// chainIntegrationFixture wires the BOTH new agent-network
+// middlewares against a live in-process management stack: real
+// sqlite store + real Manager + real gRPC server. The proxy chain
+// framework itself isn't constructed (its dispatcher / accumulator /
+// metadata gate are tested separately); we exercise the middleware
+// pair as the proxy runtime would, by invoking each with a crafted
+// Input and asserting the wire path between them.
+//
+// This is the regression cover for item 16 in the design review:
+// real LLM request → cost stamped → consumption row in the table.
+type chainIntegrationFixture struct {
+	store    store.Store
+	manager  agentnetwork.Manager
+	gatecase *llm_limit_check.Middleware
+	recorder *llm_limit_record.Middleware
+}
+
+func newChainIntegration(t *testing.T) *chainIntegrationFixture {
+	t.Helper()
+	if runtime.GOOS == "windows" {
+		t.Skip("sqlite store not properly supported on Windows yet")
+	}
+	t.Setenv("NETBIRD_STORE_ENGINE", string(nbtypes.SqliteStoreEngine))
+
+	st, cleanUp, err := store.NewTestStoreFromSQL(context.Background(), "", t.TempDir())
+	require.NoError(t, err)
+	t.Cleanup(cleanUp)
+
+	manager := agentnetwork.NewManager(st, nil, nil, nil)
+
+	server := &mgmtgrpc.ProxyServiceServer{}
+	server.SetAgentNetworkLimitsService(manager)
+
+	const bufSize = 1024 * 1024
+	lis := bufconn.Listen(bufSize)
+	srv := grpc.NewServer()
+	proto.RegisterProxyServiceServer(srv, server)
+	go func() { _ = srv.Serve(lis) }()
+	t.Cleanup(srv.Stop)
+
+	conn, err := grpc.NewClient("passthrough:///bufnet",
+		grpc.WithContextDialer(func(_ context.Context, _ string) (net.Conn, error) { return lis.Dial() }),
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+	)
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = conn.Close() })
+
+	mgmtClient := proto.NewProxyServiceClient(conn)
+	return &chainIntegrationFixture{
+		store:    st,
+		manager:  manager,
+		gatecase: llm_limit_check.New(mgmtClient, nil),
+		recorder: llm_limit_record.New(mgmtClient, nil),
+	}
+}
+
+// chainInput builds a middleware Input that mirrors what the proxy
+// framework would synthesise for a tunnel-peer LLM request. The
+// gate consumes the resolved provider id from upstream metadata
+// (set by llm_router); the recorder consumes the attribution
+// metadata stamped by the gate plus tokens / cost from
+// llm_response_parser + cost_meter.
+func chainInput(account, user, group, providerID string, requestMeta []middleware.KV) *middleware.Input {
+	_ = providerID // packed into requestMeta by the caller as KeyLLMResolvedProviderID
+	return &middleware.Input{
+		AccountID:  account,
+		UserID:     user,
+		UserGroups: []string{group},
+		Metadata:   requestMeta,
+	}
+}
+
+// chainCapPolicy builds a tight token-cap policy fixture for the
+// chain integration tests. Inlined here (rather than imported) because
+// the equivalent helper in the management gRPC package is unexported
+// and this is a different package boundary.
+func chainCapPolicy(id, account string, sourceGroups []string, providerID string, tokenCap, windowSec int64) *agentNetworkTypes.Policy {
+	return &agentNetworkTypes.Policy{
+		ID:                     id,
+		AccountID:              account,
+		Enabled:                true,
+		Name:                   id,
+		SourceGroups:           sourceGroups,
+		DestinationProviderIDs: []string{providerID},
+		Limits: agentNetworkTypes.PolicyLimits{
+			TokenLimit: agentNetworkTypes.PolicyTokenLimit{
+				Enabled:       true,
+				GroupCap:      tokenCap,
+				WindowSeconds: windowSec,
+			},
+		},
+		CreatedAt: time.Now().UTC(),
+		UpdatedAt: time.Now().UTC(),
+	}
+}
+
+// TestChain_AllowPath_StampsAttributionAndRecordsCounter walks the
+// full happy path: gate calls CheckLLMPolicyLimits → stamps
+// attribution metadata → recorder reads metadata + tokens / cost →
+// calls RecordLLMUsage → counters land in sqlite. Asserting on the
+// store at the end proves every leg of the wire works together,
+// not just each leg in isolation (which the unit tests already cover).
+func TestChain_AllowPath_StampsAttributionAndRecordsCounter(t *testing.T) {
+	f := newChainIntegration(t)
+
+	const account = "acc-1"
+	const user = "user-bob"
+	const group = "grp-engineers"
+	const provider = "prov-1"
+
+	// Seed a policy with token + budget caps; both halves carry
+	// real ceilings so the request stays within headroom.
+	require.NoError(t, f.store.SaveAgentNetworkPolicy(context.Background(),
+		chainCapPolicy("pol-1", account, []string{group}, provider, 10_000, 86_400)))
+
+	// ── Stage 1 — gate: pre-flight check ──────────────────────
+	gateIn := chainInput(account, user, group, provider, []middleware.KV{
+		{Key: middleware.KeyLLMResolvedProviderID, Value: provider},
+		{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
+	})
+	gateOut, err := f.gatecase.Invoke(context.Background(), gateIn)
+	require.NoError(t, err)
+	require.Equal(t, middleware.DecisionAllow, gateOut.Decision, "fresh policy must allow")
+
+	// Verify attribution metadata was stamped — the recorder
+	// depends on these keys.
+	metaMap := map[string]string{}
+	for _, kv := range gateOut.Metadata {
+		metaMap[kv.Key] = kv.Value
+	}
+	assert.Equal(t, "pol-1", metaMap[middleware.KeyLLMSelectedPolicyID])
+	assert.Equal(t, group, metaMap[middleware.KeyLLMAttributionGroupID])
+	assert.Equal(t, "86400", metaMap[middleware.KeyLLMAttributionWindowS])
+
+	// ── Stage 2 — recorder: post-flight write ─────────────────
+	// Build the response-leg Input the framework would synthesise
+	// for the recorder: gate's emitted attribution metadata + the
+	// tokens / cost stamped by llm_response_parser + cost_meter.
+	const tokensIn = int64(123)
+	const tokensOut = int64(45)
+	const costUSD = 0.0042
+	recordIn := chainInput(account, user, group, provider, append([]middleware.KV{},
+		gateOut.Metadata...))
+	recordIn.Metadata = append(recordIn.Metadata,
+		middleware.KV{Key: middleware.KeyLLMInputTokens, Value: strconv.FormatInt(tokensIn, 10)},
+		middleware.KV{Key: middleware.KeyLLMOutputTokens, Value: strconv.FormatInt(tokensOut, 10)},
+		middleware.KV{Key: middleware.KeyCostUSDTotal, Value: strconv.FormatFloat(costUSD, 'f', 6, 64)},
+	)
+	recordOut, err := f.recorder.Invoke(context.Background(), recordIn)
+	require.NoError(t, err)
+	assert.Equal(t, middleware.DecisionAllow, recordOut.Decision, "recorder always allows; its only side effect is the counter write")
+
+	// ── Stage 3 — assert state in sqlite ──────────────────────
+	windowStart := agentNetworkTypes.WindowStart(time.Now(), 86_400)
+	userRow, err := f.store.GetAgentNetworkConsumption(
+		context.Background(), store.LockingStrengthNone, account,
+		agentNetworkTypes.DimensionUser, user, int64(86_400), windowStart,
+	)
+	require.NoError(t, err)
+	assert.Equal(t, tokensIn, userRow.TokensInput, "user counter must hold the input tokens the recorder posted")
+	assert.Equal(t, tokensOut, userRow.TokensOutput)
+	assert.InDelta(t, costUSD, userRow.CostUSD, 1e-6)
+
+	groupRow, err := f.store.GetAgentNetworkConsumption(
+		context.Background(), store.LockingStrengthNone, account,
+		agentNetworkTypes.DimensionGroup, group, int64(86_400), windowStart,
+	)
+	require.NoError(t, err)
+	assert.Equal(t, tokensIn, groupRow.TokensInput, "group counter mirrors the user counter — single Record posts both dims")
+}
+
+// TestChain_DenyPath_GateRejectsAndNoConsumptionWritten covers the
+// negative side: when the gate denies, the recorder is never
+// invoked (the proxy framework short-circuits on Decision=Deny).
+// We assert no consumption row materialises after the gate-deny
+// path, even though the test technically calls the recorder
+// afterwards — the recorder must skip on missing attribution
+// metadata so the framework's short-circuit isn't load-bearing for
+// data integrity.
+func TestChain_DenyPath_GateRejectsAndNoConsumptionWritten(t *testing.T) {
+	f := newChainIntegration(t)
+
+	const account = "acc-1"
+	const user = "user-bob"
+	const group = "grp-tight"
+	const provider = "prov-1"
+
+	policy := chainCapPolicy("pol-tight", account, []string{group}, provider, 100, 86_400)
+	require.NoError(t, f.store.SaveAgentNetworkPolicy(context.Background(), policy))
+
+	// Pre-burn the counter to the cap so the gate denies.
+	require.NoError(t, f.store.IncrementAgentNetworkConsumption(
+		context.Background(), account,
+		agentNetworkTypes.DimensionGroup, group, int64(86_400),
+		agentNetworkTypes.WindowStart(time.Now(), 86_400),
+		100, 0, 0,
+	))
+
+	gateOut, err := f.gatecase.Invoke(context.Background(), chainInput(account, user, group, provider,
+		[]middleware.KV{
+			{Key: middleware.KeyLLMResolvedProviderID, Value: provider},
+			{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
+		},
+	))
+	require.NoError(t, err)
+	assert.Equal(t, middleware.DecisionDeny, gateOut.Decision, "policy at-cap must deny on the gate")
+	require.NotNil(t, gateOut.DenyReason)
+	assert.Equal(t, "llm_policy.token_cap_exceeded", gateOut.DenyReason.Code)
+
+	// On deny, the gate emits no attribution metadata. If the
+	// proxy framework still invokes the recorder (defense in
+	// depth), the recorder's "no attribution window = skip" guard
+	// prevents a phantom counter increment.
+	recordOut, err := f.recorder.Invoke(context.Background(), chainInput(account, user, group, provider,
+		gateOut.Metadata, // no llm.attribution_window_seconds stamped
+	))
+	require.NoError(t, err)
+	assert.Equal(t, middleware.DecisionAllow, recordOut.Decision)
+
+	// The pre-burned 100 tokens are the only counter movement —
+	// the recorder must NOT have added a fresh row for the user
+	// dimension on this denied request.
+	windowStart := agentNetworkTypes.WindowStart(time.Now(), 86_400)
+	userRow, err := f.store.GetAgentNetworkConsumption(
+		context.Background(), store.LockingStrengthNone, account,
+		agentNetworkTypes.DimensionUser, user, int64(86_400), windowStart,
+	)
+	require.NoError(t, err)
+	assert.Zero(t, userRow.TokensInput, "user dimension must not gain tokens from a denied request — recorder skip is the safety net")
+}
+
+// TestChain_CapExhaustTransition exercises the allow→deny boundary
+// the operator cares most about: a request just under cap allows
+// AND records, the next request post-record at-cap denies. This is
+// the same lifecycle 50-grpc-allow-record-deny.sh runs in bash, but
+// against the actual middleware pair rather than the smoke binary
+// driving the gRPC RPCs directly.
+func TestChain_CapExhaustTransition(t *testing.T) {
+	f := newChainIntegration(t)
+
+	const account = "acc-1"
+	const user = "user-alice"
+	const group = "grp-cap-edge"
+	const provider = "prov-1"
+	const tightCap = int64(100)
+
+	require.NoError(t, f.store.SaveAgentNetworkPolicy(context.Background(),
+		chainCapPolicy("pol-edge", account, []string{group}, provider, tightCap, 86_400)))
+
+	// Pre-burn 99 tokens so we're at the very edge.
+	require.NoError(t, f.store.IncrementAgentNetworkConsumption(
+		context.Background(), account,
+		agentNetworkTypes.DimensionGroup, group, int64(86_400),
+		agentNetworkTypes.WindowStart(time.Now(), 86_400),
+		99, 0, 0,
+	))
+
+	// Gate at 99/100 — must allow (one token of headroom).
+	gateOut, err := f.gatecase.Invoke(context.Background(), chainInput(account, user, group, provider,
+		[]middleware.KV{
+			{Key: middleware.KeyLLMResolvedProviderID, Value: provider},
+			{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
+		},
+	))
+	require.NoError(t, err)
+	require.Equal(t, middleware.DecisionAllow, gateOut.Decision, "99/100 must allow — one token of headroom")
+
+	// Record one more input token — pushes us to 100/100.
+	recordIn := chainInput(account, user, group, provider, append([]middleware.KV{},
+		gateOut.Metadata...))
+	recordIn.Metadata = append(recordIn.Metadata,
+		middleware.KV{Key: middleware.KeyLLMInputTokens, Value: "1"},
+		middleware.KV{Key: middleware.KeyLLMOutputTokens, Value: "0"},
+		middleware.KV{Key: middleware.KeyCostUSDTotal, Value: "0.000001"},
+	)
+	_, err = f.recorder.Invoke(context.Background(), recordIn)
+	require.NoError(t, err)
+
+	// Next gate call must deny — counter is exactly at cap.
+	gateOut2, err := f.gatecase.Invoke(context.Background(), chainInput(account, user, group, provider,
+		[]middleware.KV{
+			{Key: middleware.KeyLLMResolvedProviderID, Value: provider},
+			{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
+		},
+	))
+	require.NoError(t, err)
+	assert.Equal(t, middleware.DecisionDeny, gateOut2.Decision,
+		"once recorder pushed the group counter to 100/100, the next gate call must deny — allow→deny transition is the operator-visible product semantic")
+	require.NotNil(t, gateOut2.DenyReason)
+	assert.Equal(t, "llm_policy.token_cap_exceeded", gateOut2.DenyReason.Code)
+}

proxy/internal/middleware/builtin/all_test.go

@@ -0,0 +1,40 @@
+package builtin_test
+
+import (
+	"sort"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	mwbuiltin "github.com/netbirdio/netbird/proxy/internal/middleware/builtin"
+
+	_ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/cost_meter"
+	_ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_guardrail"
+	_ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_identity_inject"
+	_ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_limit_check"
+	_ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_limit_record"
+	_ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_request_parser"
+	_ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_response_parser"
+	_ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_router"
+)
+
+// TestDefaultRegistry_BuiltinIDs locks the set of middleware IDs that
+// the default builtin registry exposes once every sub-package's init()
+// has run. The list is the source of truth wired by the synthesiser
+// in management; adding a new built-in middleware should consciously
+// extend this list.
+func TestDefaultRegistry_BuiltinIDs(t *testing.T) {
+	got := mwbuiltin.DefaultRegistry().IDs()
+	sort.Strings(got)
+	want := []string{
+		"cost_meter",
+		"llm_guardrail",
+		"llm_identity_inject",
+		"llm_limit_check",
+		"llm_limit_record",
+		"llm_request_parser",
+		"llm_response_parser",
+		"llm_router",
+	}
+	assert.Equal(t, want, got, "default registry must expose every built-in middleware after anonymous imports")
+}

proxy/internal/middleware/builtin/builtin.go

@@ -0,0 +1,93 @@
+// Package builtin holds the package-level middleware registry that
+// concrete middleware packages register themselves into via init().
+// Server boot anonymous-imports each middleware sub-package; the
+// resolver attached to the middleware Manager pulls factories out of
+// this registry.
+package builtin
+
+import (
+	"context"
+	"sync"
+
+	log "github.com/sirupsen/logrus"
+	"go.opentelemetry.io/otel/metric"
+	"google.golang.org/grpc"
+
+	"github.com/netbirdio/netbird/proxy/internal/middleware"
+	"github.com/netbirdio/netbird/shared/management/proto"
+)
+
+// MgmtClient is the narrow slice of proto.ProxyServiceClient that
+// builtin middlewares may use during request / response handling.
+// Only the agent-network limit pair (llm_limit_check + llm_limit_record)
+// uses this today; declaring the surface here keeps the dependency
+// explicit at boot time.
+//
+// proto.ProxyServiceClient already satisfies this interface so server
+// boot just forwards its existing client.
+type MgmtClient interface {
+	CheckLLMPolicyLimits(ctx context.Context, in *proto.CheckLLMPolicyLimitsRequest, opts ...grpc.CallOption) (*proto.CheckLLMPolicyLimitsResponse, error)
+	RecordLLMUsage(ctx context.Context, in *proto.RecordLLMUsageRequest, opts ...grpc.CallOption) (*proto.RecordLLMUsageResponse, error)
+}
+
+// defaultRegistry is the package-level registry that concrete builtin
+// middlewares register themselves into via init().
+var defaultRegistry = middleware.NewRegistry()
+
+// FactoryContext is the per-process bag that concrete factories may
+// consult during construction. It carries the proxy-lifetime context,
+// the data directory used for static config files (pricing tables,
+// allowlists), the OTel meter, and the proxy logger.
+//
+// Configure must be called once at boot before any chain build calls
+// Resolve. Calling it twice overwrites the prior value; tests may rely
+// on this to reset state.
+type FactoryContext struct {
+	Context    context.Context
+	DataDir    string
+	Meter      metric.Meter
+	Logger     *log.Logger
+	MgmtClient MgmtClient
+}
+
+var (
+	ctxStore FactoryContext
+	ctxMu    sync.RWMutex
+)
+
+// Configure stores the per-process FactoryContext. Concrete factories
+// reach for it via Context(). mgmt may be nil on tests / standalone
+// builds with no management server; consumers must guard.
+func Configure(ctx context.Context, dataDir string, meter metric.Meter, logger *log.Logger, mgmt MgmtClient) {
+	ctxMu.Lock()
+	defer ctxMu.Unlock()
+	ctxStore = FactoryContext{
+		Context:    ctx,
+		DataDir:    dataDir,
+		Meter:      meter,
+		Logger:     logger,
+		MgmtClient: mgmt,
+	}
+}
+
+// Context returns the stored FactoryContext. Returns a zero value when
+// Configure was never called; consumers must guard against nil
+// Context/Meter/Logger if they care.
+func Context() FactoryContext {
+	ctxMu.RLock()
+	defer ctxMu.RUnlock()
+	return ctxStore
+}
+
+// Register adds a factory to the default registry. Called from init()
+// blocks of concrete middleware packages. Panics on collision so
+// duplicate IDs surface at startup.
+func Register(f middleware.Factory) {
+	defaultRegistry.MustRegister(f)
+}
+
+// DefaultRegistry returns the shared registry. The proxy server
+// constructs the Resolver from it at boot.
+func DefaultRegistry() *middleware.Registry {
+	return defaultRegistry
+}

proxy/internal/middleware/builtin/cost_calculation_matrix_test.go

@@ -0,0 +1,329 @@
+package builtin_test
+
+import (
+	"bytes"
+	"context"
+	"encoding/base64"
+	"encoding/json"
+	"strconv"
+	"testing"
+
+	"github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/proxy/internal/middleware"
+	"github.com/netbirdio/netbird/proxy/internal/middleware/builtin"
+	"github.com/netbirdio/netbird/proxy/internal/middleware/builtin/cost_meter"
+	"github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_request_parser"
+	"github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_response_parser"
+)
+
+// Drives the real pipeline (llm_request_parser → llm_response_parser → cost_meter) on the embedded default pricing
+// table and asserts exact USD amounts hardcoded from the vendors' published prices, including the cache split.
+func TestCostCalculation_ProviderMatrix(t *testing.T) {
+	// Empty data dir → embedded defaults, like a proxy with no pricing override.
+	builtin.Configure(context.Background(), t.TempDir(), nil, nil, nil)
+
+	reqMW, err := llm_request_parser.Factory{}.New(nil)
+	require.NoError(t, err, "build llm_request_parser")
+	respMW, err := llm_response_parser.Factory{}.New(nil)
+	require.NoError(t, err, "build llm_response_parser")
+	costMW, err := cost_meter.Factory{}.New(nil)
+	require.NoError(t, err, "build cost_meter")
+	t.Cleanup(func() { _ = costMW.Close() })
+
+	const jsonCT = "application/json"
+	const sseCT = "text/event-stream"
+	const awsCT = "application/vnd.amazon.eventstream"
+
+	cases := []struct {
+		name     string
+		url      string
+		reqBody  []byte
+		respCT   string
+		respBody []byte
+
+		wantProvider  string
+		wantModel     string
+		wantCost      float64 // exact expected USD; ignored when wantSkip is set
+		wantCacheCost float64 // expected cost.usd_cache portion of wantCost
+		wantSkip      string  // expected cost.skipped reason, "" when priced
+	}{
+		{
+			// gpt-4o-mini $0.15/$0.60 per MTok: 1000×0.15/1M + 500×0.60/1M.
+			name:         "openai chat completions",
+			url:          "https://api.openai.com/v1/chat/completions",
+			reqBody:      []byte(`{"model":"gpt-4o-mini","messages":[{"role":"user","content":"hi"}]}`),
+			respCT:       jsonCT,
+			respBody:     []byte(`{"choices":[{"message":{"content":"pong"}}],"usage":{"prompt_tokens":1000,"completion_tokens":500,"total_tokens":1500}}`),
+			wantProvider: "openai",
+			wantModel:    "gpt-4o-mini",
+			wantCost:     0.00045,
+		},
+		{
+			// OpenAI cached tokens are a SUBSET of prompt_tokens at a discount; gpt-4o $2.50/$10 per MTok, cached $1.25/M:
+			// 250×2.5/1M + 750×1.25/1M + 500×10/1M.
+			name:          "openai cached subset discount",
+			url:           "https://api.openai.com/v1/chat/completions",
+			reqBody:       []byte(`{"model":"gpt-4o","messages":[{"role":"user","content":"hi"}]}`),
+			respCT:        jsonCT,
+			respBody:      []byte(`{"usage":{"prompt_tokens":1000,"completion_tokens":500,"prompt_tokens_details":{"cached_tokens":750}}}`),
+			wantProvider:  "openai",
+			wantModel:     "gpt-4o",
+			wantCost:      0.0065625,
+			wantCacheCost: 0.0009375,
+		},
+		{
+			// OpenAI streaming: usage rides the final SSE frame.
+			name:         "openai chat SSE stream",
+			url:          "https://api.openai.com/v1/chat/completions",
+			reqBody:      []byte(`{"model":"gpt-4o-mini","stream":true,"messages":[{"role":"user","content":"hi"}]}`),
+			respCT:       sseCT,
+			respBody:     sseBody(`{"choices":[{"delta":{"content":"po"}}]}`, `{"choices":[{"delta":{"content":"ng"}}]}`, `{"choices":[],"usage":{"prompt_tokens":1000,"completion_tokens":500}}`, "[DONE]"),
+			wantProvider: "openai",
+			wantModel:    "gpt-4o-mini",
+			wantCost:     0.00045,
+		},
+		{
+			// Mistral speaks the OpenAI shape: mistral-large-latest $0.50/$1.50 per MTok.
+			name:         "mistral via openai shape",
+			url:          "https://api.mistral.ai/v1/chat/completions",
+			reqBody:      []byte(`{"model":"mistral-large-latest","messages":[{"role":"user","content":"hi"}]}`),
+			respCT:       jsonCT,
+			respBody:     []byte(`{"usage":{"prompt_tokens":1000,"completion_tokens":1000}}`),
+			wantProvider: "openai",
+			wantModel:    "mistral-large-latest",
+			wantCost:     0.002,
+		},
+		{
+			// The field report, minus caching: Bedrock Sonnet 4.6 $3/$15 per MTok, 3×3/1M + 1514×15/1M = $0.022719.
+			// Also covers inference-profile normalization of the region-prefixed versioned id in the URL.
+			name:         "bedrock invoke — reported scenario, no cache",
+			url:          "https://bedrock-runtime.eu-central-1.amazonaws.com/model/global.anthropic.claude-sonnet-4-6-20260115-v1:0/invoke",
+			reqBody:      []byte(`{"messages":[{"role":"user","content":"hi"}]}`),
+			respCT:       jsonCT,
+			respBody:     []byte(`{"content":[{"type":"text","text":"pong"}],"usage":{"input_tokens":3,"output_tokens":1514}}`),
+			wantProvider: "bedrock",
+			wantModel:    "anthropic.claude-sonnet-4-6",
+			wantCost:     0.022719,
+		},
+		{
+			// The field report as observed: the FIRST call of a session also wrote a 30,528-token prompt cache at
+			// 1.25× input ($3.75/M): 0.022719 + 30528×3.75/1M = $0.137199 — the reported $0.1372.
+			name:          "bedrock invoke — reported scenario with cache write",
+			url:           "https://bedrock-runtime.eu-central-1.amazonaws.com/model/global.anthropic.claude-sonnet-4-6-20260115-v1:0/invoke",
+			reqBody:       []byte(`{"messages":[{"role":"user","content":"hi"}]}`),
+			respCT:        jsonCT,
+			respBody:      []byte(`{"usage":{"input_tokens":3,"output_tokens":1514,"cache_creation_input_tokens":30528,"cache_read_input_tokens":0}}`),
+			wantProvider:  "bedrock",
+			wantModel:     "anthropic.claude-sonnet-4-6",
+			wantCost:      0.137199,
+			wantCacheCost: 0.11448,
+		},
+		{
+			// Same numbers over the InvokeModel event-stream: message_start carries input + cache, message_delta the output.
+			name:          "bedrock invoke stream with cache write",
+			url:           "https://bedrock-runtime.eu-central-1.amazonaws.com/model/global.anthropic.claude-sonnet-4-6-20260115-v1:0/invoke-with-response-stream",
+			reqBody:       []byte(`{"messages":[{"role":"user","content":"hi"}]}`),
+			respCT:        awsCT,
+			respBody:      bedrockInvokeStream(t, `{"type":"message_start","message":{"usage":{"input_tokens":3,"output_tokens":1,"cache_creation_input_tokens":30528}}}`, `{"type":"content_block_delta","delta":{"type":"text_delta","text":"pong"}}`, `{"type":"message_delta","usage":{"output_tokens":1514}}`),
+			wantProvider:  "bedrock",
+			wantModel:     "anthropic.claude-sonnet-4-6",
+			wantCost:      0.137199,
+			wantCacheCost: 0.11448,
+		},
+		{
+			// Converse camelCase usage incl. cache buckets. Haiku 4.5 $1/$5 per MTok, read $0.10/M, write $1.25/M:
+			// 50×1/1M + 100×5/1M + 2000×0.1/1M + 1000×1.25/1M = $0.002.
+			name:          "bedrock converse with cache buckets",
+			url:           "https://bedrock-runtime.eu-central-1.amazonaws.com/model/eu.anthropic.claude-haiku-4-5-20251001-v1:0/converse",
+			reqBody:       []byte(`{"messages":[{"role":"user","content":[{"text":"hi"}]}]}`),
+			respCT:        jsonCT,
+			respBody:      []byte(`{"output":{"message":{"content":[{"text":"pong"}]}},"usage":{"inputTokens":50,"outputTokens":100,"totalTokens":3150,"cacheReadInputTokens":2000,"cacheWriteInputTokens":1000}}`),
+			wantProvider:  "bedrock",
+			wantModel:     "anthropic.claude-haiku-4-5",
+			wantCost:      0.002,
+			wantCacheCost: 0.00145,
+		},
+		{
+			// Same numbers over converse-stream: usage rides the trailing metadata frame.
+			name:    "bedrock converse stream with cache buckets",
+			url:     "https://bedrock-runtime.eu-central-1.amazonaws.com/model/eu.anthropic.claude-haiku-4-5-20251001-v1:0/converse-stream",
+			reqBody: []byte(`{"messages":[{"role":"user","content":[{"text":"hi"}]}]}`),
+			respCT:  awsCT,
+			respBody: bedrockConverseStream(t,
+				`{"delta":{"text":"pong"}}`,
+				`{"usage":{"inputTokens":50,"outputTokens":100,"totalTokens":3150,"cacheReadInputTokens":2000,"cacheWriteInputTokens":1000}}`,
+			),
+			wantProvider:  "bedrock",
+			wantModel:     "anthropic.claude-haiku-4-5",
+			wantCost:      0.002,
+			wantCacheCost: 0.00145,
+		},
+		{
+			// First-party Anthropic, additive cache buckets. Sonnet 4.6:
+			// 256×3/1M + 200×15/1M + 768×0.3/1M + 512×3.75/1M.
+			name:          "anthropic messages with cache buckets",
+			url:           "https://api.anthropic.com/v1/messages",
+			reqBody:       []byte(`{"model":"claude-sonnet-4-6","messages":[{"role":"user","content":"hi"}]}`),
+			respCT:        jsonCT,
+			respBody:      []byte(`{"content":[{"type":"text","text":"pong"}],"usage":{"input_tokens":256,"output_tokens":200,"cache_read_input_tokens":768,"cache_creation_input_tokens":512}}`),
+			wantProvider:  "anthropic",
+			wantModel:     "claude-sonnet-4-6",
+			wantCost:      0.0059184,
+			wantCacheCost: 0.0021504,
+		},
+		{
+			// Anthropic SSE: input from message_start, output from message_delta. Haiku 4.5: 1000×1/1M + 2000×5/1M.
+			name:         "anthropic SSE stream",
+			url:          "https://api.anthropic.com/v1/messages",
+			reqBody:      []byte(`{"model":"claude-haiku-4-5","stream":true,"messages":[{"role":"user","content":"hi"}]}`),
+			respCT:       sseCT,
+			respBody:     sseBody(`{"type":"message_start","message":{"usage":{"input_tokens":1000,"output_tokens":2}}}`, `{"type":"content_block_delta","delta":{"type":"text_delta","text":"pong"}}`, `{"type":"message_delta","usage":{"output_tokens":2000}}`, `{"type":"message_stop"}`),
+			wantProvider: "anthropic",
+			wantModel:    "claude-haiku-4-5",
+			wantCost:     0.011,
+		},
+		{
+			// Kimi's Anthropic-compatible endpoint: kimi-k3 $3/$15 per MTok under the anthropic table.
+			name:         "kimi anthropic shape",
+			url:          "https://api.moonshot.ai/anthropic/v1/messages",
+			reqBody:      []byte(`{"model":"kimi-k3","messages":[{"role":"user","content":"hi"}]}`),
+			respCT:       jsonCT,
+			respBody:     []byte(`{"usage":{"input_tokens":1000,"output_tokens":1000}}`),
+			wantProvider: "anthropic",
+			wantModel:    "kimi-k3",
+			wantCost:     0.018,
+		},
+		{
+			// Vertex path-routed model with "@version" stripped; Anthropic-on-Vertex priced under the anthropic table.
+			name:         "vertex anthropic path-routed",
+			url:          "https://aiplatform.googleapis.com/v1/projects/p/locations/global/publishers/anthropic/models/claude-sonnet-4-6@20260115:rawPredict",
+			reqBody:      []byte(`{"anthropic_version":"vertex-2023-10-16","messages":[{"role":"user","content":"hi"}]}`),
+			respCT:       jsonCT,
+			respBody:     []byte(`{"usage":{"input_tokens":200,"output_tokens":100}}`),
+			wantProvider: "anthropic",
+			wantModel:    "claude-sonnet-4-6",
+			wantCost:     0.0021,
+		},
+		{
+			// Gateway-prefixed model ids are not in the pricing table: the meter must SKIP, never guess a rate.
+			name:         "gateway-prefixed model skips pricing",
+			url:          "https://gateway.example.com/v1/chat/completions",
+			reqBody:      []byte(`{"model":"openai/gpt-4o-mini","messages":[{"role":"user","content":"hi"}]}`),
+			respCT:       jsonCT,
+			respBody:     []byte(`{"usage":{"prompt_tokens":1000,"completion_tokens":500}}`),
+			wantProvider: "openai",
+			wantModel:    "openai/gpt-4o-mini",
+			wantSkip:     "unknown_model",
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			in := &middleware.Input{
+				Method:  "POST",
+				URL:     tc.url,
+				Headers: []middleware.KV{{Key: "Content-Type", Value: "application/json"}},
+				Body:    tc.reqBody,
+			}
+
+			reqOut, err := reqMW.Invoke(context.Background(), in)
+			require.NoError(t, err, "request parser")
+			in.Metadata = append(in.Metadata, reqOut.Metadata...)
+
+			require.Equal(t, tc.wantProvider, metaKV(in.Metadata, middleware.KeyLLMProvider), "detected provider")
+			require.Equal(t, tc.wantModel, metaKV(in.Metadata, middleware.KeyLLMModel), "detected (normalized) model")
+
+			in.Status = 200
+			in.RespHeaders = []middleware.KV{{Key: "Content-Type", Value: tc.respCT}}
+			in.RespBody = tc.respBody
+
+			respOut, err := respMW.Invoke(context.Background(), in)
+			require.NoError(t, err, "response parser")
+			in.Metadata = append(in.Metadata, respOut.Metadata...)
+
+			costOut, err := costMW.Invoke(context.Background(), in)
+			require.NoError(t, err, "cost meter")
+
+			if tc.wantSkip != "" {
+				assert.Equal(t, tc.wantSkip, metaKV(costOut.Metadata, middleware.KeyCostSkipped), "expected cost skip reason")
+				assert.Empty(t, metaKV(costOut.Metadata, middleware.KeyCostUSDTotal), "no cost may be emitted on skip")
+				return
+			}
+
+			raw := metaKV(costOut.Metadata, middleware.KeyCostUSDTotal)
+			require.NotEmpty(t, raw, "cost.usd_total must be emitted; skip=%q", metaKV(costOut.Metadata, middleware.KeyCostSkipped))
+			got, err := strconv.ParseFloat(raw, 64)
+			require.NoError(t, err, "cost must be a float")
+			// cost.usd_total is rendered with %.6f: allow half of the last printed digit on top of float error.
+			assert.InDelta(t, tc.wantCost, got, 5.1e-7, "USD cost for %s", tc.name)
+
+			rawCache := metaKV(costOut.Metadata, middleware.KeyCostUSDCache)
+			require.NotEmpty(t, rawCache, "cost.usd_cache must be emitted next to cost.usd_total")
+			gotCache, err := strconv.ParseFloat(rawCache, 64)
+			require.NoError(t, err, "cache cost must be a float")
+			assert.InDelta(t, tc.wantCacheCost, gotCache, 5.1e-7, "cache USD cost for %s", tc.name)
+		})
+	}
+}
+
+// metaKV returns the value for key in kvs, or "" when absent.
+func metaKV(kvs []middleware.KV, key string) string {
+	for _, kv := range kvs {
+		if kv.Key == key {
+			return kv.Value
+		}
+	}
+	return ""
+}
+
+// sseBody renders data frames as a text/event-stream body.
+func sseBody(frames ...string) []byte {
+	var b bytes.Buffer
+	for _, f := range frames {
+		b.WriteString("data: ")
+		b.WriteString(f)
+		b.WriteString("\n\n")
+	}
+	return b.Bytes()
+}
+
+// awsFrame encodes one AWS event-stream frame with the given :event-type.
+func awsFrame(t *testing.T, eventType string, payload []byte) []byte {
+	t.Helper()
+	var buf bytes.Buffer
+	enc := eventstream.NewEncoder()
+	require.NoError(t, enc.Encode(&buf, eventstream.Message{
+		Headers: eventstream.Headers{{Name: ":event-type", Value: eventstream.StringValue(eventType)}},
+		Payload: payload,
+	}), "encode event-stream frame")
+	return buf.Bytes()
+}
+
+// bedrockInvokeStream builds an invoke-with-response-stream body: each "chunk" frame wraps a base64 Anthropic event.
+func bedrockInvokeStream(t *testing.T, events ...string) []byte {
+	t.Helper()
+	var body bytes.Buffer
+	for _, ev := range events {
+		wrap, err := json.Marshal(map[string]string{"bytes": base64.StdEncoding.EncodeToString([]byte(ev))})
+		require.NoError(t, err)
+		body.Write(awsFrame(t, "chunk", wrap))
+	}
+	return body.Bytes()
+}
+
+// bedrockConverseStream builds a converse-stream body: contentBlockDelta frames plus a trailing metadata usage frame.
+func bedrockConverseStream(t *testing.T, deltas ...string) []byte {
+	t.Helper()
+	var body bytes.Buffer
+	for i, ev := range deltas {
+		eventType := "contentBlockDelta"
+		if i == len(deltas)-1 {
+			eventType = "metadata"
+		}
+		body.Write(awsFrame(t, eventType, []byte(ev)))
+	}
+	return body.Bytes()
+}

proxy/internal/middleware/builtin/cost_meter/factory.go

@@ -0,0 +1,88 @@
+package cost_meter
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+
+	"github.com/netbirdio/netbird/proxy/internal/llm/pricing"
+	"github.com/netbirdio/netbird/proxy/internal/middleware"
+	"github.com/netbirdio/netbird/proxy/internal/middleware/builtin"
+)
+
+// defaultPricingFilename is the basename probed inside the proxy data
+// directory when no override is configured.
+const defaultPricingFilename = "pricing.yaml"
+
+// Config is the on-wire configuration for the middleware.
+type Config struct {
+	// PricingPath optionally overrides the basename of the pricing
+	// file probed inside the proxy data directory. When empty the
+	// loader falls back to "pricing.yaml".
+	PricingPath string `json:"pricing_path"`
+}
+
+// Factory builds cost_meter instances from raw config bytes.
+type Factory struct{}
+
+// ID returns the registry identifier.
+func (Factory) ID() string { return ID }
+
+// New constructs a middleware instance. Empty, null, and {} configs
+// are accepted; non-empty rawConfig that fails to unmarshal is
+// rejected so misconfigurations surface at chain build time. The
+// pricing loader is built once per instance and reused across
+// invocations.
+func (Factory) New(rawConfig []byte) (middleware.Middleware, error) {
+	cfg, err := decodeConfig(rawConfig)
+	if err != nil {
+		return nil, err
+	}
+
+	fctx := builtin.Context()
+	pricingPath := cfg.PricingPath
+	if pricingPath == "" {
+		pricingPath = defaultPricingFilename
+	}
+
+	loader, err := pricing.NewLoader(fctx.DataDir, pricingPath, ID, nil)
+	if err != nil {
+		return nil, fmt.Errorf("init pricing loader: %w", err)
+	}
+
+	cancel := startReloader(fctx.Context, loader)
+
+	return newMiddleware(loader, cancel), nil
+}
+
+// startReloader binds the loader's mtime-poll goroutine to a context
+// derived from the proxy-lifetime context and returns its cancel func so
+// the owning middleware can stop the goroutine on teardown. Returns nil
+// when there's nothing to watch (nil context or defaults-only loader), in
+// which case the middleware's Close is a no-op.
+func startReloader(ctx context.Context, loader *pricing.Loader) context.CancelFunc {
+	if ctx == nil || !loader.WatchesFile() {
+		return nil
+	}
+	cctx, cancel := context.WithCancel(ctx)
+	go loader.Reload(cctx)
+	return cancel
+}
+
+// decodeConfig accepts empty, null, and {} configs, returning a
+// zero-value Config. Non-empty payloads must parse cleanly.
+func decodeConfig(rawConfig []byte) (Config, error) {
+	var cfg Config
+	if len(bytes.TrimSpace(rawConfig)) == 0 {
+		return cfg, nil
+	}
+	if err := json.Unmarshal(rawConfig, &cfg); err != nil {
+		return cfg, fmt.Errorf("decode config: %w", err)
+	}
+	return cfg, nil
+}
+
+func init() {
+	builtin.Register(Factory{})
+}

proxy/internal/middleware/builtin/cost_meter/middleware.go

@@ -0,0 +1,218 @@
+// Package cost_meter implements the SlotOnResponse middleware that
+// converts token-usage metadata emitted by llm_response_parser into a
+// per-request USD cost estimate. The middleware uses the shared pricing
+// loader so operator pricing overrides apply to the chain.
+package cost_meter
+
+import (
+	"context"
+	"fmt"
+	"strconv"
+
+	"github.com/netbirdio/netbird/proxy/internal/llm/pricing"
+	"github.com/netbirdio/netbird/proxy/internal/middleware"
+)
+
+// ID is the registry identifier for this middleware.
+const ID = "cost_meter"
+
+// Version is the implementation version emitted via the spec merge.
+const Version = "1.0.0"
+
+// Skip reasons emitted under KeyCostSkipped. The set is closed; the
+// dashboard surfaces these verbatim.
+const (
+	skipMissingProvider = "missing_provider"
+	skipMissingModel    = "missing_model"
+	skipMissingTokens   = "missing_tokens"
+	//nolint:gosec // skip-reason label, not a credential
+	skipUnparseableTokens = "unparseable_tokens"
+	skipZeroTokens        = "zero_tokens"
+	skipUnknownModel      = "unknown_model"
+)
+
+var metadataKeys = []string{
+	middleware.KeyCostUSDInput,
+	middleware.KeyCostUSDCachedInput,
+	middleware.KeyCostUSDCacheCreation,
+	middleware.KeyCostUSDOutput,
+	middleware.KeyCostUSDTotal,
+	middleware.KeyCostUSDCache,
+	middleware.KeyCostSkipped,
+}
+
+// Middleware computes a per-response cost estimate from the token
+// counts emitted upstream by llm_response_parser.
+type Middleware struct {
+	loader *pricing.Loader
+	// cancel stops this instance's pricing-reload goroutine. Non-nil only
+	// when the loader watches an override file; Close calls it so a chain
+	// rebuild doesn't leak a poll goroutine per retired instance.
+	cancel context.CancelFunc
+}
+
+// newMiddleware constructs a Middleware bound to the given pricing loader.
+// cancel may be nil (defaults-only loader with no reloader to stop).
+func newMiddleware(loader *pricing.Loader, cancel context.CancelFunc) *Middleware {
+	return &Middleware{loader: loader, cancel: cancel}
+}
+
+// ID returns the registry identifier.
+func (m *Middleware) ID() string { return ID }
+
+// Version returns the implementation version.
+func (m *Middleware) Version() string { return Version }
+
+// Slot reports that the middleware runs after the upstream call.
+func (m *Middleware) Slot() middleware.Slot { return middleware.SlotOnResponse }
+
+// AcceptedContentTypes is empty: cost_meter never inspects bodies.
+func (m *Middleware) AcceptedContentTypes() []string { return []string{} }
+
+// MetadataKeys returns the closed allowlist of keys this middleware
+// may emit.
+func (m *Middleware) MetadataKeys() []string {
+	return append([]string(nil), metadataKeys...)
+}
+
+// MutationsSupported reports that this middleware never mutates the
+// response.
+func (m *Middleware) MutationsSupported() bool { return false }
+
+// Close stops this instance's pricing-reload goroutine, if any. Called by
+// the chain when a rebuild retires the instance, so the mtime-poll loop
+// doesn't outlive the chain it belonged to. Safe to call on a nil receiver
+// and on an instance with no reloader.
+func (m *Middleware) Close() error {
+	if m != nil && m.cancel != nil {
+		m.cancel()
+	}
+	return nil
+}
+
+// Invoke reads provider, model, and token metadata, looks up pricing,
+// and emits either KeyCostUSDTotal or KeyCostSkipped. The decision is
+// always DecisionAllow; cost metering never denies or mutates.
+func (m *Middleware) Invoke(_ context.Context, in *middleware.Input) (*middleware.Output, error) {
+	out := &middleware.Output{Decision: middleware.DecisionAllow}
+	if in == nil {
+		return out, nil
+	}
+
+	provider := lookupKV(in.Metadata, middleware.KeyLLMProvider)
+	if provider == "" {
+		out.Metadata = skip(skipMissingProvider)
+		return out, nil
+	}
+
+	model := lookupKV(in.Metadata, middleware.KeyLLMModel)
+	if model == "" {
+		out.Metadata = skip(skipMissingModel)
+		return out, nil
+	}
+
+	inRaw, hasIn := lookupKVOK(in.Metadata, middleware.KeyLLMInputTokens)
+	outRaw, hasOut := lookupKVOK(in.Metadata, middleware.KeyLLMOutputTokens)
+	if !hasIn || !hasOut {
+		out.Metadata = skip(skipMissingTokens)
+		return out, nil
+	}
+
+	inTokens, err := strconv.ParseInt(inRaw, 10, 64)
+	if err != nil || inTokens < 0 {
+		// Unparseable or negative tokens are not a runtime error: the
+		// upstream llm_response_parser emitted a non-numeric / invalid
+		// value, so we surface that as cost.skipped and continue with
+		// Allow rather than pricing a negative count.
+		out.Metadata = skip(skipUnparseableTokens)
+		return out, nil //nolint:nilerr // structured skip; not a runtime error
+	}
+	outTokens, err := strconv.ParseInt(outRaw, 10, 64)
+	if err != nil || outTokens < 0 {
+		out.Metadata = skip(skipUnparseableTokens)
+		return out, nil //nolint:nilerr // structured skip; not a runtime error
+	}
+
+	// Cache buckets are optional and silently zeroed on a missing /
+	// malformed value; they're a refinement on top of input cost,
+	// not a precondition. A buggy value falls back to 0, never aborts.
+	cachedTokens := parseOptionalInt64(in.Metadata, middleware.KeyLLMCachedInputTokens)
+	cacheCreationTokens := parseOptionalInt64(in.Metadata, middleware.KeyLLMCacheCreationTokens)
+
+	if inTokens == 0 && outTokens == 0 && cachedTokens == 0 && cacheCreationTokens == 0 {
+		out.Metadata = skip(skipZeroTokens)
+		return out, nil
+	}
+
+	table := m.loader.Get()
+	costs, ok := table.Costs(provider, model, inTokens, outTokens, cachedTokens, cacheCreationTokens)
+	if !ok {
+		out.Metadata = skip(skipUnknownModel)
+		return out, nil
+	}
+
+	// Per-bucket costs first: they're the base of the breakdown, and the two
+	// aggregates that follow are derived from exactly these four values.
+	out.Metadata = []middleware.KV{
+		{Key: middleware.KeyCostUSDInput, Value: usd(costs.InputUSD)},
+		{Key: middleware.KeyCostUSDCachedInput, Value: usd(costs.CachedInputUSD)},
+		{Key: middleware.KeyCostUSDCacheCreation, Value: usd(costs.CacheCreationUSD)},
+		{Key: middleware.KeyCostUSDOutput, Value: usd(costs.OutputUSD)},
+		{Key: middleware.KeyCostUSDTotal, Value: usd(costs.TotalUSD)},
+		{Key: middleware.KeyCostUSDCache, Value: usd(costs.CacheUSD)},
+	}
+	return out, nil
+}
+
+// usd renders a cost as the fixed-precision string every cost.usd_* key
+// carries, so the per-bucket values and the aggregates round identically.
+//
+// 9 decimals, not 6: these values are summed downstream — per request, per
+// session, and per usage bucket — so the rounding step is applied once per
+// bucket per row and then accumulated. At 6 decimals a single row loses up to
+// 2e-6 across its four buckets (enough to break a 1e-6 reconciliation against
+// published rates), and a bucket smaller than half a microdollar quantises to
+// zero outright: 16 cache-read tokens on a cheap model is 1.6e-9, so summing
+// 10k such rows reports 0.02 instead of 0.016. Nano-dollar precision keeps the
+// per-row error ~1000x below the smallest realistic bucket.
+func usd(v float64) string { return fmt.Sprintf("%.9f", v) }
+
+// skip returns a single-entry metadata slice carrying the given skip
+// reason under KeyCostSkipped.
+func skip(reason string) []middleware.KV {
+	return []middleware.KV{{Key: middleware.KeyCostSkipped, Value: reason}}
+}
+
+// lookupKV returns the value associated with key, or the empty string
+// when the key is absent.
+func lookupKV(kvs []middleware.KV, key string) string {
+	v, _ := lookupKVOK(kvs, key)
+	return v
+}
+
+// lookupKVOK returns the value associated with key plus a presence
+// flag so callers can distinguish absent from empty.
+func lookupKVOK(kvs []middleware.KV, key string) (string, bool) {
+	for _, kv := range kvs {
+		if kv.Key == key {
+			return kv.Value, true
+		}
+	}
+	return "", false
+}
+
+// parseOptionalInt64 reads a metadata value and decodes it as int64.
+// Absent or unparseable values yield 0 — the caller treats absence as
+// "no cached tokens" rather than an error, since cache buckets are a
+// refinement, not a precondition.
+func parseOptionalInt64(kvs []middleware.KV, key string) int64 {
+	raw, ok := lookupKVOK(kvs, key)
+	if !ok {
+		return 0
+	}
+	v, err := strconv.ParseInt(raw, 10, 64)
+	if err != nil || v < 0 {
+		return 0
+	}
+	return v
+}

proxy/internal/middleware/builtin/cost_meter/middleware_test.go

@@ -0,0 +1,508 @@
+package cost_meter
+
+import (
+	"context"
+	"encoding/json"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/proxy/internal/middleware"
+	"github.com/netbirdio/netbird/proxy/internal/middleware/builtin"
+)
+
+const fixturePricing = `openai:
+  gpt-4o:
+    input_per_1k: 0.0025
+    output_per_1k: 0.01
+  gpt-4o-mini:
+    input_per_1k: 0.00015
+    output_per_1k: 0.0006
+anthropic:
+  claude-sonnet-4-5:
+    input_per_1k: 0.003
+    output_per_1k: 0.015
+`
+
+// configureBuiltin points the package-level FactoryContext at a tmp
+// directory containing the test pricing fixture. Returns the path so
+// callers can override files later if needed.
+func configureBuiltin(t *testing.T) string {
+	t.Helper()
+	dir := t.TempDir()
+	require.NoError(t, os.WriteFile(filepath.Join(dir, "pricing.yaml"), []byte(fixturePricing), 0o600), "write pricing fixture")
+	builtin.Configure(context.Background(), dir, nil, nil, nil)
+	return dir
+}
+
+func metaValue(t *testing.T, kvs []middleware.KV, key string) (string, bool) {
+	t.Helper()
+	for _, kv := range kvs {
+		if kv.Key == key {
+			return kv.Value, true
+		}
+	}
+	return "", false
+}
+
+func buildMiddleware(t *testing.T, raw []byte) middleware.Middleware {
+	t.Helper()
+	mw, err := Factory{}.New(raw)
+	require.NoError(t, err, "factory must accept the supplied config")
+	return mw
+}
+
+func TestMiddleware_StaticSurface(t *testing.T) {
+	configureBuiltin(t)
+	mw := buildMiddleware(t, nil)
+
+	assert.Equal(t, ID, mw.ID(), "ID must match the registered constant")
+	assert.Equal(t, Version, mw.Version(), "Version must match the constant")
+	assert.Equal(t, middleware.SlotOnResponse, mw.Slot(), "must run in the response slot")
+	assert.Empty(t, mw.AcceptedContentTypes(), "cost_meter does not inspect bodies")
+	assert.False(t, mw.MutationsSupported(), "cost_meter never mutates")
+	assert.NoError(t, mw.Close(), "Close on stateless middleware is a no-op")
+
+	keys := mw.MetadataKeys()
+	expected := []string{
+		middleware.KeyCostUSDInput,
+		middleware.KeyCostUSDCachedInput,
+		middleware.KeyCostUSDCacheCreation,
+		middleware.KeyCostUSDOutput,
+		middleware.KeyCostUSDTotal,
+		middleware.KeyCostUSDCache,
+		middleware.KeyCostSkipped,
+	}
+	assert.Equal(t, expected, keys, "metadata key allowlist must match the spec")
+}
+
+func TestFactory_AcceptsEmptyAndJSONConfig(t *testing.T) {
+	configureBuiltin(t)
+	cases := [][]byte{nil, {}, []byte("null"), []byte("{}"), []byte("   ")}
+	for _, raw := range cases {
+		mw, err := Factory{}.New(raw)
+		require.NoError(t, err, "empty/null/object config must be accepted")
+		require.NotNil(t, mw, "factory must return a middleware instance")
+	}
+}
+
+func TestFactory_RejectsMalformedConfig(t *testing.T) {
+	configureBuiltin(t)
+	mw, err := Factory{}.New([]byte("{not json"))
+	require.Error(t, err, "malformed config must surface at construction")
+	assert.Nil(t, mw, "no instance is returned on error")
+}
+
+func TestFactory_DefaultPricingPathLoadsFixture(t *testing.T) {
+	configureBuiltin(t)
+	mw := buildMiddleware(t, nil)
+
+	out, err := mw.Invoke(context.Background(), &middleware.Input{
+		Metadata: []middleware.KV{
+			{Key: middleware.KeyLLMProvider, Value: "openai"},
+			{Key: middleware.KeyLLMModel, Value: "gpt-4o-mini"},
+			{Key: middleware.KeyLLMInputTokens, Value: "1000"},
+			{Key: middleware.KeyLLMOutputTokens, Value: "1000"},
+		},
+	})
+	require.NoError(t, err)
+	require.Equal(t, middleware.DecisionAllow, out.Decision, "cost_meter always allows")
+
+	value, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal)
+	require.True(t, ok, "cost.usd_total must be emitted for known model")
+	assert.Equal(t, "0.000750000", value, "0.00015 + 0.0006 per 1k tokens, 9-decimal format")
+}
+
+func TestFactory_PricingPathOverride(t *testing.T) {
+	dir := t.TempDir()
+	require.NoError(t, os.WriteFile(filepath.Join(dir, "custom.yaml"), []byte(fixturePricing), 0o600), "write custom pricing")
+	builtin.Configure(context.Background(), dir, nil, nil, nil)
+
+	raw, err := json.Marshal(Config{PricingPath: "custom.yaml"})
+	require.NoError(t, err)
+
+	mw := buildMiddleware(t, raw)
+	out, err := mw.Invoke(context.Background(), &middleware.Input{
+		Metadata: []middleware.KV{
+			{Key: middleware.KeyLLMProvider, Value: "openai"},
+			{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
+			{Key: middleware.KeyLLMInputTokens, Value: "2000"},
+			{Key: middleware.KeyLLMOutputTokens, Value: "1000"},
+		},
+	})
+	require.NoError(t, err)
+
+	value, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal)
+	require.True(t, ok, "cost.usd_total must be emitted with custom pricing path")
+	assert.Equal(t, "0.015000000", value, "2*0.0025 + 1*0.01 = 0.015 with 9-decimal format")
+}
+
+func TestInvoke_ComputesCostForKnownModel(t *testing.T) {
+	configureBuiltin(t)
+	mw := buildMiddleware(t, nil)
+
+	out, err := mw.Invoke(context.Background(), &middleware.Input{
+		Metadata: []middleware.KV{
+			{Key: middleware.KeyLLMProvider, Value: "anthropic"},
+			{Key: middleware.KeyLLMModel, Value: "claude-sonnet-4-5"},
+			{Key: middleware.KeyLLMInputTokens, Value: "1000"},
+			{Key: middleware.KeyLLMOutputTokens, Value: "1000"},
+		},
+	})
+	require.NoError(t, err)
+
+	value, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal)
+	require.True(t, ok, "cost.usd_total must be emitted")
+	assert.Equal(t, "0.018000000", value, "0.003 + 0.015 = 0.018 with 9-decimal format")
+	_, skipped := metaValue(t, out.Metadata, middleware.KeyCostSkipped)
+	assert.False(t, skipped, "cost.skipped must not be set when cost is computed")
+}
+
+func TestInvoke_MissingProvider(t *testing.T) {
+	configureBuiltin(t)
+	mw := buildMiddleware(t, nil)
+
+	out, err := mw.Invoke(context.Background(), &middleware.Input{
+		Metadata: []middleware.KV{
+			{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
+			{Key: middleware.KeyLLMInputTokens, Value: "10"},
+			{Key: middleware.KeyLLMOutputTokens, Value: "10"},
+		},
+	})
+	require.NoError(t, err)
+	value, ok := metaValue(t, out.Metadata, middleware.KeyCostSkipped)
+	require.True(t, ok, "cost.skipped must be set when provider is missing")
+	assert.Equal(t, skipMissingProvider, value, "skip reason matches missing_provider")
+}
+
+func TestInvoke_MissingModel(t *testing.T) {
+	configureBuiltin(t)
+	mw := buildMiddleware(t, nil)
+
+	out, err := mw.Invoke(context.Background(), &middleware.Input{
+		Metadata: []middleware.KV{
+			{Key: middleware.KeyLLMProvider, Value: "openai"},
+			{Key: middleware.KeyLLMInputTokens, Value: "10"},
+			{Key: middleware.KeyLLMOutputTokens, Value: "10"},
+		},
+	})
+	require.NoError(t, err)
+	value, ok := metaValue(t, out.Metadata, middleware.KeyCostSkipped)
+	require.True(t, ok, "cost.skipped must be set when model is missing")
+	assert.Equal(t, skipMissingModel, value, "skip reason matches missing_model")
+}
+
+func TestInvoke_MissingTokens(t *testing.T) {
+	configureBuiltin(t)
+	mw := buildMiddleware(t, nil)
+
+	cases := []struct {
+		name string
+		md   []middleware.KV
+	}{
+		{
+			name: "input only",
+			md: []middleware.KV{
+				{Key: middleware.KeyLLMProvider, Value: "openai"},
+				{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
+				{Key: middleware.KeyLLMInputTokens, Value: "10"},
+			},
+		},
+		{
+			name: "output only",
+			md: []middleware.KV{
+				{Key: middleware.KeyLLMProvider, Value: "openai"},
+				{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
+				{Key: middleware.KeyLLMOutputTokens, Value: "10"},
+			},
+		},
+		{
+			name: "neither",
+			md: []middleware.KV{
+				{Key: middleware.KeyLLMProvider, Value: "openai"},
+				{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
+			},
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			out, err := mw.Invoke(context.Background(), &middleware.Input{Metadata: tc.md})
+			require.NoError(t, err)
+			value, ok := metaValue(t, out.Metadata, middleware.KeyCostSkipped)
+			require.True(t, ok, "cost.skipped must be set when token keys are missing")
+			assert.Equal(t, skipMissingTokens, value, "skip reason matches missing_tokens")
+		})
+	}
+}
+
+func TestInvoke_UnparseableTokens(t *testing.T) {
+	configureBuiltin(t)
+	mw := buildMiddleware(t, nil)
+
+	cases := []struct {
+		name string
+		in   string
+		out  string
+	}{
+		{name: "input non-numeric", in: "abc", out: "10"},
+		{name: "output non-numeric", in: "10", out: "xyz"},
+		{name: "both garbage", in: "??", out: "??"},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			out, err := mw.Invoke(context.Background(), &middleware.Input{
+				Metadata: []middleware.KV{
+					{Key: middleware.KeyLLMProvider, Value: "openai"},
+					{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
+					{Key: middleware.KeyLLMInputTokens, Value: tc.in},
+					{Key: middleware.KeyLLMOutputTokens, Value: tc.out},
+				},
+			})
+			require.NoError(t, err)
+			value, ok := metaValue(t, out.Metadata, middleware.KeyCostSkipped)
+			require.True(t, ok, "cost.skipped must be set on unparseable tokens")
+			assert.Equal(t, skipUnparseableTokens, value, "skip reason matches unparseable_tokens")
+		})
+	}
+}
+
+func TestInvoke_ZeroTokens(t *testing.T) {
+	configureBuiltin(t)
+	mw := buildMiddleware(t, nil)
+
+	out, err := mw.Invoke(context.Background(), &middleware.Input{
+		Metadata: []middleware.KV{
+			{Key: middleware.KeyLLMProvider, Value: "openai"},
+			{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
+			{Key: middleware.KeyLLMInputTokens, Value: "0"},
+			{Key: middleware.KeyLLMOutputTokens, Value: "0"},
+		},
+	})
+	require.NoError(t, err)
+	value, ok := metaValue(t, out.Metadata, middleware.KeyCostSkipped)
+	require.True(t, ok, "cost.skipped must be set when both token counts are zero")
+	assert.Equal(t, skipZeroTokens, value, "skip reason matches zero_tokens")
+	_, hasCost := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal)
+	assert.False(t, hasCost, "cost.usd_total must not be emitted for zero tokens")
+}
+
+func TestInvoke_UnknownModel(t *testing.T) {
+	configureBuiltin(t)
+	mw := buildMiddleware(t, nil)
+
+	out, err := mw.Invoke(context.Background(), &middleware.Input{
+		Metadata: []middleware.KV{
+			{Key: middleware.KeyLLMProvider, Value: "openai"},
+			{Key: middleware.KeyLLMModel, Value: "fantasy-model-9000"},
+			{Key: middleware.KeyLLMInputTokens, Value: "10"},
+			{Key: middleware.KeyLLMOutputTokens, Value: "10"},
+		},
+	})
+	require.NoError(t, err)
+	value, ok := metaValue(t, out.Metadata, middleware.KeyCostSkipped)
+	require.True(t, ok, "cost.skipped must be set when pricing entry is absent")
+	assert.Equal(t, skipUnknownModel, value, "skip reason matches unknown_model")
+}
+
+func TestInvoke_NilInput(t *testing.T) {
+	configureBuiltin(t)
+	mw := buildMiddleware(t, nil)
+
+	out, err := mw.Invoke(context.Background(), nil)
+	require.NoError(t, err)
+	require.NotNil(t, out, "output must be returned even on nil input")
+	assert.Equal(t, middleware.DecisionAllow, out.Decision, "decision must be allow on nil input")
+	assert.Empty(t, out.Metadata, "no metadata must be emitted on nil input")
+}
+
+const fixturePricingWithCache = `openai:
+  gpt-4o:
+    input_per_1k: 0.0025
+    output_per_1k: 0.01
+    cached_input_per_1k: 0.00125
+anthropic:
+  claude-sonnet-4-5:
+    input_per_1k: 0.003
+    output_per_1k: 0.015
+    cache_read_per_1k: 0.0003
+    cache_creation_per_1k: 0.00375
+`
+
+// configureBuiltinWithCacheRates points the package-level
+// FactoryContext at a tmp directory containing pricing entries that
+// include the cache rate fields.
+func configureBuiltinWithCacheRates(t *testing.T) {
+	t.Helper()
+	dir := t.TempDir()
+	require.NoError(t, os.WriteFile(filepath.Join(dir, "pricing.yaml"), []byte(fixturePricingWithCache), 0o600), "write cache-aware pricing fixture")
+	builtin.Configure(context.Background(), dir, nil, nil, nil)
+}
+
+// TestInvoke_OpenAICachedSubsetDiscount proves the OpenAI shape end
+// to end through the middleware: cached_input_tokens is treated as a
+// SUBSET of input_tokens and discounted at the configured rate, not
+// added on top.
+func TestInvoke_OpenAICachedSubsetDiscount(t *testing.T) {
+	configureBuiltinWithCacheRates(t)
+	mw := buildMiddleware(t, nil)
+
+	out, err := mw.Invoke(context.Background(), &middleware.Input{
+		Metadata: []middleware.KV{
+			{Key: middleware.KeyLLMProvider, Value: "openai"},
+			{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
+			{Key: middleware.KeyLLMInputTokens, Value: "1000"},
+			{Key: middleware.KeyLLMOutputTokens, Value: "500"},
+			{Key: middleware.KeyLLMCachedInputTokens, Value: "750"},
+		},
+	})
+	require.NoError(t, err)
+	require.Equal(t, middleware.DecisionAllow, out.Decision)
+
+	value, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal)
+	require.True(t, ok, "cached subset path must produce a cost — never a skip")
+	// 250 non-cached at 0.0025/1k + 750 cached at 0.00125/1k + 500 output at 0.01/1k.
+	assert.Equal(t, "0.006562500", value,
+		"cached subset must be billed at the discount rate, non-cached at the full rate; never double-billed")
+
+	cache, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDCache)
+	require.True(t, ok, "cost.usd_cache must be emitted alongside cost.usd_total")
+	// 750 cached at 0.00125/1k = 0.0009375.
+	assert.Equal(t, "0.000937500", cache, "cache cost is the discounted cost of the cached subset")
+
+	// Per-bucket breakdown. On OpenAI the cached subset is carved out of the
+	// input bucket, so input covers only the 250 non-cached tokens — the two
+	// must never double-count the same 750 tokens.
+	assertBucket(t, out.Metadata, middleware.KeyCostUSDInput, "0.000625000",
+		"input bucket bills only the non-cached remainder")
+	assertBucket(t, out.Metadata, middleware.KeyCostUSDCachedInput, "0.000937500",
+		"cached-input bucket bills the discounted subset")
+	assertBucket(t, out.Metadata, middleware.KeyCostUSDCacheCreation, "0.000000000",
+		"OpenAI has no cache-write bucket")
+	assertBucket(t, out.Metadata, middleware.KeyCostUSDOutput, "0.005000000",
+		"output bucket bills 500 tokens at 0.01/1k")
+}
+
+// TestInvoke_AnthropicCacheBucketsAdditive proves the Anthropic
+// shape: cache_read and cache_creation are additive to input_tokens
+// and each carries its own rate.
+func TestInvoke_AnthropicCacheBucketsAdditive(t *testing.T) {
+	configureBuiltinWithCacheRates(t)
+	mw := buildMiddleware(t, nil)
+
+	out, err := mw.Invoke(context.Background(), &middleware.Input{
+		Metadata: []middleware.KV{
+			{Key: middleware.KeyLLMProvider, Value: "anthropic"},
+			{Key: middleware.KeyLLMModel, Value: "claude-sonnet-4-5"},
+			{Key: middleware.KeyLLMInputTokens, Value: "256"},
+			{Key: middleware.KeyLLMOutputTokens, Value: "200"},
+			{Key: middleware.KeyLLMCachedInputTokens, Value: "768"},
+			{Key: middleware.KeyLLMCacheCreationTokens, Value: "512"},
+		},
+	})
+	require.NoError(t, err)
+	require.Equal(t, middleware.DecisionAllow, out.Decision)
+
+	value, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal)
+	require.True(t, ok)
+	// 256 input * 0.003 + 768 cache_read * 0.0003 + 512 cache_creation * 0.00375 + 200 output * 0.015
+	// = 0.000768 + 0.0002304 + 0.00192 + 0.003 = 0.0059184.
+	assert.Equal(t, "0.005918400", value,
+		"each Anthropic input bucket must bill at its own rate — cache_read cheap, cache_creation expensive, regular input mid")
+
+	cache, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDCache)
+	require.True(t, ok, "cost.usd_cache must be emitted alongside cost.usd_total")
+	// 768 cache_read * 0.0003 + 512 cache_creation * 0.00375 = 0.0021504.
+	assert.Equal(t, "0.002150400", cache, "cache cost sums the read and creation buckets")
+
+	// Per-bucket breakdown: four separately-billed buckets, each at its own rate.
+	assertBucket(t, out.Metadata, middleware.KeyCostUSDInput, "0.000768000",
+		"input bucket bills 256 tokens at 0.003/1k")
+	assertBucket(t, out.Metadata, middleware.KeyCostUSDCachedInput, "0.000230400",
+		"cache-read bucket bills 768 tokens at the cheap 0.0003/1k")
+	assertBucket(t, out.Metadata, middleware.KeyCostUSDCacheCreation, "0.001920000",
+		"cache-write bucket bills 512 tokens at the expensive 0.00375/1k")
+	assertBucket(t, out.Metadata, middleware.KeyCostUSDOutput, "0.003000000",
+		"output bucket bills 200 tokens at 0.015/1k")
+}
+
+// assertBucket asserts one per-bucket cost key carries the expected
+// 6-decimal value.
+func assertBucket(t *testing.T, md []middleware.KV, key, want, msg string) {
+	t.Helper()
+	got, ok := metaValue(t, md, key)
+	require.Truef(t, ok, "%s must be emitted", key)
+	assert.Equal(t, want, got, msg)
+}
+
+// TestInvoke_CachedTokensAbsentFallsBackToBaseFormula covers the
+// "operator hasn't opted in" path: with no cached metadata keys
+// emitted, the meter must produce exactly the same cost as before
+// the feature landed. Critical so operators with the new binary but
+// no YAML changes see no behavioural drift on OpenAI requests.
+func TestInvoke_CachedTokensAbsentFallsBackToBaseFormula(t *testing.T) {
+	configureBuiltinWithCacheRates(t)
+	mw := buildMiddleware(t, nil)
+
+	out, err := mw.Invoke(context.Background(), &middleware.Input{
+		Metadata: []middleware.KV{
+			{Key: middleware.KeyLLMProvider, Value: "openai"},
+			{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
+			{Key: middleware.KeyLLMInputTokens, Value: "1000"},
+			{Key: middleware.KeyLLMOutputTokens, Value: "500"},
+			// No KeyLLMCachedInputTokens — the parser didn't see one.
+		},
+	})
+	require.NoError(t, err)
+	value, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal)
+	require.True(t, ok)
+	// 1000 input * 0.0025 + 500 output * 0.01 = 0.0025 + 0.005 = 0.0075
+	assert.Equal(t, "0.007500000", value, "no cached metadata = same cost as before the feature landed")
+}
+
+// TestInvoke_UnparseableCachedTokensSkippedSilently proves the
+// optional-bucket contract: a malformed cached_input_tokens metadata
+// value falls back to 0 (= no cached count) and continues with the
+// regular formula. Cache buckets are a refinement, never a reason to
+// abort cost computation.
+func TestInvoke_UnparseableCachedTokensSkippedSilently(t *testing.T) {
+	configureBuiltinWithCacheRates(t)
+	mw := buildMiddleware(t, nil)
+
+	out, err := mw.Invoke(context.Background(), &middleware.Input{
+		Metadata: []middleware.KV{
+			{Key: middleware.KeyLLMProvider, Value: "openai"},
+			{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
+			{Key: middleware.KeyLLMInputTokens, Value: "1000"},
+			{Key: middleware.KeyLLMOutputTokens, Value: "500"},
+			{Key: middleware.KeyLLMCachedInputTokens, Value: "not-a-number"},
+		},
+	})
+	require.NoError(t, err)
+	value, ok := metaValue(t, out.Metadata, middleware.KeyCostUSDTotal)
+	require.True(t, ok, "garbage cache metadata must NOT switch the response from a cost to a skip — fall back to 0 cached")
+	assert.Equal(t, "0.007500000", value, "same as the no-cached-metadata path")
+}
+
+// TestMiddleware_CloseCancelsReloader proves Close stops the per-instance
+// pricing-reload goroutine: a chain rebuild retires the old instance and
+// calls Close, which must invoke the cancel func startReloader handed it so
+// the mtime-poll loop doesn't outlive the chain.
+func TestMiddleware_CloseCancelsReloader(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	m := newMiddleware(nil, cancel)
+
+	require.NoError(t, m.Close(), "Close must not error")
+	require.Error(t, ctx.Err(), "Close must cancel the reloader context so the poll goroutine exits")
+}
+
+// TestMiddleware_CloseNilSafe confirms Close is a no-op (no panic) for an
+// instance with no reloader and for a nil receiver.
+func TestMiddleware_CloseNilSafe(t *testing.T) {
+	require.NoError(t, newMiddleware(nil, nil).Close(), "no-reloader Close must be a no-op")
+	var m *Middleware
+	require.NoError(t, m.Close(), "nil-receiver Close must be safe")
+}

proxy/internal/middleware/builtin/llm_guardrail/factory.go

@@ -0,0 +1,82 @@
+package llm_guardrail
+
+import (
+	"encoding/json"
+	"fmt"
+	"strings"
+
+	"github.com/netbirdio/netbird/proxy/internal/middleware"
+	"github.com/netbirdio/netbird/proxy/internal/middleware/builtin"
+)
+
+// Config is the JSON-decoded shape accepted by the factory. The
+// runtime path consumes the normalised allowlist; raw config is not
+// retained beyond construction.
+type Config struct {
+	ModelAllowlist []string      `json:"model_allowlist"`
+	PromptCapture  PromptCapture `json:"prompt_capture"`
+}
+
+// PromptCapture toggles the optional prompt capture + redaction step
+// that emits llm.request_prompt onto the metadata bag.
+type PromptCapture struct {
+	Enabled   bool `json:"enabled"`
+	RedactPii bool `json:"redact_pii"`
+}
+
+// Factory builds a configured llm_guardrail middleware instance.
+type Factory struct{}
+
+// ID returns the registry identifier matching the middleware ID.
+func (Factory) ID() string { return ID }
+
+// New decodes the raw JSON config and returns a ready Middleware. An
+// empty / null / empty-object payload yields a zero-value Config.
+func (Factory) New(rawConfig []byte) (middleware.Middleware, error) {
+	cfg := Config{}
+	if len(rawConfig) > 0 && !isEmptyJSON(rawConfig) {
+		if err := json.Unmarshal(rawConfig, &cfg); err != nil {
+			return nil, fmt.Errorf("decode config: %w", err)
+		}
+	}
+	return New(cfg), nil
+}
+
+// isEmptyJSON reports whether the payload is whitespace, null, or an
+// empty object/array. The caller skips Unmarshal in that case so the
+// zero-value Config flows through unchanged.
+func isEmptyJSON(raw []byte) bool {
+	trimmed := strings.TrimSpace(string(raw))
+	switch trimmed {
+	case "", "null", "{}", "[]":
+		return true
+	}
+	return false
+}
+
+// normaliseConfig lowercases and trims allowlist entries so the runtime
+// match is case-insensitive. Empty entries are dropped.
+func normaliseConfig(cfg Config) Config {
+	if len(cfg.ModelAllowlist) == 0 {
+		return cfg
+	}
+	cleaned := make([]string, 0, len(cfg.ModelAllowlist))
+	for _, entry := range cfg.ModelAllowlist {
+		n := normaliseModel(entry)
+		if n == "" {
+			continue
+		}
+		cleaned = append(cleaned, n)
+	}
+	cfg.ModelAllowlist = cleaned
+	return cfg
+}
+
+// normaliseModel lowercases and trims a single model identifier.
+func normaliseModel(model string) string {
+	return strings.ToLower(strings.TrimSpace(model))
+}
+
+func init() {
+	builtin.Register(Factory{})
+}

proxy/internal/middleware/builtin/llm_guardrail/middleware.go

@@ -0,0 +1,205 @@
+// Package llm_guardrail implements the SlotOnRequest middleware that
+// enforces the per-target LLM guardrail policy: a model allowlist
+// check and an opt-in prompt-capture step that may run a PII redactor
+// before emitting the prompt into the metadata bag.
+//
+// The middleware runs after llm_request_parser, which is responsible
+// for extracting the model and raw prompt onto the metadata side
+// channel. llm_guardrail consumes those keys, decides allow/deny, and
+// emits its own decision metadata plus the optional redacted prompt.
+package llm_guardrail
+
+import (
+	"context"
+	"unicode/utf8"
+
+	"github.com/netbirdio/netbird/proxy/internal/middleware"
+)
+
+// ID is the registry key for this middleware.
+const ID = "llm_guardrail"
+
+const (
+	version          = "1.0.0"
+	maxPromptBytes   = 3500
+	denyCodeModel    = "llm_policy.model_blocked"
+	denyReasonModel  = "model_blocked"
+	denyMessageModel = "model is not in the policy allowlist"
+	// Deny reason used when an allowlist is configured but the request model
+	// could not be determined. URL/path-routed providers (AWS Bedrock, Google
+	// Vertex, ...) carry the model outside the JSON body, so a request shape the
+	// parser does not recognise reaches the guardrail with no model. Such a
+	// request must be denied (fail closed), never waved through.
+	denyCodeModelUnknown    = "llm_policy.model_unknown"
+	denyReasonModelUnknown  = "model_unknown"
+	denyMessageModelUnknown = "request model could not be determined for the policy allowlist"
+)
+
+// Middleware enforces the model allowlist and optionally captures the
+// request prompt with PII redaction.
+type Middleware struct {
+	cfg Config
+}
+
+// New constructs a Middleware with the supplied configuration. Model
+// allowlist entries are normalised so the runtime check is
+// case-insensitive and trim-tolerant.
+func New(cfg Config) *Middleware {
+	return &Middleware{cfg: normaliseConfig(cfg)}
+}
+
+// ID returns the registry identifier.
+func (m *Middleware) ID() string { return ID }
+
+// Version returns the implementation version.
+func (m *Middleware) Version() string { return version }
+
+// Slot reports the chain slot the middleware lives in.
+func (m *Middleware) Slot() middleware.Slot { return middleware.SlotOnRequest }
+
+// AcceptedContentTypes lists the request body content types the
+// middleware needs. Guardrail consumes metadata produced upstream and
+// does not touch the body itself, but we keep application/json so the
+// body policy retains the parsed payload upstream when required.
+func (m *Middleware) AcceptedContentTypes() []string {
+	return []string{"application/json"}
+}
+
+// MetadataKeys is the closed set of metadata keys this middleware may
+// emit. The accumulator drops anything outside this allowlist.
+func (m *Middleware) MetadataKeys() []string {
+	return []string{
+		middleware.KeyLLMPolicyDecision,
+		middleware.KeyLLMPolicyReason,
+		middleware.KeyLLMRequestPrompt,
+	}
+}
+
+// MutationsSupported reports whether the middleware emits header / body
+// mutations. Guardrail never mutates the request.
+func (m *Middleware) MutationsSupported() bool { return false }
+
+// Invoke runs the policy. The model allowlist is the only deny path;
+// prompt capture only affects the metadata emitted alongside an allow.
+func (m *Middleware) Invoke(_ context.Context, in *middleware.Input) (*middleware.Output, error) {
+	model, modelPresent := lookupMetadata(in.Metadata, middleware.KeyLLMModel)
+
+	if denial := m.evaluateAllowlist(model, modelPresent); denial != nil {
+		return denial, nil
+	}
+
+	out := &middleware.Output{
+		Decision: middleware.DecisionAllow,
+		Metadata: []middleware.KV{
+			{Key: middleware.KeyLLMPolicyDecision, Value: "allow"},
+			{Key: middleware.KeyLLMPolicyReason, Value: ""},
+		},
+	}
+
+	if prompt, ok := m.capturePrompt(in.Metadata); ok {
+		out.Metadata = append(out.Metadata, middleware.KV{
+			Key:   middleware.KeyLLMRequestPrompt,
+			Value: prompt,
+		})
+	}
+
+	return out, nil
+}
+
+// Close releases resources owned by the middleware. Stateless, so this
+// is a no-op.
+func (m *Middleware) Close() error { return nil }
+
+// evaluateAllowlist returns a deny Output when the configured allowlist
+// rejects the model. A nil return means the request should proceed.
+func (m *Middleware) evaluateAllowlist(model string, modelPresent bool) *middleware.Output {
+	if len(m.cfg.ModelAllowlist) == 0 {
+		return nil
+	}
+	// Fail closed: with an allowlist configured, a request whose model the
+	// upstream parser could not extract (absent or empty) must be denied rather
+	// than allowed. This is what enforces the allowlist for URL/path-routed
+	// providers (Bedrock, Vertex, ...) whose model lives outside the JSON body.
+	if !modelPresent || normaliseModel(model) == "" {
+		return denyModel("", denyCodeModelUnknown, denyMessageModelUnknown, denyReasonModelUnknown)
+	}
+	if m.modelInAllowlist(model) {
+		return nil
+	}
+	return denyModel(model, denyCodeModel, denyMessageModel, denyReasonModel)
+}
+
+// denyModel builds a 403 deny Output for a model-allowlist rejection. model is
+// included in the details only when non-empty.
+func denyModel(model, code, message, reason string) *middleware.Output {
+	details := map[string]string{}
+	if model != "" {
+		details["model"] = model
+	}
+	return &middleware.Output{
+		Decision:   middleware.DecisionDeny,
+		DenyStatus: 403,
+		DenyReason: &middleware.DenyReason{
+			Code:    code,
+			Message: message,
+			Details: details,
+		},
+		Metadata: []middleware.KV{
+			{Key: middleware.KeyLLMPolicyDecision, Value: "deny"},
+			{Key: middleware.KeyLLMPolicyReason, Value: reason},
+		},
+	}
+}
+
+// modelInAllowlist reports whether the model matches any allowlist
+// entry under the case-insensitive, trim-tolerant comparison rule.
+func (m *Middleware) modelInAllowlist(model string) bool {
+	normalised := normaliseModel(model)
+	if normalised == "" {
+		return false
+	}
+	for _, allowed := range m.cfg.ModelAllowlist {
+		if allowed == normalised {
+			return true
+		}
+	}
+	return false
+}
+
+// capturePrompt returns the prompt to emit and whether it should be
+// emitted at all. The truncation guarantee is upheld here regardless of
+// whether redaction grew the string.
+func (m *Middleware) capturePrompt(meta []middleware.KV) (string, bool) {
+	if !m.cfg.PromptCapture.Enabled {
+		return "", false
+	}
+	raw, ok := lookupMetadata(meta, middleware.KeyLLMRequestPromptRaw)
+	if !ok {
+		return "", false
+	}
+	prompt := raw
+	if m.cfg.PromptCapture.RedactPii {
+		prompt = redactPII(prompt)
+	}
+	if len(prompt) > maxPromptBytes {
+		// Back off to a UTF-8 rune boundary so we never emit a string
+		// split mid-rune.
+		cut := maxPromptBytes
+		for cut > 0 && !utf8.RuneStart(prompt[cut]) {
+			cut--
+		}
+		prompt = prompt[:cut]
+	}
+	return prompt, true
+}
+
+// lookupMetadata finds the first KV with the given key. Returns the
+// value and true when present; the empty string and false otherwise.
+func lookupMetadata(meta []middleware.KV, key string) (string, bool) {
+	for _, kv := range meta {
+		if kv.Key == key {
+			return kv.Value, true
+		}
+	}
+	return "", false
+}

proxy/internal/middleware/builtin/llm_guardrail/middleware_test.go

@@ -0,0 +1,250 @@
+package llm_guardrail
+
+import (
+	"context"
+	"encoding/json"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/proxy/internal/middleware"
+)
+
+func metaValue(t *testing.T, kvs []middleware.KV, key string) (string, bool) {
+	t.Helper()
+	for _, kv := range kvs {
+		if kv.Key == key {
+			return kv.Value, true
+		}
+	}
+	return "", false
+}
+
+func newInput(meta ...middleware.KV) *middleware.Input {
+	return &middleware.Input{Slot: middleware.SlotOnRequest, Metadata: meta}
+}
+
+func TestMiddlewareIdentity(t *testing.T) {
+	mw := New(Config{})
+	assert.Equal(t, ID, mw.ID(), "middleware ID must be llm_guardrail")
+	assert.Equal(t, "1.0.0", mw.Version(), "version must be 1.0.0")
+	assert.Equal(t, middleware.SlotOnRequest, mw.Slot(), "guardrail must run in SlotOnRequest")
+	assert.False(t, mw.MutationsSupported(), "guardrail must not mutate requests")
+	assert.Equal(t, []string{"application/json"}, mw.AcceptedContentTypes(), "guardrail accepts application/json bodies")
+	assert.Equal(t,
+		[]string{
+			middleware.KeyLLMPolicyDecision,
+			middleware.KeyLLMPolicyReason,
+			middleware.KeyLLMRequestPrompt,
+		},
+		mw.MetadataKeys(),
+		"metadata key allowlist must match the spec",
+	)
+	require.NoError(t, mw.Close())
+}
+
+func TestAllowlistEmptyAllowsAnyModel(t *testing.T) {
+	mw := New(Config{})
+	out, err := mw.Invoke(context.Background(), newInput(
+		middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
+	))
+	require.NoError(t, err)
+	require.NotNil(t, out)
+	assert.Equal(t, middleware.DecisionAllow, out.Decision, "empty allowlist must allow any model")
+	v, ok := metaValue(t, out.Metadata, middleware.KeyLLMPolicyDecision)
+	require.True(t, ok, "decision metadata must be emitted")
+	assert.Equal(t, "allow", v, "decision must be allow")
+	r, ok := metaValue(t, out.Metadata, middleware.KeyLLMPolicyReason)
+	require.True(t, ok, "reason metadata must be emitted")
+	assert.Equal(t, "", r, "reason must be empty on allow")
+}
+
+func TestAllowlistMatchAllows(t *testing.T) {
+	mw := New(Config{ModelAllowlist: []string{"gpt-4o", "claude-opus-4"}})
+	out, err := mw.Invoke(context.Background(), newInput(
+		middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
+	))
+	require.NoError(t, err)
+	assert.Equal(t, middleware.DecisionAllow, out.Decision, "model in allowlist must be allowed")
+}
+
+func TestAllowlistMissDenies(t *testing.T) {
+	mw := New(Config{ModelAllowlist: []string{"gpt-4o"}})
+	out, err := mw.Invoke(context.Background(), newInput(
+		middleware.KV{Key: middleware.KeyLLMModel, Value: "claude-opus-4"},
+	))
+	require.NoError(t, err)
+	require.NotNil(t, out)
+	assert.Equal(t, middleware.DecisionDeny, out.Decision, "non-allowlisted model must be denied")
+	assert.Equal(t, 403, out.DenyStatus, "deny status must be 403")
+	require.NotNil(t, out.DenyReason, "deny reason must be populated")
+	assert.Equal(t, "llm_policy.model_blocked", out.DenyReason.Code, "deny code must match spec")
+	assert.Equal(t, "model is not in the policy allowlist", out.DenyReason.Message, "deny message must match spec")
+	assert.Equal(t, "claude-opus-4", out.DenyReason.Details["model"], "deny details must include the offending model")
+
+	dec, _ := metaValue(t, out.Metadata, middleware.KeyLLMPolicyDecision)
+	assert.Equal(t, "deny", dec, "decision metadata must be deny")
+	reason, _ := metaValue(t, out.Metadata, middleware.KeyLLMPolicyReason)
+	assert.Equal(t, "model_blocked", reason, "reason metadata must be model_blocked")
+}
+
+func TestAllowlistCaseInsensitive(t *testing.T) {
+	mw := New(Config{ModelAllowlist: []string{"  GPT-4o  ", "Claude-OPUS-4"}})
+	cases := []string{"gpt-4o", "GPT-4O", "  claude-opus-4 "}
+	for _, model := range cases {
+		out, err := mw.Invoke(context.Background(), newInput(
+			middleware.KV{Key: middleware.KeyLLMModel, Value: model},
+		))
+		require.NoError(t, err)
+		assert.Equal(t, middleware.DecisionAllow, out.Decision, "case/whitespace variants must match: %q", model)
+	}
+}
+
+func TestAllowlistMissingModelKeyDenies(t *testing.T) {
+	// Fail closed: with an allowlist configured, a request whose model the
+	// parser could not extract (URL/path-routed providers such as Bedrock or
+	// Vertex whose shape wasn't recognised) must be denied, not allowed.
+	mw := New(Config{ModelAllowlist: []string{"gpt-4o"}})
+	out, err := mw.Invoke(context.Background(), newInput())
+	require.NoError(t, err)
+	require.NotNil(t, out)
+	assert.Equal(t, middleware.DecisionDeny, out.Decision, "absent model must be denied when an allowlist is set")
+	assert.Equal(t, 403, out.DenyStatus, "deny status must be 403")
+	require.NotNil(t, out.DenyReason, "deny reason must be populated")
+	assert.Equal(t, "llm_policy.model_unknown", out.DenyReason.Code, "deny code must be model_unknown")
+	dec, _ := metaValue(t, out.Metadata, middleware.KeyLLMPolicyDecision)
+	assert.Equal(t, "deny", dec, "decision must be deny when model key is absent")
+	reason, _ := metaValue(t, out.Metadata, middleware.KeyLLMPolicyReason)
+	assert.Equal(t, "model_unknown", reason, "reason metadata must be model_unknown")
+}
+
+func TestAllowlistEmptyModelValueDenies(t *testing.T) {
+	// A present-but-empty model is as undeterminable as an absent one.
+	mw := New(Config{ModelAllowlist: []string{"gpt-4o"}})
+	out, err := mw.Invoke(context.Background(), newInput(
+		middleware.KV{Key: middleware.KeyLLMModel, Value: "   "},
+	))
+	require.NoError(t, err)
+	require.NotNil(t, out)
+	assert.Equal(t, middleware.DecisionDeny, out.Decision, "empty model must be denied when an allowlist is set")
+	require.NotNil(t, out.DenyReason, "deny reason must be populated")
+	assert.Equal(t, "llm_policy.model_unknown", out.DenyReason.Code, "deny code must be model_unknown")
+}
+
+func TestAllowlistEmptyListAllowsMissingModel(t *testing.T) {
+	// Without an allowlist there is nothing to enforce, so a missing model is
+	// still allowed — the fail-closed rule only applies when a list is set.
+	mw := New(Config{})
+	out, err := mw.Invoke(context.Background(), newInput())
+	require.NoError(t, err)
+	assert.Equal(t, middleware.DecisionAllow, out.Decision, "no allowlist must allow even without a model")
+}
+
+func TestPromptCaptureDisabledEmitsNoPrompt(t *testing.T) {
+	mw := New(Config{})
+	out, err := mw.Invoke(context.Background(), newInput(
+		middleware.KV{Key: middleware.KeyLLMRequestPromptRaw, Value: "hello world"},
+	))
+	require.NoError(t, err)
+	_, ok := metaValue(t, out.Metadata, middleware.KeyLLMRequestPrompt)
+	assert.False(t, ok, "prompt must not be emitted when capture is disabled")
+}
+
+func TestPromptCaptureNoRedactionEmitsRaw(t *testing.T) {
+	mw := New(Config{PromptCapture: PromptCapture{Enabled: true}})
+	raw := "hello world from user@example.com"
+	out, err := mw.Invoke(context.Background(), newInput(
+		middleware.KV{Key: middleware.KeyLLMRequestPromptRaw, Value: raw},
+	))
+	require.NoError(t, err)
+	prompt, ok := metaValue(t, out.Metadata, middleware.KeyLLMRequestPrompt)
+	require.True(t, ok, "prompt must be emitted when capture is enabled")
+	assert.Equal(t, raw, prompt, "prompt must pass through unchanged when redaction is off")
+}
+
+func TestPromptCaptureWithRedactionRedacts(t *testing.T) {
+	mw := New(Config{PromptCapture: PromptCapture{Enabled: true, RedactPii: true}})
+	raw := "contact me at user@example.com or +14155551234"
+	out, err := mw.Invoke(context.Background(), newInput(
+		middleware.KV{Key: middleware.KeyLLMRequestPromptRaw, Value: raw},
+	))
+	require.NoError(t, err)
+	prompt, ok := metaValue(t, out.Metadata, middleware.KeyLLMRequestPrompt)
+	require.True(t, ok, "prompt must be emitted when capture is enabled")
+	assert.Contains(t, prompt, "[REDACTED:email]", "email must be redacted")
+	assert.Contains(t, prompt, "[REDACTED:phone]", "phone must be redacted")
+	assert.NotContains(t, prompt, "user@example.com", "raw email must not leak")
+}
+
+func TestPromptCaptureRedactionTruncatesIfGrows(t *testing.T) {
+	mw := New(Config{PromptCapture: PromptCapture{Enabled: true, RedactPii: true}})
+	body := strings.Repeat("a", maxPromptBytes-10) + " user@example.com"
+	out, err := mw.Invoke(context.Background(), newInput(
+		middleware.KV{Key: middleware.KeyLLMRequestPromptRaw, Value: body},
+	))
+	require.NoError(t, err)
+	prompt, ok := metaValue(t, out.Metadata, middleware.KeyLLMRequestPrompt)
+	require.True(t, ok, "prompt must be emitted when capture is enabled")
+	assert.LessOrEqual(t, len(prompt), maxPromptBytes, "prompt must be truncated to maxPromptBytes")
+}
+
+func TestPromptCaptureMissingRawNoEmit(t *testing.T) {
+	mw := New(Config{PromptCapture: PromptCapture{Enabled: true, RedactPii: true}})
+	out, err := mw.Invoke(context.Background(), newInput())
+	require.NoError(t, err)
+	_, ok := metaValue(t, out.Metadata, middleware.KeyLLMRequestPrompt)
+	assert.False(t, ok, "prompt must not be emitted when raw key is missing")
+}
+
+func TestFactoryAcceptsZeroConfigs(t *testing.T) {
+	cases := map[string][]byte{
+		"nil":         nil,
+		"empty":       []byte(""),
+		"whitespace":  []byte("   \n  "),
+		"null":        []byte("null"),
+		"emptyObject": []byte("{}"),
+	}
+	f := Factory{}
+	for name, raw := range cases {
+		mw, err := f.New(raw)
+		require.NoError(t, err, "case %s must yield a zero-value config", name)
+		require.NotNil(t, mw)
+		assert.Equal(t, ID, mw.ID(), "case %s must build a guardrail middleware", name)
+	}
+}
+
+func TestFactoryDecodesValidConfig(t *testing.T) {
+	cfg := Config{
+		ModelAllowlist: []string{"gpt-4o"},
+		PromptCapture:  PromptCapture{Enabled: true, RedactPii: true},
+	}
+	raw, err := json.Marshal(cfg)
+	require.NoError(t, err, "marshalling test config must succeed")
+	mw, err := Factory{}.New(raw)
+	require.NoError(t, err)
+	require.NotNil(t, mw)
+}
+
+func TestFactoryRejectsMalformedJSON(t *testing.T) {
+	mw, err := Factory{}.New([]byte("{not-json"))
+	assert.Error(t, err, "malformed JSON must surface as a factory error")
+	assert.Nil(t, mw, "no middleware must be returned on malformed config")
+}
+
+func TestFactoryNormalisesAllowlist(t *testing.T) {
+	raw := []byte(`{"model_allowlist":["  GPT-4o  ","",""," Claude-3 "]}`)
+	mw, err := Factory{}.New(raw)
+	require.NoError(t, err)
+	out, err := mw.Invoke(context.Background(), newInput(
+		middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
+	))
+	require.NoError(t, err)
+	assert.Equal(t, middleware.DecisionAllow, out.Decision, "factory must lowercase + trim allowlist entries")
+	out2, err := mw.Invoke(context.Background(), newInput(
+		middleware.KV{Key: middleware.KeyLLMModel, Value: "claude-3"},
+	))
+	require.NoError(t, err)
+	assert.Equal(t, middleware.DecisionAllow, out2.Decision, "trimmed entry must still match")
+}

proxy/internal/middleware/builtin/llm_guardrail/redact.go

@@ -0,0 +1,75 @@
+package llm_guardrail
+
+import (
+	"regexp"
+	"strings"
+
+	"github.com/netbirdio/netbird/proxy/internal/middleware"
+)
+
+// PII redactor scope: redact prompt content BEFORE it lands in the metadata
+// bag. The bearer-with-keyword pass runs first so the keyword is preserved.
+// We then chain the package-level middleware.Scan to pick up PEM, JWT, AWS
+// access keys, generic bearer tokens (40+ chars), and Luhn-validated credit
+// cards — keeping prompt redaction in sync with metadata-value scanning. Email,
+// SSN (dashed form), phone (E.164 + NA), and IPv4 are prompt-shaped patterns
+// the metadata scanner intentionally leaves alone.
+var (
+	emailRegex = regexp.MustCompile(`[A-Za-z0-9._%+\-]+@[A-Za-z0-9.\-]+\.[A-Za-z]{2,}`)
+	ssnRegex   = regexp.MustCompile(`\b\d{3}-\d{2}-\d{4}\b`)
+	phoneE164  = regexp.MustCompile(`\+\d{8,15}\b`)
+	// phoneNARgx accepts the 3-3-4 North-American shape with any of the common
+	// separators (space, dot, dash, slash) or none at all between the area code
+	// and the body. The optional `\(?...\)?` wraps the area code; the separator
+	// classes use `*` (not `?`) so multi-char separators ("(202) " followed by
+	// space-and-something) and zero-separator runs ("2025550134") both match.
+	// False-positive tradeoff: 10 consecutive digits in a prompt will be
+	// treated as a phone number. For PII redaction that is the correct way to
+	// err — under-redaction leaks; over-redaction is annoying.
+	phoneNARgx  = regexp.MustCompile(`\(?\b\d{3}\)?[\s.\-/]*\d{3}[\s.\-/]*\d{4}\b`)
+	bearerRegex = regexp.MustCompile(`(?i)\b(bearer|token|api[_-]?key|authorization)([\s:=]+)(\S{20,})`)
+	ipv4Regex   = regexp.MustCompile(`\b(?:(?:25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)\.){3}(?:25[0-5]|2[0-4]\d|1\d\d|[1-9]?\d)\b`)
+)
+
+// redactPII is the package-private alias kept for internal callers; new code
+// outside the guardrail middleware should call RedactPII.
+func redactPII(value string) string { return RedactPII(value) }
+
+// RedactPII replaces high-signal PII patterns in value with
+// `[REDACTED:<kind>]`. Non-matching input is returned unchanged. Exported so
+// the request / response parsers can reuse the same coverage on raw prompts
+// and completions when the account's redact_pii toggle is on.
+func RedactPII(value string) string {
+	if value == "" {
+		return value
+	}
+	result := value
+	// Keyword-preserving bearer first so the "bearer "/"token=" prefix survives
+	// before the generic scanner gets at the same content.
+	result = bearerRegex.ReplaceAllStringFunc(result, redactBearer)
+	// Structured secrets shared with metadata-value scanning: PEM, JWT, AWS
+	// keys, generic bearer (40+), and Luhn-validated credit cards.
+	result = middleware.Scan(result)
+	// Prompt-shaped PII the metadata scanner doesn't cover.
+	result = emailRegex.ReplaceAllString(result, "[REDACTED:email]")
+	result = ssnRegex.ReplaceAllString(result, "[REDACTED:ssn]")
+	result = phoneE164.ReplaceAllString(result, "[REDACTED:phone]")
+	result = phoneNARgx.ReplaceAllString(result, "[REDACTED:phone]")
+	result = ipv4Regex.ReplaceAllString(result, "[REDACTED:ip]")
+	return result
+}
+
+// redactBearer keeps the leading keyword and its separator, replacing
+// only the secret payload so the surrounding context is preserved.
+func redactBearer(match string) string {
+	sub := bearerRegex.FindStringSubmatch(match)
+	if len(sub) < 4 {
+		return "[REDACTED:bearer]"
+	}
+	var b strings.Builder
+	b.Grow(len(sub[1]) + len(sub[2]) + len("[REDACTED:bearer]"))
+	b.WriteString(sub[1])
+	b.WriteString(sub[2])
+	b.WriteString("[REDACTED:bearer]")
+	return b.String()
+}

proxy/internal/middleware/builtin/llm_guardrail/redact_test.go

@@ -0,0 +1,217 @@
+package llm_guardrail
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRedactPIIEmptyInput(t *testing.T) {
+	assert.Equal(t, "", redactPII(""), "empty input must round-trip unchanged")
+}
+
+func TestRedactPIIPlainTextUntouched(t *testing.T) {
+	in := "the quick brown fox jumps over the lazy dog"
+	assert.Equal(t, in, redactPII(in), "non-PII text must pass through unchanged")
+}
+
+func TestRedactPIIEmail(t *testing.T) {
+	cases := []string{
+		"contact user@example.com today",
+		"first.last+tag@sub.example.co",
+		"USER_42@EXAMPLE.COM",
+	}
+	for _, in := range cases {
+		out := redactPII(in)
+		assert.Contains(t, out, "[REDACTED:email]", "email must be redacted in %q", in)
+		assert.NotContains(t, strings.ToLower(out), "@example", "raw email host must not survive in %q", in)
+	}
+}
+
+func TestRedactPIISSN(t *testing.T) {
+	in := "ssn 123-45-6789 should be hidden"
+	out := redactPII(in)
+	assert.Contains(t, out, "[REDACTED:ssn]", "SSN must be redacted")
+	assert.NotContains(t, out, "123-45-6789", "raw SSN must not survive")
+}
+
+func TestRedactPIIPhoneE164(t *testing.T) {
+	in := "call me at +14155551234 anytime"
+	out := redactPII(in)
+	assert.Contains(t, out, "[REDACTED:phone]", "E.164 phone must be redacted")
+	assert.NotContains(t, out, "+14155551234", "raw E.164 phone must not survive")
+}
+
+func TestRedactPIIPhoneNorthAmerican(t *testing.T) {
+	cases := []string{
+		"call (415) 555-1234 now",
+		"call 415-555-1234 now",
+		"call 415.555.1234 now",
+		"call 415 555 1234 now",
+	}
+	for _, in := range cases {
+		out := redactPII(in)
+		assert.Contains(t, out, "[REDACTED:phone]", "NA phone must be redacted in %q", in)
+		assert.NotContains(t, out, "555-1234", "raw NA phone must not survive in %q", in)
+	}
+}
+
+func TestRedactPIIBearerKeepsKeyword(t *testing.T) {
+	cases := []struct {
+		in      string
+		keyword string
+	}{
+		{"Authorization: Bearer abcdefghijklmnopqrstuvwxyz0123", "Bearer"},
+		{"token = abcdefghijklmnopqrstuvwxyz", "token"},
+		{"api_key=abcdefghijklmnopqrstuvwxyz0123", "api_key"},
+		{"API-KEY: abcdefghijklmnopqrstuvwxyz0123", "API-KEY"},
+		{"authorization: abcdefghijklmnopqrstuvwxyz0123", "authorization"},
+	}
+	for _, tc := range cases {
+		out := redactPII(tc.in)
+		assert.Contains(t, out, "[REDACTED:bearer]", "bearer-style secret must be redacted in %q", tc.in)
+		assert.Contains(t, out, tc.keyword, "leading keyword %q must be preserved in %q", tc.keyword, tc.in)
+		assert.NotContains(t, out, "abcdefghijklmnopqrstuvwxyz0123", "raw bearer payload must not survive in %q", tc.in)
+	}
+}
+
+func TestRedactPIIBearerShortValueUntouched(t *testing.T) {
+	in := "token=short"
+	out := redactPII(in)
+	assert.Equal(t, in, out, "short bearer-style values must not be redacted")
+}
+
+func TestRedactPIICombined(t *testing.T) {
+	in := "email user@example.com phone +14155551234 ssn 123-45-6789 token abcdefghijklmnopqrstuvwxyz0123"
+	out := redactPII(in)
+	assert.Contains(t, out, "[REDACTED:email]", "email must be redacted in combined input")
+	assert.Contains(t, out, "[REDACTED:phone]", "phone must be redacted in combined input")
+	assert.Contains(t, out, "[REDACTED:ssn]", "SSN must be redacted in combined input")
+	assert.Contains(t, out, "[REDACTED:bearer]", "bearer must be redacted in combined input")
+	assert.NotContains(t, out, "user@example.com", "raw email must not survive combined input")
+	assert.NotContains(t, out, "+14155551234", "raw phone must not survive combined input")
+	assert.NotContains(t, out, "123-45-6789", "raw SSN must not survive combined input")
+}
+
+func TestRedactPIICreditCard(t *testing.T) {
+	// 4242424242424242 is a well-known Stripe test number (Visa, Luhn-valid).
+	cases := []string{
+		"please charge 4242424242424242 now",
+		"card: 4242-4242-4242-4242",
+		"4242 4242 4242 4242 expires 12/30",
+	}
+	for _, in := range cases {
+		out := redactPII(in)
+		assert.Contains(t, out, "[REDACTED:cc]", "Luhn-valid credit card must be redacted in %q", in)
+		assert.NotContains(t, out, "4242424242424242", "raw card digits must not survive in %q", in)
+		assert.NotContains(t, out, "4242-4242-4242-4242", "raw dashed card must not survive in %q", in)
+	}
+}
+
+func TestRedactPIIIPv4(t *testing.T) {
+	cases := []string{
+		"connect to 10.0.42.7 over the tunnel",
+		"server 192.168.1.100 down",
+		"public address 203.0.113.42 was hit",
+	}
+	for _, in := range cases {
+		out := redactPII(in)
+		assert.Contains(t, out, "[REDACTED:ip]", "IPv4 must be redacted in %q", in)
+	}
+}
+
+func TestRedactPIIJWT(t *testing.T) {
+	// No "token "/"bearer " prefix here, so the bearer-with-keyword pass leaves
+	// it alone and the JWT pattern from middleware.Scan must catch it.
+	in := "session eyJhbGciOiJIUzI1NiJ9.eyJzdWIiOiJ1c2VyXzQyIn0.signaturepart expires soon"
+	out := redactPII(in)
+	assert.Contains(t, out, "[REDACTED:jwt]", "JWT must be redacted when no bearer keyword precedes it")
+	assert.NotContains(t, out, "eyJhbGciOiJIUzI1NiJ9", "raw JWT header must not survive")
+}
+
+func TestRedactPIIAWSAccessKey(t *testing.T) {
+	in := "the key AKIAIOSFODNN7EXAMPLE belongs to test user"
+	out := redactPII(in)
+	assert.Contains(t, out, "[REDACTED:aws_key]", "AWS access key must be redacted")
+	assert.NotContains(t, out, "AKIAIOSFODNN7EXAMPLE", "raw AWS key must not survive")
+}
+
+func TestRedactPIIPlainNumbersUntouched(t *testing.T) {
+	// 1234567890123 is 13 digits but fails Luhn; must NOT trip the CC redactor.
+	// We use a 13-digit value (the CC-candidate range starts at 13) so the only
+	// risk is the CC pattern firing. Phone redaction is 10-digit by design and
+	// would catch 1234567890123 as a phone — that's expected and not what this
+	// test guards against.
+	in := "order number 1234567890123 is queued"
+	out := redactPII(in)
+	assert.NotContains(t, out, "[REDACTED:cc]", "non-Luhn digit sequences must not be redacted as credit cards")
+}
+
+// piiFixture mirrors the user-supplied test fixture: each record carries one
+// email, one SSN, and one phone in a representative format. The test asserts
+// that EVERY raw token disappears after redaction and the right [REDACTED:*]
+// markers show up. Names are kept in the input and must survive — names are
+// not a pattern the redactor tries to catch.
+type piiFixture struct {
+	name  string // person name (must survive redaction)
+	email string
+	ssn   string
+	phone string
+}
+
+var fixtureRecords = []piiFixture{
+	{"Alice Johnson", "alice.johnson@example.com", "123-45-6789", "(202) 555-0147"},
+	{"Brian Smith", "brian.smith@example.org", "987-65-4321", "202-555-0163"},
+	{"Carla Nguyen", "c.nguyen@test.local", "111-22-3333", "+1-202-555-0188"},
+	{"David Martinez", "david.martinez@example.com", "222-33-4444", "202.555.0199"},
+	{"Evelyn Parker", "evelyn.parker@example.org", "333-44-5555", "1-202-555-0112"},
+	{"Frank O'Connor", "frank.oconnor@test.local", "444-55-6666", "2025550134"},
+	{"Grace Lee", "grace.lee@example.com", "555-66-7777", "(202)555-0156"},
+	{"Hassan Ali", "hassan.ali@example.org", "666-77-8888", "+1 (202) 555-0175"},
+	{"Isabella Rossi", "i.rossi@test.local", "777-88-9999", "202 555 0121"},
+	{"Jamal Thompson", "jamal.thompson@example.com", "888-99-0001", "202/555/0108"},
+}
+
+// TestRedactPII_FixtureRecord drives every record through redactPII and
+// asserts the email, SSN, and phone are all redacted, the name survives, and
+// the appropriate REDACTED markers are present. This is the spec the redactor
+// must meet for the kind of prompts operators throw at it.
+func TestRedactPII_FixtureRecord(t *testing.T) {
+	for _, rec := range fixtureRecords {
+		t.Run(rec.name, func(t *testing.T) {
+			in := "Name: " + rec.name + "\n   Email: " + rec.email + "\n   SSN: " + rec.ssn + "\n   Phone: " + rec.phone
+			out := redactPII(in)
+
+			assert.Contains(t, out, rec.name, "name must survive (not a PII pattern the redactor catches)")
+			assert.Contains(t, out, "[REDACTED:email]", "email marker must appear for %q", rec.email)
+			assert.Contains(t, out, "[REDACTED:ssn]", "ssn marker must appear for %q", rec.ssn)
+			assert.Contains(t, out, "[REDACTED:phone]", "phone marker must appear for %q", rec.phone)
+
+			assert.NotContains(t, out, rec.email, "raw email must not survive: %q", rec.email)
+			assert.NotContains(t, out, rec.ssn, "raw SSN must not survive: %q", rec.ssn)
+			// Phone: assert the local digits (last 7) are gone. Country-code
+			// remnants like "+1 " or "1-" may remain in front of the redaction
+			// because the E.164 pattern needs digits-only after '+' — that's
+			// acceptable, the personally-identifying portion is removed.
+			localDigits := lastSevenDigits(rec.phone)
+			assert.NotContains(t, out, localDigits, "raw phone local digits %q must not survive in redacted output of %q", localDigits, rec.phone)
+		})
+	}
+}
+
+// lastSevenDigits returns the last 7 digits of a phone number, ignoring
+// formatting. It's the unique "subscriber" portion that absolutely must be
+// scrubbed regardless of which prefix the redactor leaves behind.
+func lastSevenDigits(phone string) string {
+	digits := make([]byte, 0, len(phone))
+	for i := 0; i < len(phone); i++ {
+		if phone[i] >= '0' && phone[i] <= '9' {
+			digits = append(digits, phone[i])
+		}
+	}
+	if len(digits) <= 7 {
+		return string(digits)
+	}
+	return string(digits[len(digits)-7:])
+}

proxy/internal/middleware/builtin/llm_identity_inject/factory.go

@@ -0,0 +1,113 @@
+package llm_identity_inject
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"strings"
+
+	"github.com/netbirdio/netbird/proxy/internal/middleware"
+	"github.com/netbirdio/netbird/proxy/internal/middleware/builtin"
+)
+
+// ProviderInjection describes one resolved provider's injection rule.
+// Identity stamping uses one of HeaderPair / JSONMetadata; ExtraHeaders
+// is independent — each entry is a static (operator-configured) header
+// stamped on every matching request with anti-spoof. A rule with no
+// shape AND no extras is dropped at New() time as a no-op.
+type ProviderInjection struct {
+	// ProviderID is the resolved provider id — matches the value
+	// llm_router stamps under KeyLLMResolvedProviderID.
+	ProviderID string `json:"provider_id"`
+	// HeaderPair is the LiteLLM-style wire convention: separate
+	// headers for end-user id and tags CSV.
+	HeaderPair *HeaderPairRule `json:"header_pair,omitempty"`
+	// JSONMetadata is the Portkey-style wire convention: a single
+	// header carrying a JSON object keyed by reserved field names.
+	JSONMetadata *JSONMetadataRule `json:"json_metadata,omitempty"`
+	// ExtraHeaders is an operator-configured list of static headers
+	// (e.g. "x-portkey-config: pc-...") that the middleware stamps
+	// on every matching request. The synth pre-resolves the values
+	// from the provider record's ExtraValues map; the middleware
+	// just emits them. Each name is also added to HeadersRemove for
+	// anti-spoof so a client can't smuggle their own value.
+	ExtraHeaders []ExtraHeaderKV `json:"extra_headers,omitempty"`
+}
+
+// ExtraHeaderKV is one static header entry the middleware stamps as-is.
+type ExtraHeaderKV struct {
+	Name  string `json:"name"`
+	Value string `json:"value"`
+}
+
+// HeaderPairRule emits identity through dedicated per-dimension
+// headers. The two *InBody flags layer body-level identity on top: when
+// TagsInBody is set the middleware also writes the tag list into the
+// request body's metadata.tags array (required for LiteLLM tag-budget
+// enforcement, which only inspects the body); when EndUserIDInBody is
+// set the display identity is also written into the body's top-level
+// "user" field (the OpenAI-standard end-user identifier — defense-in-
+// depth and anti-spoof on top of the header path).
+type HeaderPairRule struct {
+	EndUserIDHeader string `json:"end_user_id_header,omitempty"`
+	TagsHeader      string `json:"tags_header,omitempty"`
+	TagsInBody      bool   `json:"tags_in_body,omitempty"`
+	EndUserIDInBody bool   `json:"end_user_id_in_body,omitempty"`
+}
+
+// JSONMetadataRule emits identity through a single JSON-object header.
+// Empty UserKey/GroupsKey skip that dimension at emit time. When
+// MaxValueLength > 0 each emitted JSON value is truncated to that many
+// bytes — Portkey enforces 128 chars per value.
+type JSONMetadataRule struct {
+	Header         string `json:"header"`
+	UserKey        string `json:"user_key,omitempty"`
+	GroupsKey      string `json:"groups_key,omitempty"`
+	MaxValueLength int    `json:"max_value_length,omitempty"`
+	// Sanitize replaces characters outside the destination provider's accepted
+	// set with '_' before emitting each value. AWS Bedrock's
+	// X-Amzn-Bedrock-Request-Metadata restricts values to [A-Za-z0-9 +-=._:/@];
+	// group display names with other characters would otherwise 400.
+	Sanitize bool `json:"sanitize,omitempty"`
+}
+
+// Config is the on-wire configuration accepted by the factory. An
+// empty Providers slice yields a no-op middleware (every resolved
+// provider passes through unchanged).
+type Config struct {
+	Providers []ProviderInjection `json:"providers"`
+}
+
+// Factory builds llm_identity_inject instances from raw config bytes.
+type Factory struct{}
+
+// ID returns the registry identifier.
+func (Factory) ID() string { return ID }
+
+// New constructs a middleware instance. Empty, null, and {} configs
+// yield a no-op middleware. Non-empty payloads must parse cleanly so
+// misconfigurations surface at chain build time.
+func (Factory) New(rawConfig []byte) (middleware.Middleware, error) {
+	cfg := Config{}
+	if !isEmptyJSON(rawConfig) {
+		if err := json.Unmarshal(rawConfig, &cfg); err != nil {
+			return nil, fmt.Errorf("decode config: %w", err)
+		}
+	}
+	return New(cfg), nil
+}
+
+// isEmptyJSON reports whether the payload is whitespace, null, or an
+// empty object/array.
+func isEmptyJSON(raw []byte) bool {
+	trimmed := strings.TrimSpace(string(bytes.TrimSpace(raw)))
+	switch trimmed {
+	case "", "null", "{}", "[]":
+		return true
+	}
+	return false
+}
+
+func init() {
+	builtin.Register(Factory{})
+}

proxy/internal/middleware/builtin/llm_identity_inject/middleware.go

@@ -0,0 +1,475 @@
+// Package llm_identity_inject implements the SlotOnRequest middleware
+// that stamps the caller's NetBird identity onto upstream LLM-gateway
+// requests. It runs after llm_router (which resolves the provider) and
+// looks up the resolved provider id against a per-account injection
+// table built by the synthesiser from the catalog's IdentityInjection
+// metadata.
+//
+// Two wire shapes are supported, dispatched per-rule:
+//
+//   - HeaderPair (LiteLLM-style): separate end-user-id and tags
+//     headers; tags emitted as a CSV value.
+//   - JSONMetadata (Portkey-style): one header carrying a JSON
+//     object with reserved keys for user / groups; per-value byte
+//     length capped when the rule sets MaxValueLength.
+//
+// In both cases, identity comes from Input.UserEmail (peer-attached
+// user's email or peer.Name fallback) and groups come from the
+// authorising-groups intersection llm_router emitted (with
+// id→display-name translation via Input.UserGroups / UserGroupNames
+// positional pairing). HeadersRemove runs before HeadersAdd in the
+// framework, so a client can never spoof identity by stamping these
+// headers themselves.
+package llm_identity_inject
+
+import (
+	"context"
+	"encoding/json"
+	"sort"
+	"strings"
+
+	"github.com/netbirdio/netbird/proxy/internal/middleware"
+)
+
+// ID is the registry identifier for this middleware.
+const ID = "llm_identity_inject"
+
+// Version is reported via Middleware.Version().
+const Version = "1.0.0"
+
+// Middleware stamps NetBird identity onto upstream requests for the
+// configured set of resolved providers.
+type Middleware struct {
+	cfg  Config
+	byID map[string]ProviderInjection
+}
+
+// New constructs a Middleware from the supplied configuration. A nil
+// or empty Providers slice yields a no-op middleware.
+func New(cfg Config) *Middleware {
+	byID := make(map[string]ProviderInjection, len(cfg.Providers))
+	for _, p := range cfg.Providers {
+		if p.ProviderID == "" || !injectionEmitsAnything(p) {
+			continue
+		}
+		byID[p.ProviderID] = p
+	}
+	return &Middleware{cfg: cfg, byID: byID}
+}
+
+// injectionEmitsAnything reports whether a provider injection rule would
+// stamp anything at runtime. Rules that set both identity shapes are a
+// configuration error (we refuse to guess which wins), and rules that
+// resolve to no headers are dropped to keep the runtime check tight.
+// Non-empty extras alone keep a rule alive even when neither identity
+// shape is set.
+func injectionEmitsAnything(p ProviderInjection) bool {
+	hasExtras := false
+	for _, e := range p.ExtraHeaders {
+		if e.Name != "" && e.Value != "" {
+			hasExtras = true
+			break
+		}
+	}
+	switch {
+	case p.HeaderPair != nil && p.JSONMetadata != nil:
+		return false
+	case p.HeaderPair != nil:
+		return p.HeaderPair.EndUserIDHeader != "" || p.HeaderPair.TagsHeader != "" ||
+			p.HeaderPair.TagsInBody || p.HeaderPair.EndUserIDInBody || hasExtras
+	case p.JSONMetadata != nil:
+		if p.JSONMetadata.Header == "" {
+			return false
+		}
+		return p.JSONMetadata.UserKey != "" || p.JSONMetadata.GroupsKey != "" || hasExtras
+	default:
+		return hasExtras
+	}
+}
+
+// ID returns the registry identifier.
+func (m *Middleware) ID() string { return ID }
+
+// Version returns the implementation version.
+func (m *Middleware) Version() string { return Version }
+
+// Slot reports the chain slot the middleware lives in.
+func (m *Middleware) Slot() middleware.Slot { return middleware.SlotOnRequest }
+
+// AcceptedContentTypes returns nil — this middleware reads only
+// metadata and identity fields on the Input envelope.
+func (m *Middleware) AcceptedContentTypes() []string { return nil }
+
+// MetadataKeys is empty: the middleware emits no metadata. Identity
+// stamping is a header-only operation.
+func (m *Middleware) MetadataKeys() []string { return nil }
+
+// MutationsSupported reports that the middleware emits header
+// mutations on the Output envelope.
+func (m *Middleware) MutationsSupported() bool { return true }
+
+// Close releases resources owned by the middleware. Stateless, so
+// this is a no-op.
+func (m *Middleware) Close() error { return nil }
+
+// Invoke stamps identity headers when the resolved provider has an
+// injection rule. Always Allow.
+func (m *Middleware) Invoke(_ context.Context, in *middleware.Input) (*middleware.Output, error) {
+	out := &middleware.Output{Decision: middleware.DecisionAllow}
+	if len(m.byID) == 0 || in == nil {
+		return out, nil
+	}
+	resolved, ok := lookupMetadata(in.Metadata, middleware.KeyLLMResolvedProviderID)
+	if !ok || resolved == "" {
+		return out, nil
+	}
+	rule, ok := m.byID[resolved]
+	if !ok {
+		return out, nil
+	}
+
+	var mutations *middleware.Mutations
+	switch {
+	case rule.HeaderPair != nil:
+		mutations = applyHeaderPair(rule.HeaderPair, in)
+	case rule.JSONMetadata != nil:
+		mutations = applyJSONMetadata(rule.JSONMetadata, in)
+	}
+
+	// ExtraHeaders are independent of the identity shape. Stamp each
+	// non-empty entry with anti-spoof: Remove first (frame strips it
+	// before our Add lands) so a client can't smuggle a value, then
+	// Add our trusted one.
+	if len(rule.ExtraHeaders) > 0 {
+		if mutations == nil {
+			mutations = &middleware.Mutations{}
+		}
+		for _, h := range rule.ExtraHeaders {
+			if h.Name == "" || h.Value == "" {
+				continue
+			}
+			mutations.HeadersRemove = append(mutations.HeadersRemove, h.Name)
+			mutations.HeadersAdd = append(mutations.HeadersAdd, middleware.KV{
+				Key:   h.Name,
+				Value: h.Value,
+			})
+		}
+	}
+
+	if mutations == nil || (len(mutations.HeadersAdd) == 0 && len(mutations.HeadersRemove) == 0 && len(mutations.BodyReplace) == 0) {
+		return out, nil
+	}
+	out.Mutations = mutations
+	return out, nil
+}
+
+// applyHeaderPair builds the LiteLLM-style mutations: separate per-
+// dimension headers, with anti-spoof Removes paired with trusted Adds.
+func applyHeaderPair(rule *HeaderPairRule, in *middleware.Input) *middleware.Mutations {
+	mutations := &middleware.Mutations{}
+
+	if rule.EndUserIDHeader != "" {
+		mutations.HeadersRemove = append(mutations.HeadersRemove, rule.EndUserIDHeader)
+		// Prefer the email when the auth path carried it: gateways
+		// like LiteLLM key per-user budgets and dashboards on a
+		// human-readable identifier; the user_id is an opaque
+		// management-server primary key. Fall back to user_id when
+		// no email is available (non-OIDC schemes, legacy JWTs).
+		if identity := identityFor(in); identity != "" {
+			mutations.HeadersAdd = append(mutations.HeadersAdd, middleware.KV{
+				Key:   rule.EndUserIDHeader,
+				Value: identity,
+			})
+		}
+	}
+
+	if rule.TagsHeader != "" {
+		mutations.HeadersRemove = append(mutations.HeadersRemove, rule.TagsHeader)
+		if csv := authorisingTagsCSV(in); csv != "" {
+			mutations.HeadersAdd = append(mutations.HeadersAdd, middleware.KV{
+				Key:   rule.TagsHeader,
+				Value: csv,
+			})
+		}
+	}
+
+	if rule.TagsInBody || rule.EndUserIDInBody {
+		// Body-level identity unlocks gateway behaviour the header
+		// path can't reach (LiteLLM's _tag_max_budget_check only
+		// inspects the body; OpenAI direct only reads the body's
+		// "user" field for attribution). The header path stays
+		// intact, so we still get attribution + per-end-user budget
+		// gating when body inject can't run (truncated body,
+		// non-JSON, hostile metadata shape).
+		var bodyTags []string
+		if rule.TagsInBody {
+			bodyTags = authorisingTagsSlice(in)
+		}
+		var bodyUser string
+		if rule.EndUserIDInBody {
+			bodyUser = identityFor(in)
+		}
+		if newBody, ok := injectIntoBody(in, bodyTags, bodyUser); ok {
+			mutations.BodyReplace = newBody
+		}
+	}
+
+	return mutations
+}
+
+// injectIntoBody parses the request body and writes the supplied
+// identity dimensions into it. Tags land at metadata.tags (creating
+// the metadata object when absent); the user identity lands at the
+// top-level "user" field (OpenAI-standard end-user identifier).
+// Returns the re-marshaled body and ok=true when at least one field
+// was written. Returns ok=false (no mutation) when:
+//
+//   - both inputs are empty (nothing to write);
+//   - the body is empty or truncated (we don't have the full document
+//     to safely round-trip);
+//   - the body isn't a JSON object (skip silently — this middleware
+//     only knows how to inject into OpenAI-compatible JSON payloads).
+//
+// A non-object existing `metadata` field skips the tag write but
+// still allows the user write to land — we don't clobber the client's
+// non-object metadata, but the orthogonal user field is fair game.
+// The header path emission still runs in skip cases, so spend tracking
+// + header-resolved end-user budgets continue to work without body-
+// level enforcement.
+func injectIntoBody(in *middleware.Input, tags []string, userID string) ([]byte, bool) {
+	wantTags := len(tags) > 0
+	wantUser := userID != ""
+	if !wantTags && !wantUser {
+		return nil, false
+	}
+	if in == nil || len(in.Body) == 0 || in.BodyTruncated {
+		return nil, false
+	}
+	var doc map[string]any
+	if err := json.Unmarshal(in.Body, &doc); err != nil {
+		return nil, false
+	}
+	injected := false
+	if wantTags {
+		var meta map[string]any
+		if existing, ok := doc["metadata"]; ok {
+			if typed, isObject := existing.(map[string]any); isObject {
+				meta = typed
+			}
+			// non-object metadata: leave it; tags go unwritten so we
+			// don't clobber the client's value. Header fallback covers
+			// spend tracking.
+		} else {
+			meta = map[string]any{}
+		}
+		if meta != nil {
+			meta["tags"] = tags
+			doc["metadata"] = meta
+			injected = true
+		}
+	}
+	if wantUser {
+		// Anti-spoof: overwrite any client-supplied "user" so the
+		// gateway only sees our trusted identity.
+		doc["user"] = userID
+		injected = true
+	}
+	if !injected {
+		return nil, false
+	}
+	out, err := json.Marshal(doc)
+	if err != nil {
+		return nil, false
+	}
+	return out, true
+}
+
+// applyJSONMetadata builds the Portkey-style mutations: a single header
+// carrying a JSON object keyed by the rule's reserved field names. Per-
+// value byte length is capped at MaxValueLength when set (Portkey
+// enforces 128 chars).
+func applyJSONMetadata(rule *JSONMetadataRule, in *middleware.Input) *middleware.Mutations {
+	mutations := &middleware.Mutations{}
+	mutations.HeadersRemove = append(mutations.HeadersRemove, rule.Header)
+
+	emit := func(v string) string {
+		if rule.Sanitize {
+			v = sanitizeMetadataValue(v)
+		}
+		return truncate(v, rule.MaxValueLength)
+	}
+	payload := map[string]string{}
+	if rule.UserKey != "" {
+		if identity := identityFor(in); identity != "" {
+			payload[rule.UserKey] = emit(identity)
+		}
+	}
+	if rule.GroupsKey != "" {
+		if csv := authorisingTagsCSV(in); csv != "" {
+			payload[rule.GroupsKey] = emit(csv)
+		}
+	}
+	if len(payload) == 0 {
+		return mutations
+	}
+	raw, err := json.Marshal(payload)
+	if err != nil {
+		return mutations
+	}
+	mutations.HeadersAdd = append(mutations.HeadersAdd, middleware.KV{
+		Key:   rule.Header,
+		Value: string(raw),
+	})
+	return mutations
+}
+
+// identityFor returns the caller's display identity. UserEmail wins
+// (carries the user email when peer-attached, peer.Name otherwise);
+// UserID falls in only as a defensive last resort.
+func identityFor(in *middleware.Input) string {
+	if in.UserEmail != "" {
+		return in.UserEmail
+	}
+	return in.UserID
+}
+
+// authorisingTagsSlice returns the sorted, deduplicated slice of group
+// display names the request was authorised under. Prefers the per-
+// request authorising groups emitted by llm_router (intersection of the
+// caller's UserGroups with the resolved route's AllowedGroupIDs) so the
+// tags carry only the groups that actually authorise THIS request, not
+// every group the peer happens to be in. Falls back to the full
+// UserGroups when the router metadata key is absent.
+func authorisingTagsSlice(in *middleware.Input) []string {
+	ids := tagsIDsFromAuthorising(in.Metadata)
+	if len(ids) == 0 {
+		ids = in.UserGroups
+	}
+	return tagsNamedSlice(ids, in.UserGroups, in.UserGroupNames)
+}
+
+// authorisingTagsCSV is a convenience wrapper that joins
+// authorisingTagsSlice with commas for HeaderPair-style emission.
+func authorisingTagsCSV(in *middleware.Input) string {
+	return strings.Join(authorisingTagsSlice(in), ",")
+}
+
+// truncate caps s to maxBytes bytes when maxBytes > 0. No-op when
+// maxBytes <= 0 or s already fits. Truncation is byte-wise — sufficient
+// for Portkey's 128-char ASCII limit. UTF-8 sequences could in theory
+// be split, but the gateway treats the value as opaque bytes.
+func truncate(s string, maxBytes int) string {
+	if maxBytes <= 0 || len(s) <= maxBytes {
+		return s
+	}
+	return s[:maxBytes]
+}
+
+// sanitizeMetadataValue replaces any character outside AWS Bedrock's accepted
+// request-metadata class — letters, digits, space, and + - = . _ : / @ — with
+// '_'. This keeps values (notably the groups CSV, whose commas are rejected, and
+// group display names with arbitrary characters) from making Bedrock reject the
+// request with 400. The result stays opaque to the gateway.
+func sanitizeMetadataValue(s string) string {
+	var b strings.Builder
+	b.Grow(len(s))
+	for _, r := range s {
+		if metadataCharAllowed(r) {
+			b.WriteRune(r)
+		} else {
+			b.WriteByte('_')
+		}
+	}
+	return b.String()
+}
+
+func metadataCharAllowed(r rune) bool {
+	switch {
+	case r >= 'a' && r <= 'z', r >= 'A' && r <= 'Z', r >= '0' && r <= '9':
+		return true
+	}
+	switch r {
+	case ' ', '+', '-', '=', '.', '_', ':', '/', '@':
+		return true
+	}
+	return false
+}
+
+// tagsIDsFromAuthorising reads llm_router's authorising-groups metadata
+// (a CSV of group ids) and returns the parsed slice. Returns nil when
+// the key is absent or empty so the caller can fall back to the full
+// UserGroups.
+func tagsIDsFromAuthorising(meta []middleware.KV) []string {
+	v, ok := lookupMetadata(meta, middleware.KeyLLMAuthorisingGroups)
+	if !ok {
+		return nil
+	}
+	v = strings.TrimSpace(v)
+	if v == "" {
+		return nil
+	}
+	parts := strings.Split(v, ",")
+	out := make([]string, 0, len(parts))
+	for _, p := range parts {
+		p = strings.TrimSpace(p)
+		if p != "" {
+			out = append(out, p)
+		}
+	}
+	if len(out) == 0 {
+		return nil
+	}
+	return out
+}
+
+// tagsNamedSlice returns the sorted, deduplicated list of group display
+// names. ids carries the canonical group identifiers to emit;
+// userGroups + userGroupNames provide the positional id→name
+// translation table from the Input envelope. When a name is missing
+// for a given id (slice shorter than userGroups, or id absent from the
+// table), the id is used verbatim so the tag still attributes
+// correctly. Sorted so the same caller produces the same header value
+// across requests (helps gateway-side cache hits and log correlation).
+func tagsNamedSlice(ids, userGroups, userGroupNames []string) []string {
+	if len(ids) == 0 {
+		return nil
+	}
+	idToName := make(map[string]string, len(userGroups))
+	for i, id := range userGroups {
+		if i < len(userGroupNames) {
+			idToName[id] = userGroupNames[i]
+		}
+	}
+	seen := make(map[string]struct{}, len(ids))
+	out := make([]string, 0, len(ids))
+	for _, id := range ids {
+		id = strings.TrimSpace(id)
+		if id == "" {
+			continue
+		}
+		tag := idToName[id]
+		if tag == "" {
+			tag = id
+		}
+		if _, dup := seen[tag]; dup {
+			continue
+		}
+		seen[tag] = struct{}{}
+		out = append(out, tag)
+	}
+	if len(out) == 0 {
+		return nil
+	}
+	sort.Strings(out)
+	return out
+}
+
+// lookupMetadata returns the value for key plus a presence flag.
+func lookupMetadata(meta []middleware.KV, key string) (string, bool) {
+	for _, kv := range meta {
+		if kv.Key == key {
+			return kv.Value, true
+		}
+	}
+	return "", false
+}

proxy/internal/middleware/builtin/llm_identity_inject/middleware_test.go

@@ -0,0 +1,706 @@
+package llm_identity_inject
+
+import (
+	"context"
+	"encoding/json"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/proxy/internal/middleware"
+)
+
+const (
+	litellmProvider = "ainp_litellm-test"
+	portkeyProvider = "ainp_portkey-test"
+)
+
+func newInput(resolvedProvider, userID string, groups []string) *middleware.Input {
+	return &middleware.Input{
+		Slot:       middleware.SlotOnRequest,
+		AccountID:  "acct-test",
+		UserID:     userID,
+		UserGroups: groups,
+		SourceIP:   "100.64.0.5",
+		RequestID:  "req-1",
+		Metadata: []middleware.KV{
+			{Key: middleware.KeyLLMResolvedProviderID, Value: resolvedProvider},
+		},
+	}
+}
+
+func liteLLMRule() ProviderInjection {
+	return ProviderInjection{
+		ProviderID: litellmProvider,
+		HeaderPair: &HeaderPairRule{
+			EndUserIDHeader: "x-litellm-end-user-id",
+			TagsHeader:      "x-litellm-tags",
+		},
+	}
+}
+
+func TestMiddlewareIdentity(t *testing.T) {
+	mw := New(Config{})
+	assert.Equal(t, ID, mw.ID())
+	assert.Equal(t, Version, mw.Version())
+	assert.Equal(t, middleware.SlotOnRequest, mw.Slot())
+	assert.True(t, mw.MutationsSupported())
+	assert.Empty(t, mw.MetadataKeys(), "middleware emits no metadata")
+	assert.Nil(t, mw.AcceptedContentTypes())
+	require.NoError(t, mw.Close())
+}
+
+func TestInject_MatchedProvider_StampsHeaders(t *testing.T) {
+	mw := New(Config{Providers: []ProviderInjection{liteLLMRule()}})
+	in := newInput(litellmProvider, "alice", []string{"grp-eng", "grp-it"})
+
+	out, err := mw.Invoke(context.Background(), in)
+	require.NoError(t, err)
+	require.NotNil(t, out)
+	assert.Equal(t, middleware.DecisionAllow, out.Decision)
+	require.NotNil(t, out.Mutations)
+
+	// Strips the same headers we're about to add (anti-spoof).
+	assert.ElementsMatch(t,
+		[]string{"x-litellm-end-user-id", "x-litellm-tags"},
+		out.Mutations.HeadersRemove,
+		"every injected header must also appear in HeadersRemove so client-supplied values are wiped before our trusted values land")
+
+	added := map[string]string{}
+	for _, kv := range out.Mutations.HeadersAdd {
+		added[kv.Key] = kv.Value
+	}
+	assert.Equal(t, "alice", added["x-litellm-end-user-id"])
+	assert.Equal(t, "grp-eng,grp-it", added["x-litellm-tags"], "tags CSV must be sorted")
+}
+
+func TestInject_UnmatchedProvider_NoMutations(t *testing.T) {
+	mw := New(Config{Providers: []ProviderInjection{liteLLMRule()}})
+	in := newInput("ainp_some-other-provider", "alice", []string{"grp-eng"})
+
+	out, err := mw.Invoke(context.Background(), in)
+	require.NoError(t, err)
+	require.NotNil(t, out)
+	assert.Equal(t, middleware.DecisionAllow, out.Decision)
+	assert.Nil(t, out.Mutations, "non-LiteLLM resolved provider must produce no mutations")
+}
+
+func TestInject_NoResolvedProvider_NoMutations(t *testing.T) {
+	mw := New(Config{Providers: []ProviderInjection{liteLLMRule()}})
+	in := &middleware.Input{Slot: middleware.SlotOnRequest, UserID: "alice"}
+
+	out, err := mw.Invoke(context.Background(), in)
+	require.NoError(t, err)
+	require.NotNil(t, out)
+	assert.Nil(t, out.Mutations,
+		"missing llm.resolved_provider_id metadata means the router didn't run; never stamp identity blindly")
+}
+
+func TestInject_PartialRule_StampsOnlyConfiguredHeaders(t *testing.T) {
+	mw := New(Config{Providers: []ProviderInjection{{
+		ProviderID: litellmProvider,
+		HeaderPair: &HeaderPairRule{
+			EndUserIDHeader: "x-litellm-end-user-id",
+			// TagsHeader intentionally empty.
+		},
+	}}})
+	in := newInput(litellmProvider, "alice", []string{"grp-eng"})
+
+	out, err := mw.Invoke(context.Background(), in)
+	require.NoError(t, err)
+	require.NotNil(t, out)
+	require.NotNil(t, out.Mutations)
+
+	assert.Equal(t, []string{"x-litellm-end-user-id"}, out.Mutations.HeadersRemove,
+		"only configured header should be stripped")
+	require.Len(t, out.Mutations.HeadersAdd, 1)
+	assert.Equal(t, "x-litellm-end-user-id", out.Mutations.HeadersAdd[0].Key)
+	assert.Equal(t, "alice", out.Mutations.HeadersAdd[0].Value)
+}
+
+func TestInject_EmptyIdentity_StripsButDoesNotAdd(t *testing.T) {
+	// Caller has no UserID and no groups. We still strip the headers
+	// (so the client can't inject identity) but we don't add empty
+	// values that would mislead the gateway.
+	mw := New(Config{Providers: []ProviderInjection{liteLLMRule()}})
+	in := newInput(litellmProvider, "", nil)
+	in.AccountID = ""
+	in.SourceIP = ""
+	in.RequestID = ""
+
+	out, err := mw.Invoke(context.Background(), in)
+	require.NoError(t, err)
+	require.NotNil(t, out)
+	require.NotNil(t, out.Mutations)
+
+	assert.ElementsMatch(t,
+		[]string{"x-litellm-end-user-id", "x-litellm-tags"},
+		out.Mutations.HeadersRemove,
+		"identity headers must be stripped even when we don't have values to add — anti-spoof")
+	assert.Empty(t, out.Mutations.HeadersAdd,
+		"no NetBird identity available; do not stamp empty / misleading values")
+}
+
+func TestInject_TagsCSV_DedupesAndSorts(t *testing.T) {
+	mw := New(Config{Providers: []ProviderInjection{liteLLMRule()}})
+	in := newInput(litellmProvider, "alice", []string{"grp-zzz", "grp-aaa", "grp-zzz", "", " "})
+
+	out, err := mw.Invoke(context.Background(), in)
+	require.NoError(t, err)
+	require.NotNil(t, out)
+	require.NotNil(t, out.Mutations)
+
+	for _, kv := range out.Mutations.HeadersAdd {
+		if kv.Key == "x-litellm-tags" {
+			assert.Equal(t, "grp-aaa,grp-zzz", kv.Value,
+				"tags CSV must dedupe, drop empty, and sort")
+			return
+		}
+	}
+	t.Fatalf("expected x-litellm-tags in HeadersAdd; got %v", out.Mutations.HeadersAdd)
+}
+
+func TestFactory_RejectsBadJSON(t *testing.T) {
+	_, err := Factory{}.New([]byte("{not json"))
+	require.Error(t, err)
+}
+
+func TestFactory_AcceptsEmptyShapes(t *testing.T) {
+	for _, raw := range [][]byte{nil, []byte(""), []byte(" "), []byte("null"), []byte("{}"), []byte("[]")} {
+		mw, err := Factory{}.New(raw)
+		require.NoError(t, err)
+		require.NotNil(t, mw)
+
+		out, ierr := mw.Invoke(context.Background(),
+			newInput(litellmProvider, "alice", []string{"grp-eng"}))
+		require.NoError(t, ierr)
+		assert.Equal(t, middleware.DecisionAllow, out.Decision)
+		assert.Nil(t, out.Mutations,
+			"empty config means no providers to inject for; every resolved provider passes through")
+	}
+}
+
+func TestFactory_DropsInjectionRuleWithEmptyHeaders(t *testing.T) {
+	mw, err := Factory{}.New([]byte(`{"providers":[{"provider_id":"x"}]}`))
+	require.NoError(t, err)
+	out, ierr := mw.Invoke(context.Background(), newInput("x", "alice", []string{"grp-eng"}))
+	require.NoError(t, ierr)
+	assert.Nil(t, out.Mutations,
+		"a rule with no header names is functionally a no-op and must be dropped at New() time")
+}
+
+// TestInject_TagsFromAuthorisingMetadata pins that when llm_router has
+// emitted llm.authorising_groups, the inject middleware uses THAT
+// (the per-request authorising intersection) for the tags header — not
+// the full UserGroups, which can include groups unrelated to this
+// request's routing.
+func TestInject_TagsFromAuthorisingMetadata(t *testing.T) {
+	mw := New(Config{Providers: []ProviderInjection{liteLLMRule()}})
+	in := newInput(litellmProvider, "alice", []string{"grp-eng", "grp-it", "grp-oncall"})
+	in.Metadata = append(in.Metadata, middleware.KV{
+		Key:   middleware.KeyLLMAuthorisingGroups,
+		Value: "grp-eng",
+	})
+
+	out, err := mw.Invoke(context.Background(), in)
+	require.NoError(t, err)
+	require.NotNil(t, out)
+	require.NotNil(t, out.Mutations)
+
+	for _, kv := range out.Mutations.HeadersAdd {
+		if kv.Key == "x-litellm-tags" {
+			assert.Equal(t, "grp-eng", kv.Value,
+				"tags must come from llm.authorising_groups, not the full UserGroups; unrelated peer groups must not leak")
+			return
+		}
+	}
+	t.Fatalf("expected x-litellm-tags in HeadersAdd; got %v", out.Mutations.HeadersAdd)
+}
+
+// TestInject_TagsFallsBackToUserGroups pins the defensive fallback: if
+// llm_router didn't emit authorising-groups metadata (chain
+// misconfiguration) the middleware uses UserGroups so identity is
+// still stamped, just over-broad.
+func TestInject_TagsFallsBackToUserGroups(t *testing.T) {
+	mw := New(Config{Providers: []ProviderInjection{liteLLMRule()}})
+	in := newInput(litellmProvider, "alice", []string{"grp-eng", "grp-it"})
+	// No llm.authorising_groups metadata.
+
+	out, err := mw.Invoke(context.Background(), in)
+	require.NoError(t, err)
+	require.NotNil(t, out)
+	require.NotNil(t, out.Mutations)
+
+	for _, kv := range out.Mutations.HeadersAdd {
+		if kv.Key == "x-litellm-tags" {
+			assert.Equal(t, "grp-eng,grp-it", kv.Value,
+				"absent metadata must fall back to the full UserGroups CSV")
+			return
+		}
+	}
+	t.Fatalf("expected x-litellm-tags in HeadersAdd; got %v", out.Mutations.HeadersAdd)
+}
+
+// portkeyRule is the JSONMetadata-shape analogue of liteLLMRule: a
+// single x-portkey-metadata header carrying _user and groups, with
+// Portkey's 128-byte per-value cap.
+func portkeyRule() ProviderInjection {
+	return ProviderInjection{
+		ProviderID: portkeyProvider,
+		JSONMetadata: &JSONMetadataRule{
+			Header:         "x-portkey-metadata",
+			UserKey:        "_user",
+			GroupsKey:      "groups",
+			MaxValueLength: 128,
+		},
+	}
+}
+
+// TestInject_JSONMetadata_StampsHeader pins the Portkey-style emission:
+// one header carrying a JSON envelope with reserved keys for user
+// identity and groups CSV.
+func TestInject_JSONMetadata_StampsHeader(t *testing.T) {
+	mw := New(Config{Providers: []ProviderInjection{portkeyRule()}})
+	in := newInput(portkeyProvider, "alice", []string{"grp-eng", "grp-it"})
+	in.UserEmail = "alice@example.com"
+
+	out, err := mw.Invoke(context.Background(), in)
+	require.NoError(t, err)
+	require.NotNil(t, out)
+	require.NotNil(t, out.Mutations)
+
+	assert.Equal(t, []string{"x-portkey-metadata"}, out.Mutations.HeadersRemove,
+		"the JSON header must be stripped before we add our trusted value")
+	require.Len(t, out.Mutations.HeadersAdd, 1)
+	added := out.Mutations.HeadersAdd[0]
+	assert.Equal(t, "x-portkey-metadata", added.Key)
+
+	var payload map[string]string
+	require.NoError(t, json.Unmarshal([]byte(added.Value), &payload))
+	assert.Equal(t, "alice@example.com", payload["_user"],
+		"_user reserved key carries the display identity (UserEmail)")
+	assert.Equal(t, "grp-eng,grp-it", payload["groups"],
+		"groups key carries the sorted CSV of group display names")
+}
+
+// TestInject_JSONMetadata_TruncatesValues pins the per-value byte cap.
+// Portkey rejects metadata values longer than 128 chars; oversized
+// values are truncated rather than failing the request.
+func TestInject_JSONMetadata_TruncatesValues(t *testing.T) {
+	mw := New(Config{Providers: []ProviderInjection{portkeyRule()}})
+	in := newInput(portkeyProvider, "alice", []string{"grp-eng"})
+	in.UserEmail = strings.Repeat("a", 200) + "@example.com"
+
+	out, err := mw.Invoke(context.Background(), in)
+	require.NoError(t, err)
+	require.NotNil(t, out.Mutations)
+	require.Len(t, out.Mutations.HeadersAdd, 1)
+
+	var payload map[string]string
+	require.NoError(t, json.Unmarshal([]byte(out.Mutations.HeadersAdd[0].Value), &payload))
+	assert.Len(t, payload["_user"], 128,
+		"per-value byte length must be capped at MaxValueLength")
+}
+
+// TestInject_JSONMetadata_Sanitize pins the AWS-Bedrock sanitization path: when
+// Sanitize is set, characters outside Bedrock's accepted metadata class
+// (notably the groups CSV comma and arbitrary characters in group display
+// names) are replaced with '_' so Bedrock doesn't reject the request. Allowed
+// characters (letters, digits, spaces, and @ . _ : / + - =) pass through.
+func TestInject_JSONMetadata_Sanitize(t *testing.T) {
+	rule := ProviderInjection{
+		ProviderID: portkeyProvider,
+		JSONMetadata: &JSONMetadataRule{
+			Header:         "X-Amzn-Bedrock-Request-Metadata",
+			UserKey:        "user",
+			GroupsKey:      "group",
+			MaxValueLength: 256,
+			Sanitize:       true,
+		},
+	}
+	mw := New(Config{Providers: []ProviderInjection{rule}})
+	in := newInput(portkeyProvider, "alice", []string{"g1", "g2"})
+	in.UserEmail = "alice@example.com"
+	// Group display names carry characters Bedrock rejects (comma, '#'); the CSV
+	// join adds another comma between the two groups.
+	in.UserGroupNames = []string{"Eng,Team", "Ops#1"}
+
+	out, err := mw.Invoke(context.Background(), in)
+	require.NoError(t, err)
+	require.NotNil(t, out.Mutations)
+	require.Len(t, out.Mutations.HeadersAdd, 1)
+	added := out.Mutations.HeadersAdd[0]
+	assert.Equal(t, "X-Amzn-Bedrock-Request-Metadata", added.Key,
+		"the Bedrock cost-allocation header carries the metadata JSON")
+
+	var payload map[string]string
+	require.NoError(t, json.Unmarshal([]byte(added.Value), &payload))
+	assert.Equal(t, "alice@example.com", payload["user"],
+		"'@' and '.' are in Bedrock's accepted set and must be preserved")
+	assert.NotContains(t, payload["group"], ",", "commas must be sanitized — Bedrock rejects them")
+	assert.NotContains(t, payload["group"], "#", "disallowed characters must be sanitized")
+	assert.Contains(t, payload["group"], "Eng", "allowed characters must be preserved")
+}
+
+// TestInject_JSONMetadata_EmptyIdentity_StripsButDoesNotAdd verifies the
+// anti-spoof Remove still fires when there's nothing to stamp.
+func TestInject_JSONMetadata_EmptyIdentity_StripsButDoesNotAdd(t *testing.T) {
+	mw := New(Config{Providers: []ProviderInjection{portkeyRule()}})
+	in := newInput(portkeyProvider, "", nil)
+	in.UserEmail = ""
+
+	out, err := mw.Invoke(context.Background(), in)
+	require.NoError(t, err)
+	require.NotNil(t, out.Mutations)
+
+	assert.Equal(t, []string{"x-portkey-metadata"}, out.Mutations.HeadersRemove,
+		"strip even with no payload — client can't smuggle identity headers")
+	assert.Empty(t, out.Mutations.HeadersAdd,
+		"no NetBird identity available; do not stamp empty / misleading values")
+}
+
+// TestFactory_RejectsRuleWithBothShapes pins the configuration-error
+// guard: a rule that sets both HeaderPair and JSONMetadata is dropped
+// at New() time rather than guessing which wins.
+func TestFactory_RejectsRuleWithBothShapes(t *testing.T) {
+	mw := New(Config{Providers: []ProviderInjection{{
+		ProviderID: litellmProvider,
+		HeaderPair: &HeaderPairRule{
+			EndUserIDHeader: "x-litellm-end-user-id",
+		},
+		JSONMetadata: &JSONMetadataRule{
+			Header:  "x-portkey-metadata",
+			UserKey: "_user",
+		},
+	}}})
+	in := newInput(litellmProvider, "alice", []string{"grp-eng"})
+
+	out, err := mw.Invoke(context.Background(), in)
+	require.NoError(t, err)
+	assert.Nil(t, out.Mutations,
+		"a rule that sets both shapes is ambiguous and must be dropped at New() time")
+}
+
+// liteLLMRuleWithBody is the LiteLLM-style rule with body tag injection
+// enabled (matches the catalog default).
+func liteLLMRuleWithBody() ProviderInjection {
+	return ProviderInjection{
+		ProviderID: litellmProvider,
+		HeaderPair: &HeaderPairRule{
+			EndUserIDHeader: "x-litellm-end-user-id",
+			TagsHeader:      "x-litellm-tags",
+			TagsInBody:      true,
+		},
+	}
+}
+
+// TestInject_BodyTags_AddsMetadataTags pins the body-inject path that
+// LiteLLM's _tag_max_budget_check requires. With TagsInBody set, the
+// middleware writes the authorising-groups slice into
+// request.metadata.tags (in addition to the header).
+func TestInject_BodyTags_AddsMetadataTags(t *testing.T) {
+	mw := New(Config{Providers: []ProviderInjection{liteLLMRuleWithBody()}})
+	in := newInput(litellmProvider, "alice", []string{"grp-eng", "grp-sre"})
+	in.Body = []byte(`{"model":"gpt-4o-mini","messages":[]}`)
+
+	out, err := mw.Invoke(context.Background(), in)
+	require.NoError(t, err)
+	require.NotNil(t, out.Mutations)
+	require.NotEmpty(t, out.Mutations.BodyReplace, "body must be rewritten when TagsInBody is set")
+
+	var doc map[string]any
+	require.NoError(t, json.Unmarshal(out.Mutations.BodyReplace, &doc))
+	meta, ok := doc["metadata"].(map[string]any)
+	require.True(t, ok, "metadata must be an object")
+	tags, ok := meta["tags"].([]any)
+	require.True(t, ok, "metadata.tags must be a JSON array")
+	got := make([]string, 0, len(tags))
+	for _, t := range tags {
+		s, _ := t.(string)
+		got = append(got, s)
+	}
+	assert.Equal(t, []string{"grp-eng", "grp-sre"}, got,
+		"metadata.tags must carry the sorted authorising-groups slice")
+	assert.Equal(t, "gpt-4o-mini", doc["model"],
+		"the rest of the body must be preserved verbatim")
+}
+
+// TestInject_BodyTags_PreservesExistingMetadata pins that an existing
+// metadata object on the request is merged with our tags rather than
+// clobbered — clients sometimes set metadata fields the proxy
+// shouldn't blow away (jobID, taskName, etc.).
+func TestInject_BodyTags_PreservesExistingMetadata(t *testing.T) {
+	mw := New(Config{Providers: []ProviderInjection{liteLLMRuleWithBody()}})
+	in := newInput(litellmProvider, "alice", []string{"grp-eng"})
+	in.Body = []byte(`{"model":"gpt-4o-mini","metadata":{"jobID":"j-42","tags":["should-be-replaced"]}}`)
+
+	out, err := mw.Invoke(context.Background(), in)
+	require.NoError(t, err)
+	require.NotEmpty(t, out.Mutations.BodyReplace)
+
+	var doc map[string]any
+	require.NoError(t, json.Unmarshal(out.Mutations.BodyReplace, &doc))
+	meta := doc["metadata"].(map[string]any)
+	assert.Equal(t, "j-42", meta["jobID"],
+		"client-supplied metadata fields outside `tags` must survive")
+	tags := meta["tags"].([]any)
+	require.Len(t, tags, 1)
+	assert.Equal(t, "grp-eng", tags[0],
+		"our tags overwrite any client-supplied metadata.tags so spoofing is impossible")
+}
+
+// TestInject_BodyTags_SkipsHostileMetadataShape pins the defensive
+// refusal: when the request body has a non-object metadata field
+// (string/number/array), we don't inject — header path still emits.
+func TestInject_BodyTags_SkipsHostileMetadataShape(t *testing.T) {
+	mw := New(Config{Providers: []ProviderInjection{liteLLMRuleWithBody()}})
+	in := newInput(litellmProvider, "alice", []string{"grp-eng"})
+	in.Body = []byte(`{"model":"gpt-4o-mini","metadata":"not-an-object"}`)
+
+	out, err := mw.Invoke(context.Background(), in)
+	require.NoError(t, err)
+	require.NotNil(t, out.Mutations)
+	assert.Empty(t, out.Mutations.BodyReplace,
+		"non-object metadata must skip body inject (don't clobber)")
+
+	for _, kv := range out.Mutations.HeadersAdd {
+		if kv.Key == "x-litellm-tags" {
+			assert.Equal(t, "grp-eng", kv.Value,
+				"header path must still emit so spend tracking keeps working")
+			return
+		}
+	}
+	t.Fatalf("expected x-litellm-tags header even when body inject was skipped")
+}
+
+// TestInject_BodyTags_SkipsTruncatedBody pins that we don't blindly
+// rewrite a body we don't have in full. The header path still runs.
+func TestInject_BodyTags_SkipsTruncatedBody(t *testing.T) {
+	mw := New(Config{Providers: []ProviderInjection{liteLLMRuleWithBody()}})
+	in := newInput(litellmProvider, "alice", []string{"grp-eng"})
+	in.Body = []byte(`{"model":"gpt-4o-mini","messages":[]}`)
+	in.BodyTruncated = true
+
+	out, err := mw.Invoke(context.Background(), in)
+	require.NoError(t, err)
+	assert.Empty(t, out.Mutations.BodyReplace,
+		"truncated body must skip body inject — re-marshaling would corrupt the request")
+}
+
+// TestInject_BodyTags_SkipsNonJSONBody pins graceful behavior when the
+// body isn't JSON (e.g. a streaming binary or form upload sneaking
+// through the LLM chain). Header path still runs.
+func TestInject_BodyTags_SkipsNonJSONBody(t *testing.T) {
+	mw := New(Config{Providers: []ProviderInjection{liteLLMRuleWithBody()}})
+	in := newInput(litellmProvider, "alice", []string{"grp-eng"})
+	in.Body = []byte(`not even close to json`)
+
+	out, err := mw.Invoke(context.Background(), in)
+	require.NoError(t, err)
+	assert.Empty(t, out.Mutations.BodyReplace,
+		"non-JSON body must skip body inject silently")
+}
+
+// liteLLMRuleFull mirrors the catalog default: header path + body
+// metadata.tags (groups) + body user (end-user id).
+func liteLLMRuleFull() ProviderInjection {
+	return ProviderInjection{
+		ProviderID: litellmProvider,
+		HeaderPair: &HeaderPairRule{
+			EndUserIDHeader: "x-litellm-end-user-id",
+			TagsHeader:      "x-litellm-tags",
+			TagsInBody:      true,
+			EndUserIDInBody: true,
+		},
+	}
+}
+
+// TestInject_BodyUser_WritesTopLevelUser pins the EndUserIDInBody path
+// alone: body's top-level "user" field carries the display identity.
+// Tags-in-body is OFF here so we isolate the user write.
+func TestInject_BodyUser_WritesTopLevelUser(t *testing.T) {
+	mw := New(Config{Providers: []ProviderInjection{{
+		ProviderID: litellmProvider,
+		HeaderPair: &HeaderPairRule{
+			EndUserIDHeader: "x-litellm-end-user-id",
+			EndUserIDInBody: true,
+		},
+	}}})
+	in := newInput(litellmProvider, "alice", nil)
+	in.UserEmail = "alice@example.com"
+	in.Body = []byte(`{"model":"gpt-4o-mini","messages":[]}`)
+
+	out, err := mw.Invoke(context.Background(), in)
+	require.NoError(t, err)
+	require.NotNil(t, out.Mutations)
+	require.NotEmpty(t, out.Mutations.BodyReplace)
+
+	var doc map[string]any
+	require.NoError(t, json.Unmarshal(out.Mutations.BodyReplace, &doc))
+	assert.Equal(t, "alice@example.com", doc["user"],
+		"body's top-level user field must carry the display identity")
+	_, hasMeta := doc["metadata"]
+	assert.False(t, hasMeta, "TagsInBody is off; metadata must not be added")
+}
+
+// TestInject_BodyUser_OverwritesClientSupplied pins anti-spoof: a
+// client-supplied "user" in the body is overwritten so the gateway
+// only sees our trusted identity.
+func TestInject_BodyUser_OverwritesClientSupplied(t *testing.T) {
+	mw := New(Config{Providers: []ProviderInjection{liteLLMRuleFull()}})
+	in := newInput(litellmProvider, "alice", []string{"grp-eng"})
+	in.UserEmail = "alice@example.com"
+	in.Body = []byte(`{"model":"gpt-4o-mini","user":"ceo@company.com"}`)
+
+	out, err := mw.Invoke(context.Background(), in)
+	require.NoError(t, err)
+	require.NotEmpty(t, out.Mutations.BodyReplace)
+
+	var doc map[string]any
+	require.NoError(t, json.Unmarshal(out.Mutations.BodyReplace, &doc))
+	assert.Equal(t, "alice@example.com", doc["user"],
+		"client-supplied user must be overwritten with the trusted identity")
+}
+
+// TestInject_BodyCombined_TagsAndUser pins that with both flags on,
+// the body carries both metadata.tags AND top-level user, and the
+// header path still emits.
+func TestInject_BodyCombined_TagsAndUser(t *testing.T) {
+	mw := New(Config{Providers: []ProviderInjection{liteLLMRuleFull()}})
+	in := newInput(litellmProvider, "alice", []string{"grp-eng", "grp-sre"})
+	in.UserEmail = "alice@example.com"
+	in.Body = []byte(`{"model":"gpt-4o-mini"}`)
+
+	out, err := mw.Invoke(context.Background(), in)
+	require.NoError(t, err)
+	require.NotEmpty(t, out.Mutations.BodyReplace)
+
+	var doc map[string]any
+	require.NoError(t, json.Unmarshal(out.Mutations.BodyReplace, &doc))
+	assert.Equal(t, "alice@example.com", doc["user"])
+	meta := doc["metadata"].(map[string]any)
+	tags := meta["tags"].([]any)
+	require.Len(t, tags, 2)
+	assert.Equal(t, "grp-eng", tags[0])
+	assert.Equal(t, "grp-sre", tags[1])
+
+	// Header path still emits — header end-user-id is the primary
+	// path for LiteLLM's resolver, body is defense-in-depth.
+	added := map[string]string{}
+	for _, kv := range out.Mutations.HeadersAdd {
+		added[kv.Key] = kv.Value
+	}
+	assert.Equal(t, "alice@example.com", added["x-litellm-end-user-id"])
+	assert.Equal(t, "grp-eng,grp-sre", added["x-litellm-tags"])
+}
+
+// TestInject_BodyCombined_HostileMetadataKeepsUser pins the partial-
+// success path: a hostile (non-object) metadata field skips the tag
+// write but still allows the orthogonal user write to land.
+func TestInject_BodyCombined_HostileMetadataKeepsUser(t *testing.T) {
+	mw := New(Config{Providers: []ProviderInjection{liteLLMRuleFull()}})
+	in := newInput(litellmProvider, "alice", []string{"grp-eng"})
+	in.UserEmail = "alice@example.com"
+	in.Body = []byte(`{"model":"gpt-4o-mini","metadata":"not-an-object"}`)
+
+	out, err := mw.Invoke(context.Background(), in)
+	require.NoError(t, err)
+	require.NotEmpty(t, out.Mutations.BodyReplace,
+		"user write must still go through even when metadata is hostile")
+
+	var doc map[string]any
+	require.NoError(t, json.Unmarshal(out.Mutations.BodyReplace, &doc))
+	assert.Equal(t, "alice@example.com", doc["user"])
+	assert.Equal(t, "not-an-object", doc["metadata"],
+		"hostile metadata must be left untouched, not clobbered")
+}
+
+// TestInject_ExtraHeaders_Stamped pins the extras path: with a
+// per-provider ExtraHeader configured (e.g. Portkey config id), the
+// middleware stamps it on every matching request and adds the same
+// name to HeadersRemove for anti-spoof.
+func TestInject_ExtraHeaders_Stamped(t *testing.T) {
+	mw := New(Config{Providers: []ProviderInjection{{
+		ProviderID: portkeyProvider,
+		JSONMetadata: &JSONMetadataRule{
+			Header:    "x-portkey-metadata",
+			UserKey:   "_user",
+			GroupsKey: "groups",
+		},
+		ExtraHeaders: []ExtraHeaderKV{
+			{Name: "x-portkey-config", Value: "pc-prod-3f2a"},
+		},
+	}}})
+	in := newInput(portkeyProvider, "alice", []string{"grp-eng"})
+	in.UserEmail = "alice@example.com"
+
+	out, err := mw.Invoke(context.Background(), in)
+	require.NoError(t, err)
+	require.NotNil(t, out.Mutations)
+
+	assert.Contains(t, out.Mutations.HeadersRemove, "x-portkey-config",
+		"extras must be stripped before stamping for anti-spoof")
+	added := map[string]string{}
+	for _, kv := range out.Mutations.HeadersAdd {
+		added[kv.Key] = kv.Value
+	}
+	assert.Equal(t, "pc-prod-3f2a", added["x-portkey-config"],
+		"extras must carry the operator-configured value verbatim")
+	// Identity-stamping shape (JSONMetadata header) still emitted.
+	assert.Contains(t, added, "x-portkey-metadata",
+		"extras and identity stamping are independent — both must land")
+}
+
+// TestInject_ExtraHeaders_OnlyRule pins that an extras-only rule
+// (no HeaderPair, no JSONMetadata) survives New() and stamps the
+// extras anyway. Useful for hypothetical gateways that need a static
+// routing header but no NetBird identity stamping.
+func TestInject_ExtraHeaders_OnlyRule(t *testing.T) {
+	mw := New(Config{Providers: []ProviderInjection{{
+		ProviderID: "ainp_extras-only",
+		ExtraHeaders: []ExtraHeaderKV{
+			{Name: "x-routing-key", Value: "rk-1"},
+		},
+	}}})
+	in := newInput("ainp_extras-only", "alice", nil)
+
+	out, err := mw.Invoke(context.Background(), in)
+	require.NoError(t, err)
+	require.NotNil(t, out.Mutations,
+		"extras alone keep the rule alive — middleware must emit them")
+	added := map[string]string{}
+	for _, kv := range out.Mutations.HeadersAdd {
+		added[kv.Key] = kv.Value
+	}
+	assert.Equal(t, "rk-1", added["x-routing-key"])
+}
+
+// TestInject_ExtraHeaders_EmptyValueSkipped pins that empty values are
+// dropped silently (the synth would normally not send them, but the
+// middleware is defensive).
+func TestInject_ExtraHeaders_EmptyValueSkipped(t *testing.T) {
+	mw := New(Config{Providers: []ProviderInjection{{
+		ProviderID: portkeyProvider,
+		JSONMetadata: &JSONMetadataRule{
+			Header:  "x-portkey-metadata",
+			UserKey: "_user",
+		},
+		ExtraHeaders: []ExtraHeaderKV{
+			{Name: "x-portkey-config", Value: ""},
+		},
+	}}})
+	in := newInput(portkeyProvider, "alice", []string{"grp-eng"})
+	in.UserEmail = "alice@example.com"
+
+	out, err := mw.Invoke(context.Background(), in)
+	require.NoError(t, err)
+	require.NotNil(t, out.Mutations)
+	assert.NotContains(t, out.Mutations.HeadersRemove, "x-portkey-config",
+		"empty extra value must not even strip the header")
+	for _, kv := range out.Mutations.HeadersAdd {
+		assert.NotEqual(t, "x-portkey-config", kv.Key,
+			"empty extra value must not be stamped")
+	}
+}

proxy/internal/middleware/builtin/llm_limit_check/factory.go

@@ -0,0 +1,38 @@
+// Package llm_limit_check is the SlotOnRequest middleware that asks
+// management which agent-network policy "pays" for the current LLM
+// request. On allow, it stamps the selected policy id, attribution
+// group id, and effective window length onto the metadata bag so the
+// post-flight llm_limit_record middleware can tick the right counters.
+// On deny, it returns a 403 carrying the canonical llm_policy.* deny
+// code surfaced by management.
+package llm_limit_check
+
+import (
+	"github.com/netbirdio/netbird/proxy/internal/middleware"
+	"github.com/netbirdio/netbird/proxy/internal/middleware/builtin"
+)
+
+// ID is the registry identifier for this middleware.
+const ID = "llm_limit_check"
+
+// Factory builds a configured llm_limit_check instance. The factory
+// has no per-target config — it pulls the management gRPC client from
+// the package-level FactoryContext at construction time. A nil
+// MgmtClient on the context is allowed; the middleware then becomes
+// a no-op pass-through (allow without attribution) so a partially
+// wired environment doesn't break the chain.
+type Factory struct{}
+
+// ID returns the registry identifier matching the middleware ID.
+func (Factory) ID() string { return ID }
+
+// New ignores the rawConfig payload (no per-target config today) and
+// returns a Middleware bound to the FactoryContext's MgmtClient.
+func (Factory) New(_ []byte) (middleware.Middleware, error) {
+	ctx := builtin.Context()
+	return New(ctx.MgmtClient, ctx.Logger), nil
+}
+
+func init() {
+	builtin.Register(Factory{})
+}

proxy/internal/middleware/builtin/llm_limit_check/middleware.go

@@ -0,0 +1,196 @@
+package llm_limit_check
+
+import (
+	"context"
+	"strconv"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+
+	"github.com/netbirdio/netbird/proxy/internal/middleware"
+	"github.com/netbirdio/netbird/proxy/internal/middleware/builtin"
+	"github.com/netbirdio/netbird/shared/management/proto"
+)
+
+// Version is reported via Middleware.Version().
+const Version = "1.0.0"
+
+// callTimeout caps the wall-clock budget for the pre-flight RPC. The
+// middleware sits on the request leg, so a slow management call
+// translates directly to user-visible latency. 2s is loose enough for
+// a healthy management cluster but tight enough that a stalled call
+// fails open via the same path nil-MgmtClient does — an enforcement
+// gate that adds 30s of latency is worse than a stale gate.
+const callTimeout = 2 * time.Second
+
+// Middleware is the per-target instance that runs the pre-flight check.
+type Middleware struct {
+	mgmt   builtin.MgmtClient
+	logger *log.Logger
+}
+
+// New constructs a Middleware. mgmt may be nil — that's the
+// no-management-wired case where the middleware is a pass-through
+// (allow without attribution); useful for unit tests and for
+// progressive rollout of the management RPC.
+func New(mgmt builtin.MgmtClient, logger *log.Logger) *Middleware {
+	if logger == nil {
+		logger = log.StandardLogger()
+	}
+	return &Middleware{mgmt: mgmt, logger: logger}
+}
+
+// ID returns the registry identifier.
+func (m *Middleware) ID() string { return ID }
+
+// Version returns the implementation version.
+func (m *Middleware) Version() string { return Version }
+
+// Slot reports the chain slot the middleware lives in.
+func (m *Middleware) Slot() middleware.Slot { return middleware.SlotOnRequest }
+
+// AcceptedContentTypes returns nil because the gate consults metadata
+// emitted upstream (KeyLLMResolvedProviderID) and never inspects bodies.
+func (m *Middleware) AcceptedContentTypes() []string { return nil }
+
+// MetadataKeys is the closed allowlist of keys this middleware emits.
+func (m *Middleware) MetadataKeys() []string {
+	return []string{
+		middleware.KeyLLMSelectedPolicyID,
+		middleware.KeyLLMAttributionGroupID,
+		middleware.KeyLLMAttributionWindowS,
+		middleware.KeyLLMPolicyDecision,
+		middleware.KeyLLMPolicyReason,
+	}
+}
+
+// MutationsSupported reports that the middleware never mutates the
+// request body or headers; the only outcome is allow + metadata or
+// deny.
+func (m *Middleware) MutationsSupported() bool { return false }
+
+// Close releases resources owned by the middleware. Stateless, so
+// this is a no-op.
+func (m *Middleware) Close() error { return nil }
+
+// Invoke runs the pre-flight policy check.
+func (m *Middleware) Invoke(ctx context.Context, in *middleware.Input) (*middleware.Output, error) {
+	if m.mgmt == nil {
+		// No management client wired — fall through to allow with
+		// no attribution. RecordLLMUsage on the response leg will
+		// also be a no-op so counters stay at zero. This matches
+		// the PR1 behaviour exactly so a partial wiring is
+		// indistinguishable from "no enforcement".
+		return allowNoAttribution(), nil
+	}
+
+	providerID := lookupKV(in.Metadata, middleware.KeyLLMResolvedProviderID)
+	if providerID == "" {
+		// llm_router didn't emit a resolved provider id — usually
+		// because the request didn't carry an llm.model. The
+		// router itself denied; we won't reach here in production,
+		// but defensively pass through so we never deny on top of
+		// an upstream allow.
+		return allowNoAttribution(), nil
+	}
+
+	rpcCtx, cancel := context.WithTimeout(ctx, callTimeout)
+	defer cancel()
+
+	resp, err := m.mgmt.CheckLLMPolicyLimits(rpcCtx, &proto.CheckLLMPolicyLimitsRequest{
+		AccountId:  in.AccountID,
+		UserId:     in.UserID,
+		GroupIds:   append([]string(nil), in.UserGroups...),
+		ProviderId: providerID,
+		Model:      lookupKV(in.Metadata, middleware.KeyLLMModel),
+	})
+	if err != nil {
+		// Fail-open on transport / management errors. The
+		// alternative — denying every request when management is
+		// unreachable — is worse for v1 (operational outage =
+		// total LLM outage). Operators can audit via the
+		// access-log; PR3 can switch to fail-closed under a flag.
+		m.logger.WithError(err).
+			WithField("middleware", ID).
+			Debugf("management pre-flight failed; failing open")
+		return allowNoAttribution(), nil
+	}
+
+	if resp.GetDecision() == "deny" {
+		return denyFromManagement(resp), nil
+	}
+	return allowFromManagement(resp), nil
+}
+
+// allowNoAttribution returns the no-op allow envelope used when no
+// management client is wired or no provider was resolved. Stamps
+// decision=allow but no policy / attribution metadata so
+// llm_limit_record skips its post-flight write.
+func allowNoAttribution() *middleware.Output {
+	return &middleware.Output{
+		Decision: middleware.DecisionAllow,
+		Metadata: []middleware.KV{
+			{Key: middleware.KeyLLMPolicyDecision, Value: "allow"},
+		},
+	}
+}
+
+// allowFromManagement converts a successful CheckLLMPolicyLimits
+// response into the chain's allow envelope, stamping the attribution
+// metadata the response leg consumes.
+func allowFromManagement(resp *proto.CheckLLMPolicyLimitsResponse) *middleware.Output {
+	out := &middleware.Output{
+		Decision: middleware.DecisionAllow,
+		Metadata: []middleware.KV{
+			{Key: middleware.KeyLLMPolicyDecision, Value: "allow"},
+		},
+	}
+	if id := resp.GetSelectedPolicyId(); id != "" {
+		out.Metadata = append(out.Metadata, middleware.KV{Key: middleware.KeyLLMSelectedPolicyID, Value: id})
+	}
+	if g := resp.GetAttributionGroupId(); g != "" {
+		out.Metadata = append(out.Metadata, middleware.KV{Key: middleware.KeyLLMAttributionGroupID, Value: g})
+	}
+	if w := resp.GetWindowSeconds(); w > 0 {
+		out.Metadata = append(out.Metadata, middleware.KV{Key: middleware.KeyLLMAttributionWindowS, Value: strconv.FormatInt(w, 10)})
+	}
+	return out
+}
+
+// denyFromManagement converts a deny response into the chain's deny
+// envelope. The deny code surfaces verbatim through the framework's
+// fixed JSON template; arbitrary middleware bytes can't reach the
+// wire.
+func denyFromManagement(resp *proto.CheckLLMPolicyLimitsResponse) *middleware.Output {
+	code := resp.GetDenyCode()
+	if code == "" {
+		code = "llm_policy.cap_exceeded"
+	}
+	// The canonical code is safe to surface; the management-supplied
+	// reason can name internal quota details (used amounts, caps, rule
+	// ids), so keep the public message generic and leave the detail to
+	// server-side logs.
+	return &middleware.Output{
+		Decision:   middleware.DecisionDeny,
+		DenyStatus: 403,
+		DenyReason: &middleware.DenyReason{
+			Code:    code,
+			Message: "LLM policy limit exceeded",
+		},
+		Metadata: []middleware.KV{
+			{Key: middleware.KeyLLMPolicyDecision, Value: "deny"},
+			{Key: middleware.KeyLLMPolicyReason, Value: code},
+		},
+	}
+}
+
+// lookupKV returns the value associated with key, or the empty
+// string when absent.
+func lookupKV(kvs []middleware.KV, key string) string {
+	for _, kv := range kvs {
+		if kv.Key == key {
+			return kv.Value
+		}
+	}
+	return ""
+}

proxy/internal/middleware/builtin/llm_limit_check/middleware_test.go

@@ -0,0 +1,186 @@
+package llm_limit_check
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc"
+
+	"github.com/netbirdio/netbird/proxy/internal/middleware"
+	"github.com/netbirdio/netbird/shared/management/proto"
+)
+
+// fakeMgmt is a minimal builtin.MgmtClient stub that lets the test
+// drive CheckLLMPolicyLimits responses without a real gRPC dial.
+type fakeMgmt struct {
+	checkResp *proto.CheckLLMPolicyLimitsResponse
+	checkErr  error
+	checkReq  *proto.CheckLLMPolicyLimitsRequest
+}
+
+func (f *fakeMgmt) CheckLLMPolicyLimits(_ context.Context, in *proto.CheckLLMPolicyLimitsRequest, _ ...grpc.CallOption) (*proto.CheckLLMPolicyLimitsResponse, error) {
+	f.checkReq = in
+	return f.checkResp, f.checkErr
+}
+
+func (f *fakeMgmt) RecordLLMUsage(_ context.Context, _ *proto.RecordLLMUsageRequest, _ ...grpc.CallOption) (*proto.RecordLLMUsageResponse, error) {
+	return &proto.RecordLLMUsageResponse{}, nil
+}
+
+func runInvoke(t *testing.T, m *Middleware, in *middleware.Input) *middleware.Output {
+	t.Helper()
+	out, err := m.Invoke(context.Background(), in)
+	require.NoError(t, err, "Invoke must not propagate transport errors")
+	require.NotNil(t, out, "Invoke must always return an Output")
+	return out
+}
+
+// TestInvoke_AllowStampsAttributionMetadata covers the happy path:
+// management returns an allow decision with selected_policy_id +
+// attribution_group_id + window_seconds, the middleware emits all three
+// onto the metadata bag so the post-flight llm_limit_record
+// middleware has everything it needs to tick the right counter.
+func TestInvoke_AllowStampsAttributionMetadata(t *testing.T) {
+	mgmt := &fakeMgmt{
+		checkResp: &proto.CheckLLMPolicyLimitsResponse{
+			Decision:           "allow",
+			SelectedPolicyId:   "pol-X",
+			AttributionGroupId: "grp-engineers",
+			WindowSeconds:      86_400,
+		},
+	}
+	m := New(mgmt, nil)
+
+	out := runInvoke(t, m, &middleware.Input{
+		AccountID:  "acc-1",
+		UserID:     "user-bob",
+		UserGroups: []string{"grp-engineers"},
+		Metadata: []middleware.KV{
+			{Key: middleware.KeyLLMResolvedProviderID, Value: "prov-1"},
+			{Key: middleware.KeyLLMModel, Value: "gpt-4o"},
+		},
+	})
+
+	assert.Equal(t, middleware.DecisionAllow, out.Decision)
+	assert.Equal(t, "acc-1", mgmt.checkReq.GetAccountId(), "account_id must round-trip onto the RPC")
+	assert.Equal(t, "user-bob", mgmt.checkReq.GetUserId())
+	assert.Equal(t, []string{"grp-engineers"}, mgmt.checkReq.GetGroupIds())
+	assert.Equal(t, "prov-1", mgmt.checkReq.GetProviderId(), "resolved provider id must come from metadata")
+	assert.Equal(t, "gpt-4o", mgmt.checkReq.GetModel(), "model must come from metadata")
+
+	want := map[string]string{
+		middleware.KeyLLMPolicyDecision:     "allow",
+		middleware.KeyLLMSelectedPolicyID:   "pol-X",
+		middleware.KeyLLMAttributionGroupID: "grp-engineers",
+		middleware.KeyLLMAttributionWindowS: "86400",
+	}
+	got := map[string]string{}
+	for _, kv := range out.Metadata {
+		got[kv.Key] = kv.Value
+	}
+	assert.Equal(t, want, got, "attribution metadata must land on the bag for the response leg to consume")
+}
+
+// TestInvoke_DenyConvertsToProxyDeny proves the deny envelope round-
+// trips: management's deny code becomes the proxy framework's deny
+// payload at status 403, and the deny reason text is preserved so
+// operators can debug from the access log.
+func TestInvoke_DenyConvertsToProxyDeny(t *testing.T) {
+	mgmt := &fakeMgmt{
+		checkResp: &proto.CheckLLMPolicyLimitsResponse{
+			Decision:   "deny",
+			DenyCode:   "llm_policy.token_cap_exceeded",
+			DenyReason: "group token cap exhausted on policy pol-X (used 1000 of 1000)",
+		},
+	}
+	m := New(mgmt, nil)
+
+	out := runInvoke(t, m, &middleware.Input{
+		AccountID:  "acc-1",
+		UserGroups: []string{"grp-engineers"},
+		Metadata:   []middleware.KV{{Key: middleware.KeyLLMResolvedProviderID, Value: "prov-1"}},
+	})
+
+	assert.Equal(t, middleware.DecisionDeny, out.Decision)
+	assert.Equal(t, 403, out.DenyStatus, "policy denials are 403 — same as llm_router's")
+	require.NotNil(t, out.DenyReason, "deny envelope must carry a reason payload")
+	assert.Equal(t, "llm_policy.token_cap_exceeded", out.DenyReason.Code, "canonical deny code surfaces to the caller")
+	// The public message must stay generic: the management reason names
+	// internal quota detail (used/cap, rule id) that must not leak.
+	assert.Equal(t, "LLM policy limit exceeded", out.DenyReason.Message, "public deny message must be generic")
+	assert.NotContains(t, out.DenyReason.Message, "exhausted", "internal quota detail must not reach the caller")
+	assert.NotContains(t, out.DenyReason.Message, "1000", "internal cap numbers must not reach the caller")
+}
+
+// TestInvoke_NoMgmtClientPassesThrough proves the partial-wiring
+// safety: a middleware constructed without a management client
+// allows every request without attribution. This makes a half-set-up
+// environment indistinguishable from "no enforcement" rather than
+// breaking the chain.
+func TestInvoke_NoMgmtClientPassesThrough(t *testing.T) {
+	m := New(nil, nil)
+
+	out := runInvoke(t, m, &middleware.Input{
+		AccountID:  "acc-1",
+		UserGroups: []string{"grp-engineers"},
+		Metadata:   []middleware.KV{{Key: middleware.KeyLLMResolvedProviderID, Value: "prov-1"}},
+	})
+
+	assert.Equal(t, middleware.DecisionAllow, out.Decision)
+	for _, kv := range out.Metadata {
+		assert.NotEqual(t, middleware.KeyLLMSelectedPolicyID, kv.Key,
+			"no mgmt client = no attribution metadata; record middleware then skips its write")
+	}
+}
+
+// TestInvoke_NoResolvedProviderPassesThrough covers the defensive
+// path: when llm_router didn't set llm.resolved_provider_id (which
+// only happens on the deny side of llm_router), the gate must NOT
+// stack a second deny on top — pass through and let the upstream
+// deny stand.
+func TestInvoke_NoResolvedProviderPassesThrough(t *testing.T) {
+	m := New(&fakeMgmt{}, nil)
+
+	out := runInvoke(t, m, &middleware.Input{
+		AccountID: "acc-1",
+		Metadata:  []middleware.KV{},
+	})
+
+	assert.Equal(t, middleware.DecisionAllow, out.Decision,
+		"no resolved provider = the gate has nothing to check; never deny on top of an upstream allow")
+}
+
+// TestInvoke_RPCErrorFailsOpen proves the fail-open contract: a
+// transport error from management does NOT deny the request. v1
+// trades enforcement strictness for availability — an unreachable
+// management server otherwise turns into a total LLM outage.
+func TestInvoke_RPCErrorFailsOpen(t *testing.T) {
+	m := New(&fakeMgmt{checkErr: errors.New("connection refused")}, nil)
+
+	out := runInvoke(t, m, &middleware.Input{
+		AccountID:  "acc-1",
+		UserGroups: []string{"grp-engineers"},
+		Metadata:   []middleware.KV{{Key: middleware.KeyLLMResolvedProviderID, Value: "prov-1"}},
+	})
+
+	assert.Equal(t, middleware.DecisionAllow, out.Decision,
+		"transport errors must not cascade into total LLM outages — operators audit via access log")
+}
+
+// TestMetadataKeys_Allowlist locks the closed set this middleware can
+// emit. The accumulator drops anything outside this list; adding a
+// new emission means updating both the slice and this test.
+func TestMetadataKeys_Allowlist(t *testing.T) {
+	keys := New(nil, nil).MetadataKeys()
+	want := []string{
+		middleware.KeyLLMSelectedPolicyID,
+		middleware.KeyLLMAttributionGroupID,
+		middleware.KeyLLMAttributionWindowS,
+		middleware.KeyLLMPolicyDecision,
+		middleware.KeyLLMPolicyReason,
+	}
+	assert.ElementsMatch(t, want, keys)
+}

proxy/internal/middleware/builtin/llm_limit_record/factory.go

@@ -0,0 +1,35 @@
+// Package llm_limit_record is the SlotOnResponse middleware that
+// posts the served request's token + cost deltas back to management
+// so the per-(user, group, window) consumption counters tick. Reads
+// the attribution metadata stamped by llm_limit_check on the request
+// leg + the token / cost metadata stamped by llm_response_parser and
+// cost_meter; skips the write entirely when no attribution metadata
+// is present (e.g. catch-all-allow policy with no caps configured).
+package llm_limit_record
+
+import (
+	"github.com/netbirdio/netbird/proxy/internal/middleware"
+	"github.com/netbirdio/netbird/proxy/internal/middleware/builtin"
+)
+
+// ID is the registry identifier for this middleware.
+const ID = "llm_limit_record"
+
+// Factory builds a configured llm_limit_record instance bound to the
+// FactoryContext's MgmtClient. nil-MgmtClient disables the post-flight
+// write entirely (no-op pass-through), matching the request-leg gate's
+// behaviour so a partially wired environment is consistent.
+type Factory struct{}
+
+// ID returns the registry identifier matching the middleware ID.
+func (Factory) ID() string { return ID }
+
+// New ignores the rawConfig payload (no per-target config today).
+func (Factory) New(_ []byte) (middleware.Middleware, error) {
+	ctx := builtin.Context()
+	return New(ctx.MgmtClient, ctx.Logger), nil
+}
+
+func init() {
+	builtin.Register(Factory{})
+}