[client] Cover both invalid-ID orderings in the partial-error test

The partial-failure test only placed the unknown ID after the valid
exit node, so it never verified that selection processing continues
past a leading invalid ID. Run the same assertions for both orderings.

.github/workflows/agent-network-e2e.yml

@@ -5,6 +5,13 @@ on:
   schedule:
     - cron: "0 3 * * *"
   workflow_dispatch:
+    inputs:
+      bedrock_model:
+        description: >-
+          Bedrock inference-profile id to drive the matrix with, exactly as
+          AWS issues it. Leave empty for the Sonnet 4.6 default.
+        required: false
+        default: ""
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.ref }}
@@ -51,6 +58,9 @@ jobs:
           # token (and URL, for gateways) is unset, so partial coverage is fine.
           OPENAI_TOKEN: ${{ secrets.E2E_OPENAI_TOKEN }}
           ANTHROPIC_TOKEN: ${{ secrets.E2E_ANTHROPIC_TOKEN }}
+          # Moonshot AI platform key (platform.kimi.ai); drives both Kimi wire
+          # shapes (OpenAI /v1 and Anthropic /anthropic) through kimi_api.
+          KIMI_TOKEN: ${{ secrets.E2E_KIMI_TOKEN }}
           VERCEL_URL: ${{ secrets.E2E_VERCEL_URL }}
           VERCEL_TOKEN: ${{ secrets.E2E_VERCEL_TOKEN }}
           OPENROUTER_URL: ${{ secrets.E2E_OPENROUTER_URL }}
@@ -59,6 +69,8 @@ jobs:
           CLOUDFLARE_TOKEN: ${{ secrets.E2E_CLOUDFLARE_TOKEN }}
           AWS_BEARER_TOKEN_BEDROCK: ${{ secrets.E2E_AWS_BEARER_TOKEN_BEDROCK }}
           AWS_REGION: ${{ secrets.E2E_AWS_REGION }}
+          # Bedrock model override: dispatch input wins, then the repo variable, else the test default.
+          AWS_BEDROCK_MODEL: ${{ inputs.bedrock_model || vars.E2E_AWS_BEDROCK_MODEL }}
           # Vertex (Anthropic-on-Vertex): SA + project required; region defaults
           # to "global", model to a pinned claude snapshot.
           GOOGLE_VERTEX_SA_BASE64: ${{ secrets.E2E_GOOGLE_VERTEX_SA_BASE64 }}

.github/workflows/frontend-ui.yml

@@ -86,7 +86,7 @@ jobs:
             ${{ runner.os }}-pnpm-
 
       - name: Install dependencies
-        run: pnpm install --frozen-lockfile
+        run: pnpm install --frozen-lockfile --ignore-scripts
 
       - name: Generate Wails bindings
         run: pnpm run bindings

.github/workflows/golangci-lint.yml

@@ -45,7 +45,7 @@ jobs:
             display_name: Linux
     name: ${{ matrix.display_name }}
     runs-on: ${{ matrix.os }}
-    timeout-minutes: 15
+    timeout-minutes: 25
     steps:
       - name: Checkout code
         uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
@@ -79,4 +79,4 @@ jobs:
           skip-cache: true
           skip-save-cache: true
           cache-invalidation-interval: 0
-          args: --timeout=12m
+          args: --timeout=20m

.goreleaser.yaml

@@ -273,8 +273,8 @@ dockers_v2:
        - netbirdio/netbird
        - ghcr.io/netbirdio/netbird
      tags:
-       - "v{{ .Version }}-rootless"
-       - "{{ if eq .Env.SKIP_PUBLISH \"false\" }}latest{{ end }}"
+       - "{{ .Version }}-rootless"
+       - "{{ if eq .Env.SKIP_PUBLISH \"false\" }}rootless-latest{{ end }}"
      dockerfile: client/Dockerfile-rootless
      extra_files:
        - client/netbird-entrypoint.sh

client/android/client.go

@@ -439,10 +439,6 @@ func (c *Client) RemoveConnectionListener() {
 	c.recorder.RemoveConnectionListener()
 }
 
-func (c *Client) toggleRoute(command routeCommand) error {
-	return command.toggleRoute()
-}
-
 func (c *Client) getRouteManager() (routemanager.Manager, error) {
 	client := c.getConnectClient()
 	if client == nil {
@@ -462,22 +458,22 @@ func (c *Client) getRouteManager() (routemanager.Manager, error) {
 	return manager, nil
 }
 
-func (c *Client) SelectRoute(route string) error {
+func (c *Client) SelectRoute(id string) error {
 	manager, err := c.getRouteManager()
 	if err != nil {
 		return err
 	}
 
-	return c.toggleRoute(selectRouteCommand{route: route, manager: manager})
+	return manager.SelectRoutes([]route.NetID{route.NetID(id)}, true)
 }
 
-func (c *Client) DeselectRoute(route string) error {
+func (c *Client) DeselectRoute(id string) error {
 	manager, err := c.getRouteManager()
 	if err != nil {
 		return err
 	}
 
-	return c.toggleRoute(deselectRouteCommand{route: route, manager: manager})
+	return manager.DeselectRoutes([]route.NetID{route.NetID(id)})
 }
 
 // getNetworkDomainsFromRoute extracts domains from a route and enriches each domain

client/android/route_command.go

@@ -1,70 +0,0 @@
-//go:build android
-
-package android
-
-import (
-	"fmt"
-
-	log "github.com/sirupsen/logrus"
-	"golang.org/x/exp/maps"
-
-	"github.com/netbirdio/netbird/client/internal/routemanager"
-	"github.com/netbirdio/netbird/route"
-)
-
-func executeRouteToggle(id string, manager routemanager.Manager,
-	operationName string,
-	routeOperation func(routes []route.NetID, allRoutes []route.NetID) error) error {
-	netID := route.NetID(id)
-	routes := []route.NetID{netID}
-
-	routesMap := manager.GetClientRoutesWithNetID()
-	routes = route.ExpandV6ExitPairs(routes, routesMap)
-
-	log.Debugf("%s with ids: %v", operationName, routes)
-
-	if err := routeOperation(routes, maps.Keys(routesMap)); err != nil {
-		log.Debugf("error when %s: %s", operationName, err)
-		return fmt.Errorf("error %s: %w", operationName, err)
-	}
-
-	manager.TriggerSelection(manager.GetClientRoutes())
-
-	return nil
-}
-
-type routeCommand interface {
-	toggleRoute() error
-}
-
-type selectRouteCommand struct {
-	route   string
-	manager routemanager.Manager
-}
-
-func (s selectRouteCommand) toggleRoute() error {
-	routeSelector := s.manager.GetRouteSelector()
-	if routeSelector == nil {
-		return fmt.Errorf("no route selector available")
-	}
-
-	routeOperation := func(routes []route.NetID, allRoutes []route.NetID) error {
-		return routeSelector.SelectRoutes(routes, true, allRoutes)
-	}
-
-	return executeRouteToggle(s.route, s.manager, "selecting route", routeOperation)
-}
-
-type deselectRouteCommand struct {
-	route   string
-	manager routemanager.Manager
-}
-
-func (d deselectRouteCommand) toggleRoute() error {
-	routeSelector := d.manager.GetRouteSelector()
-	if routeSelector == nil {
-		return fmt.Errorf("no route selector available")
-	}
-
-	return executeRouteToggle(d.route, d.manager, "deselecting route", routeSelector.DeselectRoutes)
-}

client/iface/iface_test.go

@@ -464,6 +464,8 @@ func Test_RemovePeer(t *testing.T) {
 }
 
 func Test_ConnectPeers(t *testing.T) {
+	t.Setenv("NB_DISABLE_EBPF_WG_PROXY", "true")
+
 	peer1ifaceName := fmt.Sprintf("utun%d", WgIntNumber+400)
 	peer1wgIP := netip.MustParsePrefix("10.99.99.17/30")
 	peer1Key, _ := wgtypes.GeneratePrivateKey()
@@ -505,12 +507,8 @@ func Test_ConnectPeers(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	localIP, err := getLocalIP()
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	peer1endpoint, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", localIP, peer1wgPort))
+	localIP1 := "127.0.0.1"
+	peer1endpoint, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", localIP1, peer1wgPort))
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -546,7 +544,8 @@ func Test_ConnectPeers(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	peer2endpoint, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", localIP, peer2wgPort))
+	localIP2 := "127.0.0.1"
+	peer2endpoint, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", localIP2, peer2wgPort))
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -569,17 +568,17 @@ func Test_ConnectPeers(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	// todo: investigate why in some tests execution we need 30s
+	// The peers use userspace WireGuard (stdnet transport). A tight busy-loop
+	// here starves the wireguard-go goroutines that process the handshake, so
+	// poll on a ticker instead and yield the CPU between checks. WireGuard also
+	// only retries a lost handshake initiation every REKEY_TIMEOUT (5s), which
+	// is why the overall wait can occasionally stretch to tens of seconds.
 	timeout := 30 * time.Second
 	timeoutChannel := time.After(timeout)
+	ticker := time.NewTicker(500 * time.Millisecond)
+	defer ticker.Stop()
 
 	for {
-		select {
-		case <-timeoutChannel:
-			t.Fatalf("waiting for peer handshake timeout after %s", timeout.String())
-		default:
-		}
-
 		peer, gpErr := getPeer(peer1ifaceName, peer2Key.PublicKey().String())
 		if gpErr != nil {
 			t.Fatal(gpErr)
@@ -588,6 +587,12 @@ func Test_ConnectPeers(t *testing.T) {
 			t.Log("peers successfully handshake")
 			break
 		}
+
+		select {
+		case <-timeoutChannel:
+			t.Fatalf("waiting for peer handshake timeout after %s", timeout.String())
+		case <-ticker.C:
+		}
 	}
 
 }
@@ -615,28 +620,3 @@ func getPeer(ifaceName, peerPubKey string) (wgtypes.Peer, error) {
 	}
 	return wgtypes.Peer{}, fmt.Errorf("peer not found")
 }
-
-func getLocalIP() (string, error) {
-	// Get all interfaces
-	addrs, err := net.InterfaceAddrs()
-	if err != nil {
-		return "", err
-	}
-
-	for _, addr := range addrs {
-		ipNet, ok := addr.(*net.IPNet)
-		if !ok {
-			continue
-		}
-		if ipNet.IP.IsLoopback() {
-			continue
-		}
-
-		if ipNet.IP.To4() == nil {
-			continue
-		}
-		return ipNet.IP.String(), nil
-	}
-
-	return "", fmt.Errorf("no local IP found")
-}

client/internal/auth/auth.go

@@ -351,6 +351,7 @@ func (a *Auth) setSystemInfoFlags(info *system.Info) {
 		a.config.BlockLANAccess,
 		a.config.BlockInbound,
 		a.config.DisableIPv6,
+		a.config.SyncMessageVersion,
 		a.config.EnableSSHRoot,
 		a.config.EnableSSHSFTP,
 		a.config.EnableSSHLocalPortForwarding,

client/internal/conn_mgr.go

@@ -34,6 +34,8 @@ const (
 // - Handling connection establishment based on peer signaling
 //
 // The implementation is not thread-safe; it is protected by engine.syncMsgMux.
+// The only exception is ActivatePeer, which is safe for concurrent use so the
+// DNS warm-up path can call it without contending on the engine mutex.
 type ConnMgr struct {
 	peerStore        *peerstore.Store
 	statusRecorder   *peer.Status
@@ -42,12 +44,26 @@ type ConnMgr struct {
 	rosenpassEnabled bool
 
 	lazyConnMgr *manager.Manager
+	// lazyConnMgrMu guards the lazyConnMgr pointer for readers outside the
+	// engine loop (ActivatePeer). Writers hold it in addition to
+	// engine.syncMsgMux; all other reads stay under engine.syncMsgMux only.
+	lazyConnMgrMu sync.RWMutex
+
+	// reconcileRoutedIPs re-applies a peer's routed allowed IPs after its lazy wake endpoint is
+	// (re)armed (Mode A at arm time). Injected by the engine; nil disables the reconcile.
+	reconcileRoutedIPs func(peerKey string) error
 
 	wg            sync.WaitGroup
 	lazyCtx       context.Context
 	lazyCtxCancel context.CancelFunc
 }
 
+// SetRoutedIPsReconciler injects the callback used to re-apply a peer's routed allowed IPs when
+// its lazy wake endpoint is (re)armed. Must be called before the lazy manager starts.
+func (e *ConnMgr) SetRoutedIPsReconciler(fn func(peerKey string) error) {
+	e.reconcileRoutedIPs = fn
+}
+
 func NewConnMgr(engineConfig *EngineConfig, statusRecorder *peer.Status, peerStore *peerstore.Store, iface lazyconn.WGIface) *ConnMgr {
 	e := &ConnMgr{
 		peerStore:        peerStore,
@@ -238,12 +254,20 @@ func (e *ConnMgr) RemovePeerConn(peerKey string) {
 	conn.Log.Infof("removed peer from lazy conn manager")
 }
 
+// ActivatePeer wakes an idle lazy connection. Unlike the rest of ConnMgr it is
+// safe for concurrent use: the lazy manager pointer is read under lazyConnMgrMu
+// and the manager itself is internally synchronized, so callers outside the
+// engine loop (DNS warm-up) do not need engine.syncMsgMux.
 func (e *ConnMgr) ActivatePeer(ctx context.Context, conn *peer.Conn) {
-	if !e.isStartedWithLazyMgr() {
+	e.lazyConnMgrMu.RLock()
+	lazyConnMgr := e.lazyConnMgr
+	started := lazyConnMgr != nil && e.lazyCtxCancel != nil
+	e.lazyConnMgrMu.RUnlock()
+	if !started {
 		return
 	}
 
-	if found := e.lazyConnMgr.ActivatePeer(conn.GetKey()); found {
+	if found := lazyConnMgr.ActivatePeer(conn.GetKey()); found {
 		if err := conn.Open(ctx); err != nil {
 			conn.Log.Errorf("failed to open connection: %v", err)
 		}
@@ -268,16 +292,22 @@ func (e *ConnMgr) Close() {
 
 	e.lazyCtxCancel()
 	e.wg.Wait()
+
+	e.lazyConnMgrMu.Lock()
 	e.lazyConnMgr = nil
+	e.lazyConnMgrMu.Unlock()
 }
 
 func (e *ConnMgr) initLazyManager(engineCtx context.Context) {
 	cfg := manager.Config{
 		InactivityThreshold: inactivityThresholdEnv(),
+		ReconcileAllowedIPs: e.reconcileRoutedIPs,
 	}
-	e.lazyConnMgr = manager.NewManager(cfg, engineCtx, e.peerStore, e.iface)
 
+	e.lazyConnMgrMu.Lock()
+	e.lazyConnMgr = manager.NewManager(cfg, engineCtx, e.peerStore, e.iface)
 	e.lazyCtx, e.lazyCtxCancel = context.WithCancel(engineCtx)
+	e.lazyConnMgrMu.Unlock()
 
 	e.wg.Add(1)
 	go func() {
@@ -316,7 +346,10 @@ func (e *ConnMgr) closeManager(ctx context.Context) {
 
 	e.lazyCtxCancel()
 	e.wg.Wait()
+
+	e.lazyConnMgrMu.Lock()
 	e.lazyConnMgr = nil
+	e.lazyConnMgrMu.Unlock()
 
 	for _, peerID := range e.peerStore.PeersPubKey() {
 		e.peerStore.PeerConnOpen(ctx, peerID)

client/internal/conn_mgr_test.go

@@ -1,10 +1,21 @@
 package internal
 
 import (
+	"context"
+	"net"
+	"net/netip"
 	"os"
+	"sync"
 	"testing"
+	"time"
 
+	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
+
+	"github.com/netbirdio/netbird/client/iface/wgaddr"
 	"github.com/netbirdio/netbird/client/internal/lazyconn"
+	"github.com/netbirdio/netbird/client/internal/peer"
+	"github.com/netbirdio/netbird/client/internal/peerstore"
+	"github.com/netbirdio/netbird/monotime"
 )
 
 func TestResolveLazyForce(t *testing.T) {
@@ -38,3 +49,58 @@ func TestResolveLazyForce(t *testing.T) {
 		})
 	}
 }
+
+type mockLazyWGIface struct{}
+
+func (mockLazyWGIface) RemovePeer(string) error { return nil }
+func (mockLazyWGIface) UpdatePeer(string, []netip.Prefix, time.Duration, *net.UDPAddr, *wgtypes.Key) error {
+	return nil
+}
+func (mockLazyWGIface) IsUserspaceBind() bool                    { return false }
+func (mockLazyWGIface) Address() wgaddr.Address                  { return wgaddr.Address{} }
+func (mockLazyWGIface) LastActivities() map[string]monotime.Time { return nil }
+func (mockLazyWGIface) MTU() uint16                              { return 1280 }
+
+// TestConnMgr_ActivatePeerConcurrentWithLifecycle exercises ActivatePeer from
+// non-engine goroutines (the DNS warm-up path) racing the manager lifecycle,
+// which stays on the engine loop. Run with -race: it fails if ActivatePeer
+// still requires engine.syncMsgMux for safety.
+func TestConnMgr_ActivatePeerConcurrentWithLifecycle(t *testing.T) {
+	t.Setenv(lazyconn.EnvLazyConn, "on")
+
+	status := peer.NewRecorder("https://mgm")
+	store := peerstore.NewConnStore()
+	connMgr := NewConnMgr(&EngineConfig{}, status, store, mockLazyWGIface{})
+
+	conn := newTestPeerConn(t, "peerA")
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	connMgr.Start(ctx)
+
+	done := make(chan struct{})
+	var wg sync.WaitGroup
+	for range 4 {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for {
+				select {
+				case <-done:
+					return
+				default:
+					connMgr.ActivatePeer(ctx, conn)
+				}
+			}
+		}()
+	}
+
+	// Let the activators spin against the started manager, then tear it down
+	// underneath them and let them spin against the stopped manager.
+	time.Sleep(100 * time.Millisecond)
+	connMgr.Close()
+	time.Sleep(50 * time.Millisecond)
+
+	close(done)
+	wg.Wait()
+}

client/internal/connect.go

@@ -621,6 +621,7 @@ func createEngineConfig(key wgtypes.Key, config *profilemanager.Config, peerConf
 		BlockLANAccess:      config.BlockLANAccess,
 		BlockInbound:        config.BlockInbound,
 		DisableIPv6:         config.DisableIPv6,
+		SyncMessageVersion:  config.SyncMessageVersion,
 
 		LazyConnection: lazyconn.ParseState(config.LazyConnection),
 
@@ -696,6 +697,7 @@ func loginToManagement(ctx context.Context, client mgm.Client, pubSSHKey []byte,
 		config.BlockLANAccess,
 		config.BlockInbound,
 		config.DisableIPv6,
+		config.SyncMessageVersion,
 		config.EnableSSHRoot,
 		config.EnableSSHSFTP,
 		config.EnableSSHLocalPortForwarding,

client/internal/debug/debug.go

@@ -676,6 +676,7 @@ func (g *BundleGenerator) addCommonConfigFields(configContent *strings.Builder)
 	configContent.WriteString(fmt.Sprintf("BlockLANAccess: %v\n", g.internalConfig.BlockLANAccess))
 	configContent.WriteString(fmt.Sprintf("BlockInbound: %v\n", g.internalConfig.BlockInbound))
 	configContent.WriteString(fmt.Sprintf("DisableIPv6: %v\n", g.internalConfig.DisableIPv6))
+	configContent.WriteString(fmt.Sprintf("SyncMessageVersion: %v\n", g.internalConfig.SyncMessageVersion))
 
 	if g.internalConfig.DisableNotifications != nil {
 		configContent.WriteString(fmt.Sprintf("DisableNotifications: %v\n", *g.internalConfig.DisableNotifications))

client/internal/debug/debug_test.go

@@ -887,6 +887,8 @@ func TestAddConfig_AllFieldsCovered(t *testing.T) {
 		ClientCertKeyPath:             "/tmp/key",
 		LazyConnection:                "on",
 		MTU:                           1280,
+		DisableIPv6:                   true,
+		SyncMessageVersion:            func(v int) *int { return &v }(1),
 	}
 
 	for _, anonymize := range []bool{false, true} {

client/internal/dns/local/local.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"net"
 	"net/netip"
+	"os"
 	"slices"
 	"strings"
 	"sync"
@@ -36,7 +37,43 @@ type resolver interface {
 // record is left alone (it points at something outside our mesh, e.g.
 // a non-peer upstream).
 type PeerConnectivity interface {
-	IsConnectedByIP(ip string) (known, connected bool)
+	IsConnectedByIP(ip netip.Addr) (known, connected bool)
+}
+
+// PeerActivator wakes lazy-connection peers on demand. The local resolver calls
+// it with the tunnel IPs an answer points at, so a peer that is idle (lazily
+// disconnected) starts connecting at DNS-resolution time rather than racing the
+// client's first request packet. nil disables warm-up.
+type PeerActivator interface {
+	// ActivatePeersByIP triggers wake-up for the peer(s) owning addrs and blocks
+	// until one is connected or ctx (a short per-query budget) expires. It is a
+	// fast no-op for unknown or already-connected addresses.
+	ActivatePeersByIP(ctx context.Context, addrs []netip.Addr)
+}
+
+const (
+	defaultLazyWarmupTimeout = 2 * time.Second
+	envLazyWarmupTimeout     = "NB_DNS_LAZY_WARMUP_TIMEOUT"
+)
+
+// lazyWarmupTimeoutFromEnv returns the per-query budget for waking a
+// lazy-connection peer a DNS answer points at. Tunable via
+// NB_DNS_LAZY_WARMUP_TIMEOUT (a Go duration). Parsed once at construction time.
+func lazyWarmupTimeoutFromEnv() time.Duration {
+	v := os.Getenv(envLazyWarmupTimeout)
+	if v == "" {
+		return defaultLazyWarmupTimeout
+	}
+	d, err := time.ParseDuration(v)
+	if err != nil {
+		log.Warnf("invalid %s value %q, using default %s: %v", envLazyWarmupTimeout, v, defaultLazyWarmupTimeout, err)
+		return defaultLazyWarmupTimeout
+	}
+	if d <= 0 {
+		log.Warnf("non-positive %s value %q, using default %s", envLazyWarmupTimeout, v, defaultLazyWarmupTimeout)
+		return defaultLazyWarmupTimeout
+	}
+	return d
 }
 
 type Resolver struct {
@@ -51,6 +88,12 @@ type Resolver struct {
 	// filter and preserves the legacy "return whatever is registered"
 	// behaviour for callers that never wire a status source.
 	peerConn PeerConnectivity
+	// peerActivator, when non-nil, is called at resolution time to warm the
+	// lazy connection to the peer(s) an answer points at. nil disables warm-up.
+	peerActivator PeerActivator
+	// warmupTimeout is the per-query budget for the lazy-connection warm-up
+	// wait, resolved from the environment once at construction time.
+	warmupTimeout time.Duration
 
 	ctx    context.Context
 	cancel context.CancelFunc
@@ -59,11 +102,12 @@ type Resolver struct {
 func NewResolver() *Resolver {
 	ctx, cancel := context.WithCancel(context.Background())
 	return &Resolver{
-		records: make(map[dns.Question][]dns.RR),
-		domains: make(map[domain.Domain]struct{}),
-		zones:   make(map[domain.Domain]bool),
-		ctx:     ctx,
-		cancel:  cancel,
+		records:       make(map[dns.Question][]dns.RR),
+		domains:       make(map[domain.Domain]struct{}),
+		zones:         make(map[domain.Domain]bool),
+		warmupTimeout: lazyWarmupTimeoutFromEnv(),
+		ctx:           ctx,
+		cancel:        cancel,
 	}
 }
 
@@ -76,6 +120,14 @@ func (d *Resolver) SetPeerConnectivity(p PeerConnectivity) {
 	d.peerConn = p
 }
 
+// SetPeerActivator wires the DNS-time lazy-connection warm-up. Pass nil to
+// disable. Safe to call multiple times; the latest value wins.
+func (d *Resolver) SetPeerActivator(a PeerActivator) {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+	d.peerActivator = a
+}
+
 func (d *Resolver) MatchSubdomains() bool {
 	return true
 }
@@ -122,6 +174,9 @@ func (d *Resolver) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
 	replyMessage.RecursionAvailable = true
 
 	result := d.lookupRecords(logger, question)
+	// Warm before filtering: activation flips a lazily-idle target to connected,
+	// which then lets it survive the disconnected-peer filter below.
+	d.warmLazyPeers(question, result.records)
 	result.records = d.filterDisconnectedPeerAnswers(logger, question, result.records)
 	replyMessage.Authoritative = !result.hasExternalData
 	replyMessage.Answer = result.records
@@ -495,8 +550,8 @@ func (d *Resolver) filterDisconnectedPeerAnswers(logger *log.Entry, question dns
 	kept := make([]dns.RR, 0, len(records))
 	var dropped int
 	for _, rr := range records {
-		ip := extractRecordIP(rr)
-		if ip == "" {
+		ip, ok := extractRecordAddr(rr)
+		if !ok {
 			kept = append(kept, rr)
 			continue
 		}
@@ -518,22 +573,57 @@ func (d *Resolver) filterDisconnectedPeerAnswers(logger *log.Entry, question dns
 	return kept
 }
 
-// extractRecordIP returns the dotted-decimal / colon-hex IP carried by
-// an A or AAAA record, or "" for any other record type.
-func extractRecordIP(rr dns.RR) string {
+// warmLazyPeers triggers lazy-connection wake-up for the peers a resolved
+// answer points at and waits briefly for one to connect, so the caller's first
+// request doesn't race the connection establishment. Warm-up is scoped to
+// match-only (non-authoritative) zones — the synthesized private-service zones
+// and user-created zones whose records point at specific peers. The account's
+// peer zone is authoritative, so plain peer-name lookups never trigger warm-up;
+// otherwise resolving any peer's name would wake its idle connection, defeating
+// laziness mesh-wide. No-op when no activator is wired (lazy connections
+// disabled) or the answer carries no peer IPs.
+func (d *Resolver) warmLazyPeers(question dns.Question, records []dns.RR) {
+	if len(records) < 2 {
+		return
+	}
+	d.mu.RLock()
+	activator := d.peerActivator
+	var nonAuth, found bool
+	if activator != nil {
+		nonAuth, found = d.findZone(question.Name)
+	}
+	d.mu.RUnlock()
+	if activator == nil || !found || !nonAuth {
+		return
+	}
+
+	var addrs []netip.Addr
+	for _, rr := range records {
+		if addr, ok := extractRecordAddr(rr); ok {
+			addrs = append(addrs, addr)
+		}
+	}
+	if len(addrs) == 0 {
+		return
+	}
+
+	ctx, cancel := context.WithTimeout(d.ctx, d.warmupTimeout)
+	defer cancel()
+	activator.ActivatePeersByIP(ctx, addrs)
+}
+
+// extractRecordAddr returns the IP address carried by an A or AAAA record.
+// ok is false for any other record type or a record with no address.
+func extractRecordAddr(rr dns.RR) (netip.Addr, bool) {
 	switch r := rr.(type) {
 	case *dns.A:
-		if r.A == nil {
-			return ""
-		}
-		return r.A.String()
+		addr, ok := netip.AddrFromSlice(r.A)
+		return addr.Unmap(), ok
 	case *dns.AAAA:
-		if r.AAAA == nil {
-			return ""
-		}
-		return r.AAAA.String()
+		addr, ok := netip.AddrFromSlice(r.AAAA)
+		return addr.Unmap(), ok
 	}
-	return ""
+	return netip.Addr{}, false
 }
 
 // Update replaces all zones and their records

client/internal/dns/local/local_test.go

@@ -37,8 +37,8 @@ type mockPeerConnectivity struct {
 	byIP map[string]struct{ known, connected bool }
 }
 
-func (m mockPeerConnectivity) IsConnectedByIP(ip string) (known, connected bool) {
-	v, ok := m.byIP[ip]
+func (m mockPeerConnectivity) IsConnectedByIP(ip netip.Addr) (known, connected bool) {
+	v, ok := m.byIP[ip.String()]
 	if !ok {
 		return false, false
 	}

client/internal/dns/local/warmup_test.go

@@ -0,0 +1,204 @@
+package local
+
+import (
+	"context"
+	"net"
+	"net/netip"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/miekg/dns"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/client/internal/dns/test"
+	nbdns "github.com/netbirdio/netbird/dns"
+)
+
+// recordingActivator records the addresses it was asked to warm and returns
+// immediately, so ServeDNS is not blocked by the test.
+type recordingActivator struct {
+	mu     sync.Mutex
+	called bool
+	addrs  []netip.Addr
+}
+
+func (r *recordingActivator) ActivatePeersByIP(_ context.Context, addrs []netip.Addr) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.called = true
+	r.addrs = append(r.addrs, addrs...)
+}
+
+func serveA(t *testing.T, resolver *Resolver, name string) *dns.Msg {
+	t.Helper()
+	var resp *dns.Msg
+	w := &test.MockResponseWriter{WriteMsgFunc: func(m *dns.Msg) error { resp = m; return nil }}
+	resolver.ServeDNS(w, new(dns.Msg).SetQuestion(name, dns.TypeA))
+	return resp
+}
+
+// serviceZone registers rec in a match-only (non-authoritative) zone, the shape
+// the synthesized private-service zones arrive in.
+func serviceZone(t *testing.T, resolver *Resolver, zone string, records ...nbdns.SimpleRecord) {
+	t.Helper()
+	resolver.Update([]nbdns.CustomZone{{
+		Domain:           zone,
+		Records:          records,
+		NonAuthoritative: true,
+	}})
+}
+
+func TestLocalResolver_WarmsLazyPeerOnResolve(t *testing.T) {
+	// Warm-up fires only for multi-record answers (the HA / round-robin shape of
+	// the synthesized private-service zones), so register two peer targets.
+	const name = "svc.proxy.netbird.cloud."
+	recs := []nbdns.SimpleRecord{
+		{Name: name, Type: 1, Class: nbdns.DefaultClass, TTL: 300, RData: "100.64.0.7"},
+		{Name: name, Type: 1, Class: nbdns.DefaultClass, TTL: 300, RData: "100.64.0.8"},
+	}
+	resolver := NewResolver()
+	serviceZone(t, resolver, "proxy.netbird.cloud", recs...)
+
+	act := &recordingActivator{}
+	resolver.SetPeerActivator(act)
+
+	resp := serveA(t, resolver, name)
+	require.NotNil(t, resp, "resolver must answer")
+	require.NotEmpty(t, resp.Answer, "answer must carry the A records")
+
+	act.mu.Lock()
+	defer act.mu.Unlock()
+	assert.True(t, act.called, "activator must be invoked for a multi-record service-zone answer")
+	assert.Contains(t, act.addrs, netip.MustParseAddr("100.64.0.7"), "activator must receive the first peer IP")
+	assert.Contains(t, act.addrs, netip.MustParseAddr("100.64.0.8"), "activator must receive the second peer IP")
+}
+
+func TestLocalResolver_NoWarmupForSingleRecord(t *testing.T) {
+	// A single-record answer does not trigger warm-up; the resolver only warms
+	// multi-record answers.
+	rec := nbdns.SimpleRecord{Name: "svc.proxy.netbird.cloud.", Type: 1, Class: nbdns.DefaultClass, TTL: 300, RData: "100.64.0.7"}
+	resolver := NewResolver()
+	serviceZone(t, resolver, "proxy.netbird.cloud", rec)
+
+	act := &recordingActivator{}
+	resolver.SetPeerActivator(act)
+
+	resp := serveA(t, resolver, rec.Name)
+	require.NotNil(t, resp, "resolver must answer")
+	require.NotEmpty(t, resp.Answer, "answer must carry the A record")
+
+	act.mu.Lock()
+	defer act.mu.Unlock()
+	assert.False(t, act.called, "activator must not be invoked for a single-record answer")
+}
+
+func TestLocalResolver_NoActivatorNoWarmup(t *testing.T) {
+	// With no activator wired the resolver behaves exactly as before.
+	rec := nbdns.SimpleRecord{Name: "svc.proxy.netbird.cloud.", Type: 1, Class: nbdns.DefaultClass, TTL: 300, RData: "100.64.0.7"}
+	resolver := NewResolver()
+	serviceZone(t, resolver, "proxy.netbird.cloud", rec)
+
+	resp := serveA(t, resolver, rec.Name)
+	require.NotNil(t, resp, "resolver must still answer without an activator")
+	require.NotEmpty(t, resp.Answer, "answer must carry the A record")
+}
+
+func TestLocalResolver_NoWarmupForMissingRecord(t *testing.T) {
+	// A query that resolves to nothing must not invoke the activator (no IPs).
+	resolver := NewResolver()
+	serviceZone(t, resolver, "proxy.netbird.cloud",
+		nbdns.SimpleRecord{Name: "svc.proxy.netbird.cloud.", Type: 1, Class: nbdns.DefaultClass, TTL: 300, RData: "100.64.0.7"})
+
+	act := &recordingActivator{}
+	resolver.SetPeerActivator(act)
+
+	serveA(t, resolver, "absent.proxy.netbird.cloud.")
+
+	act.mu.Lock()
+	defer act.mu.Unlock()
+	assert.False(t, act.called, "activator must not be invoked when there is no answer")
+}
+
+func TestLocalResolver_NoWarmupInAuthoritativeZone(t *testing.T) {
+	// The account's peer zone is authoritative; resolving a peer's name there
+	// must not wake its lazy connection — warm-up is scoped to match-only
+	// (non-authoritative) zones such as the synthesized private-service zones.
+	// Use a multi-record answer so the authoritative-zone scoping is the only
+	// reason warm-up is skipped, not the single-record guard.
+	const name = "peer.netbird.cloud."
+	recs := []nbdns.SimpleRecord{
+		{Name: name, Type: 1, Class: nbdns.DefaultClass, TTL: 300, RData: "100.64.0.9"},
+		{Name: name, Type: 1, Class: nbdns.DefaultClass, TTL: 300, RData: "100.64.0.10"},
+	}
+	resolver := NewResolver()
+	resolver.Update([]nbdns.CustomZone{{
+		Domain:  "netbird.cloud",
+		Records: recs,
+	}})
+
+	act := &recordingActivator{}
+	resolver.SetPeerActivator(act)
+
+	resp := serveA(t, resolver, name)
+	require.NotNil(t, resp, "resolver must answer")
+	require.NotEmpty(t, resp.Answer, "answer must carry the A records")
+
+	act.mu.Lock()
+	defer act.mu.Unlock()
+	assert.False(t, act.called, "activator must not be invoked for authoritative-zone answers")
+}
+
+func TestLazyWarmupTimeoutFromEnv(t *testing.T) {
+	tests := []struct {
+		name   string
+		value  string
+		envSet bool
+		want   time.Duration
+	}{
+		{name: "unset uses default", want: defaultLazyWarmupTimeout},
+		{name: "valid overrides", value: "5s", envSet: true, want: 5 * time.Second},
+		{name: "invalid falls back", value: "not-a-duration", envSet: true, want: defaultLazyWarmupTimeout},
+		{name: "negative falls back", value: "-1s", envSet: true, want: defaultLazyWarmupTimeout},
+		{name: "zero falls back", value: "0s", envSet: true, want: defaultLazyWarmupTimeout},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if tt.envSet {
+				t.Setenv(envLazyWarmupTimeout, tt.value)
+			}
+			assert.Equal(t, tt.want, lazyWarmupTimeoutFromEnv())
+			assert.Equal(t, tt.want, NewResolver().warmupTimeout, "constructor must resolve the timeout once")
+		})
+	}
+}
+
+func TestExtractRecordAddr(t *testing.T) {
+	t.Run("A record yields unmapped v4", func(t *testing.T) {
+		// net.ParseIP returns the 16-byte v4-in-v6 form, the same shape
+		// miekg/dns stores after parsing an A record; the extracted address
+		// must compare equal to a plain v4 netip.Addr.
+		addr, ok := extractRecordAddr(&dns.A{A: net.ParseIP("100.64.0.7")})
+		require.True(t, ok)
+		assert.True(t, addr.Is4())
+		assert.Equal(t, netip.MustParseAddr("100.64.0.7"), addr)
+	})
+
+	t.Run("AAAA record yields v6", func(t *testing.T) {
+		addr, ok := extractRecordAddr(&dns.AAAA{AAAA: net.ParseIP("fd00::1")})
+		require.True(t, ok)
+		assert.Equal(t, netip.MustParseAddr("fd00::1"), addr)
+	})
+
+	t.Run("A record without address", func(t *testing.T) {
+		_, ok := extractRecordAddr(&dns.A{})
+		assert.False(t, ok)
+	})
+
+	t.Run("non-address record", func(t *testing.T) {
+		_, ok := extractRecordAddr(&dns.CNAME{Target: "target.netbird.cloud."})
+		assert.False(t, ok)
+	})
+}

client/internal/dns/mock_server.go

@@ -8,6 +8,7 @@ import (
 	"github.com/miekg/dns"
 
 	dnsconfig "github.com/netbirdio/netbird/client/internal/dns/config"
+	"github.com/netbirdio/netbird/client/internal/dns/local"
 	nbdns "github.com/netbirdio/netbird/dns"
 	"github.com/netbirdio/netbird/route"
 	"github.com/netbirdio/netbird/shared/management/domain"
@@ -92,6 +93,11 @@ func (m *MockServer) SetFirewall(Firewall) {
 	// Mock implementation - no-op
 }
 
+// SetPeerActivator mock implementation of SetPeerActivator from Server interface
+func (m *MockServer) SetPeerActivator(local.PeerActivator) {
+	// Mock implementation - no-op
+}
+
 // BeginBatch mock implementation of BeginBatch from Server interface
 func (m *MockServer) BeginBatch() {
 	// Mock implementation - no-op

client/internal/dns/server.go

@@ -82,6 +82,7 @@ type Server interface {
 	PopulateManagementDomain(mgmtURL *url.URL) error
 	SetRouteSources(selected, active func() route.HAMap)
 	SetFirewall(Firewall)
+	SetPeerActivator(local.PeerActivator)
 }
 
 type nsGroupsByDomain struct {
@@ -491,6 +492,13 @@ func (s *DefaultServer) SetFirewall(fw Firewall) {
 	}
 }
 
+// SetPeerActivator wires the DNS-time lazy-connection warm-up on the local
+// resolver. Injected after the connection manager exists (it does not at
+// DNS-server construction time). Pass nil to disable.
+func (s *DefaultServer) SetPeerActivator(a local.PeerActivator) {
+	s.localResolver.SetPeerActivator(a)
+}
+
 // Stop stops the server
 func (s *DefaultServer) Stop() {
 	s.ctxCancel()
@@ -1435,11 +1443,11 @@ type localPeerConnectivity struct {
 
 // IsConnectedByIP looks the IP up in the peerstore and surfaces both
 // the known and connected bits. Used by Resolver.filterDisconnectedPeerAnswers.
-func (l localPeerConnectivity) IsConnectedByIP(ip string) (known, connected bool) {
+func (l localPeerConnectivity) IsConnectedByIP(ip netip.Addr) (known, connected bool) {
 	if l.status == nil {
 		return false, false
 	}
-	state, ok := l.status.PeerStateByIP(ip)
+	state, ok := l.status.PeerStateByIP(ip.String())
 	if !ok {
 		return false, false
 	}

client/internal/dns_peer_activator.go

@@ -0,0 +1,76 @@
+package internal
+
+import (
+	"context"
+	"net/netip"
+	"time"
+
+	"github.com/netbirdio/netbird/client/internal/peer"
+	"github.com/netbirdio/netbird/client/internal/peerstore"
+)
+
+const dnsActivationPollInterval = 50 * time.Millisecond
+
+// dnsPeerActivator wakes lazy-connection peers from the DNS resolution path. It
+// implements dns/local.PeerActivator. DNS queries run on their own goroutines,
+// so it only touches state that is safe for concurrent use — ConnMgr.ActivatePeer,
+// peerstore.Store and peer.Status — and never takes the engine's syncMsgMux,
+// keeping DNS resolution from contending with network-map processing.
+type dnsPeerActivator struct {
+	connMgr   *ConnMgr
+	peerStore *peerstore.Store
+	status    *peer.Status
+	// ctx is the engine's long-lived context. The connection dial is tied to it
+	// (not the per-query DNS wait budget) so a handshake that outlasts the wait
+	// still completes in the background rather than being cancelled at the deadline.
+	ctx context.Context
+}
+
+// ActivatePeersByIP triggers wake-up for the peer(s) owning addrs and waits
+// until one is connected or ctx (the per-query DNS wait budget) expires.
+// Activation itself is tied to the engine's long-lived context so the dial
+// survives a wait that times out. Unknown or already-connected addresses are
+// skipped, so the steady-state (warm) path adds no latency.
+func (a *dnsPeerActivator) ActivatePeersByIP(ctx context.Context, addrs []netip.Addr) {
+	if a == nil || a.connMgr == nil {
+		return
+	}
+
+	var pending []string
+	for _, addr := range addrs {
+		ip := addr.String()
+		st, ok := a.status.PeerStateByIP(ip)
+		if !ok || st.ConnStatus == peer.StatusConnected {
+			continue
+		}
+		conn, ok := a.peerStore.PeerConn(st.PubKey)
+		if !ok {
+			continue
+		}
+		a.connMgr.ActivatePeer(a.ctx, conn)
+		pending = append(pending, ip)
+	}
+
+	if len(pending) == 0 {
+		return
+	}
+	a.waitConnected(ctx, pending)
+}
+
+// waitConnected blocks until any of ips reports a connected peer or ctx expires.
+func (a *dnsPeerActivator) waitConnected(ctx context.Context, ips []string) {
+	ticker := time.NewTicker(dnsActivationPollInterval)
+	defer ticker.Stop()
+	for {
+		for _, ip := range ips {
+			if st, ok := a.status.PeerStateByIP(ip); ok && st.ConnStatus == peer.StatusConnected {
+				return
+			}
+		}
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+		}
+	}
+}

client/internal/dns_peer_activator_test.go

@@ -0,0 +1,129 @@
+package internal
+
+import (
+	"context"
+	"net/netip"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/client/internal/peer"
+	"github.com/netbirdio/netbird/client/internal/peerstore"
+)
+
+func newTestPeerConn(t *testing.T, key string) *peer.Conn {
+	t.Helper()
+	conn, err := peer.NewConn(peer.ConnConfig{
+		Key:      key,
+		LocalKey: "local",
+		WgConfig: peer.WgConfig{
+			AllowedIps: []netip.Prefix{netip.MustParsePrefix("100.64.0.1/32")},
+		},
+	}, peer.ServiceDependencies{})
+	require.NoError(t, err)
+	return conn
+}
+
+func newTestDNSPeerActivator(t *testing.T) (*dnsPeerActivator, *peer.Status, *peerstore.Store) {
+	t.Helper()
+	status := peer.NewRecorder("https://mgm")
+	store := peerstore.NewConnStore()
+	// ConnMgr without Start: the lazy manager is nil, so ActivatePeer is a
+	// no-op — these tests exercise the activator's skip/wait logic.
+	connMgr := NewConnMgr(&EngineConfig{}, status, store, nil)
+	return &dnsPeerActivator{
+		connMgr:   connMgr,
+		peerStore: store,
+		status:    status,
+		ctx:       context.Background(),
+	}, status, store
+}
+
+func TestDNSPeerActivator_NilSafe(t *testing.T) {
+	var a *dnsPeerActivator
+	a.ActivatePeersByIP(context.Background(), []netip.Addr{netip.MustParseAddr("100.64.0.1")})
+}
+
+// TestDNSPeerActivator_SkipsUnknownAndConnectedPeers verifies the steady-state
+// (warm) path adds no latency: already-connected and unknown addresses never
+// enter the wait loop.
+func TestDNSPeerActivator_SkipsUnknownAndConnectedPeers(t *testing.T) {
+	a, status, store := newTestDNSPeerActivator(t)
+
+	require.NoError(t, status.AddPeer("peerA", "a.netbird.cloud", "100.64.0.1", "fd00::1"))
+	require.NoError(t, status.UpdatePeerState(peer.State{PubKey: "peerA", ConnStatus: peer.StatusConnected}))
+	store.AddPeerConn("peerA", newTestPeerConn(t, "peerA"))
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	start := time.Now()
+	a.ActivatePeersByIP(ctx, []netip.Addr{
+		netip.MustParseAddr("100.64.0.1"),  // known, connected -> skipped
+		netip.MustParseAddr("fd00::1"),     // known via IPv6, connected -> skipped
+		netip.MustParseAddr("100.64.0.99"), // unknown -> skipped
+	})
+	require.Less(t, time.Since(start), time.Second, "no pending peer must mean no wait")
+}
+
+// TestDNSPeerActivator_WaitsForPendingPeerToConnect verifies the wait loop
+// returns as soon as a pending peer reports connected, well before the
+// per-query budget expires.
+func TestDNSPeerActivator_WaitsForPendingPeerToConnect(t *testing.T) {
+	a, status, store := newTestDNSPeerActivator(t)
+
+	require.NoError(t, status.AddPeer("peerA", "a.netbird.cloud", "100.64.0.1", ""))
+	store.AddPeerConn("peerA", newTestPeerConn(t, "peerA"))
+
+	go func() {
+		time.Sleep(150 * time.Millisecond)
+		_ = status.UpdatePeerState(peer.State{PubKey: "peerA", ConnStatus: peer.StatusConnected})
+	}()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	start := time.Now()
+	a.ActivatePeersByIP(ctx, []netip.Addr{netip.MustParseAddr("100.64.0.1")})
+	elapsed := time.Since(start)
+
+	require.GreaterOrEqual(t, elapsed, 100*time.Millisecond, "must wait for the pending peer")
+	require.Less(t, elapsed, 5*time.Second, "must return on connect, not at the deadline")
+}
+
+// TestDNSPeerActivator_ReturnsAtBudgetWhenPeerStaysIdle verifies a peer that
+// never connects releases the DNS response at the per-query budget instead of
+// blocking it indefinitely.
+func TestDNSPeerActivator_ReturnsAtBudgetWhenPeerStaysIdle(t *testing.T) {
+	a, status, store := newTestDNSPeerActivator(t)
+
+	require.NoError(t, status.AddPeer("peerA", "a.netbird.cloud", "100.64.0.1", ""))
+	store.AddPeerConn("peerA", newTestPeerConn(t, "peerA"))
+
+	ctx, cancel := context.WithTimeout(context.Background(), 300*time.Millisecond)
+	defer cancel()
+
+	start := time.Now()
+	a.ActivatePeersByIP(ctx, []netip.Addr{netip.MustParseAddr("100.64.0.1")})
+	elapsed := time.Since(start)
+
+	require.GreaterOrEqual(t, elapsed, 250*time.Millisecond, "must wait out the budget for a pending peer")
+	require.Less(t, elapsed, 5*time.Second, "must not block past the budget")
+}
+
+// TestDNSPeerActivator_NoWaitWithoutPeerConn verifies a known-but-idle peer
+// with no connection object in the store is not waited on: there is nothing to
+// activate, so waiting could only ever time out.
+func TestDNSPeerActivator_NoWaitWithoutPeerConn(t *testing.T) {
+	a, status, _ := newTestDNSPeerActivator(t)
+
+	require.NoError(t, status.AddPeer("peerA", "a.netbird.cloud", "100.64.0.1", ""))
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	start := time.Now()
+	a.ActivatePeersByIP(ctx, []netip.Addr{netip.MustParseAddr("100.64.0.1")})
+	require.Less(t, time.Since(start), time.Second, "peer without a conn must not be waited on")
+}

client/internal/ebpf/ebpf/src/dns_fwd.c

@@ -52,11 +52,14 @@ int xdp_dns_fwd(struct iphdr  *ip, struct udphdr *udp) {
 
     if (udp->dest == GENERAL_DNS_PORT && ip->daddr == dns_ip) {
         udp->dest = dns_port;
+        // Clear the now-stale checksum; zero means "not computed" for IPv4.
+        udp->check = 0;
         return XDP_PASS;
     }
 
     if (udp->source == dns_port && ip->saddr == dns_ip) {
         udp->source = GENERAL_DNS_PORT;
+        udp->check = 0;
         return XDP_PASS;
     }
 

client/internal/ebpf/ebpf/src/wg_proxy.c

@@ -50,5 +50,11 @@ int xdp_wg_proxy(struct iphdr  *ip, struct udphdr *udp) {
     __be16 new_dst_port = htons(proxy_port);
     udp->dest = new_dst_port;
     udp->source = new_src_port;
+
+    // The ports are covered by the UDP checksum. This is an IPv4 loopback hop
+    // and the payload is already integrity-protected, so clear the checksum (a
+    // zero UDP checksum means "not computed" for IPv4) rather than leave a
+    // stale value the kernel would drop as UDP_CSUM.
+    udp->check = 0;
 	return XDP_PASS;
 }

client/internal/engine.go

@@ -64,7 +64,10 @@ import (
 	"github.com/netbirdio/netbird/route"
 	mgm "github.com/netbirdio/netbird/shared/management/client"
 	"github.com/netbirdio/netbird/shared/management/domain"
+	sharedgrpc "github.com/netbirdio/netbird/shared/management/grpc"
+	nbnetworkmap "github.com/netbirdio/netbird/shared/management/networkmap"
 	mgmProto "github.com/netbirdio/netbird/shared/management/proto"
+	types "github.com/netbirdio/netbird/shared/management/types"
 	"github.com/netbirdio/netbird/shared/netiputil"
 	auth "github.com/netbirdio/netbird/shared/relay/auth/hmac"
 	relayClient "github.com/netbirdio/netbird/shared/relay/client"
@@ -147,6 +150,7 @@ type EngineConfig struct {
 	BlockLANAccess      bool
 	BlockInbound        bool
 	DisableIPv6         bool
+	SyncMessageVersion  *int
 
 	// LazyConnection is the MDM-sourced lazy-connection override; StateUnset defers to
 	// the env var and management feature flag.
@@ -220,6 +224,13 @@ type Engine struct {
 	// networkSerial is the latest CurrentSerial (state ID) of the network sent by the Management service
 	networkSerial uint64
 
+	// latestComponents is the most-recent NetworkMapComponents decoded from
+	// a NetworkMapEnvelope (capability=3 peers only). Held alongside the
+	// NetworkMap that Calculate() produced from it so future incremental
+	// updates have a base to apply changes against. nil for legacy-format
+	// peers. Guarded by syncMsgMux.
+	latestComponents *types.NetworkMapComponents
+
 	networkMonitor *networkmonitor.NetworkMonitor
 
 	sshServer sshServer
@@ -652,8 +663,24 @@ func (e *Engine) Start(netbirdConfig *mgmProto.NetbirdConfig, mgmtURL *url.URL)
 	iceCfg := e.createICEConfig()
 
 	e.connMgr = NewConnMgr(e.config, e.statusRecorder, e.peerStore, wgIface)
+	e.connMgr.SetRoutedIPsReconciler(func(peerKey string) error {
+		if e.routeManager == nil {
+			return nil
+		}
+		return e.routeManager.ReconcilePeerAllowedIPs(peerKey)
+	})
 	e.connMgr.Start(e.ctx)
 
+	// Wire DNS-time lazy-connection warm-up now that the connection manager
+	// exists (it does not at DNS-server construction time). A DNS answer that
+	// points at an idle peer then wakes it before the client's first request.
+	e.dnsServer.SetPeerActivator(&dnsPeerActivator{
+		connMgr:   e.connMgr,
+		peerStore: e.peerStore,
+		status:    e.statusRecorder,
+		ctx:       e.ctx,
+	})
+
 	e.srWatcher = guard.NewSRWatcher(e.signal, e.relayManager, e.mobileDep.IFaceDiscover, iceCfg)
 	e.srWatcher.Start(peer.IsForceRelayed())
 
@@ -963,8 +990,12 @@ func (e *Engine) handleSync(update *mgmProto.SyncResponse) error {
 
 	e.ApplySessionDeadline(update.GetSessionExpiresAt())
 
-	if update.NetworkMap != nil && update.NetworkMap.PeerConfig != nil {
-		e.handleAutoUpdateVersion(update.NetworkMap.PeerConfig.AutoUpdate)
+	// Envelope sync responses carry PeerConfig at the top level; legacy
+	// NetworkMap syncs carry it under NetworkMap.PeerConfig.
+	if pc := update.GetPeerConfig(); pc != nil {
+		e.handleAutoUpdateVersion(pc.GetAutoUpdate())
+	} else if nm := update.GetNetworkMap(); nm != nil && nm.GetPeerConfig() != nil {
+		e.handleAutoUpdateVersion(nm.GetPeerConfig().GetAutoUpdate())
 	}
 
 	done := e.phase("netbird_config")
@@ -974,12 +1005,47 @@ func (e *Engine) handleSync(update *mgmProto.SyncResponse) error {
 		return err
 	}
 
+	// Decode the network map from either the components envelope or the
+	// legacy proto.NetworkMap before the posture-check gating below, so the
+	// "is there a network map" decision covers both wire shapes.
+	var (
+		nm         *mgmProto.NetworkMap
+		components *types.NetworkMapComponents
+	)
+	if version := update.GetVersion(); version == int32(sharedgrpc.ComponentNetworkMap) {
+		// Components-format peer: decode the envelope back to typed
+		// components, run Calculate() locally, and convert to the wire
+		// NetworkMap shape the rest of the engine consumes. Components are
+		// retained so future incremental updates can apply deltas instead
+		// of doing a full reconstruction.
+		envelope := update.GetNetworkMapEnvelope()
+		if envelope == nil {
+			return fmt.Errorf("received a SyncReponse indicating use of components network map, but components are missing")
+		}
+
+		localKey := e.config.WgPrivateKey.PublicKey().String()
+		dnsName := ""
+		if pc := update.GetPeerConfig(); pc != nil {
+			// PeerConfig.Fqdn = "<dns_label>.<dns_domain>" — extract the
+			// shared domain by stripping the peer's own label prefix. Falls
+			// back to empty if the FQDN doesn't have the expected shape.
+			dnsName = extractDNSDomainFromFQDN(pc.GetFqdn())
+		}
+		result, err := nbnetworkmap.EnvelopeToNetworkMap(e.ctx, envelope, localKey, dnsName)
+		if err != nil {
+			return fmt.Errorf("decode network map envelope: %w", err)
+		}
+		nm = result.NetworkMap
+		components = result.Components
+	} else {
+		nm = update.GetNetworkMap()
+	}
+
 	// Posture checks are bound to the network map presence:
 	//   NetworkMap != nil, checks present -> apply the received checks
 	//   NetworkMap != nil, checks nil     -> posture checks were removed, clear them
 	//   NetworkMap == nil                 -> config-only update (e.g. relay token rotation),
 	//                                        leave the previously applied checks untouched
-	nm := update.GetNetworkMap()
 	if nm == nil {
 		return nil
 	}
@@ -992,6 +1058,14 @@ func (e *Engine) handleSync(update *mgmProto.SyncResponse) error {
 	}
 
 	done = e.phase("persist")
+	// Only retain the components view when the server sent the envelope
+	// path. A legacy proto.NetworkMap means components == nil; writing it
+	// here would clobber a previously-cached snapshot, breaking the
+	// incremental-delta base on a future envelope sync.
+	if components != nil {
+		e.latestComponents = components
+	}
+
 	e.persistSyncResponse(update)
 	done()
 
@@ -1005,6 +1079,19 @@ func (e *Engine) handleSync(update *mgmProto.SyncResponse) error {
 	return nil
 }
 
+// extractDNSDomainFromFQDN returns the trailing dotted domain part of the
+// receiving peer's FQDN — the same value the management server fills as
+// dnsName when it builds the legacy NetworkMap. "peer42.netbird.cloud" →
+// "netbird.cloud". An empty string is returned for unrecognized formats.
+func extractDNSDomainFromFQDN(fqdn string) string {
+	for i := 0; i < len(fqdn); i++ {
+		if fqdn[i] == '.' && i+1 < len(fqdn) {
+			return fqdn[i+1:]
+		}
+	}
+	return ""
+}
+
 // updateNetbirdConfig applies the management-provided NetBird configuration:
 // STUN/TURN and relay servers, flow logging and DNS settings. A nil config is a no-op,
 // which is the case for sync updates carrying only a network map.
@@ -1164,6 +1251,7 @@ func (e *Engine) applyInfoFlags(info *system.Info) {
 		e.config.BlockLANAccess,
 		e.config.BlockInbound,
 		e.config.DisableIPv6,
+		e.config.SyncMessageVersion,
 		e.config.EnableSSHRoot,
 		e.config.EnableSSHSFTP,
 		e.config.EnableSSHLocalPortForwarding,
@@ -2032,6 +2120,7 @@ func (e *Engine) readInitialSettings() ([]*route.Route, *nbdns.Config, bool, err
 		e.config.BlockLANAccess,
 		e.config.BlockInbound,
 		e.config.DisableIPv6,
+		e.config.SyncMessageVersion,
 		e.config.EnableSSHRoot,
 		e.config.EnableSSHSFTP,
 		e.config.EnableSSHLocalPortForwarding,

client/internal/lazyconn/manager/manager.go

@@ -29,6 +29,11 @@ type managedPeer struct {
 
 type Config struct {
 	InactivityThreshold *time.Duration
+	// ReconcileAllowedIPs re-applies a peer's routed allowed IPs after its wake endpoint is
+	// armed. The activity listener creates the wake peer with the overlay /32 only; without the
+	// routed prefixes WireGuard would not steer subnet-bound traffic to the wake endpoint, so an
+	// idle routing peer could never be woken by that traffic. Optional; nil disables the reconcile.
+	ReconcileAllowedIPs func(peerKey string) error
 }
 
 // Manager manages lazy connections
@@ -56,6 +61,9 @@ type Manager struct {
 	peerToHAGroups map[string][]route.HAUniqueID // peer ID -> HA groups they belong to
 	haGroupToPeers map[route.HAUniqueID][]string // HA group -> peer IDs in the group
 	routesMu       sync.RWMutex
+
+	// reconcileAllowedIPs re-applies a peer's routed allowed IPs after its wake endpoint is armed.
+	reconcileAllowedIPs func(peerKey string) error
 }
 
 // NewManager creates a new lazy connection manager
@@ -73,6 +81,7 @@ func NewManager(config Config, engineCtx context.Context, peerStore *peerstore.S
 		activityManager:      activity.NewManager(wgIface),
 		peerToHAGroups:       make(map[string][]route.HAUniqueID),
 		haGroupToPeers:       make(map[route.HAUniqueID][]string),
+		reconcileAllowedIPs:  config.ReconcileAllowedIPs,
 	}
 
 	if wgIface.IsUserspaceBind() {
@@ -201,7 +210,7 @@ func (m *Manager) AddPeer(peerCfg lazyconn.PeerConfig) (bool, error) {
 		return false, nil
 	}
 
-	if err := m.activityManager.MonitorPeerActivity(peerCfg); err != nil {
+	if err := m.armActivityListener(peerCfg); err != nil {
 		return false, err
 	}
 
@@ -288,7 +297,7 @@ func (m *Manager) DeactivatePeer(peerID peerid.ConnID) {
 
 	m.inactivityManager.RemovePeer(mp.peerCfg.PublicKey)
 
-	if err := m.activityManager.MonitorPeerActivity(*mp.peerCfg); err != nil {
+	if err := m.armActivityListener(*mp.peerCfg); err != nil {
 		mp.peerCfg.Log.Errorf("failed to create activity monitor: %v", err)
 		return
 	}
@@ -465,6 +474,31 @@ func (m *Manager) close() {
 }
 
 // shouldDeferIdleForHA checks if peer should stay connected due to HA group requirements
+// armRoutedAllowedIPs re-applies the peer's routed allowed IPs onto its freshly armed wake
+// endpoint. The activity listener creates the wake peer with the overlay /32 only, so without
+// this the routed prefixes would be missing and traffic to a routed subnet could not wake the
+// idle routing peer. It is a no-op when no reconciler is configured.
+// armActivityListener (re)arms the peer's wake endpoint via the activity manager and then
+// re-applies its routed allowed IPs, so traffic to a routed subnet can wake an idle routing
+// peer. The routed prefixes must be re-applied after the wake endpoint exists because the
+// listener creates it with the overlay /32 only.
+func (m *Manager) armActivityListener(peerCfg lazyconn.PeerConfig) error {
+	if err := m.activityManager.MonitorPeerActivity(peerCfg); err != nil {
+		return err
+	}
+	m.armRoutedAllowedIPs(&peerCfg)
+	return nil
+}
+
+func (m *Manager) armRoutedAllowedIPs(peerCfg *lazyconn.PeerConfig) {
+	if m.reconcileAllowedIPs == nil {
+		return
+	}
+	if err := m.reconcileAllowedIPs(peerCfg.PublicKey); err != nil {
+		peerCfg.Log.Errorf("failed to reconcile routed allowed IPs on wake endpoint: %v", err)
+	}
+}
+
 func (m *Manager) shouldDeferIdleForHA(inactivePeers map[string]struct{}, peerID string) bool {
 	m.routesMu.RLock()
 	defer m.routesMu.RUnlock()
@@ -577,7 +611,7 @@ func (m *Manager) onPeerInactivityTimedOut(peerIDs map[string]struct{}) {
 
 		mp.peerCfg.Log.Infof("start activity monitor")
 
-		if err := m.activityManager.MonitorPeerActivity(*mp.peerCfg); err != nil {
+		if err := m.armActivityListener(*mp.peerCfg); err != nil {
 			mp.peerCfg.Log.Errorf("failed to create activity monitor: %v", err)
 			continue
 		}

client/internal/profilemanager/config.go

@@ -96,6 +96,7 @@ type ConfigInput struct {
 	BlockLANAccess      *bool
 	BlockInbound        *bool
 	DisableIPv6         *bool
+	SyncMessageVersion  *int
 
 	DisableNotifications *bool
 
@@ -137,6 +138,7 @@ type Config struct {
 	BlockLANAccess      bool
 	BlockInbound        bool
 	DisableIPv6         bool
+	SyncMessageVersion  *int
 
 	DisableNotifications *bool
 
@@ -587,6 +589,12 @@ func (config *Config) apply(input ConfigInput) (updated bool, err error) {
 		updated = true
 	}
 
+	if input.SyncMessageVersion != nil && *input.SyncMessageVersion != *config.SyncMessageVersion {
+		log.Infof("setting SyncMessageVersion to %v", *input.SyncMessageVersion)
+		*config.SyncMessageVersion = *input.SyncMessageVersion
+		updated = true
+	}
+
 	if input.DisableNotifications != nil && (config.DisableNotifications == nil || *input.DisableNotifications != *config.DisableNotifications) {
 		if *input.DisableNotifications {
 			log.Infof("disabling notifications")

client/internal/routemanager/manager.go

@@ -52,6 +52,10 @@ type Manager interface {
 	UpdateRoutes(updateSerial uint64, serverRoutes map[route.ID]*route.Route, clientRoutes route.HAMap, useNewDNSRoute bool) error
 	ClassifyRoutes(newRoutes []*route.Route) (map[route.ID]*route.Route, route.HAMap)
 	TriggerSelection(route.HAMap)
+	SelectRoutes(ids []route.NetID, appendRoute bool) error
+	DeselectRoutes(ids []route.NetID) error
+	SelectAllRoutes()
+	DeselectAllRoutes()
 	GetRouteSelector() *routeselector.RouteSelector
 	GetClientRoutes() route.HAMap
 	GetSelectedClientRoutes() route.HAMap
@@ -61,6 +65,7 @@ type Manager interface {
 	InitialRouteRange() []string
 	SetFirewall(firewall.Manager) error
 	SetDNSForwarderPort(port uint16)
+	ReconcilePeerAllowedIPs(peerKey string) error
 	Stop(stateManager *statemanager.Manager)
 }
 
@@ -232,6 +237,30 @@ func (m *DefaultManager) setupRefCounters(useNoop bool) {
 	)
 }
 
+// ReconcilePeerAllowedIPs re-applies every routed allowed IP currently tracked for the peer
+// onto the WireGuard device. The allowed-IP refcounter only calls its AddFunc (which pushes to
+// the device) on a prefix's 0->1 transition, so a peer whose device entry was rebuilt without a
+// matching refcounter change — e.g. a lazy connection cycling through idle->wake, which recreates
+// the WireGuard peer with the overlay /32 only — ends up missing routed prefixes the refcounter
+// still considers installed, and nothing retries. Calling this when the peer's WireGuard entry is
+// (re)created restores convergence. It is add-only and idempotent: AddAllowedIP is update-only, so
+// prefixes are re-added to an existing peer and an absent peer is left untouched.
+func (m *DefaultManager) ReconcilePeerAllowedIPs(peerKey string) error {
+	if m.allowedIPsRefCounter == nil {
+		return nil
+	}
+
+	return m.allowedIPsRefCounter.ReapplyMatching(
+		func(out string) bool { return out == peerKey },
+		func(prefix netip.Prefix) error {
+			if err := m.wgInterface.AddAllowedIP(peerKey, prefix); err != nil {
+				return fmt.Errorf("add allowed IP %s for peer %s: %w", prefix, peerKey, err)
+			}
+			return nil
+		},
+	)
+}
+
 // Init sets up the routing
 func (m *DefaultManager) Init() error {
 	m.routeSelector = m.initSelector()
@@ -775,7 +804,7 @@ func (m *DefaultManager) collectExitNodeInfo(clientRoutes route.HAMap) exitNodeI
 	var info exitNodeInfo
 
 	for haID, routes := range clientRoutes {
-		if !m.isExitNodeRoute(routes) {
+		if !isExitNodeRoutes(routes) {
 			continue
 		}
 
@@ -795,13 +824,6 @@ func (m *DefaultManager) collectExitNodeInfo(clientRoutes route.HAMap) exitNodeI
 	return info
 }
 
-func (m *DefaultManager) isExitNodeRoute(routes []*route.Route) bool {
-	if len(routes) == 0 {
-		return false
-	}
-	return route.IsV4DefaultRoute(routes[0].Network) || route.IsV6DefaultRoute(routes[0].Network)
-}
-
 func (m *DefaultManager) categorizeUserSelection(netID route.NetID, info *exitNodeInfo) {
 	if m.routeSelector.IsSelected(netID) {
 		info.userSelected = append(info.userSelected, netID)

client/internal/routemanager/mock.go

@@ -16,6 +16,8 @@ type MockManager struct {
 	ClassifyRoutesFunc           func(routes []*route.Route) (map[route.ID]*route.Route, route.HAMap)
 	UpdateRoutesFunc             func(updateSerial uint64, serverRoutes map[route.ID]*route.Route, clientRoutes route.HAMap, useNewDNSRoute bool) error
 	TriggerSelectionFunc         func(haMap route.HAMap)
+	SelectRoutesFunc             func(ids []route.NetID, appendRoute bool) error
+	DeselectRoutesFunc           func(ids []route.NetID) error
 	GetRouteSelectorFunc         func() *routeselector.RouteSelector
 	GetClientRoutesFunc          func() route.HAMap
 	GetSelectedClientRoutesFunc  func() route.HAMap
@@ -55,6 +57,30 @@ func (m *MockManager) TriggerSelection(networks route.HAMap) {
 	}
 }
 
+// SelectRoutes mock implementation of SelectRoutes from Manager interface
+func (m *MockManager) SelectRoutes(ids []route.NetID, appendRoute bool) error {
+	if m.SelectRoutesFunc != nil {
+		return m.SelectRoutesFunc(ids, appendRoute)
+	}
+	return nil
+}
+
+// DeselectRoutes mock implementation of DeselectRoutes from Manager interface
+func (m *MockManager) DeselectRoutes(ids []route.NetID) error {
+	if m.DeselectRoutesFunc != nil {
+		return m.DeselectRoutesFunc(ids)
+	}
+	return nil
+}
+
+// SelectAllRoutes mock implementation of SelectAllRoutes from Manager interface
+func (m *MockManager) SelectAllRoutes() {
+}
+
+// DeselectAllRoutes mock implementation of DeselectAllRoutes from Manager interface
+func (m *MockManager) DeselectAllRoutes() {
+}
+
 // GetRouteSelector mock implementation of GetRouteSelector from Manager interface
 func (m *MockManager) GetRouteSelector() *routeselector.RouteSelector {
 	if m.GetRouteSelectorFunc != nil {
@@ -112,6 +138,11 @@ func (m *MockManager) SetFirewall(firewall.Manager) error {
 func (m *MockManager) SetDNSForwarderPort(port uint16) {
 }
 
+// ReconcilePeerAllowedIPs mock implementation of ReconcilePeerAllowedIPs from Manager interface
+func (m *MockManager) ReconcilePeerAllowedIPs(peerKey string) error {
+	return nil
+}
+
 // Stop mock implementation of Stop from Manager interface
 func (m *MockManager) Stop(stateManager *statemanager.Manager) {
 	if m.StopFunc != nil {

client/internal/routemanager/reconcile_test.go

@@ -0,0 +1,90 @@
+//go:build !windows
+
+package routemanager
+
+import (
+	"net"
+	"net/netip"
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"golang.zx2c4.com/wireguard/tun/netstack"
+
+	"github.com/netbirdio/netbird/client/iface/device"
+	"github.com/netbirdio/netbird/client/iface/wgaddr"
+	"github.com/netbirdio/netbird/client/internal/routemanager/refcounter"
+)
+
+// reconcileWGMock is a minimal iface.WGIface that only records AddAllowedIP calls; every other
+// method is an inert stub because ReconcilePeerAllowedIPs exercises none of them.
+type reconcileWGMock struct {
+	mu   sync.Mutex
+	adds map[string][]netip.Prefix
+}
+
+func (m *reconcileWGMock) AddAllowedIP(peerKey string, allowedIP netip.Prefix) error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	if m.adds == nil {
+		m.adds = map[string][]netip.Prefix{}
+	}
+	m.adds[peerKey] = append(m.adds[peerKey], allowedIP)
+	return nil
+}
+
+func (m *reconcileWGMock) added(peerKey string) []netip.Prefix {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	return m.adds[peerKey]
+}
+
+func (m *reconcileWGMock) RemoveAllowedIP(string, netip.Prefix) error { return nil }
+func (m *reconcileWGMock) Name() string                               { return "utun-test" }
+func (m *reconcileWGMock) Address() wgaddr.Address                    { return wgaddr.Address{} }
+func (m *reconcileWGMock) ToInterface() *net.Interface                { return nil }
+func (m *reconcileWGMock) IsUserspaceBind() bool                      { return false }
+func (m *reconcileWGMock) GetFilter() device.PacketFilter             { return nil }
+func (m *reconcileWGMock) GetDevice() *device.FilteredDevice          { return nil }
+func (m *reconcileWGMock) GetNet() *netstack.Net                      { return nil }
+
+// TestReconcilePeerAllowedIPs verifies the declarative reconcile re-applies every routed prefix
+// tracked for the peer (self-heal, independent of refcount level) and stays scoped to that peer.
+func TestReconcilePeerAllowedIPs(t *testing.T) {
+	wg := &reconcileWGMock{}
+	m := &DefaultManager{wgInterface: wg}
+	m.allowedIPsRefCounter = refcounter.New[netip.Prefix, string, string](
+		func(_ netip.Prefix, peerKey string) (string, error) { return peerKey, nil },
+		func(netip.Prefix, string) error { return nil },
+	)
+
+	peerA1 := netip.MustParsePrefix("10.0.0.0/24")
+	peerA2 := netip.MustParsePrefix("10.1.0.0/24")
+	peerB1 := netip.MustParsePrefix("10.2.0.0/24")
+
+	for prefix, peer := range map[netip.Prefix]string{peerA1: "peerA", peerA2: "peerA", peerB1: "peerB"} {
+		_, err := m.allowedIPsRefCounter.Increment(prefix, peer)
+		require.NoError(t, err)
+	}
+	// Extra reference: reconcile must still re-apply the prefix even though its refcount never
+	// hit 0 again (the exact case the plain incremental path skips).
+	_, err := m.allowedIPsRefCounter.Increment(peerA1, "peerA")
+	require.NoError(t, err)
+
+	require.NoError(t, m.ReconcilePeerAllowedIPs("peerA"))
+
+	assert.ElementsMatch(t, []netip.Prefix{peerA1, peerA2}, wg.added("peerA"),
+		"reconcile must re-apply all routed prefixes of the peer")
+	assert.Empty(t, wg.added("peerB"), "reconcile must not touch another peer's prefixes")
+}
+
+// TestReconcilePeerAllowedIPsNoCounter verifies reconcile is a safe no-op before the refcounter is
+// set up.
+func TestReconcilePeerAllowedIPsNoCounter(t *testing.T) {
+	wg := &reconcileWGMock{}
+	m := &DefaultManager{wgInterface: wg}
+
+	require.NoError(t, m.ReconcilePeerAllowedIPs("peerA"))
+	assert.Empty(t, wg.added("peerA"))
+}

client/internal/routemanager/refcounter/refcounter.go

@@ -94,6 +94,26 @@ func (rm *Counter[Key, I, O]) Get(key Key) (Ref[O], bool) {
 	return ref, ok
 }
 
+// ReapplyMatching calls apply for every key whose stored Out satisfies pred, holding the
+// counter lock for the whole pass. Running apply under the lock keeps it atomic with respect
+// to Increment/Decrement: a prefix dropped to zero is removed from the map (and had its
+// RemoveFunc called) before this pass observes it, so a stale key can never be re-applied.
+// pred and apply are invoked under the lock, so they must not call back into the counter.
+func (rm *Counter[Key, I, O]) ReapplyMatching(pred func(out O) bool, apply func(key Key) error) error {
+	rm.mu.Lock()
+	defer rm.mu.Unlock()
+
+	var merr *multierror.Error
+	for key, ref := range rm.refCountMap {
+		if pred(ref.Out) {
+			if err := apply(key); err != nil {
+				merr = multierror.Append(merr, err)
+			}
+		}
+	}
+	return nberrors.FormatErrorOrNil(merr)
+}
+
 // Increment increments the reference count for the given key.
 // If this is the first reference to the key, the AddFunc is called.
 func (rm *Counter[Key, I, O]) Increment(key Key, in I) (Ref[O], error) {

client/internal/routemanager/refcounter/refcounter_test.go

@@ -0,0 +1,47 @@
+package refcounter
+
+import (
+	"net/netip"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// TestReapplyMatching verifies ReapplyMatching invokes apply for exactly the keys whose stored
+// Out satisfies the predicate (no duplicates for multiply-referenced keys) — the primitive
+// ReconcilePeerAllowedIPs relies on to re-apply a single peer's routed prefixes.
+func TestReapplyMatching(t *testing.T) {
+	rc := New[netip.Prefix, string, string](
+		func(_ netip.Prefix, peerKey string) (string, error) { return peerKey, nil },
+		func(netip.Prefix, string) error { return nil },
+	)
+
+	peerA1 := netip.MustParsePrefix("10.0.0.0/24")
+	peerA2 := netip.MustParsePrefix("10.1.0.0/24")
+	peerB1 := netip.MustParsePrefix("10.2.0.0/24")
+
+	for prefix, peer := range map[netip.Prefix]string{peerA1: "peerA", peerA2: "peerA", peerB1: "peerB"} {
+		_, err := rc.Increment(prefix, peer)
+		require.NoError(t, err)
+	}
+	// a second reference must not make the key applied twice
+	_, err := rc.Increment(peerA1, "peerA")
+	require.NoError(t, err)
+
+	var applied []netip.Prefix
+	err = rc.ReapplyMatching(
+		func(out string) bool { return out == "peerA" },
+		func(key netip.Prefix) error { applied = append(applied, key); return nil },
+	)
+	require.NoError(t, err)
+	assert.ElementsMatch(t, []netip.Prefix{peerA1, peerA2}, applied)
+
+	var none []netip.Prefix
+	err = rc.ReapplyMatching(
+		func(out string) bool { return out == "missing" },
+		func(key netip.Prefix) error { none = append(none, key); return nil },
+	)
+	require.NoError(t, err)
+	assert.Empty(t, none)
+}

client/internal/routemanager/selection.go

@@ -0,0 +1,138 @@
+package routemanager
+
+import (
+	"fmt"
+	"slices"
+
+	"github.com/hashicorp/go-multierror"
+	log "github.com/sirupsen/logrus"
+	"golang.org/x/exp/maps"
+
+	nberrors "github.com/netbirdio/netbird/client/errors"
+	"github.com/netbirdio/netbird/route"
+)
+
+// SelectRoutes selects the routes with the given network IDs and applies the
+// new selection. V4/v6 exit-node pairs are expanded automatically. Exit nodes
+// are mutually exclusive: if the selection activates an exit node, every other
+// available exit node is deselected so two can't be active at once. With
+// appendRoute=false the previous selection is replaced instead of extended.
+func (m *DefaultManager) SelectRoutes(ids []route.NetID, appendRoute bool) error {
+	if err := m.selectRoutes(ids, appendRoute); err != nil {
+		return err
+	}
+	m.TriggerSelection(m.GetClientRoutes())
+	return nil
+}
+
+// DeselectRoutes removes the routes with the given network IDs from the
+// selection and applies the change. V4/v6 exit-node pairs are expanded
+// automatically.
+func (m *DefaultManager) DeselectRoutes(ids []route.NetID) error {
+	if err := m.deselectRoutes(ids); err != nil {
+		return err
+	}
+	m.TriggerSelection(m.GetClientRoutes())
+	return nil
+}
+
+func (m *DefaultManager) deselectRoutes(ids []route.NetID) error {
+	routesMap := m.GetClientRoutesWithNetID()
+	routes := route.ExpandV6ExitPairs(slices.Clone(ids), routesMap)
+
+	log.Debugf("deselecting routes with ids: %v", routes)
+
+	if err := m.routeSelector.DeselectRoutes(routes, maps.Keys(routesMap)); err != nil {
+		return fmt.Errorf("deselect routes: %w", err)
+	}
+
+	return nil
+}
+
+// SelectAllRoutes selects every available route and applies the selection.
+// Exit nodes stay mutually exclusive: at most one remains active.
+func (m *DefaultManager) SelectAllRoutes() {
+	m.selectAllRoutes()
+	m.TriggerSelection(m.GetClientRoutes())
+}
+
+func (m *DefaultManager) selectAllRoutes() {
+	m.routeSelector.SelectAllRoutes()
+
+	// Select-all wipes every explicit selection, so exit nodes fall back to
+	// management's auto-apply flags — which may mark several at once.
+	// Reconcile immediately so at most one exit node stays active instead of
+	// waiting for the next network map to enforce it.
+	m.mux.Lock()
+	defer m.mux.Unlock()
+	m.updateRouteSelectorFromManagement(m.clientRoutes)
+}
+
+// DeselectAllRoutes deselects every route and applies the change.
+func (m *DefaultManager) DeselectAllRoutes() {
+	m.routeSelector.DeselectAllRoutes()
+	m.TriggerSelection(m.GetClientRoutes())
+}
+
+func (m *DefaultManager) selectRoutes(ids []route.NetID, appendRoute bool) error {
+	routesMap := m.GetClientRoutesWithNetID()
+	routes := route.ExpandV6ExitPairs(slices.Clone(ids), routesMap)
+	allIDs := maps.Keys(routesMap)
+
+	log.Debugf("selecting routes with ids: %v", routes)
+
+	// A partial failure (e.g. an unknown ID in the request) still selects the
+	// valid routes, so exclusivity below must run regardless of the error.
+	var merr *multierror.Error
+	if err := m.routeSelector.SelectRoutes(routes, appendRoute, allIDs); err != nil {
+		merr = multierror.Append(merr, fmt.Errorf("select routes: %w", err))
+	}
+
+	// Exit nodes are mutually exclusive: if this selection activates an
+	// exit node, deselect every other available exit node so two can't be
+	// selected at once. Non-exit route selections are left untouched.
+	if requestActivatesExitNode(routes, routesMap) {
+		if others := otherExitNodeIDs(routesMap, routes); len(others) > 0 {
+			if err := m.routeSelector.DeselectRoutes(others, allIDs); err != nil {
+				merr = multierror.Append(merr, fmt.Errorf("deselect sibling exit nodes: %w", err))
+			}
+		}
+	}
+
+	return nberrors.FormatErrorOrNil(merr)
+}
+
+func isExitNodeRoutes(routes []*route.Route) bool {
+	return len(routes) > 0 && (route.IsV4DefaultRoute(routes[0].Network) || route.IsV6DefaultRoute(routes[0].Network))
+}
+
+// requestActivatesExitNode reports whether any requested NetID maps to an exit
+// node (default route) in the current route table.
+func requestActivatesExitNode(requested []route.NetID, routesMap map[route.NetID][]*route.Route) bool {
+	for _, id := range requested {
+		if isExitNodeRoutes(routesMap[id]) {
+			return true
+		}
+	}
+	return false
+}
+
+// otherExitNodeIDs returns every available exit-node NetID that is not in the
+// requested set — the siblings to deselect so a single exit node stays active.
+func otherExitNodeIDs(routesMap map[route.NetID][]*route.Route, requested []route.NetID) []route.NetID {
+	keep := make(map[route.NetID]struct{}, len(requested))
+	for _, id := range requested {
+		keep[id] = struct{}{}
+	}
+	var others []route.NetID
+	for id, routes := range routesMap {
+		if !isExitNodeRoutes(routes) {
+			continue
+		}
+		if _, ok := keep[id]; ok {
+			continue
+		}
+		others = append(others, id)
+	}
+	return others
+}

client/internal/routemanager/selection_test.go

@@ -0,0 +1,129 @@
+package routemanager
+
+import (
+	"net/netip"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/client/internal/routeselector"
+	"github.com/netbirdio/netbird/route"
+)
+
+func v6ExitRoute(netID, peer string) *route.Route {
+	return &route.Route{
+		NetID:   route.NetID(netID),
+		Network: netip.MustParsePrefix("::/0"),
+		Peer:    peer,
+	}
+}
+
+func newSelectionTestManager() *DefaultManager {
+	return &DefaultManager{
+		routeSelector: routeselector.NewRouteSelector(),
+		clientRoutes: route.HAMap{
+			"exitA|0.0.0.0/0":    {exitRoute("exitA", "p1", true)},
+			"exitA-v6|::/0":      {v6ExitRoute("exitA-v6", "p1")},
+			"exitB|0.0.0.0/0":    {exitRoute("exitB", "p2", true)},
+			"lan|192.168.1.0/24": {{NetID: "lan", Network: netip.MustParsePrefix("192.168.1.0/24"), Peer: "p3"}},
+		},
+	}
+}
+
+func TestSelectRoutes_ExitNodeExclusivity(t *testing.T) {
+	m := newSelectionTestManager()
+
+	// Selecting an exit node selects its v6 pair and deselects the sibling.
+	require.NoError(t, m.selectRoutes([]route.NetID{"exitA"}, true))
+	assert.True(t, m.routeSelector.IsSelected("exitA"), "exitA should be selected")
+	assert.True(t, m.routeSelector.IsSelected("exitA-v6"), "the v6 pair follows its v4 base")
+	assert.False(t, m.routeSelector.IsSelected("exitB"), "the sibling exit node must be deselected")
+
+	// Switching to the sibling deselects the previous exit node and its v6 pair.
+	require.NoError(t, m.selectRoutes([]route.NetID{"exitB"}, true))
+	assert.True(t, m.routeSelector.IsSelected("exitB"), "exitB should now be selected")
+	assert.False(t, m.routeSelector.IsSelected("exitA"), "the previous exit node must be deselected")
+	assert.False(t, m.routeSelector.IsSelected("exitA-v6"), "the previous exit node's v6 pair must be deselected")
+	assert.True(t, m.routeSelector.IsSelected("lan"), "non-exit route selection is untouched")
+
+	// Selecting a non-exit route leaves the active exit node alone.
+	require.NoError(t, m.selectRoutes([]route.NetID{"lan"}, true))
+	assert.True(t, m.routeSelector.IsSelected("exitB"), "selecting a non-exit route keeps the exit node")
+
+	// Deselecting the active exit node turns every exit node off.
+	require.NoError(t, m.deselectRoutes([]route.NetID{"exitB"}))
+	assert.False(t, m.routeSelector.IsSelected("exitB"), "exitB should be deselected")
+	assert.False(t, m.routeSelector.IsSelected("exitA"), "exitA stays deselected")
+	assert.True(t, m.routeSelector.IsSelected("lan"), "non-exit route selection is untouched")
+}
+
+func TestSelectRoutes_PartialErrorStillEnforcesExclusivity(t *testing.T) {
+	// The unknown ID must be reported, but the valid exit node in the same
+	// request is still selected — so its sibling must still be deselected.
+	// Both orderings are covered: processing must continue past the invalid
+	// ID wherever it sits in the request.
+	requests := map[string][]route.NetID{
+		"invalid id first": {"missing", "exitB"},
+		"invalid id last":  {"exitB", "missing"},
+	}
+
+	for name, ids := range requests {
+		t.Run(name, func(t *testing.T) {
+			m := newSelectionTestManager()
+
+			require.NoError(t, m.selectRoutes([]route.NetID{"exitA"}, true))
+
+			err := m.selectRoutes(ids, true)
+			assert.Error(t, err, "unknown id must be reported")
+			assert.True(t, m.routeSelector.IsSelected("exitB"), "valid exit node from the request is selected")
+			assert.False(t, m.routeSelector.IsSelected("exitA"), "sibling exit node must be deselected despite the error")
+			assert.False(t, m.routeSelector.IsSelected("exitA-v6"), "sibling's v6 pair must be deselected too")
+		})
+	}
+}
+
+func TestSelectAllRoutes_KeepsSingleExitNode(t *testing.T) {
+	// Both exit nodes are marked for auto-apply by management
+	// (SkipAutoApply=false), the state where select-all could turn on two at
+	// once without the immediate reconciliation.
+	m := &DefaultManager{
+		routeSelector: routeselector.NewRouteSelector(),
+		clientRoutes: route.HAMap{
+			"exitA|0.0.0.0/0":    {exitRoute("exitA", "p1", false)},
+			"exitB|0.0.0.0/0":    {exitRoute("exitB", "p2", false)},
+			"lan|192.168.1.0/24": {{NetID: "lan", Network: netip.MustParsePrefix("192.168.1.0/24"), Peer: "p3"}},
+		},
+	}
+
+	require.NoError(t, m.selectRoutes([]route.NetID{"exitB"}, true))
+
+	m.selectAllRoutes()
+
+	assert.True(t, m.routeSelector.IsSelected("lan"), "non-exit routes are all selected")
+	assert.True(t, m.routeSelector.IsSelected("exitA"), "the deterministic management pick stays active")
+	assert.False(t, m.routeSelector.IsSelected("exitB"), "select-all must not leave a second exit node active")
+}
+
+func TestSelectRoutes_UnknownRoute(t *testing.T) {
+	m := newSelectionTestManager()
+
+	assert.Error(t, m.selectRoutes([]route.NetID{"missing"}, true), "selecting an unavailable route must fail")
+	assert.Error(t, m.deselectRoutes([]route.NetID{"missing"}), "deselecting an unavailable route must fail")
+}
+
+func TestExitNodeSelectionHelpers(t *testing.T) {
+	routesMap := map[route.NetID][]*route.Route{
+		"exitA": {{Network: netip.MustParsePrefix("0.0.0.0/0")}},
+		"exitB": {{Network: netip.MustParsePrefix("::/0")}},
+		"lan":   {{Network: netip.MustParsePrefix("192.168.0.0/16")}},
+	}
+
+	assert.True(t, requestActivatesExitNode([]route.NetID{"exitA"}, routesMap), "v4 default route is an exit node")
+	assert.True(t, requestActivatesExitNode([]route.NetID{"exitB"}, routesMap), "v6 default route is an exit node")
+	assert.False(t, requestActivatesExitNode([]route.NetID{"lan"}, routesMap), "lan route is not an exit node")
+	assert.False(t, requestActivatesExitNode([]route.NetID{"missing"}, routesMap), "unknown id is not an exit node")
+
+	others := otherExitNodeIDs(routesMap, []route.NetID{"exitB"})
+	assert.ElementsMatch(t, []route.NetID{"exitA"}, others, "only the other exit node is a sibling; the lan route is ignored")
+}

client/ios/NetBirdSDK/client.go

@@ -13,7 +13,6 @@ import (
 	"time"
 
 	log "github.com/sirupsen/logrus"
-	"golang.org/x/exp/maps"
 
 	"github.com/netbirdio/netbird/client/internal"
 	"github.com/netbirdio/netbird/client/internal/auth"
@@ -637,23 +636,18 @@ func (c *Client) SelectRoute(id string) error {
 	}
 
 	routeManager := engine.GetRouteManager()
-	routeSelector := routeManager.GetRouteSelector()
 	if id == "All" {
 		log.Debugf("select all routes")
-		routeSelector.SelectAllRoutes()
-	} else {
-		log.Debugf("select route with id: %s", id)
-		routes := toNetIDs([]string{id})
-		routesMap := routeManager.GetClientRoutesWithNetID()
-		routes = route.ExpandV6ExitPairs(routes, routesMap)
-		if err := routeSelector.SelectRoutes(routes, true, maps.Keys(routesMap)); err != nil {
-			log.Debugf("error when selecting routes: %s", err)
-			return fmt.Errorf("select routes: %w", err)
-		}
+		routeManager.SelectAllRoutes()
+		return nil
 	}
-	routeManager.TriggerSelection(routeManager.GetClientRoutes())
-	return nil
 
+	log.Debugf("select route with id: %s", id)
+	if err := routeManager.SelectRoutes(toNetIDs([]string{id}), true); err != nil {
+		log.Debugf("error when selecting routes: %s", err)
+		return err
+	}
+	return nil
 }
 
 func (c *Client) DeselectRoute(id string) error {
@@ -667,21 +661,17 @@ func (c *Client) DeselectRoute(id string) error {
 	}
 
 	routeManager := engine.GetRouteManager()
-	routeSelector := routeManager.GetRouteSelector()
 	if id == "All" {
 		log.Debugf("deselect all routes")
-		routeSelector.DeselectAllRoutes()
-	} else {
-		log.Debugf("deselect route with id: %s", id)
-		routes := toNetIDs([]string{id})
-		routesMap := routeManager.GetClientRoutesWithNetID()
-		routes = route.ExpandV6ExitPairs(routes, routesMap)
-		if err := routeSelector.DeselectRoutes(routes, maps.Keys(routesMap)); err != nil {
-			log.Debugf("error when deselecting routes: %s", err)
-			return fmt.Errorf("deselect routes: %w", err)
-		}
+		routeManager.DeselectAllRoutes()
+		return nil
+	}
+
+	log.Debugf("deselect route with id: %s", id)
+	if err := routeManager.DeselectRoutes(toNetIDs([]string{id})); err != nil {
+		log.Debugf("error when deselecting routes: %s", err)
+		return err
 	}
-	routeManager.TriggerSelection(routeManager.GetClientRoutes())
 	return nil
 }
 

client/ios/NetBirdSDK/version.go

@@ -0,0 +1,12 @@
+//go:build ios
+
+package NetBirdSDK
+
+import "github.com/netbirdio/netbird/version"
+
+// GoClientVersion returns the NetBird Go client version that was baked into
+// the framework at compile time via
+// -ldflags "-X github.com/netbirdio/netbird/version.version=<version>".
+func GoClientVersion() string {
+	return version.NetbirdVersion()
+}

client/server/network.go

@@ -8,7 +8,6 @@ import (
 	"sort"
 	"strings"
 
-	"golang.org/x/exp/maps"
 	"google.golang.org/grpc/codes"
 	gstatus "google.golang.org/grpc/status"
 
@@ -161,30 +160,11 @@ func (s *Server) SelectNetworks(_ context.Context, req *proto.SelectNetworksRequ
 		return nil, fmt.Errorf("no route manager")
 	}
 
-	routeSelector := routeManager.GetRouteSelector()
 	if req.GetAll() {
-		routeSelector.SelectAllRoutes()
-	} else {
-		routes := toNetIDs(req.GetNetworkIDs())
-		routesMap := routeManager.GetClientRoutesWithNetID()
-		routes = route.ExpandV6ExitPairs(routes, routesMap)
-		netIdRoutes := maps.Keys(routesMap)
-		if err := routeSelector.SelectRoutes(routes, req.GetAppend(), netIdRoutes); err != nil {
-			return nil, fmt.Errorf("select routes: %w", err)
-		}
-
-		// Exit nodes are mutually exclusive: if this selection activates an
-		// exit node, deselect every other available exit node so two can't be
-		// selected at once. Non-exit route selections are left untouched.
-		if requestActivatesExitNode(routes, routesMap) {
-			if others := otherExitNodeIDs(routesMap, routes); len(others) > 0 {
-				if err := routeSelector.DeselectRoutes(others, netIdRoutes); err != nil {
-					return nil, fmt.Errorf("deselect sibling exit nodes: %w", err)
-				}
-			}
-		}
+		routeManager.SelectAllRoutes()
+	} else if err := routeManager.SelectRoutes(toNetIDs(req.GetNetworkIDs()), req.GetAppend()); err != nil {
+		return nil, err
 	}
-	routeManager.TriggerSelection(routeManager.GetClientRoutes())
 
 	s.statusRecorder.PublishEvent(
 		proto.SystemEvent_INFO,
@@ -224,19 +204,11 @@ func (s *Server) DeselectNetworks(_ context.Context, req *proto.SelectNetworksRe
 		return nil, fmt.Errorf("no route manager")
 	}
 
-	routeSelector := routeManager.GetRouteSelector()
 	if req.GetAll() {
-		routeSelector.DeselectAllRoutes()
-	} else {
-		routes := toNetIDs(req.GetNetworkIDs())
-		routesMap := routeManager.GetClientRoutesWithNetID()
-		routes = route.ExpandV6ExitPairs(routes, routesMap)
-		netIdRoutes := maps.Keys(routesMap)
-		if err := routeSelector.DeselectRoutes(routes, netIdRoutes); err != nil {
-			return nil, fmt.Errorf("deselect routes: %w", err)
-		}
+		routeManager.DeselectAllRoutes()
+	} else if err := routeManager.DeselectRoutes(toNetIDs(req.GetNetworkIDs())); err != nil {
+		return nil, err
 	}
-	routeManager.TriggerSelection(routeManager.GetClientRoutes())
 
 	s.statusRecorder.PublishEvent(
 		proto.SystemEvent_INFO,
@@ -261,37 +233,3 @@ func toNetIDs(routes []string) []route.NetID {
 	return netIDs
 }
 
-func isExitNodeRoutes(routes []*route.Route) bool {
-	return len(routes) > 0 && (route.IsV4DefaultRoute(routes[0].Network) || route.IsV6DefaultRoute(routes[0].Network))
-}
-
-// requestActivatesExitNode reports whether any requested NetID maps to an exit
-// node (default route) in the current route table.
-func requestActivatesExitNode(requested []route.NetID, routesMap map[route.NetID][]*route.Route) bool {
-	for _, id := range requested {
-		if isExitNodeRoutes(routesMap[id]) {
-			return true
-		}
-	}
-	return false
-}
-
-// otherExitNodeIDs returns every available exit-node NetID that is not in the
-// requested set — the siblings to deselect so a single exit node stays active.
-func otherExitNodeIDs(routesMap map[route.NetID][]*route.Route, requested []route.NetID) []route.NetID {
-	keep := make(map[route.NetID]struct{}, len(requested))
-	for _, id := range requested {
-		keep[id] = struct{}{}
-	}
-	var others []route.NetID
-	for id, routes := range routesMap {
-		if !isExitNodeRoutes(routes) {
-			continue
-		}
-		if _, ok := keep[id]; ok {
-			continue
-		}
-		others = append(others, id)
-	}
-	return others
-}

client/server/network_exitnode_test.go

@@ -1,26 +0,0 @@
-package server
-
-import (
-	"net/netip"
-	"testing"
-
-	"github.com/stretchr/testify/assert"
-
-	"github.com/netbirdio/netbird/route"
-)
-
-func TestExitNodeSelectionHelpers(t *testing.T) {
-	routesMap := map[route.NetID][]*route.Route{
-		"exitA": {{Network: netip.MustParsePrefix("0.0.0.0/0")}},
-		"exitB": {{Network: netip.MustParsePrefix("::/0")}},
-		"lan":   {{Network: netip.MustParsePrefix("192.168.0.0/16")}},
-	}
-
-	assert.True(t, requestActivatesExitNode([]route.NetID{"exitA"}, routesMap), "v4 default route is an exit node")
-	assert.True(t, requestActivatesExitNode([]route.NetID{"exitB"}, routesMap), "v6 default route is an exit node")
-	assert.False(t, requestActivatesExitNode([]route.NetID{"lan"}, routesMap), "lan route is not an exit node")
-	assert.False(t, requestActivatesExitNode([]route.NetID{"missing"}, routesMap), "unknown id is not an exit node")
-
-	others := otherExitNodeIDs(routesMap, []route.NetID{"exitB"})
-	assert.ElementsMatch(t, []route.NetID{"exitA"}, others, "only the other exit node is a sibling; the lan route is ignored")
-}

client/system/info.go

@@ -79,13 +79,15 @@ type Info struct {
 	EnableSSHLocalPortForwarding  bool
 	EnableSSHRemotePortForwarding bool
 	DisableSSHAuth                bool
+
+	SyncMessageVersion *int
 }
 
 func (i *Info) SetFlags(
 	rosenpassEnabled, rosenpassPermissive bool,
 	serverSSHAllowed *bool,
 	disableClientRoutes, disableServerRoutes,
-	disableDNS, disableFirewall, blockLANAccess, blockInbound, disableIPv6 bool,
+	disableDNS, disableFirewall, blockLANAccess, blockInbound, disableIPv6 bool, syncMessageVersion *int,
 	enableSSHRoot, enableSSHSFTP, enableSSHLocalPortForwarding, enableSSHRemotePortForwarding *bool,
 	disableSSHAuth *bool,
 ) {
@@ -103,6 +105,8 @@ func (i *Info) SetFlags(
 	i.BlockInbound = blockInbound
 	i.DisableIPv6 = disableIPv6
 
+	i.SyncMessageVersion = syncMessageVersion
+
 	if enableSSHRoot != nil {
 		i.EnableSSHRoot = *enableSSHRoot
 	}

client/ui/frontend/src/modules/auto-update/UpdateVersionCard.tsx

@@ -2,6 +2,7 @@ import { type ReactNode } from "react";
 import { useTranslation } from "react-i18next";
 import { Browser } from "@wailsio/runtime";
 import { DownloadIcon, NotepadText } from "lucide-react";
+import { Update as UpdateSvc } from "@bindings/services";
 import { Button } from "@/components/buttons/Button";
 import { useClientVersion } from "@/contexts/ClientVersionContext";
 import { cn } from "@/lib/cn";
@@ -14,6 +15,12 @@ function openUrl(url: string) {
     });
 }
 
+function openInstallerDownload() {
+    UpdateSvc.DownloadURL()
+        .then(openUrl)
+        .catch(() => openUrl(GITHUB_RELEASES));
+}
+
 export function UpdateVersionCard() {
     const { t } = useTranslation();
     const { updateVersion, enforced, triggerUpdate } = useClientVersion();
@@ -37,11 +44,7 @@ export function UpdateVersionCard() {
                         {t("update.card.installNow")}
                     </Button>
                 ) : (
-                    <Button
-                        variant={"primary"}
-                        size={"xs"}
-                        onClick={() => openUrl(GITHUB_RELEASES)}
-                    >
+                    <Button variant={"primary"} size={"xs"} onClick={openInstallerDownload}>
                         <DownloadIcon size={14} />
                         {t("update.card.getInstaller")}
                     </Button>

client/ui/services/update.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/netbirdio/netbird/client/proto"
 	"github.com/netbirdio/netbird/client/ui/updater"
+	"github.com/netbirdio/netbird/version"
 )
 
 // UpdateResult mirrors TriggerUpdateResponse.
@@ -33,6 +34,12 @@ func (s *Update) GetState() updater.State {
 	return s.holder.Get()
 }
 
+// DownloadURL returns the platform-appropriate installer download link for
+// manual (non-enforced) updates.
+func (s *Update) DownloadURL() string {
+	return version.DownloadUrl()
+}
+
 // Quit exits the app. Scheduled off the calling goroutine so the JS caller's
 // response returns before the runtime tears down.
 func (s *Update) Quit() {

client/ui/tray.go

@@ -32,9 +32,8 @@ const (
 
 	quitDownTimeout = 5 * time.Second
 
-	urlGitHubRepo     = "https://github.com/netbirdio/netbird"
-	urlGitHubReleases = "https://github.com/netbirdio/netbird/releases/latest"
-	urlDocs           = "https://docs.netbird.io"
+	urlGitHubRepo = "https://github.com/netbirdio/netbird"
+	urlDocs       = "https://docs.netbird.io"
 )
 
 // TrayServices bundles the services the tray menu needs, grouped so NewTray

client/ui/tray_update.go

@@ -13,6 +13,7 @@ import (
 
 	"github.com/netbirdio/netbird/client/ui/services"
 	"github.com/netbirdio/netbird/client/ui/updater"
+	"github.com/netbirdio/netbird/version"
 )
 
 // trayUpdater owns the tray UI that reacts to auto-update. Composed inside Tray.
@@ -76,15 +77,15 @@ func (u *trayUpdater) applyLanguage() {
 	u.refreshMenuItem(state)
 }
 
-// handleClick opens the GitHub releases page when not Enforced, otherwise shows
-// the progress page and asks the daemon to start the installer.
+// handleClick opens the installer download link when not Enforced, otherwise
+// shows the progress page and asks the daemon to start the installer.
 func (u *trayUpdater) handleClick() {
 	u.mu.Lock()
 	state := u.state
 	u.mu.Unlock()
 
 	if !state.Enforced {
-		_ = u.app.Browser.OpenURL(urlGitHubReleases)
+		_ = u.app.Browser.OpenURL(version.DownloadUrl())
 		return
 	}
 

combined/cmd/admin.go

@@ -0,0 +1,151 @@
+package cmd
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/dexidp/dex/storage"
+	log "github.com/sirupsen/logrus"
+	"github.com/spf13/cobra"
+
+	"github.com/netbirdio/netbird/formatter/hook"
+	admincmd "github.com/netbirdio/netbird/management/cmd/admin"
+	tokencmd "github.com/netbirdio/netbird/management/cmd/token"
+	nbconfig "github.com/netbirdio/netbird/management/internals/server/config"
+	"github.com/netbirdio/netbird/management/server/activity"
+	activitystore "github.com/netbirdio/netbird/management/server/activity/store"
+	"github.com/netbirdio/netbird/management/server/store"
+	"github.com/netbirdio/netbird/management/server/types"
+	"github.com/netbirdio/netbird/util"
+)
+
+// newAdminCommands creates the admin command tree with combined-specific resource openers.
+func newAdminCommands() *cobra.Command {
+	return admincmd.NewCommands(admincmd.Openers{
+		Resources: withAdminResources,
+		Store:     withAdminStoreOnly,
+		IDP:       withAdminIDPOnly,
+	})
+}
+
+func newLegacyTokenCommand() *cobra.Command {
+	cmd := tokencmd.NewCommands(tokencmd.StoreOpener(withAdminStoreOnly))
+	cmd.Deprecated = "use 'admin token' instead"
+	return cmd
+}
+
+// withAdminResources loads the combined YAML config, initializes stores, and calls fn.
+func withAdminResources(cmd *cobra.Command, fn func(ctx context.Context, resources admincmd.Resources) error) error {
+	return withAdminConfig(cmd, func(ctx context.Context, cfg *CombinedConfig) error {
+		mgmtConfig, err := adminManagementConfig(cfg)
+		if err != nil {
+			return err
+		}
+
+		managementStore, err := openAdminStore(ctx, cfg)
+		if err != nil {
+			return err
+		}
+		defer admincmd.CloseStore(ctx, managementStore)
+
+		idpStorage, idpStorageFile, err := admincmd.OpenIDPStorage(mgmtConfig)
+		if err != nil {
+			return err
+		}
+		defer admincmd.CloseIDPStorage(idpStorage)
+
+		eventStore, esErr := openAdminEventStore(ctx, cfg, mgmtConfig)
+		if esErr != nil {
+			_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "Warning: audit events will not be recorded: %v\n", esErr)
+		}
+		if eventStore != nil {
+			defer func() {
+				if err := eventStore.Close(ctx); err != nil {
+					log.Debugf("close activity event store: %v", err)
+				}
+			}()
+		}
+
+		return fn(ctx, admincmd.Resources{Store: managementStore, IDPStorage: idpStorage, IDPStorageFile: idpStorageFile, EventStore: eventStore})
+	})
+}
+
+// withAdminStoreOnly opens only the management store for admin subcommands that do not
+// need embedded IdP storage.
+func withAdminStoreOnly(cmd *cobra.Command, fn func(ctx context.Context, s store.Store) error) error {
+	return withAdminConfig(cmd, func(ctx context.Context, cfg *CombinedConfig) error {
+		managementStore, err := openAdminStore(ctx, cfg)
+		if err != nil {
+			return err
+		}
+		defer admincmd.CloseStore(ctx, managementStore)
+
+		return fn(ctx, managementStore)
+	})
+}
+
+func withAdminIDPOnly(cmd *cobra.Command, fn func(ctx context.Context, idpStorage storage.Storage, storageFile string) error) error {
+	return withAdminConfig(cmd, func(ctx context.Context, cfg *CombinedConfig) error {
+		mgmtConfig, err := adminManagementConfig(cfg)
+		if err != nil {
+			return err
+		}
+		idpStorage, idpStorageFile, err := admincmd.OpenIDPStorage(mgmtConfig)
+		if err != nil {
+			return err
+		}
+		defer admincmd.CloseIDPStorage(idpStorage)
+
+		return fn(ctx, idpStorage, idpStorageFile)
+	})
+}
+
+func withAdminConfig(cmd *cobra.Command, fn func(ctx context.Context, cfg *CombinedConfig) error) error {
+	if err := util.InitLog("error", "console"); err != nil {
+		return fmt.Errorf("init log: %w", err)
+	}
+
+	ctx := context.WithValue(cmd.Context(), hook.ExecutionContextKey, hook.SystemSource) //nolint:staticcheck
+
+	cfg, err := LoadConfig(configPath)
+	if err != nil {
+		return fmt.Errorf("load config: %w", err)
+	}
+	cfg.ApplyAdminDefaults()
+	applyServerStoreEnv(cfg.Server.Store)
+
+	return fn(ctx, cfg)
+}
+
+func adminManagementConfig(cfg *CombinedConfig) (*nbconfig.Config, error) {
+	mgmtConfig, err := cfg.ToManagementConfig()
+	if err != nil {
+		return nil, fmt.Errorf("create management config: %w", err)
+	}
+	return mgmtConfig, nil
+}
+
+func openAdminStore(ctx context.Context, cfg *CombinedConfig) (store.Store, error) {
+	managementStore, err := store.NewStore(ctx, types.Engine(cfg.Management.Store.Engine), cfg.Management.DataDir, nil, true)
+	if err != nil {
+		return nil, fmt.Errorf("create store: %w", err)
+	}
+	return managementStore, nil
+}
+
+func openAdminEventStore(ctx context.Context, cfg *CombinedConfig, config *nbconfig.Config) (activity.Store, error) {
+	if config.DataStoreEncryptionKey == "" {
+		return nil, fmt.Errorf("data store encryption key is not configured")
+	}
+	if err := applyActivityStoreEnv(cfg.Server.ActivityStore); err != nil {
+		return nil, fmt.Errorf("configure activity event store: %w", err)
+	}
+	eventStore, err := activitystore.NewSqlStore(ctx, config.Datadir, config.DataStoreEncryptionKey)
+	if err != nil {
+		return nil, fmt.Errorf("open activity event store: %w", err)
+	}
+	if eventStore == nil {
+		return nil, fmt.Errorf("open activity event store: returned nil store")
+	}
+	return eventStore, nil
+}

combined/cmd/admin_config_test.go

@@ -0,0 +1,47 @@
+package cmd
+
+import (
+	"context"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	nbconfig "github.com/netbirdio/netbird/management/internals/server/config"
+)
+
+func TestApplyAdminDefaultsCopiesServerStoreWithoutExposedAddress(t *testing.T) {
+	cfg := DefaultConfig()
+	cfg.Server.ExposedAddress = ""
+	cfg.Server.DataDir = "/srv/netbird"
+	cfg.Server.Store = StoreConfig{
+		Engine: "postgres",
+		DSN:    "postgres://user:pass@example.com/netbird",
+	}
+
+	cfg.ApplyAdminDefaults()
+
+	require.Equal(t, "/srv/netbird", cfg.Management.DataDir)
+	require.Equal(t, "postgres", cfg.Management.Store.Engine)
+	require.Equal(t, cfg.Server.Store.DSN, cfg.Management.Store.DSN)
+}
+
+func TestOpenAdminEventStoreMissingEncryptionKeyReturnsNilInterface(t *testing.T) {
+	eventStore, err := openAdminEventStore(context.Background(), &CombinedConfig{}, &nbconfig.Config{})
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "encryption key")
+	require.Nil(t, eventStore)
+}
+
+func TestApplyServerStoreEnv(t *testing.T) {
+	t.Setenv("NB_STORE_ENGINE_POSTGRES_DSN", "")
+	t.Setenv("NB_STORE_ENGINE_MYSQL_DSN", "")
+	t.Setenv("NB_STORE_ENGINE_SQLITE_FILE", "")
+
+	applyServerStoreEnv(StoreConfig{Engine: "postgres", DSN: "postgres-dsn", File: "store.db"})
+	require.Equal(t, "postgres-dsn", os.Getenv("NB_STORE_ENGINE_POSTGRES_DSN"))
+	require.Equal(t, "store.db", os.Getenv("NB_STORE_ENGINE_SQLITE_FILE"))
+
+	applyServerStoreEnv(StoreConfig{Engine: "mysql", DSN: "mysql-dsn"})
+	require.Equal(t, "mysql-dsn", os.Getenv("NB_STORE_ENGINE_MYSQL_DSN"))
+}

combined/cmd/config.go

@@ -6,8 +6,7 @@ import (
 	"net"
 	"net/netip"
 	"os"
-	"path"
-	"path/filepath"
+	filePath "path/filepath"
 	"strings"
 	"time"
 
@@ -74,6 +73,9 @@ type ServerConfig struct {
 	ActivityStore           StoreConfig        `yaml:"activityStore"`
 	AuthStore               StoreConfig        `yaml:"authStore"`
 	ReverseProxy            ReverseProxyConfig `yaml:"reverseProxy"`
+
+	SupportedSyncMessageVersions           *int           `yaml:"supportedSyncMessageVersions,omitempty"`
+	PerAccountSupportedSyncMessageVersions map[string]int `yaml:"perAccountSupportedSyncMessageVersions,omitempty"`
 }
 
 // TLSConfig contains TLS/HTTPS settings
@@ -300,6 +302,19 @@ func (c *CombinedConfig) ApplySimplifiedDefaults() {
 	c.autoConfigureClientSettings(exposedProto, exposedHost, exposedHostPort, hasExternalStuns, hasExternalRelay, hasExternalSignal)
 }
 
+// ApplyAdminDefaults applies the management settings needed by admin commands even
+// when the full server config is invalid and ApplySimplifiedDefaults cannot run.
+func (c *CombinedConfig) ApplyAdminDefaults() {
+	if c.Management.DataDir == "" || c.Management.DataDir == "/var/lib/netbird/" {
+		c.Management.DataDir = c.Server.DataDir
+	}
+	if c.Management.Store.Engine == "" || c.Management.Store.Engine == "sqlite" {
+		if c.Server.Store.Engine != "" || c.Server.Store.File != "" || c.Server.Store.DSN != "" {
+			c.Management.Store = c.Server.Store
+		}
+	}
+}
+
 // applyRelayDefaults configures the relay service if no external relay is configured.
 func (c *CombinedConfig) applyRelayDefaults(exposedProto, exposedHostPort string, hasExternalRelay, hasExternalStuns bool) {
 	if hasExternalRelay {
@@ -577,11 +592,11 @@ func (c *CombinedConfig) buildEmbeddedIdPConfig(mgmt ManagementConfig) (*idp.Emb
 			return nil, fmt.Errorf("authStore.dsn is required when authStore.engine is postgres")
 		}
 	} else {
-		authStorageFile = path.Join(mgmt.DataDir, "idp.db")
+		authStorageFile = filePath.Join(mgmt.DataDir, "idp.db")
 		if c.Server.AuthStore.File != "" {
 			authStorageFile = c.Server.AuthStore.File
-			if !filepath.IsAbs(authStorageFile) {
-				authStorageFile = filepath.Join(mgmt.DataDir, authStorageFile)
+			if !filePath.IsAbs(authStorageFile) {
+				authStorageFile = filePath.Join(mgmt.DataDir, authStorageFile)
 			}
 		}
 	}
@@ -696,16 +711,18 @@ func (c *CombinedConfig) ToManagementConfig() (*nbconfig.Config, error) {
 	httpConfig.AuthCallbackURL = callbackURL + types.ProxyCallbackEndpointFull
 
 	return &nbconfig.Config{
-		Stuns:                  stuns,
-		Relay:                  relayConfig,
-		Signal:                 signalConfig,
-		Datadir:                mgmt.DataDir,
-		DataStoreEncryptionKey: mgmt.Store.EncryptionKey,
-		HttpConfig:             httpConfig,
-		StoreConfig:            storeConfig,
-		ReverseProxy:           reverseProxy,
-		DisableDefaultPolicy:   mgmt.DisableDefaultPolicy,
-		EmbeddedIdP:            embeddedIdP,
+		Stuns:                              stuns,
+		Relay:                              relayConfig,
+		Signal:                             signalConfig,
+		Datadir:                            mgmt.DataDir,
+		DataStoreEncryptionKey:             mgmt.Store.EncryptionKey,
+		HttpConfig:                         httpConfig,
+		StoreConfig:                        storeConfig,
+		ReverseProxy:                       reverseProxy,
+		DisableDefaultPolicy:               mgmt.DisableDefaultPolicy,
+		EmbeddedIdP:                        embeddedIdP,
+		HighestSupportedSyncMessageVersion: c.Server.SupportedSyncMessageVersions,
+		PerAccountHighestSupportedSyncMessageVersion: c.Server.PerAccountSupportedSyncMessageVersions,
 	}, nil
 }
 
@@ -729,7 +746,7 @@ func ApplyEmbeddedIdPConfig(ctx context.Context, cfg *nbconfig.Config, mgmtPort
 		cfg.EmbeddedIdP.Storage.Type = "sqlite3"
 	}
 	if cfg.EmbeddedIdP.Storage.Config.File == "" && cfg.Datadir != "" {
-		cfg.EmbeddedIdP.Storage.Config.File = path.Join(cfg.Datadir, "idp.db")
+		cfg.EmbeddedIdP.Storage.Config.File = filePath.Join(cfg.Datadir, "idp.db")
 	}
 
 	issuer := cfg.EmbeddedIdP.Issuer

combined/cmd/root.go

@@ -31,6 +31,7 @@ import (
 	relayServer "github.com/netbirdio/netbird/relay/server"
 	"github.com/netbirdio/netbird/relay/server/listener"
 	"github.com/netbirdio/netbird/relay/server/listener/ws"
+	syncgrpc "github.com/netbirdio/netbird/shared/management/grpc"
 	sharedMetrics "github.com/netbirdio/netbird/shared/metrics"
 	"github.com/netbirdio/netbird/shared/relay/auth"
 	"github.com/netbirdio/netbird/shared/signal/proto"
@@ -64,7 +65,8 @@ func init() {
 	rootCmd.PersistentFlags().StringVarP(&configPath, "config", "c", "", "path to YAML configuration file (required)")
 	_ = rootCmd.MarkPersistentFlagRequired("config")
 
-	rootCmd.AddCommand(newTokenCommands())
+	rootCmd.AddCommand(newAdminCommands())
+	rootCmd.AddCommand(newLegacyTokenCommand())
 }
 
 func RootCmd() *cobra.Command {
@@ -122,6 +124,37 @@ func execute(cmd *cobra.Command, _ []string) error {
 }
 
 // initializeConfig loads and validates the configuration, then initializes logging.
+func applyServerStoreEnv(storeConfig StoreConfig) {
+	if dsn := storeConfig.DSN; dsn != "" {
+		switch strings.ToLower(storeConfig.Engine) {
+		case "postgres":
+			os.Setenv("NB_STORE_ENGINE_POSTGRES_DSN", dsn)
+		case "mysql":
+			os.Setenv("NB_STORE_ENGINE_MYSQL_DSN", dsn)
+		}
+	}
+	if file := storeConfig.File; file != "" {
+		os.Setenv("NB_STORE_ENGINE_SQLITE_FILE", file)
+	}
+}
+
+func applyActivityStoreEnv(storeConfig StoreConfig) error {
+	if engine := storeConfig.Engine; engine != "" {
+		engineLower := strings.ToLower(engine)
+		if engineLower == "postgres" && storeConfig.DSN == "" {
+			return fmt.Errorf("activityStore.dsn is required when activityStore.engine is postgres")
+		}
+		os.Setenv("NB_ACTIVITY_EVENT_STORE_ENGINE", engineLower)
+		if dsn := storeConfig.DSN; dsn != "" {
+			os.Setenv("NB_ACTIVITY_EVENT_POSTGRES_DSN", dsn)
+		}
+	}
+	if file := storeConfig.File; file != "" {
+		os.Setenv("NB_ACTIVITY_EVENT_SQLITE_FILE", file)
+	}
+	return nil
+}
+
 func initializeConfig() error {
 	var err error
 	config, err = LoadConfig(configPath)
@@ -137,30 +170,10 @@ func initializeConfig() error {
 		return fmt.Errorf("failed to initialize log: %w", err)
 	}
 
-	if dsn := config.Server.Store.DSN; dsn != "" {
-		switch strings.ToLower(config.Server.Store.Engine) {
-		case "postgres":
-			os.Setenv("NB_STORE_ENGINE_POSTGRES_DSN", dsn)
-		case "mysql":
-			os.Setenv("NB_STORE_ENGINE_MYSQL_DSN", dsn)
-		}
-	}
-	if file := config.Server.Store.File; file != "" {
-		os.Setenv("NB_STORE_ENGINE_SQLITE_FILE", file)
-	}
+	applyServerStoreEnv(config.Server.Store)
 
-	if engine := config.Server.ActivityStore.Engine; engine != "" {
-		engineLower := strings.ToLower(engine)
-		if engineLower == "postgres" && config.Server.ActivityStore.DSN == "" {
-			return fmt.Errorf("activityStore.dsn is required when activityStore.engine is postgres")
-		}
-		os.Setenv("NB_ACTIVITY_EVENT_STORE_ENGINE", engineLower)
-		if dsn := config.Server.ActivityStore.DSN; dsn != "" {
-			os.Setenv("NB_ACTIVITY_EVENT_POSTGRES_DSN", dsn)
-		}
-	}
-	if file := config.Server.ActivityStore.File; file != "" {
-		os.Setenv("NB_ACTIVITY_EVENT_SQLITE_FILE", file)
+	if err := applyActivityStoreEnv(config.Server.ActivityStore); err != nil {
+		return err
 	}
 
 	log.Infof("Starting combined NetBird server")
@@ -505,6 +518,16 @@ func createManagementServer(cfg *CombinedConfig, mgmtConfig *nbconfig.Config) (m
 	}
 	mgmtPort, _ := strconv.Atoi(portStr)
 
+	if err := syncgrpc.ValidateSyncMessageVersion(mgmtConfig.HighestSupportedSyncMessageVersion); err != nil {
+		return nil, err
+	}
+
+	for accountId, version := range mgmtConfig.PerAccountHighestSupportedSyncMessageVersion {
+		if err := syncgrpc.ValidateSyncMessageVersion(&version); err != nil {
+			return nil, fmt.Errorf("unrecognized sync message version in perAccountSupportedSyncMessageVersions for account %s %w", accountId, err)
+		}
+	}
+
 	mgmtSrv := newServer(
 		&mgmtServer.Config{
 			NbConfig:                mgmtConfig,

combined/cmd/token.go

@@ -1,63 +0,0 @@
-package cmd
-
-import (
-	"context"
-	"fmt"
-	"os"
-	"strings"
-
-	log "github.com/sirupsen/logrus"
-	"github.com/spf13/cobra"
-
-	"github.com/netbirdio/netbird/formatter/hook"
-	tokencmd "github.com/netbirdio/netbird/management/cmd/token"
-	"github.com/netbirdio/netbird/management/server/store"
-	"github.com/netbirdio/netbird/management/server/types"
-	"github.com/netbirdio/netbird/util"
-)
-
-// newTokenCommands creates the token command tree with combined-specific store opener.
-func newTokenCommands() *cobra.Command {
-	return tokencmd.NewCommands(withTokenStore)
-}
-
-// withTokenStore loads the combined YAML config, initializes the store, and calls fn.
-func withTokenStore(cmd *cobra.Command, fn func(ctx context.Context, s store.Store) error) error {
-	if err := util.InitLog("error", "console"); err != nil {
-		return fmt.Errorf("init log: %w", err)
-	}
-
-	ctx := context.WithValue(cmd.Context(), hook.ExecutionContextKey, hook.SystemSource) //nolint:staticcheck
-
-	cfg, err := LoadConfig(configPath)
-	if err != nil {
-		return fmt.Errorf("load config: %w", err)
-	}
-
-	if dsn := cfg.Server.Store.DSN; dsn != "" {
-		switch strings.ToLower(cfg.Server.Store.Engine) {
-		case "postgres":
-			os.Setenv("NB_STORE_ENGINE_POSTGRES_DSN", dsn)
-		case "mysql":
-			os.Setenv("NB_STORE_ENGINE_MYSQL_DSN", dsn)
-		}
-	}
-	if file := cfg.Server.Store.File; file != "" {
-		os.Setenv("NB_STORE_ENGINE_SQLITE_FILE", file)
-	}
-
-	datadir := cfg.Management.DataDir
-	engine := types.Engine(cfg.Management.Store.Engine)
-
-	s, err := store.NewStore(ctx, engine, datadir, nil, true)
-	if err != nil {
-		return fmt.Errorf("create store: %w", err)
-	}
-	defer func() {
-		if err := s.Close(ctx); err != nil {
-			log.Debugf("close store: %v", err)
-		}
-	}()
-
-	return fn(ctx, s)
-}

dns/nameserver.go

@@ -53,6 +53,7 @@ type NameServerGroup struct {
 	ID string `gorm:"primaryKey"`
 	// AccountID is a reference to Account that this object belongs
 	AccountID string `gorm:"index"`
+	PublicID  string `json:"-"`
 	// Name group name
 	Name string
 	// Description group description

docs/agent-networks/01-end-to-end-flows.md

@@ -109,7 +109,7 @@ sequenceDiagram
         Chk->>Inj: continue
         Inj->>Inj: inject NetBird identity headers per provider config
         Inj->>Grd: continue
-        Grd->>Grd: enforce model allowlist
+        Grd->>Grd: enforce per-provider allowlist (fail-closed backstop)
         Grd->>Up: forward (over WireGuard)
         Up-->>Resp: response (JSON or SSE stream)
         Resp->>Resp: parse usage tokens, completion
@@ -135,6 +135,21 @@ sequenceDiagram
   (`redact_pii = settings.RedactPii`). Phones, emails, credit cards,
   PII names — see `redact.go` for the full set. See
   [`modules/31-proxy-middleware-builtin.md`](modules/31-proxy-middleware-builtin.md).
+- The model allowlist is enforced in TWO places. `CheckLLMPolicyLimits`
+  is authoritative: it resolves the policy that governs this
+  (provider, caller-groups) and denies (`deny_code = llm_policy.model_blocked`)
+  when no applicable policy permits the model — so an allowlist scoped to
+  one group/provider never leaks to another, and an un-guardrailed policy
+  is genuinely unrestricted. `llm_guardrail` is a per-provider fail-closed
+  backstop: it only carries an allowlist for a provider every authorising
+  policy restricts, and blocks unknown/undetermined models even when
+  management is unreachable. Because that backstop allowlist is the UNION
+  of every restricting policy's models, per-group narrowing lives only in
+  the authoritative check: during a `CheckLLMPolicyLimits` outage
+  `llm_limit_check` fails open, so a caller can reach any model in the
+  provider's union — a group scoped to model A could reach model B if
+  another group restricts the same provider to B. This is the documented
+  fail-open trade-off; a future flag may switch it to fail-closed.
 - SSE streaming requires special handling on the response side; the
   parser must handle partial chunks without buffering the whole
   stream. See [`modules/32-proxy-llm-parsers.md`](modules/32-proxy-llm-parsers.md).

docs/agent-networks/modules/21-management-agentnetwork.md

@@ -122,7 +122,7 @@ At request time the path is independent: the proxy calls `SelectPolicyForRequest
   | on_request | 1 | `llm_router` | `{"providers":[{id, models[], upstream_*, auth_header_*, allowed_group_ids[]}]}` | **true** |
   | on_request | 2 | `llm_limit_check` | `{}` | – |
   | on_request | 3 | `llm_identity_inject` | `{"providers":[{provider_id, header_pair?, json_metadata?, extra_headers?}]}` | **true** |
-  | on_request | 4 | `llm_guardrail` | `{"model_allowlist"?, "prompt_capture":{enabled,redact_pii}}` | – |
+  | on_request | 4 | `llm_guardrail` | `{"provider_allowlists"?: {providerID: []model}, "prompt_capture":{enabled,redact_pii}}` | – |
   | on_response | 5 | `llm_limit_record` | `{}` (runs LAST at runtime) | – |
   | on_response | 6 | `cost_meter` | `{}` | – |
   | on_response | 7 | `llm_response_parser` | `{"capture_completion": <bool>, "redact_pii"?: true}` | – |

docs/agent-networks/modules/31-proxy-middleware-builtin.md

@@ -244,7 +244,7 @@ no mocks. Tests: `TestChain_AllowPath_StampsAttributionAndRecordsCounter`
 | `llm_router` | `{providers: [{id, models, upstream_scheme, upstream_host, upstream_path?, auth_header_name, auth_header_value, allowed_group_ids}]}` |
 | `llm_limit_check` | `{}` — pulls `MgmtClient` from `FactoryContext` |
 | `llm_identity_inject` | `{providers: [{provider_id, header_pair?|json_metadata?, extra_headers?}]}` |
-| `llm_guardrail` | `{model_allowlist: []string, prompt_capture: {enabled, redact_pii}}` |
+| `llm_guardrail` | `{provider_allowlists: {providerID: []string}, prompt_capture: {enabled, redact_pii}}` — allowlist keyed by resolved provider id; a provider absent from the map is unrestricted (fail-closed backstop; authoritative per-policy/group check is management's `CheckLLMPolicyLimits`) |
 | `llm_response_parser` | `{redact_pii?, capture_completion?: *bool}` |
 | `cost_meter` | `{pricing_path?}` (basename inside data-dir; defaults `pricing.yaml`) |
 | `llm_limit_record` | `{}` — same pattern as `llm_limit_check` |

e2e/agentnetwork/chat_test.go

@@ -9,25 +9,234 @@ import (
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"gorm.io/driver/sqlite"
+	"gorm.io/gorm"
 
 	"github.com/netbirdio/netbird/e2e/harness"
 	"github.com/netbirdio/netbird/shared/management/http/api"
 )
 
+// per1k is a model's published USD rates per 1k tokens. read is the prompt-cache read rate
+// (OpenAI: the cached-input discount rate); write is the cache-creation rate where one exists.
+type per1k struct{ in, out, read, write float64 }
+
+// publishedPer1k hardcodes the vendors' PUBLISHED rates for the models the live matrix can drive,
+// keyed by the normalized model id the proxy stamps. Deliberately independent of the proxy's
+// pricing table so a wrong embedded rate or a broken normalization fails the run.
+var publishedPer1k = map[string]per1k{
+	"gpt-4o-mini":                 {0.00015, 0.0006, 0.000075, 0},
+	"gpt-4o":                      {0.0025, 0.01, 0.00125, 0},
+	"claude-haiku-4-5":            {0.001, 0.005, 0.0001, 0.00125},
+	"claude-sonnet-4-5":           {0.003, 0.015, 0.0003, 0.00375},
+	"claude-sonnet-4-6":           {0.003, 0.015, 0.0003, 0.00375},
+	"kimi-k3":                     {0.003, 0.015, 0.0003, 0.003}, // no published write rate: bills at the input rate
+	"anthropic.claude-haiku-4-5":  {0.001, 0.005, 0.0001, 0.00125},
+	"anthropic.claude-sonnet-4-5": {0.003, 0.015, 0.0003, 0.00375},
+	"anthropic.claude-sonnet-4-6": {0.003, 0.015, 0.0003, 0.00375},
+}
+
+// rawCostVerificationSQL is the operator-facing double-check, run straight against the management
+// sqlite store: recompute each usage row's expected total and cache cost from its own persisted
+// token buckets and hardcoded published rates. OpenAI counts cached tokens as a subset of input;
+// Anthropic-shape providers count cache buckets additively.
+const rawCostVerificationSQL = `
+WITH rates(model, in_rate, out_rate, read_rate, write_rate) AS (
+  VALUES
+    ('gpt-4o-mini',                 0.00015, 0.0006, 0.000075, 0.0),
+    ('gpt-4o',                      0.0025,  0.01,   0.00125,  0.0),
+    ('claude-haiku-4-5',            0.001,   0.005,  0.0001,   0.00125),
+    ('claude-sonnet-4-5',           0.003,   0.015,  0.0003,   0.00375),
+    ('claude-sonnet-4-6',           0.003,   0.015,  0.0003,   0.00375),
+    ('kimi-k3',                     0.003,   0.015,  0.0003,   0.003),
+    ('anthropic.claude-haiku-4-5',  0.001,   0.005,  0.0001,   0.00125),
+    ('anthropic.claude-sonnet-4-5', 0.003,   0.015,  0.0003,   0.00375),
+    ('anthropic.claude-sonnet-4-6', 0.003,   0.015,  0.0003,   0.00375)
+)
+SELECT
+  u.provider,
+  u.model,
+  u.input_tokens,
+  u.output_tokens,
+  u.cached_input_tokens,
+  u.cache_creation_tokens,
+  u.input_cost_usd,
+  u.cached_input_cost_usd,
+  u.cache_creation_cost_usd,
+  u.output_cost_usd,
+  -- No cost_usd / cache_cost_usd columns are stored: both are derived from the
+  -- four per-bucket columns above, exactly as the API renders them.
+  (u.input_cost_usd + u.cached_input_cost_usd + u.cache_creation_cost_usd + u.output_cost_usd) AS cost_usd,
+  (u.cached_input_cost_usd + u.cache_creation_cost_usd) AS cache_cost_usd,
+  CASE WHEN u.provider = 'openai' THEN
+    (u.input_tokens - MIN(u.cached_input_tokens, u.input_tokens))*r.in_rate/1000.0
+  ELSE
+    u.input_tokens*r.in_rate/1000.0
+  END AS expected_input,
+  CASE WHEN u.provider = 'openai' THEN
+    MIN(u.cached_input_tokens, u.input_tokens)*r.read_rate/1000.0
+  ELSE
+    u.cached_input_tokens*r.read_rate/1000.0
+  END AS expected_cached_input,
+  CASE WHEN u.provider = 'openai' THEN
+    0.0
+  ELSE
+    u.cache_creation_tokens*r.write_rate/1000.0
+  END AS expected_cache_creation,
+  u.output_tokens*r.out_rate/1000.0 AS expected_output,
+  CASE WHEN u.provider = 'openai' THEN
+    (u.input_tokens - MIN(u.cached_input_tokens, u.input_tokens))*r.in_rate/1000.0
+      + MIN(u.cached_input_tokens, u.input_tokens)*r.read_rate/1000.0
+      + u.output_tokens*r.out_rate/1000.0
+  ELSE
+    u.input_tokens*r.in_rate/1000.0 + u.cached_input_tokens*r.read_rate/1000.0
+      + u.cache_creation_tokens*r.write_rate/1000.0 + u.output_tokens*r.out_rate/1000.0
+  END AS expected_total,
+  CASE WHEN u.provider = 'openai' THEN
+    MIN(u.cached_input_tokens, u.input_tokens)*r.read_rate/1000.0
+  ELSE
+    u.cached_input_tokens*r.read_rate/1000.0 + u.cache_creation_tokens*r.write_rate/1000.0
+  END AS expected_cache
+FROM agent_network_request_usage u
+JOIN rates r ON r.model = u.model
+ORDER BY u.timestamp`
+
+// verifyUsageRowsSQL re-checks every persisted usage row directly in the management sqlite store,
+// bypassing the API path — the same audit an operator can run on a production store.db.
+func verifyUsageRowsSQL(t *testing.T, srv *harness.Combined) {
+	t.Helper()
+
+	dbPath, err := srv.SnapshotStoreDB(t.TempDir())
+	require.NoError(t, err, "snapshot management sqlite store")
+
+	db, err := gorm.Open(sqlite.Open(dbPath), &gorm.Config{})
+	require.NoError(t, err, "open store snapshot")
+	sqlDB, err := db.DB()
+	require.NoError(t, err)
+	defer func() { _ = sqlDB.Close() }()
+
+	rows, err := db.Raw(rawCostVerificationSQL).Rows()
+	require.NoError(t, err, "run raw cost verification query")
+	defer func() { _ = rows.Close() }()
+
+	verified := 0
+	for rows.Next() {
+		var provider, model string
+		var inTok, outTok, readTok, writeTok int64
+		var inCost, cachedInCost, cacheCreateCost, outCost, cost, cacheCost float64
+		var wantInput, wantCachedInput, wantCacheCreation, wantOutput, wantTotal, wantCache float64
+		require.NoError(t, rows.Scan(&provider, &model, &inTok, &outTok, &readTok, &writeTok,
+			&inCost, &cachedInCost, &cacheCreateCost, &outCost, &cost, &cacheCost,
+			&wantInput, &wantCachedInput, &wantCacheCreation, &wantOutput, &wantTotal, &wantCache), "scan usage row")
+		t.Logf("[sql] %s/%s: in=%d out=%d cache_read=%d cache_write=%d stored in/cached/create/out=$%.6f/$%.6f/$%.6f/$%.6f total=$%.6f cache=$%.6f expected total=$%.6f cache=$%.6f",
+			provider, model, inTok, outTok, readTok, writeTok,
+			inCost, cachedInCost, cacheCreateCost, outCost, cost, cacheCost, wantTotal, wantCache)
+		assert.InDeltaf(t, wantInput, inCost, 1e-6, "stored input_cost_usd for %s/%s must match the published-rate recompute", provider, model)
+		assert.InDeltaf(t, wantCachedInput, cachedInCost, 1e-6, "stored cached_input_cost_usd for %s/%s must match the published-rate recompute", provider, model)
+		assert.InDeltaf(t, wantCacheCreation, cacheCreateCost, 1e-6, "stored cache_creation_cost_usd for %s/%s must match the published-rate recompute", provider, model)
+		assert.InDeltaf(t, wantOutput, outCost, 1e-6, "stored output_cost_usd for %s/%s must match the published-rate recompute", provider, model)
+		assert.InDeltaf(t, wantTotal, cost, 1e-6, "derived cost_usd for %s/%s must match the published-rate recompute", provider, model)
+		assert.InDeltaf(t, wantCache, cacheCost, 1e-6, "derived cache_cost_usd for %s/%s must match the published-rate recompute", provider, model)
+		assert.InDeltaf(t, inCost+cachedInCost+cacheCreateCost+outCost, cost, 1e-9,
+			"stored buckets must sum to the derived cost_usd for %s/%s", provider, model)
+		verified++
+	}
+	require.NoError(t, rows.Err(), "iterate usage rows")
+	require.Positive(t, verified, "raw SQL check must cover at least one usage row")
+	t.Logf("[sql] verified %d usage rows in store.db against published rates", verified)
+
+	gwRows, err := db.Raw(`SELECT model,
+	  (input_cost_usd + cached_input_cost_usd + cache_creation_cost_usd + output_cost_usd) AS cost_usd
+	  FROM agent_network_request_usage WHERE model LIKE '%/%'`).Rows()
+	require.NoError(t, err, "query gateway-prefixed usage rows")
+	defer func() { _ = gwRows.Close() }()
+	for gwRows.Next() {
+		var model string
+		var cost float64
+		require.NoError(t, gwRows.Scan(&model, &cost), "scan gateway usage row")
+		t.Logf("[sql] gateway %s: stored=$%.6f (must be 0 — deliberately unpriced)", model, cost)
+		assert.Zerof(t, cost, "gateway-prefixed model %q must store cost 0, never a guessed rate", model)
+	}
+	require.NoError(t, gwRows.Err(), "iterate gateway usage rows")
+}
+
+// validateAccessLogCost recomputes a live access-log row's expected total and cache cost from the
+// published per-1k rates and the row's persisted token buckets, and asserts both stored values.
+// Gateway-prefixed model ids the proxy deliberately does not price must store cost 0.
+func validateAccessLogCost(t *testing.T, pc providerCase, row api.AgentNetworkAccessLog) {
+	t.Helper()
+	model := catalogModel(pc)
+	provider := ""
+	if row.Provider != nil {
+		provider = *row.Provider
+	}
+	t.Logf("[cost] %s: provider=%s model=%s in=%d out=%d total=%d cache_read=%d cache_write=%d cost=$%.6f cache_cost=$%.6f",
+		pc.name, provider, model, row.InputTokens, row.OutputTokens, row.TotalTokens,
+		row.CachedInputTokens, row.CacheCreationTokens, row.CostUsd, row.CacheCostUsd)
+
+	rates, known := publishedPer1k[model]
+	if !known {
+		if strings.Contains(model, "/") {
+			assert.Zerof(t, row.CostUsd, "gateway-prefixed model %q is not priced so the cost meter must skip (cost 0)", model)
+			return
+		}
+		t.Logf("[cost] %s: no published rate on file for model %q (env-overridden?); skipping cost validation", pc.name, model)
+		return
+	}
+
+	// input_tokens may legitimately be 0: Moonshot/Kimi reports fully cached prompts under the cache
+	// buckets only. Output and total must always be present on a priced row.
+	require.Positive(t, row.OutputTokens, "priced row must carry output tokens")
+	require.Positive(t, row.TotalTokens, "priced row must carry total tokens")
+
+	var wantInput, wantCachedInput, wantCacheCreation float64
+	if provider == "openai" {
+		cached := min(row.CachedInputTokens, row.InputTokens) // cached is a subset of input
+		wantInput = float64(row.InputTokens-cached) / 1000 * rates.in
+		wantCachedInput = float64(cached) / 1000 * rates.read
+		// OpenAI has no cache-write bucket; wantCacheCreation stays 0.
+	} else {
+		// Anthropic / Bedrock shape: cache buckets are additive to input_tokens.
+		wantInput = float64(row.InputTokens) / 1000 * rates.in
+		wantCachedInput = float64(row.CachedInputTokens) / 1000 * rates.read
+		wantCacheCreation = float64(row.CacheCreationTokens) / 1000 * rates.write
+	}
+	wantOutput := float64(row.OutputTokens) / 1000 * rates.out
+	wantCache := wantCachedInput + wantCacheCreation
+	wantTotal := wantInput + wantCache + wantOutput
+
+	t.Logf("[cost] %s: expecting input=$%.6f cached_input=$%.6f cache_creation=$%.6f output=$%.6f total=$%.6f cache=$%.6f from published rates",
+		pc.name, wantInput, wantCachedInput, wantCacheCreation, wantOutput, wantTotal, wantCache)
+	assert.InDeltaf(t, wantInput, row.InputCostUsd, 1e-6, "stored input_cost_usd for %s (%s)", pc.name, model)
+	assert.InDeltaf(t, wantCachedInput, row.CachedInputCostUsd, 1e-6, "stored cached_input_cost_usd for %s (%s)", pc.name, model)
+	assert.InDeltaf(t, wantCacheCreation, row.CacheCreationCostUsd, 1e-6, "stored cache_creation_cost_usd for %s (%s)", pc.name, model)
+	assert.InDeltaf(t, wantOutput, row.OutputCostUsd, 1e-6, "stored output_cost_usd for %s (%s)", pc.name, model)
+	assert.InDeltaf(t, wantTotal, row.CostUsd, 1e-6, "derived cost_usd for %s (%s)", pc.name, model)
+	assert.InDeltaf(t, wantCache, row.CacheCostUsd, 1e-6, "derived cache_cost_usd for %s (%s)", pc.name, model)
+
+	// The aggregates must be exactly the sum of the stored components, not an
+	// independently-computed figure that could drift from the breakdown.
+	assert.InDeltaf(t, row.InputCostUsd+row.CachedInputCostUsd+row.CacheCreationCostUsd+row.OutputCostUsd,
+		row.CostUsd, 1e-9, "stored buckets must sum to the derived cost_usd for %s (%s)", pc.name, model)
+	assert.InDeltaf(t, row.CachedInputCostUsd+row.CacheCreationCostUsd,
+		row.CacheCostUsd, 1e-9, "stored cache buckets must sum to the derived cache_cost_usd for %s (%s)", pc.name, model)
+}
+
 // providerCase is one entry in the live provider matrix. The same scenario runs
 // for every available provider; availability is keyed off env vars so the suite
 // covers whatever credentials are present (source ~/.llm-keys locally / set the
 // Actions secrets in CI).
 type providerCase struct {
-	name      string
-	catalogID string
-	upstream  string
-	apiKey    string
-	model     string // body model (chat/messages) or path model@version (vertex)
-	kind      string // harness.WireChat, harness.WireMessages, or harness.WireVertex
-	project   string // vertex only: GCP project for the rawPredict path
-	region    string // vertex only: GCP region for the rawPredict path
+	name       string
+	catalogID  string
+	upstream   string
+	apiKey     string
+	model      string // body model (chat/messages) or path model@version (vertex)
+	kind       string // harness.WireChat, harness.WireMessages, or harness.WireVertex
+	project    string // vertex only: GCP project for the rawPredict path
+	region     string // vertex only: GCP region for the rawPredict path
+	pathPrefix string // base-URL path prefix the agent carries (e.g. "/anthropic" for Kimi)
 }
 
 // availableProviders builds the matrix from the provider env vars that are set.
@@ -39,6 +248,24 @@ func availableProviders() []providerCase {
 	if k := os.Getenv("ANTHROPIC_TOKEN"); k != "" {
 		ps = append(ps, providerCase{name: "anthropic", catalogID: "anthropic_api", upstream: "https://api.anthropic.com", apiKey: k, model: "claude-haiku-4-5", kind: harness.WireMessages})
 	}
+	if k := os.Getenv("KIMI_TOKEN"); k != "" {
+		// Kimi (Moonshot AI) serves two body shapes from the same key: OpenAI
+		// Chat Completions on the bare host (/v1/...) and the Anthropic
+		// Messages API under the /anthropic path prefix (the endpoint
+		// Moonshot's Claude Code guide uses). The provider keeps the bare
+		// default upstream and the AGENT carries the /anthropic prefix in
+		// its base URL — exactly the documented Claude Code / Kimi CLI
+		// setup (ANTHROPIC_BASE_URL=https://<endpoint>/anthropic) — so one
+		// provider serves both shapes and the prefix rides through to
+		// Moonshot. Run the Anthropic shape, the flagship Claude Code path;
+		// the OpenAI wire shape is covered live by the other chat-shaped
+		// matrix providers, and Kimi-over-chat passed with kimi-k3 before
+		// the single-model constraint surfaced (run #73 on the kimi feature
+		// branch). The platform serves this account exactly ONE model —
+		// kimi-k3 (kimi-k2-thinking and even kimi-latest return
+		// resource_not_found_error on both surfaces).
+		ps = append(ps, providerCase{name: "kimi", catalogID: "kimi_api", upstream: "https://api.moonshot.ai", apiKey: k, model: "kimi-k3", kind: harness.WireMessages, pathPrefix: "/anthropic"})
+	}
 	if k, u := os.Getenv("VERCEL_TOKEN"), os.Getenv("VERCEL_URL"); k != "" && u != "" {
 		ps = append(ps, providerCase{name: "vercel", catalogID: "vercel_ai_gateway", upstream: u, apiKey: k, model: "openai/gpt-4o-mini", kind: harness.WireChat})
 	}
@@ -84,14 +311,27 @@ func availableProviders() []providerCase {
 		}
 	}
 
-	// Bedrock: path-routed, bearer auth. Model is a cross-region inference
-	// profile id (distinct string from the first-party Anthropic case).
+	// Bedrock: path-routed, bearer auth. Model is the FULL cross-region
+	// inference-profile id exactly as AWS issues it — region-family prefix
+	// plus the date/version suffix. A bare or wrong-region id makes Bedrock
+	// reject the request with "The provided model identifier is invalid"
+	// before any inference runs. The proxy normalizes this id to the catalog
+	// key (anthropic.claude-haiku-4-5) for routing/pricing/allowlists.
+	// Defaults pair eu-central-1 with the eu.* profile; AWS_REGION overrides
+	// the region and the prefix follows its family.
 	if k := os.Getenv("AWS_BEARER_TOKEN_BEDROCK"); k != "" {
 		region := os.Getenv("AWS_REGION")
 		if region == "" {
-			region = "us-east-1"
+			region = "eu-central-1"
 		}
-		ps = append(ps, providerCase{name: "bedrock", catalogID: "bedrock_api", upstream: "https://bedrock-runtime." + region + ".amazonaws.com", apiKey: k, model: "us.anthropic.claude-haiku-4-5", kind: harness.WireBedrock})
+		// A valid Bedrock inference-profile id, overridable per account (AWS_BEDROCK_MODEL, also the
+		// workflow's bedrock_model dispatch input). `global.` profiles work from any region. Defaults to
+		// Sonnet 4.6, whose id convention dropped the -YYYYMMDD-v1:0 suffix that Haiku 4.5 still carries.
+		model := os.Getenv("AWS_BEDROCK_MODEL")
+		if model == "" {
+			model = "global.anthropic.claude-sonnet-4-6"
+		}
+		ps = append(ps, providerCase{name: "bedrock", catalogID: "bedrock_api", upstream: "https://bedrock-runtime." + region + ".amazonaws.com", apiKey: k, model: model, kind: harness.WireBedrock})
 	}
 	return ps
 }
@@ -108,8 +348,16 @@ func providerRequest(pc providerCase) api.AgentNetworkProviderRequest {
 		Enabled:     ptr(true),
 	}
 	if pc.kind != harness.WireVertex {
+		// The router matches the normalized catalog id. Bedrock's request model
+		// travels as a region-prefixed inference-profile id in the URL path
+		// (us.anthropic...), which the router strips before matching, so register
+		// the normalized form here or routing fails as model_not_routable.
+		modelID := pc.model
+		if pc.kind == harness.WireBedrock {
+			modelID = catalogModel(pc)
+		}
 		req.Models = &[]api.AgentNetworkProviderModel{
-			{Id: pc.model, InputPer1k: 0.001, OutputPer1k: 0.002},
+			{Id: modelID, InputPer1k: 0.001, OutputPer1k: 0.002},
 		}
 	}
 	return req
@@ -201,11 +449,12 @@ func TestProvidersMatrix(t *testing.T) {
 	t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
 
 	require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
+	// Probe first: the GET resolves the endpoint (DNS error fails) and its first packet wakes the lazy proxy peer, so WaitProxyPeer sees it connected; any HTTP status counts.
+	proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
+	require.NoError(t, err, "resolve agent-network endpoint to proxy IP")
 	if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
 		t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
 	}
-	proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
-	require.NoError(t, err, "resolve agent-network endpoint to proxy IP")
 
 	for _, pc := range matrix {
 		pc := pc
@@ -216,6 +465,10 @@ func TestProvidersMatrix(t *testing.T) {
 			// session id and confirm the marker propagated end-to-end.
 			sessionID := "e2e-session-" + pc.name
 
+			// A long-form prompt so completions carry realistic token counts for cost validation;
+			// max_tokens in the harness bodies (2048) lets the full answer through.
+			const matrixPrompt = "explain GitHub workflow in 1000 words"
+
 			// Retry briefly to absorb tunnel/DNS jitter on the first call.
 			var code int
 			var body string
@@ -226,11 +479,11 @@ func TestProvidersMatrix(t *testing.T) {
 				var cerr error
 				switch pc.kind {
 				case harness.WireVertex:
-					c, b, cerr = cl.Vertex(ctx, settings.Endpoint, proxyIP, pc.project, pc.region, pc.model, "Reply with exactly: pong", sessionID)
+					c, b, cerr = cl.Vertex(ctx, settings.Endpoint, proxyIP, pc.project, pc.region, pc.model, matrixPrompt, sessionID)
 				case harness.WireBedrock:
-					c, b, cerr = cl.Bedrock(ctx, settings.Endpoint, proxyIP, pc.model, "Reply with exactly: pong", sessionID)
+					c, b, cerr = cl.Bedrock(ctx, settings.Endpoint, proxyIP, pc.model, matrixPrompt, sessionID)
 				default:
-					c, b, cerr = cl.Chat(ctx, settings.Endpoint, proxyIP, pc.kind, pc.model, "Reply with exactly: pong", sessionID)
+					c, b, cerr = cl.ChatPrefixed(ctx, settings.Endpoint, proxyIP, pc.pathPrefix, pc.kind, pc.model, matrixPrompt, sessionID)
 				}
 				if cerr == nil {
 					code, body = c, b
@@ -249,6 +502,7 @@ func TestProvidersMatrix(t *testing.T) {
 
 			// The session id sent as x-session-id must round-trip into the
 			// access-log row for this provider.
+			var row api.AgentNetworkAccessLog
 			require.Eventually(t, func() bool {
 				logs, lerr := srv.ListAccessLogs(ctx)
 				if lerr != nil {
@@ -256,11 +510,15 @@ func TestProvidersMatrix(t *testing.T) {
 				}
 				for _, r := range logs.Data {
 					if r.SessionId != nil && *r.SessionId == sessionID {
+						row = r
 						return true
 					}
 				}
 				return false
 			}, 30*time.Second, 2*time.Second, "session id %q must be recorded in an access-log row for %s", sessionID, pc.name)
+
+			// Stored total and cache cost must match the published rates applied to the row's buckets.
+			validateAccessLogCost(t, pc, row)
 		})
 	}
 
@@ -281,4 +539,7 @@ func TestProvidersMatrix(t *testing.T) {
 		}
 		return false
 	}, 60*time.Second, 3*time.Second, "consumption must be recorded with positive token counts after live traffic")
+
+	// Final raw-SQL audit: bypass the API and re-verify every persisted usage row in the store.
+	verifyUsageRowsSQL(t, srv)
 }

e2e/agentnetwork/guardrail_block_test.go

@@ -0,0 +1,209 @@
+//go:build e2e
+
+package agentnetwork
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/e2e/harness"
+	"github.com/netbirdio/netbird/shared/management/http/api"
+)
+
+// pathRoutedGuardrailCase is one provider's self-contained scenario: its own
+// provider, its own guardrail whose allowlist holds ONLY that provider's
+// allowed model, and its own policy. Each case runs in isolation (its own
+// proxy + client), so the guardrail the proxy enforces contains exactly this
+// provider's model — never a mixed cross-provider list.
+type pathRoutedGuardrailCase struct {
+	name       string
+	catalogID  string // agent-network catalog provider id
+	wire       string // harness.WireVertex | harness.WireBedrock
+	allowEntry string // the single model id put on the guardrail allowlist
+	allowModel string // model id sent that MUST be served (200)
+	blockModel string // model id sent that MUST be denied (403 model_blocked)
+}
+
+// TestGuardrailBlocksUnselectedModel_PathRouted is the end-to-end regression
+// guard for the customer report that a model-allowlist guardrail attached to a
+// policy has no effect for PATH-ROUTED providers — where the model travels in
+// the URL, not the JSON body: Google Vertex (…/models/{model}:rawPredict) and
+// AWS Bedrock (/model/{id}/invoke).
+//
+// Each provider is tested in isolation with a guardrail allowlisting a single
+// model of its own: the allowed model (in the URL path) is served (200) and an
+// unselected model (in the URL path) is denied 403 by the guardrail
+// (llm_policy.model_blocked) before the upstream. The Vertex case mirrors the
+// customer verbatim — allow Sonnet, and the unselected model is the exact
+// claude-opus-4-6 they reported reaching the model unblocked. The Bedrock case
+// sends a region-prefixed, versioned inference-profile id so URL-path model
+// normalization is exercised too.
+//
+// The provider is catch-all (no models), so the router forwards any model and a
+// 403 can only come from the guardrail, never model_not_routable. Only the
+// upstream LLM is mocked (the vLLM nginx answers any path with 200); management
+// synth/reconcile, the proxy middleware chain (URL-path model extraction,
+// router, guardrail) and the tunnel are all real, and the guardrail denies
+// before the upstream is dialed so the mock cannot influence the block. A
+// static bearer api key is used so the router injects a static Authorization
+// header instead of minting a GCP token — the only reason path-routed providers
+// normally need live credentials — so the test runs with none and is always on.
+func TestGuardrailBlocksUnselectedModel_PathRouted(t *testing.T) {
+	cases := []pathRoutedGuardrailCase{
+		{
+			name:       "vertex",
+			catalogID:  "vertex_ai_api",
+			wire:       harness.WireVertex,
+			allowEntry: "claude-sonnet-4-5",
+			allowModel: "claude-sonnet-4-5",
+			blockModel: "claude-opus-4-6", // the customer-reported model
+		},
+		{
+			name:       "bedrock",
+			catalogID:  "bedrock_api",
+			wire:       harness.WireBedrock,
+			allowEntry: "anthropic.claude-sonnet-4-5", // normalized catalog id
+			allowModel: "us.anthropic.claude-sonnet-4-5-v1:0",
+			blockModel: "us.anthropic.claude-opus-4-8-v1:0",
+		},
+	}
+
+	for _, tc := range cases {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			runPathRoutedGuardrailCase(t, tc)
+		})
+	}
+}
+
+func runPathRoutedGuardrailCase(t *testing.T, tc pathRoutedGuardrailCase) {
+	t.Helper()
+
+	const (
+		vertexProject = "e2e-project"
+		vertexRegion  = "global"
+	)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
+	defer cancel()
+
+	vllm, err := harness.StartVLLM(ctx, srv)
+	require.NoError(t, err, "start mock upstream")
+	t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
+
+	grp, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-guardrail-" + tc.name})
+	require.NoError(t, err, "create group")
+	t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grp.Id) })
+
+	ephemeral := false
+	sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
+		Name:       "e2e-guardrail-" + tc.name + "-client",
+		Type:       "reusable",
+		ExpiresIn:  86400,
+		UsageLimit: 0,
+		AutoGroups: []string{grp.Id},
+		Ephemeral:  &ephemeral,
+	})
+	require.NoError(t, err, "mint setup key")
+	require.NotEmpty(t, sk.Key, "setup key plaintext")
+
+	// Catch-all provider (no models) so the router forwards any model; a static
+	// bearer key means the router injects a static auth header instead of minting
+	// a GCP token. Bootstraps the cluster if it isn't already.
+	staticKey := "static-e2e-token"
+	prov, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
+		Name:             tc.name,
+		ProviderId:       tc.catalogID,
+		UpstreamUrl:      vllm.URL,
+		ApiKey:           &staticKey,
+		Enabled:          ptr(true),
+		BootstrapCluster: ptr(harness.AgentNetworkCluster),
+	})
+	require.NoError(t, err, "create %s provider", tc.name)
+	t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) })
+
+	// Guardrail allowlisting ONLY this provider's allowed model.
+	var gr api.AgentNetworkGuardrailRequest
+	gr.Name = "e2e-guardrail-" + tc.name
+	gr.Checks.ModelAllowlist.Enabled = true
+	gr.Checks.ModelAllowlist.Models = []string{tc.allowEntry}
+	guard, err := srv.CreateGuardrail(ctx, gr)
+	require.NoError(t, err, "create guardrail")
+	t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), guard.Id) })
+
+	enabled := true
+	pol, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
+		Name:                   "e2e-guardrail-" + tc.name,
+		Enabled:                &enabled,
+		SourceGroups:           []string{grp.Id},
+		DestinationProviderIds: []string{prov.Id},
+		GuardrailIds:           &[]string{guard.Id},
+	})
+	require.NoError(t, err, "create policy")
+	t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), pol.Id) })
+
+	settings, err := srv.GetSettings(ctx)
+	require.NoError(t, err, "read settings")
+	require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned")
+
+	proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-guardrail-"+tc.name+"-proxy")
+	require.NoError(t, err, "mint proxy token")
+	px, err := harness.StartProxy(ctx, srv, proxyToken)
+	require.NoError(t, err, "start proxy")
+	t.Cleanup(func() { _ = px.Terminate(context.Background()) })
+
+	cl, err := harness.StartClient(ctx, srv, sk.Key)
+	require.NoError(t, err, "start client")
+	t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
+
+	require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
+	// Probe first: the GET resolves the endpoint and its first packet wakes the
+	// lazy proxy peer, so WaitProxyPeer then observes it connected.
+	proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
+	require.NoError(t, err, "resolve endpoint to proxy IP")
+	if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
+		t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
+	}
+
+	send := func(model string) (int, string) {
+		var code int
+		var body string
+		var cerr error
+		switch tc.wire {
+		case harness.WireVertex:
+			code, body, cerr = cl.Vertex(ctx, settings.Endpoint, proxyIP, vertexProject, vertexRegion, model, "Reply with exactly: pong", "")
+		case harness.WireBedrock:
+			code, body, cerr = cl.Bedrock(ctx, settings.Endpoint, proxyIP, model, "Reply with exactly: pong", "")
+		default:
+			t.Fatalf("unsupported wire %q", tc.wire)
+		}
+		require.NoError(t, cerr, "request must reach the proxy for %s", tc.name)
+		return code, body
+	}
+
+	// Allowed model (in the URL path) is served. Retry to absorb tunnel/DNS
+	// jitter on the first call over the freshly warmed tunnel.
+	var code int
+	var body string
+	deadline := time.Now().Add(90 * time.Second)
+	for time.Now().Before(deadline) {
+		code, body = send(tc.allowModel)
+		if code == 200 {
+			break
+		}
+		time.Sleep(5 * time.Second)
+	}
+	assert.Equal(t, 200, code,
+		"allowed %s model (URL path) must be served; body: %s\n=== proxy logs ===\n%s", tc.name, body, px.Logs(context.Background()))
+
+	// Unselected model (in the URL path) must be blocked by the guardrail.
+	code, body = send(tc.blockModel)
+	assert.Equal(t, 403, code,
+		"unselected %s model (URL path) must be denied, not served; body: %s\n=== proxy logs ===\n%s", tc.name, body, px.Logs(context.Background()))
+	assert.Contains(t, body, "llm_policy.model_blocked",
+		"%s denial must come from the guardrail allowlist, not routing; body: %s", tc.name, body)
+}

e2e/agentnetwork/guardrail_groupswitch_test.go

@@ -0,0 +1,205 @@
+//go:build e2e
+
+package agentnetwork
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/e2e/harness"
+	"github.com/netbirdio/netbird/shared/management/http/api"
+)
+
+// TestGuardrailGroupSwitchTakesEffectAfterTTL proves that moving a peer between
+// groups flips its model-allowlist decision once the proxy's tunnel-peer cache
+// expires. The peer's groups reach the guardrail via ValidateTunnelPeer, which
+// the proxy caches; the switch is invisible until that cache expires. The proxy
+// runs with a short NB_PROXY_TUNNEL_CACHE_TTL so the flip happens in seconds
+// instead of the 5-minute default.
+//
+// Setup: one catch-all provider declaring modelA + modelB; polA (grpA -> allow
+// modelA) and polB (grpB -> allow modelB). The client starts in grpA. modelA is
+// served and modelB denied; after switching the client grpA -> grpB, modelB is
+// served and modelA denied. The cross-group deny comes from management's
+// per-policy/group CheckLLMPolicyLimits (the proxy backstop carries the union).
+func TestGuardrailGroupSwitchTakesEffectAfterTTL(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
+	defer cancel()
+
+	const (
+		modelA = "e2e-model-a"
+		modelB = "e2e-model-b"
+		// Short tunnel-cache TTL so a group switch propagates in seconds.
+		// Exercises the NB_PROXY_TUNNEL_CACHE_TTL override.
+		cacheTTL = 3 * time.Second
+	)
+
+	vllm, err := harness.StartVLLM(ctx, srv)
+	require.NoError(t, err, "start mock upstream")
+	t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
+
+	grpA, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-gswitch-a"})
+	require.NoError(t, err, "create group A")
+	t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpA.Id) })
+
+	grpB, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-gswitch-b"})
+	require.NoError(t, err, "create group B")
+	t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpB.Id) })
+
+	ephemeral := false
+	sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
+		Name:       "e2e-gswitch-client",
+		Type:       "reusable",
+		ExpiresIn:  86400,
+		UsageLimit: 0,
+		AutoGroups: []string{grpA.Id}, // client starts in group A
+		Ephemeral:  &ephemeral,
+	})
+	require.NoError(t, err, "mint setup key")
+	require.NotEmpty(t, sk.Key, "setup key plaintext")
+
+	staticKey := "static-e2e-token"
+	prov, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
+		Name:        "gswitch",
+		ProviderId:  "openai_api",
+		UpstreamUrl: vllm.URL,
+		ApiKey:      &staticKey,
+		Enabled:     ptr(true),
+		Models: &[]api.AgentNetworkProviderModel{
+			{Id: modelA, InputPer1k: 0.001, OutputPer1k: 0.001},
+			{Id: modelB, InputPer1k: 0.001, OutputPer1k: 0.001},
+		},
+		BootstrapCluster: ptr(harness.AgentNetworkCluster),
+	})
+	require.NoError(t, err, "create provider")
+	t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) })
+
+	mkGuard := func(name, model string) api.AgentNetworkGuardrail {
+		var gr api.AgentNetworkGuardrailRequest
+		gr.Name = name
+		gr.Checks.ModelAllowlist.Enabled = true
+		gr.Checks.ModelAllowlist.Models = []string{model}
+		g, gerr := srv.CreateGuardrail(ctx, gr)
+		require.NoError(t, gerr, "create guardrail %s", name)
+		t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), g.Id) })
+		return g
+	}
+	gA := mkGuard("e2e-gswitch-a", modelA)
+	gB := mkGuard("e2e-gswitch-b", modelB)
+
+	enabled := true
+	polA, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
+		Name:                   "e2e-gswitch-a",
+		Enabled:                &enabled,
+		SourceGroups:           []string{grpA.Id},
+		DestinationProviderIds: []string{prov.Id},
+		GuardrailIds:           &[]string{gA.Id},
+	})
+	require.NoError(t, err, "create policy A")
+	t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polA.Id) })
+
+	polB, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
+		Name:                   "e2e-gswitch-b",
+		Enabled:                &enabled,
+		SourceGroups:           []string{grpB.Id},
+		DestinationProviderIds: []string{prov.Id},
+		GuardrailIds:           &[]string{gB.Id},
+	})
+	require.NoError(t, err, "create policy B")
+	t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polB.Id) })
+
+	settings, err := srv.GetSettings(ctx)
+	require.NoError(t, err, "read settings")
+	require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned")
+
+	proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-gswitch-proxy")
+	require.NoError(t, err, "mint proxy token")
+	px, err := harness.StartProxy(ctx, srv, proxyToken, map[string]string{
+		"NB_PROXY_TUNNEL_CACHE_TTL": cacheTTL.String(),
+	})
+	require.NoError(t, err, "start proxy")
+	t.Cleanup(func() { _ = px.Terminate(context.Background()) })
+
+	cl, err := harness.StartClient(ctx, srv, sk.Key)
+	require.NoError(t, err, "start client")
+	t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
+
+	require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
+	proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
+	require.NoError(t, err, "resolve endpoint to proxy IP")
+	if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
+		t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
+	}
+
+	send := func(model string) (int, string) {
+		code, body, cerr := cl.Chat(ctx, settings.Endpoint, proxyIP, harness.WireChat, model, "Reply with exactly: pong", "")
+		require.NoError(t, cerr, "request must reach the proxy")
+		return code, body
+	}
+	sendUntil := func(model string, want int, timeout time.Duration) (int, string) {
+		var code int
+		var body string
+		deadline := time.Now().Add(timeout)
+		for time.Now().Before(deadline) {
+			code, body = send(model)
+			if code == want {
+				return code, body
+			}
+			time.Sleep(2 * time.Second)
+		}
+		return code, body
+	}
+
+	// Phase 1 — client is in group A: modelA served, modelB denied.
+	code, body := sendUntil(modelA, 200, 90*time.Second)
+	assert.Equal(t, 200, code,
+		"group-A model must be served while the client is in group A; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
+	code, body = send(modelB)
+	assert.Equal(t, 403, code,
+		"group-B model must be denied while the client is in group A; body: %s", body)
+	assert.Contains(t, body, "llm_policy.model_blocked")
+
+	// Switch the client peer from group A to group B.
+	peerID := clientPeerInGroup(t, ctx, grpA.Id)
+	_, err = srv.API().Groups.Update(ctx, grpB.Id, api.PutApiGroupsGroupIdJSONRequestBody{
+		Name:  grpB.Name,
+		Peers: &[]string{peerID},
+	})
+	require.NoError(t, err, "add peer to group B")
+	_, err = srv.API().Groups.Update(ctx, grpA.Id, api.PutApiGroupsGroupIdJSONRequestBody{
+		Name:  grpA.Name,
+		Peers: &[]string{},
+	})
+	require.NoError(t, err, "remove peer from group A")
+
+	// Phase 2 — after the short TTL expires the proxy re-validates the peer,
+	// sees group B, and the decision flips. Poll to absorb TTL + re-validation.
+	code, body = sendUntil(modelB, 200, 60*time.Second)
+	assert.Equal(t, 200, code,
+		"after the group switch + TTL, the group-B model must be served; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
+	code, body = send(modelA)
+	assert.Equal(t, 403, code,
+		"after the switch, the old group-A model must be denied; body: %s", body)
+	assert.Contains(t, body, "llm_policy.model_blocked")
+}
+
+// clientPeerInGroup returns the id of the single peer that is a member of the
+// given group — the test client. The proxy peer is never added to test groups.
+func clientPeerInGroup(t *testing.T, ctx context.Context, groupID string) string {
+	t.Helper()
+	peers, err := srv.API().Peers.List(ctx)
+	require.NoError(t, err, "list peers")
+	for _, p := range peers {
+		for _, g := range p.Groups {
+			if g.Id == groupID {
+				return p.Id
+			}
+		}
+	}
+	t.Fatalf("no peer found in group %s", groupID)
+	return ""
+}

e2e/agentnetwork/guardrail_multipolicy_test.go

@@ -0,0 +1,201 @@
+//go:build e2e
+
+package agentnetwork
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/e2e/harness"
+	"github.com/netbirdio/netbird/shared/management/http/api"
+)
+
+// TestGuardrailMultiPolicyModelAllowlist: modelSelected served (200), grpOther's
+// modelOther denied for the grpMain client (403 model_blocked, no cross-group
+// leak), and openModel on the un-guardrailed policy's provider served (200).
+func TestGuardrailMultiPolicyModelAllowlist(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
+	defer cancel()
+
+	const (
+		modelSelected = "e2e-selected"
+		modelOther    = "e2e-other"
+		openModel     = "e2e-open"
+	)
+
+	vllm, err := harness.StartVLLM(ctx, srv)
+	require.NoError(t, err, "start mock upstream")
+	t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
+
+	grpMain, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-guardrail-mp-main"})
+	require.NoError(t, err, "create main group")
+	t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpMain.Id) })
+
+	grpOther, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-guardrail-mp-other"})
+	require.NoError(t, err, "create other group")
+	t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpOther.Id) })
+
+	ephemeral := false
+	sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
+		Name:       "e2e-guardrail-mp-client",
+		Type:       "reusable",
+		ExpiresIn:  86400,
+		UsageLimit: 0,
+		AutoGroups: []string{grpMain.Id}, // client joins grpMain only
+		Ephemeral:  &ephemeral,
+	})
+	require.NoError(t, err, "mint setup key")
+	require.NotEmpty(t, sk.Key, "setup key plaintext")
+
+	staticKey := "static-e2e-token"
+	models := func(ids ...string) *[]api.AgentNetworkProviderModel {
+		out := make([]api.AgentNetworkProviderModel, 0, len(ids))
+		for _, id := range ids {
+			out = append(out, api.AgentNetworkProviderModel{Id: id, InputPer1k: 0.001, OutputPer1k: 0.001})
+		}
+		return &out
+	}
+
+	// pRestricted declares the two guardrailed models so routing is deterministic
+	// (model -> provider). Created first, so it carries the bootstrap cluster.
+	pRestricted, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
+		Name:             "restricted",
+		ProviderId:       "openai_api",
+		UpstreamUrl:      vllm.URL,
+		ApiKey:           &staticKey,
+		Enabled:          ptr(true),
+		Models:           models(modelSelected, modelOther),
+		BootstrapCluster: ptr(harness.AgentNetworkCluster),
+	})
+	require.NoError(t, err, "create restricted provider")
+	t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), pRestricted.Id) })
+
+	pOpen, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
+		Name:        "open",
+		ProviderId:  "openai_api",
+		UpstreamUrl: vllm.URL,
+		ApiKey:      &staticKey,
+		Enabled:     ptr(true),
+		Models:      models(openModel),
+	})
+	require.NoError(t, err, "create open provider")
+	t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), pOpen.Id) })
+
+	mkGuardrail := func(name, model string) api.AgentNetworkGuardrail {
+		var gr api.AgentNetworkGuardrailRequest
+		gr.Name = name
+		gr.Checks.ModelAllowlist.Enabled = true
+		gr.Checks.ModelAllowlist.Models = []string{model}
+		g, gerr := srv.CreateGuardrail(ctx, gr)
+		require.NoError(t, gerr, "create guardrail %s", name)
+		t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), g.Id) })
+		return g
+	}
+	gMain := mkGuardrail("e2e-guardrail-mp-main", modelSelected)
+	gOther := mkGuardrail("e2e-guardrail-mp-other", modelOther)
+
+	enabled := true
+	// polMain: grpMain restricted to modelSelected on pRestricted.
+	polMain, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
+		Name:                   "e2e-guardrail-mp-main",
+		Enabled:                &enabled,
+		SourceGroups:           []string{grpMain.Id},
+		DestinationProviderIds: []string{pRestricted.Id},
+		GuardrailIds:           &[]string{gMain.Id},
+	})
+	require.NoError(t, err, "create main policy")
+	t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polMain.Id) })
+
+	// polOther: grpOther restricted to modelOther on the SAME provider. The
+	// client is not in grpOther, so modelOther must never be usable by it.
+	polOther, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
+		Name:                   "e2e-guardrail-mp-other",
+		Enabled:                &enabled,
+		SourceGroups:           []string{grpOther.Id},
+		DestinationProviderIds: []string{pRestricted.Id},
+		GuardrailIds:           &[]string{gOther.Id},
+	})
+	require.NoError(t, err, "create other policy")
+	t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polOther.Id) })
+
+	// polOpen: grpMain on pOpen with NO guardrail — unrestricted.
+	polOpen, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
+		Name:                   "e2e-guardrail-mp-open",
+		Enabled:                &enabled,
+		SourceGroups:           []string{grpMain.Id},
+		DestinationProviderIds: []string{pOpen.Id},
+	})
+	require.NoError(t, err, "create open policy")
+	t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polOpen.Id) })
+
+	settings, err := srv.GetSettings(ctx)
+	require.NoError(t, err, "read settings")
+	require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned")
+
+	proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-guardrail-mp-proxy")
+	require.NoError(t, err, "mint proxy token")
+	px, err := harness.StartProxy(ctx, srv, proxyToken)
+	require.NoError(t, err, "start proxy")
+	t.Cleanup(func() { _ = px.Terminate(context.Background()) })
+
+	cl, err := harness.StartClient(ctx, srv, sk.Key)
+	require.NoError(t, err, "start client")
+	t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
+
+	require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
+	proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
+	require.NoError(t, err, "resolve endpoint to proxy IP")
+	if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
+		t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
+	}
+
+	send := func(model string) (int, string) {
+		code, body, cerr := cl.Chat(ctx, settings.Endpoint, proxyIP, harness.WireChat, model, "Reply with exactly: pong", "")
+		require.NoError(t, cerr, "request must reach the proxy")
+		return code, body
+	}
+	// sendUntil200 absorbs first-call tunnel/DNS jitter on the freshly warmed tunnel.
+	sendUntil200 := func(model string) (int, string) {
+		var code int
+		var body string
+		deadline := time.Now().Add(90 * time.Second)
+		for time.Now().Before(deadline) {
+			code, body = send(model)
+			if code == 200 {
+				break
+			}
+			time.Sleep(5 * time.Second)
+		}
+		return code, body
+	}
+
+	t.Run("selected model allowed for its group", func(t *testing.T) {
+		code, body := sendUntil200(modelSelected)
+		assert.Equal(t, 200, code,
+			"grpMain's allowlisted model must be served; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
+	})
+
+	t.Run("other group's model does not leak", func(t *testing.T) {
+		// modelOther is allowlisted only for grpOther. The grpMain client must be
+		// denied by management's per-policy/group check — not waved through by an
+		// account-wide union. This is the security-critical wrong-ALLOW guard.
+		code, body := send(modelOther)
+		assert.Equal(t, 403, code,
+			"another group's allowlisted model must be denied for this caller; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
+		assert.Contains(t, body, "llm_policy.model_blocked",
+			"denial must be a model-allowlist decision; body: %s", body)
+	})
+
+	t.Run("unguarded policy leaves its provider unrestricted", func(t *testing.T) {
+		// polOpen carries no guardrail, so pOpen is unrestricted for grpMain. The
+		// old account-wide union would have blocked openModel (it is on no
+		// allowlist); it must now be served — the false-DENY guard.
+		code, body := sendUntil200(openModel)
+		assert.Equal(t, 200, code,
+			"an un-guardrailed policy's provider must not be blocked by another policy's allowlist; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
+	})
+}

e2e/agentnetwork/guardrail_pergroup_providers_test.go

@@ -0,0 +1,422 @@
+//go:build e2e
+
+package agentnetwork
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/e2e/harness"
+	"github.com/netbirdio/netbird/shared/management/http/api"
+)
+
+// pergroupCase describes one provider surface for the per-group allowlist matrix.
+// selectedReq/otherReq are the model identifiers as they travel in the request
+// (URL path for Bedrock/Vertex, body "model" for chat/messages). selectedAllow/
+// otherAllow are the (normalized) forms the guardrail allowlist holds — for
+// Bedrock these differ from the request form so path normalization is exercised.
+type pergroupCase struct {
+	name          string
+	catalogID     string
+	wire          string // "chat", "messages", "vertex", "bedrock"
+	models        *[]api.AgentNetworkProviderModel
+	selectedReq   string
+	selectedAllow string
+	otherReq      string
+	otherAllow    string
+
+	providerID string // filled during setup
+}
+
+// TestGuardrailPerGroupAllowlist_AllProviders proves the per-policy/group model
+// allowlist end to end across every always-on provider surface, including the
+// path-routed ones (Vertex, Bedrock) where the model travels in the URL.
+//
+// For each provider two policies target it: grpMain (the client) is allowed only
+// selectedReq; grpOther (which the client is NOT in) is allowed only otherReq.
+// The client must get selectedReq served (200) and otherReq denied (403,
+// llm_policy.model_blocked) — the cross-group no-leak property. The deny is the
+// authoritative per-policy/group decision from management (the proxy per-provider
+// backstop carries the union of both models), so this also confirms management
+// receives the correct normalized model for path-routed providers.
+func TestGuardrailPerGroupAllowlist_AllProviders(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Minute)
+	defer cancel()
+
+	const (
+		vertexProject = "e2e-project"
+		vertexRegion  = "global"
+	)
+
+	priced := func(ids ...string) *[]api.AgentNetworkProviderModel {
+		out := make([]api.AgentNetworkProviderModel, 0, len(ids))
+		for _, id := range ids {
+			out = append(out, api.AgentNetworkProviderModel{Id: id, InputPer1k: 0.001, OutputPer1k: 0.001})
+		}
+		return &out
+	}
+
+	cases := []*pergroupCase{
+		{
+			name: "openai", catalogID: "openai_api", wire: harness.WireChat,
+			models:      priced("oai-model-a", "oai-model-b"),
+			selectedReq: "oai-model-a", selectedAllow: "oai-model-a",
+			otherReq: "oai-model-b", otherAllow: "oai-model-b",
+		},
+		{
+			name: "anthropic", catalogID: "anthropic_api", wire: harness.WireMessages,
+			models:      priced("ant-model-a", "ant-model-b"),
+			selectedReq: "ant-model-a", selectedAllow: "ant-model-a",
+			otherReq: "ant-model-b", otherAllow: "ant-model-b",
+		},
+		{
+			// Vertex catalog ids travel bare in the rawPredict path.
+			name: "vertex", catalogID: "vertex_ai_api", wire: "vertex",
+			selectedReq: "claude-sonnet-4-5", selectedAllow: "claude-sonnet-4-5",
+			otherReq: "claude-opus-4-6", otherAllow: "claude-opus-4-6",
+		},
+		{
+			// Bedrock request ids are region-prefixed/versioned; the parser
+			// normalizes them to the catalog key the allowlist holds.
+			name: "bedrock", catalogID: "bedrock_api", wire: "bedrock",
+			selectedReq: "us.anthropic.claude-sonnet-4-5-v1:0", selectedAllow: "anthropic.claude-sonnet-4-5",
+			otherReq: "us.anthropic.claude-opus-4-8-v1:0", otherAllow: "anthropic.claude-opus-4-8",
+		},
+	}
+
+	vllm, err := harness.StartVLLM(ctx, srv)
+	require.NoError(t, err, "start mock upstream")
+	t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
+
+	grpMain, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-pergroup-main"})
+	require.NoError(t, err, "create main group")
+	t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpMain.Id) })
+
+	grpOther, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-pergroup-other"})
+	require.NoError(t, err, "create other group")
+	t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpOther.Id) })
+
+	ephemeral := false
+	sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
+		Name:       "e2e-pergroup-client",
+		Type:       "reusable",
+		ExpiresIn:  86400,
+		UsageLimit: 0,
+		AutoGroups: []string{grpMain.Id},
+		Ephemeral:  &ephemeral,
+	})
+	require.NoError(t, err, "mint setup key")
+	require.NotEmpty(t, sk.Key, "setup key plaintext")
+
+	staticKey := "static-e2e-token"
+	enabled := true
+
+	for i, c := range cases {
+		req := api.AgentNetworkProviderRequest{
+			Name:        "e2e-pergroup-" + c.name,
+			ProviderId:  c.catalogID,
+			UpstreamUrl: vllm.URL,
+			ApiKey:      &staticKey,
+			Enabled:     ptr(true),
+			Models:      c.models,
+		}
+		if i == 0 {
+			req.BootstrapCluster = ptr(harness.AgentNetworkCluster)
+		}
+		prov, perr := srv.CreateProvider(ctx, req)
+		require.NoError(t, perr, "create provider %s", c.name)
+		c.providerID = prov.Id
+		t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) })
+
+		gSel := mkAllowGuardrail(t, ctx, "e2e-pergroup-"+c.name+"-sel", c.selectedAllow)
+		gOth := mkAllowGuardrail(t, ctx, "e2e-pergroup-"+c.name+"-oth", c.otherAllow)
+
+		polMain, merr := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
+			Name:                   "e2e-pergroup-" + c.name + "-main",
+			Enabled:                &enabled,
+			SourceGroups:           []string{grpMain.Id},
+			DestinationProviderIds: []string{prov.Id},
+			GuardrailIds:           &[]string{gSel.Id},
+		})
+		require.NoError(t, merr, "create main policy %s", c.name)
+		t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polMain.Id) })
+
+		polOther, oerr := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
+			Name:                   "e2e-pergroup-" + c.name + "-other",
+			Enabled:                &enabled,
+			SourceGroups:           []string{grpOther.Id},
+			DestinationProviderIds: []string{prov.Id},
+			GuardrailIds:           &[]string{gOth.Id},
+		})
+		require.NoError(t, oerr, "create other policy %s", c.name)
+		t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polOther.Id) })
+	}
+
+	settings, err := srv.GetSettings(ctx)
+	require.NoError(t, err, "read settings")
+	require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned")
+
+	proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-pergroup-proxy")
+	require.NoError(t, err, "mint proxy token")
+	px, err := harness.StartProxy(ctx, srv, proxyToken)
+	require.NoError(t, err, "start proxy")
+	t.Cleanup(func() { _ = px.Terminate(context.Background()) })
+
+	cl, err := harness.StartClient(ctx, srv, sk.Key)
+	require.NoError(t, err, "start client")
+	t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
+
+	require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
+	proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
+	require.NoError(t, err, "resolve endpoint to proxy IP")
+	if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
+		t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
+	}
+
+	send := func(c *pergroupCase, model string) (int, string) {
+		var code int
+		var body string
+		var cerr error
+		switch c.wire {
+		case "vertex":
+			code, body, cerr = cl.Vertex(ctx, settings.Endpoint, proxyIP, vertexProject, vertexRegion, model, "Reply with exactly: pong", "")
+		case "bedrock":
+			code, body, cerr = cl.Bedrock(ctx, settings.Endpoint, proxyIP, model, "Reply with exactly: pong", "")
+		default:
+			code, body, cerr = cl.Chat(ctx, settings.Endpoint, proxyIP, c.wire, model, "Reply with exactly: pong", "")
+		}
+		require.NoError(t, cerr, "request must reach the proxy for %s", c.name)
+		return code, body
+	}
+
+	for _, c := range cases {
+		t.Run(c.name, func(t *testing.T) {
+			// grpMain's own model is served. Retry to absorb tunnel/DNS jitter on
+			// the first call over the freshly warmed tunnel.
+			var code int
+			var body string
+			deadline := time.Now().Add(90 * time.Second)
+			for time.Now().Before(deadline) {
+				code, body = send(c, c.selectedReq)
+				if code == 200 {
+					break
+				}
+				time.Sleep(5 * time.Second)
+			}
+			assert.Equal(t, 200, code,
+				"%s: grpMain's allowlisted model must be served; body: %s\n=== proxy logs ===\n%s", c.name, body, px.Logs(context.Background()))
+
+			// grpOther's model must NOT leak to the grpMain client.
+			code, body = send(c, c.otherReq)
+			assert.Equal(t, 403, code,
+				"%s: another group's allowlisted model must be denied for this caller; body: %s\n=== proxy logs ===\n%s", c.name, body, px.Logs(context.Background()))
+			assert.Contains(t, body, "llm_policy.model_blocked",
+				"%s: denial must be a model-allowlist decision, not routing; body: %s", c.name, body)
+		})
+	}
+}
+
+// TestGuardrailMultiGroupUser proves the per-policy/group decision for a caller
+// that belongs to MULTIPLE groups at once. Two scenarios, one shared stack:
+//
+//   - union across the user's groups: the client is in gUX and gUY, each with
+//     its own policy+guardrail on provider P1 (gUX->union-a, gUY->union-b). The
+//     client may use BOTH models (the union of its groups' allowlists) while a
+//     third, un-allowlisted model is denied.
+//   - an un-guardrailed group lifts the restriction: the client is in gMP and
+//     gMQ on provider P2, where gMP restricts to mix-a but gMQ's policy carries
+//     NO guardrail. Because one applicable policy is unrestricted, the client may
+//     use a model on no allowlist (mix-z) as well as mix-a.
+func TestGuardrailMultiGroupUser(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
+	defer cancel()
+
+	const (
+		unionA = "mg-union-a"
+		unionB = "mg-union-b"
+		unionC = "mg-union-c" // allowlisted by neither group
+		mixA   = "mg-mix-a"
+		mixZ   = "mg-mix-z" // on no allowlist; reachable only via the un-guardrailed policy
+	)
+
+	priced := func(ids ...string) *[]api.AgentNetworkProviderModel {
+		out := make([]api.AgentNetworkProviderModel, 0, len(ids))
+		for _, id := range ids {
+			out = append(out, api.AgentNetworkProviderModel{Id: id, InputPer1k: 0.001, OutputPer1k: 0.001})
+		}
+		return &out
+	}
+
+	vllm, err := harness.StartVLLM(ctx, srv)
+	require.NoError(t, err, "start mock upstream")
+	t.Cleanup(func() { _ = vllm.Terminate(context.Background()) })
+
+	mkGroup := func(name string) *api.Group {
+		g, gerr := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: name})
+		require.NoError(t, gerr, "create group %s", name)
+		t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), g.Id) })
+		return g
+	}
+	gUX := mkGroup("e2e-mg-union-x")
+	gUY := mkGroup("e2e-mg-union-y")
+	gMP := mkGroup("e2e-mg-mix-p")
+	gMQ := mkGroup("e2e-mg-mix-q")
+
+	ephemeral := false
+	sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
+		Name:       "e2e-mg-client",
+		Type:       "reusable",
+		ExpiresIn:  86400,
+		UsageLimit: 0,
+		AutoGroups: []string{gUX.Id, gUY.Id, gMP.Id, gMQ.Id}, // client in all four groups
+		Ephemeral:  &ephemeral,
+	})
+	require.NoError(t, err, "mint setup key")
+	require.NotEmpty(t, sk.Key, "setup key plaintext")
+
+	staticKey := "static-e2e-token"
+	enabled := true
+
+	// P1 — union scenario: two restricting policies, one per group.
+	p1, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
+		Name:             "e2e-mg-union",
+		ProviderId:       "openai_api",
+		UpstreamUrl:      vllm.URL,
+		ApiKey:           &staticKey,
+		Enabled:          ptr(true),
+		Models:           priced(unionA, unionB, unionC),
+		BootstrapCluster: ptr(harness.AgentNetworkCluster),
+	})
+	require.NoError(t, err, "create union provider")
+	t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), p1.Id) })
+
+	polUX, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
+		Name:                   "e2e-mg-union-x",
+		Enabled:                &enabled,
+		SourceGroups:           []string{gUX.Id},
+		DestinationProviderIds: []string{p1.Id},
+		GuardrailIds:           &[]string{mkAllowGuardrail(t, ctx, "e2e-mg-union-x", unionA).Id},
+	})
+	require.NoError(t, err, "create union policy X")
+	t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polUX.Id) })
+
+	polUY, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
+		Name:                   "e2e-mg-union-y",
+		Enabled:                &enabled,
+		SourceGroups:           []string{gUY.Id},
+		DestinationProviderIds: []string{p1.Id},
+		GuardrailIds:           &[]string{mkAllowGuardrail(t, ctx, "e2e-mg-union-y", unionB).Id},
+	})
+	require.NoError(t, err, "create union policy Y")
+	t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polUY.Id) })
+
+	// P2 — mixed scenario: one restricting policy + one un-guardrailed policy.
+	p2, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
+		Name:        "e2e-mg-mix",
+		ProviderId:  "openai_api",
+		UpstreamUrl: vllm.URL,
+		ApiKey:      &staticKey,
+		Enabled:     ptr(true),
+		Models:      priced(mixA, mixZ),
+	})
+	require.NoError(t, err, "create mix provider")
+	t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), p2.Id) })
+
+	polMP, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
+		Name:                   "e2e-mg-mix-p",
+		Enabled:                &enabled,
+		SourceGroups:           []string{gMP.Id},
+		DestinationProviderIds: []string{p2.Id},
+		GuardrailIds:           &[]string{mkAllowGuardrail(t, ctx, "e2e-mg-mix-p", mixA).Id},
+	})
+	require.NoError(t, err, "create mix policy P")
+	t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polMP.Id) })
+
+	polMQ, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
+		Name:                   "e2e-mg-mix-q",
+		Enabled:                &enabled,
+		SourceGroups:           []string{gMQ.Id},
+		DestinationProviderIds: []string{p2.Id}, // NO guardrail -> unrestricted
+	})
+	require.NoError(t, err, "create mix policy Q")
+	t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polMQ.Id) })
+
+	settings, err := srv.GetSettings(ctx)
+	require.NoError(t, err, "read settings")
+	require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned")
+
+	proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-mg-proxy")
+	require.NoError(t, err, "mint proxy token")
+	px, err := harness.StartProxy(ctx, srv, proxyToken)
+	require.NoError(t, err, "start proxy")
+	t.Cleanup(func() { _ = px.Terminate(context.Background()) })
+
+	cl, err := harness.StartClient(ctx, srv, sk.Key)
+	require.NoError(t, err, "start client")
+	t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
+
+	require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
+	proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
+	require.NoError(t, err, "resolve endpoint to proxy IP")
+	if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
+		t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
+	}
+
+	send := func(model string) (int, string) {
+		code, body, cerr := cl.Chat(ctx, settings.Endpoint, proxyIP, harness.WireChat, model, "Reply with exactly: pong", "")
+		require.NoError(t, cerr, "request must reach the proxy")
+		return code, body
+	}
+	sendUntil200 := func(model string) (int, string) {
+		var code int
+		var body string
+		deadline := time.Now().Add(90 * time.Second)
+		for time.Now().Before(deadline) {
+			code, body = send(model)
+			if code == 200 {
+				break
+			}
+			time.Sleep(5 * time.Second)
+		}
+		return code, body
+	}
+
+	t.Run("union across the user's groups", func(t *testing.T) {
+		code, body := sendUntil200(unionA)
+		assert.Equal(t, 200, code, "model allowed by group X must be served; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
+
+		code, body = sendUntil200(unionB)
+		assert.Equal(t, 200, code, "model allowed by group Y must also be served (union across the user's groups); body: %s", body)
+
+		code, body = send(unionC)
+		assert.Equal(t, 403, code, "a model on neither group's allowlist must be denied; body: %s", body)
+		assert.Contains(t, body, "llm_policy.model_blocked")
+	})
+
+	t.Run("an un-guardrailed group lifts the restriction", func(t *testing.T) {
+		code, body := sendUntil200(mixA)
+		assert.Equal(t, 200, code, "the restricted group's model must be served; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
+
+		code, body = sendUntil200(mixZ)
+		assert.Equal(t, 200, code,
+			"a non-allowlisted model must be served because the user is also in a group whose policy has no guardrail; body: %s\n=== proxy logs ===\n%s", body, px.Logs(context.Background()))
+	})
+}
+
+// mkAllowGuardrail creates a guardrail whose model allowlist is enabled and holds
+// exactly the given model, registering cleanup.
+func mkAllowGuardrail(t *testing.T, ctx context.Context, name, model string) api.AgentNetworkGuardrail {
+	t.Helper()
+	var gr api.AgentNetworkGuardrailRequest
+	gr.Name = name
+	gr.Checks.ModelAllowlist.Enabled = true
+	gr.Checks.ModelAllowlist.Models = []string{model}
+	g, err := srv.CreateGuardrail(ctx, gr)
+	require.NoError(t, err, "create guardrail %s", name)
+	t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), g.Id) })
+	return g
+}

e2e/agentnetwork/guardrail_test.go

@@ -4,6 +4,7 @@ package agentnetwork
 
 import (
 	"context"
+	"regexp"
 	"strings"
 	"testing"
 	"time"
@@ -15,13 +16,29 @@ import (
 	"github.com/netbirdio/netbird/shared/management/http/api"
 )
 
+// bedrockRegionPrefixes and bedrockVersionSuffix mirror the proxy's Bedrock
+// model normalization (region/inference-profile prefix + version suffix) so the
+// provider is registered under the same catalog key the router matches against.
+var (
+	bedrockRegionPrefixes = []string{"us.", "eu.", "apac.", "global."}
+	bedrockVersionSuffix  = regexp.MustCompile(`-(\d{8}-)?v\d+(:\d+)?$`)
+)
+
 // catalogModel returns the normalized catalog id the proxy stamps for a
-// path-routed provider's configured model — the form the guardrail allowlist is
-// compared against (region prefix / @version stripped).
+// path-routed provider's configured model — the form the router and guardrail
+// allowlist compare against (Bedrock region prefix + version stripped, Vertex
+// @version stripped).
 func catalogModel(pc providerCase) string {
 	switch pc.kind {
 	case harness.WireBedrock:
-		return strings.TrimPrefix(pc.model, "us.")
+		m := pc.model
+		for _, p := range bedrockRegionPrefixes {
+			if strings.HasPrefix(m, p) {
+				m = m[len(p):]
+				break
+			}
+		}
+		return bedrockVersionSuffix.ReplaceAllString(m, "")
 	case harness.WireVertex:
 		return strings.SplitN(pc.model, "@", 2)[0]
 	default:
@@ -35,7 +52,9 @@ func catalogModel(pc providerCase) string {
 func disallowedModel(pc providerCase) string {
 	switch pc.kind {
 	case harness.WireBedrock:
-		return "us.anthropic.claude-opus-4-8"
+		// Same profile prefix as the allowed model so only the model name
+		// differs; the guardrail must deny it before it reaches AWS.
+		return strings.SplitN(pc.model, ".", 2)[0] + ".anthropic.claude-opus-4-8"
 	case harness.WireVertex:
 		return "claude-opus-4-8@20250101"
 	default:
@@ -55,7 +74,7 @@ func sendModel(ctx context.Context, t *testing.T, cl *harness.Client, endpoint,
 	case harness.WireVertex:
 		code, _, err = cl.Vertex(ctx, endpoint, proxyIP, pc.project, pc.region, model, "Reply with exactly: pong", "")
 	default:
-		code, _, err = cl.Chat(ctx, endpoint, proxyIP, pc.kind, model, "Reply with exactly: pong", "")
+		code, _, err = cl.ChatPrefixed(ctx, endpoint, proxyIP, pc.pathPrefix, pc.kind, model, "Reply with exactly: pong", "")
 	}
 	require.NoError(t, err, "request must reach the proxy for %s", pc.name)
 	return code
@@ -147,11 +166,12 @@ func TestModelAllowlistEnforced(t *testing.T) {
 	t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
 
 	require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
+	// Probe first: the GET resolves the endpoint (DNS error fails) and its first packet wakes the lazy proxy peer, so WaitProxyPeer sees it connected; any HTTP status counts.
+	proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
+	require.NoError(t, err, "resolve agent-network endpoint to proxy IP")
 	if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
 		t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
 	}
-	proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
-	require.NoError(t, err, "resolve agent-network endpoint to proxy IP")
 
 	for _, pc := range providers {
 		pc := pc

e2e/agentnetwork/skiptls_test.go

@@ -104,11 +104,12 @@ func TestProviderSkipTLSVerification(t *testing.T) {
 	t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
 
 	require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
+	// Probe first: the GET resolves the endpoint (DNS error fails) and its first packet wakes the lazy proxy peer, so WaitProxyPeer sees it connected; any HTTP status counts.
+	proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
+	require.NoError(t, err, "resolve endpoint to proxy IP")
 	if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
 		t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
 	}
-	proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
-	require.NoError(t, err, "resolve endpoint to proxy IP")
 
 	// Positive: skip=true reaches the self-signed upstream. Retry to absorb
 	// tunnel/DNS jitter on the first call; success also proves the path works.

e2e/agentnetwork/vllm_test.go

@@ -106,11 +106,12 @@ func TestVLLMProvider(t *testing.T) {
 	t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
 
 	require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
+	// Probe first: the GET resolves the endpoint (DNS error fails) and its first packet wakes the lazy proxy peer, so WaitProxyPeer sees it connected; any HTTP status counts.
+	proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
+	require.NoError(t, err, "resolve endpoint to proxy IP")
 	if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
 		t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
 	}
-	proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
-	require.NoError(t, err, "resolve endpoint to proxy IP")
 
 	before, _ := srv.ListAccessLogs(ctx)
 	sessionID := "e2e-session-vllm"

e2e/harness/client.go

@@ -4,6 +4,7 @@ package harness
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"os/exec"
@@ -167,22 +168,54 @@ func (cl *Client) pollStatus(ctx context.Context, timeout time.Duration, want st
 	return fmt.Errorf("timed out waiting for %q; last status:\n%s", want, last)
 }
 
-// ResolveProxyIP resolves the agent-network endpoint to the proxy peer's
-// NetBird IP from inside the client (via magic DNS).
+const (
+	// curlExitCouldNotResolve is curl's exit code for a DNS resolution failure, distinct from connection-level failures.
+	curlExitCouldNotResolve = 6
+	// dnsProbeRetryWindow bounds DNS-failure retries: the synthesized zone lands a beat after management connects, so early NXDOMAIN is propagation; a zone still absent after this window is a real failure.
+	dnsProbeRetryWindow   = 30 * time.Second
+	dnsProbeRetryInterval = 2 * time.Second
+)
+
+// ResolveProxyIP GETs https://<endpoint>/ from the client's netns: any HTTP status proves DNS + tunnel and wakes the lazy proxy peer; only DNS failures retry, within dnsProbeRetryWindow. Returns the connected IP for --resolve pinning.
 func (cl *Client) ResolveProxyIP(ctx context.Context, endpoint string) (string, error) {
-	code, reader, err := cl.container.Exec(ctx, []string{"getent", "hosts", endpoint}, tcexec.Multiplexed())
-	if err != nil {
-		return "", err
+	args := []string{
+		"run", "--rm",
+		"--network", "container:" + cl.container.GetContainerID(),
+		curlImage,
+		"-ksS", "-o", "/dev/null",
+		"--connect-timeout", "30", "--max-time", "60",
+		"-w", "%{remote_ip}",
+		"https://" + endpoint + "/",
 	}
-	out, _ := io.ReadAll(reader)
-	if code != 0 {
-		return "", fmt.Errorf("getent hosts %s exited %d", endpoint, code)
+	deadline := time.Now().Add(dnsProbeRetryWindow)
+	for {
+		cmd := exec.CommandContext(ctx, "docker", args...)
+		var stdout, stderr strings.Builder
+		cmd.Stdout = &stdout
+		cmd.Stderr = &stderr
+		err := cmd.Run()
+		if err == nil {
+			ip := strings.TrimSpace(stdout.String())
+			if ip == "" {
+				return "", fmt.Errorf("got an HTTP response from %s but no remote IP", endpoint)
+			}
+			return ip, nil
+		}
+
+		var exitErr *exec.ExitError
+		if !errors.As(err, &exitErr) || exitErr.ExitCode() != curlExitCouldNotResolve {
+			return "", fmt.Errorf("no HTTP response from %s: %w (%s)", endpoint, err, strings.TrimSpace(stderr.String()))
+		}
+		dnsErr := fmt.Errorf("DNS resolution failed for %s: %s", endpoint, strings.TrimSpace(stderr.String()))
+		if time.Until(deadline) < dnsProbeRetryInterval {
+			return "", dnsErr
+		}
+		select {
+		case <-ctx.Done():
+			return "", fmt.Errorf("%w (%w)", dnsErr, ctx.Err())
+		case <-time.After(dnsProbeRetryInterval):
+		}
 	}
-	fields := strings.Fields(string(out))
-	if len(fields) == 0 {
-		return "", fmt.Errorf("no address for %s", endpoint)
-	}
-	return fields[0], nil
 }
 
 // Wire shapes for Chat.
@@ -206,18 +239,29 @@ const (
 // the wire shape: WireChat (OpenAI) or WireMessages (Anthropic). A non-empty
 // sessionID is sent as the universal x-session-id header the proxy records.
 func (cl *Client) Chat(ctx context.Context, endpoint, proxyIP, kind, model, prompt, sessionID string) (int, string, error) {
+	return cl.ChatPrefixed(ctx, endpoint, proxyIP, "", kind, model, prompt, sessionID)
+}
+
+// ChatPrefixed is Chat with a base-URL path prefix prepended to the wire
+// path, mirroring agents whose base URL carries a shape-selecting prefix that
+// rides through to the upstream — e.g. Claude Code against a Kimi provider
+// sets ANTHROPIC_BASE_URL=https://<endpoint>/anthropic so the proxy forwards
+// /anthropic/v1/messages to Moonshot's Anthropic surface while the provider's
+// upstream URL stays the bare https://api.moonshot.ai. Empty prefix is plain
+// Chat.
+func (cl *Client) ChatPrefixed(ctx context.Context, endpoint, proxyIP, pathPrefix, kind, model, prompt, sessionID string) (int, string, error) {
 	var path, body string
 	var headers []string
 	switch kind {
 	case WireMessages:
 		path = "/v1/messages"
 		headers = []string{"anthropic-version: 2023-06-01"}
-		body = fmt.Sprintf(`{"model":%q,"max_tokens":64,"messages":[{"role":"user","content":%q}]}`, model, prompt)
+		body = fmt.Sprintf(`{"model":%q,"max_tokens":2048,"messages":[{"role":"user","content":%q}]}`, model, prompt)
 	default:
 		path = "/v1/chat/completions"
 		body = fmt.Sprintf(`{"model":%q,"messages":[{"role":"user","content":%q}]}`, model, prompt)
 	}
-	return cl.post(ctx, endpoint, proxyIP, path, body, withSessionID(headers, sessionID))
+	return cl.post(ctx, endpoint, proxyIP, pathPrefix+path, body, withSessionID(headers, sessionID))
 }
 
 // Vertex issues an Anthropic-on-Vertex rawPredict POST over the tunnel. Unlike
@@ -227,7 +271,7 @@ func (cl *Client) Chat(ctx context.Context, endpoint, proxyIP, kind, model, prom
 // is sent as the universal x-session-id header the proxy records.
 func (cl *Client) Vertex(ctx context.Context, endpoint, proxyIP, project, region, model, prompt, sessionID string) (int, string, error) {
 	path := fmt.Sprintf("/v1/projects/%s/locations/%s/publishers/anthropic/models/%s:rawPredict", project, region, model)
-	body := fmt.Sprintf(`{"anthropic_version":"vertex-2023-10-16","max_tokens":64,"messages":[{"role":"user","content":%q}]}`, prompt)
+	body := fmt.Sprintf(`{"anthropic_version":"vertex-2023-10-16","max_tokens":2048,"messages":[{"role":"user","content":%q}]}`, prompt)
 	return cl.post(ctx, endpoint, proxyIP, path, body, withSessionID(nil, sessionID))
 }
 
@@ -238,7 +282,7 @@ func (cl *Client) Vertex(ctx context.Context, endpoint, proxyIP, project, region
 // header the proxy records.
 func (cl *Client) Bedrock(ctx context.Context, endpoint, proxyIP, model, prompt, sessionID string) (int, string, error) {
 	path := "/model/" + model + "/invoke"
-	body := fmt.Sprintf(`{"anthropic_version":"bedrock-2023-05-31","max_tokens":64,"messages":[{"role":"user","content":%q}]}`, prompt)
+	body := fmt.Sprintf(`{"anthropic_version":"bedrock-2023-05-31","max_tokens":2048,"messages":[{"role":"user","content":%q}]}`, prompt)
 	return cl.post(ctx, endpoint, proxyIP, path, body, withSessionID(nil, sessionID))
 }
 
@@ -297,7 +341,7 @@ func (cl *Client) Terminate(ctx context.Context) error {
 	return cl.container.Terminate(ctx)
 }
 
-// containerLogs reads up to 256 KiB of a container's logs for diagnostics.
+// containerLogs reads up to 4 MiB of a container's logs for diagnostics — enough for a whole provider-matrix run.
 func containerLogs(ctx context.Context, c testcontainers.Container) string {
 	if c == nil {
 		return ""
@@ -307,6 +351,6 @@ func containerLogs(ctx context.Context, c testcontainers.Container) string {
 		return fmt.Sprintf("<logs error: %v>", err)
 	}
 	defer r.Close()
-	b, _ := io.ReadAll(io.LimitReader(r, 256<<10))
+	b, _ := io.ReadAll(io.LimitReader(r, 4<<20))
 	return string(b)
 }

e2e/harness/combined.go

@@ -221,6 +221,29 @@ func (c *Combined) CreateProxyTokenCLI(ctx context.Context, name string) (string
 	return "", fmt.Errorf("token not found in CLI output: %s", string(out))
 }
 
+// SnapshotStoreDB copies the management sqlite store (with WAL/SHM sidecars) out of the bind-mounted
+// data dir into dstDir and returns the copy's path; reading a copy avoids locking against live writes.
+func (c *Combined) SnapshotStoreDB(dstDir string) (string, error) {
+	src := filepath.Join(c.workDir, "data", "store.db")
+	if _, err := os.Stat(src); err != nil {
+		return "", fmt.Errorf("management store not found at %s: %w", src, err)
+	}
+	dst := filepath.Join(dstDir, "store.db")
+	for _, suffix := range []string{"", "-wal", "-shm"} {
+		data, err := os.ReadFile(src + suffix)
+		if err != nil {
+			if os.IsNotExist(err) && suffix != "" {
+				continue // sidecar only exists in WAL mode
+			}
+			return "", fmt.Errorf("read %s: %w", src+suffix, err)
+		}
+		if err := os.WriteFile(dst+suffix, data, 0o600); err != nil {
+			return "", fmt.Errorf("write %s: %w", dst+suffix, err)
+		}
+	}
+	return dst, nil
+}
+
 // Logs returns the combined server container logs, for diagnostics.
 func (c *Combined) Logs(ctx context.Context) string {
 	return containerLogs(ctx, c.container)

e2e/harness/proxy.go

@@ -38,7 +38,11 @@ type Proxy struct {
 // network, registered via the given account proxy token and serving the
 // AgentNetworkCluster over a self-signed wildcard cert. It does not wait for
 // peer connectivity — callers poll management for the proxy peer.
-func StartProxy(ctx context.Context, c *Combined, proxyToken string) (*Proxy, error) {
+// StartProxy launches the reverse-proxy container. Optional envOverrides are
+// merged into the container environment after the defaults, so callers can set
+// or override any NB_PROXY_* var (e.g. NB_PROXY_TUNNEL_CACHE_TTL for tests that
+// need a short authorization-cache window).
+func StartProxy(ctx context.Context, c *Combined, proxyToken string, envOverrides ...map[string]string) (*Proxy, error) {
 	root, err := repoRoot()
 	if err != nil {
 		return nil, err
@@ -93,6 +97,12 @@ func StartProxy(ctx context.Context, c *Combined, proxyToken string) (*Proxy, er
 		WaitingFor: wait.ForLog("Initial mapping sync complete").WithStartupTimeout(90 * time.Second),
 	}
 
+	for _, ov := range envOverrides {
+		for k, v := range ov {
+			req.Env[k] = v
+		}
+	}
+
 	ctr, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
 		ContainerRequest: req,
 		Started:          true,

idp/dex/config.go

@@ -308,7 +308,7 @@ func (s *Storage) OpenStorage(logger *slog.Logger) (storage.Storage, error) {
 		if file == "" {
 			return nil, fmt.Errorf("sqlite3 storage requires 'file' config")
 		}
-		return (&sql.SQLite3{File: file}).Open(logger)
+		return newSQLite3(file).Open(logger)
 	case "postgres":
 		dsn, _ := s.Config["dsn"].(string)
 		if dsn == "" {

idp/dex/provider.go

@@ -20,7 +20,6 @@ import (
 	"github.com/dexidp/dex/server"
 	"github.com/dexidp/dex/server/signer"
 	"github.com/dexidp/dex/storage"
-	"github.com/dexidp/dex/storage/sql"
 	"github.com/go-jose/go-jose/v4"
 	"github.com/google/uuid"
 	"github.com/prometheus/client_golang/prometheus"
@@ -41,7 +40,7 @@ type Config struct {
 	GRPCAddr string
 }
 
-const localConnectorID = "local"
+const LocalConnectorID = "local"
 
 // Provider wraps a Dex server
 type Provider struct {
@@ -79,7 +78,7 @@ func NewProvider(ctx context.Context, config *Config) (*Provider, error) {
 
 	// Initialize SQLite storage
 	dbPath := filepath.Join(config.DataDir, "oidc.db")
-	sqliteConfig := &sql.SQLite3{File: dbPath}
+	sqliteConfig := newSQLite3(dbPath)
 	stor, err := sqliteConfig.Open(logger)
 	if err != nil {
 		return nil, fmt.Errorf("failed to open storage: %w", err)
@@ -495,18 +494,60 @@ func (p *Provider) Storage() storage.Storage {
 	return p.storage
 }
 
+// SetClientsMFAChain updates the MFAChain field on OAuth2 clients in Dex storage.
+// Pass a non-empty slice (e.g. []string{"default-totp"}) to enable MFA, or nil to disable it.
+func SetClientsMFAChain(ctx context.Context, st storage.Storage, clientIDs []string, mfaChain []string) error {
+	previousChains := make(map[string][]string, len(clientIDs))
+	for _, clientID := range clientIDs {
+		client, err := st.GetClient(ctx, clientID)
+		if err != nil {
+			return fmt.Errorf("failed to get client %s before MFA chain update: %w", clientID, err)
+		}
+		previousChains[clientID] = cloneMFAChain(client.MFAChain)
+	}
+
+	updatedClientIDs := make([]string, 0, len(clientIDs))
+	for _, clientID := range clientIDs {
+		if err := st.UpdateClient(ctx, clientID, func(old storage.Client) (storage.Client, error) {
+			old.MFAChain = cloneMFAChain(mfaChain)
+			return old, nil
+		}); err != nil {
+			if rollbackErr := rollbackClientsMFAChain(ctx, st, updatedClientIDs, previousChains); rollbackErr != nil {
+				return fmt.Errorf("failed to update MFA chain on client %s: %w (also failed to roll back previous MFA chains: %v)", clientID, err, rollbackErr)
+			}
+			return fmt.Errorf("failed to update MFA chain on client %s: %w", clientID, err)
+		}
+		updatedClientIDs = append(updatedClientIDs, clientID)
+	}
+	return nil
+}
+
+func rollbackClientsMFAChain(ctx context.Context, st storage.Storage, clientIDs []string, previousChains map[string][]string) error {
+	var rollbackErrs []error
+	for i := len(clientIDs) - 1; i >= 0; i-- {
+		clientID := clientIDs[i]
+		previousChain := cloneMFAChain(previousChains[clientID])
+		if err := st.UpdateClient(ctx, clientID, func(old storage.Client) (storage.Client, error) {
+			old.MFAChain = previousChain
+			return old, nil
+		}); err != nil {
+			rollbackErrs = append(rollbackErrs, fmt.Errorf("client %s: %w", clientID, err))
+		}
+	}
+	return errors.Join(rollbackErrs...)
+}
+
+func cloneMFAChain(chain []string) []string {
+	if chain == nil {
+		return nil
+	}
+	return append([]string(nil), chain...)
+}
+
 // SetClientsMFAChain updates the MFAChain field on the dashboard and CLI OAuth2 clients.
 // Pass a non-empty slice (e.g. []string{"default-totp"}) to enable MFA, or nil to disable it.
 func (p *Provider) SetClientsMFAChain(ctx context.Context, clientIDs []string, mfaChain []string) error {
-	for _, clientID := range clientIDs {
-		if err := p.storage.UpdateClient(ctx, clientID, func(old storage.Client) (storage.Client, error) {
-			old.MFAChain = mfaChain
-			return old, nil
-		}); err != nil {
-			return fmt.Errorf("failed to update MFA chain on client %s: %w", clientID, err)
-		}
-	}
-	return nil
+	return SetClientsMFAChain(ctx, p.storage, clientIDs, mfaChain)
 }
 
 // Handler returns the Dex server as an http.Handler for embedding in another server.
@@ -546,7 +587,7 @@ func (p *Provider) CreateUser(ctx context.Context, email, username, password str
 
 	// Encode the user ID in Dex's format: base64(protobuf{user_id, connector_id})
 	// This matches the format Dex uses in JWT tokens
-	encodedID := EncodeDexUserID(userID, localConnectorID)
+	encodedID := EncodeDexUserID(userID, LocalConnectorID)
 	return encodedID, nil
 }
 
@@ -625,7 +666,7 @@ func DecodeDexUserID(encodedID string) (userID, connectorID string, err error) {
 // local password connector.
 func IsLocalUserID(encodedID string) bool {
 	_, connectorID, err := DecodeDexUserID(encodedID)
-	return err == nil && connectorID == localConnectorID
+	return err == nil && connectorID == LocalConnectorID
 }
 
 // GetUser returns a user by email

idp/dex/provider_test.go

@@ -3,6 +3,8 @@ package dex
 import (
 	"context"
 	"encoding/json"
+	"errors"
+	"io"
 	"log/slog"
 	"net/http"
 	"net/http/httptest"
@@ -11,11 +13,44 @@ import (
 	"testing"
 
 	"github.com/dexidp/dex/storage"
+	"github.com/dexidp/dex/storage/memory"
 	sqllib "github.com/dexidp/dex/storage/sql"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
+type updateFailingStorage struct {
+	storage.Storage
+	failClientID string
+}
+
+func (s *updateFailingStorage) UpdateClient(ctx context.Context, id string, updater func(storage.Client) (storage.Client, error)) error {
+	if id == s.failClientID {
+		return errors.New("forced update failure")
+	}
+	return s.Storage.UpdateClient(ctx, id, updater)
+}
+
+func TestSetClientsMFAChainRollsBackUpdatedClients(t *testing.T) {
+	ctx := context.Background()
+	st := memory.New(slog.New(slog.NewTextHandler(io.Discard, nil)))
+
+	require.NoError(t, st.CreateClient(ctx, storage.Client{ID: "client-1", MFAChain: []string{"old-1"}}))
+	require.NoError(t, st.CreateClient(ctx, storage.Client{ID: "client-2", MFAChain: []string{"old-2"}}))
+
+	err := SetClientsMFAChain(ctx, &updateFailingStorage{Storage: st, failClientID: "client-2"}, []string{"client-1", "client-2"}, []string{"new"})
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "failed to update MFA chain on client client-2")
+
+	client1, err := st.GetClient(ctx, "client-1")
+	require.NoError(t, err)
+	require.Equal(t, []string{"old-1"}, client1.MFAChain)
+
+	client2, err := st.GetClient(ctx, "client-2")
+	require.NoError(t, err)
+	require.Equal(t, []string{"old-2"}, client2.MFAChain)
+}
+
 func TestUserCreationFlow(t *testing.T) {
 	ctx := context.Background()
 

idp/dex/sqlite_cgo.go

@@ -0,0 +1,15 @@
+//go:build cgo
+
+package dex
+
+import (
+	sql "github.com/dexidp/dex/storage/sql"
+)
+
+// newSQLite3 builds the dex SQLite3 config. CGO builds use the upstream
+// struct that takes a File path. Non-CGO builds get an empty stub whose
+// Open() returns the dex "SQLite not available" error — correct behaviour
+// for binaries that can't link sqlite3 (e.g. cross-compiled ARM targets).
+func newSQLite3(file string) *sql.SQLite3 {
+	return &sql.SQLite3{File: file}
+}

idp/dex/sqlite_nocgo.go

@@ -0,0 +1,15 @@
+//go:build !cgo
+
+package dex
+
+import (
+	sql "github.com/dexidp/dex/storage/sql"
+)
+
+// newSQLite3 for non-CGO builds. The dex SQLite3 stub has no fields and its
+// Open() returns an error documenting the missing CGO support — correct
+// behaviour for cross-compiled artefacts that never actually run the
+// embedded IdP. The `file` argument is ignored.
+func newSQLite3(_ string) *sql.SQLite3 {
+	return &sql.SQLite3{}
+}

infrastructure_files/getting-started-enterprise.sh

@@ -234,9 +234,6 @@ init_environment() {
 
   NETBIRD_LICENSE_KEY=$(read_secret "Enter license key (input hidden)")
 
-  GHCR_USERNAME="netbirdExtAccess1"
-  GHCR_TOKEN=$(read_secret "Enter GHCR token (input hidden)")
-
   POSTGRES_USER="netbird"
   POSTGRES_DB="netbird"
   POSTGRES_PASSWORD=$(rand_secret)
@@ -263,10 +260,6 @@ init_environment() {
   install -m 600 /dev/null config.yaml
   render_config_yaml >> config.yaml
 
-  echo "Logging in to ghcr.io ..."
-  printf '%s' "$GHCR_TOKEN" | docker login ghcr.io -u "$GHCR_USERNAME" --password-stdin
-  unset GHCR_TOKEN
-
   echo ""
   echo "Pulling images ..."
   $DOCKER_COMPOSE_COMMAND pull

infrastructure_files/getting-started.sh

@@ -556,7 +556,7 @@ start_services_and_show_instructions() {
       echo "Creating proxy access token..."
       # Use docker exec with bash to run the token command directly
       PROXY_TOKEN=$($DOCKER_COMPOSE_COMMAND exec -T netbird-server \
-        /go/bin/netbird-server token create --name "default-proxy" --config /etc/netbird/config.yaml 2>/dev/null | grep "^Token:" | awk '{print $2}')
+        /go/bin/netbird-server admin token create --name "default-proxy" --config /etc/netbird/config.yaml 2>/dev/null | grep "^Token:" | awk '{print $2}')
 
       if [[ -z "$PROXY_TOKEN" ]]; then
         echo "ERROR: Failed to create proxy token. Check netbird-server logs." > /dev/stderr

infrastructure_files/migrate-to-enterprise.sh

@@ -490,8 +490,6 @@ init_migration() {
   echo ""
   echo "Step 1: Image swap (community → Enterprise). License key required."
   NB_LICENSE_KEY=$(read_secret "  License key")
-  GHCR_USERNAME="netbirdExtAccess1"
-  GHCR_TOKEN=$(read_secret "  GHCR token (input hidden)")
 
   # Step 2 — optional
   echo ""
@@ -588,11 +586,6 @@ apply_changes() {
     fi
   } >> "$ENV_FILE"
 
-  echo ""
-  echo "Logging in to ghcr.io ..."
-  printf '%s' "$GHCR_TOKEN" | docker login ghcr.io -u "$GHCR_USERNAME" --password-stdin
-  unset GHCR_TOKEN
-
   echo ""
   echo "Pulling enterprise images ..."
   $DOCKER_COMPOSE_COMMAND pull

management/cmd/admin.go

@@ -0,0 +1,177 @@
+package cmd
+
+import (
+	"context"
+	"fmt"
+	"path"
+	"path/filepath"
+
+	"github.com/dexidp/dex/storage"
+	log "github.com/sirupsen/logrus"
+	"github.com/spf13/cobra"
+
+	"github.com/netbirdio/netbird/formatter/hook"
+	admincmd "github.com/netbirdio/netbird/management/cmd/admin"
+	tokencmd "github.com/netbirdio/netbird/management/cmd/token"
+	nbconfig "github.com/netbirdio/netbird/management/internals/server/config"
+	"github.com/netbirdio/netbird/management/server/activity"
+	activitystore "github.com/netbirdio/netbird/management/server/activity/store"
+	"github.com/netbirdio/netbird/management/server/store"
+	"github.com/netbirdio/netbird/util"
+)
+
+var adminDatadir string
+
+// newAdminCommands creates the admin command tree with management-specific resource openers.
+func newAdminCommands() *cobra.Command {
+	cmd := admincmd.NewCommands(admincmd.Openers{
+		Resources: withAdminResources,
+		Store:     withAdminStoreOnly,
+		IDP:       withAdminIDPOnly,
+	})
+	cmd.PersistentFlags().StringVar(&adminDatadir, "datadir", "", "Override the data directory from config (used for store.db and the default idp.db)")
+	return cmd
+}
+
+func newLegacyTokenCommand() *cobra.Command {
+	cmd := tokencmd.NewCommands(tokencmd.StoreOpener(withAdminStoreOnly))
+	cmd.Deprecated = "use 'admin token' instead"
+	cmd.PersistentFlags().StringVar(&nbconfig.MgmtConfigPath, "config", defaultMgmtConfig, "Netbird config file location")
+	return cmd
+}
+
+// withAdminResources initializes logging, loads config, opens the management store
+// and embedded IdP storage, and calls fn.
+func withAdminResources(cmd *cobra.Command, fn func(ctx context.Context, resources admincmd.Resources) error) error {
+	return withAdminConfig(cmd, true, func(ctx context.Context, config *nbconfig.Config, datadir string) error {
+		managementStore, err := openAdminStore(ctx, config, datadir)
+		if err != nil {
+			return err
+		}
+		defer admincmd.CloseStore(ctx, managementStore)
+
+		idpStorage, idpStorageFile, err := admincmd.OpenIDPStorage(config)
+		if err != nil {
+			return err
+		}
+		defer admincmd.CloseIDPStorage(idpStorage)
+
+		eventStore, esErr := openAdminEventStore(ctx, config, datadir)
+		if esErr != nil {
+			_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "Warning: audit events will not be recorded: %v\n", esErr)
+		}
+		if eventStore != nil {
+			defer func() {
+				if err := eventStore.Close(ctx); err != nil {
+					log.Debugf("close activity event store: %v", err)
+				}
+			}()
+		}
+
+		return fn(ctx, admincmd.Resources{Store: managementStore, IDPStorage: idpStorage, IDPStorageFile: idpStorageFile, EventStore: eventStore})
+	})
+}
+
+// withAdminStoreOnly opens only the management store for admin subcommands that do not
+// need embedded IdP storage.
+func withAdminStoreOnly(cmd *cobra.Command, fn func(ctx context.Context, s store.Store) error) error {
+	return withAdminConfig(cmd, false, func(ctx context.Context, config *nbconfig.Config, datadir string) error {
+		managementStore, err := openAdminStore(ctx, config, datadir)
+		if err != nil {
+			return err
+		}
+		defer admincmd.CloseStore(ctx, managementStore)
+
+		return fn(ctx, managementStore)
+	})
+}
+
+func withAdminIDPOnly(cmd *cobra.Command, fn func(ctx context.Context, idpStorage storage.Storage, storageFile string) error) error {
+	return withAdminConfig(cmd, true, func(ctx context.Context, config *nbconfig.Config, _ string) error {
+		idpStorage, idpStorageFile, err := admincmd.OpenIDPStorage(config)
+		if err != nil {
+			return err
+		}
+		defer admincmd.CloseIDPStorage(idpStorage)
+
+		return fn(ctx, idpStorage, idpStorageFile)
+	})
+}
+
+func withAdminConfig(cmd *cobra.Command, applyIDPDefaults bool, fn func(ctx context.Context, config *nbconfig.Config, datadir string) error) error {
+	if err := util.InitLog("error", "console"); err != nil {
+		return fmt.Errorf("init log: %w", err)
+	}
+
+	ctx := context.WithValue(cmd.Context(), hook.ExecutionContextKey, hook.SystemSource) //nolint:staticcheck
+
+	config, datadir, err := loadAdminMgmtConfig(ctx, applyIDPDefaults)
+	if err != nil {
+		return fmt.Errorf("load config: %w", err)
+	}
+
+	return fn(ctx, config, datadir)
+}
+
+func loadAdminMgmtConfig(ctx context.Context, applyIDPDefaults bool) (*nbconfig.Config, string, error) {
+	config := &nbconfig.Config{}
+	if _, err := util.ReadJsonWithEnvSub(nbconfig.MgmtConfigPath, config); err != nil {
+		return nil, "", err
+	}
+
+	if applyIDPDefaults {
+		if err := ApplyEmbeddedIdPConfig(ctx, config); err != nil {
+			return nil, "", err
+		}
+	}
+
+	datadir := config.Datadir
+	applyAdminDatadirOverride(config, &datadir)
+	return config, datadir, nil
+}
+
+func applyAdminDatadirOverride(config *nbconfig.Config, datadir *string) {
+	if adminDatadir == "" {
+		return
+	}
+
+	oldDatadir := *datadir
+	*datadir = adminDatadir
+	if config.EmbeddedIdP != nil && config.EmbeddedIdP.Storage.Type == "sqlite3" && isDefaultIDPStorageFile(config.EmbeddedIdP.Storage.Config.File, oldDatadir) {
+		config.EmbeddedIdP.Storage.Config.File = filepath.Join(*datadir, "idp.db")
+	}
+}
+
+func isDefaultIDPStorageFile(file, datadir string) bool {
+	if file == "" {
+		return true
+	}
+	defaultFile := filepath.Join(datadir, "idp.db")
+	legacyDefaultFile := path.Join(datadir, "idp.db")
+	legacySlashDefaultFile := path.Join(filepath.ToSlash(datadir), "idp.db")
+	return filepath.Clean(file) == filepath.Clean(defaultFile) ||
+		file == legacyDefaultFile ||
+		filepath.ToSlash(file) == legacySlashDefaultFile
+}
+
+func openAdminStore(ctx context.Context, config *nbconfig.Config, datadir string) (store.Store, error) {
+	managementStore, err := store.NewStore(ctx, config.StoreConfig.Engine, datadir, nil, true)
+	if err != nil {
+		return nil, fmt.Errorf("create store: %w", err)
+	}
+	return managementStore, nil
+}
+
+func openAdminEventStore(ctx context.Context, config *nbconfig.Config, datadir string) (activity.Store, error) {
+	if config.DataStoreEncryptionKey == "" {
+		return nil, fmt.Errorf("data store encryption key is not configured")
+	}
+	eventStore, err := activitystore.NewSqlStore(ctx, datadir, config.DataStoreEncryptionKey)
+	if err != nil {
+		return nil, fmt.Errorf("open activity event store: %w", err)
+	}
+	if eventStore == nil {
+		return nil, fmt.Errorf("open activity event store: returned nil store")
+	}
+	return eventStore, nil
+}

management/cmd/admin/admin.go

@@ -0,0 +1,577 @@
+// Package admincmd provides reusable cobra commands for self-hosted administrator helpers.
+// Both the management and combined binaries use these commands, each providing
+// their own opener to handle config loading and storage initialization.
+package admincmd
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"log/slog"
+	"os"
+	"strings"
+	"time"
+
+	"github.com/dexidp/dex/storage"
+	log "github.com/sirupsen/logrus"
+	"github.com/spf13/cobra"
+	"golang.org/x/crypto/bcrypt"
+
+	"github.com/netbirdio/netbird/formatter/hook"
+	nbdex "github.com/netbirdio/netbird/idp/dex"
+	"github.com/netbirdio/netbird/management/cmd/proxy"
+	"github.com/netbirdio/netbird/management/cmd/token"
+	nbconfig "github.com/netbirdio/netbird/management/internals/server/config"
+	"github.com/netbirdio/netbird/management/server"
+	"github.com/netbirdio/netbird/management/server/activity"
+	"github.com/netbirdio/netbird/management/server/idp"
+	"github.com/netbirdio/netbird/management/server/store"
+	"github.com/netbirdio/netbird/management/server/types"
+)
+
+// Resources contains the storages required by the admin commands.
+type Resources struct {
+	Store          store.Store
+	IDPStorage     storage.Storage
+	IDPStorageFile string
+	EventStore     activity.Store
+}
+
+// Opener initializes command resources from the command context and calls fn.
+type Opener func(cmd *cobra.Command, fn func(ctx context.Context, resources Resources) error) error
+
+// StoreOpener initializes only the management store from the command context and calls fn.
+type StoreOpener func(cmd *cobra.Command, fn func(ctx context.Context, s store.Store) error) error
+
+// IDPOpener initializes only the embedded IdP storage from the command context and calls fn.
+type IDPOpener func(cmd *cobra.Command, fn func(ctx context.Context, idpStorage storage.Storage, storageFile string) error) error
+
+// Openers contains the resource openers needed by the admin command tree.
+type Openers struct {
+	Resources Opener
+	Store     StoreOpener
+	IDP       IDPOpener
+}
+
+type userSelector struct {
+	email  string
+	userID string
+}
+
+func (s userSelector) normalized() userSelector {
+	return userSelector{
+		email:  strings.TrimSpace(s.email),
+		userID: strings.TrimSpace(s.userID),
+	}
+}
+
+func (s userSelector) validate() error {
+	s = s.normalized()
+	if (s.email == "") == (s.userID == "") {
+		return fmt.Errorf("provide exactly one of --email or --user-id")
+	}
+	return nil
+}
+
+// NewCommands creates the admin command tree with the given resource openers.
+func NewCommands(openers Openers) *cobra.Command {
+	adminCmd := &cobra.Command{
+		Use:   "admin",
+		Short: "Self-hosted administrator helpers",
+		Long:  "Administrative helpers for self-hosted deployments using the embedded identity provider.",
+	}
+
+	userCmd := &cobra.Command{
+		Use:   "user",
+		Short: "Manage local embedded IdP users",
+	}
+
+	var passwordSelector userSelector
+	var password string
+	var passwordFile string
+	passwordCmd := &cobra.Command{
+		Use:     "change-password (--email email | --user-id id) (--password password | --password-file path)",
+		Aliases: []string{"set-password"},
+		Short:   "Change a local user's password",
+		Args:    cobra.NoArgs,
+		RunE: func(cmd *cobra.Command, _ []string) error {
+			if err := passwordSelector.validate(); err != nil {
+				return err
+			}
+			newPassword, err := resolvePasswordInput(cmd, password, passwordFile)
+			if err != nil {
+				return err
+			}
+			return openers.IDP(cmd, func(ctx context.Context, idpStorage storage.Storage, storageFile string) error {
+				return runChangePassword(ctx, idpStorage, cmd.OutOrStdout(), passwordSelector, newPassword, storageFile)
+			})
+		},
+	}
+	addUserSelectorFlags(passwordCmd, &passwordSelector)
+	passwordCmd.Flags().StringVar(&password, "password", "", "New password for the user")
+	passwordCmd.Flags().StringVar(&passwordFile, "password-file", "", "Read new password from file ('-' for stdin)")
+
+	var resetSelector userSelector
+	resetMFACmd := &cobra.Command{
+		Use:   "reset-mfa (--email email | --user-id id)",
+		Short: "Reset a local user's MFA enrollment",
+		Args:  cobra.NoArgs,
+		RunE: func(cmd *cobra.Command, _ []string) error {
+			if err := resetSelector.validate(); err != nil {
+				return err
+			}
+			return openers.IDP(cmd, func(ctx context.Context, idpStorage storage.Storage, storageFile string) error {
+				return runResetMFA(ctx, idpStorage, cmd.OutOrStdout(), resetSelector, storageFile)
+			})
+		},
+	}
+	addUserSelectorFlags(resetMFACmd, &resetSelector)
+
+	userCmd.AddCommand(passwordCmd, resetMFACmd)
+
+	mfaCmd := &cobra.Command{
+		Use:   "mfa",
+		Short: "Manage local MFA for embedded IdP users",
+	}
+
+	enableCmd := &cobra.Command{
+		Use:   "enable",
+		Short: "Enable MFA for local embedded IdP users",
+		Args:  cobra.NoArgs,
+		RunE: func(cmd *cobra.Command, _ []string) error {
+			return openers.Resources(cmd, func(ctx context.Context, resources Resources) error {
+				return runSetMFAEnabled(ctx, resources, cmd.OutOrStdout(), true)
+			})
+		},
+	}
+
+	disableCmd := &cobra.Command{
+		Use:   "disable",
+		Short: "Disable MFA for local embedded IdP users",
+		Args:  cobra.NoArgs,
+		RunE: func(cmd *cobra.Command, _ []string) error {
+			return openers.Resources(cmd, func(ctx context.Context, resources Resources) error {
+				return runSetMFAEnabled(ctx, resources, cmd.OutOrStdout(), false)
+			})
+		},
+	}
+
+	statusCmd := &cobra.Command{
+		Use:   "status",
+		Short: "Show local MFA status",
+		Args:  cobra.NoArgs,
+		RunE: func(cmd *cobra.Command, _ []string) error {
+			return openers.Resources(cmd, func(ctx context.Context, resources Resources) error {
+				return runMFAStatus(ctx, resources, cmd.OutOrStdout())
+			})
+		},
+	}
+
+	mfaCmd.AddCommand(enableCmd, disableCmd, statusCmd)
+	adminCmd.AddCommand(userCmd, mfaCmd)
+	if openers.Store != nil {
+		adminCmd.AddCommand(tokencmd.NewCommands(tokencmd.StoreOpener(openers.Store)))
+		adminCmd.AddCommand(proxycmd.NewCommands(proxycmd.StoreOpener(openers.Store)))
+	}
+	return adminCmd
+}
+
+// OpenEmbeddedIDPStorage opens the Dex storage configured for the embedded IdP.
+func OpenEmbeddedIDPStorage(cfg *idp.EmbeddedIdPConfig) (storage.Storage, error) {
+	if cfg == nil || !cfg.Enabled {
+		return nil, fmt.Errorf("admin commands require the embedded IdP to be enabled")
+	}
+
+	yamlConfig, err := cfg.ToYAMLConfig()
+	if err != nil {
+		return nil, fmt.Errorf("build embedded IdP config: %w", err)
+	}
+
+	logger := slog.New(slog.NewTextHandler(io.Discard, nil))
+	st, err := yamlConfig.Storage.OpenStorage(logger)
+	if err != nil {
+		return nil, fmt.Errorf("open embedded IdP storage: %w", err)
+	}
+	return st, nil
+}
+
+// CloseStore closes the management store and logs cleanup errors at debug level.
+func CloseStore(ctx context.Context, s store.Store) {
+	if s == nil {
+		return
+	}
+	if err := s.Close(ctx); err != nil {
+		log.Debugf("close store: %v", err)
+	}
+}
+
+// OpenIDPStorage opens embedded IdP storage and returns its sqlite file path when applicable.
+func OpenIDPStorage(config *nbconfig.Config) (storage.Storage, string, error) {
+	if config == nil {
+		return nil, "", fmt.Errorf("management config is required")
+	}
+	idpStorage, err := OpenEmbeddedIDPStorage(config.EmbeddedIdP)
+	if err != nil {
+		return nil, "", err
+	}
+	return idpStorage, embeddedIDPStorageFile(config), nil
+}
+
+func embeddedIDPStorageFile(config *nbconfig.Config) string {
+	if config.EmbeddedIdP == nil || config.EmbeddedIdP.Storage.Type != "sqlite3" {
+		return ""
+	}
+	return config.EmbeddedIdP.Storage.Config.File
+}
+
+// CloseIDPStorage closes embedded IdP storage and logs cleanup errors at debug level.
+func CloseIDPStorage(s storage.Storage) {
+	if s == nil {
+		return
+	}
+	if err := s.Close(); err != nil {
+		log.Debugf("close embedded IdP storage: %v", err)
+	}
+}
+
+func addUserSelectorFlags(cmd *cobra.Command, selector *userSelector) {
+	cmd.Flags().StringVar(&selector.email, "email", "", "User email")
+	cmd.Flags().StringVar(&selector.userID, "user-id", "", "User ID")
+}
+
+func resolvePasswordInput(cmd *cobra.Command, password, passwordFile string) (string, error) {
+	if password != "" && passwordFile != "" {
+		return "", fmt.Errorf("provide only one of --password or --password-file")
+	}
+	if passwordFile == "" {
+		return password, nil
+	}
+
+	var data []byte
+	var err error
+	if passwordFile == "-" {
+		data, err = io.ReadAll(cmd.InOrStdin())
+	} else {
+		data, err = os.ReadFile(passwordFile)
+	}
+	if err != nil {
+		return "", fmt.Errorf("read password: %w", err)
+	}
+	return strings.TrimRight(string(data), "\r\n"), nil
+}
+
+func runChangePassword(ctx context.Context, idpStorage storage.Storage, w io.Writer, selector userSelector, password string, idpStorageFile string) error {
+	if idpStorage == nil {
+		return fmt.Errorf("embedded IdP storage is required")
+	}
+	selector = selector.normalized()
+	if err := selector.validate(); err != nil {
+		return err
+	}
+	if password == "" {
+		return fmt.Errorf("password is required")
+	}
+	if err := server.ValidatePassword(password); err != nil {
+		return fmt.Errorf("invalid password: %w", err)
+	}
+
+	user, err := findLocalUser(ctx, idpStorage, selector, idpStorageFile)
+	if err != nil {
+		return err
+	}
+
+	hash, err := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
+	if err != nil {
+		return fmt.Errorf("hash password: %w", err)
+	}
+
+	if err := idpStorage.UpdatePassword(ctx, user.Email, func(old storage.Password) (storage.Password, error) {
+		old.Hash = hash
+		return old, nil
+	}); err != nil {
+		return fmt.Errorf("update password for %s: %w", user.Email, err)
+	}
+
+	if err := deleteLocalAuthSession(ctx, idpStorage, user.UserID); err != nil {
+		return err
+	}
+
+	_, _ = fmt.Fprintf(w, "Password updated for %s.\n", user.Email)
+	return nil
+}
+
+func runResetMFA(ctx context.Context, idpStorage storage.Storage, w io.Writer, selector userSelector, idpStorageFile string) error {
+	if idpStorage == nil {
+		return fmt.Errorf("embedded IdP storage is required")
+	}
+	selector = selector.normalized()
+	if err := selector.validate(); err != nil {
+		return err
+	}
+
+	user, err := findLocalUser(ctx, idpStorage, selector, idpStorageFile)
+	if err != nil {
+		return err
+	}
+
+	reset := false
+	err = idpStorage.UpdateUserIdentity(ctx, user.UserID, idp.LocalConnectorID, func(old storage.UserIdentity) (storage.UserIdentity, error) {
+		reset = reset || len(old.MFASecrets) > 0 || len(old.WebAuthnCredentials) > 0
+		old.MFASecrets = map[string]*storage.MFASecret{}
+		old.WebAuthnCredentials = map[string][]storage.WebAuthnCredential{}
+		return old, nil
+	})
+	if errors.Is(err, storage.ErrNotFound) {
+		if err := deleteLocalAuthSession(ctx, idpStorage, user.UserID); err != nil {
+			return err
+		}
+		_, _ = fmt.Fprintf(w, "No MFA enrollment found for %s.\n", user.Email)
+		return nil
+	}
+	if err != nil {
+		return fmt.Errorf("reset MFA for %s: %w", user.Email, err)
+	}
+
+	if err := deleteLocalAuthSession(ctx, idpStorage, user.UserID); err != nil {
+		return err
+	}
+
+	if reset {
+		_, _ = fmt.Fprintf(w, "MFA reset for %s. The user will re-enroll at next login.\n", user.Email)
+	} else {
+		_, _ = fmt.Fprintf(w, "No MFA enrollment found for %s.\n", user.Email)
+	}
+	return nil
+}
+
+func runSetMFAEnabled(ctx context.Context, resources Resources, w io.Writer, enabled bool) error {
+	if resources.Store == nil {
+		return fmt.Errorf("management store is required")
+	}
+	if resources.IDPStorage == nil {
+		return fmt.Errorf("embedded IdP storage is required")
+	}
+
+	accountID, settings, err := getSingleAccountSettings(ctx, resources.Store)
+	if err != nil {
+		return err
+	}
+
+	oldEnabled := settings.LocalMfaEnabled
+	newSettings := settings.Copy()
+	newSettings.LocalMfaEnabled = enabled
+
+	if err := setIDPClientsMFA(ctx, resources.IDPStorage, enabled); err != nil {
+		return err
+	}
+
+	if err := resources.Store.SaveAccountSettings(ctx, accountID, newSettings); err != nil {
+		if rollbackErr := setIDPClientsMFA(ctx, resources.IDPStorage, oldEnabled); rollbackErr != nil {
+			return fmt.Errorf("save local MFA account setting: %w (also failed to roll back embedded IdP MFA state: %v)", err, rollbackErr)
+		}
+		return fmt.Errorf("save local MFA account setting: %w", err)
+	}
+
+	if err := storeMFAActivity(ctx, resources.EventStore, accountID, enabled); err != nil {
+		_, _ = fmt.Fprintf(w, "Warning: failed to record audit event: %v\n", err)
+	}
+
+	state := "disabled"
+	if enabled {
+		state = "enabled"
+	}
+	_, _ = fmt.Fprintf(w, "Local MFA %s.\n", state)
+	return nil
+}
+
+func runMFAStatus(ctx context.Context, resources Resources, w io.Writer) error {
+	if resources.Store == nil {
+		return fmt.Errorf("management store is required")
+	}
+	if resources.IDPStorage == nil {
+		return fmt.Errorf("embedded IdP storage is required")
+	}
+
+	_, settings, err := getSingleAccountSettings(ctx, resources.Store)
+	if err != nil {
+		return err
+	}
+	accountStatus := "disabled"
+	if settings.LocalMfaEnabled {
+		accountStatus = "enabled"
+	}
+
+	clientStatus, err := idpClientsMFAStatus(ctx, resources.IDPStorage)
+	if err != nil {
+		return err
+	}
+
+	_, _ = fmt.Fprintf(w, "Account setting: %s\n", accountStatus)
+	_, _ = fmt.Fprintf(w, "Embedded IdP clients: %s\n", clientStatus)
+	return nil
+}
+
+func getSingleAccountSettings(ctx context.Context, s store.Store) (string, *types.Settings, error) {
+	count, err := s.GetAccountsCounter(ctx)
+	if err != nil {
+		return "", nil, fmt.Errorf("count accounts: %w", err)
+	}
+	if count != 1 {
+		return "", nil, fmt.Errorf("expected exactly one account, got %d; local MFA is supported only in single-account embedded IdP deployments", count)
+	}
+
+	accountID, err := s.GetAnyAccountID(ctx)
+	if err != nil {
+		return "", nil, fmt.Errorf("get account ID: %w", err)
+	}
+
+	settings, err := s.GetAccountSettings(ctx, store.LockingStrengthNone, accountID)
+	if err != nil {
+		return "", nil, fmt.Errorf("get account settings: %w", err)
+	}
+	if settings == nil {
+		settings = &types.Settings{}
+	}
+	return accountID, settings, nil
+}
+
+func storeMFAActivity(ctx context.Context, eventStore activity.Store, accountID string, enabled bool) error {
+	if eventStore == nil {
+		return nil
+	}
+	event := activity.AccountLocalMfaDisabled
+	if enabled {
+		event = activity.AccountLocalMfaEnabled
+	}
+	_, err := eventStore.Save(ctx, &activity.Event{
+		Timestamp:   time.Now().UTC(),
+		Activity:    event,
+		InitiatorID: string(hook.SystemSource),
+		TargetID:    accountID,
+		AccountID:   accountID,
+	})
+	if err != nil {
+		return fmt.Errorf("save local MFA audit event: %w", err)
+	}
+	return nil
+}
+
+func findLocalUser(ctx context.Context, idpStorage storage.Storage, selector userSelector, idpStorageFile string) (storage.Password, error) {
+	selector = selector.normalized()
+	if err := selector.validate(); err != nil {
+		return storage.Password{}, err
+	}
+
+	if selector.email != "" {
+		user, err := idpStorage.GetPassword(ctx, selector.email)
+		if errors.Is(err, storage.ErrNotFound) {
+			if empty, listErr := localUsersEmpty(ctx, idpStorage); listErr != nil {
+				return storage.Password{}, listErr
+			} else if empty {
+				return storage.Password{}, noLocalUsersError(idpStorageFile)
+			}
+			return storage.Password{}, fmt.Errorf("local user with email %q not found", selector.email)
+		}
+		if err != nil {
+			return storage.Password{}, fmt.Errorf("get local user by email %q: %w", selector.email, err)
+		}
+		return user, nil
+	}
+
+	rawUserID := selector.userID
+	if decodedUserID, _, err := nbdex.DecodeDexUserID(selector.userID); err == nil && decodedUserID != "" {
+		rawUserID = decodedUserID
+	}
+
+	users, err := idpStorage.ListPasswords(ctx)
+	if err != nil {
+		return storage.Password{}, fmt.Errorf("list local users: %w", err)
+	}
+	for _, user := range users {
+		if user.UserID == rawUserID || user.UserID == selector.userID {
+			return user, nil
+		}
+	}
+
+	if len(users) == 0 {
+		return storage.Password{}, noLocalUsersError(idpStorageFile)
+	}
+
+	return storage.Password{}, fmt.Errorf("local user with ID %q not found", selector.userID)
+}
+
+func localUsersEmpty(ctx context.Context, idpStorage storage.Storage) (bool, error) {
+	users, err := idpStorage.ListPasswords(ctx)
+	if err != nil {
+		return false, fmt.Errorf("list local users: %w", err)
+	}
+	return len(users) == 0, nil
+}
+
+func noLocalUsersError(idpStorageFile string) error {
+	location := ""
+	if idpStorageFile != "" {
+		location = fmt.Sprintf(" (%s)", idpStorageFile)
+	}
+	return fmt.Errorf("no local users exist in the embedded IdP storage%s; the management server may never have started with this config, or --datadir points at the wrong location", location)
+}
+
+func deleteLocalAuthSession(ctx context.Context, idpStorage storage.Storage, userID string) error {
+	err := idpStorage.DeleteAuthSession(ctx, userID, idp.LocalConnectorID)
+	if err == nil || errors.Is(err, storage.ErrNotFound) {
+		return nil
+	}
+	return fmt.Errorf("delete local auth session for user %s: %w", userID, err)
+}
+
+func setIDPClientsMFA(ctx context.Context, idpStorage storage.Storage, enabled bool) error {
+	var mfaChain []string
+	if enabled {
+		mfaChain = []string{idp.DefaultTOTPAuthenticatorID}
+	}
+
+	clientIDs := []string{idp.StaticClientCLI, idp.StaticClientDashboard}
+	if err := nbdex.SetClientsMFAChain(ctx, idpStorage, clientIDs, mfaChain); err != nil {
+		if errors.Is(err, storage.ErrNotFound) {
+			return fmt.Errorf("embedded IdP client not found; start the management server once before toggling MFA: %w", err)
+		}
+		return fmt.Errorf("update MFA chain on embedded IdP clients: %w", err)
+	}
+	return nil
+}
+
+func idpClientsMFAStatus(ctx context.Context, idpStorage storage.Storage) (string, error) {
+	clientIDs := []string{idp.StaticClientCLI, idp.StaticClientDashboard}
+	enabledCount := 0
+	for _, clientID := range clientIDs {
+		client, err := idpStorage.GetClient(ctx, clientID)
+		if errors.Is(err, storage.ErrNotFound) {
+			return "unknown", fmt.Errorf("embedded IdP client %q not found", clientID)
+		}
+		if err != nil {
+			return "unknown", fmt.Errorf("get embedded IdP client %q: %w", clientID, err)
+		}
+		if hasAuthenticator(client.MFAChain, idp.DefaultTOTPAuthenticatorID) {
+			enabledCount++
+		}
+	}
+
+	switch enabledCount {
+	case 0:
+		return "disabled", nil
+	case len(clientIDs):
+		return "enabled", nil
+	default:
+		return "partially enabled", nil
+	}
+}
+
+func hasAuthenticator(chain []string, authenticatorID string) bool {
+	for _, id := range chain {
+		if id == authenticatorID {
+			return true
+		}
+	}
+	return false
+}

management/cmd/admin/admin_test.go

@@ -0,0 +1,250 @@
+package admincmd
+
+import (
+	"bytes"
+	"context"
+	"io"
+	"log/slog"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/dexidp/dex/storage"
+	"github.com/dexidp/dex/storage/memory"
+	"github.com/spf13/cobra"
+	"github.com/stretchr/testify/require"
+	"golang.org/x/crypto/bcrypt"
+
+	nbdex "github.com/netbirdio/netbird/idp/dex"
+	"github.com/netbirdio/netbird/management/server/idp"
+	mgmtstore "github.com/netbirdio/netbird/management/server/store"
+	"github.com/netbirdio/netbird/management/server/types"
+)
+
+func newTestIDPStorage(t *testing.T) storage.Storage {
+	t.Helper()
+
+	st := memory.New(slog.New(slog.NewTextHandler(io.Discard, nil)))
+	hash, err := bcrypt.GenerateFromPassword([]byte("OldPass1!"), bcrypt.DefaultCost)
+	require.NoError(t, err)
+
+	require.NoError(t, st.CreatePassword(context.Background(), storage.Password{
+		Email:    "user@example.com",
+		Username: "User",
+		UserID:   "user-1",
+		Hash:     hash,
+	}))
+	require.NoError(t, st.CreateUserIdentity(context.Background(), storage.UserIdentity{
+		UserID:      "user-1",
+		ConnectorID: idp.LocalConnectorID,
+		MFASecrets: map[string]*storage.MFASecret{
+			idp.DefaultTOTPAuthenticatorID: {
+				AuthenticatorID: idp.DefaultTOTPAuthenticatorID,
+				Type:            "TOTP",
+				Secret:          "otpauth://totp/NetBird:user@example.com?secret=ABC",
+				Confirmed:       true,
+				CreatedAt:       time.Now(),
+			},
+		},
+		WebAuthnCredentials: map[string][]storage.WebAuthnCredential{
+			"webauthn": {{CredentialID: []byte("credential")}},
+		},
+	}))
+	require.NoError(t, st.CreateAuthSession(context.Background(), storage.AuthSession{
+		UserID:      "user-1",
+		ConnectorID: idp.LocalConnectorID,
+		Nonce:       "nonce",
+	}))
+	require.NoError(t, st.CreateClient(context.Background(), storage.Client{ID: idp.StaticClientCLI, Name: "CLI"}))
+	require.NoError(t, st.CreateClient(context.Background(), storage.Client{ID: idp.StaticClientDashboard, Name: "Dashboard"}))
+
+	return st
+}
+
+func TestRunChangePassword(t *testing.T) {
+	ctx := context.Background()
+	st := newTestIDPStorage(t)
+	var out bytes.Buffer
+
+	err := runChangePassword(ctx, st, &out, userSelector{email: "user@example.com"}, "NewPass1!", "")
+	require.NoError(t, err)
+	require.Contains(t, out.String(), "Password updated")
+
+	user, err := st.GetPassword(ctx, "user@example.com")
+	require.NoError(t, err)
+	require.NoError(t, bcrypt.CompareHashAndPassword(user.Hash, []byte("NewPass1!")))
+
+	_, err = st.GetAuthSession(ctx, "user-1", idp.LocalConnectorID)
+	require.ErrorIs(t, err, storage.ErrNotFound)
+}
+
+func TestRunChangePasswordValidatesPassword(t *testing.T) {
+	st := newTestIDPStorage(t)
+	err := runChangePassword(context.Background(), st, io.Discard, userSelector{email: "user@example.com"}, "short", "")
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "invalid password")
+}
+
+func TestRunResetMFA(t *testing.T) {
+	ctx := context.Background()
+	st := newTestIDPStorage(t)
+	var out bytes.Buffer
+
+	encodedUserID := nbdex.EncodeDexUserID("user-1", idp.LocalConnectorID)
+	err := runResetMFA(ctx, st, &out, userSelector{userID: encodedUserID}, "")
+	require.NoError(t, err)
+	require.Contains(t, out.String(), "MFA reset")
+
+	identity, err := st.GetUserIdentity(ctx, "user-1", idp.LocalConnectorID)
+	require.NoError(t, err)
+	require.Empty(t, identity.MFASecrets)
+	require.Empty(t, identity.WebAuthnCredentials)
+
+	_, err = st.GetAuthSession(ctx, "user-1", idp.LocalConnectorID)
+	require.ErrorIs(t, err, storage.ErrNotFound)
+}
+
+func TestRunResetMFAWithoutEnrollment(t *testing.T) {
+	ctx := context.Background()
+	st := newTestIDPStorage(t)
+	require.NoError(t, st.UpdateUserIdentity(ctx, "user-1", idp.LocalConnectorID, func(old storage.UserIdentity) (storage.UserIdentity, error) {
+		old.MFASecrets = nil
+		old.WebAuthnCredentials = nil
+		return old, nil
+	}))
+
+	var out bytes.Buffer
+	err := runResetMFA(ctx, st, &out, userSelector{email: "user@example.com"}, "")
+	require.NoError(t, err)
+	require.Contains(t, out.String(), "No MFA enrollment found")
+}
+
+func TestSetIDPClientsMFA(t *testing.T) {
+	ctx := context.Background()
+	st := newTestIDPStorage(t)
+
+	require.NoError(t, setIDPClientsMFA(ctx, st, true))
+	status, err := idpClientsMFAStatus(ctx, st)
+	require.NoError(t, err)
+	require.Equal(t, "enabled", status)
+
+	require.NoError(t, setIDPClientsMFA(ctx, st, false))
+	status, err = idpClientsMFAStatus(ctx, st)
+	require.NoError(t, err)
+	require.Equal(t, "disabled", status)
+}
+
+func newTestManagementStore(t *testing.T, localMFAEnabled bool) mgmtstore.Store {
+	t.Helper()
+	ctx := context.Background()
+	st, err := mgmtstore.NewStore(ctx, types.SqliteStoreEngine, t.TempDir(), nil, false)
+	require.NoError(t, err)
+	t.Cleanup(func() { require.NoError(t, st.Close(ctx)) })
+	require.NoError(t, st.SaveAccount(ctx, &types.Account{
+		Id:       "account-1",
+		Settings: &types.Settings{LocalMfaEnabled: localMFAEnabled},
+	}))
+	return st
+}
+
+func TestRunSetMFAEnabledDoesNotSaveWhenIDPUpdateFails(t *testing.T) {
+	ctx := context.Background()
+	managementStore := newTestManagementStore(t, false)
+	idpStorage := memory.New(slog.New(slog.NewTextHandler(io.Discard, nil)))
+
+	err := runSetMFAEnabled(ctx, Resources{Store: managementStore, IDPStorage: idpStorage}, io.Discard, true)
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "embedded IdP client")
+
+	settings, err := managementStore.GetAccountSettings(ctx, mgmtstore.LockingStrengthNone, "account-1")
+	require.NoError(t, err)
+	require.False(t, settings.LocalMfaEnabled)
+}
+
+func TestRunSetMFAEnabledUpdatesSettingsAfterIDP(t *testing.T) {
+	ctx := context.Background()
+	managementStore := newTestManagementStore(t, false)
+	idpStorage := newTestIDPStorage(t)
+
+	err := runSetMFAEnabled(ctx, Resources{Store: managementStore, IDPStorage: idpStorage}, io.Discard, true)
+	require.NoError(t, err)
+
+	settings, err := managementStore.GetAccountSettings(ctx, mgmtstore.LockingStrengthNone, "account-1")
+	require.NoError(t, err)
+	require.True(t, settings.LocalMfaEnabled)
+	clientStatus, err := idpClientsMFAStatus(ctx, idpStorage)
+	require.NoError(t, err)
+	require.Equal(t, "enabled", clientStatus)
+}
+
+func TestRunSetMFAEnabledSucceedsWithNilEventStore(t *testing.T) {
+	ctx := context.Background()
+	managementStore := newTestManagementStore(t, false)
+	idpStorage := newTestIDPStorage(t)
+	var out bytes.Buffer
+	var err error
+
+	require.NotPanics(t, func() {
+		err = runSetMFAEnabled(ctx, Resources{Store: managementStore, IDPStorage: idpStorage, EventStore: nil}, &out, true)
+	})
+	require.NoError(t, err)
+	require.Contains(t, out.String(), "Local MFA enabled")
+
+	settings, err := managementStore.GetAccountSettings(ctx, mgmtstore.LockingStrengthNone, "account-1")
+	require.NoError(t, err)
+	require.True(t, settings.LocalMfaEnabled)
+}
+
+func TestUserSelectorValidate(t *testing.T) {
+	require.NoError(t, userSelector{email: " user@example.com "}.validate())
+	require.NoError(t, userSelector{userID: "user-1"}.validate())
+	require.Error(t, userSelector{}.validate())
+	require.Error(t, userSelector{email: "user@example.com", userID: "user-1"}.validate())
+}
+
+func TestFindLocalUserNotFound(t *testing.T) {
+	st := newTestIDPStorage(t)
+	_, err := findLocalUser(context.Background(), st, userSelector{email: "missing@example.com"}, "")
+	require.Error(t, err)
+	require.True(t, strings.Contains(err.Error(), "not found"))
+}
+
+func TestFindLocalUserZeroUsersIncludesStoragePath(t *testing.T) {
+	st := memory.New(slog.New(slog.NewTextHandler(io.Discard, nil)))
+	_, err := findLocalUser(context.Background(), st, userSelector{email: "missing@example.com"}, "/var/lib/netbird/idp.db")
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "no local users exist")
+	require.Contains(t, err.Error(), "/var/lib/netbird/idp.db")
+}
+
+func TestUserCommandValidatesSelectorBeforeOpeningStorage(t *testing.T) {
+	opened := false
+	cmd := NewCommands(Openers{
+		IDP: func(cmd *cobra.Command, fn func(ctx context.Context, idpStorage storage.Storage, storageFile string) error) error {
+			opened = true
+			return nil
+		},
+	})
+	cmd.SetArgs([]string{"user", "change-password", "--password", "NewPass1!"})
+	cmd.SetOut(io.Discard)
+	cmd.SetErr(io.Discard)
+
+	err := cmd.Execute()
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "provide exactly one")
+	require.False(t, opened)
+}
+
+func TestResolvePasswordInputFromStdin(t *testing.T) {
+	cmd := &cobra.Command{}
+	cmd.SetIn(strings.NewReader("NewPass1!\n"))
+
+	password, err := resolvePasswordInput(cmd, "", "-")
+	require.NoError(t, err)
+	require.Equal(t, "NewPass1!", password)
+}
+
+func TestResolvePasswordInputRejectsMultipleSources(t *testing.T) {
+	_, err := resolvePasswordInput(&cobra.Command{}, "NewPass1!", "-")
+	require.Error(t, err)
+}

management/cmd/admin_config_test.go

@@ -0,0 +1,80 @@
+package cmd
+
+import (
+	"context"
+	"path"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	nbconfig "github.com/netbirdio/netbird/management/internals/server/config"
+	"github.com/netbirdio/netbird/management/server/idp"
+)
+
+func TestApplyAdminDatadirOverrideRelocatesDefaultIDPStorage(t *testing.T) {
+	oldDatadir := filepath.Join(t.TempDir(), "old")
+	newDatadir := filepath.Join(t.TempDir(), "new")
+
+	for _, defaultFile := range []string{
+		"",
+		filepath.Join(oldDatadir, "idp.db"),
+		path.Join(oldDatadir, "idp.db"),
+	} {
+		t.Run(defaultFile, func(t *testing.T) {
+			cfg := &nbconfig.Config{
+				EmbeddedIdP: &idp.EmbeddedIdPConfig{
+					Enabled: true,
+					Storage: idp.EmbeddedStorageConfig{
+						Type: "sqlite3",
+						Config: idp.EmbeddedStorageTypeConfig{
+							File: defaultFile,
+						},
+					},
+				},
+			}
+			datadir := oldDatadir
+			oldAdminDatadir := adminDatadir
+			adminDatadir = newDatadir
+			t.Cleanup(func() { adminDatadir = oldAdminDatadir })
+
+			applyAdminDatadirOverride(cfg, &datadir)
+
+			require.Equal(t, newDatadir, datadir)
+			require.Equal(t, filepath.Join(newDatadir, "idp.db"), cfg.EmbeddedIdP.Storage.Config.File)
+		})
+	}
+}
+
+func TestOpenAdminEventStoreMissingEncryptionKeyReturnsNilInterface(t *testing.T) {
+	eventStore, err := openAdminEventStore(context.Background(), &nbconfig.Config{}, t.TempDir())
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "encryption key")
+	require.Nil(t, eventStore)
+}
+
+func TestApplyAdminDatadirOverrideKeepsExplicitIDPStorage(t *testing.T) {
+	oldDatadir := filepath.Join(t.TempDir(), "old")
+	newDatadir := filepath.Join(t.TempDir(), "new")
+	explicitFile := filepath.Join(t.TempDir(), "custom-idp.db")
+	cfg := &nbconfig.Config{
+		EmbeddedIdP: &idp.EmbeddedIdPConfig{
+			Enabled: true,
+			Storage: idp.EmbeddedStorageConfig{
+				Type: "sqlite3",
+				Config: idp.EmbeddedStorageTypeConfig{
+					File: explicitFile,
+				},
+			},
+		},
+	}
+	datadir := oldDatadir
+	oldAdminDatadir := adminDatadir
+	adminDatadir = newDatadir
+	t.Cleanup(func() { adminDatadir = oldAdminDatadir })
+
+	applyAdminDatadirOverride(cfg, &datadir)
+
+	require.Equal(t, newDatadir, datadir)
+	require.Equal(t, explicitFile, cfg.EmbeddedIdP.Storage.Config.File)
+}

management/cmd/management.go

@@ -13,6 +13,7 @@ import (
 	"os"
 	"os/signal"
 	"path"
+	"path/filepath"
 	"strings"
 	"syscall"
 
@@ -25,6 +26,7 @@ import (
 	"github.com/netbirdio/netbird/management/internals/server"
 	nbconfig "github.com/netbirdio/netbird/management/internals/server/config"
 	nbdomain "github.com/netbirdio/netbird/shared/management/domain"
+	"github.com/netbirdio/netbird/shared/management/grpc"
 	"github.com/netbirdio/netbird/util"
 	"github.com/netbirdio/netbird/util/crypt"
 )
@@ -153,8 +155,20 @@ func LoadMgmtConfig(ctx context.Context, mgmtConfigPath string) (*nbconfig.Confi
 
 	ApplyCommandLineOverrides(loadedConfig)
 
+	err := grpc.ValidateSyncMessageVersion(loadedConfig.HighestSupportedSyncMessageVersion)
+	if err != nil {
+		return nil, err
+	}
+
+	for account, version := range loadedConfig.PerAccountHighestSupportedSyncMessageVersion {
+		err := grpc.ValidateSyncMessageVersion(&version)
+		if err != nil {
+			return nil, fmt.Errorf("unrecognized sync message version for account %s, %w", account, err)
+		}
+	}
+
 	// Apply EmbeddedIdP config to HttpConfig if embedded IdP is enabled
-	err := ApplyEmbeddedIdPConfig(ctx, loadedConfig)
+	err = ApplyEmbeddedIdPConfig(ctx, loadedConfig)
 	if err != nil {
 		return nil, err
 	}
@@ -209,7 +223,7 @@ func ApplyEmbeddedIdPConfig(ctx context.Context, cfg *nbconfig.Config) error {
 		cfg.EmbeddedIdP.Storage.Type = "sqlite3"
 	}
 	if cfg.EmbeddedIdP.Storage.Config.File == "" && cfg.Datadir != "" {
-		cfg.EmbeddedIdP.Storage.Config.File = path.Join(cfg.Datadir, "idp.db")
+		cfg.EmbeddedIdP.Storage.Config.File = filepath.Join(cfg.Datadir, "idp.db")
 	}
 
 	issuer := cfg.EmbeddedIdP.Issuer

management/cmd/management_test.go

@@ -4,6 +4,9 @@ import (
 	"context"
 	"os"
 	"testing"
+
+	"github.com/netbirdio/netbird/shared/management/grpc"
+	"github.com/stretchr/testify/assert"
 )
 
 const (
@@ -20,34 +23,49 @@ const (
 		"AuthAudience": "https://stageapp/",
 		"AuthIssuer": "https://something.eu.auth0.com/",
 		"OIDCConfigEndpoint": "https://something.eu.auth0.com/.well-known/openid-configuration"
+	  },
+	  "HighestSupportedSyncMessageVersion": 1,
+	  "PerAccountHighestSupportedSyncMessageVersion": {
+	    "1": 0,
+		"2": 1
 	  }
 	}`
 )
 
-func Test_loadMgmtConfig(t *testing.T) {
-	tmpFile, err := createConfig()
-	if err != nil {
-		t.Fatalf("failed to create config: %s", err)
-	}
+func Test_LoadMgmtConfig(t *testing.T) {
+	tmpFile, err := createConfig(exampleConfig)
+	assert.NoError(t, err)
 
 	cfg, err := LoadMgmtConfig(context.Background(), tmpFile)
-	if err != nil {
-		t.Fatalf("failed to load management config: %s", err)
-	}
-	if cfg.Relay == nil {
-		t.Fatalf("config is nil")
-	}
-	if len(cfg.Relay.Addresses) == 0 {
-		t.Fatalf("relay address is empty")
-	}
+	assert.NoError(t, err)
+	assert.NotEmpty(t, cfg.Relay)
+	assert.NotEmpty(t, cfg.Relay.Addresses)
+	assert.Equal(t, int(grpc.ComponentNetworkMap), *cfg.HighestSupportedSyncMessageVersion)
+	assert.Equal(t, map[string]int{"1": int(grpc.Base), "2": int(grpc.ComponentNetworkMap)}, cfg.PerAccountHighestSupportedSyncMessageVersion)
 }
 
-func createConfig() (string, error) {
+func Test_LoadMgmtConfig_Empty(t *testing.T) {
+	tmpFile, err := createConfig(`{
+		"HttpConfig": {
+			"AuthAudience": "https://stageapp/",
+			"AuthIssuer": "https://something.eu.auth0.com/",
+			"OIDCConfigEndpoint": "https://something.eu.auth0.com/.well-known/openid-configuration"
+		}
+	}`)
+	assert.NoError(t, err)
+
+	cfg, err := LoadMgmtConfig(context.Background(), tmpFile)
+	assert.NoError(t, err)
+	assert.Nil(t, cfg.HighestSupportedSyncMessageVersion)
+	assert.Nil(t, cfg.PerAccountHighestSupportedSyncMessageVersion)
+}
+
+func createConfig(config string) (string, error) {
 	tmpfile, err := os.CreateTemp("", "config.json")
 	if err != nil {
 		return "", err
 	}
-	_, err = tmpfile.Write([]byte(exampleConfig))
+	_, err = tmpfile.Write([]byte(config))
 	if err != nil {
 		return "", err
 	}

management/cmd/proxy/proxy.go

@@ -0,0 +1,141 @@
+// Package proxycmd provides reusable cobra commands for managing reverse proxy instances.
+// Both the management and combined binaries use these commands, each providing
+// their own StoreOpener to handle config loading and store initialization.
+package proxycmd
+
+import (
+	"bufio"
+	"context"
+	"fmt"
+	"io"
+	"strings"
+	"text/tabwriter"
+
+	"github.com/spf13/cobra"
+
+	rpproxy "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy"
+	"github.com/netbirdio/netbird/management/server/store"
+)
+
+// StoreOpener initializes a store from the command context and calls fn.
+type StoreOpener func(cmd *cobra.Command, fn func(ctx context.Context, s store.Store) error) error
+
+const disconnectAllConfirmation = "disconnect all proxies"
+
+// NewCommands creates the proxy command tree with the given store opener.
+// Returns the parent "proxy" command with the disconnect-all subcommand.
+func NewCommands(opener StoreOpener) *cobra.Command {
+	var dryRun bool
+	var force bool
+
+	proxyCmd := &cobra.Command{
+		Use:   "proxy",
+		Short: "Manage reverse proxy instances",
+		Long:  "Commands for inspecting and repairing the reverse proxy instances registered with the management server.",
+	}
+
+	disconnectAllCmd := &cobra.Command{
+		Use:   "disconnect-all",
+		Short: "Force-mark all reverse proxy instances as disconnected",
+		Long: "Lists all reverse proxy instances and force-marks them as disconnected, regardless of their session state. " +
+			"Use this to repair stale connection state, e.g. after an unclean management server shutdown. " +
+			"By default, it asks for manual confirmation before changing state. Use --dry-run to preview without changing state, or --force to skip confirmation. " +
+			"Run during a maintenance window; affected live proxies may stay hidden until their next heartbeat or reconnect/re-register.",
+		Args: cobra.NoArgs,
+		RunE: func(cmd *cobra.Command, _ []string) error {
+			return opener(cmd, func(ctx context.Context, s store.Store) error {
+				return runDisconnectAll(ctx, s, cmd.OutOrStdout(), cmd.InOrStdin(), dryRun, force)
+			})
+		},
+	}
+	disconnectAllCmd.Flags().BoolVar(&dryRun, "dry-run", false, "List reverse proxy instances that would be disconnected without changing state")
+	disconnectAllCmd.Flags().BoolVar(&force, "force", false, "Skip the confirmation prompt and apply the repair")
+
+	proxyCmd.AddCommand(disconnectAllCmd)
+	return proxyCmd
+}
+
+func runDisconnectAll(ctx context.Context, s store.Store, out io.Writer, in io.Reader, dryRun, force bool) error {
+	proxies, err := s.GetAllProxies(ctx)
+	if err != nil {
+		return fmt.Errorf("list proxies: %w", err)
+	}
+
+	if len(proxies) == 0 {
+		_, _ = fmt.Fprintln(out, "No reverse proxy instances found.")
+		return nil
+	}
+
+	toDisconnect := 0
+	w := tabwriter.NewWriter(out, 0, 0, 2, ' ', 0)
+	_, _ = fmt.Fprintln(w, "ID\tCLUSTER\tIP\tACCOUNT\tSTATUS\tLAST SEEN")
+	_, _ = fmt.Fprintln(w, "--\t-------\t--\t-------\t------\t---------")
+
+	for _, p := range proxies {
+		if p.Status != rpproxy.StatusDisconnected {
+			toDisconnect++
+		}
+
+		account := "-"
+		if p.AccountID != nil {
+			account = *p.AccountID
+		}
+
+		_, _ = fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\t%s\n",
+			p.ID,
+			p.ClusterAddress,
+			p.IPAddress,
+			account,
+			p.Status,
+			p.LastSeen.Format("2006-01-02 15:04:05"),
+		)
+	}
+	if err := w.Flush(); err != nil {
+		return fmt.Errorf("write proxy list: %w", err)
+	}
+
+	if dryRun {
+		_, _ = fmt.Fprintf(out, "\nDry run: would force-mark %d of %d reverse proxy instance(s) as disconnected.\n", toDisconnect, len(proxies))
+		return nil
+	}
+
+	if !force {
+		confirmed, err := confirmDisconnectAll(out, in)
+		if err != nil {
+			return err
+		}
+		if !confirmed {
+			_, _ = fmt.Fprintln(out, "Aborted. No reverse proxy instances were changed.")
+			return nil
+		}
+	}
+
+	disconnected, err := s.DisconnectAllProxies(ctx)
+	if err != nil {
+		return fmt.Errorf("disconnect proxies: %w", err)
+	}
+
+	_, _ = fmt.Fprintf(out, "\nForce-marked %d of %d reverse proxy instance(s) as disconnected.\n", disconnected, len(proxies))
+	return nil
+}
+
+func confirmDisconnectAll(out io.Writer, in io.Reader) (bool, error) {
+	if in == nil {
+		in = strings.NewReader("")
+	}
+
+	_, _ = fmt.Fprintln(out, "\nWARNING: This command changes stored reverse proxy state for every non-disconnected instance.")
+	_, _ = fmt.Fprintln(out, "Run it during a maintenance window; affected live proxies may stay hidden until "+
+		"their next heartbeat or reconnect/re-register.")
+	_, _ = fmt.Fprintf(out, "Type %q to continue: ", disconnectAllConfirmation)
+
+	scanner := bufio.NewScanner(in)
+	if !scanner.Scan() {
+		if err := scanner.Err(); err != nil {
+			return false, fmt.Errorf("read confirmation: %w", err)
+		}
+		return false, nil
+	}
+
+	return strings.EqualFold(strings.TrimSpace(scanner.Text()), disconnectAllConfirmation), nil
+}

management/cmd/proxy/proxy_test.go

@@ -0,0 +1,180 @@
+package proxycmd
+
+import (
+	"bytes"
+	"context"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/spf13/cobra"
+	"github.com/stretchr/testify/require"
+
+	rpproxy "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy"
+	"github.com/netbirdio/netbird/management/server/store"
+)
+
+func newTestStore(t *testing.T) store.Store {
+	t.Helper()
+
+	s, cleanup, err := store.NewTestStoreFromSQL(context.Background(), "", t.TempDir())
+	require.NoError(t, err)
+	t.Cleanup(cleanup)
+
+	return s
+}
+
+func seedProxies(t *testing.T, ctx context.Context, s store.Store) {
+	t.Helper()
+
+	accountID := "account-1"
+	alreadyDisconnectedAt := time.Now().Add(-time.Hour)
+	seed := []*rpproxy.Proxy{
+		{
+			ID:             "proxy-1",
+			SessionID:      "session-1",
+			ClusterAddress: "cluster-a.example.com",
+			IPAddress:      "10.0.0.1",
+			LastSeen:       time.Now(),
+			Status:         rpproxy.StatusConnected,
+		},
+		{
+			ID:             "proxy-2",
+			SessionID:      "session-2",
+			ClusterAddress: "cluster-b.example.com",
+			IPAddress:      "10.0.0.2",
+			AccountID:      &accountID,
+			LastSeen:       time.Now(),
+			Status:         rpproxy.StatusConnected,
+		},
+		{
+			ID:             "proxy-3",
+			SessionID:      "session-3",
+			ClusterAddress: "cluster-a.example.com",
+			IPAddress:      "10.0.0.3",
+			LastSeen:       time.Now().Add(-time.Hour),
+			Status:         rpproxy.StatusDisconnected,
+			DisconnectedAt: &alreadyDisconnectedAt,
+		},
+	}
+	for _, p := range seed {
+		require.NoError(t, s.SaveProxy(ctx, p))
+	}
+}
+
+func proxiesByID(t *testing.T, ctx context.Context, s store.Store) map[string]*rpproxy.Proxy {
+	t.Helper()
+
+	proxies, err := s.GetAllProxies(ctx)
+	require.NoError(t, err)
+	require.Len(t, proxies, 3)
+
+	byID := make(map[string]*rpproxy.Proxy, len(proxies))
+	for _, p := range proxies {
+		byID[p.ID] = p
+	}
+	return byID
+}
+
+func TestRunDisconnectAllWithConfirmation(t *testing.T) {
+	ctx := context.Background()
+	s := newTestStore(t)
+	seedProxies(t, ctx, s)
+
+	var out bytes.Buffer
+	require.NoError(t, runDisconnectAll(ctx, s, &out, strings.NewReader(disconnectAllConfirmation+"\n"), false, false))
+
+	output := out.String()
+	require.Contains(t, output, "proxy-1")
+	require.Contains(t, output, "proxy-2")
+	require.Contains(t, output, "proxy-3")
+	require.Contains(t, output, "cluster-a.example.com")
+	require.Contains(t, output, "account-1")
+	require.Contains(t, output, "Type \"disconnect all proxies\" to continue")
+	require.Contains(t, output, "Force-marked 2 of 3 reverse proxy instance(s) as disconnected.")
+
+	for _, p := range proxiesByID(t, ctx, s) {
+		require.Equal(t, rpproxy.StatusDisconnected, p.Status, "proxy %s should be disconnected", p.ID)
+		require.NotNil(t, p.DisconnectedAt, "proxy %s should have a disconnected timestamp", p.ID)
+	}
+}
+
+func TestRunDisconnectAllForceSkipsConfirmation(t *testing.T) {
+	ctx := context.Background()
+	s := newTestStore(t)
+	seedProxies(t, ctx, s)
+
+	var out bytes.Buffer
+	require.NoError(t, runDisconnectAll(ctx, s, &out, strings.NewReader(""), false, true))
+
+	output := out.String()
+	require.NotContains(t, output, "Type \"disconnect all proxies\" to continue")
+	require.Contains(t, output, "Force-marked 2 of 3 reverse proxy instance(s) as disconnected.")
+}
+
+func TestRunDisconnectAllAbortLeavesProxiesUnchanged(t *testing.T) {
+	ctx := context.Background()
+	s := newTestStore(t)
+	seedProxies(t, ctx, s)
+
+	var out bytes.Buffer
+	require.NoError(t, runDisconnectAll(ctx, s, &out, strings.NewReader("no\n"), false, false))
+
+	output := out.String()
+	require.Contains(t, output, "Type \"disconnect all proxies\" to continue")
+	require.Contains(t, output, "Aborted. No reverse proxy instances were changed.")
+
+	byID := proxiesByID(t, ctx, s)
+	require.Equal(t, rpproxy.StatusConnected, byID["proxy-1"].Status)
+	require.Equal(t, rpproxy.StatusConnected, byID["proxy-2"].Status)
+	require.Equal(t, rpproxy.StatusDisconnected, byID["proxy-3"].Status)
+}
+
+func TestRunDisconnectAllDryRunLeavesProxiesUnchanged(t *testing.T) {
+	ctx := context.Background()
+	s := newTestStore(t)
+	seedProxies(t, ctx, s)
+
+	var out bytes.Buffer
+	require.NoError(t, runDisconnectAll(ctx, s, &out, strings.NewReader(""), true, false))
+
+	output := out.String()
+	require.Contains(t, output, "Dry run: would force-mark 2 of 3 reverse proxy instance(s) as disconnected.")
+	require.NotContains(t, output, "Type \"disconnect all proxies\" to continue")
+
+	byID := proxiesByID(t, ctx, s)
+	require.Equal(t, rpproxy.StatusConnected, byID["proxy-1"].Status)
+	require.Equal(t, rpproxy.StatusConnected, byID["proxy-2"].Status)
+	require.Equal(t, rpproxy.StatusDisconnected, byID["proxy-3"].Status)
+}
+
+func TestNewCommandsDisconnectAllDryRun(t *testing.T) {
+	ctx := context.Background()
+	s := newTestStore(t)
+	seedProxies(t, ctx, s)
+
+	opened := false
+	cmd := NewCommands(func(cmd *cobra.Command, fn func(ctx context.Context, s store.Store) error) error {
+		opened = true
+		return fn(cmd.Context(), s)
+	})
+
+	var out bytes.Buffer
+	cmd.SetOut(&out)
+	cmd.SetErr(&out)
+	cmd.SetIn(strings.NewReader(""))
+	cmd.SetArgs([]string{"disconnect-all", "--dry-run"})
+
+	require.NoError(t, cmd.ExecuteContext(ctx))
+	require.True(t, opened)
+	require.Contains(t, out.String(), "Dry run: would force-mark 2 of 3 reverse proxy instance(s) as disconnected.")
+}
+
+func TestRunDisconnectAllEmpty(t *testing.T) {
+	ctx := context.Background()
+	s := newTestStore(t)
+
+	var out bytes.Buffer
+	require.NoError(t, runDisconnectAll(ctx, s, &out, strings.NewReader(""), false, false))
+	require.Contains(t, out.String(), "No reverse proxy instances found.")
+}

management/cmd/root.go

@@ -83,7 +83,8 @@ func init() {
 
 	rootCmd.AddCommand(migrationCmd)
 
-	tc := newTokenCommands()
-	tc.PersistentFlags().StringVar(&nbconfig.MgmtConfigPath, "config", defaultMgmtConfig, "Netbird config file location")
-	rootCmd.AddCommand(tc)
+	ac := newAdminCommands()
+	ac.PersistentFlags().StringVar(&nbconfig.MgmtConfigPath, "config", defaultMgmtConfig, "Netbird config file location")
+	rootCmd.AddCommand(ac)
+	rootCmd.AddCommand(newLegacyTokenCommand())
 }

management/cmd/token.go

@@ -1,55 +0,0 @@
-package cmd
-
-import (
-	"context"
-	"fmt"
-
-	log "github.com/sirupsen/logrus"
-	"github.com/spf13/cobra"
-
-	"github.com/netbirdio/netbird/formatter/hook"
-	tokencmd "github.com/netbirdio/netbird/management/cmd/token"
-	nbconfig "github.com/netbirdio/netbird/management/internals/server/config"
-	"github.com/netbirdio/netbird/management/server/store"
-	"github.com/netbirdio/netbird/util"
-)
-
-var tokenDatadir string
-
-// newTokenCommands creates the token command tree with management-specific store opener.
-func newTokenCommands() *cobra.Command {
-	cmd := tokencmd.NewCommands(withTokenStore)
-	cmd.PersistentFlags().StringVar(&tokenDatadir, "datadir", "", "Override the data directory from config (where store.db is located)")
-	return cmd
-}
-
-// withTokenStore initializes logging, loads config, opens the store, and calls fn.
-func withTokenStore(cmd *cobra.Command, fn func(ctx context.Context, s store.Store) error) error {
-	if err := util.InitLog("error", "console"); err != nil {
-		return fmt.Errorf("init log: %w", err)
-	}
-
-	ctx := context.WithValue(cmd.Context(), hook.ExecutionContextKey, hook.SystemSource) //nolint:staticcheck
-
-	config, err := LoadMgmtConfig(ctx, nbconfig.MgmtConfigPath)
-	if err != nil {
-		return fmt.Errorf("load config: %w", err)
-	}
-
-	datadir := config.Datadir
-	if tokenDatadir != "" {
-		datadir = tokenDatadir
-	}
-
-	s, err := store.NewStore(ctx, config.StoreConfig.Engine, datadir, nil, true)
-	if err != nil {
-		return fmt.Errorf("create store: %w", err)
-	}
-	defer func() {
-		if err := s.Close(ctx); err != nil {
-			log.Debugf("close store: %v", err)
-		}
-	}()
-
-	return fn(ctx, s)
-}

management/internals/controllers/network_map/controller/controller.go

@@ -29,6 +29,7 @@ import (
 	"github.com/netbirdio/netbird/management/server/store"
 	"github.com/netbirdio/netbird/management/server/telemetry"
 	"github.com/netbirdio/netbird/management/server/types"
+	sharedgrpc "github.com/netbirdio/netbird/shared/management/grpc"
 	"github.com/netbirdio/netbird/shared/management/proto"
 	"github.com/netbirdio/netbird/shared/management/status"
 	"github.com/netbirdio/netbird/util"
@@ -56,6 +57,10 @@ type Controller struct {
 	proxyController port_forwarding.Controller
 
 	integratedPeerValidator integrated_validator.IntegratedValidator
+
+	serverSupportedSyncMessageVersion sharedgrpc.SyncMessageVersion
+
+	perAccountServerSupportedSyncMessageVersions map[string]sharedgrpc.SyncMessageVersion
 }
 
 type bufferUpdate struct {
@@ -90,8 +95,10 @@ func NewController(ctx context.Context, store store.Store, metrics telemetry.App
 		dnsDomain:               dnsDomain,
 		config:                  config,
 
-		proxyController:       proxyController,
-		EphemeralPeersManager: ephemeralPeersManager,
+		proxyController:                              proxyController,
+		EphemeralPeersManager:                        ephemeralPeersManager,
+		serverSupportedSyncMessageVersion:            sharedgrpc.SyncMessageVersionFromConfig(config.HighestSupportedSyncMessageVersion),
+		perAccountServerSupportedSyncMessageVersions: sharedgrpc.SyncMessageVersionsFromMap(config.PerAccountHighestSupportedSyncMessageVersion),
 	}
 }
 
@@ -222,18 +229,53 @@ func (c *Controller) sendUpdateAccountPeers(ctx context.Context, accountID strin
 			c.metrics.CountCalcPostureChecksDuration(time.Since(start))
 			start = time.Now()
 
-			remotePeerNetworkMap := account.GetPeerNetworkMapFromComponents(ctx, p.ID, peersCustomZone, accountZones, approvedPeersMap, resourcePolicies, routers, c.accountManagerMetrics, groupIDToUserIDs)
+			peerGroups := account.GetPeerGroups(p.ID)
+			proxyNetworkMap := proxyNetworkMaps[p.ID]
+			var update *proto.SyncResponse
+
+			commonSyncMessageVersion := sharedgrpc.HighestCommonSyncMessageVersion(
+				c.perAccountOrGlobalSupportedSyncMessageVersions(accountID),
+				sharedgrpc.SyncMessageVersionFromConfig(&peer.Meta.SyncMessageVersion))
+
+			log.WithContext(ctx).
+				WithFields(log.Fields{
+					"sync_message_version":        commonSyncMessageVersion,
+					"server_sync_message_version": c.perAccountOrGlobalSupportedSyncMessageVersions(peer.AccountID),
+					"peer_sync_message_version":   sharedgrpc.SyncMessageVersionFromConfig(&peer.Meta.SyncMessageVersion),
+				}).Debug("common highest sync message version")
+
+			if commonSyncMessageVersion == sharedgrpc.ComponentNetworkMap {
+				components := account.GetPeerNetworkMapComponents(
+					ctx, p.ID, peersCustomZone, accountZones, approvedPeersMap, resourcePolicies, routers, groupIDToUserIDs)
+
+				c.metrics.CountCalcPeerNetworkMapDuration(time.Since(start))
+
+				start = time.Now()
+				// proxyNetworkMap rides the envelope as a ProxyPatch sidecar;
+				// the client merges it into Calculate()'s output the same
+				// way the legacy server did via NetworkMap.Merge.
+				update = grpc.ToComponentSyncResponse(ctx, nil, c.config.HttpConfig, c.config.DeviceAuthorizationFlow, p, nil, nil, components, proxyNetworkMap, dnsDomain, postureChecks, account.Settings, extraSetting, maps.Keys(peerGroups), dnsFwdPort)
+				c.metrics.CountToComponentSyncResponseDuration(time.Since(start))
+
+				c.peersUpdateManager.SendUpdate(ctx, p.ID, &network_map.UpdateMessage{
+					Update:      update,
+					MessageType: network_map.MessageTypeNetworkMap,
+				})
+
+				return
+			}
+
+			nmap := account.GetPeerNetworkMapFromComponents(
+				ctx, p.ID, peersCustomZone, accountZones, approvedPeersMap, resourcePolicies, routers, c.accountManagerMetrics, groupIDToUserIDs)
 
 			c.metrics.CountCalcPeerNetworkMapDuration(time.Since(start))
 
-			proxyNetworkMap, ok := proxyNetworkMaps[p.ID]
-			if ok {
-				remotePeerNetworkMap.Merge(proxyNetworkMap)
+			if proxyNetworkMap != nil {
+				nmap.Merge(proxyNetworkMap)
 			}
 
-			peerGroups := account.GetPeerGroups(p.ID)
 			start = time.Now()
-			update := grpc.ToSyncResponse(ctx, nil, c.config.HttpConfig, c.config.DeviceAuthorizationFlow, p, nil, nil, remotePeerNetworkMap, dnsDomain, postureChecks, dnsCache, account.Settings, extraSetting, maps.Keys(peerGroups), dnsFwdPort)
+			update = grpc.ToSyncResponse(ctx, nil, c.config.HttpConfig, c.config.DeviceAuthorizationFlow, p, nil, nil, nmap, dnsDomain, postureChecks, dnsCache, account.Settings, extraSetting, maps.Keys(peerGroups), dnsFwdPort)
 			c.metrics.CountToSyncResponseDuration(time.Since(start))
 
 			c.peersUpdateManager.SendUpdate(ctx, p.ID, &network_map.UpdateMessage{
@@ -251,6 +293,13 @@ func (c *Controller) sendUpdateAccountPeers(ctx context.Context, accountID strin
 	return nil
 }
 
+func (c *Controller) perAccountOrGlobalSupportedSyncMessageVersions(accountId string) sharedgrpc.SyncMessageVersion {
+	if perAccount, ok := c.perAccountServerSupportedSyncMessageVersions[accountId]; ok {
+		return perAccount
+	}
+	return c.serverSupportedSyncMessageVersion
+}
+
 // UpdatePeers updates all peers that belong to an account.
 // Should be called when changes have to be synced to peers.
 func (c *Controller) UpdateAccountPeers(ctx context.Context, accountID string, reason types.UpdateReason) error {
@@ -352,18 +401,53 @@ func (c *Controller) sendUpdateForAffectedPeers(ctx context.Context, accountID s
 			c.metrics.CountCalcPostureChecksDuration(time.Since(start))
 			start = time.Now()
 
-			remotePeerNetworkMap := account.GetPeerNetworkMapFromComponents(ctx, p.ID, peersCustomZone, accountZones, approvedPeersMap, resourcePolicies, routers, c.accountManagerMetrics, groupIDToUserIDs)
+			peerGroups := account.GetPeerGroups(p.ID)
+			proxyNetworkMap := proxyNetworkMaps[p.ID]
+			var update *proto.SyncResponse
+
+			commonSyncMessageVersion := sharedgrpc.HighestCommonSyncMessageVersion(
+				c.perAccountOrGlobalSupportedSyncMessageVersions(accountID),
+				sharedgrpc.SyncMessageVersionFromConfig(&peer.Meta.SyncMessageVersion))
+
+			log.WithContext(ctx).
+				WithFields(log.Fields{
+					"sync_message_version":        commonSyncMessageVersion,
+					"server_sync_message_version": c.perAccountOrGlobalSupportedSyncMessageVersions(peer.AccountID),
+					"peer_sync_message_version":   sharedgrpc.SyncMessageVersionFromConfig(&peer.Meta.SyncMessageVersion),
+				}).Debug("common highest sync message version")
+
+			if commonSyncMessageVersion == sharedgrpc.ComponentNetworkMap {
+				components := account.GetPeerNetworkMapComponents(
+					ctx, p.ID, peersCustomZone, accountZones, approvedPeersMap, resourcePolicies, routers, groupIDToUserIDs)
+
+				c.metrics.CountCalcPeerNetworkMapDuration(time.Since(start))
+
+				start = time.Now()
+				// proxyNetworkMap rides the envelope as a ProxyPatch sidecar;
+				// the client merges it into Calculate()'s output the same
+				// way the legacy server did via NetworkMap.Merge.
+				update = grpc.ToComponentSyncResponse(ctx, nil, c.config.HttpConfig, c.config.DeviceAuthorizationFlow, p, nil, nil, components, proxyNetworkMap, dnsDomain, postureChecks, account.Settings, extraSetting, maps.Keys(peerGroups), dnsFwdPort)
+				c.metrics.CountToComponentSyncResponseDuration(time.Since(start))
+
+				c.peersUpdateManager.SendUpdate(ctx, p.ID, &network_map.UpdateMessage{
+					Update:      update,
+					MessageType: network_map.MessageTypeNetworkMap,
+				})
+
+				return
+			}
+
+			nmap := account.GetPeerNetworkMapFromComponents(
+				ctx, p.ID, peersCustomZone, accountZones, approvedPeersMap, resourcePolicies, routers, c.accountManagerMetrics, groupIDToUserIDs)
 
 			c.metrics.CountCalcPeerNetworkMapDuration(time.Since(start))
 
-			proxyNetworkMap, ok := proxyNetworkMaps[p.ID]
-			if ok {
-				remotePeerNetworkMap.Merge(proxyNetworkMap)
+			if proxyNetworkMap != nil {
+				nmap.Merge(proxyNetworkMap)
 			}
 
-			peerGroups := account.GetPeerGroups(p.ID)
 			start = time.Now()
-			update := grpc.ToSyncResponse(ctx, nil, c.config.HttpConfig, c.config.DeviceAuthorizationFlow, p, nil, nil, remotePeerNetworkMap, dnsDomain, postureChecks, dnsCache, account.Settings, extraSetting, maps.Keys(peerGroups), dnsFwdPort)
+			update = grpc.ToSyncResponse(ctx, nil, c.config.HttpConfig, c.config.DeviceAuthorizationFlow, p, nil, nil, nmap, dnsDomain, postureChecks, dnsCache, account.Settings, extraSetting, maps.Keys(peerGroups), dnsFwdPort)
 			c.metrics.CountToSyncResponseDuration(time.Since(start))
 
 			c.peersUpdateManager.SendUpdate(ctx, p.ID, &network_map.UpdateMessage{
@@ -451,13 +535,7 @@ func (c *Controller) UpdateAccountPeer(ctx context.Context, accountId string, pe
 		return err
 	}
 
-	remotePeerNetworkMap := account.GetPeerNetworkMapFromComponents(ctx, peerId, peersCustomZone, accountZones, approvedPeersMap, resourcePolicies, routers, c.accountManagerMetrics, groupIDToUserIDs)
-
-	proxyNetworkMap, ok := proxyNetworkMaps[peer.ID]
-	if ok {
-		remotePeerNetworkMap.Merge(proxyNetworkMap)
-	}
-
+	proxyNetworkMap := proxyNetworkMaps[peer.ID]
 	extraSettings, err := c.settingsManager.GetExtraSettings(ctx, peer.AccountID)
 	if err != nil {
 		return fmt.Errorf("failed to get extra settings: %v", err)
@@ -466,7 +544,45 @@ func (c *Controller) UpdateAccountPeer(ctx context.Context, accountId string, pe
 	peerGroups := account.GetPeerGroups(peerId)
 	dnsFwdPort := computeForwarderPort(maps.Values(account.Peers), network_map.DnsForwarderPortMinVersion)
 
-	update := grpc.ToSyncResponse(ctx, nil, c.config.HttpConfig, c.config.DeviceAuthorizationFlow, peer, nil, nil, remotePeerNetworkMap, dnsDomain, postureChecks, dnsCache, account.Settings, extraSettings, maps.Keys(peerGroups), dnsFwdPort)
+	var update *proto.SyncResponse
+
+	commonSyncMessageVersion := sharedgrpc.HighestCommonSyncMessageVersion(
+		c.perAccountOrGlobalSupportedSyncMessageVersions(accountId),
+		sharedgrpc.SyncMessageVersionFromConfig(&peer.Meta.SyncMessageVersion))
+
+	log.WithContext(ctx).
+		WithFields(log.Fields{
+			"sync_message_version":        commonSyncMessageVersion,
+			"server_sync_message_version": c.perAccountOrGlobalSupportedSyncMessageVersions(peer.AccountID),
+			"peer_sync_message_version":   sharedgrpc.SyncMessageVersionFromConfig(&peer.Meta.SyncMessageVersion),
+		}).Debug("common highest sync message version")
+
+	if commonSyncMessageVersion == sharedgrpc.ComponentNetworkMap {
+		components := account.GetPeerNetworkMapComponents(
+			ctx, peer.ID, peersCustomZone, accountZones, approvedPeersMap, resourcePolicies, routers, groupIDToUserIDs)
+
+		// proxyNetworkMap rides the envelope as a ProxyPatch sidecar;
+		// the client merges it into Calculate()'s output the same
+		// way the legacy server did via NetworkMap.Merge.
+		update = grpc.ToComponentSyncResponse(ctx, nil, c.config.HttpConfig, c.config.DeviceAuthorizationFlow, peer, nil, nil, components, proxyNetworkMap, dnsDomain, postureChecks, account.Settings, extraSettings, maps.Keys(peerGroups), dnsFwdPort)
+
+		c.peersUpdateManager.SendUpdate(ctx, peer.ID, &network_map.UpdateMessage{
+			Update:      update,
+			MessageType: network_map.MessageTypeNetworkMap,
+		})
+
+		return nil
+	}
+
+	nmap := account.GetPeerNetworkMapFromComponents(
+		ctx, peer.ID, peersCustomZone, accountZones, approvedPeersMap, resourcePolicies, routers, c.accountManagerMetrics, groupIDToUserIDs)
+
+	if proxyNetworkMap != nil {
+		nmap.Merge(proxyNetworkMap)
+	}
+
+	update = grpc.ToSyncResponse(ctx, nil, c.config.HttpConfig, c.config.DeviceAuthorizationFlow, peer, nil, nil, nmap, dnsDomain, postureChecks, dnsCache, account.Settings, extraSettings, maps.Keys(peerGroups), dnsFwdPort)
+
 	c.peersUpdateManager.SendUpdate(ctx, peer.ID, &network_map.UpdateMessage{
 		Update:      update,
 		MessageType: network_map.MessageTypeNetworkMap,
@@ -513,6 +629,65 @@ func (c *Controller) BufferUpdateAccountPeers(ctx context.Context, accountID str
 	return nil
 }
 
+// GetValidatedPeerWithComponents is the components-format counterpart of
+// GetValidatedPeerWithMap. It returns raw NetworkMapComponents for capable
+// peers along with the proxy NetworkMap fragment (BYOP / port-forwarding
+// data the legacy server folds in via NetworkMap.Merge). The gRPC layer
+// encodes both into the wire envelope. Callers must gate on capability
+// themselves before dispatching here — this method does NOT branch on it.
+func (c *Controller) GetValidatedPeerWithComponents(ctx context.Context, isRequiresApproval bool, accountID string, peer *nbpeer.Peer) (*nbpeer.Peer, *types.NetworkMapComponents, *types.NetworkMap, []*posture.Checks, int64, error) {
+	if isRequiresApproval {
+		network, err := c.repo.GetAccountNetwork(ctx, accountID)
+		if err != nil {
+			return nil, nil, nil, nil, 0, err
+		}
+		return peer, &types.NetworkMapComponents{Network: network.Copy()}, nil, nil, 0, nil
+	}
+
+	account, err := c.requestBuffer.GetAccountWithBackpressure(ctx, accountID)
+	if err != nil {
+		return nil, nil, nil, nil, 0, err
+	}
+
+	c.injectAllProxyPolicies(ctx, account)
+
+	approvedPeersMap, err := c.integratedPeerValidator.GetValidatedPeers(ctx, account.Id, maps.Values(account.Groups), maps.Values(account.Peers), account.Settings.Extra)
+	if err != nil {
+		return nil, nil, nil, nil, 0, err
+	}
+
+	postureChecks, err := c.getPeerPostureChecks(account, peer.ID)
+	if err != nil {
+		return nil, nil, nil, nil, 0, err
+	}
+
+	accountZones, err := c.repo.GetAccountZones(ctx, account.Id)
+	if err != nil {
+		return nil, nil, nil, nil, 0, err
+	}
+
+	// Fetch the proxy network map fragment for this peer alongside the
+	// components — same single-account-load path the streaming controller
+	// uses, so initial-sync delivers BYOP/forwarding patches synchronously
+	// instead of waiting for the next streaming push.
+	proxyNetworkMaps, err := c.proxyController.GetProxyNetworkMaps(ctx, account.Id, peer.ID, account.Peers)
+	if err != nil {
+		log.WithContext(ctx).Errorf("failed to get proxy network maps: %v", err)
+		return nil, nil, nil, nil, 0, err
+	}
+
+	dnsDomain := c.GetDNSDomain(account.Settings)
+	peersCustomZone := account.GetPeersCustomZone(ctx, dnsDomain)
+
+	resourcePolicies := account.GetResourcePoliciesMap()
+	routers := account.GetResourceRoutersMap()
+	groupIDToUserIDs := account.GetActiveGroupUsers()
+	components := account.GetPeerNetworkMapComponents(ctx, peer.ID, peersCustomZone, accountZones, approvedPeersMap, resourcePolicies, routers, groupIDToUserIDs)
+	dnsFwdPort := computeForwarderPort(maps.Values(account.Peers), network_map.DnsForwarderPortMinVersion)
+
+	return peer, components, proxyNetworkMaps[peer.ID], postureChecks, dnsFwdPort, nil
+}
+
 // BufferUpdateAffectedPeers accumulates peer IDs and flushes them after the buffer interval.
 func (c *Controller) BufferUpdateAffectedPeers(ctx context.Context, accountID string, peerIDs []string, reason types.UpdateReason) error {
 	if len(peerIDs) == 0 {

management/internals/controllers/network_map/interface.go

@@ -24,6 +24,7 @@ type Controller interface {
 	UpdateAccountPeer(ctx context.Context, accountId string, peerId string) error
 	BufferUpdateAccountPeers(ctx context.Context, accountID string, reason types.UpdateReason) error
 	GetValidatedPeerWithMap(ctx context.Context, isRequiresApproval bool, accountID string, peerID string) (*types.NetworkMap, []*posture.Checks, int64, error)
+	GetValidatedPeerWithComponents(ctx context.Context, isRequiresApproval bool, accountID string, p *nbpeer.Peer) (*nbpeer.Peer, *types.NetworkMapComponents, *types.NetworkMap, []*posture.Checks, int64, error)
 	GetDNSDomain(settings *types.Settings) string
 	StartWarmup(context.Context)
 	GetNetworkMap(ctx context.Context, peerID string) (*types.NetworkMap, error)

management/internals/controllers/network_map/interface_mock.go

@@ -1,9 +1,9 @@
 // Code generated by MockGen. DO NOT EDIT.
-// Source: management/internals/controllers/network_map/interface.go
+// Source: ./interface.go
 //
 // Generated by this command:
 //
-//	mockgen -package network_map -destination=management/internals/controllers/network_map/interface_mock.go -source=management/internals/controllers/network_map/interface.go -build_flags=-mod=mod
+//	mockgen -package network_map -destination=interface_mock.go -source=./interface.go -build_flags=-mod=mod
 //
 
 // Package network_map is a generated GoMock package.
@@ -126,8 +126,27 @@ func (mr *MockControllerMockRecorder) GetNetworkMap(ctx, peerID any) *gomock.Cal
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetNetworkMap", reflect.TypeOf((*MockController)(nil).GetNetworkMap), ctx, peerID)
 }
 
+// GetValidatedPeerWithComponents mocks base method.
+func (m *MockController) GetValidatedPeerWithComponents(ctx context.Context, isRequiresApproval bool, accountID string, p *peer.Peer) (*peer.Peer, *types.NetworkMapComponents, *types.NetworkMap, []*posture.Checks, int64, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetValidatedPeerWithComponents", ctx, isRequiresApproval, accountID, p)
+	ret0, _ := ret[0].(*peer.Peer)
+	ret1, _ := ret[1].(*types.NetworkMapComponents)
+	ret2, _ := ret[2].(*types.NetworkMap)
+	ret3, _ := ret[3].([]*posture.Checks)
+	ret4, _ := ret[4].(int64)
+	ret5, _ := ret[5].(error)
+	return ret0, ret1, ret2, ret3, ret4, ret5
+}
+
+// GetValidatedPeerWithComponents indicates an expected call of GetValidatedPeerWithComponents.
+func (mr *MockControllerMockRecorder) GetValidatedPeerWithComponents(ctx, isRequiresApproval, accountID, p any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetValidatedPeerWithComponents", reflect.TypeOf((*MockController)(nil).GetValidatedPeerWithComponents), ctx, isRequiresApproval, accountID, p)
+}
+
 // GetValidatedPeerWithMap mocks base method.
-func (m *MockController) GetValidatedPeerWithMap(ctx context.Context, isRequiresApproval bool, accountID string, peerID string) (*types.NetworkMap, []*posture.Checks, int64, error) {
+func (m *MockController) GetValidatedPeerWithMap(ctx context.Context, isRequiresApproval bool, accountID, peerID string) (*types.NetworkMap, []*posture.Checks, int64, error) {
 	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "GetValidatedPeerWithMap", ctx, isRequiresApproval, accountID, peerID)
 	ret0, _ := ret[0].(*types.NetworkMap)
@@ -171,7 +190,7 @@ func (mr *MockControllerMockRecorder) OnPeerDisconnected(ctx, accountID, peerID
 }
 
 // OnPeersAdded mocks base method.
-func (m *MockController) OnPeersAdded(ctx context.Context, accountID string, peerIDs []string, affectedPeerIDs []string) error {
+func (m *MockController) OnPeersAdded(ctx context.Context, accountID string, peerIDs, affectedPeerIDs []string) error {
 	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "OnPeersAdded", ctx, accountID, peerIDs, affectedPeerIDs)
 	ret0, _ := ret[0].(error)
@@ -185,7 +204,7 @@ func (mr *MockControllerMockRecorder) OnPeersAdded(ctx, accountID, peerIDs, affe
 }
 
 // OnPeersDeleted mocks base method.
-func (m *MockController) OnPeersDeleted(ctx context.Context, accountID string, peerIDs []string, affectedPeerIDs []string) error {
+func (m *MockController) OnPeersDeleted(ctx context.Context, accountID string, peerIDs, affectedPeerIDs []string) error {
 	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "OnPeersDeleted", ctx, accountID, peerIDs, affectedPeerIDs)
 	ret0, _ := ret[0].(error)
@@ -199,7 +218,7 @@ func (mr *MockControllerMockRecorder) OnPeersDeleted(ctx, accountID, peerIDs, af
 }
 
 // OnPeersUpdated mocks base method.
-func (m *MockController) OnPeersUpdated(ctx context.Context, accountId string, peerIDs []string, affectedPeerIDs []string) error {
+func (m *MockController) OnPeersUpdated(ctx context.Context, accountId string, peerIDs, affectedPeerIDs []string) error {
 	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "OnPeersUpdated", ctx, accountId, peerIDs, affectedPeerIDs)
 	ret0, _ := ret[0].(error)

management/internals/modules/agentnetwork/accesslog_ingest.go

@@ -18,21 +18,26 @@ import (
 // contract between the proxy and management; management flattens them into
 // queryable columns. Keep in sync with the proxy side.
 const (
-	metaKeyProvider           = "llm.provider"
-	metaKeyModel              = "llm.model"
-	metaKeyResolvedProviderID = "llm.resolved_provider_id"
-	metaKeySelectedPolicyID   = "llm.selected_policy_id"
-	metaKeyPolicyDecision     = "llm_policy.decision"
-	metaKeyPolicyReason       = "llm_policy.reason"
-	metaKeyInputTokens        = "llm.input_tokens"  //nolint:gosec // metadata key name, not a credential
-	metaKeyOutputTokens       = "llm.output_tokens" //nolint:gosec // metadata key name, not a credential
-	metaKeyTotalTokens        = "llm.total_tokens"  //nolint:gosec // metadata key name, not a credential
-	metaKeyCostUSDTotal       = "cost.usd_total"
-	metaKeyStream             = "llm.stream"
-	metaKeySessionID          = "llm.session_id"
-	metaKeyAuthorisingGroups  = "llm.authorising_groups"
-	metaKeyRequestPrompt      = "llm.request_prompt"
-	metaKeyResponseCompletion = "llm.response_completion"
+	metaKeyProvider            = "llm.provider"
+	metaKeyModel               = "llm.model"
+	metaKeyResolvedProviderID  = "llm.resolved_provider_id"
+	metaKeySelectedPolicyID    = "llm.selected_policy_id"
+	metaKeyPolicyDecision      = "llm_policy.decision"
+	metaKeyPolicyReason        = "llm_policy.reason"
+	metaKeyInputTokens         = "llm.input_tokens"          //nolint:gosec // metadata key name, not a credential
+	metaKeyOutputTokens        = "llm.output_tokens"         //nolint:gosec // metadata key name, not a credential
+	metaKeyTotalTokens         = "llm.total_tokens"          //nolint:gosec // metadata key name, not a credential
+	metaKeyCachedInputTokens   = "llm.cached_input_tokens"   //nolint:gosec // metadata key name, not a credential
+	metaKeyCacheCreationTokens = "llm.cache_creation_tokens" //nolint:gosec // metadata key name, not a credential
+	metaKeyCostUSDInput        = "cost.usd_input"
+	metaKeyCostUSDCachedInput  = "cost.usd_cached_input"
+	metaKeyCostUSDCacheCreate  = "cost.usd_cache_creation"
+	metaKeyCostUSDOutput       = "cost.usd_output"
+	metaKeyStream              = "llm.stream"
+	metaKeySessionID           = "llm.session_id"
+	metaKeyAuthorisingGroups   = "llm.authorising_groups"
+	metaKeyRequestPrompt       = "llm.request_prompt"
+	metaKeyResponseCompletion  = "llm.response_completion"
 )
 
 // IngestAccessLog flattens the metadata-bearing reverse-proxy access-log entry
@@ -108,20 +113,25 @@ func flattenAccessLog(e *accesslogs.AccessLogEntry) (*types.AgentNetworkAccessLo
 		BytesUpload:   e.BytesUpload,
 		BytesDownload: e.BytesDownload,
 
-		Provider:           meta[metaKeyProvider],
-		Model:              meta[metaKeyModel],
-		SessionID:          meta[metaKeySessionID],
-		ResolvedProviderID: meta[metaKeyResolvedProviderID],
-		SelectedPolicyID:   meta[metaKeySelectedPolicyID],
-		Decision:           meta[metaKeyPolicyDecision],
-		DenyReason:         meta[metaKeyPolicyReason],
-		InputTokens:        parseMetaInt(meta, metaKeyInputTokens),
-		OutputTokens:       parseMetaInt(meta, metaKeyOutputTokens),
-		TotalTokens:        parseMetaInt(meta, metaKeyTotalTokens),
-		CostUSD:            parseMetaFloat(meta, metaKeyCostUSDTotal),
-		Stream:             parseMetaBool(meta, metaKeyStream),
-		RequestPrompt:      meta[metaKeyRequestPrompt],
-		ResponseCompletion: meta[metaKeyResponseCompletion],
+		Provider:             meta[metaKeyProvider],
+		Model:                meta[metaKeyModel],
+		SessionID:            meta[metaKeySessionID],
+		ResolvedProviderID:   meta[metaKeyResolvedProviderID],
+		SelectedPolicyID:     meta[metaKeySelectedPolicyID],
+		Decision:             meta[metaKeyPolicyDecision],
+		DenyReason:           meta[metaKeyPolicyReason],
+		InputTokens:          parseMetaInt(meta, metaKeyInputTokens),
+		OutputTokens:         parseMetaInt(meta, metaKeyOutputTokens),
+		TotalTokens:          parseMetaInt(meta, metaKeyTotalTokens),
+		CachedInputTokens:    parseMetaInt(meta, metaKeyCachedInputTokens),
+		CacheCreationTokens:  parseMetaInt(meta, metaKeyCacheCreationTokens),
+		InputCostUSD:         parseMetaFloat(meta, metaKeyCostUSDInput),
+		CachedInputCostUSD:   parseMetaFloat(meta, metaKeyCostUSDCachedInput),
+		CacheCreationCostUSD: parseMetaFloat(meta, metaKeyCostUSDCacheCreate),
+		OutputCostUSD:        parseMetaFloat(meta, metaKeyCostUSDOutput),
+		Stream:               parseMetaBool(meta, metaKeyStream),
+		RequestPrompt:        meta[metaKeyRequestPrompt],
+		ResponseCompletion:   meta[metaKeyResponseCompletion],
 	}
 
 	var groups []types.AgentNetworkAccessLogGroup
@@ -140,18 +150,23 @@ func flattenAccessLog(e *accesslogs.AccessLogEntry) (*types.AgentNetworkAccessLo
 // log's ID so the two correlate.
 func usageFromFlattenedLog(e *types.AgentNetworkAccessLog, groups []types.AgentNetworkAccessLogGroup) (*types.AgentNetworkUsage, []types.AgentNetworkUsageGroup) {
 	usage := &types.AgentNetworkUsage{
-		ID:                 e.ID,
-		AccountID:          e.AccountID,
-		Timestamp:          e.Timestamp,
-		UserID:             e.UserID,
-		ResolvedProviderID: e.ResolvedProviderID,
-		Provider:           e.Provider,
-		Model:              e.Model,
-		SessionID:          e.SessionID,
-		InputTokens:        e.InputTokens,
-		OutputTokens:       e.OutputTokens,
-		TotalTokens:        e.TotalTokens,
-		CostUSD:            e.CostUSD,
+		ID:                   e.ID,
+		AccountID:            e.AccountID,
+		Timestamp:            e.Timestamp,
+		UserID:               e.UserID,
+		ResolvedProviderID:   e.ResolvedProviderID,
+		Provider:             e.Provider,
+		Model:                e.Model,
+		SessionID:            e.SessionID,
+		InputTokens:          e.InputTokens,
+		OutputTokens:         e.OutputTokens,
+		TotalTokens:          e.TotalTokens,
+		CachedInputTokens:    e.CachedInputTokens,
+		CacheCreationTokens:  e.CacheCreationTokens,
+		InputCostUSD:         e.InputCostUSD,
+		CachedInputCostUSD:   e.CachedInputCostUSD,
+		CacheCreationCostUSD: e.CacheCreationCostUSD,
+		OutputCostUSD:        e.OutputCostUSD,
 	}
 
 	usageGroups := make([]types.AgentNetworkUsageGroup, 0, len(groups))

management/internals/modules/agentnetwork/accesslog_ingest_realstore_test.go

@@ -28,17 +28,22 @@ func newIngestTestEntry() *accesslogs.AccessLogEntry {
 		UserId:       "user-1",
 		AgentNetwork: true,
 		Metadata: map[string]string{
-			metaKeyProvider:           "openai",
-			metaKeyModel:              "gpt-5.4",
-			metaKeyResolvedProviderID: "prov-1",
-			metaKeySessionID:          "sess-1",
-			metaKeyInputTokens:        "100",
-			metaKeyOutputTokens:       "50",
-			metaKeyTotalTokens:        "150",
-			metaKeyCostUSDTotal:       "0.0123",
-			metaKeyStream:             "true",
-			metaKeyRequestPrompt:      "hello",
-			metaKeyResponseCompletion: "world",
+			metaKeyProvider:            "openai",
+			metaKeyModel:               "gpt-5.4",
+			metaKeyResolvedProviderID:  "prov-1",
+			metaKeySessionID:           "sess-1",
+			metaKeyInputTokens:         "100",
+			metaKeyOutputTokens:        "50",
+			metaKeyTotalTokens:         "1174",
+			metaKeyCachedInputTokens:   "256",
+			metaKeyCacheCreationTokens: "768",
+			metaKeyCostUSDInput:        "0.0071",
+			metaKeyCostUSDCachedInput:  "0.0009",
+			metaKeyCostUSDCacheCreate:  "0.0020",
+			metaKeyCostUSDOutput:       "0.0023",
+			metaKeyStream:              "true",
+			metaKeyRequestPrompt:       "hello",
+			metaKeyResponseCompletion:  "world",
 			// repeated id must be de-duplicated before the group rows insert.
 			metaKeyAuthorisingGroups: "grp-eng,grp-eng,grp-ops",
 		},
@@ -65,7 +70,19 @@ func TestIngestAccessLog_RealStore_LogCollectionOff(t *testing.T) {
 	require.Len(t, usage, 1, "usage row must be written even with log collection off")
 	assert.Equal(t, int64(100), usage[0].InputTokens, "input tokens must round-trip from metadata")
 	assert.Equal(t, int64(50), usage[0].OutputTokens, "output tokens must round-trip from metadata")
-	assert.InDelta(t, 0.0123, usage[0].CostUSD, 1e-9, "cost must round-trip from metadata")
+	assert.Equal(t, int64(256), usage[0].CachedInputTokens, "cache-read tokens must round-trip from metadata")
+	assert.Equal(t, int64(768), usage[0].CacheCreationTokens, "cache-write tokens must round-trip from metadata")
+	// The per-bucket breakdown is the only cost state stored, and must survive
+	// the write/read cycle as real columns — usage rows are the only cost
+	// record for accounts with log collection off, so a dropped column here
+	// loses the split permanently.
+	assert.InDelta(t, 0.0071, usage[0].InputCostUSD, 1e-9, "input cost must round-trip from metadata")
+	assert.InDelta(t, 0.0009, usage[0].CachedInputCostUSD, 1e-9, "cache-read cost must round-trip from metadata")
+	assert.InDelta(t, 0.0020, usage[0].CacheCreationCostUSD, 1e-9, "cache-write cost must round-trip from metadata")
+	assert.InDelta(t, 0.0023, usage[0].OutputCostUSD, 1e-9, "output cost must round-trip from metadata")
+	// Aggregates are derived from the stored columns, never stored themselves.
+	assert.InDelta(t, 0.0123, usage[0].TotalCostUSD(), 1e-9, "total is derived from the stored buckets")
+	assert.InDelta(t, 0.0029, usage[0].CacheCostUSD(), 1e-9, "cache cost is derived from the two cache buckets")
 
 	logs, _, err := s.GetAgentNetworkAccessLogs(ctx, store.LockingStrengthNone, testAccountID, types.AgentNetworkAccessLogFilter{})
 	require.NoError(t, err)
@@ -96,6 +113,14 @@ func TestIngestAccessLog_RealStore_LogCollectionOn(t *testing.T) {
 	require.Equal(t, int64(1), total, "exactly one access-log row expected")
 	require.Len(t, logs, 1, "full access-log row must be written when log collection is on")
 	assert.Equal(t, "gpt-5.4", logs[0].Model, "model must flatten from metadata")
+	assert.Equal(t, int64(256), logs[0].CachedInputTokens, "cache-read tokens must flatten from metadata")
+	assert.Equal(t, int64(768), logs[0].CacheCreationTokens, "cache-write tokens must flatten from metadata")
+	assert.InDelta(t, 0.0029, logs[0].CacheCostUSD(), 1e-9, "cache cost is derived from the two cache buckets")
+	assert.InDelta(t, 0.0123, logs[0].TotalCostUSD(), 1e-9, "total is derived from the stored buckets")
+	assert.InDelta(t, 0.0071, logs[0].InputCostUSD, 1e-9, "input cost must flatten from metadata")
+	assert.InDelta(t, 0.0009, logs[0].CachedInputCostUSD, 1e-9, "cache-read cost must flatten from metadata")
+	assert.InDelta(t, 0.0020, logs[0].CacheCreationCostUSD, 1e-9, "cache-write cost must flatten from metadata")
+	assert.InDelta(t, 0.0023, logs[0].OutputCostUSD, 1e-9, "output cost must flatten from metadata")
 	assert.Equal(t, "hello", logs[0].RequestPrompt, "prompt must be retained when log collection is on")
 	assert.Equal(t, "world", logs[0].ResponseCompletion, "completion must be retained when log collection is on")
 	assert.True(t, logs[0].Stream, "stream flag must flatten from metadata")

management/internals/modules/agentnetwork/accesslog_sessions_realstore_test.go

@@ -38,7 +38,7 @@ func accessLogRow(id, sessionID string, ts time.Time, opts ...func(*types.AgentN
 		InputTokens:        100,
 		OutputTokens:       50,
 		TotalTokens:        150,
-		CostUSD:            0.01,
+		InputCostUSD:       0.01,
 	}
 	for _, o := range opts {
 		o(e)
@@ -74,7 +74,7 @@ func withTokens(in, out, total int64, cost float64) func(*types.AgentNetworkAcce
 		e.InputTokens = in
 		e.OutputTokens = out
 		e.TotalTokens = total
-		e.CostUSD = cost
+		e.InputCostUSD = cost
 	}
 }
 
@@ -155,7 +155,7 @@ func TestAccessLogSessions_FoldAndAggregate(t *testing.T) {
 	assert.Equal(t, int64(310), a.InputTokens, "input tokens summed")
 	assert.Equal(t, int64(135), a.OutputTokens, "output tokens summed")
 	assert.Equal(t, int64(445), a.TotalTokens, "total tokens summed")
-	assert.InDelta(t, 0.031, a.CostUSD, 1e-9, "cost summed")
+	assert.InDelta(t, 0.031, a.TotalCostUSD(), 1e-9, "cost summed")
 	assert.Equal(t, "deny", a.Decision, "any deny makes the session a deny")
 	assert.ElementsMatch(t, []string{"openai", "anthropic"}, a.Providers, "distinct providers")
 	assert.ElementsMatch(t, []string{"gpt-5.4", "claude-haiku-4-5"}, a.Models, "distinct models")

management/internals/modules/agentnetwork/catalog/catalog.go

@@ -420,6 +420,47 @@ var providers = []Provider{
 			{ID: "mistral-embed", Label: "Mistral Embed", InputPer1k: 0.0001, OutputPer1k: 0, ContextWindow: 8192},
 		},
 	},
+	{
+		ID:                 "kimi_api",
+		Kind:               KindProvider,
+		Name:               "Kimi (Moonshot AI) API",
+		Description:        "Kimi K3 / K2 models via the Moonshot AI platform",
+		DefaultHost:        "api.moonshot.ai",
+		AuthHeaderName:     "Authorization",
+		AuthHeaderTemplate: "Bearer ${API_KEY}",
+		DefaultContentType: "application/json",
+		BrandColor:         "#1A1A2E",
+		// ParserID empty on purpose: Moonshot serves two body shapes on
+		// the same host and key, and the proxy's URL sniffer dispatches
+		// both (same pattern as Bifrost). /v1/chat/completions matches
+		// OpenAIParser; the Anthropic-compatible endpoint the official
+		// Claude Code guide uses (/anthropic/v1/messages) contains
+		// "/v1/messages" and matches AnthropicParser. Pinning "openai"
+		// here would misparse the Claude Code path — the primary way
+		// teams consume Kimi for coding today. Both endpoints accept the
+		// same Moonshot key via Authorization: Bearer (Claude Code's
+		// ANTHROPIC_AUTH_TOKEN rides that header too).
+		//
+		// api.moonshot.ai is the international platform; mainland-China
+		// accounts live on api.moonshot.cn with separate billing —
+		// operators there override the host on the provider record. The
+		// kimi.com subscription coding endpoint (api.kimi.com/coding,
+		// model id "k3") is account-bound seat licensing rather than a
+		// meterable platform key, so it's deliberately not the default.
+		ParserID: "",
+		// Pricing per Moonshot's platform rates at K3 launch (July 2026):
+		// $3/$15 per MTok with $0.30 cached input, flat across the 1M-token
+		// window. kimi-k3 is the ONLY model the platform serves newer
+		// accounts — K2-era ids (kimi-k2-thinking) and even the kimi-latest
+		// alias return resource_not_found_error, verified live 2026-07-21 —
+		// so it's the only catalog entry. Grandfathered accounts with K2
+		// access can still type those ids on the provider's model rows.
+		// The consumer app's "K3 Swarm Max" mode is not an API SKU, so it
+		// doesn't appear here.
+		Models: []Model{
+			{ID: "kimi-k3", Label: "Kimi K3", InputPer1k: 0.003, OutputPer1k: 0.015, ContextWindow: 1000000},
+		},
+	},
 	{
 		ID:                 "litellm_proxy",
 		Kind:               KindGateway,

management/internals/modules/agentnetwork/manager.go

@@ -86,12 +86,17 @@ type Manager interface {
 
 // PolicySelectionInput is the per-request selection envelope. The
 // proxy populates it from CapturedData (account, user, groups) plus
-// the provider llm_router resolved.
+// the provider llm_router resolved and the model it extracted.
 type PolicySelectionInput struct {
 	AccountID  string
 	UserID     string
 	GroupIDs   []string
 	ProviderID string
+	// Model is the already-normalised upstream model id the proxy extracted
+	// (parser strips Bedrock region/version, Vertex @version), so a
+	// case-insensitive compare suffices. Empty = undetermined → not permitted
+	// (fail closed).
+	Model string
 }
 
 // PolicySelectionResult names the policy that "pays" for this request

management/internals/modules/agentnetwork/policyselect.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"math"
 	"sort"
+	"strings"
 	"time"
 
 	"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
@@ -35,6 +36,10 @@ const (
 	denyCodeAccountTokenCapExceeded = "llm_account.token_cap_exceeded"
 	//nolint:gosec // account deny code label, not a credential
 	denyCodeAccountBudgetCapExceeded = "llm_account.budget_cap_exceeded"
+	// denyCodeModelBlocked is returned when policies govern the request's
+	// (provider, caller-groups) but none permits the model. Matches the proxy
+	// guardrail's code so both layers surface the same label.
+	denyCodeModelBlocked = "llm_policy.model_blocked"
 )
 
 // consumptionCache holds the consumption counters prefetched for one
@@ -159,6 +164,25 @@ func (m *managerImpl) SelectPolicyForRequest(ctx context.Context, in PolicySelec
 	}
 	candidates := filterApplicablePolicies(policies, in)
 
+	// Model-allowlist gate scoped to the matched policies: keep candidates whose
+	// guardrails permit the model (none enabled = unrestricted), deny when
+	// policies apply but none permits it. Skip the load when none has a guardrail.
+	if len(candidates) > 0 && anyPolicyHasGuardrails(candidates) {
+		guardrailsByID, gErr := m.loadGuardrailsByID(ctx, in.AccountID)
+		if gErr != nil {
+			return nil, gErr
+		}
+		permitted := filterModelPermittedPolicies(candidates, guardrailsByID, in.Model)
+		if len(permitted) == 0 {
+			return &PolicySelectionResult{
+				Allow:      false,
+				DenyCode:   denyCodeModelBlocked,
+				DenyReason: modelBlockedReason(in.Model),
+			}, nil
+		}
+		candidates = permitted
+	}
+
 	// Prefetch every consumption counter the ceiling + candidate policies will
 	// read, in a single store round-trip, then score against the cache.
 	cache, err := m.prefetchConsumption(ctx, in, rules, candidates, now)
@@ -250,6 +274,90 @@ func filterApplicablePolicies(policies []*types.Policy, in PolicySelectionInput)
 	return out
 }
 
+// anyPolicyHasGuardrails reports whether any policy references at least one
+// guardrail, so the selector can skip loading guardrails when none do.
+func anyPolicyHasGuardrails(policies []*types.Policy) bool {
+	for _, p := range policies {
+		if p != nil && len(p.GuardrailIDs) > 0 {
+			return true
+		}
+	}
+	return false
+}
+
+// loadGuardrailsByID loads the account's guardrails indexed by ID. Used by the
+// model-allowlist gate to resolve each candidate policy's attached guardrails.
+func (m *managerImpl) loadGuardrailsByID(ctx context.Context, accountID string) (map[string]*types.Guardrail, error) {
+	guardrails, err := m.store.GetAccountAgentNetworkGuardrails(ctx, store.LockingStrengthNone, accountID)
+	if err != nil {
+		return nil, fmt.Errorf("list account guardrails: %w", err)
+	}
+	byID := make(map[string]*types.Guardrail, len(guardrails))
+	for _, g := range guardrails {
+		if g != nil {
+			byID[g.ID] = g
+		}
+	}
+	return byID, nil
+}
+
+// filterModelPermittedPolicies returns the subset of policies whose guardrails
+// permit the model. Order is preserved so downstream scoring is unaffected.
+func filterModelPermittedPolicies(policies []*types.Policy, byID map[string]*types.Guardrail, model string) []*types.Policy {
+	out := make([]*types.Policy, 0, len(policies))
+	for _, p := range policies {
+		if policyPermitsModel(p, byID, model) {
+			out = append(out, p)
+		}
+	}
+	return out
+}
+
+// policyPermitsModel reports whether a policy permits the model. No
+// allowlist-enabled guardrail = unrestricted (permits any, incl. empty);
+// otherwise the model must be in the union of its allowlists, so an
+// empty/undetermined model fails closed.
+func policyPermitsModel(p *types.Policy, byID map[string]*types.Guardrail, model string) bool {
+	if p == nil {
+		return false
+	}
+	wanted := normaliseModelID(model)
+	restricted := false
+	for _, gID := range p.GuardrailIDs {
+		g, ok := byID[gID]
+		if !ok || g == nil || !g.Checks.ModelAllowlist.Enabled {
+			continue
+		}
+		restricted = true
+		if wanted == "" {
+			continue
+		}
+		for _, allowed := range g.Checks.ModelAllowlist.Models {
+			if normaliseModelID(allowed) == wanted {
+				return true
+			}
+		}
+	}
+	return !restricted
+}
+
+// normaliseModelID lowercases and trims a model identifier so the allowlist
+// compare is case-insensitive and trim-tolerant. Mirrors the proxy guardrail's
+// normaliseModel so both layers agree on what "same model" means.
+func normaliseModelID(model string) string {
+	return strings.ToLower(strings.TrimSpace(model))
+}
+
+// modelBlockedReason builds the human-readable deny reason for a model-allowlist
+// rejection. The model is quoted when known; an undetermined model is reported
+// as such so the access log distinguishes "wrong model" from "no model".
+func modelBlockedReason(model string) string {
+	if normaliseModelID(model) == "" {
+		return "request model could not be determined for the policy allowlist"
+	}
+	return fmt.Sprintf("model %q is not permitted by any applicable policy allowlist", model)
+}
+
 // candidate is the per-policy intermediate the selector ranks. A
 // policy that's been exhausted on any enabled cap never makes it
 // into this slice; the selector's deny envelope carries the latest

management/internals/modules/agentnetwork/policyselect_model_test.go

@@ -0,0 +1,329 @@
+package agentnetwork
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/golang/mock/gomock"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
+	"github.com/netbirdio/netbird/management/server/store"
+)
+
+// guardedPolicy builds an enabled, uncapped policy that authorises sourceGroups
+// to reach providerID under the given guardrails. Uncapped keeps the selector's
+// headroom scoring trivial so these tests isolate the model-allowlist gate.
+func guardedPolicy(id, account string, sourceGroups []string, providerID string, guardrailIDs ...string) *types.Policy {
+	return &types.Policy{
+		ID:                     id,
+		AccountID:              account,
+		Enabled:                true,
+		SourceGroups:           sourceGroups,
+		DestinationProviderIDs: []string{providerID},
+		GuardrailIDs:           guardrailIDs,
+		CreatedAt:              time.Now().UTC(),
+	}
+}
+
+// allowlistGuardrail builds a guardrail whose model allowlist is enabled and
+// carries the given models.
+func allowlistGuardrail(id, account string, models ...string) *types.Guardrail {
+	return &types.Guardrail{
+		ID:        id,
+		AccountID: account,
+		Checks: types.GuardrailChecks{
+			ModelAllowlist: types.GuardrailModelAllowlist{Enabled: true, Models: models},
+		},
+	}
+}
+
+func expectPolicies(mockStore *store.MockStore, account string, policies ...*types.Policy) {
+	mockStore.EXPECT().
+		GetAccountAgentNetworkPolicies(gomock.Any(), gomock.Any(), account).
+		Return(policies, nil)
+}
+
+func expectGuardrails(mockStore *store.MockStore, account string, guardrails ...*types.Guardrail) {
+	mockStore.EXPECT().
+		GetAccountAgentNetworkGuardrails(gomock.Any(), gomock.Any(), account).
+		Return(guardrails, nil)
+}
+
+// TestSelectPolicy_ModelBlockedByAllowlist proves the authoritative allowlist
+// decision: a policy authorises the (provider, group) but restricts the model,
+// and the requested model isn't on the list, so the request is denied.
+func TestSelectPolicy_ModelBlockedByAllowlist(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	mgr, mockStore := newSelectorMgr(t, ctrl)
+
+	policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
+	expectPolicies(mockStore, "acc-1", policy)
+	expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o"))
+
+	res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
+		AccountID:  "acc-1",
+		UserID:     "user-1",
+		GroupIDs:   []string{"grp-eng"},
+		ProviderID: "prov-1",
+		Model:      "claude-opus-4",
+	})
+	require.NoError(t, err)
+	assert.False(t, res.Allow, "a model outside the only applicable policy's allowlist must be denied")
+	assert.Equal(t, denyCodeModelBlocked, res.DenyCode, "deny code must be model_blocked")
+	assert.NotEmpty(t, res.DenyReason, "deny reason must be populated")
+}
+
+// TestSelectPolicy_ModelAllowedByAllowlist is the allow counterpart: the model
+// is on the applicable policy's allowlist, so selection proceeds normally.
+func TestSelectPolicy_ModelAllowedByAllowlist(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	mgr, mockStore := newSelectorMgr(t, ctrl)
+
+	policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
+	expectPolicies(mockStore, "acc-1", policy)
+	expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o", "claude-opus-4"))
+	expectConsumptionBatch(mockStore, nil)
+
+	res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
+		AccountID:  "acc-1",
+		UserID:     "user-1",
+		GroupIDs:   []string{"grp-eng"},
+		ProviderID: "prov-1",
+		Model:      "claude-opus-4",
+	})
+	require.NoError(t, err)
+	assert.True(t, res.Allow, "a model on the applicable policy's allowlist must be allowed")
+	assert.Equal(t, "pol-A", res.SelectedPolicyID)
+}
+
+// TestSelectPolicy_CaseInsensitiveModelMatch proves the compare tolerates case
+// and surrounding whitespace, matching the proxy guardrail's normalisation.
+func TestSelectPolicy_CaseInsensitiveModelMatch(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	mgr, mockStore := newSelectorMgr(t, ctrl)
+
+	policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
+	expectPolicies(mockStore, "acc-1", policy)
+	expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "  GPT-4o  "))
+	expectConsumptionBatch(mockStore, nil)
+
+	res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
+		AccountID:  "acc-1",
+		GroupIDs:   []string{"grp-eng"},
+		ProviderID: "prov-1",
+		Model:      "gpt-4o",
+	})
+	require.NoError(t, err)
+	assert.True(t, res.Allow, "case/whitespace variants must match the allowlist entry")
+}
+
+// TestSelectPolicy_UnguardedPolicyIsUnrestricted is the false-deny fix: when two
+// policies authorise the same (provider, group) and one has no guardrail, that
+// policy makes the request unrestricted — not caught by the other's allowlist.
+func TestSelectPolicy_UnguardedPolicyIsUnrestricted(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	mgr, mockStore := newSelectorMgr(t, ctrl)
+
+	restricted := guardedPolicy("pol-restricted", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
+	open := guardedPolicy("pol-open", "acc-1", []string{"grp-eng"}, "prov-1") // no guardrail
+	expectPolicies(mockStore, "acc-1", restricted, open)
+	expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o"))
+	expectConsumptionBatch(mockStore, nil)
+
+	res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
+		AccountID:  "acc-1",
+		GroupIDs:   []string{"grp-eng"},
+		ProviderID: "prov-1",
+		Model:      "claude-opus-4",
+	})
+	require.NoError(t, err)
+	assert.True(t, res.Allow, "an un-guardrailed policy for the same (provider, group) must leave the request unrestricted")
+	assert.Equal(t, "pol-open", res.SelectedPolicyID, "the unrestricted policy must be the one that pays")
+}
+
+// TestSelectPolicy_AllowlistDoesNotLeakAcrossGroups is the false-allow fix: a
+// model allowlisted only for grp-b must not be usable by a grp-a caller. The
+// selector considers only policies applicable to the caller's groups.
+func TestSelectPolicy_AllowlistDoesNotLeakAcrossGroups(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	mgr, mockStore := newSelectorMgr(t, ctrl)
+
+	polA := guardedPolicy("pol-a", "acc-1", []string{"grp-a"}, "prov-1", "g-a")
+	polB := guardedPolicy("pol-b", "acc-1", []string{"grp-b"}, "prov-1", "g-b")
+	expectPolicies(mockStore, "acc-1", polA, polB)
+	expectGuardrails(mockStore, "acc-1",
+		allowlistGuardrail("g-a", "acc-1", "gpt-4o"),
+		allowlistGuardrail("g-b", "acc-1", "claude-opus-4"),
+	)
+
+	res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
+		AccountID:  "acc-1",
+		GroupIDs:   []string{"grp-a"},
+		ProviderID: "prov-1",
+		Model:      "claude-opus-4", // only allowed for grp-b
+	})
+	require.NoError(t, err)
+	assert.False(t, res.Allow, "grp-b's allowlisted model must not leak to a grp-a caller")
+	assert.Equal(t, denyCodeModelBlocked, res.DenyCode)
+}
+
+// TestSelectPolicy_UndeterminedModelFailsClosed proves the fail-closed contract
+// mirrors the proxy: with a restricted applicable policy and an empty model
+// (e.g. a path-routed shape the parser couldn't map), the request is denied.
+func TestSelectPolicy_UndeterminedModelFailsClosed(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	mgr, mockStore := newSelectorMgr(t, ctrl)
+
+	policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
+	expectPolicies(mockStore, "acc-1", policy)
+	expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o"))
+
+	res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
+		AccountID:  "acc-1",
+		GroupIDs:   []string{"grp-eng"},
+		ProviderID: "prov-1",
+		Model:      "", // undetermined
+	})
+	require.NoError(t, err)
+	assert.False(t, res.Allow, "an undetermined model must fail closed against a restricted policy")
+	assert.Equal(t, denyCodeModelBlocked, res.DenyCode)
+}
+
+// TestSelectPolicy_DisabledAllowlistDoesNotRestrict proves a guardrail whose
+// model allowlist is disabled imposes no model restriction, even though the
+// policy references it.
+func TestSelectPolicy_DisabledAllowlistDoesNotRestrict(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	mgr, mockStore := newSelectorMgr(t, ctrl)
+
+	policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
+	disabled := &types.Guardrail{
+		ID:        "g-1",
+		AccountID: "acc-1",
+		Checks: types.GuardrailChecks{
+			ModelAllowlist: types.GuardrailModelAllowlist{Enabled: false, Models: []string{"gpt-4o"}},
+		},
+	}
+	expectPolicies(mockStore, "acc-1", policy)
+	expectGuardrails(mockStore, "acc-1", disabled)
+	expectConsumptionBatch(mockStore, nil)
+
+	res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
+		AccountID:  "acc-1",
+		GroupIDs:   []string{"grp-eng"},
+		ProviderID: "prov-1",
+		Model:      "anything-goes",
+	})
+	require.NoError(t, err)
+	assert.True(t, res.Allow, "a disabled allowlist must not restrict the model")
+	assert.Equal(t, "pol-A", res.SelectedPolicyID)
+}
+
+// TestSelectPolicy_UnionAcrossPolicyGuardrails proves a policy with multiple
+// allowlist guardrails permits the union of their models (not just the first).
+func TestSelectPolicy_UnionAcrossPolicyGuardrails(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	mgr, mockStore := newSelectorMgr(t, ctrl)
+
+	policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1", "g-2")
+	expectPolicies(mockStore, "acc-1", policy)
+	expectGuardrails(mockStore, "acc-1",
+		allowlistGuardrail("g-1", "acc-1", "gpt-4o"),
+		allowlistGuardrail("g-2", "acc-1", "claude-opus-4"),
+	)
+	expectConsumptionBatch(mockStore, nil)
+
+	res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
+		AccountID:  "acc-1",
+		GroupIDs:   []string{"grp-eng"},
+		ProviderID: "prov-1",
+		Model:      "claude-opus-4", // only in the second guardrail's list
+	})
+	require.NoError(t, err)
+	assert.True(t, res.Allow, "a model in any of the policy's allowlist guardrails must be permitted")
+	assert.Equal(t, "pol-A", res.SelectedPolicyID)
+}
+
+// TestSelectPolicy_GuardrailLookupErrorPropagates proves a store failure while
+// resolving the candidate policies' guardrails surfaces as an error, not a
+// silent allow/deny.
+func TestSelectPolicy_GuardrailLookupErrorPropagates(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	mgr, mockStore := newSelectorMgr(t, ctrl)
+
+	policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
+	expectPolicies(mockStore, "acc-1", policy)
+	mockStore.EXPECT().
+		GetAccountAgentNetworkGuardrails(gomock.Any(), gomock.Any(), "acc-1").
+		Return(nil, errors.New("store unavailable"))
+
+	res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
+		AccountID:  "acc-1",
+		GroupIDs:   []string{"grp-eng"},
+		ProviderID: "prov-1",
+		Model:      "gpt-4o",
+	})
+	require.Error(t, err, "a guardrail-lookup failure must surface as an error")
+	assert.Nil(t, res)
+}
+
+// TestSelectPolicy_MissingGuardrailReferenceTreatedAsUnrestricted proves a
+// policy referencing a guardrail ID absent from the account's set (a stale
+// reference) imposes no model restriction — same as no guardrail.
+func TestSelectPolicy_MissingGuardrailReferenceTreatedAsUnrestricted(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	mgr, mockStore := newSelectorMgr(t, ctrl)
+
+	policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-missing")
+	expectPolicies(mockStore, "acc-1", policy)
+	expectGuardrails(mockStore, "acc-1")
+	expectConsumptionBatch(mockStore, nil)
+
+	res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
+		AccountID:  "acc-1",
+		GroupIDs:   []string{"grp-eng"},
+		ProviderID: "prov-1",
+		Model:      "anything-goes",
+	})
+	require.NoError(t, err)
+	assert.True(t, res.Allow, "an orphaned guardrail reference must not restrict the model")
+	assert.Equal(t, "pol-A", res.SelectedPolicyID)
+}
+
+// TestSelectPolicy_PartialCandidatesPermittedAfterModelFilter proves the model
+// gate narrows candidates before cap scoring: the permitting policy is selected
+// even though the blocked one has a larger, more attractive cap.
+func TestSelectPolicy_PartialCandidatesPermittedAfterModelFilter(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	mgr, mockStore := newSelectorMgr(t, ctrl)
+
+	polBig := guardedPolicy("pol-big", "acc-1", []string{"grp-eng"}, "prov-1", "g-restrict")
+	polBig.Limits = types.PolicyLimits{
+		TokenLimit: types.PolicyTokenLimit{Enabled: true, GroupCap: 1_000_000, WindowSeconds: 3600},
+	}
+	polSmall := guardedPolicy("pol-small", "acc-1", []string{"grp-eng"}, "prov-1", "g-permit")
+	polSmall.Limits = types.PolicyLimits{
+		TokenLimit: types.PolicyTokenLimit{Enabled: true, GroupCap: 100, WindowSeconds: 3600},
+	}
+	expectPolicies(mockStore, "acc-1", polBig, polSmall)
+	expectGuardrails(mockStore, "acc-1",
+		allowlistGuardrail("g-restrict", "acc-1", "gpt-4o"),
+		allowlistGuardrail("g-permit", "acc-1", "claude-opus-4"),
+	)
+	expectConsumptionBatch(mockStore, nil)
+
+	res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
+		AccountID:  "acc-1",
+		GroupIDs:   []string{"grp-eng"},
+		ProviderID: "prov-1",
+		Model:      "claude-opus-4", // only pol-small's guardrail permits this
+	})
+	require.NoError(t, err)
+	assert.True(t, res.Allow)
+	assert.Equal(t, "pol-small", res.SelectedPolicyID,
+		"the model filter must exclude pol-big before cap scoring")
+}

management/internals/modules/agentnetwork/synthesizer.go

@@ -235,7 +235,12 @@ func SynthesizeServices(ctx context.Context, s store.Store, accountID string) ([
 
 	mergedGuardrails := mergeGuardrails(enabledPolicies, guardrailsByID)
 	applyAccountCollectionControls(&mergedGuardrails, settings)
-	guardrailJSON, err := marshalGuardrailConfig(mergedGuardrails)
+	// The proxy guardrail is a per-provider fail-closed backstop; the
+	// authoritative per-policy/group decision is management's
+	// SelectPolicyForRequest. A provider lands in this map only when every
+	// authorising policy restricts models.
+	providerAllowlists := buildProviderAllowlists(enabledPolicies, guardrailsByID)
+	guardrailJSON, err := marshalGuardrailConfig(providerAllowlists, mergedGuardrails.PromptCapture)
 	if err != nil {
 		return nil, err
 	}
@@ -780,10 +785,12 @@ func buildMiddlewareChain(routerCfgJSON, identityInjectJSON, guardrailJSON []byt
 
 // guardrailConfig is the JSON shape the proxy-side llm_guardrail
 // middleware expects. Mirrors the proxy registration documented in
-// the management→proxy contract.
+// the management→proxy contract. provider_allowlists is keyed by the
+// resolved provider id llm_router stamps; a provider absent from the map is
+// unrestricted at the proxy layer.
 type guardrailConfig struct {
-	ModelAllowlist []string               `json:"model_allowlist,omitempty"`
-	PromptCapture  guardrailPromptCapture `json:"prompt_capture"`
+	ProviderAllowlists map[string][]string    `json:"provider_allowlists,omitempty"`
+	PromptCapture      guardrailPromptCapture `json:"prompt_capture"`
 }
 
 type guardrailPromptCapture struct {
@@ -828,13 +835,10 @@ func applyAccountCollectionControls(merged *MergedGuardrails, settings *types.Se
 	merged.PromptCapture.RedactPii = settings.RedactPii || merged.PromptCapture.RedactPii
 }
 
-func marshalGuardrailConfig(merged MergedGuardrails) ([]byte, error) {
+func marshalGuardrailConfig(providerAllowlists map[string][]string, capture MergedPromptCapture) ([]byte, error) {
 	cfg := guardrailConfig{
-		ModelAllowlist: merged.ModelAllowlist,
-		PromptCapture: guardrailPromptCapture{
-			Enabled:   merged.PromptCapture.Enabled,
-			RedactPii: merged.PromptCapture.RedactPii,
-		},
+		ProviderAllowlists: providerAllowlists,
+		PromptCapture:      guardrailPromptCapture(capture),
 	}
 	out, err := json.Marshal(cfg)
 	if err != nil {
@@ -843,6 +847,74 @@ func marshalGuardrailConfig(merged MergedGuardrails) ([]byte, error) {
 	return out, nil
 }
 
+// buildProviderAllowlists returns the proxy's per-provider backstop: a provider
+// is included only when every authorising policy restricts models (their union);
+// if any leaves it unrestricted it is omitted, so management decides per group.
+func buildProviderAllowlists(policies []*types.Policy, byID map[string]*types.Guardrail) map[string][]string {
+	type providerAcc struct {
+		models          map[string]struct{}
+		anyUnrestricted bool
+	}
+	accs := make(map[string]*providerAcc)
+	for _, p := range policies {
+		if p == nil {
+			continue
+		}
+		restricted, models := policyModelAllowlist(p, byID)
+		for _, providerID := range p.DestinationProviderIDs {
+			if providerID == "" {
+				continue
+			}
+			acc, ok := accs[providerID]
+			if !ok {
+				acc = &providerAcc{models: make(map[string]struct{})}
+				accs[providerID] = acc
+			}
+			if !restricted {
+				acc.anyUnrestricted = true
+				continue
+			}
+			for _, m := range models {
+				acc.models[m] = struct{}{}
+			}
+		}
+	}
+	out := make(map[string][]string, len(accs))
+	for providerID, acc := range accs {
+		if acc.anyUnrestricted {
+			continue
+		}
+		models := make([]string, 0, len(acc.models))
+		for m := range acc.models {
+			models = append(models, m)
+		}
+		sort.Strings(models)
+		out[providerID] = models
+	}
+	return out
+}
+
+// policyModelAllowlist reports whether a policy restricts models (has an
+// allowlist-enabled guardrail) and the union of allowed models. Models are
+// verbatim; the proxy factory lowercases/trims them at decode time.
+func policyModelAllowlist(p *types.Policy, byID map[string]*types.Guardrail) (bool, []string) {
+	restricted := false
+	var models []string
+	for _, gID := range p.GuardrailIDs {
+		g, ok := byID[gID]
+		if !ok || g == nil || !g.Checks.ModelAllowlist.Enabled {
+			continue
+		}
+		restricted = true
+		for _, m := range g.Checks.ModelAllowlist.Models {
+			if m != "" {
+				models = append(models, m)
+			}
+		}
+	}
+	return restricted, models
+}
+
 // buildAccountService composes the per-account gateway Service. The
 // target carries the noop placeholder URL — the router middleware
 // rewrites every request to the matched provider's upstream before the
@@ -986,38 +1058,11 @@ func unionSourceGroups(policies []*types.Policy) []string {
 	return out
 }
 
-// MergedGuardrails is the JSON shape passed to the proxy via the
-// guardrail middleware's config_json. Mirrors the proxy-side
-// expectations and is intentionally distinct from
-// types.GuardrailChecks so we can evolve either side independently.
+// MergedGuardrails is the synthesiser's fold target. Only prompt capture is
+// merged here — the model allowlist is emitted per-provider, and
+// token/budget/retention moved onto Policy.Limits and account Settings.
 type MergedGuardrails struct {
-	ModelAllowlist []string            `json:"model_allowlist,omitempty"`
-	TokenLimits    MergedTokenLimits   `json:"token_limits"`
-	Budget         MergedBudget        `json:"budget"`
-	PromptCapture  MergedPromptCapture `json:"prompt_capture"`
-	Retention      MergedRetention     `json:"retention"`
-}
-
-type MergedTokenLimits struct {
-	Hourly  *MergedTokenWindow `json:"hourly,omitempty"`
-	Daily   *MergedTokenWindow `json:"daily,omitempty"`
-	Monthly *MergedTokenWindow `json:"monthly,omitempty"`
-}
-
-type MergedTokenWindow struct {
-	MaxInputTokens  int `json:"max_input_tokens,omitempty"`
-	MaxOutputTokens int `json:"max_output_tokens,omitempty"`
-}
-
-type MergedBudget struct {
-	Hourly  *MergedBudgetWindow `json:"hourly,omitempty"`
-	Daily   *MergedBudgetWindow `json:"daily,omitempty"`
-	Monthly *MergedBudgetWindow `json:"monthly,omitempty"`
-}
-
-type MergedBudgetWindow struct {
-	SoftCapUSD float64 `json:"soft_cap_usd,omitempty"`
-	HardCapUSD float64 `json:"hard_cap_usd,omitempty"`
+	PromptCapture MergedPromptCapture
 }
 
 type MergedPromptCapture struct {
@@ -1025,64 +1070,31 @@ type MergedPromptCapture struct {
 	RedactPii bool `json:"redact_pii"`
 }
 
-type MergedRetention struct {
-	Enabled bool `json:"enabled"`
-	Days    int  `json:"days"`
-}
-
-// mergeGuardrails computes the effective guardrail spec applied at the
-// proxy, given the referencing policies and the account's guardrail
-// catalogue. Policy enabled-ness is the caller's responsibility — only
-// enabled policies should be passed in.
+// mergeGuardrails folds the referencing policies' guardrails into the
+// prompt-capture decision only. The model allowlist is enforced per-policy/group
+// in management and shipped per-provider; token/budget/retention live off
+// guardrails now.
 //
-// Merge rules:
-//   - Model allowlist:   union of allowlists across policies that enable it.
-//   - Token / Budget:    most-restrictive (min of non-zero caps) per window.
-//   - Prompt capture:    enabled if any policy enables it; redact_pii sticks
-//     if any enabling policy turns it on.
-//   - Retention:         enabled if any enables it; smallest non-zero days wins.
+// Merge rule — prompt capture: enabled if any policy enables it; redact_pii
+// sticks if any enabling policy turns it on.
 func mergeGuardrails(policies []*types.Policy, byID map[string]*types.Guardrail) MergedGuardrails {
 	merged := MergedGuardrails{}
-	allowlist := make(map[string]struct{})
-	allowlistEnabled := false
-
 	for _, policy := range policies {
 		for _, gID := range policy.GuardrailIDs {
 			g, ok := byID[gID]
 			if !ok || g == nil {
 				continue
 			}
-			mergeGuardrail(g, &merged, allowlist, &allowlistEnabled)
+			mergeGuardrail(g, &merged)
 		}
 	}
-
-	if allowlistEnabled {
-		merged.ModelAllowlist = make([]string, 0, len(allowlist))
-		for m := range allowlist {
-			merged.ModelAllowlist = append(merged.ModelAllowlist, m)
-		}
-		sort.Strings(merged.ModelAllowlist)
-	}
 	return merged
 }
 
-// mergeGuardrail folds a single guardrail's enabled checks into the
-// running merge: model-allowlist models join the shared set (and flip
-// allowlistEnabled), and prompt-capture / redact-pii stick once any
-// enabling guardrail turns them on.
-//
-// TokenLimits, Budget, and Retention have moved off guardrails — token
-// and budget caps now live on the Policy itself (Policy.Limits) and
-// retention moves to account-level Settings — so they are not merged here.
-func mergeGuardrail(g *types.Guardrail, merged *MergedGuardrails, allowlist map[string]struct{}, allowlistEnabled *bool) {
-	if g.Checks.ModelAllowlist.Enabled {
-		*allowlistEnabled = true
-		for _, m := range g.Checks.ModelAllowlist.Models {
-			if m != "" {
-				allowlist[m] = struct{}{}
-			}
-		}
-	}
+// mergeGuardrail folds a single guardrail's prompt-capture settings into the
+// running merge: enabled / redact-pii stick once any enabling guardrail turns
+// them on.
+func mergeGuardrail(g *types.Guardrail, merged *MergedGuardrails) {
 	if g.Checks.PromptCapture.Enabled {
 		merged.PromptCapture.Enabled = true
 		if g.Checks.PromptCapture.RedactPii {

management/internals/modules/agentnetwork/synthesizer_guardrail_realstore_test.go

@@ -81,8 +81,8 @@ func TestSynthesizeServices_RealStore_PromptCaptureAccountIsSoleControl(t *testi
 	require.Len(t, services, 1)
 
 	cfg := decodeServiceGuardrailConfig(t, services[0])
-	assert.Equal(t, []string{"gpt-5.4"}, cfg.ModelAllowlist,
-		"model allowlist is a pure policy guardrail and must always reach the config")
+	assert.Equal(t, map[string][]string{"prov-1": {"gpt-5.4"}}, cfg.ProviderAllowlists,
+		"model allowlist is a pure policy guardrail and must reach the per-provider config")
 	assert.False(t, cfg.PromptCapture.Enabled,
 		"prompt capture must be off when the account toggle is off, even with a capture-enabled guardrail")
 }
@@ -172,7 +172,7 @@ func TestSynthesizeServices_RealStore_NoGuardrail_CaptureOff(t *testing.T) {
 	require.Len(t, services, 1, "exactly one synth service expected")
 
 	cfg := decodeServiceGuardrailConfig(t, services[0])
-	assert.Empty(t, cfg.ModelAllowlist, "no guardrail → no allowlist")
+	assert.Empty(t, cfg.ProviderAllowlists, "no guardrail → provider unrestricted (absent from map)")
 	assert.False(t, cfg.PromptCapture.Enabled, "no guardrail → prompt capture off by default")
 	assert.False(t, cfg.PromptCapture.RedactPii, "no guardrail → redact off by default")
 }

management/internals/modules/agentnetwork/synthesizer_provider_allowlist_test.go

@@ -0,0 +1,95 @@
+package agentnetwork
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
+)
+
+// policyForProviders builds an enabled policy authorising the given providers
+// under the given guardrails (both optional). Groups are irrelevant to
+// buildProviderAllowlists, which keys purely on destination provider.
+func policyForProviders(id string, guardrailIDs []string, providerIDs ...string) *types.Policy {
+	return &types.Policy{
+		ID:                     id,
+		Enabled:                true,
+		DestinationProviderIDs: providerIDs,
+		GuardrailIDs:           guardrailIDs,
+	}
+}
+
+func TestBuildProviderAllowlists(t *testing.T) {
+	byID := map[string]*types.Guardrail{
+		"g-4o":       allowlistGuardrail("g-4o", "acc-1", "gpt-4o"),
+		"g-opus":     allowlistGuardrail("g-opus", "acc-1", "claude-opus-4"),
+		"g-disabled": {ID: "g-disabled", Checks: types.GuardrailChecks{ModelAllowlist: types.GuardrailModelAllowlist{Enabled: false, Models: []string{"gpt-4o"}}}},
+	}
+
+	t.Run("all authorising policies restrict yields per-provider union", func(t *testing.T) {
+		policies := []*types.Policy{
+			policyForProviders("p1", []string{"g-4o"}, "prov-x"),
+			policyForProviders("p2", []string{"g-opus"}, "prov-x"),
+		}
+		got := buildProviderAllowlists(policies, byID)
+		assert.Equal(t, map[string][]string{"prov-x": {"claude-opus-4", "gpt-4o"}}, got,
+			"a provider every policy restricts carries the sorted union of their models")
+	})
+
+	t.Run("any un-guardrailed policy leaves the provider unrestricted (omitted)", func(t *testing.T) {
+		policies := []*types.Policy{
+			policyForProviders("p1", []string{"g-4o"}, "prov-x"),
+			policyForProviders("p2", nil, "prov-x"), // no guardrail
+		}
+		got := buildProviderAllowlists(policies, byID)
+		assert.NotContains(t, got, "prov-x",
+			"a provider reachable by an un-guardrailed policy must be omitted so the proxy treats it as unrestricted")
+	})
+
+	t.Run("a disabled allowlist counts as unrestricted", func(t *testing.T) {
+		policies := []*types.Policy{
+			policyForProviders("p1", []string{"g-disabled"}, "prov-x"),
+		}
+		got := buildProviderAllowlists(policies, byID)
+		assert.NotContains(t, got, "prov-x",
+			"a policy whose only guardrail has a disabled allowlist is unrestricted")
+	})
+
+	t.Run("providers are isolated from one another", func(t *testing.T) {
+		policies := []*types.Policy{
+			policyForProviders("p1", []string{"g-4o"}, "prov-x"),
+			policyForProviders("p2", []string{"g-opus"}, "prov-y"),
+		}
+		got := buildProviderAllowlists(policies, byID)
+		assert.Equal(t, []string{"gpt-4o"}, got["prov-x"], "prov-x keeps only its own model")
+		assert.Equal(t, []string{"claude-opus-4"}, got["prov-y"], "prov-y keeps only its own model")
+	})
+
+	t.Run("one policy authorising two providers restricts both", func(t *testing.T) {
+		policies := []*types.Policy{
+			policyForProviders("p1", []string{"g-4o"}, "prov-x", "prov-y"),
+		}
+		got := buildProviderAllowlists(policies, byID)
+		assert.Equal(t, []string{"gpt-4o"}, got["prov-x"])
+		assert.Equal(t, []string{"gpt-4o"}, got["prov-y"])
+	})
+
+	t.Run("union across a single policy's guardrails", func(t *testing.T) {
+		policies := []*types.Policy{
+			policyForProviders("p1", []string{"g-4o", "g-opus"}, "prov-x"),
+		}
+		got := buildProviderAllowlists(policies, byID)
+		assert.ElementsMatch(t, []string{"claude-opus-4", "gpt-4o"}, got["prov-x"],
+			"a policy's own multiple allowlist guardrails union together")
+	})
+
+	t.Run("an enabled allowlist with no models denies everything", func(t *testing.T) {
+		empty := map[string]*types.Guardrail{"g-empty": allowlistGuardrail("g-empty", "acc-1")}
+		got := buildProviderAllowlists([]*types.Policy{
+			policyForProviders("p1", []string{"g-empty"}, "prov-x"),
+		}, empty)
+		assert.Equal(t, map[string][]string{"prov-x": {}}, got,
+			"an enabled-but-empty allowlist is restricted with an empty set, not unrestricted")
+	})
+}

management/internals/modules/agentnetwork/synthesizer_test.go

@@ -1031,8 +1031,12 @@ func TestSynthesizeServices_GuardrailMerge_AllowlistUnion_LimitsRestrictive(t *t
 
 	var cfg guardrailConfig
 	require.NoError(t, json.Unmarshal(guardrailJSON, &cfg), "guardrail config must unmarshal cleanly")
-	assert.ElementsMatch(t, []string{"gpt-5.4-mini", "gpt-5.4-pro"}, cfg.ModelAllowlist,
-		"model allowlist union must keep both models")
+	// Both policies restrict the same provider, so the per-provider backstop
+	// carries the union of their models — a coarse gate that management's
+	// per-policy/group check narrows; it only blocks models outside the union
+	// when management is down.
+	assert.ElementsMatch(t, []string{"gpt-5.4-mini", "gpt-5.4-pro"}, cfg.ProviderAllowlists["prov-1"],
+		"per-provider allowlist union must keep both models")
 }
 
 func TestSynthesizeServices_BackfillsMissingSessionKeys(t *testing.T) {

management/internals/modules/agentnetwork/types/accesslog.go

@@ -41,8 +41,24 @@ type AgentNetworkAccessLog struct {
 	InputTokens        int64
 	OutputTokens       int64
 	TotalTokens        int64
-	CostUSD            float64
-	Stream             bool
+	// Prompt-cache buckets: read + write token counts.
+	CachedInputTokens   int64
+	CacheCreationTokens int64
+	// Per-bucket cost breakdown — one column per token bucket the provider
+	// bills separately. These four are the only cost state stored: the total
+	// and the cache portion are derived on read (TotalCostUSD / CacheCostUSD)
+	// rather than stored alongside, so a stored aggregate can never drift out
+	// of step with the components it summarises.
+	//
+	// default:0 matters on upgrade: these columns are ALTER TABLE ADD COLUMN
+	// on an existing table, and without it every historical row holds NULL —
+	// which a raw SUM()/scan into float64 can't read. The default backfills
+	// them as 0, so pre-upgrade rows report an unknown split, not an error.
+	InputCostUSD         float64 `gorm:"not null;default:0"`
+	CachedInputCostUSD   float64 `gorm:"not null;default:0"`
+	CacheCreationCostUSD float64 `gorm:"not null;default:0"`
+	OutputCostUSD        float64 `gorm:"not null;default:0"`
+	Stream               bool
 
 	// Prompt capture. Only populated when prompt collection is enabled
 	// (account master switch AND policy guardrail). Heavy free text.
@@ -60,19 +76,44 @@ type AgentNetworkAccessLog struct {
 // the reverse-proxy AccessLogEntry table.
 func (AgentNetworkAccessLog) TableName() string { return "agent_network_access_log" }
 
+// CostUSDSQLExpr is the SQL sum of the per-bucket cost columns — the total cost
+// of a row. Used wherever a query has to sort or aggregate on total cost now
+// that no cost_usd column is stored. Plain arithmetic over NOT NULL columns, so
+// it stays portable across SQLite and Postgres.
+const CostUSDSQLExpr = "(input_cost_usd + cached_input_cost_usd + cache_creation_cost_usd + output_cost_usd)"
+
+// TotalCostUSD is the request's total cost: the sum of the four per-bucket
+// costs. Derived rather than stored so it cannot disagree with the breakdown.
+func (a *AgentNetworkAccessLog) TotalCostUSD() float64 {
+	return a.InputCostUSD + a.CachedInputCostUSD + a.CacheCreationCostUSD + a.OutputCostUSD
+}
+
+// CacheCostUSD is the portion of the total billed for prompt-cache buckets:
+// cache reads plus cache writes.
+func (a *AgentNetworkAccessLog) CacheCostUSD() float64 {
+	return a.CachedInputCostUSD + a.CacheCreationCostUSD
+}
+
 // ToAPIResponse renders the flattened entry as the API representation.
 func (a *AgentNetworkAccessLog) ToAPIResponse() api.AgentNetworkAccessLog {
 	out := api.AgentNetworkAccessLog{
-		Id:           a.ID,
-		ServiceId:    a.ServiceID,
-		Timestamp:    a.Timestamp,
-		StatusCode:   a.StatusCode,
-		DurationMs:   int(a.Duration.Milliseconds()),
-		InputTokens:  a.InputTokens,
-		OutputTokens: a.OutputTokens,
-		TotalTokens:  a.TotalTokens,
-		CostUsd:      a.CostUSD,
-		Stream:       &a.Stream,
+		Id:                   a.ID,
+		ServiceId:            a.ServiceID,
+		Timestamp:            a.Timestamp,
+		StatusCode:           a.StatusCode,
+		DurationMs:           int(a.Duration.Milliseconds()),
+		InputTokens:          a.InputTokens,
+		OutputTokens:         a.OutputTokens,
+		TotalTokens:          a.TotalTokens,
+		CachedInputTokens:    a.CachedInputTokens,
+		CacheCreationTokens:  a.CacheCreationTokens,
+		InputCostUsd:         a.InputCostUSD,
+		CachedInputCostUsd:   a.CachedInputCostUSD,
+		CacheCreationCostUsd: a.CacheCreationCostUSD,
+		OutputCostUsd:        a.OutputCostUSD,
+		CostUsd:              a.TotalCostUSD(),
+		CacheCostUsd:         a.CacheCostUSD(),
+		Stream:               &a.Stream,
 	}
 
 	out.UserId = strPtr(a.UserID)
@@ -112,20 +153,36 @@ func strPtr(s string) *string {
 // summary plus its ordered entries. Assembled in Go from a page of entries — it
 // is not a stored table.
 type AgentNetworkAccessLogSession struct {
-	SessionID    string // empty for a session-less (singleton) request
-	UserID       string
-	GroupIDs     []string // union of the entries' authorising groups
-	StartedAt    time.Time
-	EndedAt      time.Time
-	RequestCount int
-	InputTokens  int64
-	OutputTokens int64
-	TotalTokens  int64
-	CostUSD      float64
-	Providers    []string // distinct vendors seen in the session
-	Models       []string // distinct models seen in the session
-	Decision     string   // "deny" if any entry was denied, else "allow"
-	Entries      []*AgentNetworkAccessLog
+	SessionID            string // empty for a session-less (singleton) request
+	UserID               string
+	GroupIDs             []string // union of the entries' authorising groups
+	StartedAt            time.Time
+	EndedAt              time.Time
+	RequestCount         int
+	InputTokens          int64
+	OutputTokens         int64
+	TotalTokens          int64
+	CachedInputTokens    int64
+	CacheCreationTokens  int64
+	InputCostUSD         float64
+	CachedInputCostUSD   float64
+	CacheCreationCostUSD float64
+	OutputCostUSD        float64
+	Providers            []string // distinct vendors seen in the session
+	Models               []string // distinct models seen in the session
+	Decision             string   // "deny" if any entry was denied, else "allow"
+	Entries              []*AgentNetworkAccessLog
+}
+
+// TotalCostUSD is the session's total cost: the sum of the four per-bucket
+// costs accumulated across its entries.
+func (sess *AgentNetworkAccessLogSession) TotalCostUSD() float64 {
+	return sess.InputCostUSD + sess.CachedInputCostUSD + sess.CacheCreationCostUSD + sess.OutputCostUSD
+}
+
+// CacheCostUSD is the session's prompt-cache spend: cache reads plus writes.
+func (sess *AgentNetworkAccessLogSession) CacheCostUSD() float64 {
+	return sess.CachedInputCostUSD + sess.CacheCreationCostUSD
 }
 
 // sessionKey is the grouping key for an entry: its session id, or — when the
@@ -205,7 +262,12 @@ func (sess *AgentNetworkAccessLogSession) foldEntry(sk *sessionSeen, e *AgentNet
 	sess.InputTokens += e.InputTokens
 	sess.OutputTokens += e.OutputTokens
 	sess.TotalTokens += e.TotalTokens
-	sess.CostUSD += e.CostUSD
+	sess.CachedInputTokens += e.CachedInputTokens
+	sess.CacheCreationTokens += e.CacheCreationTokens
+	sess.InputCostUSD += e.InputCostUSD
+	sess.CachedInputCostUSD += e.CachedInputCostUSD
+	sess.CacheCreationCostUSD += e.CacheCreationCostUSD
+	sess.OutputCostUSD += e.OutputCostUSD
 	if e.Timestamp.Before(sess.StartedAt) {
 		sess.StartedAt = e.Timestamp
 	}
@@ -248,15 +310,22 @@ func (sess *AgentNetworkAccessLogSession) ToAPIResponse() api.AgentNetworkAccess
 	}
 
 	out := api.AgentNetworkAccessLogSession{
-		StartedAt:    sess.StartedAt,
-		EndedAt:      sess.EndedAt,
-		RequestCount: sess.RequestCount,
-		InputTokens:  sess.InputTokens,
-		OutputTokens: sess.OutputTokens,
-		TotalTokens:  sess.TotalTokens,
-		CostUsd:      sess.CostUSD,
-		Decision:     sess.Decision,
-		Entries:      entries,
+		StartedAt:            sess.StartedAt,
+		EndedAt:              sess.EndedAt,
+		RequestCount:         sess.RequestCount,
+		InputTokens:          sess.InputTokens,
+		OutputTokens:         sess.OutputTokens,
+		TotalTokens:          sess.TotalTokens,
+		CachedInputTokens:    sess.CachedInputTokens,
+		CacheCreationTokens:  sess.CacheCreationTokens,
+		InputCostUsd:         sess.InputCostUSD,
+		CachedInputCostUsd:   sess.CachedInputCostUSD,
+		CacheCreationCostUsd: sess.CacheCreationCostUSD,
+		OutputCostUsd:        sess.OutputCostUSD,
+		CostUsd:              sess.TotalCostUSD(),
+		CacheCostUsd:         sess.CacheCostUSD(),
+		Decision:             sess.Decision,
+		Entries:              entries,
 	}
 	out.SessionId = strPtr(sess.SessionID)
 	out.UserId = strPtr(sess.UserID)

management/internals/modules/agentnetwork/types/accesslogfilter.go

@@ -54,7 +54,7 @@ var accessLogSortFields = map[string]string{
 	"provider":     "provider",
 	"status_code":  "status_code",
 	"duration":     "duration",
-	"cost_usd":     "cost_usd",
+	"cost_usd":     CostUSDSQLExpr,
 	"total_tokens": "total_tokens",
 	"user_id":      "user_id",
 	"decision":     "decision",
@@ -70,7 +70,7 @@ var accessLogSortFields = map[string]string{
 var sessionSortExprs = map[string]string{ //nolint:gosec // G101 false positive: "total_tokens" sort key, not a credential
 	"timestamp":     "MAX(timestamp)",
 	"started_at":    "MIN(timestamp)",
-	"cost_usd":      "SUM(cost_usd)",
+	"cost_usd":      "SUM" + CostUSDSQLExpr,
 	"total_tokens":  "SUM(total_tokens)",
 	"duration":      "SUM(duration)",
 	"request_count": "COUNT(*)",

management/internals/modules/agentnetwork/types/cost_test.go

@@ -0,0 +1,124 @@
+package types
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// costRow builds an access-log entry carrying only a cost breakdown — the rest
+// of the row is irrelevant to the summation identities under test.
+func costRow(id, session string, ts time.Time, in, cachedIn, cacheCreate, out float64) *AgentNetworkAccessLog {
+	return &AgentNetworkAccessLog{
+		ID:                   id,
+		SessionID:            session,
+		Timestamp:            ts,
+		InputCostUSD:         in,
+		CachedInputCostUSD:   cachedIn,
+		CacheCreationCostUSD: cacheCreate,
+		OutputCostUSD:        out,
+	}
+}
+
+// TestAPIResponse_CostComponentsSumToAggregates is the contract a client adding
+// up an API response depends on: within a single rendered object, the four
+// per-bucket fields sum to cost_usd, and the two cache fields sum to
+// cache_cost_usd. Uses rates that are not exactly representable in binary
+// floating point, so the identity is checked against real arithmetic rather
+// than round numbers.
+func TestAPIResponse_CostComponentsSumToAggregates(t *testing.T) {
+	row := costRow("r1", "s1", time.Now(), 0.000768, 0.0002304, 0.00192, 0.003)
+
+	api := row.ToAPIResponse()
+	assert.InDelta(t, api.InputCostUsd+api.CachedInputCostUsd+api.CacheCreationCostUsd+api.OutputCostUsd,
+		api.CostUsd, 1e-12, "rendered buckets must sum to the rendered cost_usd")
+	assert.InDelta(t, api.CachedInputCostUsd+api.CacheCreationCostUsd, api.CacheCostUsd, 1e-12,
+		"rendered cache buckets must sum to the rendered cache_cost_usd")
+	assert.InDelta(t, 0.0059184, api.CostUsd, 1e-12, "total is the exact sum, not a separately rounded figure")
+	assert.InDelta(t, 0.0021504, api.CacheCostUsd, 1e-12, "cache cost is the exact sum of the two cache buckets")
+}
+
+// TestSessionSummary_SumsMatchSummedEntries proves a session summary equals the
+// sum of the entries it renders: a client that adds up the entries itself must
+// land on the same number the summary reports, per bucket and in total.
+func TestSessionSummary_SumsMatchSummedEntries(t *testing.T) {
+	base := time.Date(2026, 5, 5, 10, 0, 0, 0, time.UTC)
+	entries := []*AgentNetworkAccessLog{
+		costRow("r1", "s1", base, 0.000768, 0.0002304, 0.00192, 0.003),
+		costRow("r2", "s1", base.Add(time.Minute), 0.000625, 0.0009375, 0, 0.005),
+		costRow("r3", "s1", base.Add(2*time.Minute), 0.0000016, 0, 0, 0.0000032),
+	}
+
+	sessions := FoldAccessLogSessions([]string{"s1"}, entries)
+	require.Len(t, sessions, 1)
+	sess := sessions[0].ToAPIResponse()
+
+	var wantInput, wantCachedInput, wantCacheCreation, wantOutput float64
+	for _, e := range entries {
+		wantInput += e.InputCostUSD
+		wantCachedInput += e.CachedInputCostUSD
+		wantCacheCreation += e.CacheCreationCostUSD
+		wantOutput += e.OutputCostUSD
+	}
+
+	assert.InDelta(t, wantInput, sess.InputCostUsd, 1e-12, "session input cost is the sum of its entries")
+	assert.InDelta(t, wantCachedInput, sess.CachedInputCostUsd, 1e-12, "session cache-read cost is the sum of its entries")
+	assert.InDelta(t, wantCacheCreation, sess.CacheCreationCostUsd, 1e-12, "session cache-write cost is the sum of its entries")
+	assert.InDelta(t, wantOutput, sess.OutputCostUsd, 1e-12, "session output cost is the sum of its entries")
+	assert.InDelta(t, wantInput+wantCachedInput+wantCacheCreation+wantOutput, sess.CostUsd, 1e-12,
+		"session total equals the summed entry buckets")
+
+	// Summing the rendered entries must give the same answer as reading the
+	// summary — the property a UI relies on when it totals a table itself.
+	var fromEntries float64
+	for _, e := range sess.Entries {
+		fromEntries += e.CostUsd
+	}
+	assert.InDelta(t, sess.CostUsd, fromEntries, 1e-12, "summary total must match the summed rendered entries")
+
+	// The sub-microdollar row must still contribute; it would vanish under
+	// 6-decimal quantisation.
+	assert.Greater(t, sess.InputCostUsd, 0.001393, "small-cost rows must not be quantised away")
+}
+
+// TestUsageBuckets_SumsMatchSummedRows proves the same identity one level up:
+// a usage bucket equals the sum of the ledger rows folded into it, and the
+// buckets together equal the whole range.
+func TestUsageBuckets_SumsMatchSummedRows(t *testing.T) {
+	day1 := time.Date(2026, 5, 5, 9, 0, 0, 0, time.UTC)
+	day2 := time.Date(2026, 5, 6, 9, 0, 0, 0, time.UTC)
+	rows := []*AgentNetworkUsage{
+		{ID: "u1", Timestamp: day1, InputCostUSD: 0.000768, CachedInputCostUSD: 0.0002304, CacheCreationCostUSD: 0.00192, OutputCostUSD: 0.003},
+		{ID: "u2", Timestamp: day1.Add(time.Hour), InputCostUSD: 0.000625, CachedInputCostUSD: 0.0009375, OutputCostUSD: 0.005},
+		{ID: "u3", Timestamp: day2, InputCostUSD: 0.0000016, OutputCostUSD: 0.0000032},
+	}
+
+	buckets := AggregateUsageByGranularity(rows, UsageGranularityDay)
+	require.Len(t, buckets, 2, "two distinct days expected")
+
+	var total, cache float64
+	for _, b := range buckets {
+		api := b.ToAPIResponse()
+		assert.InDelta(t, api.InputCostUsd+api.CachedInputCostUsd+api.CacheCreationCostUsd+api.OutputCostUsd,
+			api.CostUsd, 1e-12, "each bucket's components must sum to its cost_usd")
+		total += api.CostUsd
+		cache += api.CacheCostUsd
+	}
+
+	var wantTotal, wantCache float64
+	for _, r := range rows {
+		wantTotal += r.TotalCostUSD()
+		wantCache += r.CacheCostUSD()
+	}
+	assert.InDelta(t, wantTotal, total, 1e-12, "buckets must sum to the total across all ledger rows")
+	assert.InDelta(t, wantCache, cache, 1e-12, "buckets must sum to the cache spend across all ledger rows")
+
+	// A month bucket over the same rows must total identically — regrouping
+	// changes the partition, never the sum.
+	monthly := AggregateUsageByGranularity(rows, UsageGranularityMonth)
+	require.Len(t, monthly, 1)
+	assert.InDelta(t, wantTotal, monthly[0].ToAPIResponse().CostUsd, 1e-12,
+		"re-bucketing at a different granularity must preserve the total")
+}

management/internals/modules/agentnetwork/types/usage.go

@@ -25,8 +25,19 @@ type AgentNetworkUsage struct {
 	InputTokens        int64
 	OutputTokens       int64
 	TotalTokens        int64
-	CostUSD            float64
-	CreatedAt          time.Time
+	// Prompt-cache buckets: read + write token counts.
+	CachedInputTokens   int64
+	CacheCreationTokens int64
+	// Per-bucket cost breakdown, mirroring AgentNetworkAccessLog — the only
+	// cost state stored; total and cache portion are derived on read. Kept on
+	// the usage ledger too so spend can be attributed per bucket even for
+	// accounts with log collection turned off. See AgentNetworkAccessLog for
+	// why the columns carry a zero default.
+	InputCostUSD         float64 `gorm:"not null;default:0"`
+	CachedInputCostUSD   float64 `gorm:"not null;default:0"`
+	CacheCreationCostUSD float64 `gorm:"not null;default:0"`
+	OutputCostUSD        float64 `gorm:"not null;default:0"`
+	CreatedAt            time.Time
 }
 
 // TableName keeps usage records in their own stripped table. Named
@@ -34,6 +45,17 @@ type AgentNetworkUsage struct {
 // agent_network_usage table in a shared database.
 func (AgentNetworkUsage) TableName() string { return "agent_network_request_usage" }
 
+// TotalCostUSD is the request's total cost: the sum of the four per-bucket
+// costs. Derived rather than stored so it cannot disagree with the breakdown.
+func (u *AgentNetworkUsage) TotalCostUSD() float64 {
+	return u.InputCostUSD + u.CachedInputCostUSD + u.CacheCreationCostUSD + u.OutputCostUSD
+}
+
+// CacheCostUSD is the portion of the total billed for prompt-cache buckets.
+func (u *AgentNetworkUsage) CacheCostUSD() float64 {
+	return u.CachedInputCostUSD + u.CacheCreationCostUSD
+}
+
 // AgentNetworkUsageGroup is the normalised many-to-many row linking a usage
 // record to one authorising group, mirroring AgentNetworkAccessLogGroup so the
 // usage overview can filter by group with a `group_id IN (...)` join.

management/internals/modules/agentnetwork/types/usageoverview.go

@@ -33,21 +33,45 @@ func ParseUsageGranularity(s string) UsageGranularity {
 // AgentNetworkUsageBucket is one aggregated usage time bucket. PeriodStart is
 // the UTC start of the bucket as YYYY-MM-DD.
 type AgentNetworkUsageBucket struct {
-	PeriodStart  string
-	InputTokens  int64
-	OutputTokens int64
-	TotalTokens  int64
-	CostUSD      float64
+	PeriodStart          string
+	InputTokens          int64
+	OutputTokens         int64
+	TotalTokens          int64
+	CachedInputTokens    int64
+	CacheCreationTokens  int64
+	InputCostUSD         float64
+	CachedInputCostUSD   float64
+	CacheCreationCostUSD float64
+	OutputCostUSD        float64
+}
+
+// TotalCostUSD is the bucket's total spend: the sum of the four per-bucket
+// costs. Derived rather than accumulated separately so it cannot disagree with
+// the components.
+func (b *AgentNetworkUsageBucket) TotalCostUSD() float64 {
+	return b.InputCostUSD + b.CachedInputCostUSD + b.CacheCreationCostUSD + b.OutputCostUSD
+}
+
+// CacheCostUSD is the bucket's prompt-cache spend: cache reads plus writes.
+func (b *AgentNetworkUsageBucket) CacheCostUSD() float64 {
+	return b.CachedInputCostUSD + b.CacheCreationCostUSD
 }
 
 // ToAPIResponse renders the bucket as the API representation.
 func (b *AgentNetworkUsageBucket) ToAPIResponse() api.AgentNetworkUsageBucket {
 	return api.AgentNetworkUsageBucket{
-		PeriodStart:  b.PeriodStart,
-		InputTokens:  b.InputTokens,
-		OutputTokens: b.OutputTokens,
-		TotalTokens:  b.TotalTokens,
-		CostUsd:      b.CostUSD,
+		PeriodStart:          b.PeriodStart,
+		InputTokens:          b.InputTokens,
+		OutputTokens:         b.OutputTokens,
+		TotalTokens:          b.TotalTokens,
+		CachedInputTokens:    b.CachedInputTokens,
+		CacheCreationTokens:  b.CacheCreationTokens,
+		InputCostUsd:         b.InputCostUSD,
+		CachedInputCostUsd:   b.CachedInputCostUSD,
+		CacheCreationCostUsd: b.CacheCreationCostUSD,
+		OutputCostUsd:        b.OutputCostUSD,
+		CostUsd:              b.TotalCostUSD(),
+		CacheCostUsd:         b.CacheCostUSD(),
 	}
 }
 
@@ -84,7 +108,12 @@ func AggregateUsageByGranularity(rows []*AgentNetworkUsage, g UsageGranularity)
 		b.InputTokens += r.InputTokens
 		b.OutputTokens += r.OutputTokens
 		b.TotalTokens += r.TotalTokens
-		b.CostUSD += r.CostUSD
+		b.CachedInputTokens += r.CachedInputTokens
+		b.CacheCreationTokens += r.CacheCreationTokens
+		b.InputCostUSD += r.InputCostUSD
+		b.CachedInputCostUSD += r.CachedInputCostUSD
+		b.CacheCreationCostUSD += r.CacheCreationCostUSD
+		b.OutputCostUSD += r.OutputCostUSD
 	}
 
 	out := make([]*AgentNetworkUsageBucket, 0, len(byPeriod))