Merge remote-tracking branch 'origin/main' into refactor/permissions-manager

management/server/store/file_store.go

@@ -274,3 +274,15 @@ func (s *FileStore) SetFieldEncrypt(_ *crypt.FieldEncrypt) {
 func (s *FileStore) GetCustomDomainsCounts(_ context.Context) (int64, int64, error) {
 	return 0, 0, nil
 }
+
+// GetProxyMetrics is a no-op for FileStore — proxy/cluster state isn't
+// persisted in the JSON file format.
+func (s *FileStore) GetProxyMetrics(_ context.Context) (ProxyMetrics, error) {
+	return ProxyMetrics{}, nil
+}
+
+// GetAgentNetworkMetrics is a no-op for FileStore — agent-network state isn't
+// persisted in the JSON file format.
+func (s *FileStore) GetAgentNetworkMetrics(_ context.Context) (AgentNetworkMetrics, error) {
+	return AgentNetworkMetrics{}, nil
+}

management/server/store/sql_store_activity_test.go

@@ -0,0 +1,122 @@
+package store
+
+import (
+	"context"
+	"net/netip"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	nbpeer "github.com/netbirdio/netbird/management/server/peer"
+	"github.com/netbirdio/netbird/management/server/types"
+)
+
+const activityAccountID = "activityAccountId"
+
+func newActivityTestStore(t *testing.T) Store {
+	t.Helper()
+
+	store, cleanUp, err := NewTestStoreFromSQL(context.Background(), "", t.TempDir())
+	require.NoError(t, err)
+	t.Cleanup(cleanUp)
+
+	require.NoError(t, store.SaveAccount(context.Background(), &types.Account{
+		Id:        activityAccountID,
+		Domain:    "activity.example.com",
+		CreatedAt: time.Now().UTC(),
+	}))
+
+	return store
+}
+
+func TestRefreshPeerLastSeen(t *testing.T) {
+	ctx := context.Background()
+	store := newActivityTestStore(t)
+	stored := time.Now().UTC().Add(-3 * time.Hour)
+	require.NoError(t, store.AddPeerToAccount(ctx, activityPeer(stored)))
+
+	refreshed, err := store.RefreshPeerLastSeen(ctx, activityAccountID, "activityPeer", time.Now().UTC().Add(-time.Hour))
+	require.NoError(t, err)
+	assert.True(t, refreshed, "a peer seen three hours ago is stale enough to refresh")
+
+	peer, err := store.GetPeerByID(ctx, LockingStrengthNone, activityAccountID, "activityPeer")
+	require.NoError(t, err)
+	assert.WithinDuration(t, time.Now().UTC(), peer.Status.LastSeen.UTC(), time.Minute, "last seen should be stamped at write time")
+	assert.True(t, peer.Status.LastSeen.After(stored), "last seen must move forward")
+}
+
+// TestRefreshPeerLastSeenHonoursCutoff covers the throttle the caller relies on:
+// two concurrent requests both read the same stale peer, but only the statement
+// that still finds LastSeen behind the cutoff writes.
+func TestRefreshPeerLastSeenHonoursCutoff(t *testing.T) {
+	ctx := context.Background()
+	store := newActivityTestStore(t)
+	stored := time.Now().UTC().Add(-10 * time.Minute)
+	require.NoError(t, store.AddPeerToAccount(ctx, activityPeer(stored)))
+
+	refreshed, err := store.RefreshPeerLastSeen(ctx, activityAccountID, "activityPeer", time.Now().UTC().Add(-time.Hour))
+	require.NoError(t, err)
+	assert.False(t, refreshed, "a peer seen inside the interval must not be written")
+
+	peer, err := store.GetPeerByID(ctx, LockingStrengthNone, activityAccountID, "activityPeer")
+	require.NoError(t, err)
+	assert.WithinDuration(t, stored, peer.Status.LastSeen.UTC(), time.Second, "last seen must be left where it was")
+}
+
+// TestRefreshPeerLastSeenRecordsNeverSeenPeer covers the nullable column. Status
+// is an embedded pointer, so a peer stored without one leaves last seen NULL,
+// and NULL loses the cutoff comparison — such a peer would never record its
+// first activity.
+func TestRefreshPeerLastSeenRecordsNeverSeenPeer(t *testing.T) {
+	ctx := context.Background()
+	store := newActivityTestStore(t)
+	stored := activityPeer(time.Time{})
+	stored.Status = nil
+	require.NoError(t, store.AddPeerToAccount(ctx, stored))
+
+	refreshed, err := store.RefreshPeerLastSeen(ctx, activityAccountID, "activityPeer", time.Now().UTC().Add(-time.Hour))
+	require.NoError(t, err)
+	assert.True(t, refreshed, "a peer that was never seen must record its first activity")
+
+	peer, err := store.GetPeerByID(ctx, LockingStrengthNone, activityAccountID, "activityPeer")
+	require.NoError(t, err)
+	assert.WithinDuration(t, time.Now().UTC(), peer.Status.LastSeen.UTC(), time.Minute, "last seen should be stamped at write time")
+}
+
+// TestRefreshPeerLastSeenLeavesSessionStateAlone pins the column boundary: the
+// connected flag and the session token belong to the sync stream that owns the
+// peer's session, and a blind write here would corrupt its fencing. This is why
+// SavePeerStatus is not reused for an activity bump.
+func TestRefreshPeerLastSeenLeavesSessionStateAlone(t *testing.T) {
+	ctx := context.Background()
+	store := newActivityTestStore(t)
+
+	stored := activityPeer(time.Date(2026, 3, 1, 9, 0, 0, 0, time.UTC))
+	stored.Status.Connected = true
+	stored.Status.SessionStartedAt = 1234567890
+	require.NoError(t, store.AddPeerToAccount(ctx, stored))
+
+	refreshed, err := store.RefreshPeerLastSeen(ctx, activityAccountID, "activityPeer", time.Now().UTC().Add(-time.Hour))
+	require.NoError(t, err)
+	require.True(t, refreshed, "the peer is stale enough to refresh")
+
+	peer, err := store.GetPeerByID(ctx, LockingStrengthNone, activityAccountID, "activityPeer")
+	require.NoError(t, err)
+	assert.WithinDuration(t, time.Now().UTC(), peer.Status.LastSeen.UTC(), time.Minute, "last seen should move forward")
+	assert.True(t, peer.Status.Connected, "connected flag must survive an activity write")
+	assert.Equal(t, int64(1234567890), peer.Status.SessionStartedAt, "session token must survive an activity write")
+}
+
+func activityPeer(lastSeen time.Time) *nbpeer.Peer {
+	return &nbpeer.Peer{
+		ID:        "activityPeer",
+		AccountID: activityAccountID,
+		Key:       "activityPeerKey",
+		IP:        netip.MustParseAddr("100.64.0.9"),
+		Name:      "activity-peer",
+		DNSLabel:  "activity-peer",
+		Status:    &nbpeer.PeerStatus{LastSeen: lastSeen},
+	}
+}

management/server/store/sql_store_agentnetwork.go

@@ -0,0 +1,723 @@
+package store
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"math"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+	"gorm.io/gorm"
+	"gorm.io/gorm/clause"
+
+	agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
+	"github.com/netbirdio/netbird/shared/management/status"
+)
+
+// GetAllAgentNetworkProviders returns Agent Network providers across
+// every account. Used by the synthesizer to build the global service map.
+func (s *SqlStore) GetAllAgentNetworkProviders(ctx context.Context, lockStrength LockingStrength) ([]*agentNetworkTypes.Provider, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var providers []*agentNetworkTypes.Provider
+	if result := tx.Find(&providers); result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to get all agent network providers from store: %v", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get all agent network providers from store")
+	}
+
+	for _, provider := range providers {
+		if err := provider.DecryptSensitiveData(s.fieldEncrypt); err != nil {
+			log.WithContext(ctx).Errorf("failed to decrypt agent network provider %s: %v", provider.ID, err)
+			return nil, status.Errorf(status.Internal, "failed to decrypt agent network provider")
+		}
+	}
+
+	return providers, nil
+}
+
+// GetAgentNetworkMetrics returns aggregated agent-network adoption + usage
+// counts for the self-hosted metrics worker. Each value is a single cheap
+// aggregate; token/cost are summed over the always-collected per-request usage
+// ledger (independent of the log-collection toggle) so they reflect real usage.
+func (s *SqlStore) GetAgentNetworkMetrics(ctx context.Context) (AgentNetworkMetrics, error) {
+	var m AgentNetworkMetrics
+	db := s.db.WithContext(ctx)
+
+	// Providers + distinct adopting accounts in one round-trip.
+	provRow := db.Model(&agentNetworkTypes.Provider{}).
+		Select("COUNT(*) AS providers, COUNT(DISTINCT account_id) AS accounts").Row()
+	if err := provRow.Scan(&m.Providers, &m.Accounts); err != nil {
+		return AgentNetworkMetrics{}, fmt.Errorf("scan agent network provider metrics: %w", err)
+	}
+
+	if err := db.Model(&agentNetworkTypes.Policy{}).Count(&m.Policies).Error; err != nil {
+		return AgentNetworkMetrics{}, fmt.Errorf("count agent network policies: %w", err)
+	}
+
+	if err := db.Model(&agentNetworkTypes.AccountBudgetRule{}).Count(&m.BudgetRules).Error; err != nil {
+		return AgentNetworkMetrics{}, fmt.Errorf("count agent network budget rules: %w", err)
+	}
+
+	if err := db.Model(&agentNetworkTypes.Settings{}).
+		Where("enable_log_collection = ?", true).Count(&m.LogCollectionEnabled).Error; err != nil {
+		return AgentNetworkMetrics{}, fmt.Errorf("count agent network log-collection accounts: %w", err)
+	}
+
+	// COALESCE so an empty ledger scans as 0 instead of NULL.
+	usageRow := db.Model(&agentNetworkTypes.AgentNetworkUsage{}).
+		Select("COALESCE(SUM(input_tokens), 0) AS input_tokens, " +
+			"COALESCE(SUM(output_tokens), 0) AS output_tokens, " +
+			"COALESCE(SUM" + agentNetworkTypes.CostUSDSQLExpr + ", 0) AS cost_usd").Row()
+	if err := usageRow.Scan(&m.InputTokens, &m.OutputTokens, &m.CostUSD); err != nil {
+		return AgentNetworkMetrics{}, fmt.Errorf("scan agent network usage metrics: %w", err)
+	}
+
+	return m, nil
+}
+
+func (s *SqlStore) GetAccountAgentNetworkProviders(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.Provider, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var providers []*agentNetworkTypes.Provider
+	result := tx.Find(&providers, accountIDCondition, accountID)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to get agent network providers from store: %v", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get agent network providers from store")
+	}
+
+	for _, provider := range providers {
+		if err := provider.DecryptSensitiveData(s.fieldEncrypt); err != nil {
+			log.WithContext(ctx).Errorf("failed to decrypt agent network provider %s: %v", provider.ID, err)
+			return nil, status.Errorf(status.Internal, "failed to decrypt agent network provider")
+		}
+	}
+
+	return providers, nil
+}
+
+func (s *SqlStore) GetAgentNetworkProviderByID(ctx context.Context, lockStrength LockingStrength, accountID, providerID string) (*agentNetworkTypes.Provider, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var provider *agentNetworkTypes.Provider
+	result := tx.Take(&provider, accountAndIDQueryCondition, accountID, providerID)
+	if result.Error != nil {
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return nil, status.NewAgentNetworkProviderNotFoundError(providerID)
+		}
+
+		log.WithContext(ctx).Errorf("failed to get agent network provider from store: %v", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get agent network provider from store")
+	}
+
+	if err := provider.DecryptSensitiveData(s.fieldEncrypt); err != nil {
+		log.WithContext(ctx).Errorf("failed to decrypt agent network provider %s: %v", provider.ID, err)
+		return nil, status.Errorf(status.Internal, "failed to decrypt agent network provider")
+	}
+
+	return provider, nil
+}
+
+func (s *SqlStore) SaveAgentNetworkProvider(ctx context.Context, provider *agentNetworkTypes.Provider) error {
+	providerCopy := provider.Copy()
+	if err := providerCopy.EncryptSensitiveData(s.fieldEncrypt); err != nil {
+		log.WithContext(ctx).Errorf("failed to encrypt agent network provider %s: %v", provider.ID, err)
+		return status.Errorf(status.Internal, "failed to encrypt agent network provider")
+	}
+
+	result := s.db.Save(providerCopy)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to save agent network provider to store: %v", result.Error)
+		return status.Errorf(status.Internal, "failed to save agent network provider to store")
+	}
+
+	return nil
+}
+
+func (s *SqlStore) DeleteAgentNetworkProvider(ctx context.Context, accountID, providerID string) error {
+	result := s.db.Delete(&agentNetworkTypes.Provider{}, accountAndIDQueryCondition, accountID, providerID)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to delete agent network provider from store: %v", result.Error)
+		return status.Errorf(status.Internal, "failed to delete agent network provider from store")
+	}
+
+	if result.RowsAffected == 0 {
+		return status.NewAgentNetworkProviderNotFoundError(providerID)
+	}
+
+	return nil
+}
+
+func (s *SqlStore) GetAccountAgentNetworkPolicies(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.Policy, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var policies []*agentNetworkTypes.Policy
+	result := tx.Find(&policies, accountIDCondition, accountID)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to get agent network policies from store: %v", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get agent network policies from store")
+	}
+
+	return policies, nil
+}
+
+func (s *SqlStore) GetAgentNetworkPolicyByID(ctx context.Context, lockStrength LockingStrength, accountID, policyID string) (*agentNetworkTypes.Policy, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var policy *agentNetworkTypes.Policy
+	result := tx.Take(&policy, accountAndIDQueryCondition, accountID, policyID)
+	if result.Error != nil {
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return nil, status.NewAgentNetworkPolicyNotFoundError(policyID)
+		}
+
+		log.WithContext(ctx).Errorf("failed to get agent network policy from store: %v", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get agent network policy from store")
+	}
+
+	return policy, nil
+}
+
+func (s *SqlStore) SaveAgentNetworkPolicy(ctx context.Context, policy *agentNetworkTypes.Policy) error {
+	result := s.db.Save(policy)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to save agent network policy to store: %v", result.Error)
+		return status.Errorf(status.Internal, "failed to save agent network policy to store")
+	}
+
+	return nil
+}
+
+func (s *SqlStore) DeleteAgentNetworkPolicy(ctx context.Context, accountID, policyID string) error {
+	result := s.db.Delete(&agentNetworkTypes.Policy{}, accountAndIDQueryCondition, accountID, policyID)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to delete agent network policy from store: %v", result.Error)
+		return status.Errorf(status.Internal, "failed to delete agent network policy from store")
+	}
+
+	if result.RowsAffected == 0 {
+		return status.NewAgentNetworkPolicyNotFoundError(policyID)
+	}
+
+	return nil
+}
+
+func (s *SqlStore) GetAccountAgentNetworkGuardrails(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.Guardrail, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var guardrails []*agentNetworkTypes.Guardrail
+	result := tx.Find(&guardrails, accountIDCondition, accountID)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to get agent network guardrails from store: %v", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get agent network guardrails from store")
+	}
+
+	return guardrails, nil
+}
+
+func (s *SqlStore) GetAgentNetworkGuardrailByID(ctx context.Context, lockStrength LockingStrength, accountID, guardrailID string) (*agentNetworkTypes.Guardrail, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var guardrail *agentNetworkTypes.Guardrail
+	result := tx.Take(&guardrail, accountAndIDQueryCondition, accountID, guardrailID)
+	if result.Error != nil {
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return nil, status.NewAgentNetworkGuardrailNotFoundError(guardrailID)
+		}
+
+		log.WithContext(ctx).Errorf("failed to get agent network guardrail from store: %v", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get agent network guardrail from store")
+	}
+
+	return guardrail, nil
+}
+
+func (s *SqlStore) SaveAgentNetworkGuardrail(ctx context.Context, guardrail *agentNetworkTypes.Guardrail) error {
+	result := s.db.Save(guardrail)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to save agent network guardrail to store: %v", result.Error)
+		return status.Errorf(status.Internal, "failed to save agent network guardrail to store")
+	}
+
+	return nil
+}
+
+func (s *SqlStore) DeleteAgentNetworkGuardrail(ctx context.Context, accountID, guardrailID string) error {
+	result := s.db.Delete(&agentNetworkTypes.Guardrail{}, accountAndIDQueryCondition, accountID, guardrailID)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to delete agent network guardrail from store: %v", result.Error)
+		return status.Errorf(status.Internal, "failed to delete agent network guardrail from store")
+	}
+
+	if result.RowsAffected == 0 {
+		return status.NewAgentNetworkGuardrailNotFoundError(guardrailID)
+	}
+
+	return nil
+}
+
+// GetAgentNetworkSettings returns the per-account Agent Network
+// settings row. Returns status.NotFound when no row exists.
+func (s *SqlStore) GetAgentNetworkSettings(ctx context.Context, lockStrength LockingStrength, accountID string) (*agentNetworkTypes.Settings, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var settings agentNetworkTypes.Settings
+	result := tx.Take(&settings, "account_id = ?", accountID)
+	if result.Error != nil {
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return nil, status.Errorf(status.NotFound, "agent network settings for account %s not found", accountID)
+		}
+
+		log.WithContext(ctx).Errorf("failed to get agent network settings from store: %v", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get agent network settings from store")
+	}
+
+	return &settings, nil
+}
+
+// GetAllAgentNetworkSettings returns every account's settings row. Used by the
+// access-log retention sweep to learn each account's retention window.
+func (s *SqlStore) GetAllAgentNetworkSettings(ctx context.Context, lockStrength LockingStrength) ([]*agentNetworkTypes.Settings, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var settings []*agentNetworkTypes.Settings
+	if err := tx.Find(&settings).Error; err != nil {
+		log.WithContext(ctx).Errorf("failed to list agent network settings: %v", err)
+		return nil, status.Errorf(status.Internal, "failed to list agent network settings")
+	}
+	return settings, nil
+}
+
+// GetAgentNetworkSettingsByProxyAddress returns every Settings row whose
+// gateway is served by the proxy declaring the given cluster address. Used by
+// cluster-scoped synthesis to find the accounts a shared proxy serves.
+func (s *SqlStore) GetAgentNetworkSettingsByProxyAddress(ctx context.Context, lockStrength LockingStrength, proxyAddress string) ([]*agentNetworkTypes.Settings, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var settings []*agentNetworkTypes.Settings
+	result := tx.Find(&settings, "proxy_address = ?", proxyAddress)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to get agent network settings by proxy address from store: %v", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get agent network settings by proxy address from store")
+	}
+
+	return settings, nil
+}
+
+// GetAgentNetworkSettingsByDomain resolves the single Settings row holding the
+// given endpoint hostname — a point query on the domain unique index. Returns
+// status.NotFound when no account owns the domain.
+func (s *SqlStore) GetAgentNetworkSettingsByDomain(ctx context.Context, lockStrength LockingStrength, domain string) (*agentNetworkTypes.Settings, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var settings agentNetworkTypes.Settings
+	result := tx.Take(&settings, "domain = ?", domain)
+	if result.Error != nil {
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return nil, status.Errorf(status.NotFound, "agent network settings for domain %s not found", domain)
+		}
+
+		log.WithContext(ctx).Errorf("failed to get agent network settings by domain from store: %v", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get agent network settings by domain from store")
+	}
+
+	return &settings, nil
+}
+
+// CreateAgentNetworkSettings inserts a new settings row.
+//
+// Unlike SaveAgentNetworkSettings (an upsert) this is a plain INSERT, and it
+// returns the driver error unwrapped. Both properties are required by the
+// bootstrap allocator: an upsert would overwrite whichever row it collided
+// with, and the allocator classifies the rejection by matching the driver's
+// message — a unique violation on the account primary key means a concurrent
+// bootstrap for the same account won, and one on the domain index means the
+// hostname is taken.
+func (s *SqlStore) CreateAgentNetworkSettings(ctx context.Context, settings *agentNetworkTypes.Settings) error {
+	if err := s.db.Create(settings).Error; err != nil {
+		log.WithContext(ctx).Debugf("failed to create agent network settings: %v", err)
+		return err
+	}
+	return nil
+}
+
+// SaveAgentNetworkSettings upserts the per-account Agent Network
+// settings row.
+func (s *SqlStore) SaveAgentNetworkSettings(ctx context.Context, settings *agentNetworkTypes.Settings) error {
+	result := s.db.Save(settings)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to save agent network settings to store: %v", result.Error)
+		return status.Errorf(status.Internal, "failed to save agent network settings to store")
+	}
+
+	return nil
+}
+
+// DeleteAgentNetworkSettings removes the per-account Agent Network settings
+// row, releasing the account's endpoint. Returns status.NotFound when no row
+// exists. The guards on the delete (no providers, no proxy actively serving
+// the endpoint) live in the manager, which runs this inside a transaction
+// after re-checking them under a row lock.
+func (s *SqlStore) DeleteAgentNetworkSettings(ctx context.Context, accountID string) error {
+	result := s.db.Delete(&agentNetworkTypes.Settings{}, "account_id = ?", accountID)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to delete agent network settings from store: %v", result.Error)
+		return status.Errorf(status.Internal, "failed to delete agent network settings from store")
+	}
+
+	if result.RowsAffected == 0 {
+		return status.Errorf(status.NotFound, "agent network settings for account %s not found", accountID)
+	}
+
+	return nil
+}
+
+// IncrementAgentNetworkConsumption atomically upserts the consumption
+// row keyed on (account, dim_kind, dim_id, window_seconds, window_start)
+// and adds the supplied deltas. Concurrent calls from multiple proxy
+// nodes converge — the database performs the increment server-side via
+// ON CONFLICT DO UPDATE so no read-modify-write race exists.
+func (s *SqlStore) IncrementAgentNetworkConsumption(
+	ctx context.Context,
+	accountID string,
+	kind agentNetworkTypes.ConsumptionDimension,
+	dimID string,
+	windowSeconds int64,
+	windowStart time.Time,
+	tokensIn, tokensOut int64,
+	costUSD float64,
+) error {
+	if accountID == "" || dimID == "" || windowSeconds <= 0 {
+		return status.Errorf(status.InvalidArgument, "account_id, dim_id and window_seconds must be set")
+	}
+	// Deltas are added server-side via ON CONFLICT; a negative or non-finite
+	// value would silently decrement / poison the persisted totals.
+	if tokensIn < 0 || tokensOut < 0 || costUSD < 0 || math.IsNaN(costUSD) || math.IsInf(costUSD, 0) {
+		return status.Errorf(status.InvalidArgument, "consumption deltas must be non-negative and finite")
+	}
+	row := agentNetworkTypes.Consumption{
+		AccountID:      accountID,
+		DimensionKind:  kind,
+		DimensionID:    dimID,
+		WindowSeconds:  windowSeconds,
+		WindowStartUTC: windowStart.UTC(),
+		TokensInput:    tokensIn,
+		TokensOutput:   tokensOut,
+		CostUSD:        costUSD,
+		UpdatedAt:      time.Now().UTC(),
+	}
+	const tbl = "agent_network_consumption"
+	err := s.db.Clauses(clause.OnConflict{
+		Columns: []clause.Column{
+			{Name: "account_id"},
+			{Name: "dim_kind"},
+			{Name: "dim_id"},
+			{Name: "window_seconds"},
+			{Name: "window_start_utc"},
+		},
+		DoUpdates: clause.Assignments(map[string]any{
+			"tokens_input":  gorm.Expr(tbl+".tokens_input + ?", tokensIn),
+			"tokens_output": gorm.Expr(tbl+".tokens_output + ?", tokensOut),
+			"cost_usd":      gorm.Expr(tbl+".cost_usd + ?", costUSD),
+			"updated_at":    time.Now().UTC(),
+		}),
+	}).Create(&row).Error
+	if err != nil {
+		log.WithContext(ctx).Errorf("failed to increment agent network consumption: %v", err)
+		return status.Errorf(status.Internal, "failed to increment agent network consumption")
+	}
+	return nil
+}
+
+// GetAgentNetworkConsumption returns the consumption row for the exact
+// window key. Returns a zero-valued row (not found mapped to zero) so
+// callers can use the result as the headroom basis without nil checks.
+func (s *SqlStore) GetAgentNetworkConsumption(
+	ctx context.Context,
+	lockStrength LockingStrength,
+	accountID string,
+	kind agentNetworkTypes.ConsumptionDimension,
+	dimID string,
+	windowSeconds int64,
+	windowStart time.Time,
+) (*agentNetworkTypes.Consumption, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+	var row agentNetworkTypes.Consumption
+	result := tx.Take(&row,
+		"account_id = ? AND dim_kind = ? AND dim_id = ? AND window_seconds = ? AND window_start_utc = ?",
+		accountID, kind, dimID, windowSeconds, windowStart.UTC())
+	if result.Error != nil {
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return &agentNetworkTypes.Consumption{
+				AccountID:      accountID,
+				DimensionKind:  kind,
+				DimensionID:    dimID,
+				WindowSeconds:  windowSeconds,
+				WindowStartUTC: windowStart.UTC(),
+			}, nil
+		}
+		log.WithContext(ctx).Errorf("failed to get agent network consumption: %v", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get agent network consumption")
+	}
+	return &row, nil
+}
+
+// GetAgentNetworkConsumptionBatch reads many consumption counters for one
+// account in a single query, returning a map keyed by the exact
+// ConsumptionKey. Missing counters are simply absent from the map (callers
+// treat absence as a zero counter). Replaces the per-cap point reads the
+// policy selector previously issued one at a time.
+func (s *SqlStore) GetAgentNetworkConsumptionBatch(
+	ctx context.Context,
+	lockStrength LockingStrength,
+	accountID string,
+	keys []agentNetworkTypes.ConsumptionKey,
+) (map[agentNetworkTypes.ConsumptionKey]*agentNetworkTypes.Consumption, error) {
+	out := make(map[agentNetworkTypes.ConsumptionKey]*agentNetworkTypes.Consumption, len(keys))
+	if len(keys) == 0 {
+		return out, nil
+	}
+
+	// Collect the distinct dim ids, windows and window starts so a single
+	// query scopes to exactly the current windows in play, then filter the
+	// returned rows down to the exact requested keys.
+	wanted := make(map[agentNetworkTypes.ConsumptionKey]struct{}, len(keys))
+	dimSet := make(map[string]struct{})
+	winSet := make(map[int64]struct{})
+	startSet := make(map[time.Time]struct{})
+	for _, k := range keys {
+		k.WindowStartUTC = k.WindowStartUTC.UTC()
+		wanted[k] = struct{}{}
+		dimSet[k.DimID] = struct{}{}
+		winSet[k.WindowSeconds] = struct{}{}
+		startSet[k.WindowStartUTC] = struct{}{}
+	}
+	dimIDs := make([]string, 0, len(dimSet))
+	for d := range dimSet {
+		dimIDs = append(dimIDs, d)
+	}
+	windows := make([]int64, 0, len(winSet))
+	for w := range winSet {
+		windows = append(windows, w)
+	}
+	starts := make([]time.Time, 0, len(startSet))
+	for t := range startSet {
+		starts = append(starts, t)
+	}
+
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+	var rows []*agentNetworkTypes.Consumption
+	result := tx.Find(&rows,
+		"account_id = ? AND dim_id IN ? AND window_seconds IN ? AND window_start_utc IN ?",
+		accountID, dimIDs, windows, starts)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to batch-get agent network consumption: %v", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get agent network consumption")
+	}
+	for _, row := range rows {
+		k := agentNetworkTypes.ConsumptionKey{
+			Kind:           row.DimensionKind,
+			DimID:          row.DimensionID,
+			WindowSeconds:  row.WindowSeconds,
+			WindowStartUTC: row.WindowStartUTC.UTC(),
+		}
+		if _, ok := wanted[k]; ok {
+			out[k] = row
+		}
+	}
+	return out, nil
+}
+
+// IncrementAgentNetworkConsumptionBatch applies the same usage delta to every
+// supplied counter inside a single transaction, so all per-(dimension, window)
+// counters a served request books are written atomically in one round-trip
+// instead of one upsert per counter. Keys are deduplicated by the caller.
+func (s *SqlStore) IncrementAgentNetworkConsumptionBatch(
+	ctx context.Context,
+	accountID string,
+	keys []agentNetworkTypes.ConsumptionKey,
+	tokensIn, tokensOut int64,
+	costUSD float64,
+) error {
+	if accountID == "" {
+		return status.Errorf(status.InvalidArgument, "account_id must be set")
+	}
+	if tokensIn < 0 || tokensOut < 0 || costUSD < 0 || math.IsNaN(costUSD) || math.IsInf(costUSD, 0) {
+		return status.Errorf(status.InvalidArgument, "consumption deltas must be non-negative and finite")
+	}
+	if len(keys) == 0 {
+		return nil
+	}
+
+	const tbl = "agent_network_consumption"
+	err := s.db.Transaction(func(tx *gorm.DB) error {
+		for _, k := range keys {
+			if k.DimID == "" || k.WindowSeconds <= 0 {
+				return status.Errorf(status.InvalidArgument, "dim_id and window_seconds must be set")
+			}
+			now := time.Now().UTC()
+			row := agentNetworkTypes.Consumption{
+				AccountID:      accountID,
+				DimensionKind:  k.Kind,
+				DimensionID:    k.DimID,
+				WindowSeconds:  k.WindowSeconds,
+				WindowStartUTC: k.WindowStartUTC.UTC(),
+				TokensInput:    tokensIn,
+				TokensOutput:   tokensOut,
+				CostUSD:        costUSD,
+				UpdatedAt:      now,
+			}
+			if err := tx.Clauses(clause.OnConflict{
+				Columns: []clause.Column{
+					{Name: "account_id"},
+					{Name: "dim_kind"},
+					{Name: "dim_id"},
+					{Name: "window_seconds"},
+					{Name: "window_start_utc"},
+				},
+				DoUpdates: clause.Assignments(map[string]any{
+					"tokens_input":  gorm.Expr(tbl+".tokens_input + ?", tokensIn),
+					"tokens_output": gorm.Expr(tbl+".tokens_output + ?", tokensOut),
+					"cost_usd":      gorm.Expr(tbl+".cost_usd + ?", costUSD),
+					"updated_at":    now,
+				}),
+			}).Create(&row).Error; err != nil {
+				return err
+			}
+		}
+		return nil
+	})
+	if err != nil {
+		log.WithContext(ctx).Errorf("failed to batch-increment agent network consumption: %v", err)
+		return status.Errorf(status.Internal, "failed to increment agent network consumption")
+	}
+	return nil
+}
+
+// ListAgentNetworkConsumption returns every consumption row recorded
+// for the account, ordered by window_start descending. Backs the
+// dashboard's basic counter view.
+func (s *SqlStore) ListAgentNetworkConsumption(
+	ctx context.Context,
+	lockStrength LockingStrength,
+	accountID string,
+) ([]*agentNetworkTypes.Consumption, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+	var rows []*agentNetworkTypes.Consumption
+	result := tx.
+		Order("window_start_utc DESC").
+		Find(&rows, accountIDCondition, accountID)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to list agent network consumption: %v", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to list agent network consumption")
+	}
+	return rows, nil
+}
+
+// GetAccountAgentNetworkBudgetRules returns every account-level budget rule for
+// the account.
+func (s *SqlStore) GetAccountAgentNetworkBudgetRules(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.AccountBudgetRule, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var rules []*agentNetworkTypes.AccountBudgetRule
+	result := tx.Find(&rules, accountIDCondition, accountID)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to get agent network budget rules from store: %v", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get agent network budget rules from store")
+	}
+
+	return rules, nil
+}
+
+// GetAgentNetworkBudgetRuleByID returns a single budget rule scoped to the
+// account, or a NotFound error.
+func (s *SqlStore) GetAgentNetworkBudgetRuleByID(ctx context.Context, lockStrength LockingStrength, accountID, ruleID string) (*agentNetworkTypes.AccountBudgetRule, error) {
+	tx := s.db
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var rule *agentNetworkTypes.AccountBudgetRule
+	result := tx.Take(&rule, accountAndIDQueryCondition, accountID, ruleID)
+	if result.Error != nil {
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return nil, status.NewAgentNetworkBudgetRuleNotFoundError(ruleID)
+		}
+
+		log.WithContext(ctx).Errorf("failed to get agent network budget rule from store: %v", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get agent network budget rule from store")
+	}
+
+	return rule, nil
+}
+
+// SaveAgentNetworkBudgetRule upserts a budget rule.
+func (s *SqlStore) SaveAgentNetworkBudgetRule(ctx context.Context, rule *agentNetworkTypes.AccountBudgetRule) error {
+	result := s.db.Save(rule)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to save agent network budget rule to store: %v", result.Error)
+		return status.Errorf(status.Internal, "failed to save agent network budget rule to store")
+	}
+
+	return nil
+}
+
+// DeleteAgentNetworkBudgetRule removes a budget rule scoped to the account.
+func (s *SqlStore) DeleteAgentNetworkBudgetRule(ctx context.Context, accountID, ruleID string) error {
+	result := s.db.Delete(&agentNetworkTypes.AccountBudgetRule{}, accountAndIDQueryCondition, accountID, ruleID)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to delete agent network budget rule from store: %v", result.Error)
+		return status.Errorf(status.Internal, "failed to delete agent network budget rule from store")
+	}
+
+	if result.RowsAffected == 0 {
+		return status.NewAgentNetworkBudgetRuleNotFoundError(ruleID)
+	}
+
+	return nil
+}

management/server/store/sql_store_agentnetwork_accesslog_test.go

@@ -0,0 +1,302 @@
+package store
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
+)
+
+// TestAgentNetworkUsage_RealStore_RoundTrip drives CreateAgentNetworkUsage and
+// CreateAgentNetworkAccessLog through a real sqlite store to prove the schema
+// migrates and the inserts succeed for both a populated (allowed) entry and a
+// stripped (denied) entry.
+func TestAgentNetworkUsage_RealStore_RoundTrip(t *testing.T) {
+	ctx := context.Background()
+	s, cleanup, err := NewTestStoreFromSQL(ctx, "", t.TempDir())
+	require.NoError(t, err, "real sqlite test store must come up")
+	defer cleanup()
+
+	const accountID = "acc-anet-usage-1"
+	now := time.Now().UTC()
+
+	// Populated (allowed) usage row with two authorising groups.
+	usage := &agentNetworkTypes.AgentNetworkUsage{
+		ID:                 "log-allowed-1",
+		AccountID:          accountID,
+		Timestamp:          now,
+		UserID:             "user-alice",
+		ResolvedProviderID: "prov-openai-1",
+		Provider:           "openai",
+		Model:              "gpt-4o",
+		SessionID:          "sess-round-trip-1",
+		InputTokens:        1200,
+		OutputTokens:       640,
+		TotalTokens:        1840,
+		InputCostUSD:       0.0231,
+	}
+	usageGroups := []agentNetworkTypes.AgentNetworkUsageGroup{
+		{UsageID: usage.ID, GroupID: "grp-eng", AccountID: accountID},
+		{UsageID: usage.ID, GroupID: "grp-oncall", AccountID: accountID},
+	}
+	require.NoError(t, s.CreateAgentNetworkUsage(ctx, usage, usageGroups), "populated usage insert must succeed")
+
+	// Stripped (denied / 403) usage row: no provider/model/tokens, no groups.
+	denied := &agentNetworkTypes.AgentNetworkUsage{
+		ID:        "log-denied-1",
+		AccountID: accountID,
+		Timestamp: now,
+		UserID:    "user-bob",
+	}
+	require.NoError(t, s.CreateAgentNetworkUsage(ctx, denied, nil), "stripped usage insert must succeed")
+
+	// Idempotency: re-inserting the same id must not error.
+	require.NoError(t, s.CreateAgentNetworkUsage(ctx, usage, usageGroups), "duplicate usage insert must be idempotent")
+
+	// Access-log row + group children.
+	entry := &agentNetworkTypes.AgentNetworkAccessLog{
+		ID:           "log-allowed-1",
+		AccountID:    accountID,
+		ServiceID:    "agent-net-svc-1",
+		Timestamp:    now,
+		UserID:       "user-alice",
+		StatusCode:   200,
+		Provider:     "openai",
+		Model:        "gpt-4o",
+		SessionID:    "sess-round-trip-1",
+		InputTokens:  1200,
+		OutputTokens: 640,
+		TotalTokens:  1840,
+		InputCostUSD: 0.0231,
+	}
+	entryGroups := []agentNetworkTypes.AgentNetworkAccessLogGroup{
+		{LogID: entry.ID, GroupID: "grp-eng", AccountID: accountID},
+		{LogID: entry.ID, GroupID: "grp-oncall", AccountID: accountID},
+	}
+	require.NoError(t, s.CreateAgentNetworkAccessLog(ctx, entry, entryGroups), "access-log insert must succeed")
+
+	// Read back through the filtered list + verify group hydration.
+	logs, total, err := s.GetAgentNetworkAccessLogs(ctx, LockingStrengthNone, accountID, agentNetworkTypes.AgentNetworkAccessLogFilter{Page: 1, PageSize: 50})
+	require.NoError(t, err, "list must succeed")
+	assert.Equal(t, int64(1), total, "one access-log row expected")
+	require.Len(t, logs, 1)
+	assert.ElementsMatch(t, []string{"grp-eng", "grp-oncall"}, logs[0].GroupIDs, "group ids must hydrate")
+	assert.Equal(t, "sess-round-trip-1", logs[0].SessionID, "session id must persist and read back on the access-log row")
+
+	// Session filter narrows the access-log listing to one conversation.
+	sessionID := "sess-round-trip-1"
+	sessLogs, sessTotal, err := s.GetAgentNetworkAccessLogs(ctx, LockingStrengthNone, accountID,
+		agentNetworkTypes.AgentNetworkAccessLogFilter{Page: 1, PageSize: 50, SessionID: &sessionID})
+	require.NoError(t, err)
+	assert.Equal(t, int64(1), sessTotal, "session filter must match the one row with that session id")
+	require.Len(t, sessLogs, 1)
+	assert.Equal(t, entry.ID, sessLogs[0].ID, "session filter must return the matching log row")
+
+	bogus := "no-such-session"
+	_, emptyTotal, err := s.GetAgentNetworkAccessLogs(ctx, LockingStrengthNone, accountID,
+		agentNetworkTypes.AgentNetworkAccessLogFilter{Page: 1, PageSize: 50, SessionID: &bogus})
+	require.NoError(t, err)
+	assert.Equal(t, int64(0), emptyTotal, "unknown session id must match nothing")
+
+	// Session filter also narrows the always-on usage rows.
+	sessUsage, err := s.GetAgentNetworkUsageRows(ctx, LockingStrengthNone, accountID,
+		agentNetworkTypes.AgentNetworkAccessLogFilter{SessionID: &sessionID})
+	require.NoError(t, err)
+	require.Len(t, sessUsage, 1, "session filter must narrow usage rows to the matching session")
+	assert.Equal(t, "sess-round-trip-1", sessUsage[0].SessionID, "usage row must carry the session id")
+}
+
+// TestAgentNetworkUsageOverview_DailyAggregation drives GetAgentNetworkUsageRows
+// + AggregateUsageByGranularity end-to-end against a real sqlite store, with
+// two rows on the same day and one on another, plus a model filter.
+func TestAgentNetworkUsageOverview_DailyAggregation(t *testing.T) {
+	ctx := context.Background()
+	s, cleanup, err := NewTestStoreFromSQL(ctx, "", t.TempDir())
+	require.NoError(t, err, "real sqlite test store must come up")
+	defer cleanup()
+
+	const accountID = "acc-anet-overview-1"
+	day1 := time.Date(2026, 5, 5, 10, 0, 0, 0, time.UTC)
+	day1b := time.Date(2026, 5, 5, 22, 0, 0, 0, time.UTC)
+	day2 := time.Date(2026, 5, 6, 9, 0, 0, 0, time.UTC)
+
+	mk := func(id string, ts time.Time, model string, in, out int64, cost float64) *agentNetworkTypes.AgentNetworkUsage {
+		return &agentNetworkTypes.AgentNetworkUsage{
+			ID: id, AccountID: accountID, Timestamp: ts, Model: model,
+			InputTokens: in, OutputTokens: out, TotalTokens: in + out, InputCostUSD: cost,
+		}
+	}
+	require.NoError(t, s.CreateAgentNetworkUsage(ctx, mk("u1", day1, "gpt-4o", 100, 50, 0.10), nil))
+	require.NoError(t, s.CreateAgentNetworkUsage(ctx, mk("u2", day1b, "gpt-4o", 200, 80, 0.20), nil))
+	require.NoError(t, s.CreateAgentNetworkUsage(ctx, mk("u3", day2, "claude-3", 10, 5, 0.01), nil))
+
+	rows, err := s.GetAgentNetworkUsageRows(ctx, LockingStrengthNone, accountID, agentNetworkTypes.AgentNetworkAccessLogFilter{})
+	require.NoError(t, err)
+	require.Len(t, rows, 3, "all three usage rows expected")
+
+	buckets := agentNetworkTypes.AggregateUsageByGranularity(rows, agentNetworkTypes.UsageGranularityDay)
+	require.Len(t, buckets, 2, "two distinct days expected")
+	assert.Equal(t, "2026-05-05", buckets[0].PeriodStart, "oldest-first ordering")
+	assert.Equal(t, int64(300), buckets[0].InputTokens, "same-day input tokens summed")
+	assert.Equal(t, int64(130), buckets[0].OutputTokens)
+	assert.InDelta(t, 0.30, buckets[0].TotalCostUSD(), 1e-9, "same-day cost summed")
+	assert.Equal(t, "2026-05-06", buckets[1].PeriodStart)
+	assert.Equal(t, int64(15), buckets[1].TotalTokens)
+
+	// Model filter narrows to a single day.
+	model := "claude-3"
+	filtered, err := s.GetAgentNetworkUsageRows(ctx, LockingStrengthNone, accountID, agentNetworkTypes.AgentNetworkAccessLogFilter{Models: []string{model}})
+	require.NoError(t, err)
+	require.Len(t, filtered, 1, "model filter must narrow rows")
+	assert.Equal(t, "u3", filtered[0].ID)
+}
+
+// TestAgentNetworkAccessLogSessions_RealStore drives GetAgentNetworkAccessLogSessions
+// against a real sqlite store: session grouping + aggregation, recency ordering,
+// singleton groups for session-less requests, session pagination, the model
+// filter narrowing sessions, and aggregate sorting.
+func TestAgentNetworkAccessLogSessions_RealStore(t *testing.T) {
+	ctx := context.Background()
+	s, cleanup, err := NewTestStoreFromSQL(ctx, "", t.TempDir())
+	require.NoError(t, err, "real sqlite test store must come up")
+	defer cleanup()
+
+	const accountID = "acc-anet-sessions-1"
+	base := time.Date(2026, 5, 5, 10, 0, 0, 0, time.UTC)
+	at := func(h int) time.Time { return base.Add(time.Duration(h) * time.Hour) }
+
+	mk := func(id, session, user, provider, model, decision string, ts time.Time, cost float64) *agentNetworkTypes.AgentNetworkAccessLog {
+		return &agentNetworkTypes.AgentNetworkAccessLog{
+			ID: id, AccountID: accountID, ServiceID: "svc", Timestamp: ts,
+			UserID: user, StatusCode: 200, Provider: provider, Model: model,
+			SessionID: session, Decision: decision,
+			InputTokens: 100, OutputTokens: 50, TotalTokens: 150, InputCostUSD: cost,
+		}
+	}
+
+	// Two-request session s1 (alice), a one-request denied session s2 (bob), and
+	// two session-less requests (empty session id) that must each form their own
+	// singleton group.
+	require.NoError(t, s.CreateAgentNetworkAccessLog(ctx, mk("s1-a", "s1", "alice", "openai", "gpt-4o", "allow", at(1), 0.10),
+		[]agentNetworkTypes.AgentNetworkAccessLogGroup{{LogID: "s1-a", GroupID: "grp-eng", AccountID: accountID}}))
+	require.NoError(t, s.CreateAgentNetworkAccessLog(ctx, mk("s1-b", "s1", "alice", "openai", "gpt-4o", "allow", at(2), 0.20),
+		[]agentNetworkTypes.AgentNetworkAccessLogGroup{{LogID: "s1-b", GroupID: "grp-oncall", AccountID: accountID}}))
+	require.NoError(t, s.CreateAgentNetworkAccessLog(ctx, mk("s2-a", "s2", "bob", "anthropic", "claude-3", "deny", at(3), 0.05), nil))
+	require.NoError(t, s.CreateAgentNetworkAccessLog(ctx, mk("se-old", "", "carol", "openai", "o1", "allow", at(0), 0.01), nil))
+	require.NoError(t, s.CreateAgentNetworkAccessLog(ctx, mk("se-new", "", "dave", "mistral", "mistral-large", "allow", at(4), 0.02), nil))
+
+	// Default sort: last activity (MAX timestamp) descending.
+	sessions, total, err := s.GetAgentNetworkAccessLogSessions(ctx, LockingStrengthNone, accountID,
+		agentNetworkTypes.AgentNetworkAccessLogFilter{Page: 1, PageSize: 50})
+	require.NoError(t, err)
+	assert.Equal(t, int64(4), total, "four sessions: s1, s2, and two singletons")
+	require.Len(t, sessions, 4)
+
+	// se-new(t4) > s2(t3) > s1(t2) > se-old(t0)
+	assert.Equal(t, "", sessions[0].SessionID, "newest is a session-less singleton")
+	assert.Equal(t, "se-new", sessions[0].Entries[0].ID)
+	assert.Equal(t, "s2", sessions[1].SessionID)
+	assert.Equal(t, "s1", sessions[2].SessionID)
+	assert.Equal(t, "se-old", sessions[3].Entries[0].ID)
+
+	// s1 aggregation.
+	s1 := sessions[2]
+	assert.Equal(t, 2, s1.RequestCount, "s1 has two requests")
+	assert.Equal(t, int64(300), s1.TotalTokens, "tokens summed across the session")
+	assert.InDelta(t, 0.30, s1.TotalCostUSD(), 1e-9, "cost summed across the session")
+	assert.Equal(t, "alice", s1.UserID)
+	assert.Equal(t, "allow", s1.Decision)
+	// SQLite hands times back in time.Local; normalise to UTC so the instant is
+	// compared, not the (differing) *Location pointer.
+	assert.Equal(t, at(1), s1.StartedAt.UTC(), "started = earliest entry")
+	assert.Equal(t, at(2), s1.EndedAt.UTC(), "ended = latest entry")
+	assert.ElementsMatch(t, []string{"openai"}, s1.Providers)
+	assert.ElementsMatch(t, []string{"gpt-4o"}, s1.Models)
+	assert.ElementsMatch(t, []string{"grp-eng", "grp-oncall"}, s1.GroupIDs, "union of the entries' authorising groups")
+
+	// Denied session rolls up to deny.
+	assert.Equal(t, "deny", sessions[1].Decision, "any denied request makes the session deny")
+
+	// Pagination over sessions: 2 per page.
+	page1, total, err := s.GetAgentNetworkAccessLogSessions(ctx, LockingStrengthNone, accountID,
+		agentNetworkTypes.AgentNetworkAccessLogFilter{Page: 1, PageSize: 2})
+	require.NoError(t, err)
+	assert.Equal(t, int64(4), total, "total still counts all sessions")
+	require.Len(t, page1, 2)
+	assert.Equal(t, "se-new", page1[0].Entries[0].ID)
+	assert.Equal(t, "s2", page1[1].SessionID)
+
+	page2, _, err := s.GetAgentNetworkAccessLogSessions(ctx, LockingStrengthNone, accountID,
+		agentNetworkTypes.AgentNetworkAccessLogFilter{Page: 2, PageSize: 2})
+	require.NoError(t, err)
+	require.Len(t, page2, 2)
+	assert.Equal(t, "s1", page2[0].SessionID)
+	assert.Equal(t, "se-old", page2[1].Entries[0].ID)
+
+	// Model filter narrows to the session(s) with matching entries.
+	model := "claude-3"
+	filtered, fTotal, err := s.GetAgentNetworkAccessLogSessions(ctx, LockingStrengthNone, accountID,
+		agentNetworkTypes.AgentNetworkAccessLogFilter{Page: 1, PageSize: 50, Models: []string{model}})
+	require.NoError(t, err)
+	assert.Equal(t, int64(1), fTotal, "only s2 has a claude-3 request")
+	require.Len(t, filtered, 1)
+	assert.Equal(t, "s2", filtered[0].SessionID)
+
+	// Sort by total session cost, descending: s1 (0.30) leads despite not being
+	// the most recent.
+	byCost, _, err := s.GetAgentNetworkAccessLogSessions(ctx, LockingStrengthNone, accountID,
+		agentNetworkTypes.AgentNetworkAccessLogFilter{Page: 1, PageSize: 50, SortBy: "cost_usd", SortOrder: "desc"})
+	require.NoError(t, err)
+	require.Len(t, byCost, 4)
+	assert.Equal(t, "s1", byCost[0].SessionID, "highest-cost session sorts first")
+}
+
+// TestDeleteOldAgentNetworkAccessLogs verifies the retention sweep removes only
+// access-log rows (and their group children) older than the cutoff, leaving
+// recent rows — and never touching usage records.
+func TestDeleteOldAgentNetworkAccessLogs(t *testing.T) {
+	ctx := context.Background()
+	s, cleanup, err := NewTestStoreFromSQL(ctx, "", t.TempDir())
+	require.NoError(t, err, "real sqlite test store must come up")
+	defer cleanup()
+
+	const accountID = "acc-anet-retention-1"
+	old := time.Now().UTC().AddDate(0, 0, -40)
+	recent := time.Now().UTC().AddDate(0, 0, -1)
+
+	mkLog := func(id string, ts time.Time) (*agentNetworkTypes.AgentNetworkAccessLog, []agentNetworkTypes.AgentNetworkAccessLogGroup) {
+		return &agentNetworkTypes.AgentNetworkAccessLog{
+				ID: id, AccountID: accountID, ServiceID: "svc", Timestamp: ts, StatusCode: 200, Model: "gpt-4o",
+			}, []agentNetworkTypes.AgentNetworkAccessLogGroup{
+				{LogID: id, GroupID: "grp-eng", AccountID: accountID},
+			}
+	}
+	oldEntry, oldGroups := mkLog("old-1", old)
+	recentEntry, recentGroups := mkLog("recent-1", recent)
+	require.NoError(t, s.CreateAgentNetworkAccessLog(ctx, oldEntry, oldGroups))
+	require.NoError(t, s.CreateAgentNetworkAccessLog(ctx, recentEntry, recentGroups))
+	// A usage row for the old request must survive the access-log sweep.
+	require.NoError(t, s.CreateAgentNetworkUsage(ctx, &agentNetworkTypes.AgentNetworkUsage{
+		ID: "old-1", AccountID: accountID, Timestamp: old, Model: "gpt-4o", InputTokens: 10, TotalTokens: 10,
+	}, nil))
+
+	cutoff := time.Now().UTC().AddDate(0, 0, -30)
+	deleted, err := s.DeleteOldAgentNetworkAccessLogs(ctx, accountID, cutoff)
+	require.NoError(t, err)
+	assert.Equal(t, int64(1), deleted, "only the 40-day-old log is deleted")
+
+	logs, total, err := s.GetAgentNetworkAccessLogs(ctx, LockingStrengthNone, accountID, agentNetworkTypes.AgentNetworkAccessLogFilter{Page: 1, PageSize: 50})
+	require.NoError(t, err)
+	assert.Equal(t, int64(1), total, "the recent log remains")
+	require.Len(t, logs, 1)
+	assert.Equal(t, "recent-1", logs[0].ID)
+
+	// Usage is untouched by the access-log retention sweep.
+	usage, err := s.GetAgentNetworkUsageRows(ctx, LockingStrengthNone, accountID, agentNetworkTypes.AgentNetworkAccessLogFilter{})
+	require.NoError(t, err)
+	require.Len(t, usage, 1, "usage record for the deleted log must survive")
+}

management/server/store/sql_store_agentnetwork_budgetrule_test.go

@@ -0,0 +1,112 @@
+package store
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
+)
+
+// TestAgentNetworkBudgetRule_RealStore_RoundTrip is the GC-0 no-mock guard: it
+// drives the budget-rule CRUD through a real sqlite store and asserts the full
+// object — targets and the reused PolicyLimits cap shape — survives the
+// save → gorm/JSON serialize → reload round-trip, then that delete removes it
+// and a second delete reports NotFound.
+func TestAgentNetworkBudgetRule_RealStore_RoundTrip(t *testing.T) {
+	ctx := context.Background()
+	s, cleanup, err := NewTestStoreFromSQL(ctx, "", t.TempDir())
+	require.NoError(t, err, "real sqlite test store must come up")
+	defer cleanup()
+
+	const accountID = "acc-budgetrule-1"
+	rule := agentNetworkTypes.NewAccountBudgetRule(accountID)
+	rule.Name = "eng-monthly"
+	rule.TargetGroups = []string{"grp-eng", "grp-oncall"}
+	rule.TargetUsers = []string{"user-alice"}
+	rule.Limits = agentNetworkTypes.PolicyLimits{
+		TokenLimit: agentNetworkTypes.PolicyTokenLimit{
+			Enabled: true, GroupCap: 100_000, UserCap: 10_000, WindowSeconds: 2_592_000,
+		},
+		BudgetLimit: agentNetworkTypes.PolicyBudgetLimit{
+			Enabled: true, GroupCapUsd: 500, UserCapUsd: 50, WindowSeconds: 2_592_000,
+		},
+	}
+	require.NoError(t, s.SaveAgentNetworkBudgetRule(ctx, rule), "save must succeed")
+
+	got, err := s.GetAgentNetworkBudgetRuleByID(ctx, LockingStrengthNone, accountID, rule.ID)
+	require.NoError(t, err, "get by id must succeed after save")
+	assert.Equal(t, rule.Name, got.Name, "name must round-trip")
+	assert.Equal(t, []string{"grp-eng", "grp-oncall"}, got.TargetGroups, "target groups must round-trip")
+	assert.Equal(t, []string{"user-alice"}, got.TargetUsers, "target users must round-trip")
+	assert.Equal(t, rule.Limits, got.Limits, "the reused PolicyLimits cap shape must round-trip intact")
+	assert.True(t, got.Enabled, "enabled must round-trip")
+
+	list, err := s.GetAccountAgentNetworkBudgetRules(ctx, LockingStrengthNone, accountID)
+	require.NoError(t, err, "list must succeed")
+	require.Len(t, list, 1, "exactly the one saved rule must be listed")
+	assert.Equal(t, rule.ID, list[0].ID, "listed rule id must match")
+
+	require.NoError(t, s.DeleteAgentNetworkBudgetRule(ctx, accountID, rule.ID), "delete must succeed")
+
+	_, err = s.GetAgentNetworkBudgetRuleByID(ctx, LockingStrengthNone, accountID, rule.ID)
+	assert.Error(t, err, "get after delete must report not found")
+
+	err = s.DeleteAgentNetworkBudgetRule(ctx, accountID, rule.ID)
+	assert.Error(t, err, "deleting an absent rule must report not found")
+}
+
+// TestAgentNetworkBudgetRule_RealStore_ScopedByAccount pins that rules are
+// account-scoped: a rule under one account is invisible to another.
+func TestAgentNetworkBudgetRule_RealStore_ScopedByAccount(t *testing.T) {
+	ctx := context.Background()
+	s, cleanup, err := NewTestStoreFromSQL(ctx, "", t.TempDir())
+	require.NoError(t, err)
+	defer cleanup()
+
+	ruleA := agentNetworkTypes.NewAccountBudgetRule("acc-A")
+	require.NoError(t, s.SaveAgentNetworkBudgetRule(ctx, ruleA))
+
+	list, err := s.GetAccountAgentNetworkBudgetRules(ctx, LockingStrengthNone, "acc-B")
+	require.NoError(t, err)
+	assert.Empty(t, list, "account B must not see account A's budget rule")
+
+	_, err = s.GetAgentNetworkBudgetRuleByID(ctx, LockingStrengthNone, "acc-B", ruleA.ID)
+	assert.Error(t, err, "cross-account get by id must not resolve")
+}
+
+// TestAgentNetworkSettings_RealStore_CollectionTogglesRoundTrip pins the GC-0
+// additive settings columns: the three collection toggles default off on a
+// fresh row and survive a save/reload at their set values.
+func TestAgentNetworkSettings_RealStore_CollectionTogglesRoundTrip(t *testing.T) {
+	ctx := context.Background()
+	s, cleanup, err := NewTestStoreFromSQL(ctx, "", t.TempDir())
+	require.NoError(t, err)
+	defer cleanup()
+
+	const accountID = "acc-settings-toggles"
+	require.NoError(t, s.SaveAgentNetworkSettings(ctx, &agentNetworkTypes.Settings{
+		AccountID:    accountID,
+		Domain:       "violet.eu.proxy.netbird.io",
+		ProxyAddress: "eu.proxy.netbird.io",
+	}))
+
+	got, err := s.GetAgentNetworkSettings(ctx, LockingStrengthNone, accountID)
+	require.NoError(t, err)
+	assert.False(t, got.EnableLogCollection, "log collection must default off")
+	assert.False(t, got.EnablePromptCollection, "prompt collection must default off")
+	assert.False(t, got.RedactPii, "redact pii must default off")
+
+	got.EnableLogCollection = true
+	got.EnablePromptCollection = true
+	got.RedactPii = true
+	require.NoError(t, s.SaveAgentNetworkSettings(ctx, got))
+
+	reloaded, err := s.GetAgentNetworkSettings(ctx, LockingStrengthNone, accountID)
+	require.NoError(t, err)
+	assert.True(t, reloaded.EnableLogCollection, "log collection must round-trip on")
+	assert.True(t, reloaded.EnablePromptCollection, "prompt collection must round-trip on")
+	assert.True(t, reloaded.RedactPii, "redact pii must round-trip on")
+}

management/server/store/sql_store_get_account_test.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"net"
 	"net/netip"
+	"os"
+	"runtime"
 	"testing"
 	"time"
 
@@ -11,7 +13,6 @@ import (
 	"github.com/stretchr/testify/require"
 
 	nbdns "github.com/netbirdio/netbird/dns"
-	"github.com/netbirdio/netbird/management/server/integration_reference"
 	resourceTypes "github.com/netbirdio/netbird/management/server/networks/resources/types"
 	routerTypes "github.com/netbirdio/netbird/management/server/networks/routers/types"
 	networkTypes "github.com/netbirdio/netbird/management/server/networks/types"
@@ -19,8 +20,66 @@ import (
 	"github.com/netbirdio/netbird/management/server/posture"
 	"github.com/netbirdio/netbird/management/server/types"
 	"github.com/netbirdio/netbird/route"
+	"github.com/netbirdio/netbird/shared/management/integration_reference"
 )
 
+// TestGetAccount_LoadsCustomDomains verifies GetAccount populates account.Domains.
+// SynthesizePrivateServiceZones depends on this relation to anchor a custom-domain
+// private service's DNS zone; without the preload the relation is empty and the
+// service is silently skipped, so a custom domain never resolves on clients.
+func TestGetAccount_LoadsCustomDomains(t *testing.T) {
+	if runtime.GOOS == "windows" {
+		t.Skip("The SQLite store is not properly supported by Windows yet")
+	}
+
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "", t.TempDir())
+	require.NoError(t, err)
+	defer cleanup()
+
+	assertGetAccountLoadsCustomDomains(t, store)
+}
+
+func TestPostgresql_GetAccount_LoadsCustomDomains(t *testing.T) {
+	if (os.Getenv("CI") == "true" && runtime.GOOS == "darwin") || runtime.GOOS == "windows" {
+		t.Skip("skip CI tests on darwin and windows")
+	}
+
+	t.Setenv("NETBIRD_STORE_ENGINE", string(types.PostgresStoreEngine))
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "", t.TempDir())
+	require.NoError(t, err)
+	t.Cleanup(cleanup)
+
+	assertGetAccountLoadsCustomDomains(t, store)
+}
+
+// assertGetAccountLoadsCustomDomains exercises both the gorm and pgx GetAccount
+// paths: it persists two custom domains and asserts the relation comes back
+// populated, which SynthesizePrivateServiceZones relies on.
+func assertGetAccountLoadsCustomDomains(t *testing.T, store Store) {
+	t.Helper()
+	ctx := context.Background()
+
+	accountID := "acct-custom-domains"
+	require.NoError(t, store.SaveAccount(ctx, newAccountWithId(ctx, accountID, "user-1", "")))
+
+	_, err := store.CreateCustomDomain(ctx, accountID, "example.com", "eu.proxy.netbird.io", true)
+	require.NoError(t, err, "creating the first custom domain must succeed")
+	_, err = store.CreateCustomDomain(ctx, accountID, "apps.acme.io", "us.proxy.netbird.io", false)
+	require.NoError(t, err, "creating the second custom domain must succeed")
+
+	account, err := store.GetAccount(ctx, accountID)
+	require.NoError(t, err)
+	require.Len(t, account.Domains, 2, "GetAccount must preload the account's custom domains")
+
+	byDomain := map[string]string{}
+	for _, d := range account.Domains {
+		require.NotNil(t, d)
+		byDomain[d.Domain] = d.TargetCluster
+	}
+	assert.Equal(t, "eu.proxy.netbird.io", byDomain["example.com"], "custom domain must carry its target cluster")
+	assert.Equal(t, "us.proxy.netbird.io", byDomain["apps.acme.io"], "custom domain must carry its target cluster")
+}
+
 // TestGetAccount_ComprehensiveFieldValidation validates that GetAccount properly loads
 // all fields and nested objects from the database, including deeply nested structures.
 func TestGetAccount_ComprehensiveFieldValidation(t *testing.T) {
@@ -148,7 +207,8 @@ func TestGetAccount_ComprehensiveFieldValidation(t *testing.T) {
 		AccountID: accountID,
 		Key:       "peer-key-1-AAAA",
 		Name:      "Peer 1",
-		IP:        net.ParseIP("100.64.0.1"),
+		IP:        netip.MustParseAddr("100.64.0.1"),
+		IPv6:      netip.MustParseAddr("fd00::1"),
 		Meta: nbpeer.PeerSystemMeta{
 			Hostname:      "peer1.example.com",
 			GoOS:          "linux",
@@ -195,7 +255,8 @@ func TestGetAccount_ComprehensiveFieldValidation(t *testing.T) {
 		AccountID: accountID,
 		Key:       "peer-key-2-BBBB",
 		Name:      "Peer 2",
-		IP:        net.ParseIP("100.64.0.2"),
+		IP:        netip.MustParseAddr("100.64.0.2"),
+		IPv6:      netip.MustParseAddr("fd00::2"),
 		Meta: nbpeer.PeerSystemMeta{
 			Hostname:  "peer2.example.com",
 			GoOS:      "darwin",
@@ -232,7 +293,8 @@ func TestGetAccount_ComprehensiveFieldValidation(t *testing.T) {
 		AccountID: accountID,
 		Key:       "peer-key-3-CCCC",
 		Name:      "Peer 3 (Ephemeral)",
-		IP:        net.ParseIP("100.64.0.3"),
+		IP:        netip.MustParseAddr("100.64.0.3"),
+		IPv6:      netip.MustParseAddr("fd00::3"),
 		Meta: nbpeer.PeerSystemMeta{
 			Hostname: "peer3.example.com",
 			GoOS:     "windows",
@@ -710,7 +772,7 @@ func TestGetAccount_ComprehensiveFieldValidation(t *testing.T) {
 		require.True(t, exists, "Peer 1 should exist")
 		assert.Equal(t, "Peer 1", p1.Name, "Peer 1 name mismatch")
 		assert.Equal(t, "peer-key-1-AAAA", p1.Key, "Peer 1 key mismatch")
-		assert.True(t, p1.IP.Equal(net.ParseIP("100.64.0.1")), "Peer 1 IP mismatch")
+		assert.Equal(t, netip.MustParseAddr("100.64.0.1"), p1.IP, "Peer 1 IP mismatch")
 		assert.Equal(t, userID1, p1.UserID, "Peer 1 user ID mismatch")
 		assert.True(t, p1.SSHEnabled, "Peer 1 SSH should be enabled")
 		assert.Equal(t, "ssh-rsa AAAAB3NzaC1...", p1.SSHKey, "Peer 1 SSH key mismatch")

management/server/store/sql_store_pgx_parity_test.go

@@ -0,0 +1,74 @@
+package store
+
+import (
+	"strings"
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"gorm.io/gorm/schema"
+
+	rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
+)
+
+// TestPgxServiceColumnsMatchGorm guards the Postgres pgx read path against
+// drifting from the gorm model. The SQLite/MySQL gorm path loads rows by struct,
+// so a new column on a model is picked up automatically, but the hand-written
+// pgx SELECT in sql_store.go must be updated by hand. This test fails when a
+// gorm column is missing from the pgx column list, which otherwise silently
+// returns zero-valued on Postgres with no compile error.
+func TestPgxServiceColumnsMatchGorm(t *testing.T) {
+	tests := []struct {
+		name          string
+		model         any
+		selectColumns string
+		// excluded lists gorm columns intentionally not loaded by the pgx path.
+		excluded map[string]struct{}
+	}{
+		{
+			name:          "service",
+			model:         &rpservice.Service{},
+			selectColumns: serviceSelectColumns,
+		},
+		{
+			name:          "target",
+			model:         &rpservice.Target{},
+			selectColumns: targetSelectColumns,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			selected := parseColumnList(tc.selectColumns)
+			for _, col := range gormColumnNames(t, tc.model) {
+				if _, ok := tc.excluded[col]; ok {
+					continue
+				}
+				_, ok := selected[col]
+				assert.Truef(t, ok,
+					"gorm column %q is not read by the Postgres pgx SELECT; add it to %sSelectColumns in sql_store.go (or to the test's excluded set if it is intentionally not loaded)",
+					col, tc.name)
+			}
+		})
+	}
+}
+
+func parseColumnList(cols string) map[string]struct{} {
+	set := make(map[string]struct{})
+	for _, c := range strings.Split(cols, ",") {
+		if c = strings.TrimSpace(c); c != "" {
+			set[c] = struct{}{}
+		}
+	}
+	return set
+}
+
+// gormColumnNames returns the DB column names gorm would migrate for the model,
+// using the same default naming strategy the store configures.
+func gormColumnNames(t *testing.T, model any) []string {
+	t.Helper()
+	sch, err := schema.Parse(model, &sync.Map{}, schema.NamingStrategy{})
+	require.NoError(t, err)
+	return sch.DBNames
+}

management/server/store/sql_store_proxy_clusters_test.go

@@ -0,0 +1,109 @@
+package store
+
+import (
+	"context"
+	"os"
+	"runtime"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	rpproxy "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy"
+)
+
+// TestSqlStore_GetProxyClusters_DerivesOnlineAndType guards the
+// account-visible cluster list against silent regressions in two
+// dimensions:
+//
+//  1. Online derivation: a cluster with one stale and one fresh proxy
+//     is online and counts only the fresh proxy; a cluster whose
+//     proxies all heartbeated outside the 2-min window appears offline
+//     with connected_proxies = 0 (rather than disappearing, which is
+//     what the old query did).
+//  2. Type derivation: a cluster scoped to the calling account is
+//     reported as `account`; a cluster with account_id IS NULL is
+//     reported as `shared`. Clusters scoped to other accounts must not
+//     leak into the result.
+//
+// Capability flags are intentionally not asserted here — they're filled
+// by the manager (handler) layer from the per-cluster capability
+// lookups, not by the store query.
+func TestSqlStore_GetProxyClusters_DerivesOnlineAndType(t *testing.T) {
+	if (os.Getenv("CI") == "true" && runtime.GOOS == "darwin") || runtime.GOOS == "windows" {
+		t.Skip("skip CI tests on darwin and windows")
+	}
+
+	runTestForAllEngines(t, "", func(t *testing.T, store Store) {
+		ctx := context.Background()
+		accountID := "acct-clusters"
+		require.NoError(t, store.SaveAccount(ctx, newAccountWithId(ctx, accountID, "user-1", "")))
+
+		otherAccountID := "acct-other"
+		require.NoError(t, store.SaveAccount(ctx, newAccountWithId(ctx, otherAccountID, "user-2", "")))
+
+		acctID := accountID
+		otherID := otherAccountID
+
+		fresh := time.Now().Add(-30 * time.Second)
+		stale := time.Now().Add(-30 * time.Minute)
+
+		mustSave := func(id, cluster string, accID *string, status string, lastSeen time.Time) {
+			require.NoError(t, store.SaveProxy(ctx, &rpproxy.Proxy{
+				ID:             id,
+				SessionID:      id + "-sess",
+				ClusterAddress: cluster,
+				IPAddress:      "10.0.0.1",
+				AccountID:      accID,
+				LastSeen:       lastSeen,
+				Status:         status,
+			}))
+		}
+
+		// shared-mixed: one fresh + one stale proxy → online, connected=1
+		mustSave("p-shared-fresh", "shared-mixed.netbird.io", nil, rpproxy.StatusConnected, fresh)
+		mustSave("p-shared-stale", "shared-mixed.netbird.io", nil, rpproxy.StatusConnected, stale)
+
+		// shared-offline: only stale proxies → offline, connected=0,
+		// but row must still appear (this is the new semantic — old
+		// query would have dropped it entirely).
+		mustSave("p-shared-off", "shared-offline.netbird.io", nil, rpproxy.StatusConnected, stale)
+
+		// account-online: BYOP cluster owned by acctID, fresh
+		mustSave("p-acct-fresh", "byop.acct.example", &acctID, rpproxy.StatusConnected, fresh)
+
+		// other-account: must not surface for acctID
+		mustSave("p-other", "byop.other.example", &otherID, rpproxy.StatusConnected, fresh)
+
+		clusters, err := store.GetProxyClusters(ctx, accountID)
+		require.NoError(t, err)
+
+		byAddr := map[string]rpproxy.Cluster{}
+		for _, c := range clusters {
+			byAddr[c.Address] = c
+		}
+
+		assert.NotContains(t, byAddr, "byop.other.example",
+			"another account's BYOP cluster must not leak into this account's listing")
+
+		require.Contains(t, byAddr, "shared-mixed.netbird.io")
+		mixed := byAddr["shared-mixed.netbird.io"]
+		assert.Equal(t, rpproxy.ClusterTypeShared, mixed.Type, "shared cluster (account_id IS NULL) must be reported as Type=shared")
+		assert.True(t, mixed.Online, "cluster with a fresh proxy must be online")
+		assert.Equal(t, 1, mixed.ConnectedProxies, "connected_proxies must count only fresh proxies; the stale one should not bump the count")
+
+		require.Contains(t, byAddr, "shared-offline.netbird.io",
+			"offline clusters must still appear so the dashboard can render them — the old GetActiveProxyClusters would have dropped this row, which is the regression this test guards against")
+		offline := byAddr["shared-offline.netbird.io"]
+		assert.Equal(t, rpproxy.ClusterTypeShared, offline.Type)
+		assert.False(t, offline.Online, "no fresh heartbeat → offline")
+		assert.Equal(t, 0, offline.ConnectedProxies, "no fresh proxies → connected_proxies=0")
+
+		require.Contains(t, byAddr, "byop.acct.example")
+		acct := byAddr["byop.acct.example"]
+		assert.Equal(t, rpproxy.ClusterTypeAccount, acct.Type, "BYOP cluster owned by the account must be reported as Type=account")
+		assert.True(t, acct.Online)
+		assert.Equal(t, 1, acct.ConnectedProxies)
+	})
+}

management/server/store/sql_store_proxy_disconnect_test.go

@@ -0,0 +1,156 @@
+package store
+
+import (
+	"context"
+	"os"
+	"runtime"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	rpproxy "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy"
+)
+
+// TestSqlStore_DisconnectAllProxies guards the administrative
+// force-disconnect helper:
+//
+//  1. Every proxy that is not already disconnected is marked
+//     disconnected regardless of its session ID (unlike
+//     DisconnectProxy, which is session-guarded).
+//  2. Rows that are already disconnected are left untouched, so their
+//     original disconnected_at is preserved and the returned count
+//     reflects only the rows that actually changed.
+//  3. last_seen is not modified — the stale-proxy reaper keeps working
+//     off the real last heartbeat.
+func TestSqlStore_DisconnectAllProxies(t *testing.T) {
+	if (os.Getenv("CI") == "true" && runtime.GOOS == "darwin") || runtime.GOOS == "windows" {
+		t.Skip("skip CI tests on darwin and windows")
+	}
+
+	runTestForAllEngines(t, "", func(t *testing.T, store Store) {
+		ctx := context.Background()
+
+		lastSeenFresh := time.Now().Add(-30 * time.Second)
+		lastSeenStale := time.Now().Add(-30 * time.Minute)
+		oldDisconnectedAt := time.Now().Add(-time.Hour)
+
+		accountID := "acct-disconnect"
+		proxies := []*rpproxy.Proxy{
+			{
+				ID:             "p-connected-fresh",
+				SessionID:      "sess-1",
+				ClusterAddress: "cluster-a.example.com",
+				IPAddress:      "10.0.0.1",
+				LastSeen:       lastSeenFresh,
+				Status:         rpproxy.StatusConnected,
+			},
+			{
+				ID:             "p-connected-stale",
+				SessionID:      "sess-2",
+				ClusterAddress: "cluster-b.example.com",
+				IPAddress:      "10.0.0.2",
+				AccountID:      &accountID,
+				LastSeen:       lastSeenStale,
+				Status:         rpproxy.StatusConnected,
+			},
+			{
+				ID:             "p-already-disconnected",
+				SessionID:      "sess-3",
+				ClusterAddress: "cluster-a.example.com",
+				IPAddress:      "10.0.0.3",
+				LastSeen:       lastSeenStale,
+				Status:         rpproxy.StatusDisconnected,
+				DisconnectedAt: &oldDisconnectedAt,
+			},
+		}
+		for _, p := range proxies {
+			require.NoError(t, store.SaveProxy(ctx, p))
+		}
+
+		all, err := store.GetAllProxies(ctx)
+		require.NoError(t, err)
+		require.Len(t, all, 3)
+
+		disconnected, err := store.DisconnectAllProxies(ctx)
+		require.NoError(t, err)
+		assert.Equal(t, int64(2), disconnected)
+
+		all, err = store.GetAllProxies(ctx)
+		require.NoError(t, err)
+		require.Len(t, all, 3)
+
+		byID := make(map[string]*rpproxy.Proxy, len(all))
+		for _, p := range all {
+			byID[p.ID] = p
+		}
+
+		for id, p := range byID {
+			assert.Equal(t, rpproxy.StatusDisconnected, p.Status, "proxy %s should be disconnected", id)
+			require.NotNil(t, p.DisconnectedAt, "proxy %s should have disconnected_at set", id)
+		}
+
+		// force-marked rows carry a fresh disconnected_at; the untouched row keeps its original one
+		assert.WithinDuration(t, time.Now(), *byID["p-connected-fresh"].DisconnectedAt, 10*time.Second)
+		assert.WithinDuration(t, time.Now(), *byID["p-connected-stale"].DisconnectedAt, 10*time.Second)
+		assert.WithinDuration(t, oldDisconnectedAt, *byID["p-already-disconnected"].DisconnectedAt, time.Second)
+
+		// last_seen is preserved so the stale reaper schedule is unaffected
+		assert.WithinDuration(t, lastSeenFresh, byID["p-connected-fresh"].LastSeen, time.Second)
+		assert.WithinDuration(t, lastSeenStale, byID["p-connected-stale"].LastSeen, time.Second)
+
+		// idempotent: a second run has nothing left to update
+		disconnected, err = store.DisconnectAllProxies(ctx)
+		require.NoError(t, err)
+		assert.Equal(t, int64(0), disconnected)
+	})
+}
+
+func TestSqlStore_UpdateProxyHeartbeatRestoresDisconnectedCurrentSession(t *testing.T) {
+	if (os.Getenv("CI") == "true" && runtime.GOOS == "darwin") || runtime.GOOS == "windows" {
+		t.Skip("skip CI tests on darwin and windows")
+	}
+
+	runTestForAllEngines(t, "", func(t *testing.T, store Store) {
+		ctx := context.Background()
+		proxy := &rpproxy.Proxy{
+			ID:             "p-heartbeat",
+			SessionID:      "sess-heartbeat",
+			ClusterAddress: "cluster-heartbeat.example.com",
+			IPAddress:      "10.0.0.10",
+			LastSeen:       time.Now().Add(-30 * time.Second),
+			Status:         rpproxy.StatusConnected,
+		}
+		require.NoError(t, store.SaveProxy(ctx, proxy))
+
+		disconnected, err := store.DisconnectAllProxies(ctx)
+		require.NoError(t, err)
+		require.Equal(t, int64(1), disconnected)
+
+		require.NoError(t, store.UpdateProxyHeartbeat(ctx, &rpproxy.Proxy{ID: proxy.ID, SessionID: proxy.SessionID}))
+
+		all, err := store.GetAllProxies(ctx)
+		require.NoError(t, err)
+		require.Len(t, all, 1)
+		assert.Equal(t, rpproxy.StatusConnected, all[0].Status)
+		assert.Nil(t, all[0].DisconnectedAt)
+		assert.WithinDuration(t, time.Now(), all[0].LastSeen, 10*time.Second)
+
+		addresses, err := store.GetActiveProxyClusterAddresses(ctx)
+		require.NoError(t, err)
+		assert.Contains(t, addresses, proxy.ClusterAddress)
+	})
+}
+
+func TestSqlStore_GetAllProxies_Empty(t *testing.T) {
+	if (os.Getenv("CI") == "true" && runtime.GOOS == "darwin") || runtime.GOOS == "windows" {
+		t.Skip("skip CI tests on darwin and windows")
+	}
+
+	runTestForAllEngines(t, "", func(t *testing.T, store Store) {
+		all, err := store.GetAllProxies(context.Background())
+		require.NoError(t, err)
+		assert.Empty(t, all)
+	})
+}

management/server/store/sql_store_service_test.go

@@ -0,0 +1,135 @@
+package store
+
+import (
+	"context"
+	"os"
+	"runtime"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
+)
+
+func TestSqlStore_GetAccount_PrivateServiceRoundtrip(t *testing.T) {
+	if os.Getenv("CI") == "true" && (runtime.GOOS == "darwin" || runtime.GOOS == "windows") {
+		t.Skip("skip CI tests on darwin and windows")
+	}
+
+	runTestForAllEngines(t, "", func(t *testing.T, store Store) {
+		ctx := context.Background()
+		account := newAccountWithId(ctx, "account_private_svc", "testuser", "")
+		require.NoError(t, store.SaveAccount(ctx, account))
+
+		svc := &rpservice.Service{
+			ID:           "svc-private",
+			AccountID:    account.Id,
+			Name:         "private-svc",
+			Domain:       "private.example",
+			ProxyCluster: "cluster.example",
+			Enabled:      true,
+			Mode:         rpservice.ModeHTTP,
+			Private:      true,
+			AccessGroups: []string{"grp-admins", "grp-ops"},
+		}
+		require.NoError(t, store.CreateService(ctx, svc))
+
+		loaded, err := store.GetAccount(ctx, account.Id)
+		require.NoError(t, err)
+		require.Len(t, loaded.Services, 1)
+
+		got := loaded.Services[0]
+		assert.True(t, got.Private)
+		assert.Equal(t, []string{"grp-admins", "grp-ops"}, got.AccessGroups)
+	})
+}
+
+// TestSqlStore_GetAccount_ServiceTargetOptionsRoundtrip guards the Postgres pgx
+// read path (getServices) against silently dropping columns present on the gorm
+// model. Before the fix these fields loaded correctly on SQLite but came back
+// zero-valued on Postgres because the hand-written SELECT and scan omitted them.
+func TestSqlStore_GetAccount_ServiceTargetOptionsRoundtrip(t *testing.T) {
+	if os.Getenv("CI") == "true" && (runtime.GOOS == "darwin" || runtime.GOOS == "windows") {
+		t.Skip("skip CI tests on darwin and windows")
+	}
+
+	runTestForAllEngines(t, "", func(t *testing.T, store Store) {
+		ctx := context.Background()
+		account := newAccountWithId(ctx, "account_svc_opts", "testuser", "")
+		require.NoError(t, store.SaveAccount(ctx, account))
+
+		renewedAt := time.Now().UTC().Truncate(time.Second)
+		targetPath := "/api"
+		svc := &rpservice.Service{
+			ID:        "svc-opts",
+			AccountID: account.Id,
+			Name:      "opts-svc",
+			Domain:    "opts.example",
+			Enabled:   true,
+			Mode:      rpservice.ModeHTTP,
+			Restrictions: rpservice.AccessRestrictions{
+				AllowedCIDRs:     []string{"10.0.0.0/8"},
+				BlockedCountries: []string{"XX"},
+				CrowdSecMode:     "block",
+			},
+			Meta: rpservice.Meta{
+				LastRenewedAt: &renewedAt,
+			},
+			Targets: []*rpservice.Target{
+				{
+					AccountID:     account.Id,
+					ServiceID:     "svc-opts",
+					Path:          &targetPath,
+					Host:          "backend.internal",
+					Port:          8080,
+					Protocol:      "http",
+					TargetId:      "tgt-1",
+					Enabled:       true,
+					ProxyProtocol: true,
+					Options: rpservice.TargetOptions{
+						SkipTLSVerify:           true,
+						RequestTimeout:          30 * time.Second,
+						SessionIdleTimeout:      5 * time.Minute,
+						PathRewrite:             rpservice.PathRewritePreserve,
+						CustomHeaders:           map[string]string{"X-Foo": "bar"},
+						DirectUpstream:          true,
+						CaptureMaxRequestBytes:  1024,
+						CaptureMaxResponseBytes: 2048,
+						CaptureContentTypes:     []string{"application/json"},
+						AgentNetwork:            true,
+						DisableAccessLog:        true,
+					},
+				},
+			},
+		}
+		require.NoError(t, store.CreateService(ctx, svc))
+
+		loaded, err := store.GetAccount(ctx, account.Id)
+		require.NoError(t, err)
+		require.Len(t, loaded.Services, 1)
+
+		got := loaded.Services[0]
+		assert.Equal(t, []string{"10.0.0.0/8"}, got.Restrictions.AllowedCIDRs, "restrictions allowed CIDRs")
+		assert.Equal(t, []string{"XX"}, got.Restrictions.BlockedCountries, "restrictions blocked countries")
+		assert.Equal(t, "block", got.Restrictions.CrowdSecMode, "restrictions crowdsec mode")
+		require.NotNil(t, got.Meta.LastRenewedAt, "meta last renewed at")
+		assert.WithinDuration(t, renewedAt, *got.Meta.LastRenewedAt, time.Second, "meta last renewed at")
+
+		require.Len(t, got.Targets, 1)
+		tg := got.Targets[0]
+		assert.True(t, tg.ProxyProtocol, "target proxy protocol")
+		assert.True(t, tg.Options.SkipTLSVerify, "options skip TLS verify")
+		assert.Equal(t, 30*time.Second, tg.Options.RequestTimeout, "options request timeout")
+		assert.Equal(t, 5*time.Minute, tg.Options.SessionIdleTimeout, "options session idle timeout")
+		assert.Equal(t, rpservice.PathRewritePreserve, tg.Options.PathRewrite, "options path rewrite")
+		assert.Equal(t, map[string]string{"X-Foo": "bar"}, tg.Options.CustomHeaders, "options custom headers")
+		assert.True(t, tg.Options.DirectUpstream, "options direct upstream")
+		assert.Equal(t, int64(1024), tg.Options.CaptureMaxRequestBytes, "options capture max request bytes")
+		assert.Equal(t, int64(2048), tg.Options.CaptureMaxResponseBytes, "options capture max response bytes")
+		assert.Equal(t, []string{"application/json"}, tg.Options.CaptureContentTypes, "options capture content types")
+		assert.True(t, tg.Options.AgentNetwork, "options agent network")
+		assert.True(t, tg.Options.DisableAccessLog, "options disable access log")
+	})
+}

management/server/store/sql_store_test.go

@@ -6,10 +6,10 @@ import (
 	b64 "encoding/base64"
 	"encoding/binary"
 	"fmt"
-	"math/rand"
 	"net"
 	"net/netip"
 	"os"
+	"reflect"
 	"runtime"
 	"sort"
 	"sync"
@@ -35,6 +35,7 @@ import (
 	"github.com/netbirdio/netbird/management/server/util"
 	nbroute "github.com/netbirdio/netbird/route"
 	"github.com/netbirdio/netbird/shared/management/status"
+	"github.com/netbirdio/netbird/shared/testing_helpers"
 	"github.com/netbirdio/netbird/util/crypt"
 )
 
@@ -46,6 +47,7 @@ func runTestForAllEngines(t *testing.T, testDataFile string, f func(t *testing.T
 		}
 		t.Setenv("NETBIRD_STORE_ENGINE", string(engine))
 		store, cleanUp, err := NewTestStoreFromSQL(context.Background(), testDataFile, t.TempDir())
+		assert.NoError(t, err, "engine: ", string(engine))
 		t.Cleanup(cleanUp)
 		assert.NoError(t, err)
 		t.Run(string(engine), func(t *testing.T) {
@@ -92,13 +94,14 @@ func runLargeTest(t *testing.T, store Store) {
 	account.SetupKeys[setupKey.Key] = setupKey
 	const numPerAccount = 6000
 	for n := 0; n < numPerAccount; n++ {
-		netIP := randomIPv4()
+		netIP := sequentialIPv4(n)
 		peerID := fmt.Sprintf("%s-peer-%d", account.Id, n)
+		addr, _ := netip.AddrFromSlice(netIP)
 
 		peer := &nbpeer.Peer{
 			ID:         peerID,
 			Key:        peerID,
-			IP:         netIP,
+			IP:         addr.Unmap(),
 			Name:       peerID,
 			DNSLabel:   peerID,
 			UserID:     "testuser",
@@ -215,12 +218,12 @@ func runLargeTest(t *testing.T, store Store) {
 	}
 }
 
-func randomIPv4() net.IP {
-	rand.New(rand.NewSource(time.Now().UnixNano()))
+// sequentialIPv4 returns a unique IPv4 address for the given index, avoiding
+// the random collisions that would otherwise violate the unique (account_id, ip)
+// index when generating a large number of peers.
+func sequentialIPv4(n int) net.IP {
 	b := make([]byte, 4)
-	for i := range b {
-		b[i] = byte(rand.Intn(256))
-	}
+	binary.BigEndian.PutUint32(b, 0x0A000000+uint32(n))
 	return net.IP(b)
 }
 
@@ -235,7 +238,8 @@ func Test_SaveAccount(t *testing.T) {
 		account.SetupKeys[setupKey.Key] = setupKey
 		account.Peers["testpeer"] = &nbpeer.Peer{
 			Key:    "peerkey",
-			IP:     net.IP{127, 0, 0, 1},
+			IP:     netip.AddrFrom4([4]byte{127, 0, 0, 1}),
+			IPv6:   netip.MustParseAddr("fd00::1"),
 			Meta:   nbpeer.PeerSystemMeta{},
 			Name:   "peer name",
 			Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
@@ -249,7 +253,8 @@ func Test_SaveAccount(t *testing.T) {
 		account2.SetupKeys[setupKey.Key] = setupKey
 		account2.Peers["testpeer2"] = &nbpeer.Peer{
 			Key:    "peerkey2",
-			IP:     net.IP{127, 0, 0, 2},
+			IP:     netip.AddrFrom4([4]byte{127, 0, 0, 2}),
+			IPv6:   netip.MustParseAddr("fd00::2"),
 			Meta:   nbpeer.PeerSystemMeta{},
 			Name:   "peer name 2",
 			Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
@@ -294,6 +299,53 @@ func Test_SaveAccount(t *testing.T) {
 	})
 }
 
+func Test_AccountSettings_SaveAndRetrieve(t *testing.T) {
+	if runtime.GOOS == "windows" {
+		t.Skip("The SQLite store is not properly supported by Windows yet")
+	}
+
+	populateFields := testing_helpers.NewPopulateFields().WithCustomFieldSetter(
+		reflect.PointerTo(reflect.TypeOf(types.ExtraSettings{})), func(this *testing_helpers.PopulateFields, field reflect.Value) (int, error) {
+			es := types.ExtraSettings{}
+			reflectedEs := reflect.ValueOf(&es).Elem()
+			n, err := this.PopulateAll(reflectedEs)
+			if err != nil {
+				return n, err
+			}
+			field.Set(reflectedEs.Addr())
+			return n, nil
+		}).WithCustomFieldSetter(
+		reflect.PointerTo(reflect.TypeOf(types.DashboardFeatures{})), func(this *testing_helpers.PopulateFields, field reflect.Value) (int, error) {
+			t := true
+			df := types.DashboardFeatures{AgentNetwork: &t}
+			reflectedDf := reflect.ValueOf(&df).Elem()
+			field.Set(reflectedDf.Addr())
+			return 1, nil
+		}).WithSkippedTag("gorm", "-")
+
+	runTestForAllEngines(t, "", func(t *testing.T, store Store) {
+		account := newAccountWithId(context.Background(), "account_id", "testuser", "")
+		setupKey, _ := types.GenerateDefaultSetupKey()
+		account.SetupKeys[setupKey.Key] = setupKey
+
+		settings := types.Settings{}
+		numOfExportedFields, err := populateFields.PopulateAll(reflect.ValueOf(&settings).Elem())
+		assert.NoError(t, err)
+		assert.Equal(t, 27, numOfExportedFields)
+		account.Settings = &settings
+
+		err = store.SaveAccount(context.Background(), account)
+		assert.NoError(t, err)
+
+		accountFromDb, err := store.GetAccount(context.Background(), account.Id)
+		assert.NoError(t, err)
+		assert.NotNil(t, accountFromDb)
+		assert.NotNil(t, accountFromDb.Settings)
+
+		assert.True(t, reflect.DeepEqual(&settings, accountFromDb.Settings), "created settings and settings retrieved from the db should match")
+	})
+}
+
 func TestSqlite_DeleteAccount(t *testing.T) {
 	if runtime.GOOS == "windows" {
 		t.Skip("The SQLite store is not properly supported by Windows yet")
@@ -316,7 +368,8 @@ func TestSqlite_DeleteAccount(t *testing.T) {
 	account.SetupKeys[setupKey.Key] = setupKey
 	account.Peers["testpeer"] = &nbpeer.Peer{
 		Key:    "peerkey",
-		IP:     net.IP{127, 0, 0, 1},
+		IP:     netip.AddrFrom4([4]byte{127, 0, 0, 1}),
+		IPv6:   netip.MustParseAddr("fd00::1"),
 		Meta:   nbpeer.PeerSystemMeta{},
 		Name:   "peer name",
 		Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
@@ -487,53 +540,82 @@ func Test_GetAccount(t *testing.T) {
 	})
 }
 
+// TestSqlStore_GetPeerByIP_NotFound pins the not-found semantics the
+// proxy's ValidateTunnelPeer relies on: a tunnel-IP that isn't in the
+// account roster must surface as a NotFound error (not a generic
+// Internal) so callers can distinguish an expected miss from a real
+// store failure. A known IP still resolves.
+func TestSqlStore_GetPeerByIP_NotFound(t *testing.T) {
+	runTestForAllEngines(t, "../testdata/store.sql", func(t *testing.T, store Store) {
+		const accountID = "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+
+		peer, err := store.GetPeerByIP(context.Background(), LockingStrengthNone, accountID, net.ParseIP("192.168.0.0"))
+		require.NoError(t, err, "known tunnel IP must resolve")
+		require.NotNil(t, peer)
+
+		_, err = store.GetPeerByIP(context.Background(), LockingStrengthNone, accountID, net.ParseIP("100.65.0.99"))
+		require.Error(t, err, "unknown tunnel IP must error")
+		parsedErr, ok := status.FromError(err)
+		require.True(t, ok, "error must be a status error")
+		require.Equal(t, status.NotFound, parsedErr.Type(), "tunnel-IP miss must be NotFound, not Internal")
+	})
+}
+
 func TestSqlStore_SavePeer(t *testing.T) {
-	store, cleanUp, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
-	t.Cleanup(cleanUp)
-	assert.NoError(t, err)
+	populateFields := testing_helpers.NewPopulateFields()
 
-	account, err := store.GetAccount(context.Background(), "bf1c8084-ba50-4ce7-9439-34653001fc3b")
-	require.NoError(t, err)
+	runTestForAllEngines(t, "../testdata/store.sql", func(t *testing.T, store Store) {
+		account, err := store.GetAccount(context.Background(), "bf1c8084-ba50-4ce7-9439-34653001fc3b")
+		require.NoError(t, err)
 
-	// save status of non-existing peer
-	peer := &nbpeer.Peer{
-		Key:       "peerkey",
-		ID:        "testpeer",
-		IP:        net.IP{127, 0, 0, 1},
-		Meta:      nbpeer.PeerSystemMeta{Hostname: "testingpeer"},
-		Name:      "peer name",
-		Status:    &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
-		CreatedAt: time.Now().UTC(),
-	}
-	ctx := context.Background()
-	err = store.SavePeer(ctx, account.Id, peer)
-	assert.Error(t, err)
-	parsedErr, ok := status.FromError(err)
-	require.True(t, ok)
-	require.Equal(t, status.NotFound, parsedErr.Type(), "should return not found error")
+		metadata := nbpeer.PeerSystemMeta{}
+		reflectedMetadata := reflect.ValueOf(&metadata).Elem()
 
-	// save new status of existing peer
-	account.Peers[peer.ID] = peer
+		numOfFields, err := populateFields.PopulateAll(reflectedMetadata)
+		assert.NoError(t, err)
+		assert.Equal(t, 33, numOfFields)
 
-	err = store.SaveAccount(context.Background(), account)
-	require.NoError(t, err)
+		// save status of non-existing peer
+		peer := &nbpeer.Peer{
+			Key:       "peerkey",
+			ID:        "testpeer",
+			IP:        netip.AddrFrom4([4]byte{127, 0, 0, 1}),
+			IPv6:      netip.MustParseAddr("fd00::1"),
+			Meta:      metadata, //nbpeer.PeerSystemMeta{Hostname: "testingpeer"},
+			Name:      "peer name",
+			Status:    &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
+			CreatedAt: time.Now().UTC(),
+		}
+		ctx := context.Background()
+		err = store.SavePeer(ctx, account.Id, peer)
+		assert.Error(t, err)
+		parsedErr, ok := status.FromError(err)
+		require.True(t, ok)
+		require.Equal(t, status.NotFound, parsedErr.Type(), "should return not found error")
 
-	updatedPeer := peer.Copy()
-	updatedPeer.Status.Connected = false
-	updatedPeer.Meta.Hostname = "updatedpeer"
+		// save new status of existing peer
+		account.Peers[peer.ID] = peer
 
-	err = store.SavePeer(ctx, account.Id, updatedPeer)
-	require.NoError(t, err)
+		err = store.SaveAccount(context.Background(), account)
+		require.NoError(t, err)
 
-	account, err = store.GetAccount(context.Background(), account.Id)
-	require.NoError(t, err)
+		updatedPeer := peer.Copy()
+		updatedPeer.Status.Connected = false
+		updatedPeer.Meta.Hostname = "updatedpeer"
 
-	actual := account.Peers[peer.ID]
-	assert.Equal(t, updatedPeer.Meta, actual.Meta)
-	assert.Equal(t, updatedPeer.Status.Connected, actual.Status.Connected)
-	assert.Equal(t, updatedPeer.Status.LoginExpired, actual.Status.LoginExpired)
-	assert.Equal(t, updatedPeer.Status.RequiresApproval, actual.Status.RequiresApproval)
-	assert.WithinDurationf(t, updatedPeer.Status.LastSeen, actual.Status.LastSeen.UTC(), time.Millisecond, "LastSeen should be equal")
+		err = store.SavePeer(ctx, account.Id, updatedPeer)
+		require.NoError(t, err)
+
+		account, err = store.GetAccount(context.Background(), account.Id)
+		require.NoError(t, err)
+
+		actual := account.Peers[peer.ID]
+		assert.Equal(t, updatedPeer.Meta, actual.Meta)
+		assert.Equal(t, updatedPeer.Status.Connected, actual.Status.Connected)
+		assert.Equal(t, updatedPeer.Status.LoginExpired, actual.Status.LoginExpired)
+		assert.Equal(t, updatedPeer.Status.RequiresApproval, actual.Status.RequiresApproval)
+		assert.WithinDurationf(t, updatedPeer.Status.LastSeen, actual.Status.LastSeen.UTC(), time.Millisecond, "LastSeen should be equal")
+	})
 }
 
 func TestSqlStore_SavePeerStatus(t *testing.T) {
@@ -556,7 +638,8 @@ func TestSqlStore_SavePeerStatus(t *testing.T) {
 	account.Peers["testpeer"] = &nbpeer.Peer{
 		Key:    "peerkey",
 		ID:     "testpeer",
-		IP:     net.IP{127, 0, 0, 1},
+		IP:     netip.AddrFrom4([4]byte{127, 0, 0, 1}),
+		IPv6:   netip.MustParseAddr("fd00::1"),
 		Meta:   nbpeer.PeerSystemMeta{},
 		Name:   "peer name",
 		Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
@@ -592,56 +675,6 @@ func TestSqlStore_SavePeerStatus(t *testing.T) {
 	assert.WithinDurationf(t, newStatus.LastSeen, actual.LastSeen.UTC(), time.Millisecond, "LastSeen should be equal")
 }
 
-func TestSqlStore_SavePeerLocation(t *testing.T) {
-	store, cleanUp, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
-	t.Cleanup(cleanUp)
-	assert.NoError(t, err)
-
-	account, err := store.GetAccount(context.Background(), "bf1c8084-ba50-4ce7-9439-34653001fc3b")
-	require.NoError(t, err)
-
-	peer := &nbpeer.Peer{
-		AccountID: account.Id,
-		ID:        "testpeer",
-		Location: nbpeer.Location{
-			ConnectionIP: net.ParseIP("0.0.0.0"),
-			CountryCode:  "YY",
-			CityName:     "City",
-			GeoNameID:    1,
-		},
-		CreatedAt: time.Now().UTC(),
-		Meta:      nbpeer.PeerSystemMeta{},
-	}
-	// error is expected as peer is not in store yet
-	err = store.SavePeerLocation(context.Background(), account.Id, peer)
-	assert.Error(t, err)
-
-	account.Peers[peer.ID] = peer
-	err = store.SaveAccount(context.Background(), account)
-	require.NoError(t, err)
-
-	peer.Location.ConnectionIP = net.ParseIP("35.1.1.1")
-	peer.Location.CountryCode = "DE"
-	peer.Location.CityName = "Berlin"
-	peer.Location.GeoNameID = 2950159
-
-	err = store.SavePeerLocation(context.Background(), account.Id, account.Peers[peer.ID])
-	assert.NoError(t, err)
-
-	account, err = store.GetAccount(context.Background(), account.Id)
-	require.NoError(t, err)
-
-	actual := account.Peers[peer.ID].Location
-	assert.Equal(t, peer.Location, actual)
-
-	peer.ID = "non-existing-peer"
-	err = store.SavePeerLocation(context.Background(), account.Id, peer)
-	assert.Error(t, err)
-	parsedErr, ok := status.FromError(err)
-	require.True(t, ok)
-	require.Equal(t, status.NotFound, parsedErr.Type(), "should return not found error")
-}
-
 func Test_TestGetAccountByPrivateDomain(t *testing.T) {
 	if runtime.GOOS == "windows" {
 		t.Skip("The SQLite store is not properly supported by Windows yet")
@@ -784,7 +817,8 @@ func newAccount(store Store, id int) error {
 	account.SetupKeys[setupKey.Key] = setupKey
 	account.Peers["p"+str] = &nbpeer.Peer{
 		Key:    "peerkey" + str,
-		IP:     net.IP{127, 0, 0, 1},
+		IP:     netip.AddrFrom4([4]byte{127, 0, 0, 1}),
+		IPv6:   netip.MustParseAddr("fd00::1"),
 		Meta:   nbpeer.PeerSystemMeta{},
 		Name:   "peer name",
 		Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
@@ -823,7 +857,8 @@ func TestPostgresql_SaveAccount(t *testing.T) {
 	account.SetupKeys[setupKey.Key] = setupKey
 	account.Peers["testpeer"] = &nbpeer.Peer{
 		Key:    "peerkey",
-		IP:     net.IP{127, 0, 0, 1},
+		IP:     netip.AddrFrom4([4]byte{127, 0, 0, 1}),
+		IPv6:   netip.MustParseAddr("fd00::1"),
 		Meta:   nbpeer.PeerSystemMeta{},
 		Name:   "peer name",
 		Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
@@ -837,7 +872,8 @@ func TestPostgresql_SaveAccount(t *testing.T) {
 	account2.SetupKeys[setupKey.Key] = setupKey
 	account2.Peers["testpeer2"] = &nbpeer.Peer{
 		Key:    "peerkey2",
-		IP:     net.IP{127, 0, 0, 2},
+		IP:     netip.AddrFrom4([4]byte{127, 0, 0, 2}),
+		IPv6:   netip.MustParseAddr("fd00::2"),
 		Meta:   nbpeer.PeerSystemMeta{},
 		Name:   "peer name 2",
 		Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
@@ -903,7 +939,8 @@ func TestPostgresql_DeleteAccount(t *testing.T) {
 	account.SetupKeys[setupKey.Key] = setupKey
 	account.Peers["testpeer"] = &nbpeer.Peer{
 		Key:    "peerkey",
-		IP:     net.IP{127, 0, 0, 1},
+		IP:     netip.AddrFrom4([4]byte{127, 0, 0, 1}),
+		IPv6:   netip.MustParseAddr("fd00::1"),
 		Meta:   nbpeer.PeerSystemMeta{},
 		Name:   "peer name",
 		Status: &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
@@ -1010,37 +1047,39 @@ func TestSqlite_GetTakenIPs(t *testing.T) {
 
 	takenIPs, err := store.GetTakenIPs(context.Background(), LockingStrengthNone, existingAccountID)
 	require.NoError(t, err)
-	assert.Equal(t, []net.IP{}, takenIPs)
+	assert.Equal(t, []netip.Addr{}, takenIPs)
 
 	peer1 := &nbpeer.Peer{
 		ID:        "peer1",
 		AccountID: existingAccountID,
 		Key:       "key1",
 		DNSLabel:  "peer1",
-		IP:        net.IP{1, 1, 1, 1},
+		IP:        netip.AddrFrom4([4]byte{1, 1, 1, 1}),
+		IPv6:      netip.MustParseAddr("fd00::1:1:1:1"),
 	}
 	err = store.AddPeerToAccount(context.Background(), peer1)
 	require.NoError(t, err)
 
 	takenIPs, err = store.GetTakenIPs(context.Background(), LockingStrengthNone, existingAccountID)
 	require.NoError(t, err)
-	ip1 := net.IP{1, 1, 1, 1}.To16()
-	assert.Equal(t, []net.IP{ip1}, takenIPs)
+	ip1 := netip.AddrFrom4([4]byte{1, 1, 1, 1})
+	assert.Equal(t, []netip.Addr{ip1}, takenIPs)
 
 	peer2 := &nbpeer.Peer{
 		ID:        "peer1second",
 		AccountID: existingAccountID,
 		Key:       "key2",
 		DNSLabel:  "peer1-1",
-		IP:        net.IP{2, 2, 2, 2},
+		IP:        netip.AddrFrom4([4]byte{2, 2, 2, 2}),
+		IPv6:      netip.MustParseAddr("fd00::2:2:2:2"),
 	}
 	err = store.AddPeerToAccount(context.Background(), peer2)
 	require.NoError(t, err)
 
 	takenIPs, err = store.GetTakenIPs(context.Background(), LockingStrengthNone, existingAccountID)
 	require.NoError(t, err)
-	ip2 := net.IP{2, 2, 2, 2}.To16()
-	assert.Equal(t, []net.IP{ip1, ip2}, takenIPs)
+	ip2 := netip.AddrFrom4([4]byte{2, 2, 2, 2})
+	assert.Equal(t, []netip.Addr{ip1, ip2}, takenIPs)
 }
 
 func TestSqlite_GetPeerLabelsInAccount(t *testing.T) {
@@ -1060,7 +1099,8 @@ func TestSqlite_GetPeerLabelsInAccount(t *testing.T) {
 			AccountID: existingAccountID,
 			Key:       "key1",
 			DNSLabel:  "peer1",
-			IP:        net.IP{1, 1, 1, 1},
+			IP:        netip.AddrFrom4([4]byte{1, 1, 1, 1}),
+			IPv6:      netip.MustParseAddr("fd00::1:1:1:1"),
 		}
 		err = store.AddPeerToAccount(context.Background(), peer1)
 		require.NoError(t, err)
@@ -1074,7 +1114,8 @@ func TestSqlite_GetPeerLabelsInAccount(t *testing.T) {
 			AccountID: existingAccountID,
 			Key:       "key2",
 			DNSLabel:  "peer1-1",
-			IP:        net.IP{2, 2, 2, 2},
+			IP:        netip.AddrFrom4([4]byte{2, 2, 2, 2}),
+			IPv6:      netip.MustParseAddr("fd00::2:2:2:2"),
 		}
 		err = store.AddPeerToAccount(context.Background(), peer2)
 		require.NoError(t, err)
@@ -1127,7 +1168,8 @@ func Test_AddPeerWithSameIP(t *testing.T) {
 			ID:        "peer1",
 			AccountID: existingAccountID,
 			Key:       "key1",
-			IP:        net.IP{1, 1, 1, 1},
+			IP:        netip.AddrFrom4([4]byte{1, 1, 1, 1}),
+			IPv6:      netip.MustParseAddr("fd00::1:1:1:1"),
 		}
 		err = store.AddPeerToAccount(context.Background(), peer1)
 		require.NoError(t, err)
@@ -1136,7 +1178,8 @@ func Test_AddPeerWithSameIP(t *testing.T) {
 			ID:        "peer1second",
 			AccountID: existingAccountID,
 			Key:       "key2",
-			IP:        net.IP{1, 1, 1, 1},
+			IP:        netip.AddrFrom4([4]byte{1, 1, 1, 1}),
+			IPv6:      netip.MustParseAddr("fd00::2:2:2:2"),
 		}
 		err = store.AddPeerToAccount(context.Background(), peer2)
 		require.Error(t, err)
@@ -1259,6 +1302,61 @@ func TestSqlite_CreateAndGetObjectInTransaction(t *testing.T) {
 	assert.NoError(t, err)
 }
 
+func TestSqlStore_SaveAccountPersistsAgentNetworkOnly(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+	account, err := store.GetAccount(context.Background(), accountID)
+	require.NoError(t, err)
+	require.False(t, account.Settings.AgentNetworkOnly, "setting should default to false")
+
+	account.Settings.AgentNetworkOnly = true
+	require.NoError(t, store.SaveAccount(context.Background(), account))
+
+	reloaded, err := store.GetAccount(context.Background(), accountID)
+	require.NoError(t, err)
+	require.True(t, reloaded.Settings.AgentNetworkOnly, "setting should survive a save/load round-trip")
+
+	reloaded.Settings.AgentNetworkOnly = false
+	require.NoError(t, store.SaveAccount(context.Background(), reloaded))
+
+	disabled, err := store.GetAccount(context.Background(), accountID)
+	require.NoError(t, err)
+	require.False(t, disabled.Settings.AgentNetworkOnly, "disabling should persist")
+}
+
+func TestSqlStore_SaveAccountPersistsDashboardFeatures(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+	account, err := store.GetAccount(context.Background(), accountID)
+	require.NoError(t, err)
+	require.Nil(t, account.Settings.DashboardFeatures, "dashboard features should default to unset")
+
+	agentNetwork := true
+	account.Settings.DashboardFeatures = &types.DashboardFeatures{AgentNetwork: &agentNetwork}
+	require.NoError(t, store.SaveAccount(context.Background(), account))
+
+	reloaded, err := store.GetAccount(context.Background(), accountID)
+	require.NoError(t, err)
+	require.NotNil(t, reloaded.Settings.DashboardFeatures, "dashboard features should survive a save/load round-trip")
+	require.NotNil(t, reloaded.Settings.DashboardFeatures.AgentNetwork, "agent network flag should be set")
+	require.True(t, *reloaded.Settings.DashboardFeatures.AgentNetwork, "agent network flag should persist as true")
+
+	disabled := false
+	reloaded.Settings.DashboardFeatures = &types.DashboardFeatures{AgentNetwork: &disabled}
+	require.NoError(t, store.SaveAccount(context.Background(), reloaded))
+
+	reloadedDisabled, err := store.GetAccount(context.Background(), accountID)
+	require.NoError(t, err)
+	require.NotNil(t, reloadedDisabled.Settings.DashboardFeatures.AgentNetwork, "agent network flag should remain set")
+	require.False(t, *reloadedDisabled.Settings.DashboardFeatures.AgentNetwork, "explicit false should persist")
+}
+
 func TestSqlStore_GetAccountUsers(t *testing.T) {
 	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
 	t.Cleanup(cleanup)
@@ -2383,7 +2481,7 @@ func TestSqlStore_GetNetworkRouterByID(t *testing.T) {
 	}
 }
 
-func TestSqlStore_SaveNetworkRouter(t *testing.T) {
+func TestSqlStore_CreateNetworkRouter(t *testing.T) {
 	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
 	t.Cleanup(cleanup)
 	require.NoError(t, err)
@@ -2394,7 +2492,7 @@ func TestSqlStore_SaveNetworkRouter(t *testing.T) {
 	netRouter, err := routerTypes.NewNetworkRouter(accountID, networkID, "", []string{"net-router-grp"}, true, 0, true)
 	require.NoError(t, err)
 
-	err = store.SaveNetworkRouter(context.Background(), netRouter)
+	err = store.CreateNetworkRouter(context.Background(), netRouter)
 	require.NoError(t, err)
 
 	savedNetRouter, err := store.GetNetworkRouterByID(context.Background(), LockingStrengthNone, accountID, netRouter.ID)
@@ -2402,6 +2500,39 @@ func TestSqlStore_SaveNetworkRouter(t *testing.T) {
 	require.Equal(t, netRouter, savedNetRouter)
 }
 
+func TestSqlStore_UpdateNetworkRouter(t *testing.T) {
+	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
+	t.Cleanup(cleanup)
+	require.NoError(t, err)
+
+	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+	networkID := "ct286bi7qv930dsrrug0"
+	routerID := "ctc20ji7qv9ck2sebc80"
+
+	netRouter := &routerTypes.NetworkRouter{
+		ID:         routerID,
+		AccountID:  accountID,
+		NetworkID:  networkID,
+		Peer:       "",
+		PeerGroups: []string{"net-router-grp"},
+		Masquerade: true,
+		Metric:     42,
+		Enabled:    true,
+	}
+
+	err = store.UpdateNetworkRouter(context.Background(), netRouter)
+	require.NoError(t, err)
+
+	savedNetRouter, err := store.GetNetworkRouterByID(context.Background(), LockingStrengthNone, accountID, routerID)
+	require.NoError(t, err)
+	require.Equal(t, netRouter, savedNetRouter)
+
+	// Updating a router under a different account must not match any row.
+	netRouter.AccountID = "non-existent-account"
+	err = store.UpdateNetworkRouter(context.Background(), netRouter)
+	require.Error(t, err)
+}
+
 func TestSqlStore_DeleteNetworkRouter(t *testing.T) {
 	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/store.sql", t.TempDir())
 	t.Cleanup(cleanup)
@@ -2640,7 +2771,8 @@ func TestSqlStore_AddPeerToAccount(t *testing.T) {
 		ID:        "peer1",
 		AccountID: accountID,
 		Key:       "key",
-		IP:        net.IP{1, 1, 1, 1},
+		IP:        netip.AddrFrom4([4]byte{1, 1, 1, 1}),
+		IPv6:      netip.MustParseAddr("fd00::1:1:1:1"),
 		Meta: nbpeer.PeerSystemMeta{
 			Hostname:  "hostname",
 			GoOS:      "linux",
@@ -3815,10 +3947,10 @@ func BenchmarkGetAccountPeers(b *testing.B) {
 	}
 }
 
-func intToIPv4(n uint32) net.IP {
-	ip := make(net.IP, 4)
-	binary.BigEndian.PutUint32(ip, n)
-	return ip
+func intToIPv4(n uint32) netip.Addr {
+	var b [4]byte
+	binary.BigEndian.PutUint32(b[:], n)
+	return netip.AddrFrom4(b)
 }
 
 func TestSqlStore_GetPeersByGroupIDs(t *testing.T) {
@@ -3945,7 +4077,8 @@ func TestSqlStore_GetUserIDByPeerKey(t *testing.T) {
 		Key:       peerKey,
 		AccountID: existingAccountID,
 		UserID:    userID,
-		IP:        net.IP{10, 0, 0, 1},
+		IP:        netip.AddrFrom4([4]byte{10, 0, 0, 1}),
+		IPv6:      netip.MustParseAddr("fd00::a00:1"),
 		DNSLabel:  "test-peer-1",
 	}
 
@@ -3982,7 +4115,8 @@ func TestSqlStore_GetUserIDByPeerKey_NoUserID(t *testing.T) {
 		Key:       peerKey,
 		AccountID: existingAccountID,
 		UserID:    "",
-		IP:        net.IP{10, 0, 0, 1},
+		IP:        netip.AddrFrom4([4]byte{10, 0, 0, 1}),
+		IPv6:      netip.MustParseAddr("fd00::a00:1"),
 		DNSLabel:  "test-peer-1",
 	}
 
@@ -4009,7 +4143,8 @@ func TestSqlStore_ApproveAccountPeers(t *testing.T) {
 				AccountID: accountID,
 				DNSLabel:  "peer1.netbird.cloud",
 				Key:       "peer1-key",
-				IP:        net.ParseIP("100.64.0.1"),
+				IP:        netip.MustParseAddr("100.64.0.1"),
+				IPv6:      netip.MustParseAddr("fd00::1"),
 				Status: &nbpeer.PeerStatus{
 					RequiresApproval: true,
 					LastSeen:         time.Now().UTC(),
@@ -4020,7 +4155,8 @@ func TestSqlStore_ApproveAccountPeers(t *testing.T) {
 				AccountID: accountID,
 				DNSLabel:  "peer2.netbird.cloud",
 				Key:       "peer2-key",
-				IP:        net.ParseIP("100.64.0.2"),
+				IP:        netip.MustParseAddr("100.64.0.2"),
+				IPv6:      netip.MustParseAddr("fd00::2"),
 				Status: &nbpeer.PeerStatus{
 					RequiresApproval: true,
 					LastSeen:         time.Now().UTC(),
@@ -4031,7 +4167,8 @@ func TestSqlStore_ApproveAccountPeers(t *testing.T) {
 				AccountID: accountID,
 				DNSLabel:  "peer3.netbird.cloud",
 				Key:       "peer3-key",
-				IP:        net.ParseIP("100.64.0.3"),
+				IP:        netip.MustParseAddr("100.64.0.3"),
+				IPv6:      netip.MustParseAddr("fd00::3"),
 				Status: &nbpeer.PeerStatus{
 					RequiresApproval: false,
 					LastSeen:         time.Now().UTC(),
@@ -4570,3 +4707,55 @@ func TestSqlStore_DeleteZoneDNSRecords(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, 0, len(remainingRecords))
 }
+
+// TestNewSqliteStore_BusyTimeoutApplied opens a fresh SQLite store and verifies
+// that the _busy_timeout DSN parameter took effect at the driver level. Without
+// this, lock contention on the single SQLite connection waits indefinitely on
+// the Go side and can be hidden behind the 5-minute transactionTimeout.
+func TestNewSqliteStore_BusyTimeoutApplied(t *testing.T) {
+	dir := t.TempDir()
+	store, err := NewSqliteStore(context.Background(), dir, nil, true)
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		_ = store.Close(context.Background())
+	})
+
+	sqlDB, err := store.db.DB()
+	require.NoError(t, err)
+	row := sqlDB.QueryRow("PRAGMA busy_timeout")
+	var busyTimeout int
+	require.NoError(t, row.Scan(&busyTimeout))
+	assert.Equal(t, 30000, busyTimeout, "SQLite busy_timeout must be set via DSN so it survives connection recycling")
+}
+
+// TestNewSqliteStore_BusyTimeoutRespectsUserOverride confirms that an operator
+// passing _busy_timeout or its mattn alias _timeout via NB_STORE_ENGINE_SQLITE_FILE
+// wins over our 30s default. This guards the DSN merge logic in NewSqliteStore.
+func TestNewSqliteStore_BusyTimeoutRespectsUserOverride(t *testing.T) {
+	cases := []struct {
+		name     string
+		envFile  string
+		expected int
+	}{
+		{name: "explicit _busy_timeout wins", envFile: "store.db?_busy_timeout=5000", expected: 5000},
+		{name: "alias _timeout wins", envFile: "store.db?_timeout=7000", expected: 7000},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Setenv("NB_STORE_ENGINE_SQLITE_FILE", tc.envFile)
+			dir := t.TempDir()
+			store, err := NewSqliteStore(context.Background(), dir, nil, true)
+			require.NoError(t, err)
+			t.Cleanup(func() {
+				_ = store.Close(context.Background())
+			})
+
+			sqlDB, err := store.db.DB()
+			require.NoError(t, err)
+			row := sqlDB.QueryRow("PRAGMA busy_timeout")
+			var busyTimeout int
+			require.NoError(t, row.Scan(&busyTimeout))
+			assert.Equal(t, tc.expected, busyTimeout)
+		})
+	}
+}

management/server/store/sqlstore_bench_test.go

@@ -344,7 +344,8 @@ func setupBenchmarkDB(b testing.TB) (*SqlStore, func(), string) {
 			ID:        fmt.Sprintf("peer-%d", i),
 			AccountID: accountID,
 			Key:       fmt.Sprintf("peerkey-%d", i),
-			IP:        net.ParseIP(fmt.Sprintf("100.64.0.%d", i+1)),
+			IP:        netip.MustParseAddr(fmt.Sprintf("100.64.0.%d", i+1)),
+			IPv6:      netip.MustParseAddr(fmt.Sprintf("fd00::%d", i+1)),
 			Name:      fmt.Sprintf("peer-name-%d", i),
 			Status:    &nbpeer.PeerStatus{Connected: i%2 == 0, LastSeen: time.Now()},
 		})

management/server/store/store.go

@@ -1,6 +1,6 @@
 package store
 
-//go:generate go run github.com/golang/mock/mockgen -package store -destination=store_mock.go -source=./store.go -build_flags=-mod=mod
+//go:generate go tool mockgen -package store -destination=store_mock.go -source=./store.go -build_flags=-mod=mod
 
 import (
 	"context"
@@ -37,6 +37,7 @@ import (
 	"github.com/netbirdio/netbird/util"
 	"github.com/netbirdio/netbird/util/crypt"
 
+	agentNetworkTypes "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
 	"github.com/netbirdio/netbird/management/server/migration"
 	resourceTypes "github.com/netbirdio/netbird/management/server/networks/resources/types"
 	routerTypes "github.com/netbirdio/netbird/management/server/networks/routers/types"
@@ -114,6 +115,9 @@ type Store interface {
 
 	GetProxyAccessTokenByHashedToken(ctx context.Context, lockStrength LockingStrength, hashedToken types.HashedProxyToken) (*types.ProxyAccessToken, error)
 	GetAllProxyAccessTokens(ctx context.Context, lockStrength LockingStrength) ([]*types.ProxyAccessToken, error)
+	GetProxyAccessTokensByAccountID(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*types.ProxyAccessToken, error)
+	GetProxyAccessTokenByID(ctx context.Context, lockStrength LockingStrength, tokenID string) (*types.ProxyAccessToken, error)
+	IsProxyAccessTokenValid(ctx context.Context, tokenID string) (bool, error)
 	SaveProxyAccessToken(ctx context.Context, token *types.ProxyAccessToken) error
 	RevokeProxyAccessToken(ctx context.Context, tokenID string) error
 	MarkProxyAccessTokenUsed(ctx context.Context, tokenID string) error
@@ -159,12 +163,37 @@ type Store interface {
 	GetPeerByID(ctx context.Context, lockStrength LockingStrength, accountID string, peerID string) (*nbpeer.Peer, error)
 	GetPeersByIDs(ctx context.Context, lockStrength LockingStrength, accountID string, peerIDs []string) (map[string]*nbpeer.Peer, error)
 	GetPeersByGroupIDs(ctx context.Context, accountID string, groupIDs []string) ([]*nbpeer.Peer, error)
+	GetPeerIDsByGroups(ctx context.Context, accountID string, groupIDs []string) ([]string, error)
+	GetGroupIDsByPeerIDs(ctx context.Context, accountID string, peerIDs []string) ([]string, error)
+	GetEmbeddedProxyPeerIDsByCluster(ctx context.Context, accountID string) (map[string][]string, error)
 	GetAccountPeersWithExpiration(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*nbpeer.Peer, error)
 	GetAccountPeersWithInactivity(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*nbpeer.Peer, error)
 	GetAllEphemeralPeers(ctx context.Context, lockStrength LockingStrength) ([]*nbpeer.Peer, error)
 	SavePeer(ctx context.Context, accountID string, peer *nbpeer.Peer) error
 	SavePeerStatus(ctx context.Context, accountID, peerID string, status nbpeer.PeerStatus) error
-	SavePeerLocation(ctx context.Context, accountID string, peer *nbpeer.Peer) error
+	// MarkPeerConnectedIfNewerSession sets the peer to connected with the
+	// given session token, but only when the stored SessionStartedAt is
+	// strictly less than newSessionStartedAt (the sentinel zero counts as
+	// "older"). LastSeen is recorded by the database at the moment the
+	// row is updated — never by the caller — so it always reflects the
+	// real write time even under lock contention.
+	// Returns true when the update happened, false when this stream lost
+	// the race against a newer session.
+	MarkPeerConnectedIfNewerSession(ctx context.Context, accountID, peerID string, newSessionStartedAt int64) (bool, error)
+	// RefreshPeerLastSeen records that a peer was just seen, stamping the
+	// database clock like the other status writers. Connected and
+	// SessionStartedAt are left alone, so this never interferes with the
+	// session-ownership protocol MarkPeerConnectedIfNewerSession implements.
+	// The write only lands when the stored LastSeen is older than
+	// staleBefore, which keeps a caller's throttle atomic under concurrent
+	// requests for the same peer. Returns true when the update happened.
+	RefreshPeerLastSeen(ctx context.Context, accountID, peerID string, staleBefore time.Time) (bool, error)
+	// MarkPeerDisconnectedIfSameSession sets the peer to disconnected and
+	// resets SessionStartedAt to zero, but only when the stored
+	// SessionStartedAt equals the given sessionStartedAt. LastSeen is
+	// recorded by the database. Returns true when the update happened,
+	// false when a newer session has taken over.
+	MarkPeerDisconnectedIfSameSession(ctx context.Context, accountID, peerID string, sessionStartedAt int64) (bool, error)
 	ApproveAccountPeers(ctx context.Context, accountID string) (int, error)
 	DeletePeer(ctx context.Context, accountID string, peerID string) error
 
@@ -185,7 +214,7 @@ type Store interface {
 	SaveNameServerGroup(ctx context.Context, nameServerGroup *dns.NameServerGroup) error
 	DeleteNameServerGroup(ctx context.Context, accountID, nameServerGroupID string) error
 
-	GetTakenIPs(ctx context.Context, lockStrength LockingStrength, accountId string) ([]net.IP, error)
+	GetTakenIPs(ctx context.Context, lockStrength LockingStrength, accountId string) ([]netip.Addr, error)
 	IncrementNetworkSerial(ctx context.Context, accountId string) error
 	GetAccountNetwork(ctx context.Context, lockStrength LockingStrength, accountId string) (*types.Network, error)
 
@@ -210,7 +239,8 @@ type Store interface {
 	GetNetworkRoutersByNetID(ctx context.Context, lockStrength LockingStrength, accountID, netID string) ([]*routerTypes.NetworkRouter, error)
 	GetNetworkRoutersByAccountID(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*routerTypes.NetworkRouter, error)
 	GetNetworkRouterByID(ctx context.Context, lockStrength LockingStrength, accountID, routerID string) (*routerTypes.NetworkRouter, error)
-	SaveNetworkRouter(ctx context.Context, router *routerTypes.NetworkRouter) error
+	CreateNetworkRouter(ctx context.Context, router *routerTypes.NetworkRouter) error
+	UpdateNetworkRouter(ctx context.Context, router *routerTypes.NetworkRouter) error
 	DeleteNetworkRouter(ctx context.Context, accountID, routerID string) error
 
 	GetNetworkResourcesByNetID(ctx context.Context, lockStrength LockingStrength, accountID, netID string) ([]*resourceTypes.NetworkResource, error)
@@ -225,6 +255,7 @@ type Store interface {
 	IsPrimaryAccount(ctx context.Context, accountID string) (bool, string, error)
 	MarkAccountPrimary(ctx context.Context, accountID string) error
 	UpdateAccountNetwork(ctx context.Context, accountID string, ipNet net.IPNet) error
+	UpdateAccountNetworkV6(ctx context.Context, accountID string, ipNet net.IPNet) error
 	GetPolicyRulesByResourceID(ctx context.Context, lockStrength LockingStrength, accountID string, peerID string) ([]*types.PolicyRule, error)
 
 	// SetFieldEncrypt sets the field encryptor for encrypting sensitive user data.
@@ -271,6 +302,7 @@ type Store interface {
 	GetCustomDomain(ctx context.Context, accountID string, domainID string) (*domain.Domain, error)
 	ListFreeDomains(ctx context.Context, accountID string) ([]string, error)
 	ListCustomDomains(ctx context.Context, accountID string) ([]*domain.Domain, error)
+	GetCustomDomainByName(ctx context.Context, domainName string) (*domain.Domain, error)
 	CreateCustomDomain(ctx context.Context, accountID string, domainName string, targetCluster string, validated bool) (*domain.Domain, error)
 	UpdateCustomDomain(ctx context.Context, accountID string, d *domain.Domain) (*domain.Domain, error)
 	DeleteCustomDomain(ctx context.Context, accountID string, domainID string) error
@@ -278,28 +310,135 @@ type Store interface {
 	CreateAccessLog(ctx context.Context, log *accesslogs.AccessLogEntry) error
 	GetAccountAccessLogs(ctx context.Context, lockStrength LockingStrength, accountID string, filter accesslogs.AccessLogFilter) ([]*accesslogs.AccessLogEntry, int64, error)
 	DeleteOldAccessLogs(ctx context.Context, olderThan time.Time) (int64, error)
+	CreateAgentNetworkAccessLog(ctx context.Context, entry *agentNetworkTypes.AgentNetworkAccessLog, groups []agentNetworkTypes.AgentNetworkAccessLogGroup) error
+	CreateAgentNetworkUsage(ctx context.Context, usage *agentNetworkTypes.AgentNetworkUsage, groups []agentNetworkTypes.AgentNetworkUsageGroup) error
+	GetAgentNetworkAccessLogs(ctx context.Context, lockStrength LockingStrength, accountID string, filter agentNetworkTypes.AgentNetworkAccessLogFilter) ([]*agentNetworkTypes.AgentNetworkAccessLog, int64, error)
+	GetAgentNetworkAccessLogSessions(ctx context.Context, lockStrength LockingStrength, accountID string, filter agentNetworkTypes.AgentNetworkAccessLogFilter) ([]*agentNetworkTypes.AgentNetworkAccessLogSession, int64, error)
+	GetAgentNetworkUsageRows(ctx context.Context, lockStrength LockingStrength, accountID string, filter agentNetworkTypes.AgentNetworkAccessLogFilter) ([]*agentNetworkTypes.AgentNetworkUsage, error)
+	DeleteOldAgentNetworkAccessLogs(ctx context.Context, accountID string, olderThan time.Time) (int64, error)
 	GetServiceTargetByTargetID(ctx context.Context, lockStrength LockingStrength, accountID string, targetID string) (*rpservice.Target, error)
 	GetTargetsByServiceID(ctx context.Context, lockStrength LockingStrength, accountID string, serviceID string) ([]*rpservice.Target, error)
 	DeleteTarget(ctx context.Context, accountID string, serviceID string, targetID uint) error
 	DeleteServiceTargets(ctx context.Context, accountID string, serviceID string) error
 
 	SaveProxy(ctx context.Context, proxy *proxy.Proxy) error
-	UpdateProxyHeartbeat(ctx context.Context, proxyID, clusterAddress, ipAddress string) error
+	DisconnectProxy(ctx context.Context, proxyID, sessionID string) error
+	UpdateProxyHeartbeat(ctx context.Context, p *proxy.Proxy) error
 	GetActiveProxyClusterAddresses(ctx context.Context) ([]string, error)
-	GetActiveProxyClusters(ctx context.Context) ([]proxy.Cluster, error)
+	GetActiveProxyClusterAddressesForAccount(ctx context.Context, accountID string) ([]string, error)
+	GetProxyClusters(ctx context.Context, accountID string) ([]proxy.Cluster, error)
 	GetClusterSupportsCustomPorts(ctx context.Context, clusterAddr string) *bool
 	GetClusterRequireSubdomain(ctx context.Context, clusterAddr string) *bool
 	GetClusterSupportsCrowdSec(ctx context.Context, clusterAddr string) *bool
+	GetClusterSupportsPrivate(ctx context.Context, clusterAddr string) *bool
 	CleanupStaleProxies(ctx context.Context, inactivityDuration time.Duration) error
+	GetAllProxies(ctx context.Context) ([]*proxy.Proxy, error)
+	DisconnectAllProxies(ctx context.Context) (int64, error)
+	GetProxyByAccountID(ctx context.Context, accountID string) (*proxy.Proxy, error)
+	CountProxiesByAccountID(ctx context.Context, accountID string) (int64, error)
+	IsClusterAddressConflicting(ctx context.Context, clusterAddress, accountID string) (bool, error)
+	HasActiveProxyAtClusterAddress(ctx context.Context, clusterAddress string) (bool, error)
+	DeleteAccountCluster(ctx context.Context, clusterAddress, accountID string) error
 
 	GetCustomDomainsCounts(ctx context.Context) (total int64, validated int64, err error)
 
+	// GetProxyMetrics returns aggregated proxy / cluster counts for the
+	// self-hosted metrics worker. Self-hosted only — file-based stores
+	// return a zero-valued struct.
+	GetProxyMetrics(ctx context.Context) (ProxyMetrics, error)
+
+	// GetAgentNetworkMetrics returns aggregated agent-network adoption + usage
+	// counts for the self-hosted metrics worker. Self-hosted only — file-based
+	// stores return a zero-valued struct.
+	GetAgentNetworkMetrics(ctx context.Context) (AgentNetworkMetrics, error)
+
 	GetRoutingPeerNetworks(ctx context.Context, accountID, peerID string) ([]string, error)
+
+	// Agent Network persistence (providers, policies, guardrails, settings).
+	GetAllAgentNetworkProviders(ctx context.Context, lockStrength LockingStrength) ([]*agentNetworkTypes.Provider, error)
+	GetAccountAgentNetworkProviders(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.Provider, error)
+	GetAgentNetworkProviderByID(ctx context.Context, lockStrength LockingStrength, accountID, providerID string) (*agentNetworkTypes.Provider, error)
+	SaveAgentNetworkProvider(ctx context.Context, provider *agentNetworkTypes.Provider) error
+	DeleteAgentNetworkProvider(ctx context.Context, accountID, providerID string) error
+	GetAccountAgentNetworkPolicies(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.Policy, error)
+	GetAgentNetworkPolicyByID(ctx context.Context, lockStrength LockingStrength, accountID, policyID string) (*agentNetworkTypes.Policy, error)
+	SaveAgentNetworkPolicy(ctx context.Context, policy *agentNetworkTypes.Policy) error
+	DeleteAgentNetworkPolicy(ctx context.Context, accountID, policyID string) error
+	GetAccountAgentNetworkGuardrails(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.Guardrail, error)
+	GetAgentNetworkGuardrailByID(ctx context.Context, lockStrength LockingStrength, accountID, guardrailID string) (*agentNetworkTypes.Guardrail, error)
+	SaveAgentNetworkGuardrail(ctx context.Context, guardrail *agentNetworkTypes.Guardrail) error
+	DeleteAgentNetworkGuardrail(ctx context.Context, accountID, guardrailID string) error
+	GetAgentNetworkSettings(ctx context.Context, lockStrength LockingStrength, accountID string) (*agentNetworkTypes.Settings, error)
+	GetAllAgentNetworkSettings(ctx context.Context, lockStrength LockingStrength) ([]*agentNetworkTypes.Settings, error)
+	GetAgentNetworkSettingsByProxyAddress(ctx context.Context, lockStrength LockingStrength, proxyAddress string) ([]*agentNetworkTypes.Settings, error)
+	GetAgentNetworkSettingsByDomain(ctx context.Context, lockStrength LockingStrength, domain string) (*agentNetworkTypes.Settings, error)
+	CreateAgentNetworkSettings(ctx context.Context, settings *agentNetworkTypes.Settings) error
+	SaveAgentNetworkSettings(ctx context.Context, settings *agentNetworkTypes.Settings) error
+	DeleteAgentNetworkSettings(ctx context.Context, accountID string) error
+	IncrementAgentNetworkConsumption(ctx context.Context, accountID string, kind agentNetworkTypes.ConsumptionDimension, dimID string, windowSeconds int64, windowStart time.Time, tokensIn, tokensOut int64, costUSD float64) error
+	IncrementAgentNetworkConsumptionBatch(ctx context.Context, accountID string, keys []agentNetworkTypes.ConsumptionKey, tokensIn, tokensOut int64, costUSD float64) error
+	GetAgentNetworkConsumption(ctx context.Context, lockStrength LockingStrength, accountID string, kind agentNetworkTypes.ConsumptionDimension, dimID string, windowSeconds int64, windowStart time.Time) (*agentNetworkTypes.Consumption, error)
+	GetAgentNetworkConsumptionBatch(ctx context.Context, lockStrength LockingStrength, accountID string, keys []agentNetworkTypes.ConsumptionKey) (map[agentNetworkTypes.ConsumptionKey]*agentNetworkTypes.Consumption, error)
+	ListAgentNetworkConsumption(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.Consumption, error)
+	GetAccountAgentNetworkBudgetRules(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*agentNetworkTypes.AccountBudgetRule, error)
+	GetAgentNetworkBudgetRuleByID(ctx context.Context, lockStrength LockingStrength, accountID, ruleID string) (*agentNetworkTypes.AccountBudgetRule, error)
+	SaveAgentNetworkBudgetRule(ctx context.Context, rule *agentNetworkTypes.AccountBudgetRule) error
+	DeleteAgentNetworkBudgetRule(ctx context.Context, accountID, ruleID string) error
+}
+
+// ProxyMetrics aggregates self-hosted proxy + cluster usage signals
+// surfaced to the telemetry payload. Each field is best-effort: when a
+// store cannot answer (e.g. FileStore) all fields are zero.
+type ProxyMetrics struct {
+	// Clusters counts distinct cluster_address values across the proxies
+	// table — every cluster the management server has heard from, online or not.
+	Clusters int64
+	// ClustersBYOP counts distinct cluster_address values that are owned
+	// by an account (account_id IS NOT NULL). These are bring-your-own-proxy
+	// installations as opposed to NetBird-operated shared clusters.
+	ClustersBYOP int64
+	// ClustersPrivate counts distinct cluster_address values where at
+	// least one proxy reported the private capability (embedded
+	// `netbird proxy` running inside a client).
+	ClustersPrivate int64
+	// Proxies is the total number of proxy rows currently persisted.
+	Proxies int64
+	// ProxiesConnected is the subset of proxies whose status is
+	// "connected" AND last_seen falls within the active heartbeat window
+	// (~2 * heartbeat interval). Proxies the controller hasn't pruned
+	// yet but that are visibly stale don't count.
+	ProxiesConnected int64
+}
+
+// AgentNetworkMetrics aggregates self-hosted agent-network adoption + usage
+// signals surfaced to the telemetry payload. Each field is best-effort: when a
+// store cannot answer (e.g. FileStore) all fields are zero.
+type AgentNetworkMetrics struct {
+	// Accounts is the number of distinct accounts with at least one provider
+	// configured (agent-network adoption).
+	Accounts int64
+	// Providers is the total number of configured providers across all accounts.
+	Providers int64
+	// Policies is the total number of agent-network policies across all accounts.
+	Policies int64
+	// BudgetRules is the total number of account-level budget rules ("budget
+	// limits") across all accounts.
+	BudgetRules int64
+	// LogCollectionEnabled is the number of accounts that have agent-network
+	// log collection turned on.
+	LogCollectionEnabled int64
+	// InputTokens / OutputTokens / CostUSD are summed over the always-collected
+	// per-request usage ledger (agent_network_request_usage), independent of the
+	// log-collection toggle. They reflect total metered LLM usage served through
+	// agent networks.
+	InputTokens  int64
+	OutputTokens int64
+	CostUSD      float64
 }
 
 const (
-	postgresDsnEnv       = "NB_STORE_ENGINE_POSTGRES_DSN"
-	postgresDsnEnvLegacy = "NETBIRD_STORE_ENGINE_POSTGRES_DSN"
+	PostgresDsnEnv       = "NB_STORE_ENGINE_POSTGRES_DSN"
+	PostgresDsnEnvLegacy = "NETBIRD_STORE_ENGINE_POSTGRES_DSN"
 	mysqlDsnEnv          = "NB_STORE_ENGINE_MYSQL_DSN"
 	mysqlDsnEnvLegacy    = "NETBIRD_STORE_ENGINE_MYSQL_DSN"
 )
@@ -446,6 +585,9 @@ func getMigrationsPreAuto(ctx context.Context) []migrationFunc {
 		func(db *gorm.DB) error {
 			return migration.MigrateNewField[types.User](ctx, db, "email", "")
 		},
+		func(db *gorm.DB) error {
+			return migration.MigrateNewField[nbpeer.Peer](ctx, db, "peer_status_session_started_at", int64(0))
+		},
 		func(db *gorm.DB) error {
 			return migration.RemoveDuplicatePeerKeys(ctx, db)
 		},
@@ -455,6 +597,33 @@ func getMigrationsPreAuto(ctx context.Context) []migrationFunc {
 		func(db *gorm.DB) error {
 			return migration.CleanupOrphanedResources[domain.Domain, types.Account](ctx, db, "account_id")
 		},
+		func(db *gorm.DB) error {
+			return migration.BackfillPublicIDs[types.Policy](ctx, db)
+		},
+		func(db *gorm.DB) error {
+			return migration.BackfillPublicIDs[types.Group](ctx, db)
+		},
+		func(db *gorm.DB) error {
+			return migration.BackfillPublicIDs[route.Route](ctx, db)
+		},
+		func(db *gorm.DB) error {
+			return migration.BackfillPublicIDs[resourceTypes.NetworkResource](ctx, db)
+		},
+		func(db *gorm.DB) error {
+			return migration.BackfillPublicIDs[routerTypes.NetworkRouter](ctx, db)
+		},
+		func(db *gorm.DB) error {
+			return migration.BackfillPublicIDs[dns.NameServerGroup](ctx, db)
+		},
+		func(db *gorm.DB) error {
+			return migration.BackfillPublicIDs[networkTypes.Network](ctx, db)
+		},
+		func(db *gorm.DB) error {
+			return migration.BackfillPublicIDs[posture.Checks](ctx, db)
+		},
+		func(db *gorm.DB) error {
+			return migration.MigrateAgentNetworkSettingsToDomain(ctx, db)
+		},
 	}
 }
 
@@ -494,6 +663,17 @@ func getMigrationsPostAuto(ctx context.Context) []migrationFunc {
 		func(db *gorm.DB) error {
 			return migration.CreateIndexIfNotExists[nbpeer.Peer](ctx, db, "idx_peers_key_unique", "key")
 		},
+		func(db *gorm.DB) error {
+			return migration.DropIndex[proxy.Proxy](ctx, db, "idx_proxy_account_id_unique")
+		},
+		// Post-auto so the per-bucket cost columns already exist when the legacy
+		// aggregates are folded into them and dropped.
+		func(db *gorm.DB) error {
+			return migration.FoldCostAggregatesIntoBuckets[agentNetworkTypes.AgentNetworkAccessLog](ctx, db)
+		},
+		func(db *gorm.DB) error {
+			return migration.FoldCostAggregatesIntoBuckets[agentNetworkTypes.AgentNetworkUsage](ctx, db)
+		},
 	}
 }
 
@@ -602,7 +782,7 @@ func getSqlStoreEngine(ctx context.Context, store *SqlStore, kind types.Engine)
 }
 
 func newReusedPostgresStore(ctx context.Context, store *SqlStore, kind types.Engine) (*SqlStore, func(), error) {
-	dsn, ok := lookupDSNEnv(postgresDsnEnv, postgresDsnEnvLegacy)
+	dsn, ok := lookupDSNEnv(PostgresDsnEnv, PostgresDsnEnvLegacy)
 	if !ok || dsn == "" {
 		var err error
 		_, dsn, err = testutil.CreatePostgresTestContainer()
@@ -612,7 +792,7 @@ func newReusedPostgresStore(ctx context.Context, store *SqlStore, kind types.Eng
 	}
 
 	if dsn == "" {
-		return nil, nil, fmt.Errorf("%s is not set", postgresDsnEnv)
+		return nil, nil, fmt.Errorf("%s is not set", PostgresDsnEnv)
 	}
 
 	db, err := openDBWithRetry(dsn, kind, 5)